Schemas

This module defines pydantic schemas, which are used to validate the configuration before a training run is started. The top-level config yaml matches the BaseSchema.

The below schemas are organised below identically to the training config files,

Data

class anemoi.training.schemas.data.DatasetDataSchema(*, forcing: list[str] = <factory>, diagnostic: list[str] = <factory>, target: list[str] | None = None, processors: dict[str, ~anemoi.models.schemas.data_processor.PreprocessorSchema])

Bases: BaseModel

A class used to represent the configuration of a single dataset.

forcing: list[str]

Features that are not part of the forecast state but are used as forcing to generate the forecast state.

diagnostic: list[str]

Features that are only part of the forecast state and are not used as an input to the model.

target: list[str] | None

prognostic = diagnostic - forcing.union(target).

Type:

Features used to compute the loss against forecasted variables. Cannot be prognostic or diagnostic, can have the same name as forcing variables but have a different role. Such that

processors: dict[str, PreprocessorSchema]

Layers of model performing computation on latent space. Processors including imputers and normalizers are applied in order of definition. (single dataset mode)

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.training.schemas.data.DataSchema(*, format: str, frequency: str | None = None, datasets: dict[str, DatasetDataSchema] | None = None, num_features: int | None)

Bases: BaseModel

A class used to represent the overall configuration of the dataset(s).

format

The format of the data.

Type:

str

resolution

The resolution of the data.

Type:

str

frequency

The frequency of the data.

Type:

str

timestep

The timestep of the data.

Type:

str

datasets

“Dictionary mapping dataset names to their configurations.”

Type:

dict[str, DatasetDataSchema] | None

num_features

The number of features in the forecast state. To be set in the code.

Type:

int, optional

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

format: str

Format of the data.

frequency: str | None

Time frequency requested from the dataset. Must be null when using trajectory (forecast) datasets.

datasets: dict[str, DatasetDataSchema] | None

Dictionary mapping dataset names to their configurations.

num_features: int | None

Number of features in the forecast state. To be set in the code.

Dataloader

Diagnostics

class anemoi.training.schemas.diagnostics.GraphTrainableFeaturesPlotSchema(*, _target_: Literal['anemoi.training.diagnostics.callbacks.plot.GraphTrainableFeaturesPlot'], dataset_names: list[str], every_n_epochs: int | None)

Bases: BaseModel

target_: Literal['anemoi.training.diagnostics.callbacks.plot.GraphTrainableFeaturesPlot']

GraphTrainableFeaturesPlot object from anemoi training diagnostics callbacks.

dataset_names: list[str]

List of dataset names to plot.

every_n_epochs: int | None

Epoch frequency to plot at.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.training.schemas.diagnostics.FocusAreaSchema(*, name: str | None = None, mask_attr_name: str | None = None, latlon_bbox: Annotated[list[float] | None, MinLen(min_length=4), MaxLen(max_length=4)] = None)

Bases: BaseModel

name: str | None

Name of the focus_area, will be used for plot naming.

mask_attr_name: str | None

Name of the node attribute to use as masking. eg. cutout_mask

latlon_bbox: list[float] | None

Latitude and longitude bounds as [lat_min, lon_min, lat_max, lon_max].

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.training.schemas.diagnostics.PlotLossSchema(*, _target_: Literal['anemoi.training.diagnostics.callbacks.plot.PlotLoss'], dataset_names: list[str], parameter_groups: dict[str, list[str]], every_n_batches: int | None = None)

Bases: BaseModel

target_: Literal['anemoi.training.diagnostics.callbacks.plot.PlotLoss']

PlotLoss object from anemoi training diagnostics callbacks.

dataset_names: list[str]

List of dataset names to plot.

parameter_groups: dict[str, list[str]]

Dictionary with parameter groups with parameter names as key.

every_n_batches: int | None

Batch frequency to plot at.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.training.schemas.diagnostics.MatplotlibColormapSchema(*, _target_: Literal['anemoi.training.utils.custom_colormaps.MatplotlibColormap'], name: str, variables: list[str] | None = None)

Bases: BaseModel

target_: Literal['anemoi.training.utils.custom_colormaps.MatplotlibColormap']

CustomColormap object from anemoi training utils.

name: str

Name of the Matplotlib colormap.

variables: list[str] | None

A list of strings representing the variables for which the colormap is used, by default None.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.training.schemas.diagnostics.MatplotlibColormapClevelsSchema(*, _target_: Literal['anemoi.training.utils.custom_colormaps.MatplotlibColormapClevels'], clevels: list, variables: list[str] | None = None)

Bases: BaseModel

target_: Literal['anemoi.training.utils.custom_colormaps.MatplotlibColormapClevels']

CustomColormap object from anemoi training utils.

clevels: list

The custom color levels for the colormap.

variables: list[str] | None

A list of strings representing the variables for which the colormap is used, by default None.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.training.schemas.diagnostics.DistinctipyColormapSchema(*, _target_: Literal['anemoi.training.utils.custom_colormaps.DistinctipyColormap'], n_colors: int, variables: list[str] | None = None, colorblind_type: str | None = None)

Bases: BaseModel

target_: Literal['anemoi.training.utils.custom_colormaps.DistinctipyColormap']

CustomColormap object from anemoi training utils.

n_colors: int

The number of colors in the colormap.

variables: list[str] | None

A list of strings representing the variables for which the colormap is used, by default None.

colorblind_type: str | None

The type of colorblindness to simulate. If None, the default colorblindness from distinctipy is applied.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.training.schemas.diagnostics.PlotSampleSchema(*, _target_: Literal['anemoi.training.diagnostics.callbacks.plot.PlotSample'], dataset_names: list[str], sample_idx: int, parameters: list[str], accumulation_levels_plot: list[float], precip_and_related_fields: list[str] | None = None, per_sample: int, every_n_batches: int | None = None, colormaps: dict[str, Annotated[MatplotlibColormapSchema | MatplotlibColormapClevelsSchema | DistinctipyColormapSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]] | None = None, focus_area: FocusAreaSchema | None = None, prediction_label: str = 'pred', auxiliary_label: str = 'corrupted targets')

Bases: BaseModel

target_: Literal['anemoi.training.diagnostics.callbacks.plot.PlotSample']

PlotSample object from anemoi training diagnostics callbacks.

dataset_names: list[str]

List of dataset names to plot.

sample_idx: int

Index of sample to plot, must be inside batch size.

parameters: list[str]

List of parameters to plot.

accumulation_levels_plot: list[float]

Accumulation levels to plot.

precip_and_related_fields: list[str] | None

List of precipitation related fields, by default None.

per_sample: int

Number of plots per sample, by default 6.

every_n_batches: int | None

Batch frequency to plot at, by default None.

colormaps: dict[str, Annotated[MatplotlibColormapSchema | MatplotlibColormapClevelsSchema | DistinctipyColormapSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]] | None

List of colormaps to use, by default None.

focus_area: FocusAreaSchema | None

Region of interest to restrict plots to, specified by ‘mask_attr_name’ or ‘latlon_bbox’

prediction_label: str

Label used for the prediction panels.

auxiliary_label: str

Label used for the optional panel that shows the corrupted target seen by the model.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.training.schemas.diagnostics.PlotSpectrumSchema(*, _target_: Literal['anemoi.training.diagnostics.callbacks.plot.PlotSpectrum'], dataset_names: list[str], sample_idx: int, parameters: list[str], every_n_batches: int | None = None, focus_area: FocusAreaSchema | None = None)

Bases: BaseModel

target_: Literal['anemoi.training.diagnostics.callbacks.plot.PlotSpectrum']

PlotSpectrum object from anemoi training diagnostics callbacks.

dataset_names: list[str]

List of dataset names to plot.

sample_idx: int

Index of sample to plot, must be inside batch size.

parameters: list[str]

List of parameters to plot.

every_n_batches: int | None

Batch frequency to plot at, by default None.

focus_area: FocusAreaSchema | None

Region of interest to restrict plots to, specified by ‘mask_attr_name’ or ‘latlon_bbox’

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.training.schemas.diagnostics.PlotHistogramSchema(*, _target_: Literal['anemoi.training.diagnostics.callbacks.plot.PlotHistogram'], dataset_names: list[str], sample_idx: int, parameters: list[str], precip_and_related_fields: list[str] | None = None, every_n_batches: int | None = None, focus_area: FocusAreaSchema | None = None)

Bases: BaseModel

target_: Literal['anemoi.training.diagnostics.callbacks.plot.PlotHistogram']

PlotHistogram object from anemoi training diagnostics callbacks.

dataset_names: list[str]

List of dataset names to plot.

sample_idx: int

Index of sample to plot, must be inside batch size.

parameters: list[str]

List of parameters to plot.

precip_and_related_fields: list[str] | None

List of precipitation related fields, by default None.

every_n_batches: int | None

Batch frequency to plot at, by default None.

focus_area: FocusAreaSchema | None

Region of interest to restrict plots to, specified by ‘mask_attr_name’ or ‘latlon_bbox’

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.training.schemas.diagnostics.PlotEnsSampleSchema(*, _target_: Literal['anemoi.training.diagnostics.callbacks.plot.PlotEnsSample'], dataset_names: list[str], sample_idx: int, parameters: list[str], accumulation_levels_plot: list[float], precip_and_related_fields: list[str] | None = None, per_sample: int, every_n_batches: int | None = None, colormaps: dict[str, Annotated[MatplotlibColormapSchema | MatplotlibColormapClevelsSchema | DistinctipyColormapSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]] | None = None, members: list[int] | int | None = None, focus_area: FocusAreaSchema | None = None)

Bases: BaseModel

target_: Literal['anemoi.training.diagnostics.callbacks.plot.PlotEnsSample']

PlotEnsSample object from anemoi training diagnostics callbacks.

dataset_names: list[str]

List of dataset names to plot.

sample_idx: int

Index of sample to plot, must be inside batch size.

parameters: list[str]

List of parameters to plot.

accumulation_levels_plot: list[float]

Accumulation levels to plot.

precip_and_related_fields: list[str] | None

List of precipitation related fields, by default None.

per_sample: int

Number of plots per sample, by default 6.

every_n_batches: int | None

Batch frequency to plot at, by default None.

colormaps: dict[str, Annotated[MatplotlibColormapSchema | MatplotlibColormapClevelsSchema | DistinctipyColormapSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]] | None

List of colormaps to use, by default None.

members: list[int] | int | None

List of ensemble members to plot. If None, plots all members.

focus_area: FocusAreaSchema | None

Region of interest to restrict plots to, specified by ‘mask_attr_name’ or ‘latlon_bbox’

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.training.schemas.diagnostics.PlottingFrequency(*, batch: Annotated[int, Gt(gt=0)], epoch: Annotated[int, Gt(gt=0)])

Bases: BaseModel

batch: Annotated[int, Gt(gt=0)]

Frequency of the plotting in number of batches.

epoch: Annotated[int, Gt(gt=0)]

Frequency of the plotting in number of epochs.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.training.schemas.diagnostics.PlotSchema(*, asynchronous: bool, datashader: bool, projection_kind: str = 'equirectangular', callbacks: list[Annotated[GraphTrainableFeaturesPlotSchema | PlotLossSchema | PlotSampleSchema | PlotSpectrumSchema | PlotHistogramSchema | PlotEnsSampleSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]], colormaps: dict | None = None, precip_and_related_fields: list[str] | None = None, focus_areas: dict | None = None, datasets_to_plot: list[str] | None = None)

Bases: BaseModel

asynchronous: bool

Handle plotting tasks without blocking the model training.

datashader: bool

Use Datashader to plot.

projection_kind: str

Map projection for diagnostics plots.

Built-in options: 'equirectangular' (no cartopy required) and 'lambert_conformal' (auto-fitted to the data domain; requires cartopy). Any cartopy.crs class name in snake_case is also accepted (e.g. 'robinson', 'mollweide', 'orthographic'); these require cartopy and are instantiated with default constructor arguments (e.g. 'orthographic' centres on longitude/latitude 0). If you need non-default parameters, use 'lambert_conformal' (auto-fitted to the data domain) or subclass MapProjection. Must be 'equirectangular' when datashader is True.

callbacks: list[Annotated[GraphTrainableFeaturesPlotSchema | PlotLossSchema | PlotSampleSchema | PlotSpectrumSchema | PlotHistogramSchema | PlotEnsSampleSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]]

List of plotting functions to call.

colormaps: dict | None

Variable-specific colormaps keyed by ‘default’, ‘error’, or variable name group.

precip_and_related_fields: list[str] | None

Names of precipitation and related fields that use a special colormap.

focus_areas: dict | None

Named spatial focus areas (lat/lon bounding boxes or node attribute masks).

datasets_to_plot: list[str] | None

Dataset names to include in plots.

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.training.schemas.diagnostics.TimeLimitSchema(*, _target_: Literal['anemoi.training.diagnostics.callbacks.stopping.TimeLimit'], limit: int | str, record_file: str | None = None)

Bases: BaseModel

target_: Literal['anemoi.training.diagnostics.callbacks.stopping.TimeLimit']

TimeLimit object from anemoi training diagnostics callbacks.

limit: int | str

Time limit, if int, assumed to be hours, otherwise must be a string with units (e.g. ‘1h’, ‘30m’).

record_file: str | None

File to record the last checkpoint to on exit, if set.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.training.schemas.diagnostics.EarlyStoppingSchema(*, _target_: Literal['anemoi.training.diagnostics.callbacks.stopping.EarlyStopping'], monitor: str, min_delta: float = 0.0, patience: int = 3, verbose: bool = False, mode: Literal['min', 'max'] = 'min', strict: bool = True, check_finite: bool = True, stopping_threshold: float | None = None, divergence_threshold: float | None = None, check_on_train_epoch_end: bool | None = None)

Bases: BaseModel

monitor: str

Metric to monitor

min_delta: float

Minimum change in the monitored quantity to qualify as an improvement.

patience: int

Number of epochs with no improvement after which training will be stopped.

verbose: bool

If True, prints a message for each improvement.

mode: Literal['min', 'max']

One of {‘min’, ‘max’}, changes if minimisation or maximimisation of the metric is ‘good’.

strict: bool

Whether to crash the training if the monitored quantity is not found.

check_finite: bool

Whether to check for NaNs and Infs in the monitored quantity.

stopping_threshold: float | None

Stop training immediately once the monitored quantity reaches this threshold.

divergence_threshold: float | None

Stop training as soon as the monitored quantity becomes worse than this threshold..

check_on_train_epoch_end: bool | None

Whether to check the stopping criteria at the end of each training epoch.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.training.schemas.diagnostics.Debug(*, anomaly_detection: bool)

Bases: BaseModel

anomaly_detection: bool

Activate anomaly detection. This will detect and trace back NaNs/Infs, but slow down training.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.training.schemas.diagnostics.CheckpointSchema(*, save_frequency: int | None, num_models_saved: int)

Bases: BaseModel

save_frequency: int | None

Frequency at which to save the checkpoints.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

num_models_saved: int

Number of model checkpoint to save. Only the last num_models_saved checkpoints will be kept. If set to -1, all checkpoints are kept

class anemoi.training.schemas.diagnostics.WandbSchema(*, _target_: Literal['pytorch_lightning.loggers.wandb.WandbLogger'] = 'pytorch_lightning.loggers.wandb.WandbLogger', enabled: bool, offline: bool, log_model: bool | Literal['all'], project: str, gradients: bool, parameters: bool, entity: str | None = None, interval: Annotated[int, Gt(gt=0)] | None = 100)

Bases: BaseModel

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

enabled: bool

Use Weights & Biases logger.

offline: bool

Run W&B offline.

log_model: bool | Literal['all']

Log checkpoints created by ModelCheckpoint as W&B artifacts. If True, checkpoints are logged at the end of training. If ‘all’, checkpoints are logged during training.

project: str

The name of the project to which this run will belong.

gradients: bool

Whether to log the gradients.

parameters: bool

Whether to log the hyper parameters.

entity: str | None

Username or team name where to send runs. This entity must exist before you can send runs there.

interval: Annotated[int, Gt(gt=0)] | None

Logging frequency in batches.

class anemoi.training.schemas.diagnostics.MlflowSchema(*, _target_: ~typing.Literal['anemoi.training.diagnostics.mlflow.logger.AnemoiMLflowLogger'] = 'anemoi.training.diagnostics.mlflow.logger.AnemoiMLflowLogger', enabled: bool, offline: bool, authentication: bool, log_model: bool | ~typing.Literal['all'] | None = None, tracking_uri: str | None = None, experiment_name: str, project_name: str, system: bool, terminal: bool, run_name: str | None, prefix: str = '', log_hyperparams: bool = True, on_resume_create_child: bool, expand_hyperparams: list[str] = <factory>, http_max_retries: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)], max_params_length: int = 2000, save_dir: str | None = None)

Bases: BaseModel

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

enabled: bool

Use MLflow logger.

offline: bool

Run MLflow offline. Necessary if no internet access available.

authentication: bool

Whether to authenticate with server or not

log_model: bool | Literal['all'] | None

Log checkpoints created by ModelCheckpoint as MLFlow artifacts. If True, checkpoints are logged at the end of training. If ‘all’, checkpoints are logged during training.

tracking_uri: str | None

Address of local or remote tracking server.

experiment_name: str

Name of experiment.

project_name: str

Name of project.

system: bool

Activate system metrics.

terminal: bool

Log terminal logs to MLflow.

run_name: str | None

Name of run.

prefix: str

Prefix for metric keys logged to MLflow.

log_hyperparams: bool

Whether to log hyperparameters.

on_resume_create_child: bool

Whether to create a child run when resuming a run.

expand_hyperparams: list[str]

Keys to expand within params. Any key being expanded will have lists converted according to expand_iterables.

http_max_retries: Annotated[int, Gt(gt=0)]

Specifies the maximum number of retries for MLflow HTTP requests, default 35.

max_params_length: int

Maximum number of hpParams to be logged with mlflow

save_dir: str | None

Directory to save logs to when offline=True, default={system.output.root}/{system.output.logs.mlflow}

class anemoi.training.schemas.diagnostics.AzureMlflowSchema(*, _target_: ~typing.Literal['anemoi.training.diagnostics.mlflow.azureml.AnemoiAzureMLflowLogger'], enabled: bool, offline: ~typing.Literal[False], authentication: bool, log_model: bool | ~typing.Literal['all'] | None = None, tracking_uri: str | None = None, experiment_name: str, project_name: str, system: bool, terminal: ~typing.Literal[False], run_name: str | None, prefix: str = '', log_hyperparams: bool = True, on_resume_create_child: bool, expand_hyperparams: list[str] = <factory>, http_max_retries: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)], max_params_length: int = 2000, save_dir: str | None = None, identity: str | None = None, resource_group: str | None = None, workspace_name: str | None = None, subscription_id: str | None = None, azure_log_level: str = 'WARNING')

Bases: MlflowSchema

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

offline: Literal[False]

Run MLflow offline. Necessary if no internet access available.

terminal: Literal[False]

Log terminal logs to MLflow.

identity: str | None

Type of identity to use for logging in with Azure ML.

resource_group: str | None

Name of the AzureML resource group

workspace_name: str | None

Name of the AzureML workspace

subscription_id: str | None

AzureML subscription ID

azure_log_level: str

Log level for all azure packages (azure-identity, azure-core, etc)

class anemoi.training.schemas.diagnostics.LoggingSchema(*, wandb: WandbSchema | None = None, mlflow: MlflowSchema | AzureMlflowSchema | None = None, interval: Annotated[int, Gt(gt=0)])

Bases: BaseModel

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

wandb: WandbSchema | None

W&B logging schema.

mlflow: Annotated[MlflowSchema | AzureMlflowSchema | None, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]

MLflow logging schema.

interval: Annotated[int, Gt(gt=0)]

Logging frequency in batches.

class anemoi.training.schemas.diagnostics.MemorySchema(*, enabled: bool, steps: Annotated[int, Gt(gt=0)], warmup: Annotated[int, Ge(ge=0)], extra_plots: bool, trace_rank0_only: bool)

Bases: BaseModel

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

enabled: bool

Enable memory report. Default to false.

steps: Annotated[int, Gt(gt=0)]

Frequency of memory profiling. Default to 5.

warmup: Annotated[int, Ge(ge=0)]

Number of step to discard before the profiler starts to record traces. Default to 2.

extra_plots: bool

Save plots produced with torch.cuda._memory_viz.profile_plot if available. Default to false.

trace_rank0_only: bool

Trace only rank 0 from SLURM_PROC_ID. Default to false.

class anemoi.training.schemas.diagnostics.Snapshot(*, enabled: bool, steps: Annotated[int, Gt(gt=0)], warmup: Annotated[int, Ge(ge=0)])

Bases: BaseModel

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

enabled: bool

Enable memory snapshot recording. Default to false.

steps: Annotated[int, Gt(gt=0)]

Frequency of snapshot. Default to 4.

warmup: Annotated[int, Ge(ge=0)]

Number of step to discard before the profiler starts to record traces. Default to 0.

class anemoi.training.schemas.diagnostics.Profiling(*, enabled: bool, verbose: bool | None = None)

Bases: BaseModel

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

enabled: bool

Enable component profiler. Default to false.

verbose: bool | None

Set to true to include the full list of profiled action or false to keep it concise.

class anemoi.training.schemas.diagnostics.BenchmarkProfilerSchema(*, memory: MemorySchema = <factory>, time: Profiling = <factory>, speed: Profiling = <factory>, system: Profiling = <factory>, model_summary: Profiling = <factory>, snapshot: Snapshot = <factory>)

Bases: BaseModel

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

memory: MemorySchema

Schema for memory report containing metrics associated with CPU and GPU memory allocation.

time: Profiling

Report with metrics of execution time for certain steps across the code.

speed: Profiling

Report with metrics of execution speed at training and validation time.

system: Profiling

Report with metrics of GPU/CPU usage, memory and disk usage and total execution time.

model_summary: Profiling

Table summary of layers and parameters of the model.

snapshot: Snapshot

Memory snapshot if torch.cuda._record_memory_history is available.

class anemoi.training.schemas.diagnostics.ProgressBarSchema(*, _target_: Literal['pytorch_lightning.callbacks.TQDMProgressBar', 'pytorch_lightning.callbacks.RichProgressBar', 'anemoi.training.diagnostics.profilers.ProfilerProgressBar'], refresh_rate: Annotated[int, Gt(gt=0)] = 1)

Bases: BaseModel

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

target_: Literal['pytorch_lightning.callbacks.TQDMProgressBar', 'pytorch_lightning.callbacks.RichProgressBar', 'anemoi.training.diagnostics.profilers.ProfilerProgressBar']

TQDMProgressBar object from pytorch lightning.

refresh_rate: Annotated[int, Gt(gt=0)]

Refresh rate of the progress bar.

class anemoi.training.schemas.diagnostics.DiagnosticsSchema(*, plot: PlotSchema | None = None, callbacks: list = <factory>, benchmark_profiler: BenchmarkProfilerSchema, debug: Debug, log: LoggingSchema, enable_progress_bar: bool, progress_bar: ProgressBarSchema | None = None, print_memory_summary: bool, enable_checkpointing: bool, checkpoint: dict[str, ~anemoi.training.schemas.diagnostics.CheckpointSchema]=<factory>, check_val_every_n_epoch: Annotated[int, ~annotated_types.Gt(gt=0)] = 1)

Bases: BaseModel

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

plot: PlotSchema | None

Plot schema.

callbacks: list

Callbacks schema.

benchmark_profiler: BenchmarkProfilerSchema

Benchmark profiler schema for profile command.

debug: Debug

Debug schema.

log: LoggingSchema

Log schema.

enable_progress_bar: bool

Activate progress bar.

progress_bar: ProgressBarSchema | None

Progress bar schema.

print_memory_summary: bool

Print the memory summary.

enable_checkpointing: bool

Allow model to save checkpoints.

checkpoint: dict[str, CheckpointSchema]

Checkpoint schema for defined frequency (every_n_minutes, every_n_epochs, …).

check_val_every_n_epoch: Annotated[int, Gt(gt=0)]

Run validation every n epochs.

System

class anemoi.training.schemas.system.HardwareSchema(*, accelerator: Annotated[str, ~pydantic.functional_validators.AfterValidator(func=functools.partial(<function allowed_values at 0x7cbeaedfe700>, values=['cpu', 'gpu', 'auto', 'cuda', 'tpu']))]='cuda', num_gpus_per_node: Annotated[int, ~annotated_types.Ge(ge=0)] = 1, num_nodes: Annotated[int, ~annotated_types.Ge(ge=0)] = 1, num_gpus_per_model: Annotated[int, ~annotated_types.Ge(ge=0)] = 1, num_gpus_per_ensemble: Annotated[int, ~annotated_types.Ge(ge=0)] = 1)

Bases: BaseModel

accelerator: partial(<function allowed_values at 0x7cbeaedfe700>, values=['cpu', 'gpu', 'auto', 'cuda', 'tpu']))]

Accelerator to use for training.

num_gpus_per_node: Annotated[int, Ge(ge=0)]

Number of GPUs per node.

num_nodes: Annotated[int, Ge(ge=0)]

Number of nodes.

num_gpus_per_model: Annotated[int, Ge(ge=0)]

Number of GPUs per model.

num_gpus_per_ensemble: Annotated[int, Ge(ge=0)]

Number of GPUs per ensemble.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.training.schemas.system.InputSchema(*, dataset: Path | dict[str, Path] | None = None, graph: Path | None = None, truncation: Path | None = None, truncation_inv: Path | None = None, loss_matrices_path: Path | None = None, warm_start: Path | None = None)

Bases: BaseModel

dataset: Path | dict[str, Path] | None

Path to the dataset file.

graph: Path | None

Path to the graph file.

truncation: Path | None

Path to the truncation matrix file.

truncation_inv: Path | None

Path to the inverse truncation matrix file.

loss_matrices_path: Path | None

Path to the directory containing smoothing matrices for the multiscale loss

warm_start: Path | None

Path of the checkpoint file to use for warm starting the training

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.training.schemas.system.CheckpointsSchema(*, root: Path, every_n_epochs: str = 'anemoi-by_epoch-epoch_{epoch:03d}-step_{step:06d}', every_n_train_steps: str = 'anemoi-by_step-epoch_{epoch:03d}-step_{step:06d}', every_n_minutes: str = 'anemoi-by_time-epoch_{epoch:03d}-step_{step:06d}')

Bases: BaseModel

root: Path

Root directory for saving checkpoint files.

every_n_epochs: str

File name pattern for checkpoint files saved by epoch frequency.

every_n_train_steps: str

File name pattern for checkpoint files saved by step frequency.

every_n_minutes: str

File name pattern for checkpoint files saved by time frequency (minutes).

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.training.schemas.system.Logs(*, root: Path, wandb: Path | None = None, mlflow: Path | None = None, tensorboard: Path | None = None)

Bases: BaseModel

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

wandb: Path | None

Path to output wandb logs.

mlflow: Path | None

Path to output mlflow logs.

tensorboard: Path | None

Path to output tensorboard logs.

class anemoi.training.schemas.system.OutputSchema(*, root: Path | None = None, logs: Logs | None = None, checkpoints: CheckpointsSchema = <factory>, plots: Path | None = None, profiler: Path | None)

Bases: BaseModel

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

root: Path | None

Path to the output directory.

logs: Logs | None

Logging directories.

checkpoints: CheckpointsSchema

Paths to the checkpoints.

plots: Path | None

Path to the plots directory.

profiler: Path | None

Path to the profiler directory.

class anemoi.training.schemas.system.SystemSchema(*, hardware: HardwareSchema, input: InputSchema, output: OutputSchema)

Bases: BaseModel

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

hardware: HardwareSchema

Specification of hardware and compute resources available including the number of nodes, GPUs, and accelerator.

input: InputSchema

Definitions of specific input and output artifacts used relative to the directories defined in output.

output: OutputSchema

High-level directory structure describing where data is read from.

Graph

class anemoi.graphs.schemas.base_graph.NodeSchema(*, node_builder: AnemoiDatasetNodeSchema | NPZnodeSchema | TextNodeSchema | ICONMeshNodeSchema | LimitedAreaNPZFileNodesSchema | ReducedGaussianGridNodeSchema | IcosahedralandHealPixNodeSchema | LimitedAreaIcosahedralandHealPixNodeSchema | StretchedIcosahdralNodeSchema, attributes: dict[str, PlanarAreaWeightSchema | MaskedPlanarAreaWeightsSchema | SphericalAreaWeightSchema | CutOutMaskSchema | GridsMaskSchema | NonmissingAnemoiDatasetVariableSchema | BooleanOperationSchema] | None = None)

Bases: BaseModel

node_builder: Annotated[AnemoiDatasetNodeSchema | NPZnodeSchema | TextNodeSchema | ICONMeshNodeSchema | LimitedAreaNPZFileNodesSchema | ReducedGaussianGridNodeSchema | IcosahedralandHealPixNodeSchema | LimitedAreaIcosahedralandHealPixNodeSchema | StretchedIcosahdralNodeSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]

Node builder schema.

attributes: dict[str, PlanarAreaWeightSchema | MaskedPlanarAreaWeightsSchema | SphericalAreaWeightSchema | CutOutMaskSchema | GridsMaskSchema | NonmissingAnemoiDatasetVariableSchema | BooleanOperationSchema] | None

Dictionary of attributes with names as keys and anemoi.graphs.nodes.attributes objects as values.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.graphs.schemas.base_graph.EdgeSchema(*, source_name: str, target_name: str, edge_builders: list[Annotated[KNNEdgeSchema | CutoffEdgeSchema | MultiScaleEdgeSchema | HEALPixMultiScaleEdgesSchema | ICONTopologicalEdgeSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]], attributes: dict[str, BaseEdgeAttributeSchema | EdgeAttributeFromNodeSchema | DirectionalHarmonicsSchema | RadialBasisFeaturesSchema])

Bases: BaseModel

source_name: str

Source of the edges.

target_name: str

Target of the edges.

edge_builders: list[Annotated[KNNEdgeSchema | CutoffEdgeSchema | MultiScaleEdgeSchema | HEALPixMultiScaleEdgesSchema | ICONTopologicalEdgeSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]]

Edge builder schema.

attributes: dict[str, BaseEdgeAttributeSchema | EdgeAttributeFromNodeSchema | DirectionalHarmonicsSchema | RadialBasisFeaturesSchema]

Dictionary of attributes with names as keys and anemoi.graphs.edges.attributes objects as values.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.graphs.schemas.base_graph.BaseGraphSchema(*, nodes: dict[str, ~anemoi.graphs.schemas.base_graph.NodeSchema] | None=None, edges: list[EdgeSchema] | None = None, overwrite: bool, post_processors: list[~types.Annotated[~anemoi.graphs.schemas.post_processors.RemoveUnconnectedNodesSchema | ~anemoi.graphs.schemas.post_processors.SubsetNodesInAreaSchema | ~anemoi.graphs.schemas.post_processors.RestrictEdgeLengthSchema | ~anemoi.graphs.schemas.post_processors.RemoveSelfEdgesSchema | ~anemoi.graphs.schemas.post_processors.SortEdgeIndexSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]] = <factory>)

Bases: BaseModel

model_config = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

nodes: dict[str, NodeSchema] | None

Nodes schema for all types of nodes (ex. data, hidden).

edges: list[EdgeSchema] | None

List of edges schema.

overwrite: bool

whether to overwrite existing graph file. Default to True.

Model

Training