Discretization and Slicing support for smooth Waveforms

[johnzl-777]

In the Aquila whitepaper the Floquet protocol has a sine wave for its detuning. With the implementation using the Braket SDK this sine wave is discretized into individual values and durations, dependent on a user-specified run time using numpy. The user specified time is defined as a list of times considering the goal of running it on hardware and getting measurements throughout the duration of the protocol.

In the current builder methodology it isn't entirely clear how to go about this considering the run time should ideally be treated as a scalar Variable but this is incompatible with numpy (you can't plug in our Variable type into the np.linspace that gets used to generate the durations for the resulting piecewise linear function).

Ideally there'd be a way to pass in a smooth waveform and have it discretized, BUT this discretized waveform must also be sliceable (to support such operations as batch_assign where variations of the task can be produced with different run times, in line with the white paper method of generating tasks).

[weinbe58]

I think for consistency we should have a new IR objects:

Sample(waveform: Waveform, sample_points: List[Scalar])

which would allow the user to symbolically sample the waveform at a list of time points that can be determined at compile time and:

Function(python_func: Callable, args: Tuple[Scalar], kwargs: Dict[str, Scalar]) to specify an arbitrary function with parameterized arguments

[Roger-luo]

we should have a FFI node in the IR as following

@dataclass
class PythonFn(Instruction):

    fn: Callable # [[float, ...], float] # f(t) -> value
    parameters: List[Variable] # come from ast inspect
    duration: InitVar[Scalar]

    def __init__(self, fn: Callable[[Scalar], Scalar], duration: Scalar):
        pass

    def __call__(self, clock_s: float, **kwargs) -> float:
        return self.fn(clock_s, **kwargs)

@waveform(duration=2.0)
def sin(t: float, alpha: float) -> float:
    import math
    return math.sin(alpha * t)

def scan(alpha: float) -> PythonFn:
    def sin(t: float) -> float:
        import math
        return math.sin(alpha * t)
    return PythonFn(sin, 1.0)

PythonFn(sin, 1.0)

[Roger-luo]

the corresponding builder interface

start.rydberg.detuning.fn(sin).sample(dt=1e-3) # fn needs have a separate builder than Waveform that contains a post method sample available

[weinbe58]

Few things:

1. we can't support arguments like this so I guess we should error when encountering ambiguous signatures:

sin(t: float, *args, **kwargs):
    pass

2. For default arguments we should bring back the default value for variables.

[weinbe58]

The following will work after merging PR #167

def drive(time, *, omega, phi=0, amplitude):
    return amplitude * np.cos(omega * time + phi)

durations = cast([0.1, "run_time", 0.1])
total_duration = sum(durations)

task = (
    start.add_position((0, 0))
    .rydberg.detuning.uniform.fn(drive, total_duration)
    .sample(0.05, "linear")
    .rydberg.rabi.amplitude.uniform.piecewise_linear(
        durations, [0, "rabi_max", "rabi_max", 0]
    )
    .assign(omega=15, amplitude=15, rabi_max=15)
    .batch_assign(run_time=np.linspace(0, 4.0, 101))
    .braket_local_simulator(1000)
)

This should also fix a lot of the sampling issues we had in the original examples since no slicing is required 🎊