Merge pull request #1454 from qiboteam/matrix_power

Add `matrix_power` to `quantum_info.linalg_operations`

Author: renatomello

doc/source/api-reference/qibo.rst

@@ -1964,6 +1964,12 @@ Matrix exponentiation
 .. autofunction:: qibo.quantum_info.matrix_exponentiation
 
 
+Matrix power
+""""""""""""
+
+.. autofunction:: qibo.quantum_info.matrix_power
+
+
 Quantum Networks
 ^^^^^^^^^^^^^^^^
 

poetry.lock

@@ -1,4 +1,5 @@
 import abc
+from typing import Union
 
 from qibo.config import raise_error
 
@@ -355,6 +356,21 @@ def calculate_matrix_exp(
         """
         raise_error(NotImplementedError)
 
+    @abc.abstractmethod
+    def calculate_matrix_power(
+        self, matrix, power: Union[float, int]
+    ):  # pragma: no cover
+        """Calculate the (fractional) ``power`` :math:`\\alpha` of ``matrix`` :math:`A`,
+        i.e. :math:`A^{\\alpha}`.
+
+        .. note::
+            For the ``pytorch`` backend, this method relies on a copy of the original tensor.
+            This may break the gradient flow. For the GPU backends (i.e. ``cupy`` and
+            ``cuquantum``), this method falls back to CPU whenever ``power`` is not
+            an integer.
+        """
+        raise_error(NotImplementedError)
+
     @abc.abstractmethod
     def calculate_hamiltonian_matrix_product(
         self, matrix1, matrix2

src/qibo/backends/abstract.py

@@ -1,9 +1,10 @@
 import collections
 import math
+from typing import Union
 
 import numpy as np
 from scipy import sparse
-from scipy.linalg import block_diag
+from scipy.linalg import block_diag, fractional_matrix_power
 
 from qibo import __version__
 from qibo.backends import einsum_utils
@@ -768,6 +769,14 @@ def calculate_matrix_exp(self, a, matrix, eigenvectors=None, eigenvalues=None):
         ud = self.np.transpose(np.conj(eigenvectors))
         return self.np.matmul(eigenvectors, self.np.matmul(expd, ud))
 
+    def calculate_matrix_power(self, matrix, power: Union[float, int]):
+        if not isinstance(power, (float, int)):
+            raise_error(
+                TypeError,
+                f"``power`` must be either float or int, but it is type {type(power)}.",
+            )
+        return fractional_matrix_power(matrix, power)
+
     # TODO: remove this method
     def calculate_hamiltonian_matrix_product(self, matrix1, matrix2):
         return matrix1 @ matrix2

src/qibo/backends/numpy.py

@@ -199,6 +199,11 @@ def calculate_matrix_exp(self, a, matrix, eigenvectors=None, eigenvalues=None):
         ud = self.np.conj(eigenvectors).T
         return self.np.matmul(eigenvectors, self.np.matmul(expd, ud))
 
+    def calculate_matrix_power(self, matrix, power):
+        copied = self.to_numpy(self.np.copy(matrix))
+        copied = super().calculate_matrix_power(copied, power)
+        return self.cast(copied, dtype=copied.dtype)
+
     def _test_regressions(self, name):
         if name == "test_measurementresult_apply_bitflips":
             return [

src/qibo/backends/pytorch.py

@@ -3,12 +3,10 @@
 from typing import Union
 
 import numpy as np
-from scipy.linalg import fractional_matrix_power
 
 from qibo.backends import _check_backend
-from qibo.backends.pytorch import PyTorchBackend
 from qibo.config import PRECISION_TOL, raise_error
-from qibo.quantum_info.linalg_operations import partial_trace
+from qibo.quantum_info.linalg_operations import matrix_power, partial_trace
 from qibo.quantum_info.metrics import _check_hermitian, purity
 
 
@@ -724,7 +722,7 @@ def renyi_entropy(state, alpha: Union[float, int], base: float = 2, backend=None
             / np.log2(base)
         )
 
-    log = backend.np.log2(backend.np.trace(_matrix_power(state, alpha, backend)))
+    log = backend.np.log2(backend.np.trace(matrix_power(state, alpha, backend)))
 
     return (1 / (1 - alpha)) * log / np.log2(base)
 
@@ -823,17 +821,17 @@ def relative_renyi_entropy(
         return relative_von_neumann_entropy(state, target, base, backend=backend)
 
     if alpha == np.inf:
-        new_state = _matrix_power(state, 0.5, backend)
-        new_target = _matrix_power(target, 0.5, backend)
+        new_state = matrix_power(state, 0.5, backend)
+        new_target = matrix_power(target, 0.5, backend)
 
         log = backend.np.log2(
             backend.calculate_norm_density_matrix(new_state @ new_target, order=1)
         )
 
         return -2 * log / np.log2(base)
 
-    log = _matrix_power(state, alpha, backend)
-    log = log @ _matrix_power(target, 1 - alpha, backend)
+    log = matrix_power(state, alpha, backend)
+    log = log @ matrix_power(target, 1 - alpha, backend)
     log = backend.np.log2(backend.np.trace(log))
 
     return (1 / (alpha - 1)) * log / np.log2(base)
@@ -891,7 +889,7 @@ def tsallis_entropy(state, alpha: float, base: float = 2, backend=None):
         return von_neumann_entropy(state, base=base, backend=backend)
 
     return (1 / (1 - alpha)) * (
-        backend.np.trace(_matrix_power(state, alpha, backend)) - 1
+        backend.np.trace(matrix_power(state, alpha, backend)) - 1
     )
 
 
@@ -953,19 +951,3 @@ def entanglement_entropy(
     )
 
     return entropy_entanglement
-
-
-def _matrix_power(matrix, alpha, backend):
-    """Calculates ``matrix ** alpha`` according to backend."""
-    if backend.__class__.__name__ in [
-        "CupyBackend",
-        "CuQuantumBackend",
-    ]:  # pragma: no cover
-        new_matrix = backend.to_numpy(matrix)
-    else:
-        new_matrix = backend.np.copy(matrix)
-
-    if len(new_matrix.shape) == 1:
-        new_matrix = backend.np.outer(new_matrix, backend.np.conj(new_matrix))
-
-    return backend.cast(fractional_matrix_power(backend.to_numpy(new_matrix), alpha))

src/qibo/quantum_info/entropies.py

@@ -196,3 +196,21 @@ def matrix_exponentiation(
     backend = _check_backend(backend)
 
     return backend.calculate_matrix_exp(phase, matrix, eigenvectors, eigenvalues)
+
+
+def matrix_power(matrix, power: Union[float, int], backend=None):
+    """Given a ``matrix`` :math:`A` and power :math:`\\alpha`, calculate :math:`A^{\\alpha}`.
+
+    Args:
+        matrix (ndarray): matrix whose power to calculate.
+        power (float or int): power to raise ``matrix`` to.
+        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:
+        ndarray: matrix power :math:`A^{\\alpha}`.
+    """
+    backend = _check_backend(backend)
+
+    return backend.calculate_matrix_power(matrix, power)

src/qibo/quantum_info/linalg_operations.py

@@ -1,10 +1,8 @@
 import numpy as np
 import pytest
-from scipy.linalg import sqrtm
 
 from qibo.config import PRECISION_TOL
 from qibo.quantum_info.entropies import (
-    _matrix_power,
     classical_mutual_information,
     classical_relative_entropy,
     classical_relative_renyi_entropy,
@@ -19,6 +17,7 @@
     tsallis_entropy,
     von_neumann_entropy,
 )
+from qibo.quantum_info.linalg_operations import matrix_power
 from qibo.quantum_info.random_ensembles import (
     random_density_matrix,
     random_statevector,
@@ -646,8 +645,18 @@ def test_relative_renyi_entropy(backend, alpha, base, state_flag, target_flag):
             if alpha == 1.0:
                 log = relative_von_neumann_entropy(state, target, base, backend=backend)
             elif alpha == np.inf:
-                new_state = _matrix_power(state, 0.5, backend)
-                new_target = _matrix_power(target, 0.5, backend)
+                state_outer = (
+                    backend.np.outer(state, backend.np.conj(state.T))
+                    if state_flag
+                    else state
+                )
+                target_outer = (
+                    backend.np.outer(target, backend.np.conj(target.T))
+                    if target_flag
+                    else target
+                )
+                new_state = matrix_power(state_outer, 0.5, backend)
+                new_target = matrix_power(target_outer, 0.5, backend)
 
                 log = backend.np.log2(
                     backend.calculate_norm_density_matrix(
@@ -663,8 +672,8 @@ def test_relative_renyi_entropy(backend, alpha, base, state_flag, target_flag):
                 if len(target.shape) == 1:
                     target = backend.np.outer(target, backend.np.conj(target))
 
-                log = _matrix_power(state, alpha, backend)
-                log = log @ _matrix_power(target, 1 - alpha, backend)
+                log = matrix_power(state, alpha, backend)
+                log = log @ matrix_power(target, 1 - alpha, backend)
                 log = backend.np.log2(backend.np.trace(log))
 
                 log = (1 / (alpha - 1)) * log / np.log2(base)
@@ -710,7 +719,7 @@ def test_tsallis_entropy(backend, alpha, base):
         target = von_neumann_entropy(state, base=base, backend=backend)
     else:
         target = (1 / (1 - alpha)) * (
-            backend.np.trace(_matrix_power(state, alpha, backend)) - 1
+            backend.np.trace(matrix_power(state, alpha, backend)) - 1
         )
 
     backend.assert_allclose(

tests/test_quantum_info_entropies.py

@@ -5,8 +5,10 @@
 from qibo.quantum_info.linalg_operations import (
     anticommutator,
     commutator,
+    matrix_power,
     partial_trace,
 )
+from qibo.quantum_info.metrics import purity
 from qibo.quantum_info.random_ensembles import random_density_matrix, random_statevector
 
 
@@ -109,3 +111,22 @@ def test_partial_trace(backend, density_matrix):
     Id = backend.identity_density_matrix(1, normalize=True)
 
     backend.assert_allclose(traced, Id)
+
+
+@pytest.mark.parametrize("power", [2, 2.0, "2"])
+def test_matrix_power(backend, power):
+    nqubits = 2
+    dims = 2**nqubits
+
+    state = random_density_matrix(dims, backend=backend)
+
+    if isinstance(power, str):
+        with pytest.raises(TypeError):
+            test = matrix_power(state, power, backend)
+    else:
+        power = matrix_power(state, power, backend)
+
+        backend.assert_allclose(
+            float(backend.np.real(backend.np.trace(power))),
+            purity(state, backend=backend),
+        )