| |
| """Six-path smoke checks and independent derivative/gradient equivalence checks.""" |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| import time |
| from pathlib import Path |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
| import numpy as np |
| import torch |
|
|
| from scripts.common import (ROOT, build_model, gather, load_dns, predict, read_config, |
| relative_l2, sample_indices, save_json, setup, tensor) |
| from scripts.train import objective |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--device", default="cuda") |
| parser.add_argument("--benchmark", action="store_true") |
| args = parser.parse_args() |
| cfg = read_config() |
| device, dtype = setup(cfg, args.device) |
| dns = load_dns(cfg) |
| index = sample_indices(cfg, dns) |
| bq, buv = gather(dns, index["boundary"]) |
| rq, _ = gather(dns, index["residual"]) |
| model = build_model(cfg, device) |
| q = tensor(rq[:32], device, dtype) |
| checks = {} |
| output = model(q) |
| assert output.shape == (32, 2) and torch.isfinite(output).all() |
| checks["forward_pass"] = "passed" |
| values, residuals, div = model.residual_autograd(q) |
| tvalues, tresiduals, _ = model.residual_taylor(q) |
| torch.testing.assert_close(tvalues, values, atol=1e-10, rtol=1e-9) |
| torch.testing.assert_close(tresiduals, residuals, atol=1e-10, rtol=1e-9) |
| params = tuple(model.parameters()) |
| aloss = residuals.square().mean() + values[:, :2].square().mean() + 0.25 * model.slope_recovery() |
| tloss = tresiduals.square().mean() + tvalues[:, :2].square().mean() + 0.25 * model.slope_recovery() |
| ag = torch.autograd.grad(aloss, params) |
| tg = torch.autograd.grad(tloss, params) |
| for left, right in zip(ag, tg): |
| torch.testing.assert_close(left, right, atol=2e-10, rtol=2e-8) |
| assert float(div.abs().max()) < 1e-10 |
| checks["backward_pass"] = "passed: finite parameter gradients; Taylor/reverse-mode equivalence" |
| bq_t, buv_t, rq_t = (tensor(x, device, dtype) for x in (bq[:64], buv[:64], rq[:64])) |
| opt = torch.optim.Adam(model.parameters(), lr=cfg["training"]["learning_rate"], eps=1e-7) |
| losses = [] |
| for _ in range(3): |
| opt.zero_grad(set_to_none=True) |
| value = objective(model, bq_t, buv_t, rq_t, cfg, backward=True) |
| opt.step() |
| losses.append(value["loss"]) |
| assert np.isfinite(losses).all() and losses[-1] < losses[0], losses |
| checks["train_loop_dry_run"] = {"status": "passed", "losses": losses} |
| valq, valuv = gather(dns, index["validation"]) |
| valp = predict(model, valq[:128], device, dtype, 64) |
| error = relative_l2(valp[:, :2], valuv[:128]) |
| assert np.isfinite(error).all() |
| checks["validation_loop_dry_run"] = {"status": "passed", "relative_l2": error.tolist()} |
| assert np.allclose(relative_l2(dns["uv"].reshape(-1, 2), dns["uv"].reshape(-1, 2)), 0) |
| assert np.allclose(relative_l2(np.zeros_like(dns["uv"].reshape(-1, 2)), dns["uv"].reshape(-1, 2)), 1) |
| checks["domain_metric_dry_run"] = "passed: L2 self=0, zero predictor=1; streamfunction divergence <1e-10" |
| assert len(model.frequencies) == 4 and sum(f.numel() for f in model.frequencies) == 200 |
| assert [l.out_features for l in model.layers] == [50, 50, 50, 50, 2] |
| assert index["boundary"].shape == (200, 400) and index["residual"].shape == (200, 1000) |
| checks["config_consistency"] = "passed: 4x50, 200 frequencies, (x,y,t)->(psi,p), 4000 configured updates, per-time counts" |
| benchmark = None |
| if args.benchmark: |
| model = build_model(cfg, device) |
| fullb, fulluv, fullr = (tensor(x, device, dtype) for x in (bq, buv, rq)) |
| started = time.perf_counter() |
| model.zero_grad(set_to_none=True) |
| value = objective(model, fullb, fulluv, fullr, cfg, backward=True) |
| if device.type == "cuda": |
| torch.cuda.synchronize() |
| benchmark = {"seconds_one_full_objective_backward": time.perf_counter()-started, "score": value, |
| "peak_device_memory_bytes": torch.cuda.max_memory_allocated() if device.type == "cuda" else 0} |
| result = {"status": "passed", "attempt": 1, "checks": checks, "benchmark": benchmark, |
| "parameter_count": sum(p.numel() for p in model.parameters()), "torch": torch.__version__, |
| "device": torch.cuda.get_device_name() if device.type == "cuda" else "CPU"} |
| save_json(ROOT / ".onescience/smoke_result.json", result) |
| print(json.dumps(result, ensure_ascii=False, indent=2), flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
