feat(solver): add body forcing and linear drag inputs to the ns-grid solvers

mosaic/mosaic_shared/problems/navier_stokes_grid/schemas.py

@@ -126,6 +126,20 @@ class InputSchema(BaseModel):
             "Requires x_lo to be 'dirichlet' with no value."
         ),
     )
+    forcing: GridVectorField | None = Field(
+        default=None,
+        description=(
+            "Body force per unit mass (an acceleration) at the cell centers, "
+            "constant in time, same shape as v0. 2-D periodic boxes only."
+        ),
+    )
+    linear_drag: Array[(1,), Float32] | None = Field(
+        default=None,
+        description=(
+            "Coefficient α of a linear drag -αu on the right-hand side of the "
+            "momentum equation. 2-D periodic boxes only."
+        ),
+    )
 
     @property
     def extent(self) -> tuple[float, ...]:
@@ -179,6 +193,18 @@ def _check_bcs(self) -> "InputSchema":
                     f"inflow_profile has shape {tuple(self.inflow_profile.shape)}, "
                     f"expected ({ny},) to match v0"
                 )
+        if self.forcing is not None or self.linear_drag is not None:
+            if not bc.is_fully_periodic or self.obstacle is not None:
+                raise ValueError(
+                    "forcing and linear_drag require a periodic box without an obstacle"
+                )
+            if self.v0.shape[-1] != 2:
+                raise ValueError("forcing and linear_drag are only defined in 2-D")
+        if self.forcing is not None and self.forcing.shape != self.v0.shape:
+            raise ValueError(
+                f"forcing has shape {tuple(self.forcing.shape)}, expected "
+                f"{tuple(self.v0.shape)} to match v0"
+            )
         return self
 
 

mosaic/tesseracts/navier-stokes-grid/exponax/tesseract_api.py

@@ -49,6 +49,8 @@ class InputSchema(
     @model_validator(mode="after")
     def _check_periodic_bcs(self) -> "InputSchema":
         check_periodic(self, "exponax")
+        if self.forcing is not None and self.kolmogorov_forcing:
+            raise ValueError("pass either forcing or kolmogorov_forcing, not both")
         return self
 
 
@@ -113,13 +115,24 @@ def exponax_fwd(
     injection_mode: int,
     injection_scale: float,
     boundary_conditions: dict | None = None,
+    forcing: jnp.ndarray | None = None,
+    linear_drag: float | None = None,
     **_kwargs: Any,
 ) -> jnp.ndarray:
     """Run 2D or 3D incompressible Navier-Stokes using exponax.
 
     2D uses the pseudo-spectral streamfunction-vorticity stepper with an
     internal velocity<->vorticity conversion so the interface matches the
-    other CFD tesseracts.
+    other CFD tesseracts. Exponax's 2D steppers evolve q = -ω: they take the
+    stream function from Δψ = q and the velocity from (∂ᵧψ, -∂ₓψ), whose
+    vorticity is -q.
+
+    A ``forcing`` field enters as the injection term of Exponax's Kolmogorov
+    stepper, the spectral curl of the force on the grid as given. Its own
+    ``injection_mode``/``injection_scale`` cannot express the canonical force:
+    Exponax puts sample j at y = j·h, the canonical layout at (j + ½)h, so its
+    sin(κy) is half a cell off. ``linear_drag`` α maps to Exponax's drag λ,
+    which adds λω and so damps for λ < 0.
 
     3D uses the Leray-projected velocity stepper directly.
 
@@ -135,6 +148,8 @@ def exponax_fwd(
         injection_mode: Forcing wavenumber.
         injection_scale: Forcing amplitude.
         boundary_conditions: Boundary condition config (must be fully periodic).
+        forcing: Body force at the grid points, shape of v0 (2D only).
+        linear_drag: Coefficient α of a linear drag -αu (2D only).
 
     Returns:
         Final velocity field, same shape as v0.
@@ -152,15 +167,18 @@ def exponax_fwd(
     }
 
     if ndim == 2:
+        if linear_drag is not None:
+            stepper_kwargs["drag"] = drag - linear_drag
+
         # v0: (N, N, 1, 2) -> squeeze z -> (N, N, 2)
         v = v0[:, :, 0, :]
 
-        # Mean velocity is conserved by incompressible periodic NS (Galilean invariance).
-        # The vorticity formulation strips it (∂const/∂x = 0), so we save and restore it.
+        # The vorticity formulation strips the mean velocity (∂const/∂x = 0),
+        # so we save it and evolve it separately.
         v_mean = jnp.mean(v, axis=(0, 1), keepdims=True)  # (1, 1, 2)
 
-        # velocity -> vorticity: (1, N, N)  (operate on fluctuation only)
-        vort = _velocity_to_vorticity(v - v_mean, domain_extent)
+        # velocity -> Exponax's q = -ω: (1, N, N)  (operate on fluctuation only)
+        q = -_velocity_to_vorticity(v - v_mean, domain_extent)
 
         if kolmogorov_forcing:
             stepper = ex.stepper.KolmogorovFlowVorticity(
@@ -169,6 +187,18 @@ def exponax_fwd(
                 injection_scale=injection_scale,
                 **stepper_kwargs,
             )
+        elif forcing is not None:
+            stepper = ex.stepper.KolmogorovFlowVorticity(
+                num_spatial_dims=2, injection_scale=0.0, **stepper_kwargs
+            )
+            # The force drives q with -curl f. Exponax builds its injection as
+            # a real array; the force's coefficients need their phase.
+            curl = _velocity_to_vorticity(forcing[:, :, 0, :], domain_extent)
+            stepper = eqx.tree_at(
+                lambda st: st._integrator._nonlinear_fun.injection,
+                stepper,
+                jnp.fft.rfft2(-curl).astype(jnp.complex64),
+            )
         else:
             stepper = ex.stepper.NavierStokesVorticity(
                 num_spatial_dims=2, **stepper_kwargs
@@ -177,10 +207,25 @@ def exponax_fwd(
         def step_fn(carry: Any, _: Any) -> tuple[Any, None]:
             return stepper(carry).astype(carry.dtype), None
 
-        result_vort, _ = jax.lax.scan(step_fn, vort, None, length=steps)
+        result_q, _ = jax.lax.scan(step_fn, q, None, length=steps)
 
-        # vorticity -> velocity: (N, N, 2) -> (N, N, 1, 2); restore conserved mean
-        result = _vorticity_to_velocity(result_vort, domain_extent) + v_mean
+        if forcing is not None or linear_drag is not None:
+            # dū/dt = λū + f̄ under Exponax's drag λ and the mean force f̄.
+            rate = stepper_kwargs["drag"]
+            t_end = dt * steps
+            f_mean = (
+                0.0
+                if forcing is None
+                else jnp.mean(forcing[:, :, 0, :], axis=(0, 1), keepdims=True)
+            )
+            growth = jnp.exp(rate * t_end)
+            # (e^{λT} - 1)/λ, which tends to T as λ → 0 (rate may be traced).
+            safe = jnp.where(rate == 0, 1.0, rate)
+            gain = jnp.where(rate == 0, t_end, (growth - 1.0) / safe)
+            v_mean = v_mean * growth + f_mean * gain
+
+        # q -> velocity: (N, N, 2) -> (N, N, 1, 2); restore the mean
+        result = _vorticity_to_velocity(-result_q, domain_extent) + v_mean
         return result[:, :, None, :]
 
     else:  # 3D
@@ -219,13 +264,13 @@ def apply_jit(inputs: dict) -> dict:
     return {"result": exponax_fwd(**inputs)}
 
 
-_SCALAR_KEYS = ("dt", "viscosity", "drag", "injection_scale")
+_SCALAR_KEYS = ("dt", "viscosity", "drag", "injection_scale", "linear_drag")
 
 
 def _unpack_scalars(d: dict) -> dict:
     """Extract Python floats from 1-element arrays for JIT-static scalar params."""
     for key in _SCALAR_KEYS:
-        if key in d:
+        if d.get(key) is not None:
             d[key] = float(d[key][0])
     return d
 
@@ -244,7 +289,8 @@ def _diff_apply(inputs: dict) -> dict:
     """Forward pass differentiated by jax_vjp, with the scalar params traced."""
     fwd_kwargs = dict(inputs)
     for k in _SCALAR_KEYS:
-        fwd_kwargs[k] = jnp.reshape(fwd_kwargs[k], ())
+        if fwd_kwargs.get(k) is not None:
+            fwd_kwargs[k] = jnp.reshape(fwd_kwargs[k], ())
     return {"result": exponax_fwd(**fwd_kwargs)}
 
 

mosaic/tesseracts/navier-stokes-grid/ins-jl/solver.jl

@@ -209,6 +209,43 @@ function ns_forward_3d(v0::AbstractArray, rhs, setup, psolver,
 end
 
 
+# ---------------------------------------------------------------------------
+# Body force and linear drag (2-D)
+# ---------------------------------------------------------------------------
+
+"""
+Collocated body force (n,n,2) → projected, ghosted staggered force, or `nothing`.
+
+The force is averaged onto the faces like v0 and projected once, so that
+adding it to the projected right-hand side keeps the stages divergence-free.
+"""
+function staggered_force(force, setup, psolver, n::Int)
+    isnothing(force) && return nothing
+    f = add_ghosts_2d(coloc_to_stag_2d(to_device(Float32.(force)), n), n)
+    return project(f, setup; psolver)
+end
+
+"""
+`rhs` plus a constant body force `f` and a linear drag `-drag * u`.
+
+The drag acts on the stage velocity, which is divergence-free, so the sum
+stays projected. Returns `rhs` itself when there is neither.
+"""
+function forced_rhs(rhs, f, drag::Real)
+    isnothing(f) && drag == 0 && return rhs
+    return function (u, p, t)
+        du = rhs(u, p, t)
+        if !isnothing(f)
+            du = du .+ f
+        end
+        if drag != 0
+            du = du .- eltype(u)(drag) .* u
+        end
+        return du
+    end
+end
+
+
 # ---------------------------------------------------------------------------
 # Public API (called from Python via juliacall)
 # ---------------------------------------------------------------------------
@@ -226,12 +263,24 @@ used for the VJP primal. `RKMethods.RK44` applies the same per-stage BC +
 projection as the projected RHS in `ns_forward_*`, so the two paths agree to
 round-off.
 """
-function ns_apply(v0_np, nu::Float64, dt::Float64, steps::Int, n::Int, L::Float64)
+function ns_apply(v0_np, nu::Float64, dt::Float64, steps::Int, n::Int, L::Float64;
+                  force = nothing, drag::Float64 = 0.0)
     T = Float32
     v0   = to_device(T.(v0_np))
     ndim = size(v0, ndims(v0))  # last dim: 2 or 3
     setup, psolver = get_setup_and_psolver(n, L, ndim)
 
+    if !isnothing(force) || drag != 0
+        # A body force or drag runs the same RK4 loop as the VJP, so that both
+        # integrate the same right-hand side. 2-D only.
+        rhs = forced_rhs(
+            create_right_hand_side(setup, psolver),
+            staggered_force(force, setup, psolver, n),
+            drag,
+        )
+        return Array(T.(ns_forward_2d(v0, rhs, setup, psolver, T(nu), T(dt), steps, n)))
+    end
+
     u = initial_state(v0, setup, psolver, n)
 
     state, _ = solve_unsteady(;
@@ -265,13 +314,18 @@ exact), which removed the need for the finite-difference fallback (and its
 two extra forward rollouts per VJP call) that earlier versions required.
 """
 function ns_vjp(v0_np, cotangent_np, nu::Float64, dt::Float64,
-                steps::Int, n::Int, L::Float64)
+                steps::Int, n::Int, L::Float64;
+                force = nothing, drag::Float64 = 0.0)
     v0   = to_device(Float32.(v0_np))
     cot  = to_device(Float32.(cotangent_np))
     ndim = size(v0, ndims(v0))
 
     setup, psolver = get_setup_and_psolver(n, L, ndim)
-    rhs  = create_right_hand_side(setup, psolver)
+    rhs  = forced_rhs(
+        create_right_hand_side(setup, psolver),
+        staggered_force(force, setup, psolver, n),
+        drag,
+    )
 
     if ndim == 2
         fwd = (v, dt_, nu_) -> ns_forward_2d(v, rhs, setup, psolver, nu_, dt_, steps, n)

mosaic/tesseracts/navier-stokes-grid/ins-jl/tesseract_api.py

@@ -57,12 +57,26 @@ def _to_numpy(jl_arr: Any) -> np.ndarray:
     return np.asarray(jl_arr).copy()
 
 
+def _body_force(inputs: InputSchema) -> dict[str, Any]:
+    """Keyword arguments for the body force and drag, empty when there are none.
+
+    Both are 2-D only (the canonical schema rejects them in 3-D).
+    """
+    kwargs: dict[str, Any] = {}
+    if inputs.forcing is not None:
+        kwargs["force"] = _to_julia(np.asarray(inputs.forcing)[:, :, 0, :])
+    if inputs.linear_drag is not None:
+        kwargs["drag"] = float(inputs.linear_drag[0])
+    return kwargs
+
+
 def _run(
     v0: np.ndarray,
     viscosity: float,
     dt: float,
     steps: int,
     domain_extent: float,
+    **body_force: Any,
 ) -> np.ndarray:
     """Forward pass.
 
@@ -84,6 +98,7 @@ def _run(
                 int(steps),
                 int(nx),
                 float(domain_extent),
+                **body_force,
             )
         )  # (nx, ny, 2)
         return result[:, :, np.newaxis, :]  # (nx, ny, 1, 2)
@@ -108,6 +123,7 @@ def _vjp(
     dt: float,
     steps: int,
     domain_extent: float,
+    **body_force: Any,
 ) -> dict:
     """VJP: cotangent w.r.t. result -> gradients w.r.t. v0, viscosity, dt, domain_extent."""
     nx = v0.shape[0]
@@ -125,6 +141,7 @@ def _vjp(
             int(steps),
             int(nx),
             float(domain_extent),
+            **body_force,
         )
         grad_v0_2d = _to_numpy(result[0])
         grad_nu = float(result[1])
@@ -171,6 +188,7 @@ def apply(inputs: InputSchema) -> OutputSchema:
         float(inputs.dt[0]),
         inputs.steps,
         inputs.domain_extent,
+        **_body_force(inputs),
     )
     return {"result": result.astype(np.float32)}
 
@@ -193,6 +211,7 @@ def vector_jacobian_product(
         float(inputs.dt[0]),
         inputs.steps,
         inputs.domain_extent,
+        **_body_force(inputs),
     )
     return {key: grads[key] for key in vjp_inputs if key in grads}