Added the ability to read xdmf dumps from Idefix

src/yt_idefix/api.py

@@ -2,6 +2,7 @@
 from .data_structures import (
     IdefixDmpDataset,
     IdefixVtkDataset,
+    IdefixXdmfDataset,
     PlutoVtkDataset,
     PlutoXdmfDataset,
 )

src/yt_idefix/data_structures.py

@@ -31,6 +31,7 @@
 from .fields import (
     IdefixDmpFields,
     IdefixVtkFields,
+    IdefixXdmfFields,
     PlutoFields,
 )
 
@@ -263,7 +264,7 @@ def _cell_centers(self) -> tuple[XCoords, YCoords, ZCoords]:
         )
 
 
-class PlutoXdmfHierarchy(GoodBoyHierarchy):
+class XdmfHierarchy(GoodBoyHierarchy):
     @cached_property
     def _cell_widths(self) -> tuple[XSpans, YSpans, ZSpans]:
         cell_edges = h5_io.read_grid_coordinates(
@@ -842,6 +843,222 @@ def _is_valid(cls, filename: str, *args, **kwargs) -> bool:  # NOQA: ARG003
             return "Idefix" in header
 
 
+class XdmfMixin(Dataset):
+    _index_class = XdmfHierarchy
+
+
+class IdefixXdmfDataset(XdmfMixin, IdefixDataset):
+    _dataset_type = "idefix-xdmf"
+    _field_info_class = IdefixXdmfFields
+
+    @override
+    @classmethod
+    def _is_valid(cls, filename: str, *args, **kwargs) -> bool:  # NOQA: ARG003
+        if not (
+            filename.endswith((".dbl.h5", ".flt.h5"))
+            and os.path.isfile(filename.removesuffix(".h5") + ".xmf")
+        ):
+            return False
+
+        try:
+            fileh = h5py.File(filename, mode="r")
+        except (ImportError, OSError):
+            return False
+        else:
+            base_groups = list(fileh.keys())
+            key_entry = ""
+            for key in base_groups:
+                if "Timestep_" in key:
+                    key_entry = key
+                    break
+            if key_entry == "":
+                fileh.close()
+                return False
+            attributes = list(fileh[f"/{key_entry}"].attrs.keys())
+            if "version" not in attributes:
+                fileh.close()
+                return False
+            version = fileh[f"/{key_entry}"].attrs["version"][0].decode()
+            fileh.close()
+            return "Idefix" in version and "XDMF" in version
+
+    @override
+    def _parse_parameter_file(self):
+        super()._parse_parameter_file()
+
+        # parse the grid
+        coords = h5_io.read_grid_coordinates(
+            self.filename, geometry=self.attributes["geometry"]
+        )
+
+        _default_field_list = [f[0] for f in self._field_info_class.known_other_fields]
+        with h5py.File(self.filename, mode="r") as h5f:
+            root = list(h5f.keys())[0]
+            self._detected_field_list = list(h5f[f"{root}/vars/"])
+            self._field_name_map = {}
+            # some versions of Pluto define field names in lower case
+            # so we normalize standard output field names to upper case
+            # to avoid duplicating data in PlutoFields.known_other_fields
+            for varname in self._detected_field_list:
+                if varname.upper() in _default_field_list:
+                    # The key in hdf5 is case sensitive, so we have to preserve
+                    # the info of original field names for reading data
+                    self._field_name_map[varname.upper()] = varname
+                else:
+                    self._field_name_map[varname] = varname
+
+            self._detected_field_list = self._field_name_map.keys()
+
+        self.domain_dimensions = np.array(coords.array_shape)
+        self.dimensionality = np.count_nonzero(self.domain_dimensions - 1)
+
+        dle = np.array([arr.min() for arr in coords.arrays], dtype="float64")
+        dre = np.array([arr.max() for arr in coords.arrays], dtype="float64")
+
+        # temporary hack to prevent 0-width dimensions for 2D data
+        dre = np.where(dre == dle, dle + 1, dre)
+        self.domain_left_edge = dle
+        self.domain_right_edge = dre
+
+        self._periodicity = (True, True, True)
+
+    @override
+    def _read_data_header(self) -> str:
+        with h5py.File(self.filename, mode="r") as h5f:
+            base_groups = list(h5f.keys())
+            key_entry = ""
+            for key in base_groups:
+                if "Timestep_" in key:
+                    key_entry = key
+                    break
+            self.current_time = float(h5f[f"/{key_entry}"].attrs["time"])
+            version_line = h5f[f"/{key_entry}"].attrs["version"][0].decode()
+            attributes = list(h5f[f"/{key_entry}"].attrs.keys())
+            self.attributes = {}
+            for attribute in attributes:
+                if attribute == "version":
+                    continue
+                if attribute == "dump_datatype" or attribute == "geometry":
+                    self.attributes[attribute] = (
+                        h5f[f"/{key_entry}"].attrs[attribute][0].decode()
+                    )
+                else:
+                    self.attributes[attribute] = h5f[f"/{key_entry}"].attrs[attribute]
+        match = re.search(r"\d+\.\d+\.?\d*[-\w+]*", version_line)
+        if match is None:
+            warnings.warn(
+                "Could not determine code version from file HDF5 file attribute",
+                stacklevel=2,
+            )
+            return "unknown"
+
+        return match.group()
+
+    @override
+    def _setup_geometry(self) -> None:
+        self.geometry = Geometry(self.parameters["definitions"]["geometry"])
+
+    @override
+    def _set_code_unit_attributes(self):
+        """Conversion between physical units and code units."""
+
+        # Base units are length, velocity and density, but here we consider
+        # length, mass and time as base units. Since it can make us easy to calculate
+        # all units when self.units_override is not None.
+
+        # Default values of Pluto's base units which are stored in self.parameters
+        # if they can be read from definitions.h
+        # Otherwise, they are set to the following default values adopted
+        # velocity_unit = km/s
+        # density_unit = mp/cm**3
+        # length_unit = AU
+        defs = self.parameters["definitions"]
+        idefix_units = {}
+        idefix_units["length_unit"] = self.quan(defs.get("length_unit", 1.0), "cm")
+        idefix_units["velocity_unit"] = self.quan(
+            defs.get(
+                "velocity_unit",
+                defs.get("length_unit", 1.0) / defs.get("time_unit", 1.0),
+            ),
+            "cm/s",
+        )
+        idefix_units["density_unit"] = self.quan(
+            defs.get(
+                "density_unit",
+                defs.get("mass_unit", 1.0) / defs.get("length_unit", 1.0) ** 3,
+            ),
+            "g/cm**3",
+        )
+
+        uo_size = len(self.units_override)
+        if uo_size > 0 and uo_size < 3:
+            ytLogger.info(
+                "Less than 3 units were specified in units_override (got %s). "
+                "Need to rely on PLUTO's internal units to derive other units",
+                uo_size,
+            )
+
+        uo_cache = self.units_override.copy()
+        while len(uo_cache) < 3:
+            # If less than 3 units were passed into units_override,
+            # the rest will be chosen from Pluto's units
+            unit, value = idefix_units.popitem()
+            # If any Pluto's base unit is specified in units_override, it'll be preserved
+            if unit in uo_cache:
+                continue
+            uo_cache[unit] = value
+            # Make sure the combination of units are able to derive base units
+            # No need of validation and logging when no unit to be overrided
+            if uo_size > 0:
+                try:
+                    self.__class__._validate_units_override_keys(uo_cache)
+                except ValueError:
+                    # It means the combination is invalid
+                    del uo_cache[unit]
+                else:
+                    ytLogger.info("Relying on %s: %s.", unit, uo_cache[unit])
+
+        bu = _PlutoBaseUnits(uo_cache)
+        for unit, value in bu._data.items():
+            setattr(self, unit, value)
+
+        self.velocity_unit = self.length_unit / self.time_unit
+        self.density_unit = self.mass_unit / self.length_unit**3
+        self.magnetic_unit = (
+            np.sqrt(4.0 * np.pi * self.density_unit) * self.velocity_unit
+        )
+        self.magnetic_unit.convert_to_units("gauss")
+        self.temperature_unit = self.quan(1.0, "K")
+
+    invalid_unit_combinations = [
+        {"magnetic_unit", "velocity_unit", "density_unit"},
+        {"velocity_unit", "time_unit", "length_unit"},
+        {"density_unit", "length_unit", "mass_unit"},
+    ]
+
+    default_units = {
+        "length_unit": "cm",
+        "time_unit": "s",
+        "mass_unit": "g",
+        "velocity_unit": "cm/s",
+        "magnetic_unit": "gauss",
+        "temperature_unit": "K",
+        # this is the one difference with Dataset.default_units:
+        # we accept density_unit as a valid override
+        "density_unit": "g/cm**3",
+    }
+
+    @override
+    def _parse_definitions_header(self) -> None:
+        self.parameters["definitions"] = {}
+        self.parameters["definitions"]["geometry"] = self.attributes["geometry"]
+
+        for attribute in self.attributes:
+            if "unit" not in attribute:
+                continue
+            self.parameters["definitions"][attribute] = self.attributes[attribute]
+
+
 class PlutoVtkDataset(VtkMixin, StaticPlutoDataset):
     _dataset_type = "pluto-vtk"
 
@@ -860,9 +1077,9 @@ def _get_log_file(self) -> str:
         return os.path.join(self.directory, "vtk.out")
 
 
-class PlutoXdmfDataset(StaticPlutoDataset):
+class PlutoXdmfDataset(XdmfMixin, StaticPlutoDataset):
     _dataset_type = "pluto-xdmf"
-    _index_class = PlutoXdmfHierarchy
+    _index_class = XdmfHierarchy
 
     @override
     def _get_log_file(self) -> str:

src/yt_idefix/fields.py

@@ -36,6 +36,24 @@ def setup_fluid_fields(self):
         )
 
 
+class IdefixXdmfFields(FieldInfoContainer):
+    known_other_fields = (
+        ("RHO", ("code_mass / code_length**3", ["density"], None)),  # type: ignore
+        ("VX1", ("code_length / code_time", ["velocity_x"], None)),
+        ("VX2", ("code_length / code_time", ["velocity_y"], None)),
+        ("VX3", ("code_length / code_time", ["velocity_z"], None)),
+        ("BX1", ("code_magnetic", [], None)),
+        ("BX2", ("code_magnetic", [], None)),
+        ("BX3", ("code_magnetic", [], None)),
+        ("PRS", ("code_pressure", ["pressure"], None)),
+    )
+
+    def setup_fluid_fields(self):
+        setup_magnetic_field_aliases(
+            self, "idefix-vtk", [f"BX{idir}" for idir in "123"]
+        )
+
+
 class IdefixDmpFields(FieldInfoContainer):
     known_other_fields = (
         ("Vc-RHO", ("code_mass / code_length**3", ["density"], None)),  # type: ignore

src/yt_idefix/io.py

@@ -67,6 +67,43 @@ def _read_particle_fields(self, chunks, ptf, selector):
         raise NotImplementedError("Particles are not currently supported for Idefix")
 
 
+class IdefixXdmfIOHandler(BaseIOHandler):
+    _dataset_type = "idefix-xdmf"
+
+    def _read_fluid_selection(self, chunks, selector, fields, size):
+        """
+        Filenames are data.<snapnum>.<dbl/flt>.h5
+        <snapnum> needs to be parse from the filename.
+        """
+        if (match := re.search(r"\d{4}", self.ds.filename)) is not None:
+            entry = int(match.group())
+        else:
+            raise RuntimeError(f"Failed to parse output number from {self.ds.filename}")
+
+        data = {field: np.empty(size, dtype="float64") for field in fields}
+
+        with h5py.File(self.ds.filename, "r") as fh:
+            ind = 0
+            for chunk in chunks:
+                for grid in chunk.objs:
+                    nd = 0
+                    for field in fields:
+                        _, fname = field
+                        position = (
+                            f"/Timestep_{entry}/vars/{self.ds._field_name_map[fname]}"
+                        )
+                        field_data = fh[position][:].astype("=f8")
+
+                        while field_data.ndim < 3:
+                            field_data = np.expand_dims(field_data, axis=0)
+
+                        # X3 X2 X1 orderding of fields in Idefix needs to rearranged to X1 X2 X3 order in yt.
+                        values = np.transpose(field_data, axes=(2, 1, 0))
+                        nd = grid.select(selector, values, data[field], ind)
+                    ind += nd
+        return data
+
+
 class IdefixDmpIO(SingleGridIO, BaseParticleIOHandler):
     _dataset_type = "idefix-dmp"
 

tests/test_xdmf.py

@@ -3,7 +3,8 @@
 import pytest
 
 import yt
-from yt_idefix.api import PlutoXdmfDataset
+from yt_idefix.api import IdefixXdmfDataset, PlutoXdmfDataset
+from yt_idefix.data_structures import XdmfMixin
 
 pytest.importorskip("h5py")
 
@@ -15,12 +16,20 @@
 
 def test_class_validation(xdmf_file):
     file = xdmf_file
-    assert PlutoXdmfDataset._is_valid(str(file["path"]))
+    cls = {
+        "idefix": IdefixXdmfDataset,
+        "pluto": PlutoXdmfDataset,
+    }[file["kind"]]
+    assert cls._is_valid(str(file["path"]))
 
 
 def test_load_class(xdmf_file):
     file = xdmf_file
-    PlutoXdmfDataset(str(file["path"]))
+    cls = {
+        "idefix": IdefixXdmfDataset,
+        "pluto": PlutoXdmfDataset,
+    }[file["kind"]]
+    cls(str(file["path"]))
 
 
 # TODO: make this a pytest-mpl test
@@ -38,4 +47,4 @@ def test_projection_plot(xdmf_file):
 
 def test_load_magic(xdmf_file):
     ds = yt.load(xdmf_file["path"], geometry=xdmf_file["geometry"])
-    assert isinstance(ds, PlutoXdmfDataset)
+    assert isinstance(ds, XdmfMixin)