Source code for WaveformConstructor.waveform

import numpy as np
from scipy.constants import *
import abc
from .function import sine_wave
from .default import DefaultValues
from typing import List

class WaveformFunctionBase(abc.ABC):
    """Abstract class for the waveform function"""

    @abc.abstractmethod
    def evaluate(self, t: np.ndarray) -> np.ndarray:
        """Evaluate the waveform for the timeseries"""
        pass

    def __call__(self, t: np.ndarray) -> np.ndarray:
        return self.evaluate(t)

    def __add__(self, other: "WaveformFunctionBase") -> "WaveformFunctionSum":
        return WaveformFunctionSum([self, other])
    
    def Q(self) -> "WaveformFunctionBase":
        raise NotImplementedError()

class WaveformFunctionSum(WaveformFunctionBase):
    """Sum of waveform functions
    
    Parameters
    ----------
    waveform_funcs: List[WaveformFunctionBase]
        List of waveform functions to be summed
    
    """
    def __init__(self, waveform_funcs: List[WaveformFunctionBase]):
        self.waveform_funcs = waveform_funcs

    def evaluate(self, t: np.ndarray) -> np.ndarray:
        return np.sum([w(t) for w in self.waveform_funcs], axis=0)
    
    def Q(self) -> "WaveformFunctionBase":
        return WaveformFunctionSum([w.Q() for w in self.waveform_funcs])
    

class DigitalWaveformBase(WaveformFunctionBase):
    pass

class On(DigitalWaveformBase):
    def evaluate(self, t: np.ndarray) -> np.ndarray:
        return np.ones_like(t, dtype=bool)
    
class Off(DigitalWaveformBase):
    def evaluate(self, t: np.ndarray) -> np.ndarray:
        return np.zeros_like(t, dtype=bool)
    

class SineWave(WaveformFunctionBase):
    """Sinusoidal waveform

    Parameters
    ----------
    amp: float
        Amplitude of the waveform
    freq: float
        Frequency of the waveform
    phase: float
        Phase of the waveform
    """

    def __init__(self, amp: float, freq: float, phase: float):
        self.amp = amp
        self.freq = freq
        self.phase = phase

    def evaluate(self, t):
        return sine_wave(t, self.amp, self.freq, self.phase)
    
    def Q(self) -> "SineWave":
        return SineWave(self.amp, self.freq, self.phase - np.pi/2)


class Zero(WaveformFunctionBase):
    """Zero waveform"""
    def evaluate(self, t: np.ndarray):
        return np.zeros_like(t)
    
    def Q(self) -> "Zero":
        return Zero()
    
class One(WaveformFunctionBase):
    """One waveform"""
    def evaluate(self, t: np.ndarray):
        return np.ones_like(t)
     
class Constant(WaveformFunctionBase):
    """Constant waveform
    
    Parameters
    ----------
    value: float
        Value of the waveform
    """
    def __init__(self, value: float):
        self.value = value
        
    def evaluate(self, t: np.ndarray) -> np.ndarray:
        return self.value * np.ones_like(t)


class CarrierFlopping(SineWave):
    """Waveform for the carrier flopping.

    Parameters
    ----------
    amp: float
        Amplitude of the waveform
    phase: float
        Phase of the waveform
    freq: float or None
        Driving frequency. If None is given, it uses the value in DefaultValues.
    """

    def __init__(self, amp: float, phase: float, freq=None):
        self.amp = amp
        self.phase = phase

        if freq is None:
            freq = DefaultValues.carrier_freq
        super().__init__(amp, freq, phase)

    def __neg__(self):
        return CarrierFlopping(-self.amp, self.phase, freq=self.freq)


class X(CarrierFlopping):
    def __init__(self, amp, carrier_freq=None):
        super().__init__(amp, 0, freq=carrier_freq)


class Y(CarrierFlopping):
    def __init__(self, amp, carrier_freq=None):
        super().__init__(amp, np.pi / 2, freq=carrier_freq)


class FloquetBase(WaveformFunctionBase):
    """Base class for the Floquet waveform functions
    
    Parameters
    ----------
    floquet_freq: float
        Floquet frequency. Note that the floquet_freq is the frequency where the sign of the waveform changes.

    """
    def __init__(self, floquet_freq, base_waveform: WaveformFunctionBase=None):
        self.floquet_freq = floquet_freq
        if base_waveform is None:
            base_waveform = super()
        self.base_waveform = base_waveform

    def _modulation(self, t: np.ndarray) -> np.ndarray:
        t0 = t[0]
        sign =  np.sin((np.pi * self.floquet_freq) * (t - t0) )
        sign = np.where(sign >= 0, 1, -1)
        return sign

    def evaluate(self, t: np.ndarray) -> np.ndarray:
        return self.base_waveform.evaluate(t) * self._modulation(t)

    def Q(self) -> "FloquetBase":
        return FloquetBase(self.floquet_freq, self.base_waveform.Q())
    

class FloquetSinDrive(FloquetBase):
    def _modulation(self, t: np.ndarray) -> np.ndarray:
        t0 = t[0]
        return (np.pi / 2) * np.sin((np.pi * self.floquet_freq) * (t - t0) )
    
    def Q(self) -> "FloquetSinDrive":
        return FloquetSinDrive(self.floquet_freq, self.base_waveform.Q())
    
class FloquetCosDrive(FloquetBase):
    def _modulation(self, t: np.ndarray) -> np.ndarray:
        t0 = t[0]
        return (np.pi / 2) * np.cos((np.pi * self.floquet_freq) * (t - t0) )
    
    def Q(self) -> "FloquetCosDrive":
        return FloquetCosDrive(self.floquet_freq, self.base_waveform.Q())


class OneFloquet(FloquetBase, One):
    def __init__(self, floquet_freq):
        One.__init__(self)
        FloquetBase.__init__(self, floquet_freq)

class XFloquet(FloquetBase, X):
    def __init__(self, amp, floquet_freq, carrier_freq=None):
        X.__init__(self, amp, carrier_freq=carrier_freq)
        FloquetBase.__init__(self, floquet_freq)


class YFloquet(FloquetBase, Y):
    def __init__(self, amp, floquet_freq, carrier_freq=None):
        Y.__init__(self, amp, carrier_freq=carrier_freq)
        FloquetBase.__init__(self, floquet_freq)

class XFloquetSinDrive(FloquetSinDrive, XFloquet):
    pass

class YFloquetSinDrive(FloquetSinDrive, YFloquet):
    pass


class XFloquetCosDrive(FloquetCosDrive, XFloquet):
    pass

class YFloquetCosDrive(FloquetCosDrive, YFloquet):
    pass


class MS(WaveformFunctionSum):
    """Waveform for the Molmer Sorensen interaction.

    Parameters
    ----------
    amp_rsb: float
        Amplitude of the red sideband waveform.
    amp_bsb: float
        Amplitude of the blue sideband waveform.
    detuning: float
        Detuning from the carrier frequency.
    phase_rsb: float
        Phase of the red sideband waveform.
    phase_bsb: float
        Phase of the blue sideband waveform.
    carrier_freq: float or None
        Carrier frequency. If None is given, it uses the value in DefaultValues.
    """

    def __init__(self, amp_rsb, amp_bsb, detuning, phase_rsb, phase_bsb, carrier_freq=None):
        if carrier_freq is None:
            carrier_freq = DefaultValues.carrier_freq

        self.rsb = SineWave(amp_rsb, carrier_freq - detuning, phase_rsb)
        self.bsb = SineWave(amp_bsb, carrier_freq + detuning, phase_bsb)

        super().__init__([self.rsb, self.bsb])


class XX(MS):
    """Waveform for the XX interaction.
    
    Parameters
    ----------
    amp_rsb: float
        Amplitude of the red sideband waveform.
    amp_bsb: float
        Amplitude of the blue sideband waveform.
    detuning: float
        Detuning from the carrier frequency.
    carrier_freq: float or None
        Carrier frequency. If None is given, it uses the value in DefaultValues.
    """
    def __init__(self, amp_rsb, amp_bsb, detuning, carrier_freq=None):
        super().__init__(amp_rsb, amp_bsb, detuning, 0, -np.pi, carrier_freq=carrier_freq)


class YY(MS):
    """Waveform for the YY interaction.
    
    Parameters
    ----------
    amp_rsb: float
        Amplitude of the red sideband waveform.
    amp_bsb: float
        Amplitude of the blue sideband waveform.
    detuning: float
        Detuning from the carrier frequency.
    carrier_freq: float or None
        Carrier frequency. If None is given, it uses the value in DefaultValues.
    """
    def __init__(self, amp_rsb, amp_bsb, detuning, carrier_freq=None):
        super().__init__(amp_rsb, amp_bsb, detuning, np.pi/2, -np.pi/2, carrier_freq=carrier_freq)

class MS2(WaveformFunctionSum):
    """Waveform for the Molmer Sorensen interaction for two detunings.

    Parameters
    ----------
    amp_rsb1: float
        Amplitude of the red sideband waveform1.
    amp_bsb1: float
        Amplitude of the blue sideband waveform1.
    detuning1: float
        Detuning from the carrier frequency for waveform1.
    phase_rsb1: float
        Phase of the red sideband waveform1.
    phase_bsb1: float
        Phase of the blue sideband waveform1.
    amp_rsb2: float
        Amplitude of the red sideband waveform2.
    amp_bsb2: float
        Amplitude of the blue sideband waveform2.
    detuning2: float
        Detuning from the carrier frequency for waveform2.
    phase_rsb2: float
        Phase of the red sideband waveform2.
    phase_bsb2: float
        Phase of the blue sideband waveform2.
    carrier_freq: float or None
        Carrier frequency. If None is given, it uses the value in DefaultValues.
    """

    def __init__(self, amp_rsb1, amp_bsb1, detuning1, phase_rsb1, phase_bsb1,amp_rsb2, amp_bsb2, detuning2, phase_rsb2, phase_bsb2,carrier_freq=None):
        """"""
        if carrier_freq is None:
            carrier_freq = DefaultValues.carrier_freq

        self.rsb1 = SineWave(amp_rsb1, carrier_freq - detuning1, phase_rsb1)
        self.bsb1 = SineWave(amp_bsb1, carrier_freq + detuning1, phase_bsb1)
        self.rsb2 = SineWave(amp_rsb2, carrier_freq - detuning2, phase_rsb2)
        self.bsb2 = SineWave(amp_bsb2, carrier_freq + detuning2, phase_bsb2)

        super().__init__([self.rsb1, self.bsb1,self.rsb2, self.bsb2])

class XY(MS2):
    def __init__(self, amp_rsb1, amp_bsb1, detuning1,amp_rsb2, amp_bsb2, detuning2, carrier_freq=None):
        super().__init__(amp_rsb1, amp_bsb1, detuning1, 0, -np.pi,amp_rsb2, amp_bsb2, detuning2, np.pi/2, -np.pi/2, carrier_freq=carrier_freq)