diff --git a/qctrlopencontrols/driven_controls/driven_control.py b/qctrlopencontrols/driven_controls/driven_control.py index c14fc92a..87e171fc 100644 --- a/qctrlopencontrols/driven_controls/driven_control.py +++ b/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. diff --git a/tests/test_driven_controls.py b/tests/test_driven_controls.py index e14f7647..3cbc59f2 100644 --- a/tests/test_driven_controls.py +++ b/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])