Adding GraphBlas Support

docs/notebooks/basic-usage.ipynb

@@ -53,6 +53,152 @@
     "binsparse.write(z, \"csr\", csr)\n",
     "binsparse.read(z[\"csr\"])"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style>\n",
+       "table.gb-info-table {\n",
+       "    border: 1px solid black;\n",
+       "    max-width: 100%;\n",
+       "    margin-top: 0px;\n",
+       "    margin-bottom: 0px;\n",
+       "    padding-top: 0px;\n",
+       "    padding-bottom: 0px;\n",
+       "}\n",
+       "\n",
+       "td.gb-info-name-cell {\n",
+       "    white-space: nowrap;\n",
+       "}\n",
+       "\n",
+       "details.gb-arg-details {\n",
+       "    margin-top: 0px;\n",
+       "    margin-bottom: 0px;\n",
+       "    padding-top: 0px;\n",
+       "    padding-bottom: 5px;\n",
+       "    margin-left: 10px;\n",
+       "}\n",
+       "\n",
+       "summary.gb-arg-summary {\n",
+       "    display: list-item;\n",
+       "    outline: none;\n",
+       "    margin-top: 0px;\n",
+       "    margin-bottom: 0px;\n",
+       "    padding-top: 0px;\n",
+       "    padding-bottom: 0px;\n",
+       "    margin-left: -10px;\n",
+       "}\n",
+       "\n",
+       "details.gb-expr-details {\n",
+       "    margin-top: 0px;\n",
+       "    margin-bottom: 0px;\n",
+       "    padding-top: 0px;\n",
+       "    padding-bottom: 5px;\n",
+       "}\n",
+       "\n",
+       "summary.gb-expr-summary {\n",
+       "    display: list-item;\n",
+       "    outline: none;\n",
+       "    margin-top: 0px;\n",
+       "    margin-bottom: 0px;\n",
+       "    padding-top: 0px;\n",
+       "    padding-bottom: 0px;\n",
+       "}\n",
+       "\n",
+       "blockquote.gb-expr-blockquote {\n",
+       "    margin-top: 5px;\n",
+       "    margin-bottom: 0px;\n",
+       "    padding-top: 0px;\n",
+       "    padding-bottom: 0px;\n",
+       "    margin-left: 15px;\n",
+       "}\n",
+       "\n",
+       ".gb-scalar {\n",
+       "    margin-top: 0px;\n",
+       "    margin-bottom: 0px;\n",
+       "    padding-top: 0px;\n",
+       "    padding-bottom: 5px;\n",
+       "}\n",
+       "\n",
+       "/* modify pandas dataframe */\n",
+       "table.dataframe {\n",
+       "    margin-top: 0px;\n",
+       "    margin-bottom: 0px;\n",
+       "    padding-top: 0px;\n",
+       "    padding-bottom: 0px;\n",
+       "}\n",
+       "\n",
+       "/* expression tooltips */\n",
+       ".expr-tooltip .tooltip-circle {\n",
+       "    background: #9a9cc6;\n",
+       "    color: #fff;\n",
+       "    border-radius: 50%;\n",
+       "    width: 40px;\n",
+       "    height: 40px;\n",
+       "    padding-left: 4px;\n",
+       "    padding-right: 4px;\n",
+       "}\n",
+       ".expr-tooltip .tooltip-text {\n",
+       "    visibility: hidden;\n",
+       "    position: absolute;\n",
+       "    width: 450px;\n",
+       "    background: #eef;\n",
+       "    border: 1px solid #99a;\n",
+       "    text-align: left;\n",
+       "    border-radius: 6px;\n",
+       "    padding: 3px 3px 3px 8px;\n",
+       "    margin-left: 6px;\n",
+       "}\n",
+       ".expr-tooltip:hover .tooltip-text {\n",
+       "    visibility: visible;\n",
+       "}\n",
+       ".expr-tooltip code {\n",
+       "    background-color: #f8ffed;\n",
+       "}\n",
+       "</style>\n",
+       "<details class=\"gb-arg-details\"><summary class=\"gb-arg-summary\"><tt>M<sub>0</sub></tt><div>\n",
+       "<table class=\"gb-info-table\">\n",
+       "  <tr>\n",
+       "    <td rowspan=\"2\" class=\"gb-info-name-cell\"><pre>gb.Matrix</pre></td>\n",
+       "    <td><pre>nvals</pre></td>\n",
+       "    <td><pre>nrows</pre></td>\n",
+       "    <td><pre>ncols</pre></td>\n",
+       "    <td><pre>dtype</pre></td>\n",
+       "    <td><pre>format</pre></td>\n",
+       "  </tr>\n",
+       "  <tr>\n",
+       "    <td>10000</td>\n",
+       "    <td>1000</td>\n",
+       "    <td>100</td>\n",
+       "    <td>FP64</td>\n",
+       "    <td>csr</td>\n",
+       "  </tr>\n",
+       "</table>\n",
+       "</div>\n",
+       "</summary><em>(Install</em> <tt>pandas</tt> <em>to see a preview of the data)</em></details></div>"
+      ],
+      "text/plain": [
+       "\"M_0\"      nvals  nrows  ncols  dtype  format\n",
+       "gb.Matrix  10000   1000    100   FP64     csr"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "z = zarr.group()\n",
+    "\n",
+    "binsparse.write(z, \"csr\", csr)\n",
+    "binsparse.read(z[\"csr\"], struct=\"graphblas\")"
+   ]
   }
  ],
  "metadata": {
@@ -71,7 +217,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.12"
+   "version": "3.11.6"
   },
   "vscode": {
    "interpreter": {

pyproject.toml

@@ -29,12 +29,14 @@ dependencies = [
 dev = [
     # CLI for bumping the version number
     "pre-commit",
-    "twine>=4.0.2"
+    "twine>=4.0.2",
+    "python-graphblas[default]",
 ]
 doc = [
     "docutils>=0.8,!=0.18.*,!=0.19.*",
     "sphinx>=4",
     "sphinx-book-theme>=1.0.0",
+    "matplotlib",
     "myst-nb",
     "sphinxcontrib-bibtex>=1.0.0",
     "sphinx-autodoc-typehints",

src/binsparse/_io/api.py

@@ -15,7 +15,7 @@
 from binsparse._types import GroupTypes
 
 
-def read(group: GroupTypes) -> sparse.spmatrix:
+def read(group: GroupTypes, _class: str = "scipy") -> object:
     """Read a sparse matrix from a store.
 
     Parameters
@@ -25,11 +25,11 @@ def read(group: GroupTypes) -> sparse.spmatrix:
     """
     metadata = read_attr(group, "binsparse")
     if metadata["format"] == "CSR":
-        return read_csr(group)
+        return read_csr(group, _class=_class)
     elif metadata["format"] == "CSC":
-        return read_csc(group)
+        return read_csc(group, _class=_class)
     elif metadata["format"] == "COO":
-        return read_coo(group)
+        return read_coo(group, _class=_class)
     else:
         raise NotImplementedError(f"no implementation for format {metadata['format']}")
 

src/binsparse/_io/methods.py

@@ -1,6 +1,7 @@
 from collections.abc import Mapping
 from types import MappingProxyType
 
+import graphblas as gb
 import numpy as np
 from scipy import sparse
 
@@ -59,14 +60,17 @@ def write_csr(store: GroupTypes, key: str, x: sparse.csr_matrix, *, dataset_kwar
     store.create_dataset(f"{key}/values", data=x.data, **dataset_kwargs)
 
 
-def read_csr(group: GroupTypes) -> sparse.csr_matrix:
+def read_csr(group: GroupTypes, _class: str) -> gb.Matrix | sparse.csr_matrix:
     """Read a CSR matrix from a store.
 
     Parameters
     ----------
     group
         A Zarr or h5py group.
 
+    _class
+        The class we will use to create the sparse matrix. Either "scipy" or "graphblas"
+
     Returns
     -------
     x
@@ -76,14 +80,23 @@ def read_csr(group: GroupTypes) -> sparse.csr_matrix:
     assert metadata["format"] == "CSR"
     shape = tuple(metadata["shape"])
 
-    return sparse.csr_matrix(
-        (
-            group["values"][()],
-            group["indices_1"][()],
+    if _class.lower() == "scipy":
+        return sparse.csr_matrix(
+            (
+                group["values"][()],
+                group["indices_1"][()],
+                group["pointers_to_1"][()],
+            ),
+            shape=shape,
+        )
+    elif _class.lower() == "graphblas":
+        return gb.Matrix.from_csr(
             group["pointers_to_1"][()],
-        ),
-        shape=shape,
-    )
+            group["indices_1"][()],
+            group["values"][()],
+        )
+    else:
+        raise NotImplementedError(f"no implementation for returning data using {_class}")
 
 
 def write_csc(store: GroupTypes, key: str, x: sparse.csc_matrix, *, dataset_kwargs: Mapping = MappingProxyType({})):
@@ -116,14 +129,17 @@ def write_csc(store: GroupTypes, key: str, x: sparse.csc_matrix, *, dataset_kwar
     store.create_dataset(f"{key}/values", data=x.data, **dataset_kwargs)
 
 
-def read_csc(group: GroupTypes) -> sparse.csc_matrix:
+def read_csc(group: GroupTypes, _class: str) -> gb.Matrix | sparse.csc_matrix:
     """Read a CSC matrix from a store.
 
     Parameters
     ----------
     group
         A Zarr or h5py group.
 
+    _class
+        The class we will use to create the sparse matrix. Either "scipy" or "graphblas"
+
     Returns
     -------
     x
@@ -133,14 +149,23 @@ def read_csc(group: GroupTypes) -> sparse.csc_matrix:
     assert metadata["format"] == "CSC"
     shape = tuple(metadata["shape"])
 
-    return sparse.csc_matrix(
-        (
-            group["values"][()],
-            group["indices_1"][()],
+    if _class.lower() == "scipy":
+        return sparse.csc_matrix(
+            (
+                group["values"][()],
+                group["indices_1"][()],
+                group["pointers_to_1"][()],
+            ),
+            shape=shape,
+        )
+    elif _class.lower() == "graphblas":
+        return gb.Matrix.from_csc(
             group["pointers_to_1"][()],
-        ),
-        shape=shape,
-    )
+            group["indices_1"][()],
+            group["values"][()],
+        )
+    else:
+        raise NotImplementedError(f"no implementation for returning data using {_class}")
 
 
 def write_coo(store: GroupTypes, key: str, x: sparse.csc_matrix, *, dataset_kwargs: Mapping = MappingProxyType({})):
@@ -173,14 +198,17 @@ def write_coo(store: GroupTypes, key: str, x: sparse.csc_matrix, *, dataset_kwar
     store.create_dataset(f"{key}/values", data=x.data, **dataset_kwargs)
 
 
-def read_coo(group: GroupTypes):
+def read_coo(group: GroupTypes, _class: str) -> gb.Matrix | sparse.coo_matrix:
     """Read a COO matrix from a store.
 
     Parameters
     ----------
     group
         A Zarr or h5py group.
 
+    _class
+        The class we will use to create the sparse matrix. Either "scipy" or "graphblas"
+
     Returns
     -------
     x
@@ -190,13 +218,22 @@ def read_coo(group: GroupTypes):
     assert metadata["format"] == "COO"
     shape = tuple(metadata["shape"])
 
-    return sparse.coo_matrix(
-        (
-            group["values"][()],
+    if _class.lower() == "scipy":
+        return sparse.coo_matrix(
             (
-                group["indices_0"][()],
-                group["indices_1"][()],
+                group["values"][()],
+                (
+                    group["indices_0"][()],
+                    group["indices_1"][()],
+                ),
             ),
-        ),
-        shape=shape,
-    )
+            shape=shape,
+        )
+    elif _class.lower() == "graphblas":
+        return gb.Matrix.from_coo(
+            group["indices_0"][()],
+            group["indices_1"][()],
+            group["values"][()],
+        )
+    else:
+        raise NotImplementedError(f"no implementation for returning data using {_class}")

src/binsparse/_testing.py

@@ -1,10 +1,11 @@
 from functools import singledispatch
 
+import graphblas as gb
 import numpy as np
 from scipy import sparse
 
 
-def assert_equal(a: sparse.spmatrix, b):
+def assert_equal(a: gb.Matrix | sparse.spmatrix, b):
     assert type(a) == type(b), f"types differ: {type(a)} != {type(b)}"
     assert a.shape == b.shape, f"shapes differ: {a.shape} != {b.shape}"
     _assert_equal(a, b)
@@ -26,3 +27,9 @@ def _(a, b):
 def _(a, b):
     for attr in ["row", "col", "data"]:
         np.testing.assert_equal(getattr(a, attr), getattr(b, attr))
+
+
+@_assert_equal.register(gb.Matrix)
+def _(a, b):
+    equal = a.isequal(b)
+    assert equal, "GraphBLAS matrices are not equal"