text:
import Pkg; Pkg.activate(".")

using Plots
using StatsPlots
import CSV
using Distributions
using Statistics
using StatsBase
using Random

# include functions
include("common_functions.jl")

Plots.scalefontsizes(1.3)
fonts=font("Arial",7)

# get the data and filter for predation only
d = CSV.read(joinpath("data", "network_data.dat"))
d = d[d.predation .> 0 , :]

# color palette for models
pal = (
    lssl=RGB(230/255,159/255,0/255),
    cc=RGB(86/255,190/255,233/255),
    pl=RGB(0/255,158/255,115/255),
    fl=RGB(204/255,121/255,167/255)
    )

# posterior samples for the flexible links model
betab_posterior = CSV.read(joinpath("data", "posterior_distributions", "beta_binomial_posterior.csv"))

# posterior samples from previous models
lssl_posterior = CSV.read(joinpath("data", "posterior_distributions", "lssl.csv"))
const_posterior = CSV.read(joinpath("data", "posterior_distributions", "const_posterior.csv"))
powerlaw_posterior = CSV.read(joinpath("data", "posterior_distributions", "powerlaw_posterior.csv"))

# number of species
S = 3:750
mms = S.-1 # min links
Ms = S.^2 # max links
ms = mms ./ Ms  # min connectance
msl = (S .- 1) ./ (S) # min link density
total_flex = S .^ 2 .-S .+1

# map estimates
mu_map = median(betab_posterior[:mu])
phi_map = exp(median(betab_posterior[:phi]))

α = mu_map*phi_map
β = (1.0-mu_map)*phi_map

beta_map = Beta(α, β)
betabin_map = BetaBinomial.(total_flex, α, β)

# counterfactuals
betab_cf_links = betab_posterior[r"counterfactual_links"]
betab_cf_links = betab_cf_links[:, S]

lssl_cf_links = lssl_posterior[r"counterfactual_links"]
lssl_cf_links = lssl_cf_links[:, S]

const_cf_links = const_posterior[r"counterfactual_links"]
const_cf_links = const_cf_links[:, S]

powerlaw_cf_links = powerlaw_posterior[r"counterfactual_links"]
powerlaw_cf_links = powerlaw_cf_links[:, S]

# Fig -- Parameters can be estimated by Maximum Likelihood

# generate posterior draws of the Beta distribution
Random.seed!(1234)
index = rand(1:size(betab_posterior,2), 20) # 20 posterior samples
mu_rdm = betab_posterior[index, :mu]
phi_rdm = exp.(betab_posterior[index, :phi])

betab_random = Beta.(mu_rdm .* phi_rdm, (1 .- mu_rdm) .* phi_rdm)

# MLE fit
pex = (d.links .- (d.nodes .- 1)) ./  (d.nodes.^2 .- (d.nodes .- 1))

p = fit(Beta, pex)

# calculate for text
phi_MLE = p.α + p.β
mu_MLE = p.α / phi_MLE

density(pex, c=:lightgrey, fill=(:lightgrey, 0, 0.5), dpi=1000, size=(800,500),  margin=5Plots.mm, lab="Empirical data",
    foreground_color_legend=nothing, background_color_legend=:white, framestyle=:box,
    guidefont=fonts, xtickfont=fonts, ytickfont=fonts, legendfont=fonts)
plot!(betab_random[1], c=pal.fl, linewidth=1, alpha=0.3, lab="Posterior samples")
for i in 1:length(index)
    plot!(betab_random[i], c=pal.fl, linewidth=1, alpha=0.3, lab="")
end
density!(rand(p, 100_000), c=:black, ls=:dash, linewidth=2, lab="MLE fit")
yaxis!((0, 9.5), "Density")
xaxis!((0, 0.5), "p")
savefig(joinpath("figures", "beta_fit"))

# Fig -- The fexible link model ts better and makes a plausible range of predictions

# To make the maximum point apparent in plot
d.nodesmax = maximum(d.nodes)
d.linksmax = maximum(d.links)

# Function to plot the quantiles of the counterfactuals links of each model
# we use log_zeros function to plot the y-axis in log (to account for log(0))
function plot_links_quantile(model; title="", xlabel="", ylabel="", linecolor="")
    quant_015 = neg_to_zeros.(quantile.(eachcol(model), 0.015))
    quant_110 = neg_to_zeros.(quantile.(eachcol(model), 0.11))
    quant_890 = neg_to_zeros.(quantile.(eachcol(model), 0.89))
    quant_985 = neg_to_zeros.(quantile.(eachcol(model), 0.985))
    quant_500 = neg_to_zeros.(quantile.(eachcol(model), 0.5))

    plot(S, quant_985, fill=quant_015, color=:grey, alpha=0.15, label="",
        title=title, title_location=:left,
        xlabel=xlabel, ylabel=ylabel, framestyle=:box, dpi=1000,
        guidefont=fonts, xtickfont=fonts, ytickfont=fonts, titlefont=font("Arial",9)) # 97% PI
    plot!(S, quant_890, fill=quant_110, color=:grey, alpha=0.15, label="") # 89% PI
    scatter!(d[:nodes], d[:links], c=:grey, alpha=0.6, msw=0, markersize=5, label="") # Empirical links
    plot!(S, quant_500, linecolor=linecolor, linewidth=3, label="") # Median link number
    scatter!(d.nodesmax, d.linksmax, c=:grey, alpha=0.6, msw=0, markersize=5, label="") # Maximum empirical link
    plot!(S, mms, linecolor=:black, lw=1, label="") # Minimum number of links
    plot!(S, Ms, linecolor=:black, lw=2, label="") # Maximum number of links
    xaxis!(:log, xlabel=xlabel, xlims=(minimum(S), maximum(S)))
    yaxis!(:log, ylims = (1,100000), ylabel=ylabel)
end

plot_lssl = plot_links_quantile(lssl_cf_links, title="A. LSSL",
    ylabel="Number of links", linecolor=pal.lssl)
plot_const = plot_links_quantile(const_cf_links, title="B. Constant connectance",
    linecolor=pal.cc)
plot_powerlaw = plot_links_quantile(powerlaw_cf_links, title="C. Power law",
    xlabel="Species richness", ylabel="Number of links", linecolor=pal.pl)
plot_betab = plot_links_quantile(betab_cf_links, title="D. Flexible links",
    xlabel="Species richness", linecolor=pal.fl)

plot(plot_lssl, plot_const, plot_powerlaw, plot_betab, layout=(2,2), size=(700,700),  margin=5Plots.mm, dpi=200)
savefig(joinpath("figures", "models_links"))
savefig(joinpath("figures", "submission", "fig1.pdf"))

# Fig -- The shifted beta-binomial distribution can be approximated by a normal distribution

# A BetaBinomial predictions from map values

bb_rand = rand.(betabin_map, 5000)

beta_89 = quantile.(bb_rand, 0.89) .+ S .- 1
beta_11 = quantile.(bb_rand, 0.11) .+ S .- 1
beta_98 = quantile.(bb_rand, 0.985) .+ S .- 1
beta_02 = quantile.(bb_rand, 0.015) .+ S .- 1
beta_50 = quantile.(bb_rand, 0.5)  .+ S .- 1

links_beta_map = plot(S, beta_98, fill=beta_02,label="", color=:grey, alpha=0.15,
    title="A. Flexible links (MAP)", title_location=:left, framestyle=:box,
    guidefont=fonts, xtickfont=fonts, ytickfont=fonts, titlefont=font("Arial",9))
plot!(S, beta_89, fill=beta_11,label="", color=:grey, alpha=0.15)
plot!(S, beta_50, color=pal.fl, label="", linewidth=2)
#scatter!(d[:nodes], d[:links], c=:grey, msw=0, markersize=5, label="") # Empirical links
plot!(S, mms, linecolor=:black, label="", lw=1) # Minimum number of links
plot!(S, Ms, linecolor=:black, label="", lw=2) # Maximum number of links
xaxis!(:log, "Species richness", label="", xlim=(minimum(S), maximum(S)))
yaxis!(:log, "Number of links", ylims=(1,100000))

# B Normal approximation of BetaBinomial
means = (Ms .- mms) .* mu_map .+ S .- 1
vars = (Ms .- mms) .* mu_map .* (1 .- mu_map) .* (1 .+ S .* (S .- 1) .* (1 / (1 + phi_map)))

approxs = Normal.(means, sqrt.(vars))
tnormal = truncated.(approxs, 0.01, Inf)

tnormal_89 = quantile.(tnormal, 0.89)
tnormal_11 = quantile.(tnormal, 0.11)
tnormal_98 = quantile.(tnormal, 0.985)
tnormal_02 = quantile.(tnormal, 0.015)
tnormal_50 = quantile.(tnormal, 0.5)

links_normal = plot(S, tnormal_98, fill=tnormal_02,label="", color=:grey, alpha=0.15,
    title="B. Normal approximation", title_location=:left, framestyle=:box,
    guidefont=fonts, xtickfont=fonts, ytickfont=fonts, titlefont=font("Arial",9))
plot!(S, tnormal_89, fill=tnormal_11,label="", color=:grey, alpha=0.15,)
plot!(S, tnormal_50, color=pal.fl, label="", lw=2)
plot!(S, mms, linecolor=:black, label="", lw=1) # Minimum number of links
plot!(S, Ms, linecolor=:black, label="", lw=2) # Maximum number of links
#scatter!(d[:nodes], d[:links], label="", color=:grey) # Empirical links
xaxis!(:log, "Species richness", label="", xlim=(minimum(S), maximum(S)))
yaxis!(:log, "Number of links", ylims=(1,100000))

plot(links_beta_map, links_normal, layout=(1,2), size=(700,350), margin=5Plots.mm, dpi=1000)
savefig(joinpath("figures", "betabinmap_normal_links"))
savefig(joinpath("figures", "submission", "fig5.pdf"))

# Fig - Connectance and linkage density can be derived from a model for links

# A - connectance - species

## scale the "expected" distribution according to the minimum value:
bquant = LocationScale.(ms, 1 .- ms, beta_map)

# Quantiles to plot
beta015 = quantile.(bquant, 0.015)
beta985 = quantile.(bquant, 0.985)
beta11 = quantile.(bquant, 0.110)
beta89 = quantile.(bquant, 0.890)

beta500 = quantile.(bquant, 0.5)

# Empirical connectance
co_emp = d[:links] ./ (d[:nodes] .^2)

# Connectance vs species
connectance_beta = plot(S, beta985, fillrange=beta015, color=:grey, alpha=0.15,
    label="", title="A", title_location=:left, framestyle=:box,
    guidefont=fonts, xtickfont=fonts, ytickfont=fonts, titlefont=font("Arial",9)) # 97% PI
plot!(S, beta11, fillrange=beta89, color=:grey, alpha=0.15, label="") # 89% PI
scatter!(d[:nodes], co_emp, c=:grey, alpha=0.5, msw=0, markersize=5, label="") # Empirical connectance
plot!(S, beta500, linecolor=pal.fl, linewidth=2, label="") # Median connectance
plot!(S, ms, label="", linecolor=:black, lw=1) # Minimum connectance
xaxis!(:log, "Species richness", xlims=(minimum(S),maximum(S)))
yaxis!("Connectance", (0, 0.5))

# B linkage density - species

## scale the "expected" distribution according to the mimum value:
bquant_LS = LocationScale.(msl, S .- msl, beta_map)

beta015_LS = quantile.(bquant_LS, 0.015)
beta985_LS = quantile.(bquant_LS, 0.985)
beta11_LS = quantile.(bquant_LS, 0.11)
beta89_LS = quantile.(bquant_LS, 0.89)
beta50_LS = quantile.(bquant_LS, 0.50)

avg_degree_beta = plot(S, beta985_LS, fill=beta015_LS, color=:grey, alpha=0.15, lab="", title="B",
    title_location=:left, framestyle=:box,
    guidefont=fonts, xtickfont=fonts, ytickfont=fonts, titlefont=font("Arial",9))
plot!(S, beta11_LS, fill=beta89_LS, colour=:grey, alpha=0.15, lab="",)
scatter!(d.nodes, d.links ./ d.nodes, c=:grey, alpha=0.5, msw=0, markersize=5, label="") # Empirical connectance
plot!(S, beta50_LS, linecolor=pal.fl, lw=2, lab="")
plot!(S, msl, linecolor=:black, lw=1, lab="")
plot!(S, S, linecolor=:black, lw=2, lab="")
xaxis!(:log, "Species richness", xlims=(minimum(S),maximum(S)))
yaxis!(:log, "Linkage density", ylims=(0.5,1000))

plot(connectance_beta, avg_degree_beta, label=(1,2), lab="", size=(700,350),  dpi=1000,  margin=5Plots.mm)
savefig(joinpath("figures", "connectance_linkdens"))
savefig(joinpath("figures", "submission", "fig3.pdf"))

# Fig -- Only the flexible link model makes realistic predictions for small communities

function realistic_links(model_cf)
    realistic = zeros(Float64, (1, length(S)))
    for (i,s) in enumerate(S)
        belowmin = length(findall(model_cf[:,i] .< (s - 1)))
        abovemax = length(findall(model_cf[:,i] .> (s^2)))
        realistic[i] = 1 - (belowmin + abovemax) / size(model_cf, 1)
    end
    return(vec(realistic))
end

realistic_betab =  realistic_links(betab_cf_links)
realistic_lssl = realistic_links(lssl_cf_links)
realistic_const = realistic_links(const_cf_links)
realistic_powerlaw = realistic_links(powerlaw_cf_links)

medianspecies = quantile(d[:nodes], 0.5)
species05 = quantile(d[:nodes], 0.05)
species95 = quantile(d[:nodes], 0.95)

plot([medianspecies], seriestype=:vline, color=:grey, ls=:dash, lab="", ylim=(0.4,1), frame=:box,  margin=5Plots.mm,
    guidefont=fonts, xtickfont=fonts, ytickfont=fonts, legendfont=fonts, titlefont=font("Arial",9))
plot!([species05, species95], [1.0, 1.0], fill=(0, :grey, 0.12), c=:transparent, lab="")
plot!([species05], seriestype=:vline, color=:grey, ls=:dot, lab="")
plot!([species95], seriestype=:vline, color=:grey, ls=:dot, lab="")
plot!(S, realistic_lssl, color=pal.lssl, linewidth=2, label="LSSL",
    legend=:bottomright, foreground_color_legend=nothing, background_color_legend=:white)
plot!(S, realistic_const, color=pal.cc, linewidth=2, label="Constant connectance")
plot!(S, realistic_powerlaw, color=pal.pl, linewidth=2, label="Power law")
plot!(S, realistic_betab, color=pal.fl, linewidth=2, label="Flexible links")
xaxis!(:log, "Species richness", xlims=(minimum(S),maximum(S)))
yaxis!("Proportion of realistic links", (0.3, 1.01))
savefig(joinpath("figures", "real_predict"))
savefig(joinpath("figures", "submission", "fig2.pdf"))

# Fig -- Many different Network-Area Relationships are supported by the data
A = 0.0001:0.02:1.2
k,z = 200.0, 0.27
AS = convert.(Int64, ceil.(k.*A.^z))

pl_nar_sar = plot(A, AS, color=:grey, lw=2, label="", frame=:box,
    legend=:topleft, foreground_color_legend=nothing, background_color_legend=:white,
    title_location=:left, guidefont=fonts, xtickfont=fonts, ytickfont=fonts, legendfont=fonts, titlefont=font("Arial",9),
    title="A. species area relationship")

xaxis!(pl_nar_sar, (0, 1), "Relative area")
yaxis!(pl_nar_sar, (0, 200), "Species richness")

# extract counterfactual links only -- necessary so that position matches S
bb_post = betab_posterior[r"counterfactual_links"]
pl_post = powerlaw_posterior[r"counterfactual_links"]
# extract columns matching species richness from AS
pl_mod = pl_post[:,AS]
fl_mod = bb_post[:,AS]

pl500 = neg_to_zeros.(quantile.(eachcol(pl_mod./AS'), 0.5))

fl015 = neg_to_zeros.(quantile.(eachcol(fl_mod./AS'), 0.015))
fl110 = neg_to_zeros.(quantile.(eachcol(fl_mod./AS'), 0.11))
fl890 = neg_to_zeros.(quantile.(eachcol(fl_mod./AS'), 0.89))
fl985 = neg_to_zeros.(quantile.(eachcol(fl_mod./AS'), 0.985))
fl500 = neg_to_zeros.(quantile.(eachcol(fl_mod./AS'), 0.5))

pl_nar_nar = plot(A, fl985, fillrange=fl015, color=:grey, alpha=0.15, label="", frame=:box, legend=:topleft, foreground_color_legend=nothing, background_color_legend=:white,
    title_location=:left, guidefont=fonts, xtickfont=fonts, ytickfont=fonts, legendfont=fonts, titlefont=font("Arial",9),
    title="B. network area relationship")
plot!(pl_nar_nar, A, fl890, fillrange=fl110, color=:grey, alpha=0.15, label="")
plot!(pl_nar_nar, A, fl500, linecolor=pal.fl, linewidth=2, label="Flexible links")
plot!(pl_nar_nar, A, pl500, linecolor=pal.pl, linewidth=1, ls=:dot, label="Power law")
xaxis!(pl_nar_nar, xlabel="Relative area", xlims=(0.0, 1.0))
yaxis!(pl_nar_nar, ylims = (1,50), ylabel="Linkage density")

plot(pl_nar_sar, pl_nar_nar, layout=(1,2), size=(700,350), margin=5Plots.mm, dpi=1000)
savefig(joinpath("figures", "nar"))
savefig(joinpath("figures", "submission", "fig6.pdf"))

# Fig -- Stability imposes a limit on network size
# posterior samples for the flexible links model
betab_posterior_bigger = CSV.read(joinpath("data", "posterior_distributions", "beta_binomial_posterior_bigger.csv"))

# stability after May / Allesina & Tang
S = 1:1500
fl_mod = betab_posterior_bigger[:,r"counterfactual_links"]

fl015 = neg_to_zeros.(quantile.(eachcol(fl_mod./S'), 0.015))
fl110 = neg_to_zeros.(quantile.(eachcol(fl_mod./S'), 0.11))
fl890 = neg_to_zeros.(quantile.(eachcol(fl_mod./S'), 0.89))
fl985 = neg_to_zeros.(quantile.(eachcol(fl_mod./S'), 0.985))
fl500 = neg_to_zeros.(quantile.(eachcol(fl_mod./S'), 0.5))

pl_may_max = plot(S, vec(1.0./sqrt.(fl985)), fillrange=vec(1.0./sqrt.(fl015)),
    color=:grey, alpha=0.15, label="", frame=:box, size=(400, 400), margin=5Plots.mm,
    dpi=1000, legend=:topleft, foreground_color_legend=nothing,
    title_location=:left, guidefont=fonts, xtickfont=fonts, ytickfont=fonts, legendfont=fonts, titlefont=font("Arial",9),
    title="A. maximal interaction diversity",
    background_color_legend=:white,
    xlabel="Species richness", ylabel="Maximal \\sigma")
plot!(pl_may_max, S, vec(1.0./sqrt.(fl890)), fillrange=vec(1.0./sqrt.(fl110)), color=:grey, alpha=0.15, label="")
plot!(pl_may_max, S, vec(1.0./sqrt.(fl500)), lw=2, c=pal.fl, lab="")
xaxis!(pl_may_max, :log, (3,1000))
plot!(pl_may_max, S, sqrt.(S)./sqrt.(S.-1), lab="", colour = :black, lw=1)
plot!(pl_may_max, S, 1.0 ./sqrt.(S), lab="", colour = :black, lw = 2)
yaxis!(pl_may_max, (0,1.25), label="Maximal interaction diversity")
# expression for the mean.
mu_map = median(betab_posterior_bigger[:mu])
plot!(pl_may_max, S, 1 ./ sqrt.(mu_map .* S .+ (1 - mu_map) .* (S .- 1) ./ S), lab = "", color=:black, lw = 1, linestyle=:dash)

pl_may_prop = plot(frame=:box, title_location=:left,
    guidefont=fonts, xtickfont=fonts, ytickfont=fonts, legendfont=fonts, titlefont=font("Arial",9),
    title="B. probability of a network being stable",
    foreground_color_legend=nothing,
    background_color_legend=:white)

σ = 0.0:0.01:1.25

for s in [10, 30, 100, 300, 1000]
    L = fl_mod[:,s]
    m = sqrt.(s.*(L./(s^2)))
    plot!(pl_may_prop, σ, sum((σ.*m').< 1, dims=2)./2000, lab="", lw=log(s)/log(1000).*3, c=:grey, fill=(0, :darkgrey, 0.05))
end

yaxis!(pl_may_prop, (0,1), "P(stability)")
xaxis!(pl_may_prop, "\\sigma", (minimum(σ), maximum(σ)))

plot(pl_may_max, pl_may_prop, layout=(1,2), size=(700,350), margin=5Plots.mm, dpi=1000)
savefig(joinpath("figures", "may"))
savefig(joinpath("figures", "submission", "fig7.pdf"))

# Fig -- Histogram of z-scores

# Compute expected number of links and variance
L_hat = (d.nodes .^2 .- d.nodes .+ 1) .* mu_map .+ d.nodes .- 1
sigma_2_L_hat = (d.nodes .^2 .- d.nodes .+ 1) .* mu_map .* (1 .- mu_map) .* (1 .+ d.nodes .* (d.nodes .- 1) ./ (phi_map .+ 1))

# Compute z-scores
z_scores = (d.links .- L_hat) ./ sqrt.(sigma_2_L_hat)

# Non abnormal z-scores (below 1.96 in absolute values)
z_scores_normals = z_scores[abs.(z_scores) .< 1.96]

# Abnormal z-scores (above 1.96 in absolute values)
z_scores_abnormals = z_scores[abs.(z_scores) .> 1.96]
z_scores_abnormals_pct = length(z_scores_abnormals) / size(d, 1)

# Plot histogram
histogram(z_scores_normals, lab="", fill=:lightgrey,
         xlims=(-6,6), xticks=-6:1:6,
         ylims=(0,30), yticks=0:5:30,
         frame=:box, dpi=1000, bins=30,
         guidefont=fonts, xtickfont=fonts, ytickfont=fonts, legendfont=fonts, titlefont=font("Arial",9),
         xlabel="z-score", ylabel="Frequency")
histogram!(z_scores_abnormals, fill=pal.fl, lab="", bins=18)
plot!([-2, 2], [30, 30], fill=(0, :grey, 0.12), c=:transparent, lab="")
plot!([-2], seriestype=:vline, color=:grey, ls=:dot, lab="")
plot!([2], seriestype=:vline, color=:grey, ls=:dot, lab="")
plot!([0], seriestype=:vline, color=:grey, ls=:dot, lab="")
savefig(joinpath("figures", "z-scores"))
savefig(joinpath("figures", "submission", "fig4.pdf"))


text:
using IEEEFloats
using Test

@test true

ext: jl


text:
# Get information from learners while they run

# For the Categorical task (Section 2.3.2)
mutable struct CallbackMultinomial
    state_history::Array{Int,1}
    true_history::Array{Array{Float64,1}, 1}
    estimated_history::Array{Array{Float64,1}, 1}
    switchflag::Array{Bool, 1}
    nposteriors::Array{Int,1}
end
CallbackMultinomial() = CallbackMultinomial(Array{Int}(undef, 1),
                    Array{Array{Float64,1}}(undef, 1),
                    Array{Array{Float64,1}}(undef, 1), Bool[],
                    Array{Int, 1}(undef, 1))
function callback!(p::CallbackMultinomial, learner, env)

    if !isassigned(p.true_history) # If very first step
        p.state_history[1] = copy(env.state)
        p.true_history[1] = deepcopy(env.trans_probs)
        p.estimated_history[1] = deepcopy(learner.Ps1)
    else
        push!(p.state_history, copy(env.state))
        push!(p.true_history, deepcopy(env.trans_probs))
        push!(p.estimated_history, deepcopy(learner.Ps1))
    end
    push!(p.switchflag, deepcopy(env.switchflag))

    # if in(:nposteriors, fieldnames(typeof(learner)))
    #     push!(p.posteriors, copy(length(f.posteriors)))
    # end
end

# For the Gaussian task (Section 2.3.1)
mutable struct CallbackGaussian
    state_history::Array{Float64,1}
    true_history::Array{Float64,1}
    estimated_history::Array{Float64,1}
    switchflag::Array{Bool, 1}
    nposteriors::Array{Int,1}
    Sgm_history::Array{Float64,1}
    var_history::Array{Float64,1}
    γ_history::Array{Float64,1}
    Ssh_history::Array{Float64, 1}
    nresamplings_history::Array{Float64, 1}
    Py_0_history::Array{Float64, 1}
end
CallbackGaussian() = CallbackGaussian(Array{Int}(undef, 1),
                    Array{Float64,1}(undef, 1),
                    Array{Float64,1}(undef, 1), Bool[],
                    Array{Int,1}(undef, 1),
                    Array{Float64,1}(undef, 1),
                    Array{Float64,1}(undef, 1),
                    Array{Float64,1}(undef, 1),
                    Array{Float64,1}(undef, 1),
                    Array{Float64,1}(undef, 1),
                    Array{Float64,1}(undef, 1))
function callback!(p::CallbackGaussian, learner, env)
    push!(p.state_history, copy(env.state))
    push!(p.true_history, deepcopy(env.mu_n))
    push!(p.estimated_history, deepcopy(learner.mu_n[1]))
    push!(p.switchflag, deepcopy(env.switchflag))

    # if in(:nposteriors, fieldnames(typeof(learner)))
    #     push!(p.posteriors, copy(length(f.posteriors)))
    # end

    if in(:Sgm_n, fieldnames(typeof(learner)))
        push!(p.Sgm_history, deepcopy(learner.Sgm_n[1]))
    end
    if in(:var_n, fieldnames(typeof(learner)))
        push!(p.var_history, deepcopy(learner.var_n[1]))
    end

    if in(:γ_n, fieldnames(typeof(learner)))
        push!(p.γ_history, deepcopy(learner.γ_n[1]))
    end

    if in(:Ssh_n, fieldnames(typeof(learner)))
        push!(p.Ssh_history, deepcopy(learner.Ssh_n[1]))
    end

    if in(:nresamplings, fieldnames(typeof(learner)))
        push!(p.nresamplings_history, deepcopy(learner.nresamplings[1]))
    end

    if in(:Py_0_n, fieldnames(typeof(learner)))
        push!(p.Py_0_history, deepcopy(learner.Py_0_n[1]))
    end
end


text:
mutable struct LSTDLearner{AF<:AbstractFloat} <: Learner

    update::LearningUpdate
    b::Vector{Vector{AF}}
    A_inv::Vector{Matrix{AF}}
    t::Int
    e::Vector{Vector{AF}}

    w::Vector{Matrix{AF}}
    w_real::Vector{Vector{AF}}

    feature_size::Int
    num_actions::Int
    num_demons::Int
end

function LSTDLearner{F}(update, eta, feature_size, num_actions, num_demons) where {F<:Number}
    fsna = feature_size*num_actions
    b = [zeros(F, fsna) for i in 1:num_demons]
    A_inv = [zeros(F, fsna, fsna) .+ F(eta)*I(fsna) for i in 1:num_demons]
    e = [zeros(F, fsna) for i in 1:num_demons]
    w = [zeros(F, num_actions, feature_size) for i in 1:num_demons]
    w_real = [zeros(F, fsna) for i in 1:num_demons]
    LSTDLearner(update,
                b,
                A_inv,
                0,
                e,
                w,
                w_real,
                feature_size,
                num_actions,
                num_demons)
end

LSTDLearner(args...) = LSTDLearner{Float64}(args...)
update(l::LSTDLearner) = l.update

function get_demon_parameters(learner::LSTDLearner,
                              demons,
                              obs,
                              state,
                              action,
                              next_obs,
                              next_state,
                              next_action)
    C, next_discounts, π_t = get(demons, obs, action, next_obs, next_action)
    next_target_pis = get_demon_pis(demons, learner.num_actions, next_state, next_obs)
    C, next_discounts, π_t, next_target_pis
end

function get_weights(l::LSTDLearner)
    return l.w_real
end

(learner::LSTDLearner)(state, action) = predict(learner::LSTDLearner, state, action)

function predict(learner::LSTDLearner, state, action)
    ϕ = get_active_action_state_vector(state,
                                       action,
                                       learner.feature_size,
                                       learner.num_actions)
    return [dot(learner.w_real[i], ϕ) for i ∈ 1:learner.num_demons]
end

max_stars_repo_name: MatthewMcLeod/curiosity

max_issues_repo_name: MatthewMcLeod/curiosity

max_forks_repo_name: MatthewMcLeod/curiosity


text:
# const TRAVIS_SKIP = VERSION.minor != 4 && !isnothing(get(ENV, "TRAVIS_BRANCH", nothing))

function test_fixed_size(A, At, B, Bt, Aa, Aat, Ba, Bat)::NTuple{4,Float64}
    t0 = @elapsed(@test Aa * Ba ≈ A * B)
    t1 = @elapsed(@test Aa * Bat ≈ A * Bt)
    t2 = @elapsed(@test Aat * Ba ≈ At * B)
    t3 = @elapsed(@test Aat * Bat ≈ At * Bt)
    t0, t1, t2, t3
end

function test_fixed_size(M, K, N)
    A = @StrideArray rand(M,K);
    B = @StrideArray rand(K,N);
    At = (@StrideArray rand(K,M))';
    Bt = (@StrideArray rand(N,K))';
    Aa = Array(A); Ba = Array(B);
    Aat = Array(At); Bat = Array(Bt);
    t0 = test_fixed_size(A, At, B, Bt, Aa, Aat, Ba, Bat)
    t1 = test_fixed_size(StrideArrays.make_dynamic(A), StrideArrays.make_dynamic(At), B, Bt, Aa, Aat, Ba, Bat)
    t2 = test_fixed_size(A, At, StrideArrays.make_dynamic(B), StrideArrays.make_dynamic(Bt), Aa, Aat, Ba, Bat)
    t3 = test_fixed_size(StrideArrays.make_dynamic(A), StrideArrays.make_dynamic(At), StrideArrays.make_dynamic(B), StrideArrays.make_dynamic(Bt), Aa, Aat, Ba, Bat)
    gflops = let gflop = 2e-9M*K*N
        map((t0, t1, t2, t3)) do t
            gflop ./ t
        end
    end
    @show (M,K,N), gflops
    nothing
end

test_fixed_size(M) = test_fixed_size(M, M, M)

# gmul(A, B) = LinearAlgebra.generic_matmatmul('N','N', A, B)

@testset "MatrixMultiply" begin
    r = 2:12
    for M ∈ r, K ∈ r, N ∈ r
        test_fixed_size(M, K, N)
    end
    r = 13:33
    for M ∈ r
        test_fixed_size(M)
    end
    M = K = N = 80
    A = @StrideArray rand(M,K);
    B = @StrideArray rand(K,N);
    C = StrideArray{Float64}(undef, (StaticInt(M),StaticInt(N)));
    M, K, N = 23, 37, 19
    @views begin
        Av = A[1:M, 1:K];
        Bv = B[1:K, 1:N];
        Cv = C[1:M, 1:N];
        Avsl = A[StaticInt(1):M, StaticInt(1):K];
        Bvsl = B[StaticInt(1):K, StaticInt(1):N];
        Cvsl = C[StaticInt(1):M, StaticInt(1):N];
        Avsr = A[StaticInt(1):StaticInt(M), StaticInt(1):StaticInt(K)];
        Bvsr = B[StaticInt(1):StaticInt(K), StaticInt(1):StaticInt(N)];
        Cvsr = C[StaticInt(1):StaticInt(M), StaticInt(1):StaticInt(N)];
    end
    Creference = Array(Av) * Array(Bvsl);
    time = @elapsed mul!(Cv, Av, Bv)
    @test Creference ≈ Cv
    @show M, K, N, time
    time = @elapsed mul!(Cvsl, Avsl, Bvsl)
    @test Creference ≈ Cv
    @show M, K, N, time
    time = @elapsed mul!(Cvsr, Avsr, Bvsr)
    @test Creference ≈ Cv
    @show M, K, N, time

    Av = A[1:M, 1:K];
    Bv = B[1:K, 1:N];
    Cv = C[1:M, 1:N];
    Avsl = A[StaticInt(1):M, StaticInt(1):K];
    Bvsl = B[StaticInt(1):K, StaticInt(1):N];
    Cvsl = C[StaticInt(1):M, StaticInt(1):N];
    Avsr = A[StaticInt(1):StaticInt(M), StaticInt(1):StaticInt(K)];
    Bvsr = B[StaticInt(1):StaticInt(K), StaticInt(1):StaticInt(N)];
    Cvsr = C[StaticInt(1):StaticInt(M), StaticInt(1):StaticInt(N)];
    Creference = Array(Av) * Array(Bvsl);
    time = @elapsed mul!(Cv, Av, Bv)
    @test Creference ≈ Cv
    @show M, K, N, time
    time = @elapsed mul!(Cvsl, Avsl, Bvsl)
    @test Creference ≈ Cv
    @show M, K, N, time
    time = @elapsed mul!(Cvsr, Avsr, Bvsr)
    @test Creference ≈ Cv
    @show M, K, N, time
end


text:
using DelimitedFiles
using WriteVTK
using MeshCore: AbsShapeDesc, P1, L2, T3, Q4, T4, H8
using MeshCore: VecAttrib
using MeshCore: shapedesc, nshapes,  IncRel
using MeshCore: ShapeColl
using LinearAlgebra: norm
using StaticArrays

"""
    Q4blockx(xs::Vector{T}, ys::Vector{T}; intbytes = 8) where {T}

Generate a graded quadrilateral mesh  of a 2D block.

Keyword argument `intbytes` controls the size of the integer indexes.

# Return
By convention the function returns an incidence relation (`connectivity`)
between the elements (`connectivity.left`) and the vertices
(`connectivity.right`). The geometry is stored as the attribute "geom" of the
vertices.
"""
function Q4blockx(xs::Vector{T}, ys::Vector{T}; intbytes = 8) where {T}
	inttype = Int64
    if intbytes == 4
        inttype = Int32
    end
    nL = length(xs)-1;
    nW = length(ys)-1;
    nnodes = (nL+1)*(nW+1);
    ncells = (nL)*(nW);
    xys = zeros(T, nnodes, 2);
    conns = zeros(inttype, ncells, 4);

    f=1;
    for j in 1:(nW+1)
        for i in 1:(nL+1)
            xys[f, 1] = xs[i]
            xys[f, 2] = ys[j]
            f=f+1;
        end
    end

    gc=1;
    for i in 1:nL
        for j in 1:nW
        	f=(j-1)*(nL+1)+i;
        	f = (j-1) * (nL+1) + i;
        	conns[gc, 1] = f
        	conns[gc, 2] = (f+1)
        	conns[gc, 3] = f+(nL+1)+1
        	conns[gc, 4] = f+(nL+1)
        	gc=gc+1;
        end
    end

    C = inttype.(conns[1:gc-1, :]);
    N, TC = size(xys, 2), eltype(xys)
    locs =  VecAttrib([SVector{N, TC}(xys[i, :]) for i in 1:size(xys, 1)])
    vertices = ShapeColl(P1, length(locs), "vertices")
    vertices.attributes["geom"] = locs
    elements = ShapeColl(Q4, size(C, 1), "elements")
    return IncRel(elements, vertices, C)
end

"""
    Q4block(Length, Width, nL, nW; kwargs...)

Generate a quadrilateral mesh  of the 2D block.

See also: Q4blockx
"""
function Q4block(Length, Width, nL, nW; kwargs...)
    return Q4blockx(collect(linearspace(0.0, Length, nL+1)),
                    collect(linearspace(0.0, Width, nW+1)); kwargs...);
end

"""
    Q4quadrilateral(xyz::Matrix{T}, nL, nW) where {T}

Mesh of a general quadrilateral given by the location of the vertices.
"""
function Q4quadrilateral(xyz::Matrix{T}, nL, nW; intbytes = 8) where {T}
    npts = size(xyz,1);
    if npts == 2 # In this case the quadrilateral must be defined in two dimensions
        lo = minimum(xyz, dims = 1);
        hi = maximum(xyz, dims = 1);
        xyz = [[lo[1] lo[2]];
            [hi[1] lo[2]];
            [hi[1] hi[2]];
            [lo[1] hi[2]]];
    elseif npts != 4
        error("Need 2 or 4 points");
    end

    # Generate elements and vertices in a 2D square
    ir = Q4block(2.0, 2.0, nL, nW; intbytes = intbytes);

    function bfun(param_coords)
        return SMatrix{4, 1}([0.25 * (1. - param_coords[1]) * (1. - param_coords[2]);
               0.25 * (1. + param_coords[1]) * (1. - param_coords[2]);
               0.25 * (1. + param_coords[1]) * (1. + param_coords[2]);
               0.25 * (1. - param_coords[1]) * (1. + param_coords[2])]);
    end

    # Remap the quadrilateral mesh to the physical coordinates
    locs = ir.right.attributes["geom"];
    N, TC = size(xyz, 2), eltype(locs[1])
    v = SVector{N, TC}[]
    for i = 1:length(locs)
        bfv = bfun([locs[i][1]-1.0, locs[i][2]-1.0]);# shift coordinates by -1
        push!(v, SVector{N, TC}(bfv'*xyz));
    end
    nlocs =  VecAttrib(v)
    ir.right.attributes["geom"] = nlocs
    return ir
end

"""
    Q4blockwdistortion(Length, Width, nL, nW; intbytes = 8)

Generate mesh of a rectangular block with distorted quadrilaterals.
"""
function Q4blockwdistortion(Length, Width, nL, nW; intbytes = 8)
    nL2 = max(1, Int32(round(nL/2)))
    nW2 = max(1, Int32(round(nW/2)))
    c1 = Q4quadrilateral([-1 -1; -0.2 -1; -0.1 -0.2; -1 0.8], nL2, nW2; intbytes = intbytes)
    c2 = Q4quadrilateral([-0.2 -1; 1 -1; 1 -0.5; -0.1 -0.2], nL2, nW2; intbytes = intbytes)
    c3 = Q4quadrilateral([-0.1 -0.2; 1 -0.5; 1 1; 0.3 1], nL2, nW2; intbytes = intbytes)
    c4 = Q4quadrilateral([-1 0.8; -0.1 -0.2; 0.3 1; -1 1], nL2, nW2; intbytes = intbytes)
    c  = mergeirs(c1, c2, 0.001/max(nL2, nW2))
    c  = mergeirs(c, c3, 0.001/max(nL2, nW2))
    c  = mergeirs(c, c4, 0.001/max(nL2, nW2))
    transform(c, x -> begin
        [(x[1]+1.0)/2*Length, (x[2]+1.0)/2*Width]
    end)
    return c
end


text:
try
    using AA222Testing
catch
    using Pkg
    Pkg.add(PackageSpec(url="https://github.com/sisl/AA222Testing.git"))
    using AA222Testing
end

using AA222Testing: Test, localtest

include(joinpath("project0_jl", "project0.jl"))

test_f(a, b) = () -> (f(a, b) == (a + b))

tests = [Test(test_f(1, 1)),
         Test(test_f(7.2, 9.0))]

localtest(tests, show_errors = get(ARGS, 1, true))


text:
####### SIMULATION OF MARGINAL ###########################################

### +++ INITIALIZATION +++

### import packages
using Random, LinearAlgebra, RDatasets, MixedModels, Distributions, Optim, CSV, Dates, CategoricalArrays, GLM

### define functions
vcov2 = function(x, xbar, σ, τ, n, γ)
    m = length(xbar)
    a = γ / τ^2 * m # 1-1 entry
    d = (sum(x.^2) - n * γ * sum(xbar.^2) ) / σ^2 # 2-2 entry
    b = γ / τ^2 / n * sum(x) # 2-1 entry
    [d -b; -b a] / (a * d - b^2) # return inverse (X^t V^-1 X)^-1
end

beta = function(vcov1, x, xbar, y, ybar, σ, τ, n, γ)
    c1 = γ / τ^2 * sum(ybar) # = 1^t V^-1 y
    c2 = (x' * y - n * γ * sum(ybar .* xbar))/σ^2 # = x^t V^-1 y
    vcov1 * [c1 c2]'  # (X^tV^-1X)^-1 X^tV^-1y
end

###     SELECT PARAMETERS
m = 10
n = 1
τ0 = 8^0.5
σ0 = 4^0.5

### fix parameters
p = 2
nsim = 10000
Δ = 0
Random.seed!(123)

### +++ PER SPECIFICATION SETTING +++
REdistr = "normal"

### variable parameters
N = m * n

β0 = [-4.9, 0.03]
x = 10 * randn(N)
xbar = vec(sum(reshape(x, n, :), dims=1) / n)
X = hcat(fill(1, N), x)

### conditional evaluation
if REdistr == "normal"
    v = τ0 * randn(m)
elseif REdistr == "tdist"
    v = rand(TDist(2), m)
elseif REdistr == "chi2"
    v = rand(Chisq(τ0^2 / 2), m) .- τ0^2 / 2 - m
end
u = β0[1] .+ β0[2] * x + repeat(v, inner = n)
μ0 = β0[1] .+ β0[2] .* xbar + v

hatmatrix = inv(X' * X) * X'
detX = det(X' * X)
loglik = function(α, y)
    ξ = exp(α[1])
    beta = hatmatrix * y
    residual = y - X * beta
    (p * log(detX / ξ) + N * log(ξ) + residual' * residual / ξ)
end

τ1 = zeros(10000, 1)
for i in 1:10000
    y = u + σ0 * randn(N)
    τ1[i] = sqrt(maximum([exp(optimize(α -> loglik(α,y), [1.0], BFGS()).minimizer[1]) - σ0^2 0]))
end
τ1 = mean(τ1) # should be (close to) τ0, but theory for conditional convergence is missing
σ1 = σ0  #mean(σ1)

### +++ KNOWN DELTA +++

###     MARGINAL COVARIANCE MATRIX

### helpful parameters
γ0 = τ1^2 / (τ1^2 + σ1^2 / n)
vcov0 = vcov2(x, xbar, σ1, τ1, n, γ0) # = (X^tV^-1X)^-1
d0 = (1 - γ0) .* hcat(fill(1, m), xbar) # d' = l' - b'X
dd0 = d0' * d0

### assemble covariance matrix parts
g1 = γ0 * σ1^2 / n

### assemble inverse of covariance matrix with Woodbury
Σ0diag = g1
Σ0i = (I - d0 * inv(Σ0diag * inv(vcov0) + dd0) * d0' ) / Σ0diag

###     CONDITIONAL COVARIANCE MATRIX
l1 = γ0^2 / n * σ1^2

### assemble inverse of covariance matrix with Woodbury
cΣ0diag = l1
cΣ0i = (I - d0 * inv(cΣ0diag * σ1^2 *
            inv(vcov0 * (2 * γ0 * I + n^2 * dd0 * vcov0)) + dd0) * d0' ) / cΣ0diag

###     NONCENTRALITY PARAMETER
AZv = (γ0 - 1) .* v + n / τ1^2 .* d0 * vcov0 * d0' * v
λ0 = AZv' * cΣ0i * AZv

###     QUANTILES
qc0 = cquantile(NoncentralChisq(m, λ0), 0.05) # quantile of alpha = 0.05
qm = cquantile(Chisq(m), 0.05) # quantile of alpha = 0.05

### simulation
Coverage = zeros(5, length(Δ))
λ = zeros(nsim, 2)
Volume = ones(nsim, 5)
Τ = zeros(nsim, 1)
Statistic = zeros(nsim, 4)

for i in 1:nsim

    # i = 1
    ### +++ EMPIRICAL +++
    y = u + σ0 * randn(N)
    τ = sqrt(maximum([exp(optimize(α -> loglik(α,y), [1.0], BFGS()).minimizer[1]) - σ0^2 0.01]))
    Τ[i,1] = τ

    ybar = vec(sum(reshape(y, n, :), dims=1) / n)

    ### make sure the variance parameters are positive
    σ = σ1

    ###     MARGINAL

    ### helpful parameters
    γ = τ^2 / (τ^2 + σ^2 / n)
    vcov1 = vcov2(x, xbar, σ, τ, n, γ) # = (X^tV^-1X)^-1
    d = (1-γ) .* hcat(fill(1, m), xbar) # d' = l' - b'X
    dd = d' * d
    a = n * τ^2 + σ^2

    ### assemble diagonal matrix parts
    g1 = γ * σ^2 / n
    g3 = 2 * σ^4 / τ^2 * γ / N

    ### assemble inverse of covariance matrix with Woodbury
    Σdiag = g1 + 2 * g3
    Σ = Σdiag * I + d * vcov1 * d'
    Σi = (I - d * inv(Σdiag * inv(vcov1) + dd) * d' ) / Σdiag

    ###     CONDITIONAL

    ### assemble diagonal matrix parts
    l1 = γ^2 / n * σ^2

    l3_1 = 2 * σ^2 * τ^2 / N/ a^3
    l3_2 = 8 / N / a^4 / n^2 * τ^2 * σ^2 * (
        + (a^2 - σ^2) / n * (
            + σ^6 / a^2 * n
            - (σ^4  / a^2) * (τ + a / 2 * n))
        + σ^10 / a^2 * (n^3 + 1 / n)
        + a * σ^4 * (n - 1 + 1 / n / 2)  +
        - σ^6 * (2 * n - 1 + 1 / n) +
        + σ^8 / a * (1 + (n-1) * (n^2 + 1) / n) / 2)
    l3_3 = 0
    l3_4 = - 8 / N/ n * τ^2 * σ^4 / a^3 * (
        - a
        + σ^2 * (n + 1))
    l3 = 2 * (l3_1 + l3_2 + l3_3 + l3_4)

    l4 = 2 * σ^4 / N / a^2 * ( 1 / n + 1)

    l5 = 2 * n / N * σ^4 * (σ^2 - 2 * n * τ^2) / a^2

    ### assemble inverse of covariance matrix with Woodbury
    cΣdiag = maximum([l1 + l3 + l4 - l5, 1/100000])
    cΣi = (I - d * inv(cΣdiag * σ^2 *
                inv(vcov1 * (2 * γ * I + n^2 * dd * vcov1)) + dd) * d' ) / cΣdiag

    ###     COMPUTE MEAN
    β1 = beta(vcov0, x, xbar, y, ybar, σ1, τ1, n, γ0)
    μ1 = (1 - γ0) * (β1[1] .+ β1[2] .* xbar) + γ0 * ybar

    β2 = beta(vcov1, x, xbar, y, ybar, σ, τ, n, γ) # = fixef(model)
    μ2 = (1 - γ) * (β2[1] .+ β2[2] .* xbar) + γ * ybar

    ###     NONCENTRALITY PARAMETER
    ay = d0 * β1 - (1 - γ0) * ybar
    λ[i, 1] = maximum([(ay' * cΣ0i * ay)[1] - tr(cΣ0i) * γ0^2 / n^3 / τ1^4 * σ1^2, 0])

    ay = d * β2 - (1 - γ) * ybar
    λ[i, 2] = maximum([(ay' * cΣi * ay)[1] - tr(cΣi) * σ^2 / a^2 / n, 0])

    ###     QUANTILES
    qc1 = cquantile(NoncentralChisq(m, λ[i, 1]), 0.05) # quantile of alpha = 0.05
    qc2 = cquantile(NoncentralChisq(m, λ[i, 2]), 0.05) # quantile of alpha = 0.05

    ###     COMPUTE STATISTICS
    t1 = (μ1 - μ0) .+ Δ'
    t2 = (μ2 - μ0) .+ Δ'

    Statistic[i,:] = [sum(t1' .* (t1' * Σ0i), dims = 2)[1],  #marginal known
                      sum(t2' .* (t2' * Σi), dims = 2)[1], #marginal REML
                      sum(t1' .* (t1' * cΣ0i), dims = 2)[1], #conditional known
                      sum(t2' .* (t2' * cΣi), dims = 2)[1]] #conditional REML

    Coverage[1,1] += (Statistic[i,1] .< qm) / nsim
    Coverage[2,1] += (Statistic[i,2] .< qm) / nsim
    Coverage[3,1] += (Statistic[i,3] .< qc0) / nsim
    Coverage[4,1] += (Statistic[i,3] .< qc1) / nsim
    Coverage[5,1] += (Statistic[i,4] .< qc2) / nsim

    ###     COMPUTE RELATIVE VOLUME
    # (via approximations as the determinant of large matrices is numerically unstable)
    Volume[i, 3] = (Σ0diag / cΣ0diag * qm / qc0)^(-m / 2)
    Volume[i, 4] = (Σ0diag / cΣ0diag * qm / qc1)^(-m / 2)
    Volume[i, 5] = (Σdiag / cΣdiag * qm / qc2)^(-m / 2)
end


text:
using
    BenchmarkTools,
    LinearAlgebra,
    KalmanFilter,
    Plots,
    LazyArrays,
    Colors

function init_mu(num_states, num_measures)
    x = randn(num_states)
    PL = randn(num_states, num_states)
    P = PL'PL
    RL = randn(num_measures, num_measures)
    R = RL'RL
    y = randn(num_measures)
    H = randn(num_measures, num_states)
    P_chol = cholesky(P)
    R_chol = cholesky(R)
    h(x) = H * x
    h!(y, x) = mul!(y, H, x)
    h(x, noise) = H * x .+ noise
    h!(y, x, noise) = y .= @~ H * x .+ noise
    return x, y, P, H, R, P_chol, R_chol, h, h!
end

function init_tu(num_states)
    x = randn(num_states)
    A = randn(num_states, num_states)
    P = A'A
    B = randn(num_states, num_states)
    Q = B'B
    F = randn(num_states, num_states)
    P_chol = cholesky(P)
    Q_chol = cholesky(Q)
    f(x) = F * x
    f!(y, x) = mul!(y, F, x)
    f(x, noise) = F * x .+ noise
    f!(y, x, noise) = y.= @~ F * x .+ noise
    return x, P, Q, F, P_chol, Q_chol, f, f!
end

function run_measurement_update_benchmarks(num_state_tests, num_measurement_tests; allocation = false)

    num_measurements = length(num_measurement_tests)
    kf_types = (:kf, :srkf, :ukf, :srukf, :aukf, :sraukf)
    buffers = [(inplace = zeros(length(num_state_tests), num_measurements), allocating = zeros(length(num_state_tests), num_measurements)) for _ = 1:length(kf_types)]
    results = NamedTuple{kf_types}(buffers)

    for (i, num_states) in enumerate(num_state_tests)
        for (j, num_measures) in enumerate(num_measurement_tests)
            x, y, P, H, R, P_chol, R_chol, h, h! = init_mu(num_states, num_measures)

            results.kf.allocating[i,j] = if allocation
                @allocated measurement_update(x, P, y, H, R)
            else
                @belapsed measurement_update($x, $P, $y, $H, $R)
            end
            kf_inter = KFMUIntermediate(num_states, num_measures)
            results.kf.inplace[i,j] = if allocation
                @allocated measurement_update!(kf_inter, x, P, y, H, R)
            else
                @belapsed measurement_update!($kf_inter, $x, $P, $y, $H, $R)
            end

            results.srkf.allocating[i,j] = if allocation
                @allocated measurement_update(x, P_chol, y, H, R_chol)
            else
                @belapsed measurement_update($x, $P_chol, $y, $H, $R_chol)
            end
            srkf_inter = SRKFMUIntermediate(num_states, num_measures)
            results.srkf.inplace[i,j] = if allocation
                @allocated measurement_update!(srkf_inter, x, P_chol, y, H, R_chol)
            else
                @belapsed measurement_update!($srkf_inter, $x, $P_chol, $y, $H, $R_chol)
            end

            results.ukf.allocating[i,j] = if allocation
                @allocated measurement_update(x, P, y, h, R)
            else
                @belapsed measurement_update($x, $P, $y, $h, $R)
            end
            ukf_inter = UKFMUIntermediate(num_states, num_measures)
            results.ukf.inplace[i,j] = if allocation
                @allocated measurement_update!(ukf_inter, x, P, y, h!, R)
            else
                @belapsed measurement_update!($ukf_inter, $x, $P, $y, $h!, $R)
            end

            results.srukf.allocating[i,j] = if allocation
                @allocated measurement_update(x, P_chol, y, h, R_chol)
            else
                @belapsed measurement_update($x, $P_chol, $y, $h, $R_chol)
            end
            srukf_inter = SRUKFMUIntermediate(num_states, num_measures)
            results.srukf.inplace[i,j] = if allocation
                @allocated measurement_update!(srukf_inter, x, P_chol, y, h!, R_chol)
            else
                @belapsed measurement_update!($srukf_inter, $x, $P_chol, $y, $h!, $R_chol)
            end

            results.aukf.allocating[i,j] = if allocation
                @allocated measurement_update(x, P, y, h, Augment(R))
            else
                @belapsed measurement_update($x, $P, $y, $h, $(Augment(R)))
            end
            aukf_inter = AUKFMUIntermediate(num_states, num_measures)
            results.aukf.inplace[i,j] = if allocation
                @allocated measurement_update!(aukf_inter, x, P, y, h!, Augment(R))
            else
                @belapsed measurement_update!($aukf_inter, $x, $P, $y, $h!, $(Augment(R)))
            end

            results.sraukf.allocating[i,j] = if allocation
                @allocated measurement_update(x, P_chol, y, h, Augment(R_chol))
            else
                @belapsed measurement_update($x, $P_chol, $y, $h, $(Augment(R_chol)))
            end
            sraukf_inter = SRAUKFMUIntermediate(num_states, num_measures)
            results.sraukf.inplace[i,j] = if allocation
                @allocated measurement_update!(sraukf_inter, x, P_chol, y, h!, Augment(R_chol))
            else
                @belapsed measurement_update!($sraukf_inter, $x, $P_chol, $y, $h!, $(Augment(R_chol)))
            end
        end
    end
    results
end

function run_time_update_benchmarks(num_state_tests; allocation = false)

    kf_types = (:kf, :srkf, :ukf, :srukf, :aukf, :sraukf)
    buffers = [(inplace = zeros(length(num_state_tests)), allocating = zeros(length(num_state_tests))) for _ = 1:length(kf_types)]
    results = NamedTuple{kf_types}(buffers)

    for (i, num_states) in enumerate(num_state_tests)
        x, P, Q, F, P_chol, Q_chol, f, f! = init_tu(num_states)

        results.kf.allocating[i] = if allocation
            @allocated time_update(x, P, F, Q)
        else
            @belapsed time_update($x, $P, $F, $Q)
        end
        kf_inter = KFTUIntermediate(num_states)
        results.kf.inplace[i] = if allocation
            @allocated time_update!(kf_inter, x, P, F, Q)
        else
            @belapsed time_update!($kf_inter, $x, $P, $F, $Q)
        end

        results.srkf.allocating[i] = if allocation
            @allocated time_update(x, P_chol, F, Q_chol)
        else
            @belapsed time_update($x, $P_chol, $F, $Q_chol)
        end
        srkf_inter = SRKFTUIntermediate(num_states)
        results.srkf.inplace[i] = if allocation
            @allocated time_update!(srkf_inter, x, P_chol, F, Q_chol)
        else
            @belapsed time_update!($srkf_inter, $x, $P_chol, $F, $Q_chol)
        end

        results.ukf.allocating[i] = if allocation
            @allocated time_update(x, P, f, Q)
        else
            @belapsed time_update($x, $P, $f, $Q)
        end
        ukf_inter = UKFTUIntermediate(num_states)
        results.ukf.inplace[i] = if allocation
            @allocated time_update!(ukf_inter, x, P, f!, Q)
        else
            @belapsed time_update!($ukf_inter, $x, $P, $f!, $Q)
        end

        results.srukf.allocating[i] = if allocation
            @allocated time_update(x, P_chol, f, Q_chol)
        else
            @belapsed time_update($x, $P_chol, $f, $Q_chol)
        end
        srukf_inter = SRUKFTUIntermediate(num_states)
        results.srukf.inplace[i] = if allocation
            @allocated time_update!(srukf_inter, x, P_chol, f!, Q_chol)
        else
            @belapsed time_update!($srukf_inter, $x, $P_chol, $f!, $Q_chol)
        end

        results.aukf.allocating[i] = if allocation
            @allocated time_update(x, P, f, (Augment(Q)))
        else
            @belapsed time_update($x, $P, $f, $(Augment(Q)))
        end
        aukf_inter = AUKFTUIntermediate(num_states)
        results.aukf.inplace[i] = if allocation
            @allocated time_update!(aukf_inter, x, P, f!, Augment(Q))
        else
            @belapsed time_update!($aukf_inter, $x, $P, $f!, $(Augment(Q)))
        end

        results.sraukf.allocating[i] = if allocation
            @allocated time_update(x, P_chol, f, Augment(Q_chol))
        else
            @belapsed time_update($x, $P_chol, $f, $(Augment(Q_chol)))
        end
        sraukf_inter = SRAUKFTUIntermediate(num_states)
        results.sraukf.inplace[i] = if allocation
            @allocated time_update!(sraukf_inter, x, P_chol, f!, Augment(Q_chol))
        else
            @belapsed time_update!($sraukf_inter, $x, $P_chol, $f!, $(Augment(Q_chol)))
        end
    end
    results
end

function plot_benchmarks(results, num_state_tests; num_measurement_tests = 1:1, ylabel = "Time (μs)", scale = 10^-6, logscale = false)
    colors = distinguishable_colors(length(results), [RGB(1,1,1), RGB(0,0,0)], dropseed = true)
    num_rows = ceil(Int, length(num_measurement_tests) / 2)
    p = plot(
        num_state_tests,
        results.kf.allocating / scale,
        layout = length(num_measurement_tests) > 1 ? (num_rows, 2) : 1,
        seriescolor = colors[1],
        legend = false,
        title = length(num_measurement_tests) > 1 ? permutedims(map(x -> "$x measurements", num_measurement_tests)) : "Time update",
        xlabel = "# States",
        ylabel = ylabel,
        yscale = logscale ? :log10 : :identity
    )
    plot!(num_state_tests, results.kf.inplace / scale, layout = length(num_measurement_tests), seriescolor = colors[1], linestyle = :dash)

    plot!(num_state_tests, results.srkf.allocating / scale, layout = length(num_measurement_tests), seriescolor = colors[2])
    plot!(num_state_tests, results.srkf.inplace / scale, layout = length(num_measurement_tests), seriescolor = colors[2], linestyle = :dash)

    plot!(num_state_tests, results.ukf.allocating / scale, layout = length(num_measurement_tests), seriescolor = colors[3])
    plot!(num_state_tests, results.ukf.inplace / scale, layout = length(num_measurement_tests), seriescolor = colors[3], linestyle = :dash)

    plot!(num_state_tests, results.srukf.allocating / scale, layout = length(num_measurement_tests), seriescolor = colors[4])
    plot!(num_state_tests, results.srukf.inplace / scale, layout = length(num_measurement_tests), seriescolor = colors[4], linestyle = :dash)

    plot!(num_state_tests, results.aukf.allocating / scale, layout = length(num_measurement_tests), seriescolor = colors[5])
    plot!(num_state_tests, results.aukf.inplace / scale, layout = length(num_measurement_tests), seriescolor = colors[5], linestyle = :dash)

    plot!(num_state_tests, results.sraukf.allocating / scale, layout = length(num_measurement_tests), seriescolor = colors[6], label = "SRAUKF")
    plot!(num_state_tests, results.sraukf.inplace / scale, layout = length(num_measurement_tests), seriescolor = colors[6], linestyle = :dash, label = "SRAUKF inplace")

    # Hack for single legend
    legend_plot = plot([1], framestyle = :none, seriescolor = colors[1], label = "KF")
    plot!([1], seriescolor = colors[1], linestyle = :dash, label = "KF inplace")

    plot!([1], seriescolor = colors[2], label = "SRKF")
    plot!([1], seriescolor = colors[2], linestyle = :dash, label = "SRKF inplace")

    plot!([1], seriescolor = colors[3], label = "UKF")
    plot!([1], seriescolor = colors[3], linestyle = :dash, label = "UKF inplace")

    plot!([1], seriescolor = colors[4], label = "SRUKF")
    plot!([1], seriescolor = colors[4], linestyle = :dash, label = "SRUKF inplace")

    plot!([1], seriescolor = colors[5], label = "AUKF")
    plot!([1], seriescolor = colors[5], linestyle = :dash, label = "AUKF inplace")

    plot!([1], seriescolor = colors[6], label = "SRAUKF")
    plot!([1], seriescolor = colors[6], linestyle = :dash, label = "SRAUKF inplace")
    plot(p, legend_plot, layout = grid(2, 1, heights = [num_rows, 1.2] / (num_rows + 1.2)), size = (800, num_rows * 440))
end

num_state_tests = [1, 5, 10, 20, 30, 40, 50, 60]
num_measurement_tests = [2, 4, 8, 16, 32, 64]
tu_time = run_time_update_benchmarks(num_state_tests)
tu_allocations = run_time_update_benchmarks(num_state_tests, allocation = true)

tu_time_plot = plot_benchmarks(tu_time, num_state_tests)
tu_alloc_plot = plot_benchmarks(tu_allocations, num_state_tests, ylabel = "Allocations (kB)", scale = 10^3)

png(tu_time_plot, "tu_time")
png(tu_alloc_plot, "tu_alloc")

mu_time = run_measurement_update_benchmarks(num_state_tests, num_measurement_tests)
mu_allocations = run_measurement_update_benchmarks(num_state_tests, num_measurement_tests, allocation = true)

mu_time_plot = plot_benchmarks(mu_time, num_state_tests, num_measurement_tests = num_measurement_tests)
mu_alloc_plot = plot_benchmarks(mu_allocations, num_state_tests, num_measurement_tests = num_measurement_tests, ylabel = "Allocations (kB)", scale = 10^3)

png(mu_time_plot, "mu_time")
png(mu_alloc_plot, "mu_alloc")


text:
using FourierFlows: stepforward!
using Printf
using NPZ

include("nondimlinearonehalflayerqg.jl")
set_q! = NondimLinearOneHalfLayerQG.set_q!
updatevars! = NondimLinearOneHalfLayerQG.updatevars!
energies = NondimLinearOneHalfLayerQG.energies
get_bickley = NondimLinearOneHalfLayerQG.get_bickley
get_σ = NondimLinearOneHalfLayerQG.get_σ

include("common_params_onehalflay.jl")

for F in all_Fs

println("")
println("F = ", string(F))
println("")
F = float(F)

prob = NondimLinearOneHalfLayerQG.Problem(dev;
                                          nx = N,
                                          ny = N,
                                          Lx = L,
                                          Ly = L,
                                          x0 = 0.0,
                                           F = F,
                                           U = U,
                                         Uyy = Uyy,
                                           β = β,
                                          dt = dt,
                                     stepper = stepper,
                                      effort = effort)

sol, clock, params, vars, grid = prob.sol, prob.clock, prob.params, prob.vars, prob.grid
set_q!(prob, q₀)

if PSI_HOVMOLLER thovm, ψhovm = nothing, nothing; fm = Int(round(N/2)) end
startwalltime = time()
j = 0
R = 1.0
cyc = 0
Eo = energies(prob) # Reference KE level for renormalization.
keo, peo = Eo[1][1], Eo[2][1]
t = Array([clock.t])
KE = Array([keo])
PE = Array([peo])

while cyc<cycles
  stepforward!(prob)
  updatevars!(prob)
  E = energies(prob)
  ke, pe = E[1][1], E[2][1]

  if j % nlog == 0
    cfl = clock.dt * maximum([maximum(vars.u) / grid.dx, maximum(vars.v) / grid.dy])
    log = @sprintf("step: %06d, t: %2.1f, R: %1.2f, cfl: %1.2f, KE/PE: %1.2f, walltime: %1.2f min", clock.step, clock.t, R, cfl, ke/pe, (time()-startwalltime)/60)
    println(log)

    if cyc==cyclesm && PSI_HOVMOLLER
      if isnothing(thovm)
        ψhovm = vars.ψ[:, fm]
        thovm = Array([clock.t])
      else
        ψhovm = cat(ψhovm, vars.ψ[:, fm], dims=2)
        append!(thovm, clock.t)
      end
    end
  end

  append!(t, clock.t)
  append!(KE, ke)
  append!(PE, pe)

  j += 1
  R = keo/ke # renormalize (KE_{ref}/KE).
  if R<Rthresh
    set_q!(prob, vars.q*R)
    cyc += 1
    println("")
    println("Completed renormalization cycle", " ", cyc, "/", cycles)
    println("")
  end
end

# Get growth rate from exponential fit to KE time series.
σ = get_σ(t, KE, PE)

if SAVE_OUTPUT
  println("Saving output data.")

  npzname = "lin_onehalflay_N" * string(N) * "_F" * string(Int(F))
  npzname = "../simulations/" * npzname * ".npz"
  npzdata = Dict("p" => vars.ψ', "N" => N, "L" => L, "dt" => dt, "bet" => β, "sigma" => σ, "F" => F)

  if PSI_HOVMOLLER
    npzdata = merge(npzdata, Dict("phovm" => ψhovm', "thovm" => thovm))
  end

  npzwrite(npzname, npzdata)
end

end # all_Fs

max_stars_repo_name: apaloczy/RoughOceanQGEnergetics

max_issues_repo_name: apaloczy/RoughOceanQGEnergetics

max_forks_repo_name: apaloczy/RoughOceanQGEnergetics


text:
module SlepcWrap
    using PetscWrap
    using MPI

    include("load.jl")
    include("init.jl")
    export SlepcInitialize, SlepcFinalize

    include("eps.jl")
    export  SlepcEPS, CEPS,
            create_eps,
            get_eig, get_eigs,
            get_eigenvalue, get_eigenvalues,
            get_eigenpair,
            get_tolerances,
            neigs,
            EPSCreate,
            EPSDestroy,
            EPSGetConverged,
            EPSGetTolerances,
            EPSGetEigenpair,
            EPSGetEigenvalue,
            EPSGetOperators,
            EPSSetFromOptions,
            EPSSetOperators,
            EPSSetUp,
            EPSSolve,
            EPSView

    include("fancy/eps.jl")
    export  create_eps,
            get_eig, get_eigs,
            get_eigenvalue, get_eigenvalues,
            get_eigenpair,
            get_tolerances,
            neigs
end


text:
## -------------------------------------------------------------------
# marginal distributions
let
    objider = iJR.BIOMASS_IDER
    size = [300, 250]
    method2 = ME_MAX_POL

    # Iders
    model_iders, Fd_iders = [objider], ["D"]
    for Fd_met in EXCH_FLX_IDERS
        model_met = Fd_mets_map[Fd_met]
        model_exch = Fd_rxns_map[Fd_met]
        push!(model_iders, model_exch)
        push!(Fd_iders, string("u", Fd_met))
    end

    for (model_ider, Fd_ider) in zip(model_iders, Fd_iders)
        ps = Plots.Plot[]
        ps_bs = Plots.Plot[]
        for exp in EXPS
            p = plot(title = string(Fd_ider, " exp: ", exp))
            p_bs = plot(title = string(Fd_ider, " exp: ", exp))
            margin, m, M = -Inf, Inf, -Inf
            Fd_av = Fd.val(Fd_ider, exp)

            # EP
            for method in ALL_METHODS
                color = method_colors[method]

                datfile = INDEX[method, :DFILE, exp]
                dat = deserialize(datfile)
                model = dat[:model]
                objidx = ChU.rxnindex(model, objider)
                epouts = dat[:epouts]
                exp_beta = maximum(keys(epouts))
                epout = epouts[exp_beta]
                ep_av = ChU.av(model, epout, model_ider)
                ep_va = sqrt(ChU.va(model, epout, model_ider))

                ChP.plot_marginal!(p, model, [epout], model_ider;
                    legend = false, color, alpha = 0.6, lw = 5)

                m = minimum([m, ep_av, Fd_av])
                M = maximum([M, ep_av, Fd_av])
                margin = maximum([margin, 3 * ep_va])

                if method == method2
                    for (beta, epout) in sort(epouts; by = first)
                        ep_av = ChU.av(model, epout, model_ider)
                        ep_va = sqrt(ChU.va(model, epout, model_ider))

                        alpha = 0.15
                        color = method_colors[method]
                        ChP.plot_marginal!(p_bs, model, epout, model_ider;
                            legend = false, color, alpha, lw = 1)

                        if beta == exp_beta
                            ChP.plot_marginal!(p_bs, model, epout, model_ider;
                                legend = false, color,
                                alpha = 1.0, lw = 3
                            )
                            break
                        end
                    end
                    push!(ps_bs, p_bs)
                end

            end
            # Experimental
            vline!(p, [Fd_av]; label = "", lw = 6, color = :black, alpha = 0.3)
            vline!(p_bs, [Fd_av]; label = "", lw = 6, color = :black, alpha = 0.3)

            plot!(p; xlim = [m - margin, M + margin], size)
            plot!(p_bs; xlim = [m - margin, M + margin], size)
            push!(ps, p)
        end

        for k in [:xi, :D, :uGLC]
            p = plot(;title = Fd_ider, size)
            xticks =  (EXPS, string.(EXPS))
            vals = [Fd.val(k, exp) for exp in EXPS]
            p = bar!(p, EXPS, vals; title = k, label = "", xticks)
            push!(ps, p)
            push!(ps_bs, p)
        end

        # legend
        p = plot(;title = "Legend", size)
        for (method, color) in method_colors
            color = method_colors[method]
            bar!(p, [string(method)], [1]; yaxis = nothing,
                color, xrotation = 35, label = ""
            )
        end
        push!(ps, p)
        push!(ps_bs, p)

        pname = string(Fd_ider, "_marginals")
        mysavefig(ps, pname)

        method = method2
        pname = string(Fd_ider, "_marginals_vs_beta")
        mysavefig(ps_bs, pname; method)
    end

end

## -------------------------------------------------------------------
# marginals v2
let
    objider = iJR.BIOMASS_IDER
    size = [300, 250]
    Fd_mets_map = iJR.load_mets_map()
    exch_met_map = iJR.load_exch_met_map()

    # Iders
    model_iders, Fd_iders = [objider], ["D"]
    for Fd_met in EXCH_FLX_IDERS
        model_met = Fd_mets_map[Fd_met]
        model_exch = Fd_rxns_map[Fd_met]
        push!(model_iders, model_exch)
        push!(Fd_iders, string("u", Fd_met))
    end

    for (model_ider, Fd_ider) in zip(model_iders, Fd_iders)
        marg_params = (;xlabel = string(Fd_ider), yaxis = nothing, ylabel = "prob")

        epps = Plots.Plot[]
        exps = Plots.Plot[]
        for method in ALL_METHODS
            expp = plot(;title = string("Experimental"), marg_params...)
            epp = plot(;title = string(" MaxEnt: ", method), marg_params...)
            margin, m, M = -Inf, Inf, -Inf

            # EP
            for exp in EXPS
                Fd_av = Fd.val(Fd_ider, exp)
                color = exp_colors[exp]

                datfile = INDEX[method, :DFILE, exp]
                dat = deserialize(datfile)
                model = dat[:model]
                objidx = ChU.rxnindex(model, objider)
                epouts = dat[:epouts]
                exp_beta = maximum(keys(epouts))
                epout = epouts[exp_beta]
                ep_av = ChU.av(model, epout, model_ider)
                ep_va = sqrt(ChU.va(model, epout, model_ider))

                ChP.plot_marginal!(epp, model, [epout], model_ider;
                    legend = false, color, alpha = 0.8, lw = 3)

                m = minimum([m, ep_av, Fd_av])
                M = maximum([M, ep_av, Fd_av])
                margin = maximum([margin, 3 * ep_va])

                # Experimental
                vline!(expp, [Fd_av]; label = "", lw = 3, color, alpha = 0.8)

            end

            map([expp, epp]) do p
                plot!(p; xlim = [m - margin, M + margin], size)
            end

            push!(epps, epp)
            push!(exps, expp)
        end

        extras = Plots.Plot[]
        for k in [:xi, :X, Fd_ider] |> unique
            p = plot(;title = "Experimental", size,
                xlabel = "rep", ylabel = string(k))
            xticks =  (EXPS, string.(EXPS))
            vals = [Fd.val(k, exp) for exp in EXPS]
            color = [exp_colors[exp] for exp in EXPS]
            p = bar!(p, EXPS, vals; label = "", xticks, color)
            push!(extras, p)
        end

        # legend
        p = plot(;title = "Legend", size)
        for (method, color) in method_colors
            color = method_colors[method]
            bar!(p, [string(method)], [1]; yaxis = nothing,
                color, xrotation = 35, label = ""
            )
        end
        push!(extras, p)

        ps = Plots.Plot[exps; epps; extras]
        layout = (3, 4)
        pname = string(Fd_ider, "_marginals_v2")
        mysavefig(ps, pname; layout)

    end # for (model_ider, Fd_ider)

end


text:
heat_eq_dir = joinpath(output_root, "HeatEquation")

struct DifussionImplicit end
function TridiagonalMatrix(
    q::StateVec,
    aux::StateVec,
    Δt::T,
    ρ::S,
    a::S,
    K::S,
    grid::Grid{T},
    ::DifussionImplicit,
    i::Int,
) where {S <: Symbol, T}
    #! format: off
    a_ = [   -Δt * grid.Δzi2 *      (aux[ρ, Dual(k)] .* q[a, Dual(k), i] .* aux[K, Dual(k), i])[2] for k in over_elems_real(grid)[1:(end - 1)]]
    b_ = [1 + Δt * grid.Δzi2 * (sum((aux[ρ, Dual(k)] .* q[a, Dual(k), i] .* aux[K, Dual(k), i]))) for k in over_elems_real(grid)]
    c_ = [   -Δt * grid.Δzi2 *      (aux[ρ, Dual(k)] .* q[a, Dual(k), i] .* aux[K, Dual(k), i])[1] for k in over_elems_real(grid)[2:end]]
    #! format: on
    return Tridiagonal(a_, b_, c_)
end

"""
    solve_tdma!(q::StateVec, tendencies::StateVec, aux::StateVec, x::S, ρ::S, a_τp1::S, a_τ::S, K::S, grid::Grid{T}, Δt::T, i::Int=1)

Solve tri-diagonal system using Diffusion-implicit time-marching.
"""
function solve_tdma!(
    q::StateVec,
    tendencies::StateVec,
    aux::StateVec,
    x::S,
    ρ::S,
    a_τp1::S,
    a_τ::S,
    K::S,
    grid::Grid{T},
    Δt::T,
    i::Int = 1,
) where {S <: Symbol, T}
    B = [
        aux[a_τ, k, i] / q[a_τp1, k, i] * q[x, k, i] +
        Δt * tendencies[x, k, i] / (aux[ρ, k] * q[a_τp1, k, i])
        for k in over_elems_real(grid)
    ]
    A = TridiagonalMatrix(q, aux, Δt, ρ, a_τp1, K, grid, DifussionImplicit(), i)
    X = inv(A) * B
    assign_real!(q, x, grid, X, i)
end

@testset "∂_t T = K ΔT + 1, T = 0 ∈ ∂Ω, diffusion-explicit" begin
    domain_set = DomainSet(gm = 1)
    domain_subset = DomainSubSet(gm = true)
    K, maxiter, Δt = 1.0, 1000, 0.005
    grid = UniformGrid(0.0, 1.0, 10)
    q = StateVec(((:T, domain_subset),), grid, domain_set)
    aux = StateVec(((:ΔT, domain_subset),), grid, domain_set)
    rhs = deepcopy(q)
    for i in 1:maxiter
        for k in over_elems_real(grid)
            aux[:ΔT, k] = Δ_z(q[:T, Cut(k)], grid)
            rhs[:T, k] = K * aux[:ΔT, k] + 1
        end
        for k in over_elems(grid)
            q[:T, k] += Δt * rhs[:T, k]
        end
        apply_Dirichlet!(q, :T, grid, 0.0, Zmax())
        apply_Dirichlet!(q, :T, grid, 0.0, Zmin())
    end
    sol_analtyic = 1 / 2 * grid.zc - grid.zc .^ 2 / 2
    sol_error = [abs(q[:T, k] - sol_analtyic[k]) for k in over_elems_real(grid)]
    @test all(sol_error .< grid.Δz)
    # Skip plotting...
    # mkpath(heat_eq_dir)
    # nc = NetCDFWriter(joinpath(heat_eq_dir, "T"))
    # export_state(nc, grid, q)
    # rm(full_name(nc))
end

@testset "∂_t T = K ΔT, T={1, 0} ∈ {z_{min}, z_{max}}, diffusion-explicit" begin
    domain_set = DomainSet(gm = 1)
    domain_subset = DomainSubSet(gm = true)
    K, maxiter, Δt = 1.0, 1000, 0.002
    grid = UniformGrid(0.0, 1.0, 10)
    q = StateVec(((:T, domain_subset),), grid, domain_set)
    aux = StateVec(((:ΔT, domain_subset),), grid, domain_set)
    rhs = deepcopy(q)
    for i in 1:maxiter
        for k in over_elems_real(grid)
            aux[:ΔT, k] = Δ_z(q[:T, Cut(k)], grid)
            rhs[:T, k] = K * aux[:ΔT, k]
        end
        for k in over_elems(grid)
            q[:T, k] += Δt * rhs[:T, k]
        end
        apply_Dirichlet!(q, :T, grid, 1.0, Zmin())
        apply_Dirichlet!(q, :T, grid, 0.0, Zmax())
    end
    sol_analtyic = 1 .- grid.zc
    sol_error = [abs(q[:T, k] - sol_analtyic[k]) for k in over_elems_real(grid)]
    @test all(sol_error .< grid.Δz)
end

@testset "∂_t T = K ΔT, T = {1, 0} ∈ {z_{min}, z_{max}}, diffusion implicit" begin
    domain_set = DomainSet(gm = 1)
    domain_subset = DomainSubSet(gm = true)
    K, maxiter, Δt = 1.0, 1000, 0.1
    grid = UniformGrid(0.0, 1.0, 10)
    k_star1 = first_interior(grid, Zmin())
    k_star2 = first_interior(grid, Zmax())
    vars = ((:T, domain_subset), (:a, domain_subset))
    q = StateVec(vars, grid, domain_set)
    vars = (
        (:ΔT, domain_subset),
        (:a_tau, domain_subset),
        (:ρ_0, domain_subset),
        (:K, domain_subset),
        (:a, domain_subset),
    )
    aux = StateVec(vars, grid, domain_set)
    rhs = deepcopy(q)
    for k in over_elems(grid)
        q[:a, k] = 1.0
        aux[:ρ_0, k] = 1.0
        aux[:a_tau, k] = 1.0
        aux[:K, k] = 1.0
    end
    for i in 1:maxiter
        for k in over_elems_real(grid)
            rhs[:T, k] = 0.0
            if k == k_star1
                gv = 2 * 1 - q[:T, k]
                grad_T = ∇_z_flux([gv, 0], grid)
                bc_surf = ∇_z_flux([grad_T, 0], grid)
                rhs[:T, k] += bc_surf
            end
            if k == k_star2
                gv = -q[:T, k]
                grad_T = ∇_z_flux([0, gv], grid)
                bc_surf = ∇_z_flux([0, grad_T], grid)
                rhs[:T, k] += bc_surf
            end
        end
        solve_tdma!(q, rhs, aux, :T, :ρ_0, :a, :a_tau, :K, grid, Δt)
        extrap!(q, :T, grid)
    end
    sol_analtyic = 1 .- grid.zc
    sol_numerical = [q[:T, k] for k in over_elems_real(grid)]
    sol_error = [abs(q[:T, k] - sol_analtyic[k]) for k in over_elems_real(grid)]
    @test all(sol_error .< grid.Δz)
end

@testset "∂_t T = K ΔT, -K∇T = F ∈ z_{min}, T = 0 ∈ z_{max}, explicit Euler" begin
    domain_set = DomainSet(gm = 1)
    domain_subset = DomainSubSet(gm = true)
    K, maxiter, Δt = 1.0, 10, 0.001, 10
    grid = UniformGrid(0.0, 1.0, 10)
    k_star1 = first_interior(grid, Zmin())
    k_star2 = first_interior(grid, Zmax())
    q = StateVec(((:T1, domain_subset), (:T2, domain_subset)), grid, domain_set)
    aux = StateVec(
        ((:ΔT1, domain_subset), (:ΔT2, domain_subset)),
        grid,
        domain_set,
    )
    rhs = deepcopy(q)
    for i in 1:maxiter
        q_out = -K * 1
        q_in = -q_out
        for k in over_elems_real(grid)
            aux[:ΔT1, k] = Δ_z(q[:T1, Cut(k)], grid)
            aux[:ΔT2, k] = Δ_z(q[:T2, Cut(k)], grid)
            rhs[:T1, k] = K * aux[:ΔT1, k]
            rhs[:T2, k] = K * aux[:ΔT2, k]
            if k == k_star1 && i > 1
                rhs[:T2, k] += -∇_z_flux([q_in, 0.0], grid)
            end
        end
        for k in over_elems(grid)
            q[:T1, k] += Δt * rhs[:T1, k]
            q[:T2, k] += Δt * rhs[:T2, k]
        end
        apply_Neumann!(q, :T1, grid, q_out, Zmin())
        apply_Dirichlet!(q, :T1, grid, 0.0, Zmax())
        apply_Neumann!(q, :T2, grid, 0.0, Zmin())
        apply_Dirichlet!(q, :T2, grid, 0.0, Zmax())
        @test all([q[:T1, k] ≈ q[:T2, k] for k in over_elems_real(grid)])
    end
end

@testset "∂_t T = ∇ • (K(z) ∇T) + 1, T = 0 ∈ ∂Ω, explicit Euler" begin
    domain_set = DomainSet(gm = 1)
    domain_subset = DomainSubSet(gm = true)
    K, maxiter, Δt = 1.0, 1000, 0.001
    grid = UniformGrid(0.0, 1.0, 10)
    unknowns = ((:T, domain_subset),)
    vars = ((:ΔT, domain_subset), (:K_thermal, domain_subset))
    q = StateVec(unknowns, grid, domain_set)
    aux = StateVec(vars, grid, domain_set)
    rhs = deepcopy(q)
    cond_thermal(z) = z > 0.5 ? 1 : 0.1
    for k in over_elems(grid)
        aux[:K_thermal, k] = cond_thermal(grid.zc[k])
    end
    for i in 1:maxiter
        for k in over_elems_real(grid)
            aux[:ΔT, k] = Δ_z(q[:T, Cut(k)], grid, aux[:K_thermal, Cut(k)])
            rhs[:T, k] = K * aux[:ΔT, k] + 1
        end
        for k in over_elems(grid)
            q[:T, k] += Δt * rhs[:T, k]
        end
        apply_Dirichlet!(q, :T, grid, 0.0, Zmax())
        apply_Dirichlet!(q, :T, grid, 0.0, Zmin())
    end
end

@testset "∂_t T = ∇ • (K(z) ∇T) + 1, T = {0, 0} ∈ {z_{min}, z_{max}}, comparison" begin
    domain_set = DomainSet(gm = 1)
    domain_subset = DomainSubSet(gm = true)
    VS, TS, F, K = 1.0, 1.0, 1.0, 1.0
    maxiter, Δt = 5000, 0.005
    grid = UniformGrid(0.0, 1.0, 10)
    k_star1 = first_interior(grid, Zmin())
    k_star2 = first_interior(grid, Zmax())
    vars = (
        (:T_explicit_surf, domain_subset),
        (:T_explicit_vol, domain_subset),
        (:T_implicit, domain_subset),
        (:a, domain_subset),
    )
    q = StateVec(vars, grid, domain_set)
    vars = (
        (:ΔT, domain_subset),
        (:a_tau, domain_subset),
        (:ρ_0, domain_subset),
        (:K_thermal, domain_subset),
        (:a, domain_subset),
    )
    aux = StateVec(vars, grid, domain_set)
    rhs = deepcopy(q)
    for k in over_elems(grid)
        q[:a, k] = 1.0
        aux[:ρ_0, k] = 1.0
        aux[:a_tau, k] = 1.0
    end
    cond_thermal(z) = z > 0.5 ? 1 : 0.1
    for k in over_elems(grid)
        aux[:K_thermal, k] = cond_thermal(grid.zc[k])
    end
    for i in 1:maxiter
        for k in over_elems_real(grid)
            aux[:ΔT, k] =
                Δ_z(q[:T_explicit_surf, Cut(k)], grid, aux[:K_thermal, Cut(k)])
            rhs[:T_explicit_surf, k] = aux[:ΔT, k] + VS
            aux[:ΔT, k] =
                Δ_z(q[:T_explicit_vol, Cut(k)], grid, aux[:K_thermal, Cut(k)])
            rhs[:T_explicit_vol, k] = aux[:ΔT, k] + VS
            rhs[:T_implicit, k] = VS
            k == k_star1 && (
                rhs[:T_implicit, k] += bc_source(
                    q,
                    grid,
                    aux,
                    :T_implicit,
                    :ρ_0,
                    :a,
                    :K_thermal,
                    Zmin(),
                    TS,
                    Dirichlet(),
                    DiffusionAbsorbed(),
                )
            )
            k == k_star2 && (
                rhs[:T_implicit, k] += bc_source(
                    q,
                    grid,
                    aux,
                    :T_implicit,
                    :ρ_0,
                    :a,
                    :K_thermal,
                    Zmax(),
                    TS,
                    Dirichlet(),
                    DiffusionAbsorbed(),
                )
            )
            k == k_star1 && (
                rhs[:T_explicit_vol, k] += bc_source(
                    q,
                    grid,
                    aux,
                    :T_explicit_vol,
                    :ρ_0,
                    :a,
                    :K_thermal,
                    Zmin(),
                    TS,
                    Dirichlet(),
                    DiffusionAbsorbed(),
                )
            )
            k == k_star2 && (
                rhs[:T_explicit_vol, k] += bc_source(
                    q,
                    grid,
                    aux,
                    :T_explicit_vol,
                    :ρ_0,
                    :a,
                    :K_thermal,
                    Zmax(),
                    TS,
                    Dirichlet(),
                    DiffusionAbsorbed(),
                )
            )
        end

        for k in over_elems_real(grid)
            q[:T_explicit_surf, k] += Δt * rhs[:T_explicit_surf, k]
            q[:T_explicit_vol, k] += Δt * rhs[:T_explicit_vol, k]
        end
        solve_tdma!(
            q,
            rhs,
            aux,
            :T_implicit,
            :ρ_0,
            :a,
            :a_tau,
            :K_thermal,
            grid,
            Δt,
        )

        apply_Dirichlet!(q, :T_explicit_surf, grid, TS, Zmax())
        apply_Dirichlet!(q, :T_explicit_surf, grid, TS, Zmin())

        assign_ghost!(q, :T_explicit_vol, grid, 0.0, Zmin())
        assign_ghost!(q, :T_explicit_vol, grid, 0.0, Zmax())

        assign_ghost!(q, :T_implicit, grid, 0.0, Zmin())
        assign_ghost!(q, :T_implicit, grid, 0.0, Zmax())
    end
    extrap!(q, :T_explicit_vol, grid, TS, Zmin())
    extrap!(q, :T_explicit_vol, grid, TS, Zmax())
    extrap!(q, :T_implicit, grid, TS, Zmin())
    extrap!(q, :T_implicit, grid, TS, Zmax())

    @test q[:T_explicit_vol, Dual(k_star1)][1] ≈ TS
    @test q[:T_explicit_vol, Dual(k_star2)][2] ≈ TS
    @test all([
        q[:T_explicit_vol, k] ≈ q[:T_explicit_surf, k]
        for k in over_elems_real(grid)
    ])
    @test all([
        q[:T_implicit, k] ≈ q[:T_explicit_surf, k]
        for k in over_elems_real(grid)
    ])

end

@testset "∂_t T = ∇ • (K(z) ∇T) + VS, -K(z) ∇T = 1 ∈ z_{min}, T = 0 ∈ z_{max}, comparison" begin
    domain_set = DomainSet(gm = 1)
    domain_subset = DomainSubSet(gm = true)
    VS, TS, F = 1.0, 0.0, 2.0
    IC = 0.0
    maxiter, Δt = 10000, 0.0001
    grid = UniformGrid(0.0, 1.0, 20)
    k_star1 = first_interior(grid, Zmin())
    k_star2 = first_interior(grid, Zmax())
    vars = (
        (:T_explicit_surf, domain_subset),
        (:T_explicit_vol, domain_subset),
        (:T_implicit, domain_subset),
        (:a, domain_subset),
    )
    q = StateVec(vars, grid, domain_set)
    vars = (
        (:ΔT, domain_subset),
        (:a_tau, domain_subset),
        (:ρ_0, domain_subset),
        (:K_thermal, domain_subset),
    )
    aux = StateVec(vars, grid, domain_set)
    rhs = deepcopy(q)
    for k in over_elems(grid)
        q[:a, k] = 1.0
        aux[:ρ_0, k] = 1.0
        aux[:a_tau, k] = 1.0
    end
    cond_thermal(z) = 5 + 3 * cos(6 * π * z)
    for k in over_elems(grid)
        aux[:K_thermal, k] = cond_thermal(grid.zc[k])
        q[:T_explicit_surf, k] = IC
        q[:T_explicit_vol, k] = IC
        q[:T_implicit, k] = IC
    end
    K_surf = aux[:K_thermal, Dual(k_star1)][1]
    apply_Neumann!(q, :T_explicit_surf, grid, F / K_surf, Zmin())
    apply_Dirichlet!(q, :T_explicit_surf, grid, TS, Zmax())
    for i in 1:maxiter
        assign_ghost!(q, :T_explicit_vol, grid, 0.0, Zmin())
        assign_ghost!(q, :T_explicit_vol, grid, 0.0, Zmax())
        for k in over_elems_real(grid)
            aux[:ΔT, k] = Δ_z_dual(
                q[:T_explicit_surf, Cut(k)],
                grid,
                aux[:K_thermal, Dual(k)],
            )
            rhs[:T_explicit_surf, k] = aux[:ΔT, k] + VS

            aux[:ΔT, k] = Δ_z_dual(
                q[:T_explicit_vol, Cut(k)],
                grid,
                aux[:K_thermal, Dual(k)],
            )
            rhs[:T_explicit_vol, k] = aux[:ΔT, k] + VS
            k == k_star1 && (
                rhs[:T_explicit_vol, k] += bc_source(
                    q,
                    grid,
                    aux,
                    :T_explicit_vol,
                    :ρ_0,
                    :a,
                    :K_thermal,
                    Zmin(),
                    F,
                    Neumann(),
                    DiffusionAbsorbed(),
                )
            )
            k == k_star2 && (
                rhs[:T_explicit_vol, k] += bc_source(
                    q,
                    grid,
                    aux,
                    :T_explicit_vol,
                    :ρ_0,
                    :a,
                    :K_thermal,
                    Zmax(),
                    TS,
                    Dirichlet(),
                    DiffusionAbsorbed(),
                )
            )

            rhs[:T_implicit, k] = VS

            k == k_star1 && (
                rhs[:T_implicit, k] += bc_source(
                    q,
                    grid,
                    aux,
                    :T_implicit,
                    :ρ_0,
                    :a,
                    :K_thermal,
                    Zmin(),
                    F,
                    Neumann(),
                    DiffusionAbsorbed(),
                )
            )
            k == k_star2 && (
                rhs[:T_implicit, k] += bc_source(
                    q,
                    grid,
                    aux,
                    :T_implicit,
                    :ρ_0,
                    :a,
                    :K_thermal,
                    Zmax(),
                    TS,
                    Dirichlet(),
                    DiffusionAbsorbed(),
                )
            )
        end

        for k in over_elems_real(grid)
            q[:T_explicit_surf, k] += Δt * rhs[:T_explicit_surf, k]
            q[:T_explicit_vol, k] += Δt * rhs[:T_explicit_vol, k]
        end
        solve_tdma!(
            q,
            rhs,
            aux,
            :T_implicit,
            :ρ_0,
            :a,
            :a_tau,
            :K_thermal,
            grid,
            Δt,
        )

        apply_Neumann!(q, :T_explicit_surf, grid, F / K_surf, Zmin())
        apply_Dirichlet!(q, :T_explicit_surf, grid, TS, Zmax())

        apply_Neumann!(q, :T_explicit_vol, grid, F / K_surf, Zmin())
        apply_Dirichlet!(q, :T_explicit_vol, grid, TS, Zmax())

        apply_Neumann!(q, :T_implicit, grid, F / K_surf, Zmin())
        apply_Dirichlet!(q, :T_implicit, grid, TS, Zmax())
    end

    apply_Neumann!(q, :T_explicit_vol, grid, F / K_surf, Zmin())
    extrap!(q, :T_explicit_vol, grid, TS, Zmax())

    apply_Neumann!(q, :T_implicit, grid, F / K_surf, Zmin())
    extrap!(q, :T_implicit, grid, TS, Zmax())

    @test q[:T_explicit_vol, Dual(k_star2)][2] ≈ TS
    @test q[:T_implicit, Dual(k_star2)][2] ≈ TS

    @test all([
        abs(q[:T_explicit_vol, k] - q[:T_explicit_surf, k]) < 10 * Δt
        for k in over_elems(grid)
    ])
    @test all([
        abs(q[:T_explicit_vol, k] - q[:T_explicit_surf, k]) < 10 * Δt
        for k in over_elems(grid)
    ])
    @test all([
        abs(q[:T_implicit, k] - q[:T_explicit_surf, k]) < 10 * Δt
        for k in over_elems(grid)
    ])
end


text:
function prepareTravelTimeDataFiles(m,Minv::RegularMesh,mref,boundsHigh,boundsLow,filenamePrefix::String,pad::Int64,jump::Int64,offset::Int64,useFilesForFields::Bool = false)
### Here we generate data files for sources and receivers. This can be replaced with real sources/receivers files.
########################## m is in Velocity here. ###################################

RCVfile = string(filenamePrefix,"_rcvMap.dat");
SRCfile = string(filenamePrefix,"_srcMap.dat");
writeSrcRcvLocFile(SRCfile,Minv,pad,jump);
writeSrcRcvLocFile(RCVfile,Minv,pad,1);

dataFullFilename = string(filenamePrefix,"_travelTime");
HO = false;
prepareTravelTimeDataFiles(m, Minv, filenamePrefix,dataFullFilename,offset,HO,useFilesForFields);

file = matopen(string(filenamePrefix,"_PARAM.mat"), "w");
write(file,"MinvOmega",Minv.domain);
write(file,"boundsLow",boundsLow);
write(file,"boundsHigh",boundsHigh);
write(file,"mref",mref);
write(file,"MinvN",Minv.n);
write(file,"HO",HO);
close(file);
return;
end

function prepareTravelTimeDataFiles(m, Minv::RegularMesh, filenamePrefix::String,dataFullFilename::String, offset::Int64,HO::Bool,useFilesForFields::Bool = false)

########################## m is in Velocity here. ###################################

RCVfile = string(filenamePrefix,"_rcvMap.dat");
SRCfile = string(filenamePrefix,"_srcMap.dat");
srcNodeMap = readSrcRcvLocationFile(SRCfile,Minv);
rcvNodeMap = readSrcRcvLocationFile(RCVfile,Minv);

Q = generateSrcRcvProjOperators(Minv.n+1,srcNodeMap);
Q = Q.*(1/(norm(Minv.h)^2));
P = generateSrcRcvProjOperators(Minv.n+1,rcvNodeMap);

# compute observed data

println("~~~~~~~ Getting data Eikonal: ~~~~~~~");
(pForEIK,contDivEIK,SourcesSubIndEIK) = getEikonalInvParam(Minv,Q,P,HO,nworkers(),useFilesForFields);

(D,pForEIK) = getData(velocityToSlowSquared(m[:])[1],pForEIK,ones(length(pForEIK)),true);

Dobs = Array{Array{Float64,2}}(length(pForEIK))
for k = 1:length(pForEIK)
	Dobs[k] = fetch(D[k]);
end
Dobs = arrangeRemoteCallDataIntoLocalData(Dobs);

# D should be of length 1 becasue constMUSTBeOne = 1;
Dobs += 0.01*mean(abs.(Dobs))*randn(size(Dobs,1),size(Dobs,2));
Wd = (1.0./(abs.(Dobs)+ 0.1*mean(abs.(Dobs))));
Wd = limitDataToOffset(Wd,srcNodeMap,rcvNodeMap,offset);
writeDataFile(string(dataFullFilename,".dat"),Dobs,Wd,srcNodeMap,rcvNodeMap);

return (Dobs,Wd)
end


text:
# simulate.jl: functions to simulate the model
# include("setup.jl")
function evolve!( pop::Population, generations::Int64 = 1 )
    # from ReputationSets https://github.com/tkessinger/reputation_sets
    # the main function to evolve the population
    # generations specifies the total number of generations for which to run the simulation
    for i::Int64 in 1:generations
        if pop.verbose println("\ninitiating generation $(pop.generation)") end

        # first, update payoffs
        if pop.verbose println("updating actions and payoffs") end
        update_payoffs_using_matrix!(pop)  # method that stores payoffs in a matrix

        # then, select a learner and a role model and perform update
        if pop.verbose println("") end
        if pop.verbose println("evolving, generation $(pop.generation)") end
        update_strategies_and_opinions_db!(pop)

        if pop.verbose println("") end
        if pop.verbose println("ending generation $(pop.generation)") end
        pop.generation += 1
    end
end

function update_payoffs_using_matrix!( pop::Population )
    # update individual payoffs by computing pairwise payoffs and storing them in a matrix

    # to make the code run faster, we only update payoffs in a given generation
    # if strategies or sets have changed. sets_changed and strategies_changed are
    # initialized as true so that payoffs are computed in the first round
    if pop.sets_changed
        # if (previous) learner's sets have changed, re-compute all payoffs
        if pop.verbose println("updating all payoffs because set_changed = $(pop.sets_changed)") end
        pop.payoffs_mat  .= zero(Float64) # reset all payoffs_mat
        pop.prev_actions .= zero(Int64)   # reset all prev_actions

        # then, recompute all payoffs
        for k in 1:pop.sets.M  # for each set
            for (i, j) in pop.sets.set_pairs[k]   # for each pair within each set
                compute_payoffs!( pop, i, j, k )  # compute payoffs for i & j when intereacting in set k
                                                  # this also updates pop.prev_actions accordingly
            end
        end

    elseif pop.strategies_changed
        # if (previous) learner's strategy has changed, recompute payoffs involving that individual
        if pop.verbose println("updating payoffs involving $(pop.prev_learner) because strategies_changed = $(pop.strategies_changed)") end
        pop.payoffs_mat[pop.prev_learner,:]    .= zero(Float64) # first, reset payoffs involving prev_learner
        pop.payoffs_mat[:,pop.prev_learner]    .= zero(Float64)
        pop.prev_actions[:,pop.prev_learner,:] .= zero(Int64)
        pop.prev_actions[:,:,pop.prev_learner] .= zero(Int64)

        # then, recompute payoffs involving prev_learner
        for k in 1:pop.sets.M  # for each set
            for (i, j) in pop.sets.set_pairs[k]  # for each pair within each set
                if i == pop.prev_learner || j == pop.prev_learner
                    compute_payoffs!( pop, i, j, k )  # compute payoffs for i & j when intereacting in set k
                                                      # this also updates pop.prev_actions accordingly
                end
            end
        end
    else
        if pop.verbose println("not updating payoffs b/c sets_changed = $(pop.sets_changed) and strategies_changed = $(pop.strategies_changed)") end
        return # no need to update prev_interactions or payoffs
    end

    # sum rows to get total individual payoffs
    pop.payoffs = [round( x, digits=13 ) for x in vec( sum(pop.payoffs_mat, dims=2) )]  # rounding avoids precision errors
    pop.prev_interactions = 2 * sum([length(x) for x in pop.sets.set_pairs])

    if pop.verbose println( "\npayoff matrix is now $(pop.payoffs_mat)" ) end
    if pop.verbose println( "payoffs are now $(pop.payoffs)" ) end
    if pop.verbose println( "there were $(pop.prev_interactions) interactions, $(sum(pop.prev_actions)) of which were C" ) end
end

function compute_payoffs!( pop::Population, i::Int64, j::Int64, k::Int64 )
    # function to compute payoffs for individuals i and j when they interact in set k
    pop.sim.i_action, pop.sim.j_action = -1, -1  # set to -1 to catch any errors in action assignment

    # actions depend on whether i and j agree on issue k or not
    if pop.sets.h[i,k] == 0 || pop.sets.h[j,k] == 0
        # if either i or j does not belong to set k, something's wrong
        # leave pop.sim.i_action, pop.sim.j_action as -1 so that the compute_payoffs! throws an error when assigning payoffs
        @warn("$i and $j should not be interacting in set $k"); return
    elseif pop.sets.h[i,k] == pop.sets.h[j,k]
        pop.sim.i_action, pop.sim.j_action = pop.strategies[i,1], pop.strategies[j,1] # if they agree, play s_agree
    else
        pop.sim.i_action, pop.sim.j_action = pop.strategies[i,2], pop.strategies[j,2] # if they disagree, play s_disagree
    end

    if pop.verbose
        println("\n\tupdating actions of $i and $j in set $k; agree on issue $k? $(pop.sets.h[i,k] == pop.sets.h[j,k])")
        println("\t\t$i has strats $(pop.strategies[i,:]); $i plays strat $(pop.sim.i_action), earns a payoff of $(pop.game.A[pop.sim.i_action+1, pop.sim.j_action+1])")
        println("\t\t$j has strats $(pop.strategies[j,:]); $j plays strat $(pop.sim.j_action), earns a payoff of $(pop.game.A[pop.sim.j_action+1, pop.sim.i_action+1])")
    end

    # add i and j's payoffs to the payoff matrix according to their strategies and the game matrix
    pop.payoffs_mat[i,j] += pop.game.A[pop.sim.i_action+1, pop.sim.j_action+1] # +1 because julia is 1-indexed
    pop.payoffs_mat[j,i] += pop.game.A[pop.sim.j_action+1, pop.sim.i_action+1]

    # track actions
    pop.prev_actions[k,i,j] = pop.sim.i_action
    pop.prev_actions[k,j,i] = pop.sim.j_action
end

function update_strategies_and_opinions_db!( pop::Population )
    # function to perform death-birth update
    if pop.verbose println("Strategies before updating:\t $(pop.strategies)") end
    if pop.verbose println("Set memberships before updating: $(pop.sets.h)") end

    # randomly choose a learner and store a list of non_learners
    pop.sim.learner = sample( 1:pop.sets.N )

    # compute the fitness of every individual in the population
    for i in 1:pop.sets.N pop.sim.fitnesses[i] = 1.0 + pop.game.β * pop.payoffs[i] end
    if any(x->x<0.0, pop.sim.fitnesses) throw("!!! negative fitness detected, aborting !!!"); end

    # choose a role model from the rest of the population, weighted by fitness
    # OLD: pop.sim.role_model = sample( pop.sim.non_learners, Weights(pop.sim.fitnesses) )
    pop.sim.role_model = sample( 1:pop.sets.N, Weights(pop.sim.fitnesses) )

    # learner updates
    if pop.verbose println("\n\trandomly chosen learner is $(pop.sim.learner), role_model is $(pop.sim.role_model)") end
    if pop.verbose println("\t\tall payoffs are $(pop.payoffs)") end
    if pop.verbose println("\t\tfitnesses are $(pop.sim.fitnesses)") end

    # compute probability of imitation: 1 if same party, 1-p if different parties
    pop.sets.affils[pop.sim.learner] == pop.sets.affils[pop.sim.role_model] ? imit_prob = 1.0 : imit_prob = 1-pop.game.ps[pop.sets.affils[pop.sim.learner]]

    if pop.verbose
        println("\t\tlearner and role model in the same party? $(pop.sets.affils[ pop.sim.learner ] == pop.sets.affils[ pop.sim.role_model ]), imitate with prob $(imit_prob)")
    end

    # update strategies and sets: imitation / mutation occurs with probability imit_prob
    if rand() < imit_prob
        if pop.verbose println("\twill consider imitating") end
        update_strategies!( pop ) # includes mutation
        update_opinions!( pop )   # includes mutation
        pop.num_imitating += 1
    else
        if pop.verbose println("\tnah won't imitate") end
        no_strategy_or_opinion_update!( pop )
        pop.num_not_imitating += 1
    end

    # track updates
    pop.prev_learner = pop.sim.learner
    if pop.sets_changed == true
        # h is updated above; need to update h_bar, h_base10, h_bar_base10, set_members, set_pairs
        pop.sets.set_members, pop.sets.set_pairs = set_members_and_pairs(pop.sets.h)  # update set_members and set_pairs
        pop.sets.h_bar        = [abs(x) for x in pop.sets.h]
        pop.sets.h_base10     = [vectobase3(pop.sets.h[i,:]) for i in 1:pop.sets.N]
        pop.sets.h_bar_base10 = [vectobase3(pop.sets.h_bar[i,:]) for i in 1:pop.sets.N]
        pop.num_sets_changed += 1
    end
    if pop.strategies_changed == true
        # strategies are updated above; need to update strategies_base10
        pop.strategies_base10 = [2 * pop.strategies[i,1] + pop.strategies[i,2] for i in 1:pop.sets.N]
        pop.num_strategies_changed += 1
    end
    if pop.strategies_changed && pop.sets_changed
        pop.num_both_changed += 1
    end

    if pop.verbose println("\nStrategies after updating:\t $(pop.strategies)") end
    if pop.verbose println("Set memberships after updating:\t $(pop.sets.h)") end
    if pop.verbose
        println("tracking: prev_learner = $(pop.prev_learner), strategies_changed = $(pop.strategies_changed), sets_changed = $(pop.sets_changed)")
    end
end

function update_strategies!( pop::Population )
    # function to update strategy of learner based on role_model
    pop.sim.strategy_copy = copy(pop.strategies[pop.sim.learner,:])  # save a copy to track changes
    # copy suffices here because the assignments below replace elements of pop.sets.h with Int64 (immutable)

    if rand() < pop.game.u
        # strategy mutation occurs
        pop.strategies[pop.sim.learner,:] = rand(collect(0:1), 1, 2) # conditional case
        pop.num_random_strategies += 1
        pop.num_random_sa_flipped += (pop.sim.strategy_copy[1] != pop.strategies[pop.sim.learner,1])
        pop.num_random_sd_flipped += (pop.sim.strategy_copy[2] != pop.strategies[pop.sim.learner,2])
        if pop.verbose println("\t\tmutating $(pop.sim.learner) to strategy $(pop.strategies[pop.sim.learner,:])") end
    else
        # learner imitates the strategy of role model
        pop.strategies[pop.sim.learner,:] = pop.strategies[pop.sim.role_model,:]
        if pop.verbose println("\t\tadopting $(pop.sim.role_model)'s strategy $(pop.strategies[pop.sim.role_model,:])") end
    end
    pop.strategies_changed = (pop.sim.strategy_copy != pop.strategies[pop.sim.learner,:]) # track changes
end

function update_opinions!( pop::Population )
    # function to update set memberships / opinions of learner based on role_model
    pop.sim.sets_copy = copy(pop.sets.h[pop.sim.learner,:])     # save a copy to track changes
    # copy suffices here because the assignments below replace elements of pop.sets.h with Int64 (immutable)

    # with probability v, set mutation occurs
    if rand() < pop.game.v
        if rand() < 1 - pop.game.ϵ * pop.game.ps[pop.sets.affils[pop.sim.learner]]
            # with probability 1-p*ϵ, select opinions without bias
            # select a random arrangement of K opinions
            pop.sets.h[pop.sim.learner,:] = shuffle( vcat( zeros(Int64, pop.sets.M-pop.sets.K), rand([-1,1], pop.sets.K) ) )

            pop.num_random_opinions += 1
            pop.num_random_opinions_flipped += (pop.sim.sets_copy != pop.sets.h[ pop.sim.learner,: ]) # new
            if pop.verbose println("\t\tmutating $(pop.sim.learner) to random opinions $(pop.sets.h[pop.sim.learner,:])") end
        else
            # with probability p*ϵ, select opinions with bias
            # select a random set of K issues, select opinions corresponding to party
            if pop.sets.affils[pop.sim.learner] == 1
                pop.sets.h[pop.sim.learner,:] = shuffle( vcat( zeros(Int64, pop.sets.M-pop.sets.K), -ones(Int64, pop.sets.K) ) )
            elseif pop.sets.affils[pop.sim.learner] == 2
                pop.sets.h[pop.sim.learner,:] = shuffle( vcat( zeros(Int64, pop.sets.M-pop.sets.K), ones(Int64, pop.sets.K) ) )
            end
            pop.num_biased_opinions += 1
            pop.num_biased_opinions_flipped += (pop.sim.sets_copy != pop.sets.h[ pop.sim.learner,: ]) # new
            if pop.verbose
                println("\t\t$(pop.sim.learner) is in party $(pop.sets.affils[ pop.sim.learner ]), mutating to biased opinions $(pop.sets.h[ pop.sim.learner,: ])")
            end
        end
    else
        # with probability 1-v, learner imitates the strategy of role model
        pop.sets.h[ pop.sim.learner,: ] = pop.sets.h[ pop.sim.role_model,: ]
        if pop.verbose println("\t\tadopting $( pop.sim.role_model )'s set membership $(pop.sets.h[ pop.sim.role_model,: ])") end
    end
    # h is updated above; need to update h_bar, sets_changed, set_members, set_pairs
    # pop.sets.set_members, pop.sets.set_pairs = set_members_and_pairs(pop.sets.h)  # update set_members and set_pairs # moved up
    pop.sets_changed = (pop.sim.sets_copy != pop.sets.h[ pop.sim.learner,: ]) # track changes
end

function no_strategy_or_opinion_update!( pop::Population )
    # function to keep track of changes when no imitation/mutation occurs
    pop.strategies_changed = false  # track changes
    pop.sets_changed       = false
end