eigvariableznohelm.jl

include("./eigvfft-for-oliver.jl");
IJulia.clear_output();

dz=ones(4,1)[:];dz[1]=4.;dz
dz=[1.,2.,3,4.]
# dz=[ones(1,10) ones(1,20)*2 ones(1,100)*5 ]

Az=mkA_N(dz);Nz=size(Az,1)
Ah,Nx,Ny,Nzz,NN=mkA_PPN(3,3,1);
Nh=size(Ah,1);
AhExp=kron(Ah,Matrix{Float32}(I,Nz,Nz))
AzExp=kron(Matrix{Float32}(I,Nh,Nh),Az)
A=AzExp+AhExp;
# Check A is symmetric
show(stdout,"text/plain",A[1:10,1:10]);println()

# Set up f as the integrated divergence of boundary flux terms, with zero flux at top and bottom
G=rand(Nz+1,Nx+1,Ny+1);
G[1,:,:]=G[end,:,:].=0; # zero flux top and bottom
G[:,1,:]=G[:,end,  :];  # periodic in X
G[:,:,1]=G[:,  :,end];  # periodic in Y
g_z=(G[1:end-1,:,:]-G[2:end,:,:])
g_x=(G[:,2:end,:]-G[:,1:end-1,:])
g_y=(G[:,:,2:end]-G[:,:,1:end-1])
divG=g_z[:,1:end-1,1:end-1]+g_x[1:end-1,:,1:end-1]+g_y[1:end-1,1:end-1,:]
# Multiply by dz needed on RHS so A is stays symmetric
for k=1:Nz; divG[k,:,:].*dz[k]; end
sum(divG)

# Solve FFT style and FFT + tridiag style
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

Lx=Nx;
Ly=Ny;
Lz=Nz;

# fz  = FFTW.r2r(f,FFTW.REDFT10,3)
# fxyz= FFTW.fft(fz,[1,2])

# fF1=FFTW.fft(reshape(f,Nx,Ny,Nz),[1,2]);
# fF2= FFTW.r2r(real.(fF1),FFTW.REDFT10)

fF1=FFTW.r2r(divG,FFTW.REDFT10,1)
fF2= FFTW.fft(fF1,[2,3])

####
fF22=FFTW.fft(divG,[2,3])
####
# display(fF22[:,2,1])

sxcyc, sxneu=mkwaves(Nx,Lx);sx=-sxcyc;
sycyc, syneu=mkwaves(Ny,Ly);sy=-sycyc;
szcyc, szneu=mkwaves(Nz,Lz);sz=-szneu;

fFi=fF2;

for i=1:Nx
 for j=1:Ny
  for k=1:Nz
   s=sx[i]+sy[j]+sz[k]
   if abs(s) < 1.e-12
    fFi[k,i,j]=0;
   else
    fFi[k,i,j]=fFi[k,i,j]./s
   end
  end
 end
end
    
#### BEGIN TD
ld=[0 diag(Az,-1)']'; 
md=diag(Az,0);
ud=[diag(Az,1)' 0]';
xx=complex.(zeros(Nz,Nx,Ny))
s=0
for i=1:Nx
 for j=1:Ny
  s=sx[i]+sy[j]
  if abs(s) < 1.e-12
   xx[:,i,j]=tdsolve(complex.(ld),complex.(md.+s),complex.(ud),complex.(fF22[:,i,j]))
  else
   xx[:,i,j]=tdsolve(complex.(ld),complex.(md.+s),complex.(ud),complex.(fF22[:,i,j]))
  end
 end
end
#### END TD
        
####
fiF11=real.(FFTW.ifft(xx,[2,3]))
####

fiF1=FFTW.ifft(fFi,[2,3])
fiF2=FFTW.r2r(real.(fiF1),FFTW.REDFT01,1)/(2*Nz);
fiF=reshape(fiF2,Nx*Ny*Nz);
# println("divG         ",divG[:])
# println("A*fiF        ",A*fiF)
println("A*fiF11      ",A*reshape(fiF11,Nx*Ny*Nz)-divG[:])