(*Dispersion relation plotter that employs the Discretized Argument Theorem (DART) first proposed by Choueiri. *)
(*Coded in Mathematica by Benjamin Jorns, Electric Propulsion and Plasmadynamics Lab, Princeton University.  email: bjorns@princeton.edu *)
 (*Last Update:  6/12/08*)

(*Rutine to display the 'vortex' plot indicating pole locations*)
rtgraph[xmin_,xmax_,ymin_,ymax_, wmin_,N1_]:=(
N2=N1; (*Number of grid points on a side, N2*N1 = total number of boxes in vortex plot*)
Δx =(xmax-xmin)/(N1-1);
Δy =(ymax-ymin)/(N2-1);
Δw=Min[Δx,Δx]/2;
h[z_]:=f[z,wmin];
xarr=Table[xmin+Δx*(i-1/2),{i,1,(N1-1)},{j,1,N2-1}];  (*Build grid for vortex plot*)
yarr=Table[ymin+Δy*(j-1/2),{i,1,(N1-1)},{j,1,N2-1}];
zarr=xarr+I*yarr;
I0=Im[Log[h[zarr]]];     (*Plot argument of function at center of each box*)
I1=
Table[Line[{{xarr[[i,j]],yarr[[i,j]]}-Δw*{Cos[I0[[i,j]]],Sin[I0[[i,j]]]},{xarr[[i,j]],yarr[[i,j]]}+Δw*{Cos[I0[[i,j]]],Sin[I0[[i,j]]]}}],{i,1,(N1-1)},{j,1,N2-1}];
g12=Graphics[I1, Axes->True, AxesLabel->{Style["Re[k]",Blue,Bold],Style["Im[k]",Red,Bold]},PlotRange->{{xmin,xmax},{ymin,ymax}}];
Return[g12]);

(*Routine that applies DART to an initial grid and then traces the selected roots*)
dart[wmin_,wmax_, xmin_,xmax_,ymin_,ymax_]:=(t1=0;x=0;wrt=1;
crp=0;(*initialization constants for graphic display*)
Nk=Round[(wmax-wmin)*2/.025]*1.; (*Number of steps *)
Δα=(wmax-wmin)/(Nk-1);  (*step size*)
int1=rt[xmin,xmax,ymin,ymax,wmin];  (*Calls algorithm that applies DART to selected area*)
For[l=1,l≤ Nk,l++,If[reset==1, Break[]]; h[k_]:=f[k,(l-1)*Δα+wmin];  If[l==1, rts=int1;Sort[rts,Re[#2[[1]]]>Re[#1[[1]]]&]];
If[intcount≥5 && l≠1,  intlength=Length[totarr];  (*intcount tracks number of steps*)For[j=2,j≤ intlength,j++, Quiet[tp=Interpolation[Table[{totarr[[1,n]],totarr[[j,n]]},{n,totcount-4,totcount-1}]]]; (*applies Lagrangian interpolation to guess at next root location*)
Quiet[rts[[j-1,1]]=tp[(l-1)*Δα+wmin]]]];
If[intcount>=2, For[j=2,j≤ Length[totarr],j++, rts[[j-1,2]]=4*((Re[totarr[[j,totcount]]]-Re[totarr[[j,totcount-1]]])^2+ (Im[totarr[[j,totcount]]]-Im[totarr[[j,totcount-1]]])^2)^(1/2); (*creates a search window around the projected route for the tracer algorithm.  The window is equal to 4 times the distance between the current root and the previous root*)rts[[j-1,3]]=4*((Re[totarr[[j,totcount]]]-Re[totarr[[j,totcount-1]]])^2+ (Im[totarr[[j,totcount]]]-Im[totarr[[j,totcount-1]]])^2)^(1/2)]]; 
If[l>1, intres=tracer[rts]];  (*after incriminating the function a step, this calls the tracer that looks for a root in the neighborhood of the previously detected root*)
If[l>1 && Norm[intres]≠ 0,rts=intres; test=1]; If[l>1 && Norm[intres]== 0 && test≠1,  crp=1;x=1 ;(Subscript[ω, min])=(Subscript[ω, min])+(l-1)*Δα; sollent=Length[solarr]; dx=xh-xl; dy=yh-yl; xl= solarr[[sollent,1]]-dx/2; xh=solarr[[sollent,1]]+dx/2; yl=solarr[[sollent,2]]-dy/2; yh=solarr[[sollent,2]]+dy/2;  Break[]]; (*breaks the loop if no solution is returned by the tracer algorithm*)
intcount++;
totcount++;
If[l>1, Sort[rts,Re[#2[[1]]]>Re[#1[[1]]]&]]; 
test=0;
temp1=Transpose[{Prepend[rts[[All,1]],(l-1)*Δα+wmin ]}];
soll=Length[temp1];
If[(Norm[grarr]≠0 || Norm[grarri]≠0),totarr=Join[totarr,temp1,2];  (*plots solutions in real time*)solarr=Join[solarr,Table[{Re[temp1[[i,1]]],Im[temp1[[i,1]]]},{i,2,soll}]];  grarr=Join[grarr, Table[{Re[temp1[[i,1]]],temp1[[1,1]]},{i,2,soll}]]; grarri=Join[grarri, Table[{Im[temp1[[i,1]]],temp1[[1,1]]},{i,2,soll}]]];
If[ Norm[grarr]==0 && Norm[grarri]==0, totarr=temp1; solarr=Table[{Re[temp1[[i,1]]],Im[temp1[[i,1]]]},{i,2,soll}]; grarr=Table[{Re[temp1[[i,1]]],temp1[[1,1]]},{i,2,soll}];  grarri=Table[{Im[temp1[[i,1]]],temp1[[1,1]]},{i,2,soll}]];
]; count=0;app2=False;app3=True; (*update definitions for graphic display*)
);


(*Initial Root Finder that employs DART*)
rt[Xmin_,Xmax_,Ymin_,Ymax_, wmin_]:=(
g[xmin_,xmax_,ymin_,ymax_,tes_]:=(
N1=6.; (*Length of grid in X direction*)
N2=6.; (*Length of grid in Y direction*)
Δx =Abs[(xmax-xmin)/(N1-1)];
Δy =Abs[(ymax-ymin)/(N2-1)];
xarr=Table[xmin+Δx*(i-1/2),{i,1,(N1-1)},{j,1,N2-1}];
yarr=Table[ymin+Δy*(j-1/2),{i,1,(N1-1)},{j,1,N2-1}];
zarr=xarr+I*yarr;
h[z_]:=f[z,wmin];
F[z_,t_]:=z^t*N[h'[z]]/h[z](*(f[z+γ*Δx]-f[z-γ*Δx])/(2*γ*Δx*f[z]);*);
I0=Table[0,{i,1,(N1-1)},{j,1,(N2-1)}];
I1=Table[0,{i,1,(N1-1)},{j,1,(N2-1)}];
For[i=1,i≤(N1-1),i++, For[j=1,j≤(N2-1),j++, 
Quiet[I0[[i,j]]=(1/(2*Pi*I))*NIntegrate[F[z,0], {z,zarr[[i,j]]-Δy/2*I+Δx/2, zarr[[i,j]]+Δy/2*I+Δx/2, zarr[[i,j]]+Δy/2*I-Δx/2,zarr[[i,j]]-Δy/2*I-Δx/2,zarr[[i,j]]-Δy/2*I+Δx/2}, AccuracyGoal->4, MaxPoints->100]];
Quiet[I1[[i,j]]=(1/(2*Pi*I))*NIntegrate[F[z,1], {z,zarr[[i,j]]-Δy/2*I+Δx/2, zarr[[i,j]]+Δy/2*I+Δx/2, zarr[[i,j]]+Δy/2*I-Δx/2,zarr[[i,j]]-Δy/2*I-Δx/2,zarr[[i,j]]-Δy/2*I+Δx/2},AccuracyGoal->4,MaxPoints->100]]]];
(*I0 and I1 are  matrices that contain the characteristic integrals, (Subscript[I, 0]) and (Subscript[I, 1]), for each grid square*)
ϵ=10^(-2); (* the epsilons are numerical tolerances for detecting the presence of a root*)
(Subscript[ϵ, 0])=10^(-2);
(Subscript[ϵ, 1])=10^(-2);

(*tests for single roots*)
count = 1;
N4 =1;
For[i=1,i<N1,i++,For[j=1,j<N2,j++,If[Abs[1-Re[I0[[i,j]]]]<ϵ,N4=N4+1]]];
psing=Table[0,{i,1,N4-1}];
For[i=1,i<N1,i++,For[j=1,j<N2,j++,If[Abs[1-Re[I0[[i,j]]]]<ϵ,psing[[count]]={I1[[i,j]],Δx,Δy} ;count = count+1]]];

(*tests for suspect cells*)
count2 =1;
N5=1;
For[i=1,i<N1,i++,For[j=1,j<N2,j++,If[Abs[1-Re[I0[[i,j]]]]>(Subscript[ϵ, 1]) && (Subscript[ϵ, 1])<Abs[Re[I0[[i,j]]]],N5=N5+1]]];
M=Table[{0,0},{i,1,N5-1}];
For[i=1,i<N1,i++,
For[j=1,j<N2,j++,If[Abs[1-Re[I0[[i,j]]]]>(Subscript[ϵ, 1]) && (Subscript[ϵ, 1])<Abs[ Re[I0[[i,j]]]],M[[count2]]={i,j};count2 = count2+1]]];

(*tests for degenerate roots*)
N6=1;
For[n=1,n<(N5),n++,If[Abs[h[I1[[M[[n,1]],M[[n,2]]]]/Re[I0[[M[[n,1]],M[[n,2]]]]]]]<(Subscript[ϵ, 0]), N6=N6+1]];
pdeg=Table[{0,0},{i,1,N6-1}];
For[n=1,n<(N5),n++,If[Abs[h[I1[[M[[n,1]],M[[n,2]]]]/Re[I0[[M[[n,1]],M[[n,2]]]]]]]<(Subscript[ϵ, 0]), pdeg[[n]]={I1[[M[[n,1]],M[[n,2]]]]/Re[I0[[M[[n,1]],M[[n,2]]]]],Δx,Δy}; M[[n]]={0,0}]];

(*establishes array to contain information on suspsect cells*)
N7=(N5-1)-(N6-1);
count4=1;
If[N7==0,psusp=Table[{0,0,0,0},{i,1,1}]];
If[N7≠0,psusp=Table[{0,0,0,0},{i,1,N7}]];

(*Identifies suspect cells for another iteration of the DART routine and stores information on the boundaries of these cells*)
For[n=1,n<(N5),n++,If[M[[n]]≠{0,0},psusp[[count4,1]]=(xarr[[M[[n,1]],M[[n,2]]]]-Δx/2);psusp[[count4,2]]=(xarr[[M[[n,1]],M[[n,2]] ]]+Δx/2);psusp[[count4,3]]=(yarr[[M[[n,1]],M[[n,2]]]]-Δy/2);psusp[[count4,4]]=(yarr[[M[[n,1]],M[[n,2]] ]]+Δy/2);Print["Resolving suspect cells:  redrawing grid"];count4=count4+1 ]];
temp=Join[psing,pdeg];
If[tes==1,ptot=temp];
If[tes>1,ptot=Join[ptot,temp]];
Return[psusp]);
L=g[Xmin,Xmax,Ymin,Ymax,1];
lenth=Length[L];
n1=1;
(*series of steps redrawing a 6 x 6 grid over each suspect cell and reapplying the DART algorithm*)
While[ n1≤lenth && Norm[L]≠0, If[n1==1,  G=g[L[[n1,1]],L[[n1,2]],L[[n1,3]],L[[n1,4]],n1+1] ];If[1<n1<lenth ,  tat1= g[L[[n1,1]],L[[n1,2]],L[[n1,3]],L[[n1,4]],n1]; norm3=Norm[tat1];  If[norm3≠0 && Norm[G]≠0, G=Join[G, tat1]]; If[norm3≠0 && Norm[G]==0, G=tat1] ];If[n1≥lenth,L=G;  lenth=Length[L];  n1=0]; n1++ ] ; 
Return[ptot]);

(* This routine is called after DART has been applied to the entire grid and roots found.   It calculates (Subscript[I, \[Placeholder]]) and (Subscript[I, 1]) for a single cell suspected of containing the root after the function is incremented*)
tracer[tes_]:=(
P=tes; (*the array of cells in the neighborhood of where roots were previously found and where they are now suspected to be*)
N1=Length[P];
N2=1.;
Δx =Min[P[[All,2]]]/N2;
Δy =Min[P[[All,3]]]/N2;
F[z_,t_]:=z^t*N[h'[z]]/h[z];
I0=Table[0,{i,1,N1}];
I1=Table[0,{i,1,N1}];
For[i=1,i≤N1,i++, 
Quiet[I0[[i]]=(1/(2*Pi*I))*NIntegrate[F[z,0], {z,P[[i,1]]-Δy/2*I+Δx/2, P[[i,1]]+Δy/2*I+Δx/2, P[[i,1]]+Δy/2*I-Δx/2,P[[i,1]]-Δy/2*I-Δx/2,P[[i,1]]-Δy/2*I+Δx/2}, AccuracyGoal->4, MaxPoints->20]];
Quiet[I1[[i]]=(1/(2*Pi*I))*NIntegrate[F[z,1], {z,P[[i,1]]-Δy/2*I+Δx/2, P[[i,1]]+Δy/2*I+Δx/2, P[[i,1]]+Δy/2*I-Δx/2,P[[i,1]]-Δy/2*I-Δx/2,P[[i,1]]-Δy/2*I+Δx/2}, AccuracyGoal-> 4, MaxPoints->20]]]; 
ϵ=10^(-2);
(Subscript[ϵ, 0])=10^(-2);
(Subscript[ϵ, 1])=10^(-2);

count = 1;
N4 =1;
For[i=1,i≤N1,i++,If[Abs[1-Re[I0[[i]]]]<ϵ,N4=N4+1]];
If[N4-1==0, psing=Table[{0,0,0,0},{i,1,1}]];
If[N4-1≠ 0, psing=Table[0,{i,1,N4-1}]];
For[i=1,i≤ N1,i++,If[Abs[1-Re[I0[[i]]]]<ϵ,psing[[i]]={I1[[i]],Δx,Δy};count = count+1]];


count2 =1;
N5=1;
For[i=1,i≤ N1,i++,If[Abs[1-Re[I0[[i]]]]>ϵ && ϵ<Abs[Re[I0[[i]]]],N5=N5+1]];
If[N5-1==0 ,M=Table[0,{i,1,1}]];
If[N5-1≠ 0,M=Table[0,{i,1,N5-1}]];
For[i=1,i≤ N1,i++,If[Abs[1-Re[I0[[i]]]]>ϵ && ϵ<Abs[ Re[I0[[i]]]],M[[count2]]=i;count2 = count2+1]];


N6=1;
For[n=1,n<N5,n++,If[Abs[h[I1[[ M[[n]] ]]/Re[I0[[ M[[n]] ]] ]]]<(Subscript[ϵ, 0]), N6=N6+1]];
If[N6-1==0 && N4-1==0,pdeg=Table[{0,0,0,0},{i,1,N6-1}]];
If[N6-1==0 && N4-1 ≠ 0,pdeg=Table[{0,0,0,0},{i,1,N6-1}]];
If[N6-1≠0, pdeg=Table[{0,0,0,0},{i,1,N6-1}]];
If[N6-1≠0&& N4-1==0, psing=Table[0,{i,1,N4-1}]];
For[n=1,n<N5,n++,If[Abs[h[I1[[ M[[n]] ]]/Re[I0[[ M[[n]] ]] ]]]<(Subscript[ϵ, 0]), pdeg[[n]]={I1[[M[[n]]]]/Re[I0[[M[[n]]]]],Δx,Δy}; M[[n]]=0]];


N7=(N5-1)-(N6-1);
count4=1;
If[N7==0,psusp=Table[{0,0,0,0},{i,1,1}]];
If[N7≠0, psusp=Table[{0,0,0,0},{i,1,N7}]];
For[n=1,n<(N5),n++,If[Norm[M[[n]]]≠0,psusp[[count4,1]]=(Re[P[[n,1]]]-Δx/2);psusp[[count4,2]]=(Re[P[[n,1]]]+Δx/2);psusp[[count4,3]]=(Im[P[[n,1]]]-Δy/2);psusp[[count4,4]]=(Im[P[[n,1]]]+Δy/2);count4=count4+1 ]];

temp=Join[psing,pdeg];
If[Norm[psusp]≠0, temp=Table[{0,0,0,0},{i,1,1}]];
Return[temp]);

(*small routine to display the cells where DART located roots*)
disproot[sols_]:=
(lenth1=Length[sols]; 
I1=
Table[Graphics[{Red,Thickness[1],Opacity[0.3],Rectangle[{Re[sols[[i,1]]]-sols[[i,2]]/2,Im[sols[[i,1]]]-sols[[i,3]]/2},{Re[sols[[i,1]]]+sols[[i,2]]/2,Im[sols[[i,1]]]+sols[[i,3]]/2}]}],{i,1,lenth1}];count=0;
Return[I1]);

Needs["PlotLegends`"]

f[k_,w_]:=k-w; 
xlow=-10.1; (*initial window*)
xhigh=10;
ylow=-10.1;
yhigh=10;

xl=-10.1; (*initial grid boundaries*)
xh=10;
yl=-10.1;
yh=10;

(*the following are a number of initializations for the above routines and graphical display*)
count=0;
x=1;
t1=1;
(Subscript[ω, min])=2.1; 
(Subscript[ω, m])=2.1;
(Subscript[ω, max]) =5;
crp=0;
app=False;
app2=True;
app3=True;
crp2=1;
t2=1;
solarr={0,0};
grarr={0,0};
grarri={0,0};
wrt=0;
sc=100;
dmark=0;
intcount=0;
totcount=0;
 tag=RandomReal[];

(*Graphic display*)
CreateWindow[DocumentNotebook[Panel[Column[{Grid[{{Style["Function f(k,ω)", Bold],Quiet[InputField[Dynamic["function"]]],Button["Enter", dmark=1;func=RandomReal[]] , , , , , , , ,,,,,,,DisplayForm[ButtonBox["Save to File",Evaluator->Automatic,ButtonFunction->(Export[SystemDialogInput["FileSave", "dispersiondata.txt"],Join[Transpose[{SetAccuracy[grarr[[All,2]],5]}],Transpose[{SetAccuracy[grarr[[All,1]],5]}],Transpose[{SetAccuracy[grarri[[All,1]],5]}],2],"Table"]&),ButtonFrame->"DialogBox",ImageSize->100,ButtonFrame->"DialogBox", Enabled->True]]DisplayForm[ButtonBox["Trace",Evaluator->Automatic,ButtonFunction->(dart[(Subscript[ω, min]),(Subscript[ω, max]),xl,xh,yl,yh]&),ButtonFrame->"DialogBox",ImageSize->100,ButtonFrame->"DialogBox", Enabled->True]]Button["Reset",reset=0;crp=0;x=1;wrt=0;t1=1;grarr={0,0};t2=1;totcount=0;intcount=0; count=0;
grarri={0,0}; solarr={0,0};xl=xlow; xh=xhigh;yl=ylow;yh=yhigh; (Subscript[ω, min])= (Subscript[ω, m]);h[k_]:=f[k,(Subscript[ω, min])];   app=False; app2=True; app3=True;  tag=RandomReal[], ImageSize->100]}}],Grid[{{Style["ω range",Bold],, InputField[Dynamic[(Subscript[ω, min])], FieldSize->Tiny],"-", InputField[Dynamic[(Subscript[ω, max])], FieldSize->Tiny] , , , , , , , ,,,,,,,,,,,,,,,Button["Stop", reset=1, ImageSize->100]Button["Leave trace and reset", crp=0;reset=0;wrt=0;intres=0;intcount=0;t2=1;
 (Subscript[ω, min])= (Subscript[ω, m]);h[k_]:=f[k,(Subscript[ω, min])];x=1; count=0; xl=xlow; xh=xhigh;yl=ylow;yh=yhigh; app=False; app2=True; app3=True; t1=1; tag=RandomReal[],ImageSize->200]}}],Row[{" "}],Row[{" "}],Row[{Column[{Button["+", xl=(xh+xl)/2.-Max[(xh-xl)/2, (yh-yl)/2]*.8; yl=(yl+yh)/2.-Max[(xh-xl)/2, (yh-yl)/2]*.8; xh=(xh+xl)/2.+Max[(xh-xl)/2, (yh-yl)/2]*.8; yh=(yl+yh)/2.+Max[(xh-xl)/2, (yh-yl)/2]*.8, ImageSize->Medium],Button["-", xl=(xh+xl)/2.-Max[(xh-xl)/2, (yh-yl)/2]*1.2; yl=(yl+yh)/2.-Max[(xh-xl)/2, (yh-yl)/2]*1.2; xh=(xh+xl)/2.+Max[(xh-xl)/2, (yh-yl)/2]*1.2; yh=(yl+yh)/2.+Max[(xh-xl)/2, (yh-yl)/2]*1.2, ImageSize->Medium]}],Grid[{{Style["Argument Plot", Bold,Larger],Style["Frequency vs. wave number",Bold,Larger]},{DynamicModule[{p1={-1,-1}, p2={-1,-1}},
p1=Dynamic[MousePosition["Graphics"]];
p2=Dynamic[MousePosition["Graphics"]];
EventHandler[Dynamic[If[x==1 ,If[crp==1, Print["Lost the root"]]; wmin=(Subscript[ω, min]);wmax=(Subscript[ω, max]); Dynamic[If[x==1,Refresh[Show[{rtgraph[xl,xh,yl,yh, (Subscript[ω, min]),40.],Graphics[{Red,Thickness[1], Opacity[0.3], Rectangle[Dynamic[p1],p2]}]}, ImageSize->Medium],TrackedSymbols->{p1,p2,x,crp,func,tag, xl,yh,yl,yh}]], SynchronousUpdating->False],Refresh[Show[{ListPlot[solarr,  PlotRange->{{xl,xh},{yl,yh}}, PlotStyle->{Purple,Thick},  AxesLabel->{Style["Re[k]",Blue,Bold],Style["Im[k]",Red,Bold]}]}, ImageSize->Medium],TrackedSymbols->{x,func,l}]]],{"MouseUp":> ( If[x==1,count++;If[dmark==0,count=0; dmark=1]; If[count<2, p1=MousePosition["Graphics"]];If[count≥ 2 ,p2=MousePosition["Graphics"];xl=Min[p1[[1]],p2[[1]]];xh=Max[p1[[1]],p2[[1]]]; yl=Min[p1[[2]],p2[[2]]];yh=Max[p1[[2]],p2[[2]]]; tag= RandomReal[];
p1=Dynamic[MousePosition["Graphics"]];
p2=Dynamic[MousePosition["Graphics"]];count=0] ])}]],Dynamic[ If[wrt==0, (Subscript[ω, m])=(Subscript[ω, min])];Refresh[Dynamic[ShowLegend[Show[{ListPlot[grarr,  PlotRange->All, PlotStyle->Thick,  AxesLabel->{Style["k",Bold],Style["ω",Bold]}],ListPlot[grarri, PlotRange->All, PlotStyle->{Red, Thick},  AxesLabel->{Style["k",Bold,Larger],Style["ω",Bold,Larger]}]},ImageSize->Medium, PlotRange->{{Min[xl,yl],Max[xh,yh]},{(Subscript[ω, m])-(Subscript[ω, m])*.2,(Subscript[ω, max])+(Subscript[ω, max])*.2}}],{{{Graphics[{Blue,Line[{{0,0},{4,0}}]}],Style["Real",Black, Bold]},{Graphics[{Red,Line[{{0,0},{4,0}}]}],Style["Imaginary",Black,Bold]}},LegendPosition->{.25,.5},LegendTextSpace->2,LegendSize->{0.8,0.2}, LegendShadow->None}]],TrackedSymbols->{x,xh}] ]}}, Dividers->Center]}]}]]],WindowSize->All, WindowTitle->"DART"]