More tidying

Author: utf

src/atomate2/lobster/jobs.py

@@ -26,9 +26,8 @@ class LobsterMaker(Maker):
     """
     LOBSTER job maker.
 
-    The maker copies DFT output files necessary for the LOBSTER run.
-    It will create all lobsterin files, run LOBSTER, zip the outputs and parse the
-    LOBSTER outputs.
+    The maker copies DFT output files necessary for the LOBSTER run. It will create all
+    lobsterin files, run LOBSTER, zip the outputs and parse the LOBSTER outputs.
 
     Parameters
     ----------
@@ -91,8 +90,7 @@ def make(
         gzip_dir(".")
 
         # parse lobster outputs
-        task_doc = LobsterTaskDocument.from_directory(
+        return LobsterTaskDocument.from_directory(
             Path.cwd(),
             **self.task_document_kwargs,
         )
-        return task_doc

src/atomate2/lobster/run.py

@@ -19,10 +19,7 @@
 
 from atomate2 import SETTINGS
 
-__all__ = [
-    "JobType",
-    "run_lobster",
-]
+__all__ = ["JobType", "run_lobster"]
 
 _DEFAULT_VALIDATORS = (LobsterFilesValidator(), EnoughBandsValidator())
 _DEFAULT_HANDLERS = ()

src/atomate2/lobster/schemas.py

@@ -31,7 +31,13 @@
     Analysis = None
     Description = None
 
-__all__ = ["LobsterTaskDocument"]
+__all__ = [
+    "LobsterTaskDocument",
+    "LobsteroutModel",
+    "LobsterinModel",
+    "CondensedBondingAnalysis",
+    "StrongestBonds",
+]
 
 logger = logging.getLogger(__name__)
 
@@ -432,9 +438,9 @@ def from_directory(
 
         icohplist_path = dir_name / "ICOHPLIST.lobster.gz"
         cohpcar_path = dir_name / "COHPCAR.lobster.gz"
+        charge_path = dir_name / "CHARGE.lobster.gz"
         cobicar_path = dir_name / "COBICAR.lobster.gz"
         coopcar_path = dir_name / "COOPCAR.lobster.gz"
-        charge_path = dir_name / "CHARGE.lobster.gz"
         doscar_path = dir_name / "DOSCAR.lobster.gz"
         structure_path = dir_name / "POSCAR.gz"
         madelung_energies_path = dir_name / "MadelungEnergies.lobster.gz"
@@ -638,72 +644,37 @@ def _get_strong_bonds(
         lengths.append(l)
 
     bond_labels_unique = list(set(bonds))
-    sep_blabels: List[List[str]] = [[] for _ in range(len(bond_labels_unique))]
     sep_icohp: List[List[float]] = [[] for _ in range(len(bond_labels_unique))]
     sep_lengths: List[List[float]] = [[] for _ in range(len(bond_labels_unique))]
 
     for i, val in enumerate(bond_labels_unique):
         for j, val2 in enumerate(bonds):
             if val == val2:
-                sep_blabels[i].append(val2)
                 sep_icohp[i].append(icohp_all[j])
                 sep_lengths[i].append(lengths[j])
-    if not are_cobis and not are_coops:
-        bond_dict = {}
-        for i, lab in enumerate(bond_labels_unique):
-            label = lab.split("-")
-            label.sort()
-            for rel_bnd in relevant_bonds:
-                rel_bnd_list = rel_bnd.split("-")
-                rel_bnd_list.sort()
-                if label == rel_bnd_list:
-                    index = np.argmin(sep_icohp[i])
-                    bond_dict.update(
-                        {
-                            rel_bnd: {
-                                "ICOHP": min(sep_icohp[i]),
-                                "length": sep_lengths[i][index],
-                            }
-                        }
-                    )
-        return bond_dict
 
     if are_cobis and not are_coops:
-        bond_dict = {}
-        for i, lab in enumerate(bond_labels_unique):
-            label = lab.split("-")
-            label.sort()
-            for rel_bnd in relevant_bonds:
-                rel_bnd_list = rel_bnd.split("-")
-                rel_bnd_list.sort()
-                if label == rel_bnd_list:
-                    index = np.argmax(sep_icohp[i])
-                    bond_dict.update(
-                        {
-                            rel_bnd: {
-                                "ICOBI": max(sep_icohp[i]),
-                                "length": sep_lengths[i][index],
-                            }
+        prop = "ICOBI"
+    elif not are_cobis and are_coops:
+        prop = "ICOOP"
+    else:
+        prop = "ICOHP"
+
+    bond_dict = {}
+    for i, lab in enumerate(bond_labels_unique):
+        label = lab.split("-")
+        label.sort()
+        for rel_bnd in relevant_bonds:
+            rel_bnd_list = rel_bnd.split("-")
+            rel_bnd_list.sort()
+            if label == rel_bnd_list:
+                index = np.argmin(sep_icohp[i])
+                bond_dict.update(
+                    {
+                        rel_bnd: {
+                            prop: min(sep_icohp[i]),
+                            "length": sep_lengths[i][index],
                         }
-                    )
-        return bond_dict
-
-    if not are_cobis and are_coops:
-        bond_dict = {}
-        for i, lab in enumerate(bond_labels_unique):
-            label = lab.split("-")
-            label.sort()
-            for rel_bnd in relevant_bonds:
-                rel_bnd_list = rel_bnd.split("-")
-                rel_bnd_list.sort()
-                if label == rel_bnd_list:
-                    index = np.argmax(sep_icohp[i])
-                    bond_dict.update(
-                        {
-                            rel_bnd: {
-                                "ICOOP": max(sep_icohp[i]),
-                                "length": sep_lengths[i][index],
-                            }
-                        }
-                    )
-        return bond_dict
+                    }
+                )
+    return bond_dict

src/atomate2/vasp/flows/lobster.py

@@ -47,19 +47,19 @@ class VaspLobsterMaker(Maker):
     relax_maker : .BaseVaspMaker or None
         A maker to perform a relaxation on the bulk. Set to ``None`` to skip the
         bulk relaxation.
-    preconverge_static_maker: .BaseVaspMaker or None
+    preconverge_static_maker : .BaseVaspMaker or None
         A maker to perform a preconvergence run before the wavefunction computation
         without symmetry
     lobster_static_maker : .BaseVaspMaker
         A maker to perform the computation of the wavefunction before the static run.
         Cannot be skipped.
     lobster_maker : .LobsterMaker
         A maker to perform the Lobster run.
-    delete_all_wavecars : bool
+    delete_wavecars : bool
         If true, all WAVECARs will be deleted after the run.
     address_min_basis : str
         A path to a yaml file including basis set information.
-    address_max_basis: str
+    address_max_basis : str
        A path to a yaml file including basis set information.
     """
 
@@ -79,7 +79,7 @@ class VaspLobsterMaker(Maker):
         )
     )
     lobster_maker: LobsterMaker | None = field(default_factory=lambda: LobsterMaker())
-    delete_all_wavecars: bool = True
+    delete_wavecars: bool = True
     address_min_basis: str | None = None
     address_max_basis: str | None = None
 
@@ -151,7 +151,6 @@ def make(
         lobster_jobs = get_lobster_jobs(
             lobster_maker=self.lobster_maker,
             basis_dict=basis_infos.output["basis_dict"],
-            wavefunction_dir=lobster_static.output.dir_name,
             optimization_dir=optimization_dir,
             optimization_uuid=optimization_uuid,
             static_dir=lobster_static_dir,
@@ -162,7 +161,7 @@ def make(
         jobs.append(lobster_jobs)
 
         # delete all WAVECARs that have been copied
-        if self.delete_all_wavecars:
+        if self.delete_wavecars:
             delete_wavecars = delete_lobster_wavecar(
                 dirs=lobster_jobs.output["lobster_dirs"],
                 lobster_static_dir=lobster_static.output.dir_name,

src/atomate2/vasp/jobs/lobster.py

@@ -133,7 +133,7 @@ def get_basis_infos(
 
 @job
 def update_user_incar_settings_maker(
-    vasp_maker: LobsterStaticMaker,
+    vasp_maker: BaseVaspMaker,
     nbands: int,
     structure: Structure,
     prev_vasp_dir: Path | str,
@@ -167,7 +167,6 @@ def update_user_incar_settings_maker(
 def get_lobster_jobs(
     lobster_maker: LobsterMaker,
     basis_dict: dict,
-    wavefunction_dir: Path | str,
     optimization_dir: Path | str,
     optimization_uuid: str,
     static_dir: Path | str,
@@ -184,16 +183,14 @@ def get_lobster_jobs(
         maker for the Lobster jobs
     basis_dict : dict
         dict including basis set information.
-    wavefunction_dir : Path or str
-        Path to VASP calculation with WAVECAR
     optimization_dir : Path or str
         Path to optimization run.
     optimization_uuid : str
         uuid of optimization run.
     static_dir : Path or str
-        Path to static VASP calculation.
+        Path to static VASP calculation containing the WAVECAR.
     static_uuid : str
-        uuid of static run.
+        Uuid of static run.
     preconverge_static_dir : Path or str
         Path to preconvergence step.
     preconverge_static_uuid : str
@@ -221,9 +218,7 @@ def get_lobster_jobs(
         lobster_maker = LobsterMaker()
 
     for i, basis in enumerate(basis_dict):
-        lobsterjob = lobster_maker.make(
-            wavefunction_dir=wavefunction_dir, basis_dict=basis
-        )
+        lobsterjob = lobster_maker.make(wavefunction_dir=static_dir, basis_dict=basis)
         lobsterjob.append_name(f"_run_{i}")
         outputs["lobster_uuids"].append(lobsterjob.output.uuid)
         outputs["lobster_dirs"].append(lobsterjob.output.dir_name)

tests/vasp/lobster/flows/test_lobster.py

@@ -57,7 +57,7 @@ def test_lobstermaker(mock_vasp, mock_lobster, clean_dir, memory_jobstore):
                 "cohpGenerator": "from 0.1 to 3.0 orbitalwise",
             }
         ),
-        delete_all_wavecars=False,
+        delete_wavecars=False,
     ).make(si_structure)
     job = update_user_incar_settings(job, {"NPAR": 4})
 
@@ -123,7 +123,7 @@ def test_lobstermaker_delete(mock_vasp, mock_lobster, clean_dir, memory_jobstore
                 "cohpGenerator": "from 0.1 to 3.0 orbitalwise",
             }
         ),
-        delete_all_wavecars=True,
+        delete_wavecars=True,
     ).make(si_structure)
     job = update_user_incar_settings(job, {"NPAR": 4})