diff --git a/doc/source/examining/loading_data.rst b/doc/source/examining/loading_data.rst index cfd80aa08e9..30df698e8c0 100644 --- a/doc/source/examining/loading_data.rst +++ b/doc/source/examining/loading_data.rst @@ -220,6 +220,31 @@ overriding a maximum of three units. Allowed unit combinations at the moment are Appropriate errors are thrown for other combinations. +.. rubric:: Stretched grids + +To inform yt that the data is represented on a stretched grid, +define ``stretch_dim`` in a parfile in ``meshlist`` + +.. code-block:: fortran + &meshlist + stretch_dim(1) = 'uni' + / + +and pass the parfile to yt when loading the data: + +.. code-block:: python + + ds = yt.load("output0010.dat", parfiles="amrvac.par") + +Both indexed statements (``stretch_dim(1)='uni'``) and array definitions +(``stretch_dim='uni','',''``) are accepted in the parfile. + +* Only uniform stretching is supported at the moment, to be defined + as ``'uni'`` or ``'uniform'``. +* At present, stretched grids are only supported on a + single level of refinement. + + .. rubric:: Partially supported and unsupported features * a maximum of 100 dust species can be read by yt at the moment. @@ -228,11 +253,6 @@ Appropriate errors are thrown for other combinations. * particle data: currently not supported (but might come later) * staggered grids (AMRVAC 2.2 and later): yt logs a warning if you load staggered datasets, but the flag is currently ignored. -* "stretched grids" are being implemented in yt, but are not yet - fully-supported. (Previous versions of this file suggested they would - "never" be supported, which we hope to prove incorrect once we finish - implementing stretched grids in AMR. At present, stretched grids are - only supported on a single level of refinement.) .. note:: diff --git a/yt/frontends/amrvac/data_structures.py b/yt/frontends/amrvac/data_structures.py index 7fb9787e84d..5aa0a5c3cea 100644 --- a/yt/frontends/amrvac/data_structures.py +++ b/yt/frontends/amrvac/data_structures.py @@ -15,7 +15,7 @@ from more_itertools import always_iterable from yt.config import ytcfg -from yt.data_objects.index_subobjects.grid_patch import AMRGridPatch +from yt.data_objects.index_subobjects.stretched_grid import StretchedGrid from yt.data_objects.static_output import Dataset from yt.funcs import mylog, setdefaultattr from yt.geometry.api import Geometry @@ -52,17 +52,18 @@ def _parse_geometry(geometry_tag: str) -> Geometry: return Geometry(geometry_str.lower()) -class AMRVACGrid(AMRGridPatch): +class AMRVACGrid(StretchedGrid): """A class to populate AMRVACHierarchy.grids, setting parent/children relations.""" _id_offset = 0 - def __init__(self, id, index, level): + def __init__(self, id, cell_widths, filename, index, level, dims): # should use yt's convention (start from 0) - super().__init__(id, filename=index.index_filename, index=index) + super().__init__(id=id, filename=filename, index=index, cell_widths=cell_widths) self.Parent = None self.Children = [] self.Level = level + self.ActiveDimensions = dims def get_global_startindex(self): """Refresh and retrieve the starting index for each dimension at current level. @@ -100,6 +101,16 @@ def __init__(self, ds, dataset_type="amrvac"): self.directory = os.path.dirname(self.index_filename) self.float_type = np.float64 + self.stretch_dim = ["none"] * self.dataset.dimensionality + if self.dataset.namelist is not None: + meshlist = self.dataset.namelist["meshlist"] + if (stretch_dim := meshlist.get("stretch_dim")) is not None: + assert isinstance(stretch_dim, list) + assert len(stretch_dim) <= self.dataset.dimensionality + stretch_dim = ['none' if v is None else v for v in stretch_dim] + assert all(isinstance(x, str) for x in stretch_dim) + self.stretch_dim[:len(stretch_dim)] = stretch_dim + super().__init__(ds, dataset_type) def _detect_output_fields(self): @@ -139,20 +150,121 @@ def _parse_index(self): dx0 = ( domain_width / self.dataset.parameters["domain_nx"] ) # dx at coarsest grid level (YT level 0) - dim = self.dataset.dimensionality + ndim = self.dataset.dimensionality self.grids = np.empty(self.num_grids, dtype="object") + stretched_dims = [not (x == "none" or x == "") for x in self.stretch_dim] + base_stretch = np.ones(3, dtype="float64") + if np.any(stretched_dims): + meshlist = self.dataset.namelist["meshlist"] + stretch_baselevel = meshlist.get("qstretch_baselevel") + match stretch_baselevel: + case (): + assert len(stretch_baselevel) >= ndim + base_stretch[:ndim] = ( + float(b) for b in stretch_baselevel[:ndim] + ) + case float() | int(): + assert sum(stretched_dims) == 1 # exactly one stretched direction + stretched_dim = stretched_dims.index(True) + base_stretch[stretched_dim] = float(stretch_baselevel) + case None: + # compute default values dynamically, just as done in AMRVAC + assert sum(stretched_dims) == 1 # exactly one stretched direction + stretched_dim = stretched_dims.index(True) + 1 # AMRVAC index (1 offset, Fortran convention) + base_stretch[stretched_dim-1] = ( + meshlist[f"xprobmax{stretched_dim}"] + / meshlist[f"xprobmin{stretched_dim}"] + ) ** (1.0 / meshlist[f"domain_nx{stretched_dim}"]) + case _: + raise ValueError( + f"Unknown type for qstretch_baselevel: {type(stretch_baselevel)}" + ) + + qstretch = np.zeros((self.max_level + 2, ndim), dtype="float64") + dxfirst = np.zeros((self.max_level + 2, ndim), dtype="float64") + for dim in range(ndim): + match self.stretch_dim[dim]: + case "none" | "": + continue + case "uni" | "uniform": + qstretch[1, dim] = base_stretch[dim] + dxfirst[1, dim] = ( + domain_width[dim] * (1.0 - qstretch[1, dim]) + / (1.0 - qstretch[1, dim] ** meshlist[f"domain_nx{dim + 1}"]) + ) + qstretch[0, dim] = qstretch[1, dim] ** 2 + dxfirst[0, dim] = dxfirst[1, dim] * (1.0 + qstretch[1, dim]) + if self.max_level > 0: + for ilev in range(2, self.max_level + 2): + qstretch[ilev, dim] = np.sqrt(qstretch[ilev - 1, dim]) + dxfirst[ilev, dim] = dxfirst[ilev - 1, dim] / ( + 1.0 + np.sqrt(qstretch[ilev - 1, dim]) + ) + case "symm" | "symmetric": + raise ValueError( + "Symmetric stretching is not currently supported for AMRVAC data." + ) + case _: + raise ValueError( + f"Unknown stretch_dim {self.stretch_dim[dim]!r} for dimension {dim}." + ) + for igrid, (ytlevel, morton_index) in enumerate( zip(ytlevels, morton_indices, strict=True) ): - dx = dx0 / self.dataset.refine_by**ytlevel - left_edge = xmin + (morton_index - 1) * block_nx * dx + left_edge = np.zeros(ndim, dtype="float64") + right_edge = np.zeros(ndim, dtype="float64") + cell_widths = [] + + for dim in range(ndim): + match self.stretch_dim[dim]: + case "none" | "": + dx = dx0[dim] / self.dataset.refine_by**ytlevel + left_edge[dim] = xmin[dim] + (morton_index[dim] - 1) * block_nx[dim] * dx + right_edge[dim] = left_edge[dim] + block_nx[dim] * dx + cell_widths.append([dx] * block_nx[dim]) + case "uni" | "uniform": + amrvac_level = ytlevel + 1 # AMRVAC uses 1-based indexing for levels + q = qstretch[amrvac_level, dim] + + base = xmin[dim] + 0.5 * dxfirst[amrvac_level, dim] + correction = (q - 1.0) / (q + 1.0) + + # left edge + left_edge[dim] = ( + base * q ** ((morton_index[dim] - 1) * block_nx[dim]) + * (1.0 - correction) + ) + # right edge + right_edge[dim] = ( + base * q ** (morton_index[dim] * block_nx[dim] - 1) + * (1.0 + correction) + ) + # cell widths + cell_widths.append([ + ( + base * q ** ((morton_index[dim] - 1) * block_nx[dim] + i) + * 2.0 * correction + ) for i in range(block_nx[dim]) + ]) + cell_widths.extend([[1.0]] * (3 - ndim)) + cell_widths = np.stack( + np.meshgrid(*cell_widths[::-1], indexing="ij") + )[::-1].reshape(3, -1) # edges and dimensions are filled in a dimensionality-agnostic way - self.grid_left_edge[igrid, :dim] = left_edge - self.grid_right_edge[igrid, :dim] = left_edge + block_nx * dx - self.grid_dimensions[igrid, :dim] = block_nx - self.grids[igrid] = self.grid(igrid, self, ytlevels[igrid]) + self.grid_left_edge[igrid, :ndim] = left_edge + self.grid_right_edge[igrid, :ndim] = right_edge + self.grid_dimensions[igrid, :ndim] = block_nx + self.grids[igrid] = self.grid( + id=igrid, + index=self, + level=ytlevels[igrid], + filename=self.index_filename, + cell_widths=cell_widths, + dims=self.grid_dimensions[igrid], + ) def _populate_grid_objects(self): # required method diff --git a/yt/frontends/amrvac/io.py b/yt/frontends/amrvac/io.py index 4928d832fc9..5d1f5ee2781 100644 --- a/yt/frontends/amrvac/io.py +++ b/yt/frontends/amrvac/io.py @@ -37,7 +37,10 @@ def read_amrvac_namelist(parfiles): parfiles = (os.path.expanduser(pf) for pf in always_iterable(parfiles)) # first merge the namelists - namelists = [f90nml.read(parfile) for parfile in parfiles] + parser = f90nml.Parser() + parser.global_start_index = 1 + + namelists = [parser.read(parfile) for parfile in parfiles] unified_namelist = f90nml.Namelist() for nml in namelists: unified_namelist.patch(nml) diff --git a/yt/utilities/on_demand_imports.py b/yt/utilities/on_demand_imports.py index 224a664fbee..7ab2826d59e 100644 --- a/yt/utilities/on_demand_imports.py +++ b/yt/utilities/on_demand_imports.py @@ -439,6 +439,12 @@ def Namelist(self): return Namelist + @safe_import + def Parser(self): + from f90nml import Parser + + return Parser + _f90nml = f90nml_imports()