style: reformatted by back, isort

tanlin2013 · Mar 2, 2023 · 3219df2 · 3219df2
1 parent 0cba9b6
commit 3219df2
Show file tree

Hide file tree

Showing 7 changed files with 89 additions and 138 deletions.
diff --git a/tests/model/test_random_heisenberg.py b/tests/model/test_random_heisenberg.py
@@ -1,29 +1,26 @@
-import pytest
 import numpy as np
+import pytest
 
 from tnpy.model import RandomHeisenberg
 from tnpy.operators import FullHamiltonian
 
 
 class TestRandomHeisenberg:
-
-    @pytest.fixture(scope='class')
+    @pytest.fixture(scope="class")
     def model(self):
         return RandomHeisenberg(n=6, h=0.5, seed=2022)
 
-    @pytest.fixture(scope='class')
+    @pytest.fixture(scope="class")
     def offset(self):
         return 0.5
 
-    @pytest.fixture(scope='class')
+    @pytest.fixture(scope="class")
     def shifted_model(self, model, offset):
-        return RandomHeisenberg(
-            n=model.n, h=model.h, seed=model.seed, offset=offset
-        )
+        return RandomHeisenberg(n=model.n, h=model.h, seed=model.seed, offset=offset)
 
     def test_offset(self, model, offset, shifted_model):
         np.testing.assert_allclose(
-            FullHamiltonian(model.mpo).matrix - offset * np.eye(2 ** model.n),
+            FullHamiltonian(model.mpo).matrix - offset * np.eye(2**model.n),
             FullHamiltonian(shifted_model.mpo).matrix,
-            atol=1e-12
+            atol=1e-12,
         )
diff --git a/tests/model/test_total_sz.py b/tests/model/test_total_sz.py
@@ -1,6 +1,5 @@
-import pytest
-
 import numpy as np
+import pytest
 
 from tnpy.model import TotalSz
 from tnpy.operators import FullHamiltonian

diff --git a/tests/test_exact_diagonalization.py b/tests/test_exact_diagonalization.py
@@ -1,9 +1,9 @@
-import pytest
 import numpy as np
+import pytest
 import scipy.linalg as spla
 
-from tnpy.model import RandomHeisenberg
 from tnpy.exact_diagonalization import ExactDiagonalization
+from tnpy.model import RandomHeisenberg
 from tnpy.operators import FullHamiltonian
 
 

diff --git a/tests/test_finite_dmrg.py b/tests/test_finite_dmrg.py
@@ -1,5 +1,5 @@
-import pytest
 import numpy as np
+import pytest
 
 from tnpy.exact_diagonalization import ExactDiagonalization
 from tnpy.finite_dmrg import FiniteDMRG, ShiftInvertDMRG
@@ -38,9 +38,7 @@ def ed(self, model):
 
     @pytest.fixture(scope="function")
     def sidmrg(self, model, offset):
-        shifted_model = RandomHeisenberg(
-            n=model.n, h=model.h, seed=model.seed, offset=offset
-        )
+        shifted_model = RandomHeisenberg(n=model.n, h=model.h, seed=model.seed, offset=offset)
         return ShiftInvertDMRG(shifted_model.mpo, bond_dim=2**6, offset=offset)
 
     @pytest.fixture(scope="class")
@@ -61,9 +59,7 @@ def test_run(self, ed, sidmrg, nearest_eval):
     def test_restored_mps(self, sidmrg, nearest_evec):
         sidmrg.run(tol=1e-8)
         restored_vec = (
-            sidmrg.restored_mps.contract()
-            .fuse({"k": sidmrg.restored_mps.outer_inds()})
-            .data
+            sidmrg.restored_mps.contract().fuse({"k": sidmrg.restored_mps.outer_inds()}).data
         )
         # Note: They can differ up to a global phase (+1 or -1)
         if not np.allclose(restored_vec, nearest_evec, atol=1e-6):

diff --git a/tests/test_matrix_product_state.py b/tests/test_matrix_product_state.py
@@ -1,11 +1,11 @@
 import os
 from contextlib import nullcontext as does_not_raise
 
-import pytest
 import numpy as np
+import pytest
 
+from tnpy.matrix_product_state import Direction, Environment, MatrixProductState
 from tnpy.model import RandomHeisenberg
-from tnpy.matrix_product_state import Direction, MatrixProductState, Environment
 
 
 class Helper:
@@ -84,9 +84,7 @@ def test_load(self, filename, expectation):
     def test_split_tensor(self, site, mps):
         two_site_mps = mps[site] @ mps[site + 1]
         mps.split_tensor(site, direction=Direction.RIGHTWARD)
-        np.testing.assert_allclose(
-            two_site_mps.data, (mps[site] @ mps[site + 1]).data, atol=1e-12
-        )
+        np.testing.assert_allclose(two_site_mps.data, (mps[site] @ mps[site + 1]).data, atol=1e-12)
         assert mps[site].tags == {f"I{site}"}
         assert mps[site + 1].tags == {f"I{site + 1}"}
 
@@ -95,9 +93,7 @@ class TestEnvironment:
     @pytest.fixture(scope="class")
     def env(self):
         model = RandomHeisenberg(n=6, h=0, penalty=100.0)
-        return Environment(
-            model.mpo, MatrixProductState.random(n=model.n, bond_dim=16, phys_dim=2)
-        )
+        return Environment(model.mpo, MatrixProductState.random(n=model.n, bond_dim=16, phys_dim=2))
 
     def test_left(self, env):
         print(env.mps)

diff --git a/tests/test_operators.py b/tests/test_operators.py
@@ -1,119 +1,98 @@
-import pytest
 import numpy as np
+import pytest
 
-from tnpy.operators import SpinOperators, FullHamiltonian
 from tnpy.model import XXZ, RandomHeisenberg
+from tnpy.operators import FullHamiltonian, SpinOperators
 
 
 class TestSpinOperators:
-
     def test_spin_half_ops(self):
         spin_half_ops = SpinOperators()
         np.testing.assert_array_equal(
-            np.array(
-                [[0, 1],
-                 [1, 0]]
-            ),
-            spin_half_ops.Sp + spin_half_ops.Sm
+            np.array([[0, 1], [1, 0]]), spin_half_ops.Sp + spin_half_ops.Sm
         )
         np.testing.assert_array_equal(
-            np.array(
-                [[0, -1j],
-                 [1j, 0]]
-            ),
-            -1j * (spin_half_ops.Sp - spin_half_ops.Sm)
+            np.array([[0, -1j], [1j, 0]]), -1j * (spin_half_ops.Sp - spin_half_ops.Sm)
         )
 
 
 class TestMatrixProductOperator:
-
-    @pytest.mark.parametrize("model", [
-        RandomHeisenberg(n=4, h=0),
-        RandomHeisenberg(n=4, h=0.5)
-    ])
+    @pytest.mark.parametrize("model", [RandomHeisenberg(n=4, h=0), RandomHeisenberg(n=4, h=0.5)])
     def test_square(self, model):
         bilayer_mpo = model.mpo.square()
         assert bilayer_mpo[0].shape == (25, 2, 2)
         assert bilayer_mpo[1].shape == (25, 25, 2, 2)
         assert bilayer_mpo[2].shape == (25, 25, 2, 2)
         assert bilayer_mpo[3].shape == (25, 2, 2)
         ham = FullHamiltonian(model.mpo).matrix
-        np.testing.assert_allclose(
-            ham @ ham,
-            FullHamiltonian(bilayer_mpo).matrix,
-            atol=1e-12
-        )
+        np.testing.assert_allclose(ham @ ham, FullHamiltonian(bilayer_mpo).matrix, atol=1e-12)
 
     @pytest.mark.parametrize("n", [2, 4, 6])
     @pytest.mark.parametrize("h", [0, 0.5, 1])
     def test_multiply_scalar(self, n, h):
         mpo = RandomHeisenberg(n=n, h=h).mpo
         np.testing.assert_array_equal(
-            -1 * FullHamiltonian(mpo).matrix,
-            FullHamiltonian(-1 * mpo).matrix
+            -1 * FullHamiltonian(mpo).matrix, FullHamiltonian(-1 * mpo).matrix
         )
 
 
 class TestFullHamiltonian:
-
-    @pytest.fixture(scope='class', params=[
-        {
-            "n": 2,
-            "ham": FullHamiltonian(RandomHeisenberg(n=2, h=0).mpo),
-            "data": np.array(
-                [[0.25, 0, 0, 0],
-                 [0, -0.25, 0.5, 0],
-                 [0, 0.5, -0.25, 0],
-                 [0, 0, 0, 0.25]]
-            )
-        },
-        {
-            "n": 3,
-            "ham": FullHamiltonian(RandomHeisenberg(n=3, h=0).mpo),
-            "data": np.array(
-                [[0.5, 0, 0, 0, 0, 0, 0, 0],
-                 [0, 0, 0.5, 0, 0, 0, 0, 0],
-                 [0, 0.5, -0.5, 0, 0.5, 0, 0, 0],
-                 [0, 0, 0, 0, 0, 0.5, 0, 0],
-                 [0, 0, 0.5, 0, 0, 0, 0, 0],
-                 [0, 0, 0, 0.5, 0, -0.5, 0.5, 0],
-                 [0, 0, 0, 0, 0, 0.5, 0, 0],
-                 [0, 0, 0, 0, 0, 0, 0, 0.5]]
-            )
-        },
-        {
-            "n": 2,
-            "ham": FullHamiltonian(XXZ(n=2, delta=0.5).mpo),
-            "data": np.array(
-                [[-0.125, 0, 0, 0],
-                 [0, 0.125, -0.5, 0],
-                 [0, -0.5, 0.125, 0],
-                 [0, 0, 0, -0.125]]
-            )
-        },
-        {
-            "n": 3,
-            "ham": FullHamiltonian(XXZ(n=3, delta=0.5).mpo),
-            "data": np.array(
-                [[-0.25, 0, 0, 0, 0, 0, 0, 0],
-                 [0, 0, -0.5, 0, 0, 0, 0, 0],
-                 [0, -0.5, 0.25, 0, -0.5, 0, 0, 0],
-                 [0, 0, 0, 0, 0, -0.5, 0, 0],
-                 [0, 0, -0.5, 0, 0, 0, 0, 0],
-                 [0, 0, 0, -0.5, 0, 0.25, -0.5, 0],
-                 [0, 0, 0, 0, 0, -0.5, 0, 0],
-                 [0, 0, 0, 0, 0, 0, 0, -0.25]]
-            )
-        }
-    ])
+    @pytest.fixture(
+        scope="class",
+        params=[
+            {
+                "n": 2,
+                "ham": FullHamiltonian(RandomHeisenberg(n=2, h=0).mpo),
+                "data": np.array(
+                    [[0.25, 0, 0, 0], [0, -0.25, 0.5, 0], [0, 0.5, -0.25, 0], [0, 0, 0, 0.25]]
+                ),
+            },
+            {
+                "n": 3,
+                "ham": FullHamiltonian(RandomHeisenberg(n=3, h=0).mpo),
+                "data": np.array(
+                    [
+                        [0.5, 0, 0, 0, 0, 0, 0, 0],
+                        [0, 0, 0.5, 0, 0, 0, 0, 0],
+                        [0, 0.5, -0.5, 0, 0.5, 0, 0, 0],
+                        [0, 0, 0, 0, 0, 0.5, 0, 0],
+                        [0, 0, 0.5, 0, 0, 0, 0, 0],
+                        [0, 0, 0, 0.5, 0, -0.5, 0.5, 0],
+                        [0, 0, 0, 0, 0, 0.5, 0, 0],
+                        [0, 0, 0, 0, 0, 0, 0, 0.5],
+                    ]
+                ),
+            },
+            {
+                "n": 2,
+                "ham": FullHamiltonian(XXZ(n=2, delta=0.5).mpo),
+                "data": np.array(
+                    [[-0.125, 0, 0, 0], [0, 0.125, -0.5, 0], [0, -0.5, 0.125, 0], [0, 0, 0, -0.125]]
+                ),
+            },
+            {
+                "n": 3,
+                "ham": FullHamiltonian(XXZ(n=3, delta=0.5).mpo),
+                "data": np.array(
+                    [
+                        [-0.25, 0, 0, 0, 0, 0, 0, 0],
+                        [0, 0, -0.5, 0, 0, 0, 0, 0],
+                        [0, -0.5, 0.25, 0, -0.5, 0, 0, 0],
+                        [0, 0, 0, 0, 0, -0.5, 0, 0],
+                        [0, 0, -0.5, 0, 0, 0, 0, 0],
+                        [0, 0, 0, -0.5, 0, 0.25, -0.5, 0],
+                        [0, 0, 0, 0, 0, -0.5, 0, 0],
+                        [0, 0, 0, 0, 0, 0, 0, -0.25],
+                    ]
+                ),
+            },
+        ],
+    )
     def model(self, request):
         return request.param
 
     def test_n_sites(self, model):
         assert model["ham"].n_sites == model["n"]
 
     def test_matrix(self, model):
-        np.testing.assert_array_equal(
-            model["ham"].matrix,
-            model["data"]
-        )
+        np.testing.assert_array_equal(model["ham"].matrix, model["data"])
diff --git a/tests/test_tsdrg.py b/tests/test_tsdrg.py
@@ -1,15 +1,15 @@
-import pytest
 import numpy as np
+import pytest
 
-from tnpy.model import RandomHeisenberg, TotalSz
-from tnpy.operators import SpinOperators, FullHamiltonian
 from tnpy.exact_diagonalization import ExactDiagonalization
+from tnpy.model import RandomHeisenberg, TotalSz
+from tnpy.operators import FullHamiltonian, SpinOperators
 from tnpy.tsdrg import (
-    TensorTree,
-    TreeTensorNetworkSDRG as tSDRG,
     HighEnergyTreeTensorNetworkSDRG,
     ShiftInvertTreeTensorNetworkSDRG,
+    TensorTree,
 )
+from tnpy.tsdrg import TreeTensorNetworkSDRG as tSDRG
 
 
 class TestTensorTree:
@@ -21,25 +21,19 @@ def tree(self):
             left_id=1,
             right_id=2,
             new_id=4,
-            data=ExactDiagonalization(RandomHeisenberg(n=2, h=0).mpo).evecs.reshape(
-                (2, 2, 4)
-            ),
+            data=ExactDiagonalization(RandomHeisenberg(n=2, h=0).mpo).evecs.reshape((2, 2, 4)),
         )
         tree.fuse(
             left_id=0,
             right_id=4,
             new_id=5,
-            data=ExactDiagonalization(RandomHeisenberg(n=3, h=0).mpo).evecs.reshape(
-                (2, 4, 8)
-            ),
+            data=ExactDiagonalization(RandomHeisenberg(n=3, h=0).mpo).evecs.reshape((2, 4, 8)),
         )
         tree.fuse(
             left_id=5,
             right_id=3,
             new_id=6,
-            data=ExactDiagonalization(RandomHeisenberg(n=4, h=0).mpo).evecs.reshape(
-                (8, 2, 16)
-            ),
+            data=ExactDiagonalization(RandomHeisenberg(n=4, h=0).mpo).evecs.reshape((8, 2, 16)),
         )
         return tree
 
@@ -150,9 +144,7 @@ def test_run(self, ed, tsdrg):
 
         np.set_printoptions(threshold=sys.maxsize, linewidth=sys.maxsize)
         tsdrg.run()
-        np.testing.assert_allclose(
-            ed.evals[: tsdrg.chi], np.sort(tsdrg.evals), atol=1e-8
-        )
+        np.testing.assert_allclose(ed.evals[: tsdrg.chi], np.sort(tsdrg.evals), atol=1e-8)
 
     def test_measurements(self, ed, tsdrg):
         tsdrg.run()
@@ -218,9 +210,7 @@ def test_reduced_density_matrix(self, site, level_idx, ed, tsdrg):
         np.testing.assert_allclose(
             ed.reduced_density_matrix(site=site, level_idx=level_idx),
             np.linalg.eigvalsh(
-                tsdrg.measurements.reduced_density_matrix(
-                    site=site, level_idx=level_idx
-                )
+                tsdrg.measurements.reduced_density_matrix(site=site, level_idx=level_idx)
             )[::-1],
             atol=1e-12,
         )
@@ -252,9 +242,7 @@ def test_two_point_function(self, site1, site2, level_idx, ed, tsdrg):
         Sp, Sm, Sz, I2, O2 = SpinOperators()
         np.testing.assert_allclose(
             ed.two_point_function(Sz, Sz, site1, site2, level_idx=level_idx),
-            tsdrg.measurements.two_point_function(
-                Sz, Sz, site1, site2, level_idx=level_idx
-            ),
+            tsdrg.measurements.two_point_function(Sz, Sz, site1, site2, level_idx=level_idx),
             atol=1e-12,
         )
 
@@ -319,12 +307,8 @@ def model(self):
 
     @pytest.fixture(scope="function")
     def sitsdrg(self, model, offset):
-        shifted_model = RandomHeisenberg(
-            n=model.n, h=model.h, seed=model.seed, offset=offset
-        )
-        return ShiftInvertTreeTensorNetworkSDRG(
-            shifted_model.mpo, chi=2**5, offset=offset
-        )
+        shifted_model = RandomHeisenberg(n=model.n, h=model.h, seed=model.seed, offset=offset)
+        return ShiftInvertTreeTensorNetworkSDRG(shifted_model.mpo, chi=2**5, offset=offset)
 
     @pytest.fixture(scope="class")
     def ed(self, model):