Interface

class anemoi.models.interface.AnemoiModelInterface(*, config: DictConfig, n_step_input: int, n_step_output: int, graph_data: HeteroData, statistics: dict, data_indices: dict, metadata: dict, statistics_tendencies: dict | None = None, supporting_arrays: dict | None = None)

Bases: Module

An interface for Anemoi models.

This class is a wrapper around the Anemoi model that includes pre-processing and post-processing steps. It inherits from the PyTorch Module class.

config

Configuration settings for the model.

Type:

DictConfig

id

A unique identifier for the model instance.

Type:

str

n_step_input

Number of input timesteps provided to the model.

Type:

int

graph_data

Graph data for the model.

Type:

HeteroData

statistics

Statistics for the data.

Type:

dict

metadata

Metadata for the model.

Type:

dict

statistics_tendencies

Statistics for the tendencies of the data.

Type:

dict

supporting_arrays

Numpy arraysto store in the checkpoint.

Type:

dict

data_indices

Indices for the data.

Type:

dict

pre_processors

Pre-processing steps to apply to the data before passing it to the model.

Type:

Processors

post_processors

Post-processing steps to apply to the model’s output.

Type:

Processors

model

The underlying Anemoi model.

Type:

AnemoiModelEncProcDec

predict_step(batch: dict[str, Tensor], model_comm_group: ProcessGroup | None = None, gather_out: bool = True, **kwargs) → dict[str, Tensor]

Prediction step for the model.

Parameters:

• batch (dict[str, torch.Tensor]) – Input batched data.

• model_comm_group (Optional[ProcessGroup], optional) – Model communication group, specifies which GPUs work together.

• gather_out (bool, optional) – Whether to gather the output, by default True.

• **kwargs – Additional prediction keyword arguments.

Returns:

Predicted data.

Return type:

dict[str, torch.Tensor]