Added unit tests based on examples for testing compilation (#890)

* Added unit tests based on examples for testing compilation

* Run black

* ran ruff, fixed import problems

* Ran black again

---------

Co-authored-by: Phillip Weinberg <weinbe58@gmail.com>

tests/test_2d_state_compile.py

@@ -0,0 +1,36 @@
+from bloqade.atom_arrangement import Square
+
+
+def test_2d_state_compile():
+    # Have atoms separated by 5.9 micrometers
+    L = 3
+    lattice_spacing = 5.9
+
+    rabi_amplitude_values = [0.0, 15.8, 15.8, 0.0]
+    rabi_detuning_values = [-16.33, -16.33, "delta_end", "delta_end"]
+    durations = [0.8, "sweep_time", 0.8]
+
+    prog = (
+        Square(L, lattice_spacing=lattice_spacing)
+        .rydberg.rabi.amplitude.uniform.piecewise_linear(
+            durations, rabi_amplitude_values
+        )
+        .detuning.uniform.piecewise_linear(durations, rabi_detuning_values)
+    )
+
+    batch = prog.assign(delta_end=42.66, sweep_time=2.4)
+
+    bloqade_emu_target = batch.bloqade.python()
+    braket_emu_target = batch.braket.local_emulator()
+    quera_aquila_target = batch.parallelize(24).quera.aquila()
+    braket_aquila_target = batch.parallelize(24).braket.aquila()
+
+    targets = [
+        bloqade_emu_target,
+        braket_emu_target,
+        quera_aquila_target,
+        braket_aquila_target,
+    ]
+
+    for target in targets:
+        target._compile(10)

tests/test_adiabatic_compile.py

@@ -0,0 +1,42 @@
+from bloqade import start, cast, var
+import numpy as np
+
+
+def test_adiabatic_compile():
+    detuning_value = var("detuning_value")
+    durations = cast(["ramp_time", "run_time", "ramp_time"])
+    prog = (
+        start.add_position([(0, 0), (0, "atom_distance")])
+        .rydberg.rabi.amplitude.uniform.piecewise_linear(
+            durations=durations, values=[0, "rabi_value", "rabi_value", 0]
+        )
+        .detuning.uniform.piecewise_linear(
+            durations=durations,
+            values=[
+                -detuning_value,
+                -detuning_value,
+                detuning_value,
+                detuning_value,
+            ],
+        )
+    )
+
+    distances = np.arange(4, 11, 1)
+    batch = prog.assign(
+        ramp_time=1.0, run_time=2.0, rabi_value=15.0, detuning_value=15.0
+    ).batch_assign(atom_distance=distances)
+
+    bloqade_emu_target = batch.bloqade.python()
+    braket_emu_target = batch.braket.local_emulator()
+    quera_aquila_target = batch.parallelize(24).quera.aquila()
+    braket_aquila_target = batch.parallelize(24).braket.aquila()
+
+    targets = [
+        bloqade_emu_target,
+        braket_emu_target,
+        quera_aquila_target,
+        braket_aquila_target,
+    ]
+
+    for target in targets:
+        target._compile(10)

tests/test_floquet_compile.py

@@ -0,0 +1,43 @@
+from bloqade import start, cast
+import numpy as np
+
+
+def test_floquet_compile():
+    durations = cast(["ramp_time", "run_time", "ramp_time"])
+
+    def detuning_wf(t, drive_amplitude, drive_frequency):
+        return drive_amplitude * np.sin(drive_frequency * t)
+
+    floquet_program = (
+        start.add_position((0, 0))
+        .rydberg.rabi.amplitude.uniform.piecewise_linear(
+            durations, [0, "rabi_max", "rabi_max", 0]
+        )
+        .detuning.uniform.fn(detuning_wf, sum(durations))
+        .sample("min_time_step", "linear")
+    )
+
+    run_times = np.linspace(0.05, 3.0, 101)
+
+    floquet_job = floquet_program.assign(
+        ramp_time=0.06,
+        min_time_step=0.05,
+        rabi_max=15,
+        drive_amplitude=15,
+        drive_frequency=15,
+    ).batch_assign(run_time=run_times)
+
+    bloqade_emu_target = floquet_job.bloqade.python()
+    braket_emu_target = floquet_job.braket.local_emulator()
+    quera_aquila_target = floquet_job.parallelize(24).quera.aquila()
+    braket_aquila_target = floquet_job.parallelize(24).braket.aquila()
+
+    targets = [
+        bloqade_emu_target,
+        braket_emu_target,
+        quera_aquila_target,
+        braket_aquila_target,
+    ]
+
+    for target in targets:
+        target._compile(10)

tests/test_mis_compile.py

@@ -0,0 +1,32 @@
+from bloqade.atom_arrangement import Square
+import numpy as np
+
+
+def test_mis_compile():
+    rng = np.random.default_rng(1234)
+
+    durations = [0.3, 1.6, 0.3]
+
+    mis_udg_program = (
+        Square(15, lattice_spacing=5.0)
+        .apply_defect_density(0.3, rng=rng)
+        .rydberg.rabi.amplitude.uniform.piecewise_linear(
+            durations, [0.0, 15.0, 15.0, 0.0]
+        )
+        .detuning.uniform.piecewise_linear(
+            durations, [-30, -30, "final_detuning", "final_detuning"]
+        )
+    )
+
+    mis_udg_job = mis_udg_program.batch_assign(final_detuning=np.linspace(0, 80, 41))
+
+    # skip emulation targets considering not feasible to emulate
+    # bloqade_emu_target = mis_udg_job.bloqade.python()
+    # braket_emu_target = mis_udg_job.braket.local_emulator()
+    quera_aquila_target = mis_udg_job.quera.aquila()
+    braket_aquila_target = mis_udg_job.braket.aquila()
+
+    targets = [quera_aquila_target, braket_aquila_target]
+
+    for target in targets:
+        target._compile(10)

tests/test_multi_rabi_compile.py

@@ -0,0 +1,45 @@
+from bloqade import start
+import numpy as np
+
+
+def test_multi_rabi_compile():
+    distance = 4.0
+
+    # Example defaults to option 7
+    geometry = start.add_position(
+        [
+            (0, 0),
+            (distance, 0),
+            (-0.5 * distance, distance),
+            (0.5 * distance, distance),
+            (1.5 * distance, distance),
+            (0, 2 * distance),
+            (distance, 2 * distance),
+        ]
+    )
+
+    sequence = start.rydberg.rabi.amplitude.uniform.piecewise_linear(
+        durations=["ramp_time", "run_time", "ramp_time"],
+        values=[0.0, "rabi_drive", "rabi_drive", 0.0],
+    ).parse_sequence()
+
+    batch = (
+        geometry.apply(sequence)
+        .assign(ramp_time=0.06, rabi_drive=5)
+        .batch_assign(run_time=0.05 * np.arange(21))
+    )
+
+    bloqade_emu_target = batch.bloqade.python()
+    braket_emu_target = batch.braket.local_emulator()
+    quera_aquila_target = batch.parallelize(24).quera.aquila()
+    braket_aquila_target = batch.parallelize(24).braket.aquila()
+
+    targets = [
+        bloqade_emu_target,
+        braket_emu_target,
+        quera_aquila_target,
+        braket_aquila_target,
+    ]
+
+    for target in targets:
+        target._compile(10)

tests/test_noneq_dynamics_compile.py

@@ -0,0 +1,31 @@
+from bloqade.atom_arrangement import Chain
+import numpy as np
+
+
+def test_non_eq_compile():
+    initial_geometry = Chain(2, lattice_spacing="distance")
+    program_waveforms = (
+        initial_geometry.rydberg.rabi.amplitude.uniform.piecewise_linear(
+            durations=["ramp_time", "run_time", "ramp_time"],
+            values=[0.0, "rabi_ampl", "rabi_ampl", 0.0],
+        )
+    )
+    program_assigned_vars = program_waveforms.assign(
+        ramp_time=0.06, rabi_ampl=15, distance=8.5
+    )
+    batch = program_assigned_vars.batch_assign(run_time=0.05 * np.arange(31))
+
+    bloqade_emu_target = batch.bloqade.python()
+    braket_emu_target = batch.braket.local_emulator()
+    quera_aquila_target = batch.parallelize(24).quera.aquila()
+    braket_aquila_target = batch.parallelize(24).braket.aquila()
+
+    targets = [
+        bloqade_emu_target,
+        braket_emu_target,
+        quera_aquila_target,
+        braket_aquila_target,
+    ]
+
+    for target in targets:
+        target._compile(10)

tests/test_ramsey_compile.py

@@ -0,0 +1,39 @@
+from bloqade import start, cast
+from decimal import Decimal
+import numpy as np
+
+
+def test_ramsey_compile():
+    plateau_time = (np.pi / 2 - 0.625) / 12.5
+    wf_durations = cast(
+        [0.05, plateau_time, 0.05, "run_time", 0.05, plateau_time, 0.05]
+    )
+    rabi_wf_values = [0.0, 12.5, 12.5, 0.0] * 2  # repeat values twice
+
+    ramsey_program = (
+        start.add_position((0, 0))
+        .rydberg.rabi.amplitude.uniform.piecewise_linear(wf_durations, rabi_wf_values)
+        .detuning.uniform.constant(10.5, sum(wf_durations))
+    )
+
+    n_steps = 100
+    max_time = Decimal("3.0")
+    dt = (max_time - Decimal("0.05")) / n_steps
+    run_times = [Decimal("0.05") + dt * i for i in range(101)]
+
+    ramsey_job = ramsey_program.batch_assign(run_time=run_times)
+
+    bloqade_emu_target = ramsey_job.bloqade.python()
+    braket_emu_target = ramsey_job.braket.local_emulator()
+    quera_aquila_target = ramsey_job.parallelize(24).quera.aquila()
+    braket_aquila_target = ramsey_job.parallelize(24).braket.aquila()
+
+    targets = [
+        bloqade_emu_target,
+        braket_emu_target,
+        quera_aquila_target,
+        braket_aquila_target,
+    ]
+
+    for target in targets:
+        target._compile(10)

tests/test_scar_compile.py

@@ -0,0 +1,52 @@
+from bloqade import var
+from bloqade.atom_arrangement import Chain
+import numpy as np
+
+
+def test_scar_compile():
+    n_atoms = 11
+    lattice_spacing = 6.1
+    run_time = var("run_time")
+
+    quantum_scar_program = (
+        Chain(n_atoms, lattice_spacing=lattice_spacing)
+        # define detuning waveform
+        .rydberg.detuning.uniform.piecewise_linear(
+            [0.3, 1.6, 0.3], [-18.8, -18.8, 16.3, 16.3]
+        )
+        .piecewise_linear([0.2, 1.6], [16.3, 0.0, 0.0])
+        # slice the detuning waveform
+        .slice(start=0, stop=run_time)
+        # define rabi waveform
+        .amplitude.uniform.piecewise_linear([0.3, 1.6, 0.3], [0.0, 15.7, 15.7, 0.0])
+        .piecewise_linear([0.2, 1.4, 0.2], [0, 15.7, 15.7, 0])
+        # slice waveform, add padding for the linear segment
+        .slice(start=0, stop=run_time - 0.065)
+        # record the value of the waveform at the end of the slice to "rabi_value"
+        .record("rabi_value")
+        # append segment to waveform that fixes the value of the waveform to 0
+        # at the end of the waveform
+        .linear("rabi_value", 0, 0.065)
+    )
+
+    # get run times via the following:
+    prep_times = np.arange(0.2, 2.2, 0.2)
+    scar_times = np.arange(2.2, 4.01, 0.01)
+    run_times = np.unique(np.hstack((prep_times, scar_times)))
+
+    batch = quantum_scar_program.batch_assign(run_time=run_times)
+
+    bloqade_emu_target = batch.bloqade.python()
+    braket_emu_target = batch.braket.local_emulator()
+    quera_aquila_target = batch.parallelize(24).quera.aquila()
+    braket_aquila_target = batch.parallelize(24).braket.aquila()
+
+    targets = [
+        bloqade_emu_target,
+        braket_emu_target,
+        quera_aquila_target,
+        braket_aquila_target,
+    ]
+
+    for target in targets:
+        target._compile(10)

tests/test_single_rabi_compile.py

@@ -0,0 +1,35 @@
+from bloqade import start, cast
+import numpy as np
+
+
+def test_single_rabi_compile():
+    durations = cast(["ramp_time", "run_time", "ramp_time"])
+
+    rabi_oscillations_program = (
+        start.add_position((0, 0))
+        .rydberg.rabi.amplitude.uniform.piecewise_linear(
+            durations=durations, values=[0, "rabi_ampl", "rabi_ampl", 0]
+        )
+        .detuning.uniform.constant(duration=sum(durations), value="detuning_value")
+    )
+
+    run_times = np.linspace(0, 3, 101)
+
+    rabi_oscillation_job = rabi_oscillations_program.assign(
+        ramp_time=0.06, rabi_ampl=15, detuning_value=0.0
+    ).batch_assign(run_time=run_times)
+
+    bloqade_emu_target = rabi_oscillation_job.bloqade.python()
+    braket_emu_target = rabi_oscillation_job.braket.local_emulator()
+    quera_aquila_target = rabi_oscillation_job.parallelize(24).quera.aquila()
+    braket_aquila_target = rabi_oscillation_job.parallelize(24).braket.aquila()
+
+    targets = [
+        bloqade_emu_target,
+        braket_emu_target,
+        quera_aquila_target,
+        braket_aquila_target,
+    ]
+
+    for target in targets:
+        target._compile(10)