DataReplay#

class earth2studio.models.px.DataReplay(
source,
variable,
domain_coords,
step=numpy.timedelta64(6, 'h'),
)[source]#

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 (np.timedelta64, optional) – Time between frames, by default np.timedelta64(6, “h”)

__call__(x, coords)[source]#

Advance the source by one step.

Parameters:

  • x (torch.Tensor) – Input tensor.

  • coords (CoordSystem) – Input coordinate system.

Returns:

Source data and coordinates one step ahead.

Return type:

tuple[torch.Tensor, CoordSystem]

create_iterator(x, coords)[source]#

Create an iterator over source frames.

Parameters:

  • x (torch.Tensor) – Initial condition tensor.

  • coords (CoordSystem) – Initial condition coordinates.

Yields:

tuple[torch.Tensor, CoordSystem] – Initial condition followed by source frames.

Return type:

Iterator[tuple[Tensor, OrderedDict[str, ndarray]]]