segy_to_mdio_v1

src/mdio/api/convenience.py

@@ -79,7 +79,11 @@ def copy_mdio( # noqa: PLR0913
 
  writer.live_mask[:] = reader.live_mask[:]
 
- iterator = ChunkIterator(reader._traces, chunk_samples=False)
+ shape = reader._traces.shape
+ chunks = reader._traces.chunks
+ chunks = chunks[:-1] + (shape[-1],) # don't chunk samples
+
+ iterator = ChunkIterator(shape=shape, chunks=chunks)
  progress = tqdm(iterator, unit="block")
  progress.set_description(desc=f"Copying data for '{access_pattern=}'")
  for slice_ in progress:
@@ -177,7 +181,10 @@ def create_rechunk_plan(
 
  n_dimension = len(data_array.shape)
  dummy_array = zarr.empty(shape=data_array.shape, chunks=(MAX_BUFFER,) * n_dimension)
- iterator = ChunkIterator(dummy_array)
+
+ shape = dummy_array.shape
+ chunks = dummy_array.chunks
+ iterator = ChunkIterator(shape=shape, chunks=chunks)
 
  return metadata_arrs, data_arrs, live_mask, iterator
 

src/mdio/converters/numpy.py

@@ -7,7 +7,6 @@
 import numpy as np
 
 from mdio.api.accessor import MDIOWriter
-from mdio.converters.segy import get_compressor
 from mdio.core.dimension import Dimension
 from mdio.core.factory import MDIOCreateConfig
 from mdio.core.factory import MDIOVariableConfig
@@ -137,6 +136,11 @@ def numpy_to_mdio( # noqa: PLR0913
  suffix = [str(idx) for idx, value in enumerate(suffix) if value is not None]
  suffix = "".join(suffix)
 
+ # TODO(Dmitrit Repin): Implement Numpy converted in MDIO v1
+ # https://github.com/TGSAI/mdio-python/issues/596
+ def get_compressor(lossless: bool, tolerance: float) -> list[str]:
+ pass
+
  compressors = get_compressor(lossless, compression_tolerance)
  mdio_var = MDIOVariableConfig(
  name=f"chunked_{suffix}",

src/mdio/converters/type_converter.py

@@ -0,0 +1,85 @@
+"""A module for converting numpy dtypes to MDIO scalar and structured types."""
+
+from numpy import dtype as np_dtype
+
+from mdio.schemas.dtype import ScalarType
+from mdio.schemas.dtype import StructuredField
+from mdio.schemas.dtype import StructuredType
+
+
+def to_scalar_type(data_type: np_dtype) -> ScalarType:
+ """Convert numpy dtype to MDIO ScalarType.
+
+ Out of the 24 built-in numpy scalar type objects
+ (see https://numpy.org/doc/stable/reference/arrays.dtypes.html)
+ this function supports only a limited subset:
+ ScalarType.INT8 <-> int8
+ ScalarType.INT16 <-> int16
+ ScalarType.INT32 <-> int32
+ ScalarType.INT64 <-> int64
+ ScalarType.UINT8 <-> uint8
+ ScalarType.UINT16 <-> uint16
+ ScalarType.UINT32 <-> uint32
+ ScalarType.UINT64 <-> uint64
+ ScalarType.FLOAT32 <-> float32
+ ScalarType.FLOAT64 <-> float64
+ ScalarType.COMPLEX64 <-> complex64
+ ScalarType.COMPLEX128 <-> complex128
+ ScalarType.BOOL <-> bool
+
+ Args:
+ data_type: numpy dtype to convert
+
+ Returns:
+ ScalarType: corresponding MDIO scalar type
+
+ Raises:
+ ValueError: if dtype is not supported
+ """
+ try:
+ return ScalarType(data_type.name)
+ except ValueError as exc:
+ err = f"Unsupported numpy dtype '{data_type.name}' for conversion to ScalarType."
+ raise ValueError(err) from exc
+
+
+def to_structured_type(data_type: np_dtype) -> StructuredType:
+ """Convert numpy dtype to MDIO StructuredType.
+
+ This function supports only a limited subset of structured types.
+ In particular:
+ It does not support nested structured types.
+ It supports fields of only 13 out of 24 built-in numpy scalar types.
+ (see `to_scalar_type` for details).
+
+ Args:
+ data_type: numpy dtype to convert
+
+ Returns:
+ StructuredType: corresponding MDIO structured type
+
+ Raises:
+ ValueError: if dtype is not structured or has no fields
+
+ """
+ if data_type is None or len(data_type.names or []) == 0:
+ err = "None or empty dtype provided, cannot convert to StructuredType."
+ raise ValueError(err)
+
+ fields = []
+ for field_name in data_type.names:
+ field_dtype = data_type.fields[field_name][0]
+ scalar_type = to_scalar_type(field_dtype)
+ structured_field = StructuredField(name=field_name, format=scalar_type)
+ fields.append(structured_field)
+ return StructuredType(fields=fields)
+
+
+def to_numpy_dtype(data_type: ScalarType | StructuredType) -> np_dtype:
+ """Get the numpy dtype for a variable."""
+ if isinstance(data_type, ScalarType):
+ return np_dtype(data_type.value)
+ if isinstance(data_type, StructuredType):
+ return np_dtype([(f.name, f.format.value) for f in data_type.fields])
+ msg = f"Expected ScalarType or StructuredType, got '{type(data_type).__name__}'"
+ raise ValueError(msg)

src/mdio/core/indexing.py

@@ -4,78 +4,83 @@
 from math import ceil
 
 import numpy as np
-from zarr import Array
 
 
 class ChunkIterator:
- """Iterator for traversing a Zarr array in chunks.
+ """Chunk iterator for multi-dimensional arrays.
 
- This iterator yields tuples of slices corresponding to the chunk boundaries of a Zarr array.
- It supports chunking all dimensions or taking the full extent of the last dimension.
+ This iterator takes an array shape and chunks and every time it is iterated, it returns
+ a dictionary (if dimensions are provided) or a tuple of slices that align with
+ chunk boundaries. When dimensions are provided, they are used as the dictionary keys.
 
  Args:
- array: The Zarr array to iterate, providing shape and chunk sizes.
- chunk_samples: If True, chunks all dimensions. If False, takes the full extent of the
- last dimension. Defaults to True.
-
-
- Example:
- >>> import zarr
- >>> arr = zarr.array(np.zeros((10, 20)), chunks=(3, 4))
- >>> it = ChunkIterator(arr)
- >>> for slices in it:
- ... print(slices)
- (slice(0, 3, None), slice(0, 4, None))
- (slice(0, 3, None), slice(4, 8, None))
- ...
- >>> it = ChunkIterator(arr, chunk_samples=False)
- >>> for slices in it:
- ... print(slices)
- (slice(0, 3, None), slice(0, 20, None))
- (slice(3, 6, None), slice(0, 20, None))
- ...
+ shape: The shape of the array.
+ chunks: The chunk sizes for each dimension.
+ dim_names: The names of the array dimensions, to be used with DataArray.isel().
+ If the dim_names are not provided, a tuple of the slices will be returned.
+
+ Attributes: # noqa: DOC602
+ arr_shape: Shape of the array.
+ len_chunks: Length of chunks in each dimension.
+ dim_chunks: Number of chunks in each dimension.
+ num_chunks: Total number of chunks.
+
+ Examples:
+ >> chunks = (3, 4, 5)
+ >> shape = (5, 11, 19)
+ >> dims = ["inline", "crossline", "depth"]
+ >>
+ >> iter = ChunkIterator(shape=shape, chunks=chunks, dim_names=dims)
+ >> for i in range(13):
+ >> region = iter.__next__()
+ >> print(region)
+ { "inline": slice(3,6, None), "crossline": slice(0,4, None), "depth": slice(0,5, None) }
+
+ >> iter = ChunkIterator(shape=shape, chunks=chunks, dim_names=None)
+ >> for i in range(13):
+ >> region = iter.__next__()
+ >> print(region)
+ (slice(3,6,None), slice(0,4,None), slice(0,5,None))
  """
 
- def __init__(self, array: Array, chunk_samples: bool = True):
- self.arr_shape = array.shape
- self.len_chunks = array.chunks
-
- # If chunk_samples is False, set the last dimension's chunk size to its full extent
- if not chunk_samples:
- self.len_chunks = self.len_chunks[:-1] + (self.arr_shape[-1],)
-
- # Calculate the number of chunks per dimension
- self.dim_chunks = [
- ceil(len_dim / chunk)
- for len_dim, chunk in zip(self.arr_shape, self.len_chunks, strict=True)
- ]
+ def __init__(
+ self, shape: tuple[int, ...], chunks: tuple[int, ...], dim_names: tuple[str, ...] = None
+ ):
+ self.arr_shape = tuple(shape) # Deep copy to ensure immutability
+ self.len_chunks = tuple(chunks) # Deep copy to ensure immutability
+ self.dims = dim_names
+
+ # Compute number of chunks per dimension, and total number of chunks
+ self.dim_chunks = tuple(
+ [
+ ceil(len_dim / chunk)
+ for len_dim, chunk in zip(self.arr_shape, self.len_chunks, strict=True)
+ ]
+ )
  self.num_chunks = np.prod(self.dim_chunks)
 
- # Set up chunk index combinations using ranges for each dimension
+ # Under the hood stuff for the iterator. This generates C-ordered
+ # permutation of chunk indices.
  dim_ranges = [range(dim_len) for dim_len in self.dim_chunks]
  self._ranges = itertools.product(*dim_ranges)
  self._idx = 0
 
  def __iter__(self) -> "ChunkIterator":
- """Return the iterator object itself."""
+ """Iteration context."""
  return self
 
  def __len__(self) -> int:
- """Return the total number of chunks."""
+ """Get total number of chunks."""
  return self.num_chunks
 
- def __next__(self) -> tuple[slice, ...]:
- """Yield the next set of chunk slices.
-
- Returns:
- A tuple of slice objects for each dimension.
-
- Raises:
- StopIteration: When all chunks have been iterated over.
- """
- if self._idx < self.num_chunks:
+ def __next__(self) -> dict[str, slice]:
+ """Iteration logic."""
+ if self._idx <= self.num_chunks:
+ # We build slices here. It is dimension agnostic
  current_start = next(self._ranges)
 
+ # TODO (Dmitriy Repin): Enhance ChunkIterator to make the last slice, if needed, smaller
+ # https://github.com/TGSAI/mdio-python/issues/586
  start_indices = tuple(
  dim * chunk for dim, chunk in zip(current_start, self.len_chunks, strict=True)
  )
@@ -88,7 +93,17 @@ def __next__(self) -> tuple[slice, ...]:
  slice(start, stop) for start, stop in zip(start_indices, stop_indices, strict=True)
  )
 
+ if self.dims: # noqa SIM108
+ # Example
+ # {"inline":slice(3,6,None), "crossline":slice(0,4,None), "depth":slice(0,5,None)}
+ result = dict(zip(self.dims, slices, strict=False))
+ else:
+ # Example
+ # (slice(3,6,None), slice(0,4,None), slice(0,5,None))
+ result = slices
+
  self._idx += 1
- return slices
+
+ return result
 
  raise StopIteration

src/mdio/core/storage_location.py

@@ -0,0 +1,87 @@
+"""StorageLocation class for managing local and cloud storage locations."""
+
+from pathlib import Path
+from typing import Any
+
+import fsspec
+
+
+# TODO(Dmitriy Repin): Reuse fsspec functions for some methods we implemented here
+# https://github.com/TGSAI/mdio-python/issues/597
+class StorageLocation:
+ """A class to represent a local or cloud storage location for SEG-Y or MDIO files.
+
+ This class abstracts the storage location, allowing for both local file paths and
+ cloud storage URIs (e.g., S3, GCS). It uses fsspec to check existence and manage options.
+ Note, we do not want to make it a dataclass because we want the uri and the options to
+ be read-only immutable properties.
+
+ uri: The URI of the storage location (e.g., '/path/to/file', 'file:///path/to/file',
+ 's3://bucket/path', 'gs://bucket/path').
+ options: Optional dictionary of options for the cloud, such as credentials.
+
+ """
+
+ def __init__(self, uri: str = "", options: dict[str, Any] = None):
+ self._uri = uri
+ self._options = options or {}
+ self._fs = None
+
+ if uri.startswith(("s3://", "gs://")):
+ return
+
+ if uri.startswith("file://"):
+ self._uri = self._uri.removeprefix("file://")
+ # For local paths, ensure they are absolute and resolved
+ self._uri = str(Path(self._uri).resolve())
+ return
+
+ @property
+ def uri(self) -> str:
+ """Get the URI (read-only)."""
+ return self._uri
+
+ @property
+ def options(self) -> dict[str, Any]:
+ """Get the options (read-only)."""
+ # Return a copy to prevent external modification
+ return self._options.copy()
+
+ @property
+ def _filesystem(self) -> fsspec.AbstractFileSystem:
+ """Get the fsspec filesystem instance for this storage location."""
+ if self._fs is None:
+ self._fs = fsspec.filesystem(self._protocol, **self._options)
+ return self._fs
+
+ @property
+ def _path(self) -> str:
+ """Extract the path portion from the URI."""
+ if "://" in self._uri:
+ return self._uri.split("://", 1)[1]
+ return self._uri # For local paths without file:// prefix
+
+ @property
+ def _protocol(self) -> str:
+ """Extract the protocol/scheme from the URI."""
+ if "://" in self._uri:
+ return self._uri.split("://", 1)[0]
+ return "file" # Default to file protocol
+
+ def exists(self) -> bool:
+ """Check if the storage location exists using fsspec."""
+ try:
+ return self._filesystem.exists(self._path)
+ except Exception as e:
+ # Log the error and return False for safety
+ # In a production environment, you might want to use proper logging
+ print(f"Error checking existence of {self._uri}: {e}")
+ return False
+
+ def __str__(self) -> str:
+ """String representation of the storage location."""
+ return self._uri
+
+ def __repr__(self) -> str:
+ """Developer representation of the storage location."""
+ return f"StorageLocation(uri='{self._uri}', options={self._options})"

src/mdio/api/convenience.py

@@ -79,7 +79,11 @@ def copy_mdio( # noqa: PLR0913
 
  writer.live_mask[:] = reader.live_mask[:]
 
- iterator = ChunkIterator(reader._traces, chunk_samples=False)
+ shape = reader._traces.shape
+ chunks = reader._traces.chunks
+ chunks = chunks[:-1] + (shape[-1],) # don't chunk samples
+
+ iterator = ChunkIterator(shape=shape, chunks=chunks)
  progress = tqdm(iterator, unit="block")
  progress.set_description(desc=f"Copying data for '{access_pattern=}'")
  for slice_ in progress:
@@ -177,7 +181,10 @@ def create_rechunk_plan(
 
  n_dimension = len(data_array.shape)
  dummy_array = zarr.empty(shape=data_array.shape, chunks=(MAX_BUFFER,) * n_dimension)
- iterator = ChunkIterator(dummy_array)
+
+ shape = dummy_array.shape
+ chunks = dummy_array.chunks
+ iterator = ChunkIterator(shape=shape, chunks=chunks)
 
  return metadata_arrs, data_arrs, live_mask, iterator
 

src/mdio/converters/numpy.py

@@ -7,7 +7,6 @@
 import numpy as np
 
 from mdio.api.accessor import MDIOWriter
-from mdio.converters.segy import get_compressor
 from mdio.core.dimension import Dimension
 from mdio.core.factory import MDIOCreateConfig
 from mdio.core.factory import MDIOVariableConfig
@@ -137,6 +136,11 @@ def numpy_to_mdio( # noqa: PLR0913
  suffix = [str(idx) for idx, value in enumerate(suffix) if value is not None]
  suffix = "".join(suffix)
 
+ # TODO(Dmitrit Repin): Implement Numpy converted in MDIO v1
+ # https://github.com/TGSAI/mdio-python/issues/596
+ def get_compressor(lossless: bool, tolerance: float) -> list[str]:
+ pass
+
  compressors = get_compressor(lossless, compression_tolerance)
  mdio_var = MDIOVariableConfig(
  name=f"chunked_{suffix}",

src/mdio/converters/type_converter.py

@@ -0,0 +1,85 @@
+"""A module for converting numpy dtypes to MDIO scalar and structured types."""
+
+from numpy import dtype as np_dtype
+
+from mdio.schemas.dtype import ScalarType
+from mdio.schemas.dtype import StructuredField
+from mdio.schemas.dtype import StructuredType
+
+
+def to_scalar_type(data_type: np_dtype) -> ScalarType:
+ """Convert numpy dtype to MDIO ScalarType.
+
+ Out of the 24 built-in numpy scalar type objects
+ (see https://numpy.org/doc/stable/reference/arrays.dtypes.html)
+ this function supports only a limited subset:
+ ScalarType.INT8 <-> int8
+ ScalarType.INT16 <-> int16
+ ScalarType.INT32 <-> int32
+ ScalarType.INT64 <-> int64
+ ScalarType.UINT8 <-> uint8
+ ScalarType.UINT16 <-> uint16
+ ScalarType.UINT32 <-> uint32
+ ScalarType.UINT64 <-> uint64
+ ScalarType.FLOAT32 <-> float32
+ ScalarType.FLOAT64 <-> float64
+ ScalarType.COMPLEX64 <-> complex64
+ ScalarType.COMPLEX128 <-> complex128
+ ScalarType.BOOL <-> bool
+
+ Args:
+ data_type: numpy dtype to convert
+
+ Returns:
+ ScalarType: corresponding MDIO scalar type
+
+ Raises:
+ ValueError: if dtype is not supported
+ """
+ try:
+ return ScalarType(data_type.name)
+ except ValueError as exc:
+ err = f"Unsupported numpy dtype '{data_type.name}' for conversion to ScalarType."
+ raise ValueError(err) from exc
+
+
+def to_structured_type(data_type: np_dtype) -> StructuredType:
+ """Convert numpy dtype to MDIO StructuredType.
+
+ This function supports only a limited subset of structured types.
+ In particular:
+ It does not support nested structured types.
+ It supports fields of only 13 out of 24 built-in numpy scalar types.
+ (see `to_scalar_type` for details).
+
+ Args:
+ data_type: numpy dtype to convert
+
+ Returns:
+ StructuredType: corresponding MDIO structured type
+
+ Raises:
+ ValueError: if dtype is not structured or has no fields
+
+ """
+ if data_type is None or len(data_type.names or []) == 0:
+ err = "None or empty dtype provided, cannot convert to StructuredType."
+ raise ValueError(err)
+
+ fields = []
+ for field_name in data_type.names:
+ field_dtype = data_type.fields[field_name][0]
+ scalar_type = to_scalar_type(field_dtype)
+ structured_field = StructuredField(name=field_name, format=scalar_type)
+ fields.append(structured_field)
+ return StructuredType(fields=fields)
+
+
+def to_numpy_dtype(data_type: ScalarType | StructuredType) -> np_dtype:
+ """Get the numpy dtype for a variable."""
+ if isinstance(data_type, ScalarType):
+ return np_dtype(data_type.value)
+ if isinstance(data_type, StructuredType):
+ return np_dtype([(f.name, f.format.value) for f in data_type.fields])
+ msg = f"Expected ScalarType or StructuredType, got '{type(data_type).__name__}'"
+ raise ValueError(msg)

src/mdio/core/indexing.py

@@ -4,78 +4,83 @@
 from math import ceil
 
 import numpy as np
-from zarr import Array
 
 
 class ChunkIterator:
- """Iterator for traversing a Zarr array in chunks.
+ """Chunk iterator for multi-dimensional arrays.
 
- This iterator yields tuples of slices corresponding to the chunk boundaries of a Zarr array.
- It supports chunking all dimensions or taking the full extent of the last dimension.
+ This iterator takes an array shape and chunks and every time it is iterated, it returns
+ a dictionary (if dimensions are provided) or a tuple of slices that align with
+ chunk boundaries. When dimensions are provided, they are used as the dictionary keys.
 
  Args:
- array: The Zarr array to iterate, providing shape and chunk sizes.
- chunk_samples: If True, chunks all dimensions. If False, takes the full extent of the
- last dimension. Defaults to True.
-
-
- Example:
- >>> import zarr
- >>> arr = zarr.array(np.zeros((10, 20)), chunks=(3, 4))
- >>> it = ChunkIterator(arr)
- >>> for slices in it:
- ... print(slices)
- (slice(0, 3, None), slice(0, 4, None))
- (slice(0, 3, None), slice(4, 8, None))
- ...
- >>> it = ChunkIterator(arr, chunk_samples=False)
- >>> for slices in it:
- ... print(slices)
- (slice(0, 3, None), slice(0, 20, None))
- (slice(3, 6, None), slice(0, 20, None))
- ...
+ shape: The shape of the array.
+ chunks: The chunk sizes for each dimension.
+ dim_names: The names of the array dimensions, to be used with DataArray.isel().
+ If the dim_names are not provided, a tuple of the slices will be returned.
+
+ Attributes: # noqa: DOC602
+ arr_shape: Shape of the array.
+ len_chunks: Length of chunks in each dimension.
+ dim_chunks: Number of chunks in each dimension.
+ num_chunks: Total number of chunks.
+
+ Examples:
+ >> chunks = (3, 4, 5)
+ >> shape = (5, 11, 19)
+ >> dims = ["inline", "crossline", "depth"]
+ >>
+ >> iter = ChunkIterator(shape=shape, chunks=chunks, dim_names=dims)
+ >> for i in range(13):
+ >> region = iter.__next__()
+ >> print(region)
+ { "inline": slice(3,6, None), "crossline": slice(0,4, None), "depth": slice(0,5, None) }
+
+ >> iter = ChunkIterator(shape=shape, chunks=chunks, dim_names=None)
+ >> for i in range(13):
+ >> region = iter.__next__()
+ >> print(region)
+ (slice(3,6,None), slice(0,4,None), slice(0,5,None))
  """
 
- def __init__(self, array: Array, chunk_samples: bool = True):
- self.arr_shape = array.shape
- self.len_chunks = array.chunks
-
- # If chunk_samples is False, set the last dimension's chunk size to its full extent
- if not chunk_samples:
- self.len_chunks = self.len_chunks[:-1] + (self.arr_shape[-1],)
-
- # Calculate the number of chunks per dimension
- self.dim_chunks = [
- ceil(len_dim / chunk)
- for len_dim, chunk in zip(self.arr_shape, self.len_chunks, strict=True)
- ]
+ def __init__(
+ self, shape: tuple[int, ...], chunks: tuple[int, ...], dim_names: tuple[str, ...] = None
+ ):
+ self.arr_shape = tuple(shape) # Deep copy to ensure immutability
+ self.len_chunks = tuple(chunks) # Deep copy to ensure immutability
+ self.dims = dim_names
+
+ # Compute number of chunks per dimension, and total number of chunks
+ self.dim_chunks = tuple(
+ [
+ ceil(len_dim / chunk)
+ for len_dim, chunk in zip(self.arr_shape, self.len_chunks, strict=True)
+ ]
+ )
  self.num_chunks = np.prod(self.dim_chunks)
 
- # Set up chunk index combinations using ranges for each dimension
+ # Under the hood stuff for the iterator. This generates C-ordered
+ # permutation of chunk indices.
  dim_ranges = [range(dim_len) for dim_len in self.dim_chunks]
  self._ranges = itertools.product(*dim_ranges)
  self._idx = 0
 
  def __iter__(self) -> "ChunkIterator":
- """Return the iterator object itself."""
+ """Iteration context."""
  return self
 
  def __len__(self) -> int:
- """Return the total number of chunks."""
+ """Get total number of chunks."""
  return self.num_chunks
 
- def __next__(self) -> tuple[slice, ...]:
- """Yield the next set of chunk slices.
-
- Returns:
- A tuple of slice objects for each dimension.
-
- Raises:
- StopIteration: When all chunks have been iterated over.
- """
- if self._idx < self.num_chunks:
+ def __next__(self) -> dict[str, slice]:
+ """Iteration logic."""
+ if self._idx <= self.num_chunks:
+ # We build slices here. It is dimension agnostic
  current_start = next(self._ranges)
 
+ # TODO (Dmitriy Repin): Enhance ChunkIterator to make the last slice, if needed, smaller
+ # https://github.com/TGSAI/mdio-python/issues/586
  start_indices = tuple(
  dim * chunk for dim, chunk in zip(current_start, self.len_chunks, strict=True)
  )
@@ -88,7 +93,17 @@ def __next__(self) -> tuple[slice, ...]:
  slice(start, stop) for start, stop in zip(start_indices, stop_indices, strict=True)
  )
 
+ if self.dims: # noqa SIM108
+ # Example
+ # {"inline":slice(3,6,None), "crossline":slice(0,4,None), "depth":slice(0,5,None)}
+ result = dict(zip(self.dims, slices, strict=False))
+ else:
+ # Example
+ # (slice(3,6,None), slice(0,4,None), slice(0,5,None))
+ result = slices
+
  self._idx += 1
- return slices
+
+ return result
 
  raise StopIteration

src/mdio/core/storage_location.py

@@ -0,0 +1,87 @@
+"""StorageLocation class for managing local and cloud storage locations."""
+
+from pathlib import Path
+from typing import Any
+
+import fsspec
+
+
+# TODO(Dmitriy Repin): Reuse fsspec functions for some methods we implemented here
+# https://github.com/TGSAI/mdio-python/issues/597
+class StorageLocation:
+ """A class to represent a local or cloud storage location for SEG-Y or MDIO files.
+
+ This class abstracts the storage location, allowing for both local file paths and
+ cloud storage URIs (e.g., S3, GCS). It uses fsspec to check existence and manage options.
+ Note, we do not want to make it a dataclass because we want the uri and the options to
+ be read-only immutable properties.
+
+ uri: The URI of the storage location (e.g., '/path/to/file', 'file:///path/to/file',
+ 's3://bucket/path', 'gs://bucket/path').
+ options: Optional dictionary of options for the cloud, such as credentials.
+
+ """
+
+ def __init__(self, uri: str = "", options: dict[str, Any] = None):
+ self._uri = uri
+ self._options = options or {}
+ self._fs = None
+
+ if uri.startswith(("s3://", "gs://")):
+ return
+
+ if uri.startswith("file://"):
+ self._uri = self._uri.removeprefix("file://")
+ # For local paths, ensure they are absolute and resolved
+ self._uri = str(Path(self._uri).resolve())
+ return
+
+ @property
+ def uri(self) -> str:
+ """Get the URI (read-only)."""
+ return self._uri
+
+ @property
+ def options(self) -> dict[str, Any]:
+ """Get the options (read-only)."""
+ # Return a copy to prevent external modification
+ return self._options.copy()
+
+ @property
+ def _filesystem(self) -> fsspec.AbstractFileSystem:
+ """Get the fsspec filesystem instance for this storage location."""
+ if self._fs is None:
+ self._fs = fsspec.filesystem(self._protocol, **self._options)
+ return self._fs
+
+ @property
+ def _path(self) -> str:
+ """Extract the path portion from the URI."""
+ if "://" in self._uri:
+ return self._uri.split("://", 1)[1]
+ return self._uri # For local paths without file:// prefix
+
+ @property
+ def _protocol(self) -> str:
+ """Extract the protocol/scheme from the URI."""
+ if "://" in self._uri:
+ return self._uri.split("://", 1)[0]
+ return "file" # Default to file protocol
+
+ def exists(self) -> bool:
+ """Check if the storage location exists using fsspec."""
+ try:
+ return self._filesystem.exists(self._path)
+ except Exception as e:
+ # Log the error and return False for safety
+ # In a production environment, you might want to use proper logging
+ print(f"Error checking existence of {self._uri}: {e}")
+ return False
+
+ def __str__(self) -> str:
+ """String representation of the storage location."""
+ return self._uri
+
+ def __repr__(self) -> str:
+ """Developer representation of the storage location."""
+ return f"StorageLocation(uri='{self._uri}', options={self._options})"

src/mdio/api/convenience.py

@@ -79,7 +79,11 @@ def copy_mdio( # noqa: PLR0913
 
  writer.live_mask[:] = reader.live_mask[:]
 
- iterator = ChunkIterator(reader._traces, chunk_samples=False)
+ shape = reader._traces.shape
+ chunks = reader._traces.chunks
+ chunks = chunks[:-1] + (shape[-1],) # don't chunk samples
+
+ iterator = ChunkIterator(shape=shape, chunks=chunks)
  progress = tqdm(iterator, unit="block")
  progress.set_description(desc=f"Copying data for '{access_pattern=}'")
  for slice_ in progress:
@@ -177,7 +181,10 @@ def create_rechunk_plan(
 
  n_dimension = len(data_array.shape)
  dummy_array = zarr.empty(shape=data_array.shape, chunks=(MAX_BUFFER,) * n_dimension)
- iterator = ChunkIterator(dummy_array)
+
+ shape = dummy_array.shape
+ chunks = dummy_array.chunks
+ iterator = ChunkIterator(shape=shape, chunks=chunks)
 
  return metadata_arrs, data_arrs, live_mask, iterator
 

src/mdio/converters/numpy.py

@@ -7,7 +7,6 @@
 import numpy as np
 
 from mdio.api.accessor import MDIOWriter
-from mdio.converters.segy import get_compressor
 from mdio.core.dimension import Dimension
 from mdio.core.factory import MDIOCreateConfig
 from mdio.core.factory import MDIOVariableConfig
@@ -137,6 +136,11 @@ def numpy_to_mdio( # noqa: PLR0913
  suffix = [str(idx) for idx, value in enumerate(suffix) if value is not None]
  suffix = "".join(suffix)
 
+ # TODO(Dmitrit Repin): Implement Numpy converted in MDIO v1
+ # https://github.com/TGSAI/mdio-python/issues/596
+ def get_compressor(lossless: bool, tolerance: float) -> list[str]:
+ pass
+
  compressors = get_compressor(lossless, compression_tolerance)
  mdio_var = MDIOVariableConfig(
  name=f"chunked_{suffix}",

src/mdio/converters/type_converter.py

@@ -0,0 +1,85 @@
+"""A module for converting numpy dtypes to MDIO scalar and structured types."""
+
+from numpy import dtype as np_dtype
+
+from mdio.schemas.dtype import ScalarType
+from mdio.schemas.dtype import StructuredField
+from mdio.schemas.dtype import StructuredType
+
+
+def to_scalar_type(data_type: np_dtype) -> ScalarType:
+ """Convert numpy dtype to MDIO ScalarType.
+
+ Out of the 24 built-in numpy scalar type objects
+ (see https://numpy.org/doc/stable/reference/arrays.dtypes.html)
+ this function supports only a limited subset:
+ ScalarType.INT8 <-> int8
+ ScalarType.INT16 <-> int16
+ ScalarType.INT32 <-> int32
+ ScalarType.INT64 <-> int64
+ ScalarType.UINT8 <-> uint8
+ ScalarType.UINT16 <-> uint16
+ ScalarType.UINT32 <-> uint32
+ ScalarType.UINT64 <-> uint64
+ ScalarType.FLOAT32 <-> float32
+ ScalarType.FLOAT64 <-> float64
+ ScalarType.COMPLEX64 <-> complex64
+ ScalarType.COMPLEX128 <-> complex128
+ ScalarType.BOOL <-> bool
+
+ Args:
+ data_type: numpy dtype to convert
+
+ Returns:
+ ScalarType: corresponding MDIO scalar type
+
+ Raises:
+ ValueError: if dtype is not supported
+ """
+ try:
+ return ScalarType(data_type.name)
+ except ValueError as exc:
+ err = f"Unsupported numpy dtype '{data_type.name}' for conversion to ScalarType."
+ raise ValueError(err) from exc
+
+
+def to_structured_type(data_type: np_dtype) -> StructuredType:
+ """Convert numpy dtype to MDIO StructuredType.
+
+ This function supports only a limited subset of structured types.
+ In particular:
+ It does not support nested structured types.
+ It supports fields of only 13 out of 24 built-in numpy scalar types.
+ (see `to_scalar_type` for details).
+
+ Args:
+ data_type: numpy dtype to convert
+
+ Returns:
+ StructuredType: corresponding MDIO structured type
+
+ Raises:
+ ValueError: if dtype is not structured or has no fields
+
+ """
+ if data_type is None or len(data_type.names or []) == 0:
+ err = "None or empty dtype provided, cannot convert to StructuredType."
+ raise ValueError(err)
+
+ fields = []
+ for field_name in data_type.names:
+ field_dtype = data_type.fields[field_name][0]
+ scalar_type = to_scalar_type(field_dtype)
+ structured_field = StructuredField(name=field_name, format=scalar_type)
+ fields.append(structured_field)
+ return StructuredType(fields=fields)
+
+
+def to_numpy_dtype(data_type: ScalarType | StructuredType) -> np_dtype:
+ """Get the numpy dtype for a variable."""
+ if isinstance(data_type, ScalarType):
+ return np_dtype(data_type.value)
+ if isinstance(data_type, StructuredType):
+ return np_dtype([(f.name, f.format.value) for f in data_type.fields])
+ msg = f"Expected ScalarType or StructuredType, got '{type(data_type).__name__}'"
+ raise ValueError(msg)

src/mdio/core/indexing.py

@@ -4,78 +4,83 @@
 from math import ceil
 
 import numpy as np
-from zarr import Array
 
 
 class ChunkIterator:
- """Iterator for traversing a Zarr array in chunks.
+ """Chunk iterator for multi-dimensional arrays.
 
- This iterator yields tuples of slices corresponding to the chunk boundaries of a Zarr array.
- It supports chunking all dimensions or taking the full extent of the last dimension.
+ This iterator takes an array shape and chunks and every time it is iterated, it returns
+ a dictionary (if dimensions are provided) or a tuple of slices that align with
+ chunk boundaries. When dimensions are provided, they are used as the dictionary keys.
 
  Args:
- array: The Zarr array to iterate, providing shape and chunk sizes.
- chunk_samples: If True, chunks all dimensions. If False, takes the full extent of the
- last dimension. Defaults to True.
-
-
- Example:
- >>> import zarr
- >>> arr = zarr.array(np.zeros((10, 20)), chunks=(3, 4))
- >>> it = ChunkIterator(arr)
- >>> for slices in it:
- ... print(slices)
- (slice(0, 3, None), slice(0, 4, None))
- (slice(0, 3, None), slice(4, 8, None))
- ...
- >>> it = ChunkIterator(arr, chunk_samples=False)
- >>> for slices in it:
- ... print(slices)
- (slice(0, 3, None), slice(0, 20, None))
- (slice(3, 6, None), slice(0, 20, None))
- ...
+ shape: The shape of the array.
+ chunks: The chunk sizes for each dimension.
+ dim_names: The names of the array dimensions, to be used with DataArray.isel().
+ If the dim_names are not provided, a tuple of the slices will be returned.
+
+ Attributes: # noqa: DOC602
+ arr_shape: Shape of the array.
+ len_chunks: Length of chunks in each dimension.
+ dim_chunks: Number of chunks in each dimension.
+ num_chunks: Total number of chunks.
+
+ Examples:
+ >> chunks = (3, 4, 5)
+ >> shape = (5, 11, 19)
+ >> dims = ["inline", "crossline", "depth"]
+ >>
+ >> iter = ChunkIterator(shape=shape, chunks=chunks, dim_names=dims)
+ >> for i in range(13):
+ >> region = iter.__next__()
+ >> print(region)
+ { "inline": slice(3,6, None), "crossline": slice(0,4, None), "depth": slice(0,5, None) }
+
+ >> iter = ChunkIterator(shape=shape, chunks=chunks, dim_names=None)
+ >> for i in range(13):
+ >> region = iter.__next__()
+ >> print(region)
+ (slice(3,6,None), slice(0,4,None), slice(0,5,None))
  """
 
- def __init__(self, array: Array, chunk_samples: bool = True):
- self.arr_shape = array.shape
- self.len_chunks = array.chunks
-
- # If chunk_samples is False, set the last dimension's chunk size to its full extent
- if not chunk_samples:
- self.len_chunks = self.len_chunks[:-1] + (self.arr_shape[-1],)
-
- # Calculate the number of chunks per dimension
- self.dim_chunks = [
- ceil(len_dim / chunk)
- for len_dim, chunk in zip(self.arr_shape, self.len_chunks, strict=True)
- ]
+ def __init__(
+ self, shape: tuple[int, ...], chunks: tuple[int, ...], dim_names: tuple[str, ...] = None
+ ):
+ self.arr_shape = tuple(shape) # Deep copy to ensure immutability
+ self.len_chunks = tuple(chunks) # Deep copy to ensure immutability
+ self.dims = dim_names
+
+ # Compute number of chunks per dimension, and total number of chunks
+ self.dim_chunks = tuple(
+ [
+ ceil(len_dim / chunk)
+ for len_dim, chunk in zip(self.arr_shape, self.len_chunks, strict=True)
+ ]
+ )
  self.num_chunks = np.prod(self.dim_chunks)
 
- # Set up chunk index combinations using ranges for each dimension
+ # Under the hood stuff for the iterator. This generates C-ordered
+ # permutation of chunk indices.
  dim_ranges = [range(dim_len) for dim_len in self.dim_chunks]
  self._ranges = itertools.product(*dim_ranges)
  self._idx = 0
 
  def __iter__(self) -> "ChunkIterator":
- """Return the iterator object itself."""
+ """Iteration context."""
  return self
 
  def __len__(self) -> int:
- """Return the total number of chunks."""
+ """Get total number of chunks."""
  return self.num_chunks
 
- def __next__(self) -> tuple[slice, ...]:
- """Yield the next set of chunk slices.
-
- Returns:
- A tuple of slice objects for each dimension.
-
- Raises:
- StopIteration: When all chunks have been iterated over.
- """
- if self._idx < self.num_chunks:
+ def __next__(self) -> dict[str, slice]:
+ """Iteration logic."""
+ if self._idx <= self.num_chunks:
+ # We build slices here. It is dimension agnostic
  current_start = next(self._ranges)
 
+ # TODO (Dmitriy Repin): Enhance ChunkIterator to make the last slice, if needed, smaller
+ # https://github.com/TGSAI/mdio-python/issues/586
  start_indices = tuple(
  dim * chunk for dim, chunk in zip(current_start, self.len_chunks, strict=True)
  )
@@ -88,7 +93,17 @@ def __next__(self) -> tuple[slice, ...]:
  slice(start, stop) for start, stop in zip(start_indices, stop_indices, strict=True)
  )
 
+ if self.dims: # noqa SIM108
+ # Example
+ # {"inline":slice(3,6,None), "crossline":slice(0,4,None), "depth":slice(0,5,None)}
+ result = dict(zip(self.dims, slices, strict=False))
+ else:
+ # Example
+ # (slice(3,6,None), slice(0,4,None), slice(0,5,None))
+ result = slices
+
  self._idx += 1
- return slices
+
+ return result
 
  raise StopIteration

src/mdio/core/storage_location.py

@@ -0,0 +1,87 @@
+"""StorageLocation class for managing local and cloud storage locations."""
+
+from pathlib import Path
+from typing import Any
+
+import fsspec
+
+
+# TODO(Dmitriy Repin): Reuse fsspec functions for some methods we implemented here
+# https://github.com/TGSAI/mdio-python/issues/597
+class StorageLocation:
+ """A class to represent a local or cloud storage location for SEG-Y or MDIO files.
+
+ This class abstracts the storage location, allowing for both local file paths and
+ cloud storage URIs (e.g., S3, GCS). It uses fsspec to check existence and manage options.
+ Note, we do not want to make it a dataclass because we want the uri and the options to
+ be read-only immutable properties.
+
+ uri: The URI of the storage location (e.g., '/path/to/file', 'file:///path/to/file',
+ 's3://bucket/path', 'gs://bucket/path').
+ options: Optional dictionary of options for the cloud, such as credentials.
+
+ """
+
+ def __init__(self, uri: str = "", options: dict[str, Any] = None):
+ self._uri = uri
+ self._options = options or {}
+ self._fs = None
+
+ if uri.startswith(("s3://", "gs://")):
+ return
+
+ if uri.startswith("file://"):
+ self._uri = self._uri.removeprefix("file://")
+ # For local paths, ensure they are absolute and resolved
+ self._uri = str(Path(self._uri).resolve())
+ return
+
+ @property
+ def uri(self) -> str:
+ """Get the URI (read-only)."""
+ return self._uri
+
+ @property
+ def options(self) -> dict[str, Any]:
+ """Get the options (read-only)."""
+ # Return a copy to prevent external modification
+ return self._options.copy()
+
+ @property
+ def _filesystem(self) -> fsspec.AbstractFileSystem:
+ """Get the fsspec filesystem instance for this storage location."""
+ if self._fs is None:
+ self._fs = fsspec.filesystem(self._protocol, **self._options)
+ return self._fs
+
+ @property
+ def _path(self) -> str:
+ """Extract the path portion from the URI."""
+ if "://" in self._uri:
+ return self._uri.split("://", 1)[1]
+ return self._uri # For local paths without file:// prefix
+
+ @property
+ def _protocol(self) -> str:
+ """Extract the protocol/scheme from the URI."""
+ if "://" in self._uri:
+ return self._uri.split("://", 1)[0]
+ return "file" # Default to file protocol
+
+ def exists(self) -> bool:
+ """Check if the storage location exists using fsspec."""
+ try:
+ return self._filesystem.exists(self._path)
+ except Exception as e:
+ # Log the error and return False for safety
+ # In a production environment, you might want to use proper logging
+ print(f"Error checking existence of {self._uri}: {e}")
+ return False
+
+ def __str__(self) -> str:
+ """String representation of the storage location."""
+ return self._uri
+
+ def __repr__(self) -> str:
+ """Developer representation of the storage location."""
+ return f"StorageLocation(uri='{self._uri}', options={self._options})"

src/mdio/api/convenience.py

@@ -79,7 +79,11 @@ def copy_mdio( # noqa: PLR0913
 
  writer.live_mask[:] = reader.live_mask[:]
 
- iterator = ChunkIterator(reader._traces, chunk_samples=False)
+ shape = reader._traces.shape
+ chunks = reader._traces.chunks
+ chunks = chunks[:-1] + (shape[-1],) # don't chunk samples
+
+ iterator = ChunkIterator(shape=shape, chunks=chunks)
  progress = tqdm(iterator, unit="block")
  progress.set_description(desc=f"Copying data for '{access_pattern=}'")
  for slice_ in progress:
@@ -177,7 +181,10 @@ def create_rechunk_plan(
 
  n_dimension = len(data_array.shape)
  dummy_array = zarr.empty(shape=data_array.shape, chunks=(MAX_BUFFER,) * n_dimension)
- iterator = ChunkIterator(dummy_array)
+
+ shape = dummy_array.shape
+ chunks = dummy_array.chunks
+ iterator = ChunkIterator(shape=shape, chunks=chunks)
 
  return metadata_arrs, data_arrs, live_mask, iterator
 

src/mdio/converters/numpy.py

@@ -7,7 +7,6 @@
 import numpy as np
 
 from mdio.api.accessor import MDIOWriter
-from mdio.converters.segy import get_compressor
 from mdio.core.dimension import Dimension
 from mdio.core.factory import MDIOCreateConfig
 from mdio.core.factory import MDIOVariableConfig
@@ -137,6 +136,11 @@ def numpy_to_mdio( # noqa: PLR0913
  suffix = [str(idx) for idx, value in enumerate(suffix) if value is not None]
  suffix = "".join(suffix)
 
+ # TODO(Dmitrit Repin): Implement Numpy converted in MDIO v1
+ # https://github.com/TGSAI/mdio-python/issues/596
+ def get_compressor(lossless: bool, tolerance: float) -> list[str]:
+ pass
+
  compressors = get_compressor(lossless, compression_tolerance)
  mdio_var = MDIOVariableConfig(
  name=f"chunked_{suffix}",

src/mdio/converters/type_converter.py

@@ -0,0 +1,85 @@
+"""A module for converting numpy dtypes to MDIO scalar and structured types."""
+
+from numpy import dtype as np_dtype
+
+from mdio.schemas.dtype import ScalarType
+from mdio.schemas.dtype import StructuredField
+from mdio.schemas.dtype import StructuredType
+
+
+def to_scalar_type(data_type: np_dtype) -> ScalarType:
+ """Convert numpy dtype to MDIO ScalarType.
+
+ Out of the 24 built-in numpy scalar type objects
+ (see https://numpy.org/doc/stable/reference/arrays.dtypes.html)
+ this function supports only a limited subset:
+ ScalarType.INT8 <-> int8
+ ScalarType.INT16 <-> int16
+ ScalarType.INT32 <-> int32
+ ScalarType.INT64 <-> int64
+ ScalarType.UINT8 <-> uint8
+ ScalarType.UINT16 <-> uint16
+ ScalarType.UINT32 <-> uint32
+ ScalarType.UINT64 <-> uint64
+ ScalarType.FLOAT32 <-> float32
+ ScalarType.FLOAT64 <-> float64
+ ScalarType.COMPLEX64 <-> complex64
+ ScalarType.COMPLEX128 <-> complex128
+ ScalarType.BOOL <-> bool
+
+ Args:
+ data_type: numpy dtype to convert
+
+ Returns:
+ ScalarType: corresponding MDIO scalar type
+
+ Raises:
+ ValueError: if dtype is not supported
+ """
+ try:
+ return ScalarType(data_type.name)
+ except ValueError as exc:
+ err = f"Unsupported numpy dtype '{data_type.name}' for conversion to ScalarType."
+ raise ValueError(err) from exc
+
+
+def to_structured_type(data_type: np_dtype) -> StructuredType:
+ """Convert numpy dtype to MDIO StructuredType.
+
+ This function supports only a limited subset of structured types.
+ In particular:
+ It does not support nested structured types.
+ It supports fields of only 13 out of 24 built-in numpy scalar types.
+ (see `to_scalar_type` for details).
+
+ Args:
+ data_type: numpy dtype to convert
+
+ Returns:
+ StructuredType: corresponding MDIO structured type
+
+ Raises:
+ ValueError: if dtype is not structured or has no fields
+
+ """
+ if data_type is None or len(data_type.names or []) == 0:
+ err = "None or empty dtype provided, cannot convert to StructuredType."
+ raise ValueError(err)
+
+ fields = []
+ for field_name in data_type.names:
+ field_dtype = data_type.fields[field_name][0]
+ scalar_type = to_scalar_type(field_dtype)
+ structured_field = StructuredField(name=field_name, format=scalar_type)
+ fields.append(structured_field)
+ return StructuredType(fields=fields)
+
+
+def to_numpy_dtype(data_type: ScalarType | StructuredType) -> np_dtype:
+ """Get the numpy dtype for a variable."""
+ if isinstance(data_type, ScalarType):
+ return np_dtype(data_type.value)
+ if isinstance(data_type, StructuredType):
+ return np_dtype([(f.name, f.format.value) for f in data_type.fields])
+ msg = f"Expected ScalarType or StructuredType, got '{type(data_type).__name__}'"
+ raise ValueError(msg)

src/mdio/core/indexing.py

@@ -4,78 +4,83 @@
 from math import ceil
 
 import numpy as np
-from zarr import Array
 
 
 class ChunkIterator:
- """Iterator for traversing a Zarr array in chunks.
+ """Chunk iterator for multi-dimensional arrays.
 
- This iterator yields tuples of slices corresponding to the chunk boundaries of a Zarr array.
- It supports chunking all dimensions or taking the full extent of the last dimension.
+ This iterator takes an array shape and chunks and every time it is iterated, it returns
+ a dictionary (if dimensions are provided) or a tuple of slices that align with
+ chunk boundaries. When dimensions are provided, they are used as the dictionary keys.
 
  Args:
- array: The Zarr array to iterate, providing shape and chunk sizes.
- chunk_samples: If True, chunks all dimensions. If False, takes the full extent of the
- last dimension. Defaults to True.
-
-
- Example:
- >>> import zarr
- >>> arr = zarr.array(np.zeros((10, 20)), chunks=(3, 4))
- >>> it = ChunkIterator(arr)
- >>> for slices in it:
- ... print(slices)
- (slice(0, 3, None), slice(0, 4, None))
- (slice(0, 3, None), slice(4, 8, None))
- ...
- >>> it = ChunkIterator(arr, chunk_samples=False)
- >>> for slices in it:
- ... print(slices)
- (slice(0, 3, None), slice(0, 20, None))
- (slice(3, 6, None), slice(0, 20, None))
- ...
+ shape: The shape of the array.
+ chunks: The chunk sizes for each dimension.
+ dim_names: The names of the array dimensions, to be used with DataArray.isel().
+ If the dim_names are not provided, a tuple of the slices will be returned.
+
+ Attributes: # noqa: DOC602
+ arr_shape: Shape of the array.
+ len_chunks: Length of chunks in each dimension.
+ dim_chunks: Number of chunks in each dimension.
+ num_chunks: Total number of chunks.
+
+ Examples:
+ >> chunks = (3, 4, 5)
+ >> shape = (5, 11, 19)
+ >> dims = ["inline", "crossline", "depth"]
+ >>
+ >> iter = ChunkIterator(shape=shape, chunks=chunks, dim_names=dims)
+ >> for i in range(13):
+ >> region = iter.__next__()
+ >> print(region)
+ { "inline": slice(3,6, None), "crossline": slice(0,4, None), "depth": slice(0,5, None) }
+
+ >> iter = ChunkIterator(shape=shape, chunks=chunks, dim_names=None)
+ >> for i in range(13):
+ >> region = iter.__next__()
+ >> print(region)
+ (slice(3,6,None), slice(0,4,None), slice(0,5,None))
  """
 
- def __init__(self, array: Array, chunk_samples: bool = True):
- self.arr_shape = array.shape
- self.len_chunks = array.chunks
-
- # If chunk_samples is False, set the last dimension's chunk size to its full extent
- if not chunk_samples:
- self.len_chunks = self.len_chunks[:-1] + (self.arr_shape[-1],)
-
- # Calculate the number of chunks per dimension
- self.dim_chunks = [
- ceil(len_dim / chunk)
- for len_dim, chunk in zip(self.arr_shape, self.len_chunks, strict=True)
- ]
+ def __init__(
+ self, shape: tuple[int, ...], chunks: tuple[int, ...], dim_names: tuple[str, ...] = None
+ ):
+ self.arr_shape = tuple(shape) # Deep copy to ensure immutability
+ self.len_chunks = tuple(chunks) # Deep copy to ensure immutability
+ self.dims = dim_names
+
+ # Compute number of chunks per dimension, and total number of chunks
+ self.dim_chunks = tuple(
+ [
+ ceil(len_dim / chunk)
+ for len_dim, chunk in zip(self.arr_shape, self.len_chunks, strict=True)
+ ]
+ )
  self.num_chunks = np.prod(self.dim_chunks)
 
- # Set up chunk index combinations using ranges for each dimension
+ # Under the hood stuff for the iterator. This generates C-ordered
+ # permutation of chunk indices.
  dim_ranges = [range(dim_len) for dim_len in self.dim_chunks]
  self._ranges = itertools.product(*dim_ranges)
  self._idx = 0
 
  def __iter__(self) -> "ChunkIterator":
- """Return the iterator object itself."""
+ """Iteration context."""
  return self
 
  def __len__(self) -> int:
- """Return the total number of chunks."""
+ """Get total number of chunks."""
  return self.num_chunks
 
- def __next__(self) -> tuple[slice, ...]:
- """Yield the next set of chunk slices.
-
- Returns:
- A tuple of slice objects for each dimension.
-
- Raises:
- StopIteration: When all chunks have been iterated over.
- """
- if self._idx < self.num_chunks:
+ def __next__(self) -> dict[str, slice]:
+ """Iteration logic."""
+ if self._idx <= self.num_chunks:
+ # We build slices here. It is dimension agnostic
  current_start = next(self._ranges)
 
+ # TODO (Dmitriy Repin): Enhance ChunkIterator to make the last slice, if needed, smaller
+ # https://github.com/TGSAI/mdio-python/issues/586
  start_indices = tuple(
  dim * chunk for dim, chunk in zip(current_start, self.len_chunks, strict=True)
  )
@@ -88,7 +93,17 @@ def __next__(self) -> tuple[slice, ...]:
  slice(start, stop) for start, stop in zip(start_indices, stop_indices, strict=True)
  )
 
+ if self.dims: # noqa SIM108
+ # Example
+ # {"inline":slice(3,6,None), "crossline":slice(0,4,None), "depth":slice(0,5,None)}
+ result = dict(zip(self.dims, slices, strict=False))
+ else:
+ # Example
+ # (slice(3,6,None), slice(0,4,None), slice(0,5,None))
+ result = slices
+
  self._idx += 1
- return slices
+
+ return result
 
  raise StopIteration

src/mdio/core/storage_location.py

@@ -0,0 +1,87 @@
+"""StorageLocation class for managing local and cloud storage locations."""
+
+from pathlib import Path
+from typing import Any
+
+import fsspec
+
+
+# TODO(Dmitriy Repin): Reuse fsspec functions for some methods we implemented here
+# https://github.com/TGSAI/mdio-python/issues/597
+class StorageLocation:
+ """A class to represent a local or cloud storage location for SEG-Y or MDIO files.
+
+ This class abstracts the storage location, allowing for both local file paths and
+ cloud storage URIs (e.g., S3, GCS). It uses fsspec to check existence and manage options.
+ Note, we do not want to make it a dataclass because we want the uri and the options to
+ be read-only immutable properties.
+
+ uri: The URI of the storage location (e.g., '/path/to/file', 'file:///path/to/file',
+ 's3://bucket/path', 'gs://bucket/path').
+ options: Optional dictionary of options for the cloud, such as credentials.
+
+ """
+
+ def __init__(self, uri: str = "", options: dict[str, Any] = None):
+ self._uri = uri
+ self._options = options or {}
+ self._fs = None
+
+ if uri.startswith(("s3://", "gs://")):
+ return
+
+ if uri.startswith("file://"):
+ self._uri = self._uri.removeprefix("file://")
+ # For local paths, ensure they are absolute and resolved
+ self._uri = str(Path(self._uri).resolve())
+ return
+
+ @property
+ def uri(self) -> str:
+ """Get the URI (read-only)."""
+ return self._uri
+
+ @property
+ def options(self) -> dict[str, Any]:
+ """Get the options (read-only)."""
+ # Return a copy to prevent external modification
+ return self._options.copy()
+
+ @property
+ def _filesystem(self) -> fsspec.AbstractFileSystem:
+ """Get the fsspec filesystem instance for this storage location."""
+ if self._fs is None:
+ self._fs = fsspec.filesystem(self._protocol, **self._options)
+ return self._fs
+
+ @property
+ def _path(self) -> str:
+ """Extract the path portion from the URI."""
+ if "://" in self._uri:
+ return self._uri.split("://", 1)[1]
+ return self._uri # For local paths without file:// prefix
+
+ @property
+ def _protocol(self) -> str:
+ """Extract the protocol/scheme from the URI."""
+ if "://" in self._uri:
+ return self._uri.split("://", 1)[0]
+ return "file" # Default to file protocol
+
+ def exists(self) -> bool:
+ """Check if the storage location exists using fsspec."""
+ try:
+ return self._filesystem.exists(self._path)
+ except Exception as e:
+ # Log the error and return False for safety
+ # In a production environment, you might want to use proper logging
+ print(f"Error checking existence of {self._uri}: {e}")
+ return False
+
+ def __str__(self) -> str:
+ """String representation of the storage location."""
+ return self._uri
+
+ def __repr__(self) -> str:
+ """Developer representation of the storage location."""
+ return f"StorageLocation(uri='{self._uri}', options={self._options})"

src/mdio/api/convenience.py

@@ -79,7 +79,11 @@ def copy_mdio( # noqa: PLR0913
 
  writer.live_mask[:] = reader.live_mask[:]
 
- iterator = ChunkIterator(reader._traces, chunk_samples=False)
+ shape = reader._traces.shape
+ chunks = reader._traces.chunks
+ chunks = chunks[:-1] + (shape[-1],) # don't chunk samples
+
+ iterator = ChunkIterator(shape=shape, chunks=chunks)
  progress = tqdm(iterator, unit="block")
  progress.set_description(desc=f"Copying data for '{access_pattern=}'")
  for slice_ in progress:
@@ -177,7 +181,10 @@ def create_rechunk_plan(
 
  n_dimension = len(data_array.shape)
  dummy_array = zarr.empty(shape=data_array.shape, chunks=(MAX_BUFFER,) * n_dimension)
- iterator = ChunkIterator(dummy_array)
+
+ shape = dummy_array.shape
+ chunks = dummy_array.chunks
+ iterator = ChunkIterator(shape=shape, chunks=chunks)
 
  return metadata_arrs, data_arrs, live_mask, iterator
 

src/mdio/converters/numpy.py

@@ -7,7 +7,6 @@
 import numpy as np
 
 from mdio.api.accessor import MDIOWriter
-from mdio.converters.segy import get_compressor
 from mdio.core.dimension import Dimension
 from mdio.core.factory import MDIOCreateConfig
 from mdio.core.factory import MDIOVariableConfig
@@ -137,6 +136,11 @@ def numpy_to_mdio( # noqa: PLR0913
  suffix = [str(idx) for idx, value in enumerate(suffix) if value is not None]
  suffix = "".join(suffix)
 
+ # TODO(Dmitrit Repin): Implement Numpy converted in MDIO v1
+ # https://github.com/TGSAI/mdio-python/issues/596
+ def get_compressor(lossless: bool, tolerance: float) -> list[str]:
+ pass
+
  compressors = get_compressor(lossless, compression_tolerance)
  mdio_var = MDIOVariableConfig(
  name=f"chunked_{suffix}",

src/mdio/converters/type_converter.py

@@ -0,0 +1,85 @@
+"""A module for converting numpy dtypes to MDIO scalar and structured types."""
+
+from numpy import dtype as np_dtype
+
+from mdio.schemas.dtype import ScalarType
+from mdio.schemas.dtype import StructuredField
+from mdio.schemas.dtype import StructuredType
+
+
+def to_scalar_type(data_type: np_dtype) -> ScalarType:
+ """Convert numpy dtype to MDIO ScalarType.
+
+ Out of the 24 built-in numpy scalar type objects
+ (see https://numpy.org/doc/stable/reference/arrays.dtypes.html)
+ this function supports only a limited subset:
+ ScalarType.INT8 <-> int8
+ ScalarType.INT16 <-> int16
+ ScalarType.INT32 <-> int32
+ ScalarType.INT64 <-> int64
+ ScalarType.UINT8 <-> uint8
+ ScalarType.UINT16 <-> uint16
+ ScalarType.UINT32 <-> uint32
+ ScalarType.UINT64 <-> uint64
+ ScalarType.FLOAT32 <-> float32
+ ScalarType.FLOAT64 <-> float64
+ ScalarType.COMPLEX64 <-> complex64
+ ScalarType.COMPLEX128 <-> complex128
+ ScalarType.BOOL <-> bool
+
+ Args:
+ data_type: numpy dtype to convert
+
+ Returns:
+ ScalarType: corresponding MDIO scalar type
+
+ Raises:
+ ValueError: if dtype is not supported
+ """
+ try:
+ return ScalarType(data_type.name)
+ except ValueError as exc:
+ err = f"Unsupported numpy dtype '{data_type.name}' for conversion to ScalarType."
+ raise ValueError(err) from exc
+
+
+def to_structured_type(data_type: np_dtype) -> StructuredType:
+ """Convert numpy dtype to MDIO StructuredType.
+
+ This function supports only a limited subset of structured types.
+ In particular:
+ It does not support nested structured types.
+ It supports fields of only 13 out of 24 built-in numpy scalar types.
+ (see `to_scalar_type` for details).
+
+ Args:
+ data_type: numpy dtype to convert
+
+ Returns:
+ StructuredType: corresponding MDIO structured type
+
+ Raises:
+ ValueError: if dtype is not structured or has no fields
+
+ """
+ if data_type is None or len(data_type.names or []) == 0:
+ err = "None or empty dtype provided, cannot convert to StructuredType."
+ raise ValueError(err)
+
+ fields = []
+ for field_name in data_type.names:
+ field_dtype = data_type.fields[field_name][0]
+ scalar_type = to_scalar_type(field_dtype)
+ structured_field = StructuredField(name=field_name, format=scalar_type)
+ fields.append(structured_field)
+ return StructuredType(fields=fields)
+
+
+def to_numpy_dtype(data_type: ScalarType | StructuredType) -> np_dtype:
+ """Get the numpy dtype for a variable."""
+ if isinstance(data_type, ScalarType):
+ return np_dtype(data_type.value)
+ if isinstance(data_type, StructuredType):
+ return np_dtype([(f.name, f.format.value) for f in data_type.fields])
+ msg = f"Expected ScalarType or StructuredType, got '{type(data_type).__name__}'"
+ raise ValueError(msg)

src/mdio/core/indexing.py

@@ -4,78 +4,83 @@
 from math import ceil
 
 import numpy as np
-from zarr import Array
 
 
 class ChunkIterator:
- """Iterator for traversing a Zarr array in chunks.
+ """Chunk iterator for multi-dimensional arrays.
 
- This iterator yields tuples of slices corresponding to the chunk boundaries of a Zarr array.
- It supports chunking all dimensions or taking the full extent of the last dimension.
+ This iterator takes an array shape and chunks and every time it is iterated, it returns
+ a dictionary (if dimensions are provided) or a tuple of slices that align with
+ chunk boundaries. When dimensions are provided, they are used as the dictionary keys.
 
  Args:
- array: The Zarr array to iterate, providing shape and chunk sizes.
- chunk_samples: If True, chunks all dimensions. If False, takes the full extent of the
- last dimension. Defaults to True.
-
-
- Example:
- >>> import zarr
- >>> arr = zarr.array(np.zeros((10, 20)), chunks=(3, 4))
- >>> it = ChunkIterator(arr)
- >>> for slices in it:
- ... print(slices)
- (slice(0, 3, None), slice(0, 4, None))
- (slice(0, 3, None), slice(4, 8, None))
- ...
- >>> it = ChunkIterator(arr, chunk_samples=False)
- >>> for slices in it:
- ... print(slices)
- (slice(0, 3, None), slice(0, 20, None))
- (slice(3, 6, None), slice(0, 20, None))
- ...
+ shape: The shape of the array.
+ chunks: The chunk sizes for each dimension.
+ dim_names: The names of the array dimensions, to be used with DataArray.isel().
+ If the dim_names are not provided, a tuple of the slices will be returned.
+
+ Attributes: # noqa: DOC602
+ arr_shape: Shape of the array.
+ len_chunks: Length of chunks in each dimension.
+ dim_chunks: Number of chunks in each dimension.
+ num_chunks: Total number of chunks.
+
+ Examples:
+ >> chunks = (3, 4, 5)
+ >> shape = (5, 11, 19)
+ >> dims = ["inline", "crossline", "depth"]
+ >>
+ >> iter = ChunkIterator(shape=shape, chunks=chunks, dim_names=dims)
+ >> for i in range(13):
+ >> region = iter.__next__()
+ >> print(region)
+ { "inline": slice(3,6, None), "crossline": slice(0,4, None), "depth": slice(0,5, None) }
+
+ >> iter = ChunkIterator(shape=shape, chunks=chunks, dim_names=None)
+ >> for i in range(13):
+ >> region = iter.__next__()
+ >> print(region)
+ (slice(3,6,None), slice(0,4,None), slice(0,5,None))
  """
 
- def __init__(self, array: Array, chunk_samples: bool = True):
- self.arr_shape = array.shape
- self.len_chunks = array.chunks
-
- # If chunk_samples is False, set the last dimension's chunk size to its full extent
- if not chunk_samples:
- self.len_chunks = self.len_chunks[:-1] + (self.arr_shape[-1],)
-
- # Calculate the number of chunks per dimension
- self.dim_chunks = [
- ceil(len_dim / chunk)
- for len_dim, chunk in zip(self.arr_shape, self.len_chunks, strict=True)
- ]
+ def __init__(
+ self, shape: tuple[int, ...], chunks: tuple[int, ...], dim_names: tuple[str, ...] = None
+ ):
+ self.arr_shape = tuple(shape) # Deep copy to ensure immutability
+ self.len_chunks = tuple(chunks) # Deep copy to ensure immutability
+ self.dims = dim_names
+
+ # Compute number of chunks per dimension, and total number of chunks
+ self.dim_chunks = tuple(
+ [
+ ceil(len_dim / chunk)
+ for len_dim, chunk in zip(self.arr_shape, self.len_chunks, strict=True)
+ ]
+ )
  self.num_chunks = np.prod(self.dim_chunks)
 
- # Set up chunk index combinations using ranges for each dimension
+ # Under the hood stuff for the iterator. This generates C-ordered
+ # permutation of chunk indices.
  dim_ranges = [range(dim_len) for dim_len in self.dim_chunks]
  self._ranges = itertools.product(*dim_ranges)
  self._idx = 0
 
  def __iter__(self) -> "ChunkIterator":
- """Return the iterator object itself."""
+ """Iteration context."""
  return self
 
  def __len__(self) -> int:
- """Return the total number of chunks."""
+ """Get total number of chunks."""
  return self.num_chunks
 
- def __next__(self) -> tuple[slice, ...]:
- """Yield the next set of chunk slices.
-
- Returns:
- A tuple of slice objects for each dimension.
-
- Raises:
- StopIteration: When all chunks have been iterated over.
- """
- if self._idx < self.num_chunks:
+ def __next__(self) -> dict[str, slice]:
+ """Iteration logic."""
+ if self._idx <= self.num_chunks:
+ # We build slices here. It is dimension agnostic
  current_start = next(self._ranges)
 
+ # TODO (Dmitriy Repin): Enhance ChunkIterator to make the last slice, if needed, smaller
+ # https://github.com/TGSAI/mdio-python/issues/586
  start_indices = tuple(
  dim * chunk for dim, chunk in zip(current_start, self.len_chunks, strict=True)
  )
@@ -88,7 +93,17 @@ def __next__(self) -> tuple[slice, ...]:
  slice(start, stop) for start, stop in zip(start_indices, stop_indices, strict=True)
  )
 
+ if self.dims: # noqa SIM108
+ # Example
+ # {"inline":slice(3,6,None), "crossline":slice(0,4,None), "depth":slice(0,5,None)}
+ result = dict(zip(self.dims, slices, strict=False))
+ else:
+ # Example
+ # (slice(3,6,None), slice(0,4,None), slice(0,5,None))
+ result = slices
+
  self._idx += 1
- return slices
+
+ return result
 
  raise StopIteration

src/mdio/core/storage_location.py

@@ -0,0 +1,87 @@
+"""StorageLocation class for managing local and cloud storage locations."""
+
+from pathlib import Path
+from typing import Any
+
+import fsspec
+
+
+# TODO(Dmitriy Repin): Reuse fsspec functions for some methods we implemented here
+# https://github.com/TGSAI/mdio-python/issues/597
+class StorageLocation:
+ """A class to represent a local or cloud storage location for SEG-Y or MDIO files.
+
+ This class abstracts the storage location, allowing for both local file paths and
+ cloud storage URIs (e.g., S3, GCS). It uses fsspec to check existence and manage options.
+ Note, we do not want to make it a dataclass because we want the uri and the options to
+ be read-only immutable properties.
+
+ uri: The URI of the storage location (e.g., '/path/to/file', 'file:///path/to/file',
+ 's3://bucket/path', 'gs://bucket/path').
+ options: Optional dictionary of options for the cloud, such as credentials.
+
+ """
+
+ def __init__(self, uri: str = "", options: dict[str, Any] = None):
+ self._uri = uri
+ self._options = options or {}
+ self._fs = None
+
+ if uri.startswith(("s3://", "gs://")):
+ return
+
+ if uri.startswith("file://"):
+ self._uri = self._uri.removeprefix("file://")
+ # For local paths, ensure they are absolute and resolved
+ self._uri = str(Path(self._uri).resolve())
+ return
+
+ @property
+ def uri(self) -> str:
+ """Get the URI (read-only)."""
+ return self._uri
+
+ @property
+ def options(self) -> dict[str, Any]:
+ """Get the options (read-only)."""
+ # Return a copy to prevent external modification
+ return self._options.copy()
+
+ @property
+ def _filesystem(self) -> fsspec.AbstractFileSystem:
+ """Get the fsspec filesystem instance for this storage location."""
+ if self._fs is None:
+ self._fs = fsspec.filesystem(self._protocol, **self._options)
+ return self._fs
+
+ @property
+ def _path(self) -> str:
+ """Extract the path portion from the URI."""
+ if "://" in self._uri:
+ return self._uri.split("://", 1)[1]
+ return self._uri # For local paths without file:// prefix
+
+ @property
+ def _protocol(self) -> str:
+ """Extract the protocol/scheme from the URI."""
+ if "://" in self._uri:
+ return self._uri.split("://", 1)[0]
+ return "file" # Default to file protocol
+
+ def exists(self) -> bool:
+ """Check if the storage location exists using fsspec."""
+ try:
+ return self._filesystem.exists(self._path)
+ except Exception as e:
+ # Log the error and return False for safety
+ # In a production environment, you might want to use proper logging
+ print(f"Error checking existence of {self._uri}: {e}")
+ return False
+
+ def __str__(self) -> str:
+ """String representation of the storage location."""
+ return self._uri
+
+ def __repr__(self) -> str:
+ """Developer representation of the storage location."""
+ return f"StorageLocation(uri='{self._uri}', options={self._options})"

src/mdio/api/convenience.py

@@ -79,7 +79,11 @@ def copy_mdio( # noqa: PLR0913
 
  writer.live_mask[:] = reader.live_mask[:]
 
- iterator = ChunkIterator(reader._traces, chunk_samples=False)
+ shape = reader._traces.shape
+ chunks = reader._traces.chunks
+ chunks = chunks[:-1] + (shape[-1],) # don't chunk samples
+
+ iterator = ChunkIterator(shape=shape, chunks=chunks)
  progress = tqdm(iterator, unit="block")
  progress.set_description(desc=f"Copying data for '{access_pattern=}'")
  for slice_ in progress:
@@ -177,7 +181,10 @@ def create_rechunk_plan(
 
  n_dimension = len(data_array.shape)
  dummy_array = zarr.empty(shape=data_array.shape, chunks=(MAX_BUFFER,) * n_dimension)
- iterator = ChunkIterator(dummy_array)
+
+ shape = dummy_array.shape
+ chunks = dummy_array.chunks
+ iterator = ChunkIterator(shape=shape, chunks=chunks)
 
  return metadata_arrs, data_arrs, live_mask, iterator
 

src/mdio/converters/numpy.py

@@ -7,7 +7,6 @@
 import numpy as np
 
 from mdio.api.accessor import MDIOWriter
-from mdio.converters.segy import get_compressor
 from mdio.core.dimension import Dimension
 from mdio.core.factory import MDIOCreateConfig
 from mdio.core.factory import MDIOVariableConfig
@@ -137,6 +136,11 @@ def numpy_to_mdio( # noqa: PLR0913
  suffix = [str(idx) for idx, value in enumerate(suffix) if value is not None]
  suffix = "".join(suffix)
 
+ # TODO(Dmitrit Repin): Implement Numpy converted in MDIO v1
+ # https://github.com/TGSAI/mdio-python/issues/596
+ def get_compressor(lossless: bool, tolerance: float) -> list[str]:
+ pass
+
  compressors = get_compressor(lossless, compression_tolerance)
  mdio_var = MDIOVariableConfig(
  name=f"chunked_{suffix}",

src/mdio/converters/type_converter.py

@@ -0,0 +1,85 @@
+"""A module for converting numpy dtypes to MDIO scalar and structured types."""
+
+from numpy import dtype as np_dtype
+
+from mdio.schemas.dtype import ScalarType
+from mdio.schemas.dtype import StructuredField
+from mdio.schemas.dtype import StructuredType
+
+
+def to_scalar_type(data_type: np_dtype) -> ScalarType:
+ """Convert numpy dtype to MDIO ScalarType.
+
+ Out of the 24 built-in numpy scalar type objects
+ (see https://numpy.org/doc/stable/reference/arrays.dtypes.html)
+ this function supports only a limited subset:
+ ScalarType.INT8 <-> int8
+ ScalarType.INT16 <-> int16
+ ScalarType.INT32 <-> int32
+ ScalarType.INT64 <-> int64
+ ScalarType.UINT8 <-> uint8
+ ScalarType.UINT16 <-> uint16
+ ScalarType.UINT32 <-> uint32
+ ScalarType.UINT64 <-> uint64
+ ScalarType.FLOAT32 <-> float32
+ ScalarType.FLOAT64 <-> float64
+ ScalarType.COMPLEX64 <-> complex64
+ ScalarType.COMPLEX128 <-> complex128
+ ScalarType.BOOL <-> bool
+
+ Args:
+ data_type: numpy dtype to convert
+
+ Returns:
+ ScalarType: corresponding MDIO scalar type
+
+ Raises:
+ ValueError: if dtype is not supported
+ """
+ try:
+ return ScalarType(data_type.name)
+ except ValueError as exc:
+ err = f"Unsupported numpy dtype '{data_type.name}' for conversion to ScalarType."
+ raise ValueError(err) from exc
+
+
+def to_structured_type(data_type: np_dtype) -> StructuredType:
+ """Convert numpy dtype to MDIO StructuredType.
+
+ This function supports only a limited subset of structured types.
+ In particular:
+ It does not support nested structured types.
+ It supports fields of only 13 out of 24 built-in numpy scalar types.
+ (see `to_scalar_type` for details).
+
+ Args:
+ data_type: numpy dtype to convert
+
+ Returns:
+ StructuredType: corresponding MDIO structured type
+
+ Raises:
+ ValueError: if dtype is not structured or has no fields
+
+ """
+ if data_type is None or len(data_type.names or []) == 0:
+ err = "None or empty dtype provided, cannot convert to StructuredType."
+ raise ValueError(err)
+
+ fields = []
+ for field_name in data_type.names:
+ field_dtype = data_type.fields[field_name][0]
+ scalar_type = to_scalar_type(field_dtype)
+ structured_field = StructuredField(name=field_name, format=scalar_type)
+ fields.append(structured_field)
+ return StructuredType(fields=fields)
+
+
+def to_numpy_dtype(data_type: ScalarType | StructuredType) -> np_dtype:
+ """Get the numpy dtype for a variable."""
+ if isinstance(data_type, ScalarType):
+ return np_dtype(data_type.value)
+ if isinstance(data_type, StructuredType):
+ return np_dtype([(f.name, f.format.value) for f in data_type.fields])
+ msg = f"Expected ScalarType or StructuredType, got '{type(data_type).__name__}'"
+ raise ValueError(msg)

src/mdio/core/indexing.py

@@ -4,78 +4,83 @@
 from math import ceil
 
 import numpy as np
-from zarr import Array
 
 
 class ChunkIterator:
- """Iterator for traversing a Zarr array in chunks.
+ """Chunk iterator for multi-dimensional arrays.
 
- This iterator yields tuples of slices corresponding to the chunk boundaries of a Zarr array.
- It supports chunking all dimensions or taking the full extent of the last dimension.
+ This iterator takes an array shape and chunks and every time it is iterated, it returns
+ a dictionary (if dimensions are provided) or a tuple of slices that align with
+ chunk boundaries. When dimensions are provided, they are used as the dictionary keys.
 
  Args:
- array: The Zarr array to iterate, providing shape and chunk sizes.
- chunk_samples: If True, chunks all dimensions. If False, takes the full extent of the
- last dimension. Defaults to True.
-
-
- Example:
- >>> import zarr
- >>> arr = zarr.array(np.zeros((10, 20)), chunks=(3, 4))
- >>> it = ChunkIterator(arr)
- >>> for slices in it:
- ... print(slices)
- (slice(0, 3, None), slice(0, 4, None))
- (slice(0, 3, None), slice(4, 8, None))
- ...
- >>> it = ChunkIterator(arr, chunk_samples=False)
- >>> for slices in it:
- ... print(slices)
- (slice(0, 3, None), slice(0, 20, None))
- (slice(3, 6, None), slice(0, 20, None))
- ...
+ shape: The shape of the array.
+ chunks: The chunk sizes for each dimension.
+ dim_names: The names of the array dimensions, to be used with DataArray.isel().
+ If the dim_names are not provided, a tuple of the slices will be returned.
+
+ Attributes: # noqa: DOC602
+ arr_shape: Shape of the array.
+ len_chunks: Length of chunks in each dimension.
+ dim_chunks: Number of chunks in each dimension.
+ num_chunks: Total number of chunks.
+
+ Examples:
+ >> chunks = (3, 4, 5)
+ >> shape = (5, 11, 19)
+ >> dims = ["inline", "crossline", "depth"]
+ >>
+ >> iter = ChunkIterator(shape=shape, chunks=chunks, dim_names=dims)
+ >> for i in range(13):
+ >> region = iter.__next__()
+ >> print(region)
+ { "inline": slice(3,6, None), "crossline": slice(0,4, None), "depth": slice(0,5, None) }
+
+ >> iter = ChunkIterator(shape=shape, chunks=chunks, dim_names=None)
+ >> for i in range(13):
+ >> region = iter.__next__()
+ >> print(region)
+ (slice(3,6,None), slice(0,4,None), slice(0,5,None))
  """
 
- def __init__(self, array: Array, chunk_samples: bool = True):
- self.arr_shape = array.shape
- self.len_chunks = array.chunks
-
- # If chunk_samples is False, set the last dimension's chunk size to its full extent
- if not chunk_samples:
- self.len_chunks = self.len_chunks[:-1] + (self.arr_shape[-1],)
-
- # Calculate the number of chunks per dimension
- self.dim_chunks = [
- ceil(len_dim / chunk)
- for len_dim, chunk in zip(self.arr_shape, self.len_chunks, strict=True)
- ]
+ def __init__(
+ self, shape: tuple[int, ...], chunks: tuple[int, ...], dim_names: tuple[str, ...] = None
+ ):
+ self.arr_shape = tuple(shape) # Deep copy to ensure immutability
+ self.len_chunks = tuple(chunks) # Deep copy to ensure immutability
+ self.dims = dim_names
+
+ # Compute number of chunks per dimension, and total number of chunks
+ self.dim_chunks = tuple(
+ [
+ ceil(len_dim / chunk)
+ for len_dim, chunk in zip(self.arr_shape, self.len_chunks, strict=True)
+ ]
+ )
  self.num_chunks = np.prod(self.dim_chunks)
 
- # Set up chunk index combinations using ranges for each dimension
+ # Under the hood stuff for the iterator. This generates C-ordered
+ # permutation of chunk indices.
  dim_ranges = [range(dim_len) for dim_len in self.dim_chunks]
  self._ranges = itertools.product(*dim_ranges)
  self._idx = 0
 
  def __iter__(self) -> "ChunkIterator":
- """Return the iterator object itself."""
+ """Iteration context."""
  return self
 
  def __len__(self) -> int:
- """Return the total number of chunks."""
+ """Get total number of chunks."""
  return self.num_chunks
 
- def __next__(self) -> tuple[slice, ...]:
- """Yield the next set of chunk slices.
-
- Returns:
- A tuple of slice objects for each dimension.
-
- Raises:
- StopIteration: When all chunks have been iterated over.
- """
- if self._idx < self.num_chunks:
+ def __next__(self) -> dict[str, slice]:
+ """Iteration logic."""
+ if self._idx <= self.num_chunks:
+ # We build slices here. It is dimension agnostic
  current_start = next(self._ranges)
 
+ # TODO (Dmitriy Repin): Enhance ChunkIterator to make the last slice, if needed, smaller
+ # https://github.com/TGSAI/mdio-python/issues/586
  start_indices = tuple(
  dim * chunk for dim, chunk in zip(current_start, self.len_chunks, strict=True)
  )
@@ -88,7 +93,17 @@ def __next__(self) -> tuple[slice, ...]:
  slice(start, stop) for start, stop in zip(start_indices, stop_indices, strict=True)
  )
 
+ if self.dims: # noqa SIM108
+ # Example
+ # {"inline":slice(3,6,None), "crossline":slice(0,4,None), "depth":slice(0,5,None)}
+ result = dict(zip(self.dims, slices, strict=False))
+ else:
+ # Example
+ # (slice(3,6,None), slice(0,4,None), slice(0,5,None))
+ result = slices
+
  self._idx += 1
- return slices
+
+ return result
 
  raise StopIteration

src/mdio/core/storage_location.py

@@ -0,0 +1,87 @@
+"""StorageLocation class for managing local and cloud storage locations."""
+
+from pathlib import Path
+from typing import Any
+
+import fsspec
+
+
+# TODO(Dmitriy Repin): Reuse fsspec functions for some methods we implemented here
+# https://github.com/TGSAI/mdio-python/issues/597
+class StorageLocation:
+ """A class to represent a local or cloud storage location for SEG-Y or MDIO files.
+
+ This class abstracts the storage location, allowing for both local file paths and
+ cloud storage URIs (e.g., S3, GCS). It uses fsspec to check existence and manage options.
+ Note, we do not want to make it a dataclass because we want the uri and the options to
+ be read-only immutable properties.
+
+ uri: The URI of the storage location (e.g., '/path/to/file', 'file:///path/to/file',
+ 's3://bucket/path', 'gs://bucket/path').
+ options: Optional dictionary of options for the cloud, such as credentials.
+
+ """
+
+ def __init__(self, uri: str = "", options: dict[str, Any] = None):
+ self._uri = uri
+ self._options = options or {}
+ self._fs = None
+
+ if uri.startswith(("s3://", "gs://")):
+ return
+
+ if uri.startswith("file://"):
+ self._uri = self._uri.removeprefix("file://")
+ # For local paths, ensure they are absolute and resolved
+ self._uri = str(Path(self._uri).resolve())
+ return
+
+ @property
+ def uri(self) -> str:
+ """Get the URI (read-only)."""
+ return self._uri
+
+ @property
+ def options(self) -> dict[str, Any]:
+ """Get the options (read-only)."""
+ # Return a copy to prevent external modification
+ return self._options.copy()
+
+ @property
+ def _filesystem(self) -> fsspec.AbstractFileSystem:
+ """Get the fsspec filesystem instance for this storage location."""
+ if self._fs is None:
+ self._fs = fsspec.filesystem(self._protocol, **self._options)
+ return self._fs
+
+ @property
+ def _path(self) -> str:
+ """Extract the path portion from the URI."""
+ if "://" in self._uri:
+ return self._uri.split("://", 1)[1]
+ return self._uri # For local paths without file:// prefix
+
+ @property
+ def _protocol(self) -> str:
+ """Extract the protocol/scheme from the URI."""
+ if "://" in self._uri:
+ return self._uri.split("://", 1)[0]
+ return "file" # Default to file protocol
+
+ def exists(self) -> bool:
+ """Check if the storage location exists using fsspec."""
+ try:
+ return self._filesystem.exists(self._path)
+ except Exception as e:
+ # Log the error and return False for safety
+ # In a production environment, you might want to use proper logging
+ print(f"Error checking existence of {self._uri}: {e}")
+ return False
+
+ def __str__(self) -> str:
+ """String representation of the storage location."""
+ return self._uri
+
+ def __repr__(self) -> str:
+ """Developer representation of the storage location."""
+ return f"StorageLocation(uri='{self._uri}', options={self._options})"

src/mdio/api/convenience.py

@@ -79,7 +79,11 @@ def copy_mdio( # noqa: PLR0913
 
  writer.live_mask[:] = reader.live_mask[:]
 
- iterator = ChunkIterator(reader._traces, chunk_samples=False)
+ shape = reader._traces.shape
+ chunks = reader._traces.chunks
+ chunks = chunks[:-1] + (shape[-1],) # don't chunk samples
+
+ iterator = ChunkIterator(shape=shape, chunks=chunks)
  progress = tqdm(iterator, unit="block")
  progress.set_description(desc=f"Copying data for '{access_pattern=}'")
  for slice_ in progress:
@@ -177,7 +181,10 @@ def create_rechunk_plan(
 
  n_dimension = len(data_array.shape)
  dummy_array = zarr.empty(shape=data_array.shape, chunks=(MAX_BUFFER,) * n_dimension)
- iterator = ChunkIterator(dummy_array)
+
+ shape = dummy_array.shape
+ chunks = dummy_array.chunks
+ iterator = ChunkIterator(shape=shape, chunks=chunks)
 
  return metadata_arrs, data_arrs, live_mask, iterator
 

src/mdio/converters/numpy.py

@@ -7,7 +7,6 @@
 import numpy as np
 
 from mdio.api.accessor import MDIOWriter
-from mdio.converters.segy import get_compressor
 from mdio.core.dimension import Dimension
 from mdio.core.factory import MDIOCreateConfig
 from mdio.core.factory import MDIOVariableConfig
@@ -137,6 +136,11 @@ def numpy_to_mdio( # noqa: PLR0913
  suffix = [str(idx) for idx, value in enumerate(suffix) if value is not None]
  suffix = "".join(suffix)
 
+ # TODO(Dmitrit Repin): Implement Numpy converted in MDIO v1
+ # https://github.com/TGSAI/mdio-python/issues/596
+ def get_compressor(lossless: bool, tolerance: float) -> list[str]:
+ pass
+
  compressors = get_compressor(lossless, compression_tolerance)
  mdio_var = MDIOVariableConfig(
  name=f"chunked_{suffix}",

src/mdio/converters/type_converter.py

@@ -0,0 +1,85 @@
+"""A module for converting numpy dtypes to MDIO scalar and structured types."""
+
+from numpy import dtype as np_dtype
+
+from mdio.schemas.dtype import ScalarType
+from mdio.schemas.dtype import StructuredField
+from mdio.schemas.dtype import StructuredType
+
+
+def to_scalar_type(data_type: np_dtype) -> ScalarType:
+ """Convert numpy dtype to MDIO ScalarType.
+
+ Out of the 24 built-in numpy scalar type objects
+ (see https://numpy.org/doc/stable/reference/arrays.dtypes.html)
+ this function supports only a limited subset:
+ ScalarType.INT8 <-> int8
+ ScalarType.INT16 <-> int16
+ ScalarType.INT32 <-> int32
+ ScalarType.INT64 <-> int64
+ ScalarType.UINT8 <-> uint8
+ ScalarType.UINT16 <-> uint16
+ ScalarType.UINT32 <-> uint32
+ ScalarType.UINT64 <-> uint64
+ ScalarType.FLOAT32 <-> float32
+ ScalarType.FLOAT64 <-> float64
+ ScalarType.COMPLEX64 <-> complex64
+ ScalarType.COMPLEX128 <-> complex128
+ ScalarType.BOOL <-> bool
+
+ Args:
+ data_type: numpy dtype to convert
+
+ Returns:
+ ScalarType: corresponding MDIO scalar type
+
+ Raises:
+ ValueError: if dtype is not supported
+ """
+ try:
+ return ScalarType(data_type.name)
+ except ValueError as exc:
+ err = f"Unsupported numpy dtype '{data_type.name}' for conversion to ScalarType."
+ raise ValueError(err) from exc
+
+
+def to_structured_type(data_type: np_dtype) -> StructuredType:
+ """Convert numpy dtype to MDIO StructuredType.
+
+ This function supports only a limited subset of structured types.
+ In particular:
+ It does not support nested structured types.
+ It supports fields of only 13 out of 24 built-in numpy scalar types.
+ (see `to_scalar_type` for details).
+
+ Args:
+ data_type: numpy dtype to convert
+
+ Returns:
+ StructuredType: corresponding MDIO structured type
+
+ Raises:
+ ValueError: if dtype is not structured or has no fields
+
+ """
+ if data_type is None or len(data_type.names or []) == 0:
+ err = "None or empty dtype provided, cannot convert to StructuredType."
+ raise ValueError(err)
+
+ fields = []
+ for field_name in data_type.names:
+ field_dtype = data_type.fields[field_name][0]
+ scalar_type = to_scalar_type(field_dtype)
+ structured_field = StructuredField(name=field_name, format=scalar_type)
+ fields.append(structured_field)
+ return StructuredType(fields=fields)
+
+
+def to_numpy_dtype(data_type: ScalarType | StructuredType) -> np_dtype:
+ """Get the numpy dtype for a variable."""
+ if isinstance(data_type, ScalarType):
+ return np_dtype(data_type.value)
+ if isinstance(data_type, StructuredType):
+ return np_dtype([(f.name, f.format.value) for f in data_type.fields])
+ msg = f"Expected ScalarType or StructuredType, got '{type(data_type).__name__}'"
+ raise ValueError(msg)

src/mdio/core/indexing.py

@@ -4,78 +4,83 @@
 from math import ceil
 
 import numpy as np
-from zarr import Array
 
 
 class ChunkIterator:
- """Iterator for traversing a Zarr array in chunks.
+ """Chunk iterator for multi-dimensional arrays.
 
- This iterator yields tuples of slices corresponding to the chunk boundaries of a Zarr array.
- It supports chunking all dimensions or taking the full extent of the last dimension.
+ This iterator takes an array shape and chunks and every time it is iterated, it returns
+ a dictionary (if dimensions are provided) or a tuple of slices that align with
+ chunk boundaries. When dimensions are provided, they are used as the dictionary keys.
 
  Args:
- array: The Zarr array to iterate, providing shape and chunk sizes.
- chunk_samples: If True, chunks all dimensions. If False, takes the full extent of the
- last dimension. Defaults to True.
-
-
- Example:
- >>> import zarr
- >>> arr = zarr.array(np.zeros((10, 20)), chunks=(3, 4))
- >>> it = ChunkIterator(arr)
- >>> for slices in it:
- ... print(slices)
- (slice(0, 3, None), slice(0, 4, None))
- (slice(0, 3, None), slice(4, 8, None))
- ...
- >>> it = ChunkIterator(arr, chunk_samples=False)
- >>> for slices in it:
- ... print(slices)
- (slice(0, 3, None), slice(0, 20, None))
- (slice(3, 6, None), slice(0, 20, None))
- ...
+ shape: The shape of the array.
+ chunks: The chunk sizes for each dimension.
+ dim_names: The names of the array dimensions, to be used with DataArray.isel().
+ If the dim_names are not provided, a tuple of the slices will be returned.
+
+ Attributes: # noqa: DOC602
+ arr_shape: Shape of the array.
+ len_chunks: Length of chunks in each dimension.
+ dim_chunks: Number of chunks in each dimension.
+ num_chunks: Total number of chunks.
+
+ Examples:
+ >> chunks = (3, 4, 5)
+ >> shape = (5, 11, 19)
+ >> dims = ["inline", "crossline", "depth"]
+ >>
+ >> iter = ChunkIterator(shape=shape, chunks=chunks, dim_names=dims)
+ >> for i in range(13):
+ >> region = iter.__next__()
+ >> print(region)
+ { "inline": slice(3,6, None), "crossline": slice(0,4, None), "depth": slice(0,5, None) }
+
+ >> iter = ChunkIterator(shape=shape, chunks=chunks, dim_names=None)
+ >> for i in range(13):
+ >> region = iter.__next__()
+ >> print(region)
+ (slice(3,6,None), slice(0,4,None), slice(0,5,None))
  """
 
- def __init__(self, array: Array, chunk_samples: bool = True):
- self.arr_shape = array.shape
- self.len_chunks = array.chunks
-
- # If chunk_samples is False, set the last dimension's chunk size to its full extent
- if not chunk_samples:
- self.len_chunks = self.len_chunks[:-1] + (self.arr_shape[-1],)
-
- # Calculate the number of chunks per dimension
- self.dim_chunks = [
- ceil(len_dim / chunk)
- for len_dim, chunk in zip(self.arr_shape, self.len_chunks, strict=True)
- ]
+ def __init__(
+ self, shape: tuple[int, ...], chunks: tuple[int, ...], dim_names: tuple[str, ...] = None
+ ):
+ self.arr_shape = tuple(shape) # Deep copy to ensure immutability
+ self.len_chunks = tuple(chunks) # Deep copy to ensure immutability
+ self.dims = dim_names
+
+ # Compute number of chunks per dimension, and total number of chunks
+ self.dim_chunks = tuple(
+ [
+ ceil(len_dim / chunk)
+ for len_dim, chunk in zip(self.arr_shape, self.len_chunks, strict=True)
+ ]
+ )
  self.num_chunks = np.prod(self.dim_chunks)
 
- # Set up chunk index combinations using ranges for each dimension
+ # Under the hood stuff for the iterator. This generates C-ordered
+ # permutation of chunk indices.
  dim_ranges = [range(dim_len) for dim_len in self.dim_chunks]
  self._ranges = itertools.product(*dim_ranges)
  self._idx = 0
 
  def __iter__(self) -> "ChunkIterator":
- """Return the iterator object itself."""
+ """Iteration context."""
  return self
 
  def __len__(self) -> int:
- """Return the total number of chunks."""
+ """Get total number of chunks."""
  return self.num_chunks
 
- def __next__(self) -> tuple[slice, ...]:
- """Yield the next set of chunk slices.
-
- Returns:
- A tuple of slice objects for each dimension.
-
- Raises:
- StopIteration: When all chunks have been iterated over.
- """
- if self._idx < self.num_chunks:
+ def __next__(self) -> dict[str, slice]:
+ """Iteration logic."""
+ if self._idx <= self.num_chunks:
+ # We build slices here. It is dimension agnostic
  current_start = next(self._ranges)
 
+ # TODO (Dmitriy Repin): Enhance ChunkIterator to make the last slice, if needed, smaller
+ # https://github.com/TGSAI/mdio-python/issues/586
  start_indices = tuple(
  dim * chunk for dim, chunk in zip(current_start, self.len_chunks, strict=True)
  )
@@ -88,7 +93,17 @@ def __next__(self) -> tuple[slice, ...]:
  slice(start, stop) for start, stop in zip(start_indices, stop_indices, strict=True)
  )
 
+ if self.dims: # noqa SIM108
+ # Example
+ # {"inline":slice(3,6,None), "crossline":slice(0,4,None), "depth":slice(0,5,None)}
+ result = dict(zip(self.dims, slices, strict=False))
+ else:
+ # Example
+ # (slice(3,6,None), slice(0,4,None), slice(0,5,None))
+ result = slices
+
  self._idx += 1
- return slices
+
+ return result
 
  raise StopIteration

src/mdio/core/storage_location.py

@@ -0,0 +1,87 @@
+"""StorageLocation class for managing local and cloud storage locations."""
+
+from pathlib import Path
+from typing import Any
+
+import fsspec
+
+
+# TODO(Dmitriy Repin): Reuse fsspec functions for some methods we implemented here
+# https://github.com/TGSAI/mdio-python/issues/597
+class StorageLocation:
+ """A class to represent a local or cloud storage location for SEG-Y or MDIO files.
+
+ This class abstracts the storage location, allowing for both local file paths and
+ cloud storage URIs (e.g., S3, GCS). It uses fsspec to check existence and manage options.
+ Note, we do not want to make it a dataclass because we want the uri and the options to
+ be read-only immutable properties.
+
+ uri: The URI of the storage location (e.g., '/path/to/file', 'file:///path/to/file',
+ 's3://bucket/path', 'gs://bucket/path').
+ options: Optional dictionary of options for the cloud, such as credentials.
+
+ """
+
+ def __init__(self, uri: str = "", options: dict[str, Any] = None):
+ self._uri = uri
+ self._options = options or {}
+ self._fs = None
+
+ if uri.startswith(("s3://", "gs://")):
+ return
+
+ if uri.startswith("file://"):
+ self._uri = self._uri.removeprefix("file://")
+ # For local paths, ensure they are absolute and resolved
+ self._uri = str(Path(self._uri).resolve())
+ return
+
+ @property
+ def uri(self) -> str:
+ """Get the URI (read-only)."""
+ return self._uri
+
+ @property
+ def options(self) -> dict[str, Any]:
+ """Get the options (read-only)."""
+ # Return a copy to prevent external modification
+ return self._options.copy()
+
+ @property
+ def _filesystem(self) -> fsspec.AbstractFileSystem:
+ """Get the fsspec filesystem instance for this storage location."""
+ if self._fs is None:
+ self._fs = fsspec.filesystem(self._protocol, **self._options)
+ return self._fs
+
+ @property
+ def _path(self) -> str:
+ """Extract the path portion from the URI."""
+ if "://" in self._uri:
+ return self._uri.split("://", 1)[1]
+ return self._uri # For local paths without file:// prefix
+
+ @property
+ def _protocol(self) -> str:
+ """Extract the protocol/scheme from the URI."""
+ if "://" in self._uri:
+ return self._uri.split("://", 1)[0]
+ return "file" # Default to file protocol
+
+ def exists(self) -> bool:
+ """Check if the storage location exists using fsspec."""
+ try:
+ return self._filesystem.exists(self._path)
+ except Exception as e:
+ # Log the error and return False for safety
+ # In a production environment, you might want to use proper logging
+ print(f"Error checking existence of {self._uri}: {e}")
+ return False
+
+ def __str__(self) -> str:
+ """String representation of the storage location."""
+ return self._uri
+
+ def __repr__(self) -> str:
+ """Developer representation of the storage location."""
+ return f"StorageLocation(uri='{self._uri}', options={self._options})"

src/mdio/api/convenience.py

@@ -79,7 +79,11 @@ def copy_mdio( # noqa: PLR0913
 
  writer.live_mask[:] = reader.live_mask[:]
 
- iterator = ChunkIterator(reader._traces, chunk_samples=False)
+ shape = reader._traces.shape
+ chunks = reader._traces.chunks
+ chunks = chunks[:-1] + (shape[-1],) # don't chunk samples
+
+ iterator = ChunkIterator(shape=shape, chunks=chunks)
  progress = tqdm(iterator, unit="block")
  progress.set_description(desc=f"Copying data for '{access_pattern=}'")
  for slice_ in progress:
@@ -177,7 +181,10 @@ def create_rechunk_plan(
 
  n_dimension = len(data_array.shape)
  dummy_array = zarr.empty(shape=data_array.shape, chunks=(MAX_BUFFER,) * n_dimension)
- iterator = ChunkIterator(dummy_array)
+
+ shape = dummy_array.shape
+ chunks = dummy_array.chunks
+ iterator = ChunkIterator(shape=shape, chunks=chunks)
 
  return metadata_arrs, data_arrs, live_mask, iterator
 

src/mdio/converters/numpy.py

@@ -7,7 +7,6 @@
 import numpy as np
 
 from mdio.api.accessor import MDIOWriter
-from mdio.converters.segy import get_compressor
 from mdio.core.dimension import Dimension
 from mdio.core.factory import MDIOCreateConfig
 from mdio.core.factory import MDIOVariableConfig
@@ -137,6 +136,11 @@ def numpy_to_mdio( # noqa: PLR0913
  suffix = [str(idx) for idx, value in enumerate(suffix) if value is not None]
  suffix = "".join(suffix)
 
+ # TODO(Dmitrit Repin): Implement Numpy converted in MDIO v1
+ # https://github.com/TGSAI/mdio-python/issues/596
+ def get_compressor(lossless: bool, tolerance: float) -> list[str]:
+ pass
+
  compressors = get_compressor(lossless, compression_tolerance)
  mdio_var = MDIOVariableConfig(
  name=f"chunked_{suffix}",

src/mdio/converters/type_converter.py

@@ -0,0 +1,85 @@
+"""A module for converting numpy dtypes to MDIO scalar and structured types."""
+
+from numpy import dtype as np_dtype
+
+from mdio.schemas.dtype import ScalarType
+from mdio.schemas.dtype import StructuredField
+from mdio.schemas.dtype import StructuredType
+
+
+def to_scalar_type(data_type: np_dtype) -> ScalarType:
+ """Convert numpy dtype to MDIO ScalarType.
+
+ Out of the 24 built-in numpy scalar type objects
+ (see https://numpy.org/doc/stable/reference/arrays.dtypes.html)
+ this function supports only a limited subset:
+ ScalarType.INT8 <-> int8
+ ScalarType.INT16 <-> int16
+ ScalarType.INT32 <-> int32
+ ScalarType.INT64 <-> int64
+ ScalarType.UINT8 <-> uint8
+ ScalarType.UINT16 <-> uint16
+ ScalarType.UINT32 <-> uint32
+ ScalarType.UINT64 <-> uint64
+ ScalarType.FLOAT32 <-> float32
+ ScalarType.FLOAT64 <-> float64
+ ScalarType.COMPLEX64 <-> complex64
+ ScalarType.COMPLEX128 <-> complex128
+ ScalarType.BOOL <-> bool
+
+ Args:
+ data_type: numpy dtype to convert
+
+ Returns:
+ ScalarType: corresponding MDIO scalar type
+
+ Raises:
+ ValueError: if dtype is not supported
+ """
+ try:
+ return ScalarType(data_type.name)
+ except ValueError as exc:
+ err = f"Unsupported numpy dtype '{data_type.name}' for conversion to ScalarType."
+ raise ValueError(err) from exc
+
+
+def to_structured_type(data_type: np_dtype) -> StructuredType:
+ """Convert numpy dtype to MDIO StructuredType.
+
+ This function supports only a limited subset of structured types.
+ In particular:
+ It does not support nested structured types.
+ It supports fields of only 13 out of 24 built-in numpy scalar types.
+ (see `to_scalar_type` for details).
+
+ Args:
+ data_type: numpy dtype to convert
+
+ Returns:
+ StructuredType: corresponding MDIO structured type
+
+ Raises:
+ ValueError: if dtype is not structured or has no fields
+
+ """
+ if data_type is None or len(data_type.names or []) == 0:
+ err = "None or empty dtype provided, cannot convert to StructuredType."
+ raise ValueError(err)
+
+ fields = []
+ for field_name in data_type.names:
+ field_dtype = data_type.fields[field_name][0]
+ scalar_type = to_scalar_type(field_dtype)
+ structured_field = StructuredField(name=field_name, format=scalar_type)
+ fields.append(structured_field)
+ return StructuredType(fields=fields)
+
+
+def to_numpy_dtype(data_type: ScalarType | StructuredType) -> np_dtype:
+ """Get the numpy dtype for a variable."""
+ if isinstance(data_type, ScalarType):
+ return np_dtype(data_type.value)
+ if isinstance(data_type, StructuredType):
+ return np_dtype([(f.name, f.format.value) for f in data_type.fields])
+ msg = f"Expected ScalarType or StructuredType, got '{type(data_type).__name__}'"
+ raise ValueError(msg)

src/mdio/core/indexing.py

@@ -4,78 +4,83 @@
 from math import ceil
 
 import numpy as np
-from zarr import Array
 
 
 class ChunkIterator:
- """Iterator for traversing a Zarr array in chunks.
+ """Chunk iterator for multi-dimensional arrays.
 
- This iterator yields tuples of slices corresponding to the chunk boundaries of a Zarr array.
- It supports chunking all dimensions or taking the full extent of the last dimension.
+ This iterator takes an array shape and chunks and every time it is iterated, it returns
+ a dictionary (if dimensions are provided) or a tuple of slices that align with
+ chunk boundaries. When dimensions are provided, they are used as the dictionary keys.
 
  Args:
- array: The Zarr array to iterate, providing shape and chunk sizes.
- chunk_samples: If True, chunks all dimensions. If False, takes the full extent of the
- last dimension. Defaults to True.
-
-
- Example:
- >>> import zarr
- >>> arr = zarr.array(np.zeros((10, 20)), chunks=(3, 4))
- >>> it = ChunkIterator(arr)
- >>> for slices in it:
- ... print(slices)
- (slice(0, 3, None), slice(0, 4, None))
- (slice(0, 3, None), slice(4, 8, None))
- ...
- >>> it = ChunkIterator(arr, chunk_samples=False)
- >>> for slices in it:
- ... print(slices)
- (slice(0, 3, None), slice(0, 20, None))
- (slice(3, 6, None), slice(0, 20, None))
- ...
+ shape: The shape of the array.
+ chunks: The chunk sizes for each dimension.
+ dim_names: The names of the array dimensions, to be used with DataArray.isel().
+ If the dim_names are not provided, a tuple of the slices will be returned.
+
+ Attributes: # noqa: DOC602
+ arr_shape: Shape of the array.
+ len_chunks: Length of chunks in each dimension.
+ dim_chunks: Number of chunks in each dimension.
+ num_chunks: Total number of chunks.
+
+ Examples:
+ >> chunks = (3, 4, 5)
+ >> shape = (5, 11, 19)
+ >> dims = ["inline", "crossline", "depth"]
+ >>
+ >> iter = ChunkIterator(shape=shape, chunks=chunks, dim_names=dims)
+ >> for i in range(13):
+ >> region = iter.__next__()
+ >> print(region)
+ { "inline": slice(3,6, None), "crossline": slice(0,4, None), "depth": slice(0,5, None) }
+
+ >> iter = ChunkIterator(shape=shape, chunks=chunks, dim_names=None)
+ >> for i in range(13):
+ >> region = iter.__next__()
+ >> print(region)
+ (slice(3,6,None), slice(0,4,None), slice(0,5,None))
  """
 
- def __init__(self, array: Array, chunk_samples: bool = True):
- self.arr_shape = array.shape
- self.len_chunks = array.chunks
-
- # If chunk_samples is False, set the last dimension's chunk size to its full extent
- if not chunk_samples:
- self.len_chunks = self.len_chunks[:-1] + (self.arr_shape[-1],)
-
- # Calculate the number of chunks per dimension
- self.dim_chunks = [
- ceil(len_dim / chunk)
- for len_dim, chunk in zip(self.arr_shape, self.len_chunks, strict=True)
- ]
+ def __init__(
+ self, shape: tuple[int, ...], chunks: tuple[int, ...], dim_names: tuple[str, ...] = None
+ ):
+ self.arr_shape = tuple(shape) # Deep copy to ensure immutability
+ self.len_chunks = tuple(chunks) # Deep copy to ensure immutability
+ self.dims = dim_names
+
+ # Compute number of chunks per dimension, and total number of chunks
+ self.dim_chunks = tuple(
+ [
+ ceil(len_dim / chunk)
+ for len_dim, chunk in zip(self.arr_shape, self.len_chunks, strict=True)
+ ]
+ )
  self.num_chunks = np.prod(self.dim_chunks)
 
- # Set up chunk index combinations using ranges for each dimension
+ # Under the hood stuff for the iterator. This generates C-ordered
+ # permutation of chunk indices.
  dim_ranges = [range(dim_len) for dim_len in self.dim_chunks]
  self._ranges = itertools.product(*dim_ranges)
  self._idx = 0
 
  def __iter__(self) -> "ChunkIterator":
- """Return the iterator object itself."""
+ """Iteration context."""
  return self
 
  def __len__(self) -> int:
- """Return the total number of chunks."""
+ """Get total number of chunks."""
  return self.num_chunks
 
- def __next__(self) -> tuple[slice, ...]:
- """Yield the next set of chunk slices.
-
- Returns:
- A tuple of slice objects for each dimension.
-
- Raises:
- StopIteration: When all chunks have been iterated over.
- """
- if self._idx < self.num_chunks:
+ def __next__(self) -> dict[str, slice]:
+ """Iteration logic."""
+ if self._idx <= self.num_chunks:
+ # We build slices here. It is dimension agnostic
  current_start = next(self._ranges)
 
+ # TODO (Dmitriy Repin): Enhance ChunkIterator to make the last slice, if needed, smaller
+ # https://github.com/TGSAI/mdio-python/issues/586
  start_indices = tuple(
  dim * chunk for dim, chunk in zip(current_start, self.len_chunks, strict=True)
  )
@@ -88,7 +93,17 @@ def __next__(self) -> tuple[slice, ...]:
  slice(start, stop) for start, stop in zip(start_indices, stop_indices, strict=True)
  )
 
+ if self.dims: # noqa SIM108
+ # Example
+ # {"inline":slice(3,6,None), "crossline":slice(0,4,None), "depth":slice(0,5,None)}
+ result = dict(zip(self.dims, slices, strict=False))
+ else:
+ # Example
+ # (slice(3,6,None), slice(0,4,None), slice(0,5,None))
+ result = slices
+
  self._idx += 1
- return slices
+
+ return result
 
  raise StopIteration

src/mdio/core/storage_location.py

@@ -0,0 +1,87 @@
+"""StorageLocation class for managing local and cloud storage locations."""
+
+from pathlib import Path
+from typing import Any
+
+import fsspec
+
+
+# TODO(Dmitriy Repin): Reuse fsspec functions for some methods we implemented here
+# https://github.com/TGSAI/mdio-python/issues/597
+class StorageLocation:
+ """A class to represent a local or cloud storage location for SEG-Y or MDIO files.
+
+ This class abstracts the storage location, allowing for both local file paths and
+ cloud storage URIs (e.g., S3, GCS). It uses fsspec to check existence and manage options.
+ Note, we do not want to make it a dataclass because we want the uri and the options to
+ be read-only immutable properties.
+
+ uri: The URI of the storage location (e.g., '/path/to/file', 'file:///path/to/file',
+ 's3://bucket/path', 'gs://bucket/path').
+ options: Optional dictionary of options for the cloud, such as credentials.
+
+ """
+
+ def __init__(self, uri: str = "", options: dict[str, Any] = None):
+ self._uri = uri
+ self._options = options or {}
+ self._fs = None
+
+ if uri.startswith(("s3://", "gs://")):
+ return
+
+ if uri.startswith("file://"):
+ self._uri = self._uri.removeprefix("file://")
+ # For local paths, ensure they are absolute and resolved
+ self._uri = str(Path(self._uri).resolve())
+ return
+
+ @property
+ def uri(self) -> str:
+ """Get the URI (read-only)."""
+ return self._uri
+
+ @property
+ def options(self) -> dict[str, Any]:
+ """Get the options (read-only)."""
+ # Return a copy to prevent external modification
+ return self._options.copy()
+
+ @property
+ def _filesystem(self) -> fsspec.AbstractFileSystem:
+ """Get the fsspec filesystem instance for this storage location."""
+ if self._fs is None:
+ self._fs = fsspec.filesystem(self._protocol, **self._options)
+ return self._fs
+
+ @property
+ def _path(self) -> str:
+ """Extract the path portion from the URI."""
+ if "://" in self._uri:
+ return self._uri.split("://", 1)[1]
+ return self._uri # For local paths without file:// prefix
+
+ @property
+ def _protocol(self) -> str:
+ """Extract the protocol/scheme from the URI."""
+ if "://" in self._uri:
+ return self._uri.split("://", 1)[0]
+ return "file" # Default to file protocol
+
+ def exists(self) -> bool:
+ """Check if the storage location exists using fsspec."""
+ try:
+ return self._filesystem.exists(self._path)
+ except Exception as e:
+ # Log the error and return False for safety
+ # In a production environment, you might want to use proper logging
+ print(f"Error checking existence of {self._uri}: {e}")
+ return False
+
+ def __str__(self) -> str:
+ """String representation of the storage location."""
+ return self._uri
+
+ def __repr__(self) -> str:
+ """Developer representation of the storage location."""
+ return f"StorageLocation(uri='{self._uri}', options={self._options})"