Update classical action of addition gates and fix classical action bu…

…g in Join (#1174)

* Update classical action of addition gates and fix classical action bug in Join

* address comments

* fix typo

---------

Co-authored-by: Matthew Harrigan <[email protected]>

qualtran/bloqs/arithmetic/addition.py

@@ -11,7 +11,6 @@
 #  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.
-import math
 from functools import cached_property
 from typing import (
     Dict,
@@ -59,6 +58,7 @@
 from qualtran.bloqs.mcmt.multi_control_multi_target_pauli import MultiControlX
 from qualtran.cirq_interop import decompose_from_cirq_style_method
 from qualtran.drawing import directional_text_box, Text
+from qualtran.simulation.classical_sim import add_ints
 
 if TYPE_CHECKING:
     from qualtran.drawing import WireSymbol
@@ -129,19 +129,10 @@ def on_classical_vals(
     ) -> Dict[str, 'ClassicalValT']:
         unsigned = isinstance(self.a_dtype, (QUInt, QMontgomeryUInt))
         b_bitsize = self.b_dtype.bitsize
-        N = 2**b_bitsize
-        if unsigned:
-            return {'a': a, 'b': int((a + b) % N)}
-
-        # Addition of signed integers can result in overflow. In most classical programming languages (e.g. C++)
-        # what happens when an overflow happens is left as an implementation detail for compiler designers.
-        # However for quantum subtraction the operation should be unitary and that means that the unitary of
-        # the bloq should be a permutation matrix.
-        # If we hold `a` constant then the valid range of values of `b` [-N/2, N/2) gets shifted forward or backwards
-        # by `a`. to keep the operation unitary overflowing values wrap around. this is the same as moving the range [0, N)
-        # by the same amount modulu $N$. that is add N/2 before addition and then remove it.
-        half_n = N >> 1
-        return {'a': a, 'b': int(a + b + half_n) % N - half_n}
+        return {
+            'a': a,
+            'b': add_ints(int(a), int(b), num_bits=int(b_bitsize), is_signed=not unsigned),
+        }
 
     def _circuit_diagram_info_(self, _) -> cirq.CircuitDiagramInfo:
         wire_symbols = ["In(x)"] * int(self.a_dtype.bitsize)
@@ -302,7 +293,13 @@ def adjoint(self) -> 'OutOfPlaceAdder':
     def on_classical_vals(
         self, *, a: 'ClassicalValT', b: 'ClassicalValT'
     ) -> Dict[str, 'ClassicalValT']:
-        return {'a': a, 'b': b, 'c': a + b}
+        if isinstance(self.bitsize, sympy.Expr):
+            raise ValueError(f'Classical simulation is not support for symbolic bloq {self}')
+        return {
+            'a': a,
+            'b': b,
+            'c': add_ints(int(a), int(b), num_bits=self.bitsize + 1, is_signed=False),
+        }
 
     def with_registers(self, *new_registers: Union[int, Sequence[int]]):
         raise NotImplementedError("no need to implement with_registers.")
@@ -421,14 +418,18 @@ def signature(self) -> 'Signature':
     def on_classical_vals(
         self, x: 'ClassicalValT', **vals: 'ClassicalValT'
     ) -> Dict[str, 'ClassicalValT']:
+        if isinstance(self.k, sympy.Expr) or isinstance(self.bitsize, sympy.Expr):
+            raise ValueError(f"Classical simulation isn't supported for symbolic block {self}")
         N = 2**self.bitsize
         if len(self.cvs) > 0:
             ctrls = vals['ctrls']
         else:
-            return {'x': int(math.fmod(x + self.k, N))}
+            return {
+                'x': add_ints(int(x), int(self.k), num_bits=self.bitsize, is_signed=self.signed)
+            }
 
         if np.all(self.cvs == ctrls):
-            x = int(math.fmod(x + self.k, N))
+            x = add_ints(int(x), int(self.k), num_bits=self.bitsize, is_signed=self.signed)
 
         return {'ctrls': ctrls, 'x': x}
 

qualtran/bloqs/arithmetic/addition_test.py

@@ -247,10 +247,12 @@ def test_add_classical():
 def test_out_of_place_adder():
     basis_map = {}
     gate = OutOfPlaceAdder(bitsize=3)
+    cbloq = gate.decompose_bloq()
     for x in range(2**3):
         for y in range(2**3):
             basis_map[int(f'0b_{x:03b}_{y:03b}_0000', 2)] = int(f'0b_{x:03b}_{y:03b}_{x+y:04b}', 2)
             assert gate.call_classically(a=x, b=y, c=0) == (x, y, x + y)
+            assert cbloq.call_classically(a=x, b=y, c=0) == (x, y, x + y)
     op = GateHelper(gate).operation
     op_inv = cirq.inverse(op)
     cirq.testing.assert_equivalent_computational_basis_map(basis_map, cirq.Circuit(op))
@@ -323,9 +325,9 @@ def test_classical_add_signed_overflow(bitsize):
     assert bloq.call_classically(a=mx, b=mx) == (mx, -2)
 
 
-# TODO: write tests for signed integer addition (subtraction)
-# https://github.com/quantumlib/Qualtran/issues/606
-@pytest.mark.parametrize('bitsize,k,x,cvs,ctrls,result', [(5, 2, 0, (1, 0), (1, 0), 2)])
+@pytest.mark.parametrize(
+    'bitsize,k,x,cvs,ctrls,result', [(5, 2, 0, (1, 0), (1, 0), 2), (6, -3, 2, (), (), -1)]
+)
 def test_classical_add_k_signed(bitsize, k, x, cvs, ctrls, result):
     bloq = AddK(bitsize=bitsize, k=k, cvs=cvs, signed=True)
     cbloq = bloq.decompose_bloq()

qualtran/bloqs/arithmetic/negate.py

@@ -61,7 +61,7 @@ def signature(self) -> 'Signature':
 
     def build_composite_bloq(self, bb: 'BloqBuilder', x: 'SoquetT') -> dict[str, 'SoquetT']:
         x = bb.add(BitwiseNot(self.dtype), x=x)  # ~x
-        x = bb.add(AddK(self.dtype.num_qubits, k=1), x=x)  # -x
+        x = bb.add(AddK(self.dtype.num_qubits, k=1, signed=isinstance(self.dtype, QInt)), x=x)  # -x
         return {'x': x}
 
 

qualtran/bloqs/bookkeeping/join.py

@@ -26,6 +26,7 @@
     DecomposeTypeError,
     QBit,
     QDType,
+    QFxp,
     QUInt,
     Register,
     Side,
@@ -95,7 +96,10 @@ def my_tensors(
         ]
 
     def on_classical_vals(self, reg: 'NDArray[np.uint]') -> Dict[str, int]:
-        return {'reg': bits_to_ints(reg)[0]}
+        if isinstance(self.dtype, QFxp):
+            # TODO(#1095): support QFxp in classical simulation
+            return {'reg': bits_to_ints(reg)[0]}
+        return {'reg': self.dtype.from_bits(reg.tolist())}
 
     def wire_symbol(self, reg: Optional[Register], idx: Tuple[int, ...] = tuple()) -> 'WireSymbol':
         if reg is None:

qualtran/simulation/classical_sim.py

@@ -14,7 +14,7 @@
 
 """Functionality for the `Bloq.call_classically(...)` protocol."""
 import itertools
-from typing import Any, Dict, Iterable, List, Mapping, Sequence, Tuple, Type, Union
+from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple, Type, Union
 
 import networkx as nx
 import numpy as np
@@ -265,3 +265,33 @@ def format_classical_truth_table(
         for invals, outvals in truth_table
     ]
     return '\n'.join([heading] + entries)
+
+
+def add_ints(a: int, b: int, *, num_bits: Optional[int] = None, is_signed: bool = False) -> int:
+    r"""Performs addition modulo $2^\mathrm{num\_bits}$ of (un)signed in a reversible way.
+
+    Addition of signed integers can result in an overflow. In most classical programming languages (e.g. C++)
+    what happens when an overflow happens is left as an implementation detail for compiler designers. However,
+    for quantum subtraction, the operation should be unitary and that means that the unitary of the bloq should
+    be a permutation matrix.
+
+    If we hold `a` constant then the valid range of values of $b \in [-2^{\mathrm{num\_bits}-1}, 2^{\mathrm{num\_bits}-1})$
+    gets shifted forward or backward by `a`. To keep the operation unitary overflowing values wrap around. This is the same
+    as moving the range $2^\mathrm{num\_bits}$ by the same amount modulo $2^\mathrm{num\_bits}$. That is add
+    $2^{\mathrm{num\_bits}-1})$ before addition modulo and then remove it.
+
+    Args:
+        a: left operand of addition.
+        b: right operand of addition.
+        num_bits: optional num_bits. When specified addition is done in the interval [0, 2**num_bits) or
+            [-2**(num_bits-1), 2**(num_bits-1)) based on the value of `is_signed`.
+        is_signed: boolean whether the numbers are unsigned or signed ints. This value is only used when
+            `num_bits` is provided.
+    """
+    c = a + b
+    if num_bits is not None:
+        N = 2**num_bits
+        if is_signed:
+            return (c + N // 2) % N - N // 2
+        return c % N
+    return c

qualtran/simulation/classical_sim_test.py

@@ -12,6 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+import itertools
 from typing import Dict
 
 import cirq
@@ -24,6 +25,7 @@
 from qualtran.bloqs.basic_gates import CNOT
 from qualtran.simulation.classical_sim import (
     _update_assign_from_vals,
+    add_ints,
     bits_to_ints,
     call_cbloq_classically,
     ClassicalValT,
@@ -168,6 +170,35 @@ def test_apply_classical_cbloq():
     np.testing.assert_array_equal(z, xarr)
 
 
+@pytest.mark.parametrize(
+    ['x', 'y', 'n_bits'],
+    [
+        (x, y, n_bits)
+        for n_bits in range(1, 5)
+        for x, y in itertools.product(range(1 << n_bits), repeat=2)
+    ],
+)
+def test_add_ints_unsigned(x, y, n_bits):
+    assert add_ints(x, y, num_bits=n_bits, is_signed=False) == (x + y) % (1 << n_bits)
+
+
+@pytest.mark.parametrize(
+    ['x', 'y', 'n_bits'],
+    [
+        (x, y, n_bits)
+        for n_bits in range(2, 5)
+        for x, y in itertools.product(range(-(2 ** (n_bits - 1)), 2 ** (n_bits - 1)), repeat=2)
+    ],
+)
+def test_add_ints_signed(x, y, n_bits):
+    half_n = 1 << (n_bits - 1)
+    # Addition of signed ints `x` and `y` is a cyclic rotation of the interval [-2^(n-1), 2^(n-1)) by `y`.
+    interval = [*range(-(2 ** (n_bits - 1)), 2 ** (n_bits - 1))]
+    i = x + half_n  # position of `x` in the interval
+    z = interval[(i + y) % len(interval)]  # rotate by `y`
+    assert add_ints(x, y, num_bits=n_bits, is_signed=True) == z
+
+
 @pytest.mark.notebook
 def test_notebook():
     execute_notebook('classical_sim')