make structure/molecule aliases of mol_or_struct, allow for removal from ionic step properties

Author: esoteric-ephemera

src/atomate2/ase/schemas.py

@@ -20,7 +20,7 @@
 from emmet.core.structure import MoleculeMetadata, StructureMetadata
 from emmet.core.trajectory import AtomTrajectory
 from emmet.core.types.enums import StoreTrajectoryOption, TaskState, ValueEnum
-from pydantic import BaseModel, Field
+from pydantic import AliasChoices, BaseModel, Field
 from pymatgen.core import Molecule, Structure
 from pymatgen.core.trajectory import Trajectory as PmgTrajectory
 from pymatgen.entries.computed_entries import ComputedEntry
@@ -116,17 +116,24 @@ class AseBaseModel(BaseModel):
     """Base document class for ASE input and output."""
 
     mol_or_struct: Structure | Molecule | None = Field(
-        None, description="The molecule or structure at this step."
+        None,
+        description="The molecule or structure at this step.",
+        validation_alias=AliasChoices("mol_or_struct", "structure", "molecule"),
     )
-    structure: Structure | None = Field(None, description="The structure at this step.")
-    molecule: Molecule | None = Field(None, description="The molecule at this step.")
 
-    def model_post_init(self, context: Any, /) -> None:
-        """Establish alias to structure and molecule fields."""
-        if self.structure is None and isinstance(self.mol_or_struct, Structure):
-            self.structure = self.mol_or_struct
-        elif self.molecule is None and isinstance(self.mol_or_struct, Molecule):
-            self.molecule = self.mol_or_struct
+    @property
+    def structure(self) -> Structure | None:
+        """Retrieve the structure associated with this document, if applicable."""
+        if isinstance(self.mol_or_struct, Structure):
+            return self.mol_or_struct
+        return None
+
+    @property
+    def molecule(self) -> Molecule | None:
+        """Retrieve the molecule associated with this document, if applicable."""
+        if isinstance(self.mol_or_struct, Molecule):
+            return self.mol_or_struct
+        return None
 
 
 class IonicStep(AseBaseModel):
@@ -476,8 +483,18 @@ def from_ase_compatible_result(
         final_stress = None
         ionic_steps = None
 
+        if "mol_or_struct" not in (
+            user_ionic_step_data := set(ionic_step_data or tuple())
+        ):
+            for ms_alias in ("molecule", "structure"):
+                if ms_alias in user_ionic_step_data:
+                    user_ionic_step_data.add("mol_or_struct")
+
         if trajectory:
             ionic_step_props = {"energy", "forces"}
+            if save_atoms := "mol_or_struct" in user_ionic_step_data:
+                user_ionic_step_data.remove("mol_or_struct")
+
             if isinstance(trajectory, AtomTrajectory):
                 final_energy = trajectory.energy[-1]
                 final_forces = trajectory.forces[-1]
@@ -501,21 +518,17 @@ def from_ase_compatible_result(
                     ionic_step_props.add("magmoms")
 
             ionic_steps = []
-            if (
-                len(
-                    use_ionic_step_props := ionic_step_props.intersection(
-                        ionic_step_data or set()
-                    )
-                )
-                > 0
-            ):
+            use_ionic_step_props = ionic_step_props.intersection(user_ionic_step_data)
+            if len(use_ionic_step_props) > 0:
                 if isinstance(trajectory, AtomTrajectory):
                     ionic_steps = [
                         IonicStep(
                             mol_or_struct=trajectory.to_pmg(
                                 frame_props=tuple(),
                                 indices=idx,
-                            )[0],
+                            )[0]
+                            if save_atoms
+                            else None,
                             **{
                                 key: getattr(trajectory, key)[idx]
                                 for key in use_ionic_step_props
@@ -527,7 +540,7 @@ def from_ase_compatible_result(
                 else:
                     ionic_steps = [
                         IonicStep(
-                            mol_or_struct=atoms,
+                            mol_or_struct=atoms if save_atoms else None,
                             **{
                                 key: convert_stress_from_voigt_to_symm(
                                     trajectory.frame_properties[idx].get(key)