julia code

4 files changed, 397 insertions, 0 deletions

diff --git a/julia/driving.jl b/julia/driving.jl
new file mode 100644
index 0000000..657f252
--- /dev/null
+++ b/julia/driving.jl
@@ -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
+
+
diff --git a/julia/main.jl b/julia/main.jl
new file mode 100644
index 0000000..33444f3
--- /dev/null
+++ b/julia/main.jl
@@ -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()
diff --git a/julia/navier_stokes.jl b/julia/navier_stokes.jl
new file mode 100644
index 0000000..5f8a9ca
--- /dev/null
+++ b/julia/navier_stokes.jl
@@ -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
diff --git a/julia/print.jl b/julia/print.jl
new file mode 100644
index 0000000..7e04013
--- /dev/null
+++ b/julia/print.jl
@@ -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