""" Loaders for The Well data. CONTRACT (see provenance.py for shared error types): - get_dataset() loads REAL data only. It NEVER silently substitutes synthetic data. On any failure it raises DataLoadError or SchemaValidationError. This is a deliberate design decision: retraining on the wrong data (or on synthetic data believed to be real) is worse than a loud failure. - Synthetic data is only available via the separate, explicitly-named get_synthetic_dataset() — callers must opt in on purpose. - Tensor layout is never guessed from shape magnitude. Callers must declare the expected channel_layout; a mismatch is a hard error, not a silent transpose. Setup for real data: 1. Run: the-well-download --base-path ./data/well --dataset active_matter --split train 2. Or place .hdf5 files under data/real/ with a known, declared schema. """ from __future__ import annotations import os from pathlib import Path from typing import Optional, List, Literal import torch from torch.utils.data import Dataset import numpy as np from .provenance import DataLoadError, SchemaValidationError ChannelLayout = Literal["channels_first", "channels_last"] class LocalWellHDF5(Dataset): KNOWN_KEYS = ("t0_fields", "fields", "data", "x", "trajectory") def __init__( self, root: str, max_samples: int = 256, n_steps: Optional[int] = None, expected_channels: int = 2, channel_layout: ChannelLayout = "channels_first", strict: bool = True, ): self.root = Path(root) self.files = sorted(self.root.rglob("*.hdf5")) + sorted(self.root.rglob("*.h5")) if not self.files: raise DataLoadError( f"no .hdf5/.h5 files found under {root}", outcome_code="NO_FILES_FOUND", ) self.max_samples = max_samples self.n_steps = n_steps self.expected_channels = expected_channels self.channel_layout = channel_layout self.strict = strict self.samples: List[torch.Tensor] = [] self._load() def _load(self): try: import h5py except ImportError as e: raise DataLoadError( "h5py is required to load local Well HDF5 files but is not installed", outcome_code="MISSING_DEPENDENCY", ) from e count = 0 rejected = [] for fp in self.files: if count >= self.max_samples: break with h5py.File(fp, "r") as f: arr = None for k in self.KNOWN_KEYS: if k in f: arr = f[k] break if arr is None and list(f.keys()): arr = f[list(f.keys())[0]] if arr is None: rejected.append((str(fp), "no recognizable dataset key")) continue data = np.array(arr) try: if data.ndim == 5: # N, T, C/H, H/C, W or similar for i in range(min(data.shape[0], self.max_samples - count)): traj = self._validate_and_orient(data[i], fp) self.samples.append(torch.from_numpy(traj).float()) count += 1 elif data.ndim == 4: traj = self._validate_and_orient(data, fp) self.samples.append(torch.from_numpy(traj).float()) count += 1 else: rejected.append((str(fp), f"unsupported ndim={data.ndim}")) except SchemaValidationError as e: rejected.append((str(fp), e.detail)) if self.strict: raise if rejected and not self.strict: print(f"[LocalWellHDF5] WARNING: {len(rejected)} file(s) rejected: {rejected}") if not self.samples: raise SchemaValidationError( f"no valid trajectories loaded from {self.root}; " f"rejected files: {rejected}", outcome_code="NO_VALID_TRAJECTORIES", ) print(f"[LocalWellHDF5] loaded {len(self.samples)} trajectories " f"from {len(self.files)} files (schema={self.channel_layout}, " f"channels={self.expected_channels})") def _validate_and_orient(self, data: np.ndarray, fp: Path) -> np.ndarray: if data.ndim != 4: raise SchemaValidationError( f"{fp.name}: expected 4D (T,C,H,W)-like array, got ndim={data.ndim}", outcome_code="WRONG_NDIM", ) T, A, B, C_ = data.shape if self.channel_layout == "channels_first": channel_axis_size = A oriented = data elif self.channel_layout == "channels_last": channel_axis_size = C_ oriented = np.transpose(data, (0, 3, 1, 2)) else: raise SchemaValidationError( f"unknown channel_layout '{self.channel_layout}'", outcome_code="INVALID_LAYOUT_SPEC", ) if channel_axis_size != self.expected_channels: raise SchemaValidationError( f"{fp.name}: declared channel_layout='{self.channel_layout}' " f"implies {channel_axis_size} channels, but " f"expected_channels={self.expected_channels}. " f"Refusing to guess a different layout for this file — " f"pass the correct channel_layout/expected_channels explicitly.", outcome_code="CHANNEL_COUNT_MISMATCH", ) return oriented def __len__(self): return len(self.samples) def __getitem__(self, idx): traj = self.samples[idx] if self.n_steps is not None and traj.size(0) > self.n_steps: traj = traj[: self.n_steps] return {"fields": traj, "idx": idx} def get_dataset( max_samples: int = 128, n_steps: int = 14, expected_channels: int = 2, channel_layout: ChannelLayout = "channels_first", search_roots: Optional[List[str]] = None, ): roots = search_roots or ["./data/real", "./data/well"] existing_dirs = [r for r in roots if os.path.isdir(r)] if not existing_dirs: raise DataLoadError( f"no real-data directory found among candidates: {roots}. " f"Run `the-well-download ...` or place .hdf5 files under one " f"of these paths before training.", outcome_code="NO_DATA_DIRECTORY", ) last_error = None for root in existing_dirs: try: ds = LocalWellHDF5( root, max_samples=max_samples, n_steps=n_steps, expected_channels=expected_channels, channel_layout=channel_layout, strict=True, ) print(f"[data] using REAL local Well data from {root} ({len(ds)} trajs)") return ds, "REAL_LOCAL" except (DataLoadError, SchemaValidationError) as e: last_error = e continue raise last_error or DataLoadError( f"no usable real data found in {existing_dirs}", outcome_code="NO_DATA_DIRECTORY", ) def get_synthetic_dataset(max_samples: int = 128, n_steps: int = 14): """ Explicit, opt-in synthetic data. Never called implicitly by get_dataset(). Returns: (dataset, provenance) where provenance == "SYNTHETIC". """ from .synthetic_fields import SyntheticWellLike print("[data] EXPLICIT synthetic mode requested — not real Well data") ds = SyntheticWellLike(n_samples=max_samples, n_steps=n_steps, height=32, width=32) return ds, "SYNTHETIC"