Source code for semi_cr.core.trial.shapes

from abc import ABC, abstractmethod

import numpy as np

TWO_DIMENSIONS = 2




class PulseShape(ABC):
    """
    Abstract base class for pulse shapes.
    """



    @abstractmethod
    def get_shape(self) -> np.typing.NDArray[np.floating]:
        """
        Get the value of the pulse shape.
        :return: The value of the pulse shape at the given time.
        """
        pass




    def integrate(self):
        return self.get_shape().sum()






class YizhiPulse(PulseShape):
    """
    A pulse shape that linearly ramps up to a specified amplitude, then "asymptotically" goes on from that amplitude's
    complement to the negative of that amplitude, then back to zero. Sort of like this:
    """

    def __init__(self, amplitude: float, steps: int):
        if not np.isfinite(amplitude) or amplitude < 0:
            raise ValueError("Amplitude must be non-negative and finite.")

        if steps <= 0:
            raise ValueError("Steps must be a positive integer.")

        self.amplitude = amplitude
        self.steps = steps



    def get_shape(self) -> np.typing.NDArray[np.floating]:
        amplitude_shape = np.linspace(-self.amplitude, self.amplitude, self.steps)
        return amplitude_shape






class YizhiPulse2D(PulseShape):
    """
    A 2D pulse shape that linearly ramps up to a specified amplitude, then "asymptotically" goes on from that amplitude's
    complement to the negative of that amplitude, then back to zero. Sort of like this:
    """

    def __init__(self, amplitude: float | tuple[float, float], steps: int):
        amplitudes = (amplitude, amplitude) if isinstance(amplitude, (int, float)) else amplitude
        if len(amplitudes) != TWO_DIMENSIONS or any(not np.isfinite(value) or value < 0 for value in amplitudes):
            raise ValueError("Each amplitude must be non-negative and finite.")

        if steps <= 0:
            raise ValueError("Steps must be a positive integer.")

        self.amplitude = amplitudes
        self.steps = steps



    def get_shape(self) -> np.typing.NDArray[np.floating]:
        gate_1_amplitude_sequences, gate_2_amplitude_sequences = self.get_axes()
        return np.column_stack((gate_1_amplitude_sequences, gate_2_amplitude_sequences))




    def get_axes(self) -> tuple[np.typing.NDArray[np.floating], np.typing.NDArray[np.floating]]:
        """
        Get the axis values for the 2D pulse shape.
        :return: The axis values for the 2D pulse shape.
        """
        return tuple(np.linspace(-amplitude, amplitude, self.steps) for amplitude in self.amplitude)  # type: ignore[return-value]