Fix statistics calculations (TGSAI#679)

* fix stats calculation and tolerances

* switch allclose assertion at attribute comparison

* make values list

* drop histogram from value assertion

Author: tasansal

src/mdio/segy/_workers.py

@@ -134,14 +134,15 @@ def trace_worker(  # noqa: PLR0913
     ds_to_write = dataset[worker_variables]
 
     if header_key in worker_variables:
-        # Create temporary array for headers with the correct shape
-        # TODO(BrianMichell): Implement this better so that we can enable fill values without changing the code. #noqa: TD003
+        # TODO(BrianMichell): Implement this better so that we can enable fill values without changing the code
+        # https://github.com/TGSAI/mdio-python/issues/584
         tmp_headers = np.zeros_like(dataset[header_key])
         tmp_headers[not_null] = traces.header
         # Create a new Variable object to avoid copying the temporary array
         # The ideal solution is to use `ds_to_write[header_key][:] = tmp_headers`
         # but Xarray appears to be copying memory instead of doing direct assignment.
         # TODO(BrianMichell): #614 Look into this further.
+        # https://github.com/TGSAI/mdio-python/issues/584
         ds_to_write[header_key] = Variable(
             ds_to_write[header_key].dims,
             tmp_headers,
@@ -153,8 +154,9 @@ def trace_worker(  # noqa: PLR0913
     fill_value = _get_fill_value(ScalarType(data_variable.dtype.name))
     tmp_samples = np.full_like(data_variable, fill_value=fill_value)
     tmp_samples[not_null] = traces.sample
-    # Create a new Variable object to avoid copying the temporary array
+
     # TODO(BrianMichell): #614 Look into this further.
+    # https://github.com/TGSAI/mdio-python/issues/584
     ds_to_write[data_variable_name] = Variable(
         ds_to_write[data_variable_name].dims,
         tmp_samples,
@@ -164,12 +166,13 @@ def trace_worker(  # noqa: PLR0913
 
     to_mdio(ds_to_write, output_path=output_path, region=region, mode="r+")
 
+    nonzero_samples = np.ma.masked_values(traces.sample, 0, copy=False)
     histogram = CenteredBinHistogram(bin_centers=[], counts=[])
     return SummaryStatistics(
-        count=traces.sample.size,
-        min=traces.sample.min(),
-        max=traces.sample.max(),
-        sum=traces.sample.sum(),
-        sum_squares=(traces.sample**2).sum(),
+        count=nonzero_samples.count(),
+        min=nonzero_samples.min(),
+        max=nonzero_samples.max(),
+        sum=nonzero_samples.sum(dtype="float64"),
+        sum_squares=(np.ma.power(nonzero_samples, 2).sum(dtype="float64")),
         histogram=histogram,
     )

tests/integration/test_segy_roundtrip_teapot.py

@@ -181,15 +181,18 @@ def test_variable_metadata(self, zarr_tmp: Path) -> None:
         """Metadata reading tests."""
         ds = open_mdio(zarr_tmp)
         expected_attrs = {
-            "count": 97354860,
-            "sum": -8594.551666259766,
-            "sumSquares": 40571291.6875,
+            "count": 46854270,
+            "sum": -8594.551589292674,
+            "sumSquares": 40571285.42351971,
             "min": -8.375323295593262,
             "max": 7.723702430725098,
             "histogram": {"counts": [], "binCenters": []},
         }
-        actual_attrs_json = json.loads(ds["amplitude"].attrs["statsV1"])
-        assert actual_attrs_json == expected_attrs
+        actual_attrs = json.loads(ds["amplitude"].attrs["statsV1"])
+        assert actual_attrs.keys() == expected_attrs.keys()
+        actual_attrs.pop("histogram")
+        expected_attrs.pop("histogram")
+        np.testing.assert_allclose(list(actual_attrs.values()), list(expected_attrs.values()))
 
     def test_grid(self, zarr_tmp: Path, teapot_segy_spec: SegySpec) -> None:
         """Test validating MDIO variables."""
@@ -237,23 +240,22 @@ def test_inline_reads(self, zarr_tmp: Path) -> None:
         """Read and compare every 75 inlines' mean and std. dev."""
         ds = open_mdio(zarr_tmp)
         inlines = ds["amplitude"][::75, :, :]
-        mean, std = inlines.mean(), inlines.std()
-        npt.assert_allclose([mean, std], [1.0555277e-04, 6.0027051e-01])
+        mean, std = inlines.mean(dtype="float64"), inlines.std(dtype="float64")
+        npt.assert_allclose([mean, std], [0.00010555267, 0.60027058412])  # 11 precision
 
     def test_crossline_reads(self, zarr_tmp: Path) -> None:
         """Read and compare every 75 crosslines' mean and std. dev."""
         ds = open_mdio(zarr_tmp)
         xlines = ds["amplitude"][:, ::75, :]
-        mean, std = xlines.mean(), xlines.std()
-
-        npt.assert_allclose([mean, std], [-5.0329847e-05, 5.9406823e-01])
+        mean, std = xlines.mean(dtype="float64"), xlines.std(dtype="float64")
+        npt.assert_allclose([mean, std], [-5.03298501828e-05, 0.59406807762])  # 11 precision
 
     def test_zslice_reads(self, zarr_tmp: Path) -> None:
         """Read and compare every 225 z-slices' mean and std. dev."""
         ds = open_mdio(zarr_tmp)
         slices = ds["amplitude"][:, :, ::225]
-        mean, std = slices.mean(), slices.std()
-        npt.assert_allclose([mean, std], [0.005236923, 0.61279935])
+        mean, std = slices.mean(dtype="float64"), slices.std(dtype="float64")
+        npt.assert_allclose([mean, std], [0.00523692339, 0.61279943571])  # 11 precision
 
     @pytest.mark.dependency("test_3d_import")
     def test_3d_export(