ctree.py 7.54 KB
Newer Older
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
1
#!/usr/bin/python
2
from __future__ import print_function
3
import sys
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
4
import csnapshot as cs
5
6
7
8
9
try:
    import py_unstools as ut         # import py_unstools package
except ImportError:    
    print("WARNING !!!, failed to import module [py_unstools]",file=sys.stderr)

10
import numpy as np
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
11
import cfalcon as cf
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
12
import unsio as unsio
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
13

14
15
16
# -----------------------------------------------------
#
class CTree:
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
17
18
19
    """
    methods imported from CTree C++ class
    """
20
    __tree=None
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
21
    # data given to constructor
22
23
24
    __pos=None
    __mass=None
    __vel=None
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
25
    __time=0.0
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
26
27
28
29
30
31
32
33
34
    # data computed in tree
    __tree_pos=None
    __tree_vel=None
    __tree_mass=None
    __tree_rho=None
    __tree_hsml=None
    #
    # constructor
    # 
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
35
    def __init__(self,pos,vel=None,mass=None,fcells=0.9,rsize=0.4,time=0.0,seed=None):
36
37
38
        self.__pos=pos
        self.__mass=mass
        self.__vel=vel
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
39
        self.__time=float(time)
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
40
41
        # initialyse random generator number
        np.random.seed(seed)
42
43
        if (type(pos[0])==np.float32):
            self.__tree=ut.CTreeF(pos,mass,fcells,rsize)
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
44
            
45
46
        else: # assume 64 bits (double)
            self.__tree=ut.CTreeD(pos,mass,fcells,rsize)
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
47
            
48

LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
49
50
51
52

    #
    # getLevels
    # 
53
54
55
56
57
    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

LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
58
59
60
61

    #
    # fastCod
    # 
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
    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())

79
        #print("Threshold = ",threshold)
80
81
82
83
84
85
86
87

        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],]
88
        #print(p.shape,p.size)
89
90
91
        m=self.__mass[idx[0:threshold]]

        p=np.reshape(p,(-1,))
92
93
        #print(p.shape,p.size)
        #print (p[:,])
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
94
        c=cf.CFalcon() # new falcon object
95
        ok,rho,hsml=c.getDensity(p,m) # compute density
96
        #print (ok)
97
98
99
100
101
102

        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,))
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
103
        cxv=cs.CSnapshot(None).center(p,v,m*rho)
104
        
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
105
        unso=unsio.CunsOut("/home/jcl/x_x","gadget2")
106
107
108
109
110
111
112
113
        
        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
114
115
116
    #
    # fastCod2
    # 
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
    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())

134
        #print("Threshold = ",threshold)
135
136
137
138
139
140
141
142

        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],]
143
        #print(p.shape,p.size)
144
145
146
        m=self.__mass[idx[0:threshold]]

        p=np.reshape(p,(-1,))
147
148
        #print(p.shape,p.size)
        #print (p[:,])
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
149
        c=cf.CFalcon() # new falcon object
150
        ok,rho,hsml=c.getDensity(p,m) # compute density
151
        #print (ok)
152
153
154
155
156
157

        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,))
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
158
        cxv=cs.CSnapshot(None).center(p,v,m*rho)
159
160


LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
161
        unso=unsio.CunsOut("/home/jcl/x_x","gadget2")
162
163
164
165
166
167
168
169
        
        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
170
171
172
    #
    # fastCod3
    # 
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
173
    def fastCod3(self,threshold=10000,outfile=None):
174
175
176
177
178
179
180
        """
        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

LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
181
182
        outfile : if not None save 'outfile' to NEMO format 

183
        Return:
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
184
185
        cxv : COD, 1D numpy array of size 6
        rpos : positions of particles used to compute cod 1D num array
186
187
188
189
190
191
192
        """
        
        if threshold<0: # it's percentage of nbodies
            threshold = (abs(threshold)*self.__tree.getNbody())/100
        else:
            threshold = min(threshold,self.__tree.getNbody())

193
        #print("Threshold = ",threshold)
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208

        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
209
        #print(p.shape,p.size)
210
211
212
        m=self.__mass[shuffle_idx_threshold] # get randomized thresholded indexes for mass

        p=np.reshape(p,(-1,)) # flatten pos array
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
213

214
215
        #print(p.shape,p.size)
        #print (p[:,])
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
216
        c=cf.CFalcon() # new falcon object
217
        ok,rho,hsml=c.getDensity(p,m) # compute density
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
218
219
220
221
        self.__tree_pos  = p
        self.__tree_mass = m
        self.__tree_rho  = rho
        self.__tree_hsml = hsml
222
        #print (ok)
223
224
225
226
227
228

        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,))
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
229
230
231
             self.__tree_vel=v
        cxv=cs.CSnapshot(None).center(p,v,m*rho)
        
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
232

LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
233
234
235
236
237
        if v is None: # no velocities
            cxv[3]=0.0
            cxv[4]=0.0
            cxv[5]=0.0
        if outfile is not None:
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
238
            unso=unsio.CunsOut(outfile,"nemo")
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
239
240
            unso.setValueF("time",self.__time)
            unso.setArrayF("all","pos",p)
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
241
242
            if self.__vel is not None:
                unso.setArrayF("all","vel",v)
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
243
244
245
246
247
            unso.setArrayF("all","mass",m)
            unso.setArrayF("all","rho",rho)
            unso.setArrayF("all","hsml",hsml)
            unso.save()
        return cxv
LAMBERT Jean-charles's avatar
LAMBERT Jean-charles committed
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265

    #
    # getTreeDensity
    # 
    def getTreeDensity(self):
        """
        return rho and hsml computed in tree
        """
        return self.__tree_rho,self.__tree_hsml

    #
    # getTreePart
    # 
    def getTreePart(self):
        """
        return pos,vel and mass used to computed the tree
        """
        return self.__tree_pos,self.__tree_vel, self.__tree_mass