(Re)implement dynamic chunking for live mask and non-dimension coordinates (TGSAI#660)

* Re-implement on-the-fly chunking for non-dimension coordinates and live-mask

* refactor a little

* make chunks a little smaller

* ensure non-dim coordinates are compressed

* readjust sizes

---------

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

Author: BrianMichell

src/mdio/builder/templates/abstract_dataset_template.py

@@ -160,22 +160,17 @@ def _add_coordinates(self) -> None:
             )
 
         # Add non-dimension coordinates
-        # TODO(Dmitriy Repin): do chunked write for non-dimensional coordinates and trace_mask
-        # https://github.com/TGSAI/mdio-python/issues/587
-        # The chunk size used for trace mask will be different from the _var_chunk_shape
         for i in range(len(self._coord_names)):
             self._builder.add_coordinate(
                 self._coord_names[i],
                 dimensions=self._coord_dim_names,
                 data_type=ScalarType.FLOAT64,
+                compressor=compressors.Blosc(cname=compressors.BloscCname.zstd),
                 metadata=CoordinateMetadata(units_v1=self._horizontal_coord_unit),
             )
 
     def _add_trace_mask(self) -> None:
         """Add trace mask variables."""
-        # TODO(Dmitriy Repin): do chunked write for non-dimensional coordinates and trace_mask
-        # https://github.com/TGSAI/mdio-python/issues/587
-        # The chunk size used for trace mask will be different from the _var_chunk_shape
         self._builder.add_variable(
             name="trace_mask",
             dimensions=self._dim_names[:-1],  # All dimensions except vertical (the last one)

src/mdio/builder/templates/seismic_2d_prestack_shot.py

@@ -2,6 +2,7 @@
 
 from typing import Any
 
+from mdio.builder.schemas import compressors
 from mdio.builder.schemas.dtype import ScalarType
 from mdio.builder.schemas.v1.variable import CoordinateMetadata
 from mdio.builder.templates.abstract_dataset_template import AbstractDatasetTemplate
@@ -32,33 +33,39 @@ def _add_coordinates(self) -> None:
             self._builder.add_coordinate(name, dimensions=(name,), data_type=ScalarType.INT32)
 
         # Add non-dimension coordinates
+        compressor = compressors.Blosc(cname=compressors.BloscCname.zstd)
         coordinate_metadata = CoordinateMetadata(units_v1=self._horizontal_coord_unit)
         self._builder.add_coordinate(
             "gun",
             dimensions=("shot_point",),
             data_type=ScalarType.UINT8,
+            compressor=compressor,
         )
         self._builder.add_coordinate(
             "source_coord_x",
             dimensions=("shot_point",),
             data_type=ScalarType.FLOAT64,
+            compressor=compressor,
             metadata=coordinate_metadata,
         )
         self._builder.add_coordinate(
             "source_coord_y",
             dimensions=("shot_point",),
             data_type=ScalarType.FLOAT64,
+            compressor=compressor,
             metadata=coordinate_metadata,
         )
         self._builder.add_coordinate(
             "group_coord_x",
             dimensions=("shot_point", "channel"),
             data_type=ScalarType.FLOAT64,
+            compressor=compressor,
             metadata=coordinate_metadata,
         )
         self._builder.add_coordinate(
             "group_coord_y",
             dimensions=("shot_point", "channel"),
             data_type=ScalarType.FLOAT64,
+            compressor=compressor,
             metadata=coordinate_metadata,
         )

src/mdio/builder/templates/seismic_3d_prestack_coca.py

@@ -2,6 +2,7 @@
 
 from typing import Any
 
+from mdio.builder.schemas import compressors
 from mdio.builder.schemas.dtype import ScalarType
 from mdio.builder.schemas.v1.units import AngleUnitModel
 from mdio.builder.schemas.v1.variable import CoordinateMetadata
@@ -59,15 +60,18 @@ def _add_coordinates(self) -> None:
         )
 
         # Add non-dimension coordinates
+        compressor = compressors.Blosc(cname=compressors.BloscCname.zstd)
         self._builder.add_coordinate(
             "cdp_x",
             dimensions=("inline", "crossline"),
             data_type=ScalarType.FLOAT64,
+            compressor=compressor,
             metadata=CoordinateMetadata(units_v1=self._horizontal_coord_unit),
         )
         self._builder.add_coordinate(
             "cdp_y",
             dimensions=("inline", "crossline"),
             data_type=ScalarType.FLOAT64,
+            compressor=compressor,
             metadata=CoordinateMetadata(units_v1=self._horizontal_coord_unit),
         )

src/mdio/builder/templates/seismic_3d_prestack_shot.py

@@ -2,6 +2,7 @@
 
 from typing import Any
 
+from mdio.builder.schemas import compressors
 from mdio.builder.schemas.dtype import ScalarType
 from mdio.builder.schemas.v1.variable import CoordinateMetadata
 from mdio.builder.templates.abstract_dataset_template import AbstractDatasetTemplate
@@ -32,32 +33,38 @@ def _add_coordinates(self) -> None:
             self._builder.add_coordinate(name, dimensions=(name,), data_type=ScalarType.INT32)
 
         # Add non-dimension coordinates
+        compressor = compressors.Blosc(cname=compressors.BloscCname.zstd)
         self._builder.add_coordinate(
             "gun",
             dimensions=("shot_point",),
             data_type=ScalarType.UINT8,
+            compressor=compressor,
         )
         self._builder.add_coordinate(
             "source_coord_x",
             dimensions=("shot_point",),
             data_type=ScalarType.FLOAT64,
+            compressor=compressor,
             metadata=CoordinateMetadata(units_v1=self._horizontal_coord_unit),
         )
         self._builder.add_coordinate(
             "source_coord_y",
             dimensions=("shot_point",),
             data_type=ScalarType.FLOAT64,
+            compressor=compressor,
             metadata=CoordinateMetadata(units_v1=self._horizontal_coord_unit),
         )
         self._builder.add_coordinate(
             "group_coord_x",
             dimensions=("shot_point", "cable", "channel"),
             data_type=ScalarType.FLOAT64,
+            compressor=compressor,
             metadata=CoordinateMetadata(units_v1=self._horizontal_coord_unit),
         )
         self._builder.add_coordinate(
             "group_coord_y",
             dimensions=("shot_point", "cable", "channel"),
             data_type=ScalarType.FLOAT64,
+            compressor=compressor,
             metadata=CoordinateMetadata(units_v1=self._horizontal_coord_unit),
         )

src/mdio/converters/segy.py

@@ -32,6 +32,9 @@
 from mdio.converters.exceptions import GridTraceSparsityError
 from mdio.converters.type_converter import to_structured_type
 from mdio.core.grid import Grid
+from mdio.core.utils_write import MAX_COORDINATES_BYTES
+from mdio.core.utils_write import MAX_SIZE_LIVE_MASK
+from mdio.core.utils_write import get_constrained_chunksize
 from mdio.segy import blocked_io
 from mdio.segy.utilities import get_grid_plan
 
@@ -429,6 +432,32 @@ def enhanced_add_variables() -> None:
     return mdio_template
 
 
+def _chunk_variable(ds: Dataset, target_variable_name: str) -> None:
+    """Determines and sets the chunking for a specific Variable in the Dataset."""
+    # Find variable index by name
+    index = next((i for i, obj in enumerate(ds.variables) if obj.name == target_variable_name), None)
+
+    def determine_target_size(var_type: str) -> int:
+        """Determines the target size (in bytes) for a Variable based on its type."""
+        if var_type == "bool":
+            return MAX_SIZE_LIVE_MASK
+        return MAX_COORDINATES_BYTES
+
+    # Create the chunk grid metadata
+    var_type = ds.variables[index].data_type
+    full_shape = tuple(dim.size for dim in ds.variables[index].dimensions)
+    target_size = determine_target_size(var_type)
+
+    chunk_shape = get_constrained_chunksize(full_shape, var_type, target_size)
+    chunk_grid = RegularChunkGrid(configuration=RegularChunkShape(chunk_shape=chunk_shape))
+
+    # Create variable metadata if it doesn't exist
+    if ds.variables[index].metadata is None:
+        ds.variables[index].metadata = VariableMetadata()
+
+    ds.variables[index].metadata.chunk_grid = chunk_grid
+
+
 def segy_to_mdio(  # noqa PLR0913
     segy_spec: SegySpec,
     mdio_template: AbstractDatasetTemplate,
@@ -487,6 +516,11 @@ def segy_to_mdio(  # noqa PLR0913
 
     _add_grid_override_to_metadata(dataset=mdio_ds, grid_overrides=grid_overrides)
 
+    # Dynamically chunk the variables based on their type
+    _chunk_variable(ds=mdio_ds, target_variable_name="trace_mask")  # trace_mask is a Variable and not a Coordinate
+    for coord in mdio_template.coordinate_names:
+        _chunk_variable(ds=mdio_ds, target_variable_name=coord)
+
     xr_dataset: xr_Dataset = to_xarray_dataset(mdio_ds=mdio_ds)
 
     xr_dataset, drop_vars_delayed = _populate_coordinates(

src/mdio/core/utils_write.py

@@ -9,9 +9,8 @@
     from numpy.typing import DTypeLike
 
 
-MAX_SIZE_LIVE_MASK = 512 * 1024**2
-
-JsonSerializable = str | int | float | bool | None | dict[str, "JsonSerializable"] | list["JsonSerializable"]
+MAX_SIZE_LIVE_MASK = 256 * 1024**2
+MAX_COORDINATES_BYTES = 32 * 1024**2
 
 
 def get_constrained_chunksize(

tests/unit/test_auto_chunking.py

@@ -48,9 +48,9 @@ class TestAutoChunkLiveMask:
         [
             ((100,), (100,)),  # small 1d
             ((100, 100), (100, 100)),  # small 2d
-            ((50000, 50000), (25000, 25000)),  # large 2d
+            ((50000, 50000), (16667, 16667)),  # large 2d
             ((1500, 1500, 1500), (750, 750, 750)),  # large 3d
-            ((1000, 1000, 100, 36), (334, 334, 100, 36)),  # large 4d
+            ((1000, 1000, 100, 36), (250, 250, 100, 36)),  # large 4d
         ],
     )
     def test_auto_chunk_live_mask(
@@ -65,16 +65,14 @@ def test_auto_chunk_live_mask(
     @pytest.mark.parametrize(
         "shape",
         [
-            # Below are >500MiB. Smaller ones tested above
+            # Below are >250MiB. Smaller ones tested above
             (32768, 32768),
             (46341, 46341),
             (86341, 96341),
             (55000, 97500),
             (100000, 100000),
-            (1024, 1024, 1024),
-            (215, 215, 215, 215),
             (512, 216, 512, 400),
-            (74, 74, 74, 74, 74),
+            (64, 128, 64, 32, 64),
             (512, 17, 43, 200, 50),
         ],
     )
@@ -84,6 +82,6 @@ def test_auto_chunk_live_mask_nbytes(self, shape: tuple[int, ...]) -> None:
         result = get_live_mask_chunksize(shape)
         chunk_elements = np.prod(result)
 
-        # We want them to be 500MB +/- 25%
+        # We want them to be 250MB +/- 50%
         assert chunk_elements > MAX_SIZE_LIVE_MASK * 0.75
         assert chunk_elements < MAX_SIZE_LIVE_MASK * 1.25