File size: 4,186 Bytes
ae73c7f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | """
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) # (T, H, W, C)
fields = combined.permute(0, 3, 1, 2).contiguous() # -> (T, C, H, W)
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}
|