XanaduAI / MrMustard

A differentiable bridge between phase space and Fock space
https://mrmustard.readthedocs.io/
Apache License 2.0
76 stars 23 forks source link

Calculate Wigner functions and marginals directly from ABC #419

Open JacobHast opened 2 months ago

JacobHast commented 2 months ago

Before posting a feature request

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 commented 2 months ago

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)
Screenshot 2024-06-22 at 6 24 30 PM

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

JacobHast commented 1 month ago

BtoPS = Bargmann to Phase Space?

JacobHast commented 1 month ago

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]))

newplot image

JacobHast commented 1 month ago

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.