radiobject 0.1.0__py3-none-any.whl

radiobject/__init__.py

@@ -0,0 +1,24 @@
+"""RadiObject - TileDB-backed data structure for radiology data at scale."""
+
+from radiobject._types import AttrValue, LabelSource, TransformFn
+from radiobject.ctx import ReadConfig, WriteConfig, configure, ctx, get_config
+from radiobject.dataframe import Dataframe
+from radiobject.radi_object import RadiObject
+from radiobject.volume import Volume
+from radiobject.volume_collection import VolumeCollection
+
+__version__ = "0.1.0"
+__all__ = [
+    "RadiObject",
+    "Volume",
+    "VolumeCollection",
+    "Dataframe",
+    "ctx",
+    "configure",
+    "get_config",
+    "WriteConfig",
+    "ReadConfig",
+    "TransformFn",
+    "AttrValue",
+    "LabelSource",
+]

radiobject/_types.py

@@ -0,0 +1,19 @@
+"""Shared type aliases for the radiobject package."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+import pandas as pd
+
+# Volume transform function: (X, Y, Z) -> (X', Y', Z')
+TransformFn = Callable[[npt.NDArray[np.floating]], npt.NDArray[np.floating]]
+
+# Scalar values storable in TileDB obs attributes
+AttrValue = int | float | bool | str
+
+# Flexible label specification for ML datasets
+LabelSource = str | pd.DataFrame | dict[str, Any] | Callable[[str], Any] | None

radiobject/ctx.py

@@ -0,0 +1,359 @@
+"""TileDB context configuration for radiology data."""
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Self
+
+import boto3
+import tiledb
+from pydantic import BaseModel, Field, model_validator
+
+
+class SliceOrientation(str, Enum):
+    """Preferred slicing orientation for tile optimization."""
+
+    AXIAL = "axial"  # X-Y slices, vary Z (most common for neuro)
+    SAGITTAL = "sagittal"  # Y-Z slices, vary X
+    CORONAL = "coronal"  # X-Z slices, vary Y
+    ISOTROPIC = "isotropic"  # Balanced 64³ chunks for 3D ROI
+
+
+class Compressor(str, Enum):
+    """Compression algorithms suited for radiology data."""
+
+    ZSTD = "zstd"  # Good balance of speed and ratio
+    LZ4 = "lz4"  # Fast, lower ratio
+    GZIP = "gzip"  # High ratio, slower
+    NONE = "none"
+
+
+class TileConfig(BaseModel):
+    """Tile dimensions for chunked storage."""
+
+    orientation: SliceOrientation = Field(
+        default=SliceOrientation.AXIAL,
+        description="Primary slicing orientation for tile optimization",
+    )
+    x: int | None = Field(default=None, ge=1, description="Tile extent in X (None = auto)")
+    y: int | None = Field(default=None, ge=1, description="Tile extent in Y (None = auto)")
+    z: int | None = Field(default=None, ge=1, description="Tile extent in Z (None = auto)")
+    t: int = Field(default=1, ge=1, description="Tile extent in T dimension")
+
+    def extents_for_shape(self, shape: tuple[int, ...]) -> tuple[int, ...]:
+        """Compute optimal tile extents based on orientation and volume shape."""
+        sx, sy, sz = shape[0], shape[1], shape[2]
+
+        match self.orientation:
+            case SliceOrientation.AXIAL:
+                extents = (self.x or sx, self.y or sy, self.z or 1)
+            case SliceOrientation.SAGITTAL:
+                extents = (self.x or 1, self.y or sy, self.z or sz)
+            case SliceOrientation.CORONAL:
+                extents = (self.x or sx, self.y or 1, self.z or sz)
+            case SliceOrientation.ISOTROPIC:
+                extents = (self.x or 64, self.y or 64, self.z or 64)
+
+        if len(shape) == 4:
+            extents = (*extents, self.t)
+        return extents
+
+
+class CompressionConfig(BaseModel):
+    """Compression settings for volume data."""
+
+    algorithm: Compressor = Field(
+        default=Compressor.ZSTD,
+        description="Compression algorithm",
+    )
+    level: int = Field(
+        default=3,
+        ge=-1,
+        le=22,
+        description="Compression level (algorithm-dependent)",
+    )
+
+    def as_filter(self) -> tiledb.Filter | None:
+        match self.algorithm:
+            case Compressor.ZSTD:
+                return tiledb.ZstdFilter(level=self.level)
+            case Compressor.LZ4:
+                return tiledb.LZ4Filter(level=self.level)
+            case Compressor.GZIP:
+                return tiledb.GzipFilter(level=self.level)
+            case Compressor.NONE:
+                return None
+
+
+class WriteConfig(BaseModel):
+    """Settings applied when creating new TileDB arrays (immutable after creation)."""
+
+    tile: TileConfig = Field(default_factory=TileConfig)
+    compression: CompressionConfig = Field(default_factory=CompressionConfig)
+    orientation: "OrientationConfig" = Field(default_factory=lambda: OrientationConfig())
+
+
+class ReadConfig(BaseModel):
+    """Settings for reading TileDB arrays."""
+
+    memory_budget_mb: int = Field(
+        default=1024,
+        ge=64,
+        description="Memory budget for TileDB operations (MB)",
+    )
+    concurrency: int = Field(
+        default=4,
+        ge=1,
+        le=64,
+        description="Number of TileDB internal I/O threads",
+    )
+    max_workers: int = Field(
+        default=4,
+        ge=1,
+        le=32,
+        description="Max parallel workers for volume I/O operations",
+    )
+
+
+class IOConfig(BaseModel):
+    """I/O and memory settings (deprecated, use ReadConfig)."""
+
+    memory_budget_mb: int = Field(
+        default=1024,
+        ge=64,
+        description="Memory budget for TileDB operations (MB)",
+    )
+    concurrency: int = Field(
+        default=4,
+        ge=1,
+        le=64,
+        description="Number of TileDB internal I/O threads",
+    )
+    max_workers: int = Field(
+        default=4,
+        ge=1,
+        le=32,
+        description="Max parallel workers for volume I/O operations",
+    )
+
+
+class S3Config(BaseModel):
+    """S3/cloud storage settings."""
+
+    region: str = Field(default="us-east-1")
+    endpoint: str | None = Field(default=None, description="Custom S3 endpoint")
+    use_virtual_addressing: bool = Field(default=True)
+    max_parallel_ops: int = Field(default=8, ge=1)
+    multipart_part_size_mb: int = Field(default=50, ge=5)
+    include_credentials: bool = Field(
+        default=True, description="Include AWS credentials from boto3 session"
+    )
+
+
+class OrientationConfig(BaseModel):
+    """Orientation detection and standardization settings."""
+
+    auto_detect: bool = Field(
+        default=True,
+        description="Automatically detect orientation from file headers",
+    )
+    canonical_target: str = Field(
+        default="RAS",
+        description="Target canonical orientation (RAS, LAS, or LPS)",
+    )
+    reorient_on_load: bool = Field(
+        default=False,
+        description="Reorient to canonical when loading (preserves original by default)",
+    )
+    store_original_affine: bool = Field(
+        default=True,
+        description="Store original affine in metadata when reorienting",
+    )
+
+    @model_validator(mode="after")
+    def validate_canonical_target(self) -> Self:
+        """Ensure canonical target is valid."""
+        valid_targets = {"RAS", "LAS", "LPS"}
+        if self.canonical_target not in valid_targets:
+            raise ValueError(
+                f"canonical_target must be one of {valid_targets}, got {self.canonical_target}"
+            )
+        return self
+
+
+class RadiObjectConfig(BaseModel):
+    """Configuration for RadiObject TileDB context."""
+
+    write: WriteConfig = Field(default_factory=WriteConfig)
+    read: ReadConfig = Field(default_factory=ReadConfig)
+    s3: S3Config = Field(default_factory=S3Config)
+
+    # Deprecated flat access (for backwards compatibility)
+    _tile: TileConfig | None = None
+    _compression: CompressionConfig | None = None
+    _io: IOConfig | None = None
+    _orientation: OrientationConfig | None = None
+
+    @property
+    def tile(self) -> TileConfig:
+        """Tile configuration (deprecated, use write.tile)."""
+        return self.write.tile
+
+    @property
+    def compression(self) -> CompressionConfig:
+        """Compression configuration (deprecated, use write.compression)."""
+        return self.write.compression
+
+    @property
+    def orientation(self) -> OrientationConfig:
+        """Orientation configuration (deprecated, use write.orientation)."""
+        return self.write.orientation
+
+    @property
+    def io(self) -> ReadConfig:
+        """I/O configuration (deprecated, use read)."""
+        return self.read
+
+    @model_validator(mode="after")
+    def validate_compression_level(self) -> Self:
+        """Ensure compression level is valid for the algorithm."""
+        max_levels = {
+            Compressor.ZSTD: 22,
+            Compressor.LZ4: 16,
+            Compressor.GZIP: 9,
+            Compressor.NONE: 0,
+        }
+        max_level = max_levels[self.write.compression.algorithm]
+        if self.write.compression.level > max_level:
+            self.write.compression.level = max_level
+        return self
+
+    def to_tiledb_config(self, include_s3_credentials: bool = False) -> tiledb.Config:
+        """Convert to TileDB Config object.
+
+        Args:
+            include_s3_credentials: If True, fetch AWS credentials from boto3.
+                This is off by default to avoid expensive/failing credential
+                lookups for local-only operations.
+        """
+        cfg = tiledb.Config()
+
+        # Memory settings from read config
+        cfg["sm.memory_budget"] = str(self.read.memory_budget_mb * 1024 * 1024)
+        cfg["sm.compute_concurrency_level"] = str(self.read.concurrency)
+        cfg["sm.io_concurrency_level"] = str(self.read.concurrency)
+
+        # S3 settings (configuration only, no credential lookup)
+        cfg["vfs.s3.region"] = self.s3.region
+        cfg["vfs.s3.use_virtual_addressing"] = "true" if self.s3.use_virtual_addressing else "false"
+        cfg["vfs.s3.max_parallel_ops"] = str(self.s3.max_parallel_ops)
+        cfg["vfs.s3.multipart_part_size"] = str(self.s3.multipart_part_size_mb * 1024 * 1024)
+        if self.s3.endpoint:
+            cfg["vfs.s3.endpoint_override"] = self.s3.endpoint
+
+        # Fetch AWS credentials if configured or explicitly requested
+        if include_s3_credentials or self.s3.include_credentials:
+            self._add_s3_credentials(cfg)
+
+        return cfg
+
+    def _add_s3_credentials(self, cfg: tiledb.Config) -> None:
+        """Add AWS credentials to config from boto3 session."""
+        try:
+            session = boto3.Session()
+            credentials = session.get_credentials()
+            if credentials:
+                frozen = credentials.get_frozen_credentials()
+                cfg["vfs.s3.aws_access_key_id"] = frozen.access_key
+                cfg["vfs.s3.aws_secret_access_key"] = frozen.secret_key
+                if frozen.token:
+                    cfg["vfs.s3.aws_session_token"] = frozen.token
+        except Exception:
+            pass  # AWS credentials unavailable; S3 operations will fail
+
+    def to_tiledb_ctx(self, include_s3_credentials: bool = False) -> tiledb.Ctx:
+        """Convert to TileDB Ctx object.
+
+        Args:
+            include_s3_credentials: If True, fetch AWS credentials from boto3.
+        """
+        return tiledb.Ctx(self.to_tiledb_config(include_s3_credentials))
+
+
+# Global mutable configuration
+_config: RadiObjectConfig = RadiObjectConfig()
+_ctx: tiledb.Ctx | None = None
+
+
+def get_config() -> RadiObjectConfig:
+    """Get the current global configuration."""
+    return _config
+
+
+def ctx() -> tiledb.Ctx:
+    """Get the global TileDB context (lazily built from config)."""
+    global _ctx
+    if _ctx is None:
+        _ctx = _config.to_tiledb_ctx()
+    return _ctx
+
+
+def configure(
+    *,
+    write: WriteConfig | None = None,
+    read: ReadConfig | None = None,
+    s3: S3Config | None = None,
+    # Deprecated flat arguments for backwards compatibility
+    tile: TileConfig | None = None,
+    compression: CompressionConfig | None = None,
+    io: IOConfig | ReadConfig | None = None,
+    orientation: OrientationConfig | None = None,
+) -> None:
+    """Update global configuration.
+
+    Example (new API):
+        configure(write=WriteConfig(tile=TileConfig(orientation=SliceOrientation.AXIAL)))
+        configure(read=ReadConfig(memory_budget_mb=2048))
+
+    Example (deprecated flat API, still supported):
+        configure(tile=TileConfig(x=128, y=128, z=32))
+        configure(compression=CompressionConfig(algorithm=Compressor.LZ4))
+    """
+    global _config, _ctx
+
+    updates: dict = {}
+
+    # Handle new nested API
+    if write is not None:
+        updates["write"] = write
+    if read is not None:
+        updates["read"] = read
+    if s3 is not None:
+        updates["s3"] = s3
+
+    # Handle deprecated flat API by mapping to nested structure
+    if tile is not None or compression is not None or orientation is not None:
+        write_updates = {}
+        if tile is not None:
+            write_updates["tile"] = tile
+        if compression is not None:
+            write_updates["compression"] = compression
+        if orientation is not None:
+            write_updates["orientation"] = orientation
+        if write_updates:
+            current_write = _config.write.model_copy(update=write_updates)
+            updates["write"] = current_write
+
+    if io is not None:
+        # Map deprecated io to read config
+        if isinstance(io, ReadConfig):
+            updates["read"] = io
+        else:
+            updates["read"] = ReadConfig(
+                memory_budget_mb=io.memory_budget_mb,
+                concurrency=io.concurrency,
+                max_workers=io.max_workers,
+            )
+
+    _config = _config.model_copy(update=updates)
+    _ctx = None

radiobject/dataframe.py

@@ -0,0 +1,186 @@
+"""Dataframe - a 2D heterogeneous array backed by TileDB."""
+
+from __future__ import annotations
+
+from functools import cached_property
+
+import numpy as np
+import pandas as pd
+import tiledb
+
+from radiobject.ctx import ctx as global_ctx
+from radiobject.ctx import get_config
+
+# Mandatory index columns for all Dataframes
+INDEX_COLUMNS = ("obs_subject_id", "obs_id")
+
+
+class Dataframe:
+    """TileDB-backed sparse dataframe for observation metadata.
+
+    Used internally for obs_meta (subject-level) and obs (volume-level) storage.
+    Indexed by (obs_subject_id, obs_id) with user-defined attribute columns.
+
+    Example:
+        df = dataframe.read(columns=["age"], value_filter="age > 40")
+    """
+
+    def __init__(self, uri: str, ctx: tiledb.Ctx | None = None):
+        self.uri: str = uri
+        self._ctx: tiledb.Ctx | None = ctx
+
+    def _effective_ctx(self) -> tiledb.Ctx:
+        return self._ctx if self._ctx else global_ctx()
+
+    @cached_property
+    def _schema(self) -> tiledb.ArraySchema:
+        """Cached schema loaded once from disk."""
+        return tiledb.ArraySchema.load(self.uri, ctx=self._effective_ctx())
+
+    @property
+    def shape(self) -> tuple[int, int]:
+        """(n_rows, n_columns) dimensions."""
+        with tiledb.open(self.uri, "r", ctx=self._effective_ctx()) as arr:
+            n_rows = arr.nonempty_domain()[0][1] if arr.nonempty_domain()[0] else 0
+            if isinstance(n_rows, str):
+                n_rows = len(arr.query(attrs=[])[:][INDEX_COLUMNS[0]])
+        n_cols = self._schema.nattr
+        return (n_rows, n_cols)
+
+    @property
+    def index_columns(self) -> tuple[str, str]:
+        """Index column names (dimension names)."""
+        return INDEX_COLUMNS
+
+    @property
+    def columns(self) -> list[str]:
+        """Attribute column names (excluding index columns)."""
+        return [self._schema.attr(i).name for i in range(self._schema.nattr)]
+
+    @property
+    def all_columns(self) -> list[str]:
+        """All column names including index columns."""
+        return list(INDEX_COLUMNS) + self.columns
+
+    @cached_property
+    def dtypes(self) -> dict[str, np.dtype]:
+        """Column data types (attributes only)."""
+        schema = self._schema
+        return {schema.attr(i).name: schema.attr(i).dtype for i in range(schema.nattr)}
+
+    def __len__(self) -> int:
+        with tiledb.open(self.uri, "r", ctx=self._effective_ctx()) as arr:
+            result = arr.query(attrs=[])[:][INDEX_COLUMNS[0]]
+            return len(result)
+
+    def __repr__(self) -> str:
+        return f"Dataframe(uri={self.uri!r}, shape={self.shape}, columns={self.columns})"
+
+    def read(
+        self,
+        columns: list[str] | None = None,
+        value_filter: str | None = None,
+        include_index: bool = True,
+    ) -> pd.DataFrame:
+        """Read data with optional column selection and value filtering."""
+        # Filter out index columns from requested columns (they're dimensions, not attributes)
+        if columns is not None:
+            attrs = [c for c in columns if c not in INDEX_COLUMNS]
+        else:
+            attrs = self.columns
+        with tiledb.open(self.uri, "r", ctx=self._effective_ctx()) as arr:
+            if value_filter is not None:
+                result = arr.query(cond=value_filter, attrs=attrs)[:]
+            else:
+                result = arr.query(attrs=attrs)[:]
+        data = {col: result[col] for col in attrs}
+        if include_index:
+            for idx_col in INDEX_COLUMNS:
+                # Convert bytes to strings for index columns
+                raw = result[idx_col]
+                data[idx_col] = np.array(
+                    [v.decode() if isinstance(v, bytes) else str(v) for v in raw]
+                )
+        df = pd.DataFrame(data)
+        if include_index:
+            col_order = list(INDEX_COLUMNS) + attrs
+            df = df[col_order]
+        return df
+
+    @staticmethod
+    def _validate_schema(schema: dict[str, np.dtype]) -> None:
+        """Validate column names and types (for non-index attributes)."""
+        for name, dtype in schema.items():
+            if "\x00" in name:
+                raise ValueError(f"Column name contains null byte: {name!r}")
+            if name in INDEX_COLUMNS:
+                raise ValueError(f"Column name conflicts with index column: {name!r}")
+            if not isinstance(dtype, np.dtype):
+                try:
+                    np.dtype(dtype)
+                except TypeError as e:
+                    raise TypeError(f"Invalid dtype for column {name!r}: {dtype}") from e
+
+    @classmethod
+    def create(
+        cls,
+        uri: str,
+        schema: dict[str, np.dtype],
+        ctx: tiledb.Ctx | None = None,
+    ) -> Dataframe:
+        """Create an empty sparse Dataframe indexed by obs_subject_id and obs_id."""
+        cls._validate_schema(schema)
+        effective_ctx = ctx if ctx else global_ctx()
+
+        dims = [
+            tiledb.Dim(name=INDEX_COLUMNS[0], dtype="ascii", ctx=effective_ctx),
+            tiledb.Dim(name=INDEX_COLUMNS[1], dtype="ascii", ctx=effective_ctx),
+        ]
+        domain = tiledb.Domain(*dims, ctx=effective_ctx)
+
+        config = get_config()
+        compression_filter = config.compression.as_filter()
+        compression = (
+            tiledb.FilterList([compression_filter]) if compression_filter else tiledb.FilterList()
+        )
+        attrs = [
+            tiledb.Attr(name=name, dtype=dtype, filters=compression, ctx=effective_ctx)
+            for name, dtype in schema.items()
+        ]
+
+        array_schema = tiledb.ArraySchema(
+            domain=domain,
+            attrs=attrs,
+            sparse=True,
+            ctx=effective_ctx,
+        )
+        tiledb.Array.create(uri, array_schema, ctx=effective_ctx)
+
+        return cls(uri, ctx=ctx)
+
+    @classmethod
+    def from_pandas(
+        cls,
+        uri: str,
+        df: pd.DataFrame,
+        ctx: tiledb.Ctx | None = None,
+    ) -> Dataframe:
+        """Create a new Dataframe from a pandas DataFrame with obs_subject_id and obs_id columns."""
+        for idx_col in INDEX_COLUMNS:
+            if idx_col not in df.columns:
+                raise ValueError(f"DataFrame must contain index column: {idx_col!r}")
+
+        attr_cols = [col for col in df.columns if col not in INDEX_COLUMNS]
+        schema = {col: df[col].to_numpy().dtype for col in attr_cols}
+        dataframe = cls.create(uri, schema=schema, ctx=ctx)
+
+        effective_ctx = ctx if ctx else global_ctx()
+        with tiledb.open(uri, mode="w", ctx=effective_ctx) as arr:
+            coords = {
+                INDEX_COLUMNS[0]: df[INDEX_COLUMNS[0]].astype(str).to_numpy(),
+                INDEX_COLUMNS[1]: df[INDEX_COLUMNS[1]].astype(str).to_numpy(),
+            }
+            data = {col: df[col].to_numpy() for col in attr_cols}
+            arr[coords[INDEX_COLUMNS[0]], coords[INDEX_COLUMNS[1]]] = data
+
+        return dataframe