crash_solver.py

"""Option 2 - temporal crash-field SURROGATE (CarCrashNet's CrashSolver idea).

Predicts the full crash DEFORMATION TRAJECTORY of a design from its geometry +
impact parameters, the way CarCrashNet's CrashSolver does, so the viewer can
animate a learned time-history (displacement + von Mises) without an FE solve.

Architecture (compact, single-GPU):
    geometry point cloud (N,3)  --PointNet encoder-->  g (latent)
    impact params (v, pole, offset, ...) --MLP-->      p
    [g ; p]  --GRU temporal decoder over T steps-->     per-step latent
              --per-point MLP head-->                   disp_t (N,3), vm_t (N,)

Trains on CarCrashNet VTKHDF trajectories (fetch_carcrashnet.py). It is fully
self-supervised on the released data: input = t0 geometry + scalars, target =
the saved displacement / von-Mises sequence.

Status gate: available() and train() require the VTKHDF dataset to be present.
This module is wired and importable now; it trains once the data is downloaded.
"""
from __future__ import annotations

import glob
import math
from pathlib import Path
from typing import Optional

import numpy as np

HERE = Path(__file__).parent
DATA_DIR = HERE / "DrivAerNet" / "CrashTrajectories"
WEIGHTS_PATH = HERE / "weights" / "crash_solver.pt"

N_POINTS = 2048
N_TIME = 24          # decoded timesteps
N_PARAM = 6          # v, pole, offset, box_t, beam_t, mass (normalised)


def _have_torch() -> bool:
    try:
        import torch  # noqa: F401
        return True
    except Exception:
        return False


def dataset_ready() -> bool:
    for ext in ("*.vtkhdf", "*.hdf", "*.h5"):
        if glob.glob(str(DATA_DIR / "**" / ext), recursive=True):
            return True
    return False


def available() -> bool:
    return WEIGHTS_PATH.exists() and _have_torch()


# ─────────────────────────────────────────────────────────────────
# Model
# ─────────────────────────────────────────────────────────────────
def _build_model():
    import torch
    import torch.nn as nn

    class PointEnc(nn.Module):
        def __init__(self, out=512):
            super().__init__()
            self.mlp = nn.Sequential(
                nn.Conv1d(3, 64, 1), nn.BatchNorm1d(64), nn.ReLU(),
                nn.Conv1d(64, 128, 1), nn.BatchNorm1d(128), nn.ReLU(),
                nn.Conv1d(128, out, 1), nn.BatchNorm1d(out), nn.ReLU())

        def forward(self, x):                       # x (B,N,3)
            x = x.transpose(1, 2)
            return self.mlp(x).max(dim=-1)[0]       # (B,out) global feature

    class CrashSolver(nn.Module):
        """geometry + impact params -> per-node (disp,vm) trajectory over T."""
        def __init__(self, n_time=N_TIME, n_param=N_PARAM):
            super().__init__()
            self.n_time = n_time
            self.enc = PointEnc(512)
            self.pmlp = nn.Sequential(nn.Linear(n_param, 64), nn.ReLU(),
                                      nn.Linear(64, 128), nn.ReLU())
            self.gru = nn.GRU(input_size=1, hidden_size=640, batch_first=True)
            # per-point head: [global(640) ; local point(3)] -> (disp3 + vm1)
            self.head = nn.Sequential(
                nn.Conv1d(640 + 3, 256, 1), nn.ReLU(),
                nn.Conv1d(256, 128, 1), nn.ReLU(),
                nn.Conv1d(128, 4, 1))               # dx,dy,dz,vm

        def forward(self, pts, params):
            B, N, _ = pts.shape
            g = self.enc(pts)                       # (B,512)
            p = self.pmlp(params)                   # (B,128)
            h0 = torch.cat([g, p], dim=1).unsqueeze(0)   # (1,B,640)
            # decode T steps from a ramp input (time embedding)
            ramp = torch.linspace(0, 1, self.n_time, device=pts.device)
            ramp = ramp.view(1, self.n_time, 1).repeat(B, 1, 1)
            seq, _ = self.gru(ramp, h0.contiguous())     # (B,T,640)
            outs = []
            pf = pts.transpose(1, 2)                # (B,3,N)
            for t in range(self.n_time):
                ht = seq[:, t, :].unsqueeze(-1).repeat(1, 1, N)   # (B,640,N)
                feat = torch.cat([ht, pf], dim=1)                 # (B,643,N)
                outs.append(self.head(feat).transpose(1, 2))      # (B,N,4)
            return torch.stack(outs, dim=1)         # (B,T,N,4)

    return CrashSolver


# ─────────────────────────────────────────────────────────────────
# Dataset (VTKHDF trajectory reader)
# ─────────────────────────────────────────────────────────────────
class CrashTrajDataset:
    """Reads CarCrashNet VTKHDF trajectories -> (pts0, params, disp_seq, vm_seq)."""

    def __init__(self, data_dir=None, n_points=N_POINTS, n_time=N_TIME,
                 split="train", val_frac=0.15, seed=42):
        self.n_points = n_points
        self.n_time = n_time
        root = Path(data_dir) if data_dir else DATA_DIR
        files = []
        for ext in ("*.vtkhdf", "*.hdf", "*.h5"):
            files += sorted(glob.glob(str(root / "**" / ext), recursive=True))
        rng = np.random.default_rng(seed)
        order = rng.permutation(len(files))
        nv = max(1, int(round(len(files) * val_frac))) if files else 0
        val = set(order[:nv].tolist())
        self.files = [files[i] for i in range(len(files))
                      if (i in val) == (split == "val")]

    def __len__(self):
        return len(self.files)

    def _read(self, path):
        """Read one VTKHDF trajectory: points0 (P,3), disp (T,P,3), vm (T,P),
        params (n_param,). Uses h5py (VTKHDF is HDF5)."""
        import h5py
        with h5py.File(path, "r") as f:
            # VTKHDF stores PolyData under /VTKHDF; field trajectories under
            # /VTKHDF/PointData/<name> with a /Steps group. Layout varies by
            # export; this reader pulls Points + displacement + vonMises.
            g = f["VTKHDF"] if "VTKHDF" in f else f
            pts0 = np.asarray(g["Points"])[: ]
            pd = g["PointData"]
            disp = np.asarray(pd["displacement"]) if "displacement" in pd else None
            vm = np.asarray(pd["vonMises"]) if "vonMises" in pd else None
        return pts0, disp, vm

    def __getitem__(self, i):
        import torch
        pts0, disp, vm = self._read(self.files[i])
        # subsample points consistently
        P = pts0.shape[0]
        idx = np.random.default_rng(0).choice(P, self.n_points, replace=P < self.n_points)
        p0 = pts0[idx].astype(np.float32)
        c = p0.mean(0); s = (np.abs(p0 - c).max() + 1e-9)
        p0n = (p0 - c) / s
        # build/resample sequences to n_time
        def _seq(arr, last):
            if arr is None:
                return np.zeros((self.n_time, self.n_points, last), np.float32)
            arr = arr.reshape(arr.shape[0], P, -1)[:, idx, :last]
            ti = np.linspace(0, arr.shape[0] - 1, self.n_time).round().astype(int)
            return (arr[ti] / s).astype(np.float32)
        d = _seq(disp, 3)
        v = _seq(vm, 1)[..., 0]
        params = np.zeros(N_PARAM, np.float32)      # filled from file metadata if present
        return (torch.from_numpy(p0n), torch.from_numpy(params),
                torch.from_numpy(d), torch.from_numpy(v))


def train(data_dir=None, epochs=60, batch=4, lr=1e-3, device="auto"):
    if not _have_torch():
        raise RuntimeError("torch required")
    if not dataset_ready():
        raise RuntimeError(
            "No CarCrashNet VTKHDF trajectories found. Download them "
            "(python fetch_carcrashnet.py) into DrivAerNet/CrashTrajectories first.")
    import torch
    import torch.nn as nn
    from torch.utils.data import DataLoader
    if device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"
    tr = CrashTrajDataset(data_dir, split="train")
    va = CrashTrajDataset(data_dir, split="val")
    print(f"[crash-solver] train {len(tr)} val {len(va)} on {device}")
    tl = DataLoader(tr, batch_size=batch, shuffle=True, drop_last=True)
    vl = DataLoader(va, batch_size=batch)
    m = _build_model()().to(device)
    opt = torch.optim.AdamW(m.parameters(), lr=lr, weight_decay=1e-4)
    huber = nn.SmoothL1Loss()
    best = math.inf
    WEIGHTS_PATH.parent.mkdir(parents=True, exist_ok=True)
    for ep in range(1, epochs + 1):
        m.train(); tot = 0; nb = 0
        for p0, pr, d, v in tl:
            p0, pr, d, v = [x.to(device) for x in (p0, pr, d, v)]
            opt.zero_grad()
            out = m(p0, pr)                          # (B,T,N,4)
            loss = huber(out[..., :3], d) + 0.5 * huber(out[..., 3], v)
            loss.backward(); opt.step()
            tot += loss.item(); nb += 1
        # val
        m.eval(); vtot = 0; vb = 0
        with torch.no_grad():
            for p0, pr, d, v in vl:
                p0, pr, d, v = [x.to(device) for x in (p0, pr, d, v)]
                out = m(p0, pr)
                vtot += (huber(out[..., :3], d) + 0.5 * huber(out[..., 3], v)).item(); vb += 1
        vl_loss = vtot / max(vb, 1)
        print(f"[crash-solver] ep {ep:3d} train {tot/max(nb,1):.5f} val {vl_loss:.5f}")
        if vl_loss < best:
            best = vl_loss
            torch.save(m.state_dict(), WEIGHTS_PATH)
    print(f"[crash-solver] done. best val {best:.5f} -> {WEIGHTS_PATH}")
    return {"best_val": best}


if __name__ == "__main__":
    import argparse
    p = argparse.ArgumentParser()
    sub = p.add_subparsers(dest="cmd", required=True)
    pt = sub.add_parser("train"); pt.add_argument("--data", default=None)
    pt.add_argument("--epochs", type=int, default=60); pt.add_argument("--batch", type=int, default=4)
    sub.add_parser("info")
    a = p.parse_args()
    if a.cmd == "train":
        train(a.data, a.epochs, a.batch)
    else:
        print("dataset ready:", dataset_ready(), " weights:", WEIGHTS_PATH.exists())
        if _have_torch():
            n = sum(x.numel() for x in _build_model()().parameters())
            print(f"CrashSolver params: {n:,}")