text:
# expm.jl
# This file is based on
# https://gist.github.com/sdewaele/2c176cb634280cf8a23c5970739cea0e
# with minor modifications to allow ForwardDiff.jl to differentiate through
# `frexp`, and major modifications to use in-place operations and
# BlochSim.jl-specific structs

function expmchk()
# EXPMCHK Check the class of input A and
#    initialize M_VALS and THETA accordingly.
    m_vals = (3, 5, 7, 9, 13)
    theta = (
        # 3.650024139523051e-008
        # 5.317232856892575e-004
        1.495585217958292e-002, # m_vals = 3
        # 8.536352760102745e-002
        2.539398330063230e-001, # m_vals = 5
        # 5.414660951208968e-001
        9.504178996162932e-001, # m_vals = 7
        # 1.473163964234804e+000
        2.097847961257068e+000, # m_vals = 9
        # 2.811644121620263e+000
        # 3.602330066265032e+000
        # 4.458935413036850e+000
        5.371920351148152e+000  # m_vals = 13
    )

    return m_vals, theta
end

function getPadeCoefficients(m)
# GETPADECOEFFICIENTS Coefficients of numerator P of Pade approximant
#    C = GETPADECOEFFICIENTS returns coefficients of numerator
#    of [M/M] Pade approximant, where M = 3,5,7,9,13.
    if m == 3
        c = (120, 60, 12, 1)
    elseif m == 5
        c = (30240, 15120, 3360, 420, 30, 1)
    elseif m == 7
        c = (17297280, 8648640, 1995840, 277200, 25200, 1512, 56, 1)
    elseif m == 9
        c = (17643225600, 8821612800, 2075673600, 302702400, 30270240,
             2162160, 110880, 3960, 90, 1)
    elseif m == 13
        c = (64764752532480000, 32382376266240000, 7771770303897600,
             1187353796428800,  129060195264000,   10559470521600,
             670442572800,      33522128640,       1323241920,
             40840800,          960960,            16380,  182,  1)
    end

    return c
end

struct MatrixExponentialWorkspace{T<:Real,N}
    expA2::BlochMcConnellMatrix{T,N}
    A2::BlochMcConnellMatrix{T,N}
    A4::BlochMcConnellMatrix{T,N}
    A6::BlochMcConnellMatrix{T,N}
    A8::BlochMcConnellMatrix{T,N}
    tmp1::BlochMcConnellMatrix{T,N}
    tmp2::BlochMcConnellMatrix{T,N}
    U::BlochMcConnellMatrix{T,N}
    V::BlochMcConnellMatrix{T,N}
    mat1::Matrix{T}
    mat2::Matrix{T}
end

MatrixExponentialWorkspace{T}(N) where {T} =
    MatrixExponentialWorkspace(BlochMcConnellMatrix{T}(N),
                               BlochMcConnellMatrix{T}(N),
                               BlochMcConnellMatrix{T}(N),
                               BlochMcConnellMatrix{T}(N),
                               BlochMcConnellMatrix{T}(N),
                               BlochMcConnellMatrix{T}(N),
                               BlochMcConnellMatrix{T}(N),
                               BlochMcConnellMatrix{T}(N),
                               BlochMcConnellMatrix{T}(N),
                               Matrix{T}(undef, 3N, 3N),
                               Matrix{T}(undef, 3N, 3N))

"""
    expm!(expA, A, [workspace])

Compute the matrix exponential of `A`, storing it in `expA`.

`workspace isa MatrixExponentialWorkspace`.
"""
function expm!(
    expA::BlochMcConnellMatrix{T1,N},
    A::BlochMcConnellDynamicsMatrix{T2,N,M},
    workspace::MatrixExponentialWorkspace{T3,N} = MatrixExponentialWorkspace{T1}(N)
) where {T1,T2,T3,N,M}
# EXPM   Matrix exponential.
#   EXPM(X) is the matrix exponential of X.  EXPM is computed using
#   a scaling and squaring algorithm with a Pade approximation.
#
# Julia implementation closely based on MATLAB code by Nicholas Higham

    # Initialization
    (m_vals, theta) = expmchk()

    normA = absolutesum(A)

    if normA <= theta[end]
        # no scaling and squaring is required
        for i = 1:length(m_vals)
            if normA <= theta[i]
                PadeApproximantOfDegree!(expA, A, workspace, m_vals[i])
                break
            end
        end
    else
        tmp = normA / theta[end]
        t = frexp1(tmp)
        s = frexp2(tmp)
        s = s - (t == 0.5) # adjust s if normA / theta[end] is a power of 2
        mul!(A, 1 / 2^s) # Scaling
        PadeApproximantOfDegree!(expA, A, workspace, m_vals[end])

        for i = 1:s
            mul!(workspace.expA2, expA, expA) # Squaring
            copyto!(expA, workspace.expA2)
        end
    end

end

function PadeApproximantOfDegree!(
    expA::BlochMcConnellMatrix{T1,N},
    A::BlochMcConnellDynamicsMatrix{T2,N,M},
    workspace::MatrixExponentialWorkspace{T3,N},
    m::Integer
) where {T1,T2,T3,N,M}
#PADEAPPROXIMANTOFDEGREE  Pade approximant to exponential.
#   F = PADEAPPROXIMANTOFDEGREE(M) is the degree M diagonal
#   Pade approximant to EXP(A), where M = 3, 5, 7, 9 or 13.
#   Series are evaluated in decreasing order of powers, which is
#   in approx. increasing order of maximum norms of the terms.

    n = 3N
    c = getPadeCoefficients(m)

    mul!(workspace.A2, A, A)
    mul!(workspace.A4, workspace.A2, workspace.A2)
    mul!(workspace.A6, workspace.A2, workspace.A4)

    # Evaluate Pade approximant
    if m == 13
        # For optimal evaluation need different formula for m >= 12
        mul!(workspace.tmp1, workspace.A6, c[14])
        muladd!(workspace.tmp1, workspace.A4, c[12])
        muladd!(workspace.tmp1, workspace.A2, c[10])
        mul!(workspace.tmp2, workspace.A6, workspace.tmp1)
        muladd!(workspace.tmp2, workspace.A6, c[8])
        muladd!(workspace.tmp2, workspace.A4, c[6])
        muladd!(workspace.tmp2, workspace.A2, c[4])
        muladd!(workspace.tmp2, I, c[2])
        mul!(workspace.U, A, workspace.tmp2)

        mul!(workspace.tmp1, workspace.A6, c[13])
        muladd!(workspace.tmp1, workspace.A4, c[11])
        muladd!(workspace.tmp1, workspace.A2, c[9])
        mul!(workspace.V, workspace.A6, workspace.tmp1)
        muladd!(workspace.V, workspace.A6, c[7])
        muladd!(workspace.V, workspace.A4, c[5])
        muladd!(workspace.V, workspace.A2, c[3])
        muladd!(workspace.V, I, c[1])
    else # m == 3, 5, 7, 9
        fill!(workspace.tmp1, zero(T3))
        fill!(workspace.V, zero(T3))

        if m >= 9
            mul!(workspace.A8, workspace.A2, workspace.A6)
            muladd!(workspace.tmp1, workspace.A8, c[10])
            muladd!(workspace.V, workspace.A8, c[9])
        end
        if m >= 7
            muladd!(workspace.tmp1, workspace.A6, c[8])
            muladd!(workspace.V, workspace.A6, c[7])
        end
        if m >= 5
            muladd!(workspace.tmp1, workspace.A4, c[6])
            muladd!(workspace.V, workspace.A4, c[5])
        end
        muladd!(workspace.tmp1, workspace.A2, c[4])
        muladd!(workspace.V, workspace.A2, c[3])
        muladd!(workspace.tmp1, I, c[2])
        muladd!(workspace.V, I, c[1])
        mul!(workspace.U, A, workspace.tmp1)
    end

    subtract!(workspace.mat1, workspace.V, workspace.U)
    add!(workspace.mat2, workspace.V, workspace.U)
    F = lu!(workspace.mat1)
    ldiv!(F, workspace.mat2)
    copyto!(expA, workspace.mat2)

end

frexp1(x) = frexp(x)[1]
frexp2(x) = frexp(x)[2]
dfrexp1(x) = 2.0^(-floor(log2(abs(x))) - 1)
dfrexp2(x) = 0

for (f, df) in ((:frexp1, :dfrexp1), (:frexp2, :dfrexp2))
@eval begin
    function ($f)(d::ForwardDiff.Dual{T}) where T

        x = ForwardDiff.value(d)
        y = ($f)(x)
        dy = ($df)(x)
        return ForwardDiff.Dual{T}(y, dy * ForwardDiff.partials(d))

    end
end
end


text:
# export fourier_nodes, fourier_wavenumbers, cheb
using LinearAlgebra, FFTW

"""
fourier_nodes(n; a = 0, b = 2π)
# Description
- Create a uniform grid of points for periodic functions
# Arguments
- `N`: integer | number of evenly spaced points
# Keyword Arguments
- `a`: number | starting point of interval [a, b)
- `b`: number | ending point of interval [a, b)
# Return
- `g`: array | an array of points of evenly spaced points from [a, b)
"""
function fourier_nodes(N; a = 0, b = 2π)
    return (b - a) .* collect(0:(N - 1)) / N .+ a
end

"""
fourier_wavenumbers(N; L = 2π)
# Description
- Create wavenumbers associated with the domain of length L
# Arguments
- `N`: integer | number of wavevectors
# Keyword Arguments
- `L`: number | length of interval [a, b), L = b-a
# Return
- `wavenumbers`: array | an array of wavevectors
"""
function fourier_wavenumbers(N; L = 2π)
    up = collect(0:1:(N - 1))
    down = collect((-N):1:-1)
    indices = up
    indices[(div(N, 2) + 1):end] = down[(div(N, 2) + 1):end]
    wavenumbers = 2π / L .* indices
    return wavenumbers
end

"""
# Description
Spectral differentiation matrix and nodes for periodic domains
# Argument
- 'N': number of gridpoints
# Keyword Argument
a and b specify the interval [a, b)
default assumes x ∈ [0, 2π)
# Return
- 'D': Fourier differentiation matrix
- 'x': Fourier grid points
"""
function fourier(N; a = 0, b = 2π)
    if N == 0
        return [0], [0]
    else
        k = fourier_wavenumbers(N, L = b - a)
        x = fourier_nodes(N, a = a, b = b)
        ℱ = fft(I + zeros(N, N), 1)
        ℱ⁻¹ = ifft(I + zeros(N, N), 1)
        D = real.(ℱ⁻¹ * Diagonal(im .* k) * ℱ)
        return D, x
    end
end

"""
# Description
Julia version of Spectral Methods in Matlab
# Argument
- 'N': polynomial order
# Keyword Argument
a and b specify the interval [a, b]
default assumes x ∈ [-1, 1]
# Return
- 'D': Chebyshev differentiation matrix
- 'x': Guass-Lobatto points
"""
function chebyshev(N; a = -1, b = 1)
    if N == 0
        return [0], [(a + b) / 2]
    else
        x = @. cos(pi * (0:N) / N)
        c = [2; ones(N - 1); 2] .* (-1) .^ (0:N)
        dX = x .- x'
        D = (c ./ c') ./ (dX + I)                # off-diagonal entries
        D = D - Diagonal(sum(D', dims = 1)[1:(N + 1)]) # diagonal entries
        return 2 / (b - a) * D, (b - a) * (x .+ 1) / 2 .+ a
    end
end

ext: jl

max_stars_repo_name: sandreza/StatisticalNonlocality

max_issues_repo_name: sandreza/StatisticalNonlocality

max_forks_repo_name: sandreza/StatisticalNonlocality


text:
using Flux
using Flux: glorot_uniform, normalise, @functor# , destructure
using Zygote: @adjoint, @nograd
using LinearAlgebra, SparseArrays
using Statistics
using ChemistryFeaturization
using DeepUncertainty

struct MCAGNConv{W,F}
    selfweight::W
    convweight::W
    bias::W
    σ::F
    dropout
end

function MCAGNConv(
    ch::Pair{<:Integer,<:Integer},
    dropout_rate = 0.1,
    σ = softplus;
    initW = glorot_uniform,
    initb = zeros,
    T::DataType = Float64,
)
    selfweight = T.(initW(ch[2], ch[1]))
    convweight = T.(initW(ch[2], ch[1]))
    b = T.(initb(ch[2], 1))
    dropout = (x; k...) -> Flux.dropout(x, dropout_rate; k...)
    MCAGNConv(selfweight, convweight, b, σ, dropout)
end

@functor MCAGNConv

"""
 Define action of layer on inputs: do a graph convolution, add this (weighted by convolutional weight) to the features themselves (weighted by self weight) and the per-feature bias (concatenated to match number of nodes in graph).
# Arguments
- input: a FeaturizedAtoms object, or graph_laplacian, encoded_features
# Note
In the case of providing two matrices, the following conditions must hold:
- `lapl` must be square and of dimension N x N where N is the number of nodes in the graph
- `X` (encoded features) must be of dimension M x N, where M is `size(l.convweight)[2]` (or equivalently, `size(l.selfweight)[2]`)
"""
function (l::MCAGNConv)(lapl, X)
    # should we put dimension checks here? Could allow more informative errors, but would likely introduce performance penalty. For now it's just in docstring.
    out_mat = normalise(
        l.σ.(
            l.convweight * X * lapl +
            l.selfweight * X +
            reduce(hcat, l.bias for i = 1:size(X, 2)),
        ),
        dims = [1, 2],
    )
    out_mat = l.dropout(out_mat)
    lapl, out_mat
end

# alternate signature so FeaturizedAtoms can be fed into first layer
(l::MCAGNConv)(a::FeaturizedAtoms{AtomGraph,GraphNodeFeaturization}) =
    l(a.atoms.laplacian, a.encoded_features)

# signature to splat appropriately
(l::MCAGNConv)(t::Tuple{Matrix{R1},Matrix{R2}}) where {R1<:Real,R2<:Real} = l(t...)

# Bayesian Atomic Graph CNN
struct VariationalAGNConv{W,F}
    selfweight_sampler::W
    convweight_sampler::W
    bias_smapler::W
    σ::F
end

function VariationalAGNConv(
    ch::Pair{<:Integer,<:Integer},
    σ = softplus;
    initW = glorot_uniform,
    initb = zeros,
    T::DataType = Float64,
)
    selfweight = T.(initW(ch[2], ch[1]))
    selfweight_sampler = TrainableMvNormal(size(selfweight))
    convweight = T.(initW(ch[2], ch[1]))
    convweight_sampler = TrainableMvNormal(size(convweight))
    b = T.(initb(ch[2], 1))
    bias_sampler = TrainableMvNormal(size(b))
    VariationalAGNConv(selfweight_sampler, convweight_sampler, bias_sampler, σ)
end

@functor VariationalAGNConv

"""
 Define action of layer on inputs: do a graph convolution, add this (weighted by convolutional weight) to the features themselves (weighted by self weight) and the per-feature bias (concatenated to match number of nodes in graph).
# Arguments
- input: a FeaturizedAtoms object, or graph_laplacian, encoded_features
# Note
In the case of providing two matrices, the following conditions must hold:
- `lapl` must be square and of dimension N x N where N is the number of nodes in the graph
- `X` (encoded features) must be of dimension M x N, where M is `size(l.convweight)[2]` (or equivalently, `size(l.selfweight)[2]`)
"""
function (l::VariationalAGNConv)(lapl, X)
    # should we put dimension checks here? Could allow more informative errors, but would likely introduce performance penalty. For now it's just in docstring.
    convweight = l.convweight_sampler()
    selfweight = l.selfweight_sampler()
    bias = l.bias_smapler()
    out_mat = normalise(
        l.σ.(
            convweight * X * lapl +
            selfweight * X +
            reduce(hcat, bias for i = 1:size(X, 2)),
        ),
        dims = [1, 2],
    )
    lapl, out_mat
end

# alternate signature so FeaturizedAtoms can be fed into first layer
(l::VariationalAGNConv)(a::FeaturizedAtoms{AtomGraph,GraphNodeFeaturization}) =
    l(a.atoms.laplacian, a.encoded_features)

# signature to splat appropriately
(l::VariationalAGNConv)(t::Tuple{Matrix{R1},Matrix{R2}}) where {R1<:Real,R2<:Real} = l(t...)


text:
module ErgodicControlPlots

const T2F = NTuple{2, Float64}    # x, y
const MF = Matrix{Float64}
const VMF = Vector{MF}
const VF = Vector{Float64}
const VVF = Vector{VF}
const VVVF = Vector{VVF}

# package code goes here
using ErgodicControl

const VTM = Vector{TrajectoryManager}

include("plots.jl")
include("gif.jl")

end # module


text:
function finddifference(z, p, s, dmax)
    Δz = z
    d = 0
    while d < dmax && pvalue(ADFTest(Δz, :none, p)) > 0.05 d += 1; Δz = difference(z, d=d, s=s)  end
    return d
end

function boxjenkins(z, c::Bool, pmax::Integer, dmax::Integer, qmax::Integer, s::AbstractVector{<:Integer}; criterium=aicc)
    λ = guerrero(z)
    λ = λ < 0.2 ? 0.0 : λ
    z = boxcox(z, λ)
    d = [finddifference(z, pmax, s, dmax) for s in s]
    params = AbstractParams[]
    S = length(s)
    for i in 0:pmax, j in 0:qmax
        push!(params, MSARIMAParams(c, fill(i, S), d, fill(j, S), s))
    end
    return fit(params, z, criterium=criterium)
end


text:
function great_painter(r::Robot)
    num=0
    actions=[]
    while ((isborder(r,Sud))&&(isborder(r,West)))==false
        push!(actions,moves!(r,West))
        push!(actions,moves!(r,Sud))
        num+=2
    end
    num_ver=0
    num_hor=0
    for i in 1:num
        if isodd(i)==true
            num_hor+=actions[i]
        else
            num_ver+=actions[i]
        end
    end
    side=Nord
    for _ in 1:2
        num_ver=mark_kord(r,side,num_ver)
        side=next(side)
        num_hor=mark_kord(r,side,num_hor)
        side=next(side)
    end
    while (num>0)==true
        side=isodd(num) ? Ost : Nord
        move_steps(r,side,actions[num])
        num-=1
    end
end

function moves!(r::Robot,side::HorizonSide)
    num_steps=0
    while isborder(r,side)==false
        move!(r,side)
        num_steps+=1
    end
    return num_steps
end

function mark_kord(r::Robot, side::HorizonSide, num_steps::Int)
    move_steps(r,side,num_steps)
    putmarker!(r)
    num_steps=moves!(r,side)
    return num_steps
end

function move_steps(r::Robot, side::HorizonSide, num_steps::Int)
    for _ in 1:num_steps
        move!(r,side)
    end
end

next(side::HorizonSide)=HorizonSide(mod(Int(side)+3,4))

max_stars_repo_name: TheDarkMoon/mirea-progs

max_issues_repo_name: TheDarkMoon/mirea-progs

max_forks_repo_name: TheDarkMoon/mirea-progs


text:
### A Pluto.jl notebook ###
# v0.12.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

# ╔═╡ 9c8a7e5a-12dd-11eb-1b99-cd1d52aefa1d
begin
	import Pkg
	Pkg.activate(mktempdir())
	Pkg.add([
		"Plots",
		"PlutoUI",
		"Images",
		"FileIO",
		"ImageMagick",
		"ImageIO",
		"OffsetArrays"
	])
	using Statistics
	using Plots
	using PlutoUI
	using Images
	using OffsetArrays
end

# ╔═╡ 0f8db6f4-2113-11eb-18b4-21a469c67f3a
md"""
### Lecture 23: Solving Partial Differential Equations (PDEs) Numerically
**Part II: Heat transport by ocean currents (two-dimensional advection and diffusion)**

Guest Lecturer: Henri F. Drake (MIT Climate Science and Oceanography PhD Student)
"""

# ╔═╡ ed741ec6-1f75-11eb-03be-ad6284abaab8
html"""
<iframe width="700" height="394" src="https://www.youtube-nocookie.com/embed/6_GQuVopmUM?start=15" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture" allowfullscreen></iframe>
"""

# ╔═╡ ac759b96-2114-11eb-24cb-d50b556f4142
md"""
### 1) Background: two-dimensional advection-diffusion

##### 1.1) The two-dimensional advection-diffusion equation
Recall from **Lecture 22** that the one-dimensional advection-diffusion equation is written as

$\frac{\partial T(x,t)}{\partial t} = -U \frac{\partial T}{\partial x} + \kappa \frac{\partial^{2} T}{\partial x^{2}},$

where $T(x, t)$ is the temperature, $U$ is a constant *advective velocity* and $\kappa$ is the *diffusivity*.

The two-dimensional advection diffusion equation adds advection and diffusion operators acting in a *second spatial dimension* $y$ (orthogonal to $x$):

$\frac{\partial T(x,y,t)}{\partial t} = u(x,y) \frac{\partial T}{\partial x} + v(x,y) \frac{\partial T}{\partial y} + \kappa \left( \frac{\partial^{2} T}{\partial x^{2}} + \frac{\partial^{2} T}{\partial y^{2}} \right),$

where $\vec{u}(x,y) = (u, v) = u\,\mathbf{\hat{x}} + v\,\mathbf{\hat{y}}$ is a velocity vector field.

Throughout the rest of the Climate Modelling module, we will consider $x$ to be the *longitundinal* direction (positive from west to east) and $y$ to the be the *latitudinal* direction (positive from south to north).
"""

# ╔═╡ 3a4a1aea-2118-11eb-30a9-57b87f2ddfae
begin
	reviewBox = @bind show_review CheckBox(default=false)
	md"""
	##### 1.2) Review of multivariable calculus identities and notation

	*Check the box for a review* $(reviewBox)

	"""
end

# ╔═╡ 023779a0-2a95-11eb-35b5-7be93c43afaf
if show_review
	md"""
	Conventionally, the two-dimensional advection-diffusion equation is written more succintly as

	$\frac{\partial T(x,y,t)}{\partial t} = - \vec{u} \cdot \nabla T + \kappa \nabla^{2} T,$

	using the following shorthand notation.

	The **gradient** operator is defined as

	$\nabla \equiv (\frac{\partial}{\partial x}, \frac{\partial }{\partial y})$

	such that

	$\nabla T = (\frac{\partial T}{\partial x}, \frac{\partial T}{\partial y})$ and

	$\vec{u} \cdot \nabla T = (u, v) \cdot (\frac{\partial T}{\partial x}, \frac{\partial T}{\partial y}) = u \frac{\partial T}{\partial x} + v\frac{\partial T}{\partial y}.$

	The **Laplacian** operator $\nabla^{2}$ (sometimes denoted $\Delta$) is defined as

	$\nabla^{2} = \frac{\partial^{2}}{\partial x^{2}} + \frac{\partial^{2}}{\partial y^{2}}$

	such that

	$\nabla^{2} T = \frac{\partial^{2} T}{\partial x^{2}} + \frac{\partial^{2} T}{\partial y^{2}}.$

	The **divergence** operator is defined as $\nabla \cdot [\quad]$, such that

	$\nabla \cdot \vec{u} = \left(\frac{\partial}{\partial x}, \frac{\partial}{\partial x} \right) \cdot (u,v) = \frac{\partial u}{\partial x} + \frac{\partial v}{\partial y}.$

	**Note:** Since seawater is largely incompressible, we can approximate ocean currents as a *non-divergent flow*, with $\nabla \cdot \vec{u} = \frac{\partial u}{\partial x} + \frac{\partial v}{\partial y} = 0$. Among other implications, this allows us to write:

	\begin{align}
	\vec{u} \cdot \nabla T&=
	u\frac{\partial T(x,y,t)}{\partial x} + v\frac{\partial T(x,y,t)}{\partial y}\newline &=
	u\frac{\partial T}{\partial x} + v\frac{\partial T}{\partial y} + T\left(\frac{\partial u}{\partial x} + \frac{\partial v}{\partial y}\right)\newline &=
	\left( u\frac{\partial T}{\partial x} + T\frac{\partial u}{\partial x} \right) +
	\left( v\frac{\partial T}{\partial y} + \frac{\partial v}{\partial y} \right)
	\newline &=
	\frac{\partial (uT)}{\partial x} + \frac{\partial (vT)}{\partial x}\newline &=
	\nabla \cdot (\vec{u}T)
	\end{align}

	using the product rule (separately in both $x$ and $y$).

	##### 1.3) The flux-form two-dimensional advection-diffusion equation

	This lets us finally re-write the two-dimensional advection-diffusion equation as:

	$\frac{\partial T}{\partial t} = - \nabla \cdot (\vec{u}T) + \kappa \nabla^{2} T$

	which is the form we will use in our numerical algorithm below.
	"""
end

# ╔═╡ b1b5625e-211a-11eb-3ee1-3ba9c9cc375a
md"""
### 2) Numerical implementation

##### 2.1) Discretizing advection in two dimensions

In Lecture XX we saw that in one dimension we can discretize a first-partial derivative in space using the *centered finite difference*

$\frac{\partial T(x_{i}, t_{n})}{\partial x} \approx \frac{T_{i+1}^{n} - T_{i-1}^{n}}{2 \Delta x}.$

In two dimensions, we discretize the partial derivative the exact same way, except that we also need to keep track of the cell index $j$ in the $y$ dimension:

$\frac{\partial T(x_{i}, y_{j}, t_{n})}{\partial x} \approx \frac{T_{i+1,\, j}^{n} - T_{i-1,\,j}^{n}}{2 \Delta x}.$

The *x-gradient kernel* below, implemented using the `OffsetArray` type, is shown below, and is reminiscent of the *edge-detection* or *sharpening* kernels used in image processing and machine learning:
"""

# ╔═╡ 3578b158-2a97-11eb-0771-bf6d82d3b6d1
md"""
The first-order partial derivative in $y$ is similarly discretized as:

$\frac{\partial T(x_{i}, y_{j}, t_{n})}{\partial y} \approx \frac{T_{i,\, j+1}^{n} - T_{i,\,j-1}^{n}}{2 \Delta y}.$

Its kernel is shown below.
"""

# ╔═╡ 7f3c9550-2a97-11eb-1549-455009025872
md"""
Now that we have discretized the two derivate terms, we can write out the *advective tendency* for computing $T_{i, j, n+1}$ as

$u\frac{\partial T}{\partial x} + v\frac{\partial T}{\partial y} \approx u_{i,\, j}^{n} \frac{T_{i,\, j+1}^{n} - T_{i,\,j-1}^{n}}{2 \Delta y} + v_{i,\, j}^{n} \frac{T_{i,\, j+1}^{n} - T_{i,\,j-1}^{n}}{2 \Delta y}.$

We implement this in julia as a series of methods for the `advect` function. The first method computes the advective tendency (as a single `Float64` type) for the $(i,j)$ grid cell while the second method returns an array of the tendencies for each grid cell using two nested for loops.
"""

# ╔═╡ 0127bca6-2a99-11eb-16a0-8d7af66694f8
md"""
##### 2.2) Discretizing diffusion in two dimensions

Just as with advection, the process for discretizing the diffusion operators effectively consists of repeating the one-dimensional process for $x$ in a second dimension $y$, separately:

$\kappa \left( \frac{\partial^{2} T}{\partial x^{2}} + \frac{\partial^{2} T}{\partial y^{2}} \right) = \kappa \left(
\frac{T_{i+1,\;j}^{n} - 2T_{i,\;j}^{n} + T_{i-1,\;j}^{n}}{\left( \Delta x \right)^{2}} +
\frac{T_{i,\;j+1}^{n} - 2T_{i,\;j}^{n} + T_{i,\;j-1}^{n}}{\left( \Delta y \right)^{2}}
\right)$

The corresponding $x$ and $y$ second-derivative kernels are shown below:
"""

# ╔═╡ eac507ce-2a99-11eb-3eba-0780a4a7e078
md"""
Just as we did with advection, we implement a `diffuse` function using multiple dispatch:
"""

# ╔═╡ 09f179c0-2a9a-11eb-1d0f-e59012f9e77b
md"""##### 2.3) No-flux boundary conditions

We want to impose the no-flux boundary conditions, which states that

$u\frac{\partial T}{\partial x} = \kappa \frac{\partial T}{\partial x} = 0$ at $x$ boundaries and

$v\frac{\partial T}{\partial y} = \kappa \frac{\partial T}{\partial y} = 0$ at the $y$ boundaries.

To impose this, we treat $i=1$ and $i=N_{x}$ as *ghost cells*, which do not do anything expect help us impose these boundaries conditions. Discretely, the boundary fluxes between $i=1$ and $i=2$ vanish if

$\dfrac{T_{2,\,j}^{n} -T_{1,\,j}^{n}}{\Delta x} = 0$ or

$T_{1,\,j}^{n} = T_{2,\,j}^{n}.$

Thus, we can implement the boundary conditions by updating the temperature of the ghost cells to match their interior-point neighbors:
"""

# ╔═╡ 7caca2fa-2a9a-11eb-373f-156a459a1637
function update_ghostcells!(A::Array{Float64,2}; option="no-flux")
	Atmp = @view A[:,:]
	if option=="no-flux"
		A[1, :] = Atmp[2, :]; Atmp[end, :] = Atmp[end-1, :]
		A[:, 1] = Atmp[:, 2]; Atmp[:, end] = Atmp[:, end-1]
	end
end

# ╔═╡ 1f06bc34-2a9b-11eb-1030-ff12d103176c
md"Let's get a feel for what this is actually doing"

# ╔═╡ 2ef2f0cc-2a9b-11eb-03c6-5d8b4c6ae822
begin
	A = rand(Float64, 6,6);
	heatmap(A, size=(200,200))
end |> as_svg

# ╔═╡ 4558d4f8-2a9b-11eb-1f56-416975bcd180
begin
	Acopy = copy(A);
	update_ghostcells!(Acopy)
	heatmap(Acopy, size=(200,200))
end |> as_svg

# ╔═╡ 74aa7512-2a9c-11eb-118c-c7a5b60eac1b
md"""
##### 2.4) Timestepping
"""

# ╔═╡ 13eb3966-2a9a-11eb-086c-05510a3f5b80
md"""
##### 2.5) Data structures
"""

# ╔═╡ cd2ee4ca-2a06-11eb-0e61-e9a2ecf72bd6
struct Grid
	N::Int64
	L::Float64

	Δx::Float64
	Δy::Float64

	x::Array{Float64, 2}
	y::Array{Float64, 2}

	Nx::Int64
	Ny::Int64

	function Grid(N, L)
		Δx = L/N # [m]
		Δy = L/N # [m]

		x = 0. -Δx/2.:Δx:L +Δx/2.
		x = reshape(x, (1, size(x,1)))
		y = -L -Δy/2.:Δy:L +Δy/2.
		y = reshape(y, (size(y,1), 1))

		Nx, Ny = size(x, 2), size(y, 1)

		return new(N, L, Δx, Δy, x, y, Nx, Ny)
	end
end

# ╔═╡ 2a93145e-2a09-11eb-323b-01817062aa89
struct Parameters
	κ::Float64
end

# ╔═╡ 32663184-2a81-11eb-0dd1-dd1e10ed9ec6
abstract type ClimateModel end

# ╔═╡ d3796644-2a05-11eb-11b8-87b6e8c311f9
begin
	struct OceanModel <: ClimateModel
		G::Grid
		P::Parameters

		u::Array{Float64, 2}
		v::Array{Float64, 2}
	end

	OceanModel(G, P) = OceanModel(G, P, zeros(G), zeros(G))
	OceanModel(G) = OceanModel(G, Parameters(1.e4), zeros(G), zeros(G))
end;

# ╔═╡ f92086c4-2a74-11eb-3c72-a1096667183b
begin
	mutable struct ClimateModelSimulation{ModelType<:ClimateModel}
		model::ModelType

		T::Array{Float64, 2}
		Δt::Float64

		iteration::Int64
	end

	ClimateModelSimulation(C::ModelType, T, Δt) where ModelType =
		ClimateModelSimulation{ModelType}(C, T, Δt, 0)
end

# ╔═╡ 31cb0c2c-2a9a-11eb-10ba-d90a00d8e03a
md"""
##### 3) Simulating heat transport by advective & diffusive ocean currents
"""

# ╔═╡ 981ef38a-2a8b-11eb-08be-b94be2924366
md"**Simulation controls**"

# ╔═╡ d042d25a-2a62-11eb-33fe-65494bb2fad5
begin
	quiverBox = @bind show_quiver CheckBox(default=false)
	anomalyBox = @bind show_anomaly CheckBox(default=false)
	md"""
	*Click to show the velocity field* $(quiverBox) *or to show temperature **anomalies** instead of absolute values* $(anomalyBox)
	"""
end

# ╔═╡ c20b0e00-2a8a-11eb-045d-9db88411746f
begin
	U_ex_Slider = @bind U_ex Slider(-4:1:8, default=0, show_value=false);
	md"""
	$(U_ex_Slider)
	"""
end

# ╔═╡ 6dbc3d34-2a89-11eb-2c80-75459a8e237a
begin
	md"*Vary the current speed U:*  $(2. ^U_ex) [× reference]"
end

# ╔═╡ 933d42fa-2a67-11eb-07de-61cab7567d7d
begin
	κ_ex_Slider = @bind κ_ex Slider(0.:1.e3:1.e5, default=1.e4, show_value=true)
	md"""
	*Vary the diffusivity κ:* $(κ_ex_Slider) [m²/s]
	"""
end

# ╔═╡ c9ea0f72-2a67-11eb-20ba-376ca9c8014f
@bind go_ex Clock(0.1)

# ╔═╡ c3f086f4-2a9a-11eb-0978-27532cbecebf
md"""
**Some unit tests for verification**
"""

# ╔═╡ ad7b7ed6-2a9c-11eb-06b7-0f5595167575
function CFL_adv(sim::ClimateModelSimulation)
	maximum(sqrt.(sim.model.u.^2 + sim.model.v.^2)) * sim.Δt / sim.model.G.Δx
end

# ╔═╡ a04d3dee-2a9c-11eb-040e-7bd2facb2eaa
md"""
# Appendix
"""

# ╔═╡ 16905a6a-2a78-11eb-19ea-81adddc21088
Nvec = 1:25

# ╔═╡ c0e46442-27fb-11eb-2c94-15edbda3f84d
function plot_state(sim::ClimateModelSimulation; clims=(-1.1,1.1),
		show_quiver=true, show_anomaly=false, IC=nothing)

	model = sim.model
	grid = sim.model.G

	p = plot(;
		xlabel="longitudinal distance [km]", ylabel="latitudinal distance [km]",
		clabel="Temperature",
		yticks=( (-grid.L:1000e3:grid.L), Int64.(1e-3*(-grid.L:1000e3:grid.L)) ),
		xticks=( (0:1000e3:grid.L), Int64.(1e-3*(0:1000e3:grid.L)) ),
		xlims=(0., grid.L), ylims=(-grid.L, grid.L),
		)

	X = repeat(grid.x, grid.Ny, 1)
	Y = repeat(grid.y, 1, grid.Nx)
	if show_anomaly
		arrow_col = :black
		maxdiff = maximum(abs.(sim.T .- IC))
		heatmap!(p, grid.x[:], grid.y[:], sim.T .- IC, clims=(-1.1, 1.1),
			color=:balance, colorbar_title="Temperature anomaly [°C]", linewidth=0.,
			size=(400,530)
		)
	else
		arrow_col = :white
		heatmap!(p, grid.x[:], grid.y[:], sim.T,
			color=:thermal, levels=clims[1]:(clims[2]-clims[1])/21.:clims[2],
			colorbar_title="Temperature [°C]", clims=clims,
			linewidth=0., size=(400,520)
		)
	end

	annotate!(p,
		50e3, 6170e3,
		text(
			string("t = ", Int64(round(sim.iteration*sim.Δt/(60*60*24))), " days"),
			color=:black, :left, 9
		)
	)

	if show_quiver
		Nq = grid.N ÷ 5
		quiver!(p,
			X[(Nq+1)÷2:Nq:end], Y[(Nq+1)÷2:Nq:end],
			quiver=grid.L*4 .*(model.u[(Nq+1)÷2:Nq:end], model.v[(Nq+1)÷2:Nq:end]),
			color=arrow_col, alpha=0.7
		)
	end

	as_png(p)
end

# ╔═╡ c0298712-2a88-11eb-09af-bf2c39167aa6
md"""##### Computing the velocity field for a single circular vortex
"""

# ╔═╡ e3ee80c0-12dd-11eb-110a-c336bb978c51
begin
	∂x(ϕ, Δx) = (ϕ[:,2:end] - ϕ[:,1:end-1])/Δx
	∂y(ϕ, Δy) = (ϕ[2:end,:] - ϕ[1:end-1,:])/Δy

	xpad(ϕ) = hcat(zeros(size(ϕ,1)), ϕ, zeros(size(ϕ,1)))
	ypad(ϕ) = vcat(zeros(size(ϕ,2))', ϕ, zeros(size(ϕ,2))')

	xitp(ϕ) = 0.5*(ϕ[:,2:end]+ϕ[:,1:end-1])
	yitp(ϕ) = 0.5*(ϕ[2:end,:]+ϕ[1:end-1,:])

	function diagnose_velocities(ψ, G)
		u = xitp(∂y(ψ, G.Δy/G.L))
		v = yitp(-∂x(ψ, G.Δx/G.L))
		return u,v
	end
end

# ╔═╡ df706ebc-2a63-11eb-0b09-fd9f151cb5a8
function impose_no_flux!(u, v)
	u[1,:] .= 0.; v[1,:] .= 0.;
	u[end,:] .= 0.; v[end,:] .= 0.;
	u[:,1] .= 0.; v[:,1] .= 0.;
	u[:,end] .= 0.; v[:,end] .= 0.;
end

# ╔═╡ e2e4cfac-2a63-11eb-1b7f-9d8d5d304b43
function PointVortex(G; Ω=1., a=0.2, x0=0.5, y0=0.)
	x = reshape(0. -G.Δx/(G.L):G.Δx/G.L:1. +G.Δx/(G.L), (1, G.Nx+1))
	y = reshape(-1. -G.Δy/(G.L):G.Δy/G.L:1. +G.Δy/(G.L), (G.Ny+1, 1))

	function ψ̂(x,y)
		r = sqrt.((y .-y0).^2 .+ (x .-x0).^2)

		stream = -Ω/4*r.^2
		stream[r .> a] = -Ω*a^2/4*(1. .+ 2*log.(r[r .> a]/a))

		return stream
	end

	u, v = diagnose_velocities(ψ̂(x, y), G)
	impose_no_flux!(u, v)

	return u,v
end

# ╔═╡ bb084ace-12e2-11eb-2dfc-111e90eabfdd
md"""##### Computing a quasi-realistic ocean velocity field $\vec{u} = (u, v)$
Our velocity field is given by an analytical solution to the classic wind-driven gyre
problem, which is given by solving the fourth-order partial differential equation:

$- \epsilon_{M} \hat{\nabla}^{4} \hat{\Psi} + \frac{\partial \hat{\Psi} }{ \partial \hat{x}} = \nabla \times \hat{\tau} \mathbf{z},$

where the hats denote that all of the variables have been non-dimensionalized and all of their constant coefficients have been bundles into the single parameter $\epsilon_{M} \equiv \dfrac{\nu}{\beta L^3}$.

The solution makes use of an advanced *asymptotic method* (valid in the limit that $\epsilon \ll 1$) known as *boundary layer analysis* (see MIT course 18.305 to learn more).
"""

# ╔═╡ ecaab27e-2a16-11eb-0e99-87c91e659cf3
function DoubleGyre(G; β=2e-11, τ₀=0.1, ρ₀=1.e3, ν=1.e5, κ=1.e5, H=1000.)
	ϵM = ν/(β*G.L^3)
	ϵ = ϵM^(1/3.)
	x = reshape(0. -G.Δx/(G.L):G.Δx/G.L:1. +G.Δx/(G.L), (1, G.Nx+1))
	y = reshape(-1. -G.Δy/(G.L):G.Δy/G.L:1. +G.Δy/(G.L), (G.Ny+1, 1))

	ψ̂(x,y) = π*sin.(π*y) * (
		1 .- x - exp.(-x/(2*ϵ)) .* (
			cos.(√3*x/(2*ϵ)) .+
			(1. /√3)*sin.(√3*x/(2*ϵ))
			)
		.+ ϵ*exp.((x .- 1.)/ϵ)
	)

	u, v = (τ₀/ρ₀)/(β*G.L*H) .* diagnose_velocities(ψ̂(x, y), G)
	impose_no_flux!(u, v)

	return u, v
end

# ╔═╡ e59d869c-2a88-11eb-2511-5d5b4b380b80
md"""
##### Some simple initial temperature fields
"""

# ╔═╡ c4424838-12e2-11eb-25eb-058344b39c8b
linearT(G) = 0.5*(1. .+[ -(y/G.L) for y in G.y[:, 1], x in G.x[1, :] ])

# ╔═╡ 3d12c114-2a0a-11eb-131e-d1a39b4f440b
function InitBox(G; value=1., nx=2, ny=2, xspan=false, yspan=false)
	T = zeros(G)
	T[G.Ny÷2-ny:G.Ny÷2+ny, G.Nx÷2-nx:G.Nx÷2+nx] .= value
	if xspan
		T[G.Ny÷2-ny:G.Ny÷2+ny, :] .= value
	end
	if yspan
		T[:, G.Nx÷2-nx:G.Nx÷2+nx] .= value
	end
	return T
end

# ╔═╡ 863a6330-2a08-11eb-3992-c3db439fb624
begin
	G = Grid(10, 6.e6);
	P = Parameters(κ_ex);

	#u, v = zeros(G), zeros(G)
	u, v = PointVortex(G, Ω=0.5)
	# u, v = DoubleGyre(G)

	# IC = InitBox(G)
	IC = InitBox(G, xspan=true)
	# IC = linearT(G)

	model = OceanModel(G, P, u*2. ^U_ex, v*2. ^U_ex)
	Δt = 12*60*60

	ocean_sim = ClimateModelSimulation(model, copy(IC), Δt)
end;

# ╔═╡ dc9d12d0-2a9a-11eb-3dae-85b3b6029658
begin
	heat_capacity = 51.
	total_heat_content = sum(heat_capacity*ocean_sim.T*(ocean_sim.model.G.Δx*ocean_sim.model.G.Δy))*1e-15
	mean_temp = mean(ocean_sim.T)
end;

# ╔═╡ bff89550-2a9a-11eb-3038-d70249c96219
begin
	#go_ex
	md"""
	Let's make sure our model conserves energy. We have not added any energy to the system: advection and diffusion just move the energy around. The total heat content is $(round(total_heat_content, digits=3)) peta-Joules and the average temperature is $(round(mean_temp, digits=2)) °C.
	"""
end

# ╔═╡ d9e23a5a-2a8b-11eb-23f1-73ff28be9f12
md"**The CFL condition**

The CFL condition is defined by $\text{CFL} = \dfrac{\max\left(\sqrt{u² + v²}\right)Δt}{Δx} =$ $(round(CFL_adv(ocean_sim), digits=2))
"

# ╔═╡ 2908988e-2a9a-11eb-2cf7-494972f93152
Base.zeros(G::Grid) = zeros(G.Ny, G.Nx)

# ╔═╡ 6b3b6030-2066-11eb-3343-e19284638efb
plot_kernel(A) = heatmap(
	collect(A),
	color=:bluesreds, clims=(-maximum(abs.(A)), maximum(abs.(A))), colorbar=false,
	xticks=false, yticks=false, size=(30+30*size(A, 2), 30+30*size(A, 1)), xaxis=false, yaxis=false
)

# ╔═╡ 1e8d37ee-2a97-11eb-1d45-6b426b25d4eb
begin
	xgrad_kernel = OffsetArray(reshape([-1., 0, 1.], 1, 3),  0:0, -1:1)
	plot_kernel(xgrad_kernel)
end

# ╔═╡ 682f2530-2a97-11eb-3ee6-99a7c79b3767
begin
	ygrad_kernel = OffsetArray(reshape([-1., 0, 1.], 3, 1),  -1:1, 0:0)
	plot_kernel(ygrad_kernel)
end

# ╔═╡ f4c884fc-2a97-11eb-1ba9-01bf579f8b43
begin
	function advect(T, u, v, Δy, Δx, j, i)
		return .-(
			u[j, i].*sum(xgrad_kernel[0, -1:1].*T[j, i-1:i+1])/(2Δx) .+
			v[j, i].*sum(ygrad_kernel[-1:1, 0].*T[j-1:j+1, i])/(2Δy)
		)
	end
	advect(T, u, v, Δy, Δx) = [
		advect(T, u, v, Δy, Δx, j, i)
		for j=2:size(T, 1)-1, i=2:size(T, 2)-1
	]

	advect(T, O::OceanModel) = advect(T, O.u, O.v, O.G.Δy, O.G.Δx)
end

# ╔═╡ b629d89a-2a95-11eb-2f27-3dfa45789be4
begin
	xdiff_kernel = OffsetArray(reshape([1., -2., 1.], 1, 3),  0:0, -1:1)
	ydiff_kernel = OffsetArray(reshape([1., -2., 1.], 3, 1),  -1:1, 0:0)

	[plot_kernel(xdiff_kernel), plot_kernel(ydiff_kernel)]
end

# ╔═╡ ee6716c8-2a95-11eb-3a00-319ee69dd37f
begin
	function diffuse(T, κ, Δy, Δx, j, i)
		return κ.*(
			sum(xdiff_kernel[0, -1:1].*T[j, i-1:i+1])/(Δx^2) +
			sum(ydiff_kernel[-1:1, 0].*T[j-1:j+1, i])/(Δy^2)
		)
	end
	diffuse(T, κ, Δy, Δx) = [
		diffuse(T, κ, Δy, Δx, j, i) for j=2:size(T, 1)-1, i=2:size(T, 2)-1
	]

	diffuse(T, O::OceanModel) = diffuse(T, O.P.κ, O.G.Δy, O.G.Δx)
end

# ╔═╡ 81bb6a4a-2a9c-11eb-38bb-f7701c79afa2
function timestep!(sim::ClimateModelSimulation{OceanModel})
	update_ghostcells!(sim.T)
	tendencies = advect(sim.T, sim.model) .+ diffuse(sim.T, sim.model)
	sim.T[2:end-1, 2:end-1] .+= sim.Δt*tendencies
	sim.iteration += 1
end;

# ╔═╡ c98f4680-2b45-11eb-2fbe-ada2cbc080ca
for i in 1:50
	timestep!(ocean_sim)
end

# ╔═╡ 3b24e1b0-2b46-11eb-383b-c57cbf3e68f1
let
	go_ex
	if ocean_sim.iteration == 0
		timestep!(ocean_sim)
	else
		for i in 1:50
			timestep!(ocean_sim)
		end
	end
	plot_state(ocean_sim, clims=(-0.1, 1), show_quiver=show_quiver, show_anomaly=show_anomaly, IC=IC)
end

# ╔═╡ 8346b590-2b41-11eb-0bc1-1ba79bb77dfb
tvec = map(Nvec) do Npower
	G = Grid(8*Npower, 6.e6);
	P = Parameters(κ_ex);

	#u, v = DoubleGyre(G)
	#u, v = PointVortex(G, Ω=0.5)
	u, v = zeros(G), zeros(G)

	model = OceanModel(G, P, u, v)

	IC = InitBox(G)
	#IC = InitBox(G, nx=G.Nx÷2-1)
	#IC = linearT(G)

	Δt = 6*60*60
	S = ClimateModelSimulation(model, copy(IC), Δt)

	return @elapsed timestep!(S)
end

# ╔═╡ 794c2148-2a78-11eb-2756-5bd28b7726fa
begin
	plot(8*Nvec, tvec, xlabel="Number of Grid Cells (in x-direction)", ylabel="elapsed time per timestep [s]")
end |> as_svg

# ╔═╡ Cell order:
# ╟─0f8db6f4-2113-11eb-18b4-21a469c67f3a
# ╠═ed741ec6-1f75-11eb-03be-ad6284abaab8
# ╟─ac759b96-2114-11eb-24cb-d50b556f4142
# ╟─3a4a1aea-2118-11eb-30a9-57b87f2ddfae
# ╟─023779a0-2a95-11eb-35b5-7be93c43afaf
# ╟─b1b5625e-211a-11eb-3ee1-3ba9c9cc375a
# ╠═1e8d37ee-2a97-11eb-1d45-6b426b25d4eb
# ╟─3578b158-2a97-11eb-0771-bf6d82d3b6d1
# ╟─682f2530-2a97-11eb-3ee6-99a7c79b3767
# ╟─7f3c9550-2a97-11eb-1549-455009025872
# ╠═f4c884fc-2a97-11eb-1ba9-01bf579f8b43
# ╟─0127bca6-2a99-11eb-16a0-8d7af66694f8
# ╟─b629d89a-2a95-11eb-2f27-3dfa45789be4
# ╟─eac507ce-2a99-11eb-3eba-0780a4a7e078
# ╠═ee6716c8-2a95-11eb-3a00-319ee69dd37f
# ╟─09f179c0-2a9a-11eb-1d0f-e59012f9e77b
# ╠═7caca2fa-2a9a-11eb-373f-156a459a1637
# ╟─1f06bc34-2a9b-11eb-1030-ff12d103176c
# ╠═2ef2f0cc-2a9b-11eb-03c6-5d8b4c6ae822
# ╠═4558d4f8-2a9b-11eb-1f56-416975bcd180
# ╟─74aa7512-2a9c-11eb-118c-c7a5b60eac1b
# ╠═81bb6a4a-2a9c-11eb-38bb-f7701c79afa2
# ╟─13eb3966-2a9a-11eb-086c-05510a3f5b80
# ╠═cd2ee4ca-2a06-11eb-0e61-e9a2ecf72bd6
# ╠═2a93145e-2a09-11eb-323b-01817062aa89
# ╠═32663184-2a81-11eb-0dd1-dd1e10ed9ec6
# ╠═d3796644-2a05-11eb-11b8-87b6e8c311f9
# ╠═f92086c4-2a74-11eb-3c72-a1096667183b
# ╟─31cb0c2c-2a9a-11eb-10ba-d90a00d8e03a
# ╠═863a6330-2a08-11eb-3992-c3db439fb624
# ╠═c98f4680-2b45-11eb-2fbe-ada2cbc080ca
# ╟─981ef38a-2a8b-11eb-08be-b94be2924366
# ╟─d042d25a-2a62-11eb-33fe-65494bb2fad5
# ╟─6dbc3d34-2a89-11eb-2c80-75459a8e237a
# ╟─c20b0e00-2a8a-11eb-045d-9db88411746f
# ╟─933d42fa-2a67-11eb-07de-61cab7567d7d
# ╟─c9ea0f72-2a67-11eb-20ba-376ca9c8014f
# ╟─3b24e1b0-2b46-11eb-383b-c57cbf3e68f1
# ╟─c3f086f4-2a9a-11eb-0978-27532cbecebf
# ╟─bff89550-2a9a-11eb-3038-d70249c96219
# ╟─dc9d12d0-2a9a-11eb-3dae-85b3b6029658
# ╟─d9e23a5a-2a8b-11eb-23f1-73ff28be9f12
# ╠═ad7b7ed6-2a9c-11eb-06b7-0f5595167575
# ╟─a04d3dee-2a9c-11eb-040e-7bd2facb2eaa
# ╠═16905a6a-2a78-11eb-19ea-81adddc21088
# ╠═8346b590-2b41-11eb-0bc1-1ba79bb77dfb
# ╠═794c2148-2a78-11eb-2756-5bd28b7726fa
# ╠═c0e46442-27fb-11eb-2c94-15edbda3f84d
# ╟─c0298712-2a88-11eb-09af-bf2c39167aa6
# ╟─e2e4cfac-2a63-11eb-1b7f-9d8d5d304b43
# ╟─e3ee80c0-12dd-11eb-110a-c336bb978c51
# ╟─df706ebc-2a63-11eb-0b09-fd9f151cb5a8
# ╟─bb084ace-12e2-11eb-2dfc-111e90eabfdd
# ╟─ecaab27e-2a16-11eb-0e99-87c91e659cf3
# ╟─e59d869c-2a88-11eb-2511-5d5b4b380b80
# ╟─c4424838-12e2-11eb-25eb-058344b39c8b
# ╟─3d12c114-2a0a-11eb-131e-d1a39b4f440b
# ╟─2908988e-2a9a-11eb-2cf7-494972f93152
# ╠═9c8a7e5a-12dd-11eb-1b99-cd1d52aefa1d
# ╟─6b3b6030-2066-11eb-3343-e19284638efb

max_stars_repo_name: mathijsvdv/ComputationalThinking

max_issues_repo_name: mathijsvdv/ComputationalThinking

max_forks_repo_name: mathijsvdv/ComputationalThinking


text:
using Random
include("ED.jl")

################################################################################
### Verify ED
################################################################################

@testset "ED checks" begin
    void = State(0)
    up = State(1)
    down = State(2)
    updown = State(3)

    # create up at site 1
    _sign, s = create(void, 1, 1)
    @test _sign == 1.0 &&  s == up
    _sign, s = create(up, 1, 1)
    @test _sign == 0
    _sign, s = create(down, 1, 1)
    @test _sign == 1.0 &&  s == updown
    _sign, s = create(updown, 1, 1)
    @test _sign == 0

    # create down at site 1
    _sign, s = create(void, 1, 2)
    @test _sign == 1.0 &&  s == down
    _sign, s = create(up, 1, 2)
    @test _sign == -1.0 &&  s == updown
    _sign, s = create(down, 1, 2)
    @test _sign == 0
    _sign, s = create(updown, 1, 2)
    @test _sign == 0

    # annihilate up at site 1
    _sign, s = annihilate(void, 1, 1)
    @test _sign == 0
    _sign, s = annihilate(up, 1, 1)
    @test _sign == 1.0 &&  s == void
    _sign, s = annihilate(down, 1, 1)
    @test _sign == 0
    _sign, s = annihilate(updown, 1, 1)
    @test _sign == 1.0 &&  s == down

    # annihilate down at site 1
    _sign, s = annihilate(void, 1, 2)
    @test _sign == 0
    _sign, s = annihilate(up, 1, 2)
    @test _sign == 0
    _sign, s = annihilate(down, 1, 2)
    @test _sign == 1.0 &&  s == void
    _sign, s = annihilate(updown, 1, 2)
    @test _sign == -1.0 &&  s == up

    # check Greens consistency
    model = HubbardModel(2, 2, U = rand(), mu = rand(), t = rand()
    )
    H = HamiltonMatrix(model)
    for substate1 in 1:2, substate2 in 1:2
        for site1 in 1:model.l.sites, site2 in 1:model.l.sites
            G = expectation_value(
                Greens(site1, site2, substate1, substate2),
                H,
                N_sites = model.l.sites,
            )
            G_perm = expectation_value(
                Greens_permuted(site1, site2, substate1, substate2),
                H,
                N_sites = model.l.sites,
            )
            @test G ≈ G_perm

            # Test G(t, t) = G(0, 0) = G
            UTG = expectation_value(
                s -> annihilate(s, site2, substate2),
                s -> create(s, site1, substate1),
                H, 0.1, 0.1, N_sites=model.l.sites
            )
            @test check(UTG, real(G), 1e-14, 0.0)
            UTG = expectation_value(
                s -> annihilate(s, site2, substate2),
                s -> create(s, site1, substate1),
                H, 0.7, 0.7, N_sites=model.l.sites
            )
            @test check(UTG, real(G), 1e-14, 0.0)
        end
    end

    # Test G(t, 0) + G(0, beta - t) = 0
    UTG1 = calculate_Greens_matrix(H, 0.7, 0.0, model.l, beta=1.0)
    UTG2 = calculate_Greens_matrix(H, 0.0, 1.0 - 0.7, model.l, beta=1.0)
    @test UTG1 ≈ -UTG2 atol=1e-13
end

################################################################################
### Checks with U = 0 (should be numerically exact)
################################################################################

@testset "Exact Greens comparison (ED, tiny systems)" begin
    # These are theoretically the same but their implementation differs on
    # some level. To make sure both are correct it makes sense to check both here.
    models = (
        HubbardModel(2, 2, U = 0.0, t = 1.0),
    )

    println("Exact Greens comparison (ED)")
    for model in models, beta in (1.0, 8.9)
        @testset "$(typeof(model))" begin
            dqmc = DQMC(
                model, beta=beta, delta_tau = 0.1, safe_mult=5, recorder = Discarder(),
                thermalization = 1, sweeps = 2, measure_rate=1#, field = field
            )
            print(
                "  Running DQMC ($(nameof(typeof(model))) $(nameof(field))) " *
                "β=$(dqmc.parameters.beta)\n    "
            )

            dqmc[:G]    = greens_measurement(dqmc, model)
            dqmc[:E]    = total_energy(dqmc, model)
            dqmc[:Occs] = occupation(dqmc, model)
            dqmc[:CDC]  = charge_density_correlation(dqmc, model)
            dqmc[:Mx]   = magnetization(dqmc, model, :x)
            dqmc[:My]   = magnetization(dqmc, model, :y)
            dqmc[:Mz]   = magnetization(dqmc, model, :z)
            dqmc[:SDCx] = spin_density_correlation(dqmc, model, :x)
            dqmc[:SDCy] = spin_density_correlation(dqmc, model, :y)
            dqmc[:SDCz] = spin_density_correlation(dqmc, model, :z)
            dqmc[:PC]   = pairing_correlation(
                dqmc, model, lattice_iterator = EachLocalQuadByDistance(1:4)
            )

            # Unequal time
            l1s = [0, 3, 5, 7, 3, 0, MonteCarlo.nslices(dqmc), 0]
            l2s = [1, 7, 5, 2, 2, MonteCarlo.nslices(dqmc), 0, 0]
            for i in eachindex(l1s)
                dqmc[Symbol(:UTG, i)] = greens_measurement(dqmc, model, GreensAt(l2s[i], l1s[i]))
            end

            dqmc[:CDS]  = charge_density_susceptibility(dqmc, model)
            dqmc[:SDSx] = spin_density_susceptibility(dqmc, model, :x)
            dqmc[:SDSy] = spin_density_susceptibility(dqmc, model, :y)
            dqmc[:SDSz] = spin_density_susceptibility(dqmc, model, :z)
            dqmc[:PS]   = pairing_susceptibility(
                dqmc, model, lattice_iterator = EachLocalQuadByDistance(1:4)
            )
            dqmc[:CCS]  = current_current_susceptibility(
                dqmc, model, lattice_iterator = EachLocalQuadByDistance(1:4)
            )

            @time run!(dqmc, verbose=false)

            # error tolerance
            atol = 1e-13
            rtol = 1e-13
            N = length(lattice(model))

            # Direct calculation simialr to what DQMC should be doing
            T = Matrix(MonteCarlo.hopping_matrix(dqmc, model))
            # Doing an eigenvalue decomposition makes this pretty stable
            vals, U = eigen(exp(-T))
            D = Diagonal(vals)^(dqmc.parameters.beta)

            # G = I - U * inv(I + D) * adjoint(U)
            G = U * inv(I + D) * adjoint(U)

            print("    Running ED and checking (tolerance: $atol, $(100rtol)%)\n    ")
            @time begin
                H = HamiltonMatrix(model)

                @testset "(total) energy" begin
                    dqmc_E = mean(dqmc[:E])
                    ED_E = energy(H, beta = dqmc.parameters.beta)
                    @test dqmc_E ≈ ED_E atol=atol rtol=rtol
                end

                # G_DQMC is smaller because it doesn't differentiate between spin up/down
                @testset "Greens" begin
                    G_DQMC = mean(dqmc.measurements[:G])
                    occs = mean(dqmc.measurements[:Occs])
                    G_ED = calculate_Greens_matrix(H, model.l, beta = dqmc.parameters.beta)
                    @test check(G_DQMC, G_ED[1:size(G_DQMC, 1), 1:size(G_DQMC, 2)], atol, rtol)
                    @test check(G, G_ED[1:size(G_DQMC, 1), 1:size(G_DQMC, 2)], atol, rtol)
                    @test check(G_DQMC, G, atol, rtol)
                    for i in 1:size(G_DQMC, 1)
                        @test check(occs[i],  1 - G_ED[i, i], atol, rtol)
                    end
                end

                @testset "Charge Density Correlation" begin
                    CDC = mean(dqmc.measurements[:CDC])
                    ED_CDC = zeros(size(CDC))
                    for (dir, src, trg) in MonteCarlo.lattice_iterator(dqmc[:CDC], dqmc, model)
                        ED_CDC[dir] += expectation_value(
                            charge_density_correlation(trg, src),
                            H, beta = dqmc.parameters.beta, N_sites = N
                        )
                    end
                    @test check(ED_CDC/N, CDC, atol, rtol)
                end

                @testset "Magnetization x" begin
                    Mx = mean(dqmc.measurements[:Mx])
                    for site in 1:length(Mx)
                        ED_Mx = expectation_value(
                            m_x(site), H, beta = dqmc.parameters.beta, N_sites = N
                        )
                        @test check(ED_Mx, Mx[site], atol, rtol)
                    end
                end
                @testset "Magnetization y" begin
                    My = mean(dqmc.measurements[:My])
                    for site in 1:length(My)
                        ED_My = expectation_value(
                            m_y(site), H, beta = dqmc.parameters.beta, N_sites = N
                        ) |> imag
                        @test check(ED_My, My[site], atol, rtol)
                    end
                end
                @testset "Magnetization z" begin
                    Mz = mean(dqmc.measurements[:Mz])
                    for site in 1:length(Mz)
                        ED_Mz = expectation_value(
                            m_z(site), H, beta = dqmc.parameters.beta, N_sites = N
                        )
                        @test check(ED_Mz, Mz[site], atol, rtol)
                    end
                end

                @testset "Spin density correlation x" begin
                    SDCx = mean(dqmc.measurements[:SDCx])
                    ED_SDCx = zeros(ComplexF64, size(SDCx))
                    for (dir, src, trg) in MonteCarlo.lattice_iterator(dqmc[:SDCx], dqmc, model)
                        ED_SDCx[dir] += expectation_value(
                            spin_density_correlation_x(trg, src),
                            H, beta = dqmc.parameters.beta, N_sites = N
                        )
                    end
                    @test check(ED_SDCx/N, SDCx, atol, rtol)
                end
                @testset "Spin density correlation y" begin
                    SDCy = mean(dqmc.measurements[:SDCy])
                    ED_SDCy = zeros(ComplexF64, size(SDCy))
                    for (dir, src, trg) in MonteCarlo.lattice_iterator(dqmc[:SDCy], dqmc, model)
                        ED_SDCy[dir] += expectation_value(
                            spin_density_correlation_y(trg, src),
                            H, beta = dqmc.parameters.beta, N_sites = N
                        )
                    end
                    @test check(ED_SDCy/N, SDCy, atol, rtol)
                end
                @testset "Spin density correlation z" begin
                    SDCz = mean(dqmc.measurements[:SDCz])
                    ED_SDCz = zeros(ComplexF64, size(SDCz))
                    for (dir, src, trg) in MonteCarlo.lattice_iterator(dqmc[:SDCz], dqmc, model)
                        ED_SDCz[dir] += expectation_value(
                            spin_density_correlation_z(trg, src),
                            H, beta = dqmc.parameters.beta, N_sites = N
                        )
                    end
                    @test check(ED_SDCz/N, SDCz, atol, rtol)
                end

                @testset "Pairing Correlation" begin
                    PC = mean(dqmc.measurements[:PC])
                    ED_PC = zeros(ComplexF64, size(PC))
                    for (dirs, src1, trg1, src2, trg2) in
                            MonteCarlo.lattice_iterator(dqmc[:PC], dqmc, model)
                        ED_PC[dirs] += expectation_value(
                            pairing_correlation(src1, trg1, src2, trg2),
                            H, beta = dqmc.parameters.beta, N_sites = N
                        )
                    end
                    @test check(ED_PC/N, PC, atol, rtol)
                end

                ################################################################
                ### Unequal Time
                ################################################################

                @testset "Unequal Time Greens" begin
                    for (i, tau1, tau2) in zip(eachindex(l1s), 0.1l1s, 0.1l2s)
                        UTG = mean(dqmc.measurements[Symbol(:UTG, i)])
                        M = size(UTG, 1)
                        ED_UTG = calculate_Greens_matrix(
                            H, tau2, tau1, model.l, beta=dqmc.parameters.beta
                        )

                        @testset "[$i] $tau1 -> $tau2" begin
                            for k in 1:M, l in 1:M
                                @test check(UTG[k, l], ED_UTG[k, l], atol, rtol)
                            end
                        end
                    end
                end

                # CDC
                @testset "Charge Density Susceptibility" begin
                    CDS = mean(dqmc.measurements[:CDS])
                    ED_CDS = zeros(size(CDS))
                    for (dir, src, trg) in MonteCarlo.lattice_iterator(dqmc[:CDS], dqmc, model)
                        ED_CDS[dir] += expectation_value_integrated(
                            number_operator(trg), number_operator(src), H,
                            step = dqmc.parameters.delta_tau, beta = dqmc.parameters.beta, N_sites = N
                        )
                    end
                    @test check(ED_CDS/N, CDS, atol, rtol)
                end

                # SDS
                @testset "Spin density Susceptibility x" begin
                    SDSx = mean(dqmc.measurements[:SDSx])
                    ED_SDSx = zeros(ComplexF64, size(SDSx))
                    for (dir, src, trg) in MonteCarlo.lattice_iterator(dqmc[:SDSx], dqmc, model)
                        ED_SDSx[dir] += expectation_value_integrated(
                            m_x(trg), m_x(src), H, step = dqmc.parameters.delta_tau,
                            beta = dqmc.parameters.beta, N_sites = N
                        )
                    end
                    @test check(ED_SDSx/N, SDSx, atol, rtol)
                end
                @testset "Spin density Susceptibility y" begin
                    SDSy = mean(dqmc.measurements[:SDSy])
                    ED_SDSy = zeros(ComplexF64, size(SDSy))
                    for (dir, src, trg) in MonteCarlo.lattice_iterator(dqmc[:SDSy], dqmc, model)
                        ED_SDSy[dir] += expectation_value_integrated(
                            m_y(trg), m_y(src), H, step = dqmc.parameters.delta_tau,
                            beta = dqmc.parameters.beta, N_sites = N
                        )
                    end
                    @test check(ED_SDSy/N, SDSy, atol, rtol)
                end
                @testset "Spin density Susceptibility z" begin
                    SDSz = mean(dqmc.measurements[:SDSz])
                    ED_SDSz = zeros(ComplexF64, size(SDSz))
                    for (dir, src, trg) in MonteCarlo.lattice_iterator(dqmc[:SDSz], dqmc, model)
                        ED_SDSz[dir] += expectation_value_integrated(
                            m_z(trg), m_z(src), H, step = dqmc.parameters.delta_tau,
                            beta = dqmc.parameters.beta, N_sites = N
                        )
                    end
                    @test check(ED_SDSz/N, SDSz, atol, rtol)
                end

                @testset "Pairing Susceptibility" begin
                    PS = mean(dqmc.measurements[:PS])
                    ED_PS = zeros(Float64, size(PS))
                    for (dirs, src1, trg1, src2, trg2) in
                            MonteCarlo.lattice_iterator(dqmc[:PS], dqmc, model)
                        ED_PS[dirs] += expectation_value_integrated(
                            state -> begin
                                sign1, _state  = annihilate(state, trg1, DOWN)
                                sign2, _state1 = annihilate(_state, src1, UP)
                                sign3, _state  = create(state, src1, UP)
                                sign4, _state2 = create(_state, trg1, DOWN)
                                (sign1*sign2, sign3*sign4), (_state1, _state2)
                            end,
                            state -> begin
                                sign1, _state  = create(state, src2, UP)
                                sign2, _state1 = create(_state, trg2, DOWN)
                                sign3, _state  = annihilate(state, trg2, DOWN)
                                sign4, _state2 = annihilate(_state, src2, UP)
                                (sign1*sign2, sign3*sign4), (_state1, _state2)
                            end,
                            H, step = dqmc.parameters.delta_tau, beta = dqmc.parameters.beta, N_sites = N
                        )
                    end
                    @test check(ED_PS/N, PS, atol, rtol)
                end

                @testset "Current Current Susceptibility" begin
                    CCS = mean(dqmc.measurements[:CCS])
                    ED_CCS = zeros(Float64, size(CCS))
                    T = dqmc.stack.hopping_matrix
                    for (dirs, src1, trg1, src2, trg2) in
                            MonteCarlo.lattice_iterator(dqmc[:CCS], dqmc, model)
                        ED_CCS[dirs] -= expectation_value_integrated(
                            # actually the order of this doesn't seem to matter
                            current_density(src2, trg2, T),
                            current_density(src1, trg1, T),
                            H, step = dqmc.parameters.delta_tau, beta = dqmc.parameters.beta, N_sites = N
                        )
                    end
                    @test check(ED_CCS/N, CCS, atol, rtol)
                end
            end

        end
    end
end

################################################################################
### Checks with finite U ≠ 0 (Potentially large Trotter error)
################################################################################

@testset "Repulsive/Attractive Hubbard Model (ED)" begin
    models = (
        HubbardModel(2, 2, U = -1.0, t = 1.0),
        HubbardModel(2, 2, U = 1.0, mu = 1.0, t = 1.0)
    )
    # fields = (MagneticHirschField, DensityHirschField, MagneticGHQField)

    println("Finite U ED Comparison")
    for model in models
        @testset "$(typeof(model))" begin
            Random.seed!(123)
            dqmc = DQMC(
                model, beta=1.0, delta_tau = 0.1, safe_mult=5, recorder = Discarder(),
                thermalization = 5_000, sweeps = 5_000, print_rate=1000,
                scheduler = AdaptiveScheduler(
                    (LocalSweep(10), Adaptive(),), (GlobalShuffle(), GlobalFlip())
                )
            )
            print(
                "  Running DQMC ($(nameof(typeof(model))) $(nameof(field))) " *
                "β=$(dqmc.parameters.beta), 5k + 5k sweeps\n    "
            )

            dqmc[:G]    = greens_measurement(dqmc, model)
            dqmc[:E]    = total_energy(dqmc, model)
            dqmc[:Occs] = occupation(dqmc, model)
            dqmc[:CDC]  = charge_density_correlation(dqmc, model)
            dqmc[:Mx]   = magnetization(dqmc, model, :x)
            dqmc[:My]   = magnetization(dqmc, model, :y)
            dqmc[:Mz]   = magnetization(dqmc, model, :z)
            dqmc[:SDCx] = spin_density_correlation(dqmc, model, :x)
            dqmc[:SDCy] = spin_density_correlation(dqmc, model, :y)
            dqmc[:SDCz] = spin_density_correlation(dqmc, model, :z)
            dqmc[:PC]   = pairing_correlation(
                dqmc, model, lattice_iterator = EachLocalQuadByDistance(1:4)
            )

            # Unequal time
            l1s = [0, 3, 5, 7, 3, 0]
            l2s = [1, 7, 5, 2, 2, MonteCarlo.nslices(dqmc)]
            # l1s = [0, 3, 5, 0]
            # l2s = [1, 7, 5, MonteCarlo.nslices(dqmc)]
            dqmc[:UTG1] = greens_measurement(dqmc, model, GreensAt(l2s[1], l1s[1]))
            dqmc[:UTG2] = greens_measurement(dqmc, model, GreensAt(l2s[2], l1s[2]))
            dqmc[:UTG3] = greens_measurement(dqmc, model, GreensAt(l2s[3], l1s[3]))
            dqmc[:UTG4] = greens_measurement(dqmc, model, GreensAt(l2s[4], l1s[4]))
            dqmc[:UTG5] = greens_measurement(dqmc, model, GreensAt(l2s[5], l1s[5]))
            dqmc[:UTG6] = greens_measurement(dqmc, model, GreensAt(l2s[6], l1s[6]))

            dqmc[:CDS]  = charge_density_susceptibility(dqmc, model)
            dqmc[:SDSx] = spin_density_susceptibility(dqmc, model, :x)
            dqmc[:SDSy] = spin_density_susceptibility(dqmc, model, :y)
            dqmc[:SDSz] = spin_density_susceptibility(dqmc, model, :z)
            dqmc[:PS]   = pairing_susceptibility(
                dqmc, model, lattice_iterator = EachLocalQuadByDistance(1:4)
            )
            dqmc[:CCS]  = current_current_susceptibility(
                dqmc, model, lattice_iterator = EachLocalQuadByDistance(1:4)
            )

            # MonteCarlo.enable_benchmarks()

            @time run!(dqmc, verbose=!true)

            # Absolute tolerance from Trotter decompositon
            atol = 2.5dqmc.parameters.delta_tau^2
            rtol = 2dqmc.parameters.delta_tau^2
            N = length(lattice(model))

            print("    Running ED and checking (tolerance: $atol, $(100rtol)%)\n    ")
            @time begin
                H = HamiltonMatrix(model)

                @testset "(total) energy" begin
                    dqmc_E = mean(dqmc[:E])
                    ED_E = energy(H, beta = dqmc.parameters.beta)
                    @test dqmc_E ≈ ED_E atol=atol rtol=rtol
                end

                # G_DQMC is smaller because it doesn't differentiate between spin up/down
                @testset "Greens" begin
                    G_DQMC = mean(dqmc.measurements[:G])
                    occs = mean(dqmc.measurements[:Occs])
                    G_ED = calculate_Greens_matrix(H, model.l, beta = dqmc.parameters.beta)
                    for i in 1:size(G_DQMC, 1), j in 1:size(G_DQMC, 2)
                        @test check(G_DQMC[i, j], G_ED[i, j], atol, rtol)
                    end
                    for i in 1:size(G_DQMC, 1)
                        @test check(occs[i],  1 - G_ED[i, i], atol, rtol)
                    end
                end

                @testset "Charge Density Correlation" begin
                    CDC = mean(dqmc.measurements[:CDC])
                    ED_CDC = zeros(size(CDC))
                    for (dir, src, trg) in MonteCarlo.lattice_iterator(dqmc[:CDC], dqmc, model)
                        ED_CDC[dir] += expectation_value(
                            charge_density_correlation(trg, src),
                            H, beta = dqmc.parameters.beta, N_sites = N
                        )
                    end
                    @test check(ED_CDC/N, CDC, atol, rtol)
                end

                @testset "Magnetization x" begin
                    Mx = mean(dqmc.measurements[:Mx])
                    for site in 1:length(Mx)
                        ED_Mx = expectation_value(
                            m_x(site), H, beta = dqmc.parameters.beta, N_sites = N
                        )
                        @test check(ED_Mx, Mx[site], atol, rtol)
                    end
                end
                @testset "Magnetization y" begin
                    My = mean(dqmc.measurements[:My])
                    for site in 1:length(My)
                        ED_My = expectation_value(
                            m_y(site), H, beta = dqmc.parameters.beta, N_sites = N
                        ) |> imag
                        @test check(ED_My, My[site], atol, rtol)
                    end
                end
                @testset "Magnetization z" begin
                    Mz = mean(dqmc.measurements[:Mz])
                    for site in 1:length(Mz)
                        ED_Mz = expectation_value(
                            m_z(site), H, beta = dqmc.parameters.beta, N_sites = N
                        )
                        @test check(ED_Mz, Mz[site], atol, rtol)
                    end
                end

                @testset "Spin density correlation x" begin
                    SDCx = mean(dqmc.measurements[:SDCx])
                    ED_SDCx = zeros(ComplexF64, size(SDCx))
                    for (dir, src, trg) in MonteCarlo.lattice_iterator(dqmc[:SDCx], dqmc, model)
                        ED_SDCx[dir] += expectation_value(
                            spin_density_correlation_x(trg, src),
                            H, beta = dqmc.parameters.beta, N_sites = N
                        )
                    end
                    @test check(ED_SDCx/N, SDCx, atol, rtol)
                end
                @testset "Spin density correlation y" begin
                    SDCy = mean(dqmc.measurements[:SDCy])
                    ED_SDCy = zeros(ComplexF64, size(SDCy))
                    for (dir, src, trg) in MonteCarlo.lattice_iterator(dqmc[:SDCy], dqmc, model)
                        ED_SDCy[dir] += expectation_value(
                            spin_density_correlation_y(trg, src),
                            H, beta = dqmc.parameters.beta, N_sites = N
                        )
                    end
                    @test check(ED_SDCy/N, SDCy, atol, rtol)
                end
                @testset "Spin density correlation z" begin
                    SDCz = mean(dqmc.measurements[:SDCz])
                    ED_SDCz = zeros(ComplexF64, size(SDCz))
                    for (dir, src, trg) in MonteCarlo.lattice_iterator(dqmc[:SDCz], dqmc, model)
                        ED_SDCz[dir] += expectation_value(
                            spin_density_correlation_z(trg, src),
                            H, beta = dqmc.parameters.beta, N_sites = N
                        )
                    end
                    @test check(ED_SDCz/N, SDCz, atol, rtol)
                end

                @testset "Pairing Correlation" begin
                    PC = mean(dqmc.measurements[:PC])
                    ED_PC = zeros(ComplexF64, size(PC))
                    for (dirs, src1, trg1, src2, trg2) in
                            MonteCarlo.lattice_iterator(dqmc[:PC], dqmc, model)
                        ED_PC[dirs] += expectation_value(
                            pairing_correlation(src1, trg1, src2, trg2),
                            H, beta = dqmc.parameters.beta, N_sites = N
                        )
                    end
                    @test check(ED_PC/N, PC, atol, rtol)
                end

                ################################################################
                ### Unequal Time
                ################################################################

                @testset "Unequal Time Greens" begin
                    dt = dqmc.parameters.delta_tau
                    for (i, tau1, tau2) in zip(eachindex(l1s), dt * l1s, dt * l2s)
                        UTG = mean(dqmc.measurements[Symbol(:UTG, i)])
                        # ΔUTG = std_error(dqmc.measurements[Symbol(:UTG, i)])
                        M = size(UTG, 1)
                        ED_UTG = calculate_Greens_matrix(
                            H, tau2, tau1, model.l, beta=dqmc.parameters.beta
                        )

                        @testset "[$i] $tau1 -> $tau2" begin
                            for k in 1:M, l in 1:M
                                @test check(UTG[k, l], ED_UTG[k, l], atol, rtol)
                            end
                        end
                    end
                end

                # CDC
                @testset "Charge Density Susceptibility" begin
                    CDS = mean(dqmc.measurements[:CDS])
                    ED_CDS = zeros(size(CDS))
                    for (dir, src, trg) in MonteCarlo.lattice_iterator(dqmc[:CDS], dqmc, model)
                        ED_CDS[dir] += expectation_value_integrated(
                            number_operator(trg), number_operator(src), H,
                            step = dqmc.parameters.delta_tau, beta = dqmc.parameters.beta, N_sites = N
                        )
                    end
                    @test check(ED_CDS/N, CDS, atol, rtol)
                end

                # SDS
                @testset "Spin density Susceptibility x" begin
                    SDSx = mean(dqmc.measurements[:SDSx])
                    ED_SDSx = zeros(ComplexF64, size(SDSx))
                    for (dir, src, trg) in MonteCarlo.lattice_iterator(dqmc[:SDSx], dqmc, model)
                        ED_SDSx[dir] += expectation_value_integrated(
                            m_x(trg), m_x(src), H, step = dqmc.parameters.delta_tau,
                            beta = dqmc.parameters.beta, N_sites = N
                        )
                    end
                    @test check(ED_SDSx/N, SDSx, atol, rtol)
                end
                @testset "Spin density Susceptibility y" begin
                    SDSy = mean(dqmc.measurements[:SDSy])
                    ED_SDSy = zeros(ComplexF64, size(SDSy))
                    for (dir, src, trg) in MonteCarlo.lattice_iterator(dqmc[:SDSy], dqmc, model)
                        ED_SDSy[dir] += expectation_value_integrated(
                            m_y(trg), m_y(src), H, step = dqmc.parameters.delta_tau,
                            beta = dqmc.parameters.beta, N_sites = N
                        )
                    end
                    @test check(ED_SDSy/N, SDSy, atol, rtol)
                end
                @testset "Spin density Susceptibility z" begin
                    SDSz = mean(dqmc.measurements[:SDSz])
                    ED_SDSz = zeros(ComplexF64, size(SDSz))
                    for (dir, src, trg) in MonteCarlo.lattice_iterator(dqmc[:SDSz], dqmc, model)
                        ED_SDSz[dir] += expectation_value_integrated(
                            m_z(trg), m_z(src), H, step = dqmc.parameters.delta_tau,
                            beta = dqmc.parameters.beta, N_sites = N
                        )
                    end
                    @test check(ED_SDSz/N, SDSz, atol, rtol)
                end

                @testset "Pairing Susceptibility" begin
                    PS = mean(dqmc.measurements[:PS])
                    ED_PS = zeros(Float64, size(PS))
                    for (dirs, src1, trg1, src2, trg2) in
                            MonteCarlo.lattice_iterator(dqmc[:PS], dqmc, model)
                        ED_PS[dirs] += expectation_value_integrated(
                            state -> begin
                                sign1, _state  = annihilate(state, trg1, DOWN)
                                sign2, _state1 = annihilate(_state, src1, UP)
                                sign3, _state  = create(state, src1, UP)
                                sign4, _state2 = create(_state, trg1, DOWN)
                                (sign1*sign2, sign3*sign4), (_state1, _state2)
                            end,
                            state -> begin
                                sign1, _state  = create(state, src2, UP)
                                sign2, _state1 = create(_state, trg2, DOWN)
                                sign3, _state  = annihilate(state, trg2, DOWN)
                                sign4, _state2 = annihilate(_state, src2, UP)
                                (sign1*sign2, sign3*sign4), (_state1, _state2)
                            end,
                            H, step = dqmc.parameters.delta_tau, beta = dqmc.parameters.beta, N_sites = N
                        )
                    end
                    @test check(ED_PS/N, PS, atol, rtol)
                end

                @testset "Current Current Susceptibility" begin
                    CCS = mean(dqmc.measurements[:CCS])
                    ED_CCS = zeros(Float64, size(CCS))
                    T = dqmc.stack.hopping_matrix
                    for (dirs, src1, trg1, src2, trg2) in
                            MonteCarlo.lattice_iterator(dqmc[:CCS], dqmc, model)
                        ED_CCS[dirs] -= expectation_value_integrated(
                            # actually the order of this doesn't seem to matter
                            current_density(src2, trg2, T),
                            current_density(src1, trg1, T),
                            H, step = dqmc.parameters.delta_tau, beta = dqmc.parameters.beta, N_sites = N
                        )
                    end
                    @test check(ED_CCS/N, CCS, atol, rtol)
                end
            end
        end
    end
end

################################################################################
### Field checks
################################################################################

@testset "Repulsive/Attractive Hubbard Model (ED)" begin
    models = (
        HubbardModel(2, 2, U = -1.0, t = 1.0),
        HubbardModel(2, 2, U = 1.0, mu = 1.0, t = 1.0)
    )
    fields = (MagneticHirschField, DensityHirschField, MagneticGHQField)

    println("Finite U ED Comparison")
    for model in models, field in fields
        if field == MonteCarlo.choose_field(model)
            continue
        end

        @testset "$(nameof(typeof(model))) + $(nameof(field))" begin
            Random.seed!(123)
            dqmc = DQMC(
                model, beta=1.0, delta_tau = 0.1, safe_mult=5, recorder = Discarder(),
                thermalization = 5_000, sweeps = 5_000, print_rate=1000, field = field,
                scheduler = AdaptiveScheduler(
                    (LocalSweep(10), Adaptive(),), (GlobalShuffle(), GlobalFlip())
                )
            )
            str = model.U >= 0 ? "attractive" : "repulsive"
            print(
                "  Running DQMC ($mod $(nameof(field))) " *
                "β=$(dqmc.parameters.beta), 5k + 5k sweeps\n    "
            )

            dqmc[:G] = greens_measurement(dqmc, model)
            dqmc[:E] = total_energy(dqmc, model)

            @time run!(dqmc, verbose=!true)

            # Absolute tolerance from Trotter decompositon
            atol = 2.5dqmc.parameters.delta_tau^2
            rtol = 2dqmc.parameters.delta_tau^2
            N = length(lattice(model))

            print("    Running ED and checking (tolerance: $atol, $(100rtol)%)\n    ")
            @time begin
                H = HamiltonMatrix(model)

                @testset "(total) energy" begin
                    dqmc_E = mean(dqmc[:E])
                    ED_E = energy(H, beta = dqmc.parameters.beta)
                    @test dqmc_E ≈ ED_E atol=atol rtol=rtol
                end

                # G_DQMC is smaller because it doesn't differentiate between spin up/down
                @testset "Greens" begin
                    G_DQMC = mean(dqmc.measurements[:G])
                    G_ED = calculate_Greens_matrix(H, model.l, beta = dqmc.parameters.beta)
                    for i in 1:size(G_DQMC, 1), j in 1:size(G_DQMC, 2)
                        @test check(G_DQMC[i, j], G_ED[i, j], atol, rtol)
                    end
                end

            end
        end
    end
end


text:
# ============================================================
# Structs of Solution
# Solver stores data in structs of arrays (SoA)
# ============================================================

mutable struct Solution1D{A} <: AbstractSolution1D

    w::A
    prim::A
    sw::A

    function Solution1D(
        w::AbstractArray,
        prim::AbstractArray,
        sw = [zeros(axes(w[1])) for i in axes(w, 1)]::AbstractArray,
    )
        new{typeof(w)}(w, prim, sw)
    end

end

mutable struct Solution1D1F{A,B} <: AbstractSolution1D

    w::A
    prim::A
    sw::A
    f::B
    sf::B

    function Solution1D1F(w::AbstractArray, prim::AbstractArray, f::AbstractArray)
        sw = [zeros(axes(w[1])) for i in axes(w, 1)]
        sf = [zeros(axes(f[1])) for i in axes(f, 1)]

        new{typeof(w),typeof(f)}(w, prim, sw, f, sf)
    end

    function Solution1D1F(
        w::AbstractArray,
        prim::AbstractArray,
        sw::AbstractArray,
        f::AbstractArray,
        sf::AbstractArray,
    )
        new{typeof(w),typeof(f)}(w, prim, sw, f, sf)
    end

end

mutable struct Solution1D2F{A,B} <: AbstractSolution1D

    w::A
    prim::A
    sw::A
    h::B
    b::B
    sh::B
    sb::B

    function Solution1D2F(
        w::AbstractArray,
        prim::AbstractArray,
        h::AbstractArray,
        b::AbstractArray,
    )
        sw = [zeros(axes(w[1])) for i in axes(w, 1)]
        sh = [zeros(axes(h[1])) for i in axes(h, 1)]
        sb = [zeros(axes(b[1])) for i in axes(b, 1)]

        new{typeof(w),typeof(h)}(w, prim, sw, h, b, sh, sb)
    end

    function Solution1D2F(
        w::AbstractArray,
        prim::AbstractArray,
        sw::AbstractArray,
        h::AbstractArray,
        b::AbstractArray,
        sh::AbstractArray,
        sb::AbstractArray,
    )
        new{typeof(w),typeof(h)}(w, prim, sw, h, b, sh, sb)
    end

end

mutable struct Solution2D{A,B} <: AbstractSolution2D

    w::A
    prim::A
    sw::B

    function Solution2D(
        w::AbstractArray,
        prim::AbstractArray,
        sw = [
            zeros((axes(w[1])..., Base.OneTo(2))) for i in axes(w, 1), j in axes(w, 2)
        ]::AbstractArray,
    )

        new{typeof(w),typeof(sw)}(w, prim, sw)
    end

end

mutable struct Solution2D1F{A,B,C,D} <: AbstractSolution2D

    w::A
    prim::A
    sw::B
    f::C
    sf::D

    function Solution2D1F(w::AbstractArray, prim::AbstractArray, f::AbstractArray)
        sw = [zeros((axes(w[1])..., Base.OneTo(2))) for i in axes(w, 1), j in axes(w, 2)]
        sf = [zeros((axes(f[1])..., Base.OneTo(2))) for i in axes(f, 1), j in axes(f, 2)]

        new{typeof(w),typeof(sw),typeof(f),typeof(sf)}(w, prim, sw, f, sf)
    end

    function Solution2D1F(
        w::AbstractArray,
        prim::AbstractArray,
        sw::AbstractArray,
        f::AbstractArray,
        sf::AbstractArray,
    )
        new{typeof(w),typeof(sw),typeof(f),typeof(sf)}(w, prim, sw, f, sf)
    end

end

mutable struct Solution2D2F{A,B,C,D} <: AbstractSolution2D

    w::A
    prim::A
    sw::B
    h::C
    b::C
    sh::D
    sb::D

    function Solution2D2F(
        w::AbstractArray,
        prim::AbstractArray,
        h::AbstractArray,
        b::AbstractArray,
    )
        sw = [zeros((axes(w[1])..., Base.OneTo(2))) for i in axes(w, 1), j in axes(w, 2)]
        sh = [zeros((axes(h[1])..., Base.OneTo(2))) for i in axes(h, 1), j in axes(h, 2)]
        sb = [zeros((axes(b[1])..., Base.OneTo(2))) for i in axes(b, 1), j in axes(b, 2)]

        new{typeof(w),typeof(sw),typeof(h),typeof(sh)}(w, prim, sw, h, b, sh, sb)
    end

    function Solution2D2F(
        w::AbstractArray,
        prim::AbstractArray,
        sw::AbstractArray,
        h::AbstractArray,
        b::AbstractArray,
        sh::AbstractArray,
        sb::AbstractArray,
    )
        new{typeof(w),typeof(sw),typeof(h),typeof(sh)}(w, prim, sw, h, b, sh, sb)
    end

end


text:
using BandedMatrices, MatrixFactorizations, LinearAlgebra, Test, Random

Random.seed!(0)

@testset "QR tests" begin
    for T in (Float64,ComplexF64,Float32,ComplexF32)
        A=brand(T,10,10,3,2)
        Q,R=qr(A)
        @test Matrix(Q)*Matrix(R) ≈ A
        b=rand(T,10)
        @test mul!(similar(b),Q,mul!(similar(b),Q',b)) ≈ b
        for j=1:size(A,2)
            @test Q' * A[:,j] ≈ R[:,j]
        end
        A=brand(T,14,10,3,2)

        Q,R=qr(A)
        @test Matrix(Q)*Matrix(R) ≈ A

        for k=1:size(A,1),j=1:size(A,2)
            @test Q[k,j] ≈ Matrix(Q)[k,j]
        end

        A=brand(T,10,14,3,2)
        Q,R=qr(A)
        @test Matrix(Q)*Matrix(R) ≈ A

        for k=1:size(Q,1),j=1:size(Q,2)
            @test Q[k,j] ≈ Matrix(Q)[k,j]
        end

        A=brand(T,100,100,3,4)
        @test qr(A).factors ≈ LinearAlgebra.qrfactUnblocked!(Matrix(A)).factors
        @test qr(A).τ ≈ LinearAlgebra.qrfactUnblocked!(Matrix(A)).τ
        b=rand(T,100)
        @test qr(A)\b ≈ Matrix(A)\b
        b=rand(T,100,2)
        @test qr(A)\b ≈ Matrix(A)\b
        @test_throws DimensionMismatch qr(A) \ randn(3)
        @test_throws DimensionMismatch qr(A).Q'randn(3)

        A=brand(T,102,100,3,4)
        @test qr(A).factors ≈ LinearAlgebra.qrfactUnblocked!(Matrix(A)).factors
        @test qr(A).τ ≈ LinearAlgebra.qrfactUnblocked!(Matrix(A)).τ
        b=rand(T,102)
        @test qr(A)\b ≈ Matrix(A)\b
        b=rand(T,102,2)
        @test qr(A)\b ≈ Matrix(A)\b
        @test_throws DimensionMismatch qr(A) \ randn(3)
        @test_throws DimensionMismatch qr(A).Q'randn(3)

        A=brand(T,100,102,3,4)
        @test qr(A).factors ≈ LinearAlgebra.qrfactUnblocked!(Matrix(A)).factors
        @test qr(A).τ ≈ LinearAlgebra.qrfactUnblocked!(Matrix(A)).τ
        b=rand(T,100)
        @test_broken qr(A)\b ≈ Matrix(A)\b

        A = Tridiagonal(randn(T,99), randn(T,100), randn(T,99))
        @test qr(A).factors ≈ LinearAlgebra.qrfactUnblocked!(Matrix(A)).factors
        @test qr(A).τ ≈ LinearAlgebra.qrfactUnblocked!(Matrix(A)).τ
        b=rand(T,100)
        @test qr(A)\b ≈ Matrix(A)\b
        b=rand(T,100,2)
        @test qr(A)\b ≈ Matrix(A)\b
        @test_throws DimensionMismatch qr(A) \ randn(3)
        @test_throws DimensionMismatch qr(A).Q'randn(3)
    end

    @testset "Mixed types" begin
        A=brand(10,10,3,2)
        b=rand(ComplexF64,10)
        Q,R=qr(A)
        @test R\(Q'*b) ≈ qr(A)\b ≈ Matrix(A)\b

        A=brand(ComplexF64,10,10,3,2)
        b=rand(10)
        Q,R=qr(A)
        @test R\(Q'*b) ≈ qr(A)\b ≈ Matrix(A)\b

        A = BandedMatrix{Int}(undef, (2,1), (4,4))
        A.data .= 1:length(A.data)
        Q, R = qr(A)
        @test Q*R ≈ A
    end
end

@testset "QL tests" begin
    for T in (Float64,ComplexF64,Float32,ComplexF32)
        A=brand(T,10,10,3,2)
        Q,L=ql(A)
        @test ql(A).factors ≈ ql!(Matrix(A)).factors
        @test ql(A).τ ≈ ql!(Matrix(A)).τ
        @test Matrix(Q)*Matrix(L) ≈ A
        b=rand(T,10)
        @test mul!(similar(b),Q,mul!(similar(b),Q',b)) ≈ b
        for j=1:size(A,2)
            @test Q' * A[:,j] ≈ L[:,j]
        end

        A=brand(T,14,10,3,2)
        Q,L=ql(A)
        @test ql(A).factors ≈ ql!(Matrix(A)).factors
        @test ql(A).τ ≈ ql!(Matrix(A)).τ
        @test_broken Matrix(Q)*Matrix(L) ≈ A

        for k=1:size(A,1),j=1:size(A,2)
            @test Q[k,j] ≈ Matrix(Q)[k,j]
        end

        A=brand(T,10,14,3,2)
        Q,L=ql(A)
        @test ql(A).factors ≈ ql!(Matrix(A)).factors
        @test ql(A).τ ≈ ql!(Matrix(A)).τ
        @test Matrix(Q)*Matrix(L) ≈ A

        for k=1:size(Q,1),j=1:size(Q,2)
            @test Q[k,j] ≈ Matrix(Q)[k,j]
        end

        A=brand(T,10,14,3,6)
        Q,L=ql(A)
        @test ql(A).factors ≈ ql!(Matrix(A)).factors
        @test ql(A).τ ≈ ql!(Matrix(A)).τ
        @test Matrix(Q)*Matrix(L) ≈ A

        for k=1:size(Q,1),j=1:size(Q,2)
            @test Q[k,j] ≈ Matrix(Q)[k,j]
        end

        A=brand(T,100,100,3,4)
        @test ql(A).factors ≈ ql!(Matrix(A)).factors
        @test ql(A).τ ≈ ql!(Matrix(A)).τ
        b=rand(T,100)
        @test ql(A)\b ≈ Matrix(A)\b
        b=rand(T,100,2)
        @test ql(A)\b ≈ Matrix(A)\b
        @test_throws DimensionMismatch ql(A) \ randn(3)
        @test_throws DimensionMismatch ql(A).Q'randn(3)

        A=brand(T,102,100,3,4)
        @test ql(A).factors ≈ ql!(Matrix(A)).factors
        @test ql(A).τ ≈ ql!(Matrix(A)).τ
        b=rand(T,102)
        @test_broken ql(A)\b ≈ Matrix(A)\b
        b=rand(T,102,2)
        @test_broken ql(A)\b ≈ Matrix(A)\b
        @test_throws DimensionMismatch ql(A) \ randn(3)
        @test_throws DimensionMismatch ql(A).Q'randn(3)

        A=brand(T,100,102,3,4)
        @test ql(A).factors ≈ ql!(Matrix(A)).factors
        @test ql(A).τ ≈ ql!(Matrix(A)).τ
        b=rand(T,100)
        @test_broken ql(A)\b ≈ Matrix(A)\b

        A = Tridiagonal(randn(T,99), randn(T,100), randn(T,99))
        @test ql(A).factors ≈ ql!(Matrix(A)).factors
        @test ql(A).τ ≈ ql!(Matrix(A)).τ
        b=rand(T,100)
        @test ql(A)\b ≈ Matrix(A)\b
        b=rand(T,100,2)
        @test ql(A)\b ≈ Matrix(A)\b
        @test_throws DimensionMismatch ql(A) \ randn(3)
        @test_throws DimensionMismatch ql(A).Q'randn(3)
    end

    @testset "lmul!/rmul!" begin
        A = brand(100,100,3,4)
        Q,R = qr(A)
        x = randn(100)
        b = randn(100,2)
        @test lmul!(Q, copy(x)) ≈ Matrix(Q)*x
        @test lmul!(Q, copy(b)) ≈ Matrix(Q)*b
        @test lmul!(Q', copy(x)) ≈ Matrix(Q)'*x
        @test lmul!(Q', copy(b)) ≈ Matrix(Q)'*b
        c = randn(2,100)
        @test rmul!(copy(c), Q) ≈ c*Matrix(Q)
        @test rmul!(copy(c), Q') ≈ c*Matrix(Q')

        A = brand(100,100,3,4)
        Q,L = ql(A)
        x = randn(100)
        b = randn(100,2)
        @test lmul!(Q, copy(x)) ≈ Matrix(Q)*x
        @test lmul!(Q, copy(b)) ≈ Matrix(Q)*b
        @test lmul!(Q', copy(x)) ≈ Matrix(Q)'*x
        @test lmul!(Q', copy(b)) ≈ Matrix(Q)'*b
        c = randn(2,100)
        @test rmul!(copy(c), Q) ≈ c*Matrix(Q)
        @test rmul!(copy(c), Q') ≈ c*Matrix(Q')
    end

    @testset "Mixed types" begin
        A=brand(10,10,3,2)
        b=rand(ComplexF64,10)
        Q,L=ql(A)
        @test L\(Q'*b) ≈ ql(A)\b ≈ Matrix(A)\b
        @test Q*L ≈ A

        A=brand(ComplexF64,10,10,3,2)
        b=rand(10)
        Q,L=ql(A)
        @test Q*L ≈ A
        @test L\(Q'*b) ≈ ql(A)\b ≈ Matrix(A)\b

        A = BandedMatrix{Int}(undef, (2,1), (4,4))
        A.data .= 1:length(A.data)
        Q, L = ql(A)
        @test_broken Q*L ≈ A
    end
end


text:
module UtilTest

using Test
using DiffPrecTest
using ProximalBase
using Random, LinearAlgebra, SparseArrays, Statistics

@testset "kron_sub" begin
    A = randn(10, 5)
    B = randn(7, 9)

    k = kron(A, B)

    I = [1]
    @test DiffPrecTest.kron_sub(A, B, I, I) == k[I, I]

    I = [1, 2, 3]
    @test DiffPrecTest.kron_sub(A, B, I, I) == k[I, I]
end

@testset "skron_sub" begin
    A = randn(10, 5)
    B = randn(7, 9)

    k = (kron(A, B) + kron(B, A)) / 2.

    I = [1]
    @test DiffPrecTest.skron_sub(A, B, I, I) == k[I, I]

    I = [1, 2, 3]
    @test DiffPrecTest.skron_sub(A, B, I, I) == k[I, I]

    I1 = [40, 41, 42]
    @test DiffPrecTest.skron_sub(A, B, I, I1) == k[I, I1]
end

@testset "getelem symkron" begin
    S = Symmetric( cov( randn(100, 5) ) )
    X = randn(10, 5)

    for i=1:10
        sKron = ( kron(S, X[i,:] * X[i, :]') + kron(X[i,:] * X[i, :]', S) ) / 2.

        for ci=1:25
            for ri=1:25
                v = DiffPrecTest._getElemSKron(S, view(X, i, :), ri, ci)
                @test sKron[ri, ci] ≈ v
            end
        end
    end

    S2 = Symmetric( cov( randn(100, 6) ) )
    sKron = ( kron(S, S2) + kron(S2, S) ) / 2.
    for ci=1:30
        for ri=1:30
            v = DiffPrecTest._getElemSKron(S, S2, ri, ci)
            @test sKron[ri, ci] ≈ v
        end
    end
end

@testset "getelem kron" begin
    S = Symmetric( cov( randn(100, 5) ) )
    X = randn(10, 5)

    for i=1:10
        k  = kron(S, X[i,:] * X[i, :]')
        k1 = kron(X[i,:] * X[i, :]', S)

        for ci=1:25
            for ri=1:25
                v  = DiffPrecTest._getElemKron(S, view(X, i, :), ri, ci)
                v1 = DiffPrecTest._getElemKron(view(X, i, :), S, ri, ci)
                @test k[ri, ci] ≈ v
                @test k1[ri, ci] ≈ v1
            end
        end
    end
end

end


text:
module JSCall

using WebIO, JSExpr
import WebIO: tojs
using JSON
abstract type AbstractJSObject end

"""
References objects stored in Javascript.
Maps the following expressions to actual calls on the Javascript side:
```julia
jso = JSObject(name, scope, typ)
# getfield:
x = jso.property # returns a new JSObject
# setfield
jso.property = "newval" # Works with JSObjects or Julia objects as newval
# call:
result = jso.func(args...) # works with julia objects and other JSObjects as args

# constructors are wrapped this way:
scene = jso.new.Scene() # same as in JS: scene = new jso.Scene()
```
"""
mutable struct JSObject <: AbstractJSObject
    # fields are private and not accessible via jsobject.field
    name::Symbol
    scope::Scope
    typ::Symbol
    # transporting the UUID allows us to have a uuid different from the objectid
    # which will help to better capture === equivalence on the js side.
    uuid::UInt64

    function JSObject(name::Symbol, scope::Scope, typ::Symbol)
        obj = new(name, scope, typ)
        setfield!(obj, :uuid, objectid(obj))
        finalizer(remove_js_reference, obj)
        return obj
    end

    function JSObject(name::Symbol, scope::Scope, typ::Symbol, uuid::UInt64)
        obj = new(name, scope, typ, uuid)
        finalizer(remove_js_reference, obj)
        return obj
    end
end

# define accessors
for name in (:name, :scope, :typ, :uuid)
    @eval $(name)(jso::JSObject) = getfield(jso, $(QuoteNode(name)))
end

"""
WebIO seems to purposefully not show symbols as variables, so we need our own Sym
"""
struct Sym
    symbol::Symbol
    Sym(x) = new(Symbol(x))
end

const object_pool_identifier = Sym("WebIO.object_pool")

"""
Removes an JSObject from the object pool!
"""
function remove_js_reference(jso::JSObject)
    enqueue_javascript!(
        scope(jso),
        js"delete $(object_pool_identifier)[$(uuid(jso))]"
    )
end

"""
    uuidstr(jso::JSObject)

Returns the uuid as a string
"""
uuidstr(jso::JSObject) = string(uuid(jso))

"""
    JSObject(jso::JSObject, typ::Symbol)

Copy constructor with a new `typ`
"""
function JSObject(jso::JSObject, typ::Symbol)
    jsonew = JSObject(name(jso), scope(jso), typ)
    # point new object to old one on the javascript side:
    evaljs(jso, js"$jsonew = $jso; undefined;")
    return jsonew
end

tojsexpr(x::String) = JSString(x)
tojsexpr(x::JSString) = x
tojsexpr(x) = error("$(typeof(x)) is not convertible to a JS string / expression")

"""
    evaljs(
        jso::JSObject, js::Union{JSString, String};
        try_fetch = false, try_seconds = 2
    )

Evaluates Javascript in the scope of JSObject
"""
function WebIO.evaljs(
        jso::JSObject, js::Union{JSString, AbstractString};
    )
    jse = tojsexpr(js)
    eval_javascript!(scope(jso), jse)
    return nothing
end

WebIO.showjs(io::IO, sym::Sym) = print(io, sym.symbol)
# Make interpolation work for JSObject
function WebIO.tojs(jso::JSObject)
    return js"$(object_pool_identifier)[$(uuidstr(jso))]"
end
function JSON.lower(jso::JSObject)
    return js"$(object_pool_identifier)[$(uuidstr(jso))]"
end

"""
Overloading getproperty to allow the same semantic as Javascript.
Since there is no `new` keyword in Julia like in JS, we missuse
jsobject.new, to return an instance of jsobject with a new modifier.

So this Javascript:
```js
obj = new Module.Constructor()
```

Will translates to the following Julia code:
```Julia
obj = Module.new.Constructor()
```
"""
function Base.getproperty(jso::JSObject, field::Symbol)
    if field === :new
        # Create a new instance of jso, with the `new` modifier
        return JSObject(jso, :new)
    else
        # allocate result object:
        result = JSObject(field, scope(jso), typ(jso))
        js = js"""
            var object = $jso
            var result = object.$(Sym(field))
            // if result is a class method, we need to bind it to the parent object
            if(result.bind != undefined){
                result = result.bind(object)
            }
            $(object_pool_identifier)[$(uuidstr(result))] = result
            undefined
        """
        evaljs(jso, js)
        return result
    end
end

function Base.setproperty!(jso::JSObject, name::Symbol, val)
    evaljs(jso, js"$(jso).$(Sym(name)) = $val; undefined;")
    return val
end

modifier(jso::JSObject) = getfield(jso, :typ) === :new ? "new " : ""

"""
Constructs the arguments for a JS call
"""
function get_args(args, kw_args)
    if isempty(kw_args)
        return join(map(x-> sprint(io-> WebIO.showjs(io, tojs(x))), args), ", ")
    elseif isempty(args)
        # tojs isn't recursive bug:
        return sprint() do io
            print(io, '{')
            for (k, v) in kw_args
                print(io, k, ':')
                WebIO.showjs(io, tojs(v))
                print(io, ',')
            end
            print(io, '}')
        end
    else
        # TODO: I'm not actually sure about this :D
        error("""
        Javascript only supports keyword arguments OR arguments.
        Found posititional arguments and keyword arguments
        """)
    end
end

"""
    jsobject(args...; kw_args...)

Call overload for JSObjects.
Only supports keyword arguments OR positional arguments.
"""
function (jso::JSObject)(args...; kw_args...)
    result = JSObject(:result, scope(jso), :call)
    input_args = get_args(args, kw_args)
    js = """
        var func = $(tojs(jso))
        var result = $(modifier(jso))func(
            $input_args
        )
        $(object_pool_identifier.symbol)[\"$(uuid(result))\"] = result
        undefined;
    """
    evaljs(jso, js)
    return result
end

"""
Indicates that no value is in an observable.
Sentinel needed to check if an Observable got a new value.
Can't use nothing, since the value itself could be nothing.
"""
struct NoValue end

JSON.lower(x::NoValue) = ""

const eval_queue = Dict{UInt64, Tuple{Bool, Vector{JSString}}}()

function set_batchmode!(scope::Scope, mode::Bool)
    bm, queue = get_queue!(scope)
    eval_queue[objectid(scope)] = (mode, queue)
    return bm
end
function get_queue!(scope::Scope; init_batchmode = true)
    id = objectid(scope)
    batchmode, queue = get!(eval_queue, id) do
        finalizer(scope) do s
            delete!(eval_queue, id)
        end
        (init_batchmode, JSString[])
    end
    return batchmode, queue
end

"""
    eval_queue!(scope::Scope)
Evals all the js code that queued up, empties queue & sets batchmode to false.
"""
function eval_queue!(scope::Scope)
    jss = queued_javascript!(scope)
    return evaljs(scope, jss)
end

"""
    fused(f, scope::Scope)

Executes function `f` with batchmode of scope set to batched!
This means, when called like:
```Julia
fused(scope) do
    res = JSObject.some_func(bla)
    res.some_val = 42
end
```
The Javascript execution inside the closure will get executed fused & in one go.
"""
function fused(f, scope::Scope)
    oldmode = set_batchmode!(scope, true)
    result = f()
    set_batchmode!(scope, oldmode)
    if !oldmode
        evaljs(scope, queued_javascript!(scope))
    end
    return result
end

fused(f, jsm::AbstractJSObject) = fused(f, scope(jsm))

"""
    queued_javascript(scope::Scope)
Get's one JSString containing all javascript that got queued up untill
calling this function. Empties the queue and sets mode to not batched!
"""
function queued_javascript!(scope::Scope)
    bm, queue = get_queue!(scope)
    return JSString(sprint() do io
        for elem in queue
            println(io, elem)
        end
        empty!(queue)
        set_batchmode!(scope, false)
    end)
end

function enqueue_javascript!(scope::Scope, js::JSString)
    bm, queue = get_queue!(scope)
    push!(queue, js)
    return scope
end

"""
    eval_javascript(scope::Scope, js::JSString)
Evals the javascript code `js` in `scope`. If batchmode is true,
it will just queue the execution. If not, it will execute all js up to this point.
"""
function eval_javascript!(scope::Scope, js::JSString)
    batchmode, queue = get_queue!(scope)
    enqueue_javascript!(scope, js)
    if !batchmode
        evaljs(scope, queued_javascript!(scope))
    end
end

queued_javascript!(jsm::AbstractJSObject) = queued_javascript!(scope(jsm))

struct JSModule <: AbstractJSObject
    scope::Scope
    mod::JSObject
    document::JSObject
    window::JSObject
    this::JSObject
    display_func # A function that gets called on show with the modules Scope
end
scope(x::JSModule) = getfield(x, :scope)

function make_renderable!(jsm::JSModule)
    jss = queued_javascript!(jsm)
    jss = js"""
        function (mod){
            $(object_pool_identifier) = {}
            $(object_pool_identifier)[$(uuidstr(jsm.mod))] = mod
            $(object_pool_identifier)[$(uuidstr(jsm.document))] = document
            $(object_pool_identifier)[$(uuidstr(jsm.window))] = window
            $(object_pool_identifier)[$(uuidstr(jsm.this))] = this
            $(jss) // execute all queued statements in onimport
        }
    """
    onimport(jsm.scope, jss)
    return jsm.display_func(jsm.scope)
end

function Base.show(
        io::IO, m::WebIO.WEBIO_NODE_MIME, x::JSModule
    )
    return show(io, m, WebIO.render(x))
end

function WebIO.render(jsm::JSModule)
    WebIO.render(make_renderable!(jsm))
end

"""
    Module(name::Symbol, url::String)

Wraps a Javascript library.
"""
function JSModule(name::Symbol, url::String)
    JSModule(identity, name, url)
end

function JSModule(display_func, name::Symbol, url::String)
    scope = Scope(imports = [url])
    # insert scope into our execution queue, witch batchmode set to true
    # Like this, all evaljs calls on this scope will get batched untill
    # display and then executed in one go in the onimport function.
    get_queue!(scope, init_batchmode = true)
    Observable{Any}(scope, "_jscall_value_comm", NoValue())
    mod = JSObject(name, scope, :module)
    document = JSObject(:document, scope, :module)
    window = JSObject(:window, scope, :module)
    this = JSObject(:window, scope, :module)
    return JSModule(scope, mod, document, window, this, display_func)
end

using JSExpr: jsexpr

macro jsfun(expr)
    func = copy(expr.args[1].args)
    arguments = func[2:end]
    jss = JSCall.JSExpr.jsstring(expr)
    expr = quote
        args = ($(arguments...),)
        idx = findfirst(x-> x isa JSObject, args)
        idx === nothing && error("At least one argument needs to be of type JSObject")
        jso = args[idx]
        s = scope(jso)
        any(x-> x isa JSObject && (scope(x) !== s), args) && error("All JSObjects need to come from the same Scope")
        argstr = join(map(x-> sprint(io-> WebIO.showjs(io, WebIO.tojs(x))), args), ", ")
        result = JSObject(:result, scope(jso), :object)
        result_js = WebIO.tojs(result)
        jss = JSString(string($(jss...)))
        jss = "$result_js = ($jss)($argstr)"
        evaljs(jso, jss)
        return result
    end
    return Expr(:function, Expr(:call, esc(func[1]), func[2:end]...), expr)
end

"""
    jlvalue(jso::JSObject)

Fetches the Julia value representing a certin JSObject
"""
function jlvalue(jso::JSObject)
    obs, sync = scope(jso).observs["_jscall_value_comm"]
    obs[] = NoValue()
    evaljs(jso, @js($obs[] = $jso))
    timeout = 5 # wait at most 5 seconds
    start = time()
    yield() # maybe we can already get the value on first iteration
    while time() - start < timeout
        obs[] !== NoValue() && return obs[]
        sleep(0.01)
    end
    error("Timed out while trying to fetch value")
end

export JSObject, JSModule, scope, @jsfun, jlvalue

end # module