MPI-AMRVAC support for uniformly stretched grids

yt/frontends/amrvac/data_structures.py

@@ -12,10 +12,11 @@
 from pathlib import Path
 
 import numpy as np
+from itertools import product
 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 +53,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):
         # <level> 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 +102,14 @@ 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
+                self.stretch_dim = stretch_dim
+
         super().__init__(ds, dataset_type)
 
     def _detect_output_fields(self):
@@ -139,20 +149,117 @@ 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 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 list() | tuple():
+                    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] = stretch_baselevel
+                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(
+                        f"Symmetric stretching is not currently supported for AMRVAC data."
+                    )
+                case _:
+                    raise ValueError(
+                        f"Unknown stretch_dim '{self.stretch_dim[dim]}' 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")
+            cw_aux = []
+
+            dim_count = 0
+            for dim in range(ndim):
+                match self.stretch_dim[dim]:
+                    case "none" | "":
+                        dx = dx0 / self.dataset.refine_by**ytlevel
+                        left_edge[dim] = xmin[dim] + (morton_index[dim] - 1) * block_nx[dim] * dx[dim]
+                        right_edge[dim] = left_edge[dim] + block_nx[dim] * dx[dim]
+                        cw_aux.append([dx[dim]] * block_nx[dim])
+                    case "uni" | "uniform":
+                        # left edge
+                        center1 = (xmin[dim] + 0.5 * dxfirst[ytlevel+1, dim]) * qstretch[ytlevel+1, dim] ** ((morton_index[dim] - 1) * block_nx[dim])
+                        dcenter1 = 2.0 * center1 * (qstretch[ytlevel+1, dim] - 1.0) / (qstretch[ytlevel+1, dim] + 1.0)
+                        left_edge[dim] = center1 - 0.5 * dcenter1
+                        # right edge
+                        center2 = (xmin[dim] + 0.5 * dxfirst[ytlevel+1, dim]) * qstretch[ytlevel+1, dim] ** (morton_index[dim] * block_nx[dim] - 1)
+                        dcenter2 = 2.0 * center2 * (qstretch[ytlevel+1, dim] - 1.0) / (qstretch[ytlevel+1, dim] + 1.0)
+                        right_edge[dim] = center2 + 0.5 * dcenter2
+                        # cell widths
+                        dcenter = [
+                            (
+                                (xmin[dim] + 0.5 * dxfirst[ytlevel+1, dim]) * qstretch[ytlevel+1, dim] ** ((morton_index[dim] - 1) * block_nx[dim] + i)
+                                * 2.0 * (qstretch[ytlevel+1, dim] - 1.0) / (qstretch[ytlevel+1, dim] + 1.0)
+                            ) for i in range(block_nx[dim])
+                        ]
+                        cw_aux.append(dcenter)
+                dim_count += 1
+            while dim_count < 3:
+                cw_aux.append([1.0])
+                dim_count += 1
+            prod = np.array(list(product(*cw_aux[::-1])))
+            cell_widths = (prod.T)[::-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