add svd

tensornetwork/block_tensor/block_tensor.py

@@ -744,6 +744,40 @@ def init_random():
 
     return cls(data=init_random(), indices=indices)
 
+  def __sub__(self, other: "BlockSparseTensor"):
+    if self.dense_shape != other.dense_shape:
+      raise ValueError("cannot subtract tensors with shapes {}and {}".format(
+          self.dense_shape, other.dense_shape))
+    if len(self.indices) != len(other.indices):
+      raise ValueError(
+          "cannot subtract tensors with different index-lengths {} and {}"
+          .format(len(self.indices), len(other.indices)))
+
+    if not np.all(
+        self.indices[n] == other.indices[n] for n in range(len(self.indices))):
+      raise ValueError("cannot subtract tensors non-matching indices")
+    return BlockSparseTensor(data=self.data - other.data, indices=self.indices)
+
+  def __add__(self, other: "BlockSparseTensor"):
+    if self.dense_shape != other.dense_shape:
+      raise ValueError("cannot add tensors with shapes {}and {}".format(
+          self.dense_shape, other.dense_shape))
+    if len(self.indices) != len(other.indices):
+      raise ValueError(
+          "cannot add tensors with different index-lengths {} and {}".format(
+              len(self.indices), len(other.indices)))
+
+    if not np.all(
+        self.indices[n] == other.indices[n] for n in range(len(self.indices))):
+      raise ValueError("cannot add tensors non-matching indices")
+    return BlockSparseTensor(data=self.data + other.data, indices=self.indices)
+
+  def __mul__(self, number: np.number):
+    return BlockSparseTensor(data=self.data * number, indices=self.indices)
+
+  def __rmul__(self, number: np.number):
+    return BlockSparseTensor(data=self.data * number, indices=self.indices)
+
   @property
   def rank(self):
     return len(self.indices)
@@ -770,14 +804,6 @@ def shape(self) -> Tuple:
   def dtype(self) -> Type[np.number]:
     return self.data.dtype
 
-  @property
-  def flows(self):
-    return [i.flow for i in self.indices]
-
-  @property
-  def charges(self):
-    return [i.charges for i in self.indices]
-
   @property
   def flat_charges(self):
     flat = []
@@ -957,7 +983,7 @@ def transpose(tensor: BlockSparseTensor,
   Returns:
     BlockSparseTensor: The transposed tensor.
   """
-  return tensor.transpose()
+  return tensor.transpose(order)
 
 
 def outerproduct(tensor1: BlockSparseTensor,
@@ -1176,3 +1202,78 @@ def tensordot(
               shapes_1[:, n1])] @ tensor2.data[tr_sparse_blocks_2[n2].reshape(
                   shapes_2[:, n2])]).ravel()
     return BlockSparseTensor(data=data, indices=free_indices)
+
+
+def svd(tensor: BlockSparseTensor,
+        full_matrices: Optional[bool] = True,
+        compute_uv: Optional[bool] = True,
+        hermitian: Optional[bool] = False):
+  if tensor.rank != 2:
+    raise NotImplementedError("SVD currently supports only rank-2 tensors.")
+
+  flat_charges = tensor.indices[0]._charges + tensor.indices[1]._charges
+  flat_flows = tensor.flat_flows
+  partition = len(tensor.indices[0].flat_charges)
+  blocks, charges, shapes = _find_diagonal_sparse_blocks(
+      flat_charges, flat_flows, partition)
+
+  u_blocks = []
+  singvals = []
+  v_blocks = []
+  for n in range(len(blocks)):
+    u, s, v = np.linalg.svd(
+        np.reshape(tensor.data[blocks[n]], shapes[:, n]), full_matrices,
+        compute_uv, hermitian)
+    u_blocks.append(u)
+    v_blocks.append(v)
+    singvals.append(np.diag(s))
+
+  #define the new charges on the two central bonds
+  new_left_charge = charges.__new__(type(charges))
+  new_right_charge = charges.__new__(type(charges))
+  left_charge_labels = np.concatenate([
+      np.full(u_blocks[n].shape[1], fill_value=n, dtype=np.int16)
+      for n in range(len(u_blocks))
+  ])
+  right_charge_labels = np.concatenate([
+      np.full(v_blocks[n].shape[0], fill_value=n, dtype=np.int16)
+      for n in range(len(v_blocks))
+  ])
+  left_singval_charge = charges.__new__(type(charges))
+  right_singval_charge = charges.__new__(type(charges))
+  left_singval_charge_labels = np.concatenate([
+      np.full(singvals[n].shape[0], fill_value=n, dtype=np.int16)
+      for n in range(len(singvals))
+  ])
+  right_singval_charge_labels = np.concatenate([
+      np.full(singvals[n].shape[1], fill_value=n, dtype=np.int16)
+      for n in range(len(singvals))
+  ])
+
+  new_left_charge.__init__(charges.unique_charges, left_charge_labels,
+                           charges.charge_types)
+  new_right_charge.__init__(charges.unique_charges, right_charge_labels,
+                            charges.charge_types)
+  left_singval_charge.__init__(charges.unique_charges,
+                               left_singval_charge_labels, charges.charge_types)
+  right_singval_charge.__init__(
+      charges.unique_charges, right_singval_charge_labels, charges.charge_types)
+
+  #get the indices of the new tensors U,S and V
+  indices_u = [Index(new_left_charge, True), tensor.indices[0]]
+  indices_v = [Index(new_right_charge, False), tensor.indices[1]]
+  indices_s = [
+      Index(left_singval_charge, False),
+      Index(right_singval_charge, True)
+  ]
+
+  #We fill in data into the transposed U
+  #TODO: reuse data from _find_diagonal_sparse_blocks above
+  #to avoid the transpose
+  return BlockSparseTensor(
+      np.concatenate([np.ravel(u.T) for u in u_blocks]), indices_u).transpose(
+          (1, 0)), BlockSparseTensor(
+              np.concatenate([np.ravel(s) for s in singvals]),
+              indices_s), BlockSparseTensor(
+                  np.concatenate([np.ravel(v) for v in v_blocks]),
+                  indices_v), u_blocks, v_blocks, charges

tensornetwork/block_tensor/index.py

@@ -24,21 +24,22 @@
 class Index:
   """
   An index class to store indices of a symmetric tensor.
-  An index keeps track of all its childs by storing references
-  to them (i.e. it is a binary tree).
   """
 
   def __init__(self,
                charges: Union[List[BaseCharge], BaseCharge],
                flow: Union[List[int], int],
                name: Optional[Union[List[Text], Text]] = None) -> None:
+    """
+    Initialize an `Index` object.
+    """
     if isinstance(charges, BaseCharge):
       charges = [charges]
     self._charges = charges
-    if isinstance(flow, (np.bool_, bool, np.bool)):
+    if np.isscalar(flow):
       flow = [flow]
     if not all([isinstance(f, (np.bool_, np.bool, bool)) for f in flow]):
-      raise TypeError("flows have to be boolean")
+      raise TypeError("flows have to be boolean. Found flow = {}".format(flow))
     self.flow = flow
     if isinstance(name, str):
       name = [name]
@@ -51,6 +52,18 @@ def __repr__(self):
   def dim(self):
     return np.prod([i.dim for i in self._charges])
 
+  def __eq__(self, other):
+    if len(other._charges) != len(self._charges):
+      return False
+    for n in range(len(self._charges)):
+      if not np.all(
+          self._charges[n].unique_charges == other._charges[n].unique_charges):
+        return False
+      if not np.all(
+          self._charges[n].charge_labels == other._charges[n].charge_labels):
+        return False
+    return True
+
   def copy(self):
     """
     Returns: