Address Matt's comments

Author: tanujkhattar

cirq-ft/cirq_ft/algos/and_gate.py

@@ -49,7 +49,9 @@ class And(infra.GateWithRegisters):
         ValueError: If number of control values (i.e. `len(self.cv)`) is less than 2.
     """
 
-    cv: Tuple[int, ...] = attr.field(default=(1, 1), converter=infra.to_tuple)
+    cv: Tuple[int, ...] = attr.field(
+        default=(1, 1), converter=lambda v: (v,) if isinstance(v, int) else tuple(v)
+    )
     adjoint: bool = False
 
     @cv.validator

cirq-ft/cirq_ft/algos/arithmetic_gates.py

@@ -611,7 +611,9 @@ class AddMod(cirq.ArithmeticGate):
     bitsize: int
     mod: int = attr.field()
     add_val: int = 1
-    cv: Tuple[int, ...] = attr.field(converter=infra.to_tuple, default=())
+    cv: Tuple[int, ...] = attr.field(
+        converter=lambda v: (v,) if isinstance(v, int) else tuple(v), default=()
+    )
 
     @mod.validator
     def _validate_mod(self, attribute, value):

cirq-ft/cirq_ft/algos/mean_estimation/mean_estimation_operator.py

@@ -80,7 +80,9 @@ class MeanEstimationOperator(infra.GateWithRegisters):
     """
 
     code: CodeForRandomVariable
-    cv: Tuple[int, ...] = attr.field(converter=infra.to_tuple, default=())
+    cv: Tuple[int, ...] = attr.field(
+        converter=lambda v: (v,) if isinstance(v, int) else tuple(v), default=()
+    )
     power: int = 1
     arctan_bitsize: int = 32
 

cirq-ft/cirq_ft/algos/multi_control_multi_target_pauli.py

@@ -73,7 +73,7 @@ class MultiControlPauli(infra.GateWithRegisters):
         (https://algassert.com/circuits/2015/06/05/Constructing-Large-Controlled-Nots.html)
     """
 
-    cvs: Tuple[int, ...] = attr.field(converter=infra.to_tuple)
+    cvs: Tuple[int, ...] = attr.field(converter=lambda v: (v,) if isinstance(v, int) else tuple(v))
     target_gate: cirq.Pauli = cirq.X
 
     @cached_property

cirq-ft/cirq_ft/algos/prepare_uniform_superposition.py

@@ -46,7 +46,9 @@ class PrepareUniformSuperposition(infra.GateWithRegisters):
     """
 
     n: int
-    cv: Tuple[int, ...] = attr.field(converter=infra.to_tuple, default=())
+    cv: Tuple[int, ...] = attr.field(
+        converter=lambda v: (v,) if isinstance(v, int) else tuple(v), default=()
+    )
 
     @cached_property
     def registers(self) -> infra.Registers:

cirq-ft/cirq_ft/infra/__init__.py

@@ -25,4 +25,3 @@
 from cirq_ft.infra.qubit_management_transformers import map_clean_and_borrowable_qubits
 from cirq_ft.infra.qubit_manager import GreedyQubitManager
 from cirq_ft.infra.t_complexity_protocol import TComplexity, t_complexity
-from cirq_ft.infra.type_convertors import to_tuple

cirq-ft/cirq_ft/infra/gate_with_registers.py

@@ -43,7 +43,7 @@ def all_idxs(self) -> Iterable[Tuple[int, ...]]:
     def total_bits(self) -> int:
         """The total number of bits in this register.
 
-        This is the product of bitsize and each of the dimensions in `shape`.
+        This is the product of each of the dimensions in `shape`.
         """
         return int(np.product(self.shape))
 
@@ -52,12 +52,16 @@ def __repr__(self):
 
 
 def total_bits(registers: Iterable[Register]) -> int:
+    """Sum of `reg.total_bits()` for each register `reg` in input `registers`."""
+
     return sum(reg.total_bits() for reg in registers)
 
 
 def split_qubits(
     registers: Iterable[Register], qubits: Sequence[cirq.Qid]
 ) -> Dict[str, NDArray[cirq.Qid]]:  # type: ignore[type-var]
+    """Splits the flat list of qubits into a dictionary of appropriately shaped qubit arrays."""
+
     qubit_regs = {}
     base = 0
     for reg in registers:
@@ -70,21 +74,25 @@ def merge_qubits(
     registers: Iterable[Register],
     **qubit_regs: Union[cirq.Qid, Sequence[cirq.Qid], NDArray[cirq.Qid]],
 ) -> List[cirq.Qid]:
+    """Merges the dictionary of appropriately shaped qubit arrays into a flat list of qubits."""
+
     ret: List[cirq.Qid] = []
     for reg in registers:
-        assert (
-            reg.name in qubit_regs
-        ), f"All qubit registers must be present. {reg.name} not in qubit_regs"
+        if reg.name not in qubit_regs:
+            raise ValueError(f"All qubit registers must be present. {reg.name} not in qubit_regs")
         qubits = qubit_regs[reg.name]
         qubits = np.array([qubits] if isinstance(qubits, cirq.Qid) else qubits)
-        assert (
-            qubits.shape == reg.shape
-        ), f'{reg.name} register must of shape {reg.shape} but is of shape {qubits.shape}'
+        if qubits.shape != reg.shape:
+            raise ValueError(
+                f'{reg.name} register must of shape {reg.shape} but is of shape {qubits.shape}'
+            )
         ret += qubits.flatten().tolist()
     return ret
 
 
 def get_named_qubits(registers: Iterable[Register]) -> Dict[str, NDArray[cirq.Qid]]:
+    """Returns a dictionary of appropriately shaped named qubit registers for input `registers`."""
+
     def _qubit_array(reg: Register):
         qubits = np.empty(reg.shape, dtype=object)
         for ii in reg.all_idxs():
@@ -169,6 +177,43 @@ class SelectionRegister(Register):
 
     `SelectionRegister` extends the `Register` class to store the iteration length
     corresponding to that register along with its size.
+
+    LCU methods often make use of coherent for-loops via UnaryIteration, iterating over a range
+    of values stored as a superposition over the `SELECT` register. Such (nested) coherent
+    for-loops can be represented using a `Tuple[SelectionRegister, ...]` where the i'th entry
+    stores the bitsize and iteration length of i'th nested for-loop.
+
+    One useful feature when processing such nested for-loops is to flatten out a composite index,
+    represented by a tuple of indices (i, j, ...), one for each selection register into a single
+    integer that can be used to index a flat target register. An example of such a mapping
+    function is described in Eq.45 of https://arxiv.org/abs/1805.03662. A general version of this
+    mapping function can be implemented using `numpy.ravel_multi_index` and `numpy.unravel_index`.
+
+    For example:
+        1) We can flatten a 2D for-loop as follows
+        >>> import numpy as np
+        >>> N, M = 10, 20
+        >>> flat_indices = set()
+        >>> for x in range(N):
+        ...     for y in range(M):
+        ...         flat_idx = x * M + y
+        ...         assert np.ravel_multi_index((x, y), (N, M)) == flat_idx
+        ...         assert np.unravel_index(flat_idx, (N, M)) == (x, y)
+        ...         flat_indices.add(flat_idx)
+        >>> assert len(flat_indices) == N * M
+
+        2) Similarly, we can flatten a 3D for-loop as follows
+        >>> import numpy as np
+        >>> N, M, L = 10, 20, 30
+        >>> flat_indices = set()
+        >>> for x in range(N):
+        ...     for y in range(M):
+        ...         for z in range(L):
+        ...             flat_idx = x * M * L + y * L + z
+        ...             assert np.ravel_multi_index((x, y, z), (N, M, L)) == flat_idx
+        ...             assert np.unravel_index(flat_idx, (N, M, L)) == (x, y, z)
+        ...             flat_indices.add(flat_idx)
+        >>> assert len(flat_indices) == N * M * L
     """
 
     iteration_length: int = attr.field()

cirq-ft/cirq_ft/infra/gate_with_registers_test.py

@@ -57,6 +57,13 @@ def test_registers():
     assert qregs["r3"].tolist() == [cirq.LineQubit(7)]
 
     qubits = qubits[::-1]
+
+    with pytest.raises(ValueError, match="qubit registers must be present"):
+        _ = merge_qubits(regs, r1=qubits[:5], r2=qubits[5:7], r4=qubits[-1])
+
+    with pytest.raises(ValueError, match="register must of shape"):
+        _ = merge_qubits(regs, r1=qubits[:4], r2=qubits[5:7], r3=qubits[-1])
+
     merged_qregs = merge_qubits(regs, r1=qubits[:5], r2=qubits[5:7], r3=qubits[-1])
     assert merged_qregs == qubits