runners

default

class anemoi.inference.runners.default.DefaultRunner(config: RunConfiguration, *, classes: RunnerClasses | None = None)

Bases: Runner

The default runner is a forecaster.

external_graph

anemoi.inference.runners.external_graph.contains_any(key, specifications)
anemoi.inference.runners.external_graph.update_state_dict(model, external_state_dict, keywords='', ignore_mismatched_layers=False, ignore_additional_layers=False)

Update the model’s stated_dict with entries from an external state_dict. Only entries whose keys contain the specified keywords are considered.

anemoi.inference.runners.external_graph.get_updated_supporting_arrays(update_supporting_arrays: dict[Literal['graph', 'file'], dict[str, str]], graph: Any) → dict

Get an update dict for the supporting arrays.

Allows to update the supporting arrays in the checkpoint metadata from the graph data, and from files, which will be loaded using torch.load() or np.load().

parallel

anemoi.inference.runners.parallel.create_parallel_runner(config: Configuration, client_factory: ComputeClientFactory) → None

Creates and runs a parallel runner.

Parameters:

  • config (Configuration) – The configuration object for the runner.

  • client_factory (ComputeClientFactory) – The compute client factory to use for distributed inference.

class anemoi.inference.runners.parallel.NoOp

Bases: object

No operation class used when returning after spawning processes.

execute(*a, **k) → None
class anemoi.inference.runners.parallel.ParallelRunnerMixin(config: Any, compute_client: ComputeClient | None = None, **kwargs)

Bases: Runner

Runner which splits a model over multiple devices. Should be mixed in with a base runner class.

seed(comm_group: ProcessGroup | None) → None

Seed all processes in the cluster to ensure reproducibility.

predict_step(model: Any, input_tensor_torch: Tensor, **kwargs: Any) → Tensor

Performs a prediction step.

Parameters:

  • model (Any) – The model to use for prediction.

  • input_tensor_torch (torch.Tensor) – The input tensor for the model.

  • **kwargs (Any) – Additional arguments for the prediction step.

Returns:

The prediction result.

Return type:

torch.Tensor

complete_forecast_hook() → None

Hook called at the end of the forecast.

create_output(*args, **kwargs) → Output

Creates the real output on rank 0 and a none on the others.

simple

class anemoi.inference.runners.simple.SimpleTensorHandler(context: Runner, metadata: Metadata, constant_forcings_input: Input, dynamic_forcings_input: Input, boundary_forcings_input: Input, trace_path: str | None = None)

Bases: TensorHandler

This tensor handler only supports computed forcings.

create_constant_coupled_forcings(variables: list[str], mask: ndarray[tuple[Any, ...], dtype[Any]]) → list[Forcings]
create_dynamic_coupled_forcings(variables: list[str], mask: ndarray[tuple[Any, ...], dtype[Any]]) → list[Forcings]
class anemoi.inference.runners.simple.SimpleRunner(checkpoint: str, **kwargs)

Bases: Runner

Runner for the low-level API. The user provides a prepared State object and directly calls the Runner.run generator. The input State must contain all fields expected by the model, except for computed forcings which will be loaded by the runner. The user is responsible for providing any constant and dynamic forcings in the input State and during rollout.

execute() → None

Execute the runner.

run(*, input_states: dict[str, dict[str, Any]] | dict[str, Any], **kwargs) → Generator[dict[str, dict[str, Any]] | dict[str, Any], None, None]

Run the model.

Parameters:

  • input_states (dict[str, State]) – The input states for each dataset.

  • lead_time (Union[str, int, datetime.timedelta]) – The lead time.

  • return_numpy (bool, optional) – Whether to return the output state fields as numpy arrays, by default True. Otherwise, it will return torch tensors.

Returns:

The forecasted states.

Return type:

Generator[dict[str, State], None, None]

temporal_downscaler

class anemoi.inference.runners.temporal_downscaler.TemporalDownscalerMetadata(metadata: dict[str, Any], supporting_arrays: dict[str, ndarray[tuple[Any, ...], dtype[Any]]] = {})

Bases: Metadata

property lagged: list[timedelta]

Return the list of steps for the multi_step_input fields. A temporal downscaler looks ahead instead of backwards.

class anemoi.inference.runners.temporal_downscaler.TemporalDownscalerMultiOutRunner(config: RunConfiguration)

Bases: Runner

A runner to be used for inference of a trained temporal downscaler with multiple output steps. Unlike the single output, the temporal downscaling is all done as one step. Can be applied directly on analysis/forecast data without being coupled to a forecasting model.

This runner makes the following assumptions:

  • The model was trained with two input states: (t and t+temporal_downscaling_window)

  • The output states are between these two states and are set by “frequency” in the config

patch_data_request(request: dict, dataset_name: str) → dict

Set sensible defaults when this runner is used with the retrieve command.

execute() → None

Execute the temporal downscaler runner with support for multiple temporal downscaling periods.

temporal_downscaler_stepper(start_date: datetime) → Generator[tuple[timedelta, datetime], None, None]

Generate step and date variables for the temporal downscaling loop.

Parameters:

start_date (datetime.datetime) – Input start date

Returns:

  • step (datetime.timedelta) – Time delta between the target index date and the start date

  • date (datetime.datetime) – Date of the zeroth index of the input tensor

forecast(lead_time: timedelta, input_tensors_numpy: dict[str, ndarray[tuple[Any, ...], dtype[Any]]], input_states: dict[str, dict[str, Any]]) → Generator[dict[str, dict[str, Any]], None, None]

Temporally downscale between the current and future state in the input tensor.

Parameters:

  • lead_time (datetime.timedelta) – Unused. This method processes one temporal downscaling window at a time.

  • input_tensors_numpy (dict[str, FloatArray]) – The input tensors for each dataset, as numpy arrays with shape (multi_step_input, variables, values).

  • input_states (dict[str, State]) – The input states for each dataset. It contains both input dates defined by the config explicit_times.input

Returns:

The temporally downscaled states for each dataset.

Return type:

dict[str, State]

testing

class anemoi.inference.runners.testing.NoModelMixing

Bases: object

property model: Module
class anemoi.inference.runners.testing.NoModelRunner(config: RunConfiguration, *, classes: RunnerClasses | None = None)

Bases: NoModelMixing, DefaultRunner

Runner for running tests.

Inherits from DefaultRunner.