Implement .generators.ego.ego_graph (#61)

* Implement `.generators.ego.ego_graph`

Also, clean up shared BFS functions and move to `_bfs.py`.

* use external images in README so they render on PyPI

Support and test against Python 3.11
Change development status to Beta (was Alpha).

Author: eriknw

Diff for: .github/workflows/test.yml

@@ -15,7 +15,7 @@ jobs:
       fail-fast: true
       matrix:
         os: ["ubuntu-latest", "macos-latest", "windows-latest"]
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
     steps:
       - name: Checkout
         uses: actions/checkout@v3

Diff for: .pre-commit-config.yaml

@@ -55,7 +55,7 @@ repos:
       - id: black
       # - id: black-jupyter
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: v0.0.263
+    rev: v0.0.264
     hooks:
       - id: ruff
         args: [--fix-only, --show-fixes]
@@ -81,7 +81,7 @@ repos:
         additional_dependencies: [tomli]
         files: ^(graphblas_algorithms|docs)/
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: v0.0.263
+    rev: v0.0.264
     hooks:
       - id: ruff
   # `pyroma` may help keep our package standards up to date if best practices change.

Diff for: MANIFEST.in

@@ -4,6 +4,3 @@ include setup.py
 include README.md
 include LICENSE
 include MANIFEST.in
-docs/_static/img/logo-name-medium.png
-docs/_static/img/graphblas-vs-igraph.png
-docs/_static/img/graphblas-vs-networkx.png

Diff for: README.md

@@ -1,4 +1,4 @@
-![GraphBLAS Algorithms](docs/_static/img/logo-name-medium.svg)
+![GraphBLAS Algorithms](https://raw.githubusercontent.com/python-graphblas/graphblas-algorithms/main/docs/_static/img/logo-name-medium.svg)
 <br>
 [![conda-forge](https://img.shields.io/conda/vn/conda-forge/graphblas-algorithms.svg)](https://anaconda.org/conda-forge/graphblas-algorithms)
 [![pypi](https://img.shields.io/pypi/v/graphblas-algorithms.svg)](https://pypi.python.org/pypi/graphblas-algorithms/)
@@ -21,9 +21,9 @@ Why use GraphBLAS Algorithms? Because it is *fast*, *flexible*, and *familiar* b
 Are we missing any [algorithms](#Plugin-Algorithms) that you want?
 [Please let us know!](https://github.com/python-graphblas/graphblas-algorithms/issues)
 <br>
-<img src="docs/_static/img/graphblas-vs-networkx.png" alt="GraphBLAS vs NetworkX" title="Even faster than scipy.sparse!" width="640" />
+<img src="https://raw.githubusercontent.com/python-graphblas/graphblas-algorithms/main/docs/_static/img/graphblas-vs-networkx.png" alt="GraphBLAS vs NetworkX" title="Even faster than scipy.sparse!" width="640" />
 <br>
-<img src="docs/_static/img/graphblas-vs-igraph.png" alt="GraphBLAS vs igraph" title="igraph may use different algorithms for PageRank" width="600" />
+<img src="https://raw.githubusercontent.com/python-graphblas/graphblas-algorithms/main/docs/_static/img/graphblas-vs-igraph.png" alt="GraphBLAS vs igraph" title="igraph may use different algorithms for PageRank" width="600" />
 
 ### Installation
 ```
@@ -151,6 +151,8 @@ dispatch pattern shown above.
   - descendants
 - Dominating
   - is_dominating_set
+- Generators
+  - ego_graph
 - Isolate
   - is_isolate
   - isolates

Diff for: graphblas_algorithms/__init__.py

@@ -3,6 +3,7 @@
 from .classes import *
 
 from .algorithms import *  # isort:skip
+from .generators import *  # isort:skip
 
 try:
     __version__ = importlib.metadata.version("graphblas-algorithms")

Diff for: graphblas_algorithms/algorithms/_bfs.py

@@ -0,0 +1,150 @@
+"""BFS routines used by other algorithms"""
+
+import numpy as np
+from graphblas import Matrix, Vector, binary, replace, unary
+from graphblas.semiring import any_pair
+
+
+def _get_cutoff(n, cutoff):
+    if cutoff is None or cutoff >= n:
+        return n  # Everything
+    return cutoff + 1  # Inclusive
+
+
+def _plain_bfs(G, source, *, cutoff=None):
+    index = G._key_to_id[source]
+    A = G.get_property("offdiag")
+    n = A.nrows
+    v = Vector(bool, n, name="bfs_plain")
+    q = Vector(bool, n, name="q")
+    v[index] = True
+    q[index] = True
+    any_pair_bool = any_pair[bool]
+    cutoff = _get_cutoff(n, cutoff)
+    for _i in range(1, cutoff):
+        q(~v.S, replace) << any_pair_bool(q @ A)
+        if q.nvals == 0:
+            break
+        v(q.S) << True
+    return v
+
+
+def _bfs_level(G, source, cutoff=None, *, transpose=False, dtype=int):
+    if dtype == bool:
+        dtype = int
+    index = G._key_to_id[source]
+    A = G.get_property("offdiag")
+    if transpose and G.is_directed():
+        A = A.T  # TODO: should we use "AT" instead?
+    n = A.nrows
+    v = Vector(dtype, n, name="bfs_level")
+    q = Vector(bool, n, name="q")
+    v[index] = 0
+    q[index] = True
+    any_pair_bool = any_pair[bool]
+    cutoff = _get_cutoff(n, cutoff)
+    for i in range(1, cutoff):
+        q(~v.S, replace) << any_pair_bool(q @ A)
+        if q.nvals == 0:
+            break
+        v(q.S) << i
+    return v
+
+
+def _bfs_levels(G, nodes, cutoff=None, *, dtype=int):
+    if dtype == bool:
+        dtype = int
+    A = G.get_property("offdiag")
+    n = A.nrows
+    if nodes is None:
+        # TODO: `D = Vector.from_scalar(0, n, dtype).diag()`
+        D = Vector(dtype, n, name="bfs_levels_vector")
+        D << 0
+        D = D.diag(name="bfs_levels")
+    else:
+        ids = G.list_to_ids(nodes)
+        D = Matrix.from_coo(
+            np.arange(len(ids), dtype=np.uint64),
+            ids,
+            0,
+            dtype,
+            nrows=len(ids),
+            ncols=n,
+            name="bfs_levels",
+        )
+    Q = unary.one[bool](D).new(name="Q")
+    any_pair_bool = any_pair[bool]
+    cutoff = _get_cutoff(n, cutoff)
+    for i in range(1, cutoff):
+        Q(~D.S, replace) << any_pair_bool(Q @ A)
+        if Q.nvals == 0:
+            break
+        D(Q.S) << i
+    return D
+
+
+# TODO: benchmark this and the version commented out below
+def _plain_bfs_bidirectional(G, source):
+    # Bi-directional BFS w/o symmetrizing the adjacency matrix
+    index = G._key_to_id[source]
+    A = G.get_property("offdiag")
+    # XXX: should we use `AT` if available?
+    n = A.nrows
+    v = Vector(bool, n, name="bfs_plain")
+    q_out = Vector(bool, n, name="q_out")
+    q_in = Vector(bool, n, name="q_in")
+    v[index] = True
+    q_in[index] = True
+    any_pair_bool = any_pair[bool]
+    is_out_empty = True
+    is_in_empty = False
+    for _i in range(1, n):
+        # Traverse out-edges from the most recent `q_in` and `q_out`
+        if is_out_empty:
+            q_out(~v.S) << any_pair_bool(q_in @ A)
+        else:
+            q_out << binary.any(q_out | q_in)
+            q_out(~v.S, replace) << any_pair_bool(q_out @ A)
+        is_out_empty = q_out.nvals == 0
+        if not is_out_empty:
+            v(q_out.S) << True
+        elif is_in_empty:
+            break
+        # Traverse in-edges from the most recent `q_in` and `q_out`
+        if is_in_empty:
+            q_in(~v.S) << any_pair_bool(A @ q_out)
+        else:
+            q_in << binary.any(q_out | q_in)
+            q_in(~v.S, replace) << any_pair_bool(A @ q_in)
+        is_in_empty = q_in.nvals == 0
+        if not is_in_empty:
+            v(q_in.S) << True
+        elif is_out_empty:
+            break
+    return v
+
+
+"""
+def _plain_bfs_bidirectional(G, source):
+    # Bi-directional BFS w/o symmetrizing the adjacency matrix
+    index = G._key_to_id[source]
+    A = G.get_property("offdiag")
+    n = A.nrows
+    v = Vector(bool, n, name="bfs_plain")
+    q = Vector(bool, n, name="q")
+    q2 = Vector(bool, n, name="q_2")
+    v[index] = True
+    q[index] = True
+    any_pair_bool = any_pair[bool]
+    for _i in range(1, n):
+        q2(~v.S, replace) << any_pair_bool(q @ A)
+        v(q2.S) << True
+        q(~v.S, replace) << any_pair_bool(A @ q)
+        if q.nvals == 0:
+            if q2.nvals == 0:
+                break
+            q, q2 = q2, q
+        elif q2.nvals != 0:
+            q << binary.any(q | q2)
+    return v
+"""

Diff for: graphblas_algorithms/algorithms/centrality/eigenvector.py

@@ -1,11 +1,7 @@
 from graphblas import Vector
 
-from graphblas_algorithms.algorithms._helpers import is_converged, normalize
-from graphblas_algorithms.algorithms.exceptions import (
-    ConvergenceFailure,
-    GraphBlasAlgorithmException,
-    PointlessConcept,
-)
+from .._helpers import is_converged, normalize
+from ..exceptions import ConvergenceFailure, GraphBlasAlgorithmException, PointlessConcept
 
 __all__ = ["eigenvector_centrality"]
 

Diff for: graphblas_algorithms/algorithms/centrality/katz.py

@@ -2,11 +2,8 @@
 from graphblas.core.utils import output_type
 from graphblas.semiring import plus_first, plus_times
 
-from graphblas_algorithms.algorithms._helpers import is_converged, normalize
-from graphblas_algorithms.algorithms.exceptions import (
-    ConvergenceFailure,
-    GraphBlasAlgorithmException,
-)
+from .._helpers import is_converged, normalize
+from ..exceptions import ConvergenceFailure, GraphBlasAlgorithmException
 
 __all__ = ["katz_centrality"]
 

Diff for: graphblas_algorithms/algorithms/components/connected.py

@@ -1,7 +1,5 @@
-from graphblas import Vector, replace
-from graphblas.semiring import any_pair
-
-from graphblas_algorithms.algorithms.exceptions import PointlessConcept
+from .._bfs import _plain_bfs
+from ..exceptions import PointlessConcept
 
 
 def is_connected(G):
@@ -12,20 +10,3 @@ def is_connected(G):
 
 def node_connected_component(G, n):
     return _plain_bfs(G, n)
-
-
-def _plain_bfs(G, source):
-    index = G._key_to_id[source]
-    A = G.get_property("offdiag")
-    n = A.nrows
-    v = Vector(bool, n, name="bfs_plain")
-    q = Vector(bool, n, name="q")
-    v[index] = True
-    q[index] = True
-    any_pair_bool = any_pair[bool]
-    for _i in range(1, n):
-        q(~v.S, replace) << any_pair_bool(q @ A)
-        if q.nvals == 0:
-            break
-        v(q.S) << True
-    return v

Diff for: graphblas_algorithms/algorithms/components/weakly_connected.py

@@ -1,77 +1,8 @@
-from graphblas import Vector, binary, replace
-from graphblas.semiring import any_pair
-
-from graphblas_algorithms.algorithms.exceptions import PointlessConcept
+from .._bfs import _plain_bfs_bidirectional
+from ..exceptions import PointlessConcept
 
 
 def is_weakly_connected(G):
     if len(G) == 0:
         raise PointlessConcept("Connectivity is undefined for the null graph.")
-    return _plain_bfs(G, next(iter(G))).nvals == len(G)
-
-
-# TODO: benchmark this and the version commented out below
-def _plain_bfs(G, source):
-    # Bi-directional BFS w/o symmetrizing the adjacency matrix
-    index = G._key_to_id[source]
-    A = G.get_property("offdiag")
-    # XXX: should we use `AT` if available?
-    n = A.nrows
-    v = Vector(bool, n, name="bfs_plain")
-    q_out = Vector(bool, n, name="q_out")
-    q_in = Vector(bool, n, name="q_in")
-    v[index] = True
-    q_in[index] = True
-    any_pair_bool = any_pair[bool]
-    is_out_empty = True
-    is_in_empty = False
-    for _i in range(1, n):
-        # Traverse out-edges from the most recent `q_in` and `q_out`
-        if is_out_empty:
-            q_out(~v.S) << any_pair_bool(q_in @ A)
-        else:
-            q_out << binary.any(q_out | q_in)
-            q_out(~v.S, replace) << any_pair_bool(q_out @ A)
-        is_out_empty = q_out.nvals == 0
-        if not is_out_empty:
-            v(q_out.S) << True
-        elif is_in_empty:
-            break
-        # Traverse in-edges from the most recent `q_in` and `q_out`
-        if is_in_empty:
-            q_in(~v.S) << any_pair_bool(A @ q_out)
-        else:
-            q_in << binary.any(q_out | q_in)
-            q_in(~v.S, replace) << any_pair_bool(A @ q_in)
-        is_in_empty = q_in.nvals == 0
-        if not is_in_empty:
-            v(q_in.S) << True
-        elif is_out_empty:
-            break
-    return v
-
-
-"""
-def _plain_bfs(G, source):
-    # Bi-directional BFS w/o symmetrizing the adjacency matrix
-    index = G._key_to_id[source]
-    A = G.get_property("offdiag")
-    n = A.nrows
-    v = Vector(bool, n, name="bfs_plain")
-    q = Vector(bool, n, name="q")
-    q2 = Vector(bool, n, name="q_2")
-    v[index] = True
-    q[index] = True
-    any_pair_bool = any_pair[bool]
-    for _i in range(1, n):
-        q2(~v.S, replace) << any_pair_bool(q @ A)
-        v(q2.S) << True
-        q(~v.S, replace) << any_pair_bool(A @ q)
-        if q.nvals == 0:
-            if q2.nvals == 0:
-                break
-            q, q2 = q2, q
-        elif q2.nvals != 0:
-            q << binary.any(q | q2)
-    return v
-"""
+    return _plain_bfs_bidirectional(G, next(iter(G))).nvals == len(G)

Diff for: graphblas_algorithms/algorithms/core.py

@@ -1,11 +1,12 @@
 from graphblas import Matrix, monoid, replace, select, semiring
 
-from graphblas_algorithms.classes.graph import Graph
+from graphblas_algorithms import Graph
 
 __all__ = ["k_truss"]
 
 
 def k_truss(G: Graph, k) -> Graph:
+    # TODO: should we have an option to keep the output matrix the same size?
     # Ignore self-edges
     S = G.get_property("offdiag")
 
@@ -32,6 +33,5 @@ def k_truss(G: Graph, k) -> Graph:
     Ktruss = C[indices, indices].new()
 
     # Convert back to networkx graph with correct node ids
-    keys = G.list_to_keys(indices)
-    key_to_id = dict(zip(keys, range(len(indices))))
+    key_to_id = G.renumber_key_to_id(indices.tolist())
     return Graph(Ktruss, key_to_id=key_to_id)

Diff for: graphblas_algorithms/algorithms/link_analysis/hits_alg.py

@@ -1,7 +1,7 @@
 from graphblas import Vector
 
-from graphblas_algorithms.algorithms._helpers import is_converged, normalize
-from graphblas_algorithms.algorithms.exceptions import ConvergenceFailure
+from .._helpers import is_converged, normalize
+from ..exceptions import ConvergenceFailure
 
 __all__ = ["hits"]