Schemas

This module defines pydantic schemas, which are used to validate the configuration.

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.

class anemoi.graphs.schemas.node_schemas.AnemoiDatasetNodeSchema(*, _target_: Literal['anemoi.graphs.nodes.AnemoiDatasetNodes', 'anemoi.graphs.nodes.ZarrDatasetNodes'], dataset: str | list | dict | None)

Bases: BaseModel

target_: Literal['anemoi.graphs.nodes.AnemoiDatasetNodes', 'anemoi.graphs.nodes.ZarrDatasetNodes']: Nodes from Anemoi dataset class implementation from anemoi.graphs.nodes.

dataset: str | list | dict | None: The dataset containing the nodes.

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.node_schemas.NPZnodeSchema(*, _target_: Literal['anemoi.graphs.nodes.NPZFileNodes'], npz_file: str, lat_key: str, lon_key: str)

Bases: BaseModel

target_: Literal['anemoi.graphs.nodes.NPZFileNodes']: Nodes from NPZ grids class implementation from anemoi.graphs.nodes.

npz_file: str: Path to the NPZ file.

lat_key: str: The key name of the latitude field.

lon_key: str: The key name of the longitude 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.graphs.schemas.node_schemas.TextNodeSchema(*, _target_: Literal['anemoi.graphs.nodes.TextNodes'], dataset: str | Path, idx_lon: int, idx_lat: int)

Bases: BaseModel

target_: Literal['anemoi.graphs.nodes.TextNodes']: Nodes from text file class implementation from anemoi.graphs.nodes.

dataset: str | Path: The path to text file containing the coordinates of the nodes.

idx_lon: int: The index of the longitude in the dataset.

idx_lat: int: The index of the latitude in 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.graphs.schemas.node_schemas.XArrayNodeSchema(*, _target_: Literal['anemoi.graphs.nodes.XArrayNodes'], dataset: str | Path, lon_key: str, lat_key: str)

Bases: BaseModel

target_: Literal['anemoi.graphs.nodes.XArrayNodes']: Nodes from xarray dataset class implementation from anemoi.graphs.nodes.

dataset: str | Path: The path to xarray dataset containing the coordinates of the nodes.

lon_key: str: The key name of the longitude field.

lat_key: str: The key name of the latitude 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.graphs.schemas.node_schemas.ReducedGaussianGridNodeSchema(*, _target_: Literal['anemoi.graphs.nodes.ReducedGaussianGridNodes'], grid: Literal['o16', 'o32', 'o48', 'o96', 'o160', 'o256', 'o320', 'n320', 'o1280'])

Bases: BaseModel

target_: Literal['anemoi.graphs.nodes.ReducedGaussianGridNodes']: Nodes from NPZ grids class implementation from anemoi.graphs.nodes.

grid: Literal['o16', 'o32', 'o48', 'o96', 'o160', 'o256', 'o320', 'n320', 'o1280']: Reduced gaussian grid.

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.node_schemas.ICONMeshNodeSchema(*, _target_: Literal['anemoi.graphs.nodes.ICONMultiMeshNodes', 'anemoi.graphs.nodes.ICONCellGridNodes'], grid_filename: str, max_level: int)

Bases: BaseModel

target_: Literal['anemoi.graphs.nodes.ICONMultiMeshNodes', 'anemoi.graphs.nodes.ICONCellGridNodes']: Mesh based on ICON grid class implementation from anemoi.graphs.nodes.

grid_filename: str: Name of NetCDF ICON grid file.

max_level: int: Maximum refinement level of the multi mesh / cell grid.

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.node_schemas.LimitedAreaNPZFileNodesSchema(*, _target_: Literal['anemoi.graphs.nodes.LimitedAreaNPZFileNodes'], grid_definition_path: str, resolution: str, reference_node_name: str, mask_attr_name: str | None, margin_radius_km: Annotated[float, Gt(gt=0)])

Bases: BaseModel

target_: Literal['anemoi.graphs.nodes.LimitedAreaNPZFileNodes']: Class implementation for nodes from NPZ grids using an area of interest from anemoi.graphs.nodes.

grid_definition_path: str: Path to the folder containing the grid definition files.

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

resolution: str: The grid resolution.

reference_node_name: str: Name of the reference nodes in the graph to consider for the Area Mask.

mask_attr_name: str | None: Name of a node to attribute to mask the reference nodes, if desired. Defaults to consider all reference nodes.

margin_radius_km: Annotated[float, Gt(gt=0)]: Maximum distance to the reference nodes to consider a node as valid, in kilometers. Defaults to 100 km.

class anemoi.graphs.schemas.node_schemas.IcosahedralandHealPixNodeSchema(*, _target_: Literal['anemoi.graphs.nodes.TriNodes', 'anemoi.graphs.nodes.HexNodes', 'anemoi.graphs.nodes.HEALPixNodes'], resolution: Annotated[int, Gt(gt=0)])

Bases: BaseModel

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

target_: Literal['anemoi.graphs.nodes.TriNodes', 'anemoi.graphs.nodes.HexNodes', 'anemoi.graphs.nodes.HEALPixNodes']: Icohedral and HEAL Pix nodes class implementations from anemoi.graphs.nodes.

resolution: Annotated[int, Gt(gt=0)]: Refinement level of the mesh.

class anemoi.graphs.schemas.node_schemas.LimitedAreaIcosahedralandHealPixNodeSchema(*, _target_: Literal['anemoi.graphs.nodes.LimitedAreaTriNodes', 'anemoi.graphs.nodes.LimitedAreaHexNodes', 'anemoi.graphs.nodes.LimitedAreaHEALPixNodes'], resolution: Annotated[int, Gt(gt=0)], reference_node_name: str, mask_attr_name: str | None, margin_radius_km: Annotated[float, Gt(gt=0)])

Bases: BaseModel

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

target_: Literal['anemoi.graphs.nodes.LimitedAreaTriNodes', 'anemoi.graphs.nodes.LimitedAreaHexNodes', 'anemoi.graphs.nodes.LimitedAreaHEALPixNodes']: Class implementations for Icosahedral and HEAL Pix nodes using an area of interest from anemoi.graphs.nodes.

resolution: Annotated[int, Gt(gt=0)]: Refinement level of the mesh.

reference_node_name: str: Name of the reference nodes in the graph to consider for the Area Mask.

mask_attr_name: str | None: Name of a node to attribute to mask the reference nodes, if desired. Defaults to consider all reference nodes.

margin_radius_km: Annotated[float, Gt(gt=0)]: Maximum distance to the reference nodes to consider a node as valid, in kilometers. Defaults to 100 km.

class anemoi.graphs.schemas.node_schemas.StretchedIcosahdralNodeSchema(*, _target_: Literal['anemoi.graphs.nodes.StretchedTriNodes'], global_resolution: Annotated[int, Gt(gt=0)], lam_resolution: Annotated[int, Gt(gt=0)], reference_node_name: str, mask_attr_name: str | None, margin_radius_km: Annotated[float, Gt(gt=0)])

Bases: BaseModel

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

target_: Literal['anemoi.graphs.nodes.StretchedTriNodes']: Class implementation for nodes based on iterative refinements of an icosahedron with 2 different resolutions.

global_resolution: Annotated[int, Gt(gt=0)]: Refinement level of the mesh on the global area.

lam_resolution: Annotated[int, Gt(gt=0)]: Refinement level of the mesh on the local area.

reference_node_name: str: Name of the reference nodes in the graph to consider for the Area Mask.

mask_attr_name: str | None: Name of a node to attribute to mask the reference nodes, if desired. Defaults to consider all reference nodes.

margin_radius_km: Annotated[float, Gt(gt=0)]: Maximum distance to the reference nodes to consider a node as valid, in kilometers. Defaults to 100 km.

class anemoi.graphs.schemas.edge_schemas.KNNEdgeSchema(*, _target_: Literal['anemoi.graphs.edges.KNNEdges', 'anemoi.graphs.edges.ReversedKNNEdges'], num_nearest_neighbours: Annotated[int, Gt(gt=0)], source_mask_attr_name: str | None = None, target_mask_attr_name: str | None = None)

Bases: BaseModel

target_: Literal['anemoi.graphs.edges.KNNEdges', 'anemoi.graphs.edges.ReversedKNNEdges']: KNN based edges implementation from anemoi.graphs.edges.

num_nearest_neighbours: Annotated[int, Gt(gt=0)]: Number of nearest neighbours. Default to 3.

source_mask_attr_name: str | None: Mask to apply to source nodes of the edges. Default to None.

target_mask_attr_name: str | None: Mask to apply to target nodes of the edges. Default to None.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.graphs.schemas.edge_schemas.MutualKNNEdgeSchema(*, _target_: Literal['anemoi.graphs.edges.MutualKNNEdges'], num_nearest_neighbours: Annotated[int, Gt(gt=0)], reversed_num_nearest_neighbours: Annotated[int, Gt(gt=0)] | None = None, source_mask_attr_name: str | None = None, target_mask_attr_name: str | None = None)

Bases: BaseModel

target_: Literal['anemoi.graphs.edges.MutualKNNEdges']: Mutual KNN based edges implementation from anemoi.graphs.edges.

num_nearest_neighbours: Annotated[int, Gt(gt=0)]: Number of nearest source nodes considered for each target node.

reversed_num_nearest_neighbours: Annotated[int, Gt(gt=0)] | None: Number of nearest target nodes considered for each source node. Defaults to num_nearest_neighbours.

source_mask_attr_name: str | None: Mask to apply to source nodes of the edges. Default to None.

target_mask_attr_name: str | None: Mask to apply to target nodes of the edges. Default to None.

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class anemoi.graphs.schemas.edge_schemas.CutoffEdgeSchema(*, _target_: Literal['anemoi.graphs.edges.CutOffEdges', 'anemoi.graphs.edges.ReversedCutOffEdges'], cutoff_factor: Annotated[float, Gt(gt=0)] | None = None, cutoff_distance_km: Annotated[float, Gt(gt=0)] | None = None, source_mask_attr_name: str | None = None, target_mask_attr_name: str | None = None, max_num_neighbours: Annotated[int, Gt(gt=0)] = 64)

Bases: BaseModel

target_: Literal['anemoi.graphs.edges.CutOffEdges', 'anemoi.graphs.edges.ReversedCutOffEdges']: Cut-off based edges implementation from anemoi.graphs.edges.

cutoff_factor: Annotated[float, Gt(gt=0)] | None: Factor to multiply the grid reference distance to get the cut-off radius. Mutually exclusive with cutoff_distance_km.

cutoff_distance_km: Annotated[float, Gt(gt=0)] | None: Cutoff radius in kilometers. Mutually exclusive with cutoff_factor.

source_mask_attr_name: str | None: Mask to apply to source nodes of the edges. Default to None.

target_mask_attr_name: str | None: Mask to apply to target nodes of the edges. Default to None.

max_num_neighbours: Annotated[int, Gt(gt=0)]: Maximum number of nearest neighbours to consider when building edges. Default to 64.

validate_cutoff_params(): Validate that exactly one of cutoff_factor or cutoff_distance_km is provided.

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.edge_schemas.MultiScaleEdgeSchema(*, _target_: Literal['anemoi.graphs.edges.MultiScaleEdges'] = 'anemoi.graphs.edges.MultiScaleEdges', x_hops: Annotated[int, Gt(gt=0)], scale_resolutions: Annotated[int, Gt(gt=0)] | list[Annotated[int, Gt(gt=0)]] | None, source_mask_attr_name: str | None = None, target_mask_attr_name: str | None = None)

Bases: BaseModel

target_: Literal['anemoi.graphs.edges.MultiScaleEdges']: Multi-scale edges implementation from anemoi.graphs.edges.

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

x_hops: Annotated[int, Gt(gt=0)]: Number of hops (in the refined icosahedron) between two nodes to connect them with an edge. Default to 1.

scale_resolutions: Annotated[int, Gt(gt=0)] | list[Annotated[int, Gt(gt=0)]] | None: Specifies the resolution scales for computing the hop neighbourhood.

source_mask_attr_name: str | None: Mask to apply to source nodes of the edges. Default to None.

target_mask_attr_name: str | None: Mask to apply to target nodes of the edges. Default to None.

class anemoi.graphs.schemas.edge_schemas.HEALPixMultiScaleEdgesSchema(*, _target_: Literal['anemoi.graphs.edges.HEALPixMultiScaleEdges'], scale_resolutions: Annotated[int, Gt(gt=0)] | list[Annotated[int, Gt(gt=0)]] | None, source_mask_attr_name: str | None = None, target_mask_attr_name: str | None = None)

Bases: BaseModel

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

target_: Literal['anemoi.graphs.edges.HEALPixMultiScaleEdges']: HEALPix multi-scale edges implementation from anemoi.graphs.edges.

scale_resolutions: Annotated[int, Gt(gt=0)] | list[Annotated[int, Gt(gt=0)]] | None: Specifies the resolution scales for computing the hop neighbourhood.

source_mask_attr_name: str | None: Mask to apply to source nodes of the edges. Default to None.

target_mask_attr_name: str | None: Mask to apply to target nodes of the edges. Default to None.

class anemoi.graphs.schemas.edge_schemas.ICONTopologicalEdgeSchema(*, _target_: Literal['anemoi.graphs.edges.ICONTopologicalProcessorEdges', 'anemoi.graphs.edges.ICONTopologicalEncoderEdges', 'anemoi.graphs.edges.ICONTopologicalDecoderEdges'] = 'anemoi.graphs.edges.ICONTopologicalProcessorEdges', source_mask_attr_name: str | None = None, target_mask_attr_name: str | None = None)

Bases: BaseModel

model_config = {'extra': 'forbid', 'use_attribute_docstrings': True, 'use_enum_values': True, 'validate_assignment': True, 'validate_default': True}: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

source_mask_attr_name: str | None: Mask to apply to source nodes of the edges. Default to None.

target_mask_attr_name: str | None: Mask to apply to target nodes of the edges. Default to None.

class anemoi.graphs.schemas.edge_schemas.EdgeAttributeSchema(*, _target_: Literal['anemoi.graphs.edges.attributes.EdgeLength', 'anemoi.graphs.edges.attributes.EdgeDirection'] = 'anemoi.graphs.edges.attributes.EdgeLength', norm: Literal['unit-max', 'l1', 'l2', 'unit-sum', 'unit-std'])

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.graphs.edges.attributes.EdgeLength', 'anemoi.graphs.edges.attributes.EdgeDirection']: Edge attributes object from anemoi.graphs.edges.

norm: Literal['unit-max', 'l1', 'l2', 'unit-sum', 'unit-std']: Normalisation method applied to the edge attribute.