Instabilities of displacement operator (#351)

* impl [fock_gradientes] - displacement with laguerre

* fix [fock_gradients] - displacement

* refactor [fock_gradients] - displacement: calculate factorial once

* test [fock_gradients] - laguerre polys

* fix [fock_gradients] - displacement: missing conjugate

Co-authored-by: JacobHast <jacobhastrup@gmail.com>

* Laguerre renorm avoids division by large factorial

* breaking a long line into two shorter ones

* removed factorial

* boom

* better docstring and variable names

* adds gradient of displacement gate

* format - black

* re-adds  to test

* remove coverage from jitted `_laguerre`

* add changelog entry

Co-authored-by: JacobHast <jacobhastrup@gmail.com>
Co-authored-by: ziofil <miatto@gmail.com>

.github/CHANGELOG.md

@@ -8,6 +8,9 @@
 
 * Added function to extend single mode symplectic to act on multiple modes. [(#347)](https://github.com/XanaduAI/thewalrus/pull/347)
 
+* The displacement function now calculates the displacement operator's matrix elements with respect to the Fock basis using the Laguerre
+  polynomials. This provides enhanced numerical stability for larger cutoff values. [(#351)](https://github.com/XanaduAI/thewalrus/pull/351)
+
 ### Bug fixes
 
 * Remove redundant call of `Qmat`, `Amat` from `generate_hafnian_sample`. [(#343)](https://github.com/XanaduAI/thewalrus/pull/343)
@@ -22,7 +25,7 @@
 
 This release contains contributions from (in alphabetical order):
 
-Mikhail Andrenkov, Antonín Hoskovec
+Mikhail Andrenkov, Sebastián Duque, Jacon Hastrup, Antonín Hoskovec, Filippo Miatto
 
 ---
 

thewalrus/fock_gradients.py

@@ -40,13 +40,15 @@
 ------------
 """
 import numpy as np
-
 from numba import jit
 
+SQRT = np.sqrt(np.arange(1e6))
+
 
 @jit(nopython=True)
 def displacement(r, phi, cutoff, dtype=np.complex128):  # pragma: no cover
     r"""Calculates the matrix elements of the displacement gate using a recurrence relation.
+    Uses the log of the matrix elements to avoid numerical issues and then takes the exponential.
 
     Args:
         r (float): displacement magnitude
@@ -58,18 +60,40 @@ def displacement(r, phi, cutoff, dtype=np.complex128):  # pragma: no cover
         array[complex]: matrix representing the displacement operation.
     """
     D = np.zeros((cutoff, cutoff), dtype=dtype)
-    sqrt = np.sqrt(np.arange(cutoff, dtype=dtype))
-    mu = np.array([r * np.exp(1j * phi), -r * np.exp(-1j * phi)])
+    rng = np.arange(cutoff)
+    rng[0] = 1
+    log_k_fac = np.cumsum(np.log(rng))
+    for n_minus_m in range(cutoff):
+        m_max = cutoff - n_minus_m
+        logL = np.log(_laguerre(r**2.0, m_max, n_minus_m))
+        for m in range(m_max):
+            n = n_minus_m + m
+            D[n, m] = np.exp(
+                0.5 * (log_k_fac[m] - log_k_fac[n])
+                + n_minus_m * np.log(r)
+                - (r**2.0) / 2.0
+                + 1j * phi * n_minus_m
+                + logL[m]
+            )
+            D[m, n] = (-1.0) ** (n_minus_m) * np.conj(D[n, m])
+    return D
 
-    D[0, 0] = np.exp(-0.5 * r**2)
-    for m in range(1, cutoff):
-        D[m, 0] = mu[0] / sqrt[m] * D[m - 1, 0]
 
-    for m in range(cutoff):
-        for n in range(1, cutoff):
-            D[m, n] = mu[1] / sqrt[n] * D[m, n - 1] + sqrt[m] / sqrt[n] * D[m - 1, n - 1]
+@jit(nopython=True, cache=True)
+def _laguerre(x, N, alpha, dtype=np.complex128):  # pragma: no cover
+    r"""Returns the N first generalized Laguerre polynomials evaluated at x.
 
-    return D
+    Args:
+        x (float): point at which to evaluate the polynomials
+        N (int): maximum Laguerre polynomial to calculate
+        alpha (float): continuous parameter for the generalized Laguerre polynomials
+    """
+    L = np.zeros(N, dtype=dtype)
+    L[0] = 1.0
+    if N > 1:
+        for m in range(0, N - 1):
+            L[m + 1] = ((2 * m + 1 + alpha - x) * L[m] - (m + alpha) * L[m - 1]) / (m + 1)
+    return L
 
 
 @jit(nopython=True)

thewalrus/tests/test_fock_gradients.py

@@ -24,8 +24,10 @@
     grad_beamsplitter,
     mzgate,
     grad_mzgate,
+    _laguerre,
 )
 import numpy as np
+from scipy.special import genlaguerre
 import pytest
 
 
@@ -177,6 +179,17 @@ def test_displacement_values(tol):
     assert np.allclose(T, expected, atol=tol, rtol=0)
 
 
+@pytest.mark.parametrize("alpha", np.linspace(-0.9, 10, 15))
+@pytest.mark.parametrize("x", np.linspace(-5, 10, 15))
+def test_laguerre_values(tol, x, alpha):
+    """Test recursive calculation of laguerre polynomials give the correct result"""
+    N = 10
+    res = _laguerre(x, N, alpha)
+    expected = np.array([genlaguerre(n=ni, alpha=alpha)(x) for ni in range(N)])
+
+    assert np.allclose(res, expected, atol=tol)
+
+
 def test_squeezing_values(tol):
     """Tests the correct construction of the single mode squeezing operation"""
     r = 0.5