Add ctrl system for LinearCombination (#1595)

* add failing tests

* add ctrl flag and system for `LinearCombination`

---------

Co-authored-by: Matthew Harrigan <mpharrigan@google.com>

Author: anurudhp

qualtran/bloqs/block_encoding/chebyshev_polynomial.py

@@ -11,14 +11,15 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-
+from collections import Counter
 from functools import cached_property
 from typing import Dict, Tuple, TYPE_CHECKING, Union
 
 import attrs
 import numpy as np
 
 from qualtran import (
+    Bloq,
     bloq_example,
     BloqBuilder,
     BloqDocSpec,
@@ -36,11 +37,7 @@
 from qualtran.symbolics import is_symbolic, SymbolicFloat, SymbolicInt
 
 if TYPE_CHECKING:
-    from qualtran.resource_counting import (
-        BloqCountDictT,
-        MutableBloqCountDictT,
-        SympySymbolAllocator,
-    )
+    from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator
 
 
 @attrs.frozen
@@ -144,12 +141,11 @@ def build_composite_bloq(self, bb: BloqBuilder, **soqs: SoquetT) -> Dict[str, So
 
     def build_call_graph(self, ssa: 'SympySymbolAllocator') -> 'BloqCountDictT':
         n = self.order
-        s: 'MutableBloqCountDictT' = {
-            self.block_encoding: n // 2 + n % 2,
-            self.block_encoding.adjoint(): n // 2,
-        }
+        s = Counter[Bloq]()
+        s[self.block_encoding] += n // 2 + n % 2
+        s[self.block_encoding.adjoint()] += n // 2
         if is_symbolic(self.ancilla_bitsize) or self.ancilla_bitsize > 0:
-            s[self.reflection_bloq] = n - n % 2
+            s[self.reflection_bloq] += n - n % 2
         return s
 
 

qualtran/bloqs/block_encoding/chebyshev_polynomial_test.py

@@ -34,6 +34,7 @@
 from qualtran.bloqs.reflections.prepare_identity import PrepareIdentity
 from qualtran.bloqs.state_preparation.black_box_prepare import BlackBoxPrepare
 from qualtran.linalg.matrix import random_hermitian_matrix
+from qualtran.resource_counting import get_cost_value, QECGatesCost
 from qualtran.symbolics import is_symbolic, SymbolicFloat, SymbolicInt
 from qualtran.testing import assert_equivalent_bloq_example_counts, execute_notebook
 
@@ -156,6 +157,11 @@ def test_scaled_chebyshev_even_tensors():
     np.testing.assert_allclose(from_gate, from_tensors, atol=0.06)
 
 
+def test_scaled_chebyshev_even_cost():
+    bloq = _scaled_chebyshev_poly_even()
+    _ = get_cost_value(bloq, QECGatesCost())
+
+
 @pytest.mark.slow
 def test_scaled_chebyshev_odd_tensors():
     from_gate = t5(Hadamard().tensor_contract() * 3.14)

qualtran/bloqs/block_encoding/linear_combination.ipynb

@@ -56,7 +56,8 @@
     " - `lambd`: Corresponding coefficients.\n",
     " - `lambd_bits`: Number of bits needed to represent coefficients precisely.\n",
     " - `prepare`: If specified, oracle preparing $\\sum_i \\sqrt{|\\lambda_i|} |i\\rangle$ (state should be normalized and can have junk).\n",
-    " - `select`: If specified, oracle taking $|i\\rangle|\\psi\\rangle \\mapsto \\text{sgn}(\\lambda_i) |i\\rangle U_i|\\psi\\rangle$. \n",
+    " - `select`: If specified, oracle taking $|i\\rangle|\\psi\\rangle \\mapsto \\text{sgn}(\\lambda_i) |i\\rangle U_i|\\psi\\rangle$.\n",
+    " - `is_controlled`: if True, implements a controlled version. Defaults to False. \n",
     "\n",
     "#### Registers\n",
     " - `system`: The system register.\n",

qualtran/bloqs/block_encoding/linear_combination.py

@@ -20,17 +20,20 @@
 from typing_extensions import Self
 
 from qualtran import (
+    AddControlledT,
+    Bloq,
     bloq_example,
     BloqBuilder,
     BloqDocSpec,
     BQUInt,
+    CtrlSpec,
+    DecomposeTypeError,
     QAny,
     Register,
     Signature,
     Soquet,
     SoquetT,
 )
-from qualtran._infra.bloq import DecomposeTypeError
 from qualtran.bloqs.block_encoding import BlockEncoding
 from qualtran.bloqs.block_encoding.lcu_block_encoding import BlackBoxPrepare, BlackBoxSelect
 from qualtran.bloqs.block_encoding.phase import Phase
@@ -68,6 +71,7 @@ class LinearCombination(BlockEncoding):
             (state should be normalized and can have junk).
         select: If specified, oracle taking
             $|i\rangle|\psi\rangle \mapsto \text{sgn}(\lambda_i) |i\rangle U_i|\psi\rangle$.
+        is_controlled: if True, implements a controlled version. Defaults to False.
 
     Registers:
         system: The system register.
@@ -88,6 +92,8 @@ class LinearCombination(BlockEncoding):
     _prepare: Optional[BlackBoxPrepare] = None
     _select: Optional[BlackBoxSelect] = None
 
+    is_controlled: bool = False
+
     def __attrs_post_init__(self):
         if len(self._block_encodings) != len(self._lambd):
             raise ValueError("Must provide the same number of block encodings and coefficients.")
@@ -138,6 +144,7 @@ def rescaled_lambd(self):
     @cached_property
     def signature(self) -> Signature:
         return Signature.build_from_dtypes(
+            ctrl=QAny(1 if self.is_controlled else 0),
             system=QAny(self.system_bitsize),
             ancilla=QAny(self.ancilla_bitsize),
             resource=QAny(self.resource_bitsize),
@@ -298,11 +305,26 @@ def build_composite_bloq(
         be_part = Partition(self.select.system_bitsize, tuple(be_regs))
 
         prepare_soqs = bb.add_d(self.prepare, **prepare_in_soqs)
-        select_out_soqs = bb.add_d(
-            self.select,
-            selection=prepare_soqs.pop("selection"),
-            system=cast(Soquet, bb.add(evolve(be_part, partition=False), **be_system_soqs)),
-        )
+
+        if not self.is_controlled:
+            select_out_soqs = bb.add_d(
+                self.select,
+                selection=prepare_soqs.pop("selection"),
+                system=cast(Soquet, bb.add(evolve(be_part, partition=False), **be_system_soqs)),
+            )
+        else:
+            _, add_ctrl_select = self.select.get_ctrl_system(CtrlSpec())
+            (ctrl,), select_out_soqs_t = add_ctrl_select(
+                bb,
+                [soqs.pop('ctrl')],
+                dict(
+                    selection=prepare_soqs.pop("selection"),
+                    system=cast(Soquet, bb.add(evolve(be_part, partition=False), **be_system_soqs)),
+                ),
+            )
+            select_out_soqs = dict(zip(["selection", "system"], select_out_soqs_t))
+            select_out_soqs["ctrl"] = ctrl
+
         prep_adj_soqs = bb.add_d(
             self.prepare.adjoint(), selection=select_out_soqs.pop("selection"), **prepare_soqs
         )
@@ -311,6 +333,9 @@ def build_composite_bloq(
         be_soqs = bb.add_d(be_part, x=select_out_soqs.pop("system"))
         out: Dict[str, SoquetT] = {"system": be_soqs.pop("system")}
 
+        if self.is_controlled:
+            out["ctrl"] = select_out_soqs.pop("ctrl")
+
         # merge ancilla registers of block encoding and Prepare oracle
         anc_soqs = {"selection": prep_adj_soqs.pop("selection")}
         if self.be_ancilla_bitsize > 0:
@@ -331,6 +356,20 @@ def build_composite_bloq(
     def __str__(self) -> str:
         return f"B[{'+'.join(str(be)[2:-1] for be in self.signed_block_encodings)}]"
 
+    def get_ctrl_system(self, ctrl_spec: 'CtrlSpec') -> tuple['Bloq', 'AddControlledT']:
+        from qualtran.bloqs.mcmt.specialized_ctrl import get_ctrl_system_1bit_cv_from_bloqs
+
+        return get_ctrl_system_1bit_cv_from_bloqs(
+            self,
+            ctrl_spec,
+            current_ctrl_bit=1 if self.is_controlled else None,
+            bloq_with_ctrl=evolve(self, is_controlled=True),
+            ctrl_reg_name='ctrl',
+        )
+
+    def adjoint(self) -> 'LinearCombination':
+        return self
+
 
 @bloq_example
 def _linear_combination_block_encoding() -> LinearCombination:

qualtran/bloqs/block_encoding/linear_combination_test.py

@@ -39,6 +39,7 @@
 from qualtran.bloqs.block_encoding.unitary import Unitary
 from qualtran.bloqs.for_testing.matrix_gate import MatrixGate
 from qualtran.bloqs.reflections.prepare_identity import PrepareIdentity
+from qualtran.resource_counting import get_cost_value, QECGatesCost
 from qualtran.testing import execute_notebook
 
 
@@ -103,6 +104,11 @@ def test_linear_combination_tensors():
     np.testing.assert_allclose(from_gate, from_tensors)
 
 
+def test_linear_combination_cost():
+    bloq = _linear_combination_block_encoding()
+    _ = get_cost_value(bloq.controlled(), QECGatesCost())
+
+
 def run_gate_test(gates, lambd, lambd_bits=1, atol=1e-07):
     bloq = LinearCombination(tuple(Unitary(g) for g in gates), lambd, lambd_bits)
     from_gate = sum(l * g.tensor_contract() for l, g in zip(lambd, gates))