Commits

Matthew Turk committed 77b261f

Adding spatial chunking to particle frontends.

  • Participants
  • Parent commits 1babf36

Comments (0)

Files changed (3)

File yt/data_objects/static_output.py

     particle_types = ("all",)
     geometry = "cartesian"
     coordinates = None
+    max_level = 99
 
     class __metaclass__(type):
         def __init__(cls, name, b, d):

File yt/frontends/sph/data_structures.py

     GeometryHandler, YTDataChunk
 from yt.data_objects.static_output import \
     StaticOutput
+from yt.data_objects.octree_subset import \
+    OctreeSubset
 from yt.utilities.definitions import \
     mpc_conversion, sec_conversion
 from .fields import \
     def _calculate_offsets(self, fields):
         pass
 
-class ParticleDomainSubset(object):
-    def __init__(self, domain, mask, count):
-        self.domain = domain
-        self.mask = mask
-        self.cell_count = count
-        self.oct_handler = domain.pf.h.oct_handler
-        level_counts = self.oct_handler.count_levels(
-            99, self.domain.domain_id, mask)
-        level_counts[1:] = level_counts[:-1]
-        level_counts[0] = 0
-        self.level_counts = np.add.accumulate(level_counts)
-
-    def select_icoords(self, dobj):
-        return self.oct_handler.icoords(self.domain.domain_id, self.mask,
-                                        self.cell_count)
-
-    def select_fcoords(self, dobj):
-        return self.oct_handler.fcoords(self.domain.domain_id, self.mask,
-                                        self.cell_count)
-
-    def select_fwidth(self, dobj):
-        # Recall domain_dimensions is the number of cells, not octs
-        base_dx = (self.domain.pf.domain_width /
-                   self.domain.pf.domain_dimensions)
-        widths = np.empty((self.cell_count, 3), dtype="float64")
-        dds = (2**self.ires(dobj))
-        for i in range(3):
-            widths[:,i] = base_dx[i] / dds
-        return widths
-
-    def select_ires(self, dobj):
-        return self.oct_handler.ires(self.domain.domain_id, self.mask,
-                                     self.cell_count)
-
+class ParticleDomainSubset(OctreeSubset):
+    pass
 
 class ParticleGeometryHandler(OctreeGeometryHandler):
 
         oobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info)
         yield YTDataChunk(dobj, "all", oobjs, dobj.size)
 
-    def _chunk_spatial(self, dobj, ngz):
-        raise NotImplementedError
+    def _chunk_spatial(self, dobj, ngz, sort = None):
+        sobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info)
+        for i,og in enumerate(sobjs):
+            if ngz > 0:
+                g = og.retrieve_ghost_zones(ngz, [], smoothed=True)
+            else:
+                g = og
+            size = og.cell_count
+            if size == 0: continue
+            yield YTDataChunk(dobj, "spatial", [g], size)
 
     def _chunk_io(self, dobj):
         oobjs = getattr(dobj._current_chunk, "objs", dobj._chunk_info)
         self.domain_right_edge = np.ones(3, "float64") * hvals["BoxSize"]
         self.domain_dimensions = np.ones(3, "int32") * self._root_dimensions
         self.cosmological_simulation = 1
+        self.periodicity = (True, True, True)
         self.current_redshift = hvals["Redshift"]
         self.omega_lambda = hvals["OmegaLambda"]
         self.omega_matter = hvals["Omega0"]
         self.domain_left_edge = np.zeros(3, "float64")
         self.domain_right_edge = np.ones(3, "float64") * hvals["BoxSize"]
         self.domain_dimensions = np.ones(3, "int32") * self._root_dimensions
+        self.periodicity = (True, True, True)
 
         self.cosmological_simulation = 1
 
         self.domain_left_edge = np.zeros(3, "float64") - 0.5
         self.domain_right_edge = np.ones(3, "float64") + 0.5
         self.domain_dimensions = np.ones(3, "int32") * self._root_dimensions
+        self.periodicity = (True, True, True)
 
         self.cosmological_simulation = 1
 

File yt/geometry/oct_container.pyx

     @cython.cdivision(True)
     def icoords(self, int domain_id,
                 np.ndarray[np.uint8_t, ndim=2, cast=True] mask,
-                np.int64_t cell_count):
+                np.int64_t cell_count,
+                np.ndarray[np.int64_t, ndim=1] level_counts):
         #Return the integer positions of the cells
         #Limited to this domain and within the mask
         #Positions are binary; aside from the root mesh
     @cython.cdivision(True)
     def ires(self, int domain_id,
                 np.ndarray[np.uint8_t, ndim=2, cast=True] mask,
-                np.int64_t cell_count):
+                np.int64_t cell_count,
+                np.ndarray[np.int64_t, ndim=1] level_counts):
         #Return the 'resolution' of each cell; ie the level
         cdef np.ndarray[np.int64_t, ndim=1] res
         res = np.empty(cell_count, dtype="int64")
     @cython.cdivision(True)
     def fcoords(self, int domain_id,
                 np.ndarray[np.uint8_t, ndim=2, cast=True] mask,
-                np.int64_t cell_count):
+                np.int64_t cell_count,
+                np.ndarray[np.int64_t, ndim=1] level_counts):
         #Return the floating point unitary position of every cell
         cdef np.ndarray[np.float64_t, ndim=2] coords
         coords = np.empty((cell_count, 3), dtype="float64")
                 count[o.domain] += mask[oi,i]
         return count
 
+    def domain_and(self, np.ndarray[np.uint8_t, ndim=2, cast=True] mask,
+                   int domain_id):
+        cdef np.int64_t i, oi, n, 
+        cdef Oct *o
+        cdef np.ndarray[np.uint8_t, ndim=2] m2 = \
+                np.zeros((mask.shape[0], 8), 'uint8')
+        n = mask.shape[0]
+        for oi in range(n):
+            o = self.oct_list[oi]
+            if o.domain != domain_id: continue
+            for i in range(8):
+                m2[o.local_ind, i] = mask[o.local_ind, i]
+        return m2