feat(graphs): add lat weighted attribute

graphs/src/anemoi/graphs/nodes/attributes/__init__.py

@@ -9,6 +9,7 @@
 
 from .area_weights import CosineLatWeightedAttribute
 from .area_weights import IsolatitudeAreaWeights
+from .area_weights import MaskedPlanarAreaWeights
 from .area_weights import PlanarAreaWeights
 from .area_weights import SphericalAreaWeights
 from .area_weights import UniformWeights

graphs/src/anemoi/graphs/nodes/attributes/area_weights.py

@@ -377,7 +377,7 @@ def compute_latitude_weight(self, latitudes: np.ndarray) -> np.ndarray:
         # Compute the latitude band area
         lat_1 = divisory_lats[1:]
         lat_2 = divisory_lats[:-1]
-        ring_area_km = 2 * np.pi * EARTH_RADIUS * (np.sin(lat_2) - np.sin(lat_1))
+        ring_area_km = 2 * np.pi * EARTH_RADIUS * (np.sin(lat_1) - np.sin(lat_2))
 
         # Compute the number of points/nodes at each latitude band
         lat_to_ring = {lat: idx for idx, lat in enumerate(unique_lats)}

graphs/tests/nodes/attributes/test_weights.py

@@ -15,6 +15,7 @@
 
 from anemoi.graphs.nodes.attributes import CosineLatWeightedAttribute
 from anemoi.graphs.nodes.attributes import IsolatitudeAreaWeights
+from anemoi.graphs.nodes.attributes import MaskedPlanarAreaWeights
 from anemoi.graphs.nodes.attributes import PlanarAreaWeights
 from anemoi.graphs.nodes.attributes import SphericalAreaWeights
 from anemoi.graphs.nodes.attributes import UniformWeights
@@ -71,12 +72,35 @@ def test_spherical_area_weights_wrong_fill_value(fill_value: str):
 
 
 @pytest.mark.parametrize("attr_class", [IsolatitudeAreaWeights, CosineLatWeightedAttribute])
-def test_latweighted(attr_class: Type[BaseNodeAttribute], graph_with_rectilinear_nodes):
+@pytest.mark.parametrize("norm", [None, "l1", "unit-max"])
+def test_latweighted(attr_class: Type[BaseNodeAttribute], graph_with_rectilinear_nodes, norm: str):
     """Test attribute builder for Lat with different fill values."""
-    node_attr_builder = attr_class(norm="l1")
+    node_attr_builder = attr_class(norm=norm)
     weights = node_attr_builder.compute(graph_with_rectilinear_nodes, "test_nodes")
 
     assert weights is not None
     assert isinstance(weights, torch.Tensor)
+    assert torch.all(weights >= 0)
     assert weights.shape[0] == graph_with_rectilinear_nodes["test_nodes"].x.shape[0]
     assert weights.dtype == node_attr_builder.dtype
+
+
+def test_masked_planar_area_weights(graph_with_nodes: HeteroData):
+    """Test attribute builder for PlanarAreaWeights."""
+    node_attr_builder = MaskedPlanarAreaWeights(mask_node_attr_name="interior_mask")
+    weights = node_attr_builder.compute(graph_with_nodes, "test_nodes")
+
+    assert weights is not None
+    assert isinstance(weights, torch.Tensor)
+    assert weights.shape[0] == graph_with_nodes["test_nodes"].x.shape[0]
+    assert weights.dtype == node_attr_builder.dtype
+
+    mask = graph_with_nodes["test_nodes"]["interior_mask"]
+    assert torch.all(weights[~mask] == 0)
+
+
+def test_masked_planar_area_weights_fail(graph_with_nodes: HeteroData):
+    """Test attribute builder for AreaWeights with invalid radius."""
+    with pytest.raises(AssertionError):
+        node_attr_builder = MaskedPlanarAreaWeights(mask_node_attr_name="nonexisting")
+        node_attr_builder.compute(graph_with_nodes, "test_nodes")