% This script plots the nondimensional dispersion curves for Matsuno waves
% as well as the eigenfunctions of modes chosen by the user.  March, 2021
%-------------------------------------------------------------------------
n=-1:5;         % Range of n values
k=-5:0.01:5;    % Range of zonal wavenumber
xmax=2*pi;      % Maximum nondimensional zonal extent
ymax=3;         % Maximum nondimensional y in eigenfunction plots
cmp='jet';      % Colormap for eigenfunction plot 
%-------------------------------------------------------------------------
qn=max(size(n));
qk=max(size(k));
%
clr=zeros(qn,3);   % Colors of dispersion curves
clr(1,:)=[0 0 1];
clr(2,:)=[0.133 0.0545 0.133];
for i=3:qn
   clr(i,:)=[1 0 0];
end   
clr_rossby=[1 0 1];
%
om=zeros(qn,qk,3);
close all
figure('units','normalized','position',[0.55,0.45,0.4,0.45])
hold on
for i=1:qn
    if n(i) == -1       % Kelvin
        dum=k;
        dum(k<0)=NaN;
        om(i,:,1)=dum;
        om(i,:,2:3)=NaN;
    elseif n(i) == 0    % Yanai
        om(i,:,1)=0.5*(k+sqrt(k.^2+4));
        om(i,:,2:3)=NaN;
    else                % Poincare and Rossby
        for j=1:qk
            p=[1,0,-(k(j)^2+2*n(i)+1),-k(j)];
            dum=roots(p);
            dum(dum<=0)=NaN;
            om(i,j,:)=dum;
        end 
    end 
    %
    for m=1:3
        if n(i) > 0 && max(om(i,:,m)) < 1  % Identify as Rossby if n > 0 and  omega < 1
            clr(i,:)=clr_rossby; % 
        end    
        plot(k,squeeze(om(i,:,m)),'color',clr(i,:),'linewidth',1.2);
    end    
end
hold off
%
set(gca,'fontsize',13,'fontweight','bold')
xlabel('k','fontsize',15)
ylabel('\omega','rotation',0,'fontsize',15)
set(gca,'xtick',(min(k):max(k)))
title('Matsuno Modes')
%
% Now make a plot of user-selected eigenfunction
%
nu=input('Input mode number n (-1 for Kelvin)');
kmin=min(k);kmax=max(k);kmins=num2str(kmin);kmaxs=num2str(kmax);
if nu > 0
    rossby=input('Rossby (''R'') or Poincare (''P'') mode?','s');
end    
istring=['Enter zonal wavenumber between ',kmins,' and ', kmaxs,' '];
ku=input(istring);
[~,nindex]=min(abs(n-nu));
[~,kindex]=min(abs(k-ku));
%
% Find wave frequency
%
if nu == -1
    omg=ku;
elseif nu == 0
    omg=0.5*(ku+sqrt(ku.^2+4));
else
    if strcmpi(rossby,'P')
        omg=max(om(nindex,kindex,:));
    else
        omg=min(om(nindex,kindex,:));
    end
end    
%
%
y=-ymax:0.1*ymax:ymax; % Values of y for mapping eigenfunction
x=0:0.1*xmax:xmax;         % Values of x for mapping eigenfunction
%
i=sqrt(-1);  % Restore definition of i
%
ys=max(size(y));
xs=max(size(x));
ey=exp(-0.5.*y.^2);
eyd=-y.*ey;
fx=exp(i.*x);
if ku == 0
    fx=zeros(1,xs)+1-i;
end    
xmat=repmat(fx',[1,ys]);
ymat=repmat(ey,[xs,1]);
%
if nu == -1
    v=zeros(xs,ys);
    u=real(xmat.*ymat);
    h=real(omg*xmat.*ymat/ku);
    w=real(i*omg^2*xmat.*ymat/ku);
else
    clear ha hap
    ha(1,1:ys)=1;        % Hermite polynomials
    ha(2,1:ys)=2.*y;
    ha(3,1:ys)=4.*y.^2-2;
    ha(4,1:ys)=8.*y.^3-12.*y;
    %
    hap(1,1:ys)=0;       % First derivatives of Hermite polynomials
    hap(2,1:ys)=2;
    hap(3,1:ys)=8.*y;
    hap(4,1:ys)=24.*y.^2-12;
    %
    vf=ey.*ha(nu+1,:);
    dvdy=ey.*(squeeze(hap(nu+1,:))-y.*squeeze(ha(nu+1,:)));
    hf=(i*omg.*dvdy-i*y.*ku.*vf)/(ku^2-omg^2);
    uf=(ku*hf+i.*y.*vf)./omg;
    wf=i*omg*hf;
    %
    v=real(repmat(vf,[xs,1]).*xmat);
    h=real(repmat(hf,[xs,1]).*xmat);
    u=real(repmat(uf,[xs,1]).*xmat);
    w=real(repmat(wf,[xs,1]).*xmat);
end    
%
% Plot eigenfunction
%
figure('units','normalized','position',[0.55,0.05,0.4,0.45])
contourf(x,y,h');
us=colormap(cmp);
us(1:100,:)=[];
colormap(us)
hold on
if min(min(w)) < 0
    contour(x,y,min(w',0),'--k')
end
if max(max(w)) > 0
    contour(x,y,max(w',0),'k')
end    
set(gca,'CLim',[-max(max(abs(h))) max(max(abs(h)))])
g=quiver(x,y,u',v','k','linewidth',1.5);
hold off
%
set(gca,'ylim',[-ymax ymax])
set(gca,'xlim',[0 xmax])
set(gca,'fontsize',13,'fontweight','bold')
xlabel('X','fontsize',14)
ylabel('Y','fontsize',14,'rotation',0)
colorbar
tit=['h (shading),u,v, and w (contours); n = ',num2str(nu),', k = ',num2str(ku)];
title(tit)