aragog.solver

def apply_flux_boundary_conditions(self, state: State) -> None:
    """Applies the boundary conditions to the state.

    Args:
        state: The state to apply the boundary conditions to
    """
    self.apply_flux_inner_boundary_condition(state)
    self.apply_flux_outer_boundary_condition(state)
    logger.debug('temperature = %s', state.temperature_basic)
    logger.debug('heat_flux = %s', state.heat_flux)

def apply_flux_inner_boundary_condition(self, state: State) -> None:
    """Applies the flux boundary condition to the state at the inner boundary.

    Args:
        state: The state to apply the boundary conditions to

    Equivalent to CORE_BC in C code.
        1: Core cooling (EntropySolver only; raises here)
        2: Prescribed heat flux
        3: Prescribed temperature
    """
    if self._settings.inner_boundary_condition == 1:
        msg = 'inner_boundary_condition = 1 (core cooling) is applied by EntropySolver'
        raise NotImplementedError(msg)
    elif self._settings.inner_boundary_condition == 2:
        state.heat_flux[0, :] = self._settings.inner_boundary_value
    elif self._settings.inner_boundary_condition == 3:
        pass
        # raise NotImplementedError
    else:
        msg: str = f'inner_boundary_condition = {self._settings.inner_boundary_condition} is unknown'
        raise ValueError(msg)

def apply_flux_outer_boundary_condition(self, state: State) -> None:
    """Applies the flux boundary condition to the state at the outer boundary.

    Args:
        state: The state to apply the boundary conditions to

    Equivalent to SURFACE_BC in C code.
        1: Grey-body atmosphere
        2: Zahnle steam atmosphere (not implemented)
        4: Prescribed surface heat flux (atmosphere coupling)
        5: Prescribed surface temperature
    """
    if self._settings.outer_boundary_condition == 1:
        self.grey_body(state)
    elif self._settings.outer_boundary_condition == 2:
        raise NotImplementedError
    elif self._settings.outer_boundary_condition == 4:
        state.heat_flux[-1, :] = self._settings.outer_boundary_value
    elif self._settings.outer_boundary_condition == 5:
        pass
    else:
        msg: str = f'outer_boundary_condition = {self._settings.outer_boundary_condition} is unknown'
        raise ValueError(msg)

def apply_temperature_boundary_conditions(
    self, temperature: npt.NDArray, temperature_basic: npt.NDArray, dTdr: npt.NDArray
) -> None:
    """Conforms the temperature and dTdr at the basic nodes to temperature boundary conditions.

    Args:
        temperature: Temperature at the staggered nodes
        temperature_basic: Temperature at the basic nodes
        dTdr: Temperature gradient at the basic nodes
    """
    # Core-mantle boundary
    if self._settings.inner_boundary_condition == 3:
        temperature_basic[0, :] = self._settings.inner_boundary_value
        dTdr[0, :] = (
            2
            * (temperature[0, :] - temperature_basic[0, :])
            / self._mesh.basic.delta_mesh[0]
            * self._mesh.dxidr[0]
        )
    # Surface
    if self._settings.outer_boundary_condition == 5:
        temperature_basic[-1, :] = self._settings.outer_boundary_value
        dTdr[-1, :] = (
            2
            * (temperature_basic[-1, :] - temperature[-1, :])
            / self._mesh.basic.delta_mesh[-1]
            * self._mesh.dxidr[-1]
        )

def apply_temperature_boundary_conditions_melt(
    self, melt_fraction: npt.NDArray, melt_fraction_basic: npt.NDArray, dphidr: npt.NDArray
) -> None:
    """Conforms the melt fraction gradient dphidr at the basic nodes
       to temperature boundary conditions.

    Args:
        melt_fraction: Melt fraction at the staggered nodes
        melt_fraction_basic: Melt fraction at the basic nodes
        dphidr: Melt fraction gradient at the basic nodes
    """
    # Core-mantle boundary
    if self._settings.inner_boundary_condition == 3:
        dphidr[0, :] = (
            2
            * (melt_fraction[0, :] - melt_fraction_basic[0, :])
            / self._mesh.basic.delta_mesh[0]
            * self._mesh.dxidr[0]
        )
    # Surface
    if self._settings.outer_boundary_condition == 5:
        dphidr[-1, :] = (
            2
            * (melt_fraction_basic[-1, :] - melt_fraction[-1, :])
            / self._mesh.basic.delta_mesh[-1]
            * self._mesh.dxidr[-1]
        )

def grey_body(self, state: State) -> None:
    """Applies a grey body flux at the surface.

    When param_utbl is enabled, the surface radiating temperature is reduced
    to account for the ultra-thin thermal boundary layer at the magma ocean
    surface. The temperature drop across this unresolved boundary layer is
    parameterized as dT = b * T_surf^3 (Bower et al. 2018, Eq. 18), giving
    the cubic relation T_interior = T_surf + b * T_surf^3. The analytical
    solution (Cardano's formula) gives T_surf < T_interior.

    Parameters
    ----------
    state : State
        The state to apply the boundary conditions to.
    """
    t_top = state.top_temperature
    if self._settings.param_utbl:
        t_surf = self._utbl_tsurf(t_top)
    else:
        t_surf = t_top
    state.heat_flux[-1, :] = (
        self._settings.emissivity
        * sp_constants.Stefan_Boltzmann
        * (np.power(t_surf, 4) - self._settings.equilibrium_temperature**4)
    )

def dSdt(
    self,
    time: npt.NDArray | float,
    state_vec: npt.NDArray,
) -> npt.NDArray:
    """Time derivative of the full state vector.

    The layout and the length of the state vector depend on the
    core BC mode (``_dSdt_single`` lists the five), and this
    returns the derivative in the same layout.

    The integrator passes ``vectorized=False``; this RHS handles
    the 1D path only.

    Parameters
    ----------
    time : float
        Time [yr].
    state_vec : array
        Solver state vector [J/kg/K for entropy, K for T_core],
        one column of the length of the active core BC mode.

    Returns
    -------
    array
        d(state_vec)/dt with the same shape as the input.
    """
    return self._dSdt_single(time, state_vec)

@classmethod
def from_file(cls, filename: str, eos_dir: str, root: str = '') -> 'EntropySolver':
    """Create EntropySolver from a config file and EOS directory.

    Parameters
    ----------
    filename : str
        Path to TOML configuration file.
    eos_dir : str
        Path to directory with SPIDER-format P-S tables.
    root : str
        Root directory for the config file.
    """
    config_path = Path(root) / Path(filename)
    parameters = Parameters.from_file(config_path)
    entropy_eos = EntropyEOS(Path(eos_dir))
    return cls(parameters, entropy_eos)

def get_current_core_temperature(self) -> float | None:
    """Return the most recent core temperature from the solver state.

    Reads ``self._solution.y[T_core_slot, -1]`` (the final T_core from the
    last accepted solve). Returns ``None`` when no solution exists yet,
    or when the state vector lacks the T_core slot (core_bc other than
    ``bower2018`` or ``core_module``).

    Used by PROTEUS's retry ladder to snapshot the pre-solve
    T_core before a sequence of retry attempts and restore it on
    each retry, mirroring ``get_current_dSdr_cmb``.
    """
    return self._final_extra_state('T_core')

def get_current_dSdr_cmb(self) -> float | None:
    """Return the most recent CMB entropy gradient from the solver state.

    Reads ``self._solution.y[n_stag, -1]`` (the final dSdr_cmb from the
    last accepted solve). Returns ``None`` when no solution exists yet,
    or when the state vector lacks the dSdr_cmb slot (core_bc other
    than ``energy_balance`` or ``core_module``).

    Used by PROTEUS's retry ladder to snapshot the pre-solve
    dSdr_cmb before a sequence of retry attempts and restore it on
    each retry, breaking the positive-feedback loop where a
    rejected attempt's final dSdr_cmb would otherwise become the
    next retry's hot-start IC and drive the boundary state further
    from the pre-solve value each time.
    """
    return self._final_extra_state('dSdr_cmb')

def get_state(self) -> SolverOutput:
    """Extract the solver state as a clean output dataclass.

    This is the primary API for callers to retrieve results. It
    avoids the need to access solver internals (evaluator, mesh,
    state, phase objects).

    Returns
    -------
    SolverOutput
        Dataclass containing all quantities needed by PROTEUS.
    """
    self._sync_top_node()
    sol = self._solution
    eos = self.entropy_eos
    mesh = self.evaluator.mesh

    n_stag = self._n_stag
    energy_balance = self._core_bc == 'energy_balance'
    bower = self._core_bc == 'bower2018'
    gradient_mode = self._core_bc == 'gradient'
    is_ext = energy_balance or bower or self._core_bc == 'core_module'

    # solve() stores the integrals; a result set without solve() gets them here, before
    # the final-state refresh below that the rest of get_state() reads.
    step_integrals = sol.get('energy_integrals') or self._compute_step_energy_integrals()

    # Slice the final state vector.
    if gradient_mode:
        n_basic = n_stag + 1
        dSdr_final = sol.y[:n_basic, -1]
        S_surf_final = float(sol.y[n_basic, -1])
        S_final, _ = self._reconstruct_entropy(dSdr_final, S_surf_final)
    elif is_ext:
        S_final = sol.y[:n_stag, -1]
    else:
        S_final = sol.y[:, -1]

    P_stag = self._P_stag_flat
    r_basic = self._r_basic_flat
    r_stag = np.asarray(mesh.staggered.radii).ravel()
    vol = self._volume_flat

    if eos is not None:
        T_stag = np.asarray(eos.temperature(P_stag, S_final)).ravel()
        phi_stag = np.asarray(eos.melt_fraction(P_stag, S_final)).ravel()
        rho_stag = np.asarray(eos.density(P_stag, S_final)).ravel()
    else:
        # const_properties: analytical T, phi=1, rho=const
        pm = self.parameters.phase_mixed
        T_stag = pm.const_T_ref * np.exp((S_final - pm.const_S_ref) / pm.const_Cp)
        phi_stag = np.ones_like(S_final)
        rho_stag = np.full_like(S_final, pm.const_rho)

    # Refresh the state at the final entropy through the RHS, so the boundary
    # fluxes are the ones the integrator applied; core_module zeroes the gradient
    # slot, so the CMB node is the bottom cell carried to the CMB pressure.
    y_diag = sol.y[:, -1].copy()
    if self._core_bc == 'core_module':
        y_diag[n_stag] = 0.0
    self._dSdt_single(sol.t[-1], y_diag)
    visc_stag = np.asarray(self.state.phase_staggered.viscosity()).ravel()
    heat_flux = self.state.heat_flux.copy()
    heating = self.state.heating.copy()
    eddy_diff = self.state.eddy_diffusivity.copy()
    cap_stag = np.asarray(self.state.capacitance_staggered()).ravel()

    # Diagnostic snapshot from the post-integration state.update().
    jcond_b = self.state.jcond.copy()
    jconv_b = self.state.jconv.copy()
    jgrav_b = self.state.jgrav_heat.copy()
    jmix_b = self.state.jmix_heat.copy()
    dSdr_b = self.state.dSdr.copy()
    phi_basic_diag = self.state.phi_basic_diag.copy()
    T_basic_diag = self.state.T_basic_diag.copy()
    cp_basic_diag = self.state.cp_basic_diag.copy()
    rho_basic_diag = self.state.rho_basic_diag.copy()
    if self._core_bc == 'core_module':
        # Node 0 carries the applied closure: conduction across the thinner of the half
        # cell and the boundary layer.
        jcond_b[0], jconv_b[0], jgrav_b[0], jmix_b[0] = heat_flux[0], 0.0, 0.0, 0.0
    # Phase-boundary quantities at the basic nodes, for consumers that compute melt drainage.
    ph = self.state.phase_basic
    porosity_b, rho_solid_b, rho_melt_b, g_b = (
        np.array(x, dtype=float).ravel()
        for x in (
            ph.porosity(),
            *ph.phase_boundary_densities(),
            ph.gravitational_acceleration(),
        )
    )

    # Scalar quantities. M_mantle is the structural mass between the CMB
    # and the surface from the EOS mass integral, not the PALEOS rho_stag
    # used just below for per-cell output.
    mesh = self.evaluator.mesh
    M_mantle = float(
        np.diff(mesh.eos.get_mass_within_radii(self._r_basic_flat[[0, -1]]))[0]
    )
    mass_stag = rho_stag * vol  # PALEOS density for per-cell output
    # T_magma = top basic-node temperature, evaluated at
    # r = outer_boundary where P = surface_pressure. This matches
    # SPIDER's `atmosphere/temperature_surface` definition.
    # PROTEUS passes T_magma to the atmosphere module as the
    # interior-side surface boundary condition.
    T_basic_final = np.asarray(self.state.phase_basic.temperature()).ravel()
    T_magma = float(T_basic_final[-1])
    # Core temperature: bottom staggered cell (T_stag[0]).
    # SPIDER reports T_core = interior_o.temp[-1], the last
    # staggered node in its surface-to-CMB ordering. Aragog
    # orders CMB-to-surface, so [0] = CMB cell. Reading T_stag[0]
    # rather than evaluating T at the CMB basic node avoids a
    # systematic ~10 K offset from the half-cell pressure
    # difference and matches SPIDER's definition.
    T_core = self._core_temperature_from_column(sol.y[:, -1])
    # Per-solve core-temperature excursion: the largest change of the
    # core temperature from its solve-entry value over the returned
    # grid, with a flag set when it exceeds the configured limit or
    # when any sampled core temperature is non-finite.
    tcore_change_max, tcore_change_exceeded = self._core_temperature_excursion(sol)
    # Mass-weighted melt fraction = M_mantle_liquid / M_mantle.
    # MUST be mass-weighted, not volume-weighted, when
    # ``mass_coordinates = true``: the mesh is uniform in mass
    # coordinate, so deep high-density cells span small radial
    # intervals and have small volumes while surface low-density
    # cells are large. During bottom-up crystallisation the
    # surface stays liquid longest, so volume-weighted Phi stays
    # anchored near the surface value while the actual mantle has
    # crystallised, freezing the helpfile Phi and breaking
    # PROTEUS's stop-criterion + structure-update bookkeeping.
    # Matches the rootfn formula bit-for-bit.
    mass_total_for_phi = float(np.sum(mass_stag))
    if mass_total_for_phi > 0.0:
        Phi_global = float(np.sum(phi_stag * mass_stag) / mass_total_for_phi)
    else:
        Phi_global = float(np.mean(phi_stag))

    # Rheological front depth. ``phi_basic_stag_interp`` is the
    # staggered-phi interpolated to basic nodes -- distinct from the
    # EOS-evaluated ``phi_basic_diag`` captured for the diagnostic
    # output above, which is ``phase_basic.melt_fraction()``.
    phi_rheo = self.parameters.phase_mixed.rheological_transition_melt_fraction
    phi_basic_stag_interp = mesh.quantity_at_basic_nodes(phi_stag).ravel()
    # 0.99 / 0.01 bypass thresholds: short-circuit ``argmin`` when no
    # transition exists in the mantle (essentially-fully-liquid or
    # essentially-fully-solid). These mirror ``rheological_front`` in
    # ``output/diagnostics.py`` (now exposed as ``phi_high`` /
    # ``phi_low`` kwargs there). The ``argmin``-based search assumes
    # bottom-up crystallisation — a middle-out solidification mode
    # would need a redefined RF concept (multiple crossings).
    if Phi_global > 0.99:
        rf = float(r_basic[0])
    elif Phi_global < 0.01:
        rf = float(r_basic[-1])
    else:
        idx = np.argmin(np.abs(phi_basic_stag_interp - phi_rheo))
        rf = float(r_basic[idx])
    R_outer = float(r_basic[-1])
    RF_depth = 1.0 - rf / R_outer if R_outer > 0 else 0.0

    # Thermal energy (sensible, for comparison with SPIDER).
    # Uses the real heat capacity Cp(P, S) from the EntropyEOS
    # phase evaluator so that E_th tracks the EOS internal energy
    # consistently with the solver state, rather than a fixed
    # reference Cp.
    Cp_stag = np.asarray(self.state.phase_staggered.heat_capacity()).ravel()
    E_th = float(np.sum(mass_stag * Cp_stag * T_stag))

    # Effective heat capacity (mass-weighted mean Cp)
    Cp_eff = float(np.sum(mass_stag * Cp_stag)) / max(M_mantle, 1.0)

    # EOS-consistent mantle integrated enthalpy [J]. Falls back
    # to NaN when no entropy EOS is loaded (e.g. const_properties
    # path), signalling to PROTEUS that the conservation diagnostic
    # is not available for this run.
    if eos is not None:
        from aragog.output.diagnostics import total_enthalpy

        E_state = total_enthalpy(eos, P_stag, S_final, mass_stag)
        # Conservation-grade integrated enthalpy with FROZEN
        # structural mass per shell. Same enthalpy table, but the
        # mass weighting matches the entropy ODE's frame so that
        # d/dt[E_state_cons] equals the divergence-of-flux budget
        # (with frozen-mass Q_*_cons) to numerical precision.
        mass_struct_stag = np.asarray(mesh.staggered_effective_density).ravel() * vol
        E_state_cons = total_enthalpy(eos, P_stag, S_final, mass_struct_stag)
    else:
        E_state = float('nan')
        E_state_cons = float('nan')

    # CMB heat flux reported as the energy-conserving step-average
    # (trapezoidal time-mean over the call; integral equals step_dE_F_cmb_J).
    # Falls back to the end-of-step basic-node-0 value when the call has no
    # integrable trajectory. Sign positive-out-of-core into the mantle.
    F_cmb = step_integrals['F_cmb_step_avg']
    if F_cmb is None:
        F_cmb = float(heat_flux[0])

    # Mantle-integrated source powers [W] for the closed-mantle
    # energy balance dE/dt = -F_int*A_int + F_cmb*A_cmb + Q_radio +
    # Q_tidal. Each per-source heating array is in W/kg at staggered
    # nodes; mass-weighting recovers the total power.
    Q_radio_total = float(np.dot(np.asarray(self.state.heating_radio).ravel(), mass_stag))
    Q_tidal_total = float(np.dot(np.asarray(self.state.heating_tidal).ravel(), mass_stag))

    # Volumetric melt fraction (porosity-based)
    if eos is not None:
        rho_sol = np.asarray(
            eos._lookup_at_phase_boundary('density', P_stag, 'solid')
        ).ravel()
        rho_liq = np.asarray(
            eos._lookup_at_phase_boundary('density', P_stag, 'melt')
        ).ravel()
        rho_bulk = 1.0 / (
            phi_stag / np.where(rho_liq > 0, rho_liq, 1.0)
            + (1 - phi_stag) / np.where(rho_sol > 0, rho_sol, 1.0)
        )
        drho = rho_sol - rho_liq
        safe_drho = np.where(np.abs(drho) > 1e-10, drho, 1.0)
        porosity = np.clip((rho_sol - rho_bulk) / safe_drho, 0, 1)
        Phi_global_vol = float(np.sum(porosity * vol) / np.sum(vol))
    else:
        Phi_global_vol = 1.0  # const_properties: fully liquid

    # Heating flux
    area_surf = 4 * np.pi * float(r_basic[-1]) ** 2
    F_heat_total = float(np.dot(heating, mass_stag)) / area_surf

    shell = self._core_shell()
    if shell is not None:
        y_end = np.asarray(sol.y[:, -1])
        t_shell_end = y_end[n_stag + 2 :]
        layer_base = float(shell.layer_base(t_shell_end, y_end[n_stag + 1]))
        t_top = float(shell.top_temperature(t_shell_end))
    return SolverOutput(
        S_final=S_final,
        T_stag=T_stag,
        phi_stag=phi_stag,
        rho_stag=rho_stag,
        visc_stag=visc_stag,
        P_stag=P_stag,
        r_basic=r_basic,
        r_stag=r_stag,
        vol=vol,
        mass_stag=mass_stag,
        heat_flux=heat_flux,
        heating=heating,
        eddy_diff=eddy_diff,
        cap_stag=cap_stag,
        T_magma=T_magma,
        T_core=T_core,
        Phi_global=Phi_global,
        Phi_global_vol=Phi_global_vol,
        M_mantle=M_mantle,
        M_mantle_liquid=float(Phi_global * M_mantle),
        M_mantle_solid=float((1.0 - Phi_global) * M_mantle),
        RF_depth=RF_depth,
        E_th=E_th,
        E_state=E_state,
        E_state_cons=E_state_cons,
        Cp_eff=Cp_eff,
        F_heat_total=F_heat_total,
        F_cmb=F_cmb,
        Q_radio_total=Q_radio_total,
        Q_tidal_total=Q_tidal_total,
        step_dE_F_int_J=step_integrals['F_int'],
        step_dE_F_cmb_J=step_integrals['F_cmb'],
        step_dE_core_J=step_integrals['core'],
        step_dE_Q_radio_J=step_integrals['Q_radio'],
        step_dE_Q_tidal_J=step_integrals['Q_tidal'],
        step_dE_Q_radio_cons_J=step_integrals['Q_radio_cons'],
        step_dE_Q_tidal_cons_J=step_integrals['Q_tidal_cons'],
        step_solver_residual_J=step_integrals['solver_residual'],
        # state_heat is the conservation-budget state term (Sum rho T dS);
        # compression is an informational PdV diagnostic from the preceding
        # structure re-solve, deliberately kept out of the predicted energy
        # to avoid double-counting the same work (see the SolverOutput field
        # comments and docs/Explanations/energy_diagnostics.md).
        step_dE_compression_J=float(getattr(self, '_last_compression_J', 0.0)),
        step_dE_state_heat_J=step_integrals['state_heat'],
        dt_actual=float(sol.t[-1] - sol.t[0]),
        status=sol.status,
        cvode_flag=int(getattr(sol, 'cvode_flag', 0)),
        cvode_flag_name=str(getattr(sol, 'cvode_flag_name', 'N/A')),
        tcore_change_max=tcore_change_max,
        tcore_change_exceeded=tcore_change_exceeded,
        core_T_shell=None if shell is None else t_shell_end.copy(),
        core_layer_base=np.nan if shell is None else layer_base,
        core_T_top=np.nan if shell is None else t_top,
        jcond_b=jcond_b,
        jconv_b=jconv_b,
        jgrav_b=jgrav_b,
        jmix_b=jmix_b,
        dSdr_b=dSdr_b,
        phi_basic=phi_basic_diag,
        T_basic=T_basic_diag,
        cp_basic=cp_basic_diag,
        rho_basic=rho_basic_diag,
        porosity_b=porosity_b,
        rho_solid_b=rho_solid_b,
        rho_melt_b=rho_melt_b,
        g_b=g_b,
    )

def initialize(self) -> None:
    """Initialize mesh, boundary conditions, and entropy state.

    Unlike the T-based Solver, we only need the mesh and BCs from
    the Evaluator. The T-based phase evaluators (which require
    solidus/liquidus files) are replaced by EntropyPhaseEvaluator.
    """
    logger.info('Initializing EntropySolver')
    self._initialize_internals()

def reset(self) -> None:
    """Reset for a new integration (PROTEUS coupling loop).

    Re-reads the EOS mesh file if eos_method=2, then rebuilds
    the mesh, BCs, and entropy state. Matches Solver.reset().
    """
    logger.info('Resetting EntropySolver')
    # Snapshot the pre-re-solve pressure, carried entropy and frozen
    # mass so the compression work across the mesh change can be
    # measured once the new mesh is built (see ``_last_compression_J``),
    # and the staggered mass coordinates so the carried entropy can be
    # remapped onto the new mass grid per parcel.
    P_old, S_old, mass_old = self._snapshot_for_compression()
    try:
        self._xi_pre_resolve = self.staggered_mass_coordinates.copy()
    except (AttributeError, TypeError):
        self._xi_pre_resolve = None
    if self.parameters.mesh.eos_method == 2 and self.parameters.mesh.eos_file:
        arr = np.loadtxt(self.parameters.mesh.eos_file)
        self.parameters.mesh.eos_radius = arr[:, 0]
        self.parameters.mesh.eos_pressure = arr[:, 1]
        self.parameters.mesh.eos_density = arr[:, 2]
        self.parameters.mesh.eos_gravity = arr[:, 3]
        _validate_eos_radius_range(self.parameters.mesh)
    self._initialize_internals()
    self._last_compression_J = self._compute_compression_work(P_old, S_old, mass_old)

def set_initial_core_temperature(self, T_core_init: float | None) -> None:
    """Set the initial core temperature (``bower2018`` / ``core_module``).

    Must be called BEFORE ``set_initial_entropy``. If not called,
    the initial T_core defaults to the previous solution's final value
    (if any), else to the EOS temperature of the bottom cell's entropy
    in S_init, at the CMB pressure for ``core_module`` (no CMB flux at
    the start without a shell) and at the bottom-cell pressure for ``bower2018``.

    The value applies to every later call of ``set_initial_entropy``.
    Pass ``None`` to clear it, restoring the hot-start behaviour; only a
    set value is checked against the mantle temperature at the CMB.
    """
    self._T_core_init = None if T_core_init is None else float(T_core_init)

def set_initial_dSdr_cmb(self, dSdr_cmb_init: float | None) -> None:
    """Set the initial CMB entropy gradient (energy_balance / core_module).

    Must be called BEFORE ``set_initial_entropy``. If not called,
    the initial ``dSdr_cmb`` is taken from the previous solution
    (if any), else from a one-sided FD of the staggered S_init at
    the bottom (which is zero for a uniform isentrope).

    Pass ``None`` to clear a previously-set override, restoring
    the hot-start behaviour on the next call to
    ``set_initial_entropy``. Used by PROTEUS's retry ladder to
    restore the pre-solve dSdr_cmb after a rejection, then release
    the override so the subsequent coupling step can hot-start
    from its own ``_solution`` normally.
    """
    self._dSdr_cmb_init = None if dSdr_cmb_init is None else float(dSdr_cmb_init)

def set_initial_entropy(self, S_init: npt.NDArray | float) -> None:
    """Set the initial entropy profile and (if used) initial T_core.

    Parameters
    ----------
    S_init : array or float
        Entropy at staggered nodes [J/kg/K]. If scalar, sets uniform
        (isentropic) profile.

    Notes
    -----
    State vector length depends on the active core BC mode:
    ``quasi_steady`` uses length N; ``energy_balance`` and
    ``bower2018`` use length N+1 (entropy block plus one extra
    state variable: ``dSdr_cmb`` or ``T_core`` respectively);
    ``gradient`` uses length N+2; ``core_module`` uses length N+2
    (``dSdr_cmb`` and ``T_core``) plus the shell temperatures of a
    stratified core. For the bower2018 mode the
    initial ``T_core`` is taken from the bottom-cell mantle
    temperature derived from ``S_init`` via the EOS unless
    ``set_initial_core_temperature`` has been called first.
    """
    # Prefer the cached _n_stag from _initialize_internals; fall
    # back to the mesh accessor for legacy callers that bypass it.
    if hasattr(self, '_n_stag') and self._n_stag is not None:
        n_stag = self._n_stag
    else:
        n_stag = self.evaluator.mesh.staggered.radii.shape[0]
        self._n_stag = n_stag

    if np.isscalar(S_init):
        S_arr = np.full(n_stag, float(S_init))
    else:
        S_arr = np.asarray(S_init, dtype=float)
        if len(S_arr) != n_stag:
            raise ValueError(f'S_init length {len(S_arr)} != mesh staggered nodes {n_stag}')

    # Carry the entropy onto the current mesh by mass parcel when a
    # structure re-solve changed the mass grid (no-op otherwise, and
    # always consumes the cached pre-resolve grid).
    S_arr = self._remap_entropy_to_current_mesh(S_arr)

    # Prefer the cached _core_bc from _initialize_internals.
    if hasattr(self, '_core_bc') and self._core_bc is not None:
        core_bc = self._core_bc
    else:
        core_bc = getattr(self.parameters.boundary_conditions, 'core_bc', 'quasi_steady')
        self._core_bc = core_bc
    logger.debug('set_initial_entropy: core_bc=%r, n_stag=%d', core_bc, n_stag)

    if core_bc == 'energy_balance':
        # SPIDER bit-parity core BC (energy_balance mode).
        # State = [S_0, ..., S_{N-1}, dSdr_cmb]
        # The boundary state dSdr_cmb is the entropy gradient at
        # the CMB basic node (mirror of SPIDER's dSdxi[ind_cmb]).
        # Its time derivative is set by the bc.c:76-131 formula
        # in _dSdt_single from the actual physical heat flux.
        #
        # Cold-start dSdr_cmb_init: use the finite-difference
        # estimate from the staggered cells (one-sided forward
        # difference, since there's no cell below the CMB).
        # The boundary state will then evolve from this to its
        # quasi-equilibrium value over the first few coupling
        # steps as the energy balance constraint kicks in.
        #
        # Hot-start dSdr_cmb_init: preserve from the previous
        # solution if available, so the integrated boundary state
        # survives PROTEUS coupling resets.
        dSdr_cmb_init = self._resolve_dSdr_cmb_init(S_arr, n_stag)
        self._S0 = np.empty(n_stag + 1)
        self._S0[:n_stag] = S_arr
        self._S0[n_stag] = float(dSdr_cmb_init)
        logger.info(
            'Initial state (energy_balance): S_min=%.0f, S_max=%.0f, dSdr_cmb_init=%.3e',
            S_arr.min(),
            S_arr.max(),
            dSdr_cmb_init,
        )
    elif core_bc in ('bower2018', 'core_module'):
        # Core temperature as an ODE state: bower2018 [S, T_core] (parity testing only),
        # core_module [S, dSdr_cmb, T_core] with T_core the core budget's state. The CMB
        # flux law takes T_core (the top shell cell with a shell) against the bottom cell.
        slots = EXTRA_STATE_SLOTS[core_bc]
        n_extra = len(slots)
        T_core_slot = n_stag + slots.index('T_core')
        T_core_init = getattr(self, '_T_core_init', None)
        t_core_override = T_core_init is not None
        if T_core_init is None:
            T_core_init = self.get_current_core_temperature()

        dSdr_cmb_init = None
        if 'dSdr_cmb' in slots:
            dSdr_cmb_init = self._resolve_dSdr_cmb_init(S_arr, n_stag)

        T_bottom_eos = None
        if (
            T_core_init is None or core_bc == 'core_module'
        ) and self.entropy_eos is not None:
            # core_module: the bottom cell at the CMB pressure, the mantle side of the
            # contrast of the flux law, so the default start without a shell has no flux.
            P_bot, S_bot = float(self._P_stag_flat[0]), float(S_arr[0])
            if core_bc == 'core_module':
                P_bot = float(getattr(self, '_P_basic_flat', self._P_stag_flat)[0])
            T_bottom_eos = float(
                np.asarray(
                    self.entropy_eos.temperature(np.array([P_bot]), np.array([S_bot]))
                ).item()
            )

        if T_core_init is None:
            if T_bottom_eos is not None:
                T_core_init = T_bottom_eos
            else:
                raise ValueError(
                    f"Cannot initialize core temperature for core_bc='{core_bc}': "
                    'entropy EOS is not available and no initial core temperature was provided.'
                )
        if (
            core_bc == 'core_module'
            and t_core_override
            and not getattr(self, '_t_core_offset_warned', False)
            and T_bottom_eos is not None
            and abs(T_core_init - T_bottom_eos) > 0.2 * max(T_bottom_eos, 1.0)
        ):
            # Once, for a supplied start: a hot start continues the run's own contrast.
            # Under a stiff base the flux relaxes a large offset only slowly.
            self._t_core_offset_warned = True
            logger.warning(
                'core_module: initial T_core=%.0f K differs from the '
                'mantle temperature at the CMB %.0f K by more than 20%%; '
                'the CMB flux relaxes the offset at the rate the basal '
                'boundary layer allows, which shifts inner-core nucleation timing.',
                T_core_init,
                T_bottom_eos,
            )
        shell = self._core_shell() if core_bc == 'core_module' else None
        self._S0 = np.empty(n_stag + n_extra + (0 if shell is None else shell.n_cells))
        self._S0[:n_stag] = S_arr
        self._S0[T_core_slot] = T_core_init
        if 'dSdr_cmb' in slots:
            self._S0[n_stag + slots.index('dSdr_cmb')] = float(dSdr_cmb_init)
        if shell is not None:
            self._S0[n_stag + n_extra :] = self._initial_shell(
                shell, T_core_init, t_core_override
            )
        logger.info(
            'Initial state (%s): S_min=%.0f, S_max=%.0f, T_core_init=%.0f K',
            core_bc,
            S_arr.min(),
            S_arr.max(),
            T_core_init,
        )
    elif core_bc == 'gradient':
        # Gradient-based formulation mirroring SPIDER's dS/dxi state.
        # State = [dS/dr at N+1 basic nodes, S_surf], length N+2.
        # S at staggered nodes is reconstructed by cumulative sum
        # at each RHS evaluation, providing the global smoothing
        # that prevents BDF overshoot at the solidus.
        mesh = self.evaluator.mesh
        dSdr_init = mesh.d_dr_at_basic_nodes(S_arr).ravel()
        n_basic = len(dSdr_init)
        S_surf = float(S_arr[-1])
        self._S0 = np.empty(n_basic + 1)
        self._S0[:n_basic] = dSdr_init
        self._S0[n_basic] = S_surf
        # Verify roundtrip
        S_check, _ = self._reconstruct_entropy(dSdr_init, S_surf)
        roundtrip_err = float(np.max(np.abs(S_check - S_arr)))
        logger.info(
            'Initial state (gradient): dSdr range [%.3e, %.3e] J/kg/K/m, '
            'S_surf=%.1f, roundtrip err=%.2e J/kg/K',
            dSdr_init.min(),
            dSdr_init.max(),
            S_surf,
            roundtrip_err,
        )
    else:
        # quasi_steady BC: state = [S_0, ..., S_{N-1}]
        self._S0 = S_arr
        logger.info(
            'Initial entropy (quasi_steady BC): S_min=%.0f, S_max=%.0f J/kg/K',
            S_arr.min(),
            S_arr.max(),
        )

def set_initial_shell_temperature(self, t_shell: npt.NDArray | None) -> None:
    """Set the shell temperatures [K] of a stratified core (a resume).

    Must be called before ``set_initial_entropy``. The profile applies to every later
    call of ``set_initial_entropy``; ``None`` clears it, so the shell starts from the
    previous solution, or on the adiabat of the core without one.
    """
    self._T_shell_init = None if t_shell is None else np.asarray(t_shell, dtype=float)

def set_jax_cvode_factory(self, factory) -> None:
    """Register a factory that produces JAX-derived CVODE callbacks.

    The factory is called inside ``solve()`` when ``use_jax_jacobian``
    is enabled in the config. It is given the current nondim scaling
    spec (NonDimScales) and core-BC mode and must return the
    ``(rhs_fn, jac_fn)`` pair accepted by scikits.odes CVODE.

    Parameters
    ----------
    factory : callable
        ``factory(scales, core_bc_mode) -> (rhs_fn, jac_fn)`` where
        ``scales`` is an ``aragog.jax.nondim.NonDimScales`` instance.
        ``rhs_fn`` may optionally carry a ``step_powers(t_nodes, Y_nodes, aux)``
        attribute for batched per-call energy power evaluation.
        May be None to disable the Option Z path even if the flag
        is on.
    """
    self._jax_cvode_factory = factory

def solve(self, _floor_discard: str | None = None) -> None:
    """Run the BDF time integration.

    A successful solve also stores the per-call energy integrals on the result as
    ``energy_integrals`` (see ``_compute_step_energy_integrals``); a failed one stores
    zeros. ``_floor_discard`` (internal) refuses the call with that warning: a call whose
    floor root returns a state below the EOS table, or stops at the call start, is solved
    again this way, so it returns its start state.
    """
    wall_start = time.perf_counter()
    self._sync_top_node()  # the EOS may have changed since the reset
    start_time = self.parameters.solver.start_time
    end_time = self.parameters.solver.end_time
    # Absolute tolerance floor (1e-8) matches SPIDER's atol=rtol.
    # Tight tolerance is necessary to resolve the crystallisation
    # front and prevent cumulative integration error during long
    # solidification runs.
    atol_base = max(self.parameters.solver.atol, 1.0e-8)
    rtol = self.parameters.solver.rtol

    # Phase-aware atol: tight during crystallization, relaxed
    # only when fully solid.
    try:
        n_stag = self._n_stag
        if self._core_bc == 'gradient':
            # Gradient mode: reconstruct S to compute phi0
            n_basic = n_stag + 1
            dSdr0 = self._S0[:n_basic]
            S_surf0 = float(self._S0[n_basic])
            S0_block, _ = self._reconstruct_entropy(dSdr0, S_surf0)
        elif self._state_is_extended:
            S0_block = self._S0[:n_stag]
        else:
            S0_block = self._S0
        phi0 = float(
            np.asarray(self.entropy_eos.melt_fraction(self._P_stag_flat, S0_block)).mean()
        )
        phi0 = max(0.0, min(1.0, phi0))
    except Exception:
        phi0 = 1.0

    if phi0 > 0.01:
        atol_scale = 1.0
        max_step = 100.0  # years
    else:
        atol_scale = 10.0
        max_step = np.inf

    # Step-cap anchor of the root or event, None when no cap is armed (the floor arms the
    # root without it); set here for the paths without an EOS, such as const_properties.
    phi_cap_anchor = None
    floor, floor_refused = None, False
    rate_mode = False

    # Runs whenever the entropy EOS is loaded: the T and S caps stay armed in the deep solid,
    # where the phi cap is blind. A cell within the margin of a boundary or in the mushy band
    # tightens max_step: 1 yr for 'fixed', CVODE segments for 'rate'.
    if self.entropy_eos is not None:
        entropy_margin = _resolve_entropy_margin(
            getattr(self.parameters.energy, 'phase_boundary_entropy_margin', None)
        )
        # S0_block is staggered (length n_stag). For per-cell
        # phase-boundary check we use the staggered S directly
        # against the staggered-pressure phase boundaries.
        P_stag = self._P_stag_flat
        S_liq_stag = np.asarray(self.entropy_eos.liquidus_entropy(P_stag)).ravel()
        S_sol_stag = np.asarray(self.entropy_eos.solidus_entropy(P_stag)).ravel()
        S_arr_stag = np.asarray(S0_block).ravel()
        margin_to_liq = S_arr_stag - S_liq_stag  # > 0 means above liquidus
        margin_to_sol = S_arr_stag - S_sol_stag  # > 0 means above solidus
        near_liq = np.any(np.abs(margin_to_liq) < entropy_margin)
        near_sol = np.any(np.abs(margin_to_sol) < entropy_margin)
        in_mushy = np.any((margin_to_liq < 0.0) & (margin_to_sol > 0.0))
        # Keep the original CMB-only check too as a backstop
        P_cmb = float(self._P_basic_flat[0])
        S_liq = float(self.entropy_eos.liquidus_entropy(np.array([P_cmb])).item())
        S_sol = float(self.entropy_eos.solidus_entropy(np.array([P_cmb])).item())
        S0_block_cmb = float(S0_block[0])
        cap_set = getattr(self.parameters.energy, 'phase_boundary_cap', None)
        cap_mode = cap_set or 'rate'
        # Staggered cells plus the CMB cell at the CMB-pressure boundaries.
        rate_liq = np.append(S_liq_stag, S_liq)
        rate_sol = np.append(S_sol_stag, S_sol)
        w = float(getattr(self.parameters.phase_mixed, 'matprop_smooth_width', 0.0))
        rate_delta = max(3.0 * w * float(np.max(rate_liq - rate_sol)), _RATE_CAP_MIN_DELTA)
        use_cvode = (
            getattr(self.parameters.energy, 'solver_method', 'cvode') == 'cvode'
            and _CVODE_AVAILABLE
        )
        segment_margin = _segment_arming_margin(
            entropy_margin, rate_delta, cap_mode, self._core_bc, use_cvode
        )

        def clamps(margin):
            return _phase_boundary_max_step_clamp(
                np.any(np.abs(margin_to_liq) < margin),
                np.any(np.abs(margin_to_sol) < margin),
                in_mushy,
                cmb_margin_to_liq=S0_block_cmb - S_liq,
                cmb_margin_to_sol=S0_block_cmb - S_sol,
                entropy_margin=margin,
            )

        if clamps(entropy_margin) or (
            segment_margin > entropy_margin and clamps(segment_margin)
        ):
            rate_mode = cap_mode == 'rate' and self._core_bc != 'gradient'
            max_step = 1.0
            if cap_mode == 'rate' and not rate_mode:
                self._warn_once(
                    'gradient',
                    'phase_boundary_cap="rate" is not used by the gradient core; '
                    'max_step 1 yr',
                    logging.WARNING if cap_set else logging.INFO,
                )

        # Per-call step caps as a SUNDIALS root function. Build the
        # anchor metadata here (per-cell melt fraction, temperature and
        # entropy at solve entry); the rootfn instance is constructed
        # below once the nondim state_scale is known. CVODE returns at
        # the exact time the tightest active cap is reached, enforced by
        # the integrator's own trajectory rather than a start-time rate
        # estimate. The per-cell melt-fraction term catches a deep cell
        # crossing the two-phase window in one step (a negligible move in
        # the global mean but the source of the crystallisation-onset
        # core-temperature jump). The melt-fraction derivative saturates
        # once a cell is fully solid, so it cannot bound the temperature
        # drop on the solid adiabat just below the solidus; the optional
        # temperature and entropy caps bound |ΔT| and |ΔS| per cell
        # directly. All armed caps share one rootfn (g is the minimum
        # margin). Armed as soon as any cell is inside, near, or below the
        # two-phase window so the caps stay active through the solid
        # adiabat. The caps are skipped in gradient core_bc mode, where the
        # state vector holds entropy gradients rather than entropies and
        # the per-cell EOS lookups would be dimensionally meaningless.
        phi_step_cap = _resolve_step_cap(
            getattr(self.parameters.energy, 'phi_step_cap', None)
        )
        temp_step_cap = _resolve_step_cap(
            getattr(self.parameters.energy, 'temperature_step_cap', None)
        )
        entropy_step_cap = _resolve_step_cap(
            getattr(self.parameters.energy, 'entropy_step_cap', None)
        )
        any_cap = phi_step_cap > 0.0 or temp_step_cap > 0.0 or entropy_step_cap > 0.0
        below_solidus = bool(np.any(margin_to_sol < 0.0))
        gradient_mode = getattr(self, '_core_bc', None) == 'gradient'
        if gradient_mode and any_cap:
            logger.info('step caps are not supported in gradient core_bc mode; not armed')
        if (
            any_cap
            and not gradient_mode
            and (near_liq or near_sol or in_mushy or below_solidus)
        ):
            try:
                rho_stag_init = np.asarray(
                    self.entropy_eos.density(self._P_stag_flat, S_arr_stag)
                ).ravel()
                phi_stag_init = np.asarray(
                    self.entropy_eos.melt_fraction(self._P_stag_flat, S_arr_stag)
                ).ravel()
                mass_stag_init = rho_stag_init * self._volume_flat
                mass_total_init = float(np.sum(mass_stag_init))
                if mass_total_init > 0.0:
                    phi0_global = float(
                        np.sum(mass_stag_init * phi_stag_init) / mass_total_init
                    )
                    T0_stag_init = None
                    if temp_step_cap > 0.0:
                        T0_stag_init = np.asarray(
                            self.entropy_eos.temperature(self._P_stag_flat, S_arr_stag)
                        ).ravel()
                    phi_cap_anchor = {
                        'cap_phi': phi_step_cap,
                        'phi0_global': phi0_global,
                        'phi0_cell': phi_stag_init,
                        'cap_T': temp_step_cap,
                        'T0_cell': T0_stag_init,
                        'cap_S': entropy_step_cap,
                        'S0_cell': S_arr_stag.copy(),
                    }
                    logger.debug(
                        'step caps armed (CVODE rootfn): phi=%.3g, T=%.3g K, S=%.3g '
                        'J/kg/K (global Φ0=%.4f)',
                        phi_step_cap,
                        temp_step_cap,
                        entropy_step_cap,
                        phi0_global,
                    )
            except Exception as exc:
                logger.warning(
                    'step-cap rootfn anchor failed (%s); proceeding without cap',
                    exc,
                )

        # EOS floor: a root ends the call at S_min + m, or S_min + m/2 for a start below
        # S_min + m; closer starts are decided by _floor_start_level.
        floor_margin = _resolve_step_cap(
            getattr(self.parameters.energy, 'eos_floor_margin', EOS_FLOOR_MARGIN_DEFAULT)
        )
        S_tab = float(self.entropy_eos.S_min)
        if floor_margin >= float(self.entropy_eos.S_max) - S_tab:
            raise ValueError(
                f'eos_floor_margin {floor_margin!r} J/kg/K must be below the entropy '
                f'span of the EOS tables, {float(self.entropy_eos.S_max) - S_tab!r}'
            )
        if _floor_discard is not None:
            floor_refused, end_time = True, start_time
            logger.warning(_floor_discard)
        elif floor_margin > 0.0 and not gradient_mode and end_time > start_time:
            mesh = self.evaluator.mesh
            floor_w = np.asarray(mesh._quantity_transform)[-1]
            if mesh.top_node == 'limited':  # the node stays in the table with its cells
                floor_w = np.eye(1, floor_w.size, floor_w.size - 1).ravel()
            level, refuse = _floor_start_level(
                floor_w,
                S_arr_stag,
                lambda: np.ravel(self._dSdt_single(start_time, self._S0))[: self._n_stag],
                S_tab,
                floor_margin,
                end_time - start_time,
            )
            if level is not None:
                floor = dict(
                    floor_weights=floor_w, floor_level=level, floor_margin=floor_margin
                )
            elif refuse:
                floor_refused, end_time = True, start_time
                s_floor = _floor_entropy(floor_w, S_arr_stag)
                band, half = 2.0e-6 * floor_margin, 0.5 * (s_floor + S_tab)
                if s_floor <= S_tab:
                    why = 'at or below the table minimum'
                elif s_floor <= S_tab + band and not S_tab < half < s_floor:
                    why = 'no level fits between it and the table minimum'
                elif s_floor <= S_tab + band:
                    why = f'counts as at the table minimum (within {band:.2g} J/kg/K) and falls'
                else:
                    why = (
                        f'at most {0.5 * floor_margin:.1f} J/kg/K above it, falls too fast'
                    )
                logger.warning(
                    'EOS floor: entropy %.9g J/kg/K at the call start, %+.3e J/kg/K from the '
                    'table minimum %.9g J/kg/K, %s; the call does not advance',
                    s_floor,
                    s_floor - S_tab,
                    S_tab,
                    why,
                )

    # External per-call atol scale factor (PROTEUS retry ladder
    # opts in by setting solver._atol_sf before solve()). Default
    # 1.0 leaves behaviour unchanged for callers that don't set it.
    atol_sf_external = float(getattr(self, '_atol_sf', 1.0))
    atol = atol_base * atol_scale * atol_sf_external

    logger.info(
        'EntropySolver: integrating from %.2e to %.2e yr '
        '(Phi_init=%.3f, atol_scale=%.1fx, atol_sf=%.1fx, atol=%.2e, rtol=%.2e)',
        start_time,
        end_time,
        phi0,
        atol_scale,
        atol_sf_external,
        atol,
        rtol,
    )

    # ── Nondimensionalise state, time, and tolerances ──
    # All state components become O(1) after dividing by their
    # reference scales, so scalar atol works uniformly and the
    # BDF solver no longer needs to resolve 9+ significant digits
    # on an O(1000) state variable. The physics code in
    # _dSdt_single stays in physical units; only the solver
    # interface scales in and out.
    # NonDimScales is the single source of truth for nondim
    # scaling, consumed by both the scipy/CVODE wrapper here and
    # the JAX factory. ``__post_init__`` enforces the contract
    # ``rhs_scale = t_ref / state_scale``.
    scales = self._build_nondim_scales()
    _offset = scales.state_offset
    S_ref = self._S_ref
    t_ref = scales.t_ref
    n_s = self._n_stag
    _state_scale = scales.state_scale
    _rhs_scale = scales.rhs_scale

    S0_nd = (self._S0 - _offset) / _state_scale
    start_nd = start_time / t_ref
    end_nd = end_time / t_ref
    max_step_nd = max_step / t_ref if np.isfinite(max_step) else max_step

    # Per-component atol in nondim units. Each state component
    # has its own nondim scale (S_ref for entropy, dSdr_ref for
    # dSdr_cmb in energy_balance mode, dSdr_ref for all gradient
    # components in gradient mode). Dividing atol by _state_scale
    # element-wise gives the same physical tolerance on every
    # component, regardless of its nondim scale. A scalar
    # atol_nd = atol/S_ref would force ~10^8x tighter physical
    # tolerance on dSdr_cmb, choking the step size and
    # suppressing the entropy cooling rate enough to prevent
    # solidification.
    atol_nd = atol / _state_scale

    def _rhs_nondim(t_nd, y_nd):
        """Nondim wrapper: scale state to physical, call physics, scale RHS back."""
        dydt_phys = self._dSdt_single(
            t_nd * t_ref,
            y_nd * _state_scale + _offset,
        )
        return dydt_phys * _rhs_scale

    logger.info(
        'Nondimensionalisation: S_ref=%.3f t_ref=%.3e yr '
        'atol_nd min/max=[%.2e, %.2e] S0_nd range [%.4f, %.4f]',
        S_ref,
        t_ref,
        float(np.min(atol_nd)),
        float(np.max(atol_nd)),
        S0_nd[:n_s].min(),
        S0_nd[:n_s].max(),
    )

    # BDF integration with phase-aware max_step constraint.
    jac_sparsity = self._build_jac_sparsity()

    # ── melt-fraction cap rootfn / event construction ──
    # Build the SUNDIALS rootfn (CVODE path) and the equivalent
    # solve_ivp event (scipy fallback) only when the cap is armed.
    # Both consume nondim ``y`` and rescale via ``_state_scale``
    # before reading the EOS. The phase-boundary max_step=1 yr
    # clamp set above is preserved; the cap is an ADDITIONAL
    # guardrail anchored at the per-cell and global melt fractions at
    # solve entry.
    events = None
    phi_cap_rootfn = None
    if phi_cap_anchor is not None or floor is not None:
        a = phi_cap_anchor or {}
        try:
            cap_kw = dict(
                eos=self.entropy_eos,
                P_stag=self._P_stag_flat,
                volume=self._volume_flat,
                n_stag=self._n_stag,
                phi0_global=a.get('phi0_global', 0.0),
                cap=a.get('cap_phi', 0.0),
                state_scale=_state_scale,
                phi0_per_cell=a.get('phi0_cell'),
                cap_temperature=a.get('cap_T', 0.0),
                T0_per_cell=a.get('T0_cell'),
                cap_entropy=a.get('cap_S', 0.0),
                S0_per_cell=a.get('S0_cell'),
                **(floor or {}),
            )
            phi_cap_rootfn = _PhiCapRootFunction(**cap_kw)
            events = [_phi_cap_event_factory(**cap_kw)]
        except Exception as exc:
            logger.warning(
                'step-cap rootfn instantiation failed (%s); integrating without cap',
                exc,
            )
            phi_cap_rootfn = None
            events = None

    # Integrator dispatch. SUNDIALS CVODE (via scikits.odes) is
    # the same solver SPIDER uses and is the recommended choice
    # for production-grade stiff integration: scipy's BDF and
    # Radau collapse their step size to machine epsilon at the
    # crystallisation front and abort with `Required step size
    # is less than spacing between numbers` because scipy's
    # Newton iterator cannot converge through the stiff phase
    # transition. CVODE's modified-Newton with cached Jacobian
    # factorisation handles the discontinuity cleanly. Scipy
    # solve_ivp (Radau or BDF) is kept as a fallback for systems
    # without scikits.odes available.
    solver_method = getattr(self.parameters.energy, 'solver_method', 'cvode')
    # Warn when CVODE was requested but scikits.odes is not
    # importable: the fallback to scipy Radau is a substantial
    # change in solver behaviour (no modified-Newton, no cached
    # Jacobian factorisation) and a silent fallback would let a
    # broken environment masquerade as a physics regression.
    if solver_method == 'cvode' and not _CVODE_AVAILABLE:
        logger.warning(
            'EntropySolver: solver_method="cvode" requested but '
            'scikits.odes is not installed; falling back to scipy '
            'Radau. Install scikits-odes to enable the production '
            'CVODE path (same solver SPIDER uses).'
        )
    use_cvode = solver_method == 'cvode' and _CVODE_AVAILABLE
    self._cvode_step_powers_batch = None
    if use_cvode:
        # Build JAX-derived CVODE callbacks when the factory is registered
        # and use_jax_jacobian is active. Callbacks reuse cached functions.
        cvode_rhs_override = None
        cvode_jacfn = None
        use_jax_jac = (
            getattr(self.parameters.energy, 'use_jax_jacobian', False)
            and self._jax_cvode_factory is not None
        )
        if use_jax_jac:
            try:
                # Pass the NonDimScales instance, which bundles
                # the internal nondim contract.
                cvode_rhs_override, cvode_jacfn = self._jax_cvode_factory(
                    scales,
                    self._core_bc,
                )
                self._cvode_step_powers_batch = getattr(
                    cvode_rhs_override, 'step_powers', None
                )
                logger.info(
                    'EntropySolver: option Z active '
                    '(JAX analytic Jacobian + JAX RHS via scikits.odes)'
                )
            except Exception as exc:  # pragma: no cover - fallback path
                logger.warning(
                    'EntropySolver: JAX CVODE factory failed (%s); '
                    'falling back to numpy RHS + FD Jacobian',
                    exc,
                )
                cvode_rhs_override = None
                cvode_jacfn = None

        if cvode_rhs_override is None:
            logger.info('EntropySolver: using CVODE (solver_method=cvode)')
        cvode_kw = dict(
            atol=atol_nd,
            rtol=rtol,
            rhs=_rhs_nondim,
            cvode_rhs_fn_override=cvode_rhs_override,
            cvode_jacfn=cvode_jacfn,
            phi_cap_rootfn=phi_cap_rootfn,
        )
        spans = end_time > start_time
        D = self.state._convective_threshold
        if spans and D > 0.0:
            self._warn_once(
                'threshold',
                f'convective_threshold {D:.1e} > 0: kappa_h jumps at dS/dr = 0, which can '
                'stall CVODE at a node whose gradient stays near 0, also at rtol 1e-8; '
                'see the heat transport docs',
            )
        if spans and rtol > _CVODE_TOL_LIMIT:
            atol_cfg = self.parameters.solver.atol
            loose = atol_cfg > _CVODE_TOL_LIMIT and self._core_bc == 'energy_balance'
            note = (
                f'; the configured solver.atol {atol_cfg:.1e} adds to this' if loose else ''
            )
            self._warn_once(
                'cvode_tol',
                f'CVODE at solver.rtol {rtol:.1e}: an rtol above 1e-8 can lock CVODE '
                'at the convective switch (convective mask and kappa_h floor; in test runs '
                'at convective_threshold = 0 '
                'with atol = rtol and the initial entropy drawn per run: 12 of 1133 at 1e-6, '
                '0 of 1026 at 1e-8, and 0 of 80 at 1e-5 but about 12 times more RHS '
                f'evaluations than at 1e-8; see the solver tuning docs){note}. '
                'Use 1e-8, the Aragog default, for rtol and atol',
            )
        if rate_mode:
            if spans and rtol > _RATE_CAP_RTOL_LIMIT:
                self._warn_once(
                    'rtol',
                    f'phase_boundary_cap="rate" at rtol {rtol:.1e}: its accuracy is measured '
                    'at rtol 1e-8',
                )

            def _rate_h(t_nd, y_nd):
                try:
                    y = np.asarray(y_nd, dtype=float) * _state_scale + _offset
                    S = y[:n_s]
                    dSdt = np.asarray(self._dSdt_single(t_nd * t_ref, y)).ravel()[:n_s]
                    rho = np.asarray(self.entropy_eos.density(self._P_stag_flat, S)).ravel()
                    h = _rate_phase_boundary_max_step(
                        np.append(S, S[0]),
                        np.append(dSdt, dSdt[0]),
                        rate_liq,
                        rate_sol,
                        np.append(rho * self._volume_flat, 0.0),
                        rate_delta,
                    )
                except Exception as exc:
                    logger.warning(
                        'rate cap: dS/dt evaluation failed (%s); max_step 1 yr', exc
                    )
                    h = 1.0
                return h / t_ref

            def _rate_roots(y_nd, inside):
                return _PhaseBoundarySegmentRoot(
                    rate_liq,
                    rate_sol,
                    y_nd,
                    rate_delta,
                    inside,
                    _state_scale,
                    n_s,
                    cap=phi_cap_rootfn,
                )

            self._solution = self._solve_cvode_segments(
                start_time=start_nd,
                end_time=end_nd,
                y0=S0_nd,
                h_at=_rate_h,
                roots_at=_rate_roots,
                t_ref=t_ref,
                h_min=1.0 / t_ref,
                **cvode_kw,
            )
        else:
            self._solution = self._solve_cvode(
                start_time=start_nd,
                end_time=end_nd,
                y0=S0_nd,
                max_step=max_step_nd,
                **cvode_kw,
            )
    else:
        shell = self._core_shell() is not None
        if shell and solver_method == 'radau':
            raise ValueError(
                "a stratified core_module needs solver_method 'cvode' or 'bdf': Radau "
                'takes minutes per call on the stiff mixing of the resolved shell'
            )
        method = 'BDF' if solver_method == 'bdf' or shell else 'Radau'
        logger.info('EntropySolver: using scipy %s', method)
        if rate_mode:
            self._warn_once(
                'scipy',
                'phase_boundary_cap="rate" needs CVODE; max_step 1 yr with scipy',
                logging.WARNING if cap_set else logging.INFO,
            )
        self._solution = solve_ivp(
            _rhs_nondim,
            (start_nd, end_nd),
            S0_nd,
            method=method,
            vectorized=False,
            dense_output=False,
            atol=atol_nd,
            rtol=rtol,
            jac_sparsity=jac_sparsity,
            max_step=max_step_nd,
            events=events,
        )

    # ── Restore physical units ──
    sol = self._solution
    if sol.t is not None:
        sol.t = np.asarray(sol.t, dtype=float) * t_ref
    if sol.y is not None:
        sol_y = np.asarray(sol.y, dtype=float)
        if sol_y.ndim == 2:
            sol.y = sol_y * _state_scale[:, np.newaxis] + _offset[:, np.newaxis]
        else:
            sol.y = sol_y * _state_scale + _offset
    trace = sol.get('energy_trace')
    if trace is not None:
        sol.energy_trace = self._physical_trace(trace, t_ref, _state_scale, _offset)
    if floor_refused:
        sol.status = 1
        sol.message = 'EOS floor: the call does not advance'
    elif floor is not None and sol.status >= 0 and sol.t is not None and len(sol.t):
        t_end = float(np.asarray(sol.t)[-1])
        ys = [np.asarray(sol.y, dtype=float)]
        if trace is not None:
            ys.append(np.asarray(sol.energy_trace[1], dtype=float))
        S_all = np.concatenate([y[:n_s] for y in ys], axis=1)
        low = min(float(np.min(floor['floor_weights'] @ S_all)), float(S_all.min()))
        if low < S_tab:
            return self.solve(
                _floor_discard=(
                    f'EOS floor: the call with level S_min {floor["floor_level"] - S_tab:+.3e} '
                    f'J/kg/K put the floor entropy at S_min {low - S_tab:+.3e} J/kg/K at '
                    f'{t_end - start_time:.6g} yr after the call start {start_time:.9g} yr '
                    f'(table minimum {S_tab:.9g} J/kg/K); the call is discarded and does not '
                    'advance'
                )
            )
        if t_end <= float(np.asarray(sol.t)[0]):  # the returned start, after the time scale
            return self.solve(
                _floor_discard=(
                    f'EOS floor: the root of level S_min {floor["floor_level"] - S_tab:+.3e} '
                    f'J/kg/K fell at the call start t={start_time:.9g} yr; the call does not '
                    'advance'
                )
            )

    # Step-cap-fire log, in physical time (after the t_ref restoration
    # above) and naming whichever margin actually bound: read from
    # ``sol.cap_label`` on the CVODE path, from the event closure's
    # ``binding_cap`` on the scipy fallback path.
    cap_display_names = {
        'phi': 'ΔΦ_global',
        'temperature': 'ΔT',
        'entropy': 'ΔS',
        'eos_floor': 'EOS floor',
    }
    if getattr(sol, 'cap_fired', False) and sol.t is not None:
        t_root_phys = float(sol.t[-1])
        cap_label = getattr(sol, 'cap_label', None)
        if cap_label is None:
            # The terminating rootfn call raised, so which cap bound is
            # unknown; report that rather than guessing 'phi'.
            logger.info(
                'step cap fired at t=%.3e yr after %d evals; '
                'binding cap unknown (EOS lookup failed on the terminating call)',
                t_root_phys,
                int(getattr(sol, 'cap_evals', 0)),
            )
        else:
            cap_name = cap_display_names.get(cap_label, 'ΔΦ_global')
            if cap_label == 'phi':
                logger.info(
                    '%s cap: CVODE rootfn fired at '
                    't=%.3e yr after %d evals; cap=%.3g, '
                    'Φ_global(start)=%.4f',
                    cap_name,
                    t_root_phys,
                    int(getattr(sol, 'cap_evals', 0)),
                    float(getattr(sol, 'cap_value', 0.0)),
                    float(getattr(sol, 'cap_phi0', 0.0)),
                )
            else:
                floor_fired = cap_label == 'eos_floor'
                logger.log(
                    logging.WARNING if floor_fired else logging.INFO,
                    '%s cap: CVODE rootfn fired at t=%.3e yr after %d evals; %s',
                    cap_name,
                    t_root_phys,
                    int(getattr(sol, 'cap_evals', 0)),
                    f'level={sol.cap_value:.6g}'
                    if floor_fired
                    else f'cap={sol.cap_value:.3g}',
                )
    elif (
        getattr(sol, 't_events', None) is not None
        and len(sol.t_events) > 0
        and len(sol.t_events[0]) > 0
    ):
        t_event_phys = float(sol.t_events[0][0]) * t_ref
        event_label = getattr(events[0], 'binding_cap', None) if events else None
        if event_label is None:
            logger.info(
                'step cap fired at t=%.3e yr (terminal=True); '
                'binding cap unknown (EOS lookup failed on the terminating call)',
                t_event_phys,
            )
        else:
            cap_name = cap_display_names.get(event_label, 'ΔΦ_global')
            logger.log(
                logging.WARNING if event_label == 'eos_floor' else logging.INFO,
                '%s cap: scipy event fired at t=%.3e yr (terminal=True)%s',
                cap_name,
                t_event_phys,
                f'; level={floor["floor_level"]:.6g}' if event_label == 'eos_floor' else '',
            )

    # Diagnostic logging: internal BDF step statistics (in physical yr)
    if sol.t is not None and len(sol.t) > 1:
        dt_internal = np.diff(sol.t)
        cvode_nst = getattr(sol, 'cvode_nst', None)
        cvode_nfe = getattr(sol, 'cvode_nfe', None)
        if use_cvode and cvode_nst is not None:
            logger.info(
                'EntropySolver: %d internal steps, %d RHS evals, '
                'output grid dt_min=%.2e yr, dt_max=%.2e yr, dt_med=%.2e yr',
                cvode_nst,
                cvode_nfe,
                dt_internal.min(),
                dt_internal.max(),
                np.median(dt_internal),
            )
        else:
            logger.info(
                'EntropySolver: %d internal steps, %d RHS evals, '
                'dt_min=%.2e yr, dt_max=%.2e yr, dt_med=%.2e yr',
                len(sol.t),
                sol.nfev,
                dt_internal.min(),
                dt_internal.max(),
                np.median(dt_internal),
            )
        logger.info(
            'EntropySolver: phase-boundary cache hits=%d misses=%d',
            self.state._pb_cache_hits,
            self.state._pb_cache_misses,
        )
        self.state._pb_cache_hits = 0
        self.state._pb_cache_misses = 0

    self._log_solution_outcome(self.parameters.solver.end_time)
    if self.stop_early:
        sol.energy_integrals = dict(_ZERO_ENERGY_INTEGRALS)
        return
    budget = getattr(self, '_core_module_budget', None)
    if self._core_bc == 'core_module' and budget is not None:
        from aragog.core.regime import refuse_unmodelled_regime

        # every accepted step of the call, not only the output points
        t_core = (sol.get('energy_trace') or (sol.t, sol.y))[1][self._n_stag + 1]
        refuse_unmodelled_regime(budget, t_core)
    t_solved = time.perf_counter()
    sol.energy_integrals = self._compute_step_energy_integrals()
    logger.info(
        'EntropySolver: energy integrals in %.2f s (solve %.2f s)',
        time.perf_counter() - t_solved,
        t_solved - wall_start,
    )