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: Simple core cooling
        2: Prescribed heat flux
        3: Prescribed temperature
    """
    if self._settings.inner_boundary_condition == 1:
        self.core_cooling(state)
    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 core_cooling(self, state: State) -> None:
    """Applies a core cooling heat flux according to Eq. (37) of Bower et al., 2018.

    The core is modelled as a well-mixed reservoir with an effective
    temperature T_core = tfac_core_avg * T_cmb. The factor tfac_core_avg
    accounts for the adiabatic temperature gradient within the core
    (mass-weighted average core temperature / CMB temperature).
    Default 1.147 is for Earth-like parameters (Bower+2018, Table 2).

    Parameters
    ----------
    state : State
        The state to apply the boundary condition to.
    """
    # Core thermal capacity: C_core = rho_core * cp_core * V_core
    r_cmb = np.asarray(self._mesh.basic.radii).flat[0]
    core_capacity = (
        4
        / 3
        * np.pi
        * r_cmb**3
        * self._mesh.settings.core_density
        * self._settings.core_heat_capacity
    )

    # First mantle cell thermal capacity: C_cell = rho * cp * V_cell
    cap_stag = state.capacitance_staggered()  # rho * cp, may be float or array
    cap_first = np.asarray(cap_stag).flat[0]
    cell_capacity = np.asarray(self._mesh.basic.volume).flat[0] * cap_first

    # Geometric correction: area ratio between first interior face and CMB
    r_above = np.asarray(self._mesh.basic.radii).flat[1]
    radius_ratio = r_above / r_cmb

    # Core buffering factor (Bower+2018 Eq. 37):
    # alpha = (R_1/R_0)^2 / (1 + C_cell / (C_core * tfac))
    # When C_core >> C_cell: alpha -> (R_1/R_0)^2 (core tracks mantle)
    # When C_core << C_cell: alpha -> 0 (core absorbs all heat)
    tfac = self._settings.tfac_core_avg
    alpha = radius_ratio**2 / (cell_capacity / (core_capacity * tfac) + 1)

    state.heat_flux[0, :] = alpha * state.heat_flux[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.

    For the ``bower2018`` core BC the state vector is
    ``[S_0, ..., S_{N-1}, T_core]`` of length N+1, and this returns
    ``[dS/dt, dT_core/dt]`` of the same length.

    For the quasi_steady BC the state vector is just
    ``[S_0, ..., S_{N-1}]`` of length N.

    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].
        Shape (N,) or (N+1,) only.

    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_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``).

    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.
    """
    n_stag = getattr(self, '_n_stag', None)
    prev_sol = getattr(self, '_solution', None)
    if (
        n_stag is None
        or prev_sol is None
        or getattr(prev_sol, 'y', None) is None
        or prev_sol.y.size == 0
        or prev_sol.y.shape[0] != n_stag + 1
    ):
        return None
    return float(prev_sol.y[n_stag, -1])

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.
    """
    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

    # Compute per-call energy integrals BEFORE refreshing state at
    # the final entropy. ``_compute_step_energy_integrals`` walks
    # the CVODE trajectory and calls state.update() at each
    # accepted sub-step; doing it before the final refresh keeps
    # the end-of-call snapshot (heat_flux, heating arrays, etc.)
    # consistent with what callers see in the rest of get_state().
    step_integrals = 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, S_basic_final = self._reconstruct_entropy(dSdr_final, S_surf_final)
        extra_final = None
    elif is_ext:
        S_final = sol.y[:n_stag, -1]
        extra_final = float(sol.y[n_stag, -1])
    else:
        S_final = sol.y[:, -1]
        extra_final = None

    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 for derived quantities.
    if gradient_mode:
        self.state.update(S_final, sol.t[-1], dSdr=dSdr_final, entropy_basic=S_basic_final)
    elif energy_balance:
        self.state.update(S_final, sol.t[-1], dSdr_cmb=extra_final)
    else:
        self.state.update(S_final, sol.t[-1])
    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()

    # Scalar quantities.
    # M_mantle uses the analytic A-W mass integral (matching
    # SPIDER's EOSAdamsWilliamson_GetMassWithinShell) when
    # eos_method=1. The discrete sum (rho_stag * vol) has O(h^2)
    # quadrature error vs the analytic integral, causing the
    # structure root finder to converge to a different R_int.
    # For eos_method=2 (external mesh), fall back to discrete sum.
    mesh = self.evaluator.mesh
    if hasattr(mesh.eos, 'get_mass_within_radii'):
        r_cmb = float(self._r_basic_flat[0])
        r_surf = float(self._r_basic_flat[-1])
        M_mantle = (
            mesh.eos.get_mass_within_radii(np.array([r_surf]))
            - mesh.eos.get_mass_within_radii(np.array([r_cmb]))
        ).item()
    else:
        rho_struct_stag = np.asarray(mesh.staggered_effective_density).ravel()
        mass_stag = rho_struct_stag * vol
        M_mantle = float(np.sum(mass_stag))
    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.
    if bower:
        T_core = extra_final  # bower2018: T_core integrated as ODE state
    else:
        T_core = float(T_stag[0])
    # 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: lower-boundary value of the basic-node heat
    # flux array. Sign convention follows the rest of Aragog,
    # positive-out-of-core when entering the mantle.
    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

    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(np.sum(phi_stag * mass_stag)),
        M_mantle_solid=float(M_mantle - np.sum(phi_stag * mass_stag)),
        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_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'],
        dt_actual=float(sol.t[-1] - sol.t[0]),
        status=sol.status,
        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,
    )

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')
    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()

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

    Must be called BEFORE ``set_initial_entropy``. If not called,
    the initial T_core defaults to the bottom-cell mantle
    temperature derived from S_init via the EOS.
    """
    self._T_core_init = float(T_core_init)

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

    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. 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}')

    # 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 = getattr(self, '_dSdr_cmb_init', None)
        if dSdr_cmb_init is None:
            # Hot start: preserve from previous solution if shape matches
            prev_sol = getattr(self, '_solution', None)
            if (
                prev_sol is not None
                and getattr(prev_sol, 'y', None) is not None
                and prev_sol.y.size > 0
                and prev_sol.y.shape[0] == n_stag + 1
            ):
                dSdr_cmb_init = float(prev_sol.y[n_stag, -1])
                logger.info(
                    'Preserved dSdr_cmb from previous solve: %.3e J/kg/K/m',
                    dSdr_cmb_init,
                )
        if dSdr_cmb_init is None:
            # Cold start: one-sided FD of S_init at the bottom.
            # dSdr_cmb ≈ (S_stag[1] - S_stag[0]) / (r_stag[1] - r_stag[0])
            # For a uniform S_init this is exactly zero, which
            # is the correct neutral-buoyancy starting point.
            if n_stag >= 2:
                r_basic = np.asarray(self.evaluator.mesh.basic.radii).ravel()
                r_stag_0 = 0.5 * (r_basic[0] + r_basic[1])
                r_stag_1 = 0.5 * (r_basic[1] + r_basic[2])
                dSdr_cmb_init = (float(S_arr[1]) - float(S_arr[0])) / max(
                    r_stag_1 - r_stag_0, 1.0
                )
            else:
                dSdr_cmb_init = 0.0
            logger.info(
                'Cold-start dSdr_cmb from FD: %.3e J/kg/K/m',
                dSdr_cmb_init,
            )
        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 == 'bower2018':
        # Core temperature as ODE state variable (conduction-only
        # flux). State = [S, T_core]. Available for parity testing
        # only; not recommended for production runs.
        T_core_init = getattr(self, '_T_core_init', None)
        if T_core_init is None:
            prev_sol = getattr(self, '_solution', None)
            if (
                prev_sol is not None
                and getattr(prev_sol, 'y', None) is not None
                and prev_sol.y.size > 0
                and prev_sol.y.shape[0] == n_stag + 1
            ):
                T_core_init = float(prev_sol.y[n_stag, -1])
        if T_core_init is None:
            P_bottom = float(self._P_stag_flat[0])
            T_core_init = float(
                np.asarray(
                    self.entropy_eos.temperature(np.array([P_bottom]), np.array([S_arr[0]]))
                ).item()
            )
        self._S0 = np.empty(n_stag + 1)
        self._S0[:n_stag] = S_arr
        self._S0[n_stag] = T_core_init
        logger.info(
            'Initial state (bower2018): S_min=%.0f, S_max=%.0f, T_core_init=%.0f K',
            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_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.
        May be None to disable the Option Z path even if the flag
        is on.
    """
    self._jax_cvode_factory = factory

def solve(self) -> None:
    """Run the BDF time integration."""
    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

    # phi_cap_anchor stores the (cap, phi0_global) tuple for the
    # SUNDIALS rootfn / scipy event, or None when the cap is not
    # armed. Initialised here so the consumer at the bottom of
    # solve() never hits an UnboundLocalError when the
    # ``phi0 > 0.01 and self.entropy_eos is not None`` branch
    # below is skipped (e.g. const_properties path with no EOS).
    phi_cap_anchor = None

    # Tighten max_step when ANY cell is near a phase boundary.
    # Extends the previous CMB-only check to all cells, catching
    # phase transitions that propagate through upper cells during
    # the rheological-front sweep, not just CMB crystallisation.
    # When a cell's entropy is within 200 J/kg/K of either the
    # solidus or liquidus, OR sits inside the mushy band, max_step
    # is reduced to 1 yr to give CVODE enough resolution to handle
    # the phase-boundary stiffness gradually.
    if phi0 > 0.01 and self.entropy_eos is not None:
        P_basic = self._P_basic_flat
        np.asarray(self.entropy_eos.liquidus_entropy(P_basic)).ravel()
        np.asarray(self.entropy_eos.solidus_entropy(P_basic)).ravel()
        # 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) < 200.0)
        near_sol = np.any(np.abs(margin_to_sol) < 200.0)
        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])
        margin = S0_block_cmb - S_liq
        if (
            near_liq
            or near_sol
            or in_mushy
            or margin < 200.0
            or (S0_block_cmb < S_liq and S0_block_cmb > S_sol)
        ):
            max_step = 1.0

        # Per-call mass-weighted ΔΦ_global cap as a SUNDIALS root
        # function. Build the rootfn metadata here (anchor:
        # Φ_global at solve entry); the actual rootfn instance is
        # constructed below once the nondim state_scale is known.
        # CVODE returns at the exact time t* where
        # |Φ_global(t*) − Φ_global(start)| = cap, so the cap is
        # enforced by the integrator's own trajectory rather than
        # by a start-time rate estimate. A rate-estimate cap can
        # truncate end_time and lock PROTEUS's adaptive dt onto
        # the truncated value at the rheological transition.
        phi_step_cap = float(getattr(self.parameters.energy, 'phi_step_cap', 0.0))
        if phi_step_cap > 0.0 and (near_liq or near_sol or in_mushy):
            in_mushy_arr = (margin_to_liq < 0.0) & (margin_to_sol > 0.0)
            if np.any(in_mushy_arr):
                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
                        )
                        phi_cap_anchor = (phi_step_cap, phi0_global)
                        logger.info(
                            'ΔΦ_global cap %.3g: '
                            'arming CVODE rootfn anchored at '
                            'Φ_global(start)=%.4f',
                            phi_step_cap,
                            phi0_global,
                        )
                except Exception as exc:
                    logger.warning(
                        'ΔΦ_global cap rootfn anchor failed (%s); proceeding without cap',
                        exc,
                    )

    # 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()
    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 / _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,
        )
        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()

    # ── ΔΦ_global 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 Φ_global(start).
    events = None
    phi_cap_rootfn = None
    if phi_cap_anchor is not None:
        cap_value, phi0_global = phi_cap_anchor
        try:
            phi_cap_rootfn = _PhiCapRootFunction(
                eos=self.entropy_eos,
                P_stag=self._P_stag_flat,
                volume=self._volume_flat,
                n_stag=self._n_stag,
                phi0_global=phi0_global,
                cap=cap_value,
                state_scale=_state_scale,
            )
            events = [
                _phi_cap_event_factory(
                    eos=self.entropy_eos,
                    P_stag=self._P_stag_flat,
                    volume=self._volume_flat,
                    n_stag=self._n_stag,
                    phi0_global=phi0_global,
                    cap=cap_value,
                    state_scale=_state_scale,
                )
            ]
        except Exception as exc:
            logger.warning(
                'ΔΦ_global 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
    if use_cvode:
        # Option Z: build JAX-derived CVODE callbacks when the
        # factory is registered AND the config flag is on. The
        # factory re-creates the (rhs_fn, jac_fn) pair per solve
        # call (JIT recompile cost applies; the same JAX tracing
        # cache is hit on matching signatures, so cost collapses
        # after the first call within a Python process lifetime).
        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,
                )
                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)')
        self._solution = self._solve_cvode(
            start_time=start_nd,
            end_time=end_nd,
            y0=S0_nd,
            atol=atol_nd,
            rtol=rtol,
            max_step=max_step_nd,
            rhs=_rhs_nondim,
            cvode_rhs_fn_override=cvode_rhs_override,
            cvode_jacfn=cvode_jacfn,
            phi_cap_rootfn=phi_cap_rootfn,
        )
    else:
        method = 'Radau' if solver_method != 'bdf' else 'BDF'
        logger.info('EntropySolver: using scipy %s', method)
        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]
        else:
            sol.y = sol_y * _state_scale

    # ΔΦ_global cap-fire log, emitted with PHYSICAL time. The
    # CVODE path attaches ``cap_fired`` etc. to the result inside
    # ``_solve_cvode``; the scipy fallback exposes the same signal
    # via ``sol.t_events`` (non-empty when the terminal event
    # fired). Logging here (after the t_ref restoration) keeps
    # operator-facing messages in physical years rather than the
    # nondim time the rootfn sees internally.
    if getattr(sol, 'cap_fired', False) and sol.t is not None:
        t_root_phys = float(sol.t[-1])
        logger.info(
            'ΔΦ_global cap: CVODE rootfn fired at '
            't=%.3e yr after %d evals; cap=%.3g, '
            'Φ_global(start)=%.4f',
            t_root_phys,
            int(getattr(sol, 'cap_evals', 0)),
            float(getattr(sol, 'cap_value', 0.0)),
            float(getattr(sol, 'cap_phi0', 0.0)),
        )
    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
        logger.info(
            'ΔΦ_global cap: scipy event fired at t=%.3e yr (terminal=True)',
            t_event_phys,
        )

    # 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)
        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

    if self._solution.status == 0:
        logger.info('EntropySolver: integration completed successfully.')
        self.stop_early = False
    elif self._solution.status == 1:
        # Termination event (liquidus crossing at CMB cell).
        # Integration succeeded up to the event time.
        t_event = self._solution.t[-1]
        logger.info(
            'EntropySolver: liquidus-crossing event at t=%.2e yr '
            '(stopped %.1f yr before end_time). Bottom cell reached '
            'onset of crystallization.',
            t_event,
            end_time - t_event,
        )
        self.stop_early = False
    else:
        logger.error(
            'EntropySolver: integration failed (status=%d): %s',
            self._solution.status,
            self._solution.message,
        )
        self.stop_early = True

def write_netcdf(
    self,
    path: str | Path,
    *,
    time: float | None = None,
    description: str | None = None,
) -> None:
    """Convenience wrapper: ``self.get_state().to_netcdf(path)``.

    Designed for the standalone solver script: build the solver,
    call ``solve()``, then dump the final state as a NetCDF4 file.

    Parameters
    ----------
    path : str or Path
        Destination NetCDF4 file (overwrites if it exists).
    time : float, optional
        Simulation time at which this snapshot was taken [yr].
        Defaults to ``SolverOutput.dt_actual`` if omitted.
    description : str, optional
        Free-form description string written as the dataset's
        ``description`` global attribute.

    See Also
    --------
    SolverOutput.to_netcdf : underlying writer with the full schema.
    """
    # Forward only when the caller supplied a description; otherwise
    # let ``to_netcdf``'s default take effect so the two entry points
    # stamp the same string into ``ds.description``.
    kwargs: dict[str, str] = {}
    if description is not None:
        kwargs['description'] = description
    self.get_state().to_netcdf(path, time=time, **kwargs)