Hamming weight phasing with configurable number of ancilla

qualtran/bloqs/rotations/hamming_weight_phasing.py

@@ -1,3 +1,3 @@
 #  Copyright 2023 Google LLC
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,11 +14,12 @@
 
 from functools import cached_property
 from typing import Dict, Optional, Tuple, TYPE_CHECKING
+from collections import Counter
 
 import attrs
 import numpy as np
 
-from qualtran import bloq_example, BloqDocSpec, GateWithRegisters, QFxp, QUInt, Register, Signature
+from qualtran import bloq_example, BloqDocSpec, GateWithRegisters, QFxp, QUInt, Register, Signature, Bloq
 from qualtran.bloqs.arithmetic import HammingWeightCompute
 from qualtran.bloqs.basic_gates import ZPowGate
 from qualtran.bloqs.rotations.quantum_variable_rotation import QvrPhaseGradient
@@ -209,3 +210,101 @@
     bloq_cls=HammingWeightPhasingViaPhaseGradient,
     examples=(_hamming_weight_phasing_via_phase_gradient,),
 )
+
+
+@attrs.frozen
+class HammingWeightPhasingWithConfigurableAncilla(Bloq):
+    r""""Applies $Z^{\text{exponent}}$ to every qubit of an input register of size `bitsize`.
+
+    See docstring of 'HammingWeightPhasing' for more details about how hamming weight phasing works, and the docstring of 'HammingWeightPhasingViaPhaseGradient' for more details about how the rotations are synthesized.
+
+    This is a variant of hamming weight phasing via phase gradient which uses a constant number of ancilla specified by the user. See Appendix A.2 of [1] (pages 19-21) for details. Note that this method has increased T-complexity compared to hamming weight phasing via phase gradient, since the hamming weight of the entire register cannot be calculated at once due to the limited number of ancilla available. Instead, we calculate the hamming weight in parts, performing the full rotation over the course of $\lceil n/(r+1)\rceil$ repetitions. Again, see [1] for a detailed analysis of the resource costs of this technique compared to vanilla hamming weight phasing and compard with hamming weight phasing via phase gradient. Also see the note in the 'HammingWeightPhasingViaPhaseGradient' docstring for information about when these methods are actually practical to use over vannilla hamming weight phasing.
+
+
+
+    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:
+        [Improved Fault-Tolerant Quantum Simulation of Condensed-Phase Correlated Electrons via Trotterization](https://arxiv.org/abs/1902.10673)
+        Appendix A.2: Hamming weight phasing with limited ancilla
+    """
+
+    bitsize: SymbolicInt
+    ancillasize: SymbolicInt # TODO: verify that ancillasize is always < bitsize-1
+    exponent: SymbolicFloat = 1
+    eps: SymbolicFloat = 1e-10
+
+    @cached_property
+    def signature(self) -> 'Signature':
+        return Signature.build_from_dtypes(x=QUInt(self.bitsize))
+
+    def __attrs_post_init__(self):
+        if self.ancillasize >= self.bitsize - 1:
+            raise ValueError('ancillasize should be less than bitsize - 1.')
+
+
+    def build_composite_bloq(self, bb: 'BloqBuilder', *, x: 'SoquetT') -> Dict[str, 'SoquetT']:
+        '''
+        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.
+        '''
+        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
+
+        counts = Counter[Bloq]()
+        counts[HammingWeightPhasing(self.ancillasize+1, self.exponent, self.eps)] += num_iters
+
+        if remainder > 1:
+            counts[HammingWeightPhasing(remainder, self.exponent, self.eps)] += 1
+        elif remainder:
+            counts[ZPowGate(exponent=self.exponent, eps=self.eps)] += 1
+
+        return counts
+
+
+@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

@@ -1,3 +1,3 @@
 #  Copyright 2023 Google LLC
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
@@ -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
@@ -31,6 +32,12 @@
     generalize_rotation_angle,
     ignore_split_join,
 )
+from qualtran.resource_counting import (
+    QECGatesCost, 
+    GateCounts, 
+    get_cost_value,
+)
+
 from qualtran.symbolics import SymbolicInt
 
 if TYPE_CHECKING:
@@ -127,3 +134,39 @@
     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]
+    )
+
+    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)
+
+
+@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_bloq_counts(n: int, ancillasize: int, theta: float):
+    gate = HammingWeightPhasingWithConfigurableAncilla(n, ancillasize, theta)
+    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.
+#    Outdated method of doing gatecount tests. Also tests T-count rather than Toffoli count.
+#    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)
+
+    gc = get_cost_value(gate, QECGatesCost())
+    assert gc.rotation <= (-(-n // (ancillasize+1))) * (ancillasize+1).bit_length() + remainder.bit_length()
+    assert gc.toffoli + gc.and_bloq + gc.cswap <= ancillasize * -(-n // (ancillasize+1))