add optional run_type kwarg to get_trajectory

Author: tsmathis

mp_api/client/routes/materials/tasks.py

@@ -8,6 +8,7 @@
 from emmet.core.mpid import MPID, AlphaID
 from emmet.core.tasks import CoreTaskDoc
 from emmet.core.trajectory import RelaxTrajectory
+from emmet.core.vasp.calc_types import RunType
 
 from mp_api.client.core import BaseRester, MPRestError
 from mp_api.client.core.utils import validate_ids
@@ -24,18 +25,28 @@ class TaskRester(BaseRester):
     primary_key: str = "task_id"
     delta_backed = True
 
-    def get_trajectory(self, task_id: MPID | AlphaID | str) -> dict[str, Any]:
+    def get_trajectory(
+        self, task_id: MPID | AlphaID | str, run_type: str | RunType | None = None
+    ) -> dict[str, Any]:
         """Returns a Trajectory object containing the geometry of the
         material throughout a calculation. This is most useful for
         observing how a material relaxes during a geometry optimization.
 
         Args:
             task_id (str, MPID, AlphaID): Task ID
+            run_type (str, RunType): Task run type
 
         Returns:
             dict representing emmet.core.trajectory.RelaxTrajectory
         """
         as_alpha = str(AlphaID(task_id, padlen=8)).split("-")[-1]
+
+        predicate = (
+            f"WHERE run_type='{str(run_type)}' AND identifier='{as_alpha}'"
+            if run_type
+            else f"WHERE identifier='{as_alpha}'"
+        )
+
         traj_tbl = DeltaTable(
             "s3a://materialsproject-parsed/core/trajectories/",
             storage_options={"AWS_SKIP_SIGNATURE": "true", "AWS_REGION": "us-east-1"},
@@ -48,7 +59,7 @@ def get_trajectory(self, task_id: MPID | AlphaID | str) -> dict[str, Any]:
                 f"""
                 SELECT *
                 FROM   traj
-                WHERE  identifier='{as_alpha}'
+                {predicate};
                 """
             )
             .read_all()