passive transformation (#249)

* first commit for moving torontonian sampling to new probabilities function which includes displacement@ * fixed code, tests now pass * ran black on files * hbar passed to prob calcs * code factor stuff * try to fix failing test on the repo * remove forcing probs to be real * try to fix broken test * remove seed from test * new linear transformation function and tests@ * multimode squeezing function * run black on symplectic * update changelog * links in changelog * try to make code factor happy * expand passive transformation added * try to make code factor happy again * trailing whitespace * test for not specifying phi in squeezing * Update .github/CHANGELOG.md Co-authored-by: Nicolas Quesada <nicolas@xanadu.ai> * add docstring sentence to say squeezing can now be multimode * fixed expand_passive docstring * Update thewalrus/symplectic.py Co-authored-by: Christina Lee <chrissie.c.l@gmail.com> * Update thewalrus/symplectic.py Co-authored-by: Christina Lee <chrissie.c.l@gmail.com> * Update thewalrus/symplectic.py Co-authored-by: Christina Lee <chrissie.c.l@gmail.com> * Update thewalrus/symplectic.py Co-authored-by: Christina Lee <chrissie.c.l@gmail.com> * Update thewalrus/symplectic.py Co-authored-by: Christina Lee <chrissie.c.l@gmail.com> * Update thewalrus/symplectic.py Co-authored-by: Christina Lee <chrissie.c.l@gmail.com> * Update thewalrus/symplectic.py Co-authored-by: Christina Lee <chrissie.c.l@gmail.com> * new test and autosummary * hbar test * run black on symplectic.py * _ for unused variable * Apply suggestions from code review Co-authored-by: Christina Lee <chrissie.c.l@gmail.com> * Update thewalrus/symplectic.py Co-authored-by: Nicolas Quesada <nicolas@xanadu.ai> Co-authored-by: Nicolas Quesada <zeitus@gmail.com> Co-authored-by: Christina Lee <chrissie.c.l@gmail.com>
XanaduAI · Jun 22, 2021 · e32c40a · e32c40a
1 parent 355b8ff
commit e32c40a
Show file tree

Hide file tree

Showing 3 changed files with 253 additions and 13 deletions.
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
@@ -2,8 +2,16 @@
 
 ### New features
 
+
+* ``symplectic.squeezing`` function is now generalized to multiple modes of single mode squeezing [#249](https://github.com/XanaduAI/thewalrus/pull/249)
+
+* adds a function ``symplectic.passive_transformation`` which allows for Gaussian states to be transformed by arbitrary non-unitary, non-square linear optical transformations [#249](https://github.com/XanaduAI/thewalrus/pull/249)
+
+* ``torontonian_sample_state`` now can sample displaced Gaussian states [#248](https://github.com/XanaduAI/thewalrus/pull/248)
+
 * Adds the function `hafnian_banded` to calculate the hafnian of a banded matrix [#246](https://github.com/XanaduAI/thewalrus/pull/246)
 
+
 ### Improvements
 
 * Speeds up the calculation of photon number variances/covariances [#224](https://github.com/XanaduAI/thewalrus/pull/224)
@@ -16,7 +24,8 @@
 
 This release contains contributions from (in alphabetical order):
 
-Timjan Kalajdzievski, Nicolas Quesada
+Jake Bulmer, Timjan Kalajdzievski, Nicolas Quesada
+
 
 ---
 

diff --git a/thewalrus/symplectic.py b/thewalrus/symplectic.py
@@ -27,6 +27,7 @@
 .. autosummary::
     expand
     expand_vector
+    expand_passive
     reduced_state
     is_symplectic
     sympmat
@@ -45,6 +46,7 @@
 
 .. autosummary::
     two_mode_squeezing
+    squeezing
     interferometer
     loss
     mean_photon_number
@@ -99,8 +101,33 @@ def expand_vector(alpha, mode, N, hbar=2.0):
     return r
 
 
+def expand_passive(T, modes, N):
+    r"""Returns the expanded linear optical transformation
+    acting on specified modes, with identity acting on all other modes
+
+    Args:
+        T (array): square :math:`M \times M` matrix of linear optical transformation
+        modes (array): the :math:`M` modes of the transformation
+        N (int): number of modes in the new expanded transformation
+
+    Returns:
+        array: :math:`N \times N` array of expanded passive transformation
+    """
+
+    if T.shape[0] != T.shape[1]:
+        raise ValueError("The input matrix is not square")
+
+    if len(modes) != T.shape[0]:
+        raise ValueError("length of modes must match the shape of T")
+
+    T_expand = np.eye(N, dtype=T.dtype)
+    T_expand[np.ix_(modes, modes)] = T
+
+    return T_expand
+
+
 def reduced_state(mu, cov, modes):
-    r""" Returns the vector of means and the covariance matrix of the specified modes.
+    r"""Returns the vector of means and the covariance matrix of the specified modes.
 
     Args:
         modes (int of Sequence[int]): indices of the requested modes
@@ -146,22 +173,40 @@ def vacuum_state(modes, hbar=2.0, dtype=np.float64):
     return state
 
 
-def squeezing(r, phi, dtype=np.float64):
-    r"""Squeezing. In fock space this corresponds to \exp(\tfrac{1}{2}r e^{i \phi} (a^2 - a^{\dagger 2}) ).
+def squeezing(r, phi=None, dtype=np.float64):
+    r"""Squeezing. In fock space this corresponds to:
+
+    .. math::
+
+         \exp(\tfrac{1}{2}r e^{i \phi} (a^2 - a^{\dagger 2}) ).
+
+
+    By passing an array of squeezing parameters and phases, it applies a tensor product of squeezing operations.
 
     Args:
-        r (float): squeezing magnitude
-        phi (float): rotation parameter
-        dtype (numpy.dtype): datatype to represent the Symplectic matrix
+        r (Union[array, float]): squeezing magnitude
+        phi (Union[array, float]): rotation parameter. If ``None``, then the function uses zeros of the same shape as ``r``.
+        dtype (numpy.dtype): datatype to represent the Symplectic matrix.  Defaults to ``numpy.float64``.
     Returns:
         array: symplectic transformation matrix
     """
     # pylint: disable=assignment-from-no-return
-    cp = np.cos(phi, dtype=dtype)
-    sp = np.sin(phi, dtype=dtype)
-    ch = np.cosh(r, dtype=dtype)
-    sh = np.sinh(r, dtype=dtype)
-    S = np.array([[ch - cp * sh, -sp * sh], [-sp * sh, ch + cp * sh]])
+
+    r = np.atleast_1d(r)
+
+    if phi is None:
+        phi = np.zeros_like(r)
+
+    phi = np.atleast_1d(phi)
+
+    M = len(r)
+    S = np.identity(2 * M, dtype=dtype)
+
+    for i, (r_i, phi_i) in enumerate(zip(r, phi)):
+        S[i, i] = np.cosh(r_i) - np.sinh(r_i) * np.cos(phi_i)
+        S[i, i + M] = -np.sinh(r_i) * np.sin(phi_i)
+        S[i + M, i] = -np.sinh(r_i) * np.sin(phi_i)
+        S[i + M, i + M] = np.cosh(r_i) + np.sinh(r_i) * np.cos(phi_i)
 
     return S
 
@@ -211,6 +256,32 @@ def interferometer(U):
     return S
 
 
+def passive_transformation(mu, cov, T, hbar=2):
+    r"""Perform a covariance matrix transformation for an arbitrary linear optical channel
+    on an :math:`N` modes state mapping it to a to an :math:`M` modes state.
+
+    Args:
+        mu (array): :math:`2N`-length means vector
+        cov (array): :math:`2N \times 2N` covariance matrix
+        T (array): :math:`M \times N` linear optical transformation
+
+    Keyword Args:
+        hbar (float)=2: the value to use for hbar
+
+    Returns:
+        array: :math:`2M`-length transformed means vector
+        array :math:`2M \times 2M` tranformed covariance matrix
+    """
+
+    P = interferometer(T)
+    L = (hbar / 2) * (np.eye(P.shape[0]) - P @ P.T)
+
+    cov = P @ cov @ P.T + L
+    mu = P @ mu
+
+    return mu, cov
+
+
 # pylint: disable=too-many-arguments
 def loss(mu, cov, T, mode, nbar=0, hbar=2):
     r"""Loss channel acting on a Gaussian state.
@@ -241,6 +312,7 @@ def loss(mu, cov, T, mode, nbar=0, hbar=2):
 
     return mu_res, cov_res
 
+
 ### Comment: This function strongly overlaps with `quantum.photon_number_mean`
 ### Wonder if it is worth removing it.
 def mean_photon_number(mu, cov, hbar=2):
@@ -314,7 +386,7 @@ def sympmat(N, dtype=np.float64):
 
 
 def is_symplectic(S, rtol=1e-05, atol=1e-08):
-    r""" Checks if matrix S is a symplectic matrix
+    r"""Checks if matrix S is a symplectic matrix
 
     Args:
         S (array): a square matrix
@@ -332,6 +404,7 @@ def is_symplectic(S, rtol=1e-05, atol=1e-08):
 
     return np.allclose(S.T @ Omega @ S, Omega, rtol=rtol, atol=atol)
 
+
 def autonne(A, rtol=1e-05, atol=1e-08, svd_order=True):
     r"""Autonne-Takagi decomposition of a complex symmetric (not Hermitian!) matrix.
 
@@ -364,6 +437,7 @@ def autonne(A, rtol=1e-05, atol=1e-08, svd_order=True):
         return (vals[n : 2 * n])[::-1], U[:, ::-1]
     return vals[n : 2 * n], U
 
+
 def xxpp_to_xpxp(S):
     """Permutes the entries of the input from xxpp ordering to xpxp ordering.
 
@@ -389,6 +463,7 @@ def xxpp_to_xpxp(S):
 
     return S[ind]
 
+
 def xpxp_to_xxpp(S):
     """Permutes the entries of the input from xpxp ordering to xxpp ordering.
 

diff --git a/thewalrus/tests/test_symplectic.py b/thewalrus/tests/test_symplectic.py
@@ -19,6 +19,7 @@
 from scipy.linalg import block_diag
 
 from thewalrus import symplectic
+from thewalrus.quantum import is_valid_cov
 
 
 # pylint: disable=too-few-public-methods
@@ -59,6 +60,53 @@ def test_squeezing(self, tol):
         expected = rotation @ np.diag(np.exp([-2*r, 2*r])) @ rotation.T
         assert np.allclose(out, expected, atol=tol, rtol=0)
 
+    def test_squeezing_no_phi(self, tol):
+        """Test the squeezing symplectic transform without specifying phi"""
+        r = 0.543
+        phi = 0.
+        S = symplectic.squeezing(r)
+        out = S @ S.T
+
+        # apply to an identity covariance matrix
+        rotation = np.array([[np.cos(phi/2), -np.sin(phi/2)], [np.sin(phi/2), np.cos(phi/2)]])
+        expected = rotation @ np.diag(np.exp([-2*r, 2*r])) @ rotation.T
+        assert np.allclose(out, expected, atol=tol, rtol=0)
+
+    def test_squeezing_no_phi_array(self, tol):
+        """Test multimode squeezing symplectic transform without specifying phi"""
+        r = np.random.randn(6)
+        phi = np.zeros_like(r)
+
+        S = symplectic.squeezing(r)
+        S_phi = symplectic.squeezing(r, phi)
+
+        assert np.allclose(S, S_phi, atol=tol, rtol=0)
+
+    def test_symplectic_multimode(self, tol):
+        """Test multimode version gives symplectic matrix"""
+        r = [0.543] * 4
+        phi = [0.123] * 4
+        S = symplectic.squeezing(r, phi)
+
+        # the symplectic matrix
+        O = symplectic.sympmat(4)
+
+        assert np.allclose(S @ O @ S.T, O, atol=tol, rtol=0)
+
+    def test_dtype(self, tol):
+        """Test multimode version gives symplectic matrix"""
+        r = [0.543] * 4
+        phi = [0.123] * 4
+        S = symplectic.squeezing(r, phi)
+
+        S32_bit = symplectic.squeezing(r, phi, dtype=np.float32)
+
+        # the symplectic matrix
+        O = symplectic.sympmat(4)
+
+        assert np.allclose(S32_bit @ O @ S32_bit.T, O, atol=tol, rtol=0)
+        assert np.allclose(S, S32_bit, atol=tol, rtol=0)
+
 
 class TestTwoModeSqueezing:
     """Tests for the TMS symplectic"""
@@ -176,6 +224,90 @@ def test_rotation(self, tol):
         np.allclose(S, expected, atol=tol, rtol=0)
 
 
+class TestPassiveTransformation:
+    """ tests for linear transformation """
+    def test_transformation(self, tol):
+        """Test that an transformation returns the correct state"""
+
+        M = 4
+        cov = np.arange(4 * M ** 2, dtype=np.float64).reshape((2*M, 2*M))
+        mu = np.arange(2 * M, dtype=np.float64)
+
+        T = np.sqrt(0.9) * M ** (-0.5) * np.ones((6,M), dtype=np.float64)
+
+        mu_out, cov_out = symplectic.passive_transformation(mu, cov, T)
+        # fmt:off
+        expected_mu = np.array([ 2.84604989,  2.84604989,  2.84604989,  2.84604989,  2.84604989,
+                                 2.84604989, 10.43551628, 10.43551628, 10.43551628, 10.43551628,
+                                 10.43551628, 10.43551628])
+        expected_cov = np.array([
+            [ 48.7,  47.7,  47.7,  47.7,  47.7,  47.7,  63. ,  63. ,  63. , 63. ,  63. ,  63. ],
+            [ 47.7,  48.7,  47.7,  47.7,  47.7,  47.7,  63. ,  63. ,  63. , 63. ,  63. ,  63. ],
+            [ 47.7,  47.7,  48.7,  47.7,  47.7,  47.7,  63. ,  63. ,  63. , 63. ,  63. ,  63. ],
+            [ 47.7,  47.7,  47.7,  48.7,  47.7,  47.7,  63. ,  63. ,  63. , 63. ,  63. ,  63. ],
+            [ 47.7,  47.7,  47.7,  47.7,  48.7,  47.7,  63. ,  63. ,  63. , 63. ,  63. ,  63. ],
+            [ 47.7,  47.7,  47.7,  47.7,  47.7,  48.7,  63. ,  63. ,  63. , 63. ,  63. ,  63. ],
+            [163.8, 163.8, 163.8, 163.8, 163.8, 163.8, 178.3, 177.3, 177.3, 177.3, 177.3, 177.3],
+            [163.8, 163.8, 163.8, 163.8, 163.8, 163.8, 177.3, 178.3, 177.3, 177.3, 177.3, 177.3],
+            [163.8, 163.8, 163.8, 163.8, 163.8, 163.8, 177.3, 177.3, 178.3, 177.3, 177.3, 177.3],
+            [163.8, 163.8, 163.8, 163.8, 163.8, 163.8, 177.3, 177.3, 177.3, 178.3, 177.3, 177.3],
+            [163.8, 163.8, 163.8, 163.8, 163.8, 163.8, 177.3, 177.3, 177.3, 177.3, 178.3, 177.3],
+            [163.8, 163.8, 163.8, 163.8, 163.8, 163.8, 177.3, 177.3, 177.3, 177.3, 177.3, 178.3]])
+        # fmt:on
+
+        assert np.allclose(mu_out, expected_mu, atol=tol, rtol=0)
+        assert np.allclose(cov_out, expected_cov, atol=tol, rtol=0)
+
+    @pytest.mark.parametrize("M", range(1,10))
+    def test_valid_cov(self, M, tol):
+        """test that the output is a valid covariance matrix, even when not square"""
+        a = np.arange(4 * M ** 2, dtype=np.float64).reshape((2*M, 2*M))
+        cov = a @ a.T + np.eye(2*M)
+        mu = np.arange(2 * M, dtype=np.float64)
+
+        T = np.sqrt(0.9) * M ** (-0.5) * np.ones((6,M), dtype=np.float64)
+
+        mu_out, cov_out = symplectic.passive_transformation(mu, cov, T)
+
+        assert cov_out.shape == (12, 12)
+        assert len(mu_out) == 12
+        assert is_valid_cov(cov_out, atol=tol, rtol=0)
+
+    @pytest.mark.parametrize("M", range(1, 6))
+    def test_unitary(self, M, tol):
+        """
+        test that the outputs agree with the interferometer class when
+        transformation is unitary
+        """
+        a = np.arange(4 * M ** 2, dtype=np.float64).reshape((2 * M, 2 * M))
+        cov = a @ a.T + np.eye(2 * M)
+        mu = np.arange(2 * M, dtype=np.float64)
+
+        U = M ** (-0.5) * np.fft.fft(np.eye(M))
+        S_U = symplectic.interferometer(U)
+        cov_U = S_U @ cov @ S_U.T
+        mu_U = S_U @ mu
+
+        mu_T, cov_T = symplectic.passive_transformation(mu, cov, U)
+
+        assert np.allclose(mu_U, mu_T, atol=tol, rtol=0)
+        assert np.allclose(cov_U, cov_T, atol=tol, rtol=0)
+
+    @pytest.mark.parametrize("hbar", [1,2, 1.05e-34])
+    def test_hbar(self, hbar, tol):
+        """test that the output is a valid covariance matrix, even when not square"""
+
+        M = 4
+        a = np.arange(4 * M ** 2, dtype=np.float64).reshape((2*M, 2*M))
+        cov = a @ a.T + np.eye(2*M)
+        mu = np.arange(2 * M, dtype=np.float64)
+
+        T = np.sqrt(0.9) * M ** (-0.5) * np.ones((6,M), dtype=np.float64)
+
+        _, cov_out = symplectic.passive_transformation(mu, cov, T, hbar=hbar)
+
+        assert is_valid_cov(cov_out, hbar=hbar, atol=tol, rtol=0)
+
 class TestReducedState:
     """Tests for the reduced state function"""
 
@@ -465,6 +597,30 @@ def test_expand_one(self, hbar, tol):
         expected[mode + N] = np.sqrt(2 * hbar) * alpha.imag
         assert np.allclose(r, expected, atol=tol, rtol=0)
 
+class TestExpandPassive:
+    """Tests for expanding a displacement operation into a phase-space displacement vector"""
+
+    def test_expand_one(self, tol):
+        """Test that a 1x1 matrix is expanded correctly"""
+        T = np.array([[0.5]])
+
+        T_expand = symplectic.expand_passive(T, [1], 3)
+
+        expected = np.array([[1,0,0],[0,0.5,0],[0,0,1]])
+
+        assert np.allclose(T_expand, expected, atol=tol, rtol=0)
+
+    def test_expend_not_square(self):
+        """ test that error is raised for non square input"""
+        with pytest.raises(ValueError, match="The input matrix is not square"):
+            symplectic.expand_passive(np.ones((3,2)), [0,1,2], 5)
+
+    def test_modes_length(self):
+        """ test that error is raised when length of modes array is incorrect"""
+        with pytest.raises(ValueError, match="length of modes must match the shape of T"):
+            symplectic.expand_passive(np.ones((3,3)), [0,1,2,3,4], 8)
+
+
 class TestSymplecticExpansion:
     """Tests for the expanding a symplectic matrix"""