Repository: scipy/scipy

Issue:
Faster implementation of scipy.sparse.block_diag
Before making a move on this, I would like to put this potential enhancement out on GitHub to discuss, to gauge interest first. I'm hoping this discussion reveals something fruitful, perhaps even something new that I can learn, and it doesn't necessarily have to result in a PR.

For one of my projects involving message passing neural networks, which operate on graphs, I use block diagonal sparse matrices a lot. I found that the bottleneck is in the `block_diag` creation step.

To verify this, I created a minimal reproducible example, using NetworkX.

### Reproducing code example:

```python
import networkx as nx
import numpy as np
from scipy.sparse import block_diag, coo_matrix
import matplotlib.pyplot as plt

Gs = []
As = []
for i in range(50000):
    n = np.random.randint(10, 30)
    p = 0.2
    G = nx.erdos_renyi_graph(n=n, p=p)
    Gs.append(G)
    A = nx.to_numpy_array(G)
    As.append(A)
```

In a later Jupyter notebook cell, I ran the following code block.

```python
%%time
block_diag(As)
```

If we vary the number of graphs that we create a block diagonal matrix on, we can time the `block_diag` function, which I did, resulting in the following chart:

<img width="371" alt="download" src="https://user-images.githubusercontent.com/2631566/45048634-6c8c4280-b04a-11e8-9460-75dd323dc08e.png">

Some additional detail:

- 3.5 seconds for 5,000 graphs, 97,040 nodes in total.
- 13.3 seconds for 10,000 graphs, 195,367 nodes in total. (2x the inputs, ~4x the time)
- 28.2 seconds for 15,000 graphs, 291,168 nodes in total. (3x the inputs, ~8x the time)
- 52.9 seconds for 20,000 graphs, 390,313 nodes in total. (4x the inputs, ~16x the time)

My expectation was that `block_diag` should scale linearly with the number of graphs. However, the timings clearly don't show that.

I looked into the implementation of [`block_diag` in `scipy.sparse`](https://github.com/scipy/scipy/blob/v1.1.0/scipy/sparse/construct.py#L625-L676), and found that it makes a call to `bmat`, which considers multiple cases for array creation. That said, I wasn't quite clear how to make an improvement to `bmat`, thus, I embarked on an implementation that I thought might be faster. The code looks like such:

```python
def block_diag_nojit(As: List[np.array]):
    row = []
    col = []
    data = []
    start_idx = 0
    for A in As:
        nrows, ncols = A.shape
        for r in range(nrows):
            for c in range(ncols):
                if A[r, c] != 0:
                    row.append(r + start_idx)
                    col.append(c + start_idx)
                    data.append(A[r, c])
        start_idx = start_idx + nrows
    return row, col, data
                    
row, col, data = block_diag_nojit(As)
```

Profiling the time for this:

```
CPU times: user 7.44 s, sys: 146 ms, total: 7.59 s
Wall time: 7.58 s
```

And just to be careful, I ensured that the resulting row, col and data were correct:

```python
row, col, data = block_diag_nojit(As[0:10])
amat = block_diag(As[0:10])

assert np.allclose(row, amat.row)
assert np.allclose(col, amat.col)
assert np.allclose(data, amat.data)
```

I then tried JIT-ing the function using `numba`:

```python
%%time
row, col, data = jit(block_diag_nojit)(As)
```

Profiling the time:

```
CPU times: user 1.27 s, sys: 147 ms, total: 1.42 s
Wall time: 1.42 s
```

Creating the `coo_matrix` induces a flat overhead of about 900ms, thus it is most efficient to call `coo_matrix` outside of the `block_diag_nojit` function:

```python
# This performs at ~2 seconds with the scale of the graphs above.
def block_diag(As):
    row, col, data = jit(block_diag_nojit)(As)
    return coo_matrix((data, (row, col)))
```

For a full characterization:

<img width="384" alt="download" src="https://user-images.githubusercontent.com/2631566/45049207-eec93680-b04b-11e8-97b4-c6576941a12c.png">

(This is linear in the number of graphs.)

I'm not quite sure why the `bmat`-based `block_diag` implementation is super-linear, but my characterization of this re-implementation shows that it is up to ~230x faster with JIT-ing, and ~60x faster without JIT-ing (at the current scale of data). More importantly, the time taken is linear with the number of graphs, and not raised to some power.

I am quite sure that the `block_diag` matrix implementation, as it currently stands, was built with a goal in mind that I'm not aware of, which necessitated the current implementation design (using `bmat`). That said, I'm wondering if there might be interest in a PR contribution with the aforementioned faster `block_diag` function? I am open to the possibility that this might not be the right place, in which I'm happy to find a different outlet. 

Thank you for taking the time to read this!


Discussion:
Generally, as long as the function retains its API unchanged, performance improvements can be made without many restrictions.

In this case, the bad performance probably comes from the large number of blocks. It should be straightforward to write `block_diag` in the same way as the inner loop of `bmat`, so that it does not have to do the O(nblocks**2) loop.

(Adding relevant tags.)
Hey, I'm fairly new to Python and would like to contribute to this. Is that allowed?
@aadithpm yes, absolutely! I’m not a core dev of SciPy, but I’ve done some PRs into the matplotlib codebase. Happy to guide you with this, if the core devs are ok with that. (I will obviously defer to the core devs on when this is suitable for merger.)
Thank you! What should I be doing and how can I ensure that there are no API changes? Any specific ways I could test whatever changes I make?
@aadithpm my pleasure!

In terms of API, we have to make sure that the function signature does not change, and that whatever is returned remains the same.

I'd design the tests first. I'll list a (non-comprehensive) set of tests you might want to include, inspired by the docstrings; feel free to add more. When it comes to testing, it's usually practically possible to cover every corner case, but it does take some time to get there.

1. Create a list of `numpy` binary matrices (let's call it `A`), and pass them through the `block_diag` function. You can create these matrices using `numpy.random.binomial(n=1, p=p)` (substitute value of p appropriately).
1. Convert each binary matrix into a sparse matrix, and then pass in a list of sparse matrices.

I'm assuming that you are comfortable with the PR workflow. Is this assumption of mine correct? If not, please let me know, I'm happy to provide more detail on that too.
@ericmjl Thank you! I'll get started on this in a while. Thanks for the comprehensive list as well. 

As for the PR workflow; I make my changes on my repo then create a PR referencing this issue, correct?
@aadithpm That's right! The workflow I was taught to use was to fork the scipy repository, clone the fork locally, then work off a branch on the local copy. Finally, PR back into the original scipy repository from the branch. Hope this helps, and if you need more guidance, or if I was unclear about something, do ping me here.
@ericmjl if something is misunderstood, it's mostly because I'm new to this. Your replies were very clear. Thank you for taking the time to help me with this :)
Contributions are always welcome! A few more notes for getting started:

 - Work in a new branch on your fork of scipy. Using `master` can cause issues when merging.
 - Be sure to run the existing scipy test suite: `python runtests.py -s sparse` will build a local version from your current repo state and run the `scipy.sparse` tests.
 - For interactive access to the built code, run `python runtests.py -i` to get an IPython shell.
I'm sorry for the horrible delay, but I would still like to contribute to this issue if allowed. Will try to get it done as soon as possible.
@aadithpm no worries, you're experiencing first-hand how open source works - there's many things that I also want to work on, but can only find pockets of time to selectively choose :smile:. Feel free to go ahead with this issue!

Resolving patch:
diff --git a/benchmarks/benchmarks/sparse.py b/benchmarks/benchmarks/sparse.py
--- a/benchmarks/benchmarks/sparse.py
+++ b/benchmarks/benchmarks/sparse.py
@@ -192,6 +192,38 @@ def time_construction(self, name, format):
             T[i, j] = v
 
 
+class BlockDiagDenseConstruction(Benchmark):
+    param_names = ['num_matrices']
+    params = [1000, 5000, 10000, 15000, 20000]
+
+    def setup(self, num_matrices):
+        self.matrices = []
+        for i in range(num_matrices):
+            rows = np.random.randint(1, 4)
+            columns = np.random.randint(1, 4)
+            mat = np.random.randint(0, 10, (rows, columns))
+            self.matrices.append(mat)
+
+    def time_block_diag(self, num_matrices):
+        sparse.block_diag(self.matrices)
+
+
+class BlockDiagSparseConstruction(Benchmark):
+    param_names = ['num_matrices']
+    params = [100, 500, 1000, 1500, 2000]
+
+    def setup(self, num_matrices):
+        self.matrices = []
+        for i in range(num_matrices):
+            rows = np.random.randint(1, 20)
+            columns = np.random.randint(1, 20)
+            mat = np.random.randint(0, 10, (rows, columns))
+            self.matrices.append(mat)
+
+    def time_block_diag(self, num_matrices):
+        sparse.block_diag(self.matrices)
+
+
 class Conversion(Benchmark):
     params = [
         ['csr', 'csc', 'coo', 'dia', 'lil', 'dok', 'bsr'],
diff --git a/scipy/sparse/construct.py b/scipy/sparse/construct.py
--- a/scipy/sparse/construct.py
+++ b/scipy/sparse/construct.py
@@ -6,7 +6,7 @@
 __all__ = ['spdiags', 'eye', 'identity', 'kron', 'kronsum',
            'hstack', 'vstack', 'bmat', 'rand', 'random', 'diags', 'block_diag']
 
-
+import numbers
 from functools import partial
 import numpy as np
 
@@ -667,16 +667,33 @@ def block_diag(mats, format=None, dtype=None):
            [0, 0, 0, 7]])
 
     """
-    nmat = len(mats)
-    rows = []
-    for ia, a in enumerate(mats):
-        row = [None]*nmat
+    row = []
+    col = []
+    data = []
+    r_idx = 0
+    c_idx = 0
+    for a in mats:
+        if isinstance(a, (list, numbers.Number)):
+            a = coo_matrix(a)
+        nrows, ncols = a.shape
         if issparse(a):
-            row[ia] = a
+            a = a.tocoo()
+            row.append(a.row + r_idx)
+            col.append(a.col + c_idx)
+            data.append(a.data)
         else:
-            row[ia] = coo_matrix(a)
-        rows.append(row)
-    return bmat(rows, format=format, dtype=dtype)
+            a_row, a_col = np.divmod(np.arange(nrows*ncols), ncols)
+            row.append(a_row + r_idx)
+            col.append(a_col + c_idx)
+            data.append(a.ravel())
+        r_idx += nrows
+        c_idx += ncols
+    row = np.concatenate(row)
+    col = np.concatenate(col)
+    data = np.concatenate(data)
+    return coo_matrix((data, (row, col)),
+                      shape=(r_idx, c_idx),
+                      dtype=dtype).asformat(format)
 
 
 def random(m, n, density=0.01, format='coo', dtype=None,


Test patch:
diff --git a/scipy/sparse/tests/test_construct.py b/scipy/sparse/tests/test_construct.py
--- a/scipy/sparse/tests/test_construct.py
+++ b/scipy/sparse/tests/test_construct.py
@@ -423,6 +423,21 @@ def test_block_diag_1(self):
         assert_equal(construct.block_diag([1]).todense(),
                      matrix([[1]]))
 
+    def test_block_diag_sparse_matrices(self):
+        """ block_diag with sparse matrices """
+
+        sparse_col_matrices = [coo_matrix(([[1, 2, 3]]), shape=(1, 3)),
+                               coo_matrix(([[4, 5]]), shape=(1, 2))]
+        block_sparse_cols_matrices = construct.block_diag(sparse_col_matrices)
+        assert_equal(block_sparse_cols_matrices.todense(),
+                     matrix([[1, 2, 3, 0, 0], [0, 0, 0, 4, 5]]))
+
+        sparse_row_matrices = [coo_matrix(([[1], [2], [3]]), shape=(3, 1)),
+                               coo_matrix(([[4], [5]]), shape=(2, 1))]
+        block_sparse_row_matrices = construct.block_diag(sparse_row_matrices)
+        assert_equal(block_sparse_row_matrices.todense(),
+                     matrix([[1, 0], [2, 0], [3, 0], [0, 4], [0, 5]]))
+
     def test_random_sampling(self):
         # Simple sanity checks for sparse random sampling.
         for f in sprand, _sprandn: