feat: Add function to sample a driven control

qctrlopencontrols/driven_controls/driven_control.py

@@ -20,10 +20,13 @@
 from typing import (
     Dict,
     Optional,
+    Tuple,
+    Union,
 )
 
 import numpy as np
 
+from ..exceptions import ArgumentsValueError
 from ..utils import (
     Coordinate,
     FileFormat,
@@ -330,6 +333,72 @@ def duration(self) -> float:
 
         return np.sum(self.durations)
 
+    def sample(
+        self,
+        sample_times: Union[np.ndarray, float],
+        coordinates: str = Coordinate.CYLINDRICAL.value,
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        r"""
+        Returns samples from the control.
+
+        Parameters
+        ----------
+        sample_times : np.ndarray or float
+            The times at which to sample, :math:`\{t_n\}`. You can pass either an array of times or
+            a single number :math:`dt`. In the latter case, the sample times are taken as
+            :math:`0, dt, 2dt, \dots`.
+        coordinates : str, optional
+            The coordinate system in which to return the sampled control. Must be 'cylindrical' or
+            'cartesian'. Defaults to 'cylindrical'.
+
+        Returns
+        -------
+        amplitude_x or rabi_rates : np.ndarray
+            For cylindrical coordinates, the sampled Rabi rates :math:`\{\Omega(t_n)\}`. For
+            Cartesian coordinates, the sampled x-amplitudes, :math:`\{\Omega(t_n) \cos \phi(t_n)\}`.
+        amplitude_y or azimuthal_angles : np.ndarray
+            For cylindrical coordinates, the sampled azimuthal angles :math:\{\phi(t_n)\}`. For
+            Cartesian coordinates, The sampled y-amplitudes, :math:`\{\Omega(t_n) \sin \phi(t_n)\}`.
+        detunings : np.ndarray
+            The sampled detunings, :math:`\{\Delta(t_n)\}`.
+
+        Raises
+        ------
+        ArgumentsValueError
+            If the inputs are invalid.
+        """
+        times = np.asarray(sample_times)
+        if times.shape == ():
+            times = np.arange(0, self.duration, times)
+
+        check_arguments(
+            len(times.shape) == 1,
+            "Sample times must be a 1D array or a single number.",
+            {"sample_times": sample_times},
+        )
+
+        indices = np.digitize(times, bins=np.cumsum(self.durations), right=True)
+
+        if coordinates == Coordinate.CARTESIAN.value:
+            return (
+                self.amplitude_x[indices],
+                self.amplitude_y[indices],
+                self.detunings[indices],
+            )
+
+        if coordinates == Coordinate.CYLINDRICAL.value:
+            return (
+                self.rabi_rates[indices],
+                self.azimuthal_angles[indices],
+                self.detunings[indices],
+            )
+
+        raise ArgumentsValueError(
+            "Requested coordinate type is not supported. Please use "
+            f"one of {Coordinate.CARTESIAN.value!r} and {Coordinate.CYLINDRICAL.value!r}",
+            {"coordinates": coordinates},
+        )
+
     def _qctrl_expanded_export_content(self, coordinates: str) -> Dict:
         """
         Prepare the content to be saved in Q-CTRL expanded format.

tests/test_driven_controls.py

@@ -406,3 +406,87 @@ def test_pretty_print():
     expected_string = "\n".join(_pretty_string)
 
     assert str(driven_control) == expected_string
+
+
+def test_sample_dt_cylindrical():
+    driven_control = DrivenControl(
+        rabi_rates=[0, 2],
+        azimuthal_angles=[1.5, 0.5],
+        detunings=[1.3, 2.3],
+        durations=[1, 1],
+        name="control",
+    )
+
+    rabi_rates, azimuthal_angles, detunings = driven_control.sample(0.3)
+
+    assert len(rabi_rates) == 7
+    assert len(azimuthal_angles) == 7
+    assert len(detunings) == 7
+
+    assert np.allclose(rabi_rates, [0, 0, 0, 0, 2, 2, 2])
+    assert np.allclose(azimuthal_angles, [1.5, 1.5, 1.5, 1.5, 0.5, 0.5, 0.5])
+    assert np.allclose(detunings, [1.3, 1.3, 1.3, 1.3, 2.3, 2.3, 2.3])
+
+
+def test_sample_dt_cartesian():
+    driven_control = DrivenControl(
+        rabi_rates=[3, 2],
+        azimuthal_angles=[0, np.pi / 2],
+        detunings=[1.3, 2.3],
+        durations=[1, 1],
+        name="control",
+    )
+
+    x_amplitudes, y_amplitudes, detunings = driven_control.sample(0.3, "cartesian")
+
+    assert len(x_amplitudes) == 7
+    assert len(y_amplitudes) == 7
+    assert len(detunings) == 7
+
+    assert np.allclose(x_amplitudes, [3, 3, 3, 3, 0, 0, 0])
+    assert np.allclose(y_amplitudes, [0, 0, 0, 0, 2, 2, 2])
+    assert np.allclose(detunings, [1.3, 1.3, 1.3, 1.3, 2.3, 2.3, 2.3])
+
+
+def test_sample_cylindrical():
+    driven_control = DrivenControl(
+        rabi_rates=[0, 2],
+        azimuthal_angles=[1.5, 0.5],
+        detunings=[1.3, 2.3],
+        durations=[1, 1],
+        name="control",
+    )
+
+    rabi_rates, azimuthal_angles, detunings = driven_control.sample(
+        np.array([0.5, 0.8, 1.5])
+    )
+
+    assert len(rabi_rates) == 3
+    assert len(azimuthal_angles) == 3
+    assert len(detunings) == 3
+
+    assert np.allclose(rabi_rates, [0, 0, 2])
+    assert np.allclose(azimuthal_angles, [1.5, 1.5, 0.5])
+    assert np.allclose(detunings, [1.3, 1.3, 2.3])
+
+
+def test_sample_cartesian():
+    driven_control = DrivenControl(
+        rabi_rates=[3, 2],
+        azimuthal_angles=[0, np.pi / 2],
+        detunings=[1.3, 2.3],
+        durations=[1, 1],
+        name="control",
+    )
+
+    x_amplitudes, y_amplitudes, detunings = driven_control.sample(
+        np.array([0.5, 0.8, 1.5]), "cartesian"
+    )
+
+    assert len(x_amplitudes) == 3
+    assert len(y_amplitudes) == 3
+    assert len(detunings) == 3
+
+    assert np.allclose(x_amplitudes, [3, 3, 0])
+    assert np.allclose(y_amplitudes, [0, 0, 2])
+    assert np.allclose(detunings, [1.3, 1.3, 2.3])