In the quantum mean computation circuit, how to apply the controlled rotation to the extra qubit?

Question

I'm trying to implement quantum mean computation via amplitude estimation, as suggested in Example 8.3.5 in Kaye, LaFlamme, Mosca. I wonder about the actual circuit to accomplish this. This previous answer suggested to construct an operator $U_a$ for a (normalized) $F(x)$:

$$ U_a:|0\rangle|0\rangle\mapsto \frac{1}{2^{n/2}} \sum_x|x\rangle(\sqrt{1-F(x)}|0\rangle+\sqrt{F(x)}|1\rangle). $$

It is important to note that this unitary is easily implemented. You start with a Hadamard transform on the first register, perform a computation of f(x) on an ancilla register, use this to implement a controlled-rotation of the second register, and then uncompute the ancilla register.

I have trouble implementing this. It is clear that in the formula above the first $|0\rangle$ is the multi-qubit state $|00...0\rangle$ which I have to put in equal superposition with $H^{\otimes n}$. I call the second $|0\rangle$ the 'extra' qubit.

For each state $x$, I compute $v = F(x)$. So if my state vector is $[0.1026, 0.2052, 0.3078, 0.9234]$ then I compute, for example, $F(0) = v = 0.1026$. I rotate an ancilla via $2 R_y(\arcsin(v))$. So far, so good.

The next step is confusing. How do I apply the controlled rotation of the extra qubit? Just having something like a Controlled-Not from the ancilla to the extra qubit will not work and just leave the state in an equal superposition. But even have a fully controlled rotation, controlled by the bit patterns of $x$ doesn't seem to be working?

Is there perhaps a pointer to a working implementation? (Thx)

Answer 1

Self-answer: This is actually a bit simpler than I thought. For two qubits (and this can be easily generalized to more qubits), the circuit would have to look like the following, where the $\theta_i$ corresponding to the results of calls to $F(x_i)$ as $\theta_i = 2 \arcsin(F(x_i))$:

With this circuit one can estimate the mean (of the $F(x_i)$) by measuring the probability of the extra qubit to be in state $|1\rangle$. Expressing the circuit as a single big-matrix operator would allow amplitude estimation to find the same results.

I've put a working version in Python (here).