Configurable Grid Sparsity Ratio Limit (#361)

* Add option to change grid sparsity ratio limit by setting env MDIO_GRID_SPARSITY_RATIO_LIMIT.

* Add option to change grid sparsity ratio limit by setting env MDIO_GRID_SPARSITY_RATIO_LIMIT.

* Provide more context in GridTraceSparsityError rather than as a warning.  Rename MDIO__GRID__SPARSITY_RATIO_LIMIT.

* Add option to change grid sparsity ratio limit by setting env MDIO_GRID_SPARSITY_RATIO_LIMIT.

* Provide more context in GridTraceSparsityError rather than as a warning.  Rename MDIO__GRID__SPARSITY_RATIO_LIMIT.

* FixTestImport4DSparse.

* make env vars consistent to upcoming standard

---------

Co-authored-by: Altay Sansal <[email protected]>

Author: markspec

src/mdio/converters/exceptions.py

@@ -1,6 +1,18 @@
 """Custom exceptions for MDIO converters."""
 
 
+class EnvironmentFormatError(Exception):
+    """Raised when environment variable is of the wrong format."""
+
+    def __init__(self, name, format, msg: str = ""):
+        """Initialize error."""
+        self.message = (
+            f"Environment variable: {name} not of expected format: {format}. "
+        )
+        self.message += f"\n{msg}" if msg else ""
+        super().__init__(self.message)
+
+
 class GridTraceCountError(Exception):
     """Raised when grid trace counts don't match the SEG-Y trace count."""
 
@@ -20,11 +32,11 @@ def __init__(self, grid_traces, segy_traces):
 class GridTraceSparsityError(Exception):
     """Raised when mdio grid will be sparsely populated from SEG-Y traces."""
 
-    def __init__(self, shape, num_traces):
+    def __init__(self, shape, num_traces, msg: str = ""):
         """Initialize error."""
         self.message = (
             f"Grid shape: {shape} but SEG-Y tracecount: {num_traces}. "
             "This grid is very sparse and most likely user error with indexing."
         )
-
+        self.message += f"\n{msg}" if msg else ""
         super().__init__(self.message)

src/mdio/converters/segy.py

@@ -16,6 +16,7 @@
 import zarr
 
 from mdio.api.io_utils import process_url
+from mdio.converters.exceptions import EnvironmentFormatError
 from mdio.converters.exceptions import GridTraceCountError
 from mdio.converters.exceptions import GridTraceSparsityError
 from mdio.core import Grid
@@ -66,42 +67,61 @@ def grid_density_qc(grid: Grid, num_traces: int) -> None:
     logged.  If density is less than 10% an exception is raised. To ignore
     trace sparsity check set environment variable:
         MDIO_IGNORE_CHECKS = True
+    To change the ratio set the environment variable:
+        MDIO__GRID__SPARSITY_RATIO_LIMIT = 10
 
     Args:
         grid: The grid instance to check.
         num_traces: Expected number of traces.
 
     Raises:
-        GridTraceSparsityError: When the grid is too sparse.
+        GridTraceSparsityError: Raised if the grid is significantly larger
+            than the number of traces in the SEG-Y file. By default the error
+            is raised if the grid is more than 10 times larger than the number
+            of traces in the SEG-Y file. This can be disabled by setting the
+            environment variable `MDIO_IGNORE_CHECKS` to `True`. The limit can
+            be changed by setting the environment variable
+            `MDIO__GRID__SPARSITY_RATIO_LIMIT`.
+        EnvironmentFormatError: Raised if the environment variable
+            MDIO__GRID__SPARSITY_RATIO_LIMIT is not a float.
     """
     grid_traces = np.prod(grid.shape[:-1], dtype=np.uint64)  # Exclude sample
     dims = {k: v for k, v in zip(grid.dim_names, grid.shape)}  # noqa: B905
 
     logger.debug(f"Dimensions: {dims}")
     logger.debug(f"num_traces = {num_traces}")
 
-    # Extreme case where the grid is very sparse (usually user error)
-    if grid_traces > 10 * num_traces:
-        logger.warning("WARNING: Sparse mdio grid detected!")
-        for dim_name in grid.dim_names:
-            dim_min = grid.get_min(dim_name)
-            dim_max = grid.get_max(dim_name)
-            logger.warning(f"{dim_name} min: {dim_min} max: {dim_max}")
-        if os.getenv("MDIO_IGNORE_CHECKS", False):
-            # Do not raise an exception if MDIO_IGNORE_CHECK is False
-            pass
-        else:
-            raise GridTraceSparsityError(grid.shape, num_traces)
+    grid_sparsity_ratio_limit = os.getenv("MDIO__GRID__SPARSITY_RATIO_LIMIT", 10)
+    try:
+        grid_sparsity_ratio_limit_ = float(grid_sparsity_ratio_limit)
+    except ValueError:
+        raise EnvironmentFormatError(
+            "MDIO__GRID__SPARSITY_RATIO_LIMIT", "float"
+        ) from None
 
     # Warning if we have above 50% sparsity.
-    if grid_traces > 2 * num_traces:
+    msg = ""
+    if grid_traces > min(2, grid_sparsity_ratio_limit_) * num_traces:
         msg = (
             f"Proposed ingestion grid is sparse. Ingestion grid: {dims}. "
             f"SEG-Y trace count:{num_traces}, grid trace count: {grid_traces}."
         )
+        for dim_name in grid.dim_names:
+            dim_min = grid.get_min(dim_name)
+            dim_max = grid.get_max(dim_name)
+            msg += f"\n{dim_name} min: {dim_min} max: {dim_max}"
 
         logger.warning(msg)
 
+    # Extreme case where the grid is very sparse (usually user error)
+    if grid_traces > grid_sparsity_ratio_limit_ * num_traces:
+        logger.warning("WARNING: Sparse mdio grid detected!")
+        if os.getenv("MDIO__IGNORE_CHECKS", False):
+            # Do not raise an exception if MDIO_IGNORE_CHECK is False
+            pass
+        else:
+            raise GridTraceSparsityError(grid.shape, num_traces, msg)
+
 
 def segy_to_mdio(
     segy_path: str,

tests/integration/test_segy_import_export.py

@@ -139,6 +139,58 @@ def test_import_4d_segy(
         assert grid.select_dim("sample") == samples_exp
 
 
+@pytest.mark.parametrize("header_locations", [(17, 137, 13)])
+@pytest.mark.parametrize("header_names", [("shot_point", "cable", "channel")])
+@pytest.mark.parametrize("header_types", [("int32", "int16", "int32")])
+@pytest.mark.parametrize("endian", ["big"])
+@pytest.mark.parametrize(
+    "grid_overrides",
+    [
+        None,
+    ],
+)
+@pytest.mark.parametrize(
+    "chan_header_type",
+    [
+        StreamerShotGeometryType.A,
+    ],
+)
+class TestImport4DSparse:
+    """Test for 4D segy import with grid overrides."""
+
+    def test_import_4d_segy(
+        self,
+        segy_mock_4d_shots,
+        zarr_tmp,
+        header_locations,
+        header_names,
+        header_types,
+        endian,
+        grid_overrides,
+        chan_header_type,
+    ):
+        """Test importing a SEG-Y file to MDIO."""
+        import os
+
+        from mdio.converters.exceptions import GridTraceSparsityError
+
+        segy_path = segy_mock_4d_shots[chan_header_type]
+        os.environ["MDIO__GRID__SPARSITY_RATIO_LIMIT"] = "1.1"
+
+        with pytest.raises(GridTraceSparsityError):
+            segy_to_mdio(
+                segy_path=segy_path,
+                mdio_path_or_buffer=zarr_tmp.__str__(),
+                index_bytes=header_locations,
+                index_names=header_names,
+                index_types=header_types,
+                chunksize=(8, 2, 128, 1024),
+                overwrite=True,
+                endian=endian,
+                grid_overrides=grid_overrides,
+            )
+
+
 @pytest.mark.parametrize("header_locations", [(17, 13)])
 @pytest.mark.parametrize("header_names", [("inline", "crossline")])
 @pytest.mark.parametrize("endian", ["big"])