make mps benchmark uptodate

Author: refraction-ray

benchmarks/scripts/utils.py

@@ -6,6 +6,8 @@
 import numpy as np
 import tensorflow as tf
 from pathlib import Path
+import optax
+import tensorcircuit as tc
 
 
 qml_data = {}
@@ -218,12 +220,11 @@ def save(data, _uuid, path):
 
 def timing(f, nitrs, timeLimit):
     t0 = time.time()
-    print(f())
+    a = f()
     t1 = time.time()
     Nitrs = 1e-8
     for i in range(nitrs):
         a = f()
-        print(a)
         # if a != None:
         #    print(a)
         if time.time() - t1 > timeLimit:
@@ -263,15 +264,19 @@ def qml_timing(f, nbatch, nitrs, timeLimit, tfq=False):
 
 
 class Opt:
-    def __init__(self, f, params, lr=0.01, tuning=True):
+    def __init__(self, f, params, lr=0.002, tuning=True, backend="tensorflow"):
         self.f = f
         self.params = params
-        self.adam = tf.keras.optimizers.Adam(lr)
+        if backend == "tensorflow":
+            self.adam = tc.backend.optimizer(tf.keras.optimizers.Adam(lr))
+        else:
+            self.adam = tc.backend.optimizer(optax.adam(lr))
         self.tuning = tuning
 
     def step(self):
         e, grad = self.f(*self.params)
         if self.tuning:
-            grad = [tf.convert_to_tensor(g) for g in grad]
-            self.adam.apply_gradients(zip(grad, self.params))
+            self.params = self.adam.update(grad, self.params)
+            # grad = [tf.convert_to_tensor(g) for g in grad]
+            # self.adam.apply_gradients(zip(grad, self.params))
         return e[()]

benchmarks/scripts/vqe_tc.py

@@ -109,7 +109,7 @@ def energy_raw(paramx, paramzz):
             c = tc.Circuit(n)
         else:
             c = tc.MPSCircuit(n)
-            c.set_truncation_rule(max_singular_values=mpsd)
+            c.set_split_rules({"max_singular_values": mpsd})
         paramx = tc.backend.cast(paramx, dtype)
         paramzz = tc.backend.cast(paramzz, dtype)
 
@@ -132,8 +132,8 @@ def energy_raw(paramx, paramzz):
             e = tfi_energy(c, n)
             fd = c._fidelity
         # tensorflow only works for complex case, while jax only works for real case, don't know how to solve it
-        if tcbackend != "tensorflow":
-            e = tc.backend.real(e)
+        # if tcbackend != "tensorflow":
+        e = tc.backend.real(e)
 
         return e, fd
 
@@ -144,12 +144,13 @@ def energy(paramx, paramzz):
         # paramx = tc.backend.convert_to_tensor(paramx)
         # paramzz = tc.backend.convert_to_tensor(paramzz)
         (value, f), grads = energy_raw(paramx, paramzz)
-        print(tc.backend.numpy(f), tc.backend.numpy(value))
+        print("fidelity: ", f, value)
+        # print(tc.backend.numpy(f), tc.backend.numpy(value))
         # value = tc.backend.numpy(tc.backend.real(value))
         # grads = [tc.backend.numpy(tc.backend.real(g)) for g in grads]
         return value, grads
 
-    opt = utils.Opt(energy, [paramx, paramzz], tuning=False)
+    opt = utils.Opt(energy, [paramx, paramzz], tuning=True, backend=tcbackend)
     ct, it, Nitrs = utils.timing(opt.step, nitrs, timeLimit)
     meta["Results"]["with jit"] = {
         "Construction time": ct,

docs/source/advance.rst

@@ -7,6 +7,16 @@ MPS Simulator
 
 (Still experimental support)
 
+Very simple, we provide the same set of API for ``MPSCircuit`` as ``Circuit``, 
+the only new line is to set the bond dimension for the new simulator.
+
+.. code-block:: python
+
+    c = tc.MPSCircuit(n)
+    c.set_split_rules({"max_singular_values": 50})
+
+The larger bond dimension we set, the better approximation ratio (of course the most computational cost we pay)
+
 Split Two-qubit Gates
 -------------------------