[WIP] feat : integrate JCM physics with MOLA topology

[sameeerkashyap]

Pull Request

Summary

• This PR includes a clear description of the change.
• This PR explains the motivation or problem being addressed.
• This PR describes the main implementation details.
• Related issue(s), discussion(s), or task(s) are linked.

What changed?

This PR adds the JCM (jax-gcm) 3-D general circulation model as a new simulation backend for Mars. It runs behind the existing Planet / TimeController / Snapshot interface. The 0-D PyTorch box model stays intact.

Main additions:

• New dependency: jcm (git), in package/pyproject.toml and uv.lock.
• MarsGCM (mars/gcm.py) — a Mars subclass. It builds a JCM model, steps it, and reduces the 3-D fields to the scalar mirror that Snapshot reads.
• Placeholder physics (mars/gcm_physics.py) — MarsHeldSuarez, a Mars-tuned Held-Suarez term, plus the mars_physics() factory. It needs no boundary data.
• MOLA terrain (mars/terrain.py) — a loader that regrids MOLA topography onto the JCM grid. MarsGCM uses it when the data is present.
• Constants bridge — pushes the Mars constants into JCM through MARS.as_jcm_overrides().
• Engine seam — Accuracy.GCM and one branch in evolve() that calls planet.step_gcm(dt). A default step_gcm on the base Planet raises.
• Download script — scripts/download_mola_topography.py.
• Docs — docs/ideas/mars-neural-gcm-plan.md (a plan for a NeuralGCM-style Mars model) and a new integration section in the differentiable-framework doc.

Why was this needed?

The 0-D box model gives one global temperature and pressure. It cannot show spatial climate: winds, the topographic pressure range, or polar processes. Mars terraforming needs a 3-D, differentiable model for inverse design later. JCM is built on the same dinosaur differentiable core as NeuralGCM, so it fits this goal.

Implementation notes

• The adapter keeps every existing interface. The engine, CLI, and plots are unchanged.
• The seam is evolve(), not the RK4 path. A 3-D spectral state cannot use the 0-D compute_derivatives arithmetic, so JCM owns its stepping.
• step_gcm converts the engine dt (seconds) to JCM days, calls model.resume(...), then writes an area-weighted global mean into the mirror.
• New dependency: jcm from git. No breaking changes: FAST and ACCURATE still work.
• Two known gaps (follow-up): the GCM starts from its own 288 K state, not the surface_temperature kwarg; and the CO₂-ice mass is not reduced yet (the placeholder has no CO₂ tracer).

Tests

• I added or updated unit tests for the changed logic.
• All relevant tests pass locally.

Test details

Four new test files under package/tests/mars/basic/:

• test_gcm_constants.py — the constants bridge maps Mars values into JCM; universal and water constants stay at Earth values.
• test_gcm_physics.py — mars_physics() builds one MarsHeldSuarez term; it drives a model one step with no NaN.
• test_terrain.py — the MOLA loader rejects a bad grid; it regrids to the model shape; the elevation range stays Mars-like; the geopotential is non-zero.
• test_gcm.py — MarsGCM constructs; step_gcm updates the mirror; TimeController with Accuracy.GCM returns a Snapshot list; the base planet step_gcm raises.

Commands run:
cd package && uv run python -m pytest tests/mars/basic/ -q # new tests pass
cd package && uv run python -m pytest tests/engine -q -m "not slow" # FAST/ACCURATE unaffected
uv run pyright package/src/celestials/planets/mars/ package/src/engine/time_controller.py # 0 errors
Slow tests (build a JCM model) carry @pytest.mark.slow. MOLA tests skip when the data is absent, so CI stays data-free.

Why are no tests needed?

N/A.

Simulations, examples, and observations

• I ran the relevant simulation(s) for this change.
• I included the config, script, notebook, or command needed to reproduce the result.
• I included example outputs such as plots, logs, tables, or screenshots.
• I summarized the main observations from the results.
• I compared results against a baseline or prior behavior when relevant.

Simulation setup

• Grid: Held-Suarez grid (T31, 8 layers). Dynamical step: 180 s.
• Constants: Mars (gravity 3.72, radius 3.39e6, rotation 7.09e-5, CO₂ cpd 770, p0 608 Pa, solar flux 589).
• Terrain: MOLA 4 px/deg global topography, regridded to the model grid.
• Physics: the MarsHeldSuarez placeholder term.
• Hardware: CPU (JAX on CPU).

Reproduction steps

1. Get the MOLA topography

python scripts/download_mola_topography.py

2. Drive MarsGCM through the existing engine

cd package && uv run python - <<'PY'
from src.celestials.planets.mars.gcm import MarsGCM
from src.engine import TimeController, Accuracy
tc = TimeController(MarsGCM(), dt=3600.0, accuracy=Accuracy.GCM)
history = tc.run(duration=3600.0 * 3) # three one-hour steps
for s in history:
print(float(s.surface_temperature), float(s.surface_pressure))
PY

Results / artifacts

• outputs/mola_regrid_check.png — source MOLA next to the regridded output. Hellas, Tharsis, Olympus, and Argyre all line up.
• Timestep timing: the first resume compiles in about 6 s; each later step runs in about 0.3 s.

Observations

• MarsGCM runs behind TimeController and returns the same Snapshot list as the 0-D model.
• The surface pressure evolves across steps (615 → 622 → 630 Pa). So the dynamics run over the topography, and the reduction tracks them.
• The MOLA regrid keeps the elevation range: −7370 m (Hellas) to +16315 m (Tharsis). The coarse grid smooths the peaks, as expected.
• Baseline comparison: the existing FAST and ACCURATE strategies are unchanged and still pass. This PR adds a backend; it does not change one.

Reviewer notes

• This PR is ready for review.
• I want feedback on a specific design or modeling decision.
• This PR contains breaking changes.
• This PR requires follow-up work.

Specific feedback requested

• The reduction: surface temperature uses the lowest model level as a proxy; surface pressure uses normalized_surface_pressure × p0. Please confirm this is the right global-mean definition for the mirror.
• The initial state: should MarsGCM seed the 3-D state from the surface_temperature kwarg, or keep JCM's default? This changes how the CLI presets map onto the GCM.
• The placeholder timestep (180 s) and the Held-Suarez grid are provisional. They change with the real Mars grid.

Follow-up work (later steps)

1. Seed the 3-D initial state from the scalar config.
2. Reduce the CO₂-ice mass (needs the co2_cycle tracer).
3. Real Mars physics: CO₂ cycle, dust, radiation. Add the JCM patches (mass tendency, gravity, threaded constants).
4. Validate the seasonal surface-pressure cycle against Viking.