Issue 2281

nose_ignores

@@ -47,6 +47,7 @@
 --ignore-file=test_field_parsing\.py
 --ignore-file=test_disks\.py
 --ignore-file=test_offaxisprojection_pytestonly\.py
+--ignore-file=test_process_octree_locality_pytest\.py
 --ignore-file=test_sph_pixelization_pytestonly\.py
 --ignore-file=test_time_series\.py
 --ignore-file=test_cf_radial_pytest\.py

yt/geometry/fake_octree.pyx

@@ -36,13 +36,13 @@ def create_fake_octree(SparseOctreeContainer oct_handler,
     for i in range(3):
         ind[i] = 0
         dd[i] = ndd[i]
-    oct_handler.allocate_domains([max_noct])
+    oct_handler.allocate_domains([max_noct], 1)
     parent = oct_handler.next_root(1, ind)
     parent.domain = 1
     cur_leaf = 8 #we've added one parent...
     mask = np.ones((max_noct,8),dtype='uint8')
-    while oct_handler.domains[0].n_assigned < max_noct:
-        print("root: nocts ", oct_handler.domains[0].n_assigned)
+    while oct_handler.nocts < max_noct:
+        print("root: nocts ", oct_handler.nocts)
         cur_leaf = subdivide(oct_handler, parent, ind, dd, cur_leaf, 0,
                              max_noct, max_level, fsubdivide, mask)
     return cur_leaf
@@ -61,15 +61,15 @@ cdef long subdivide(SparseOctreeContainer oct_handler,
     cdef float rf #random float from 0-1
     if cur_level >= max_level:
         return cur_leaf
-    if oct_handler.domains[0].n_assigned >= max_noct:
+    if oct_handler.nocts >= max_noct:
         return cur_leaf
     for i in range(3):
         ind[i] = <int> ((rand() * 1.0 / RAND_MAX) * dd[i])
         ddr[i] = 2
     rf = rand() * 1.0 / RAND_MAX
     if rf > fsubdivide:
         ii = cind(ind[0], ind[1], ind[2])
-        if parent.children[ii] == NULL:
+        if parent.children == NULL or parent.children[ii] == NULL:
             cur_leaf += 7
         oct = oct_handler.next_child(1, ind, parent)
         oct.domain = 1

yt/geometry/oct_container.pxd

@@ -64,6 +64,9 @@ cdef class OctreeContainer:
     cdef public int num_domains
     cdef Oct *get(self, np.float64_t ppos[3], OctInfo *oinfo = ?,
                   int max_level = ?) noexcept nogil
+    cdef Oct *get_near(self, Oct *start_oct, OctInfo *start_oi, 
+                      np.float64_t ppos[3], OctInfo *oinfo = ?,
+                      int max_level = ?) noexcept nogil
     cdef int get_root(self, int ind[3], Oct **o) noexcept nogil
     cdef Oct **neighbors(self, OctInfo *oinfo, np.int64_t *nneighbors,
                          Oct *o, bint periodicity[3])

yt/geometry/oct_container.pyx

@@ -126,6 +126,7 @@ cdef class OctreeContainer:
                     o = &cur.my_objs[cur.n_assigned]
                     o.domain_ind = o.file_ind = 0
                     o.domain = 1
+                    o.parent = NULL
                     obj.root_mesh[i][j][k] = o
                     cur.n_assigned += 1
                     visitor.pos[0] = i
@@ -265,6 +266,79 @@ cdef class OctreeContainer:
             oinfo.left_edge[i] = oinfo.ipos[i] * (oinfo.dds[i] * self.nz[i]) + self.DLE[i]
         oinfo.level = level
         return cur
+
+    @cython.boundscheck(False)
+    @cython.wraparound(False)
+    @cython.cdivision(True)
+    cdef Oct *get_near(self, Oct *start_oct, OctInfo *start_oi,
+                        np.float64_t ppos[3], OctInfo *oinfo = NULL,
+                        int max_level = 99) noexcept nogil:
+        cdef int i
+        cdef Oct *cur
+        cdef Oct *next
+        cdef np.float64_t dds[3]
+        cdef np.float64_t left_edge[3]
+        cdef np.float64_t cand_dds[3]
+        cdef np.float64_t cand_left_edge[3]
+        cdef np.float64_t mid
+        cdef np.int64_t ind[3]
+        cdef np.int64_t level
+        cdef bint contains
+        #instead of starting at root, we start at given oct
+        cur = start_oct
+        for i in range(3):
+            dds[i] = start_oi.dds[i] * self.nz[i]
+            left_edge[i] = start_oi.left_edge[i]
+        level = start_oi.level
+
+        #walk up towards root, doubling each step until box contains ppos
+        while level > 0:
+            contains = True
+            for i in range(3):
+                if ppos[i] < left_edge[i] or ppos[i] >= left_edge[i] + dds[i]:
+                    contains = False
+                    break
+            if contains:
+                break
+            cur = cur.parent
+            for i in range(3):
+                dds[i] = dds[i] * 2.0
+                left_edge[i] = floor((left_edge[i] - self.DLE[i]) / dds[i]) * dds[i] + self.DLE[i]
+            level -= 1
+        #if still not in same root cell as before, start over
+        contains = True
+        for i in range(3):
+            if ppos[i] < left_edge[i] or ppos[i] >= left_edge[i] + dds[i]:
+                contains = False
+                break
+        if not contains:
+            return self.get(ppos, oinfo, max_level)
+        #walk back down toward ppos from wherever we ended up
+        while cur.children != NULL and level < max_level:
+            for i in range(3):
+                cand_dds[i] = dds[i] / 2.0
+                mid = left_edge[i] + cand_dds[i]
+                if mid > ppos[i]:
+                    ind[i] = 0
+                    cand_left_edge[i] = left_edge[i]
+                else:
+                    ind[i] = 1
+                    cand_left_edge[i] = mid
+            next = cur.children[cind(ind[0], ind[1], ind[2])]
+            if next == NULL:
+                break
+            cur = next
+            level += 1
+            for i in range(3):
+                dds[i] = cand_dds[i]
+                left_edge[i] = cand_left_edge[i]
+        if oinfo == NULL: return cur
+        for i in range(3):
+            oinfo.dds[i] = dds[i] / self.nz[i]
+            oinfo.left_edge[i] = left_edge[i]
+            oinfo.ipos[i] = 0 #not used by particle_deposit loop
+        oinfo.level = level
+        return cur
 
     def locate_positions(self, np.float64_t[:,:] positions):
         """
@@ -665,6 +739,7 @@ cdef class OctreeContainer:
         next = &cont.my_objs[cont.n_assigned]
         cont.n_assigned += 1
         self.root_mesh[ind[0]][ind[1]][ind[2]] = next
+        next.parent = NULL
         self.nocts += 1
         return next
 
@@ -684,6 +759,7 @@ cdef class OctreeContainer:
         next = &cont.my_objs[cont.n_assigned]
         cont.n_assigned += 1
         parent.children[cind(ind[0],ind[1],ind[2])] = next
+        next.parent = parent
         self.nocts += 1
         return next
 
@@ -1097,6 +1173,7 @@ cdef class SparseOctreeContainer(OctreeContainer):
         tsearch(<void*>ikey, &self.tree_root, root_node_compare)
         self.num_root += 1
         self.nocts += 1
+        next.parent = NULL
         return next
 
     def allocate_domains(self, domain_counts, int root_nodes):
@@ -1217,6 +1294,7 @@ cdef class OctObjectPool(ObjectPool):
             octs[n].file_ind = octs[n].domain = - 1
             octs[n].domain_ind = n + offset
             octs[n].children = NULL
+            octs[n].parent = NULL
 
     cdef void teardown_objs(self, void *obj, np.uint64_t n, np.uint64_t offset,
                            np.int64_t con_id):

yt/geometry/oct_visitors.pxd

@@ -17,6 +17,7 @@ cdef struct Oct:
     np.int64_t domain_ind   # index within the global set of domains
     np.int64_t domain       # (opt) addl int index
     Oct **children          # Up to 8 long
+    Oct *parent             # NULL for root octs
 
 cdef struct OctInfo:
     np.float64_t left_edge[3]

yt/geometry/oct_visitors.pyx

@@ -320,6 +320,7 @@ cdef class LoadOctree(OctVisitor):
                 o.children[ii + i].file_ind = -1
                 o.children[ii + i].domain = -1
                 o.children[ii + i].children = NULL
+                o.children[ii + i].parent = o
                 self.nocts[0] += 1
         else:
             print("SOMETHING IS AMISS", self.index)

yt/geometry/particle_deposit.pyx

@@ -67,6 +67,10 @@ cdef class ParticleDepositOperation:
         cdef np.int64_t offset, moff
         cdef Oct *oct
         cdef np.int8_t use_lvlmax
+        cdef int this_max_level
+        cdef Oct *cached_oct = NULL
+        cdef OctInfo cached_oi
+        cdef bint in_cached_oct
         moff = octree.get_domain_offset(domain_id + domain_offset)
         if lvlmax is None:
             use_lvlmax = False
@@ -87,10 +91,33 @@ cdef class ParticleDepositOperation:
             # previously generated.  This way we can support not knowing the
             # full octree structure.  All we *really* care about is some
             # arbitrary offset into a field value for deposition.
-            if not use_lvlmax:
-                oct = octree.get(pos, &oi)
+
+            #only cache particles if max level is equal or looser with no children
+            #comparing the limit avoids wrong answers where refinement may be deeper.
+            this_max_level = lvlmaxval[i] if use_lvlmax else 99
+            in_cached_oct = False
+            if cached_oct != NULL and this_max_level >= cached_oi.level and \
+                not (this_max_level > cached_oi.level and cached_oct.children != NULL):
+
+                in_cached_oct = True
+                for j in range(3):
+                    if pos[j] < cached_oi.left_edge[j] or \
+                        pos[j] >= cached_oi.left_edge[j] + cached_oi.dds[j] * dims[j]:
+                        in_cached_oct = False
+                        break
+            if in_cached_oct:
+                oct = cached_oct
+                oi = cached_oi
+            elif cached_oct != NULL:
+                #different oct, find it but reuse last traversal as starting point
+                oct = octree.get_near(cached_oct, &cached_oi, pos, &oi, this_max_level)
+                cached_oct = oct
+                cached_oi = oi
             else:
-                oct = octree.get(pos, &oi, max_level=lvlmaxval[i])
+                #should be first particle
+                oct = octree.get(pos, &oi, max_level=this_max_level)
+                cached_oct = oct
+                cached_oi = oi
             # This next line is unfortunate.  Basically it says, sometimes we
             # might have particles that belong to octs outside our domain.
             # For the distributed-memory octrees, this will manifest as a NULL