testing

checks

anemoi.inference.testing.checks.check_grib(*, file: Path, expected_variables: list[Variable], grib_keys: dict = {}, check_accum: str | None = None, check_nans=False, reference_date: str = None, **kwargs) → None

anemoi.inference.testing.checks.check_grib_cutout(*, file: Path, metadata: Metadata, reference_grib: str, reference_datetime: str | None = None, mask='lam_0', **kwargs)

check shape and values of inner region against a reference GRIB

anemoi.inference.testing.checks.check_with_xarray(*, file: Path, expected_variables: list[Variable], check_accum: str | None = None, check_nans=False, expected_time: int | None = None, **kwargs) → None

anemoi.inference.testing.checks.check_cutout_with_xarray(*, file: Path, metadata: Metadata, mask='lam_0', reference_date: str = None, reference_dataset={}, reference_file=None, **kwargs) → None

anemoi.inference.testing.checks.check_boundary_forcings_with_xarray(*, file: Path, metadata: Metadata, reference_dataset={}, reference_file=None, **kwargs) → None

anemoi.inference.testing.checks.check_file_exist(*, file: Path, **kwargs) → None

anemoi.inference.testing.checks.check_files_in_directory(*, file: Path, expected_files: int | None = None, **kwargs) → None

mock_checkpoint

anemoi.inference.testing.mock_checkpoint.mock_load_metadata(path: str | None, *, supporting_arrays: bool = True) → tuple[dict[str, Any], dict[str, Any]]

Load metadata from a YAML file.

Parameters:

• path (str, optional) – The path to the checkpoint file.

• supporting_arrays (bool, optional) – Whether to include supporting arrays, by default True.

Returns:

The loaded metadata.

Return type:

dict

anemoi.inference.testing.mock_checkpoint.minimum_mock_checkpoint(metadata: dict[str, Any]) → dict[str, Any]

Create a minimum mock checkpoint from the given metadata.

Parameters:

metadata (dict) – The metadata to create the mock checkpoint from.

Returns:

The minimum mock checkpoint.

Return type:

dict

anemoi.inference.testing.mock_checkpoint.mock_torch_load(path: str, map_location: Any, weights_only: bool) → Any

Load a mock torch model for testing purposes.

Parameters:

• path (str) – The path to the model file.

• map_location (Any) – The device on which to load the model.

• weights_only (bool) – Whether to load only the weights.

Returns:

The mock torch model.

Return type:

Any

class anemoi.inference.testing.mock_checkpoint.MockRunConfiguration

Bases: object

Mock class for loading run configurations from files_for_tests.

classmethod load(path: str, *args, **kwargs) → dict[str, Any]

Load a run configuration from a given path.

Parameters:

• path (str) – The path to the run configuration file.

• *args (tuple) – Additional positional arguments.

• **kwargs (dict) – Additional keyword arguments.

Returns:

The loaded run configuration.

Return type:

dict

mock_model

class anemoi.inference.testing.mock_model.SimpleMetadata(input_shape: tuple, output_shape: tuple, prognostic_variables: list[str], variable_to_input_index: dict[str, int], output_index_to_variable: dict[int, str], variable_to_output_index: dict[str, int], prognostic_input_indices: list[int], prognostic_output_indices: list[int])

Bases: object

Some Metadata attributes are not serialisable, so we store only a subset of metadata in here.

input_shape: tuple

output_shape: tuple

prognostic_variables: list[str]

variable_to_input_index: dict[str, int]

output_index_to_variable: dict[int, str]

variable_to_output_index: dict[str, int]

prognostic_input_indices: list[int]

prognostic_output_indices: list[int]

class anemoi.inference.testing.mock_model.SimpleMockModel(raw_metadata: dict, supporting_arrays: dict)

Bases: Module

Simple mock model that copies prognostics from input to output. Diagnostics are filled with ones. In the case of multi-step input, only the last input step is copied to the output. In the case of multi-step output, the same output is repeated for each output step.

This model has support for legacy models and newer multi-dataset models. When loaded with legacy metadata, the predict_step will take a single tensor and output a single output tensor. When loaded with multi-dataset metadata, the predict_step will take a dict of tensors and output a dict of tensors, with the dataset names, as defined in the metadata, as keys.

metadata: dict[str, SimpleMetadata]

predict_step(input_tensor: Tensor | dict[str, Tensor], date: datetime = None, step: timedelta = None, **kwargs: Any) → Tensor | dict[str, Tensor]

class anemoi.inference.testing.mock_model.LegacyMockModel(raw_metadata: dict[str, Any], supporting_arrays: dict[str, Any])

Bases: Module

Mock model with internal sanity checks. Assumes the input comes from the dummy input source.

constant_in_time: dict[int, Tensor]

predict_step(x: Tensor, date: datetime, step: timedelta, **kwargs: Any) → Tensor

variables