Zalmoxis.solver

def main(
    config_params,
    material_dictionaries,
    melting_curves_functions,
    input_dir,
    layer_mixtures=None,
    volatile_profile=None,
    temperature_function=None,
    temperature_arrays=None,
    p_center_hint=None,
    initial_density=None,
    initial_radii=None,
):
    """Run the exoplanet internal structure model with automatic retry.

    Calls `_solve()` with mass-adaptive default parameters. If the first
    attempt does not converge, automatically retries once with tighter
    tolerances and doubled iteration limits.

    Parameters
    ----------
    config_params : dict
        Configuration parameters for the model. Numerical solver
        parameters (tolerances, iteration limits, step sizes) are
        optional; sensible mass-adaptive defaults are used when absent.
    material_dictionaries : dict
        EOS registry dict keyed by EOS identifier string.
    melting_curves_functions : tuple or None
        (solidus_func, liquidus_func) for EOS needing external melting curves.
    input_dir : str
        Directory containing input files.
    layer_mixtures : dict or None, optional
        Per-layer LayerMixture objects. If None, parsed from
        ``config_params['layer_eos_config']``. PROTEUS/CALLIOPE can provide
        pre-built mixtures with runtime-updated fractions.
    volatile_profile : VolatileProfile or None, optional
        Volatile profile for binodal-aware structure. Passed through to
        ``_solve()`` but not used directly by the retry logic.
    temperature_function : callable or None, optional
        External temperature function with signature ``f(r, P) -> T``
        where ``r`` is radius in m, ``P`` is pressure in Pa, and ``T``
        is temperature in K. When provided, bypasses the internal
        temperature mode dispatch (isothermal/linear/prescribed/adiabatic).
        Used by PROTEUS to pass SPIDER/Aragog T(r) profiles directly
        in memory.
    temperature_arrays : tuple[ndarray, ndarray] or None, optional
        Explicit r-indexed T profile ``(r_arr, T_arr)``. Only consumed
        by the JAX path (``config_params['use_jax']=True``). Preferred
        over ``temperature_function`` when the caller's T is naturally
        r-indexed (e.g. SPIDER/Aragog-coupled runs): the P-indexed
        tabulation inside ``jax_eos.wrapper`` collapses to a constant
        for P-ignoring callables and breaks convergence. See
        ``tools/benchmarks/bench_coupled_tempfunc.py``.
    initial_density : numpy.ndarray or None, optional
        Density profile from a previous solve, used to seed the Picard
        iteration. Must be paired with ``initial_radii``. Interpolated
        onto the current radial grid. Dramatically accelerates convergence
        when the structure changes incrementally between coupling steps.
    initial_radii : numpy.ndarray or None, optional
        Radial grid corresponding to ``initial_density``.

    Returns
    -------
    dict
        Model results including radii, density, gravity, pressure, temperature,
        mass enclosed, convergence status, and timing.
    """
    # Validate outer-solver choice. Default is 'picard' (the damped
    # fixed-point loop inside `_solve()`); 'newton' dispatches to
    # `_solve_newton_outer()`.
    outer_solver_explicit = 'outer_solver' in config_params
    outer_solver = config_params.get('outer_solver', 'picard')
    if outer_solver not in _VALID_OUTER_SOLVERS:
        raise ValueError(
            f'outer_solver must be one of {_VALID_OUTER_SOLVERS!r}, got {outer_solver!r}.'
        )

    # Advisory: when outer_solver is implicitly picard AND the profile
    # is hot fully-molten (T_surf > 3000 K) or super-Earth scale
    # (planet_mass > 2 M_E), log a one-line INFO suggesting Newton.
    # Damped Picard has a basin attractor on hot fully-molten and
    # super-Earth-scale profiles where the outer R-search oscillates
    # without converging; Newton + brentq is robust across 1-10 M_E.
    # The default stays at picard so that opting into Newton (which
    # requires tighter integrator tolerances) is an explicit user choice.
    if not outer_solver_explicit:
        _planet_mass = float(config_params.get('planet_mass') or 0.0)
        _surface_temperature = float(config_params.get('surface_temperature') or 0.0)
        _is_hot = _surface_temperature > 3000.0
        _is_massive = _planet_mass > 2.0 * earth_mass
        if _is_hot or _is_massive:
            _trigger = []
            if _is_massive:
                _trigger.append(f'planet_mass={_planet_mass / earth_mass:.2f} M_E > 2 M_E')
            if _is_hot:
                _trigger.append(f'surface_temperature={_surface_temperature:.0f} K > 3000 K')
            logger.info(
                "Zalmoxis using outer_solver='picard' (default). "
                "For this regime (%s), 'newton' is recommended: damped "
                'Picard has a known basin attractor on hot fully-molten '
                "and super-Earth profiles. Set outer_solver='newton' "
                'in your config to opt in. This advisory does not '
                'change behaviour.',
                ', '.join(_trigger),
            )
    if outer_solver == 'newton':
        return _solve_newton_outer(
            config_params,
            material_dictionaries,
            melting_curves_functions,
            input_dir,
            layer_mixtures=layer_mixtures,
            volatile_profile=volatile_profile,
            temperature_function=temperature_function,
            temperature_arrays=temperature_arrays,
            p_center_hint=p_center_hint,
            initial_density=initial_density,
            initial_radii=initial_radii,
        )

    result = _solve(
        config_params,
        material_dictionaries,
        melting_curves_functions,
        input_dir,
        layer_mixtures=layer_mixtures,
        volatile_profile=volatile_profile,
        temperature_function=temperature_function,
        temperature_arrays=temperature_arrays,
        p_center_hint=p_center_hint,
        initial_density=initial_density,
        initial_radii=initial_radii,
    )

    # Skip retry when the first attempt is "good enough": density and
    # pressure converged, mass error within `_RETRY_SKIP_MASS_ERR`. On
    # hot fully-molten T(r) profiles (early CHILI coupled iters) the
    # retry path doubles every iteration cap and wall_timeout, takes
    # 600+ s, and routinely lands at the SAME or SLIGHTLY-WORSE mass
    # error as the first attempt. The outer mass-radius loop will
    # converge across PROTEUS iters anyway; one resolve ending at 5-7%
    # mass error is fine when the next resolve corrects it.
    _RETRY_SKIP_MASS_ERR = 0.07
    _first_mass_err = result.get('best_mass_error')
    _retry_skipped = (
        not result['converged']
        and result.get('converged_density', False)
        and result.get('converged_pressure', False)
        and _first_mass_err is not None
        and _first_mass_err < _RETRY_SKIP_MASS_ERR
    )
    if (
        _retry_skipped
    ):  # pragma: no cover - retry-skip diagnostic; reached only on borderline first-attempt
        logger.info(
            'Skipping retry: density and pressure converged, mass error '
            '%.4f%% < %.2f%% threshold. Accepting first-attempt solution.',
            _first_mass_err * 100,
            _RETRY_SKIP_MASS_ERR * 100,
        )
        result['converged'] = True
        result['converged_mass'] = True

    if not result['converged']:
        # Build tightened params for retry
        planet_mass = config_params['planet_mass']
        defaults = _default_solver_params(planet_mass)
        # Merge: explicit caller values override defaults
        effective = dict(defaults)
        for key in defaults:
            if key in config_params:
                effective[key] = config_params[key]
        tightened = _tighten_solver_params(effective)

        logger.warning(
            'Structure solve did not converge (mass=%s, density=%s, pressure=%s). '
            'Retrying with tighter parameters.',
            result.get('converged_mass', False),
            result.get('converged_density', False),
            result.get('converged_pressure', False),
        )

        # Seed retry with the first attempt's results
        retry_params = copy.copy(config_params)
        retry_params.update(tightened)
        retry_density = None
        retry_radii = None
        if 'radii' in result and result['radii'] is not None and len(result['radii']) > 0:
            r_last = result['radii'][-1]
            mass_in_earth = max(planet_mass, 0.01 * earth_mass) / earth_mass
            r_seager = earth_radius * mass_in_earth**0.282
            if np.isfinite(r_last) and r_last > 0 and 0.2 * r_seager < r_last < 5.0 * r_seager:
                retry_params['_initial_radius_guess'] = r_last
            # Use first attempt's density as seed for retry
            if result.get('density') is not None and np.any(result['density'] > 0):
                retry_density = result['density']
                retry_radii = result['radii']

        result = _solve(
            retry_params,
            material_dictionaries,
            melting_curves_functions,
            input_dir,
            layer_mixtures=layer_mixtures,
            volatile_profile=volatile_profile,
            temperature_function=temperature_function,
            temperature_arrays=temperature_arrays,
            p_center_hint=p_center_hint,
            initial_density=retry_density,
            initial_radii=retry_radii,
        )

        if not result['converged']:  # pragma: no cover - retry also failed; defensive
            logger.warning(
                'Structure solve did not converge after retry. '
                'Returning best result with converged=False.'
            )

    return result