CAMulator

GlobalCM202524 GBNCARPyTorch

Import path: earth2studio.models.px.CAMulator

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin, PrognosticMixin

CAMulator: NSF NCAR's autoregressive emulator of the CAM6 atmosphere, trained in the CREDIT framework. A CrossFormer encoder-decoder advances a 1 degree (192x288), 32 hybrid-sigma-level atmospheric state by 6 hours given prescribed sea-surface temperature, sea-ice fraction, TOA insolation and CO2 forcing, for climate-length rollouts. The inference-time post-processing of the CREDIT toolbox is reproduced: tracer clipping, the jet wind-artifact filter and the global dry-air mass, water and total-energy conservation fixers.

The 130 prognostic variables are the model input; the output additionally holds 17 output-only diagnostics (precipitation, surface temperature, cloud fractions, surface stresses, 10 m wind speed, evaporation and radiative/heat fluxes). In create_iterator the initial condition is yielded in the output schema with the diagnostics NaN-filled, and only the prognostic slice is fed back at each step. Vertical levels use the {var}{k}k naming with k the CAMulator hybrid level index (0 = top of model, 31 = lowest layer). Fluxes CAMulator accumulates over the 6 h step (J m-2) are returned as mean rates (W m-2); precipitation and evaporation are 6 h accumulations (m).

Note

Forcing is read at the input valid time from forcing_data_source on the CAMulator grid. The default CAMulatorForcing serves the shipped climatological (cyclic) year; the forcing files use a 365-day calendar, so leap days reuse the 28 February forcing.

Note

For more information see:

• arxiv.org/abs/2504.06007
• huggingface.co/willychap/camulator
• NCAR/miles-credit

Parameters:

• core_model (Module) –

  CAMulator CrossFormer network mapping (batch, 136, 1, lat, lon) to (batch, 147, 1, lat, lon) in normalized units.

• center (Tensor) –

  Normalization mean of the 147 output channels, shape (147, lat, lon).

• scale (Tensor) –

  Normalization std of the 147 output channels, shape (147, lat, lon).

• forcing_center (Tensor) –

  Normalization mean of the 4 forcing channels, shape (4,).

• forcing_scale (Tensor) –

  Normalization std of the 4 forcing channels, shape (4,).

• tracer_center (Tensor) –

  Normalization mean used by the tracer clipping, one per tracer channel.

• tracer_scale (Tensor) –

  Normalization std used by the tracer clipping, one per tracer channel.

• statics (Tensor) –

  Static input fields z_norm and LANDM_COSLAT, shape (2, lat, lon).

• hyai (Tensor) –

  Hybrid a interface coefficients (Pa), shape (33,).

• hybi (Tensor) –

  Hybrid b interface coefficients, shape (33,).

• area (Tensor) –

  Grid-cell area (m^2), shape (lat, lon).

• phis (Tensor) –

  Surface geopotential (m^2 s-2), shape (lat, lon).

• forcing_data_source (DataSource, default: None ) –

  Data source providing mtdwswrf, sst, sic and global_mean_co2 on the CAMulator grid, by default CAMulatorForcing().

• conservation_fixers (bool, default: True ) –

  Apply the global mass, water and energy fixers, by default True.

• wind_filter (bool, default: True ) –

  Apply the wind-artifact filter, by default True.

__call__

__call__(
    x: Tensor, coords: CoordSystem
) -> tuple[Tensor, CoordSystem]

Runs prognostic model 1 step (6 hours)

Parameters:

• x (Tensor) –

  Input tensor

• coords (CoordSystem) –

  Input coordinate system

Returns:

• tuple[Tensor, CoordSystem] –

  Output tensor and coordinate system 6 hours in the future

create_iterator

create_iterator(
    x: Tensor, coords: CoordSystem
) -> Iterator[tuple[Tensor, CoordSystem]]

Creates a iterator which can be used to perform time-integration of the prognostic model. Will return the initial condition first (0th step) in the output variable schema with the diagnostic variables set to NaN.

Parameters:

• x (Tensor) –

  Input tensor

• coords (CoordSystem) –

  Input coordinate system

Yields:

• Iterator[tuple[Tensor, CoordSystem]] –

  Iterator that generates time-steps of the prognostic model container the output data tensor and coordinate system dictionary.

load_default_package classmethod

load_default_package() -> Package

Load the default CAMulator package from HuggingFace

Returns:

• Package –

  Package

load_model classmethod

load_model(
    package: Package,
    checkpoint: str = "checkpoint.pt00069.pt",
    forcing_data_source: DataSource | None = None,
    conservation_fixers: bool = True,
    wind_filter: bool = True,
) -> PrognosticModel

Load prognostic model from package

Parameters:

• package (Package) –

  Package to load model from

• checkpoint (str, default: 'checkpoint.pt00069.pt' ) –

  Checkpoint file in the package. The HuggingFace repository hosts several training epochs (checkpoint.pt000NN.pt); epoch 69 is the model card default, by default "checkpoint.pt00069.pt"

• forcing_data_source (DataSource, default: None ) –

  Forcing data source, by default CAMulatorForcing()

• conservation_fixers (bool, default: True ) –

  Apply the global mass, water and energy fixers, by default True

• wind_filter (bool, default: True ) –

  Apply the wind-artifact filter, by default True

Returns:

• PrognosticModel –

  Prognostic model