import matplotlib.pyplot as plt import numpy as np from scipy.special import airy from tqdm import tqdm

def f(x, p, B=1, t=0):

 return (1/(2**(1/3) * np.pi * B)) * airy(2**(2/3) * (B*x + (p**2)/(B**2) - 2*B*p*t))[0]

xmax = 4.5 q = np.linspace(-xmax,xmax, 500) p = np.linspace(-xmax,xmax, 500)

Q, P = np.meshgrid(q, p)

import os movie_name = "airy_wavepacket" dir_path = f"./{movie_name}" tmax = 4 frames_per_t = 24 t_values = np.linspace(-tmax, tmax, frames_per_t * (tmax * 2))

for N, t in tqdm(enumerate(t_values)):

   W = f(Q, P, B=1, t = t)

   plt.gca().set_aspect('equal')
   fig, ax = plt.subplots(figsize=(18, 16))
   wmap = wigner_cmap(W)
   contour = ax.contourf(Q, P, W, 40, cmap=wmap)
   plt.colorbar(contour)

   plt.xlabel('q')
   plt.ylabel('p')
   plt.grid(True)
   if not os.path.exists(dir_path):
       os.makedirs(dir_path)

   plt.savefig(f"{dir_path}/{N:03d}.png",bbox_inches='tight')
   plt.close()

1. !ffmpeg -framerate 24 -i %03d.png output.webm