tensorcircuit-nightly 1.3.0.dev20250728__py3-none-any.whl → 1.4.0.dev20251103__py3-none-any.whl

tensorcircuit/basecircuit.py

@@ -1,5 +1,9 @@
 """
 Quantum circuit: common methods for all circuit classes as MixIn
+
+Note:
+  - Supports qubit (d = 2) and qudit (d >= 2) systems.
+  - For string-encoded samples/counts when d <= 36, digits use base-d characters 0-9A-Z (A = 10, ..., Z = 35).
 """
 
 # pylint: disable=invalid-name
@@ -19,8 +23,10 @@ from .quantum import (
     correlation_from_counts,
     measurement_counts,
     sample_int2bin,
-    sample_bin2int,
     sample2all,
+    _infer_num_sites,
+    _decode_basis_label,
+    onehot_d_tensor,
 )
 from .abstractcircuit import AbstractCircuit
 from .cons import npdtype, backend, dtypestr, contractor, rdtypestr
@@ -41,8 +47,9 @@ class BaseCircuit(AbstractCircuit):
     is_mps = False
 
     @staticmethod
-    def all_zero_nodes(n: int, d: int = 2, prefix: str = "qb-") -> List[tn.Node]:
-        l = [0.0 for _ in range(d)]
+    def all_zero_nodes(n: int, prefix: str = "qb-", dim: int = 2) -> List[tn.Node]:
+        prefix = "qd-" if dim > 2 else prefix
+        l = [0.0 for _ in range(dim)]
         l[0] = 1.0
         nodes = [
             tn.Node(
@@ -289,7 +296,7 @@ class BaseCircuit(AbstractCircuit):
         for op, index in ops:
             if not isinstance(op, tn.Node):
                 # op is only a matrix
-                op = backend.reshape2(op)
+                op = backend.reshaped(op, d=self._d)
                 op = backend.cast(op, dtype=dtypestr)
                 op = gates.Gate(op)
             else:
@@ -355,12 +362,12 @@ class BaseCircuit(AbstractCircuit):
 
     def perfect_sampling(self, status: Optional[Tensor] = None) -> Tuple[str, float]:
         """
-        Sampling bistrings from the circuit output based on quantum amplitudes.
+        Sampling base-d strings (0-9A-Z when d <= 36) from the circuit output based on quantum amplitudes.
         Reference: arXiv:1201.3974.
 
         :param status: external randomness, with shape [nqubits], defaults to None
         :type status: Optional[Tensor]
-        :return: Sampled bit string and the corresponding theoretical probability.
+        :return: Sampled base-d string and the corresponding theoretical probability.
         :rtype: Tuple[str, float]
         """
         return self.measure_jit(*range(self._nqubits), with_prob=True, status=status)
@@ -369,10 +376,10 @@ class BaseCircuit(AbstractCircuit):
         self, *index: int, with_prob: bool = False, status: Optional[Tensor] = None
     ) -> Tuple[Tensor, Tensor]:
         """
-        Take measurement to the given quantum lines.
+        Take measurement on the given site indices (computational basis).
         This method is jittable is and about 100 times faster than unjit version!
 
-        :param index: Measure on which quantum line.
+        :param index: Measure on which site (wire) index.
         :type index: int
         :param with_prob: If true, theoretical probability is also returned.
         :type with_prob: bool, optional
@@ -383,9 +390,8 @@ class BaseCircuit(AbstractCircuit):
         """
         # finally jit compatible ! and much faster than unjit version ! (100x)
         sample: List[Tensor] = []
-        p = 1.0
-        p = backend.convert_to_tensor(p)
-        p = backend.cast(p, dtype=rdtypestr)
+        one_r = backend.cast(backend.convert_to_tensor(1.0), rdtypestr)
+        p = one_r
         for k, j in enumerate(index):
             if self.is_dm is False:
                 nodes1, edge1 = self._copy()
@@ -400,40 +406,71 @@ class BaseCircuit(AbstractCircuit):
                 if i != j:
                     e ^ edge2[i]
             for i in range(k):
-                m = (1 - sample[i]) * gates.array_to_tensor(np.array([1, 0])) + sample[
-                    i
-                ] * gates.array_to_tensor(np.array([0, 1]))
-                newnodes.append(Gate(m))
-                newnodes[-1].id = id(newnodes[-1])
-                newnodes[-1].is_dagger = False
-                newnodes[-1].flag = "measurement"
-                newnodes[-1].get_edge(0) ^ edge1[index[i]]
-                newnodes.append(Gate(m))
-                newnodes[-1].id = id(newnodes[-1])
-                newnodes[-1].is_dagger = True
-                newnodes[-1].flag = "measurement"
-                newnodes[-1].get_edge(0) ^ edge2[index[i]]
+                if self._d == 2:
+                    m = (1 - sample[i]) * gates.array_to_tensor(
+                        np.array([1, 0])
+                    ) + sample[i] * gates.array_to_tensor(np.array([0, 1]))
+                else:
+                    m = onehot_d_tensor(sample[i], d=self._d)
+                g1 = Gate(m)
+                g1.id = id(g1)
+                g1.is_dagger = False
+                g1.flag = "measurement"
+                newnodes.append(g1)
+                g1.get_edge(0) ^ edge1[index[i]]
+                g2 = Gate(m)
+                g2.id = id(g2)
+                g2.is_dagger = True
+                g2.flag = "measurement"
+                newnodes.append(g2)
+                g2.get_edge(0) ^ edge2[index[i]]
+
             rho = (
                 1
                 / backend.cast(p, dtypestr)
                 * contractor(newnodes, output_edge_order=[edge1[j], edge2[j]]).tensor
             )
-            pu = backend.real(rho[0, 0])
-            if status is None:
-                r = backend.implicit_randu()[0]
+            if self._d == 2:
+                pu = backend.real(rho[0, 0])
+                if status is None:
+                    r = backend.implicit_randu()[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)
             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)
-
-        sample = backend.stack(sample)
-        sample = backend.real(sample)
+                pu = backend.clip(
+                    backend.real(backend.diagonal(rho)),
+                    backend.convert_to_tensor(0.0),
+                    backend.convert_to_tensor(1.0),
+                )
+                pu = pu / backend.sum(pu)
+                if status is None:
+                    ind = backend.implicit_randc(
+                        a=backend.arange(self._d),
+                        shape=1,
+                        p=backend.cast(pu, rdtypestr),
+                    )
+                else:
+                    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(pu, rdtypestr),
+                        status=one_r - st,
+                    )
+                k_out = backend.cast(ind[0], "int32")
+                sample.append(backend.cast(k_out, rdtypestr))
+                p = p * backend.cast(pu[k_out], rdtypestr)
+        sample = backend.real(backend.stack(sample))
         if with_prob:
             return sample, p
         else:
@@ -453,25 +490,22 @@ class BaseCircuit(AbstractCircuit):
         :return: The tensornetwork nodes for the amplitude of the circuit.
         :rtype: List[Gate]
         """
+
         no, d_edges = self._copy()
         ms = []
         if self.is_dm:
             msconj = []
         if isinstance(l, str):
-            for s in l:
-                if s == "1":
-                    endn = np.array([0, 1], dtype=npdtype)
-                elif s == "0":
-                    endn = np.array([1, 0], dtype=npdtype)
-                ms.append(tn.Node(endn))
+            symbols = _decode_basis_label(l, n=self._nqubits, dim=self._d)
+            for k in symbols:
+                n = onehot_d_tensor(k, d=self._d)
+                ms.append(tn.Node(n))
                 if self.is_dm:
-                    msconj.append(tn.Node(endn))
-        else:  # l is Tensor
+                    msconj.append(tn.Node(n))
+        else:
             l = backend.cast(l, dtype=dtypestr)
             for i in range(self._nqubits):
-                endn = l[i] * gates.array_to_tensor(np.array([0, 1])) + (
-                    1 - l[i]
-                ) * gates.array_to_tensor(np.array([1, 0]))
+                endn = onehot_d_tensor(l[i], d=self._d)
                 ms.append(tn.Node(endn))
                 if self.is_dm:
                     msconj.append(tn.Node(endn))
@@ -522,17 +556,18 @@ class BaseCircuit(AbstractCircuit):
 
     def probability(self) -> Tensor:
         """
-        get the 2^n length probability vector over computational basis
+        get the d^n length probability vector over computational basis
 
-        :return: probability vector
+        :return: probability vector of shape [dim**n]
         :rtype: Tensor
         """
         s = self.state()  # type: ignore
         if self.is_dm is False:
-            p = backend.abs(s) ** 2
-
+            amp = backend.reshape(s, [-1])
+            p = backend.real(backend.abs(amp) ** 2)
         else:
-            p = backend.abs(backend.diagonal(s))
+            diag = backend.diagonal(s)
+            p = backend.real(backend.reshape(diag, [-1]))
         return p
 
     @partial(arg_alias, alias_dict={"format": ["format_"]})
@@ -546,7 +581,7 @@ class BaseCircuit(AbstractCircuit):
         status: Optional[Tensor] = None,
         jittable: bool = True,
     ) -> Any:
-        """
+        r"""
         batched sampling from state or circuit tensor network directly
 
         :param batch: number of samples, defaults to None
@@ -569,6 +604,7 @@ class BaseCircuit(AbstractCircuit):
                 "count_tuple": # (np.array([0]), np.array([2]))
 
                 "count_dict_bin": # {"00": 2, "01": 0, "10": 0, "11": 0}
+                    for cases d\in [11, 36], use 0-9A-Z digits (e.g., 'A' -> 10, ..., 'Z' -> 35);
 
                 "count_dict_int": # {0: 2, 1: 0, 2: 0, 3: 0}
 
@@ -576,7 +612,8 @@ class BaseCircuit(AbstractCircuit):
         :param random_generator: random generator,  defaults to None
         :type random_generator: Optional[Any], optional
         :param status: external randomness given by tensor uniformly from [0, 1],
-            if set, can overwrite random_generator
+            if set, can overwrite random_generator, shape [batch] for `allow_state=True`
+            and shape [batch, nqubits] for `allow_state=False` using perfect sampling implementation
         :type status: Optional[Tensor]
         :param jittable: when converting to count, whether keep the full size. if false, may be conflict
             external jit, if true, may fail for large scale system with actual limited count results
@@ -597,25 +634,29 @@ class BaseCircuit(AbstractCircuit):
                     return r
                 r = [r]  # type: ignore
             else:
+                r = []  # type: ignore
+                if status is not None:
+                    assert backend.shape_tuple(status)[0] == batch
+                    for seed in status:
+                        r.append(self.perfect_sampling(seed))  # type: ignore
 
-                @backend.jit
-                def perfect_sampling(key: Any) -> Any:
-                    backend.set_random_state(key)
-                    return self.perfect_sampling()
+                else:
 
-                # TODO(@refraction-ray): status is not used here
+                    @backend.jit
+                    def perfect_sampling(key: Any) -> Any:
+                        backend.set_random_state(key)
+                        return self.perfect_sampling()
 
-                r = []  # type: ignore
+                    subkey = random_generator
+                    for _ in range(batch):
+                        key, subkey = backend.random_split(subkey)
+                        r.append(perfect_sampling(key))  # type: ignore
 
-                subkey = random_generator
-                for _ in range(batch):
-                    key, subkey = backend.random_split(subkey)
-                    r.append(perfect_sampling(key))  # type: ignore
             if format is None:
                 return r
             r = backend.stack([ri[0] for ri in r])  # type: ignore
-            r = backend.cast(r, "int32")
-            ch = sample_bin2int(r, self._nqubits)
+            ch = backend.cast(r, "int32")
+            # ch = sample_bin2int(r, self._nqubits, dim=self._d)
         else:  # allow_state
             if batch is None:
                 nbatch = 1
@@ -640,7 +681,7 @@ class BaseCircuit(AbstractCircuit):
             #     2,
             # )
             if format is None:  # for backward compatibility
-                confg = sample_int2bin(ch, self._nqubits)
+                confg = sample_int2bin(ch, self._nqubits, dim=self._d)
                 prob = backend.gather1d(p, ch)
                 r = list(zip(confg, prob))  # type: ignore
                 if batch is None:
@@ -648,7 +689,9 @@ class BaseCircuit(AbstractCircuit):
                 return r
         if self._nqubits > 35:
             jittable = False
-        return sample2all(sample=ch, n=self._nqubits, format=format, jittable=jittable)
+        return sample2all(
+            sample=ch, n=self._nqubits, format=format, jittable=jittable, dim=self._d
+        )
 
     def sample_expectation_ps(
         self,
@@ -848,9 +891,9 @@ class BaseCircuit(AbstractCircuit):
         """
         inputs = backend.reshape(inputs, [-1])
         N = inputs.shape[0]
-        n = int(np.log(N) / np.log(2))
+        n = _infer_num_sites(N, self._d)
         assert n == self._nqubits
-        inputs = backend.reshape(inputs, [2 for _ in range(n)])
+        inputs = backend.reshape(inputs, [self._d for _ in range(n)])
         if self.inputs is not None:
             self._nodes[0].tensor = inputs
             if self.is_dm:
@@ -881,9 +924,9 @@ class BaseCircuit(AbstractCircuit):
 
 
 
-        :param index: the qubit for the z-basis measurement
+        :param index: the site index for the Z-basis measurement
         :type index: int
-        :return: 0 or 1 for z measurement on up and down freedom
+        :return: 0 or 1 for Z-basis measurement outcome
         :rtype: Tensor
         """
         return self.general_kraus(  # type: ignore
@@ -962,8 +1005,8 @@ class BaseCircuit(AbstractCircuit):
 
     def projected_subsystem(self, traceout: Tensor, left: Tuple[int, ...]) -> Tensor:
         """
-        remaining wavefunction or density matrix on qubits in left, with other qubits
-        fixed in 0 or 1 indicated by traceout
+        remaining wavefunction or density matrix on sites in ``left``, with other sites
+        fixed to given digits (0..d-1) as indicated by ``traceout``
 
         :param traceout: can be jitted
         :type traceout: Tensor
@@ -972,15 +1015,14 @@ class BaseCircuit(AbstractCircuit):
         :return: _description_
         :rtype: Tensor
         """
-        end0, end1 = gates.array_to_tensor(np.array([1.0, 0]), np.array([0, 1.0]))
+
         traceout = backend.cast(traceout, dtypestr)
         nodes, front = self._copy()
         L = self._nqubits
         edges = []
         for i in range(len(traceout)):
             if i not in left:
-                b = traceout[i]
-                n = gates.Gate((1 - b) * end0 + b * end1)
+                n = Gate(onehot_d_tensor(traceout[i], d=self._d))
                 nodes.append(n)
                 front[i] ^ n[0]
             else:
@@ -989,8 +1031,7 @@ class BaseCircuit(AbstractCircuit):
         if self.is_dm:
             for i in range(len(traceout)):
                 if i not in left:
-                    b = traceout[i]
-                    n = gates.Gate((1 - b) * end0 + b * end1)
+                    n = Gate(onehot_d_tensor(traceout[i], d=self._d))
                     nodes.append(n)
                     front[i + L] ^ n[0]
                 else: