""" 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. ------------------------------------------------------------------------------- Finite element reference for the steady lid-driven cavity (incompressible Navier-Stokes). u u_x + v u_y + p_x = (1/Re) (u_xx + u_yy) u v_x + v v_y + p_y = (1/Re) (v_xx + v_yy) u_x + v_y = 0 on the unit square Taylor-Hood elements: Q2 (biquadratic) velocity and Q1 (bilinear) pressure on a uniform N x N mesh of squares. The nonlinear system is solved with Newton's method. The pressure is fixed at one node (it is only defined up to a constant); compare pressures after subtracting their mean over the cavity. Lid options: "regularised" u = 16 x^2 (1 - x)^2 (zero at the corners, 1 at the middle) or "classic" u = 1 (the usual benchmark, with a corner singularity). Command line (needs numpy and scipy only): python cavity_reference.py solve on 128 x 128 elements, write cavity_reference_N128.npz python cavity_reference.py --n 96 another mesh python cavity_reference.py --verify Kovasznay flow: error against the exact solution python cavity_reference.py --refine the cavity on N = 16 ... 96 against N = 128 (a few minutes) Licence: CC BY-NC 4.0, (c) 2026 Scientific ML Studio. Credit "Scientific ML Studio", https://scimlstudio.com/ """ from __future__ import annotations import numpy as np 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 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", bc=None): """bc: optional function (x, y) -> (u, v) giving the velocity on the whole boundary (for verification tests).""" 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] = lid_profile(self.xn[top], lid) self.v_bc = np.zeros(self.Nv) if bc is not None: ub, vb = bc(self.xn, self.yn) self.u_bc, self.v_bc = np.where(self.boundary, ub, 0.0), np.where(self.boundary, vb, 0.0) 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 mean_pressure(self, n=201): g = np.linspace(0, 1, n); X, Y = np.meshgrid(g, g, indexing="ij") return float(np.trapezoid(np.trapezoid(self.eval(X, Y, ("p",))["p"].reshape(n, n), g, axis=1), g)) def streamfunction(self, n=401): """psi(x, y) = integral of u dy from the bottom wall; returns grid and psi (n x n, x-major).""" g = np.linspace(0, 1, n); X, Y = np.meshgrid(g, g, indexing="ij") U = self.eval(X, Y, ("u",))["u"].reshape(n, n) psi = np.concatenate([np.zeros((n, 1)), np.cumsum(0.5 * (U[:, 1:] + U[:, :-1]) * (g[1] - g[0]), axis=1)], axis=1) return g, psi def vortex_centre(self, n=401): g, psi = self.streamfunction(n) i, j = np.unravel_index(psi.argmin(), psi.shape) return float(g[i]), float(g[j]), float(psi[i, j]) # ---------------------------------------------------------------------------------------- # Checks and command line # ---------------------------------------------------------------------------------------- def kovasznay(re=40.0): """Exact Kovasznay flow: a solution of the steady Navier-Stokes equations with no forcing.""" lam = re / 2 - np.sqrt(re ** 2 / 4 + 4 * np.pi ** 2) def exact(x, y): u = 1 - np.exp(lam * x) * np.cos(2 * np.pi * y) v = lam / (2 * np.pi) * np.exp(lam * x) * np.sin(2 * np.pi * y) p = 0.5 * (1 - np.exp(2 * lam * x)) return u, v, p return exact def verify_kovasznay(meshes=(4, 8, 16, 32, 64), re=40.0, n_eval=161): """Solve Kovasznay flow on the unit square and report the error against the exact solution.""" exact = kovasznay(re) g = np.linspace(0, 1, n_eval) X, Y = np.meshgrid(g, g, indexing="ij") ue, ve, pe = exact(X.ravel(), Y.ravel()) pe = pe - pe.mean() rows = [] print(f"Kovasznay flow, Re = {re:g}, unit square. Relative L2 error against the exact solution.") print(f"{'N':>4} {'velocity':>11} {'order':>6} {'pressure':>11} {'order':>6}") prev = None for n in meshes: f = Cavity(n_el=n, re=re, bc=lambda x, y: exact(x, y)[:2]).solve() o = f.eval(X.ravel(), Y.ravel(), ("u", "v", "p")) pn = o["p"] - o["p"].mean() ev = np.sqrt(np.sum((o["u"] - ue) ** 2 + (o["v"] - ve) ** 2) / np.sum(ue ** 2 + ve ** 2)) ep = np.linalg.norm(pn - pe) / np.linalg.norm(pe) ov = "" if prev is None else f"{np.log2(prev[0] / ev):.2f}" op = "" if prev is None else f"{np.log2(prev[1] / ep):.2f}" print(f"{n:>4} {ev:>11.3e} {ov:>6} {ep:>11.3e} {op:>6}") rows.append((n, ev, ep)) prev = (ev, ep) return rows def reference_fields(fem, n=201): """u, v, p (mean removed) and vorticity on an n x n grid, x-major: A[i, j] is at (x_i, y_j).""" g = np.linspace(0, 1, n) X, Y = np.meshgrid(g, g, indexing="ij") o = fem.eval(X, Y, ("u", "v", "p", "omega")) out = {k: o[k].reshape(n, n) for k in o} out["p"] = out["p"] - float(np.trapezoid(np.trapezoid(out["p"], g, axis=1), g, axis=0)) out["x"], out["y"] = g, g return out def refine_cavity(meshes=(16, 24, 32, 48, 64, 96), n_ref=128, re=100.0): """Mesh convergence of the regularised-lid cavity: each mesh against the n_ref mesh.""" ref = reference_fields(Cavity(n_el=n_ref, re=re).solve()) print(f"Cavity, Re = {re:g}, regularised lid. Relative L2 difference from the N = {n_ref} mesh.") print(f"{'N':>4} {'u':>10} {'v':>10} {'p':>10} {'vorticity':>10}") for n in meshes: r = reference_fields(Cavity(n_el=n, re=re).solve()) e = {k: np.linalg.norm(r[k] - ref[k]) / np.linalg.norm(ref[k]) for k in ("u", "v", "p", "omega")} print(f"{n:>4} {e['u']:>10.2e} {e['v']:>10.2e} {e['p']:>10.2e} {e['omega']:>10.2e}") if __name__ == "__main__": import argparse import time ap = argparse.ArgumentParser(description="Finite element reference for the lid-driven cavity") ap.add_argument("--n", type=int, default=128, help="elements per side") ap.add_argument("--re", type=float, default=100.0, help="Reynolds number") ap.add_argument("--verify", action="store_true", help="check the solver on Kovasznay flow") ap.add_argument("--refine", action="store_true", help="mesh convergence of the cavity") a = ap.parse_args() if a.verify: verify_kovasznay() elif a.refine: refine_cavity(re=a.re) else: t0 = time.time() fem = Cavity(n_el=a.n, re=a.re).solve() fields = reference_fields(fem) out = f"cavity_reference_N{a.n}.npz" np.savez(out, **fields) print(f"{a.n} x {a.n} Q2-Q1 elements, {len(fem.newton_history) - 1} Newton steps, " f"final residual {fem.newton_history[-1]:.1e}, {time.time() - t0:.0f} s -> {out}") x, y, psi = fem.vortex_centre() print(f"primary vortex centre ({x:.4f}, {y:.4f}), stream-function minimum {psi:.5f}")