# Finite Difference Solver for a fully-developed flow inside a rectangular
# pipe of dimensions Lx and Ly (UNSTEADY VERSION).
# A theta-scheme for time discretization leads to the following system:
#
#  (1/dt*Am + theta*Af).W(n+1) = (1/dt*Am - (1-theta)*Af).W(n) + Bp(n+theta)
#
# W(x,y) is prescribed to zero (Dirichlet boundary conditions) at the four
# walls of the cross section: ((0,y),(Lx,y),(x,0),(x,Ly))
#
# Author: Gustavo Buscaglia
# Last Modification: 08/31/2015

#BEGIN PROBLEM DEFINITION
#---------------------------------------------------------------------------
#Convention used
# N2   x  x  x  ... x 
# .    .
# .    .
# .    . 
# 3    x  x  x  ... x  
# 2    x  x  x  ... x  
# 1    x  x  x  ... x  
#
# j/i  1  2  3  ... N1

# Global index if(N1 < N2) n = ij2n(i,j) = i + (j-1)*N1
# Global index if(N1 > N2) n = ij2n(i,j) = j + (i-1)*N2

global N1 = 100;  #X Direction
global N2 = 50;   #Y Direction
L1 = 2.0;
L2 = 1.0;
rho = 1.;
visc = 1.;
#uniform grid spacing
dx = L1/(N1-1);
dy = L2/(N2-1);
#uniform viscosity
mu = visc;
#temporal integration
theta = 0.5;
dt = .05;
T0 = 0.;
NTS = 40;
#plotting
IP = 2;                 #timesteps increment for plots
NP = NTS+1;      #number of stored timesteps
#working space
nunk = N1*N2;
fprintf(stdout, "  * PROBLEM DEFINITION DONE -> Unknowns = %d\n", nunk);
fflush(stdout);
#matrices and arrays
Af = zeros(nunk,nunk);   # flux matrix
Am = zeros(nunk,nunk);   # mass matrix
bm = zeros(nunk,1);      # mass RHS
#arrays for diagnostic purposes
W0 = zeros(nunk,1);      # initial condition (code it!)
WN = zeros(nunk,NP/IP);  # stored solutions
TN = zeros(NP/IP);       # times for stored solutions
diag = zeros(10,NP);     # diagnostics
#---------------------------------------------------------------------------
#END PROBLEM DEFINITION

#BEGIN PRE-PROCESSING STEP
#---------------------------------------------------------------------------
#-- Assembly: loop over nodes
for i=1:N1
   for j=1:N2
	if (i==1 || i==N1 || j==1 || j==N2)
		continue;
   	else
# viscous matrix
	     pP=ij2n(i,j); pN=ij2n(i,j+1); pE=ij2n(i+1,j); 
             pS=ij2n(i,j-1); pW=ij2n(i-1,j);
             aux1 = mu/dx^2; aux2 = mu/dy^2;
             Af(pP,pP) = 2*(aux1+aux2); 
             Af(pP,pN)=-aux2; Af(pP,pS)=-aux2; 
             Af(pP,pE)=-aux1; Af(pP,pW)=-aux1;
# mass matrix
             Am(pP,pP)=rho/dt;
	     bm(pP)=dx*dy;
        endif
   endfor
endfor
#-- Timestepping Matrices: M, R
M = Am + theta*Af;
R = Am - (1-theta)*Af;
#-- Correct M for no-slip boundary conditions
vones=ones(nunk,1);
for i=1:N1
   for j=1:N2
	if (i==1 || i==N1 || j==1 || j==N2)
	   pP=ij2n(i,j); 
           M(pP,pP)=1;
           vones(pP)=0;
	endif
   endfor
endfor

#-- Switch to sparse storage ?
flag_sparse = 1;
fprintf(stdout, "  * CREATING FINAL MATRICES -> flag_sparse = %d\n",
        flag_sparse);
fflush(stdout);

if(flag_sparse == 1) 
  M = sparse(M);
  R = sparse(R);
endif
#---------------------------------------------------------------------------
#END PRE-PROCESSING STEP

#BEGIN SOLUTION STEP
#---------------------------------------------------------------------------
#-- Temporal integration loop
nt = 1;
np = 0;
tn = T0;
wn = W0;
anorm = 0;
fprintf(stdout, "     SOLVING ...\n");
exit_flag = 0;
while 1
   #-- Diagnostics

   #store 'some' solutions for visualization
   if rem(nt-1,IP) == 0
      np=np+1;
      TN(np)=tn;
      WN(:,np)=wn;
   endif

   #time(1)
   diag(1,nt)=tn;

   #kinetic energy(2)
   diag(2,nt)=0.5*rho*(bm.*wn)'*wn;

   #flow rate(3)
   diag(3,nt)=bm'*wn;

   #screen output
   wnorm=norm(wn);
   fprintf(stdout, "t= %12.5E |w|= %12.5E |a|= %12.5E\n", tn, wnorm, anorm);
   fflush(stdout);

   #should exit ?
   if exit_flag > 0
     break;
   endif 

   #-- Advance in time
   tm=tn+dt;

   #assemble RHS
   aux=-(theta*dpdz(tm)+(1.-theta)*dpdz(tn));
   b=R*wn + aux*vones;

   #solve linear system
   wm = M\b;

   #acceleration
   am = (wm.-wn)/dt;
   anorm=norm(am);

   #end of integration ?
   if nt >= NTS
     exit_flag = 1;
   endif

   #end timestep
   wn = wm;
   tn = tm;
   nt = nt + 1;
endwhile
#---------------------------------------------------------------------------
#END SOLUTION STEP

#BEGIN POST-PROCESSING STEP
#---------------------------------------------------------------------------
#-- Integrals
plot(diag(1,:),diag(2,:))
title("Kinetic Energy")
pause;
figure()
plot(diag(1,:),diag(3,:))
title("Flow Rate")
pause;
figure()

#-- Meshgrid: 
X=zeros(N1,1);
X(1)=0;
for i=2:N1
  X(i)=X(i-1)+dx;
endfor
Y=zeros(N2,1);
Y(1)=0;
for j=2:N2
  Y(j)=Y(j-1)+dy;
endfor

#-- Limits
xmin=min(X);
xmax=max(X);
ymin=min(Y);
ymax=max(Y);
wmin=min(WN(:,np));
wmax=max(WN(:,np));
for i=1:np
  wmin = min(wmin,min(WN(:,i)));
  wmax = max(wmax,max(WN(:,i)));
endfor

#-- Plot Solution 
for i=1:np
  for ii=1:N1
   for jj=1:N2
    ip=ij2n(ii,jj);
    Wmat(ii,jj)=WN(ip,i);
   endfor
  endfor
#  surf(X,Y,Wmat');
#  axis([xmin xmax ymin ymax wmin wmax]);
  title(sprintf("T=%g",TN(i)))
  #figure()
  contourf(X,Y,Wmat');
  pause();
endfor
#---------------------------------------------------------------------------
#END POST-PROCESSING STEP