uv.jl

using Plots
using PyPlot, FourierFlows

using Random: seed!
using Printf: @printf

import GeophysicalFlows.TwoDTurb
import GeophysicalFlows.TwoDTurb: energy, enstrophy, dissipation, work, drag

nx = 128        # Resolution
Lx = 2π         # Domain size
nu = 1e-4       # Viscosity
nnu = 1         # Order of (hyper-)viscosity. nnu=1 means Laplacian
dt = 0.1        # Timestep
nint = 300      # Number of steps between plots
ntot = 50nint # Number of total timesteps

kf, dkf = 4.0, 2.0     # forcing central wavenumber, wavenumber width
ε = 0.000001 # energy injection rate

mygrid  = TwoDGrid(nx, Lx)
x, y = gridpoints(mygrid)
Kr = [ mygrid.kr[i] for i=1:mygrid.nkr, j=1:mygrid.nl]
force2k = @. exp(-(sqrt(mygrid.Krsq)-kf)^2/(2*dkf^2))
force2k[mygrid.Krsq .< 2.0^2 ] .= 0
force2k[mygrid.Krsq .> 20.0^2 ] .= 0
force2k[Kr .< 2π/Lx] .= 0
ε0 = FourierFlows.parsevalsum(force2k.*mygrid.invKrsq/2.0, mygrid)/(mygrid.Lx*mygrid.Ly)
force2k .= ε/ε0 * force2k

seed!(1234)

function calcF!(Fh, sol, t, cl, v, p, g)
  eta = exp.(2π*im*rand(Float64, size(sol)))/sqrt(cl.dt)
  eta[1, 1] = 0
  @. Fh = eta*sqrt(force2k)
  nothing
end

prob = TwoDTurb.Problem(nx=nx, Lx=Lx, nu=nu, nnu=nnu, dt=dt, stepper="RK4",
  calcF=calcF!,stochastic=true)
TwoDTurb.set_zeta!(prob, 2*rand(nx, nx))


cl, vs, mygrid = prob.clock, prob.vars, prob.grid
x, y = gridpoints(mygrid)

"Plot the vorticity of the twodturb problem `prob`."
function makeplot!(ax, prob)
  Plots.display(Plots.heatmap(vs.zeta))
  nothing
end

# Step forward

all_vs = zeros(128,128,20)
all_us = zeros(128,128,20)
all_zetas = zeros(128,128,500)

while true
  @time begin
    stepforward!(prob, 1000)
    @printf("step: %04d, t: %6.1f", cl.step, cl.t)
  end
  TwoDTurb.updatevars!(prob)
  if cl.t > 1200
    break
  end
  makeplot!(1,1)
end

for i in 1:10
  @time begin
    stepforward!(prob, 20)
    @printf("step: %04d, t: %6.1f", cl.step, cl.t)
  end
  all_us[:,:,i] = vs.u
  all_vs[:,:,i] = vs.v
  all_zetas[:,:,i] = vs.zeta
  TwoDTurb.updatevars!(prob)
end