% Program to demonstrate the RK3 scheme for fluctuating Landau-Lifshitz 
% Navier-Stokes equations with periodic boundary conditions
% See Bell, et al, Phys. Rev. E, Vol. 76, 016708 (2007)
clear all;
%% Set physical parameters and system constants
k_B = 1.38054e-16;  % Boltzmann's constant
gamma = 5/3;        % Specific heat ratio
gammaM1 = gamma-1;
mass = 6.63e-23;    % Particle mass
diam = 3.66e-8;     % Particle diameter
Rgas = k_B/mass;    % Gas constant
c_V = 1.5*Rgas;     % Specific heat per particle
rhoEq = 1.78e-3;    % Equilibrium mass density
tempEq = 273;       % Equilibrium temperature
cSound = sqrt(gamma*Rgas*tempEq);   % Sound speed
velEq = 0.5*cSound;     % Equilibrium fluid velocity (can be zero)
momEq = rhoEq*velEq;    % Equilibrium momentum density
presEq = rhoEq * Rgas * tempEq ;        % Equilibrium pressure
enrEq = presEq/(gamma-1) + 0.5*rhoEq*velEq^2;   % Energy density
viscC = 5/(16*diam^2)*sqrt(mass*k_B/pi);  % Viscosity coefficient
lambC = 15*k_B*viscC/(4*mass);      % Thermal conductivity coefficient
alpha1 = 0.25 * (sqrt(7)+1);        % Interpolation coefficient
alpha2 = 0.25 * (sqrt(7)-1);        % Interpolation coefficient

%% Initialize the grid values
sysLen = 1.25e-4;               % System length
sysVol = 1.96e-16;              % System volume
nCells = 40;                    % Number of grid points (or cells)
cellVol = sysVol / nCells;      % Volume of single grid point
dx = sysLen / nCells;           % Grid spacing
rho = rhoEq * ones(nCells,1);   % Mass density (\rho)
mom = momEq * ones(nCells,1);   % Momentum density (J)
enr = enrEq * ones(nCells,1);   % Energy density (E)
ip = [2:nCells, 1];             % i+1, with periodic b.c.
im = [nCells, 1:(nCells-1)];    % i-1, with periodic b.c.
ipp = [3:nCells, 1, 2];         % i+2, with periodic b.c.
imm = [(nCells-1), nCells, 1:(nCells-2)];   % i-2, with periodic b.c.

%% Initialize statistics variables
rhoMean = zeros(nCells,1);     rhoVar = zeros(nCells,1);
momMean = zeros(nCells,1);     momVar = zeros(nCells,1);
enrMean = zeros(nCells,1);     enrVar = zeros(nCells,1);
rhoMomCov = zeros(nCells,1);  
rhoAveInit = mean(rho);  momAveInit = mean(mom);  enrAveInit = mean(enr);

%%%%%% Main loop %%%%%%
dt = 1.0e-12;   % Time step
dtOdx = dt/dx;
sCoeff = sqrt(2)*sqrt(4*k_B/(3*dt*cellVol));  % Momentum noise coefficient
qCoeff = sqrt(2)*sqrt(k_B/(dt*cellVol));      % Energy noise coefficient
nSteps = input('Enter number of steps (try 100000): ');
nSkip = nSteps/4; nSamp = 0;    % Start taking statistics after nSkip steps
tic;    % Start MATLAB timer
for iStep=1:nSteps

    for iStage=1:3  % RK3 has three stages
        
        %% Set current values for this stage
        switch( iStage ) 
            case {1}
                rhoI = rho;     momI = mom;     enrI = enr;
            case {2}
                rhoI = rho13;     momI = mom13;     enrI = enr13;
            case {3}
                rhoI = rho23;     momI = mom23;     enrI = enr23;
        end

        %% Compute conserved variables on half-grid by special interpolation
        %% Notation:  rhoI(i) = \rho_i (Full grid); rhoP = \rho_{i+1/2} (Half-grid)
        rhoP = alpha1*( rhoI(ip) + rhoI(:) ) - alpha2*( rhoI(ipp) + rhoI(im) );
        momP = alpha1*( momI(ip) + momI(:) ) - alpha2*( momI(ipp) + momI(im) );
        enrP = alpha1*( enrI(ip) + enrI(:) ) - alpha2*( enrI(ipp) + enrI(im) );

        %% Compute hydrodynamic variables
        velP = momP ./ rhoP;                        % Half-grid velocity
        presP = gammaM1*(enrP - 0.5*momP.*velP);    % Half-grid pressure
        tempP = presP./(Rgas*rhoP);                 % Half-grid temperature
        velI = momI ./ rhoI;                        % Full-grid velocity
        tempI = (enrI./rhoI - 0.5*velI.^2)/c_V;     % Full-grid temperature
        
        %% Compute transport coefficients
        sqrtTI = sqrt(tempI);       sqrtTP = sqrt(tempP);
        viscI = viscC*sqrtTI;       viscP = viscC*sqrtTP;
        lambI = lambC*sqrtTI;       lambP = lambC*sqrtTP;

        %% Compute gradients of hydrodynamic variables
        velGradP = ( velI(ip) - velI(:) )/dx;     % Note: velI(ip) = v_{i+1} with p.b.c.
        tempGradP = ( tempI(ip) - tempI(:) )/dx;

        %% Compute stochastic fluxes
        sTildeP = sCoeff * sqrt(viscI(ip).*tempI(ip) + viscI.*tempI) .* randn(nCells,1);
        qTildeP = qCoeff * sqrt(lambI(ip).*tempI(ip).^2 + lambI.*tempI.^2) .* randn(nCells,1);

        %% Compute total fluxes
        rhoFluxP = momP;
        momFluxP = momP.*velP + presP - (4/3)*viscP.*velGradP - sTildeP;
        enrFluxP = (enrP + presP).*velP - ((4/3)*viscP.*velGradP).*velP ...
                            - lambP.*tempGradP - sTildeP.*velP - qTildeP;

        %% Compute update for this stage
        switch( iStage ) 
            case {1}
                rho13 = rho - dtOdx * ( rhoFluxP(:) - rhoFluxP(im) );
                mom13 = mom - dtOdx * ( momFluxP(:) - momFluxP(im) );
                enr13 = enr - dtOdx * ( enrFluxP(:) - enrFluxP(im) );
            case {2}
                rho23 = 0.75*rho + 0.25*rho13 - 0.25*dtOdx * ( rhoFluxP(:) - rhoFluxP(im) );
                mom23 = 0.75*mom + 0.25*mom13 - 0.25*dtOdx * ( momFluxP(:) - momFluxP(im) );
                enr23 = 0.75*enr + 0.25*enr13 - 0.25*dtOdx * ( enrFluxP(:) - enrFluxP(im) );
            case {3}
                rho = (1/3)*rho + (2/3)*rho23 - (2/3)*dtOdx * ( rhoFluxP(:) - rhoFluxP(im) );
                mom = (1/3)*mom + (2/3)*mom23 - (2/3)*dtOdx * ( momFluxP(:) - momFluxP(im) );
                enr = (1/3)*enr + (2/3)*enr23 - (2/3)*dtOdx * ( enrFluxP(:) - enrFluxP(im) );
        end
    end % End of RK3 stages loop

    %% Accumulate statistics
    if( iStep > nSkip )
        nSamp = nSamp + 1;
        rhoMean = rhoMean + rho;             rhoVar = rhoVar + rho.^2;
        momMean = momMean + mom;             momVar = momVar + mom.^2;
        enrMean = enrMean + enr;             enrVar = enrVar + enr.^2;
        rhoMomCov = rhoMomCov + rho.*mom;  
    end
    
    %% Periodically display diagnostics
    if( rem(iStep,(nSteps/20)) < 1 )
        xPlot = 1:nCells;
        figure(1); clf;
        subplot(3,1,1), plot(xPlot,rho,'*'); ylabel('\rho');
        subplot(3,1,2), plot(xPlot,mom,'*'); ylabel('J');
        subplot(3,1,3), plot(xPlot,enr,'*'); ylabel('E');
        drawnow;
        timeElapsed = toc;
        fprintf('Finished %g of %g samples; about %g seconds to completion\n',...
            iStep,nSteps,timeElapsed/iStep*(nSteps-iStep));       
    end
end % End of Main Loop

%% Normailze statistics
rhoMean = rhoMean/nSamp;
rhoVar = rhoVar/nSamp - rhoMean.^2;
momMean = momMean/nSamp;
momVar = momVar/nSamp - momMean.^2;
enrMean = enrMean/nSamp;
enrVar = enrVar/nSamp - enrMean.^2;
rhoMomCov = rhoMomCov/nSamp - rhoMean.*momMean;

%% Verify conservation
rhoAveFinal = mean(rho);  momAveFinal = mean(mom);  enrAveFinal = mean(enr);
fprintf('Conserved variables: (Initial)/(Final)\n');
fprintf('Density            :  %g  /  %g \n',rhoAveInit,rhoAveFinal);
fprintf('X-momentum         :  %g  /  %g \n',momAveInit,momAveFinal);
fprintf('Energy             :  %g  /  %g \n',enrAveInit,enrAveFinal);

%% Compute theory values for variances and print results
numEq = rhoEq*cellVol/mass;
CT2 = k_B*tempEq/mass;
rho2Th = rhoEq^2/numEq;
mom2Th = (momEq^2 + rhoEq^2*CT2)/numEq;
enr2Th = (enrEq^2 + momEq^2*CT2 + 2/3*c_V^2*rhoEq^2*tempEq^2)/numEq;
rhoMomTh = rhoEq*momEq/numEq;
fprintf('Variances and Covariances\n');
fprintf('            (Theory) : (Measured) +/- (Error Bar)\n');
fprintf('Density   = %g : %g +/- %g \n',rho2Th,mean(rhoVar),std(rhoVar)/sqrt(nCells));
fprintf('Momentum  = %g : %g +/- %g \n',mom2Th,mean(momVar),std(momVar)/sqrt(nCells));
fprintf('Energy    = %g : %g +/- %g \n',enr2Th,mean(enrVar),std(enrVar)/sqrt(nCells));
fprintf('Rho-Jx       = %g : %g +/- %g\n',rhoMomTh,mean(rhoMomCov),std(rhoMomCov)/sqrt(nCells));

%% Plot results and theory
xPlot = 1:nCells;
figure(1); clf;     % Plot mean values
subplot(3,1,1)
plot(xPlot,rhoMean,'*',xPlot,rhoEq*ones(nCells,1)); ylabel('< \rho >');
subplot(3,1,2)
plot(xPlot,momMean,'*',xPlot,momEq*ones(nCells,1)); ylabel('< J >');
subplot(3,1,3)
plot(xPlot,enrMean,'*',xPlot,enrEq*ones(nCells,1)); ylabel('< E >');
figure(2); clf;     % Plot variances
subplot(3,1,1)
plot(xPlot,rhoVar,'*',xPlot,rho2Th*ones(nCells,1)); ylabel('< \delta\rho^2 >');
subplot(3,1,2)
plot(xPlot,momVar,'*',xPlot,mom2Th*ones(nCells,1)); ylabel('< \delta J^2 >');
subplot(3,1,3)
plot(xPlot,enrVar,'*',xPlot,enr2Th*ones(nCells,1)); ylabel('< \delta E^2 >');
