diff --git a/src/braket/program_sets/circuit_binding.py b/src/braket/program_sets/circuit_binding.py index 320309f40..79819c68a 100644 --- a/src/braket/program_sets/circuit_binding.py +++ b/src/braket/program_sets/circuit_binding.py @@ -16,7 +16,16 @@ import warnings from collections import defaultdict from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from braket.default_simulator.openqasm.parser.openqasm_ast import ( + Identifier, + IntegerLiteral, + QASMNode, + QuantumMeasurementStatement, + QubitDeclaration, +) +from braket.default_simulator.openqasm.parser.openqasm_parser import parse from braket.ir.openqasm import Program from braket.circuits import Circuit, Gate, Observable @@ -29,10 +38,79 @@ from braket.registers import QubitSet +@dataclass +class _AngleInjectionPlan: + declarations_offset: int + measure_offset: int + qubit_format: str + qubits: QubitSet + + +def _span_offset(node: QASMNode, line_offsets: Sequence[int]) -> int: + return line_offsets[node.span.start_line - 1] + node.span.start_column + + +def _plan_injection(source: str) -> _AngleInjectionPlan: + program = parse(source) + + line_offsets = [0] + offset = 0 + for line in source.split("\n")[:-1]: + offset += len(line) + 1 + line_offsets.append(offset) + + declarations_offset = ( + _span_offset(program.statements[0], line_offsets) if program.statements else len(source) + ) + + measure_offset = len(source) + register_name = None + register_size = 0 + for stmt in program.statements: + if isinstance(stmt, QubitDeclaration) and register_name is None: + # TODO: support multiple qubit registers + register_name = stmt.qubit.name + # stmt.size is None for a single unindexed qubit + register_size = stmt.size.value if isinstance(stmt.size, IntegerLiteral) else 1 + + measure_offset = ( + min(measure_offset, _span_offset(stmt, line_offsets)) + if isinstance(stmt, QuantumMeasurementStatement) + else len(source) + ) + + qubit_format, qubits = ( + (f"{register_name}[{{}}]", QubitSet(range(register_size))) + if register_name + else ( + "${}", + QubitSet( + int(node.name[1:]) + for node in _walk(program) + if isinstance(node, Identifier) and node.name.startswith("$") + ), + ) + ) + return _AngleInjectionPlan( + declarations_offset=declarations_offset, + measure_offset=measure_offset, + qubit_format=qubit_format, + qubits=qubits, + ) + + +def _walk(node: QASMNode): + yield node + for value in vars(node).values(): + for child in value if isinstance(value, list) else [value]: + if isinstance(child, QASMNode): + yield from _walk(child) + + class CircuitBinding: def __init__( self, - circuit: Circuit, + circuit: Circuit | str, input_sets: ParameterSetsLike | None = None, observables: Sequence[Observable | PauliString | str] | Sum | None = None, ): @@ -51,7 +129,8 @@ def __init__( Note: Circuits cannot have result types attached. Args: - circuit (Circuit): The parametrized circuit + circuit (Circuit | str): The parametrized circuit, either as a Circuit object or as + an OpenQASM string. input_sets (ParameterSetsLike | None): The inputs to the circuit, if specified. observables (Sequence[Observable | PauliString | str] | Sum | None): The observables or Hamiltonian to measure, if specified. @@ -70,11 +149,14 @@ def __init__( and any(isinstance(obs, Sum) for obs in observables) ): raise TypeError("Cannot have Sum Hamiltonian in list of observables") - if circuit.result_types: + if isinstance(circuit, Circuit) and circuit.result_types: raise ValueError("Circuit cannot have result types") self._circuit = circuit self._input_sets = ParameterSets(input_sets) if input_sets else None self._observables = CircuitBinding._to_observables(observables) + self._injection_plan = ( + _plan_injection(circuit) if isinstance(circuit, str) and self._observables else None + ) @staticmethod def _to_observables( @@ -96,9 +178,9 @@ def _to_observables( return obs @property - def circuit(self) -> Circuit: + def circuit(self) -> Circuit | str: """ - Circuit: The parametrized circuit + Circuit | str: The parametrized circuit, either as a Circuit object or an OpenQASM string. """ return self._circuit @@ -135,23 +217,60 @@ def to_ir( """ if not self._observables: return Program( - source=self._circuit.to_ir( - IRType.OPENQASM, gate_definitions=gate_definitions - ).source, + source=self._circuit_openqasm(gate_definitions), inputs=self._input_sets.as_dict() if self._input_sets else None, ) - # with_euler_angles validates that the observable has valid Euler angle gates - circuit_with_euler_angles = self._circuit.with_euler_angles(self._observables) euler_angles = self._get_euler_angles() + if isinstance(self._circuit, Circuit): + source = ( + self._circuit + .with_euler_angles(self._observables) + .to_ir(IRType.OPENQASM, gate_definitions=gate_definitions) + .source + ) + else: + source = _inject_euler_angles( + self._circuit, + self._injection_plan, + self._euler_rotation_targets(), + euler_angles.keys(), + ) return Program( - source=circuit_with_euler_angles.to_ir( - IRType.OPENQASM, gate_definitions=gate_definitions - ).source, + source=source, inputs=( self._input_sets * euler_angles if self._input_sets else ParameterSets(euler_angles) ).as_dict(), ) + def _circuit_openqasm( + self, + gate_definitions: Mapping[tuple[Gate, QubitSet], PulseSequence] | None, + ) -> str: + if isinstance(self._circuit, Circuit): + return self._circuit.to_ir(IRType.OPENQASM, gate_definitions=gate_definitions).source + return self._circuit + + def _circuit_qubits(self) -> QubitSet: + if isinstance(self._circuit, Circuit): + return self._circuit.qubits + return self._injection_plan.qubits + + def _euler_rotation_targets(self) -> QubitSet: + observables = self._observables + circuit_qubits = self._circuit_qubits() + if isinstance(observables, Sum): + if observables.targets: + # Sum.targets is a per-summand list of QubitSets + return QubitSet().union(*observables.targets) + return circuit_qubits + targets = QubitSet() + for obs in observables: + if obs.targets: + targets |= obs.targets + else: + targets |= circuit_qubits + return targets + def _get_euler_angles(self) -> dict[str, float] | None: observables = self._observables return ( @@ -164,7 +283,7 @@ def _get_euler_angles_sum(self, observables: Sum) -> dict[str, float]: euler_angles = defaultdict(list) summands = observables.summands if not observables.targets: - targets = self._circuit.qubits + targets = self._circuit_qubits() for obs in summands: for param, angle in obs.get_euler_angles(targets).items(): euler_angles[param].append(angle) @@ -179,7 +298,7 @@ def _get_euler_angles_sum(self, observables: Sum) -> dict[str, float]: def _get_euler_angles_list(self, observables: Sequence[Observable]) -> dict[str, float]: euler_angles = defaultdict(list) - circuit_qubits = self._circuit.qubits + circuit_qubits = self._circuit_qubits() targets = QubitSet(q for obs in observables for q in (obs.targets or circuit_qubits)) for obs in observables: if not obs.targets: @@ -207,12 +326,17 @@ def bind_observables_to_inputs( well as CompositeEntry.expectation. Kwargs: - inplace (bool): whether or not to return a new circuit binding or use the same one - add_measure (bool): whether or not to apply Measure instructions to the circuit + inplace (bool): Whether to return a new circuit binding or use the same one + add_measure (bool): Whether to apply Measure instructions to the circuit. Only + applies when the underlying circuit is a `Circuit`; for OpenQASM string + circuits, the source is preserved verbatim aside from injected Euler-angle + rotations. Returns: CircuitBinding: A new circuit binding with the observables bound. """ + if isinstance(self._circuit, str): + return self._bind_observables_to_inputs_str(inplace) measure = Circuit() parameters = self._input_sets.as_dict() if self._input_sets else None if observables := self._observables: @@ -236,6 +360,33 @@ def bind_observables_to_inputs( return self return CircuitBinding(self._circuit + measure, input_sets=parameters) + def _bind_observables_to_inputs_str(self, inplace: bool) -> CircuitBinding: + source = self._circuit + parameters = self._input_sets.as_dict() if self._input_sets else None + if observables := self._observables: + if isinstance(observables, Sum): + warnings.warn( + "Binding a Sum discards information on observable weights; please " + "distribute your observable in advance using observable.summands.", + stacklevel=2, + ) + euler_angles = self._get_euler_angles() + source = _inject_euler_angles( + source, + self._injection_plan, + self._euler_rotation_targets(), + euler_angles.keys(), + ) + parameters = self._input_sets * euler_angles if parameters else euler_angles + if inplace: + self._circuit = source + # Observables are now bound into the source, so no further injection plan is needed. + self._injection_plan = None + self._observables = None + self._input_sets = parameters + return self + return CircuitBinding(source, input_sets=parameters) + def __len__(self): input_sets = self._input_sets observables = self._observables @@ -260,3 +411,28 @@ def __repr__(self): f"input_sets={self._input_sets}, " f"observables={self._observables})" ) + + +def _inject_euler_angles( + source: str, + plan: _AngleInjectionPlan, + targets: QubitSet, + parameter_names: Sequence[str], +) -> str: + rotations = [] + for q in targets: + theta, phi, omega = euler_angle_parameter_names(q) + formatted = plan.qubit_format.format(int(q)) + rotations.extend([ + f"rz({theta}) {formatted};", + f"rx({phi}) {formatted};", + f"rz({omega}) {formatted};", + ]) + for offset, statements in ( + (plan.measure_offset, rotations), + (plan.declarations_offset, [f"input float {name};" for name in parameter_names]), + ): + block = "\n".join(statements) + prefix = "" if offset == 0 or source[offset - 1] == "\n" else "\n" + source = f"{source[:offset]}{prefix}{block}\n{source[offset:]}" + return source diff --git a/test/unit_tests/braket/program_sets/test_circuit_binding.py b/test/unit_tests/braket/program_sets/test_circuit_binding.py index 9af083fad..22784d0fe 100644 --- a/test/unit_tests/braket/program_sets/test_circuit_binding.py +++ b/test/unit_tests/braket/program_sets/test_circuit_binding.py @@ -15,11 +15,17 @@ from braket.circuits import Circuit from braket.circuits.observables import X, Y, Z +from braket.circuits.serialization import IRType +from braket.default_simulator.openqasm.parser.openqasm_parser import parse from braket.parametric import FreeParameter from braket.program_sets import CircuitBinding from braket.quantum_information import PauliString +def _source(circuit): + return circuit.to_ir(IRType.OPENQASM).source + + def test_equality(circuit_rx_parametrized): input_sets = {"theta": [1.23, 3.21]} observable = X(0) @ Z(1) + 3 * Y(0) @@ -168,3 +174,242 @@ def test_binding_without_measure(circuit_rx_parametrized): circ = cb2.circuit circ.measure(range(2)) assert circ == cb3.circuit + + +def test_string_circuit_no_observables(circuit_rx_parametrized): + src = _source(circuit_rx_parametrized) + cb = CircuitBinding(src, input_sets={"theta": [1.23, 3.21]}) + program = cb.to_ir() + assert program.source == src + assert program.inputs == {"theta": [1.23, 3.21]} + # With no observables to inject, the source is emitted verbatim and never parsed. + assert cb._injection_plan is None + + +def test_string_circuit_matches_circuit_to_ir(circuit_rx_parametrized): + circuit = Circuit(circuit_rx_parametrized).cnot(0, 1) + src = _source(circuit) + observable = [X(0) @ Z(1)] + cb_circ = CircuitBinding(circuit, {"theta": [1.23]}, observable) + cb_str = CircuitBinding(src, {"theta": [1.23]}, observable) + circ_program = cb_circ.to_ir() + str_program = cb_str.to_ir() + assert circ_program.inputs == str_program.inputs + # Same set of statements; declaration ordering may differ between paths. + assert sorted(circ_program.source.splitlines()) == sorted(str_program.source.splitlines()) + + +def test_string_circuit_targetless_observable(circuit_rx_parametrized): + circuit = Circuit(circuit_rx_parametrized).cnot(0, 1) + src = _source(circuit) + cb_circ = CircuitBinding(circuit, observables=[X() @ Y()]) + cb_str = CircuitBinding(src, observables=[X() @ Y()]) + assert cb_str.to_ir().inputs == cb_circ.to_ir().inputs + + +def test_string_circuit_sum_observable(circuit_rx_parametrized): + circuit = Circuit(circuit_rx_parametrized).cnot(0, 1) + src = _source(circuit) + h = 0.5 * X(0) @ Z(1) + 2 * Y(0) + cb_circ = CircuitBinding(circuit, observables=h) + cb_str = CircuitBinding(src, observables=h) + assert cb_str.to_ir().inputs == cb_circ.to_ir().inputs + + +def test_string_circuit_targetless_sum_observable(circuit_rx_parametrized): + circuit = Circuit(circuit_rx_parametrized).cnot(0, 1) + src = _source(circuit) + h_targetless = X() @ Y() - 3 * Z() @ X() + cb_circ = CircuitBinding(circuit, observables=h_targetless) + cb_str = CircuitBinding(src, observables=h_targetless) + assert cb_str.to_ir().inputs == cb_circ.to_ir().inputs + + +def test_string_circuit_bind_sum_warning(circuit_rx_parametrized): + src = _source(circuit_rx_parametrized) + cb = CircuitBinding(src, observables=0.5 * X(0) @ Z(1) + 2 * Y(0)) + with pytest.warns(UserWarning): + cb.bind_observables_to_inputs() + + +def test_string_circuit_custom_register_name(): + src = ( + "OPENQASM 3.0;\n" + "input float theta;\n" + "bit[2] b;\n" + "qubit[2] foo;\n" + "rx(theta) foo[0];\n" + "cnot foo[0], foo[1];\n" + "b[0] = measure foo[0];\n" + "b[1] = measure foo[1];" + ) + cb = CircuitBinding(src, input_sets={"theta": [0.5]}, observables=[X(0) @ Y(1)]) + serialized = cb.to_ir().source + assert "rz(_OBSERVABLE_THETA_0) foo[0];" in serialized + assert "rz(_OBSERVABLE_THETA_1) foo[1];" in serialized + assert "q[0]" not in serialized + assert "q[1]" not in serialized + + +def test_string_circuit_no_measure_or_version_line(): + # Source with no `OPENQASM 3.0;` line and no measurement — exercises the fallback offsets + # of _plan_injection for declarations_offset and measure_offset. + src = "qubit[2] q;\nh q[0];\ncnot q[0], q[1];" + cb = CircuitBinding(src, observables=[X(0) @ Y(1)]) + serialized = cb.to_ir().source + # Declarations should be prepended (no OPENQASM line to anchor on). + assert serialized.startswith("input float ") + # Rotations should land at the end (no measurement to anchor on). + assert serialized.rstrip().endswith("rz(_OBSERVABLE_OMEGA_1) q[1];") + + +def test_string_circuit_single_line_source(): + # A newline-free program must inject the same statements as its multi-line equivalent. + multiline = ( + "OPENQASM 3.0;\n" + "input float theta;\n" + "bit[2] b;\n" + "qubit[2] q;\n" + "rx(theta) q[0];\n" + "cnot q[0], q[1];\n" + "b[0] = measure q[0];\n" + "b[1] = measure q[1];" + ) + single_line = multiline.replace("\n", " ") + observable = [X(0) @ Z(1)] + ml_program = CircuitBinding(multiline, {"theta": [0.5]}, observable).to_ir() + sl_program = CircuitBinding(single_line, {"theta": [0.5]}, observable).to_ir() + assert sl_program.inputs == ml_program.inputs + # The injected statements match; only the original (un-split) lines differ in packing. + injected = "input float _OBSERVABLE" + assert sl_program.source.count(injected) == ml_program.source.count(injected) + # Both forms parse back into the same circuit. + assert Circuit.from_ir(sl_program.source) == Circuit.from_ir(ml_program.source) + + +def test_string_circuit_single_line_reparses(): + # The injected single-line result must be valid OpenQASM that round-trips through from_ir. + src = "OPENQASM 3.0; bit[1] b; qubit[1] q; rx(0.5) q[0]; b[0] = measure q[0];" + serialized = CircuitBinding(src, observables=[X(0)]).to_ir().source + assert "input float _OBSERVABLE_THETA_0;" in serialized + assert "rz(_OBSERVABLE_THETA_0) q[0];" in serialized + # Rotations are injected before the measurement. + assert serialized.index("rz(_OBSERVABLE_THETA_0)") < serialized.index("measure q[0]") + Circuit.from_ir(serialized) + + +def test_string_circuit_midcircuit_measurement(): + # With a mid-circuit measurement, rotations must anchor before the *terminal* measurement + # block, not the first measurement — otherwise the basis change lands ahead of a readout it + # is not meant to affect and no-ops for the terminal readout. + src = ( + "OPENQASM 3.0;\n" + "bit[2] b;\n" + "qubit[2] q;\n" + "h q[0];\n" + "b[0] = measure q[0];\n" # mid-circuit measurement + "if (b[0]) x q[1];\n" + "b[1] = measure q[1];" # terminal measurement + ) + serialized = CircuitBinding(src, observables=[X(0) @ Z(1)]).to_ir().source + # Rotations sit after the mid-circuit measurement + branch, before the terminal measurement. + assert ( + serialized.index("b[0] = measure") + < serialized.index("if (b[0])") + < serialized.index("rz(_OBSERVABLE_THETA_0)") + < serialized.index("b[1] = measure") + ) + # The result is still valid OpenQASM. + parse(serialized) + + +def test_string_circuit_measurement_then_gate_is_not_terminal(): + # A lone measurement followed by a gate is not a terminal-measurement block; rotations + # append at the end (fallback), leaving the earlier measurement untouched. + src = "OPENQASM 3.0;\nqubit[1] q;\nh q[0];\nb[0] = measure q[0];\nx q[0];" + serialized = CircuitBinding(src, observables=[X(0)]).to_ir().source + assert serialized.index("x q[0];") < serialized.index("rz(_OBSERVABLE_THETA_0)") + assert serialized.rstrip().endswith("rz(_OBSERVABLE_OMEGA_0) q[0];") + + +def test_string_circuit_no_instructions(): + # An instruction-free (header-only) program has no statements; injected declarations and + # rotations append after the header, and the result stays valid OpenQASM. + serialized = CircuitBinding("OPENQASM 3.0;", observables=[X(0)]).to_ir().source + assert serialized.startswith("OPENQASM 3.0;") + assert "input float _OBSERVABLE_THETA_0;" in serialized + assert "rz(_OBSERVABLE_THETA_0) $0;" in serialized + parse(serialized) + + +def test_string_circuit_broadcast_gate(): + # A register-broadcast gate (`h q;`) should still yield Euler rotations on every qubit, + # driven by the declared register size rather than explicit `q[i]` references. + src = "OPENQASM 3.0;\nqubit[2] q;\nh q;" + serialized = CircuitBinding(src, observables=[X(0) @ Y(1)]).to_ir().source + assert "rz(_OBSERVABLE_THETA_0) q[0];" in serialized + assert "rz(_OBSERVABLE_THETA_1) q[1];" in serialized + + +def test_string_circuit_single_unindexed_qubit(): + # `qubit q;` (no size) declares a single qubit at index 0. + src = "OPENQASM 3.0;\nqubit q;\nh q;" + serialized = CircuitBinding(src, observables=[X(0)]).to_ir().source + assert "rz(_OBSERVABLE_THETA_0) q[0];" in serialized + assert "q[1]" not in serialized + + +def test_circuit_binding_dunders(circuit_rx_parametrized): + # Exercises the input_sets/observables properties and __len__/__repr__. + input_sets = {"theta": [1.23, 3.21, 0.5]} + observable = [X(0), Y(0)] + cb = CircuitBinding(circuit_rx_parametrized, input_sets, observable) + assert cb.input_sets.as_dict() == input_sets + assert list(cb.observables) == observable + assert len(cb) == 6 + assert len(CircuitBinding(circuit_rx_parametrized, input_sets)) == 3 + assert len(CircuitBinding(circuit_rx_parametrized, observables=observable)) == 2 + assert "CircuitBinding(circuit=" in repr(cb) + + +def test_string_circuit_physical_qubits(circuit_rx_parametrized): + verbatim = Circuit().add_verbatim_box(Circuit().rx(1, FreeParameter("theta")).cnot(1, 0)) + src = _source(verbatim) + cb = CircuitBinding(src, input_sets={"theta": [0.5]}, observables=[X(0) @ Y(1)]) + serialized = cb.to_ir().source + # Inserted Euler rotations target physical qubits, not virtual q[i]. + assert "rz(_OBSERVABLE_THETA_0) $0;" in serialized + assert "rz(_OBSERVABLE_THETA_1) $1;" in serialized + assert "q[" not in serialized.replace("bit[", "") + + +def test_string_circuit_bind_observables_to_inputs(circuit_rx_parametrized): + src = _source(circuit_rx_parametrized) + observable = [X(0) @ Z(1), Y(0), Z(0)] + cb1 = CircuitBinding(src, input_sets={"theta": [1.35, 1.58]}, observables=observable) + + cb2 = cb1.bind_observables_to_inputs(inplace=False) + assert cb1 != cb2 + assert cb1.to_ir() == cb2.to_ir() + + cb1.bind_observables_to_inputs(inplace=True) + assert cb1 == cb2 + # Observables are now folded into the source, so no injection plan remains. + assert cb1._injection_plan is None + + +def test_string_circuit_bind_no_observables(circuit_rx_parametrized): + src = _source(circuit_rx_parametrized) + cb1 = CircuitBinding(src, input_sets={"theta": [1.35, 1.58]}) + cb2 = cb1.bind_observables_to_inputs(inplace=False) + assert cb1 == cb2 + + +def test_string_circuit_equality(): + src = "OPENQASM 3.0;\ninput float theta;\nbit[1] b;\nqubit[1] q;\nrx(theta) q[0];" + input_sets = {"theta": [1.23, 3.21]} + observable = [X(0)] + cb = CircuitBinding(src, input_sets, observable) + assert cb == CircuitBinding(src, input_sets, observable) + assert cb != CircuitBinding(src, observables=observable) + assert cb != CircuitBinding(src, input_sets) diff --git a/tox.ini b/tox.ini index c43008c4b..52f78bde2 100644 --- a/tox.ini +++ b/tox.ini @@ -17,7 +17,7 @@ extras = test deps = {[test-deps]deps} commands = - pytest {posargs} --cov=braket --cov-report=term-missing --cov-report=html --cov-report=xml --cov-append + pytest {posargs} --cov --cov-report=term-missing --cov-report=html --cov-report=xml --cov-append [testenv:integ-tests] extras = test