Commits

Matthew Turk committed 05a9950

First draft of .deposit() for OctreeSubset objects.

This includes a few fixes to how particles are assigned to ParticleOctree
elements.

Comments (0)

Files changed (5)

yt/data_objects/octree_subset.py

     NeedsDataField, \
     NeedsProperty, \
     NeedsParameter
+import yt.geometry.particle_deposit as particle_deposit
 
 class OctreeSubset(YTSelectionContainer):
     _spatial = True
             return tr
         finfo = self.pf._get_field_info(*fields[0])
         if not finfo.particle_type:
-            nz = self._num_zones + 2*self._num_ghost_zones
             # We may need to reshape the field, if it is being queried from
             # field_data.  If it's already cached, it just passes through.
-            if len(tr.shape) < 4: 
-                n_oct = tr.shape[0] / (nz**3.0)
-                tr.shape = (n_oct, nz, nz, nz)
-                tr = np.rollaxis(tr, 0, 4)
+            if len(tr.shape) < 4:
+                tr = self._reshape_vals(tr)
             return tr
         return tr
 
+    def _reshape_vals(self, arr):
+        nz = self._num_zones + 2*self._num_ghost_zones
+        n_oct = arr.shape[0] / (nz**3.0)
+        arr.shape = (n_oct, nz, nz, nz)
+        arr = np.rollaxis(arr, 0, 4)
+        return arr
+
     _domain_ind = None
 
     @property
             self._domain_ind = di
         return self._domain_ind
 
-    def deposit(self, positions, fields, method):
+    def deposit(self, positions, fields = None, method = None):
         # Here we perform our particle deposition.
-        pass
+        cls = getattr(particle_deposit, "deposit_%s" % method, None)
+        if cls is None:
+            raise YTParticleDepositionNotImplemented(method)
+        nvals = (self.domain_ind >= 0).sum() * 8
+        op = cls(nvals) # We allocate number of zones, not number of octs
+        op.initialize()
+        op.process_octree(self.oct_handler, self.domain_ind, positions, fields)
+        vals = op.finalize()
+        return self._reshape_vals(vals)
 
     def select(self, selector):
         if id(selector) == self._last_selector_id:

yt/geometry/oct_container.pxd

     Oct *children[2][2][2]
     Oct *parent
 
+cdef struct OctInfo:
+    np.float64_t left_edge[3]
+    np.float64_t dds[3]
+
 cdef struct OctAllocationContainer
 cdef struct OctAllocationContainer:
     np.int64_t n
     cdef np.float64_t DLE[3], DRE[3]
     cdef public int nocts
     cdef public int max_domain
-    cdef Oct* get(self, np.float64_t ppos[3], int *ii = ?)
+    cdef Oct* get(self, np.float64_t ppos[3], OctInfo *oinfo = ?)
     cdef void neighbors(self, Oct *, Oct **)
     cdef void oct_bounds(self, Oct *, np.float64_t *, np.float64_t *)
 

yt/geometry/oct_container.pyx

     @cython.boundscheck(False)
     @cython.wraparound(False)
     @cython.cdivision(True)
-    cdef Oct *get(self, np.float64_t ppos[3], int *ii = NULL):
+    cdef Oct *get(self, np.float64_t ppos[3], OctInfo *oinfo = NULL):
         #Given a floating point position, retrieve the most
         #refined oct at that time
         cdef np.int64_t ind[3]
         cdef Oct *cur
         cdef int i
         for i in range(3):
-            pp[i] = ppos[i] - self.DLE[i]
             dds[i] = (self.DRE[i] - self.DLE[i])/self.nn[i]
-            ind[i] = <np.int64_t> ((pp[i] - self.DLE[i])/dds[i])
-            cp[i] = (ind[i] + 0.5) * dds[i]
+            ind[i] = <np.int64_t> ((ppos[i] - self.DLE[i])/dds[i])
+            cp[i] = (ind[i] + 0.5) * dds[i] + self.DLE[i]
         cur = self.root_mesh[ind[0]][ind[1]][ind[2]]
         while cur.children[0][0][0] != NULL:
             for i in range(3):
                 dds[i] = dds[i] / 2.0
-                if cp[i] > pp[i]:
+                if cp[i] > ppos[i]:
                     ind[i] = 0
                     cp[i] -= dds[i] / 2.0
                 else:
                     ind[i] = 1
                     cp[i] += dds[i]/2.0
             cur = cur.children[ind[0]][ind[1]][ind[2]]
-        if ii != NULL: return cur
+        if oinfo == NULL: return cur
         for i in range(3):
-            if cp[i] > pp[i]:
-                ind[i] = 0
-            else:
-                ind[i] = 1
-        ii[0] = ((ind[2]*2)+ind[1])*2+ind[0]
+            oinfo.dds[i] = dds[i] # Cell width
+            oinfo.left_edge[i] = cp[i] - dds[i]
         return cur
 
     @cython.boundscheck(False)
 
 
 cdef class ARTOctreeContainer(RAMSESOctreeContainer):
-    #this class is specifically for the NMSU ART
-    @cython.boundscheck(False)
-    @cython.wraparound(False)
-    @cython.cdivision(True)
-    def deposit_particle_cumsum(self,
-                                np.ndarray[np.float64_t, ndim=2] ppos, 
-                                np.ndarray[np.float64_t, ndim=1] pdata,
-                                np.ndarray[np.float64_t, ndim=1] mask,
-                                np.ndarray[np.float64_t, ndim=1] dest,
-                                fields, int domain):
-        cdef Oct *o
-        cdef OctAllocationContainer *dom = self.domains[domain - 1]
-        cdef np.float64_t pos[3]
-        cdef int ii
-        cdef int no = ppos.shape[0]
-        for n in range(no):
-            for j in range(3):
-                pos[j] = ppos[n,j]
-            o = self.get(pos, &ii) 
-            if mask[o.local_ind,ii]==0: continue
-            dest[o.ind+ii] += pdata[n]
-        return dest
 
     @cython.boundscheck(True)
     @cython.wraparound(False)
         cdef int max_level = 0
         self.oct_list = <Oct**> malloc(sizeof(Oct*)*self.nocts)
         cdef np.int64_t i = 0
+        cdef np.int64_t dom_ind
         cdef ParticleArrays *c = self.first_sd
         while c != NULL:
             self.oct_list[i] = c.oct
         self.dom_offsets = <np.int64_t *>malloc(sizeof(np.int64_t) *
                                                 (self.max_domain + 3))
         self.dom_offsets[0] = 0
+        dom_ind = 0
         for i in range(self.nocts):
             self.oct_list[i].local_ind = i
+            self.oct_list[i].ind = dom_ind
+            dom_ind += 1
             if self.oct_list[i].domain > cur_dom:
                 cur_dom = self.oct_list[i].domain
                 self.dom_offsets[cur_dom + 1] = i
+                dom_ind = 0
         self.dom_offsets[cur_dom + 2] = self.nocts
 
     cdef Oct* allocate_oct(self):
         # index into the domain subset array for deposition.
         cdef np.int64_t i, j, k, oi, noct, n, nm, use, offset
         cdef Oct *o
-        offset = self.dom_offsets[domain_id]
-        noct = self.dom_offsets[domain_id + 1] - offset
+        # For particle octrees, domain 0 is special and means non-leaf nodes.
+        offset = self.dom_offsets[domain_id + 1]
+        noct = self.dom_offsets[domain_id + 2] - offset
         cdef np.ndarray[np.int64_t, ndim=1] ind = np.zeros(noct, 'int64')
         nm = 0
         for oi in range(noct):

yt/geometry/particle_deposit.pyx

 cimport cython
 
 from fp_utils cimport *
-from oct_container cimport Oct, OctAllocationContainer, OctreeContainer
+from oct_container cimport Oct, OctAllocationContainer, \
+    OctreeContainer, OctInfo
 
 cdef class ParticleDepositOperation:
     def __init__(self, nvals):
                      np.ndarray[np.int64_t, ndim=1] dom_ind,
                      np.ndarray[np.float64_t, ndim=2] positions,
                      fields = None):
-        raise NotImplementedError
-
+        cdef int nf, i, j
+        if fields is None:
+            fields = []
+        nf = len(fields)
+        cdef np.float64_t **field_pointers, *field_vals, pos[3]
+        cdef np.ndarray[np.float64_t, ndim=1] tarr
+        field_pointers = <np.float64_t**> alloca(sizeof(np.float64_t *) * nf)
+        field_vals = <np.float64_t*>alloca(sizeof(np.float64_t) * nf)
+        for i in range(nf):
+            tarr = fields[i]
+            field_pointers[i] = <np.float64_t *> tarr.data
+        cdef int dims[3]
+        dims[0] = dims[1] = dims[2] = 2
+        cdef OctInfo oi
+        cdef np.int64_t offset
+        cdef Oct *oct
+        for i in range(positions.shape[0]):
+            # We should check if particle remains inside the Oct here
+            for j in range(nf):
+                field_vals[j] = field_pointers[j][i]
+            for j in range(3):
+                pos[j] = positions[i, j]
+            oct = octree.get(pos, &oi)
+            #print oct.local_ind, oct.pos[0], oct.pos[1], oct.pos[2]
+            offset = dom_ind[oct.ind]
+            self.process(dims, oi.left_edge, oi.dds,
+                         offset, pos, field_vals)
+        
     def process_grid(self, gobj,
                      np.ndarray[np.float64_t, ndim=2] positions,
                      fields = None):
 
 cdef class CountParticles(ParticleDepositOperation):
     cdef np.float64_t *count # float, for ease
-    cdef object ocount
+    cdef public object ocount
     def initialize(self):
         self.ocount = np.zeros(self.nvals, dtype="float64")
         cdef np.ndarray arr = self.ocount
         self.count = <np.float64_t*> arr.data
 
+    @cython.cdivision(True)
     cdef void process(self, int dim[3],
                       np.float64_t left_edge[3], 
                       np.float64_t dds[3],
         cdef int ii[3], i
         for i in range(3):
             ii[i] = <int>((ppos[i] - left_edge[i])/dds[i])
-        self.count[gind(ii[0], ii[1], ii[2], dim)] += 1
+        #print "Depositing into", offset,
+        #print gind(ii[0], ii[1], ii[2], dim)
+        self.count[gind(ii[0], ii[1], ii[2], dim) + offset] += 1
         
     def finalize(self):
         return self.ocount
 
+deposit_count = CountParticles
+
 """
 # Mode functions
 ctypedef np.float64_t (*type_opt)(np.float64_t, np.float64_t)

yt/utilities/exceptions.py

     def __str__(self):
         return "Data selector '%s' not implemented." % (self.class_name)
 
+class YTParticleDepositionNotImplemented(YTException):
+    def __init__(self, class_name):
+        self.class_name = class_name
+
+    def __str__(self):
+        return "Particle deposition method '%s' not implemented." % (self.class_name)
+
 class YTDomainOverflow(YTException):
     def __init__(self, mi, ma, dle, dre):
         self.mi = mi