Block sparse (#417)

mganahl · web-flow · commit c2399c88c8f7 · 2019-12-18T22:39:50.000-05:00
* started implementing block-sparse tensors * removed files * working on AbelianIndex * working in block sparisty * added reshape and lots of other stuff * added Index, an index type for symmetric tensors * added small tutorial * added docstring * fixed bug in retrieve_diagonal_blocks * TODO added * improved initialization a bit * more efficient initialization * just formatting * added random * added fuse_degeneracies * fixed bug in reshape * dosctring, typing * removed TODO * removed confusing code line * bug removed * comment * added __mul__ to Index * added sparse_shape and updated reshape to accept both int and Index lists * more in tutorial * comment * added new test function * testing function hacking * docstring * small speed up * Remove gui directory (migrated to another repo) (#399) * a slightly more elegant code * use one more np function * removed some crazy slow code * faster code * Update README.md (#404) * add return_data * doc * bug fix * a little faster * substantial speedup * renaming * removed todo * some comments * comments * fixed some bug in reshape * comments * default value changed * fixed bug, old version is now faster again * cleaned up reshape * started adding tests * replace kron with broadcasting * column-major -> row-major * documentation * added function to compute unique charges and charge degeneracies Function avoids explicit full fusion of all legs, and instead only keeps track of the unique charges and their degeneracies upon fusion * improved block finding, fixed bug in reshape re-intorduced BlockSparseTensor.dense_shape new method for fusing charges and degeneracies (faster for very rectangular matrices) * fuse_charge_pair added fuse_charges added * use is_leave * new tests * removed TODO, BlockSparseTensor.shape returns ref instead of copy * added tests * added tests * column-major -> row-major forgot to fix fusing order of charges and degeneracies * fix broken tests * test added * mostly docstring * docstring
diff --git a/tensornetwork/block_tensor/block_tensor.py b/tensornetwork/block_tensor/block_tensor.py
@@ -160,24 +160,35 @@ def compute_nonzero_block_shapes(charges: List[np.ndarray],
   return charge_shape_dict
 
 
-def retrieve_non_zero_diagonal_blocks_old_version(
+def retrieve_non_zero_diagonal_blocks(
     data: np.ndarray,
-    charges: List[np.ndarray],
-    flows: List[Union[bool, int]],
+    row_charges: List[Union[List, np.ndarray]],
+    column_charges: List[Union[List, np.ndarray]],
+    row_flows: List[Union[bool, int]],
+    column_flows: List[Union[bool, int]],
     return_data: Optional[bool] = True) -> Dict:
   """
-  Deprecated: this version is about 2 times slower (worst case) than the current used
-  implementation
   Given the meta data and underlying data of a symmetric matrix, compute 
   all diagonal blocks and return them in a dict.
+  `row_charges` and `column_charges` are lists of np.ndarray. The tensor
+  is viewed as a matrix with rows given by fusing `row_charges` and 
+  columns given by fusing `column_charges`. Note that `column_charges`
+  are never explicitly fused (`row_charges` are).
   Args: 
     data: An np.ndarray of the data. The number of elements in `data`
       has to match the number of non-zero elements defined by `charges` 
       and `flows`
-    charges: List of np.ndarray, one for each leg. 
-      Each np.ndarray `charges[leg]` is of shape `(D[leg],)`.
+    row_charges: List of np.ndarray, one for each leg of the row-indices.
+      Each np.ndarray `row_charges[leg]` is of shape `(D[leg],)`.
       The bond dimension `D[leg]` can vary on each leg.
-    flows: A list of integers, one for each leg,
+    column_charges: List of np.ndarray, one for each leg of the column-indices.
+      Each np.ndarray `row_charges[leg]` is of shape `(D[leg],)`.
+      The bond dimension `D[leg]` can vary on each leg.
+    row_flows: A list of integers, one for each entry in `row_charges`.
+      with values `1` or `-1`, denoting the flow direction
+      of the charges on each leg. `1` is inflowing, `-1` is outflowing
+      charge.
+    column_flows: A list of integers, one for each entry in `column_charges`.
       with values `1` or `-1`, denoting the flow direction
       of the charges on each leg. `1` is inflowing, `-1` is outflowing
       charge.
@@ -193,20 +204,25 @@ def retrieve_non_zero_diagonal_blocks_old_version(
     dict: Dictionary mapping quantum numbers (integers) to either an np.ndarray 
       or a python list of locations and shapes, depending on the value of `return_data`.
   """
-  if len(charges) != 2:
-    raise ValueError("input has to be a two-dimensional symmetric matrix")
+  flows = row_flows.copy()
+  flows.extend(column_flows)
   check_flows(flows)
-  if len(flows) != len(charges):
-    raise ValueError("`len(flows)` is different from `len(charges)`")
+  if len(flows) != (len(row_charges) + len(column_charges)):
+    raise ValueError(
+        "`len(flows)` is different from `len(row_charges) + len(column_charges)`"
+    )
 
-  #we multiply the flows into the charges
-  row_charges = flows[0] * charges[0]  # a list of charges on each row
-  column_charges = flows[1] * charges[1]  # a list of charges on each column
+  #since we are using row-major we have to fuse the row charges anyway.
+  fused_row_charges = fuse_charges(row_charges, row_flows)
+  #get the unique row-charges
+  unique_row_charges, row_dims = np.unique(
+      fused_row_charges, return_counts=True)
 
-  #get the unique charges
-  unique_row_charges, row_dims = np.unique(row_charges, return_counts=True)
-  unique_column_charges, column_dims = np.unique(
-      column_charges, return_counts=True)
+  #get the unique column-charges
+  #we only care about their degeneracies, not their order; that's much faster
+  #to compute since we don't have to fuse all charges explicitly
+  unique_column_charges, column_dims = compute_fused_charge_degeneracies(
+      column_charges, column_flows)
   #get the charges common to rows and columns (only those matter)
   common_charges = np.intersect1d(
       unique_row_charges, -unique_column_charges, assume_unique=True)
@@ -217,8 +233,8 @@ def retrieve_non_zero_diagonal_blocks_old_version(
   column_degeneracies = dict(zip(unique_column_charges, column_dims))
 
   # we only care about charges common to row and columns
-  mask = np.isin(row_charges, common_charges)
-  relevant_row_charges = row_charges[mask]
+  mask = np.isin(fused_row_charges, common_charges)
+  relevant_row_charges = fused_row_charges[mask]
 
   #some numpy magic to get the index locations of the blocks
   #we generate a vector of `len(relevant_row_charges) which,
@@ -261,35 +277,24 @@ def retrieve_non_zero_diagonal_blocks_old_version(
   return blocks
 
 
-def retrieve_non_zero_diagonal_blocks(
+def retrieve_non_zero_diagonal_blocks_old_version(
     data: np.ndarray,
-    row_charges: List[Union[List, np.ndarray]],
-    column_charges: List[Union[List, np.ndarray]],
-    row_flows: List[Union[bool, int]],
-    column_flows: List[Union[bool, int]],
+    charges: List[np.ndarray],
+    flows: List[Union[bool, int]],
     return_data: Optional[bool] = True) -> Dict:
   """
+  Deprecated: this version is about 2 times slower (worst case) than the current used
+  implementation
   Given the meta data and underlying data of a symmetric matrix, compute 
   all diagonal blocks and return them in a dict.
-  `row_charges` and `column_charges` are lists of np.ndarray. The tensor
-  is viewed as a matrix with rows given by fusing `row_charges` and 
-  columns given by fusing `column_charges`. Note that `column_charges`
-  are never explicitly fused (`row_charges` are).
   Args: 
     data: An np.ndarray of the data. The number of elements in `data`
       has to match the number of non-zero elements defined by `charges` 
       and `flows`
-    row_charges: List of np.ndarray, one for each leg of the row-indices.
-      Each np.ndarray `row_charges[leg]` is of shape `(D[leg],)`.
-      The bond dimension `D[leg]` can vary on each leg.
-    column_charges: List of np.ndarray, one for each leg of the column-indices.
-      Each np.ndarray `row_charges[leg]` is of shape `(D[leg],)`.
+    charges: List of np.ndarray, one for each leg. 
+      Each np.ndarray `charges[leg]` is of shape `(D[leg],)`.
       The bond dimension `D[leg]` can vary on each leg.
-    row_flows: A list of integers, one for each entry in `row_charges`.
-      with values `1` or `-1`, denoting the flow direction
-      of the charges on each leg. `1` is inflowing, `-1` is outflowing
-      charge.
-    column_flows: A list of integers, one for each entry in `column_charges`.
+    flows: A list of integers, one for each leg,
       with values `1` or `-1`, denoting the flow direction
       of the charges on each leg. `1` is inflowing, `-1` is outflowing
       charge.
@@ -305,25 +310,20 @@ def retrieve_non_zero_diagonal_blocks(
     dict: Dictionary mapping quantum numbers (integers) to either an np.ndarray 
       or a python list of locations and shapes, depending on the value of `return_data`.
   """
-  flows = row_flows.copy()
-  flows.extend(column_flows)
+  if len(charges) != 2:
+    raise ValueError("input has to be a two-dimensional symmetric matrix")
   check_flows(flows)
-  if len(flows) != (len(row_charges) + len(column_charges)):
-    raise ValueError(
-        "`len(flows)` is different from `len(row_charges) + len(column_charges)`"
-    )
+  if len(flows) != len(charges):
+    raise ValueError("`len(flows)` is different from `len(charges)`")
 
-  #since we are using row-major we have to fuse the row charges anyway.
-  fused_row_charges = fuse_charges(row_charges, row_flows)
-  #get the unique row-charges
-  unique_row_charges, row_dims = np.unique(
-      fused_row_charges, return_counts=True)
+  #we multiply the flows into the charges
+  row_charges = flows[0] * charges[0]  # a list of charges on each row
+  column_charges = flows[1] * charges[1]  # a list of charges on each column
 
-  #get the unique column-charges
-  #we only care about their degeneracies, not their order; that's much faster
-  #to compute since we don't have to fuse all charges explicitly
-  unique_column_charges, column_dims = compute_fused_charge_degeneracies(
-      column_charges, column_flows)
+  #get the unique charges
+  unique_row_charges, row_dims = np.unique(row_charges, return_counts=True)
+  unique_column_charges, column_dims = np.unique(
+      column_charges, return_counts=True)
   #get the charges common to rows and columns (only those matter)
   common_charges = np.intersect1d(
       unique_row_charges, -unique_column_charges, assume_unique=True)
@@ -334,8 +334,8 @@ def retrieve_non_zero_diagonal_blocks(
   column_degeneracies = dict(zip(unique_column_charges, column_dims))
 
   # we only care about charges common to row and columns
-  mask = np.isin(fused_row_charges, common_charges)
-  relevant_row_charges = fused_row_charges[mask]
+  mask = np.isin(row_charges, common_charges)
+  relevant_row_charges = row_charges[mask]
 
   #some numpy magic to get the index locations of the blocks
   #we generate a vector of `len(relevant_row_charges) which,
@@ -585,8 +585,8 @@ def compute_mapping_table(charges: List[np.ndarray],
   # is moving quickest when iterating through the linear data
   # transposing is done taking, for each value of the indices i_0 to i_N-2
   # the junk i_N-1 that gives non-zero
-  tables = np.meshgrid([np.arange(c.shape[0]) for c in charges], indexing='ij')
-  tables = tables[::-1]  #reverse the order
+
+  #for example
   raise NotImplementedError()
 
 
diff --git a/tensornetwork/block_tensor/block_tensor_test.py b/tensornetwork/block_tensor/block_tensor_test.py
@@ -0,0 +1,32 @@
+import numpy as np
+import pytest
+# pylint: disable=line-too-long
+from tensornetwork.block_tensor.block_tensor import BlockSparseTensor, compute_num_nonzero
+from index import Index
+
+np_dtypes = [np.float32, np.float16, np.float64, np.complex64, np.complex128]
+
+
+@pytest.mark.parametrize("dtype", np_dtypes)
+def test_block_sparse_init(dtype):
+  D = 10  #bond dimension
+  B = 10  #number of blocks
+  rank = 4
+  flows = np.asarray([1 for _ in range(rank)])
+  flows[-2::] = -1
+  charges = [
+      np.random.randint(-B // 2, B // 2 + 1, D).astype(np.int16)
+      for _ in range(rank)
+  ]
+  indices = [
+      Index(charges=charges[n], flow=flows[n], name='index{}'.format(n))
+      for n in range(rank)
+  ]
+  num_elements = compute_num_nonzero([i.charges for i in indices],
+                                     [i.flow for i in indices])
+  A = BlockSparseTensor.random(indices=indices, dtype=dtype)
+  assert A.dtype == dtype
+  for r in range(rank):
+    assert A.indices[r].name == 'index{}'.format(r)
+  assert A.dense_shape == tuple([D] * rank)
+  assert len(A.data) == num_elements
diff --git a/tensornetwork/block_tensor/index.py b/tensornetwork/block_tensor/index.py
@@ -130,9 +130,10 @@ def fuse_charge_pair(q1: Union[List, np.ndarray], flow1: int,
   for U(1) charges). `q1` and `q2` typically belong to two consecutive
   legs of `BlockSparseTensor`.
   Given `q1 = [0,1,2]` and `q2 = [10,100]`, this returns
-  `[10, 11, 12, 100, 101, 102]`.
+  `[10, 100, 11, 101, 12, 102]`.
   When using row-major ordering of indices in `BlockSparseTensor`, 
   the position of q1 should be "to the left" of the position of q2.
+
   Args:
     q1: Iterable of integers
     flow1: Flow direction of charge `q1`.
@@ -142,15 +143,17 @@ def fuse_charge_pair(q1: Union[List, np.ndarray], flow1: int,
     np.ndarray: The result of fusing `q1` with `q2`.
   """
   return np.reshape(
-      flow2 * np.asarray(q2)[:, None] + flow1 * np.asarray(q1)[None, :],
+      flow1 * np.asarray(q1)[:, None] + flow2 * np.asarray(q2)[None, :],
       len(q1) * len(q2))
 
 
 def fuse_charges(charges: List[Union[List, np.ndarray]],
                  flows: List[int]) -> np.ndarray:
   """
   Fuse all `charges` by simple addition (valid
-  for U(1) charges). 
+  for U(1) charges). Charges are fused from "right to left", 
+  in accordance with row-major order (see `fuse_charges_pair`).
+
   Args:
     chargs: A list of charges to be fused.
     flows: A list of flows, one for each element in `charges`.
@@ -173,19 +176,17 @@ def fuse_degeneracies(degen1: Union[List, np.ndarray],
   Fuse degeneracies `degen1` and `degen2` of two leg-charges 
   by simple kronecker product. `degen1` and `degen2` typically belong to two 
   consecutive legs of `BlockSparseTensor`.
-  Given `q1 = [0,1,2]` and `q2 = [10,100]`, this returns
-  `[10, 11, 12, 100, 101, 102]`.
+  Given `degen1 = [1, 2, 3]` and `degen2 = [10, 100]`, this returns
+  `[10, 100, 20, 200, 30, 300]`.
   When using row-major ordering of indices in `BlockSparseTensor`, 
-  the position of q1 should be "to the left" of the position of q2.
+  the position of `degen1` should be "to the left" of the position of `degen2`.
   Args:
-    q1: Iterable of integers
-    flow1: Flow direction of charge `q1`.
-    q2: Iterable of integers
-    flow2: Flow direction of charge `q2`.
+    degen1: Iterable of integers
+    degen2: Iterable of integers
   Returns:
-    np.ndarray: The result of fusing `q1` with `q2`.
+    np.ndarray: The result of fusing `dege1` with `degen2`.
   """
-  return np.reshape(degen2[:, None] * degen1[None, :],
+  return np.reshape(degen1[:, None] * degen2[None, :],
                     len(degen1) * len(degen2))
 
 
diff --git a/tensornetwork/block_tensor/index_test.py b/tensornetwork/block_tensor/index_test.py
