Make the MACE fit more flexible and build a finetuning workflow (#182)

* make the mace interface more flexibel, start

* attempt a fix for virial but more is necessary

* fix typo and add more handling of xyz names

* more flexibility to mace fitting

* add more test files

* update installation

* commit latest changes

* add finetuning test and add a preliminary dependency

* fix precommit

* fix linting

* fix linting

* fix linting

* fix pyproject

* fix pyproject

* fix pyproject

* update torch dependencies

* clear uv cache before installs

* fix stupid typo

* add finetuning option and test data, document it

* add input set to mock vasp

* uncomment one test

* fix finetuning

* fix linting

* fix linting

* fix linting

* fix linting

* fix documentation

* fix documentation

* remove test files

* remove test files

* fix test folders

* mention that one might need to download model

* fix doc

* fix dataset

---------

Co-authored-by: Aakash Ashok Naik <91958822+naik-aakash@users.noreply.github.com>

autoplex/auto/phonons/flows.py

@@ -5,6 +5,16 @@
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
+from atomate2.vasp.flows.mp import (
+    MPGGADoubleRelaxMaker,
+    MPGGARelaxMaker,
+    MPGGAStaticMaker,
+)
+from pymatgen.io.vasp.sets import (
+    MPRelaxSet,
+    MPStaticSet,
+)
+
 from autoplex.data.phonons.flows import TightDFTStaticMaker
 from autoplex.fitting.common.utils import (
     MLIP_PHONON_DEFAULTS_FILE_PATH,
@@ -33,7 +43,11 @@
 from autoplex.benchmark.phonons.jobs import write_benchmark_metrics
 from autoplex.fitting.common.flows import MLIPFitMaker
 
-__all__ = ["CompleteDFTvsMLBenchmarkWorkflow", "DFTSupercellSettingsMaker"]
+__all__ = [
+    "CompleteDFTvsMLBenchmarkWorkflow",
+    "CompleteDFTvsMLBenchmarkWorkflowMPSettings",
+    "DFTSupercellSettingsMaker",
+]
 
 
 # Volker's idea: provide several default flows with different setting/setups
@@ -194,6 +208,7 @@ def make(
         dft_references: list[PhononBSDOSDoc] | None = None,
         benchmark_structures: list[Structure] | None = None,
         benchmark_mp_ids: list[str] | None = None,
+        use_defaults_fitting: bool = True,
         **fit_kwargs,
     ):
         """
@@ -314,6 +329,7 @@ def make(
             add_data_fit = MLIPFitMaker(
                 mlip_type=ml_model,
                 glue_file_path=self.glue_file_path,
+                use_defaults=use_defaults_fitting,
             ).make(
                 species_list=isoatoms.output["species"],
                 isolated_atoms_energies=isoatoms.output["energies"],
@@ -598,6 +614,224 @@ def add_dft_random(
         return additonal_dft_random
 
 
+@dataclass
+class CompleteDFTvsMLBenchmarkWorkflowMPSettings(CompleteDFTvsMLBenchmarkWorkflow):
+    """
+    Maker to construct a DFT (VASP) based dataset, composed of the following two configuration types.
+
+    (1) single atom displaced supercells (based on the atomate2 PhononMaker subroutines)
+    (2) supercells with randomly displaced atoms (based on the ase rattled function).
+
+    Machine-learned interatomic potential(s) are then fitted on the dataset, followed by
+    benchmarking the resulting potential(s) to DFT (VASP) level using the provided benchmark
+    structure(s) and comparing the respective DFT and MLIP-based Phonon calculations.
+    The benchmark metrics are provided in form of a phonon band structure comparison and
+    q-point-wise phonons RMSE plots, as well as a summary text file.
+
+    Parameters
+    ----------
+    name : str
+        Name of the flow produced by this maker.
+    add_dft_phonon_struct: bool.
+        If True, will add displaced supercells via phonopy for DFT calculation.
+    add_dft_random_struct: bool.
+        If True, will add randomly distorted structures for DFT calculation.
+    add_rss_struct: bool.
+        If True, will add RSS generated structures for DFT calculation.
+        n_structures: int.
+        The total number of randomly displaced structures to be generated.
+    displacement_maker: BaseVaspMaker
+        Maker used for a static calculation for a supercell.
+    phonon_bulk_relax_maker: BaseVaspMaker
+        Maker used for the bulk relax unit cell calculation.
+    rattled_bulk_relax_maker: BaseVaspMaker
+        Maker used for the bulk relax unit cell calculation.
+    phonon_static_energy_maker: BaseVaspMaker
+        Maker used for the static energy unit cell calculation.
+    isolated_atom_maker: IsoAtomStaticMaker
+        VASP maker for the isolated atom calculation.
+    n_structures : int.
+        Total number of distorted structures to be generated.
+        Must be provided if distorting volume without specifying a range, or if distorting angles.
+        Default=10.
+    displacements: list[float]
+        displacement distances for phonons
+    symprec: float
+        Symmetry precision to use in the
+        reduction of symmetry to find the primitive/conventional cell
+        (use_primitive_standard_structure, use_conventional_standard_structure)
+        and to handle all symmetry-related tasks in phonopy.
+    uc: bool.
+        If True, will generate randomly distorted structures (unitcells)
+        and add static computation jobs to the flow.
+    distort_type : int.
+        0- volume distortion, 1- angle distortion, 2- volume and angle distortion. Default=0.
+    volume_scale_factor_range : list[float]
+        [min, max] of volume scale factors.
+        e.g. [0.90, 1.10] will distort volume +-10%.
+    volume_custom_scale_factors : list[float]
+        Specify explicit scale factors (if range is not specified).
+        If None, will default to [0.90, 0.95, 0.98, 0.99, 1.01, 1.02, 1.05, 1.10].
+    min_distance: float
+        Minimum separation allowed between any two atoms.
+        Default= 1.5A.
+    angle_percentage_scale: float
+        Angle scaling factor.
+        Default= 10 will randomly distort angles by +-10% of original value.
+    angle_max_attempts: int.
+        Maximum number of attempts to distort structure before aborting.
+        Default=1000.
+    w_angle: list[float]
+        List of angle indices to be changed i.e. 0=alpha, 1=beta, 2=gamma.
+        Default= [0, 1, 2].
+    rattle_type: int.
+        0- standard rattling, 1- Monte-Carlo rattling. Default=0.
+    rattle_std: float.
+        Rattle amplitude (standard deviation in normal distribution).
+        Default=0.01.
+        Note that for MC rattling, displacements generated will roughly be
+        rattle_mc_n_iter**0.5 * rattle_std for small values of n_iter.
+    rattle_seed: int.
+        Seed for setting up NumPy random state from which random numbers are generated.
+        Default=42.
+    rattle_mc_n_iter: int.
+        Number of Monte Carlo iterations.
+        Larger number of iterations will generate larger displacements.
+        Default=10.
+    ml_models: list[str]
+        list of the ML models to be used. Default is GAP.
+    hyper_para_loop: bool
+        making it easier to loop through several hyperparameter sets.
+    atomwise_regularization_list: list
+        List of atom-wise regularization parameters that are checked.
+    soap_delta_list: list
+        List of SOAP delta values that are checked.
+    n_sparse_list: list
+        List of GAP n_sparse values that are checked.
+    supercell_settings: dict
+        settings for supercell generation
+    benchmark_kwargs: dict
+        kwargs for the benchmark flows
+    summary_filename_prefix: str
+        Prefix of the result summary file.
+    glue_file_path: str
+        Name of the glue.xml file path.
+    """
+
+    phonon_bulk_relax_maker: BaseVaspMaker = field(
+        default_factory=lambda: MPGGADoubleRelaxMaker.from_relax_maker(
+            MPGGARelaxMaker(
+                run_vasp_kwargs={"handlers": ()},
+                input_set_generator=MPRelaxSet(
+                    force_gamma=True,
+                    auto_metal_kpoints=True,
+                    inherit_incar=False,
+                    user_incar_settings={
+                        "NPAR": 4,
+                        "EDIFF": 1e-7,
+                        "EDIFFG": 1e-6,
+                        "ALGO": "NORMAL",
+                        "ISPIN": 1,
+                        "LREAL": False,
+                        "LCHARG": False,
+                        "ISMEAR": 0,
+                        "KSPACING": 0.2,
+                    },
+                ),
+            )
+        )
+    )
+    rattled_bulk_relax_maker: BaseVaspMaker = field(
+        default_factory=lambda: MPGGADoubleRelaxMaker.from_relax_maker(
+            MPGGARelaxMaker(
+                run_vasp_kwargs={"handlers": ()},
+                input_set_generator=MPRelaxSet(
+                    force_gamma=True,
+                    auto_metal_kpoints=True,
+                    inherit_incar=False,
+                    user_incar_settings={
+                        "NPAR": 4,
+                        "EDIFF": 1e-7,
+                        "EDIFFG": 1e-6,
+                        "ALGO": "NORMAL",
+                        "ISPIN": 1,
+                        "LREAL": False,
+                        "LCHARG": False,
+                        "ISMEAR": 0,
+                        "KSPACING": 0.2,
+                    },
+                ),
+            )
+        )
+    )
+    displacement_maker: BaseVaspMaker = field(
+        default_factory=lambda: MPGGAStaticMaker(
+            run_vasp_kwargs={"handlers": ()},
+            name="dft phonon static",
+            input_set_generator=MPStaticSet(
+                force_gamma=True,
+                auto_metal_kpoints=True,
+                inherit_incar=False,
+                user_incar_settings={
+                    "NPAR": 4,
+                    "EDIFF": 1e-7,
+                    "EDIFFG": 1e-6,
+                    "ALGO": "NORMAL",
+                    "ISPIN": 1,
+                    "LREAL": False,
+                    "LCHARG": False,
+                    "ISMEAR": 0,
+                    "KSPACING": 0.2,
+                },
+            ),
+        )
+    )
+    isolated_atom_maker: BaseVaspMaker = field(
+        default_factory=lambda: MPGGAStaticMaker(
+            run_vasp_kwargs={"handlers": ()},
+            input_set_generator=MPStaticSet(
+                user_kpoints_settings={"reciprocal_density": 1},
+                force_gamma=True,
+                auto_metal_kpoints=True,
+                inherit_incar=False,
+                user_incar_settings={
+                    "NPAR": 4,
+                    "EDIFF": 1e-7,
+                    "EDIFFG": 1e-6,
+                    "ALGO": "NORMAL",
+                    "ISPIN": 1,
+                    "LREAL": False,
+                    "LCHARG": False,
+                    "ISMEAR": 0,
+                },
+            ),
+        )
+    )
+
+    phonon_static_energy_maker: BaseVaspMaker = field(
+        default_factory=lambda: MPGGAStaticMaker(
+            run_vasp_kwargs={"handlers": ()},
+            name="dft phonon static",
+            input_set_generator=MPStaticSet(
+                force_gamma=True,
+                auto_metal_kpoints=True,
+                inherit_incar=False,
+                user_incar_settings={
+                    "NPAR": 4,
+                    "EDIFF": 1e-7,
+                    "EDIFFG": 1e-6,
+                    "ALGO": "NORMAL",
+                    "ISPIN": 1,
+                    "LREAL": False,
+                    "LCHARG": False,
+                    "ISMEAR": 0,
+                    "KSPACING": 0.2,
+                },
+            ),
+        )
+    )
+
+
 @dataclass
 class DFTSupercellSettingsMaker(Maker):
     """

autoplex/auto/phonons/jobs.py

@@ -135,8 +135,12 @@ def complete_benchmark(  # this function was put here to prevent circular import
         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"
-
+            # treat finetuned potentials
+            ml_potential_fine = Path(ml_path) / f"MACE_final{suffix}.model"
+            if ml_potential_fine.exists():
+                ml_potential = ml_potential_fine
+            else:
+                ml_potential = Path(ml_path) / f"MACE_model{suffix}.model"
         if Path(ml_potential).exists():
             add_data_ml_phonon = MLPhononMaker(
                 relax_maker_kwargs=relax_maker_kwargs,

autoplex/data/common/jobs.py

@@ -69,6 +69,7 @@ def convert_to_extxyz(job_output, pkl_file, config_type, factor):
     data[-1].write("tmp.xyz")
     file = read("tmp.xyz", index=":")
     for i in file:
+        # TODO: enable switching to stress
         virial_list = -voigt_6_to_full_3x3_stress(i.get_stress()) * i.get_volume()
         i.info["REF_virial"] = " ".join(map(str, virial_list.flatten()))
         del i.calc.results["stress"]

autoplex/fitting/common/flows.py

@@ -50,6 +50,8 @@ class MLIPFitMaker(Maker):
         Reference virial name.
     glue_file_path: str
         Name of the glue.xml file path.
+    use_defaults: bool
+        if true, uses default fit parameters
     """
 
     name: str = "MLpotentialFit"
@@ -59,6 +61,7 @@ class MLIPFitMaker(Maker):
     ref_force_name: str = "REF_forces"
     ref_virial_name: str = "REF_virial"
     glue_file_path: str = "glue.xml"
+    use_defaults: bool = True
 
     # TO DO: Combine parameters used only for gap into one category (as noted below),
     # otherwise it will be too specific.
@@ -172,6 +175,7 @@ def make(
                 ref_energy_name=self.ref_energy_name,
                 ref_force_name=self.ref_force_name,
                 ref_virial_name=self.ref_virial_name,
+                use_defaults=self.use_defaults,
                 device=device,
                 species_list=species_list,
                 **fit_kwargs,

autoplex/fitting/common/jobs.py

@@ -32,6 +32,7 @@ def machine_learning_fit(
     ref_energy_name: str = "REF_energy",
     ref_force_name: str = "REF_forces",
     ref_virial_name: str = "REF_virial",
+    use_defaults: bool = True,
     device: str = "cuda",
     hyper_param_optimization: bool = False,
     **fit_kwargs,
@@ -64,6 +65,8 @@ def machine_learning_fit(
         Reference force name.
     ref_virial_name : str, optional
         Reference virial name.
+    use_defaults: bool
+        If True, use default fitting parameters
     device: str
         specify device to use cuda or cpu.
     hyper_param_optimization: bool
@@ -145,8 +148,9 @@ def machine_learning_fit(
             ref_energy_name=ref_energy_name,
             ref_force_name=ref_force_name,
             ref_virial_name=ref_virial_name,
-            fit_kwargs=fit_kwargs,
+            use_defaults=use_defaults,
             device=device,
+            fit_kwargs=fit_kwargs,
         )
 
     check_conv = check_convergence(train_test_error["test_error"])

autoplex/fitting/common/mlip-rss-defaults.json

@@ -39,6 +39,7 @@
   },
   "MACE": {
     "model": "MACE",
+    "name": "MACE_model",
     "config_type_weights": "{'Default':1.0}",
     "hidden_irreps": "128x0e + 128x1o",
     "r_max": 5.0,
@@ -48,6 +49,12 @@
     "ema_decay": 0.99,
     "correlation": 3,
     "loss": "huber",
-    "default_dtype": "float32"
+    "default_dtype": "float32",
+    "swa": true,
+    "ema": true,
+    "amsgrad": true,
+    "restart_latest": true,
+    "seed": 123,
+    "device": "cpu"
   }
 }