Source code for semi_cr.core.trial.shapes

from abc import ABC, abstractmethod

import numpy as np




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 amplitude > 1.0 or amplitude < 0:
            raise ValueError("Amplitude must be between 0 and 1.")

        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]:
        positive_steps = self.steps // 2
        negative_steps = self.steps - positive_steps

        amplitude_sequences_up = np.linspace(0, self.amplitude, positive_steps)
        amplitude_sequences_down = np.linspace(-self.amplitude, 0, negative_steps)
        amplitude_shape = np.concatenate((amplitude_sequences_up, amplitude_sequences_down))

        # If the number of steps is odd, we need to subtract the residual amplitude in order to get a zero integral.
        amplitude_shape -= amplitude_shape.mean()

        return amplitude_shape