ctree.py 6.74 KB
Newer Older
1
from __future__ import print_function
2
3
4
5
6
7
import sys
try:
    import py_unstools as ut         # import py_unstools package
except ImportError:    
    print("WARNING !!!, failed to import module [py_unstools]",file=sys.stderr)

8
9
10
11
12
13
14
15
16
17
18
19
import numpy as np
from cfalcon import *
from py_unsio import *
from csnapshot import *
# -----------------------------------------------------
#
class CTree:
    """methods imported from CTree C++ class"""
    __tree=None
    __pos=None
    __mass=None
    __vel=None
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
20
    __time=0.0
21
22
# -----------------------------------------------------
#
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
23
    def __init__(self,pos,vel=None,mass=None,fcells=0.9,rsize=0.4,time=0.0,seed=None):
24
25
26
        self.__pos=pos
        self.__mass=mass
        self.__vel=vel
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
27
        self.__time=float(time)
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
28
29
        # initialyse random generator number
        np.random.seed(seed)
30
31
        if (type(pos[0])==np.float32):
            self.__tree=ut.CTreeF(pos,mass,fcells,rsize)
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
32

33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
        else: # assume 64 bits (double)
            self.__tree=ut.CTreeD(pos,mass,fcells,rsize)

# -----------------------------------------------------
#
    def getLevels(self):
        """Return an numpy array with for each particles its level in the octree"""
        ok,levels=self.__tree.get_levels(self.__tree.getNbody())
        return ok,levels

# -----------------------------------------------------
#
    def fastCod(self,threshold=10000):
        """
        compute cod using octree

        Argument:
        threshold > 0 : use this values as number of particles selected for computing cod
        threshold < 0 : use this |values| as a percentage of the total particles for computing cod

        Return:
        
        """
        
        if threshold<0: # it's percentage of nbodies
            threshold = (abs(threshold)*self.__tree.getNbody())/100
        else:
            threshold = min(threshold,self.__tree.getNbody())

62
        #print("Threshold = ",threshold)
63
64
65
66
67
68
69
70

        ok,levels=self.getLevels()
        #idx=np.argsort(levels)
        idx=(-levels).argsort()
        print("idx=",idx.size,idx,levels[idx])

        p=np.reshape(self.__pos,(-1,3))
        p=p[idx[0:threshold],]
71
        #print(p.shape,p.size)
72
73
74
        m=self.__mass[idx[0:threshold]]

        p=np.reshape(p,(-1,))
75
76
        #print(p.shape,p.size)
        #print (p[:,])
77
78
        c=CFalcon() # new falcon object
        ok,rho,hsml=c.getDensity(p,m) # compute density
79
        #print (ok)
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115

        v=None
        if self.__vel is not None:
             v=np.reshape(self.__vel,(-1,3))
             v=v[idx[0:threshold],]
             v=np.reshape(v,(-1,))
        cxv=CSnapshot(None).center(p,v,m*rho)
        
        unso=CunsOut("/home/jcl/x_x","gadget2")
        
        unso.setArrayF("gas","pos",p)
        unso.setArrayF("gas","mass",m)
        unso.setArrayF("gas","rho",rho)
        unso.setArrayF("gas","hsml",hsml)
        unso.save()
        return cxv[0:3]

# -----------------------------------------------------
#
    def fastCod2(self,threshold=10000):
        """
        compute cod using octree

        Argument:
        threshold > 0 : use this values as number of particles selected for computing cod
        threshold < 0 : use this |values| as a percentage of the total particles for computing cod

        Return:
        
        """
        
        if threshold<0: # it's percentage of nbodies
            threshold = (abs(threshold)*self.__tree.getNbody())/100
        else:
            threshold = min(threshold,self.__tree.getNbody())

116
        #print("Threshold = ",threshold)
117
118
119
120
121
122
123
124

        ok,radius=self.__tree.get_closest_distance_to_mesh(self.__tree.getNbody())
        #idx=np.argsort(radius)
        idx=(radius).argsort()
        print("idx=",idx.size,idx,radius[idx])

        p=np.reshape(self.__pos,(-1,3))
        p=p[idx[0:threshold],]
125
        #print(p.shape,p.size)
126
127
128
        m=self.__mass[idx[0:threshold]]

        p=np.reshape(p,(-1,))
129
130
        #print(p.shape,p.size)
        #print (p[:,])
131
132
        c=CFalcon() # new falcon object
        ok,rho,hsml=c.getDensity(p,m) # compute density
133
        #print (ok)
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153

        v=None
        if self.__vel is not None:
             v=np.reshape(self.__vel,(-1,3))
             v=v[idx[0:threshold],]
             v=np.reshape(v,(-1,))
        cxv=CSnapshot(None).center(p,v,m*rho)


        unso=CunsOut("/home/jcl/x_x","gadget2")
        
        unso.setArrayF("gas","pos",p)
        unso.setArrayF("gas","mass",m)
        unso.setArrayF("gas","rho",rho)
        unso.setArrayF("gas","hsml",hsml)
        unso.save()
        return cxv[0:3]

# -----------------------------------------------------
#  
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
154
    def fastCod3(self,threshold=10000,outfile=None):
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
        """
        compute cod using octree

        Argument:
        threshold > 0 : use this values as number of particles selected for computing cod
        threshold < 0 : use this |values| as a percentage of the total particles for computing cod

        Return:
        
        """
        
        if threshold<0: # it's percentage of nbodies
            threshold = (abs(threshold)*self.__tree.getNbody())/100
        else:
            threshold = min(threshold,self.__tree.getNbody())

171
        #print("Threshold = ",threshold)
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186

        ok,levels=self.getLevels() # get octree levels for each particles

        # we randomize indexes bc simulation code save particles processor per processor
        # in kinda block structures
        shuffle=np.arange(levels.size)
        np.random.shuffle(shuffle)

        idx=(-levels[shuffle]).argsort()
        #print("idx=",idx.size,idx,levels[shuffle[idx]])

        shuffle_idx_threshold=shuffle[idx[0:threshold]] # to save some speed

        p=np.reshape(self.__pos,(-1,3)) # convert flat array to [nbody,3]
        p=p[shuffle_idx_threshold,]     # get randomized thresholded indexes for pos
187
        #print(p.shape,p.size)
188
189
190
        m=self.__mass[shuffle_idx_threshold] # get randomized thresholded indexes for mass

        p=np.reshape(p,(-1,)) # flatten pos array
191
192
        #print(p.shape,p.size)
        #print (p[:,])
193
194
        c=CFalcon() # new falcon object
        ok,rho,hsml=c.getDensity(p,m) # compute density
195
        #print (ok)
196
197
198
199
200
201
202

        v=None
        if self.__vel is not None:
             v=np.reshape(self.__vel,(-1,3))
             v=v[shuffle_idx_threshold, ]
             v=np.reshape(v,(-1,))
        cxv=CSnapshot(None).center(p,v,m*rho)
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
203

LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
204
205
206
207
208
209
210
211
        if v is None: # no velocities
            cxv[3]=0.0
            cxv[4]=0.0
            cxv[5]=0.0
        if outfile is not None:
            unso=CunsOut(outfile,"nemo")
            unso.setValueF("time",self.__time)
            unso.setArrayF("all","pos",p)
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
212
213
            if self.__vel is not None:
                unso.setArrayF("all","vel",v)
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
214
215
216
217
218
            unso.setArrayF("all","mass",m)
            unso.setArrayF("all","rho",rho)
            unso.setArrayF("all","hsml",hsml)
            unso.save()
        return cxv