% Last modified: 30/12/2012 by Pete

%-------------------------------------------------------------------------%
% Pete's modifications to Matt's original code indicated as such.
%-------------------------------------------------------------------------%

function loglikes = Pete_StochSolver( app )

% function [P u v rho NodeP Nodeu Nodev] =
% NavierStokesInc_DO_rho(app,SetupScript,PlotScript)
% Solves the Boussinesq equations with coriolis forces using a simple Finite 
% Volume Incremental Pressure Correction in Rotational form projection method.
% This function uses a first-order implicit Euler integration scheme, and
% uses LU factorization to solve the matrices.The advection uses the quick 
% scheme when the mean is involved, Central difference when only the modes are
% invovled. Advection routines are mexed.
%
%INPUTS:
%       app:	    The application data structured. At a minimum it should 
%		     contain the following fields (which may be added through
%                    the SetupScript).
%          app.Nx:         Number of interior nodes in x-direction 
%                           (integer number)
%          app.Ny:         Number of interior nodes in y-direction 
%                           (integer number)
%          app.t:          Final time
%          dt:         Timestep size
%          nu:         Viscosity
%          app.kappa:      Diffusivity of density (optional)
%          S:          If length(S)==1 Size of the stochastic
%                           dimension. 
%                          If length(S)==2, then we have adaptive 
%                           number of stochastic dimensions:
%             S:       Minimum number of stochastic dimensions
%             S(2):       Maximum number of stochastic dimensions
%          app.var           Variance of largest mode at initialization
%                             (Optional) for DO solvers
%          app.addvar        If smallest variance mode has this variance,
%                              a new mode is added for adaptive DO solvers    
%                             (Optional) 
%          app.subbvar       If smallest variance mode has this variance, &
%                            second smallest variance is less than twice
%                            this variance, then the smallest variance mode
%                            is subtraced for adaptive DO solvers 
%                            (Optional)
%          MC:         Number of monte-carlo samples for stochastic 
%                           coefficients (Default = 10,000)
%          app.DOAdapt:    Number of time integration intervals between 
%                          checking if the number of stochastic modes 
%                          should be adjusted for adaptive stochastic
%                          dimension case (length(S)==2)
%                          Default value is 10.
%          app.PlotIntrvl: Number of time integration intervals between 
%                          calling "plotscript."
%       PlotScript: (String input) of filename used to plot specific problem.
%                   Can also be used to save outputs and perform other tasks.
%                       Note: Variables u, v, and P, will be available for
%                       plotting.
%       SetupScript: String input of filename used to setup the problem.
%           This file needs to correctly create the following variables:
%	    Nt:     Number of timesteps (integer number)
%           dx:     Size of control volume in x-direction
%           dy:     Size of control volume in y-direction
%           bcsDrho:Density Dirichlet boundary conditions (optional)
%           bcsDP:  Pressure Dirichlet boundary conditions
%           bcsDu:  U-velocity Dirichlet boundary conditions
%           bcsDv:  V-velocity Dirichlet boundary conditions
%           bcsNrho:Density Neumann boundary conditions (optional)
%           bcsNP:  Pressure Neumann boundary conditions
%           bcsNu:  U-velocity Neumann boundary conditions
%           bcsNv:  V-velocity Neumann boundary conditions
%           NodeP:  Pressure node-numbering matrix (see NodePad.m)
%           Nodeu:  U-velocity node-numbering matrix (see NodePad.m)
%           Nodev:  V-velocity node-numbering matrix (see NodePad.m)
%           idsP:   Interior Pressure node id's (see NodePad.m)
%           idsu:   Interior U-velocity node id's (see NodePad.m)
%           idsv:   Interior V-velocity node id's (see NodePad.m)
%           Fu(time,p):   Forcing function for U-velocity
%                         p =[ XU(idsu),YU(idsu) ]
%           Fv(time,p):   Forcing function for V-velocity
%                         p =[ XV(idsv),YV(idsv) ]
%           XU:     Matrix of x-coordinates corresponding to U-velocity CV
%                    centers. XU(idsu) used as 'x' input to Fu. (see meshgrid)
%           XV:     Matrix of x-coordinates corresponding to V-velocity CV
%                    centers. YU(idsu) used as 'y' input to Fv. (see meshgrid)
%           YU:     Matrix of y-coordinates corresponding to U-velocity CV
%                    centers. XV(idsv) used as 'x' input to Fu. (see meshgrid)
%           YV:     Matrix of y-coordinates corresponding to V-velocity CV
%                    centers. YV(idsv) used as 'y' input to Fv. (see meshgrid)
%           P:      Initial Pressure vector. The first #Nbc terms contain
%                   the boundary condition values.
%           u:      Initial U-velocity vector. The first #Nbc terms contain
%                   the boundary condition values.
%           v:      Initial V-velocity vector. The first #Nbc terms contain
%                   the boundary condition values.
%
%OUTPUTS:
%        P:      Final Pressure vector, including the boundary
%                 conditions.
%        u:      Final U-velocity vector, including the boundary
%                 conditions.
%        v:      Final V-velocity vector, including the boundary
%                 conditions.
%        rho:    Final density vector, including the boundary
%                 conditions.
%        NodeP:  Pressure node-numbering matrix (useful for plotting)
%        Nodeu:  U-velocity node-numbering matrix (useful for plotting)
%        Nodev:  V-velocity node-numbering matrix (useful for plotting)
%
%   Uses Matlab functions: symamd, lu
%   See also: mk_DiffusionOperator, FixDiffuvDbcs, Grad2D, Div2D
%
% Authors:   Matt Ueckermann, Themis Sapsis,
%            Pierre Lermusiaux, Pat Haley for MIT course 2.29

%-------------------------------------------------------------------------%
% find trunk directory
trunk = regexp(pwd,'.*trunk','match','once') ;

% define scripts
SetupScript  = 'Pete_stochsetup' ;
SaveScript   = 'Pete_stochsave'  ;
ObsScript = 'Pete_obs' ;
%-------------------------------------------------------------------------%

%Run setup script
eval(SetupScript)

if ~isfield(app,'IP'),app.IP=[1,1,1];end
if ~exist('bcsDrho','var'),bcsDrho=zeros(size(bcsDP));end
if ~exist('bcsNrho','var'),bcsNrho=zeros(size(bcsNP));end
if ~exist('NodeRho','var'),NodeRho=NodeP;end
if ~exist('rho','var'),rho=zeros(size(P));end
if ~exist('rhoi','var'),rhoi=zeros(size(P,1),S);end
if ~isfield(app,'kappa'),app.kappa=0;end
if ~isfield(app,'PrRa')
    app.PrRa=1;
end
if ~exist('bcsOP','var'),bcsOP=[];end
if ~exist('bcsOu','var'),bcsOu=[];end
if ~exist('bcsOv','var'),bcsOv=[];end
if ~isfield(app,'S'),    S=0;end
if ~exist('YYt','var'), YYt=0;end
if ~isfield(app,'DOAdapt')
    if length(S)==2
        app.DOAdapt=10;
    else
        app.DOAdapt=0;
    end
end
%Options related to the Coriolis forces
if ~isfield(app,'Fcor'), app.Fcor = @(t, p) zeros(size(p, 1), 1);end

if ~isfield(app,'DO_cor_iter'), app.DO_cor_iter=1; end

%Find number of boundary conditions
Nbcs = length(bcsDP) + length(bcsNP);

%Check to see how many of the modes have boundary conditions
tmp = (sum(abs(ui(1:Nbcs,:)) + abs(vi(1:Nbcs,:)) +...
    abs(Pi(1:Nbcs,:)) + abs(rhoi(1:Nbcs,:))) > sqrt(eps));
Sb = length(find(tmp));
%arrange so that the boundary modes are always first
[tmp, tmp] = sort(tmp, 'descend');
ui = ui(:, tmp);
vi = vi(:, tmp);
Pi = Pi(:, tmp);
rhoi = rhoi(:, tmp);
if Sb > 0
   if ~exist('rhsYibcs', 'var')
       rhsYibcs = @(yi,t) zeros(size(yi));
   end
end

%Set tolerances etc. for GMRES solvers
TOL=sqrt(eps);
MAXIT=20;
RESTARTS=10;

% Low storage Runge-Kutta coefficients
rk4a = [            0.0 ...
        -567301805773.0/1357537059087.0 ...
        -2404267990393.0/2016746695238.0 ...
        -3550918686646.0/2091501179385.0  ...
        -1275806237668.0/842570457699.0];
rk4b = [ 1432997174477.0/9575080441755.0 ...
         5161836677717.0/13612068292357.0 ...
         1720146321549.0/2090206949498.0  ...
         3134564353537.0/4481467310338.0  ...
         2277821191437.0/14882151754819.0];
rk4c = [             0.0  ...
         1432997174477.0/9575080441755.0 ...
         2526269341429.0/6820363962896.0 ...
         2006345519317.0/3224310063776.0 ...
         2802321613138.0/2924317926251.0];
     
%% Constructing A and RHS
[Arho rhsrho] =  mk_DiffusionOperator(Nx, Ny, dx, dy,...
    bcsDrho,bcsNrho, NodeRho,'P');
[AP rhsPbc] =  mk_DiffusionOperator(Nx, Ny, dx, dy,...
    bcsDP,bcsNP, NodeP,'P');
[Au rhsubc] =  mk_DiffusionOperator(Nx, Ny, dx, dy,...
    bcsDu, bcsNu, Nodeu,'u');
[Av rhsvbc] =  mk_DiffusionOperator(Nx, Ny, dx, dy, ...
    bcsDv, bcsNv, Nodev,'v');
[AP rhsPbc] = FixDiffPObcs(AP, bcsOP, Nbcs,...
    rhsPbc, NodeP, dx, dy); bcsDP(bcsOP) = [];
[Au rhsubc] = FixDiffPObcs(Au, bcsOu, Nbcs,...
    rhsubc, Nodeu, dx, dy); bcsDu(bcsOu) = [];
[Av rhsvbc] = FixDiffPObcs(Av, bcsOv, Nbcs,...
    rhsvbc, Nodev, dx, dy); bcsDv(bcsOv) = [];
[Au Av rhsubc rhsvbc] = FixDiffuvDbcs(Au, Av, bcsDu, bcsDv,...
    rhsubc, rhsvbc, Nodeu, Nodev, dx, dy, Nbcs);

if (sum(abs(app.Fcor(0, [XV(idsv) YV(idsv)])))>0)
    [Au2v Av2u] = CoriolisAVG(Nx, Ny, Nodeu, Nodev,...
        idsu, idsv, Nbcs);
    Au2v = spdiags(app.Fcor(0, [XV(idsv) YV(idsv)]), 0, ...
        length(idsv), length(idsv)) * Au2v;
    Av2u = spdiags(app.Fcor(0, [XU(idsu) YU(idsu)]), 0, ...
        length(idsu), length(idsu)) * Av2u;
else
    Au2v = spalloc(length(v) - Nbcs,length(u), 0);
    Av2u = spalloc(length(u)-Nbcs,length(v),0);
end

AU = -nu * Au;
AV = -nu * Av;
ARHO = -app.kappa * Arho;

Au = nu * Au + 1 / dt * speye(size(Au));
Av = nu * Av + 1 / dt * speye(size(Av));
Arho = app.kappa * Arho + 1 / dt * speye(size(Arho));

%% LU factorization of matrices
imposemean = false;
if isempty(bcsDP)
    imposemean = true;
end
%First order matrix to give best spartsity pattern
rhoperm = symamd(Arho);
Pperm = symamd(AP);
uperm = symamd(Au);
vperm = symamd(Av);

AP=AP(Pperm, Pperm);

% % LU factorization
[Lrho Urho] = lu(Arho(rhoperm, rhoperm));
[LP UP] = lu(AP);
[Lu Uu] = lu(Au(uperm, uperm));
[Lv Uv] = lu(Av(vperm, vperm));

% Get interior ids
pid = Nbcs + 1:length(P);
uid = Nbcs + 1:length(u);
vid = Nbcs + 1:length(v);
srcid     = 2:Nx + 1;
srcTopid  = 2:Ny;
srcBotid  = 3:Ny + 1;
Neu = [length(bcsDu) + 1,Nbcs];
Nev = [length(bcsDv) + 1,Nbcs];
NbcsDrho = length(bcsDrho); %Needed for advection
NbcsDu = length(bcsDu); %Needed for advection
NbcsDv = length(bcsDv); %Needed for advection

%% Initialize variables and sizes
frho = zeros(length(pid),1);
fu = zeros(length(uid),1);
fv = zeros(length(vid),1);

if S>0 || length(S)>1
    frhoi = zeros(length(pid), S);
    fui = zeros(length(uid), S);
    fvi = zeros(length(vid), S);
    frhoij = zeros(length(pid), S^2);
    fuij= zeros(length(uid), S^2);
    fvij= zeros(length(vid), S^2);
else
    frhoi = 0;
    frhoij = 0;
    fui = 0;
    fvi = 0;
    fuij= 0;
    fvij= 0;
end

%If stochastic modes were not initialized in Setup script, 
%initialize automatically.
DOInit;
if ~exist('rhoi', 'var') && (S>0 || length(S)>1)
   rhoi = zeros(length(rho), S); 
end

%% Solving system
q = zeros(size(P));
CYY = cov(YYt);

if app.DOAdapt > 0
    DOAdapt_mex;
end
drawnow;

%-------------------------------------------------------------------------%
t = 0 ; k = 0 ; runtime = toc(app.toptic) ;
eval(SaveScript);
%-------------------------------------------------------------------------%

for k = 1:Nt
    t=dt*k;
    
    %Need to save ui00 etc to do the projections correctly
    ui00 = ui;
    vi00 = vi;
    rhoi00 = rhoi;
    
    %Calculate statistics
    MY = mean(YYt);
    for i=1:S
        YYt(:,i) = YYt(:,i) - MY(i);
    end
    CYY = cov(YYt);
    CYYINV = pinv(CYY, max(CYY(:)) * sqrt(eps));    
    MYYY = mom_thrd(YYt);
    if S>0
        MYYY = reshape(MYYY,S * S, S);
    end
    if S ==0
       CYY = 0; 
    end
    
    %%%%Solve mean field
    
    %Calculate Pressure gradient
    [dpdx dpdy]   = Grad2D(P, NodeP, dx, dy, idsu, idsv);
    
    % Calculate RHS vectors for mean u, v, and rho
    [fu fv] = UVadvect_TVD_mex(Nx, Ny, dx, dy, dt, ...
        u, v, u, v, Nodeu, Nodev, idsu, idsv, Nbcs, NbcsDu, NbcsDv);
    fu = fu + Av2u * v;
    src = app.PrRa * 0.5*( rho(NodeRho(srcTopid, srcid))...
        + rho(NodeRho(srcBotid, srcid)) );
    src = src(idsv);
    fv = fv - Au2v * u - src;
    frho = SCAadvect_TVD_mex(Nx, Ny, dx, dy, dt, rho, u, v,...
        NodeRho, Nodeu, Nodev, idsP, Nbcs, NbcsDrho, NbcsDu, NbcsDv);
        
    %calculate fui and fvi forcings, as well as fuij, and fvij
    for i=1:S
        [fui(:,i) fvi(:,i)] = UVadvect_TVD_mex(Nx, Ny, dx, dy,...
            dt, ui(:,i), vi(:,i),u      , v      , ...
            Nodeu, Nodev, idsu, idsv, Nbcs, NbcsDu, NbcsDv);
        [tmp tmp1]          = UVadvect_TVD_mex(Nx, Ny, dx, dy,...
            dt, u      , v      ,ui(:,i), vi(:,i),... 
            Nodeu, Nodev, idsu, idsv, Nbcs, NbcsDu, NbcsDv);         
        [tmp2 tmp3]          = UVadvect_TVD_mex(Nx, Ny, dx, dy, ...
            dt, u      , v      ,-ui(:,i), -vi(:,i),... 
            Nodeu, Nodev, idsu, idsv, Nbcs, NbcsDu, NbcsDv);         

        src = app.PrRa * 0.5*( rhoi(NodeRho(srcTopid,srcid),i) +...
            rhoi(NodeRho(srcBotid,srcid),i));
        src = src(idsv);
        fui(:,i) = fui(:,i) + 0.5 * (tmp - tmp2)  + Av2u * vi(:,i);
        fvi(:,i) = fvi(:,i) + 0.5 * (tmp1 - tmp3) - Au2v * ui(:,i) - src;

        frhoi(:,i) = SCAadvect_TVD_mex(Nx, Ny, dx, dy, dt,...
            rhoi(:,i), u, v, NodeRho, Nodeu, Nodev, idsP, Nbcs, NbcsDrho,...
            NbcsDu,NbcsDv)...
            + 0.5 * SCAadvect_TVD_mex(Nx, Ny, dx, dy, dt,...
            rho, ui(:,i), vi(:,i), NodeRho, Nodeu, Nodev, idsP, Nbcs, ...
            NbcsDrho,NbcsDu,NbcsDv) ...
             - 0.5 * SCAadvect_TVD_mex(Nx, Ny, dx, dy, dt,...
             rho, -ui(:,i), -vi(:,i), NodeRho, Nodeu, Nodev, idsP, Nbcs,...
             NbcsDrho,NbcsDu,NbcsDv);
        
         %For stochastic boundary forcing, need to change above two lines
        for j=1:S
            [fuij(:, i + (j - 1) * S) fvij(:, i + (j - 1) * S)] = ...
                UVadvect_TVD_mex(Nx, Ny, dx, dy, dt,...
                ui(:,i), vi(:,i), ui(:,j), vi(:,j), ...
                Nodeu, Nodev, idsu, idsv, Nbcs, NbcsDu, NbcsDv);
            [tmp tmp1] = ...
                UVadvect_TVD_mex(Nx, Ny, dx, dy, dt,...
                ui(:,i), vi(:,i), -ui(:,j), -vi(:,j), ...
                Nodeu, Nodev, idsu, idsv, Nbcs, NbcsDu, NbcsDv);
            frhoij(:,i+(j-1)*S) = ...
                0.5 * SCAadvect_TVD_mex(Nx, Ny, dx, dy, dt,...
                rhoi(:,i), ui(:,j), vi(:,j), NodeRho, Nodeu, Nodev, idsP,...
                Nbcs, NbcsDrho, NbcsDu, NbcsDv)...
                - 0.5 * SCAadvect_TVD_mex(Nx, Ny, dx, dy, dt,...
                rhoi(:,i), -ui(:,j), -vi(:,j), NodeRho, Nodeu, Nodev, idsP,...
                Nbcs, NbcsDrho, NbcsDu, NbcsDv);
            
            %%To be consistent with the notation in the paper, let us take
            %%the negative sign of the above expression:
            frhoij(:, i + (j - 1) * S) =...
                -frhoij(:, i + (j - 1) * S);
            fuij(:, i + (j - 1) * S) = ...
                0.5 * (tmp - fuij(:, i + (j - 1) * S));
            fvij(:, i + (j - 1) * S) =...
                0.5 * (tmp1 - fvij(:, i + (j - 1) * S));
        end
    end
    
    %Advance Yi's using RK4
    if S>0
        fui_diffusion = fui + AU * ui(uid, :);
        fvi_diffusion = fvi + AV * vi(vid, :);
        frhoi_diffusion = frhoi + ARHO * rhoi(pid, :);
    end
    RHSY = zeros(size(YYt));
    resYYt = zeros(size(YYt));
    for irk=1:5
        Tlocal0 = t + dt * rk4c(irk);
        %Calculate YiYj's
        for i=1:S
            for j=1:S
                YiYj(:, i + (j - 1) * S) = ...
                    YYt(: , i) .* YYt(:, j);
            end
        end
        RHSY(:, 1:Sb) = rhsYibcs(YYt(:, 1:Sb), Tlocal0);
        for i=Sb+1:S
            A1j =  dx * dy * (app.IP(1) * (fui_diffusion)' * ui(uid, i) + ...
                app.IP(2) * (fvi_diffusion)' * vi(vid, i) + ...
                app.IP(3) * (frhoi_diffusion)' * rhoi(pid, i)); %column vector
            A2j =  dx * dy * (app.IP(1) * fuij' * ui(uid, i) + ...
                app.IP(2) * fvij' * vi(vid, i) + ...
                app.IP(3) * frhoij' * rhoi(pid, i));
            RHSY(:, i) = YYt * A1j + CYY(:)' * A2j(:) - YiYj * A2j(:);
        end
        %RK time stepping
        resYYt = resYYt * rk4a(irk) + dt * RHSY;
        YYt = YYt + rk4b(irk) * resYYt;
    end
        
    % Advance mean fields
    RHSq=1/dt*rho(pid) + frho - frhoij*CYY(:) + app.kappa*rhsrho;
    RHSu = 1/dt*u(uid) + fu - fuij*CYY(:) + rhsubc*nu - dpdx + ...
        Fu(t,[XU(idsu) YU(idsu)]);
    RHSv = 1/dt*v(vid) + fv - fvij*CYY(:) + rhsvbc*nu - dpdy +...
        Fv(t,[XV(idsv) YV(idsv)]);
    
    rho(pid(rhoperm)) = Urho\(Lrho\RHSq(rhoperm));
    u(uid(uperm))=Uu\(Lu\RHSu(uperm));
    v(vid(vperm))=Uv\(Lv\RHSv(vperm));
    
    frhoijmyyy = frhoij*MYYY;
    fuijmyyy = fuij*MYYY;
    fvijmyyy = fvij*MYYY;
    
    %Advance ui's, vi's, and rhoi's
    RHSY(:, 1:Sb) = rhsYibcs(YYt(:, 1:Sb), T - dt);
    for i=1:S
        f1rho = frhoi(:,i) - frhoijmyyy*CYYINV(:,i);
        f1u = fui(:,i) - fuijmyyy*CYYINV(:,i);
        f1v = fvi(:,i) - fvijmyyy*CYYINV(:,i);
        f2rho = f1rho + ARHO*rhoi(pid,i);
        f2u = f1u + AU*ui(uid,i);
        f2v = f1v + AV*vi(vid,i);
        %Calculate inner product terms <ffiUj> -- there has to be S of them
        InProdi=zeros(S,1);
        
        %Do the inner products for the boundary nodes
        for j = 1:Sb
            InProdi(j) = CYYINV(i,:) * dt *...
                (sum(spdiags(RHSY(:, j), 0, MC, MC) * YYt))';
        end
        for j=Sb+1:S
            InProdi(j) = (app.IP(1) * f2u' * ui00(uid,j) + ...
                app.IP(2) * f2v' * vi00(vid,j) + ...
                app.IP(3) * f2rho'*rhoi00(pid,j)) * dx * dy;
        end
        
        %We're going to treat the Laplacian for THIS mode implicitly.
        %but the inner product treats the Laplcian EXplicitly
        [dpdx dpdy]   = Grad2D(Pi(:,i),NodeP,dx,dy,idsu,idsv);
        RHSu = 1/dt*ui(uid,i) + f1u - dpdx...
            -ui00(uid,:)*InProdi;
        %NOTE! I am assuming that ui's are zero on the boundary. To
        %introduce stochastic boundary forcings, add to the above line.
        RHSv = 1/dt*vi(vid,i) + f1v - dpdy...
            -vi00(vid,:)*InProdi;
        %NOTE! I am assuming that vi's are zero on the boundary. To
        %introduce stochastic boundary forcings, add to the above line.
        ui(uid(uperm),i)=Uu\(Lu\RHSu(uperm));
        vi(vid(vperm),i)=Uv\(Lv\RHSv(vperm));
        
        %And now the tracers
        RHSq=1/dt*rhoi(pid,i) + f1rho - rhoi00(pid,:)*InProdi;
        rhoi(pid(rhoperm),i) = Urho\(Lrho\RHSq(rhoperm));
    end
    
    %Calculate RHS vector for P
    divU = Div2D(u,v,Nodeu,Nodev,dx,dy,idsP,Neu,Nev);
    RHSq = rhsPbc - 1/dt*divU;
    q(pid(Pperm))=UP\(LP\RHSq(Pperm));
    
    %Correct Velocity
    [dqdx dqdy]   = Grad2D(q,NodeP,dx,dy,idsu,idsv);
    u(uid) = u(uid) - dt*dqdx;
    v(vid) = v(vid) - dt*dqdy;
    
    %correct pressure
    P(pid) = P(pid) + q(pid) - nu*divU;
    %remove mean pressure
    if imposemean
        P(pid)=P(pid)-mean(P(pid));
    else
        P([1:length(bcsDP),pid])=...
            P([1:length(bcsDP),pid])-mean(P(pid));
    end
    
    %Correct pressure for ui, vi's
    for i=1:S
        %Calculate RHS vector for P
        divU = Div2D(ui(:,i),vi(:,i),Nodeu,Nodev,dx,dy,idsP,Neu,Nev);
        RHSq = rhsPbc - 1/dt*divU;
        q(pid(Pperm))=UP\(LP\RHSq(Pperm));
        
        %Correct Velocity
        [dqdx dqdy]   = Grad2D(q,NodeP,dx,dy,idsu,idsv);
        ui(uid,i) = ui(uid,i) - dt*dqdx;
        vi(vid,i) = vi(vid,i) - dt*dqdy;
        
        %correct pressure
        Pi(pid,i) = Pi(pid,i) + q(pid) - nu*divU;
        %remove mean pressure
        if imposemean
            Pi(pid,i)=Pi(pid,i)-mean(Pi(pid,i));
        else
            if i <= Sb
                Pi([1:length(bcsDP),pid], i)=...
                    Pi([1:length(bcsDP),pid], i)-mean(Pi(pid, i));
            else            
                Pi(pid, i) = ...
                    Pi(pid, i) - mean(Pi(pid, i));
            end
        end
    end

    %%% CORRECT DO CONDITION FOR IMPLICIT TERMS
    %%%The number of iterations
    for iter = 1:app.DO_cor_iter
        %Calculate the correction for each ui, vi, and rhoi
        %use fui, fvi, and frhoi to save some memory
        for i = 1:S
            f2rho = ARHO * rhoi00(pid, i);
            f2u = AU * ui00(uid, i);
            f2v = AV * vi00(vid, i);
            %Calculate inner product terms <ffiUj> -- there has to be S of them
            InProdi = zeros(S,1);
            for j = 1:S
                InProdi(j) = (app.IP(1) * f2u' * ui00(uid, j) + ...
                    app.IP(2) * f2v' * vi00(vid, j) + ...
                    app.IP(3) * f2rho' * rhoi00(pid, j)) * dx * dy;
            end
            fui(:, i) = ui00(uid, :) * InProdi;
            fvi(:, i) = vi00(vid, :) * InProdi;
            frhoi(:, i) = rhoi00(pid, :) * InProdi;

            %Now repeat at the new time-level
            f2rho = ARHO * rhoi(pid, i);
            f2u = AU * ui(uid, i);
            f2v = AV * vi(vid, i);
            %Calculate inner product terms <ffiUj> -- there has to be S of them
            InProdi = zeros(S,1);
            for j = 1:S
                InProdi(j) = (app.IP(1) * f2u' * ui(uid, j) + ...
                    app.IP(2) * f2v' * vi(vid, j) + ...
                    app.IP(3) * f2rho' * rhoi(pid, j)) * dx * dy;
            end
            fui(:, i) = fui(:, i) - ui(uid, :) * InProdi;
            fvi(:, i) = fvi(:, i) - vi(vid, :) * InProdi;
            frhoi(:, i) = frhoi(:, i) - rhoi(pid, :) * InProdi;
        end

        %Now do the corrections
        for i = 1:S
            ui(uid, i) = ui(uid, i) + dt * fui(:, i);
            vi(vid, i) = vi(vid, i) + dt * fvi(:, i);
            rhoi(pid, i) = rhoi(pid, i) + dt * frhoi(:, i);
        end
        %Record old basis for next iteration
        if iter < app.DO_cor_iter
            ui00 = ui;
            vi00 = vi;
            rhoi00 = rhoi;
        end
        %done, repeat if necessary
    end
    
%-------------------------------------------------------------------------%
    [YYt,ui,vi,Pi,rhoi] = ...
        Pete_DOorthnorm(app,dx,dy,YYt,ui,uid,vi,vid,Pi,pid,S,Sb,rhoi) ;
    
    runtime = toc(app.toptic) ;
    
    fprintf('Simulation step %04d of %04d. Time elapsed: %05.0f s.\n', ...
        k, T/dt, runtime)
    
    % call saving script
    if ( ~mod(k,SaveIntrvl) )
        
        % get observation at designated observation times
        if ( obsstep <= n_obssteps && t == ObsTimes(obsstep) )
            eval(ObsScript)
            obsstep = obsstep + 1 ;
        end
        
        eval(SaveScript)
        drawnow
        
    end
%-------------------------------------------------------------------------%
    
end

