Skip to content

Turn on Zalmoxis Anderson acceleration by default and require aragog 26.10.08 - #971

Merged
timlichtenberg merged 14 commits into
mainfrom
tl/speedup-963
Oct 8, 2026
Merged

timlichtenberg merged 14 commits into
mainfrom
tl/speedup-963

Conversation

@timlichtenberg

Copy link
Copy Markdown
Member

Description

Option 2 of #963: Zalmoxis Anderson acceleration (interior_struct.zalmoxis.use_anderson) is now on by default, and the structure call without temperature data (the setup solve and the equilibration) keeps it on. That call still turns use_jax off, because the internal temperature mode converges to a different structure on the JAX path; Anderson needs no such exclusion on either path. The four input files that pinned use_anderson = false as the old default now carry true, and the schema and the config docs follow.

This PR also raises the fwl-aragog floor to 26.10.08, the release that caches the jitted CVODE RHS and Jacobian and the parsed EOS tables (FormingWorlds/aragog#73).

Together these are options 1, 2 and 8 of #963; option 3 (vectorised energy integrals) follows separately. Outputs are not bit-identical to main: the changes move the coupled trajectory at the level of a 1 % change of the Aragog rtol, which is the acceptance level set in #963.

Validation of changes

Paired coupled Earth runs on the same machine, started together, this branch with the paired aragog branch against main and against main with rtol x 1.01.

Dynamic-structure Earth coupled run: T_magma, Phi_global, F_int, R_int, M_int and P_cmb against time for the Anderson default, main, and main with rtol x 1.01

Anderson default alone (option 2), dynamic structure:

  • Phi 0.05 crossing time +0.26 % against main; the median trajectory differences stay below 1.13 times the rtol spread, the maxima below the rtol spread.
  • Structure parity on a 20-step dynamic run with 3 re-solves, same step sequence on both arms: R_int agrees to 2e-7 and M_int to 1e-6 (JAX path), R_int to 6e-8 and M_int to 3e-6 (numpy path). All structure solves converge on both arms.
  • Wall time of that 20-step run: 24.3 to 22.2 min on the JAX path, 36.7 to 12.9 min on the numpy path (use_jax = false), where every re-solve runs with Anderson and the density seed.

All three options together: the static run goes from 26 to 11 min. In the dynamic runs, every median difference on a uniform time grid stays within twice the spread of two rtol perturbations, and the Phi 0.05 crossing (+0.15 to +0.17 %) lies inside their range (+0.004 % and +1.64 %). Single-column maxima exceed that spread at local features: an F_cmb sawtooth at 1.8 to 3.8 kyr in one run that does not repeat in a second run, and the onset time of the T_magma drop near 50 kyr.

Tests: tests/interior_struct/test_zalmoxis.py (the guard keeps Anderson and turns off JAX for a call without temperature data, over use_anderson x outer_solver), tests/interior_struct/test_zalmoxis_outer_solver_config.py (use_anderson reaches the solver configuration), tests/config/test_struct.py (default). The nightly science validation on this branch with the aragog 26.10.08 release from PyPI is green (25 of 25 jobs): the unit tier reports 5081 passed and the integration tier 206 passed on both runners.

Tested on macOS 26 (Apple silicon), Python 3.12.

Checklist

  • I have followed the contributing guidelines
  • My code follows the style guidelines of this project
  • I have performed a self-review of my code
  • My changes generate no new warnings or errors
  • I have checked that the tests still pass on my computer
  • I have updated the docs, as appropriate
  • I have added tests for these changes, as appropriate
  • I have checked that all dependencies have been updated, as required

Set use_anderson to true by default in interior structure configuration. Preserve Anderson acceleration setting when temperature profile data is absent.
Update solver configuration comments to describe the active defaults for Anderson acceleration and JAX dispatch. Clarify that the temperature profile guard disables JAX across all outer solver modes. Document density caching behavior on the NumPy and JAX execution paths.
Clarify that both JAX and Anderson acceleration default on.
Verify that calls without temperature profiles disable JAX while preserving Anderson acceleration under both Newton and Picard outer solvers.
Assert that calls without temperature data preserve Anderson acceleration across both outer solvers and Anderson configuration settings.
Verify that the Zalmoxis solver guard preserves outer solver settings across test configurations without temperature data. Clarify docstring to note that the guard disables the JAX path for every outer solver.
@timlichtenberg
timlichtenberg requested a review from a team as a code owner October 8, 2026 08:22
@codecov

codecov Bot commented Oct 8, 2026 •

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 93.66%. Comparing base (4fc66f8) to head (bef40f8).

Additional details and impacted files
@@            Coverage Diff             @@
##             main     #971      +/-   ##
==========================================
- Coverage   93.97%   93.66%   -0.32%     
==========================================
  Files         128      128              
  Lines       20891    20890       -1     
  Branches     3515     3515              
==========================================
- Hits        19633    19567      -66     
- Misses        908      914       +6     
- Partials      350      409      +59     
Flag Coverage Δ
nightly 93.63% <100.00%> (+0.08%) ⬆️
unit-tests 92.28% <100.00%> (+0.01%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

The guard test that checks calls without temperature data now runs with use_jax True and False, so the gate in zalmoxis_solver is exercised on both arms. The gate test helper takes a use_jax argument to set the config value.
@timlichtenberg
timlichtenberg merged commit e020e13 into main Oct 8, 2026
48 checks passed
@timlichtenberg
timlichtenberg deleted the tl/speedup-963 branch October 8, 2026 14:45
IKisvardai added a commit that referenced this pull request Oct 9, 2026
Brings in the prescribed orbital migration (#937), the CALLIOPE cold-start
seed (#958), the fwl-io dataset fetch and stellar spectra changes (#960,
#964, #967, #968), the PETSc download (#965), Zalmoxis Anderson acceleration
with fwl-aragog 26.10.08 (#971) and the CI changes (#955, #970). The schema
invariant tests keep both the trapping and the migration validators, and the
golden trajectory is main's re-recorded one with the trapping columns added.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant