Christopher Moody avatar Christopher Moody committed a3b9d69

changes to allow for particles gridded everywere or on the root grid

Comments (0)

Files changed (3)

yt/frontends/art/data_structures.py

     _id_offset = 0
 
     def __init__(self, id, hierarchy, level, locations,start_index, le,re,gd,
-            child_mask=None,nop=0):
+            child_mask=None,nop=None):
         AMRGridPatch.__init__(self, id, filename = hierarchy.hierarchy_filename,
                               hierarchy = hierarchy)
         start_index =start_index 
         self.Children = []
         self.locations = locations
         self.start_index = start_index.copy()
-        
         self.LeftEdge = le
         self.RightEdge = re
         self.ActiveDimensions = gd
-        self.NumberOfParticles=nop
+        self.NumberOfParticles= nop 
         for particle_field in particle_fields:
             setattr(self,particle_field,np.array([]))
 
         self.float_type = np.float64
         AMRHierarchy.__init__(self,pf,data_style)
         self._setup_field_list()
-        if not pf.skip_particles and pf.file_particle_data:
-            self._grid_particles()
-            mylog.info("Gridded particles")
+        self._grid_particles()
+
 
     def _grid_particles(self):
-        #fill out self.particle_grid_indices
-        x,y,z = [self.pf.particle_position[:,i] for i in range(3)]
-        point_grids, self.particle_grid_indices = self.find_points(x,y,z)
+        if not self.pf.skip_particles and self.pf.file_particle_data\
+                and not self.pf.skip_gridding_particles:
+            #fill out self.particle_grid_indices
+            x = self.pf.particle_position_x
+            y = self.pf.particle_position_y
+            z = self.pf.particle_position_z
+            point_grids, self.pf.particle_grid_indices = self.find_points(x,y,z)
+            gindices = np.arange(len(self.grids)+0)-0.5
+            counts,indices= np.histogram(self.pf.particle_grid_indices,
+                                         bins=gindices)
+            for g,c in zip(self.grids,counts):
+                g.NumberOfParticles = c
+            mylog.info("Gridded particles")
+        elif not self.pf.skip_particles and self.pf.file_particle_data:
+            g = self.grids[0]
+            for particle_field in particle_fields:
+                source = getattr(self.pf,particle_field,None)
+                if source is not None:
+                    mylog.info("Attaching %s to the root grid",
+                                particle_field)
+                    g.NumberOfParticles = source.shape[0]
+                    setattr(g,particle_field,source)
+            g.particle_index = np.arange(g.NumberOfParticles)
+            self.pf.particle_grid_indices = None
+            mylog.info("Attached particles to root grid")
 
     def _initialize_data_storage(self):
         pass
         if not self.pf.skip_particles and self.pf.file_particle_data:
             lspecies = self.pf.parameters['lspecies']
             wspecies = self.pf.parameters['wspecies']
-            self.pf.particle_position,self.pf.particle_velocity= \
+            particle_position,particle_velocity= \
                 read_particles(self.pf.file_particle_data,
                         self.pf.parameters['Nrow'])
             nparticles = lspecies[-1]
-            if not np.all(self.pf.particle_position[nparticles:]==0.0):
+            if not np.all(particle_position[nparticles:]==0.0):
                 mylog.info('WARNING: unused particles discovered from lspecies')
-            self.pf.particle_position = self.pf.particle_position[:nparticles]
-            self.pf.particle_velocity = self.pf.particle_velocity[:nparticles]
-            self.pf.particle_position  /= self.pf.domain_dimensions 
+            particle_position /= self.pf.domain_dimensions 
+            self.pf.particle_position_x = particle_position[:nparticles,0]
+            self.pf.particle_position_y = particle_position[:nparticles,1]
+            self.pf.particle_position_z = particle_position[:nparticles,2]
+            self.pf.particle_velocity_x = particle_velocity[:nparticles,0]
+            self.pf.particle_velocity_y = particle_velocity[:nparticles,1]
+            self.pf.particle_velocity_z = particle_velocity[:nparticles,2]
+            del particle_position, particle_velocity
             self.pf.particle_type = np.zeros(nparticles,dtype='int')
             self.pf.particle_mass = np.zeros(nparticles,dtype='float64')
             self.pf.particle_star_index = len(wspecies)-1
             g._prepare_grid()
             g._setup_dx()
             #instead of gridding particles assign them all to the root grid
-            if gi==0:
-                for particle_field in particle_fields:
-                    source = getattr(self.pf,particle_field,None)
-                    if source is None:
-                        for i,ax in enumerate('xyz'):
-                            pf = particle_field.replace('_%s'%ax,'')
-                            source = getattr(self.pf,pf,None)
-                            if source is not None:
-                                source = source[:,i]
-                                break
-                    if source is not None:
-                        mylog.info("Attaching %s to the root grid",
-                                    particle_field)
-                        g.NumberOfParticles = source.shape[0]
-                        setattr(g,particle_field,source)
-                g.particle_index = np.arange(g.NumberOfParticles)
         pb.finish()
         self.max_level = self.grid_levels.max()
 
     _fieldinfo_known = KnownARTFields
     
     def __init__(self, file_amr, storage_filename = None,
-            skip_particles=False,skip_stars=False,limit_level=None,
+            skip_particles=False,skip_stars=False,
+            skip_gridding_particles = True,
+            limit_level=None,
             spread_age=True,data_style='art'):
         self.data_style = data_style
         self._find_files(file_amr)
         self.skip_particles = skip_particles
         self.skip_stars = skip_stars
+        self.skip_gridding_particles = skip_gridding_particles
         self.file_amr = file_amr
         self.parameter_filename = file_amr
         self.limit_level = limit_level

yt/frontends/art/fields.py

 
 
 #Particle fields
-
-def ParticleMass(field,data):
-    return data['particle_mass']
-add_field("ParticleMass",function=ParticleMass,units=r"\rm{g}",particle_type=True)
-
-
-#Derived particle fields
-
-def ParticleMassMsun(field,data):
-    return data['particle_mass']*data.pf['Msun']
-add_field("ParticleMassMsun",function=ParticleMassMsun,units=r"\rm{g}",particle_type=True)
-
-def _creation_time(field,data):
-    pa = data["particle_age"]
-    tr = np.zeros(pa.shape,dtype='float')-1.0
-    tr[pa>0] = pa[pa>0]
-    return tr
-add_field("creation_time",function=_creation_time,units=r"\rm{s}",particle_type=True)
-
-def mass_dm(field, data):
-    tr = np.ones(data.ActiveDimensions, dtype='float32')
-    idx = data["particle_type"]<5
-    #make a dumb assumption that the mass is evenly spread out in the grid
-    #must return an array the shape of the grid cells
-    if np.sum(idx)>0:
-        tr /= np.prod(data['CellVolumeCode']*data.pf['mpchcm']**3.0) #divide by the volume
-        tr *= np.sum(data['particle_mass'][idx])*data.pf['Msun'] #Multiply by total contaiend mass
-        print tr.shape
-        return tr
-    else:
-        return tr*1e-9
-
-add_field("particle_cell_mass_dm", function=mass_dm, units = r"\mathrm{M_{sun}}",
-        validators=[ValidateSpatial(0)],        
-        take_log=False,
-        projection_conversion="1")
-
 def _spdensity(field, data):
-    grid_mass = np.zeros(data.ActiveDimensions, dtype='float32')
-    if data.star_mass.shape[0] ==0 : return grid_mass 
-    amr_utils.CICDeposit_3(data.star_position_x,
-                           data.star_position_y,
-                           data.star_position_z,
-                           data.star_mass.astype('float32'),
-                           data.star_mass.shape[0],
-                           grid_mass, 
-                           np.array(data.LeftEdge).astype(np.float64),
+    blank = np.zeros(data.ActiveDimensions, dtype='float32')
+    idx = data['particle_type']==data.pf.particle_star_index
+    if not idx.any(): return blank
+    amr_utils.CICDeposit_3(data["particle_position_x"][idx].astype(np.float64),
+                           data["particle_position_y"][idx].astype(np.float64),
+                           data["particle_position_z"][idx].astype(np.float64),
+                           data["particle_mass"][idx].astype(np.float32),
+                           np.int64(np.where(idx)[0].size),
+                           blank, np.array(data.LeftEdge).astype(np.float64),
                            np.array(data.ActiveDimensions).astype(np.int32), 
                            np.float64(data['dx']))
-    return grid_mass 
+    return blank
+add_field("star_density", function=_spdensity,
+          validators=[ValidateSpatial(0)], convert_function=_convertDensity)
 
-#add_field("star_density", function=_spdensity,
-#          validators=[ValidateSpatial(0)], convert_function=_convertDensity)
 
-def _simple_density(field,data):
-    mass = np.sum(data.star_mass)
-    volume = data['dx']*data.ActiveDimensions.prod().astype('float64')
-    return mass/volume
-
-add_field("star_density", function=_simple_density,
-          validators=[ValidateSpatial(0)], convert_function=_convertDensity)

yt/frontends/art/io.py

         self.level_data.pop(level, None)
 
     def _read_particle_field(self, grid, field):
-        dat = getattr(grid,field,None)
-        if dat is not None: 
-            return dat
-        starfield = field.replace('star','particle')
-        dat = getattr(grid,starfield,None)
-        if dat is not None:
-            psi = grid.pf.particle_star_index
-            idx = grid.particle_type==psi
-            return dat[idx]
+        if grid.pf.particle_grid_indices is not None:
+            alldat = getattr(grid.pf,field,None)
+            if alldat is not None: 
+                idx = grid.pf.particle_grid_indices==grid.id
+                return alldat[idx]
+        else:
+            alldat = getattr(grid,field,None)
+            if alldat is not None: 
+                return alldat
+        if 'star' in field:
+            pfield = field.replace('star','particle')
+            if grid.pf.particle_grid_indices is not None:
+                alldat = getattr(grid.pf,field,None)
+                if alldat is not None: 
+                    idx = grid.pf.particle_grid_indices==grid.id
+                    idx &= grid.pf.particle_grid_indices==grid.id
+                    return alldat[idx]
+            else:
+                alldat = getattr(grid,field,None)
+                if alldat is not None: 
+                    return alldat
         raise KeyError
         
     def _read_data_set(self, grid, field):
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.