VQE on 1D TFIM#
Overview#
The main aim of this tutorial is not about the physics perspective of VQE, instead, we demonstrate the main ingredients of TensorCircuit by this simple VQE toy model.
Background#
Basically, we train a parameterized quantum circuit with repetitions of \(e^{i\theta} ZZ\) and \(e^{i\theta X}\) layers as \(U(\rm{\theta})\). And the objective to be minimized is this task is \(\mathcal{L}(\rm{\theta})=\langle 0^n\vert U(\theta)^\dagger H U(\theta)\vert 0^n\rangle\). The Hamiltonian is from TFIM as \(H = \sum_{i} Z_iZ_{i+1} -\sum_i X_i\).
Setup#
[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
To enable automatic differentiation support, we should set the TensorCircuit backend beyond the default one “NumPy”. And we can also set the high precision complex128 for the simulation.
[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 gate matrix to be utilized
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.]]
Higher-level API#
We first design the Hamiltonian energy expectation function with the input as quantum circuit.
[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)
Now we make the quantum function with \(\rm{\theta}\) as input and energy expectation \(\mathcal{L}\) as output.
[6]:
def vqe_tfim(param, n, nlayers):
c = tc.Circuit(n)
paramc = tc.backend.cast(
param, tc.dtypestr
) # We assume the input param with 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
To train the parameterized circuit, we should utilize the gradient information \(\frac{\partial \mathcal{L}}{\partial \rm{\theta}}\) with gradient descent. We also use jit to wrap the value and grad function for a substantial speedup. Note how (1, 2) args of vqe_tfim are labeled as static since they are just integers for qubit number and layer number instead of tensors.
[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>
Batched VQE Example#
We can even run a batched version of VQE optimization, namely, we simultaneously optimize parameterized circuits for different random initializations, so that we can try our best to avoid local minimums and locate the best of the converged energies.
[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])>
Different Backends#
We can change the backends at runtime without even changing one line of the code!
However, in normal user cases, we strongly recommend the users stick to one backend in one jupyter or python script. One can enjoy the facility provided by other backends by changing the set_backend line and running the same script again. This approach is much safer than using multiple backends in the same file unless you know the lower-level details of TensorCircuit enough.
[11]:
tc.set_backend("jax") # change to jax backend
[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)
Lower-level API#
The higher-level API under the namespace of TensorCircuit provides a unified framework to do linear algebra and automatic differentiation which is backend agnostic.
One may also use the related APIs (ops, AD-related, jit-related) directly provided by TensorFlow or Jax, as long as one is ok to stick with one fixed backend. See the tensorflow backend example below.
[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>