% Set-up matrix to solve for convection-diffusion equation

N = (Nx+1)*(Ny+1);
A = zeros(N, N);
r = zeros(N, 1);

ND = 2;
drdX = zeros(N,ND);

% First, discretize the interior points

for ii = 2:Nx,
  for jj = 2:Ny,

    kk0  = ii   + (jj-1)*(Nx+1);
    kkE  = ii+1 + (jj-1)*(Nx+1);
    kkW  = ii-1 + (jj-1)*(Nx+1);
    kkWW = ii-2 + (jj-1)*(Nx+1);
    kkN  = ii   + (jj  )*(Nx+1);
    kkS  = ii   + (jj-2)*(Nx+1);

    yp  = 0.5*(y(ii,jj+1) + y(ii,jj  ));
    ym  = 0.5*(y(ii,jj  ) + y(ii,jj-1));
    y0  = y(ii,jj);

    dyp = y(ii,jj+1) - y(ii,jj);
    dym = y(ii,jj  ) - y(ii,jj-1);
    dy0 = 0.5*(dyp + dym);

    ddyp_dh = dy_dh(ii,jj+1) - dy_dh(ii,jj);
    ddym_dh = dy_dh(ii,jj  ) - dy_dh(ii,jj-1);
    ddy0_dh = 0.5*(ddyp_dh + ddym_dh);
    
    dx = 0.5*(x(ii+1,jj) - x(ii-1,jj));

    Kp = Kinlet(yp,params);
    Km = Kinlet(ym,params);
    K0 = Kinlet(y0,params);
    
    if (ii == 2),
      
      r(kk0) = 0.5*u(ii,jj)/dx*(T(ii+1, jj  ) - T(ii-1, jj  )) ...
	     - K0*(T(ii+1, jj  ) - 2*T(ii  , jj  ) + T(ii-1, jj  ))/dx^2 ...
	     -( Kp*(T(ii  , jj+1) - T(ii  , jj  ))/dyp ...
             -  Km*(T(ii  , jj  ) - T(ii  , jj-1))/dym )/dy0;      
      
      A(kk0, kk0) =                    2*K0/dx^2 + (Kp/dyp+Km/dym)/dy0;
      A(kk0, kkE) =  0.5*u(ii,jj)/dx -   K0/dx^2;
      A(kk0, kkW) = -0.5*u(ii,jj)/dx -   K0/dx^2;
      A(kk0, kkN) =                              -  Kp/dyp/dy0;
      A(kk0, kkS) =                              -         Km/dym/dy0;
      
      % Calculate residuals derivative w.r.t. to design variables
      drdX(kk0,1) = 0.5*du_dUcool(ii,jj)/dx*(T(ii+1, jj  ) - T(ii-1, jj  ));

      drdX(kk0,2) = (  Kp*(T(ii  , jj+1) - T(ii  , jj  ))*ddyp_dh/dyp^2 ...
                     - Km*(T(ii  , jj  ) - T(ii  , jj-1))*ddym_dh/dym^2 )/dy0 ...
         	  + (  Kp*(T(ii  , jj+1) - T(ii  , jj  ))/dyp ...
                     - Km*(T(ii  , jj  ) - T(ii  , jj-1))/dym )*ddy0_dh/dy0^2;      

      
    else

      r(kk0) = u(ii,jj)/dx*(1.5*T(ii, jj  ) - 2*T(ii-1, jj  ) + 0.5*T(ii-2, jj )) ...
	     - K0*(T(ii+1, jj  ) - 2*T(ii  , jj  ) + T(ii-1, jj  ))/dx^2 ...
	     -( Kp*(T(ii  , jj+1) - T(ii  , jj  ))/dyp ...
             -  Km*(T(ii  , jj  ) - T(ii  , jj-1))/dym )/dy0;

      A(kk0, kk0)  =  1.5*u(ii,jj)/dx + 2*K0/dx^2 + (Kp/dyp+Km/dym)/dy0;
      A(kk0, kkE)  =                  -   K0/dx^2;
      A(kk0, kkW)  = -2.0*u(ii,jj)/dx -   K0/dx^2;
      A(kk0, kkWW) =  0.5*u(ii,jj)/dx;
      A(kk0, kkN)  =                              -  Kp/dyp/dy0;
      A(kk0, kkS)  =                              -         Km/dym/dy0;

      % Calculate residuals derivative w.r.t. to design variables
      drdX(kk0,1) = du_dUcool(ii,jj)/dx*(1.5*T(ii, jj  ) - 2*T(ii-1, jj  ) + 0.5*T(ii-2, jj ));
      
      drdX(kk0,2) = (  Kp*(T(ii  , jj+1) - T(ii  , jj  ))*ddyp_dh/dyp^2 ...
                     - Km*(T(ii  , jj  ) - T(ii  , jj-1))*ddym_dh/dym^2 )/dy0 ...
         	  + (  Kp*(T(ii  , jj+1) - T(ii  , jj  ))/dyp ...
                     - Km*(T(ii  , jj  ) - T(ii  , jj-1))/dym )*ddy0_dh/dy0^2;      
      
    end

  end
end

% Set-up Dirichlet conditions at inlet

ii = 1;
for jj = 1:Ny+1,
 
  kk0 = ii   + (jj-1)*(Nx+1);

  r(kk0     ) = T(ii,jj) - Tinlet(y(ii,jj),params);
  A(kk0, kk0) = 1.0;

end

% Set-up Dirichlet conditions at bottom

jj = 1;
for ii = 2:Nx+1,

  kk0 = ii   + (jj-1)*(Nx+1);

  r(kk0     ) = T(ii,jj) - Tinlet(y(ii,jj),params);
  A(kk0, kk0) = 1.0;

end

% Set-up Dirichlet conditions at top boundary

jj = Ny+1;
for ii = 2:Nx+1,
 
  kk0 = ii   + (jj-1)*(Nx+1);

  r(kk0     ) = T(ii,jj) - Tinlet(y(ii,jj),params);
  A(kk0, kk0) = 1.0;

end

% Set-up parabolized boundary conditions at outlet

ii = Nx+1;
for jj = 2:Ny,
 
  kk0  = ii   + (jj-1)*(Nx+1);
  kkW  = ii-1 + (jj-1)*(Nx+1);
  kkWW = ii-2 + (jj-1)*(Nx+1);
  kkN  = ii   + (jj  )*(Nx+1);
  kkS  = ii   + (jj-2)*(Nx+1);

  yp  = 0.5*(y(ii,jj+1) + y(ii,jj  ));
  ym  = 0.5*(y(ii,jj  ) + y(ii,jj-1));
  y0  = y(ii,jj);
  
  dyp = y(ii,jj+1) - y(ii,jj);
  dym = y(ii,jj  ) - y(ii,jj-1);
  dy0 = 0.5*(dyp + dym);
       
  ddyp_dh = dy_dh(ii,jj+1) - dy_dh(ii,jj);
  ddym_dh = dy_dh(ii,jj  ) - dy_dh(ii,jj-1);
  ddy0_dh = 0.5*(ddyp_dh + ddym_dh);
  
  dx = x(ii,jj) - x(ii-1,jj);

  Kp = Kinlet(yp,params);
  Km = Kinlet(ym,params);
  K0 = Kinlet(y0,params);
  
  r(kk0) = u(ii,jj)/dx*(1.5*T(ii, jj  ) - 2*T(ii-1, jj  ) + 0.5*T(ii-2, jj )) ...
         -( Kp*(T(ii  , jj+1) - T(ii  , jj  ))/dyp ...
         -  Km*(T(ii  , jj  ) - T(ii  , jj-1))/dym )/dy0;

  A(kk0, kk0)  =  1.5*u(ii,jj)/dx + (Kp/dyp+Km/dym)/dy0;
  A(kk0, kkW)  = -2.0*u(ii,jj)/dx;
  A(kk0, kkWW) =  0.5*u(ii,jj)/dx;
  A(kk0, kkN)  =                  -  Kp/dyp/dy0;
  A(kk0, kkS)  =                  -         Km/dym/dy0;

  % Calculate residuals derivative w.r.t. to design variables
  drdX(kk0,1) = du_dUcool(ii,jj)/dx*(1.5*T(ii, jj  ) - 2*T(ii-1, jj  ) + 0.5*T(ii-2, jj ));

  drdX(kk0,2) = (  Kp*(T(ii  , jj+1) - T(ii  , jj  ))*ddyp_dh/dyp^2 ...
                 - Km*(T(ii  , jj  ) - T(ii  , jj-1))*ddym_dh/dym^2 )/dy0 ...
              + (  Kp*(T(ii  , jj+1) - T(ii  , jj  ))/dyp ...
                 - Km*(T(ii  , jj  ) - T(ii  , jj-1))/dym )*ddy0_dh/dy0^2;      

end