Merge pull request #6 from tencent-quantum-lab/master

local sync

Author: PeilinZHENG

CHANGELOG.md

@@ -14,6 +14,10 @@
 
 - Add circuit copy method that avoid shallow copy issue `Circuit.copy()`
 
+### Fixed
+
+- improve the `adaptive_vmap` to support internal jit and pytree output
+
 ### Changed
 
 - The static method `BaseCircuit.copy` is renamed as `BaseCircuit.copy_nodes`

requirements/requirements-extra.txt

@@ -1,5 +1,5 @@
 # extra dependencies for ci
-qiskit
+qiskit==0.43
 qiskit-nature
 cirq
 torch

tensorcircuit/applications/physics/fss.py

@@ -1,7 +1,7 @@
 """
 finite size scaling tools
 """
-from typing import List, Tuple
+from typing import List, Tuple, Optional
 
 import numpy as np
 
@@ -14,6 +14,8 @@ def data_collapse(
     nu: float,
     beta: float = 0,
     obs_type: int = 1,
+    fit_type: int = 0,
+    dobs: Optional[List[List[float]]] = None,
 ) -> Tuple[List[float], List[List[float]], List[List[float]], float]:
     xL: List[List[float]] = []  # x=(p-pc)L^(1/\nv)
     yL: List[List[float]] = []  # y=S(p,L)-S(pc,L) or S(p,L)
@@ -33,27 +35,61 @@ def data_collapse(
                 yL[n0].append(obs[n0][p0] * n[n0] ** beta)
                 # tripartite mutual information
         pc_list.append(pc_L)
-    xL_all = []
-    for i in range(len(xL)):
-        xL_all.extend(xL[i])
-    yL_ave = []
-    loss = []
-    for x0 in range(len(xL_all)):
-        ybar_list = []
-        for n0 in range(len(n)):
-            if xL_all[x0] >= xL[n0][0] and xL_all[x0] <= xL[n0][-1]:
-                yinsert = pc_linear_interpolation(xL[n0], yL[n0], xL_all[x0])
-                ybar_list.append(yinsert)
+    if fit_type == 0:
+        xL_all = []
+        for i in range(len(xL)):
+            xL_all.extend(xL[i])
+        yL_ave = []
+        loss = []
+        for x0 in range(len(xL_all)):
+            ybar_list = []
+            for n0 in range(len(n)):
+                if xL_all[x0] >= xL[n0][0] and xL_all[x0] <= xL[n0][-1]:
+                    yinsert = pc_linear_interpolation(xL[n0], yL[n0], xL_all[x0])
+                    ybar_list.append(yinsert)
+
+            ybar = np.mean(ybar_list)
+            mean_squared = [(ybar_list[i] - ybar) ** 2 for i in range(len(ybar_list))]
+            loss.append(np.sum(mean_squared))
+            yL_ave.append(ybar)
+        loss = np.sum(loss)
 
-        ybar = np.mean(ybar_list)
-        mean_squared = [(ybar_list[i] - ybar) ** 2 for i in range(len(ybar_list))]
-        loss.append(np.sum(mean_squared))
-        yL_ave.append(ybar)
-    loss = np.sum(loss)
+        return pc_list, xL, yL, loss  # type: ignore
 
+    # fit_type == 1
+    if dobs is None:
+        raise ValueError("uncertainty of each y has to be specified in `dobs`")
+
+    datas = []
+    for n0 in range(len(n)):
+        for i in range(len(xL[n0])):
+            datas.append([xL[n0][i], yL[n0][i], dobs[n0][i]])
+    datas = sorted(datas, key=lambda x: x[0])
+    loss = _quality_objective_v2(datas)  # type: ignore
     return pc_list, xL, yL, loss  # type: ignore
 
 
+def _quality_objective_v2(datas: List[List[float]]) -> float:
+    #  https://journals.aps.org/prb/supplemental/10.1103/PhysRevB.101.060301/Supplement.pdf
+    loss = []
+    for i in range(len(datas) - 2):
+        # i, i+1, i+2
+        x, y, d = datas[i + 1]
+        x1, y1, d1 = datas[i]
+        x2, y2, d2 = datas[i + 2]
+        if np.abs(x - x1) < 1e-4 or np.abs(x - x2) < 1e-4:
+            continue
+        ybar = ((x2 - x) * y1 - (x1 - x) * y2) / (x2 - x1)
+        delta = (
+            d**2
+            + d1**2 * (x2 - x) ** 2 / (x2 - x1) ** 2
+            + d2**2 * (x1 - x) ** 2 / (x2 - x1) ** 2
+        )
+        w = (y - ybar) ** 2 / delta
+        loss.append(w)
+    return np.mean(loss)  # type: ignore
+
+
 def pc_linear_interpolation(p: List[float], SA: List[float], pc_input: float) -> float:
     if pc_input in p:
         pc_index = p.index(pc_input)

tensorcircuit/experimental.py

@@ -17,6 +17,7 @@
 def adaptive_vmap(
     f: Callable[..., Any],
     vectorized_argnums: Union[int, Sequence[int]] = 0,
+    static_argnums: Optional[Union[int, Sequence[int]]] = None,
     chunk_size: Optional[int] = None,
 ) -> Callable[..., Any]:
     if chunk_size is None:
@@ -46,7 +47,10 @@ def wrapper(*args: Any, **kws: Any) -> Tensor:
             reshape_args.append(arg)
             if s2 != 0:
                 rest_args.append(arg_rest)
-        _vmap = backend.vmap(f, vectorized_argnums)
+        _vmap = backend.jit(
+            backend.vmap(f, vectorized_argnums=vectorized_argnums),
+            static_argnums=static_argnums,
+        )
         r = []
         for i in range(s1):
             # currently using naive python loop for simplicity
@@ -55,16 +59,16 @@ def wrapper(*args: Any, **kws: Any) -> Tensor:
                 for j, a in enumerate(reshape_args)
             ]
             r.append(_vmap(*nreshape_args, **kws))
-        r = backend.stack(r)
-        rshape = list(backend.shape_tuple(r))
-        if len(rshape) == 2:
-            nshape = [rshape[0] * rshape[1]]
-        else:
-            nshape = [rshape[0] * rshape[1], -1]
-        r = backend.reshape(r, nshape)
+        r = backend.tree_map(lambda *x: backend.concat(x), *r)
+        # rshape = list(backend.shape_tuple(r))
+        # if len(rshape) == 2:
+        #     nshape = [rshape[0] * rshape[1]]
+        # else:
+        #     nshape = [rshape[0] * rshape[1], -1]
+        # r = backend.reshape(r, nshape)
         if s2 != 0:
             rest_r = _vmap(*rest_args, **kws)
-            return backend.concat([r, rest_r])
+            return backend.tree_map(lambda *x: backend.concat(x), r, rest_r)
         return r
 
     return wrapper