From ede3511e72b01d6b0cea0b859ac2fb797bf60a82 Mon Sep 17 00:00:00 2001
From: Talley Lambert <talley.lambert@gmail.com>
Date: Thu, 21 Nov 2024 13:34:21 -0500
Subject: [PATCH] feat: add MDAEvent.slm_image (#204)

* feat: add MDAEvent.slm_image

* new assertion

* fix 3.8

* add exposure
---
 src/useq/_mda_event.py | 48 +++++++++++++++++++++++++++++++++++++++++-
 tests/test_sequence.py | 19 +++++++++++++++++
 2 files changed, 66 insertions(+), 1 deletion(-)

diff --git a/src/useq/_mda_event.py b/src/useq/_mda_event.py
index f4b0327..5c69c30 100644
--- a/src/useq/_mda_event.py
+++ b/src/useq/_mda_event.py
@@ -14,7 +14,9 @@
     Tuple,
 )
 
-from pydantic import Field, field_validator
+import numpy as np
+import numpy.typing as npt
+from pydantic import Field, field_validator, model_validator
 
 from useq._actions import AcquireImage, AnyAction
 from useq._base_model import UseqModel
@@ -52,6 +54,43 @@ def __eq__(self, _value: object) -> bool:
         return super().__eq__(_value)
 
 
+class SLMImage(UseqModel):
+    """SLM Image in a MDA event.
+
+    This object can be cast to a numpy.array using `np.asarray` or `np.array`.
+
+    Attributes
+    ----------
+    data: npt.ArrayLike
+        Image data. Anything that can be cast to a numpy array. For pydantic simplicity,
+        we mark this as Any, but in practice it should be numpy.typing.ArrayLike (which
+        is anything that can be cast to a numpy array using `np.asarray`).
+    device: Optional[str]
+        Optional name of the SLM device to use. If not provided, the "default" SLM
+        device should be used. (It is left to the backend to determine what device that
+        is). By default, `None`.
+    exposure: Optional[float]
+        Exposure time for the SLM specifically (if different from the detector), in
+        milliseconds. If not provided, the exposure on the owning MDAEvent should be
+        used. By default, `None`.
+    """
+
+    data: Any
+    device: Optional[str] = None
+    exposure: Optional[float] = None
+
+    @model_validator(mode="before")
+    def _cast_data(cls, v: Any) -> Any:
+        """Can single, non-dict values to be the data."""
+        if not isinstance(v, dict):
+            v = {"data": v}
+        return v
+
+    def __array__(self, *args: Any, **kwargs: Any) -> npt.NDArray:
+        """Cast the image data to a numpy array."""
+        return np.asarray(self.data, *args, **kwargs)
+
+
 class PropertyTuple(NamedTuple):
     """Three-tuple capturing a device, property, and value.
 
@@ -110,6 +149,12 @@ class MDAEvent(UseqModel):
     z_pos : float | None
         Z position in microns. If not provided, implies use current position. By
         default, `None`.
+    slm_image : SLMImage | None
+        Image data to display on an SLM device. `SLMImage` is a simple pydantic object
+        with two attributes: `data` and `device`. `data` is the image data (anything
+        that can be cast to a numpy array), `device` is the name of the SLM device to
+        use. If not provided, the "default" SLM device should be used. By default,
+        `None`.
     sequence : MDASequence | None
         A reference to the [`useq.MDASequence`][] this event belongs to. This is a
         read-only attribute. By default, `None`.
@@ -146,6 +191,7 @@ class MDAEvent(UseqModel):
     x_pos: Optional[float] = None
     y_pos: Optional[float] = None
     z_pos: Optional[float] = None
+    slm_image: Optional[SLMImage] = None
     sequence: Optional["MDASequence"] = Field(default=None, repr=False)
     properties: Optional[List[PropertyTuple]] = None
     metadata: Dict[str, Any] = Field(default_factory=dict)
diff --git a/tests/test_sequence.py b/tests/test_sequence.py
index 71efd45..3c7a846 100644
--- a/tests/test_sequence.py
+++ b/tests/test_sequence.py
@@ -22,6 +22,7 @@
     ZRangeAround,
     ZRelativePositions,
 )
+from useq._mda_event import SLMImage
 from useq._position import RelativePosition
 
 _T = List[Tuple[Any, Sequence[float]]]
@@ -462,3 +463,21 @@ def test_reset_event_timer() -> None:
     assert not events[1].reset_event_timer
     assert events[2].reset_event_timer
     assert not events[3].reset_event_timer
+
+
+def test_slm_image() -> None:
+    data = [[0, 0], [1, 1]]
+
+    # directly passing data
+    event = MDAEvent(slm_image=data)
+    assert isinstance(event.slm_image, SLMImage)
+
+    # we can cast SLMIamge to a numpy array
+    assert isinstance(np.asarray(event.slm_image), np.ndarray)
+    np.testing.assert_array_equal(event.slm_image, np.array(data))
+
+    # variant that also specifies device label
+    event2 = MDAEvent(slm_image={"data": data, "device": "SLM"})
+    assert event2.slm_image is not None
+    np.testing.assert_array_equal(event2.slm_image, np.array(data))
+    assert event2.slm_image.device == "SLM"