"""Lid-driven cavity, stationary Navier-Stokes, solved by a NONLINEAR multigrid (FAS). (u . grad) u - nu Delta u + grad p = 0, -div u = 0 on the unit square, u = (1, 0) on the lid y = 1 (corners excluded), u = 0 on the other walls, Taylor-Hood Q2/Q1, against Ghia, Ghia & Shin (J. Comput. Phys. 48, 1982). The same problem as ``solve_lid_driven_cavity_navier_stokes_2d.py``, without Newton on the fine level: the multigrid cycle IS the nonlinear solver. Usage: ``python solve_lid_driven_cavity_fas_2d.py [Re ...] [--n64] [--no-plot]`` -- by default 64x64, continuation 100, 200, ..., 1000. The solver, each piece the textbook one: * the Full Approximation Scheme (Brandt): the cycle carries the SOLUTION, the coarse levels solve ``F_H(u_H) = F_H(R u_h) + R (b - F_h(u_h))`` (the tau correction), so the fixed point is the fine-level solution, whatever the coarse discretisation. Coarse levels are therefore REDISCRETISED (``coarse_blocks="rediscretised"``): a Galerkin product is a linear operator, a FAS coarse level needs a nonlinear residual; * the smoother: multiplicative Vanka (``SchwarzSmoother(nonlinear=True)``) on vertex stars (``vanka_star_patches``): one block per pressure node, its own pressure DOF plus every velocity DOF of the cells around it. The blocks are taken from the Jacobian at the level's point, the sweep updates the nonlinear residual. It is THE Vanka for a continuous pressure: the cell Vanka that suits a discontinuous pressure does not smooth Taylor-Hood (two-grid rate 1.02 on Stokes, measured), the star with ONE pressure does, independently of h (0.22 per V(2, 2) at 16x16 and at 32x32); * the coarsest level (16x16) solved by Newton with a DENSE LU on the system bordered by the gauge mean(p) = 0 (``GaugedNewtonSolver``): the pressure constant is free in an enclosed cavity; * continuation in Re = 1/nu, each Re starting from the previous solution; the first from a nested iteration (``MG.nested_start``). The viscosity is a leaf of the weak form: every Re runs the same executables. The damping ``omega`` of the sweep is lowered past Re = 400: the Vanka blocks are frozen at the start of the cycle, and at a high cell Reynolds number the iterate moves too much within the cycle for an undamped update. Measured (CPU, 64x64, 37 507 DOFs, 3 levels down to 16x16, V(2, 2), until ||F|| < 1e-9, as printed; largest gap to Ghia's u(0.5, y)): Re omega cycles mean rate solve gap to Ghia 100 0.5 17 0.29 19.5 s 5.1e-3 (compiles the cycle) 200 0.5 17 0.36 7.8 s 300 0.5 16 0.37 7.1 s 400 0.5 16 0.37 7.2 s 2.9e-3 600 0.3 31 0.60 14.1 s 800 0.3 50 0.75 23.3 s 1000 0.3 147 0.91 66.3 s 6.5e-3 Setup: 5.7 s for the first Re, 0.3 s for the next ones. The cycle compiles ONCE for the whole continuation (viscosity and omega are leaves, the bases module-level functions); the first solve's extra ~12 s is the trace and compilation of ``_solve_step`` and ``_nested_start``. At Re = 1000 the contraction is steady (0.93 per cycle from cycle 20 on, no stagnation) but slow. The gap to Ghia stops moving long before convergence: it is the discretisation error at 64x64. ⚠ To improve: the coarse levels are unstabilised Galerkin at a cell Reynolds number of ~60. The same accuracy as the Newton-Krylov example (2.9e-3 at 64x64 for Re = 400), and the cycle counts do not grow with h (Re = 100 / 600 / 800: 18 / 36 / 71 cycles at 32x32, 17 / 31 / 50 at 64x64). Re = 1000, which the Newton-Krylov example does not reach, diverges at 32x32 (coarsest 8x8, at a cell Reynolds number of 125) and converges at 64x64. Printed per Re: cycles, residual, mean contraction per cycle, time, gap to Ghia. """ import sys import time import jax.numpy as jnp import matplotlib.pyplot as plt import numpy as np from scimba_jax.linear_approximation.basis.analytic_bases import local_lagrange_basis from scimba_jax.linear_approximation.basis.dof_map import LagrangeDofMap from scimba_jax.linear_approximation.basis.general_bases import AnalyticBasis from scimba_jax.linear_approximation.galerkin.fem.elliptic_saddle_scheme import ( EllipticSaddleScheme, ) from scimba_jax.linear_approximation.meshes.mesh import Mesh from scimba_jax.linear_approximation.quad.gauss_quad import UnitSquareTensorized from scimba_jax.linear_approximation.solvers import GaugedNewtonSolver from scimba_jax.linear_approximation.solvers.multigrid import MG from scimba_jax.linear_approximation.solvers.smoothers import ( SchwarzSmoother, vanka_star_patches, ) from scimba_jax.linear_approximation.transfer.hierarchy import ( build_hierarchy_structured, ) from scimba_jax.linear_approximation.variables.variables_fe import VariablesFE from scimba_jax.mapping.mapping import InvertibleFunction, Mapping from scimba_jax.physical_models.abstract_physical_weak_model import ( AbstractPhysicalWeakModel, ) from scimba_jax.physical_models.classical_weakform.navier_stokes_weak_form import ( NavierStokesWeakForm, ) from scimba_jax.physical_models.weak_boundary_conditions import Dirichlet from scimba_jax.plots.plots_galerkin import plot_flow_2d N_CELLS = int(next((a[3:] for a in sys.argv if a[:3] == "--n" and a[3:].isdigit()), 64)) REYNOLDS = tuple(float(r) for r in sys.argv[1:] if not r.startswith("-")) or ( 100.0, 200.0, 300.0, 400.0, 600.0, 800.0, 1000.0, ) N_COARSE = 16 SWEEPS = 2 # V(2, 2) # ||F|| to reach: the last Re only needs its algebraic error well below the # discretisation error, an intermediate one only a good start for the next # (see the measurements in the docstring). TOL = 1e-7 TOL_CONTINUATION = 1e-6 MAX_CYCLES = 200 U, P = 0, 1 # Ghia, Ghia & Shin (1982): u(0.5, y) and v(x, 0.5). GHIA_Y = np.array([1.0, 0.9766, 0.9688, 0.9609, 0.9531, 0.8516, 0.7344, 0.6172, 0.5, 0.4531, 0.2813, 0.1719, 0.1016, 0.0703, 0.0625, 0.0547, 0.0]) # fmt: skip GHIA_U = { 100: [1.0, 0.84123, 0.78871, 0.73722, 0.68717, 0.23151, 0.00332, -0.13641, -0.20581, -0.21090, -0.15662, -0.10150, -0.06434, -0.04775, -0.04192, -0.03717, 0.0], 400: [1.0, 0.75837, 0.68439, 0.61756, 0.55892, 0.29093, 0.16256, 0.02135, -0.11477, -0.17119, -0.32726, -0.24299, -0.14612, -0.10338, -0.09266, -0.08186, 0.0], 1000: [1.0, 0.65928, 0.57492, 0.51117, 0.46604, 0.33304, 0.18719, 0.05702, -0.06080, -0.10648, -0.27805, -0.38289, -0.29730, -0.22220, -0.20196, -0.18109, 0.0], } # fmt: skip GHIA_X = np.array([1.0, 0.9688, 0.9609, 0.9531, 0.9453, 0.9063, 0.8594, 0.8047, 0.5, 0.2344, 0.2266, 0.1563, 0.0938, 0.0781, 0.0703, 0.0625, 0.0]) # fmt: skip GHIA_V = { 100: [0.0, -0.05906, -0.07391, -0.08864, -0.10313, -0.16914, -0.22445, -0.24533, 0.05454, 0.17527, 0.17507, 0.16077, 0.12317, 0.10890, 0.10091, 0.09233, 0.0], 400: [0.0, -0.12146, -0.15663, -0.19254, -0.22847, -0.23827, -0.44993, -0.38598, 0.05186, 0.30174, 0.30203, 0.28124, 0.22965, 0.20920, 0.19713, 0.18360, 0.0], 1000: [0.0, -0.21388, -0.27669, -0.33714, -0.39188, -0.51550, -0.42665, -0.31966, 0.02526, 0.32235, 0.33075, 0.37095, 0.32627, 0.30353, 0.29012, 0.27485, 0.0], } # fmt: skip # ⚠ Module-level functions, not lambdas: a basis and a mapping keep their # function in the pytree's aux_data, compared by IDENTITY. A lambda built per # call makes every level of every Reynolds number a new treedef, and the FAS # cycle recompiles at each step of the continuation (measured: 1.9 s per Re at # 32x32). def identity(x): return x def velocity_basis(y, i, m): return local_lagrange_basis(y, i, m, order=2, out_dim=2) def pressure_basis(y, i, m): return local_lagrange_basis(y, i, m, order=1, out_dim=1) def make_mesh(n_cells): return Mesh( dim=2, n_cells=(n_cells, n_cells), ref_quad=UnitSquareTensorized(dim=2, order=4), mapping=Mapping(mappings=[InvertibleFunction(identity, identity)]), ) def make_space(mesh, local_basis, nb_basis, out_dim, basis_type): basis = AnalyticBasis( nb_basis=nb_basis, out_dim=out_dim, mesh=mesh, local_basis=local_basis, basis_type=basis_type, ) return VariablesFE(basis=basis, nb_variables=out_dim, dof_map=LagrangeDofMap) def lid(x): """u = (1, 0) on the lid, corners excluded; zero on the other walls.""" on_lid = (x[1] > 1.0 - 1e-12) & (x[0] > 1e-12) & (x[0] < 1.0 - 1e-12) return jnp.array([jnp.where(on_lid, 1.0, 0.0), 0.0]) def omega(re): """Damping of the Vanka sweep: 0.5 up to Re = 400, 0.3 beyond.""" return 0.5 if re <= 400.0 else 0.3 def navier_stokes_factory(nu): """The scheme on an ``n_cells`` grid: one per level of the hierarchy.""" def factory(n_cells): mesh = make_mesh(n_cells[0]) model = AbstractPhysicalWeakModel(dim=2) model.add_weak_form( "main", NavierStokesWeakForm(dim=2, viscosity=jnp.asarray(nu)) ) model.add_boundary_condition("0/boundary", Dirichlet(lid)) return EllipticSaddleScheme( model, [ make_space(mesh, velocity_basis, 9, 2, "field"), make_space(mesh, pressure_basis, 4, 1, "scalar"), ], ) return factory def fas(re, n_levels): """The FAS multigrid for one Reynolds number.""" hierarchy = build_hierarchy_structured( navier_stokes_factory(1.0 / re), (N_CELLS, N_CELLS), n_levels, coarse_blocks="rediscretised", ) coarsest = hierarchy.levels[0] n_velocity = int(coarsest.scheme.variables_list[U].dofsl.size) return MG( hierarchy, SchwarzSmoother( blocks=vanka_star_patches, nonlinear=True, # A jnp value is a LEAF: changing omega does not recompile the cycle # (a Python float is a compile-time constant; measured 8 s at 64x64). omega=jnp.asarray(omega(re)), ), nu_pre=SWEEPS, nu_post=SWEEPS, fas=True, coarse_solver=GaugedNewtonSolver(gauge=(n_velocity, coarsest.n_dofs)), ) def plot(dofsl, scheme, re): """The PINN figure (streamlines, vorticity, pressure) and Ghia's centrelines.""" velocity, pressure = scheme.variables_list velocity.dofsl, pressure.dofsl = scheme._split_dofs(dofsl) fig, axes = plot_flow_2d( velocity, pressure, title=f"Lid-driven cavity, FAS, Q2/Q1 {N_CELLS}x{N_CELLS}, Re = {re:g}", n_extra=1, show=False, ) s = np.linspace(1e-6, 1.0 - 1e-6, 200) half = 0.5 * np.ones_like(s) u = np.asarray(velocity.evaluate(jnp.asarray(np.stack([half, s], -1))))[:, 0] v = np.asarray(velocity.evaluate(jnp.asarray(np.stack([s, half], -1))))[:, 1] ax = axes[3] ax.plot(s, u, "C0", label="u(0.5, s)") ax.plot(s, v, "C3", label="v(s, 0.5)") ax.plot(GHIA_Y, GHIA_U[int(re)], "o", mfc="none", c="C0", label="Ghia et al.") ax.plot(GHIA_X, GHIA_V[int(re)], "s", mfc="none", c="C3") ax.set_xlabel("s") ax.set_title("Centrelines") ax.grid(alpha=0.3) ax.legend() fig.tight_layout() def gap_to_ghia(dofsl, scheme, re): """Largest gap to Ghia's u(0.5, y), walls excluded; nan without a reference.""" if int(re) not in GHIA_U: return float("nan") velocity = scheme.variables_list[U] velocity.dofsl = scheme._split_dofs(dofsl)[U] inside = np.clip(GHIA_Y, 1e-9, 1.0 - 1e-9) points = jnp.stack([0.5 * jnp.ones(inside.size), jnp.asarray(inside)], -1) u_line = np.asarray(velocity.evaluate(points))[:, 0] return float(np.abs(u_line - np.array(GHIA_U[int(re)]))[1:-1].max()) def main(): n_levels = int(np.log2(N_CELLS // N_COARSE)) + 1 print( f"Lid-driven cavity, FAS, Q2/Q1 {N_CELLS}x{N_CELLS}, {n_levels} levels " f"(coarsest {N_COARSE}x{N_COARSE}, gauged LU), Vanka star V({SWEEPS}, {SWEEPS})" ) dofsl = None for re in REYNOLDS: tol = TOL if re == REYNOLDS[-1] else TOL_CONTINUATION started = time.perf_counter() mg = fas(re, n_levels) zero = jnp.zeros(mg.hierarchy.finest.n_dofs) setup = time.perf_counter() - started started = time.perf_counter() if dofsl is None: dofsl = mg.nested_start(zero) first = float(jnp.linalg.norm(mg.hierarchy.finest.residual_fn(dofsl))) dofsl, cycles, norm = mg.solve( zero, u0=dofsl, tol=tol, max_cycles=MAX_CYCLES, relinearise=True ) dofsl.block_until_ready() seconds = time.perf_counter() - started rate = (float(norm) / first) ** (1.0 / max(cycles, 1)) scheme = mg.hierarchy.finest.scheme print( f" Re {re:6.0f} (omega {omega(re)}): {cycles:3d} cycles, " f"||F|| {first:.1e} -> {float(norm):.1e}, mean rate {rate:.2f}, " f"setup {setup:.1f} s, solve {seconds:.1f} s | " f"gap to Ghia u {gap_to_ghia(dofsl, scheme, re):.2e}", flush=True, ) if not np.isfinite(float(norm)): print(" diverged: stopping the continuation") break if "--no-plot" not in sys.argv and int(re) in GHIA_U: plot(dofsl, scheme, re) plt.show() if __name__ == "__main__": main()