"""Geo-FNO Elasticity dataset. Reconciled against the authoritative reference files (see ``docs/RECONCILIATION.md``): - Geo-FNO ``elasticity/elas_geofno.py`` (raw .npy axis layout, filenames) - Transolver ``exp_elas.py`` + ``utils/normalizer.py`` (split, output normalizer, loss scale) Raw array layout (the **sample axis is LAST** in every raw ``.npy``): Random_UnitCell_sigma_10.npy : (972, 2000) -> per-node von Mises stress (target) Random_UnitCell_XY_10.npy : (972, 2, 2000) -> node coordinates (input) Random_UnitCell_rr_10.npy : (42, 2000) -> void-boundary radii (geometry params; OOD) Split (verbatim from ``exp_elas.py``): train = first ``ntrain`` samples, test = LAST ``ntest``, out of 2000 total. The middle samples are unused. """ from __future__ import annotations import os from dataclasses import dataclass from typing import Optional import numpy as np import torch from torch.utils.data import Dataset SIGMA_FILE = "Random_UnitCell_sigma_10.npy" XY_FILE = "Random_UnitCell_XY_10.npy" RR_FILE = "Random_UnitCell_rr_10.npy" N_NODES = 972 N_TOTAL = 2000 class UnitTransformer: """Global scalar z-score normalizer for the stress target. Verbatim behavior of Transolver ``utils/normalizer.py::UnitTransformer``: mean/std are reduced over dims ``(0, 1)`` (samples and nodes), giving a single scalar mean and std. ``decode`` is applied to predictions before the metric. """ def __init__(self, x: torch.Tensor): self.mean = x.mean(dim=(0, 1), keepdim=True) self.std = x.std(dim=(0, 1), keepdim=True) + 1e-8 def to(self, device) -> "UnitTransformer": self.mean = self.mean.to(device) self.std = self.std.to(device) return self def encode(self, x: torch.Tensor) -> torch.Tensor: return (x - self.mean) / self.std def decode(self, x: torch.Tensor) -> torch.Tensor: return x * self.std + self.mean def _find_file(data_dir: str, name: str) -> str: """Locate ``name`` under ``data_dir`` (gdown may nest files in a subfolder).""" direct = os.path.join(data_dir, name) if os.path.isfile(direct): return direct for root, _dirs, files in os.walk(data_dir): if name in files: return os.path.join(root, name) raise FileNotFoundError( f"Could not find {name} under {data_dir!r}. " f"Run `make data` (or python -m stress_operator.data.download --out {data_dir})." ) def load_raw_arrays(data_dir: str): """Load and reorient the raw arrays so the sample axis is first. Returns float32 tensors: coords : (N_TOTAL, 972, 2) sigma : (N_TOTAL, 972) physical (de-normalized) stress rr : (N_TOTAL, 42) geometry params (may be absent -> None) """ sigma_np = np.load(_find_file(data_dir, SIGMA_FILE)) # (972, 2000) xy_np = np.load(_find_file(data_dir, XY_FILE)) # (972, 2, 2000) sigma = torch.tensor(sigma_np, dtype=torch.float).permute(1, 0).contiguous() # (2000, 972) coords = torch.tensor(xy_np, dtype=torch.float).permute(2, 0, 1).contiguous() # (2000, 972, 2) rr: Optional[torch.Tensor] = None try: rr_np = np.load(_find_file(data_dir, RR_FILE)) # (42, 2000) rr = torch.tensor(rr_np, dtype=torch.float).permute(1, 0).contiguous() # (2000, 42) except FileNotFoundError: rr = None return coords, sigma, rr @dataclass class Splits: train_coords: torch.Tensor train_sigma: torch.Tensor # physical test_coords: torch.Tensor test_sigma: torch.Tensor # physical normalizer: UnitTransformer train_rr: Optional[torch.Tensor] = None test_rr: Optional[torch.Tensor] = None def build_splits(data_dir: str, ntrain: int = 1000, ntest: int = 200) -> Splits: """Build the train/test split exactly as ``exp_elas.py`` does. train = first ``ntrain``; test = LAST ``ntest``. The normalizer is fit on the (physical) train stress only. """ coords, sigma, rr = load_raw_arrays(data_dir) train_coords = coords[:ntrain] train_sigma = sigma[:ntrain] test_coords = coords[-ntest:] test_sigma = sigma[-ntest:] normalizer = UnitTransformer(train_sigma) train_rr = rr[:ntrain] if rr is not None else None test_rr = rr[-ntest:] if rr is not None else None return Splits( train_coords=train_coords, train_sigma=train_sigma, test_coords=test_coords, test_sigma=test_sigma, normalizer=normalizer, train_rr=train_rr, test_rr=test_rr, ) def build_splits_from_indices(data_dir: str, train_idx, test_idx) -> Splits: """Build a split from explicit sample indices over all 2000 samples (used by the OOD eval). The normalizer is fit on the (physical) train stress only. """ coords, sigma, rr = load_raw_arrays(data_dir) train_idx = torch.as_tensor(np.asarray(train_idx), dtype=torch.long) test_idx = torch.as_tensor(np.asarray(test_idx), dtype=torch.long) train_sigma = sigma[train_idx] normalizer = UnitTransformer(train_sigma) return Splits( train_coords=coords[train_idx], train_sigma=train_sigma, test_coords=coords[test_idx], test_sigma=sigma[test_idx], normalizer=normalizer, train_rr=(rr[train_idx] if rr is not None else None), test_rr=(rr[test_idx] if rr is not None else None), ) class ElasticityDataset(Dataset): """Yields ``(coords, sigma)`` per sample. ``coords`` : (972, 2) node coordinates (model input) ``sigma`` : (972, 1) **physical** stress target Normalization is intentionally *not* applied here: the training loop predicts in normalized space and decodes before the relative-L2 loss (identical to ``exp_elas.py``; see ``docs/RECONCILIATION.md`` ยง4). Keeping physical targets in the dataset makes the de-normalize-before-metric contract explicit and impossible to forget. """ def __init__(self, coords: torch.Tensor, sigma: torch.Tensor): assert coords.shape[0] == sigma.shape[0] self.coords = coords # ensure (n, 972, 1) self.sigma = sigma if sigma.dim() == 3 else sigma.unsqueeze(-1) def __len__(self) -> int: return self.coords.shape[0] def __getitem__(self, idx: int): return self.coords[idx], self.sigma[idx]