Merge pull request #30 from QuantumChemist/main

JaGeo · web-flow · commit a4f2802354cd · 2024-02-03T11:12:41.000+01:00
added config_type keyword to gapfit
diff --git a/autoplex/auto/flows.py b/autoplex/auto/flows.py
@@ -36,9 +36,7 @@
 
 
 @dataclass
-class CompleteDFTvsMLBenchmarkWorkflow(
-    Maker
-):  # merge with complete wf and set another flag for adding data
+class CompleteDFTvsMLBenchmarkWorkflow(Maker):
     """
     Maker to add more data to existing dataset (.xyz file).
 
@@ -103,21 +101,18 @@ def make(
         fit_input = {}
         collect = []
 
-        # if xyz_file is None:
-        #    raise Exception("Error. Please provide an existing xyz file.")
-
-        for i, structure in enumerate(structure_list):
+        for structure, mp_id in zip(structure_list, mp_ids):
             if self.add_dft_random_struct:
                 addDFTrand = self.add_dft_random(
                     structure,
-                    mp_ids[i],
+                    mp_id,
                     self.phonon_displacement_maker,
                     self.n_struct,
                     self.uc,
                     self.supercell_matrix,
                 )
                 flows.append(addDFTrand)
-                fit_input.update({mp_ids[i]: addDFTrand.output})
+                fit_input.update({mp_id: addDFTrand.output})
             if self.add_dft_phonon_struct:
                 addDFTphon = self.add_dft_phonons(
                     structure,
@@ -127,17 +122,20 @@ def make(
                     self.min_length,
                 )
                 flows.append(addDFTphon)
-                fit_input.update({mp_ids[i]: addDFTphon.output})
+                fit_input.update({mp_id: addDFTphon.output})
             if self.add_dft_random_struct and self.add_dft_phonon_struct:
-                fit_input.update(
-                    {mp_ids[i]: {**addDFTrand.output, **addDFTphon.output}}
-                )
+                fit_input.update({mp_id: {**addDFTrand.output, **addDFTphon.output}})
             if self.add_rss_struct:
                 raise NotImplementedError
 
         isoatoms = get_iso_atom(structure_list)
         flows.append(isoatoms)
 
+        if xyz_file is None:
+            fit_input.update(
+                {"isolated_atom": {"iso_atoms_dir": [isoatoms.output["dirs"]]}}
+            )
+
         add_data_fit = PhononDFTMLFitFlow().make(
             species=isoatoms.output["species"],
             isolated_atoms_energy=isoatoms.output["energies"],
@@ -153,7 +151,6 @@ def make(
             for ibenchmark_structure, benchmark_structure in enumerate(
                 benchmark_structures
             ):
-                # not sure if it would make sense to put everything from here in its own flow?
                 add_data_ml_phonon = get_phonon_ml_calculation_jobs(
                     structure=benchmark_structure,
                     min_length=self.min_length,
@@ -379,7 +376,9 @@ def make(self, structure: Structure, mp_id):
 @dataclass
 class PhononDFTMLFitFlow(Maker):
     """
-    Maker to fit ML potentials based on DFT data.
+    Maker to fit several types of ML potentials (GAP, ACE etc.) based on DFT data.
+
+    (Currently only the subroutines for GAP are implemented).
 
     Parameters
     ----------
diff --git a/autoplex/auto/jobs.py b/autoplex/auto/jobs.py
@@ -314,5 +314,12 @@ def get_iso_atom(structure_list: list[Structure]):
     isoatoms = IsoAtomMaker().make(all_species=all_species)
     jobs.append(isoatoms)
 
-    flow = Flow(jobs, {"species": all_species, "energies": isoatoms.output})
+    flow = Flow(
+        jobs,
+        {
+            "species": all_species,
+            "energies": isoatoms.output["energies"],
+            "dirs": isoatoms.output["dirs"],
+        },
+    )
     return Response(replace=flow)
diff --git a/autoplex/data/flows.py b/autoplex/data/flows.py
@@ -357,7 +357,8 @@ def make(self, all_species: list[Species]):
             list of pymatgen specie object.
         """
         jobs = []
-        isoatoms = []
+        isoatoms_energy = []
+        isoatoms_dirs = []
         for species in all_species:
             site = Site(species=species, coords=[0, 0, 0])
             mol = Molecule.from_sites([site])
@@ -370,6 +371,8 @@ def make(self, all_species: list[Species]):
             ).make(iso_atom)
 
             jobs.append(isoatom_calcs)
-            isoatoms.append(isoatom_calcs.output.output.energy_per_atom)
+            isoatoms_energy.append(isoatom_calcs.output.output.energy_per_atom)
+            isoatoms_dirs.append(isoatom_calcs.output.dir_name)
+
         # create a flow including all jobs
-        return Flow(jobs, isoatoms)
+        return Flow(jobs, {"energies": isoatoms_energy, "dirs": isoatoms_dirs})
diff --git a/autoplex/fitting/gap-defaults.json b/autoplex/fitting/gap-defaults.json
@@ -1,11 +1,10 @@
 {
   "general": {
     "at_file": "trainGAP.xyz",
-    "e0": "2",
     "default_sigma": "{0.01 0.2 0.2 0.0}",
-    "energy_parameter_name": "energy",
-    "force_parameter_name": "forces",
-    "virial_parameter_name": "virial",
+    "energy_parameter_name": "REF_energy",
+    "force_parameter_name": "REF_forces",
+    "virial_parameter_name": "REF_virial",
     "sparse_jitter": 1.0e-8,
     "do_copy_at_file": "F",
     "openmp_chunk_size": 10000,
diff --git a/autoplex/fitting/jobs.py b/autoplex/fitting/jobs.py
@@ -52,6 +52,8 @@ def gapfit(
         bool indicating whether to include soap hyperparameters
     xyz_file: str or None
         a possibly already existing xyz file
+    config_types: list[str] or None
+            list of config_types.
     fit_kwargs : dict.
         dict including gap fit keyword args.
 
@@ -60,13 +62,24 @@ def gapfit(
     Response.output
         Path to the gap fit file.
     """
+    config_types = []
     if fit_kwargs is None:
         fit_kwargs = field(default_factory=dict)
 
     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, xyz_file=xyz_file
+        path_to_vasp_static_calcs=list_of_vasp_calc_dirs,
+        config_types=config_types,
+        xyz_file=xyz_file,
     )
 
     gap_default_hyperparameters = load_gap_hyperparameter_defaults(
diff --git a/autoplex/fitting/utils.py b/autoplex/fitting/utils.py
@@ -7,7 +7,7 @@
 import shutil
 from pathlib import Path
 
-import numpy as np
+from ase.constraints import voigt_6_to_full_3x3_stress
 from ase.io import read, write
 from atomate2.utils.path import strip_hostname
 
@@ -62,16 +62,16 @@ def gap_hyperparameter_constructor(
            gap fit input parameter string.
     """
     # convert gap_parameter_dict to representation compatible with gap
-    if atoms_energies and atoms_symbols is not None:
-        e0 = ":".join(
-            [
-                f"{iso_atom}:{iso_energy}"
-                for iso_atom, iso_energy in zip(atoms_symbols, atoms_energies)
-            ]
-        )
+    # if atoms_energies and atoms_symbols is not None:
+    #     e0 = ":".join(
+    #         [
+    #             f"{iso_atom}:{iso_energy}"
+    #             for iso_atom, iso_energy in zip(atoms_symbols, atoms_energies)
+    #         ]
+    #     )
 
-        # Update the isolated atom energy argument
-        gap_parameter_dict["general"].update({"e0": e0})
+    # Update the isolated atom energy argument
+    # gap_parameter_dict["general"].update({"e0": e0})
 
     general = [f"{key}={value}" for key, value in gap_parameter_dict["general"].items()]
 
@@ -136,7 +136,11 @@ def get_list_of_vasp_calc_dirs(flow_output):
     return list_of_vasp_calc_dirs
 
 
-def outcar_2_extended_xyz(path_to_vasp_static_calcs: list, xyz_file: str | None = None):
+def outcar_2_extended_xyz(
+    path_to_vasp_static_calcs: list,
+    config_types: list[str] | None = None,
+    xyz_file: str | None = None,
+):
     """
     Parse all VASP OUTCARs and generates a trainGAP.xyz.
 
@@ -149,16 +153,27 @@ def outcar_2_extended_xyz(path_to_vasp_static_calcs: list, xyz_file: str | None
         List of VASP static calculation directories.
     xyz_file: str or None
         a possibly already existing xyz file.
+    config_types: list[str] or None
+            list of config_types.
     """
-    for path in path_to_vasp_static_calcs:
+    if config_types is None:
+        config_types = ["bulk"] * len(path_to_vasp_static_calcs)
+
+    for path, config_type in zip(path_to_vasp_static_calcs, config_types):
         # strip hostname if it exists in the path
-        path_without_hostname = Path(strip_hostname(path)).joinpath("OUTCAR.gz")
+        path_without_hostname = Path(strip_hostname(path)).joinpath("vasprun.xml.gz")
         # read the outcar
         file = read(path_without_hostname, index=":")
         for i in file:
-            xx, yy, zz, yz, xz, xy = -i.calc.results["stress"] * i.get_volume()
-            i.info["virial"] = np.array([(xx, xy, xz), (xy, yy, yz), (xz, yz, zz)])
+            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"]
+            i.arrays["REF_forces"] = i.calc.results["forces"]
+            del i.calc.results["forces"]
+            i.info["REF_energy"] = i.calc.results["free_energy"]
+            del i.calc.results["energy"]
+            del i.calc.results["free_energy"]
+            i.info["config_type"] = config_type
             i.pbc = True
         if xyz_file is not None:
             shutil.copy2(xyz_file, os.getcwd())
diff --git a/tests/fitting/test_fitting_flows.py b/tests/fitting/test_fitting_flows.py
@@ -65,14 +65,33 @@ def test_mlip_fit_maker(test_dir, clean_dir, memory_jobstore, vasp_test_dir):
                 ]
             ],
             "phonon_data": [],
+        },
+        "isolated_atom": {"iso_atoms_dir": [[
+            (
+                    vasp_test_dir
+                    / "Li_iso_atoms"
+                    / "Li-statisoatom"
+                    / "outputs"
+            )
+            .absolute()
+            .as_posix(),
+            (
+                    vasp_test_dir
+                    / "Cl_iso_atoms"
+                    / "Cl-statisoatom"
+                    / "outputs"
+            )
+            .absolute()
+            .as_posix(),
+        ]]
         }
     }
 
     # Test to check if gap fit runs with default hyperparameter sets (i.e. include_two_body and include_soap is True)
     gapfit = MLIPFitMaker().make(
         species_list=["Li", "Cl"],
         iso_atom_energy=[-0.28649227, -0.25638457],
-        fit_input=fit_input_dict,
+        fit_input=fit_input_dict
     )
 
     responses = run_locally(
@@ -107,26 +126,26 @@ def test_mlip_fit_maker_with_kwargs(
             "rand_struc_dir": [
                 [
                     (
-                        vasp_test_dir
-                        / "dft_ml_data_generation"
-                        / "rand_static_1"
-                        / "outputs"
+                            vasp_test_dir
+                            / "dft_ml_data_generation"
+                            / "rand_static_1"
+                            / "outputs"
                     )
                     .absolute()
                     .as_posix(),
                     (
-                        vasp_test_dir
-                        / "dft_ml_data_generation"
-                        / "rand_static_2"
-                        / "outputs"
+                            vasp_test_dir
+                            / "dft_ml_data_generation"
+                            / "rand_static_2"
+                            / "outputs"
                     )
                     .absolute()
                     .as_posix(),
                     (
-                        vasp_test_dir
-                        / "dft_ml_data_generation"
-                        / "rand_static_3"
-                        / "outputs"
+                            vasp_test_dir
+                            / "dft_ml_data_generation"
+                            / "rand_static_3"
+                            / "outputs"
                     )
                     .absolute()
                     .as_posix(),
@@ -135,24 +154,43 @@ def test_mlip_fit_maker_with_kwargs(
             "phonon_dir": [
                 [
                     (
-                        vasp_test_dir
-                        / "dft_ml_data_generation"
-                        / "phonon_static_1"
-                        / "outputs"
+                            vasp_test_dir
+                            / "dft_ml_data_generation"
+                            / "phonon_static_1"
+                            / "outputs"
                     )
                     .absolute()
                     .as_posix(),
                     (
-                        vasp_test_dir
-                        / "dft_ml_data_generation"
-                        / "phonon_static_2"
-                        / "outputs"
+                            vasp_test_dir
+                            / "dft_ml_data_generation"
+                            / "phonon_static_2"
+                            / "outputs"
                     )
                     .absolute()
                     .as_posix(),
                 ]
             ],
             "phonon_data": [],
+        },
+        "isolated_atom": {"iso_atoms_dir": [[
+            (
+                    vasp_test_dir
+                    / "Li_iso_atoms"
+                    / "Li-statisoatom"
+                    / "outputs"
+            )
+            .absolute()
+            .as_posix(),
+            (
+                    vasp_test_dir
+                    / "Cl_iso_atoms"
+                    / "Cl-statisoatom"
+                    / "outputs"
+            )
+            .absolute()
+            .as_posix(),
+        ]]
         }
     }
 
diff --git a/tests/fitting/test_fitting_utils.py b/tests/fitting/test_fitting_utils.py
diff --git a/tests/test_data/fitting/ref_files/trainGAP.xyz b/tests/test_data/fitting/ref_files/trainGAP.xyz
