%
%  Solve Navier-Stokes problem in driven cavity
%  nx, ny = grid size
%  m = number of time steps
%  T = final time
%  Re = Reynolds number
%
% Example:
%   ns(41,41,3000,35,1000)
%
function ns(nx,ny,m,T,Re)
ome=zeros(nx,ny);
psi=zeros(size(ome));
u = zeros(size(ome));
v = zeros(size(ome));
u(:,ny)=1.0;

clf;
dt = T/m;
hx=1.0/(nx-1);
hy=1.0/(ny-1);
hx2=hx*hx; hy2=hy*hy;
dtmax = 0.5*Re^2*min(hx2,hy2);
if dt > dtmax
    fprintf(1,'time step too long dt=%f dtmax=%f\n',dt,dtmax);
    return
end
c1=dt/(Re*hx2);
c2=dt/(Re*hy2);
c3=dt/(2*hx);
c4=dt/(2*hy);
xx = (0:nx-1)/(nx-1);
yy = (0:ny-1)/(ny-1);
[X,Y]=ndgrid(xx,yy);

for k=1:m
    t = (k-1)*dt;
    vx = 1 - exp(-10*t);
    %
    % boundary values of ome:
    %
    ome(1:nx,1)  = -2*psi(1:nx,2)/hy2;
    ome(1:nx,ny) = -2*(psi(1:nx,ny-1)+hy*vx)/hy2;
    ome(1,1:ny-1)  = -2*psi(2,1:ny-1)/hx2;
    ome(nx,1:ny-1) = -2*psi(nx-1,1:ny-1)/hx2;
    %
    ome_new = ome;
    %
    % update internal nodes of ome:
    %
    ome_new(2:nx-1,2:ny-1) = ome(2:nx-1,2:ny-1) + ...
       - c3*u(2:nx-1,2:ny-1).*(ome(3:nx,2:ny-1)-ome(1:nx-2,2:ny-1)) ...
       - c4*v(2:nx-1,2:ny-1).*(ome(2:nx-1,3:ny)-ome(2:nx-1,1:ny-2)) ...
       + c1*(ome(3:nx,2:ny-1)-2*ome(2:nx-1,2:ny-1)+ome(1:nx-2,2:ny-1)) ...
       + c2*(ome(2:nx-1,3:ny)-2*ome(2:nx-1,2:ny-1)+ome(2:nx-1,1:ny-2));
    %
    ome = ome_new;
    %
    % Solve  -Laplace(psi) = ome     
    %
    psi(2:nx-1,2:ny-1) = poisson_solver(nx-2,ny-2,ome(2:nx-1,2:ny-1));
    %
    % Compute velocity components from stream function psi:
    %
    u(2:nx-1,2:ny-1) = (psi(2:nx-1,3:ny)-psi(2:nx-1,1:ny-2))/(2*hy);
    v(2:nx-1,2:ny-1) = -(psi(3:nx,2:ny-1)-psi(1:nx-2,2:ny-1))/(2*hx);
    %
    if mod(k,nx)==0 % mod(k,m/50)==0      
        subplot(1,3,1)
        i1=round(0.75*nx); i2=nx;
        j1=1; j2=round(0.25*ny);
        quiver(X(i1:i2,j1:j2),Y(i1:i2,j1:j2),u(i1:i2,j1:j2),v(i1:i2,j1:j2),5); 
        axis([0.75 1 0 .25]);  
        title('detail of lower right corner');
        subplot(1,3,2)
        clevel=linspace(-0.1,0,20); clevel=[clevel 1e-5 1e-4 5e-4 1e-3];
        contour(X,Y,psi,clevel);
        title('streamlines');
        subplot(1,3,3)
        clevel=linspace(0,1,20);
        unorm=sqrt(u.*u+v.*v);
        [cplot,handu]=contourf(X,Y,unorm,clevel); 
        set(handu,'EdgeColor','none');
        axis([0 1 0 1]); axis off
        caxis([min(clevel) max(clevel)]); %colorbar;
        title('velocity');
        pause(0.01)
    end
end


%
% Solve  -Delta(u)=f with homog. Dirichlet conditions
%
function u = poisson_solver(nx,ny,f)
h1=1/(nx+1);
h2=1/(ny+1);
ff = zeros(size(f));
v  = zeros(size(f));
u  = zeros(size(f));
for i=1:nx
   ff(i,:) = dst(f(i,:));
end
tri=zeros(3,nx);
for i=1:ny
    tri(2,:) = (2/h1^2)+2/h2^2-(2/h2^2)*cos(i*pi/(ny+1));
    tri(1,:) = -1/h1^2;
    tri(3,:) = -1/h1^2;
    vv = tridiag_solve(tri(1,:),tri(2,:),tri(3,:),ff(:,i));
    v(:,i) = vv';
end
for i=1:nx
    u(i,:) = idst(v(i,:));
end
