BUG: Update GIZMO frontend to handle newer GIZMO versions

yt/frontends/gizmo/data_structures.py

@@ -1,6 +1,10 @@
 import os
 
+import numpy as np
+
 from yt.frontends.gadget.data_structures import GadgetHDF5Dataset
+from yt.utilities.cosmology import Cosmology
+from yt.utilities.logger import ytLogger as mylog
 from yt.utilities.on_demand_imports import _h5py as h5py
 
 from .fields import GizmoFieldInfo
@@ -12,7 +16,7 @@ class GizmoDataset(GadgetHDF5Dataset):
     @classmethod
     def _is_valid(cls, filename, *args, **kwargs):
         need_groups = ["Header"]
-        veto_groups = ["FOF", "Group", "Subhalo"]
+        veto_groups = ["Config", "Constants", "FOF", "Group", "Subhalo"]
         valid = True
         valid_fname = filename
         # If passed arg is a directory, look for the .0 file in that dir
@@ -38,9 +42,16 @@ def _is_valid(cls, filename, *args, **kwargs):
             valid = all(ng in fh["/"] for ng in need_groups) and not any(
                 vg in fh["/"] for vg in veto_groups
             )
-            dmetal = "/PartType0/Metallicity"
-            if dmetal not in fh or fh[dmetal].shape[1] < 11:
-                valid = False
+            # From Apr 2021, 7f1f06f, public gizmo includes a header variable
+            # GIZMO_version, which is set to the year of the most recent commit
+            # We should prefer this to checking the metallicity, which might
+            # not exist
+            if "GIZMO_version" not in fh["/Header"].attrs:
+                dmetal = "/PartType0/Metallicity"
+                if dmetal not in fh or (
+                    fh[dmetal].ndim > 1 and fh[dmetal].shape[1] < 11
+                ):
+                    valid = False
             fh.close()
         except Exception:
             valid = False
@@ -49,3 +60,91 @@ def _is_valid(cls, filename, *args, **kwargs):
 
     def _set_code_unit_attributes(self):
         super()._set_code_unit_attributes()
+
+    def _parse_parameter_file(self):
+        hvals = self._get_hvals()
+
+        self.dimensionality = 3
+        self.refine_by = 2
+        self.parameters["HydroMethod"] = "sph"
+        # Set standard values
+
+        # We may have an overridden bounding box.
+        if self.domain_left_edge is None and hvals["BoxSize"] != 0:
+            self.domain_left_edge = np.zeros(3, "float64")
+            self.domain_right_edge = np.ones(3, "float64") * hvals["BoxSize"]
+
+        self.domain_dimensions = np.ones(3, "int32")
+        self._periodicity = (True, True, True)
+
+        self.cosmological_simulation = 1
+
+        self.current_redshift = hvals.get("Redshift", 0.0)
+        if "Redshift" not in hvals:
+            mylog.info("Redshift is not set in Header. Assuming z=0.")
+
+        if "ComovingIntegrationOn" in hvals:
+            # In 1d8479, Nov 2020, public GIZMO updated the names of the Omegas
+            # to include an _, added baryons and radiation and added the
+            # ComovingIntegrationOn field. ComovingIntegrationOn is always set,
+            # but the Omega's are only included if ComovingIntegrationOn is true
+            mylog.debug("Reading cosmological parameters using post-1d8479 format")
+            self.cosmological_simulation = hvals["ComovingIntegrationOn"]
+            if self.cosmological_simulation:
+                self.omega_lambda = hvals["Omega_Lambda"]
+                self.omega_matter = hvals["Omega_Matter"]
+                self.omega_baryon = hvals["Omega_Baryon"]
+                self.omega_radiation = hvals["Omega_Radiation"]
+            self.hubble_constant = hvals["HubbleParam"]
+        elif "OmegaLambda" in hvals:
+            # Should still support GIZMO versions prior to 1d8479 too
+            mylog.info(
+                "ComovingIntegrationOn does not exist, falling back to OmegaLambda",
+            )
+            self.omega_lambda = hvals["OmegaLambda"]
+            self.omega_matter = hvals["Omega0"]
+            self.hubble_constant = hvals["HubbleParam"]
+            self.cosmological_simulation = self.omega_lambda != 0.0
+        else:
+            # If these are not set it is definitely not a cosmological dataset.
+            mylog.debug("No cosmological information found, assuming defaults")
+            self.omega_lambda = 0.0
+            self.omega_matter = 0.0  # Just in case somebody asks for it.
+            self.cosmological_simulation = 0
+            # Hubble is set below for Omega Lambda = 0.
+
+        if not self.cosmological_simulation:
+            mylog.info(
+                "ComovingIntegrationOn != 1 or (not found "
+                "and OmegaLambda is 0.0), so we are turning off Cosmology.",
+            )
+            self.hubble_constant = 1.0  # So that scaling comes out correct
+            self.current_redshift = 0.0
+            # This may not be correct.
+            self.current_time = hvals["Time"]
+        else:
+            # Now we calculate our time based on the cosmology, because in
+            # ComovingIntegration hvals["Time"] will in fact be the expansion
+            # factor, not the actual integration time, so we re-calculate
+            # global time from our Cosmology.
+            cosmo = Cosmology(
+                hubble_constant=self.hubble_constant,
+                omega_matter=self.omega_matter,
+                omega_lambda=self.omega_lambda,
+            )
+            self.current_time = cosmo.lookback_time(self.current_redshift, 1e6)
+            mylog.info(
+                "Calculating time from %0.3e to be %0.3e seconds",
+                hvals["Time"],
+                self.current_time,
+            )
+        self.parameters = hvals
+
+        prefix = os.path.join(self.directory, self.basename.split(".", 1)[0])
+
+        if hvals["NumFiles"] > 1:
+            self.filename_template = f"{prefix}.%(num)s{self._suffix}"
+        else:
+            self.filename_template = self.parameter_filename
+
+        self.file_count = hvals["NumFiles"]

yt/frontends/gizmo/tests/test_outputs.py

@@ -3,7 +3,8 @@
 import yt
 from yt.frontends.gizmo.api import GizmoDataset
 from yt.frontends.gizmo.fields import metal_elements
-from yt.testing import requires_file, requires_module
+from yt.testing import assert_allclose_units, requires_file, requires_module
+from yt.units import Myr
 from yt.utilities.answer_testing.framework import requires_ds, sph_answer
 
 # This maps from field names to weight field names to use for projections
@@ -20,6 +21,8 @@
 g64 = "gizmo_64/output/snap_N64L16_135.hdf5"
 gmhd = "gizmo_mhd_mwdisk/gizmo_mhd_mwdisk.hdf5"
 gmhd_bbox = [[-400, 400]] * 3
+zeld_wg = "gizmo_zeldovich/snapshot_076_wi_gizver.hdf5"
+zeld_ng = "gizmo_zeldovich/snapshot_076_no_gizver.hdf5"
 
 
 @requires_module("h5py")
@@ -32,6 +35,32 @@ def test_gizmo_64():
         yield test
 
 
+@requires_module("h5py")
+@requires_file(zeld_wg)
+@requires_file(zeld_ng)
+def test_gizmo_zeldovich():
+    """
+    Test loading a recent gizmo snapshot that doesn't have cooling/metallicity
+
+    The gizmo_zeldovich file has no metallicity field on the gas particles
+    but is a cosmological dataset run using GIZMO_version=2022. There are two
+    versions of the file, with GIZMO_version (_wg) and without GIZMO_version
+    (_ng). Check that both load as gizmo datasets and correctly pull the
+    cosmological variables. This test should get simpler when the file switches
+    to pytest.
+    """
+    for fn in [zeld_wg, zeld_ng]:
+        ds = yt.load(fn)
+        assert isinstance(ds, GizmoDataset)
+
+        assert ds.cosmological_simulation
+        assert ds.omega_matter == 1.0
+        assert ds.omega_lambda == 0.0
+        # current_time is calculated from the cosmology so this checks if that
+        # was calculated correctly
+        assert_allclose_units(ds.current_time, 1672.0678 * Myr)
+
+
 @requires_module("h5py")
 @requires_file(gmhd)
 def test_gizmo_mhd():

yt/sample_data_registry.json

@@ -724,6 +724,12 @@
     "load_name": "gizmo_mhd_mwdisk.hdf5",
     "url": "https://yt-project.org/data/gizmo_mhd_mwdisk.tar.gz"
   },
+  "gizmo_zeldovich.tar.gz": {
+    "hash": "5f1a73ead736024ffb6c099f84e834ec927d2818c93d7c775d9ff625adaa7c51",
+    "load_kwargs": {},
+    "load_name": "snapshot_076_wi_gizver.hdf5",
+    "url": "https://yt-project.org/data/gizmo_zeldovich.tar.gz"
+  },
   "grs-50-cube.fits.gz": {
     "hash": "1aff152f616d2626c8515c1493e65f07843d9b5f2ba73b7e4271a2dc6a9e6da7",
     "load_kwargs": {},