% lab5.m - Solving the Navier-Stokes equations in vorticity-stream function
% formulation
%
% Equations:
%
%  Vorticity-transport:
%
%  omega_t + u omega_x + v omega_y = nu (omega_xx + omega_yy)
%
%  Stream function:
%
%  psi_xx+psi_yy = -omega
%
%  Velocities in terms of stream function:
%
%  u = psi_y   v=-psi_x
%
% Boundary conditions: periodic in x and y
%
%
% Problem: Decay of an initial random vorticity distribution
%
clear all;
% Set up problem:
%
nu=1.0e-3; I=sqrt(-1);
nx=32; ny=32; hx=2*pi/nx; hy=2*pi/ny;
x=(0:nx-1)*hx; y=(0:ny-1)*hy; 
omega=zeros([nx ny]); omega_x=omega; omega_y=omega; del_omega=omega;
omega_hat=omega; omega_x_hat=omega; omega_y_hat=omega; 
u=omega; v=omega; u_hat=omega; v_hat=omega; 
conv=omega; conv_hat=omega;
psi_hat=omega;
omega=random('unif',-1,1,nx,ny);

t=0.;
tfinal=10.0;
dt=1.0e-3;

% Form matrices of wavenumbers so we can use matrix operations
kx=zeros([nx ny]); ky=zeros([nx ny]); kmod2=zeros([nx ny]);
for l=1:nx    
   for m=1:ny
       kx(l,m)=I*wavenumber(l,nx);
       ky(l,m)=I*wavenumber(m,ny);
       kmod2(l,m)=kx(l,m)^2+ky(l,m)^2;
       if (kx(l,m)>=2/3*nx) | (ky(l,m)>=2/3*ny)
           alias(l,m)=0;
       else
           alias(l,m)=1;
       end
   end
end
kmod2(1,1)=1;

% Enter time loop

while t<tfinal
    % Compute stream function from current vorticity distribution
    % 1. Take fft of omega
    omega_hat=fft2(omega); psi_hat=zeros([nx ny]);
    % 2. Divide by wavenumber squared to find psi_hat
    psi_hat = -omega_hat ./ kmod2;
    % 3. From psi_hat compute u_hat, v_hat
    u_hat =  ky.*psi_hat;
    v_hat = -kx.*psi_hat;
    % 4. Compute derivatives of omega in Fourier space
    omega_x_hat = kx*omega_hat;
    omega_y_hat = ky*omega_hat;
    % 5. Inverse fft to find u, v, omega_x, omega_y
    u=real(ifft2(u_hat)); v=real(ifft2(v_hat));
    omega_x=real(ifft2(omega_x_hat)); omega_y=real(ifft2(omega_y_hat));
    % 6. Compute convective term in real space
    conv = u.*omega_x + v.*omega_y;    
    % 7. Compute convective term in Fourier space
    conv_hat = fft2(conv);
    % 8. Remove components affected by aliasing
    conv_hat = alias.*conv_hat;
    % 9. Advance omega in time in Fourier space
    omega_hat_new = omega_hat + dt*( nu*kmod2.*omega_hat-conv_hat);
    % 10. Find omega in real space
    omega=real(ifft2(omega_hat_new));
    t=t+dt
    % 11. Plot current vorticity field
    colormap('jet');
    [c,h]=contourf(omega,-2:0.1:2,'w.'); colorbar
    title(num2str(t));
    pause(0.05)    
end