stress_operator/data/dataset.py

"""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]