Merge pull request #74 from QuantumChemist/main

Fully Integrate all MLIPs

Author: JaGeo

autoplex/auto/phonons/flows.py

@@ -150,6 +150,7 @@ class CompleteDFTvsMLBenchmarkWorkflow(Maker):
     atomwise_regularization_list: list | None = None
     soap_delta_list: list | None = None
     n_sparse_list: list | None = None
+    benchmark_kwargs: dict = field(default_factory=dict)
 
     def make(
         self,
@@ -263,7 +264,7 @@ def make(
                 mlip_hyper=ml_hyper,
             ).make(
                 species_list=isoatoms.output["species"],
-                isolated_atoms_energy=isoatoms.output["energies"],
+                isolated_atoms_energies=isoatoms.output["energies"],
                 fit_input=fit_input,
                 split_ratio=split_ratio,
                 f_max=f_max,
@@ -311,6 +312,7 @@ def make(
                         symprec=self.symprec,
                         phonon_displacement_maker=self.phonon_displacement_maker,
                         dft_references=dft_references,
+                        **self.benchmark_kwargs,
                     )
                     flows.append(complete_bm)
                     bm_outputs.append(complete_bm.output)
@@ -343,7 +345,7 @@ def make(
                                 mlip_type=ml_model, mlip_hyper=ml_hyper
                             ).make(
                                 species_list=isoatoms.output["species"],
-                                isolated_atoms_energy=isoatoms.output["energies"],
+                                isolated_atoms_energies=isoatoms.output["energies"],
                                 fit_input=fit_input,
                                 split_ratio=split_ratio,
                                 f_max=f_max,
@@ -379,6 +381,7 @@ def make(
                                         symprec=self.symprec,
                                         phonon_displacement_maker=self.phonon_displacement_maker,
                                         dft_references=dft_references,
+                                        **self.benchmark_kwargs,
                                     )
                                     flows.append(complete_bm)
                                     bm_outputs.append(complete_bm.output)

autoplex/auto/phonons/jobs.py

@@ -36,6 +36,9 @@ def complete_benchmark(  # this function was put here to prevent circular import
     symprec,
     phonon_displacement_maker: BaseVaspMaker,
     dft_references=None,
+    relax_maker_kwargs: dict | None = None,
+    static_maker_kwargs: dict | None = None,
+    **ml_phonon_maker_kwargs,
 ):
     """
     Construct a complete flow for benchmarking the MLIP fit quality using a DFT based phonon structure.
@@ -76,6 +79,12 @@ def complete_benchmark(  # this function was put here to prevent circular import
         Maker used to compute the forces for a supercell.
     dft_references:
         a list of DFT reference files containing the PhononBSDOCDoc object. Default None.
+    relax_maker_kwargs: dict
+        Keyword arguments that can be passed to the RelaxMaker.
+    static_maker_kwargs: dict
+        Keyword arguments that can be passed to the StaticMaker.
+    ml_phonon_maker_kwargs: dict
+        Keyword arguments that can be passed to the MLPhononMaker.
     """
     jobs = []
     collect_output = []
@@ -84,16 +93,34 @@ def complete_benchmark(  # this function was put here to prevent circular import
     if min_length >= 18:
         phonon_displacement_maker = TightDFTStaticMakerBigSupercells()
     for suffix in ["", "_wo_sigma", "_phonon", "_rand_struc"]:
-        if Path(Path(ml_path) / f"gap_file{suffix}.xml").exists():
-            # TODO: this needs to beextended for the other MLIPs
+        # _wo_sigma", "_phonon", "_rand_struc" only available for GAP at the moment
+        if ml_model == "GAP":
+            ml_potential = Path(ml_path) / f"gap_file{suffix}.xml"
+        elif ml_model == "J-ACE":
+            raise UserWarning("No atomate2 ACE.jl PhononMaker implemented.")
+        elif ml_model in ["M3GNET"]:
+            ml_potential = Path(ml_path.join(suffix)) / "training"
+            # M3GNet requires path
+            # also need to find a different solution for separated fit then
+        elif ml_model in ["NEQUIP"]:
+            ml_potential = Path(ml_path) / f"deployed_nequip_model{suffix}.pth"
+        else:  # MACE
+            ml_potential = Path(ml_path) / f"MACE_model{suffix}.model"
+
+        if Path(ml_potential).exists():
             add_data_ml_phonon = MLPhononMaker(
                 min_length=min_length,
+                relax_maker_kwargs=relax_maker_kwargs,
+                static_maker_kwargs=static_maker_kwargs,
             ).make_from_ml_model(
                 structure=benchmark_structure,
-                ml_model=ml_path,
-                suffix=suffix,
+                ml_model=ml_model,
+                potential_file=ml_potential,
+                **ml_phonon_maker_kwargs,
             )
             jobs.append(add_data_ml_phonon)
+
+            # DFT benchmark reference preparations
             if dft_references is None and benchmark_mp_ids is not None:
                 if (
                     benchmark_mp_ids[ibenchmark_structure] in mp_ids
@@ -322,18 +349,22 @@ def get_iso_atom(structure_list: list[Structure]):
         list of pymatgen Structure objects
     """
     jobs = []
+    iso_atoms_dict = {}
     all_species = list(
         {specie for s in structure_list for specie in s.types_of_species}
     )
 
     isoatoms = IsoAtomMaker().make(all_species=all_species)
     jobs.append(isoatoms)
 
+    for i, species in enumerate(all_species):
+        iso_atoms_dict.update({species.number: isoatoms.output["energies"][i]})
+
     flow = Flow(
         jobs,
         {
             "species": all_species,
-            "energies": isoatoms.output["energies"],
+            "energies": iso_atoms_dict,
             "dirs": isoatoms.output["dirs"],
         },
     )

autoplex/data/phonons/flows.py

@@ -17,6 +17,12 @@
     ForceFieldStaticMaker,
     GAPRelaxMaker,
     GAPStaticMaker,
+    M3GNetRelaxMaker,
+    M3GNetStaticMaker,
+    MACERelaxMaker,
+    MACEStaticMaker,
+    NequipRelaxMaker,
+    NequipStaticMaker,
 )
 from atomate2.vasp.flows.core import DoubleRelaxMaker
 from atomate2.vasp.flows.phonons import PhononMaker
@@ -27,6 +33,7 @@
 from pymatgen.core import Molecule, Site
 
 from autoplex.data.common.jobs import generate_randomized_structures
+from autoplex.data.phonons.utils import ml_phonon_maker_preparation
 
 __all__ = [
     "DFTPhononMaker",
@@ -332,17 +339,17 @@ class MLPhononMaker(FFPhononMaker):
           High-throughput electronic band structure calculations:
           Challenges and tools. Computational Materials Science,
           49(2), 299-312. doi:10.1016/j.commatsci.2010.05.010.
-          We will however use seekpath and primitive structures
+          We will, however, use seekpath and primitive structures
           as determined by from phonopy to compute the phonon band structure
-    bulk_relax_maker : .ForceFieldRelaxMaker or None
+    bulk_relax_maker: .ForceFieldRelaxMaker or None
         A maker to perform a tight relaxation on the bulk.
         Set to ``None`` to skip the
         bulk relaxation
-    static_energy_maker : .ForceFieldStaticMaker or None
+    static_energy_maker: .ForceFieldStaticMaker or None
         A maker to perform the computation of the DFT energy on the bulk.
         Set to ``None`` to skip the
         static energy computation
-    phonon_displacement_maker : .ForceFieldStaticMaker or None
+    phonon_displacement_maker: .ForceFieldStaticMaker or None
         Maker used to compute the forces for a supercell.
     generate_frequencies_eigenvectors_kwargs : dict
         Keyword arguments passed to :obj:`generate_frequencies_eigenvectors`.
@@ -364,6 +371,10 @@ class MLPhononMaker(FFPhononMaker):
         in the future
     store_force_constants: bool
         if True, force constants will be stored
+    relax_maker_kwargs: dict
+        Keyword arguments that can be passed to the RelaxMaker.
+    static_maker_kwargs: dict
+        Keyword arguments that can be passed to the StaticMaker.
     """
 
     name: str = "ml phonon"
@@ -374,7 +385,7 @@ class MLPhononMaker(FFPhononMaker):
         )
     )
     phonon_displacement_maker: ForceFieldStaticMaker | None = field(
-        default_factory=lambda: GAPStaticMaker(name="ml phonon static")
+        default_factory=lambda: GAPStaticMaker(name="gap phonon static")
     )
     static_energy_maker: ForceFieldStaticMaker | None = field(
         default_factory=lambda: GAPStaticMaker()
@@ -383,11 +394,18 @@ class MLPhononMaker(FFPhononMaker):
     generate_frequencies_eigenvectors_kwargs: dict = field(
         default_factory=lambda: {"units": "THz", "tol_imaginary_modes": 1e-1}
     )
-    relax_maker_kwargs: dict = field(default_factory=dict)
-    static_maker_kwargs: dict = field(default_factory=dict)
+    relax_maker_kwargs: dict | None = field(default_factory=dict)
+    static_maker_kwargs: dict | None = field(default_factory=dict)
 
     @job
-    def make_from_ml_model(self, structure, ml_model, suffix: str = "", **make_kwargs):
+    def make_from_ml_model(
+        self,
+        structure,
+        potential_file,
+        ml_model: str = "GAP",
+        calculator_kwargs: dict | None = None,
+        **make_kwargs,
+    ):
         """
         Maker for GAP phonon jobs.
 
@@ -397,10 +415,13 @@ def make_from_ml_model(self, structure, ml_model, suffix: str = "", **make_kwarg
             A pymatgen structure. Please start with a structure
             that is nearly fully optimized as the internal optimizers
             have very strict settings!
-        ml_model : str
-            Complete path to MLIP file(s).
-        suffix:
-            Train, test and MLIP suffix ("", "_wo_sigma", "_phonon", "_rand_struc").
+        ml_model: str
+            ML model to be used. Default is GAP.
+        potential_file :
+            Complete path to MLIP file(s)
+            Train, test and MLIP files (+ suffixes "", "_wo_sigma", "_phonon", "_rand_struc").
+        calculator_kwargs :
+            Keyword arguments for the ASE Calculator.
         make_kwargs :
             Keyword arguments for the PhononMaker.
 
@@ -409,41 +430,93 @@ def make_from_ml_model(self, structure, ml_model, suffix: str = "", **make_kwarg
         PhononMaker jobs.
 
         """
-        ml_model = ml_model + f"/gap_file{suffix}.xml"
-        if self.bulk_relax_maker is not None:
-            br = self.bulk_relax_maker
-            self.bulk_relax_maker = br.update_kwargs(
-                update={
-                    "calculator_kwargs": {
-                        "args_str": "IP GAP",
-                        "param_filename": str(ml_model),
-                    },
-                    **self.relax_maker_kwargs,
+        if ml_model == "GAP":
+            if calculator_kwargs is None:
+                calculator_kwargs = {
+                    "args_str": "IP GAP",
+                    "param_filename": str(potential_file),
                 }
+
+            ml_prep = ml_phonon_maker_preparation(
+                bulk_relax_maker=self.bulk_relax_maker,
+                phonon_displacement_maker=self.phonon_displacement_maker,
+                static_energy_maker=self.static_energy_maker,
+                calculator_kwargs=calculator_kwargs,
+                relax_maker_kwargs=self.relax_maker_kwargs,
+                static_maker_kwargs=self.static_maker_kwargs,
             )
-        if self.phonon_displacement_maker is not None:
-            ph_disp = self.phonon_displacement_maker
-            self.phonon_displacement_maker = ph_disp.update_kwargs(
-                update={
-                    "calculator_kwargs": {
-                        "args_str": "IP GAP",
-                        "param_filename": str(ml_model),
-                    },
-                    **self.static_maker_kwargs,
+
+        elif ml_model == "J-ACE":
+            raise UserWarning("No atomate2 ACE.jl PhononMaker implemented.")
+
+        elif ml_model == "NEQUIP":
+            if calculator_kwargs is None:
+                calculator_kwargs = {
+                    "model_path": str(potential_file),
+                    "device": "cuda",
                 }
+            else:
+                calculator_kwargs.update({"model_path": str(potential_file)})
+
+            ml_prep = ml_phonon_maker_preparation(
+                bulk_relax_maker=NequipRelaxMaker(
+                    relax_cell=True, relax_kwargs={"interval": 500}
+                ),
+                phonon_displacement_maker=NequipStaticMaker(
+                    name="nequip phonon static"
+                ),
+                static_energy_maker=NequipStaticMaker(),
+                calculator_kwargs=calculator_kwargs,
+                relax_maker_kwargs=self.relax_maker_kwargs,
+                static_maker_kwargs=self.static_maker_kwargs,
             )
-        if self.static_energy_maker is not None:
-            stat_en = self.static_energy_maker
-            self.static_energy_maker = stat_en.update_kwargs(
-                update={
-                    "calculator_kwargs": {
-                        "args_str": "IP GAP",
-                        "param_filename": str(ml_model),
-                    },
-                    **self.static_maker_kwargs,
-                }
+
+        elif ml_model == "M3GNET":
+            if calculator_kwargs is None:
+                calculator_kwargs = {"path": str(potential_file)}
+
+            ml_prep = ml_phonon_maker_preparation(
+                bulk_relax_maker=M3GNetRelaxMaker(
+                    relax_cell=True, relax_kwargs={"interval": 500}
+                ),
+                phonon_displacement_maker=M3GNetStaticMaker(
+                    name="m3gnet phonon static"
+                ),
+                static_energy_maker=M3GNetStaticMaker(),
+                calculator_kwargs=calculator_kwargs,
+                relax_maker_kwargs=self.relax_maker_kwargs,
+                static_maker_kwargs=self.static_maker_kwargs,
             )
 
+        else:  # MACE
+            if calculator_kwargs is None:
+                calculator_kwargs = {"model": str(potential_file), "device": "cuda"}
+            elif "model" in calculator_kwargs:
+                calculator_kwargs.update(
+                    {"default_dtype": "float64"}
+                )  # Use float64 for geometry optimization.
+            else:
+                calculator_kwargs.update(
+                    {"model": str(potential_file), "default_dtype": "float64"}
+                )
+
+            ml_prep = ml_phonon_maker_preparation(
+                bulk_relax_maker=MACERelaxMaker(
+                    relax_cell=True, relax_kwargs={"interval": 500}
+                ),
+                phonon_displacement_maker=MACEStaticMaker(name="mace phonon static"),
+                static_energy_maker=MACEStaticMaker(),
+                calculator_kwargs=calculator_kwargs,
+                relax_maker_kwargs=self.relax_maker_kwargs,
+                static_maker_kwargs=self.static_maker_kwargs,
+            )
+
+        (
+            self.bulk_relax_maker,
+            self.phonon_displacement_maker,
+            self.static_energy_maker,
+        ) = ml_prep
+
         flow = self.make(structure=structure, **make_kwargs)
         return Response(replace=flow, output=flow.output)