python - How to apply same operation on every block of block matrices in numpy efficiently? -

i have big 2d array following shape:

b = [b_0, b_1, b_2, b_n] 

where b_0, b_1, ..., b_n have same number of rows, different number of columns , n may large. have 1d array idx shape (n+1,), ,

b_i = b[:, idx[i]:idx[i+1]] 

and idx[-1] (the last elements of idx) total number of columns of b.

i want same matrix operation every b_i, example:

b_i.t()@b_i 

or 2d array:

d = [[d_0], [d_1], ..., [d_n]] 

with d_0, d_1, ..., d_n have same number of columns equal number of rows of b, different number of rows, ,

d_i = d[idx[i]:idx[i+1], :] 

and want compute d_i@b_i.

so question how implement in python , avoid using for loop?

the following example:

import numpy np timeit import default_timer timer # prepare test data n = 1000000 # number of small matrix   idx = np.zeros(n+1, dtype=np.int) idx[1:] = np.random.randint(1, 10, size=n) idx = np.cumsum(idx)  b = np.random.rand(3, idx[-1])  # computation start = timer() c = [] in range(n):     b_i = b[:, idx[i]:idx[i+1]]     c_i = b_i.t@b_i     c.append(c_i) end = timer() print('total time:', end - start) 

if add code:

print(b.shape) print(idx) print([x.shape x in c])  bnn = np.zeros((n, 3, idx[-1])) in range(n):     s = np.s_[idx[i]:idx[i+1]]     bnn[i,:,s] = b[:, s] bnn = bnn.reshape(3*n,-1) cnn = bnn.t @ bnn print(bnn.shape, cnn.shape) print(cnn.sum(), sum([x.sum() x in c])) 

and change n=5, get

2115:~/mypy$ python3 stack46209231.py  (3, 31)    # b shape [ 0  9 17 18 25 31] [(9, 9), (8, 8), (1, 1), (7, 7), (6, 6)]  # shapes of c elements (15, 31) (31, 31)     # shapes of diagonalized b , c 197.407879357 197.407879357   # c sums 2 routes 

so idea of making diagonalized version of b, , performing dot product works. modest sized arrays should faster, though iteration create bnn take time, extracting blocks cnn.

but bnn , cnn large, , bogged down memory swaps, , memory errors.

---

with block_diag function, turning b sparse matrix quite easy:

from scipy import sparse  blist = [b[:, idx[i]:idx[i+1]] in range(n)] bs = sparse.block_diag(blist, format='bsr') print(repr(bs)) cs = bs.t@bs print(repr(cs)) print(cs.sum()) 

and sample run

2158:~/mypy$ python3 stack46209231.py  (3, 20) [ 0  1  5  9 17 20] [(1, 1), (4, 4), (4, 4), (8, 8), (3, 3)] (15, 20) (20, 20) 94.4190125992 94.4190125992 <15x20 sparse matrix of type '<class 'numpy.float64'>'     60 stored elements (blocksize = 1x1) in block sparse row format> <20x20 sparse matrix of type '<class 'numpy.float64'>'     106 stored elements (blocksize = 1x1) in block sparse row format> 

and shapes , checksums match.

for n = 10000, bnn large memory. sparse bs creation slow, matrix multiplication fast.