from __future__ import annotations
import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, Mapping, Sequence
import numpy as np
from transcriptml.data.schemas import SequenceSchema, get_schema
[docs]
@dataclass
class DatasetBundle:
"""Self-describing processed dataset.
``X`` and optional ``y`` retain TranscriptML's original compact contract.
Workflows with additional aligned targets (for example, RBPNet count
profiles) may use ``arrays``. Every named array must share ``X``'s first
dimension and is serialized as an ordinary ``<name>.npy`` file.
"""
X: np.ndarray
y: np.ndarray | None = None
ids: Sequence[str] | None = None
schema: SequenceSchema | str = "rna4"
metadata: Sequence[Mapping[str, Any]] | None = None
splits: Mapping[str, Sequence[int]] | None = None
config: Mapping[str, Any] = field(default_factory=dict)
arrays: Mapping[str, np.ndarray] = field(default_factory=dict)
def __post_init__(self) -> None:
"""Normalize schema and validate array-aligned fields."""
self.schema = get_schema(self.schema)
if self.ids is None:
self.ids = [str(i) for i in range(int(self.X.shape[0]))]
if len(self.ids) != int(self.X.shape[0]):
raise ValueError("ids length must match X.shape[0]")
if self.y is not None and int(self.y.shape[0]) != int(self.X.shape[0]):
raise ValueError("y length must match X.shape[0]")
if self.metadata is not None and len(self.metadata) != int(self.X.shape[0]):
raise ValueError("metadata length must match X.shape[0]")
reserved = {"X", "y"}
for name, array in self.arrays.items():
if not name or name in reserved or not name.replace("_", "").isalnum():
raise ValueError(f"invalid named array key: {name!r}")
if int(array.shape[0]) != int(self.X.shape[0]):
raise ValueError(f"named array {name!r} length must match X.shape[0]")
def _json_default(obj: Any) -> Any:
"""Convert common NumPy and path objects to JSON-serializable values.
Args:
obj: Object passed by ``json.dumps`` when the standard encoder cannot
serialize it.
"""
if isinstance(obj, np.ndarray):
return obj.tolist()
if isinstance(obj, (np.integer, np.floating)):
return obj.item()
if isinstance(obj, Path):
return str(obj)
raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")
[docs]
def save_bundle(bundle: DatasetBundle, out_dir: str | Path) -> None:
"""Write a complete dataset bundle, including arrays and sidecars.
Args:
bundle: Dataset bundle containing ``X`` and optional ``y`` arrays plus
sidecar metadata.
out_dir: Destination directory for the complete on-disk bundle.
"""
out = Path(out_dir)
out.mkdir(parents=True, exist_ok=True)
np.save(out / "X.npy", bundle.X)
if bundle.y is not None:
np.save(out / "y.npy", bundle.y)
for name, array in bundle.arrays.items():
np.save(out / f"{name}.npy", array)
save_bundle_metadata(bundle, out)
[docs]
def load_bundle(path: str | Path, *, mmap_mode: str | None = None) -> DatasetBundle:
"""Load a processed dataset bundle from disk.
Args:
path: Directory containing ``X.npy`` and TranscriptML sidecar files.
mmap_mode: Optional NumPy memory-map mode to pass when loading arrays.
"""
root = Path(path)
X = np.load(root / "X.npy", mmap_mode=mmap_mode)
y_path = root / "y.npy"
y = np.load(y_path, mmap_mode=mmap_mode) if y_path.exists() else None
ids = (root / "ids.txt").read_text(encoding="utf-8").splitlines()
schema = SequenceSchema.from_dict(json.loads((root / "schema.json").read_text(encoding="utf-8")))
metadata_path = root / "metadata.json"
metadata = json.loads(metadata_path.read_text(encoding="utf-8")) if metadata_path.exists() else None
splits_path = root / "splits.json"
splits = json.loads(splits_path.read_text(encoding="utf-8")) if splits_path.exists() else None
config_path = root / "config.json"
config = json.loads(config_path.read_text(encoding="utf-8")) if config_path.exists() else {}
arrays = {}
for name, spec in config.get("named_arrays", {}).items():
filename = Path(spec["file"])
if filename.is_absolute() or len(filename.parts) != 1 or filename.suffix != ".npy":
raise ValueError(f"invalid named array file for {name!r}: {filename}")
array = np.load(root / filename, mmap_mode=mmap_mode)
expected_shape = tuple(int(value) for value in spec.get("shape", array.shape))
expected_dtype = np.dtype(spec.get("dtype", array.dtype))
if array.shape != expected_shape or array.dtype != expected_dtype:
raise ValueError(
f"named array {name!r} does not match config metadata: "
f"found {array.shape}/{array.dtype}, expected {expected_shape}/{expected_dtype}"
)
arrays[name] = array
return DatasetBundle(
X=X,
y=y,
ids=ids,
schema=schema,
metadata=metadata,
splits=splits,
config=config,
arrays=arrays,
)