function U=diffusion2(J,aLEFT,aDOWN,f,g);
% DIFFUSION2 Solve an equilibrium diffusion equation
%    div(a grad u) + f = (a u_x)_x + (a u_y)_y + f = 0,
% where diffusion a=a(x,y) and source f=f(x,y) depend on location.  Uses
% Dirichlet condition u=g on boundary of domain is (0,1) x (0,1).  Uses 
% inline function or function handle for g.  
% However, diffusion  a  is input by two matrices aLEFT,aDOWN which are the 
% values of a(x,y) on the staggered grid:
%    aLEFT_{r,s} = a(x_r-dx/2,y_s),   1 <= r <= J, 1 <= s <= J-1,
%    aDOWN_{r,s} = a(x_r,y_s-dy/2),   1 <= r <= J-1, 1 <= s <= J,
% for   x_r=r*dx  and  y_s=s*dy  (where 0 <= r,s <= J) and dx=dy=1/J.
% Similarly, f is input by a (J-1) x (J-1) matrix giving the values of 
% f(x,y) at the interior points:
%    f_{r,s} = f(x_r,y_s),   1 <= r,s <= J-1.
% Note a(x,y) must be everywhere positive.
%   See Morton & Mayers section 6.2 for the method used here.  Note this 
% code is designed to be compatible with the solution procedure for some 
% nonlinear equations.  For example see MINIMAL.
%
% Form:
%      U = diffusion2(J,aa,f,g);
% where 
%        U = approx soln as (J+1)x(J+1) matrix; includes bdry values;
%                 U_{r,s} approximates u(x_r,y_s)
%        J = number of grid spaces (in both x and y)
%    aLEFT = J x (J-1) matrix with aLEFT_{r,s} = a(x_r-dx/2,y_s) as above
%    aDOWN = (J-1) x J matrix with aDOWN_{r,s} = a(x_r,y_s-dy/2) as above
%        f = (J-1) x (J-1) matrix giving f(x_r,y_s) at interior pts
%        g = g(x,y) = inline function or fcn handle; value of u on bdry
%
% Example I:  For linear example, including comparison to manufactured 
%   solution see USEDIFF2.
% Example II:  For nonlinear, iterative example, see MINIMAL.
%
% ELB 4/24/05 update of DIFFUSION to make indexing match discussion in
%   class  
% See also: USEDIFF2, POTENTIAL2, MINIMAL

if any(any(aLEFT<=0))|any(any(aDOWN<=0))
    error('diffusivity a(x,y) must be positive'), end
dx=1/J;  dy=dx;  dx2=dx*dx;

% fill in A, b; solve
N=(J-1)^2;   A=sparse(N,N);  b=zeros(N,1);
for r=1:J-1
    for s=1:J-1
        k = kk(J,r,s);  x=r*dx;  y=s*dy;
        A(k,k) = -( aLEFT(r,s)+aLEFT(r+1,s)+aDOWN(r,s)+aDOWN(r,s+1) );
        b(k) = - f(r,s) * dx2;
        if r==J-1, b(k)=b(k) - aLEFT(J,s)*feval(g,1,y); % right neighbor
        else, A(k,kk(J,r+1,s))=aLEFT(r+1,s); end  % otherwise, add to A
        if r==1,   b(k)=b(k) - aLEFT(1,s)*feval(g,0,y); % left
        else, A(k,kk(J,r-1,s))=aLEFT(r,s); end
        if s==J-1, b(k)=b(k) - aDOWN(r,J)*feval(g,x,1); % upper
        else, A(k,kk(J,r,s+1))=aDOWN(r,s+1); end
        if s==1,   b(k)=b(k) - aDOWN(r,1)*feval(g,x,0); % lower
        else, A(k,kk(J,r,s-1))=aDOWN(r,s); end
    end
end
w=A\b;    % compute soln w (in column vector form)

% put result into grid, with boundary values
U=zeros(J+1,J+1);
for r=1:J-1, for s=1:J-1,  U(r+1,s+1)=w(kk(J,r,s));  end, end  % interior
% lower and upper, the left and right boundaries
for r=0:J, U(r+1,1)=feval(g,r*dx,0); U(r+1,J+1)=feval(g,r*dx,1); end  
for s=1:J-1, U(1,s+1)=feval(g,0,s*dy); U(J+1,s+1)=feval(g,1,s*dy); end

% plot approx. soln; see "help mesh" on axis orientation
xx=0:dx:1;  yy=xx;  mesh(xx,yy,U')
xlabel x,  ylabel y,  hidden off,  title('approximate solution'),  drawnow

function k=kk(J,r,s)  % go from grid (r,s) (for U) to vector index k
k=(s-1)*(J-1)+r;