From d896e4a7b94b70cae3b2ef339a3c0ddaaa390e52 Mon Sep 17 00:00:00 2001
From: Erik Welch <erik.n.welch@gmail.com>
Date: Sat, 29 Oct 2022 12:02:01 -0500
Subject: [PATCH] Update `matrix_to_dicts` and fix `generalized_degree`

Now `matrix_to_dicts` can return either the node id or row or column index.
---
 graphblas_algorithms/classes/_utils.py | 40 ++++++++++++++++++++------
 graphblas_algorithms/nxapi/cluster.py  |  2 +-
 2 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/graphblas_algorithms/classes/_utils.py b/graphblas_algorithms/classes/_utils.py
index 4c60998..f46ea14 100644
--- a/graphblas_algorithms/classes/_utils.py
+++ b/graphblas_algorithms/classes/_utils.py
@@ -167,13 +167,19 @@ def vector_to_set(self, v):
     return {id_to_key[index] for index in indices}
 
 
-def matrix_to_dicts(self, A):
-    """{row: {col: val}}"""
+def matrix_to_dicts(self, A, *, use_row_index=False, use_column_index=False):
+    """Convert a Matrix to a dict of dicts of the form ``{row: {col: val}}``
+
+    Use ``use_row_index=True`` to return the row index as keys in the dict,
+    and likewise for `use_column_index=True``.
+
+    """
     if isinstance(A, TransposedMatrix):
         # Not covered
         d = A.T.ss.export("hypercsc")
         rows = d["cols"].tolist()
         col_indices = d["row_indices"].tolist()
+        use_row_index, use_column_index = use_column_index, use_row_index
     else:
         d = A.ss.export("hypercsr")
         rows = d["rows"].tolist()
@@ -181,14 +187,30 @@ def matrix_to_dicts(self, A):
     indptr = d["indptr"]
     values = d["values"].tolist()
     id_to_key = self.id_to_key
-    return {
-        id_to_key[row]: {
-            id_to_key[col]: val for col, val in zip(col_indices[start:stop], values[start:stop])
+    it = zip(rows, np.lib.stride_tricks.sliding_window_view(indptr, 2).tolist())
+    if use_row_index and use_column_index:
+        return {
+            row: dict(zip(col_indices[start:stop], values[start:stop])) for row, (start, stop) in it
+        }
+    elif use_row_index:
+        return {
+            row: {
+                id_to_key[col]: val for col, val in zip(col_indices[start:stop], values[start:stop])
+            }
+            for row, (start, stop) in it
+        }
+    elif use_column_index:
+        return {
+            id_to_key[row]: dict(zip(col_indices[start:stop], values[start:stop]))
+            for row, (start, stop) in it
+        }
+    else:
+        return {
+            id_to_key[row]: {
+                id_to_key[col]: val for col, val in zip(col_indices[start:stop], values[start:stop])
+            }
+            for row, (start, stop) in it
         }
-        for row, (start, stop) in zip(
-            rows, np.lib.stride_tricks.sliding_window_view(indptr, 2).tolist()
-        )
-    }
 
 
 def to_networkx(self, edge_attribute="weight"):
diff --git a/graphblas_algorithms/nxapi/cluster.py b/graphblas_algorithms/nxapi/cluster.py
index 7f8bd96..3b3ee58 100644
--- a/graphblas_algorithms/nxapi/cluster.py
+++ b/graphblas_algorithms/nxapi/cluster.py
@@ -138,4 +138,4 @@ def generalized_degree(G, nodes=None):
         return G.vector_to_nodemap(result)
     mask = G.list_to_mask(nodes)
     result = algorithms.generalized_degree(G, mask=mask)
-    return G.matrix_to_dicts(result)
+    return G.matrix_to_dicts(result, use_column_index=True)