Implement logarithm.

Author: fedimser

notebooks/accuracy/log_fbe.ipynb

@@ -0,0 +1,64 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "87da278b",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "OK\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Evaluating log2(x), x in [1,2). \n",
+    "# Reference: https://arxiv.org/pdf/2001.00807 (section 3.1.1)\n",
+    "\n",
+    "import numpy as np\n",
+    "\n",
+    "# r - precision\n",
+    "def log2_fbe(x, r=40):\n",
+    "    w=[0]*r\n",
+    "    assert 1<=x<2\n",
+    "    a=x**2\n",
+    "    for i in range(r):\n",
+    "        assert 1<=a<4\n",
+    "        if a>=2:\n",
+    "            w[i]=1\n",
+    "            a=a/2\n",
+    "        a=a**2\n",
+    "    return sum(w[i]*(2**(-(i+1))) for i in range(r))\n",
+    "\n",
+    "for x in np.linspace(1, 1.99, 100):\n",
+    "    assert np.isclose(np.log2(x), log2_fbe(x))\n",
+    "\n",
+    "print(\"OK\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

notebooks/classic/log2_qfbe.ipynb

@@ -0,0 +1,22 @@
+"""Bitwise functions."""
+
+from psiqworkbench import QFixed, QInt, QUInt, Qubits
+from psiqworkbench.qubricks import Qubrick
+
+
+class HighestSetBit(Qubrick):
+    """Finds most significant set bit in a and sets it in ans."""
+
+    def _compute(self, a: Qubits, ans: Qubits):
+        flag: Qubits = self.alloc_temp_qreg(1, "flag")
+
+        # For each input qubit i compute which output qubit must be set if i is MSB.
+        for i in range(a.num_qubits - 1, -1, -1):
+            # Copy a[i] to ans[j], but only if flag is unset.
+            flag.x()
+            ans[i].lelbow(a[i] | flag)
+            flag.x()
+
+            # If ans[i]=1 (which implies flag was unset), set the flag.
+            # All less significant qubits will be ignored.
+            flag.x(ans[i])

qmath/func/bits.py

@@ -8,15 +8,21 @@
     https://arxiv.org/abs/2001.00807
 """
 
+import math
+
 import psiqworkbench.qubricks as qbk
-from psiqworkbench import QFixed, QInt, Qubits, QUInt
+from psiqworkbench import QFixed, QInt, Qubits, QUFixed, QUInt
 from psiqworkbench.qubits.base_qubits import BaseQubits
 from psiqworkbench.qubricks import Qubrick
 from psiqworkbench.symbolics.qubrick_costs import QubrickCosts
 
-from .common import AddConst, Negate
-from .sqrt import Sqrt
+from ..utils.gates import ParallelCnot, ParallelCnotCtrl, write_int
+from ..utils.rotate import Div2, rotate_left, rotate_right
 from ..utils.symbolic import alloc_temp_qreg_like
+from .bits import HighestSetBit
+from .common import AddConst, Negate, MultiplyConstAdd
+from .sqrt import Sqrt
+from .square import Square, SquareOptimized
 
 
 def _sqrt_half(x: QFixed) -> QFixed:
@@ -119,3 +125,97 @@ def _compute(self, x: QFixed):
         cos_op = CosFbe(result_radix=self.result_radix)
         cos_op.compute(x)
         self.set_result_qreg(cos_op.get_result_qreg())
+
+
+class Log2FbeSegment(Qubrick):
+    """Computes log2(x) where 1<=x<2.
+
+    Reference: https://arxiv.org/abs/2001.00807, section 3.1.1.
+    """
+
+    def _square(self, x: QUFixed) -> QUFixed:
+        result = QUFixed(self.alloc_temp_qreg(x.num_qubits, name="a"), radix=x.radix)
+        SquareOptimized(signed=False).compute(x, result)
+        return result
+
+    def _compute(self, x: QUFixed, result: QUFixed):
+        assert x.num_qubits == 2 + x.radix
+        assert result.num_qubits == result.radix
+        a = self._square(x)
+
+        for i in range(result.radix):
+            result_bit = result[result.radix - 1 - i]
+            result_bit.x(a[-1])
+            Div2().compute(a, ctrl=result_bit)
+            a = self._square(a)
+
+
+class Log2Fbe(Qubrick):
+    """Computes log2(x) where x>0."""
+
+    def __init__(
+        self,
+        *,
+        result_radix: None | int = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.result_radix = result_radix
+
+    def _compute(self, x: QUFixed):
+        # Find most significant bit of `x`, set it it `msb`.
+        xn = x.num_qubits
+        msb = self.alloc_temp_qreg(xn, "msb")
+        HighestSetBit().compute(x, msb)
+
+        # Make shifted copy of input, such that second most significnat bit in
+        # the copy corresponds to highest set bit in input.
+        # This way value in x_copy is in range [1, 2).
+        x_copy_qubits = self.alloc_temp_qreg(xn, name="x_copy")
+        x_copy = QUFixed(x_copy_qubits, radix=xn - 2)
+        for i in range(xn):
+            # Controlled shift-copy.
+            if i == xn - 1:
+                ParallelCnotCtrl().compute(msb[i], x[1:], x_copy_qubits[0 : xn - 1])
+            else:
+                shift_left = xn - 2 - i
+                assert shift_left >= 0
+                ParallelCnotCtrl().compute(msb[i], x[0 : xn - shift_left], x_copy_qubits[shift_left:])
+
+        # Compute logarithm for shifted copy.
+        r = self.result_radix or x.radix
+        result_fract_part = QUFixed(self.alloc_temp_qreg(r, "result_frac"), radix=r)
+        Log2FbeSegment().compute(x_copy, result_fract_part)
+
+        # Add integer to result, corresponding to input's shift.
+        int_part_size = math.ceil(math.log2(max(x.radix, xn - x.radix))) + 1
+        result_int_part = QInt(self.alloc_temp_qreg(int_part_size, "result_int"))
+        for i in range(xn):
+            write_int(result_int_part, i - x.radix, ctrl=msb[i])
+
+        self.set_result_qreg(QFixed(result_fract_part | result_int_part, radix=r))
+
+
+class LogFbe(Qubrick):
+    """Computes logarithm in given base."""
+
+    def __init__(
+        self,
+        base: float,
+        *,
+        result_radix: None | int = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.ans_multiplier = 1.0 / math.log2(base)
+        self.result_radix = result_radix
+
+    def _compute(self, x: QUFixed):
+        op = Log2Fbe(result_radix=self.result_radix)
+        with op.computed(x):
+            _, ans = alloc_temp_qreg_like(self, op.get_result_qreg())
+            if self.ans_multiplier == 1.0:
+                ParallelCnot().compute(op.get_result_qreg(), ans)
+            else:
+                MultiplyConstAdd(self.ans_multiplier).compute(ans, op.get_result_qreg())
+        self.set_result_qreg(ans)

qmath/func/fbe.py

@@ -1,7 +1,12 @@
-from qmath.func.fbe import CosFbe, SinFbe
-from qmath.utils.test_utils import QPUTestHelper
+import os
 
 import numpy as np
+import pytest
+
+from qmath.func.fbe import CosFbe, Log2Fbe, SinFbe
+from qmath.utils.test_utils import QPUTestHelper
+
+RUN_SLOW_TESTS = os.getenv("RUN_SLOW_TESTS") == "1"
 
 
 def test_cos():
@@ -30,3 +35,19 @@ def test_sin():
         result = qpu_helper.apply_op([x])
         expected = np.sin(np.pi * x)
         assert abs(result - expected) < 1e-4
+
+
+@pytest.mark.skipif(not RUN_SLOW_TESTS, reason="slow test")
+def test_log2():
+    qpu_helper = QPUTestHelper(num_inputs=1, num_qubits=1000, qubits_per_reg=30, radix=24)
+    qs_x = qpu_helper.inputs[0]
+    op = Log2Fbe(result_radix=11)
+    op.compute(qs_x)
+    result = op.get_result_qreg()
+    qpu_helper.record_op(op.get_result_qreg())
+
+    x_range = list(np.linspace(1, 2, 21)) + [1e-5, 1e-3, 0.5, 3, 4, 10, 20]
+    for x in x_range:
+        result = qpu_helper.apply_op([x])
+        expected = np.log2(x)
+        assert abs(result - expected) < 2e-3

qmath/func/fbe_test.py

@@ -7,29 +7,16 @@
 from .square import Square
 from .common import Subtract, MultiplyAdd
 from ..utils.symbolic import alloc_temp_qreg_like
+from .bits import HighestSetBit
 
 
 class _InitialGuess(Qubrick):
-    def _msb(self, a: Qubits, ans: Qubits):
-        """Finds most significant bit in a and sets it in ans."""
-        flag: Qubits = self.alloc_temp_qreg(1, "flag")
-
-        # For each input qubit i compute which output qubit must be set if i is MSB.
-        for i in range(a.num_qubits - 1, -1, -1):
-            # Copy a[i] to ans[j], but only if flag is unset.
-            flag.x()
-            ans[i].lelbow(a[i] | flag)
-            flag.x()
-
-            # If ans[i]=1 (which implies flag was unset), set the flag.
-            # All less significant qubits will be ignored.
-            flag.x(ans[i])
 
     def _compute(self, a: QFixed, ans: QFixed):
         """Computes ans := 2**(-(floor(log2(a)))//2)."""
         # TODO: can this be optimized to compute result directly into ans?
         r = self.alloc_temp_qreg(a.num_qubits, "r")
-        self._msb(a, r)
+        HighestSetBit().compute(a, r)
 
         for i in range(a.num_qubits - 1, -1, -1):
             pos1 = i - a.radix

qmath/func/inv_sqrt.py

@@ -44,6 +44,10 @@ class SquareOptimized(Qubrick):
     While this precision issue is not fixed, it's recommedned to use Square instead.
     """
 
+    def __init__(self, *, signed=True, **kwargs):
+        super().__init__(**kwargs)
+        self.signed = signed
+
     def _compute_unsigned(self, x: QUFixed, target: QUFixed):
         """Computes square assuming x is unsigned."""
         anc: Qubits = self.alloc_temp_qreg(target.num_qubits, "anc")
@@ -61,10 +65,13 @@ def _compute_unsigned(self, x: QUFixed, target: QUFixed):
         anc.release()
 
     def _compute(self, x: QFixed, target: QFixed):
-        with AbsInPlace().computed(x):
-            x_unsigned = QUFixed(x[0 : x.num_qubits - 1], radix=x.radix)
-            target_unsigned = QUFixed(target[0 : target.num_qubits - 1], radix=target.radix)
-            self._compute_unsigned(x_unsigned, target_unsigned)
+        if self.signed:
+            with AbsInPlace().computed(x):
+                x_unsigned = QUFixed(x[0 : x.num_qubits - 1], radix=x.radix)
+                target_unsigned = QUFixed(target[0 : target.num_qubits - 1], radix=target.radix)
+                self._compute_unsigned(x_unsigned, target_unsigned)
+        else:
+            self._compute_unsigned(x, target)
 
     def _estimate(self, x: SymbolicQFixed, target: SymbolicQFixed):
         n = x.num_qubits

qmath/func/square.py

@@ -1,4 +1,4 @@
-from psiqworkbench import Qubits, QUInt, Qubrick, QFixed
+from psiqworkbench import Qubits, QUInt, Qubrick, QFixed, QInt
 from psiqworkbench.symbolics.qubrick_costs import QubrickCosts
 
 from typing import Optional
@@ -14,6 +14,13 @@ def ccnot(a: Qubits, b: Qubits, c: Qubits, ctrl: Optional[Qubits] = None):
     c.x(a | b | ctrl)
 
 
+def swap(x: Qubits, y: Qubits, ctrl: Optional[Qubits] = None):
+    if ctrl is None:
+        x.swap(y)
+    else:
+        x.swap(y, condition_mask=ctrl)
+
+
 def write_uint(target: QUInt, number: int, ctrl: Optional[Qubits] = None):
     """Writes target ⊕= number*ctrl."""
     assert 0 <= number < 2**target.num_qubits
@@ -22,6 +29,12 @@ def write_uint(target: QUInt, number: int, ctrl: Optional[Qubits] = None):
             target[i].x(ctrl)
 
 
+def write_int(target: QInt, number: int, ctrl: Optional[Qubits] = None):
+    n = target.num_qubits
+    assert -(2 ** (n - 1)) <= number < 2 ** (n - 1)
+    write_uint(target, number % (2**n), ctrl=ctrl)
+
+
 def write_qfixed(target: QFixed, number: float, ctrl: Optional[Qubits] = None):
     """Writes target ⊕= number*ctrl."""
     assert 0 <= number < 2**target.num_qubits
@@ -41,3 +54,17 @@ def _estimate(self, a: Qubits, b: Qubits):
         n = a.num_qubits
         assert b.num_qubits == n
         self.get_qc().add_cost_event(QubrickCosts(active_volume=4 * n))
+
+
+class ParallelCnotCtrl(Qubrick):
+    def _compute(
+        self,
+        ctrl: Qubits,
+        src: Qubits,
+        dst: Qubits,
+    ):
+        assert len(ctrl) == 1
+        n = src.num_qubits
+        assert dst.num_qubits == n
+        for i in range(n):
+            dst[i].x(src[i] | ctrl)