Update matrix_to_dicts and fix generalized_degree

graphblas_algorithms/classes/_utils.py

@@ -167,28 +167,50 @@ 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()
         col_indices = d["col_indices"].tolist()
     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"):

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)