Commits

Matthew Turk committed 05d417a

Adding domain_mask for particle octrees.

  • Participants
  • Parent commits b7949c2

Comments (0)

Files changed (1)

yt/geometry/oct_container.pyx

         # 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 np.int64_t i, j, k, oi, n, nm, use
         cdef OctAllocationContainer *cur = self.domains[domain_id - 1]
         cdef Oct *o
         n = mask.shape[0]
                 m2[o.local_ind, i] = mask[o.local_ind, i]
         return m2
 
+    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, use
+        cdef Oct *o
+        n = mask.shape[0]
+        nm = 0
+        for oi in range(n):
+            o = self.oct_list[oi]
+            if o.domain != domain_id: continue
+            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(n):
+            o = self.oct_list[oi]
+            if o.domain != domain_id: continue
+            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")