Add total_variation_distance to quantum_info.utils

doc/source/api-reference/qibo.rst

@@ -2398,6 +2398,12 @@ Hellinger shot error
 .. autofunction:: qibo.quantum_info.hellinger_shot_error
 
 
+Total variation distance
+""""""""""""""""""""""""
+
+.. autofunction:: qibo.quantum_info.total_variation_distance
+
+
 Haar integral
 """""""""""""
 

src/qibo/quantum_info/utils.py

@@ -280,7 +280,7 @@ def hellinger_fidelity(prob_dist_p, prob_dist_q, validate: bool = False, backend
             :class:`qibo.backends.GlobalBackend`. Defaults to ``None``.
 
     Returns:
-        (float): Hellinger fidelity.
+        float: Hellinger fidelity.
 
     """
     backend = _check_backend(backend)
@@ -317,7 +317,7 @@ def hellinger_shot_error(
             :class:`qibo.backends.GlobalBackend`. Defaults to ``None``.
 
     Returns:
-        (float): Hellinger fidelity error.
+        float: Hellinger fidelity error.
 
     """
     backend = _check_backend(backend)
@@ -339,6 +339,58 @@ def hellinger_shot_error(
     return hellinger_error
 
 
+def total_variation_distance(
+    prob_dist_p, prob_dist_q, validate: bool = False, backend=None
+):
+    """Calculates the total variation distance between two discrete probability distributions.
+
+    For probabilities :math:`p` and :math:`q`, the total variation distance is defined as
+
+    .. math::
+        \\operatorname{TVD}(p, q) = \frac{1}{2} \\, \\norm{p - q}_{1}
+            = \\frac{1}{2} \\, \\sum_{x} \\, \\abs{p(x) - q(x)} \\, ,
+
+    where :math:`\\norm{\\cdot}_{1}` detones the :math:`\\ell_{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 measure.
+    """
+    backend = _check_backend(backend)
+
+    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 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 backend.np.abs(backend.np.sum(prob_dist_p) - 1.0) > PRECISION_TOL:
+            raise_error(ValueError, "First probability array must sum to 1.")
+
+        if backend.np.abs(backend.np.sum(prob_dist_q) - 1.0) > PRECISION_TOL:
+            raise_error(ValueError, "Second probability array must sum to 1.")
+
+    tvd = backend.calculate_norm(prob_dist_p - prob_dist_q, order=1)
+
+    return tvd / 2
+
+
 def haar_integral(
     nqubits: int,
     power_t: int,

tests/test_quantum_info_utils.py

@@ -17,6 +17,7 @@
     hellinger_fidelity,
     hellinger_shot_error,
     pqc_integral,
+    total_variation_distance,
 )
 
 
@@ -212,6 +213,52 @@ def test_hellinger_shot_error(backend, validate, kind):
     assert 2 * hellinger_error < hellinger_fid
 
 
+@pytest.mark.parametrize("kind", [None, list])
+@pytest.mark.parametrize("validate", [False, True])
+def test_total_variation_distance(backend, validate, kind):
+    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)
+    prob_p = backend.cast(prob_p, dtype=prob_p.dtype)
+    prob_q = backend.cast(prob_q, dtype=prob_q.dtype)
+
+    target = float(backend.calculate_norm(prob_p - prob_q, order=1) / 2)
+
+    prob_p = (
+        kind(prob_p) if kind is not None else backend.cast(prob_p, dtype=prob_p.dtype)
+    )
+    prob_q = (
+        kind(prob_q) if kind is not None else backend.cast(prob_q, dtype=prob_q.dtype)
+    )
+
+    tvd = total_variation_distance(prob_p, prob_q, validate, backend)
+    distance = hellinger_distance(prob_p, prob_q, validate, backend)
+
+    assert tvd == target
+    assert tvd <= np.sqrt(2) * distance
+    assert tvd >= distance**2
+
+
 def test_haar_integral_errors(backend):
     with pytest.raises(TypeError):
         nqubits, power_t, samples = 0.5, 2, 10