function U = a5prob4(J,tf,f,u0);
% A5PROB4 solves the heat equation
%   u_t = (1+2x) u_xx + 3 u + f(x,t)
% on the interval 0 < x < 1 with boundary conditions
% u_x(0,t)=0 and u(0,t)=0 and initial condition u(x,0) = u0(x).
% Note f and u0 are function handles; see below.
% Example I: to solve A#5 problem 4a
%   >> tf = 1.0,  myf = @(x,t) (0),   myu0 = @(x) 1-x
%   >> U20 = a5prob4(20,tf,myf,myu0);
%   >> U200 = a5prob4(200,tf,myf,myu0);
% Example II: to solve A#5 problem 4c, see script a5prob4c;

% ELB 3/9/07

dx = 1/J;  x = 0:dx:1;
nu = 1/2;  dt = nu * dx;
N = ceil(tf/dt);  dt = tf/N;  % assure N dt = tf
disp(['[J = ' num2str(J) ', N = ' num2str(N) ', dt = ' num2str(dt) ...
   ', mu = dt / dx^2 = ' num2str(dt/(dx*dx)) ']'])

% build A in linear problem to solve at each time step:   A U^{n+1} = b
muj = (dt / (2 * dx * dx)) * (1 + 2 * x);
A = sparse(J,J);
A(1,1:2) = [1 + 2 * muj(1) - (3/2) * dt, -2 * muj(1)]; % insert row 1
for j=2:J
   A(j,j-1) = - muj(j);
   A(j,j) =  1 + 2 * muj(j) - (3/2) * dt;
   if j < J % last row is exception
      A(j,j+1) = - muj(j);
   end
end
% full(A), figure, spy(A) % uncomment to see entries and pattern in A

% fill in initial condition
U = zeros(size(x));  % note U(J+2) = 0, which is right boundary condition
for k = 1:J+1,   U(k) = u0(x(k));   end
figure, plot(x,U,'g--'), hold on

% run
for n = 0:N-1
   tstar = (n + (1/2)) * dt;
   b = zeros(J,1); % build b = RHS for this time step
   b(1) = (1-2*muj(1)+(3/2)*dt) * U(1) + ...
         2 * muj(1) * U(2) + dt * f(x(1),tstar);
   for j = 2:J
      b(j) = muj(j) * U(j-1) + (1-2*muj(j)+(3/2)*dt) * U(j) + ...
         muj(j) * U(j+1) + dt * f(x(j),tstar);
   end
   % size(A), size(b), size(U), return % uncomment to check sizes
   U(1:J) = (A \ b)'; % do time step
   % plot(x,U)  % uncomment to see every step
end

plot(x,U,'r'), hold off
xlabel x, ylabel u, title('solid red is u(x,tf); dashed green is u(x,0)')