Fluent API

The fluent API of the earthkit-workflows-anemoi package provides a convenient way to create and manage workflows for anemoi tasks. It defines high level functions to begin inference from a variety of initial conditions, and ways to make anemoi-datasets in a parallel manner.

Fluent API for anemoi inference.

earthkit.workflows.plugins.anemoi.fluent.from_config(config: os.PathLike | dict[str, Any] | RunConfiguration, overrides: dict[str, Any] | None = None, *, date: DATE | None = None, ensemble_members: ENSEMBLE_MEMBER_SPECIFICATION = None, environment: ENVIRONMENT = None, **kwargs: Any) → fluent.Action

Run an anemoi inference model from a configuration file

Parameters:

  • config (os.PathLike | dict[str, Any]) – Path to the configuration file, or dictionary of configuration

  • overrides (Optional[dict[str, Any]], optional) – Override for the config, by default None

  • date (Optional[DATE], optional) – Specific override for date, by default None

  • ensemble_members (ENSEMBLE_MEMBER_SPECIFICATION , optional) – Number of ensemble members to run, None will run a single instance, by default None

  • environment (ENVIRONMENT, optional) – Environment to run the model in, by default None If None, will use the current environment Should be set to strings, as if used in pip install, e.g. [“anemoi-models==0.3.1”] Can be dict[str, list[str]] with keys inference and initial_conditions to set the environment for each part of the run.

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

Returns:

earthkit.workflows action of the model results

Return type:

fluent.Action

Examples

>>> from earthkit.workflows.plugins.anemoi.fluent import from_config
>>> from_config("config.yaml", date = "2021-01-01T00:00:00")
earthkit.workflows.plugins.anemoi.fluent.from_input(ckpt: VALID_CKPT, input: str | dict[str, Any], date: DATE, lead_time: LEAD_TIME, *, ensemble_members: ENSEMBLE_MEMBER_SPECIFICATION = None, environment: ENVIRONMENT = None, **kwargs: Any) → fluent.Action

Run an anemoi inference model from a given input source

Parameters:

  • ckpt (VALID_CKPT) – Checkpoint to load

  • input (str | dict[str, Any]) – anemoi.inference input source. Can be mars, grib, etc or a dictionary of input configuration,

  • date (DATE) – Date to get initial conditions for

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

  • ensemble_members (ENSEMBLE_MEMBER_SPECIFICATION, optional) – Number of ensemble members to run, None will run a single instance, by default None

  • environment (ENVIRONMENT, optional) – Environment to run the model in, by default None If None, will use the current environment Should be set to strings, as if used in pip install, e.g. [“anemoi-models==0.3.1”] Can be dict[str, list[str]] with keys inference and initial_conditions to set the environment for each part of the run.

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

Returns:

earthkit.workflows action of the model results

Return type:

fluent.Action

Examples

>>> from earthkit.workflows.plugins.anemoi.fluent import from_input
>>> from_input("anemoi_model.ckpt", "mars", date = "2021-01-01T00:00:00", lead_time = "10D")
earthkit.workflows.plugins.anemoi.fluent.from_initial_conditions(ckpt: VALID_CKPT, initial_conditions: State | fluent.Action | fluent.Payload | Callable, lead_time: LEAD_TIME, configuration_kwargs: dict[str, Any] | None = None, *, ensemble_members: ENSEMBLE_MEMBER_SPECIFICATION = None, environment: list[str] | None = None, **kwargs: Any) → fluent.Action

Run an anemoi inference model from initial conditions

Parameters:

  • ckpt (VALID_CKPT) – Checkpoint to load

  • initial_conditions (State | fluent.Action | fluent.Payload | Callable) – Initial conditions for the model Can be other fluent actions, payloads, or a callable, or a State. If a fluent action and multiple ensemble member initial conditions are included, the dimension must be named ensemble_member.

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

  • configuration_kwargs (dict[str, Any]:) – kwargs for anemoi.inference configuration

  • ensemble_members (Optional[ENSEMBLE_MEMBER_SPECIFICATION], optional) – Number of ensemble members to run, If initial_conditions is a fluent action, with multiple ensemble members, this argument can be set to None, and the number of ensemble members will be inferred from the action. by default None.

  • environment (Optional[list[str]], optional) – Environment to run the model in, by default None If None, will use the current environment Should be set to strings, as if used in pip install, e.g. [“anemoi-models==0.3.1”]

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

Returns:

earthkit.workflows action of the model results

Return type:

fluent.Action

Examples

>>> from earthkit.workflows.plugins.anemoi.fluent import from_initial_conditions
>>> from_initial_conditions("anemoi_model.ckpt", init_conditions, lead_time = "10D")
earthkit.workflows.plugins.anemoi.fluent.create_dataset(config: dict[str, Any] | PathLike, path: PathLike, *, number_of_tasks: int | None = None, overwrite: bool = False, test: bool = False, environment: list[str] | None = None) → Action

Create an anemoi dataset from a configuration.

Parameters:

  • config (dict[str, Any] | os.PathLike) – Configuration to use

  • path (os.PathLike) – Path to save the dataset to

  • number_of_tasks (Optional[int], optional) – Number of tasks to run in parallel, by default None If None, will use a heurisitic based on date groups

  • overwrite (bool, optional) – Whether to overwrite the dataset if it exists, by default False

  • test (bool, optional) – Build a small dataset, using only the first dates. And, when possible, using low resolution and less ensemble members, by default False

  • environment (Optional[list[str]], optional) – Environment to run the model in, by default None If None, will use the current environment Should be set to strings, as if used in pip install, e.g. [“anemoi-datasets==0.3.1”]

Returns:

earthkit.workflows action to create the dataset

Return type:

fluent.Action

Examples

>>> from earthkit.workflows.plugins.anemoi.fluent import create_dataset
>>> create_dataset("dataset_recipe.yaml", "output_dir/dataset.zarr")
earthkit.workflows.plugins.anemoi.fluent.from_dataset(ckpt: VALID_CKPT, dataset_config: dict[str, Any] | os.PathLike, date: DATE, lead_time: LEAD_TIME, *, ensemble_members: ENSEMBLE_MEMBER_SPECIFICATION = None, input_template: dict[str, Any] | None = None, number_of_dataset_tasks: int | None = None, environment: ENVIRONMENT = None, **kwargs: Any) → fluent.Action

Run an anemoi inference model after creating a dataset from a recipe.

NOTE: This will create a dataset in the current working directory and will not delete it after the run. #TODO Fix.

Parameters:

  • ckpt (VALID_CKPT) – Checkpoint to load

  • dataset_config (dict[str, Any] | os.PathLike) – Dataset Configuration to use

  • date (DATE) – Date to get initial conditions for

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

  • ensemble_members (ENSEMBLE_MEMBER_SPECIFICATION, optional) – Number of ensemble members to run, None will run a single instance, by default None

  • input_template (Optional[dict[str, Any]], optional) – Template of input to use for inference, by default None Use “%DATASET_PATH%” to mark where the dataset path should be inserted Default is {“dataset”: “%DATASET_PATH%”}

  • number_of_dataset_tasks (Optional[int], optional) – Number of tasks to run in parallel, by default None If None, will use a heurisitic based on date groups

  • environment (ENVIRONMENT, optional) – Environment to run the model in, by default None If None, will use the current environment Should be set to strings, as if used in pip install, e.g. [“anemoi-models==0.3.1”] Can be dict[str, list[str]] with keys inference, initial_conditions, and dataset to set the environment for each part of the run.

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

Returns:

earthkit.workflows actions of the dataset, and then model.

Return type:

fluent.Action

Examples

>>> from earthkit.workflows.plugins.anemoi.fluent import from_dataset
>>> from_dataset("anemoi_model.ckpt", "dataset_recipe.yaml", date = "2021-01-01T00:00:00", lead_time = "10D")