add onehot_d_tensor in quantum.py and apply.

Author: WeiguoMa

tensorcircuit/basecircuit.py

@@ -26,6 +26,7 @@
     sample2all,
     _infer_num_sites,
     _decode_basis_label,
+    onehot_d_tensor,
 )
 from .abstractcircuit import AbstractCircuit
 from .cons import npdtype, backend, dtypestr, contractor, rdtypestr
@@ -410,8 +411,7 @@ def measure_jit(
                         np.array([1, 0])
                     ) + sample[i] * gates.array_to_tensor(np.array([0, 1]))
                 else:
-                    vec = backend.one_hot(backend.cast(sample[i], "int32"), self._d)
-                    m = backend.cast(vec, dtypestr)
+                    m = onehot_d_tensor(sample[i], d=self._d)
                 g1 = Gate(m)
                 g1.id = id(g1)
                 g1.is_dagger = False
@@ -507,29 +507,21 @@ def amplitude_before(self, l: Union[str, Tensor]) -> List[Gate]:
         :rtype: List[Gate]
         """
 
-        def _basis_nod(_k: int) -> Tensor:
-            _vec = np.zeros((self._d,), dtype=npdtype)
-            _vec[_k] = 1.0
-            return _vec
-
         no, d_edges = self._copy()
         ms = []
         if self.is_dm:
             msconj = []
         if isinstance(l, str):
             symbols = _decode_basis_label(l, n=self._nqubits, dim=self._d)
             for k in symbols:
-                n = _basis_nod(k)
+                n = onehot_d_tensor(k, d=self._d)
                 ms.append(tn.Node(n))
                 if self.is_dm:
                     msconj.append(tn.Node(n))
         else:
             l = backend.cast(l, dtype=dtypestr)
             for i in range(self._nqubits):
-                endn = backend.cast(
-                    backend.one_hot(backend.cast(l[i], "int32"), self._d),
-                    dtype=dtypestr,
-                )
+                endn = onehot_d_tensor(l[i], d=self._d)
                 ms.append(tn.Node(endn))
                 if self.is_dm:
                     msconj.append(tn.Node(endn))
@@ -1040,18 +1032,13 @@ def projected_subsystem(self, traceout: Tensor, left: Tuple[int, ...]) -> Tensor
         :rtype: Tensor
         """
 
-        def _basis_gate(k_tensor: Any) -> Gate:
-            vec = backend.one_hot(backend.cast(k_tensor, "int32"), self._d)
-            vec = backend.cast(vec, dtypestr)
-            return Gate(vec)
-
         traceout = backend.cast(traceout, dtypestr)
         nodes, front = self._copy()
         L = self._nqubits
         edges = []
         for i in range(len(traceout)):
             if i not in left:
-                n = _basis_gate(traceout[i])
+                n = Gate(onehot_d_tensor(traceout[i], d=self._d))
                 nodes.append(n)
                 front[i] ^ n[0]
             else:
@@ -1060,7 +1047,7 @@ def _basis_gate(k_tensor: Any) -> Gate:
         if self.is_dm:
             for i in range(len(traceout)):
                 if i not in left:
-                    n = _basis_gate(traceout[i])
+                    n = Gate(onehot_d_tensor(traceout[i], d=self._d))
                     nodes.append(n)
                     front[i + L] ^ n[0]
                 else:

tensorcircuit/quantum.py

@@ -57,7 +57,44 @@ def get_all_nodes(edges: Iterable[Edge]) -> List[Node]:
     return nodes
 
 
+def onehot_d_tensor(_k: Union[int, Tensor], d: int = 2) -> Tensor:
+    """
+    Construct a one-hot vector (or matrix) of local dimension ``d``.
+
+    :param _k: index or indices to set as 1. Can be an int or a backend Tensor.
+    :type _k: int or Tensor
+    :param d: local dimension (number of categories), defaults to 2
+    :type d: int, optional
+    :return: one-hot encoded vector (shape [d]) or matrix (shape [len(_k), d])
+    :rtype: Tensor
+    """
+    if isinstance(_k, int):
+        vec = backend.one_hot(_k, d)
+    else:
+        vec = backend.one_hot(backend.cast(_k, "int32"), d)
+    return backend.cast(vec, dtypestr)
+
+
 def _decode_basis_label(label: str, n: int, dim: int) -> List[int]:
+    """
+    Decode a string basis label into a list of integer digits.
+
+    The label is interpreted in base-``dim`` using characters ``0–9A–Z``.
+    Only dimensions up to 36 are supported.
+
+    :param label: basis label string, e.g. "010" or "A9F"
+    :type label: str
+    :param n: number of sites (expected length of the label)
+    :type n: int
+    :param dim: local dimension (2 <= dim <= 36)
+    :type dim: int
+    :return: list of integer digits of length ``n``, each in ``[0, dim-1]``
+    :rtype: List[int]
+
+    :raises NotImplementedError: if ``dim > 36``
+    :raises ValueError: if the label length mismatches ``n``,
+                        or contains invalid/out-of-range characters
+    """
     if dim > 36:
         raise NotImplementedError(
             f"String basis label supports d<=36 (0–9A–Z). Got dim={dim}. "