fix override parameter bug and add tests

tasansal · tasansal · commit fad687ed8472 · 2023-06-06T12:57:17.000-05:00
diff --git a/src/mdio/segy/geometry.py b/src/mdio/segy/geometry.py
@@ -110,12 +110,21 @@ class GridOverrideCommand(ABC):
     def required_keys(self) -> set:
         """Get the set of required keys for the grid override command."""
 
+    @property
+    @abstractmethod
+    def required_parameters(self) -> set:
+        """Get the set of required parameters for the grid override command."""
+
     @abstractmethod
-    def validate(self, index_headers: npt.NDArray, **kwargs) -> None:
+    def validate(
+        self, index_headers: npt.NDArray, grid_overrides: dict[str, bool | int]
+    ) -> None:
         """Validate if this transform should run on the type of data."""
 
     @abstractmethod
-    def transform(self, index_headers: npt.NDArray, **kwargs) -> npt.NDArray:
+    def transform(
+        self, index_headers: npt.NDArray, grid_overrides: dict[str, bool | int]
+    ) -> dict[str, npt.NDArray]:
         """Perform the grid transform."""
 
     @property
@@ -129,25 +138,42 @@ def check_required_keys(self, index_headers: npt.NDArray) -> None:
         if not self.required_keys.issubset(index_names):
             raise GridOverrideKeysError(self.name, self.required_keys)
 
+    def check_required_params(self, grid_overrides: dict[str, str | int]) -> None:
+        """Check if all required keys are present in the index headers."""
+        if self.required_parameters is None:
+            return
+
+        passed_parameters = set(grid_overrides.keys())
+
+        if not self.required_parameters.issubset(passed_parameters):
+            missing_params = self.required_parameters - passed_parameters
+            raise GridOverrideMissingParameterError(self.name, missing_params)
+
 
 class AutoChannelWrap(GridOverrideCommand):
     """Automatically determine Streamer acquisition type."""
 
     required_keys = {"shot", "cable", "channel"}
+    required_parameters = None
 
-    def validate(self, index_headers: npt.NDArray, **kwargs) -> None:
+    def validate(
+        self, index_headers: npt.NDArray, grid_overrides: dict[str, bool | int]
+    ) -> None:
         """Validate if this transform should run on the type of data."""
-        self.check_required_keys(index_headers)
-
-        if "ChannelWrap" in kwargs:
+        if "ChannelWrap" in grid_overrides:
             raise GridOverrideIncompatibleError(self.name, "ChannelWrap")
 
-        if "CalculateCable" in kwargs:
+        if "CalculateCable" in grid_overrides:
             raise GridOverrideIncompatibleError(self.name, "CalculateCable")
 
-    def transform(self, index_headers: npt.NDArray, **kwargs):
+        self.check_required_keys(index_headers)
+        self.check_required_params(grid_overrides)
+
+    def transform(
+        self, index_headers: npt.NDArray, grid_overrides: dict[str, bool | int]
+    ) -> dict[str, npt.NDArray]:
         """Perform the grid transform."""
-        self.validate(index_headers, **kwargs)
+        self.validate(index_headers, grid_overrides)
 
         result = analyze_streamer_headers(index_headers)
         unique_cables, cable_chan_min, cable_chan_max, geom_type = result
@@ -179,22 +205,25 @@ class ChannelWrap(GridOverrideCommand):
     """Wrap channels to start from one at cable boundaries."""
 
     required_keys = {"shot", "cable", "channel"}
+    required_parameters = {"ChannelsPerCable"}
 
-    def validate(self, index_headers: npt.NDArray, **kwargs) -> None:
+    def validate(
+        self, index_headers: npt.NDArray, grid_overrides: dict[str, bool | int]
+    ) -> None:
         """Validate if this transform should run on the type of data."""
-        self.check_required_keys(index_headers)
-
-        if "ChannelsPerCable" not in kwargs:
-            raise GridOverrideMissingParameterError(self.name, "ChannelsPerCable")
-
-        if "AutoCableChannel" in kwargs:
+        if "AutoChannelWrap" in grid_overrides:
             raise GridOverrideIncompatibleError(self.name, "AutoCableChannel")
 
-    def transform(self, index_headers: npt.NDArray, **kwargs) -> npt.NDArray:
+        self.check_required_keys(index_headers)
+        self.check_required_params(grid_overrides)
+
+    def transform(
+        self, index_headers: npt.NDArray, grid_overrides: dict[str, bool | int]
+    ) -> dict[str, npt.NDArray]:
         """Perform the grid transform."""
-        self.validate(index_headers, **kwargs)
+        self.validate(index_headers, grid_overrides)
 
-        channels_per_cable = kwargs["ChannelsPerCable"]
+        channels_per_cable = grid_overrides["ChannelsPerCable"]
         index_headers["channel"] = (
             index_headers["channel"] - 1
         ) % channels_per_cable + 1
@@ -206,20 +235,25 @@ class CalculateCable(GridOverrideCommand):
     """Calculate cable numbers from unwrapped channels."""
 
     required_keys = {"shot", "cable", "channel"}
+    required_parameters = {"ChannelsPerCable"}
 
-    def validate(self, index_headers: npt.NDArray, **kwargs) -> None:
+    def validate(
+        self, index_headers: npt.NDArray, grid_overrides: dict[str, bool | int]
+    ) -> None:
         """Validate if this transform should run on the type of data."""
-        self.check_required_keys(index_headers)
-
-        if "ChannelsPerCable" not in kwargs:
-            raise GridOverrideMissingParameterError(self.name, "ChannelsPerCable")
-
-        if "AutoCableChannel" in kwargs:
+        if "AutoChannelWrap" in grid_overrides:
             raise GridOverrideIncompatibleError(self.name, "AutoCableChannel")
 
-    def transform(self, index_headers, **kwargs):
+        self.check_required_keys(index_headers)
+        self.check_required_params(grid_overrides)
+
+    def transform(
+        self, index_headers, grid_overrides: dict[str, bool | int]
+    ) -> dict[str, npt.NDArray]:
         """Perform the grid transform."""
-        channels_per_cable = kwargs["ChannelsPerCable"]
+        self.validate(index_headers, grid_overrides)
+
+        channels_per_cable = grid_overrides["ChannelsPerCable"]
         index_headers["cable"] = (
             index_headers["channel"] - 1
         ) // channels_per_cable + 1
@@ -237,24 +271,40 @@ class GridOverrider:
     """
 
     def __init__(self):
-        """Define allowed overrides here."""
+        """Define allowed overrides and parameters here."""
         self.commands = {
             "AutoChannelWrap": AutoChannelWrap(),
             "CalculateCable": CalculateCable(),
             "ChannelWrap": ChannelWrap(),
         }
 
+        self.parameters = self.get_allowed_parameters()
+
+    def get_allowed_parameters(self) -> set:
+        """Get list of allowed parameters from the allowed commands."""
+        parameters = set()
+        for command in self.commands.values():
+            if command.required_parameters is None:
+                continue
+
+            parameters.update(command.required_parameters)
+
+        return parameters
+
     def run(
         self,
         index_headers: npt.NDArray,
         grid_overrides: dict[str, bool],
     ) -> npt.NDArray:
         """Run grid overrides and return result."""
         for override in grid_overrides:
-            if override in self.commands:
-                function = self.commands[override].transform
-                index_headers = function(index_headers, grid_overrides=grid_overrides)
-            else:
+            if override in self.parameters:
+                continue
+
+            if override not in self.commands:
                 raise GridOverrideUnknownError(override)
 
+            function = self.commands[override].transform
+            index_headers = function(index_headers, grid_overrides=grid_overrides)
+
         return index_headers
diff --git a/tests/unit/test_segy_grid_overrides.py b/tests/unit/test_segy_grid_overrides.py
@@ -0,0 +1,137 @@
+"""Check grid overrides."""
+
+
+from __future__ import annotations
+
+import numpy.typing as npt
+import pytest
+from numpy import arange
+from numpy import column_stack
+from numpy import meshgrid
+from numpy import unique
+from numpy.testing import assert_array_equal
+
+from mdio.core import Dimension
+from mdio.segy.exceptions import GridOverrideIncompatibleError
+from mdio.segy.exceptions import GridOverrideMissingParameterError
+from mdio.segy.exceptions import GridOverrideUnknownError
+from mdio.segy.geometry import GridOverrider
+
+
+SHOTS = arange(100, 104, dtype="int32")
+CABLES = arange(11, 15, dtype="int32")
+RECEIVERS = arange(1, 6, dtype="int32")
+
+
+@pytest.fixture
+def mock_streamer_headers() -> dict[str, npt.NDArray]:
+    """Generate dictionary of mocked streamer index headers."""
+    grids = meshgrid(SHOTS, CABLES, RECEIVERS, indexing="ij")
+    permutations = column_stack([grid.ravel() for grid in grids])
+
+    # Make channel from receiver ids
+    for shot in SHOTS:
+        shot_mask = permutations[:, 0] == shot
+        permutations[shot_mask, -1] = arange(1, len(CABLES) * len(RECEIVERS) + 1)
+
+    result = dict(
+        shot=permutations[:, 0],
+        cable=permutations[:, 1],
+        channel=permutations[:, 2],
+    )
+
+    return result
+
+
+class TestStreamerGridOverrides:
+    """Check grid overrides for shot data with streamer acquisition."""
+
+    def test_channel_wrap(self, mock_streamer_headers: npt.NDArray) -> None:
+        """Test the ChannelWrap command."""
+        grid_overrides = {"ChannelWrap": True, "ChannelsPerCable": len(RECEIVERS)}
+
+        overrider = GridOverrider()
+        results = overrider.run(mock_streamer_headers, grid_overrides)
+
+        dims = []
+        for index_name, index_coords in results.items():
+            dim_unique = unique(index_coords)
+            dims.append(Dimension(coords=dim_unique, name=index_name))
+
+        assert_array_equal(dims[0], SHOTS)
+        assert_array_equal(dims[1], CABLES)
+        assert_array_equal(dims[2], RECEIVERS)
+
+    def test_calculate_cable(self, mock_streamer_headers: npt.NDArray) -> None:
+        """Test the CalculateCable command."""
+        grid_overrides = {
+            "CalculateCable": True,
+            "ChannelsPerCable": len(RECEIVERS),
+        }
+
+        overrider = GridOverrider()
+        results = overrider.run(mock_streamer_headers, grid_overrides)
+
+        dims = []
+        for index_name, index_coords in results.items():
+            dim_unique = unique(index_coords)
+            dims.append(Dimension(coords=dim_unique, name=index_name))
+
+        # We need channels because unwrap isn't done here
+        channels = unique(mock_streamer_headers["channel"])
+
+        # We reset the cables to start from 1.
+        cables = arange(1, len(CABLES) + 1, dtype="uint32")
+
+        assert_array_equal(dims[0], SHOTS)
+        assert_array_equal(dims[1], cables)
+        assert_array_equal(dims[2], channels)
+
+    def test_wrap_and_calc_cable(self, mock_streamer_headers: npt.NDArray) -> None:
+        """Test the combined ChannelWrap and CalculateCable commands."""
+        grid_overrides = {
+            "CalculateCable": True,
+            "ChannelWrap": True,
+            "ChannelsPerCable": len(RECEIVERS),
+        }
+
+        overrider = GridOverrider()
+        results = overrider.run(mock_streamer_headers, grid_overrides)
+
+        dims = []
+        for index_name, index_coords in results.items():
+            dim_unique = unique(index_coords)
+            dims.append(Dimension(coords=dim_unique, name=index_name))
+
+        # We reset the cables to start from 1.
+        cables = arange(1, len(CABLES) + 1, dtype="uint32")
+
+        assert_array_equal(dims[0], SHOTS)
+        assert_array_equal(dims[1], cables)
+        assert_array_equal(dims[2], RECEIVERS)
+
+    def test_missing_param(self, mock_streamer_headers: npt.NDArray) -> None:
+        """Test missing parameters for the commands."""
+        overrider = GridOverrider()
+        with pytest.raises(GridOverrideMissingParameterError):
+            overrider.run(mock_streamer_headers, {"ChannelWrap": True})
+
+        with pytest.raises(GridOverrideMissingParameterError):
+            overrider.run(mock_streamer_headers, {"CalculateCable": True})
+
+    def test_incompatible_overrides(self, mock_streamer_headers: npt.NDArray) -> None:
+        """Test commands that can't be run together."""
+        overrider = GridOverrider()
+        with pytest.raises(GridOverrideIncompatibleError):
+            grid_overrides = {"ChannelWrap": True, "AutoChannelWrap": True}
+            overrider.run(mock_streamer_headers, grid_overrides)
+
+        with pytest.raises(GridOverrideIncompatibleError):
+            grid_overrides = {"CalculateCable": True, "AutoChannelWrap": True}
+            overrider.run(mock_streamer_headers, grid_overrides)
+
+    def test_unknown_override(self, mock_streamer_headers: npt.NDArray) -> None:
+        """Test exception if user provides a command that's not allowed."""
+        overrider = GridOverrider()
+        with pytest.raises(GridOverrideUnknownError):
+            overrider.run(mock_streamer_headers, {"WrongCommand": True})
