Wigner function for states (#581)

**Context:**
Our current implementation of ``BtoPS`` is in the complex domain, and
has extra factors one should be careful about. This unconventional
description makes things less intuitive.

**Description of the Change:**
This PR creates a wigner ansatz for states that can be called by using
ordinary phase space (x,p) conventions.

**Benefits:**
Making Wigner function calculations more intuitive and simple.

**Possible Drawbacks:**
Could be unstable for ``ArraysAnsatz``.

**Related GitHub Issues:**
None.

---------

Co-authored-by: Anthony <[email protected]>
Co-authored-by: Filippo Miatto <[email protected]>

Author: 3 people

mrmustard/lab/states/base.py

@@ -43,7 +43,7 @@
 )
 
 from ..circuit_components import CircuitComponent
-from ..circuit_components_utils import BtoChar, BtoQ
+from ..circuit_components_utils import BtoChar, BtoQ, BtoPS
 from ..transformations import Transformation
 
 __all__ = ["State"]
@@ -758,3 +758,25 @@ def visualize_dm(
         if return_fig:
             return fig
         display(fig)
+
+    @property
+    def wigner(self):
+        r"""
+        Returns the Wigner function of this state in phase space as an ``Ansatz``.
+
+        .. code-block::
+
+            >>> import numpy as np
+            >>> from mrmustard.lab import Ket
+
+            >>> state = Ket.random([0])
+            >>> x = np.linspace(-5, 5, 100)
+
+            >>> assert np.all(state.wigner(x,0).real >= 0)
+        """
+        if isinstance(self.ansatz, PolyExpAnsatz):
+            return (self >> BtoPS(self.modes, s=0)).ansatz.PS
+        else:
+            raise ValueError(
+                "Wigner ansatz not implemented for Fock states. Consider calling ``.to_bargmann()`` first."
+            )

mrmustard/physics/ansatz/polyexp_ansatz.py

@@ -44,6 +44,7 @@
     join_Abc,
 )
 from mrmustard.physics.utils import generate_batch_str, verify_batch_triple
+from mrmustard.physics.fock_utils import c_in_PS
 
 from mrmustard import math, widgets, settings
 from mrmustard.math.parameters import Variable
@@ -208,6 +209,47 @@ def num_derived_vars(self) -> int:
     def num_vars(self):
         return self.A.shape[-1]
 
+    @property
+    def PS(self) -> PolyExpAnsatz:
+        r"""
+        The ansatz defined using real (i.e., phase-space) variables.
+        """
+        n = self.A.shape[-1]
+        if n % 2:
+            raise ValueError(
+                f"A phase space ansatz must have even number of indices. (n={n} is odd)"
+            )
+
+        if self.num_derived_vars == 0:
+            W = math.conj(math.rotmat(n // 2)) / math.sqrt(settings.HBAR, dtype=math.complex128)
+
+            A = math.einsum("ji, ...jk, kl-> ...il", W, self.A, W)
+            b = math.einsum("ij, ...j-> ...i", W, self.b)
+            c = self.c / (2 * settings.HBAR) ** (n // 2)
+            return PolyExpAnsatz(A, b, c, lin_sup=self._lin_sup)
+
+        else:
+            if self.num_derived_vars != 2:
+                raise ValueError(
+                    f"This transformation supports 2 core and 0 or 2 derived variables"
+                )
+            A_tmp = self.A
+
+            A_tmp = A_tmp[..., [0, 2, 1, 3], :][..., [0, 2, 1, 3]]
+            b = self.b[..., [0, 2, 1, 3]]
+            c = c_in_PS(self.c)  # implements PS transformations on ``c``
+
+            W = math.conj(math.rotmat(n // 2)) / math.sqrt(settings.HBAR, dtype=math.complex128)
+
+            A = math.einsum("ji, ...jk, kl-> ...il", W, A_tmp, W)
+            b = math.einsum("ij, ...j-> ...i", W, b)
+            c = c / (2 * settings.HBAR)
+
+            A_final = A[..., [0, 2, 1, 3], :][..., :, [0, 2, 1, 3]]
+            b_final = b[..., [0, 2, 1, 3]]
+
+            return PolyExpAnsatz(A_final, b_final, c, lin_sup=self._lin_sup)
+
     @property
     def scalar(self) -> Scalar:
         r"""
@@ -863,10 +905,10 @@ def __call__(self: PolyExpAnsatz, *z_inputs: ArrayLike | None) -> Batch[ComplexT
                - *b are the batch dimensions of the combined inputs.
                - *L is the batch shape of the ansatz.
         """
-        z_only = [arr for arr in z_inputs if arr is not None]
+        z_only = [math.cast(arr, dtype=math.complex128) for arr in z_inputs if arr is not None]
         broadcasted_z = math.broadcast_arrays(*z_only)
         z = (
-            math.cast(math.stack(broadcasted_z, axis=-1), dtype=math.complex128)
+            math.stack(broadcasted_z, axis=-1)
             if broadcasted_z
             else math.astensor([], dtype=math.complex128)
         )

mrmustard/physics/fock_utils.py

@@ -24,6 +24,7 @@
 from typing import Sequence, Iterable
 
 import numpy as np
+from scipy.special import comb, factorial
 
 from mrmustard import math, settings
 from mrmustard.math.lattice import strategies
@@ -512,3 +513,46 @@ def vjp(dLdGc):
         return math.astensor(dr, dtype=r.dtype), math.astensor(dphi, phi.dtype)
 
     return ret, vjp
+
+
+def c_ps_matrix(m, n, alpha):
+    """
+    helper function for ``c_in_PS``.
+    """
+    mu_range = range(max(0, alpha - n), min(m, alpha) + 1)
+    tmp = [comb(m, mu) * comb(n, alpha - mu) * (1j) ** (m - n - 2 * mu + alpha) for mu in mu_range]
+    return np.sum(tmp)
+
+
+def gamma_matrix(c):
+    """
+    helper function for ``c_in_PS`.
+    constructs the matrix transformation that helps transforming ``c``.
+    ``c`` here must be 2-dimensional.
+    """
+    M = c.shape[0] + c.shape[1] - 1
+    Gamma = np.zeros((M**2, c.shape[0] * c.shape[1]), dtype=np.complex128)
+
+    for m in range(c.shape[0]):
+        for n in range(c.shape[1]):
+            for alpha in range(m + n + 1):
+                factor = math.sqrt(
+                    factorial(m) * factorial(n) / (factorial(alpha) * factorial(m + n - alpha))
+                )
+                value = c_ps_matrix(m, n, alpha) * math.sqrt(settings.HBAR / 2) ** (m + n)
+                row = alpha * M + (m + n - alpha)
+                col = m * c.shape[0] + n
+                Gamma[row, col] = value / factor
+    return Gamma
+
+
+def c_in_PS(c):
+    """
+    Transforms the ``c`` matrix of a ``DM`` object from bargmann to phase-space.
+    It is a helper function used in
+
+    Args:
+        c (Tensor): the 2-dimensional ``c`` matrix of the ``DM`` object
+    """
+    M = c.shape[0] + c.shape[1] - 1
+    return np.reshape(gamma_matrix(c) @ np.reshape(c, (c.shape[0] * c.shape[1], 1)), (M, M))

tests/test_lab/test_states/test_dm.py

@@ -612,3 +612,12 @@ def test_stellar_decomposition_mixed(self):
 
         assert math.allclose(core.dm().contract(phi, mode="zip").ansatz.A, sigma.ansatz.A)
         assert math.allclose(core.dm().contract(phi, mode="zip").ansatz.b, sigma.ansatz.b)
+
+    def test_wigner(self):
+
+        ans = Vacuum(0).dm().wigner
+        x = np.linspace(0, 1, 100)
+
+        solution = np.exp(-(x**2)) / np.pi
+
+        assert math.allclose(ans(x, 0), solution)

tests/test_lab/test_states/test_ket.py

@@ -29,6 +29,7 @@
     CircuitComponent,
     Coherent,
     Dgate,
+    Ggate,
     Identity,
     Ket,
     Number,
@@ -39,6 +40,7 @@
 )
 from mrmustard.physics.representations import Representation
 from mrmustard.physics.triples import coherent_state_Abc
+from mrmustard.physics.wigner import wigner_discretized
 from mrmustard.physics.wires import Wires
 from mrmustard.widgets import state as state_widget
 
@@ -618,3 +620,20 @@ def test_formal_stellar_decomposition(self):
         core, U = sigma.formal_stellar_decomposition([0])
 
         assert sigma == core.contract(U, mode="zip")
+
+    def test_wigner(self):
+
+        ans = Vacuum(0).wigner
+        x = np.linspace(0, 1, 100)
+        solution = np.exp(-(x**2)) / np.pi
+
+        assert math.allclose(ans(x, 0), solution)
+
+    @pytest.mark.parametrize("n", [1, 2, 3])
+    def test_wigner_poly_exp(self, n):
+
+        psi = (Number(0, n).dm().to_bargmann()) >> Ggate(0)
+        xs = np.linspace(-5, 5, 100)
+        poly_exp_wig = math.real(psi.wigner(xs, 0))
+        wig = wigner_discretized(psi.fock_array(), xs, 0)
+        assert math.allclose(poly_exp_wig[:, None], wig[0], atol=3e-3)

tests/test_physics/test_ansatz/test_polyexp_ansatz.py

@@ -29,6 +29,7 @@
     complex_gaussian_integral_1,
     complex_gaussian_integral_2,
 )
+from mrmustard.lab.transformations import Identity
 
 from ...random import Abc_triple
 
@@ -624,3 +625,11 @@ def test_and_with_lin_sup_other(self):
         assert ansatz_and.A.shape == (2, 5, 3, 4, 3, 3)
         assert ansatz_and.b.shape == (2, 5, 3, 4, 3)
         assert ansatz_and.c.shape == (2, 5, 3, 4)
+
+    def test_PS(self):
+        ans = Identity(0).ansatz
+
+        x = np.linspace(0, 1, 10)
+        gaussian = np.exp((x**2) / 2) / 2
+
+        assert math.allclose(ans.PS(x, 0), gaussian)