| """Scientific invariants and six required smoke checks; never writes formal weights.""" |
| from __future__ import annotations |
|
|
| import json |
| import time |
|
|
| import numpy as np |
| import torch |
|
|
| from common import (DarcyEnergy, analytical_coefficient_tensor, arguments, build_model, configuration, |
| element_corners, evaluate, initialize, relative_l2, read_data, output_path, write_json) |
|
|
|
|
| def gauss_energy(p, u, h): |
| """Independent high-order quadrature, only for verification of the analytical loss.""" |
| nodes, weights = np.polynomial.legendre.leggauss(3) |
| corners_p, corners_u = element_corners(p), element_corners(u) |
| answer = torch.zeros(p.shape[0], dtype=torch.float64) |
| for gx, wx in zip((nodes + 1)/2, weights/2): |
| for gy, wy in zip((nodes + 1)/2, weights/2): |
| shape = torch.tensor([(1-gx)*(1-gy), gx*(1-gy), gx*gy, (1-gx)*gy]) |
| dx = torch.tensor([-(1-gy), 1-gy, gy, -gy]) / h |
| dy = torch.tensor([-(1-gx), -gx, gx, 1-gx]) / h |
| coefficient = (corners_p * shape).sum(-1) |
| value = 0.5 * coefficient * ((corners_u*dx).sum(-1)**2+(corners_u*dy).sum(-1)**2) - (corners_u*shape).sum(-1) |
| answer += value.sum((1,2)) * h**2 * wx * wy |
| return answer |
|
|
|
|
| def main(): |
| args = arguments(__doc__).parse_args() |
| cfg = configuration(args) |
| device = initialize(cfg) |
| fields, metadata = read_data(cfg) |
| checks = {} |
| start = time.perf_counter() |
| |
| exact = DarcyEnergy(5).double() |
| p = torch.rand(2,5,5,1,dtype=torch.float64)+1 |
| u = torch.rand_like(p) |
| analytical = exact(u, p, reduction="none") |
| quadrature = gauss_energy(p, u, 0.25) |
| torch.testing.assert_close(analytical, quadrature, rtol=2e-7, atol=1e-8) |
| constant_stiffness = analytical_coefficient_tensor(1,1).sum(0) |
| expected = torch.tensor([[4,-1,-2,-1],[-1,4,-1,-2],[-2,-1,4,-1],[-1,-2,-1,4]],dtype=torch.float64)/6 |
| torch.testing.assert_close(constant_stiffness, expected) |
| grad_u = torch.rand(1,3,3,1,dtype=torch.float64,requires_grad=True) |
| torch.autograd.gradcheck(lambda z: DarcyEnergy(3).double()(z,torch.ones_like(z)), (grad_u,)) |
| checks["analytical_energy"] = {"status":"passed", "quadrature_max_difference": (analytical-quadrature).abs().max().item(), "constant_stiffness":"passed", "autograd_gradcheck":"passed"} |
| model = build_model(cfg,device) |
| objective = DarcyEnergy(64).to(device) |
| coefficient = fields["P_train"][:2].to(device) |
| output = model(coefficient) |
| assert output.shape == coefficient.shape and torch.isfinite(output).all() |
| assert output[:,0].abs().max() == 0 and output[:,-1].abs().max() == 0 |
| assert output[:,:,0].abs().max() == 0 and output[:,:,-1].abs().max() == 0 |
| checks["forward_pass"] = "passed" |
| objective(output,coefficient).backward() |
| assert all(param.grad is not None and torch.isfinite(param.grad).all() for param in model.parameters()) |
| checks["backward_pass"] = "passed" |
| optimizer = torch.optim.Adam(model.parameters(),lr=cfg["training"]["learning_rate"]) |
| model.train() |
| losses=[] |
| for batch in fields["P_train"].split(cfg["training"]["batch_size"]): |
| batch = batch.to(device) |
| optimizer.zero_grad(set_to_none=True) |
| loss = objective(model(batch), batch) |
| assert torch.isfinite(loss) |
| loss.backward() |
| optimizer.step() |
| losses.append(loss.item()) |
| assert losses[-1] < losses[0] |
| checks["train_loop_dry_run"] = {"status":"passed", "samples":1000, "first_energy":losses[0],"last_energy":losses[-1]} |
| _, errors, _ = evaluate(model, fields["P_test"].to(device),fields["U_test"].to(device),objective,100) |
| assert len(errors)==100 and torch.isfinite(errors).all() |
| checks["validation_loop_dry_run"]="passed" |
| reference=fields["U_test"][:3] |
| torch.testing.assert_close(relative_l2(reference,reference),torch.zeros(3)) |
| torch.testing.assert_close(relative_l2(2*reference,reference),torch.ones(3)) |
| checks["domain_metric_dry_run"]="passed" |
| assert cfg["training"]["epochs"]==500 and cfg["model"]["depth"]==4 and cfg["model"]["width"]==32 and cfg["model"]["modes"]==12 |
| checks["config_consistency"]="passed" |
| report={"status":"passed","attempts":1,"checks":checks,"data":metadata,"seconds":time.perf_counter()-start, |
| "device":str(device),"device_name":torch.cuda.get_device_name() if device.type=="cuda" else "cpu"} |
| write_json(output_path(cfg["output"]["results_dir"])/"verification.json",report) |
| print(json.dumps(report,indent=2),flush=True) |
|
|
|
|
| if __name__=="__main__": |
| main() |
|
|
