File size: 6,149 Bytes
2c93889 | 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 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | """Tensile2d (PLAID-datasets/Tensile2d, HF, CC-BY-SA-4.0) adapter — SECOND benchmark.
Why: unlike Geo-FNO Elasticity (scalar von Mises target -> our tensor is latent), Tensile2d directly
supervises the full 2D Cauchy stress tensor (sig11, sig22, sig12) per node on an irregular mesh, so our
equilibrium regularizer acts on a SUPERVISED tensor and we can validate the predicted tensor itself.
VERIFIED facts (golden rule, by reading real samples):
- Each sample is a pickled CGNS tree (no PLAID lib needed): GridCoordinates {CoordinateX, CoordinateY}
and PointData {sig11, sig22, sig12, U1, U2, q}; plus input scalars {P, p1..p5} and output scalars.
- 702 rows: samples 0-499 are the labeled `train_500` set (full stress); samples 500-701 (test 200 +
OOD 2) have stress fields WITHHELD (competition) -> only coords + input scalars. So we use the 500
labeled samples and hold out a local test split (no official test ground truth is public).
- Node count varies per sample (6143-11801, mean ~9400). Static, plane-strain, quasistatic.
`build_cache()` writes data/tensile2d/samples/*.npz once. `build_tensile_splits()` loads them, fits
z-score normalizers on the train split, and returns per-sample lists (variable N -> no stacking).
"""
from __future__ import annotations
import glob
import os
from dataclasses import dataclass
from typing import List
import numpy as np
import torch
INPUT_SCALARS = ["P", "p1", "p2", "p3", "p4", "p5"]
def _field(tree, name):
stack = [tree]
while stack:
n = stack.pop()
if isinstance(n, list) and len(n) == 4:
if n[0] == name and isinstance(n[1], np.ndarray):
return n[1]
stack.extend(n[2] or [])
return None
def build_cache(out_dir: str = "data/tensile2d/samples") -> int:
"""Download Tensile2d parquet shards and cache the 500 labeled samples as .npz. Returns count."""
from huggingface_hub import hf_hub_download
import pyarrow.parquet as pq
import pickle
os.makedirs(out_dir, exist_ok=True)
shards = [hf_hub_download("PLAID-datasets/Tensile2d",
f"data/all_samples-0000{i}-of-00002.parquet", repo_type="dataset")
for i in (0, 1)]
idx = 0
saved = 0
for sh in shards:
col = pq.read_table(sh).column("sample")
for i in range(len(col)):
try:
obj = pickle.loads(col[i].as_py())
tree = list(obj["meshes"].values())[0]
x, y = _field(tree, "CoordinateX"), _field(tree, "CoordinateY")
s11, s22, s12 = _field(tree, "sig11"), _field(tree, "sig22"), _field(tree, "sig12")
if any(v is None for v in (x, y, s11, s22, s12)):
idx += 1
continue # withheld test/OOD sample
coords = np.stack([x, y], 1).astype(np.float32)
sigma = np.stack([s11, s22, s12], 1).astype(np.float32)
vm = np.sqrt(s11 ** 2 - s11 * s22 + s22 ** 2 + 3 * s12 ** 2).astype(np.float32)
sc = obj.get("scalars", {})
scin = np.array([float(sc[k]) for k in INPUT_SCALARS], np.float32)
np.savez(os.path.join(out_dir, f"{saved:04d}.npz"),
coords=coords, sigma=sigma, vm=vm, scalars_in=scin)
saved += 1
except Exception:
pass
idx += 1
return saved
@dataclass
class TensileSplits:
train_coords: List[torch.Tensor]
train_sigma: List[torch.Tensor] # scaled by 1/S (N,3)
train_scalars: List[torch.Tensor] # z-scored (6,)
test_coords: List[torch.Tensor]
test_sigma: List[torch.Tensor]
test_scalars: List[torch.Tensor]
scale_S: float # SINGLE global stress scale; de-normalize: sigma*S
scalar_mean: torch.Tensor # (6,)
scalar_std: torch.Tensor
def build_tensile_splits(samples_dir: str = "data/tensile2d/samples",
ntrain: int = 400, ntest: int = 100, seed: int = 0) -> TensileSplits:
"""Load cached samples; deterministic train/test split; z-score normalizers fit on TRAIN only.
Coords are kept RAW (physical units) so the MLS divergence operator measures true distances; the
model conditions on raw coords + normalized input scalars (the encoder MLP handles the coord scale).
Stress is z-scored per component; input scalars z-scored per dim.
"""
files = sorted(glob.glob(os.path.join(samples_dir, "*.npz")))
if not files:
raise FileNotFoundError(f"no cached Tensile2d samples in {samples_dir} — run build_cache() first")
rng = np.random.default_rng(seed)
order = rng.permutation(len(files))
tr_idx, te_idx = order[:ntrain], order[ntrain:ntrain + ntest]
def load(i):
d = np.load(files[i])
return d["coords"], d["sigma"], d["scalars_in"]
tr = [load(i) for i in tr_idx]
te = [load(i) for i in te_idx]
# SINGLE global stress scale S (RMS of all train stress components) so the divergence regularizer
# de-normalizes with one scalar: sigma_phys = S * sigma_norm => div(sigma_phys) = S * div(sigma_norm).
# (Per-component z-score would scale each channel differently and corrupt the physical divergence.)
sig_all = np.concatenate([s for _, s, _ in tr], 0) # (sum N, 3)
S = float(np.sqrt((sig_all ** 2).mean())) # global RMS scale
sc_all = np.stack([sc for _, _, sc in tr], 0) # (ntrain, 6)
sc_mean, sc_std = sc_all.mean(0), sc_all.std(0) + 1e-8
def pack(rows):
cs, ss, scs = [], [], []
for c, s, sc in rows:
cs.append(torch.tensor(c))
ss.append(torch.tensor(s / S))
scs.append(torch.tensor((sc - sc_mean) / sc_std))
return cs, ss, scs
trc, trs, trsc = pack(tr)
tec, tes, tesc = pack(te)
return TensileSplits(
train_coords=trc, train_sigma=trs, train_scalars=trsc,
test_coords=tec, test_sigma=tes, test_scalars=tesc,
scale_S=S,
scalar_mean=torch.tensor(sc_mean), scalar_std=torch.tensor(sc_std),
)
|