tensorcircuit-nightly 1.2.0.dev20250326__py3-none-any.whl → 1.4.0.dev20251128__py3-none-any.whl

tensorcircuit/mpscircuit.py

@@ -4,21 +4,25 @@ Quantum circuit: MPS state simulator
 
 # 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 types
 
 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, _decode_basis_label
 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 +81,10 @@ class MPSCircuit(AbstractCircuit):
 
     is_mps = True
 
+    @partial(
+        arg_alias,
+        alias_dict={"wavefunction": ["inputs"]},
+    )
     def __init__(
         self,
         nqubits: int,
@@ -84,12 +92,16 @@ class MPSCircuit(AbstractCircuit):
         tensors: Optional[Sequence[Tensor]] = None,
         wavefunction: Optional[Union[QuVector, Tensor]] = None,
         split: Optional[Dict[str, Any]] = None,
+        dim: Optional[int] = None,
     ) -> None:
         """
         MPSCircuit object based on state simulator.
+        Do not use this class with d!=2 directly
 
         :param nqubits: The number of qubits in the circuit.
         :type nqubits: int
+        :param dim: The local Hilbert space dimension per site. Qudit is supported for 2 <= d <= 36.
+        :type dim: If None, the dimension of the circuit will be `2`, which is a qubit system.
         :param center_position: The center position of MPS, default to 0
         :type center_position: int, optional
         :param tensors: If not None, the initial state of the circuit is taken as ``tensors``
@@ -102,6 +114,7 @@ class MPSCircuit(AbstractCircuit):
         :param split: Split rules
         :type split: Any
         """
+        self._d = 2 if dim is None else dim
         self.circuit_param = {
             "nqubits": nqubits,
             "center_position": center_position,
@@ -118,8 +131,21 @@ class MPSCircuit(AbstractCircuit):
             ), "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, dim_phys=self._d, split=self.split
+                )
             assert len(tensors) == nqubits
             self._mps = FiniteMPS(tensors, canonicalize=False)
             self._mps.center_position = 0
@@ -133,8 +159,13 @@ class MPSCircuit(AbstractCircuit):
                 self._mps = FiniteMPS(tensors, canonicalize=True, center_position=0)
         else:
             tensors = [
-                np.array([1.0, 0.0], dtype=npdtype)[None, :, None]
-                for i in range(nqubits)
+                np.concatenate(
+                    [
+                        np.array([1.0], dtype=npdtype),
+                        np.zeros((self._d - 1,), dtype=npdtype),
+                    ]
+                )[None, :, None]
+                for _ in range(nqubits)
             ]
             self._mps = FiniteMPS(tensors, canonicalize=False)
             if center_position is not None:
@@ -260,6 +291,17 @@ class MPSCircuit(AbstractCircuit):
         index_to: int,
         split: Optional[Dict[str, Any]] = None,
     ) -> None:
+        """
+        Apply a series of SWAP gates to move a qubit from ``index_from`` to ``index_to``.
+
+        :param index_from: The starting index of the qubit.
+        :type index_from: int
+        :param index_to: The destination index of the qubit.
+        :type index_to: int
+        :param split: Truncation options for the SWAP gates. Defaults to None.
+            consistent with the split option of the class.
+        :type split: Optional[Dict[str, Any]], optional
+        """
         if split is None:
             split = self.split
         self.position(index_from)
@@ -341,42 +383,51 @@ class MPSCircuit(AbstractCircuit):
         #      b
 
         # index must be ordered
-        assert np.all(np.diff(index) > 0)
-        index_left = np.min(index)
+        if len(index) == 0:
+            raise ValueError("`index` must contain at least one site.")
+        if not all(index[i] < index[i + 1] for i in range(len(index) - 1)):
+            raise ValueError("`index` must be strictly increasing.")
+
+        index_left = int(np.min(index))
         if isinstance(gate, tn.Node):
             gate = backend.copy(gate.tensor)
-        index = np.array(index) - index_left
+
         nindex = len(index)
+        in_dims = tuple(backend.shape_tuple(gate))[:nindex]
+        dim = int(in_dims[0])
+        dim_phys_mpo = dim * dim
+        gate = backend.reshape(gate, (dim,) * nindex + (dim,) * nindex)
         # transform gate from (in1, in2, ..., out1, out2 ...) to
         # (in1, out1, in2, out2, ...)
-        order = tuple(np.arange(2 * nindex).reshape((2, nindex)).T.flatten())
-        shape = (4,) * nindex
-        gate = backend.reshape(backend.transpose(gate, order), shape)
-        argsort = np.argsort(index)
+        order = tuple(np.arange(2 * nindex).reshape(2, nindex).T.flatten().tolist())
+        gate = backend.transpose(gate, order)
         # reorder the gate according to the site positions
-        gate = backend.transpose(gate, tuple(argsort))
-        index = index[argsort]  # type: ignore
+        gate = backend.reshape(gate, (dim_phys_mpo,) * nindex)
         # split the gate into tensors assuming they are adjacent
-        main_tensors = cls.wavefunction_to_tensors(gate, dim_phys=4, norm=False)
+        main_tensors = cls.wavefunction_to_tensors(
+            gate, dim_phys=dim_phys_mpo, norm=False
+        )
         # each tensor is in shape of (i, a, b, j)
-        tensors = []
-        previous_i = None
-        for i, main_tensor in zip(index, main_tensors):
-            # insert identites in the middle
+        tensors: list[Tensor] = []
+        previous_i: Optional[int] = None
+        index_arr = np.array(index, dtype=int) - index_left
+
+        for i, main_tensor in zip(index_arr, main_tensors):
             if previous_i is not None:
-                for _ in range(previous_i + 1, i):
-                    bond_dim = tensors[-1].shape[-1]
-                    I = (
-                        np.eye(bond_dim * 2)
-                        .reshape((bond_dim, 2, bond_dim, 2))
-                        .transpose((0, 1, 3, 2))
-                        .astype(dtypestr)
+                for _gap_site in range(int(previous_i) + 1, int(i)):
+                    bond_dim = int(backend.shape_tuple(tensors[-1])[-1])
+                    eye2d = backend.eye(
+                        bond_dim * dim, dtype=backend.dtype(tensors[-1])
                     )
-                    tensors.append(backend.convert_to_tensor(I))
-            nleft, _, nright = main_tensor.shape
-            tensor = backend.reshape(main_tensor, (nleft, 2, 2, nright))
+                    I4 = backend.reshape(eye2d, (bond_dim, dim, bond_dim, dim))
+                    I4 = backend.transpose(I4, (0, 1, 3, 2))
+                    tensors.append(I4)
+
+            nleft, _, nright = backend.shape_tuple(main_tensor)
+            tensor = backend.reshape(main_tensor, (int(nleft), dim, dim, int(nright)))
             tensors.append(tensor)
-            previous_i = i
+            previous_i = int(i)
+
         return tensors, index_left
 
     @classmethod
@@ -419,15 +470,15 @@ class MPSCircuit(AbstractCircuit):
         """
         if split is None:
             split = {}
-        ni = tensor_left.shape[0]
-        nk = tensor_right.shape[-1]
+        ni, di = tensor_left.shape[0], tensor_right.shape[1]
+        nk, dk = tensor_right.shape[-1], tensor_right.shape[-2]
         T = backend.einsum("iaj,jbk->iabk", tensor_left, tensor_right)
-        T = backend.reshape(T, (ni * 2, nk * 2))
+        T = backend.reshape(T, (ni * di, nk * dk))
         new_tensor_left, new_tensor_right = split_tensor(
             T, center_left=center_left, split=split
         )
-        new_tensor_left = backend.reshape(new_tensor_left, (ni, 2, -1))
-        new_tensor_right = backend.reshape(new_tensor_right, (-1, 2, nk))
+        new_tensor_left = backend.reshape(new_tensor_left, (ni, di, -1))
+        new_tensor_right = backend.reshape(new_tensor_right, (-1, dk, nk))
         return new_tensor_left, new_tensor_right
 
     def reduce_dimension(
@@ -437,7 +488,15 @@ class MPSCircuit(AbstractCircuit):
         split: Optional[Dict[str, Any]] = None,
     ) -> None:
         """
-        Reduce the bond dimension between two adjacent sites by SVD
+        Reduce the bond dimension between two adjacent sites using SVD.
+
+        :param index_left: The index of the left tensor of the bond to be truncated.
+        :type index_left: int
+        :param center_left: If True, the orthogonality center will be on the left tensor after truncation.
+                            Otherwise, it will be on the right tensor. Defaults to True.
+        :type center_left: bool, optional
+        :param split: Truncation options for the SVD. Defaults to None.
+        :type split: Optional[Dict[str, Any]], optional
         """
         if split is None:
             split = self.split
@@ -463,7 +522,22 @@ class MPSCircuit(AbstractCircuit):
         split: Optional[Dict[str, Any]] = None,
     ) -> None:
         """
-        Apply a MPO to the MPS
+        Apply a Matrix Product Operator (MPO) to the MPS.
+
+        The application involves three main steps:
+        1. Contract the MPO tensors with the corresponding MPS tensors.
+        2. Canonicalize the resulting tensors by moving the orthogonality center.
+        3. Truncate the bond dimensions to control complexity.
+
+        :param tensors: A sequence of tensors representing the MPO.
+        :type tensors: Sequence[Tensor]
+        :param index_left: The starting index on the MPS where the MPO is applied.
+        :type index_left: int
+        :param center_left: If True, the final orthogonality center will be at the left end of the MPO.
+                            Otherwise, it will be at the right end. Defaults to True.
+        :type center_left: bool, optional
+        :param split: Truncation options for bond dimension reduction. Defaults to None.
+        :type split: Optional[Dict[str, Any]], optional
         """
         # step 1:
         #     contract tensor
@@ -498,10 +572,11 @@ class MPSCircuit(AbstractCircuit):
         for i, idx in zip(i_list, idx_list):
             O = tensors[i]
             T = self._mps.tensors[idx]
-            ni, _, _, nj = O.shape
+            ni, d_in, _, nj = O.shape
             nk, _, nl = T.shape
             OT = backend.einsum("iabj,kbl->ikajl", O, T)
-            OT = backend.reshape(OT, (ni * nk, 2, nj * nl))
+            OT = backend.reshape(OT, (ni * nk, d_in, nj * nl))
+
             self._mps.tensors[idx] = OT
 
         # canonicalize
@@ -521,10 +596,17 @@ class MPSCircuit(AbstractCircuit):
         *index: int,
         split: Optional[Dict[str, Any]] = None,
     ) -> None:
-        # TODO(@SUSYUSTC): jax autograd is wrong on this function
         """
-        Apply a n-qubit gate by transforming the gate to MPO
+        Apply an n-qubit gate to the MPS by converting it to an MPO.
+
+        :param gate: The n-qubit gate to apply.
+        :type gate: Gate
+        :param index: The indices of the qubits to apply the gate to.
+        :type index: int
+        :param split: Truncation options for the MPO application. Defaults to None.
+        :type split: Optional[Dict[str, Any]], optional
         """
+        # TODO(@SUSYUSTC): jax autograd is wrong on this function
         ordered = np.all(np.diff(index) > 0)
         if not ordered:
             order = np.argsort(index)
@@ -601,8 +683,7 @@ class MPSCircuit(AbstractCircuit):
         :type keep: int, optional
         """
         # normalization not guaranteed
-        assert keep in [0, 1]
-        gate = backend.zeros((2, 2), dtype=dtypestr)
+        gate = backend.zeros((self._d, self._d), dtype=dtypestr)
         gate = backend.scatter(
             gate,
             backend.convert_to_tensor([[keep, keep]]),
@@ -633,7 +714,7 @@ class MPSCircuit(AbstractCircuit):
     def wavefunction_to_tensors(
         cls,
         wavefunction: Tensor,
-        dim_phys: int = 2,
+        dim_phys: Optional[int] = None,
         norm: bool = True,
         split: Optional[Dict[str, Any]] = None,
     ) -> List[Tensor]:
@@ -651,6 +732,7 @@ class MPSCircuit(AbstractCircuit):
         :return: The tensors
         :rtype: List[Tensor]
         """
+        dim_phys = dim_phys if dim_phys is not None else 2
         if split is None:
             split = {}
         wavefunction = backend.reshape(wavefunction, (-1, 1))
@@ -709,10 +791,16 @@ class MPSCircuit(AbstractCircuit):
         for key in vars(self):
             if key == "_mps":
                 continue
-            if backend.is_tensor(info[key]):
-                copied_value = backend.copy(info[key])
+            val = info[key]
+            if backend.is_tensor(val):
+                copied_value = backend.copy(val)
+            elif isinstance(val, types.ModuleType):
+                copied_value = val
             else:
-                copied_value = copy(info[key])
+                try:
+                    copied_value = copy(val)
+                except TypeError:
+                    copied_value = val
             setattr(result, key, copied_value)
         return result
 
@@ -756,7 +844,8 @@ class MPSCircuit(AbstractCircuit):
 
     def amplitude(self, l: str) -> Tensor:
         assert len(l) == self._nqubits
-        tensors = [self._mps.tensors[i][:, int(s), :] for i, s in enumerate(l)]
+        idx_list = _decode_basis_label(l, n=self._nqubits, dim=self._d)
+        tensors = [self._mps.tensors[i][:, idx, :] for i, idx in enumerate(idx_list)]
         return reduce(backend.matmul, tensors)[0, 0]
 
     def proj_with_mps(self, other: "MPSCircuit", conj: bool = True) -> Tensor:
@@ -814,6 +903,7 @@ class MPSCircuit(AbstractCircuit):
 
         mps = self.__class__(
             nqubits,
+            dim=self._d,
             tensors=tensors,
             center_position=center_position,
             split=self.split.copy(),
@@ -941,36 +1031,35 @@ class MPSCircuit(AbstractCircuit):
         # set the center to the left side, then gradually move to the right and do measurement at sites
         """
         mps = self.copy()
-        up = backend.convert_to_tensor(np.array([1, 0]).astype(dtypestr))
-        down = backend.convert_to_tensor(np.array([0, 1]).astype(dtypestr))
 
         p = 1.0
         p = backend.convert_to_tensor(p)
         p = backend.cast(p, dtype=rdtypestr)
-        sample = []
+        sample: Tensor = []
         for k, site in enumerate(index):
             mps.position(site)
-            # do measurement
             tensor = mps._mps.tensors[site]
             ps = backend.real(
                 backend.einsum("iaj,iaj->a", tensor, backend.conj(tensor))
             )
             ps /= backend.sum(ps)
-            pu = ps[0]
             if status is None:
-                r = backend.implicit_randu()[0]
+                outcome = backend.implicit_randc(
+                    self._d, shape=1, p=backend.cast(ps, rdtypestr)
+                )[0]
             else:
-                r = status[k]
-            r = backend.real(backend.cast(r, dtypestr))
-            eps = 0.31415926 * 1e-12
-            sign = backend.sign(r - pu + eps) / 2 + 0.5  # in case status is exactly 0.5
-            sign = backend.convert_to_tensor(sign)
-            sign = backend.cast(sign, dtype=rdtypestr)
-            sign_complex = backend.cast(sign, dtypestr)
-            sample.append(sign_complex)
-            p = p * (pu * (-1) ** sign + sign)
-            m = (1 - sign_complex) * up + sign_complex * down
+                one_r = backend.cast(backend.convert_to_tensor(1.0), rdtypestr)
+                st = backend.cast(status[k : k + 1], rdtypestr)
+                ind = backend.probability_sample(
+                    shots=1, p=backend.cast(ps, rdtypestr), status=one_r - st
+                )
+                outcome = backend.cast(ind[0], "int32")
+
+            p = p * ps[outcome]
+            basis = backend.convert_to_tensor(np.eye(self._d).astype(dtypestr))
+            m = basis[outcome]
             mps._mps.tensors[site] = backend.einsum("iaj,a->ij", tensor, m)[:, None, :]
+            sample.append(outcome)
         sample = backend.stack(sample)
         sample = backend.real(sample)
         if with_prob: