custom dimensions for a multimode density matrix

.github/CHANGELOG.md

@@ -7,6 +7,7 @@
 ### Improvements
 
 * Improved speed of `_hermite_multidimensional_renorm` by 4x [(#404)](https://github.com/XanaduAI/thewalrus/pull/404).
+* Allow 'cutoff' to be a list of ints for the 'density_matrix(mu, cov, post_select, normalize, cutoff, hbar)' function from 'fock_tensors.py'. It enables to perform larger simulations by decreasing required RAM [(#406)](https://github.com/XanaduAI/thewalrus/pull/406).
 
 ### Bug fixes
 
@@ -20,7 +21,7 @@
 
 This release contains contributions from (in alphabetical order):
 
-T.H. Dodd, L.G. Helt, F. Miatto, A. Motamedi
+T.H. Dodd, L.G. Helt, F. Miatto, A. Motamedi, Ivan Solomakhin
 
 ---
 

thewalrus/quantum/fock_tensors.py

@@ -30,7 +30,11 @@
 from ..symplectic import expand, is_symplectic, reduced_state
 
 from .._hafnian import hafnian, hafnian_repeated, reduction
-from .._hermite_multidimensional import hermite_multidimensional, hafnian_batched, interferometer
+from .._hermite_multidimensional import (
+    hermite_multidimensional,
+    hafnian_batched,
+    interferometer,
+)
 
 from .conversions import (
     Amat,
@@ -257,33 +261,41 @@ def density_matrix(mu, cov, post_select=None, normalize=False, cutoff=5, hbar=2)
             is calculated directly using (multidimensional) Hermite polynomials, which is significantly faster
             than calculating one hafnian at a time.
         normalize (bool): If ``True``, a post-selected density matrix is re-normalized.
-        cutoff (dim): the final length (i.e., Hilbert space dimension) of each
-            mode in the density matrix.
+        cutoff (list[int] or int): the final length (i.e., Hilbert space dimension) of each
+            mode in the density matrix. if int, all modes have the same final length. If list[int], each mode can have
+            a custom final length
         hbar (float): the value of :math:`\hbar` in the commutation
             relation :math:`[\x,\p]=i\hbar`.
 
     Returns:
         np.array[complex]: the density matrix of the Gaussian state
     """
     N = len(mu) // 2
+    if isinstance(cutoff, int):
+        cutoff = [cutoff] * N
     pref = _prefactor(mu, cov, hbar=hbar)
 
     if post_select is None:
         A = Amat(cov, hbar=hbar).conj()
         sf_order = tuple(chain.from_iterable([[i, i + N] for i in range(N)]))
-
+        cutoff_modified = [elem for _ in range(2) for elem in cutoff]
         if np.allclose(mu, np.zeros_like(mu)):
-            tensor = pref * hermite_multidimensional(-A, cutoff, renorm=True, modified=True)
+            tensor = pref * hermite_multidimensional(
+                -A, cutoff_modified, renorm=True, modified=True
+            )
             return tensor.transpose(sf_order)
         beta = complex_to_real_displacements(mu, hbar=hbar)
         y = beta - A @ beta.conj()
-        tensor = pref * hermite_multidimensional(-A, cutoff, y=y, renorm=True, modified=True)
+        tensor = pref * hermite_multidimensional(
+            -A, cutoff_modified, y=y, renorm=True, modified=True
+        )
         return tensor.transpose(sf_order)
 
     M = N - len(post_select)
-    rho = np.zeros([cutoff] * (2 * M), dtype=np.complex128)
+    cutoff = [cutoff[i] for i in range(N) if not i in post_select]
+    rho = np.zeros([elem for _ in range(2) for elem in cutoff], dtype=np.complex128)
 
-    for idx in product(range(cutoff), repeat=2 * M):
+    for idx in product(*[range(n) for n in cutoff], repeat=2):
         el = []
 
         counter = count(0)
@@ -304,7 +316,8 @@ def density_matrix(mu, cov, post_select=None, normalize=False, cutoff=5, hbar=2)
     if normalize:
         # construct the standard 2D density matrix, and take the trace
         new_ax = np.arange(2 * M).reshape([M, 2]).T.flatten()
-        tr = np.trace(rho.transpose(new_ax).reshape([cutoff**M, cutoff**M])).real
+        rho = rho.transpose(new_ax).reshape([np.prod(cutoff), np.prod(cutoff)])
+        tr = np.trace(rho).real
         # renormalize
         rho /= tr
 

thewalrus/tests/test_quantum.py

@@ -93,7 +93,10 @@ def test_reduced_gaussian(n):
     assert np.all(
         res[1]
         == np.array(
-            [[(N + 1) * n, (N + 1) * n + m], [(N + 1) * n + N * m, (N + 1) * n + N * m + m]]
+            [
+                [(N + 1) * n, (N + 1) * n + m],
+                [(N + 1) * n + N * m, (N + 1) * n + N * m + m],
+            ]
         )
     )
 
@@ -363,6 +366,16 @@ def test_density_matrix_squeezed():
     assert np.allclose(res, expected)
 
 
+def test_density_matrix_vacuum_custom_cutoff():
+    """Test custom cutoff"""
+    mu = np.zeros([4])
+    V = np.identity(4)
+
+    res = density_matrix(mu, V, cutoff=[5, 3])
+
+    assert res.shape == (5, 5, 3, 3)
+
+
 def test_coherent_squeezed():
     """Test density matrix for a squeezed displaced state"""
     r = 0.43
@@ -895,11 +908,35 @@ def test_single_mode_squeezing(choi_r, tol):
     # np.array(squeeze(40,r).data.todense())[0:5,0:5]
     expected = np.array(
         [
-            [0.80501818 + 0.0j, 0.0 + 0.0j, 0.43352515 + 0.0j, 0.0 + 0.0j, 0.2859358 + 0.0j],
+            [
+                0.80501818 + 0.0j,
+                0.0 + 0.0j,
+                0.43352515 + 0.0j,
+                0.0 + 0.0j,
+                0.2859358 + 0.0j,
+            ],
             [0.0 + 0.0j, 0.52169547 + 0.0j, 0.0 + 0.0j, 0.48661591 + 0.0j, 0.0 + 0.0j],
-            [-0.43352515 + 0.0j, 0.0 + 0.0j, 0.10462138 + 0.0j, 0.0 + 0.0j, 0.29199268 + 0.0j],
-            [0.0 + 0.0j, -0.48661591 + 0.0j, 0.0 + 0.0j, -0.23479643 + 0.0j, 0.0 + 0.0j],
-            [0.2859358 + 0.0j, 0.0 + 0.0j, -0.29199268 + 0.0j, 0.0 + 0.0j, -0.34474749 + 0.0j],
+            [
+                -0.43352515 + 0.0j,
+                0.0 + 0.0j,
+                0.10462138 + 0.0j,
+                0.0 + 0.0j,
+                0.29199268 + 0.0j,
+            ],
+            [
+                0.0 + 0.0j,
+                -0.48661591 + 0.0j,
+                0.0 + 0.0j,
+                -0.23479643 + 0.0j,
+                0.0 + 0.0j,
+            ],
+            [
+                0.2859358 + 0.0j,
+                0.0 + 0.0j,
+                -0.29199268 + 0.0j,
+                0.0 + 0.0j,
+                -0.34474749 + 0.0j,
+            ],
         ]
     )
     T = fock_tensor(S, alphas, cutoff, choi_r=choi_r)
@@ -1223,7 +1260,8 @@ def test_loss_value_error(eta):
     """Tests the correct error is raised"""
     n = 50
     with pytest.raises(
-        ValueError, match="The transmission parameter eta should be a number between 0 and 1."
+        ValueError,
+        match="The transmission parameter eta should be a number between 0 and 1.",
     ):
         loss_mat(eta, n)
 
@@ -1471,7 +1509,18 @@ def test_expt_two_mode_squeezed(r, phi):
         [2, 2, 0, 0],
         [0, 0, 2, 2],
     ]
-    expected = [1, a, np.conj(a), adxa, adxa, a2, np.conj(a2), adxbdxab, ad2bd2, np.conj(ad2bd2)]
+    expected = [
+        1,
+        a,
+        np.conj(a),
+        adxa,
+        adxa,
+        a2,
+        np.conj(a2),
+        adxbdxab,
+        ad2bd2,
+        np.conj(ad2bd2),
+    ]
     for pattern, value in zip(patterns, expected):
         result = normal_ordered_expectation(means, cov, pattern, hbar=hbar)
         assert np.allclose(result, value)
@@ -1590,7 +1639,8 @@ def test_non_physical_cov_in_tvd():
     """Test the correct error is raised when an unphysical covariance matrix is passed to tvd_cutoff_bound"""
     A = np.array([[1, 2], [3, 4]])
     with pytest.raises(
-        ValueError, match="The input covariance matrix violates the uncertainty relation."
+        ValueError,
+        match="The input covariance matrix violates the uncertainty relation.",
     ):
         tvd_cutoff_bounds(np.zeros(2), A, 10)
 
@@ -1786,7 +1836,8 @@ def test_n_body_marginals_mismatch_shape():
     cov = squeezing([2 * r1, 0, 2 * r2])
     mu = np.zeros([4])
     with pytest.raises(
-        ValueError, match="The covariance matrix and vector of means have incompatible dimensions"
+        ValueError,
+        match="The covariance matrix and vector of means have incompatible dimensions",
     ):
         n_body_marginals(mu, cov, 4, 3)
 
@@ -1808,7 +1859,8 @@ def test_n_body_marginals_too_high_correlation():
     cov = squeezing([2 * r1, 0, 2 * r2])
     mu = np.zeros([6])
     with pytest.raises(
-        ValueError, match="The order of the correlations is higher than the number of modes"
+        ValueError,
+        match="The order of the correlations is higher than the number of modes",
     ):
         n_body_marginals(mu, cov, 4, 4)