refact(_math): xpxp-xxpp basis changing rewritten

Sometimes we need to switch between the xxpp and xpxp basis. In Piquasso, this
is implemented by a basis changing matrix. However, as this matrix only
reorders rows/columns, it would be much nicer if the vectors/matrices would
just be reindexed. Therefore, this approach is deleted in favour of basis
changing indices.

.codecov.yml

@@ -6,5 +6,6 @@ ignore:
   - piquasso/_math/combinatorics.py
   - piquasso/_math/fock.py
   - piquasso/_math/indices.py
+  - piquasso/_math/transformations.py
   - piquasso/_simulators/fock/calculations.py
   - piquasso/_simulators/connectors/numpy_/interferometer.py

piquasso/_math/decompositions.py

@@ -18,7 +18,7 @@
 from scipy.optimize import root_scalar
 
 from piquasso._math.symplectic import xp_symplectic_form
-from piquasso._math.transformations import from_xxpp_to_xpxp_transformation_matrix
+from piquasso._math.transformations import xpxp_to_xxpp_indices
 
 from piquasso.api.exceptions import InvalidParameter
 from piquasso.api.connector import BaseConnector
@@ -176,9 +176,10 @@ def williamson(matrix: np.ndarray, connector: BaseConnector) -> tuple:
         output="real",
     )
 
+    indices = xpxp_to_xxpp_indices(d)
     basis_change = _rotation_to_positive_above_diagonals(
         block_diagonal_part, connector
-    ) @ from_xxpp_to_xpxp_transformation_matrix(d)
+    )[:, indices]
     ordered_block_diagonal = basis_change.T @ block_diagonal_part @ basis_change
 
     inverse_diagonal_matrix = connector.block_diag(

piquasso/_math/transformations.py

@@ -13,40 +13,47 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import functools
 import numpy as np
+import numba as nb
 
 
-@functools.lru_cache()
-def from_xxpp_to_xpxp_transformation_matrix(d: int) -> np.ndarray:
+@nb.njit(cache=True)
+def xxpp_to_xpxp_indices(d: int) -> np.ndarray:
     r"""
-    Basis changing with the basis change operator.
+    Indices for basis changing from the xxpp to the xpxp basis.
 
-    This transformation will change the basis from xxpp-basis to xpxp-basis.
+    Args:
+        d (int): The number of modes.
 
-    .. math::
+    Returns:
+        numpy.ndarray: The basis changing indices.
+    """
 
-        T_{ij} = \delta_{j, 2i-1} + \delta_{j + 2d, 2i}
+    indices = np.empty(2 * d, dtype=nb.int32)
 
-    Intuitively, it changes the basis as
+    for i in range(d):
+        indices[2 * i] = i
+        indices[2 * i + 1] = d + i
 
-    .. math::
+    return indices
 
-        T Y = T (x_1, \dots, x_d, p_1, \dots, p_d)^T
-            = (x_1, p_1, \dots, x_d, p_d)^T,
 
-    which is very helpful in :mod:`piquasso._simulators.gaussian.state`.
+@nb.njit(cache=True)
+def xpxp_to_xxpp_indices(d: int) -> np.ndarray:
+    r"""
+    Indices for basis changing from the xpxp to the xxpp basis.
 
     Args:
         d (int): The number of modes.
 
     Returns:
-        numpy.ndarray: The basis changing matrix.
+        numpy.ndarray: The basis changing indices.
     """
 
-    T = np.zeros((2 * d, 2 * d), dtype=int)
+    indices = np.empty(2 * d, dtype=nb.int32)
+
     for i in range(d):
-        T[2 * i, i] = 1
-        T[2 * i + 1, i + d] = 1
+        indices[i] = 2 * i
+        indices[d + i] = 2 * i + 1
 
-    return T
+    return indices

piquasso/_simulators/gaussian/state.py

@@ -31,7 +31,7 @@
 )
 from piquasso._math.symplectic import symplectic_form
 from piquasso._math.fock import get_fock_space_basis
-from piquasso._math.transformations import from_xxpp_to_xpxp_transformation_matrix
+from piquasso._math.transformations import xxpp_to_xpxp_indices, xpxp_to_xxpp_indices
 
 from piquasso._math.decompositions import williamson
 
@@ -178,8 +178,7 @@ def xxpp_mean_vector(self) -> np.ndarray:
 
     @xxpp_mean_vector.setter
     def xxpp_mean_vector(self, value):
-        T = from_xxpp_to_xpxp_transformation_matrix(self.d)
-        self.xpxp_mean_vector = T @ value
+        self.xpxp_mean_vector = value[xxpp_to_xpxp_indices(self.d)]
 
     @property
     def xxpp_covariance_matrix(self) -> np.ndarray:
@@ -218,8 +217,8 @@ def xxpp_covariance_matrix(self) -> np.ndarray:
 
     @xxpp_covariance_matrix.setter
     def xxpp_covariance_matrix(self, value):
-        T = from_xxpp_to_xpxp_transformation_matrix(self.d)
-        self.xpxp_covariance_matrix = T @ value @ T.transpose()
+        indices = xxpp_to_xpxp_indices(self.d)
+        self.xpxp_covariance_matrix = value[np.ix_(indices, indices)]
 
     @property
     def xxpp_correlation_matrix(self) -> np.ndarray:
@@ -271,8 +270,7 @@ def xpxp_mean_vector(self) -> np.ndarray:
             operators in `xxpp`-ordering, i.e. :math:`\operatorname{Tr} \rho R`, where
             :math:`R = (x_1, p_1, \dots, x_d, p_d)^T`.
         """
-        T = from_xxpp_to_xpxp_transformation_matrix(self.d)
-        return T @ self.xxpp_mean_vector
+        return self.xxpp_mean_vector[xxpp_to_xpxp_indices(self.d)]
 
     @xpxp_mean_vector.setter
     def xpxp_mean_vector(self, value: np.ndarray) -> None:
@@ -307,8 +305,8 @@ def xpxp_covariance_matrix(self) -> np.ndarray:
                 xpxp-ordered basis.
         """
 
-        T = from_xxpp_to_xpxp_transformation_matrix(self.d)
-        return T @ self.xxpp_covariance_matrix @ T.transpose()
+        indices = xxpp_to_xpxp_indices(self.d)
+        return self.xxpp_covariance_matrix[np.ix_(indices, indices)]
 
     @xpxp_covariance_matrix.setter
     def xpxp_covariance_matrix(self, new_cov: np.ndarray) -> None:
@@ -318,10 +316,10 @@ def xpxp_covariance_matrix(self, new_cov: np.ndarray) -> None:
 
         self._validate_cov(new_cov, d)
 
-        T = from_xxpp_to_xpxp_transformation_matrix(d)
+        indices = xpxp_to_xxpp_indices(self.d)
 
         dimensionless_cov = new_cov / self._config.hbar
-        dimensionless_xp_cov = T.transpose() @ dimensionless_cov @ T
+        dimensionless_xp_cov = dimensionless_cov[np.ix_(indices, indices)]
 
         blocks = (dimensionless_xp_cov - np.identity(2 * d)) / 4
 
@@ -359,8 +357,8 @@ def xpxp_correlation_matrix(self) -> np.ndarray:
             numpy.ndarray: The :math:`2d \times 2d` `xpxp`-ordered correlation matrix.
         """
 
-        T = from_xxpp_to_xpxp_transformation_matrix(self.d)
-        return T @ self.xxpp_correlation_matrix @ T.transpose()
+        indices = xxpp_to_xpxp_indices(self.d)
+        return self.xxpp_correlation_matrix[np.ix_(indices, indices)]
 
     @property
     def xpxp_representation(self) -> Tuple[np.ndarray, np.ndarray]:
@@ -983,22 +981,25 @@ def purify(self) -> "GaussianState":
             GaussianState: The purified Gaussian state.
         """  # noqa: E501
         np = self._connector.np
+        fallback_np = self._connector.fallback_np
 
         hbar = self._config.hbar
         cov = self.xpxp_covariance_matrix / hbar
 
         d = self.d
 
-        T = from_xxpp_to_xpxp_transformation_matrix(d)
+        indices = xpxp_to_xxpp_indices(d)
 
-        S, D = williamson(T.T @ cov @ T, self._connector)
+        S, D = williamson(cov[fallback_np.ix_(indices, indices)], self._connector)
 
         beta_diagonals = np.diag(np.sqrt(np.abs(np.diag(D)[:d] ** 2 - 1)))
         zeros = np.zeros_like(beta_diagonals)
 
         beta_skeleton = np.block([[zeros, beta_diagonals], [beta_diagonals, zeros]])
 
-        TS = T @ S
+        indices = xxpp_to_xpxp_indices(d)
+
+        TS = S[indices, :]
 
         beta = TS @ beta_skeleton @ TS.T
 

piquasso/fermionic/gaussian/calculations.py

@@ -22,7 +22,7 @@
 
 from piquasso._math.validations import all_zero_or_one
 from piquasso._math.transformations import (
-    from_xxpp_to_xpxp_transformation_matrix,
+    xxpp_to_xpxp_indices,
 )
 
 from ._misc import validate_fermionic_gaussian_hamiltonian
@@ -80,6 +80,7 @@ def gaussian_hamiltonian(
     validate_fermionic_gaussian_hamiltonian(H)
 
     np = state._connector.np
+    fallback_np = state._connector.fallback_np
 
     d = state.d
 
@@ -88,18 +89,14 @@ def gaussian_hamiltonian(
 
     A_plus_B = A + B
     A_minus_B = A - B
-    T = from_xxpp_to_xpxp_transformation_matrix(d)
-
-    h = (
-        T
-        @ np.block(
-            [
-                [A_plus_B.imag, A_plus_B.real],
-                [-A_minus_B.real, A_minus_B.imag],
-            ]
-        )
-        @ T.T
-    )
+    indices = xxpp_to_xpxp_indices(d)
+
+    h = np.block(
+        [
+            [A_plus_B.imag, A_plus_B.real],
+            [-A_minus_B.real, A_minus_B.imag],
+        ]
+    )[fallback_np.ix_(indices, indices)]
 
     SO = state._connector.expm(-2 * h)
 

piquasso/fermionic/gaussian/state.py

@@ -17,9 +17,7 @@
 
 from piquasso._math.validations import all_in_interval
 from piquasso._math.linalg import is_selfadjoint, is_skew_symmetric
-from piquasso._math.transformations import (
-    from_xxpp_to_xpxp_transformation_matrix,
-)
+from piquasso._math.transformations import xpxp_to_xxpp_indices, xxpp_to_xpxp_indices
 from piquasso._math.combinatorics import sort_and_get_parity
 
 from piquasso.api.exceptions import InvalidState, PiquassoException
@@ -85,26 +83,24 @@ def covariance_matrix(self):
         two_D_plus_E = 2 * (D + E)
         two_D_minus_E = 2 * (D - E)
 
-        T = from_xxpp_to_xpxp_transformation_matrix(d)
+        indices = xxpp_to_xpxp_indices(d)
 
-        return (
-            T
-            @ np.block(
-                [
-                    [two_D_plus_E.imag, -two_D_minus_E.real + ident],
-                    [two_D_plus_E.real - ident, two_D_minus_E.imag],
-                ]
-            )
-            @ T.T
-        )
+        return np.block(
+            [
+                [two_D_plus_E.imag, -two_D_minus_E.real + ident],
+                [two_D_plus_E.real - ident, two_D_minus_E.imag],
+            ]
+        )[np.ix_(indices, indices)]
 
     @covariance_matrix.setter
     def covariance_matrix(self, input_covariance_matrix_xpxp):
         d = self.d
 
-        T = from_xxpp_to_xpxp_transformation_matrix(d)
+        indices = xpxp_to_xxpp_indices(d)
 
-        input_covariance_matrix_xxpp = T.T @ input_covariance_matrix_xpxp @ T
+        input_covariance_matrix_xxpp = input_covariance_matrix_xpxp[
+            self._connector.fallback_np.ix_(indices, indices)
+        ]
 
         np = self._connector.np
 
@@ -380,9 +376,9 @@ def density_matrix(self) -> "np.ndarray":
         # 2. Unitary which diagonalizes the original density matrix
         omega = get_omega(self.d, connector)
 
-        T = from_xxpp_to_xpxp_transformation_matrix(d)
+        indices = xxpp_to_xpxp_indices(d)
 
-        basis_transformation = T @ omega
+        basis_transformation = omega[indices]
 
         i2H = (
             basis_transformation.conj().T
@@ -398,6 +394,7 @@ def density_matrix(self) -> "np.ndarray":
 
     def _set_occupation_numbers(self, occupation_numbers):
         dtype = self._config.dtype
+        fallback_np = self._connector.fallback_np
         np = self._connector.np
 
         plus_minus = np.diag(1 - 2 * np.array(occupation_numbers))
@@ -406,10 +403,12 @@ def _set_occupation_numbers(self, occupation_numbers):
 
         d = len(occupation_numbers)
 
-        T = from_xxpp_to_xpxp_transformation_matrix(d)
+        indices = xxpp_to_xpxp_indices(d)
 
         self.covariance_matrix = (
-            T @ np.block([[zeros, plus_minus], [-plus_minus, zeros]]) @ T.T
+            np.block([[zeros, plus_minus], [-plus_minus, zeros]])[
+                fallback_np.ix_(indices, indices)
+            ]
         ).astype(dtype)
 
     @classmethod

tests/_math/test_torontonian.py

@@ -23,7 +23,7 @@
 
 
 from piquasso._math.torontonian import torontonian, loop_torontonian
-from piquasso._math.transformations import from_xxpp_to_xpxp_transformation_matrix
+from piquasso._math.transformations import xpxp_to_xxpp_indices
 
 
 def powerset(iterable):
@@ -38,9 +38,9 @@ def torontonian_naive(A: np.ndarray) -> complex:
     """
     d = A.shape[0] // 2
 
-    T = from_xxpp_to_xpxp_transformation_matrix(d)
+    indices = xpxp_to_xxpp_indices(d)
 
-    A = T.T @ A @ T
+    A = A[np.ix_(indices, indices)]
 
     if d == 0:
         return 1.0 + 0j
@@ -78,10 +78,10 @@ def loop_torontonian_naive(A, gamma):
 
     d = A.shape[0] // 2
 
-    T = from_xxpp_to_xpxp_transformation_matrix(d)
+    indices = xpxp_to_xxpp_indices(d)
 
-    A = T.T @ A @ T
-    gamma = T.T @ gamma
+    A = A[np.ix_(indices, indices)]
+    gamma = gamma[indices]
 
     if d == 0:
         return 1.0

tests/fermionic/gaussian/test_state.py

@@ -23,9 +23,7 @@
 
 from scipy.linalg import block_diag
 
-from piquasso._math.transformations import (
-    from_xxpp_to_xpxp_transformation_matrix,
-)
+from piquasso._math.transformations import xxpp_to_xpxp_indices, xpxp_to_xxpp_indices
 
 from piquasso._math.linalg import is_selfadjoint, is_skew_symmetric
 from piquasso._math.validations import all_in_interval
@@ -163,10 +161,10 @@ def test_GaussianHamiltonian_covariance_and_correlation_matrix_equivalence(
 
     omega = get_omega(d, connector)
 
-    T = from_xxpp_to_xpxp_transformation_matrix(d)
+    indices = xpxp_to_xxpp_indices(d)
 
     assert np.allclose(
-        T.T @ state.covariance_matrix @ T,
+        state.covariance_matrix[np.ix_(indices, indices)],
         -1j
         * omega
         @ (2 * state.correlation_matrix - np.identity(2 * d))
@@ -1097,9 +1095,11 @@ def test_covariance_matrix_GaussianHamiltonian_equivalence_from_Vacuum(
 
     gate_hamiltonian_majorana = -1j * omega @ gate_hamiltonian @ omega.conj().T
 
-    T = from_xxpp_to_xpxp_transformation_matrix(d)
+    indices = xxpp_to_xpxp_indices(d)
 
-    covariance_unitary = T @ connector.expm(-2 * gate_hamiltonian_majorana) @ T.T
+    covariance_unitary = connector.expm(-2 * gate_hamiltonian_majorana)[
+        np.ix_(indices, indices)
+    ]
 
     assert np.allclose(
         final_state.covariance_matrix,