r"""Stationary Boussinesq NS in a differentially heated square cavity. Geometry: unit square [−1, 1]². Equations (Boussinesq approximation): .. math:: (u \cdot \nabla)u - \nu \Delta u + \nabla p &= \alpha (T - T_\text{ref}) g \\ \nabla \cdot u &= 0 \\ (u \cdot \nabla) T - \kappa \Delta T &= 0 Boundary conditions — same structure as the Torch reference: * west wall: u = 0, T = +1 (hot) — **hard constraint** via post-processing * east wall: u = 0, T = −1 (cold) — **hard constraint** via post-processing * north/south: u = 0 — **hard constraint** via post-processing * ∂T/∂n = 0 — **weak** (boundary residual) Hard constraint implementation (post-processing on network output): * u, v → multiplied by φ = (x+1)(1−x)(y+1)(1−y), which vanishes on all walls * T → T_net · (1−x²) − x which gives T(−1)=+1, T(+1)=−1 * p → unconstrained (pressure is determined up to a constant) A single ``MODE`` flag selects the approximation-space layout: * ``"vec"`` — one network, output ``[u_x, u_y, p, T]`` * ``"field_scalar"`` — three networks: one field network for ``u``, one scalar network for ``p``, one scalar network for ``T`` * ``"scalar"`` — four scalar networks: one per variable ``u_x, u_y, p, T`` Uses :class:`NavierStokesThermalND` from ``scimba_jax.physical_models.elliptic_pde.navier_stokes_stationary``. """ import timeit import jax import jax.numpy as jnp import matplotlib.pyplot as plt from scimba_jax.domains.meshless_domains.domains_2d import Square2D from scimba_jax.nonlinear_approximation.approximation_spaces.approximation_spaces import ( ApproximationSpace, ) from scimba_jax.nonlinear_approximation.integration.monte_carlo import ( DomainSampler, TensorizedSampler, ) from scimba_jax.nonlinear_approximation.model_class.funcparam_vectorial import ( ParamFieldFunction, ParamVecFunction, ) from scimba_jax.nonlinear_approximation.networks.mlp import MLP from scimba_jax.nonlinear_approximation.numerical_solvers.projectors import Projector from scimba_jax.physical_models.abstract_residuals import BoundaryResidual, static from scimba_jax.physical_models.elliptic_pde.navier_stokes_stationary import ( NavierStokesThermalND, _extract_u_p_t, ) from scimba_jax.plots.plots_nd import plot_abstract_approx_space # ── configuration ───────────────────────────────────────────────────────────── # "vec" → one network, output [u_x, u_y, p, T] # "field_scalar" → three networks: field net for u, scalar nets for p and T # "scalar" → four scalar networks: u_x, u_y, p, T MODE = "vec" # ── hyper-parameters ────────────────────────────────────────────────────────── N_COLLOC = 5000 N_BC_COLLOC = 3000 N_EPOCHS = 100 # physical parameters nu = 0.01 kappa = 0.01 alpha = 0.1 T_ref = 0.0 # Loss weights — interior: [NS_x, NS_y, div, energy]; bc ns: [u_x, u_y, ∂T/∂n] WEIGHTS = { "interior": [1.0, 1.0, 2.0, 1.0], "bc ns": [20.0, 20.0, 20.0], } # ── hard-constraint post-processing ────────────────────────────────────────── # Box bounds: x ∈ [a, b], y ∈ [a, b] _a, _b = -1.0, 1.0 # phi_sq = (x-a)(b-x)(y-a)(b-y) vanishes on all four walls → u = v = 0 exactly # T = T_net*(x-a)(b-x) + (-x) satisfies T(-1)=+1, T(+1)=-1 def post_processing_uv(output, x): x_, y_ = x[0], x[1] phi = (x_ - _a) * (_b - x_) * (y_ - _a) * (_b - y_) return jnp.array([output[0] * phi, output[1] * phi]) def post_processing_p(output, x): return output # pressure: unconstrained def post_processing_t(output, x): x_ = x[0] return jnp.array([output[0] * (x_ - _a) * (_b - x_) + (-x_)]) def post_processing_vec(output, x): x_, y_ = x[0], x[1] phi_sq = (x_ - _a) * (_b - x_) * (y_ - _a) * (_b - y_) u = output[0] * phi_sq v = output[1] * phi_sq p = output[2] T = output[3] * (x_ - _a) * (_b - x_) + (-x_) return jnp.array([u, v, p, T]) def _post_proc_u_scalar(output, x): x_, y_ = x[0], x[1] phi = (x_ - _a) * (_b - x_) * (y_ - _a) * (_b - y_) return jnp.array([output[0] * phi]) # ── boundary-condition residual (weak: ∂T/∂n = 0 on north/south) ───────────── class NeumannThermalWallBC(BoundaryResidual): """No-slip velocity + insulated temperature (∂T/∂n = 0). Residual: (u₁, …, u_dim, ∂T/∂n) with f_rhs = 0. u components are already zero by hard constraint; they appear here only for structural consistency (size = dim + 1). """ _ns_dim: int = static(2) _ns_mode: str = static("vec") def __init__(self, domain, dim: int = 2, mode: str = "vec"): super().__init__(domain=domain, size=dim + 1, model_type="x", f_rhs=None) self._ns_dim = dim self._ns_mode = mode def construct_residual(self, *vars): u, _, T = _extract_u_p_t(vars, self._ns_mode, self._ns_dim) grad_T = T.gradient_x() n_model = ParamFieldFunction( T.dims, lambda *args: args[T.argnums["n"]], T.f_type, ) dTdn = grad_T.dot(n_model) return ParamVecFunction.cat( [u.component(i) for i in range(self._ns_dim)] + [dTdn] ) # ── domain ──────────────────────────────────────────────────────────────────── key = jax.random.PRNGKey(0) box = [(_a, _b), (_a, _b)] dx = Square2D(box, is_main_domain=True) dx.set_boundaries_dict({"bc ns": ["bc north", "bc south"]}) # ── physics model ───────────────────────────────────────────────────────────── boundaries = dx.get_all_boundaries() bc_residuals = { "bc ns": NeumannThermalWallBC(boundaries["bc ns"], dim=2, mode=MODE), } def _f_rhs_gravity(_x): return jnp.array([0.0, 9.81, 0.0, 0.0]) model = NavierStokesThermalND( dx, dim=2, nu=nu, kappa=kappa, alpha=alpha, T_ref=T_ref, gravity=(0.0, -9.81), mode=MODE, bc_residuals=bc_residuals, f_rhs=_f_rhs_gravity, ) # ── approximation space (depends on MODE) ───────────────────────────────────── if MODE == "vec": nn = MLP(in_size=2, out_size=4, hidden_sizes=[20, 20, 20], key=key) space = ApproximationSpace( {"x": 2}, [(nn, "vec", 4)], model_type="x", post_processing=post_processing_vec, ) elif MODE == "field_scalar": nn_uv = MLP(in_size=2, out_size=2, hidden_sizes=[18, 18, 18, 18], key=key) nn_p = MLP(in_size=2, out_size=1, hidden_sizes=[16, 16, 16], key=key) nn_T = MLP(in_size=2, out_size=1, hidden_sizes=[16, 16, 16], key=key) space = ApproximationSpace( {"x": 2}, [(nn_uv, "field", 2), (nn_p, "scalar", 1), (nn_T, "scalar", 1)], model_type="x", post_processing=[post_processing_uv, post_processing_p, post_processing_t], ) else: # "scalar" nn_ux = MLP(in_size=2, out_size=1, hidden_sizes=[15, 15, 15], key=key) nn_uy = MLP(in_size=2, out_size=1, hidden_sizes=[15, 15, 15], key=key) nn_p = MLP(in_size=2, out_size=1, hidden_sizes=[15, 15, 15], key=key) nn_T = MLP(in_size=2, out_size=1, hidden_sizes=[15, 15, 15], key=key) space = ApproximationSpace( {"x": 2}, [ (nn_ux, "scalar", 1), (nn_uy, "scalar", 1), (nn_p, "scalar", 1), (nn_T, "scalar", 1), ], model_type="x", post_processing=[ _post_proc_u_scalar, _post_proc_u_scalar, post_processing_p, post_processing_t, ], ) # ── sampler ─────────────────────────────────────────────────────────────────── sampler = TensorizedSampler( [DomainSampler(dx)], bc=True, ) # ── training ────────────────────────────────────────────────────────────────── pinn = Projector(model, space, sampler, weights=WEIGHTS, one_loss_by_equation=True) key, sample_dict = sampler.sample(key, N_COLLOC, N_BC_COLLOC) loss0 = pinn.evaluate_loss(space, sample_dict) print("initial loss:", loss0) start = timeit.default_timer() key, pinn = pinn.project(key, space, N_EPOCHS, N_COLLOC, N_BC_COLLOC) end = timeit.default_timer() print(f"best loss: {pinn.best_loss}") print(f"time for {N_EPOCHS} epochs: {end - start:.1f}s") # ── visualisation ───────────────────────────────────────────────────────────── plot_abstract_approx_space( pinn.space, dx, components=[0, 1, 2, 3], loss=pinn.losses, residual=pinn.model, draw_contours=True, n_drawn_contours=20, loss_groups=["bc", "interior"], ) plt.show()