Add stats to validate_xr_variable

Author: dmitriyrepin

tests/conftest.py

@@ -78,9 +78,7 @@ def empty_mdio_dir(tmp_path_factory: pytest.TempPathFactory) -> Path:
     return tmp_path_factory.mktemp(r"empty_mdio_dir")
 
 
-#
 # Uncomment the function below for local debugging
-#
 # @pytest.fixture(scope="session")
 # def tmp_path_factory() -> pytest.TempPathFactory:
 #     """Custom tmp_path_factory implementation for local debugging."""

tests/integration/test_create_empty.py

@@ -92,57 +92,39 @@ def _validate_empty_mdio_dataset(cls, ds: xr_Dataset, has_headers: bool, is_velo
         assert ds.sizes == {"inline": 345, "crossline": 188, "time": 1501}
 
         # Validate the dimension coordinate variables
-        validate_xr_variable(ds, "inline", {"inline": 345}, UNITS_NONE, np.int32, range(1, 346), get_values)
-        validate_xr_variable(ds, "crossline", {"crossline": 188}, UNITS_NONE, np.int32, range(1, 189), get_values)
-        validate_xr_variable(ds, "time", {"time": 1501}, UNITS_SECOND, np.int32, range(0, 3002, 2), get_values)
+        validate_xr_variable(ds, "inline", {"inline": 345}, UNITS_NONE, np.int32, False, range(1, 346), get_values)
+        validate_xr_variable(
+            ds, "crossline", {"crossline": 188}, UNITS_NONE, np.int32, False, range(1, 189), get_values
+        )
+        validate_xr_variable(ds, "time", {"time": 1501}, UNITS_SECOND, np.int32, False, range(0, 3002, 2), get_values)
 
         # Validate the non-dimensional coordinate variables (should be empty for empty dataset)
-        validate_xr_variable(ds, "cdp_x", {"inline": 345, "crossline": 188}, UNITS_METER, np.float64, None, None)
-        validate_xr_variable(ds, "cdp_y", {"inline": 345, "crossline": 188}, UNITS_METER, np.float64, None, None)
+        validate_xr_variable(ds, "cdp_x", {"inline": 345, "crossline": 188}, UNITS_METER, np.float64)
+        validate_xr_variable(ds, "cdp_y", {"inline": 345, "crossline": 188}, UNITS_METER, np.float64)
 
         if has_headers:
             # Validate the headers (should be empty for empty dataset)
             # Infer the dtype from segy_spec and ignore endianness
             header_dtype = get_teapot_segy_spec().trace.header.dtype.newbyteorder("native")
-            validate_xr_variable(ds, "headers", {"inline": 345, "crossline": 188}, UNITS_NONE, header_dtype, None, None)
-            validate_xr_variable(
-                ds,
-                "segy_file_header",
-                dims={},
-                units=UNITS_NONE,
-                data_type=np.dtype("U1"),
-                expected_values=None,
-                actual_value_generator=None,
-            )
+            validate_xr_variable(ds, "headers", {"inline": 345, "crossline": 188}, UNITS_NONE, header_dtype)
+            validate_xr_variable(ds, "segy_file_header", dims={}, units=UNITS_NONE, data_type=np.dtype("U1"))
         else:
             assert "headers" not in ds.variables
             assert "segy_file_header" not in ds.variables
 
         # Validate the trace mask (should be all True for empty dataset)
-        validate_xr_variable(ds, "trace_mask", {"inline": 345, "crossline": 188}, UNITS_NONE, np.bool_, None, None)
+        validate_xr_variable(ds, "trace_mask", {"inline": 345, "crossline": 188}, UNITS_NONE, np.bool_)
         trace_mask = ds["trace_mask"].values
         assert not np.any(trace_mask), "All traces should be marked as dead in empty dataset"
 
         # Validate the velocity or amplitude data (should be empty)
         if is_velocity:
             validate_xr_variable(
-                ds,
-                "velocity",
-                {"inline": 345, "crossline": 188, "time": 1501},
-                UNITS_METER_PER_SECOND,
-                np.float32,
-                None,
-                None,
+                ds, "velocity", {"inline": 345, "crossline": 188, "time": 1501}, UNITS_METER_PER_SECOND, np.float32
             )
         else:
             validate_xr_variable(
-                ds,
-                "amplitude",
-                {"inline": 345, "crossline": 188, "time": 1501},
-                UNITS_NONE,
-                np.float32,
-                None,
-                None,
+                ds, "amplitude", {"inline": 345, "crossline": 188, "time": 1501}, UNITS_NONE, np.float32
             )
 
     @classmethod
@@ -273,17 +255,19 @@ def test_overwrite_behavior(self, empty_mdio_dir: Path) -> None:
     def test_populate_empty_dataset(self, mdio_with_headers: Path) -> None:
         """Test showing how to populate empty dataset."""
         # Open an empty PostStack3DVelocityTime dataset with SEG-Y 1.0 headers
-        # NOTES:
+        # 
         # When this empty dataset was created from the 'PostStack3DVelocityTime' template and dimensions,
         # * 'inline', 'crossline', and 'time' dimension coordinate variables were created and pre-populated
+        #   NOTE: the 'time' units are specified in the template, so they are not None in this case.
         # * 'cdp_x', 'cdp_y' non-dimensional coordinate variables were created
-        # * 'amplitude' variable was created (the name of this variable is specified in the template)
-        #   HACK: in this example, we will use this variable to store the velocity data
+        #   NOTE: the 'cdp_x' and 'cdp_y' units are specified in the template, so they are not None in this case.
+        # * 'velocity' variable was created (the name of this default variable is specified in the template)
+        #   NOTE: the 'velocity' units are specified in the template, so they are not None in this case.
         # * 'trace_mask' variable was created and pre-populated with 'False' fill values
         #   (all traces are marked as dead)
         # * 'headers' and 'segy_file_header' variables were created (if the dataset was created with
         #   headers not None). The 'headers' variable structured datatype is defined by the HeaderSpec
-        # that was used to create the empty MDIO
+        #   that was used to create the empty MDIO
         # * dataset attribute called 'attributes' was created
         ds = open_mdio(mdio_with_headers)
 
@@ -320,7 +304,9 @@ def test_populate_empty_dataset(self, mdio_with_headers: Path) -> None:
         # 4) Populate dataset's trace mask (optional)
         ds.trace_mask[:] = ~np.isnan(velocity[:, :, 0])
 
-        # 5) Set coordinate and data variable units (optional)
+        # 5) If the units were not set in the template or you want to change the coordinate and data variable units
+        #    you can set the unitsV1 attribute for the coordinate and data variables (optional).
+        #    If you are happy with the units specified in the template, you should skip this step.
         ds.time.attrs["unitsV1"] = TimeUnitModel(time=TimeUnitEnum.MILLISECOND).model_dump(mode="json")
 
         ds.cdp_x.attrs["unitsV1"] = LengthUnitModel(length=LengthUnitEnum.FOOT).model_dump(mode="json")
@@ -376,7 +362,7 @@ def test_populate_empty_dataset(self, mdio_with_headers: Path) -> None:
         output_path_mdio = mdio_with_headers.parent / "populated_empty.mdio"
         to_mdio(ds, output_path=output_path_mdio, mode="w", compute=True)
 
-        # 9) Convert the populated emptyMDIO to SEG-Y
+        # 9) Convert the populated empty MDIO to SEG-Y
         if "headers" in ds.variables:
             # Select the SEG-Y standard to use for the conversion
             custom_segy_spec = get_segy_standard(1.0)

tests/integration/test_segy_roundtrip_teapot.py

@@ -243,13 +243,15 @@ def test_grid(self, teapot_mdio_tmp: Path, teapot_segy_spec: SegySpec) -> None:
         ds = open_mdio(teapot_mdio_tmp)
 
         # Validate the dimension coordinate variables
-        validate_xr_variable(ds, "inline", {"inline": 345}, UNITS_NONE, np.int32, range(1, 346), get_values)
-        validate_xr_variable(ds, "crossline", {"crossline": 188}, UNITS_NONE, np.int32, range(1, 189), get_values)
-        validate_xr_variable(ds, "time", {"time": 1501}, UNITS_SECOND, np.int32, range(0, 3002, 2), get_values)
+        validate_xr_variable(ds, "inline", {"inline": 345}, UNITS_NONE, np.int32, False, range(1, 346), get_values)
+        validate_xr_variable(
+            ds, "crossline", {"crossline": 188}, UNITS_NONE, np.int32, False, range(1, 189), get_values
+        )
+        validate_xr_variable(ds, "time", {"time": 1501}, UNITS_SECOND, np.int32, False, range(0, 3002, 2), get_values)
 
         # Validate the non-dimensional coordinate variables
-        validate_xr_variable(ds, "cdp_x", {"inline": 345, "crossline": 188}, UNITS_METER, np.float64, None, None)
-        validate_xr_variable(ds, "cdp_y", {"inline": 345, "crossline": 188}, UNITS_METER, np.float64, None, None)
+        validate_xr_variable(ds, "cdp_x", {"inline": 345, "crossline": 188}, UNITS_METER, np.float64)
+        validate_xr_variable(ds, "cdp_y", {"inline": 345, "crossline": 188}, UNITS_METER, np.float64)
 
         # Validate the headers
         # We have a custom set of headers since we used customize_segy_specs()
@@ -262,22 +264,17 @@ def test_grid(self, teapot_mdio_tmp: Path, teapot_segy_spec: SegySpec) -> None:
             {"inline": 345, "crossline": 188},
             UNITS_NONE,
             data_type.newbyteorder("native"),  # mdio saves with machine endian, spec could be different endian
+            False,
             range(1, 346),
             get_inline_header_values,
         )
 
         # Validate the trace mask
-        validate_xr_variable(ds, "trace_mask", {"inline": 345, "crossline": 188}, UNITS_NONE, np.bool, None, None)
+        validate_xr_variable(ds, "trace_mask", {"inline": 345, "crossline": 188}, UNITS_NONE, np.bool)
 
         # validate the amplitude data
         validate_xr_variable(
-            ds,
-            "amplitude",
-            {"inline": 345, "crossline": 188, "time": 1501},
-            UNITS_NONE,
-            np.float32,
-            None,
-            None,
+            ds, "amplitude", {"inline": 345, "crossline": 188, "time": 1501}, UNITS_NONE, np.float32, True
         )
 
     def test_inline_reads(self, teapot_mdio_tmp: Path) -> None:

tests/integration/testing_helpers.py

@@ -24,7 +24,8 @@ def validate_xr_variable(  # noqa PLR0913
     dims: dict[int],
     units: AllUnitModel,
     data_type: np.dtype,
-    expected_values: range | None,
+    has_stats: bool = False,
+    expected_values: range | None = None,
     actual_value_generator: Callable[[xr.DataArray], np.ndarray] | None = None,
 ) -> None:
     """Validate the properties of a variable in an Xarray dataset."""
@@ -47,10 +48,16 @@ def validate_xr_variable(  # noqa PLR0913
     else:
         assert data_type == v.dtype
 
-    assert v.attrs.get("statsV1", None) is None, "StatsV1 should be empty for empty dataset variables"
+    stats = v.attrs.get("statsV1", None)
+    if has_stats:
+        assert stats is not None, "StatsV1 should not be empty for dataset variables with stats"
+    else:
+        assert stats is None, "StatsV1 should be empty for dataset variables without stats"
 
     if units is not None:
-        assert v.attrs == {"unitsV1": units.model_dump(mode="json")}
+        units_v1 = v.attrs.get("unitsV1", None)
+        assert units_v1 is not None, "UnitsV1 should not be empty for dataset variables with units"
+        assert units_v1 == units.model_dump(mode="json")
     else:
         assert "unitsV1" not in v.attrs, "UnitsV1 should not exist for unit-unaware variables"