[yt-svn] commit/yt: ngoldbaum: Merged in qobilidop/yt (pull request #1985)

commits-noreply at bitbucket.org commits-noreply at bitbucket.org
Sat Jul 9 14:49:38 PDT 2016


1 new commit in yt:

https://bitbucket.org/yt_analysis/yt/commits/98f754270a40/
Changeset:   98f754270a40
Branch:      yt
User:        ngoldbaum
Date:        2016-07-09 21:49:09+00:00
Summary:     Merged in qobilidop/yt (pull request #1985)

Add particle type aware octree construction
Affected #:  7 files

diff -r 142714a92cbb18866fbdb7b94cf064c10be848ad -r 98f754270a409f2798b0abe8cd8db13021ac9743 doc/source/examining/loading_data.rst
--- a/doc/source/examining/loading_data.rst
+++ b/doc/source/examining/loading_data.rst
@@ -884,6 +884,21 @@
 It's recommended that if you want higher-resolution, try reducing the value of
 ``n_ref`` to 32 or 16.
 
+Also yt can be set to generate the global mesh index according to a specific
+type of particles instead of all the particles through the parameter
+``index_ptype``. For example, to build the octree only according to the
+``"PartType0"`` particles, you can do:
+
+.. code-block:: python
+
+   ds = yt.load("snapshot_061.hdf5", index_ptype="PartType0")
+
+By default, ``index_ptype`` is set to ``"all"``, which means all the particles.
+Currently this feature only works for the Gadget HDF5 and OWLS datasets. To
+bring the feature to other frontends, it's recommended to refer to this
+`PR <https://bitbucket.org/yt_analysis/yt/pull-requests/1985/add-particle-type-aware-octree/diff>`_
+for implementation details.
+
 .. _gadget-field-spec:
 
 Field Specifications

diff -r 142714a92cbb18866fbdb7b94cf064c10be848ad -r 98f754270a409f2798b0abe8cd8db13021ac9743 yt/frontends/gadget/data_structures.py
--- a/yt/frontends/gadget/data_structures.py
+++ b/yt/frontends/gadget/data_structures.py
@@ -75,6 +75,7 @@
                  additional_fields=(),
                  unit_base=None, n_ref=64,
                  over_refine_factor=1,
+                 index_ptype="all",
                  bounding_box = None,
                  header_spec = "default",
                  field_spec = "default",
@@ -90,6 +91,7 @@
             ptype_spec, gadget_ptype_specs)
         self.n_ref = n_ref
         self.over_refine_factor = over_refine_factor
+        self.index_ptype = index_ptype
         self.storage_filename = None
         if unit_base is not None and "UnitLength_in_cm" in unit_base:
             # We assume this is comoving, because in the absence of comoving
@@ -343,6 +345,7 @@
     def __init__(self, filename, dataset_type="gadget_hdf5",
                  unit_base = None, n_ref=64,
                  over_refine_factor=1,
+                 index_ptype="all",
                  bounding_box = None,
                  units_override=None,
                  unit_system="cgs"):
@@ -353,7 +356,7 @@
                                "Use unit_base instead.")
         super(GadgetHDF5Dataset, self).__init__(
             filename, dataset_type, unit_base=unit_base, n_ref=n_ref,
-            over_refine_factor=over_refine_factor,
+            over_refine_factor=over_refine_factor, index_ptype=index_ptype,
             bounding_box = bounding_box, unit_system=unit_system)
 
     def _get_hvals(self):

diff -r 142714a92cbb18866fbdb7b94cf064c10be848ad -r 98f754270a409f2798b0abe8cd8db13021ac9743 yt/frontends/owls/io.py
--- a/yt/frontends/owls/io.py
+++ b/yt/frontends/owls/io.py
@@ -123,11 +123,18 @@
             f.close()
 
     def _initialize_index(self, data_file, regions):
+        index_ptype = self.index_ptype
         f = _get_h5_handle(data_file.filename)
-        pcount = f["/Header"].attrs["NumPart_ThisFile"][:].sum()
+        if index_ptype == "all":
+            pcount = f["/Header"].attrs["NumPart_ThisFile"][:].sum()
+            keys = f.keys()
+        else:
+            pt = int(index_ptype[-1])
+            pcount = f["/Header"].attrs["NumPart_ThisFile"][pt]
+            keys = [index_ptype]
         morton = np.empty(pcount, dtype='uint64')
         ind = 0
-        for key in f.keys():
+        for key in keys:
             if not key.startswith("PartType"): continue
             if "Coordinates" not in f[key]: continue
             ds = f[key]["Coordinates"]

diff -r 142714a92cbb18866fbdb7b94cf064c10be848ad -r 98f754270a409f2798b0abe8cd8db13021ac9743 yt/frontends/stream/data_structures.py
--- a/yt/frontends/stream/data_structures.py
+++ b/yt/frontends/stream/data_structures.py
@@ -299,7 +299,7 @@
         name = "InMemoryParameterFile_%s" % (uuid.uuid4().hex)
         from yt.data_objects.static_output import _cached_datasets
         _cached_datasets[name] = self
-        Dataset.__init__(self, name, self._dataset_type, 
+        Dataset.__init__(self, name, self._dataset_type,
                          unit_system=unit_system)
 
     def _parse_parameter_file(self):
@@ -400,13 +400,13 @@
     Assign particle data to the grids using MatchPointsToGrids. This
     will overwrite any existing particle data, so be careful!
     """
-    
+
     # Note: what we need to do here is a bit tricky.  Because occasionally this
     # gets called before we property handle the field detection, we cannot use
     # any information about the index.  Fortunately for us, we can generate
     # most of the GridTree utilizing information we already have from the
     # stream handler.
-    
+
     if len(ds.stream_handler.fields) > 1:
 
         if ("io", "particle_position_x") in pdata:
@@ -425,7 +425,7 @@
             np.equal(parent_ids, i, mask)
             num_children[i] = mask.sum()
         levels = ds.stream_handler.levels.astype("int64").ravel()
-        grid_tree = GridTree(num_grids, 
+        grid_tree = GridTree(num_grids,
                              ds.stream_handler.left_edges,
                              ds.stream_handler.right_edges,
                              ds.stream_handler.dimensions,
@@ -443,8 +443,8 @@
                               out=particle_indices[1:])
         else :
             particle_indices[1] = particle_grid_count.squeeze()
-    
-        pdata.pop("number_of_particles", None) 
+
+        pdata.pop("number_of_particles", None)
         grid_pdata = []
         for i, pcount in enumerate(particle_grid_count):
             grid = {}
@@ -457,12 +457,12 @@
 
     else :
         grid_pdata = [pdata]
-    
+
     for pd, gi in zip(grid_pdata, sorted(ds.stream_handler.fields)):
         ds.stream_handler.fields[gi].update(pd)
         npart = ds.stream_handler.fields[gi].pop("number_of_particles", 0)
         ds.stream_handler.particle_count[gi] = npart
-                                        
+
 def unitify_data(data):
     new_data, field_units = {}, {}
     for field, val in data.items():
@@ -507,7 +507,7 @@
     # At this point, we have arrays for all our fields
     new_data = {}
     for field in data:
-        if isinstance(field, tuple): 
+        if isinstance(field, tuple):
             new_field = field
         elif len(data[field].shape) in (1, 2):
             new_field = ("io", field)
@@ -580,7 +580,7 @@
         "spectral_cube".  Optionally, a tuple can be provided to specify the
         axis ordering -- for instance, to specify that the axis ordering should
         be z, x, y, this would be: ("cartesian", ("z", "x", "y")).  The same
-        can be done for other coordinates, for instance: 
+        can be done for other coordinates, for instance:
         ("spherical", ("theta", "phi", "r")).
 
     Examples
@@ -782,7 +782,7 @@
         "spectral_cube".  Optionally, a tuple can be provided to specify the
         axis ordering -- for instance, to specify that the axis ordering should
         be z, x, y, this would be: ("cartesian", ("z", "x", "y")).  The same
-        can be done for other coordinates, for instance: 
+        can be done for other coordinates, for instance:
         ("spherical", ("theta", "phi", "r")).
     refine_by : integer
         Specifies the refinement ratio between levels.  Defaults to 2.
@@ -998,7 +998,7 @@
 
 class StreamParticleIndex(ParticleIndex):
 
-    
+
     def __init__(self, ds, dataset_type = None):
         self.stream_handler = ds.stream_handler
         super(StreamParticleIndex, self).__init__(ds, dataset_type)
@@ -1041,7 +1041,7 @@
 
     This will initialize an Octree of data.  Note that fluid fields will not
     work yet, or possibly ever.
-    
+
     Parameters
     ----------
     data : dict
@@ -1091,7 +1091,7 @@
 
     field_units, data = unitify_data(data)
     sfh = StreamDictFieldHandler()
-    
+
     pdata = {}
     for key in data.keys() :
         if not isinstance(key, tuple):
@@ -1192,7 +1192,7 @@
     array([[-1.  , -1.  , -1.  ],
            [-1.  , -1.  , -0.25],
            [-1.  , -1.  ,  0.  ],
-           ..., 
+           ...,
            [ 1.  ,  1.  ,  0.  ],
            [ 1.  ,  1.  ,  0.25],
            [ 1.  ,  1.  ,  1.  ]])
@@ -1270,7 +1270,7 @@
 
     Particle fields are detected as one-dimensional fields. The number of particles
     is set by the "number_of_particles" key in data.
-    
+
     Parameters
     ----------
     data : dict
@@ -1303,7 +1303,7 @@
         "spectral_cube".  Optionally, a tuple can be provided to specify the
         axis ordering -- for instance, to specify that the axis ordering should
         be z, x, y, this would be: ("cartesian", ("z", "x", "y")).  The same
-        can be done for other coordinates, for instance: 
+        can be done for other coordinates, for instance:
         ("spherical", ("theta", "phi", "r")).
 
     """
@@ -1318,9 +1318,9 @@
 
     field_units, data = unitify_data(data)
     sfh = StreamDictFieldHandler()
-    
+
     particle_types = set_particle_types(data)
-    
+
     sfh.update({'connectivity': connectivity,
                 'coordinates': coordinates,
                 0: data})
@@ -1404,7 +1404,7 @@
         dest.update((field, np.empty(cell_count, dtype="float64"))
                     for field in content)
         # Make references ...
-        count = oct_handler.fill_level(0, levels, cell_inds, file_inds, 
+        count = oct_handler.fill_level(0, levels, cell_inds, file_inds,
                                        dest, content, offset)
         return count
 
@@ -1493,7 +1493,7 @@
 
     This will initialize an Octree of data.  Note that fluid fields will not
     work yet, or possibly ever.
-    
+
     Parameters
     ----------
     octree_mask : np.ndarray[uint8_t]
@@ -1826,7 +1826,7 @@
 
     sds._node_fields = node_data[0].keys()
     sds._elem_fields = elem_data[0].keys()
-    sds.default_field = [f for f in sds.field_list 
+    sds.default_field = [f for f in sds.field_list
                          if f[0] == 'connect1'][-1]
 
     return sds

diff -r 142714a92cbb18866fbdb7b94cf064c10be848ad -r 98f754270a409f2798b0abe8cd8db13021ac9743 yt/geometry/particle_geometry_handler.py
--- a/yt/geometry/particle_geometry_handler.py
+++ b/yt/geometry/particle_geometry_handler.py
@@ -38,6 +38,13 @@
         self.float_type = np.float64
         super(ParticleIndex, self).__init__(ds, dataset_type)
 
+    @property
+    def index_ptype(self):
+        if hasattr(self.dataset, "index_ptype"):
+            return self.dataset.index_ptype
+        else:
+            return "all"
+
     def _setup_geometry(self):
         mylog.debug("Initializing Particle Geometry Handler.")
         self._initialize_particle_handler()
@@ -69,14 +76,21 @@
         cls = self.dataset._file_class
         self.data_files = [cls(self.dataset, self.io, template % {'num':i}, i)
                            for i in range(ndoms)]
-        self.total_particles = sum(
-                sum(d.total_particles.values()) for d in self.data_files)
+        index_ptype = self.index_ptype
+        if index_ptype == "all":
+            self.total_particles = sum(
+                    sum(d.total_particles.values()) for d in self.data_files)
+        else:
+            self.total_particles = sum(
+                    d.total_particles[index_ptype] for d in self.data_files)
         ds = self.dataset
         self.oct_handler = ParticleOctreeContainer(
             [1, 1, 1], ds.domain_left_edge, ds.domain_right_edge,
             over_refine = ds.over_refine_factor)
         self.oct_handler.n_ref = ds.n_ref
-        only_on_root(mylog.info, "Allocating for %0.3e particles", self.total_particles)
+        only_on_root(mylog.info, "Allocating for %0.3e particles "
+                                 "(index particle type '%s')",
+                     self.total_particles, index_ptype)
         # No more than 256^3 in the region finder.
         N = min(len(self.data_files), 256) 
         self.regions = ParticleRegions(
@@ -100,10 +114,17 @@
         #   * Broadcast back a serialized octree to join
         #
         # For now we will do this in serial.
+        index_ptype = self.index_ptype
+        # Set the index_ptype attribute of self.io dynamically here, so we don't
+        # need to assume that the dataset has the attribute.
+        self.io.index_ptype = index_ptype
         morton = np.empty(self.total_particles, dtype="uint64")
         ind = 0
         for data_file in self.data_files:
-            npart = sum(data_file.total_particles.values())
+            if index_ptype == "all":
+                npart = sum(data_file.total_particles.values())
+            else:
+                npart = data_file.total_particles[index_ptype]
             morton[ind:ind + npart] = \
                 self.io._initialize_index(data_file, self.regions)
             ind += npart

diff -r 142714a92cbb18866fbdb7b94cf064c10be848ad -r 98f754270a409f2798b0abe8cd8db13021ac9743 yt/geometry/tests/test_particle_octree.py
--- a/yt/geometry/tests/test_particle_octree.py
+++ b/yt/geometry/tests/test_particle_octree.py
@@ -144,6 +144,21 @@
             cv2 = dd2["cell_volume"].sum(dtype="float64")
             yield assert_equal, cv1, cv2
 
+index_ptype_snap = "snapshot_033/snap_033.0.hdf5"
+ at requires_file(index_ptype_snap)
+def test_particle_index_ptype():
+    ds = yt.load(index_ptype_snap)
+    ds_all = yt.load(index_ptype_snap, index_ptype="all")
+    ds_pt0 = yt.load(index_ptype_snap, index_ptype="PartType0")
+    dd = ds.all_data()
+    dd_all = ds_all.all_data()
+    dd_pt0 = ds_pt0.all_data()
+    cv = dd["cell_volume"]
+    cv_all = dd_all["cell_volume"]
+    cv_pt0 = dd_pt0["cell_volume"]
+    yield assert_equal, cv.shape, cv_all.shape
+    yield assert_equal, cv.sum(dtype="float64"), cv_pt0.sum(dtype="float64")
+
 class FakeDS:
     domain_left_edge = None
     domain_right_edge = None

Repository URL: https://bitbucket.org/yt_analysis/yt/

--

This is a commit notification from bitbucket.org. You are receiving
this because you have the service enabled, addressing the recipient of
this email.



More information about the yt-svn mailing list