diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index bbc225d..8f820a8 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -40,7 +40,7 @@ repos: - --force-single-line-imports - --profile black - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.7.2 + rev: v0.8.1 hooks: - id: ruff args: @@ -64,7 +64,7 @@ repos: hooks: - id: pyproject-fmt - repo: https://github.com/jshwi/docsig # Check docstrings against function sig - rev: v0.64.0 + rev: v0.65.0 hooks: - id: docsig args: @@ -74,6 +74,5 @@ repos: - --check-protected # Check protected methods - --check-class # Check class docstrings - --disable=E113 # Disable empty docstrings - - --summary # Print a summary ci: autoupdate_schedule: monthly diff --git a/CHANGELOG.md b/CHANGELOG.md index eeb48ae..02a96e9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,16 @@ Keep it human-readable, your future self will thank you! ## [Unreleased](https://github.com/ecmwf/anemoi-graphs/compare/0.4.1...HEAD) +### Added + +- feat: Support for providing lon/lat coordinates from a text file (loaded with numpy loadtxt method) to build the graph `TextNodes` (#93) +- feat: Build 2D graphs with `Voronoi` in case `SphericalVoronoi` does not work well/is an overkill (LAM). Set `flat=true` in the nodes attributes to compute area weight using Voronoi with a qhull options preventing the empty region creation (#93)  +- feat: Support for defining nodes from lat& lon NumPy arrays (#98) +- feat: new transform functions to map from sin&cos values to latlon (#98) + +### Changed +- fix: faster edge builder for tri icosahedron. (#92) + ## [0.4.1 - ICON graphs, multiple edge builders and post processors](https://github.com/ecmwf/anemoi-graphs/compare/0.4.0...0.4.1) - 2024-11-26 ### Added diff --git a/docs/graphs/node_coordinates.rst b/docs/graphs/node_coordinates.rst index 76973d2..d120a87 100644 --- a/docs/graphs/node_coordinates.rst +++ b/docs/graphs/node_coordinates.rst @@ -25,6 +25,8 @@ a file: node_coordinates/zarr_dataset node_coordinates/npz_file node_coordinates/icon_mesh + node_coordinates/text_file + node_coordinates/latlon_arrays or based on other algorithms. A commonn approach is to use an icosahedron to project the earth's surface, and refine it iteratively to diff --git a/docs/graphs/node_coordinates/latlon_arrays.rst b/docs/graphs/node_coordinates/latlon_arrays.rst new file mode 100644 index 0000000..d21d036 --- /dev/null +++ b/docs/graphs/node_coordinates/latlon_arrays.rst @@ -0,0 +1,18 @@ +####################################### + From latitude & longitude coordinates +####################################### + +Nodes can also be created directly using latitude and longitude +coordinates. Below is an example demonstrating how to add these nodes to +a graph: + +.. code:: python + + from anemoi.graphs.nodes import LatLonNodes + + ... + + lats = np.array([45.0, 45.0, 40.0, 40.0]) + lons = np.array([5.0, 10.0, 10.0, 5.0]) + + graph = LatLonNodes(latitudes=lats, longitudes=lons, name="my_nodes").update_graph(graph) diff --git a/docs/graphs/node_coordinates/text_file.rst b/docs/graphs/node_coordinates/text_file.rst new file mode 100644 index 0000000..c83710c --- /dev/null +++ b/docs/graphs/node_coordinates/text_file.rst @@ -0,0 +1,20 @@ +################ + From text file +################ + +To define the `node coordinates` based on a `.txt` file, you can +configure the `.yaml` as follows: + +.. code:: yaml + + nodes: + data: # name of the nodes + node_builder: + _target_: anemoi.graphs.nodes.TextNodes + dataset: my_file.txt + idx_lon: 0 + idx_lat: 1 + +Here, dataset refers to the path of the `.txt` file that contains the +latitude and longitude values in the columns specified by `idx_lat` and +`idx_lon`, respectively. diff --git a/src/anemoi/graphs/generate/transforms.py b/src/anemoi/graphs/generate/transforms.py index 9392241..094f9c9 100644 --- a/src/anemoi/graphs/generate/transforms.py +++ b/src/anemoi/graphs/generate/transforms.py @@ -52,6 +52,42 @@ def cartesian_to_latlon_rad(xyz: np.ndarray) -> np.ndarray: return np.array((lat, lon), dtype=np.float32).transpose() +def sincos_to_latlon_rad(sincos: np.ndarray) -> np.ndarray: + """Sine & cosine components to lat-lon coordinates. + + Parameters + ---------- + sincos : np.ndarray + The sine and cosine componenets of the latitude and longitude. Shape: (N, 4). + The dimensions correspond to: sin(lat), cos(lat), sin(lon) and cos(lon). + + Returns + ------- + np.ndarray + A 2D array of the coordinates of shape (N, 2) in radians. + """ + latitudes = np.arctan2(sincos[:, 0], sincos[:, 1]) + longitudes = np.arctan2(sincos[:, 2], sincos[:, 3]) + return np.stack([latitudes, longitudes], axis=-1) + + +def sincos_to_latlon_degrees(sincos: np.ndarray) -> np.ndarray: + """Sine & cosine components to lat-lon coordinates. + + Parameters + ---------- + sincos : np.ndarray + The sine and cosine componenets of the latitude and longitude. Shape: (N, 4). + The dimensions correspond to: sin(lat), cos(lat), sin(lon) and cos(lon). + + Returns + ------- + np.ndarray + A 2D array of the coordinates of shape (N, 2) in degrees. + """ + return np.rad2deg(sincos_to_latlon_rad(sincos)) + + def latlon_rad_to_cartesian(loc: tuple[np.ndarray, np.ndarray], radius: float = 1) -> np.ndarray: """Convert planar coordinates to 3D coordinates in a sphere. diff --git a/src/anemoi/graphs/generate/tri_icosahedron.py b/src/anemoi/graphs/generate/tri_icosahedron.py index 72cd5cb..d387649 100644 --- a/src/anemoi/graphs/generate/tri_icosahedron.py +++ b/src/anemoi/graphs/generate/tri_icosahedron.py @@ -183,9 +183,8 @@ def add_edges_to_nx_graph( node_neighbours = get_neighbours_within_hops(r_sphere, x_hops, valid_nodes=valid_nodes) _, vertex_mapping_index = tree.query(r_vertices_rad, k=1) - for idx_node, idx_neighbours in node_neighbours.items(): - graph = add_neigbours_edges(graph, idx_node, idx_neighbours, vertex_mapping_index=vertex_mapping_index) - + neighbour_pairs = create_node_neighbours_list(graph, node_neighbours, vertex_mapping_index) + graph.add_edges_from(neighbour_pairs) return graph @@ -270,3 +269,42 @@ def add_neigbours_edges( graph.add_edge(node_neighbour, node) return graph + + +def create_node_neighbours_list( + graph: nx.Graph, + node_neighbours: dict[int, set[int]], + vertex_mapping_index: np.ndarray | None = None, + self_loops: bool = False, +) -> list[tuple]: + """Preprocesses the dict of node neighbours. + + Parameters: + ----------- + graph: nx.Graph + The graph. + node_neighbours: dict[int, set[int]] + dictionairy with key: node index and value: set of neighbour node indices + vertex_mapping_index: np.ndarry + Index to map the vertices from the refined sphere to the original one, by default None. + self_loops: bool + Whether is supported to add self-loops, by default False. + + Returns: + -------- + list: tuple + A list with containing node neighbour pairs in tuples + """ + graph_nodes_idx = list(sorted(graph.nodes)) + + if vertex_mapping_index is None: + vertex_mapping_index = np.arange(len(graph.nodes)).reshape(len(graph.nodes), 1) + + neighbour_pairs = [ + (graph_nodes_idx[vertex_mapping_index[node_neighbour][0]], graph_nodes_idx[vertex_mapping_index[node][0]]) + for node, neighbours in node_neighbours.items() + for node_neighbour in neighbours + if node != node_neighbour or (self_loops and node == node_neighbour) + ] + + return neighbour_pairs diff --git a/src/anemoi/graphs/nodes/__init__.py b/src/anemoi/graphs/nodes/__init__.py index 228a282..0a9eaf9 100644 --- a/src/anemoi/graphs/nodes/__init__.py +++ b/src/anemoi/graphs/nodes/__init__.py @@ -9,6 +9,7 @@ from .builders.from_file import LimitedAreaNPZFileNodes from .builders.from_file import NPZFileNodes +from .builders.from_file import TextNodes from .builders.from_file import ZarrDatasetNodes from .builders.from_healpix import HEALPixNodes from .builders.from_healpix import LimitedAreaHEALPixNodes @@ -20,6 +21,7 @@ from .builders.from_refined_icosahedron import LimitedAreaTriNodes from .builders.from_refined_icosahedron import StretchedTriNodes from .builders.from_refined_icosahedron import TriNodes +from .builders.from_vectors import LatLonNodes __all__ = [ "ZarrDatasetNodes", @@ -27,6 +29,7 @@ "TriNodes", "HexNodes", "HEALPixNodes", + "LatLonNodes", "LimitedAreaHEALPixNodes", "LimitedAreaNPZFileNodes", "LimitedAreaTriNodes", @@ -35,4 +38,5 @@ "ICONMultimeshNodes", "ICONCellGridNodes", "ICONNodes", + "TextNodes", ] diff --git a/src/anemoi/graphs/nodes/attributes.py b/src/anemoi/graphs/nodes/attributes.py index 13ec00b..9f6676b 100644 --- a/src/anemoi/graphs/nodes/attributes.py +++ b/src/anemoi/graphs/nodes/attributes.py @@ -18,7 +18,9 @@ import numpy as np import torch from anemoi.datasets import open_dataset +from scipy.spatial import ConvexHull from scipy.spatial import SphericalVoronoi +from scipy.spatial import Voronoi from torch_geometric.data import HeteroData from torch_geometric.data.storage import NodeStorage @@ -101,6 +103,68 @@ def get_raw_values(self, nodes: NodeStorage, **kwargs) -> np.ndarray: class AreaWeights(BaseNodeAttribute): """Implements the area of the nodes as the weights. + Attributes + ---------- + flat: bool + If True, the area is computed in 2D, otherwise in 3D. + **other: Any + Additional keyword arguments, see PlanarAreaWeights and SphericalAreaWeights + for details. + + Methods + ------- + compute(self, graph, nodes_name) + Compute the area attributes for each node. + """ + + def __new__(cls, flat: bool = False, **kwargs): + logging.warning( + "Creating %s with flat=%s and kwargs=%s. In a future release, AreaWeights will be deprecated: please use directly PlanarAreaWeights or SphericalAreaWeights.", + cls.__name__, + flat, + kwargs, + ) + if flat: + return PlanarAreaWeights(**kwargs) + return SphericalAreaWeights(**kwargs) + + +class PlanarAreaWeights(BaseNodeAttribute): + """Implements the 2D area of the nodes as the weights. + + Attributes + ---------- + norm : str + Normalisation of the weights. + + Methods + ------- + compute(self, graph, nodes_name) + Compute the area attributes for each node. + """ + + def __init__( + self, + norm: str | None = None, + dtype: str = "float32", + ) -> None: + super().__init__(norm, dtype) + + def get_raw_values(self, nodes: NodeStorage, **kwargs) -> np.ndarray: + latitudes, longitudes = nodes.x[:, 0], nodes.x[:, 1] + points = np.stack([latitudes, longitudes], -1) + v = Voronoi(points, qhull_options="QJ Pp") + areas = [] + for r in v.regions: + area = ConvexHull(v.vertices[r, :]).volume + areas.append(area) + result = np.asarray(areas) + return result + + +class SphericalAreaWeights(BaseNodeAttribute): + """Implements the 3D area of the nodes as the weights. + Attributes ---------- norm : str @@ -132,22 +196,6 @@ def __init__( self.fill_value = fill_value def get_raw_values(self, nodes: NodeStorage, **kwargs) -> np.ndarray: - """Compute the area associated to each node. - - It uses Voronoi diagrams to compute the area of each node. - - Parameters - ---------- - nodes : NodeStorage - Nodes of the graph. - kwargs : dict - Additional keyword arguments. - - Returns - ------- - np.ndarray - Attributes. - """ latitudes, longitudes = nodes.x[:, 0], nodes.x[:, 1] points = latlon_rad_to_cartesian((np.asarray(latitudes), np.asarray(longitudes))) sv = SphericalVoronoi(points, self.radius, self.centre) diff --git a/src/anemoi/graphs/nodes/builders/from_file.py b/src/anemoi/graphs/nodes/builders/from_file.py index c6bdc88..069b4ab 100644 --- a/src/anemoi/graphs/nodes/builders/from_file.py +++ b/src/anemoi/graphs/nodes/builders/from_file.py @@ -63,6 +63,37 @@ def get_coordinates(self) -> torch.Tensor: return self.reshape_coords(dataset.latitudes, dataset.longitudes) +class TextNodes(BaseNodeBuilder): + """Nodes from text file. + + Attributes + ---------- + dataset : str | DictConfig + The path to txt 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. + """ + + def __init__(self, dataset, name: str, idx_lon: int = 0, idx_lat: int = 1) -> None: + LOGGER.info("Reading the dataset from %s.", dataset) + self.dataset = np.loadtxt(dataset) + self.idx_lon = idx_lon + self.idx_lat = idx_lat + super().__init__(name) + + def get_coordinates(self) -> torch.Tensor: + """Get the coordinates of the nodes. + + Returns + ------- + torch.Tensor of shape (num_nodes, 2) + A 2D tensor with the coordinates, in radians. + """ + return self.reshape_coords(self.dataset[self.idx_lat, :], self.dataset[self.idx_lon, :]) + + class NPZFileNodes(BaseNodeBuilder): """Nodes from NPZ defined grids. @@ -146,7 +177,10 @@ def get_coordinates(self) -> np.ndarray: ) area_mask = self.area_mask_builder.get_mask(coords) - LOGGER.info("Dropping %d nodes from the processor mesh.", len(area_mask) - area_mask.sum()) + LOGGER.info( + "Dropping %d nodes from the processor mesh.", + len(area_mask) - area_mask.sum(), + ) coords = coords[area_mask] return coords diff --git a/src/anemoi/graphs/nodes/builders/from_vectors.py b/src/anemoi/graphs/nodes/builders/from_vectors.py new file mode 100644 index 0000000..f58767c --- /dev/null +++ b/src/anemoi/graphs/nodes/builders/from_vectors.py @@ -0,0 +1,67 @@ +# (C) Copyright 2024 Anemoi contributors. +# +# 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 __future__ import annotations + +import logging + +import numpy as np +import torch + +from anemoi.graphs.nodes.builders.base import BaseNodeBuilder + +LOGGER = logging.getLogger(__name__) + + +class LatLonNodes(BaseNodeBuilder): + """Nodes from its latitude and longitude positions (in numpy arrays). + + Attributes + ---------- + latitudes : list | np.ndarray + The latitude of the nodes, in degrees. + longitudes : list | np.ndarray + The longitude of the nodes, in degrees. + + Methods + ------- + get_coordinates() + Get the lat-lon coordinates of the nodes. + register_nodes(graph, name) + Register the nodes in the graph. + register_attributes(graph, name, config) + Register the attributes in the nodes of the graph specified. + update_graph(graph, name, attrs_config) + Update the graph with new nodes and attributes. + """ + + def __init__(self, latitudes: list[float] | np.ndarray, longitudes: list[float] | np.ndarray, name: str) -> None: + super().__init__(name) + self.latitudes = latitudes if isinstance(latitudes, np.ndarray) else np.array(latitudes) + self.longitudes = longitudes if isinstance(longitudes, np.ndarray) else np.array(longitudes) + + assert len(self.latitudes) == len( + self.longitudes + ), f"Lenght of latitudes and longitudes must match but {len(self.latitudes)}!={len(self.longitudes)}." + assert self.latitudes.ndim == 1 or ( + self.latitudes.ndim == 2 and self.latitudes.shape[1] == 1 + ), "latitudes must have shape (N, ) or (N, 1)." + assert self.longitudes.ndim == 1 or ( + self.longitudes.ndim == 2 and self.longitudes.shape[1] == 1 + ), "longitudes must have shape (N, ) or (N, 1)." + + def get_coordinates(self) -> torch.Tensor: + """Get the coordinates of the nodes. + + Returns + ------- + torch.Tensor of shape (num_nodes, 2) + A 2D tensor with the coordinates, in radians. + """ + return self.reshape_coords(self.latitudes, self.longitudes) diff --git a/tests/nodes/test_arrays.py b/tests/nodes/test_arrays.py new file mode 100644 index 0000000..77e6f20 --- /dev/null +++ b/tests/nodes/test_arrays.py @@ -0,0 +1,64 @@ +# (C) Copyright 2024 Anemoi contributors. +# +# 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 pytest +import torch +from torch_geometric.data import HeteroData + +from anemoi.graphs.nodes.attributes import AreaWeights +from anemoi.graphs.nodes.attributes import UniformWeights +from anemoi.graphs.nodes.builders.from_vectors import LatLonNodes + +lats = [45.0, 45.0, 40.0, 40.0] +lons = [5.0, 10.0, 10.0, 5.0] + + +def test_init(): + """Test LatLonNodes initialization.""" + node_builder = LatLonNodes(latitudes=lats, longitudes=lons, name="test_nodes") + assert isinstance(node_builder, LatLonNodes) + + +def test_fail_init_length_mismatch(): + """Test LatLonNodes initialization with invalid argument.""" + lons = [5.0, 10.0, 10.0, 5.0, 5.0] + + with pytest.raises(AssertionError): + LatLonNodes(latitudes=lats, longitudes=lons, name="test_nodes") + + +def test_fail_init_missing_argument(): + """Test NPZFileNodes initialization with missing argument.""" + with pytest.raises(TypeError): + LatLonNodes(name="test_nodes") + + +def test_register_nodes(): + """Test LatLonNodes register correctly the nodes.""" + graph = HeteroData() + node_builder = LatLonNodes(latitudes=lats, longitudes=lons, name="test_nodes") + graph = node_builder.register_nodes(graph) + + assert graph["test_nodes"].x is not None + assert isinstance(graph["test_nodes"].x, torch.Tensor) + assert graph["test_nodes"].x.shape == (len(lats), 2) + assert graph["test_nodes"].node_type == "LatLonNodes" + + +@pytest.mark.parametrize("attr_class", [UniformWeights, AreaWeights]) +def test_register_attributes(graph_with_nodes: HeteroData, attr_class): + """Test LatLonNodes register correctly the weights.""" + node_builder = LatLonNodes(latitudes=lats, longitudes=lons, name="test_nodes") + config = {"test_attr": {"_target_": f"anemoi.graphs.nodes.attributes.{attr_class.__name__}"}} + + graph = node_builder.register_attributes(graph_with_nodes, config) + + assert graph["test_nodes"]["test_attr"] is not None + assert isinstance(graph["test_nodes"]["test_attr"], torch.Tensor) + assert graph["test_nodes"]["test_attr"].shape[0] == graph["test_nodes"].x.shape[0]