chore: ban module-level import of pymatgen (#1597)

njzjz · web-flow · commit 3a1ed00ac12e · 2024-07-19T22:40:26.000+08:00
Pymatgen is still imcompatbile with NumPy 2. Move the import into
runtime to fix the documentation builds and prevent potential errors
when one uses dpgen without using pymatgen.

Signed-off-by: Jinzhe Zeng &lt;jinzhe.zeng@rutgers.edu&gt;


&lt;!-- This is an auto-generated comment: release notes by coderabbit.ai
--&gt;

## Summary by CodeRabbit

- **Chores**
- Reorganized import statements across multiple files to improve code
readability and maintainability.
- Localized imports within specific methods for better encapsulation and
efficient module loading.
- Updated linting configurations in `pyproject.toml` to include new
rules and per-file exceptions.

&lt;!-- end of auto-generated comment: release notes by coderabbit.ai --&gt;

Signed-off-by: Jinzhe Zeng &lt;jinzhe.zeng@rutgers.edu&gt;
diff --git a/dpgen/auto_test/Elastic.py b/dpgen/auto_test/Elastic.py
@@ -4,11 +4,6 @@
 from shutil import copyfile
 
 from monty.serialization import dumpfn, loadfn
-from pymatgen.analysis.elasticity.elastic import ElasticTensor
-from pymatgen.analysis.elasticity.strain import DeformedStructureSet, Strain
-from pymatgen.analysis.elasticity.stress import Stress
-from pymatgen.core.structure import Structure
-from pymatgen.io.vasp import Incar, Kpoints
 
 import dpgen.auto_test.lib.abacus as abacus
 import dpgen.auto_test.lib.vasp as vasp
@@ -53,6 +48,9 @@ def __init__(self, parameter, inter_param=None):
         self.inter_param = inter_param if inter_param is not None else {"type": "vasp"}
 
     def make_confs(self, path_to_work, path_to_equi, refine=False):
+        from pymatgen.analysis.elasticity.strain import DeformedStructureSet, Strain
+        from pymatgen.core.structure import Structure
+
         path_to_work = os.path.abspath(path_to_work)
         if os.path.exists(path_to_work):
             dlog.warning(f"{path_to_work} already exists")
@@ -189,6 +187,8 @@ def make_confs(self, path_to_work, path_to_equi, refine=False):
         return task_list
 
     def post_process(self, task_list):
+        from pymatgen.io.vasp import Incar, Kpoints
+
         if self.inter_param["type"] == "abacus":
             POSCAR = "STRU"
             INCAR = "INPUT"
@@ -250,6 +250,9 @@ def task_param(self):
         return self.parameter
 
     def _compute_lower(self, output_file, all_tasks, all_res):
+        from pymatgen.analysis.elasticity.elastic import ElasticTensor
+        from pymatgen.analysis.elasticity.stress import Stress
+
         output_file = os.path.abspath(output_file)
         res_data = {}
         ptr_data = os.path.dirname(output_file) + "\n"
diff --git a/dpgen/auto_test/Gamma.py b/dpgen/auto_test/Gamma.py
@@ -8,8 +8,6 @@
 from ase.lattice.cubic import BodyCenteredCubic as bcc
 from ase.lattice.cubic import FaceCenteredCubic as fcc
 from monty.serialization import dumpfn, loadfn
-from pymatgen.core.structure import Structure
-from pymatgen.io.ase import AseAtomsAdaptor
 
 import dpgen.auto_test.lib.abacus as abacus
 import dpgen.auto_test.lib.vasp as vasp
@@ -94,6 +92,8 @@ def __init__(self, parameter, inter_param=None):
         self.inter_param = inter_param if inter_param is not None else {"type": "vasp"}
 
     def make_confs(self, path_to_work, path_to_equi, refine=False):
+        from pymatgen.core.structure import Structure
+
         path_to_work = os.path.abspath(path_to_work)
         if os.path.exists(path_to_work):
             dlog.warning(f"{path_to_work} already exists")
@@ -287,6 +287,8 @@ def return_direction(self):
         return directions
 
     def __gen_slab_ase(self, symbol, lat_param):
+        from pymatgen.io.ase import AseAtomsAdaptor
+
         if not self.lattice_type:
             raise RuntimeError("Error! Please provide the input lattice type!")
         elif self.lattice_type == "bcc":
diff --git a/dpgen/auto_test/Interstitial.py b/dpgen/auto_test/Interstitial.py
@@ -5,8 +5,6 @@
 
 import numpy as np
 from monty.serialization import dumpfn, loadfn
-from pymatgen.analysis.defects.generators import InterstitialGenerator
-from pymatgen.core.structure import Structure
 
 import dpgen.auto_test.lib.abacus as abacus
 import dpgen.auto_test.lib.lammps as lammps
@@ -78,6 +76,9 @@ def __init__(self, parameter, inter_param=None):
         self.inter_param = inter_param if inter_param is not None else {"type": "vasp"}
 
     def make_confs(self, path_to_work, path_to_equi, refine=False):
+        from pymatgen.analysis.defects.generators import InterstitialGenerator
+        from pymatgen.core.structure import Structure
+
         path_to_work = os.path.abspath(path_to_work)
         path_to_equi = os.path.abspath(path_to_equi)
 
diff --git a/dpgen/auto_test/Surface.py b/dpgen/auto_test/Surface.py
@@ -6,8 +6,6 @@
 import dpdata
 import numpy as np
 from monty.serialization import dumpfn, loadfn
-from pymatgen.core.structure import Structure
-from pymatgen.core.surface import generate_all_slabs
 
 import dpgen.auto_test.lib.abacus as abacus
 import dpgen.auto_test.lib.vasp as vasp
@@ -85,6 +83,9 @@ def __init__(self, parameter, inter_param=None):
         self.inter_param = inter_param if inter_param is not None else {"type": "vasp"}
 
     def make_confs(self, path_to_work, path_to_equi, refine=False):
+        from pymatgen.core.structure import Structure
+        from pymatgen.core.surface import generate_all_slabs
+
         path_to_work = os.path.abspath(path_to_work)
         if os.path.exists(path_to_work):
             dlog.warning(f"{path_to_work} already exists")
diff --git a/dpgen/auto_test/VASP.py b/dpgen/auto_test/VASP.py
@@ -2,8 +2,6 @@
 
 from dpdata import LabeledSystem
 from monty.serialization import dumpfn
-from pymatgen.core.structure import Structure
-from pymatgen.io.vasp import Incar, Kpoints
 
 import dpgen.auto_test.lib.vasp as vasp
 from dpgen import dlog
@@ -22,6 +20,8 @@ def __init__(self, inter_parameter, path_to_poscar):
         self.path_to_poscar = path_to_poscar
 
     def make_potential_files(self, output_dir):
+        from pymatgen.core.structure import Structure
+
         potcar_not_link_list = ["vacancy", "interstitial"]
         task_type = output_dir.split("/")[-2].split("_")[0]
 
@@ -69,6 +69,8 @@ def make_potential_files(self, output_dir):
         dumpfn(self.inter, os.path.join(output_dir, "inter.json"), indent=4)
 
     def make_input_file(self, output_dir, task_type, task_param):
+        from pymatgen.io.vasp import Incar, Kpoints
+
         sepline(ch=output_dir)
         dumpfn(task_param, os.path.join(output_dir, "task.json"), indent=4)
 
diff --git a/dpgen/auto_test/Vacancy.py b/dpgen/auto_test/Vacancy.py
@@ -5,8 +5,6 @@
 
 import numpy as np
 from monty.serialization import dumpfn, loadfn
-from pymatgen.analysis.defects.generators import VacancyGenerator
-from pymatgen.core.structure import Structure
 
 import dpgen.auto_test.lib.abacus as abacus
 from dpgen import dlog
@@ -77,6 +75,9 @@ def __init__(self, parameter, inter_param=None):
         self.inter_param = inter_param if inter_param is not None else {"type": "vasp"}
 
     def make_confs(self, path_to_work, path_to_equi, refine=False):
+        from pymatgen.analysis.defects.generators import VacancyGenerator
+        from pymatgen.core.structure import Structure
+
         path_to_work = os.path.abspath(path_to_work)
         if os.path.exists(path_to_work):
             dlog.warning(f"{path_to_work} already exists")
diff --git a/dpgen/auto_test/gen_confs.py b/dpgen/auto_test/gen_confs.py
@@ -4,9 +4,6 @@
 import os
 import re
 
-from pymatgen.analysis.structure_matcher import StructureMatcher
-from pymatgen.ext.matproj import MPRester
-
 import dpgen.auto_test.lib.crys as crys
 
 global_std_crystal = {
@@ -20,6 +17,8 @@
 
 
 def test_fit(struct, data):
+    from pymatgen.analysis.structure_matcher import StructureMatcher
+
     m = StructureMatcher()
     for ii in data:
         if m.fit(ii["structure"], struct):
@@ -50,6 +49,9 @@ def gen_ele_std(ele_name, ctype):
 
 
 def gen_element(ele_name, key):
+    from pymatgen.analysis.structure_matcher import StructureMatcher
+    from pymatgen.ext.matproj import MPRester
+
     assert isinstance(ele_name, str)
     mpr = MPRester(key)
     data = mpr.query(
@@ -93,6 +95,8 @@ def gen_element_std(ele_name):
 
 
 def gen_alloy(eles, key):
+    from pymatgen.ext.matproj import MPRester
+
     mpr = MPRester(key)
 
     data = mpr.query(
diff --git a/dpgen/auto_test/lib/abacus.py b/dpgen/auto_test/lib/abacus.py
@@ -7,7 +7,6 @@
 from dpdata.abacus.scf import make_unlabeled_stru
 from dpdata.utils import uniq_atom_names
 from dpdata.vasp import poscar as dpdata_poscar
-from pymatgen.core.structure import Structure
 
 import dpgen.generator.lib.abacus_scf as abacus_scf
 
@@ -343,6 +342,8 @@ def final_stru(abacus_path):
 
 
 def stru2Structure(struf):
+    from pymatgen.core.structure import Structure
+
     stru = dpdata.System(struf, fmt="stru")
     stru.to("poscar", "POSCAR.tmp")
     ss = Structure.from_file("POSCAR.tmp")
diff --git a/dpgen/auto_test/lib/crys.py b/dpgen/auto_test/lib/crys.py
@@ -1,15 +1,18 @@
 import numpy as np
-from pymatgen.core.lattice import Lattice
-from pymatgen.core.structure import Structure
 
 
 def fcc(ele_name="ele", a=4.05):
+    from pymatgen.core.structure import Structure
+
     box = np.array([[0.0, 0.5, 0.5], [0.5, 0.0, 0.5], [0.5, 0.5, 0.0]])
     box *= a
     return Structure(box, [ele_name], [[0, 0, 0]])
 
 
 def fcc1(ele_name="ele", a=4.05):
+    from pymatgen.core.lattice import Lattice
+    from pymatgen.core.structure import Structure
+
     latt = Lattice.cubic(a)
     return Structure(
         latt,
@@ -19,11 +22,17 @@ def fcc1(ele_name="ele", a=4.05):
 
 
 def sc(ele_name="ele", a=2.551340126037118):
+    from pymatgen.core.lattice import Lattice
+    from pymatgen.core.structure import Structure
+
     latt = Lattice.cubic(a)
     return Structure(latt, [ele_name], [[0, 0, 0]])
 
 
 def bcc(ele_name="ele", a=3.2144871302356037):
+    from pymatgen.core.lattice import Lattice
+    from pymatgen.core.structure import Structure
+
     latt = Lattice.cubic(a)
     return Structure(
         latt,
@@ -38,6 +47,9 @@ def bcc(ele_name="ele", a=3.2144871302356037):
 def hcp(
     ele_name="ele", a=4.05 / np.sqrt(2), c=4.05 / np.sqrt(2) * 2.0 * np.sqrt(2.0 / 3.0)
 ):
+    from pymatgen.core.lattice import Lattice
+    from pymatgen.core.structure import Structure
+
     box = np.array([[1, 0, 0], [0.5, 0.5 * np.sqrt(3), 0], [0, 0, 1]])
     box[0] *= a
     box[1] *= a
@@ -51,6 +63,9 @@ def hcp(
 def dhcp(
     ele_name="ele", a=4.05 / np.sqrt(2), c=4.05 / np.sqrt(2) * 4.0 * np.sqrt(2.0 / 3.0)
 ):
+    from pymatgen.core.lattice import Lattice
+    from pymatgen.core.structure import Structure
+
     box = np.array([[1, 0, 0], [0.5, 0.5 * np.sqrt(3), 0], [0, 0, 1]])
     box[0] *= a
     box[1] *= a
@@ -69,6 +84,9 @@ def dhcp(
 
 
 def diamond(ele_name="ele", a=2.551340126037118):
+    from pymatgen.core.lattice import Lattice
+    from pymatgen.core.structure import Structure
+
     box = np.array([[0.0, 1.0, 1.0], [1.0, 0.0, 1.0], [1.0, 1.0, 0.0]])
     box *= a
     latt = Lattice(box)
diff --git a/dpgen/auto_test/lib/vasp.py b/dpgen/auto_test/lib/vasp.py
@@ -3,7 +3,6 @@
 import warnings
 
 import numpy as np
-from pymatgen.io.vasp import Incar, Kpoints
 
 import dpgen.auto_test.lib.util as util
 from dpgen.generator.lib.vasp import incar_upper
@@ -503,6 +502,8 @@ def make_vasp_kpoints(kpoints, kgamma=False):
 
 
 def make_vasp_kpoints_from_incar(work_dir, jdata):
+    from pymatgen.io.vasp import Incar, Kpoints
+
     cwd = os.getcwd()
     fp_aniso_kspacing = jdata.get("fp_aniso_kspacing")
     os.chdir(work_dir)
diff --git a/dpgen/auto_test/mpdb.py b/dpgen/auto_test/mpdb.py
@@ -1,13 +1,13 @@
 import os
 
-from pymatgen.ext.matproj import MPRester, MPRestError
-
 from dpgen import dlog
 
 web = "materials.org"
 
 
 def check_apikey():
+    from pymatgen.ext.matproj import MPRester, MPRestError
+
     try:
         apikey = os.environ["MAPI_KEY"]
     except KeyError:
diff --git a/dpgen/data/gen.py b/dpgen/data/gen.py
@@ -10,8 +10,6 @@
 
 import dpdata
 import numpy as np
-from pymatgen.core import Structure
-from pymatgen.io.vasp import Incar
 
 import dpgen.data.tools.bcc as bcc
 import dpgen.data.tools.diamond as diamond
@@ -300,6 +298,8 @@ def make_unit_cell_ABACUS(jdata):
 
 
 def make_super_cell(jdata):
+    from pymatgen.core import Structure
+
     out_dir = jdata["out_dir"]
     super_cell = jdata["super_cell"]
     path_uc = os.path.join(out_dir, global_dirname_02)
@@ -342,6 +342,8 @@ def make_super_cell_ABACUS(jdata, stru_data):
 
 
 def make_super_cell_poscar(jdata):
+    from pymatgen.core import Structure
+
     out_dir = jdata["out_dir"]
     super_cell = jdata["super_cell"]
     path_sc = os.path.join(out_dir, global_dirname_02)
@@ -1444,6 +1446,8 @@ def run_abacus_md(jdata, mdata):
 
 
 def gen_init_bulk(args):
+    from pymatgen.io.vasp import Incar
+
     jdata = load_file(args.PARAM)
     if args.MACHINE is not None:
         mdata = load_file(args.MACHINE)
diff --git a/dpgen/data/surf.py b/dpgen/data/surf.py
@@ -12,12 +12,7 @@
 from ase.build import general_surface
 
 # -----ASE-------
-from pymatgen.core import Element, Structure
-from pymatgen.io.ase import AseAtomsAdaptor
-
 # -----PMG---------
-from pymatgen.io.vasp import Poscar
-
 import dpgen.data.tools.bcc as bcc
 import dpgen.data.tools.diamond as diamond
 import dpgen.data.tools.fcc as fcc
@@ -168,6 +163,8 @@ def poscar_scale_direct(str_in, scale):
 
 
 def poscar_elong(poscar_in, poscar_out, elong, shift_center=True):
+    from pymatgen.core import Structure
+
     with open(poscar_in) as fin:
         lines = list(fin)
     if lines[7][0].upper() != "C":
@@ -215,6 +212,9 @@ def make_unit_cell(jdata):
 
 
 def make_super_cell_pymatgen(jdata):
+    from pymatgen.core import Element, Structure
+    from pymatgen.io.ase import AseAtomsAdaptor
+
     make_unit_cell(jdata)
     out_dir = jdata["out_dir"]
     path_uc = os.path.join(out_dir, global_dirname_02)
@@ -401,6 +401,8 @@ def poscar_scale_cartesian(str_in, scale):
 
 
 def poscar_scale(poscar_in, poscar_out, scale):
+    from pymatgen.io.vasp import Poscar
+
     with open(poscar_in) as fin:
         lines = list(fin)
     if "D" == lines[7][0] or "d" == lines[7][0]:
diff --git a/dpgen/database/entry.py b/dpgen/database/entry.py
diff --git a/dpgen/database/vasp.py b/dpgen/database/vasp.py
diff --git a/dpgen/generator/lib/ele_temp.py b/dpgen/generator/lib/ele_temp.py
diff --git a/dpgen/generator/lib/vasp.py b/dpgen/generator/lib/vasp.py
diff --git a/dpgen/generator/run.py b/dpgen/generator/run.py
diff --git a/pyproject.toml b/pyproject.toml
