add templates submodule; add extra large vqe example

Author: refraction-ray

check_all.sh

@@ -7,7 +7,7 @@ mypy tensorcircuit
 echo "pylint check"
 pylint tensorcircuit tests
 echo "pytest check"
-pytest --cov=tensorcircuit -vv -W ignore::DeprecationWarning
+pytest -n 4 --cov=tensorcircuit -vv -W ignore::DeprecationWarning
 echo "sphinx check"
 cd docs && make html
-echo "all checks passed, congratulates! 💐"
+echo "all checks passed, congratulates! 💐"

examples/vqe_extra.py

@@ -0,0 +1,129 @@
+"""
+demonstration of TFIM VQE on V100 with lager qubit number counts
+"""
+
+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 tensorcircuit as tc
+from tensorcircuit import keras
+import cotengra as ctg
+
+optr = ctg.ReusableHyperOptimizer(
+    methods=["greedy", "kahypar"],
+    parallel=True,
+    minimize="flops",
+    max_time=120,
+    max_repeats=4096,
+    progbar=True,
+)
+tc.set_contractor("custom", optimizer=optr, preprocessing=True)
+# tc.set_contractor("custom_stateful", optimizer=oem.RandomGreedy, max_time=60, max_repeats=128, minimize="size")
+tc.set_dtype("complex64")
+tc.set_backend("tensorflow")
+dtype = np.complex64
+
+import time
+
+nwires, nlayers = 50, 7
+
+
+def vqe_forward(param, structures):
+    split_conf = {
+        "max_singular_values": 2,
+        "fixed_choice": 1,
+    }
+    structuresc = tc.backend.cast(structures, dtype="complex64")
+    paramc = tc.backend.cast(param, dtype="complex64")
+    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=paramc[2 * j, i],
+                unitary=tc.gates._zz_matrix,
+            )
+
+        for i in range(nwires):
+            c.rx(i, theta=paramc[2 * j + 1, i])
+
+    obs = []
+    for i in range(nwires):
+        obs.append(
+            [
+                tc.gates.Gate(
+                    sum(
+                        [
+                            structuresc[i, k] * g.tensor
+                            for k, g in enumerate(tc.gates.pauli_gates)
+                        ]
+                    )
+                ),
+                (i,),
+            ]
+        )
+    loss = c.expectation(*obs, reuse=False)
+    return tc.backend.real(loss)
+
+
+slist = []
+for i in range(nwires):
+    t = np.zeros(nwires)
+    t[i] = 1
+    slist.append(t)
+for i in range(nwires):
+    t = np.zeros(nwires)
+    t[i] = 3
+    t[(i + 1) % nwires] = 3
+    slist.append(t)
+structures = np.array(slist, dtype=np.int32)
+structures = tc.backend.onehot(structures, num=4)
+structures = tc.backend.reshape(structures, [-1, nwires, 4])
+print(structures.shape)
+time0 = time.time()
+
+tc_vag = tc.backend.jit(
+    tc.backend.vectorized_value_and_grad(vqe_forward, argnums=0, vectorized_argnums=1)
+)
+param = tf.Variable(tf.random.normal(stddev=0.1, shape=[2 * nlayers, nwires]))
+
+print(tc_vag(param, structures[:50]))
+
+time1 = time.time()
+print("staging time: ", time1 - time0)
+
+try:
+    keras.save_func(tc_vag, "./funcs/%s_%s_10_tfim" % (nwires, nlayers))
+except ValueError as e:
+    print(e)  # keras.save_func now has issues to be resolved
+
+
+def train_step(param):
+    vag_list = []
+    for i in range(2):
+        vag_list.append(tc_vag(param, structures[i * nwires : i * nwires + nwires]))
+    loss = tc.backend.sum(vag_list[0][0] - vag_list[1][0])
+    gr = vag_list[0][1] - vag_list[1][1]
+    return loss, gr
+
+
+if __name__ == "__main__":
+    opt = tf.keras.optimizers.Adam(0.02)
+    for j in range(5000):
+        loss, gr = train_step(param)
+        opt.apply_gradients([(gr, param)])
+        if j % 20 == 0:
+            print("loss", loss.numpy())

tensorcircuit/__init__.py

@@ -12,3 +12,4 @@
 from .densitymatrix import DMCircuit
 from .densitymatrix2 import DMCircuit2
 from .gates import num_to_tensor, array_to_tensor
+from . import templates

tensorcircuit/applications/utils.py

@@ -386,3 +386,20 @@ def repr2array(inputs: str) -> Array:
         o = [float(c.strip()) for c in l.split(" ") if c.strip()]
         outputs.append(o)
     return np.array(outputs)
+
+
+def TFIM1Denergy(
+    L: int, Jzz: float = 1.0, Jx: float = 1.0, Pauli: bool = True
+) -> float:
+    # PBC
+    # nice tutorial: https://arxiv.org/pdf/2009.09208.pdf
+    # further investigation on 1 TFIM solution structure is required
+    # will fail on AFM phase Jzz>Jx and Jzz>0 and odd sites (frustration in boundary)
+    e = 0
+    if Pauli:
+        Jx *= 2
+        Jzz *= 4
+    for i in range(L):
+        q = np.pi * (2 * i - (1 + (-1) ** L) / 2) / L
+        e -= np.abs(Jx) / 2 * np.sqrt(1 + Jzz ** 2 / 4 / Jx ** 2 - Jzz / Jx * np.cos(q))
+    return e

tensorcircuit/backends/jax_backend.py

@@ -45,7 +45,7 @@ def _svd_jax(
     max_truncation_error: Optional[float] = None,
     relative: Optional[bool] = False,
 ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
-    from .ops import adaware_svd
+    from .ops import adaware_svd_jit as adaware_svd
 
     left_dims = tensor.shape[:pivot_axis]
     right_dims = tensor.shape[pivot_axis:]

tensorcircuit/backends/ops.py

@@ -67,3 +67,5 @@ def jaxsvd_bwd(r: Sequence[Array], tangents: Sequence[Array]) -> Tuple[Array]:
 
 
 adaware_svd.defvjp(jaxsvd_fwd, jaxsvd_bwd)
+
+adaware_svd_jit = jax.jit(adaware_svd)

tensorcircuit/keras.py

@@ -141,6 +141,10 @@ def save_func(f: Callable[..., Any], path: str) -> None:
     tf.saved_model.save(m, path)
 
 
+# TODO(@refraction-ray): sometimes, save_func complains about ``return tuple(tensor.shape.as_list())``
+# ``ValueError: as_list() is not defined on an unknown TensorShape.``
+
+
 def load_func(
     *path: str, fallback: Optional[Callable[..., Any]] = None
 ) -> Callable[..., Any]:

tensorcircuit/simplify.py

@@ -78,7 +78,7 @@ def _split_two_qubit_gate(
     if (s1 >= 4) and (s2 >= 4):
         # jax jit unspport split_node with trun_err anyway
         # tf function doesn't work either, though I believe it may work on tf side
-        # TODO(@refraction-ray): tf.function version with trun_err set
+        # CANNOT DONE(@refraction-ray): tf.function version with trun_err set
         return None
     if s1 <= s2:  # equal is necessary for max values to pick on unswap one
         return n1, n2, False  # no swap

tensorcircuit/templates/__init__.py

@@ -0,0 +1 @@
+from . import measurements

tensorcircuit/templates/measurements.py

@@ -0,0 +1,53 @@
+"""
+shortcuts for measurement patterns on circuit
+"""
+
+from typing import Any
+
+from ..circuit import Circuit
+from ..cons import backend, dtypestr
+from .. import gates
+
+Tensor = Any
+
+
+def any_measurements(c: Circuit, structures: Tensor, onehot: bool = False) -> Tensor:
+    """
+    This measurements pattern is specifically suitable for vmap. Parameterize the Pauli string
+    to be measured
+
+    :param c: [description]
+    :type c: Circuit
+    :param structures: parameter tensors determines what Pauli string to be measured,
+        shape is [nwires, 4] if onehot is False.
+    :type structures: Tensor
+    :param onehot: [description], defaults to False
+    :type onehot: bool, optional
+    :return: [description]
+    :rtype: Tensor
+    """
+    if onehot is True:
+        structuresc = backend.cast(structures, dtype="int32")
+        structuresc = backend.onehot(structuresc, num=4)
+        structuresc = backend.cast(structuresc, dtype=dtypestr)
+    else:
+        structuresc = structures
+    nwires = c._nqubits
+    obs = []
+    for i in range(nwires):
+        obs.append(
+            [
+                gates.Gate(
+                    sum(
+                        [
+                            structuresc[i, k] * g.tensor
+                            for k, g in enumerate(gates.pauli_gates)
+                        ]
+                    )
+                ),
+                (i,),
+            ]
+        )
+    loss = c.expectation(*obs, reuse=False)  # type: ignore
+    # TODO(@refraction-ray): is reuse=True in this setup has user case?
+    return backend.real(loss)

tests/test_templates.py

@@ -0,0 +1,21 @@
+import sys
+import os
+import numpy as np
+
+thisfile = os.path.abspath(__file__)
+modulepath = os.path.dirname(os.path.dirname(thisfile))
+
+sys.path.insert(0, modulepath)
+import tensorcircuit as tc
+
+
+def test_any_measurement():
+    c = tc.Circuit(2)
+    c.H(0)
+    c.H(1)
+    mea = np.array([1, 1])
+    r = tc.templates.measurements.any_measurements(c, mea, onehot=True)
+    np.testing.assert_allclose(r, 1.0, atol=1e-5)
+    mea2 = np.array([3, 0])
+    r2 = tc.templates.measurements.any_measurements(c, mea2, onehot=True)
+    np.testing.assert_allclose(r2, 0.0, atol=1e-5)