Schemas

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.

class anemoi.models.schemas.processor.NoOpProcessorSchema(*, _target_: Literal['anemoi.models.layers.processor.NoOpProcessor'])

Bases: BaseModel

target_: Literal['anemoi.models.layers.processor.NoOpProcessor']: No-op processor, used for ablations.

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

class anemoi.models.schemas.processor.GNNProcessorSchema(*, _convert_: str = 'all', trainable_size: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], num_chunks: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], cpu_offload: bool, gradient_checkpointing: bool = True, sub_graph_edge_attributes: list[str] = <factory>, mlp_extra_layers: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], mlp_hidden_ratio: ~typing.Annotated[float, ~annotated_types.Gt(gt=0)] = 1.0, mlp_implementation: ~typing.Literal['mlp', 'glu', 'swiglu', 'geglu', 'reglu'] = 'mlp', layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.processor.GNNProcessor'], num_channels: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], num_layers: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)])

Bases: GNNModelComponent

target_: Literal['anemoi.models.layers.processor.GNNProcessor']: GNN Processor object from anemoi.models.layers.processor.

num_channels: Annotated[int, Ge(ge=0)]: Number of channels in the GNN processor. Default to 512.

num_layers: Annotated[int, Ge(ge=0)]: Number of layers of GNN processor. Default to 16.

num_chunks: Annotated[int, Ge(ge=0)]: Number of chunks to divide the layer into. Default to 2.

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.processor.GraphTransformerProcessorSchema(*, _convert_: str = 'all', cpu_offload: bool, gradient_checkpointing: bool = True, num_chunks: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], mlp_hidden_ratio: ~typing.Annotated[float, ~annotated_types.Gt(gt=0)], mlp_implementation: ~typing.Literal['mlp', 'glu', 'swiglu', 'geglu', 'reglu'] = 'mlp', num_heads: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], attn_channels: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] | None = None, layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.processor.GraphTransformerProcessor'], num_channels: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], trainable_size: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], sub_graph_edge_attributes: list[str], num_layers: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], qk_norm: bool, **extra_data: ~typing.Any)

Bases: TransformerModelComponent

target_: Literal['anemoi.models.layers.processor.GraphTransformerProcessor']: Graph transformer processor object from anemoi.models.layers.processor.

num_channels: Annotated[int, Ge(ge=0)]: Number of channels in the Graph Transformer processor. Default to 512.

trainable_size: Annotated[int, Ge(ge=0)]: Size of trainable parameters vector. Default to 8.

sub_graph_edge_attributes: list[str]: Edge attributes to consider in the processor features. Default [edge_length, endge_dirs].

num_layers: Annotated[int, Ge(ge=0)]: Number of layers of Graph Transformer processor. Default to 16.

num_chunks: Annotated[int, Ge(ge=0)]: Number of chunks to divide the layer into. Default to 2.

qk_norm: bool: Normalize the query and key vectors. Default to False.

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.models.schemas.processor.TransformerProcessorSchema(*, _convert_: str = 'all', cpu_offload: bool, gradient_checkpointing: bool = True, num_chunks: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], mlp_hidden_ratio: ~typing.Annotated[float, ~annotated_types.Gt(gt=0)], mlp_implementation: ~typing.Literal['mlp', 'glu', 'swiglu', 'geglu', 'reglu'] = 'mlp', num_heads: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], attn_channels: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] | None = None, layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.processor.TransformerProcessor'], num_channels: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], num_layers: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], window_size: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)] | None, dropout_p: ~typing.Annotated[float, ~annotated_types.Ge(ge=0)], attention_implementation: str, qk_norm: bool, softcap: ~typing.Annotated[float, ~annotated_types.Ge(ge=0)], use_alibi_slopes: bool, **extra_data: ~typing.Any)

Bases: TransformerModelComponent

target_: Literal['anemoi.models.layers.processor.TransformerProcessor']: Transformer processor object from anemoi.models.layers.processor.

num_channels: Annotated[int, Ge(ge=0)]: Number of channels in the Transformer processor. Default to 512.

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

num_layers: Annotated[int, Ge(ge=0)]: Number of layers of Transformer processor. Default to 16.

num_chunks: Annotated[int, Ge(ge=0)]: Number of chunks to divide the layer into. Default to 2.

window_size: Annotated[int, Ge(ge=0)] | None: Attention window size along the longitude axis. Default to 512.

dropout_p: Annotated[float, Ge(ge=0)]: Dropout probability used for multi-head self attention, default 0.0

attention_implementation: str: Attention implementation to use. Default to ‘flash_attention’.

qk_norm: bool: Normalize the query and key vectors. Default to False.

softcap: Annotated[float, Ge(ge=0)]: Softcap value for attention. Default to 0.0.

use_alibi_slopes: bool: Use alibi slopes for attention implementation. Default to False.

class anemoi.models.schemas.processor.PointWiseMLPProcessorSchema(*, _convert_: str = 'all', num_chunks: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], cpu_offload: bool, gradient_checkpointing: bool = True, layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.processor.PointWiseMLPProcessor'], num_channels: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], num_layers: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], mlp_hidden_ratio: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], dropout_p: ~typing.Annotated[float, ~annotated_types.Ge(ge=0)] = 0.0)

Bases: PointWiseModelComponent

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.processor.PointWiseMLPProcessor']: PointWise MLP processor object from anemoi.models.layers.processor.

num_channels: Annotated[int, Ge(ge=0)]: Number of channels in the PointWise MLP processor. Default to 512.

num_layers: Annotated[int, Ge(ge=0)]: Number of layers of the PointWise MLP processor.

mlp_hidden_ratio: Annotated[int, Ge(ge=0)]: Ratio of the hidden dimension to the processor channel dimension.

dropout_p: Annotated[float, Ge(ge=0)]: Dropout probability, default 0.0

class anemoi.models.schemas.encoder.GNNEncoderSchema(*, _convert_: str = 'all', trainable_size: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)] = 0, num_chunks: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], cpu_offload: bool, gradient_checkpointing: bool = True, sub_graph_edge_attributes: list[str] = <factory>, mlp_extra_layers: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], mlp_hidden_ratio: ~typing.Annotated[float, ~annotated_types.Gt(gt=0)] = 1.0, mlp_implementation: ~typing.Literal['mlp', 'glu', 'swiglu', 'geglu', 'reglu'] = 'mlp', layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.mapper.GNNForwardMapper'], num_channels: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)])

Bases: GNNModelComponent

target_: Literal['anemoi.models.layers.mapper.GNNForwardMapper']: GNN encoder object from anemoi.models.layers.mapper.

num_channels: Annotated[int, Ge(ge=0)]: Hidden dimension of the GNN encoder. Default to 512.

trainable_size: Annotated[int, Ge(ge=0)]: Size of trainable parameters vector. Default to 0.

sub_graph_edge_attributes: list[str]: Edge attributes to consider in the model component features.

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.encoder.GraphTransformerEncoderSchema(*, _convert_: str = 'all', cpu_offload: bool, gradient_checkpointing: bool = True, num_chunks: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], mlp_hidden_ratio: ~typing.Annotated[float, ~annotated_types.Gt(gt=0)], mlp_implementation: ~typing.Literal['mlp', 'glu', 'swiglu', 'geglu', 'reglu'] = 'mlp', num_heads: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], attn_channels: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] | None = None, layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.mapper.GraphTransformerForwardMapper'], num_channels: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], trainable_size: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)] = 0, sub_graph_edge_attributes: list[str] = <factory>, qk_norm: bool, **extra_data: ~typing.Any)

Bases: TransformerModelComponent

target_: Literal['anemoi.models.layers.mapper.GraphTransformerForwardMapper']: Graph Transfromer Encoder object from anemoi.models.layers.mapper.

num_channels: Annotated[int, Ge(ge=0)]: Hidden dimension of the Graph Transformer encoder. Default to 512.

trainable_size: Annotated[int, Ge(ge=0)]: Size of trainable parameters vector. Default to 0.

sub_graph_edge_attributes: list[str]: Edge attributes to consider in the model component features.

qk_norm: bool: Normalize the query and key vectors. Default to False.

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.models.schemas.encoder.TransformerEncoderSchema(*, _convert_: str = 'all', cpu_offload: bool, gradient_checkpointing: bool = True, num_chunks: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], mlp_hidden_ratio: ~typing.Annotated[float, ~annotated_types.Gt(gt=0)], mlp_implementation: ~typing.Literal['mlp', 'glu', 'swiglu', 'geglu', 'reglu'] = 'mlp', num_heads: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], attn_channels: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] | None = None, layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.mapper.TransformerForwardMapper'], num_channels: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], window_size: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)] | None, dropout_p: ~typing.Annotated[float, ~annotated_types.Ge(ge=0)], attention_implementation: str, softcap: ~typing.Annotated[float, ~annotated_types.Ge(ge=0)], use_alibi_slopes: bool, **extra_data: ~typing.Any)

Bases: TransformerModelComponent

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

target_: Literal['anemoi.models.layers.mapper.TransformerForwardMapper']: Transformer Encoder object from anemoi.models.layers.mapper.

num_channels: Annotated[int, Ge(ge=0)]: Hidden dimension of the Transformer encoder. Default to 512.

window_size: Annotated[int, Ge(ge=0)] | None: Attention window size along the longitude axis. Default to 512.

dropout_p: Annotated[float, Ge(ge=0)]: Dropout probability used for multi-head self attention, default 0.0

attention_implementation: str: Attention implementation to use. Default to ‘flash_attention’.

softcap: Annotated[float, Ge(ge=0)]: Softcap value for attention. Default to 0.0.

use_alibi_slopes: bool: Use alibi slopes for attention implementation. Default to False.

class anemoi.models.schemas.encoder.PointWiseForwardMapperSchema(*, _convert_: str = 'all', cpu_offload: bool, gradient_checkpointing: bool = True, layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.mapper.PointWiseForwardMapper'], num_channels: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)])

Bases: PointWiseMapperComponent

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.mapper.PointWiseForwardMapper']: Point-wise encoder object from anemoi.models.layers.mapper.

num_channels: Annotated[int, Ge(ge=0)]: Hidden dimension of the Point-wise encoder. Default to 512.

class anemoi.models.schemas.decoder.GNNDecoderSchema(*, _convert_: str = 'all', trainable_size: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)] = 0, num_chunks: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], cpu_offload: bool, gradient_checkpointing: bool = True, sub_graph_edge_attributes: list[str] = <factory>, mlp_extra_layers: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], mlp_hidden_ratio: ~typing.Annotated[float, ~annotated_types.Gt(gt=0)] = 1.0, mlp_implementation: ~typing.Literal['mlp', 'glu', 'swiglu', 'geglu', 'reglu'] = 'mlp', layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.mapper.GNNBackwardMapper'], num_channels: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)])

Bases: GNNModelComponent

target_: Literal['anemoi.models.layers.mapper.GNNBackwardMapper']: GNN decoder object from anemoi.models.layers.mapper.

num_channels: Annotated[int, Ge(ge=0)]: Hidden dimension of the GNN decoder. Default to 512.

trainable_size: Annotated[int, Ge(ge=0)]: Size of trainable parameters vector. Default to 0.

sub_graph_edge_attributes: list[str]: Edge attributes to consider in the model component features.

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.decoder.GraphTransformerDecoderSchema(*, _convert_: str = 'all', cpu_offload: bool, gradient_checkpointing: bool = True, num_chunks: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], mlp_hidden_ratio: ~typing.Annotated[float, ~annotated_types.Gt(gt=0)], mlp_implementation: ~typing.Literal['mlp', 'glu', 'swiglu', 'geglu', 'reglu'] = 'mlp', num_heads: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], attn_channels: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] | None = None, layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.mapper.GraphTransformerBackwardMapper'], num_channels: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], trainable_size: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)] = 0, sub_graph_edge_attributes: list[str] = <factory>, qk_norm: bool, initialise_data_extractor_zero: bool, **extra_data: ~typing.Any)

Bases: TransformerModelComponent

target_: Literal['anemoi.models.layers.mapper.GraphTransformerBackwardMapper']: Graph Transformer Decoder object from anemoi.models.layers.mapper.

num_channels: Annotated[int, Ge(ge=0)]: Hidden dimension of the Graph Transformer decoder. Default to 512.

trainable_size: Annotated[int, Ge(ge=0)]: Size of trainable parameters vector. Default to 0.

sub_graph_edge_attributes: list[str]: Edge attributes to consider in the model component features.

qk_norm: bool: Normalize the query and key vectors. Default to False.

initialise_data_extractor_zero: bool: Initialise the data extractor with zeros. Default to False.

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.models.schemas.decoder.TransformerDecoderSchema(*, _convert_: str = 'all', cpu_offload: bool, gradient_checkpointing: bool = True, num_chunks: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], mlp_hidden_ratio: ~typing.Annotated[float, ~annotated_types.Gt(gt=0)], mlp_implementation: ~typing.Literal['mlp', 'glu', 'swiglu', 'geglu', 'reglu'] = 'mlp', num_heads: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], attn_channels: ~typing.Annotated[int, ~annotated_types.Gt(gt=0)] | None = None, layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.mapper.TransformerBackwardMapper'], num_channels: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], window_size: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)] | None, dropout_p: ~typing.Annotated[float, ~annotated_types.Ge(ge=0)], attention_implementation: str, softcap: ~typing.Annotated[float, ~annotated_types.Ge(ge=0)], use_alibi_slopes: bool, use_rotary_embeddings: bool, **extra_data: ~typing.Any)

Bases: TransformerModelComponent

target_: Literal['anemoi.models.layers.mapper.TransformerBackwardMapper']: Transformer Decoder object from anemoi.models.layers.mapper.

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

num_channels: Annotated[int, Ge(ge=0)]: Hidden dimension of the Transformer decoder. Default to 512.

window_size: Annotated[int, Ge(ge=0)] | None: Attention window size along the longitude axis. Default to 512.

dropout_p: Annotated[float, Ge(ge=0)]: Dropout probability used for multi-head self attention, default 0.0

attention_implementation: str: Attention implementation to use. Default to ‘flash_attention’.

softcap: Annotated[float, Ge(ge=0)]: Softcap value for attention. Default to 0.0.

use_alibi_slopes: bool: Use alibi slopes for attention implementation. Default to False.

use_rotary_embeddings: bool: Use rotary embeddings for attention implementation. Default to False.

class anemoi.models.schemas.decoder.PointWiseBackwardMapperSchema(*, _convert_: str = 'all', cpu_offload: bool, gradient_checkpointing: bool = True, layer_kernels: dict[str, dict] | None = <factory>, _target_: ~typing.Literal['anemoi.models.layers.mapper.PointWiseBackwardMapper'], num_channels: ~typing.Annotated[int, ~annotated_types.Ge(ge=0)], initialise_data_extractor_zero: bool = False)

Bases: PointWiseMapperComponent

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.mapper.PointWiseBackwardMapper']: Point-wise decoder object from anemoi.models.layers.mapper.

num_channels: Annotated[int, Ge(ge=0)]: Hidden dimension of the Point-wise decoder. Default to 512.

initialise_data_extractor_zero: bool: Initialise the data extractor with zeros. Default to False.