anemoi-graphs 0.2.0__py3-none-any.whl

anemoi/graphs/__init__.py

@@ -0,0 +1,10 @@
+# (C) Copyright 2023 European Centre for Medium-Range Weather Forecasts.
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+from ._version import __version__
+
+EARTH_RADIUS = 6371.0  # km

anemoi/graphs/__main__.py

@@ -0,0 +1,28 @@
+#!/usr/bin/env python
+# (C) Copyright 2024 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+
+from anemoi.utils.cli import cli_main
+from anemoi.utils.cli import make_parser
+
+from . import __version__
+from .commands import COMMANDS
+
+
+# For read-the-docs
+def create_parser():
+    return make_parser(__doc__, COMMANDS)
+
+
+def main():
+    cli_main(__version__, __doc__, COMMANDS)
+
+
+if __name__ == "__main__":
+    main()

anemoi/graphs/_version.py

@@ -0,0 +1,16 @@
+# file generated by setuptools_scm
+# don't change, don't track in version control
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple, Union
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.2.0'
+__version_tuple__ = version_tuple = (0, 2, 0)

anemoi/graphs/commands/__init__.py

@@ -0,0 +1,24 @@
+#!/usr/bin/env python
+# (C) Copyright 2024 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+#
+
+import os
+
+from anemoi.utils.cli import Command
+from anemoi.utils.cli import Failed
+from anemoi.utils.cli import register_commands
+
+__all__ = ["Command"]
+
+COMMANDS = register_commands(
+    os.path.dirname(__file__),
+    __name__,
+    lambda x: x.command(),
+    lambda name, error: Failed(name, error),
+)

anemoi/graphs/commands/create.py

@@ -0,0 +1,28 @@
+from anemoi.graphs.create import GraphCreator
+
+from . import Command
+
+
+class Create(Command):
+    """Create a graph."""
+
+    internal = True
+    timestamp = True
+
+    def add_arguments(self, command_parser):
+        command_parser.add_argument(
+            "--overwrite",
+            action="store_true",
+            help="Overwrite existing files. This will delete the target graph if it already exists.",
+        )
+        command_parser.add_argument("config", help="Configuration yaml file defining the recipe to create the graph.")
+        command_parser.add_argument("path", help="Path to store the created graph.")
+
+    def run(self, args):
+        kwargs = vars(args)
+
+        c = GraphCreator(**kwargs)
+        c.create()
+
+
+command = Create

anemoi/graphs/create.py

@@ -0,0 +1,100 @@
+import logging
+import os
+
+import torch
+from anemoi.utils.config import DotDict
+from hydra.utils import instantiate
+from torch_geometric.data import HeteroData
+
+LOGGER = logging.getLogger(__name__)
+
+
+class GraphCreator:
+    """Graph creator."""
+
+    def __init__(
+        self,
+        path,
+        config=None,
+        cache=None,
+        print=print,
+        overwrite=False,
+        **kwargs,
+    ):
+        if isinstance(config, str) or isinstance(config, os.PathLike):
+            self.config = DotDict.from_file(config)
+        else:
+            self.config = config
+
+        self.path = path  # Output path
+        self.cache = cache
+        self.print = print
+        self.overwrite = overwrite
+
+    def init(self):
+        if self._path_readable() and not self.overwrite:
+            raise Exception(f"{self.path} already exists. Use overwrite=True to overwrite.")
+
+    def generate_graph(self) -> HeteroData:
+        """Generate the graph.
+
+        It instantiates the node builders and edge builders defined in the configuration
+        file and applies them to the graph.
+
+        Returns
+        -------
+            HeteroData: The generated graph.
+        """
+        graph = HeteroData()
+
+        for nodes_name, nodes_cfg in self.config.nodes.items():
+            graph = instantiate(nodes_cfg.node_builder, name=nodes_name).update_graph(
+                graph, nodes_cfg.get("attributes", {})
+            )
+
+        for edges_cfg in self.config.edges:
+            graph = instantiate(edges_cfg.edge_builder, edges_cfg.source_name, edges_cfg.target_name).update_graph(
+                graph, edges_cfg.get("attributes", {})
+            )
+
+        return graph
+
+    def clean(self, graph: HeteroData) -> HeteroData:
+        """Clean the hidden attributes of the nodes and edges."""
+        for nodes_name in graph.node_types:
+            node_attrs = list(graph[nodes_name].keys())
+            for node_attr_name in node_attrs:
+                if node_attr_name.startswith("_"):
+                    del graph[nodes_name][node_attr_name]
+
+        for edge_key in graph.edge_types:
+            edge_attrs = graph[edge_key].keys()
+            for edge_attr_name in edge_attrs:
+                if edge_attr_name.startswith("_"):
+                    del graph[edge_key][edge_attr_name]
+
+        return graph
+
+    def save(self, graph: HeteroData) -> None:
+        """Save the graph to the output path."""
+        if not os.path.exists(self.path) or self.overwrite:
+            torch.save(graph, self.path)
+            self.print(f"Graph saved at {self.path}.")
+
+    def create(self) -> HeteroData:
+        """Create the graph and save it to the output path."""
+        self.init()
+        graph = self.generate_graph()
+        graph = self.clean(graph)
+        self.save(graph)
+        return graph
+
+    def _path_readable(self) -> bool:
+        """Check if the output path is readable."""
+        import torch
+
+        try:
+            torch.load(self.path)
+            return True
+        except FileNotFoundError:
+            return False

anemoi/graphs/edges/__init__.py

@@ -0,0 +1,5 @@
+from .builder import CutOffEdges
+from .builder import KNNEdges
+from .builder import MultiScaleEdges
+
+__all__ = ["KNNEdges", "CutOffEdges", "MultiScaleEdges"]

anemoi/graphs/edges/attributes.py

@@ -0,0 +1,149 @@
+import logging
+from abc import ABC
+from abc import abstractmethod
+from typing import Optional
+
+import numpy as np
+import torch
+from torch_geometric.data import HeteroData
+
+from anemoi.graphs.edges.directional import directional_edge_features
+from anemoi.graphs.normalizer import NormalizerMixin
+from anemoi.graphs.utils import haversine_distance
+
+LOGGER = logging.getLogger(__name__)
+
+
+class BaseEdgeAttribute(ABC, NormalizerMixin):
+    """Base class for edge attributes."""
+
+    def __init__(self, norm: Optional[str] = None) -> None:
+        self.norm = norm
+
+    @abstractmethod
+    def get_raw_values(self, graph: HeteroData, source_name: str, target_name: str, *args, **kwargs) -> np.ndarray: ...
+
+    def post_process(self, values: np.ndarray) -> torch.Tensor:
+        """Post-process the values."""
+        if values.ndim == 1:
+            values = values[:, np.newaxis]
+
+        normed_values = self.normalize(values)
+
+        return torch.tensor(normed_values, dtype=torch.float32)
+
+    def compute(self, graph: HeteroData, edges_name: tuple[str, str, str], *args, **kwargs) -> torch.Tensor:
+        """Compute the edge attributes."""
+        source_name, _, target_name = edges_name
+        assert (
+            source_name in graph.node_types
+        ), f"Node \"{source_name}\" not found in graph. Optional nodes are {', '.join(graph.node_types)}."
+        assert (
+            target_name in graph.node_types
+        ), f"Node \"{target_name}\" not found in graph. Optional nodes are {', '.join(graph.node_types)}."
+
+        values = self.get_raw_values(graph, source_name, target_name, *args, **kwargs)
+        return self.post_process(values)
+
+
+class EdgeDirection(BaseEdgeAttribute):
+    """Edge direction feature.
+
+    If using the rotated features, the direction of the edge is computed
+    rotating the target nodes to the north pole. If not, it is computed
+    as the diference in latitude and longitude between the source and
+    target nodes.
+
+    Attributes
+    ----------
+    norm : Optional[str]
+        Normalization method.
+    luse_rotated_features : bool
+        Whether to use rotated features.
+
+    Methods
+    -------
+    get_raw_values(graph, source_name, target_name)
+        Compute directions between nodes connected by edges.
+    compute(graph, source_name, target_name)
+        Compute directional attributes.
+    """
+
+    def __init__(self, norm: Optional[str] = None, luse_rotated_features: bool = True) -> None:
+        super().__init__(norm)
+        self.luse_rotated_features = luse_rotated_features
+
+    def get_raw_values(self, graph: HeteroData, source_name: str, target_name: str) -> np.ndarray:
+        """Compute directional features for edges.
+
+        Parameters
+        ----------
+        graph : HeteroData
+            The graph.
+        source_name : str
+            The name of the source nodes.
+        target_name : str
+            The name of the target nodes.
+
+        Returns
+        -------
+        np.ndarray
+            The directional features.
+        """
+        edge_index = graph[(source_name, "to", target_name)].edge_index
+        source_coords = graph[source_name].x.numpy()[edge_index[0]].T
+        target_coords = graph[target_name].x.numpy()[edge_index[1]].T
+        edge_dirs = directional_edge_features(source_coords, target_coords, self.luse_rotated_features).T
+        return edge_dirs
+
+
+class EdgeLength(BaseEdgeAttribute):
+    """Edge length feature.
+
+    Attributes
+    ----------
+    norm : str
+        Normalization method.
+    invert : bool
+        Whether to invert the edge lengths, i.e. 1 - edge_length.
+
+    Methods
+    -------
+    get_raw_values(graph, source_name, target_name)
+        Compute haversine distance between nodes connected by edges.
+    compute(graph, source_name, target_name)
+        Compute edge lengths attributes.
+    """
+
+    def __init__(self, norm: Optional[str] = None, invert: bool = False) -> None:
+        super().__init__(norm)
+        self.invert = invert
+
+    def get_raw_values(self, graph: HeteroData, source_name: str, target_name: str) -> np.ndarray:
+        """Compute haversine distance (in kilometers) between nodes connected by edges.
+
+        Parameters
+        ----------
+        graph : HeteroData
+            The graph.
+        source_name : str
+            The name of the source nodes.
+        target_name : str
+            The name of the target nodes.
+
+        Returns
+        -------
+        np.ndarray
+            The edge lengths.
+        """
+        edge_index = graph[(source_name, "to", target_name)].edge_index
+        source_coords = graph[source_name].x.numpy()[edge_index[0]]
+        target_coords = graph[target_name].x.numpy()[edge_index[1]]
+        edge_lengths = haversine_distance(source_coords, target_coords)
+        return edge_lengths
+
+    def post_process(self, values: np.ndarray) -> torch.Tensor:
+        """Post-process edge lengths."""
+        if self.invert:
+            values = 1 - values
+        return super().post_process(values)

anemoi/graphs/edges/builder.py

@@ -0,0 +1,328 @@
+import logging
+from abc import ABC
+from abc import abstractmethod
+from typing import Optional
+
+import networkx as nx
+import numpy as np
+import torch
+from anemoi.utils.config import DotDict
+from hydra.utils import instantiate
+from sklearn.neighbors import NearestNeighbors
+from torch_geometric.data import HeteroData
+from torch_geometric.data.storage import NodeStorage
+
+from anemoi.graphs import EARTH_RADIUS
+from anemoi.graphs.generate import hexagonal
+from anemoi.graphs.generate import icosahedral
+from anemoi.graphs.nodes.builder import HexNodes
+from anemoi.graphs.nodes.builder import TriNodes
+from anemoi.graphs.utils import get_grid_reference_distance
+
+LOGGER = logging.getLogger(__name__)
+
+
+class BaseEdgeBuilder(ABC):
+    """Base class for edge builders."""
+
+    def __init__(self, source_name: str, target_name: str):
+        self.source_name = source_name
+        self.target_name = target_name
+
+    @property
+    def name(self) -> tuple[str, str, str]:
+        """Name of the edge subgraph."""
+        return self.source_name, "to", self.target_name
+
+    @abstractmethod
+    def get_adjacency_matrix(self, source_nodes: NodeStorage, target_nodes: NodeStorage): ...
+
+    def prepare_node_data(self, graph: HeteroData) -> tuple[NodeStorage, NodeStorage]:
+        """Prepare node information and get source and target nodes."""
+        return graph[self.source_name], graph[self.target_name]
+
+    def get_edge_index(self, graph: HeteroData) -> torch.Tensor:
+        """Get the edge indices of source and target nodes.
+
+        Parameters
+        ----------
+        graph : HeteroData
+            The graph.
+
+        Returns
+        -------
+        torch.Tensor of shape (2, num_edges)
+            The edge indices.
+        """
+        source_nodes, target_nodes = self.prepare_node_data(graph)
+
+        adjmat = self.get_adjacency_matrix(source_nodes, target_nodes)
+
+        # Get source & target indices of the edges
+        edge_index = np.stack([adjmat.col, adjmat.row], axis=0)
+
+        return torch.from_numpy(edge_index).to(torch.int32)
+
+    def register_edges(self, graph: HeteroData) -> HeteroData:
+        """Register edges in the graph.
+
+        Parameters
+        ----------
+        graph : HeteroData
+            The graph to register the edges.
+
+        Returns
+        -------
+        HeteroData
+            The graph with the registered edges.
+        """
+        graph[self.name].edge_index = self.get_edge_index(graph)
+        graph[self.name].edge_type = type(self).__name__
+        return graph
+
+    def register_attributes(self, graph: HeteroData, config: DotDict) -> HeteroData:
+        """Register attributes in the edges of the graph specified.
+
+        Parameters
+        ----------
+        graph : HeteroData
+            The graph to register the attributes.
+        config : DotDict
+            The configuration of the attributes.
+
+        Returns
+        -------
+        HeteroData
+            The graph with the registered attributes.
+        """
+        for attr_name, attr_config in config.items():
+            graph[self.name][attr_name] = instantiate(attr_config).compute(graph, self.name)
+        return graph
+
+    def update_graph(self, graph: HeteroData, attrs_config: Optional[DotDict] = None) -> HeteroData:
+        """Update the graph with the edges.
+
+        Parameters
+        ----------
+        graph : HeteroData
+            The graph.
+        attrs_config : DotDict
+            The configuration of the edge attributes.
+
+        Returns
+        -------
+        HeteroData
+            The graph with the edges.
+        """
+        graph = self.register_edges(graph)
+
+        if attrs_config is None:
+            return graph
+
+        graph = self.register_attributes(graph, attrs_config)
+
+        return graph
+
+
+class KNNEdges(BaseEdgeBuilder):
+    """Computes KNN based edges and adds them to the graph.
+
+    Attributes
+    ----------
+    source_name : str
+        The name of the source nodes.
+    target_name : str
+        The name of the target nodes.
+    num_nearest_neighbours : int
+        Number of nearest neighbours.
+
+    Methods
+    -------
+    get_adjacency_matrix(source_nodes, target_nodes)
+        Compute the adjacency matrix for the KNN method.
+    register_edges(graph)
+        Register the edges in the graph.
+    register_attributes(graph, config)
+        Register attributes in the edges of the graph.
+    update_graph(graph, attrs_config)
+        Update the graph with the edges.
+    """
+
+    def __init__(self, source_name: str, target_name: str, num_nearest_neighbours: int):
+        super().__init__(source_name, target_name)
+        assert isinstance(num_nearest_neighbours, int), "Number of nearest neighbours must be an integer"
+        assert num_nearest_neighbours > 0, "Number of nearest neighbours must be positive"
+        self.num_nearest_neighbours = num_nearest_neighbours
+
+    def get_adjacency_matrix(self, source_nodes: np.ndarray, target_nodes: np.ndarray):
+        """Compute the adjacency matrix for the KNN method.
+
+        Parameters
+        ----------
+        source_nodes : np.ndarray
+            The source nodes.
+        target_nodes : np.ndarray
+            The target nodes.
+        """
+        assert self.num_nearest_neighbours is not None, "number of neighbors required for knn encoder"
+        LOGGER.info(
+            "Using KNN-Edges (with %d nearest neighbours) between %s and %s.",
+            self.num_nearest_neighbours,
+            self.source_name,
+            self.target_name,
+        )
+
+        nearest_neighbour = NearestNeighbors(metric="haversine", n_jobs=4)
+        nearest_neighbour.fit(source_nodes.x.numpy())
+        adj_matrix = nearest_neighbour.kneighbors_graph(
+            target_nodes.x.numpy(),
+            n_neighbors=self.num_nearest_neighbours,
+            mode="distance",
+        ).tocoo()
+        return adj_matrix
+
+
+class CutOffEdges(BaseEdgeBuilder):
+    """Computes cut-off based edges and adds them to the graph.
+
+    Attributes
+    ----------
+    source_name : str
+        The name of the source nodes.
+    target_name : str
+        The name of the target nodes.
+    cutoff_factor : float
+        Factor to multiply the grid reference distance to get the cut-off radius.
+
+    Methods
+    -------
+    get_cutoff_radius(graph, mask_attr)
+        Compute the cut-off radius.
+    get_adjacency_matrix(source_nodes, target_nodes)
+        Get the adjacency matrix for the cut-off method.
+    register_edges(graph)
+        Register the edges in the graph.
+    register_attributes(graph, config)
+        Register attributes in the edges of the graph.
+    update_graph(graph, attrs_config)
+        Update the graph with the edges.
+    """
+
+    def __init__(self, source_name: str, target_name: str, cutoff_factor: float):
+        super().__init__(source_name, target_name)
+        assert isinstance(cutoff_factor, (int, float)), "Cutoff factor must be a float"
+        assert cutoff_factor > 0, "Cutoff factor must be positive"
+        self.cutoff_factor = cutoff_factor
+
+    def get_cutoff_radius(self, graph: HeteroData, mask_attr: Optional[torch.Tensor] = None):
+        """Compute the cut-off radius.
+
+        The cut-off radius is computed as the product of the target nodes reference distance and the cut-off factor.
+
+        Parameters
+        ----------
+        graph : HeteroData
+            The graph.
+        mask_attr : torch.Tensor
+            The mask attribute.
+
+        Returns
+        -------
+        float
+            The cut-off radius.
+        """
+        target_nodes = graph[self.target_name]
+        mask = target_nodes[mask_attr] if mask_attr is not None else None
+        target_grid_reference_distance = get_grid_reference_distance(target_nodes.x, mask)
+        radius = target_grid_reference_distance * self.cutoff_factor
+        return radius
+
+    def prepare_node_data(self, graph: HeteroData) -> tuple[NodeStorage, NodeStorage]:
+        """Prepare node information and get source and target nodes."""
+        self.radius = self.get_cutoff_radius(graph)
+        return super().prepare_node_data(graph)
+
+    def get_adjacency_matrix(self, source_nodes: NodeStorage, target_nodes: NodeStorage):
+        """Get the adjacency matrix for the cut-off method.
+
+        Parameters
+        ----------
+        source_nodes : NodeStorage
+            The source nodes.
+        target_nodes : NodeStorage
+            The target nodes.
+        """
+        LOGGER.info(
+            "Using CutOff-Edges (with radius = %.1f km) between %s and %s.",
+            self.radius * EARTH_RADIUS,
+            self.source_name,
+            self.target_name,
+        )
+
+        nearest_neighbour = NearestNeighbors(metric="haversine", n_jobs=4)
+        nearest_neighbour.fit(source_nodes.x)
+        adj_matrix = nearest_neighbour.radius_neighbors_graph(target_nodes.x, radius=self.radius).tocoo()
+        return adj_matrix
+
+
+class MultiScaleEdges(BaseEdgeBuilder):
+    """Base class for multi-scale edges in the nodes of a graph."""
+
+    def __init__(self, source_name: str, target_name: str, x_hops: int):
+        super().__init__(source_name, target_name)
+        assert source_name == target_name, f"{self.__class__.__name__} requires source and target nodes to be the same."
+        assert isinstance(x_hops, int), "Number of x_hops must be an integer"
+        assert x_hops > 0, "Number of x_hops must be positive"
+        self.x_hops = x_hops
+
+    def adjacency_from_tri_nodes(self, source_nodes: NodeStorage):
+        source_nodes["_nx_graph"] = icosahedral.add_edges_to_nx_graph(
+            source_nodes["_nx_graph"],
+            resolutions=source_nodes["_resolutions"],
+            x_hops=self.x_hops,
+        )  # HeteroData refuses to accept None
+
+        adjmat = nx.to_scipy_sparse_array(
+            source_nodes["_nx_graph"], nodelist=list(range(len(source_nodes["_nx_graph"]))), format="coo"
+        )
+        return adjmat
+
+    def adjacency_from_hex_nodes(self, source_nodes: NodeStorage):
+
+        source_nodes["_nx_graph"] = hexagonal.add_edges_to_nx_graph(
+            source_nodes["_nx_graph"],
+            resolutions=source_nodes["_resolutions"],
+            x_hops=self.x_hops,
+        )
+
+        adjmat = nx.to_scipy_sparse_array(source_nodes["_nx_graph"], format="coo")
+        return adjmat
+
+    def get_adjacency_matrix(self, source_nodes: NodeStorage, target_nodes: NodeStorage):
+        if self.node_type == TriNodes.__name__:
+            adjmat = self.adjacency_from_tri_nodes(source_nodes)
+        elif self.node_type == HexNodes.__name__:
+            adjmat = self.adjacency_from_hex_nodes(source_nodes)
+        else:
+            raise ValueError(f"Invalid node type {self.node_type}")
+
+        adjmat = self.post_process_adjmat(source_nodes, adjmat)
+
+        return adjmat
+
+    def post_process_adjmat(self, nodes: NodeStorage, adjmat):
+        graph_sorted = {node_pos: i for i, node_pos in enumerate(nodes["_node_ordering"])}
+        sort_func = np.vectorize(graph_sorted.get)
+        adjmat.row = sort_func(adjmat.row)
+        adjmat.col = sort_func(adjmat.col)
+        return adjmat
+
+    def update_graph(self, graph: HeteroData, attrs_config: DotDict | None = None) -> HeteroData:
+        assert (
+            graph[self.source_name].node_type == TriNodes.__name__
+            or graph[self.source_name].node_type == HexNodes.__name__
+        ), f"{self.__class__.__name__} requires {TriNodes.__name__} or {HexNodes.__name__}."
+
+        self.node_type = graph[self.source_name].node_type
+
+        return super().update_graph(graph, attrs_config)