diff --git a/src/yt_idefix/api.py b/src/yt_idefix/api.py index 0da76a94..5ff13b0e 100644 --- a/src/yt_idefix/api.py +++ b/src/yt_idefix/api.py @@ -2,6 +2,7 @@ from .data_structures import ( IdefixDmpDataset, IdefixVtkDataset, + IdefixXdmfDataset, PlutoVtkDataset, PlutoXdmfDataset, ) diff --git a/src/yt_idefix/data_structures.py b/src/yt_idefix/data_structures.py index a53e258a..5781ff14 100644 --- a/src/yt_idefix/data_structures.py +++ b/src/yt_idefix/data_structures.py @@ -32,6 +32,7 @@ from .fields import ( IdefixDmpFields, IdefixVtkFields, + IdefixXdmfFields, PlutoFields, ) @@ -264,7 +265,7 @@ def _cell_centers(self) -> tuple[XCoords, YCoords, ZCoords]: ) -class PlutoXdmfHierarchy(GoodBoyHierarchy): +class XdmfHierarchy(GoodBoyHierarchy): _load_requirements = ["inifix", "h5py"] @cached_property @@ -845,36 +846,8 @@ def _is_valid(cls, filename: str, *args, **kwargs) -> bool: # NOQA: ARG003 return "Idefix" in header -class PlutoVtkDataset(VtkMixin, StaticPlutoDataset): - _dataset_type = "pluto-vtk" - - @override - @classmethod - def _is_valid(cls, filename: str, *args, **kwargs) -> bool: # NOQA: ARG003 - try: - header = vtk_io.read_header(filename) - except Exception: - return False - else: - return "PLUTO" in header - - @override - def _get_log_file(self) -> str: - return os.path.join(self.directory, "vtk.out") - - -class PlutoXdmfDataset(StaticPlutoDataset): - _dataset_type = "pluto-xdmf" - _index_class = PlutoXdmfHierarchy - - @override - def _get_log_file(self) -> str: - if (suffix := re.search(r"(flt|dbl)\.h5$", self.filename)) is not None: - return os.path.join(self.directory, f"{suffix.group()}.out") - else: - raise RuntimeError( - f"Failed to detect log file associated with {self.filename}" - ) +class XdmfMixin(Dataset): + _index_class = XdmfHierarchy @override def _parse_parameter_file(self): @@ -893,9 +866,7 @@ def _parse_parameter_file(self): super()._parse_parameter_file() # parse the grid - coords = h5_io.read_grid_coordinates( - self.filename, geometry=self.parameters["definitions"]["geometry"] - ) + coords = h5_io.read_grid_coordinates(self.filename, geometry=self.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: @@ -928,6 +899,210 @@ def _parse_parameter_file(self): self._periodicity = (True, True, True) + +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 _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" + + @override + @classmethod + def _is_valid(cls, filename: str, *args, **kwargs) -> bool: # NOQA: ARG003 + try: + header = vtk_io.read_header(filename) + except Exception: + return False + else: + return "PLUTO" in header + + @override + def _get_log_file(self) -> str: + return os.path.join(self.directory, "vtk.out") + + +class PlutoXdmfDataset(XdmfMixin, StaticPlutoDataset): + _dataset_type = "pluto-xdmf" + _index_class = XdmfHierarchy + + @override + def _get_log_file(self) -> str: + if (suffix := re.search(r"(flt|dbl)\.h5$", self.filename)) is not None: + return os.path.join(self.directory, f"{suffix.group()}.out") + else: + raise RuntimeError( + f"Failed to detect log file associated with {self.filename}" + ) + def _read_data_header(self) -> str: grid_file = os.path.join(self.directory, "grid.out") if not os.path.isfile(grid_file): diff --git a/src/yt_idefix/fields.py b/src/yt_idefix/fields.py index 50bed767..ac4976a5 100644 --- a/src/yt_idefix/fields.py +++ b/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 diff --git a/src/yt_idefix/io.py b/src/yt_idefix/io.py index 53462f23..198251fd 100644 --- a/src/yt_idefix/io.py +++ b/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...h5 + 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" diff --git a/tests/test_xdmf.py b/tests/test_xdmf.py index 650e43c0..a2cac2e8 100644 --- a/tests/test_xdmf.py +++ b/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)