add extra size vqe example using mpo hamiltonian

Author: refraction-ray

examples/vqe_extra_mpo.py

@@ -0,0 +1,146 @@
+"""
+Demonstration of TFIM VQE with extra size in MPO formulation
+"""
+
+import time
+import logging
+import sys
+import numpy as np
+
+logger = logging.getLogger("tensorcircuit")
+logger.setLevel(logging.INFO)
+ch = logging.StreamHandler()
+ch.setLevel(logging.DEBUG)
+logger.addHandler(ch)
+
+sys.setrecursionlimit(10000)
+
+import tensorflow as tf
+import tensornetwork as tn
+import cotengra as ctg
+
+import optax
+import tensorcircuit as tc
+
+opt = ctg.ReusableHyperOptimizer(
+    methods=["greedy", "kahypar"],
+    parallel="ray",
+    minimize="combo",
+    max_time=360,
+    max_repeats=4096,
+    progbar=True,
+)
+
+
+def opt_reconf(inputs, output, size, **kws):
+    tree = opt.search(inputs, output, size)
+    tree_r = tree.subtree_reconfigure_forest(
+        parallel="ray",
+        progbar=True,
+        num_trees=20,
+        num_restarts=20,
+        subtree_weight_what=("size",),
+    )
+    return tree_r.path()
+
+
+tc.set_contractor("custom", optimizer=opt_reconf, preprocessing=True)
+tc.set_dtype("complex64")
+tc.set_backend("tensorflow")
+# jax backend is incompatible with keras.save
+
+dtype = np.complex64
+
+nwires, nlayers = 150, 7
+
+
+Jx = np.array([1.0 for _ in range(nwires - 1)])  # strength of xx interaction (OBC)
+Bz = np.array([-1.0 for _ in range(nwires)])  # strength of transverse field
+hamiltonian_mpo = tn.matrixproductstates.mpo.FiniteTFI(
+    Jx, Bz, dtype=dtype
+)  # matrix product operator
+hamiltonian_mpo = tc.quantum.tn2qop(hamiltonian_mpo)
+
+
+def vqe_forward(param):
+    print("compiling")
+    split_conf = {
+        "max_singular_values": 2,
+        "fixed_choice": 1,
+    }
+    c = tc.Circuit(nwires, split=split_conf)
+    for i in range(nwires):
+        c.H(i)
+    for j in range(nlayers):
+        for i in range(0, nwires - 1):
+            c.exp1(
+                i,
+                (i + 1) % nwires,
+                theta=param[4 * j, i],
+                unitary=tc.gates._xx_matrix,
+            )
+
+        for i in range(nwires):
+            c.rz(i, theta=param[4 * j + 1, i])
+        for i in range(nwires):
+            c.ry(i, theta=param[4 * j + 2, i])
+        for i in range(nwires):
+            c.rz(i, theta=param[4 * j + 3, i])
+    return tc.templates.measurements.mpo_expectation(c, hamiltonian_mpo)
+
+
+if __name__ == "__main__":
+    refresh = False
+
+    time0 = time.time()
+    if refresh:
+        tc_vag = tf.function(
+            tc.backend.value_and_grad(vqe_forward),
+            input_signature=[tf.TensorSpec([4 * nlayers, nwires], tf.float32)],
+        )
+        tc.keras.save_func(tc_vag, "./funcs/%s_%s_tfim_mpo" % (nwires, nlayers))
+        time1 = time.time()
+        print("staging time: ", time1 - time0)
+
+    tc_vag_loaded = tc.keras.load_func("./funcs/%s_%s_tfim_mpo" % (nwires, nlayers))
+
+    lr1 = 0.008
+    lr2 = 0.06
+    steps = 2000
+    switch = 400
+    debug_steps = 20
+
+    if tc.backend.name == "jax":
+        opt = tc.backend.optimizer(optax.adam(lr1))
+        opt2 = tc.backend.optimizer(optax.sgd(lr2))
+    else:
+        opt = tc.backend.optimizer(tf.keras.optimizers.Adam(lr1))
+        opt2 = tc.backend.optimizer(tf.keras.optimizers.SGD(lr2))
+
+    times = []
+    param = tc.backend.implicit_randn(stddev=0.1, shape=[4 * nlayers, nwires])
+
+    for j in range(steps):
+        loss, gr = tc_vag_loaded(param)
+        if j < switch:
+            param = opt.update(gr, param)
+        else:
+            if j == switch:
+                print("switching the optimizer")
+            param = opt2.update(gr, param)
+        if j % debug_steps == 0 or j == steps - 1:
+            times.append(time.time())
+            print("loss", tc.backend.numpy(loss))
+            if j > 0:
+                print("running time:", (times[-1] - times[0]) / j)
+
+
+"""
+# Baseline code: obtained from DMRG using quimb
+
+import quimb
+
+h = quimb.tensor.tensor_gen.MPO_ham_ising(nwires, 4, 2, cyclic=False)
+dmrg = quimb.tensor.tensor_dmrg.DMRG2(m, bond_dims=[10, 20, 100, 100, 200], cutoffs=1e-10)
+dmrg.solve(tol=1e-9, verbosity=1) # may require repetition of this API
+"""