feat(graphs): New Edge Attribute: AttributeFromNode (ecmwf#62)

* Implemented new attribute

* changelog update

* changelog update

* Refactored following review

* Update src/anemoi/graphs/edges/attributes.py

Co-authored-by: Mario Santa Cruz <48736305+JPXKQX@users.noreply.github.com>

* Update src/anemoi/graphs/edges/attributes.py

Co-authored-by: Mario Santa Cruz <48736305+JPXKQX@users.noreply.github.com>

* Update src/anemoi/graphs/edges/attributes.py

Co-authored-by: Mario Santa Cruz <48736305+JPXKQX@users.noreply.github.com>

* Update src/anemoi/graphs/edges/attributes.py

Co-authored-by: Mario Santa Cruz <48736305+JPXKQX@users.noreply.github.com>

* Refactor after review

* Docstring done

* Docstring done

* Docstring done

* Fixed ABC issue

* addressed changes in docs and exception error

* Changed changelog

* TMP

* Addressed comments

* gpc

* Removed update to changelog

* Added test for fail copy

* Gpc passed

* Fixed test

* Fixed test

* Test working now

Author: icedoom888

graphs/docs/graphs/edge_attributes.rst

@@ -4,8 +4,8 @@
  Edges - Attributes
 ####################
 
-There are two main edge attributes implemented in the
-:ref:`anemoi-graphs <anemoi-graphs:index-page>` package:
+There are few edge attributes implemented in the `anemoi-graphs`
+package:
 
 *************
  Edge length
@@ -44,3 +44,49 @@ latitude and longitude coordinates of the source and target nodes.
        attributes:
          edge_length:
             _target_: anemoi.graphs.edges.attributes.EdgeDirection
+
+*********************
+ Attribute from Node
+*********************
+
+Attributes can also be copied from nodes to edges. This is done using
+the `AttributeFromNode` base class, with specialized versions for source
+and target nodes.
+
+From Source
+===========
+
+This attribute copies a specific property of the source node to the
+edge. Example usage for copying the cutout mask from nodes to edges in
+the encoder:
+
+.. code:: yaml
+
+   edges:
+     # Encoder
+   - source_name: data
+     target_name: hidden
+     edge_builders: ...
+     attributes:
+       comes_from_cutout: # Assigned name to the edge attribute, can be different than node_attr_name
+         _target_: anemoi.graphs.edges.attributes.AttributeFromSourceNode
+         node_attr_name: cutout
+
+From Target
+===========
+
+This attribute copies a specific property of the target node to the
+edge. Example usage for copying the coutout mask from nodes to edges in
+the decoder:
+
+.. code:: yaml
+
+   edges:
+      # Decoder
+    - source_name: hidden
+      target_name: data
+      edge_builders: ...
+      attributes:
+        comes_from_cutout: # Assigned name to the edge attribute, can be different than node_attr_name
+          _target_: anemoi.graphs.edges.attributes.AttributeFromTargetNode
+          node_attr_name: cutout

graphs/src/anemoi/graphs/edges/attributes.py

@@ -24,8 +24,9 @@
 class BaseEdgeAttribute(ABC, NormaliserMixin):
     """Base class for edge attributes."""
 
-    def __init__(self, norm: str | None = None) -> None:
+    def __init__(self, norm: str | None = None, dtype: str = "float32") -> None:
         self.norm = norm
+        self.dtype = dtype
 
     @abstractmethod
     def get_raw_values(self, graph: HeteroData, source_name: str, target_name: str, *args, **kwargs) -> np.ndarray: ...
@@ -35,9 +36,9 @@ def post_process(self, values: np.ndarray) -> torch.Tensor:
         if values.ndim == 1:
             values = values[:, np.newaxis]
 
-        normed_values = self.normalise(values)
+        normalised_values = self.normalise(values)
 
-        return torch.tensor(normed_values, dtype=torch.float32)
+        return torch.tensor(normalised_values.astype(self.dtype))
 
     def compute(self, graph: HeteroData, edges_name: tuple[str, str, str], *args, **kwargs) -> torch.Tensor:
         """Compute the edge attributes."""
@@ -155,3 +156,81 @@ def post_process(self, values: np.ndarray) -> torch.Tensor:
             values = 1 - values
 
         return values
+
+
+class BooleanBaseEdgeAttribute(BaseEdgeAttribute, ABC):
+    """Base class for boolean edge attributes."""
+
+    def __init__(self) -> None:
+        super().__init__(norm=None, dtype="bool")
+
+
+class BaseAttributeFromNode(BooleanBaseEdgeAttribute, ABC):
+    """
+    Base class for Attribute from Node.
+
+    Copy an attribute of either the source or target node to the edge.
+    Accesses source/target node attribute and propagates it to the edge.
+    Used for example to identify if an encoder edge originates from a LAM or global node.
+
+    Attributes
+    ----------
+    node_attr_name : str
+        Name of the node attribute to propagate.
+
+    Methods
+    -------
+    get_node_name(source_name, target_name)
+        Return the name of the node to copy.
+
+    get_raw_values(graph, source_name, target_name)
+        Computes the edge attribute from the source or target node attribute.
+
+    """
+
+    def __init__(self, node_attr_name: str) -> None:
+        super().__init__()
+        self.node_attr_name = node_attr_name
+        self.idx = None
+
+    @abstractmethod
+    def get_node_name(self, source_name: str, target_name: str): ...
+
+    def get_raw_values(self, graph: HeteroData, source_name: str, target_name: str) -> np.ndarray:
+
+        node_name = self.get_node_name(source_name, target_name)
+
+        edge_index = graph[(source_name, "to", target_name)].edge_index
+        try:
+            return graph[node_name][self.node_attr_name].numpy()[edge_index[self.idx]]
+
+        except KeyError:
+            raise KeyError(
+                f"{self.__class__.__name__} failed because the attribute '{self.node_attr_name}' is not defined for the nodes."
+            )
+
+
+class AttributeFromSourceNode(BaseAttributeFromNode):
+    """
+    Copy an attribute of the source node to the edge.
+    """
+
+    def __init__(self, node_attr_name: str) -> None:
+        super().__init__(node_attr_name)
+        self.idx = 0
+
+    def get_node_name(self, source_name: str, target_name: str):
+        return source_name
+
+
+class AttributeFromTargetNode(BaseAttributeFromNode):
+    """
+    Copy an attribute of the target node to the edge.
+    """
+
+    def __init__(self, node_attr_name: str) -> None:
+        super().__init__(node_attr_name)
+        self.idx = 1
+
+    def get_node_name(self, source_name: str, target_name: str):
+        return target_name

graphs/tests/edges/test_edge_attributes.py

@@ -7,9 +7,13 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+from functools import partial
+
 import pytest
 import torch
 
+from anemoi.graphs.edges.attributes import AttributeFromSourceNode
+from anemoi.graphs.edges.attributes import AttributeFromTargetNode
 from anemoi.graphs.edges.attributes import EdgeDirection
 from anemoi.graphs.edges.attributes import EdgeLength
 
@@ -36,3 +40,17 @@ def test_fail_edge_features(attribute_builder, graph_nodes_and_edges):
     """Test edge attribute builder fails with unknown nodes."""
     with pytest.raises(AssertionError):
         attribute_builder.compute(graph_nodes_and_edges, ("test_nodes", "to", "unknown_nodes"))
+
+
+@pytest.mark.parametrize(
+    "attribute_builder",
+    [
+        partial(AttributeFromSourceNode, node_attr_name="example_attr"),
+        partial(AttributeFromTargetNode, node_attr_name="example_attr"),
+    ],
+)
+def test_fail_edge_attribute_from_node(attribute_builder, graph_nodes_and_edges):
+    """Test edge attribute builder fails with unknown nodes."""
+    with pytest.raises(KeyError):
+        builder_instance = attribute_builder()
+        builder_instance.compute(graph_nodes_and_edges, ("test_nodes", "to", "test_nodes"))