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.

using Pkg
pkg.add(["PassiveTracerFlows", "CairoMakie", "JLD2"])

Let's begin

First load packages needed to run this example.

using PassiveTracerFlows, Printf, CairoMakie, JLD2
using Random: seed!

Choosing a device: CPU or GPU

dev = CPU()

Setting up a MultiLayerQG.Problem to generate a turbulent flow

The tubulent flow we use to advect the passive tracer is generated using the MultiLayerQG module from the GeophysicalFlows.jl package. A more detailed setup of this two layer system is found at the GeophysicalFlows Documentation.

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₀ = 1                   # Coriolis parameter
H = [0.2, 0.8]           # the rest depths of each layer
b = [-1.0, -1.2]         # Boussinesq buoyancy 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₀, H, b, U, μ, β,
                               dt, stepper, aliased_fraction=0)

nx, ny = MQGprob.grid.nx, MQGprob.grid.ny

(128, 128)

Initial conditions

seed!(1234) # reset of the random number generator for reproducibility
q₀  = 1e-2 * device_array(dev)(randn((nx, ny, nlayers)))
q₀h = MQGprob.timestepper.filter .* rfft(q₀, (1, 2)) # apply rfft  only in dims=1, 2
q₀  = irfft(q₀h, nx, (1, 2))                         # apply irfft only in dims=1, 2

MultiLayerQG.set_q!(MQGprob, q₀)

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 and also on the same device (CPU/GPU). 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 let the flow evolve up to t = tracer_release_time and then release the tracer and let it evolve with the flow.

κ = 0.002                        # constant diffusivity
nsteps = 4000                    # total number of time-steps
tracer_release_time = 25.0       # run flow for some time before releasing tracer

ADprob = TracerAdvectionDiffusion.Problem(MQGprob; κ, stepper, tracer_release_time)

Problem
  ├─────────── grid: grid (on CPU)
  ├───── parameters: params
  ├────── variables: vars
  ├─── state vector: sol
  ├─────── equation: eqn
  ├────────── clock: clock
  │                  └──── dt: 0.0025
  └──── timestepper: FilteredRK4TimeStepper

Some shortcuts for the advection-diffusion problem:

sol, clock, vars, params, grid = ADprob.sol, ADprob.clock, ADprob.vars, ADprob.params, ADprob.grid
x, y = grid.x, grid.y

(-3.141592653589793:0.04908738521234052:3.0925052683774528, -3.141592653589793:0.04908738521234052:3.0925052683774528)

Initial condition for concentration in both layers

We have a two layer system so we advect-diffuse the tracer in both layers. To do this we set the initial condition for tracer concetration as a Gaussian centered at the origin.

gaussian(x, y, σ) = exp(-(x^2 + y^2) / 2σ^2)

amplitude, spread = 10, 0.15
c₀ = [amplitude * gaussian(x[i], y[j], spread) for j=1:grid.ny, i=1:grid.nx]

TracerAdvectionDiffusion.set_c!(ADprob, c₀)

Saving output

The parent package FourierFlows.jl provides the functionality to save the output from our simulation. To do this we write a function get_concentration and pass this to the Output constructor along with the TracerAdvectionDiffusion.Problem and the name of the output file.

function get_concentration(prob)
  invtransform!(prob.vars.c, deepcopy(prob.sol), prob.params.MQGprob.params)

  return prob.vars.c
end

function get_streamfunction(prob)
  params, vars, grid = prob.params.MQGprob.params, prob.params.MQGprob.vars, prob.grid

  @. vars.qh = prob.params.MQGprob.sol

  streamfunctionfrompv!(vars.ψh, vars.qh, params, grid)

  invtransform!(vars.ψ, vars.ψh, params)

  return vars.ψ
end

output = Output(ADprob, "advection-diffusion.jld2",
                (:concentration, get_concentration), (:streamfunction, get_streamfunction))

Output
  ├──── prob: FourierFlows.Problem{DataType, Array{ComplexF64, 3}, Float64, Array{Float64, 3}}
  ├──── path: advection-diffusion_1.jld2
  └── fields: Dict{Symbol, Function}(:streamfunction => Main.var"Main".get_streamfunction, :concentration => Main.var"Main".get_concentration)

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 save_frequency timesteps; then we step the problem forward.

save_frequency = 50 # frequency 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: %.2f, walltime: %.2f min",
                   clock.step, clock.t, (time()-startwalltime) / 60)

    println(log)
  end

  stepforward!(ADprob)
  stepforward!(params.MQGprob)
  MultiLayerQG.updatevars!(params.MQGprob)
end

Output saved, step: 0000, t: 0.00, walltime: 0.00 min
Output saved, step: 0050, t: 0.13, walltime: 0.01 min
Output saved, step: 0100, t: 0.25, walltime: 0.01 min
Output saved, step: 0150, t: 0.38, walltime: 0.02 min
Output saved, step: 0200, t: 0.50, walltime: 0.02 min
Output saved, step: 0250, t: 0.62, walltime: 0.02 min
Output saved, step: 0300, t: 0.75, walltime: 0.03 min
Output saved, step: 0350, t: 0.87, walltime: 0.03 min
Output saved, step: 0400, t: 1.00, walltime: 0.04 min
Output saved, step: 0450, t: 1.12, walltime: 0.04 min
Output saved, step: 0500, t: 1.25, walltime: 0.05 min
Output saved, step: 0550, t: 1.37, walltime: 0.05 min
Output saved, step: 0600, t: 1.50, walltime: 0.06 min
Output saved, step: 0650, t: 1.62, walltime: 0.06 min
Output saved, step: 0700, t: 1.75, walltime: 0.07 min
Output saved, step: 0750, t: 1.87, walltime: 0.07 min
Output saved, step: 0800, t: 2.00, walltime: 0.08 min
Output saved, step: 0850, t: 2.12, walltime: 0.08 min
Output saved, step: 0900, t: 2.25, walltime: 0.08 min
Output saved, step: 0950, t: 2.37, walltime: 0.09 min
Output saved, step: 1000, t: 2.50, walltime: 0.09 min
Output saved, step: 1050, t: 2.62, walltime: 0.10 min
Output saved, step: 1100, t: 2.75, walltime: 0.10 min
Output saved, step: 1150, t: 2.87, walltime: 0.11 min
Output saved, step: 1200, t: 3.00, walltime: 0.11 min
Output saved, step: 1250, t: 3.12, walltime: 0.12 min
Output saved, step: 1300, t: 3.25, walltime: 0.12 min
Output saved, step: 1350, t: 3.37, walltime: 0.13 min
Output saved, step: 1400, t: 3.50, walltime: 0.13 min
Output saved, step: 1450, t: 3.62, walltime: 0.14 min
Output saved, step: 1500, t: 3.75, walltime: 0.14 min
Output saved, step: 1550, t: 3.87, walltime: 0.15 min
Output saved, step: 1600, t: 4.00, walltime: 0.15 min
Output saved, step: 1650, t: 4.12, walltime: 0.15 min
Output saved, step: 1700, t: 4.25, walltime: 0.16 min
Output saved, step: 1750, t: 4.37, walltime: 0.16 min
Output saved, step: 1800, t: 4.50, walltime: 0.17 min
Output saved, step: 1850, t: 4.63, walltime: 0.17 min
Output saved, step: 1900, t: 4.75, walltime: 0.18 min
Output saved, step: 1950, t: 4.88, walltime: 0.18 min
Output saved, step: 2000, t: 5.00, walltime: 0.19 min
Output saved, step: 2050, t: 5.13, walltime: 0.19 min
Output saved, step: 2100, t: 5.25, walltime: 0.20 min
Output saved, step: 2150, t: 5.38, walltime: 0.20 min
Output saved, step: 2200, t: 5.50, walltime: 0.21 min
Output saved, step: 2250, t: 5.63, walltime: 0.21 min
Output saved, step: 2300, t: 5.75, walltime: 0.22 min
Output saved, step: 2350, t: 5.88, walltime: 0.22 min
Output saved, step: 2400, t: 6.00, walltime: 0.22 min
Output saved, step: 2450, t: 6.13, walltime: 0.23 min
Output saved, step: 2500, t: 6.25, walltime: 0.23 min
Output saved, step: 2550, t: 6.38, walltime: 0.24 min
Output saved, step: 2600, t: 6.50, walltime: 0.24 min
Output saved, step: 2650, t: 6.63, walltime: 0.25 min
Output saved, step: 2700, t: 6.75, walltime: 0.25 min
Output saved, step: 2750, t: 6.88, walltime: 0.26 min
Output saved, step: 2800, t: 7.00, walltime: 0.26 min
Output saved, step: 2850, t: 7.13, walltime: 0.27 min
Output saved, step: 2900, t: 7.25, walltime: 0.27 min
Output saved, step: 2950, t: 7.38, walltime: 0.28 min
Output saved, step: 3000, t: 7.50, walltime: 0.28 min
Output saved, step: 3050, t: 7.63, walltime: 0.29 min
Output saved, step: 3100, t: 7.75, walltime: 0.29 min
Output saved, step: 3150, t: 7.88, walltime: 0.30 min
Output saved, step: 3200, t: 8.00, walltime: 0.30 min
Output saved, step: 3250, t: 8.13, walltime: 0.31 min
Output saved, step: 3300, t: 8.25, walltime: 0.31 min
Output saved, step: 3350, t: 8.38, walltime: 0.32 min
Output saved, step: 3400, t: 8.50, walltime: 0.32 min
Output saved, step: 3450, t: 8.63, walltime: 0.32 min
Output saved, step: 3500, t: 8.75, walltime: 0.33 min
Output saved, step: 3550, t: 8.88, walltime: 0.33 min
Output saved, step: 3600, t: 9.00, walltime: 0.34 min
Output saved, step: 3650, t: 9.13, walltime: 0.34 min
Output saved, step: 3700, t: 9.25, walltime: 0.35 min
Output saved, step: 3750, t: 9.38, walltime: 0.35 min
Output saved, step: 3800, t: 9.50, walltime: 0.36 min
Output saved, step: 3850, t: 9.63, walltime: 0.36 min
Output saved, step: 3900, t: 9.75, walltime: 0.37 min
Output saved, step: 3950, t: 9.88, walltime: 0.37 min
Output saved, step: 4000, t: 10.00, walltime: 0.38 min

Visualizing the output

We now have output from our simulation saved in advection-diffusion.jld2. As a demonstration, we load the JLD2 output and create a time series for the tracer in the lower layer of our fluid along with the flow streamlines.

Create time series for the concentration and streamfunction in the bottom layer, layer = 2.

file = jldopen(output.path)

iterations = parse.(Int, keys(file["snapshots/t"]))
t = [file["snapshots/t/$i"] for i ∈ iterations]

layer = 2

c = [file["snapshots/concentration/$i"][:, :, layer] for i ∈ iterations]
ψ = [file["snapshots/streamfunction/$i"][:, :, layer] for i ∈ iterations]

We normalize all streamfunctions to have maximum absolute value 1.

for i in 1:lastindex(ψ)
  ψ[i] /= maximum(abs, ψ[i])
end

x,  y  = file["grid/x"],  file["grid/y"]
Lx, Ly = file["grid/Lx"], file["grid/Ly"]

n = Observable(1)

c_anim = @lift Array(c[$n])
ψ_anim = @lift Array(ψ[$n])
title = @lift @sprintf("concentration, t = %.2f", t[$n])

fig = Figure(size = (600, 600))
ax = Axis(fig[1, 1],
          xlabel = "x",
          ylabel = "y",
          aspect = 1,
          title = title,
          limits = ((-Lx/2, Lx/2), (-Ly/2, Ly/2)))

hm = heatmap!(ax, x, y, c_anim;
              colormap = :balance, colorrange = (-1, 1))
contour!(ax, x, y, ψ_anim;
         levels = 0.1:0.2:1, color = :grey, linestyle = :solid)
contour!(ax, x, y, ψ_anim;
         levels = -0.1:-0.2:-1, color = (:grey, 0.8), linestyle = :dash)

Create a movie of the tracer with the streamlines.

frames = 1:length(t)
record(fig, "turbulentflow_advection-diffusion.mp4", frames, framerate = 12) do i
    n[] = i
end

This page was generated using Literate.jl.