Merge pull request #35 from QuantumChemist/merge_mlip_fitting_from_yb

[WIP] Merge mlip_fitting code from Yuanbin into autoplex

Author: JaGeo

autoplex/auto/phonons/flows.py

@@ -154,7 +154,7 @@ def make(
                 add_data_ml_phonon = get_phonon_ml_calculation_jobs(
                     structure=benchmark_structure,
                     min_length=self.min_length,
-                    ml_dir=add_data_fit.output,
+                    ml_dir=add_data_fit.output["mlip_xml"],
                 )
                 flows.append(add_data_ml_phonon)
 
@@ -407,7 +407,7 @@ def make(
         isolated_atoms_energy: list.
             Isolated atoms energy list
         fit_input: list.
-            Mixed list of dictionary and lists
+            Mixed list of dictionary and lists of the fit input data.
         xyz_file: str or None
             a possibly already existing xyz file
         fit_kwargs : dict.

autoplex/benchmark/phonons/utils.py

@@ -114,6 +114,7 @@ def compare_plot(
     plotter = PhononBSPlotter(bs=ml_bs)
     plotter2 = PhononBSPlotter(bs=dft_bs)
     new_plotter = plotter.plot_compare(plotter2)
-    new_plotter.figure.savefig(file_name)
+
+    new_plotter.figure.savefig(fname=file_name)
 
     return new_plotter.figure

autoplex/fitting/flows.py

@@ -9,7 +9,8 @@
     from pathlib import Path
 from jobflow import Flow, Maker
 
-from autoplex.fitting.jobs import gapfit
+from autoplex.mlip_fitting.fitting.mlip_fitting import MLIPFitMaker as YbMLIPFitMaker
+from autoplex.mlip_fitting.fitting.mlip_fitting import data_preprocessing
 
 __all__ = ["MLIPFitMaker"]
 
@@ -57,12 +58,23 @@ def make(
             dict including gap fit keyword args.
         """
         jobs = []
-        gap_fit_job = gapfit(
-            fit_input=fit_input,
-            isolated_atoms=species_list,
-            isolated_atoms_energy=iso_atom_energy,
-            xyz_file=xyz_file,
-            fit_kwargs=fit_kwargs,
+        # gap_fit_job = gapfit(
+        #     fit_input=fit_input,
+        #     isolated_atoms=species_list,
+        #     isolated_atoms_energy=iso_atom_energy,
+        #     xyz_file=xyz_file,
+        #     fit_kwargs=fit_kwargs,
+        # )
+        data_prep_job = data_preprocessing(
+            split_ratio=0.4, regularization=True, distillation=True, f_max=40
+        ).make(
+            fit_input=fit_input, pre_database_dir=None, xyz_file=xyz_file
+        )  # CE: I think pre_database_dir and xyz_file
+        # are meant to be the same
+        jobs.append(data_prep_job)
+        gap_fit_job = YbMLIPFitMaker(mlip_type="GAP").make(
+            database_dir=data_prep_job.output,
+            isol_es=None,  # nequip={},
         )
         jobs.append(gap_fit_job)  # type: ignore
 

autoplex/fitting/utils.py

@@ -177,4 +177,4 @@ def outcar_2_extended_xyz(
             i.pbc = True
         if xyz_file is not None:
             shutil.copy2(xyz_file, os.getcwd())
-        write("trainGAP.xyz", file, append=True)
+        write("vasp_ref.extxyz", file, append=True)

autoplex/mlip_fitting/README.md

@@ -0,0 +1,12 @@
+For using this code for fitting MLIPs, one could use the following two jobs:
+
+
+    from autoplex.fitting.mlip_fitting import data_preprocessing, mlip_fit
+    job1 = data_preprocessing(split_ratio=0.1, regularization=True, distillation=True, f_max=40).make(fit_input=fit_input, pre_database_dir=None)
+    job2 = mlip_fit(mlip_type='GAP').make(database_dir=job1.output, isol_es=None)
+
+The "split_ratio" parameter is used to divide the training set and the test set. A value of 0.1 means that the ratio of the training set to the test set is 9:1.
+
+Here 'isol_es' is for the ACE fitting, because the ACE model prefers to use cohesive energy for fitting instead of free energy. The format of 'isol_es' for Si is like isol_es={14: -0.81432699}.
+
+For GAP fitting, the isol_es will be automatically accepted in the "train.extxyz" as a configuration with the "config_type=isolated_atom".

autoplex/mlip_fitting/fitting/__init__.py

@@ -0,0 +1 @@
+"""Some description goes here."""

autoplex/mlip_fitting/fitting/mlip_fitting.py

@@ -0,0 +1,222 @@
+"""The mlip_fitting part."""
+
+from __future__ import annotations
+
+import os
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+
+import ase.io
+from jobflow import Maker, job
+
+from autoplex.fitting.utils import get_list_of_vasp_calc_dirs, outcar_2_extended_xyz
+from autoplex.mlip_fitting.fitting.mlip_models import gap_fitting  # , ace_fitting
+from autoplex.mlip_fitting.fitting.regularization import set_sigma
+from autoplex.mlip_fitting.fitting.utilities import data_distillation, split_dataset
+
+
+@dataclass
+class data_preprocessing(Maker):
+    """
+    Data preprocessing function.
+
+    Parameters
+    ----------
+    name : str
+        Name of the flows produced by this maker.
+    split_ratio: float
+        Parameter to divide the training set and the test set.
+        A value of 0.1 means that the ratio of the training set to the test set is 9:1
+    regularization: bool
+        For using regularization.
+    distillation: bool
+        For using distillation.
+    f_max: float
+        Maximally allowed force in the data set.
+
+    """
+
+    name: str = "data_preprocessing_for_fitting"
+    split_ratio: float = 0.5
+    regularization: bool = False
+    distillation: bool = False
+    f_max: float = 40.0
+
+    @job
+    def make(
+        self,
+        fit_input: dict,
+        pre_database_dir: str,
+        xyz_file: str | None = None,
+    ):
+        """
+        Maker for data preprocessing.
+
+        Parameters
+        ----------
+        fit_input:
+            Mixed list of dictionary and lists of the fit input data.
+        pre_database_dir:
+            the pre-database directory.
+        xyz_file:
+            the already existing training datasets labeled by VASP.
+
+        """
+        config_types = []
+
+        list_of_vasp_calc_dirs = get_list_of_vasp_calc_dirs(flow_output=fit_input)
+
+        config_types = [
+            key
+            for key, value in fit_input.items()
+            for key2, value2 in value.items()
+            if key2 != "phonon_data"
+            for _ in value2[0]
+        ]
+
+        outcar_2_extended_xyz(
+            path_to_vasp_static_calcs=list_of_vasp_calc_dirs,
+            config_types=config_types,
+            xyz_file=xyz_file,
+        )
+
+        # reject structures with large force components
+        atoms = (
+            data_distillation("vasp_ref.extxyz", self.f_max)
+            if self.distillation
+            else ase.io.read("vasp_ref.extxyz", index=":")
+        )
+
+        # split dataset into training and testing datasets with a ratio of 9:1
+        (train_structures, test_structures) = split_dataset(atoms, self.split_ratio)
+
+        # Merging database
+        if pre_database_dir and os.path.exists(pre_database_dir):
+            files_to_copy = ["train.extxyz", "test.extxyz"]
+            current_working_directory = os.getcwd()
+
+            for file_name in files_to_copy:
+                source_file_path = os.path.join(pre_database_dir, file_name)
+                destination_file_path = os.path.join(
+                    current_working_directory, file_name
+                )
+                shutil.copy(source_file_path, destination_file_path)
+                print(f"File {file_name} has been copied to {destination_file_path}")
+
+        ase.io.write("train.extxyz", train_structures, format="extxyz", append=True)
+        ase.io.write("test.extxyz", test_structures, format="extxyz", append=True)
+
+        if self.regularization:
+            atoms = ase.io.read("train.extxyz", index=":")
+            atom_with_sigma = set_sigma(
+                atoms,
+                etup=[(0.1, 1), (0.001, 0.1), (0.0316, 0.316), (0.0632, 0.632)],
+            )
+            ase.io.write("train_with_sigma.extxyz", atom_with_sigma, format="extxyz")
+
+        # database_path = Path.cwd()
+
+        return Path.cwd()  # database_path
+
+
+@dataclass
+class MLIPFitMaker(Maker):
+    """
+    Maker to fitting potential.
+
+    Parameters
+    ----------
+    name : str
+        Name of the flows produced by this maker.
+    mlip_type: str
+        Choose one specific MLIP type:
+        'GAP' | 'SNAP' | 'ACE' | 'Nequip' | 'Allegro' | 'MACE'
+    HPO: bool
+        call hyperparameter optimization (HPO) or not
+
+    """
+
+    name: str = "MLIP_FIT"
+    mlip_type: str | None = None
+    HPO: bool = False
+
+    @job
+    def make(
+        self,
+        database_dir: str,
+        # nequip: dict,
+        gap_para=None,
+        # ace_para={
+        #     "energy_name": "REF_energy",
+        #     "force_name": "REF_forces",
+        #     "virial_name": "REF_virials",
+        #     "order": 3,
+        #     "totaldegree": 6,
+        #     "cutoff": 2.0,
+        #     "solver": "BLR",
+        # },
+        isol_es: None = None,
+        num_of_threads: int = 128,
+        **kwargs,
+    ):
+        """
+        Maker for data preprocessing.
+
+        Parameters
+        ----------
+        database_dir:
+            the database directory.
+        nequip:
+            nequip parameters.
+        gap_para: dict
+            gap fit parameters.
+        isol_es:
+            isolated es.
+        num_of_threads: int
+            number of threads to be used.
+
+        """
+        if gap_para is None:
+            gap_para = {"two_body": True, "three_body": False, "soap": True}
+
+        mlip_path = Path.cwd()
+        if os.path.join(database_dir, "train_with_sigma.extxyz"):
+            shutil.copy(
+                os.path.join(database_dir, "train_with_sigma.extxyz"),
+                os.path.join(mlip_path, "train_with_sigma.extxyz"),
+            )
+        shutil.copy(
+            os.path.join(database_dir, "test.extxyz"),
+            os.path.join(mlip_path, "test.extxyz"),
+        )
+        shutil.copy(
+            os.path.join(database_dir, "train.extxyz"),
+            os.path.join(mlip_path, "train.extxyz"),
+        )
+
+        if self.mlip_type is None:
+            raise ValueError(
+                "MLIP type is not defined! "
+                "The current version supports the fitting of GAP, SNAP, ACE, Nequip, Allegro, or MACE."
+            )
+
+        if self.mlip_type == "GAP":
+            train_error, test_error = gap_fitting(
+                dir=database_dir,
+                two_body=gap_para["two_body"],
+                three_body=gap_para["three_body"],
+                soap=["soap"],
+            )
+
+        convergence = False
+        if test_error < 0.01:
+            convergence = True
+
+        return {
+            "mlip_path": mlip_path,
+            "mlip_xml": mlip_path.joinpath("gap_file.xml"),
+            "train_error": train_error,
+            "test_error": test_error,
+            "convergence": convergence,
+        }