tensorcircuit-nightly 1.0.2.dev20250108__py3-none-any.whl → 1.4.0.dev20251103__py3-none-any.whl

tensorcircuit/results/counts.py

@@ -6,14 +6,26 @@ from typing import Any, Dict, Optional, Sequence
 
 import numpy as np
 
+from ..quantum import _decode_basis_label
 
 Tensor = Any
 ct = Dict[str, int]
 
-# TODO(@refraction-ray): merge_count
-
 
 def reverse_count(count: ct) -> ct:
+    """
+    Reverse the bit string keys in a count dictionary.
+
+    :param count: A dictionary mapping bit strings to counts
+    :type count: ct
+    :return: A new dictionary with reversed bit string keys
+    :rtype: ct
+
+    :Example:
+
+    >>> reverse_count({"01": 10, "10": 20})
+    {'10': 10, '01': 20}
+    """
     ncount = {}
     for k, v in count.items():
         ncount[k[::-1]] = v
@@ -21,15 +33,56 @@ def reverse_count(count: ct) -> ct:
 
 
 def sort_count(count: ct) -> ct:
+    """
+    Sort the count dictionary by counts in descending order.
+
+    :param count: A dictionary mapping bit strings to counts
+    :type count: ct
+    :return: A new dictionary sorted by count values (descending)
+    :rtype: ct
+
+    :Example:
+
+    >>> sort_count({"00": 5, "01": 15, "10": 10})
+    {'01': 15, '10': 10, '00': 5}
+    """
     return {k: v for k, v in sorted(count.items(), key=lambda item: -item[1])}
 
 
 def normalized_count(count: ct) -> Dict[str, float]:
+    """
+    Normalize the count dictionary to represent probabilities.
+
+    :param count: A dictionary mapping bit strings to counts
+    :type count: ct
+    :return: A new dictionary with probabilities instead of counts
+    :rtype: Dict[str, float]
+
+    :Example:
+
+    >>> normalized_count({"00": 5, "01": 15})
+    {'00': 0.25, '01': 0.75}
+    """
     shots = sum([v for k, v in count.items()])
     return {k: v / shots for k, v in count.items()}
 
 
 def marginal_count(count: ct, keep_list: Sequence[int]) -> ct:
+    """
+    Compute the marginal distribution of a count dictionary over specified qubits.
+
+    :param count: A dictionary mapping bit strings to counts
+    :type count: ct
+    :param keep_list: List of qubit indices to keep in the marginal distribution
+    :type keep_list: Sequence[int]
+    :return: A new count dictionary with marginal distribution
+    :rtype: ct
+
+    :Example:
+
+    >>> marginal_count({"001": 10, "110": 20}, [0, 2])
+    {'01': 10, '10': 20}
+    """
     import qiskit
 
     count = reverse_count(count)
@@ -37,34 +90,87 @@ def marginal_count(count: ct, keep_list: Sequence[int]) -> ct:
     return reverse_count(ncount)
 
 
-def count2vec(count: ct, normalization: bool = True) -> Tensor:
-    nqubit = len(list(count.keys())[0])
-    probability = [0] * 2**nqubit
-    shots = sum([v for k, v in count.items()])
+def count2vec(
+    count: ct, normalization: bool = True, dim: Optional[int] = None
+) -> Tensor:
+    """
+    Convert a dictionary of counts (with string keys) to a probability/count vector.
+
+    Support:
+      - base-d string (d <= 36), characters taken from 0-9A-Z (case-insensitive)
+        For example:
+          qubit: '0101'
+          qudit: '012' or '09A' (A represents 10, which means [0, 9, 10])
+
+    :param count: A dictionary mapping bit strings to counts
+    :type count: ct
+    :param normalization: Whether to normalize the counts to probabilities, defaults to True
+    :type normalization: bool, optional
+    :param dim: Dimensionality of the vector, defaults to 2
+    :type dim: int, optional
+    :return: Probability vector as numpy array
+    :rtype: Tensor
+
+    :Example:
+
+    >>> count2vec({"00": 2, "10": 3, "11": 5})
+    array([0.2, 0. , 0.3, 0.5])
+    """
+
+    if not count:
+        return np.array([], dtype=float)
+
+    dim = 2 if dim is None else dim
+
+    n = len(next(iter(count)).upper())
+    prob = np.zeros(dim**n, dtype=float)
+    shots = float(sum(count.values())) if normalization else 1.0
+    if shots == 0:
+        return prob
+
+    powers = [dim**p for p in range(n)][::-1]
     for k, v in count.items():
-        if normalization is True:
-            v /= shots  # type: ignore
-        probability[int(k, 2)] = v
-    return np.array(probability)
+        digits = _decode_basis_label(k, n, dim)
+        idx = sum(dig * p for dig, p in zip(digits, powers))
+        prob[idx] = (v / shots) if normalization else v
 
+    return prob
+
+
+def vec2count(vec: Tensor, prune: bool = False, dim: Optional[int] = None) -> ct:
+    """
+    Map a count/probability vector of length D to a dictionary with base-d string keys (0-9A-Z).
+    Only generate string keys when d <= 36; if d is inferred to be > 36, raise a NotImplementedError.
 
-def vec2count(vec: Tensor, prune: bool = False) -> ct:
-    from ..quantum import count_vector2dict
+    :param vec: A one-dimensional vector of length D = d**n
+    :param prune: Whether to prune near-zero elements (threshold 1e-8)
+    :param dim: Dimensionality of the vector, defaults to 2
+    :return: {base-d string key: value}, key length n
+    """
+    from ..quantum import count_vector2dict, _infer_num_sites
 
+    dim = 2 if dim is None else dim
     if isinstance(vec, list):
         vec = np.array(vec)
-    n = int(np.log(vec.shape[0]) / np.log(2) + 1e-9)
-    c = count_vector2dict(vec, n, key="bin")
-    if prune is True:
-        nc = c.copy()
-        for k, v in c.items():
-            if np.abs(v) < 1e-8:
-                del nc[k]
-        return nc
+    n = _infer_num_sites(int(vec.shape[0]), dim)
+    c: ct = count_vector2dict(vec, n, key="bin", dim=dim)  # type: ignore
+    if prune:
+        c = {k: v for k, v in c.items() if np.abs(v) >= 1e-8}
+
     return c
 
 
 def kl_divergence(c1: ct, c2: ct) -> float:
+    """
+    Compute the Kullback-Leibler divergence between two count distributions.
+
+    :param c1: First count dictionary
+    :type c1: ct
+    :param c2: Second count dictionary
+    :type c2: ct
+    :return: KL divergence value
+    :rtype: float
+    """
     eps = 1e-4  # typical value for inverse of the total shots
     c1 = normalized_count(c1)  # type: ignore
     c2 = normalized_count(c2)  # type: ignore
@@ -107,6 +213,27 @@ def expectation(
     return r / shots
 
 
+def merge_count(*counts: ct) -> ct:
+    """
+    Merge multiple count dictionaries by summing up their counts
+
+    :param counts: Variable number of count dictionaries
+    :type counts: ct
+    :return: Merged count dictionary
+    :rtype: ct
+
+    :Example:
+
+    >>> merge_count({"00": 10, "01": 20}, {"00": 5, "10": 15})
+    {'00': 15, '01': 20, '10': 15}
+    """
+    merged: ct = {}
+    for count in counts:
+        for k, v in count.items():
+            merged[k] = merged.get(k, 0) + v
+    return merged
+
+
 def plot_histogram(data: Any, **kws: Any) -> Any:
     """
     See ``qiskit.visualization.plot_histogram``:

tensorcircuit/results/readout_mitigation.py

@@ -723,7 +723,10 @@ class ReadoutMit:
         cals = self._form_cals(qubits)
         M = M3MatVec(dict(counts), cals, distance)
         L = spla.LinearOperator(
-            (M.num_elems, M.num_elems), matvec=M.matvec, rmatvec=M.rmatvec
+            (M.num_elems, M.num_elems),
+            matvec=M.matvec,
+            rmatvec=M.rmatvec,
+            dtype=np.float64,
         )
         diags = M.get_diagonal()
 

tensorcircuit/shadows.py

@@ -334,7 +334,7 @@ def entropy_shadow(
 
 def renyi_entropy_2(snapshots: Tensor, sub: Optional[Sequence[int]] = None) -> Tensor:
     r"""To calculate the second order Renyi entropy of a subsystem from snapshot, please refer to
-    Brydges, T. et al. Science 364, 260–263 (2019). This function is not jitable.
+    Brydges, T. et al. Science 364, 260-263 (2019). This function is not jitable.
 
     :param snapshots: shape = (ns, repeat, nq)
     :type: Tensor

tensorcircuit/simplify.py

@@ -121,7 +121,9 @@ def _split_two_qubit_gate(
     if fixed_choice == 2:  # swap one
         return n3, n4, True  # swap
     s2 = n3.tensor.shape[-1]
-    if (s1 >= 4) and (s2 >= 4):
+    if (s1 >= n[0].dimension * n[2].dimension) and (
+        s2 >= n[1].dimension * n[3].dimension
+    ):
         # jax jit unspport split_node with trun_err anyway
         # tf function doesn't work either, though I believe it may work on tf side
         # CANNOT DONE(@refraction-ray): tf.function version with trun_err set