"""Multigrille sur un système MULTI-SPACE à maillages DIFFÉRENTS. Même physique que ``solve_2d_system_diffusion_advection_multi_space`` -- les constantes sont reprises telles quelles : -ε Δu + b₁·∇u + c₁·∇v = f sur Ω, u|∂Ω = 0 -ε Δv + b₂·∇v + c₂·∇u = f sur Ω, v|∂Ω = 0 ``u`` et ``v`` vivent chacun sur SON maillage, de finesse différente, et le couplage croisé passe par une évaluation d'un espace dans les mailles de l'autre. Ce que cet exemple ajoute, et qui n'était pas possible avant-hier, c'est de le résoudre par MULTIGRILLE. Ce qu'il faut savoir du multigrille sur multi-space --------------------------------------------------- Le transfert est **bloc-diagonal, un bloc par espace** : deux espaces d'un même schéma n'ont jamais de ddl commun -- ce sont des inconnues différentes, pas deux morceaux d'une même inconnue -- donc le vecteur est une concaténation et rien ne lie les blocs. Le couplage vit dans la forme faible, pas dans la numérotation. Et chaque espace apporte SON transfert, choisi par la même règle que s'il était seul. Ici les deux sont cartésiens nodaux, donc deux tensoriels de tailles différentes ; un espace spline ou non structuré aurait le sien, sans que le cycle s'en aperçoive. ⚠ **Les comptes de mailles ne sont pas ceux de l'exemple d'origine** (40 et 30), et c'est la seule chose qui change. Une hiérarchie divise les comptes par deux à chaque niveau et REFUSE de le faire à la louche : arrondir donnerait des niveaux non emboîtés, donc un transfert faux qu'aucune vérification de forme n'attraperait. 48 et 32 descendent proprement sur trois niveaux (48→24→12, 32→16→8), 30 non. ⚠ **BiCGSTAB et non CG** : l'advection et le couplage croisé rendent la jacobienne non symétrique, et le CG n'y converge pas lentement -- il stagne ou diverge, en rendant un itéré qui a l'air convergé. Lancer : python solve_2d_system_diffusion_advection_multi_space_mg.py """ import time 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.fem.elliptic_fe_scheme_multi_space import ( # noqa: E501 EllipticFEschemeMultipleSpaces, ) 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.multigrid import MG from scimba_jax.linear_approximation.solvers.newton import NewtonKrylov from scimba_jax.linear_approximation.solvers.preconditioners import ( coloured_jacobi, ) from scimba_jax.linear_approximation.solvers.smoothers import DampedJacobiSmoother 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.nonlinear_approximation.model_class.funcparam_vectorial import ( ParamVecFunction, ) from scimba_jax.physical_models.abstract_physical_weak_model import ( AbstractPhysicalWeakModel, ) from scimba_jax.physical_models.abstract_weak_form import AbstractWeakForm from scimba_jax.physical_models.weak_boundary_conditions import Dirichlet # ── Paramètres : ceux de l'exemple d'origine, sauf les comptes de mailles ──── physical_dim = 2 quad_order = 3 poly_order = 1 #: ⚠ 48 et 32 au lieu de 40 et 30 : voir l'en-tête. Les deux espaces gardent #: des finesses DIFFÉRENTES, qui est tout l'intérêt du cas. n_cells_u = 48 n_cells_v = 32 n_levels = 3 sigma_gauss = 0.07 x0 = jnp.array([0.5, 0.5]) eps = 0.06 b1_vec = jnp.array([1.0, 0.0]) b2_vec = jnp.array([-1.0, 0.0]) c1_vec = jnp.array([0.0, 0.3]) c2_vec = jnp.array([0.0, -0.3]) mapping_id = InvertibleFunction(lambda x: x, lambda y: y) def gaussian_source(x): return jnp.exp(-jnp.sum((x - x0) ** 2) / (2.0 * sigma_gauss**2)) / ( 2.0 * jnp.pi * sigma_gauss**2 ) # ── Forme faible : identique à l'exemple d'origine ─────────────────────────── class SystemDiffAdvWeakForm(AbstractWeakForm): """``u`` et ``v`` couplés par leurs gradients croisés.""" def __init__(self, dim, eps, b1, b2, c1, c2, f): super().__init__(dim=dim) self.A = lambda x: eps * jnp.eye(dim) self.b1 = b1 self.b2 = b2 self.c1 = c1 self.c2 = c2 self.f = f def bilinear_form(self, u: ParamVecFunction, v: ParamVecFunction): u_1, u_2 = u # un argument par espace v_1, v_2 = v fields = self.get_fields() grad_u1, grad_u2 = u_1.gradient("x"), u_2.gradient("x") grad_v1, grad_v2 = v_1.gradient("x"), v_2.gradient("x") equation_u = ( grad_v1.dot(fields["A"] @ grad_u1) + grad_u1.dot(fields["b1"]) * v_1 + grad_u2.dot(fields["c1"]) * v_1 ) equation_v = ( grad_v2.dot(fields["A"] @ grad_u2) + grad_u2.dot(fields["b2"]) * v_2 + grad_u1.dot(fields["c2"]) * v_2 ) return ParamVecFunction.cat([equation_u, equation_v]) def linear_form(self, v: ParamVecFunction): v_1, v_2 = v f = self.get_fields()["f"] return ParamVecFunction.cat([f * v_1, f * v_2]) def make_variables(n_cells): """Un espace CG-FEM Q1 sur sa propre grille.""" mesh = Mesh( dim=physical_dim, n_cells=(int(n_cells), int(n_cells)), ref_quad=UnitSquareTensorized(dim=physical_dim, order=quad_order), mapping=Mapping(mappings=[mapping_id]), ) basis = AnalyticBasis( nb_basis=(poly_order + 1) ** physical_dim, out_dim=1, mesh=mesh, local_basis=lambda coords, i, m: local_lagrange_basis( coords, i, m, order=poly_order, out_dim=1 ), basis_type="scalar", ) return VariablesFE(basis=basis, nb_variables=1) def make_scheme(counts): """Le schéma des DEUX espaces, à partir d'un compte PAR ESPACE. ⚠ ``build_hierarchy_structured`` ne fait que diviser les nombres qu'on lui donne et laisse la fabrique décider de ce qu'ils désignent -- ici un compte par espace et non deux directions d'un même maillage. C'est ce qui permet à deux espaces de finesses différentes de descendre chacun la sienne. """ count_u, count_v = (int(n) for n in np.atleast_1d(counts)) pde = SystemDiffAdvWeakForm( dim=physical_dim, eps=eps, b1=lambda x: b1_vec, b2=lambda x: b2_vec, c1=lambda x: c1_vec, c2=lambda x: c2_vec, f=gaussian_source, ) model = AbstractPhysicalWeakModel.from_weak_form(pde) model.add_boundary_condition("0/boundary", Dirichlet(lambda x: jnp.zeros(1))) model.add_boundary_condition("1/boundary", Dirichlet(lambda x: jnp.zeros(1))) return EllipticFEschemeMultipleSpaces( pde=model, variables_list=[make_variables(count_u), make_variables(count_v)], equation_spaces=[0, 1], ) # ── Le banc : trois préconditionneurs sur le même problème ─────────────────── def run(name, preconditioner, scheme): """Résout et rapporte, depuis rien -- setup et compilation compris. Returns: ``(n_iterations, secondes, residu)``. """ start = time.perf_counter() _, report = EllipticFEschemeMultipleSpaces.solve( scheme, solver=NewtonKrylov( max_iter=1, tol=1e-9, cg_solver="bicgstab", preconditioner=preconditioner, ), return_report=True, ) return int(report.n_linear), time.perf_counter() - start, float(report.residual) def main(): """Le même problème à trois finesses : c'est la CROISSANCE qui parle. ⚠ Un tableau à une seule taille ne dirait rien d'un multigrille -- il paie un setup que les autres n'ont pas, et sur un petit cas ce setup domine. Ce qu'on veut voir est que son compte d'itérations ne bouge PAS quand le problème grossit, là où celui des autres suit la taille. """ print( f"\n-eps Lap(u) + b1.grad u + c1.grad v = f, et son symetrique en v" f"\nDeux espaces Q{poly_order} sur des maillages DIFFERENTS," f" {n_levels} niveaux\n" ) header = f" {'u x v':<12} {'ddl':>7} " + " ".join( f"{name:>20}" for name in ("sans precond", "jacobi", "multigrille") ) print(header) print(" " + "-" * (len(header) - 2)) for count_u, count_v in ((24, 16), (48, 32), (96, 64)): hierarchy = build_hierarchy_structured( make_scheme, (count_u, count_v), n_levels ) scheme = hierarchy.finest.scheme n_dofs = sum(v.ndof_linear for v in scheme.variables_list) results = [ run("sans precond", None, scheme), run("jacobi", coloured_jacobi(scheme), scheme), run( "multigrille", MG( hierarchy, DampedJacobiSmoother(omega=2.0 / 3.0), nu_pre=2, nu_post=2, ), scheme, ), ] cells = f"{count_u}^2 x {count_v}^2" cols = " ".join(f"{n:>8d} it {t:6.2f}s" for n, t, _ in results) print(f" {cells:<12} {n_dofs:>7} {cols}") print( "\n it = iterations de BiCGSTAB ; s = temps depuis rien, setup et" "\n compilation compris." "\n" "\n ⚠ Ce qu'il faut lire est la COLONNE. Mesure : 51, 101, 197" "\n iterations sans preconditionneur -- le compte DOUBLE avec la" "\n taille -- contre 5, 4, 4 en multigrille, qui ne bouge pas. C'est" "\n la seule propriete qu'on lui demande, et il l'a sur un multi-space" "\n a maillages differents comme ailleurs." "\n" "\n ⚠ Deux choses que ce tableau dit et qu'on aimerait taire." "\n D'abord, le Jacobi n'achete RIEN ici : 51, 101, 196, c'est-a-dire" "\n le compte du cas non preconditionne. Une diagonale ne voit pas un" "\n operateur dont le mal est global." "\n Ensuite, le multigrille est PLUS LENT en temps total aux trois" "\n tailles. Son temps est quasi constant (4.8, 4.9, 5.0 s) parce" "\n qu'il est domine par la COMPILATION du cycle et non par le" "\n probleme ; celui des autres croit lentement (1.6 -> 1.9 s) parce" "\n qu'une iteration est bon marche a ces tailles. Le croisement est" "\n donc plus loin -- ou, plus surement, dans un second solve sur le" "\n meme operateur, ou la compilation et le setup sont deja payes." "\n C'est exactement le regime d'une boucle en temps ou d'un Newton." ) if __name__ == "__main__": main()