Matthew Turk avatar Matthew Turk committed 711e95a Merge

Merged in jzuhone/yt (pull request #361)

Comments (0)

Files changed (3)

yt/frontends/stream/data_structures.py

             self.stream_handler.particle_types[key] = particle_types[key]
             if key not in self.field_list:
                 self.field_list.append(key)
+                
+        self._setup_unknown_fields()
 
         for i, grid in enumerate(self.grids) :
             if data[i].has_key("number_of_particles") :
                 grid.NumberOfParticles = data[i].pop("number_of_particles")
             for key in data[i].keys() :
+                if key in grid.keys() : del grid[key]
                 self.stream_handler.fields[grid.id][key] = data[i][key]
             
-        self._setup_unknown_fields()
         self._detect_fields()
         
+        
 class StreamStaticOutput(StaticOutput):
     _hierarchy_class = StreamHierarchy
     _fieldinfo_fallback = StreamFieldInfo

yt/utilities/particle_generator.py

+import numpy as np
+import h5py
+from yt.utilities.lib import CICSample_3
+from yt.funcs import *
+
+class ParticleGenerator(object) :
+
+    default_fields = ["particle_position_x",
+                      "particle_position_y",
+                      "particle_position_z",
+                      "particle_index"]
+
+    def __init__(self, pf, num_particles, field_list) :
+        """
+        Base class for generating particle fields which may be applied to
+        streams. Normally this would not be called directly, since it doesn't
+        really do anything except allocate memory. Takes a *pf* to serve as the
+        basis for determining grids, the number of particles *num_particles*,
+        and a list of fields, *field_list*.
+        """
+        self.pf = pf
+        self.num_particles = num_particles
+        self.field_list = field_list
+            
+        try :
+            self.posx_index = self.field_list.index(self.default_fields[0])
+            self.posy_index = self.field_list.index(self.default_fields[1])
+            self.posz_index = self.field_list.index(self.default_fields[2])
+            self.index_index = self.field_list.index(self.default_fields[3])
+        except :
+            raise KeyError("Field list must contain the following fields: " +
+                           "\'particle_position_x\', \'particle_position_y\'" +
+                           ", \'particle_position_z\', \'particle_index\' ")
+
+        self.num_grids = self.pf.h.num_grids
+        self.NumberOfParticles = np.zeros((self.num_grids), dtype='int64')
+        self.ParticleIndices = np.zeros(self.num_grids + 1, dtype='int64')
+        
+        self.num_fields = len(self.field_list)
+        
+        self.particles = np.zeros((self.num_particles, self.num_fields),
+                                  dtype='float64')
+
+    def has_key(self, key) :
+        """
+        Check to see if *key* is in the particle field list.
+        """
+        return (key in self.field_list)
+            
+    def keys(self) :
+        """
+        Return the list of particle fields.
+        """
+        return self.field_list
+    
+    def __getitem__(self, key) :
+        """
+        Get the field associated with key.
+        """
+        return self.particles[:,self.field_list.index(key)]
+    
+    def __setitem__(self, key, val) :
+        """
+        Sets a field to be some other value. Note that we assume
+        that the particles have been sorted by grid already, so
+        make sure the setting of the field is consistent with this.
+        """
+        self.particles[:,self.field_list.index(key)] = val[:]
+                
+    def __len__(self) :
+        """
+        The number of particles
+        """
+        return self.num_particles
+
+    def get_from_grid(self, grid) :
+        """
+        Return a dict containing all of the particle fields in the specified grid.
+        """
+        ind = grid.id-grid._id_offset
+        start = self.ParticleIndices[ind]
+        end = self.ParticleIndices[ind+1]
+        return dict([(field, self.particles[start:end,self.field_list.index(field)])
+                     for field in self.field_list])
+    
+    def _setup_particles(self,x,y,z,setup_fields=None) :
+        """
+        Assigns grids to particles and sets up particle positions. *setup_fields* is
+        a dict of fields other than the particle positions to set up. 
+        """
+        particle_grids, particle_grid_inds = self.pf.h.find_points(x,y,z)
+        idxs = np.argsort(particle_grid_inds)
+        self.particles[:,self.posx_index] = x[idxs]
+        self.particles[:,self.posy_index] = y[idxs]
+        self.particles[:,self.posz_index] = z[idxs]
+        self.NumberOfParticles = np.bincount(particle_grid_inds,
+                                             minlength=self.num_grids)
+        if self.num_grids > 1 :
+            np.add.accumulate(self.NumberOfParticles.squeeze(),
+                              out=self.ParticleIndices[1:])
+        else :
+            self.ParticleIndices[1] = self.NumberOfParticles.squeeze()
+        if setup_fields is not None:
+            for key, value in setup_fields.items():
+                if key not in self.default_fields:
+                    self.particles[:,self.field_list.index(key)] = value[idxs]
+    
+    def assign_indices(self, function=None, **kwargs) :
+        """
+        Assign unique indices to the particles. The default is to just use
+        numpy.arange, but any function may be supplied with keyword arguments.
+        """
+        if function is None :
+            self.particles[:,self.index_index] = np.arange((self.num_particles))
+        else :
+            self.particles[:,self.index_index] = function(**kwargs)
+            
+    def map_grid_fields_to_particles(self, mapping_dict) :
+        r"""
+        For the fields in  *mapping_dict*, map grid fields to the particles
+        using CIC sampling.
+
+        Examples
+        --------
+        >>> field_map = {'Density':'particle_density',
+        >>>              'Temperature':'particle_temperature'}
+        >>> particles.map_grid_fields_to_particles(field_map)
+        """
+        pbar = get_pbar("Mapping fields to particles", self.num_grids)
+        for i, grid in enumerate(self.pf.h.grids) :
+            pbar.update(i)
+            if self.NumberOfParticles[i] > 0:
+                start = self.ParticleIndices[i]
+                end = self.ParticleIndices[i+1]
+                # Note we add one ghost zone to the grid!
+                cube = grid.retrieve_ghost_zones(1, mapping_dict.keys())
+                le = np.array(grid.LeftEdge).astype(np.float64)
+                dims = np.array(grid.ActiveDimensions).astype(np.int32)
+                for gfield, pfield in mapping_dict.items() :
+                    field_index = self.field_list.index(pfield)
+                    CICSample_3(self.particles[start:end,self.posx_index],
+                                self.particles[start:end,self.posy_index],
+                                self.particles[start:end,self.posz_index],
+                                self.particles[start:end,field_index],
+                                np.int64(self.NumberOfParticles[i]),
+                                cube[gfield], le, dims,
+                                np.float64(grid['dx']))
+        pbar.finish()
+
+    def apply_to_stream(self, clobber=False) :
+        """
+        Apply the particles to a stream parameter file. If particles already exist,
+        and clobber=False, do not overwrite them, but add the new ones to them. 
+        """
+        grid_data = []
+        for i,g in enumerate(self.pf.h.grids) :
+            data = {}
+            if clobber :
+                data["number_of_particles"] = self.NumberOfParticles[i]
+            else :
+                data["number_of_particles"] = self.NumberOfParticles[i] + \
+                                              g.NumberOfParticles
+            grid_particles = self.get_from_grid(g)
+            for field in self.field_list :
+                if data["number_of_particles"] > 0 :
+                    # We have particles in this grid
+                    if g.NumberOfParticles > 0 and not clobber:
+                        # Particles already exist
+                        if field in self.pf.h.field_list :
+                            # This field already exists
+                            #prev_particles = self.pf.h.io._read_data_set(g,field)
+                            prev_particles = g[field]
+                        else :
+                            # This one doesn't, set the previous particles' field
+                            # values to zero
+                            prev_particles = np.zeros((g.NumberOfParticles))
+                        data[field] = np.concatenate((prev_particles,
+                                                      grid_particles[field]))
+                    else :
+                        # Particles do not already exist or we're clobbering
+                        data[field] = grid_particles[field]
+                else :
+                    # We don't have particles in this grid
+                    data[field] = np.array([], dtype='float64')
+            grid_data.append(data)
+        self.pf.h.update_data(grid_data)
+
+class FromListParticleGenerator(ParticleGenerator) :
+
+    def __init__(self, pf, num_particles, data) :
+        r"""
+        Generate particle fields from array-like lists contained in a dict.
+
+        Parameters
+        ----------
+        pf : `StaticOutput`
+            The parameter file which will serve as the base for these particles.
+        num_particles : int
+            The number of particles in the dict.
+        data : dict of NumPy arrays
+            The particle fields themselves.
+
+        Examples
+        --------
+        >>> num_p = 100000
+        >>> posx = np.random.random((num_p))
+        >>> posy = np.random.random((num_p))
+        >>> posz = np.random.random((num_p))
+        >>> mass = np.ones((num_p))
+        >>> data = {'particle_position_x': posx, 'particle_position_y': posy,
+        >>>         'particle_position_z': posz, 'particle_mass': mass}
+        >>> particles = FromListParticleGenerator(pf, num_p, data)
+        """
+
+        field_list = data.keys()
+        x = data.pop("particle_position_x")
+        y = data.pop("particle_position_y")
+        z = data.pop("particle_position_z")
+
+        xcond = np.logical_or(x < pf.domain_left_edge[0],
+                              x >= pf.domain_right_edge[0])
+        ycond = np.logical_or(y < pf.domain_left_edge[1],
+                              y >= pf.domain_right_edge[1])
+        zcond = np.logical_or(z < pf.domain_left_edge[2],
+                              z >= pf.domain_right_edge[2])
+        cond = np.logical_or(xcond, ycond)
+        cond = np.logical_or(zcond, cond)
+
+        if np.any(cond) :
+            raise ValueError("Some particles are outside of the domain!!!")
+
+        ParticleGenerator.__init__(self, pf, num_particles, field_list)
+        self._setup_particles(x,y,z,setup_fields=data)
+        
+class LatticeParticleGenerator(ParticleGenerator) :
+
+    def __init__(self, pf, particles_dims, particles_left_edge,
+                 particles_right_edge, field_list) :
+        r"""
+        Generate particles in a lattice arrangement. 
+
+        Parameters
+        ----------
+        pf : `StaticOutput`
+            The parameter file which will serve as the base for these particles.
+        particles_dims : int, array-like 
+            The number of particles along each dimension
+        particles_left_edge : float, array-like
+            The 'left-most' starting positions of the lattice.
+        particles_right_edge : float, array-like
+             The 'right-most' ending positions of the lattice.
+        field_list : list of strings
+             A list of particle fields
+             
+        Examples
+        --------
+        >>> dims = (128,128,128)
+        >>> le = np.array([0.25,0.25,0.25])
+        >>> re = np.array([0.75,0.75,0.75])
+        >>> fields = ["particle_position_x","particle_position_y",
+        >>>           "particle_position_z","particle_index",
+        >>>           "particle_density","particle_temperature"]
+        >>> particles = LatticeParticleGenerator(pf, dims, le, re, fields)
+        """
+
+        num_x = particles_dims[0]
+        num_y = particles_dims[1]
+        num_z = particles_dims[2]
+        xmin = particles_left_edge[0]
+        ymin = particles_left_edge[1]
+        zmin = particles_left_edge[2]
+        xmax = particles_right_edge[0]
+        ymax = particles_right_edge[1]
+        zmax = particles_right_edge[2]
+
+        xcond = (xmin < pf.domain_left_edge[0]) or \
+                (xmax >= pf.domain_right_edge[0])
+        ycond = (ymin < pf.domain_left_edge[1]) or \
+                (ymax >= pf.domain_right_edge[1])
+        zcond = (zmin < pf.domain_left_edge[2]) or \
+                (zmax >= pf.domain_right_edge[2])
+        cond = xcond or ycond or zcond
+
+        if cond :
+            raise ValueError("Proposed bounds for particles are outside domain!!!")
+
+        ParticleGenerator.__init__(self, pf, num_x*num_y*num_z, field_list)
+
+        dx = (xmax-xmin)/(num_x-1)
+        dy = (ymax-ymin)/(num_y-1)
+        dz = (zmax-zmin)/(num_z-1)
+        inds = np.indices((num_x,num_y,num_z))
+        xpos = inds[0]*dx + xmin
+        ypos = inds[1]*dy + ymin
+        zpos = inds[2]*dz + zmin
+        
+        self._setup_particles(xpos.flat[:], ypos.flat[:], zpos.flat[:])
+        
+class WithDensityParticleGenerator(ParticleGenerator) :
+
+    def __init__(self, pf, data_source, num_particles, field_list,
+                 density_field="Density") :
+        r"""
+        Generate particles based on a density field.
+
+        Parameters
+        ----------
+        pf : `StaticOutput`
+            The parameter file which will serve as the base for these particles.
+        data_source : `yt.data_objects.api.AMRData`
+            The data source containing the density field.
+        num_particles : int
+            The number of particles to be generated
+        field_list : list of strings
+            A list of particle fields
+        density_field : string, optional
+            A density field which will serve as the distribution function for the
+            particle positions. Theoretically, this could be any 'per-volume' field. 
+            
+        Examples
+        --------
+        >>> sphere = pf.h.sphere(pf.domain_center, 0.5)
+        >>> num_p = 100000
+        >>> fields = ["particle_position_x","particle_position_y",
+        >>>           "particle_position_z","particle_index",
+        >>>           "particle_density","particle_temperature"]
+        >>> particles = WithDensityParticleGenerator(pf, sphere, num_particles,
+        >>>                                          fields, density_field='Dark_Matter_Density')
+        """
+
+        ParticleGenerator.__init__(self, pf, num_particles, field_list)
+
+        num_cells = len(data_source["x"].flat)
+        max_density = data_source[density_field].max()
+        num_particles_left = num_particles
+        all_x = []
+        all_y = []
+        all_z = []
+        
+        pbar = get_pbar("Generating Particles", num_particles)
+        tot_num_accepted = int(0)
+        
+        while num_particles_left > 0:
+
+            rho = np.random.uniform(high=1.01*max_density,
+                                    size=num_particles_left)
+            idxs = np.random.random_integers(low=0, high=num_cells-1,
+                                             size=num_particles_left)
+            rho_true = data_source[density_field].flat[idxs]
+            accept = rho <= rho_true
+            num_accepted = accept.sum()
+            accepted_idxs = idxs[accept]
+            
+            xpos = data_source["x"].flat[accepted_idxs] + \
+                   np.random.uniform(low=-0.5, high=0.5, size=num_accepted) * \
+                   data_source["dx"].flat[accepted_idxs]
+            ypos = data_source["y"].flat[accepted_idxs] + \
+                   np.random.uniform(low=-0.5, high=0.5, size=num_accepted) * \
+                   data_source["dy"].flat[accepted_idxs]
+            zpos = data_source["z"].flat[accepted_idxs] + \
+                   np.random.uniform(low=-0.5, high=0.5, size=num_accepted) * \
+                   data_source["dz"].flat[accepted_idxs]
+
+            all_x.append(xpos)
+            all_y.append(ypos)
+            all_z.append(zpos)
+
+            num_particles_left -= num_accepted
+            tot_num_accepted += num_accepted
+            pbar.update(tot_num_accepted)
+
+        pbar.finish()
+
+        x = np.concatenate(all_x)
+        y = np.concatenate(all_y)
+        z = np.concatenate(all_z)
+
+        self._setup_particles(x,y,z)
+        

yt/utilities/tests/test_particle_generator.py

+import numpy as np
+from yt.testing import *
+from yt.utilities.particle_generator import *
+from yt.frontends.stream.api import load_uniform_grid, refine_amr
+import yt.utilities.initial_conditions as ic
+import yt.utilities.flagging_methods as fm
+from IPython import embed
+
+def setup() :
+    pass
+
+def test_particle_generator() :
+    
+    # First generate our pf
+    domain_dims = (128, 128, 128)
+    dens = np.zeros(domain_dims) + 0.1
+    temp = 4.*np.ones(domain_dims)
+    fields = {"Density": dens, "Temperature": temp}
+    ug = load_uniform_grid(fields, domain_dims, 1.0)
+    fo = [ic.BetaModelSphere(1.0,0.1,0.5,[0.5,0.5,0.5],{"Density":(10.0)})]
+    rc = [fm.flagging_method_registry["overdensity"](4.0)]
+    pf = refine_amr(ug, rc, fo, 3)
+
+    # Now generate particles from density
+
+    field_list = ["particle_position_x","particle_position_y",
+                  "particle_position_z","particle_index",
+                  "particle_gas_density"]
+    num_particles = 1000000
+    field_dict = {"Density": "particle_gas_density"}
+    sphere = pf.h.sphere(pf.domain_center, 0.45)
+
+    particles1 = WithDensityParticleGenerator(pf, sphere, num_particles, field_list)
+    particles1.assign_indices()
+    particles1.map_grid_fields_to_particles(field_dict)
+    
+    # Test to make sure we ended up with the right number of particles per grid
+    particles1.apply_to_stream()
+    particles_per_grid1 = [grid.NumberOfParticles for grid in pf.h.grids]
+    assert_equal(particles_per_grid1, particles1.NumberOfParticles)
+    particles_per_grid1 = [len(grid["particle_position_x"]) for grid in pf.h.grids]
+    assert_equal(particles_per_grid1, particles1.NumberOfParticles)
+
+    # Set up a lattice of particles
+    pdims = np.array([64,64,64])
+    def new_indices() :
+        # We just add new indices onto the existing ones
+        return np.arange((np.product(pdims)))+num_particles
+    le = np.array([0.25,0.25,0.25])
+    re = np.array([0.75,0.75,0.75])
+    new_field_list = field_list + ["particle_gas_temperature"]
+    new_field_dict = {"Density": "particle_gas_density",
+                      "Temperature": "particle_gas_temperature"}
+
+    particles2 = LatticeParticleGenerator(pf, pdims, le, re, new_field_list)
+    particles2.assign_indices(function=new_indices)
+    particles2.map_grid_fields_to_particles(new_field_dict)
+
+    #Test lattice positions
+    xpos = np.unique(particles2["particle_position_x"])
+    ypos = np.unique(particles2["particle_position_y"])
+    zpos = np.unique(particles2["particle_position_z"])
+
+    xpred = np.linspace(le[0],re[0],num=pdims[0],endpoint=True)
+    ypred = np.linspace(le[1],re[1],num=pdims[1],endpoint=True)
+    zpred = np.linspace(le[2],re[2],num=pdims[2],endpoint=True)
+
+    assert_almost_equal(xpos, xpred)
+    assert_almost_equal(ypos, ypred)
+    assert_almost_equal(zpos, zpred)
+
+    #Test the number of particles again
+    particles2.apply_to_stream()
+    particles_per_grid2 = [grid.NumberOfParticles for grid in pf.h.grids]
+    assert_equal(particles_per_grid2, particles1.NumberOfParticles+particles2.NumberOfParticles)
+    particles_per_grid2 = [len(grid["particle_position_x"]) for grid in pf.h.grids]
+    assert_equal(particles_per_grid2, particles1.NumberOfParticles+particles2.NumberOfParticles)
+
+    #Test the uniqueness of tags
+    tags = np.concatenate([grid["particle_index"] for grid in pf.h.grids])
+    tags.sort()
+    assert_equal(tags, np.arange((np.product(pdims)+num_particles)))
+
+    # Test that the old particles have zero for the new field
+    old_particle_temps = [grid["particle_gas_temperature"][:particles_per_grid1[i]]
+                          for i, grid in enumerate(pf.h.grids)]
+    test_zeros = [np.zeros((particles_per_grid1[i])) 
+                  for i, grid in enumerate(pf.h.grids)]
+    assert_equal(old_particle_temps, test_zeros)
+
+    #Now dump all of these particle fields out into a dict
+    pdata = {}
+    dd = pf.h.all_data()
+    for field in new_field_list :
+        pdata[field] = dd[field]
+
+    #Test the "from-list" generator and particle field clobber
+    particles3 = FromListParticleGenerator(pf, num_particles+np.product(pdims), pdata)
+    particles3.apply_to_stream(clobber=True)
+    
+    #Test the number of particles again
+    particles_per_grid3 = [grid.NumberOfParticles for grid in pf.h.grids]
+    assert_equal(particles_per_grid3, particles1.NumberOfParticles+particles2.NumberOfParticles)
+    particles_per_grid2 = [len(grid["particle_position_z"]) for grid in pf.h.grids]
+    assert_equal(particles_per_grid3, particles1.NumberOfParticles+particles2.NumberOfParticles)
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.