| """ |
| Adapter between the real `the_well` package's WellDataset sample format |
| and this project's `{"fields": (T, C, H, W)}` contract. |
| |
| VERIFIED against the actual installed `the_well` package source this |
| session (not guessed): WellDataset._construct_sample() returns a dict |
| with `input_fields`/`output_fields` keys, shaped `(T, H, W, C)` -- |
| channels LAST -- confirmed both by the library's own docstring comment |
| and by reading `_postprocess_data()`, which flattens all field types |
| (scalar/vector/tensor-order) into one combined channel axis via |
| `.unsqueeze(-1).flatten(...)` then `torch.concatenate(..., dim=-1)`, |
| i.e. the new combined channel dimension is the LAST dimension. |
| |
| Every consumer in this project (FieldNormalizer, MultiScaleEncoder, |
| ReplayBuffer, the collate functions in run_full.py/continual_demo.py) |
| assumes `{"fields": (T, C, H, W)}` -- channels-FIRST, one contiguous |
| trajectory, not split into separate input/output tensors. Without this |
| adapter, real streamed Well data would either crash on first use or |
| (worse) silently broadcast-mismatch through FieldNormalizer's |
| `(-1, 1, 1)` per-channel view, since channels-last (T,H,W,C) sliced per |
| timestep gives (H,W,C), which a channels-first-only broadcast can |
| sometimes "succeed" against incorrectly if H, W, or C happen to share a |
| size -- silent wrong results, not a crash. This adapter exists |
| specifically to prevent that: hard-fail on any shape it doesn't |
| recognize, never guess a layout. |
| """ |
| from __future__ import annotations |
| from typing import Any, Dict, Optional |
|
|
| import torch |
| from torch.utils.data import Dataset |
|
|
| from .provenance import SchemaValidationError |
|
|
|
|
| def well_sample_to_fields(sample: Dict[str, Any], include_output: bool = True) -> torch.Tensor: |
| if "input_fields" not in sample: |
| raise SchemaValidationError( |
| f"WellDataset sample is missing 'input_fields'. Keys present: " |
| f"{sorted(sample.keys())}. This adapter was built against " |
| f"the_well's documented input_fields/output_fields contract; " |
| f"if that has changed, this needs updating -- not guessing " |
| f"a different key.", |
| outcome_code="WELL_SAMPLE_MISSING_INPUT_FIELDS", |
| ) |
|
|
| parts = [sample["input_fields"]] |
| if include_output: |
| if "output_fields" not in sample: |
| raise SchemaValidationError( |
| f"include_output=True but sample has no 'output_fields'. " |
| f"Keys present: {sorted(sample.keys())}.", |
| outcome_code="WELL_SAMPLE_MISSING_OUTPUT_FIELDS", |
| ) |
| parts.append(sample["output_fields"]) |
|
|
| for i, part in enumerate(parts): |
| if not torch.is_tensor(part) or part.dim() != 4: |
| raise SchemaValidationError( |
| f"WellDataset field tensor #{i} has shape " |
| f"{tuple(part.shape) if torch.is_tensor(part) else type(part)}, " |
| f"expected 4D (T, H, W, C). Refusing to guess how to " |
| f"reinterpret it.", |
| outcome_code="WELL_SAMPLE_UNEXPECTED_SHAPE", |
| ) |
|
|
| if len(parts) > 1: |
| hw_c_shapes = {tuple(p.shape[1:]) for p in parts} |
| if len(hw_c_shapes) > 1: |
| raise SchemaValidationError( |
| f"input_fields and output_fields have mismatched (H,W,C): " |
| f"{hw_c_shapes}. Cannot concatenate along time without " |
| f"either dropping data or inventing values -- refusing both.", |
| outcome_code="WELL_SAMPLE_SHAPE_MISMATCH", |
| ) |
|
|
| combined = torch.cat(parts, dim=0) |
| fields = combined.permute(0, 3, 1, 2).contiguous() |
| return fields |
|
|
|
|
| class WellStreamAdapter(Dataset): |
| def __init__(self, well_dataset, include_output: bool = True): |
| self.well_dataset = well_dataset |
| self.include_output = include_output |
| self.provenance = getattr(well_dataset, "provenance", None) |
|
|
| def __len__(self): |
| return len(self.well_dataset) |
|
|
| def __getitem__(self, idx): |
| raw = self.well_dataset[idx] |
| fields = well_sample_to_fields(raw, include_output=self.include_output) |
| return {"fields": fields, "idx": idx} |
|
|
