This is a proof-of-concept Pluto notebook using AIBECS to run tiny simulations of a tracer forward and backwards in time, around the Gradients cruises. The goal is to provide a tool to investigate the movement/supply of nutrients.

Gradients_AllHands_Aug2020_AIBECS.jl

### A Pluto.jl notebook ###
# v0.11.10

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 5bf0f6f2-e8c7-11ea-2d18-af0006cfea10
begin
	using LinearAlgebra, Plots, PlutoUI
	Core.eval(Main, :(using AIBECS))
	using Unitful: yr, s, mol, m, d
	md"Alright, let's go! (Loading some packages in the background...)"
end

# ╔═╡ a0b3b796-e903-11ea-0820-55b8a1a4fb0f
begin
	Circulations = Dict(
		"OCCA" => OCCA,
		"OCIM0" => OCIM0,
		"OCIM1" => OCIM1,
		"OCIM2" => OCIM2
	)
	md"""
	# Step 1: Chose the circulation
	
	Please select one of these circulations:
	$(@bind Circ Select(collect(keys(Circulations)), default="OCCA"))
	"""
end

# ╔═╡ 825fe1f4-e8c7-11ea-3833-7736c1c23e65
begin
	grd, T, = Circulations[Circ].load()
	md"""
	You have selected the $Circ circulation. Below is a map of the surface grid, to give you an idea of its resolution
	"""
end

# ╔═╡ 89934ce4-e909-11ea-3353-0b279b8b42f1
surfacemap(zeros(count(iswet(grd))), grd, color=:blues, title="The $Circ grid (at the surface)")

# ╔═╡ 8a4a9276-e8fa-11ea-3135-13e615b6b14e
md"# Step 2: Chose the origin/destination"

# ╔═╡ c2d94974-e8f8-11ea-2065-15bca551c273
begin
	ilon = findmin(abs.(grd.lon .- 200u"°"))[2]
	ilats = findfirst(grd.lat .≥ 20u"°"):findlast(grd.lat .≤ 45u"°")
	md"""
	$(@bind ilat Slider(ilats, default=findmin(abs.(grd.lat .- 30u"°"))[2]))
	$(@bind idepth Slider(eachindex(grd.depth))) 
	"""
end

# ╔═╡ b0fac100-e8f9-11ea-0a2b-2b7987e5ce34
begin
	isim = findall((abs.(latvec(grd) .- ustrip(grd.lat[ilat])) .≤ 15) .&
	               (abs.(lonvec(grd) .- ustrip(grd.lon[ilon])) .≤ 30) .&
		           (depthvec(grd) .≤ 1000))

	md"""
	Gradients cruises 1 and 2 were ≈ ((20–45)°N, 200°E), so we fix the longitude at $(grd.lon[ilon]).

	Because we are not interested in the very long timescales, we can restrict the simulation to a smaller region to speed up computations.

	What we will do is simulate a tracer coming from (or going to) a given model box using the $Circ model.

	You can chose the latitude and depth of the origin/destination box with these sliders:
	"""
end

# ╔═╡ 9c9b7170-e8fc-11ea-3f98-e7a2b464a7b6
md"""
# Step 3: Simulate!

Chose the number of days to run the simulation for: $(@bind days Slider(0:10:1000, default=100, show_value=true))
"""

# ╔═╡ a6735abc-eb51-11ea-20a6-cf2258b372d8
begin
	latlim=(20, 65)
	lonlim=(100, 300)
	depthlim=(0, 1000)
	md"(some extras for the plot lims)"
end

# ╔═╡ d6692a5c-e8c8-11ea-3f93-618e8cf1f077
md"""

# The math and code behind this notebook

Let's start from the tracer equation,

$$\partial_t x +  \nabla \cdot \left[ \boldsymbol{u} - \mathbf{K} \nabla \right] x = 0,$$

for a passive, conservative tracer, $x(\boldsymbol{r}, t)$ transported by advection with mean velocity $\boldsymbol{u}(\boldsymbol{r},t)$ and by diffusive eddies parameterized via a the diffusivity tensor $\mathbf{K}(\boldsymbol{r},t)$.

In the framework of transport matrices, this equation is converted into a system of linear equations

$$\partial_t \boldsymbol{x} + \mathbf{T} \boldsymbol{x} = 0.$$

We then want to propagate a given initial (or final) condition, $\boldsymbol{x}_0$, forward (or backward) in time.

## One step forward

To simulate the tracer forward in time, we can use the Crank-Nicolson method (a type of Runge-Kutta) that is uncondionally stable, so that we don't have to worry too much about the time step size. We write it as

$$\frac{\boldsymbol{x}(t+\Delta t) - \boldsymbol{x}(t)}{\Delta t} + \mathbf{T} \frac{\boldsymbol{x}(t+\Delta t) + \boldsymbol{x}(t)}{2} = 0$$

such that

$$\boldsymbol{x}(t+\Delta t) = - \mathbf{A}^{-1} \, \mathbf{B} \, \boldsymbol{x}(t)$$

with 

$$\mathbf{A} = \mathbf{T} + \frac{2}{Δ t}\mathbf{I}$$

and

$$\mathbf{B} = \mathbf{T} - \frac{2}{Δ t}\mathbf{I}$$


## One step backwards

To go backwards in time, little care must be taken, because diffusion is entropic.
We have to propagate things in the time-reversed *adjoint* flow (or TRAF, emphasis on "adjoint").
Skipping the details here, but the adjoint of $\partial_t$ is $-\partial_t$, and the adjoint of $\mathbf{T}$ is $\tilde{\mathbf{T}} = \mathbf{V}^{-1} \mathbf{T}^\mathsf{T} \mathbf{V}$ where $\mathbf{V}$ is the diagonal matrix of the volumes of the ocean boxes (the inner product being the volume integral).
In other words, we must propagate things back in time using the adjoint tracer equation,

$$(-\partial_t + \tilde{\mathbf{T}}) \tilde{\boldsymbol{x}} = 0$$

We apply the Crank-Nicolson algorithm again:

$$-\frac{\tilde{\boldsymbol{x}}(t) - \tilde{\boldsymbol{x}}(t-\Delta t)}{\Delta t} + \tilde{\mathbf{T}} \frac{\tilde{\boldsymbol{x}}(t) + \tilde{\boldsymbol{x}}(t- \Delta t)}{2} = 0$$

such that

$$\tilde{\boldsymbol{x}}(t- \Delta t) = - \tilde{\mathbf{A}}^{-1} \tilde{\mathbf{B}} \tilde{\boldsymbol{x}}(t)$$

with 

$$\tilde{\mathbf{A}} = \tilde{\mathbf{T}} + \frac{2}{Δ t}\mathbf{I}$$

and

$$\tilde{\mathbf{B}} =\tilde{\mathbf{T}} - \frac{2}{Δ t}\mathbf{I}$$

which is pretty much the same as the forward step except $\mathbf{T}$ has been replaced by its adjoint $\tilde{\mathbf{T}}$.

## The code

The main reason we can leverage the stability of Crank-Nicolson without too much of a computation time penalty is because for a given time step $Δ t$, the matrices $\mathbf{A}$ and $\tilde{\mathbf{A}}$ can be factorized (LU) once, and then solving systems is fairly fast (about an order or two faster than the factorization). 

In Julia, this is as simple as

```julia
Af = factorize(A)
A \ (B * x)
```

for the forward case. You can reveal the code cells below to see the functions used here in this Pluto notebook!

The computation speed is also increased by removing a good fraction of the ocean when time-stepping, by restricting the simulation to just a smaller portion of the Pacific basin.

"""

# ╔═╡ e2685686-e9b1-11ea-3ca3-d761133ebe71
begin
	# Define the CN matrices with a function (notation different from text)
	function CrankNicolsonMatrices(T, v, isim, Δt)
		# Restrict to simulation region indices
		T = T[isim, isim]
		v = v[isim]
		V = sparse(Diagonal(v))
		V⁻¹	= sparse(Diagonal(1 ./ v))
		Tᵃ = V⁻¹ * T' * V
		T⁺ = T + 2I / Δt
		T⁻ = T - 2I / Δt
		Tᵃ⁺ = Tᵃ + 2I / Δt
		Tᵃ⁻ = Tᵃ - 2I / Δt
		T⁺f = factorize(T⁺)
		Tᵃ⁺f = factorize(Tᵃ⁺)
		return T⁺f, T⁻, Tᵃ⁺f, Tᵃ⁻
	end
	# functions to step forward and backward in place
	backstep!(f) = f .= Tᵃ⁺f \ (Tᵃ⁻ * -f)
	forwardstep!(x) = x .= T⁺f \ (T⁻ * -x)
	Δt = 1.0d
	T⁺f, T⁻, Tᵃ⁺f, Tᵃ⁻ = CrankNicolsonMatrices(T, volumevec(grd), isim, ustrip(upreferred(Δt)))
	# initial/final condition
	function x0()
		f03D = fill(0.0, size(grd))
		f03D[ilat, ilon, idepth] = 1.0
		return f03D[iswet(grd)]
	end
end

# ╔═╡ 6e54c012-eb56-11ea-073c-4be030b27efa
md"""
# Caveats

It's important to note that this is **just a tool** to get a feeling for where nutrients go and where they come from. And in the real ocean, it could be different.

You might also notice some negative values, that arise from numerical errors that can be caused by multiple factors. First is floating-point arithmetic errors. Second is the spatial resolution. Third is the temporal resolution (`Δt =` $Δt here). And I probably missed others!
"""

# ╔═╡ 7cc7f966-e8fb-11ea-183c-1336e499e093
# Forward and backward loops
begin
	function backsim(d)
		fout = x0()
		f = fout[isim]
		for i in 1:d
			backstep!(f)
		end
		fout[isim] .= f
		return fout
	end
	function forwardsim(d)
		xout = x0()
		x = xout[isim]
		for i in 1:d
			forwardstep!(x)
		end
		xout[isim] .= x
		return xout
	end
end

# ╔═╡ f6f44696-e9b1-11ea-2824-cb19e49025ef
# Cell actually running the simulations forward and backward,
# and storing the result in `x` and `f`.
begin
	f = backsim(days)
	x = forwardsim(days)
end

# ╔═╡ 2a772f4c-e9b2-11ea-31b3-7725da7fac93
# quick hacky code to draw rectangle of simulation region
begin
	rectangle(ex, ey) = rectangle(ex[2]-ex[1], ey[2]-ey[1], ex[1], ey[1])
	rectangle(w, h, x, y) = Shape(x .+ [0,w,w,0], y .+ [0,0,h,h])
	function addsimregion!(plt; kwargs...)
		δlon, δlat = ustrip(grd.δlon[1]), ustrip(grd.δlat[1])
		area = rectangle(extrema(lonvec(grd)[isim]) .+ 0.5 .* (-δlon, δlon), extrema(latvec(grd)[isim]) .+ 0.5 .* (-δlat, δlat))
		plot!(plt, area, label="Simulation region", la=1, lw=1, α=0; kwargs...)
	end
end

# ╔═╡ 9ab2737e-e8fa-11ea-3962-51c2e36b8205
begin
	plt = surfacemap(zeros(count(iswet(grd))), grd, color=:blues, title="Destination/origin at $(round(grd.lon[ilon], digits=1))E, $(round(grd.lat[ilat], digits=1))N (depth=$(round(grd.depth[idepth], digits=1)))")
	vspan!(plt, grd.lon[ilon] .+ [-1, 1] * 0.5grd.δlon[ilon], lab="$(round(grd.lon[ilon], digits=1))E", c=:red, α=0.3)
	hspan!(plt, grd.lat[ilat] .+ [-1, 1] * 0.5grd.δlat[ilat], lab="$(round(grd.lat[ilat], digits=1))N", c=:green, α=0.3)
	hspan!(plt, [grd.lat[ilats[1]] - 0.5grd.δlat[ilats[1]], grd.lat[ilats[end]] + 0.5grd.δlat[ilats[end]]], lab="Gradients lat range", c=:green, α=0.1)
	addsimregion!(plt)
end

# ╔═╡ 1b1d7bb0-e8fd-11ea-14a0-d1f54e291c36
begin
	pltforward = plotverticalmean(x, grd, title="Going forward $days days", xlim=lonlim, ylim=latlim, cbar=false)
	vspan!(pltforward, grd.lon[ilon] .+ [-1, 1] * 0.5grd.δlon[ilon], lab="", c=:green, α=0.3)
	hspan!(pltforward, grd.lat[ilat] .+ [-1, 1] * 0.5grd.δlat[ilat], lab="", c=:green, α=0.3)
	hspan!(pltforward, [grd.lat[ilats[1]] - 0.5grd.δlat[ilats[1]], grd.lat[ilats[end]] + 0.5grd.δlat[ilats[end]]], lab="", c=:green, α=0.1)
	addsimregion!(pltforward, lc=:green)
	MSforward = plotmeridionalslice(x, grd, lon=grd.lon[ilon], title="Meridional slice", ylim=depthlim, xlim=latlim, cbar=false)
	vspan!(MSforward, ustrip.(grd.lat[ilat] .+ [-1, 1] * 0.5grd.δlat[ilat]), lab="", c=:green, α=0.3)
	hspan!(MSforward, [0; ustrip.(cumsum(grd.δdepth))][idepth:idepth+1], lab="", c=:green, α=0.3)
	ZSforward = plotzonalslice(x, grd, lat=grd.lat[ilat], title="Zonal slice", ylim=depthlim, xlim=lonlim, cbar=false)
	vspan!(ZSforward, ustrip.(grd.lon[ilon] .+ [-1, 1] * 0.5grd.δlon[ilon]), lab="", c=:green, α=0.3)
	hspan!(ZSforward, [0; ustrip.(cumsum(grd.δdepth))][idepth:idepth+1], lab="", c=:green, α=0.3)
	plot(pltforward, MSforward, ZSforward, layout=@layout [a{0.7w} [b; c]])
end

# ╔═╡ 4cdbd200-e9b1-11ea-0671-49886101557e
begin
	pltback = plotverticalmean(f, grd, title="Going back $days days", xlim=lonlim, ylim=latlim, cbar=false)
	vspan!(pltback, grd.lon[ilon] .+ [-1, 1] * 0.5grd.δlon[ilon], lab="", c=:green, α=0.3)
	hspan!(pltback, grd.lat[ilat] .+ [-1, 1] * 0.5grd.δlat[ilat], lab="", c=:green, α=0.3)
	hspan!(pltback, [grd.lat[ilats[1]] - 0.5grd.δlat[ilats[1]], grd.lat[ilats[end]] + 0.5grd.δlat[ilats[end]]], lab="", c=:green, α=0.1)
	addsimregion!(pltback, lc=:green)
	MSback = plotmeridionalslice(f, grd, lon=grd.lon[ilon], title="Meridional slice", ylim=depthlim, xlim=latlim, cbar=false)
	vspan!(MSback, ustrip.(grd.lat[ilat] .+ [-1, 1] * 0.5grd.δlat[ilat]), lab="", c=:green, α=0.3)
	hspan!(MSback, [0; ustrip.(cumsum(grd.δdepth))][idepth:idepth+1], lab="", c=:green, α=0.3)
	ZSback = plotzonalslice(f, grd, lat=grd.lat[ilat], title="Zonal slice", ylim=depthlim, xlim=lonlim, cbar=false)
	vspan!(ZSback, ustrip.(grd.lon[ilon] .+ [-1, 1] * 0.5grd.δlon[ilon]), lab="", c=:green, α=0.3)
	hspan!(ZSback, [0; ustrip.(cumsum(grd.δdepth))][idepth:idepth+1], lab="", c=:green, α=0.3)
	plot(pltback, MSback, ZSback, layout=@layout [a{0.7w} [b; c]])
end

# ╔═╡ Cell order:
# ╠═5bf0f6f2-e8c7-11ea-2d18-af0006cfea10
# ╟─a0b3b796-e903-11ea-0820-55b8a1a4fb0f
# ╟─825fe1f4-e8c7-11ea-3833-7736c1c23e65
# ╟─89934ce4-e909-11ea-3353-0b279b8b42f1
# ╟─8a4a9276-e8fa-11ea-3135-13e615b6b14e
# ╟─b0fac100-e8f9-11ea-0a2b-2b7987e5ce34
# ╟─c2d94974-e8f8-11ea-2065-15bca551c273
# ╟─9ab2737e-e8fa-11ea-3962-51c2e36b8205
# ╟─9c9b7170-e8fc-11ea-3f98-e7a2b464a7b6
# ╟─1b1d7bb0-e8fd-11ea-14a0-d1f54e291c36
# ╟─4cdbd200-e9b1-11ea-0671-49886101557e
# ╟─a6735abc-eb51-11ea-20a6-cf2258b372d8
# ╟─d6692a5c-e8c8-11ea-3f93-618e8cf1f077
# ╟─6e54c012-eb56-11ea-073c-4be030b27efa
# ╟─e2685686-e9b1-11ea-3ca3-d761133ebe71
# ╟─7cc7f966-e8fb-11ea-183c-1336e499e093
# ╟─f6f44696-e9b1-11ea-2824-cb19e49025ef
# ╟─2a772f4c-e9b2-11ea-31b3-7725da7fac93