diff --git a/examples/Project.toml b/examples/Project.toml index 1669ce76..3e974ded 100644 --- a/examples/Project.toml +++ b/examples/Project.toml @@ -1,5 +1,6 @@ [deps] CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" GeophysicalFlows = "44ee3b1c-bc02-53fa-8355-8e347616e15e" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" diff --git a/examples/twodnavierstokes_decaying_tracer.jl b/examples/twodnavierstokes_decaying_tracer.jl new file mode 100644 index 00000000..44cd3e53 --- /dev/null +++ b/examples/twodnavierstokes_decaying_tracer.jl @@ -0,0 +1,259 @@ +# # 2D decaying turbulence with tracers +# +# A simulation of decaying two-dimensional turbulence with tracers (in this case, 3 tracers with different initial conditions). +# +# ## Install dependencies +# +# First let's make sure we have all required packages installed. + +# ```julia +# using Pkg +# pkg"add GeophysicalFlows, Printf, Random, Plots" +# ``` + +# ## Let's begin +# Let's load `GeophysicalFlows.jl` and some other needed packages. +# +using GeophysicalFlows, Printf, Random, Plots + +using Random: seed! +using GeophysicalFlows: peakedisotropicspectrum + + +# ## Choosing a device: CPU or GPU + +dev = CPU() # Device (CPU/GPU) +nothing # hide + + +# ## Numerical, domain, and simulation parameters +# +# First, we pick some numerical and physical parameters for our model. + +n, L = 128, 2π # grid resolution and domain length +nothing # hide + +# Then we pick the time-stepper parameters + dt = 1e-2 # timestep +nsteps = 4000 # total number of steps + nsubs = 20 # number of steps between each plot + ntracers = 3 # number of tracers +nothing # hide + + +# ## Problem setup +# We initialize a `Problem` by providing a set of keyword arguments. The +# `stepper` keyword defines the time-stepper to be used. +prob = TwoDNavierStokesTracer.Problem(ntracers, dev; nx=n, Lx=L, ny=n, Ly=L, dt=dt, stepper="FilteredRK4") +nothing # hide + +# Next we define some shortcuts for convenience. +sol, clock, vars, grid = prob.sol, prob.clock, prob.vars, prob.grid +x, y = grid.x, grid.y +nothing # hide + + +# ## Setting initial conditions + +# Our initial condition tries to reproduce the initial condition used by McWilliams (_JFM_, 1984). +seed!(1234) +k₀, E₀ = 6, 0.5 +ζ₀ = peakedisotropicspectrum(grid, k₀, E₀, mask=prob.timestepper.filter[:,:,1]) + +# Our initial condition for tracers (``c``). +# For the first tracer we use a gaussian centered at ``(x, y) = (L_x/5, 0)``. +gaussian(x, y, σ) = exp(-(x^2 + y^2) / (2σ^2)) + +amplitude, spread = 0.5, 0.15 +for j in 1:ntracers + if j>1 + global c₀ = cat(c₀, 0.0.*ζ₀, dims=3) + else + global c₀ = 0.0.*ζ₀ + end +end +c₀[:,:,1] = [amplitude * gaussian(x[i] - 0.1 * grid.Lx, y[j], spread) for i=1:grid.nx, j=1:grid.ny] +# For the second tracer we use a wider gaussian centered at a different location +c₀[:,:,2] = [amplitude * gaussian(x[i] - 0.2 * grid.Lx, y[j] + 0.1 * grid.Ly, 4*spread) for i=1:grid.nx, j=1:grid.ny] +# For the third tracer we use a straight band in the x-direction that is an eighth as wide as the domain +width = 1/8 +[c₀[i,j,3] = amplitude for i=1:grid.nx, j=Int( (1/2) * grid.ny ):Int( (1/2 + width) * grid.ny )] + +TwoDNavierStokesTracer.set_ζ_and_tracers!(prob, ζ₀, c₀) +nothing # hide + +# Let's plot the initial vorticity field. Note that when plotting, we decorate the variable +# to be plotted with `Array()` to make sure it is brought back on the CPU when `vars` live on +# the GPU. +heatmap(x, y, Array(vars.ζ[:,:,1]'), + aspectratio = 1, + c = :balance, + clim = (-40, 40), + xlims = (-L/2, L/2), + ylims = (-L/2, L/2), + xticks = -3:3, + yticks = -3:3, + xlabel = "x", + ylabel = "y", + title = "initial vorticity", + framestyle = :box) + + +# ## Diagnostics + +# Create Diagnostics -- `energy` and `enstrophy` functions are imported at the top. +E = Diagnostic(TwoDNavierStokesTracer.energy, prob; nsteps=nsteps) +Z = Diagnostic(TwoDNavierStokesTracer.enstrophy, prob; nsteps=nsteps) +diags = [E, Z] # A list of Diagnostics types passed to "stepforward!" will be updated every timestep. +nothing # hide + + +# ## Output + +# We choose folder for outputing `.jld2` files and snapshots (`.png` files). +filepath = "." +plotpath = "./plots_decayingTwoDNavierStokesTracer" +plotname = "snapshots" +filename = joinpath(filepath, "decayingTwoDNavierStokesTracer.jld2") +nothing # hide + +# Do some basic file management +if isfile(filename); rm(filename); end +if !isdir(plotpath); mkdir(plotpath); end +nothing # hide + +# And then create Output +get_sol(prob) = prob.sol # extracts the Fourier-transformed solution +get_u(prob) = irfft(im * grid.l .* grid.invKrsq .* sol, grid.nx) +out = Output(prob, filename, (:sol, get_sol), (:u, get_u)) +saveproblem(out) +nothing # hide + + +# ## Visualizing the simulation + +# We initialize a plot with the vorticity field and the time-series of +# energy and enstrophy diagnostics. + +p1 = heatmap(x, y, Array(vars.ζ[:,:,1]'), + aspectratio = 1, + c = :balance, + clim = (-40, 40), + xlims = (-L/2, L/2), + ylims = (-L/2, L/2), + xticks = -3:3, + yticks = -3:3, + xlabel = "x", + ylabel = "y", + title = "vorticity, t=" * @sprintf("%.2f", clock.t), + framestyle = :box) + +p2 = plot(2, # this means "a plot with two series" + label = ["energy E(t)/E(0)" "enstrophy Z(t)/Z(0)"], + legend = :right, + linewidth = 2, + alpha = 0.7, + xlabel = "t", + xlims = (0, 41), + ylims = (0, 1.1)) + +ptimeϕ₁ = heatmap(x, y, Array(vars.ζ[:,:,2]'), + aspectratio = 1, + c = :balance, + clim = (-0.5, 0.5), + xlims = (-L/2, L/2), + ylims = (-L/2, L/2), + xticks = -3:3, + yticks = -3:3, + xlabel = "x", + ylabel = "y", + title = "ϕ₁(x, y, t=" * @sprintf("%.2f", clock.t) * ")", + framestyle = :box) + + ptimeϕ₂ = heatmap(x, y, Array(vars.ζ[:,:,3]'), + aspectratio = 1, + c = :balance, + clim = (-0.5, 0.5), + xlims = (-L/2, L/2), + ylims = (-L/2, L/2), + xticks = -3:3, + yticks = -3:3, + xlabel = "x", + ylabel = "y", + title = "ϕ₂(x, y, t=" * @sprintf("%.2f", clock.t) * ")", + framestyle = :box) + + ptimeϕ₃ = heatmap(x, y, Array(vars.ζ[:,:,4]'), + aspectratio = 1, + c = :balance, + clim = (-0.5, 0.5), + xlims = (-L/2, L/2), + ylims = (-L/2, L/2), + xticks = -3:3, + yticks = -3:3, + xlabel = "x", + ylabel = "y", + title = "ϕ₃(x, y, t=" * @sprintf("%.2f", clock.t) * ")", + framestyle = :box) + +p = plot(p1, p2, ptimeϕ₁, ptimeϕ₂, ptimeϕ₃, layout=5, size = (1500, 800)) +#l = @layout Plots.grid(1, 2) +#p = plot(p1, p2, layout = l, size = (800, 360)) + + +# ## Time-stepping the `Problem` forward + +# We time-step the `Problem` forward in time. + +startwalltime = time() + +anim = @animate for j = 0:Int(nsteps/nsubs) + if j % (1000 / nsubs) == 0 + cfl = clock.dt * maximum([maximum(vars.u) / grid.dx, maximum(vars.v) / grid.dy]) + + log = @sprintf("step: %04d, t: %d, cfl: %.2f, ΔE: %.4f, ΔZ: %.4f, walltime: %.2f min", + clock.step, clock.t, cfl, E.data[E.i]/E.data[1], Z.data[Z.i]/Z.data[1], (time()-startwalltime)/60) + + println(log) + end + + p[1][1][:z] = Array(vars.ζ[:,:,1]) + p[1][:title] = "vorticity, t=" * @sprintf("%.2f", clock.t) + push!(p[2][1], E.t[E.i], E.data[E.i]/E.data[1]) + push!(p[2][2], Z.t[Z.i], Z.data[Z.i]/Z.data[1]) + p[3][1][:z] = Array(vars.ζ[:,:,2]) + p[4][1][:z] = Array(vars.ζ[:,:,3]) + p[5][1][:z] = Array(vars.ζ[:,:,4]) + + stepforward!(prob, diags, nsubs) + TwoDNavierStokesTracer.updatevars!(prob) +end + +mp4(anim, "twodturbtracer.mp4", fps=18) + + +# Last we can save the output by calling +# ```julia +# saveoutput(out)` +# ``` + + +# ## Radial energy spectrum + +# After the simulation is done we plot the instantaneous radial energy spectrum to illustrate +# how `FourierFlows.radialspectrum` can be used, + +E = @. 0.5 * (vars.u^2 + vars.v^2) # energy density +Eh = rfft(E) # Fourier transform of energy density +kr, Ehr = FourierFlows.radialspectrum(Eh, grid, refinement=1) # compute radial specturm of `Eh` +nothing # hide + +# and we plot it. +plot(kr, abs.(Ehr), + linewidth = 2, + alpha = 0.7, + xlabel = "kᵣ", ylabel = "∫ |Ê| kᵣ dk_θ", + xlims = (5e-1, grid.nx), + xscale = :log10, yscale = :log10, + title = "Radial energy spectrum", + legend = false) diff --git a/examples/twodnavierstokes_stochasticforcing_tracer.jl b/examples/twodnavierstokes_stochasticforcing_tracer.jl new file mode 100644 index 00000000..4dae2afe --- /dev/null +++ b/examples/twodnavierstokes_stochasticforcing_tracer.jl @@ -0,0 +1,261 @@ +# # 2D forced-dissipative turbulence +# +#md # This example can be viewed as a Jupyter notebook via [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/literated/twodnavierstokes_stochasticforcing.ipynb). +# +# A simulation of forced-dissipative two-dimensional turbulence. We solve the +# two-dimensional vorticity equation with stochastic excitation and dissipation in +# the form of linear drag and hyperviscosity. +# +# ## Install dependencies +# +# First let's make sure we have all required packages installed. + +# ```julia +# using Pkg +# pkg"add GeophysicalFlows, Random, Printf, Plots" +# ``` + +# ## Let's begin +# Let's load `GeophysicalFlows.jl` and some other needed packages. +# +using GeophysicalFlows, Random, Printf, Plots +parsevalsum = FourierFlows.parsevalsum + +# ## Choosing a device: CPU or GPU + +dev = CPU() # Device (CPU/GPU) +nothing # hide + + +# ## Numerical, domain, and simulation parameters +# +# First, we pick some numerical and physical parameters for our model. + + n, L = 256, 2π # grid resolution and domain length + ν, nν = 2e-7, 2 # hyperviscosity coefficient and hyperviscosity order + μ, nμ = 1e-1, 0 # linear drag coefficient + κ, nκ = 2e-7, 2 + dt = 0.005 # timestep +nsteps = 4000 # total number of steps + nsubs = 20 # number of steps between each plot + ntracers = 3 # number of tracers +nothing # hide + + +# ## Forcing + +# We force the vorticity equation with stochastic excitation that is delta-correlated in time +# and while spatially homogeneously and isotropically correlated. The forcing has a spectrum +# with power in a ring in wavenumber space of radius ``k_f`` (`forcing_wavenumber`) and width +# ``δ_f`` (`forcing_bandwidth`), and it injects energy per unit area and per unit time +# equal to ``\varepsilon``. That is, the forcing covariance spectrum is proportional to +# ``\exp{[-(|\bm{k}| - k_f)^2 / (2 δ_f^2)]}``. + +forcing_wavenumber = 14.0 * 2π/L # the forcing wavenumber, `k_f`, for a spectrum that is a ring in wavenumber space +forcing_bandwidth = 1.5 * 2π/L # the width of the forcing spectrum, `δ_f` +ε = 0.1 # energy input rate by the forcing + +grid = TwoDGrid(dev, n, L) + +K = @. sqrt(grid.Krsq) # a 2D array with the total wavenumber + +forcing_spectrum = @. exp(-(K - forcing_wavenumber)^2 / (2 * forcing_bandwidth^2)) +@CUDA.allowscalar forcing_spectrum[grid.Krsq .== 0] .= 0 # ensure forcing has zero domain-average + +ε0 = parsevalsum(forcing_spectrum .* grid.invKrsq / 2, grid) / (grid.Lx * grid.Ly) +@. forcing_spectrum *= ε/ε0 # normalize forcing to inject energy at rate ε +nothing # hide + + +# We reset of the random number generator for reproducibility +if dev==CPU(); Random.seed!(1234); else; CUDA.seed!(1234); end +nothing # hide + + +# Next we construct function `calcF!` that computes a forcing realization every timestep. +# First we make sure that if `dev=GPU()`, then `CUDA.rand()` function is called for random +# numbers uniformly distributed between 0 and 1. +random_uniform = dev==CPU() ? rand : CUDA.rand + +function calcF!(Fh, sol, t, clock, vars, params, grid) + Fh .= sqrt.(forcing_spectrum) .* exp.(2π * im * random_uniform(eltype(grid), size(sol[:,:,1]))) ./ sqrt(clock.dt) + + return nothing +end +nothing # hide + + +# ## Problem setup +# We initialize a `Problem` by providing a set of keyword arguments. The +# `stepper` keyword defines the time-stepper to be used. +prob = TwoDNavierStokesTracer.Problem(ntracers, dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, nμ=nμ, κ=κ, nκ=nκ, dt=dt, stepper="ETDRK4", + calcF=calcF!, stochastic=true) +nothing # hide + +# Define some shortcuts for convenience. +sol, clock, vars, params, grid = prob.sol, prob.clock, prob.vars, prob.params, prob.grid + +x, y = grid.x, grid.y +nothing # hide + + +# First let's see how a forcing realization looks like. Function `calcF!()` computes +# the forcing in Fourier space and saves it into variable `vars.Fh`, so we first need to +# go back to physical space. + +# Note that when plotting, we decorate the variable to be plotted with `Array()` to make sure +# it is brought back on the CPU when the variable lives on the GPU. +calcF!(vars.Fh, sol, 0.0, clock, vars, params, grid) + +heatmap(x, y, Array(irfft(vars.Fh, grid.nx)'), + aspectratio = 1, + c = :balance, + clim = (-200, 200), + xlims = (-L/2, L/2), + ylims = (-L/2, L/2), + xticks = -3:3, + yticks = -3:3, + xlabel = "x", + ylabel = "y", + title = "a forcing realization", + framestyle = :box) + + +# ## Setting initial conditions + +# Our initial condition is a fluid at rest. + +ζ₀ = ArrayType(dev)(zeros(grid.nx, grid.ny)) + +# Our initial condition for tracers (``c``). +# For the first tracer we use a gaussian centered at ``(x, y) = (L_x/5, 0)``. +gaussian(x, y, σ) = exp(-(x^2 + y^2) / (2σ^2)) + +amplitude, spread = 0.5, 0.15 +for j in 1:ntracers + if j>1 + global c₀ = cat(c₀, 0.0.*ζ₀, dims=3) + else + global c₀ = 0.0.*ζ₀ + end +end +c₀[:,:,1] = [amplitude * gaussian(x[i] - 0.1 * grid.Lx, y[j], spread) for i=1:grid.nx, j=1:grid.ny] +# For the second tracer we use a wider gaussian centered at a different location +c₀[:,:,2] = [amplitude * gaussian(x[i] - 0.2 * grid.Lx, y[j] + 0.1 * grid.Ly, 4*spread) for i=1:grid.nx, j=1:grid.ny] +# For the third tracer we use a straight band in the x-direction that is an eighth as wide as the domain +width = 1/8 +[c₀[i,j,3] = amplitude for i=1:grid.nx, j=Int( (1/2) * grid.ny ):Int( (1/2 + width) * grid.ny )] + +TwoDNavierStokesTracer.set_ζ_and_tracers!(prob, ζ₀, c₀) +nothing # hide + + +# ## Diagnostics + +# Create Diagnostics; the diagnostics are aimed to probe the energy budget. +E = Diagnostic(TwoDNavierStokesTracer.energy, prob, nsteps=nsteps) # energy +Z = Diagnostic(TwoDNavierStokesTracer.enstrophy, prob, nsteps=nsteps) # enstrophy +diags = [E, Z] # a list of Diagnostics passed to `stepforward!` will be updated every timestep. +nothing # hide + + +# ## Visualizing the simulation + +# We initialize a plot with the vorticity field and the time-series of +# energy and enstrophy diagnostics. To plot energy and enstrophy on the same +# axes we scale enstrophy with ``k_f^2``. + +p1 = heatmap(x, y, Array(vars.ζ[:,:,1]'), + aspectratio = 1, + c = :balance, + clim = (-40, 40), + xlims = (-L/2, L/2), + ylims = (-L/2, L/2), + xticks = -3:3, + yticks = -3:3, + xlabel = "x", + ylabel = "y", + title = "vorticity, t=" * @sprintf("%.2f", clock.t), + framestyle = :box) + +p2 = plot(2, # this means "a plot with two series" + label = ["energy E(t)" "enstrophy Z(t) / k_f²"], + legend = :right, + linewidth = 2, + alpha = 0.7, + xlabel = "μ t", + xlims = (0, 1.1 * μ * nsteps * dt), + ylims = (0, 0.55)) + +ptimeϕ₁ = heatmap(x, y, Array(vars.ζ[:,:,2]'), + aspectratio = 1, + c = :balance, + clim = (-0.5, 0.5), + xlims = (-L/2, L/2), + ylims = (-L/2, L/2), + xticks = -3:3, + yticks = -3:3, + xlabel = "x", + ylabel = "y", + title = "ϕ₁(x, y, t=" * @sprintf("%.2f", clock.t) * ")", + framestyle = :box) + + ptimeϕ₂ = heatmap(x, y, Array(vars.ζ[:,:,3]'), + aspectratio = 1, + c = :balance, + clim = (-0.5, 0.5), + xlims = (-L/2, L/2), + ylims = (-L/2, L/2), + xticks = -3:3, + yticks = -3:3, + xlabel = "x", + ylabel = "y", + title = "ϕ₂(x, y, t=" * @sprintf("%.2f", clock.t) * ")", + framestyle = :box) + + ptimeϕ₃ = heatmap(x, y, Array(vars.ζ[:,:,4]'), + aspectratio = 1, + c = :balance, + clim = (-0.5, 0.5), + xlims = (-L/2, L/2), + ylims = (-L/2, L/2), + xticks = -3:3, + yticks = -3:3, + xlabel = "x", + ylabel = "y", + title = "ϕ₃(x, y, t=" * @sprintf("%.2f", clock.t) * ")", + framestyle = :box) + +p = plot(p1, p2, ptimeϕ₁, ptimeϕ₂, ptimeϕ₃, layout=5, size = (1500, 800)) +#l = @layout Plots.grid(1, 2) +#p = plot(p1, p2, layout = l, size = (900, 420)) + + +# ## Time-stepping the `Problem` forward + +# Finally, we time-step the `Problem` forward in time. + +startwalltime = time() + +anim = @animate for j = 0:round(Int, nsteps / nsubs) + if j % (1000/nsubs) == 0 + cfl = clock.dt * maximum([maximum(vars.u) / grid.dx, maximum(vars.v) / grid.dy]) + + log = @sprintf("step: %04d, t: %d, cfl: %.2f, E: %.4f, Z: %.4f, walltime: %.2f min", + clock.step, clock.t, cfl, E.data[E.i], Z.data[Z.i], (time()-startwalltime)/60) + println(log) + end + + p[1][1][:z] = Array(vars.ζ[:,:,1]) + p[1][:title] = "vorticity, μt = " * @sprintf("%.2f", μ * clock.t) + push!(p[2][1], μ * E.t[E.i], E.data[E.i]) + push!(p[2][2], μ * Z.t[Z.i], Z.data[Z.i] / forcing_wavenumber^2) + p[3][1][:z] = Array(vars.ζ[:,:,2]) + p[4][1][:z] = Array(vars.ζ[:,:,3]) + p[5][1][:z] = Array(vars.ζ[:,:,4]) + + stepforward!(prob, diags, nsubs) + TwoDNavierStokesTracer.updatevars!(prob) +end + +mp4(anim, "twodturb_forced_tracer.mp4", fps=18) diff --git a/src/GeophysicalFlows.jl b/src/GeophysicalFlows.jl index b3401da9..6c4971f1 100644 --- a/src/GeophysicalFlows.jl +++ b/src/GeophysicalFlows.jl @@ -16,12 +16,14 @@ using include("utils.jl") include("twodnavierstokes.jl") +include("twodnavierstokeswithtracer.jl") include("singlelayerqg.jl") include("multilayerqg.jl") include("surfaceqg.jl") include("barotropicqgql.jl") @reexport using GeophysicalFlows.TwoDNavierStokes +@reexport using GeophysicalFlows.TwoDNavierStokesTracer @reexport using GeophysicalFlows.SingleLayerQG @reexport using GeophysicalFlows.MultiLayerQG @reexport using GeophysicalFlows.SurfaceQG diff --git a/src/twodnavierstokes.jl b/src/twodnavierstokes.jl index be691a7a..8d96fcbb 100644 --- a/src/twodnavierstokes.jl +++ b/src/twodnavierstokes.jl @@ -386,7 +386,7 @@ Returns the domain-averaged enstrophy, end """ - energy_dissipation(prob, ξ, νξ) + energy_dissipation(prob, ξ, nξ) Return the domain-averaged energy dissipation rate done by the viscous term, @@ -406,21 +406,21 @@ where ``ξ`` and ``nξ`` could be either the (hyper)-viscosity coefficient ``ν` end """ - energy_dissipation_hyperviscosity(prob, ξ, νξ) + energy_dissipation_hyperviscosity(prob, ξ, nξ) Return the domain-averaged energy dissipation rate done by the ``ν`` (hyper)-viscosity. """ energy_dissipation_hyperviscosity(prob) = energy_dissipation(prob, prob.params.ν, prob.params.nν) """ - energy_dissipation_hypoviscosity(prob, ξ, νξ) + energy_dissipation_hypoviscosity(prob, ξ, nξ) Return the domain-averaged energy dissipation rate done by the ``μ`` (hypo)-viscosity. """ energy_dissipation_hypoviscosity(prob) = energy_dissipation(prob, prob.params.μ, prob.params.nμ) """ - enstrophy_dissipation(prob, ξ, νξ) + enstrophy_dissipation(prob, ξ, nξ) Return the domain-averaged enstrophy dissipation rate done by the viscous term, @@ -441,14 +441,14 @@ where ``ξ`` and ``nξ`` could be either the (hyper)-viscosity coefficient ``ν` end """ - enstrophy_dissipation_hyperviscosity(prob, ξ, νξ) + enstrophy_dissipation_hyperviscosity(prob, ξ, nξ) Return the domain-averaged enstrophy dissipation rate done by the ``ν`` (hyper)-viscosity. """ enstrophy_dissipation_hyperviscosity(prob) = enstrophy_dissipation(prob, prob.params.ν, prob.params.nν) """ - enstrophy_dissipation_hypoviscosity(prob, ξ, νξ) + enstrophy_dissipation_hypoviscosity(prob, ξ, nξ) Return the domain-averaged enstrophy dissipation rate done by the ``μ`` (hypo)-viscosity. """ diff --git a/src/twodnavierstokeswithtracer.jl b/src/twodnavierstokeswithtracer.jl new file mode 100644 index 00000000..f46d54fb --- /dev/null +++ b/src/twodnavierstokeswithtracer.jl @@ -0,0 +1,585 @@ +module TwoDNavierStokesTracer + +export + fwdtransform!, + invtransform!, + Problem, + set_ζ!, + updatevars!, + + energy, + energy_dissipation_hyperviscosity, + energy_dissipation_hypoviscosity, + energy_work, + enstrophy, + enstrophy_dissipation_hyperviscosity, + enstrophy_dissipation_hypoviscosity, + enstrophy_work + +using + FFTW, + CUDA, + Reexport, + DocStringExtensions + +@reexport using FourierFlows + +using LinearAlgebra: mul!, ldiv! +using FourierFlows: parsevalsum, plan_flows_rfft + +nothingfunction(args...) = nothing + +""" + Problem(ntracers::Int, + dev::Device=CPU(); + nx = 256, + ny = nx, + Lx = 2π, + Ly = Lx, + ν = 0, + nν = 1, + μ = 0, + nμ = 0, + κ = 0, + nκ = 1, + dt = 0.01, + stepper = "RK4", + calcF = nothingfunction, + stochastic = false, + aliased_fraction = 1/3, + T = Float64) + +Construct a two-dimensional Navier-Stokes with tracers `problem` on device `dev`. + +Keyword arguments +================= + - `ntracers`: (required) Number of tracers. + - `dev`: (required) `CPU()` or `GPU()`; computer architecture used to time-step `problem`. + - `nx`: Number of grid points in ``x``-domain. + - `ny`: Number of grid points in ``y``-domain. + - `Lx`: Extent of the ``x``-domain. + - `Ly`: Extent of the ``y``-domain. + - `ν`: Small-scale (hyper)-viscosity coefficient. + - `nν`: (Hyper)-viscosity order, `nν```≥ 1``. + - `μ`: Large-scale (hypo)-viscosity coefficient. + - `nμ`: (Hypo)-viscosity order, `nμ```≤ 0``. + - `κ`: Small-scale diffusivity coefficient. + - `nκ`: Diffusivity order, `nκ```≥ 1``. + - `dt`: Time-step. + - `stepper`: Time-stepping method. + - `calcF`: Function that calculates the Fourier transform of the forcing, ``F̂``. + - `stochastic`: `true` or `false`; boolean denoting whether `calcF` is temporally stochastic. + - `aliased_fraction`: the fraction of high-wavenumbers that are zero-ed out by `dealias!()`. + - `T`: `Float32` or `Float64`; floating point type used for `problem` data. +""" +function Problem(ntracers :: Int, # number of tracers + dev :: Device=CPU(); + # Numerical parameters + nx = 256, + ny = nx, + Lx = 2π, + Ly = Lx, + # Drag and/or hyper-/hypo-viscosity + ν = 0.0, + nν = 1, + μ = 0.0, + nμ = 0, + κ = 0.0, + nκ = 0, + # Timestepper and equation options + dt = 0.01, + stepper = "RK4", + calcF = nothingfunction, + stochastic = false, + # Float type and dealiasing + aliased_fraction = 1/3, + T = Float64) + + grid = TwoDGrid(dev, nx, Lx, ny, Ly; aliased_fraction=aliased_fraction, T=T) + + params = Params(ntracers, ν, nν, μ, nμ, κ, nκ, grid; calcF, dev) + + vars = calcF == nothingfunction ? DecayingVars(dev, grid, params) : (stochastic ? StochasticForcedVars(dev, grid, params) : ForcedVars(dev, grid, params)) + + equation = Equation(dev, params, grid) + + return FourierFlows.Problem(equation, stepper, dt, grid, vars, params, dev) +end + + +# ---------- +# Parameters +# ---------- + +""" + Params{T}(ν, nν, μ, nμ, κ, nκ calcF!) + +A container for the parameters for the two-dimensional Navier-Stokes problem with tracers. + +$(TYPEDFIELDS) +""" +struct Params{T, Trfft} <: AbstractParams + "number of tracers" + ntracers :: Int + "small-scale (hyper)-viscosity coefficient" + ν :: T + "(hyper)-viscosity order, `nν```≥ 1``" + nν :: Int + "large-scale (hypo)-viscosity coefficient" + μ :: T + "(hypo)-viscosity order, `nμ```≤ 0``" + nμ :: Int + "tracer diffusivity coefficient" + κ :: T + "tracer diffusivity order" + nκ :: Int + "function that calculates the Fourier transform of the forcing, ``F̂``" + calcF! :: Function + "rfft plan for FFTs (Derived parameter)" + rfftplan :: Trfft +end + +""" + Params(ntracers, ν, nν, μ, nμ, κ, nκ, grid; + calcF=nothingfunction, effort=FFTW.MEASURE, dev::Device=CPU()) + +Return the parameters for a the two-dimensional Navier-Stokes problem with tracers. +""" +function Params(ntracers, ν, nν, μ, nμ, κ, nκ, grid; + calcF=nothingfunction, effort=FFTW.MEASURE, dev::Device=CPU()) where TU + T = eltype(grid) + A = ArrayType(dev) + + ny, nx = grid.ny , grid.nx + nkr, nl = grid.nkr, grid.nl + kr, l = grid.kr , grid.l + + rfftplanlayered = plan_flows_rfft(A{T, 3}(undef, grid.nx, grid.ny, ntracers + 1), [1, 2]; flags=effort) + + return Params(ntracers, ν, nν, μ, nμ, κ, nκ, calcF, rfftplanlayered) +end + +""" + fwdtransform!(varh, var, params) + +Compute the Fourier transform of `var` and store it in `varh`. +""" +fwdtransform!(varh, var, params::AbstractParams) = mul!(varh, params.rfftplan, var) + +""" + invtransform!(var, varh, params) + +Compute the inverse Fourier transform of `varh` and store it in `var`. +""" +invtransform!(var, varh, params::AbstractParams) = ldiv!(var, params.rfftplan, varh) + + +# --------- +# Equations +# --------- + +"Create a variable for number of layers. First layer describes motion, the other layers are tracers." +numberoflayers(params) = params.ntracers + 1 + +""" + Equation(dev, params, grid) + Return the `equation` for two-dimensional dynamics and the `equations`` for tracer evolution under these dynamics, + using `params` and `grid`. The first layer of the linear operator ``L`` describes the dynamics and + includes (hyper)-viscosity of order ``n_ν`` with coefficient ``ν`` and hypo-viscocity of order ``n_μ`` + with coefficient ``μ``. The second layer onwards of ``L`` describe the tracer diffusion of order ``n_κ`` + with coefficient ``κ``. + + ```math + L[:, :, 1] = - ν |𝐤|^{2 n_ν} - μ |𝐤|^{2 n_μ} . + ``` + ```math + L[:, :, 2:ntracers+1] = - κ |𝐤|^{2 n_κ}. + ``` + Plain old viscocity corresponds to ``n_ν=1`` while ``n_μ=0`` corresponds to linear drag. + The nonlinear term is computed via function `calcN!()`. +""" +function Equation(dev, params, grid) + #L = hyperviscosity(dev, params, grid) + nlayers = numberoflayers(params) + T = eltype(grid) + L = ArrayType(dev){T}(undef, (grid.nkr, grid.nl, nlayers)) + @views @. L[:,:,1] = - params.ν * grid.Krsq^params.nν - params.μ * grid.Krsq^params.nμ + @views @. L[:,:,2:nlayers] = - params.κ * grid.Krsq^params.nκ + # Need to add diffusivities for different layers + @views @. L[1, 1, :] = 0 + + return FourierFlows.Equation(L, calcN!, grid) +end + +# ---- +# Vars +# ---- + +""" + Vars{Aphys, Atrans, F, P}(ζ, u, v, ζh, uh, vh, Fh, prevsol) + +The variables for two-dimensional Navier-Stokes: + +$(FIELDS) +""" +struct Vars{Aphys3D, Aphys, Atrans3D, Atrans, F, P} <: AbstractVars + "relative vorticity ([:, :, 1] layer) and tracers (other layers)" + ζ :: Aphys3D + "x-component of velocity" + u :: Aphys + "y-component of velocity" + v :: Aphys + "Fourier transform of relative vorticity ([:, :, 1] layer) and tracers (other layers)" + ζh :: Atrans3D + "Fourier transform of x-component of velocity" + uh :: Atrans + "Fourier transform of y-component of velocity" + vh :: Atrans + "Fourier transform of forcing" + Fh :: F + "`sol` at previous time-step" + prevsol :: P +end + +const DecayingVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing, Nothing} +const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing} +const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray} + +""" + DecayingVars(dev, grid) + +Return the `vars` for unforced two-dimensional Navier-Stokes problem on device `dev` and +with `grid`. +""" +function DecayingVars(::Dev, grid::AbstractGrid, params) where Dev + nlayers = numberoflayers(params) + T = eltype(grid) + + @devzeros Dev T (grid.nx, grid.ny, nlayers) ζ + @devzeros Dev T (grid.nx, grid.ny) u v + @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) ζh + @devzeros Dev Complex{T} (grid.nkr, grid.nl) uh vh + + return Vars(ζ, u, v, ζh, uh, vh, nothing, nothing) +end + +""" + ForcedVars(dev, grid) + +Return the `vars` for forced two-dimensional Navier-Stokes on device `dev` and with `grid`. +""" +function ForcedVars(dev::Dev, grid::AbstractGrid, params) where Dev + nlayers = numberoflayers(params) + T = eltype(grid) + + @devzeros Dev T (grid.nx, grid.ny, nlayers) ζ + @devzeros Dev T (grid.nx, grid.ny) u v + @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) ζh + @devzeros Dev Complex{T} (grid.nkr, grid.nl) uh vh Fh + + return Vars(ζ, u, v, ζh, uh, vh, Fh, nothing) +end + +""" + StochasticForcedVars(dev, grid) + +Return the `vars` for stochastically forced two-dimensional Navier-Stokes on device `dev` and +with `grid`. +""" +function StochasticForcedVars(dev::Dev, grid::AbstractGrid, params) where Dev + nlayers = numberoflayers(params) + T = eltype(grid) + + @devzeros Dev T (grid.nx, grid.ny, nlayers) ζ + @devzeros Dev T (grid.nx, grid.ny) u v + @devzeros Dev Complex{T} (grid.nkr, grid.nl, nlayers) ζh + @devzeros Dev Complex{T} (grid.nkr, grid.nl) uh vh Fh prevsol + + return Vars(ζ, u, v, ζh, uh, vh, Fh, prevsol) +end + + +# ------- +# Solvers +# ------- + +""" + calcN_advection!(N, sol, t, clock, vars, params, grid) + +Calculate the Fourier transform of the advection term, ``- 𝖩(ψ, ζ)`` in conservative form, +i.e., ``- ∂_x[(∂_y ψ)ζ] - ∂_y[(∂_x ψ)ζ]`` and store it in `N`: + +```math +N = - \\widehat{𝖩(ψ, ζ)} = - i k_x \\widehat{u ζ} - i k_y \\widehat{v ζ} . +``` + +Note the first layer ``N[:, :, 1]`` of this advection term corresponds to advection of vorticity ``ζ``. +The subsequent layers of ``N[:, :, 2:ntracers+1]`` correspond to advection of the each tracer. +""" +function calcN_advection!(N, sol, t, clock, vars, params, grid) + nlayers = numberoflayers(params) + @. vars.ζh = sol + + for j in 1:nlayers + @. vars.uh = im * grid.l * grid.invKrsq * sol[:, :, 1] + @. vars.vh = - im * grid.kr * grid.invKrsq * sol[:, :, 1] + ldiv!(vars.u, grid.rfftplan, vars.ζh[:, :, j]) + vars.ζ[:, : ,j] = vars.u + ldiv!(vars.u, grid.rfftplan, vars.uh) + ldiv!(vars.v, grid.rfftplan, vars.vh) + + uζ, uζh = vars.u, vars.uh # use vars.u, vars.uh as scratch variables + vζ, vζh = vars.v, vars.vh # use vars.v, vars.vh as scratch variables + + @. uζ *= vars.ζ[:, :, j] # u*ζ [note u is 2D array and ζ is 3D array] + @. vζ *= vars.ζ[:, :, j] # v*ζ [note v is 2D array and ζ is 3D array] + + mul!(uζh, grid.rfftplan, uζ) # \hat{u*ζ} + mul!(vζh, grid.rfftplan, vζ) # \hat{v*ζ} + + @views @. N[:, :, j] = - im * grid.kr * uζh - im * grid.l * vζh + end + + return nothing +end + +""" + calcN!(N, sol, t, clock, vars, params, grid) + +Calculate the nonlinear term, that is the advection term and the forcing, + +```math +N = - \\widehat{𝖩(ψ, ζ)} + F̂ . +``` +""" +function calcN!(N, sol, t, clock, vars, params, grid) + dealias!(sol, grid) + + calcN_advection!(N, sol, t, clock, vars, params, grid) + + addforcing!(N, sol, t, clock, vars, params, grid) + + return nothing +end + +""" + addforcing!(N, sol, t, clock, vars, params, grid) + +When the problem includes forcing, calculate the forcing term ``F̂`` and add it to the +nonlinear term ``N``. +""" +addforcing!(N, sol, t, clock, vars::DecayingVars, params, grid) = nothing + +function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid) + params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) + @views @. N[:, :, 1] += vars.Fh + + return nothing +end + +function addforcing!(N, sol, t, clock, vars::StochasticForcedVars, params, grid) + if t == clock.t # not a substep + @. vars.prevsol = sol[:, :, 1] # sol at previous time-step is needed to compute budgets for stochastic forcing + params.calcF!(vars.Fh, sol, t, clock, vars, params, grid) + end + @views @. N[:, :, 1] += vars.Fh + + return nothing +end + + +# ---------------- +# Helper functions +# ---------------- + +""" + updatevars!(prob) + +Update variables in `vars` with solution in `sol`. +""" +function updatevars!(prob) + vars, grid, sol = prob.vars, prob.grid, prob.sol + + dealias!(sol, grid) + + @. vars.ζh = sol + @. vars.uh = im * grid.l * grid.invKrsq * sol[:,:,1] + @. vars.vh = - im * grid.kr * grid.invKrsq * sol[:,:,1] + + invtransform!(vars.ζ, deepcopy(vars.ζh), prob.params) # deepcopy() since inverse real-fft destroys its input + ldiv!(vars.u, grid.rfftplan, deepcopy(vars.uh)) # deepcopy() since inverse real-fft destroys its input + ldiv!(vars.v, grid.rfftplan, deepcopy(vars.vh)) # deepcopy() since inverse real-fft destroys its input + + return nothing +end + +""" + set_ζ_and_tracers!(prob, ζ) + +Set the first solution layer `sol[:,:,1]` as the transform of `ζ` and lower layers as transform of `tracers`. +Then update variables in `vars`. +""" +function set_ζ_and_tracers!(prob, ζ, tracers) + full_ζ = cat(ζ, tracers, dims=3) # append b and tracers for use in sol + fwdtransform!(prob.sol, full_ζ, prob.params) + @views CUDA.@allowscalar prob.sol[1, 1, 1] = 0 # zero domain average + + updatevars!(prob) + + return nothing +end + +""" + energy(prob) + +Return the domain-averaged kinetic energy. Since ``u² + v² = |{\\bf ∇} ψ|²``, the domain-averaged +kinetic energy is + +```math +\\int \\frac1{2} |{\\bf ∇} ψ|² \\frac{𝖽x 𝖽y}{L_x L_y} = \\sum_{𝐤} \\frac1{2} |𝐤|² |ψ̂|² . +``` +""" +@inline function energy(prob) + sol, vars, grid = prob.sol, prob.vars, prob.grid + energyh = vars.uh # use vars.uh as scratch variable + + @. energyh = 1 / 2 * grid.invKrsq * abs2(sol[:,:,1]) + return 1 / (grid.Lx * grid.Ly) * parsevalsum(energyh, grid) +end + +""" + enstrophy(prob) + +Returns the domain-averaged enstrophy, + +```math +\\int \\frac1{2} ζ² \\frac{𝖽x 𝖽y}{L_x L_y} = \\sum_{𝐤} \\frac1{2} |ζ̂|² . +``` +""" +@inline function enstrophy(prob) + sol, grid = prob.sol, prob.grid + return 1 / (2 * grid.Lx * grid.Ly) * parsevalsum(abs2.(sol[:,:,1]), grid) +end + +""" + energy_dissipation(prob, ξ, nξ) + +Return the domain-averaged energy dissipation rate done by the viscous term, + +```math +- ξ (-1)^{n_ξ+1} \\int ψ ∇^{2n_ξ} ζ \\frac{𝖽x 𝖽y}{L_x L_y} = - ξ \\sum_{𝐤} |𝐤|^{2(n_ξ-1)} |ζ̂|² , +``` +where ``ξ`` and ``nξ`` could be either the (hyper)-viscosity coefficient ``ν`` and its order +``n_ν``, or the hypo-viscocity coefficient ``μ`` and its order ``n_μ``. +""" +@inline function energy_dissipation(prob, ξ, nξ) + sol, vars, grid = prob.sol, prob.vars, prob.grid + energy_dissipationh = vars.uh # use vars.uh as scratch variable + + @. energy_dissipationh = - ξ * grid.Krsq^(nξ - 1) * abs2(sol[:,:,1]) + CUDA.@allowscalar energy_dissipationh[1, 1] = 0 + return 1 / (grid.Lx * grid.Ly) * parsevalsum(energy_dissipationh, grid) +end + +""" + energy_dissipation_hyperviscosity(prob, ξ, νξ) + +Return the domain-averaged energy dissipation rate done by the ``ν`` (hyper)-viscosity. +""" +energy_dissipation_hyperviscosity(prob) = energy_dissipation(prob, prob.params.ν, prob.params.nν) + +""" + energy_dissipation_hypoviscosity(prob, ξ, νξ) + +Return the domain-averaged energy dissipation rate done by the ``μ`` (hypo)-viscosity. +""" +energy_dissipation_hypoviscosity(prob) = energy_dissipation(prob, prob.params.μ, prob.params.nμ) + +""" + enstrophy_dissipation(prob, ξ, νξ) + +Return the domain-averaged enstrophy dissipation rate done by the viscous term, + +```math +ξ (-1)^{n_ξ+1} \\int ζ ∇^{2n_ξ} ζ \\frac{𝖽x 𝖽y}{L_x L_y} = - ξ \\sum_{𝐤} |𝐤|^{2n_ξ} |ζ̂|² , +``` + +where ``ξ`` and ``nξ`` could be either the (hyper)-viscosity coefficient ``ν`` and its order +``n_ν``, or the hypo-viscocity coefficient ``μ`` and its order ``n_μ``. +""" +@inline function enstrophy_dissipation(prob, ξ, nξ) + sol, vars, grid = prob.sol, prob.vars, prob.grid + enstrophy_dissipationh = vars.uh # use vars.uh as scratch variable + + @. enstrophy_dissipationh = - ξ * grid.Krsq^nξ * abs2(sol[:,:,1]) + CUDA.@allowscalar enstrophy_dissipationh[1, 1] = 0 + return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_dissipationh, grid) +end + +""" + enstrophy_dissipation_hyperviscosity(prob, ξ, νξ) + +Return the domain-averaged enstrophy dissipation rate done by the ``ν`` (hyper)-viscosity. +""" +enstrophy_dissipation_hyperviscosity(prob) = enstrophy_dissipation(prob, prob.params.ν, prob.params.nν) + +""" + enstrophy_dissipation_hypoviscosity(prob, ξ, νξ) + +Return the domain-averaged enstrophy dissipation rate done by the ``μ`` (hypo)-viscosity. +""" +enstrophy_dissipation_hypoviscosity(prob) = enstrophy_dissipation(prob, prob.params.μ, prob.params.nμ) + +""" + energy_work(prob) + energy_work(sol, vars, grid) + +Return the domain-averaged rate of work of energy by the forcing ``F``, + +```math +- \\int ψ F \\frac{𝖽x 𝖽y}{L_x L_y} = - \\sum_{𝐤} ψ̂ F̂^* . +``` +""" +@inline function energy_work(sol, vars::ForcedVars, grid) + energy_workh = vars.uh # use vars.uh as scratch variable + + @. energy_workh = grid.invKrsq * sol[:,:,1] * conj(vars.Fh) + return 1 / (grid.Lx * grid.Ly) * parsevalsum(energy_workh, grid) +end + +@inline function energy_work(sol, vars::StochasticForcedVars, grid) + energy_workh = vars.uh # use vars.uh as scratch variable + + @. energy_workh = grid.invKrsq * (vars.prevsol + sol[:,:,1]) / 2 * conj(vars.Fh) + return 1 / (grid.Lx * grid.Ly) * parsevalsum(energy_workh, grid) +end + +@inline energy_work(prob) = energy_work(prob.sol, prob.vars, prob.grid) + +""" + enstrophy_work(prob) + enstrophy_work(sol, vars, grid) + +Return the domain-averaged rate of work of enstrophy by the forcing ``F``, + +```math +\\int ζ F \\frac{𝖽x 𝖽y}{L_x L_y} = \\sum_{𝐤} ζ̂ F̂^* . +``` +""" +@inline function enstrophy_work(sol, vars::ForcedVars, grid) + enstrophy_workh = vars.uh # use vars.uh as scratch variable + + @. enstrophy_workh = sol[:,:,1] * conj(vars.Fh) + return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_workh, grid) +end + +@inline function enstrophy_work(sol, vars::StochasticForcedVars, grid) + enstrophy_workh = vars.uh # use vars.uh as scratch variable + + @. enstrophy_workh = (vars.prevsol + sol[:,:,1]) / 2 * conj(vars.Fh) + return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_workh, grid) +end + +@inline enstrophy_work(prob) = enstrophy_work(prob.sol, prob.vars, prob.grid) + +end # module