Matthew Turk avatar Matthew Turk committed b7949c2

Spatial chunking within data objects for Octree codes now works.

Comments (0)

Files changed (3)

yt/data_objects/octree_subset.py

         level_counts[1:] = level_counts[:-1]
         level_counts[0] = 0
         self.level_counts = np.add.accumulate(level_counts)
+        self._last_mask = None
+        self._last_selector_id = None
         self._current_particle_type = 'all'
         self._current_fluid_type = self.pf.default_fluid_type
 
         if not finfo.particle_type:
             nz = self._num_zones + 2*self._num_ghost_zones
             n_oct = tr.shape[0] / (nz**3.0)
-            dest_shape = (nz, nz, nz, n_oct)
-            return tr.reshape(dest_shape)
+            tr.shape = (n_oct, nz, nz, nz)
+            tr = np.rollaxis(tr, 0, 4)
+            return tr
         return tr
 
+    def deposit(self, positions, fields, method):
+        pass
 
+    def select(self, selector):
+        if id(selector) == self._last_selector_id:
+            return self._last_mask
+        self._last_mask = self.oct_handler.domain_mask(
+                self.mask, self.domain.domain_id)
+        if self._last_mask.sum() == 0: return None
+        self._last_selector_id = id(selector)
+        return self._last_mask
+
+    def count(self, selector):
+        if id(selector) == self._last_selector_id:
+            if self._last_mask is None: return 0
+            return self._last_mask.sum()
+        self.select(selector)
+        return self.count(selector)
+
+    def count_particles(self, selector, x, y, z):
+        # We don't cache the selector results
+        count = selector.count_points(x,y,z)
+        return count
+
+    def select_particles(self, selector, x, y, z):
+        mask = selector.select_points(x,y,z)
+        return mask

yt/geometry/oct_container.pyx

             o = &cur.my_octs[oi]
             for i in range(8):
                 m2[o.local_ind, i] = mask[o.local_ind, i]
-        return m2
+        return m2 # NOTE: This is uint8_t
+
+    def domain_mask(self,
+                    # mask is the base selector's *global* mask
+                    np.ndarray[np.uint8_t, ndim=2, cast=True] mask,
+                    int domain_id):
+        # What distinguishes this one from domain_and is that we have a mask,
+        # which covers the whole domain, but our output will only be of a much
+        # smaller subset of octs that belong to a given domain *and* the mask.
+        # Note also that typically when something calls domain_and, they will 
+        # use a logical_any along the oct axis.  Here we don't do that.
+        # Note also that we change the shape of the returned array.
+        cdef np.int64_t i, j, k, oi, n, nm
+        cdef OctAllocationContainer *cur = self.domains[domain_id - 1]
+        cdef Oct *o
+        n = mask.shape[0]
+        nm = 0
+        for oi in range(cur.n_assigned):
+            o = &cur.my_octs[oi]
+            use = 0
+            for i in range(8):
+                if mask[o.local_ind, i] == 1: use = 1
+            nm += use
+        cdef np.ndarray[np.uint8_t, ndim=4] m2 = \
+                np.zeros((2, 2, 2, nm), 'uint8')
+        nm = 0
+        for oi in range(cur.n_assigned):
+            o = &cur.my_octs[oi]
+            use = 0
+            for i in range(2):
+                for j in range(2):
+                    for k in range(2):
+                        ii = ((k*2)+j)*2+i
+                        if mask[o.local_ind, ii] == 0: continue
+                        use = m2[i, j, k, nm] = 1
+            nm += use
+        return m2.astype("bool")
 
     def check(self, int curdom):
         cdef int dind, pi

yt/geometry/selection_routines.pyx

                          int eterm[3]) nogil:
         return 1
 
-
 octree_subset_selector = OctreeSubsetSelector
 
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.