%%%%%%%%%%%%%%%%%%
% WARNING: Need to add the line:
% app.Frhoi = @(t, XY, rhoi) -kappa*(rhoi).*repmat(((XY(:,1) <= a2)&(XY(:,2) >= b1)),1,size(rhoi,2));
% preferably immediately after initializing app.Frho in the setup scripts.
%%%%%%%%%%%%%%%%%%


function [P, u, v, rho, NodeP, Nodeu, Nodev] =...
    FluidsSolver2_DO3_force_GMM(app, SetupScript, PlotScript)
% function [P, u, v, rho, NodeP, Nodeu, Nodev] =...
%     FluidsSolver2_DO(app, SetupScript, PlotScript)
%
% Solves the stochastic Navier-Stokes/Stokes equations with Boussinesq buoyancy
% approximation by a simple Finite Volume discretization on a uniform Cartesian mesh using
% an Incremental Pressure Correction in Rotational Form projection method (by default).
% This function uses Dynamically Orthogonal (DO) Equations for uncertainty quantification.
% For time marching, a first-order backward difference scheme is used.
% LU factorization is used to solve the linear systems. The advection is mexed.
%
%
% INPUTS:
%
% app:	          The application data structure of which the fields can also be
%                 added through the SetupScript.
% SetupScript:    String input of filename used to setup the problem.
% 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.
%
%
% At a minimum, the structure app should contain the following fields:
%
% app.Nx:         Number of interior P cells in x-direction (integer number)
% app.Ny:         Number of interior P cells in y-direction (integer number)
% app.T:          Final time
% app.dt:         Timestep size
% app.PlotIntrvl: Number of time steps between two "PlotScript" calls
%
% Other optional fields of app:
%
% app.nu:         Effective viscosity in x direction
% app.nu2:        Effective viscosity in y direction
% app.kappa:      Effective diffusivity of density in x direction
% app.kappa2:     Effective diffusivity of density in y direction
% app.g:          Effective gravity acceleration along negative y direction
% app.Stokes:     If true, this function solves the Stokes equation
% app.NonInc:     If true, solves the equations using the
%                 non-incremental Projection Method
% app.NonRot:     If true, solves the equations using the
%                 Incremental Projection method in the non-rotational form
% app.Advect:     'TVD'   -- Use the Total Variation Diminishing
%                            advection scheme (DEFAULT)
%                 'QUICK' -- Quadratic Upwind Interpolation scheme for
%                            for Convective Kinetics
%                 'UW'    -- Use the UpWind advection scheme
%                 'CDS'   -- Use the Central Differencing Scheme
% app.Fcor:       Function that sets the coriolis forces.
%                 Called in the codes as: app.Fcor(time, [X, Y])
% app.updateBcs:  (Optional) Function that updates the time dependent boundary conditions.
%
%
% This solver needs the following variables correctly created in the SetupScript.
%
% 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
% bcsDq:     q (for Projection Method) Dirichlet boundary conditions
% bcsDu:     U-velocity Dirichlet boundary conditions
% bcsDv:     V-velocity Dirichlet boundary conditions
% bcsNrho:   Density Neumann boundary conditions
% bcsNq:     q (for Projection Method) 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 cell id's (see NodePad.m)
% idsu:      Interior U-velocity cell id's (see NodePad.m)
% idsv:      Interior V-velocity cell id's (see NodePad.m)
% Fu(t, XY): Forcing function for U-velocity, XY =[XU(idsu), YU(idsu)]
% Fv(t, XY): Forcing function for V-velocity, XY =[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 for q.
% 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 for q.
% 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: Diffusion_Operator, Grad2D_P, Grad2D_uv
%
% Date for Last Modification: Jan. 2014 (Last modified by Jing Lin)
%
% Authors: Matt Ueckermann, Themis Sapsis, Pierre Lermusiaux, Pat Haley


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%% Preliminary setup

tic;

% Run setup script
eval(SetupScript);

% Default options related to the solver
if ~isfield(app, 'Stokes'), app.Stokes = 0; end;
if ~isfield(app, 'NonInc'), app.NonInc = 0; end;
if ~isfield(app, 'NonRot'), app.NonRot = 0; end;

% Default options related to the tracer equations
if ~exist('bcsDrho', 'var'), bcsDrho = zeros(size([bcsDq; bcsNq])); end;
if ~exist('bcsOrho', 'var'), bcsOrho = []; end;
if ~exist('bcsNrho', 'var'), bcsNrho = []; end;
if ~exist('NodeRho', 'var'), NodeRho = NodeP; end;
if ~exist('rho', 'var'), rho = zeros(size(P)); end;
if ~isfield(app, 'kappa'), app.kappa = 0; end;
if ~isfield(app, 'kappa2'), app.kappa2 = app.kappa; end;
if ~isfield(app,'g'), app.g = 1; end;
if ~isfield(app,'Frho'), app.Frho = @(t, XY, rho) zeros(size(XY, 1), 1); end;

% Default options related to advection
if ~isfield(app, 'Advect'), app.Advect = 'TVD'; end;

% Default options related to diffusion
if ~isfield(app, 'nu'), app.nu = 1; end;
if ~isfield(app, 'nu2'), app.nu2 = app.nu; end;

% Default options related to the Coriolis forces
if ~isfield(app,'Fcor'), app.Fcor = @(t, XY) zeros(size(XY, 1), 1); end;

% Default options related to open boundary conditions
if ~exist('bcsOq','var'), bcsOq = []; end;
if ~exist('bcsOu','var'), bcsOu = []; end;
if ~exist('bcsOv','var'), bcsOv = []; end;
if ~exist('bcsOrho','var'), bcsOrho = []; end;
if (~isempty(bcsOu))||(~isempty(bcsOv))||(~isempty(bcsOrho))
    error('Open boundary conditions can only be set for pressure!');
end;
if ~isfield(app, 'RmSing'), app.RmSing = 0; end;

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


% Find number of boundary conditions and number of interior cells
Nbcs = length(bcsDq) + length(bcsNq);
Num_IntCell_P = max(NodeP(:)) - Nbcs;
Num_IntCell_u = max(Nodeu(:)) - Nbcs;
Num_IntCell_v = max(Nodev(:)) - Nbcs;

% Default options related to DO
if ~isfield(app, 'IP'), app.IP = [1, 1, 1]; end;
if ~isfield(app, 'S'), app.S = 0; end;
if ~exist('Yi', 'var'), Yi = 0; end;
if ~isfield(app,'DOAdaptIntrvl')
    if length(app.S) == 2
        app.DOAdaptIntrvl = 10;
    else
        app.DOAdaptIntrvl = 0;
    end;
end;
if ~isfield(app,'DO_cor_iter'), app.DO_cor_iter = 1; end;

% 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];


%% Construct diffusion operator matrices and corresponding RHS vectors

[Dfs_rho, Dfs_bc_rho, Dfs_bc_rhs_rho] = Diffusion_Operator...
    ('P', NodeRho, bcsDrho, bcsOrho, bcsNrho, dx, dy, app.kappa, app.kappa2);
[Dfs_q, Dfs_bc_q, Dfs_bc_rhs_q] = Diffusion_Operator...
    ('P', NodeP, bcsDq, bcsOq, bcsNq, dx, dy, 1, 1, app.RmSing);
[Dfs_u, Dfs_bc_u, Dfs_bc_rhs_u] = Diffusion_Operator...
    ('u', Nodeu, bcsDu, bcsOu, bcsNu, dx, dy, app.nu, app.nu2);
[Dfs_v, Dfs_bc_v, Dfs_bc_rhs_v] = Diffusion_Operator...
    ('v', Nodev, bcsDv, bcsOv, bcsNv, dx, dy, app.nu, app.nu2);

Dfs_u_TmMch = Dfs_u + (1/app.dt)*speye(size(Dfs_u));
Dfs_v_TmMch = Dfs_v + (1/app.dt)*speye(size(Dfs_v));
Dfs_rho_TmMch = Dfs_rho + (1/app.dt)*speye(size(Dfs_rho));


%% Construct the Coriolis operators

if (sum(abs(app.Fcor(0, [XV(idsv) YV(idsv)]))) > 0)
    [Au2v, Av2u] = CoriolisAVG(app.Nx, app.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


%% LU factorization of matrices

% Compute permutations that give best LU spartsity pattern
rhoperm = amd(Dfs_rho_TmMch);
qperm = amd(Dfs_q);
uperm = amd(Dfs_u_TmMch);
vperm = amd(Dfs_v_TmMch);

% LU factorization
[Lrho, Urho] = lu(Dfs_rho_TmMch(rhoperm, rhoperm));
[Lq, Uq] = lu(Dfs_q(qperm, qperm));
[Lu, Uu] = lu(Dfs_u_TmMch(uperm, uperm));
[Lv, Uv] = lu(Dfs_v_TmMch(vperm, vperm));

% Get interior ids
pid = Nbcs+1:length(P);
uid = Nbcs+1:length(u);
vid = Nbcs+1:length(v);
srcid     = 2:app.Nx+1;
srcTopid  = 2:app.Ny;
srcBotid  = 3:app.Ny+1;
bcN_ID_u = [length(bcsDu)+1, Nbcs];
bcN_ID_v = [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 (app.S(1) > 0)
    frhoi = zeros(length(pid), app.S(1));
    fui = zeros(length(uid), app.S(1));
    fvi = zeros(length(vid), app.S(1));
    frhoij = zeros(length(pid), app.S(1)^2);
    fuij = zeros(length(uid), app.S(1)^2);
    fvij = zeros(length(vid), app.S(1)^2);
else
    frhoi = 0; fui = 0; fvi = 0; frhoij = 0; fuij= 0; fvij= 0;
end;

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


%% Time marching to solve the time-evolving system

if isfield(app,'k');
    k = app.k;
else
    k = 0;
end;

T = app.dt*k; CYY = cov(Yi);
eval(PlotScript);
if (obsstep <= n_obssteps)&&((k/app.PlotIntrvl) == ObsTimes(obsstep))
    eval(app.ObsScript);
    obsstep = obsstep + 1 ;
end;
display('Paused at initial condition. Hit any key to continue');
pause;

if (app.DOAdaptIntrvl > 0), DOAdapt_mex; end;

q = zeros(size(P));
kstart = k + 1;
for k = kstart : Nt
    T = app.dt*k;
    disp(['   T=',num2str(T)]);
    
    % Need to save old modes to do the projections correctly
    ui_old = ui;
    vi_old = vi;
    rhoi_old = rhoi;
    
    % Calculate statistics
    MY = mean(Yi);
    for i = 1 : app.S(1)
        Yi(:,i) = Yi(:,i) - MY(i);
    end
    CYY = cov(Yi);
    CYY_Inv = pinv(CYY, max(abs(CYY(:)))*sqrt(eps));
    MYYY = mom_thrd(Yi);
    if (app.S(1) > 0), MYYY = reshape(MYYY, app.S(1)^2, app.S(1)); end;
    if (app.S(1) == 0), CYY = 0; end;
    
    % Solve mean field
    
    % Calculate pressure gradient
    [dpdx, dpdy] = Grad2D_P(P, NodeP, dx, dy, idsu, idsv);
    
    % Calculate Boussinesq gravity forcing term.
    src = app.g*0.5*(rho(NodeRho(srcTopid, srcid)) + rho(NodeRho(srcBotid, srcid)));
    src = src(idsv);
    
    % Calculate RHS vectors for mean u, v, and rho
    [fu, fv] = UVadvect_TVD_mex(app.Nx, app.Ny, dx, dy, app.dt, ...
        u, v, u, v, Nodeu, Nodev, idsu, idsv, Nbcs, NbcsDu, NbcsDv);
    fu = fu + Av2u * v;
    fv = fv - Au2v*u - src;
    frho = SCAadvect_TVD_mex(app.Nx, app.Ny, dx, dy, app.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 : app.S(1)
        [fui(:,i), fvi(:,i)] = UVadvect_TVD_mex(app.Nx, app.Ny, dx, dy, app.dt,...
            ui(:,i), vi(:,i), u, v, Nodeu, Nodev, idsu, idsv, Nbcs, NbcsDu, NbcsDv);
        
        % Compute two fluxes of convection by modes for "average TVD"
        [tmp, tmp1] = UVadvect_TVD_mex(app.Nx, app.Ny, dx, dy, app.dt,...
            u, v, ui(:,i), vi(:,i), Nodeu, Nodev, idsu, idsv, Nbcs, NbcsDu, NbcsDv);
        [tmp2, tmp3] = UVadvect_TVD_mex(app.Nx, app.Ny, dx, dy, app.dt,...
            u, v, -ui(:,i), -vi(:,i), Nodeu, Nodev, idsu, idsv, Nbcs, NbcsDu, NbcsDv);
        
        src = app.g*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(app.Nx, app.Ny, dx, dy, app.dt, rhoi(:,i),...
            u, v, NodeRho, Nodeu, Nodev, idsP, Nbcs, NbcsDrho, NbcsDu, NbcsDv)...
            + 0.5 * ( SCAadvect_TVD_mex(app.Nx, app.Ny, dx, dy, app.dt, rho, ui(:,i),...
            vi(:,i), NodeRho, Nodeu, Nodev, idsP, Nbcs, NbcsDrho, NbcsDu, NbcsDv)...
            - SCAadvect_TVD_mex(app.Nx, app.Ny, dx, dy, app.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 : app.S(1)
            [fuij(:, i + (j - 1) * app.S(1)), fvij(:, i + (j - 1) * app.S(1))] = ...
                UVadvect_TVD_mex(app.Nx, app.Ny, dx, dy, app.dt,...
                ui(:,i), vi(:,i), ui(:,j), vi(:,j), ...
                Nodeu, Nodev, idsu, idsv, Nbcs, NbcsDu, NbcsDv);
            [tmp, tmp1] = ...
                UVadvect_TVD_mex(app.Nx, app.Ny, dx, dy, app.dt,...
                ui(:,i), vi(:,i), -ui(:,j), -vi(:,j), ...
                Nodeu, Nodev, idsu, idsv, Nbcs, NbcsDu, NbcsDv);
            frhoij(:,i+(j-1)*app.S(1)) = ...
                0.5 * ( SCAadvect_TVD_mex(app.Nx, app.Ny, dx, dy, app.dt,...
                rhoi(:,i), ui(:,j), vi(:,j), NodeRho, Nodeu, Nodev, idsP,...
                Nbcs, NbcsDrho, NbcsDu, NbcsDv)...
                -  SCAadvect_TVD_mex(app.Nx, app.Ny, dx, dy, app.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)*app.S(1)) = -frhoij(:, i + (j-1)*app.S(1));
            fuij(:, i + (j-1)*app.S(1)) = 0.5*(tmp - fuij(:, i + (j-1)*app.S(1)));
            fvij(:, i + (j-1)*app.S(1)) = 0.5*(tmp1 - fvij(:, i + (j-1)*app.S(1)));
        end;
    end;
    
    % Advance Yi's using RK4
    if (app.S(1) > 0)
        fui_diffusion = fui - Dfs_u * ui(uid, :);
        fvi_diffusion = fvi - Dfs_v * vi(vid, :);
        frhoi_diffusion = frhoi - Dfs_rho * rhoi(pid, :) + app.Frhoi(T,[XP(idsP) YP(idsP)], rhoi_old(pid,:));
    end;
    RHSY = zeros(size(Yi));
    resYYt = zeros(size(Yi));
    for iRK = 1 : 5
        Tlocal = T + app.dt * rk4c(iRK);
        % Calculate YiYj's
        for i=1:app.S(1)
            for j=1:app.S(1)
                YiYj(:, i + (j - 1)*app.S(1)) = Yi(:, i).*Yi(:, j);
            end;
        end;
        A1j =  (dx*dy) * (app.IP(1)*(fui_diffusion'*ui(uid,:)) ...
            + app.IP(2)*(fvi_diffusion'*vi(vid, :))...
            + app.IP(3)*(frhoi_diffusion'* rhoi(pid, :)));
        A2j =  (dx*dy) * (app.IP(1)*(fuij'*ui(uid, :)) + ...
            app.IP(2)*(fvij'*vi(vid, :)) + app.IP(3)*(frhoij'*rhoi(pid, :)));
        RHSY = Yi*A1j + (repmat(CYY(:)',[app.MC 1]) - YiYj)*A2j;
        % RK time stepping
        resYYt = resYYt * rk4a(iRK) + app.dt*RHSY;
        Yi = Yi + rk4b(iRK) * resYYt;
    end;
    
    % Advance mean fields
    RHSq = (1/app.dt)*rho(pid) + frho - frhoij*CYY(:) + Dfs_bc_rhs_rho + app.Frho(T,[XP(idsP) YP(idsP)], rho(pid));
    RHSu = (1/app.dt)*u(uid) + fu - fuij*CYY(:) + Dfs_bc_rhs_u - dpdx + ...
        Fu(T,[XU(idsu) YU(idsu)]);
    RHSv = (1/app.dt)*v(vid) + fv - fvij*CYY(:) + Dfs_bc_rhs_v - 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
    for i = 1 : app.S(1)
        f1rho = frhoi(:,i) - frhoijmyyy*CYY_Inv(:,i)+ app.Frhoi(T,[XP(idsP) YP(idsP)], rhoi_old(pid,i));
        f1u = fui(:,i) - fuijmyyy*CYY_Inv(:,i);
        f1v = fvi(:,i) - fvijmyyy*CYY_Inv(:,i);
        f2rho = f1rho - Dfs_rho*rhoi(pid,i);
        f2u = f1u - Dfs_u*ui(uid,i);
        f2v = f1v - Dfs_v*vi(vid,i);
        
        %Calculate inner product terms <ffiUj> -- there has to be app.S of them
        InProdi = (dx*dy) * (app.IP(1)*(f2u'* ui_old(uid,:)) + ...
            app.IP(2)*(f2v'*vi_old(vid,:)) + app.IP(3)*(f2rho'*rhoi_old(pid,:)))';
        
        %We're going to treat the Laplacian for THIS mode implicitly.
        %but the inner product treats the Laplcian EXplicitly
        [dpdx, dpdy] = Grad2D_P(Pi(:,i),NodeP,dx,dy,idsu,idsv);
        RHSu = (1/app.dt)*ui(uid,i) + f1u - dpdx - ui_old(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/app.dt)*vi(vid,i) + f1v - dpdy - vi_old(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/app.dt)*rhoi(pid,i) + f1rho - rhoi_old(pid,:)*InProdi;
        rhoi(pid(rhoperm),i) = Urho \ (Lrho \ RHSq(rhoperm));
    end;
    
    % Calculate RHS vector for P
    [dudx, dvdy] = Grad2D_uv(u, v, Nodeu, Nodev, idsP, dx, dy, bcN_ID_u, bcN_ID_v);
    RHSq = Dfs_bc_rhs_q - (1/app.dt)*(dudx + dvdy);
    q(pid(qperm)) = Uq \ (Lq \ RHSq(qperm));
    
    % Correct Velocity
    [dqdx, dqdy] = Grad2D_P(q, NodeP, dx, dy, idsu, idsv);
    u(uid) = u(uid) - app.dt*dqdx;
    v(vid) = v(vid) - app.dt*dqdy;
    
    % Correct pressure
    if ~app.NonInc
        P(pid) = P(pid) + q(pid);
        if ~app.NonRot
            P(pid) = P(pid) - (app.nu*dudx + app.nu2*dvdy);
            [dudx, dvdy] =...
                Grad2D_uv(u, v, Nodeu, Nodev, idsP, dx, dy, bcN_ID_u, bcN_ID_v);
            P(pid) = P(pid) + (app.nu*dudx + app.nu2*dvdy);
        end;
    end;
    
    % Remove the mean of the pressure (determined up to a constant)
    P(pid) = P(pid) - mean(P(pid));
    
    % Correct pressure for ui, vi's
    for i = 1:app.S(1)
        % Calculate RHS vector for P
        [dudx, dvdy] =...
            Grad2D_uv(ui(:,i), vi(:,i), Nodeu, Nodev, idsP, dx, dy, bcN_ID_u, bcN_ID_v);
        RHSq = Dfs_bc_rhs_q - (1/app.dt)*(dudx + dvdy);
        q(pid(qperm)) = Uq \ (Lq \ RHSq(qperm));
        
        % Correct Velocity
        [dqdx, dqdy] = Grad2D_P(q, NodeP, dx, dy, idsu, idsv);
        ui(uid,i) = ui(uid,i) - app.dt*dqdx;
        vi(vid,i) = vi(vid,i) - app.dt*dqdy;
        
        % Correct pressure
        if ~app.NonInc
            Pi(pid,i) = Pi(pid,i) + q(pid);
            if ~app.NonRot
                Pi(pid,i) = Pi(pid,i) - (app.nu*dudx + app.nu2*dvdy);
                [dudx, dvdy] = Grad2D_uv(ui(:,i), vi(:,i),...
                    Nodeu, Nodev, idsP, dx, dy, bcN_ID_u, bcN_ID_v);
                Pi(pid,i) = Pi(pid,i) + (app.nu*dudx + app.nu2*dvdy);
            end;
        end;
        
        % Remove the mean of the pressure (determined up to a constant)
        Pi(pid,i) = Pi(pid,i) - mean(Pi(pid,i));
    end;
    
    % CORRECT DO CONDITION FOR IMPLICIT TERMS
    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:app.S
            f2rho = -Dfs_rho * rhoi_old(pid, i);
            f2u = -Dfs_u * ui_old(uid, i);
            f2v = -Dfs_v * vi_old(vid, i);
            %Calculate inner product terms <ffiUj> -- there has to be app.S of them
            InProdi = (app.IP(1) * (f2u' * ui_old(uid, :)) + ...
                app.IP(2) * (f2v' * vi_old(vid, :)) + ...
                app.IP(3) * (f2rho' * rhoi_old(pid, :)))' * (dx * dy);
            fui(:, i) = ui_old(uid, :) * InProdi;
            fvi(:, i) = vi_old(vid, :) * InProdi;
            frhoi(:, i) = rhoi_old(pid, :) * InProdi;
            
            % Now repeat at the new time-level
            f2rho = -Dfs_rho * rhoi(pid, i);
            f2u = -Dfs_u * ui(uid, i);
            f2v = -Dfs_v * vi(vid, i);
            % Calculate inner product terms <ffiUj> -- there has to be app.S of them
            InProdi = (app.IP(1) * (f2u' * ui(uid, :)) + ...
                app.IP(2) * (f2v' * vi(vid, :)) + ...
                app.IP(3) * (f2rho' * rhoi(pid, :)))' * (dx * dy);
            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:app.S
            ui(uid, i) = ui(uid, i) + app.dt * fui(:, i);
            vi(vid, i) = vi(vid, i) + app.dt * fvi(:, i);
            rhoi(pid, i) = rhoi(pid, i) + app.dt * frhoi(:, i);
        end;
        %Record old basis for next iteration
        ui_old = ui; vi_old = vi; rhoi_old = rhoi;
    end;
    
    if (app.ReOrthonorm)
        [Yi, ui, vi, Pi, rhoi] = DOorthnorm(app, dx, dy, Yi, ui, uid, vi, vid, Pi, pid, rhoi);
    end;
    
    %Adapt stochastic dimension if required
    if ~mod(k, app.DOAdaptIntrvl), DOAdapt_mex; end;
    
    %Plot solution
    if ~mod(k,app.PlotIntrvl)
        eval(PlotScript);
        if (obsstep <= n_obssteps)&&((k/app.PlotIntrvl) == ObsTimes(obsstep))
            eval(app.ObsScript);
            obsstep = obsstep + 1 ;
        end;
    end;
end;

%% Clean up any temporary files
if exist('existbasis','var')
    delete([d0,sprintf('/Save/DOAdaptSave/Basis/orth_basis%0.3d.mat',existbasis)]);
end
