%iterative(dx,dt,gridSize,initcond,sources,
%          dirichlet,tolsq,Dconst,solver,omega)

% Alan Davidson
% SciComp final project

% This file will use the Jacobi iterative method to solve the
% time-dependent diffusion equation,
%    u_t=D*(u_xx + u_yy + u_zz) + a(x,y,z)
% where a is a source/sink term
%       u is the temperature of that spot
%       t is time
%       x, y, and z are spacial dimensions (possibly not all are used)
%       D is the diffusion constant
%
% We use a first-degree Backward Euler FDA, so we are trying to solve
%
% u(x,y,z,t+dt)+
%      dt*D/dx^2*((2*dim)*u(x,y,z,t+dt)-sum_neighbors(u(x',y',z',t+dt)))
%      =u(x,y,z,t)+dt*a(x,y,z)

% When we have a Dirichlet boundary condition, note that we instead want
%
% u(x,y,z,t+dt) = u(x,y,z,t)

% We use the Jacobi iterative method, which is the following: let us
% suppose we are trying to solve Ax=b, and that x(0) denotes the initial
% condition of the system. Now, let
%
% x(i+1) = D^-1 * (-L-U)*x(i) + D^-1 * b  or in other words
% x(i+1) = P*x(i) + Q
%
% where D, L, and U are the diagonal, lower triangular, and upper
% triangular parts of A, respectively, and P=D^-1*(-L-U) and Q=D^-1*b.
% We repeat this process until ||x(i+1)-x(i)|| is less than our tolerance,
% in which case we stop, and call x(i+1) a good approximation of the 
% answer. 

% It should be noted that I started this as part of my clinic. Cris Cecka
% did much of the work in getting the clinic code working, which was
% written in C++. I have adapted it myself for Matlab. I have also adapted
% it to work in fewer than 3 dimensions. 

% The function prints out the number of iterations required to converge. 
% It returns the array of chemical concentrations. 

function y=iterative(dx,dt,gridSize,initcond,sources,dirichlet,...
                     tolsq,Dconst,solver,omega)
% dx is the amount by which the grid spaces are separated
% dt is the amount by which we step through time
% gridSize is a vector with 1, 2, or 3 elements, denoting the 
%      size/dimension of the grid
% initcond is the initial condition of the system. The number of elements
%      in initcond must be equal to the product of the elements of size. 
% dirichlet is a matrix of the Dirichlet conditions, with a 
%      negative number for a place without a Dirichlet BC (this is ok,
%      since the concentration of a diffused chemical should never be
%      negative, much less held at that value, so we should never actualy
%      require a Dirichlet BC to actually be negative).
% tolsq is the square of the tolerance to be used at any given timestep
% Dconst is the diffusion constant
% solver is a string - either 'jacobi' or 'ssor' for the different methods
% omega is the scalar in the SSOR (if solver='jacobi' it is ignored)

% temp should be 1
[temp,dim]=size(gridSize);
totalSize=prod(gridSize);

if(temp ~= 1 || dim<1 || dim > 3)
    disp('Too many (or too few?) dimensions.')
    y=0;
    return
end

[numSites,temp]=size(initcond);
if(temp ~= 1 || numSites ~= totalSize)
    disp('Size of initial conditions is wrong.')
    y=0;
    return
end

[numSites,temp]=size(sources);
if(temp ~= 1 || numSites ~= totalSize)
    disp('Size of source/sink terms is wrong.')
    y=0;
    return
end

[numSites,temp]=size(dirichlet);
if(temp ~= 1 || numSites ~= totalSize)
    disp('Size of boundary conditions is wrong.')
    y=0;
    return
end

% Every dimension of the grid should have at least 4 elements in it. 
if(any(gridSize<4))
    disp('Region is too small; quitting');
    % OK, this is only necessary if there are places along the edge
    % of the region that aren't boundary conditions, but I'm going
    % with it anyway. If the region is too small, the FDA on the boundary
    % won't work correctly. 
    y=0;
    return
end

% Where the BC exists, we want elements of x to be identically
% equal to the apposite elements of b. 
A=speye(totalSize);

switch dim
    case 1
        % Where the BC doesn't exist, we want to add 2D*dt/dx^2 to the
        % diagonal term
        [i,j]=find(dirichlet<0);
        A=A+sparse(i,i,2*Dconst*dt/(dx*dx),totalSize,totalSize);
        % We also want the off-diagonal terms to be -D*dt/dx^2 in these
        % cases, but not the ones on the ends!
        
        % We start out by making an array with the ends marked
        edges=zeros(totalSize,1);
        edges(1)=1;
        edges(totalSize)=1;
        % Now, make temp equal dirichlet, except with 1's forced on the
        % ends. This makes temp be -1 at all interior non-BC points
        temp=dirichlet-(edges.*dirichlet)+edges;
        [i,j]=find(temp<0);
        % Now, we can finally make the off-diagonal terms. 
        A=A+sparse(i,i+1,(-Dconst*dt/(dx*dx)),totalSize,totalSize);
        A=A+sparse(i,i-1,(-Dconst*dt/(dx*dx)),totalSize,totalSize);
        
        % Lastly, we need to worry about non-BC sites on the edges. 
        % Here, we want y''(0)=(2y(0)-5y(1)+4y(2)-y(3))/dx^2. 
        % We have already added on the 2y(0)/dx^2 term. 
        if(dirichlet(1)<0)
            A=A+sparse([1;1;1],[2;3;4],[-5;4;-1].*(Dconst*dt/(dx*dx)),...
                totalSize,totalSize);
        end
        if(dirichlet(totalSize)<0)
            A=A+sparse([totalSize;totalSize;totalSize],...
                       [totalSize-1;totalSize-2;totalSize-3],...
                       [-5;4;-1].*(Dconst*dt/(dx*dx)),totalSize,totalSize);
        end
        
    case 2
        % The space (x,y) will be indexed into the array at location
        % (x+gridSize(1)*(y-1)). 
        
        % We will handle the interior and the edges of the area separately:
        indices=1:1:totalSize;
        notTheTop=indices>gridSize(1);
        notTheBottom=indices<=(totalSize-gridSize(1));
        notTheLeft=mod(indices,gridSize(1))~=1;
        notTheRight=mod(indices,gridSize(1))~=0;
        interior=notTheTop.*notTheBottom.*notTheLeft.*notTheRight;
       
        % In places without BCs, we want a total of 1+4Dconst*dt/dx^2 on 
        % the diagonal terms.
        [i,j]=find(dirichlet<0);
        A=A+sparse(i,i,4*Dconst*dt/(dx*dx),totalSize,totalSize);
        % In the interior, we subtract Dconst*dt/dx^2 from the off-diagonal
        % terms if we're not on the BC. 
        [i,j]=find(interior.*(dirichlet<0));
        A=A+sparse(i,i+1,-Dconst*dt/(dx*dx),totalSize,totalSize);
        A=A+sparse(i,i-1,-Dconst*dt/(dx*dx),totalSize,totalSize);
        A=A+sparse(i,i+gridSize(1),-Dconst*dt/(dx*dx),totalSize,totalSize);
        A=A+sparse(i,i-gridSize(1),-Dconst*dt/(dx*dx),totalSize,totalSize);
        % This should handle the top edge for non-BCs:
        [i,j]=find((indices<=gridSize(1)).*(dirichlet<0));
        A=A+sparse(i,i+gridSize(1),-5*Dconst*dt/(dx*dx),...
                   totalSize,totalSize);
        A=A+sparse(i,i+2*gridSize(1),4*Dconst*dt/(dx*dx),...
                   totalSize,totalSize);
        A=A+sparse(i,i+3*gridSize(1),-Dconst*dt/(dx*dx),...
                   totalSize,totalSize);
        % This should handle the bottom edge for non-BCs:
        [i,j]=find((indices>totalSize-gridSize(1)).*(dirichlet<0));
        A=A+sparse(i,i-gridSize(1),-5*Dconst*dt/(dx*dx),...
                   totalSize,totalSize);
        A=A+sparse(i,i-2*gridSize(1),4*Dconst*dt/(dx*dx),...
                   totalSize,totalSize);
        A=A+sparse(i,i-3*gridSize(1),-Dconst*dt/(dx*dx),...
                   totalSize,totalSize);
        % This should handle the left edge for non-BCs:
        [i,j]=find((mod(indices,gridSize(1))==1).*(dirichlet<0));
        A=A+sparse(i,i+1,-5*Dconst*dt/(dx*dx),totalSize,totalSize);
        A=A+sparse(i,i+2,4*Dconst*dt/(dx*dx),totalSize,totalSize);
        A=A+sparse(i,i+3,-Dconst*dt/(dx*dx),totalSize,totalSize);
        % This should handle the right edge for non-BCs:
        [i,j]=find((mod(indices,gridSize(1))==0).*(dirichlet<0));
        A=A+sparse(i,i-1,-5*Dconst*dt/(dx*dx),totalSize,totalSize);
        A=A+sparse(i,i-2,4*Dconst*dt/(dx*dx),totalSize,totalSize);
        A=A+sparse(i,i-3,-Dconst*dt/(dx*dx),totalSize,totalSize);
       
    case 3
        % The space (x,y,z) will be indexed into the array at location
        % (x+gridSize(1)*(y-1)+gridSize(1)*gridSize(2)*(z-1)). 
        
        % We will handle the interior and the edges of the area separately:
        indices=1:1:totalSize;
        notTheTop=indices>gridSize(1)*gridSize(2);
        notTheBottom=indices<=(totalSize-gridSize(1));
        notTheLeft=mod(indices,gridSize(1))~=1;
        notTheRight=mod(indices,gridSize(1))~=0;
        notTheFront=mod(indices,gridSize(1)*gridSize(2))<=gridSize(1);
        notTheFront=mod(indices,gridSize(1)*gridSize(2))>=...
                    gridSize(1)*gridSize(2)-gridSize(1);
        interior=notTheTop.*notTheBottom.*notTheLeft.*notTheRight.*...
                 notTheFront.*notTheBack;
        % In places without BCs, we want a total of 1+6Dconst*dt/dx^2 on 
        % the diagonal terms.
        [i,j]=find(dirichlet<0);
        A=A+sparse(i,i,6*Dconst*dt/(dx*dx),totalSize,totalSize);
        % In the interior, we subtract Dconst*dt/dx^2 from the off-diagonal
        % terms if we're not on the BC. 
        [i,j]=find(interior.*(dirichlet<0));
        A=A+sparse(i,i+1,-Dconst*dt/(dx*dx),totalSize,totalSize);
        A=A+sparse(i,i-1,-Dconst*dt/(dx*dx),totalSize,totalSize);
        A=A+sparse(i,i+gridSize(1),-Dconst*dt/(dx*dx),totalSize,totalSize);
        A=A+sparse(i,i-gridSize(1),-Dconst*dt/(dx*dx),totalSize,totalSize);
        A=A+sparse(i,i+gridSize(1)*gridSize(2),...
                   -Dconst*dt/(dx*dx),totalSize,totalSize);
        A=A+sparse(i,i-gridSize(1)*gridSize(2),...
                   -Dconst*dt/(dx*dx),totalSize,totalSize);
        % This should handle the top face for non-BCs:
        [i,j]=find((indices<=gridSize(1)*gridSize(2)).*(dirichlet<0));
        A=A+sparse(i,i+gridSize(1)*gridSize(2),-5*Dconst*dt/(dx*dx),...
                   totalSize,totalSize);
        A=A+sparse(i,i+2*gridSize(1)*gridSize(2),4*Dconst*dt/(dx*dx),...
                   totalSize,totalSize);
        A=A+sparse(i,i+3*gridSize(1)*gridSize(2),-Dconst*dt/(dx*dx),...
                   totalSize,totalSize);
        % This should handle the bottom face for non-BCs:
        [i,j]=find((indices>totalSize-gridSize(1)*gridSize(2))...
                   .*(dirichlet<0));
        A=A+sparse(i,i-gridSize(1)*gridSize(2),-5*Dconst*dt/(dx*dx),...
                   totalSize,totalSize);
        A=A+sparse(i,i-2*gridSize(1)*gridSize(2),4*Dconst*dt/(dx*dx),...
                   totalSize,totalSize);
        A=A+sparse(i,i-3*gridSize(1)*gridSize(2),-Dconst*dt/(dx*dx),...
                   totalSize,totalSize);
        % This should handle the left face for non-BCs:
        [i,j]=find((mod(indices,gridSize(1))==1).*(dirichlet<0));
        A=A+sparse(i,i+1,-5*Dconst*dt/(dx*dx),totalSize,totalSize);
        A=A+sparse(i,i+2,4*Dconst*dt/(dx*dx),totalSize,totalSize);
        A=A+sparse(i,i+3,-Dconst*dt/(dx*dx),totalSize,totalSize);
        % This should handle the right face for non-BCs:
        [i,j]=find((mod(indices,gridSize(1))==0).*(dirichlet<0));
        A=A+sparse(i,i-1,-5*Dconst*dt/(dx*dx),totalSize,totalSize);
        A=A+sparse(i,i-2,4*Dconst*dt/(dx*dx),totalSize,totalSize);
        A=A+sparse(i,i-3,-Dconst*dt/(dx*dx),totalSize,totalSize);
        % This should handle the front face for non-BCs:
        [i,j]=find((mod(indices,gridSize(1)*gridSize(2))<=gridSize(1))...
                   .*(dirichlet<0));
        A=A+sparse(i,i+gridSize(1),-5*Dconst*dt/(dx*dx),...
                   totalSize,totalSize);
        A=A+sparse(i,i+2*gridSize(1),4*Dconst*dt/(dx*dx),...
                   totalSize,totalSize);
        A=A+sparse(i,i+3*gridSize(1),-Dconst*dt/(dx*dx),...
                   totalSize,totalSize);
       % This should handle the back face for non-BCs:
        [i,j]=find((mod(indices,gridSize(1)*gridSize(2))<=...
                   gridSize(1)*gridSize(2)-gridSize(1))...
                   .*(dirichlet<0));
        A=A+sparse(i,i-gridSize(1),-5*Dconst*dt/(dx*dx),...
                   totalSize,totalSize);
        A=A+sparse(i,i-2*gridSize(1),4*Dconst*dt/(dx*dx),...
                   totalSize,totalSize);
        A=A+sparse(i,i-3*gridSize(1),-Dconst*dt/(dx*dx),...
                   totalSize,totalSize);

    otherwise
        disp('Not yet implemented');
        y=0;
        return
end

% used for debugging purposes
%full(A)
        
% b is the boundary conditions if they exist, and the other
% conditions, otherwise. 
b=(dirichlet>=0).*dirichlet+(dirichlet<0).*(initcond+dt*sources);
        
D=speye(totalSize).*A;
L=tril(A,-1);
U=triu(A,1);

switch lower(solver)
    case 'jacobi'
        P=D^-1*(-L-U);
        Q=D^-1*sparse(b);
           
        if(norm(full(P))>=1)
            disp('System is not stable.')
            norm(full(P))
            y=0;
            return
        end
    case 'ssor'
        E=-L;
        F=-U;
        P=(D-omega*F)^-1*(omega*E+(1-omega)*D)*...
          (D-omega*E)^-1*(omega*F+(1-omega)*D);
        Q=omega*(2-omega)*(D-omega*F)^-1*D*(D-omega*E)^-1*b;
    otherwise 
        disp('Invalid solving method');
        y=0;
        return
end

% Used for debugging purposes
%fullD=full(D)
%fullL=full(L)
%fullU=full(U)
%fullP=full(P)
%fullQ=full(Q)

if(norm(full(P))>=1)
    disp('System is not stable.')
    % It's possible in a few cases that the system is actually stable, but
    % I don't want to take chances on this sort of thing.
    norm(full(P))
    y=0;
    return
end        

% Now, we actually want to compute the new concentrations. 
xold=initcond;
% square of the magnitude of xold-xnew is ressq
ressq = tolsq+1;
iterations=0;

while(ressq>tolsq)
    xnew=P*xold+Q;
    ressq=dot(xnew-xold,xnew-xold);
    xold=xnew;
    iterations=iterations+1;
    
    % For debugging purposes
    %if(mod(iterations,50)==0)
    %    iterations
    %    ressq
    %end
end
% print out the number of iterations we needed. 
iterations=iterations
y=xnew;
return