add fromdense to BlockSparseTensor (#595)

* add fromdense

* test added

* *** empty log message ***

Co-authored-by: Chase Roberts <chaseriley@google.com>

Author: mganahl

tensornetwork/block_sparse/blocksparsetensor.py

@@ -18,7 +18,7 @@
 import numpy as np
 from tensornetwork.block_sparse.index import Index
 # pylint: disable=line-too-long
-from tensornetwork.block_sparse.utils import _find_transposed_diagonal_sparse_blocks, _find_diagonal_sparse_blocks, flatten, get_flat_meta_data, compute_num_nonzero, _find_best_partition
+from tensornetwork.block_sparse.utils import _find_transposed_diagonal_sparse_blocks, _find_diagonal_sparse_blocks, flatten, get_flat_meta_data, compute_num_nonzero, _find_best_partition, reduce_charges
 from tensornetwork.block_sparse.charge import fuse_charges, BaseCharge, intersect, charge_equal
 import copy
 # pylint: disable=line-too-long
@@ -452,6 +452,48 @@ def copy(self) -> "BlockSparseTensor":
                              self._flows.copy(), copy.deepcopy(self._order),
                              False)
 
+  @classmethod
+  def fromdense(cls, indices: List[Index],
+                array: np.ndarray) -> "BlockSparseTensor":
+    """
+    Initialize a BlockSparseTensor from a dense array.
+    Args:
+      indices: A list of `Index` objects.
+      array: A numpy array.
+    Returns:
+      BlockSparseTensors: A Tensor initialized from the elements
+        of `array` at the positions where `indices` fuse to 
+        the identity charge.
+    """
+    shape = [i.dim for i in indices]
+    if not np.array_equal(shape, array.shape):
+      raise ValueError(
+          f"Cannot initialize an BlockSparseTensor of shape {shape}"
+          f" from an array of shape {array.shape}")
+    tmp = np.append(0, np.cumsum([len(i.flat_charges) for i in indices]))
+    order = [list(np.arange(tmp[n], tmp[n + 1])) for n in range(len(tmp) - 1)]
+
+    charges = []
+    flows = []
+    for i in indices:
+      charges.extend(i.flat_charges)
+      flows.extend(i.flat_flows)
+
+    _, locs = reduce_charges(
+        charges=charges,
+        flows=flows,
+        target_charges=charges[0].identity_charges.unique_charges,
+        return_locations=True)
+
+    ar = np.ravel(array)
+    data = ar[locs]
+    return cls(
+        data=data,
+        charges=charges,
+        flows=flows,
+        order=order,
+        check_consistency=False)
+
   def todense(self) -> np.ndarray:
     """
     Map the sparse tensor to dense storage.

tensornetwork/block_sparse/blocksparsetensor_test.py

@@ -399,6 +399,41 @@ def test_todense(num_charges, chargetype):
   np.testing.assert_allclose(dense[inds2], 0)
 
 
+@pytest.mark.parametrize('chargetype', ["U1", "Z2", "mixed"])
+@pytest.mark.parametrize('num_charges', [1, 2, 3])
+def test_fromdense(num_charges, chargetype):
+  np.random.seed(10)
+  Ds = [8, 9, 10, 11]
+  rank = 4
+  flows = np.random.choice([True, False], size=rank, replace=True)
+  charges = [get_charge(chargetype, num_charges, Ds[n]) for n in range(rank)]
+  fused = fuse_charges(charges, flows)
+  mask = fused == np.zeros((num_charges, 1))
+  inds = np.nonzero(mask)[0]
+  inds2 = np.nonzero(np.logical_not(mask))[0]
+  indices = [Index(charges[n], flows[n]) for n in range(rank)]
+
+  dense = np.random.random_sample(Ds)
+  arr = BlockSparseTensor.fromdense(indices, dense)
+  dense_arr = arr.todense()
+
+  np.testing.assert_allclose(np.ravel(dense)[inds], arr.data)
+  np.testing.assert_allclose(np.ravel(dense_arr)[inds2], 0)
+
+
+def test_fromdense_raises():
+  np.random.seed(10)
+  Ds = [8, 9, 10, 11]
+  rank = len(Ds)
+  flows = np.random.choice([True, False], size=rank, replace=True)
+  charges = [U1Charge.random(Ds[n], -5, 5) for n in range(rank)]
+  indices = [Index(charges[n], flows[n]) for n in range(rank)]
+
+  dense = np.random.random_sample([8, 9, 9, 11])
+  with pytest.raises(ValueError):
+    _ = BlockSparseTensor.fromdense(indices, dense)
+
+
 @pytest.mark.parametrize('chargetype', ["U1", "Z2", "mixed"])
 @pytest.mark.parametrize('num_charges', [1, 2, 3])
 @pytest.mark.parametrize('op', [np.add, np.subtract])