Inference

This module provides the underlying inference functionality for the earthkit-workflows-anemoi package, and is at times a thin wrapper around the anemoi-inference package.

For simple tasks, the fluent API is recommended, which provides a more user-friendly interface to create and manage workflows.

earthkit.workflows.plugins.anemoi.inference.get_initial_conditions_source(config: RunConfiguration | fluent.Action, date: DATE, ensemble_members: ENSEMBLE_MEMBER_SPECIFICATION = None, *, initial_condition_perturbation: bool = False, payload_metadata: dict[str, Any] | None = None) → fluent.Action

Get the initial conditions for the model

Parameters:

  • config (RunConfiguration | fluent.Action) – Configuration object, must contain checkpoint and input. If is a fluent action, the action must return the RunConfiguration object.

  • date (str | tuple[int, int, int]) – Date to get initial conditions for

  • ensemble_members (ENSEMBLE_MEMBER_SPECIFICATION, optional) – Number of ensemble members to get, by default None

  • initial_condition_perturbation (bool, optional) – Whether to get perturbed initial conditions, by default False If False, only one initial condition is returned, and the ensemble members are simulated by wrapping the action.

  • payload_metadata (Optional[dict[str, Any]], optional) – Metadata to add to the payload, by default None

Returns:

Fluent action of the initial conditions

Return type:

fluent.Action

earthkit.workflows.plugins.anemoi.inference.run_model(runner: CascadeRunner, config: RunConfiguration, input_state_source: fluent.Action, lead_time: LEAD_TIME, payload_metadata: dict[str, Any] | None = None, **kwargs) → fluent.Action

Run the model, expanding the results to the correct dimensions.

Parameters:

  • runner (Runner) – anemoi.inference runner

  • config (RunConfiguration) – Configuration object

  • input_state_source (fluent.Action) – Fluent action of initial conditions

  • lead_time (LEAD_TIME) – Lead time to run out to. Can be a string, i.e. 1H, 1D, int, or a datetime.timedelta

  • payload_metadata (Optional[dict[str, Any]], optional) – Metadata to add to the payload, by default None

  • kwargs (dict) – Additional arguments to pass to the runner

Returns:

Cascade action of the model results

Return type:

fluent.Action

earthkit.workflows.plugins.anemoi.inference.run(input_state: dict, runner: CascadeRunner, lead_time: LEAD_TIME) → Generator[Any]

Run the model.

Parameters:

  • input_state (dict) – Initial conditions for the model

  • runner (CascadeRunner) – CascadeRunner object

  • lead_time (LEAD_TIME) – Lead time for the model

Returns:

State of the model at each time step

Return type:

Generator[Any, None, None]

earthkit.workflows.plugins.anemoi.inference.convert_to_fieldlist(state: dict, initial_date: datetime, runner: CascadeRunner, ensemble_member: int | None, **kwargs) → SimpleFieldList

Convert the state to an earthkit FieldList.

Parameters:

  • state – State of the model at a given time step

  • initial_date (datetime.datetime) – Initial date of the model run

  • runner (CascadeRunner) – Runner object

  • ensemble_member (int | None) – Ensemble member number

  • kwargs (dict) – Additional metadata to add to the fields

Returns:

Earthkit FieldList with the model results

Return type:

ekd.FieldList

earthkit.workflows.plugins.anemoi.inference.run_as_earthkit(input_state: dict, runner: CascadeRunner, lead_time: LEAD_TIME, extra_metadata: dict[str, Any] | None = None) → Generator[ekd.SimpleFieldList]

Run the model and yield the results as earthkit FieldList

Parameters:

  • input_state (dict) – Initial Conditions for the model

  • runner (CascadeRunner) – CascadeRunner Object

  • lead_time (LEAD_TIME) – Lead time for the model

  • extra_metadata (dict[str, Any], optional) – Extra metadata to add to the fields, by default None

Returns:

State of the model at each time step

Return type:

Generator[SimpleFieldList, None, None]

earthkit.workflows.plugins.anemoi.inference.run_as_earthkit_from_config(input_state: dict, runner: CascadeRunner, lead_time: LEAD_TIME, extra_metadata: dict[str, Any] | None = None) → Generator[ekd.SimpleFieldList]

Run the model and yield the results as earthkit FieldList

Parameters:

  • input_state (dict) – Initial Conditions for the model

  • runner (CascadeRunner) – CascadeRunner Object

  • lead_time (LEAD_TIME) – Lead time for the model

  • extra_metadata (dict[str, Any], optional) – Extra metadata to add to the fields, by default None

Returns:

State of the model at each time step

Return type:

Generator[SimpleFieldList, None, None]

earthkit.workflows.plugins.anemoi.inference.collect_as_earthkit(input_state: dict, runner: CascadeRunner, lead_time: LEAD_TIME, extra_metadata: dict[str, Any] | None = None) → ekd.SimpleFieldList

Collect the results of the model run as earthkit FieldList

Parameters:

  • input_state (dict) – Initial conditions for the model

  • runner (CascadeRunner) – CascadeRunner object

  • lead_time (LEAD_TIME) – Lead time for the model

  • extra_metadata (dict[str, Any], optional) – Extra metadata to add to the fields, by default None

Returns:

Combined FieldList of the model run

Return type:

ekd.SimpleFieldList

earthkit.workflows.plugins.anemoi.inference.collect_as_earthkit_from_config(input_state: dict, runner: CascadeRunner, lead_time: LEAD_TIME, extra_metadata: dict[str, Any] | None = None) → ekd.SimpleFieldList

Collect the results of the model run as earthkit FieldList

Parameters:

  • input_state (dict) – Initial conditions for the model

  • runner (CascadeRunner) – CascadeRunner object

  • lead_time (LEAD_TIME) – Lead time for the model

  • extra_metadata (dict[str, Any], optional) – Extra metadata to add to the fields, by default None

Returns:

Combined FieldList of the model run

Return type:

ekd.SimpleFieldList