adjust imports in tests

pymf/tests/test_graphene.py

@@ -3,11 +3,13 @@ import numpy as np
 import pytest
 
 from pymf.kwant_helper import kwant_examples, utils
-from pymf.model import Model
-from pymf.solvers import solver
-from pymf.tb.tb import add_tb
-from pymf.tb.transforms import tb_to_khamvector
-from pymf.tb.utils import generate_guess
+from pymf import (
+    Model,
+    solver,
+    tb_to_khamvector,
+    generate_guess,
+    add_tb,
+)
 
 
 def compute_gap(tb, fermi_energy=0, nk=100):

pymf/tests/test_hat.py

 # %%
 import numpy as np
-from pymf.solvers import solver
-from pymf.tb import utils
-from pymf.model import Model
-from pymf.tb.tb import add_tb, scale_tb
-from pymf import mf
-from pymf import observables
 import pytest
 
+from pymf import (
+    Model,
+    solver,
+    generate_guess,
+    scale_tb,
+    add_tb,
+    expectation_value,
+    construct_density_matrix,
+)
+
 
 # %%
 def total_energy(ham_tb, rho_tb):
-    return np.real(observables.expectation_value(rho_tb, ham_tb))
+    return np.real(expectation_value(rho_tb, ham_tb))
 
 
 # %%
@@ -31,7 +35,7 @@ h_int_U0 = {
 @np.vectorize
 def mf_rescaled(alpha, mf0):
     hamiltonian = add_tb(h_0, scale_tb(mf0, alpha))
-    rho, _ = mf.construct_density_matrix(hamiltonian, filling=filling, nk=nk)
+    rho, _ = construct_density_matrix(hamiltonian, filling=filling, nk=nk)
     hamiltonian = add_tb(h_0, scale_tb(mf0, np.sign(alpha)))
     return total_energy(hamiltonian, rho)
 
@@ -39,7 +43,7 @@ def mf_rescaled(alpha, mf0):
 @pytest.mark.parametrize("seed", range(repeat_number))
 def test_mexican_hat(seed):
     np.random.seed(seed)
-    guess = utils.generate_guess(frozenset(h_int_U0), len(h_int_U0[(0,)]))
+    guess = generate_guess(frozenset(h_int_U0), len(h_int_U0[(0,)]))
     _model = Model(h_0, h_int_U0, filling=filling)
     mf_sol_groundstate = solver(
         _model, mf_guess=guess, nk=nk, optimizer_kwargs={"M": 0}

pymf/tests/test_hubbard.py

@@ -2,11 +2,13 @@
 import numpy as np
 import pytest
 
-from pymf.model import Model
-from pymf.solvers import solver
-from pymf.tb import utils
-from pymf.tb.tb import add_tb
 from pymf.tests.test_graphene import compute_gap
+from pymf import (
+    Model,
+    solver,
+    generate_guess,
+    add_tb,
+)
 
 repeat_number = 10
 
@@ -33,7 +35,7 @@ def gap_relation_hubbard(Us, nk, nk_dense, tol=1e-3):
         h_int = {
             (0,): U * np.kron(np.ones((2, 2)), np.eye(2)),
         }
-        guess = utils.generate_guess(frozenset(h_int), len(list(h_0.values())[0]))
+        guess = generate_guess(frozenset(h_int), len(list(h_0.values())[0]))
         full_model = Model(h_0, h_int, filling=2)
         mf_sol = solver(full_model, guess, nk=nk)
         _gap = compute_gap(add_tb(h_0, mf_sol), fermi_energy=0, nk=nk_dense)

pymf/tests/test_params.py

@@ -3,7 +3,7 @@ import pytest
 import numpy as np
 from pymf.params.rparams import rparams_to_tb, tb_to_rparams
 from pymf.tb.tb import compare_dicts
-from pymf.tb.utils import generate_guess
+from pymf import generate_guess
 
 repeat_number = 10
 

pymf/tests/test_zero_hint.py

@@ -2,10 +2,9 @@
 import numpy as np
 import pytest
 
-from pymf.model import Model
-from pymf.solvers import solver
 from pymf.tb import utils
-from pymf.tb.tb import add_tb, compare_dicts
+from pymf.tb.tb import compare_dicts
+from pymf import Model, solver, generate_guess, add_tb, calculate_fermi_energy
 
 # %%
 repeat_number = 10
@@ -24,13 +23,13 @@ def test_zero_hint(seed):
     random_hopping_vecs = utils.generate_tb_keys(cutoff, dim)
 
     zero_key = tuple([0] * dim)
-    h_0_random = utils.generate_guess(random_hopping_vecs, ndof, scale=1)
-    h_int_only_phases = utils.generate_guess(random_hopping_vecs, ndof, scale=0)
-    guess = utils.generate_guess(random_hopping_vecs, ndof, scale=1)
+    h_0_random = generate_guess(random_hopping_vecs, ndof, scale=1)
+    h_int_only_phases = generate_guess(random_hopping_vecs, ndof, scale=0)
+    guess = generate_guess(random_hopping_vecs, ndof, scale=1)
     model = Model(h_0_random, h_int_only_phases, filling=filling)
 
     mf_sol = solver(model, guess, nk=40, optimizer_kwargs={"M": 0, "f_tol": 1e-10})
-    h_fermi = utils.calculate_fermi_energy(mf_sol, filling=filling, nk=20)
+    h_fermi = calculate_fermi_energy(mf_sol, filling=filling, nk=20)
     mf_sol[zero_key] -= h_fermi * np.eye(mf_sol[zero_key].shape[0])
 
     compare_dicts(add_tb(mf_sol, h_0_random), h_0_random, atol=1e-10)