diff --git a/tests/tests.yaml b/tests/tests.yaml index 7cddf71720a..5249c5794ff 100644 --- a/tests/tests.yaml +++ b/tests/tests.yaml @@ -142,7 +142,7 @@ answer_tests: - yt/frontends/boxlib/tests/test_outputs.py:test_NyxDataset - yt/frontends/boxlib/tests/test_outputs.py:test_WarpXDataset - local_ramses_002: + local_ramses_003: - yt/frontends/ramses/tests/test_outputs.py:test_output_00080 local_ytdata_007: diff --git a/yt/data_objects/data_containers.py b/yt/data_objects/data_containers.py index 62e30e0dc5f..e9a0f5cf6ac 100644 --- a/yt/data_objects/data_containers.py +++ b/yt/data_objects/data_containers.py @@ -353,7 +353,9 @@ def _generate_spatial_fluid(self, field, ngz): outputs.append(rv) ind = 0 # Does this work with mesh? with o._activate_cache(): - ind += o.select(self.selector, self[field], rv, ind) + ind += o.select( + self.selector, source=self[field], dest=rv, offset=ind + ) else: chunks = self.index._chunk(self, "spatial", ngz=ngz) for i, chunk in enumerate(chunks): @@ -366,15 +368,12 @@ def _generate_spatial_fluid(self, field, ngz): np.empty(wogz.ires.size, dtype="float64"), units ) outputs.append(rv) - if gz._type_name == "octree_subset": - raise NotImplementedError - else: - ind += wogz.select( - self.selector, - gz[field][ngz:-ngz, ngz:-ngz, ngz:-ngz], - rv, - ind, - ) + ind += wogz.select( + self.selector, + source=gz[field][ngz:-ngz, ngz:-ngz, ngz:-ngz], + dest=rv, + offset=ind, + ) if accumulate: rv = uconcatenate(outputs) return rv diff --git a/yt/data_objects/octree_subset.py b/yt/data_objects/octree_subset.py index 0fa27fd9524..a72772753e9 100644 --- a/yt/data_objects/octree_subset.py +++ b/yt/data_objects/octree_subset.py @@ -37,11 +37,14 @@ class OctreeSubset(YTSelectionContainer): _con_args = ("base_region", "domain", "ds") _domain_offset = 0 _cell_count = -1 - _block_reorder = None + _block_order = "C" - def __init__(self, base_region, domain, ds, over_refine_factor=1): + def __init__( + self, base_region, domain, ds, over_refine_factor=1, num_ghost_zones=0 + ): super(OctreeSubset, self).__init__(ds, None) self._num_zones = 1 << (over_refine_factor) + self._num_ghost_zones = num_ghost_zones self._oref = over_refine_factor self.domain = domain self.domain_id = domain.domain_id @@ -528,9 +531,10 @@ def __init__(self, ind, block_slice): def __getitem__(self, key): bs = self.block_slice - rv = bs.octree_subset[key][:, :, :, self.ind].T - if bs.octree_subset._block_reorder: - rv = rv.copy(order=bs.octree_subset._block_reorder) + rv = np.require( + bs.octree_subset[key][:, :, :, self.ind].T, + requirements=bs.octree_subset._block_order, + ) return rv @property diff --git a/yt/fields/fluid_fields.py b/yt/fields/fluid_fields.py index 03287a1b67a..ef8588f706a 100644 --- a/yt/fields/fluid_fields.py +++ b/yt/fields/fluid_fields.py @@ -221,8 +221,8 @@ def setup_gradient_fields(registry, grad_field, field_units, slice_info=None): geom = registry.ds.geometry if is_curvilinear(geom): mylog.warning( - "In %s geometry, gradient fields may contain artifacts near cartesian axes." - % geom + "In %s geometry, gradient fields may contain artifacts near cartesian axes.", + geom, ) assert isinstance(grad_field, tuple) @@ -240,16 +240,27 @@ def grad_func(axi, ax): slice_3dr = slice_3d[:axi] + (sl_right,) + slice_3d[axi + 1 :] def func(field, data): - ds = div_fac * data[ftype, "d%s" % ax] + block_order = getattr(data, "_block_order", "C") + if block_order == "F": + # Fortran-ordering: we need to swap axes here and + # reswap below + field_data = data[grad_field].swapaxes(0, 2) + else: + field_data = data[grad_field] + dx = div_fac * data[ftype, f"d{ax}"] if ax == "theta": - ds *= data[ftype, "r"] + dx *= data[ftype, "r"] if ax == "phi": - ds *= data[ftype, "r"] * np.sin(data[ftype, "theta"]) - f = data[grad_field][slice_3dr] / ds[slice_3d] - f -= data[grad_field][slice_3dl] / ds[slice_3d] + dx *= data[ftype, "r"] * np.sin(data[ftype, "theta"]) + f = field_data[slice_3dr] / dx[slice_3d] + f -= field_data[slice_3dl] / dx[slice_3d] new_field = np.zeros_like(data[grad_field], dtype=np.float64) - new_field = data.ds.arr(new_field, f.units) + new_field = data.ds.arr(new_field, field_data.units / dx.units) new_field[slice_3d] = f + + if block_order == "F": + new_field = new_field.swapaxes(0, 2) + return new_field return func diff --git a/yt/fields/geometric_fields.py b/yt/fields/geometric_fields.py index cf7f0a6b972..0d355f0f3f4 100644 --- a/yt/fields/geometric_fields.py +++ b/yt/fields/geometric_fields.py @@ -101,10 +101,11 @@ def _zeros(field, data): def _ones(field, data): """Returns one for all cells""" - arr = np.ones(data.ires.shape, dtype="float64") + tmp = np.ones(data.ires.shape, dtype="float64") + arr = data.apply_units(tmp, field.units) if data._spatial: return data._reshape_vals(arr) - return data.apply_units(arr, field.units) + return arr registry.add_field( ("index", "ones"), diff --git a/yt/frontends/ramses/data_structures.py b/yt/frontends/ramses/data_structures.py index 93ebaa4fa0b..a862424f114 100644 --- a/yt/frontends/ramses/data_structures.py +++ b/yt/frontends/ramses/data_structures.py @@ -185,16 +185,47 @@ def included(self, selector): class RAMSESDomainSubset(OctreeSubset): _domain_offset = 1 - _block_reorder = "F" + _block_order = "F" - def fill(self, fd, fields, selector, file_handler): + _base_domain = None + + def __init__( + self, + base_region, + domain, + ds, + over_refine_factor=1, + num_ghost_zones=0, + base_grid=None, + ): + super(RAMSESDomainSubset, self).__init__( + base_region, domain, ds, over_refine_factor, num_ghost_zones + ) + + self._base_grid = base_grid + + if num_ghost_zones > 0: + if not all(ds.periodicity): + mylog.warn("Ghost zones will wrongly assume the domain to be periodic.") + # Create a base domain *with no self._base_domain.fwidth + base_domain = RAMSESDomainSubset( + ds.all_data(), domain, ds, over_refine_factor + ) + self._base_domain = base_domain + elif num_ghost_zones < 0: + raise RuntimeError( + "Cannot initialize a domain subset with a negative number of ghost zones," + " was called with num_ghost_zones=%s" % num_ghost_zones + ) + + def _fill_no_ghostzones(self, fd, fields, selector, file_handler): ndim = self.ds.dimensionality # Here we get a copy of the file, which we skip through and read the # bits we want. oct_handler = self.oct_handler all_fields = [f for ft, f in file_handler.field_list] fields = [f for ft, f in fields] - tr = {} + data = {} cell_count = selector.count_oct_cells(self.oct_handler, self.domain_id) levels, cell_inds, file_inds = self.oct_handler.file_index_octs( @@ -203,12 +234,59 @@ def fill(self, fd, fields, selector, file_handler): # Initializing data container for field in fields: - tr[field] = np.zeros(cell_count, "float64") + data[field] = np.zeros(cell_count, "float64") + + cpu_list = [self.domain_id - 1] + fill_hydro( + fd, + file_handler.offset, + file_handler.level_count, + cpu_list, + levels, + cell_inds, + file_inds, + ndim, + all_fields, + fields, + data, + oct_handler, + ) + return data + + def _fill_with_ghostzones( + self, fd, fields, selector, file_handler, num_ghost_zones + ): + ndim = self.ds.dimensionality + ncpu = self.ds.parameters["ncpu"] + # Here we get a copy of the file, which we skip through and read the + # bits we want. + oct_handler = self.oct_handler + all_fields = [f for ft, f in file_handler.field_list] + fields = [f for ft, f in fields] + tr = {} + + cell_count = ( + selector.count_octs(self.oct_handler, self.domain_id) * self.nz ** ndim + ) + + ( + levels, + cell_inds, + file_inds, + domains, + ) = self.oct_handler.file_index_octs_with_ghost_zones( + selector, self.domain_id, cell_count + ) + # Initializing data container + for field in fields: + tr[field] = np.zeros(cell_count, "float64") + cpu_list = list(range(ncpu)) fill_hydro( fd, file_handler.offset, file_handler.level_count, + cpu_list, levels, cell_inds, file_inds, @@ -217,9 +295,59 @@ def fill(self, fd, fields, selector, file_handler): fields, tr, oct_handler, + domains=domains, ) return tr + @property + def fwidth(self): + fwidth = super(RAMSESDomainSubset, self).fwidth + if self._num_ghost_zones > 0: + fwidth = fwidth.reshape(-1, 8, 3) + n_oct = fwidth.shape[0] + # 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[:, :, :] = fwidth[:, 0:1, :] + fwidth = new_fwidth.reshape(-1, 3) + return fwidth + + @property + def fcoords(self): + num_ghost_zones = self._num_ghost_zones + if num_ghost_zones == 0: + return super(RAMSESDomainSubset, self).fcoords + + oh = self.oct_handler + + indices = oh.fill_index(self.selector).reshape(-1, 8) + oct_inds, cell_inds = oh.fill_octcellindex_neighbours(self.selector) + + oct_inds = oct_inds.reshape(-1, 64) + cell_inds = cell_inds.reshape(-1, 64) + + inds = indices[oct_inds, cell_inds] + + fcoords = self.ds.arr(oh.fcoords(self.selector)[inds].reshape(-1, 3), "unitary") + + return fcoords + + def fill(self, fd, fields, selector, file_handler): + if self._num_ghost_zones == 0: + return self._fill_no_ghostzones(fd, fields, selector, file_handler) + else: + return self._fill_with_ghostzones( + fd, fields, selector, file_handler, self._num_ghost_zones + ) + + def retrieve_ghost_zones(self, ngz, fields, smoothed=False): + new_subset = RAMSESDomainSubset( + self.base_region, self.domain, self.ds, num_ghost_zones=ngz, base_grid=self + ) + + return new_subset + class RAMSESIndex(OctreeIndex): def __init__(self, ds, dataset_type="ramses"): @@ -275,7 +403,12 @@ def _identify_base_chunk(self, dobj): if len(domains) > 1: mylog.debug("Identified %s intersecting domains", len(domains)) subsets = [ - RAMSESDomainSubset(base_region, domain, self.dataset) + RAMSESDomainSubset( + base_region, + domain, + self.dataset, + num_ghost_zones=dobj._num_ghost_zones, + ) for domain in domains ] dobj._chunk_info = subsets diff --git a/yt/frontends/ramses/io_utils.pyx b/yt/frontends/ramses/io_utils.pyx index 3e73e6bfed3..1b29199e176 100644 --- a/yt/frontends/ramses/io_utils.pyx +++ b/yt/frontends/ramses/io_utils.pyx @@ -85,7 +85,7 @@ def read_amr(FortranFile f, dict headers, @cython.nonecheck(False) cpdef read_offset(FortranFile f, INT64_t min_level, INT64_t domain_id, INT64_t nvar, dict headers): - cdef np.ndarray[np.int64_t, ndim=1] offset, level_count + cdef np.ndarray[np.int64_t, ndim=2] offset, level_count cdef INT64_t ndim, twotondim, nlevelmax, n_levels, nboundary, ncpu, ncpu_and_bound cdef INT64_t ilevel, icpu, skip_len cdef INT32_t file_ilevel, file_ncache @@ -103,12 +103,11 @@ cpdef read_offset(FortranFile f, INT64_t min_level, INT64_t domain_id, INT64_t n skip_len = twotondim * nvar # It goes: level, CPU, 8-variable (1 oct) - offset = np.zeros(n_levels, dtype=np.int64) - offset -= 1 - level_count = np.zeros(n_levels, dtype=np.int64) + offset = np.full((ncpu_and_bound, n_levels), -1, dtype=np.int64) + level_count = np.zeros((ncpu_and_bound, n_levels), dtype=np.int64) - cdef np.int64_t[:] level_count_view = level_count - cdef np.int64_t[:] offset_view = offset + cdef np.int64_t[:, ::1] level_count_view = level_count + cdef np.int64_t[:, ::1] offset_view = offset for ilevel in range(nlevelmax): for icpu in range(ncpu_and_bound): @@ -123,9 +122,9 @@ cpdef read_offset(FortranFile f, INT64_t min_level, INT64_t domain_id, INT64_t n 'from data (%s) is not coherent with the expected (%s)', f.name, file_ilevel, ilevel) - if icpu + 1 == domain_id and ilevel >= min_level: - offset[ilevel - min_level] = f.tell() - level_count[ilevel - min_level] = file_ncache + if ilevel >= min_level: + offset_view[icpu, ilevel - min_level] = f.tell() + level_count_view[icpu, ilevel - min_level] = file_ncache f.skip(skip_len) return offset, level_count @@ -135,45 +134,58 @@ cpdef read_offset(FortranFile f, INT64_t min_level, INT64_t domain_id, INT64_t n @cython.cdivision(True) @cython.nonecheck(False) def fill_hydro(FortranFile f, - np.ndarray[np.int64_t, ndim=1] offsets, - np.ndarray[np.int64_t, ndim=1] level_count, + np.ndarray[np.int64_t, ndim=2] offsets, + np.ndarray[np.int64_t, ndim=2] level_count, + list cpu_enumerator, np.ndarray[np.uint8_t, ndim=1] levels, np.ndarray[np.uint8_t, ndim=1] cell_inds, np.ndarray[np.int64_t, ndim=1] file_inds, INT64_t ndim, list all_fields, list fields, dict tr, - RAMSESOctreeContainer oct_handler): - cdef INT64_t ilevel, ifield, nfields, noffset + RAMSESOctreeContainer oct_handler, + np.ndarray[np.int32_t, ndim=1] domains=np.array([], dtype='int32')): cdef INT64_t offset cdef dict tmp cdef str field - cdef INT64_t twotondim, i + cdef INT64_t twotondim + cdef int ilevel, icpu, ifield, nfields, nlevels, nc, ncpu_selected cdef np.ndarray[np.uint8_t, ndim=1] mask twotondim = 2**ndim nfields = len(all_fields) - noffset = len(offsets) + ncpu = offsets.shape[0] + nlevels = offsets.shape[1] + ncpu_selected = len(cpu_enumerator) mask = np.array([(field in fields) for field in all_fields], dtype=np.uint8) # Loop over levels - for ilevel in range(noffset): - offset = offsets[ilevel] - if offset == -1: - continue - f.seek(offset) - nc = level_count[ilevel] - tmp = {} - # Initalize temporary data container for io - for field in all_fields: - tmp[field] = np.empty((nc, twotondim), dtype="float64") - - for i in range(twotondim): - # Read the selected fields - for ifield in range(nfields): - if not mask[ifield]: - f.skip() - else: - tmp[all_fields[ifield]][:, i] = f.read_vector('d') # i-th cell - - oct_handler.fill_level(ilevel, levels, cell_inds, file_inds, tr, tmp) + for ilevel in range(nlevels): + # Loop over cpu domains + for icpu in cpu_enumerator: + nc = level_count[icpu, ilevel] + if nc == 0: + continue + offset = offsets[icpu, ilevel] + if offset == -1: + continue + f.seek(offset) + tmp = {} + # Initalize temporary data container for io + # note: we use Fortran ordering to reflect the in-file ordering + for field in all_fields: + tmp[field] = np.empty((nc, twotondim), dtype="float64", order='F') + + for i in range(twotondim): + # Read the selected fields + for ifield in range(nfields): + if not mask[ifield]: + f.skip() + else: + tmp[all_fields[ifield]][:, i] = f.read_vector('d') # i-th cell + if ncpu_selected > 1: + oct_handler.fill_level_with_domain( + ilevel, levels, cell_inds, file_inds, domains, tr, tmp, domain=icpu+1) + else: + oct_handler.fill_level( + ilevel, levels, cell_inds, file_inds, tr, tmp) diff --git a/yt/frontends/ramses/tests/test_outputs.py b/yt/frontends/ramses/tests/test_outputs.py index 42b3b54de60..58752a15d83 100644 --- a/yt/frontends/ramses/tests/test_outputs.py +++ b/yt/frontends/ramses/tests/test_outputs.py @@ -21,7 +21,12 @@ ) from yt.utilities.on_demand_imports import _f90nml as f90nml -_fields = ("temperature", "density", "velocity_magnitude") +_fields = ( + "temperature", + "density", + "velocity_magnitude", + "pressure_gradient_magnitude", +) output_00080 = "output_00080/info_00080.txt" @@ -465,9 +470,6 @@ def check_unit(array, unit): check_unit(ds.r[("gas", "Electron_number_density")], "cm**(-3)") -ramses_rt = "ramses_rt_00088/output_00088/info_00088.txt" - - @requires_file(ramses_rt) def test_ramses_mixed_files(): # Test that one can use derived fields that depend on different @@ -540,3 +542,22 @@ def test_magnetic_field_aliasing(): ]: assert ("gas", field) in ds.derived_field_list ad[("gas", field)] + + +@requires_file(output_00080) +def test_field_accession(): + ds = yt.load(output_00080) + fields = [ + ("gas", "density"), # basic ones + ("gas", "pressure"), + ("gas", "pressure_gradient_magnitude"), # requires ghost zones + ] + # Check accessing gradient works for a variety of spatial domains + for reg in ( + ds.all_data(), + ds.sphere([0.1] * 3, 0.01), + ds.sphere([0.5] * 3, 0.05), + ds.box([0.1] * 3, [0.2] * 3), + ): + for field in fields: + reg[field] diff --git a/yt/geometry/fake_octree.pyx b/yt/geometry/fake_octree.pyx index 9c972bc1648..bba59f80dc6 100644 --- a/yt/geometry/fake_octree.pyx +++ b/yt/geometry/fake_octree.pyx @@ -15,11 +15,11 @@ from oct_visitors cimport cind import numpy as np cimport cython -from oct_container cimport Oct, RAMSESOctreeContainer +from oct_container cimport Oct, SparseOctreeContainer # Create a balanced octree by a random walk that recursively # subdivides -def create_fake_octree(RAMSESOctreeContainer oct_handler, +def create_fake_octree(SparseOctreeContainer oct_handler, long max_noct, long max_level, np.ndarray[np.int32_t, ndim=1] ndd, @@ -46,7 +46,7 @@ def create_fake_octree(RAMSESOctreeContainer oct_handler, return cur_leaf -cdef long subdivide(RAMSESOctreeContainer oct_handler, +cdef long subdivide(SparseOctreeContainer oct_handler, Oct *parent, int ind[3], int dd[3], long cur_leaf, long cur_level, diff --git a/yt/geometry/oct_container.pxd b/yt/geometry/oct_container.pxd index 4f7c482f075..75ada2b97e7 100644 --- a/yt/geometry/oct_container.pxd +++ b/yt/geometry/oct_container.pxd @@ -12,7 +12,7 @@ cimport numpy as np from yt.utilities.lib.fp_utils cimport * cimport oct_visitors cimport selection_routines -from .oct_visitors cimport OctVisitor, Oct, cind +from .oct_visitors cimport OctVisitor, Oct, cind, OctInfo from libc.stdlib cimport bsearch, qsort, realloc, malloc, free from libc.math cimport floor from yt.utilities.lib.allocation_container cimport \ @@ -27,12 +27,6 @@ cdef struct OctKey: np.int64_t *indices np.int64_t pcount -cdef struct OctInfo: - np.float64_t left_edge[3] - np.float64_t dds[3] - np.int64_t ipos[3] - np.int32_t level - cdef struct OctList cdef struct OctList: diff --git a/yt/geometry/oct_container.pyx b/yt/geometry/oct_container.pyx index 0f38a8a7a01..db2e13490bf 100644 --- a/yt/geometry/oct_container.pyx +++ b/yt/geometry/oct_container.pyx @@ -13,10 +13,10 @@ Oct container cimport cython cimport numpy as np import numpy as np -from selection_routines cimport SelectorObject +from selection_routines cimport SelectorObject, AlwaysSelector from libc.math cimport floor, ceil -cimport selection_routines -from yt.geometry.oct_visitors cimport OctPadded +from yt.geometry.oct_visitors cimport OctPadded, NeighbourCellVisitor, StoreIndex, NeighbourCellIndexVisitor + ORDER_MAX = 20 _ORDER_MAX = ORDER_MAX @@ -79,7 +79,7 @@ cdef class OctreeContainer: header['right_edge'], over_refine = header['over_refine'], partial_coverage = header['partial_coverage']) # NOTE: We do not allow domain/file indices to be specified. - cdef SelectorObject selector = selection_routines.AlwaysSelector(None) + cdef SelectorObject selector = AlwaysSelector(None) cdef oct_visitors.LoadOctree visitor visitor = oct_visitors.LoadOctree(obj, -1) cdef int i, j, k, n @@ -213,7 +213,7 @@ cdef class OctreeContainer: ind32[i] = ind[i] self.get_root(ind32, &next) # We want to stop recursing when there's nowhere else to go - while next != NULL and level <= max_level: + while next != NULL and level < max_level: level += 1 for i in range(3): ipos[i] = (ipos[i] << 1) + ind[i] @@ -474,7 +474,7 @@ cdef class OctreeContainer: right_edge = (self.DRE[0], self.DRE[1], self.DRE[2]), over_refine = self.oref, partial_coverage = self.partial_coverage) - cdef SelectorObject selector = selection_routines.AlwaysSelector(None) + 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) @@ -551,7 +551,7 @@ cdef class OctreeContainer: def domain_ind(self, selector, int domain_id = -1): cdef np.ndarray[np.int64_t, ndim=1] ind # Here's where we grab the masked items. - ind = np.zeros(self.nocts, 'int64') - 1 + ind = np.full(self.nocts, -1, 'int64') cdef oct_visitors.IndexOcts visitor visitor = oct_visitors.IndexOcts(self, domain_id) visitor.oct_index = ind @@ -664,13 +664,10 @@ cdef class OctreeContainer: cdef np.ndarray[np.int64_t, ndim=1] file_inds if num_cells < 0: num_cells = selector.count_oct_cells(self, domain_id) - levels = np.zeros(num_cells, dtype="uint8") - file_inds = np.zeros(num_cells, dtype="int64") - cell_inds = np.zeros(num_cells, dtype="uint8") - for i in range(num_cells): - levels[i] = 100 - file_inds[i] = -1 - cell_inds[i] = 9 + # 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") cdef oct_visitors.FillFileIndicesO visitor_o cdef oct_visitors.FillFileIndicesR visitor_r if self.fill_style == "r": @@ -731,17 +728,197 @@ cdef class OctreeContainer: cdef np.ndarray[np.float64_t, ndim=2] source cdef np.ndarray[np.float64_t, ndim=1] dest cdef int i - cdef np.int64_t local_filled = 0 + + for key in dest_fields: + dest = dest_fields[key] + source = source_fields[key] + for i, lvl in enumerate(levels): + if lvl != level: continue + if file_inds[i] < 0: + dest[i + offset] = np.nan + else: + dest[i + offset] = source[file_inds[i], cell_inds[i]] + + def fill_index(self, SelectorObject selector = AlwaysSelector(None)): + """Get the on-file index of each cell""" + cdef StoreIndex visitor + + cdef np.int64_t[:, :, :, :] cell_inds + + cell_inds = np.full((self.nocts, 2, 2, 2), -1, dtype=np.int64) + + visitor = StoreIndex(self, -1) + visitor.cell_inds = cell_inds + + self.visit_all_octs(selector, visitor) + + return np.asarray(cell_inds) + + def fill_octcellindex_neighbours(self, SelectorObject selector, int num_octs = -1, domain_id = -1): + """Compute the oct and cell indices of all the cells within all selected octs, extended + by one cell in all directions (for ghost zones computations). + + Parameters + ---------- + selector : SelectorObject + Selector for the octs to compute neighbour of + num_octs : int, optional + The number of octs to read in + domain_id : int, optional + The domain to perform the selection over + + Returns + ------- + oct_inds : int64 ndarray (nocts*8, ) + The on-domain index of the octs containing each cell + cell_inds : uint8 ndarray (nocts*8, ) + The index of the cell in its parent oct + + Note + ---- + oct_inds/cell_inds + """ + if num_octs == -1: + num_octs = selector.count_octs(self, domain_id) + + cdef NeighbourCellIndexVisitor visitor + + cdef np.uint8_t[::1] cell_inds + cdef np.int64_t[::1] oct_inds + + cell_inds = np.full(num_octs*4**3, 8, dtype=np.uint8) + oct_inds = np.full(num_octs*4**3, -1, dtype=np.int64) + + visitor = NeighbourCellIndexVisitor(self, -1) + visitor.cell_inds = cell_inds + visitor.domain_inds = oct_inds + + self.visit_all_octs(selector, visitor) + + return np.asarray(oct_inds), np.asarray(cell_inds) + + @cython.boundscheck(False) + @cython.wraparound(False) + @cython.cdivision(True) + def fill_level_with_domain( + self, int level, + np.uint8_t[:] levels, + np.uint8_t[:] cell_inds, + np.int64_t[:] file_inds, + np.int32_t[:] domains, + dict dest_fields, + dict source_fields, + np.int32_t domain, + np.int64_t offset = 0 + ): + """Similar to fill_level but accepts a domain argument. + + This is particularly useful for frontends that have buffer zones at CPU boundaries. + These buffer oct cells have a different domain than the local one and + are usually not read, but one has to read them e.g. to compute ghost zones. + """ + cdef np.ndarray[np.float64_t, ndim=2] source + cdef np.ndarray[np.float64_t, ndim=1] dest + cdef int i, count + for key in dest_fields: dest = dest_fields[key] source = source_fields[key] - for i in range(levels.shape[0]): - if levels[i] != level: continue - dest[i + offset] = source[file_inds[i], cell_inds[i]] - local_filled += 1 + count = 0 + for i, (lev, dom) in enumerate(zip(levels, domains)): + if lev != level or dom != domain: continue + count += 1 + if file_inds[i] < 0: + dest[i + offset] = np.nan + else: + dest[i + offset] = source[file_inds[i], cell_inds[i]] + return count + + @cython.boundscheck(False) + @cython.wraparound(False) + @cython.cdivision(True) + def file_index_octs_with_ghost_zones( + self, SelectorObject selector, int domain_id, + int num_cells = -1): + """Similar as file_index_octs, but returns the level, cell index, + file index and domain of the neighbouring cells as well. + + Arguments + --------- + selector : SelectorObject + The selector object. It is expected to select all cells for a given oct. + domain_id : int + The domain to select. Set to -1 to select all domains. + num_cells : int, optional + The total number of cells (accounting for a 1-cell thick ghost zone layer). + + Returns + ------- + levels : uint8, shape (num_cells,) + The level of each cell of the super oct + cell_inds : uint8, 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. + domains : int32, shape (num_cells) + The domain to which the cells belongs. See notes below. + + Notes + ----- + + The algorithm constructs a "super-oct" around each oct (see sketch below, + where the original oct cells are marked with an x). + + Note that for sparse octrees (such as RAMSES'), the neighbouring cells + may belong to another domain (this is stored in `domains`). If the dataset + provides buffer zones between domains (such as RAMSES), this may be stored + locally and can be accessed directly. + + + +---+---+---+---+ + | | | | | + |---+---+---+---| + | | x | x | | + |---+---+---+---| + | | x | x | | + |---+---+---+---| + | | | | | + +---+---+---+---+ + + """ + cdef np.int64_t i + cdef int num_octs + if num_cells < 0: + num_octs = selector.count_octs(self, domain_id) + num_cells = num_octs * 4**3 + 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.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") + domains = np.full(num_cells, -1, dtype="int32") + + visitor = NeighbourCellVisitor(self, -1) + # output: level, file_ind and cell_ind of the neighbouring cells + visitor.levels = levels + visitor.file_inds = file_inds + visitor.cell_inds = cell_inds + visitor.domains = domains + # direction to explore and extra parameters of the visitor + visitor.octree = self + visitor.last = -1 + + # Compute indices + self.visit_all_octs(selector, visitor) + + return levels, cell_inds, file_inds, domains def finalize(self): - cdef SelectorObject selector = selection_routines.AlwaysSelector(None) + cdef SelectorObject selector = AlwaysSelector(None) cdef oct_visitors.AssignDomainInd visitor visitor = oct_visitors.AssignDomainInd(self, 1) self.visit_all_octs(selector, visitor) @@ -891,9 +1068,6 @@ cdef class SparseOctreeContainer(OctreeContainer): # called. if self.root_nodes != NULL: free(self.root_nodes) -cdef class RAMSESOctreeContainer(SparseOctreeContainer): - pass - cdef class ARTOctreeContainer(OctreeContainer): def __init__(self, oct_domain_dimensions, domain_left_edge, domain_right_edge, partial_coverage = 0, diff --git a/yt/geometry/oct_visitors.pxd b/yt/geometry/oct_visitors.pxd index 125c69a523e..63920480cf6 100644 --- a/yt/geometry/oct_visitors.pxd +++ b/yt/geometry/oct_visitors.pxd @@ -17,6 +17,12 @@ cdef struct Oct: np.int64_t domain # (opt) addl int index Oct **children # Up to 8 long +cdef struct OctInfo: + np.float64_t left_edge[3] + np.float64_t dds[3] + np.int64_t ipos[3] + np.int32_t level + cdef struct OctPadded: np.int64_t file_ind np.int64_t domain_ind @@ -136,3 +142,34 @@ cdef inline int cind(int i, int j, int k) nogil: # THIS ONLY WORKS FOR CHILDREN. It is not general for zones. return (((i*2)+j)*2+k) +from oct_container cimport OctreeContainer + +cdef class StoreIndex(OctVisitor): + cdef np.int64_t[:,:,:,:] cell_inds + +# cimport oct_container +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 bint other_oct + cdef Oct *neighbour + cdef OctreeContainer octree + cdef OctInfo oi + + cdef void set_neighbour_info(self, Oct *o, int ishift[3]) + + cdef inline np.uint8_t neighbour_rind(self): + cdef int d = (1 << self.oref) + return (((self.neigh_ind[2]*d)+self.neigh_ind[1])*d+self.neigh_ind[0]) + +cdef class NeighbourCellIndexVisitor(BaseNeighbourVisitor): + cdef np.uint8_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.int32_t[::1] domains + diff --git a/yt/geometry/oct_visitors.pyx b/yt/geometry/oct_visitors.pyx index e27f57bd280..cd951d26271 100644 --- a/yt/geometry/oct_visitors.pyx +++ b/yt/geometry/oct_visitors.pyx @@ -10,11 +10,11 @@ Oct visitor functions cimport cython -cimport numpy -import numpy +cimport numpy as np +import numpy as np from yt.utilities.lib.fp_utils cimport * from libc.stdlib cimport malloc, free -from yt.geometry.oct_container cimport OctreeContainer +from yt.geometry.oct_container cimport OctreeContainer, OctInfo from yt.utilities.lib.geometry_utils cimport encode_morton_64bit # Now some visitor functions @@ -338,3 +338,192 @@ cdef class MortonIndexOcts(OctVisitor): np.uint64(coord[1]), np.uint64(coord[2])) self.index += 1 + +# Store cell index +cdef class StoreIndex(OctVisitor): + @cython.boundscheck(False) + @cython.wraparound(False) + @cython.initializedcheck(False) + cdef void visit(self, Oct* o, np.uint8_t selected): + if not selected: return + self.cell_inds[o.domain_ind, self.ind[2], self.ind[1], self.ind[0]] = self.index + + self.index += 1 + +cdef class BaseNeighbourVisitor(OctVisitor): + def __init__(self, OctreeContainer octree, int domain_id = -1): + self.octree = octree + self.neigh_ind[0] = 0 + self.neigh_ind[1] = 0 + self.neigh_ind[2] = 0 + super(BaseNeighbourVisitor, self).__init__(octree, domain_id) + + @cython.boundscheck(False) + @cython.wraparound(False) + @cython.initializedcheck(False) + cdef void set_neighbour_info(self, Oct *o, int ishift[3]): + cdef int i + cdef np.float64_t c, dx + cdef np.int64_t ipos + cdef np.float64_t fcoords[3] + cdef Oct *neighbour + cdef bint local_oct + cdef bint other_oct + dx = 1.0 / ((1 << self.oref) << self.level) + local_oct = True + + # Compute position of neighbouring cell + for i in range(3): + c = (self.pos[i] << self.oref) + fcoords[i] = (c + 0.5 + ishift[i]) * dx / self.octree.nn[i] + # Assuming periodicity + if fcoords[i] < 0: + fcoords[i] += 1 + elif fcoords[i] > 1: + fcoords[i] -= 1 + local_oct &= (0 <= ishift[i] <= 1) + other_oct = not local_oct + + # Use octree to find neighbour + if other_oct: + neighbour = self.octree.get(fcoords, &self.oi, max_level=self.level) + else: + neighbour = o + self.oi.level = self.level + for i in range(3): + self.oi.ipos[i] = (self.pos[i] << self.oref) + 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.oref) + ishift[i])) >> 1 + if (self.oi.ipos[i] << 1) == ipos: + self.oi.ipos[i] = 0 + else: + self.oi.ipos[i] = 1 + self.neighbour = neighbour + + # Index of neighbouring cell within its oct + for i in range(3): + self.neigh_ind[i] = (ishift[i]) % 2 + + self.other_oct = other_oct + if other_oct: + if self.neighbour != NULL: + if self.oi.level == self.level - 1: + # Position within neighbouring oct is stored in oi.ipos + for i in range(3): + self.neigh_ind[i] = self.oi.ipos[i] + elif self.oi.level != self.level: + print('This should not happen! %s %s' % (self.oi.level, self.level)) + self.neighbour = NULL + +# Store neighbouring cell index in current cell +cdef class NeighbourCellIndexVisitor(BaseNeighbourVisitor): + # This piece of code is very much optimizable. Here are possible routes to achieve + # much better performance: + + # - Work oct by oct, which would reduce the number of neighbor lookup + # from 4³=64 to 3³=27, + # - Use faster neighbor lookup method(s). For now, all searches are started from + # the root mesh down to leaf nodes, but we could instead go up the tree from the + # central oct then down to find all neighbors (see e.g. + # https://geidav.wordpress.com/2017/12/02/advanced-octrees-4-finding-neighbor-nodes/). + # - Pre-compute the face-neighbors of all octs. + + # Note that for the last point, algorithms exist that generate the neighbors of all + # octs in O(1) time (https://link.springer.com/article/10.1007/s13319-015-0060-9) + # during the octree construction. + # Another possible solution would be to keep a unordered hash map of all the octs + # indexed by their (3-integers) position. With such structure, finding a neighbor + # takes O(1) time. This could even come as a replacement of the current + # pointer-based octree structure. + + @cython.boundscheck(False) + @cython.wraparound(False) + @cython.initializedcheck(False) + 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.int64_t neigh_domain_ind + if selected == 0: return + # Work at oct level + if self.last == o.domain_ind: return + + self.last = o.domain_ind + + # Loop over cells in and directly around oct + for i in range(-1, 3): + ishift[0] = i + for j in range(-1, 3): + ishift[1] = j + for k in range(-1, 3): + ishift[2] = k + self.set_neighbour_info(o, ishift) + + if not self.other_oct: + neigh_domain_ind = o.domain_ind + neigh_cell_ind = self.neighbour_rind() + elif self.neighbour != NULL: + neigh_domain_ind = self.neighbour.domain_ind + neigh_cell_ind = self.neighbour_rind() + else: + neigh_domain_ind = -1 + neigh_cell_ind = 8 + + self.cell_inds[self.index] = neigh_cell_ind + self.domain_inds[self.index] = neigh_domain_ind + + self.index += 1 + +# Store file position + cell of neighbouring cell in current cell +cdef class NeighbourCellVisitor(BaseNeighbourVisitor): + + @cython.boundscheck(False) + @cython.wraparound(False) + @cython.initializedcheck(False) + cdef void visit(self, Oct* o, np.uint8_t selected): + 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.int32_t neigh_domain + cdef np.uint8_t neigh_level + if selected == 0: return + # Work at oct level + if self.last == o.domain_ind: return + + self.last = o.domain_ind + + # Loop over cells in and directly around oct + for i in range(-1, 3): + ishift[0] = i + for j in range(-1, 3): + ishift[1] = j + for k in range(-1, 3): + ishift[2] = k + self.set_neighbour_info(o, ishift) + + if not self.other_oct: + neigh_level = self.level + neigh_domain = o.domain + neigh_file_ind = o.file_ind + neigh_cell_ind = self.neighbour_rind() + elif self.neighbour != NULL: + neigh_level = self.oi.level + neigh_domain = self.neighbour.domain + neigh_file_ind = self.neighbour.file_ind + neigh_cell_ind = self.neighbour_rind() + else: + neigh_level = 255 + neigh_domain = -1 + neigh_file_ind = -1 + neigh_cell_ind = 8 + + self.levels[self.index] = neigh_level + self.file_inds[self.index] = neigh_file_ind + self.cell_inds[self.index] = neigh_cell_ind + self.domains[self.index] = neigh_domain + + self.index += 1