support one_to_many transformations, and add test for both behaviors

Author: esoteric-ephemera

src/atomate2/common/jobs/transform.py

@@ -25,18 +25,42 @@ class Transformer(Maker):
     """Apply a pymatgen transformation, as a job.
 
     For many of the standard and advanced transformations,
-    this will work just by supplying the transformation.
+    this should "just work" by supplying the transformation.
     """
 
     transformation: AbstractTransformation
     name: str = "pymatgen transformation maker"
 
     @job
-    def make(self, structure: Structure, **kwargs) -> TransformTask:
-        """Run the transformation."""
+    def make(
+        self, structure: Structure, **kwargs
+    ) -> TransformTask | list[TransformTask]:
+        """Evaluate the transformation.
+
+        Parameters
+        ----------
+        structure : Structure to transform
+        **kwargs : to pass to the `apply_transformation` method
+
+        Returns
+        -------
+        list of TransformTask, if `self.transformation.is_one_to_many`
+        (many structures are produced from a single transformation)
+
+        TransformTask, otherwise
+        """
         transformed_structure = self.transformation.apply_transformation(
             structure, **kwargs
         )
+        if self.transformation.is_one_to_many:
+            return [
+                TransformTask(
+                    input_structure=structure,
+                    final_structure=dct["structure"],
+                    transformation=dct.get("transformation") or self.transformation,
+                )
+                for dct in transformed_structure
+            ]
         return TransformTask(
             input_structure=structure,
             final_structure=transformed_structure,

tests/common/jobs/test_transform.py

@@ -7,12 +7,16 @@
 
 import numpy as np
 import pytest
-from jobflow import run_locally
+from jobflow import Flow, run_locally
 from pymatgen.core import Structure
 from pymatgen.transformations.advanced_transformations import SQSTransformation
+from pymatgen.transformations.standard_transformations import (
+    OrderDisorderedStructureTransformation,
+    OxidationStateDecorationTransformation,
+)
 
-from atomate2.common.jobs.transform import SQS
-from atomate2.common.schemas.transform import SQSTask
+from atomate2.common.jobs.transform import SQS, Transformer
+from atomate2.common.schemas.transform import SQSTask, TransformTask
 
 
 @pytest.fixture(scope="module")
@@ -32,6 +36,43 @@ def simple_alloy() -> Structure:
     )
 
 
+def test_simple_and_advanced():
+    # simple disordered zincblende structure
+    structure = Structure(
+        3.8 * np.array([[0.0, 0.5, 0.5], [0.5, 0.0, 0.5], [0.5, 0.5, 0.0]]),
+        ["Zn", {"S": 0.75, "Se": 0.25}],
+        [[0.0, 0.0, 0.0], [0.25, 0.25, 0.25]],
+    ).to_conventional()
+
+    oxi_dict = {"Zn": 2, "S": -2, "Se": -2}
+    oxi_job = Transformer(
+        name="oxistate", transformation=OxidationStateDecorationTransformation(oxi_dict)
+    ).make(structure)
+
+    odst_job = Transformer(
+        name="odst", transformation=OrderDisorderedStructureTransformation()
+    ).make(oxi_job.output.final_structure, return_ranked_list=2)
+
+    flow = Flow([oxi_job, odst_job])
+    resp = run_locally(flow)
+
+    oxi_state_output = resp[oxi_job.uuid][1].output
+    assert isinstance(oxi_state_output, TransformTask)
+
+    # check correct assignment of oxidation states
+    assert all(
+        specie.oxi_state == oxi_dict.get(specie.element.value)
+        for site in oxi_state_output.final_structure
+        for specie in site.species
+    )
+
+    odst_output = resp[odst_job.uuid][1].output
+    # return_ranked_list = 2, so should get 2 output docs
+    assert len(odst_output) == 2
+    assert all(isinstance(doc, TransformTask) for doc in odst_output)
+    assert all(doc.final_structure.is_ordered for doc in odst_output)
+
+
 @pytest.mark.skipif(
     icet is None, reason="`icet` must be installed to perform this test."
 )