stress_operator/eval.py

"""Evaluation: mean test relative-L2 on a split (de-normalized).

Mirrors the eval loop in Transolver ``exp_elas.py``: predictions are decoded with the
(train-fitted) normalizer before the metric; targets are physical.
"""
from __future__ import annotations

import torch
from torch.utils.data import DataLoader

from .losses.relative_l2 import relative_l2


@torch.no_grad()
def evaluate(model, loader: DataLoader, normalizer, device) -> float:
    """Return mean relative-L2 over the loader (physical units)."""
    model.eval()
    total = 0.0
    n = 0
    for coords, sigma in loader:
        coords = coords.to(device)
        sigma = sigma.to(device)  # physical (B, N, 1)
        out = model(coords, None)  # (B, N, 1) normalized
        out = normalizer.decode(out)  # -> physical
        total += relative_l2(out, sigma, reduction="sum").item()
        n += coords.shape[0]
    return total / max(n, 1)