Allow MatPesStaticFlowMaker(static2=None) (#997)

* allow setting MatPesStaticFlowMaker(static2=None)

* expand test_matpes_static_flow_maker to check maker=None cases

Author: janosh

src/atomate2/vasp/flows/matpes.py

@@ -50,7 +50,7 @@ class MatPesStaticFlowMaker(Maker):
             ),
         )
     )
-    static2: Maker = field(
+    static2: Maker | None = field(
         default_factory=lambda: MatPesMetaGGAStaticMaker(
             # start from pre-conditioned WAVECAR from static1 to speed up convergence
             # could copy CHGCAR too but is redundant since VASP can reconstruct it from
@@ -69,6 +69,11 @@ class MatPesStaticFlowMaker(Maker):
         )
     )
 
+    def __post_init__(self) -> None:
+        """Validate flow."""
+        if (self.static1, self.static2, self.static3) == (None, None, None):
+            raise ValueError("Must provide at least one StaticMaker")
+
     def make(self, structure: Structure, prev_dir: str | Path | None = None) -> Flow:
         """Create a flow with MatPES statics.
 
@@ -89,10 +94,20 @@ def make(self, structure: Structure, prev_dir: str | Path | None = None) -> Flow
         Flow
             A flow containing 2 or 3 statics.
         """
-        static1 = self.static1.make(structure, prev_dir=prev_dir)
-        static2 = self.static2.make(structure, prev_dir=static1.output.dir_name)
-        output = {"static1": static1.output, "static2": static2.output}
-        jobs = [static1, static2]
+        jobs = []
+        output = {}
+
+        if self.static1 is not None:
+            static1 = self.static1.make(structure, prev_dir=prev_dir)
+            jobs += [static1]
+            output["static1"] = static1.output
+
+        prev_dir = static1.output.dir_name if self.static1 is not None else prev_dir
+
+        if self.static2 is not None:
+            static2 = self.static2.make(structure, prev_dir=prev_dir)
+            jobs += [static2]
+            output["static2"] = static2.output
 
         # only run 3rd static if set generator not None and structure contains at least
         # one element with Hubbard +U corrections
@@ -104,7 +119,7 @@ def make(self, structure: Structure, prev_dir: str | Path | None = None) -> Flow
                 anion in elems and elems & {*cations}
                 for anion, cations in u_corrections.items()
             ):
-                static3 = self.static3.make(structure, prev_dir=static1.output.dir_name)
+                static3 = self.static3.make(structure, prev_dir=prev_dir)
                 output["static3"] = static3.output
                 jobs += [static3]
 

tests/vasp/flows/test_matpes.py

@@ -46,3 +46,29 @@ def test_matpes_static_flow_maker(mock_vasp, clean_dir, vasp_test_dir):
 
     assert isinstance(flow.output, dict)
     assert {*flow.output} == {"static1", "static2"}
+
+    # Test setting single maker to None
+    flow_maker = MatPesStaticFlowMaker(static1=None)
+    flow = flow_maker.make(si_struct)
+    assert len(flow) == 1
+    assert flow[0].name == "MatPES meta-GGA static"
+
+    flow_maker = MatPesStaticFlowMaker(static2=None)
+    flow = flow_maker.make(si_struct)
+    assert len(flow) == 1
+    assert flow[0].name == "MatPES GGA static"
+
+    # Test setting two makers to None
+    flow_maker = MatPesStaticFlowMaker(static1=None, static2=None)
+    flow = flow_maker.make(si_struct)
+    assert len(flow) == 0
+
+    # Test setting all three makers to None
+    with pytest.raises(ValueError, match="Must provide at least one StaticMaker"):
+        MatPesStaticFlowMaker(static1=None, static2=None, static3=None)
+
+    # Test no static3 if structure requires no Hubbard U corrections
+    flow_maker = MatPesStaticFlowMaker(static1=None, static2=None)
+    flow = flow_maker.make(si_struct)
+    assert len(flow) == 0
+    assert flow.output == {}