Added loopable arguments to AM and PM modulation

pulse_lib/segments/data_classes/data_IQ.py

-"""
-data class to store IQ based signals.
-"""
+from collections.abc import Callable
+
 from scipy import signal
 import numpy as np
 
 from dataclasses import dataclass
 
+
+EnvelopeFunction = Callable[[float, float, ...], np.ndarray]
+
+
 class envelope_generator():
-    """
-    Object that handles envelope functions that can be used in spin qubit experiments.
-    Key properties
-        * Makes sure average amplitude of the evelope is the one expressed
-        * Executes some subsampling functions to give greater time resolution than the sample rate of the AWG.
-        * Allows for plotting the FT of the envelope function.
-    """
-    def __init__(self, AM_envelope_function = None, PM_envelope_function = None):
+    def __init__(self,
+                 AM_envelope_function: str | tuple[str, ...] | EnvelopeFunction | None = None,
+                 PM_envelope_function: str | tuple[str, ...] | EnvelopeFunction | None = None,
+                 kwargs: dict[str, any] | None = None):
         """
         define envelope funnctions.
         Args
-            AM_envelope_function (str/tuple) : spec of a windowing as defined for the scipy.signal.windows.get_window function
-            AM_envelope_function (lamba M) : (function overload) function where M is the number of points where the evelopen should be rendered for.
-            AM_envelope_function (lamba M) : (function overload) function where M is the number of points where the evelopen should be rendered for.
+            AM_envelope_function ('str/tuple/function') : function describing an amplitude modulation (see examples in pulse_lib.segments.data_classes.data_IQ)
+            PM_envelope_function ('str/tuple/function') : function describing an phase modulation (see examples in pulse_lib.segments.data_classes.data_IQ)
+            kwargs: keyword arguments passed into the AM and PM functions.
         """
         self.AM_envelope_function = AM_envelope_function
         self.PM_envelope_function = PM_envelope_function
+        self.kwargs = kwargs
 
-    def get_AM_envelope(self, delta_t, sample_rate=1):
+    def get_AM_envelope(self, delta_t: float, sample_rate: float = 1.0):
         """
-        Render the envelope for the given waveshape (in init).
+        Render the amplitude envelope for the wave.
 
         Args:
-            delta_t (float) : time of the pulse (5.6 ns)
-            sample_rate (float) : number of samples per second (e.g. 1GS/s)
+            delta_t: time of the pulse [ns]
+            sample_rate: number of samples per second [GSa/s]
 
         Returns:
             envelope (np.ndarray[ndim=1, dtype=double]) : envelope function in DC
@@ -46,17 +46,17 @@ class envelope_generator():
         elif isinstance(self.AM_envelope_function, tuple) or isinstance(self.AM_envelope_function, str):
             envelope = signal.get_window(self.AM_envelope_function, int(n_points*10))[::10][:int(n_points)] #ugly fix
         else:
-            envelope = self.AM_envelope_function(delta_t, sample_rate) # user reponsible to do good subsampling him/herself.
+            envelope = self.AM_envelope_function(delta_t, sample_rate, **self.kwargs)
 
         return envelope
 
-    def get_PM_envelope(self, delta_t, sample_rate=1):
+    def get_PM_envelope(self, delta_t: float, sample_rate: float = 1):
         """
-        Render the envelope for the given waveshape (in init).
+        Return the phase modulation values for the wave.
 
         Args:
-            delta_t (float) : time of the pulse (5.6 ns)
-            sample_rate (float) : number of samples per second (e.g. 1GS/s)
+            delta_t: time of the pulse [ns]
+            sample_rate: number of samples per second [GSa/s]
 
         Returns:
             envelope (np.ndarray[ndim=1, dtype=double]) : envelope function in DC
@@ -71,11 +71,12 @@ class envelope_generator():
         elif isinstance(self.PM_envelope_function, tuple) or isinstance(self.PM_envelope_function, str):
             envelope = signal.get_window(self.PM_envelope_function, int(n_points*10))[::10]
         else:
-            envelope = self.PM_envelope_function(delta_t, sample_rate) # user reponsible to do good subsampling him/herself.
+            envelope = self.PM_envelope_function(delta_t, sample_rate, **self.kwargs)
 
         return envelope
 
-def make_chirp(f_start, f_stop, time0, time1):
+
+def make_chirp(f_start: float, f_stop: float, time0: float, time1: float):
     '''
     Make a chirp.
 
@@ -86,7 +87,7 @@ def make_chirp(f_start, f_stop, time0, time1):
 
     chirp_constant = (f_stop - f_start)/(time1*1e-9-time0*1e-9)/2
 
-    def my_chirp(delta_t, sample_rate = 1):
+    def my_chirp(delta_t: float, sample_rate: float = 1.0):
         """
         Function that makes a phase envelope to make a chirped pulse
 
@@ -105,13 +106,14 @@ def make_chirp(f_start, f_stop, time0, time1):
 
     return my_chirp
 
+
 @dataclass
 class IQ_data_single:
-    start : float = 0
-    stop : float = 0
-    amplitude : float = 1
-    frequency : float = 0
-    phase_offset : float = 0
+    start : float = 0.0
+    stop : float = 0.0
+    amplitude : float = 1.0
+    frequency : float = 0.0
+    phase_offset : float = 0.0
     ''' offset from coherent pulse  '''
     envelope : envelope_generator = None
     ref_channel : str = None
@@ -122,6 +124,7 @@ class IQ_data_single:
         ''' Old name for phase_offset '''
         return self.phase_offset
 
+
 @dataclass
 class Chirp:
     start : float
@@ -138,6 +141,7 @@ class Chirp:
                           self.stop_frequency,
                           self.start, self.stop)
 
+
 if __name__ == '__main__':
 
     """

pulse_lib/segments/segment_IQ.py

@@ -54,7 +54,7 @@ class segment_IQ(segment_base):
         return self.data_tmp
 
     @loop_controller
-    def add_MW_pulse(self, t0, t1, amp, freq, phase = 0, AM = None, PM = None):
+    def add_MW_pulse(self, t0, t1, amp, freq, phase = 0, AM = None, PM = None, **kwargs):
         '''
         Make a sine pulse (generic constructor)
 
@@ -66,12 +66,13 @@ class segment_IQ(segment_base):
             phase (float) : phase of the microwave.
             AM ('str/tuple/function') : function describing an amplitude modulation (see examples in pulse_lib.segments.data_classes.data_IQ)
             PM ('str/tuple/function') : function describing an phase modulation (see examples in pulse_lib.segments.data_classes.data_IQ)
+            kwargs: keyword arguments passed into the AM and PM functions.
         '''
         MW_data = IQ_data_single(t0 + self.data_tmp.start_time,
                                  t1 + self.data_tmp.start_time,
                                  amp, freq,
                                  phase,
-                                 envelope_generator(AM, PM),
+                                 envelope_generator(AM, PM, kwargs),
                                  self.name)
         self.data_tmp.add_MW_data(MW_data)
         return self.data_tmp

pulse_lib/tests/test_iq/test_MW_mod.py

+
+from pulse_lib.tests.configurations.test_configuration import context
+
+#%%
+import numpy as np
+
+from pulse_lib.segments.utility import looping as lp
+
+def get_AM_envelope(delta_t: float, sample_rate: float):
+    npt = int(delta_t*sample_rate + 0.5)
+    return np.linspace(0, 1.0, npt)
+
+
+def get_AM_envelope2(delta_t: float, sample_rate: float, alpha: float):
+    npt = int(delta_t*sample_rate + 0.5)
+    return np.linspace(alpha, 1.0, npt)
+
+
+
+def test1():
+    pulse = context.init_pulselib(n_gates=1, n_qubits=1)
+
+    f_q1 = pulse.qubit_channels['q1'].resonance_frequency
+
+    s = pulse.mk_segment()
+
+    s.q1.add_MW_pulse(0, 200, 100, f_q1, AM=get_AM_envelope)
+
+    sequence = pulse.mk_sequence([s])
+    sequence.n_rep = 1
+    context.plot_awgs(sequence)
+
+    return None
+
+
+def test2():
+    pulse = context.init_pulselib(n_gates=1, n_qubits=1)
+
+    f_q1 = pulse.qubit_channels['q1'].resonance_frequency
+    alpha = lp.linspace(0.2, 1.0, 5, "alpha", axis=0)
+
+    s = pulse.mk_segment()
+
+    s.q1.add_MW_pulse(0, 200, 100, f_q1, AM=get_AM_envelope2, alpha=alpha)
+
+    sequence = pulse.mk_sequence([s])
+    sequence.n_rep = 1
+    for i in range(len(alpha)):
+        context.plot_awgs(sequence, index=(i,))
+
+    return None
+
+
+class IQ_modulation:
+    def __init__(
+            self,
+            I: list[lp.loop_obj] | list[float],
+            Q: list[lp.loop_obj] | list[float],
+            segment_length: int):
+        self.I = I
+        self.Q = Q
+        self.segment_length = segment_length
+        if len(I) != len(Q):
+            raise Exception("I and Q list must have same length")
+        self.n_segments = len(I)
+        self.length = self.n_segments * segment_length
+
+        self.indexed = hasattr(I[0], 'axis')
+        if self.indexed:
+            axis = I[0].axis[0]
+            n = I[0].shape[0]
+            for loop in I + Q:
+                if loop.axis != [axis] or loop.shape != (n,):
+                    raise Exception(f"All loops must have axis {axis} and shape {(n,)}. "
+                                    f"Error on loop {loop}")
+            self.loop = lp.arange(0, n, name="index", axis=axis)
+        else:
+            self.loop = None
+
+    def get_index_loop(self):
+        return self.loop
+
+    def _render_IQ(self, delta_t: float, sample_rate: float, index: float = 0.0):
+        # NOTE: looping arguments are always float. Cast to int.
+        index = int(index)
+        if round(delta_t*sample_rate) != self.length:
+            raise Exception("Duration of wave doesn't match modulation spec. "
+                            f"Expected {self.length}, but got {round(delta_t*sample_rate)} ns")
+        if self.indexed:
+            i_values = [i[index] for i in self.I]
+            q_values = [q[index] for q in self.Q]
+        else:
+            i_values = self.I
+            q_values = self.Q
+
+        iq = np.zeros(self.length, dtype=complex)
+        size = self.segment_length
+        for i in range(self.n_segments):
+            iq[i*size : (i+1)*size] = i_values[i] + 1j*q_values[i]
+
+        return iq
+
+    def get_AM_envelope_indexed(self, delta_t: float, sample_rate: float, index: float = 0.0):
+        iq_data = self._render_IQ(delta_t, sample_rate, index)
+        return np.abs(iq_data)
+
+    def get_PM_envelope_indexed(self, delta_t: float, sample_rate: float, index: float = 0.0):
+        iq_data = self._render_IQ(delta_t, sample_rate, index)
+        return np.angle(iq_data)
+
+
+def testIQloop():
+    pulse = context.init_pulselib(n_gates=1, n_qubits=1)
+
+    f_q1 = pulse.qubit_channels['q1'].resonance_frequency
+
+    iq_mod = IQ_modulation(
+        I=[
+            lp.array([0.5, 0.2, 0.1], "i1", axis=0),
+            lp.array([0.5, 0.8, -0.1], "i2", axis=0),
+        ],
+        Q=[
+            lp.array([0.0, 0.2, 0.5], "q1", axis=0),
+            lp.array([0.0, 0.8, -0.5], "q2", axis=0),
+        ],
+        segment_length=100,
+        )
+
+    s = pulse.mk_segment()
+
+    s.q1.add_MW_pulse(0, 200, 100, f_q1,
+                      AM=iq_mod.get_AM_envelope_indexed,
+                      PM=iq_mod.get_PM_envelope_indexed,
+                      index=iq_mod.get_index_loop())
+
+    sequence = pulse.mk_sequence([s])
+    sequence.n_rep = 1
+    for i in range(len(iq_mod.get_index_loop())):
+        context.plot_awgs(sequence, index=(i,))
+
+    return None
+
+
+
+#%%
+
+if __name__ == '__main__':
+    # ds1 = test1()
+    # ds2 = test2()
+    testIQloop()