Stable dgate fock

mrmustard/lab_dev/transformations/dgate.py

@@ -19,10 +19,13 @@
 from __future__ import annotations
 
 from typing import Sequence
+from mrmustard.utils.typing import ComplexTensor
+
+from mrmustard import math
 
 from .base import Unitary
 from ...physics.representations import Bargmann
-from ...physics import triples
+from ...physics import triples, fock
 from ..utils import make_parameter, reshape_params
 
 __all__ = ["Dgate"]
@@ -95,3 +98,47 @@ def __init__(
         self._representation = Bargmann.from_function(
             fn=triples.displacement_gate_Abc, x=self.x, y=self.y
         )
+
+    def fock(self, shape: int | Sequence[int] = None, batched=False) -> ComplexTensor:
+        r""", shape: Optional[int | Sequence[int]] = None, batched=False) -> CircuitComponent:
+        Returns the unitary representation of the Displacement gate using the Laguerre polynomials.
+        If the shape is not given, it defaults to the ``auto_shape`` of the component if it is
+        available, otherwise it defaults to the value of ``AUTOSHAPE_MAX`` in the settings.
+        Args:
+            shape: The shape of the returned representation. If ``shape`` is given as an ``int``,
+                it is broadcasted to all the dimensions. If not given, it is estimated.
+            batched: Whether the returned representation is batched or not. If ``False`` (default)
+                it will squeeze the batch dimension if it is 1.
+        Returns:
+            array: The Fock representation of this component.
+        """
+        if isinstance(shape, int):
+            shape = (shape,) * self.representation.ansatz.num_vars
+        auto_shape = self.auto_shape()
+        shape = shape or auto_shape
+        if len(shape) != len(auto_shape):
+            raise ValueError(
+                f"Expected Fock shape of length {len(auto_shape)}, got length {len(shape)}"
+            )
+        N = self.n_modes
+        x = self.x.value * math.ones(N, dtype=self.x.value.dtype)
+        y = self.y.value * math.ones(N, dtype=self.y.value.dtype)
+
+        if N > 1:
+            # calculate displacement unitary for each mode and concatenate with outer product
+            Ud = None
+            for idx, out_in in enumerate(zip(shape[:N], shape[N:])):
+                if Ud is None:
+                    Ud = fock.displacement(x[idx], y[idx], shape=out_in)
+                else:
+                    U_next = fock.displacement(x[idx], y[idx], shape=out_in)
+                    Ud = math.outer(Ud, U_next)
+
+            array = math.transpose(
+                Ud,
+                list(range(0, 2 * N, 2)) + list(range(1, 2 * N, 2)),
+            )
+        else:
+            array = fock.displacement(x[0], y[0], shape=shape)
+        arrays = math.expand_dims(array, 0) if batched else array
+        return arrays

tests/test_lab_dev/test_transformations/test_dgate.py

@@ -17,9 +17,9 @@
 # pylint: disable=protected-access, missing-function-docstring, expression-not-assigned
 
 import pytest
-
+import numpy as np
 from mrmustard import math
-from mrmustard.lab_dev.transformations import Dgate
+from mrmustard.lab_dev import Dgate, SqueezedVacuum
 from mrmustard.physics.representations import Fock
 
 
@@ -46,6 +46,14 @@ def test_init_error(self):
         with pytest.raises(ValueError, match="y"):
             Dgate(modes=[0, 1], x=1, y=[2, 3, 4])
 
+    def test_to_fock_method(self):
+        # test stable Dgate in fock basis
+        state = SqueezedVacuum([0], r=1.0)
+        # displacement gate in fock representation for large displacement
+        dgate = Dgate([0], x=10.0).to_fock(150)
+        assert (state.to_fock() >> dgate).probability < 1
+        assert np.all(math.abs(dgate.fock(150)) < 1)
+
     def test_representation(self):
         rep1 = Dgate(modes=[0], x=0.1, y=0.1).representation
         assert math.allclose(rep1.A, [[[0, 1], [1, 0]]])