DataReplay

GlobalNVIDIAPyTorch

Import path: earth2studio.models.px.DataReplay

View source on GitHub

Documentation

Bases: Module, PrognosticMixin

Replay a data source through the prognostic model interface.

Data sources are queried at successive valid times. Forecast sources are queried at successive lead times from the initial time.

Parameters:

• source (DataSource | ForecastSource) –

  Source to replay.

• variable (str | list[str]) –

  Variables to fetch.

• domain_coords (CoordSystem) –

  Spatial coordinates expected from the source.

• step (timedelta64, default: timedelta64(6, 'h') ) –

  Time between frames, by default np.timedelta64(6, "h")

__call__

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

Advance the source by one step.

Parameters:

• x (Tensor) –

  Input tensor.

• coords (CoordSystem) –

  Input coordinate system.

Returns:

• tuple[Tensor, CoordSystem] –

  Source data and coordinates one step ahead.

create_iterator

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

Create an iterator over source frames.

Parameters:

• x (Tensor) –

  Initial condition tensor.

• coords (CoordSystem) –

  Initial condition coordinates.

Yields:

• tuple[Tensor, CoordSystem] –

  Initial condition followed by source frames.