proteus.interior_energetics.wrapper¶

def equilibrate_initial_state(dirs: dict, config: Config, hf_row: dict, outdir: str):
    """Iterate CALLIOPE + Zalmoxis until structure and composition converge.

    Runs volatile partitioning (CALLIOPE with optional binodal H2) followed
    by structure re-computation (Zalmoxis) in a loop. No SPIDER call.
    Convergence is checked on relative changes in R_int and P_surf.

    Called from proteus.py after the initial structure solve and volatile
    inventory setup. The converged state provides a self-consistent starting
    point for SPIDER's first time step.

    Parameters
    ----------
    dirs : dict
        Directory paths (including spider_mesh, spider_eos_dir).
    config : Config
        PROTEUS configuration.
    hf_row : dict
        Current helpfile row with volatile masses, R_int, P_surf, T_magma, etc.
    outdir : str
        Output directory for Zalmoxis files.
    """
    from proteus.interior_struct.zalmoxis import generate_spider_tables, zalmoxis_solver
    from proteus.outgas.wrapper import calc_target_elemental_inventories, run_outgassing

    max_iter = config.interior_struct.zalmoxis.equilibrate_max_iter
    tol = config.interior_struct.zalmoxis.equilibrate_tol
    nlev_b = get_nlevb(config)
    num_spider_nodes = nlev_b if config.interior_energetics.module == 'spider' else 0

    log.info(
        'Starting init equilibration loop (max %d iter, tol %.1f%%)',
        max_iter,
        tol * 100,
    )

    # Initialize convergence metrics (used in the warning if max_iter reached)
    delta_R = 1.0
    delta_P = 1.0

    # liquidus_super: re-solve the structure against the true super-liquidus
    # adiabat T(P), the same hand-off determine_interior_radius_with_zalmoxis
    # uses, so equilibration does not overwrite the isentropic, hydrostatically
    # consistent IC with the colder linear-in-r guess. The molten adiabat is set
    # by the surface superheat, not by the volatile composition this loop
    # iterates, so it is built once from the structure's current P_cmb. None =>
    # the linear guess (unchanged behaviour for SPIDER and non-liquidus_super
    # runs). P_cmb_grid records the upper pressure the grid was sized for so a
    # later iteration whose converged P_cmb exceeds it triggers a rebuild rather
    # than clipping the deep mantle to a flat T_cmb.
    adiabat_tfunc = None
    P_cmb_grid = 0.0
    use_adiabat = _use_superliquidus_adiabat_ic(config)
    if use_adiabat:
        P_cmb_seed = float(hf_row.get('P_cmb', 0.0) or 0.0)
        if P_cmb_seed > 0.0:
            P_cmb_grid = P_cmb_seed * _ADIABAT_IC_PCMB_MARGIN
            built = _build_superliquidus_adiabat_tp(config, hf_row, P_cmb_grid)
            if built is not None:
                adiabat_tfunc = built[0]

    for i in range(max_iter):
        R_old = float(hf_row.get('R_int', 0.0))
        P_old = float(hf_row.get('P_surf', 0.0))

        # 1. Volatile partitioning: recompute elemental targets and
        #    run CALLIOPE to get atmosphere/melt distribution
        calc_target_elemental_inventories(dirs, config, hf_row)
        run_outgassing(dirs, config, hf_row, initial=True)

        # 2. Re-compute structure with updated composition (volatile_profile is
        #    built inside zalmoxis_solver from hf_row). When the super-liquidus
        #    adiabat is available, the structure is solved against it with a
        #    mass-anchor guard and rollback; on failure the iteration degrades
        #    to the linear-guess solve so init never aborts where the linear
        #    path would have converged. SPIDER and non-liquidus_super runs take
        #    the linear-guess solve unchanged.
        if adiabat_tfunc is not None:
            spider_mesh_file, ok = _solve_structure_with_adiabat_or_rollback(
                config,
                outdir,
                hf_row,
                num_spider_nodes,
                adiabat_tfunc,
            )
            if not ok:
                log.warning(
                    'Equilibration iter %d/%d: adiabat re-solve failed; '
                    'falling back to the linear-guess solve for this iteration.',
                    i + 1,
                    max_iter,
                )
                _cmb_radius, spider_mesh_file = zalmoxis_solver(
                    config,
                    outdir,
                    hf_row,
                    num_spider_nodes=num_spider_nodes,
                    temperature_function=None,
                )
        else:
            _cmb_radius, spider_mesh_file = zalmoxis_solver(
                config,
                outdir,
                hf_row,
                num_spider_nodes=num_spider_nodes,
                temperature_function=None,
            )

        # Update M_mantle from Zalmoxis results (M_int and M_core are set
        # by zalmoxis_solver, but M_mantle is not). run_outgassing needs
        # an up-to-date M_mantle for dissolved fraction calculations.
        hf_row['M_mantle'] = float(hf_row.get('M_int', 0.0)) - float(hf_row.get('M_core', 0.0))

        # If this iteration's converged P_cmb now exceeds the pressure span the
        # adiabat grid was tabulated over, rebuild the adiabat so the deepest
        # mantle is not clipped to a flat T_cmb on the next iteration.
        if adiabat_tfunc is not None:
            P_cmb_now = float(hf_row.get('P_cmb', 0.0) or 0.0)
            if P_cmb_now > P_cmb_grid > 0.0:
                P_cmb_grid_new = P_cmb_now * _ADIABAT_IC_PCMB_MARGIN
                log.info(
                    'Equilibration iter %d/%d: converged P_cmb=%.1f GPa exceeds '
                    'the adiabat grid ceiling %.1f GPa; rebuilding the adiabat '
                    'to %.1f GPa.',
                    i + 1,
                    max_iter,
                    P_cmb_now / 1e9,
                    P_cmb_grid / 1e9,
                    P_cmb_grid_new / 1e9,
                )
                rebuilt = _build_superliquidus_adiabat_tp(config, hf_row, P_cmb_grid_new)
                if rebuilt is not None:
                    adiabat_tfunc = rebuilt[0]
                    P_cmb_grid = P_cmb_grid_new

        # Update mesh path if written
        if spider_mesh_file:
            dirs['spider_mesh'] = spider_mesh_file
            prev_path = spider_mesh_file + '.prev'
            shutil.copy2(spider_mesh_file, prev_path)
            dirs['spider_mesh_prev'] = prev_path

        # 3. Check convergence
        R_new = float(hf_row.get('R_int', 0.0))
        P_new = float(hf_row.get('P_surf', 0.0))

        delta_R = abs(R_new - R_old) / R_old if R_old > 0 else 1.0
        delta_P = abs(P_new - P_old) / P_old if P_old > 0 else 1.0

        log.info(
            'Equilibration iter %d/%d: dR/R=%.4f, dP/P=%.4f (R=%.3e m, P=%.2f bar)',
            i + 1,
            max_iter,
            delta_R,
            delta_P,
            R_new,
            P_new,
        )

        if delta_R < tol and delta_P < tol:
            log.info('Equilibration converged after %d iterations', i + 1)
            break
    else:
        log.warning(
            'Equilibration did not converge after %d iterations (dR/R=%.4f, dP/P=%.4f)',
            max_iter,
            delta_R,
            delta_P,
        )

    # 4. Regenerate SPIDER EOS tables with final composition
    if config.interior_energetics.module in ('spider', 'aragog'):
        spider_tables = generate_spider_tables(config, outdir)
        if spider_tables is not None:
            dirs['spider_eos_dir'] = spider_tables['eos_dir']
            dirs['spider_solidus_ps'] = spider_tables['solidus_path']
            dirs['spider_liquidus_ps'] = spider_tables['liquidus_path']

def solve_structure(
    dirs: dict,
    config: Config,
    hf_all: pd.DataFrame,
    hf_row: dict,
    outdir: str,
    thermal_solve: bool = True,
):
    """Solve for the planet structure based on the method set in the configuration file.

    If the structure is set by the radius, then this is trivial because the radius is used
    as an input to the interior modules anyway. If the structure is set by mass, then it is
    solved as an inverse problem for now.

    Parameters
    ----------
    dirs : dict
        Directories dictionary.
    config : Config
        Model configuration.
    hf_all : pd.DataFrame
        Historical helpfile dataframe.
    hf_row : dict
        Current step helpfile row.
    outdir : str
        Output directory path.
    thermal_solve : bool, optional
        Whether to run an interior thermal solve or update mechanical mass and
        structure only (default is True).
    """
    # Set by total mass (mantle + core + volatiles)
    if config.planet.mass_tot is not None:
        # Choose the method to determine the interior radius
        match config.interior_struct.module:
            case 'dummy':
                return determine_interior_radius_with_dummy(
                    dirs, config, hf_all, hf_row, outdir, thermal_solve=thermal_solve
                )
            case 'spider':
                if not thermal_solve:
                    raise ValueError(
                        "interior_struct.module = 'spider' does not support thermal_solve=False"
                    )
                return determine_interior_radius(dirs, config, hf_all, hf_row, outdir)
            case 'zalmoxis':
                # Zalmoxis computes its own radius; temporarily disable orbital
                # feedback during structure solve (restored after return).
                _orig_orbit_module = config.orbit.module
                config.orbit.module = 'dummy'
                try:
                    if config.params.stop.solid.phi_crit < 0.01:
                        log.warning(
                            'phi_crit=%.4f is below 0.01. Zalmoxis cases may plateau '
                            'at ~0.9%% melt fraction, so phi_crit < 0.01 can prevent '
                            'the simulation from terminating. Consider phi_crit >= 0.01.',
                            config.params.stop.solid.phi_crit,
                        )
                    return determine_interior_radius_with_zalmoxis(
                        dirs, config, hf_all, hf_row, outdir, thermal_solve=thermal_solve
                    )
                finally:
                    config.orbit.module = _orig_orbit_module
        raise ValueError(
            f"Invalid structure interior module selected '{config.interior_struct.module}'"
        )

    else:
        raise ValueError('planet.mass_tot must be set to solve for the interior structure')