一维横场 Ising 模型 VQE#

概述#

本教程的主要目的不是关于 VQE 物理层面的讨论,而是我们通过演示 这个简单的 VQE 玩具模型来了解张量电路的主要技术组件和用法。

背景#

基本上,我们训练一个参数化的量子电路,其线路结构为重复的 \(e^{i\theta} ZZ\)\(e^{i\theta X}\) 层的 \(U(\rm{\theta})\)。 而要最小化的目标是这个任务 \(\mathcal{L}(\rm{\theta})=\langle 0^n\vert U(\theta)^\dagger HU(\theta)\vert 0^n \rangle\)。 哈密顿量来自 TFIM,\(H = \sum_{i} Z_iZ_{i+1} -\sum_i X_i\)

设置#

[1]:
from functools import partial
import numpy as np
import tensorflow as tf
import jax
from jax.config import config

config.update("jax_enable_x64", True)
from jax import numpy as jnp
from jax.experimental import optimizers
import tensorcircuit as tc

为了启用自动微分支持,我们应该将 TensorCircuit 设置为非 “numpy” 后端。 而且我们还可以设置高精度 complex128 进行模拟。

[2]:
tc.set_backend("tensorflow")
tc.set_dtype("complex128")
[3]:
print(
    "complex dtype of simulation:",
    tc.dtypestr,
    "\nreal dtype of simulation:",
    tc.rdtypestr,
    "\nbackend package of simulation:",
    tc.backend.name,
)
complex dtype of simulation: complex128
real dtype of simulation: float64
backend package of simulation: tensorflow
[4]:
# 要使用的 zz 门矩阵
zz = np.kron(tc.gates._z_matrix, tc.gates._z_matrix)
print(zz)
[[ 1.  0.  0.  0.]
 [ 0. -1.  0. -0.]
 [ 0.  0. -1. -0.]
 [ 0. -0. -0.  1.]]

更高层的 API#

我们首先设计了以量子电路为输入的哈密顿能量期望函数。

[5]:
def tfi_energy(c: tc.Circuit, j: float = 1.0, h: float = -1.0):
    e = 0.0
    n = c._nqubits
    for i in range(n):
        e += h * c.expectation((tc.gates.x(), [i]))  # <X_i>
    for i in range(n - 1):  # OBC
        e += j * c.expectation(
            (tc.gates.z(), [i]), (tc.gates.z(), [(i + 1) % n])
        )  # <Z_iZ_{i+1}>
    return tc.backend.real(e)

现在我们以 \(\rm{\theta}\) 作为输入;并将能量期望 \(\mathcal{L}\) 作为输出来制作量子函数。

[6]:
def vqe_tfim(param, n, nlayers):
    c = tc.Circuit(n)
    paramc = tc.backend.cast(param, tc.dtypestr)  # 我们假设输入参数的 dtype 为 float64
    for i in range(n):
        c.H(i)
    for j in range(nlayers):
        for i in range(n - 1):
            c.exp1(i, i + 1, unitary=zz, theta=paramc[2 * j, i])
        for i in range(n):
            c.rx(i, theta=paramc[2 * j + 1, i])
    e = tfi_energy(c)
    return e

为了训练参数化电路,我们应该利用梯度下降的梯度信息 \(\frac{\partial \mathcal{L}}{\partial \rm{\theta}}\)。 我们还使用 jit 来包装 value 和 grad 函数以显著加快速度。 注意 vqe_tfim 的 (1, 2) args 是如何被标记为静态的,因为它们只是量子比特数和层数的整数,而不是张量。

[7]:
vqe_tfim_vag = tc.backend.jit(
    tc.backend.value_and_grad(vqe_tfim), static_argnums=(1, 2)
)
[8]:
def train_step_tf(n, nlayers, maxiter=10000):
    param = tf.Variable(
        initial_value=tf.random.normal(
            shape=[nlayers * 2, n], stddev=0.1, dtype=getattr(tf, tc.rdtypestr)
        )
    )
    opt = tf.keras.optimizers.Adam(1e-2)
    for i in range(maxiter):
        e, grad = vqe_tfim_vag(param, n, nlayers)
        opt.apply_gradients([(grad, param)])
        if i % 200 == 0:
            print(e)
    return e


train_step_tf(6, 3, 2000)
tf.Tensor(-5.2044235531710905, shape=(), dtype=float64)
tf.Tensor(-7.168460907768175, shape=(), dtype=float64)
tf.Tensor(-7.229007202330065, shape=(), dtype=float64)
tf.Tensor(-7.2368387790165105, shape=(), dtype=float64)
tf.Tensor(-7.246179597523659, shape=(), dtype=float64)
tf.Tensor(-7.262673966580785, shape=(), dtype=float64)
tf.Tensor(-7.286129364991173, shape=(), dtype=float64)
tf.Tensor(-7.291252895716095, shape=(), dtype=float64)
tf.Tensor(-7.2930457160020765, shape=(), dtype=float64)
tf.Tensor(-7.293225326335964, shape=(), dtype=float64)
[8]:
<tf.Tensor: shape=(), dtype=float64, numpy=-7.293297606006469>

批处理 VQE 示例#

我们甚至可以运行批量版本的 VQE 优化,即我们针对不同的随机初始化同时优化参数化电路,这样我们就可以尽量避免局部最小值,从而找到收敛能量的最佳值。

[9]:
vqe_tfim_vvag = tc.backend.jit(
    tc.backend.vectorized_value_and_grad(vqe_tfim), static_argnums=(1, 2)
)
[10]:
def batched_train_step_tf(batch, n, nlayers, maxiter=10000):
    param = tf.Variable(
        initial_value=tf.random.normal(
            shape=[batch, nlayers * 2, n], stddev=0.1, dtype=getattr(tf, tc.rdtypestr)
        )
    )
    opt = tf.keras.optimizers.Adam(1e-2)
    for i in range(maxiter):
        e, grad = vqe_tfim_vvag(param, n, nlayers)
        opt.apply_gradients([(grad, param)])
        if i % 200 == 0:
            print(e)
    return e


batched_train_step_tf(16, 6, 3, 2000)
tf.Tensor(
[-4.56780182 -5.32411397 -5.34948039 -5.49728838 -5.51974631 -4.89464895
 -5.23113926 -5.70097167 -5.4384308  -5.27898261 -4.73926061 -5.43748391
 -5.02246224 -4.46749643 -5.34320604 -5.29828815], shape=(16,), dtype=float64)
tf.Tensor(
[-7.15906597 -7.20867528 -7.16615816 -7.16164269 -7.15427498 -7.17176534
 -7.15677645 -7.19769858 -7.1876547  -7.17160745 -7.14313137 -7.16458417
 -7.12556993 -7.1043696  -7.17233218 -7.17955502], shape=(16,), dtype=float64)
tf.Tensor(
[-7.22332735 -7.28775096 -7.22854626 -7.28800389 -7.22006811 -7.2773814
 -7.22241623 -7.23446324 -7.23115651 -7.23081143 -7.25399986 -7.26564648
 -7.16463543 -7.27854832 -7.23574558 -7.28935649], shape=(16,), dtype=float64)
tf.Tensor(
[-7.23956454 -7.29093555 -7.23464822 -7.2914774  -7.22326999 -7.29014637
 -7.24891067 -7.2505597  -7.23879431 -7.23826618 -7.28737831 -7.29193732
 -7.22649018 -7.29136679 -7.25276205 -7.29214669], shape=(16,), dtype=float64)
tf.Tensor(
[-7.24561853 -7.29413883 -7.23950499 -7.29230127 -7.22749993 -7.29051998
 -7.28702174 -7.289441   -7.25016979 -7.26370483 -7.29320874 -7.29451577
 -7.22882824 -7.29213765 -7.27040912 -7.29358236], shape=(16,), dtype=float64)
tf.Tensor(
[-7.24997971 -7.294748   -7.25420008 -7.29271584 -7.24577837 -7.29082466
 -7.29171805 -7.29016935 -7.28645108 -7.29170429 -7.29499124 -7.29520514
 -7.23115011 -7.29305292 -7.28793637 -7.2949226 ], shape=(16,), dtype=float64)
tf.Tensor(
[-7.25300306 -7.29512508 -7.28240557 -7.29287622 -7.28264095 -7.29125472
 -7.29399162 -7.29066326 -7.29233232 -7.29290676 -7.29521188 -7.29530935
 -7.23475933 -7.29429836 -7.29053038 -7.29559969], shape=(16,), dtype=float64)
tf.Tensor(
[-7.25706762 -7.29527205 -7.29168082 -7.29292216 -7.29221412 -7.29183888
 -7.29474989 -7.29119684 -7.29306204 -7.29300514 -7.29525079 -7.29538907
 -7.24226472 -7.29544118 -7.29086393 -7.29576418], shape=(16,), dtype=float64)
tf.Tensor(
[-7.26218683 -7.29529443 -7.29438024 -7.29294969 -7.29426403 -7.29243576
 -7.29491073 -7.29171828 -7.29384631 -7.29301505 -7.29527304 -7.29545727
 -7.25772904 -7.29581457 -7.2910381  -7.29582192], shape=(16,), dtype=float64)
tf.Tensor(
[-7.26654138 -7.29529938 -7.29515899 -7.29297612 -7.29499748 -7.29281659
 -7.29500531 -7.29227399 -7.29455691 -7.29302087 -7.29529063 -7.29551808
 -7.2892432  -7.29593328 -7.29120759 -7.29585771], shape=(16,), dtype=float64)
[10]:
<tf.Tensor: shape=(16,), dtype=float64, numpy=
array([-7.29011428, -7.29530356, -7.29549915, -7.29300424, -7.29529224,
       -7.29296047, -7.29508027, -7.29270021, -7.29499648, -7.29302725,
       -7.29530574, -7.2955733 , -7.29593608, -7.29604964, -7.29138482,
       -7.29589095])>

不同的后端#

我们可以在运行时更改后端,甚至无需更改一行代码! 但是,在普通用户情况下,我们强烈建议用户在一个 jupyter 或 python 脚本中坚持使用一个后端。 通过更改set_backend行并再次运行相同的脚本,可以享受其他后端提供的便利。 这种方法比在同一个文件中使用多个后端更安全,除非你足够了解 TensorCircuit 的底层细节。

[11]:
tc.set_backend("jax")  # 更改为 jax 后端
[12]:
vqe_tfim_vvag = tc.backend.jit(
    tc.backend.vectorized_value_and_grad(vqe_tfim), static_argnums=(1, 2)
)


def batched_train_step_jax(batch, n, nlayers, maxiter=10000):
    key = jax.random.PRNGKey(42)
    param = jax.random.normal(key, shape=[batch, nlayers * 2, n]) * 0.1
    opt_init, opt_update, get_params = optimizers.adam(step_size=1e-2)
    opt_state = opt_init(param)

    def update(i, opt_state):
        param = get_params(opt_state)
        (value, gradient) = vqe_tfim_vvag(param, n, nlayers)
        return value, opt_update(i, gradient, opt_state)

    for i in range(maxiter):
        value, opt_state = update(i, opt_state)
        param = get_params(opt_state)
        if i % 200 == 0:
            print(value)
    return value


batched_train_step_jax(16, 6, 3, 2000)
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
[-5.67575948 -5.44768444 -5.7821556  -5.36699503 -5.00485098 -5.59416181
 -5.13421084 -5.70462279 -5.73699416 -5.25819658 -4.70729299 -5.82823766
 -5.69154358 -5.51112311 -5.46091316 -5.31649863]
[-7.16831387 -7.17873365 -7.21905991 -7.17714641 -7.21910053 -7.17729778
 -7.23594046 -7.1978075  -7.2311691  -7.18566164 -7.15141273 -7.1760751
 -7.20727055 -7.22174427 -7.15227955 -7.15343225]
[-7.24047827 -7.23486717 -7.26382185 -7.25267406 -7.23938877 -7.24135079
 -7.28655961 -7.24413064 -7.28070556 -7.24825735 -7.23400189 -7.25234153
 -7.25756263 -7.2505181  -7.22647645 -7.2589444 ]
[-7.28642159 -7.23707926 -7.28988032 -7.28627451 -7.28716418 -7.25068739
 -7.29122589 -7.2510777  -7.2906953  -7.25976327 -7.23891735 -7.29227009
 -7.28973637 -7.26238069 -7.245065   -7.29155041]
[-7.29198674 -7.24196434 -7.29188725 -7.29243688 -7.2926968  -7.26254168
 -7.29233808 -7.26729904 -7.29277165 -7.28066403 -7.24315235 -7.29344766
 -7.2920645  -7.26717433 -7.26959622 -7.29307748]
[-7.29320541 -7.27162341 -7.29245991 -7.2934821  -7.29360574 -7.27103573
 -7.29311302 -7.29213505 -7.29356392 -7.29162927 -7.24981922 -7.29384423
 -7.2930642  -7.27323089 -7.29156252 -7.293684  ]
[-7.29384132 -7.29130123 -7.29333191 -7.29442748 -7.29396942 -7.29011734
 -7.2936549  -7.2929216  -7.29430937 -7.29288465 -7.27758063 -7.29446428
 -7.2939575  -7.27958774 -7.29260011 -7.29408982]
[-7.29433271 -7.29273603 -7.29467579 -7.29507754 -7.29438517 -7.29298445
 -7.2940647  -7.29297896 -7.29494014 -7.29348052 -7.29142477 -7.29524322
 -7.29455267 -7.28904449 -7.29319333 -7.29437335]
[-7.29480095 -7.29294578 -7.29507043 -7.29530148 -7.29479745 -7.29390014
 -7.29439439 -7.29304737 -7.29520739 -7.29397645 -7.29276852 -7.2955135
 -7.29492853 -7.29145529 -7.29390729 -7.29462868]
[-7.29530498 -7.29299558 -7.29520643 -7.29537379 -7.29505127 -7.2945667
 -7.29466477 -7.29312996 -7.29529573 -7.29447499 -7.29327567 -7.29560796
 -7.29520605 -7.29179062 -7.29484071 -7.29489978]
[12]:
DeviceArray([-7.29575181, -7.29302602, -7.29529976, -7.29541094,
             -7.29517778, -7.29488194, -7.29487651, -7.29323608,
             -7.29532772, -7.29494698, -7.29369784, -7.29567791,
             -7.29534388, -7.29187906, -7.29536221, -7.29516005],            dtype=float64)

更低层的 API#

TensorCircuit 命名空间下的更高级别 API 提供了一个统一的框架来进行线性代数和自动微分,这与后端无关。 也可以使用 TensorFlow 或 Jax 直接提供的相关 API(ops、自动微分相关、可即时编译相关),只要坚持一个固定后端即可。 请参阅下面的 TensorFlow 后端示例。

[13]:
tc.set_backend("tensorflow")
[14]:
def tfi_energy(c: tc.Circuit, j: float = 1.0, h: float = -1.0):
    e = 0.0
    n = c._nqubits
    for i in range(n):
        e += h * c.expectation((tc.gates.x(), [i]))  # <X_i>
    for i in range(n - 1):  # OBC
        e += j * c.expectation(
            (tc.gates.z(), [i]), (tc.gates.z(), [(i + 1) % n])
        )  # <Z_iZ_{i+1}>
    return tf.math.real(e)


def vqe_tfim(param, n, nlayers):
    c = tc.Circuit(n)
    paramc = tf.cast(param, tf.complex128)
    for i in range(n):
        c.H(i)
    for j in range(nlayers):
        for i in range(n - 1):
            c.exp1(i, i + 1, unitary=zz, theta=paramc[2 * j, i])
        for i in range(n):
            c.rx(i, theta=paramc[2 * j + 1, i])
    e = tfi_energy(c)
    return e


@tf.function
def vqe_tfim_vag(param, n, nlayers):
    with tf.GradientTape() as tape:
        tape.watch(param)
        v = vqe_tfim(param, n, nlayers)
    grad = tape.gradient(v, param)
    return v, grad
[15]:
train_step_tf(6, 3, 2000)
tf.Tensor(-5.5454151788179376, shape=(), dtype=float64)
tf.Tensor(-7.167693061786028, shape=(), dtype=float64)
tf.Tensor(-7.254761404891117, shape=(), dtype=float64)
tf.Tensor(-7.290050014550046, shape=(), dtype=float64)
tf.Tensor(-7.29133881232428, shape=(), dtype=float64)
tf.Tensor(-7.2918048286324915, shape=(), dtype=float64)
tf.Tensor(-7.292590929769901, shape=(), dtype=float64)
tf.Tensor(-7.294195015132205, shape=(), dtype=float64)
tf.Tensor(-7.295013538531699, shape=(), dtype=float64)
tf.Tensor(-7.2951174084838835, shape=(), dtype=float64)
[15]:
<tf.Tensor: shape=(), dtype=float64, numpy=-7.295170441938537>