% IMP2D  Script to solve the heat problem
%   u_t = u_xx + u_yy
% on the square [0,1]x[0,1].  Has zero boundary value around the edge
% and has initial condition u(x,y,t=0)=x(1-x)y(1-y).

% ELB 2/26/07

J = 12;
N = 50;
tf = 0.1;

dx = 1/J;  x = 0:dx:1;  y = x;
mu = (tf/N) / (dx*dx);

JJ = (J-1)*(J-1);
A = sparse(JJ,JJ);
U = zeros(JJ,1);  % no need to be sparse

for j=1:J-1
    for k=1:J-1
        % for each point (j,k) in grid, build row of A
        n = (k-1)*(J-1) + j;   % from local coords to global ordering
        A(n,n) = 1 + 4 * mu;
        if j > 1
            A(n,(k-1)*(J-1) + (j-1)) = -mu;
        end
        if j < J-1
            A(n,(k-1)*(J-1) + (j+1)) = -mu;
        end
        if k > 1
            A(n,(k-2)*(J-1) + j) = -mu;
        end
        if k < J-1
            A(n,k*(J-1) + j) = -mu;
        end
        
        % for each point (j,k) in grid, set the initial condition
        U(n) = x(j)*(1-x(j)) * y(k)*(1-y(k));
    end
end

figure
spy(A)

% time-stepping loop
for n = 1:N
    U = A \ U;
end

% plot initial condition and solution using mesh
figure
[xx yy] = meshgrid(x,y);
subplot(2,1,1)
mesh(x,y,xx.*(1-xx).*yy.*(1-yy))
xlabel x, ylabel y, zlabel u
title('initial condition')
subplot(2,1,2)
% go through grid an put U back onto grid
UU = zeros(J+1,J+1);
for j=1:J-1
    for k=1:J-1
        n = (k-1)*(J-1) + j;   % from local coords to global ordering
        UU(j+1,k+1) = U(n);
    end
end
mesh(x,y,UU)
xlabel x, ylabel y, zlabel u
title(['solution at time ' num2str(tf)])