def test_multiply():
x = sparse.brandom((4, 2, 6), (2, 1, 2), 0.5, format='bcoo')
x_d = x.todense()
y = sparse.brandom((4, 2, 6), (2, 1, 2), 0.5, format='bcoo')
y_d = y.todense()
z = x * y
z_d = x_d * y_d
assert_eq(z, z_d)
data = np.arange(1, 4).repeat(4).reshape(3, 2, 2)
coords = np.array([[1, 1, 0], [0, 1, 1]])
x = BCOO(coords, data=data, shape=(4, 4), block_shape=(2, 2))
x_d = x.todense()
y = x * x
y_d = x_d * x_d
assert_eq(y, y_d)
data_x = np.arange(1, 4).repeat(4).reshape(3, 2, 2)
coords_x = np.array([[1, 1, 0], [0, 1, 1]])
x = BCOO(coords_x, data=data_x, shape=(4, 4), block_shape=(2, 2))
x_d = x.todense()
data_y = np.array([[[-2.0, -2.5], [-2, -2]], [[-2, -2], [-2, -2]]])
coords_y = np.array([[0, 1], [0, 1]])
y = BCOO(coords_y, data=data_y, shape=(4, 4), block_shape=(2, 2))
y_d = y.todense()
z = x * y
z_d = x_d * y_d
assert_eq(z, z_d)
def test_add_mismatch_block():
data_x = np.arange(1, 4).repeat(4).reshape(3, 2, 2)
coords_x = np.array([[1, 1, 0], [0, 1, 1]])
x = BCOO(coords_x, data=data_x, shape=(4, 4), block_shape=(2, 2))
x_d = x.todense()
data_y = np.array([[[-2.0, -2.5], [-2, -2]], [[-2, -2], [-2, -2]]])
coords_y = np.array([[0, 1], [0, 1]])
y = BCOO(coords_y, data=data_y, shape=(4, 4), block_shape=(2, 2))
y_d = y.todense()
z = x + y
z_d = x_d + y_d
assert_eq(z, z_d)
def test_tobsr():
data = np.arange(1, 7).repeat(4).reshape((-1, 2, 2))
coords = np.array([[0, 0, 0, 2, 1, 2], [0, 1, 1, 0, 2, 2]])
block_shape = (2, 2)
shape = (8, 6)
print coords
print data
#x = BCOO(coords, data=data, shape=shape, block_shape=block_shape, sorted = True, has_duplicates = False)
x = BCOO(coords,
data=data,
shape=shape,
block_shape=block_shape,
sorted=True,
has_duplicates=False)
y = x.todense()
#print y
#print x.coords
z = x.tobsr()
#print z.data
#print z.has_canonical_format
#print z.has_sorted_indices
x2 = BCOO(coords.copy(),
data=data.copy(),
shape=shape,
block_shape=block_shape,
sorted=False,
has_duplicates=True)
y2 = x2.todense()
print y
print y2
print x.coords
print x2.coords
#print y
#print x.coords
z2 = x2.tobsr()
#print z2.data
print z.has_canonical_format
print z.has_sorted_indices
print z2.has_canonical_format
print z2.has_sorted_indices
#print "z"
print z.toarray()
#print "z2"
print z2.toarray()
#assert_eq(z,z2)
exit()
assert_eq(z, y)
def test_invalid_data_input():
data = np.array([[-1, -2.5], [-3, -4]])
coords = np.array([[0, 1], [0, 1]])
block_shape = (2, 2)
shape = (4, 4)
with pytest.raises(AssertionError):
x = BCOO(coords, data=data, shape=shape, block_shape=block_shape)
def test_add_simple():
data = np.arange(1, 4).repeat(4).reshape(3, 2, 2)
coords = np.array([[1, 1, 0], [0, 1, 1]])
x = BCOO(coords, data=data, shape=(4, 4), block_shape=(2, 2))
x_d = x.todense()
y = x + x
y_d = x_d + x_d
assert_eq(y, y_d)
def test_tobsr():
data = np.arange(1, 7).repeat(4).reshape((-1, 2, 2))
coords = np.array([[0, 0, 0, 2, 1, 2], [0, 1, 1, 0, 2, 2]])
block_shape = (2, 2)
shape = (8, 6)
x = BCOO(coords, data=data, shape=shape, block_shape=block_shape)
y = x.todense()
z = x.tobsr()
assert_eq(z, y)
def test_create_with_lists_of_tuples():
L = [((0, 0, 0), np.random.random((2, 4, 3))),
((1, 2, 1), np.random.random((2, 4, 3))),
((1, 1, 1), np.random.random((2, 4, 3))),
((1, 3, 2), np.random.random((2, 4, 3)))]
s = BCOO(L, block_shape=(2, 4, 3))
x = np.zeros((2, 4, 3, 2, 4, 3))
for ind, value in L:
x[ind] = value
x = x.transpose(0, 3, 1, 4, 2, 5).reshape(2 * 2, 4 * 4, 3 * 3)
assert_eq(s, x)
def test_subtraction():
x = sparse.brandom((4, 2, 6), (2, 1, 2), 0.5, format='bcoo')
x_d = x.todense()
y = sparse.brandom((4, 2, 6), (2, 1, 2), 0.5, format='bcoo')
y_d = y.todense()
z = x - y
z_d = x_d - y_d
assert_eq(z, z_d)
data = np.arange(1, 4).repeat(4).reshape(3, 2, 2)
coords = np.array([[1, 1, 0], [0, 1, 1]])
x = BCOO(coords, data=data, shape=(4, 4), block_shape=(2, 2))
x_d = x.todense()
y = x - x
y_d = x_d - x_d
assert_eq(y, y_d)
def test_add_zero():
# the result is zero
data_x = np.arange(3, 6).repeat(4).reshape(3, 2, 2).astype(np.double)
coords_x = np.array([[1, 1, 0], [0, 1, 1]])
x = BCOO(coords_x, data=data_x, shape=(4, 4), block_shape=(2, 2))
x_d = x.todense()
y = -x
y_d = y.todense()
z = x + y
z_d = x_d + y_d
assert_eq(z, z_d)
# all add numbers are zero
a = np.zeros((6, 5, 4, 1), dtype=np.complex)
x = BCOO.from_numpy(a, block_shape=(2, 5, 2, 1))
x_d = x.todense()
y = BCOO.from_numpy(a, block_shape=(2, 5, 2, 1))
y_d = y.todense()
z = x + y
z_d = x_d + y_d
assert_eq(z, z_d)
def test_scaling():
data_x = np.arange(3, 6).repeat(4).reshape(3, 2, 2).astype(np.double)
coords_x = np.array([[1, 1, 0], [0, 1, 1]])
x = BCOO(coords_x, data=data_x, shape=(4, 4), block_shape=(2, 2))
x_d = x.todense()
z = x * -3.1
z_d = x_d * -3.1
assert_eq(z, z_d)
x = sparse.brandom((4, 2, 6), (2, 1, 2), 0.5, format='bcoo')
x_d = x.todense()
scaling_factor = np.random.random()
z = x * scaling_factor
z_d = x_d * scaling_factor
assert_eq(z, z_d)
x = sparse.brandom((4, 2, 6), (2, 1, 2), 0.5, format='bcoo')
x_d = x.todense()
scaling_factor = 0
z = x * scaling_factor
z_d = x_d * scaling_factor
assert_eq(z, z_d)
def einsum(idx_str, *tensors, **kwargs):
'''Perform a more efficient einsum via reshaping to a matrix multiply.
Current differences compared to numpy.einsum:
This assumes that each repeated index is actually summed (i.e. no 'i,i->i')
and appears only twice (i.e. no 'ij,ik,il->jkl'). The output indices must
be explicitly specified (i.e. 'ij,j->i' and not 'ij,j').
'''
DEBUG = kwargs.get('DEBUG', False)
idx_str = idx_str.replace(' ', '')
indices = "".join(re.split(',|->', idx_str))
if '->' not in idx_str or any(indices.count(x) > 2 for x in set(indices)):
#return np.einsum(idx_str,*tensors)
raise NotImplementedError
if idx_str.count(',') > 1:
indices = re.split(',|->', idx_str)
indices_in = indices[:-1]
idx_final = indices[-1]
n_shared_max = 0
for i in range(len(indices_in)):
for j in range(i):
tmp = list(set(indices_in[i]).intersection(indices_in[j]))
n_shared_indices = len(tmp)
if n_shared_indices > n_shared_max:
n_shared_max = n_shared_indices
shared_indices = tmp
[a, b] = [i, j]
tensors = list(tensors)
A, B = tensors[a], tensors[b]
idxA, idxB = indices[a], indices[b]
idx_out = list(idxA + idxB)
idx_out = "".join([x for x in idx_out if x not in shared_indices])
C = einsum(idxA + "," + idxB + "->" + idx_out, A, B)
indices_in.pop(a)
indices_in.pop(b)
indices_in.append(idx_out)
tensors.pop(a)
tensors.pop(b)
tensors.append(C)
return einsum(",".join(indices_in) + "->" + idx_final, *tensors)
A, B = tensors
# Call numpy.asarray because A or B may be HDF5 Datasets
# A = numpy.asarray(A, order='A')
# B = numpy.asarray(B, order='A')
# if A.size < 2000 or B.size < 2000:
# return numpy.einsum(idx_str, *tensors)
# Split the strings into a list of idx char's
idxA, idxBC = idx_str.split(',')
idxB, idxC = idxBC.split('->')
idxA, idxB, idxC = [list(x) for x in [idxA, idxB, idxC]]
assert (len(idxA) == A.ndim)
assert (len(idxB) == B.ndim)
if DEBUG:
print("*** Einsum for", idx_str)
print(" idxA =", idxA)
print(" idxB =", idxB)
print(" idxC =", idxC)
# Get the range for each index and put it in a dictionary
rangeA = dict()
rangeB = dict()
block_rangeA = dict()
block_rangeB = dict()
for idx, rnge in zip(idxA, A.outer_shape): # ZHC NOTE
rangeA[idx] = rnge
for idx, rnge in zip(idxB, B.outer_shape):
rangeB[idx] = rnge
for idx, rnge in zip(idxA, A.block_shape):
block_rangeA[idx] = rnge
for idx, rnge in zip(idxB, B.block_shape):
block_rangeB[idx] = rnge
if DEBUG:
print("rangeA =", rangeA)
print("rangeB =", rangeB)
print("block_rangeA =", block_rangeA)
print("block_rangeB =", block_rangeB)
# Find the shared indices being summed over
shared_idxAB = list(set(idxA).intersection(idxB))
#if len(shared_idxAB) == 0:
# return np.einsum(idx_str,A,B)
idxAt = list(idxA)
idxBt = list(idxB)
inner_shape = 1
block_inner_shape = 1
insert_B_loc = 0
for n in shared_idxAB:
if rangeA[n] != rangeB[n]:
err = ('ERROR: In index string %s, the range of index %s is '
'different in A (%d) and B (%d)' %
(idx_str, n, rangeA[n], rangeB[n]))
raise RuntimeError(err)
# Bring idx all the way to the right for A
# and to the left (but preserve order) for B
idxA_n = idxAt.index(n)
idxAt.insert(len(idxAt) - 1, idxAt.pop(idxA_n))
idxB_n = idxBt.index(n)
idxBt.insert(insert_B_loc, idxBt.pop(idxB_n))
insert_B_loc += 1
inner_shape *= rangeA[n]
block_inner_shape *= block_rangeA[n]
if DEBUG:
print("shared_idxAB =", shared_idxAB)
print("inner_shape =", inner_shape)
print("block_inner_shape =", block_inner_shape)
# Transpose the tensors into the proper order and reshape into matrices
new_orderA = [idxA.index(idx) for idx in idxAt]
new_orderB = [idxB.index(idx) for idx in idxBt]
if DEBUG:
print("Transposing A as", new_orderA)
print("Transposing B as", new_orderB)
print("Reshaping A as (-1,", inner_shape, ")")
print("Reshaping B as (", inner_shape, ",-1)")
print("Reshaping block A as (-1,", block_inner_shape, ")")
print("Reshaping block B as (", block_inner_shape, ",-1)")
shapeCt = list()
block_shapeCt = list()
idxCt = list()
for idx in idxAt:
if idx in shared_idxAB:
break
shapeCt.append(rangeA[idx])
block_shapeCt.append(block_rangeA[idx])
idxCt.append(idx)
for idx in idxBt:
if idx in shared_idxAB:
continue
shapeCt.append(rangeB[idx])
block_shapeCt.append(block_rangeB[idx])
idxCt.append(idx)
new_orderCt = [idxCt.index(idx) for idx in idxC]
np_shapeCt = tuple(np.multiply(shapeCt, block_shapeCt))
if A.nnz == 0 or B.nnz == 0:
shapeCt = [shapeCt[i] for i in new_orderCt]
block_shapeCt = [block_shapeCt[i] for i in new_orderCt]
return BCOO(np.array([],dtype = np.int), data = np.array([], dtype = \
np.result_type(A.dtype,B.dtype)), shape=np_shapeCt,\
block_shape = block_shapeCt, has_duplicates=False,\
sorted=True).transpose(new_orderCt)
At = A.transpose(new_orderA)
Bt = B.transpose(new_orderB)
# ZHC TODO optimize
# if At.flags.f_contiguous:
# At = numpy.asarray(At.reshape((-1,inner_shape), (-1,block_inner_shape)), order='F')
# else:
At = At.block_reshape((-1, inner_shape),
block_shape=(-1, block_inner_shape))
# if Bt.flags.f_contiguous:
# Bt = numpy.asarray(Bt.reshape((inner_shape,-1), (block_inner_shape,-1)), order='F')
# else:
Bt = Bt.block_reshape((inner_shape, -1),
block_shape=(block_inner_shape, -1))
#AdotB = At.tobsr().dot(Bt.tobsr())
At = At.tobsr()
Bt = Bt.tobsr()
AdotB = At.dot(Bt)
AdotB_bcoo = BCOO.from_bsr(AdotB)
if DEBUG:
print("AdotB bsr format indptr, indices")
print(AdotB.indptr)
print(AdotB.indices)
print("AdotB bcoo format coords")
print(AdotB_bcoo.coords)
return AdotB_bcoo.block_reshape(
shapeCt, block_shape=block_shapeCt).transpose(new_orderCt)
# c = ab
# else:
# if beta == 0:
# c[:] = 0
# else:
# c *= beta
# c += ab
# return c
if __name__ == '__main__':
data_x = np.arange(1, 7).repeat(4).reshape((-1, 2, 2))
coords_x = np.array([[0, 0, 0, 2, 1, 2], [0, 1, 1, 0, 2, 2]])
shape_x = (8, 6)
block_shape_x = (2, 2)
x = BCOO(coords_x, data=data_x, shape=shape_x, block_shape=block_shape_x)
x_d = x.todense()
shape_y = (8, 4, 6)
block_shape_y = (2, 2, 2)
y = sparse.brandom(shape_y, block_shape_y, 0.5, format='bcoo')
y_d = y.todense()
c = einsum("ij,ikj->k", x, y, DEBUG=True)
elemC = np.einsum("ij,ikj->k", x_d, y_d)
# another test
'''
shape_x = (8,4,9)
block_shape_x = (1,2,3)
x = sparse.brandom(shape_x, block_shape_x, 0.2, format='bcoo')
def _contract(subscripts, *tensors, **kwargs):
DEBUG = kwargs.get('DEBUG', False)
idx_str = subscripts.replace(' ', '')
indices = idx_str.replace(',', '').replace('->', '')
if '->' not in idx_str or any(indices.count(x) > 2 for x in set(indices)):
# TODO 1. No ->, contract over repeated indices 2. more than 2 indices need to contract.
raise NotImplementedError
A, B = tensors
# mix type, transfer to dense case
if not (isinstance(A, BCOO) and isinstance(B, BCOO)):
print(
"Warning: the block einsum takes non-BCOO objects, try to transfer to dense..."
)
if hasattr(A, 'todense'):
A = A.todense()
if hasattr(B, 'todense'):
B = B.todense()
return np.einsum(idx_str, A, B)
# ZHC NOTE threshold to determine which lib to use here?
# Split the strings into a list of idx char's
idxA, idxBC = idx_str.split(',')
idxB, idxC = idxBC.split('->')
#idxA, idxB, idxC = [list(x) for x in [idxA,idxB,idxC]]
assert (len(idxA) == A.ndim)
assert (len(idxB) == B.ndim)
if DEBUG:
print("*** Einsum for", idx_str)
print(" idxA =", idxA)
print(" idxB =", idxB)
print(" idxC =", idxC)
# Get the range for each index and put it in a dictionary
rangeA = dict(zip(idxA, A.outer_shape))
rangeB = dict(zip(idxB, B.outer_shape))
block_rangeA = dict(zip(idxA, A.block_shape))
block_rangeB = dict(zip(idxB, B.block_shape))
if DEBUG:
print("rangeA =", rangeA)
print("rangeB =", rangeB)
print("block_rangeA =", block_rangeA)
print("block_rangeB =", block_rangeB)
# duplicated indices 'in,ijj->n' # TODO: first index out the repeated indices.
if len(rangeA) != A.ndim or len(rangeB) != B.ndim:
raise NotImplementedError
# Find the shared indices being summed over
shared_idxAB = list(set(idxA).intersection(idxB))
if len(shared_idxAB) == 0: # TODO Indices must overlap
raise NotImplementedError
idxAt = list(idxA)
idxBt = list(idxB)
inner_shape = 1
block_inner_shape = 1
insert_B_loc = 0
for n in shared_idxAB:
if rangeA[n] != rangeB[n]:
err = (
'ERROR: In index string %s, the outer_shape range of index %s is '
'different in A (%d) and B (%d)' %
(idx_str, n, rangeA[n], rangeB[n]))
raise ValueError(err)
if block_rangeA[n] != block_rangeB[n]:
err = (
'ERROR: In index string %s, the block_shape range of index %s is '
'different in A (%d) and B (%d)' %
(idx_str, n, block_rangeA[n], block_rangeB[n]))
raise RuntimeError(err)
# Bring idx all the way to the right for A
# and to the left (but preserve order) for B
idxA_n = idxAt.index(n)
idxAt.insert(len(idxAt) - 1, idxAt.pop(idxA_n))
idxB_n = idxBt.index(n)
idxBt.insert(insert_B_loc, idxBt.pop(idxB_n))
insert_B_loc += 1
inner_shape *= rangeA[n]
block_inner_shape *= block_rangeA[n]
if DEBUG:
print("shared_idxAB =", shared_idxAB)
print("inner_shape =", inner_shape)
print("block_inner_shape =", block_inner_shape)
# Transpose the tensors into the proper order and reshape into matrices
new_orderA = [idxA.index(idx) for idx in idxAt]
new_orderB = [idxB.index(idx) for idx in idxBt]
if DEBUG:
print("Transposing A as", new_orderA)
print("Transposing B as", new_orderB)
print("Reshaping A as (-1,", inner_shape, ")")
print("Reshaping B as (", inner_shape, ",-1)")
print("Reshaping block A as (-1,", block_inner_shape, ")")
print("Reshaping block B as (", block_inner_shape, ",-1)")
shapeCt = list()
block_shapeCt = list()
idxCt = list()
for idx in idxAt:
if idx in shared_idxAB:
break
shapeCt.append(rangeA[idx])
block_shapeCt.append(block_rangeA[idx])
idxCt.append(idx)
for idx in idxBt:
if idx in shared_idxAB:
continue
shapeCt.append(rangeB[idx])
block_shapeCt.append(block_rangeB[idx])
idxCt.append(idx)
new_orderCt = [idxCt.index(idx) for idx in idxC]
if A.nnz == 0 or B.nnz == 0:
shapeCt = [shapeCt[i] for i in new_orderCt]
block_shapeCt = [block_shapeCt[i] for i in new_orderCt]
np_shapeCt = tuple(np.multiply(shapeCt, block_shapeCt))
return BCOO(np.array([],dtype = np.int), data = np.array([], dtype = \
np.result_type(A.dtype,B.dtype)), shape=np_shapeCt,\
block_shape = block_shapeCt, has_duplicates=False,\
sorted=True)
At = A.transpose(new_orderA)
Bt = B.transpose(new_orderB)
At = At.block_reshape((-1, inner_shape),
block_shape=(-1, block_inner_shape))
Bt = Bt.block_reshape((inner_shape, -1),
block_shape=(block_inner_shape, -1))
#AdotB = At.tobsr().dot(Bt.tobsr())
At = At.tobsr()
Bt = Bt.tobsr()
AdotB = At.dot(Bt)
AdotB_bcoo = BCOO.from_bsr(AdotB)
if DEBUG:
print("AdotB bsr format indptr, indices")
print(AdotB.indptr)
print(AdotB.indices)
print("AdotB bcoo format coords")
print(AdotB_bcoo.coords)
return AdotB_bcoo.block_reshape(
shapeCt, block_shape=block_shapeCt).transpose(new_orderCt)