%% PS2: Fluid flow in a channel

function ps2()


    channelflow(21, 3.0, 0.5, 1, true);
    channelflow(41, 3.0, 0, 1, true); 
    channelflow(41, 3.0, 0.5, 1, true);
    channelflow(41, 3.0, 1, 1, true);    
    
    %% check for convergence of both flowrates and velocity in L2 norm,
    %% using a 81x81 point discretization as reference
    refgridN = 81;
    [refgrideta, refgridxi] = ndgrid(0:1/(refgridN-1):1, ...
                                     0:1/(refgridN-1):1);
    gridsizes = [11, 21, 41];
    bvals = [0 0.5 1];
    
    for b=bvals
        a = sqrt((3-b)^2/4 - 1);
        jacobian = b/2 + a*refgrideta;
        [refQ, refu] = channelflow(refgridN, 3, b, 1, false);
        for idx=1:length(gridsizes)
            dxi(idx) = 1/(gridsizes(idx) - 1);
            [qtemp, utemp] = channelflow(gridsizes(idx), 3, b, 1, false);
            Ql2diff(idx) = abs(qtemp - refQ);
            [grideta, gridxi] = ndgrid(0:dxi(idx):1, 0:dxi(idx):1);
            udiff = refu - interp2(gridxi, grideta, utemp, refgridxi, ...
                                   refgrideta);
            ul2diff(idx) = sqrt(numint((udiff.^2).*jacobian));
        end
        figure();
        loglog(dxi, ul2diff, dxi, ul2diff(end)/dxi(end)^2 * dxi.^2);
        legend('||u_{ref} - u||_2', 'C d\xi^2');
        xlabel('d\xi = d\eta');
        title(sprintf('L_2 error in u vs grid size, b=%g', b));
        figure();
        loglog(dxi, Ql2diff, dxi, Ql2diff(end)/dxi(end)^2 * dxi.^2);
        legend('|Q_{ref} - Q|', 'C d\xi^2');
        xlabel('d\xi = d\eta');
        title(sprintf('L_2 error in Q vs grid size, b=%g', b));
    end
    
end

% Solves Poisson equation for fluidflow in a channel with irregular
% cross-section by transforming to the unit square and solving there.
% Calculates total flowrate Q and velocity u

function [Q, u] = channelflow(N, l, b, h, qshow)
    
    a = sqrt((l-b)^2/4 -h^2);
    u = solvesys(N, a, b, h);
    
    [eta, xi] = ndgrid(0:1/(N-1):1, 0:1/(N-1):1);
    x = xi * b/2 + a*xi.*eta;
    y = eta*h;
       
    jacobian = h*(b/2 + a*eta);
    Q = numint(u.*jacobian);

    if qshow
        figure();
        patch([0, b/2, b/2 + a, 0], [0,0,h,h], -ones(1,4), 0, ...
              'facecolor', [.8,.8,.8]);
        hold on;
        [cc,hh]=contour(x,y,u,0.02:0.02:max(u(:)) );
        clabel(cc, hh);
        axis equal;
        title(sprintf(['Contour diagram for the flow when ' ...
                       'l=%g, b=%g, h=%g, N=%d'], l, b, h, N));
        text(b/2+a/2+.1,1/2, sprintf('Q = %f', Q))
    end
end

% Given the channel parameters, solves the transformed Poisson equation
% with transformed BCs; returns the system matrix and the solution in the
% transformed space.
function u = solvesys(N, a, b, h)

    dxi = 1/(N-1);
    A = spalloc(N^2, N^2, 7*N^2);
    f = zeros(N^2,1);
    
    %% setup the zero boundary conditions on borders eta = 0 & xi = 1
    self = gridtovect(N, 1:N, N);
    A( tensorind(self, self, N) ) = 1;
    f(self) = 0;
    
    self = gridtovect(1:N, 1, N);
    A( tensorind(self, self, N) ) = 1;
       
    %% setup the left boundary conditions (xi = 0)
    %% du/dxi = 0
    A = forwardFDd1x(A, 1, 2:N, dxi, 1, N);
    
    %% setup the top boundary conditions (on the border eta = 1)
    %% du/deta + c(xi)*du/dxi = 0
    
    % du/deta
    A = backwardFDd1y(A, 2:N-1, N, dxi, 1, N);
   
    % c(xi)*du/dxi
    xi = dxi:dxi:1-dxi;
    c = -(2*a*xi)/(b+2*a);
    A = centerFDd1x(A, 2:N-1, N, dxi, c, N);
    
    %% setup the interior PDE
    %% u_{xi, xi} c1(xi, eta) + u_{xi, eta} c2(xi, eta) 
    %%    + u_{xi} c3(xi,eta) + u_{eta, eta} = -h^2
    xi = dxi:dxi:1-dxi;
    for grideta = 2:N-1
        eta = (grideta -1)*dxi;
        c1 = 4*(h^2 + a^2*xi.^2) / (b + 2*a*eta)^2;
        c2 = -(4*a*xi) / (b + 2*a*eta);
        c3 = 8*a^2*xi / (b + 2*a*eta)^2;
        
        % u_{xi, xi} c1(xi, eta)
        A = centerFDd2x(A, 2:N-1, grideta, dxi, c1, N);
        % u_{xi, eta} c2(xi, eta)
        A = asymFDdxy(A, 2:N-1, grideta, dxi, c2, N);
        % u_{xi} c3(xi, eta)
        A = centerFDd1x(A, 2:N-1, grideta, dxi, c3, N);
        % u_{eta, eta}
        A = centerFDd2y(A, 2:N-1, grideta, dxi, 1, N);
        
        self = gridtovect(2:N-1, grideta, N);
        f(self) = -h^2;
    end
    
    for row=1:N^2
        f(row) = f(row)/max(abs(A(row,:)));
        A(row, :) = A(row, :)/max(abs(A(row,:)));
    end
    
    u = A\f;
    u = reshape(u, N, N);
end

%% UTILITY FUNCTIONS -- worth reusing in the future

% given a system matrix A, and x,y pairs in column vectors X,Y, modifies
% the entries in A to add a 2nd order accurate centered finite difference
% term to approximate the x-deriv, centered at each x,y pair; accepts a
% vector field c to multiply each FD by, and the step size dx
function A = centerFDd1x(A, X, Y, dx, c, N)
    self = gridtovect(X, Y, N);
    lft = gridtovect(X-1, Y, N);
    rt = gridtovect(X+1, Y, N);
    
    A( tensorind(self, lft, N) ) = ...
        A( tensorind(self, lft, N) ) + c * -1/(2*dx);
    A( tensorind(self, rt, N) ) = ...
        A( tensorind(self, rt, N) ) + c * 1/(2*dx);
end

% given a system matrix A, and x,y pairs in column vectors X,Y, modifies
% the entries in A to add a 2nd order accurate centered finite difference
% term to approximate the 2nd x-deriv, centered at each x,y pair; accepts a
% vector field c to multiply each FD by, and the step size dx
function A = centerFDd2x(A, X, Y, dx, c, N)
    self = gridtovect(X, Y, N);
    lft = gridtovect(X-1, Y, N);
    rt = gridtovect(X+1, Y, N);
    
    A( tensorind(self, self, N) ) = ...
        A( tensorind(self, self, N) ) + c * -2/(dx^2);
    A( tensorind(self, lft, N) ) = ...
        A( tensorind(self, lft, N) ) + c * 1/dx^2;
    A( tensorind(self, rt, N) ) = ...
        A( tensorind(self, rt, N) ) + c * 1/dx^2;
end

% given a system matrix A, and x,y pairs in column vectors X,Y, modifies
% the entries in A to add a 2nd order accurate centered finite difference
% term to approximate the 2nd y-deriv, centered at each x,y pair; accepts a
% vector field c to multiply each FD by, and the step size dy
function A = centerFDd2y(A, X, Y, dy, c, N)
    self = gridtovect(X, Y, N);
    bot = gridtovect(X, Y-1, N);
    top = gridtovect(X, Y+1, N);
    
    A( tensorind(self, self, N) ) = ...
        A( tensorind(self, self, N) ) + c * -2/(dy^2);
    A( tensorind(self, top, N) ) = ...
        A( tensorind(self, top, N) ) + c * 1/dy^2;
    A( tensorind(self, bot, N) ) = ...
        A( tensorind(self, bot, N) ) + c * 1/dy^2;
end

% given a system matrix A, and x,y pairs in column vectors X,Y, modifies
% the entries in A to add a 2nd order accurate forward finite difference
% term to approximate the x-deriv, centered at each x,y pair; accepts a
% vector field c to multiply each FD by, and the step size dx
function A = forwardFDd1x(A, X, Y, dx, c, N)
    self = gridtovect(X, Y, N);
    onert = gridtovect(X+1, Y, N);
    twort = gridtovect(X+2, Y, N);
    
    A( tensorind(self, self, N) ) = ...
        A( tensorind(self, self, N) ) + c * -3/(2*dx);
    A( tensorind(self, onert, N) ) = ...
        A( tensorind(self, onert, N) ) + c * 4/(2*dx);
    A( tensorind(self, twort, N) ) = ...
        A( tensorind(self, twort, N) ) + c * -1/(2*dx);
end

% given a system matrix A, and x,y pairs in column vectors X,Y, modifies
% the entries in A to add a 2nd order accurate backward finite difference
% term to approximate the y-deriv, centered at each x,y pair; accepts a
% vector field c to multiply each FD by, and the step size dy
function A = backwardFDd1y(A, X, Y, dy, c, N)
    self = gridtovect(X, Y, N);
    onebot = gridtovect(X, Y-1, N);
    twobot = gridtovect(X, Y-2, N);
    
    A( tensorind(self, self, N) ) = ...
        A( tensorind(self, self, N) ) + c * 3/(2*dy);
    A( tensorind(self, onebot, N) ) = ...
        A( tensorind(self, onebot, N) ) + c * -4/(2*dy);
    A( tensorind(self, twobot, N) ) = ...
        A( tensorind(self, twobot, N) ) + c * 1/(2*dy);
end

% given a system matrix A, and x,y pairs in column vectors X,Y, modifies
% the entries in A to add a 2nd order accurate finite difference
% term to approximate the xy-deriv, centered at each x,y pair; accepts a
% vector field c to multiply each FD by, and the grid size h
function A = asymFDdxy(A, X, Y, h, c, N)
    self = gridtovect(X, Y, N);
    bot = gridtovect(X, Y-1, N);
    top = gridtovect(X, Y+1, N);
    lft = gridtovect(X-1, Y, N);
    rt = gridtovect(X+1, Y, N);
    lftbot = gridtovect(X-1, Y-1, N);
    rttop = gridtovect(X+1, Y+1, N);
    
    A( tensorind(self, self, N) ) = ...
        A( tensorind(self, self, N) ) + c * 2/(2*h^2);
    A( tensorind(self, bot, N) ) = ...
        A( tensorind(self, bot, N) ) + c * -1/(2*h^2);
    A( tensorind(self, top, N) ) = ... 
        A( tensorind(self, top, N) ) + c * -1/(2*h^2);
    A( tensorind(self, lft, N) ) = ...
        A( tensorind(self, lft, N) ) + c * -1/(2*h^2);
    A( tensorind(self, rt, N) ) = ...
        A( tensorind(self, rt, N) ) + c * -1/(2*h^2);
    A( tensorind(self, lftbot, N) ) = ...
        A( tensorind(self, lftbot, N) ) + c * 1/(2*h^2);
    A( tensorind(self, rttop, N) ) = ...
        A( tensorind(self, rttop, N) ) + c * 1/(2*h^2);
end

% given the number of grid points, N, and x coords I, y coords J, returns
% corresponding linear indices into the long vector corresponding to the
% grid.
function indices = gridtovect(I, J, N)
    if length(I) == 1
        I = ones(size(J))*I;
    elseif length(J) == 1
        J = ones(size(I))*J;
    end
    indices = J + (I-1)*N;
end

% given the number of grid points N, and I, J x and y coordinate vectors,
% returns the linear indices corresponding to (x,y)
function indices = tensorind(I, J, N)
    indices = sub2ind([N^2,N^2], I, J);
end

% numerically integrate a uniform-grid function defined on the unit square;
% currently uses the trapezoid rule for 2D domains: weight 1 for interior
% points, 1/2 for boundary points, 1/4 for corner pts.
function A = numint(gridfun)
    A = sum(sum(gridfun(2:end-1, 2:end-1))) + ...
        1/2*( sum(gridfun(1, 2:end-1)) + sum(gridfun(end, 2:end-1)) + ...
              sum(gridfun(2:end-1, 1)) + sum(gridfun(2:end-1, end)) ) + ...
        1/4*( gridfun(1,1) + gridfun(1,end) + gridfun(end,1) + ...
              gridfun(end,end));
     dx = 1/(size(gridfun,1)-1);
     A = dx^2 * A;
end