From 4c72a230b29ddfc9315a12d786c070cc3f15eb26 Mon Sep 17 00:00:00 2001 From: Phillip Weinberg Date: Mon, 25 Sep 2023 10:57:01 -0400 Subject: [PATCH] Fixing more bugs in emulator. (#614) * making detuning matrix elements negative. * mising factor of 2 in rabi amplitude. * add integration test comparing braket to bloqade emulator. * removing extra sorting in KS test impl. * adding check for bounds in locations. * adding visitor to check if sequence has hyperfine coupling. * fixing some bugs in local drives. * making tests shorter, adding scipy KS test * reduce shots counts * reducing test time for emulator. * adding integration tests for emulator ir codegen (#623) * update `LocalBatch` annotation. --- src/bloqade/codegen/common/is_hyperfine.py | 28 +++ src/bloqade/codegen/emulator_ir.py | 43 ++-- src/bloqade/emulate/codegen/hamiltonian.py | 4 +- src/bloqade/emulate/ir/state_vector.py | 2 +- src/bloqade/task/batch.py | 3 +- tests/test_emulator_codegen.py | 274 ++++++++++++++++++++- tests/test_python_emulator.py | 95 +++++-- 7 files changed, 413 insertions(+), 36 deletions(-) create mode 100644 src/bloqade/codegen/common/is_hyperfine.py diff --git a/src/bloqade/codegen/common/is_hyperfine.py b/src/bloqade/codegen/common/is_hyperfine.py new file mode 100644 index 000000000..6243a18aa --- /dev/null +++ b/src/bloqade/codegen/common/is_hyperfine.py @@ -0,0 +1,28 @@ +from typing import Any +import bloqade.ir.analog_circuit as analog_circuit +import bloqade.ir.control.sequence as sequence +from bloqade.ir.visitor.analog_circuit import AnalogCircuitVisitor + + +class IsHyperfineSequence(AnalogCircuitVisitor): + def __init__(self): + self.is_hyperfine = False + + def visit_analog_circuit(self, ast: analog_circuit.AnalogCircuit) -> Any: + self.visit(ast.sequence) + + def visit_append_sequence(self, ast: sequence.Append) -> Any: + list(map(self.visit, ast.sequences)) + + def visit_slice_sequence(self, ast: sequence.Slice) -> Any: + self.visit(ast.sequence) + + def visit_named_sequence(self, ast: sequence.NamedSequence) -> Any: + self.visit(ast.sequence) + + def visit_sequence(self, ast: sequence.Sequence) -> Any: + self.is_hyperfine = self.is_hyperfine or sequence.hyperfine in ast.pulses + + def emit(self, ast: analog_circuit.AnalogCircuit) -> bool: + self.visit(ast) + return self.is_hyperfine diff --git a/src/bloqade/codegen/emulator_ir.py b/src/bloqade/codegen/emulator_ir.py index eaa146b8d..392d67cfd 100644 --- a/src/bloqade/codegen/emulator_ir.py +++ b/src/bloqade/codegen/emulator_ir.py @@ -12,9 +12,9 @@ import bloqade.ir.control.sequence as sequence import bloqade.ir.control.pulse as pulse import bloqade.ir.control.waveform as waveform -import bloqade.ir.scalar as scalar +import bloqade.ir.control.field as field import bloqade.ir as ir - +from bloqade.codegen.common.is_hyperfine import IsHyperfineSequence from bloqade.emulate.ir.atom_type import ThreeLevelAtom, TwoLevelAtom from bloqade.emulate.ir.emulator import ( DetuningOperatorData, @@ -59,20 +59,18 @@ def __init__( self.original_index = [] def visit_analog_circuit(self, ast: ir.AnalogCircuit): - self.n_atoms = ast.register.n_atoms - - self.visit(ast.sequence) self.visit(ast.register) + self.visit(ast.sequence) def visit_register(self, ast: AtomArrangement) -> Any: positions = [] - for original_index, loc_info in enumerate(ast.enumerate()): + for org_index, loc_info in enumerate(ast.enumerate()): if loc_info.filling == SiteFilling.filled: position = tuple([pos(**self.assignments) for pos in loc_info.position]) positions.append(position) - self.original_index.append(original_index) + self.original_index.append(org_index) - if sequence.hyperfine in self.level_couplings: + if self.is_hyperfine: self.register = Register( ThreeLevelAtom, positions, @@ -91,7 +89,6 @@ def visit_sequence(self, ast: sequence.Sequence) -> None: sequence.rydberg: LevelCoupling.Rydberg, } for level_coupling, sub_pulse in ast.pulses.items(): - self.level_couplings.add(level_coupling) self.visit(sub_pulse) self.pulses[level_coupling_mapping[level_coupling]] = Fields( detuning=self.detuning_terms, @@ -144,9 +141,18 @@ def visit_assigned_run_time_vector( def visit_scaled_locations(self, ast: ScaledLocations) -> Dict[int, Decimal]: target_atoms = {} + + for location in ast.value.keys(): + if location.value >= self.n_sites or location.value < 0: + raise ValueError( + f"Location {location.value} is out of bounds for register with " + f"{self.n_sites} sites." + ) + for new_index, original_index in enumerate(self.original_index): - value = ast.value.get(original_index, scalar.Literal(0)) - target_atoms[new_index] = value(**self.assignments) + value = ast.value.get(field.Location(original_index)) + if value is not None: + target_atoms[new_index] = value(**self.assignments) return target_atoms @@ -174,8 +180,9 @@ def visit_detuning(self, ast: Optional[Field]): for atom in range(self.n_atoms): wf = sum( ( - target_atom_dict[sm].get(atom, 0.0) * wf + target_atom_dict[sm][atom] * wf for sm, wf in ast.value.items() + if atom in target_atom_dict[sm] ), start=waveform.Constant(0.0, 0.0), ) @@ -221,8 +228,9 @@ def visit_rabi(self, amplitude: Optional[Field], phase: Optional[Field]): for atom in range(self.n_atoms): amplitude_wf = sum( ( - amplitude_target_atoms_dict[sm].get(atom, 0.0) * wf + amplitude_target_atoms_dict[sm][atom] * wf for sm, wf in amplitude.value.items() + if atom in amplitude_target_atoms_dict[sm] ), start=waveform.Constant(0.0, 0.0), ) @@ -274,16 +282,18 @@ def visit_rabi(self, amplitude: Optional[Field], phase: Optional[Field]): for atom in range(self.n_atoms): phase_wf = sum( ( - phase_target_atoms_dict[sm].get(atom, 0.0) * wf + phase_target_atoms_dict[sm][atom] * wf for sm, wf in phase.value.items() + if atom in phase_target_atoms_dict[sm] ), start=waveform.Constant(0.0, 0.0), ) amplitude_wf = sum( ( - amplitude_target_atoms_dict[sm].get(atom, 0.0) * wf + amplitude_target_atoms_dict[sm][atom] * wf for sm, wf in amplitude.value.items() + if atom in amplitude_target_atoms_dict[sm] ), start=waveform.Constant(0.0, 0.0), ) @@ -310,6 +320,9 @@ def visit_rabi(self, amplitude: Optional[Field], phase: Optional[Field]): def emit(self, circuit: ir.AnalogCircuit) -> EmulatorProgram: self.assignments = AssignmentScan(self.assignments).emit(circuit.sequence) + self.is_hyperfine = IsHyperfineSequence().emit(circuit) + self.n_atoms = circuit.register.n_atoms + self.n_sites = circuit.register.n_sites self.visit(circuit) return EmulatorProgram( diff --git a/src/bloqade/emulate/codegen/hamiltonian.py b/src/bloqade/emulate/codegen/hamiltonian.py index 0ea27fcde..ee52dbeff 100644 --- a/src/bloqade/emulate/codegen/hamiltonian.py +++ b/src/bloqade/emulate/codegen/hamiltonian.py @@ -98,7 +98,7 @@ def visit_fields(self, fields: Fields): def visit_detuning_operator_data(self, detuning_data: DetuningOperatorData): if (self.register, detuning_data) in self.compile_cache.operator_cache: - return self.compile_cache.operator_cache[(self.space, detuning_data)] + return self.compile_cache.operator_cache[(self.register, detuning_data)] diagonal = np.zeros(self.space.size, dtype=np.float64) if self.space.atom_type == TwoLevelAtomType(): @@ -110,7 +110,7 @@ def visit_detuning_operator_data(self, detuning_data: DetuningOperatorData): state = ThreeLevelAtomType.State.Hyperfine for atom_index, value in detuning_data.target_atoms.items(): - diagonal[self.space.is_state_at(atom_index, state)] += float(value) + diagonal[self.space.is_state_at(atom_index, state)] -= float(value) self.compile_cache.operator_cache[(self.register, detuning_data)] = diagonal return diagonal diff --git a/src/bloqade/emulate/ir/state_vector.py b/src/bloqade/emulate/ir/state_vector.py index 3fdc1e93f..a5dd400e5 100644 --- a/src/bloqade/emulate/ir/state_vector.py +++ b/src/bloqade/emulate/ir/state_vector.py @@ -34,7 +34,7 @@ class RabiOperator: phase: Optional[Callable[[float], float]] = None def dot(self, register: NDArray, time: float): - amplitude = self.amplitude(time) + amplitude = self.amplitude(time) / 2 if self.phase is None: return self.op.dot(register) * amplitude diff --git a/src/bloqade/task/batch.py b/src/bloqade/task/batch.py index 4307b3d0c..d28b38f47 100644 --- a/src/bloqade/task/batch.py +++ b/src/bloqade/task/batch.py @@ -3,6 +3,7 @@ from bloqade.task.quera import QuEraTask from bloqade.task.braket import BraketTask from bloqade.task.braket_simulator import BraketEmulatorTask +from bloqade.task.bloqade import BloqadeTask from bloqade.builder.base import Builder @@ -45,7 +46,7 @@ def json(self, **options) -> str: @Serializer.register class LocalBatch(Serializable): source: Optional[Builder] - tasks: OrderedDict[int, BraketEmulatorTask] + tasks: OrderedDict[int, Union[BraketEmulatorTask, BloqadeTask]] name: Optional[str] = None def report(self) -> Report: diff --git a/tests/test_emulator_codegen.py b/tests/test_emulator_codegen.py index ef57d1b21..5c53d479a 100644 --- a/tests/test_emulator_codegen.py +++ b/tests/test_emulator_codegen.py @@ -17,11 +17,13 @@ from bloqade import start from bloqade.ir.control.sequence import rydberg from bloqade.ir.control.pulse import detuning, rabi -from bloqade.ir.control.field import Uniform +from bloqade.ir.control.field import Location, RunTimeVector, ScaledLocations, Uniform from decimal import Decimal +from bloqade.ir.scalar import cast -def test_codegen_detuning(): + +def test_codegen_global_detuning(): program = ( start.add_position((0, 0)) .add_position((0, 5)) @@ -60,7 +62,7 @@ def test_codegen_detuning(): assert emulator_ir == expected_emulator_ir -def test_codegen_detuning_and_rabi(): +def test_codegen_global_detuning_and_rabi(): program = ( start.add_position((0, 0)) .add_position((0, 5)) @@ -111,3 +113,269 @@ def test_codegen_detuning_and_rabi(): print(expected_emulator_ir) assert emulator_ir == expected_emulator_ir + + +def test_codegen_detuning_max_terms(): + program = ( + start.add_position((0, 0)) + .add_position((0, 5)) + .rydberg.detuning.uniform.piecewise_linear([0.1, 0.8, 0.1], [-10, -10, 10, 10]) + .var("mask_1") + .piecewise_linear([0.1, 0.8, 0.1], [-20, -20, 10, 10]) + .location(0) + .scale(0.4) + .piecewise_linear([0.1, 0.8, 0.1], [-10, -10, 20, 20]) + .parse_circuit() + ) + mask_1_value = [0.5, 2.0] + + mask_1 = RunTimeVector("mask_1") + location_0 = ScaledLocations({Location(0): cast(0.4)}) + uniform_wf = program.sequence.pulses[rydberg].fields[detuning].value[Uniform] + mask_1_wf = program.sequence.pulses[rydberg].fields[detuning].value[mask_1] + location_0_wf = program.sequence.pulses[rydberg].fields[detuning].value[location_0] + + wf_0 = (1.0 * uniform_wf + mask_1_value[0] * mask_1_wf) + 0.4 * location_0_wf + wf_1 = 1.0 * uniform_wf + mask_1_value[1] * mask_1_wf + + assignments = {"mask_1": mask_1_value} + + compiled_wf_0 = CompiledWaveform(assignments, wf_0) + compiled_wf_1 = CompiledWaveform(assignments, wf_1) + + geometry = Register( + TwoLevelAtom, + [(Decimal("0"), Decimal("0")), (Decimal("0"), Decimal("5"))], + Decimal("0"), + ) + + detuning_term_0 = DetuningTerm( + DetuningOperatorData({0: Decimal("1")}), + compiled_wf_0, + ) + + detuning_term_1 = DetuningTerm( + DetuningOperatorData({1: Decimal("1")}), + compiled_wf_1, + ) + + rydberg_drive = Fields(detuning=[detuning_term_0, detuning_term_1], rabi=[]) + drives = {LevelCoupling.Rydberg: rydberg_drive} + + expected_emulator_ir = EmulatorProgram( + register=geometry, duration=1.0, pulses=drives + ) + + emulator_ir = EmulatorProgramCodeGen(assignments=assignments).emit(program) + + print(emulator_ir) + print(expected_emulator_ir) + + assert emulator_ir == expected_emulator_ir + + +def test_codegen_rabi_max_terms(): + program = ( + start.add_position((0, 0)) + .add_position((0, 5)) + .rydberg.rabi.amplitude.uniform.piecewise_linear( + [0.1, 0.8, 0.1], [-10, -10, 10, 10] + ) + .var("mask_1") + .piecewise_linear([0.1, 0.8, 0.1], [-20, -20, 10, 10]) + .location(0) + .scale(0.4) + .piecewise_linear([0.1, 0.8, 0.1], [-10, -10, 20, 20]) + .parse_circuit() + ) + mask_1_value = [0.5, 2.0] + + mask_1 = RunTimeVector("mask_1") + location_0 = ScaledLocations({Location(0): cast(0.4)}) + uniform_wf = program.sequence.pulses[rydberg].fields[rabi.amplitude].value[Uniform] + mask_1_wf = program.sequence.pulses[rydberg].fields[rabi.amplitude].value[mask_1] + location_0_wf = ( + program.sequence.pulses[rydberg].fields[rabi.amplitude].value[location_0] + ) + + wf_0 = (1.0 * uniform_wf + mask_1_value[0] * mask_1_wf) + 0.4 * location_0_wf + wf_1 = 1.0 * uniform_wf + mask_1_value[1] * mask_1_wf + + assignments = {"mask_1": mask_1_value} + + compiled_wf_0 = CompiledWaveform(assignments, wf_0) + compiled_wf_1 = CompiledWaveform(assignments, wf_1) + + geometry = Register( + TwoLevelAtom, + [(Decimal("0"), Decimal("0")), (Decimal("0"), Decimal("5"))], + Decimal("0"), + ) + + rabi_term_0 = RabiTerm( + RabiOperatorData( + RabiOperatorType.RabiSymmetric, + {0: Decimal("1")}, + ), + compiled_wf_0, + None, + ) + + rabi_term_1 = RabiTerm( + RabiOperatorData( + RabiOperatorType.RabiSymmetric, + {1: Decimal("1")}, + ), + compiled_wf_1, + None, + ) + + rydberg_drive = Fields(detuning=[], rabi=[rabi_term_0, rabi_term_1]) + + drives = {LevelCoupling.Rydberg: rydberg_drive} + + expected_emulator_ir = EmulatorProgram( + register=geometry, duration=1.0, pulses=drives + ) + + emulator_ir = EmulatorProgramCodeGen(assignments=assignments).emit(program) + + print(emulator_ir) + print(expected_emulator_ir) + + assert emulator_ir == expected_emulator_ir + + +def test_codegen_rabi_uniform_phase(): + program = ( + start.add_position((0, 0)) + .add_position((0, 5)) + .rydberg.rabi.amplitude.uniform.piecewise_linear( + [0.1, 0.8, 0.1], [-10, -10, 10, 10] + ) + .rydberg.rabi.phase.uniform.piecewise_linear( + [0.1, 0.8, 0.1], [-10, -10, 10, 10] + ) + .parse_circuit() + ) + + uniform_amp = program.sequence.pulses[rydberg].fields[rabi.amplitude].value[Uniform] + uniform_phase = program.sequence.pulses[rydberg].fields[rabi.phase].value[Uniform] + + assignments = {} + + compiled_amp = CompiledWaveform(assignments, uniform_amp) + compiled_phase = CompiledWaveform(assignments, uniform_phase) + + geometry = Register( + TwoLevelAtom, + [(Decimal("0"), Decimal("0")), (Decimal("0"), Decimal("5"))], + Decimal("0"), + ) + + rabi_term = RabiTerm( + RabiOperatorData( + RabiOperatorType.RabiAsymmetric, + {i: Decimal("1") for i in range(len(geometry))}, + ), + compiled_amp, + compiled_phase, + ) + + rydberg_drive = Fields(detuning=[], rabi=[rabi_term]) + + drives = {LevelCoupling.Rydberg: rydberg_drive} + + expected_emulator_ir = EmulatorProgram( + register=geometry, duration=1.0, pulses=drives + ) + + emulator_ir = EmulatorProgramCodeGen(assignments=assignments).emit(program) + + print(emulator_ir) + print(expected_emulator_ir) + + assert emulator_ir == expected_emulator_ir + + +def test_codegen_uniform_phase_rabi_max_terms(): + program = ( + start.add_position((0, 0)) + .add_position((0, 5)) + .rydberg.rabi.amplitude.uniform.piecewise_linear( + [0.1, 0.8, 0.1], [-10, -10, 10, 10] + ) + .var("mask_1") + .piecewise_linear([0.1, 0.8, 0.1], [-20, -20, 10, 10]) + .location(0) + .scale(0.4) + .piecewise_linear([0.1, 0.8, 0.1], [-10, -10, 20, 20]) + .phase.uniform.piecewise_linear([0.1, 0.8, 0.1], [-10, -10, 10, 10]) + .parse_circuit() + ) + + mask_1_value = [0.5, 2.0] + mask_1 = RunTimeVector("mask_1") + location_0 = ScaledLocations({Location(0): cast(0.4)}) + + uniform_amp = program.sequence.pulses[rydberg].fields[rabi.amplitude].value[Uniform] + mask_1_amp = program.sequence.pulses[rydberg].fields[rabi.amplitude].value[mask_1] + location_0_amp = ( + program.sequence.pulses[rydberg].fields[rabi.amplitude].value[location_0] + ) + phase_wf = program.sequence.pulses[rydberg].fields[rabi.phase].value[Uniform] + + wf_0 = (1.0 * uniform_amp + mask_1_value[0] * mask_1_amp) + 0.4 * location_0_amp + wf_1 = 1.0 * uniform_amp + mask_1_value[1] * mask_1_amp + wf_0_phase = 1.0 * phase_wf + wf_1_phase = 1.0 * phase_wf + + assignments = {"mask_1": mask_1_value} + + compiled_wf_0 = CompiledWaveform(assignments, wf_0) + compiled_wf_1 = CompiledWaveform(assignments, wf_1) + compiled_wf_0_phase = CompiledWaveform(assignments, wf_0_phase) + compiled_wf_1_phase = CompiledWaveform(assignments, wf_1_phase) + + geometry = Register( + TwoLevelAtom, + [(Decimal("0"), Decimal("0")), (Decimal("0"), Decimal("5"))], + Decimal("0"), + ) + + rabi_term_0 = RabiTerm( + RabiOperatorData( + RabiOperatorType.RabiAsymmetric, + {0: Decimal("1")}, + ), + compiled_wf_0, + compiled_wf_0_phase, + ) + + rabi_term_1 = RabiTerm( + RabiOperatorData( + RabiOperatorType.RabiAsymmetric, + {1: Decimal("1")}, + ), + compiled_wf_1, + compiled_wf_1_phase, + ) + + rydberg_drive = Fields(detuning=[], rabi=[rabi_term_0, rabi_term_1]) + + drives = {LevelCoupling.Rydberg: rydberg_drive} + + expected_emulator_ir = EmulatorProgram( + register=geometry, duration=1.0, pulses=drives + ) + + emulator_ir = EmulatorProgramCodeGen(assignments=assignments).emit(program) + + print(emulator_ir) + print(expected_emulator_ir) + + assert emulator_ir == expected_emulator_ir + + +if __name__ == "__main__": + test_codegen_uniform_phase_rabi_max_terms() diff --git a/tests/test_python_emulator.py b/tests/test_python_emulator.py index 3f01cb5ac..d0c8b3568 100644 --- a/tests/test_python_emulator.py +++ b/tests/test_python_emulator.py @@ -1,6 +1,9 @@ -from bloqade import start, var +from bloqade import start, var, cast +from bloqade.atom_arrangement import Chain from bloqade.serialize import dumps, loads import numpy as np +from beartype.typing import Dict +from scipy.stats import ks_2samp def test_integration_1(): @@ -12,8 +15,7 @@ def test_integration_1(): [0.1, "ramp_time", 0.1], [-100, -100, 100, 100] ) .amplitude.uniform.piecewise_linear([0.1, "ramp_time", 0.1], [0, 10, 10, 0]) - .assign(ramp_time=3.0) - .batch_assign(r=np.linspace(4, 10, 11).tolist()) + .assign(ramp_time=3.0, r=8) .bloqade.python() .run(10000, cache_matrices=True, blockade_radius=6.0, interaction_picture=True) .report() @@ -34,8 +36,7 @@ def test_integration_2(): .phase.uniform.piecewise_constant( [0.1, ramp_time / 2, ramp_time / 2, 0.1], [0, 0, np.pi, np.pi] ) - .assign(ramp_time=3.0) - .batch_assign(r=np.linspace(4, 10, 11).tolist()) + .assign(ramp_time=3.0, r=6) .bloqade.python() .run(10000, cache_matrices=True, blockade_radius=6.0) .report() @@ -58,8 +59,7 @@ def test_integration_3(): ) .amplitude.var("rabi_mask") .fn(lambda t: 4 * np.sin(3 * t), ramp_time + 0.2) - .assign(ramp_time=3.0, rabi_mask=[10.0, 0.1]) - .batch_assign(r=np.linspace(4, 10, 11).tolist()) + .assign(ramp_time=3.0, rabi_mask=[10.0, 0.1], r=6) .bloqade.python() .run(10000, cache_matrices=True, blockade_radius=6.0) .report() @@ -81,8 +81,7 @@ def test_integration_4(): .fn(lambda t: 4 * np.sin(3 * t), ramp_time + 0.2) .amplitude.location(1) .linear(0.0, 1.0, ramp_time + 0.2) - .assign(ramp_time=3.0, rabi_mask=[10.0, 0.1]) - .batch_assign(r=np.linspace(4, 10, 11).tolist()) + .assign(ramp_time=3.0, rabi_mask=[10.0, 0.1], r=6) .bloqade.python() .run(10000, cache_matrices=True, blockade_radius=6.0) .report() @@ -102,8 +101,7 @@ def test_integration_5(): .amplitude.uniform.piecewise_linear([0.1, ramp_time, 0.1], [0, 10, 10, 0]) .phase.location(1) .linear(0.0, 1.0, ramp_time + 0.2) - .assign(ramp_time=3.0) - .batch_assign(r=np.linspace(4, 10, 11).tolist()) + .assign(ramp_time=3.0, r=6) .bloqade.python() .run(10000, cache_matrices=True, blockade_radius=6.0) .report() @@ -127,8 +125,7 @@ def test_integration_6(): .amplitude.uniform.piecewise_linear([0.1, ramp_time, 0.1], [0, 10, 10, 0]) .phase.location(1) .linear(0.0, 1.0, ramp_time + 0.2) - .assign(ramp_time=3.0) - .batch_assign(r=np.linspace(4, 10, 11).tolist()) + .assign(ramp_time=3.0, r=6) .bloqade.python() .run(10000, cache_matrices=True, blockade_radius=6.0) .report() @@ -151,7 +148,7 @@ def test_serialization(): .amplitude.location(1) .linear(0.0, 1.0, ramp_time + 0.2) .assign(ramp_time=3.0, rabi_mask=[10.0, 0.1]) - .batch_assign(r=np.linspace(4, 10, 11).tolist()) + .batch_assign(r=np.linspace(0.1, 4, 5).tolist()) .bloqade.python() ._compile(100) ) @@ -159,3 +156,73 @@ def test_serialization(): obj_str = dumps(batch) batch2 = loads(obj_str) assert isinstance(batch2, type(batch)) + + +def KS_test( + lhs_counts: Dict[str, int], rhs_counts: Dict[str, int], alpha: float = 0.05 +) -> None: + lhs_samples = [] + rhs_samples = [] + + for bitstring, count in lhs_counts.items(): + lhs_samples += [int(bitstring, 2)] * count + + for bitstring, count in rhs_counts.items(): + rhs_samples += [int(bitstring, 2)] * count + + result = ks_2samp(lhs_samples, rhs_samples, method="exact") + + assert result.pvalue > alpha + + +def test_bloqade_against_braket(): + np.random.seed(9123892) + durations = cast([0.1, 0.1, 0.1]) + + prog = ( + Chain(3, lattice_spacing=6.1) + .rydberg.detuning.uniform.piecewise_linear(durations, [-20, -20, "d", "d"]) + .amplitude.uniform.piecewise_linear(durations, [0, 15, 15, 0]) + .phase.uniform.constant(0.3, sum(durations)) + .batch_assign(d=[10, 20]) + ) + + nshots = 1000 + a = prog.bloqade.python().run(nshots, cache_matrices=True).report().counts + b = prog.braket.local_emulator().run(nshots).report().counts + + for lhs, rhs in zip(a, b): + KS_test(lhs, rhs) + + +def test_bloqade_against_braket_2(): + np.random.seed(192839812) + durations = cast([0.1, 0.1, 0.1]) + values = [0, 15, 15, 0] + + prog_1 = ( + Chain(3, lattice_spacing=6.1) + .rydberg.detuning.uniform.piecewise_linear(durations, [-20, -20, "d", "d"]) + .amplitude.uniform.piecewise_linear(durations, values) + .batch_assign(d=[10, 20]) + ) + prog_2 = ( + Chain(3, lattice_spacing=6.1) + .rydberg.detuning.uniform.piecewise_linear(durations, [-20, -20, "d", "d"]) + .amplitude.location(0) + .piecewise_linear(durations, values) + .amplitude.location(1) + .piecewise_linear(durations, values) + .amplitude.location(2) + .piecewise_linear(durations, values) + .phase.location(0) + .constant(0.0, sum(durations)) + .batch_assign(d=[10, 20]) + ) + + nshots = 1000 + a = prog_2.bloqade.python().run(nshots, cache_matrices=True).report().counts + b = prog_1.braket.local_emulator().run(nshots).report().counts + + for lhs, rhs in zip(a, b): + KS_test(lhs, rhs)