FuXiS2S

GlobalS2S202440 GBONNX

Import path: earth2studio.models.px.FuXiS2S

View source on GitHub View install commands

Documentation

Bases: Module, AutoModelMixin, PrognosticMixin

FuXi-S2S global daily-mean prognostic model.

FuXi-S2S consumes daily means from two consecutive UTC calendar days and
predicts the following daily mean. A timestamp at 00:00 UTC labels the
corresponding calendar-day aggregate; it is not an instantaneous midnight
state.

Note

This model uses the ONNX checkpoint from the original publication repository. For
additional information see the following resources:

- https://www.nature.com/articles/s41467-024-50714-1
- https://github.com/tpys/FuXi-S2S
- https://zenodo.org/records/15718402
- https://huggingface.co/datasets/FudanFuXi/FuXi-S2S

Parameters

onnx_path : str
    Path to the FuXi-S2S ONNX graph. Its external weight file named
    ``fuxi_s2s`` must be in the same directory.

Note

Initial conditions must contain two consecutive UTC daily means on the
model's 1.5-degree grid. Instantaneous fields use calendar-day averages
from 00--23 UTC. Accumulated ``tp`` and ``ttr`` fields use the 24
interval-ending values from 01 UTC through 00 UTC of the following day.
Their daily means retain the units of each one-hour accumulation; for
example, multiply predicted ``tp`` by 24 to obtain a daily total.
Sea-surface temperature must retain ``NaN`` values over land. The wrapper
does not aggregate hourly fields or regrid initial conditions; callers must
provide these prepared daily inputs through an Earth2Studio data source.

The official ONNX graph samples flow-dependent perturbations internally, so
each forecast trajectory is one stochastic ensemble member. Member ``00`` in
the official inference script is the first stochastic member, not a
deterministic control.

Warning

We encourage users to familiarize themselves with the license restrictions of this
model's checkpoints.

Example

The following shows how to get approximate inputs from standard data
sources.  For formal input data consult

zenodo.org/records/15718402

```python
class DailyMeanARCO:
    def __init__(self) -> None:
        self.arco = ARCO_ERA5()

    def __call__(
        self,
        time: datetime | list[datetime] | TimeArray,
        variable: str | list[str] | VariableArray,
    ) -> xr.DataArray:
        time, variable = prep_data_inputs(time, variable)
        daily_arrays = []
        for day in time:
            # Approximate daily means from 6-hourly ARCO ERA5
            sub_daily = [day + timedelta(hours=h) for h in (0, 6, 12, 18)]
            da_mean = self.arco(sub_daily, variable).mean("time", skipna=True)
            daily_arrays.append(
                da_mean.interp(lat=S2S_LAT, lon=S2S_LON, method="linear")
            )
        da_out = xr.concat(daily_arrays, dim="time")
        da_out = da_out.assign_coords(
            time=np.array(time, dtype="datetime64[ns]")
        )
        return da_out


# Load the model
package = FuXiS2S.load_default_package()
model = FuXiS2S.load_model(package).to("cuda:0")

# Run a single deterministic step
import earth2studio.run as run
from earth2studio.io import ZarrBackend

io = run.deterministic(["2024-01-15"], 1, model, DailyMeanARCO(), ZarrBackend())
```

__call__

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

Run FuXi-S2S one daily step.

Parameters:

• x (Tensor) –

  Two consecutive UTC daily means.

• coords (CoordSystem) –

  Coordinates describing the input tensor.

Returns:

• tuple[Tensor, CoordSystem] –

  Predicted next daily mean and its coordinates.

create_iterator

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

Create a daily FuXi-S2S forecast iterator.

Parameters:

• x (Tensor) –

  Two consecutive UTC daily means.

• coords (CoordSystem) –

  Coordinates describing the input tensor.

Yields:

• Iterator[tuple[Tensor, CoordSystem]] –

  Initial current day followed by successive daily predictions.

load_default_package classmethod

load_default_package() -> Package

Load the FuXi-S2S package from an immutable Hugging Face mirror.

Returns:

• Package –

  Package pointing to the mirrored FuXi-S2S checkpoint.

Note

The mirror contains unchanged assets from the official Zenodo record. The checkpoint is licensed CC BY-NC-ND 4.0 and restricted to non-commercial research use by its authors.

load_model classmethod

load_model(package: Package) -> PrognosticModel

Load FuXi-S2S from an Earth2Studio package.

Parameters:

• package (Package) –

  Package containing fuxi_s2s.onnx and its external data file.

Returns:

• PrognosticModel –

  Loaded FuXi-S2S prognostic wrapper.