fix: sparse export (#686)

Fixes for the export of a sparse matrix

---------
Co-authored-by: Ana Prieto Nemesio <[email protected]>
Co-authored-by: Mario Santa Cruz <[email protected]>

Author: 3 people

graphs/docs/usage/create_sparse_matrices.rst

@@ -74,7 +74,7 @@ For example, to export only the ``gauss_weight`` attribute for the
 
    % anemoi-graphs export_to_sparse graph_recipe.yaml output_dir/ \
        --edges-attributes-name gauss_weight \
-       --edges-name data->down
+       --edges-name data down
 
 You can specify multiple attributes or subgraphs by repeating the
 arguments:
@@ -84,8 +84,8 @@ arguments:
    % anemoi-graphs export_to_sparse graph_recipe.yaml output_dir/ \
        --edges-attributes-name gauss_weight \
        --edges-attributes-name another_weight \
-       --edges-name data->down \
-       --edges-name down->data
+       --edges-name data down \
+       --edges-name down data
 
 This flexibility allows you to generate only the sparse matrices you
 need for your application, reducing storage and processing time.

graphs/src/anemoi/graphs/commands/export_to_sparse.py

@@ -18,7 +18,7 @@ class ExportToSparse(Command):
 
     Example usage specifying an edge attribute:
     ```
-    anemoi-graphs export-to-sparse graph.pt output_path --edge_attribute_name edge_attr
+    anemoi-graphs export-to-sparse graph.pt output_path --edge-attribute-name edge_attr
     ```
 
     Example usage specifying a subset of edges:
@@ -33,7 +33,7 @@ class ExportToSparse(Command):
     def add_arguments(self, command_parser):
         command_parser.add_argument("graph", help="Path to the graph (a .PT file) or a config file defining the graph.")
         command_parser.add_argument("output_path", help="Path to store the inspection results.")
-        command_parser.add_argument("--edge_attribute_name", default=None, help="Name of the edge attribute to export.")
+        command_parser.add_argument("--edge-attribute-name", default=None, help="Name of the edge attribute to export.")
         command_parser.add_argument(
             "--edges-name",
             nargs=2,
@@ -46,8 +46,7 @@ def run(self, args):
         kwargs = vars(args)
         edges_name = kwargs.get("edges_name", None)
         if edges_name is not None:
-            # Convert list of lists to list of tuples
-            kwargs["edges_name"] = [tuple(pair) for pair in edges_name]
+            kwargs["edges_name"] = [(pair[0], "to", pair[1]) for pair in edges_name]
 
         GraphExporter(
             graph=kwargs["graph"],

graphs/src/anemoi/graphs/export.py

@@ -63,19 +63,46 @@ def get_nodes_info(self, source_name: str, target_name: str) -> tuple[int, int]:
 
     @staticmethod
     def get_sparse_matrix(edge_index, edge_attribute, num_source_nodes, num_target_nodes):
-        # Create sparse matrix
+        """Create sparse matrix for y = A x.
+
+        x is defined on source nodes, and output is on target nodes.
+
+        Arguments
+        ---------
+        edge_index : torch.Tensor
+            (2, E) tensor with rows: target nodes, cols: source nodes
+        edge_attribute : torch.Tensor
+            Edge weights/attributes
+        num_source_nodes : int
+            Number of source nodes (n_in)
+        num_target_nodes : int
+            Number of target nodes (n_out)
+
+        Returns
+        -------
+        torch.sparse_coo_tensor
+            Sparse COO tensor on target nodes
+        """
+        rows = edge_index[1]
+        cols = edge_index[0]
+        indices = torch.stack([rows, cols])
+
         A = torch.sparse_coo_tensor(
-            edge_index, edge_attribute, (num_source_nodes, num_target_nodes), device=edge_index.device
+            indices,
+            edge_attribute,
+            (num_target_nodes, num_source_nodes),
+            device=edge_index.device,
         )
         return A.coalesce()
 
     @staticmethod
     def convert_to_scipy_sparse(A):
         """Convert PyTorch sparse tensor to SciPy sparse matrix and save.
 
-        Args:
-            A: PyTorch sparse COO tensor
-            filename: Output filename (.npz extension)
+        Arguments
+        ---------
+        A : torch.sparse_coo_tensor
+            Sparse tensor
         """
         # Get indices and values from PyTorch sparse tensor
         indices = A.indices().cpu().numpy()