# Copyright (c) 2014-2015, Alan Saul
# Licensed under the BSD 3-clause license (see LICENSE.txt)
import numpy as np
[docs]def get_blocks_3d(A, blocksizes, pagesizes=None):
"""
Given a 3d matrix, make a block matrix, where the first and second dimensions are blocked according
to blocksizes, and the pages are blocked using pagesizes
"""
assert (A.shape[0]==A.shape[1]) and len(A.shape)==3, "can't blockify this non-square matrix, may need to use 2d version"
N = np.sum(blocksizes)
assert A.shape[0] == N, "bad blocksizes"
num_blocks = len(blocksizes)
if pagesizes == None:
#Assume each page of A should be its own dimension
pagesizes = range(A.shape[2])#[0]*A.shape[2]
num_pages = len(pagesizes)
B = np.empty(shape=(num_blocks, num_blocks, num_pages), dtype=np.object)
count_k = 0
#for Bk, k in enumerate(pagesizes):
for Bk in pagesizes:
count_i = 0
for Bi, i in enumerate(blocksizes):
count_j = 0
for Bj, j in enumerate(blocksizes):
#We want to have it count_k:count_k + k but its annoying as it makes a NxNx1 array is page sizes are set to 1
B[Bi, Bj, Bk] = A[count_i:count_i + i, count_j:count_j + j, Bk]
count_j += j
count_i += i
#count_k += k
return B
[docs]def get_blocks(A, blocksizes):
assert (A.shape[0]==A.shape[1]) and len(A.shape)==2, "can't blockify this non-square matrix"
N = np.sum(blocksizes)
assert A.shape[0] == N, "bad blocksizes"
num_blocks = len(blocksizes)
B = np.empty(shape=(num_blocks, num_blocks), dtype=np.object)
count_i = 0
for Bi, i in enumerate(blocksizes):
count_j = 0
for Bj, j in enumerate(blocksizes):
B[Bi, Bj] = A[count_i:count_i + i, count_j : count_j + j]
count_j += j
count_i += i
return B
[docs]def get_block_shapes_3d(B):
assert B.dtype is np.dtype('object'), "Must be a block matrix"
#FIXME: This isn't general AT ALL...
return get_block_shapes(B[:,:,0]), B.shape[2]
[docs]def get_block_shapes(B):
assert B.dtype is np.dtype('object'), "Must be a block matrix"
return [B[b,b].shape[0] for b in range(0, B.shape[0])]
[docs]def unblock(B):
assert B.dtype is np.dtype('object'), "Must be a block matrix"
block_shapes = get_block_shapes(B)
num_elements = np.sum(block_shapes)
A = np.empty(shape=(num_elements, num_elements))
count_i = 0
for Bi, i in enumerate(block_shapes):
count_j = 0
for Bj, j in enumerate(block_shapes):
A[count_i:count_i + i, count_j:count_j + j] = B[Bi, Bj]
count_j += j
count_i += i
return A
[docs]def block_dot(A, B, diagonal=False):
"""
Element wise dot product on block matricies
+------+------+ +------+------+ +-------+-------+
| | | | | | |A11.B11|B12.B12|
| A11 | A12 | | B11 | B12 | | | |
+------+------+ o +------+------| = +-------+-------+
| | | | | | |A21.B21|A22.B22|
| A21 | A22 | | B21 | B22 | | | |
+-------------+ +------+------+ +-------+-------+
..Note
If any block of either (A or B) are stored as 1d vectors then we assume
that it denotes a diagonal matrix efficient dot product using numpy
broadcasting will be used, i.e. A11*B11
If either (A or B) of the diagonal matrices are stored as vectors then a more
efficient dot product using numpy broadcasting will be used, i.e. A11*B11
"""
#Must have same number of blocks and be a block matrix
assert A.dtype is np.dtype('object'), "Must be a block matrix"
assert B.dtype is np.dtype('object'), "Must be a block matrix"
assert A.shape == B.shape
def f(C,D):
"""
C is an element of A, D is the associated element of B
"""
Cshape = C.shape
Dshape = D.shape
if diagonal and (len(Cshape) == 1 or len(Dshape) == 1\
or C.shape[0] != C.shape[1] or D.shape[0] != D.shape[1]):
print("Broadcasting, C: {} D:{}".format(C.shape, D.shape))
return C*D
else:
print("Dotting, C: {} C:{}".format(C.shape, D.shape))
return np.dot(C,D)
dot = np.vectorize(f, otypes = [np.object])
return dot(A,B)
if __name__=='__main__':
A = np.zeros((5,5))
B = get_blocks(A,[2,3])
B[0,0] += 7
print(B)
assert np.all(unblock(B) == A)