r""" Scientific ML Studio -- reference code (written by hand, not generated by the Studio) Licence : Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) https://creativecommons.org/licenses/by-nc/4.0/ (c) 2026 Scientific ML Studio. All rights are reserved to the Scientific ML Studio platform except the permissions granted by the licence above. You may share and adapt this code for NON-COMMERCIAL purposes, provided you credit "Scientific ML Studio" with a link to https://scimlstudio.com/, and indicate any changes you make. Commercial use, paid client work and production deployment need prior written permission, which you can ask for through the site's Contact page. ------------------------------------------------------------------------------- Scientific ML Studio -- lid-driven cavity, final script (training + post-processing) Licence : Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) https://creativecommons.org/licenses/by-nc/4.0/ (c) 2026 Scientific ML Studio. All rights are reserved to the Scientific ML Studio platform except the permissions granted by the licence above. You may share and adapt this code for NON-COMMERCIAL purposes, provided you credit "Scientific ML Studio" with a link to https://scimlstudio.com/, and indicate any changes you make. Commercial use, paid client work and production deployment need prior written permission, which you can ask for through the site's Contact page. ------------------------------------------------------------------------------- Problem Steady incompressible flow in the unit square, Reynolds number Re = 100. u u_x + v u_y + p_x = (u_xx + u_yy) / Re u v_x + v v_y + p_y = (v_xx + v_yy) / Re u_x + v_y = 0 No-slip (u = v = 0) on the left, right and bottom walls. The lid (y = 1) moves with u = 16 x^2 (1 - x)^2, v = 0, a smooth profile that is zero at both corners and 1 at the middle. The pressure has no boundary condition: it is only defined up to a constant, so every pressure comparison below is made after subtracting the mean over the cavity. What this file does 1. Trains a PINN (inputs x, y; outputs u, v, p) with Adam, then with L-BFGS. 2. Computes a finite element reference (Taylor-Hood Q2-Q1, Newton iteration). 3. Compares the two: contours of u, v, p, speed, vorticity and the stream function, error maps, streamlines, centreline and lid profiles, parity plots, residual maps, loss curves, error against epoch, and GIFs of the training and of tracer particles. 4. Writes the metrics to metrics.json, metrics.csv and metrics_table.md, and bundles everything into cavity_results.zip. Start from the Studio file lid-driven-cavity.py. Changes made to it: * The three placeholder residuals are replaced by the Navier-Stokes equations above. * torch.manual_seed is called before the network is built (the Studio file built it first). * The boundary terms are labelled correctly (the Studio labels were doubled and shifted; the physics was the same). * A second L-BFGS stage is added, with the boundary points frozen during that stage. * The field is saved every SNAP_EVERY epochs (snapshots.npz) and the weights are also saved as plain numpy arrays (weights.npz). * plot_field and write_report only used output 0 and the first residual; they are replaced by the post-processing below. Run it Colab : paste the whole file into one cell and run it (set the runtime to GPU). First a quick check: main(["--smoke"]) (about a minute) Local : python lid_driven_cavity_final.py [--smoke] [--epochs N] [--lbfgs-rounds N] ... Again without training (reuse model.pt): python ... --postprocess-only This file has not been run with PyTorch by its author. The finite element reference and the whole post-processing were tested on their own; use --smoke first. """ from __future__ import annotations import argparse import csv import json import math import time import zipfile from pathlib import Path import numpy as np import torch import torch.nn as nn # TARGET_ENV is fixed at generation time from the Train block; IN_COLAB is # decided at import time by whether `google.colab` is importable, so this # same file runs unmodified as `python this_file.py` on a local machine and # as a pasted-in cell on Colab, without anything to edit by hand. TARGET_ENV = 'auto' # "auto" | "local" | "colab" try: import google.colab # noqa: F401 _COLAB_MODULE_PRESENT = True except ImportError: _COLAB_MODULE_PRESENT = False IN_COLAB = _COLAB_MODULE_PRESENT if TARGET_ENV == "auto" else TARGET_ENV == "colab" try: # figures are optional; training is not from matplotlib.figure import Figure from matplotlib.backends.backend_agg import FigureCanvasAgg from matplotlib.colors import Normalize from matplotlib.collections import LineCollection from matplotlib import colormaps HAVE_MPL = True except ImportError: HAVE_MPL = False # -------------------------------------------------------------------------- # Configuration # -------------------------------------------------------------------------- TITLE = 'Lid-driven cavity' SEED = 0 CHECKPOINT = 'model.pt' COORDS = ['x', 'y'] DIM = 2 SPATIAL_DIM = 2 X_MIN = 0.0 X_MAX = 1.0 Y_MIN = 0.0 Y_MAX = 1.0 RE = 100.0 # Reynolds number: lid speed 1, cavity size 1, viscosity 1 / RE N_INTERIOR = 10000 INTERIOR_STRATEGY = 'uniform' N_BC = 200 # boundary points per wall and per epoch N_INITIAL = 200 EPOCHS = 20000 # Adam epochs LEARNING_RATE = 0.001 LBFGS_ROUNDS = 300 # L-BFGS stage: each round is up to LBFGS_INNER iterations LBFGS_INNER = 20 LOG_EVERY = 200 RESAMPLE_EVERY = 0 BC_WEIGHT = 1.0 # weight on the eight boundary terms (1.0 = as in the Studio file) SNAP_EVERY = 500 # Adam epochs between field snapshots SNAP_EVERY_LBFGS = 10 # L-BFGS rounds between snapshots SNAP_N = 101 # snapshot grid is SNAP_N x SNAP_N FEM_N = 128 # reference mesh: FEM_N x FEM_N squares (Q2-Q1 elements) POST_N = 201 # post-processing grid is POST_N x POST_N (must be odd) DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") torch.manual_seed(SEED) # before the network is built, so the initial weights are reproducible np.random.seed(SEED) OUTPUT_DIR = Path(".") CHECKPOINT = str(OUTPUT_DIR / CHECKPOINT) FIGURE_DIR = OUTPUT_DIR / 'figures' FIGURE_DPI = 150 # -------------------------------------------------------------------------- # Sampling # # Every strategy produces points in the unit cube; the geometry maps them. # -------------------------------------------------------------------------- def unit_samples(n: int, dim: int, strategy: str = "uniform") -> torch.Tensor: """n x dim points in [0, 1]^dim.""" n = max(int(n), 1) if strategy == "grid": per = max(int(round(n ** (1.0 / dim))), 2) axes = [torch.linspace(0.0, 1.0, per) for _ in range(dim)] mesh = torch.meshgrid(*axes, indexing="ij") return torch.stack([m.reshape(-1) for m in mesh], dim=1) if strategy == "sobol": engine = torch.quasirandom.SobolEngine(dimension=dim, scramble=True, seed=SEED) return engine.draw(n).to(torch.get_default_dtype()) if strategy == "latin_hypercube": cols = [] for _ in range(dim): perm = torch.randperm(n).to(torch.get_default_dtype()) cols.append((perm + torch.rand(n)) / n) return torch.stack(cols, dim=1) return torch.rand(n, dim) def as_points(columns: list[torch.Tensor], grad: bool = True) -> torch.Tensor: """Assemble coordinate columns into one leaf tensor on the run device.""" pts = torch.cat([c.reshape(-1, 1) for c in columns], dim=1).to(DEVICE) return pts.requires_grad_(grad) def sample_interior(n: int = N_INTERIOR, strategy: str = INTERIOR_STRATEGY) -> torch.Tensor: """Collocation points strictly inside the domain.""" u = unit_samples(n, 2, strategy) x = X_MIN + (X_MAX - X_MIN) * u[:, 0] y = Y_MIN + (Y_MAX - Y_MIN) * u[:, 1] return as_points([x, y]) # During the L-BFGS stage the boundary points must not change between function evaluations # (a line search needs a fixed objective), so the samplers can be frozen to one draw. _BC_FROZEN = {"on": False, "store": {}} def _freezable(fn): name = fn.__name__ def wrapper(*args, **kwargs): if not _BC_FROZEN["on"]: return fn(*args, **kwargs) store = _BC_FROZEN["store"] if name not in store: pts, normals = fn(*args, **kwargs) store[name] = (pts.detach().clone().requires_grad_(True), normals) return store[name] wrapper.__name__ = name wrapper.__doc__ = fn.__doc__ return wrapper @_freezable def sample_bc1(n: int = N_BC, strategy: str = 'uniform'): """BC1: x = x_min (left wall). Returns (points, outward normals).""" u = unit_samples(n, 1, strategy) x = torch.full((u.shape[0],), X_MIN) y = Y_MIN + (Y_MAX - Y_MIN) * u[:, 0] normals = torch.tensor([[-1.0, 0.0]], device=DEVICE).expand(u.shape[0], SPATIAL_DIM) return as_points([x, y]), normals @_freezable def sample_bc2(n: int = N_BC, strategy: str = 'uniform'): """BC2: x = x_max (right wall). Returns (points, outward normals).""" u = unit_samples(n, 1, strategy) x = torch.full((u.shape[0],), X_MAX) y = Y_MIN + (Y_MAX - Y_MIN) * u[:, 0] normals = torch.tensor([[1.0, 0.0]], device=DEVICE).expand(u.shape[0], SPATIAL_DIM) return as_points([x, y]), normals @_freezable def sample_bc3(n: int = N_BC, strategy: str = 'uniform'): """BC3: y = y_min (bottom wall). Returns (points, outward normals).""" u = unit_samples(n, 1, strategy) x = X_MIN + (X_MAX - X_MIN) * u[:, 0] y = torch.full((u.shape[0],), Y_MIN) normals = torch.tensor([[0.0, -1.0]], device=DEVICE).expand(u.shape[0], SPATIAL_DIM) return as_points([x, y]), normals @_freezable def sample_bc4(n: int = N_BC, strategy: str = 'uniform'): """BC4: y = y_max (lid). Returns (points, outward normals).""" u = unit_samples(n, 1, strategy) x = X_MIN + (X_MAX - X_MIN) * u[:, 0] y = torch.full((u.shape[0],), Y_MAX) normals = torch.tensor([[0.0, 1.0]], device=DEVICE).expand(u.shape[0], SPATIAL_DIM) return as_points([x, y]), normals # -------------------------------------------------------------------------- # Network # -------------------------------------------------------------------------- LOWER = torch.tensor([0.0, 0.0]) UPPER = torch.tensor([1.0, 1.0]) HIDDEN = [50, 50, 50, 50] class PINN(nn.Module): """Trial function (x, y) -> (u, v, p).""" def __init__(self): super().__init__() self.encoding = None widths = [DIM] + HIDDEN + [3] layers = [] for i in range(len(widths) - 1): layers.append(nn.Linear(widths[i], widths[i + 1])) if i < len(widths) - 2: layers.append(nn.Tanh()) self.net = nn.Sequential(*layers) self.apply(self._init_weights) @staticmethod def _init_weights(module): if not isinstance(module, nn.Linear): return nn.init.xavier_normal_(module.weight) if module.bias is not None: nn.init.zeros_(module.bias) def forward(self, pts): lower = LOWER.to(pts.device) upper = UPPER.to(pts.device) z = 2.0 * (pts - lower) / torch.clamp(upper - lower, min=1e-12) - 1.0 return self.net(z) net = PINN().to(DEVICE) def grad_of(y: torch.Tensor, pts: torch.Tensor) -> torch.Tensor: """d y / d pts, as one tensor with a column per coordinate.""" if not y.requires_grad: return pts * 0.0 gradient = torch.autograd.grad( y, pts, grad_outputs=torch.ones_like(y), create_graph=True, retain_graph=True, allow_unused=True )[0] return pts * 0.0 if gradient is None else gradient + pts * 0.0 def field(pts: torch.Tensor, component: int = 0) -> torch.Tensor: """One scalar output of the network, shaped (n, 1).""" return net(pts)[:, component:component + 1] # -------------------------------------------------------------------------- # Loss terms # -------------------------------------------------------------------------- def reduce_term(r: torch.Tensor) -> torch.Tensor: """Mean square residual.""" return torch.mean(r ** 2) def navier_stokes(pts: torch.Tensor, full: bool = False): """Continuity, x-momentum and y-momentum residuals at pts = (x, y). With full=True a fourth value is returned: the fields and first derivatives used. """ out = net(pts) u, v, p = out[:, 0:1], out[:, 1:2], out[:, 2:3] gu, gv, gp = grad_of(u, pts), grad_of(v, pts), grad_of(p, pts) u_x, u_y = gu[:, 0:1], gu[:, 1:2] v_x, v_y = gv[:, 0:1], gv[:, 1:2] u_xx = grad_of(u_x, pts)[:, 0:1] u_yy = grad_of(u_y, pts)[:, 1:2] v_xx = grad_of(v_x, pts)[:, 0:1] v_yy = grad_of(v_y, pts)[:, 1:2] continuity = u_x + v_y momentum_x = u * u_x + v * u_y + gp[:, 0:1] - (u_xx + u_yy) / RE momentum_y = u * v_x + v * v_y + gp[:, 1:2] - (v_xx + v_yy) / RE if not full: return continuity, momentum_x, momentum_y aux = dict(u=u, v=v, p=p, u_x=u_x, u_y=u_y, v_x=v_x, v_y=v_y) return continuity, momentum_x, momentum_y, aux def lid_profile(pts: torch.Tensor) -> torch.Tensor: x = pts[:, 0:1] return (16.0 * (x ** 2.0)) * ((1.0 - x) ** 2.0) # term name -> (sampler, network output, target); target None means zero BC_TERMS = { 'bc1': (sample_bc1, 0, None), # left wall, u = 0 'bc2': (sample_bc1, 1, None), # left wall, v = 0 'bc3': (sample_bc2, 0, None), # right wall, u = 0 'bc4': (sample_bc2, 1, None), # right wall, v = 0 'bc5': (sample_bc3, 0, None), # bottom wall, u = 0 'bc6': (sample_bc3, 1, None), # bottom wall, v = 0 'bc7': (sample_bc4, 0, lid_profile), # lid, u = 16 x^2 (1 - x)^2 'bc8': (sample_bc4, 1, None), # lid, v = 0 } def loss_bc(name: str) -> torch.Tensor: sampler, component, target = BC_TERMS[name] pts, _normals = sampler() goal = 0.0 if target is None else target(pts) return reduce_term(field(pts, component) - goal) # -------------------------------------------------------------------------- # Objective # -------------------------------------------------------------------------- TERM_GROUPS = {'pde1': 'pde', 'pde2': 'pde', 'pde3': 'pde', 'bc1': 'bc', 'bc2': 'bc', 'bc3': 'bc', 'bc4': 'bc', 'bc5': 'bc', 'bc6': 'bc', 'bc7': 'bc', 'bc8': 'bc'} TERM_LABELS = {'pde1': 'continuity', 'pde2': 'x-momentum', 'pde3': 'y-momentum', 'bc1': 'left wall u', 'bc2': 'left wall v', 'bc3': 'right wall u', 'bc4': 'right wall v', 'bc5': 'bottom wall u', 'bc6': 'bottom wall v', 'bc7': 'lid u', 'bc8': 'lid v'} TERM_SCALES = {name: 1.0 for name in TERM_GROUPS} GROUP_WEIGHTS = {"pde": 1.0, "bc": BC_WEIGHT, "ic": 1.0} def loss_terms(pts: torch.Tensor) -> dict[str, torch.Tensor]: """Every term of the objective, unweighted, keyed by name.""" continuity, momentum_x, momentum_y = navier_stokes(pts) terms = {'pde1': reduce_term(continuity), 'pde2': reduce_term(momentum_x), 'pde3': reduce_term(momentum_y)} for name in BC_TERMS: terms[name] = loss_bc(name) return terms class Balancer: """Fixed weights: the numbers from the Loss block, unchanged.""" def __init__(self): self.weights = dict(GROUP_WEIGHTS) def update(self, epoch: int, terms: dict[str, torch.Tensor]) -> None: return balancer = Balancer() def total_loss(pts: torch.Tensor): """The weighted objective, plus the individual terms for logging.""" terms = loss_terms(pts) total = torch.zeros((), device=DEVICE) for name, value in terms.items(): group = TERM_GROUPS.get(name, "pde") total = total + balancer.weights.get(group, 1.0) * TERM_SCALES.get(name, 1.0) * value return total, terms # -------------------------------------------------------------------------- # Snapshots of the field during training (for the GIFs and the error-versus-epoch plot) # -------------------------------------------------------------------------- _SNAPS = {"epoch": [], "uvp": [], "grid": None} def _snapshot(epoch: int, adam_epochs: int = 0) -> None: if _SNAPS["grid"] is None: g = np.linspace(0.0, 1.0, SNAP_N, dtype=np.float32) gx, gy = np.meshgrid(g, g, indexing="ij") _SNAPS["grid"] = (g, np.stack([gx.reshape(-1), gy.reshape(-1)], axis=1)) g, flat = _SNAPS["grid"] with torch.no_grad(): out = net(torch.tensor(flat, device=DEVICE)).cpu().numpy().reshape(SNAP_N, SNAP_N, 3) _SNAPS["epoch"].append(epoch) _SNAPS["uvp"].append(out.astype(np.float32)) np.savez_compressed(str(OUTPUT_DIR / "snapshots.npz"), epoch=np.array(_SNAPS["epoch"]), uvp=np.stack(_SNAPS["uvp"]), x=g, y=g, adam_epochs=np.array(adam_epochs)) # -------------------------------------------------------------------------- # Training # -------------------------------------------------------------------------- def build_optimiser(lr: float = LEARNING_RATE): return torch.optim.Adam(net.parameters(), lr=lr, weight_decay=0.0) def train(epochs: int = EPOCHS, lbfgs_rounds: int = LBFGS_ROUNDS, log_every: int = LOG_EVERY, verbose: bool = True, snap_every: int = SNAP_EVERY, snap_every_lbfgs: int = SNAP_EVERY_LBFGS) -> dict: """Adam for `epochs` epochs, then L-BFGS for `lbfgs_rounds` rounds. Returns the loss history.""" torch.manual_seed(SEED) pts = sample_interior() history = {"epoch": [], "total": [], "time": [], "stage": []} for name in TERM_GROUPS: history[name] = [] history["adam_epochs"] = int(epochs) _SNAPS["epoch"].clear() _SNAPS["uvp"].clear() started = time.time() def record(epoch: int, total: torch.Tensor, terms: dict, stage: int) -> float: value = float(total.detach()) history["epoch"].append(epoch) history["total"].append(value) history["time"].append(time.time() - started) history["stage"].append(stage) for name, term in terms.items(): history[name].append(float(term.detach())) return value def show(epoch: int, value: float, terms: dict, tag: str) -> None: parts = " ".join(f"{TERM_LABELS.get(k, k)}={float(v.detach()):.2e}" for k, v in terms.items()) print(f"{tag} {epoch:>7d} total={value:.4e} {parts}") _snapshot(0, epochs) # ---- stage 1: Adam ---------------------------------------------------- optimiser = build_optimiser() for epoch in range(1, epochs + 1): if RESAMPLE_EVERY and epoch % RESAMPLE_EVERY == 0: pts = sample_interior() optimiser.zero_grad(set_to_none=True) total, terms_now = total_loss(pts) if total.requires_grad: total.backward() optimiser.step() balancer.update(epoch, terms_now) value = record(epoch, total, terms_now, 0) if not math.isfinite(value): print(f"Loss is {value} at epoch {epoch}: stopping. Try a smaller LEARNING_RATE.") history["diverged"] = True return history if verbose and (epoch % log_every == 0 or epoch == 1): show(epoch, value, terms_now, "adam ") if snap_every and epoch % snap_every == 0: _snapshot(epoch, epochs) # ---- stage 2: L-BFGS, full batch, boundary points frozen ---------------------- if lbfgs_rounds > 0: _BC_FROZEN["store"].clear() _BC_FROZEN["on"] = True lbfgs = torch.optim.LBFGS(net.parameters(), lr=1.0, max_iter=LBFGS_INNER, history_size=50, tolerance_grad=1e-12, tolerance_change=1e-14, line_search_fn="strong_wolfe") def closure(): lbfgs.zero_grad(set_to_none=True) loss, _ = total_loss(pts) loss.backward() return loss for r in range(1, lbfgs_rounds + 1): lbfgs.step(closure) epoch = epochs + r total, terms_now = total_loss(pts) # loss at the accepted point (no backward) value = record(epoch, total, terms_now, 1) if not math.isfinite(value): print(f"Loss is {value} in L-BFGS round {r}: stopping.") history["diverged"] = True break if verbose and (r % max(log_every // 20, 1) == 0 or r == 1): show(epoch, value, terms_now, "lbfgs") if snap_every_lbfgs and (r % snap_every_lbfgs == 0 or r == lbfgs_rounds): _snapshot(epoch, epochs) _BC_FROZEN["on"] = False _BC_FROZEN["store"].clear() if verbose: print(f"finished in {time.time() - started:.1f}s") return history def save_checkpoint(history: dict, path: str = CHECKPOINT) -> str: """Weights plus enough metadata to reload without the diagram; also weights.npz (numpy only).""" payload = { "model": net.state_dict(), "coords": COORDS, "title": TITLE, "generator": "Scientific ML Studio Lab 2.0, modified by lid_driven_cavity_final.py", "history": history, } torch.save(payload, path) Path(path).with_suffix(".history.json").write_text(json.dumps(history)) np.savez(str(OUTPUT_DIR / "weights.npz"), **{k: v.detach().cpu().numpy() for k, v in net.state_dict().items()}) return path def load_checkpoint(path: str = CHECKPOINT) -> dict: try: payload = torch.load(path, map_location=DEVICE, weights_only=False) except TypeError: # very old PyTorch has no weights_only payload = torch.load(path, map_location=DEVICE) net.load_state_dict(payload["model"]) return payload.get("history", {}) def evaluate_network(n: int = POST_N, chunk: int = 8192) -> dict: """The trained network on an n x n grid (x-major: array[i, j] is at x_i, y_j). Returns numpy arrays: u, v, p, omega (= v_x - u_y) and the three Navier-Stokes residuals. """ g = np.linspace(0.0, 1.0, n) gx, gy = np.meshgrid(g, g, indexing="ij") flat = np.stack([gx.ravel(), gy.ravel()], axis=1) keys = ["u", "v", "p", "omega", "r_continuity", "r_momentum_x", "r_momentum_y"] parts = {k: [] for k in keys} for i in range(0, flat.shape[0], chunk): pts = torch.tensor(flat[i:i + chunk], dtype=torch.get_default_dtype(), device=DEVICE, requires_grad=True) cont, mom_x, mom_y, aux = navier_stokes(pts, full=True) values = {"u": aux["u"], "v": aux["v"], "p": aux["p"], "omega": aux["v_x"] - aux["u_y"], "r_continuity": cont, "r_momentum_x": mom_x, "r_momentum_y": mom_y} for k in keys: parts[k].append(values[k].detach().cpu().numpy().reshape(-1)) out = {k: np.concatenate(parts[k]).astype(np.float64).reshape(n, n) for k in keys} out["x"] = g out["y"] = g return out def write_report(history: dict | None = None, path: str | None = None) -> str: """Tabulate the fields and the three residuals over a fresh set of interior points.""" path = path or str(OUTPUT_DIR / 'results.csv') pts = sample_interior() cont, mom_x, mom_y, aux = navier_stokes(pts, full=True) columns = ['x', 'y', 'u', 'v', 'p', 'r_continuity', 'r_momentum_x', 'r_momentum_y'] data = [pts[:, 0], pts[:, 1], aux["u"][:, 0], aux["v"][:, 0], aux["p"][:, 0], cont[:, 0], mom_x[:, 0], mom_y[:, 0]] data = [d.detach().cpu().numpy() for d in data] lines = [",".join(columns)] lines += [",".join(f"{col[i]:.6g}" for col in data) for i in range(pts.shape[0])] lines += ["", "# summary", f"# title,{TITLE}", f"# rows,{pts.shape[0]}"] for key in ("total", *TERM_GROUPS): series = (history or {}).get(key) if series: lines.append(f"# final_{key},{series[-1]:.6e}") Path(path).write_text("\n".join(lines) + "\n") return path # ===== BEGIN FEM REFERENCE ==================================================== # Finite element reference for the steady lid-driven cavity: Taylor-Hood elements # (Q2 biquadratic velocity, Q1 bilinear pressure) on a uniform N x N mesh of squares, # Newton iteration, pressure fixed at one node. Needs scipy. Checked against the Kovasznay # flow (exact solution: velocity L2 error converges at order 3, pressure at order 2). import scipy.sparse as sp import scipy.sparse.linalg as spla _G = np.array([-np.sqrt(3 / 5), 0.0, np.sqrt(3 / 5)]) _W = np.array([5 / 9, 8 / 9, 5 / 9]) def _q2(xi): xi = np.asarray(xi, float) f = np.stack([0.5 * xi * (xi - 1), 1 - xi ** 2, 0.5 * xi * (xi + 1)], axis=-1) d = np.stack([xi - 0.5, -2 * xi, xi + 0.5], axis=-1) return f, d def _q1(xi): xi = np.asarray(xi, float) f = np.stack([0.5 * (1 - xi), 0.5 * (1 + xi)], axis=-1) d = np.stack([-0.5 * np.ones_like(xi), 0.5 * np.ones_like(xi)], axis=-1) return f, d def _fem_lid_profile(x, kind="regularised"): x = np.asarray(x, float) return 16 * x ** 2 * (1 - x) ** 2 if kind == "regularised" else np.ones_like(x) class Cavity: def __init__(self, n_el=32, re=100.0, lid="regularised"): self.N, self.Re, self.lid = n_el, float(re), lid N = n_el self.nv = 2 * N + 1 # velocity nodes per side self.npn = N + 1 # pressure nodes per side self.Nv, self.Np = self.nv ** 2, self.npn ** 2 h = 1.0 / N self.h = h # reference-element tables at the 3 x 3 Gauss points (q = i * 3 + j, i along x, j along y) f2, d2 = _q2(_G); f1, d1 = _q1(_G) nq = 9 self.Nq = np.zeros((nq, 9)); self.dNx = np.zeros((nq, 9)); self.dNy = np.zeros((nq, 9)) self.Pq = np.zeros((nq, 4)); self.dPx = np.zeros((nq, 4)); self.dPy = np.zeros((nq, 4)) self.Wq = np.zeros(nq) for i in range(3): for j in range(3): q = i * 3 + j self.Wq[q] = _W[i] * _W[j] * (h / 2) ** 2 for a in range(3): for b in range(3): l = a * 3 + b self.Nq[q, l] = f2[i, a] * f2[j, b] self.dNx[q, l] = d2[i, a] * f2[j, b] * (2 / h) self.dNy[q, l] = f2[i, a] * d2[j, b] * (2 / h) for c in range(2): for d in range(2): l = c * 2 + d self.Pq[q, l] = f1[i, c] * f1[j, d] # element connectivity ex, ey = np.meshgrid(np.arange(N), np.arange(N), indexing="ij") ex, ey = ex.ravel(), ey.ravel() self.Ev = np.stack([(2 * ex + a) * self.nv + (2 * ey + b) for a in range(3) for b in range(3)], axis=1) self.Ep = np.stack([(ex + c) * self.npn + (ey + d) for c in range(2) for d in range(2)], axis=1) # boundary data ix, iy = np.meshgrid(np.arange(self.nv), np.arange(self.nv), indexing="ij") self.xn, self.yn = ix.ravel() / (2 * N), iy.ravel() / (2 * N) self.boundary = (ix.ravel() == 0) | (ix.ravel() == 2 * N) | (iy.ravel() == 0) | (iy.ravel() == 2 * N) self.u_bc = np.zeros(self.Nv) top = iy.ravel() == 2 * N self.u_bc[top] = _fem_lid_profile(self.xn[top], lid) self.v_bc = np.zeros(self.Nv) free_v = np.where(~self.boundary)[0] free_p = np.arange(1, self.Np) # pressure pinned at node 0 self.free = np.concatenate([free_v, self.Nv + free_v, 2 * self.Nv + free_p]) self.sol = None # -------------------------------------------------------------- element assembly def _fields(self, x): u, v, p = x[:self.Nv], x[self.Nv:2 * self.Nv], x[2 * self.Nv:] ue, ve, pe = u[self.Ev], v[self.Ev], p[self.Ep] return dict(u=ue @ self.Nq.T, ux=ue @ self.dNx.T, uy=ue @ self.dNy.T, v=ve @ self.Nq.T, vx=ve @ self.dNx.T, vy=ve @ self.dNy.T, p=pe @ self.Pq.T) def residual(self, x): f = self._fields(x); nu = 1.0 / self.Re; W = self.Wq Nq, dNx, dNy = self.Nq, self.dNx, self.dNy conv_u = f["u"] * f["ux"] + f["v"] * f["uy"] conv_v = f["u"] * f["vx"] + f["v"] * f["vy"] Ru = np.einsum("q,eq,qi->ei", W, conv_u, Nq) + nu * (np.einsum("q,eq,qi->ei", W, f["ux"], dNx) + np.einsum("q,eq,qi->ei", W, f["uy"], dNy)) - np.einsum("q,eq,qi->ei", W, f["p"], dNx) Rv = np.einsum("q,eq,qi->ei", W, conv_v, Nq) + nu * (np.einsum("q,eq,qi->ei", W, f["vx"], dNx) + np.einsum("q,eq,qi->ei", W, f["vy"], dNy)) - np.einsum("q,eq,qi->ei", W, f["p"], dNy) div = f["ux"] + f["vy"] Rp = np.einsum("q,eq,qi->ei", W, div, self.Pq) R = np.zeros(2 * self.Nv + self.Np) R[:self.Nv] = np.bincount(self.Ev.ravel(), Ru.ravel(), self.Nv) R[self.Nv:2 * self.Nv] = np.bincount(self.Ev.ravel(), Rv.ravel(), self.Nv) R[2 * self.Nv:] = np.bincount(self.Ep.ravel(), Rp.ravel(), self.Np) return R def jacobian(self, x): f = self._fields(x); nu = 1.0 / self.Re; W = self.Wq Nq, dNx, dNy, Pq = self.Nq, self.dNx, self.dNy, self.Pq E = lambda c, A, B: np.einsum("q,eq,qi,qj->eij", W, c, A, B) diff = np.einsum("q,qi,qj->ij", W, dNx, dNx) + np.einsum("q,qi,qj->ij", W, dNy, dNy) conv = E(f["u"], Nq, dNx) + E(f["v"], Nq, dNy) one = np.ones_like(f["u"]) Juu = conv + E(f["ux"], Nq, Nq) + nu * diff[None] Jvv = conv + E(f["vy"], Nq, Nq) + nu * diff[None] Juv = E(f["uy"], Nq, Nq) Jvu = E(f["vx"], Nq, Nq) Jup = -E(one, dNx, Pq); Jvp = -E(one, dNy, Pq) Jpu = E(one, Pq, dNx); Jpv = E(one, Pq, dNy) def coo(Ke, rows, cols, shape): r = np.broadcast_to(rows[:, :, None], Ke.shape).ravel() c = np.broadcast_to(cols[:, None, :], Ke.shape).ravel() return sp.coo_matrix((Ke.ravel(), (r, c)), shape=shape).tocsr() nv, npn = self.Nv, self.Np; Ev, Ep = self.Ev, self.Ep blocks = [[coo(Juu, Ev, Ev, (nv, nv)), coo(Juv, Ev, Ev, (nv, nv)), coo(Jup, Ev, Ep, (nv, npn))], [coo(Jvu, Ev, Ev, (nv, nv)), coo(Jvv, Ev, Ev, (nv, nv)), coo(Jvp, Ev, Ep, (nv, npn))], [coo(Jpu, Ep, Ev, (npn, nv)), coo(Jpv, Ep, Ev, (npn, nv)), None]] return sp.bmat(blocks, format="csr") # ------------------------------------------------------------------------ solve def solve(self, tol=1e-10, max_iter=30, verbose=False, x0=None): x = np.zeros(2 * self.Nv + self.Np) if x0 is None else x0.copy() x[:self.Nv][self.boundary] = self.u_bc[self.boundary] x[self.Nv:2 * self.Nv][self.boundary] = self.v_bc[self.boundary] history = [] for it in range(max_iter): R = self.residual(x) rn = np.abs(R[self.free]).max(); history.append(rn) if verbose: print(f" Newton {it}: max residual {rn:.2e}") if rn < tol: break J = self.jacobian(x)[self.free][:, self.free].tocsc() dx = spla.splu(J).solve(-R[self.free]) x[self.free] += dx self.x, self.newton_history = x, history self.u, self.v, self.p = x[:self.Nv], x[self.Nv:2 * self.Nv], x[2 * self.Nv:] return self # --------------------------------------------------------------------- evaluate def eval(self, px, py, what=("u", "v", "p")): """Values (and u_y, v_x for vorticity) at arbitrary points.""" px = np.clip(np.asarray(px, float).ravel(), 0, 1); py = np.clip(np.asarray(py, float).ravel(), 0, 1) N = self.N ex = np.minimum((px * N).astype(int), N - 1); ey = np.minimum((py * N).astype(int), N - 1) xi = 2 * (px * N - ex) - 1; eta = 2 * (py * N - ey) - 1 fx, dfx = _q2(xi); fy, dfy = _q2(eta) gx, gy = _q1(xi)[0], _q1(eta)[0] out = {} def q2sum(nodal, fa, fb): s = np.zeros(px.shape) for a in range(3): for b in range(3): s += nodal[(2 * ex + a) * self.nv + (2 * ey + b)] * fa[:, a] * fb[:, b] return s if "u" in what: out["u"] = q2sum(self.u, fx, fy) if "v" in what: out["v"] = q2sum(self.v, fx, fy) if "p" in what: s = np.zeros(px.shape) for c in range(2): for d in range(2): s += self.p[(ex + c) * self.npn + (ey + d)] * gx[:, c] * gy[:, d] out["p"] = s if "omega" in what: vx = q2sum(self.v, dfx, fy) * 2 * N uy = q2sum(self.u, fx, dfy) * 2 * N out["omega"] = vx - uy return out def solve_reference(n_el: int = 128, grid_n: int = 201, verbose: bool = True) -> dict: """Solve the cavity on an n_el x n_el mesh and sample u, v, p, omega on a grid_n x grid_n grid.""" t0 = time.time() fem = Cavity(n_el=n_el, re=RE, lid="regularised").solve(verbose=False) g = np.linspace(0.0, 1.0, grid_n) gx, gy = np.meshgrid(g, g, indexing="ij") out = fem.eval(gx, gy, ("u", "v", "p", "omega")) ref = {k: out[k].reshape(grid_n, grid_n) for k in ("u", "v", "p", "omega")} ref["x"], ref["y"] = g, g ref["n_el"] = np.array(n_el) ref["newton_residual"] = np.array(fem.newton_history[-1]) if verbose: print(f"FEM reference: {n_el} x {n_el} Q2-Q1 elements, {len(fem.newton_history) - 1} Newton steps, " f"final residual {fem.newton_history[-1]:.1e}, {time.time() - t0:.0f} s") return ref def load_or_solve_reference(n_el: int = 128, grid_n: int = 201, path: str | None = None, verbose: bool = True) -> dict: """Use a saved reference file when one exists (cavity_reference_N.npz), otherwise solve and save.""" path = Path(path) if path else OUTPUT_DIR / f"cavity_reference_N{n_el}.npz" if path.exists(): data = np.load(path) ref = {k: data[k] for k in data.files} if ref["u"].shape == (grid_n, grid_n) and all(k in ref for k in ("u", "v", "p", "omega", "x", "y")): ref.setdefault("n_el", np.array(n_el)) if verbose: print(f"FEM reference loaded from {path}") return ref print(f"{path} does not match the {grid_n} x {grid_n} grid; solving again.") ref = solve_reference(n_el, grid_n, verbose) np.savez(path, **ref) return ref # ===== END FEM REFERENCE ====================================================== # ===== BEGIN POSTPROCESS (numerics and metrics) =============================== # Everything from here to END POSTPROCESS uses numpy and matplotlib only. # Arrays are x-major: A[i, j] is the value at (x_i, y_j). _trapz = getattr(np, "trapezoid", None) or np.trapz FIELD_LABELS = {"u": "u", "v": "v", "p": "p (mean removed)", "speed": "speed |u|", "omega": "vorticity", "psi": "stream function"} def domain_mean(f: np.ndarray, g: np.ndarray) -> float: """Mean of f over the square (trapezoid rule).""" area = (g[-1] - g[0]) ** 2 return float(_trapz(_trapz(f, g, axis=1), g, axis=0) / area) def stream_function(u: np.ndarray, g: np.ndarray) -> np.ndarray: """psi(x, y) = integral of u dy from the bottom wall, so that u = psi_y and v = -psi_x.""" inc = 0.5 * (u[:, 1:] + u[:, :-1]) * np.diff(g)[None, :] return np.concatenate([np.zeros((u.shape[0], 1)), np.cumsum(inc, axis=1)], axis=1) def _extremum_1d(line: np.ndarray, g: np.ndarray, kind: str): """Location and value of the minimum or maximum of a sampled line (parabola through 3 points).""" k = int(np.argmin(line) if kind == "min" else np.argmax(line)) if 0 < k < len(line) - 1: fm, f0, fp = line[k - 1], line[k], line[k + 1] d = fm - 2.0 * f0 + fp if d != 0.0: h = g[1] - g[0] return float(g[k] + 0.5 * h * (fm - fp) / d), float(f0 - (fp - fm) ** 2 / (8.0 * d)) return float(g[k]), float(line[k]) def _min_2d(f: np.ndarray, g: np.ndarray): i, j = np.unravel_index(int(np.argmin(f)), f.shape) xc, v1 = _extremum_1d(f[:, j], g, "min") yc, v2 = _extremum_1d(f[i, :], g, "min") return xc, yc, float(f[i, j] + (v1 - f[i, j]) + (v2 - f[i, j])) def prepare_fields(raw: dict, g: np.ndarray) -> dict: """Add the derived fields: pressure with its mean removed, speed and stream function.""" out = {k: np.asarray(v) for k, v in raw.items()} out["p"] = raw["p"] - domain_mean(raw["p"], g) out["speed"] = np.hypot(raw["u"], raw["v"]) out["psi"] = stream_function(raw["u"], g) return out def flow_features(F: dict, g: np.ndarray) -> dict: """Numbers usually quoted for the cavity: primary vortex and centreline extremes.""" mid = len(g) // 2 # index of 0.5 (the grid size is odd) xc, yc, psi_min = _min_2d(F["psi"], g) y_umin, u_min = _extremum_1d(F["u"][mid, :], g, "min") x_vmax, v_max = _extremum_1d(F["v"][:, mid], g, "max") x_vmin, v_min = _extremum_1d(F["v"][:, mid], g, "min") return {"vortex_x": xc, "vortex_y": yc, "psi_min": psi_min, "u_min_on_x0.5": u_min, "u_min_y": y_umin, "v_max_on_y0.5": v_max, "v_max_x": x_vmax, "v_min_on_y0.5": v_min, "v_min_x": x_vmin} def error_stats(a: np.ndarray, b: np.ndarray, g: np.ndarray) -> dict: """Error of a against the reference b.""" d = a - b k = np.unravel_index(int(np.argmax(np.abs(d))), d.shape) return {"rel_l2": float(np.linalg.norm(d) / np.linalg.norm(b)), "rmse": float(np.sqrt(np.mean(d ** 2))), "mae": float(np.mean(np.abs(d))), "max_abs": float(np.abs(d).max()), "max_at_x": float(g[k[0]]), "max_at_y": float(g[k[1]]), "bias": float(d.mean()), "ref_rms": float(np.sqrt(np.mean(b ** 2))), "ref_max_abs": float(np.abs(b).max())} def _residual_stats(r: np.ndarray) -> dict: a = np.abs(r) return {"rms": float(np.sqrt(np.mean(r ** 2))), "median_abs": float(np.median(a)), "p99_abs": float(np.percentile(a, 99)), "max_abs": float(a.max())} def compute_metrics(P: dict, R: dict, g: np.ndarray, history: dict | None = None, setup: dict | None = None) -> dict: """All the numbers reported for the lesson. P = PINN fields, R = reference fields.""" m: dict = {"setup": setup or {}} m["fields"] = {k: error_stats(P[k], R[k], g) for k in ("u", "v", "p", "speed", "omega", "psi")} dv = np.sqrt((P["u"] - R["u"]) ** 2 + (P["v"] - R["v"]) ** 2) m["velocity_vector"] = { "rel_l2": float(np.sqrt(np.sum(dv ** 2) / np.sum(R["u"] ** 2 + R["v"] ** 2))), "rmse": float(np.sqrt(np.mean(dv ** 2))), "max_abs": float(dv.max())} inner = (g >= 0.05) & (g <= 0.95) mask = inner[:, None] & inner[None, :] m["interior_rel_l2"] = {k: float(np.linalg.norm((P[k] - R[k])[mask]) / np.linalg.norm(R[k][mask])) for k in ("u", "v", "p")} edges = np.linspace(0.0, 1.0, 6) bands = [] for lo, hi in zip(edges[:-1], edges[1:]): cols = (g >= lo) & ((g < hi) if hi < 1.0 else (g <= hi)) rmse = float(np.sqrt(np.mean(dv[:, cols] ** 2))) ref_rms = float(np.sqrt(np.mean(R["speed"][:, cols] ** 2))) bands.append({"y_from": float(lo), "y_to": float(hi), "rmse_velocity": rmse, "ref_speed_rms": ref_rms, "relative": rmse / ref_rms}) m["bands"] = bands fp, fr = flow_features(P, g), flow_features(R, g) m["features"] = {"pinn": fp, "reference": fr, "difference": {k: fp[k] - fr[k] for k in fp}} m["residuals"] = {"continuity": _residual_stats(P["r_continuity"]), "momentum_x": _residual_stats(P["r_momentum_x"]), "momentum_y": _residual_stats(P["r_momentum_y"])} if "r_continuity" in P else {} lid_target = 16.0 * g ** 2 * (1.0 - g) ** 2 m["boundary"] = { "walls_max_abs_u": float(max(np.abs(P["u"][0, :]).max(), np.abs(P["u"][-1, :]).max(), np.abs(P["u"][:, 0]).max())), "walls_max_abs_v": float(max(np.abs(P["v"][0, :]).max(), np.abs(P["v"][-1, :]).max(), np.abs(P["v"][:, 0]).max())), "lid_u_rmse": float(np.sqrt(np.mean((P["u"][:, -1] - lid_target) ** 2))), "lid_v_max_abs": float(np.abs(P["v"][:, -1]).max())} if history and history.get("epoch"): n_adam = int(history.get("adam_epochs", len(history["epoch"]))) m["training"] = {"epochs_recorded": len(history["epoch"]), "adam_epochs": n_adam, "lbfgs_rounds": max(len(history["epoch"]) - n_adam, 0), "seconds": float(history["time"][-1]) if history.get("time") else None, "final_total": float(history["total"][-1]), "diverged": bool(history.get("diverged", False)), "final_terms": {k: float(history[k][-1]) for k in history if k.startswith(("pde", "bc")) and history[k]}} return m def load_snapshots(path: str | Path) -> dict | None: path = Path(path) if not path.exists(): return None data = np.load(path) return {k: data[k] for k in data.files} def error_vs_epoch(snaps: dict, R: dict, g: np.ndarray) -> dict: """Relative L2 error of u, v, p and the velocity vector at each saved snapshot.""" uvp = snaps["uvp"].astype(np.float64) m = uvp.shape[1] stride = (len(g) - 1) // (m - 1) gs = g[::stride] Rs = {k: R[k][::stride, ::stride] for k in ("u", "v", "p")} Rs["p"] = Rs["p"] - domain_mean(Rs["p"], gs) out = {"epoch": [int(e) for e in snaps["epoch"]], "u": [], "v": [], "p": [], "velocity": []} for s in range(uvp.shape[0]): u, v, p = uvp[s, :, :, 0], uvp[s, :, :, 1], uvp[s, :, :, 2] p = p - domain_mean(p, gs) out["u"].append(float(np.linalg.norm(u - Rs["u"]) / np.linalg.norm(Rs["u"]))) out["v"].append(float(np.linalg.norm(v - Rs["v"]) / np.linalg.norm(Rs["v"]))) out["p"].append(float(np.linalg.norm(p - Rs["p"]) / np.linalg.norm(Rs["p"]))) out["velocity"].append(float(np.sqrt((np.sum((u - Rs["u"]) ** 2) + np.sum((v - Rs["v"]) ** 2)) / (np.sum(Rs["u"] ** 2) + np.sum(Rs["v"] ** 2))))) return out def _flatten(d: dict, prefix: str = ""): for k, v in d.items(): key = f"{prefix}{k}" if isinstance(v, dict): yield from _flatten(v, key + ".") elif isinstance(v, list): for i, item in enumerate(v): if isinstance(item, dict): yield from _flatten(item, f"{key}[{i}].") else: yield f"{key}[{i}]", item else: yield key, v def metrics_markdown(m: dict) -> str: """The tables for the lesson, as markdown.""" s = m.get("setup", {}) lines = ["## Setup", ""] for k, v in s.items(): lines.append(f"- {k}: {v}") lines += ["", "## Error against the finite element reference", "", "| field | relative L2 | RMS error | mean abs error | largest error | where (x, y) |", "|---|---|---|---|---|---|"] for k, e in m["fields"].items(): lines.append(f"| {FIELD_LABELS[k]} | {e['rel_l2']:.3e} | {e['rmse']:.3e} | {e['mae']:.3e} | " f"{e['max_abs']:.3e} | ({e['max_at_x']:.3f}, {e['max_at_y']:.3f}) |") vv = m["velocity_vector"] lines += ["", f"Velocity vector (u, v) together: relative L2 = {vv['rel_l2']:.3e}, " f"RMS = {vv['rmse']:.3e}, largest = {vv['max_abs']:.3e}.", "", "Relative L2 away from the walls (0.05 <= x, y <= 0.95): " + ", ".join(f"{k} = {v:.3e}" for k, v in m["interior_rel_l2"].items()) + ".", "", "## Velocity error by height", "", "| y range | RMS velocity error | RMS reference speed | ratio |", "|---|---|---|---|"] for b in m["bands"]: lines.append(f"| {b['y_from']:.1f} to {b['y_to']:.1f} | {b['rmse_velocity']:.3e} | " f"{b['ref_speed_rms']:.3e} | {b['relative']:.3e} |") f = m["features"] lines += ["", "## Flow features", "", "| quantity | PINN | FEM | difference |", "|---|---|---|---|"] for k in f["pinn"]: lines.append(f"| {k} | {f['pinn'][k]:.5f} | {f['reference'][k]:.5f} | {f['difference'][k]:+.2e} |") if m.get("residuals"): lines += ["", "## Equation residuals of the trained network (on the post-processing grid)", "", "| residual | RMS | median abs | 99th percentile | largest |", "|---|---|---|---|---|"] for k, r in m["residuals"].items(): lines.append(f"| {k} | {r['rms']:.3e} | {r['median_abs']:.3e} | {r['p99_abs']:.3e} | {r['max_abs']:.3e} |") b = m["boundary"] lines += ["", "## Boundary conditions", "", f"- largest |u| on the left, right and bottom walls: {b['walls_max_abs_u']:.3e}", f"- largest |v| on the left, right and bottom walls: {b['walls_max_abs_v']:.3e}", f"- RMS difference between u on the lid and 16 x^2 (1 - x)^2: {b['lid_u_rmse']:.3e}", f"- largest |v| on the lid: {b['lid_v_max_abs']:.3e}"] t = m.get("training") if t: lines += ["", "## Training", "", f"- Adam epochs: {t['adam_epochs']}, L-BFGS rounds: {t['lbfgs_rounds']}" + (f", time: {t['seconds']:.0f} s" if t.get("seconds") else ""), f"- final total loss: {t['final_total']:.3e}"] for k, v in t["final_terms"].items(): lines.append(f" - {TERM_LABELS.get(k, k)}: {v:.3e}") lines += ["", "These numbers come from one training run (one seed).", ""] return "\n".join(lines) def write_metrics(m: dict, evo: dict | None, outdir: Path) -> None: outdir.mkdir(parents=True, exist_ok=True) (outdir / "metrics.json").write_text(json.dumps(m, indent=2, default=float)) with open(outdir / "metrics.csv", "w", newline="") as fh: w = csv.writer(fh) w.writerow(["metric", "value"]) for key, value in _flatten(m): w.writerow([key, value]) (outdir / "metrics_table.md").write_text(metrics_markdown(m)) if evo: with open(outdir / "error_vs_epoch.csv", "w", newline="") as fh: w = csv.writer(fh) w.writerow(["epoch", "rel_l2_u", "rel_l2_v", "rel_l2_p", "rel_l2_velocity"]) for row in zip(evo["epoch"], evo["u"], evo["v"], evo["p"], evo["velocity"]): w.writerow(row) # ---------------------------------------------------------------------------- # Figures # ---------------------------------------------------------------------------- PINN_COLOUR = "#1f77b4" FEM_COLOUR = "#222222" FIELD_CMAPS = {"u": ("RdBu_r", True), "v": ("RdBu_r", True), "p": ("RdBu_r", True), "speed": ("viridis", False), "omega": ("RdBu_r", True), "psi": ("viridis", False)} def _new_fig(w, h, nrows=1, ncols=1, **kw): fig = Figure(figsize=(w, h), layout="constrained") FigureCanvasAgg(fig) return fig, fig.subplots(nrows, ncols, **kw) def _save(fig, name): FIGURE_DIR.mkdir(parents=True, exist_ok=True) fig.savefig(FIGURE_DIR / f"{name}.png", dpi=FIGURE_DPI) def _frame(fig): fig.canvas.draw() return np.asarray(fig.canvas.buffer_rgba())[:, :, :3].copy() def _square(ax, xlabel="x", ylabel="y"): ax.set_aspect("equal") ax.set_xlabel(xlabel) ax.set_ylabel(ylabel) def _limits(a, b, sym, pct=100.0): both = np.concatenate([np.ravel(a), np.ravel(b)]) if sym: v = float(np.percentile(np.abs(both), pct)) v = v if v > 0 else 1e-12 return -v, v lo, hi = float(both.min()), float(both.max()) return lo, (hi if hi > lo else lo + 1e-12) def fig_triplet(key, P, R, g): """PINN | FEM reference | absolute error, for one field.""" cmap, sym = FIELD_CMAPS[key] X, Y = np.meshgrid(g, g, indexing="ij") a, b = P[key], R[key] lo, hi = _limits(a, b, sym, 99.0 if key == "omega" else 100.0) levels = np.linspace(lo, hi, 41) err = np.abs(a - b) emax = max(float(err.max()), 1e-15) st = error_stats(a, b, g) fig, axs = _new_fig(13.5, 4.3, 1, 3) for ax, z, title in zip(axs[:2], (a, b), ("PINN", "FEM reference")): cs = ax.contourf(X, Y, z, levels=levels, cmap=cmap, extend="both") ax.set_title(title) _square(ax) fig.colorbar(cs, ax=list(axs[:2]), shrink=0.8, label=FIELD_LABELS[key]) ce = axs[2].contourf(X, Y, err, levels=np.linspace(0, emax, 41), cmap="magma") fig.colorbar(ce, ax=axs[2], shrink=0.8, label=f"|error| in {FIELD_LABELS[key]}") axs[2].set_title(f"absolute error (relative L2 {st['rel_l2']:.2e})") _square(axs[2]) _save(fig, f"cavity_{key}") def fig_error_maps(P, R, g): """All the error maps on one sheet.""" X, Y = np.meshgrid(g, g, indexing="ij") fig, axs = _new_fig(13, 8, 2, 3) for ax, key in zip(axs.ravel(), ("u", "v", "p", "speed", "omega", "psi")): err = np.abs(P[key] - R[key]) top = max(float(np.percentile(err, 99.5)), 1e-15) cs = ax.contourf(X, Y, err, levels=np.linspace(0, top, 31), cmap="magma", extend="max") fig.colorbar(cs, ax=ax, shrink=0.8) ax.set_title(f"{FIELD_LABELS[key]}: rel. L2 {error_stats(P[key], R[key], g)['rel_l2']:.2e}", fontsize=10) _square(ax) fig.suptitle("Absolute error against the FEM reference (colour scale ends at the 99.5th percentile)") _save(fig, "cavity_error_maps") def fig_streamlines(P, R, g, feats): fig, axs = _new_fig(11.5, 5.2, 1, 2) norm = Normalize(0.0, 1.0) for ax, F, title, f in zip(axs, (P, R), ("PINN", "FEM reference"), (feats["pinn"], feats["reference"])): sp = ax.streamplot(g, g, F["u"].T, F["v"].T, color=F["speed"].T, cmap="viridis", norm=norm, density=1.7, linewidth=0.9, arrowsize=0.8) ax.plot([f["vortex_x"]], [f["vortex_y"]], "r+", ms=10, mew=2) ax.set_title(f"{title}: vortex centre ({f['vortex_x']:.3f}, {f['vortex_y']:.3f})") ax.set_xlim(0, 1) ax.set_ylim(0, 1) _square(ax) fig.colorbar(sp.lines, ax=list(axs), shrink=0.8, label="speed") _save(fig, "cavity_streamlines") def fig_psi_lines(P, R, g, feats): """Stream-function contour lines, PINN solid and FEM dashed, on one plot.""" X, Y = np.meshgrid(g, g, indexing="ij") lo = float(R["psi"].min()) levels = np.sort(np.concatenate([np.linspace(lo, 0, 14)[:-1], [-1e-3 * abs(lo), 1e-4 * abs(lo)]])) fig, ax = _new_fig(6.4, 6.0) ax.contour(X, Y, R["psi"], levels=levels, colors=FEM_COLOUR, linestyles="--", linewidths=1.3) ax.contour(X, Y, P["psi"], levels=levels, colors=PINN_COLOUR, linestyles="-", linewidths=1.0) ax.plot([feats["reference"]["vortex_x"]], [feats["reference"]["vortex_y"]], "k+", ms=10, mew=2) ax.plot([feats["pinn"]["vortex_x"]], [feats["pinn"]["vortex_y"]], "x", color=PINN_COLOUR, ms=9, mew=2) ax.plot([], [], color=PINN_COLOUR, label="PINN") ax.plot([], [], color=FEM_COLOUR, ls="--", label="FEM reference") ax.legend(loc="upper left", fontsize=9) ax.set_title("Stream function: same contour levels for both") _square(ax) _save(fig, "cavity_psi_lines") def fig_profiles(P, R, g): """u along the vertical centreline and v along the horizontal centreline.""" mid = len(g) // 2 fig, axs = _new_fig(11, 6.8, 2, 2, gridspec_kw={"height_ratios": [3, 1]}) lines = [(P["u"][mid, :], R["u"][mid, :], "u on x = 0.5", "y"), (P["v"][:, mid], R["v"][:, mid], "v on y = 0.5", "x")] for col, (pv, rv, xlabel, along) in enumerate(lines): top, bottom = axs[0, col], axs[1, col] if col == 0: top.plot(rv, g, "--", color=FEM_COLOUR, label="FEM reference") top.plot(pv, g, color=PINN_COLOUR, label="PINN") top.set_xlabel(xlabel) top.set_ylabel("y") else: top.plot(g, rv, "--", color=FEM_COLOUR, label="FEM reference") top.plot(g, pv, color=PINN_COLOUR, label="PINN") top.set_xlabel("x") top.set_ylabel(xlabel) top.grid(alpha=0.25) top.legend(fontsize=9) bottom.plot(g, pv - rv, color=PINN_COLOUR) bottom.axhline(0, color="0.6", lw=0.8) bottom.set_xlabel(along) bottom.set_ylabel("PINN - FEM") bottom.grid(alpha=0.25) _save(fig, "cavity_profiles") def fig_lid(P, R, g): """Quantities along the moving lid.""" target = 16.0 * g ** 2 * (1.0 - g) ** 2 fig, axs = _new_fig(14, 4.0, 1, 3) axs[0].plot(g, target, "--", color=FEM_COLOUR, label="prescribed 16 x$^2$(1-x)$^2$") axs[0].plot(g, P["u"][:, -1], color=PINN_COLOUR, label="PINN") axs[0].set_ylabel("u on the lid") axs[1].plot(g, R["omega"][:, -1], "--", color=FEM_COLOUR, label="FEM reference") axs[1].plot(g, P["omega"][:, -1], color=PINN_COLOUR, label="PINN") axs[1].set_ylabel("vorticity on the lid") axs[2].plot(g, R["p"][:, -1], "--", color=FEM_COLOUR, label="FEM reference") axs[2].plot(g, P["p"][:, -1], color=PINN_COLOUR, label="PINN") axs[2].set_ylabel("pressure on the lid (mean removed)") for ax in axs: ax.set_xlabel("x") ax.grid(alpha=0.25) ax.legend(fontsize=9) _save(fig, "cavity_lid") def fig_parity(P, R, g): rng = np.random.default_rng(0) fig, axs = _new_fig(13, 4.4, 1, 3) for ax, key in zip(axs, ("u", "v", "p")): a, b = P[key].ravel(), R[key].ravel() idx = rng.choice(a.size, size=min(6000, a.size), replace=False) ax.scatter(b[idx], a[idx], s=3, alpha=0.35, color=PINN_COLOUR, rasterized=True) lim = [min(b.min(), a.min()), max(b.max(), a.max())] ax.plot(lim, lim, "k--", lw=1) ax.set_xlabel(f"{FIELD_LABELS[key]}, FEM reference") ax.set_ylabel(f"{FIELD_LABELS[key]}, PINN") ax.set_title(f"{key}: relative L2 {error_stats(P[key], R[key], g)['rel_l2']:.2e}") ax.set_aspect("equal") ax.grid(alpha=0.25) _save(fig, "cavity_parity") def fig_error_cdf(P, R, g): """Fraction of the cavity where the error is below a given size.""" fig, ax = _new_fig(6.8, 4.6) for key in ("u", "v", "p", "speed", "omega"): e = np.sort(np.abs(P[key] - R[key]).ravel()) / float(np.abs(R[key]).max()) ax.semilogx(np.maximum(e, 1e-9), np.arange(1, e.size + 1) / e.size, label=FIELD_LABELS[key]) ax.set_xlabel("|error| / largest |reference value|") ax.set_ylabel("fraction of the cavity with smaller error") ax.grid(alpha=0.25, which="both") ax.legend(fontsize=9) ax.set_title("Error distribution") _save(fig, "cavity_error_cdf") def fig_bands(m): fig, ax = _new_fig(6.4, 4.2) labels = [f"{b['y_from']:.1f}-{b['y_to']:.1f}" for b in m["bands"]] ax.bar(labels, [b["rmse_velocity"] for b in m["bands"]], color=PINN_COLOUR) ax.set_xlabel("height y") ax.set_ylabel("RMS velocity error") ax.set_title("Where the velocity error sits (the lid is at the right)") ax.grid(alpha=0.25, axis="y") _save(fig, "cavity_error_by_height") def fig_residuals(P, g): X, Y = np.meshgrid(g, g, indexing="ij") fig, axs = _new_fig(13.5, 4.3, 1, 3) names = (("r_continuity", "continuity u_x + v_y"), ("r_momentum_x", "x-momentum"), ("r_momentum_y", "y-momentum")) for ax, (key, title) in zip(axs, names): z = np.log10(np.abs(P[key]) + 1e-12) cs = ax.contourf(X, Y, z, levels=30, cmap="viridis") fig.colorbar(cs, ax=ax, shrink=0.8, label="log10 |residual|") ax.set_title(title) _square(ax) _save(fig, "cavity_residuals") def fig_loss(history): ep = np.asarray(history["epoch"]) n_adam = int(history.get("adam_epochs", len(ep))) names = [k for k in history if k.startswith(("pde", "bc")) and isinstance(history[k], list)] fig, axs = _new_fig(13, 4.6, 1, 2) axs[0].semilogy(ep, history["total"], color="k", lw=2, label="total (weighted)") for k in names: axs[0].semilogy(ep, history[k], lw=1, alpha=0.8, ls="-" if k.startswith("pde") else "--", label=TERM_LABELS.get(k, k)) pde = np.sum([history[k] for k in names if k.startswith("pde")], axis=0) bc = np.sum([history[k] for k in names if k.startswith("bc")], axis=0) axs[1].semilogy(ep, pde, label="equation terms (sum)") axs[1].semilogy(ep, bc, label="boundary terms (sum)") axs[1].semilogy(ep, history["total"], color="k", lw=2, label="total (weighted)") for ax in axs: if ep.size and ep[-1] > n_adam: ax.axvline(n_adam, color="0.5", ls=":", lw=1.2) ax.text(n_adam, ax.get_ylim()[1], " L-BFGS", va="top", fontsize=8, color="0.4") ax.set_xlabel("epoch") ax.set_ylabel("loss") ax.grid(alpha=0.25, which="both") axs[0].legend(fontsize=7, ncol=2) axs[1].legend(fontsize=9) axs[0].set_title("Every loss term") axs[1].set_title("Equations against boundaries") _save(fig, "cavity_loss") def fig_error_vs_epoch(evo, n_adam): fig, ax = _new_fig(7.2, 4.6) ep = np.asarray(evo["epoch"]) for key, label in (("u", "u"), ("v", "v"), ("p", "p (mean removed)"), ("velocity", "velocity (u, v)")): ax.semilogy(ep, evo[key], marker="o", ms=3, lw=1.4, label=label) if ep.size and ep[-1] > n_adam: ax.axvline(n_adam, color="0.5", ls=":", lw=1.2) ax.text(n_adam, ax.get_ylim()[1], " L-BFGS", va="top", fontsize=8, color="0.4") ax.set_xlabel("epoch") ax.set_ylabel("relative L2 error against FEM") ax.grid(alpha=0.25, which="both") ax.legend(fontsize=9) ax.set_title("Error during training") _save(fig, "cavity_error_vs_epoch") # ---------------------------------------------------------------------------- # GIFs # ---------------------------------------------------------------------------- def _save_gif(frames, path, frame_ms=90, hold_ms=1800): from PIL import Image adaptive = getattr(getattr(Image, "Palette", Image), "ADAPTIVE") imgs = [Image.fromarray(f).convert("P", palette=adaptive, colors=256) for f in frames] durations = [frame_ms] * len(imgs) durations[-1] = hold_ms imgs[0].save(path, save_all=True, append_images=imgs[1:], duration=durations, loop=0, disposal=2) def _pick(n, most=80): return np.unique(np.linspace(0, n - 1, min(n, most)).round().astype(int)) def _snapshot_view(snaps, R, g): uvp = snaps["uvp"].astype(np.float64) stride = (len(g) - 1) // (uvp.shape[1] - 1) gs = g[::stride] ref = {k: R[k][::stride, ::stride] for k in ("u", "v", "p", "speed")} return uvp, gs, ref def gif_evolution(snaps, R, g, path): """Speed field during training: PINN | FEM reference | absolute error.""" uvp, gs, ref = _snapshot_view(snaps, R, g) X, Y = np.meshgrid(gs, gs, indexing="ij") n_adam = int(snaps.get("adam_epochs", 10 ** 9)) lv, lve = np.linspace(0, 1.0, 31), np.linspace(0, 0.3, 31) frames = [] for s in _pick(uvp.shape[0]): u, v = uvp[s, :, :, 0], uvp[s, :, :, 1] speed = np.hypot(u, v) rel = np.sqrt((np.sum((u - ref["u"]) ** 2) + np.sum((v - ref["v"]) ** 2)) / (np.sum(ref["u"] ** 2) + np.sum(ref["v"] ** 2))) fig, axs = _new_fig(11, 3.8, 1, 3) c0 = axs[0].contourf(X, Y, np.clip(speed, 0, 1.0), levels=lv, cmap="viridis") axs[1].contourf(X, Y, ref["speed"], levels=lv, cmap="viridis") c2 = axs[2].contourf(X, Y, np.abs(speed - ref["speed"]), levels=lve, cmap="magma", extend="max") fig.colorbar(c0, ax=list(axs[:2]), shrink=0.8, ticks=[0, 0.5, 1.0], label="speed") fig.colorbar(c2, ax=axs[2], shrink=0.8, ticks=[0, 0.15, 0.3], label="|speed error|") stage = "Adam" if int(snaps["epoch"][s]) <= n_adam else "L-BFGS" axs[0].set_title(f"PINN, epoch {int(snaps['epoch'][s])} ({stage})", fontsize=10) axs[1].set_title("FEM reference", fontsize=10) axs[2].set_title(f"error, velocity rel. L2 = {rel:.3f}", fontsize=10) for ax in axs: _square(ax) ax.set_xticks([0, 0.5, 1]) ax.set_yticks([0, 0.5, 1]) frames.append(_frame(fig)) _save_gif(frames, path) def gif_streamlines(snaps, R, g, path): """Streamlines of the network during training, next to the reference.""" uvp, gs, ref = _snapshot_view(snaps, R, g) n_adam = int(snaps.get("adam_epochs", 10 ** 9)) norm = Normalize(0.0, 1.0) frames = [] for s in _pick(uvp.shape[0], 70): u, v = uvp[s, :, :, 0], uvp[s, :, :, 1] fig, axs = _new_fig(8.6, 4.4, 1, 2) for ax, uu, vv, title in ((axs[0], u, v, f"PINN, epoch {int(snaps['epoch'][s])}"), (axs[1], ref["u"], ref["v"], "FEM reference")): try: sp = ax.streamplot(gs, gs, uu.T, vv.T, color=np.hypot(uu, vv).T, cmap="viridis", norm=norm, density=1.4, linewidth=0.9, arrowsize=0.7) except Exception: # a flow that is still all zero has no streamlines sp = None ax.set_xlim(0, 1) ax.set_ylim(0, 1) ax.set_title(title, fontsize=10) _square(ax) if sp is not None: fig.colorbar(sp.lines, ax=list(axs), shrink=0.8, label="speed") frames.append(_frame(fig)) _save_gif(frames, path, frame_ms=110) def _bilinear(A, g, px, py): n, h = len(g), g[1] - g[0] fx = np.clip(px / h, 0, n - 1 - 1e-9) fy = np.clip(py / h, 0, n - 1 - 1e-9) i, j = fx.astype(int), fy.astype(int) tx, ty = fx - i, fy - j return (A[i, j] * (1 - tx) * (1 - ty) + A[i + 1, j] * tx * (1 - ty) + A[i, j + 1] * (1 - tx) * ty + A[i + 1, j + 1] * tx * ty) def _advect(F, g, p0, dt, steps): """Tracer particles moved with the steady velocity field (midpoint rule).""" def vel(p): return np.stack([_bilinear(F["u"], g, p[:, 0], p[:, 1]), _bilinear(F["v"], g, p[:, 0], p[:, 1])], axis=1) path = np.empty((steps + 1, p0.shape[0], 2)) path[0] = p0 for s in range(steps): x = path[s] path[s + 1] = np.clip(x + dt * vel(x + 0.5 * dt * vel(x)), 0.0, 1.0) return path def gif_tracers(P, R, g, path, n_particles=250, steps=160, dt=0.04, trail=12): """The same particles released in the PINN flow and in the FEM flow.""" rng = np.random.default_rng(1) p0 = rng.uniform(0.03, 0.97, size=(n_particles, 2)) paths = [_advect(F, g, p0, dt, steps) for F in (P, R)] base = np.array(colormaps["viridis"](p0[:, 1])) # colour by release height frames = [] for k in range(1, steps + 1, 2): fig, axs = _new_fig(9.4, 4.7, 1, 2) for ax, pa, title in zip(axs, paths, ("PINN", "FEM reference")): k0 = max(0, k - trail) seg = np.stack([pa[k0:k], pa[k0 + 1:k + 1]], axis=2) # (m, particles, 2 points, 2) m = seg.shape[0] colours = np.tile(base, (m, 1)) colours[:, 3] = np.repeat((np.arange(m) + 1) / m, n_particles) ax.add_collection(LineCollection(seg.reshape(-1, 2, 2), colors=colours, linewidths=1.2)) ax.scatter(pa[k, :, 0], pa[k, :, 1], s=5, color=base, zorder=3) ax.set_xlim(0, 1) ax.set_ylim(0, 1) ax.set_title(title, fontsize=10) _square(ax) sep = float(np.mean(np.linalg.norm(paths[0][k] - paths[1][k], axis=1))) fig.suptitle(f"tracer particles, t = {k * dt:.1f} mean distance between the two = {sep:.3f}", fontsize=10) frames.append(_frame(fig)) _save_gif(frames, path, frame_ms=70) # ---------------------------------------------------------------------------- # Orchestration # ---------------------------------------------------------------------------- def _safe(label, fn, *args, **kwargs) -> bool: try: fn(*args, **kwargs) print(f" done {label}") return True except Exception as exc: # one broken figure must not lose the others print(f" FAILED {label}: {type(exc).__name__}: {exc}") return False def postprocess(raw: dict, ref: dict, history: dict | None = None, snaps: dict | None = None, setup: dict | None = None, figures: bool = True, gifs: bool = True) -> dict: """Metrics, figures and GIFs from the network's fields and the FEM reference.""" g = np.asarray(raw["x"], dtype=float) assert len(g) % 2 == 1, "the post-processing grid must have an odd number of points" P, R = prepare_fields(raw, g), prepare_fields(ref, g) evo = error_vs_epoch(snaps, R, g) if snaps is not None else None metrics = compute_metrics(P, R, g, history, setup) write_metrics(metrics, evo, OUTPUT_DIR) np.savez_compressed(str(OUTPUT_DIR / "postprocess_fields.npz"), x=g, **{f"pinn_{k}": v for k, v in P.items() if np.ndim(v) == 2}, **{f"fem_{k}": v for k, v in R.items() if np.ndim(v) == 2}) print("metrics written: metrics.json, metrics.csv, metrics_table.md") if not HAVE_MPL: print("matplotlib is not installed: figures and GIFs skipped") return metrics feats = metrics["features"] if figures: print("figures:") for key in ("u", "v", "p", "speed", "omega", "psi"): _safe(f"cavity_{key}.png", fig_triplet, key, P, R, g) _safe("cavity_error_maps.png", fig_error_maps, P, R, g) _safe("cavity_streamlines.png", fig_streamlines, P, R, g, feats) _safe("cavity_psi_lines.png", fig_psi_lines, P, R, g, feats) _safe("cavity_profiles.png", fig_profiles, P, R, g) _safe("cavity_lid.png", fig_lid, P, R, g) _safe("cavity_parity.png", fig_parity, P, R, g) _safe("cavity_error_cdf.png", fig_error_cdf, P, R, g) _safe("cavity_error_by_height.png", fig_bands, metrics) if "r_continuity" in P: _safe("cavity_residuals.png", fig_residuals, P, g) if history and history.get("epoch"): _safe("cavity_loss.png", fig_loss, history) if evo: _safe("cavity_error_vs_epoch.png", fig_error_vs_epoch, evo, int(snaps.get("adam_epochs", 10 ** 9))) if gifs: FIGURE_DIR.mkdir(parents=True, exist_ok=True) print("GIFs:") if snaps is not None: _safe("cavity_evolution.gif", gif_evolution, snaps, R, g, FIGURE_DIR / "cavity_evolution.gif") _safe("cavity_streamlines_evolution.gif", gif_streamlines, snaps, R, g, FIGURE_DIR / "cavity_streamlines_evolution.gif") _safe("cavity_tracers.gif", gif_tracers, P, R, g, FIGURE_DIR / "cavity_tracers.gif") return metrics # ===== END POSTPROCESS ======================================================== # -------------------------------------------------------------------------- # Bundle and entry point # -------------------------------------------------------------------------- def make_bundle(checkpoint: str = CHECKPOINT, name: str = "cavity_results.zip") -> str: """Everything worth sending back, in one zip.""" files = [checkpoint, str(Path(checkpoint).with_suffix(".history.json")), "weights.npz", "snapshots.npz", "results.csv", "metrics.json", "metrics.csv", "metrics_table.md", "error_vs_epoch.csv", "postprocess_fields.npz"] target = OUTPUT_DIR / name with zipfile.ZipFile(target, "w", zipfile.ZIP_DEFLATED) as z: for f in files: p = Path(f) if Path(f).is_absolute() or Path(f).exists() else OUTPUT_DIR / f if p.exists(): z.write(p, p.name) for f in sorted(FIGURE_DIR.glob("*")): z.write(f, f"figures/{f.name}") try: script = Path(__file__) if script.exists(): z.write(script, script.name) except NameError: # pasted into a notebook cell: no file pass return str(target) def main(argv: list[str] | None = None) -> None: parser = argparse.ArgumentParser(description=TITLE) parser.add_argument("--epochs", type=int, default=EPOCHS, help="Adam epochs") parser.add_argument("--lbfgs-rounds", type=int, default=LBFGS_ROUNDS, help="L-BFGS rounds (0 = skip)") parser.add_argument("--device", default=None, choices=["cpu", "cuda"]) parser.add_argument("--checkpoint", default=CHECKPOINT) parser.add_argument("--log-every", type=int, default=LOG_EVERY) parser.add_argument("--bc-weight", type=float, default=BC_WEIGHT, help="weight on the boundary terms") parser.add_argument("--fem-n", type=int, default=FEM_N, help="reference mesh: N x N elements") parser.add_argument("--reference", default=None, help="saved FEM reference (.npz) to use instead of solving") parser.add_argument("--smoke", action="store_true", help="a very short run to check that everything works") parser.add_argument("--postprocess-only", action="store_true", help="reload the checkpoint, do not train") parser.add_argument("--no-plots", action="store_true", help="metrics only: no figures and no GIFs") parser.add_argument("--no-gifs", action="store_true") parser.add_argument("--quiet", action="store_true") args, _unknown = parser.parse_known_args(argv) epochs, rounds, fem_n = args.epochs, args.lbfgs_rounds, args.fem_n snap_every, snap_every_lbfgs = SNAP_EVERY, SNAP_EVERY_LBFGS if args.smoke: epochs, rounds, fem_n, snap_every, snap_every_lbfgs = 300, 3, 24, 50, 1 print("SMOKE RUN: 300 Adam epochs, 3 L-BFGS rounds, 24 x 24 reference mesh. The results are not meaningful.") global DEVICE if args.device: DEVICE = torch.device(args.device) net.to(DEVICE) balancer.weights["bc"] = args.bc_weight print(f"device: {DEVICE}") if args.postprocess_only: history = load_checkpoint(args.checkpoint) print(f"loaded {args.checkpoint}") else: history = train(epochs=epochs, lbfgs_rounds=rounds, log_every=args.log_every, verbose=not args.quiet, snap_every=snap_every, snap_every_lbfgs=snap_every_lbfgs) save_checkpoint(history, args.checkpoint) write_report(history) print(f"saved {args.checkpoint}, weights.npz, snapshots.npz, results.csv") ref = load_or_solve_reference(fem_n, POST_N, args.reference) raw = evaluate_network(POST_N) snaps = load_snapshots(OUTPUT_DIR / "snapshots.npz") n_adam = int(history.get("adam_epochs", epochs)) if history else epochs n_lbfgs = max(len(history.get("epoch", [])) - n_adam, 0) if history else rounds setup = { "problem": "steady lid-driven cavity, unit square, lid u = 16 x^2 (1 - x)^2", "Reynolds number": RE, "network": f"{DIM} inputs, hidden layers {HIDDEN}, 3 outputs (u, v, p), tanh, Xavier normal start", "interior points": N_INTERIOR, "boundary points per wall per epoch": N_BC, "optimiser": f"Adam, learning rate {LEARNING_RATE}, {n_adam} epochs; then L-BFGS, {n_lbfgs} rounds " f"of up to {LBFGS_INNER} iterations", "loss weights": f"equations 1, boundaries {args.bc_weight}", "seed": SEED, "device": str(DEVICE), "reference": f"finite elements, Taylor-Hood Q2-Q1, {int(ref.get('n_el', fem_n))} x {int(ref.get('n_el', fem_n))} elements", "comparison grid": f"{POST_N} x {POST_N} points", } print("post-processing:") metrics = postprocess(raw, ref, history, snaps, setup, figures=not args.no_plots, gifs=not (args.no_plots or args.no_gifs)) bundle = make_bundle(args.checkpoint) vv, f = metrics["velocity_vector"], metrics["features"] print("\nsummary") print(f" velocity relative L2 error : {vv['rel_l2']:.3e}") for k in ("u", "v", "p", "omega"): print(f" {k:<5s} relative L2 error : {metrics['fields'][k]['rel_l2']:.3e}") print(f" vortex centre PINN ({f['pinn']['vortex_x']:.4f}, {f['pinn']['vortex_y']:.4f}) " f"FEM ({f['reference']['vortex_x']:.4f}, {f['reference']['vortex_y']:.4f})") print(f"\neverything is in {bundle}") if IN_COLAB: print("To download it in Colab, run in a new cell:\n" " from google.colab import files; files.download('cavity_results.zip')") if __name__ == "__main__": main()