eigvfft-for-oliver.jl

## Functions for eigvfft notebook. 
## to use: include("./eigvfft.jl")

# Set environment
using Pkg

Pkg.add("LinearAlgebra")
using LinearAlgebra

Pkg.add("FFTW")
using FFTW

Pkg.add("LaTeXStrings")
using LaTeXStrings
Pkg.add("PyCall")
using PyCall

Pkg.add("PyPlot")
using PyPlot

# Pkg.add("Interact")
# using Interact

Pkg.add("SparseArrays")
using SparseArrays

## Function to create a 3d laplace operator matrix A
## with size given as input. First two dimensions have
## cyclic boundary conditions. Last dimension has homogeneous 
## Neumann boundary conditions. 

# Make 3-d PPN operator 
### function mkA_PPN(Nx,Ny,Nz)
### NN=Nx*Ny*Nz;
### A=zeros(NN,NN)
### # Modulo and offset for 1 based index
### MOD(i,n)=mod(i-1,n)+1
### OFF(i,j,k,ni,nj,nk)= (k-1)*ni*nj + (j-1)*ni + (i-1) + 1 
### for k=1:Nz
###  for j=1:Ny
###   for i=1:Nx
###    ic=i ; iw=MOD(i-1,Nx); ie=MOD(i+1,Nx)
###    jc=j ; js=MOD(j-1,Ny); jn=MOD(j+1,Ny)
###    kc=k ; ku=MOD(k-1,Nz); kd=MOD(k+1,Nz)
###    offc=OFF(i , j, k, Nx, Ny, Nz)
###    offw=OFF(iw, j, k, Nx, Ny, Nz)
###    offe=OFF(ie, j, k, Nx, Ny, Nz)
###    offs=OFF( i,js, k, Nx, Ny, Nz)
###    offn=OFF( i,jn, k, Nx, Ny, Nz)
###    offu=OFF( i, j,ku, Nx, Ny, Nz)
###    offd=OFF( i, j,kd, Nx, Ny, Nz)
###    A[offc,offc]=-6
###    A[offc,offw]= A[offc,offw]+1
###    A[offc,offe]= A[offc,offe]+1
###    A[offc,offs]= A[offc,offs]+1
###    A[offc,offn]= A[offc,offn]+1
###    if k == 1
###     A[offc,offu]= A[offc,offu]+0
###     A[offc,offc]= A[offc,offc]+1
###     A[offc,offd]= A[offc,offd]+1   
###    elseif k == Nz
###     A[offc,offu]= A[offc,offu]+1
###     A[offc,offc]= A[offc,offc]+1
###     A[offc,offd]= A[offc,offd]+0   
###    else
###     A[offc,offu]= A[offc,offu]+1
###     A[offc,offd]= A[offc,offd]+1 
###    end
###   end
###  end
### end
###
### # show(IOContext(stdout), "text/plain", Matrix(A))
### return A, Nx, Ny, Nz, NN
### end
        
function mkA_PPN(Nx,Ny,Nz)
 return mkA_PPN(Nx,Ny,Nz,Nx,Ny,Nz)
end

# Make 3-d PPN operator with domain lengths specified
function mkA_PPN(Nx,Ny,Nz,Lx,Ly,Lz)
 NN=Nx*Ny*Nz;
 A=zeros(NN,NN)
 dx=Lx/Nx;rdx2=1. / (dx^2.)
 dy=Ly/Ny;rdy2=1. / (dy^2.)
 dz=Lz/Nz;rdz2=1. / (dz^2.)
 # Modulo and offset for 1 based index
 MOD(i,n)=mod(i-1,n)+1
 OFF(i,j,k,ni,nj,nk)= (k-1)*ni*nj + (j-1)*ni + (i-1) + 1 
 for k=1:Nz
  for j=1:Ny
   for i=1:Nx
    ic=i ; iw=MOD(i-1,Nx); ie=MOD(i+1,Nx)
    jc=j ; js=MOD(j-1,Ny); jn=MOD(j+1,Ny)
    kc=k ; ku=MOD(k-1,Nz); kd=MOD(k+1,Nz)
    offc=OFF(i , j, k, Nx, Ny, Nz)
    offw=OFF(iw, j, k, Nx, Ny, Nz)
    offe=OFF(ie, j, k, Nx, Ny, Nz)
    offs=OFF( i,js, k, Nx, Ny, Nz)
    offn=OFF( i,jn, k, Nx, Ny, Nz)
    offu=OFF( i, j,ku, Nx, Ny, Nz)
    offd=OFF( i, j,kd, Nx, Ny, Nz)
    A[offc,offc]=-1. * ( 2. * rdz2 + 
                         2. * rdy2 + 
                         2. * rdz2 
                        )
    A[offc,offw]= A[offc,offw]+rdx2
    A[offc,offe]= A[offc,offe]+rdx2
    A[offc,offs]= A[offc,offs]+rdy2
    A[offc,offn]= A[offc,offn]+rdy2
    if k == 1
     A[offc,offu]= A[offc,offu]+0
     A[offc,offc]= A[offc,offc]+rdz2
     A[offc,offd]= A[offc,offd]+rdz2
    elseif k == Nz
     A[offc,offu]= A[offc,offu]+rdz2
     A[offc,offc]= A[offc,offc]+rdz2
     A[offc,offd]= A[offc,offd]+0   
    else
     A[offc,offu]= A[offc,offu]+rdz2
     A[offc,offd]= A[offc,offd]+rdz2
    end
   end
  end
 end

 # show(IOContext(stdout), "text/plain", Matrix(A))
 return A, Nx, Ny, Nz, NN
end
                

        
# Make a 3d cube plot
function plot3dcube(Nx,Ny,Nz)            
 Lx=Ly=Lz=10
 xp=LinRange(1,10,Nx+1);
 yp=LinRange(1,10,Ny+1);
 zp=LinRange(1,10,Nz+1);
 xc=fill(0,Nx+1);
 yc=fill(0,Ny+1);
 zc=fill(0,Nz+1);
 PyPlot.figure(figsize=(20.0,20.0)) 
 for zoff in zp
  for yoff in yp
   yc=fill(0,Nx+1);
   zc=fill(0,Nx+1);
   PyPlot.plot3D(xp,yc.+yoff,zc.+zoff,color=:gray); 
  end
 end
 for xoff in xp
  for yoff in yp
   yc=fill(0,Nz+1);
   xc=fill(0,Nz+1);
   PyPlot.plot3D(xc.+xoff,yc.+yoff,zp,color=:gray); 
  end
 end
 for zoff in zp
  for xoff in xp
   zc=fill(0,Ny+1);
   xc=fill(0,Ny+1);
   PyPlot.plot3D(xc.+xoff,yp,zc.+zoff,color=:gray); 
  end
 end
 ax = PyPlot.gca();
 ax[:axis]("off");
 
 PyPlot.text3D(Lx/2,Ly/2,Lz/2,L"\phi_{i,j,k}",FontSize=36,Weight="bold")

 PyPlot.text3D(Lx/2-Lx/Nx,Ly/2,Lz/2,L"\phi_{i-1,j,k}",FontSize=36,Weight="bold")
 PyPlot.text3D(Lx/2+Lx/Nx,Ly/2,Lz/2,L"\phi_{i+1,j,k}",FontSize=36,Weight="bold")
 
 PyPlot.text3D(Lx/2,Ly/2-Ly/Ny,Lz/2,L"\phi_{i,j-1,k}",FontSize=36,Weight="bold")
 PyPlot.text3D(Lx/2,Ly/2+Ly/Ny,Lz/2,L"\phi_{i,j+1,k}",FontSize=36,Weight="bold")

 PyPlot.text3D(Lx/2,Ly/2,Lz/2-Lz/Nz,L"\phi_{i,j,k-1}",FontSize=36,Weight="bold")
 PyPlot.text3D(Lx/2,Ly/2,Lz/2+Lz/Nz,L"\phi_{i,j,k+1}",FontSize=36,Weight="bold")

 PyPlot.plot3D([Lx/2-Lx/Nx, Lx/2+Lx/Nx],[Ly/2,Ly/2],[Lz/2,Lz/2],
    color=:black,LineWidth=3,LineStyle="--")
 PyPlot.plot3D([Lx/2, Lx/2],[Ly/2-Ly/Ny,Ly/2+Ly/Ny],[Lz/2,Lz/2],
    color=:black,LineWidth=3,LineStyle="--")
 PyPlot.plot3D([Lx/2, Lx/2],[Ly/2,Ly/2],[Lz/2-Lz/Nz,Lz/2+Lz/Nz],
    color=:black,LineWidth=3,LineStyle="--")
end
                
# Make 1-d N operator with variable grid spacing in Z specified
# To keep the operator symmetric we multiply ther righthand side by dZ
function mkA_N(dzArr::Array)
 N=length(dzArr)
 A=zeros(N,N)
 for k=1:N
  if k == 1
   upTerm=0
  else
   dK=dzArr[k]
   dKm1=dzArr[k-1]
   upTerm=2.0/(dK*dKm1+dK*dK);
   upTerm=2.0/(dKm1+dK);
  end
  if k == N
   dnTerm=0
  else
   dK=dzArr[k]
   dKp1=dzArr[k+1]
   dnTerm=2.0/(dK*dKp1+dK*dK);
   dnTerm=2.0/(dKp1+dK);
  end
  if k == 1
   A[k  ,k  ]=-dnTerm
   A[k  ,k+1]= dnTerm
  elseif k == N
   A[k  ,k-1]= upTerm
   A[k  ,k  ]=-upTerm
  else
   A[k  ,k-1]= upTerm
   A[k  ,k  ]=-dnTerm-upTerm
   A[k  ,k+1]= dnTerm
  end
 end
 return A
end
      
function tdsolve(ld,md,ud,rhs)
# Tridoagonal solve per Numerical Recipes, Press et. al 1992 (Sec  2.4 )
# ld[2:N  ] - lower diagonal
# md[1:N  ] - main diagonal
# ud[1:N-1] - upper diagonal
# phi       - solution vector
# rhs       - right hand side
 # Get length and allocate memory
 # return rhs
 N=length(rhs)
 phi=rhs.*typeof(rhs[1])(0)
 gamma=rhs.*typeof(rhs[1])(0)
 #       
 beta=md[1]
 phi[1]=rhs[1]/beta

 for j=2:N
  gamma[j]=ud[j-1]/beta
  beta=md[j]-ld[j]*gamma[j]
  # beta small should only happen on last element of forward pass for 
  # problem with zero eigenvalue. In that case algorithmn is still stable.
  if abs(beta) < 1.e-12
   break
  end
  phi[j]=(rhs[j]-ld[j]*phi[j-1])/beta
 end
 for j=1:N-1
  k=N-j
  phi[k]=phi[k]-gamma[k+1]*phi[k+1]
 end
 return phi
end
      
function mkwaves(N,L)
 scyc=zeros(N,1); sneu=zeros(N,1);
 for i in 1:N
  scyc[i]=(2*sin((i-1)*π/N)/(L/N)).^2
  sneu[i]=(2*sin((i-1)*π/(2*(N)))/(L/N)).^2 
 end   
 return scyc, sneu
end