From fc88fce76965f3ee77063a4eda45032712c38cd1 Mon Sep 17 00:00:00 2001
From: Codrut <codrut@google.com>
Date: Sun, 16 Mar 2025 11:54:27 +0100
Subject: [PATCH 1/4] Optimize the HHL algorithm using amplitude amplification.

I changed the example to simulate two circuits: the first one is the
previous HHL circuit in the example without amplitude amplification,
and the second one is the HHL circuit with amplitude amplification.

The example demonstrates that amplitude amplification obtains an estimate
faster (using fewer total calls to simulate.run).
---
 examples/hhl.py | 297 +++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 231 insertions(+), 66 deletions(-)

diff --git a/examples/hhl.py b/examples/hhl.py
index c8f952d79ea..89937a7578e 100644
--- a/examples/hhl.py
+++ b/examples/hhl.py
@@ -49,6 +49,12 @@
 Coles, Eidenbenz et al. Quantum Algorithm Implementations for Beginners
 https://arxiv.org/abs/1804.03719
 
+Brassard, Gilles et al. Quantum Amplitude Amplification and Estimation
+https://arxiv.org/abs/quant-ph/0005055
+
+Kaye, Phillip et al. An Introduction to Quantum Computing,
+Section 8.4: Searching without knowing the success probability
+
 === CIRCUIT ===
 Example of circuit with 2 register qubits.
 
@@ -65,6 +71,7 @@
 """
 
 import math
+import random
 import numpy as np
 import sympy
 import cirq
@@ -194,73 +201,219 @@ def _ancilla_rotation(self, k):
         return cirq.ry(theta)
 
 
-def hhl_circuit(A, C, t, register_size, *input_prep_gates):
-    """Constructs the HHL circuit.
+class HHLAlgorithm:
+    """A class for the HHL algorithm.
 
-    Args:
-        A: The input Hermitian matrix.
-        C: Algorithm parameter, see above.
-        t: Algorithm parameter, see above.
-        register_size: The size of the eigenvalue register.
-        *input_prep_gates: A list of gates to be applied to |0> to generate the desired input
-            state |b>.
-
-    Returns:
-        The HHL circuit. The ancilla measurement has key 'a' and the memory measurement is in key
-        'm'.  There are two parameters in the circuit, `exponent` and `phase_exponent` corresponding
-        to a possible rotation  applied before the measurement on the memory with a
-        `cirq.PhasedXPowGate`.
+    It provides two versions of the algorithm: one with amplitude amplification,
+    and one without.
     """
 
-    ancilla = cirq.LineQubit(0)
-    # to store eigenvalues of the matrix
-    register = [cirq.LineQubit(i + 1) for i in range(register_size)]
-    # to store input and output vectors
-    memory = cirq.LineQubit(register_size + 1)
-
-    c = cirq.Circuit()
-    hs = HamiltonianSimulation(A, t)
-    pe = PhaseEstimation(register_size + 1, hs)
-    c.append([gate(memory) for gate in input_prep_gates])
-    c.append(
-        [
-            pe(*(register + [memory])),
-            EigenRotation(register_size + 1, C, t)(*(register + [ancilla])),
-            pe(*(register + [memory])) ** -1,
-            cirq.measure(ancilla, key='a'),
-        ]
-    )
+    def __init__(self, A, C, t, register_size, *input_prep_gates):
+        """Initializes the HHL algorithm.
+
+        Args:
+            A: The input Hermitian matrix.
+            C: Algorithm parameter, see above.
+            t: Algorithm parameter, see above.
+            register_size: The size of the eigenvalue register.
+            *input_prep_gates: A list of gates to be applied to |0> to generate the desired input
+                state |b>.
+        """
+        self.A = A
+        self.C = C
+        self.t = t
+        self.register_size = register_size
+        self.input_prep_gates = input_prep_gates
+
+        # ancilla qubit
+        self.ancilla = cirq.LineQubit(0)
+        # to store eigenvalues of the matrix
+        self.register = [cirq.LineQubit(i + 1) for i in range(register_size)]
+        # to store input and output vectors
+        self.memory = cirq.LineQubit(register_size + 1)
+
+    def hhl_circuit(self):
+        """Constructs the HHL circuit.
+
+        Returns:
+            The HHL circuit. The ancilla measurement has key 'a' and the memory measurement
+            is in key 'm'. If the ancilla is measured as 1, the memory is in the desired
+            state |x>.
+        """
+        c = cirq.Circuit()
+        hs = HamiltonianSimulation(self.A, self.t)
+        pe = PhaseEstimation(self.register_size + 1, hs)
+        c.append([gate(self.memory) for gate in self.input_prep_gates])
+        c.append(
+            [
+                pe(*(self.register + [self.memory])),
+                EigenRotation(self.register_size + 1, self.C, self.t)(
+                    *(self.register + [self.ancilla])
+                ),
+                pe(*(self.register + [self.memory])) ** -1,
+            ]
+        )
+        return c
+
+    def diffusion_operator(self):
+        """Creates the diffusion operator I - 2|0><0| which reflects about the initial state
+        with all qubits in |0>.
+        """
+        qubits = [self.ancilla] + self.register + [self.memory]
+        c = cirq.Circuit()
+        c.append(cirq.X.on_each(qubits))
+        c.append(cirq.ControlledGate(cirq.Z, num_controls=self.register_size + 1).on(*qubits))
+        c.append(cirq.X.on_each(qubits))
+        return c
+
+    def amplitude_amplification(self, num_iterations):
+        """Applies amplitude amplification to increase the probability of measuring the ancilla
+        qubit as 1.
+
+        Args:
+            num_iterations: The number of times to apply amplitude amplification.
+
+        Returns:
+            The circuit with amplitude amplification applied.
+        """
+        hhl_circuit = self.hhl_circuit()
+        c = cirq.Circuit()
+        c += hhl_circuit
+
+        r_succ = cirq.Z(self.ancilla)
+        r_init = self.diffusion_operator()
+
+        for _ in range(num_iterations):
+            c.append(r_succ)
+            c += hhl_circuit**-1
+            c += r_init
+            c += hhl_circuit
+
+        return c
+
+    def measure_circuit(self, circuit):
+        """Measures the ancilla qubit and the memory qubit.
+
+        There are two parameters in the circuit, `exponent` and `phase_exponent` corresponding
+        to a possible rotation applied before the measurement on the memory with a
+        `cirq.PhasedXPowGate`.
 
-    c.append(
-        [
-            cirq.PhasedXPowGate(
-                exponent=sympy.Symbol('exponent'), phase_exponent=sympy.Symbol('phase_exponent')
-            )(memory),
-            cirq.measure(memory, key='m'),
+        Args:
+            circuit: The circuit to measure.
+
+        Returns:
+            The circuit with the ancilla and memory qubits measured.
+        """
+        circuit.append(cirq.measure(self.ancilla, key='a'))
+        circuit.append(
+            [
+                cirq.PhasedXPowGate(
+                    exponent=sympy.Symbol('exponent'), phase_exponent=sympy.Symbol('phase_exponent')
+                )(self.memory),
+                cirq.measure(self.memory, key='m'),
+            ]
+        )
+
+        return circuit
+
+    def simulate_without_amplification(self, total_runs):
+        """Simulates the HHL circuit without amplitude amplification.
+
+        Measures X, Y and Z on the memory qubit.
+        Only the runs where the ancilla is measured as 1 are considered.
+
+        Args:
+            total_runs: The number of simulation runs.
+
+        Returns:
+            The success probability of measuring the ancilla qubit as 1.
+        """
+        simulator = cirq.Simulator()
+
+        # Cases for measuring X, Y, and Z (respectively) on the memory qubit.
+        params = [
+            {'exponent': 0.5, 'phase_exponent': -0.5},
+            {'exponent': 0.5, 'phase_exponent': 0},
+            {'exponent': 0, 'phase_exponent': 0},
         ]
-    )
-
-    return c
-
 
-def simulate(circuit):
-    simulator = cirq.Simulator()
-
-    # Cases for measuring X, Y, and Z (respectively) on the memory qubit.
-    params = [
-        {'exponent': 0.5, 'phase_exponent': -0.5},
-        {'exponent': 0.5, 'phase_exponent': 0},
-        {'exponent': 0, 'phase_exponent': 0},
-    ]
-
-    results = simulator.run_sweep(circuit, params, repetitions=5000)
+        results = simulator.run_sweep(
+            self.measure_circuit(self.hhl_circuit()), params, repetitions=total_runs
+        )
+
+        success_probability = []
+        for label, result in zip(('X', 'Y', 'Z'), list(results)):
+            # Only select cases where the ancilla is 1.
+            expectation = 1 - 2 * np.mean(result.measurements['m'][result.measurements['a'] == 1])
+            estimates = (result.measurements['a'] == 1).sum()
+            print(
+                f'{label} = {expectation} from {estimates} estimates '
+                f'out of {total_runs} simulation runs'
+            )
+            success_probability.append(estimates / total_runs)
+
+        return np.mean(success_probability)
+
+    def simulate_with_amplification(self, total_estimates):
+        """Simulates the HHL circuit with amplitude amplification.
+
+        Amplitude amplification is applied to the HHL circuit until the ancilla is measured as 1.
+        Afterwards X, Y and Z are measured on the memory qubit.
+
+        Args:
+            total_estimates: The number of estimates to obtain for each X, Y and Z measurement.
+        """
+        simulator = cirq.Simulator()
+
+        # Cases for measuring X, Y, and Z (respectively) on the memory qubit.
+        params = [
+            {'exponent': 0.5, 'phase_exponent': -0.5},
+            {'exponent': 0.5, 'phase_exponent': 0},
+            {'exponent': 0, 'phase_exponent': 0},
+        ]
 
-    for label, result in zip(('X', 'Y', 'Z'), list(results)):
-        # Only select cases where the ancilla is 1.
-        # TODO: optimize using amplitude amplification algorithm.
-        # Github issue: https://github.com/quantumlib/Cirq/issues/2216
-        expectation = 1 - 2 * np.mean(result.measurements['m'][result.measurements['a'] == 1])
-        print(f'{label} = {expectation}')
+        results = []
+        runs = []
+        for param in params:
+            resolver = cirq.ParamResolver(param)
+            results_for_param = []
+            repetitions = total_estimates
+            num_runs = 0  # number of calls to simulator.run
+            # We use the algorithm "Quantum Searching Without Knowing Success Probability II"
+            # from Kaye et al., Section 8.4.
+            for l in range(1, 10):
+                M = 2**l
+                # Apply amplitude amplification a random number of times between 0 and M-1.
+                # For M large enough, doing this twice increases the success probability to
+                # 3/4 - O(1/(M*θ)), where θ is the angle such that sin^2(θ) equals the probability
+                # of success of a single run without amplification.
+                for _ in range(2):
+                    j = random.randint(0, M - 1)
+                    result = simulator.run(
+                        self.measure_circuit(self.amplitude_amplification(num_iterations=j)),
+                        resolver,
+                        repetitions,
+                    )
+                    num_runs += repetitions
+                    results_for_param.append(
+                        result.measurements['m'][result.measurements['a'] == 1]
+                    )
+                    repetitions -= (result.measurements['a'] == 1).sum()
+                    if repetitions <= 0:
+                        break
+                if repetitions <= 0:
+                    break
+
+            results.append(np.concatenate(results_for_param, axis=0))
+            runs.append(num_runs)
+
+        for label, result, num_runs in zip(('X', 'Y', 'Z'), results, runs):
+            expectation = 1 - 2 * np.mean(result)
+            num_estimates = result.shape[0]
+            print(
+                f'{label} = {expectation} from {num_estimates} estimates '
+                f'out of {num_runs} simulation runs'
+            )
 
 
 def main():
@@ -291,12 +444,24 @@ def main():
     C = 2 * math.pi / (2**register_size * t)
 
     # Simulate circuit.
-    print("Expected observable outputs:")
-    print("X =", expected[0])
-    print("Y =", expected[1])
-    print("Z =", expected[2])
-    print("Actual: ")
-    simulate(hhl_circuit(A, C, t, register_size, *input_prep_gates))
+    print('Expected observable outputs:')
+    print('X =', expected[0])
+    print('Y =', expected[1])
+    print('Z =', expected[2])
+
+    hhl = HHLAlgorithm(A, C, t, register_size, *input_prep_gates)
+    print('Actual values without amplitude amplification: ')
+    success_probability = hhl.simulate_without_amplification(total_runs=5000)
+    print(
+        f'Success probability of a single run without amplitude amplification: '
+        f'{success_probability}'
+    )
+    # We show that amplitude amplification obtains similar values using fewer simulations.
+    # Because amplitude amplification continues until the algorithm succeeds, the number of
+    # estimates is set to 5000 * success_probability to be roughly equivalent to 5000 runs
+    # without amplitude amplification.
+    print('Actual values with amplitude amplification: ')
+    hhl.simulate_with_amplification(total_estimates=int(5000 * success_probability))
 
 
 if __name__ == '__main__':

From 96f24b448252600c8fd3290d22032da6b0b22352 Mon Sep 17 00:00:00 2001
From: Codrut <codrut@google.com>
Date: Sun, 16 Mar 2025 12:13:42 +0100
Subject: [PATCH 2/4] Add an explanation for the number of iterations in
 amplitude amplification.

---
 examples/hhl.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/examples/hhl.py b/examples/hhl.py
index 89937a7578e..d6ebd2d05a4 100644
--- a/examples/hhl.py
+++ b/examples/hhl.py
@@ -381,11 +381,13 @@ def simulate_with_amplification(self, total_estimates):
             num_runs = 0  # number of calls to simulator.run
             # We use the algorithm "Quantum Searching Without Knowing Success Probability II"
             # from Kaye et al., Section 8.4.
+            # We perform 9 iterations, which suffices if the success probability p of a single run
+            # without amplification is at least 1/2^10.
             for l in range(1, 10):
                 M = 2**l
                 # Apply amplitude amplification a random number of times between 0 and M-1.
                 # For M large enough, doing this twice increases the success probability to
-                # 3/4 - O(1/(M*θ)), where θ is the angle such that sin^2(θ) equals the probability
+                # 3/4 - O(1/(M*θ)), where θ is the angle such that sin^2(θ) = p, the probability
                 # of success of a single run without amplification.
                 for _ in range(2):
                     j = random.randint(0, M - 1)

From 5e7884fee7f7a0b5d821583f02cc2c07713b6be3 Mon Sep 17 00:00:00 2001
From: Codrut <codrut@google.com>
Date: Sun, 6 Apr 2025 19:42:53 +0200
Subject: [PATCH 3/4] Add parameter types and factor the output print to a
 separate method.

---
 examples/hhl.py | 92 ++++++++++++++++++++++++++++++++-----------------
 1 file changed, 60 insertions(+), 32 deletions(-)

diff --git a/examples/hhl.py b/examples/hhl.py
index d6ebd2d05a4..150c9c9ca24 100644
--- a/examples/hhl.py
+++ b/examples/hhl.py
@@ -72,6 +72,9 @@
 
 import math
 import random
+from collections.abc import Mapping
+from typing import Any
+
 import numpy as np
 import sympy
 import cirq
@@ -208,7 +211,9 @@ class HHLAlgorithm:
     and one without.
     """
 
-    def __init__(self, A, C, t, register_size, *input_prep_gates):
+    def __init__(
+        self, A: np.ndarray, C: float, t: float, register_size: int, *input_prep_gates: cirq.Gate
+    ):
         """Initializes the HHL algorithm.
 
         Args:
@@ -266,7 +271,7 @@ def diffusion_operator(self):
         c.append(cirq.X.on_each(qubits))
         return c
 
-    def amplitude_amplification(self, num_iterations):
+    def amplitude_amplification(self, num_iterations: int):
         """Applies amplitude amplification to increase the probability of measuring the ancilla
         qubit as 1.
 
@@ -291,7 +296,7 @@ def amplitude_amplification(self, num_iterations):
 
         return c
 
-    def measure_circuit(self, circuit):
+    def measure_circuit(self, circuit: cirq.Circuit):
         """Measures the ancilla qubit and the memory qubit.
 
         There are two parameters in the circuit, `exponent` and `phase_exponent` corresponding
@@ -316,22 +321,47 @@ def measure_circuit(self, circuit):
 
         return circuit
 
-    def simulate_without_amplification(self, total_runs):
+    def print_results(self, results: list[np.ndarray], num_runs: list[int]) -> np.floating[Any]:
+        """Prints the measurement results for X, Y, and Z measurements.
+
+        Args:
+            results: List containing memory qubit measurements for X, Y and Z.
+            num_runs: List containing the total number of simulation runs to obtain X, Y and Z.
+
+        Returns:
+            the success probability, as the mean of (# of successful runs) / (# of runs) needed
+            to obtain an estimate for X, Y and Z.
+        """
+        success_probability = []
+
+        for label, result, runs in zip(('X', 'Y', 'Z'), results, num_runs):
+            expectation = 1 - 2 * np.mean(result)
+            num_estimates = result.shape[0]
+            print(
+                f'{label} = {expectation} from {num_estimates} estimates '
+                f'out of {runs} simulation runs'
+            )
+            success_probability.append(num_estimates / runs)
+
+        return np.mean(success_probability)
+
+    def simulate_without_amplification(self, total_runs: int) -> list[np.ndarray]:
         """Simulates the HHL circuit without amplitude amplification.
 
-        Measures X, Y and Z on the memory qubit.
-        Only the runs where the ancilla is measured as 1 are considered.
+        At the end measures X, Y and Z on the memory qubit.
+        Only the runs where the ancilla is measured as 1 are returned.
 
         Args:
-            total_runs: The number of simulation runs.
+            total_runs: The total number of simulations to run.
 
         Returns:
-            The success probability of measuring the ancilla qubit as 1.
+            A list containing the simulation results where the ancilla is 1, for each
+            X, Y and Z measurement.
         """
         simulator = cirq.Simulator()
 
         # Cases for measuring X, Y, and Z (respectively) on the memory qubit.
-        params = [
+        params: list[Mapping[str, float]] = [
             {'exponent': 0.5, 'phase_exponent': -0.5},
             {'exponent': 0.5, 'phase_exponent': 0},
             {'exponent': 0, 'phase_exponent': 0},
@@ -341,20 +371,11 @@ def simulate_without_amplification(self, total_runs):
             self.measure_circuit(self.hhl_circuit()), params, repetitions=total_runs
         )
 
-        success_probability = []
-        for label, result in zip(('X', 'Y', 'Z'), list(results)):
-            # Only select cases where the ancilla is 1.
-            expectation = 1 - 2 * np.mean(result.measurements['m'][result.measurements['a'] == 1])
-            estimates = (result.measurements['a'] == 1).sum()
-            print(
-                f'{label} = {expectation} from {estimates} estimates '
-                f'out of {total_runs} simulation runs'
-            )
-            success_probability.append(estimates / total_runs)
-
-        return np.mean(success_probability)
+        return [result.measurements['m'][result.measurements['a'] == 1] for result in results]
 
-    def simulate_with_amplification(self, total_estimates):
+    def simulate_with_amplification(
+        self, total_estimates: int
+    ) -> tuple[list[np.ndarray], list[int]]:
         """Simulates the HHL circuit with amplitude amplification.
 
         Amplitude amplification is applied to the HHL circuit until the ancilla is measured as 1.
@@ -362,11 +383,16 @@ def simulate_with_amplification(self, total_estimates):
 
         Args:
             total_estimates: The number of estimates to obtain for each X, Y and Z measurement.
+
+        Returns:
+            A pair of lists, containing for each X, Y and Z measurement goal:
+                - a numpy array with the results for successful simulations.
+                - the total number of simulation runs.
         """
         simulator = cirq.Simulator()
 
         # Cases for measuring X, Y, and Z (respectively) on the memory qubit.
-        params = [
+        params: list[Mapping[str, float]] = [
             {'exponent': 0.5, 'phase_exponent': -0.5},
             {'exponent': 0.5, 'phase_exponent': 0},
             {'exponent': 0, 'phase_exponent': 0},
@@ -409,13 +435,7 @@ def simulate_with_amplification(self, total_estimates):
             results.append(np.concatenate(results_for_param, axis=0))
             runs.append(num_runs)
 
-        for label, result, num_runs in zip(('X', 'Y', 'Z'), results, runs):
-            expectation = 1 - 2 * np.mean(result)
-            num_estimates = result.shape[0]
-            print(
-                f'{label} = {expectation} from {num_estimates} estimates '
-                f'out of {num_runs} simulation runs'
-            )
+        return results, runs
 
 
 def main():
@@ -453,7 +473,8 @@ def main():
 
     hhl = HHLAlgorithm(A, C, t, register_size, *input_prep_gates)
     print('Actual values without amplitude amplification: ')
-    success_probability = hhl.simulate_without_amplification(total_runs=5000)
+    results = hhl.simulate_without_amplification(total_runs=5000)
+    success_probability = hhl.print_results(results, num_runs=[5000, 5000, 5000])
     print(
         f'Success probability of a single run without amplitude amplification: '
         f'{success_probability}'
@@ -463,7 +484,14 @@ def main():
     # estimates is set to 5000 * success_probability to be roughly equivalent to 5000 runs
     # without amplitude amplification.
     print('Actual values with amplitude amplification: ')
-    hhl.simulate_with_amplification(total_estimates=int(5000 * success_probability))
+    results, num_runs = hhl.simulate_with_amplification(
+        total_estimates=int(5000 * success_probability)
+    )
+    success_probability = hhl.print_results(results, num_runs)
+    print(
+        f'Success probability of a single run with amplitude amplification: '
+        f'{success_probability}'
+    )
 
 
 if __name__ == '__main__':

From f925b355fa8187aa4b228e688883e6bd78029830 Mon Sep 17 00:00:00 2001
From: Codrut <codrut@google.com>
Date: Sun, 6 Apr 2025 20:26:19 +0200
Subject: [PATCH 4/4] Delete empty line.

---
 examples/hhl.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/examples/hhl.py b/examples/hhl.py
index 3b7fa18eada..99322434da3 100644
--- a/examples/hhl.py
+++ b/examples/hhl.py
@@ -75,7 +75,6 @@
 from collections.abc import Mapping
 from typing import Any
 
-
 import numpy as np
 import sympy