New cft

mrmustard/lab_dev/states/base.py

@@ -329,7 +329,9 @@ def phase_space(self, s: float) -> tuple:
             raise ValueError("Can calculate phase space only for Bargmann states.")
 
         new_state = self >> BtoPS(self.modes, s=s)
-        return bargmann_Abc_to_phasespace_cov_means(*new_state.bargmann_triple(batched=True))
+        return bargmann_Abc_to_phasespace_cov_means(
+            *new_state.bargmann_triple(batched=True), batched=True
+        )
 
     def quadrature_distribution(self, quad: Vector, phi: float = 0.0) -> tuple | ComplexTensor:
         r"""

mrmustard/lab_dev/transformations/__init__.py

@@ -26,3 +26,4 @@
 from .rgate import *
 from .s2gate import *
 from .sgate import *
+from .cft import *

mrmustard/lab_dev/transformations/cft.py

@@ -0,0 +1,52 @@
+# Copyright 2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+The class representing a complex fourier transform.
+"""
+
+from typing import Sequence
+from mrmustard.lab_dev.transformations.base import Map
+from mrmustard.physics.representations import Bargmann
+from mrmustard.physics import triples
+
+__all__ = ["CFT"]
+
+
+class CFT(Map):
+    r"""The Complex Fourier Transformation as a channel.
+    The main use is to convert between Characteristic functions and phase space functions.
+
+    Args:
+        num_modes: number of modes of this channel.
+
+    .. details::
+        For a function f defined on the complex domain, we can define the following Complex Fourier Transform (aka symplectic Fourier Transform):
+        .. math::
+            \hat f(\zeta) = \int f(z) \exp(z \bar{\zeta} - \bar{z} \zeta)
+        this formulation can be extended to multi-variables by taking the transform over many (i.e., by as many zeta variables as needed)
+    """
+
+    def __init__(
+        self,
+        modes: Sequence[int],
+    ):
+        super().__init__(
+            modes_out=modes,
+            modes_in=modes,
+            name="CFT",
+        )
+        self._representation = Bargmann.from_function(
+            fn=triples.complex_fourier_transform_Abc, n_modes=len(modes)
+        )

mrmustard/physics/ansatze.py

@@ -1184,7 +1184,7 @@ def __truediv__(self, other: Scalar | ArrayAnsatz) -> ArrayAnsatz:
 
 
 def bargmann_Abc_to_phasespace_cov_means(
-    A: Matrix, b: Vector, c: Scalar
+    A: Matrix, b: Vector, c: Scalar, batched: bool = False
 ) -> tuple[Matrix, Vector, Scalar]:
     r"""
     Function to derive the covariance matrix and mean vector of a Gaussian state from its Wigner characteristic function in ABC form.
@@ -1208,6 +1208,10 @@ def bargmann_Abc_to_phasespace_cov_means(
     Returns:
         The covariance matrix, mean vector and coefficient of the state in phase space.
     """
+    # batched = len(A.shape) == 3 and len(b.shape) == 2 and len(c.shape) == 1
+    A = math.atleast_3d(A)
+    b = math.atleast_2d(b)
+    c = math.atleast_1d(c)
     num_modes = A.shape[-1] // 2
     Omega = math.cast(math.transpose(math.J(num_modes)), dtype=math.complex128)
     W = math.transpose(math.conj(math.rotmat(num_modes)))
@@ -1219,4 +1223,6 @@ def bargmann_Abc_to_phasespace_cov_means(
         1j * math.matvec(Omega @ W, bvec) * math.sqrt(settings.HBAR, dtype=math.complex128)
         for bvec in b
     ]
-    return math.astensor(cov), math.astensor(mean), coeff
+    if batched:
+        return math.astensor(cov), math.astensor(mean), coeff
+    return cov[0], mean[0], coeff[0]

mrmustard/physics/triples.py

@@ -677,10 +677,32 @@ def displacement_map_s_parametrized_Abc(s: int, n_modes: int) -> Union[Matrix, V
 
     A = math.astensor(math.asnumpy(A)[order_list, :][:, order_list])
     b = _vacuum_B_vector(4 * n_modes)
-    c = 1.0 + 0j
+    c = 1.0 / (2 * np.pi) ** n_modes + 0.0j
     return math.astensor(A), b, c
 
 
+def complex_fourier_transform_Abc(n_modes: int) -> tuple[Matrix, Vector, Scalar]:
+    r"""
+    The ``(A, b, c)`` triple of the complex Fourier transform between two pairs of complex variables.
+    Given a function :math:`f(z^*, z)`, the complex Fourier transform is defined as
+    :math:
+        \hat{f} (y^*, y) = \int_{\mathbb{C}} \frac{d^2 z}{\pi} e^{yz^* - y^*z} f(z^*, z).
+    The indices of this triple correspond to the variables :math:`(y^*, z^*, y, z)`.
+
+    Args:
+        n_modes: the number of modes for this map.
+
+    Returns:
+        The ``(A, b, c)`` triple of the complex fourier transform.
+    """
+    O2n = math.zeros((2 * n_modes, 2 * n_modes))
+    Omega = math.J(n_modes)
+    A = math.block([[O2n, -Omega], [Omega, O2n]])
+    b = _vacuum_B_vector(4 * n_modes)
+    c = 1.0 + 0j
+    return A, b, c
+
+
 # ~~~~~~~~~~~~~~~~
 # Kraus operators
 # ~~~~~~~~~~~~~~~~

tests/test_lab_dev/test_states/test_states_base.py

@@ -144,7 +144,7 @@ def test_to_from_fock(self, modes):
     @pytest.mark.parametrize("modes", [[0], [0, 1], [3, 19, 2]])
     def test_to_from_phase_space(self, modes):
         cov, means, coeff = Coherent([0], x=1, y=2).phase_space(s=0)
-        assert math.allclose(coeff[0], 1.0)
+        assert math.allclose(coeff[0], 1.0 / (2 * np.pi))
         assert math.allclose(cov[0], np.eye(2) * settings.HBAR / 2)
         assert math.allclose(means[0], np.array([1.0, 2.0]) * np.sqrt(2 * settings.HBAR))
         n_modes = len(modes)
@@ -552,7 +552,7 @@ def test_to_from_fock(self, modes):
     def test_to_from_phase_space(self):
         state0 = Coherent([0], x=1, y=2) >> Attenuator([0], 1.0)
         cov, means, coeff = state0.phase_space(s=0)  # batch = 1
-        assert coeff[0] == 1.0
+        assert coeff == 1.0 / (2 * np.pi)
         assert math.allclose(cov[0], np.eye(2) * settings.HBAR / 2)
         assert math.allclose(means[0], np.array([1.0, 2.0]) * np.sqrt(settings.HBAR * 2))
 

tests/test_lab_dev/test_transformations/test_cft.py

@@ -0,0 +1,54 @@
+# Copyright 2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the ``CFT`` class."""
+
+import numpy as np
+from mrmustard import math, settings
+from mrmustard.lab_dev import BtoPS, Ket, Dgate
+from mrmustard.lab_dev.transformations.cft import CFT
+from mrmustard.physics.wigner import wigner_discretized
+
+
+class TestCFT:
+    r"""
+    Tests the CFT gate
+    """
+
+    def test_init(self):
+        "tests the initialization of the CFT gate"
+        cft = CFT([0])
+        assert cft.name == "CFT"
+        assert cft.modes == [0]
+
+    def test_wigner_function(self):
+        r"""
+        Tests that the characteristic function is converted to the Wigner function
+        for a single-mode squeezed state.
+        """
+
+        state = Ket.random([0]) >> Dgate([0], x=1.0, y=0.1)
+
+        dm = math.sum(state.to_fock(100).dm().representation.array, axes=[0])
+        vec = np.linspace(-5, 5, 100)
+        wigner, _, _ = wigner_discretized(dm, vec, vec)
+
+        Wigner = (state >> CFT([0]).inverse() >> BtoPS([0], s=0)).representation.ansatz
+        X, Y = np.meshgrid(
+            vec * np.sqrt(2 / settings.HBAR), vec * np.sqrt(2 / settings.HBAR)
+        )  # scaling to take care of HBAR
+        Z = np.array([X - 1j * Y, X + 1j * Y]).transpose((1, 2, 0))
+        assert math.allclose(
+            2 / settings.HBAR * (np.real(Wigner(Z))), (np.real(wigner.T)), atol=1e-8
+        )  # scaling to take care of HBAR

tests/test_physics/test_ansatz.py

@@ -233,15 +233,15 @@ def test_bargmann_Abc_to_phasespace_cov_means(self):
         state_means = np.array([0.2, 0.3])
         state = DM.from_bargmann([0], wigner_to_bargmann_rho(state_cov, state_means))
         state_after = state >> BtoPS(modes=[0], s=0)  # pylint: disable=protected-access
-        A1, b1, c1 = state_after.bargmann_triple(batched=True)
+        A1, b1, c1 = state_after.bargmann_triple()
         (
             new_state_cov,
             new_state_means,
             new_state_coeff,
         ) = bargmann_Abc_to_phasespace_cov_means(A1, b1, c1)
-        assert np.allclose(state_cov, new_state_cov[0])
-        assert np.allclose(state_means, new_state_means[0])
-        assert np.allclose(1.0, new_state_coeff[0])
+        assert np.allclose(state_cov, new_state_cov)
+        assert np.allclose(state_means, new_state_means)
+        assert np.allclose(1.0 / (2 * np.pi), new_state_coeff)
 
         state_cov = np.array(
             [
@@ -256,27 +256,25 @@ def test_bargmann_Abc_to_phasespace_cov_means(self):
         state = DM.from_bargmann(modes=[0, 1], triple=(A, b, c))
 
         state_after = state >> BtoPS(modes=[0, 1], s=0)  # pylint: disable=protected-access
-        A1, b1, c1 = state_after.bargmann_triple(batched=True)
+        A1, b1, c1 = state_after.bargmann_triple()
         (
             new_state_cov1,
             new_state_means1,
             new_state_coeff1,
         ) = bargmann_Abc_to_phasespace_cov_means(A1, b1, c1)
 
-        A22, b22, c22 = (state >> BtoPS([0], 0) >> BtoPS([1], 0)).bargmann_triple(
-            batched=True
-        )  # pylint: disable=protected-access
+        A22, b22, c22 = (state >> BtoPS([0], 0) >> BtoPS([1], 0)).bargmann_triple()
         (
             new_state_cov22,
             new_state_means22,
             new_state_coeff22,
         ) = bargmann_Abc_to_phasespace_cov_means(A22, b22, c22)
-        assert math.allclose(new_state_cov22[0], state_cov)
-        assert math.allclose(new_state_cov1[0], state_cov)
-        assert math.allclose(new_state_means1[0], state_means)
-        assert math.allclose(new_state_means22[0], state_means)
-        assert math.allclose(new_state_coeff1[0], 1.0)
-        assert math.allclose(new_state_coeff22[0], 1.0)
+        assert math.allclose(new_state_cov22, state_cov)
+        assert math.allclose(new_state_cov1, state_cov)
+        assert math.allclose(new_state_means1, state_means)
+        assert math.allclose(new_state_means22, state_means)
+        assert math.allclose(new_state_coeff1, 1 / (2 * np.pi) ** 2)
+        assert math.allclose(new_state_coeff22, 1 / (2 * np.pi) ** 2)
 
 
 class TestPolyExpAnsatz:

tests/test_physics/test_triples.py

@@ -317,19 +317,19 @@ def test_displacement_gate_s_parametrized_Abc(self):
         A1_correct = np.array([[0, -0.5, -1, 0], [-0.5, 0, 0, 1], [-1, 0, 0, 1], [0, 1, 1, 0]])
         assert math.allclose(A1, A1_correct[[0, 3, 1, 2], :][:, [0, 3, 1, 2]])
         assert math.allclose(b1, math.zeros(4))
-        assert math.allclose(c1, 1)
+        assert math.allclose(c1, 1 / (2 * np.pi))
 
         A2, b2, c2 = triples.displacement_map_s_parametrized_Abc(s=1, n_modes=1)
         A2_correct = np.array([[0, 0, -1, 0], [0, 0, 0, 1], [-1, 0, 0, 1], [0, 1, 1, 0]])
         assert math.allclose(A2, A2_correct[[0, 3, 1, 2], :][:, [0, 3, 1, 2]])
         assert math.allclose(b2, math.zeros(4))
-        assert math.allclose(c2, 1)
+        assert math.allclose(c2, 1 / (2 * np.pi))
 
         A3, b3, c3 = triples.displacement_map_s_parametrized_Abc(s=-1, n_modes=1)
         A3_correct = np.array([[0, -1, -1, 0], [-1, 0, 0, 1], [-1, 0, 0, 1], [0, 1, 1, 0]])
         assert math.allclose(A3, A3_correct[[0, 3, 1, 2], :][:, [0, 3, 1, 2]])
         assert math.allclose(b3, math.zeros(4))
-        assert math.allclose(c3, 1)
+        assert math.allclose(c3, 1 / (2 * np.pi))
 
     @pytest.mark.parametrize("eta", [0.0, 0.1, 0.5, 0.9, 1.0])
     def test_attenuator_kraus_Abc(self, eta):