"""
********************************************************************************
* Copyright (c) 2026 the Qrisp authors
*
* This program and the accompanying materials are made available under the
* terms of the Eclipse Public License 2.0 which is available at
* http://www.eclipse.org/legal/epl-2.0.
*
* This Source Code may also be made available under the following Secondary
* Licenses when the conditions for such availability set forth in the Eclipse
* Public License, v. 2.0 are satisfied: GNU General Public License, version 2
* with the GNU Classpath Exception which is
* available at https://www.gnu.org/software/classpath/license.html.
*
* SPDX-License-Identifier: EPL-2.0 OR GPL-2.0 WITH Classpath-exception-2.0
********************************************************************************
"""
import numpy as np
from scipy.special import comb
import jax.numpy as jnp
from qrisp import (
QuantumVariable,
QuantumFloat,
x,
z,
ry,
control,
conjugate,
invert,
)
from qrisp.alg_primitives.reflection import reflection
from qrisp.block_encodings import BlockEncoding
from qrisp.jasp import jrange
import numpy.typing as npt
from typing import Tuple, TYPE_CHECKING
if TYPE_CHECKING:
from jax.typing import ArrayLike
[docs]
def cks_params(eps: float, kappa: float, max_beta: int = None) -> Tuple[int, int]:
"""
Computes the complexity parameter :math:`\\beta` and the truncation order :math:`j_0` for the
truncated Chebyshev approximation of :math:`1/x`, as described in the `Childs–Kothari–Somma paper <https://arxiv.org/abs/1511.02306>`_.
Given the condition number :math:`\\kappa` and the target precision
:math:`\\epsilon`, the parameters are computed as:
.. math::
\\beta = \kappa^2 \log\!\left(\\frac{\kappa}{\epsilon}\\right),
\quad
j_0 = \sqrt{\\beta \log\!\left(\\frac{4\\beta}{\epsilon}\\right)}.
If ``max_beta`` is provided, :math:`\\beta` is capped to
:math:`\\min(\\beta, \\beta_{\\max})`. The returned values are cast to integers
via the floor operation.
Parameters
----------
eps : float
Target precision :math:`\epsilon`.
kappa : float
An upper bound for the condition number :math:`\\kappa` of :math:`A`. This value defines the "gap"
around zero where the function :math:`1/x` is not approximated.
max_beta : float, optional
Optional upper bound on the complexity parameter :math:`\\beta`.
Returns
-------
j0 : int
Truncation order of the Chebyshev expansion
:math:`j_0 = \\lfloor\sqrt{\\beta \log(4\\beta/\epsilon)}\\rfloor`.
beta : int
Complexity parameter :math:`\\beta = \\lfloor\kappa^2 \log(\kappa/\epsilon)\\rfloor`.
"""
if max_beta == None:
beta = kappa**2 * np.log(kappa / eps)
else:
beta = min(kappa**2 * np.log(kappa / eps), max_beta)
j0 = np.sqrt(beta * np.log(4 * beta / eps))
return int(j0), int(beta)
[docs]
def cks_coeffs(j0: int, b: int) -> npt.NDArray[float]:
"""
Computes the positive coefficients :math:`\\alpha_i` for the truncated
Chebyshev expansion of :math:`1/x` up to order :math:`2j_0+1`.
The approximation is expressed as a linear combination
of odd Chebyshev polynomials truncated at index :math:`j_0` (Lemma 14):
.. math::
g(x) = 4 \\sum_{j=0}^{j_0} (-1)^j
\\left[ \\sum_{i=j+1}^{b} \\frac{\\binom{2b}{b+i}}{2^{2b}} \\right]
T_{2j+1}(x)
The Linear Combination of Unitaries (LCU) lemma requires strictly
positive coefficients :math:`\\alpha_i > 0`, their absolute values are
used. The alternating factor :math:`(-1)^j` is later implemented as a set of
Z-gates within the CKS circuit.
Parameters
----------
j0 : int
Truncation order of the Chebyshev expansion,
:math:`j_0 = \\lfloor\sqrt{\\beta \log(4\\beta/\epsilon)}\\rfloor`.
b : float
Complexity parameter :math:`\\beta = \\lfloor\kappa^2 \log(\kappa/\epsilon)\\rfloor`.
Returns
-------
coeffs : ndarray
1-D array of positive Chebyshev coefficients :math:`{\\alpha_{2j+1}}`,
corresponding to the odd degree Chebyshev polynomials of the first kind :math:`T_1, T_3, \\dotsc, T_{2j_0+1}`.
"""
coeffs = []
for j in range(j0 + 1):
sum_i = 0
for i in range(j + 1, b + 1):
sum_i += comb(2 * b, b + i)
coeff = 4 * (2 ** (-2 * b)) * sum_i
coeffs.append(coeff)
return np.array(coeffs)
def _unary_angles(coeffs: "ArrayLike") -> "ArrayLike":
"""
Computes rotation angles :math:`\\phi_i` to prepare the unary state :math:`\ket{\\text{unary}}`,
corresponding to the square-root-amplitude encoding of the Chebyshev coefficients.
The prepared unary state is
.. math::
\ket{\\text{unary}} \propto
\\sqrt{\\alpha_1}\ket{100\\dots00} + \\sqrt{\\alpha_3}\ket{110\\dots00} +
\\cdots + \\sqrt{\\alpha_{2j_0+1}}\ket{111\\dots11}
Parameters
----------
coeffs : ArrayLike
1-D array of positive Chebyshev coefficients :math:`\\alpha_i`.
Returns
-------
ArrayLike
1-D array of rotation angles :math:`\\phi_i` for unary state preparation.
"""
alpha = jnp.sqrt(
coeffs
) # coeffs need not to be normalized since unary angles only depend on their ratio
phi = jnp.zeros(len(alpha) - 1)
phi = phi.at[-1].set(
jnp.arctan(alpha[-1] / alpha[-2])
) # Last angle is determined by ratio of final two amplitudes
for i in range(len(phi) - 2, -1, -1):
phi = phi.at[i].set(
jnp.arctan(alpha[i + 1] / alpha[i] / jnp.cos(phi[i + 1]))
) # Recursively compute previous rotation angles
return 2 * phi # Compensate for factor 1/2 in ry
[docs]
def unary_prep(case: QuantumVariable, coeffs: "ArrayLike") -> None:
"""
Prepares the unary-encoded state :math:`\ket{\\text{unary}}` of Chebyshev coefficients.
When applied to a variable in state :math:`\ket{0}`, the resulting superposition is
.. math::
\ket{\\text{unary}} \propto
\\sqrt{\\alpha_1}\ket{100\\dots00} + \\sqrt{\\alpha_3}\ket{110\\dots00} +
\\cdots + \\sqrt{\\alpha_{2j_0+1}}\ket{111\\dots11}
Parameters
----------
case : QuantumVariable
Variable with :math:`j_0` qubits on which the unary state preparation will be performed.
If the variable is in state :math:`\ket{0}`, the state :math:`\ket{\\text{unary}}` is prepared.
coeffs : ArrayLike
1-D array of :math:`j_0` Chebyshev coefficients :math:`\\alpha_1,\\alpha_3,\dotsc,\\alpha_{2j_0+1}`.
"""
phi = _unary_angles(coeffs)
x(case[0])
ry(phi[0], case[1])
for i in jrange(1, case.size - 1):
with control(case[i]):
ry(phi[i], case[i + 1])
[docs]
def CKS(
A: BlockEncoding, eps: float, kappa: float, max_beta: float = None
) -> BlockEncoding:
"""
Performs the `Childs–Kothari–Somma (CKS) quantum algorithm <https://arxiv.org/abs/1511.02306>`_ to solve the Quantum Linear System Problem (QLSP)
:math:`A \\vec{x} = \\vec{b}`, using the Chebyshev approximation of :math:`1/x`.
When applied to a state $\ket{b}$, the algorithm prepares a state :math:`\ket{\\tilde{x}} \propto A^{-1} \ket{b}`
within target precision :math:`\epsilon` of the ideal solution :math:`\ket{x}`.
For a block-encoded **Hermitian** matrix :math:`A`, this function returns a BlockEncoding of an
operator :math:`\\tilde{A}^{-1}` such that :math:`\|\\tilde{A}^{-1} - A^{-1}\| \leq \epsilon`.
The inversion is implemented using a polynomial approximation of :math:`1/x` over the domain :math:`D_{\\kappa} = [-1, -1/\\kappa] \\cup [1/\\kappa, 1]`.
The asymptotic complexity is
:math:`\\mathcal{O}\\!\\left(\\log(N)s \\kappa^2 \\text{polylog}\\!\\frac{s\\kappa}{\\epsilon}\\right)`, where :math:`N` is the matrix size, :math:`s` its sparsity, and
:math:`\kappa` its condition number. This represents an exponentially
better precision scaling compared to the HHL algorithm.
This function integrates all core components of the CKS approach:
Chebyshev polynomial approximation, Linear Combination of Unitaries (LCU),
`qubitization with reflection operators <https://arxiv.org/abs/2208.00567>`_, and the :ref:`Repeat-Until-Success (RUS) <RUS>` protocol.
The semantics of the approach can be illustrated with the following circuit schematics:
.. image:: /_static/CKS_schematics.png
:align: center
Implementation overview:
1. Compute the CKS parameters :math:`j_0` and :math:`\\beta` (:func:`cks_params`).
2. Generate Chebyshev coefficients and the auxiliary unary state (:func:`cks_coeffs`, :func:`unary_prep`).
3. Build the core LCU structure via qubitization operator.
The goal of this algorithm is to apply the non-unitary operator :math:`A^{-1}` to the
input state :math:`\ket{b}`. Following the Chebyshev approach introduced in the CKS paper, we express :math:`A^{-1}` as
a linear combination of odd Chebyshev polynomials:
.. math::
A^{-1}\propto\sum_{j=0}^{j_0}\\alpha_{2j+1}T_{2j+1}(A),
where :math:`T_k(A)` are Chebyshev polynomials of the first kind and :math:`\\alpha_{2j+1} > 0` are computed
via :func:`cks_coeffs`. These operators can be efficiently implemented with qubitization, which relies on a unitary
block encoding :math:`U` of the matrix :math:`A`, and a :ref:`reflection operator <reflection>` :math:`R`.
If the block encoding unitary is $U$ is Hermitian (:math:`U^2=I`),
the fundamental iteration step is defined as :math:`(RU)`, where :math:`R` reflects
around the auxiliary block-encoding state :math:`\ket{G}`, prepared as the ``inner_case`` QuantumFloat.
Repeated applications of these unitaries, :math:`(RU)^k`, yield a block encoding of the :math:`k`-th Chebyshev polynomial
of the first kind :math:`T_k(A)`.
We then construct a linear combination of these block-encoded polynomials using the LCU structure
.. math::
\\text{LCU}\ket{0}\ket{\psi}=\\text{PREP}^{\dagger}\cdot \\text{SEL}\cdot \\text{PREP}\ket{0}\ket{\psi}=\\tilde{A}\ket{0}\ket{\psi}.
Here, the :math:`\\text{PREP}` operation prepares an auxiliary ``out_case`` Quantumfloat in the unary state :math:`\ket{\\text{unary}}`
that encodes the square root of the Chebyshev coefficients :math:`\sqrt{\\alpha_j}`. The :math:`\\text{SEL}` operation selects and applies the
appropriate Chebyshev polynomial operator :math:`T_k(A)`, implemented by :math:`(RU)^k`, controlled on ``out_case`` in the unary
state :math:`\ket{\\text{unary}}`. Based on the Hamming-weight :math:`k` of :math:`\ket{\\text{unary}}`,
the polynomial :math:`T_{2k-1}` is block encoded and applied to the circuit.
To construct a linear combination of Chebyshev polynomials up to the :math:`2j_0+1`-th order, as in the original paper,
our implementation requires :math:`j_0+1` qubits in the ``out_case`` state :math:`\ket{\\text{unary}}`.
The Chebyshev coefficients alternate in sign :math:`(-1)^j\\alpha_j`.
Since the LCU lemma requires :math:`\\alpha_j>0`, negative terms are accounted for
by applying Z-gates on each index qubit in ``out_case``.
Parameters
----------
A : BlockEncoding
The block-encoded Hermitian matrix to be inverted.
eps : float
Target precision :math:`\epsilon`.
kappa : float
An upper bound for the condition number :math:`\\kappa` of :math:`A`. This value defines the "gap"
around zero where the function :math:`1/x` is not approximated.
max_beta : float, optional
Optional upper bound on the complexity parameter :math:`\\beta`.
Returns
-------
BlockEncoding
A new BlockEncoding instance representing an approximation of the inverse $A^{-1}$.
Examples
--------
The following examples demonstrate how the CKS algorithm can be applied to solve the
quantum linear systems problem (QLSP)
:math:`A \\vec{x} = \\vec{b}`, using either a direct Hermitian matrix input or
a preconstructed block-encoding representation.
**Example 1: Solving a 4×4 Hermitian system**
First, we define a small Hermitian matrix :math:`A` and a right-hand side vector :math:`\\vec{b}`:
::
import numpy as np
A = np.array([[0.73255474, 0.14516978, -0.14510851, -0.0391581],
[0.14516978, 0.68701415, -0.04929867, -0.00999921],
[-0.14510851, -0.04929867, 0.76587818, -0.03420339],
[-0.0391581, -0.00999921, -0.03420339, 0.58862043]])
b = np.array([0, 1, 1, 1])
kappa = np.linalg.cond(A)
print("Condition number of A: ", kappa)
# Condition number of A: 1.8448536035491883
Next, we solve this linear system using the CKS quantum algorithm:
::
from qrisp import prepare, QuantumFloat
from qrisp.algorithms.cks import CKS
from qrisp.block_encodings import BlockEncoding
from qrisp.jasp import terminal_sampling
def prep_b():
operand = QuantumFloat(2)
prepare(operand, b)
return operand
@terminal_sampling
def main():
BA = BlockEncoding.from_array(A)
x = CKS(BA, 0.01, kappa).apply_rus(prep_b)()
return x
res_dict = main()
# The resulting dictionary contains the measurement probabilities
# for each computational basis state.
# To extract the corresponding quantum amplitudes (up to sign):
amps = np.sqrt([res_dict.get(i, 0) for i in range(len(b))])
print("QUANTUM SIMULATION\\n", amps)
# QUANTUM SIMULATION
# [0.02714082 0.55709921 0.53035395 0.63845794]
We compare this to the classical normalized solution:
::
c = (np.linalg.inv(A) @ b) / np.linalg.norm(np.linalg.inv(A) @ b)
print("CLASSICAL SOLUTION\\n", c)
# CLASSICAL SOLUTION
# [0.02944539 0.55423278 0.53013239 0.64102936]
We see that we obtained the same result in our quantum simulation up to precision :math:`\epsilon`!
To perform quantum resource estimation, replace the ``@terminal_sampling``
decorator with ``@count_ops(meas_behavior="0")``:
::
from qrisp.jasp import count_ops
@count_ops(meas_behavior="0")
def main():
BA = BlockEncoding.from_array(A)
x = CKS(BA, 0.01, kappa).apply_rus(prep_b)()
return x
res_dict = main()
print(res_dict)
# {'gphase': 51, 't_dg': 1548, 'cz': 150, 'ry': 2, 'z': 12, 'h': 890, 'cy': 50, 'cx': 3516, 'x': 377, 's_dg': 70, 'p': 193, 's': 70, 't': 1173, 'u3': 753, 'measure': 18}
The printed dictionary lists the estimated quantum gate counts required for the CKS algorithm. Since the simulation itself is not executed,
this approach enables scalable gate count estimation for linear systems of arbitrary size.
**Example 2: Using a custom block encoding**
The previous example displays how to solve the linear system for any Hermitian matrix :math:`A`,
where the block-encoding is constructed from the Pauli decomposition of :math:`A`. While this approach is efficient when :math:`A`
corresponds to, e.g., an Ising or a Heisenberg Hamiltonian, the number of Pauli terms may not scale favorably in general.
For certain sparse matrices, their structure can be exploited to construct more efficient block-encodings.
In this example, we construct a block-encoding representation for the following following tridiagonal sparse matrix:
::
import numpy as np
def tridiagonal_shifted(n, mu=1.0, dtype=float):
I = np.eye(n, dtype=dtype)
return (2 + mu) * I - 2*np.eye(n, k=n//2, dtype=dtype) - 2*np.eye(n, k=-n//2, dtype=dtype)
N = 8
A = tridiagonal_shifted(N, mu=3)
b = np.random.randint(0, 2, size=N)
print(A)
# [[ 5. 0. 0. 0. -2. 0. 0. 0.]
# [ 0. 5. 0. 0. 0. -2. 0. 0.]
# [ 0. 0. 5. 0. 0. 0. -2. 0.]
# [ 0. 0. 0. 5. 0. 0. 0. -2.]
# [-2. 0. 0. 0. 5. 0. 0. 0.]
# [ 0. -2. 0. 0. 0. 5. 0. 0.]
# [ 0. 0. -2. 0. 0. 0. 5. 0.]
# [ 0. 0. 0. -2. 0. 0. 0. 5.]]
This matrix can be decomposed using three unitaries: the identity $I$, and two shift operators $V\colon\ket{k}\\rightarrow-\ket{k+N/2 \mod N}$ and $V^{\dagger}\colon\ket{k}\\rightarrow-\ket{k-N/2 \mod N}$.
We define their corresponding functions and the block-encoding using Linear Combination of Unitaries (LCU):
::
from qrisp import gphase
def I(qv):
pass
def V(qv):
qv += N//2
gphase(np.pi, qv[0])
def V_dg(qv):
qv -= N//2
gphase(np.pi, qv[0])
unitaries = [I, V, V_dg]
coeffs = np.array([5,1,1,0])
BA = BlockEncoding.from_lcu(coeffs, unitaries)
Next, we solve this linear system using the CKS quantum algorithm:
::
def prep_b():
operand = QuantumFloat(3)
prepare(operand, b)
return operand
@terminal_sampling
def main():
x = CKS(BA, 0.01, kappa).apply_rus(prep_b)()
return x
res_dict = main()
amps = np.sqrt([res_dict.get(i, 0) for i in range(len(b))])
We, again, compare the solution obtained by quantum simulation with the classical solution
::
c = (np.linalg.inv(A) @ b) / np.linalg.norm(np.linalg.inv(A) @ b)
print("QUANTUM SIMULATION\\n", amps, "\\nCLASSICAL SOLUTION\\n", c)
# QUANTUM SIMULATION
# [0.43917486 0.31395935 0.12522139 0.43917505 0.43917459 0.12522104 0.31395943 0.43917502]
# CLASSICAL SOLUTION
# [0.43921906 0.3137279 0.12549116 0.43921906 0.43921906 0.12549116 0.3137279 0.43921906]
Quantum resource estimation can be performed in the same way as in the first example:
::
@count_ops(meas_behavior="0")
def main():
x = CKS(BA, 0.01, kappa).apply_rus(prep_b)()
return x
res_dict = main()
print(res_dict)
# {'gphase': 51, 's_dg': 70, 's': 120, 'z': 12, 'x': 127, 't': 298, 'u3': 157, 'cx': 1194, 'p': 118, 'ry': 2, 'h': 390, 't_dg': 348, 'measure': 116}
"""
j_0, beta = cks_params(eps, kappa, max_beta=max_beta)
cheb_coeffs = cks_coeffs(j_0, beta)
m = len(A._anc_templates)
# Following https://math.berkeley.edu/~linlin/qasc/qasc_notes.pdf (page 104):
# T1 = (U R), T2 = (U R U_dg R) -> T2^k T1 is block encoding T_{2k+1}(A)
# This is valid even if the block encoding unitary is not Hermitian.
# The first control-(U R) can be simplified to U since:
# 1) The first qubit of the |unary> state is always 1.
# 2) The R acts as identity on |0>.
def inv_unitary(*args):
with invert():
A.unitary(*args)
def T2(*args):
reflection(args[:m])
with conjugate(inv_unitary)(*args):
reflection(args[:m])
def new_unitary(*args):
# Core LCU protocol: PREP, SELECT, PREP^†
with conjugate(unary_prep)(args[0], cheb_coeffs):
A.unitary(*args[1:])
for i in jrange(1, j_0 + 1):
z(args[0][i])
with control(args[0][i]):
T2(*args[1:])
new_anc_templates = [QuantumFloat(j_0 + 1).template()] + A._anc_templates
new_alpha = np.sum(np.abs(cheb_coeffs)) / A.alpha
return BlockEncoding(new_alpha, new_anc_templates, new_unitary)
