diff --git a/src/proteus/config/_params.py b/src/proteus/config/_params.py index 88ca4d95d..3bb94eb2c 100644 --- a/src/proteus/config/_params.py +++ b/src/proteus/config/_params.py @@ -141,6 +141,16 @@ class TimeStepParams: Replacement speed-up factor applied while the hysteresis counter is active. Must be ``>= 1.0`` and ``<= SFINC`` (1.6). Default 1.1 (gentle ramp-up). + impact_maximum: float + Maximum time-step size [yr] for the step that lands on a + scheduled giant impact. The landing step is otherwise + clamped only to however much simulated time remains before + the impact, so after a long quiescent phase has let ``dt`` + coarsen, that remaining time can itself be large and the + step absorbing the impact's melt-fraction jump inherits the + same coarseness. Set to 0 (default) to disable, in which + case the remaining-time clamp applies with no independent + ceiling. """ starspec: float = field(default=1e8, validator=ge(0)) @@ -170,6 +180,7 @@ class TimeStepParams: mushy_upper: float = field(default=0.99, validator=(gt(0), lt(1))) hysteresis_iters: int = field(default=0, validator=ge(0)) hysteresis_sfinc: float = field(default=1.1, validator=ge(1.0)) + impact_maximum: float = field(default=0.0, validator=ge(0)) # Cap on dt growth ratio between consecutive steps. Bounds # dtswitch / dtprev to at most max_growth_factor, preventing diff --git a/src/proteus/interior_energetics/aragog.py b/src/proteus/interior_energetics/aragog.py index 6d4556816..496895202 100644 --- a/src/proteus/interior_energetics/aragog.py +++ b/src/proteus/interior_energetics/aragog.py @@ -62,13 +62,23 @@ def _eos_content_key(eos_dir_str: str) -> str: The PROTEUS test fixture materialises the EOS tables into a fresh per-test ``outdir/data/spider_eos`` directory each time, so a path - based cache key misses across tests. The content fingerprint is a - sorted tuple of ``(filename, file size)`` pairs for every regular - file in the directory; it is stable across distinct on-disk copies - of the same tables but cheap to compute (one ``os.listdir`` + one - ``getsize`` per file). + based cache key misses across tests. + + The generator writes the parameters that define the tables into + ``.cache_info.txt``: the pressure ceiling, the grid shape, the mushy-zone + factor and the EOS identity. That marker is the key when present. Sizes + alone are not enough on the accretion path: a giant impact grows the planet + and the tables are rewritten to a higher pressure ceiling on the same grid, + so every file keeps its length and a size-based key cannot see that the + tables now describe a different planet. """ try: + marker = os.path.join(eos_dir_str, '.cache_info.txt') + if os.path.isfile(marker): + with open(marker) as f: + key = f.read().strip() + if key: + return key pairs = [] for name in sorted(os.listdir(eos_dir_str)): full = os.path.join(eos_dir_str, name) @@ -586,6 +596,10 @@ def setup_or_update_solver( else: AragogRunner.update_structure(config, hf_row, interior_o) AragogRunner.update_solver(dt, hf_row, interior_o) + # Refresh before reset(): the compression-work diagnostic inside + # reset() evaluates the new mesh pressures against the installed + # table, which clamps at a stale ceiling on impact steps. + AragogRunner._refresh_entropy_eos(config, interior_o) interior_o.aragog_solver.reset() # Restore entropy IC from previous solve if hasattr(interior_o, '_last_entropy') and interior_o._last_entropy is not None: @@ -1223,6 +1237,34 @@ def _append_radnuc(_iso, _cnc): _t_post_solver - _t_post_eos, ) + @staticmethod + def _refresh_entropy_eos(config: Config, interior_o: Interior_t) -> None: + """Point the solver at the P-S tables as they stand now. + + The solver keeps whatever table object it was built with, and the tables + are rewritten whenever the structure solve reruns, with a pressure + ceiling that grows with the planet. ``solver.entropy_eos`` is read live + throughout the solve (the RHS, the per-call energy integrals, and the + compression-work diagnostic inside ``reset()``), so a stale object + misreports the energy budget on exactly the runs that outgrow their + starting table. The loader is cached on the table parameters, so an + unchanged table costs one small read. + + Parameters + ---------- + config : Config + PROTEUS configuration; a const-properties run carries no tables. + interior_o : Interior_t + Interior state holding the live Aragog solver. + """ + if config.interior_energetics.const_properties: + return + solver = getattr(interior_o, 'aragog_solver', None) + eos_dir = getattr(interior_o, '_spider_eos_dir', '') + if solver is None or not eos_dir or not os.path.isdir(str(eos_dir)): + return + solver.entropy_eos = _cached_entropy_eos(str(eos_dir)) + @staticmethod def _maybe_install_jax_cvode_factory(config: Config, interior_o: Interior_t) -> None: """Install a JAX CVODE callback factory on the solver (option Z). @@ -1280,7 +1322,9 @@ def _maybe_install_jax_cvode_factory(config: Config, interior_o: Interior_t) -> try: eos_dir = interior_o._spider_eos_dir _t_pre_jax_eos = time.perf_counter() - eos_jax = _cached_entropy_eos_jax(str(eos_dir)) + # Build once here so an unreadable EOS directory fails the install + # rather than the first solve. The factory reloads it per call. + _cached_entropy_eos_jax(str(eos_dir)) _t_post_jax_eos = time.perf_counter() if nightly_strict: log.info( @@ -1313,14 +1357,10 @@ def _maybe_install_jax_cvode_factory(config: Config, interior_o: Interior_t) -> phase_smoothing_width=0.01, ) - _t_pre_mesh = time.perf_counter() - mesh_jax = MeshArrays.from_numpy_mesh(solver.evaluator.mesh) - _t_post_mesh = time.perf_counter() - n_stag = solver._n_stag if nightly_strict: log.info( - 'aragog diag: jax_cvode_factory phases params_jax+mesh=%.2fs', - _t_post_mesh - _t_post_jax_eos, + 'aragog diag: jax_cvode_factory phases params_jax=%.2fs', + time.perf_counter() - _t_post_jax_eos, ) def factory(scales, core_bc_mode): @@ -1329,6 +1369,16 @@ def factory(scales, core_bc_mode): # consumed the analytic Jacobian rather than silently falling # back to the FD path. solver._jax_factory_call_count += 1 + # Rebuild the mesh from live solver state every solve() call: + # impacts and structure re-solves replace it, and a copy taken + # at install time would keep integrating the pre-change planet. + mesh_jax = MeshArrays.from_numpy_mesh(solver.evaluator.mesh) + n_stag = solver._n_stag + # Same for the EOS tables: regeneration raises their pressure + # ceiling with the planet's mass, and an install-time copy would + # clamp the deep mantle at the smaller planet's table edge. The + # loader is cached, so an unchanged table costs one small read. + eos_jax = _cached_entropy_eos_jax(str(interior_o._spider_eos_dir)) # ``scales`` is an aragog.jax.nondim.NonDimScales single # source of truth. # Rebuild BoundaryParams from live solver state every @@ -1400,17 +1450,27 @@ def factory(scales, core_bc_mode): ) return rhs_fn, jac_fn - solver.set_jax_cvode_factory(factory) - log.info( + # Read the diagnostic geometry before installing, so that the + # install is the last thing here that can fail. Anything raising + # after it would send a working factory into the handler below. + r_basic = np.asarray(solver._r_basic_flat).ravel() + installed = ( 'Option Z: JAX CVODE factory installed on aragog solver ' - '(core_bc=%s, n_stag=%d).', - solver._core_bc, - n_stag, + f'(core_bc={solver._core_bc}, n_stag={int(solver._n_stag)}, ' + f'r_cmb={float(r_basic[0]):.6e} m, r_surf={float(r_basic[-1]):.6e} m). ' + 'The mesh is read from the solver on every solve, so this is the ' + 'geometry at install time, not for the run.' ) + solver.set_jax_cvode_factory(factory) + log.info(installed) except Exception as exc: msg = f'Option Z factory install failed ({exc}); falling back to FD Jacobian.' if nightly_strict: raise RuntimeError(msg) from exc + # Leave nothing half-installed: the solve-time check is only that a + # factory is present, so a partial install would run this path on + # state the failure above left incomplete. + solver.set_jax_cvode_factory(None) log.warning(msg) @staticmethod @@ -2065,6 +2125,15 @@ def _solve_with_retry(self, hf_row, interior_o) -> SolverOutput: sanity_dT_core = max( 3000.0, 1500.0 * mass_tot ) # max plausible T_core change per retry [K] + # A giant impact re-melts the mantle between solves, so the T_core jump + # it produces is real and is identical at every step size. Retrying + # cannot shrink it, so the guard would spend the whole ladder and kill + # the run. Skip it on that one step; every other step keeps it. + impact_step = bool(getattr(interior_o, 'impact_reset_this_step', False)) + + # Immediately before the solve, so the state-heat integral this step + # books is taken against the tables the step actually runs on. + AragogRunner._refresh_entropy_eos(self._config, interior_o) # Capture IC for restoration on retry, and pre-call T_core for # the sanity check on retry success. @@ -2153,7 +2222,15 @@ def _solve_with_retry(self, hf_row, interior_o) -> SolverOutput: # converged core temperature exists to compare against, so # the jump guard is necessarily inactive on that one step. dT = abs(T_core_post - T_core_pre) if T_core_pre > 0 else 0.0 - if dT > sanity_dT_core: + if dT > sanity_dT_core and impact_step: + log.info( + 'T_core jumped %.1f K (>%.0f K threshold) on the ' + 'step a giant impact re-melted the mantle. The ' + 'jump is the impact, so the guard is skipped here.', + dT, + sanity_dT_core, + ) + if dT > sanity_dT_core and not impact_step: log.warning( 'Aragog attempt %d returned status=0 but T_core ' 'jumped %.1f K (>%.0f K threshold). Treating as ' diff --git a/src/proteus/interior_energetics/common.py b/src/proteus/interior_energetics/common.py index 6073b5ce1..e403bebf5 100644 --- a/src/proteus/interior_energetics/common.py +++ b/src/proteus/interior_energetics/common.py @@ -615,6 +615,12 @@ def __init__(self, nlev_b: int, spider_dir=None, eos_dir=None): # solver anomaly. Consumed and cleared on that one step. self.impact_reset = False + # ``impact_reset`` as read at the top of this step, kept readable for + # the whole step because the flag itself is cleared before the interior + # solver runs. The solvers' jump guards read this to tell an impact's + # deliberate temperature step from a corrupted solve. + self.impact_reset_this_step = False + # True when the most recent call to next_step() had its step size # clamped. For example, by `_estimate_bolscale()`. self.timestep_clamped = False diff --git a/src/proteus/interior_energetics/timestep.py b/src/proteus/interior_energetics/timestep.py index c9df81b85..a8503a955 100644 --- a/src/proteus/interior_energetics/timestep.py +++ b/src/proteus/interior_energetics/timestep.py @@ -413,15 +413,16 @@ def next_step( ) dtswitch = dt_capped - # Land exactly on the next scheduled giant impact. An impact re-melts - # the mantle and grows the planet, so it has to be applied at the - # state the timeline says it happens at, not at whatever state a step - # that jumped over it produced. The clamp only ever shortens dt, and - # it is floored at the minimum step so a nearby impact cannot collapse - # dt to zero; the event handler uses a half-open time window, so an - # impact inside a floored step is still applied exactly once. + # Land exactly on the next scheduled giant impact, since it remelts the + # mantle and grows the planet: the state must match the timeline, not + # whatever a step that overshot it produced. The clamp only shortens + # dt and floors at the minimum step; impact_maximum further bounds the + # landing step independent of how coarse dt had grown beforehand. if interior_o is not None and np.isfinite(interior_o.t_next_impact): dt_to_impact = interior_o.t_next_impact - hf_row['Time'] + impact_ceiling = float(config.params.dt.impact_maximum) + if impact_ceiling > 0.0: + dt_to_impact = min(dt_to_impact, impact_ceiling) dtfloor = config.params.dt.minimum + config.params.dt.minimum_rel * hf_row['Time'] dt_to_impact = max(dt_to_impact, dtfloor) if dtswitch > dt_to_impact: diff --git a/src/proteus/interior_energetics/wrapper.py b/src/proteus/interior_energetics/wrapper.py index 1d37f18dd..41d80431a 100644 --- a/src/proteus/interior_energetics/wrapper.py +++ b/src/proteus/interior_energetics/wrapper.py @@ -2165,6 +2165,9 @@ def run_interior( # cannot wrongly suppress the clip on a later, ordinary step. impact_reset = getattr(interior_o, 'impact_reset', False) interior_o.impact_reset = False + # The interior solvers run below, after the flag is cleared, so keep the + # value readable for the rest of this step. + interior_o.impact_reset_this_step = impact_reset # Write tidal heating file if config.interior_energetics.heat_tidal: @@ -3339,17 +3342,15 @@ def temperature_function(r, P): # MgSiO3 is a planet-state-invariant material EOS, so the pre-built # tables are stable for the entire evolution. The comp_changed path # is reached in wet runs where binodal redistribution or degassing - # shifts mantle volatile fractions by > 5% (SPIDER reads the fresh - # file on next call; Aragog's in-memory EntropyEOS, built once during - # AragogRunner.setup_solver, is NOT invalidated here, so Aragog would - # silently use the stale in-memory tables). + # shifts mantle volatile fractions by > 5%. SPIDER reads the fresh file on + # its next call. Aragog's JAX right-hand side reloads the tables per solve + # through a loader cached on the table parameters, so it follows them here + # and at the ungated regeneration a giant impact triggers. # - # KNOWN GAP: for Aragog + wet runs we would need to (i) reload - # EntropyEOS from the regenerated files, (ii) re-install the JAX - # CVODE factory so its captured eos_jax pytree matches the new - # tables, (iii) bounds-check the cached _last_entropy against the - # new [S_min, S_max] range. Dry runs do not need this; it is a - # precondition for quantitative wet-run work. + # REMAINING GAP: the cached _last_entropy is not bounds-checked against the + # regenerated [S_min, S_max]. Both a composition change and a raised + # pressure ceiling can move that range (it is scanned up to the ceiling), + # so out-of-range carried entropy clamps at the table edge in the solve. if comp_changed and config.interior_energetics.module in ('spider', 'aragog'): from proteus.interior_struct.zalmoxis import generate_spider_tables @@ -3360,10 +3361,10 @@ def temperature_function(r, P): dirs['spider_liquidus_ps'] = spider_tables['liquidus_path'] log.info('Regenerated SPIDER EOS tables (composition change)') if config.interior_energetics.module == 'aragog': - log.warning( - 'Aragog: regenerated P-S tables on composition change, ' - 'but Aragog in-memory EntropyEOS is not refreshed. ' - 'Known gap for wet runs. Dry runs are not affected.' + log.info( + 'Aragog reloads the regenerated tables on its next solve; ' + 'the entropy carried over from the previous step is not ' + 'bounds-checked against their new range.' ) # Update composition sentinels for next trigger check diff --git a/tests/interior_energetics/test_aragog.py b/tests/interior_energetics/test_aragog.py index 79df4a96a..545df2d18 100644 --- a/tests/interior_energetics/test_aragog.py +++ b/tests/interior_energetics/test_aragog.py @@ -3,7 +3,10 @@ Tests the Zalmoxis-specific branches in AragogRunner.setup_solver() that set inner_radius from zalmoxis_solver and configure temperature-dependent initial -conditions. +conditions, plus the contracts that keep the solver on the planet as it grows: +the retry ladder and its giant-impact exemption, per-solve mesh re-reads, the +EOS-table reload and its content-keyed cache, and the JAX factory install and +failure paths. Testing standards and documentation: - docs/How-to/testing.md: Running, writing, and marking tests; coverage and CI @@ -11,6 +14,9 @@ Functions tested: - AragogRunner.setup_solver(): Zalmoxis branches for inner_radius, EOS fallback +- AragogRunner._solve_with_retry(): ladder policy, guards, table refresh wiring +- AragogRunner._refresh_entropy_eos() and _eos_content_key(): reload discipline +- AragogRunner._maybe_install_jax_cvode_factory(): install-last and clear-on-fail """ from __future__ import annotations @@ -979,6 +985,54 @@ def test_a_step_stopped_by_the_terminal_event_is_accepted_as_it_stands(): assert out_full.dt_actual == pytest.approx(100.0, rel=1e-12) +@pytest.mark.unit +def test_the_ladder_refreshes_the_tables_before_the_first_solve(monkeypatch, tmp_path): + """The retry ladder points the solver at the current tables before solving. + + Verifies: + - The ladder swaps ``solver.entropy_eos`` to the freshly loaded table + object before the first ``solve()`` call, so the step integrates on the + tables as regenerated, not on the object the solver was built with. + - The loader is handed the interior's table directory, not a cached path. + - The const-properties guard holds: a run with no tables refreshes + nothing, so the exemption cannot silently load a table set. + """ + runner, interior_o, attempts = _retry_ladder_runner(status=0, dt_actual=100.0) + runner._config.interior_energetics.const_properties = False + solver = interior_o.aragog_solver = runner.aragog_solver + interior_o._spider_eos_dir = str(tmp_path) + + order: list[str] = [] + sentinel = object() + orig_solve = solver.solve + solver.solve = lambda: (order.append('solve'), orig_solve())[1] + + def fake_loader(path): + order.append('refresh') + assert path == str(tmp_path) + return sentinel + + monkeypatch.setattr('proteus.interior_energetics.aragog._cached_entropy_eos', fake_loader) + runner._solve_with_retry({'Time': 2.15e5, 'T_cmb': 4000.0}, interior_o) + + assert solver.entropy_eos is sentinel + assert order.index('refresh') < order.index('solve'), ( + 'the tables were refreshed after the solve had already run on the stale object' + ) + assert order.count('refresh') == 1 + + # Guard: a const-properties run carries no tables, so nothing is loaded + # and no entropy_eos is installed on the solver. + guarded, guarded_interior, _ = _retry_ladder_runner(status=0, dt_actual=100.0) + guarded._config.interior_energetics.const_properties = True + guarded_interior.aragog_solver = guarded.aragog_solver + guarded_interior._spider_eos_dir = str(tmp_path) + order.clear() + guarded._solve_with_retry({'Time': 2.15e5, 'T_cmb': 4000.0}, guarded_interior) + assert order == [] + assert not hasattr(guarded.aragog_solver, 'entropy_eos') + + @pytest.mark.unit def test_a_step_that_never_advanced_is_still_refused(): """The ladder still refuses results that carry no usable state. @@ -1252,3 +1306,398 @@ def test_progress_is_weighed_against_what_the_coupling_asked_for(): # under the threshold: against the halved interval it would not be. assert advanced / asked < _STEP_PROGRESS_MIN_SHARE assert advanced / (0.5 * asked) > _STEP_PROGRESS_MIN_SHARE + + +@pytest.mark.unit +def test_the_core_temperature_guard_stands_aside_for_a_giant_impact(): + """A giant impact's core-temperature jump is kept, not retried away. + + Physical scenario: an impactor merges with the planet and re-melts the + mantle between two interior solves, so the core temperature moves by + thousands of kelvin in one coupling step. That jump is the impact, applied + outside the solver, and it is identical at every step size. + + Contract clause: the jump guard exists to reject a solve that returned + garbage, which a smaller step can fix. It cannot fix an impact, so on the + step a re-melt fires the guard stands aside; on every other step it keeps + its full strength. + + Verifies: + - The same 8000 K jump is accepted on the first attempt with the impact + flag raised and rejected down the whole ladder without it, which is the + discriminating pair: only the flag differs. + - The exemption is scoped to the jump guard, so a solve that actually + failed is still retried even on an impact step. + """ + prior = {'Time': 7.68e5, 'T_cmb': 4000.0} + + # 12000 K against a 4000 K prior state is an 8000 K jump, well past the + # 3000 K floor the guard applies at 1 M_earth. + impacted, impacted_interior, impacted_attempts = _retry_ladder_runner( + status=0, dt_actual=100.0, T_core=12000.0 + ) + impacted_interior.impact_reset_this_step = True + out = impacted._solve_with_retry(prior, impacted_interior) + + assert len(impacted_attempts) == 1, ( + 'the impact jump cannot shrink with the step, so retrying it burns the ' + 'ladder and kills the run at the impact' + ) + assert out.T_core == pytest.approx(12000.0, rel=1e-12) + + # Same solver result, same prior state, flag down: the guard must reject. + ordinary, ordinary_interior, ordinary_attempts = _retry_ladder_runner( + status=0, dt_actual=100.0, T_core=12000.0 + ) + with pytest.raises(RuntimeError, match='T_core jump'): + ordinary._solve_with_retry(prior, ordinary_interior) + assert len(ordinary_attempts) == 6, ( + 'without an impact to explain it, a jump of this size is a corrupted ' + 'solve and has to go down the ladder' + ) + + # The exemption covers the jump guard only. A solver that reports failure + # is still retried on an impact step, or a genuinely broken solve would be + # waved through whenever it landed on an impact. + failed, failed_interior, failed_attempts = _retry_ladder_runner( + status=-1, dt_actual=0.0, T_core=12000.0 + ) + failed_interior.impact_reset_this_step = True + with pytest.raises(RuntimeError): + failed._solve_with_retry(prior, failed_interior) + assert len(failed_attempts) == 6 + + +def _jax_factory_config(): + """Config carrying the numeric fields the option Z factory install reads.""" + config = MagicMock() + ie = config.interior_energetics + ie.rfront_loc = 0.5 + ie.rfront_wid = 0.2 + ie.solid_log10visc = 22.0 + ie.melt_log10visc = 2.0 + ie.grain_size = 0.1 + ie.solid_cond = 4.0 + ie.melt_cond = 4.0 + ie.spider.matprop_smooth_width = 0.0 + ie.trans_conduction = True + ie.trans_convection = True + ie.trans_grav_sep = True + ie.trans_mixing = True + ie.eddy_diffusivity_thermal = 1.0 + ie.eddy_diffusivity_chemical = 1.0 + ie.kappah_floor = 10.0 + ie.aragog.phase_smoothing = 'tanh' + ie.aragog.backend = 'jax' + return config + + +@pytest.mark.unit +def test_the_jax_right_hand_side_reads_the_mesh_on_every_solve(monkeypatch): + """The JAX right-hand side follows the structure it is asked to integrate. + + The factory is called once per solve. It reads the mesh from the solver at + that moment, for the same reason it rereads the boundary conditions: a + giant impact grows the planet, and Zalmoxis re-solves the structure as the + mantle freezes. A mesh copied once when the factory was installed would + leave the right-hand side integrating the planet from before the change, + while every other consumer sees the new one. + + Verifies: + - The mesh is read once per factory call, not once per install, so two + solves read it twice. + - The second read sees the replaced mesh object, not the one present when + the factory was installed. + """ + pytest.importorskip('jax') + pytest.importorskip('aragog.jax.phase') + from proteus.interior_energetics.aragog import AragogRunner + + monkeypatch.delenv('PROTEUS_CI_NIGHTLY', raising=False) + + before_mesh, after_mesh = object(), object() + solver = SimpleNamespace( + _n_stag=79, + _r_basic_flat=np.linspace(2.86e6, 5.84e6, 80), + _core_bc='energy_balance', + evaluator=SimpleNamespace(mesh=before_mesh), + parameters=SimpleNamespace( + boundary_conditions=MagicMock(), + energy=SimpleNamespace(tidal_array=np.zeros(79)), + radionuclides=[], + mesh=SimpleNamespace(core_density=10800.0), + ), + ) + installed = {} + solver.set_jax_cvode_factory = lambda f: installed.update(factory=f) + interior_o = SimpleNamespace(aragog_solver=solver, _spider_eos_dir='/nonexistent') + + with ( + patch('aragog.jax.phase.MeshArrays') as mesh_arrays, + patch('aragog.jax.phase.PhaseParams'), + patch('aragog.jax.solver.BoundaryParams'), + patch( + 'aragog.solver.cvode_jax.build_jax_rhs_and_jacobian', + return_value=('rhs', 'jac', {}), + ), + patch('proteus.interior_energetics.aragog._cached_entropy_eos_jax'), + ): + AragogRunner._maybe_install_jax_cvode_factory(_jax_factory_config(), interior_o) + factory = installed.get('factory') + assert factory is not None, 'the factory was not installed' + + # Installing must not read the mesh: reading it there is what froze the + # geometry, and a copy taken at install time is the defect itself. + assert mesh_arrays.from_numpy_mesh.call_count == 0 + + factory(MagicMock(), 'energy_balance') + assert mesh_arrays.from_numpy_mesh.call_count == 1 + + # A structure change replaces the mesh between solves. + solver.evaluator.mesh = after_mesh + factory(MagicMock(), 'energy_balance') + assert mesh_arrays.from_numpy_mesh.call_count == 2 + + meshes = [c.args[0] for c in mesh_arrays.from_numpy_mesh.call_args_list] + assert meshes == [before_mesh, after_mesh] + + +@pytest.mark.unit +def test_an_interior_that_moves_under_fixed_radii_is_still_followed(monkeypatch): + """A structure change that leaves both bounding radii untouched is followed. + + With mass coordinates the mesh pins its first and last node to the core and + surface radii and solves every interior node from the density profile, so a + Zalmoxis re-solve can redistribute the whole interior, and with it pressure, + gravity, area and volume, while both bounding radii and the cell count stay + bit-identical. Comparing geometry by those three numbers reports nothing has + changed and leaves the right-hand side on the previous structure. + + Verifies: + - The second solve is handed the moved mesh even though cell count and both + bounding radii are unchanged, which is what a fingerprint on those three + would miss. + - The interior really does differ, so the case is not vacuous. + """ + pytest.importorskip('jax') + pytest.importorskip('aragog.jax.phase') + from proteus.interior_energetics.aragog import AragogRunner + + monkeypatch.delenv('PROTEUS_CI_NIGHTLY', raising=False) + + n = 8 + r_cmb, r_surf = 2.86e6, 5.84e6 + # Same endpoints and same count; only the interior node placement differs, + # as a denser mantle would produce after a re-solve. + before = SimpleNamespace(radii=np.linspace(r_cmb, r_surf, n)) + moved = np.linspace(r_cmb, r_surf, n) ** 1.02 + moved *= (r_surf - r_cmb) / (moved[-1] - moved[0]) + moved += r_cmb - moved[0] + after = SimpleNamespace(radii=moved) + + assert after.radii[0] == pytest.approx(before.radii[0], rel=1e-15) + assert after.radii[-1] == pytest.approx(before.radii[-1], rel=1e-15) + assert len(after.radii) == len(before.radii) + # The interior genuinely moved, well beyond any tolerance a check could use. + assert np.max(np.abs(after.radii[1:-1] - before.radii[1:-1])) > 1.0e3 + + solver = SimpleNamespace( + _n_stag=n, + _r_basic_flat=before.radii, + _core_bc='energy_balance', + evaluator=SimpleNamespace(mesh=before), + parameters=SimpleNamespace( + boundary_conditions=MagicMock(), + energy=SimpleNamespace(tidal_array=np.zeros(n)), + radionuclides=[], + mesh=SimpleNamespace(core_density=10800.0), + ), + ) + installed = {} + solver.set_jax_cvode_factory = lambda f: installed.update(factory=f) + interior_o = SimpleNamespace(aragog_solver=solver, _spider_eos_dir='/nonexistent') + + with ( + patch('aragog.jax.phase.MeshArrays') as mesh_arrays, + patch('aragog.jax.phase.PhaseParams'), + patch('aragog.jax.solver.BoundaryParams'), + patch( + 'aragog.solver.cvode_jax.build_jax_rhs_and_jacobian', + return_value=('rhs', 'jac', {}), + ), + patch('proteus.interior_energetics.aragog._cached_entropy_eos_jax'), + ): + AragogRunner._maybe_install_jax_cvode_factory(_jax_factory_config(), interior_o) + factory = installed['factory'] + + factory(MagicMock(), 'energy_balance') + solver.evaluator.mesh = after + factory(MagicMock(), 'energy_balance') + + seen = [c.args[0] for c in mesh_arrays.from_numpy_mesh.call_args_list] + assert seen == [before, after] + np.testing.assert_allclose(seen[1].radii, moved) + + +@pytest.mark.unit +def test_the_solver_is_pointed_at_the_current_tables_before_each_solve(tmp_path): + """The energy diagnostic integrates the tables the step actually runs on. + + The solver keeps the table object it was built with, and `_step_heat_content` + integrates that object to produce the state side of the energy budget. The + tables are rewritten with a higher pressure ceiling whenever the planet + grows, so a solver left on the startup tables misreports the budget on + exactly the runs that outgrow them. + + Verifies: + - The solver is repointed when the tables have been rewritten. + - A const-properties run, which has no tables at all, is left alone rather + than being handed one. + - A missing table directory is a no-op, not a crash mid-run. + """ + from proteus.interior_energetics.aragog import AragogRunner + + d = tmp_path / 'spider_eos' + d.mkdir() + (d / '.cache_info.txt').write_text('P_max=2.750000e+11_nP=1350_nS=280') + (d / 'density_melt.dat').write_bytes(b'x' * 512) + + startup = object() + solver = SimpleNamespace(entropy_eos=startup) + interior_o = SimpleNamespace(aragog_solver=solver, _spider_eos_dir=str(d)) + config = MagicMock() + config.interior_energetics.const_properties = False + + loaded = object() + with patch( + 'proteus.interior_energetics.aragog._cached_entropy_eos', return_value=loaded + ) as loader: + AragogRunner._refresh_entropy_eos(config, interior_o) + assert loader.call_count == 1 + assert loader.call_args.args[0] == str(d) + assert solver.entropy_eos is loaded + assert solver.entropy_eos is not startup + + # const_properties carries no tables, so nothing may be attached. + const_cfg = MagicMock() + const_cfg.interior_energetics.const_properties = True + solver.entropy_eos = None + with patch('proteus.interior_energetics.aragog._cached_entropy_eos') as loader: + AragogRunner._refresh_entropy_eos(const_cfg, interior_o) + assert loader.call_count == 0 + assert solver.entropy_eos is None + + # A directory that is not there is a no-op: the run keeps whatever it had. + solver.entropy_eos = startup + gone = SimpleNamespace(aragog_solver=solver, _spider_eos_dir=str(tmp_path / 'absent')) + with patch('proteus.interior_energetics.aragog._cached_entropy_eos') as loader: + AragogRunner._refresh_entropy_eos(config, gone) + assert loader.call_count == 0 + assert solver.entropy_eos is startup + + +@pytest.mark.unit +def test_regenerated_eos_tables_are_seen_even_at_identical_file_sizes(tmp_path): + """Tables rewritten to a higher pressure ceiling are treated as new tables. + + A giant impact grows the planet, and the P-S tables are rewritten with a + ceiling scaled to the new mass on the same entropy and pressure grid. Every + file therefore keeps its length, so a key made of file sizes reports the + tables unchanged and the solver keeps evaluating the deepest cells against a + table built for the smaller planet, clamping at its edge. + + Verifies: + - The key changes when only the recorded ceiling changes, with byte counts + held equal, which is what the size-based key could not see. + - It still changes for a genuine size change, so the marker has not simply + replaced one blind spot with another. + - A directory with no marker still yields a usable key rather than raising. + """ + from proteus.interior_energetics.aragog import _eos_content_key + + def write(ceiling, pad=0): + d = tmp_path / f'eos_{ceiling}_{pad}' + d.mkdir() + (d / '.cache_info.txt').write_text( + f'P_max={ceiling:.6e}_nP=1350_nS=280_mzf=0.8_layout=2phase_eos=PALEOS-2phase' + ) + # Same grid shape means the same byte count, which is the whole trap. + (d / 'density_melt.dat').write_bytes(b'x' * (4096 + pad)) + return d + + before = write(2.750e11) # 0.5 M_earth embryo + after = write(8.750e11) # the same planet after growing to 4.5 M_earth + + sizes = {p.name: p.stat().st_size for p in before.iterdir() if p.name != '.cache_info.txt'} + after_sizes = { + p.name: p.stat().st_size for p in after.iterdir() if p.name != '.cache_info.txt' + } + assert sizes == after_sizes, 'the table files must match in size for this to bite' + + k_before = _eos_content_key(str(before)) + k_after = _eos_content_key(str(after)) + assert k_before != k_after + # The ceiling is what moved, so it must be what the key carries. + assert '2.750000e+11' in k_before + assert '8.750000e+11' in k_after + + # The marker fully describes the tables, so identical markers are the same + # tables however the bytes fall. This is deliberate, not a second blind spot. + grown = write(2.750e11, pad=512) + assert _eos_content_key(str(grown)) == k_before + + # Without a marker the fallback is the file listing, and it still separates + # two directories that differ only in size. + bare = tmp_path / 'bare' + bare.mkdir() + (bare / 'density_melt.dat').write_bytes(b'y' * 2048) + bigger = tmp_path / 'bigger' + bigger.mkdir() + (bigger / 'density_melt.dat').write_bytes(b'y' * 4096) + bare_key = _eos_content_key(str(bare)) + assert 'density_melt.dat' in bare_key + assert bare_key != _eos_content_key(str(bigger)) + + # A missing directory yields the path itself rather than raising. + assert _eos_content_key(str(tmp_path / 'missing')) == str(tmp_path / 'missing') + + +@pytest.mark.unit +def test_a_failed_factory_install_leaves_no_factory_behind(monkeypatch): + """A failed install must not leave a previous factory in charge. + + A first install has nothing to leave behind, so reporting a fallback to the + finite-difference Jacobian is accurate. A later one does: the solver still + carries the factory from the earlier install, and keeping it would run the + option Z path while the log reports a fallback that did not happen. + + Verifies: + - The factory is cleared, so the solve-time gate (factory is not None) + turns the path off rather than leaving the stale one installed. + """ + pytest.importorskip('jax') + pytest.importorskip('aragog.jax.phase') + from proteus.interior_energetics.aragog import AragogRunner + + # Nightly escalates every fallback to a hard failure; this test is about the + # non-strict path that a production run actually takes. + monkeypatch.delenv('PROTEUS_CI_NIGHTLY', raising=False) + + stale_factory = object() + solver = SimpleNamespace( + _n_stag=79, + _r_basic_flat=np.linspace(2.86e6, 5.84e6, 80), + _jax_cvode_factory=stale_factory, + ) + solver.set_jax_cvode_factory = lambda f: setattr(solver, '_jax_cvode_factory', f) + + config = MagicMock() + config.interior_energetics.aragog.backend = 'jax' + # interior_o carries no _spider_eos_dir, so the EOS lookup raises part-way + # through the install: the failure a live solver has to survive. + interior_o = SimpleNamespace(aragog_solver=solver) + + AragogRunner._maybe_install_jax_cvode_factory(config, interior_o) + + assert solver._jax_cvode_factory is not stale_factory + assert solver._jax_cvode_factory is None diff --git a/tests/interior_energetics/test_timestep.py b/tests/interior_energetics/test_timestep.py index 4640a7120..6643f9289 100644 --- a/tests/interior_energetics/test_timestep.py +++ b/tests/interior_energetics/test_timestep.py @@ -39,6 +39,7 @@ def _make_config( bol_scale: float = 1.0, bol_scale_start: float | None = None, bol_scale_duration: float = 0.0, + impact_maximum: float = 0.0, ): """Build a minimal duck-typed config that ``next_step`` reads from. @@ -65,6 +66,7 @@ def _make_config( hysteresis_iters=hysteresis_iters, hysteresis_sfinc=hysteresis_sfinc, max_growth_factor=max_growth_factor, + impact_maximum=impact_maximum, ) stop_solid = SimpleNamespace(enabled=True, phi_crit=phi_crit) stop_radeqm = SimpleNamespace(enabled=False) @@ -782,6 +784,98 @@ def test_an_imminent_impact_is_floored_at_the_minimum_step(self): assert dt > 0.0 assert hf_row['Time'] + dt > t_impact + @pytest.mark.physics_invariant + def test_impact_maximum_bounds_the_landing_step_below_the_remaining_time(self): + """A positive ceiling cuts the landing step further than the time + remaining to the impact would on its own. + + This is the case the ceiling exists for: a coarse phase leaves a + lot of time remaining when the impact clamp first engages, and + without a ceiling that whole remaining time becomes the landing + step. + """ + from proteus.interior_energetics.timestep import next_step + + config = _make_config(impact_maximum=3.0e3) + hf_all = _make_hf_all(n_rows=12, dt_prev=5.0e3, phi=1.0) + time_now = 1.0e5 + hf_row = {'Time': time_now, 'F_atm': 1.0e4, 'Phi_global': 1.0} + t_impact = time_now + 6.0e3 + + dt = next_step( + config, + {}, + hf_row, + hf_all, + 1.0, + interior_o=_make_interior_o(t_next_impact=t_impact), + ) + + assert dt == pytest.approx(3.0e3, rel=1e-9), f'Expected the 3e3 ceiling, got {dt}' + # Discrimination: the remaining-time clamp alone would land the step + # at 6e3 (the impact is nearer than the controller's 8e3 choice), so + # only an active ceiling can produce 3e3 here. + assert hf_row['Time'] + dt < t_impact + + def test_impact_maximum_does_not_shorten_a_step_already_below_it(self): + """The ceiling never lengthens the step and stays inert once the + remaining-time clamp has already produced something smaller.""" + from proteus.interior_energetics.timestep import next_step + + time_now = 1.0e5 + hf_row = {'Time': time_now, 'F_atm': 1.0e4, 'Phi_global': 1.0} + t_impact = time_now + 2.0e3 + hf_all = _make_hf_all(n_rows=12, dt_prev=5.0e3, phi=1.0) + + dt_with_ceiling = next_step( + _make_config(impact_maximum=5.0e3), + {}, + dict(hf_row), + hf_all, + 1.0, + interior_o=_make_interior_o(t_next_impact=t_impact), + ) + dt_without_ceiling = next_step( + _make_config(impact_maximum=0.0), + {}, + dict(hf_row), + hf_all, + 1.0, + interior_o=_make_interior_o(t_next_impact=t_impact), + ) + + assert dt_with_ceiling == pytest.approx(dt_without_ceiling, rel=1e-12) + assert dt_with_ceiling == pytest.approx(2.0e3, rel=1e-9) + + @pytest.mark.physics_invariant + def test_impact_maximum_never_beats_the_minimum_floor(self): + """A ceiling set below the minimum-step floor must not win. + + The floor exists so an imminent impact cannot collapse dt to zero; + a misconfigured ceiling smaller than the floor must not reopen + that hazard. + """ + from proteus.interior_energetics.timestep import next_step + + config = _make_config(impact_maximum=50.0) + hf_all = _make_hf_all(n_rows=12, dt_prev=5.0e3, phi=1.0) + time_now = 1.0e5 + hf_row = {'Time': time_now, 'F_atm': 1.0e4, 'Phi_global': 1.0} + t_impact = time_now + 10.0 + + dt = next_step( + config, + {}, + hf_row, + hf_all, + 1.0, + interior_o=_make_interior_o(t_next_impact=t_impact), + ) + + # The 600 yr floor wins over the 50 yr ceiling. + assert dt == pytest.approx(600.0, rel=1e-6), f'Expected the 600 yr floor, got {dt}' + assert dt > 0.0 + class TestImpactAndBolscaleClampsTogether: """Verify the step when a giant impact and a stellar-scaling edge compete. diff --git a/tests/interior_energetics/test_wrapper.py b/tests/interior_energetics/test_wrapper.py index c81d57dbc..16c86005b 100644 --- a/tests/interior_energetics/test_wrapper.py +++ b/tests/interior_energetics/test_wrapper.py @@ -2120,6 +2120,52 @@ def _run_interior_with_dummy(config, hf_all, hf_row, *, ic: int, output: dict): return hf_row +@pytest.mark.unit +def test_run_interior_consumes_the_impact_flag_into_the_step_flag(): + """run_interior translates the one-shot impact flag into the per-step flag. + + Verifies: + - An armed ``impact_reset`` is consumed (cleared) and surfaces as + ``impact_reset_this_step`` for the rest of the step, which is what the + temperature-jump clip and the solver's core-temperature guard read. + - The very next step reads False again, so one impact cannot exempt two + steps from the guards. + """ + from proteus.interior_energetics.common import Interior_t + from proteus.interior_energetics.wrapper import run_interior + + config = _make_run_interior_config(prevent_warming=False) + hf_all, hf_row = _make_run_interior_state(prev_f_int=1.0) + out = { + 'T_magma': 3005.0, + 'T_surf': 2805.0, + 'Phi_global': 0.7, + 'F_int': 2.0, + 'M_mantle': 4.0e24, + 'M_mantle_liquid': 1.0e24, + 'M_mantle_solid': 3.0e24, + 'M_core': 2.0e24, + } + interior_o = Interior_t(nlev_b=10) + interior_o.ic = 2 + interior_o.impact_reset = True + + with ( + patch( + 'proteus.interior_energetics.dummy.run_dummy_int', + return_value=(110.0, out), + ), + patch('proteus.interior_energetics.wrapper.update_planet_mass'), + ): + run_interior({}, config, hf_all, hf_row, interior_o, MagicMock(), verbose=False) + assert interior_o.impact_reset_this_step is True + assert interior_o.impact_reset is False, 'the one-shot flag was not consumed' + + # The following step is ordinary again: nothing re-armed the flag. + run_interior({}, config, hf_all, hf_row, interior_o, MagicMock(), verbose=False) + assert interior_o.impact_reset_this_step is False + + @pytest.mark.unit def test_run_interior_prevent_warming_clamps_T_and_Fint_on_ic2(): """prevent_warming=True + ic=2 must clip Phi/T_magma/T_surf to previous values