File size: 4,624 Bytes
029ac7a | 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 | """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()
# Validate the energy implementation independently in float64.
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()
|