mpscircuit now try first convert quvector directly into mps if possible

Author: refraction-ray

CHANGELOG.md

@@ -18,6 +18,8 @@
 
 - Add `ode_evol_local` and `ode_evol_global` methods for local and global ODE evolution.
 
+- Add transformation method between tensornetwork, quimb, tenpy and QuOperator in tc-ng including `qop2tenpy`, `qop2quimb`, `qop2tn`, `tenpy2qop`, support both MPS and MPO formats.
+
 ### Fixed
 
 - Fixed `one_hot` in numpy backend.
@@ -26,6 +28,8 @@
 
 - Fix potential np.matrix return from `PaulistringSum2Dense`.
 
+- `MPSCircuit` now will first try to transform `QuVector` input to tensors directly instead of evaluating it to dense wavefunction first.
+
 ### Changed
 
 - The order of arguments of `tc.timeevol.ed_evol` are changed for consistent interface with other evolution methods.

examples/tenpy_sz_convention.py

@@ -120,7 +120,7 @@
 
 print("--- Scenario 3: Correcting XXZChain DMRG state with X-gates ---")
 
-L = 10
+L = 30
 xxz_model_params = {"L": L, "bc_MPS": "finite", "Jxx": 1.0, "Jz": 1.0, "conserve": None}
 xxz_M = XXZChain(xxz_model_params)
 psi0_xxz = MPS.from_product_state(

tensorcircuit/mpscircuit.py

@@ -4,21 +4,24 @@
 
 # pylint: disable=invalid-name
 
-from functools import reduce
+from functools import reduce, partial
 from typing import Any, List, Optional, Sequence, Tuple, Dict, Union
 from copy import copy
+import logging
 
 import numpy as np
 import tensornetwork as tn
-from tensorcircuit.quantum import QuOperator, QuVector
 
 from . import gates
 from .cons import backend, npdtype, contractor, rdtypestr, dtypestr
+from .quantum import QuOperator, QuVector, extract_tensors_from_qop
 from .mps_base import FiniteMPS
 from .abstractcircuit import AbstractCircuit
+from .utils import arg_alias
 
 Gate = gates.Gate
 Tensor = Any
+logger = logging.getLogger(__name__)
 
 
 def split_tensor(
@@ -77,6 +80,10 @@ class MPSCircuit(AbstractCircuit):
 
     is_mps = True
 
+    @partial(
+        arg_alias,
+        alias_dict={"wavefunction": ["inputs"]},
+    )
     def __init__(
         self,
         nqubits: int,
@@ -118,8 +125,19 @@ def __init__(
             ), "tensors and wavefunction cannot be used at input simutaneously"
             # TODO(@SUSYUSTC): find better way to address QuVector
             if isinstance(wavefunction, QuVector):
-                wavefunction = wavefunction.eval()
-            tensors = self.wavefunction_to_tensors(wavefunction, split=self.split)
+                try:
+                    nodes, is_mps, _ = extract_tensors_from_qop(wavefunction)
+                    if not is_mps:
+                        raise ValueError("wavefunction is not a valid MPS")
+                    tensors = [node.tensor for node in nodes]
+                except ValueError as e:
+                    logger.warning(repr(e))
+                    wavefunction = wavefunction.eval()
+                    tensors = self.wavefunction_to_tensors(
+                        wavefunction, split=self.split
+                    )
+            else:  # full wavefunction
+                tensors = self.wavefunction_to_tensors(wavefunction, split=self.split)
             assert len(tensors) == nqubits
             self._mps = FiniteMPS(tensors, canonicalize=False)
             self._mps.center_position = 0

tensorcircuit/quantum.py

@@ -1181,7 +1181,6 @@ def extract_tensors_from_qop(qop: QuOperator) -> Tuple[List[Node], bool, int]:
     # Find endpoint nodes
     endpoint_nodes = set()
     physical_edges = set(qop.out_edges) if is_mps else set(qop.in_edges + qop.out_edges)
-
     if is_mps:
         rank_2_nodes = {node for node in nodes_for_sorting if len(node.edges) == 2}
         if len(rank_2_nodes) == 2:
@@ -1340,6 +1339,7 @@ def qop2tenpy(qop: QuOperator) -> Any:
     - Cyclic boundary conditions NOT supported
 
     :param qop: The corresponding state/operator as a QuOperator.
+    :type qop: QuOperator
     :return: MPO or MPS object from the TeNPy package.
     :rtype: Union[tenpy.networks.mpo.MPO, tenpy.networks.mps.MPS]
     """
@@ -1480,6 +1480,7 @@ def qop2quimb(qop: QuOperator) -> Any:
     - Cyclic boundary conditions NOT supported
 
     :param qop: MPO in the form of QuOperator
+    :type qop: QuOperator
     :return: MPO in the form of Quimb package
     :rtype: quimb.tensor.tensor_gen.MatrixProductOperator
     """

tests/test_mpscircuit.py

@@ -309,6 +309,7 @@ def test_circuits(backend, dtype):
     do_test_measure(circuits)
 
 
+# TODO(@refraction-ray): fails  (lf("jaxb"), lf("highp"))
 @pytest.mark.parametrize("backend, dtype", [(lf("tfb"), lf("highp"))])
 def test_circuits_jit(backend, dtype):
     def expec(params):
@@ -325,15 +326,48 @@ def expec(params):
     expec_vg_jit = tc.backend.jit(expec_vg)
     exp = expec(params)
     exp_jit, exp_grad_jit = expec_vg_jit(params)
-    dir = tc.backend.convert_to_tensor(np.array([1.0, 2.0, 3.0], dtype=tc.dtypestr))
+    dir_ = tc.backend.convert_to_tensor(np.array([1.0, 2.0, 3.0], dtype=tc.dtypestr))
     epsilon = 1e-6
-    exp_p = expec(params + dir * epsilon)
-    exp_m = expec(params - dir * epsilon)
+    exp_p = expec(params + dir_ * epsilon)
+    exp_m = expec(params - dir_ * epsilon)
     exp_grad_dir_numerical = (exp_p - exp_m) / (epsilon * 2)
-    exp_grad_dir_jit = tc.backend.real(tc.backend.sum(exp_grad_jit * dir))
+    exp_grad_dir_jit = tc.backend.real(tc.backend.sum(exp_grad_jit * dir_))
+    np.testing.assert_allclose(exp, exp_jit, atol=1e-10)
     np.testing.assert_allclose(
-        tc.backend.numpy(exp), tc.backend.numpy(exp_jit), atol=1e-10
+        tc.backend.numpy(exp_grad_dir_numerical),
+        tc.backend.numpy(exp_grad_dir_jit),
+        atol=1e-6,
     )
+
+
+@pytest.mark.parametrize(
+    "backend, dtype", [(lf("tfb"), lf("highp")), (lf("jaxb"), lf("highp"))]
+)
+def test_simple_circuits_ad(backend, dtype):
+    def expec(params):
+        mps = tc.MPSCircuit(N, split=split)
+        mps.rx(0, theta=params[0])
+        mps.cx(0, 1)
+        mps.cx(1, 2)
+        mps.ry(2, theta=params[1])
+        mps.rzz(1, 3, theta=params[2])
+        x = [0, 2]
+        z = [1]
+        exp = mps.expectation_ps(x=x, z=z)
+        return tc.backend.real(exp)
+
+    params = tc.backend.ones((3,), dtype=tc.dtypestr)
+    expec_vg = tc.backend.value_and_grad(expec)
+    expec_vg_jit = tc.backend.jit(expec_vg)
+    exp = expec(params)
+    exp_jit, exp_grad_jit = expec_vg_jit(params)
+    dir_ = tc.backend.convert_to_tensor(np.array([1.0, 2.0, 3.0], dtype=tc.dtypestr))
+    epsilon = 1e-6
+    exp_p = expec(params + dir_ * epsilon)
+    exp_m = expec(params - dir_ * epsilon)
+    exp_grad_dir_numerical = (exp_p - exp_m) / (epsilon * 2)
+    exp_grad_dir_jit = tc.backend.real(tc.backend.sum(exp_grad_jit * dir_))
+    np.testing.assert_allclose(exp, exp_jit, atol=1e-6)
     np.testing.assert_allclose(
         tc.backend.numpy(exp_grad_dir_numerical),
         tc.backend.numpy(exp_grad_dir_jit),