From 803be0ad104e2f6b8a5b60dc8186ecae15607087 Mon Sep 17 00:00:00 2001 From: timlichtenberg Date: Wed, 5 Aug 2026 11:41:25 +0200 Subject: [PATCH 1/9] Keep the interior solver on the planet an impact just made A giant impact grows the planet between two interior solves. Two things did not follow it. The JAX right-hand side captures the mesh by value when its factory is installed, and the factory was installed once per process. After an impact the numpy mesh was rebuilt for the new structure while the JAX side kept integrating the old one, so the boundary-flux budget and the state-heat integral stopped describing the same planet. The gap shows up as a step in their ratio at every impact, flat in between, and it disappears if the run is restarted, because a restart rebuilds the factory. Record the geometry the factory was built against and rebuild it when the mesh no longer matches, which for a run without impacts never fires. The core-temperature jump guard also treated the impact as a bad solve. The re-melt moves T_core by thousands of kelvin outside the solver, so the jump is identical at every step size and the retry ladder cannot reduce it; the run died at its first impact after burning six attempts. The guard now stands aside on the step a re-melt fires and keeps its full strength everywhere else, reusing the impact_reset flag the temperature clamps already honour. That flag is cleared before the interior solver runs, so the value is kept for the rest of the step. --- src/proteus/interior_energetics/aragog.py | 89 +++++++++++++++- src/proteus/interior_energetics/common.py | 6 ++ src/proteus/interior_energetics/wrapper.py | 3 + tests/interior_energetics/test_aragog.py | 114 +++++++++++++++++++++ 4 files changed, 210 insertions(+), 2 deletions(-) diff --git a/src/proteus/interior_energetics/aragog.py b/src/proteus/interior_energetics/aragog.py index 6d4556816..9b667281a 100644 --- a/src/proteus/interior_energetics/aragog.py +++ b/src/proteus/interior_energetics/aragog.py @@ -586,11 +586,54 @@ def setup_or_update_solver( else: AragogRunner.update_structure(config, hf_row, interior_o) AragogRunner.update_solver(dt, hf_row, interior_o) + AragogRunner._refresh_jax_cvode_factory_if_mesh_moved(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: interior_o.aragog_solver.set_initial_entropy(interior_o._last_entropy) + @staticmethod + def _refresh_jax_cvode_factory_if_mesh_moved( + config: Config, interior_o: Interior_t + ) -> None: + """Rebuild the JAX CVODE factory when the mesh no longer matches it. + + The factory captures the mesh by value at install time, so a structure + change (a giant impact grows the planet, and Zalmoxis re-solves the + radius and the core/mantle split) leaves its right-hand side integrating + the old geometry while every other consumer sees the new one. Rebuilding + costs a JAX retrace, so it happens only when the geometry actually moved, + which for a run without impacts is never. + + Parameters + ---------- + config : Config + PROTEUS configuration, forwarded to the factory installer. + interior_o : Interior_t + Interior state holding the live Aragog solver. + """ + solver = interior_o.aragog_solver + installed = getattr(solver, '_jax_mesh_key', None) + if not isinstance(installed, tuple): + return # option Z inactive, or the install failed and fell back + + current = AragogRunner._jax_mesh_key(solver) + if current == installed: + return + + log.info( + 'Option Z: mesh moved under the JAX factory (r_surf %.6e -> %.6e m, ' + 'r_cmb %.6e -> %.6e m, n_stag %d -> %d); rebuilding it so the ' + 'right-hand side integrates the current structure.', + installed[2], + current[2], + installed[1], + current[1], + installed[0], + current[0], + ) + AragogRunner._maybe_install_jax_cvode_factory(config, interior_o) + @staticmethod def setup_solver(config: Config, hf_row: dict, interior_o: Interior_t, outdir: str): solver = _SolverParameters( @@ -1223,6 +1266,28 @@ def _append_radnuc(_iso, _cnc): _t_post_solver - _t_post_eos, ) + @staticmethod + def _jax_mesh_key(solver) -> tuple[int, float, float]: + """Geometry fingerprint of the solver's current mesh. + + The JAX CVODE factory captures the mesh by value, so this is what has + to match for its right-hand side to describe the planet the rest of the + step describes. Cell count plus the two bounding radii is enough: the + mesh is rebuilt whole on a structure change, never edited in place. + + Parameters + ---------- + solver : EntropySolver + Aragog solver whose mesh is being fingerprinted. + + Returns + ------- + tuple of (int, float, float) + Staggered cell count, CMB radius [m], surface radius [m]. + """ + r_basic = np.asarray(solver._r_basic_flat).ravel() + return (int(solver._n_stag), float(r_basic[0]), float(r_basic[-1])) + @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). @@ -1401,11 +1466,18 @@ def factory(scales, core_bc_mode): return rhs_fn, jac_fn solver.set_jax_cvode_factory(factory) + # The factory closes over mesh_jax by value, so the geometry it + # integrates is frozen here. Record it so a later structure change + # (a giant impact grows the planet) can be detected and the factory + # rebuilt, rather than integrating the old planet silently. + solver._jax_mesh_key = AragogRunner._jax_mesh_key(solver) log.info( 'Option Z: JAX CVODE factory installed on aragog solver ' - '(core_bc=%s, n_stag=%d).', + '(core_bc=%s, n_stag=%d, r_cmb=%.6e m, r_surf=%.6e m).', solver._core_bc, n_stag, + solver._jax_mesh_key[1], + solver._jax_mesh_key[2], ) except Exception as exc: msg = f'Option Z factory install failed ({exc}); falling back to FD Jacobian.' @@ -2065,6 +2137,11 @@ 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)) # Capture IC for restoration on retry, and pre-call T_core for # the sanity check on retry success. @@ -2153,7 +2230,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/wrapper.py b/src/proteus/interior_energetics/wrapper.py index 1d37f18dd..088128be4 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: diff --git a/tests/interior_energetics/test_aragog.py b/tests/interior_energetics/test_aragog.py index 79df4a96a..6f6e337c0 100644 --- a/tests/interior_energetics/test_aragog.py +++ b/tests/interior_energetics/test_aragog.py @@ -1252,3 +1252,117 @@ 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 +@pytest.mark.physics_invariant +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 + + +@pytest.mark.unit +def test_the_jax_factory_is_rebuilt_only_when_the_mesh_moves(): + """The JAX right-hand side follows the structure across a giant impact. + + The factory captures the mesh by value, so a structure change leaves it + integrating the old planet while every other consumer sees the new one. + It is rebuilt when the geometry moves and left alone when it has not, + because the rebuild costs a JAX retrace on every step that triggers it. + + Verifies: + - An unchanged mesh triggers no rebuild, so a run without impacts never + pays the retrace. + - A grown planet triggers exactly one rebuild. + - A solver with no recorded key (option Z inactive, or its install fell + back) is left alone rather than raising. + """ + from proteus.interior_energetics.aragog import AragogRunner + + def _solver(n_stag, r_cmb, r_surf, key): + s = SimpleNamespace( + _n_stag=n_stag, _r_basic_flat=np.linspace(r_cmb, r_surf, n_stag + 1) + ) + if key is not None: + s._jax_mesh_key = key + return s + + config = MagicMock() + settled = _solver(79, 2.86e6, 5.84e6, (79, 2.86e6, 5.84e6)) + assert AragogRunner._jax_mesh_key(settled) == (79, 2.86e6, 5.84e6) + + with patch.object(AragogRunner, '_maybe_install_jax_cvode_factory') as install: + AragogRunner._refresh_jax_cvode_factory_if_mesh_moved( + config, SimpleNamespace(aragog_solver=settled) + ) + assert install.call_count == 0 + + # An impact grows the planet: both radii move and the key no longer matches. + grown = _solver(79, 3.46e6, 7.16e6, (79, 2.86e6, 5.84e6)) + with patch.object(AragogRunner, '_maybe_install_jax_cvode_factory') as install: + AragogRunner._refresh_jax_cvode_factory_if_mesh_moved( + config, SimpleNamespace(aragog_solver=grown) + ) + assert install.call_count == 1 + + # No key recorded: nothing to compare against, and nothing to rebuild. + absent = _solver(79, 2.86e6, 5.84e6, None) + with patch.object(AragogRunner, '_maybe_install_jax_cvode_factory') as install: + AragogRunner._refresh_jax_cvode_factory_if_mesh_moved( + config, SimpleNamespace(aragog_solver=absent) + ) + assert install.call_count == 0 From 9c7b64d521006ac63cf06915e3aa144e80e550e2 Mon Sep 17 00:00:00 2001 From: timlichtenberg Date: Wed, 5 Aug 2026 13:08:42 +0200 Subject: [PATCH 2/9] Check the mesh against the factory just before the solve The rebuild check sat at solver setup, where the mesh has not yet reached the geometry the step will integrate: the structure update propagates a step later, so the first solve after an impact still ran on the old right-hand side and only the second picked up the new one. Moving the check to immediately before the solve puts it where the mesh is final. Measured on a case whose two impacts land inside the first 2000 yr: the booked-flux to state-heat ratio now reads 0.99999 on every ordinary step either side of both impacts, against 0.688 and 0.425 held indefinitely before. The one step after each impact still reads 0.74 and 0.60, which is the CMB term the per-call integrals over-count there, not this. --- src/proteus/interior_energetics/aragog.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/proteus/interior_energetics/aragog.py b/src/proteus/interior_energetics/aragog.py index 9b667281a..d6dd88ccb 100644 --- a/src/proteus/interior_energetics/aragog.py +++ b/src/proteus/interior_energetics/aragog.py @@ -586,7 +586,6 @@ def setup_or_update_solver( else: AragogRunner.update_structure(config, hf_row, interior_o) AragogRunner.update_solver(dt, hf_row, interior_o) - AragogRunner._refresh_jax_cvode_factory_if_mesh_moved(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: @@ -612,7 +611,7 @@ def _refresh_jax_cvode_factory_if_mesh_moved( interior_o : Interior_t Interior state holding the live Aragog solver. """ - solver = interior_o.aragog_solver + solver = getattr(interior_o, 'aragog_solver', None) installed = getattr(solver, '_jax_mesh_key', None) if not isinstance(installed, tuple): return # option Z inactive, or the install failed and fell back @@ -2143,6 +2142,12 @@ def _solve_with_retry(self, hf_row, interior_o) -> SolverOutput: # the run. Skip it on that one step; every other step keeps it. impact_step = bool(getattr(interior_o, 'impact_reset_this_step', False)) + # Checked here rather than at solver setup: the mesh reaches its final + # geometry for the step only once the structure update has propagated, + # which lags setup by a step, and the factory has to match the mesh this + # solve actually integrates. + AragogRunner._refresh_jax_cvode_factory_if_mesh_moved(self._config, interior_o) + # Capture IC for restoration on retry, and pre-call T_core for # the sanity check on retry success. t_start = float(solver.parameters.solver.start_time) From cb2f3fb7f232d3b5384b571398d0612cc0e65911 Mon Sep 17 00:00:00 2001 From: timlichtenberg Date: Thu, 6 Aug 2026 09:22:45 +0200 Subject: [PATCH 3/9] Clear the JAX factory when a rebuild of it fails The option Z factory is now built more than once: it is rebuilt whenever the mesh moves under it. Its failure path was only ever correct for the first build, where there is nothing installed yet, so reporting a fall back to the finite-difference Jacobian is accurate. On a rebuild the solver still carries the factory built against the geometry the rebuild was meant to replace, and that factory stayed installed and in use while the log said the run had fallen back. The result is the failure this whole path exists to prevent: the solve keeps integrating the planet from before the structure changed, silently. Clear the factory and the recorded geometry key together, so the solve-time check turns the path off rather than running it on a stale mesh, and so a later geometry comparison cannot measure itself against a factory that is no longer there. The new test drives a rebuild to failure on a solver that already has a factory installed, and fails without this change. --- src/proteus/interior_energetics/aragog.py | 6 +++ tests/interior_energetics/test_aragog.py | 49 +++++++++++++++++++++++ 2 files changed, 55 insertions(+) diff --git a/src/proteus/interior_energetics/aragog.py b/src/proteus/interior_energetics/aragog.py index d6dd88ccb..4218be3fa 100644 --- a/src/proteus/interior_energetics/aragog.py +++ b/src/proteus/interior_energetics/aragog.py @@ -1482,6 +1482,12 @@ def factory(scales, core_bc_mode): msg = f'Option Z factory install failed ({exc}); falling back to FD Jacobian.' if nightly_strict: raise RuntimeError(msg) from exc + # On a rebuild the solver already carries a factory built against the + # previous geometry. Leaving it installed would keep the solve on that + # geometry while this message claims the opposite, so clear both it and + # the key that records what it was built against. + solver.set_jax_cvode_factory(None) + solver._jax_mesh_key = None log.warning(msg) @staticmethod diff --git a/tests/interior_energetics/test_aragog.py b/tests/interior_energetics/test_aragog.py index 6f6e337c0..4b521ad53 100644 --- a/tests/interior_energetics/test_aragog.py +++ b/tests/interior_energetics/test_aragog.py @@ -1366,3 +1366,52 @@ def _solver(n_stag, r_cmb, r_surf, key): config, SimpleNamespace(aragog_solver=absent) ) assert install.call_count == 0 + + +def test_a_failed_rebuild_stops_the_solver_integrating_the_old_geometry(monkeypatch): + """A rebuild that fails must not leave the previous factory in charge. + + A first install has no factory to leave behind, so reporting a fallback to + the finite-difference Jacobian is accurate. A rebuild does have one: the + solver still carries the factory built against the geometry the rebuild was + meant to replace, and keeping it would integrate the old planet 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 option Z path off rather than running it on stale geometry. + - The recorded mesh key is cleared with it, so the geometry check cannot + later compare against a key describing a factory that is not 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, + _jax_mesh_key=(79, 2.86e6, 5.84e6), + ) + 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 rebuild: 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 + assert solver._jax_mesh_key is None + + with patch.object(AragogRunner, '_maybe_install_jax_cvode_factory') as install: + AragogRunner._refresh_jax_cvode_factory_if_mesh_moved(config, interior_o) + assert install.call_count == 0 From 9d829c29b7082737d883ba620d3d81d073cc4d2c Mon Sep 17 00:00:00 2001 From: timlichtenberg Date: Thu, 6 Aug 2026 10:28:34 +0200 Subject: [PATCH 4/9] Read the interior mesh on every solve, not once The JAX right-hand side was built from a copy of the mesh taken when its factory was installed, so it kept describing the planet as it stood at that moment. A giant impact grows the planet, and Zalmoxis re-solves the structure as the mantle freezes; either one replaces the mesh, and the right-hand side went on integrating the old one while every other consumer saw the new one. Comparing a stored geometry fingerprint against the current one, which is how I first approached this, only closes part of it. 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 structure change can leave both of those radii untouched and still move the whole interior, along with the pressure, gravity, area and volume arrays the right-hand side is built from. The fingerprint would report nothing had changed. Read the mesh inside the factory instead, alongside the boundary conditions and the heating arrays, which are already reread there for the same reason. The factory runs once per solve, so the mesh can no longer be older than the step it describes, and the fingerprint and its refresh check are no longer needed. The test asserts the mesh is read once per solve rather than once per install, and that the second solve sees a replaced mesh. It fails if the read is moved back out of the factory. --- src/proteus/interior_energetics/aragog.py | 109 +++------------- tests/interior_energetics/test_aragog.py | 146 +++++++++++++--------- 2 files changed, 108 insertions(+), 147 deletions(-) diff --git a/src/proteus/interior_energetics/aragog.py b/src/proteus/interior_energetics/aragog.py index 4218be3fa..5fd1549fd 100644 --- a/src/proteus/interior_energetics/aragog.py +++ b/src/proteus/interior_energetics/aragog.py @@ -591,48 +591,6 @@ def setup_or_update_solver( if hasattr(interior_o, '_last_entropy') and interior_o._last_entropy is not None: interior_o.aragog_solver.set_initial_entropy(interior_o._last_entropy) - @staticmethod - def _refresh_jax_cvode_factory_if_mesh_moved( - config: Config, interior_o: Interior_t - ) -> None: - """Rebuild the JAX CVODE factory when the mesh no longer matches it. - - The factory captures the mesh by value at install time, so a structure - change (a giant impact grows the planet, and Zalmoxis re-solves the - radius and the core/mantle split) leaves its right-hand side integrating - the old geometry while every other consumer sees the new one. Rebuilding - costs a JAX retrace, so it happens only when the geometry actually moved, - which for a run without impacts is never. - - Parameters - ---------- - config : Config - PROTEUS configuration, forwarded to the factory installer. - interior_o : Interior_t - Interior state holding the live Aragog solver. - """ - solver = getattr(interior_o, 'aragog_solver', None) - installed = getattr(solver, '_jax_mesh_key', None) - if not isinstance(installed, tuple): - return # option Z inactive, or the install failed and fell back - - current = AragogRunner._jax_mesh_key(solver) - if current == installed: - return - - log.info( - 'Option Z: mesh moved under the JAX factory (r_surf %.6e -> %.6e m, ' - 'r_cmb %.6e -> %.6e m, n_stag %d -> %d); rebuilding it so the ' - 'right-hand side integrates the current structure.', - installed[2], - current[2], - installed[1], - current[1], - installed[0], - current[0], - ) - AragogRunner._maybe_install_jax_cvode_factory(config, interior_o) - @staticmethod def setup_solver(config: Config, hf_row: dict, interior_o: Interior_t, outdir: str): solver = _SolverParameters( @@ -1265,28 +1223,6 @@ def _append_radnuc(_iso, _cnc): _t_post_solver - _t_post_eos, ) - @staticmethod - def _jax_mesh_key(solver) -> tuple[int, float, float]: - """Geometry fingerprint of the solver's current mesh. - - The JAX CVODE factory captures the mesh by value, so this is what has - to match for its right-hand side to describe the planet the rest of the - step describes. Cell count plus the two bounding radii is enough: the - mesh is rebuilt whole on a structure change, never edited in place. - - Parameters - ---------- - solver : EntropySolver - Aragog solver whose mesh is being fingerprinted. - - Returns - ------- - tuple of (int, float, float) - Staggered cell count, CMB radius [m], surface radius [m]. - """ - r_basic = np.asarray(solver._r_basic_flat).ravel() - return (int(solver._n_stag), float(r_basic[0]), float(r_basic[-1])) - @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). @@ -1377,14 +1313,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): @@ -1393,6 +1325,14 @@ 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, + # for the same reason the boundary conditions below are. A giant + # impact grows the planet and Zalmoxis re-solves the structure as + # the mantle freezes; either one replaces the mesh, and a copy + # taken once at install time would keep integrating the planet + # from before the change. + mesh_jax = MeshArrays.from_numpy_mesh(solver.evaluator.mesh) + n_stag = solver._n_stag # ``scales`` is an aragog.jax.nondim.NonDimScales single # source of truth. # Rebuild BoundaryParams from live solver state every @@ -1465,29 +1405,24 @@ def factory(scales, core_bc_mode): return rhs_fn, jac_fn solver.set_jax_cvode_factory(factory) - # The factory closes over mesh_jax by value, so the geometry it - # integrates is frozen here. Record it so a later structure change - # (a giant impact grows the planet) can be detected and the factory - # rebuilt, rather than integrating the old planet silently. - solver._jax_mesh_key = AragogRunner._jax_mesh_key(solver) + r_basic = np.asarray(solver._r_basic_flat).ravel() log.info( 'Option Z: JAX CVODE factory installed on aragog solver ' - '(core_bc=%s, n_stag=%d, r_cmb=%.6e m, r_surf=%.6e m).', + '(core_bc=%s, n_stag=%d, r_cmb=%.6e m, r_surf=%.6e m). The mesh ' + 'is read from the solver on every solve, so these are the ' + 'geometry at install time, not for the run.', solver._core_bc, - n_stag, - solver._jax_mesh_key[1], - solver._jax_mesh_key[2], + solver._n_stag, + float(r_basic[0]), + float(r_basic[-1]), ) 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 - # On a rebuild the solver already carries a factory built against the - # previous geometry. Leaving it installed would keep the solve on that - # geometry while this message claims the opposite, so clear both it and - # the key that records what it was built against. + # A failed rebuild would otherwise leave the previously installed + # factory in charge while this message claims the run fell back. solver.set_jax_cvode_factory(None) - solver._jax_mesh_key = None log.warning(msg) @staticmethod @@ -2148,12 +2083,6 @@ def _solve_with_retry(self, hf_row, interior_o) -> SolverOutput: # the run. Skip it on that one step; every other step keeps it. impact_step = bool(getattr(interior_o, 'impact_reset_this_step', False)) - # Checked here rather than at solver setup: the mesh reaches its final - # geometry for the step only once the structure update has propagated, - # which lags setup by a step, and the factory has to match the mesh this - # solve actually integrates. - AragogRunner._refresh_jax_cvode_factory_if_mesh_moved(self._config, interior_o) - # Capture IC for restoration on retry, and pre-call T_core for # the sanity check on retry success. t_start = float(solver.parameters.solver.start_time) diff --git a/tests/interior_energetics/test_aragog.py b/tests/interior_energetics/test_aragog.py index 4b521ad53..e7c3c00a0 100644 --- a/tests/interior_energetics/test_aragog.py +++ b/tests/interior_energetics/test_aragog.py @@ -1316,72 +1316,110 @@ def test_the_core_temperature_guard_stands_aside_for_a_giant_impact(): @pytest.mark.unit -def test_the_jax_factory_is_rebuilt_only_when_the_mesh_moves(): - """The JAX right-hand side follows the structure across a giant impact. +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 + + +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 captures the mesh by value, so a structure change leaves it - integrating the old planet while every other consumer sees the new one. - It is rebuilt when the geometry moves and left alone when it has not, - because the rebuild costs a JAX retrace on every step that triggers it. + 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: - - An unchanged mesh triggers no rebuild, so a run without impacts never - pays the retrace. - - A grown planet triggers exactly one rebuild. - - A solver with no recorded key (option Z inactive, or its install fell - back) is left alone rather than raising. + - 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 - def _solver(n_stag, r_cmb, r_surf, key): - s = SimpleNamespace( - _n_stag=n_stag, _r_basic_flat=np.linspace(r_cmb, r_surf, n_stag + 1) - ) - if key is not None: - s._jax_mesh_key = key - return s + monkeypatch.delenv('PROTEUS_CI_NIGHTLY', raising=False) - config = MagicMock() - settled = _solver(79, 2.86e6, 5.84e6, (79, 2.86e6, 5.84e6)) - assert AragogRunner._jax_mesh_key(settled) == (79, 2.86e6, 5.84e6) + 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.object(AragogRunner, '_maybe_install_jax_cvode_factory') as install: - AragogRunner._refresh_jax_cvode_factory_if_mesh_moved( - config, SimpleNamespace(aragog_solver=settled) - ) - assert install.call_count == 0 + 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' - # An impact grows the planet: both radii move and the key no longer matches. - grown = _solver(79, 3.46e6, 7.16e6, (79, 2.86e6, 5.84e6)) - with patch.object(AragogRunner, '_maybe_install_jax_cvode_factory') as install: - AragogRunner._refresh_jax_cvode_factory_if_mesh_moved( - config, SimpleNamespace(aragog_solver=grown) - ) - assert install.call_count == 1 + # 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 - # No key recorded: nothing to compare against, and nothing to rebuild. - absent = _solver(79, 2.86e6, 5.84e6, None) - with patch.object(AragogRunner, '_maybe_install_jax_cvode_factory') as install: - AragogRunner._refresh_jax_cvode_factory_if_mesh_moved( - config, SimpleNamespace(aragog_solver=absent) - ) - assert install.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 -def test_a_failed_rebuild_stops_the_solver_integrating_the_old_geometry(monkeypatch): - """A rebuild that fails must not leave the previous factory in charge. + meshes = [c.args[0] for c in mesh_arrays.from_numpy_mesh.call_args_list] + assert meshes == [before_mesh, after_mesh] - A first install has no factory to leave behind, so reporting a fallback to - the finite-difference Jacobian is accurate. A rebuild does have one: the - solver still carries the factory built against the geometry the rebuild was - meant to replace, and keeping it would integrate the old planet while the - log reports a fallback that did not happen. + +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 option Z path off rather than running it on stale geometry. - - The recorded mesh key is cleared with it, so the geometry check cannot - later compare against a key describing a factory that is not installed. + turns the path off rather than leaving the stale one installed. """ pytest.importorskip('jax') pytest.importorskip('aragog.jax.phase') @@ -1396,22 +1434,16 @@ def test_a_failed_rebuild_stops_the_solver_integrating_the_old_geometry(monkeypa _n_stag=79, _r_basic_flat=np.linspace(2.86e6, 5.84e6, 80), _jax_cvode_factory=stale_factory, - _jax_mesh_key=(79, 2.86e6, 5.84e6), ) 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 rebuild: the failure a live solver has to survive. + # 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 - assert solver._jax_mesh_key is None - - with patch.object(AragogRunner, '_maybe_install_jax_cvode_factory') as install: - AragogRunner._refresh_jax_cvode_factory_if_mesh_moved(config, interior_o) - assert install.call_count == 0 From 4c4689afef4619ff53294a0af827a0299e2f6f4f Mon Sep 17 00:00:00 2001 From: timlichtenberg Date: Thu, 6 Aug 2026 15:54:25 +0200 Subject: [PATCH 5/9] Install the JAX factory last, so a diagnostic cannot undo it The factory was installed and then a purely diagnostic log line read the mesh radii, still inside the same try. Anything raising in that line reached the handler, which clears the factory, so a working install would have been torn down and the run quietly moved to the finite-difference Jacobian while reporting an install failure. Read the diagnostic values first and install last, so nothing that can fail runs after the install. Also moves a test marker that was left on a helper rather than on the two tests it belonged to, and adds the case that defeated the earlier approach: a mesh whose cell count and both bounding radii are bit-identical while its interior has moved. Mass coordinates pin the first and last node and solve the rest from the density profile, so that case is what a fingerprint on those three numbers cannot see. --- src/proteus/interior_energetics/aragog.py | 24 ++++--- tests/interior_energetics/test_aragog.py | 81 ++++++++++++++++++++++- 2 files changed, 93 insertions(+), 12 deletions(-) diff --git a/src/proteus/interior_energetics/aragog.py b/src/proteus/interior_energetics/aragog.py index 5fd1549fd..6b0230642 100644 --- a/src/proteus/interior_energetics/aragog.py +++ b/src/proteus/interior_energetics/aragog.py @@ -1404,24 +1404,26 @@ def factory(scales, core_bc_mode): ) return rhs_fn, jac_fn - solver.set_jax_cvode_factory(factory) + # 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() - log.info( + installed = ( 'Option Z: JAX CVODE factory installed on aragog solver ' - '(core_bc=%s, n_stag=%d, r_cmb=%.6e m, r_surf=%.6e m). The mesh ' - 'is read from the solver on every solve, so these are the ' - 'geometry at install time, not for the run.', - solver._core_bc, - solver._n_stag, - float(r_basic[0]), - float(r_basic[-1]), + 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 - # A failed rebuild would otherwise leave the previously installed - # factory in charge while this message claims the run fell back. + # 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) diff --git a/tests/interior_energetics/test_aragog.py b/tests/interior_energetics/test_aragog.py index e7c3c00a0..367ebcfa7 100644 --- a/tests/interior_energetics/test_aragog.py +++ b/tests/interior_energetics/test_aragog.py @@ -1315,7 +1315,6 @@ def test_the_core_temperature_guard_stands_aside_for_a_giant_impact(): assert len(failed_attempts) == 6 -@pytest.mark.unit def _jax_factory_config(): """Config carrying the numeric fields the option Z factory install reads.""" config = MagicMock() @@ -1340,6 +1339,7 @@ def _jax_factory_config(): 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. @@ -1409,6 +1409,85 @@ def test_the_jax_right_hand_side_reads_the_mesh_on_every_solve(monkeypatch): assert meshes == [before_mesh, after_mesh] +@pytest.mark.unit +@pytest.mark.physics_invariant +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_a_failed_factory_install_leaves_no_factory_behind(monkeypatch): """A failed install must not leave a previous factory in charge. From 00ca73bf5f9019d1bd615e71d910259d24974dee Mon Sep 17 00:00:00 2001 From: timlichtenberg Date: Thu, 6 Aug 2026 16:58:42 +0200 Subject: [PATCH 6/9] Reload the EOS tables when they are regenerated The P-S tables are rebuilt whenever the structure solve reruns, and a giant impact raises their pressure ceiling along with the planet's mass. Aragog kept the copy it loaded at startup, so once the planet outgrew that table the deepest cells were evaluated at its edge instead of at their own pressure. On a probe run growing 0.5 to 4.5 Earth masses the core-mantle boundary reaches 1.67 times the starting table's ceiling, 60 percent of the rows sit beyond it, and density there is understated by 15 percent. Two things kept the stale copy in place. The loader's cache key was a list of file names and sizes, and a regenerated table keeps the same grid shape, so every file keeps its length and the key could not tell the new tables from the old. The key now uses the marker the generator already writes beside the tables, which records the pressure ceiling, the grid shape and the material. Second, the tables were read once when the solver was built; they are now read where the mesh is read, once per solve, so both follow the planet the step is about. This closes the reload the structure-update path documented as missing for composition changes. The remaining piece there is unchanged: the entropy carried over from the previous step is still not bounds-checked against a regenerated range, which a composition change can move and a pressure-ceiling change cannot. --- src/proteus/interior_energetics/aragog.py | 31 ++++++++-- src/proteus/interior_energetics/wrapper.py | 25 ++++---- tests/interior_energetics/test_aragog.py | 67 ++++++++++++++++++++++ 3 files changed, 103 insertions(+), 20 deletions(-) diff --git a/src/proteus/interior_energetics/aragog.py b/src/proteus/interior_energetics/aragog.py index 6b0230642..3341308ba 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) @@ -1280,7 +1290,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( @@ -1333,6 +1345,13 @@ def factory(scales, core_bc_mode): # from before the change. mesh_jax = MeshArrays.from_numpy_mesh(solver.evaluator.mesh) n_stag = solver._n_stag + # Same for the EOS tables. They are regenerated whenever the + # structure solve reruns, and a giant impact raises their + # pressure ceiling with the planet's mass, so a copy taken at + # install time would have the deep mantle clamped to the edge of + # a table built for a smaller planet. The loader is cached on the + # table parameters, 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 diff --git a/src/proteus/interior_energetics/wrapper.py b/src/proteus/interior_energetics/wrapper.py index 088128be4..36e197e74 100644 --- a/src/proteus/interior_energetics/wrapper.py +++ b/src/proteus/interior_energetics/wrapper.py @@ -3342,17 +3342,14 @@ 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]. A composition change can move that range; a + # pressure-ceiling change alone does not. if comp_changed and config.interior_energetics.module in ('spider', 'aragog'): from proteus.interior_struct.zalmoxis import generate_spider_tables @@ -3363,10 +3360,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 367ebcfa7..d3f6e3ef7 100644 --- a/tests/interior_energetics/test_aragog.py +++ b/tests/interior_energetics/test_aragog.py @@ -1487,6 +1487,73 @@ def test_an_interior_that_moves_under_fixed_radii_is_still_followed(monkeypatch) np.testing.assert_allclose(seen[1].radii, moved) +@pytest.mark.unit +@pytest.mark.physics_invariant +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. From 8fb954070726ecb89932a967b8f0a96008e061e2 Mon Sep 17 00:00:00 2001 From: timlichtenberg Date: Thu, 6 Aug 2026 21:11:29 +0200 Subject: [PATCH 7/9] Point the solver at the current tables before each solve The interior solver kept whichever P-S tables it was built with at startup. Those tables are rewritten whenever the structure solve reruns, with a pressure ceiling that grows with the planet, so a run that outgrows its starting table went on integrating the old one. The right-hand side already follows the regenerated tables; this is the same reload for the state-heat integral, which is the other half of the energy budget and the quantity the conservation residual is measured against. That makes it a correctness fix for the diagnostic rather than for the trajectory: on a planet that stays inside its starting table nothing changes, and on one that outgrows it the reported budget was being taken against tables that no longer describe the planet. The reload sits immediately before the solve, so the integral a step books is taken against the tables that step runs on, and the loader is cached on the table parameters so an unchanged table costs one small read. --- src/proteus/interior_energetics/aragog.py | 31 ++++++++++++ tests/interior_energetics/test_aragog.py | 58 +++++++++++++++++++++++ 2 files changed, 89 insertions(+) diff --git a/src/proteus/interior_energetics/aragog.py b/src/proteus/interior_energetics/aragog.py index 3341308ba..43bab8f06 100644 --- a/src/proteus/interior_energetics/aragog.py +++ b/src/proteus/interior_energetics/aragog.py @@ -1233,6 +1233,33 @@ 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. That object is what + ``_step_heat_content`` integrates, which is the state side of the energy + budget, so a stale one misreports that 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). @@ -2104,6 +2131,10 @@ def _solve_with_retry(self, hf_row, interior_o) -> SolverOutput: # 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. t_start = float(solver.parameters.solver.start_time) diff --git a/tests/interior_energetics/test_aragog.py b/tests/interior_energetics/test_aragog.py index d3f6e3ef7..93dc3d63a 100644 --- a/tests/interior_energetics/test_aragog.py +++ b/tests/interior_energetics/test_aragog.py @@ -1487,6 +1487,64 @@ def test_an_interior_that_moves_under_fixed_radii_is_still_followed(monkeypatch) np.testing.assert_allclose(seen[1].radii, moved) +@pytest.mark.unit +@pytest.mark.physics_invariant +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 @pytest.mark.physics_invariant def test_regenerated_eos_tables_are_seen_even_at_identical_file_sizes(tmp_path): From 38ec624d9bfd18bac7feee599e2b87c3bf2f14a6 Mon Sep 17 00:00:00 2001 From: timlichtenberg Date: Sat, 8 Aug 2026 11:20:27 +0200 Subject: [PATCH 8/9] Refresh the tables before reset and test the reload wiring The compression-work diagnostic inside solver.reset() evaluates the re-read mesh pressures against whatever table object the solver holds, so on the step after an impact regenerates the tables it was evaluated against the outgrown set, which clamps at the old ceiling. The refresh now runs before reset(), and the later refresh inside the retry ladder stays as a cheap cache hit. Two pieces of wiring had no test: the ladder's refresh call (every ladder test ran with no solver wired, so removing the call changed nothing) and the translation of the one-shot impact flag into the per-step flag inside run_interior. Both are now pinned, the first with an ordering assertion that the refresh lands before the first solve, the second across two steps so one impact cannot exempt two. The docstring on the refresh helper now names the consumers that actually read the live table object (the RHS, the per-call energy integrals, and the compression diagnostic) rather than the impact heat booking, which deliberately runs on the pre-impact tables. A wrapper comment claimed a raised pressure ceiling cannot move the entropy range of the regenerated tables; the range is scanned up to the ceiling, so it can, and the comment now says so. Four structural tests carried the physics-invariant marker without asserting an invariant and lose it, and the test file docstring now covers the contract areas the file grew. --- src/proteus/interior_energetics/aragog.py | 34 ++++++------ src/proteus/interior_energetics/wrapper.py | 5 +- tests/interior_energetics/test_aragog.py | 60 ++++++++++++++++++++-- tests/interior_energetics/test_wrapper.py | 46 +++++++++++++++++ 4 files changed, 121 insertions(+), 24 deletions(-) diff --git a/src/proteus/interior_energetics/aragog.py b/src/proteus/interior_energetics/aragog.py index 43bab8f06..496895202 100644 --- a/src/proteus/interior_energetics/aragog.py +++ b/src/proteus/interior_energetics/aragog.py @@ -596,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: @@ -1239,11 +1243,12 @@ def _refresh_entropy_eos(config: Config, interior_o: Interior_t) -> None: 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. That object is what - ``_step_heat_content`` integrates, which is the state side of the energy - budget, so a stale one misreports that 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. + 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 ---------- @@ -1364,20 +1369,15 @@ 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, - # for the same reason the boundary conditions below are. A giant - # impact grows the planet and Zalmoxis re-solves the structure as - # the mantle freezes; either one replaces the mesh, and a copy - # taken once at install time would keep integrating the planet - # from before the change. + # 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. They are regenerated whenever the - # structure solve reruns, and a giant impact raises their - # pressure ceiling with the planet's mass, so a copy taken at - # install time would have the deep mantle clamped to the edge of - # a table built for a smaller planet. The loader is cached on the - # table parameters, so an unchanged table costs one small read. + # 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. diff --git a/src/proteus/interior_energetics/wrapper.py b/src/proteus/interior_energetics/wrapper.py index 36e197e74..41d80431a 100644 --- a/src/proteus/interior_energetics/wrapper.py +++ b/src/proteus/interior_energetics/wrapper.py @@ -3348,8 +3348,9 @@ def temperature_function(r, P): # and at the ungated regeneration a giant impact triggers. # # REMAINING GAP: the cached _last_entropy is not bounds-checked against the - # regenerated [S_min, S_max]. A composition change can move that range; a - # pressure-ceiling change alone does not. + # 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 diff --git a/tests/interior_energetics/test_aragog.py b/tests/interior_energetics/test_aragog.py index 93dc3d63a..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. @@ -1255,7 +1309,6 @@ def test_progress_is_weighed_against_what_the_coupling_asked_for(): @pytest.mark.unit -@pytest.mark.physics_invariant def test_the_core_temperature_guard_stands_aside_for_a_giant_impact(): """A giant impact's core-temperature jump is kept, not retried away. @@ -1410,7 +1463,6 @@ def test_the_jax_right_hand_side_reads_the_mesh_on_every_solve(monkeypatch): @pytest.mark.unit -@pytest.mark.physics_invariant def test_an_interior_that_moves_under_fixed_radii_is_still_followed(monkeypatch): """A structure change that leaves both bounding radii untouched is followed. @@ -1488,7 +1540,6 @@ def test_an_interior_that_moves_under_fixed_radii_is_still_followed(monkeypatch) @pytest.mark.unit -@pytest.mark.physics_invariant 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. @@ -1546,7 +1597,6 @@ def test_the_solver_is_pointed_at_the_current_tables_before_each_solve(tmp_path) @pytest.mark.unit -@pytest.mark.physics_invariant 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. 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 From 0be0464784c5b872d5138ef376b098eb9cee9855 Mon Sep 17 00:00:00 2001 From: timlichtenberg Date: Sun, 23 Aug 2026 18:34:18 +0200 Subject: [PATCH 9/9] Bound the time step that lands on a coarse-phase giant impact The step that lands exactly on a scheduled giant impact is clamped only to whatever simulated time remains before it. After a long quiescent phase has let dt coarsen, that remaining time can itself be large, so the step absorbing the impact's melt-fraction jump inherits the same coarseness the run had just grown into. campaign-2 hit this on probe18_capsoff_guardtip: an impact landed while dt had coarsened to 1.06e6 yr, the remaining-time clamp only cut it to 2.14e5 yr, and CVODE burned two failed attempts before succeeding at 4.27e4 yr. I added dt.impact_maximum, an optional absolute ceiling on the landing step, applied before the existing minimum-step floor so a misconfigured ceiling smaller than the floor still can't collapse dt to zero near an impact. It defaults to 0 (disabled), so no existing run's behaviour changes unless I set it. Reproduced the coarse-phase landing cheaply with the dummy interior and the analytical accretion module: a proportional-dt run reaching an impact after dt had grown to ~8e5 yr lands on it at ~3.35e5 yr uncapped, and at the configured ceiling once impact_maximum is set. Added unit tests for the new ceiling to test_timestep.py, including that it never beats the minimum-step floor, and confirmed each one fails against the unfixed clamp before passing against this change. I have not reproduced the CVODE retry-ladder itself against the real Aragog solver; the tests above verify the dt-sizing mechanism a coarse-phase impact triggers, not the solver's response to it. --- src/proteus/config/_params.py | 11 +++ src/proteus/interior_energetics/timestep.py | 15 ++-- tests/interior_energetics/test_timestep.py | 94 +++++++++++++++++++++ 3 files changed, 113 insertions(+), 7 deletions(-) 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/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/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.