FourierFlows · navidcy · Jun 7, 2022 · May 30, 2022 · May 30, 2022 · May 30, 2022
diff --git a/.gitignore b/.gitignore
@@ -10,10 +10,14 @@
 **/*.cov
 **/Manifest.toml
 
+# vs code
+.vscode
+
 # Documenter.jl and MkDocs
 docs/build/
 docs/site/
 docs/src/generated
+docs/src/literated
 docs/src/examples
 docs/Manifest.toml
 docs/build/

diff --git a/Project.toml b/Project.toml
@@ -10,6 +10,7 @@ version = "0.6.0"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 FourierFlows = "2aec4490-903f-5c70-9b11-9bed06a700e1"
+GeophysicalFlows = "44ee3b1c-bc02-53fa-8355-8e347616e15e"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -19,7 +20,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 [compat]
 CUDA = "^1.3, ^2, ^3"
 DocStringExtensions = "^0.8"
-FourierFlows = "^0.6.17, ^0.7"
+FourierFlows = "^0.6.17, ^0.7, 0.8, 0.9"
 JLD2 = "^0.1, ^0.2, ^0.3, ^0.4"
 Reexport = "^0.2, ^1"
 julia = "^1.6"

diff --git a/docs/Project.toml b/docs/Project.toml
@@ -1,5 +1,6 @@
 [deps]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"

diff --git a/docs/make.jl b/docs/make.jl
@@ -18,6 +18,7 @@ const OUTPUT_DIR   = joinpath(@__DIR__, "src/literated")
 
 examples = [
     "cellularflow.jl",
+    "turbulent_advection-diffusion.jl"
 ]
 
 for example in examples
@@ -56,6 +57,7 @@ makedocs(
            "Examples" => Any[
              "TracerAdvectionDiffusion" => Any[
                "literated/cellularflow.md",
+               "literated/turbulent_advection-diffusion.md",
                ]
            ],
            "Modules" => Any[

diff --git a/examples/turbulent_advection-diffusion.jl b/examples/turbulent_advection-diffusion.jl
@@ -0,0 +1,169 @@
+# # Advection-diffusion of tracer by a turbulent flow
+#
+# This is an example demonstrating the advection-diffusion of a tracer using a 
+# turbulent flow generated by the `GeophysicalFlows.jl` package.
+#
+# ## Install dependencies
+#
+# First let's make sure we have all the required packages installed
+
+# ```julia
+# using Pkg
+# pkg.add(["PassiveTracerFlows", "Printf", "Plots", "JLD2"])
+#
+# ## Let's begin
+# First load `PassiveTracerFlows.jl` and the other packages needed to run this example.
+using PassiveTracerFlows, Printf, Plots, JLD2
+using Random: seed!
+
+# ## Choosing a device: CPU or GPU
+dev = CPU()
+nothing # hide
+
+# ## Setting up a `MultiLayerQG.Problem` to generate a turbulent flow
+#
+# The tubulent flow we will use to advect the passive tracer is generated using the 
+# [`MultiLayerQG`](https://fourierflows.github.io/GeophysicalFlowsDocumentation/stable/modules/multilayerqg/) module from the [`GeophysicalFlows.jl`](https://fourierflows.github.io/GeophysicalFlowsDocumentation/stable/) package. A more detailed setup  
+# of this two layer system can be found [here](https://fourierflows.github.io/GeophysicalFlowsDocumentation/stable/literated/multilayerqg_2layer/).
+#
+# ### Numerical and time stepping parameters for the flow
+
+      n = 128            # 2D resolution = n²
+stepper = "FilteredRK4"  # timestepper
+     dt = 2.5e-3         # timestep
+
+# Physical parameters 
+L = 2π                   # domain size
+μ = 5e-2                 # bottom drag
+β = 5                    # the y-gradient of planetary PV
+
+nlayers = 2              # number of layers
+f₀, g = 1, 1             # Coriolis parameter and gravitational constant
+ H = [0.2, 0.8]          # the rest depths of each layer
+ ρ = [4.0, 5.0]          # the density of each layer
+
+ U = zeros(nlayers) # the imposed mean zonal flow in each layer
+ U[1] = 1.0
+ U[2] = 0.0
+
+# ### `MultiLayerQG.Problem` setup, shortcuts and initial conditions
+MQGprob = MultiLayerQG.Problem(nlayers, dev;
+                        nx=n, Lx=L, f₀=f₀, g=g, H=H, ρ=ρ, U=U, μ=μ, β=β,
+                        dt=dt, stepper=stepper, aliased_fraction=0)
-                        nx=n, Lx=L, f₀=f₀, g=g, H=H, ρ=ρ, U=U, μ=μ, β=β,
-                        dt=dt, stepper=stepper, aliased_fraction=0)
+                               nx=n, Lx=L, f₀, g, H, ρ, U, μ, β,
+                               dt, stepper, aliased_fraction=0)
-                        nx=n, Lx=L, f₀=f₀, g=g, H=H, ρ=ρ, U=U, μ=μ, β=β,
-                        dt=dt, stepper=stepper, aliased_fraction=0)
+                               nx=n, Lx=L, f₀, g, H, ρ, U, μ, β,
+                               dt, stepper, aliased_fraction=0)
+grid = MQGprob.grid
+x, y = grid.x, grid.y
+
+# Initial conditions                        
+seed!(1234) # reset of the random number generator for reproducibility
+q₀  = 1e-2 * ArrayType(dev)(randn((grid.nx, grid.ny, nlayers)))
+q₀h = MQGprob.timestepper.filter .* rfft(q₀, (1, 2)) # apply rfft  only in dims=1, 2
+q₀  = irfft(q₀h, grid.nx, (1, 2))                 # apply irfft only in dims=1, 2
-q₀  = irfft(q₀h, grid.nx, (1, 2))                 # apply irfft only in dims=1, 2
+q₀  = irfft(q₀h, grid.nx, (1, 2))                    # apply irfft only in dims=1, 2
-q₀  = irfft(q₀h, grid.nx, (1, 2))                 # apply irfft only in dims=1, 2
+q₀  = irfft(q₀h, grid.nx, (1, 2))                    # apply irfft only in dims=1, 2
+
+MultiLayerQG.set_q!(MQGprob, q₀)
+nothing 
+
+# ## Tracer advection-diffusion setup
+# Now that we have a `MultiLayerQG.Problem` setup to generate our turbulent flow, we
+# setup an advection-diffusion simulation. This is done by passing the `MultiLayerQG.Problem`
+# as an argument to `TracerAdvectionDiffusion.Problem` which sets up an advection-diffusion problem
+# with same parameters where applicable. We also need to pass a value for the constant diffusivity `κ`,
+# the `stepper` used to step the problem forward and when we want the tracer released into the flow.
+# We will let the flow run until it reaches a statistical equilibrium and then advect-diffuse the tracer.
+
+κ = 0.002
+nsteps = 4000               # total number of time-steps
+nsubs = 1                   # number of steps the simulation takes at each iteration 
+tracer_release = dt * 8000  # run flow for some time before releasing tracer
+
+ADprob = TracerAdvectionDiffusion.Problem(dev, MQGprob; κ, stepper, tracer_release)
+nothing
+
+# ## Initial condition for concentration in both layers
+# We have a two layer system so we will advect-diffuse the tracer in both layers.
+# To do this we set the initial condition for tracer concetration as a Gaussian centred at the origin.
+# Then we create some shortcuts for the `TracerAdvectionDiffusion.Problem`.
+gaussian(x, y, σ) = exp(-(x^2 + y^2) / (2σ^2))
+
+amplitude, spread = 10, 0.15
+c₀ = [amplitude * gaussian(x[i], y[j], spread) for i=1:grid.nx, j=1:grid.ny]
+
+TracerAdvectionDiffusion.set_c!(ADprob, c₀, nlayers)
+
+# Shortcuts for advection-diffusion problem
+sol, clock, vars, params, grid = ADprob.sol, ADprob.clock, ADprob.vars, ADprob.params, ADprob.grid
+x, y = grid.x, grid.y
+
+# ## Saving output
+# The parent package `FourierFlows.jl` provides a way to save information from a simulation.
+# To do this we write a function `GetConcentration` and pass this to the `Output` function along 
+# with the `TracerAdvectionDiffusion.Problem` and the name of the output file.
+
+function GetConcentration(prob)
+    Concentration = @. prob.vars.c
+    return Concentration
+end
+output = Output(ADprob, "advection-diffusion.jld2", (:Concentration, GetConcentration))
+# This saves information that we will use for plotting later on
+saveproblem(output)
+
+# ## Step the problem forward and save the output
+# We specify that we would like to save the concentration every 50 timesteps using `save_frequency`
+# then step the problem forward.
+save_frequency = 50 # Freqeuncy at which output is saved
+
+startwalltime = time()
+while clock.step <= nsteps
+
+    if clock.step % save_frequency == 0
+
+       saveoutput(output)
+       log = @sprintf("Output saved, step: %04d, t: %d, walltime: %.2f min",
+                      clock.step, clock.t, (time()-startwalltime)/60)
+
+       println(log)
+     end
+
+     stepforward!(ADprob)
+     TracerAdvectionDiffusion.MQGupdatevars!(ADprob)
+end
+
+# Append this information to our saved data for plotting later on
+jldopen(output.path, "a+") do path
+    path["save_frequency"] = save_frequency
+    path["final_step"] = ADprob.clock.step - 1
+end
+
+# ## Visualising the output
+# We now have output from our simulation saved in `advection-diffusion.jld2`.
+# From this we can create a time series for the tracer that has been advected-diffused
+# in the lower layer of our turbulent flow (the simulation from the upper layer can be obtained in a similar manner).
+
+# Create time series for the concentration in the upper layer
+conc_data = load("advection-diffusion.jld2")
+
+saved_data = 0:conc_data["save_frequency"]:conc_data["final_step"]
+t = [conc_data["snapshots/t/"*string(i)] for i ∈ saved_data]
+# Concentration time series in the lower layer
+Cₗ = [abs.(conc_data["snapshots/Concentration/"*string(i)][:, :, 2]) for i ∈ saved_data]
+
+x, y,  = conc_data["grid/x"], conc_data["grid/y"]
+Lx, Ly = conc_data["grid/Lx"], conc_data["grid/Ly"]
+plot_args = (xlabel = "x",
+            ylabel = "y",
+            aspectratio = 1,
+            framestyle = :box,
+            xlims = (-Lx/2, Lx/2),
+            ylims = (-Ly/2, Ly/2),
+            colorbar = true,
+            colorbar_title = " \nConcentration",
+            color = :deep)
+
+p = heatmap(x, y, Cₗ[1]', title = "Concentration, t = $(t[1])"; plot_args...)
+conc_anim = @animate for i ∈ 2:length(t)
+
+    heatmap!(p, x, y, Cₗ[i]', title = "Concentration, t = $(t[i])"; plot_args...)
+
+end
+
+# Create a movie of the tracer
+mp4(conc_anim, "conc_adv-diff.mp4", fps = 12)