WIP: New Lobster workflow with additional speedup due to two subsequent Lobster runs

src/atomate2/lobster/flows.py

@@ -0,0 +1,91 @@
+"""Module defining lobster jobs."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from jobflow import Maker, job, Flow
+from pymatgen.electronic_structure.cohp import CompleteCohp
+from pymatgen.electronic_structure.dos import LobsterCompleteDos
+from pymatgen.io.lobster import Bandoverlaps, Icohplist, Lobsterin
+
+from atomate2 import SETTINGS
+from atomate2.common.files import gzip_output_folder
+from atomate2.lobster.files import (
+    LOBSTEROUTPUT_FILES,
+    VASP_OUTPUT_FILES,
+    copy_lobster_files,
+)
+from atomate2.lobster.jobs import LobsterMaker, retrieve_relevant_bonds
+from atomate2.lobster.run import run_lobster
+from atomate2.lobster.schemas import LobsterTaskDocument
+
+
+logger = logging.getLogger(__name__)
+
+
+_FILES_TO_ZIP = [*LOBSTEROUTPUT_FILES, "lobsterin", *VASP_OUTPUT_FILES]
+
+
+
+@dataclass
+class AdvancedLobsterMaker(Maker):
+    """
+    LOBSTER job maker with additional speedup.
+
+    1. The maker copies DFT output files necessary for the LOBSTER run.
+    2. It will create all lobsterin files, run LOBSTER several times,
+        zip the outputs and parse the LOBSTER outputs. In this step, the COHP/COBI/COOP
+        curves will only be written with a limited accuracy.
+    3. After an analysis of the most important bonds, COHP/COBI/COOP curves for those will
+        be computed with high accuracy.
+
+    In the future, this workflow could also benefit from further symmetry considerations.
+
+    Parameters
+    ----------
+    name : str
+        Name of jobs produced by this maker.
+    task_document_kwargs : dict
+        Keyword arguments passed to :obj:`.LobsterTaskDocument.from_directory`.
+    user_lobsterin_settings : dict
+        Dict including additional information on the Lobster settings.
+    run_lobster_kwargs : dict
+        Keyword arguments that will get passed to :obj:`.run_lobster`.
+    calculation_type : str
+        Type of calculation for the Lobster run that will get passed to
+        :obj:`.Lobsterin.standard_calculations_from_vasp_files`.
+    """
+    name: str = "lobster"
+    lobster_maker_1: LobsterMaker=field(default_factory=lambda:LobsterMaker(user_lobsterin_settings={"cohpsteps":1}))
+    lobster_maker_2: LobsterMaker=field(default_factory=LobsterMaker)
+
+    def make(
+            self,
+            wavefunction_dir: str | Path = None,
+            basis_dict: dict | None = None,
+    ) -> LobsterTaskDocument:
+        """Run a LOBSTER calculation.
+
+        Parameters
+        ----------
+        wavefunction_dir : str or Path
+            A directory containing a WAVEFUNCTION and other outputs needed for Lobster
+        basis_dict: dict
+            A dict including information on the basis set
+        """
+        # TODO: why is calc quality failing?
+        jobs=[]
+        lobster_1=self.lobster_maker_1.make(wavefunction_dir, basis_dict)
+        jobs.append(lobster_1)
+
+        # code to postprocess the lobster data and identify relevant bonds
+        # switch between cation, anion modes and potentially different ways to indentify the bonds
+        cohp_between_dict=retrieve_relevant_bonds(lobster_1.output.lobsterpy_data)
+
+        # think about how to modify the input dict during the run time
+        lobster_2 = self.lobster_maker_2.make(wavefunction_dir, basis_dict, cohp_between_dict)
+        jobs.append(lobster_2)
+        return Flow(jobs=jobs, output=lobster_2.output)

src/atomate2/lobster/jobs.py

@@ -27,6 +27,7 @@
 _FILES_TO_ZIP = [*LOBSTEROUTPUT_FILES, "lobsterin", *VASP_OUTPUT_FILES]
 
 
+
 @dataclass
 class LobsterMaker(Maker):
     """
@@ -69,6 +70,7 @@ def make(
         self,
         wavefunction_dir: str | Path = None,
         basis_dict: dict | None = None,
+        cohp_between_dict: dict | None = None,
     ) -> LobsterTaskDocument:
         """Run a LOBSTER calculation.
 
@@ -78,6 +80,8 @@ def make(
             A directory containing a WAVEFUNCTION and other outputs needed for Lobster
         basis_dict: dict
             A dict including information on the basis set
+        cohp_between_dict: dict
+            A dict including information on the bonds that should be analysed.
         """
         # copy previous inputs # VASP for example
         copy_lobster_files(wavefunction_dir)
@@ -93,6 +97,13 @@ def make(
                 if key != "basisfunctions":
                     lobsterin[key] = parameter
 
+        if cohp_between_dict:
+            # add code to only compute specific interactions
+            # ideally used without cohpgenerator for speedup
+            # cohpbetween
+            pass
+
+
         lobsterin.write_lobsterin("lobsterin")
 
         # run lobster
@@ -111,3 +122,11 @@ def make(
             Path.cwd(),
             **self.task_document_kwargs,
         )
+
+
+@job
+def retrieve_relevant_bonds(condensed_bonding_analysis):
+    logging.log("test")
+    print(condensed_bonding_analysis.sites)
+
+    return None

src/atomate2/lobster/schemas.py

@@ -745,6 +745,7 @@ def from_directory(
         additional_fields: dict = None,
         add_coxxcar_to_task_document: bool = False,
         analyze_outputs: bool = True,
+        skip_calc_quality: bool = False,
         calc_quality_kwargs: dict = None,
         lobsterpy_kwargs: dict = None,
         plot_kwargs: dict = None,
@@ -766,6 +767,8 @@ def from_directory(
             to the task document.
         analyze_outputs: bool.
             If True, will enable lobsterpy analysis.
+        skip_calc_quality: bool. 
+            The calc quality analysis will be skipped.
         calc_quality_kwargs : dict.
             kwargs to change calc quality summary options in lobsterpy.
         lobsterpy_kwargs : dict.
@@ -873,15 +876,15 @@ def from_directory(
                     which_bonds="cation-anion",
                 )
             # Get lobster calculation quality summary data
-
-            calc_quality_summary = CalcQualitySummary.from_directory(
-                dir_name,
-                calc_quality_kwargs=calc_quality_kwargs,
-            )
-
-            calc_quality_text = Description.get_calc_quality_description(
-                calc_quality_summary.model_dump()
-            )
+            if not skip_calc_quality: 
+                calc_quality_summary = CalcQualitySummary.from_directory(
+                    dir_name,
+                    calc_quality_kwargs=calc_quality_kwargs,
+                )
+    
+                calc_quality_text = Description.get_calc_quality_description(
+                    calc_quality_summary.model_dump()
+                )
 
         # Read in charges
         charges = None
@@ -1054,14 +1057,15 @@ def from_directory(
                     data, allow_bson=True, strict=True, enum_values=True
                 )
                 json.dump(monty_encoded_json_doc, file)
-                file.write(",")
-                data = {
-                    "calc_quality_text": ["".join(doc.calc_quality_text)]  # type: ignore[dict-item]
-                }  # add calc quality summary dict
-                monty_encoded_json_doc = jsanitize(
-                    data, allow_bson=True, strict=True, enum_values=True
-                )
-                json.dump(monty_encoded_json_doc, file)
+                if not skip_calc_quality:
+                    file.write(",")
+                    data = {
+                        "calc_quality_text": ["".join(doc.calc_quality_text)]  # type: ignore[dict-item]
+                    }  # add calc quality summary dict
+                    monty_encoded_json_doc = jsanitize(
+                        data, allow_bson=True, strict=True, enum_values=True
+                    )
+                    json.dump(monty_encoded_json_doc, file)
                 file.write(",")
                 data = {"dos": doc.dos}  # add NON LSO of lobster
                 monty_encoded_json_doc = jsanitize(

tests/vasp/lobster/flows/test_lobster.py

@@ -2,6 +2,7 @@
 from pymatgen.core.structure import Structure
 
 from atomate2.lobster.jobs import LobsterMaker
+from atomate2.lobster.flows import AdvancedLobsterMaker
 from atomate2.lobster.schemas import LobsterTaskDocument
 from atomate2.vasp.flows.lobster import VaspLobsterMaker
 from atomate2.vasp.flows.mp import MPVaspLobsterMaker
@@ -327,3 +328,99 @@ def test_mp_vasp_lobstermaker(
     )
 
     assert isinstance(task_doc, LobsterTaskDocument)
+
+def test_improved_lobster_maker(mock_vasp, mock_lobster, clean_dir, memory_jobstore, si_structure: Structure):
+    # mapping from job name to directory containing test files
+    ref_paths = {
+        "relax 1": "Si_lobster_uniform/relax_1",
+        "relax 2": "Si_lobster_uniform/relax_2",
+        "static": "Si_lobster_uniform/static",
+        "non-scf uniform": "Si_lobster_uniform/non-scf_uniform",
+    }
+
+    # settings passed to fake_run_vasp; adjust these to check for certain INCAR settings
+    fake_run_vasp_kwargs = {
+        "relax 1": {"incar_settings": ["NSW", "ISMEAR"]},
+        "relax 2": {"incar_settings": ["NSW", "ISMEAR"]},
+        "static": {
+            "incar_settings": [
+                "NSW",
+                "LWAVE",
+                "ISMEAR",
+                "ISYM",
+                "NBANDS",
+                "ISPIN",
+                "LCHARG",
+            ],
+            # TODO restore POSCAR input checking e.g. when next updating test files
+            "check_inputs": ["potcar", "kpoints", "incar"],
+        },
+        "non-scf uniform": {
+            "incar_settings": [
+                "NSW",
+                "LWAVE",
+                "ISMEAR",
+                "ISYM",
+                "NBANDS",
+                "ISPIN",
+                "ICHARG",
+            ],
+            # TODO restore POSCAR input checking e.g. when next updating test files
+            "check_inputs": ["potcar", "kpoints", "incar"],
+        },
+    }
+
+    ref_paths_lobster = {
+        "lobster_run_0": "Si_lobster/lobster_0",
+    }
+
+    # settings passed to fake_run_vasp; adjust these to check for certain INCAR settings
+    fake_run_lobster_kwargs = {
+        "lobster_run_0": {"lobsterin_settings": ["basisfunctions"]},
+    }
+
+    # automatically use fake VASP and write POTCAR.spec during the test
+    mock_vasp(ref_paths, fake_run_vasp_kwargs)
+    mock_lobster(ref_paths_lobster, fake_run_lobster_kwargs)
+
+    job = VaspLobsterMaker(
+        lobster_maker=AdvancedLobsterMaker(lobster_maker_1=LobsterMaker(
+            task_document_kwargs={
+                "calc_quality_kwargs": {"potcar_symbols": ["Si"], "n_bins": 10},
+                "add_coxxcar_to_task_document": True,
+                "skip_calc_quality":True,
+            },
+            user_lobsterin_settings={
+                "COHPstartEnergy": -5.0,
+                "COHPEndEnergy": 5.0,
+                "cohpGenerator": "from 0.1 to 3.0 orbitalwise",
+            },
+        ),
+        lobster_maker_2=LobsterMaker(
+            task_document_kwargs={
+                "calc_quality_kwargs": {"potcar_symbols": ["Si"], "n_bins": 10},
+                "add_coxxcar_to_task_document": True,
+                "skip_calc_quality": True,
+            },
+            user_lobsterin_settings={
+                "COHPstartEnergy": -5.0,
+                "COHPEndEnergy": 5.0,
+                "cohpGenerator": "from 0.1 to 3.0 orbitalwise",
+            },
+        )),
+    delete_wavecars = False).make(si_structure)
+
+    job = update_user_incar_settings(job, {"NPAR": 4})
+
+    # run the flow or job and ensure that it finished running successfully
+    responses = run_locally(
+        job, store=memory_jobstore, create_folders=True, ensure_success=True
+    )
+
+    assert isinstance(
+        responses[job.jobs[-1].uuid][1]
+        .replace.output["lobster_task_documents"][0]
+        .resolve(memory_jobstore),
+        LobsterTaskDocument,
+    )
+