Skip to content

SamudrACE

Import path: earth2studio.models.px.SamudrACE

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin, PrognosticMixin

SamudrACE coupled climate emulator.

Couples the ACE2 atmosphere emulator with the Samudra ocean emulator, driven by FME's CoupledStepper. The atmosphere advances in 6 hour steps and the ocean advances once per coupled (5 day) cycle; all coupling logic (SST prescription, flux exchange and averaging, ocean-fraction prediction, masking) is owned by FME via CoupledStepper.predict, which this wrapper calls exactly once per coupled cycle.

The primary interface is create_iterator, which yields the initial condition followed by one atmosphere step at a time. Atmosphere output fields update on every step; ocean output fields update once per coupled cycle and are held constant between cycle boundaries. Ocean diagnostic (non-prognostic) fields are NaN until the first cycle boundary, as are atmosphere diagnostic fields at the initial condition step.

The model is iterator-only: it has no single-step forward call. Its time-step is the coupled cycle, which spans n_inner_steps atmosphere steps, so advancing the model is not expressible as one 6 hour forward pass; calling the model raises NotImplementedError.

Times are CM4 model years (e.g. year 151), which are outside the range of nanosecond-precision timestamps; provide time coordinates as second-precision np.datetime64 values. Trajectories that cross a February 29 of the proleptic Gregorian time coordinate raise an error when the forcing is looked up, since that date has no counterpart on the no-leap forcing calendar.

cuDNN autotuning (torch.backends.cudnn.benchmark) is left at its default (disabled). Toggling it on and off around each coupled cycle was found to make inference non-reproducible run-to-run: with a run-specific, GPU-timing-based algorithm search re-triggered every cycle, cuDNN can select a different (numerically non-identical) algorithm on different runs, which the coupled system's chaotic dynamics then amplify over the forecast.

Note

For more information see the following references:

Parameters:

  • stepper (CoupledStepper) –

    FME coupled stepper holding the atmosphere and ocean component steppers and the coupling configuration. Both components must be on the same latitude-longitude grid.

  • forcing_data_source (DataSource) –

    Data source providing all exogenous forcing variables required by the coupled stepper (both atmosphere and ocean), on the model grid and with the model's Earth2Studio variable names (see SamudrACELexicon).

__call__ 

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

Not supported; SamudrACE is an iterator-only prognostic.

The model's time-step is the coupled (ocean) step, which spans n_inner_steps atmosphere steps, so a single-step forward call cannot express advancing the model: advancing a full coupled cycle to return one 6 hour step would discard the remainder of the cycle and report ocean fields that have not advanced. Use create_iterator instead, which yields one atmosphere step at a time and runs the coupled stepper at each cycle boundary.

Parameters:

  • x (Tensor) –

    Input tensor

  • coords (CoordSystem) –

    Input coordinate system

Returns:

Raises:

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).

The iterator yields one atmosphere (6 hour) step at a time. The coupled stepper runs lazily at each coupled cycle boundary; ocean output fields update once per cycle and are held constant between boundaries.

Parameters:

  • x (Tensor) –

    Input tensor

  • coords (CoordSystem) –

    Input coordinate system

Yields:

  • Iterator[tuple[Tensor, CoordSystem]] –

    Iterator of output tensors and coordinate systems

load_default_package classmethod 

load_default_package() -> Package

Load the default SamudrACE model package from HuggingFace.

Returns:

  • Package –

    Package holding the allenai/SamudrACE-CM4-piControl checkpoint, pinned to a specific repository revision

load_model classmethod 

load_model(
    package: Package,
    forcing_data_source: DataSource | None = None,
    scenario: str = "0311",
) -> PrognosticMixin

Load SamudrACE prognostic from package.

Parameters:

  • package (Package) –

    Package holding the coupled model checkpoint

  • forcing_data_source (DataSource | None, default: None ) –

    External data source providing all exogenous forcing variables required by the coupled stepper, on the model grid and with the model's Earth2Studio variable names. If None, a SamudrACEForcingData source for the selected scenario is used, by default None

  • scenario (str, default: '0311' ) –

    Forcing scenario for the default forcing data source, either "0151" or "0311"; ignored when forcing_data_source is provided, by default "0311"

Returns:

  • PrognosticMixin –

    Prognostic model