% Solving Ostrovsky equation in Matlab, pseudospectral, Trefethen modification
% of Runge-Kutta method for time-stepping

% Clear variable list
clear;

% Parameters in the equation
alpha = 1.0;
beta = 1.0;
gamma = 1.0;

% Small, non-zero value for zero mode
k0 = 1e-2;

% Create directory with relevant name and navigate to created directory
mkdir('Ost_PS');
cd('Ost_PS');

% Set number of points and domain size
L = 1000;
h = 0.1;	% Discretisation in x
N = round(2*L/h); % Round to obtain integer

% Pre-allocated arrays
u1 = zeros(1, N);

% Calculate s
s = pi/L;

% Time step
dt = 0.0001;

% Set tmax
tmax = 5;
nt = tmax/dt;

% Set number of outputs
fout = 0.1;
nout = tmax/fout;
fo = fout/dt;

% Spatial discretisations
x = (-L:h:L-h);

% Anti-aliasing array coefficients (Orszag 2/3 rule)
AA = ones(1,N);
AA(round(N/3):round(2*N/3)) = 0;

% Calculate initial condition
A = 32;
Lambda = sqrt(12/abs(A));
U = A*sech(x/Lambda).^2;

% Set k array
k = [0:N/2,-(N/2-1):-1];
k(1) = k0;	% Small, non-zero value to avoid division by 0 error

% Parameters in sponge layer
nu = 1.0;
kappa = 12.0/L;

% Sponge function
spg = (nu/2)*(2 + tanh(kappa*(x - 3*L/4)) - tanh(kappa*(x + 3*L/4)));

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Solve Ostrovsky equation
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% Output initial condition
h = figure('Units','Normalized','Position',[.1 .1 .55 .55]);
plot(x, U, '-b','LineWidth',1.5);
xlim([-300 300]);
title(['Solution of the Ostrovsky equation at t = 0']);
set(gca, 'FontSize', 32);
xlabel('x');
ylabel('u','Rotation',0);
ax = gca;
ax.LineWidth = 2.5;
savefig('Ostn0.fig');
close(h);

% Convert initial condition to Fourier space
U = fft(U);

% Anti-aliasing on initial condition
U = U.*AA;

% Progress bar
h1 = waitbar(0,sprintf('Ostrovsky equation'));

% Define E and RKCoef (for Trefethen method)
E = exp(-(1.0/2.0)*(1i)*dt*(gamma./(s*k) - beta*((s*k).^3)));
RKCoef = -(1i)*s*alpha*dt*k;

% Solve using Runge-Kutta 4th order method
% Two loops - one for inner calculation, one for number of outputs
for i = 1:nout
    for j = 1:fo

        % Calculate RK4 coefficients for both leading order and higher order terms
        for z = 1:4
            switch z
                case 1
                    utemp = U;
                case 2
                    utemp = E.*(U + u1(1,:)/2);
                case 3
                    utemp = E.*U + u1(2,:)/2;
                case 4
                    utemp = (E.^2).*U + E.*u1(3,:);
            end
            
            % Calculate sponge layer
            
            
            % Compute u1
            u1(z,:) = (RKCoef/2).*fft(ifft(U).^2) - fft(spg.*ifft(AA.*U));

            % Anti-aliasing
            u1(z,:) = u1(z,:).*AA;
			
			% Set u1 at k = 0 to zero, so no evolution of U
			u1(z,1) = 0;
        end

        % Find next solution
        U = (E.^2).*U + (1/6)*((E.^2).*u1(1,:) + 2*E.*(u1(2,:) + u1(3,:)) + u1(4,:));

        % Update progress bar
        waitbar(((i-1)*fo + j)/(nout*fo));
    end

    % Create output array (separate step here for division by zero problem)
	k(1) = 0;
	UPlot = real(ifft(U));
	k(1) = k0;

	% Comparison plot - closed immediately as now saved in directory
    h = figure('Units','Normalized','Position',[.1 .1 .55 .55]);
    plot(x, UPlot, '-b','LineWidth',1.5);
    xlim([-100 100]); % Manually adjust to appropriate value
    title(['Solution of the Ostrovsky equation at t = ' num2str(i*fo*dt)]);
    set(gca, 'FontSize', 32);
    xlabel('x');
    ylabel('u','Rotation',0);
    ax = gca;
    ax.LineWidth = 2.5;

    % Save figure with relevant name
    savefig(['Ostn' num2str(i*fo) '.fig']);
    close(h);
end
close(h1);