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.

anemoi.training.schemas.base_schema.BaseSchema.task: Task configuration.

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

Data

Dataloader

class anemoi.training.schemas.dataloader.Frequency(root: RootModelRootType = PydanticUndefined)

Bases: RootModel

model_config = {}: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

Bases: BaseModel

Dictionary-style dataset config passed directly to open_dataset.

model_config = {'arbitrary_types_allowed': True, 'extra': 'allow'}: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

dataset: str | Path | dict | list[dict]: Dataset source identifier.

frequency: Frequency | None: Optional frequency requested from open_dataset.

drop: list[str] | None: Optional list of variables to drop from the dataset.

select: list[str] | None: Optional list of variables to select from the dataset.

statistics: str | Path | None: Optional path to custom statistics file.

step_start: Annotated[int, Ge(ge=0)] | None: First forecast step to include, in hours (inclusive). Selects ForecastStepDataset when set.

step_end: Annotated[int, Gt(gt=0)] | None: Last forecast step to include, in hours (inclusive). Selects ForecastStepDataset when set.

step_frequency: Annotated[int, Gt(gt=0)] | None: Stride between selected forecast steps, in hours. Selects ForecastStepDataset when set.

Bases: BaseModel

Dataset configuration schema.

dataset_config: str | DatasetConfigSchema | Path | list[dict] | None: Dataset definition passed to open_dataset.

start: str | int | None: Starting datetime for sample of the dataset.

end: str | int | None: Ending datetime [inclusive] for sample of the dataset.

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.dataloader.TrajectorySamplingSchema(*, stride: int | None = None)

Bases: BaseModel

Trajectory anchor sampling configuration.

stride: int | None: Stride between anchor positions. None = window size (non-overlapping); 1 = every position.

model_config = {}: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.training.schemas.dataloader.TrajectorySchema(*, sampling: TrajectorySamplingSchema | None = None)

Bases: BaseModel

Trajectory (forecast) reader configuration for 5-D trajectory zarr datasets.

sampling: TrajectorySamplingSchema | None: Anchor sampling config. stride=None → non-overlapping; stride=1 → all; stride=N → step N.

model_config = {}: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

Bases: NativeDatasetSchema

Dataset configuration schema.

trajectory: TrajectorySchema | None: Trajectory configuration.

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.dataloader.LoaderSet(*, training: Annotated[int, Ge(ge=0)] | None, validation: Annotated[int, Ge(ge=0)] | None, test: Annotated[int, Ge(ge=0)] | None)

Bases: BaseModel

training: Annotated[int, Ge(ge=0)] | None: Value for training dataset

validation: Annotated[int, Ge(ge=0)] | None: Value for validation dataset

test: Annotated[int, Ge(ge=0)] | None: Value for test dataset

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.dataloader.DataLoadingPolicySchema(*, _target_: str = 'anemoi.training.data.loading.SynchronizedLoadingPolicy', **extra_data: Any)

Bases: GenericSchema

Hydra instantiation config for the DataLoadingPolicy that builds the dataset and batch sampler for each stage.

model_config = {'extra': 'allow', 'populate_by_name': True, '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_: str: Hydra target class of the policy.

class anemoi.training.schemas.dataloader.DataLoaderSchema(*, prefetch_factor: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], pin_memory: bool, persistent_workers: bool = False, fake_dataloading: bool = False, loading_policy: ~anemoi.training.schemas.dataloader.DataLoadingPolicySchema = <factory>, num_workers: ~anemoi.training.schemas.dataloader.LoaderSet, batch_size: ~anemoi.training.schemas.dataloader.LoaderSet, limit_batches: ~anemoi.training.schemas.dataloader.LoaderSet, training: dict[~typing.Literal['datasets'], dict[str, ~anemoi.training.schemas.dataloader.NativeDatasetSchema | ~anemoi.training.schemas.dataloader.TrajectoryDatasetSchema]] | ~anemoi.training.schemas.dataloader.NativeDatasetSchema | ~anemoi.training.schemas.dataloader.TrajectoryDatasetSchema, validation: dict[~typing.Literal['datasets'], dict[str, ~anemoi.training.schemas.dataloader.NativeDatasetSchema | ~anemoi.training.schemas.dataloader.TrajectoryDatasetSchema]] | ~anemoi.training.schemas.dataloader.NativeDatasetSchema | ~anemoi.training.schemas.dataloader.TrajectoryDatasetSchema, test: dict[~typing.Literal['datasets'], dict[str, ~anemoi.training.schemas.dataloader.NativeDatasetSchema | ~anemoi.training.schemas.dataloader.TrajectoryDatasetSchema]] | ~anemoi.training.schemas.dataloader.NativeDatasetSchema | ~anemoi.training.schemas.dataloader.TrajectoryDatasetSchema, read_group_size: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)], trajectory_sampling: ~anemoi.training.schemas.dataloader.TrajectorySamplingSchema | None = None, multiprocessing_context: str | None = None, **extra_data: ~typing.Any)

Bases: BaseModel

model_config = {'arbitrary_types_allowed': True, 'extra': 'allow'}: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

prefetch_factor: int: Number of batches loaded in advance by each worker.

pin_memory: bool: If True, the data loader will copy Tensors into device/CUDA pinned memory before returning them.

persistent_workers: bool: Keep dataloader workers alive between epochs instead of starting new ones every epoch.

fake_dataloading: bool: Load one real sample per worker and reuse it for subsequent accesses.

loading_policy: DataLoadingPolicySchema: Hydra config for the DataLoadingPolicy that builds the dataset and batch sampler for each stage.

num_workers: LoaderSet: Number of process per-GPU for batch distribution.

batch_size: LoaderSet: Per-GPU batch size.

limit_batches: LoaderSet: Limit number of batches to run. Default value null, will run on all the batches.

training: dict[Literal['datasets'], dict[str, NativeDatasetSchema | TrajectoryDatasetSchema]] | NativeDatasetSchema | TrajectoryDatasetSchema: Training DatasetSchema.

validation: dict[Literal['datasets'], dict[str, NativeDatasetSchema | TrajectoryDatasetSchema]] | NativeDatasetSchema | TrajectoryDatasetSchema: Validation DatasetSchema.

test: dict[Literal['datasets'], dict[str, NativeDatasetSchema | TrajectoryDatasetSchema]] | NativeDatasetSchema | TrajectoryDatasetSchema: Test DatasetSchema.

read_group_size: Annotated[int, Gt(gt=0)]: Number of GPUs per reader group. Defaults to number of GPUs (see BaseSchema validators).

trajectory_sampling: TrajectorySamplingSchema | None: Default trajectory anchor sampling used across all splits. stride=None → non-overlapping; stride=N → step N.

multiprocessing_context: str | None: Multiprocessing context to use for workers. If None, the default context will be used

Diagnostics

System

class anemoi.training.schemas.system.HardwareSchema(*, accelerator: Annotated[str, ~pydantic.functional_validators.AfterValidator(func=functools.partial(<function allowed_values at 0x7a66c13020c0>, 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 0x7a66c13020c0>, 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].

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

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

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 | MutualKNNEdgeSchema | 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 | MutualKNNEdgeSchema | 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

class anemoi.models.schemas.models.DefinedModels(value): Bases: str, Enum

class anemoi.models.schemas.models.Model(*, _target_: DefinedModels, hidden_nodes_name: str | list[str], latent_skip: bool = True, _convert_: str = 'none')

Bases: BaseModel

target_: DefinedModels: Model object defined in anemoi.models.model.

hidden_nodes_name: str | list[str]: Name of the hidden nodes. If the model is hierarchical, it can be a list of names for each level.

latent_skip: bool: Add skip connection in latent space before/after processor.

convert_: str: Keep OmegaConf containers when instantiating — model code uses attribute-style access throughout.

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.models.schemas.models.SparseProjectorSchema(*, num_chunks: Annotated[int, Gt(gt=0)] = 1)

Bases: BaseModel

num_chunks: PositiveInt: Number of chunks to use for sparse projection matmuls.

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.models.schemas.models.TransportSourceConfig(*, kind: Literal['default', 'zero', 'gaussian', 'reference_state'] = 'default', scale: Annotated[float, Ge(ge=0)] = 1.0, noise_scale: Annotated[float, Ge(ge=0)] = 0.0)

Bases: BaseModel

Configuration of the starting/source field for transport objectives.

The defaults map 1:1 to TransportSourceSettings in anemoi.models.transport.settings and are restated here so that importing the schemas does not pull in the model code.

kind: Literal['default', 'zero', 'gaussian', 'reference_state']: Starting field used before the transport objective moves toward the target.

scale: NonNegativeFloat: Multiplier applied to the starting/source field.

noise_scale: NonNegativeFloat: Additional additive Gaussian noise applied to the starting/source field.

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.models.schemas.models.TransportConfig(*, objective: Literal['edm_diffusion', 'stochastic_interpolant']='edm_diffusion', sigma_data: Annotated[float, ~annotated_types.Gt(gt=0)] = 1.0, noise_channels: Annotated[int, ~annotated_types.Gt(gt=0)] = 32, noise_cond_dim: Annotated[int, ~annotated_types.Gt(gt=0)] = 16, sigma_max: Annotated[float, ~annotated_types.Gt(gt=0)] = 100.0, sigma_min: Annotated[float, ~annotated_types.Gt(gt=0)] = 0.02, rho: Annotated[float, ~annotated_types.Gt(gt=0)] = 7.0, si_alpha_schedule: Literal['linear'] = 'linear', si_beta_schedule: Literal['linear', 'quadratic']='linear', si_sigma_schedule: Literal['brownian_bridge', 'quadratic_bridge']='brownian_bridge', source: TransportSourceConfig = <factory>, si_noise_scale: Annotated[float, ~annotated_types.Ge(ge=0)] = 1.0, training_condition: dict = <factory>, noise_embedder: dict = <factory>, inference_defaults: dict = <factory>)

Bases: BaseModel

Configuration of the transport objective, path, conditioning and inference.

The defaults map 1:1 to EdmSettings, NoiseConditioningSettings and StochasticInterpolantSettings in anemoi.models.transport.settings and are restated here so that importing the schemas does not pull in the model code.

objective: Literal['edm_diffusion', 'stochastic_interpolant']: Training and sampling objective used by the transport model.

sigma_data: PositiveFloat: Typical data scale used by EDM diffusion.

noise_channels: PositiveInt: Number of channels in the noise or bridge-time embedding.

noise_cond_dim: PositiveInt: Size of the conditioning vector passed to conditional layers.

sigma_max: PositiveFloat: Maximum EDM diffusion noise level used during training.

sigma_min: PositiveFloat: Minimum EDM diffusion noise level used during training.

rho: PositiveFloat: Shape parameter for the Karras EDM noise schedule.

si_alpha_schedule: Literal['linear']: Schedule for how strongly the SI bridge keeps the source field.

si_beta_schedule: Literal['linear', 'quadratic']: Schedule for how strongly the SI bridge moves toward the target field.

si_sigma_schedule: Literal['brownian_bridge', 'quadratic_bridge']: Schedule for the SI bridge-noise amplitude.

source: TransportSourceConfig: Configuration for the starting/source field.

si_noise_scale: NonNegativeFloat: Overall scale of the stochastic-interpolant bridge noise.

training_condition: dict: Distribution used to sample one training noise level or bridge time per sample.

noise_embedder: dict: Hydra configuration for embedding the current noise level or bridge time.

inference_defaults: dict: Default sampler parameters used during inference.

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.models.schemas.models.TransportModel(*, _target_: ~anemoi.models.schemas.models.DefinedModels, hidden_nodes_name: str | list[str], latent_skip: bool = True, _convert_: str = 'none', transport: ~anemoi.models.schemas.models.TransportConfig = <factory>)

Bases: Model

transport: TransportConfig: Transport model objective, path, conditioning, and inference configuration.

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.models.schemas.models.NoOutputMaskSchema(*, _target_: Literal['anemoi.training.utils.masks.NoOutputMask'])

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].

class anemoi.models.schemas.models.Boolean1DSchema(*, _target_: Literal['anemoi.training.utils.masks.Boolean1DMask'], attribute_name: str)

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].

class anemoi.models.schemas.models.EncodersSchema(*, source_datasets: list[str], dataset_fusing_strategy: Literal['not_supported'] = 'not_supported', mapper: GNNEncoderSchema | GraphTransformerEncoderSchema | TransformerEncoderSchema | PointWiseForwardMapperSchema)

Bases: BaseModel

Encoder schema

source_datasets: list[str]: List of datasets for which the encoder is applicable.

dataset_fusing_strategy: Literal['not_supported']: Dataset fusing strategy. Default to ‘not_supported’.

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.models.schemas.models.DecodersSchema(*, target_datasets: list[str], target_node_features: list[~typing.Literal['coordinates', 'encoded_data', 'forcings', 'prognostics', 'trainable_parameters']] = <factory>, mapper: ~anemoi.models.schemas.decoder.GNNDecoderSchema | ~anemoi.models.schemas.decoder.GraphTransformerDecoderSchema | ~anemoi.models.schemas.decoder.TransformerDecoderSchema | ~anemoi.models.schemas.decoder.PointWiseBackwardMapperSchema)

Bases: BaseModel

Decoder schema

target_datasets: list[str]: List of datasets for which the decoder is applicable.

target_node_features: list[Literal[tuple(sorted(VALID_TARGET_FEATURES))]]: Whether to use the encoded latents from the encoder.

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.models.schemas.models.BaseModelSchema(*, keep_batch_sharded: bool = True, sparse_projector: ~anemoi.models.schemas.models.SparseProjectorSchema = <factory>, model: ~anemoi.models.schemas.models.Model = <factory>, node_trainable_parameters: dict[str, ~typing.Annotated[int, ~annotated_types.Ge(ge=0)]], bounding: dict[~typing.Literal['datasets'], dict[str, list[~typing.Annotated[~anemoi.models.schemas.bounding.ReluBoundingSchema | ~anemoi.models.schemas.bounding.LeakyReluBoundingSchema | ~anemoi.models.schemas.bounding.FractionBoundingSchema | ~anemoi.models.schemas.bounding.LeakyFractionBoundingSchema | ~anemoi.models.schemas.bounding.HardtanhBoundingSchema | ~anemoi.models.schemas.bounding.LeakyHardtanhBoundingSchema | ~anemoi.models.schemas.bounding.NormalizedReluBoundingSchema | ~anemoi.models.schemas.bounding.NormalizedLeakyReluBoundingSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]]]] | list[~typing.Annotated[~anemoi.models.schemas.bounding.ReluBoundingSchema | ~anemoi.models.schemas.bounding.LeakyReluBoundingSchema | ~anemoi.models.schemas.bounding.FractionBoundingSchema | ~anemoi.models.schemas.bounding.LeakyFractionBoundingSchema | ~anemoi.models.schemas.bounding.HardtanhBoundingSchema | ~anemoi.models.schemas.bounding.LeakyHardtanhBoundingSchema | ~anemoi.models.schemas.bounding.NormalizedReluBoundingSchema | ~anemoi.models.schemas.bounding.NormalizedLeakyReluBoundingSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]], output_mask: dict[~typing.Literal['datasets'], dict[str, ~anemoi.models.schemas.models.NoOutputMaskSchema | ~anemoi.models.schemas.models.Boolean1DSchema]] | ~anemoi.models.schemas.models.NoOutputMaskSchema | ~anemoi.models.schemas.models.Boolean1DSchema, latent_skip: bool = True, latent_aggregator: ~anemoi.models.schemas.aggregator.AggregatorSchema, processor: ~anemoi.models.schemas.processor.NoOpProcessorSchema | ~anemoi.models.schemas.processor.GNNProcessorSchema | ~anemoi.models.schemas.processor.GraphTransformerProcessorSchema | ~anemoi.models.schemas.processor.TransformerProcessorSchema | ~anemoi.models.schemas.processor.PointWiseMLPProcessorSchema, encoders: dict[str, ~anemoi.models.schemas.models.EncodersSchema], decoders: dict[str, ~anemoi.models.schemas.models.DecodersSchema], residual: dict[~typing.Literal['datasets'], dict[str, ~types.Annotated[~anemoi.models.schemas.residual.SkipConnectionSchema | ~anemoi.models.schemas.residual.TruncatedConnectionSchema | ~anemoi.models.schemas.residual.ScalarOrnsteinConnectionSchema | ~anemoi.models.schemas.residual.SpectralOrnsteinConnectionSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]]] | ~types.Annotated[~anemoi.models.schemas.residual.SkipConnectionSchema | ~anemoi.models.schemas.residual.TruncatedConnectionSchema | ~anemoi.models.schemas.residual.ScalarOrnsteinConnectionSchema | ~anemoi.models.schemas.residual.SpectralOrnsteinConnectionSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')], compile: list[dict[str, ~typing.Any]] | None = None, recompile_limit: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 8)

Bases: BaseModel

keep_batch_sharded: bool: Keep the input batch and the output of the model sharded

sparse_projector: SparseProjectorSchema: Sparse projection settings.

model: Model: Model schema.

node_trainable_parameters: dict[str, NonNegativeInt]: Learnable node and edge parameters.

bounding: DatasetDict[list[BoundingSchema]]: List of bounding configuration applied in order to the specified variables.

output_mask: DatasetDict[OutputMaskSchemas]: Output mask

latent_skip: bool: Add skip connection in latent space before/after processor.

latent_aggregator: AggregatorSchema: Latent aggregator schema.

processor: NoOpProcessorSchema | GNNProcessorSchema | GraphTransformerProcessorSchema | TransformerProcessorSchema | PointWiseMLPProcessorSchema: Model processor schema.

encoders: dict[str, EncodersSchema]: Model encoders schemas.

decoders: dict[str, DecodersSchema]: Model decoders schemas.

model_config = {}: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

residual: DatasetDict[ResidualConnectionSchema]: Residual connection schema.

compile: list[dict[str, Any]] | None: Modules to be compiled

recompile_limit: PositiveInt: How many times torch.compile will recompile a function for a given input shape.

classmethod cast_encoder_decoder_keys_to_str(data: Any) → Any: Cast encoder/decoder dict keys to str (YAML may parse them as int).

class anemoi.models.schemas.models.NoOpNoiseInjectorSchema(*, _target_: Literal['anemoi.models.layers.ensemble.NoOpNoiseInjector'])

Bases: BaseModel

Schema for NoOpNoiseInjector - passes input through unchanged.

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['anemoi.models.layers.ensemble.NoOpNoiseInjector']: No-op noise injector class

class anemoi.models.schemas.models.NoiseConditioningSchema(*, _target_: ~typing.Literal['anemoi.models.layers.ensemble.NoiseConditioning'], noise_std: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], noise_channels_dim: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], noise_mlp_hidden_dim: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], layer_kernels: dict[str, dict] | None = <factory>, noise_matrix: str | None = None, noise_edges_name: tuple[str, str, str] | None = None, edge_weight_attribute: str | None = None, row_normalize_noise_matrix: bool = False, autocast: bool = False)

Bases: BaseModel

Schema for NoiseConditioning - generates noise for conditioning.

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['anemoi.models.layers.ensemble.NoiseConditioning']: Noise conditioning layer class

noise_std: NonNegativeInt: Standard deviation of the noise to be injected.

noise_channels_dim: NonNegativeInt: Number of channels in the noise tensor.

noise_mlp_hidden_dim: NonNegativeInt: Hidden dimension of the MLP used to process the noise.

layer_kernels: dict[str, dict] | None: Settings related to custom kernels for encoder processor and decoder blocks

noise_matrix: str | None: Path to the noise projection matrix file (.npz). If None, no projection is applied.

noise_edges_name: tuple[str, str, str] | None: Edge type identifier (src, relation, dst) for graph-based noise projection.

edge_weight_attribute: str | None: Optional edge attribute name for graph-based noise projection weights.

row_normalize_noise_matrix: bool: Whether to row-normalize the noise projection matrix weights.

autocast: bool: Whether to use autocast for the noise projection matrix operations.

class anemoi.models.schemas.models.NoiseInjectorSchema(*, _target_: ~typing.Literal['anemoi.models.layers.ensemble.NoiseInjector'], noise_std: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], noise_channels_dim: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], noise_mlp_hidden_dim: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], layer_kernels: dict[str, dict] | None = <factory>)

Bases: BaseModel

Schema for NoiseInjector - injects noise directly into input tensor.

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['anemoi.models.layers.ensemble.NoiseInjector']: Noise injector layer class

noise_std: NonNegativeInt: Standard deviation of the noise to be injected.

noise_channels_dim: NonNegativeInt: Number of channels in the noise tensor.

noise_mlp_hidden_dim: NonNegativeInt: Hidden dimension of the MLP used to process the noise.

layer_kernels: dict[str, dict] | None: Settings related to custom kernels for encoder processor and decoder blocks

class anemoi.models.schemas.models.EnsModelSchema(*, keep_batch_sharded: bool = True, sparse_projector: ~anemoi.models.schemas.models.SparseProjectorSchema = <factory>, model: ~anemoi.models.schemas.models.Model = <factory>, node_trainable_parameters: dict[str, ~typing.Annotated[int, ~annotated_types.Ge(ge=0)]], bounding: dict[~typing.Literal['datasets'], dict[str, list[~typing.Annotated[~anemoi.models.schemas.bounding.ReluBoundingSchema | ~anemoi.models.schemas.bounding.LeakyReluBoundingSchema | ~anemoi.models.schemas.bounding.FractionBoundingSchema | ~anemoi.models.schemas.bounding.LeakyFractionBoundingSchema | ~anemoi.models.schemas.bounding.HardtanhBoundingSchema | ~anemoi.models.schemas.bounding.LeakyHardtanhBoundingSchema | ~anemoi.models.schemas.bounding.NormalizedReluBoundingSchema | ~anemoi.models.schemas.bounding.NormalizedLeakyReluBoundingSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]]]] | list[~typing.Annotated[~anemoi.models.schemas.bounding.ReluBoundingSchema | ~anemoi.models.schemas.bounding.LeakyReluBoundingSchema | ~anemoi.models.schemas.bounding.FractionBoundingSchema | ~anemoi.models.schemas.bounding.LeakyFractionBoundingSchema | ~anemoi.models.schemas.bounding.HardtanhBoundingSchema | ~anemoi.models.schemas.bounding.LeakyHardtanhBoundingSchema | ~anemoi.models.schemas.bounding.NormalizedReluBoundingSchema | ~anemoi.models.schemas.bounding.NormalizedLeakyReluBoundingSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]], output_mask: dict[~typing.Literal['datasets'], dict[str, ~anemoi.models.schemas.models.NoOutputMaskSchema | ~anemoi.models.schemas.models.Boolean1DSchema]] | ~anemoi.models.schemas.models.NoOutputMaskSchema | ~anemoi.models.schemas.models.Boolean1DSchema, latent_skip: bool = True, latent_aggregator: ~anemoi.models.schemas.aggregator.AggregatorSchema, processor: ~anemoi.models.schemas.processor.NoOpProcessorSchema | ~anemoi.models.schemas.processor.GNNProcessorSchema | ~anemoi.models.schemas.processor.GraphTransformerProcessorSchema | ~anemoi.models.schemas.processor.TransformerProcessorSchema | ~anemoi.models.schemas.processor.PointWiseMLPProcessorSchema, encoders: dict[str, ~anemoi.models.schemas.models.EncodersSchema], decoders: dict[str, ~anemoi.models.schemas.models.DecodersSchema], residual: dict[~typing.Literal['datasets'], dict[str, ~types.Annotated[~anemoi.models.schemas.residual.SkipConnectionSchema | ~anemoi.models.schemas.residual.TruncatedConnectionSchema | ~anemoi.models.schemas.residual.ScalarOrnsteinConnectionSchema | ~anemoi.models.schemas.residual.SpectralOrnsteinConnectionSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]]] | ~types.Annotated[~anemoi.models.schemas.residual.SkipConnectionSchema | ~anemoi.models.schemas.residual.TruncatedConnectionSchema | ~anemoi.models.schemas.residual.ScalarOrnsteinConnectionSchema | ~anemoi.models.schemas.residual.SpectralOrnsteinConnectionSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')], compile: list[dict[str, ~typing.Any]] | None = None, recompile_limit: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 8, noise_injector: ~anemoi.models.schemas.models.NoOpNoiseInjectorSchema | ~anemoi.models.schemas.models.NoiseConditioningSchema | ~anemoi.models.schemas.models.NoiseInjectorSchema, condition_on_residual: bool = False)

Bases: BaseModelSchema

model_config = {}: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

noise_injector: NoiseInjectorUnion: Noise injection configuration. Use NoOpNoiseInjector to disable, NoiseConditioning for conditioning, or NoiseInjector for direct injection.

condition_on_residual: bool: Whether to condition the noise injection on the residual connection.

class anemoi.models.schemas.models.TransportModelSchema(*, keep_batch_sharded: bool = True, sparse_projector: ~anemoi.models.schemas.models.SparseProjectorSchema = <factory>, model: ~anemoi.models.schemas.models.TransportModel = <factory>, node_trainable_parameters: dict[str, ~typing.Annotated[int, ~annotated_types.Ge(ge=0)]], bounding: dict[~typing.Literal['datasets'], dict[str, list[~typing.Annotated[~anemoi.models.schemas.bounding.ReluBoundingSchema | ~anemoi.models.schemas.bounding.LeakyReluBoundingSchema | ~anemoi.models.schemas.bounding.FractionBoundingSchema | ~anemoi.models.schemas.bounding.LeakyFractionBoundingSchema | ~anemoi.models.schemas.bounding.HardtanhBoundingSchema | ~anemoi.models.schemas.bounding.LeakyHardtanhBoundingSchema | ~anemoi.models.schemas.bounding.NormalizedReluBoundingSchema | ~anemoi.models.schemas.bounding.NormalizedLeakyReluBoundingSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]]]] | list[~typing.Annotated[~anemoi.models.schemas.bounding.ReluBoundingSchema | ~anemoi.models.schemas.bounding.LeakyReluBoundingSchema | ~anemoi.models.schemas.bounding.FractionBoundingSchema | ~anemoi.models.schemas.bounding.LeakyFractionBoundingSchema | ~anemoi.models.schemas.bounding.HardtanhBoundingSchema | ~anemoi.models.schemas.bounding.LeakyHardtanhBoundingSchema | ~anemoi.models.schemas.bounding.NormalizedReluBoundingSchema | ~anemoi.models.schemas.bounding.NormalizedLeakyReluBoundingSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]], output_mask: dict[~typing.Literal['datasets'], dict[str, ~anemoi.models.schemas.models.NoOutputMaskSchema | ~anemoi.models.schemas.models.Boolean1DSchema]] | ~anemoi.models.schemas.models.NoOutputMaskSchema | ~anemoi.models.schemas.models.Boolean1DSchema, latent_skip: bool = True, latent_aggregator: ~anemoi.models.schemas.aggregator.AggregatorSchema, processor: ~anemoi.models.schemas.processor.NoOpProcessorSchema | ~anemoi.models.schemas.processor.GNNProcessorSchema | ~anemoi.models.schemas.processor.GraphTransformerProcessorSchema | ~anemoi.models.schemas.processor.TransformerProcessorSchema | ~anemoi.models.schemas.processor.PointWiseMLPProcessorSchema, encoders: dict[str, ~anemoi.models.schemas.models.EncodersSchema], decoders: dict[str, ~anemoi.models.schemas.models.DecodersSchema], residual: dict[~typing.Literal['datasets'], dict[str, ~types.Annotated[~anemoi.models.schemas.residual.SkipConnectionSchema | ~anemoi.models.schemas.residual.TruncatedConnectionSchema | ~anemoi.models.schemas.residual.ScalarOrnsteinConnectionSchema | ~anemoi.models.schemas.residual.SpectralOrnsteinConnectionSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]]] | ~types.Annotated[~anemoi.models.schemas.residual.SkipConnectionSchema | ~anemoi.models.schemas.residual.TruncatedConnectionSchema | ~anemoi.models.schemas.residual.ScalarOrnsteinConnectionSchema | ~anemoi.models.schemas.residual.SpectralOrnsteinConnectionSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')], compile: list[dict[str, ~typing.Any]] | None = None, recompile_limit: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 8)

Bases: BaseModelSchema

model_config = {}: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model: TransportModel: Transport model schema.

class anemoi.models.schemas.models.TransportTendModelSchema(*, keep_batch_sharded: bool = True, sparse_projector: ~anemoi.models.schemas.models.SparseProjectorSchema = <factory>, model: ~anemoi.models.schemas.models.TransportModel = <factory>, node_trainable_parameters: dict[str, ~typing.Annotated[int, ~annotated_types.Ge(ge=0)]], bounding: dict[~typing.Literal['datasets'], dict[str, list[~typing.Annotated[~anemoi.models.schemas.bounding.ReluBoundingSchema | ~anemoi.models.schemas.bounding.LeakyReluBoundingSchema | ~anemoi.models.schemas.bounding.FractionBoundingSchema | ~anemoi.models.schemas.bounding.LeakyFractionBoundingSchema | ~anemoi.models.schemas.bounding.HardtanhBoundingSchema | ~anemoi.models.schemas.bounding.LeakyHardtanhBoundingSchema | ~anemoi.models.schemas.bounding.NormalizedReluBoundingSchema | ~anemoi.models.schemas.bounding.NormalizedLeakyReluBoundingSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]]]] | list[~typing.Annotated[~anemoi.models.schemas.bounding.ReluBoundingSchema | ~anemoi.models.schemas.bounding.LeakyReluBoundingSchema | ~anemoi.models.schemas.bounding.FractionBoundingSchema | ~anemoi.models.schemas.bounding.LeakyFractionBoundingSchema | ~anemoi.models.schemas.bounding.HardtanhBoundingSchema | ~anemoi.models.schemas.bounding.LeakyHardtanhBoundingSchema | ~anemoi.models.schemas.bounding.NormalizedReluBoundingSchema | ~anemoi.models.schemas.bounding.NormalizedLeakyReluBoundingSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]], output_mask: dict[~typing.Literal['datasets'], dict[str, ~anemoi.models.schemas.models.NoOutputMaskSchema | ~anemoi.models.schemas.models.Boolean1DSchema]] | ~anemoi.models.schemas.models.NoOutputMaskSchema | ~anemoi.models.schemas.models.Boolean1DSchema, latent_skip: bool = True, latent_aggregator: ~anemoi.models.schemas.aggregator.AggregatorSchema, processor: ~anemoi.models.schemas.processor.NoOpProcessorSchema | ~anemoi.models.schemas.processor.GNNProcessorSchema | ~anemoi.models.schemas.processor.GraphTransformerProcessorSchema | ~anemoi.models.schemas.processor.TransformerProcessorSchema | ~anemoi.models.schemas.processor.PointWiseMLPProcessorSchema, encoders: dict[str, ~anemoi.models.schemas.models.EncodersSchema], decoders: dict[str, ~anemoi.models.schemas.models.DecodersSchema], residual: dict[~typing.Literal['datasets'], dict[str, ~types.Annotated[~anemoi.models.schemas.residual.SkipConnectionSchema | ~anemoi.models.schemas.residual.TruncatedConnectionSchema | ~anemoi.models.schemas.residual.ScalarOrnsteinConnectionSchema | ~anemoi.models.schemas.residual.SpectralOrnsteinConnectionSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]]] | ~types.Annotated[~anemoi.models.schemas.residual.SkipConnectionSchema | ~anemoi.models.schemas.residual.TruncatedConnectionSchema | ~anemoi.models.schemas.residual.ScalarOrnsteinConnectionSchema | ~anemoi.models.schemas.residual.SpectralOrnsteinConnectionSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')], compile: list[dict[str, ~typing.Any]] | None = None, recompile_limit: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 8, condition_on_residual: bool = False)

Bases: TransportModelSchema

model_config = {}: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

condition_on_residual: bool: Whether to condition the noise injection on the residual connection.

class anemoi.models.schemas.models.HierarchicalModelSchema(*, keep_batch_sharded: bool = True, sparse_projector: ~anemoi.models.schemas.models.SparseProjectorSchema = <factory>, model: ~anemoi.models.schemas.models.Model = <factory>, node_trainable_parameters: dict[str, ~typing.Annotated[int, ~annotated_types.Ge(ge=0)]], bounding: dict[~typing.Literal['datasets'], dict[str, list[~typing.Annotated[~anemoi.models.schemas.bounding.ReluBoundingSchema | ~anemoi.models.schemas.bounding.LeakyReluBoundingSchema | ~anemoi.models.schemas.bounding.FractionBoundingSchema | ~anemoi.models.schemas.bounding.LeakyFractionBoundingSchema | ~anemoi.models.schemas.bounding.HardtanhBoundingSchema | ~anemoi.models.schemas.bounding.LeakyHardtanhBoundingSchema | ~anemoi.models.schemas.bounding.NormalizedReluBoundingSchema | ~anemoi.models.schemas.bounding.NormalizedLeakyReluBoundingSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]]]] | list[~typing.Annotated[~anemoi.models.schemas.bounding.ReluBoundingSchema | ~anemoi.models.schemas.bounding.LeakyReluBoundingSchema | ~anemoi.models.schemas.bounding.FractionBoundingSchema | ~anemoi.models.schemas.bounding.LeakyFractionBoundingSchema | ~anemoi.models.schemas.bounding.HardtanhBoundingSchema | ~anemoi.models.schemas.bounding.LeakyHardtanhBoundingSchema | ~anemoi.models.schemas.bounding.NormalizedReluBoundingSchema | ~anemoi.models.schemas.bounding.NormalizedLeakyReluBoundingSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]], output_mask: dict[~typing.Literal['datasets'], dict[str, ~anemoi.models.schemas.models.NoOutputMaskSchema | ~anemoi.models.schemas.models.Boolean1DSchema]] | ~anemoi.models.schemas.models.NoOutputMaskSchema | ~anemoi.models.schemas.models.Boolean1DSchema, latent_skip: bool = True, latent_aggregator: ~anemoi.models.schemas.aggregator.AggregatorSchema, processor: ~anemoi.models.schemas.processor.NoOpProcessorSchema | ~anemoi.models.schemas.processor.GNNProcessorSchema | ~anemoi.models.schemas.processor.GraphTransformerProcessorSchema | ~anemoi.models.schemas.processor.TransformerProcessorSchema | ~anemoi.models.schemas.processor.PointWiseMLPProcessorSchema, encoders: dict[str, ~anemoi.models.schemas.models.EncodersSchema], decoders: dict[str, ~anemoi.models.schemas.models.DecodersSchema], residual: dict[~typing.Literal['datasets'], dict[str, ~types.Annotated[~anemoi.models.schemas.residual.SkipConnectionSchema | ~anemoi.models.schemas.residual.TruncatedConnectionSchema | ~anemoi.models.schemas.residual.ScalarOrnsteinConnectionSchema | ~anemoi.models.schemas.residual.SpectralOrnsteinConnectionSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')]]] | ~types.Annotated[~anemoi.models.schemas.residual.SkipConnectionSchema | ~anemoi.models.schemas.residual.TruncatedConnectionSchema | ~anemoi.models.schemas.residual.ScalarOrnsteinConnectionSchema | ~anemoi.models.schemas.residual.SpectralOrnsteinConnectionSchema, FieldInfo(annotation=NoneType, required=True, discriminator='target_')], compile: list[dict[str, ~typing.Any]] | None = None, recompile_limit: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 8, enable_hierarchical_level_processing: bool = False, level_process_num_layers: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)] = 1, upscale_mapper: ~anemoi.models.schemas.encoder.GNNEncoderSchema | ~anemoi.models.schemas.encoder.GraphTransformerEncoderSchema | ~anemoi.models.schemas.encoder.TransformerEncoderSchema | ~anemoi.models.schemas.encoder.PointWiseForwardMapperSchema, downscale_mapper: ~anemoi.models.schemas.decoder.GNNDecoderSchema | ~anemoi.models.schemas.decoder.GraphTransformerDecoderSchema | ~anemoi.models.schemas.decoder.TransformerDecoderSchema | ~anemoi.models.schemas.decoder.PointWiseBackwardMapperSchema)

Bases: BaseModelSchema

model_config = {}: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

enable_hierarchical_level_processing: bool: Toggle to do message passing at every downscaling and upscaling step

level_process_num_layers: NonNegativeInt: Number of message passing steps at each level

upscale_mapper: GNNEncoderSchema | GraphTransformerEncoderSchema | TransformerEncoderSchema | PointWiseForwardMapperSchema: Mapper used to upscale from a lower level to a higher level in the hierarchy.

downscale_mapper: GNNDecoderSchema | GraphTransformerDecoderSchema | TransformerDecoderSchema | PointWiseBackwardMapperSchema: Mapper used to downscale from a higher level to a lower level in the hierarchy.

classmethod default_num_channels_in_hierarchical_mapper(data: Any) → Any

Allow num_channels to be omitted.

It will be set at model build time.

Training

class anemoi.training.schemas.training.GenericSchema(*, _target_: str, **extra_data: Any)

Bases: BaseModel

Generic Hydra instantiation schema with a required _target_ and arbitrary extra fields.

model_config = {'extra': 'allow', 'populate_by_name': True, '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_: str: Hydra target class or function to instantiate.

class anemoi.training.schemas.training.OptimizerSchema(*, _target_: str, **extra_data: Any)

Bases: GenericSchema

Hydra instantiation config for a PyTorch optimizer.

model_config = {'extra': 'allow', 'populate_by_name': True, '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.training.LRSchedulerSchema(*, _target_: str, **extra_data: Any)

Bases: GenericSchema

Hydra instantiation config for a learning rate scheduler.

model_config = {'extra': 'allow', 'populate_by_name': True, '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.training.PLSchedulerSchema(*, interval: str = 'step', **extra_data: Any)

Bases: BaseModel

PyTorch Lightning LRSchedulerConfig wrapper fields (interval, monitor, etc.).

model_config = {'extra': 'allow', '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].

interval: str: Interval at which Lightning calls lr_scheduler_step (‘step’ or ‘epoch’). Defaults to ‘step’.

class anemoi.training.schemas.training.GradientClip(*, val: float = 32.0, algorithm: ~typing.Annotated[str, ~pydantic.functional_validators.AfterValidator(func=functools.partial(<function allowed_values at 0x7a66c13020c0>, values=['value', 'norm']))])

Bases: BaseModel

Gradient clipping configuration.

val: float: Gradient clipping value.

algorithm: partial(<function allowed_values at 0x7a66c13020c0>, values=['value', 'norm']))]: The gradient clipping algorithm to use

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.training.WeightAveragingSchema(*, _target_: str, **extra_data: Any)

Bases: GenericSchema

Hydra instantiation config for a weight averaging callback (EMA or SWA).

Example

weight_averaging:: _target_: anemoi.training.diagnostics.callbacks.weight_averaging.EMAWeightAveraging decay: 0.999 update_starting_at_step: 1000

Stock pytorch_lightning.callbacks.*WeightAveraging classes can also be used. Set use_buffers=False when the model contains non-floating-point buffers.

model_config = {'extra': 'allow', 'populate_by_name': True, '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.training.OptimizationSchema(*, lr: ~typing.Annotated[float, ~annotated_types.Ge(ge=0)], optimizer: ~anemoi.training.schemas.training.OptimizerSchema, lr_scheduler: ~anemoi.training.schemas.training.LRSchedulerSchema | None = None, pl_lr_scheduler: ~anemoi.training.schemas.training.PLSchedulerSchema = <factory>)

Bases: BaseModel

Optimizer and LR scheduler configuration.

lr: Annotated[float, Ge(ge=0)]: Base learning rate per GPU. Scaled by hardware config at runtime.

optimizer: OptimizerSchema: Hydra instantiation config for the optimizer.

lr_scheduler: LRSchedulerSchema | None: Hydra instantiation config for the LR scheduler. If None, no scheduler is used.

pl_lr_scheduler: PLSchedulerSchema: PyTorch Lightning LRSchedulerConfig wrapper fields (interval, monitor, etc.).

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.training.ExplicitTimes(*, input: list[Annotated[int, Ge(ge=0)]], target: list[Annotated[int, Ge(ge=0)]])

Bases: BaseModel

Time indices for input and output.

Starts at index 0. Input and output can overlap.

input: list[Annotated[int, Ge(ge=0)]]: Input time indices.

target: list[Annotated[int, Ge(ge=0)]]: Target time indices.

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.training.TargetForcing(*, data: list[str], time_fraction: bool)

Bases: BaseModel

Forcing parameters for target output times.

Extra forcing parameters to use as input to distinguish between different target times.

data: list[str]: List of forcing parameters to use as input to the model at the interpolated step.

time_fraction: bool: Use target time as a fraction between input boundary times as input.

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.training.LossScalingSchema(*, default: int = 1, pl: dict[str, Annotated[float, Ge(ge=0)]], sfc: dict[str, Annotated[float, Ge(ge=0)]])

Bases: BaseModel

default: int: Default scaling value applied to the variables loss. Default to 1.

pl: dict[str, Annotated[float, Ge(ge=0)]]: Scaling value associated to each pressure level variable loss.

sfc: dict[str, Annotated[float, Ge(ge=0)]]: Scaling value associated to each surface variable loss.

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.training.GeneralVariableLossScalerSchema(*, _target_: Literal['anemoi.training.losses.scalers.GeneralVariableLossScaler'], weights: dict[str, float])

Bases: BaseModel

weights: dict[str, float]: Weight of each variable.

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.training.VariableMaskingScalerSchema(*, _target_: Literal['anemoi.training.losses.scalers.VariableMaskingLossScaler'], variables: list[str], invert: bool)

Bases: BaseModel

variables: list[str]: Variables to compute the loss over.

invert: bool: Flag to invert the variable mask.

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.training.NaNMaskScalerSchema(*, _target_: Literal['anemoi.training.losses.scalers.NaNMaskScaler'], use_processors_tendencies: bool = False)

Bases: BaseModel

use_processors_tendencies: bool: Flag to include processors for tendencies when building the loss mask.

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.training.TendencyScalerTargets(value): Bases: StrEnum

class anemoi.training.schemas.training.TendencyScalerSchema(*, _target_: TendencyScalerTargets, timestep: str | None = None)

Bases: BaseModel

timestep: str | None: Timestep key used to select tendency statistics for scalers.

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.training.VariableLevelScalerTargets(value): Bases: StrEnum

class anemoi.training.schemas.training.VariableLevelScalerSchema(*, _target_: VariableLevelScalerTargets, group: str, slope: float, y_intercept: float)

Bases: BaseModel

group: str: Group of variables to scale.

slope: float: Slope of scaling function.

y_intercept: float: Y-axis shift of scaling function.

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.training.GraphNodeAttributeScalerSchema(*, _target_: Literal['anemoi.training.losses.scalers.GraphNodeAttributeScaler'], nodes_attribute_name: str, norm: Literal['unit-max', 'unit-sum'] | None)

Bases: BaseModel

nodes_attribute_name: str: Name of the node attribute to return.

norm: Literal['unit-max', 'unit-sum'] | None: Normalisation method applied to the node attribute.

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.training.TimeStepScalerSchema(*, _target_: Literal['anemoi.training.losses.scalers.TimeStepScaler'], norm: Literal['unit-max', 'unit-sum'] | None = 'unit-sum', weights: list[float])

Bases: BaseModel

norm: Literal['unit-max', 'unit-sum'] | None: Normalisation method applied to the weights.

weights: list[float]: Weights for each time step.

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.training.UniformTimeStepScalerSchema(*, _target_: Literal['anemoi.training.losses.scalers.UniformTimeStepScaler'])

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].

class anemoi.training.schemas.training.LeadTimeDecayScalerSchema(*, _target_: Literal['anemoi.training.losses.scalers.LeadTimeDecayScaler'], output_lead_times: list[int], decay_factor: float, max_lead_time: int, decay_type: Literal['linear', 'exponential'] | None = 'linear', inverse: bool | None = False, norm: Literal['unit-max', 'unit-sum'] | None = 'unit-sum')

Bases: BaseModel

output_lead_times: list[int]: Lead times corresponding to each output step.

decay_factor: float: Decay factor for the lead time weights.

max_lead_time: int: Maximum lead time for decay calculation.

decay_type: Literal['linear', 'exponential'] | None: Type of decay to apply.

inverse: bool | None: If true, weights increase with lead time.

norm: Literal['unit-max', 'unit-sum'] | None: Normalisation method applied to the weights.

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.training.ReweightedGraphNodeAttributeScalerSchema(*, _target_: Literal['anemoi.training.losses.scalers.ReweightedGraphNodeAttributeScaler'], nodes_attribute_name: str, scaling_mask_attribute_name: str, weight_frac_of_total: float, norm: Literal['unit-max', 'unit-sum'] | None)

Bases: BaseModel

nodes_attribute_name: str: Name of the node attribute to return.

scaling_mask_attribute_name: str: Name of the node attribute to use as a mask to reweight the reference values.

weight_frac_of_total: float: Fraction of total weight to assign to nodes within the scaling mask. The remaining weight is distributed among

norm: Literal['unit-max', 'unit-sum'] | None: Normalisation method applied to the node attribute.

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.training.SpectralDimensionScalerSchema(*, _target_: Literal['anemoi.training.losses.scalers.SpectralDimensionScaler'], n_spectral_modes: Annotated[int, Gt(gt=0)], spectral_dims: Annotated[int, Gt(gt=0)] | None = None, norm: Literal['unit-sum', 'unit-mean', 'l1'] | None = None)

Bases: BaseModel

n_spectral_modes: Annotated[int, Gt(gt=0)]: Number of total wavenumbers (L dimension). For SHT-based losses this is truncation + 1.

spectral_dims: Annotated[int, Gt(gt=0)] | None: Length of the spectral dimension as seen by the loss. Defaults to n_spectral_modes.

norm: Literal['unit-sum', 'unit-mean', 'l1'] | None: Normalisation method applied to the scaler 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.training.schemas.training.ImplementedLossesUsingBaseLossSchema(value): Bases: StrEnum

class anemoi.training.schemas.training.CheckVariablesCompatibilitySchema(*, ignore_units: bool | list[str] = False, ignore_processing_period: bool | list[str] = False, ignore_time_processing: bool | list[str] = False, ignore_type_of_level: bool | list[str] = False)

Bases: BaseModel

Options forwarded to Variable.check_compatibility / Variable.compatible.

Each field may be set to True to skip the check for all variables, or to a list of variable names to skip only those specific variables.

ignore_units: bool | list[str]: Ignore unit mismatches. True skips all unit checks; a list of variable names skips only those variables.

ignore_processing_period: bool | list[str]: Ignore accumulation-period mismatches. True skips all period checks; a list of variable names skips only those variables.

ignore_time_processing: bool | list[str]: Ignore time-processing type mismatches. True skips all time-processing checks; a list of variable names skips only those variables.

ignore_type_of_level: bool | list[str]: Ignore mismatches in pressure level status. True skips all type-of-level checks; a list of variable names skips only those variables.

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.training.BaseLossSchema(*, _target_: ~anemoi.training.schemas.training.ImplementedLossesUsingBaseLossSchema, scalers: list[str], ignore_nans: bool = False, predicted_variables: list[str] | None = None, target_variables: list[str] | None = None, check_variables_compatibility: ~anemoi.training.schemas.training.CheckVariablesCompatibilitySchema = <factory>)

Bases: BaseModel

target_: ImplementedLossesUsingBaseLossSchema: Loss function object from anemoi.training.losses.

scalers: list[str]: Scalers to include in loss calculation

ignore_nans: bool: Allow nans in the loss and apply methods ignoring nans for measuring the loss.

check_variables_compatibility: CheckVariablesCompatibilitySchema: Options forwarded to Variable.check_compatibility when checking predicted vs. target variable units.

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.training.CRPSSchema(*, _target_: ~anemoi.training.schemas.training.ImplementedLossesUsingBaseLossSchema, scalers: list[str], ignore_nans: bool = False, predicted_variables: list[str] | None = None, target_variables: list[str] | None = None, check_variables_compatibility: ~anemoi.training.schemas.training.CheckVariablesCompatibilitySchema = <factory>, alpha: ~typing.Annotated[float, ~annotated_types.Ge(ge=0.0), ~annotated_types.Le(le=1.0)] = 0.95, backend: ~typing.Literal['naive', 'stable'] = 'stable', no_autocast: bool = True)

Bases: BaseLossSchema

alpha: float: Factor for linear combination of fair CRPS and standard CRPS. Values between 0 and 1 give the almost fair CRPS formulation.

backend: Literal['naive', 'stable']: Backend used for the point-wise CRPS calculation.

no_autocast: bool: Deactivate autocast for the kernel CRPS calculation

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.training.GraphLossMatrixSchema(*, edges_name: tuple[str, str, str], edge_weight_attribute: str | None = None, src_node_weight_attribute: str | None = None, row_normalize: bool = False)

Bases: BaseModel

One graph-backed smoothing matrix definition for multiscale loss.

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.training.GraphScoreGraphSchema(*, edges_name: tuple[str, str, str], edge_weight_attribute: str | None = None, src_node_weight_attribute: str | None = None, row_normalize: bool = False, validate_row_sums: bool = False)

Bases: BaseModel

Graph edge definition used by graph score losses.

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.training.EnsembleScoreLossSchema(*, scalers: list[str] = <factory>, ignore_nans: bool = False, predicted_variables: list[str] | None = None, target_variables: list[str] | None = None, check_variables_compatibility: CheckVariablesCompatibilitySchema = <factory>, no_autocast: bool = True)

Bases: BaseModel

Configuration shared by ensemble score losses.

model_config = {'extra': 'forbid', 'populate_by_name': True, '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.training.GraphScoreLossSchema(*, scalers: list[str] = <factory>, ignore_nans: bool = False, predicted_variables: list[str] | None = None, target_variables: list[str] | None = None, check_variables_compatibility: CheckVariablesCompatibilitySchema = <factory>, no_autocast: bool = True)

Bases: EnsembleScoreLossSchema

Configuration shared by graph score losses.

model_config = {'extra': 'forbid', 'populate_by_name': True, '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.training.GraphEnergyScoreLossSchema(*, scalers: list[str] = <factory>, ignore_nans: bool = False, predicted_variables: list[str] | None = None, target_variables: list[str] | None = None, check_variables_compatibility: ~anemoi.training.schemas.training.CheckVariablesCompatibilitySchema = <factory>, no_autocast: bool = True, _target_: ~typing.Literal['anemoi.training.losses.GraphEnergyScoreLoss'], fair: bool = True, loss_graph: ~anemoi.training.schemas.training.GraphScoreGraphSchema | None = None)

Bases: GraphScoreLossSchema

model_config = {'extra': 'forbid', 'populate_by_name': True, '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.training.EnergyScoreLossSchema(*, scalers: list[str] = <factory>, ignore_nans: bool = False, predicted_variables: list[str] | None = None, target_variables: list[str] | None = None, check_variables_compatibility: ~anemoi.training.schemas.training.CheckVariablesCompatibilitySchema = <factory>, no_autocast: bool = True, _target_: ~typing.Literal['anemoi.training.losses.EnergyScoreLoss'], fair: bool = True, norm_over: ~typing.Literal['spatial', 'variables', 'spatial_and_variables'] = 'spatial')

Bases: EnsembleScoreLossSchema

model_config = {'extra': 'forbid', 'populate_by_name': True, '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.training.GraphVariogramScoreLossSchema(*, scalers: list[str] = <factory>, ignore_nans: bool = False, predicted_variables: list[str] | None = None, target_variables: list[str] | None = None, check_variables_compatibility: ~anemoi.training.schemas.training.CheckVariablesCompatibilitySchema = <factory>, no_autocast: bool = True, _target_: ~typing.Literal['anemoi.training.losses.GraphVariogramScoreLoss'], loss_graph: ~anemoi.training.schemas.training.GraphScoreGraphSchema, p: ~typing.Annotated[float, ~annotated_types.Gt(gt=0)] = 1.0, fair: bool = True)

Bases: GraphScoreLossSchema

model_config = {'extra': 'forbid', 'populate_by_name': True, '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.training.GraphEdgeCRPSLossSchema(*, scalers: list[str] = <factory>, ignore_nans: bool = False, predicted_variables: list[str] | None = None, target_variables: list[str] | None = None, check_variables_compatibility: ~anemoi.training.schemas.training.CheckVariablesCompatibilitySchema = <factory>, no_autocast: bool = True, _target_: ~typing.Literal['anemoi.training.losses.GraphEdgeCRPSLoss'], loss_graph: ~anemoi.training.schemas.training.GraphScoreGraphSchema, alpha: ~typing.Annotated[float, ~annotated_types.Ge(ge=0.0), ~annotated_types.Le(le=1.0)] = 1.0)

Bases: GraphScoreLossSchema

model_config = {'extra': 'forbid', 'populate_by_name': True, '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.training.GraphEdgeEnergyScoreLossSchema(*, scalers: list[str] = <factory>, ignore_nans: bool = False, predicted_variables: list[str] | None = None, target_variables: list[str] | None = None, check_variables_compatibility: ~anemoi.training.schemas.training.CheckVariablesCompatibilitySchema = <factory>, no_autocast: bool = True, _target_: ~typing.Literal['anemoi.training.losses.GraphEdgeEnergyScoreLoss'], loss_graph: ~anemoi.training.schemas.training.GraphScoreGraphSchema, fair: bool = True)

Bases: GraphScoreLossSchema

model_config = {'extra': 'forbid', 'populate_by_name': True, '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.training.MultiscaleConfigDiskSchema(*, loss_matrices_path: str | None = None, loss_matrices: list[str | None], scalers: list[str] | None = None)

Bases: BaseModel

File-based multiscale config: smoothing matrices loaded from .npz files.

model_config = {'extra': 'forbid', 'populate_by_name': True, '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].

scalers: list[str] | None: Scalers to apply to the wrapped loss (delegated to inner per_scale_loss).

class anemoi.training.schemas.training.MultiscaleConfigOnTheFlySchema(*, num_scales: int | None = None, base_num_nearest_neighbours: int | None = None, base_sigma: float | None = None, scale_factor: int | None = None, smoothers: dict[str, dict] | None = None)

Bases: BaseModel

On-the-fly multiscale config: smoothing subgraphs built from the main graph.

model_config = {'extra': 'forbid', 'populate_by_name': True, '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.training.TimeAggregateLossWrapperSchema(*, _target_: Literal['anemoi.training.losses.aggregate.TimeAggregateLossWrapper'], time_aggregation_types: Annotated[list[Literal['diff', 'mean', 'min', 'max']], MinLen(min_length=1)], loss_fn: BaseLossSchema | CRPSSchema, scalers: list[str] | None = None)

Bases: BaseModel

Schema for TimeAggregateLossWrapper used inside CombinedLoss.

time_aggregation_types: list[Literal['diff', 'mean', 'min', 'max']]: Time aggregation operations to apply over the time dimension before computing the loss.

loss_fn: BaseLossSchema | CRPSSchema: Inner loss function applied to each time-aggregated output.

scalers: list[str] | None: Scalers to apply to the wrapped loss (delegated to inner loss_fn).

classmethod add_empty_scalers_to_inner(v: Any) → Any

Inject empty scalers for inner loss if missing; scalers flow through the wrapper.

This is needed to avoid validation errors on the inner loss when scalers are only defined at the wrapper level.

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.training.MultiScaleLossSchema(*, _target_: Literal['anemoi.training.losses.MultiscaleLossWrapper'], per_scale_loss: CRPSSchema | TimeAggregateLossWrapperSchema | GraphEnergyScoreLossSchema | EnergyScoreLossSchema | GraphVariogramScoreLossSchema | GraphEdgeCRPSLossSchema | GraphEdgeEnergyScoreLossSchema | BaseLossSchema, weights: list[float], multiscale_config: MultiscaleConfigDiskSchema | MultiscaleConfigOnTheFlySchema | None = None, sparse_projector_num_chunks: Annotated[int, Gt(gt=0)] = 1, loss_matrices_path: str | None = None, loss_matrices: list[str | None] | None = None)

Bases: BaseModel

classmethod add_empty_scalers_to_inner(v: Any) → Any

Inject empty scalers for inner loss if missing; scalers flow through the wrapper.

This is needed to avoid validation errors on the inner loss when scalers are only defined at the wrapper level.

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.training.HuberLossSchema(*, _target_: ~anemoi.training.schemas.training.ImplementedLossesUsingBaseLossSchema, scalers: list[str], ignore_nans: bool = False, predicted_variables: list[str] | None = None, target_variables: list[str] | None = None, check_variables_compatibility: ~anemoi.training.schemas.training.CheckVariablesCompatibilitySchema = <factory>, delta: float = 1.0)

Bases: BaseLossSchema

delta: float: Threshold for Huber loss.

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.training.SpectralProjectionConfigSchema(*, matrix_path: str | None = None, edges_name: tuple[str, str, str] | None = None, num_nearest_neighbours: int | None = None, sigma: float | None = None, grid: str | None = None, node_builder: dict | None = None, target_node_name: str = 'target_grid', edge_weight_attribute: str | None = None, src_node_weight_attribute: str | None = None, row_normalize: bool = False)

Bases: BaseModel

Configuration for sparse projection applied before the spectral transform.

Exactly one mode must be configured:

File mode: provide matrix_path.
Edge mode: provide edges_name.
Target-grid mode: provide num_nearest_neighbours and sigma together with either grid or node_builder.

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.training.SpectralLossSchema(*, _target_: ~anemoi.training.schemas.training.ImplementedLossesUsingBaseLossSchema, scalers: list[str], ignore_nans: bool = False, predicted_variables: list[str] | None = None, target_variables: list[str] | None = None, check_variables_compatibility: ~anemoi.training.schemas.training.CheckVariablesCompatibilitySchema = <factory>, transform: ~typing.Literal['fft2d', 'dct2d', 'regular_sht', 'reduced_sht', 'octahedral_sht'], subgrid: tuple[int, int | None] | ~typing.Literal['output_mask'] | None = None, projection_config: ~anemoi.training.schemas.training.SpectralProjectionConfigSchema | None = None, coefficient_magnitude: bool = False, **extra_data: ~typing.Any)

Bases: BaseLossSchema

Spectral loss class.

transform: Literal['fft2d', 'dct2d', 'regular_sht', 'reduced_sht', 'octahedral_sht']: Type of spectral transform to use.

subgrid: tuple[int, int | None] | Literal['output_mask'] | None: Optional slice or string to select a subgrid before the transform.

projection_config: SpectralProjectionConfigSchema | None: Optional sparse projection applied to the data before the spectral transform.

coefficient_magnitude: bool: For spectral CRPS, compare the moduli of the coefficients rather than their complex values.

class Config

Bases: Config

Override to allow extra parameters for spectral transforms.

model_config = {'extra': 'allow', '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.training.CombinedLossSchema(*, _target_: ~typing.Literal['anemoi.training.losses.combined.CombinedLoss'], scalers: list[str] = <factory>, ignore_nans: bool = False, predicted_variables: list[str] | None = None, target_variables: list[str] | None = None, check_variables_compatibility: ~anemoi.training.schemas.training.CheckVariablesCompatibilitySchema = <factory>, losses: ~typing.Annotated[list[~typing.Annotated[~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.BaseLossSchema, ~pydantic.types.Tag(tag=base)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.HuberLossSchema, ~pydantic.types.Tag(tag=huber)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CRPSSchema, ~pydantic.types.Tag(tag=crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.EnergyScoreLossSchema, ~pydantic.types.Tag(tag=energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphVariogramScoreLossSchema, ~pydantic.types.Tag(tag=graph_variogram_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeCRPSLossSchema, ~pydantic.types.Tag(tag=graph_edge_crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_edge_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.SpectralLossSchema, ~pydantic.types.Tag(tag=spectral)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.MultiScaleLossSchema, ~pydantic.types.Tag(tag=multiscale)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.TimeAggregateLossWrapperSchema, ~pydantic.types.Tag(tag=time_aggregate)], ~pydantic.types.Discriminator(discriminator=~anemoi.training.schemas.training._loss_discriminator, custom_error_type=None, custom_error_message=None, custom_error_context=None)]], ~annotated_types.MinLen(min_length=1)], loss_weights: list[int | float] | None = None)

Bases: BaseLossSchema

Schema for CombinedLoss.

Top-level scalers act as defaults for sub-losses that don’t specify their own. Sub-losses that explicitly set scalers override the parent value.

model_config = {'extra': 'forbid', 'populate_by_name': True, '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['anemoi.training.losses.combined.CombinedLoss']: CombinedLoss target.

scalers: list[str]: Optional top-level scalers propagated to sub-losses that don’t define their own.

losses: list[Annotated[Annotated[BaseLossSchema, Tag(tag=base)] | Annotated[HuberLossSchema, Tag(tag=huber)] | Annotated[CRPSSchema, Tag(tag=crps)] | Annotated[GraphEnergyScoreLossSchema, Tag(tag=graph_energy_score)] | Annotated[EnergyScoreLossSchema, Tag(tag=energy_score)] | Annotated[GraphVariogramScoreLossSchema, Tag(tag=graph_variogram_score)] | Annotated[GraphEdgeCRPSLossSchema, Tag(tag=graph_edge_crps)] | Annotated[GraphEdgeEnergyScoreLossSchema, Tag(tag=graph_edge_energy_score)] | Annotated[SpectralLossSchema, Tag(tag=spectral)] | Annotated[MultiScaleLossSchema, Tag(tag=multiscale)] | Annotated[TimeAggregateLossWrapperSchema, Tag(tag=time_aggregate)], Discriminator(discriminator=_loss_discriminator, custom_error_type=None, custom_error_message=None, custom_error_context=None)]]: Losses to combine, can be any of the normal losses.

loss_weights: list[int | float] | None: Weightings of losses, if not set, all losses are weighted equally.

classmethod propagate_scalers_to_children(data: Any) → Any

Propagate parent scalers to sub-losses that don’t specify their own.

MultiscaleLossWrapper is skipped because it manages scalers via per_scale_loss.

check_length_of_weights_and_losses() → Self: Check that the number of losses and weights match, or if not set, skip.

class anemoi.training.schemas.training.ImplementedStrategiesUsingBaseDDPStrategySchema(value): Bases: StrEnum

class anemoi.training.schemas.training.SingleDeviceStrategySchema(*, _target_: Literal['anemoi.training.distributed.single.SingleDeviceStrategy'], device: str = 'auto')

Bases: BaseModel

target_: Literal['anemoi.training.distributed.single.SingleDeviceStrategy']: Hydra target for the single device strategy.

device: str: Device to use for training.

model_config = {'extra': 'ignore', '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.training.BaseDDPStrategySchema(*, _target_: ImplementedStrategiesUsingBaseDDPStrategySchema, num_gpus_per_model: Annotated[int, Gt(gt=0)], read_group_size: Annotated[int, Gt(gt=0)], use_local_synchronization: bool = True, broadcast_buffers: bool = False, static_graph: bool = True, find_unused_parameters: bool = False)

Bases: BaseModel

Strategy configuration.

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_gpus_per_model: Annotated[int, Gt(gt=0)]: Number of GPUs per model.

read_group_size: Annotated[int, Gt(gt=0)]: Number of GPUs per reader group. Defaults to number of GPUs.

use_local_synchronization: bool: Use synchronization local to the group when creating process groups.

broadcast_buffers: bool: Broadcast model buffers at the start of each iteration. Defaults to False.

static_graph: bool: Tell DDP that the same parameters are trained in every step. Must be false when training.accum_grad_batches > 1.

find_unused_parameters: bool

false, DDP sums the gradients in the same groups in every step, so a resumed run continues bitwise identically.

Type:: Let DDP check every step for parameters that receive no gradient. Together with static_graph

class anemoi.training.schemas.training.DDPEnsGroupStrategyStrategySchema(*, _target_: ImplementedStrategiesUsingBaseDDPStrategySchema, num_gpus_per_model: Annotated[int, Gt(gt=0)], read_group_size: Annotated[int, Gt(gt=0)], use_local_synchronization: bool = True, broadcast_buffers: bool = False, static_graph: bool = True, find_unused_parameters: bool = False, num_gpus_per_ensemble: Annotated[int, Gt(gt=0)])

Bases: BaseDDPStrategySchema

Strategy object from anemoi.training.strategy.

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_gpus_per_ensemble: Annotated[int, Gt(gt=0)]: Number of GPUs per ensemble.

class anemoi.training.schemas.training.UpdateDsStatsOnCkptLoadSchema(*, states: bool = False, tendencies: bool = True)

Bases: BaseModel

Configuration for updating processor statistics on checkpoint load.

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].

states: bool: Rebuild state pre/post-processing statistics from the current dataset.

tendencies: bool: Rebuild tendency pre/post-processing statistics from the current dataset.

class anemoi.training.schemas.training.CheckpointPipelineSchema(*, source: GenericSchema | None = None, loading: GenericSchema | None = None, modifiers: list[GenericSchema] | None = None)

Bases: BaseModel

Declarative checkpoint pipeline consumed by build_checkpoint_pipeline.

Each block is a Hydra _target_ config: at most one acquisition source and one loading strategy, plus an ordered list of modifiers applied after loading. A None field configures no checkpoint pipeline (a fresh run).

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].

source: GenericSchema | None: Acquisition source stage (LocalSource / RunIdSource / S3Source / HTTPSource).

loading: GenericSchema | None: Loading strategy stage (WeightsOnlyLoader / TransferLearningLoader / WarmStartLoader / ColdStartLoader).

modifiers: list[GenericSchema] | None: Ordered model-modifier stages applied after loading (e.g. FreezingModifierStage).

class anemoi.training.schemas.training.BaseTrainingSchema(*, seed: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)] | None = None, check_variables_compatibility: ~anemoi.training.schemas.training.CheckVariablesCompatibilitySchema = <factory>, allow_variable_subset: bool = False, update_ds_stats_on_ckpt_load: ~anemoi.training.schemas.training.UpdateDsStatsOnCkptLoadSchema = <factory>, checkpoint: ~anemoi.training.schemas.training.CheckpointPipelineSchema | None = None, deterministic: bool = False, precision: str = '16-mixed', preferred_blas_backend: str | None = None, accum_grad_batches: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 1, num_sanity_val_steps: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], gradient_clip: ~anemoi.training.schemas.training.GradientClip, strategy: ~anemoi.training.schemas.training.BaseDDPStrategySchema | ~anemoi.training.schemas.training.DDPEnsGroupStrategyStrategySchema | ~anemoi.training.schemas.training.SingleDeviceStrategySchema, training_loss: dict[~typing.Literal['datasets'], dict[str, ~typing.Annotated[~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.BaseLossSchema, ~pydantic.types.Tag(tag=base)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.HuberLossSchema, ~pydantic.types.Tag(tag=huber)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CombinedLossSchema, ~pydantic.types.Tag(tag=combined)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CRPSSchema, ~pydantic.types.Tag(tag=crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.EnergyScoreLossSchema, ~pydantic.types.Tag(tag=energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphVariogramScoreLossSchema, ~pydantic.types.Tag(tag=graph_variogram_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeCRPSLossSchema, ~pydantic.types.Tag(tag=graph_edge_crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_edge_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.SpectralLossSchema, ~pydantic.types.Tag(tag=spectral)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.TimeAggregateLossWrapperSchema, ~pydantic.types.Tag(tag=time_aggregate)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.MultiScaleLossSchema, ~pydantic.types.Tag(tag=multiscale)], ~pydantic.types.Discriminator(discriminator=~anemoi.training.schemas.training._loss_discriminator, custom_error_type=None, custom_error_message=None, custom_error_context=None)]]] | ~typing.Annotated[~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.BaseLossSchema, ~pydantic.types.Tag(tag=base)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.HuberLossSchema, ~pydantic.types.Tag(tag=huber)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CombinedLossSchema, ~pydantic.types.Tag(tag=combined)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CRPSSchema, ~pydantic.types.Tag(tag=crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.EnergyScoreLossSchema, ~pydantic.types.Tag(tag=energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphVariogramScoreLossSchema, ~pydantic.types.Tag(tag=graph_variogram_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeCRPSLossSchema, ~pydantic.types.Tag(tag=graph_edge_crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_edge_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.SpectralLossSchema, ~pydantic.types.Tag(tag=spectral)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.TimeAggregateLossWrapperSchema, ~pydantic.types.Tag(tag=time_aggregate)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.MultiScaleLossSchema, ~pydantic.types.Tag(tag=multiscale)], ~pydantic.types.Discriminator(discriminator=~anemoi.training.schemas.training._loss_discriminator, custom_error_type=None, custom_error_message=None, custom_error_context=None)], weight_averaging: ~anemoi.training.schemas.training.WeightAveragingSchema | None = None, loss_gradient_scaling: bool = False, scalers: dict[~typing.Literal['datasets'], dict[str, dict[str, ~anemoi.training.schemas.training.GeneralVariableLossScalerSchema | ~anemoi.training.schemas.training.VariableLevelScalerSchema | ~anemoi.training.schemas.training.VariableMaskingScalerSchema | ~anemoi.training.schemas.training.TendencyScalerSchema | ~anemoi.training.schemas.training.NaNMaskScalerSchema | ~anemoi.training.schemas.training.GraphNodeAttributeScalerSchema | ~anemoi.training.schemas.training.TimeStepScalerSchema | ~anemoi.training.schemas.training.UniformTimeStepScalerSchema | ~anemoi.training.schemas.training.LeadTimeDecayScalerSchema | ~anemoi.training.schemas.training.ReweightedGraphNodeAttributeScalerSchema | ~anemoi.training.schemas.training.SpectralDimensionScalerSchema]]] | dict[str, ~anemoi.training.schemas.training.GeneralVariableLossScalerSchema | ~anemoi.training.schemas.training.VariableLevelScalerSchema | ~anemoi.training.schemas.training.VariableMaskingScalerSchema | ~anemoi.training.schemas.training.TendencyScalerSchema | ~anemoi.training.schemas.training.NaNMaskScalerSchema | ~anemoi.training.schemas.training.GraphNodeAttributeScalerSchema | ~anemoi.training.schemas.training.TimeStepScalerSchema | ~anemoi.training.schemas.training.UniformTimeStepScalerSchema | ~anemoi.training.schemas.training.LeadTimeDecayScalerSchema | ~anemoi.training.schemas.training.ReweightedGraphNodeAttributeScalerSchema | ~anemoi.training.schemas.training.SpectralDimensionScalerSchema], validation_metrics: dict[~typing.Literal['datasets'], dict[str, dict[str, ~typing.Annotated[~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.BaseLossSchema, ~pydantic.types.Tag(tag=base)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.HuberLossSchema, ~pydantic.types.Tag(tag=huber)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CombinedLossSchema, ~pydantic.types.Tag(tag=combined)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CRPSSchema, ~pydantic.types.Tag(tag=crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.EnergyScoreLossSchema, ~pydantic.types.Tag(tag=energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphVariogramScoreLossSchema, ~pydantic.types.Tag(tag=graph_variogram_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeCRPSLossSchema, ~pydantic.types.Tag(tag=graph_edge_crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_edge_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.SpectralLossSchema, ~pydantic.types.Tag(tag=spectral)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.TimeAggregateLossWrapperSchema, ~pydantic.types.Tag(tag=time_aggregate)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.MultiScaleLossSchema, ~pydantic.types.Tag(tag=multiscale)], ~pydantic.types.Discriminator(discriminator=~anemoi.training.schemas.training._loss_discriminator, custom_error_type=None, custom_error_message=None, custom_error_context=None)]] | None]] | dict[str, ~typing.Annotated[~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.BaseLossSchema, ~pydantic.types.Tag(tag=base)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.HuberLossSchema, ~pydantic.types.Tag(tag=huber)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CombinedLossSchema, ~pydantic.types.Tag(tag=combined)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CRPSSchema, ~pydantic.types.Tag(tag=crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.EnergyScoreLossSchema, ~pydantic.types.Tag(tag=energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphVariogramScoreLossSchema, ~pydantic.types.Tag(tag=graph_variogram_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeCRPSLossSchema, ~pydantic.types.Tag(tag=graph_edge_crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_edge_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.SpectralLossSchema, ~pydantic.types.Tag(tag=spectral)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.TimeAggregateLossWrapperSchema, ~pydantic.types.Tag(tag=time_aggregate)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.MultiScaleLossSchema, ~pydantic.types.Tag(tag=multiscale)], ~pydantic.types.Discriminator(discriminator=~anemoi.training.schemas.training._loss_discriminator, custom_error_type=None, custom_error_message=None, custom_error_context=None)]] | None, variable_groups: dict[~typing.Literal['datasets'], dict[str, dict[str, str | list[str] | dict[str, str | bool | list[str | int]]] | None]] | dict[str, str | list[str] | dict[str, str | bool | list[str | int]]] | None, max_epochs: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] | None = None, max_steps: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 150000, optimization: ~anemoi.training.schemas.training.OptimizationSchema, metrics: dict[~typing.Literal['datasets'], dict[str, list[str]]] | list[str], ensemble_size_per_device: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 1)

Bases: BaseModel

Training configuration.

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].

seed: Annotated[int, Ge(ge=0)] | None: Base seed of all random number generators. If null, a run continuing from a checkpoint takes the checkpoint’s seed and a new run draws one. A seed that differs from the checkpoint’s changes the rest of the interrupted epoch.

check_variables_compatibility: CheckVariablesCompatibilitySchema: Options forwarded to Variable.check_compatibility when checking checkpoint vs. current dataset (fine-tuning).

allow_variable_subset: bool: Allow fine-tuning into a model with FEWER variables (issue #838). When True, a checkpoint

update_ds_stats_on_ckpt_load: UpdateDsStatsOnCkptLoadSchema: Rebuild pre/post-processing statistics from the current dataset when loading a checkpoint.

checkpoint: CheckpointPipelineSchema | None

an

Type:: Checkpoint pipeline configuration consumed by build_checkpoint_pipeline

deterministic: bool: This flag sets torch.backends.cudnn.deterministic. It may reduce nondeterminism, but does not guarantee exact

precision: str: Precision

preferred_blas_backend: str | None: Optionally override PyTorch’s default BLAS backend.

accum_grad_batches: Annotated[int, Gt(gt=0)]: Accumulates gradients over k batches before stepping the optimizer. K >= 1 (if K == 1 then no accumulation). The effective bacthsize becomes num-device * k.

num_sanity_val_steps: Annotated[int, Ge(ge=0)]: Sanity check runs n batches of val before starting the training routine.

gradient_clip: GradientClip: Config for gradient clipping.

strategy: BaseDDPStrategySchema | DDPEnsGroupStrategyStrategySchema | SingleDeviceStrategySchema: Strategy to use.

training_loss: dict[Literal['datasets'], dict[str, Annotated[Annotated[BaseLossSchema, Tag(tag=base)] | Annotated[HuberLossSchema, Tag(tag=huber)] | Annotated[CombinedLossSchema, Tag(tag=combined)] | Annotated[CRPSSchema, Tag(tag=crps)] | Annotated[GraphEnergyScoreLossSchema, Tag(tag=graph_energy_score)] | Annotated[EnergyScoreLossSchema, Tag(tag=energy_score)] | Annotated[GraphVariogramScoreLossSchema, Tag(tag=graph_variogram_score)] | Annotated[GraphEdgeCRPSLossSchema, Tag(tag=graph_edge_crps)] | Annotated[GraphEdgeEnergyScoreLossSchema, Tag(tag=graph_edge_energy_score)] | Annotated[SpectralLossSchema, Tag(tag=spectral)] | Annotated[TimeAggregateLossWrapperSchema, Tag(tag=time_aggregate)] | Annotated[MultiScaleLossSchema, Tag(tag=multiscale)], Discriminator(discriminator=_loss_discriminator, custom_error_type=None, custom_error_message=None, custom_error_context=None)]]] | Annotated[Annotated[BaseLossSchema, Tag(tag=base)] | Annotated[HuberLossSchema, Tag(tag=huber)] | Annotated[CombinedLossSchema, Tag(tag=combined)] | Annotated[CRPSSchema, Tag(tag=crps)] | Annotated[GraphEnergyScoreLossSchema, Tag(tag=graph_energy_score)] | Annotated[EnergyScoreLossSchema, Tag(tag=energy_score)] | Annotated[GraphVariogramScoreLossSchema, Tag(tag=graph_variogram_score)] | Annotated[GraphEdgeCRPSLossSchema, Tag(tag=graph_edge_crps)] | Annotated[GraphEdgeEnergyScoreLossSchema, Tag(tag=graph_edge_energy_score)] | Annotated[SpectralLossSchema, Tag(tag=spectral)] | Annotated[TimeAggregateLossWrapperSchema, Tag(tag=time_aggregate)] | Annotated[MultiScaleLossSchema, Tag(tag=multiscale)], Discriminator(discriminator=_loss_discriminator, custom_error_type=None, custom_error_message=None, custom_error_context=None)]: Training loss configuration.

weight_averaging: WeightAveragingSchema | None: Config for weight averaging (SWA or EMA). Set to null to disable.

loss_gradient_scaling: bool: Dynamic rescaling of the loss gradient. Not yet tested.

scalers: dict[Literal['datasets'], dict[str, dict[str, GeneralVariableLossScalerSchema | VariableLevelScalerSchema | VariableMaskingScalerSchema | TendencyScalerSchema | NaNMaskScalerSchema | GraphNodeAttributeScalerSchema | TimeStepScalerSchema | UniformTimeStepScalerSchema | LeadTimeDecayScalerSchema | ReweightedGraphNodeAttributeScalerSchema | SpectralDimensionScalerSchema]]] | dict[str, GeneralVariableLossScalerSchema | VariableLevelScalerSchema | VariableMaskingScalerSchema | TendencyScalerSchema | NaNMaskScalerSchema | GraphNodeAttributeScalerSchema | TimeStepScalerSchema | UniformTimeStepScalerSchema | LeadTimeDecayScalerSchema | ReweightedGraphNodeAttributeScalerSchema | SpectralDimensionScalerSchema]: Scalers to use in the computation of the loss and validation scores.

validation_metrics: dict[Literal['datasets'], dict[str, dict[str, Annotated[Annotated[BaseLossSchema, Tag(tag=base)] | Annotated[HuberLossSchema, Tag(tag=huber)] | Annotated[CombinedLossSchema, Tag(tag=combined)] | Annotated[CRPSSchema, Tag(tag=crps)] | Annotated[GraphEnergyScoreLossSchema, Tag(tag=graph_energy_score)] | Annotated[EnergyScoreLossSchema, Tag(tag=energy_score)] | Annotated[GraphVariogramScoreLossSchema, Tag(tag=graph_variogram_score)] | Annotated[GraphEdgeCRPSLossSchema, Tag(tag=graph_edge_crps)] | Annotated[GraphEdgeEnergyScoreLossSchema, Tag(tag=graph_edge_energy_score)] | Annotated[SpectralLossSchema, Tag(tag=spectral)] | Annotated[TimeAggregateLossWrapperSchema, Tag(tag=time_aggregate)] | Annotated[MultiScaleLossSchema, Tag(tag=multiscale)], Discriminator(discriminator=_loss_discriminator, custom_error_type=None, custom_error_message=None, custom_error_context=None)]] | None]] | dict[str, Annotated[Annotated[BaseLossSchema, Tag(tag=base)] | Annotated[HuberLossSchema, Tag(tag=huber)] | Annotated[CombinedLossSchema, Tag(tag=combined)] | Annotated[CRPSSchema, Tag(tag=crps)] | Annotated[GraphEnergyScoreLossSchema, Tag(tag=graph_energy_score)] | Annotated[EnergyScoreLossSchema, Tag(tag=energy_score)] | Annotated[GraphVariogramScoreLossSchema, Tag(tag=graph_variogram_score)] | Annotated[GraphEdgeCRPSLossSchema, Tag(tag=graph_edge_crps)] | Annotated[GraphEdgeEnergyScoreLossSchema, Tag(tag=graph_edge_energy_score)] | Annotated[SpectralLossSchema, Tag(tag=spectral)] | Annotated[TimeAggregateLossWrapperSchema, Tag(tag=time_aggregate)] | Annotated[MultiScaleLossSchema, Tag(tag=multiscale)], Discriminator(discriminator=_loss_discriminator, custom_error_type=None, custom_error_message=None, custom_error_context=None)]] | None: List of validation metrics configurations.

variable_groups: dict[Literal['datasets'], dict[str, dict[str, str | list[str] | dict[str, str | bool | list[str | int]]] | None]] | dict[str, str | list[str] | dict[str, str | bool | list[str | int]]] | None: Groups for variable loss scaling

max_epochs: Annotated[int, Gt(gt=0)] | None: Maximum number of epochs, stops earlier if max_steps is reached first.

max_steps: Annotated[int, Gt(gt=0)]: Maximum number of steps, stops earlier if max_epochs is reached first.

optimization: OptimizationSchema: Optimizer and LR scheduler configuration.

metrics: dict[Literal['datasets'], dict[str, list[str]]] | list[str]: List of metrics

ensemble_size_per_device: Annotated[int, Gt(gt=0)]: Number of ensemble members per device. Default is 1 for non-ensemble forecasting.

check_static_graph_without_accumulation() → Self: DDP’s static graph does not support accumulating gradients over several batches.

class anemoi.training.schemas.training.SingleTrainingMethodSchema(*, _target_: Literal['anemoi.training.train.methods.SingleTraining'])

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['anemoi.training.train.methods.SingleTraining']: Hydra target for the single training method.

class anemoi.training.schemas.training.EnsembleTrainingMethodSchema(*, _target_: Literal['anemoi.training.train.methods.EnsembleTraining'])

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['anemoi.training.train.methods.EnsembleTraining']: Hydra target for the ensemble training method.

class anemoi.training.schemas.training.TransportTrainingMethodSchema(*, _target_: Literal['anemoi.training.train.methods.TransportTraining'])

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['anemoi.training.train.methods.TransportTraining']: Hydra target for the transport training method.

class anemoi.training.schemas.training.SingleTrainingSchema(*, seed: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)] | None = None, check_variables_compatibility: ~anemoi.training.schemas.training.CheckVariablesCompatibilitySchema = <factory>, allow_variable_subset: bool = False, update_ds_stats_on_ckpt_load: ~anemoi.training.schemas.training.UpdateDsStatsOnCkptLoadSchema = <factory>, checkpoint: ~anemoi.training.schemas.training.CheckpointPipelineSchema | None = None, deterministic: bool = False, precision: str = '16-mixed', preferred_blas_backend: str | None = None, accum_grad_batches: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 1, num_sanity_val_steps: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], gradient_clip: ~anemoi.training.schemas.training.GradientClip, strategy: ~anemoi.training.schemas.training.BaseDDPStrategySchema | ~anemoi.training.schemas.training.DDPEnsGroupStrategyStrategySchema | ~anemoi.training.schemas.training.SingleDeviceStrategySchema, training_loss: dict[~typing.Literal['datasets'], dict[str, ~typing.Annotated[~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.BaseLossSchema, ~pydantic.types.Tag(tag=base)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.HuberLossSchema, ~pydantic.types.Tag(tag=huber)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CombinedLossSchema, ~pydantic.types.Tag(tag=combined)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CRPSSchema, ~pydantic.types.Tag(tag=crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.EnergyScoreLossSchema, ~pydantic.types.Tag(tag=energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphVariogramScoreLossSchema, ~pydantic.types.Tag(tag=graph_variogram_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeCRPSLossSchema, ~pydantic.types.Tag(tag=graph_edge_crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_edge_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.SpectralLossSchema, ~pydantic.types.Tag(tag=spectral)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.TimeAggregateLossWrapperSchema, ~pydantic.types.Tag(tag=time_aggregate)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.MultiScaleLossSchema, ~pydantic.types.Tag(tag=multiscale)], ~pydantic.types.Discriminator(discriminator=~anemoi.training.schemas.training._loss_discriminator, custom_error_type=None, custom_error_message=None, custom_error_context=None)]]] | ~typing.Annotated[~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.BaseLossSchema, ~pydantic.types.Tag(tag=base)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.HuberLossSchema, ~pydantic.types.Tag(tag=huber)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CombinedLossSchema, ~pydantic.types.Tag(tag=combined)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CRPSSchema, ~pydantic.types.Tag(tag=crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.EnergyScoreLossSchema, ~pydantic.types.Tag(tag=energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphVariogramScoreLossSchema, ~pydantic.types.Tag(tag=graph_variogram_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeCRPSLossSchema, ~pydantic.types.Tag(tag=graph_edge_crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_edge_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.SpectralLossSchema, ~pydantic.types.Tag(tag=spectral)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.TimeAggregateLossWrapperSchema, ~pydantic.types.Tag(tag=time_aggregate)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.MultiScaleLossSchema, ~pydantic.types.Tag(tag=multiscale)], ~pydantic.types.Discriminator(discriminator=~anemoi.training.schemas.training._loss_discriminator, custom_error_type=None, custom_error_message=None, custom_error_context=None)], weight_averaging: ~anemoi.training.schemas.training.WeightAveragingSchema | None = None, loss_gradient_scaling: bool = False, scalers: dict[~typing.Literal['datasets'], dict[str, dict[str, ~anemoi.training.schemas.training.GeneralVariableLossScalerSchema | ~anemoi.training.schemas.training.VariableLevelScalerSchema | ~anemoi.training.schemas.training.VariableMaskingScalerSchema | ~anemoi.training.schemas.training.TendencyScalerSchema | ~anemoi.training.schemas.training.NaNMaskScalerSchema | ~anemoi.training.schemas.training.GraphNodeAttributeScalerSchema | ~anemoi.training.schemas.training.TimeStepScalerSchema | ~anemoi.training.schemas.training.UniformTimeStepScalerSchema | ~anemoi.training.schemas.training.LeadTimeDecayScalerSchema | ~anemoi.training.schemas.training.ReweightedGraphNodeAttributeScalerSchema | ~anemoi.training.schemas.training.SpectralDimensionScalerSchema]]] | dict[str, ~anemoi.training.schemas.training.GeneralVariableLossScalerSchema | ~anemoi.training.schemas.training.VariableLevelScalerSchema | ~anemoi.training.schemas.training.VariableMaskingScalerSchema | ~anemoi.training.schemas.training.TendencyScalerSchema | ~anemoi.training.schemas.training.NaNMaskScalerSchema | ~anemoi.training.schemas.training.GraphNodeAttributeScalerSchema | ~anemoi.training.schemas.training.TimeStepScalerSchema | ~anemoi.training.schemas.training.UniformTimeStepScalerSchema | ~anemoi.training.schemas.training.LeadTimeDecayScalerSchema | ~anemoi.training.schemas.training.ReweightedGraphNodeAttributeScalerSchema | ~anemoi.training.schemas.training.SpectralDimensionScalerSchema], validation_metrics: dict[~typing.Literal['datasets'], dict[str, dict[str, ~typing.Annotated[~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.BaseLossSchema, ~pydantic.types.Tag(tag=base)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.HuberLossSchema, ~pydantic.types.Tag(tag=huber)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CombinedLossSchema, ~pydantic.types.Tag(tag=combined)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CRPSSchema, ~pydantic.types.Tag(tag=crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.EnergyScoreLossSchema, ~pydantic.types.Tag(tag=energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphVariogramScoreLossSchema, ~pydantic.types.Tag(tag=graph_variogram_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeCRPSLossSchema, ~pydantic.types.Tag(tag=graph_edge_crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_edge_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.SpectralLossSchema, ~pydantic.types.Tag(tag=spectral)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.TimeAggregateLossWrapperSchema, ~pydantic.types.Tag(tag=time_aggregate)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.MultiScaleLossSchema, ~pydantic.types.Tag(tag=multiscale)], ~pydantic.types.Discriminator(discriminator=~anemoi.training.schemas.training._loss_discriminator, custom_error_type=None, custom_error_message=None, custom_error_context=None)]] | None]] | dict[str, ~typing.Annotated[~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.BaseLossSchema, ~pydantic.types.Tag(tag=base)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.HuberLossSchema, ~pydantic.types.Tag(tag=huber)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CombinedLossSchema, ~pydantic.types.Tag(tag=combined)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CRPSSchema, ~pydantic.types.Tag(tag=crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.EnergyScoreLossSchema, ~pydantic.types.Tag(tag=energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphVariogramScoreLossSchema, ~pydantic.types.Tag(tag=graph_variogram_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeCRPSLossSchema, ~pydantic.types.Tag(tag=graph_edge_crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_edge_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.SpectralLossSchema, ~pydantic.types.Tag(tag=spectral)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.TimeAggregateLossWrapperSchema, ~pydantic.types.Tag(tag=time_aggregate)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.MultiScaleLossSchema, ~pydantic.types.Tag(tag=multiscale)], ~pydantic.types.Discriminator(discriminator=~anemoi.training.schemas.training._loss_discriminator, custom_error_type=None, custom_error_message=None, custom_error_context=None)]] | None, variable_groups: dict[~typing.Literal['datasets'], dict[str, dict[str, str | list[str] | dict[str, str | bool | list[str | int]]] | None]] | dict[str, str | list[str] | dict[str, str | bool | list[str | int]]] | None, max_epochs: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] | None = None, max_steps: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 150000, optimization: ~anemoi.training.schemas.training.OptimizationSchema, metrics: dict[~typing.Literal['datasets'], dict[str, list[str]]] | list[str], ensemble_size_per_device: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 1, method: ~anemoi.training.schemas.training.SingleTrainingMethodSchema)

Bases: BaseTrainingSchema

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].

method: SingleTrainingMethodSchema: Training method.

class anemoi.training.schemas.training.EnsembleTrainingSchema(*, seed: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)] | None = None, check_variables_compatibility: ~anemoi.training.schemas.training.CheckVariablesCompatibilitySchema = <factory>, allow_variable_subset: bool = False, update_ds_stats_on_ckpt_load: ~anemoi.training.schemas.training.UpdateDsStatsOnCkptLoadSchema = <factory>, checkpoint: ~anemoi.training.schemas.training.CheckpointPipelineSchema | None = None, deterministic: bool = False, precision: str = '16-mixed', preferred_blas_backend: str | None = None, accum_grad_batches: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 1, num_sanity_val_steps: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], gradient_clip: ~anemoi.training.schemas.training.GradientClip, strategy: ~anemoi.training.schemas.training.BaseDDPStrategySchema | ~anemoi.training.schemas.training.DDPEnsGroupStrategyStrategySchema | ~anemoi.training.schemas.training.SingleDeviceStrategySchema, training_loss: dict[~typing.Literal['datasets'], dict[str, ~typing.Annotated[~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.BaseLossSchema, ~pydantic.types.Tag(tag=base)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.HuberLossSchema, ~pydantic.types.Tag(tag=huber)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CombinedLossSchema, ~pydantic.types.Tag(tag=combined)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CRPSSchema, ~pydantic.types.Tag(tag=crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.EnergyScoreLossSchema, ~pydantic.types.Tag(tag=energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphVariogramScoreLossSchema, ~pydantic.types.Tag(tag=graph_variogram_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeCRPSLossSchema, ~pydantic.types.Tag(tag=graph_edge_crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_edge_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.SpectralLossSchema, ~pydantic.types.Tag(tag=spectral)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.TimeAggregateLossWrapperSchema, ~pydantic.types.Tag(tag=time_aggregate)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.MultiScaleLossSchema, ~pydantic.types.Tag(tag=multiscale)], ~pydantic.types.Discriminator(discriminator=~anemoi.training.schemas.training._loss_discriminator, custom_error_type=None, custom_error_message=None, custom_error_context=None)]]] | ~typing.Annotated[~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.BaseLossSchema, ~pydantic.types.Tag(tag=base)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.HuberLossSchema, ~pydantic.types.Tag(tag=huber)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CombinedLossSchema, ~pydantic.types.Tag(tag=combined)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CRPSSchema, ~pydantic.types.Tag(tag=crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.EnergyScoreLossSchema, ~pydantic.types.Tag(tag=energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphVariogramScoreLossSchema, ~pydantic.types.Tag(tag=graph_variogram_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeCRPSLossSchema, ~pydantic.types.Tag(tag=graph_edge_crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_edge_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.SpectralLossSchema, ~pydantic.types.Tag(tag=spectral)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.TimeAggregateLossWrapperSchema, ~pydantic.types.Tag(tag=time_aggregate)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.MultiScaleLossSchema, ~pydantic.types.Tag(tag=multiscale)], ~pydantic.types.Discriminator(discriminator=~anemoi.training.schemas.training._loss_discriminator, custom_error_type=None, custom_error_message=None, custom_error_context=None)], weight_averaging: ~anemoi.training.schemas.training.WeightAveragingSchema | None = None, loss_gradient_scaling: bool = False, scalers: dict[~typing.Literal['datasets'], dict[str, dict[str, ~anemoi.training.schemas.training.GeneralVariableLossScalerSchema | ~anemoi.training.schemas.training.VariableLevelScalerSchema | ~anemoi.training.schemas.training.VariableMaskingScalerSchema | ~anemoi.training.schemas.training.TendencyScalerSchema | ~anemoi.training.schemas.training.NaNMaskScalerSchema | ~anemoi.training.schemas.training.GraphNodeAttributeScalerSchema | ~anemoi.training.schemas.training.TimeStepScalerSchema | ~anemoi.training.schemas.training.UniformTimeStepScalerSchema | ~anemoi.training.schemas.training.LeadTimeDecayScalerSchema | ~anemoi.training.schemas.training.ReweightedGraphNodeAttributeScalerSchema | ~anemoi.training.schemas.training.SpectralDimensionScalerSchema]]] | dict[str, ~anemoi.training.schemas.training.GeneralVariableLossScalerSchema | ~anemoi.training.schemas.training.VariableLevelScalerSchema | ~anemoi.training.schemas.training.VariableMaskingScalerSchema | ~anemoi.training.schemas.training.TendencyScalerSchema | ~anemoi.training.schemas.training.NaNMaskScalerSchema | ~anemoi.training.schemas.training.GraphNodeAttributeScalerSchema | ~anemoi.training.schemas.training.TimeStepScalerSchema | ~anemoi.training.schemas.training.UniformTimeStepScalerSchema | ~anemoi.training.schemas.training.LeadTimeDecayScalerSchema | ~anemoi.training.schemas.training.ReweightedGraphNodeAttributeScalerSchema | ~anemoi.training.schemas.training.SpectralDimensionScalerSchema], validation_metrics: dict[~typing.Literal['datasets'], dict[str, dict[str, ~typing.Annotated[~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.BaseLossSchema, ~pydantic.types.Tag(tag=base)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.HuberLossSchema, ~pydantic.types.Tag(tag=huber)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CombinedLossSchema, ~pydantic.types.Tag(tag=combined)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CRPSSchema, ~pydantic.types.Tag(tag=crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.EnergyScoreLossSchema, ~pydantic.types.Tag(tag=energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphVariogramScoreLossSchema, ~pydantic.types.Tag(tag=graph_variogram_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeCRPSLossSchema, ~pydantic.types.Tag(tag=graph_edge_crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_edge_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.SpectralLossSchema, ~pydantic.types.Tag(tag=spectral)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.TimeAggregateLossWrapperSchema, ~pydantic.types.Tag(tag=time_aggregate)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.MultiScaleLossSchema, ~pydantic.types.Tag(tag=multiscale)], ~pydantic.types.Discriminator(discriminator=~anemoi.training.schemas.training._loss_discriminator, custom_error_type=None, custom_error_message=None, custom_error_context=None)]] | None]] | dict[str, ~typing.Annotated[~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.BaseLossSchema, ~pydantic.types.Tag(tag=base)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.HuberLossSchema, ~pydantic.types.Tag(tag=huber)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CombinedLossSchema, ~pydantic.types.Tag(tag=combined)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CRPSSchema, ~pydantic.types.Tag(tag=crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.EnergyScoreLossSchema, ~pydantic.types.Tag(tag=energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphVariogramScoreLossSchema, ~pydantic.types.Tag(tag=graph_variogram_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeCRPSLossSchema, ~pydantic.types.Tag(tag=graph_edge_crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_edge_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.SpectralLossSchema, ~pydantic.types.Tag(tag=spectral)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.TimeAggregateLossWrapperSchema, ~pydantic.types.Tag(tag=time_aggregate)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.MultiScaleLossSchema, ~pydantic.types.Tag(tag=multiscale)], ~pydantic.types.Discriminator(discriminator=~anemoi.training.schemas.training._loss_discriminator, custom_error_type=None, custom_error_message=None, custom_error_context=None)]] | None, variable_groups: dict[~typing.Literal['datasets'], dict[str, dict[str, str | list[str] | dict[str, str | bool | list[str | int]]] | None]] | dict[str, str | list[str] | dict[str, str | bool | list[str | int]]] | None, max_epochs: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] | None = None, max_steps: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 150000, optimization: ~anemoi.training.schemas.training.OptimizationSchema, metrics: dict[~typing.Literal['datasets'], dict[str, list[str]]] | list[str], ensemble_size_per_device: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 1, method: ~anemoi.training.schemas.training.EnsembleTrainingMethodSchema)

Bases: BaseTrainingSchema

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].

method: EnsembleTrainingMethodSchema: Training method.

class anemoi.training.schemas.training.TransportTrainingConfigSchema(*, prediction_mode: Literal['state', 'tendency'] = 'state', objective: Literal['edm_diffusion', 'stochastic_interpolant'] = 'edm_diffusion')

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].

prediction_mode: Literal['state', 'tendency']: Endpoint semantics for the transport objective.

objective: Literal['edm_diffusion', 'stochastic_interpolant']: Transport objective used to perturb targets and train the model.

class anemoi.training.schemas.training.TransportTrainingSchema(*, seed: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)] | None = None, check_variables_compatibility: ~anemoi.training.schemas.training.CheckVariablesCompatibilitySchema = <factory>, allow_variable_subset: bool = False, update_ds_stats_on_ckpt_load: ~anemoi.training.schemas.training.UpdateDsStatsOnCkptLoadSchema = <factory>, checkpoint: ~anemoi.training.schemas.training.CheckpointPipelineSchema | None = None, deterministic: bool = False, precision: str = '16-mixed', preferred_blas_backend: str | None = None, accum_grad_batches: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 1, num_sanity_val_steps: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], gradient_clip: ~anemoi.training.schemas.training.GradientClip, strategy: ~anemoi.training.schemas.training.BaseDDPStrategySchema | ~anemoi.training.schemas.training.DDPEnsGroupStrategyStrategySchema | ~anemoi.training.schemas.training.SingleDeviceStrategySchema, training_loss: dict[~typing.Literal['datasets'], dict[str, ~typing.Annotated[~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.BaseLossSchema, ~pydantic.types.Tag(tag=base)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.HuberLossSchema, ~pydantic.types.Tag(tag=huber)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CombinedLossSchema, ~pydantic.types.Tag(tag=combined)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CRPSSchema, ~pydantic.types.Tag(tag=crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.EnergyScoreLossSchema, ~pydantic.types.Tag(tag=energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphVariogramScoreLossSchema, ~pydantic.types.Tag(tag=graph_variogram_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeCRPSLossSchema, ~pydantic.types.Tag(tag=graph_edge_crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_edge_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.SpectralLossSchema, ~pydantic.types.Tag(tag=spectral)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.TimeAggregateLossWrapperSchema, ~pydantic.types.Tag(tag=time_aggregate)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.MultiScaleLossSchema, ~pydantic.types.Tag(tag=multiscale)], ~pydantic.types.Discriminator(discriminator=~anemoi.training.schemas.training._loss_discriminator, custom_error_type=None, custom_error_message=None, custom_error_context=None)]]] | ~typing.Annotated[~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.BaseLossSchema, ~pydantic.types.Tag(tag=base)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.HuberLossSchema, ~pydantic.types.Tag(tag=huber)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CombinedLossSchema, ~pydantic.types.Tag(tag=combined)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CRPSSchema, ~pydantic.types.Tag(tag=crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.EnergyScoreLossSchema, ~pydantic.types.Tag(tag=energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphVariogramScoreLossSchema, ~pydantic.types.Tag(tag=graph_variogram_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeCRPSLossSchema, ~pydantic.types.Tag(tag=graph_edge_crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_edge_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.SpectralLossSchema, ~pydantic.types.Tag(tag=spectral)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.TimeAggregateLossWrapperSchema, ~pydantic.types.Tag(tag=time_aggregate)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.MultiScaleLossSchema, ~pydantic.types.Tag(tag=multiscale)], ~pydantic.types.Discriminator(discriminator=~anemoi.training.schemas.training._loss_discriminator, custom_error_type=None, custom_error_message=None, custom_error_context=None)], weight_averaging: ~anemoi.training.schemas.training.WeightAveragingSchema | None = None, loss_gradient_scaling: bool = False, scalers: dict[~typing.Literal['datasets'], dict[str, dict[str, ~anemoi.training.schemas.training.GeneralVariableLossScalerSchema | ~anemoi.training.schemas.training.VariableLevelScalerSchema | ~anemoi.training.schemas.training.VariableMaskingScalerSchema | ~anemoi.training.schemas.training.TendencyScalerSchema | ~anemoi.training.schemas.training.NaNMaskScalerSchema | ~anemoi.training.schemas.training.GraphNodeAttributeScalerSchema | ~anemoi.training.schemas.training.TimeStepScalerSchema | ~anemoi.training.schemas.training.UniformTimeStepScalerSchema | ~anemoi.training.schemas.training.LeadTimeDecayScalerSchema | ~anemoi.training.schemas.training.ReweightedGraphNodeAttributeScalerSchema | ~anemoi.training.schemas.training.SpectralDimensionScalerSchema]]] | dict[str, ~anemoi.training.schemas.training.GeneralVariableLossScalerSchema | ~anemoi.training.schemas.training.VariableLevelScalerSchema | ~anemoi.training.schemas.training.VariableMaskingScalerSchema | ~anemoi.training.schemas.training.TendencyScalerSchema | ~anemoi.training.schemas.training.NaNMaskScalerSchema | ~anemoi.training.schemas.training.GraphNodeAttributeScalerSchema | ~anemoi.training.schemas.training.TimeStepScalerSchema | ~anemoi.training.schemas.training.UniformTimeStepScalerSchema | ~anemoi.training.schemas.training.LeadTimeDecayScalerSchema | ~anemoi.training.schemas.training.ReweightedGraphNodeAttributeScalerSchema | ~anemoi.training.schemas.training.SpectralDimensionScalerSchema], validation_metrics: dict[~typing.Literal['datasets'], dict[str, dict[str, ~typing.Annotated[~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.BaseLossSchema, ~pydantic.types.Tag(tag=base)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.HuberLossSchema, ~pydantic.types.Tag(tag=huber)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CombinedLossSchema, ~pydantic.types.Tag(tag=combined)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CRPSSchema, ~pydantic.types.Tag(tag=crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.EnergyScoreLossSchema, ~pydantic.types.Tag(tag=energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphVariogramScoreLossSchema, ~pydantic.types.Tag(tag=graph_variogram_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeCRPSLossSchema, ~pydantic.types.Tag(tag=graph_edge_crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_edge_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.SpectralLossSchema, ~pydantic.types.Tag(tag=spectral)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.TimeAggregateLossWrapperSchema, ~pydantic.types.Tag(tag=time_aggregate)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.MultiScaleLossSchema, ~pydantic.types.Tag(tag=multiscale)], ~pydantic.types.Discriminator(discriminator=~anemoi.training.schemas.training._loss_discriminator, custom_error_type=None, custom_error_message=None, custom_error_context=None)]] | None]] | dict[str, ~typing.Annotated[~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.BaseLossSchema, ~pydantic.types.Tag(tag=base)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.HuberLossSchema, ~pydantic.types.Tag(tag=huber)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CombinedLossSchema, ~pydantic.types.Tag(tag=combined)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.CRPSSchema, ~pydantic.types.Tag(tag=crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.EnergyScoreLossSchema, ~pydantic.types.Tag(tag=energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphVariogramScoreLossSchema, ~pydantic.types.Tag(tag=graph_variogram_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeCRPSLossSchema, ~pydantic.types.Tag(tag=graph_edge_crps)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.GraphEdgeEnergyScoreLossSchema, ~pydantic.types.Tag(tag=graph_edge_energy_score)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.SpectralLossSchema, ~pydantic.types.Tag(tag=spectral)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.TimeAggregateLossWrapperSchema, ~pydantic.types.Tag(tag=time_aggregate)] | ~anemoi.training.schemas.training.Annotated[~anemoi.training.schemas.training.MultiScaleLossSchema, ~pydantic.types.Tag(tag=multiscale)], ~pydantic.types.Discriminator(discriminator=~anemoi.training.schemas.training._loss_discriminator, custom_error_type=None, custom_error_message=None, custom_error_context=None)]] | None, variable_groups: dict[~typing.Literal['datasets'], dict[str, dict[str, str | list[str] | dict[str, str | bool | list[str | int]]] | None]] | dict[str, str | list[str] | dict[str, str | bool | list[str | int]]] | None, max_epochs: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] | None = None, max_steps: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 150000, optimization: ~anemoi.training.schemas.training.OptimizationSchema, metrics: dict[~typing.Literal['datasets'], dict[str, list[str]]] | list[str], ensemble_size_per_device: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] = 1, method: ~anemoi.training.schemas.training.TransportTrainingMethodSchema, transport: ~anemoi.training.schemas.training.TransportTrainingConfigSchema = <factory>)

Bases: BaseTrainingSchema

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].

method: TransportTrainingMethodSchema: Training method.

transport: TransportTrainingConfigSchema: Transport training configuration.