Fix imports

.ruff.toml

@@ -11,3 +11,4 @@ ignore = [
         # Prevents the use of the characters 'l', 'O', or 'I' as variable names.
         # Overly restrictive, in particular when implementing mathematical expressions.
 ]
+preview = true

src/kbe/__init__.py

@@ -1,2 +1,4 @@
 from kbe.fragment import fragpart
 from kbe.pbe import BE
+
+__all__ = ["fragpart", "BE"]

src/kbe/_opt.py

@@ -2,6 +2,7 @@
 
 import sys
 
+from kbe.misc import print_energy
 from molbe._opt import BEOPT
 
 
@@ -53,8 +54,6 @@ def optimize(
     J0 : list of list of float
        Initial Jacobian.
     """
-    from .misc import print_energy
-
     # Check if only chemical potential optimization is required
     if not only_chem:
         pot = self.pot

src/kbe/autofrag.py

@@ -3,7 +3,9 @@
 import sys
 
 import numpy
+from pyscf import lib
 
+from kbe.misc import sgeom
 from molbe.helper import get_core
 
 
@@ -324,10 +326,6 @@ def autogen(
         Weights for each fragment. Each entry contains a weight and a list of LO
         indices.
     """
-    from pyscf import lib
-
-    from .misc import sgeom
-
     if not float(unitcell).is_integer():
         print("Fractional unitcell is not supported!")
         sys.exit()

src/kbe/chain.py

@@ -3,10 +3,10 @@
 
 # Old fragmentation code - do not use other than debugging!
 
+import math
 
-def findH(mol, nh, tmphlist=[]):
-    import math
 
+def findH(mol, nh, tmphlist=[]):
     coord_ = mol.atom_coords()
 
     hidx = []

src/kbe/fragment.py

@@ -2,10 +2,9 @@
 
 import sys
 
+from kbe.autofrag import autogen
 from molbe.helper import get_core
 
-from .autofrag import autogen
-
 
 def print_mol_missing():
     print("Provide pyscf gto.Cell object in fragpart() and restart!", flush=True)
@@ -104,7 +103,7 @@ def __init__(
         self.allcen = allcen
         self.valence_basis = valence_basis
         self.kpt = kpt
-        self.molecule = False  ### remove this
+        self.molecule = False  # remove this
 
         # Check for frozen core approximation
         if frozen_core:
@@ -156,4 +155,6 @@ def __init__(
             print("exiting", flush=True)
             sys.exit()
 
-    from .chain import polychain
+    # This import makes polychain a method of the class and
+    #  cannot be moved to the top of the file
+    from kbe.chain import polychain  # noqa: PLC0415

src/kbe/helper.py

@@ -3,11 +3,10 @@
 import functools
 
 import numpy
+from pyscf import scf
 
 
 def get_veff(eri_, dm, S, TA, hf_veff, return_veff0=False):
-    from pyscf import scf
-
     """
     Calculate the effective HF potential (Veff) for a given density matrix
     and electron repulsion integrals.

src/kbe/lo.py

@@ -1,12 +1,24 @@
 # Author(s): Oinam Meitei
 #            Henry Tran
 #
+import functools
+import os
 import sys
 from functools import reduce
 
 import numpy
-
+import scipy.linalg
+from libdmet.lo import pywannier90
+
+from kbe.lo_k import (
+    get_iao_k,
+    get_pao_native_k,
+    get_xovlp_k,
+    remove_core_mo_k,
+    symm_orth_k,
+)
 from molbe.external.lo_helper import get_aoind_by_atom, reorder_by_atom_
+from molbe.helper import ncore_
 
 # iao tmp
 
@@ -51,17 +63,11 @@ def localize(
     iao_wannier : bool
        Whether to perform Wannier localization in the IAO space
     """
-    import functools
-
-    import scipy.linalg
-
-    from molbe.helper import ncore_
-
     if lo_method == "iao":
         if valence_basis == "sto-3g":
-            from .basis_sto3g_core_val import core_basis, val_basis
+            from .basis_sto3g_core_val import core_basis, val_basis  # noqa: PLC0415
         elif valence_basis == "minao":
-            from .basis_minao_core_val import core_basis, val_basis
+            from .basis_minao_core_val import core_basis, val_basis  # noqa: PLC0415
         elif iao_val_core:
             sys.exit(
                 "valence_basis="
@@ -128,18 +134,6 @@ def localize(
         self.cinv = cinv_
 
     elif lo_method == "iao":
-        import os
-
-        from libdmet.lo import pywannier90
-
-        from .lo_k import (
-            get_iao_k,
-            get_pao_native_k,
-            get_xovlp_k,
-            remove_core_mo_k,
-            symm_orth_k,
-        )
-
         if not iao_val_core or not self.frozen_core:
             Co = self.C[:, :, : self.Nocc].copy()
             S12, S2 = get_xovlp_k(self.cell, self.kpts, basis=valence_basis)
@@ -395,11 +389,6 @@ def localize(
         self.cinv = cinv_
 
     elif lo_method == "wannier":
-        # from pyscf.pbc.tools import pywannier90
-        from libdmet.lo import pywannier90
-
-        from .lo_k import remove_core_mo_k
-
         nk, nao, nmo = self.C.shape
         lorb = numpy.zeros((nk, nao, nmo), dtype=numpy.complex128)
         lorb_nocore = numpy.zeros((nk, nao, nmo - self.ncore), dtype=numpy.complex128)

src/kbe/lo_k.py

@@ -6,6 +6,7 @@
 
 import numpy
 import scipy
+from pyscf.pbc import gto as pgto
 
 
 def dot_gen(A, B, ovlp):
@@ -67,9 +68,6 @@ def get_xovlp_k(cell, kpts, basis="sto-3g"):
         S12 - Overlap of two basis sets
         S22 - Overlap in new basis set
     """
-
-    from pyscf.pbc import gto as pgto  # intor_cross
-
     cell_alt = cell.copy()
     cell_alt.basis = basis
     cell_alt.build()

src/kbe/misc.py

@@ -2,6 +2,7 @@
 
 import numpy
 from pyscf import lib
+from pyscf.pbc import tools
 
 
 def sgeom(cell, kmesh=None):
@@ -15,11 +16,7 @@ def sgeom(cell, kmesh=None):
     kmesh : list of int
        Number of k-points in each lattice vector dimension
     """
-    from pyscf.pbc import tools
-
-    scell = tools.super_cell(cell, kmesh)
-
-    return scell
+    return tools.super_cell(cell, kmesh)
 
 
 def get_phase(cell, kpts, kmesh):

src/kbe/pbe.py

@@ -3,14 +3,18 @@
 import os
 import pickle
 import sys
+from multiprocessing import Pool
 
 import h5py
 import numpy
+from libdmet.basis_transform.eri_transform import get_emb_eri_fast_gdf
+from pyscf import ao2mo
+from pyscf.pbc import df, gto
+from pyscf.pbc.df.df_jk import _ewald_exxdiv_for_G0
 
 import molbe.be_var as be_var
-
-from .misc import storePBE
-from .pfrag import Frags
+from kbe.misc import storePBE
+from kbe.pfrag import Frags
 
 
 class BE:
@@ -322,11 +326,11 @@ def __init__(
         if not restart:
             self.initialize(mf._eri, compute_hf)
 
-    # this is a molbe method not BEOPT
-    from molbe.external.optqn import get_be_error_jacobian
-
-    from ._opt import optimize
-    from .lo import localize
+    # The following import of these functions turns them into
+    # proper methods of the class.
+    from kbe._opt import optimize  # noqa: PLC0415
+    from kbe.lo import localize  # noqa: PLC0415
+    from molbe.external.optqn import get_be_error_jacobian  # noqa: PLC0415
 
     def print_ini(self):
         """
@@ -349,8 +353,6 @@ def print_ini(self):
         print(flush=True)
 
     def ewald_sum(self, kpts=None):
-        from pyscf.pbc.df.df_jk import _ewald_exxdiv_for_G0
-
         dm_ = self.mf.make_rdm1()
         nk, nao, nao = dm_.shape
 
@@ -380,13 +382,6 @@ def initialize(self, eri_, compute_hf, restart=False):
         restart : bool, optional
             Whether to restart from a previous calculation, by default False.
         """
-        import os
-        from multiprocessing import Pool
-
-        import h5py
-        from libdmet.basis_transform.eri_transform import get_emb_eri_fast_gdf
-        from pyscf import ao2mo
-
         if compute_hf:
             E_hf = 0.0
         EH1 = 0.0
@@ -609,8 +604,9 @@ def oneshot(
         clean_eri : bool, optional
             Whether to clean up ERI files after calculation, by default False.
         """
-        from .be_parallel import be_func_parallel
-        from .solver import be_func
+        # has to be here because of circular dependency
+        from .be_parallel import be_func_parallel  # noqa: PLC0415
+        from .solver import be_func  # noqa: PLC0415
 
         print("Calculating Energy by Fragment? ", calc_frag_energy)
         if nproc == 1:
@@ -768,10 +764,6 @@ def eritransform_parallel(a, atom, basis, kpts, C_ao_emb, cderi):
     """
     Wrapper for parallel eri transformation
     """
-    from pyscf.pbc import df, gto
-
-    from molbe.external.eri_transform import get_emb_eri_fast_gdf
-
     cell = gto.Cell()
     cell.a = a
     cell.atom = atom

src/kbe/pfrag.py

@@ -6,12 +6,11 @@
 import h5py
 import numpy
 
+from kbe.helper import get_veff
+from kbe.misc import get_phase, get_phase1
+from kbe.solver import schmidt_decomp_svd
 from molbe.helper import get_eri, get_scfObj
 
-from .helper import *
-from .misc import *
-from .solver import schmidt_decomp_svd
-
 
 class Frags:
     """
@@ -141,9 +140,6 @@ def sd(
         kmesh : list of int
             Number of k-points in each lattice vector direction
         """
-
-        from .misc import get_phase
-
         nk, nao, nlo = lao.shape
         rdm1_lo_k = numpy.zeros((nk, nlo, nlo), dtype=numpy.result_type(lmo, lmo))
         for k in range(nk):

src/kbe/solver.py

@@ -1,6 +1,7 @@
 # Author(s): Oinam Romesh Meitei
 
 import numpy
+import scipy.linalg
 
 
 def schmidt_decomp_svd(rdm, Frag_sites):
@@ -23,8 +24,6 @@ def schmidt_decomp_svd(rdm, Frag_sites):
     numpy.ndarray
         Transformation matrix (TA) including both fragment and entangled bath orbitals.
     """
-    import scipy.linalg
-
     thres = 1.0e-10
     Tot_sites = rdm.shape[0]
 

src/molbe/__init__.py

@@ -1,3 +1,5 @@
 from molbe.fragment import fragpart
 from molbe.mbe import BE
 from molbe.ube import UBE
+
+__all__ = ["fragpart", "BE", "UBE"]

src/molbe/_opt.py

@@ -4,8 +4,10 @@
 
 import numpy
 
-from .be_parallel import be_func_parallel
-from .solver import be_func
+from molbe.be_parallel import be_func_parallel
+from molbe.external.optqn import FrankQN
+from molbe.misc import print_energy
+from molbe.solver import be_func
 
 
 class BEOPT:
@@ -182,10 +184,6 @@ def optimize(self, method, J0=None, trust_region=False):
         trust_region : bool, optional
            Use trust-region based QN optimization, by default False
         """
-        import sys
-
-        from molbe.external.optqn import FrankQN
-
         print("-----------------------------------------------------", flush=True)
         print("             Starting BE optimization ", flush=True)
         print("             Solver : ", self.solver, flush=True)
@@ -289,8 +287,6 @@ def optimize(
     trust_region : bool, optional
        Use trust-region based QN optimization, by default False
     """
-    from .misc import print_energy
-
     # Check if only chemical potential optimization is required
     if not only_chem:
         pot = self.pot

src/molbe/autofrag.py

@@ -4,7 +4,7 @@
 
 import numpy
 
-from .helper import get_core
+from molbe.helper import get_core
 
 
 def autogen(