julia code

This commit is contained in:
Ian Jauslin 2022-05-18 09:57:44 +02:00
parent f21ab0d795
commit 03a3a0f880
4 changed files with 397 additions and 0 deletions

17
julia/driving.jl Normal file
View File

@ -0,0 +1,17 @@
# driving forces (g stands for k^\perp.g/(2i\pi|k|)
# a simple test driving force
function g_test(
kx::Int64,
ky::Int64
)
if(kx==2 && ky==-1)
return 0.5+sqrt(3)/2*1im
elseif (kx==-2 && ky==1)
return 0.5-sqrt(3)/2*1im
else
return 0.
end
end

150
julia/main.jl Normal file
View File

@ -0,0 +1,150 @@
using Printf
using FFTW
include("navier_stokes.jl")
include("driving.jl")
include("print.jl")
function main()
## defaults
delta=0.0001220703125
nsteps=10000000
nu=0.00048828125
print_freq=1000
K1=16
K2=K1
N1=4*K1+1
N2=4*K2+1
# read cli arguments
(params,driving,command)=read_args(ARGS)
# whether N was set in parameters
setN1=false
setN2=false
# read params
if params!=""
for param in split(params,";")
terms=split(param,"=")
if length(terms)!=2
print(stderr,"error: could not read parameter '",param,"'.\n")
exit(-1)
end
lhs=string(terms[1])
rhs=string(terms[2])
if lhs=="K1"
K1=parse(Int64,rhs)
elseif lhs=="K2"
K2=parse(Int64,rhs)
elseif lhs=="K"
K1=parse(Int64,rhs)
K2=K1
elseif lhs=="N1"
N1=parse(Int64,rhs)
setN1=true
elseif lhs=="N2"
N2=parse(Int64,rhs)
setN2=true
elseif lhs=="N"
N1=parse(Int64,rhs)
N2=N1
setN1=true
setN2=true
elseif lhs=="nsteps"
nsteps=parse(Int64,rhs)
elseif lhs=="nu"
nu=parse(Float64,rhs)
elseif lhs=="delta"
delta=parse(Float64,rhs)
elseif lhs=="print_freq"
print_freq=parse(Int64,rhs)
else
print(stderr,"unrecognized parameter '",lhs,"'.\n")
exit(-1)
end
end
end
## read driving
if driving=="test"
g=(kx,ky)->g_test(kx,ky)
end
## set parameters
if ! setN1
N1=4*K1+1
end
if ! setN2
N2=4*K2+1
end
## run command
if command=="uk"
navier_stokes_uk(nsteps,nu,g,N1,N2,K1,K2,delta,print_freq)
elseif command=="alpha"
navier_stokes_alpha(nsteps,nu,g,N1,N2,K1,K2,delta,print_freq)
else
print(stderr,"unrecognized command '",command,"'.\n")
exit(-1)
end
end
# read cli arguments
function read_args(ARGS)
# flag
flag=""
# output strings
command=""
params=""
# default driving force
driving="test"
# loop over arguments
for arg in ARGS
# argument that starts with a dash
if arg[1]=='-'
# go through characters after dash
for char in arg[2:length(arg)]
# set params
if char=='p'
# raise flag
flag="params"
elseif char=='g'
# raise flag
flag="driving"
else
print_usage()
exit(-1)
end
end
# arguments that do not start with a dash
else
if flag=="params"
params=arg
elseif flag=="driving"
driving=arg
else
command=arg
end
# reset flag
flag=""
end
end
if command==""
print_usage()
exit(-1)
end
return (params,driving,command)
end
# usage
function print_usage()
print(stderr,"usage: nstrophy [-p params] [-g driving] <command>\n")
end
main()

215
julia/navier_stokes.jl Normal file
View File

@ -0,0 +1,215 @@
# solve Navier-Stokes
function navier_stokes_uk(
nsteps::Int64,
nu::Float64,
g::Function,
N1::Int64,
N2::Int64,
K1::Int64,
K2::Int64,
delta::Float64,
print_freq::Int64
)
# init
u=ns_init(K1,K2)
for i in 1:nsteps
u=ns_RK4(u,nu,g,N1,N2,K1,K2,delta)
if i % print_freq==0
@printf("% .15e ",delta*i)
print_u(u,N1,N2,K1,K2)
# print to stderr
@printf(stderr,"% .8e ",delta*i)
for k1 in -K1:K1
for k2 in -K2:K2
index=indices_from_momentum(k1,2*K1+1)*(2*K2+1)+indices_from_momentum(k2,2*K2+1)+1
if (k1^2+k2^2<=1)
@printf(stderr,"% .8e % .8e ",real(u[index]),imag(u[index]))
end
end
end
@printf(stderr,"\n")
end
end
end
# compute the enstrophy
function navier_stokes_alpha(
nsteps::Int64,
nu::Float64,
g::Function,
N1::Int64,
N2::Int64,
K1::Int64,
K2::Int64,
delta::Float64,
print_freq::Int64
)
# init
u=ns_init(K1,K2)
for i in 1:nsteps
u=ns_RK4(u,nu,g,N1,N2,K1,K2,delta)
if i % print_freq==0
alpha=ns_alpha(u,g,K1,K2)
@printf("% .15e % .15e % .15e\n",i*delta,real(alpha),imag(alpha))
@printf(stderr,"% .15e % .15e % .15e\n",i*delta,real(alpha),imag(alpha))
end
end
end
# initialize u
function ns_init(
K1::Int64,
K2::Int64
)
u=zeros(Complex{Float64},(2*K1+1)*(2*K2+1))
for k1 in -K1:K1
for k2 in -K2:K2
index=indices_from_momentum(k1,2*K1+1)*(2*K2+1)+indices_from_momentum(k2,2*K2+1)+1
u[index]=exp(-sqrt(k1^2+k2^2))
end
end
return u
end
# Navier-Stokes equation (right side)
function ns_rhs(
u::Array{Complex{Float64},1},
nu::Float64,
g::Function,
N1::Int64,
N2::Int64,
K1::Int64,
K2::Int64
)
# compute the T term
out=ns_T(u,N1,N2,K1,K2)
# add the rest of the terms
for k1 in -K1:K1
for k2 in -K2:K2
index=indices_from_momentum(k1,2*K1+1)*(2*K2+1)+indices_from_momentum(k2,2*K2+1)+1
out[index]=-(4*pi^2*nu*(k1^2+k2^2))*u[index]+g(k1,k2)+(k1^2+k2^2>0 ? 4*pi^2/sqrt(k1^2+k2^2)*out[index] : 0)
end
end
return out
end
# compute k\in K from n\in N
function momentum_from_indices(
i::Int64,
K::Int64,
N::Int64,
)
if i<=K
k=i
elseif i>=N-K
k=i-N
else
k=i
end
return(k)
end
# inverse
function indices_from_momentum(
k::Int64,
N::Int64
)
if k>=0
return k
else
return N+k
end
end
# T term in Navier-Stokes
function ns_T(
u::Array{Complex{Float64},1},
N1::Int64,
N2::Int64,
K1::Int64,
K2::Int64
)
Fx1=zeros(Complex{Float64},N1,N2)
Fy1=zeros(Complex{Float64},N1,N2)
Fx2=zeros(Complex{Float64},N1,N2)
Fy2=zeros(Complex{Float64},N1,N2)
for k1 in -K1:K1
for k2 in -K2:K2
i1=indices_from_momentum(k1,N1)+1
i2=indices_from_momentum(k2,N2)+1
index=indices_from_momentum(k1,2*K1+1)*(2*K2+1)+indices_from_momentum(k2,2*K2+1)+1
# no need to divide by N: ifft does the division already
Fx1[i1,i2]=(k1==0 ? 0. : k1/sqrt(k1^2+k2^2)*u[index])
Fx2[i1,i2]=(k1==0 ? 0. : k1*sqrt(k1^2+k2^2)*u[index])
Fy1[i1,i2]=(k2==0 ? 0. : k2/sqrt(k1^2+k2^2)*u[index])
Fy2[i1,i2]=(k2==0 ? 0. : k2*sqrt(k1^2+k2^2)*u[index])
end
end
fft!(Fx1)
fft!(Fx2)
fft!(Fy1)
fft!(Fy2)
Fx1=Fx1.*Fy2.-Fx2.*Fy1
ifft!(Fx1)
out=zeros(Complex{Float64},(2*K1+1)*(2*K2+1))
for k1 in -K1:K1
for k2 in -K2:K2
i1=indices_from_momentum(k1,N1)+1
i2=indices_from_momentum(k2,N2)+1
index=indices_from_momentum(k1,2*K1+1)*(2*K2+1)+indices_from_momentum(k2,2*K2+1)+1
out[index]=Fx1[i1,i2]
end
end
return out
end
# Runge-Kutta step to solve Navier-Stokes
function ns_RK4(
u::Array{Complex{Float64},1},
nu::Float64,
g::Function,
N1::Int64,
N2::Int64,
K1::Int64,
K2::Int64,
delta::Float64
)
k1=delta*ns_rhs(u,nu,g,N1,N2,K1,K2)
k2=delta*ns_rhs(u.+(k1./2),nu,g,N1,N2,K1,K2)
k3=delta*ns_rhs(u.+(k2./2),nu,g,N1,N2,K1,K2)
k4=delta*ns_rhs(u.+k3,nu,g,N1,N2,K1,K2)
return u.+(k1+2 .*k2+2 .*k3.+k4)./6
end
# enstrophy
function ns_alpha(
u::Array{Complex{Float64},1},
g::Function,
K1::Int64,
K2::Int64
)
denom=0. *1im
num=0. *1im
for k1 in -K1:K1
for k2 in -K2:K2
index=indices_from_momentum(k1,2*K1+1)*(2*K2+1)+indices_from_momentum(k2,2*K2+1)+1
denom+=(k1^2+k2^2)^2*u[index]*conj(u[index])*(1. +(k2!=0 ? k1^2/k2^2 : 0.))
num+=(k1^2+k2^2)*u[index]*conj(g(k1,k2))*(1. +(k2!=0 ? k1^2/k2^2 : 0.))
end
end
return num/denom
end

15
julia/print.jl Normal file
View File

@ -0,0 +1,15 @@
function print_u(
u::Array{Complex{Float64},1},
N1::Int64,
N2::Int64,
K1::Int64,
K2::Int64
)
for k1 in -K1:K1
for k2 in -K2:K2
index=indices_from_momentum(k1,2*K1+1)*(K2+1)+indices_from_momentum(k2,2*K2+1)+1
@printf("% .15e % .15e ",real(u[index]),imag(u[index]))
end
end
print("\n")
end