Generalize data creation for phonon workflow and beyond

pyproject.toml

@@ -25,7 +25,7 @@ classifiers = [
 requires-python = ">=3.10,<3.12"
 dependencies = [
      "pymatgen>=2024.11.13",
-     "atomate2[strict]>=0.0.20",
+     "atomate2[strict]>=0.0.21",
      "ase==3.24.0",
      "calorine>=3.0",
      "matgl>=1.2.6",
@@ -69,7 +69,7 @@ workflow-managers = [
 strict = [
      "calorine==3.0",
      "pymatgen==2025.2.18",
-     "atomate2[strict]==0.0.20",
+     "atomate2[strict]==0.0.21",
      "matgl==1.2.6",
      "quippy-ase==0.9.14; python_version < '3.12'",
      "ase==3.24.0",

src/autoplex/auto/phonons/flows.py

@@ -4,6 +4,7 @@
 import warnings
 from dataclasses import dataclass, field
 
+from atomate2.ase.jobs import AseMaker
 from atomate2.common.schemas.phonons import PhononBSDOSDoc
 from atomate2.vasp.flows.core import DoubleRelaxMaker
 from atomate2.vasp.flows.mp import (
@@ -71,15 +72,15 @@ class CompleteDFTvsMLBenchmarkWorkflow(Maker):
         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
+    displacement_maker: BaseVaspMaker|AseMaker
         Maker used for a static calculation for a supercell.
-    phonon_bulk_relax_maker: BaseVaspMaker
+    phonon_bulk_relax_maker: BaseVaspMaker|AseMaker
         Maker used for the bulk relax unit cell calculation.
-    rattled_bulk_relax_maker: BaseVaspMaker
+    rattled_bulk_relax_maker: BaseVaspMaker|AseMaker
         Maker used for the bulk relax unit cell calculation.
-    phonon_static_energy_maker: BaseVaspMaker
+    phonon_static_energy_maker: BaseVaspMaker|AseMaker
         Maker used for the static energy unit cell calculation.
-    isolated_atom_maker: IsoAtomStaticMaker
+    isolated_atom_maker: IsoAtomStaticMaker|AseMaker
         VASP maker for the isolated atom calculation.
     n_structures : int.
         Total number of distorted structures to be generated.
@@ -180,8 +181,8 @@ class CompleteDFTvsMLBenchmarkWorkflow(Maker):
     add_dft_phonon_struct: bool = True
     add_dft_rattled_struct: bool = True
     add_rss_struct: bool = False
-    displacement_maker: BaseVaspMaker = None
-    phonon_bulk_relax_maker: BaseVaspMaker | None = field(
+    displacement_maker: BaseVaspMaker | AseMaker = None
+    phonon_bulk_relax_maker: BaseVaspMaker | AseMaker | None = field(
         default_factory=lambda: DoubleRelaxMaker.from_relax_maker(
             TightRelaxMaker(
                 name="dft tight relax",
@@ -210,9 +211,9 @@ class CompleteDFTvsMLBenchmarkWorkflow(Maker):
             )
         )
     )
-    phonon_static_energy_maker: BaseVaspMaker | None = None
+    phonon_static_energy_maker: BaseVaspMaker | AseMaker | None = None
 
-    rattled_bulk_relax_maker: BaseVaspMaker | None = field(
+    rattled_bulk_relax_maker: BaseVaspMaker | AseMaker | None = field(
         default_factory=lambda: TightRelaxMaker(
             run_vasp_kwargs={"handlers": {}},
             input_set_generator=TightRelaxSetGenerator(
@@ -236,7 +237,7 @@ class CompleteDFTvsMLBenchmarkWorkflow(Maker):
             ),
         )
     )
-    isolated_atom_maker: IsoAtomStaticMaker | None = None
+    isolated_atom_maker: IsoAtomStaticMaker | AseMaker | None = None
     n_structures: int = 10
     displacements: list[float] = field(default_factory=lambda: [0.01])
     symprec: float = 1e-4
@@ -427,6 +428,7 @@ def make(
         isoatoms = get_iso_atom(structure_list, self.isolated_atom_maker)
         flows.append(isoatoms)
 
+        # isolated atoms energy needs to be added
         if pre_xyz_files is None:
             fit_input.update(
                 {"IsolatedAtom": {"iso_atoms_dir": [isoatoms.output["dirs"]]}}

src/autoplex/auto/rss/jobs.py

@@ -190,6 +190,9 @@ def initial_rss(
         bcur_params=bcur_params,
         random_seed=random_seed,
     ).make()
+
+    # TODO: this needs to be generalized beyond VASP and instead be able to use a different dft calculator,
+    # or a force field
     do_dft_static = DFTStaticLabelling(
         e0_spin=e0_spin,
         isolatedatom_box=isolatedatom_box,

src/autoplex/fitting/common/flows.py

@@ -314,7 +314,7 @@ def make(
             self.pre_xyz_files = ["train.extxyz", "test.extxyz"]
 
         list_of_vasp_calc_dirs = get_list_of_vasp_calc_dirs(flow_output=fit_input)
-
+        print(fit_input)
         config_types = [
             key
             for key, value in fit_input.items()

src/autoplex/fitting/common/utils.py

@@ -1517,9 +1517,15 @@ def vaspoutput_2_extended_xyz(
     ):
         # strip hostname if it exists in the path
         path_without_hostname = Path(strip_hostname(path)).joinpath("vasprun.xml.gz")
+        path_without_hostname2 = Path(strip_hostname(path)).joinpath(
+            "final_atoms_object.xyz"
+        )
         try:
             # read the vasp output
-            file = read(path_without_hostname, index=":")
+            if path_without_hostname.exists():
+                file = read(path_without_hostname, index=":")
+            elif path_without_hostname2.exists():
+                file = read(path_without_hostname2, index=":")
             for i in file:
                 virial_list = (
                     -voigt_6_to_full_3x3_stress(i.get_stress()) * i.get_volume()

tests/auto/phonons/test_flows.py

@@ -800,6 +800,58 @@ def test_complete_dft_vs_ml_benchmark_workflow_gap(
             assert expected_soap_dict in results_file, f"Expected soap_dict not found in {file_path}"
 
 
+def test_complete_dft_vs_ml_benchmark_workflow_gap_ml_potential_for_data(
+        vasp_test_dir, test_dir, memory_jobstore,  clean_dir
+):
+    import glob
+
+    from atomate2.forcefields.jobs import ForceFieldRelaxMaker, ForceFieldStaticMaker
+
+    path_to_struct = vasp_test_dir / "dft_ml_data_generation" / "POSCAR"
+    structure = Structure.from_file(path_to_struct)
+
+    complete_workflow = CompleteDFTvsMLBenchmarkWorkflow(
+        symprec=1e-2, displacements=[0.01],
+        volume_custom_scale_factors=[0.975, 1.0, 1.025, 1.05],
+        supercell_settings={"min_length": 8, "min_atoms": 20},
+        apply_data_preprocessing=True,
+        displacement_maker=ForceFieldStaticMaker(force_field_name="MACE_MP_0B3"),
+        phonon_bulk_relax_maker=ForceFieldRelaxMaker(force_field_name="MACE_MP_0B3", steps=5000),
+        rattled_bulk_relax_maker=ForceFieldRelaxMaker(force_field_name="MACE_MP_0B3", steps=5000),
+        phonon_static_energy_maker=None,
+        isolated_atom_maker=ForceFieldStaticMaker(force_field_name="MACE_MP_0B3"),
+    ).make(
+        structure_list=[structure],
+        mp_ids=["test"],
+        benchmark_mp_ids=["mp-22905"],
+        benchmark_structures=[structure],
+    )
+    # run the flow or job and ensure that it finished running successfully
+    responses = run_locally(
+        complete_workflow,
+        create_folders=True,
+        ensure_success=True,
+        store=memory_jobstore,
+    )
+
+    assert complete_workflow.jobs[5].name == "complete_benchmark_mp-22905"
+    assert responses[complete_workflow.jobs[-1].output.uuid][1].output["metrics"][0][0][
+               "benchmark_phonon_rmse"] == pytest.approx(
+        0.5214141274873627, abs=1.5  # it's kinda fluctuating because of the little data
+    )
+
+    # check if soap_default_dict is correctly constructed from
+    # n_sparse and delta values in mlip_phonon_default json file
+    expected_soap_dict = "atom-wise f=0.1: n_sparse = 6000, SOAP delta = 1.0"
+    results_files = glob.glob('job*/results_LiCl.txt')
+
+    for file_path in results_files:
+        with open(file_path, 'r') as file:
+            results_file = file.read().strip()
+            assert expected_soap_dict in results_file, f"Expected soap_dict not found in {file_path}"
+
+
+
 def test_complete_dft_vs_gap_benchmark_workflow_database(
         vasp_test_dir, mock_vasp, test_dir, memory_jobstore, ref_paths4_mpid, fake_run_vasp_kwargs4_mpid, clean_dir
 ):