Implements the pre-Iwasawa and Iwasawa decompositions. (#382)

* First working version * First working version * First working version * better tests * better tests * Finally correct * Finally correct * putting the traingles in the right places * Removes unnecesary T * More tests * More tests and spell correction * More tests and spell correction * Removes unnecesary complex number usage * Apply suggestions from code review Co-authored-by: Sebastián Duque Mesa <[email protected]> * Apply suggestions from code review Co-authored-by: Sebastián Duque Mesa <[email protected]> --------- Co-authored-by: Nicolas Quesada <[email protected]> Co-authored-by: Sebastián Duque Mesa <[email protected]>
XanaduAI · Feb 29, 2024 · 25b5f08 · 25b5f08
1 parent 8759363
commit 25b5f08
Show file tree

Hide file tree

Showing 3 changed files with 209 additions and 3 deletions.
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 ### New features
 
+* Implements the pre-Iwasawa and Iwasawa decompositions for symplectic matrices [(#382)](https://github.com/XanaduAI/thewalrus/pull/382).
+
 ### Breaking changes
 
 ### Improvements
@@ -19,7 +21,7 @@
 
 This release contains contributions from (in alphabetical order):
 
-Nicolas Quesada
+Will McCutcheon, Nicolas Quesada
 
 ---
 

diff --git a/thewalrus/decompositions.py b/thewalrus/decompositions.py
@@ -30,13 +30,15 @@
     symplectic_eigenvals
     blochmessiah
     takagi
+    pre_iwasawa
+    iwasawa
 
 Code details
 ------------
 """
 import numpy as np
 
-from scipy.linalg import sqrtm, schur, polar
+from scipy.linalg import sqrtm, schur, polar, qr
 from thewalrus.symplectic import sympmat
 from thewalrus.quantum.gaussian_checks import is_symplectic
 
@@ -198,3 +200,77 @@ def takagi(A, svd_order=True):
     if svd_order is False:
         return d[::-1], U[:, ::-1]
     return d, U
+
+
+def pre_iwasawa(S):
+    """Pre-Iwasawa decomposition of a symplectic matrix.
+    See `Arvind et al. The Real Symplectic Groups in Quantum Mechanics and Optics <https://arxiv.org/pdf/quant-ph/9509002.pdf>`_
+
+
+    Args:
+        S (array): the symplectic matrix
+
+    Returns:
+        tuple[array, array, array]: (E,D,F) symplectic matrices such that E @ D @ F = S and,
+        E = np.block([[np.eye(N), np.zeros(N,N)],[X, np.eye(N)]]) with X == X.T,
+        D is block diagonal with the top left block being the inverse of the bottom right block,
+        F is symplectic orthogonal.
+    """
+
+    if not is_symplectic(S):
+        raise ValueError("Input matrix is not symplectic.")
+
+    N, _ = S.shape
+    N = N // 2
+    zerom = np.zeros([N, N])
+    idm = np.eye(N)
+    A = S[:N, :N]
+    B = S[:N, N:]
+    C = S[N:, :N]
+    D = S[N:, N:]
+    A0 = sqrtm(A @ A.T + B @ B.T)
+    A0inv = np.linalg.inv(A0)
+    X = A0inv @ A
+    Y = A0inv @ B
+    C0 = (C @ A.T + D @ B.T) @ A0inv
+    E = np.block([[idm, zerom], [C0 @ A0inv, idm]])
+    D = np.block([[A0, zerom], [zerom, A0inv]])
+    F = np.block([[X, Y], [-Y, X]])
+    return E, D, F
+
+
+def iwasawa(S):
+    """Iwasawa decomposition of a symplectic matrix.
+    See `Arvind et al. The Real Symplectic Groups in Quantum Mechanics and Optics <https://arxiv.org/pdf/quant-ph/9509002.pdf>`_
+
+
+    Args:
+        S (array): the symplectic matrix
+
+    Returns:
+        tuple[array, array, array]: (E,D,F) symplectic matrices such that E @ D @ F = S,
+        EE = np.block([[AA, np.zeros(N,N)],[CC, np.linalg.inv(A.T)]]) with A.T @ C == C.T @ A, and AA upper trinagular with ones in the diagonal
+        DD is diagonal and symplectic,
+        FF is symplectic orthogonal.
+    """
+
+    E, D, F = pre_iwasawa(S)
+    N, _ = S.shape
+    N = N // 2
+    DNN = D[:N, :N]
+    Q, R = qr(DNN)
+    R = R.T
+    Q = Q.T
+    dR = np.diag(R)
+    dd = np.abs(dR)
+    ds = np.sign(dR)
+    R = R * (1 / dR)
+    RinvT = np.linalg.inv(R).T
+    DD = np.diag(np.concatenate([dd, 1 / dd]))
+    zerom = np.zeros([N, N])
+    OO = np.block([[R, zerom], [zerom, RinvT]])
+    Q = ds[:, None] * Q
+    AA = np.block([[Q, zerom], [zerom, Q]])
+    EE = E @ OO
+    FF = AA @ F
+    return EE, DD, FF
diff --git a/thewalrus/tests/test_decompositions.py b/thewalrus/tests/test_decompositions.py
@@ -19,7 +19,7 @@
 
 from thewalrus.random import random_interferometer as haar_measure
 from thewalrus.random import random_symplectic
-from thewalrus.decompositions import williamson, blochmessiah, takagi
+from thewalrus.decompositions import williamson, blochmessiah, takagi, pre_iwasawa, iwasawa
 from thewalrus.symplectic import sympmat as omega
 from thewalrus.quantum.gaussian_checks import is_symplectic
 
@@ -441,3 +441,131 @@ def test_real_input_edge():
     # Now, reconstruct A, see
     Ar = u * l @ u.T
     assert np.allclose(A, Ar)
+
+
+@pytest.mark.parametrize("rank1", [2, 4, 5])
+@pytest.mark.parametrize("rank2", [2, 4, 5])
+@pytest.mark.parametrize("rankrand", [2, 4, 5])
+@pytest.mark.parametrize("rankzero", [2, 4, 5])
+@pytest.mark.parametrize("symmetric", [True, False])
+@pytest.mark.parametrize("unitary", [True, False])
+def test_pre_iwasawa(rank1, rank2, rankrand, rankzero, symmetric, unitary):
+    """Tests the pre_iwasawa decomposition"""
+    vals = np.array(
+        [np.random.rand(1)[0]] * rank1
+        + [np.random.rand(1)[0]] * rank2
+        + list(np.random.rand(rankrand))
+        + [1] * rankzero
+    )
+    if unitary is True:
+        vals = np.ones_like(vals)
+    dd = np.concatenate([vals, 1 / vals])
+    dim = len(vals)
+    U = haar_measure(dim)
+    O = np.block([[U.real, -U.imag], [U.imag, U.real]])
+    if symmetric is False:
+        V = haar_measure(dim)
+        P = np.block([[V.real, -V.imag], [V.imag, V.real]])
+    else:
+        P = O.T
+
+    S = (O * dd) @ P
+    EE, DD, FF = pre_iwasawa(S)
+    assert np.allclose(EE @ DD @ FF, S)
+    assert is_symplectic(EE)
+    assert is_symplectic(FF)
+    assert is_symplectic(FF)
+    assert np.allclose(FF @ FF.T, np.identity(2 * dim))
+    assert np.allclose(DD[:dim, :dim] @ DD[dim:, dim:], np.identity(dim))
+    assert np.allclose(DD[:dim, dim:], 0)
+    assert np.allclose(DD[dim:, :dim], 0)
+    A = EE[:dim, :dim]
+    B = EE[:dim, dim:]
+    C = EE[dim:, :dim]
+    D = EE[dim:, dim:]
+    assert np.allclose(A, np.eye(dim))
+    assert np.allclose(B, 0)
+    assert np.allclose(C, C.T)
+    assert np.allclose(D, np.eye(dim))
+
+
+@pytest.mark.parametrize("rank1", [2, 4, 5])
+@pytest.mark.parametrize("rank2", [2, 4, 5])
+@pytest.mark.parametrize("rankrand", [2, 4, 5])
+@pytest.mark.parametrize("rankzero", [2, 4, 5])
+@pytest.mark.parametrize("symmetric", [True, False])
+@pytest.mark.parametrize("unitary", [True, False])
+def test_iwasawa(rank1, rank2, rankrand, rankzero, symmetric, unitary):
+    """Tests the Iwasawa decomposition"""
+    vals = np.array(
+        [np.random.rand(1)[0]] * rank1
+        + [np.random.rand(1)[0]] * rank2
+        + list(np.random.rand(rankrand))
+        + [1] * rankzero
+    )
+    if unitary is True:
+        vals = np.ones_like(vals)
+    dd = np.concatenate([vals, 1 / vals])
+    dim = len(vals)
+    U = haar_measure(dim)
+    O = np.block([[U.real, -U.imag], [U.imag, U.real]])
+    if symmetric is False:
+        V = haar_measure(dim)
+        P = np.block([[V.real, -V.imag], [V.imag, V.real]])
+    else:
+        P = O.T
+    S = (O * dd) @ P
+    EE, DD, FF = iwasawa(S)
+    assert np.allclose(EE @ DD @ FF, S)
+    assert is_symplectic(EE)
+    assert is_symplectic(FF)
+    assert is_symplectic(FF)
+    assert np.allclose(FF @ FF.T, np.identity(2 * dim))
+    assert np.allclose(DD, np.diag(np.diag(DD)))
+    assert np.allclose(DD[:dim, :dim] @ DD[dim:, dim:], np.identity(dim))
+    A = EE[:dim, :dim]
+    B = EE[:dim, dim:]
+    C = EE[dim:, :dim]
+    D = EE[dim:, dim:]
+    assert np.allclose(B, 0)
+    XX = A.T @ C
+    assert np.allclose(XX, XX.T)
+    assert np.allclose(A @ D.T, np.eye(dim))
+    assert np.allclose(np.diag(EE), 1)
+    assert np.allclose(np.tril(A), A)
+    assert np.allclose(np.triu(D), D)
+
+
+def test_pre_iwasawa_error():
+    """Tests error is raised when input not symplectic"""
+    M = np.random.rand(4, 5)
+    with pytest.raises(ValueError, match="Input matrix is not symplectic."):
+        pre_iwasawa(M)
+
+
+def test_iwasawa_error():
+    """Tests error is raised when input not symplectic"""
+    M = np.random.rand(4, 5)
+    with pytest.raises(ValueError, match="Input matrix is not symplectic."):
+        iwasawa(M)
+
+
+def test_iwasawa2x2():
+    """Compares numerics against exact result for 2x2 matrices in Arvind 1995"""
+    num_tests = 100
+    for _ in range(num_tests):
+        S = random_symplectic(1)
+        A, N, K = iwasawa(S)
+        a = S[0, 0]
+        b = S[0, 1]
+        c = S[1, 0]
+        d = S[1, 1]
+        eta = a**2 + b**2
+        xi = (a * c + b * d) / eta
+        eta = np.sqrt(eta)
+        AA = np.array([[1, 0], [xi, 1]])
+        NN = np.diag([eta, 1 / eta])
+        KK = np.array([[a, b], [-b, a]]) / eta
+        assert np.allclose(A, AA)
+        assert np.allclose(K, KK)
+        assert np.allclose(N, NN)