MPMorph bug fixes

* correct unassigned variable and incorrectly-named energy kwarg in the MPMorph flow

Author: Jonas-Finkler

src/atomate2/common/flows/mpmorph.py

@@ -112,29 +112,24 @@ def make(
                 *self.initial_strain, self.postprocessor.min_data_points
             )
             working_outputs = {
-                "relax": {
-                    key: [] for key in ("energies", "volume", "stress", "pressure")
-                }
+                "relax": {key: [] for key in ("energy", "volume", "stress", "pressure")}
             }
 
         else:
             # Fit EOS to running list of energies and volumes
             self.postprocessor.fit(working_outputs)
             working_outputs = dict(self.postprocessor.results)
             flow_output = {"working_outputs": working_outputs.copy(), "structure": None}
-            for k in ("pressure", "energy"):
+            for k in ("pressure",):
                 working_outputs["relax"].pop(k, None)
 
             # Stop flow here if EOS cannot be fit
             if (v0 := working_outputs.get("V0")) is None:
                 return Response(output=flow_output, stop_children=True)
 
             # Check if equilibrium volume is in range of attempted volumes
-            v0_in_range = (
-                (vmin := working_outputs.get("Vmin"))
-                <= v0
-                <= (vmax := working_outputs.get("Vmax"))
-            )
+            vmin = working_outputs.get("Vmin")
+            vmax = working_outputs.get("Vmax")
 
             # Check if maximum number of refinement NVT runs is set,
             # and if so, if that limit has been reached
@@ -143,7 +138,7 @@ def make(
             )
 
             # Successful fit: return structure at estimated equilibrium volume
-            if v0_in_range or max_attempts_reached:
+            if (vmin <= v0 <= vmax) or max_attempts_reached:
                 flow_output["structure"] = structure.copy()
                 flow_output["structure"].scale_lattice(v0)  # type: ignore[attr-defined]
                 return flow_output
@@ -162,19 +157,19 @@ def make(
         eos_jobs = []
         for index in range(len(deformation_matrices)):
             md_job = self.md_maker.make(
-                structure=deformed_structures[index].final_structure,
+                deformed_structures[index].final_structure,
                 prev_dir=None,
             )
             relaxed_vol = len(working_outputs["relax"]["volume"])
             md_job.name = f"{self.name} {md_job.name} {relaxed_vol + 1}"
 
-            working_outputs["relax"]["energies"].append(md_job.output.output.energy)
+            working_outputs["relax"]["energy"].append(md_job.output.output.energy)
             working_outputs["relax"]["volume"].append(md_job.output.structure.volume)
             working_outputs["relax"]["stress"].append(md_job.output.output.stress)
             eos_jobs.append(md_job)
 
         recursive = self.make(
-            structure=structure,
+            structure,
             prev_dir=None,
             working_outputs=working_outputs,
         )
@@ -262,14 +257,12 @@ def make(
             structure = convergence_flow.output["structure"]
 
         self.production_md_maker.name = self.name + " production run"
-        production_run = self.production_md_maker.make(
-            structure=structure, prev_dir=prev_dir
-        )
+        production_run = self.production_md_maker.make(structure, prev_dir=prev_dir)
         flow_jobs.append(production_run)
 
         if self.quench_maker:
             quench_flow = self.quench_maker.make(
-                structure=production_run.output.structure,
+                production_run.output.structure,
                 prev_dir=production_run.output.dir_name,
             )
             flow_jobs += [quench_flow]
@@ -469,7 +462,7 @@ def make(self, structure: Structure, prev_dir: str | Path | None = None) -> Flow
             )
             if self.descent_method == "stepwise":
                 md_job = self.call_md_maker(
-                    structure=structure,
+                    structure,
                     temp=temp,
                     prev_dir=prev_dir,
                 )
@@ -478,13 +471,13 @@ def make(self, structure: Structure, prev_dir: str | Path | None = None) -> Flow
                 self.descent_method == "linear with hold"
             ):  # TODO: Work in Progress; needs testing
                 md_job_linear = self.call_md_maker(
-                    structure=structure,
+                    structure,
                     temp=[temp, temp - self.quench_temperature_step],  # type: ignore[arg-type]
                     prev_dir=prev_dir,
                 )
 
                 md_job = self.call_md_maker(
-                    structure=md_job_linear.output.structure,
+                    md_job_linear.output.structure,
                     temp=temp - self.quench_temperature_step,
                     prev_dir=md_job_linear.output.dir_name,
                 )

tests/ase/flows/test_mpmorph.py

@@ -0,0 +1,55 @@
+"""Test MPMoprh flow using Lennard-Jones."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+from jobflow import run_locally
+
+from atomate2.ase.jobs import LennardJonesStaticMaker
+from atomate2.common.flows.mpmorph import EquilibriumVolumeMaker
+
+if TYPE_CHECKING:
+    from pymatgen.core import Structure
+
+
+@pytest.mark.parametrize("vol_scale", [1.25, 0.75])
+def test_vol_scale(
+    lj_fcc_ne_pars: dict[str, float], fcc_ne_structure: Structure, vol_scale: float
+) -> None:
+    # Perform single points here instead of MD for
+    # (1) time cost and (2) reliability of output
+    # This is a coarse test to ensure that the volume up and down scaling can occur
+
+    test_structure = fcc_ne_structure.to_conventional()
+    test_structure = test_structure.scale_lattice(vol_scale * test_structure.volume)
+    test_structure = test_structure * (2, 2, 2)
+    init_vol_per_atom = test_structure.volume
+
+    init_strain = 0.05
+    vscale = ((1 - init_strain) ** 3, (1 + init_strain) ** 3)
+    flow = EquilibriumVolumeMaker(
+        md_maker=LennardJonesStaticMaker(
+            calculator_kwargs=lj_fcc_ne_pars,
+        ),
+        min_strain=0.05,
+        initial_strain=0.05,
+    ).make(test_structure)
+
+    resp = run_locally(flow, ensure_success=True)
+    output = None
+    for _output in resp.values():
+        if isinstance(_output[1].output, dict) and (
+            output := _output[1].output.get("working_outputs")
+        ):
+            break
+
+    # ensure that at least one refinement on the volume range was undertaken
+    assert len(output["relax"]["volume"]) > 3
+
+    final_v_ratio = output["V0"] / init_vol_per_atom
+    if vol_scale > 1.0:
+        assert final_v_ratio <= vscale[1]
+    elif vol_scale < 1.0:
+        assert vscale[0] <= final_v_ratio