Calculate Wigner functions and marginals directly from ABC

[JacobHast]

Before posting a feature request

• I have searched exisisting GitHub issues to make sure the feature request does not already exist.

Feature details

Right now Wigner functions and quadrature marginals are calculated in Fock basis (at least in visualize_2d), but we should be able to take advantage of the Gaussian nature of the states in Bargmann representation to get a significant speed-up when calculating these things.

Implementation

No response

How important would you say this feature is?

2: Somewhat important. Needed this quarter.

Additional information

No response

[ziofil]

If you mean the function values, we can do that already, because the Ansatz is callable and you can use it as the Wigner function, quadrature function etc:

from mrmustard.lab_dev import *

sq_wigner = SqueezedVacuum([0], 1.0) >> BtoPS([0], 0)
W = sq_wigner.representation.ansatz  # use this as Wigner function

x = np.linspace(-7, 7, 101)
y = np.linspace(-7, 7, 101)
X,Y = np.meshgrid(x,y)
Z = np.dstack([X+1j*Y, X - 1j*Y])  # for now we need complex coord

from matplotlib import pyplot as plt

plt.matshow(W(Z).real)

mind you this is is rough around the edges and the UI will improve

[JacobHast]

BtoPS = Bargmann to Phase Space?

[JacobHast]

Does your implementation assume ħ = 2? If I change to e.g. ħ = 1, it doesn't seem to give the correct wigner function. Also, I think something is off in the magnitude, i.e. the W in your implementation has max(W) = 1

import mrmustard.lab_dev as mm
import numpy as np
import matplotlib.pyplot as plt
from mrmustard import settings

settings.HBAR = 1

state_squeezed = mm.SqueezedVacuum(modes=[0], r=0)
state_ps = state_squeezed >> mm.BtoPS([0], 0)

x = np.linspace(-3, 3, 101)
y = np.linspace(-3, 3, 101)
X,Y = np.meshgrid(x,y)
Z = np.dstack([X+1j*Y, X - 1j*Y])

W = state_ps.representation.ansatz(Z)

fig, ax = plt.subplots()
pcolor = ax.pcolor(x, y, W.real)
ax.set_aspect('equal')
fig.colorbar(pcolor)

state_squeezed.visualize_2d(return_fig=True, xbounds=(x[0], x[-1]), pbounds=(y[0], y[-1]))

[JacobHast]

Anyway it seems like an easy fix, maybe each representation could have a .wigner(xvec, pvec) and .quadrature(xvec, quadrature_angle) method? With an implementation such that one can call e.g., state.wigner(xvec, pvec) and state.quadrature(xvec, quadrature_angle) for arbitrary states.