| """Configuration, exact MAT contract, sampling and checkpoint utilities.""" |
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| import random |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import yaml |
| from scipy.io import loadmat |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def read_config(path=None): |
| path = Path(path) if path else ROOT / "config/config.yaml" |
| with path.open() as handle: |
| cfg = yaml.safe_load(handle) |
| if cfg["model"]["hidden_layers"] != [50] * 4: |
| raise ValueError("Formal experiment requires exactly four 50-unit hidden layers") |
| if cfg["training"]["iterations"] != 4000: |
| raise ValueError("Formal experiment requires exactly 4000 iterations") |
| return cfg |
|
|
|
|
| def resolve(path): |
| p = Path(path) |
| return p if p.is_absolute() else ROOT / p |
|
|
|
|
| def save_json(path, value): |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| tmp = path.with_suffix(path.suffix + ".tmp") |
| tmp.write_text(json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False)) |
| tmp.replace(path) |
|
|
|
|
| def setup(cfg, device=None): |
| seed = int(cfg["seed"]) |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| torch.set_num_threads(int(cfg["runtime"]["cpu_threads"])) |
| dtype = getattr(torch, cfg["model"]["dtype"]) |
| torch.set_default_dtype(dtype) |
| device = torch.device(device or cfg["runtime"]["device"]) |
| if device.type == "cuda" and not torch.cuda.is_available(): |
| raise RuntimeError("Requested accelerator is unavailable; run inside SLURM allocation") |
| return device, dtype |
|
|
|
|
| def build_model(cfg, device): |
| from model.TSAPINN import TSAPINN |
| m = cfg["model"] |
| return TSAPINN(**{k: m[k] for k in ["hidden_layers", "sigma", "zeta_sin", "zeta_cos", "reynolds", "bias_initialization"]}).to(device) |
|
|
|
|
| def load_dns(cfg): |
| path = Path(cfg["data"]["path"]) |
| if not path.is_file(): |
| raise FileNotFoundError(f"Required author DNS data absent: {path}") |
| raw = loadmat(path) |
| for key, shape in cfg["data"]["expected_shapes"].items(): |
| if raw[key].shape != tuple(shape) or not np.isfinite(raw[key]).all(): |
| raise ValueError(f"Invalid DNS field {key}: expected finite {shape}") |
| xy = np.asarray(raw["X_star"], dtype=np.float64) |
| times = np.asarray(raw["t"], dtype=np.float64).ravel() |
| if not np.all(np.diff(times) > 0): |
| raise ValueError("DNS time must be strictly increasing") |
| return {"xy": xy, "t": times, "uv": np.asarray(raw["U_star"]).transpose(2, 0, 1), |
| "p": np.asarray(raw["p_star"]).T, "path": str(path), |
| "sha256": hashlib.sha256(path.read_bytes()).hexdigest()} |
|
|
|
|
| def sample_indices(cfg, dns): |
| """Literal paper per-time counts, geometric boundary; fixed throughout training. |
| |
| Only 296 unique boundary nodes exist. Balanced sampling with replacement |
| provides 100 direct DNS observations per side without inventing new labels. |
| Boundary corner duplication and replacement are explicit implementation choices. |
| """ |
| rng = np.random.RandomState(cfg["seed"]) |
| xy, times = dns["xy"], dns["t"] |
| xmin, ymin = xy.min(0) |
| xmax, ymax = xy.max(0) |
| sides = [np.flatnonzero(np.isclose(xy[:, 0], xmin)), np.flatnonzero(np.isclose(xy[:, 0], xmax)), |
| np.flatnonzero(np.isclose(xy[:, 1], ymin)), np.flatnonzero(np.isclose(xy[:, 1], ymax))] |
| boundary = np.unique(np.concatenate(sides)) |
| interior = np.setdiff1d(np.arange(len(xy)), boundary) |
| nt, tr = len(times), cfg["training"] |
| nb, nr = int(tr["boundary_points_per_time"]), int(tr["residual_points_per_time"]) |
| if nb % 4: |
| raise ValueError("Balanced side sampling needs boundary count divisible by 4") |
| bc = np.stack([np.concatenate([rng.choice(side, nb // 4, replace=True) for side in sides]) for _ in times]) |
| res = np.stack([rng.choice(interior, nr, replace=False) for _ in times]) |
| val = np.stack([rng.choice(np.setdiff1d(interior, res[k]), tr["validation_points_per_time"], replace=False) for k in range(nt)]) |
| return {"boundary": bc, "residual": res, "validation": val} |
|
|
|
|
| def gather(dns, indices): |
| nt, n = indices.shape |
| xy = dns["xy"][indices] |
| time = np.broadcast_to(dns["t"][:, None, None], (nt, n, 1)) |
| coordinates = np.concatenate((xy, time), 2).reshape(-1, 3) |
| labels = dns["uv"][np.arange(nt)[:, None], indices].reshape(-1, 2) |
| return coordinates, labels |
|
|
|
|
| def tensor(array, device, dtype): |
| return torch.as_tensor(np.ascontiguousarray(array), dtype=dtype, device=device) |
|
|
|
|
| def predict(model, coordinates, device, dtype, batch_size): |
| values = [] |
| model.eval() |
| with torch.no_grad(): |
| for start in range(0, len(coordinates), batch_size): |
| q = tensor(coordinates[start:start + batch_size], device, dtype) |
| values.append(model.flow_taylor(q).cpu().numpy()) |
| return np.concatenate(values) |
|
|
|
|
| def relative_l2(prediction, reference): |
| difference = np.asarray(prediction, dtype=np.float64) - reference |
| numerator = np.sum(difference * difference, axis=0, dtype=np.float64) |
| denominator = np.sum(np.asarray(reference, dtype=np.float64)**2, axis=0, dtype=np.float64) |
| return np.sqrt(numerator / denominator) |
|
|
|
|
| def atomic_checkpoint(path, payload): |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| tmp = path.with_suffix(".tmp") |
| torch.save(payload, tmp) |
| tmp.replace(path) |
|
|
|
|
| def load_checkpoint(cfg, device, path=None): |
| path = resolve(path or cfg["outputs"]["checkpoint"]) |
| checkpoint = torch.load(path, map_location=device, weights_only=False) |
| saved_config = checkpoint["config"].copy() |
| saved_config["data"] = saved_config["data"].copy() |
| saved_config["data"]["path"] = cfg["data"]["path"] |
| if saved_config != cfg: |
| raise ValueError("Checkpoint configuration does not match requested configuration") |
| model = build_model(cfg, device) |
| model.load_state_dict(checkpoint["model_state_dict"], strict=True) |
| model.eval() |
| return model, checkpoint |
|
|
