jcm.model.Model

class jcm.model.Model(dycore: DynamicalCore | None = None, *, coords=None, time_step: float = 30.0, terrain: TerrainData = None, physics: Physics = None, start_date: Datetime = jax_datetime.Datetime(delta=jax_datetime.Timedelta(days=10957, seconds=0)), calendar: str = '365_day', radiation_chunk_size: int | None = None, log_level=50)

Top level class for a JAX-GCM simulation.

The Model orchestrates the run (timestep, forcing, op-split scan, post-processing). Dynamics-specific work (state init, the per-dt step, the spectral↔gridpoint bridge) is delegated to a DynamicalCore; physics-specific work to a Physics.

__init__(dycore: DynamicalCore | None = None, *, coords=None, time_step: float = 30.0, terrain: TerrainData = None, physics: Physics = None, start_date: Datetime = jax_datetime.Datetime(delta=jax_datetime.Timedelta(days=10957, seconds=0)), calendar: str = '365_day', radiation_chunk_size: int | None = None, log_level=50) → None

Initialise the model.

Args:

dycore: The DynamicalCore driving the integration. When

None, a default DinosaurDycore is constructed from coords and terrain for convenience. Backend- specific knobs (diffusion, nudging-as-PhysicsTerm targets, IMEX stepper details) belong to the dycore’s own constructor — wire them there, then pass the dycore in.

coords: CoordinateSystem. Required when dycore is None.

To enable SPMD parallelization, pass spmd_mesh to the coords helper (e.g. get_speedy_coords()).

time_step: Model time step in minutes. terrain: TerrainData (orography, land-sea mask, etc.).

Defaults to an aquaplanet when building the default dycore.

physics: Physics describing the model physics. Defaults

to speedy_physics(). Add nudging via the jcm.nudging.NudgingTerm PhysicsTerm.

start_date: jax_datetime.Datetime for the start of the run.

Used to convert state.sim_time to a DateData that’s threaded into the physics-step diagnostics dict (so forcing-driven and date-aware terms can read it).

calendar: Calendar string ("365_day" or "gregorian") for

the same date conversion.

radiation_chunk_size: Override the RRTMGP chunked-vmap chunk size

(cells per chunk). Default None auto-detects from the JAX device’s HBM. No-op when the active radiation scheme is not RRTMGP.

log_level: Logging verbosity level.

Methods

__init__([dycore, coords, time_step, ...])	Initialise the model.
bootstrap_state([initial_state])	Populate _final_dycore_state and _final_physics_state without integrating.
resume([forcing, save_interval, total_time, ...])	Continue from end of previous run / resume.
run([initial_state, forcing, save_interval, ...])	Set the initial state and run the full simulation forward in time.
run_from_state(initial_state, forcing[, ...])	Run the simulation forward from a given dycore-native initial state.
run_from_state_with_carry(initial_state, forcing)	Lower-level run_from_state that exposes the cross-step physics carry.