fix some api in dmcircuit

Author: refraction-ray

CHANGELOG.md

@@ -2,9 +2,15 @@
 
 ## Unreleased
 
+### Added
+
+- Add `expectation_ps` method for `DMCircuit`
+
 ### Fixed
 
-- With `Circuit.vis_tex`, for the Circuit has customized input state, the default visualization is psi instead of all zeros now.
+- With `Circuit.vis_tex`, for the Circuit has customized input state, the default visualization is psi instead of all zeros now
+
+- `general_kraus` is synced with `apply_general_kraus` for `DMCircuit`
 
 ## 0.0.220402
 

docs/source/conf.py

@@ -48,6 +48,8 @@
     "nbsphinx",
 ]
 
+autosectionlabel_prefix_document = True
+
 # Add any paths that contain templates here, relative to this directory.
 templates_path = ["_templates"]
 

docs/source/faq.rst

@@ -40,6 +40,22 @@ For PyTorch, we can in pricinple wrap the corresponding quantum function into a
 In terms of Jax backend, we highly suggested to keep the functional programming paradigm for such machine learning task.
 Besides, it is worthing noting that, jit and vmap is automatically taken care of in ``QuantumLayer``.
 
+When do I need to customize the contractor and how?
+------------------------------------------------------
+
+As a rule of thumb, for circuit with qubit count larger than 16 and circuit depth larger than 8, customized contraction may outperform the deafult built-in greedy contraction strategy.
+
+To setup or not setup the customized contractor is about a trade-off between the time on contraction path finding and the time on the real contraction via matmul.
+
+The customized contractor cost much more time than default contractor in terms of contraction path searching, and via the path it finds, the real contraction can take less time and space.
+
+If the circuit simulation time is the bottleneck of the whole workflow, one can always try customized contractor to see whether there is some performance improvement.
+
+We recommend to use `cotengra library <https://cotengra.readthedocs.io/en/latest/index.html>`_ to setup the contractor, since there are lots of interesting hyperparameters to tune, we can achieve better trade-off between the time on contraction path search and the time on the real tensor network contraction.
+
+It is also worth noting that for jitted function which we usually use, the contraction path search is only called at the first run of the function, which further amortize the time and favor the use of highly customized contractor.
+
+In terms of how-to on contractor setup, please refer to :ref:`quickstart:Setup the Contractor`.
 
 Is there some API less cumbersome than ``expectation`` for Pauli string?
 ----------------------------------------------------------------------------

docs/source/tutorial.rst

@@ -13,5 +13,6 @@ Jupyter Tutorials
     tutorials/gradient_benchmark.ipynb
     tutorials/contractors.ipynb
     tutorials/operator_spreading.ipynb
+    tutorials/optimization_and_expressibility.ipynb
     tutorials/vqex_mbl.ipynb
     tutorials/dqas.ipynb

docs/source/tutorial_cn.rst

@@ -13,5 +13,6 @@
     tutorials/gradient_benchmark_cn.ipynb
     tutorials/contractors_cn.ipynb
     tutorials/operator_spreading_cn.ipynb
+    tutorials/optimization_and_expressibility_cn.ipynb
     tutorials/vqex_mbl_cn.ipynb
     tutorials/dqas_cn.ipynb

examples/vqe_extra_mpo.py

@@ -18,8 +18,8 @@
 import tensorflow as tf
 import tensornetwork as tn
 import cotengra as ctg
-
 import optax
+
 import tensorcircuit as tc
 
 opt = ctg.ReusableHyperOptimizer(

examples/vqe_extra_mpo_spopt.py

@@ -19,11 +19,9 @@
 
 sys.setrecursionlimit(10000)
 
-import tensorflow as tf
 import tensornetwork as tn
 import cotengra as ctg
 
-import optax
 import tensorcircuit as tc
 
 optc = ctg.ReusableHyperOptimizer(

tensorcircuit/densitymatrix.py

@@ -13,7 +13,7 @@
 from . import gates
 from . import channels
 from .circuit import Circuit
-from .cons import backend, contractor, npdtype
+from .cons import backend, contractor, npdtype, dtypestr
 
 Gate = gates.Gate
 Tensor = Any
@@ -50,6 +50,8 @@ def __init__(
                 lnodes.extend(nodes)
                 self._nodes = lnodes
             elif inputs is not None:
+                inputs = backend.convert_to_tensor(inputs)
+                inputs = backend.cast(inputs, dtype=dtypestr)
                 inputs = backend.reshape(inputs, [-1])
                 N = inputs.shape[0]
                 n = int(np.log(N) / np.log(2))
@@ -63,6 +65,8 @@ def __init__(
                 lnodes.extend(nodes)
                 self._nodes = lnodes
             else:  # dminputs is not None
+                dminputs = backend.convert_to_tensor(dminputs)
+                dminputs = backend.cast(dminputs, dtype=dtypestr)
                 dminputs = backend.reshape(dminputs, [2 for _ in range(2 * nqubits)])
                 dminputs = Gate(dminputs)
                 nodes = [dminputs]
@@ -269,5 +273,39 @@ def expectation(self, *ops: Tuple[tn.Node, List[int]]) -> tn.Node.tensor:
     def check_density_matrix(dm: Tensor) -> None:
         assert np.allclose(backend.trace(dm), 1.0, atol=1e-5)
 
+    def expectation_ps(
+        self,
+        x: Optional[Sequence[int]] = None,
+        y: Optional[Sequence[int]] = None,
+        z: Optional[Sequence[int]] = None,
+        reuse: bool = True,
+    ) -> Tensor:
+        """
+        Shortcut for Pauli string expectation.
+        x, y, z list are for X, Y, Z positions
+
+        :param x: _description_, defaults to None
+        :type x: Optional[Sequence[int]], optional
+        :param y: _description_, defaults to None
+        :type y: Optional[Sequence[int]], optional
+        :param z: _description_, defaults to None
+        :type z: Optional[Sequence[int]], optional
+        :param reuse: whether to cache and reuse the wavefunction, defaults to True
+        :type reuse: bool, optional
+        :return: Expectation value
+        :rtype: Tensor
+        """
+        obs = []
+        if x is not None:
+            for i in x:
+                obs.append([gates.x(), [i]])  # type: ignore
+        if y is not None:
+            for i in y:
+                obs.append([gates.y(), [i]])  # type: ignore
+        if z is not None:
+            for i in z:
+                obs.append([gates.z(), [i]])  # type: ignore
+        return self.expectation(*obs)  # type: ignore
+
 
 DMCircuit._meta_apply()

tensorcircuit/densitymatrix2.py

@@ -5,8 +5,10 @@
 
 from typing import Any, Callable, Sequence
 
+import tensornetwork as tn
+
 from . import gates
-from .cons import backend
+from .cons import backend, dtypestr
 from .channels import kraus_to_super_gate
 from .densitymatrix import DMCircuit
 
@@ -26,6 +28,12 @@ def _copy_DMCircuit(self) -> "DMCircuit2":
 
     def apply_general_kraus(self, kraus: Sequence[Gate], *index: int) -> None:  # type: ignore
         # incompatible API for now
+        kraus = [
+            k
+            if isinstance(k, tn.Node)
+            else Gate(backend.cast(backend.convert_to_tensor(k), dtypestr))
+            for k in kraus
+        ]
         self.check_kraus(kraus)
         if not isinstance(
             index[0], int
@@ -48,6 +56,8 @@ def apply_general_kraus(self, kraus: Sequence[Gate], *index: int) -> None:  # ty
         self._nodes.append(super_op)
         setattr(self, "state_tensor", None)
 
+    general_kraus = apply_general_kraus  # type: ignore
+
     @staticmethod
     def apply_general_kraus_delayed(
         krausf: Callable[..., Sequence[Gate]]

tests/test_dmcircuit.py

@@ -226,3 +226,15 @@ def test_ad_channel(tfb):
         g = tape.gradient(loss, p)
     np.testing.assert_allclose(loss.numpy(), 0.7562, atol=1e-4)
     np.testing.assert_allclose(g.numpy(), -2.388, atol=1e-4)
+
+
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb"), lf("torchb")])
+def test_inputs_pipeline(backend):
+    dm = 0.25 * np.eye(4)
+    c = tc.DMCircuit(2, dminputs=dm)
+    c.H(0)
+    c.general_kraus(
+        [np.sqrt(0.5) * tc.gates._i_matrix, np.sqrt(0.5) * tc.gates._x_matrix], 1
+    )
+    r = c.expectation_ps(z=[0])
+    np.testing.assert_allclose(r, 0.0, atol=1e-5)