add SHVQE script

Author: royess

examples/shvqe.py

@@ -0,0 +1,237 @@
+"""
+Schrodinger-Heisenberg quantum variational eigensolver (SHVQE) with DQAS-style optimization.
+
+DQAS part is modified from: examples/clifford_optimization.py
+"""
+
+import sys
+sys.path.insert(0, "../")
+
+import numpy as np
+import tensorflow as tf
+
+import tensorcircuit as tc
+from tensorcircuit.applications.vqes import construct_matrix_v3
+
+ctype, rtype = tc.set_dtype("complex64")
+K = tc.set_backend("tensorflow")
+
+n = 10 # the number of qubits (must be even for consistency later)
+ncz = 2 # number of cz layers in Schrodinger circuit
+nlayersq = ncz + 1 # Schrodinger parameter layers
+
+# training setup
+epochs = 1000
+batch = 1000
+
+# Hamiltonian
+h6h = np.load("./h6_hamiltonian.npy")  # reported in 0.99 A
+hamiltonian = construct_matrix_v3(h6h.tolist())
+
+def hybrid_ansatz(structure, paramq, preprocess="direct", train=True):
+    """_summary_
+
+    Parameters
+    ----------
+    structure : K.Tensor, (n//2, 2)
+        parameters to decide graph structure of Clifford circuits
+    paramq : K.Tensor, (nlayersq, n, 3)
+        parameters in quantum variational circuits, the last layer for Heisenberg circuits
+    preprocess : str, optional
+        preprocess, by default "direct"
+
+    Returns
+    -------
+    K.Tensor, [1,]
+        loss value
+    """
+    c = tc.Circuit(n)
+    if preprocess == "softmax":
+        structure = K.softmax(structure, axis=-1)
+    elif preprocess == "most":
+        structure = K.onehot(K.argmax(structure, axis=-1), num=2)
+    elif preprocess == "direct":
+        pass
+
+    structure = K.cast(structure, ctype)
+    structure = tf.reshape(structure, shape=[n//2, 2])
+
+    # quantum variational in Schrodinger part, first consider a ring topol
+    for j in range(nlayersq):
+        if j !=0 and j!=nlayersq-1:
+            for i in range(j%2,n,2):
+                c.cz(i, (i+1)%n)
+        for i in range(n):
+            c.rx(i, theta=paramq[j, i, 0])
+            c.ry(i, theta=paramq[j, i, 1])
+            c.rz(i, theta=paramq[j, i, 2])
+
+    # Clifford part, which is actually virtual
+    if train:
+        for j in range(0,n//2-1):
+            dis = j + 1
+            for i in range(0,n):
+                c.unitary(
+                    i,
+                    (i+dis) % n,
+                    unitary=structure[j, 0] * tc.gates.ii().tensor
+                    + structure[j, 1] * tc.gates.cz().tensor,
+                )
+
+        for i in range(0,n//2):
+            c.unitary(
+                i,
+                i + n//2,
+                unitary=structure[n//2-1, 0] * tc.gates.ii().tensor
+                + structure[n//2-1, 1] * tc.gates.cz().tensor,
+            )
+    else: # if not for training, we just put nontrivial gates
+        for j in range(0,n//2-1):
+            dis = j + 1
+            for i in range(0,n):
+                if structure[j, 1]==1:
+                    c.cz(i, (i+dis) % n)
+
+        for i in range(0,n//2):
+            if structure[j, 1]==1:
+                c.cz(i, i + n//2)
+
+    return c
+
+def hybrid_vqe(structure, paramq, preprocess="direct"):
+    """_summary_
+
+    Parameters
+    ----------
+    structure : K.Tensor, (n//2, 2)
+        parameters to decide graph structure of Clifford circuits
+    paramq : K.Tensor, (nlayersq, n, 3)
+        parameters in quantum variational circuits, the last layer for Heisenberg circuits
+    preprocess : str, optional
+        preprocess, by default "direct"
+
+    Returns
+    -------
+    K.Tensor, [1,]
+        loss value
+    """
+    c = hybrid_ansatz(structure, paramq, preprocess)
+    return tc.templates.measurements.operator_expectation(c, hamiltonian)
+
+def sampling_from_structure(structures, batch=1):
+    ch = structures.shape[-1]
+    prob = K.softmax(K.real(structures), axis=-1)
+    prob = K.reshape(prob, [-1, ch])
+    p = prob.shape[0]
+    r = np.stack(
+        np.array(
+            [np.random.choice(ch, p=K.numpy(prob[i]), size=[batch]) for i in range(p)]
+        )
+    )
+    return r.transpose()
+
+
+@K.jit
+def best_from_structure(structures):
+    return K.argmax(structures, axis=-1)
+
+
+def nmf_gradient(structures, oh):
+    """ compute the Monte Carlo gradient with respect of naive mean-field probabilistic model
+
+    Parameters
+    ----------
+    structures : K.Tensor, (n//2, ch)
+        structure parameter for single- or two-qubit gates
+    oh : K.Tensor, (n//2, ch), onehot
+        a given structure sampled via strcuture parameters (in main function)
+
+    Returns
+    -------
+    K.Tensor, (n//2 * 2, ch) == (n, ch)
+        MC gradients
+    """
+    choice = K.argmax(oh, axis=-1)
+    prob = K.softmax(K.real(structures), axis=-1)
+    indices = K.transpose(
+        K.stack([K.cast(tf.range(structures.shape[0]), "int64"), choice])
+    )
+    prob = tf.gather_nd(prob, indices)
+    prob = K.reshape(prob, [-1, 1])
+    prob = K.tile(prob, [1, structures.shape[-1]])
+
+    return K.real(
+        tf.tensor_scatter_nd_add(
+            tf.cast(-prob, dtype=ctype),
+            indices,
+            tf.ones([structures.shape[0]], dtype=ctype),
+        )
+    ) # in oh : 1-p, not in oh : -p
+
+# vmap for a batch of structures
+nmf_gradient_vmap = K.jit(
+    K.vmap(nmf_gradient, vectorized_argnums=1))
+
+# vvag for a batch of structures
+vvag_hybrid = K.jit(
+    K.vectorized_value_and_grad(hybrid_vqe, vectorized_argnums=(0,), argnums=(1,)),
+    static_argnums=(2,))
+
+def train_hybrid(stddev=0.05, lr=None, epochs=2000, debug_step=50, batch=256, verbose=False):
+    # params = K.implicit_randn([n//2, 2], stddev=stddev)
+    params = K.ones([n//2, 2], dtype=float)
+    paramq = K.implicit_randn([nlayersq, n, 3], stddev=stddev) * 2*np.pi
+    if lr is None:
+        lr = tf.keras.optimizers.schedules.ExponentialDecay(0.6, 100, 0.8)
+    structure_opt = K.optimizer(tf.keras.optimizers.Adam(lr))
+
+    avcost = 0
+    avcost2 = 0
+    loss_history = []
+    for epoch in range(epochs):  # iteration to update strcuture param
+        # random sample some structures
+        batched_stucture = K.onehot(
+            sampling_from_structure(params, batch=batch),
+            num=params.shape[-1],
+        )
+        vs, gq = vvag_hybrid(batched_stucture, paramq, "direct")
+        loss_history.append(np.min(vs))
+        gq = gq[0]
+        avcost = K.mean(vs) # average cost of the batch
+        gs = nmf_gradient_vmap(params, batched_stucture)  # \nabla lnp
+        gs = K.mean(K.reshape(vs - avcost2, [-1, 1, 1]) * gs, axis=0)
+        # avcost2 is averaged cost in the last epoch
+        avcost2 = avcost
+
+        [params, paramq] = structure_opt.update([gs, gq], [params, paramq])
+        if epoch % debug_step == 0 or epoch == epochs - 1:
+            print("----------epoch %s-----------" % epoch)
+            print(
+                "batched average loss: ",
+                np.mean(vs),
+                "minimum candidate loss: ",
+                np.min(vs),
+            )
+
+            # max over choices, min over layers and qubits
+            minp = tf.math.reduce_min(tf.math.reduce_max(tf.math.softmax(params), axis=-1))
+            if minp > 0.5:
+                print("probability converged")
+
+            if verbose:
+                print(
+                    "strcuture parameter: \n",
+                    params.numpy()
+                )
+
+            cand_preset = best_from_structure(params)
+            print(cand_preset)
+            print("current recommendation loss: ", hybrid_vqe(params, paramq, "most"))
+
+    loss_history = np.array(loss_history)
+    return hybrid_vqe(params, paramq, "most"), params, paramq, loss_history
+
+
+print('Train hybrid.')
+ee, params, paramq, loss_history = train_hybrid(epochs=epochs, batch=batch, verbose=True)
+print('Energy:', ee)