Hamming weight phasing with configurable number of ancilla

qualtran/bloqs/rotations/hamming_weight_phasing.py

@@ -209,3 +209,94 @@ def _hamming_weight_phasing_via_phase_gradient() -> HammingWeightPhasingViaPhase
     bloq_cls=HammingWeightPhasingViaPhaseGradient,
     examples=(_hamming_weight_phasing_via_phase_gradient,),
 )
+
+
+@attrs.frozen
+class HammingWeightPhasingWithConfigurableAncilla(GateWithRegisters):
+    r"""
+    Args:
+        bitsize: Size of input register to apply 'Z ** exponent' to.
+        ancillasize: Size of the ancilla register to be used to calculate the hamming weight of 'x'.
+        exponent: the exponent of 'Z ** exponent' to be applied to each qubit in the input register.
+        eps: Accuracy of synthesizing the Z rotations.
+
+    Registers:
+        x: A 'THRU' register of 'bitsize' qubits.
+
+    References:
+    """
+
+    bitsize: int
+    ancillasize: int # TODO: verify that ancillasize is always < bitsize-1
+    exponent: float = 1
+    eps: SymbolicFloat = 1e-10
+
+    @cached_property
+    def signature(self) -> 'Signature':
+        return Signature.build_from_dtypes(x=QUInt(self.bitsize))
+
+
+    '''
+    General strategy: find the max-bitsize number (n bits) we can compute the HW of using our available ancilla,
+    greedily do this on the first n bits of x, perform the rotations, then the next n bits and perform those
+    rotations, and so on until we have computed the HW of the entire input. Can express this as repeated calls to
+    HammingWeightPhasing bloqs on subsets of the input.
+    '''
+    def build_composite_bloq(self, bb: 'BloqBuilder', *, x: 'SoquetT') -> Dict[str, 'SoquetT']:
+        num_iters = self.bitsize // (self.ancillasize + 1)
+        remainder = self.bitsize % (self.ancillasize+1)
+        x = bb.split(x)
+        x_parts = []
+        for i in range(num_iters):
+            x_part = bb.join(x[i*(self.ancillasize+1):(i+1)*(self.ancillasize+1)], dtype=QUInt(self.ancillasize+1))
+            x_part = bb.add(HammingWeightPhasing(bitsize=self.ancillasize+1, exponent=self.exponent, eps=self.eps), x=x_part)
+            x_parts.extend(bb.split(x_part))
+        if remainder > 1:
+            x_part = bb.join(x[(-1*remainder):], dtype=QUInt(remainder))
+            x_part = bb.add(HammingWeightPhasing(bitsize=remainder, exponent=self.exponent, eps=self.eps), x=x_part)
+            x_parts.extend(bb.split(x_part))
+        if remainder == 1:
+            x_part = x[-1]
+            x_part = bb.add(ZPowGate(exponent=self.exponent, eps=self.eps), q=x_part)
+            x_parts.append(x_part)
+        x = bb.join(np.array(x_parts), dtype=QUInt(self.bitsize))
+        return {'x': x}
+
+
+    def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
+        if reg is None:
+            return Text(f'HWPCA_{self.bitsize}/(Z^{self.exponent})')
+        return super().wire_symbol(reg, idx)
+
+
+    def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
+        num_iters = self.bitsize // (self.ancillasize + 1)
+        remainder = self.bitsize - (self.ancillasize + 1) * num_iters
+        # TODO: Surely there is a better way of doing this
+        if remainder > 1:
+
+            return {
+                HammingWeightPhasing(self.ancillasize+1, self.exponent, self.eps): num_iters,
+                HammingWeightPhasing(remainder, self.exponent, self.eps): bool(remainder),
+            }
+        elif remainder:
+            return {
+                HammingWeightPhasing(self.ancillasize+1, self.exponent, self.eps): num_iters,
+                ZPowGate(exponent=self.exponent, eps=self.eps): 1
+            }
+        else:
+            return {
+                HammingWeightPhasing(self.ancillasize+1, self.exponent, self.eps): num_iters,
+            }
+
+
+@bloq_example
+def _hamming_weight_phasing_with_configurable_ancilla() -> HammingWeightPhasingWithConfigurableAncilla:
+    hamming_weight_phasing_with_configurable_ancilla = HammingWeightPhasingWithConfigurableAncilla(4, 2, np.pi / 2.0)
+    return hamming_weight_phasing_with_configurable_ancilla
+
+
+_HAMMING_WEIGHT_PHASING_WITH_CONFIGURABLE_ANCILLA_DOC = BloqDocSpec(
+    bloq_cls=HammingWeightPhasingWithConfigurableAncilla,
+    examples=(_hamming_weight_phasing_with_configurable_ancilla,),
+)

qualtran/bloqs/rotations/hamming_weight_phasing_test.py

@@ -23,6 +23,7 @@
 from qualtran.bloqs.rotations.hamming_weight_phasing import (
     HammingWeightPhasing,
     HammingWeightPhasingViaPhaseGradient,
+    HammingWeightPhasingWithConfigurableAncilla,
 )
 from qualtran.bloqs.rotations.phase_gradient import PhaseGradientState
 from qualtran.cirq_interop.testing import GateHelper
@@ -127,3 +128,30 @@ def test_hamming_weight_phasing_via_phase_gradient_t_complexity(n: int, theta: f
     naive_total_t = naive_hwp_t_complexity.t_incl_rotations(eps=eps / n.bit_length())
 
     assert total_t < naive_total_t
+
+@pytest.mark.parametrize('n, ancillasize', [(n, ancillasize) for n in range(3, 9) for ancillasize in range(1, n-1)])
+@pytest.mark.parametrize('theta', [1 / 10, 1 / 5, 1 / 7, np.pi / 2])
+def test_hamming_weight_phasing_with_configurable_ancilla(n: int, ancillasize: int, theta: float):
+    gate = HammingWeightPhasingWithConfigurableAncilla(n, ancillasize, theta)
+    qlt_testing.assert_valid_bloq_decomposition(gate)
+    qlt_testing.assert_equivalent_bloq_counts(
+        gate, [ignore_split_join, cirq_to_bloqs, generalize_rotation_angle]
+    )
+
+    remainder = n % (ancillasize+1)
+
+#    assert gate.t_complexity().rotations == (-(-n // (ancillasize+1))-1) * (ancillasize+1).bit_length() + remainder.bit_length() # exact, fails for remainder = 0.
+    assert gate.t_complexity().rotations <= (-(-n // (ancillasize+1))) * (ancillasize+1).bit_length() + remainder.bit_length() # upper bound
+    assert gate.t_complexity().t <= 4 * (ancillasize) * -(-n // (ancillasize+1))
+    # TODO: add an assertion that number of ancilla allocated is never > ancillasize.
+
+    gh = GateHelper(gate)
+    sim = cirq.Simulator(dtype=np.complex128)
+    initial_state = cirq.testing.random_superposition(dim=2**n, random_state=12345)
+    state_prep = cirq.Circuit(cirq.StatePreparationChannel(initial_state).on(*gh.quregs['x']))
+    brute_force_phasing = cirq.Circuit(state_prep, (cirq.Z**theta).on_each(*gh.quregs['x']))
+    expected_final_state = sim.simulate(brute_force_phasing).final_state_vector
+
+    hw_phasing = cirq.Circuit(state_prep, HammingWeightPhasingWithConfigurableAncilla(n, ancillasize, theta).on(*gh.quregs['x']))
+    hw_final_state = sim.simulate(hw_phasing).final_state_vector
+    assert np.allclose(expected_final_state, hw_final_state, atol=1e-7)