CUDA-Q で遊んでみる (4) — QAOA でピタゴラス数を求めてみる

目的

CUDA-Q の更に応用的な利用として、Higher Order Unconstrained Binary Optimization (HUBO もしくは HOBO) の概念を用いて QAOA を実行し、ピタゴラス数を求めてみたい。ピタゴラス数は

$$\begin{align*} x^2 + y^2 = z^2 \end{align*}$$

を満たすような正の整数 $(x, y, z)$ のこととする。

Higher Order Unconstrained Binary Optimization (HUBO)

大体この手の話では Quadratic Unconstrained Binary Optimization (QUBO) で定式化して QAOA なりをするのだけど、別にそんなことをしなくても計算できるので、直接 3 次以上の項を扱いたい。つまり、恐らく Rosenberg et al. の手法と言った感じで呼ばれる方法で補助量子ビットを追加して 3 次以上の項を 2 次の項に落とすのが定石ではあると思うが、それを用いない。

今回使う定式化は、$x_i \ ( 0 \leq i \leq 2)$ らをバイナリ変数とし、$x = 2^2 x_0 + 2^1 x_1 + 2^0 x_2 + 1$ で 1～8 までの整数を扱う。コスト項にピタゴラス数の条件を設定し、更に制約条件で $x \neq 1$, $y \neq 1$, $z \neq 1$ および $x \neq y$ を課している:

$$\begin{align*} H =& (x^2 + y^2 - z^2)^2 + \\ &+ \lambda \{(1 - x_0)(1 - x_1)(1 - x_2) + \\ &\ \ \ \ \ \ \ \ \ (1 - y_0)(1 - y_1)(1 - y_2) + \\ &\ \ \ \ \ \ \ \ \ (1 - z_0)(1 - z_1)(1 - z_2) + \\ &\ \ \ \ \ \ \ \ \ (1 - (x_0 - y_0)^2)(1 - (x_1 - y_1)^2)(1 - (x_2 - y_2)^2) \} \end{align*}$$

実はこの問題は準最適解が結構多いため解きにくい。つまり $x^2 + y^2 - z^2 = 1$ や $x^2 + y^2 - z^2 = 2$ が多いのである。よって制約項を設定しないとかなり際どい状態になる。

そしてこの制約項によって 6 次式になる。こういうものを扱う場合を、Higher Order を含むので HUBO と呼ぶ。これを QUBO に変換 (reduce) するのはとてもしんどい。そんなことをしなければ $x_0, x_1, \ldots, z_2$ の 9 変数で済む。

シミュレーティッドアニーリングで解いてみる

実はこの問題は「HOBO対応の疑似量子アニーリングパッケージ」である hobotan の サンプルコード2 にヒントを得ている。そして hobotan の大部分の機能は TYTAN SDK で遊んでみる (1) — 入門の入門 で触れた tytan に取り込まれているので、今回は tytan を用いる。

まずは必要なパッケージをインストールする

CUDA-Q を含め必要なものをインストールする。

!pip install -qU cudaq "cupy-cuda12x==13.6.0" tytan

アスキーアートの回路図がちゃんと表示されるようにする

ユーティリティを定義する。

from IPython.display import HTML, display

def show_fixed(text, font="Consolas, Roboto Mono, monospace", size=13):
    text = text.expandtabs(4)   # タブをスペースに変換
    esc = (text.replace("&", "&")
               .replace("<", "&lt;")
               .replace(">", "&gt;"))
    html = f'<pre style="font-family:{font}; font-size:{size}px; white-space:pre; font-variant-ligatures:none;">{esc}</pre>'
    display(HTML(html))

tytan でピタゴラス数を求める

hobotan のサンプルコードを少し書き換えて以下のようにした。

%%time

import numpy as np
from tytan import (
    symbols_list,
    Compile,
    sampler,
    Auto_array,
)

def to_symbols_nbit(xs):
    return sum([2**(len(xs)-i-1)*xs[i] for i in range(len(xs))])

#量子ビットを2進数表現で用意
xs = symbols_list([3], 'x{}')
ys = symbols_list([3], 'y{}')
zs = symbols_list([3], 'z{}')
x = to_symbols_nbit(xs) + 1
y = to_symbols_nbit(ys) + 1
z = to_symbols_nbit(zs) + 1

lam = 100
# x != 1, y != 1, z != 1
penalties = np.prod(1 - xs) + np.prod(1 - ys) + np.prod(1 - zs)
# x != y
penalties += np.prod(1 - (xs - ys)**2)

#ピタゴラス条件
H = (x**2 + y**2 - z**2)**2 + lam * penalties

#HOBOテンソルにコンパイル
hobo, offset = Compile(H,).get_hobo()
print(f'offset\n{offset}')

#サンプラー選択
solver = sampler.MIKASAmpler()

#サンプリング
result = solver.run(hobo, shots=1000)

#上位10件
for r in result[:10]:
    print(f'Energy {r[1] + offset}, Occurrence {r[2]}')

    #さくっと10進数に戻す
    print('x =', Auto_array(r[0]).get_nbit_value(x))
    print('y =', Auto_array(r[0]).get_nbit_value(y))
    print('z =', Auto_array(r[0]).get_nbit_value(z))

offset
401.0
MODE: GPU
DEVICE: cuda:0
Energy 0.0, Occurrence 376
x = 3.0
y = 4.0
z = 5.0
Energy 0.0, Occurrence 624
x = 4.0
y = 3.0
z = 5.0
CPU times: user 9.21 s, sys: 633 ms, total: 9.84 s
Wall time: 13.7 s

期待する解を得た。

以下ではこれを QAOA で実装する。

HUBO 定式化を用いた QAOA でピタゴラス数を求める

ユーティリティの準備

まずは HUBO 式を量子回路に埋め込みたいのでユーティリティを準備する。tytan のコードを少し書き換えて用意する。内容の概要はコードに下に簡単に記載する。

実はこのユーティリティの QUBO 版は QAOA を眺めてみる (3) ― グラフカラーリング問題と QAOA で実装したが、今回は HUBO 版を実装している。

import symengine
from symengine import Symbol, Expr
from tytan import symbols
from tytan.compile import replace_function


Hobo = dict[tuple[int, ...], float]
Ising = dict[tuple[int, ...], float]


def get_hobo(expr):
    """
    get hobo data
    Raises:
        TypeError: Input type is symengine.
    Returns:
        [hobo, index_map], offset. hobo is dict from indices to weight.
    """
    #symengine型のサブクラス
    if 'symengine.lib' in str(type(expr)):
        offset = 0.0
        #式を展開して同類項をまとめる
        expr = symengine.expand(expr)

        for term in expr.args:
            # 定数項はオフセットに追加
            if term.is_Number:
                offset += float(term)
                continue

        #二乗項を一乗項に変換
        expr = replace_function(expr, lambda e: isinstance(e, symengine.Pow) and e.exp == 2, lambda e, *args: e)

        #もう一度同類項をまとめる
        expr = symengine.expand(expr)

        #文字と係数の辞書
        coeff_dict = expr.as_coefficients_dict()
        # print(coeff_dict)

        #定数項を消す　{1: 25} 必ずある
        if 1 in coeff_dict:
            del coeff_dict[1]
        # print(coeff_dict)

        #シンボル対応表
        # 重複なしにシンボルを抽出
        keys = list(set(sum([str(key).split('*') for key in coeff_dict.keys()], [])))
        # print(keys)

        # 要素のソート（ただしアルファベットソート）
        keys.sort()
        # print(keys)

        # シンボルにindexを対応させる
        index_map = {key:i for i, key in enumerate(keys)}
        # print(index_map)

        #量子ビット数
        num = len(index_map)
        # print(num)

        if False:
            #HOBO行列生成
            hobo = np.zeros(num ** ho, dtype=float).reshape([num] * ho)
            for key, value in coeff_dict.items():
                qnames = str(key).split('*')
                indices = sorted([index_map[qname] for qname in qnames])
                indices = [indices[0]] * (ho - len(indices)) + indices
                hobo[tuple(indices)] = float(value)
        else:
            hobo: dict[tuple[int, ...], float] = {}
            for key, value in coeff_dict.items():
                qnames = str(key).split('*')
                indices = sorted([index_map[qname] for qname in qnames])
                indices = tuple(sorted(indices))
                hobo[tuple(indices)] = float(value)

        return [hobo, index_map], offset

    else:
        raise TypeError("Input type must be symengine.")


# calc keys for double sort
def _calc_key(num_qubits: int, key: tuple[int, ...]) -> int:
    value = 0
    for i, num in enumerate(key):
        value += num * num_qubits**i
    return value


def get_ising(
    hobo: Hobo, num_qubits: int
) -> tuple[Ising, float]:
    ising_dict: Ising = {}
    offset = 0.0

    for indices, weight in hobo.items():
        expr = 1
        names = [f"z{i}" for i in indices]
        z_tuple = symbols(",".join(names))
        if isinstance(z_tuple, Symbol):  # single Z case
            z_tuple = tuple([z_tuple])
        for z in z_tuple:
            expr *= (1 - z) / 2
        expr = weight * expr

        expr = expr.expand()

        for term in expr.args:
            # 定数項はオフセットに追加
            if term.is_Number:
                offset += float(term)
                continue

        #二乗項を一乗項に変換
        expr = replace_function(expr, lambda e: isinstance(e, symengine.Pow) and e.exp == 2, lambda e, *args: e)

        #もう一度同類項をまとめる
        expr = expr.expand()
        coeff_dict = expr.as_coefficients_dict()
        if 1 in coeff_dict:
            del coeff_dict[1]

        for key, value in coeff_dict.items():
            znames = str(key).split('*')
            indices = tuple(sorted([int(zname[1:]) for zname in znames]))
            ising_dict.setdefault(indices, 0.0)
            ising_dict[indices] += float(value)

    if False:
        ising_dict = {k: v for k, v in ising_dict.items() if not np.isclose(v, 0)}
    ising_dict = dict(
        sorted(ising_dict.items(), key=lambda k_v: _calc_key(num_qubits, k_v[0]))
    )
    return ising_dict, offset


if True:
    test_var = symbols_nbit(0, 4, 'x{}', num=2)
    test_H = test_var**2
    print(test_H)

    test_hobo, test_offset = get_hobo(test_H)
    test_ising = get_ising(test_hobo[0], len(test_hobo[1]))
    print(test_ising)

(2.0*x0 + 1.0*x1)**2
({(0,): -3.0, (1,): -1.5, (0, 1): 1.0}, 3.5)

何をやっているかと言うと、バイナリ変数 $x \in \{0, 1\}$ をイジング変数 $z = 1 - 2x \in \{-1, 1\}$ で書き換えた上で、$3 z_0 z_1 z_2$ のような項を {(0, 1, 2): 3, ... という辞書に変換して返す実装をしている。1 行目の print の結果を

$$\begin{align*} (2 x_0 + x_1)^2 &= 4 x_0^2 + 4 x_0 x_1 + x_1^2 \\ &= 4 x_0 + 4 x_0 x_1 + x_1 \\ &= 4 \frac{1 - z_0}{2} + 4 \frac{1 - z_0}{2} \frac{1 - z_1}{2} + \frac{1 - z_1}{2} \\ &= 2 - 2 z_0 + 1 - z_0 - z_1 + z_0 z_1 + \frac{1}{2} - \frac{1}{2} z_1 \\ &= -3 z_0 - \frac{3}{2} z_1 + z_0 z_1 + \frac{7}{2} \end{align*}$$

と計算して、係数をまとめると 2 行目の print に対応する。

イジング変数を用いた $3 z_0 z_1 z_2$ などはハミルトニアンとしては $3 Z_0 Z_1 Z_2$ となり、量子回路の中ではこれの回転ゲートとして $RZ^{\otimes 3}$ として埋め込まれる。この辺の技術的な詳細は RZZZ ゲートの分解 (1) や RZZZ ゲートの分解 (2) に書いていた。

HUBO 定式化をイジング式に変換する

API 名はさておき、上記で用意したユーティリティを用いて、HUBO の情報を量子回路に埋め込む準備をする。因みに、求める解は、コスト項を 0 にし、制約項も 0 にするので、一旦横に置いている定数項の和と相殺する。つまり、offset + ising[1] という値をマイナスにしたものが最適化で得られて欲しい最低エネルギーになる。

num_qubits = 9

hobo, offset = get_hobo(H)
ising = get_ising(hobo[0], num_qubits)

print(f"expected optimal_expectation=", -(offset + ising[1]))

expected optimal_expectation= -2039.0

QAOA を実装する

まずは CUDA-Q 周辺のパッケージをインストールして、GPU を有効にする。今回は cuStateVec を使う形にする。

import cudaq
from cudaq import spin
from scipy.optimize import minimize

if cudaq.num_available_gpus() > 0:
    cudaq.set_target('nvidia')

CUDA-Q カーネルの定義

以下のように CUDA-Q カーネルを定義する。基本的には CUDA-Q の Quantum Approximate Optimization Algorithm と Molecular docking via DC-QAOA を参考にしている。

num_qubits = 9


def ham_pythagoras(H, num_qubits) -> cudaq.SpinOperator:
    spin_ham = 0

    hobo, offset = get_hobo(H)
    ising = get_ising(hobo[0], num_qubits)
    for indices, coeff in ising[0].items():
        term = coeff
        for i in indices:
            term *= spin.z(i)
        spin_ham += term

    return spin_ham


# Collect coefficients from a spin operator so we can pass them to a kernel
def term_coefficients(hamiltonian : cudaq.SpinOperator) -> list[complex]:
    result = []
    for term in hamiltonian:
        result.append(term.evaluate_coefficient())
    return result

    # Collect Pauli words from a spin operator so we can pass them to a kernel


def term_words(hamiltonian : cudaq.SpinOperator) -> list[str]:
    # Our kernel uses these words to apply exp_pauli to the entire state.
    # we hence ensure that each pauli word covers the entire space.

    result = []
    for term in hamiltonian:
        result.append(term.get_pauli_word(num_qubits))
    return result

以下では Molecular docking via DC-QAOA を参考にちょっとインチキをして、純粋な QAOA ではない実装を採用している。

The increase in parameters is hopefully offset by requiring fewer layers.

とあるように、パラメータ数を増やして、代わりにブロックの繰り返し数を抑えている。実はこの方法をやらないと今回は繰り返し数 30 でも期待する解からは程遠い状態だったのである。

@cudaq.kernel
def kernel_qaoa(
    num_qubits: int, num_layers: int, thetas: list[float],
    coef: list[complex], words: list[cudaq.pauli_word]
):
    qubits = cudaq.qvector(num_qubits)

    h(qubits)

    count = 0
    for p in range(num_layers):
        # Problem unitary
        for i in range(len(coef)):
            exp_pauli(thetas[count] * coef[i].real, qubits, words[i])
            count += 1

        # Mixer unitary
        for j in range(num_qubits):
            rx(thetas[count], qubits[j])
            count += 1


try:
    kernel_qaoa.compile()
except Exception as e:
    print(e)

QAOA の準備

ハミルトニアンを定義して、QAOA の初期パラメータを決める。

hamiltonian = ham_pythagoras(H, num_qubits)
coef = term_coefficients(hamiltonian)
words = term_words(hamiltonian)


layer_count: int = 3
parameter_count: int = (num_qubits + len(coef)) * layer_count

rng = np.random.default_rng(42)
initial_parameters = rng.random(parameter_count) * np.pi

最適化を実行する

%%time

import sys
import time

# Define the objective, return `<state(params) | H | state(params)>`
def objective(parameters, *args):
    (notify_loss,) = args
    energy = cudaq.observe(kernel_qaoa, hamiltonian, num_qubits, layer_count,
                           parameters, coef, words).expectation()

    if notify_loss is not None:
        notify_loss(energy, parameters)

    return energy


min_energy = sys.maxsize
count = 0
losses = []
start = time.time()


def notify_loss(energy: float, params: np.ndarray):
    global min_energy
    global count
    global losses
    global start

    losses.append(energy)

    if energy < min_energy:
        min_energy = energy

    count += 1


method="COBYLA"
bounds = None
maxiter = 50000

result = minimize(
    objective,
    initial_parameters,
    args=(notify_loss),
    method=method,
    bounds = bounds,
    options={
        "maxiter": maxiter
    },
)

optimal_parameters = result.x.tolist()
optimal_expectation = result.fun
if isinstance(optimal_expectation, np.ndarray):
    optimal_expectation = optimal_expectation.item()


print('optimal_expectation =', optimal_expectation)
print('optimal_parameters =', optimal_parameters)

optimal_expectation = -1550.5032373344113
optimal_parameters = [2.4322273127907943, 2.3785082548084837, 3.697146928303546, ...]
CPU times: user 1h 39min 31s, sys: 2min 50s, total: 1h 42min 21s
Wall time: 58min 2s

エネルギーの目標値は -2039.0 だったのでまだまだなのだが、今回はこの程度で妥協する。

結果を確認する

最適化の状況の確認

念のため最適化の状況を確認する。恐らく標準的なグラフであろう。

import matplotlib.pyplot as plt

x_values = list(range(len(losses)))
y_values = losses

plt.plot(x_values, y_values)

plt.xlabel("Epochs")
plt.ylabel("Cost Value")
plt.show()

サンプリングする

以下は最適化が不十分な状態で解を求めるため相当なインチキだが、何度かガチャをしてそこそこ上の方で解が求まったケースを掲載する。

result = cudaq.sample(kernel_qaoa, num_qubits, layer_count, optimal_parameters, coef, words, shots_count=10000)
result: dict[str, int] = {k: v for k, v in result.items()}
for i, (k, v) in enumerate(list(sorted(result.items(), key=lambda k_v: -k_v[1]))[:50]):
    k_ = [int(v) for v in list(k)]
    x_ = 4 * k_[0] + 2 * k_[1] + k_[2] + 1
    y_ = 4 * k_[3] + 2 * k_[4] + k_[5] + 1
    z_ = 4 * k_[6] + 2 * k_[7] + k_[8] + 1
    diff = abs(x_**2 + y_**2 - z_**2)
    if diff < 3:
        print(f"[{i}] x={x_} y={y_} z={z_} diff={diff}", v)

[14] x=2 y=2 z=3 diff=1 105
[31] x=3 y=4 z=5 diff=0 75
[32] x=1 y=1 z=1 diff=1 75
[37] x=8 y=1 z=8 diff=1 72
[41] x=5 y=3 z=6 diff=2 66
[45] x=7 y=1 z=7 diff=1 62

ということで、$3^2 + 4^2 = 5^2$ が成立するピタゴラス数 $(3, 4, 5)$ が求まった。

エネルギーを -1000 程度までしか下げられていなかった時は、もっと順位が低いところでしか求まらなかったので、これでもかなりマシになったほうである・・・。

まとめ

大分インチキをしているが、CUDA-Q を用いた QAOA をやってみた。かなり苦しい内容だが、HUBO を用いてピタゴラス数を求めた。

恐らくブロックの繰り返し数を 4 や 5 にすると、よりエネルギーを下げられる可能性はあるが、ソルバが耐えるか不明だし、思ったように最適化ができるかも分からない。

参考文献

• CUDA-Q 関連
  • NVIDIA、CUDA-Q プラットフォームで、世界中の量子コンピューティング センターを加速
  • NVIDIA CUDA-Q
  • Quick Start
  • cuda-quantum
  • CUDA-Q Language Specification
  • 13. Example Programs
• CUDA-Q QAOA 関連
  • Quantum Approximate Optimization Algorithm
  • Molecular docking via DC-QAOA
• cuQuantum 関連
  • cuQuantum
  • cuStateVec
• 疑似量子アニーリング関連
  • hobotan
  • tytan