Adds improved Bloch-Messiah that gives correct results when the symplectic matrix is degenerate

.github/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ### New features
 
+* New function to produce Bloch-Messiah decomposition of symplectic matrices.
+
 ### Breaking changes
 
 ### Improvements
@@ -25,7 +27,7 @@
 
 This release contains contributions from (in alphabetical order):
 
-Mikhail Andrenkov, Sebastián Duque, Jacon Hastrup, Antonín Hoskovec, Filippo Miatto
+Mikhail Andrenkov, Sebastián Duque, Jacob Hastrup, Antonín Hoskovec, Filippo Miatto
 
 ---
 

thewalrus/decompositions.py

@@ -28,6 +28,7 @@
 .. autosummary::
     williamson
     symplectic_eigenvals
+    blochmessiah
 
 Code details
 ------------
@@ -36,6 +37,7 @@
 
 from scipy.linalg import block_diag, sqrtm, schur
 from thewalrus.symplectic import sympmat
+from thewalrus.quantum.gaussian_checks import is_symplectic
 
 
 def williamson(V, rtol=1e-05, atol=1e-08):
@@ -112,3 +114,61 @@ def symplectic_eigenvals(cov):
     M = int(len(cov) / 2)
     D, _ = williamson(cov)
     return np.diag(D)[:M]
+
+
+def blochmessiah(S):
+    """Returns the Bloch-Messiah decomposition of a symplectic matrix S = uff @ dff @ vff
+       where uff and vff are orthogonal symplectic matrices and dff is a diagonal matrix
+       of the form diag(d1,d2,...,dn,d1^-1, d2^-1,...,dn^-1),
+
+    Args:
+        S (array[float]): 2N x 2N real symplectic matrix
+
+    Returns:
+        tupple(array[float],  : orthogonal symplectic matrix uff
+               array[float],  : diagional matrix dff
+               array[float])  : orthogonal symplectic matrix vff
+    """
+
+    N, _ = S.shape
+
+    if not is_symplectic(S):
+        raise ValueError("Input matrix is not symplectic.")
+
+    # Changing Basis
+    R = (1 / np.sqrt(2)) * np.block(
+        [[np.eye(N // 2), 1j * np.eye(N // 2)], [np.eye(N // 2), -1j * np.eye(N // 2)]]
+    )
+    Sc = R @ S @ np.conjugate(R).T
+    # Polar Decomposition
+    u1, d1, v1 = np.linalg.svd(Sc)
+    Sig = u1 @ np.diag(d1) @ np.conjugate(u1).T
+    Unitary = u1 @ v1
+    # Blocks of Unitary and Hermitian symplectics
+    alpha = Unitary[0 : N // 2, 0 : N // 2]
+    beta = Sig[0 : N // 2, N // 2 : N]
+    # Bloch-Messiah in this Basis
+    u2, d2, v2 = np.linalg.svd(beta)
+    sval = np.arcsinh(d2)
+    takagibeta = u2 @ sqrtm(np.conjugate(u2).T @ (v2.T))
+    uf = np.block([[takagibeta, 0 * takagibeta], [0 * takagibeta, np.conjugate(takagibeta)]])
+    vf = np.block(
+        [
+            [np.conjugate(takagibeta).T @ alpha, 0 * takagibeta],
+            [0 * takagibeta, np.conjugate(np.conjugate(takagibeta).T @ alpha)],
+        ]
+    )
+    df = np.block(
+        [
+            [np.diag(np.cosh(sval)), np.diag(np.sinh(sval))],
+            [np.diag(np.sinh(sval)), np.diag(np.cosh(sval))],
+        ]
+    )
+    # Rotating Back to Original Basis
+    uff = np.conjugate(R).T @ uf @ R
+    vff = np.conjugate(R).T @ vf @ R
+    dff = np.conjugate(R).T @ df @ R
+    dff = np.real_if_close(dff)
+    vff = np.real_if_close(vff)
+    uff = np.real_if_close(uff)
+    return uff, dff, vff

thewalrus/tests/test_decompositions.py

@@ -18,8 +18,10 @@
 from scipy.linalg import block_diag
 
 from thewalrus.random import random_interferometer as haar_measure
-from thewalrus.decompositions import williamson
+from thewalrus.random import random_symplectic
+from thewalrus.decompositions import williamson, blochmessiah
 from thewalrus.symplectic import sympmat as omega
+from thewalrus.quantum.gaussian_checks import is_symplectic
 
 
 class TestWilliamsonDecomposition:
@@ -147,3 +149,107 @@ def test_mixed_state(self, create_cov, hbar, tol):
         assert np.allclose(sorted(nbar[:n]), sorted(nbar_in), atol=tol, rtol=0)
         # check S is symplectic
         assert np.allclose(S @ O @ S.T, O, atol=tol, rtol=0)
+
+
+class TestBlochMessiahDecomposition:
+    """Tests for the Bloch Messiah decomposition"""
+
+    @pytest.fixture
+    def create_transform(self):
+        """create a symplectic transform for use in testing.
+
+        Args:
+            n (int): number of modes
+            passive (bool): whether transform should be passive or not
+
+        Returns:
+            array: symplectic matrix
+        """
+
+        def _create_transform(n, passive=True):
+            """wrapped function"""
+            O = omega(n)
+
+            # interferometer 1
+            U1 = haar_measure(n)
+            S1 = np.vstack([np.hstack([U1.real, -U1.imag]), np.hstack([U1.imag, U1.real])])
+
+            Sq = np.identity(2 * n)
+            if not passive:
+                # squeezing
+                r = np.log(0.2 * np.arange(n) + 2)
+                Sq = block_diag(np.diag(np.exp(-r)), np.diag(np.exp(r)))
+
+            # interferometer 2
+            U2 = haar_measure(n)
+            S2 = np.vstack([np.hstack([U2.real, -U2.imag]), np.hstack([U2.imag, U2.real])])
+
+            # final symplectic
+            S_final = S2 @ Sq @ S1
+
+            # check valid symplectic transform
+            assert np.allclose(S_final.T @ O @ S_final, O)
+            return S_final
+
+        return _create_transform
+
+    @pytest.mark.parametrize("N", range(50, 500, 50))
+    def test_blochmessiah_rand(self, N):
+        """Tests blochmessiah function for different matrix sizes."""
+        S = random_symplectic(N)
+        u, d, v = blochmessiah(S)
+        assert np.allclose(u @ d @ v, S)
+        assert np.allclose(u.T @ u, np.eye(len(u)))
+        assert np.allclose(v.T @ v, np.eye(len(v)))
+        assert is_symplectic(u)
+        assert is_symplectic(v)
+
+    @pytest.mark.parametrize("M", [np.random.rand(4, 5), np.random.rand(4, 4)])
+    def test_blochmessiah_error(self, M):
+        """Tests that non-symplectic matrices raise a ValueError in blochmessiah."""
+        with pytest.raises(ValueError, match="Input matrix is not symplectic."):
+            blochmessiah(M)
+
+    def test_identity(self, tol):
+        """Test identity"""
+        n = 2
+        S_in = np.identity(2 * n)
+        O1, S, O2 = blochmessiah(S_in)
+
+        assert np.allclose(O1 @ O2, np.identity(2 * n), atol=tol, rtol=0)
+        assert np.allclose(S, np.identity(2 * n), atol=tol, rtol=0)
+
+        # test orthogonality
+        assert np.allclose(O1.T, O1, atol=tol, rtol=0)
+        assert np.allclose(O2.T, O2, atol=tol, rtol=0)
+
+        # test symplectic
+        O = omega(n)
+        assert np.allclose(O1 @ O @ O1.T, O, atol=tol, rtol=0)
+        assert np.allclose(O2 @ O @ O2.T, O, atol=tol, rtol=0)
+
+    @pytest.mark.parametrize("passive", [True, False])
+    def test_transform(self, passive, create_transform, tol):
+        """Test decomposition agrees with transform. also checks that passive transform has no squeezing.
+        Note: this test also tests the case with degenerate symplectic values"""
+        n = 3
+        S_in = create_transform(3, passive=passive)
+        O1, S, O2 = blochmessiah(S_in)
+
+        # test decomposition
+        assert np.allclose(O1 @ S @ O2, S_in, atol=tol, rtol=0)
+
+        # test no squeezing
+        if passive:
+            assert np.allclose(O1 @ O2, S_in, atol=tol, rtol=0)
+            assert np.allclose(S, np.identity(2 * n), atol=tol, rtol=0)
+
+        # test orthogonality
+        assert np.allclose(O1.T @ O1, np.identity(2 * n), atol=tol, rtol=0)
+        assert np.allclose(O2.T @ O2, np.identity(2 * n), atol=tol, rtol=0)
+
+        # test symplectic
+        O = omega(n)
+        assert np.allclose(O1.T @ O @ O1, O, atol=tol, rtol=0)
+        assert np.allclose(O2.T @ O @ O2, O, atol=tol, rtol=0)
+        assert np.allclose(S @ O @ S.T, O, atol=tol, rtol=0)