remove another layer of job flow

autoplex/benchmark/phonons/flows.py

@@ -11,8 +11,8 @@
     from atomate2.common.schemas.phonons import PhononBSDOSDoc
     from pymatgen.core.structure import Structure
 
-from autoplex.benchmark.phonons.jobs import compute_bandstructure_benchmark_metrics
-
+from autoplex.benchmark.phonons.utils import compute_bandstructure_benchmark_metrics
+from jobflow import job
 __all__ = ["PhononBenchmarkMaker"]
 
 
@@ -32,6 +32,7 @@ class PhononBenchmarkMaker(Maker):
 
     name: str = "PhononBenchmark"
 
+    @job
     def make(
         self,
         ml_model: str,
@@ -56,17 +57,14 @@ def make(
         dft_phonon_task_doc: PhononBSDOSDoc
             Phonon task doc from DFT runs consisting of pymatgen band-structure object.
         """
-        jobs = []
 
-        benchmark_job = compute_bandstructure_benchmark_metrics(
+
+        benchmark_output = compute_bandstructure_benchmark_metrics(
             ml_model=ml_model,
             ml_phonon_bs=ml_phonon_task_doc.phonon_bandstructure,
             dft_phonon_bs=dft_phonon_task_doc.phonon_bandstructure,
             dft_imag_modes=dft_phonon_task_doc.has_imaginary_modes,
             ml_imag_modes=ml_phonon_task_doc.has_imaginary_modes,
             structure=structure,
         )
-        jobs.append(benchmark_job)
-
-        # create a flow including all jobs
-        return Flow(jobs=jobs, output=benchmark_job.output, name=self.name)
+        return benchmark_output

autoplex/benchmark/phonons/jobs.py

@@ -12,68 +12,6 @@
 from autoplex.benchmark.phonons.utils import compare_plot, get_rmse, rmse_qdep_plot
 
 
-@job
-def compute_bandstructure_benchmark_metrics(
-    ml_model: str,
-    structure: Structure,
-    ml_phonon_bs: PhononBandStructureSymmLine,
-    dft_phonon_bs: PhononBandStructureSymmLine,
-    ml_imag_modes: bool,
-    dft_imag_modes: bool,
-):
-    """
-    Compute phonon band-structure benchmark metrics and generate associated plots.
-
-    Parameters
-    ----------
-    ml_model: str
-        ML model to be used. Default is GAP.
-    structure : .Structure
-        A structure object.
-    ml_phonon_bs: PhononBandStructureSymmLine.
-       ML generated pymatgen phonon band-structure object.
-    dft_phonon_bs: PhononBandStructureSymmLine.
-       DFT generated pymatgen phonon band-structure object.
-    ml_imag_modes: bool
-        Whether the ML-based phonon band structure shows imaginary modes.
-    dft_imag_modes: bool
-        Whether the DFT-based phonon band structure shows imaginary modes.
-
-    Returns
-    -------
-    Response.output
-       Overall root mean squared error between DFT and ML phonon band-structure.
-    """
-    # compute overall root mean squared error
-    overall_rmse = get_rmse(ml_bs=ml_phonon_bs, dft_bs=dft_phonon_bs)
-
-    # saves rmse k-dependent plot
-    file_name = f"{structure.composition.reduced_formula}_rmse_phonons.pdf"
-    _ = rmse_qdep_plot(
-        ml_bs=ml_phonon_bs,
-        dft_bs=dft_phonon_bs,
-        which_q_path=2,
-        file_name=file_name,
-        img_format="pdf",
-    )
-
-    # saves DFT and ML phonon band-structure overlay plot
-    file_name = f"{structure.composition.reduced_formula}_band_comparison.pdf"
-    _ = compare_plot(
-        ml_model=ml_model,
-        ml_bs=ml_phonon_bs,
-        dft_bs=dft_phonon_bs,
-        file_name=file_name,
-    )
-
-    return Response(
-        output={
-            "benchmark_phonon_rmse": overall_rmse,
-            "dft_imaginary_modes": dft_imag_modes,
-            "ml_imaginary_modes": ml_imag_modes,
-        }
-    )  # TODO TaskDoc
-
 
 @job
 def write_benchmark_metrics(

autoplex/benchmark/phonons/utils.py

@@ -11,6 +11,64 @@
     from matplotlib.figure import Figure
     from pymatgen.phonon.bandstructure import PhononBandStructureSymmLine
 from pymatgen.phonon.plotter import PhononBSPlotter
+from pymatgen.core.structure import Structure
+def compute_bandstructure_benchmark_metrics(
+    ml_model: str,
+    structure: Structure,
+    ml_phonon_bs: PhononBandStructureSymmLine,
+    dft_phonon_bs: PhononBandStructureSymmLine,
+    ml_imag_modes: bool,
+    dft_imag_modes: bool,
+):
+    """
+    Compute phonon band-structure benchmark metrics and generate associated plots.
+
+    Parameters
+    ----------
+    ml_model: str
+        ML model to be used. Default is GAP.
+    structure : .Structure
+        A structure object.
+    ml_phonon_bs: PhononBandStructureSymmLine.
+       ML generated pymatgen phonon band-structure object.
+    dft_phonon_bs: PhononBandStructureSymmLine.
+       DFT generated pymatgen phonon band-structure object.
+    ml_imag_modes: bool
+        Whether the ML-based phonon band structure shows imaginary modes.
+    dft_imag_modes: bool
+        Whether the DFT-based phonon band structure shows imaginary modes.
+
+    Returns
+    -------
+    dict including
+       Overall root mean squared error between DFT and ML phonon band-structure.
+    """
+    # compute overall root mean squared error
+    overall_rmse = get_rmse(ml_bs=ml_phonon_bs, dft_bs=dft_phonon_bs)
+
+    # saves rmse k-dependent plot
+    file_name = f"{structure.composition.reduced_formula}_rmse_phonons.pdf"
+    _ = rmse_qdep_plot(
+        ml_bs=ml_phonon_bs,
+        dft_bs=dft_phonon_bs,
+        which_q_path=2,
+        file_name=file_name,
+        img_format="pdf",
+    )
+
+    # saves DFT and ML phonon band-structure overlay plot
+    file_name = f"{structure.composition.reduced_formula}_band_comparison.pdf"
+    _ = compare_plot(
+        ml_model=ml_model,
+        ml_bs=ml_phonon_bs,
+        dft_bs=dft_phonon_bs,
+        file_name=file_name,
+    )
+    return {
+            "benchmark_phonon_rmse": overall_rmse,
+            "dft_imaginary_modes": dft_imag_modes,
+            "ml_imaginary_modes": ml_imag_modes,
+        }
 
 
 def get_rmse(

tests/auto/test_auto_jobs.py

@@ -2,18 +2,54 @@
 
 import os
 from unittest import mock
+from jobflow import Flow
 from pymatgen.core.structure import Structure
 from autoplex.auto.phonons.jobs import (
     get_iso_atom,
     dft_phonopy_gen_data,
+complete_benchmark
 )
+
 from atomate2.common.schemas.phonons import PhononBSDOSDoc
 from autoplex.data.phonons.flows import TightDFTStaticMaker
 from jobflow import run_locally
 
+from conftest import memory_jobstore
+from get_phonon_doc import response
+from pytest import approx
+
 os.environ["OMP_NUM_THREADS"] = "4"  # export OMP_NUM_THREADS=4
 os.environ["OPENBLAS_NUM_THREADS"] = "1"  # export OPENBLAS_NUM_THREADS=1
 
+def test_complete_benchmark(test_dir, memory_jobstore):
+    from monty.serialization import loadfn
+    from atomate2.common.schemas.phonons import PhononBSDOSDoc
+    from autoplex.fitting.common.flows import MLIPFitMaker
+    database_dir = test_dir / "fitting/rss_training_dataset/"
+    jobs=[]
+    gapfit = MLIPFitMaker().make(
+        auto_delta=False,
+        glue_xml=False,
+        twob={"delta": 2.0, "cutoff": 4},
+        threeb={"n_sparse": 10},
+        preprocessing_data=False,
+        database_dir=database_dir
+    )
+    dft_data = loadfn(test_dir / "benchmark" / "phonon_doc_si.json")
+    dft_doc: PhononBSDOSDoc = dft_data["output"]
+    structure=dft_doc.structure
+
+    jobs.append(gapfit)
+
+
+    bm=complete_benchmark(ibenchmark_structure=0, benchmark_structure=structure, mp_ids=["mp-82"],benchmark_mp_ids=["mp-82"], ml_path=gapfit.output["mlip_path"], ml_model="GAP", dft_references=[dft_doc], add_dft_phonon_struct=False, fit_input=None,symprec=1e-1, phonon_displacement_maker=None, supercell_settings={"min_length": 8})
+    jobs.append(bm)
+
+    response=run_locally(Flow(jobs), store=memory_jobstore)
+    output=response[bm.output.uuid][1].output[0].resolve(store=memory_jobstore)
+    assert output["benchmark_phonon_rmse"]== approx(1.0, abs=0.5)
+    assert output["dft_imaginary_modes"]== False
+    assert output["ml_imaginary_modes"]== False
 
 def test_get_iso_atom(vasp_test_dir, mock_vasp, clean_dir, memory_jobstore):
     structure_list = [

tests/benchmark/test_benchmark_flows.py

@@ -28,7 +28,6 @@ def test_benchmark(test_dir, clean_dir):
         dft_phonon_task_doc =dft_doc,
         benchmark_mp_id="test",
     )
-    assert len(benchmark_flow.jobs) == 1
 
     responses = run_locally(benchmark_flow, create_folders=False, ensure_success=True)
 

tests/benchmark/test_benchmark_jobs.py

@@ -3,11 +3,8 @@
 from pymatgen.core.structure import Structure
 from pymatgen.io.phonopy import get_ph_bs_symm_line
 from pymatgen.phonon.bandstructure import PhononBandStructureSymmLine
-from autoplex.benchmark.phonons.jobs import (
-    compute_bandstructure_benchmark_metrics,
-    write_benchmark_metrics,
-)
-
+from autoplex.benchmark.phonons.jobs import write_benchmark_metrics
+from autoplex.benchmark.phonons.utils import compute_bandstructure_benchmark_metrics
 
 def test_compute_bandstructure_benchmark_metrics_dummy(test_dir, clean_dir):
     import os
@@ -26,13 +23,12 @@ def test_compute_bandstructure_benchmark_metrics_dummy(test_dir, clean_dir):
     df_bs = PhononBandStructureSymmLine.from_dict(dummy_bs_dict)
     ml_bs = PhononBandStructureSymmLine.from_dict(dummy_bs_dict)
 
-    benchmark_job = compute_bandstructure_benchmark_metrics(
+    result=compute_bandstructure_benchmark_metrics(
         ml_model="GAP", structure=df_bs.structure, ml_phonon_bs=ml_bs, dft_phonon_bs=df_bs, ml_imag_modes=False, dft_imag_modes=False,
     )
 
-    responses = run_locally(benchmark_job, create_folders=False, ensure_success=True)
 
-    assert responses[benchmark_job.output.uuid][1].output["benchmark_phonon_rmse"] == pytest.approx(0.0)
+    assert result["benchmark_phonon_rmse"] == pytest.approx(0.0)
 
     # get list of generated plot files
     test_files_dir = Path(test_dir / "benchmark").resolve()
@@ -63,13 +59,12 @@ def test_compute_bandstructure_benchmark_metrics(test_dir, clean_dir):
     df_bs = get_ph_bs_symm_line(bands_path=dft_bs_file_path)
     ml_bs = get_ph_bs_symm_line(bands_path=ml_bs_file_path)
 
-    benchmark_job = compute_bandstructure_benchmark_metrics(
+    result= compute_bandstructure_benchmark_metrics(
         ml_model="GAP",structure=df_bs.structure, ml_phonon_bs=ml_bs, dft_phonon_bs=df_bs, ml_imag_modes=False, dft_imag_modes=False,
     )
 
-    responses = run_locally(benchmark_job, create_folders=False, ensure_success=True)
 
-    assert responses[benchmark_job.output.uuid][1].output["benchmark_phonon_rmse"] == pytest.approx(
+    assert result["benchmark_phonon_rmse"] == pytest.approx(
         0.5716963823412201, abs=0.3  # TODO check results
     )