def test_tn_reshape(dtype): np.random.seed(10) Ds = [8, 9, 10, 11] indices = [ Index(U1Charge.random(dimension=Ds[n], minval=-5, maxval=5), False) for n in range(4) ] arr = BlockSparseTensor.random(indices, dtype=dtype) arr2 = reshape(arr, [72, 110]) for n in range(2): for m in range(2): assert charge_equal(arr2.charges[n][m], indices[n * 2 + m].charges) np.testing.assert_allclose(arr2.shape, [72, 110]) np.testing.assert_allclose(arr2._order, [[0, 1], [2, 3]]) np.testing.assert_allclose(arr2.flows, [[False, False], [False, False]]) assert arr2.ndim == 2 arr3 = reshape(arr, Ds) for n in range(4): assert charge_equal(arr3.charges[n][0], indices[n].charges) np.testing.assert_allclose(arr3.shape, Ds) np.testing.assert_allclose(arr3._order, [[0], [1], [2], [3]]) np.testing.assert_allclose(arr3.flows, [[False], [False], [False], [False]]) assert arr3.ndim == 4
def test_tn_transpose_reshape(): np.random.seed(10) Ds = np.array([8, 9, 10, 11]) flows = [True, False, True, False] indices = [Index(U1Charge.random(-5, 5, Ds[n]), flows[n]) for n in range(4)] arr = BlockSparseTensor.random(indices) arr2 = transpose(arr, [2, 0, 1, 3]) arr3 = reshape(arr2, [80, 99]) np.testing.assert_allclose(arr3.shape, [80, 99]) np.testing.assert_allclose(arr3._order, [[2, 0], [1, 3]]) np.testing.assert_allclose(arr3.flows, [[True, True], [False, False]]) arr4 = transpose(arr3, [1, 0]) np.testing.assert_allclose(arr4.shape, [99, 80]) np.testing.assert_allclose(arr4._order, [[1, 3], [2, 0]]) np.testing.assert_allclose(arr4.flows, [[False, False], [True, True]]) arr5 = reshape(arr4, [9, 11, 10, 8]) np.testing.assert_allclose(arr5.shape, [9, 11, 10, 8]) np.testing.assert_allclose(arr5._order, [[1], [3], [2], [0]]) np.testing.assert_allclose(arr5.flows, [[False], [False], [True], [True]])