diff --git a/yt/data_objects/index_subobjects/octree_subset.py b/yt/data_objects/index_subobjects/octree_subset.py index d450ebaa3ec..410890e1986 100644 --- a/yt/data_objects/index_subobjects/octree_subset.py +++ b/yt/data_objects/index_subobjects/octree_subset.py @@ -46,7 +46,10 @@ class OctreeSubset(YTSelectionContainer, abc.ABC): def __init__(self, base_region, domain, ds, num_zones=2, num_ghost_zones=0): super().__init__(ds, None) - self._num_zones = num_zones + if hasattr(num_zones, "__len__"): + self._num_zones = np.array(num_zones, dtype="int64") + else: + self._num_zones = np.array([num_zones, num_zones, num_zones], dtype="int64") self._num_ghost_zones = num_ghost_zones self.domain = domain self.domain_id = domain.domain_id @@ -80,23 +83,28 @@ def __getitem__(self, key): @property def nz(self): - return self._num_zones + 2 * self._num_ghost_zones + nz = self._num_zones + 2 * self._num_ghost_zones + if hasattr(nz, "__len__"): + return nz + return np.array([nz, nz, nz], dtype="int64") def get_bbox(self): return self.base_region.get_bbox() def _reshape_vals(self, arr): nz = self.nz + nzx, nzy, nzz = nz[0], nz[1], nz[2] + nzones = nzx * nzy * nzz if len(arr.shape) <= 2: - n_oct = arr.shape[0] // (nz**3) + n_oct = arr.shape[0] // nzones elif arr.shape[-1] == 3: n_oct = arr.shape[-2] else: n_oct = arr.shape[-1] - if arr.size == nz * nz * nz * n_oct: - new_shape = (nz, nz, nz, n_oct) - elif arr.size == nz * nz * nz * n_oct * 3: - new_shape = (nz, nz, nz, n_oct, 3) + if arr.size == nzones * n_oct: + new_shape = (nzx, nzy, nzz, n_oct) + elif arr.size == nzones * n_oct * 3: + new_shape = (nzx, nzy, nzz, n_oct, 3) else: raise RuntimeError # Note that if arr is already F-contiguous, this *shouldn't* copy the @@ -172,7 +180,7 @@ def deposit(self, positions, fields=None, method=None, kernel_name="cubic"): if cls is None: raise YTParticleDepositionNotImplemented(method) nz = self.nz - nvals = (nz, nz, nz, (self.domain_ind >= 0).sum()) + nvals = (int(nz[0]), int(nz[1]), int(nz[2]), (self.domain_ind >= 0).sum()) if np.max(self.domain_ind) >= nvals[-1]: print( f"nocts, domain_ind >= 0, max {self.oct_handler.nocts} {nvals[-1]} {np.max(self.domain_ind)}" @@ -335,7 +343,7 @@ def smooth( [1, 1, 1], self.ds.domain_left_edge, self.ds.domain_right_edge, - num_zones=self._nz, + num_zones=self._num_zones, ) # This should ensure we get everything within one neighbor of home. particle_octree.n_ref = nneighbors * 2 @@ -354,7 +362,7 @@ def smooth( raise YTParticleDepositionNotImplemented(method) nz = self.nz mdom_ind = self.domain_ind - nvals = (nz, nz, nz, (mdom_ind >= 0).sum()) + nvals = (int(nz[0]), int(nz[1]), int(nz[2]), (mdom_ind >= 0).sum()) op = cls(nvals, len(fields), nneighbors, kernel_name) op.initialize() mylog.debug( @@ -455,7 +463,7 @@ def particle_operation( raise YTParticleDepositionNotImplemented(method) nz = self.nz mdom_ind = self.domain_ind - nvals = (nz, nz, nz, (mdom_ind >= 0).sum()) + nvals = (int(nz[0]), int(nz[1]), int(nz[2]), (mdom_ind >= 0).sum()) op = cls(nvals, len(fields), nneighbors, kernel_name) op.initialize() mylog.debug( @@ -548,7 +556,7 @@ def __init__(self, ind, block_slice): self.ind = ind self.block_slice = block_slice nz = self.block_slice.octree_subset.nz - self.ActiveDimensions = np.array([nz, nz, nz], dtype="int64") + self.ActiveDimensions = np.array([nz[0], nz[1], nz[2]], dtype="int64") self.ds = block_slice.ds def __getitem__(self, key): diff --git a/yt/data_objects/tests/test_octree.py b/yt/data_objects/tests/test_octree.py index 06652612f04..0c52ed337ec 100644 --- a/yt/data_objects/tests/test_octree.py +++ b/yt/data_objects/tests/test_octree.py @@ -1,6 +1,7 @@ import numpy as np from numpy.testing import assert_almost_equal, assert_equal +from yt.geometry.oct_container import OctreeContainer from yt.testing import fake_sph_grid_ds n_ref = 4 @@ -118,3 +119,26 @@ def test_octree_properties(): refined = octree["index", "refined"] refined_ans = np.array([True] + [False] * 7 + [True] + [False] * 8, dtype=np.bool_) assert_equal(refined, refined_ans) + + +def test_num_zones_tuple(): + """ + Test that OctreeContainer accepts num_zones as a scalar or a tuple (N, M, L). + Both should correctly set per-dimension zone counts. + """ + # Scalar: all dimensions equal + oct_scalar = OctreeContainer( + [1, 1, 1], [0.0, 0.0, 0.0], [1.0, 1.0, 1.0], num_zones=2 + ) + # Tuple: potentially different per-dimension + oct_tuple = OctreeContainer( + [1, 1, 1], [0.0, 0.0, 0.0], [1.0, 1.0, 1.0], num_zones=(2, 2, 2) + ) + # Non-uniform tuple + oct_nonuniform = OctreeContainer( + [1, 1, 1], [0.0, 0.0, 0.0], [1.0, 1.0, 1.0], num_zones=(2, 3, 4) + ) + # Verify that creating these containers doesn't raise exceptions + assert oct_scalar is not None + assert oct_tuple is not None + assert oct_nonuniform is not None diff --git a/yt/frontends/artio/data_structures.py b/yt/frontends/artio/data_structures.py index 2338d40e002..a9d80866287 100644 --- a/yt/frontends/artio/data_structures.py +++ b/yt/frontends/artio/data_structures.py @@ -128,7 +128,7 @@ def deposit(self, positions, fields=None, method=None, kernel_name="cubic"): if cls is None: raise YTParticleDepositionNotImplemented(method) nz = self.nz - nvals = (nz, nz, nz, self.ires.size) + nvals = (int(nz[0]), int(nz[1]), int(nz[2]), self.ires.size) # We allocate number of zones, not number of octs op = cls(nvals, kernel_name) op.initialize() @@ -241,6 +241,7 @@ def _identify_base_chunk(self, dobj): sfc_start = getattr(dobj, "sfc_start", None) sfc_end = getattr(dobj, "sfc_end", None) nz = getattr(dobj, "_num_zones", 0) + nz_scalar = int(np.asarray(nz).flat[0]) if hasattr(nz, "__len__") else nz if all_data: mylog.debug("Selecting entire artio domain") list_sfc_ranges = self.ds._handle.root_sfc_ranges_all( @@ -271,7 +272,7 @@ def _identify_base_chunk(self, dobj): ) range_handler.construct_mesh() self.range_handlers[start, end] = range_handler - if nz != 2: + if nz_scalar != 2: ci.append( ARTIORootMeshSubset( base_region, @@ -281,7 +282,7 @@ def _identify_base_chunk(self, dobj): self.ds, ) ) - if nz != 1 and range_handler.total_octs > 0: + if nz_scalar != 1 and range_handler.total_octs > 0: ci.append( ARTIOOctreeSubset( base_region, diff --git a/yt/frontends/ramses/data_structures.py b/yt/frontends/ramses/data_structures.py index 4918d31ea2a..e6c9aaa51f1 100644 --- a/yt/frontends/ramses/data_structures.py +++ b/yt/frontends/ramses/data_structures.py @@ -467,8 +467,8 @@ def _fill_with_ghostzones( fields = [f for ft, f in fields] tr = {} - cell_count = ( - selector.count_octs(self.oct_handler, self.domain_id) * self.nz**ndim + cell_count = selector.count_octs(self.oct_handler, self.domain_id) * int( + np.prod(self.nz[:ndim]) ) # Initializing data container @@ -520,7 +520,7 @@ def fwidth(self): # new_fwidth contains the fwidth of the oct+ghost zones # this is a constant array in each oct, so we simply copy # the oct value using numpy fancy-indexing - new_fwidth = np.zeros((n_oct, self.nz**3, 3), dtype=fwidth.dtype) + new_fwidth = np.zeros((n_oct, int(np.prod(self.nz)), 3), dtype=fwidth.dtype) new_fwidth[:, :, :] = fwidth[:, 0:1, :] fwidth = new_fwidth.reshape(-1, 3) return fwidth @@ -538,7 +538,7 @@ def fcoords(self): self.selector, self._num_ghost_zones ) - N_per_oct = self.nz**3 + N_per_oct = int(np.prod(self.nz)) oct_inds = oct_inds.reshape(-1, N_per_oct) cell_inds = cell_inds.reshape(-1, N_per_oct) diff --git a/yt/frontends/stream/data_structures.py b/yt/frontends/stream/data_structures.py index 241c405a55d..2214833748d 100644 --- a/yt/frontends/stream/data_structures.py +++ b/yt/frontends/stream/data_structures.py @@ -832,8 +832,8 @@ def _fill_no_ghostzones(self, content, dest, selector, offset): def _fill_with_ghostzones(self, content, dest, selector, offset): oct_handler = self.oct_handler ndim = self.ds.dimensionality - cell_count = ( - selector.count_octs(self.oct_handler, self.domain_id) * self.nz**ndim + cell_count = selector.count_octs(self.oct_handler, self.domain_id) * int( + np.prod(self.nz[:ndim]) ) gz_cache = getattr(self, "_ghost_zone_cache", None) diff --git a/yt/geometry/_selection_routines/selector_object.pxi b/yt/geometry/_selection_routines/selector_object.pxi index 7f1cd54f324..df08ffd8232 100644 --- a/yt/geometry/_selection_routines/selector_object.pxi +++ b/yt/geometry/_selection_routines/selector_object.pxi @@ -159,7 +159,7 @@ cdef class SelectorObject: visitor.pos[1] = (visitor.pos[1] >> 1) visitor.pos[2] = (visitor.pos[2] >> 1) visitor.level -= 1 - elif this_level == 1 and visitor.nz > 1: + elif this_level == 1 and (visitor.nz[0] > 1 or visitor.nz[1] > 1 or visitor.nz[2] > 1): visitor.global_index += increment increment = 0 self.visit_oct_cells(root, ch, spos, sdds, @@ -178,10 +178,22 @@ cdef class SelectorObject: cdef void visit_oct_cells(self, Oct *root, Oct *ch, np.float64_t spos[3], np.float64_t sdds[3], OctVisitor visitor, int i, int j, int k): - # We can short-circuit the whole process if data.nz == 2. - # This saves us some funny-business. + """Visit the cells in this oct. + + Parameters + ---------- + root: The oct whose cells we are visiting. + ch: The child oct, if it exists. + spos: The position of a potential cell center, assuming that the + oct contains 8 cells. + sdds: The cell size, assuming that the oct contains 8 cells. + visitor: The visitor object that is visiting the cells. + i, j, k: The indices of the cell within the oct. + """ cdef int selected - if visitor.nz == 2: + # If visitor.nz is 2 in all dimensions, then the passed spos and sdds + # are correct and we just need to call `select_cell` on them. + if visitor.nz[0] == 2 and visitor.nz[1] == 2 and visitor.nz[2] == 2: selected = self.select_cell(spos, sdds) if ch != NULL: selected *= self.overlap_cells @@ -191,34 +203,42 @@ cdef class SelectorObject: visitor.ind[2] = k visitor.visit(root, selected) return - # Okay, now that we've got that out of the way, we have to do some - # other checks here. In this case, spos[] is the position of the - # center of a *possible* oct child, which means it is the center of a - # cluster of cells. That cluster might have 1, 8, 64, ... cells in it. - # But, we can figure it out by calculating the cell dds. + # Otherwise, we have to do some work to figure out where the cell centers are. + # We assign integer index ranges to each octant using half-open bounds. cdef np.float64_t dds[3] cdef np.float64_t pos[3] + cdef np.float64_t full_left[3] cdef int ci, cj, ck - cdef int nr = (visitor.nz >> 1) + cdef int start[3] + cdef int end[3] + cdef int split + cdef int oct_ind[3] + oct_ind[0] = i + oct_ind[1] = j + oct_ind[2] = k for ci in range(3): - dds[ci] = sdds[ci] / nr - # Boot strap at the first index. - pos[0] = (spos[0] - sdds[0]/2.0) + dds[0] * 0.5 - for ci in range(nr): - pos[1] = (spos[1] - sdds[1]/2.0) + dds[1] * 0.5 - for cj in range(nr): - pos[2] = (spos[2] - sdds[2]/2.0) + dds[2] * 0.5 - for ck in range(nr): + dds[ci] = (2.0 * sdds[ci]) / visitor.nz[ci] + full_left[ci] = (spos[ci] - sdds[ci] / 2.0) - oct_ind[ci] * sdds[ci] + split = visitor.nz[ci] // 2 + if oct_ind[ci] == 0: + start[ci] = 0 + end[ci] = split + else: + start[ci] = split + end[ci] = visitor.nz[ci] + for ci in range(start[0], end[0]): + pos[0] = full_left[0] + (ci + 0.5) * dds[0] + for cj in range(start[1], end[1]): + pos[1] = full_left[1] + (cj + 0.5) * dds[1] + for ck in range(start[2], end[2]): + pos[2] = full_left[2] + (ck + 0.5) * dds[2] selected = self.select_cell(pos, dds) if ch != NULL: selected *= self.overlap_cells - visitor.ind[0] = ci + i * nr - visitor.ind[1] = cj + j * nr - visitor.ind[2] = ck + k * nr + visitor.ind[0] = ci + visitor.ind[1] = cj + visitor.ind[2] = ck visitor.visit(root, selected) - pos[2] += dds[2] - pos[1] += dds[1] - pos[0] += dds[0] @cython.boundscheck(False) @cython.wraparound(False) diff --git a/yt/geometry/oct_container.pxd b/yt/geometry/oct_container.pxd index 40e189f81b1..166981f98b7 100644 --- a/yt/geometry/oct_container.pxd +++ b/yt/geometry/oct_container.pxd @@ -57,7 +57,7 @@ cdef class OctreeContainer: cdef int partial_coverage cdef int level_offset cdef int nn[3] - cdef np.uint8_t nz + cdef np.uint8_t nz[3] cdef np.float64_t DLE[3] cdef np.float64_t DRE[3] cdef public np.int64_t nocts @@ -86,7 +86,7 @@ cdef class OctreeContainer: self, const int level, const np.uint8_t[::1] level_inds, - const np.uint8_t[::1] cell_inds, + const np.uint32_t[::1] cell_inds, const np.int64_t[::1] file_inds, dict dest_fields, dict source_fields, @@ -96,7 +96,7 @@ cdef class OctreeContainer: self, const int level, const np.uint8_t[::1] level_inds, - const np.uint8_t[::1] cell_inds, + const np.uint32_t[::1] cell_inds, const np.int64_t[::1] file_inds, const np.int32_t[::1] domain_inds, dict dest_fields, diff --git a/yt/geometry/oct_container.pyx b/yt/geometry/oct_container.pyx index 806099b056a..d28bd3082f7 100644 --- a/yt/geometry/oct_container.pyx +++ b/yt/geometry/oct_container.pyx @@ -49,9 +49,14 @@ cdef class OctreeContainer: domain_right_edge, partial_coverage = 0, num_zones = 2): # This will just initialize the root mesh octs - self.nz = num_zones - self.partial_coverage = partial_coverage cdef int i + if hasattr(num_zones, '__len__'): + for i in range(3): + self.nz[i] = num_zones[i] + else: + for i in range(3): + self.nz[i] = num_zones + self.partial_coverage = partial_coverage for i in range(3): self.nn[i] = oct_domain_dimensions[i] self.num_domains = 0 @@ -91,7 +96,8 @@ cdef class OctreeContainer: cdef int i, j, k visitor.global_index = -1 visitor.level = 0 - visitor.nz = visitor.nzones = 1 + visitor.nz[0] = visitor.nz[1] = visitor.nz[2] = 1 + visitor.nzones = 1 visitor.max_level = 0 assert(ref_mask.shape[0] / float(visitor.nzones) == (ref_mask.shape[0]/float(visitor.nzones))) @@ -133,8 +139,9 @@ cdef class OctreeContainer: pos[1] += dds[1] pos[0] += dds[0] obj.nocts = cur.n_assigned - if obj.nocts * visitor.nz != ref_mask.size: - raise KeyError(ref_mask.size, obj.nocts, obj.nz, + if obj.nocts * visitor.nzones != ref_mask.size: + raise KeyError(ref_mask.size, obj.nocts, + (obj.nz[0], obj.nz[1], obj.nz[2]), obj.partial_coverage, visitor.nzones) obj.max_level = visitor.max_level return obj @@ -248,15 +255,14 @@ cdef class OctreeContainer: else: next = NULL if oinfo == NULL: return cur - cdef np.float64_t factor = 1.0 / self.nz * 2 for i in range(3): # We don't normally need to change dds[i] as it has been halved # from the oct width, thus making it already the cell width. # But, since not everything has the cell width equal to have the # width of the oct, we need to apply "factor". - oinfo.dds[i] = dds[i] * factor # Cell width + oinfo.dds[i] = dds[i] * 2.0 / self.nz[i] # Cell width oinfo.ipos[i] = ipos[i] - oinfo.left_edge[i] = oinfo.ipos[i] * (oinfo.dds[i] * self.nz) + self.DLE[i] + oinfo.left_edge[i] = oinfo.ipos[i] * (oinfo.dds[i] * self.nz[i]) + self.DLE[i] oinfo.level = level return cur @@ -266,33 +272,36 @@ cdef class OctreeContainer: list of oct IDs and a dictionary of Oct info for all the positions supplied. Positions must be in code_length. """ - cdef np.float64_t factor = self.nz + cdef np.float64_t nz_factor[3] cdef dict all_octs = {} cdef OctInfo oi cdef Oct* o = NULL cdef np.float64_t pos[3] cdef np.ndarray[np.uint8_t, ndim=1] recorded cdef np.ndarray[np.int64_t, ndim=1] oct_id + cdef int i + for i in range(3): + nz_factor[i] = self.nz[i] oct_id = np.ones(positions.shape[0], dtype="int64") * -1 recorded = np.zeros(self.nocts, dtype="uint8") - cdef np.int64_t i, j - for i in range(positions.shape[0]): - for j in range(3): - pos[j] = positions[i,j] + cdef np.int64_t pi, pj + for pi in range(positions.shape[0]): + for pj in range(3): + pos[pj] = positions[pi,pj] o = self.get(pos, &oi) if o == NULL: raise RuntimeError if recorded[o.domain_ind] == 0: left_edge = np.asarray(oi.left_edge).copy() dds = np.asarray(oi.dds).copy() - right_edge = left_edge + dds*factor + right_edge = left_edge + dds * np.asarray(nz_factor) all_octs[o.domain_ind] = dict( left_edge = left_edge, right_edge = right_edge, level = oi.level ) recorded[o.domain_ind] = 1 - oct_id[i] = o.domain_ind + oct_id[pi] = o.domain_ind return oct_id, all_octs def domain_identify(self, SelectorObject selector): @@ -334,7 +343,7 @@ cdef class OctreeContainer: for i in range(3): ndim[i] = ((self.DRE[i] - self.DLE[i]) / oi.dds[i]) # Here we adjust for oi.dds meaning *cell* width. - ndim[i] = (ndim[i] / self.nz) + ndim[i] = (ndim[i] / self.nz[i]) my_list = olist = OctList_append(NULL, o) for i in range(3): npos[0] = (oi.ipos[0] + (1 - i)) @@ -400,8 +409,7 @@ cdef class OctreeContainer: cdef np.ndarray[np.uint8_t, ndim=4] mask cdef oct_visitors.MaskOcts visitor visitor = oct_visitors.MaskOcts(self, domain_id) - cdef int ns = self.nz - mask = np.zeros((num_cells, ns, ns, ns), dtype="uint8") + mask = np.zeros((num_cells, self.nz[0], self.nz[1], self.nz[2]), dtype="uint8") visitor.mask = mask self.visit_all_octs(selector, visitor) return mask.astype("bool") @@ -487,13 +495,13 @@ cdef class OctreeContainer: header = dict(dims = (self.nn[0], self.nn[1], self.nn[2]), left_edge = (self.DLE[0], self.DLE[1], self.DLE[2]), right_edge = (self.DRE[0], self.DRE[1], self.DRE[2]), - num_zones = self.nz, + num_zones = (self.nz[0], self.nz[1], self.nz[2]), partial_coverage = self.partial_coverage) cdef SelectorObject selector = AlwaysSelector(None) # domain_id = -1 here, because we want *every* oct cdef oct_visitors.StoreOctree visitor visitor = oct_visitors.StoreOctree(self, -1) - visitor.nz = 1 + visitor.nz[0] = visitor.nz[1] = visitor.nz[2] = 1 cdef np.ndarray[np.uint8_t, ndim=1] ref_mask ref_mask = np.zeros(self.nocts * visitor.nzones, dtype="uint8") - 1 visitor.ref_mask = ref_mask @@ -683,14 +691,14 @@ cdef class OctreeContainer: num_cells = -1): # We create oct arrays of the correct size cdef np.ndarray[np.uint8_t, ndim=1] levels - cdef np.ndarray[np.uint8_t, ndim=1] cell_inds + cdef np.ndarray[np.uint32_t, ndim=1] cell_inds cdef np.ndarray[np.int64_t, ndim=1] file_inds if num_cells < 0: num_cells = selector.count_oct_cells(self, domain_id) # Initialize variables with dummy values levels = np.full(num_cells, 255, dtype="uint8") file_inds = np.full(num_cells, -1, dtype="int64") - cell_inds = np.full(num_cells, 8, dtype="uint8") + cell_inds = np.full(num_cells, self.nz[0] * self.nz[1] * self.nz[2], dtype="uint32") cdef oct_visitors.FillFileIndicesO visitor_o cdef oct_visitors.FillFileIndicesR visitor_r if self.fill_style == "r": @@ -746,7 +754,7 @@ cdef class OctreeContainer: self, const int level, const np.uint8_t[::1] levels, - const np.uint8_t[::1] cell_inds, + const np.uint32_t[::1] cell_inds, const np.int64_t[::1] file_inds, dict dest_fields, dict source_fields, @@ -799,7 +807,7 @@ cdef class OctreeContainer: ------- oct_inds : int64 ndarray (nocts*8, ) The on-domain index of the octs containing each cell - cell_inds : uint8 ndarray (nocts*8, ) + cell_inds : uint32 ndarray (nocts*8, ) The index of the cell in its parent oct Note @@ -811,10 +819,10 @@ cdef class OctreeContainer: cdef NeighbourCellIndexVisitor visitor - cdef np.uint8_t[::1] cell_inds + cdef np.uint32_t[::1] cell_inds cdef np.int64_t[::1] oct_inds - cell_inds = np.full(num_octs*4**3, 8, dtype=np.uint8) + cell_inds = np.full(num_octs*4**3, self.nz[0] * self.nz[1] * self.nz[2], dtype=np.uint32) oct_inds = np.full(num_octs*4**3, -1, dtype=np.int64) visitor = NeighbourCellIndexVisitor(self, -1, n_ghost_zones) @@ -832,7 +840,7 @@ cdef class OctreeContainer: self, const int level, const np.uint8_t[::1] level_inds, - const np.uint8_t[::1] cell_inds, + const np.uint32_t[::1] cell_inds, const np.int64_t[::1] file_inds, const np.int32_t[::1] domain_inds, dict dest_fields, @@ -888,7 +896,7 @@ cdef class OctreeContainer: ------- levels : uint8, shape (num_cells,) The level of each cell of the super oct - cell_inds : uint8, shape (num_cells, ) + cell_inds : uint32, shape (num_cells, ) The index of each cell of the super oct within its own oct file_inds : int64, shape (num_cells, ) The on-file position of the cell. See notes below. @@ -925,12 +933,12 @@ cdef class OctreeContainer: cdef NeighbourCellVisitor visitor cdef np.ndarray[np.uint8_t, ndim=1] levels - cdef np.ndarray[np.uint8_t, ndim=1] cell_inds + cdef np.ndarray[np.uint32_t, ndim=1] cell_inds cdef np.ndarray[np.int64_t, ndim=1] file_inds cdef np.ndarray[np.int32_t, ndim=1] domains levels = np.full(num_cells, 255, dtype="uint8") file_inds = np.full(num_cells, -1, dtype="int64") - cell_inds = np.full(num_cells, 8, dtype="uint8") + cell_inds = np.full(num_cells, self.nz[0] * self.nz[1] * self.nz[2], dtype="uint32") domains = np.full(num_cells, -1, dtype="int32") visitor = NeighbourCellVisitor(self, -1, n_ghost_zones) @@ -973,7 +981,12 @@ cdef class SparseOctreeContainer(OctreeContainer): num_zones = 2): cdef int i self.partial_coverage = 1 - self.nz = num_zones + if hasattr(num_zones, '__len__'): + for i in range(3): + self.nz[i] = num_zones[i] + else: + for i in range(3): + self.nz[i] = num_zones for i in range(3): self.nn[i] = domain_dimensions[i] self.domains = OctObjectPool() diff --git a/yt/geometry/oct_visitors.pxd b/yt/geometry/oct_visitors.pxd index 76d7b76e0c9..a46dafeeb2d 100644 --- a/yt/geometry/oct_visitors.pxd +++ b/yt/geometry/oct_visitors.pxd @@ -39,8 +39,8 @@ cdef class OctVisitor: cdef int dims cdef np.int32_t domain cdef np.int8_t level - cdef np.int8_t nz # This is number of zones along each dimension. 1 => 1 zones, 2 => 8, etc. - # To calculate nzones, nz**3 + cdef np.int8_t nz[3] # This is number of zones along each dimension. 1 => 1 zones, 2 => 8, etc. + # To calculate nzones, nz[0]*nz[1]*nz[2] cdef np.int32_t nzones # There will also be overrides for the memoryviews associated with the @@ -49,12 +49,10 @@ cdef class OctVisitor: cdef void visit(self, Oct*, np.uint8_t selected) cdef inline int oind(self): - cdef int d = self.nz - return (((self.ind[0]*d)+self.ind[1])*d+self.ind[2]) + return (self.ind[0]*self.nz[1]+self.ind[1])*self.nz[2]+self.ind[2] cdef inline int rind(self): - cdef int d = self.nz - return (((self.ind[2]*d)+self.ind[1])*d+self.ind[0]) + return (self.ind[2]*self.nz[1]+self.ind[1])*self.nz[0]+self.ind[0] cdef class CountTotalOcts(OctVisitor): pass @@ -116,12 +114,12 @@ cdef class AssignDomainInd(OctVisitor): cdef class FillFileIndicesO(OctVisitor): cdef np.uint8_t[:] levels cdef np.int64_t[:] file_inds - cdef np.uint8_t[:] cell_inds + cdef np.uint32_t[:] cell_inds cdef class FillFileIndicesR(OctVisitor): cdef np.uint8_t[:] levels cdef np.int64_t[:] file_inds - cdef np.uint8_t[:] cell_inds + cdef np.uint32_t[:] cell_inds cdef class CountByDomain(OctVisitor): cdef np.int64_t[:] domain_counts @@ -154,7 +152,7 @@ cdef class StoreIndex(OctVisitor): cdef class BaseNeighbourVisitor(OctVisitor): cdef int idim # 0,1,2 for x,y,z cdef int direction # +1 for +x, -1 for -x - cdef np.uint8_t neigh_ind[3] + cdef np.uint32_t neigh_ind[3] cdef bint other_oct cdef Oct *neighbour cdef OctreeContainer octree @@ -163,16 +161,15 @@ cdef class BaseNeighbourVisitor(OctVisitor): cdef void set_neighbour_info(self, Oct *o, int ishift[3]) - cdef inline np.uint8_t neighbour_rind(self): - cdef int d = self.nz - return (((self.neigh_ind[2]*d)+self.neigh_ind[1])*d+self.neigh_ind[0]) + cdef inline np.uint32_t neighbour_rind(self): + return (self.neigh_ind[2]*self.nz[1]+self.neigh_ind[1])*self.nz[0]+self.neigh_ind[0] cdef class NeighbourCellIndexVisitor(BaseNeighbourVisitor): - cdef np.uint8_t[::1] cell_inds + cdef np.uint32_t[::1] cell_inds cdef np.int64_t[::1] domain_inds cdef class NeighbourCellVisitor(BaseNeighbourVisitor): cdef np.uint8_t[::1] levels cdef np.int64_t[::1] file_inds - cdef np.uint8_t[::1] cell_inds + cdef np.uint32_t[::1] cell_inds cdef np.int32_t[::1] domains diff --git a/yt/geometry/oct_visitors.pyx b/yt/geometry/oct_visitors.pyx index b9ee8eb1da5..e64cb6bff2f 100644 --- a/yt/geometry/oct_visitors.pyx +++ b/yt/geometry/oct_visitors.pyx @@ -35,8 +35,10 @@ cdef class OctVisitor: self.dims = 0 self.domain = domain_id self.level = -1 - self.nz = octree.nz - self.nzones = self.nz**3 + self.nz[0] = octree.nz[0] + self.nz[1] = octree.nz[1] + self.nz[2] = octree.nz[2] + self.nzones = self.nz[0] * self.nz[1] * self.nz[2] cdef void visit(self, Oct* o, np.uint8_t selected): raise NotImplementedError @@ -173,7 +175,7 @@ cdef class ICoordsOcts(OctVisitor): if selected == 0: return cdef int i for i in range(3): - self.icoords[self.index,i] = (self.pos[i] * self.nz) + self.ind[i] + self.icoords[self.index,i] = (self.pos[i] * self.nz[i]) + self.ind[i] self.index += 1 # Level @@ -197,9 +199,9 @@ cdef class FCoordsOcts(OctVisitor): if selected == 0: return cdef int i cdef np.float64_t c, dx - dx = 1.0 / ((self.nz) << self.level) for i in range(3): - c = ((self.pos[i] * self.nz) + self.ind[i]) + dx = 1.0 / ((self.nz[i]) * (1 << self.level)) + c = ((self.pos[i] * self.nz[i]) + self.ind[i]) self.fcoords[self.index,i] = (c + 0.5) * dx self.index += 1 @@ -215,8 +217,8 @@ cdef class FWidthOcts(OctVisitor): if selected == 0: return cdef int i cdef np.float64_t dx - dx = 1.0 / (self.nz << self.level) for i in range(3): + dx = 1.0 / ((self.nz[i]) * (1 << self.level)) self.fwidth[self.index,i] = dx self.index += 1 @@ -332,7 +334,7 @@ cdef class MortonIndexOcts(OctVisitor): cdef np.int64_t coord[3] cdef int i for i in range(3): - coord[i] = (self.pos[i] * self.nz) + self.ind[i] + coord[i] = (self.pos[i] * self.nz[i]) + self.ind[i] if (coord[i] < 0): raise RuntimeError("Oct coordinate in dimension {} is ".format(i)+ "negative. ({})".format(coord[i])) @@ -374,12 +376,12 @@ cdef class BaseNeighbourVisitor(OctVisitor): cdef Oct *neighbour cdef bint local_oct cdef bint other_oct - dx = 1.0 / (self.nz << self.level) local_oct = True # Compute position of neighbouring cell for i in range(3): - c = (self.pos[i] * self.nz) + dx = 1.0 / ((self.nz[i]) * (1 << self.level)) + c = (self.pos[i] * self.nz[i]) fcoords[i] = (c + 0.5 + ishift[i]) * dx / self.octree.nn[i] # Assuming periodicity if fcoords[i] < 0: @@ -396,12 +398,12 @@ cdef class BaseNeighbourVisitor(OctVisitor): neighbour = o self.oi.level = self.level for i in range(3): - self.oi.ipos[i] = (self.pos[i] * self.nz) + ishift[i] + self.oi.ipos[i] = (self.pos[i] * self.nz[i]) + ishift[i] # Extra step - compute cell position in neighbouring oct (and store in oi.ipos) if self.oi.level == self.level - 1: for i in range(3): - ipos = (((self.pos[i] * self.nz) + ishift[i])) >> 1 + ipos = (((self.pos[i] * self.nz[i]) + ishift[i])) >> 1 if (self.oi.ipos[i] << 1) == ipos: self.oi.ipos[i] = 0 else: @@ -410,7 +412,7 @@ cdef class BaseNeighbourVisitor(OctVisitor): # Index of neighbouring cell within its oct for i in range(3): - self.neigh_ind[i] = (ishift[i]) % 2 + self.neigh_ind[i] = (ishift[i]) % 2 self.other_oct = other_oct if other_oct: @@ -450,7 +452,7 @@ cdef class NeighbourCellIndexVisitor(BaseNeighbourVisitor): cdef void visit(self, Oct* o, np.uint8_t selected): cdef int i, j, k cdef int ishift[3] - cdef np.uint8_t neigh_cell_ind + cdef np.uint32_t neigh_cell_ind cdef np.int64_t neigh_domain_ind if selected == 0: return # Work at oct level @@ -495,7 +497,7 @@ cdef class NeighbourCellVisitor(BaseNeighbourVisitor): cdef int i, j, k cdef int ishift[3] cdef np.int64_t neigh_file_ind - cdef np.uint8_t neigh_cell_ind + cdef np.uint32_t neigh_cell_ind cdef np.int32_t neigh_domain cdef np.uint8_t neigh_level if selected == 0: return diff --git a/yt/geometry/particle_deposit.pyx b/yt/geometry/particle_deposit.pyx index b04a4ad72dd..d893dbf7d88 100644 --- a/yt/geometry/particle_deposit.pyx +++ b/yt/geometry/particle_deposit.pyx @@ -61,7 +61,9 @@ cdef class ParticleDepositOperation: cdef np.float64_t pos[3] cdef np.float64_t[:] field_vals = np.empty(nf, dtype="float64") cdef int dims[3] - dims[0] = dims[1] = dims[2] = octree.nz + dims[0] = octree.nz[0] + dims[1] = octree.nz[1] + dims[2] = octree.nz[2] cdef OctInfo oi cdef np.int64_t offset, moff cdef Oct *oct diff --git a/yt/geometry/particle_oct_container.pyx b/yt/geometry/particle_oct_container.pyx index d9d64843836..70fc953cc33 100644 --- a/yt/geometry/particle_oct_container.pyx +++ b/yt/geometry/particle_oct_container.pyx @@ -1988,7 +1988,7 @@ cdef class ParticleBitmapOctreeContainer(SparseOctreeContainer): cdef oct_visitors.AssignDomainInd visitor visitor = oct_visitors.AssignDomainInd(self) self.visit_all_octs(selector, visitor) - assert ((visitor.global_index+1)*visitor.nz == visitor.index) + assert ((visitor.global_index+1)*visitor.nzones == visitor.index) # Copy indexes self._ptr_index_base_octs = malloc(sizeof(np.uint8_t)*self.nocts) self._index_base_octs = self._ptr_index_base_octs diff --git a/yt/geometry/particle_smooth.pyx b/yt/geometry/particle_smooth.pyx index ea3a7e79cbd..eaffbdf206f 100644 --- a/yt/geometry/particle_smooth.pyx +++ b/yt/geometry/particle_smooth.pyx @@ -121,7 +121,9 @@ cdef class ParticleSmoothOperation: periodicity = (False, False, False) else: raise NotImplementedError - dims[0] = dims[1] = dims[2] = mesh_octree.nz + dims[0] = mesh_octree.nz[0] + dims[1] = mesh_octree.nz[1] + dims[2] = mesh_octree.nz[2] cdef int nz = dims[0] * dims[1] * dims[2] # pcount is the number of particles per oct. pcount = np.zeros_like(pdom_ind) @@ -151,7 +153,6 @@ cdef class ParticleSmoothOperation: for i in range(3): self.DW[i] = (mesh_octree.DRE[i] - mesh_octree.DLE[i]) self.periodicity[i] = periodicity[i] - cdef np.float64_t factor = particle_octree.nz for i in range(positions.shape[0]): for j in range(3): pos[j] = positions[i, j] @@ -168,7 +169,7 @@ cdef class ParticleSmoothOperation: # in octs that we know are too far away for j in range(3): oct_left_edges[offset, j] = oinfo.left_edge[j] - oct_dds[offset, j] = oinfo.dds[j] * factor + oct_dds[offset, j] = oinfo.dds[j] * particle_octree.nz[j] # Now we have oct assignments. Let's sort them. # Note that what we will be providing to our processing functions will # actually be indirectly-sorted fields. This preserves memory at the