function nma_math_228b_HW4_parblem_1_part_b
%Lax-Wendroff to solve the wave equation

%nma_math_228b_HW4_parblem_1_part_b  Solves HW4, problem 1, Math 228B
%
%This file contains the functions used to solve problem 1, HW4
%Lax-Wendroff to solve the wave equation
%
%By Nasser M. Abbasi
%3/14/2011

close all;

% initialize data and parameters used

k             = 1.2*10^5;          %bulk modulus for air in Pa
rao           = 1.22521;           %density of air Kg/meter
courantNumber = 0.8;               %courant number for CFL condition
the_speed     = sqrt(k/rao);       %sound speed in this medium
h             = 0.005;             %grid spacing
delt = courantNumber*h/the_speed;  %time step
X             = 0:h:1;             %domain size
max_t         = 0.02;            %maximum run-time in seconds
current_time  = 0;

%set the initial value of pressure and velocity
[u0,Xh] = initial_data(X,@pressure,@speed,rao,k);
clim    = 0.6*[min(u0(1,2:end-1)) max(u0(1,2:end-1))];      %for plotting

if abs(min(u0(1,2:end-1)))<eps
   clim = 0.6*[-0.1*max(u0(1,2:end-1)) max(u0(1,2:end-1))]; %for plotting
end

%center the figure in the screen
%nma_set_figure_position();

%initial plot at t=0
plot_solution(u0,Xh,current_time,0,clim,h,delt)
u = u0;  %allocate space for solution

%calculate how many steps to run and check if need to do partial step
nSteps = max_t/delt;
completeNumberOfSteps = floor(nSteps);
if (abs(nSteps-completeNumberOfSteps)>2*eps)
   partialStep = (nSteps - completeNumberOfSteps)*delt;
else
   partialStep = 0;
end

%allocate the A matrix
A=[0 k;1/rao 0];

% start the main loop, check if doing animation
ANIMATE = false;

if ANIMATE
   f        = getframe(gcf);
   [im,map] = rgb2ind(f.cdata,512,'nodither');
end


for n = 1:completeNumberOfSteps%-1 %works better this way, find out why
   u = makeStep(u,delt,h,A,Xh,k,rao);
   plot_solution(u,Xh,n*delt,n,clim,h,delt);
   
   if ANIMATE     
         f = getframe(gcf);       
         im(:,:,1,n) = rgb2ind(f.cdata,map,'nodither');     
   end
   
end

if partialStep ~=0
   n = n+1;
   u = makeStep(u,partialStep,h,A,Xh,k,rao);
   currentTime = (n-1)*delt + partialStep;
   plot_solution(u,Xh,currentTime,n,clim,h,delt);
   
   if ANIMATE
      f = getframe(gcf);
      im(:,:,1,n) = rgb2ind(f.cdata,map,'nodither');
   end
end

if ANIMATE
   imwrite(im,map,'problem1_part_b.gif','DelayTime',0.0,'LoopCount',Inf);
end

end
%--------------------------
function plot_solution(u,Xh,current_time,step_number,clim,h,delt)
% THis function is called to make a plot of the current solution
% the input u is the vector which contain both pressure and the
% speed of the wave
%

subplot(3,1,1);
A = repmat(u(1,2:end-1),5,1);  %to make a nice pressure diagram
imagesc(A,clim);
colormap(bone);
set(gca,'YTick',[]);
set(gca,'XTick',[]);
set(gca,'FontSize',8);

subplot(3,1,2);
plot(Xh(2:end-1),u(1,2:end-1),'r.-');
title(sprintf('pressure P(x,t), time=%f, step %d',current_time,step_number),...
   'FontSize',8);
xlabel('x (meters)','FontSize',8); ylabel('P(x) Pa units','FontSize',8);
ylim(2*clim);
xlim([0 1]);
set(gca,'FontSize',8);
grid;
drawnow;

%this plot for the initial conditions to allow one to compare
subplot(3,1,3);
if step_number==0
   plot(Xh(2:end-1),u(1,2:end-1),'r.-');
   title(sprintf('P(x,t) at time zero, h=%3.4f, delt=%3.8f',h,delt),'FontSize',8);
   xlabel('x (meters)','FontSize',8); ylabel('P(x) Pa units','FontSize',8);
   ylim(2*clim);
   xlim([0 1]);
   set(gca, 'FontSize', 8)
   grid;
   drawnow;
end

end

%----------------------------------------
function  u = makeStep(u,delt,h,A,Xh,k,rao)
% This function is called to make the solution step
%
% INPUT:
%   u: current solution vector (pressure and speed)
%   delt: time step
%   h: space step
%   A: Matrix A
%   Xh: space grid at half cells
%   k: bulk modulus
%   rao: density
%
% OUTPUT:
%   u: the new solution at next time step.
%

unew = zeros(size(u));

for j=2:length(Xh)-1
   unew(:,j) = u(:,j)- delt/(2*h)*A*( u(:,j+1)-u(:,j-1) )+...
      + delt^2/(2*h^2)*A^2*(u(:,j-1)-2*u(:,j)+u(:,j+1));
end

%set the boundary conditions as specified in problem
unew = assign_boundary_conditions(unew,k,rao);

u = unew;
end
%----------------------------------------
function v=pressure(X)
% THis function is called to initialize pressure for initial conditions
% there are 4 different initial conditions used for simulation.
% adjust the variable TYPE to select which one to use
%

TYPE = 1;
v    = zeros(1,length(X));

if TYPE == 0
   for i = 1:length(X)
      if X(i)>0.4 && X(i)<0.6
         v(i) = sin(10*pi*(X(i)-0.4));
      end
   end
elseif TYPE==1
   for i = 1:length(X)
      if X(i)>0.4 && X(i)<0.6
         v(i) = 1;
      end
   end
elseif TYPE == 2
   for i = 1:length(X)
      if X(i)>0.4 && X(i)<0.6
         v(i) = sin(20*pi*(X(i)-0.4));
      end
   end
elseif TYPE == 3
   v=(X > 0.1).*(X < 0.3) + exp(-200*(X-0.75).^2);
elseif TYPE == 4
   for i = 1:length(X)
      if X(i)>=0.4 && X(i)<=0.5
         v(i)=20*(X(i)-0.4);
      elseif X(i)>0.5 && X(i)<=0.6
         v(i) = 20*(0.6-X(i));
      else
         v(i) = 0;
      end
   end
end
end
%----------------------------------
function v=speed(X)
% Initial speed is set to zero
v=0.*X;
end

%----------------------------
function [U,Xh] = initial_data(X,p,u,rao,k)
%initial_data(X,p,u) function which returns initial data for the problem
%
%INPUT:
%  X: vector which contains the x coordinates of the orginal grid [0..1]
%  p: function handle to evaluate pressure
%  u: function handle to evaluate speed u
%  rao: density
%  k: bulk modulus
%
%OUTPUT:
%  U: a matix of 2 rows and N-1 columns, which represent the
%     value of the solution U at initial time. The top row is
%     pressure p, and second row is the speed u
%  Xh: The cell centered grid x-coordinates, goes from center or left
%      ghost cell to the center of the right most ghost cell.
%

h  = X(2)-X(1);                     %find the spacing h
Xh = X(1)-h/2:h:X(end)+h/2;         %generate cell centered grid
U  = zeros(2,length(Xh));           %allocate storage for U

% evaluate p(x) and u(x) at the cell centered grid points
U(1,2:end-1) = p(Xh(2:end-1));
U(2,2:end-1) = u(Xh(2:end-1));

U = assign_boundary_conditions(U,k,rao);

end
%----------------------------------
function U = assign_boundary_conditions(U,k,rao)
% Called to set the boundary conditions after each step
% 2 different boundary conditions are used to simulation

WALL_ON_BOTH_ENDS = 1;
PROBLEM_BC        = 2;

BC = WALL_ON_BOTH_ENDS ;  %change this to select which BC to use

if BC == PROBLEM_BC
   
   %note U(1,:) is pressure, and U(2,:) is speed
   % assign boundary conditions as given in problem statment
   U(1,1) =  U(1,2);                           %left ghost point updated
   U(2,1) = -U(2,2);                           %left ghost point updated
   
   U(1,end) = U(1,end-1)+U(2,end-1)*sqrt(k*rao);       %right ghost, p
   U(2,end) = U(1,end-1)/sqrt(k*rao) + U(2,end-1);     %right ghost, u
   U(:,end) = (1/2)*U(:,end);
   
else
   
   U(1,1) =  U(1,2);                           %left ghost point updated
   U(2,1) = -U(2,2);                           %left ghost point updated
   
   U(1,end) =  U(1,end-1);
   U(2,end) =  -U(2,end-1);
   
end
end