def test_split_node_qr_unitarity(dtype, num_charges): np.random.seed(10) a = tn.Node(get_square_matrix(50, num_charges, dtype=dtype), backend='symmetric') q, r = tn.split_node_qr(a, [a[0]], [a[1]]) r[0] | q[1] qbar = tn.conj(q) q[1] ^ qbar[1] u1 = q @ qbar qbar[0] ^ q[0] u2 = qbar @ q blocks, _, shapes = _find_diagonal_sparse_blocks(u1.tensor.flat_charges, u1.tensor.flat_flows, len(u1.tensor._order[0])) for n, block in enumerate(blocks): np.testing.assert_almost_equal( np.reshape(u1.tensor.data[block], shapes[:, n]), np.eye(N=shapes[0, n], M=shapes[1, n])) blocks, _, shapes = _find_diagonal_sparse_blocks(u2.tensor.flat_charges, u2.tensor.flat_flows, len(u2.tensor._order[0])) for n, block in enumerate(blocks): np.testing.assert_almost_equal( np.reshape(u2.tensor.data[block], shapes[:, n]), np.eye(N=shapes[0, n], M=shapes[1, n]))
def test_pinv(dtype, num_charges): np.random.seed(10) R = 2 D = 10 charge = BaseCharge(np.random.randint(-5, 6, (num_charges, D), dtype=np.int16), charge_types=[U1Charge] * num_charges) flows = [True, False] A = BlockSparseTensor.random([Index(charge, flows[n]) for n in range(R)], (-0.5, 0.5), dtype=dtype) invA = pinv(A) left_eye = invA @ A blocks, _, shapes = _find_diagonal_sparse_blocks(left_eye.flat_charges, left_eye.flat_flows, 1) for n, block in enumerate(blocks): t = np.reshape(left_eye.data[block], shapes[:, n]) assert np.linalg.norm(t - np.eye(t.shape[0], t.shape[1])) < 1E-12 right_eye = A @ invA blocks, _, shapes = _find_diagonal_sparse_blocks(right_eye.flat_charges, right_eye.flat_flows, 1) for n, block in enumerate(blocks): t = np.reshape(right_eye.data[block], shapes[:, n]) assert np.linalg.norm(t - np.eye(t.shape[0], t.shape[1])) < 1E-12
def test_get_diag(dtype, num_charges, Ds, flow): np.random.seed(10) np_flow = -np.int((np.int(flow) - 0.5) * 2) indices = [ Index( BaseCharge( np.random.randint(-2, 3, (num_charges, Ds[n])), charge_types=[U1Charge] * num_charges), flow) for n in range(2) ] arr = BlockSparseTensor.random(indices, dtype=dtype) fused = fuse_charges(arr.flat_charges, arr.flat_flows) inds = np.nonzero(fused == np.zeros((num_charges, 1), dtype=np.int16))[0] # pylint: disable=no-member left, _ = np.divmod(inds, Ds[1]) unique = np.unique( np_flow * (indices[0]._charges[0].charges[:, left]), axis=1) diagonal = diag(arr) sparse_blocks, _, block_shapes = _find_diagonal_sparse_blocks( arr.flat_charges, arr.flat_flows, 1) data = np.concatenate([ np.diag(np.reshape(arr.data[sparse_blocks[n]], block_shapes[:, n])) for n in range(len(sparse_blocks)) ]) np.testing.assert_allclose(data, diagonal.data) np.testing.assert_allclose(unique, diagonal.flat_charges[0].unique_charges)
def outerproduct(tensor1: BlockSparseTensor, tensor2: BlockSparseTensor) -> BlockSparseTensor: """ Compute the outer product of two `BlockSparseTensor`. The first `tensor1.ndim` indices of the resulting tensor are the indices of `tensor1`, the last `tensor2.ndim` indices are those of `tensor2`. Args: tensor1: A tensor. tensor2: A tensor. Returns: BlockSparseTensor: The result of taking the outer product. """ final_charges = tensor1._charges + tensor2._charges final_flows = tensor1.flat_flows + tensor2.flat_flows order2 = [list(np.asarray(s) + len(tensor1._charges)) for s in tensor2._order] data = np.zeros( compute_num_nonzero(final_charges, final_flows), dtype=tensor1.dtype) if ((len(tensor1.data) > 0) and (len(tensor2.data) > 0)) and (len(data) > 0): # find the location of the zero block in the output final_block_maps, final_block_charges, _ = _find_diagonal_sparse_blocks( final_charges, final_flows, len(tensor1._charges)) index = np.nonzero( final_block_charges == final_block_charges.identity_charges)[0][0] data[final_block_maps[index].ravel()] = np.outer(tensor1.data, tensor2.data).ravel() return BlockSparseTensor( data, charges=final_charges, flows=final_flows, order=tensor1._order + order2, check_consistency=False)
def contiguous(self, permutation: Optional[Union[Tuple, List, np.ndarray]] = None, inplace: Optional[bool] = False) -> Any: """ Transpose the tensor data in place such that the linear order of the elements in `BlockSparseTensor.data` corresponds to the current order of tensor indices. Consider a tensor with current order given by `_order=[[1,2],[3],[0]]`, i.e. `data` was initialized according to order [0,1,2,3], and the tensor has since been reshaped and transposed. The linear oder of `data` does not match the desired order [1,2,3,0] of the tensor. `contiguous` fixes this by permuting `data` into this order, transposing `_charges` and `_flows`, and changing `_order` to `[[0,1],[2],[3]]`. Args: permutation: An optional alternative order to be used to transposed the tensor. If `None` defaults to `BlockSparseTensor.permutation`. """ flat_charges = self._charges flat_flows = self._flows if permutation is None: permutation = self.flat_order if np.array_equal(permutation, np.arange(len(permutation))): return self tr_partition = _find_best_partition( [flat_charges[n].dim for n in permutation]) tr_sparse_blocks, tr_charges, _ = _find_transposed_diagonal_sparse_blocks( flat_charges, flat_flows, tr_partition, permutation) sparse_blocks, charges, _ = _find_diagonal_sparse_blocks( [flat_charges[n] for n in permutation], [flat_flows[n] for n in permutation], tr_partition) data = np.empty(len(self.data), dtype=self.dtype) for n, sparse_block in enumerate(sparse_blocks): ind = np.nonzero(tr_charges == charges[n])[0][0] perm = tr_sparse_blocks[ind] data[sparse_block] = self.data[perm] _, inds = np.unique(permutation, return_index=True) new_flat_order = inds[self.flat_order] tmp = np.append(0, np.cumsum([len(o) for o in self._order])) order = [ list(new_flat_order[tmp[n]:tmp[n + 1]]) for n in range(len(tmp) - 1) ] charges = [self._charges[o] for o in permutation] flows = [self._flows[o] for o in permutation] if not inplace: return BlockSparseTensor(data, charges=charges, flows=flows, order=order, check_consistency=False) self.data = data self._order = order self._charges = charges self._flows = flows return self
def test_eye(dtype, num_charges, D): charge = BaseCharge( np.random.randint(-5, 6, (num_charges, D), dtype=np.int16), charge_types=[U1Charge] * num_charges) flow = False index = Index(charge, flow) A = eye(index, dtype=dtype) blocks, _, shapes = _find_diagonal_sparse_blocks(A.flat_charges, A.flat_flows, 1) for n, block in enumerate(blocks): t = np.reshape(A.data[block], shapes[:, n]) np.testing.assert_almost_equal(t, np.eye(t.shape[0], t.shape[1]))
def test_find_diagonal_sparse_blocks(num_legs, num_charges): np.random.seed(10) np_charges = [ np.random.randint(-5, 5, (60, num_charges), dtype=np.int16) for _ in range(num_legs) ] fused = np.stack([ fuse_ndarrays([np_charges[n][:, c] for n in range(num_legs)]) for c in range(num_charges) ], axis=1) left_charges = np.stack([ fuse_ndarrays([np_charges[n][:, c] for n in range(num_legs // 2)]) for c in range(num_charges) ], axis=1) right_charges = np.stack([ fuse_ndarrays( [np_charges[n][:, c] for n in range(num_legs // 2, num_legs)]) for c in range(num_charges) ], axis=1) #pylint: disable=no-member nz = np.nonzero( np.logical_and.reduce(fused == np.zeros((1, num_charges)), axis=1))[0] linear_locs = np.arange(len(nz)) # pylint: disable=no-member left_inds, _ = np.divmod(nz, right_charges.shape[0]) left = left_charges[left_inds, :] unique_left = np.unique(left, axis=0) blocks = [] for n in range(unique_left.shape[0]): ul = unique_left[n, :][None, :] #pylint: disable=no-member blocks.append(linear_locs[np.nonzero( np.logical_and.reduce(left == ul, axis=1))[0]]) charges = [ BaseCharge(left_charges, charge_types=[U1Charge] * num_charges), BaseCharge(right_charges, charge_types=[U1Charge] * num_charges) ] print(left_charges) print(right_charges) bs, cs, ss = _find_diagonal_sparse_blocks(charges, [False, False], 1) np.testing.assert_allclose(cs.charges, unique_left) for b1, b2 in zip(blocks, bs): assert np.all(b1 == b2) assert np.sum(np.prod(ss, axis=0)) == np.sum([len(b) for b in bs]) np.testing.assert_allclose(unique_left, cs.charges)
def test_create_diag(dtype, num_charges): np.random.seed(10) D = 200 index = Index( BaseCharge(np.random.randint(-2, 3, (num_charges, D)), charge_types=[U1Charge] * num_charges), False) arr = ChargeArray.random([index], dtype=dtype) diagarr = diag(arr) dense = np.ravel(diagarr.todense()) np.testing.assert_allclose(np.sort(dense[dense != 0.0]), np.sort(diagarr.data[diagarr.data != 0.0])) sparse_blocks, charges, block_shapes = _find_diagonal_sparse_blocks( diagarr.flat_charges, diagarr.flat_flows, 1) #in range(index._charges[0].unique_charges.shape[1]): for n, block in enumerate(sparse_blocks): shape = block_shapes[:, n] block_diag = np.diag(np.reshape(diagarr.data[block], shape)) np.testing.assert_allclose( arr.data[np.squeeze(index._charges[0] == charges[n])], block_diag)
def eye(column_index: Index, row_index: Optional[Index] = None, dtype: Optional[Type[np.number]] = None) -> BlockSparseTensor: """ Return an identity matrix. Args: column_index: The column index of the matrix. row_index: The row index of the matrix. dtype: The dtype of the matrix. Returns: BlockSparseTensor """ if row_index is None: row_index = column_index.copy().flip_flow() if dtype is None: dtype = np.float64 blocks, _, shapes = _find_diagonal_sparse_blocks( column_index.flat_charges + row_index.flat_charges, column_index.flat_flows + row_index.flat_flows, len(column_index.flat_charges)) data = np.empty(np.int64(np.sum(np.prod(shapes, axis=0))), dtype=dtype) for n, block in enumerate(blocks): data[block] = np.ravel(np.eye(shapes[0, n], shapes[1, n], dtype=dtype)) order = [list(np.arange(0, len(column_index.flat_charges)))] + [ list( np.arange( len(column_index.flat_charges), len(column_index.flat_charges) + len(row_index.flat_charges))) ] return BlockSparseTensor( data=data, charges=column_index.flat_charges + row_index.flat_charges, flows=column_index.flat_flows + row_index.flat_flows, order=order, check_consistency=False)
def tensordot( tensor1: BlockSparseTensor, tensor2: BlockSparseTensor, axes: Optional[Union[Sequence[Sequence[int]], int]] = 2) -> BlockSparseTensor: """ Contract two `BlockSparseTensor`s along `axes`. Args: tensor1: First tensor. tensor2: Second tensor. axes: The axes to contract. Returns: BlockSparseTensor: The result of the tensor contraction. """ #process scalar input for `axes` if isinstance(axes, (np.integer, int)): axes = [ np.arange(tensor1.ndim - axes, tensor1.ndim, dtype=np.int16), np.arange(0, axes, dtype=np.int16) ] elif isinstance(axes[0], (np.integer, int)): if len(axes) > 1: raise ValueError( "invalid input `axes = {}` to tensordot".format(axes)) axes = [np.array(axes, dtype=np.int16), np.array(axes, dtype=np.int16)] axes1 = axes[0] axes2 = axes[1] if len(axes1) != len(axes2): raise ValueError( "`axes1 = {}` and `axes2 = {}` have to be of same length. ".format( axes1, axes2)) if len(axes1) > len(tensor1.shape): raise ValueError( "`axes1 = {}` is incompatible with `tensor1.shape = {}. ".format( axes1, tensor1.shape)) if len(axes2) > len(tensor2.shape): raise ValueError( "`axes2 = {}` is incompatible with `tensor2.shape = {}. ".format( axes2, tensor2.shape)) if not np.all(np.unique(axes1) == np.sort(axes1)): raise ValueError( "Some values in axes[0] = {} appear more than once!".format(axes1)) if not np.all(np.unique(axes2) == np.sort(axes2)): raise ValueError( "Some values in axes[1] = {} appear more than once!".format(axes2)) #special case outer product if len(axes1) == 0: return outerproduct(tensor1, tensor2) #more checks if max(axes1) >= len(tensor1.shape): raise ValueError( "rank of `tensor1` is smaller than `max(axes1) = {}.`".format( max(axes1))) if max(axes2) >= len(tensor2.shape): raise ValueError( "rank of `tensor2` is smaller than `max(axes2) = {}`".format( max(axes1))) contr_flows_1 = [] contr_flows_2 = [] contr_charges_1 = [] contr_charges_2 = [] for a in axes1: contr_flows_1.extend(tensor1._flows[tensor1._order[a]]) contr_charges_1.extend( [tensor1._charges[n] for n in tensor1._order[a]]) for a in axes2: contr_flows_2.extend(tensor2._flows[tensor2._order[a]]) contr_charges_2.extend( [tensor2._charges[n] for n in tensor2._order[a]]) if len(contr_charges_2) != len(contr_charges_1): raise ValueError( "`axes1 = {}` and `axes2 = {}` have incompatible elementary" " shapes {} and {}".format(axes1, axes2, [e.dim for e in contr_charges_1], [e.dim for e in contr_charges_2])) if not np.all( np.asarray(contr_flows_1) == np.logical_not( np.asarray(contr_flows_2))): raise ValueError( "`axes1 = {}` and `axes2 = {}` have incompatible elementary" " flows {} and {}".format(axes1, axes2, contr_flows_1, contr_flows_2)) charge_check = [ charge_equal(c1, c2) for c1, c2 in zip(contr_charges_1, contr_charges_2) ] if not np.all(charge_check): inds = np.nonzero(np.logical_not(charge_check))[0] raise ValueError( "`axes = {}` of tensor1 and `axes = {}` of tensor2 have incompatible charges" " {} and {}".format( np.array(axes1)[inds], np.array(axes2)[inds], [contr_charges_1[i] for i in inds], [contr_charges_2[i] for i in inds])) #checks finished #special case inner product if (len(axes1) == tensor1.ndim) and (len(axes2) == tensor2.ndim): t1 = tensor1.transpose(axes1).transpose_data() t2 = tensor2.transpose(axes2).transpose_data() data = np.dot(t1.data, t2.data) charge = tensor1._charges[0] final_charge = charge.__new__(type(charge)) final_charge.__init__(np.empty((charge.num_symmetries, 0), dtype=np.int16), charge_labels=np.empty(0, dtype=np.int16), charge_types=charge.charge_types) return BlockSparseTensor(data=data, charges=[final_charge], flows=[False], order=[[0]], check_consistency=False) #in all other cases we perform a regular tensordot free_axes1 = sorted(set(np.arange(tensor1.ndim)) - set(axes1)) free_axes2 = sorted(set(np.arange(tensor2.ndim)) - set(axes2)) new_order1 = [tensor1._order[n] for n in free_axes1] + [tensor1._order[n] for n in axes1] new_order2 = [tensor2._order[n] for n in axes2] + [tensor2._order[n] for n in free_axes2] flat_order_1 = flatten(new_order1) flat_order_2 = flatten(new_order2) flat_charges_1, flat_flows_1 = tensor1._charges, tensor1.flat_flows flat_charges_2, flat_flows_2 = tensor2._charges, tensor2.flat_flows left_charges = [] right_charges = [] left_flows = [] right_flows = [] left_order = [] right_order = [] s = 0 for n in free_axes1: left_charges.extend([tensor1._charges[o] for o in tensor1._order[n]]) left_order.append(list(np.arange(s, s + len(tensor1._order[n])))) s += len(tensor1._order[n]) left_flows.extend([tensor1._flows[o] for o in tensor1._order[n]]) s = 0 for n in free_axes2: right_charges.extend([tensor2._charges[o] for o in tensor2._order[n]]) right_order.append( list(len(left_charges) + np.arange(s, s + len(tensor2._order[n])))) s += len(tensor2._order[n]) right_flows.extend([tensor2._flows[o] for o in tensor2._order[n]]) tr_sparse_blocks_1, charges1, shapes_1 = _find_transposed_diagonal_sparse_blocks( flat_charges_1, flat_flows_1, len(left_charges), flat_order_1) tr_sparse_blocks_2, charges2, shapes_2 = _find_transposed_diagonal_sparse_blocks( flat_charges_2, flat_flows_2, len(contr_charges_2), flat_order_2) common_charges, label_to_common_1, label_to_common_2 = intersect( charges1.unique_charges, charges2.unique_charges, axis=1, return_indices=True) #Note: `cs` may contain charges that are not present in `common_charges` charges = left_charges + right_charges flows = left_flows + right_flows sparse_blocks, cs, _ = _find_diagonal_sparse_blocks( charges, flows, len(left_charges)) num_nonzero_elements = np.int64(np.sum([len(v) for v in sparse_blocks])) #Note that empty is not a viable choice here. data = np.zeros(num_nonzero_elements, dtype=np.result_type(tensor1.dtype, tensor2.dtype)) label_to_common_final = intersect(cs.unique_charges, common_charges, axis=1, return_indices=True)[1] for n in range(common_charges.shape[1]): n1 = label_to_common_1[n] n2 = label_to_common_2[n] nf = label_to_common_final[n] data[sparse_blocks[nf].ravel()] = np.ravel( np.matmul( tensor1.data[tr_sparse_blocks_1[n1].reshape(shapes_1[:, n1])], tensor2.data[tr_sparse_blocks_2[n2].reshape(shapes_2[:, n2])])) res = BlockSparseTensor(data=data, charges=charges, flows=flows, order=left_order + right_order, check_consistency=False) return res