Finite volumes

import jax
import jax.numpy as jnp
from scimba_jax.linear_approximation.finite_volume import (
    EllipticFVScheme,
    TPFADiffusionFlux,
)
from scimba_jax.linear_approximation.meshes.cartesian_mesh import cartesian_mesh
from scimba_jax.linear_approximation.variables.variables_fv import VariablesFV
from scimba_jax.nonlinear_approximation.model_class.funcparam_matrix import (
    ParamMatrixFunction,
)
from scimba_jax.nonlinear_approximation.model_class.funcparam_scalar import (
    ParamScalarFunction,
)
from scimba_jax.physical_models.abstract_conservative_pde import (
    AbstractConservativePDE,
)
from scimba_jax.physical_models.abstract_physical_conservative_model import (
    AbstractPhysicalConservativeModel,
)


class Poisson2D(AbstractConservativePDE):
    """``-div(grad u) = f``, written as a CONSERVATION LAW: div(flux) = source."""

    def __init__(self):
        super().__init__(dim=2)

    def construct_A(self, u):  # the diffusion tensor multiplying grad(u)
        return ParamMatrixFunction(u.dims, lambda _sp, _x: jnp.eye(2), f_type=u.f_type)

    def construct_f(self):  # the source, one value per cell
        return ParamScalarFunction(
            {"x": 2, "mu": 0},
            lambda _sp, x: 2.0 * (x[0] * (1.0 - x[0]) + x[1] * (1.0 - x[1])),
            f_type="x",
        )

    def construct_F(self, u):
        return None  # no first-order (advective) flux here

    def construct_R(self, u):
        return None  # no reaction term here


mesh = cartesian_mesh([32, 32], quad_order=3)
model = AbstractPhysicalConservativeModel.from_pde(Poisson2D(), dirichlet=lambda x: jnp.zeros(1))
# TPFA: flux through a face = (u_right - u_left) / distance, times A.
scheme = EllipticFVScheme(
    model, VariablesFV(mesh, projection="average"), TPFADiffusionFlux(),
    assemble_first_order=False, assemble_reaction=False,
)
scheme = EllipticFVScheme.solve(scheme)
# RMS error against x(1-x)y(1-y): 3.7e-05 at 32x32 cells.

class AbstractFVFlux(ScimbaPytree):
    """Every finite-volume flux answers up to four questions on a face."""

    def first_order(self, model, u_left, u_right, context):
        """Advective flux density leaving the left cell (a Riemann solver)."""
        return jnp.zeros_like(u_left)

    def second_order(self, model, u_left, u_right, context):
        """Diffusive flux density leaving the left cell (e.g. two-point flux)."""
        return jnp.zeros_like(u_left)

    def boundary_first_order(self, model, u_inside, boundary_condition, context):
        """Same, on a boundary face -- context.point replaces u_right."""
        return jnp.zeros_like(u_inside)

    def boundary_second_order(self, model, u_inside, boundary_condition, context):
        """Diffusive counterpart of boundary_first_order."""
        return jnp.zeros_like(u_inside)

# TPFADiffusionFlux only overrides the diffusive pair. A hyperbolic law is
# the mirror image: RusanovFlux overrides only first_order, a local
# Lax-Friedrichs flux built from a spectral-radius bound on the wave speed.
from scimba_jax.linear_approximation.finite_volume.hyperbolic import RusanovFlux

flux = RusanovFlux(wave_speed=lambda _model, u, _context: jnp.abs(u[0]))
# {F(uL)+F(uR)}/2 - alpha/2 (uR - uL), alpha = max wave speed on the face.

from scimba_jax.linear_approximation.finite_volume import (
    FiniteVolumeScheme,
    TimeDiscreteFVscheme,
)
from scimba_jax.linear_approximation.finite_volume.hyperbolic import RusanovFlux
from scimba_jax.linear_approximation.meshes.cartesian_mesh import cartesian_mesh
from scimba_jax.linear_approximation.variables.variables_fv import VariablesFV
from scimba_jax.nonlinear_approximation.model_class.funcparam_scalar import (
    ParamScalarFunction,
)
from scimba_jax.physical_models.abstract_conservative_pde import AbstractConservativePDE
from scimba_jax.physical_models.abstract_physical_conservative_model import (
    AbstractPhysicalConservativeModel,
)
from scimba_jax.physical_models.weak_boundary_conditions import Dirichlet
from scimba_jax.time_discrete.butcher_tableau import build_explicit_euler_tableau

# Burgers: du/dt + d(u^2/2)/dx = 0, a Riemann problem u_L=1, u_R=-0.25 that
# forms a shock. FiniteVolumeScheme becomes a right-hand side to STEP,
# the same way EllipticFVScheme became a residual to SOLVE.
class BurgersPDE(AbstractConservativePDE):
    def __init__(self):
        super().__init__(dim=1)

    def construct_F(self, u):
        return ParamScalarFunction(u.dims, lambda sp, x: 0.5 * u(sp, x) ** 2, f_type=u.f_type)

    def construct_A(self, u):
        return None

    def construct_R(self, u):
        return None


model = AbstractPhysicalConservativeModel(BurgersPDE())
model.add_boundary_condition("west", Dirichlet(lambda _x: jnp.array([1.0])))
model.add_boundary_condition("east", Dirichlet(lambda _x: jnp.array([-0.25])))
flux = RusanovFlux(wave_speed=lambda _model, u, _context: jnp.abs(u[0]))
spatial_scheme = FiniteVolumeScheme(
    model, VariablesFV(cartesian_mesh([40], quad_order=2)), flux,
    assemble_second_order=False, assemble_reaction=False, assemble_source=False,
)
scheme = TimeDiscreteFVscheme(spatial_scheme, build_explicit_euler_tableau(), dt=0.4 / 40)
initial = scheme.initialize(lambda x: jnp.where(x[0] < 0.5, jnp.array([1.0]), jnp.array([-0.25])))
final, history = scheme.solve(initial, t0=0.0, nt=20)
# Explicit Euler needs a CFL-limited dt (dt <= CFL h / max wave speed); a
# shock has no smooth stiffness for an implicit step to buy back. 40 cells,
# CFL=0.4 to t=0.2: mass conserved to 0.46875 (exact), no overshoot.