text:
#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/build_literate/Analyze_NMR_IR_Data.ipynb)

# # Inversion Recovery Experiments
# The following code replicates the NMR data analysis in Figs. 4-6 and complements the paper with additional analyses that are not shown in the paper in the interest of brevity.

# For this analysis we need the following packages:
using MRIgeneralizedBloch
using DifferentialEquations
using LinearAlgebra
using LsqFit
using Statistics
import Pingouin
using Printf
using Formatting
using Plots
plotlyjs(bg = RGBA(31/255,36/255,36/255,1.0), ticks=:native); #hide #!nb

# The raw data is stored in a separate [github repository](https://github.com/JakobAsslaender/MRIgeneralizedBloch_NMRData) and the following functions return the URL to the individual files:
MnCl2_data(TRF_scale) = string("https://github.com/JakobAsslaender/MRIgeneralizedBloch_NMRData/blob/main/20210419_1mM_MnCl2/ja_IR_v2%20(", TRF_scale, ")/1/data.2d?raw=true")
BSA_data(TRF_scale)   = string("https://github.com/JakobAsslaender/MRIgeneralizedBloch_NMRData/blob/main/20210416_15%25BSA_2ndBatch/ja_IR_v2%20(", TRF_scale, ")/1/data.2d?raw=true");

# which can be loaded with functions implemented in this file:
include(string(pathof(MRIgeneralizedBloch), "/../../docs/src/load_NMR_data.jl"));

# ## MnCl``_2`` Sample
# ### ``T_2^{*,f}`` Estimation
# We estimate ``T_2^{*,f}`` by fitting a mono-exponential decay curve to the FID of the acquisition with ``T_\text{RF} = 22.8``μs and ``T_\text{i} = 5``s.
M = load_Data(MnCl2_data(1))
M = M[:,1]; # select Tᵢ = 5s

# The data was measured at the following time points in units of seconds:
T_dwell = 100e-6 # s
TE = T_dwell * ((1:length(M)) .+ 7) # s

# Note that the signal is an FID, so the phrase *echo* time is a bit misleading.

# The function [curve_fit](https://julianlsolvers.github.io/LsqFit.jl/latest/api/#LsqFit.curve_fit) from the [LsqFit.jl](https://julianlsolvers.github.io/LsqFit.jl/latest/) package is only implemented for real-valued models. To accommodate this, we need to split the data into its real and imaginary part:
TEreal = [TE;TE]
Mreal = [real(M);imag(M)];

# Here, we are using a simple mono-exponential model with a complex-valued scaling factor `p[1] + 1im p[2]`, the decay time ``T_2^{*,f} =`` `p[3]`, and the Larmor frequency `p[4]`:
FID_model(t, p) = @. [p[1] * exp(- t[1:end ÷ 2] / p[3]) * cos(p[4] * t[1:end ÷ 2]); p[2] * exp(- t[end ÷ 2 + 1:end] / p[3]) * sin(p[4] * t[end ÷ 2 + 1:end])];

# Fitting this model to the NMR data estimates ``T_2^{*,f}``:
fit = curve_fit(FID_model, TEreal, Mreal, [1, 1, 0.1, 0])
T₂star_MnCl2 = fit.param[3] # s
# seconds and its uncertainty (also in units of seconds)
stderror(fit)[3] # s

# Visually, the plot and the data show good agreement:
Mfitted = FID_model(TEreal, fit.param)
Mfitted = Mfitted[1:end÷2] + 1im * Mfitted[end÷2+1:end]
p = plot(xlabel="TE [s]", ylabel="|FID(TE)| [a.u.]")
plot!(p, TE, abs.(M), label="data")
plot!(p, TE, abs.(Mfitted), label=@sprintf("fit with T₂* = %2.3f ms", 1e3 * T₂star_MnCl2))
#md Main.HTMLPlot(p) #hide

# The relative residual norm of the fit, i.e. ``||\text{residual}||_2/||M||_2`` is
norm(fit.resid) / norm(M)

# Despite its small ``\ell_2``-norm, the Shapiro-Wilk test indicates that the residual is not Gaussian or normal distributed at a significance level of `α=0.05`
Pingouin.normality(fit.resid, α=0.05)
# We note that mono-exponential ``T_2^*`` decays assume a Lorentzian distributed magnetic field, which is in general an assumption rather than a well-founded theory.

# ### Mono-Exponential IR Model
# We performed several experiments in which we inverted the thermal equilibrium magnetization with rectangular π-pulses with the following pulse durations (in seconds):
Tʳᶠmin = 22.8e-6 # s - shortest Tʳᶠ possible on the NMR
TRF_scale = [1;2;5:5:40] # scaling factor
Tʳᶠ = TRF_scale * Tʳᶠmin # s

# and acquired inversion recovery data at exponentially spaced inversion times (in seconds):
Tᵢ = exp.(range(log(3e-3), log(5), length=20)) # s
Tᵢ .+= 12 * Tʳᶠmin + (13 * 15.065 - 5) * 1e-6 # s - correction factors

# We calculate the Rabi frequencies of the RF pulses and a finer grid of ``T_\text{i}`` to plot the IR model:
ω₁ = π ./ Tʳᶠ # rad/s
Tᵢplot = exp.(range(log(Tᵢ[1]), log(Tᵢ[end]), length=500)); # s

# After loading and normalizing the data
M = zeros(Float64, length(Tᵢ), length(TRF_scale))
for i = 1:length(TRF_scale)
    M[:,i] = load_first_datapoint(MnCl2_data(TRF_scale[i]))
end
M ./= maximum(M);

#src #############################################################################
#src # export MnCl2 data
#src #############################################################################
io = open(expanduser(string("~/Documents/Paper/2021_MT_IDE/Figures/IR_data_MnCl2.txt")), "w") #src
write(io, "TI_s") #src
for i = 1:length(Tʳᶠ) #src
    write(io, " z_$(@sprintf("%.2e", Tʳᶠ[i]))") #src
end #src
write(io, " \n")  #src

for j = 1:length(Tᵢ) #src
    write(io, "$(@sprintf("%.2e", Tᵢ[j])) ") #src
    for i = 1:length(Tʳᶠ) #src
        write(io, "$(@sprintf("%.2e", M[j,i])) ") #src
    end #src
    write(io, " \n") #src
end #src
close(io) #src
#src #############################################################################

# we analyze each inversion recovery curve that corresponds to a different ``T_\text{RF}`` separately. This allows us to fit a simple mono-exponential model
standard_IR_model(t, p) = @. p[1] - p[3] * exp(- t * p[2]);
# where `p[1]` is the thermal equilibrium magnetization, `p[2]` ``= T_1``, and `p[1] - p[3]` is the magnetization right after the inversion pulse or, equivalently, `Minv = p[1] / p[3] - 1` is the inversion efficiency, which is 1 for an ideal π-pulse and smaller otherwise. The parameters are initialized with
p0 = [1.0, 1.0, 2.0];

# and we can loop over ``T_\text{RF}`` to perform the fits:
param = similar(M[1,:], Vector{Float64}) #src
R₁ = similar(M[1,:])
Minv = similar(R₁)
residual = similar(R₁)
p = plot(xlabel="Tᵢ [s]", ylabel="zᶠ(Tʳᶠ, Tᵢ) [a.u.]")
for i = 1:length(TRF_scale)
    Mi = @view M[:,i]

    fit = curve_fit(standard_IR_model, Tᵢ, Mi, p0)
    param[i] = fit.param #src

    R₁[i] = fit.param[2]
    Minv[i] = fit.param[3] / fit.param[1] - 1

    residual[i] = norm(fit.resid) / norm(Mi)

    scatter!(p, Tᵢ, Mi, label=@sprintf("Tʳᶠ = %1.2es - data", Tʳᶠ[i]), color=i)
    plot!(p, Tᵢplot, standard_IR_model(Tᵢplot, fit.param), label=@sprintf("fit with R₁ = %.3f/s; MInv = %.3f", R₁[i], Minv[i]), color=i)
end
display(p) #!md
#md Main.HTMLPlot(p) #hide

#src #############################################################################
#src # export fitted curves
#src #############################################################################
Mp = [standard_IR_model(Tᵢplot, param[i]) for i=1:length(Tʳᶠ)] #src
io = open(expanduser(string("~/Documents/Paper/2021_MT_IDE/Figures/IR_monoExp_fit_MnCl2.txt")), "w") #src
write(io, "TI_s") #src
for i = 1:length(Tʳᶠ) #src
    write(io, " z_$(@sprintf("%.2e", Tʳᶠ[i]))") #src
end #src
write(io, " \n")  #src

for j = 1:length(Tᵢplot) #src
    write(io, "$(@sprintf("%.2e", Tᵢplot[j])) ") #src
    for i = 1:length(Tʳᶠ) #src
        write(io, "$(@sprintf("%.2e", Mp[i][j])) ") #src
    end #src
    write(io, " \n") #src
end #src
close(io) #src
#src #############################################################################

# Here, the data measured with different ``T_\text{RF}`` are indicated by markers in different colors, and the corresponding fits are the line plots in the same color. The fitted parameters are denoted in the legend. In the paper, we highlight the estimated inversion efficiency and the relaxation rate of the dataset acquired with ``T_\text{RF}=22.8``μs:
Minv[1]
#-
R₁[1] # 1/s
# and of the dataset acquired with ``T_\text{RF}=912``μs:
Minv[end]
#-
R₁[end] # 1/s
# The mean value of all R₁ estimates is
mean(R₁) # 1/s

# 1/s and its standard deviation in units of 1/s is
std(R₁) # 1/s

# The relative residual norm of the fits is on average
mean(residual)

# ### Global IR Fit
# As an alternative to individual fits to the inversion recovery curves with different ``T_\text{RF}``, we can also perform a global fit that accounts for the ``T_2^{*,f}`` decay during the inversion pulse. The model first simulates the ``T_2^{*,f}`` decay during the inversion pulse, followed by ``T_1`` recovery:
function Bloch_IR_model(p, Tʳᶠ, Tᵢ, T2)
    (m0, m0_inv, R₁) = p
    R2 = 1 / T2

    M = zeros(Float64, length(Tᵢ), length(Tʳᶠ))
    for i = 1:length(Tʳᶠ)
        ## simulate inversion pulse
        ω₁ = π / Tʳᶠ[i]
        H = [-R2 -ω₁ 0 ;
              ω₁ -R₁ R₁;
               0   0 0 ]

        m_inv = m0_inv * (exp(H * Tʳᶠ[i]) * [0,1,1])[2]

        ## simulate T1 recovery
        H = [-R₁ R₁*m0;
               0     0]

        for j = 1:length(Tᵢ)
            M[j,i] = m0 * (exp(H .* (Tᵢ[j] - Tʳᶠ[i] / 2)) * [m_inv,1])[1]
        end
    end
    return vec(M)
end;

# We use the previously estimated ``T_2^{*,f}`` value for the fit:
fit = curve_fit((x, p) -> Bloch_IR_model(p, Tʳᶠ, Tᵢ, T₂star_MnCl2), 1:length(M), vec(M), [ 1, .8, 1])

p = plot(xlabel="Tᵢ [s]", ylabel="zᶠ(Tʳᶠ, Tᵢ) [a.u.]")
for i=1:length(Tʳᶠ)
    scatter!(p, Tᵢ, M[:,i], label=@sprintf("Tʳᶠ = %1.2es", Tʳᶠ[i]), color=i)
    plot!(p, Tᵢplot, Bloch_IR_model(fit.param, Tʳᶠ[i], Tᵢplot, T₂star_MnCl2), label=@sprintf("Tʳᶠ = %1.2es", Tʳᶠ[i]), color=i)
end
display(p) #!md
#md Main.HTMLPlot(p) #hide

# With this global fit, we get a very similar relaxation rate in units of 1/s
R₁_MnCl2 = fit.param[3] # 1/s

# with an uncertainty (also in units of 1/s) of
stderror(fit)[3] # 1/s

# Note that the relative residual norm is somewhat increased compared to individual fits to each inversion recovery curve:
norm(fit.resid) / norm(M)

# ## Bovine Serum Albumin Sample
# ### ``T_2^{*,f}`` Estimation
# We repeat the ``T_2^{*,f}`` estimation for the bovine serum albumin (BSA) sample by fitting a mono-exponential decay curve to the FID of the acquisition with ``T_\text{RF} = 22.8``μs and ``T_\text{i} = 5``s.

M = load_Data(BSA_data(1));
M = M[:,1] # select Tᵢ = 5s
Mreal = [real(M);imag(M)]

fit = curve_fit(FID_model, TEreal, Mreal, [1.0, 1.0, .1, 0.0]);

# The estimated ``T_2^{*,f}`` of the BSA sample is
T₂star_BSA = fit.param[3] # s
# seconds with an uncertainty of
stderror(fit)[3] # s
# seconds.

# Visually, the plot and the data align well for the BSA sample, too:
Mfitted = FID_model(TEreal, fit.param)
Mfitted = Mfitted[1:end÷2] + 1im * Mfitted[end÷2+1:end]
p = plot(xlabel="TE [s]", ylabel="|FID(TE)| [a.u.]")
plot!(p, TE, abs.(M), label="data")
plot!(p, TE, abs.(Mfitted), label=@sprintf("fit with T₂* = %2.3f ms", 1e3 * T₂star_BSA))
#md Main.HTMLPlot(p) #hide

# The relative residual norm (``||\text{residual}||_2/||M||_2``) is
norm(fit.resid) / norm(M)

# Despite the small residual, the Shapiro-Wilk test indicates that the residual is not normal distributed for this sample either:
Pingouin.normality(fit.resid, α=0.05)

# ### Mono-Exponential IR Model
# We also fit a mono-exponential model to each inversion recovery curve of the BSA data:
M = zeros(Float64, length(Tᵢ), length(TRF_scale))
for i = 1:length(TRF_scale)
    M[:,i] = load_first_datapoint(BSA_data(TRF_scale[i]))
end
M ./= maximum(M)
#src #############################################################################
#src # export BSA data
#src #############################################################################
io = open(expanduser(string("~/Documents/Paper/2021_MT_IDE/Figures/IR_data_BSA.txt")), "w") #src
write(io, "TI_s") #src
for i = 1:length(Tʳᶠ) #src
    write(io, " z_$(@sprintf("%.2e", Tʳᶠ[i]))") #src
end #src
write(io, " \n")  #src

for j = 1:length(Tᵢ) #src
    write(io, "$(@sprintf("%.2e", Tᵢ[j])) ") #src
    for i = 1:length(Tʳᶠ) #src
        write(io, "$(@sprintf("%.2e", M[j,i])) ") #src
    end #src
    write(io, " \n") #src
end #src
close(io) #src
#src #############################################################################

p = plot(xlabel="Tᵢ [s]", ylabel="zᶠ(Tʳᶠ, Tᵢ) [a.u.]")
for i = 1:length(TRF_scale)
    Mi = @view M[:,i]

    fit = curve_fit(standard_IR_model, Tᵢ, Mi, p0)
    param[i] = fit.param #src

    R₁[i] = fit.param[2]
    Minv[i] = fit.param[3] / fit.param[1] - 1
    residual[i] = norm(fit.resid) / norm(Mi)

    scatter!(p, Tᵢ, Mi, label=@sprintf("Tʳᶠ = %1.2es - data", Tʳᶠ[i]), color=i)
    plot!(p, Tᵢplot, standard_IR_model(Tᵢplot, fit.param), label=@sprintf("fit with R₁ = %.3f/s; MInv = %.3f", R₁[i], Minv[i]), color=i)
end
display(p) #!md
#md Main.HTMLPlot(p) #hide

#src #############################################################################
#src # export fitted curves
#src #############################################################################
Mp = [standard_IR_model(Tᵢplot, param[i]) for i=1:length(Tʳᶠ)] #src
io = open(expanduser(string("~/Documents/Paper/2021_MT_IDE/Figures/IR_monoExp_fit_BSA.txt")), "w") #src
write(io, "TI_s") #src
for i = 1:length(Tʳᶠ) #src
    write(io, " z_$(@sprintf("%.2e", Tʳᶠ[i]))") #src
end #src
write(io, " \n")  #src

for j = 1:length(Tᵢplot) #src
    write(io, "$(@sprintf("%.2e", Tᵢplot[j])) ") #src
    for i = 1:length(Tʳᶠ) #src
        write(io, "$(@sprintf("%.2e", Mp[i][j])) ") #src
    end #src
    write(io, " \n") #src
end #src
close(io) #src
#src #############################################################################

# Zooming into the early phase of the recovery curve reveals the poor fit quality, in particular for long ``T_\text{RF}``. This is also reflected by a substantially larger relative residual norm compared to the MnCl``_2`` sample:
mean(residual)

# In the paper, we highlight the estimated inversion efficiency and the relaxation rate of the dataset acquired with ``T_\text{RF}=22.8``μs
Minv[1]
#-
R₁[1] # 1/s
# and of the dataset acquired with ``T_\text{RF}=912``μs
Minv[end]
#-
R₁[end] # 1/s
# The mean value of all R₁ estimates is
mean(R₁) # 1/s

# and its standard deviation is substantially larger compared to the same fit of the MnCl``_2`` sample:
std(R₁) # 1/s

# ### Global IR Fit - Generalized Bloch Model
# In order to repeat the global fit that includes all ``T_\text{RF}`` values, we have to account for the spin dynamics in the semi-solid pool during the RF-pulse. First, we do this with the proposed generalized Bloch model:
function gBloch_IR_model(p, G, Tʳᶠ, TI, R2f)
    (m0, m0f_inv, m0s, R₁, T₂ˢ, Rx) = p
    m0f = 1 - m0s
    ω₁ = π ./ Tʳᶠ

    m0vec = [0, 0, m0f, m0s, 1]
    m_fun(p, t; idxs=nothing) = typeof(idxs) <: Number ? 0.0 : zeros(5)

    H = [-R₁-m0s*Rx     m0f*Rx R₁*m0f;
             m0s*Rx -R₁-m0f*Rx R₁*m0s;
              0          0         0 ]

    M = zeros(Float64, length(TI), length(Tʳᶠ))
    for i = 1:length(Tʳᶠ)
        param = (ω₁[i], 1, 0, m0s, R₁, R2f, Rx, R₁, T₂ˢ, G)
        prob = DDEProblem(apply_hamiltonian_gbloch!, m0vec, m_fun, (0.0, Tʳᶠ[i]), param)
        m = solve(prob)[end]

        for j = 1:length(TI)
            M[j,i] = m0 * (exp(H .* (TI[j] - Tʳᶠ[i] / 2)) * [m0f_inv * m[3],m[4],1])[1]
        end
    end
    return vec(M)
end;

# Here, we use assume a super-Lorentzian lineshape, whose Green's function is interpolated to speed up the fitting routine:
T₂ˢ_min = 5e-6 # s
G_superLorentzian = interpolate_greens_function(greens_superlorentzian, 0, maximum(Tʳᶠ)/T₂ˢ_min);

# The fit is initialized with `p0 = [m0, m0f_inv, m0_s, R₁, T2_s, Rx]` and we set some reasonable bounds to the fitted parameters:
p0   = [  1, 0.932,  0.1,   1, 10e-6, 50]
pmin = [  0, 0.100,   .0, 0.3,  1e-9, 10]
pmax = [Inf,   Inf,  1.0, Inf, 20e-6,1e3]

fit_gBloch = curve_fit((x, p) -> gBloch_IR_model(p, G_superLorentzian, Tʳᶠ, Tᵢ, 1/T₂star_BSA), [], vec(M), p0, lower=pmin, upper=pmax);

# Visually, the plot and the data align well:
p = plot(xlabel="Tᵢ [s]", ylabel="zᶠ(Tʳᶠ, Tᵢ) [a.u.]")
for i=1:length(Tʳᶠ)
    scatter!(p, Tᵢ, M[:,i], label=@sprintf("Tʳᶠ = %1.2es", Tʳᶠ[i]), color=i)
    plot!(p, Tᵢplot, gBloch_IR_model(fit_gBloch.param, G_superLorentzian, Tʳᶠ[i], Tᵢplot, 1/T₂star_BSA), label=@sprintf("Tʳᶠ = %1.2es", Tʳᶠ[i]), color=i)
end
display(p) #!md
#md Main.HTMLPlot(p) #hide

# which becomes particularly apparent when zooming into the beginning of the inversion recovery curves. Further, the relative residual norm is much smaller compared to the mono-exponential fit:
norm(fit_gBloch.resid) / norm(M)

# The estimated parameters are
m0 = fit_gBloch.param[1]
#-
Minv = fit_gBloch.param[2]
#-
m0s = fit_gBloch.param[3]
#-
R₁ = fit_gBloch.param[4] # 1/s
#-
T₂ˢ = 1e6fit_gBloch.param[5] # μs
#-
Rx = fit_gBloch.param[6] # 1/s
# with the uncertainties (in the same order)
stderror(fit_gBloch)

#src #############################################################################
#src # export fitted curves
#src #############################################################################
Mp = reshape(gBloch_IR_model(fit_gBloch.param, G_superLorentzian, Tʳᶠ, Tᵢplot, 1/T₂star_BSA), length(Tᵢplot), length(Tʳᶠ)) #src
io = open(expanduser(string("~/Documents/Paper/2021_MT_IDE/Figures/IR_gBloch_fit.txt")), "w") #src
write(io, "TI_s") #src
for i = 1:length(Tʳᶠ) #src
    write(io, " z_$(@sprintf("%.2e", Tʳᶠ[i]))") #src
end #src
write(io, " \n")  #src

for j = 1:length(Tᵢplot) #src
    write(io, "$(@sprintf("%.2e", Tᵢplot[j])) ") #src
    for i = 1:length(Tʳᶠ) #src
        write(io, "$(@sprintf("%.2e", Mp[j,i])) ") #src
    end #src
    write(io, " \n") #src
end #src
close(io) #src
#src #############################################################################

# ### Global IR Fit - Graham's Spectral Model
# For comparison, we repeat the same fit with [Graham's spectral model](http://dx.doi.org/10.1002/jmri.1880070520):
function Graham_IR_model(p, Tʳᶠ, TI, R2f)
    (m0, m0f_inv, m0s, R₁, T₂ˢ, Rx) = p
    m0f = 1 - m0s
    ω₁ = π ./ Tʳᶠ

    m0vec = [0, 0, m0f, m0s, 1]

    H = [-R₁-m0s*Rx     m0f*Rx R₁*m0f;
             m0s*Rx -R₁-m0f*Rx R₁*m0s;
              0          0         0 ]

    M = zeros(Float64, length(TI), length(Tʳᶠ))
    for i = 1:length(Tʳᶠ)
        param = (ω₁[i], 1, 0, Tʳᶠ[i], m0s, R₁, R2f, Rx, R₁, T₂ˢ)
        prob = ODEProblem(apply_hamiltonian_graham_superlorentzian!, m0vec, (0.0, Tʳᶠ[i]), param)
        m = solve(prob)[end]

        for j = 1:length(TI)
            M[j,i] = m0 * (exp(H .* (TI[j] - Tʳᶠ[i] / 2)) * [m0f_inv * m[3],m[4],1])[1]
        end
    end
    return vec(M)
end

fit_Graham = curve_fit((x, p) -> Graham_IR_model(p, Tʳᶠ, Tᵢ, 1/T₂star_BSA), [], vec(M), p0, lower=pmin, upper=pmax);

# Visually, the plot and the data align substantially worse:
p = plot(xlabel="Tᵢ [s]", ylabel="zᶠ(Tʳᶠ, Tᵢ) [a.u.]")
for i=1:length(Tʳᶠ)
    scatter!(p, Tᵢ, M[:,i], label=@sprintf("Tʳᶠ = %1.2es", Tʳᶠ[i]), color=i)
    plot!(p, Tᵢplot, Graham_IR_model(fit_Graham.param, Tʳᶠ[i], Tᵢplot, 1/T₂star_BSA), label=@sprintf("Tʳᶠ = %1.2es", Tʳᶠ[i]), color=i)
end
display(p) #!md
#md Main.HTMLPlot(p) #hide

# which becomes particularly apparent when zooming into the beginning of the inversion recovery curves. Further, the relative residual norm is much larger compared to the generalized Bloch fit:
norm(fit_Graham.resid) / norm(M)

# The estimated parameters are
m0 = fit_Graham.param[1]
#-
Minv = fit_Graham.param[2]
#-
m0s = fit_Graham.param[3]
#-
R₁ = fit_Graham.param[4] # 1/s
#-
T₂ˢ = 1e6fit_Graham.param[5] # μs
#-
Rx = fit_Graham.param[6] # 1/s
# with the uncertainties (in the same order)
stderror(fit_Graham)

#src #############################################################################
#src # export data
#src #############################################################################
#src export fitted curves
Mp = reshape(Graham_IR_model(fit_Graham.param, Tʳᶠ, Tᵢplot, 1/T₂star_BSA), length(Tᵢplot), length(Tʳᶠ)) #src
io = open(expanduser(string("~/Documents/Paper/2021_MT_IDE/Figures/IR_Graham_fit.txt")), "w") #src
write(io, "TI_s") #src
for i = 1:length(Tʳᶠ) #src
    write(io, " z_$(@sprintf("%.2e", Tʳᶠ[i]))") #src
end #src
write(io, " \n")  #src

for j = 1:length(Tᵢplot) #src
    write(io, "$(@sprintf("%.2e", Tᵢplot[j])) ") #src
    for i = 1:length(Tʳᶠ) #src
        write(io, "$(@sprintf("%.2e", Mp[j,i])) ") #src
    end #src
    write(io, " \n") #src
end #src
close(io) #src

# ### Global IR Fit - Sled's Model
# We also performed the fit with [Sled's model](http://dx.doi.org/10.1006/jmre.2000.2059):
function Sled_IR_model(p, G, Tʳᶠ, TI, R2f)
    (m0, m0f_inv, m0s, R₁, T₂ˢ, Rx) = p
    m0f = 1 - m0s
    ω₁ = π ./ Tʳᶠ

    m0vec = [0, 0, m0f, m0s, 1]

    H = [-R₁-m0s*Rx     m0f*Rx R₁*m0f;
             m0s*Rx -R₁-m0f*Rx R₁*m0s;
              0          0         0 ]

    M = zeros(Float64, length(TI), length(Tʳᶠ))
    for i = 1:length(Tʳᶠ)
        param = (ω₁[i], 1, 0, m0s, R₁, R2f, Rx, R₁, T₂ˢ, G)
        prob = ODEProblem(apply_hamiltonian_sled!, m0vec, (0.0, Tʳᶠ[i]), param)
        m = solve(prob)[end]

        for j = 1:length(TI)
            M[j,i] = m0 * (exp(H .* (TI[j] - Tʳᶠ[i] / 2)) * [m0f_inv * m[3],m[4],1])[1]
        end
    end
    return vec(M)
end

fit_Sled = curve_fit((x, p) -> Sled_IR_model(p, G_superLorentzian, Tʳᶠ, Tᵢ, 1/T₂star_BSA), [], vec(M), p0, lower=pmin, upper=pmax);

# Visually, the plot and the data do not align well either:
p = plot(xlabel="Tᵢ [s]", ylabel="zᶠ(Tʳᶠ, Tᵢ) [a.u.]")
for i=1:length(Tʳᶠ)
    scatter!(p, Tᵢ, M[:,i], label=@sprintf("Tʳᶠ = %1.2es", Tʳᶠ[i]), color=i)
    plot!(p, Tᵢplot, Sled_IR_model(fit_Sled.param, G_superLorentzian, Tʳᶠ[i], Tᵢplot, 1/T₂star_BSA), label=@sprintf("Tʳᶠ = %1.2es", Tʳᶠ[i]), color=i)
end
display(p) #!md
#md Main.HTMLPlot(p) #hide

# which becomes particularly apparent when zooming into the beginning of the inversion recovery curves. Further, the relative residual norm is also large compared to the generalized Bloch fit:
norm(fit_Sled.resid) / norm(M)

# The estimated parameters are
m0 = fit_Sled.param[1]
#-
Minv = fit_Sled.param[2]
#-
m0s = fit_Sled.param[3]
#-
R₁ = fit_Sled.param[4] # 1/s
#-
T₂ˢ = 1e6fit_Sled.param[5] # μs
#-
Rx = fit_Sled.param[6] # 1/s
# with the uncertainties (in the same order)
stderror(fit_Sled)

#src #############################################################################
#src # export data
#src #############################################################################
#src export fitted curves
Mp = reshape(Sled_IR_model(fit_Sled.param, G_superLorentzian, Tʳᶠ, Tᵢplot, 1/T₂star_BSA), length(Tᵢplot), length(Tʳᶠ)) #src
io = open(expanduser(string("~/Documents/Paper/2021_MT_IDE/Figures/IR_Sled_fit.txt")), "w") #src
write(io, "TI_s") #src
for i = 1:length(Tʳᶠ) #src
    write(io, " z_$(@sprintf("%.2e", Tʳᶠ[i]))") #src
end #src
write(io, " \n")  #src

for j = 1:length(Tᵢplot) #src
    write(io, "$(@sprintf("%.2e", Tᵢplot[j])) ") #src
    for i = 1:length(Tʳᶠ) #src
        write(io, "$(@sprintf("%.2e", Mp[j,i])) ") #src
    end #src
    write(io, " \n") #src
end #src
close(io) #src

# ### Analysis of the Residuals
# In order to visualize how well the three models align with the data at different ``T_\text{RF}``, we calculate the ``\ell_2``-norm of the residuals after subtracting the modeled from the measured signal and normalize it by the ``\ell_2``-norm of the signal:

resid_gBlo = similar(Tʳᶠ)
resid_Sled = similar(Tʳᶠ)
resid_Grah = similar(Tʳᶠ)
for i=1:length(Tʳᶠ)
    resid_gBlo[i] = norm(gBloch_IR_model(fit_gBloch.param, G_superLorentzian, Tʳᶠ[i], Tᵢ, 1/T₂star_BSA) .- M[:,i]) / norm(M[:,i])
    resid_Grah[i] = norm(Graham_IR_model(fit_Graham.param, Tʳᶠ[i], Tᵢ, 1/T₂star_BSA)                    .- M[:,i]) / norm(M[:,i])
    resid_Sled[i] = norm(Sled_IR_model(  fit_Sled.param,   G_superLorentzian, Tʳᶠ[i], Tᵢ, 1/T₂star_BSA) .- M[:,i]) / norm(M[:,i])
end

p = plot(xlabel="Tʳᶠ [s]", ylabel="relative residual")
scatter!(p, Tʳᶠ, resid_gBlo, label="generalized Bloch model")
scatter!(p, Tʳᶠ, resid_Grah, label="Graham's spectral model")
scatter!(p, Tʳᶠ, resid_Sled, label="Sled's model")
#md Main.HTMLPlot(p) #hide

#src #############################################################################
#src # export data
#src #############################################################################
io = open(expanduser(string("~/Documents/Paper/2021_MT_IDE/Figures/IR_residuum_each_fit.txt")), "w") #src
write(io, "TRF_ms Graham_percent Sled_percent gBloch_percent \n")  #src

for i = 1:length(Tʳᶠ) #src
    write(io, "$(@sprintf("%.2e", 1e3 * Tʳᶠ[i])) ") #src
    write(io, "$(@sprintf("%.2e", 1e2 * resid_Grah[i])) ") #src
    write(io, "$(@sprintf("%.2e", 1e2 * resid_Sled[i])) ") #src
    write(io, "$(@sprintf("%.2e", 1e2 * resid_gBlo[i])) ") #src
    write(io, " \n") #src
end #src
close(io) #src

# This analysis examines the residuals from the actual fits, i.e. it uses the biophysical parameters of respective fit to model the signal. The disadvantage of this approach is that residuals at long ``T_\text{RF}`` are negatively affected by the poor fits of Graham's and Sled's models at short ``T_\text{RF}``. This problem is overcome by subtracting the measured signal from signal that is simulated with the biophysical parameters that were estimated by fitting the generalized Bloch model:

for i=1:length(Tʳᶠ)
    resid_gBlo[i] = norm(gBloch_IR_model(fit_gBloch.param, G_superLorentzian, Tʳᶠ[i], Tᵢ, 1/T₂star_BSA) .- M[:,i]) / norm(M[:,i])
    resid_Grah[i] = norm(Graham_IR_model(fit_gBloch.param, Tʳᶠ[i], Tᵢ, 1/T₂star_BSA)                    .- M[:,i]) / norm(M[:,i])
    resid_Sled[i] = norm(Sled_IR_model(  fit_gBloch.param, G_superLorentzian, Tʳᶠ[i], Tᵢ, 1/T₂star_BSA) .- M[:,i]) / norm(M[:,i])
end

p = plot(xlabel="Tʳᶠ [s]", ylabel="relative residual")
scatter!(p, Tʳᶠ, resid_gBlo, label="generalized Bloch model")
scatter!(p, Tʳᶠ, resid_Grah, label="Graham's spectral model")
scatter!(p, Tʳᶠ, resid_Sled, label="Sled's model")
#md Main.HTMLPlot(p) #hide

# One can observe reduced residuals for Graham's and Sled's models for long ``T_\text{RF}`` as a trade off for larger residuals at short ``T_\text{RF}``. Yet, the residuals at long ``T_\text{RF}`` are still substantially larger compared to ones of the generalized Bloch model.

#src #############################################################################
#src # export data
#src #############################################################################
io = open(expanduser(string("~/Documents/Paper/2021_MT_IDE/Figures/IR_residuum_gBloch_est.txt")), "w") #src
write(io, "TRF_ms Graham_percent Sled_percent gBloch_percent \n")  #src

for i = 1:length(Tʳᶠ) #src
    write(io, "$(@sprintf("%.2e", 1e3 * Tʳᶠ[i])) ") #src
    write(io, "$(@sprintf("%.2e", 1e2 * resid_Grah[i])) ") #src
    write(io, "$(@sprintf("%.2e", 1e2 * resid_Sled[i])) ") #src
    write(io, "$(@sprintf("%.2e", 1e2 * resid_gBlo[i])) ") #src
    write(io, " \n") #src
end #src
close(io) #src


text:
using GLMakie
using Oceananigans
using Oceananigans.Units
using Printf

using Oceananigans.TurbulenceClosures:
    RiBasedVerticalDiffusivity,
    CATKEVerticalDiffusivity,
    ConvectiveAdjustmentVerticalDiffusivity,
    ExplicitTimeDiscretization

#####
##### Setup simulation
#####

convective_adjustment = ConvectiveAdjustmentVerticalDiffusivity(convective_κz=0.1, convective_νz=0.01)

grid = RectilinearGrid(size=32, z=(-256, 0), topology=(Flat, Flat, Bounded))
coriolis = FPlane(f=1e-4)

N² = 1e-5
Qᵇ = +1e-8
Qᵘ = -1e-3

b_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵇ))
u_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵘ))

closures_to_run = [
                   CATKEVerticalDiffusivity(),
                   RiBasedVerticalDiffusivity(),
                   convective_adjustment,
                   ]

for closure in closures_to_run

    model = HydrostaticFreeSurfaceModel(; grid, closure, coriolis,
                                        tracers = (:b, :e),
                                        buoyancy = BuoyancyTracer(),
                                        boundary_conditions = (; b=b_bcs, u=u_bcs))

    bᵢ(x, y, z) = N² * z
    set!(model, b = bᵢ)

    simulation = Simulation(model, Δt=10minute, stop_time=48hours)

    closurename = string(nameof(typeof(closure)))

    simulation.output_writers[:fields] =
        JLD2OutputWriter(model, merge(model.velocities, model.tracers),
                         schedule = TimeInterval(10minutes),
                         filename = "windy_convection_" * closurename,
                         overwrite_existing = true)

    progress(sim) = @info string("Iter: ", iteration(sim), " t: ", prettytime(sim))
    simulation.callbacks[:progress] = Callback(progress, IterationInterval(100))

    @info "Running a simulation of $model..."

    run!(simulation)
end

#####
##### Visualize
#####

b_ts = []
u_ts = []
v_ts = []
e_ts = []

for closure in closures_to_run
    closurename = string(nameof(typeof(closure)))
    filepath = "windy_convection_" * closurename * ".jld2"

    push!(b_ts, FieldTimeSeries(filepath, "b"))
    push!(u_ts, FieldTimeSeries(filepath, "u"))
    push!(v_ts, FieldTimeSeries(filepath, "v"))
    push!(e_ts, FieldTimeSeries(filepath, "e"))
end

b1 = first(b_ts)

z = znodes(b1)
Nt = length(b1.times)

fig = Figure(resolution=(1200, 800))

slider = Slider(fig[2, 1:2], range=1:Nt, startvalue=1)
n = slider.value

buoyancy_label = @lift "Buoyancy at t = " * prettytime(b1.times[$n])
velocities_label = @lift "Velocities at t = " * prettytime(b1.times[$n])
TKE_label = @lift "Turbulent kinetic energy t = " * prettytime(b1.times[$n])
ax_b = Axis(fig[1, 1], xlabel=buoyancy_label, ylabel="z")
ax_u = Axis(fig[1, 2], xlabel=velocities_label, ylabel="z")
ax_e = Axis(fig[1, 3], xlabel=TKE_label, ylabel="z")

xlims!(ax_b, -grid.Lz * N², 0)
xlims!(ax_u, -1.0, 1.0)
xlims!(ax_e, -1e-4, 3e-3)

colors = [:black, :blue, :red, :orange]

for (i, closure) in enumerate(closures_to_run)
    bn = @lift interior(b_ts[i][$n], 1, 1, :)
    un = @lift interior(u_ts[i][$n], 1, 1, :)
    vn = @lift interior(v_ts[i][$n], 1, 1, :)
    en = @lift interior(e_ts[i][$n], 1, 1, :)

    closurename = string(nameof(typeof(closure)))

    lines!(ax_b, bn, z, label=closurename, color=colors[i])
    lines!(ax_u, un, z, label="u, " * closurename, color=colors[i])
    lines!(ax_u, vn, z, label="v, " * closurename, linestyle=:dash, color=colors[i])
    lines!(ax_e, en, z, label="e, " * closurename, linestyle=:dash, color=colors[i])
end

axislegend(ax_b, position=:lb)
axislegend(ax_u, position=:rb)
axislegend(ax_e, position=:rb)

display(fig)

record(fig, "windy_convection.mp4", 1:Nt, framerate=24) do nn
    n[] = nn
end

ext: jl


text:
################################################################################
## Copyright 2016. Juri Hinz and Jeremy Yee.
## A collection of useful functions for Bellman recursion
################################################################################

## Determining the subgradient envelope on more than 1 point

function row_rearrange(subgrad::Array{Float64, 2}, grid::Array{Float64, 2})

    ## Extract sizes
    gnum = size(grid)[1] # number of grid points
    snum = size(grid)[2] # size of state

    ## Get the subgradient evelope
    result = Array{Float64}(gnum, snum)

    for i = 1:gnum
        result[i, :] = subgrad[indmax(subgrad * grid[i, :]), :]
    end

    return result

end

## Determining the subgradient envelope on only one point

function row_rearrange(subgrad::Array{Float64, 2}, grid::Array{Float64, 1})

    ## Extract sizes
    snum = length(grid) # size of state

    ## Get the subgradient evelope
    result = Array{Float64}(1, snum)
    result = subgrad[indmax(subgrad * grid), :]

    return result

end

## Compute the continuation value with just the row rearrangement operator

function expected(value::Array{Float64, 2}, grid::Array{Float64, 2},
                  disturb::Array{Float64, 3}, weight::Array{Float64, 1})

    ## Extract sizes
    gnum = size(value)[1]
    snum = size(value)[2]
    dnum = size(disturb)[3] # number of disturbances

    ## Approximate continuation value function
    result = zeros(Array{Float64}(gnum, snum))

    for i = 1:dnum
        subgrad = value * disturb[:, :, i] * weight[i]
        result += row_rearrange(subgrad, grid)
    end

    return result

end

## Compute the continuation value using k nearest neighbors using search function
## This should be used outside of the Bellman recursion

function expected_host(value::Array{Float64, 2}, grid::Array{Float64, 2},
                       disturb::Array{Float64, 3}, weight::Array{Float64, 1},
                       search::Function, k::Int64)

    ## Extract sizes
    gnum = size(value)[1]
    snum = size(value)[2]
    dnum = size(disturb)[3] # number of disturbances

    ## Construct search structure
    tree = search(transpose(grid))

    ## Approximate continuation value function
    result = zeros(Array{Float64}(gnum, snum))
    temp_subgrad = Array{Float64}(k, snum) # subgradients under consideration

    for i = 1:dnum
        temp_indices = knn(tree, disturb[:, :, i] * transpose(grid), k, true)[1]
        for g = 1:gnum
            temp_subgrad = value[temp_indices[g], :]
            result[g, :] += row_rearrange(temp_subgrad, grid[g, :]) * weight[i]
        end
    end

    return result

end

## Compute the continuation value using k nearest neighbors using search function
## This should used inside the Bellman recursion

function expected_host(value::Array{Float64, 2}, grid::Array{Float64, 2},
                       disturb::Array{Float64, 3}, weight::Array{Float64, 1},
                       indices::Array{Array{Int64, 1}, 1})

    ## Extract sizes
    gnum = size(value)[1]
    snum = size(value)[2]
    dnum = size(disturb)[3] # number of disturbances
    k = length(indices[1])

    ## Approximate continuation value function
    result = zeros(Array{Float64}(gnum, snum))
    temp_subgrad = Array{Float64}(k, snum) # subgradients under consideration

    for i = 1:dnum
        for g = 1:gnum
            temp_subgrad = value[indices[(i - 1) * gnum + g], :] * disturb[:, :, i]
            result[g, :] += row_rearrange(temp_subgrad, grid[g, :]) * weight[i]
        end
    end

    return result

end

## Compute the continuation value using the permutation methods

function expected_fast(value::Array{Float64, 2}, X::PermMat)

    ## Extract sizes
    gnum = size(value)[1]
    snum = size(value)[2]

    ## Approximate continuation value function
    result = zeros(Array{Float64}(gnum, snum))
    result = X.permMat[:, :, X.rnum + 1] * value * X.constMat

    for i = 1:X.rnum
        result[:, X.rIndex[i, 2]] += X.permMat[:, :, i] * value[:, X.rIndex[i, 1]]
    end

    return result

end


text:
using Colors
using DataFrames
using DataFramesMeta
using NearestNeighbors
using Random
using Statistics
using VegaLite

import Base.show
import Base64
import CairoMakie as MK
import ColorSchemes
import MultivariateStats

plot_molecules!(args...; kwargs...) = plot_molecules(args...; append=true, kwargs...)

function plot_molecules(df_spatial::DataFrame, polygons::Array{Matrix{Float64}, 1}=Matrix{Float64}[]; markersize=2,
        color::Union{Vector, Symbol, String}=:gene, size=(800, 800), poly_strokewidth=1, xlims=nothing, ylims=nothing, offset=(0, 0),
        is_noise::Union{Vector, BitArray, Symbol, Nothing}=nothing, annotation::Union{<:AbstractVector, Symbol, Nothing} = nothing,
        ann_colors::Union{Nothing, Dict} = nothing, legend=(annotation !== nothing), fontsize=8,
        noise_ann = nothing, shuffle_colors::Bool=false, append::Bool=false,
        polygon_kwargs::KWArgT=nothing, axis_kwargs::KWArgT=nothing, noise_kwargs::KWArgT=nothing, legend_kwargs::KWArgT=nothing, kwargs...)

    noise_args_default = (marker=:xcross, markersize=(0.75 * markersize), strokewidth=0, color="black");
    axis_args_default = (xticklabelsize=12, yticklabelsize=12, xticksvisible=false, xticklabelsvisible=false, yticksvisible=false, yticklabelsvisible=false);
    legend_args_default = (bgcolor=Colors.RGBA(1, 1, 1, 0.85),);
    polygon_args_default = (strokecolor="black", color="transparent", strokewidth=poly_strokewidth, label="")

    noise_kwargs = update_args(update_args(noise_args_default, Dict(kwargs...)), noise_kwargs)
    axis_kwargs = update_args(axis_args_default, axis_kwargs)
    legend_kwargs = update_args(legend_args_default, legend_kwargs)
    polygon_kwargs = update_args(polygon_args_default, polygon_kwargs)

    if typeof(color) === Symbol
        color = df_spatial[!,color]
    end

    if annotation !== nothing
        if typeof(annotation) === Symbol
            annotation = df_spatial[!,annotation]
        end

        annotation = ["$a" for a in annotation]
        if noise_ann !== nothing
            noise_ann = "$noise_ann"
        end
    end

    xlims = something(xlims, val_range(df_spatial.x))
    ylims = something(ylims, val_range(df_spatial.y))

    if !append
        fig = MK.Figure(resolution=size)
        fig[1, 1] = MK.Axis(fig; axis_kwargs...);
    end

    df_noise = nothing

    if typeof(is_noise) === Symbol
        is_noise = df_spatial[!,is_noise]
    end

    if is_noise !== nothing
        df_noise = df_spatial[is_noise,:]
        df_spatial = df_spatial[.!is_noise,:]
        color = color[.!is_noise]
    end

    if is_noise !== nothing
        MK.scatter!(df_noise.x .+ offset[1], df_noise.y .+ offset[2]; noise_kwargs...)
    end

    if annotation === nothing
        MK.scatter!(df_spatial.x .+ offset[1], df_spatial.y .+ offset[2]; color=color,
            strokewidth=0, markersize=markersize, kwargs...)
    else
        ann_vals = annotation[annotation .!= noise_ann] |> unique |> sort
        c_map = Colors.distinguishable_colors(length(ann_vals), Colors.colorant"#007a10", lchoices=range(20, stop=70, length=15))
        if shuffle_colors
            Random.shuffle!(c_map)
        end

        for (color, ann) in zip(c_map, ann_vals)
            style_dict = (ann_colors === nothing) ? Dict() : Dict(:color => ann_colors[ann])
            MK.scatter!(df_spatial.x[annotation .== ann] .+ offset[1], df_spatial.y[annotation .== ann] .+ offset[2];
                strokewidth=0, markersize=markersize, label=ann, color=color, style_dict..., kwargs...)
        end

        if noise_ann in annotation
            MK.scatter!(df_spatial.x[annotation .== noise_ann] .+ offset[1], df_spatial.y[annotation .== noise_ann] .+ offset[2];
                label=noise_ann, noise_kwargs...)
        end

        if legend
            MK.axislegend(;legend_kwargs...)
        end
    end

    if length(polygons) > 0
        MK.poly!([MK.Point2.(eachrow(p .+ [offset[1] offset[2]])) for p in polygons]; polygon_kwargs...)
    end

    MK.xlims!(MK.current_axis(), xlims .+ offset[1])
    MK.ylims!(MK.current_axis(), ylims .+ offset[2])

    return MK.current_figure()
end

function shuffle_labels(labels::Array{Int})
    new_labs = deepcopy(labels)
    mask = (new_labs .!= 0)
    new_labs[mask] = shuffle(1:maximum(labels))[new_labs[mask]]
    return new_labs
end

function shuffle_colors(colors::Vector)
    uniq_cols = unique(colors);
    col_ord = Dict(Pair.(uniq_cols, shuffle(1:length(uniq_cols))));
    return [uniq_cols[col_ord[tc]] for tc in colors]
end

function plot_expression_vectors(vecs...; gene_names::Vector{String}, min_expr_frac::Float64=0.025, alpha::Float64=0.5, text_offset::Float64=0.005,
        labels::Vector{String}=["y$i" for i in 1:length(vecs)], xrotation=-60, xticks::Bool=false, ylabel::String="Expression")
    y_vals = maximum(hcat(vecs...), dims=2) |> vec
    scale = sum(y_vals)

    p_df = DataFrame(:gene => gene_names)
    plt = @vlplot(x={:gene, title="Gene"}, tooltip=:gene, config={axisX={labels=xticks, ticks=false, labelAngle=xrotation}}, width=700, height=300)

    for (l,cnt) in zip(labels, vecs)
        p_df[!,l] = cnt
        plt += @vlplot({:bar, opacity=alpha}, y={l, title=ylabel}, color={datum=l})
    end
    p_df[!, :_yf] = y_vals
    p_df[!, :_yt] = y_vals .+ text_offset * scale

    return p_df |> (plt + @vlplot(:text, y=:_yt, text=:gene, transform=[{filter="datum._yf > $(min_expr_frac * scale)"}]))
end

### Tracing

function plot_num_of_cells_per_iterarion(tracer::Dict{Symbol, Any})
    if !(:n_components in keys(tracer)) || length(tracer[:n_components]) == 0
        error("No data about #components per iteration was stored")
    end

    n_components_per_iter = hcat(collect.(values.(tracer[:n_components]))...);
    labels = collect(keys(tracer[:n_components][1]));

    n_components_per_iter = n_components_per_iter[sortperm(labels),:]
    labels = sort(labels)

    fig = MK.Figure(resolution=(600, 400));
    axis_args = (xticksvisible=false, yticksvisible=false, xticklabelsize=14, yticklabelsize=14, xlabel="Iteration", ylabel="Num. cells",
        xlabelpadding=0, ylabelpadding=0, font="Helvetica")
    fig[1, 1] = MK.Axis(fig; title="Convergence", axis_args...);

    for i in 1:size(n_components_per_iter, 1)
        MK.lines!(n_components_per_iter[i,:], label="$(labels[i])", color=ColorSchemes.Dark2_8[mod(i - 1, 8) + 1])
    end
    MK.axislegend("Min #molecules"; position=(0.01, 1), nbanks = 2, titlesize=16, labelsize=14)
    MK.xlims!(MK.current_axis(), 0, size(n_components_per_iter, 2))
    MK.ylims!(MK.current_axis(), 0, maximum(n_components_per_iter))

    ax2 = MK.Axis(fig; bbox=MK.BBox(370, 550, 180, 330), axis_args...)
    subplot_start = ceil(Int, 0.75 * size(n_components_per_iter, 2))
    iter_vals = subplot_start:size(n_components_per_iter, 2)
    for i in 1:size(n_components_per_iter, 1)
        MK.lines!(iter_vals, n_components_per_iter[i, iter_vals], color=ColorSchemes.Dark2_8[mod(i - 1, 8) + 1])
    end

    MK.ylims!(MK.current_axis(), 0, maximum(n_components_per_iter[:, iter_vals]))

    return fig
end

### Diagnostics

function plot_clustering_convergence(clust_res::NamedTuple, title::String="")::VegaLite.VLSpec
    plt = DataFrame(:x => 1:length(clust_res.diffs), :diff => 100 .* clust_res.diffs, :change_frac => 100 .* clust_res.change_fracs) |>
        @vlplot(x={:x, title="Iteration"}, title=title, width=300, height=250) +
        @vlplot(:line, y={:diff, title="Change, %"}, color={datum="Max prob. difference"}) +
        @vlplot(:line, y={:change_frac}, color={datum="Molecules changed"})
end

### Colormaps

function map_to_colors(vals::Array{T, 1} where T; lims=nothing, palette::Vector=Colors.sequential_palette(0, 11))
    offset = (lims === nothing) ? minimum(vals) : lims[1]
    scale = (lims === nothing) ? maximum(vals) - offset : (lims[2] - lims[1])

    if lims !== nothing
        vals = min.(max.(vals, lims[1]), lims[2])
    end

    color_ticks = collect(range(0.0, 1.0, length=length(palette))) .* scale .+ offset
    colors = palette[floor.(Int, ((vals .- offset) ./ scale) .* (length(palette) - 1) .+ 1)]

    return Dict(:colors=>colors, :ticks=>color_ticks, :palette=>palette)
end

function distinguishable_colors(vals::T where T <: AbstractVector, args...; kwargs...)
    id_per_type = Dict(c => i for (i,c) in enumerate(unique(vals)));
    colors_uniq = Colors.distinguishable_colors(length(id_per_type), args...; kwargs...);
    colors = colors_uniq[get.(Ref(id_per_type), vals, 0)];

    return Dict(:colors=>colors, :ticks=>unique(vals), :palette=>colors_uniq)
end

plot_colorbar(colors; kwargs...) = plot_colorbar(colors[:ticks], colors[:palette]; kwargs...)

function plot_colorbar(color_ticks, palette; size=(500, 110), xticklabelsize=12,
        xlabelsize=14, xlabelpadding=0, kwargs...)
    fig = MK.Figure(resolution=size)
    fig[1, 1] = MK.Axis(fig; yticksvisible=false, yticklabelsvisible=false, xticksvisible=false, xticklabelsize=12,
        xlabelsize=14, xlabelpadding=0, kwargs...)
    MK.barplot!(color_ticks, ones(length(color_ticks)), color=palette)
    MK.xlims!(MK.current_axis(), val_range(color_ticks))

    return fig
end

### Utils

function Base.show(io::IO, m::MIME"text/html", fig::MK.Scene)
    io64 = Base64.IOBuffer();
    iob64 = Base64.Base64EncodePipe(io64)
    show(iob64, MIME("image/png"), fig);
    close(iob64)
    str = String(take!(io64))
    close(io64)

    print(io, "")
    # show(io, MIME("text/plain"), "")
end

max_stars_repo_name: repushko/Baysor

max_issues_repo_name: repushko/Baysor

max_forks_repo_name: repushko/Baysor


text:
struct SoundBuffer <: SoundObject{sfSoundBuffer}
    ptr::Ptr{sfSoundBuffer}
end

function SoundBuffer(file::AbstractString)
    ptr = sfSoundBuffer_createFromFile(file)
    ptr == C_NULL && error("Couldn't create SoundBuffer from file '$file'")
    return SoundBuffer(ptr)
end

function SoundBuffer(samples::AbstractArray{Int16}, count, channels, samplerate)
    ptr = sfSoundBuffer_createFromSamples(samples, count, channels, samplerate)
    ptr == C_NULL && error("Couldn't create SoundBuffer from samples.")
    return SoundBuffer(ptr)
end

function SoundBuffer(samples::AbstractArray{Int16})
    return SoundBuffer(samples, length(samples), 2, 44100)
end

destroy!(this::SoundBuffer) = sfSoundBuffer_destroy(this)
nchannels(buffer::SoundBuffer) = Int(sfSoundBuffer_getChannelCount(buffer))
duration(buffer::SoundBuffer) = sfSoundBuffer_getDuration(buffer).microseconds
nsamples(buffer::SoundBuffer) = sfSoundBuffer_getSampleCount(buffer)
samplerate(buffer::SoundBuffer) = Int(sfSoundBuffer_getSampleRate(buffer))

function samplesptr(buffer::SoundBuffer)
    return sfSoundBuffer_getSamples(buffer)
end

function getsamples(buffer::SoundBuffer)::Vector{Int16}
    N = nsamples(buffer)
    data = zeros(Int16, N)
    unsafe_copyto!(pointer(data), samplesptr(buffer), N)
    return data
end

function Base.getindex(buffer::SoundBuffer, idx...)
    N = nsamples(buffer)
    data = unsafe_wrap(Vector{Int16}, samplesptr(buffer), N)
    return getindex(data, idx...)
end

Base.size(buffer::SoundBuffer) = Int(nsamples(buffer))
Base.length(buffer::SoundBuffer) = size(buffer)

Base.firstindex(::SoundBuffer) = 1
Base.lastindex(buffer::SoundBuffer) = length(buffer)

Base.first(buffer::SoundBuffer) = buffer[ firstindex(buffer) ]
Base.last(buffer::SoundBuffer)  = buffer[ lastindex(buffer) ]

FFTW.fft(buffer::SoundBuffer) = fft(buffer[:])
FFTW.ifft(buffer::SoundBuffer) = ifft(buffer[:])


text:
using DatasToolbox
using DataFrames
using NullableArrays
using PyCall

df = makeTestDF(Int64, Float64, String, DateTime, Date, nrows=10^6)
df[2, 3] = Nullable()

pydf = pandas(df)

@time newdf = convertPyDF(pydf)

coltypes = [eltype(newdf[c]) for c in names(newdf)]

max_stars_repo_name: JuliaTagBot/DatasToolbox

max_issues_repo_name: JuliaTagBot/DatasToolbox

max_forks_repo_name: JuliaTagBot/DatasToolbox


text:
function rand_matrix{X<:Matrix}(A::Type{X}, N, T)
    rand(T, N, N)
end
function rand_matrix{X<:FixedSizeArrays.Mat}(A::Type{X}, N, T)
    rand(A{N, N, T})
end
function rand_matrix{X<:Tuple}(A::Type{X}, N, T)
  ntuple(N) do i
    tuple(rand(T, N)...)
  end
end
function rand_matrix{X<:VecElement}(A::Type{X}, N, T)
    map(VecElement, rand_matrix(Tuple, N, T))
end

function matrix_bench(suite)
    constructors = [
        ("Array", Matrix),
        ("FSA", FixedSizeArrays.Mat),
        ("VecTuple", VecElement),
        ("Tuple", Tuple)
    ]
    element_types = (Float32, Float64)

    for (name, c) in constructors
      suite[name] = BenchmarkGroup()
      for N in (2, 4, 9), T in element_types
          a = rand_matrix(c, N, T)
          b = rand_matrix(c, N, T)
          suite[name]["matmul", N, T] = @benchmarkable a * b
          suite[name]["sum", N, T] = @benchmarkable sum(a)
      end
    end
end


text:
abstract type AbstractCost end

include("quadratic_cost.jl")


text:
@enum STATUS INFO WARNING ERROR FATAL
::Type{enum STATUS (INFO, WARNING, ERROR, FATAL)}

INFO; # => INFO
STATUS(0); # => INFO
STATUS(1); # => WARNING
typeof(INFO); # => enum STATUS (INFO, WARNING, ERROR, FATAL)

The typeof(INFO) looks a bit strange so to see how this occurs let us look at the definition of the enum macro.

macro enum(T,syms...)
  blk = quote
    immutable $(esc(T))
            n::Int32
            $(esc(T))(n::Integer) = new(n)
    end
    Base.show(io::IO, x::$(esc(T))) = print(io, $syms[x.n+1])
    Base.show(io::IO, x::Type{$(esc(T))})
       = print(io,
       $(string("enum ", T, ' ', '(', join(syms, ", "), ')')))
  end
  for (i,sym) in enumerate(syms)
    push!(blk.args, :(const $(esc(sym)) = $(esc(T))($(i-1))))
  end
  push!(blk.args, :nothing)
  blk.head = :toplevel
  return blk
end

max_stars_repo_path: Module 3/Chapter04/enum.jl

max_issues_repo_path: Module 3/Chapter04/enum.jl

max_forks_repo_path: Module 3/Chapter04/enum.jl


text:
function sumsqrtn(n)
r = 0
for i = 1:n
r = r + sqrt(i)
end
return r
end

max_stars_repo_path: Module 2/Chapter03/sumsqrtn(n).jl

max_issues_repo_path: Module 2/Chapter03/sumsqrtn(n).jl

max_forks_repo_path: Module 2/Chapter03/sumsqrtn(n).jl


text:
# This file defines midlevel api wrappers. We include name normalization for methods that are
# applicable to different hdf5 api-layers. We still try to adhere close proximity to the underlying
# method name in the hdf5-library.

"""
    HDF5.set_extent_dims(dset::HDF5.Dataset, new_dims::Dims)

Change the current dimensions of a dataset to `new_dims`, limited by
`max_dims = get_extent_dims(dset)[2]`. Reduction is possible and leads to loss of truncated data.
"""
function set_extent_dims(dset::Dataset, size::Dims)
    checkvalid(dset)
    h5d_set_extent(dset, hsize_t[reverse(size)...])
end

"""
    HDF5.set_extent_dims(dspace::HDF5.Dataspace, new_dims::Dims, max_dims::Union{Dims,Nothing} = nothing)

Change the dimensions of a dataspace `dspace` to `new_dims`, optionally with the maximum possible
dimensions `max_dims` different from the active size `new_dims`. If not given, `max_dims` is set equal
to `new_dims`.
"""
function set_extent_dims(dspace::Dataspace, size::Dims, max_dims::Union{Dims,Nothing} = nothing)
    checkvalid(dspace)
    rank = length(size)
    current_size = hsize_t[reverse(size)...]
    maximum_size = isnothing(max_dims) ? C_NULL : [reverse(max_dims .% hsize_t)...]
    h5s_set_extent_simple(dspace, rank, current_size, maximum_size)
    return nothing
end

"""
    HDF5.get_extent_dims(obj::Union{HDF5.Dataspace, HDF5.Dataset, HDF5.Attribute}) -> dims, maxdims

Get the array dimensions from a dataspace, dataset, or attribute and return a tuple of `dims` and `maxdims`.
"""
function get_extent_dims(obj::Union{Dataspace,Dataset,Attribute})
    dspace = obj isa Dataspace ? checkvalid(obj) : dataspace(obj)
    h5_dims, h5_maxdims = h5s_get_simple_extent_dims(dspace)
    # reverse dimensions since hdf5 uses C-style order
    N = length(h5_dims)
    dims = ntuple(i -> @inbounds(Int(h5_dims[N-i+1])), N)
    maxdims = ntuple(i -> @inbounds(h5_maxdims[N-i+1]) % Int, N) # allows max_dims to be specified as -1 without triggering an overflow
    obj isa Dataspace || close(dspace)
    return dims, maxdims
end

"""
    silence_errors(f::Function)

During execution of the function `f`, disable printing of internal HDF5 library error messages.
"""
function silence_errors(f::Function)
    estack = H5E_DEFAULT
    func, client_data = h5e_get_auto(estack)
    h5e_set_auto(estack, C_NULL, C_NULL)
    try
        return f()
    finally
        h5e_set_auto(estack, func, client_data)
    end
end


text:
# src/AbstractWallets.jl
# Abstract Digital Wallets

__precompile__(true)

"""```
AbstractWallets
```
Abstractions Module for Digital Wallets.
See https://github.com/bhgomes/AbstractWallets.jl for more details.
"""
module AbstractWallets

include("core.jl")
include("portfolios.jl")
include("transactions.jl")
include("addresses.jl")
include("crypto.jl")

end  # module AbstractWallets


text:
# This file was generated, do not modify it. # hide
coefs, intercept = fitted_params(mtm.fitresult).linear_regressor
@show coefs
@show intercept


text:
#!/usr/bin/env julia

const OUTPUT = haskey(ENV, "OUTPUT")

function usage()
    prog = basename(Base.source_path())
#    println(STDERR,"Usage ",prog," <rows> <cols> <y1> <y2> <x1> <x2> <no. of threads><lamda> <no. of iter>")
    println(STDERR,"Usage ",prog," <rows> <cols> <y1> <y2> <x1> <x2> <lamda> <no. of iter>")
    println(STDERR,"\t<rows>   - number of rows")
    println(STDERR, "\t<cols>    - number of cols\n")
    println(STDERR, "\t<y1> 	 - y1 value of the speckle\n")
    println(STDERR, "\t<y2>      - y2 value of the speckle\n")
    println(STDERR, "\t<x1>       - x1 value of the speckle\n")
    println(STDERR, "\t<x2>       - x2 value of the speckle\n")
#    println(STDERR, "\t<no. of threads>  - no. of threads\n")
    println(STDERR, "\t<lamda>   - lambda (0,1)\n")
    println(STDERR, "\t<no. of iter>   - number of iterations\n")
    exit(1)
end

function main(args)
    if length(args) != 8
      usage()
    end
    rows = parse(Int64,args[1])
    cols = parse(Int64,args[2])
    if ((rows % 16) != 0) | ((cols % 16) != 0)
        println(STDERR,"rows and cols must be mutiples of 16")
        exit(1)
    end
    r1 = parse(Int32,args[3]) + 1
    r2 = parse(Int32,args[4]) + 1
    c1 = parse(Int32,args[5]) + 1
    c2 = parse(Int32,args[6]) + 1
    lambda = parse(Float32,args[7])
    niter = parse(Int32,args[8])

    size_R = (r2 - r1 + 1) * (c2 - c1 + 1);

    iN::Array{Int32,1} = [i-1 for i in 1:rows]
    iS::Array{Int32,1} = [i+1 for i in 1:rows]
    jW::Array{Int32,1} = [j-1 for j in 1:cols]
    jE::Array{Int32,1} = [j+1 for j in 1:cols]
    iN[1] = 1
    iS[rows] = rows
    jW[1] = 1
    jE[cols] = cols

    dN::SharedArray{Float32,2} = SharedArray(Float32,(rows,cols))
    dS::SharedArray{Float32,2} = SharedArray(Float32,(rows,cols))
    dW::SharedArray{Float32,2} = SharedArray(Float32,(rows,cols))
    dE::SharedArray{Float32,2} = SharedArray(Float32,(rows,cols))

    println("Randomizing the input matrix")

    srand(7)
    I = rand(Float32,rows,cols)
    J::SharedArray{Float32,2} = SharedArray(Float32,(rows,cols),init = J -> J[Base.localindexes(J)] = exp.(I[Base.localindexes(J)]))
    c::SharedArray{Float32,2} = SharedArray(Float32,(rows,cols))

    println("Start the SRAD main loop")

    tic()

    sum::Float32 = 0
    sum2::Float32 = 0
    for iter in 1:niter
        sum = 0
        sum2 = 0
        for i in r1:r2
            for j in c1:c2
                tmp = J[i,j]
                sum += tmp
                sum2 += tmp * tmp
            end
        end
    end

    meanROI::Float32 = sum / size_R
    varROI::Float32 = (sum / size_R) - meanROI * meanROI
    q0sqr::Float32 = varROI / (meanROI * meanROI)

    @sync @parallel for i in 1:size(J,1)
        for j in size(J,2)
            Jc = J[i,j]

            # directional derivates
            dN[i,j] = J[iN[i],j] - Jc
            dS[i,j] = J[iS[i],j] - Jc
            dW[i,j] = J[i,jW[j]] - Jc
            dE[i,j] = J[i,jE[j]] - Jc

            G2 = (dN[i,j] * dN[i,j] + dS[i,j] * dS[i,j] + dW[i,j] * dW[i,j] +
                  dE[i,j] * dE[i,j]) /
                 (Jc * Jc)

            L = (dN[i,j] + dS[i,j] + dW[i,j] + dE[i,j]) / Jc

            num = (0.5 * G2) - ((1.0 / 16.0) * (L * L))
            den = 1 + (0.25 * L)
            qsqr = num / (den * den)

            # diffusion coefficent (equ 33)
            den = (qsqr - q0sqr) / (q0sqr * (1 + q0sqr))
            c[i,j] = 1.0 / (1.0 + den)

            # saturate diffusion coefficent
            if (c[i,j] < 0)
                c[i,j] = 0
            elseif (c[i,j] > 1)
                c[i,j] = 1
            end
        end
    end

    @sync @parallel for i in 1:size(J,1)
        for j in 1:size(J,2)
            # diffusion coefficient
            cN = c[i,j]
            cS = c[iS[i],j]
            cW = c[i,j]
            cE = c[i,jE[j]]

            # divergence (equ 58)
            D = cN * dN[i,j] + cS * dS[i,j] + cW * dW[i,j] + cE * dE[i,j]

            # image update (equ 61)
            J[i,j] += 0.25 * lambda * D
        end
    end

    toc()

    println("Computation done")

    if OUTPUT
        f = open("output.txt","w")
        for i in 1:size(J,1)
            for j in 1:size(J,2)
               print(f,@sprintf("%.5f ",J[i,j]))
            end
            println(f)
        end
        close(f)
    end
end


text:
# A collection of smooth domains

## A kite-shaped domain

export Kite
"A kite-shaped domain in `ℝ^2`."
struct Kite{T} <: Domain{SVector{2,T}}
    a   ::  Int       # a parameter
end

Kite() = Kite{Float64}()
Kite{T}() where {T} = Kite{T}(2)
Kite(a::Integer) = Kite{float(typeof(a))}(a)

indomain(x, ::Kite) = error("Don't know how to evaluate indomain for a kite.")

"""
Smooth parameterization of the kite-shaped domain, mapping from `[0,1)]` to the kite.
"""
struct KiteMap{T} <: Parameterization{T}
    a   ::  Int       # a parameter
end

hasparameterization(d::Kite) = true
parameterization(d::Kite{T}) where T = KiteMap{T}(d.a)

domain(d::KiteMap{T}) where T = Interval{:closed,:open,T}(0, 1)
image(m::KiteMap{T}) where T = Kite{T}(m.a)

function applymap(m::KiteMap, t)
    b = 2*convert(domaintype(m), pi)
    SVector(-sin(b*t)-cos(m.a*b*t), cos(b*t))
end

function gradient(m::KiteMap, t)
    b = 2*convert(domaintype(m), pi)
    SVector(-b*cos(b*t)+m.a*b*sin(m.a*b*t), -b*sin(b*t))
end

function gradient_derivative(m::KiteMap, t)
    b = 2*convert(domaintype(m), pi)
    SVector(b^2*sin(b*t)+(m.a*b)^2*cos(m.a*b*t), -b^2*cos(b*t))
end

## An ellipse

export Ellipse
"An ellipse-shaped domain"
struct Ellipse{T} <: Domain{SVector{2,T}}
    center  ::  SVector{2,T}        # center of the ellipse
    radius  ::  SVector{2,T}        # (radius1,radius2)

    Ellipse{T}(center = SVector(zero(T),zero(T)), radius = SVector(one(T),one(T))) where {T} = new(center, radius)
end

Ellipse() = Ellipse{Float64}()

Ellipse(center::SVector{2,T}, radius::SVector{2,S}) where {S,T} = Ellipse{promote_type(S,T)}(center, radius)

Ellipse(x, y, radius1, radius2) = Ellipse(SVector(x, y), SVector(radius1, radius2))

indomain(x, d::Ellipse) = ((x[1]-d.center[1])/d.radius[1])^2 + ((x[2]-d.center[2])/d.radius[2])^2 ≈ 1

"""
Smooth parameterization of an ellipse, mapping from `[0,1)]`.
"""
struct EllipseMap{T} <: Parameterization{T}
    center  ::  SVector{2,T}        # center of the ellipse
    radius  ::  SVector{2,T}        # (radius1,radius2)
end

hasparameterization(d::Ellipse) = true
parameterization(d::Ellipse) = EllipseMap(d.center, d.radius)

domain(m::EllipseMap{T}) where T = Interval{:closed,:open,T}(0, 1)
image(m::EllipseMap) = Ellipse(m.center, m.radius)

function applymap(m::EllipseMap, t)
    b = 2*convert(domaintype(m), pi)
    m.center + SVector(m.radius[1]*cos(b*t), m.radius[2]*sin(b*t))
end

function gradient(m::EllipseMap, t)
    b = 2*convert(domaintype(m), pi)
    SVector(-m.radius[1]*b*sin(b*t),  m.radius[2]*b*cos(b*t))
end

function gradient_derivative(m::EllipseMap, t)
    b = 2*convert(domaintype(m), pi)
    SVector(m.radius[1]*b^2*cos(b*t),  m.radius[2]*b^2*sin(b*t))
end

# # function normal{T}(b::Ellipse{T}, t)
# #     n = Vec2{T}(b.radius2 * cos(2*T(pi)*t), b.radius1 * sin(2*T(pi)*t))
# #     n / norm(n)
# # end
#
# gradient_d2{T}(b::Ellipse{T}, t) = Vec2{T}(-b.radius1 * (2 * T(pi))^2 * cos(2*T(pi)*t),  -b.radius2 * (2 * T(pi))^2 * sin(2*T(pi)*t))
# gradient_d3{T}(b::Ellipse{T}, t) = Vec2{T}( b.radius1 * (2 * T(pi))^3 * sin(2*T(pi)*t),  -b.radius2 * (2 * T(pi))^3 * cos(2*T(pi)*t))


text:
#!/usr/bin/env julia

using Luxor

using Test

using Random
Random.seed!(42)

function poly_split_2(fname)
    Drawing(2500, 2500, fname)
    background(1, .8, .1)
    origin()
    setopacity(0.7)
    fontsize(2)
    sethue("black")
    s = squircle(O, 600, 600, vertices=true)
    s = box(O, 60, 60, vertices=true)
    setline(2)
    tiles = Tiler(2500, 2500, 6, 6, margin=25)
    for (pos, n) in tiles
        gsave()
        s = squircle(O, tiles.tileheight/2 - 10, tiles.tileheight/2 - 10, vertices=true)
        translate(pos)
        pt1 = Point(O.x + rand(-120:120), -tiles.tileheight/2)
        pt2 = Point(O.x + rand(-120:120), tiles.tileheight/2)
        setdash("dot")
        move(pt1)
        rline(pt2 - pt1)
        strokepath()
        setdash("solid")
        poly1, poly2 = polysplit(s, pt1, pt2)
        randomhue()
        prettypoly(poly1, :fill, close=true)
        randomhue()
        translate(5, 0)
        prettypoly(poly2, :fill, close=true)
        grestore()
    end
    @test finish() == true
end

fname = "polysplit-2.pdf"
poly_split_2(fname)
println("...finished test: output in $(fname)")


text:
module TestReconstructor
try
    using Test
catch
    using Base.Test
end
using Reconstructables: @add_kwonly, recon, constructor_of, statements

@test constructor_of(Vector) == Array
@test constructor_of([0]) == Array

struct A
    x
    @add_kwonly A(x) = new(x)
end
a1 = A(1)
a2 = recon(a1, x=2)
@test a2.x == 2

""" List of calls inferable by Julia 0.6. """
inferable_by_v06 = statements(quote
    A(1)
end)

""" List of calls *not* inferable by Julia 0.6 but (hopefully) in ≥ 0.7. """
non_inferable_by_v06 = statements(quote
    A(x=1)
    recon(a2, x=2)
end)

for (broken, expressions) in [(false, inferable_by_v06),
                              (true, non_inferable_by_v06)]
    for call in expressions
        # At the moment, it is broken even in 0.7:
        if broken  # && VERSION < v"0.7-"
            @eval @test_broken begin
                @inferred $call
                true
            end
        else
            @eval @test begin
                @inferred $call
                true
            end
        end
    end
end

end


text:
using Distributions

function dists(specs; N=50)
  X = Float64[]
  Y = Float64[]
  for (x, y, σ) in specs
    xd = rand(Normal(x, σ), N)
    yd = rand(Normal(y, σ), N)
    append!(X, xd)
    append!(Y, yd)
  end
  X, Y
end

N = 1000000
x, y = dists([(2,2,0.02), (2,-2,0.1), (-2,-2,0.5), (-2,2,1.0), (0,0,3.0)], N=N)

println("# of points: ", length(x))

using GR

setviewport(0.1, 0.95, 0.1, 0.95)
setwindow(-10, 10, -10, 10)
setcharheight(0.02)
axes2d(0.5, 0.5, -10, -10, 4, 4, -0.005)
setcolormap(GR.COLORMAP_HOT)

@time shadepoints(x, y, xform=GR.XFORM_EQUALIZED)

updatews()


text:
using QuasiCopula, LinearAlgebra, GLM
using Random, Distributions, DataFrames, ToeplitzMatrices
using Test, BenchmarkTools

BLAS.set_num_threads(1)
Threads.nthreads()

p = 3   # number of fixed effects, including intercept
m = 1    # number of variance components
# true parameter values
Random.seed!(12345)
βtrue = rand(Uniform(-2, 2), p)
σ2true = [0.5]
ρtrue = [0.5]
τtrue = 100.0
σ2 = inv(τtrue)
σ = sqrt(σ2)

function get_V(ρ, n)
    vec = zeros(n)
    vec[1] = 1.0
    for i in 2:n
        vec[i] = ρ
    end
    V = ToeplitzMatrices.SymmetricToeplitz(vec)
    V
end

#simulation parameters
samplesize = 10000
ni = 25

V = get_V(ρtrue[1], ni)

# true Gamma
Γ = σ2true[1] * V

T = Float64
gcs = Vector{GaussianCopulaCSObs{T}}(undef, samplesize)

# for reproducibility I will simulate all the design matrices here
Random.seed!(12345)
X_samplesize = [randn(ni, p - 1) for i in 1:samplesize]

for i in 1:samplesize
    X = [ones(ni) X_samplesize[i]]
    μ = X * βtrue
    vecd = Vector{ContinuousUnivariateDistribution}(undef, ni)
    for i in 1:ni
        vecd[i] = Normal(μ[i], σ)
    end
    nonmixed_multivariate_dist = NonMixedMultivariateDistribution(vecd, Γ)
    # simuate single vector y
    y = Vector{Float64}(undef, ni)
    res = Vector{Float64}(undef, ni)
    rand(nonmixed_multivariate_dist, y, res)
    gcs[i] = GaussianCopulaCSObs(y, X)
end

gcm = GaussianCopulaCSModel(gcs)
# precompile
println("precompiling Gaussian CS fit")
gcm2 = deepcopy(gcm);
QuasiCopula.fit!(gcm2, IpoptSolver(print_level = 0, max_iter = 20));

loglikelihood!(gcm, true, true)
fittime = @elapsed QuasiCopula.fit!(gcm, IpoptSolver(print_level = 5, max_iter = 100, tol = 10^-8, limited_memory_max_history = 50, accept_after_max_steps = 4, hessian_approximation = "limited-memory"))
@show fittime
@show gcm.β
@show gcm.σ2
@show gcm.ρ
@show gcm.∇β
@show gcm.∇σ2
@show gcm.∇ρ

@show gcm.τ
@show gcm.∇τ

loglikelihood!(gcm, true, true)
vcov!(gcm)
@show QuasiCopula.confint(gcm)

mseβ, mseρ, mseσ2, mseτ = MSE(gcm, βtrue, τtrue[1], ρtrue, σ2true)
@show mseβ
@show mseτ
@show mseσ2
@show mseρ

@test mseβ < 0.01
@test mseσ2 < 1
@test mseρ < 0.01
@test mseτ < 0.01


text:
##### fill halo points based on topology
@inline  fill_halo_west!(c, H::Int64, N::Int64, ::Periodic)  = @views @. c[1:H, :, :] = c[N+1:N+H, :, :]
@inline fill_halo_south!(c, H::Int64, N::Int64, ::Periodic)  = @views @. c[:, 1:H, :] = c[:, N+1:N+H, :]
@inline   fill_halo_top!(c, H::Int64, N::Int64, ::Periodic)  = @views @. c[:, :, 1:H] = c[:, :, N+1:N+H]

@inline   fill_halo_east!(c, H::Int64, N::Int64, ::Periodic) = @views @. c[N+H+1:N+2H, :, :] = c[1+H:2H, :, :]
@inline  fill_halo_north!(c, H::Int64, N::Int64, ::Periodic) = @views @. c[:, N+H+1:N+2H, :] = c[:, 1+H:2H, :]
@inline fill_halo_bottom!(c, H::Int64, N::Int64, ::Periodic) = @views @. c[:, :, N+H+1:N+2H] = c[:, :, 1+H:2H]

@inline  fill_halo_west!(c, H::Int64, N::Int64, ::Bounded)   = @views @. c[1:H, :, :] = c[H+1:H+1, :, :]
@inline fill_halo_south!(c, H::Int64, N::Int64, ::Bounded)   = @views @. c[:, 1:H, :] = c[:, H+1:H+1, :]
@inline   fill_halo_top!(c, H::Int64, N::Int64, ::Bounded)   = @views @. c[:, :, 1:H] = c[:, :, H+1:H+1]

@inline   fill_halo_east!(c, H::Int64, N::Int64, ::Bounded)  = @views @. c[N+H+1:N+2H, :, :] = c[N+H:N+H, :, :]
@inline  fill_halo_north!(c, H::Int64, N::Int64, ::Bounded)  = @views @. c[:, N+H+1:N+2H, :] = c[:, N+H:N+H, :]
@inline fill_halo_bottom!(c, H::Int64, N::Int64, ::Bounded)  = @views @. c[:, :, N+H+1:N+2H] = c[:, :, N+H:N+H]

@inline   fill_halo_east_vel!(c, H::Int64, N::Int64, ::Bounded)  = @views @. c[N+H+2:N+2H, :, :] = c[N+H+1:N+H+1, :, :]
@inline  fill_halo_north_vel!(c, H::Int64, N::Int64, ::Bounded)  = @views @. c[:, N+H+2:N+2H, :] = c[:, N+H+1:N+H+1, :]
@inline fill_halo_bottom_vel!(c, H::Int64, N::Int64, ::Bounded)  = @views @. c[:, :, N+H+2:N+2H] = c[:, :, N+H+1:N+H+1]

@inline   fill_halo_east_Gc!(c, H::Int64, N::Int64, ::Periodic) = @views @. c[N+H+1:N+2H, :, :] = c[1+H:2H, :, :]
@inline  fill_halo_north_Gc!(c, H::Int64, N::Int64, ::Periodic) = @views @. c[:, N+H+1:N+2H, :] = c[:, 1+H:2H, :]
@inline fill_halo_bottom_Gc!(c, H::Int64, N::Int64, ::Periodic) = @views @. c[:, :, N+H+1:N+2H] = c[:, :, 1+H:2H]

@inline   fill_halo_east_Gc!(c, H::Int64, N::Int64, ::Bounded) = @views @. c[N+H+1:N+2H, :, :] = 0.0
@inline  fill_halo_north_Gc!(c, H::Int64, N::Int64, ::Bounded) = @views @. c[:, N+H+1:N+2H, :] = 0.0
@inline fill_halo_bottom_Gc!(c, H::Int64, N::Int64, ::Bounded) = @views @. c[:, :, N+H+1:N+2H] = 0.0

@inline function fill_halo_nut!(nuts::NamedTuple, g::AbstractGrid{TX, TY, TZ}) where {TX, TY, TZ}
    for nut in nuts
          fill_halo_west!(nut.data, g.Hx, g.Nx, TX())
          fill_halo_east!(nut.data, g.Hx, g.Nx, TX())
         fill_halo_south!(nut.data, g.Hy, g.Ny, TY())
         fill_halo_north!(nut.data, g.Hy, g.Ny, TY())
           fill_halo_top!(nut.data, g.Hz, g.Nz, TZ())
        fill_halo_bottom!(nut.data, g.Hz, g.Nz, TZ())
    end
    return nothing
end

@inline function fill_halo_Gcs!(nuts::NamedTuple, g::AbstractGrid{TX, TY, TZ}) where {TX, TY, TZ}
    for nut in nuts
          fill_halo_east_Gc!(nut.data, g.Hx, g.Nx, TX())
         fill_halo_north_Gc!(nut.data, g.Hy, g.Ny, TY())
        fill_halo_bottom_Gc!(nut.data, g.Hz, g.Nz, TZ())
    end
    return nothing
end

@inline function fill_halo_u!(u, g::AbstractGrid{TX, TY, TZ}) where {TX, TY, TZ}
    if isa(TX(), Bounded)
        fill_halo_east_vel!(u, g.Hx, g.Nx, TX())
    else
        fill_halo_east!(u, g.Hx, g.Nx, TX())
    end
      fill_halo_west!(u, g.Hx, g.Nx, TX())
     fill_halo_south!(u, g.Hy, g.Ny, TY())
     fill_halo_north!(u, g.Hy, g.Ny, TY())
       fill_halo_top!(u, g.Hz, g.Nz, TZ())
    fill_halo_bottom!(u, g.Hz, g.Nz, TZ())
end

@inline function fill_halo_v!(v, g::AbstractGrid{TX, TY, TZ}) where {TX, TY, TZ}
    if isa(TY(), Bounded)
        fill_halo_north_vel!(v, g.Hy, g.Ny, TY())
    else
        fill_halo_north!(v, g.Hy, g.Ny, TY())
    end
      fill_halo_west!(v, g.Hx, g.Nx, TX())
      fill_halo_east!(v, g.Hx, g.Nx, TX())
     fill_halo_south!(v, g.Hy, g.Ny, TY())
       fill_halo_top!(v, g.Hz, g.Nz, TZ())
    fill_halo_bottom!(v, g.Hz, g.Nz, TZ())
end

@inline function fill_halo_w!(w, g::AbstractGrid{TX, TY, TZ}) where {TX, TY, TZ}
    if isa(TZ(), Bounded)
        fill_halo_bottom_vel!(w, g.Hz, g.Nz, TZ())
    else
        fill_halo_bottom!(w, g.Hz, g.Nz, TZ())
    end
     fill_halo_west!(w, g.Hx, g.Nx, TX())
     fill_halo_east!(w, g.Hx, g.Nx, TX())
    fill_halo_south!(w, g.Hy, g.Ny, TY())
    fill_halo_north!(w, g.Hy, g.Ny, TY())
      fill_halo_top!(w, g.Hz, g.Nz, TZ())
end

@inline function fill_halo_vel!(vels::NamedTuple, g::AbstractGrid)
    fill_halo_u!(vels.u.data, g)
    fill_halo_v!(vels.v.data, g)
    fill_halo_w!(vels.w.data, g)
    return nothing
end


text:
using Learn
using Test

@testset "Learn.jl" begin
    # Write your tests here.
    @testset "Add" begin
        @test true
        @test Learn.add(3,5) == 8
        @test Learn.add(1.1, 3.0) == 4.1
        @test Learn.add(1+5im, 4-2im) == 5+3im
        @test Learn.add(Vector([1,2,3]), Vector([4,5,6])) == Vector([5,7,9])
        @test Learn.add("asdf ", "fdsa") == "asdf fdsa"
    end
    @testset "Sum" begin
        @test Learn.∑(1,2,3,4,5) == 15
        @test Learn.∑([1,2,3,4,5]) == 15
    end

    include("testDayCounts.jl")
    include("testHolidays.jl")
end

max_stars_repo_name: binnisb/JuliaLearn

max_issues_repo_name: binnisb/JuliaLearn

max_forks_repo_name: binnisb/JuliaLearn


text:
type PlainText
  borders::Border
  content::Matrix{PlainCell}
  widths::Vector{Int}
  heights::Vector{Int}
end

immutable CellBorder
  left::Bool
  right::Bool
  top::Bool
  bottom::Bool
end

immutable PlainCell
  content::String
  colspan::Int
  rowspan::Int
end

PlainCell() = PlainCell("")
height(c::PlainCell) = 1
width(c::PlainCell) = length(c.content)


text:
const fit = [
    "fit_main",
    "em",
    "fit_exact",
    "fit_interface"
]

for dname in fit
    include(joinpath("fit", "$(dname).jl"))
end

max_stars_repo_name: sparktseung/JuliaLRMoE

max_issues_repo_name: sparktseung/JuliaLRMoE

max_forks_repo_name: sparktseung/JuliaLRMoE


text:
### A Pluto.jl notebook ###
# v0.19.3

using Markdown
using InteractiveUtils

# ╔═╡ ed0f13cc-4f7d-476b-a434-d14313d88eea
using Markdown: withtag

# ╔═╡ 6043d6c5-54e4-40c1-a8a5-aec3ad7e1aa0
md"# asdfsf"

# ╔═╡ f11f9ead-bbe9-4fa5-b99c-408cc4a69a7e
md"""

## Hello!!

# level 1

## level 2

### level 3

### level 3 again

asdf

# level 1 again

#### level 4

##### level 5

## back to 2
"""

# ╔═╡ d6940210-4f9b-47b5-af74-e53700a42417
const toc_js = toc -> """
const getParentCell = el => el.closest("pluto-cell")

const getHeaders = () => {
	const depth = Math.max(1, Math.min(6, $(toc.depth))) // should be in range 1:6
	const range = Array.from({length: depth}, (x, i) => i+1) // [1, ..., depth]

	const selector = range.map(i => `pluto-notebook pluto-cell h\${i}`).join(",")
	return Array.from(document.querySelectorAll(selector))
}

const indent = $(repr(toc.indent))
const aside = $(repr(toc.aside))

const clickHandler = (event) => {
	const path = (event.path || event.composedPath())
	const toc = path.find(elem => elem?.classList?.contains?.("toc-toggle"))
	if (toc) {
		event.stopImmediatePropagation()
		toc.closest(".plutoui-toc").classList.toggle("hide")
	}
}

document.addEventListener("click", clickHandler)

const render = (el) => html`\${el.map(h => {
	const parent_cell = getParentCell(h)

	const a = html`\${h.innerText}`
	/* a.onmouseover=()=>{
		parent_cell.firstElementChild.classList.add(
			'highlight-pluto-cell-shoulder'
		)
	}
	a.onmouseout=() => {
		parent_cell.firstElementChild.classList.remove(
			'highlight-pluto-cell-shoulder'
		)
	} */
	a.onclick=(e) => {
		e.preventDefault();
		h.scrollIntoView({
			behavior: 'smooth',
			block: 'start'
		})
	}

	return html`<div class="toc-row">\${a}</div>`
})}`

const tocNode = html`

     📖
     📕
	$(toc.title)

`

tocNode.classList.toggle("aside", aside)
tocNode.classList.toggle("indent", indent)

const updateCallback = () => {
	tocNode.querySelector("section").replaceWith(
		html`<section>\${render(getHeaders())}</section>`
	)
}
updateCallback()
setTimeout(updateCallback, 100)
setTimeout(updateCallback, 1000)
setTimeout(updateCallback, 5000)

const notebook = document.querySelector("pluto-notebook")

// We have a mutationobserver for each cell:
const observers = {
	current: [],
}

const createCellObservers = () => {
	observers.current.forEach((o) => o.disconnect())
	observers.current = Array.from(notebook.querySelectorAll("pluto-cell")).map(el => {
		const o = new MutationObserver(updateCallback)
		o.observe(el, {attributeFilter: ["class"]})
		return o
	})
}
createCellObservers()

// And one for the notebook's child list, which updates our cell observers:
const notebookObserver = new MutationObserver(() => {
	updateCallback()
	createCellObservers()
})
notebookObserver.observe(notebook, {childList: true})

// And finally, an observer for the document.body classList, to make sure that the toc also works when if is loaded during notebook initialization
const bodyClassObserver = new MutationObserver(updateCallback)
bodyClassObserver.observe(document.body, {attributeFilter: ["class"]})

// Hide/show the ToC when the screen gets small
let m = matchMedia("(max-width: 1000px)")
let match_listener = () =>
	tocNode.classList.toggle("hide", m.matches)
match_listener()
m.addListener(match_listener)

invalidation.then(() => {
	notebookObserver.disconnect()
	bodyClassObserver.disconnect()
	observers.current.forEach((o) => o.disconnect())
	document.removeEventListener("click", clickHandler)
	m.removeListener(match_listener)
})

return tocNode
"""

# ╔═╡ 731a4662-c329-42a2-ae71-7954140bb290
const toc_css = """
@media not print {

.plutoui-toc {
	--main-bg-color: unset;
	--pluto-output-color: hsl(0, 0%, 36%);
	--pluto-output-h-color: hsl(0, 0%, 21%);
}

@media (prefers-color-scheme: dark) {
	.plutoui-toc {
		--main-bg-color: hsl(0deg 0% 21%);
		--pluto-output-color: hsl(0, 0%, 90%);
		--pluto-output-h-color: hsl(0, 0%, 97%);
	}
}

.plutoui-toc.aside {
	color: var(--pluto-output-color);
	position:fixed;
	right: 1rem;
	top: 5rem;
	width: min(80vw, 300px);
	padding: 10px;
	border: 3px solid rgba(0, 0, 0, 0.15);
	border-radius: 10px;
	box-shadow: 0 0 11px 0px #00000010;
	/* That is, viewport minus top minus Live Docs */
	max-height: calc(100vh - 5rem - 56px);
	overflow: auto;
	z-index: 40;
	background-color: var(--main-bg-color);
	transition: transform 300ms cubic-bezier(0.18, 0.89, 0.45, 1.12);
}

.plutoui-toc.aside.hide {
	transform: translateX(calc(100% - 28px));
}
.plutoui-toc.aside.hide section {
	display: none;
}
.plutoui-toc.aside.hide header {
	margin-bottom: 0em;
	padding-bottom: 0em;
	border-bottom: none;
}
}  /* End of Media print query */
.plutoui-toc.aside.hide .open-toc,
.plutoui-toc.aside:not(.hide) .closed-toc,
.plutoui-toc:not(.aside) .closed-toc {
	display: none;
}

@media (prefers-reduced-motion) {
  .plutoui-toc.aside {
    transition-duration: 0s;
  }
}

.toc-toggle {
	cursor: pointer;
	padding: 1em;
	margin: -1em;
    margin-right: -0.7em;
}

.plutoui-toc header {
	display: block;
	font-size: 1.5em;
	margin-top: -0.1em;
	margin-bottom: 0.4em;
	padding-bottom: 0.4em;
	margin-left: 0;
	margin-right: 0;
	font-weight: bold;
	border-bottom: 2px solid rgba(0, 0, 0, 0.15);
}

.plutoui-toc section .toc-row {
	white-space: nowrap;
	overflow: hidden;
	text-overflow: ellipsis;
	padding-bottom: 2px;
}

.highlight-pluto-cell-shoulder {
	background: rgba(0, 0, 0, 0.05);
	background-clip: padding-box;
}

.plutoui-toc section a {
	text-decoration: none;
	font-weight: normal;
	color: var(--pluto-output-color);
}
.plutoui-toc section a:hover {
	color: var(--pluto-output-h-color);
}

.plutoui-toc.indent section a.H1 {
	font-weight: 700;
	line-height: 1em;
}

.plutoui-toc.indent section a.H1 {
	padding-left: 0px;
}
.plutoui-toc.indent section a.H2 {
	padding-left: 10px;
}
.plutoui-toc.indent section a.H3 {
	padding-left: 20px;
}
.plutoui-toc.indent section a.H4 {
	padding-left: 30px;
}
.plutoui-toc.indent section a.H5 {
	padding-left: 40px;
}
.plutoui-toc.indent section a.H6 {
	padding-left: 50px;
}
"""

# ╔═╡ 434cc67b-a1e8-4804-b7ba-f47d0f879046
begin
	"""Generate Table of Contents using Markdown cells. Headers h1-h6 are used.

	# Keyword arguments:
	`title` header to this element, defaults to "Table of Contents"

	`indent` flag indicating whether to vertically align elements by hierarchy

	`depth` value to limit the header elements, should be in range 1 to 6 (default = 3)

	`aside` fix the element to right of page, defaults to true

	# Examples:

	```julia
	TableOfContents()

	TableOfContents(title="Experiments 🔬")

	TableOfContents(title="📚 Table of Contents", indent=true, depth=4, aside=true)
	```
	"""
	Base.@kwdef struct TableOfContents
		title::AbstractString="Table of Contents"
		indent::Bool=true
		depth::Integer=3
		aside::Bool=true
	end
	function Base.show(io::IO, ::MIME"text/html", toc::TableOfContents)
		withtag(io, :script) do
			print(io, toc_js(toc))
		end
		withtag(io, :style) do
			print(io, toc_css)
		end
	end
end

# ╔═╡ 98cd39ae-a93c-40fe-a5d1-0883e1542e22
export TableOfContents

# ╔═╡ fdf8750b-653e-4f23-8f8f-9e2ef4e24e75
TableOfContents()

# ╔═╡ 7c32fd56-6cc5-420b-945b-53446833a125
TableOfContents(; aside = false)

# ╔═╡ 06ac2f13-e1f7-477a-9b3c-4d8545b777d9

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

[[Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"
"""

# ╔═╡ Cell order:
# ╠═98cd39ae-a93c-40fe-a5d1-0883e1542e22
# ╠═ed0f13cc-4f7d-476b-a434-d14313d88eea
# ╠═fdf8750b-653e-4f23-8f8f-9e2ef4e24e75
# ╠═7c32fd56-6cc5-420b-945b-53446833a125
# ╟─6043d6c5-54e4-40c1-a8a5-aec3ad7e1aa0
# ╟─f11f9ead-bbe9-4fa5-b99c-408cc4a69a7e
# ╠═434cc67b-a1e8-4804-b7ba-f47d0f879046
# ╠═d6940210-4f9b-47b5-af74-e53700a42417
# ╠═731a4662-c329-42a2-ae71-7954140bb290
# ╠═06ac2f13-e1f7-477a-9b3c-4d8545b777d9
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


text:
## Allows for lazy calculation (when requested)
## Also will have observer mixin in here

mutable struct LazyMixin
  calculated::Bool
  frozen::Bool
  observe::ObserverMixin

  LazyMixin() = new(false, false, ObserverMixin())
end

is_calculated(lazy::LazyObject) = lazy.lazyMixin.calculated
is_frozen(lazy::LazyObject) = lazy.lazyMixin.frozen

calculated!(lazy::LazyObject, setting::Bool = true) = lazy.lazyMixin.calculated = setting

function calculate!(lazy::LazyObject)
  if !is_calculated(lazy)
    calculated!(lazy)
    perform_calculations!(lazy)
  end

  return lazy
end

function recalculate!(lazy::LazyObject)
  calculated!(lazy, false)
  calculate!(lazy)

  return lazy
end

get_observers(lazy::LazyObject) = lazy.lazyMixin.observe.observers


text:
# Worst case 15 minute timeout. If the compute environment has just scaled it will wait
# 8 minutes before scaling again. Spot instance requests can take some time to be fufilled
# but are usually instant and instances take around 4 minutes before they are ready.
const TIMEOUT = Minute(15)

# Scrapes the log output to determine the worker job IDs as stated by the manager
function scrape_worker_job_ids(output::AbstractString)
    m = match(BATCH_SPAWN_REGEX, output)

    if m !== nothing
        worker_job = m[:id]

        if m[:n] !== nothing
            num_workers = parse(Int, m[:n])
            return String["$worker_job:$i" for i in 0:(num_workers - 1)]
        else
            return String["$worker_job"]
        end
    else
        return String[]
    end
end

function run_batch_job(
    image_name::AbstractString,
    num_workers::Integer;
    timeout::Period=TIMEOUT,
    should_fail::Bool=false,
)
    # TODO: Use AWS Batch job parameters to avoid re-registering the job

    timeout_secs = Dates.value(Second(timeout))

    # Will be running the HEAD revision of the code remotely
    # Note: Pkg.checkout doesn't work on untracked branches / SHAs with Julia 0.5.1
    code = """
        using AWSClusterManagers: AWSClusterManagers, AWSBatchManager
        using Dates: Second
        using Distributed
        using Memento

        Memento.config!("debug"; fmt="{msg}")
        setlevel!(getlogger(AWSClusterManagers), "debug")

        addprocs(
            AWSBatchManager(
                $num_workers;
                queue="$(STACK["WorkerJobQueueArn"])",
                memory=512,
                timeout=Second($(timeout_secs - 15))
            )
        )
        println("NumProcs: ", nprocs())

        @everywhere using AWSClusterManagers: container_id
        for i in workers()
            println("Worker container \$i: ", remotecall_fetch(container_id, i))
            println("Worker job \$i: ", remotecall_fetch(() -> ENV["AWS_BATCH_JOB_ID"], i))
        end

        println("Manager Complete")
        """

    # Note: The manager can run out of memory with enough workers:
    # - 64 workers with a manager with 1024 MB of memory
    info(LOGGER, "Submitting AWS Batch job with $num_workers workers")
    job = run_batch(;
        name=STACK["JobName"] * "-n$num_workers",
        queue=STACK["ManagerJobQueueArn"],
        definition=STACK["JobDefinitionName"],
        image=image_name,
        role=STACK["JobRoleArn"],
        vcpus=1,
        memory=2048,
        cmd=Cmd(["julia", "-e", code]),
    )

    # If no compute environment resources are available it could take around
    # 5 minutes before the manager job is running
    info(LOGGER, "Waiting for AWS Batch manager job $(job.id) to run (~5 minutes)")
    start_time = time()
    @test wait(state -> state < AWSBatch.RUNNING, job; timeout=timeout_secs) == true
    info(LOGGER, "Manager spawning duration: $(time_str(time() - start_time))")

    # Once the manager job is running it will spawn additional AWS Batch jobs as
    # the workers.
    #
    # Since compute environments only scale every 5 minutes we will definitely have
    # to wait if we scaled up for the mananager job. To reduce this wait time make
    # sure you have one VCPU available for the manager to start right away.
    info(LOGGER, "Waiting for AWS Batch workers and manager job to complete (~5 minutes)")
    start_time = time()
    if should_fail
        @test wait(job, [AWSBatch.FAILED], [AWSBatch.SUCCEEDED]; timeout=timeout_secs) ==
              true
    else
        @test wait(job, [AWSBatch.SUCCEEDED]; timeout=timeout_secs) == true
    end
    info(LOGGER, "Worker spawning duration: $(time_str(time() - start_time))")

    # Remove the job definition as it is specific to a revision
    job_definition = JobDefinition(job)
    deregister(job_definition)

    # CloudWatch can take several seconds to ingest the log record so we'll wait until we
    # find the end-of-log message.
    # Note: Do not assume that the "Manager Complete" message will be the last thing written
    # to the log as busy worker may cause additional warnings messages.
    # https://github.com/JuliaCloud/AWSClusterManagers.jl/issues/10
    output = ""
    if status(job) == AWSBatch.SUCCEEDED
        log_wait_start = time()

        while true
            events = log_events(job)
            if events !== nothing &&
               !isempty(events) &&
               any(e -> e.message == "Manager Complete", events)
                output = join(
                    [string(event.timestamp, "  ", event.message) for event in events], '\n'
                )
                break
            elseif time() - log_wait_start > 60
                error("CloudWatch logs have not completed ingestion within 1 minute")
            end

            sleep(5)
        end
    end

    return job, output
end

@testset "AWSBatchManager (online)" begin
    # Note: Start with the largest number of workers so the remaining tests don't have
    # to wait for the cluster to scale up on subsequent tests.
    @testset "Num workers ($num_workers)" for num_workers in [10, 1, 0]
        job, output = run_batch_job(TEST_IMAGE, num_workers)

        m = match(r"(?<=NumProcs: )\d+", output)
        if m !== nothing
            num_procs = parse(Int, m.match)
        else
            error("The logs do not contain the `NumProcs` for job \"$(job.id)\".")
        end

        # Spawned are the AWS Batch job IDs reported upon job submission at launch
        # while reported is the self-reported job ID of each worker.
        spawned_jobs = scrape_worker_job_ids(output)
        reported_jobs = [
            m[1] for m in eachmatch(r"Worker job \d+: ([0-9a-f\-]+(?:\:\d+)?)", output)
        ]
        reported_containers = [
            m[1] for m in eachmatch(r"Worker container \d+: ([0-9a-f]*)", output)
        ]

        @test num_procs == num_workers + 1
        if num_workers > 0
            @test length(reported_jobs) == num_workers
            @test Set(reported_jobs) == Set(spawned_jobs)
        else
            # When we request no workers the manager job will be treated as the worker
            @test length(reported_jobs) == 1
            @test reported_jobs == [job.id]
        end

        # Ensure that the container IDs were found
        @test all(.!isempty.(reported_containers))

        # Determine the image name from an AWS Batch job ID.
        job_image_name(job_id::AbstractString) = job_image_name(BatchJob(job_id))
        job_image_name(job::BatchJob) = describe(job)["container"]["image"]

        @test TEST_IMAGE == job_image_name(job)  # Manager's image
        @test all(TEST_IMAGE .== job_image_name.(spawned_jobs))

        # Report some details about the job
        d = describe(job)
        created_at = Dates.unix2datetime(d["createdAt"] / 1000)
        started_at = Dates.unix2datetime(d["startedAt"] / 1000)
        stopped_at = Dates.unix2datetime(d["stoppedAt"] / 1000)

        # TODO: Unless I'm forgetting something just extracting the seconds from the
        # milliseconds is awkward
        launch_duration = Dates.value(started_at - created_at) / 1000
        run_duration = Dates.value(stopped_at - started_at) / 1000

        info(LOGGER, "Job launch duration: $(time_str(launch_duration))")
        info(LOGGER, "Job run duration:    $(time_str(run_duration))")
    end

    @testset "exceed worker limit" begin
        num_workers = typemax(Int64)
        job, output = run_batch_job(TEST_IMAGE, num_workers; should_fail=true)

        # Spawned are the AWS Batch job IDs reported upon job submission at launch
        # while reported is the self-reported job ID of each worker.
        spawned_jobs = scrape_worker_job_ids(output)

        @test match(r"(?<=NumProcs: )\d+", output) === nothing
        @test isempty(spawned_jobs)
    end
end


text:
using ESValues
using Base.Test
using Iterators
using StatsBase

include("iterators.jl")

# basic test
P = 5
X = zeros(P, 1)
x = ones(P,1)#randn(P, 1)
f = x->sum(x, 1)
fnull,φ,φVar = esvalues(x, f, X, nsamples=10)
@test fnull == 0
@test norm(φ .- x) < 1e-5
@test norm(φVar) < 1e-5
@test length(φ) == P
@test length(φVar) == P

# basic test rand
srand(1)
P = 5
X = zeros(P, 1)
x = randn(P, 1)
f = x->sum(x, 1)
fnull,φ,φVar = esvalues(x, f, X, nsamples=10)
@test fnull == 0
@test norm(φ .- x) < 1e-5
@test norm(φVar) < 1e-5
@test length(φ) == P
@test length(φVar) == P

# check computation of groups when nothing varies
X = ones(P, 4)
x = ones(P, 1)
f = x->sum(x, 1)
groups = [[1,2],[3],[4],[5]]
fnull,φg,φVar = esvalues(x, f, X, featureGroups=groups, nsamples=8)
@test length(φg) == 4

# check computation of groups when nothing varies
X = ones(P, 4)
x = ones(P, 1)
x[1,1] = 0
f = x->sum(x, 1)
groups = [[1,2],[3],[4],[5]]
fnull,φg,φVar = esvalues(x, f, X, featureGroups=groups, nsamples=8)
@test length(φg) == 4

# check computation with many features
P = 200
X = zeros(P, 4)
x = ones(P, 1)
f = x->sum(x, 1)
fnull,φ,φVar = esvalues(x, f, X, nsamples=400)
@test length(φ) == 200

# make sure things work when only two features vary
P = 5
x = ones(P, 1)
x[1:2,1] = 0
X = ones(P, 1)
fnull,φ,φVar = esvalues(x, f, X, nsamples=8000)
@test fnull == 5
@test norm(X + φ .- x) < 1e-5
@test norm(φVar) < 1e-5

# make sure things work when only two features vary with weights
P = 5
x = ones(P, 1)
x[1:2,1] = 0
X = ones(P, 2)
X[:,2] = 2
weights = [1., 2.]
fnull,φ,φVar = esvalues(x, f, X, nsamples=8000, weights=weights)
@test abs(fnull - (5+10+10)/3) < 1e-6
@test norm(φVar) < 1e-5

# X and x are identical
srand(1)
X = randn(P, 1)
fnull,φ,φVar = esvalues(x, f, X, nsamples=10)
@test fnull == sum(X)
@test norm(X + φ .- x) < 1e-5
@test norm(φVar) < 1e-5

# X and x are identical
srand(1)
x = zeros(P,1)
x[1,1] = 1
X = zeros(P,1)
fnull,φ,φVar = esvalues(x, f, X, nsamples=8)
@test fnull == sum(X)
@test norm(X + φ .- x) < 1e-5
@test norm(φVar) < 1e-5

# non-zero reference distribution
X = ones(P, 1)
fnull,φ,φVar = esvalues(x, f, X, nsamples=8000)
@test fnull == 5
@test norm(X + φ .- x) < 1e-5
@test norm(φVar) < 1e-5

X = randn(P, 1)
fnull,φ,φVar = esvalues(x, f, X, nsamples=10)
@test fnull == f(X)[1]
@test norm(X + φ .- x) < 1e-5
@test norm(φVar) < 1e-5

function rawShapley(x, f, X, ind, g=identity; weights=nothing, featureGroups=nothing)

    weights != nothing || (weights = ones(size(X)[2]))
    weights ./= sum(weights)
    featureGroups != nothing || (featureGroups = Array{Int64,1}[Int64[i] for i in 1:size(X)[1]])
    featureGroups = convert(Array{Array{Int64,1},1}, featureGroups)

    M = length(featureGroups)
    val = 0.0
    sumw = 0.0
    for s in subsets(setdiff(1:M, ind))
        S = length(s)
        w = factorial(S)*factorial(M - S - 1)/factorial(M)
        tmp = copy(X)
        for i in 1:size(X)[2]
            for j in s
                for k in featureGroups[j]
                    tmp[k,i] = x[k]
                end
            end
        end
        y1 = g(sum(vec(f(tmp)) .* weights))
        for i in 1:size(X)[2]
            for k in featureGroups[ind]
                tmp[k,i] = x[k]
            end
        end
        y2 = g(sum(vec(f(tmp)) .* weights))
        val += w*(y2-y1)
        sumw += w
    end
    @assert abs(sumw - 1.0) < 1e-6
    val
end

# check brute force computation of groups
X = randn(P, 4)
f = x->sum(x, 1)
groups = [[1,2],[3],[4],[5]]
@test abs(rawShapley(x, f, X, 1) + rawShapley(x, f, X, 2) - rawShapley(x, f, X, 1, featureGroups=groups)) < 1e-5

# check computation of groups
X = randn(P, 4)
f = x->sum(x, 1)
groups = [[1,2],[3],[4],[5]]
fnull,φ,φVar = esvalues(x, f, X, nsamples=10)
fnull,φg,φVar = esvalues(x, f, X, featureGroups=groups, nsamples=8)
@test abs(φ[1] + φ[2] - φg[1]) < 1e-5

# check computation of groups when nothing varies
X = ones(P, 4)
x = ones(P, 1)
f = x->sum(x, 1)
groups = [[1,2],[3],[4],[5]]
fnull,φ,φVar = esvalues(x, f, X, nsamples=10)
fnull,φg,φVar = esvalues(x, f, X, featureGroups=groups, nsamples=8)
@test abs(φ[1] + φ[2] - φg[1]) < 1e-5

# check against brute force computation
X = randn(P, 4)
f = x->sum(x, 1)
fnull,φ,φVar = esvalues(x, f, X, nsamples=10)
for i in 1:length(φ)
    @test abs(φ[i] - rawShapley(x, f, X, i)) < 1e-5
end

# non-linear function
f = x->sum(x, 1).^2
fnull,φ,φVar = esvalues(x, f, X, nsamples=10)
for i in 1:length(φ)
    @test abs(φ[i] - rawShapley(x, f, X, i)) < 1e-5
end

# non-linear function that interestingly is still possible to estimate with only 2P samples
f = x->sum(x.^2, 1).^2
fnull,φ,φVar = esvalues(x, f, X, nsamples=10)
for i in 1:length(φ)
    @test abs(φ[i] - rawShapley(x, f, X, i)) < 1e-5
end

# non-linear logistic function
f = x->logistic(sum(x, 1))
fnull,φ,φVar = esvalues(x, f, X, nsamples=10000)
for i in 1:length(φ)
    @test abs(φ[i] - rawShapley(x, f, X, i)) < 1e-5
end

# non-linear logistic function with groups
f = x->logistic(sum(x, 1))
groups = [[1,2],[3],[4],[5]]
fnull,φ,φVar = esvalues(x, f, X, featureGroups=groups, nsamples=10000)
φRaw = [rawShapley(x, f, X, i, featureGroups=groups) for i in 1:length(φ)]
for i in 1:length(φ)
    @test abs(φ[i] - φRaw[i]) < 1e-5
end

# non-linear logistic function with groups
f = x->logistic(sum(x, 1))
groups = [[1,2],[3],[4],[5]]
weights = randn(4)
fnull,φ,φVar = esvalues(x, f, X, featureGroups=groups, weights=weights, nsamples=10000)
φRaw = [rawShapley(x, f, X, i, featureGroups=groups, weights=weights) for i in 1:length(φ)]
for i in 1:length(φ)
    @test abs(φ[i] - φRaw[i]) < 1e-5
end

# test many totally arbitrary functions
function gen_model(M)
    model = Dict()
    for k in subsets(collect(1:M))
        model[k] = randn()
    end
    model
end
P = 10
X = zeros(P, 2)
x = ones(P,1)
for i in 1:10
    model = gen_model(P)
    f = x->[model[find(x[:,i])] for i in 1:size(x)[2]]
    fnull,φ,φVar = esvalues(x, f, X, nsamples=1000000)
    phiRaw = [rawShapley(x, f, X, j) for j in 1:P]
    @test norm(φ .- [rawShapley(x, f, X, j) for j in 1:P]) < 1e-6
end

# non-linear logistic function with logit link
X = randn(P, 1)
x = randn(P)
f = x->logistic(sum(x, 1))
fnull,φ,φVar = esvalues(x, f, X, logit, nsamples=1000000)
for i in 1:length(φ)
    sv = rawShapley(x, f, X, i, logit)
    @test abs(φ[i] - rawShapley(x, f, X, i, logit)) < 1e-5
end
@test sum(abs(φVar)) < 1e-12 # we have exhausted the sample space so there should be no uncertainty

# non-linear logistic function with logit link and random background
P = 2
X = randn(P, 10)
x = randn(P)
f = x->logistic(sum(x, 1))
fnull,φ,φVar = esvalues(x, f, X, logit, nsamples=1000000)
φRaw = [rawShapley(x, f, X, i, logit) for i in 1:length(φ)]
@test norm(φ - φRaw) < 1e-5

# test many totally arbitrary functions with logit link
P = 10
X = zeros(P, 2)
x = ones(P,1)
for i in 1:10
    model = gen_model(P)
    f = x->[logistic(model[find(x[:,i])]) for i in 1:size(x)[2]]
    fnull,φ,φVar = esvalues(x, f, X, logit, nsamples=1000000)
    phiRaw = [rawShapley(x, f, X, j, logit) for j in 1:P]
    @test norm(φ .- phiRaw) < 1e-6
end

# test arbitrary functions with logit link and feature groups
P = 10
X = zeros(P, 2)
x = ones(P,1)
groups = [[1,2],[3],[4],[5,6],[7,8,9,10]]
for i in 1:3
    model = gen_model(P)
    f = x->[logistic(model[find(x[:,i])]) for i in 1:size(x)[2]]
    fnull,φ,φVar = esvalues(x, f, X, logit, nsamples=1000000, featureGroups=groups)
    phiRaw = [rawShapley(x, f, X, j, logit, featureGroups=groups) for j in 1:length(groups)]
    @test norm(φ .- phiRaw) < 1e-6
    @test abs(logistic(logit(fnull)+sum(φ)) - f(x)[1]) < 1e-6
end