Skip to content

Architecture — gcm3d

Component Overview

The reusable GCM implementation lives under src.framework; Mars-specific configuration, data adapters, maps, and outputs live under src.celestials.planets.mars.

Layer Modules Depends on JAX?
Body abstraction framework/gcm/body.py (BodyConstants, EARTH) No — always importable
Dependency guard framework/gcm/_dinosaur.py — (raises if extra missing)
Discretisation framework/gcm/coordinates.py (spectral grid × sigma), framework/gcm/specs.py (nondimensionalisation) Yes
Dry dynamical core framework/gcm/dynamics.py (primitive_equations, stepper, integrate) Yes
Reusable physics framework/physics/gcm.py (radiation, surface energy, drag, PBL, convection, condensable cycle), framework/physics/ames_radiation.py (correlated-k operators/tables) Yes
Mars configuration celestials/planets/mars/planet.py (MARS_BODY_3D) and celestials/planets/mars/gcm.py (forcing factories) Yes
Generic diagnostics framework/gcm/benchmarks.py, framework/gcm/restart.py Yes
Mars boundary data celestials/planets/mars/topography.py, surface.py, dust.py, mcd.py topography: Yes; surface/dust: NumPy+xarray
Mars experiments celestials/planets/mars/maps.py, seasonal.py Yes
Parameterized rollouts framework/gcm/learning.py Yes
Neural components framework/neural/ (MLP, radiation, heating, checkpoints) Yes
Calibration application apps/mars-calibration/mars_calibration/ Yes

The dependency direction is one-way: Mars imports and configures framework operators; framework code never imports src.celestials. The pre-release top-level src.gcm3d package was removed rather than retained as an alias.

Flow Diagram — the forced 3-D map run

flowchart TD
    BC[BodyConstants MARS_BODY_3D] --> SPECS[physics_specs<br/>nondimensionalisation]
    BC --> COORD[coordinate_system<br/>spectral grid × sigma]
    MOLA[MOLA MEGDR raster] --> ORO[mola_modal_orography<br/>lower boundary]
    COORD --> ORO
    COORD --> REST[initial_rest_state<br/>terrain-balanced p_s]
    SPECS --> REST

    subgraph EQ[forced_co2_primitive_equations = dinosaur ⊕ physics]
      DRY[dinosaur PrimitiveEquationsSigma<br/>implicit = gravity waves] --> EXPL[explicit_terms]
      PHYS[column physics parameterization] --> EXPL
    end
    ORO --> EQ
    REST --> STATE0[ColumnPhysicsState<br/>+ surface_T, co2_ice, ground_T]

    STATE0 --> SCAN[stepper imex_rk_sil3<br/>integrate = jax.lax.scan]
    EQ --> SCAN
    SCAN --> FILTER[optional spectral diffusion]
    FILTER -->|per step| PROJ[positivity_preserving_co2_step<br/>frost positivity + global CO2 inventory]
    PROJ --> FINAL[final ColumnPhysicsState]
    FINAL --> MAPS[MarsMapFields<br/>modal→nodal, dimensionalize]
    MAPS --> OUT[save_netcdf + plot_maps]

The physics parameterization (one step, per grid column) sums, in order: radiative surface energy + latent CO₂ heating (temperature/surface tendency), surface drag + PBL vertical diffusion (momentum), dry convective adjustment (temperature), and the CO₂ mass exchange (d ln p_s, frost, latent surface heat).

Key Architectural Choices

  1. Physics is added to explicit_terms, never the implicit side.
  2. Why: dinosaur's implicit solve is the gravity-wave semi-implicit linearisation; physics is a slow forcing. Keeping it explicit preserves the dycore's stability contract and its implicit_inverse unchanged.

  3. Alternative rejected: relaxing to a prescribed rest temperature (Held–Suarez style) — that cannot develop the model's own subsolar-hot / nightside-cold structure, which is what drives the circulation.

  4. ColumnPhysicsState = dinosaur State + non-advected surface reservoirs.

  5. Why: surface temperature, CO₂ frost, and multilayer ground temperature sit on the ground and must not be horizontally advected. Wrapping the dycore state in a JAX-pytree NamedTuple keeps the dynamics intact while carrying them.

  6. Alternative rejected: dinosaur tracers — those are transported by wind.

  7. Nondimensionalisation is body-anchored (length=radius, time=1/2Ω).

  8. Why: dinosaur's spectral Grid assumes radius 1 and Ω=0.5. Building the scale from each BodyConstants (specs.py) makes any body satisfy that without patching dinosaur's Earth-hardcoded scales.py.

  9. CO₂ mass exchange changes atmospheric mass through log_surface_pressure.

  10. Why: condensing CO₂ to frost removes atmospheric mass locally; the tendency is applied as d(ln p_s)/dt on the dycore's own prognostic surface pressure. Phase exchange balances atmosphere and frost instantaneously. After each complete energy-limited step, local frost repair and a Gaussian-quadrature global pressure rescaling restore the inventory while retaining configured escape. This corrects finite-step drift from advancing log(p_s); it is not a guarantee of exact local discrete transport.

Dependencies

  • dinosaur (JAX): spectral transforms, sigma coordinates, PrimitiveEquationsSigma, imex_rk_sil3, published initial states (isothermal_rest_atmosphere, steady_state_jw), held_suarez forcing, xarray_utils.
  • JAX: autodiff, lax.scan, vmap, jit.
  • NumPy / xarray: raster IO, NetCDF export, boundary-field regridding.
  • matplotlib (plot_maps only, Agg backend).
  • Staged assets: MOLA MEGDR raster (data/mola/…), Ames CO₂ 12-band table (framework/physics/ames_co2_12band.npz, bundled), optional TES surface + Ames dust NetCDFs.

Extension Points

  • New body: define BodyConstants for the generic dry core, then supply appropriate forcing and boundary data. Mars-specific radiation and frost parameters do not become valid for other compositions automatically.
  • New / learned column physics: implement a parameterization(state) -> ColumnPhysicsTendencies and wrap with column_primitive_equations — the NeuralGCM-style residual-tendency seam.
  • Better radiation: two_stream_radiative_fluxes isolates the optical-depth closure; correlated-k coefficients replace it without changing callers.
  • New surface/dust data: surface.py / dust.py interpolate any conforming dataset onto the grid.

Training and deployment

Use make_parameterized_step and rollout from framework.gcm.learning for parameter gradients and sampled/rematerialized trajectories. The reporting and dataset I/O in run_maps belong outside differentiated training objectives. Radiation hooks replace a flux calculation; bounded neural tendencies compose with conventional physics. See neural interfaces.

The temperature-only pipeline trains on cached forecasts and does not modify the timestep. Calibration differentiates physical controls through coupled integration. Keep these experiment contracts separate when comparing results.