% Hua-sheng XIE, huashengxie@gmail.com, IFTS-ZJU, 2014-06-01 17:00
% pdrk.m, Plasma Dispersion Relastion solver, Maxwellian distribution
% Transform to matrix eigenvalue problem lambda*X=M*X.
% Ronnmark1982, eight-pole approximation for Z(zeta)=-sum(b_j/(zeta-c_j))
% 
% sparse, eigs
%
% 2014-08-22 01:09 T_perp/T_para \neq 1,  vs0 \neq 0

close all; clear; clc;

N=2; B0=1.0;

% % read input parameters
% par = importdata('pdrk.in', ' ', 1); 
par = importdata('pdrk_beam.in', ' ', 1); 
[S, col]=size(par.data);
if(col~=6)
    disp('Wrong input data !!!');
end

for s=1:S
    qs(s)=par.data(s,1); % charge
    ms(s)=par.data(s,2); % mass
    ns0(s)=par.data(s,3); % desity
    Tzs(s)=par.data(s,4); % parallel temperature
    Tps(s)=par.data(s,5); % perp temperature
    vs0(s)=par.data(s,6); % para drift velocity
end
vtzs=sqrt(2*Tzs./ms); % para thermal velocity
lambdaDs=sqrt(Tzs./(ns0.*qs.^2)); % Debye length, Tzs or Tps?
kDs=1./lambdaDs;
wps=sqrt(ns0.*qs.^2./ms); % plasma frequency
wcs=B0*qs./ms; % cyclon frequency
rhocs=sqrt(Tps./ms)./abs(wcs); % cyclotron radius

wps2=wps.^2;
lmdT=Tzs./Tps;

c2=1e4; % light speed
epsilon0=1; mu0=1/(c2*epsilon0);

betasz=2*mu0.*ns0.*Tzs./B0^2;
betasp=2*mu0.*ns0.*Tps./B0^2;
vA=B0/sqrt(mu0*sum(ms.*ns0)); %
cS=sqrt(2*min(Tzs)/max(ms));

%%
J=8;
if(J==4)
    % Martin1980, 4-pole
    bzj(1)=0.546796859834032 + 0.037196505239277i;
    bzj(2)=-1.046796859834027 + 2.101852568038518i;
    czj(1)=1.23588765343592 - 1.21498213255731i;
    czj(2)=-0.378611612386277 - 1.350943585432730i;
    
    bzj(3:4)=conj(bzj(1:2));
    czj(3:4)=-conj(czj(1:2));
else % J==8
    % Ronnmark1982, 8-pole for Z function, and Z'
    bzj(1)=-1.734012457471826E-2-4.630639291680322E-2i;
    bzj(2)=-7.399169923225014E-1+8.395179978099844E-1i;
    bzj(3)=5.840628642184073+9.536009057643667E-1i;
    bzj(4)=-5.583371525286853-1.120854319126599E1i;
    czj(1)=2.237687789201900-1.625940856173727i;
    czj(2)=1.465234126106004-1.789620129162444i;
    czj(3)=.8392539817232638-1.891995045765206i;
    czj(4)=.2739362226285564-1.941786875844713i;

    bzj(5:8)=conj(bzj(1:4));
    czj(5:8)=-conj(czj(1:4));
end
J=length(bzj);
% sum(bzj) % = -1
% sum(bzj.*czj) % = 0

SNJ=S*(2*N+1)*J;
SNJ1=SNJ+1;
SNJ3=3*SNJ1;
NN=SNJ3+6;

sp=0; % sparse eigs() or eig()
if(sp==0)
    nw=NN; % number of roots to obtain
else
    nw=2;
end
wg=1.5; % inital guess

kk=[]; kxx=[]; kzz=[]; ww=[];
tic;
kmx=0.5; kmz=kmx;
nz=20; nx=21; dx=kmx/(nx-1); dz=kmz/nz;
xx=0:dx:kmx; zz=dz:dz:kmz;
[ZZ,XX]=ndgrid(zz,xx);
WW=0.*ZZ;
for jx=1:nx
    for jz=1:nz

        kz=zz(jz); kx=xx(jx);        
        
        k=sqrt(kz^2+kx^2);
        bs=kx*rhocs;
        bs(abs(bs)<1e-50)=1e-50;  % to avoid singular when k_perp=0       

        bs2=bs.^2;
        M=sparse(NN,NN);
        snj=0;

        % initialize
        b11=0; b12=0; b13=0; b21=0; b22=0; b23=0; b31=0; b32=0; b33=0;    
        csnj=zeros(1,3*SNJ); 
        b11snj=csnj.*0; b12snj=csnj.*0; b13snj=csnj.*0; 
        b21snj=csnj.*0; b22snj=csnj.*0; b23snj=csnj.*0; 
        b31snj=csnj.*0; b32snj=csnj.*0; b33snj=csnj.*0;

        for s=1:S % species 
            for n=-N:N % Bessel function
                Gamn=exp(-bs2(s))*besseli(n,bs2(s)); % large k_perp will NaN
                Gamnp=exp(-bs2(s))*(besseli(n+1,bs2(s))+...
                    besseli(n-1,bs2(s))-2*besseli(n,bs2(s)))/2;
                for j=1:J % poles of Z(zeta)
                    snj=snj+1;

                    csnj(snj)=czj(j)*kz*vtzs(s)+kz*vs0(s)-n*wcs(s);

                    cnj=csnj(snj);
                    bj0=vs0(s)+(1-1/lmdT(s))*czj(j)*vtzs(s);

                    %
                    bnj1=kz*bj0/cnj;  % for A_nj
                    bnj2=1-bnj1;

                    tmp=wps2(s)*bzj(j);

                    b11snj(snj)=tmp*bnj2*n^2*Gamn/bs2(s);
                    b11=b11+tmp*bnj1*n^2*Gamn/bs2(s);

                    b12snj(snj)=tmp*bnj2*1i*n*Gamnp;
                    b12=b12+tmp*bnj1*1i*n*Gamnp;
                    b21snj(snj)=-b12snj(snj);
                    b21=-b12;

                    b22snj(snj)=tmp*bnj2*(n^2*Gamn/bs2(s)-2*bs2(s)*Gamnp);
                    b22=b22+tmp*bnj1*(n^2*Gamn/bs2(s)-2*bs2(s)*Gamnp);

                    %
                    bnj1=n*wcs(s)*bj0/cnj/vtzs(s);  % for eta_n*A_nj
                    bnj2=czj(j)/lmdT(s)-bnj1;

                    b13snj(snj)=tmp*bnj2*n*sqrt(2*lmdT(s))*Gamn/bs(s);
                    b13=b13+tmp*bnj1*n*sqrt(2*lmdT(s))*Gamn/bs(s);
                    b31snj(snj)=b13snj(snj);
                    b31=b13;

                    b23snj(snj)=-1i*tmp*bnj2*sqrt(2*lmdT(s))*Gamnp*bs(s);
                    b23=b23-1i*tmp*bnj1*sqrt(2*lmdT(s))*Gamnp*bs(s);
                    b32snj(snj)=-b23snj(snj);
                    b32=-b23;

                    %
                    bnj1=n^2*bj0/cnj/vtzs(s)^2/kz;  % for eta_n^2*A_nj
                    bnj2=(vs0(s)/vtzs(s)+czj(j))*czj(j)/lmdT(s)-...
                        n*wcs(s)*bj0*(1+n*wcs(s)/cnj)/vtzs(s)^2/kz;

                    b33snj(snj)=tmp*bnj2*2*lmdT(s)*Gamn;
                    b33=b33+tmp*bnj1*2*lmdT(s)*Gamn;
                end
            end
        end

        for snj=1:SNJ % set the eigen matrix
            jjx=snj+0*SNJ1;
            jjy=snj+1*SNJ1;
            jjz=snj+2*SNJ1;
            % v_snjx
            M=M+sparse(jjx,jjx,csnj(snj),NN,NN)+...
                sparse(jjx,SNJ3+1,b11snj(snj),NN,NN)+...
                sparse(jjx,SNJ3+2,b12snj(snj),NN,NN)+...
                sparse(jjx,SNJ3+3,b13snj(snj),NN,NN);

            % v_snjy
            M=M+sparse(jjy,jjy,csnj(snj),NN,NN)+...
                sparse(jjy,SNJ3+1,b21snj(snj),NN,NN)+...
                sparse(jjy,SNJ3+2,b22snj(snj),NN,NN)+...
                sparse(jjy,SNJ3+3,b23snj(snj),NN,NN);

            % v_snjz
            M=M+sparse(jjz,jjz,csnj(snj),NN,NN)+...
                sparse(jjz,SNJ3+1,b31snj(snj),NN,NN)+...
                sparse(jjz,SNJ3+2,b32snj(snj),NN,NN)+...
                sparse(jjz,SNJ3+3,b33snj(snj),NN,NN);

        end

        % E(J), J_{x,y,z}=j_{x,y,z}+sum(v_snj{x,y,z})
        tp=-1;
        jj=(0*SNJ1+1):(1*SNJ1); ii=0.*jj+SNJ3+1; M=M+sparse(ii,jj,tp,NN,NN);
        jj=(1*SNJ1+1):(2*SNJ1); ii=0.*jj+SNJ3+2; M=M+sparse(ii,jj,tp,NN,NN);
        jj=(2*SNJ1+1):(3*SNJ1); ii=0.*jj+SNJ3+3; M=M+sparse(ii,jj,tp,NN,NN);

        % jx(E), jy(E), jz(E)
        M=M+sparse(1*SNJ1,SNJ3+1,b11,NN,NN)+...
            sparse(1*SNJ1,SNJ3+2,b12,NN,NN)+...
            sparse(1*SNJ1,SNJ3+3,b13,NN,NN)+...
            sparse(2*SNJ1,SNJ3+1,b21,NN,NN)+...
            sparse(2*SNJ1,SNJ3+2,b22,NN,NN)+...
            sparse(2*SNJ1,SNJ3+3,b23,NN,NN)+...
            sparse(3*SNJ1,SNJ3+1,b31,NN,NN)+...
            sparse(3*SNJ1,SNJ3+2,b32,NN,NN)+...
            sparse(3*SNJ1,SNJ3+3,b33,NN,NN);

        % E(B)
        M=M+sparse(SNJ3+1,SNJ3+5,c2*kz,NN,NN)+...
            sparse(SNJ3+2,SNJ3+4,-c2*kz,NN,NN)+...
            sparse(SNJ3+2,SNJ3+6,c2*kx,NN,NN)+...
            sparse(SNJ3+3,SNJ3+5,-c2*kx,NN,NN);

        % B(E)
        M=M+sparse(SNJ3+4,SNJ3+2,-kz,NN,NN)+...
            sparse(SNJ3+5,SNJ3+1,kz,NN,NN)+...
            sparse(SNJ3+5,SNJ3+3,-kx,NN,NN)+...
            sparse(SNJ3+6,SNJ3+2,kx,NN,NN);
        if(sp==1)
            d=eigs(M,nw,wg); % solve the roots
        else
            d=eig(full(M));
        end
        omega=d;
        [wi,ind]=sort(imag(omega),'descend');
    %     [wr,ind]=sort(real(omega),'descend');
        w=omega(ind);

%         kk=[kk,k];
%         kxx=[kxx,kx]; kzz=[kzz,kz];
%         ww=[ww,w];
        
        WW(jz,jx)=w(1);
    end
end
runtime=toc;


%%
rhoci=max(abs(rhocs)); mi=max(ms); me=min(ms);
wci=min(abs(wcs));
% kk=kk*rhoci; kxx=kxx*rhoci; kzz=kzz*rhoci; 
ww=ww/wci;
wp=sqrt(wps(1)^2+wps(2)^2);
% kk=kk*sqrt(c2)/wps(3); kxx=kxx*sqrt(c2)/wps(3); kzz=kzz*sqrt(c2)/wps(3);
XX=XX*(sqrt(c2)/wps(3)); ZZ=ZZ*(sqrt(c2)/wps(3)); WW=WW/wci;
WW0=WW;

%%
% XX=XX/(sqrt(c2)/wps(3)); ZZ=ZZ/(sqrt(c2)/wps(3)); WW=WW*wci;
WW(imag(WW)<=1e-5)=NaN;
WW(WW==NaN)=0;

kkp=kzz;

pltc=[0.0  1.0  0.0
    1.0  0.0  0.0
    0.2  0.2  1.0
    0.8 0.8 0.0
    1.0  0.6  0.0
    0.9  0.0  0.9
    0.0  0.8  0.0
    0.0  0.8  0.8
    0.0  0.0  0.0
    0.6  0.0  0.0
    0.4  0.7  0.4 
    0.0  0.0  0.5 
    0.6  0.0  0.6 
    0.0  0.5  1.0];

plt=1;
if (plt==1)
    h=figure('unit','normalized','Position',[0.01 0.45 0.7 0.45],...
        'DefaultAxesFontSize',15);
	subplot(121);
%     xlim([min(kkp),max(kkp)]); 
    surf(ZZ,XX,real(WW)); axis tight;
    box on;
%     ylim([-0.1,0.1]); 
    xlabel('k_zc/\omega_p'); ylabel('k_xc/\omega_p'); zlabel('\omega/\omega_{ci}');
%     title(['(a) \omega_{pe}=',num2str(abs(wp/wcs(1))),...
%         ', m_i/m_e=',num2str(mi/me),', N=',num2str(N)]);
    hold on;
    
    subplot(122);
%     xlim([min(kkp),max(kkp)]); 
    surf(ZZ,XX,imag(WW)); axis tight;
    box on;
    xlabel('k_zc/\omega_p'); ylabel('k_xc/\omega_p'); zlabel('\gamma/\omega_{ci}');
%     title(['(b) J=',num2str(J),', NN=',num2str(NN)]);
%     title(['(b) NN=',num2str(NN),', ',num2str(runtime),'s']);
    hold on;
else
    axes('position',[0.17 0.65 0.23 0.23]); box on;
    xlim([min(kk),max(kk)]);
    ylim([0.03,0.5]);
    hold on;
end

if (plt==1)
    subplot(121);
end
for n=1:5
    plot([0,max(kkp)],[n,n],'c--','LineWidth',2);
    hold on;
end
% 
% kkkp=reshape(repmat(kkp,nw,1),1,[]);
% www=reshape(ww,1,[]);
% nkk=length(kkkp);
% jj=[];
% for j=1:nkk
%     % imag(www(j))/abs(real(www(j))-90)<-4.25e-2 || 
% %     if(imag(www(j))<-1.0e-0 || abs(real(www(j))-2.4*wA)>1.8*wA)
%     if(imag(www(j))<-1.0e1) %|| imag(www(j))>1.0e-0
%         jj=[jj,j];
%     end
% end
% www(jj)=[];
% kkkp(jj)=[];

% hold on;
% plot(kkkp,real(www),'r+','LineWidth',2);
% if (plt==1)
%     subplot(122); 
%     plot((kkkp),imag(www),'r+','LineWidth',2);
% end
% xlim([log(min(kkkp)),log(max(kkkp))]);


print('-dpng',['pdrk_em_wp',num2str(abs(wps(1)/wcs(1))),...
    '_N',num2str(N),'_J',num2str(J),'_sp',num2str(sp),...
    '_nw',num2str(nw),'_vs0',num2str(vs0(1)),'_lmdT_vs0_2d.png']);

