% Set all constants
Lx = 1;
Ly = 1;
hbar = 1;
m = 1;
kx = @(nx) nx*pi/Lx;
ky = @(ny) ny*pi/Ly;
w = @(n1,n2) hbar*(kx(n1)^2+ky(n2)^2)/2*m;
A = sqrt(4/Lx*Ly);
[x,y] = meshgrid(0:Lx/100:Lx, 0:Ly/100:Ly);
t = linspace(0,10,500);

% Define the position and time functions
psi = @(x,y,nx,ny) A.*sin(kx(nx)*x).*sin(ky(ny)*y);
phi = @(t,n1,n2) exp(-1i.*w(n1,n2).*t);


% Make the figure larger to allow for subplots
scrsz = get(groot,'ScreenSize');
figure('Position',[1 scrsz(4) scrsz(3) scrsz(4)]);

% Make the animation
filename = '2DParticleALL.gif';
for j = 1:4:length(t)
    subplot(3,2,1)
    surf(x,y,real(psi(x,y,1,1)*phi(t(j),1,1)))
    xlim([0 Lx]), ylim([0 Ly]), zlim([-A A]);
    title('$(n_x,n_y) = (1,1)$','interpreter','latex')
    
    subplot(3,2,2)
    surf(x,y,real(psi(x,y,2,3)*phi(t(j),2,3)))
    xlim([0 Lx]), ylim([0 Ly]), zlim([-A A]);
    title('$(n_x,n_y) = (2,3)$','interpreter','latex')
    
    subplot(3,2,3)
    surf(x,y,real(psi(x,y,3,2)*phi(t(j),3,2)))
    xlim([0 Lx]), ylim([0 Ly]), zlim([-A A]);
    title('$(n_x,n_y) = (3,2)$','interpreter','latex')
    
    subplot(3,2,4)
    surf(x,y,real(psi(x,y,4,4)*phi(t(j),4,4)))
    xlim([0 Lx]), ylim([0 Ly]), zlim([-A A]);
    title('$(n_x,n_y) = (4,4)$','interpreter','latex')
    
    subplot(3,2,5)
    surf(x,y,real(psi(x,y,1,5)*phi(t(j),1,5)))
    xlim([0 Lx]), ylim([0 Ly]), zlim([-A A]);
    title('$(n_x,n_y) = (1,5)$','interpreter','latex')
    
    subplot(3,2,6)
    surf(x,y,real(psi(x,y,3,6)*phi(t(j),3,6)))
    xlim([0 Lx]), ylim([0 Ly]), zlim([-A A]);
    title('$(n_x,n_y) = (3,6)$','interpreter','latex')
    
    drawnow
    frame = getframe(1);
    im = frame2im(frame);
    [Q,map] = rgb2ind(im,256);
    if j == 1;
        imwrite(Q,map,filename,'gif','LoopCount',Inf,'DelayTime',0);
    else
        imwrite(Q,map,filename,'gif','WriteMode','append','DelayTime',0);
    end
end