<xarray.MDIODataset> Size: 203MB\n",
+ "Dimensions: (inline: 256, crossline: 512, depth: 384)\n",
+ "Coordinates:\n",
+ " * inline (inline) int32 1kB 0 0 0 0 0 0 0 0 0 ... 0 0 0 0 0 0 0 0 0\n",
+ " * crossline (crossline) int32 2kB 0 0 0 0 0 0 0 0 0 ... 0 0 0 0 0 0 0 0\n",
+ " * depth (depth) int32 2kB 0 0 0 0 0 0 0 0 0 0 ... 0 0 0 0 0 0 0 0 0\n",
+ "Data variables:\n",
+ " stack_amplitude (inline, crossline, depth) float32 201MB 0.0 0.0 ... 0.0\n",
+ " cdp_x (inline, crossline) float64 1MB 0.0 0.0 0.0 ... 0.0 0.0 0.0\n",
+ " cdp_y (inline, crossline) float64 1MB 0.0 0.0 0.0 ... 0.0 0.0 0.0\n",
+ "Attributes:\n",
+ " apiVersion: 1.0.0\n",
+ " createdOn: 2025-06-02 13:40:42.724997+00:00\n",
+ " name: psdm_stack_example\n",
+ " attributes: {'description': 'Example PSDM stack'}{\n", + " 'metadata': {\n", + " 'name': 'psdm_stack_example',\n", + " 'apiVersion': '1.0.0',\n", + " 'createdOn': '2025-06-02T13:40:42.724997Z',\n", + " 'attributes': {'description': 'Example PSDM stack'}\n", + " },\n", + " 'variables': [\n", + " {\n", + " 'dataType': 'int32',\n", + " 'dimensions': [{'name': 'inline', 'size': 256}],\n", + " 'name': 'inline',\n", + " 'longName': 'Inline Number'\n", + " },\n", + " {\n", + " 'dataType': 'int32',\n", + " 'dimensions': [{'name': 'crossline', 'size': 512}],\n", + " 'name': 'crossline',\n", + " 'longName': 'Crossline Number'\n", + " },\n", + " {\n", + " 'dataType': 'int32',\n", + " 'dimensions': [{'name': 'depth', 'size': 384}],\n", + " 'name': 'depth',\n", + " 'longName': 'Depth Sample'\n", + " },\n", + " {\n", + " 'dataType': 'float32',\n", + " 'dimensions': [\n", + " {'name': 'inline', 'size': 256},\n", + " {'name': 'crossline', 'size': 512},\n", + " {'name': 'depth', 'size': 384}\n", + " ],\n", + " 'compressor': {'name': 'blosc', 'algorithm': 'zstd', 'level': 3, 'shuffle': 1, 'blocksize': 0},\n", + " 'name': 'stack_amplitude',\n", + " 'coordinates': ['inline', 'crossline', 'cdp_x', 'cdp_y'],\n", + " 'metadata': {'chunkGrid': {'name': 'regular', 'configuration': {'chunkShape': [64, 64, 64]}}}\n", + " },\n", + " {\n", + " 'dataType': 'float64',\n", + " 'dimensions': [{'name': 'inline', 'size': 256}, {'name': 'crossline', 'size': 512}],\n", + " 'name': 'cdp_x',\n", + " 'longName': 'CDP X (UTM Easting)',\n", + " 'metadata': {'unitsV1': {'length': 'm'}, 'attributes': {'attributes': {'MGA': 51}}}\n", + " },\n", + " {\n", + " 'dataType': 'float64',\n", + " 'dimensions': [{'name': 'inline', 'size': 256}, {'name': 'crossline', 'size': 512}],\n", + " 'name': 'cdp_y',\n", + " 'longName': 'CDP Y (UTM Northing)',\n", + " 'metadata': {'unitsV1': {'length': 'm'}, 'attributes': {'attributes': {'MGA': 51}}}\n", + " }\n", + " ]\n", + "}\n", + "\n" + ], + "text/plain": [ + "\u001b[1m{\u001b[0m\n", + " \u001b[32m'metadata'\u001b[0m: \u001b[1m{\u001b[0m\n", + " \u001b[32m'name'\u001b[0m: \u001b[32m'psdm_stack_example'\u001b[0m,\n", + " \u001b[32m'apiVersion'\u001b[0m: \u001b[32m'1.0.0'\u001b[0m,\n", + " \u001b[32m'createdOn'\u001b[0m: \u001b[32m'2025-06-02T13:40:42.724997Z'\u001b[0m,\n", + " \u001b[32m'attributes'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'description'\u001b[0m: \u001b[32m'Example PSDM stack'\u001b[0m\u001b[1m}\u001b[0m\n", + " \u001b[1m}\u001b[0m,\n", + " \u001b[32m'variables'\u001b[0m: \u001b[1m[\u001b[0m\n", + " \u001b[1m{\u001b[0m\n", + " \u001b[32m'dataType'\u001b[0m: \u001b[32m'int32'\u001b[0m,\n", + " \u001b[32m'dimensions'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m{\u001b[0m\u001b[32m'name'\u001b[0m: \u001b[32m'inline'\u001b[0m, \u001b[32m'size'\u001b[0m: \u001b[1;36m256\u001b[0m\u001b[1m}\u001b[0m\u001b[1m]\u001b[0m,\n", + " \u001b[32m'name'\u001b[0m: \u001b[32m'inline'\u001b[0m,\n", + " \u001b[32m'longName'\u001b[0m: \u001b[32m'Inline Number'\u001b[0m\n", + " \u001b[1m}\u001b[0m,\n", + " \u001b[1m{\u001b[0m\n", + " \u001b[32m'dataType'\u001b[0m: \u001b[32m'int32'\u001b[0m,\n", + " \u001b[32m'dimensions'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m{\u001b[0m\u001b[32m'name'\u001b[0m: \u001b[32m'crossline'\u001b[0m, \u001b[32m'size'\u001b[0m: \u001b[1;36m512\u001b[0m\u001b[1m}\u001b[0m\u001b[1m]\u001b[0m,\n", + " \u001b[32m'name'\u001b[0m: \u001b[32m'crossline'\u001b[0m,\n", + " \u001b[32m'longName'\u001b[0m: \u001b[32m'Crossline Number'\u001b[0m\n", + " \u001b[1m}\u001b[0m,\n", + " \u001b[1m{\u001b[0m\n", + " \u001b[32m'dataType'\u001b[0m: \u001b[32m'int32'\u001b[0m,\n", + " \u001b[32m'dimensions'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m{\u001b[0m\u001b[32m'name'\u001b[0m: \u001b[32m'depth'\u001b[0m, \u001b[32m'size'\u001b[0m: \u001b[1;36m384\u001b[0m\u001b[1m}\u001b[0m\u001b[1m]\u001b[0m,\n", + " \u001b[32m'name'\u001b[0m: \u001b[32m'depth'\u001b[0m,\n", + " \u001b[32m'longName'\u001b[0m: \u001b[32m'Depth Sample'\u001b[0m\n", + " \u001b[1m}\u001b[0m,\n", + " \u001b[1m{\u001b[0m\n", + " \u001b[32m'dataType'\u001b[0m: \u001b[32m'float32'\u001b[0m,\n", + " \u001b[32m'dimensions'\u001b[0m: \u001b[1m[\u001b[0m\n", + " \u001b[1m{\u001b[0m\u001b[32m'name'\u001b[0m: \u001b[32m'inline'\u001b[0m, \u001b[32m'size'\u001b[0m: \u001b[1;36m256\u001b[0m\u001b[1m}\u001b[0m,\n", + " \u001b[1m{\u001b[0m\u001b[32m'name'\u001b[0m: \u001b[32m'crossline'\u001b[0m, \u001b[32m'size'\u001b[0m: \u001b[1;36m512\u001b[0m\u001b[1m}\u001b[0m,\n", + " \u001b[1m{\u001b[0m\u001b[32m'name'\u001b[0m: \u001b[32m'depth'\u001b[0m, \u001b[32m'size'\u001b[0m: \u001b[1;36m384\u001b[0m\u001b[1m}\u001b[0m\n", + " \u001b[1m]\u001b[0m,\n", + " \u001b[32m'compressor'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'name'\u001b[0m: \u001b[32m'blosc'\u001b[0m, \u001b[32m'algorithm'\u001b[0m: \u001b[32m'zstd'\u001b[0m, \u001b[32m'level'\u001b[0m: \u001b[1;36m3\u001b[0m, \u001b[32m'shuffle'\u001b[0m: \u001b[1;36m1\u001b[0m, \u001b[32m'blocksize'\u001b[0m: \u001b[1;36m0\u001b[0m\u001b[1m}\u001b[0m,\n", + " \u001b[32m'name'\u001b[0m: \u001b[32m'stack_amplitude'\u001b[0m,\n", + " \u001b[32m'coordinates'\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m'inline'\u001b[0m, \u001b[32m'crossline'\u001b[0m, \u001b[32m'cdp_x'\u001b[0m, \u001b[32m'cdp_y'\u001b[0m\u001b[1m]\u001b[0m,\n", + " \u001b[32m'metadata'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'chunkGrid'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'name'\u001b[0m: \u001b[32m'regular'\u001b[0m, \u001b[32m'configuration'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'chunkShape'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1;36m64\u001b[0m, \u001b[1;36m64\u001b[0m, \u001b[1;36m64\u001b[0m\u001b[1m]\u001b[0m\u001b[1m}\u001b[0m\u001b[1m}\u001b[0m\u001b[1m}\u001b[0m\n", + " \u001b[1m}\u001b[0m,\n", + " \u001b[1m{\u001b[0m\n", + " \u001b[32m'dataType'\u001b[0m: \u001b[32m'float64'\u001b[0m,\n", + " \u001b[32m'dimensions'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m{\u001b[0m\u001b[32m'name'\u001b[0m: \u001b[32m'inline'\u001b[0m, \u001b[32m'size'\u001b[0m: \u001b[1;36m256\u001b[0m\u001b[1m}\u001b[0m, \u001b[1m{\u001b[0m\u001b[32m'name'\u001b[0m: \u001b[32m'crossline'\u001b[0m, \u001b[32m'size'\u001b[0m: \u001b[1;36m512\u001b[0m\u001b[1m}\u001b[0m\u001b[1m]\u001b[0m,\n", + " \u001b[32m'name'\u001b[0m: \u001b[32m'cdp_x'\u001b[0m,\n", + " \u001b[32m'longName'\u001b[0m: \u001b[32m'CDP X \u001b[0m\u001b[32m(\u001b[0m\u001b[32mUTM Easting\u001b[0m\u001b[32m)\u001b[0m\u001b[32m'\u001b[0m,\n", + " \u001b[32m'metadata'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'unitsV1'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'length'\u001b[0m: \u001b[32m'm'\u001b[0m\u001b[1m}\u001b[0m, \u001b[32m'attributes'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'attributes'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'MGA'\u001b[0m: \u001b[1;36m51\u001b[0m\u001b[1m}\u001b[0m\u001b[1m}\u001b[0m\u001b[1m}\u001b[0m\n", + " \u001b[1m}\u001b[0m,\n", + " \u001b[1m{\u001b[0m\n", + " \u001b[32m'dataType'\u001b[0m: \u001b[32m'float64'\u001b[0m,\n", + " \u001b[32m'dimensions'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m{\u001b[0m\u001b[32m'name'\u001b[0m: \u001b[32m'inline'\u001b[0m, \u001b[32m'size'\u001b[0m: \u001b[1;36m256\u001b[0m\u001b[1m}\u001b[0m, \u001b[1m{\u001b[0m\u001b[32m'name'\u001b[0m: \u001b[32m'crossline'\u001b[0m, \u001b[32m'size'\u001b[0m: \u001b[1;36m512\u001b[0m\u001b[1m}\u001b[0m\u001b[1m]\u001b[0m,\n", + " \u001b[32m'name'\u001b[0m: \u001b[32m'cdp_y'\u001b[0m,\n", + " \u001b[32m'longName'\u001b[0m: \u001b[32m'CDP Y \u001b[0m\u001b[32m(\u001b[0m\u001b[32mUTM Northing\u001b[0m\u001b[32m)\u001b[0m\u001b[32m'\u001b[0m,\n", + " \u001b[32m'metadata'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'unitsV1'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'length'\u001b[0m: \u001b[32m'm'\u001b[0m\u001b[1m}\u001b[0m, \u001b[32m'attributes'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'attributes'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'MGA'\u001b[0m: \u001b[1;36m51\u001b[0m\u001b[1m}\u001b[0m\u001b[1m}\u001b[0m\u001b[1m}\u001b[0m\n", + " \u001b[1m}\u001b[0m\n", + " \u001b[1m]\u001b[0m\n", + "\u001b[1m}\u001b[0m\n" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "# Build our Dataset model from the builder\n", + "dataset = builder.build()\n", + "\n", + "# Serialize the Dataset model to JSON\n", + "contract = json.loads(dataset.json())\n", + "\n", + "# Reorder the contract so that metadata is displayed first\n", + "ordered_contract = {\n", + " \"metadata\": contract[\"metadata\"],\n", + " \"variables\": contract[\"variables\"],\n", + "}\n", + "\n", + "rprint(ordered_contract)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "jupyter", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/tutorials/builder.md b/docs/tutorials/builder.md new file mode 100644 index 00000000..9fddf377 --- /dev/null +++ b/docs/tutorials/builder.md @@ -0,0 +1,146 @@ +# Constructing a v1 Dataset with the MDIODatasetBuilder + +In this tutorial, we'll walk through how to use the `MDIODatasetBuilder` class to programmatically construct an MDIO v1 dataset. The builder enforces a specific build order to ensure a valid dataset: + +1. Add dimensions via `add_dimension()` +2. (Optional) Add coordinates via `add_coordinate()` +3. Add variables via `add_variable()` +4. Call `build()` to finalize the dataset. + +## Importing the Builder + +```python +from mdio.core.v1.builder import MDIODatasetBuilder, write_mdio_metadata +from mdio.schemas.dtype import ScalarType, StructuredType +from mdio.schemas.compressors import Blosc, ZFP +``` + +## Creating the Builder + +First, create a builder instance with a name and optional global attributes: + +```python +builder = MDIODatasetBuilder( + name="example_dataset", + attributes={ + "description": "An example MDIO v1 dataset", + "creator": "Your Name", + }, +) +``` + +## Adding Dimensions + +Dimensions define the axes of your dataset. You must add at least one dimension before adding coordinates or variables: + +```python +builder = ( + builder + .add_dimension(name="inline", size=256, long_name="Inline Number") + .add_dimension(name="crossline", size=512, long_name="Crossline Number") + .add_dimension(name="depth", size=384, long_name="Depth Sample") +) +``` + +## Adding Coordinates (Optional) + +Coordinates map grid indices to real-world positions (e.g., UTM coordinates on the inline–crossline plane): + +```python +builder = ( + builder + .add_coordinate( + name="cdp_x", + dimensions=["inline", "crossline"], + long_name="CDP X (UTM Easting)", + data_type=ScalarType.FLOAT64, + metadata={"unitsV1": {"length": "m"}}, + ) + .add_coordinate( + name="cdp_y", + dimensions=["inline", "crossline"], + long_name="CDP Y (UTM Northing)", + data_type=ScalarType.FLOAT64, + metadata={"unitsV1": {"length": "m"}}, + ) +) +``` + +If you omit `name`, the builder auto-generates names like `coord_0`. If you omit `dimensions`, it uses all defined dimensions. + +## Adding Variables + +Add one or more seismic data variables (e.g., post-stack amplitude volumes). Variables can have compressors, statistics, and more: + +```python +builder = builder.add_variable( + name="stack_amplitude", + dimensions=["inline", "crossline", "depth"], + data_type=ScalarType.FLOAT32, + compressor=Blosc(algorithm="zstd", level=3), + coordinates=["inline", "crossline", "cdp_x", "cdp_y"], + metadata={ + "chunkGrid": {"name": "regular", "configuration": {"chunkShape": [64, 64, 64]}} + }, +) +``` + +For structured dtypes, use `StructuredType`: + +```python +from mdio.schemas.dtype import StructuredType, ScalarType + +structured_dtype = StructuredType( + fields=[ + {"name": "flag", "format": ScalarType.INT8}, + {"name": "value", "format": ScalarType.FLOAT32}, + ] +) + +builder = builder.add_variable( + name="metadata", + dimensions=["x", "y"], + data_type=structured_dtype, +) +``` + +## Building the Dataset + +After adding all components, call: + +```python +dataset = builder.build() +``` + +This returns a `Dataset` object conforming to the MDIO v1 schema. + +## Writing Metadata and Writing Data + +The `.build()` method returns an in-memory Pydantic `Dataset` model (MDIO v1 schema). To serialize this model to disk, use the following approaches: + +- **Metadata only** (no array values written): + + ```python + # Write metadata structure only (no data arrays) + mds = write_mdio_metadata( + dataset, + store="path/to/output.mdio" + ) + ``` + + This writes only the metadata to the `.mdio` store and returns an `mdio.Dataset` (an xarray.Dataset subclass) with placeholder arrays. + +- **Write actual data** (array values): + + After writing metadata, call `to_mdio()` on the returned `mdio.Dataset` with `compute=True` to write the actual data arrays: + + ```python + # Write data arrays into the existing store + mds.to_mdio( + store="path/to/output.mdio", + mode="a", + compute=True, + ) + ``` + + Alternatively, skip `write_mdio_metadata()` and write both metadata and data in one call by invoking `to_mdio()` directly on the `mdio.Dataset` produced by `_construct_mdio_dataset`, if you have it available. diff --git a/noxfile.py b/noxfile.py index 281ec525..421cf931 100644 --- a/noxfile.py +++ b/noxfile.py @@ -200,6 +200,7 @@ def tests(session: Session) -> None: "pygments", "pytest-dependency", "s3fs", + "zfpy", # TODO(BrianMichell): #0 Ensure this is pulling from the pyproject.toml ], ) diff --git a/pyproject.toml b/pyproject.toml index 7e01894e..d5605839 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,8 +34,8 @@ dependencies = [ "rich (>=13.9.4,<14.0.0)", "segy (>=0.4.0,<0.5.0)", "tqdm (>=4.67.0,<5.0.0)", - "xarray>=2025.3.1", "zarr (>=3.0.8,<4.0.0)", + "xarray (>=2025.4.0)", ] [project.optional-dependencies] diff --git a/src/mdio/core/v1/__init__.py b/src/mdio/core/v1/__init__.py new file mode 100644 index 00000000..d8ec5b85 --- /dev/null +++ b/src/mdio/core/v1/__init__.py @@ -0,0 +1,26 @@ +"""MDIO core v1 package initialization. + +Exposes the MDIO overloads and core v1 functionality. +""" + +from ._overloads import mdio +from ._serializer import make_coordinate +from ._serializer import make_dataset +from ._serializer import make_dataset_metadata +from ._serializer import make_named_dimension +from ._serializer import make_variable +from .builder import MDIODatasetBuilder +from .factory import SCHEMA_TEMPLATE_MAP +from .factory import MDIOSchemaType + +__all__ = [ + "MDIODatasetBuilder", + "make_coordinate", + "make_dataset", + "make_dataset_metadata", + "make_named_dimension", + "make_variable", + "mdio", + "MDIOSchemaType", + "SCHEMA_TEMPLATE_MAP", +] diff --git a/src/mdio/core/v1/_overloads.py b/src/mdio/core/v1/_overloads.py new file mode 100644 index 00000000..a07353b1 --- /dev/null +++ b/src/mdio/core/v1/_overloads.py @@ -0,0 +1,93 @@ +"""Overloads for xarray. + +The intent of overloading here is: +1. To provide a consistent mdio.* naming scheme. +""" + +from collections.abc import Mapping + +import xarray as xr +from xarray import DataArray as _DataArray +from xarray import Dataset as _Dataset + + +class MDIODataset(_Dataset): + """xarray.Dataset subclass with MDIO v1 extensions.""" + + __slots__ = () + + def to_mdio( + self, + store: str | None = None, + *args: str | int | float | bool, + **kwargs: Mapping[str, str | int | float | bool], + ) -> None: + """Alias for `.to_zarr()`.""" + # Ensure zarr_format=2 by default unless explicitly overridden + zarr_format = kwargs.get("zarr_format", 2) + if zarr_format != 2: # noqa: PLR2004 + msg = "MDIO only supports zarr_format=2" + raise ValueError(msg) + kwargs["zarr_format"] = zarr_format + return super().to_zarr(*args, store=store, **kwargs) + + +class MDIODataArray(_DataArray): + """xarray.DataArray subclass with MDIO v1 extensions.""" + + __slots__ = () + + def to_mdio( + self, + store: str | None = None, + *args: str | int | float | bool, + **kwargs: Mapping[str, str | int | float | bool], + ) -> None: + """Alias for `.to_zarr()`, and writes to Zarr store.""" + # Ensure zarr_format=2 by default unless explicitly overridden + zarr_format = kwargs.get("zarr_format", 2) + if zarr_format != 2: # noqa: PLR2004 + msg = "MDIO only supports zarr_format=2" + raise ValueError(msg) + kwargs["zarr_format"] = zarr_format + return super().to_zarr(*args, store=store, **kwargs) + + +class MDIO: + """MDIO namespace for overloaded types and functions.""" + + Dataset = MDIODataset + DataArray = MDIODataArray + + @staticmethod + def open( + store: str, + *args: str | int | float | bool, + engine: str = "zarr", + consolidated: bool = False, + **kwargs: Mapping[str, str | int | float | bool], + ) -> MDIODataset: + """Open a Zarr store as an MDIODataset. + + Casts the returned xarray.Dataset (and its variables) to the MDIO subclasses. + """ + ds = xr.open_dataset( + store, + *args, + engine=engine, + consolidated=consolidated, + **kwargs, + ) + # Cast Dataset to MDIODataset + ds.__class__ = MDIODataset + # Cast each DataArray in data_vars and coords + + for _name, var in ds.data_vars.items(): # noqa: PERF102 .values() failed tests + var.__class__ = MDIODataArray + for _name, coord in ds.coords.items(): # noqa: PERF102 .values() failed tests + coord.__class__ = MDIODataArray + return ds + + +# Create module-level MDIO namespace +mdio = MDIO() diff --git a/src/mdio/core/v1/_serializer.py b/src/mdio/core/v1/_serializer.py new file mode 100644 index 00000000..d7c6656d --- /dev/null +++ b/src/mdio/core/v1/_serializer.py @@ -0,0 +1,293 @@ +"""Internal serialization module for MDIO v1 datasets. + +This module contains internal implementation details for serializing MDIO schema models +to Zarr storage. This API is not considered stable and may change without notice. +""" + +from datetime import datetime +from typing import Any + +import numpy as np +from numcodecs import Blosc as NumcodecsBlosc + +from mdio.core.v1._overloads import mdio +from mdio.schemas.chunk_grid import * # noqa: F403 +from mdio.schemas.compressors import ZFP +from mdio.schemas.compressors import Blosc +from mdio.schemas.dimension import NamedDimension +from mdio.schemas.dtype import ScalarType +from mdio.schemas.dtype import StructuredType +from mdio.schemas.metadata import UserAttributes +from mdio.schemas.v1.dataset import Dataset as MDIODataset +from mdio.schemas.v1.dataset import DatasetMetadata +from mdio.schemas.v1.stats import * # noqa: F403 +from mdio.schemas.v1.units import AllUnits +from mdio.schemas.v1.variable import Coordinate +from mdio.schemas.v1.variable import Variable +from mdio.schemas.v1.variable import VariableMetadata + +try: + import zfpy as zfpy_base # Base library + from numcodecs import ZFPY # Codec +except ImportError: + zfpy_base = None + ZFPY = None + + +def make_named_dimension(name: str, size: int) -> NamedDimension: + """Create a NamedDimension with the given name and size.""" + return NamedDimension(name=name, size=size) + + +def make_coordinate( + name: str, + dimensions: list[NamedDimension | str], + data_type: ScalarType | StructuredType, + long_name: str = None, + metadata: list[AllUnits | UserAttributes] | dict[str, Any] | None = None, +) -> Coordinate: + """Create a Coordinate with the given name, dimensions, data_type, and metadata.""" + # Build metadata list of AllUnits or UserAttributes to satisfy Coordinate.schema + coord_meta_list: list[AllUnits | UserAttributes] | None = None + if metadata is not None: + items: list[AllUnits | UserAttributes] = [] + # single dict input + if isinstance(metadata, dict): + if "unitsV1" in metadata: + items.append(AllUnits(unitsV1=metadata["unitsV1"])) + if "attributes" in metadata: + items.append(UserAttributes(attributes=metadata["attributes"])) + # list input may contain dict or model instances + elif isinstance(metadata, list): + for md in metadata: + if isinstance(md, AllUnits) or isinstance(md, UserAttributes): # noqa: SIM101 + items.append(md) + elif isinstance(md, dict): + if "unitsV1" in md: + items.append(AllUnits(unitsV1=md["unitsV1"])) + if "attributes" in md: + items.append(UserAttributes(attributes=md["attributes"])) + else: + msg = f"Unsupported metadata element type for coordinate: {type(md)}" + raise TypeError(msg) + else: + msg = f"Unsupported metadata type for coordinate: {type(metadata)}" + raise TypeError(msg) + coord_meta_list = items or None + return Coordinate( + name=name, + longName=long_name, + dimensions=dimensions, + dataType=data_type, + metadata=coord_meta_list, + ) + + +def make_variable( # noqa: PLR0913 PLR0912 + name: str, + dimensions: list[NamedDimension | str], + data_type: ScalarType | StructuredType, + long_name: str = None, + compressor: Blosc | ZFP | None = None, + coordinates: list[Coordinate | str] | None = None, + metadata: list[AllUnits | UserAttributes] | dict[str, Any] | VariableMetadata | None = None, +) -> Variable: + """Create a Variable with the given parameters. + + Args: + name: Name of the variable + dimensions: List of dimensions + data_type: Data type of the variable + long_name: Optional long name + compressor: Optional compressor + coordinates: Optional list of coordinates + metadata: Optional metadata + + Returns: + Variable: A Variable instance with the specified parameters. + + Raises: + TypeError: If the metadata type is not supported. + """ + + def _to_serializable(val: object) -> dict[str, Any] | object: + return val.model_dump(mode="json", by_alias=True) if hasattr(val, "model_dump") else val + + var_metadata = None + if metadata: + if isinstance(metadata, list): + metadata_dict = {} + for md in metadata: + if isinstance(md, AllUnits): + val = md.units_v1 + if isinstance(val, list) and len(val) == 1: + val = val[0] + metadata_dict["unitsV1"] = val + elif isinstance(md, UserAttributes): + attrs = _to_serializable(md) + metadata_dict["attributes"] = ( + attrs[0] if isinstance(attrs, list) and len(attrs) == 1 else attrs + ) + var_metadata = VariableMetadata(**metadata_dict) + + elif isinstance(metadata, dict): + converted_dict = {} + for key, value in metadata.items(): + if key == "unitsV1": + val = value[0] if isinstance(value, list) and len(value) == 1 else value + converted_dict["unitsV1"] = _to_serializable(val) + else: + converted_dict[key] = value + var_metadata = VariableMetadata(**converted_dict) + + elif isinstance(metadata, VariableMetadata): + # Flatten any single-element list fields in metadata + md = metadata.model_dump(by_alias=True, exclude_none=True) + for key, value in list(md.items()): + if isinstance(value, list) and len(value) == 1: + md[key] = value[0] + var_metadata = VariableMetadata(**md) + + else: + msg = f"Unsupported metadata type: {type(metadata)}" + raise TypeError(msg) + + return Variable( + name=name, + longName=long_name, + dimensions=dimensions, + dataType=data_type, + compressor=compressor, + coordinates=coordinates, + metadata=var_metadata, + ) + + +def make_dataset_metadata( + name: str, + api_version: str, + created_on: datetime, + attributes: dict[str, Any] | None = None, +) -> DatasetMetadata: + """Create a DatasetMetadata with name, api_version, created_on, and optional attributes.""" + dataset_metadata_dict = { + "name": name, + "apiVersion": api_version, + "createdOn": created_on, + "attributes": attributes, + } + return DatasetMetadata(**dataset_metadata_dict) + + +def make_dataset( + variables: list[Variable], + metadata: DatasetMetadata, +) -> MDIODataset: + """Create a Dataset with the given variables and metadata.""" + return MDIODataset( + variables=variables, + metadata=metadata, + ) + + +def _convert_compressor( + model: Blosc | ZFP | None, +) -> NumcodecsBlosc | ZFPY | None: + if isinstance(model, Blosc): + return NumcodecsBlosc( + cname=model.algorithm.value, + clevel=model.level, + shuffle=model.shuffle.value, + blocksize=model.blocksize if model.blocksize > 0 else 0, + ) + if isinstance(model, ZFP): + if zfpy_base is None or ZFPY is None: + msg = "zfpy and numcodecs are required to use ZFP compression" + raise ImportError(msg) + return ZFPY( + mode=model.mode.value, + tolerance=model.tolerance, + rate=model.rate, + precision=model.precision, + ) + if model is None: + return None + msg = f"Unsupported compressor model: {type(model)}" + raise TypeError(msg) + + +def _construct_mdio_dataset(mdio_ds: MDIODataset) -> mdio.Dataset: # noqa: PLR0912 + """Build an MDIO dataset with correct dimensions and dtypes. + + This internal function constructs the underlying data structure for an MDIO dataset, + handling dimension mapping, data types, and metadata organization. + + Args: + mdio_ds: The source MDIO dataset to construct from. + + Returns: + The constructed dataset with proper MDIO structure and metadata. + + Raises: + TypeError: If an unsupported data type is encountered. + """ + # Collect dimension sizes + dims: dict[str, int] = {} + for var in mdio_ds.variables: + for d in var.dimensions: + if isinstance(d, NamedDimension): + dims[d.name] = d.size + + # Build data variables + data_vars: dict[str, mdio.DataArray] = {} + for var in mdio_ds.variables: + dim_names = [d.name if isinstance(d, NamedDimension) else d for d in var.dimensions] + shape = tuple(dims[name] for name in dim_names) + dt = var.data_type + if isinstance(dt, ScalarType): + dtype = np.dtype(dt.value) + elif isinstance(dt, StructuredType): + dtype = np.dtype([(f.name, f.format.value) for f in dt.fields]) + else: + msg = f"Unsupported data_type: {dt}" + raise TypeError(msg) + arr = np.zeros(shape, dtype=dtype) + data_array = mdio.DataArray(arr, dims=dim_names) + data_array.encoding["fill_value"] = 0.0 + + # Set long_name if present + if var.long_name is not None: + data_array.attrs["long_name"] = var.long_name + + # Set coordinates if present, excluding dimension names + if var.coordinates is not None: + dim_set = set(dim_names) + coord_names = [ + c.name if isinstance(c, Coordinate) else c + for c in var.coordinates + if (c.name if isinstance(c, Coordinate) else c) not in dim_set + ] + if coord_names: + data_array.attrs["coordinates"] = " ".join(coord_names) + + # Attach variable metadata into DataArray attributes + if var.metadata is not None: + md = var.metadata.model_dump( + by_alias=True, + exclude_none=True, + exclude={"chunk_grid"}, + ) + for key, value in md.items(): + if isinstance(value, list) and len(value) == 1: + md[key] = value[0] + data_array.attrs.update(md) + data_vars[var.name] = data_array + + ds = mdio.Dataset(data_vars) + # Attach dataset metadata + ds.attrs["apiVersion"] = mdio_ds.metadata.api_version + ds.attrs["createdOn"] = str(mdio_ds.metadata.created_on) + ds.attrs["name"] = mdio_ds.metadata.name + if mdio_ds.metadata.attributes: + ds.attrs["attributes"] = mdio_ds.metadata.attributes + return ds diff --git a/src/mdio/core/v1/builder.py b/src/mdio/core/v1/builder.py new file mode 100644 index 00000000..8ffef60f --- /dev/null +++ b/src/mdio/core/v1/builder.py @@ -0,0 +1,299 @@ +"""Builder pattern implementation for MDIO v1 schema models.""" + +from collections.abc import Mapping +from datetime import UTC +from datetime import datetime +from enum import Enum +from enum import auto +from typing import Any + +from zarr.core.chunk_key_encodings import V2ChunkKeyEncoding # noqa: F401 + +from mdio.core.v1._overloads import mdio +from mdio.schemas.compressors import ZFP +from mdio.schemas.compressors import Blosc +from mdio.schemas.dimension import NamedDimension +from mdio.schemas.dtype import ScalarType +from mdio.schemas.dtype import StructuredType +from mdio.schemas.metadata import UserAttributes +from mdio.schemas.v1.dataset import Dataset +from mdio.schemas.v1.units import AllUnits +from mdio.schemas.v1.variable import Coordinate +from mdio.schemas.v1.variable import Variable +from mdio.schemas.v1.variable import VariableMetadata + +# Import factory functions from serializer module +from ._serializer import _construct_mdio_dataset +from ._serializer import _convert_compressor +from ._serializer import make_coordinate +from ._serializer import make_dataset +from ._serializer import make_dataset_metadata +from ._serializer import make_named_dimension +from ._serializer import make_variable + + +class _BuilderState(Enum): + """States for the template builder.""" + + INITIAL = auto() + HAS_DIMENSIONS = auto() + HAS_COORDINATES = auto() + HAS_VARIABLES = auto() + + +class MDIODatasetBuilder: + """Builder for creating MDIO datasets with enforced build order. + + This builder implements the builder pattern to create MDIO datasets with a v1 schema. + It enforces a specific build order to ensure valid dataset construction: + 1. Must add dimensions first via add_dimension() + 2. Can optionally add coordinates via add_coordinate() + 3. Must add variables via add_variable() + 4. Must call build() to create the dataset. + """ + + def __init__(self, name: str, attributes: dict[str, Any] | None = None): + self.name = name + self.api_version = "1.0.0" # TODO(BrianMichell, #0): Pull from package metadata + self.created_on = datetime.now(UTC) + self.attributes = attributes + self._dimensions: list[NamedDimension] = [] + self._coordinates: list[Coordinate] = [] + self._variables: list[Variable] = [] + self._state = _BuilderState.INITIAL + self._unnamed_variable_counter = 0 + + def add_dimension( # noqa: PLR0913 + self, + name: str, + size: int, + long_name: str = None, + data_type: ScalarType | StructuredType = ScalarType.INT32, + metadata: list[AllUnits | UserAttributes] | None | dict[str, Any] = None, + ) -> "MDIODatasetBuilder": + """Add a dimension. + + This must be called at least once before adding coordinates or variables. + + Args: + name: Name of the dimension + size: Size of the dimension + long_name: Optional long name for the dimension variable + data_type: Data type for the dimension variable (defaults to INT32) + metadata: Optional metadata for the dimension variable + + Returns: + self: Returns self for method chaining + """ + # Create the dimension + dimension = make_named_dimension(name, size) + self._dimensions.append(dimension) + + # Create a variable for the dimension + dim_var = make_variable( + name=name, + long_name=long_name, + dimensions=[dimension], + data_type=data_type, + metadata=metadata, + ) + self._variables.append(dim_var) + + self._state = _BuilderState.HAS_DIMENSIONS + return self + + def add_coordinate( # noqa: PLR0913 + self, + name: str = "", + *, + long_name: str = None, + dimensions: list[NamedDimension | str] | None = None, + data_type: ScalarType | StructuredType = ScalarType.FLOAT32, + metadata: list[AllUnits | UserAttributes] | None | dict[str, Any] = None, + ) -> "MDIODatasetBuilder": + """Add a coordinate after adding at least one dimension.""" + if self._state == _BuilderState.INITIAL: + msg = "Must add at least one dimension before adding coordinates" + raise ValueError(msg) + + if name == "": + name = f"coord_{len(self._coordinates)}" + if dimensions is None: + dimensions = self._dimensions + if isinstance(metadata, dict): + metadata = [metadata] + + # Convert string dimension names to NamedDimension objects + dim_objects = [] + for dim in dimensions: + if isinstance(dim, str): + dim_obj = next((d for d in self._dimensions if d.name == dim), None) + if dim_obj is None: + msg = f"Dimension {dim!r} not found" + raise ValueError(msg) + dim_objects.append(dim_obj) + else: + dim_objects.append(dim) + + self._coordinates.append( + make_coordinate( + name=name, + long_name=long_name, + dimensions=dim_objects, + data_type=data_type, + metadata=metadata, + ) + ) + self._state = _BuilderState.HAS_COORDINATES + return self + + def add_variable( # noqa: PLR0913 + self, + name: str = "", + *, + long_name: str = None, + dimensions: list[NamedDimension | str] | None = None, + data_type: ScalarType | StructuredType = ScalarType.FLOAT32, + compressor: Blosc | ZFP | None = None, + coordinates: list[Coordinate | str] | None = None, + metadata: VariableMetadata | None = None, + ) -> "MDIODatasetBuilder": + """Add a variable after adding at least one dimension.""" + if self._state == _BuilderState.INITIAL: + msg = "Must add at least one dimension before adding variables" + raise ValueError(msg) + + if name == "": + name = f"var_{self._unnamed_variable_counter}" + self._unnamed_variable_counter += 1 + if dimensions is None: + dimensions = self._dimensions + + # Convert string dimension names to NamedDimension objects + dim_objects = [] + for dim in dimensions: + if isinstance(dim, str): + dim_obj = next((d for d in self._dimensions if d.name == dim), None) + if dim_obj is None: + msg = f"Dimension {dim!r} not found" + raise ValueError(msg) + dim_objects.append(dim_obj) + else: + dim_objects.append(dim) + + self._variables.append( + make_variable( + name=name, + long_name=long_name, + dimensions=dim_objects, + data_type=data_type, + compressor=compressor, + coordinates=coordinates, + metadata=metadata, + ) + ) + self._state = _BuilderState.HAS_VARIABLES + return self + + def build(self) -> Dataset: + """Build the final dataset.""" + if self._state == _BuilderState.INITIAL: + msg = "Must add at least one dimension before building" + raise ValueError(msg) + + metadata = make_dataset_metadata( + self.name, self.api_version, self.created_on, self.attributes + ) + + # Add coordinates as variables to the dataset + # We make a copy so that coordinates are not duplicated if the builder is reused + all_variables = self._variables.copy() + for coord in self._coordinates: + # Convert coordinate to variable + coord_var = make_variable( + name=coord.name, + long_name=coord.long_name, + dimensions=coord.dimensions, + data_type=coord.data_type, + metadata=coord.metadata, + ) + all_variables.append(coord_var) + + return make_dataset(all_variables, metadata) + + def to_mdio( + self, + store: str, + mode: str = "w", + compute: bool = False, + **kwargs: Mapping[str, str | int | float | bool], + ) -> Dataset: + """Write the dataset to a Zarr store and return the constructed mdio.Dataset. + + This function constructs an mdio.Dataset from the MDIO dataset and writes its metadata + to a Zarr store. The actual data is not written, only the metadata structure is created. + """ + return write_mdio_metadata(self.build(), store, mode, compute, **kwargs) + + +def write_mdio_metadata( + mdio_ds: Dataset, + store: str, + mode: str = "w", + compute: bool = False, + **kwargs: Mapping[str, str | int | float | bool], +) -> mdio.Dataset: + """Write MDIO metadata to a Zarr store and return the constructed mdio.Dataset. + + This function constructs an mdio.Dataset from the MDIO dataset and writes its metadata + to a Zarr store. The actual data is not written, only the metadata structure is created. + + Args: + mdio_ds: The MDIO dataset to serialize + store: Path to the Zarr or .mdio store + mode: Write mode to pass to to_mdio(), e.g. 'w' or 'a' + compute: Whether to compute (write) array chunks (True) or only metadata (False) + **kwargs: Additional arguments to pass to to_mdio() + + Returns: + The constructed xarray Dataset with MDIO extensions + """ + ds = _construct_mdio_dataset(mdio_ds) + + def _generate_encodings() -> dict: + """Generate encodings for each variable in the MDIO dataset. + + Returns: + Dictionary mapping variable names to their encoding configurations. + """ + # TODO(Anybody, #10274): Re-enable chunk_key_encoding when supported by xarray + # dimension_separator_encoding = V2ChunkKeyEncoding(separator="/").to_dict() + global_encodings = {} + for var in mdio_ds.variables: + fill_value = 0 + if isinstance(var.data_type, StructuredType): + continue + chunks = None + if var.metadata is not None and var.metadata.chunk_grid is not None: + chunks = var.metadata.chunk_grid.configuration.chunk_shape + global_encodings[var.name] = { + "chunks": chunks, + # TODO(Anybody, #10274): Re-enable chunk_key_encoding when supported by xarray + # "chunk_key_encoding": dimension_separator_encoding, + "_FillValue": fill_value, + "dtype": var.data_type, + "compressors": _convert_compressor(var.compressor), + } + return global_encodings + + ds.to_mdio( + store, + mode=mode, + zarr_format=2, + consolidated=True, + safe_chunks=False, + compute=compute, + encoding=_generate_encodings(), + **kwargs, + ) + return ds diff --git a/src/mdio/core/v1/factory.py b/src/mdio/core/v1/factory.py new file mode 100644 index 00000000..3cc7ced7 --- /dev/null +++ b/src/mdio/core/v1/factory.py @@ -0,0 +1,259 @@ +"""MDIO factories for seismic data.""" + +# TODO(BrianMichell, #535): Add implementations for other canonical datasets. + +from __future__ import annotations + +from enum import Enum +from enum import auto +from typing import TYPE_CHECKING +from typing import Any + +from mdio.core.v1.builder import MDIODatasetBuilder +from mdio.schemas.compressors import Blosc +from mdio.schemas.dtype import ScalarType +from mdio.schemas.dtype import StructuredType + +if TYPE_CHECKING: + from mdio.schemas.v1.dataset import Dataset + +import json + +from pydantic import ValidationError + + +def from_contract(store: str, contract: str | dict) -> Dataset: + """Creates an MDIO Dataset from the contract and writes the metadata to the store. + + Args: + store: The store to write the metadata to. + contract: The contract to create the dataset from. + + Raises: + ValueError: If the contract cannot be validated successfully. + + Returns: + The created MDIO Dataset. + """ + from mdio.core.v1._serializer import _construct_mdio_dataset + from mdio.schemas.v1 import Dataset as V1Dataset + + if isinstance(contract, str): + contract = json.loads(contract) + + try: + V1Dataset.model_validate(contract) + except ValidationError as e: + msg = f"Failed to validate the input contract: {e}" + raise ValueError(msg) from e + + ds = _construct_mdio_dataset(contract) + return ds.to_mdio(store) + + +class MDIOSchemaType(Enum): + """MDIO templates for specific data types.""" + + SEISMIC_3D_POST_STACK_GENERIC = auto() + SEISMIC_3D_POST_STACK_TIME = auto() + SEISMIC_3D_POST_STACK_DEPTH = auto() + SEISMIC_3D_PRE_STACK_CDP_TIME = auto() + SEISMIC_3D_PRE_STACK_CDP_DEPTH = auto() + + +class Seismic3DPostStackGeneric: + """Generic 3D seismic post stack dataset.""" + + def __init__(self) -> None: + self._dim_names = ["inline", "crossline", "sample"] + self._chunks = [128, 128, 128] # 8 mb + self._coords = { + "cdp-x": ("float64", {"unitsV1": {"length": "m"}}, self._dim_names[:-1]), + "cdp-y": ("float64", {"unitsV1": {"length": "m"}}, self._dim_names[:-1]), + } + + def create( # noqa: PLR0913 + self, + name: str, + shape: list[int], + header_fields: dict[str, str], + create_coords: bool = False, + sample_format: str | None = None, + chunks: list[int] | None = None, + sample_units: dict[str, str] | None = None, + z_units: dict[str, str] | None = None, + attributes: dict[str, Any] | None = None, + ) -> Dataset: + """Create a generic seismic dataset schema. + + Args: + name: Name of the dataset + shape: Shape of the dataset + header_fields: Header fields to include as a dict of field_name: dtype + create_coords: Whether to create coordinates + sample_format: Format of the samples + chunks: Chunk sizes + sample_units: Units for samples + z_units: Units for z-axis + attributes: Additional attributes to include in the dataset metadata + + Returns: + Dataset: The created dataset + """ + chunks = chunks or self._chunks + sample_format = sample_format or "float32" + + builder = MDIODatasetBuilder( + name=name, + attributes=attributes, + ) + + # Add dimensions + for dim_name, dim_size in zip(self._dim_names, shape, strict=True): + builder.add_dimension( + name=dim_name, + size=dim_size, + data_type=ScalarType.UINT32, + metadata=z_units if dim_name == "sample" else None, + ) + + # Add coordinates if requested + if create_coords: + for coord_name, (format_, unit, coord_dims) in self._coords.items(): + builder.add_coordinate( + name=coord_name, + data_type=ScalarType(format_), + dimensions=coord_dims, + metadata=unit, + ) + + # Add seismic variable + builder.add_variable( + name="seismic", + data_type=ScalarType(sample_format), + dimensions=self._dim_names, + compressor=Blosc(name="blosc", algorithm="zstd"), + metadata=sample_units, + ) + + # Add header variable with structured dtype + header_dtype = StructuredType( + fields=[ + {"name": field_name, "format": field_type} + for field_name, field_type in header_fields.items() + ] + ) + builder.add_variable( + name="headers", + data_type=header_dtype, + dimensions=self._dim_names[:-1], + compressor=Blosc(name="blosc"), + ) + + # Add trace mask + builder.add_variable( + name="trace_mask", + data_type=ScalarType.BOOL, + dimensions=self._dim_names[:-1], + compressor=Blosc(name="blosc"), + ) + + return builder.build() + + +class Seismic3DPostStack(Seismic3DPostStackGeneric): + """3D seismic post stack dataset with domain-specific attributes.""" + + def __init__(self, domain: str) -> None: + super().__init__() + self._dim_names = ["inline", "crossline", domain] + + def create( # noqa: PLR0913 + self, + name: str, + shape: list[int], + header_fields: dict[str, str], + create_coords: bool = False, + sample_format: str | None = None, + chunks: list[int] | None = None, + sample_units: dict[str, str] | None = None, + z_units: dict[str, str] | None = None, + attributes: dict[str, Any] | None = None, + ) -> Dataset: + """Create a seismic dataset schema with domain-specific attributes.""" + # Add seismic-specific attributes + seismic_attrs = { + "surveyDimensionality": "3D", + "ensembleType": "line", + "processingStage": "post-stack", + } + if attributes: + seismic_attrs.update(attributes) + + return super().create( + name=name, + shape=shape, + header_fields=header_fields, + create_coords=create_coords, + sample_format=sample_format, + chunks=chunks, + sample_units=sample_units, + z_units=z_units, + attributes=seismic_attrs, + ) + + +class Seismic3DPreStack(Seismic3DPostStackGeneric): + """3D seismic pre stack dataset.""" + + def __init__(self, domain: str) -> None: + super().__init__() + self._dim_names = ["inline", "crossline", "offset", domain] + self._chunks = [1, 1, 512, 4096] # 8 mb + self._coords = { + "cdp-x": ("float64", {"length": "m"}, self._dim_names[:-2]), + "cdp-y": ("float64", {"length": "m"}, self._dim_names[:-2]), + } + + def create( # noqa: PLR0913 + self, + name: str, + shape: list[int], + header_fields: dict[str, str], + create_coords: bool = False, + sample_format: str | None = None, + chunks: list[int] | None = None, + sample_units: dict[str, str] | None = None, + z_units: dict[str, str] | None = None, + attributes: dict[str, Any] | None = None, + ) -> Dataset: + """Create a seismic dataset schema with pre-stack attributes.""" + # Add seismic-specific attributes + seismic_attrs = { + "surveyDimensionality": "3D", + "ensembleType": "cdp", + "processingStage": "pre-stack", + } + if attributes: + seismic_attrs.update(attributes) + + return super().create( + name=name, + shape=shape, + header_fields=header_fields, + create_coords=create_coords, + sample_format=sample_format, + chunks=chunks, + sample_units=sample_units, + z_units=z_units, + attributes=seismic_attrs, + ) + + +SCHEMA_TEMPLATE_MAP = { + MDIOSchemaType.SEISMIC_3D_POST_STACK_GENERIC: Seismic3DPostStackGeneric(), + MDIOSchemaType.SEISMIC_3D_POST_STACK_TIME: Seismic3DPostStack("time"), + MDIOSchemaType.SEISMIC_3D_POST_STACK_DEPTH: Seismic3DPostStack("depth"), + MDIOSchemaType.SEISMIC_3D_PRE_STACK_CDP_TIME: Seismic3DPreStack("time"), + MDIOSchemaType.SEISMIC_3D_PRE_STACK_CDP_DEPTH: Seismic3DPreStack("depth"), +} diff --git a/src/mdio/schemas/builder.py b/src/mdio/schemas/builder.py deleted file mode 100644 index 40908ff0..00000000 --- a/src/mdio/schemas/builder.py +++ /dev/null @@ -1,154 +0,0 @@ -"""Schema builders.""" - -from __future__ import annotations - -from typing import Any - -from mdio.schemas import NamedDimension -from mdio.schemas.v1.dataset import Dataset -from mdio.schemas.v1.dataset import DatasetMetadata -from mdio.schemas.v1.variable import Variable -from mdio.schemas.v1.variable import VariableMetadata - - -class VariableBuilder: - """Dataset builder.""" - - def __init__(self) -> None: - self.name = None - self.long_name = None - self.dtype = None - self.chunks = None - self.dims = None - self.coords = None - self.compressor = None - self.meta_dict = None - - def set_name(self, name: str) -> VariableBuilder: - """Set variable name.""" - self.name = name - return self - - def set_long_name(self, long_name: str) -> VariableBuilder: - """Add long, descriptive name to the variable.""" - self.long_name = long_name - return self - - def set_compressor(self, compressor: dict[str, Any]) -> VariableBuilder: - """Add long, descriptive name to the variable.""" - self.compressor = compressor - return self - - def add_dimension(self, *dimensions: str | dict[str, int]) -> VariableBuilder: - """Add a dimension to the dataset.""" - if self.dims is None: - self.dims = [] - - if isinstance(dimensions[0], str): - dims = list(dimensions) - elif isinstance(dimensions[0], dict): - dims = [ - NamedDimension(name=name, size=size) - for dim in dimensions - for name, size in dim.items() - ] - else: - raise NotImplementedError - - self.dims.extend(dims) - return self - - def add_coordinate(self, *names: str) -> VariableBuilder: - """Add a coordinate to the variable.""" - if self.coords is None: - self.coords = [] - - self.coords.extend(names) - return self - - def set_format(self, format_: str | dict[str, str]) -> VariableBuilder: - """Set variable format.""" - if isinstance(format_, dict): - fields = [{"name": n, "format": f} for n, f in format_.items()] - format_ = {"fields": fields} - - self.dtype = format_ - return self - - def set_chunks(self, chunks: list[int]) -> VariableBuilder: - """Set variable chunks.""" - if self.meta_dict is None: - self.meta_dict = {} - - self.meta_dict["chunkGrid"] = {"configuration": {"chunkShape": chunks}} - return self - - def set_units(self, units: dict[str, str]) -> VariableBuilder: - """Set variable units.""" - if self.meta_dict is None: - self.meta_dict = {} - - self.meta_dict["unitsV1"] = units - return self - - def add_attribute(self, key: str, value: Any) -> VariableBuilder: # noqa: ANN401 - """Add a user attribute to the variable metadata.""" - if self.meta_dict is None: - self.meta_dict = {} - - self.meta_dict["attributes"] = {key: value} - return self - - def build(self) -> Variable: - """Build the dataset model.""" - if self.chunks is not None and len(self.chunks) != len(self.dims): - msg = "Variable chunks must have same number of dimensions." - raise ValueError(msg) - - var_kwargs = {} - - if self.meta_dict is not None: - var_kwargs["metadata"] = VariableMetadata.model_validate(self.meta_dict) - - return Variable( - name=self.name, - long_name=self.long_name, - data_type=self.dtype, - dimensions=self.dims, - coordinates=self.coords, - compressor=self.compressor, - **var_kwargs, - ) - - -class DatasetBuilder: - """Dataset builder.""" - - def __init__(self) -> None: - self.variables = [] - self.name = None - self.metadata = None - - def set_name(self, name: str) -> DatasetBuilder: - """Set dataset name.""" - self.name = name - return self - - def add_variable(self, variable: Variable) -> DatasetBuilder: - """Add a variable to the dataset.""" - self.variables.append(variable) - return self - - def add_variables(self, variables: list[Variable]) -> DatasetBuilder: - """Add multiple variables to the dataset.""" - [self.add_variable(variable) for variable in variables] - return self - - def set_metadata(self, metadata: DatasetMetadata) -> DatasetBuilder: - """Add a metadata to the dataset.""" - self.metadata = metadata - return self - - def build(self) -> Dataset: - """Build the dataset model.""" - return Dataset(variables=self.variables, metadata=self.metadata) diff --git a/src/mdio/schemas/core.py b/src/mdio/schemas/core.py index 34a09066..7768be06 100644 --- a/src/mdio/schemas/core.py +++ b/src/mdio/schemas/core.py @@ -3,10 +3,10 @@ from __future__ import annotations from typing import Any -from typing import get_type_hints from pydantic import BaseModel from pydantic import ConfigDict +from pydantic import Field from pydantic.alias_generators import to_camel @@ -28,12 +28,16 @@ def model_fields(model: type[BaseModel]) -> dict[str, tuple[Any, Any]]: >>> model_fields(MyModel) {'name': (str,