Block sparse

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()
 
 

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

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))