fixing unit test

autoplex/data/phonons/utils.py

@@ -38,25 +38,29 @@ def ml_phonon_maker_preparation(
         Keyword arguments that can be passed to the StaticMaker.
     """
     if bulk_relax_maker is not None:
-        br = bulk_relax_maker
-        bulk_relax_maker = br.update_kwargs(
+        bulk_relax_maker = bulk_relax_maker.update_kwargs(
             update={"calculator_kwargs": calculator_kwargs}
         )
         if relax_maker_kwargs is not None:
-            bulk_relax_maker = br.update_kwargs(update={**relax_maker_kwargs})
+            bulk_relax_maker = bulk_relax_maker.update_kwargs(
+                update={**relax_maker_kwargs}
+            )
+
     if phonon_displacement_maker is not None:
-        ph_disp = phonon_displacement_maker
-        phonon_displacement_maker = ph_disp.update_kwargs(
+        phonon_displacement_maker = phonon_displacement_maker.update_kwargs(
             update={"calculator_kwargs": calculator_kwargs}
         )
         if static_maker_kwargs is not None:
-            phonon_displacement_maker = ph_disp.update_kwargs({**static_maker_kwargs})
+            phonon_displacement_maker = phonon_displacement_maker.update_kwargs(
+                {**static_maker_kwargs}
+            )
     if static_energy_maker is not None:
-        stat_en = static_energy_maker
-        static_energy_maker = stat_en.update_kwargs(
+        static_energy_maker = static_energy_maker.update_kwargs(
             update={"calculator_kwargs": calculator_kwargs}
         )
         if static_maker_kwargs is not None:
-            static_energy_maker = stat_en.update_kwargs(update={**static_maker_kwargs})
+            static_energy_maker = static_energy_maker.update_kwargs(
+                update={**static_maker_kwargs}
+            )
 
     return bulk_relax_maker, phonon_displacement_maker, static_energy_maker

tests/auto/test_auto_flows.py

@@ -169,6 +169,8 @@ def test_complete_dft_vs_ml_benchmark_workflow_gap(
         store=memory_jobstore,
     )
 
+    print("RMSE: ", responses[complete_workflow.jobs[-1].output.uuid][1].output[0][0]["benchmark_phonon_rmse"])
+
     assert complete_workflow.jobs[4].name == "complete_benchmark"
     assert responses[complete_workflow.jobs[-1].output.uuid][1].output[0][0]["benchmark_phonon_rmse"] == pytest.approx(
         2.002641337594289, abs=1.0  # it's kinda fluctuating because of the little data

tests/data/test_datagen_flows.py

@@ -5,6 +5,11 @@
 from atomate2.vasp.powerups import update_user_incar_settings
 from atomate2.common.schemas.phonons import PhononBSDOSDoc
 from pymatgen.core.structure import Structure
+from atomate2.forcefields.jobs import (
+    GAPRelaxMaker,
+    GAPStaticMaker,
+
+)
 
 from autoplex.data.common.flows import GenerateTrainingDataForTesting
 from autoplex.data.phonons.flows import IsoAtomMaker, RandomStructuresDataGenerator, MLPhononMaker
@@ -16,12 +21,17 @@
 def test_ml_phonon_maker(test_dir, clean_dir, memory_jobstore):
     from jobflow import run_locally
 
-    potential_file_dir = test_dir / "fitting" / "ref_files"
+    potential_file = test_dir / "fitting" / "ref_files" / "gap_file.xml"
     path_to_struct = test_dir / "fitting" / "ref_files" / "POSCAR"
     structure = Structure.from_file(path_to_struct)
 
-    gap_phonon_jobs = MLPhononMaker(min_length=20).make_from_ml_model(
-        structure=structure, potential_file=str(potential_file_dir),
+    gap_phonon_jobs = MLPhononMaker(
+        min_length=20,
+        bulk_relax_maker=GAPRelaxMaker(relax_cell=True, relax_kwargs={"interval": 500}),
+        phonon_displacement_maker=GAPStaticMaker(name="gap phonon static"),
+        static_energy_maker=GAPStaticMaker(),
+    ).make_from_ml_model(
+        structure=structure, potential_file=potential_file,
     )
 
     responses = run_locally(