| """Data, Sobol, PDE and metric contracts shared by the three entry points.""" |
| import hashlib |
| import json |
| import os |
| import random |
| import socket |
| from pathlib import Path |
|
|
| import h5py |
| import numpy as np |
| import torch |
| import yaml |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def load_config(path): |
| with open(path, encoding="utf-8") as stream: |
| config = yaml.safe_load(stream) |
| if config["training"]["batch_size"] != 64 or config["training"]["epochs"] != 20000: |
| raise ValueError("Formal protocol must retain batch_size=64 and epochs=20000") |
| if config["training"]["early_stopping"]: |
| raise ValueError("User requires complete formal training; early stopping disabled") |
| return config |
|
|
|
|
| def resolved(config, name): |
| path = Path(config["paths"][name]) |
| return path if path.is_absolute() else ROOT / path |
|
|
|
|
| def atomic_json(path, payload): |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| tmp = path.with_suffix(path.suffix + ".tmp") |
| tmp.write_text(json.dumps(payload, ensure_ascii=False, indent=2, allow_nan=False), encoding="utf-8") |
| tmp.replace(path) |
|
|
|
|
| def atomic_save(path, payload): |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| tmp = path.with_suffix(path.suffix + ".tmp") |
| torch.save(payload, tmp) |
| tmp.replace(path) |
|
|
|
|
| def require_compute_node(): |
| |
| jid, name = os.environ.get("SLURM_JOB_ID"), os.environ.get("SLURM_JOB_NAME", "") |
| if not jid or not name.startswith("qcpinn_cavity_") or "login" in socket.gethostname().lower(): |
| raise RuntimeError("Training/smoke requires a new case-owned SLURM compute allocation") |
| return {"job_id": jid, "job_name": name, "hostname": socket.gethostname(), |
| "array_job_id": os.environ.get("SLURM_ARRAY_JOB_ID"), |
| "array_task_id": os.environ.get("SLURM_ARRAY_TASK_ID")} |
|
|
|
|
| def seed_everything(seed, threads=1): |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| torch.set_num_threads(threads) |
|
|
|
|
| def build_model(config): |
| from model.QCPINN import QCPINN |
| m = config["model"] |
| model = QCPINN(**{key: m[key] for key in ["input_dim", "output_dim", "hidden_dim", "n_qubits", "quantum_layers", "backend", "shots"]}) |
| model = model.to(dtype=getattr(torch, m["dtype"]), device=m["device"]) |
| n = sum(p.numel() for p in model.parameters()) |
| if n != m["expected_parameters"]: |
| raise ValueError(f"Parameter count {n} != paper count {m['expected_parameters']}") |
| return model |
|
|
|
|
| def make_sobol_pools(config, seed): |
| s = config["sampling"] |
| pools = {} |
| for index, kind in enumerate(["interior", "lid", "left", "right", "bottom", "initial"]): |
| count = s["interior"] if kind == "interior" else s["initial"] if kind == "initial" else s["per_wall"] |
| engine = torch.quasirandom.SobolEngine(3, scramble=s["scramble"], seed=seed*17+index) |
| xyz = engine.draw(count).to(dtype=getattr(torch, config["model"]["dtype"])) |
| xyz[:, 0] *= config["physics"]["t_max"] |
| if kind == "lid": xyz[:, 2] = 1.0 |
| if kind == "left": xyz[:, 1] = 0.0 |
| if kind == "right": xyz[:, 1] = 1.0 |
| if kind == "bottom": xyz[:, 2] = 0.0 |
| if kind == "initial": xyz[:, 0] = 0.0 |
| pools[kind] = xyz.to(config["model"]["device"]) |
| return pools |
|
|
|
|
| def sample_batches(pools, size, rng): |
| return {key: points[torch.as_tensor(rng.choice(len(points), size=size, replace=False), device=points.device)] |
| for key, points in pools.items()} |
|
|
|
|
| def derivative(field, coords): |
| return torch.autograd.grad(field, coords, torch.ones_like(field), create_graph=True)[0] |
|
|
|
|
| def pde_residuals(model, coordinates, physics): |
| coords = coordinates.detach().clone().requires_grad_(True) |
| uvp = model(coords) |
| u, v, p = uvp.unbind(-1) |
| du, dv, dp = [derivative(field, coords) for field in (u, v, p)] |
| uxx = derivative(du[:, 1], coords)[:, 1] |
| uyy = derivative(du[:, 2], coords)[:, 2] |
| vxx = derivative(dv[:, 1], coords)[:, 1] |
| vyy = derivative(dv[:, 2], coords)[:, 2] |
| ru = du[:, 0] + u*du[:, 1] + v*du[:, 2] + dp[:, 1]/physics["rho"] - physics["mu"]*(uxx+uyy) |
| rv = dv[:, 0] + u*dv[:, 1] + v*dv[:, 2] + dp[:, 2]/physics["rho"] - physics["mu"]*(vxx+vyy) |
| rc = du[:, 1] + dv[:, 2] |
| return torch.stack((ru, rv, rc), dim=-1) |
|
|
|
|
| def loss_terms(model, batches, config): |
| physics = config["physics"] |
| if physics["loss_formula"] != "paper_literal": |
| raise ValueError("This reproduction implements the literal Appendix A.2 losses") |
| residuals = pde_residuals(model, batches["interior"], physics) |
| terms = {"physics": residuals.square().mean(0).sum()} |
| names = ["lid", "left", "right", "bottom", "initial"] |
| sizes = [len(batches[k]) for k in names] |
| fields = torch.split(model(torch.cat([batches[k] for k in names])), sizes) |
| lid, *rest = fields |
| terms["lid"] = (1-lid[:, 0]+lid[:, 1]).square().mean() |
| terms["walls"] = sum(out[:, :2].sum(-1).square().mean() for out in rest[:3]) |
| terms["initial"] = rest[3].sum(-1).square().mean() |
| terms["total"] = sum(physics["loss_weights"][key]*value for key, value in terms.items()) |
| return terms |
|
|
|
|
| def validation_indices(config, total): |
| size = config["validation"]["size"] |
| if size >= total: |
| raise ValueError("Validation must leave nonempty holdout") |
| return np.sort(np.random.default_rng(config["validation"]["seed"]).choice(total, size, replace=False)) |
|
|
|
|
| def load_validation(config): |
| key = config["validation"]["reference_key"] |
| with h5py.File(resolved(config, "reference"), "r") as source: |
| data = source[key] |
| if data.shape[0] != 6: |
| raise ValueError("Expected MATLAB v7.3 dataset of shape (6,N)") |
| ids = validation_indices(config, data.shape[1]) |
| rows = data[:, ids].T |
| if not np.isfinite(rows).all(): |
| raise ValueError("Reference validation contains NaN/Inf") |
| return rows[:, :3], rows[:, 3:], ids |
|
|
|
|
| def relative_metrics(error_squares, reference_squares): |
| metrics = {} |
| for i, key in enumerate(("u", "v", "p")): |
| den, num = float(reference_squares[i]), float(error_squares[i]) |
| value = float(np.sqrt(num/den)) if den > 0 else None |
| metrics[key] = {"relative_l2": value, "relative_l2_percent": None if value is None else 100*value, |
| "error_sum_squares": num, "reference_sum_squares": den, |
| "reason": None if value is not None else "zero_reference_norm"} |
| return metrics |
|
|
|
|
| @torch.no_grad() |
| def predict(model, coords, size): |
| param = next(model.parameters()) |
| return np.concatenate([model(torch.as_tensor(coords[i:i+size], dtype=param.dtype, device=param.device)).cpu().numpy() |
| for i in range(0, len(coords), size)], axis=0) |
|
|
|
|
| def evaluate(model, coordinates, reference, size): |
| was_training = model.training |
| model.eval() |
| predicted = predict(model, coordinates, size).astype(np.float64) |
| model.train(was_training) |
| errors = predicted-reference |
| return relative_metrics(np.sum(errors**2, 0), np.sum(reference**2, 0)) |
|
|
|
|
| def selection_score(metrics): |
| values = [metrics[k]["relative_l2"] for k in ("u", "v", "p")] |
| if any(v is None or not np.isfinite(v) for v in values): |
| raise ValueError("Validation L2 is undefined/nonfinite") |
| return float(np.mean(values)) |
|
|
|
|
| def sha256(path): |
| digest = hashlib.sha256() |
| with open(path, "rb") as stream: |
| for block in iter(lambda: stream.read(1024*1024), b""): digest.update(block) |
| return digest.hexdigest() |
|
|
|
|
| def load_checkpoint(path, model): |
| checkpoint = torch.load(path, map_location=next(model.parameters()).device, weights_only=False) |
| model.load_state_dict(checkpoint["model_state_dict"], strict=True) |
| return checkpoint |
|
|
