"""Stokes in discontinuous Galerkin, Q2/Q1, with the block preconditioners of the FEM. -nu Delta u + grad p = f, -div u = 0, u = 0 on the boundary The manufactured problem of ``fem/.../solve_stokes_block_preconditioners_2d.py``, discretised by DG: velocity in discontinuous Q2, pressure in discontinuous Q1 (the pair of Hansbo & Larson -- inf-sup stable without any pressure stabilisation, the pressure block of the Jacobian zero as for Taylor-Hood). The volume terms are the same ``StokesWeakForm`` as the FEM; the faces are ``StokesDGFlux``: interior penalty on the velocity, ``{p} [v.n]`` and ``{q} [u.n]`` for the coupling, in the symmetric convention, so MINRES applies. The point of the file is that the preconditioner is the SAME object as in the FEM -- ``BlockSchurPreconditioner`` over the two spaces of a multi-space scheme -- with the two inverses a DG space calls for: * velocity: one V-cycle of a DG multigrid (cellwise transfers), Schwarz smoother, read off the system's own element matrices, Galerkin coarse levels (the default of ``build_hierarchy_structured``); * pressure: ``-nu M_p^-1``, and a DG mass being block diagonal per cell, its inverse is EXACT and local (``CellBlockInverse.mass_of``), where the continuous one needs Chebyshev iterations. Measured (CPU, Krylov to 1e-8, iterations; in brackets, rediscretised coarse levels instead of Galerkin): ========================== ========= ========= ========= 8x8 16x16 32x32 ========================== ========= ========= ========= none, MINRES 925 1 758 -- block-diagonal + MINRES 40 (41) 42 (39) 43 (39) block-triangular + FGMRES 17 (17) 18 (17) 20 (17) ========================== ========= ========= ========= A few iterations more with Galerkin coarse levels (the known +-2: the projected SIPG penalty is twice the rediscretised one on coarse faces), and still less time overall at 32x32 -- 4.3 s against 4.7 s for the triangular row, setup and compilation included -- the coarse levels being projected rather than assembled. The counts are flat in h, the triangular one at those of Taylor-Hood. The errors do not depend on the preconditioner: velocity 2.57e-3, 3.26e-4, 4.10e-5 (order 3), pressure 7.82e-3, 1.32e-3, 2.81e-4. """ import time import jax import jax.numpy as jnp import numpy as np from scimba_jax.linear_approximation.basis.analytic_bases import local_lagrange_basis from scimba_jax.linear_approximation.basis.general_bases import AnalyticBasis from scimba_jax.linear_approximation.galerkin.dg.elliptic_dg_scheme import ( EllipticDGscheme, ) from scimba_jax.linear_approximation.galerkin.dg.elliptic_dg_scheme_multi_space import ( # noqa: E501 EllipticDGschemeMultipleSpaces, ) from scimba_jax.linear_approximation.galerkin.dg.flux import SIPGFlux, StokesDGFlux 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 LinearKrylov from scimba_jax.linear_approximation.solvers.multigrid import MG from scimba_jax.linear_approximation.solvers.schur import ( BlockSchurPreconditioner, CellBlockInverse, MultigridInverse, ScaledInverse, ) from scimba_jax.linear_approximation.solvers.smoothers import SchwarzSmoother from scimba_jax.linear_approximation.transfer.hierarchy import ( build_hierarchy_structured, ) from scimba_jax.linear_approximation.variables.variables_dg import VariablesDG 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.stokes_weak_form import ( StokesWeakForm, VectorLaplacianWeakForm, ) from scimba_jax.physical_models.weak_boundary_conditions import Dirichlet NU = 1.0 ORDER = 2 # velocity Q2, pressure Q1 SIGMA = 4.0 * ORDER * (ORDER + 1) N_CELLS = (8, 16, 32) N_COARSE = 2 U, P = 0, 1 def u_exact(x): sx, cx = jnp.sin(jnp.pi * x[0]), jnp.cos(jnp.pi * x[0]) sy, cy = jnp.sin(jnp.pi * x[1]), jnp.cos(jnp.pi * x[1]) return jnp.pi * jnp.array([sx * sx * sy * cy, -sy * sy * sx * cx]) def p_exact(x): return jnp.array([jnp.cos(jnp.pi * x[0]) * jnp.sin(jnp.pi * x[1])]) def source(x): laplacian = jnp.trace(jax.jacfwd(jax.jacfwd(u_exact))(x), axis1=1, axis2=2) return -NU * laplacian + jax.jacfwd(lambda y: p_exact(y)[0])(x) def make_mesh(n_cells): return Mesh( dim=2, n_cells=(n_cells, n_cells), ref_quad=UnitSquareTensorized(dim=2, order=ORDER + 2), mapping=Mapping(mappings=[InvertibleFunction(lambda x: x, lambda y: y)]), ) def make_space(mesh, order, out_dim, basis_type): basis = AnalyticBasis( nb_basis=(order + 1) ** 2, out_dim=out_dim, mesh=mesh, basis_type=basis_type, local_basis=lambda y, i, m: local_lagrange_basis( y, i, m, order=order, out_dim=out_dim ), ) return VariablesDG(basis=basis, nb_variables=out_dim) def stokes(n_cells): mesh = make_mesh(n_cells) model = AbstractPhysicalWeakModel(dim=2) model.add_weak_form("main", StokesWeakForm(dim=2, f=source, viscosity=NU)) model.add_boundary_condition("0/boundary", Dirichlet(lambda x: jnp.zeros(2))) # The multi-space DG scheme wants a datum on every space; the flux reads # only the velocity's. model.add_boundary_condition("1/boundary", Dirichlet(lambda x: jnp.zeros(1))) flux = StokesDGFlux(dim=2, viscosity=NU, sigma=SIGMA, h=1.0 / n_cells) return EllipticDGschemeMultipleSpaces( pde=model, variables_list=[ make_space(mesh, ORDER, 2, "field"), make_space(mesh, ORDER - 1, 1, "scalar"), ], flux_list=[flux, flux], equation_spaces=[0, 1], use_scan_quad=False, linearization="blocks", ) def velocity_scheme(n_cells): """The velocity space alone: the DG multigrid's structure.""" model = AbstractPhysicalWeakModel.from_weak_form( VectorLaplacianWeakForm(dim=2, viscosity=NU), dirichlet=lambda x: jnp.zeros(2) ) return EllipticDGscheme( model, make_space(make_mesh(n_cells), ORDER, 2, "field"), SIPGFlux(sigma=SIGMA * NU, h=1.0 / n_cells), use_scan_quad=False, linearization="blocks", ) def velocity_multigrid(n_cells): """Galerkin coarse levels (the builder's default), projected off the fine block.""" n_levels = int(np.log2(n_cells // N_COARSE)) + 1 hierarchy = build_hierarchy_structured( lambda counts: velocity_scheme(counts[0]), (n_cells, n_cells), n_levels ) return MG(hierarchy, SchwarzSmoother()) def errors(scheme, dofsl): velocity, pressure = scheme.variables_list velocity.dofsl, pressure.dofsl = scheme._split_dofs(dofsl) grid = (np.arange(60) + 0.5) / 60 points = jnp.asarray( np.stack(np.meshgrid(grid, grid, indexing="ij"), -1).reshape(-1, 2) ) u_h, p_h = velocity.evaluate(points), pressure.evaluate(points)[:, 0] u_ref, p_ref = jax.vmap(u_exact)(points), jax.vmap(p_exact)(points)[:, 0] error_u = float(jnp.sqrt(jnp.mean(jnp.sum((u_h - u_ref) ** 2, -1)))) error_p = float( jnp.sqrt(jnp.mean(((p_h - p_h.mean()) - (p_ref - p_ref.mean())) ** 2)) ) return error_u, error_p def solver_for(name, scheme, n_cells): if name == "none, MINRES": return LinearKrylov(tol=1e-8, cg_solver="minres", max_iter_linear=20000) zeros = scheme._initial_dofs() velocity = MultigridInverse.frozen_component( velocity_multigrid(n_cells), scheme, zeros, U ) mass = CellBlockInverse.mass_of(scheme.variables_list[P]) if name.startswith("block-diagonal"): pc = BlockSchurPreconditioner( sweep=[U, P], inverses={U: velocity, P: ScaledInverse(mass, NU)}, mode="jacobi", ) return LinearKrylov(tol=1e-8, cg_solver="minres", preconditioner=pc) pc = BlockSchurPreconditioner( sweep=[P, U], inverses={U: velocity, P: ScaledInverse(mass, -NU)} ) return LinearKrylov(tol=1e-8, cg_solver="fgmres", preconditioner=pc) def main(): print(f"Stokes, DG Q{ORDER}/Q{ORDER - 1}, nu = {NU}\n") for name in ( "none, MINRES", "block-diagonal + MINRES", "block-triangular + FGMRES", ): for n_cells in N_CELLS: if name == "none, MINRES" and n_cells > 16: continue scheme = stokes(n_cells) solver = solver_for(name, scheme, n_cells) started = time.perf_counter() dofsl, report = EllipticDGschemeMultipleSpaces.solve_pure( scheme, solver=solver ) dofsl.block_until_ready() error_u, error_p = errors(scheme, dofsl) print( f" {name:26s} {n_cells:2d}x{n_cells:<2d} {dofsl.size:6d} DOFs: " f"{int(report.n_linear):5d} it, {time.perf_counter() - started:5.1f} s" f" | u {error_u:.3e} p {error_p:.3e}", flush=True, ) if __name__ == "__main__": main()