From dac736e9589e5f35529a9ac1c72db82536eaf452 Mon Sep 17 00:00:00 2001
From: Renato Mello <renato.msf@gmail.com>
Date: Wed, 31 Jan 2024 10:20:23 +0400
Subject: [PATCH] tests

---
 tests/test_quantum_info_entropies.py | 98 ++++++++++++++++++++++++++++
 1 file changed, 98 insertions(+)

diff --git a/tests/test_quantum_info_entropies.py b/tests/test_quantum_info_entropies.py
index 702ce8db2e..c12532424c 100644
--- a/tests/test_quantum_info_entropies.py
+++ b/tests/test_quantum_info_entropies.py
@@ -5,6 +5,7 @@
 from qibo.quantum_info.entropies import (
     classical_relative_entropy,
     classical_renyi_entropy,
+    classical_renyi_relative_entropy,
     entanglement_entropy,
     entropy,
     relative_entropy,
@@ -115,6 +116,103 @@ def test_classical_relative_entropy(backend, base, kind):
     backend.assert_allclose(divergence, target, atol=1e-5)
 
 
+@pytest.mark.parametrize("kind", [None, list])
+@pytest.mark.parametrize("base", [2, 10, np.e, 5])
+@pytest.mark.parametrize("alpha", [0, 1 / 2, 1, 2, 3, np.inf])
+def test_classical_renyi_relative_entropy(backend, alpha, base, 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 = classical_renyi_relative_entropy(
+            prob, prob_q, alpha, base, 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 = classical_renyi_relative_entropy(
+            prob, prob_q, alpha, base, 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 = classical_renyi_relative_entropy(
+            prob, prob_q, alpha, base, 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 = classical_renyi_relative_entropy(
+            prob, prob_q, alpha, base, 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 = classical_renyi_relative_entropy(
+            prob, prob_q, alpha, base, backend=backend
+        )
+    with pytest.raises(ValueError):
+        prob = np.array([1.0, 0.0])
+        prob_q = np.array([0.0, 1.0])
+        prob = backend.cast(prob, dtype=prob.dtype)
+        prob_q = backend.cast(prob_q, dtype=prob_q.dtype)
+        test = classical_renyi_relative_entropy(
+            prob, prob_q, alpha, base=-2, backend=backend
+        )
+    with pytest.raises(TypeError):
+        prob = np.array([1.0, 0.0])
+        prob_q = np.array([0.0, 1.0])
+        prob = backend.cast(prob, dtype=prob.dtype)
+        prob_q = backend.cast(prob_q, dtype=prob_q.dtype)
+        test = classical_renyi_relative_entropy(
+            prob, prob_q, alpha="1", base=base, backend=backend
+        )
+    with pytest.raises(ValueError):
+        prob = np.array([1.0, 0.0])
+        prob_q = np.array([0.0, 1.0])
+        prob = backend.cast(prob, dtype=prob.dtype)
+        prob_q = backend.cast(prob_q, dtype=prob_q.dtype)
+        test = classical_renyi_relative_entropy(
+            prob, prob_q, alpha=-2, base=base, 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)
+
+    if alpha == 0.5:
+        target = -2 * np.log2(np.sum(np.sqrt(prob_p * prob_q))) / np.log2(base)
+    elif alpha == 1.0:
+        target = classical_relative_entropy(prob_p, prob_q, base=base, backend=backend)
+    elif alpha == np.inf:
+        target = np.log2(max(prob_p / prob_q)) / np.log2(base)
+    else:
+        target = (
+            (1 / (alpha - 1))
+            * np.log2(np.sum(prob_p**alpha * prob_q ** (1 - alpha)))
+            / np.log2(base)
+        )
+
+    if kind is not None:
+        prob_p, prob_q = kind(prob_p), kind(prob_q)
+
+    divergence = classical_renyi_relative_entropy(
+        prob_p, prob_q, alpha=alpha, base=base, backend=backend
+    )
+
+    backend.assert_allclose(divergence, target, atol=1e-5)
+
+
 @pytest.mark.parametrize("kind", [None, list])
 @pytest.mark.parametrize("base", [2, 10, np.e, 5])
 @pytest.mark.parametrize("alpha", [0, 1, 2, 3, np.inf])