Update to TGSAI#529 as main

Author: BrianMichell

docs/index.md

@@ -17,6 +17,7 @@ maxdepth: 1
 ---
 installation
 notebooks/quickstart
+notebooks/creation
 notebooks/compression
 notebooks/rechunking
 usage

docs/notebooks/creation.ipynb

@@ -7,6 +7,12 @@
 from mdio.api.convenience import copy_mdio
 from mdio.converters import mdio_to_segy
 from mdio.converters import segy_to_mdio
+from mdio.core.dimension import Dimension
+from mdio.core.factory import MDIOCreateConfig
+from mdio.core.factory import MDIOVariableConfig
+from mdio.core.factory import create_empty
+from mdio.core.factory import create_empty_like
+from mdio.core.grid import Grid
 
 
 __all__ = [
@@ -15,6 +21,12 @@
     "copy_mdio",
     "mdio_to_segy",
     "segy_to_mdio",
+    "Dimension",
+    "MDIOCreateConfig",
+    "MDIOVariableConfig",
+    "create_empty",
+    "create_empty_like",
+    "Grid",
 ]
 
 

src/mdio/__init__.py

@@ -129,6 +129,8 @@ class MDIOAccessor:
         we can directly unpack it and use it further down our code.
     """
 
+    _stats_keys = {"mean", "std", "rms", "min", "max"}
+
     _array_load_function_mapper = {
         "zarr": open_zarr_array,
         "dask": open_zarr_array_dask,
@@ -154,16 +156,13 @@ def __init__(
 
         # Set private attributes for public interface.
         # Pep8 complains because they are defined outside __init__
-        self._binary_header = None
         self._chunks = None
         self._live_mask = None
         self._root = None
         self._n_dim = None
         self._orig_chunks = None
         self._store = None
         self._shape = None
-        self._stats = None
-        self._text_header = None
         self._trace_count = None
 
         # Private attributes
@@ -212,19 +211,17 @@ def _connect(self):
             )
             raise MDIONotFoundError(msg) from e
 
+    def _consolidate_metadata(self) -> None:
+        """Flush optimized MDIO metadata, run after modifying it."""
+        zarr.consolidate_metadata(self.root.store)
+
     def _deserialize_grid(self):
         """Deserialize grid from Zarr metadata."""
         self.grid = Grid.from_zarr(self.root)
 
     def _set_attributes(self):
         """Deserialize attributes from Zarr metadata."""
         self.trace_count = self.root.attrs["trace_count"]
-        self.stats = {
-            key: self.root.attrs[key] for key in ["mean", "std", "rms", "min", "max"]
-        }
-
-        self.text_header = self._metadata_group.attrs["text_header"]
-        self.binary_header = self._metadata_group.attrs["binary_header"]
 
         # Grid based attributes
         self.shape = self.grid.shape
@@ -331,26 +328,28 @@ def trace_count(self, value: int) -> None:
     @property
     def text_header(self) -> list:
         """Get seismic text header."""
-        return self._text_header
+        return self._metadata_group.attrs["text_header"]
 
     @text_header.setter
     def text_header(self, value: list) -> None:
         """Validate and set seismic text header."""
-        if not isinstance(value, list):
+        if not isinstance(value, list) or len(value) != 40:
             raise AttributeError("Text header must be a list of str with 40 elements")
-        self._text_header = value
+        self._metadata_group.attrs["text_header"] = value
+        self._consolidate_metadata()
 
     @property
     def binary_header(self) -> dict:
         """Get seismic binary header metadata."""
-        return self._binary_header
+        return self._metadata_group.attrs["binary_header"]
 
     @binary_header.setter
     def binary_header(self, value: dict) -> None:
         """Validate and set seismic binary header metadata."""
         if not isinstance(value, dict):
             raise AttributeError("Binary header has to be a dictionary type collection")
-        self._binary_header = value
+        self._metadata_group.attrs["binary_header"] = value
+        self._consolidate_metadata()
 
     @property
     def chunks(self) -> tuple[int, ...]:
@@ -365,12 +364,16 @@ def chunks(self, value: tuple[int, ...]) -> None:
     @property
     def stats(self) -> dict:
         """Get global statistics like min/max/rms/std."""
-        return self._stats
+        return {key: self.root.attrs[key] for key in self._stats_keys}
 
     @stats.setter
     def stats(self, value: dict) -> None:
         """Set global statistics like min/max/rms/std."""
-        self._stats = value
+        if not isinstance(value, dict) or not self._stats_keys.issubset(value.keys()):
+            msg = f"For settings status, you must provide keys: {self._stats_keys}"
+            raise AttributeError(msg)
+        self.root.attrs.update(value)
+        self._consolidate_metadata()
 
     @property
     def _metadata_group(self) -> zarr.Group:
@@ -403,6 +406,7 @@ def __getitem__(self, item: int | tuple) -> npt.ArrayLike | da.Array | tuple:
     def __setitem__(self, key: int | tuple, value: npt.ArrayLike) -> None:
         """Data setter."""
         self._traces[key] = value
+        self._live_mask[key] = True
 
     def coord_to_index(
         self,
@@ -643,10 +647,10 @@ def __init__(
         memory_cache_size: int = 0,
         disk_cache: bool = False,
     ):  # TODO: Disabled all caching by default, sometimes causes performance issues
-        """Initialize super class with `r+` permission."""
+        """Initialize accessor class with `w` permission."""
         super().__init__(
             mdio_path_or_buffer=mdio_path_or_buffer,
-            mode="r+",
+            mode="w",
             access_pattern=access_pattern,
             storage_options=storage_options,
             return_metadata=return_metadata,

src/mdio/api/accessor.py

@@ -5,8 +5,6 @@
 import logging
 import os
 from collections.abc import Sequence
-from datetime import datetime
-from datetime import timezone
 from importlib import metadata
 from typing import Any
 
@@ -15,20 +13,30 @@
 from segy import SegyFile
 from segy.config import SegySettings
 from segy.schema import HeaderField
+from zarr import Blosc
 
 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
-from mdio.core.grid import _calculate_live_mask_chunksize
+from mdio.core.factory import MDIOCreateConfig
+from mdio.core.factory import MDIOVariableConfig
+from mdio.core.factory import create_empty
 from mdio.core.utils_write import write_attribute
 from mdio.segy import blocked_io
 from mdio.segy.compat import mdio_segy_spec
-from mdio.segy.helpers_segy import create_zarr_hierarchy
 from mdio.segy.utilities import get_grid_plan
 
 
+try:
+    import zfpy  # Base library
+    from zarr import ZFPY  # Codec
+except ImportError:
+    ZFPY = None
+    zfpy = None
+
+
 logger = logging.getLogger(__name__)
 
 try:
@@ -104,6 +112,28 @@ def grid_density_qc(grid: Grid, num_traces: int) -> None:
             raise GridTraceSparsityError(grid.shape, num_traces, msg)
 
 
+def get_compressor(
+    lossless: bool, compression_tolerance: float = -1
+) -> Blosc | ZFPY | None:
+    """Get the appropriate compressor for the seismic traces."""
+    if lossless:
+        compressor = Blosc("zstd")
+    else:
+        if zfpy is None or ZFPY is None:
+            msg = (
+                "Lossy compression requires the 'zfpy' library. It is "
+                "not installed in your environment. To proceed please "
+                "install 'zfpy' or install mdio with `--extras lossy`"
+            )
+            raise ImportError(msg)
+
+        compressor = ZFPY(
+            mode=zfpy.mode_fixed_accuracy,
+            tolerance=compression_tolerance,
+        )
+    return compressor
+
+
 def segy_to_mdio(  # noqa: C901
     segy_path: str,
     mdio_path_or_buffer: str,
@@ -365,14 +395,6 @@ def segy_to_mdio(  # noqa: C901
     if storage_options_output is None:
         storage_options_output = {}
 
-    store = process_url(
-        url=mdio_path_or_buffer,
-        mode="w",
-        storage_options=storage_options_output,
-        memory_cache_size=0,  # Making sure disk caching is disabled,
-        disk_cache=False,  # Making sure disk caching is disabled
-    )
-
     # Open SEG-Y with MDIO's SegySpec. Endianness will be inferred.
     mdio_spec = mdio_segy_spec()
     segy_settings = SegySettings(storage_options=storage_options_input)
@@ -407,47 +429,6 @@ def segy_to_mdio(  # noqa: C901
         logger.warning(f"Ingestion grid shape: {grid.shape}.")
         raise GridTraceCountError(np.sum(grid.live_mask), num_traces)
 
-    zarr_root = create_zarr_hierarchy(
-        store=store,
-        overwrite=overwrite,
-    )
-
-    # Get UTC time, then add local timezone information offset.
-    iso_datetime = datetime.now(timezone.utc).isoformat()
-
-    write_attribute(name="created", zarr_group=zarr_root, attribute=iso_datetime)
-    write_attribute(name="api_version", zarr_group=zarr_root, attribute=API_VERSION)
-
-    dimensions_dict = [dim.to_dict() for dim in dimensions]
-    write_attribute(name="dimension", zarr_group=zarr_root, attribute=dimensions_dict)
-
-    # Write trace count
-    trace_count = np.count_nonzero(grid.live_mask)
-    write_attribute(name="trace_count", zarr_group=zarr_root, attribute=trace_count)
-
-    live_mask_chunksize = _calculate_live_mask_chunksize(grid)
-
-    # Note, live mask is not chunked since it's bool and small.
-    zarr_root["metadata"].create_dataset(
-        data=grid.live_mask,
-        name="live_mask",
-        shape=grid.shape[:-1],
-        chunks=live_mask_chunksize,
-        dimension_separator="/",
-    )
-
-    write_attribute(
-        name="text_header",
-        zarr_group=zarr_root["metadata"],
-        attribute=text_header.split("\n"),
-    )
-
-    write_attribute(
-        name="binary_header",
-        zarr_group=zarr_root["metadata"],
-        attribute=binary_header.to_dict(),
-    )
-
     if chunksize is None:
         dim_count = len(index_names) + 1
         if dim_count == 2:
@@ -470,18 +451,59 @@ def segy_to_mdio(  # noqa: C901
         suffix = [str(idx) for idx, value in enumerate(suffix) if value is not None]
         suffix = "".join(suffix)
 
+    compressor = get_compressor(compression_tolerance, lossless)
+    header_dtype = segy.spec.trace.header.dtype.newbyteorder("=")
+    var_conf = MDIOVariableConfig(
+        name=f"chunked_{suffix}",
+        dtype="float32",
+        chunks=chunksize,
+        compressor=compressor,
+        header_dtype=header_dtype,
+    )
+    config = MDIOCreateConfig(path=mdio_path_or_buffer, grid=grid, variables=[var_conf])
+
+    zarr_root = create_empty(
+        config,
+        overwrite=overwrite,
+        storage_options=storage_options_output,
+        consolidate_meta=False,
+    )
+    data_group, meta_group = zarr_root["data"], zarr_root["metadata"]
+    data_array = data_group[f"chunked_{suffix}"]
+    header_array = meta_group[f"chunked_{suffix}_trace_headers"]
+
+    # Write actual live mask and metadata to empty MDIO
+    meta_group["live_mask"][:] = grid.live_mask
+    write_attribute(
+        name="trace_count",
+        zarr_group=zarr_root,
+        attribute=np.count_nonzero(grid.live_mask),
+    )
+    write_attribute(
+        name="text_header",
+        zarr_group=zarr_root["metadata"],
+        attribute=text_header.split("\n"),
+    )
+    write_attribute(
+        name="binary_header",
+        zarr_group=zarr_root["metadata"],
+        attribute=binary_header.to_dict(),
+    )
+
+    # Write traces
     stats = blocked_io.to_zarr(
         segy_file=segy,
         grid=grid,
-        data_root=zarr_root["data"],
-        metadata_root=zarr_root["metadata"],
+        data_array=data_array,
+        header_array=header_array,
         name="_".join(["chunked", suffix]),
         dtype="float32",
         chunks=chunksize,
         lossless=lossless,
         compression_tolerance=compression_tolerance,
     )
 
+    # Write actual stats
     for key, value in stats.items():
         write_attribute(name=key, zarr_group=zarr_root, attribute=value)