add some docstrings

Author: refraction-ray

tensorcircuit/results/counts.py

@@ -12,22 +12,76 @@
 
 
 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
     return ncount
 
 
 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)
@@ -36,6 +90,21 @@ def marginal_count(count: ct, keep_list: Sequence[int]) -> ct:
 
 
 def count2vec(count: ct, normalization: bool = True) -> Tensor:
+    """
+    Convert count dictionary to probability vector.
+
+    :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
+    :return: Probability vector as numpy array
+    :rtype: Tensor
+
+    :Example:
+
+    >>> count2vec({"00": 2, "10": 3, "11": 5})
+    array([0.2, 0. , 0.3, 0.5])
+    """
     nqubit = len(list(count.keys())[0])
     probability = [0] * 2**nqubit
     shots = sum([v for k, v in count.items()])
@@ -47,6 +116,21 @@ def count2vec(count: ct, normalization: bool = True) -> Tensor:
 
 
 def vec2count(vec: Tensor, prune: bool = False) -> ct:
+    """
+    Convert probability vector to count dictionary.
+
+    :param vec: Probability vector
+    :type vec: Tensor
+    :param prune: Whether to remove near-zero probabilities, defaults to False
+    :type prune: bool, optional
+    :return: Count dictionary
+    :rtype: ct
+
+    :Example:
+
+    >>> vec2count(np.array([0.2, 0.3, 0.1, 0.4]))
+    {'00': 0.2, '01': 0.3, '10': 0.1, '11': 0.4}
+    """
     from ..quantum import count_vector2dict
 
     if isinstance(vec, list):
@@ -63,6 +147,16 @@ def vec2count(vec: Tensor, prune: bool = False) -> ct:
 
 
 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
@@ -113,6 +207,11 @@ def merge_count(*counts: ct) -> ct:
     :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:

tensorcircuit/utils.py

@@ -11,6 +11,13 @@
 
 
 def gpu_memory_share(flag: bool = True) -> None:
+    """
+    Set the GPU memory growth mode
+
+    :param flag: whether to set the GPU memory growth mode, defaults to True
+    :type flag: bool
+    :return: None
+    """
     # TODO(@refraction-ray): the default torch behavior should be True
     # preallocate behavior for torch to be investigated
     if flag is True: