Merge pull request #1175 from qiboteam/total_variation_distance

Add `quantum_info.utils.total_variation_distance` and improve documentation of `quantum_info.utils`

doc/source/api-reference/qibo.rst

@@ -2081,6 +2081,12 @@ Shannon entropy
 .. autofunction:: qibo.quantum_info.shannon_entropy
 
 
+Total Variation distance
+""""""""""""""""""""""""
+
+.. autofunction:: qibo.quantum_info.total_variation_distance
+
+
 Hellinger distance
 """"""""""""""""""
 

src/qibo/quantum_info/utils.py

@@ -191,7 +191,7 @@ def hadamard_transform(array, implementation: str = "fast", backend=None):
 
 
 def shannon_entropy(probability_array, base: float = 2, backend=None):
-    """Calculate the Shannon entropy of a probability array :math:`\\mathbf{p}`, which is given by
+    """Calculates the Shannon entropy of a probability array :math:`\\mathbf{p}`, which is given by
 
     .. math::
         H(\\mathbf{p}) = - \\sum_{k = 0}^{d^{2} - 1} \\, p_{k} \\, \\log_{b}(p_{k}) \\, ,
@@ -214,7 +214,7 @@ def shannon_entropy(probability_array, base: float = 2, backend=None):
         backend = GlobalBackend()
 
     if isinstance(probability_array, list):
-        probability_array = backend.cast(probability_array, dtype=float)
+        probability_array = backend.cast(probability_array, dtype=np.float64)
 
     if base <= 0:
         raise_error(ValueError, "log base must be non-negative.")
@@ -256,10 +256,72 @@ def shannon_entropy(probability_array, base: float = 2, backend=None):
     return complex(entropy).real
 
 
+def total_variation_distance(
+    prob_dist_p, prob_dist_q, validate: bool = False, backend=None
+):
+    """Calculates the Total Variation (TV) distance between two discrete probability distributions.
+
+    For probabilities :math:`\\mathbf{p}` and :math:`\\mathbf{q}`, it is defined as
+
+    .. math::
+        d_{\\text{TV}}(\\mathbf{p} \\, , \\, \\mathbf{q}) = \\frac{1}{2}
+            \\, \\| \\mathbf{p} - \\mathbf{q} \\|_{1} \\, ,
+
+    where :math:`\\| \\cdot \\|_{1}` is the :math:`\\mathcal{l}_{1}`-norm.
+
+    Args:
+        prob_dist_p (ndarray or list): discrete probability distribution :math:`p`.
+        prob_dist_q (ndarray or list): discrete probability distribution :math:`q`.
+        validate (bool, optional): If ``True``, checks if :math:`p` and :math:`q` are proper
+            probability distributions. Defaults to ``False``.
+        backend (:class:`qibo.backends.abstract.Backend`, optional): backend to be
+            used in the execution. If ``None``, it uses
+            :class:`qibo.backends.GlobalBackend`. Defaults to ``None``.
+
+    Returns:
+        float: Total variation distance between :math:`\\mathbf{p}` and :math:`\\mathbf{q}`.
+    """
+    if backend is None:  # pragma: no cover
+        backend = GlobalBackend()
+
+    if isinstance(prob_dist_p, list):
+        prob_dist_p = backend.cast(prob_dist_p, dtype=np.float64)
+    if isinstance(prob_dist_q, list):
+        prob_dist_q = backend.cast(prob_dist_q, dtype=np.float64)
+
+    if (len(prob_dist_p.shape) != 1) or (len(prob_dist_q.shape) != 1):
+        raise_error(
+            TypeError,
+            "Probability arrays must have dims (k,) but have "
+            + f"dims {prob_dist_p.shape} and {prob_dist_q.shape}.",
+        )
+
+    if (len(prob_dist_p) == 0) or (len(prob_dist_q) == 0):
+        raise_error(TypeError, "At least one of the arrays is empty.")
+
+    if validate:
+        if (any(prob_dist_p < 0) or any(prob_dist_p > 1.0)) or (
+            any(prob_dist_q < 0) or any(prob_dist_q > 1.0)
+        ):
+            raise_error(
+                ValueError,
+                "All elements of the probability array must be between 0. and 1..",
+            )
+        if np.abs(np.sum(prob_dist_p) - 1.0) > PRECISION_TOL:
+            raise_error(ValueError, "First probability array must sum to 1.")
+
+        if np.abs(np.sum(prob_dist_q) - 1.0) > PRECISION_TOL:
+            raise_error(ValueError, "Second probability array must sum to 1.")
+
+    total_variation = 0.5 * np.sum(np.abs(prob_dist_p - prob_dist_q))
+
+    return total_variation
+
+
 def hellinger_distance(prob_dist_p, prob_dist_q, validate: bool = False, backend=None):
-    """Calculate the Hellinger distance :math:`H(p, q)` between
-    two discrete probability distributions, :math:`\\mathbf{p}` and :math:`\\mathbf{q}`.
-    It is defined as
+    """Calculates the Hellinger distance :math:`H` between two discrete probability distributions.
+
+    For probabilities :math:`\\mathbf{p}` and :math:`\\mathbf{q}`, it is defined as
 
     .. math::
         H(\\mathbf{p} \\, , \\, \\mathbf{q}) = \\frac{1}{\\sqrt{2}} \\, \\|
@@ -268,10 +330,10 @@ def hellinger_distance(prob_dist_p, prob_dist_q, validate: bool = False, backend
     where :math:`\\|\\cdot\\|_{2}` is the Euclidean norm.
 
     Args:
-        prob_dist_p: (discrete) probability distribution :math:`p`.
-        prob_dist_q: (discrete) probability distribution :math:`q`.
-        validate (bool): if True, checks if :math:`p` and :math:`q` are proper
-            probability distributions. Default: False.
+        prob_dist_p (ndarray or list): discrete probability distribution :math:`p`.
+        prob_dist_q (ndarray or list): discrete probability distribution :math:`q`.
+        validate (bool, optional): If ``True``, checks if :math:`p` and :math:`q` are proper
+            probability distributions. Defaults to ``False``.
         backend (:class:`qibo.backends.abstract.Backend`, optional): backend to be
             used in the execution. If ``None``, it uses
             :class:`qibo.backends.GlobalBackend`. Defaults to ``None``.
@@ -283,9 +345,9 @@ def hellinger_distance(prob_dist_p, prob_dist_q, validate: bool = False, backend
         backend = GlobalBackend()
 
     if isinstance(prob_dist_p, list):
-        prob_dist_p = backend.cast(prob_dist_p, dtype=float)
+        prob_dist_p = backend.cast(prob_dist_p, dtype=np.float64)
     if isinstance(prob_dist_q, list):
-        prob_dist_q = backend.cast(prob_dist_q, dtype=float)
+        prob_dist_q = backend.cast(prob_dist_q, dtype=np.float64)
 
     if (len(prob_dist_p.shape) != 1) or (len(prob_dist_q.shape) != 1):
         raise_error(
@@ -319,15 +381,20 @@ def hellinger_distance(prob_dist_p, prob_dist_q, validate: bool = False, backend
 
 
 def hellinger_fidelity(prob_dist_p, prob_dist_q, validate: bool = False, backend=None):
-    """Calculate the Hellinger fidelity between two discrete
-    probability distributions, :math:`p` and :math:`q`. The fidelity is
-    defined as :math:`(1 - H^{2}(p, q))^{2}`, where :math:`H(p, q)`
-    is the Hellinger distance.
+    """Calculates the Hellinger fidelity between two discrete probability distributions.
+
+    For probabilities :math:`p` and :math:`q`, the fidelity is defined as
+
+    .. math::
+        (1 - H^{2}(p, q))^{2} \\, ,
+
+    where :math:`H(p, q)` is the Hellinger distance
+    (:func:`qibo.quantum_info.utils.hellinger_distance`).
 
     Args:
-        prob_dist_p: (discrete) probability distribution :math:`p`.
-        prob_dist_q: (discrete) probability distribution :math:`q`.
-        validate (bool): if True, checks if :math:`p` and :math:`q` are proper
+        prob_dist_p (ndarray or list): discrete probability distribution :math:`p`.
+        prob_dist_q (ndarray or list): discrete probability distribution :math:`q`.
+        validate (bool, optional): if True, checks if :math:`p` and :math:`q` are proper
             probability distributions. Default: False.
         backend (:class:`qibo.backends.abstract.Backend`, optional): backend to be
             used in the execution. If ``None``, it uses

tests/test_quantum_info_utils.py

@@ -16,6 +16,7 @@
     hellinger_fidelity,
     pqc_integral,
     shannon_entropy,
+    total_variation_distance,
 )
 
 
@@ -163,7 +164,58 @@ def test_shannon_entropy(backend, base):
         backend.assert_allclose(result, 1.0)
 
 
-def test_hellinger(backend):
+@pytest.mark.parametrize("kind", [None, list])
+@pytest.mark.parametrize("validate", [False, True])
+def test_total_variation_distance(backend, validate, kind):
+    with pytest.raises(TypeError):
+        prob = np.random.rand(1, 2)
+        prob_q = np.random.rand(1, 5)
+        prob = backend.cast(prob, dtype=prob.dtype)
+        prob_q = backend.cast(prob_q, dtype=prob_q.dtype)
+        test = total_variation_distance(prob, prob_q, backend=backend)
+    with pytest.raises(TypeError):
+        prob = np.random.rand(1, 2)[0]
+        prob_q = np.array([])
+        prob = backend.cast(prob, dtype=prob.dtype)
+        prob_q = backend.cast(prob_q, dtype=prob_q.dtype)
+        test = total_variation_distance(prob, prob_q, backend=backend)
+    with pytest.raises(ValueError):
+        prob = np.array([-1, 2.0])
+        prob_q = np.random.rand(1, 5)[0]
+        prob = backend.cast(prob, dtype=prob.dtype)
+        prob_q = backend.cast(prob_q, dtype=prob_q.dtype)
+        test = total_variation_distance(prob, prob_q, validate=True, backend=backend)
+    with pytest.raises(ValueError):
+        prob = np.random.rand(1, 2)[0]
+        prob_q = np.array([1.0, 0.0])
+        prob = backend.cast(prob, dtype=prob.dtype)
+        prob_q = backend.cast(prob_q, dtype=prob_q.dtype)
+        test = total_variation_distance(prob, prob_q, validate=True, backend=backend)
+    with pytest.raises(ValueError):
+        prob = np.array([1.0, 0.0])
+        prob_q = np.random.rand(1, 2)[0]
+        prob = backend.cast(prob, dtype=prob.dtype)
+        prob_q = backend.cast(prob_q, dtype=prob_q.dtype)
+        test = total_variation_distance(prob, prob_q, validate=True, backend=backend)
+
+    prob_p = np.random.rand(10)
+    prob_q = np.random.rand(10)
+    prob_p /= np.sum(prob_p)
+    prob_q /= np.sum(prob_q)
+
+    target = 0.5 * np.sum(np.abs(prob_p - prob_q))
+
+    if kind is not None:
+        prob_p, prob_q = kind(prob_p), kind(prob_q)
+
+    distance = total_variation_distance(prob_p, prob_q, validate, backend=backend)
+
+    backend.assert_allclose(distance, target, atol=1e-5)
+
+
+@pytest.mark.parametrize("kind", [None, list])
+@pytest.mark.parametrize("validate", [False, True])
+def test_hellinger(backend, validate, kind):
     with pytest.raises(TypeError):
         prob = np.random.rand(1, 2)
         prob_q = np.random.rand(1, 5)
@@ -195,10 +247,23 @@ def test_hellinger(backend):
         prob_q = backend.cast(prob_q, dtype=prob_q.dtype)
         test = hellinger_distance(prob, prob_q, validate=True, backend=backend)
 
-    prob = [1.0, 0.0]
-    prob_q = [1.0, 0.0]
-    backend.assert_allclose(hellinger_distance(prob, prob_q, backend=backend), 0.0)
-    backend.assert_allclose(hellinger_fidelity(prob, prob_q, backend=backend), 1.0)
+    prob_p = np.random.rand(10)
+    prob_q = np.random.rand(10)
+    prob_p /= np.sum(prob_p)
+    prob_q /= np.sum(prob_q)
+
+    target = float(
+        backend.calculate_norm(np.sqrt(prob_p) - np.sqrt(prob_q)) / np.sqrt(2)
+    )
+
+    if kind is not None:
+        prob_p, prob_q = list(prob_p), list(prob_q)
+
+    distance = hellinger_distance(prob_p, prob_q, validate=validate, backend=backend)
+    fidelity = hellinger_fidelity(prob_p, prob_q, validate=validate, backend=backend)
+
+    assert distance == target
+    assert fidelity == (1 - target**2) ** 2
 
 
 def test_haar_integral_errors(backend):