aragog.jax

def compute_fluxes(
    S_stag: jax.Array,
    time: float,
    eos: EntropyEOS_JAX,
    params: PhaseParams,
    mesh: MeshArrays,
    heating_rate: jax.Array,
    S_basic_cmb_override=None,
    dSdr_cmb_override=None,
) -> FluxOutput:
    """Compute all heat and mass fluxes from the entropy profile.

    This is the physics kernel called by the ODE RHS. It is a pure
    function: no mutation, no side effects, fully JIT-compilable.

    Parameters
    ----------
    S_stag : jax.Array
        Entropy at staggered nodes [J/kg/K].
    time : float
        Current time [yr] (used for radionuclide heating).
    eos : EntropyEOS_JAX
        JAX EOS tables.
    params : PhaseParams
        Static material parameters and transport flags.
    mesh : MeshArrays
        Mesh geometry and transform matrices.
    heating_rate : jax.Array
        Internal heating rate at staggered nodes [W/kg]
        (radionuclide + tidal, pre-computed by caller).
    S_basic_cmb_override : float or None
        Optional override for the entropy at the CMB basic node
        (basic-node index 0). Used by the energy_balance core BC
        which reconstructs S_basic[0] from the state-tracked
        dSdr_cmb via S[0] + dSdr_cmb * (r_basic[0] - r_stag[0]).
        When None, the standard quantity_matrix mapping is used.
    dSdr_cmb_override : float or None
        Optional override for the entropy gradient at the CMB
        basic node (dSdr[0]). Used by the energy_balance core BC
        where dSdr_cmb is a state-tracked variable. When None,
        the standard d_dr_matrix mapping is used.

    Returns
    -------
    FluxOutput
        Heat flux, mass flux, eddy diffusivity, heating.
    """
    # Interpolate entropy to basic nodes and compute gradient
    S_basic = mesh.quantity_matrix @ S_stag
    dSdr = mesh.d_dr_matrix @ S_stag

    # SPIDER ic.c:450 boundary-copy convention for dSdr at the surface
    # (mirrors numpy entropy_state.py:390 ``dSdxi[-1] = dSdxi[-2]``).
    # The d_dr_matrix gives a linear-extrapolation gradient at the top
    # basic node, but the SPIDER convention is to copy the adjacent
    # interior value. Without this, the surface dSdr can flip sign
    # relative to numpy, blowing up kappa_h and dS/dt at the second-
    # to-last staggered cell. The CMB side has its own override path
    # via ``dSdr_cmb_override``, so the surface copy here is the only
    # one needed at this layer.
    dSdr = dSdr.at[-1].set(dSdr[-2])

    # Apply energy_balance overrides at the CMB basic node.
    # Done as separate jax.lax.cond branches so the function remains
    # JIT-compatible regardless of whether the overrides are None or
    # JAX scalars. We use jnp.where with a static bool flag passed
    # through from the caller via the optional argument.
    if S_basic_cmb_override is not None:
        S_basic = S_basic.at[0].set(S_basic_cmb_override)
    if dSdr_cmb_override is not None:
        dSdr = dSdr.at[0].set(dSdr_cmb_override)
    else:
        # Mirror numpy entropy_state.py:389 ``dSdxi[0] = dSdxi[1]`` for
        # the non-energy_balance modes (quasi_steady etc.) where there
        # is no boundary-state override. In energy_balance mode the
        # explicit override above wins and this branch is skipped.
        dSdr = dSdr.at[0].set(dSdr[1])

    # Phase properties at basic nodes only. The SPIDER-bracket Jmix
    # is built entirely from basic-node quantities, so staggered-node
    # phase properties are not materialised here.
    phase_basic = evaluate_phase(eos, params, mesh.P_basic, S_basic)

    # MLT eddy diffusivity
    kappa_h, kappa_c = compute_mlt(dSdr, phase_basic, mesh, params)

    # Basic node properties
    rho = phase_basic.density
    T = phase_basic.temperature
    Cp = phase_basic.heat_capacity
    k = phase_basic.thermal_conductivity
    dTdPs_basic = phase_basic.dTdPs

    # Heat flux components (multiply by flag to enable/disable)
    heat_flux = jnp.zeros_like(S_basic)

    # Conduction (SPIDER decomposition matching numpy entropy_state):
    #   F_cond = -k * [(T/Cp) * dS/dr + dT/dr|_adiabat]
    # where the adiabatic gradient is dT/dr|_ad = dTdPs * dPdr_basic
    # (EOS-table dT/dP|_S times the structural Adams-Williamson dP/dr).
    # This avoids the noise from finite-differencing T_stag and matches
    # SPIDER's eos-consistent conduction at phase boundaries.
    # Cp floor (100 J/kg/K) parity with numpy entropy_state.update. A
    # logger.warning inside this jit-compiled RHS would only fire
    # during tracing, not on actual data; the equivalent diagnostic
    # is emitted at load time by EntropySolver._check_eos_floors,
    # plus a once-per-instance runtime warning on the numpy path.
    Cp_safe = jnp.maximum(Cp, 100.0)
    superadiabatic = (T / Cp_safe) * dSdr
    dT_dr_adiabat = dTdPs_basic * mesh.dP_dr_basic
    heat_flux = heat_flux + params.conduction * (-k * (superadiabatic + dT_dr_adiabat))

    # Convection: F_conv = rho * T * kappa_h * (-dS/dr)
    heat_flux = heat_flux + params.convection * (rho * T * kappa_h * (-dSdr))

    # Mass flux for gravitational separation and mixing
    mass_flux = jnp.zeros_like(S_basic)

    # Gravitational separation.
    #
    # Raw Stokes/permeability-driven mass flux:
    #     jgrav_raw = rho * phi * (1 - phi) * v_rel
    # SPIDER analogue smoothing (SPIDER/energy.c:523-533,
    # JGRAV_BOTTOM_UP + get_smoothing): multiply jgrav_raw by a bounded
    # polynomial of an UN-truncated two-phase fraction
    #     gphi = (S - S_sol(P)) / (S_liq(P) - S_sol(P))
    # evaluated at the staggered cell immediately BELOW the interface.
    # The polynomial `16 * gphi^2 * (1 - gphi)^2` (clipped to [0, 1])
    # vanishes cleanly at both pure phases and has bounded derivatives
    # everywhere, unlike SPIDER's tanh smoothing. This is the scipy-
    # path fix from entropy_state.py mirrored here so the JAX backend
    # doesn't reproduce the pre-fix CMB drain at first crystallisation.
    phi_b = phase_basic.melt_fraction
    v_rel = relative_velocity(
        eos,
        params,
        mesh.P_basic,
        rho,
        phi_b,
        mesh.gravity,
    )
    jgrav_raw = rho * phi_b * (1.0 - phi_b) * v_rel

    # gphi at STAGGERED nodes (cell below each basic interface)
    S_sol_stag = eos.solidus_entropy(mesh.P_stag)
    S_liq_stag = eos.liquidus_entropy(mesh.P_stag)
    dS_stag = jnp.maximum(S_liq_stag - S_sol_stag, 1.0)
    gphi_stag = (S_stag - S_sol_stag) / dS_stag

    smth_stag = phase_boundary_smoothing(gphi_stag, params)

    # Map staggered smoothing to basic-node interfaces: interior basic
    # node i (1..N-2) sees the smoothing of staggered node i-1 (the
    # cell BELOW). Boundary basic nodes (0 and -1) use smth = 1 as a
    # placeholder because the mass flux at those indices is zeroed a
    # few lines below anyway. Lengths: staggered has N entries, basic
    # has N+1; smth_stag[:-1] supplies the N-1 interior interfaces
    # plus the two boundaries, totalling N+1.
    smth_basic = jnp.concatenate(
        [
            jnp.array([1.0]),
            smth_stag[:-1],
            jnp.array([1.0]),
        ]
    )

    # `bottom_up_grav_sep = 1.0` selects the smoothed flux, 0.0 selects
    # the raw flux (for reproducing the pre-fix drain in regression
    # tests).
    jgrav_smoothed = jgrav_raw * smth_basic
    jgrav = (
        params.bottom_up_grav_sep * jgrav_smoothed
        + (1.0 - params.bottom_up_grav_sep) * jgrav_raw
    )
    mass_flux = mass_flux + params.grav_sep * jgrav

    # Zero mass fluxes at boundaries (SPIDER convention)
    mass_flux = mass_flux.at[0].set(0.0)
    mass_flux = mass_flux.at[-1].set(0.0)

    # Add latent heat transport from mass flux
    heat_flux = heat_flux + mass_flux * phase_basic.latent_heat

    # Mixing flux (SPIDER-parity bracket form, heat flux only).
    #
    # Mirrors numpy ``entropy_state.update`` (entropy_state.py:656-692):
    #     Jmix_heat = -kappa_c * rho * T_fus * bracket * smth_basic_mix
    #     bracket = dS/dr - [phi·dS_liq/dP + (1-phi)·dS_sol/dP] · dP/dr
    # evaluated at basic nodes. The mass-flux form
    # ``mass_flux += mixing * rho * kappa_c * (-dphi/dr)`` (delivered
    # via the latent-heat term) is NOT equivalent and diverges from
    # the numpy production path.
    #
    # The smth polynomial ``16·gphi²·(1-gphi)²`` at basic nodes zeroes
    # the flux outside the mushy band; the bracket itself is finite in
    # pure phases because dS_sol/dP and dS_liq/dP are bounded.
    S_sol_basic = eos.solidus_entropy(mesh.P_basic)
    S_liq_basic = eos.liquidus_entropy(mesh.P_basic)
    dS_phase_basic = jnp.maximum(S_liq_basic - S_sol_basic, 1.0)
    dS_sol_dP_basic = eos.solidus_entropy_dP(mesh.P_basic)
    dS_liq_dP_basic = eos.liquidus_entropy_dP(mesh.P_basic)
    phi_clipped = phase_basic.melt_fraction
    bracket = (
        dSdr
        - (phi_clipped * dS_liq_dP_basic + (1.0 - phi_clipped) * dS_sol_dP_basic)
        * mesh.dP_dr_basic
    )

    # T_fus = latent_heat / dS_phase (mirrors numpy
    # _ensure_basic_phase_boundary_cache: T_fus_basic = L_basic / dS_phase).
    T_fus_basic = phase_basic.latent_heat / dS_phase_basic

    # Smoothing: gphi at basic nodes, same smoothing family as Jgrav
    # (cubic Hermite or SPIDER tanh, selected by params.phase_smoothing_tanh).
    gphi_basic = (S_basic - S_sol_basic) / dS_phase_basic
    smth_basic_mix = phase_boundary_smoothing(gphi_basic, params)

    jmix_heat = -kappa_c * rho * T_fus_basic * bracket * smth_basic_mix
    # Zero at CMB and surface (no mass/heat transfer across those boundaries)
    jmix_heat = jmix_heat.at[0].set(0.0)
    jmix_heat = jmix_heat.at[-1].set(0.0)
    heat_flux = heat_flux + params.mixing * jmix_heat

    return FluxOutput(
        heat_flux=heat_flux,
        mass_flux=mass_flux,
        eddy_diffusivity=kappa_h,
        heating=heating_rate,
        jmix_heat=jmix_heat,
    )

def compute_mlt(
    dSdr: jax.Array,
    phase_basic: PhaseProperties,
    mesh: MeshArrays,
    params: PhaseParams,
) -> tuple[jax.Array, jax.Array]:
    """Compute MLT eddy diffusivity from the entropy gradient.

    Parameters
    ----------
    dSdr : jax.Array
        Entropy gradient at basic nodes [J/kg/K/m].
    phase_basic : PhaseProperties
        Material properties at basic nodes.
    mesh : MeshArrays
        Mesh geometry.
    params : PhaseParams
        Static parameters.

    Returns
    -------
    kappa_h : jax.Array
        Thermal eddy diffusivity at basic nodes [m^2/s].
    kappa_c : jax.Array
        Chemical eddy diffusivity at basic nodes [m^2/s].
    """
    alpha = phase_basic.thermal_expansivity
    T = phase_basic.temperature
    Cp = phase_basic.heat_capacity
    nu = phase_basic.kinematic_viscosity

    # Buoyancy from entropy gradient. ``jnp.abs`` has a non-differentiable
    # kink at zero (subgradient anywhere in [-1, +1] is mathematically
    # valid, but JAX returns sign(0) = 0 in fwd and the backward pass
    # through the abs and the downstream multiplications can produce
    # NaN at exact-zero entropy gradients. Numpy uses
    # ``_smooth_abs_neg(dSdr, eps=1e-30)`` which is the smoothed
    # ``max(-x, 0)``; we use the differentiable
    # ``0.5*(|x| + sqrt(x^2 + eps^2)) ~ |x|`` instead, which keeps the
    # forward equal to ``abs(x)`` to ULP precision while delivering a
    # finite analytic gradient everywhere.
    eps_abs = 1.0e-30
    abs_dSdr_safe = 0.5 * (jnp.abs(dSdr) + jnp.sqrt(dSdr * dSdr + eps_abs * eps_abs))
    effective_superadiabatic = alpha * T * abs_dSdr_safe / jnp.maximum(Cp, 1.0)
    velocity_prefactor = mesh.gravity * effective_superadiabatic

    # Convective mask: unstable when dS/dr < 0.
    # Hard mask matching the numpy reference. Using jnp.where (not boolean
    # indexing) for JAX traceability. Produces exactly 0 at stable nodes,
    # avoiding spurious convection from a soft sigmoid.
    conv_mask = jnp.where(dSdr < 0.0, 1.0, 0.0)

    # Viscous velocity (Re <= Re_crit)
    viscous_velocity = (velocity_prefactor * mesh.mixing_length_cu / (18.0 * nu)) * conv_mask

    # Inviscid velocity (Re > Re_crit). ``jnp.sqrt(jnp.maximum(x, 0))`` is
    # the textbook NaN-safe forward, but its backward gradient at x=0 is
    # ``0.5 * 0 / sqrt(0) = 0/0 = NaN`` (the maximum's stop-at-zero
    # subgradient kills the divisor protection). Numpy avoids this with
    # ``np.sqrt(x + 1e-20)`` (always-positive argument). Mirror it here:
    # eps**2 = 1e-40 is far below any physical inviscid_velocity_sq, so
    # the forward is bit-equivalent for non-trivial x.
    eps_sqrt = 1.0e-20
    inviscid_velocity_sq = (velocity_prefactor * mesh.mixing_length_sq / 16.0) * conv_mask
    inviscid_velocity = jnp.sqrt(inviscid_velocity_sq + eps_sqrt)

    # Reynolds number
    reynolds = viscous_velocity * mesh.mixing_length / nu

    # Smooth blend between viscous and inviscid regimes. The narrow
    # blend_width (0.01 * RE_CRIT) keeps inviscid k_h confined to the
    # convecting regime; widening the blend leaks inviscid mixing
    # into the solid regime and induces T_core bistability.
    blend_width = 0.01 * RE_CRIT
    inviscid_weight = 0.5 * (
        1.0 + jnp.tanh((reynolds - RE_CRIT) / jnp.maximum(blend_width, 1e-30))
    )

    # Raw eddy diffusivity
    kh_raw = (
        (1.0 - inviscid_weight) * viscous_velocity + inviscid_weight * inviscid_velocity
    ) * mesh.mixing_length

    # Thermal eddy diffusivity (SPIDER convention: positive=scale, negative=constant)
    kappa_h = jnp.where(
        params.eddy_diff_thermal > 0,
        params.eddy_diff_thermal * kh_raw,
        jnp.full_like(kh_raw, -params.eddy_diff_thermal),
    )

    # Chemical eddy diffusivity (from raw kh, before floor)
    kappa_c = jnp.where(
        params.eddy_diff_chemical > 0,
        params.eddy_diff_chemical * kh_raw,
        jnp.full_like(kh_raw, -params.eddy_diff_chemical),
    )

    # kappa_h floor (phase-dependent, modulated by melt fraction).
    # Production PROTEUS runs use kappah_floor = 10 m^2/s; the
    # phi-modulated f_floor ramps from 0 in solid layers (no spurious
    # convective flux) to ~1 in mushy/liquid layers, where physical
    # convection is expected and MLT can otherwise numerically freeze.
    # See solver/entropy_state.py for the same comment block on the
    # numpy path. The transition is anchored on the rheological critical
    # melt fraction (``params.phi_rheo``, default 0.4) with width
    # ``params.phi_width`` (default 0.15) so the floor turns on exactly
    # where Costa-blended viscosity drops, consistent across config knobs.
    phi_basic = phase_basic.melt_fraction
    f_floor = tanh_weight(phi_basic, params.phi_rheo, params.phi_width)
    kh_floor = params.kappah_floor * f_floor
    kappa_h = jnp.maximum(kappa_h, kh_floor)

    # SPIDER energy.c:220-223 CMB fix: use kappa_h from the first interior
    # node at the CMB basic node, since kappa_h is a nonlinear function of
    # dSdr and the boundary extrapolation can over- or under-estimate it
    # relative to the interior value. Mirrors numpy entropy_state.py:533.
    kappa_h = kappa_h.at[0].set(kappa_h[1])

    return kappa_h, kappa_c

def evaluate_phase(
    eos: EntropyEOS_JAX,
    params: PhaseParams,
    P: jax.Array,
    S: jax.Array,
) -> PhaseProperties:
    """Compute all material properties at (P, S) nodes.

    Parameters
    ----------
    eos : EntropyEOS_JAX
        JAX EOS tables.
    params : PhaseParams
        Static material parameters.
    P : jax.Array
        Pressure [Pa], 1D.
    S : jax.Array
        Entropy [J/kg/K], 1D.

    Returns
    -------
    PhaseProperties
        All material properties at the given nodes.
    """
    # SPIDER-parity single-pass evaluation: T, rho, Cp, alpha, dTdPs, k
    # all derived from one shared (S_sol, S_liq, gphi, smth, phase-boundary)
    # cache, with the smth-blend mixed<->single matching numpy
    # EntropyPhaseEvaluator._update_eos.
    state = eos.compute_phase_state(
        P,
        S,
        k_solid=params.k_solid,
        k_liquid=params.k_liquid,
        matprop_smooth_width=params.matprop_smooth_width,
    )
    phi = state.melt_fraction

    # Viscosity: two-stage blend mirroring numpy entropy_phase.py:311-327
    # (used downstream by MLT -> kappa_c -> Jmix, which is why both stages
    # matter). Stage 1: tanh blend at phi_rheo (SPIDER util.c:255-259).
    # Stage 2: combine_matprop with the cached matprop_smooth_width smth
    # that compute_phase_state also uses for T/rho/Cp/alpha/k.
    w = tanh_weight(phi, params.phi_rheo, params.phi_width)
    log_visc_mixed = (1.0 - w) * params.log10_visc_solid + w * params.log10_visc_liquid
    log_visc_single = jnp.where(
        phi > 0.5,
        params.log10_visc_liquid,
        params.log10_visc_solid,
    )
    log_visc = state.smth * log_visc_mixed + (1.0 - state.smth) * log_visc_single
    viscosity = 10.0**log_visc
    kinematic_viscosity = viscosity / state.density

    # Capacitance for entropy equation
    capacitance = state.density * state.temperature

    return PhaseProperties(
        temperature=state.temperature,
        density=state.density,
        heat_capacity=state.heat_capacity,
        thermal_expansivity=state.thermal_expansivity,
        dTdPs=state.dTdPs,
        melt_fraction=phi,
        viscosity=viscosity,
        kinematic_viscosity=kinematic_viscosity,
        thermal_conductivity=state.thermal_conductivity,
        latent_heat=state.latent_heat,
        capacitance=capacitance,
    )