diff --git a/src/braket/circuits/circuit.py b/src/braket/circuits/circuit.py index b656763c3..13b6fa783 100644 --- a/src/braket/circuits/circuit.py +++ b/src/braket/circuits/circuit.py @@ -1626,6 +1626,71 @@ def _add_fixed_argument_calibrations( }) return additional_calibrations + def count( + self, + operator: str | None = None, + qubits: QubitSet | None = None, + include_noise: bool = False, + ) -> int | Counter[str]: + """Counts circuit instructions by operator name, optionally filtered by specific + operators and/or qubits. + + Args: + operator (str | None): Optional operator name to count. Matching is case-insensitive. + If not provided, returns counts for all operator names. + qubits (QubitSet | None): Optional set of qubits to consider. If not provided, + considers all qubits. + include_noise (bool): Whether to include noise instructions in the count. + Default is False. + + Returns: + int | Counter[str]: The count for ``operator`` if provided, otherwise a ``Counter`` + keyed by operator name. + + Raises: + ValueError: If any qubits in ``qubits`` are not present in the circuit + + Examples: + >>> circuit = Circuit().h(0).h(1).cnot(0, 1).amplitude_damping(0, gamma=0.1) + >>> circuit.count("cnot") + 1 + >>> circuit.count() + Counter({'cnot': 1, 'h': 2}) + >>> circuit.count(qubits={0}) + Counter({'cnot': 1, 'h': 1}) + >>> circuit.count(qubits={0}, operator="h") + 1 + >>> circuit.count(include_noise=True) + Counter({'cnot': 1, 'h': 2, 'amplitudedamping': 1}) + """ + counts: dict[str, int] = {} + + # to avoid erroring on qubits that are part of the circuit but not targeted + # by any instruction we construct the following set as opposed to just using + # :attr:`self.qubits` directly + circuit_qubits = set( + set(range(min(self.qubits), max(self.qubits) + 1)) if self.qubits else set() + ) + + if qubits and qubits.intersection(circuit_qubits) != qubits: + raise ValueError( + "All qubits in the 'qubits' argument must be present in the circuit. " + f"Invalid qubits: {qubits - circuit_qubits}" + ) + + for instruction in self.instructions: + if qubits and not set(instruction.target).intersection(qubits): + continue + if not include_noise and isinstance(instruction.operator, Noise): + continue + name = instruction.operator.name.lower() + counts[name] = counts.get(name, 0) + 1 + + if operator: + return counts.get(operator.lower(), 0) + + return Counter(counts) + def to_unitary(self) -> np.ndarray: """Returns the unitary matrix representation of the entire circuit. diff --git a/test/unit_tests/braket/circuits/test_circuit.py b/test/unit_tests/braket/circuits/test_circuit.py index aca5adc04..bc02fd575 100644 --- a/test/unit_tests/braket/circuits/test_circuit.py +++ b/test/unit_tests/braket/circuits/test_circuit.py @@ -3831,3 +3831,30 @@ def test_barrier_jaqcd_export_fails(): pytest.raises(NotImplementedError, match="Barrier is not supported in JAQCD"), ): circ.to_ir(IRType.JAQCD) + + +def test_circuit_count_ops(): + circ = Circuit().h(0).h(1).cnot(0, 1).measure([0, 1]).gphase(0.5) + assert circ.count("cnot") == 1 + assert circ.count("h") == 2 + assert circ.count("unknown") == 0 + assert circ.count() == {"cnot": 1, "h": 2, "measure": 2, "gphase": 1} + assert circ.count(qubits={0, 1}) == {"cnot": 1, "h": 2, "measure": 2} + assert circ.count(qubits={0}) == {"cnot": 1, "h": 1, "measure": 1} + assert circ.count(qubits={0}, operator="h") == 1 + + circ = Circuit().cnot(0, 2) + assert circ.count(qubits={1}) == {} + + circ = Circuit().amplitude_damping(0, gamma=0.1) + + assert circ.count() == {} + assert circ.count(include_noise=True) == {"amplitudedamping": 1} + assert circ.count(qubits={0}, include_noise=True) == {"amplitudedamping": 1} + assert circ.count(operator="amplitudedamping", include_noise=True) == 1 + + with pytest.raises( + ValueError, + match="All qubits in the 'qubits' argument must be present in the circuit", + ): + circ.count(qubits={3})