OPTIONAL: 'Enable PreStackCocaGathers3D - take two' with 'Export MDIO v1 to SEGY' merged in

src/mdio/converters/mdio.py

@@ -5,34 +5,37 @@
 import os
 from pathlib import Path
 from tempfile import TemporaryDirectory
+from typing import TYPE_CHECKING
 
+import dask.array as da
 import numpy as np
+import xarray as xr
 from psutil import cpu_count
 from tqdm.dask import TqdmCallback
 
-from mdio import MDIOReader
+from mdio.core.utils_read import open_zarr_dataset
 from mdio.segy.blocked_io import to_segy
 from mdio.segy.creation import concat_files
 from mdio.segy.creation import mdio_spec_to_segy
-from mdio.segy.utilities import segy_export_rechunker
+from mdio.segy.utilities import segy_export_rechunker  
 
 try:
     import distributed
 except ImportError:
     distributed = None
 
+if TYPE_CHECKING:
+    from segy.schema import SegySpec
+
+    from mdio.core.storage_location import StorageLocation
 
 default_cpus = cpu_count(logical=True)
 NUM_CPUS = int(os.getenv("MDIO__EXPORT__CPU_COUNT", default_cpus))
 
-
-def mdio_to_segy(  # noqa: PLR0912, PLR0913
-    mdio_path_or_buffer: str,
-    output_segy_path: str,
-    endian: str = "big",
-    access_pattern: str = "012",
-    storage_options: dict = None,
-    new_chunks: tuple[int, ...] = None,
+def mdio_to_segy(  # noqa: PLR0912, PLR0913, PLR0915
+    segy_spec: SegySpec,
+    input_location: StorageLocation,
+    output_location: StorageLocation,
     selection_mask: np.ndarray = None,
     client: distributed.Client = None,
 ) -> None:
@@ -47,13 +50,9 @@ def mdio_to_segy(  # noqa: PLR0912, PLR0913
     A `selection_mask` can be provided (same shape as spatial grid) to export a subset.
 
     Args:
-        mdio_path_or_buffer: Input path where the MDIO is located.
-        output_segy_path: Path to the output SEG-Y file.
-        endian: Endianness of the input SEG-Y. Rev.2 allows little endian. Default is 'big'.
-        access_pattern: This specificies the chunk access pattern. Underlying zarr.Array must
-            exist. Examples: '012', '01'
-        storage_options: Storage options for the cloud storage backend. Default: None (anonymous)
-        new_chunks: Set manual chunksize. For development purposes only.
+        segy_spec: The SEG-Y specification to use for the conversion.
+        input_location: Store or URL (and cloud options) for MDIO file.
+        output_location: Path to the output SEG-Y file.
         selection_mask: Array that lists the subset of traces
         client: Dask client. If `None` we will use local threaded scheduler. If `auto` is used we
             will create multiple processes (with 8 threads each).
@@ -64,8 +63,7 @@ def mdio_to_segy(  # noqa: PLR0912, PLR0913
 
     Examples:
         To export an existing local MDIO file to SEG-Y we use the code snippet below. This will
-        export the full MDIO (without padding) to SEG-Y format using IBM floats and big-endian
-        byte order.
+        export the full MDIO (without padding) to SEG-Y format.
 
         >>> from mdio import mdio_to_segy
         >>>
@@ -74,51 +72,32 @@ def mdio_to_segy(  # noqa: PLR0912, PLR0913
         ...     mdio_path_or_buffer="prefix2/file.mdio",
         ...     output_segy_path="prefix/file.segy",
         ... )
-
-        If we want to export this as an IEEE big-endian, using a selection mask, we would run:
-
-        >>> mdio_to_segy(
-        ...     mdio_path_or_buffer="prefix2/file.mdio",
-        ...     output_segy_path="prefix/file.segy",
-        ...     selection_mask=boolean_mask,
-        ... )
-
     """
-    backend = "dask"
-
-    output_segy_path = Path(output_segy_path)
+    output_segy_path = Path(output_location.uri)
 
-    mdio = MDIOReader(
-        mdio_path_or_buffer=mdio_path_or_buffer,
-        access_pattern=access_pattern,
-        storage_options=storage_options,
-    )
+    mdio_xr = open_zarr_dataset(input_location)
 
-    if new_chunks is None:
-        new_chunks = segy_export_rechunker(mdio.chunks, mdio.shape, mdio._traces.dtype)
+    trace_variable_name = mdio_xr.attrs["attributes"]["traceVariableName"]
+    amplitude = mdio_xr[trace_variable_name]
+    chunks: tuple[int, ...] = amplitude.data.chunksize
+    shape: tuple[int, ...] = amplitude.data.shape
+    dtype = amplitude.dtype
+    new_chunks = segy_export_rechunker(chunks, shape, dtype)
 
-    creation_args = [
-        mdio_path_or_buffer,
-        output_segy_path,
-        access_pattern,
-        endian,
-        storage_options,
-        new_chunks,
-        backend,
-    ]
+    creation_args = [segy_spec, input_location, output_location]
 
     if client is not None:
         if distributed is not None:
             # This is in case we work with big data
             feature = client.submit(mdio_spec_to_segy, *creation_args)
-            mdio, segy_factory = feature.result()
+            mdio_xr, segy_factory = feature.result()
         else:
             msg = "Distributed client was provided, but `distributed` is not installed"
             raise ImportError(msg)
     else:
-        mdio, segy_factory = mdio_spec_to_segy(*creation_args)
+        mdio_xr, segy_factory = mdio_spec_to_segy(*creation_args)
 
-    live_mask = mdio.live_mask.compute()
+    live_mask = mdio_xr["trace_mask"].data.compute()
 
     if selection_mask is not None:
         live_mask = live_mask & selection_mask
@@ -138,8 +117,9 @@ def mdio_to_segy(  # noqa: PLR0912, PLR0913
         dim_slices += (slice(start, stop),)
 
     # Lazily pull the data with limits now, and limit mask so its the same shape.
-    live_mask, headers, samples = mdio[dim_slices]
-    live_mask = live_mask.rechunk(headers.chunks)
+    live_mask = mdio_xr["trace_mask"].data.rechunk(new_chunks[:-1])[dim_slices]
+    headers = mdio_xr["headers"].data.rechunk(new_chunks[:-1])[dim_slices]
+    samples = mdio_xr["amplitude"].data.rechunk(new_chunks)[dim_slices]
 
     if selection_mask is not None:
         selection_mask = selection_mask[dim_slices]
@@ -170,4 +150,4 @@ def mdio_to_segy(  # noqa: PLR0912, PLR0913
         if client is not None:
             _ = client.submit(concat_files, paths=ordered_files).result()
         else:
-            concat_files(paths=ordered_files, progress=True)
+            concat_files(paths=ordered_files, progress=True)

src/mdio/converters/segy.py

@@ -12,7 +12,6 @@
 from segy.standards.codes import MeasurementSystem as segy_MeasurementSystem
 from segy.standards.fields.trace import Rev0 as TraceHeaderFieldsRev0
 
-from mdio.constants import UINT32_MAX
 from mdio.converters.exceptions import EnvironmentFormatError
 from mdio.converters.exceptions import GridTraceCountError
 from mdio.converters.exceptions import GridTraceSparsityError
@@ -214,18 +213,89 @@ def populate_dim_coordinates(
     return dataset, drop_vars_delayed
 
 
-def populate_non_dim_coordinates(
+def _check_dimensions_values_identical(arr: np.ndarray, axes_to_check: tuple[int, ...]) -> bool:
+    """Check if all values along specified dimensions are identical.
+
+    Check if all values along specified dimensions are identical for each sub-array
+    defined by the other dimensions.
+
+    Args:
+        arr: an N-dimensional array. For example, an array of all 'cdp_x' segy
+            header values for coordinates "inline", "crossline", "offset", "azimuth".
+        axes_to_check: A tuple of integers representing the axes to check for
+            identical values. For example, (2, 3) would check the "offset", "azimuth"
+            dimensions.
+
+    Returns:
+        True indicates the all values in the dimensions selected by axes_to_check
+        are identical, and False otherwise.
+    """
+    # Create a slicing tuple to get the first element along the axes to check
+    full_slice = [slice(None)] * arr.ndim
+    for axis in axes_to_check:
+        full_slice[axis] = 0
+
+    # Broadcast the first element along the specified axes for comparison
+    first_element_slice = arr[tuple(full_slice)]
+
+    # Add new axes to the slice to enable broadcasting against the original array
+    for axis in axes_to_check:
+        first_element_slice = np.expand_dims(first_element_slice, axis)
+
+    # Compare the array with the broadcasted slice and use np.all()
+    # to collapse the dimensions being checked
+    identical = np.all(np.isclose(arr, first_element_slice, rtol=1e-05, atol=1e-08), axis=axes_to_check)
+    return np.all(identical).item()
+
+
+def _populate_non_dim_coordinates(
     dataset: xr_Dataset,
     grid: Grid,
-    coordinates: dict[str, SegyHeaderArray],
+    coordinate_headers: dict[str, SegyHeaderArray],
     drop_vars_delayed: list[str],
 ) -> tuple[xr_Dataset, list[str]]:
     """Populate the xarray dataset with coordinate variables."""
-    not_null = grid.map[:] != UINT32_MAX
-    for c_name, c_values in coordinates.items():
-        c_tmp_array = dataset[c_name].values
-        c_tmp_array[not_null] = c_values
-        dataset[c_name][:] = c_tmp_array
+    # Load the grid map values into memory.
+    for c_name, coord_headers_values in coordinate_headers.items():
+        headers_dims = grid.dim_names[:-1]  # e.g.: "inline", "crossline", "offset", "azimuth"
+        coord_dims = dataset[c_name].dims  # e.g.: "inline", "crossline"
+        axes_to_check = tuple(i for i, dim in enumerate(headers_dims) if dim not in coord_dims)
+        if axes_to_check == ():
+            # In case the coordinate has the same dimensions as grid map
+            not_null = grid.map[:] != grid.map.fill_value
+            tmp_coord_values = dataset[c_name].values
+            tmp_coord_values[not_null] = coord_headers_values
+        else:
+            # In the case of Coca and some other templates, the coordinate header values,
+            # coord_headers_values, have a full set of dimensions (e.g. a 4-tuple of "inline",
+            # "crossline", "offset", "azimuth"), while the non-dimensional coordinates, (e.g.,
+            # dataset["cdp_x"]) are defined over only a subset of the dimensions (e.g. 2-tuple of
+            # "inline", "crossline").
+            # Thus, every coordinate 2-tuple has multiple duplicate values of the "cdp_x" coordinates
+            # stored in coord_headers_values. Those needs to be reduced to a unique value.
+            # We assume (and check) that all the duplicate values are (near) identical.
+            #
+            # The following will create a temporary array in memory with the same shape as the
+            # coordinate defined in the dataset. Since the coordinate variable has not yet been
+            # populated, the temporary array will be populated with _FillValue from the current
+            # coordinate values.
+            tmp_coord_values = dataset[c_name].values
+            # Create slices for the all grid dimensions that are also the coordinate dimensions.
+            # For other dimension, select the first element (with index 0)
+            slices = tuple(slice(None) if name in coord_dims else 0 for name in headers_dims)
+            # Create a boolean mask for the live trace values with the dimensions of the coordinate
+            not_null = grid.map[slices] != grid.map.fill_value
+            ch_reshaped = coord_headers_values.reshape(grid.map.shape)
+            # Select a subset of the coordinate_headers that have unique values
+            # and save the unique coordinate values for live traces only
+            tmp_coord_values[not_null] = ch_reshaped[slices].ravel()
+
+            # Validate the all reduced dimensions had identical values
+            if not _check_dimensions_values_identical(ch_reshaped, axes_to_check):
+                err = f"Coordinate '{c_name}' has non-identical values along reduced dimensions."
+                raise ValueError(err)
+
+        dataset[c_name][:] = tmp_coord_values
         drop_vars_delayed.append(c_name)
     return dataset, drop_vars_delayed
 
@@ -264,15 +334,16 @@ def _populate_coordinates(
         coords: The non-dim coordinates to populate.
 
     Returns:
-        Xarray dataset with filled coordinates and updated variables to drop after writing
+        A tuple of the Xarray dataset with filled coordinates and updated variables to drop
+        after writing
     """
     drop_vars_delayed = []
     # Populate the dimension coordinate variables (1-D arrays)
-    dataset, vars_to_drop_later = populate_dim_coordinates(dataset, grid, drop_vars_delayed=drop_vars_delayed)
+    dataset, drop_vars_delayed = populate_dim_coordinates(dataset, grid, drop_vars_delayed=drop_vars_delayed)
 
     # Populate the non-dimension coordinate variables (N-dim arrays)
-    dataset, vars_to_drop_later = populate_non_dim_coordinates(
-        dataset, grid, coordinates=coords, drop_vars_delayed=drop_vars_delayed
+    dataset, drop_vars_delayed = _populate_non_dim_coordinates(
+        dataset, grid, coordinate_headers=coords, drop_vars_delayed=drop_vars_delayed
     )
 
     return dataset, drop_vars_delayed

src/mdio/core/utils_read.py

@@ -0,0 +1,16 @@
+"""Utils for reading MDIO dataset"""
+
+from typing import TYPE_CHECKING
+
+from xarray import open_dataset as xr_open_dataset
+from mdio.core.storage_location import StorageLocation
+from xarray import Dataset as xr_Dataset
+
+def open_zarr_dataset(storage_location: StorageLocation) -> xr_Dataset:
+    """
+    Open a Zarr dataset from the specified storage location.
+    """
+    # NOTE: If mask_and_scale is not set,
+    # Xarray will convert int to float and replace _FillValue with NaN
+    # We are using chunks={} to force Xarray to create Dask arrays
+    return xr_open_dataset(storage_location.uri, engine="zarr", chunks={}, mask_and_scale=False)

src/mdio/schemas/v1/templates/abstract_dataset_template.py

@@ -84,7 +84,9 @@ def build_dataset(
         self._dim_sizes = sizes
         self._horizontal_coord_unit = horizontal_coord_unit
 
-        self._builder = MDIODatasetBuilder(name=name, attributes=self._load_dataset_attributes())
+        attr = self._load_dataset_attributes() or UserAttributes(attributes={})
+        attr.attributes["traceVariableName"] = self._trace_variable_name
+        self._builder = MDIODatasetBuilder(name=name, attributes=attr)
         self._add_dimensions()
         self._add_coordinates()
         self._add_variables()

src/mdio/schemas/v1/templates/seismic_3d_prestack_coca.py

@@ -12,8 +12,8 @@ class Seismic3DPreStackCocaTemplate(AbstractDatasetTemplate):
     def __init__(self, domain: str):
         super().__init__(domain=domain)
 
-        self._coord_dim_names = ["inline", "crossline", "offset", "azimuth"]
-        self._dim_names = [*self._coord_dim_names, self._trace_domain]
+        self._coord_dim_names = ["inline", "crossline"]
+        self._dim_names = [*self._coord_dim_names, "offset", "azimuth", self._trace_domain]
         self._coord_names = ["cdp_x", "cdp_y"]
         self._var_chunk_shape = [8, 8, 32, 1, 1024]
 
@@ -64,13 +64,13 @@ def _add_coordinates(self) -> None:
         # Add non-dimension coordinates
         self._builder.add_coordinate(
             "cdp_x",
-            dimensions=["inline", "crossline"],
+            dimensions=self._coord_dim_names,
             data_type=ScalarType.FLOAT64,
             metadata_info=[self._horizontal_coord_unit],
         )
         self._builder.add_coordinate(
             "cdp_y",
-            dimensions=["inline", "crossline"],
+            dimensions=self._coord_dim_names,
             data_type=ScalarType.FLOAT64,
             metadata_info=[self._horizontal_coord_unit],
         )