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

tensorcircuit/quantum.py

@@ -9,6 +9,7 @@ import math
 import os
 from functools import partial, reduce
 from operator import matmul, mul, or_
+from collections import Counter
 from typing import (
     Any,
     Callable,
@@ -23,6 +24,7 @@ from typing import (
 )
 
 import numpy as np
+import tensornetwork as tn
 from tensornetwork.network_components import AbstractNode, CopyNode, Edge, Node, connect
 from tensornetwork.network_operations import (
     copy,
@@ -30,7 +32,7 @@ from tensornetwork.network_operations import (
     remove_node,
 )
 
-from .cons import backend, contractor, dtypestr, npdtype, rdtypestr
+from .cons import backend, contractor, dtypestr, npdtype, rdtypestr, _ALPHABET
 from .gates import Gate, num_to_tensor
 from .utils import arg_alias
 
@@ -55,6 +57,91 @@ 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}. "
+            "Use an integer array/tensor of length n instead."
+        )
+    s = label.upper()
+    if len(s) != n:
+        raise ValueError(f"Basis label length mismatch: expect {n}, got {len(s)}")
+    digits = []
+    for ch in s:
+        if ch not in _ALPHABET:
+            raise ValueError(
+                f"Invalid character '{ch}' in basis label (allowed 0-9A-Z)."
+            )
+        v = _ALPHABET.index(ch)
+        if v >= dim:
+            raise ValueError(
+                f"Digit '{ch}' (= {v}) out of range for base-d with dim={dim}."
+            )
+        digits.append(v)
+    return digits
+
+
+def _infer_num_sites(D: int, dim: int) -> int:
+    """
+    Infer the number of sites (n) from a Hilbert space dimension D
+    and local dimension d, assuming D = d**n.
+
+    :param D: total Hilbert space dimension (int)
+    :param dim: local dimension per site (int)
+    :return: n such that D == d**n
+    :raises ValueError: if D is not an exact power of d
+    """
+    if not (isinstance(D, int) and D > 0):
+        raise ValueError(f"D must be a positive integer, got {D}")
+    if not (isinstance(dim, int) and dim >= 2):
+        raise ValueError(f"d must be an integer >= 2, got {dim}")
+
+    tmp, n = D, 0
+    while tmp % dim == 0 and tmp > 1:
+        tmp //= dim
+        n += 1
+    if tmp != 1:
+        raise ValueError(f"Dimension {D} is not a power of local dim {dim}")
+    return n
+
+
 def _reachable(nodes: List[AbstractNode]) -> List[AbstractNode]:
     if not nodes:
         raise ValueError("Reachable requires at least 1 node.")
@@ -71,7 +158,7 @@ def _reachable(nodes: List[AbstractNode]) -> List[AbstractNode]:
                 if n not in seen_nodes and n not in node_que[i + 1 :]:
                     node_que.append(n)
         i += 1
-    return seen_nodes
+    return sorted(seen_nodes, key=lambda node: getattr(node, "_stable_id_", -1))
 
 
 def reachable(
@@ -661,10 +748,10 @@ class QuOperator:
         return self.__mul__(other)
 
     def tensor_product(self, other: "QuOperator") -> "QuOperator":
-        """
+        r"""
         Tensor product with another operator.
         Given two operators `A` and `B`, produces a new operator `AB` representing
-        :math:`A ⊗ B`. The `out_edges` (`in_edges`) of `AB` is simply the
+        :math:`A \otimes B`. The `out_edges` (`in_edges`) of `AB` is simply the
         concatenation of the `out_edges` (`in_edges`) of `A.copy()` with that of
         `B.copy()`:
         `new_out_edges = [*out_edges_A_copy, *out_edges_B_copy]`
@@ -1151,33 +1238,281 @@ def generate_local_hamiltonian(
     return hop
 
 
-def tn2qop(tn_mpo: Any) -> QuOperator:
+# TODO(@Charlespkuer): Add more conversion functions for other packages
+def extract_tensors_from_qop(qop: QuOperator) -> Tuple[List[Node], bool, int]:
     """
-    Convert MPO in TensorNetwork package to QuOperator.
+    Extract and sort tensors from QuOperator for conversion to other tensor network formats.
 
-    :param tn_mpo: MPO in the form of TensorNetwork package
-    :type tn_mpo: ``tn.matrixproductstates.mpo.*``
-    :return: MPO in the form of QuOperator
+    :param qop: Input QuOperator to extract tensors from
+    :type qop: QuOperator
+    :return: Tuple containing (sorted_nodes, is_mps, nwires) where:
+    - sorted_nodes: List of Node objects sorted in linear chain order
+    - is_mps: Boolean flag indicating if the structure is MPS (True) or MPO (False)
+    - nwires: Integer number of physical edges/qubits in the system
+    :rtype: Tuple[List[Node], bool, int]
+    """
+    is_mps = len(qop.in_edges) == 0
+    nwires = len(qop.out_edges)
+    if not is_mps and len(qop.in_edges) != nwires:
+        raise ValueError(
+            "MPO must have the same number of input and output edges. "
+            f"Got {len(qop.in_edges)} and {nwires}."
+        )
+
+    # Collect all nodes from edges
+    nodes_for_sorting = qop.nodes
+    if len(nodes_for_sorting) != nwires:
+        raise ValueError(f"Number of nodes  does not match number of wires.")
+
+    # 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:
+            endpoint_nodes = rank_2_nodes
+
+    if not endpoint_nodes:
+        endpoint_nodes = {edge.node1 for edge in qop.ignore_edges if edge.node1}
+
+    if not endpoint_nodes and len(nodes_for_sorting) > 1:
+        virtual_bond_counts = {}
+        virtual_bond_dim_sums = {}
+
+        for node in nodes_for_sorting:
+            virtual_bonds = 0
+            virtual_dim_sum = 0
+
+            for edge in node.edges:
+                if edge not in physical_edges and not edge.is_dangling():
+                    virtual_bonds += 1
+                    virtual_dim_sum += edge.dimension
+
+            virtual_bond_counts[node] = virtual_bonds
+            virtual_bond_dim_sums[node] = virtual_dim_sum
+
+        min_dim_sum = min(virtual_bond_dim_sums.values())
+        min_dim_nodes = {
+            node
+            for node, dim_sum in virtual_bond_dim_sums.items()
+            if dim_sum == min_dim_sum
+        }
+
+        if len(min_dim_nodes) == 2:
+            endpoint_nodes = min_dim_nodes
+
+    if not endpoint_nodes:
+        if len(nodes_for_sorting) == 1:
+            raise ValueError("Cannot determine chain structure: only one node found.")
+        elif len(nodes_for_sorting) >= 2:
+            raise ValueError(f"Cannot identify endpoint nodes for your nodes.")
+
+    # Sort nodes along the chain
+    sorted_nodes: list[Node] = []
+    if endpoint_nodes and len(endpoint_nodes) >= 1:
+        current = next(iter(endpoint_nodes))
+        while current and len(sorted_nodes) < nwires:
+            sorted_nodes.append(current)
+            current = next(
+                (
+                    e.node2 if e.node1 is current else e.node1
+                    for e in current.edges
+                    if not e.is_dangling()
+                    and e not in physical_edges
+                    and (e.node2 if e.node1 is current else e.node1) not in sorted_nodes
+                ),
+                None,
+            )
+
+    if not sorted_nodes:
+        raise ValueError("No valid chain structure found in the QuOperator. ")
+    if len(sorted_nodes) > 0 and len(qop.ignore_edges) > 0:
+        if sorted_nodes[0] is not qop.ignore_edges[0].node1:
+            sorted_nodes = sorted_nodes[::-1]
+
+    return sorted_nodes, is_mps, nwires
+
+
+def tenpy2qop(tenpy_obj: Any) -> QuOperator:
+    """
+    Converts a TeNPy MPO or MPS to a TensorCircuit QuOperator.
+    This definitive version correctly handles axis ordering and boundary
+    conditions to be compatible with `eval_matrix`.
+
+    :param tenpy_obj: A MPO or MPS object from the TeNPy package.
+    :type tenpy_obj: Union[tenpy.networks.mpo.MPO, tenpy.networks.mps.MPS]
+    :return: The corresponding state or operator as a QuOperator.
     :rtype: QuOperator
     """
-    tn_mpo = tn_mpo.tensors
-    nwires = len(tn_mpo)
-    mpo = []
-    for i in range(nwires):
-        mpo.append(Node(tn_mpo[i]))
+    # MPO objects have _W attribute containing tensor list (documented in tenpy.networks.mpo.MPO)
+    # MPS objects have _B attribute containing tensor list (documented in tenpy.networks.mps.MPS)
+    # These are internal attributes that store the actual tensor data for each site
+    # Reference: https://tenpy.readthedocs.io/en/latest/reference/tenpy.networks.mpo.html
+    #           https://tenpy.readthedocs.io/en/latest/reference/tenpy.networks.mps.html
+    is_mpo = hasattr(tenpy_obj, "_W")
+    tenpy_tensors = tenpy_obj._W if is_mpo else tenpy_obj._B
+    nwires = len(tenpy_tensors)
+    if nwires == 0:
+        return quantum_constructor([], [], [])
+
+    nodes = []
+    if is_mpo:
+        original_tensors_obj = tenpy_tensors
+
+        for i, W_obj in enumerate(original_tensors_obj):
+            arr = W_obj.to_ndarray()
+            labels = W_obj.get_leg_labels()
+            wL_idx = labels.index("wL")
+            p_idx = labels.index("p")
+            p_star_idx = labels.index("p*")
+            wR_idx = labels.index("wR")
+
+            arr_reordered = arr.transpose((wL_idx, p_idx, p_star_idx, wR_idx))
+            if nwires == 1:
+                arr_reordered = arr_reordered[[0], :, :, :]
+                arr_reordered = arr_reordered[:, :, :, [-1]]
+            else:
+                if i == 0:
+                    arr_reordered = arr_reordered[[0], :, :, :]
+                elif i == nwires - 1:
+                    arr_reordered = arr_reordered[:, :, :, [-1]]
+
+            node = Node(
+                arr_reordered, name=f"mpo_{i}", axis_names=["wL", "p", "p*", "wR"]
+            )
+            nodes.append(node)
+
+        if nwires > 1:
+            for i in range(nwires - 1):
+                nodes[i][3] ^ nodes[i + 1][0]
+
+        out_edges = [n[2] for n in nodes]
+        in_edges = [n[1] for n in nodes]
+        ignore_edges = [nodes[0][0], nodes[-1][3]]
+    else:  # MPS
+        for i in range(nwires):
+            B_obj = tenpy_obj.get_B(i)
+            arr = B_obj.to_ndarray()
+            labels = B_obj.get_leg_labels()
+            vL_idx = labels.index("vL")
+            p_idx = labels.index("p")
+            vR_idx = labels.index("vR")
+            arr_reordered = arr.transpose((vL_idx, p_idx, vR_idx))
+            node = Node(arr_reordered, name=f"mps_{i}", axis_names=["vL", "p", "vR"])
+            nodes.append(node)
+
+        if nwires > 1:
+            for i in range(nwires - 1):
+                nodes[i][2] ^ nodes[i + 1][0]
+
+        out_edges = [n[1] for n in nodes]
+        in_edges = []
+        ignore_edges = [nodes[0][0], nodes[-1][2]]
+
+    qop = quantum_constructor(out_edges, in_edges, [], ignore_edges)
 
-    for i in range(nwires - 1):
-        connect(mpo[i][1], mpo[i + 1][0])
-    # TODO(@refraction-ray): whether in and out edge is in the correct order require further check
-    qop = quantum_constructor(
-        [mpo[i][-1] for i in range(nwires)],  # out_edges
-        [mpo[i][-2] for i in range(nwires)],  # in_edges
-        [],
-        [mpo[0][0], mpo[-1][1]],  # ignore_edges
-    )
     return qop
 
 
+def qop2tenpy(qop: QuOperator) -> Any:
+    """
+    Convert TensorCircuit QuOperator to MPO or MPS from TeNPy.
+
+    Requirements: QuOperator must represent valid MPS/MPO structure:
+    - Linear chain topology with open boundaries only
+    - MPS: no input edges, consistent virtual bonds, rank-3 or 4(with empty input edges) tensors
+    - MPO: equal input/output edges, rank-4 tensors
+    - 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]
+    """
+    try:
+        from tenpy.networks import MPO, MPS, Site
+        from tenpy.linalg import np_conserved as npc
+        from tenpy.linalg import LegCharge
+    except ImportError:
+        raise ImportError("Please install TeNPy package to use this function.")
+
+    sorted_nodes, is_mps, nwires = extract_tensors_from_qop(qop)
+
+    physical_dim = qop.out_edges[0].dimension if is_mps else qop.in_edges[0].dimension
+    sites = [Site(LegCharge.from_trivial(physical_dim), "q") for _ in range(nwires)]
+
+    # MPS Conversion
+    if is_mps:
+        tensors = []
+        for i, node in enumerate(sorted_nodes):
+            tensor = np.asarray(node.tensor)
+            if tensor.ndim == 3:
+                if i == 0:
+                    if tensor.shape[0] > 1:
+                        tensor = tensor[0:1, :, :]
+                elif i == len(sorted_nodes) - 1:
+                    if tensor.shape[2] > 1:
+                        tensor = tensor[:, :, 0:1]
+            tensors.append(
+                npc.Array.from_ndarray(
+                    tensor,
+                    legcharges=[LegCharge.from_trivial(s) for s in tensor.shape],
+                    labels=["vL", "p", "vR"],
+                )
+            )
+
+        SVs = (
+            [np.ones([1])]
+            + [np.ones(tensors[i].get_leg("vR").ind_len) for i in range(nwires - 1)]
+            + [np.ones([1])]
+        )
+        return MPS(sites, tensors, SVs, bc="finite")
+
+    # MPO Conversion
+    raw_tensors = [np.asarray(node.tensor) for node in sorted_nodes]
+
+    if nwires == 1:
+        chi = 1
+        IdL = IdR = 0
+        reconstructed_tensors = raw_tensors
+    else:
+        chi = max(
+            raw_tensors[0].shape[3] if raw_tensors[0].ndim > 3 else 1,
+            raw_tensors[-1].shape[0] if raw_tensors[-1].ndim > 3 else 1,
+        )
+        IdL = 0
+        IdR = chi - 1 if chi > 1 else 0
+
+        reconstructed_tensors = []
+        for i, tensor in enumerate(raw_tensors):
+            if i == 0 and tensor.shape[0] < chi:
+                new_shape = (chi,) + tensor.shape[1:]
+                padded_tensor = np.zeros(new_shape, dtype=tensor.dtype)
+                padded_tensor[IdL, ...] = tensor[0, ...]
+                reconstructed_tensors.append(padded_tensor)
+            elif i == nwires - 1 and len(tensor.shape) > 3 and tensor.shape[3] < chi:
+                new_shape = tensor.shape[:3] + (chi,)
+                padded_tensor = np.zeros(new_shape, dtype=tensor.dtype)
+                padded_tensor[..., IdR] = tensor[..., 0]
+                reconstructed_tensors.append(padded_tensor)
+            else:
+                reconstructed_tensors.append(tensor)
+
+    tenpy_Ws = []
+    for tensor in reconstructed_tensors:
+        labels = ["wL", "wR", "p", "p*"]
+        tensor = np.transpose(tensor, (0, 3, 1, 2))
+        tenpy_Ws.append(
+            npc.Array.from_ndarray(
+                tensor,
+                legcharges=[LegCharge.from_trivial(s) for s in tensor.shape],
+                labels=labels,
+            )
+        )
+
+    return MPO(sites, tenpy_Ws, bc="finite", IdL=IdL, IdR=IdR)
+
+
 def quimb2qop(qb_mpo: Any) -> QuOperator:
     """
     Convert MPO in Quimb package to QuOperator.
@@ -1219,6 +1554,152 @@ def quimb2qop(qb_mpo: Any) -> QuOperator:
     return qop
 
 
+def qop2quimb(qop: QuOperator) -> Any:
+    """
+    Convert QuOperator to MPO or MPS in Quimb package.
+
+    Requirements: QuOperator must represent valid MPS/MPO structure:
+    - Linear chain topology with open boundaries only
+    - MPS: no input edges, consistent virtual bonds between adjacent tensors
+    - MPO: equal input/output edges, rank-4 tensors
+    - Edge connectivity: each internal node connected to exactly 2 neighbors
+    - 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
+    """
+    try:
+        import quimb.tensor as qtn
+    except ImportError:
+        raise ImportError("Please install Quimb package to use this function.")
+
+    sorted_nodes, is_mps, _ = extract_tensors_from_qop(qop)
+
+    quimb_tensors = []
+    node_map = {node: i for i, node in enumerate(sorted_nodes)}
+
+    for i, node in enumerate(sorted_nodes):
+        tensor_data = node.tensor
+        inds: List[str] = []
+
+        for axis, edge in enumerate(node.edges):
+            if edge in qop.out_edges:
+                site_index = qop.out_edges.index(edge)
+                inds.append(f"k{site_index}")
+            elif edge in qop.in_edges and not is_mps:
+                site_index = qop.in_edges.index(edge)
+                inds.append(f"b{site_index}")
+            elif edge in qop.ignore_edges:
+                if i == 0:
+                    inds.append("_left_dangling")
+                elif i == len(sorted_nodes) - 1:
+                    inds.append("_right_dangling")
+                else:
+                    inds.append(f"_ignore_{i}_{axis}")
+            else:
+                neighbor = edge.node1 if edge.node2 is node else edge.node2
+                if neighbor in node_map:
+                    j = node_map[neighbor]
+                    left, right = min(i, j), max(i, j)
+                    inds.append(f"v{left}_{right}")
+                else:
+                    inds.append(f"_unconnected_{i}_{axis}")
+
+        quimb_tensors.append(qtn.Tensor(tensor_data, inds=inds, tags=f"I{i}"))
+
+    tn = qtn.TensorNetwork(quimb_tensors)
+
+    if is_mps:
+        return tn.as_network(qtn.MatrixProductState)
+    else:
+        return tn.as_network(qtn.MatrixProductOperator)
+
+
+def tn2qop(tn_obj: Any) -> QuOperator:
+    """
+    Convert MPO in TensorNetwork package to QuOperator.
+
+    :param tn_mpo: MPO in the form of TensorNetwork package
+    :type tn_mpo: ``tn.matrixproductstates.mpo.*``
+    :return: MPO in the form of QuOperator
+    :rtype: QuOperator
+    """
+    tn_tensors = tn_obj.tensors
+    nwires = len(tn_tensors)
+
+    if nwires == 0:
+        return quantum_constructor([], [], [])
+
+    is_mps = all(len(t.shape) <= 3 for t in tn_tensors)
+
+    nodes = []
+    for i in range(nwires):
+        nodes.append(Node(tn_tensors[i], name=f"tensor_{i}"))
+
+    if is_mps:
+        for i in range(nwires - 1):
+            connect(nodes[i][-1], nodes[i + 1][0])
+
+        out_edges = []
+        for i, node in enumerate(nodes):
+            if len(node.edges) == 2:
+                physical_edge = next(e for e in node.edges if e.is_dangling())
+                out_edges.append(physical_edge)
+            else:
+                out_edges.append(node[1])
+
+        in_edges = []
+
+        ignore_edges = []
+        left_dangling = next(
+            (e for e in nodes[0].edges if e.is_dangling() and e not in out_edges), None
+        )
+        right_dangling = next(
+            (e for e in nodes[-1].edges if e.is_dangling() and e not in out_edges), None
+        )
+
+        if left_dangling:
+            ignore_edges.append(left_dangling)
+        if right_dangling:
+            ignore_edges.append(right_dangling)
+
+    else:
+        for i in range(nwires - 1):
+            connect(nodes[i][1], nodes[i + 1][0])
+
+        out_edges = [nodes[i][-1] for i in range(nwires)]
+        in_edges = [nodes[i][-2] for i in range(nwires)]
+        ignore_edges = [nodes[0][0], nodes[-1][1]]
+
+    qop = quantum_constructor(
+        out_edges,
+        in_edges,
+        [],
+        ignore_edges,
+    )
+    return qop
+
+
+def qop2tn(qop: QuOperator) -> Any:
+    """
+    Convert QuOperator back to MPO or MPS in TensorNetwork package.
+
+    :param qop: MPO or MPS in the form of QuOperator, param in docstring
+    :return: MPO or MPS in the form of TensorNetwork
+    :rtype: Union[tn.matrixproductstates.MPO, tn.matrixproductstates.MPS]
+    """
+    sorted_nodes, is_mps, _ = extract_tensors_from_qop(qop)
+
+    tensors = [node.tensor for node in sorted_nodes]
+
+    if is_mps:
+        return tn.FiniteMPS(tensors, canonicalize=False)
+    else:
+        return tn.matrixproductstates.mpo.FiniteMPO(tensors)
+
+
 # TODO(@refraction-ray): Z2 analogy or more general analogies for the following u1 functions
 
 
@@ -1433,7 +1914,7 @@ def PauliStringSum2Dense(
         return sparsem.todense()
     sparsem = backend.coo_sparse_matrix_from_numpy(sparsem)
     densem = backend.to_dense(sparsem)
-    return densem
+    return backend.convert_to_tensor(densem)
 
 
 # already implemented as backend method
@@ -1755,7 +2236,10 @@ def entanglement_entropy(state: Tensor, cut: Union[int, List[int]]) -> Tensor:
 
 
 def reduced_wavefunction(
-    state: Tensor, cut: List[int], measure: Optional[List[int]] = None
+    state: Tensor,
+    cut: List[int],
+    measure: Optional[List[int]] = None,
+    dim: Optional[int] = None,
 ) -> Tensor:
     """
     Compute the reduced wavefunction from the quantum state ``state``.
@@ -1770,20 +2254,22 @@ def reduced_wavefunction(
     :type measure: List[int]
     :return: _description_
     :rtype: Tensor
+    :param dim: dimension of qudit system
+    :type dim: int
     """
+    dim = 2 if dim is None else dim
     if measure is None:
         measure = [0 for _ in cut]
-    s = backend.reshape2(state)
+    s = backend.reshaped(state, dim)
     n = len(backend.shape_tuple(s))
     s_node = Gate(s)
     end_nodes = []
     for c, m in zip(cut, measure):
-        rt = backend.cast(backend.convert_to_tensor(1 - m), dtypestr) * backend.cast(
-            backend.convert_to_tensor(np.array([1.0, 0.0])), dtypestr
-        ) + backend.cast(backend.convert_to_tensor(m), dtypestr) * backend.cast(
-            backend.convert_to_tensor(np.array([0.0, 1.0])), dtypestr
+        oh = backend.cast(
+            backend.one_hot(backend.cast(backend.convert_to_tensor(m), "int32"), dim),
+            dtypestr,
         )
-        end_node = Gate(rt)
+        end_node = Gate(backend.convert_to_tensor(oh))
         end_nodes.append(end_node)
         s_node[c] ^ end_node[0]
     new_node = contractor(
@@ -1798,8 +2284,9 @@ def reduced_density_matrix(
     cut: Union[int, List[int]],
     p: Optional[Tensor] = None,
     normalize: bool = True,
+    dim: Optional[int] = None,
 ) -> Union[Tensor, QuOperator]:
-    """
+    r"""
     Compute the reduced density matrix from the quantum state ``state``.
 
     :param state: The quantum state in form of Tensor or QuOperator.
@@ -1811,8 +2298,12 @@ def reduced_density_matrix(
     :type p: Optional[Tensor]
     :return: The reduced density matrix.
     :rtype: Union[Tensor, QuOperator]
-    :normalize: if True, returns a trace 1 density matrix. Otherwise does not normalize.
+    :param normalize: if True, returns a trace 1 density matrix. Otherwise, does not normalize.
+    :type normalize: bool
+    :param dim: dimension of qudit system
+    :type dim: int
     """
+    dim = 2 if dim is None else dim
     if isinstance(cut, list) or isinstance(cut, tuple) or isinstance(cut, set):
         traceout = list(cut)
     else:
@@ -1825,21 +2316,19 @@ def reduced_density_matrix(
         return state.partial_trace(traceout)
     if len(state.shape) == 2 and state.shape[0] == state.shape[1]:
         # density operator
-        freedomexp = backend.sizen(state)
-        # traceout = sorted(traceout)[::-1]
-        freedom = int(np.log2(freedomexp) / 2)
-        # traceout2 = [i + freedom for i in traceout]
+        freedom = _infer_num_sites(state.shape[0], dim)
         left = traceout + [i for i in range(freedom) if i not in traceout]
         right = [i + freedom for i in left]
-        rho = backend.reshape(state, [2 for _ in range(2 * freedom)])
+
+        rho = backend.reshape(state, [dim] * (2 * freedom))
         rho = backend.transpose(rho, perm=left + right)
         rho = backend.reshape(
             rho,
             [
-                2 ** len(traceout),
-                2 ** (freedom - len(traceout)),
-                2 ** len(traceout),
-                2 ** (freedom - len(traceout)),
+                dim ** len(traceout),
+                dim ** (freedom - len(traceout)),
+                dim ** len(traceout),
+                dim ** (freedom - len(traceout)),
             ],
         )
         if p is None:
@@ -1852,20 +2341,20 @@ def reduced_density_matrix(
             p = backend.reshape(p, [-1])
             rho = backend.einsum("a,aiaj->ij", p, rho)
         rho = backend.reshape(
-            rho, [2 ** (freedom - len(traceout)), 2 ** (freedom - len(traceout))]
+            rho, [dim ** (freedom - len(traceout)), dim ** (freedom - len(traceout))]
         )
         if normalize:
             rho /= backend.trace(rho)
 
     else:
         w = state / backend.norm(state)
-        freedomexp = backend.sizen(state)
-        freedom = int(np.log(freedomexp) / np.log(2))
+        size = int(backend.sizen(state))
+        freedom = _infer_num_sites(size, dim)
         perm = [i for i in range(freedom) if i not in traceout]
         perm = perm + traceout
-        w = backend.reshape(w, [2 for _ in range(freedom)])
+        w = backend.reshape(w, [dim for _ in range(freedom)])
         w = backend.transpose(w, perm=perm)
-        w = backend.reshape(w, [-1, 2 ** len(traceout)])
+        w = backend.reshape(w, [-1, dim ** len(traceout)])
         if p is None:
             rho = w @ backend.adjoint(w)
         else:
@@ -1914,13 +2403,13 @@ def free_energy(
 
 def renyi_entropy(rho: Union[Tensor, QuOperator], k: int = 2) -> Tensor:
     """
-    Compute the Rényi entropy of order :math:`k` by given density matrix.
+    Compute the Renyi entropy of order :math:`k` by given density matrix.
 
     :param rho: The density matrix in form of Tensor or QuOperator.
     :type rho: Union[Tensor, QuOperator]
-    :param k: The order of Rényi entropy, default is 2.
+    :param k: The order of Renyi entropy, default is 2.
     :type k: int, optional
-    :return: The :math:`k` th order of Rényi entropy.
+    :return: The :math:`k` th order of Renyi entropy.
     :rtype: Tensor
     """
     s = 1 / (1 - k) * backend.real(backend.log(trace_product(*[rho for _ in range(k)])))
@@ -1934,7 +2423,7 @@ def renyi_free_energy(
     k: int = 2,
 ) -> Tensor:
     """
-    Compute the Rényi free energy of the corresponding density matrix and Hamiltonian.
+    Compute the Renyi free energy of the corresponding density matrix and Hamiltonian.
 
     :Example:
 
@@ -1951,9 +2440,9 @@ def renyi_free_energy(
     :type h: Union[Tensor, QuOperator]
     :param beta: Constant for the optimization, default is 1.
     :type beta: float, optional
-    :param k: The order of Rényi entropy, default is 2.
+    :param k: The order of Renyi entropy, default is 2.
     :type k: int, optional
-    :return: The :math:`k` th order of Rényi entropy.
+    :return: The :math:`k` th order of Renyi entropy.
     :rtype: Tensor
     """
     energy = backend.real(trace_product(rho, h))
@@ -2008,7 +2497,9 @@ def truncated_free_energy(
 
 
 @op2tensor
-def partial_transpose(rho: Tensor, transposed_sites: List[int]) -> Tensor:
+def partial_transpose(
+    rho: Tensor, transposed_sites: List[int], dim: Optional[int] = None
+) -> Tensor:
     """
     _summary_
 
@@ -2016,10 +2507,13 @@ def partial_transpose(rho: Tensor, transposed_sites: List[int]) -> Tensor:
     :type rho: Tensor
     :param transposed_sites: sites int list to be transposed
     :type transposed_sites: List[int]
+    :param dim: dimension of qudit system
+    :type dim: int
     :return: _description_
     :rtype: Tensor
     """
-    rho = backend.reshape2(rho)
+    dim = 2 if dim is None else dim
+    rho = backend.reshaped(rho, dim)
     rho_node = Gate(rho)
     n = len(rho.shape) // 2
     left_edges = []
@@ -2037,7 +2531,9 @@ def partial_transpose(rho: Tensor, transposed_sites: List[int]) -> Tensor:
 
 
 @op2tensor
-def entanglement_negativity(rho: Tensor, transposed_sites: List[int]) -> Tensor:
+def entanglement_negativity(
+    rho: Tensor, transposed_sites: List[int], dim: Optional[int] = None
+) -> Tensor:
     """
     _summary_
 
@@ -2045,17 +2541,21 @@ def entanglement_negativity(rho: Tensor, transposed_sites: List[int]) -> Tensor:
     :type rho: Tensor
     :param transposed_sites: _description_
     :type transposed_sites: List[int]
+    :param dim: dimension of qudit system
+    :type dim: int
     :return: _description_
     :rtype: Tensor
     """
-    rhot = partial_transpose(rho, transposed_sites)
+    rhot = partial_transpose(rho, transposed_sites, dim=dim)
     es = backend.eigvalsh(rhot)
     rhot_m = backend.sum(backend.abs(es))
     return (rhot_m - 1.0) / 2.0
 
 
 @op2tensor
-def log_negativity(rho: Tensor, transposed_sites: List[int], base: str = "e") -> Tensor:
+def log_negativity(
+    rho: Tensor, transposed_sites: List[int], base: str = "e", dim: Optional[int] = None
+) -> Tensor:
     """
     _summary_
 
@@ -2065,10 +2565,13 @@ def log_negativity(rho: Tensor, transposed_sites: List[int], base: str = "e") ->
     :type transposed_sites: List[int]
     :param base: whether use 2 based log or e based log, defaults to "e"
     :type base: str, optional
+    :param dim: dimension of qudit system
+    :type dim: int
     :return: _description_
     :rtype: Tensor
     """
-    rhot = partial_transpose(rho, transposed_sites)
+    dim = 2 if dim is None else dim
+    rhot = partial_transpose(rho, transposed_sites, dim)
     es = backend.eigvalsh(rhot)
     rhot_m = backend.sum(backend.abs(es))
     een = backend.log(rhot_m)
@@ -2154,7 +2657,9 @@ def double_state(h: Tensor, beta: float = 1) -> Tensor:
 
 
 @op2tensor
-def mutual_information(s: Tensor, cut: Union[int, List[int]]) -> Tensor:
+def mutual_information(
+    s: Tensor, cut: Union[int, List[int]], dim: Optional[int] = None
+) -> Tensor:
     """
     Mutual information between AB subsystem described by ``cut``.
 
@@ -2162,9 +2667,12 @@ def mutual_information(s: Tensor, cut: Union[int, List[int]]) -> Tensor:
     :type s: Tensor
     :param cut: The AB subsystem.
     :type cut: Union[int, List[int]]
+    :param dim: The diagonal matrix in form of Tensor.
+    :type dim: Tensor
     :return: The mutual information between AB subsystem described by ``cut``.
     :rtype: Tensor
     """
+    dim = 2 if dim is None else dim
     if isinstance(cut, list) or isinstance(cut, tuple) or isinstance(cut, set):
         traceout = list(cut)
     else:
@@ -2172,22 +2680,22 @@ def mutual_information(s: Tensor, cut: Union[int, List[int]]) -> Tensor:
 
     if len(s.shape) == 2 and s.shape[0] == s.shape[1]:
         # mixed state
-        n = int(np.log2(backend.sizen(s)) / 2)
+        n = _infer_num_sites(s.shape[0], dim=dim)
         hab = entropy(s)
 
         # subsystem a
-        rhoa = reduced_density_matrix(s, traceout)
+        rhoa = reduced_density_matrix(s, traceout, dim=dim)
         ha = entropy(rhoa)
 
         # need subsystem b as well
         other = tuple(i for i in range(n) if i not in traceout)
-        rhob = reduced_density_matrix(s, other)  # type: ignore
+        rhob = reduced_density_matrix(s, other, dim=dim)  # type: ignore
         hb = entropy(rhob)
 
     # pure system
     else:
         hab = 0.0
-        rhoa = reduced_density_matrix(s, traceout)
+        rhoa = reduced_density_matrix(s, traceout, dim=dim)
         ha = hb = entropy(rhoa)
 
     return ha + hb - hab
@@ -2196,7 +2704,9 @@ def mutual_information(s: Tensor, cut: Union[int, List[int]]) -> Tensor:
 # measurement results and transformations and correlations below
 
 
-def count_s2d(srepr: Tuple[Tensor, Tensor], n: int) -> Tensor:
+def count_s2d(
+    srepr: Tuple[Tensor, Tensor], n: int, dim: Optional[int] = None
+) -> Tensor:
     """
     measurement shots results, sparse tuple representation to dense representation
     count_vector to count_tuple
@@ -2205,11 +2715,14 @@ def count_s2d(srepr: Tuple[Tensor, Tensor], n: int) -> Tensor:
     :type srepr: Tuple[Tensor, Tensor]
     :param n: number of qubits
     :type n: int
+    :param dim: [description], defaults to None
+    :type dim: int, optional
     :return: [description]
     :rtype: Tensor
     """
+    dim = 2 if dim is None else dim
     return backend.scatter(
-        backend.cast(backend.zeros([2**n]), srepr[1].dtype),
+        backend.cast(backend.zeros([dim**n]), srepr[1].dtype),
         backend.reshape(srepr[0], [-1, 1]),
         srepr[1],
     )
@@ -2252,117 +2765,146 @@ def count_d2s(drepr: Tensor, eps: float = 1e-7) -> Tuple[Tensor, Tensor]:
 count_t2v = count_d2s
 
 
-def sample_int2bin(sample: Tensor, n: int) -> Tensor:
+def sample_int2bin(sample: Tensor, n: int, dim: Optional[int] = None) -> Tensor:
     """
-    int sample to bin sample
+    Convert linear-index samples to per-site digits (base-d).
 
-    :param sample: in shape [trials] of int elements in the range [0, 2**n)
+    :param sample: shape [trials], integers in [0, d**n)
     :type sample: Tensor
-    :param n: number of qubits
+    :param n: number of sites
     :type n: int
-    :return: in shape [trials, n] of element (0, 1)
+    :param dim: local dimension, defaults to 2
+    :type dim: int, optional
+    :return: shape [trials, n], entries in [0, d-1]
     :rtype: Tensor
     """
-    confg = backend.mod(
-        backend.right_shift(sample[..., None], backend.reverse(backend.arange(n))),
-        2,
-    )
-    return confg
+    dim = 2 if dim is None else dim
+    if dim == 2:
+        return backend.mod(
+            backend.right_shift(sample[..., None], backend.reverse(backend.arange(n))),
+            2,
+        )
+    else:
+        pos = backend.reverse(backend.arange(n))
+        base = backend.power(dim, pos)
+        digits = backend.mod(
+            backend.floor_divide(sample[..., None], base),  # ⌊sample / d**pos⌋
+            dim,
+        )
+        return backend.cast(digits, "int32")
 
 
-def sample_bin2int(sample: Tensor, n: int) -> Tensor:
+def sample_bin2int(sample: Tensor, n: int, dim: Optional[int] = None) -> Tensor:
     """
     bin sample to int sample
 
     :param sample: in shape [trials, n] of elements (0, 1)
     :type sample: Tensor
-    :param n: number of qubits
+    :param n: number of sites
     :type n: int
+    :param dim: local dimension, defaults to 2
+    :type dim: int, optional
     :return: in shape [trials]
     :rtype: Tensor
     """
-    power = backend.convert_to_tensor([2**j for j in reversed(range(n))])
+    dim = 2 if dim is None else dim
+    power = backend.convert_to_tensor([dim**j for j in reversed(range(n))])
     return backend.sum(sample * power, axis=-1)
 
 
 def sample2count(
-    sample: Tensor, n: int, jittable: bool = True
+    sample: Tensor,
+    n: int,
+    jittable: bool = True,
+    dim: Optional[int] = None,
 ) -> Tuple[Tensor, Tensor]:
     """
-    sample_int to count_tuple
+    sample_int to count_tuple (indices, counts), size = d**n
 
-    :param sample: _description_
+    :param sample: linear-index samples, shape [shots]
     :type sample: Tensor
-    :param n: _description_
+    :param n: number of sites
     :type n: int
-    :param jittable: _description_, defaults to True
-    :type jittable: bool, optional
-    :return: _description_
+    :param jittable: whether to return fixed-size outputs (backend dependent)
+    :type jittable: bool
+    :param dim: local dimension per site, default 2 (qubit)
+    :type dim: int, optional
+    :return: (unique_indices, counts)
     :rtype: Tuple[Tensor, Tensor]
     """
-    d = 2**n
+    dim = 2 if dim is None else dim
+    size = dim**n
     if not jittable:
         results = backend.unique_with_counts(sample)  # non-jittable
-    else:  # jax specified
-        results = backend.unique_with_counts(sample, size=d, fill_value=-1)
+    else:  # jax specified / fixed-size
+        results = backend.unique_with_counts(sample, size=size, fill_value=-1)
     return results
 
 
-def count_vector2dict(count: Tensor, n: int, key: str = "bin") -> Dict[Any, int]:
+def count_vector2dict(
+    count: Tensor, n: int, key: str = "bin", dim: Optional[int] = None
+) -> Dict[Any, int]:
     """
-    convert_vector to count_dict_bin or count_dict_int
+    Convert count_vector to count_dict_bin or count_dict_int.
+    For d>10 cases, a base-d string (0-9A-Z) is used.
 
-    :param count: tensor in shape [2**n]
+    :param count: tensor in shape [d**n]
     :type count: Tensor
-    :param n: number of qubits
+    :param n: number of sites
     :type n: int
     :param key: can be "int" or "bin", defaults to "bin"
     :type key: str, optional
-    :return: _description_
-    :rtype: _type_
+    :param dim: local dimension (default 2)
+    :type dim: int, optional
+    :return: mapping from configuration to count
+    :rtype: Dict[Any, int]
     """
     from .interfaces import which_backend
 
+    dim = 2 if dim is None else dim
     b = which_backend(count)
-    d = {i: b.numpy(count[i]).item() for i in range(2**n)}
+    out_int = {i: b.numpy(count[i]).item() for i in range(dim**n)}
     if key == "int":
-        return d
+        return out_int
     else:
-        dn = {}
-        for k, v in d.items():
-            kn = str(bin(k))[2:].zfill(n)
-            dn[kn] = v
-        return dn
+        out_str = {}
+        for k, v in out_int.items():
+            kn = np.base_repr(k, base=dim).zfill(n)
+            out_str[kn] = v
+        return out_str
 
 
 def count_tuple2dict(
-    count: Tuple[Tensor, Tensor], n: int, key: str = "bin"
+    count: Tuple[Tensor, Tensor], n: int, key: str = "bin", dim: Optional[int] = None
 ) -> Dict[Any, int]:
     """
     count_tuple to count_dict_bin or count_dict_int
 
-    :param count: count_tuple format
+    :param count: count_tuple format (indices, counts)
     :type count: Tuple[Tensor, Tensor]
-    :param n: number of qubits
+    :param n: number of sites (qubits or qudits)
     :type n: int
     :param key: can be "int" or "bin", defaults to "bin"
     :type key: str, optional
+    :param dim: local dimension, defaults to 2
+    :type dim: int, optional
     :return: count_dict
-    :rtype: _type_
+    :rtype: Dict[Any, int]
     """
-    d = {
+    dim = 2 if dim is None else dim
+    out_int = {
         backend.numpy(i).item(): backend.numpy(j).item()
         for i, j in zip(count[0], count[1])
         if i >= 0
     }
     if key == "int":
-        return d
+        return out_int
     else:
-        dn = {}
-        for k, v in d.items():
-            kn = str(bin(k))[2:].zfill(n)
-            dn[kn] = v
-        return dn
+        out_str = {}
+        for k, v in out_int.items():
+            kn = np.base_repr(k, base=dim).zfill(n)
+            out_str[kn] = v
+        return out_str
 
 
 @partial(arg_alias, alias_dict={"counts": ["shots"], "format": ["format_"]})
@@ -2374,8 +2916,9 @@ def measurement_counts(
     random_generator: Optional[Any] = None,
     status: Optional[Tensor] = None,
     jittable: bool = False,
+    dim: Optional[int] = None,
 ) -> Any:
-    """
+    r"""
     Simulate the measuring of each qubit of ``p`` in the computational basis,
     thus producing output like that of ``qiskit``.
 
@@ -2390,6 +2933,7 @@ def measurement_counts(
     "count_tuple": # (np.array([0]), np.array([2]))
 
     "count_dict_bin": # {"00": 2, "01": 0, "10": 0, "11": 0}
+    / for cases d\in [10, 36], "10" -> "A", ..., "35" -> "Z"
 
     "count_dict_int": # {0: 2, 1: 0, 2: 0, 3: 0}
 
@@ -2441,21 +2985,22 @@ def measurement_counts(
             state /= backend.norm(state)
             pi = backend.real(backend.conj(state) * state)
         pi = backend.reshape(pi, [-1])
-    d = int(backend.shape_tuple(pi)[0])
-    n = int(np.log(d) / np.log(2) + 1e-8)
+
+    local_d = 2 if dim is None else dim
+    total_dim = int(backend.shape_tuple(pi)[0])
+    n = _infer_num_sites(total_dim, local_d)
+
     if (counts is None) or counts <= 0:
         if format == "count_vector":
             return pi
         elif format == "count_tuple":
             return count_d2s(pi)
         elif format == "count_dict_bin":
-            return count_vector2dict(pi, n, key="bin")
+            return count_vector2dict(pi, n, key="bin", dim=local_d)
         elif format == "count_dict_int":
-            return count_vector2dict(pi, n, key="int")
+            return count_vector2dict(pi, n, key="int", dim=local_d)
         else:
-            raise ValueError(
-                "unsupported format %s for analytical measurement" % format
-            )
+            raise ValueError(f"unsupported format {format} for analytical measurement")
     else:
         raw_counts = backend.probability_sample(
             counts, pi, status=status, g=random_generator
@@ -2466,7 +3011,7 @@ def measurement_counts(
         # raw_counts = backend.stateful_randc(
         # random_generator, a=drange, shape=counts, p=pi
         # )
-        return sample2all(raw_counts, n, format=format, jittable=jittable)
+        return sample2all(raw_counts, n, format=format, jittable=jittable, dim=local_d)
 
 
 measurement_results = measurement_counts
@@ -2474,45 +3019,64 @@ measurement_results = measurement_counts
 
 @partial(arg_alias, alias_dict={"format": ["format_"]})
 def sample2all(
-    sample: Tensor, n: int, format: str = "count_vector", jittable: bool = False
+    sample: Tensor,
+    n: int,
+    format: str = "count_vector",
+    jittable: bool = False,
+    dim: Optional[int] = None,
 ) -> Any:
     """
-    transform ``sample_int`` or ``sample_bin`` form results to other forms specified by ``format``
+    transform ``sample_int`` or ``sample_bin`` results to other forms specified by ``format``
 
-    :param sample: measurement shots results in ``sample_int`` or ``sample_bin`` format
+    :param sample: measurement shots results in ``sample_int`` (shape [shots]) or ``sample_bin`` (shape [shots, n])
     :type sample: Tensor
-    :param n: number of qubits
+    :param n: number of sites
     :type n: int
-    :param format: see the doc in the doc in :py:meth:`tensorcircuit.quantum.measurement_results`,
-        defaults to "count_vector"
+    :param format: see :py:meth:`tensorcircuit.quantum.measurement_results`, defaults to "count_vector"
     :type format: str, optional
     :param jittable: only applicable to count transformation in jax backend, defaults to False
     :type jittable: bool, optional
+    :param dim: local dimension (2 for qubit; >2 for qudit), defaults to 2
+    :type dim: Optional[int]
     :return: measurement results specified as ``format``
     :rtype: Any
     """
+    dim = 2 if dim is None else int(dim)
+    n_max_d = int(32 / np.log2(dim))
+    if n > n_max_d:
+        assert (
+            len(backend.shape_tuple(sample)) == 2
+        ), f"n>{n_max_d} is only supported for ``sample_bin``"
+        if format == "sample_bin":
+            return sample
+        if format == "count_dict_bin":
+            binary_strings = ["".join(map(str, shot)) for shot in sample]
+            return dict(Counter(binary_strings))
+        raise ValueError(f"n={n} is too large for measurement representaion: {format}")
+
     if len(backend.shape_tuple(sample)) == 1:
         sample_int = sample
-        sample_bin = sample_int2bin(sample, n)
+        sample_bin = sample_int2bin(sample, n, dim=dim)
     elif len(backend.shape_tuple(sample)) == 2:
-        sample_int = sample_bin2int(sample, n)
+        sample_int = sample_bin2int(sample, n, dim=dim)
         sample_bin = sample
     else:
         raise ValueError("unrecognized tensor shape for sample")
+
     if format == "sample_int":
         return sample_int
     elif format == "sample_bin":
         return sample_bin
     else:
-        count_tuple = sample2count(sample_int, n, jittable)
+        count_tuple = sample2count(sample_int, n, jittable=jittable, dim=dim)
         if format == "count_tuple":
             return count_tuple
         elif format == "count_vector":
-            return count_s2d(count_tuple, n)
+            return count_s2d(count_tuple, n, dim=dim)
         elif format == "count_dict_bin":
-            return count_tuple2dict(count_tuple, n, key="bin")
+            return count_tuple2dict(count_tuple, n, key="bin", dim=dim)
         elif format == "count_dict_int":
-            return count_tuple2dict(count_tuple, n, key="int")
+            return count_tuple2dict(count_tuple, n, key="int", dim=dim)
         else:
             raise ValueError(
                 "unsupported format %s for finite shots measurement" % format