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

using Markdown
using InteractiveUtils

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

# ╔═╡ 489b13da-7d23-11eb-39de-57a12469721f
using Plots, LaTeXStrings

# ╔═╡ 82c5386a-7d2f-11eb-3992-e985ee16933c
using ColorSchemes

# ╔═╡ 4bd45848-7d23-11eb-1a8b-b36f7461f368
begin
	import WGLMakie
	using JSServe, PlutoUI
end

# ╔═╡ fa6e4914-7d24-11eb-0ff2-1974e4ec8df2
using StatsBase, Statistics, LinearAlgebra

# ╔═╡ 2ef712f2-7d2e-11eb-228c-2d023dc6ab10
using Random

# ╔═╡ a1b684a0-7d40-11eb-2a4c-e990f12c6cbb
using Distributions

# ╔═╡ 4256e210-7d23-11eb-2f7a-1d8aceab0682
function ingredients(path::String)
	# this is from the Julia source code (evalfile in base/loading.jl)
	# but with the modification that it returns the module instead of the last object
	name = Symbol(basename(path))
	m = Module(name)
	Core.eval(m,
        Expr(:toplevel,
             :(eval(x) = $(Expr(:core, :eval))($name, x)),
             :(include(x) = $(Expr(:top, :include))($name, x)),
             :(include(mapexpr::Function, x) = $(Expr(:top, :include))(mapexpr, $name, x)),
             :(include($path))))
	m
end

# ╔═╡ 3d2bf622-7d23-11eb-33fb-3f8a4f0f309e
function makie(f::Function)
    scene = WGLMakie.Scene(resolution = (600, 400))
	f(scene)

    return scene
end

# ╔═╡ 5e1fa2de-7d23-11eb-18da-ddf0ec5714cc
M = ingredients("../src/main.jl")

# ╔═╡ a0c3a118-7d3b-11eb-0309-a9540887541f
I = ingredients("../src/infer.jl")

# ╔═╡ fce24760-7d22-11eb-2474-a5efdd176182
scrna, genes = M.SeqSpace.expression();

# ╔═╡ 330fd846-7d25-11eb-3933-b5f8871eedc5
size(scrna)

# ╔═╡ 4ddffd5a-7d24-11eb-2c70-675569787f53
"""
This section generates figures for linear dimensional reduction
"""

# ╔═╡ 72a6c262-7d25-11eb-1a05-756c421933f5
pvals(matrix) = (λ=reverse(eigvals(cov(matrix,dims=2)))[1:1297])/sum(λ)

# ╔═╡ 85650856-7d24-11eb-055d-75a99e549182
function sample_from!(null, matrix)
	@assert size(null) == size(matrix)
	for g ∈ 1:size(null,1)
		sample!(matrix[g,:], view(null,g,:); replace=false)
	end
end

# ╔═╡ c3fdc5ce-7d26-11eb-235a-33e432adcd91
function null_pvals(matrix; niter=10)
	null = zeros(size(matrix))
	λ 	 = zeros(minimum(size(matrix)), niter)
	for i ∈ 1:niter
		sample_from!(null, matrix)
		λ[:,i] = pvals(null)
	end

	return λ
end

# ╔═╡ 7bf6e4e4-7d27-11eb-0d11-a343a5ead1fe
λ = (
	data = pvals(scrna),
	null = null_pvals(scrna),
);

# ╔═╡ 564474de-7d2a-11eb-107a-094fc5b153f6
λ.null

# ╔═╡ de3c3612-7d25-11eb-3233-393941f590e5
begin
	plot(λ.data[1:end-1], label="data", color=:cyan3, linewidth=3)
	for i ∈ 1:size(λ.null,2)
		plot!(λ.null[1:end-1,i], label="", color=:orangered2, alpha=0.5, linewidth=1)
	end
	plot!(λ.null[1:end-1,1], label="null", color=:orangered2, alpha=1, linewidth=1)

	xaxis!("principal value", :log10)
	yaxis!("fraction of variance", :log10)

	#savefig("../figs/principal_values.svg")
end

# ╔═╡ 59156b62-7d2c-11eb-1f14-3139307e312e
"""
This section generates figures analyzing state localization
"""

# ╔═╡ 713cebde-7d2c-11eb-3918-31d392278150
ipr(ψ) = sum(ψ.^4, dims=1) ./ (sum(ψ.^2, dims=1)).^2

# ╔═╡ 151d45be-7d2d-11eb-00ae-495dff2916f2
function scan_ipr(x, Ns; niter=5)
	IPR = zeros(size(x,2), niter, length(Ns))
	for (n,N) ∈ enumerate(Ns)
		for i ∈ 1:niter
			ι = randperm(size(x,1))[1:N]
			IPR[:,i,n] = ipr(x[ι,:])
		end
	end
	IPR
end

# ╔═╡ f940b7ee-7d2d-11eb-3ecf-930c0eb1b938
F = svd(scrna);

# ╔═╡ 1ab88ad2-7d2e-11eb-2327-3f2015c5220c
Ns = [100,250,500,750,1000,2000,3000,4000,5000,6000,7000,8000,size(F.U,1)]; IPR = scan_ipr(F.U, Ns; niter=50);

# ╔═╡ 476449e0-7d2e-11eb-2d91-11e5a358d060
ĪP̄R̄ = dropdims(mean(IPR,dims=2);dims=2)

# ╔═╡ 4a089ff2-7d2e-11eb-0e3b-9d580d220c45
begin
	cmap = ColorSchemes.plasma
	plot(Ns, 1 ./ ĪP̄R̄[1,:],color=get(cmap,0.0),linewidth=2,label="")
	for i in 2:49
		plot!(Ns, 1 ./ ĪP̄R̄[i,:],color=get(cmap,(i-1)/50),linewidth=2,label="")
	end
	plot!(Ns, 1 ./ ĪP̄R̄[50,:],color=get(cmap,1.0),linewidth=2,label="")

	xaxis!("number of genes subsampled", :log10)
	yaxis!(L"1 / \sum \psi^4", :log10)

	#savefig("../figs/participation_ratio_subsample_scaling.svg")
end

# ╔═╡ 49dc2956-7d31-11eb-03c8-f5e29d9d0798
ψ = ipr(F.U)

# ╔═╡ 58a0f57a-7d31-11eb-3b59-293c95b4aa07
begin
	plot(1 ./ ψ')

	xaxis!("principal component", :log10)
	yaxis!("participation ratio", :log10)

	#savefig("../figs/participation_ratio.svg")
end

# ╔═╡ b499f6bc-7d34-11eb-106d-53126cd7a5c6
"""
This section generates figures analyzing point cloud scaling.
"""

# ╔═╡ c115efae-7d34-11eb-0d58-bba15e9922bb
G = F.U[:, 1:25] * Diagonal(F.S[1:25]) * F.Vt[1:25,:];

# ╔═╡ 1a663bc4-7d35-11eb-20e2-99f298281450
Dₑ = M.SeqSpace.PointCloud.distance(G);

# ╔═╡ 3838406e-7d35-11eb-31fb-256eea2a72cb
Dₚ = M.SeqSpace.PointCloud.geodesics(G, 6);

# ╔═╡ 629180aa-7d35-11eb-2bd8-47deefc64f91
function scaling(D, N)
	Rₘᵢₙ = minimum(D[D .> 0])
	Rₘₐₓ = maximum(D)
	Rs   = range(Rₘᵢₙ,Rₘₐₓ,length=N)

	ϕ = zeros(size(D,1), N)
	for (i,R) ∈ enumerate(Rs)
		ϕ[:,i] = sum(D .<= R, dims=1)
	end

	return ϕ, Rs
end

# ╔═╡ 4c497cfc-7d36-11eb-013a-07b8d183909f
ρₑ, Rₑ = scaling(Dₑ, 1000);

# ╔═╡ f3da8bec-7d37-11eb-06c3-e58f9312c561
ρₚ, Rₚ = scaling(Dₚ, 1000);

# ╔═╡ 61689ce4-7d36-11eb-206f-a362c5f1b24b
begin
	plot(Rₑ, ρₑ', alpha=0.03, color=:cyan3, label="", title="Euclidean")
	plot!(Rₑ, mean(ρₑ, dims=1)', linewidth=5, alpha=1, color=:cyan3, label="mean")
	plot!(Rₑ, 5e-10*Rₑ.^7, linestyle=:dashdot, linewidth=2, alpha=1, color=:black, label=L"\alpha R^7")

	xaxis!("radius", :log10, (25, 400))
	yaxis!("number of points", :log10, (10, 1297))

	#savefig("../figs/euclidean_ball_scaling.svg")
end

# ╔═╡ f1e88e9a-7d37-11eb-0529-a3ab42f90864
begin
	plot(Rₚ, ρₚ', alpha=0.02, color=:cyan3, label="", title="Geodesic")
	plot!(Rₚ, mean(ρₚ, dims=1)', linewidth=5, alpha=1, color=:cyan3, label="mean")
	plot!(Rₚ, 9e-5*Rₚ.^3.2, linestyle=:dashdot, linewidth=2, alpha=1, color=:black, label=L"\alpha R^3")

	xaxis!("radius", :log10, (25, 400))
	yaxis!("number of points", :log10, (10, 1297))

	#savefig("../figs/geodesic_ball_scaling.svg")
end

# ╔═╡ e2a9cb1c-7d3a-11eb-088a-7d51b807d714
"""
This section generates figures analyzing isomap
"""

# ╔═╡ 8c321e78-7d3b-11eb-375b-5dd081921104
invert, embryo = I.Inference.inversion(); Ψ = invert(0.5);

# ╔═╡ a7c62f6e-7d3b-11eb-1ef0-f3743d0e58e1
begin
	ψᵣ = Ψ  ./ sum(Ψ, dims=2)
	ψₗ = (Ψ ./ sum(Ψ, dims=1))'
end;

# ╔═╡ b0bc877e-7d3b-11eb-1110-913567d545da
r̂ = (ψₗ * embryo)'; AP = r̂[1,:]; DV = atan.(r̂[2,:], r̂[3,:])

# ╔═╡ ed76e2e6-7d3a-11eb-04ed-69f2bdd9f68a
ξ = M.SeqSpace.PointCloud.mds(Dₚ.^2, 10)

# ╔═╡ 0c2a0f80-7d3e-11eb-17b6-71185c47896d
begin

residual = zeros(size(ξ,2))
for d ∈ 1:size(ξ,2)
	D̄ = .√M.SeqSpace.PointCloud.distance²(ξ[:,1:d]')

	residual[d] = cor(M.SeqSpace.PointCloud.upper_tri(D̄), M.SeqSpace.PointCloud.upper_tri(Dₚ))
end

end;

# ╔═╡ 694eb1e4-7d42-11eb-0ac6-8fe4f3896971
begin
	d̄ = M.SeqSpace.PointCloud.upper_tri(.√M.SeqSpace.PointCloud.distance²(ξ[:,1:6]'));
	d = M.SeqSpace.PointCloud.upper_tri(Dₚ)

	ι = sample(1:length(d), 5000)

	scatter(d[ι], d̄[ι], legend=false)

	xaxis!("estimated geodesic distance")
	yaxis!("pairwise euclidean distance in latent space")
	#savefig("../figs/distance_isomap_correlation_eg_d3.svg")
end

# ╔═╡ 8cc4e066-7d3e-11eb-37b7-312b8d2a4a73
begin
	plot(residual, linewidth=2, label="")

	xaxis!("latent dimension")
	yaxis!("correlation of distance matrices", (0.7,0.95))

	#savefig("../figs/distance_correlation_vs_isomap_dimension.svg")
end

# ╔═╡ 43b9ec00-7d3d-11eb-2973-cbb9af2197a8
@bind THETA PlutoUI.Slider(range(0,90,length=100))

# ╔═╡ 6cca0e84-7d3d-11eb-1a00-77473b588de7
@bind PHI PlutoUI.Slider(range(0,90,length=100))

# ╔═╡ ffdcaae6-7d3b-11eb-3ddb-558d8e1638d6
begin
	scatter(ξ[:,2],-ξ[:,1],ξ[:,3],
		marker_z=AP,
		markersize=5,
		xlims=(-150,150),
		ylims=(-150,150),
		zlims=(-125,75),
		title="predicted AP",
		camera=(THETA,PHI))

	#savefig("../figs/isomap_3d_AP.svg")
end

# ╔═╡ a3fd873e-7d3d-11eb-391e-cfc5d3c80797
begin
	scatter(ξ[:,2],-ξ[:,1],ξ[:,3],
		marker_z=DV,
		markersize=5,
		xlims=(-150,150),
		ylims=(-150,150),
		zlims=(-125,75),
		title="predicted DV",
		camera=(THETA,PHI))

	#savefig("../figs/isomap_3d_DV.svg")
end

# ╔═╡ 66270b84-7d46-11eb-2e09-c97d61f6425b
begin
scatter(F.Vt[3,:],F.Vt[4,:],F.Vt[5,:],
		marker_z=AP,
		markersize=5,
		title="predicted AP",
		camera=(THETA,PHI)
)

#savefig("../figs/svd_3d_AP.svg")

end

# ╔═╡ 927cd290-7d46-11eb-057c-c14cffce4ce5
begin

scatter(F.Vt[2,:],F.Vt[3,:],F.Vt[4,:],
		marker_z=DV,
		markersize=5,
		title="predicted DV",
		camera=(THETA,PHI)
)

#savefig("../figs/svd_3d_DV.svg")

end

# ╔═╡ d1c09536-7d46-11eb-2d39-79a39782580b
Dₓ = M.SeqSpace.PointCloud.geodesics(r̂,6)

# ╔═╡ ead3da76-7d46-11eb-0d7c-cb4496b2c27a
begin
	dx = M.SeqSpace.PointCloud.upper_tri(Dₓ)
	dp = M.SeqSpace.PointCloud.upper_tri(Dₚ)

	ι₂ = sample(1:length(dx), 5000)

	scatter(dx[ι₂], dp[ι₂])
	xlabel!("spatial geodesic distance")
	ylabel!("expression geodesic distance")

	#savefig("../figs/spatial_vs_expression_geodesic.svg")
end

# ╔═╡ Cell order:
# ╠═489b13da-7d23-11eb-39de-57a12469721f
# ╠═82c5386a-7d2f-11eb-3992-e985ee16933c
# ╠═4bd45848-7d23-11eb-1a8b-b36f7461f368
# ╠═4256e210-7d23-11eb-2f7a-1d8aceab0682
# ╠═3d2bf622-7d23-11eb-33fb-3f8a4f0f309e
# ╠═5e1fa2de-7d23-11eb-18da-ddf0ec5714cc
# ╠═a0c3a118-7d3b-11eb-0309-a9540887541f
# ╠═fce24760-7d22-11eb-2474-a5efdd176182
# ╠═330fd846-7d25-11eb-3933-b5f8871eedc5
# ╠═4ddffd5a-7d24-11eb-2c70-675569787f53
# ╠═fa6e4914-7d24-11eb-0ff2-1974e4ec8df2
# ╠═72a6c262-7d25-11eb-1a05-756c421933f5
# ╠═85650856-7d24-11eb-055d-75a99e549182
# ╠═c3fdc5ce-7d26-11eb-235a-33e432adcd91
# ╠═7bf6e4e4-7d27-11eb-0d11-a343a5ead1fe
# ╠═564474de-7d2a-11eb-107a-094fc5b153f6
# ╠═de3c3612-7d25-11eb-3233-393941f590e5
# ╠═59156b62-7d2c-11eb-1f14-3139307e312e
# ╠═713cebde-7d2c-11eb-3918-31d392278150
# ╠═151d45be-7d2d-11eb-00ae-495dff2916f2
# ╠═f940b7ee-7d2d-11eb-3ecf-930c0eb1b938
# ╠═2ef712f2-7d2e-11eb-228c-2d023dc6ab10
# ╠═1ab88ad2-7d2e-11eb-2327-3f2015c5220c
# ╠═476449e0-7d2e-11eb-2d91-11e5a358d060
# ╠═4a089ff2-7d2e-11eb-0e3b-9d580d220c45
# ╠═49dc2956-7d31-11eb-03c8-f5e29d9d0798
# ╠═58a0f57a-7d31-11eb-3b59-293c95b4aa07
# ╠═b499f6bc-7d34-11eb-106d-53126cd7a5c6
# ╠═c115efae-7d34-11eb-0d58-bba15e9922bb
# ╠═1a663bc4-7d35-11eb-20e2-99f298281450
# ╠═3838406e-7d35-11eb-31fb-256eea2a72cb
# ╠═629180aa-7d35-11eb-2bd8-47deefc64f91
# ╠═4c497cfc-7d36-11eb-013a-07b8d183909f
# ╠═f3da8bec-7d37-11eb-06c3-e58f9312c561
# ╠═61689ce4-7d36-11eb-206f-a362c5f1b24b
# ╠═f1e88e9a-7d37-11eb-0529-a3ab42f90864
# ╠═e2a9cb1c-7d3a-11eb-088a-7d51b807d714
# ╠═8c321e78-7d3b-11eb-375b-5dd081921104
# ╠═a7c62f6e-7d3b-11eb-1ef0-f3743d0e58e1
# ╠═b0bc877e-7d3b-11eb-1110-913567d545da
# ╠═ed76e2e6-7d3a-11eb-04ed-69f2bdd9f68a
# ╠═a1b684a0-7d40-11eb-2a4c-e990f12c6cbb
# ╠═0c2a0f80-7d3e-11eb-17b6-71185c47896d
# ╠═694eb1e4-7d42-11eb-0ac6-8fe4f3896971
# ╠═8cc4e066-7d3e-11eb-37b7-312b8d2a4a73
# ╠═43b9ec00-7d3d-11eb-2973-cbb9af2197a8
# ╠═6cca0e84-7d3d-11eb-1a00-77473b588de7
# ╠═ffdcaae6-7d3b-11eb-3ddb-558d8e1638d6
# ╠═a3fd873e-7d3d-11eb-391e-cfc5d3c80797
# ╠═66270b84-7d46-11eb-2e09-c97d61f6425b
# ╠═927cd290-7d46-11eb-057c-c14cffce4ce5
# ╠═d1c09536-7d46-11eb-2d39-79a39782580b
# ╠═ead3da76-7d46-11eb-0d7c-cb4496b2c27a

max_stars_repo_name: nnoll/seqspace

max_issues_repo_name: nnoll/seqspace

max_forks_repo_name: nnoll/seqspace


text:
function difftensor(data::AbstractArray, debug_info, backs::Pair...)
    # filter out not required AD rules.
    mask = [requires_grad(pair.first) for pair in backs]
    t = DiffTensor(data, any(mask))
    push!(GLOBAL_TAPE.instructs, Instruction(debug_info, t,backs[mask]...))
    return t
end
requires_grad(t::DiffTensor) = t.requires_grad
requires_grad(t::Tuple) = any(x->requires_grad(x), t)

for EC in [:DynamicEinCode, :StaticEinCode]
    @eval function OMEinsum.einsum(code::$EC, @nospecialize(xs::NTuple{N,DiffTensor} where N), size_dict::Dict)
        y = einsum(code, getdata.(xs), size_dict)
        return difftensor(y, debug_info(einsum, code, xs, size_dict), ntuple(i->(xs[i]=>
            dy->OMEinsum.einsum_grad(OMEinsum.getixs(code), xs, OMEinsum.getiy(code), size_dict, conj(dy), i)
            ), length(xs))...)
    end
end

function Base.:(+)(x::DiffTensor, y::DiffTensor)
    difftensor(x.data + y.data, debug_info(+, x, y), x=>projectto(x), y=>projectto(y))
end
function Base.:(-)(x::DiffTensor, y::DiffTensor)
    difftensor(x.data - y.data, debug_info(-, x, y), x=>projectto(x), y=>dz->projectto(y,-dz))
end
function Base.:(-)(x::DiffTensor)
    difftensor(-x.data, debug_info(-, x), x=>(-))
end
function Base.broadcasted(::typeof(/), x::DiffTensor{T1,N}, y::DiffTensor{T2,N}) where {T1,T2,N}
    res = asarray(x.data ./ y.data, x.data)
    difftensor(res, debug_info("./", x, y), x=>dz->projectto(x, dz ./ conj(y)), y=>dz->projectto(y, -dz .* conj(x ./ y .^ 2)))
end
function Base.broadcasted(::typeof(inv), x::DiffTensor{T1,N}) where {T1,N}
    res = asarray(inv.(x.data), x.data)
    difftensor(res, debug_info("inv.", x), x=>dz->projectto(x, -dz .* conj(inv.(x) .^ 2)))
end
function Base.broadcasted(::typeof(^), x::DiffTensor, y::Number)
    res = asarray(x.data .^ y, x.data)
    difftensor(res, debug_info(".^", x, y), x=>dz-> dz .* conj(y * x .^ (y-1)))
end
# literal and variables are called into different functions.
function Base.broadcasted(::typeof(Base.literal_pow), ::typeof(^), x::DiffTensor, ::Val{y}) where y
    res = asarray(Base.broadcast(Base.literal_pow, ^, x.data, Val(y)), x.data)
    difftensor(res, debug_info(".^", x, y), x=>dz-> dz .* conj(y * x .^ (y-1)))
end
function Base.broadcasted(::typeof(sqrt), x::DiffTensor)
    res = asarray(sqrt.(x.data), x.data)
    difftensor(res, debug_info("sqrt.", x), x=>dz-> 0.5 * dz ./ conj(sqrt.(x)))
end
function Base.copy(x::DiffTensor)
    difftensor(copy(x.data), debug_info(copy, x), x=>identity)
end
function Base.getindex(x::DiffTensor, index1::Int, indices::Int...)
    error("get element is forbidden!")
end
function Base.getindex(x::DiffTensor, indices::Union{Int,AbstractRange,Colon}...)
    difftensor(Base.getindex(x.data, indices...), debug_info(getindex, x, indices...),
        x=>dy->accum(zero(x), indices, dy))
end

function accum(x::AbstractArray, indices, y::AbstractArray)
    z = copy(x)
    z[indices...] .+= y
    return z
end
function accum(x::DiffTensor, indices, y::DiffTensor)
    difftensor(accum(x.data, indices, y.data),
        debug_info(accum, x, indices, y), x=>projectto(x), y=>dz->projectto(y, dz[indices...]))
end
function Base.transpose(x::DiffTensor)
    difftensor(projectto(x.data, transpose(x.data)),
        debug_info(transpose, x), x=>transpose)
end

function Base.reshape(x::DiffTensor, sz::NTuple{N,Int}) where N
    size0 = size(x)
    difftensor(reshape(x.data, sz), debug_info(reshape, x, sz),
        x=>dy->reshape(dy, size0))
end

function Base.:(*)(a::Number, b::DiffTensor{T,N}) where {T,N}
    res = asarray(a * b.data, b.data)
    difftensor(res, debug_info("*", a, b), b=>dy->projectto(b, dy * conj(a)))
end
Base.:(*)(a::DiffTensor{T,N}, b::Number) where {T,N} = b * a
Base.:(*)(a::DiffTensor{T1,2}, b::DiffTensor{T2,2}) where {T1,T2} = ein"ij,jk->ik"(a, b)

# scalar-array multiply
function Base.:(*)(a::DiffTensor{T1,0}, b::DiffTensor{T2,N}) where {T1,T2,N}
    res = asarray(a.data[] * b.data, b.data)
    difftensor(res, debug_info("*", a, b), a=>dy->projectto(a, sum(dy .* conj(b))), b=>dy->projectto(b, dy * conj(a)))
end
# array-scalar multiply
Base.:(*)(a::DiffTensor{T1,N}, b::DiffTensor{T2,0}) where {T1,T2,N} = b * a
# scalar multiply
function Base.:(*)(a::DiffTensor{T1,0}, b::DiffTensor{T2,0}) where {T1,T2}
    difftensor(asarray(a.data[] * b.data[], a.data), debug_info("*", a, b), a=>dz->projectto(a, conj(b) * dz), b=>dz->projectto(b, conj(a) * dz))
end

# sum and fill
function Base.sum(a::DiffTensor{T,N}) where {T,N}
    difftensor(asarray(sum(a.data), a.data), debug_info("sum", a), a=>dz->fill(dz, size(a)...))
end
function Base.fill(a::DiffTensor{T,0}, size::Union{Integer, AbstractUnitRange}...) where {T}
    difftensor(fill(a.data[], size...), debug_info("fill", a), a=>dz->sum(dz))
end

function Base.broadcasted(::typeof(*), a::DiffTensor{T1,N}, b::DiffTensor{T2,N}) where {T1,T2,N}
    res = asarray(a.data .* b.data, a.data)
    difftensor(res, debug_info(".*", a, b), a=>dy->dy .* conj(b), b=>dy->dy .* conj(a))
end
function Base.broadcasted(::typeof(sin), a::DiffTensor{T}) where T
    res = asarray(sin.(a.data), a.data)
    difftensor(res, debug_info("sin.", a), a=>dy -> dy .* conj(cos.(a)))
end
function Base.broadcasted(::typeof(cos), a::DiffTensor{T}) where T
    res = asarray(cos.(a.data), a.data)
    difftensor(res, debug_info("cos.", a), a=>dy -> -dy .* conj(sin.(a)))
end
function Base.cat(a::DiffTensor{T}, B::DiffTensor{T}...; dims) where T
    A = (a, B...)
    ends = ntuple(i->cumsum(size.(A, i)), ndims(a))
    difftensor(cat(getdata.(A)...; dims=dims),
        debug_info(cat, a, B...; dims),
        ntuple(i->A[i]=>dy->getindex(dy, ntuple(idim->idim ∈ dims ? ((i==1 ? 1 : ends[idim][i-1]+1):ends[idim][i]) : (:), ndims(dy))...), length(A))...
    )
end
function Base.vcat(a::DiffTensor{T, 1}...) where T
    cat(a...; dims=1)
end
function Base.hcat(a::DiffTensor{T, 1}...) where T
    cat(a...; dims=2)
end
function projectto(::DiffTensor{<:Real}, x::DiffTensor{<:Complex})
    return real(x)
end
function projectto(::DiffTensor{<:Real}, x::DiffTensor{<:Real})
    return x
end
function projectto(::DiffTensor{<:Complex}, x::DiffTensor{<:Real})
    return Complex.(x)
end
function projectto(::DiffTensor{<:Complex}, x::DiffTensor{<:Complex})
    return x
end
projectto(x::DiffTensor) = y->projectto(x, y)
function projectto(x::AbstractArray, y::LinearAlgebra.Transpose)
    return typeof(x)(y)
end

### COMPLEX ###
function Base.real(x::DiffTensor)
    difftensor(real(x.data), debug_info(real, x), x=>projectto(x))
end
function Base.imag(x::DiffTensor{T}) where T<:Complex
    difftensor(imag(x.data), debug_info(imag, x), x=>dz->(im)*dz)
end
function Base.conj(x::DiffTensor)
    difftensor(conj(x.data), debug_info(conj, x), x=>conj)
end
function Base.broadcasted(::typeof(abs), x::DiffTensor)
    res = asarray(abs.(x.data), x.data)
    difftensor(res, debug_info("abs.", x), x=>dy->sign.(x) .* dy)
end
function Base.broadcasted(::typeof(sign), x::DiffTensor)
    x ./ abs.(x)
end
Base.broadcasted(::Type{T}, x::DiffTensor{<:Complex}) where T<:Complex = x
function Base.broadcasted(::Type{T}, x::DiffTensor{<:Real}) where T<:Complex
    res = asarray(T.(x.data), x.data)
    difftensor(res, debug_info("Complex.", x), x=>dy->real(dy))
end

##### maps
Base.map(::OMEinsum.ProjectTo{T}, x::DiffTensor) where T<:Complex = T.(x)
Base.map(::OMEinsum.ProjectTo{T}, x::DiffTensor) where T<:Real = real(x)

# TODO
# complex.
# accum!

ext: jl


text:
using Test
using NQCBase

atoms = Atoms{Float64}([:C, :H])
@test atoms isa Atoms{Float64}
@test length(atoms) == 2
@test atoms.numbers == [6, 1]
@test atoms.types == [:C, :H]
@test atoms.masses isa Vector

@test range(atoms) == 1:2

@test atoms[1] isa Atoms{Float64}
@test atoms[1:2] isa Atoms{Float64}

@test Atoms(2000) isa Atoms{Float64}


text:
module LinearSolversTests

using Gridap

using LinearAlgebra
using SparseArrays

ls = LUSolver()

diff1(M) = [ [1.0 zeros(1,M-1)]; diagm(1=>ones(M-1)) - Matrix(1.0I,M,M) ]

sdiff1(M) = sparse(diff1(M))

function Laplacian(Nx, Ny, Lx, Ly)
   dx = Lx / (Nx+1)
   dy = Ly / (Ny+1)
   Dx = sdiff1(Nx) / dx
   Dy = sdiff1(Ny) / dy
   Ax = Dx' * Dx
   Ay = Dy' * Dy
   return kron(sparse(1.0I,Ny,Ny), Ax) + kron(Ay, sparse(1.0I,Nx,Nx))
end

n = 10
A = Laplacian(n,n,1,1)
x = rand(n^2)
b = A*x

test_linear_solver(ls,A,b,x)

n = 10
A = Laplacian(n,n,1,1)
x = rand(n^2)
b = A*x

ls = BackslashSolver()

test_linear_solver(ls,A,b,x)

end # module LinearSolversTests


text:
"""
    classnames() -> Vector{String}

Return the 10 names for the Fashion-MNIST classes as a vector of
strings.
"""
classnames() = CLASSES

"""
    traintensor([T = N0f8], [indices]; [dir]) -> Array{T}

Returns the Fashion-MNIST **training** images corresponding to
the given `indices` as a multi-dimensional array of eltype `T`.

The image(s) is/are returned in the native horizontal-major
memory layout as a single numeric array. If `T <: Integer`, then
all values will be within `0` and `255`, otherwise the values are
scaled to be between `0` and `1`.

If the parameter `indices` is omitted or an `AbstractVector`, the
images are returned as a 3D array (i.e. a `Array{T,3}`), in which
the first dimension corresponds to the pixel *rows* (x) of the
image, the second dimension to the pixel *columns* (y) of the
image, and the third dimension denotes the index of the image.

```julia-repl
julia> FashionMNIST.traintensor() # load all training images
28×28×60000 Array{N0f8,3}:
[...]

julia> FashionMNIST.traintensor(Float32, 1:3) # first three images as Float32
28×28×3 Array{Float32,3}:
[...]
```

If `indices` is an `Integer`, the single image is returned as
`Matrix{T}` in horizontal-major layout, which means that the
first dimension denotes the pixel *rows* (x), and the second
dimension denotes the pixel *columns* (y) of the image.

```julia-repl
julia> FashionMNIST.traintensor(1) # load first training image
28×28 Array{N0f8,2}:
[...]
```

As mentioned above, the images are returned in the native
horizontal-major layout to preserve the original feature
ordering. You can use the utility function
[`convert2image`](@ref) to convert an FashionMNIST array into a
vertical-major Julia image with the corrected color values.

```julia-repl
julia> FashionMNIST.convert2image(FashionMNIST.traintensor(1)) # convert to column-major colorant array
28×28 Array{Gray{N0f8},2}:
[...]
```

$(download_docstring("FashionMNIST", DEPNAME))
"""
function traintensor(::Type{T}, args...; dir = nothing) where T
    path = datafile(DEPNAME, TRAINIMAGES, dir)
    images = Reader.readimages(path, args...)
    bytes_to_type(T, images)
end

function traintensor(args...; dir = nothing)
    traintensor(N0f8, args...; dir = dir)
end

"""
    testtensor([T = N0f8], [indices]; [dir]) -> Array{T}

Returns the Fashion-MNIST **test** images corresponding to the
given `indices` as a multi-dimensional array of eltype `T`.

The image(s) is/are returned in the native horizontal-major
memory layout as a single numeric array. If `T <: Integer`, then
all values will be within `0` and `255`, otherwise the values are
scaled to be between `0` and `1`.

If the parameter `indices` is omitted or an `AbstractVector`, the
images are returned as a 3D array (i.e. a `Array{T,3}`), in which
the first dimension corresponds to the pixel *rows* (x) of the
image, the second dimension to the pixel *columns* (y) of the
image, and the third dimension denotes the index of the image.

```julia-repl
julia> FashionMNIST.testtensor() # load all test images
28×28×10000 Array{N0f8,3}:
[...]

julia> FashionMNIST.testtensor(Float32, 1:3) # first three images as Float32
28×28×3 Array{Float32,3}:
[...]
```

If `indices` is an `Integer`, the single image is returned as
`Matrix{T}` in horizontal-major layout, which means that the
first dimension denotes the pixel *rows* (x), and the second
dimension denotes the pixel *columns* (y) of the image.

```julia-repl
julia> FashionMNIST.testtensor(1) # load first test image
28×28 Array{N0f8,2}:
[...]
```

As mentioned above, the images are returned in the native
horizontal-major layout to preserve the original feature
ordering. You can use the utility function
[`convert2image`](@ref) to convert an FashionMNIST array into a
vertical-major Julia image with the corrected color values.

```julia-repl
julia> FashionMNIST.convert2image(FashionMNIST.testtensor(1)) # convert to column-major colorant array
28×28 Array{Gray{N0f8},2}:
[...]
```

$(download_docstring("FashionMNIST", DEPNAME))
"""
function testtensor(::Type{T}, args...; dir = nothing) where T
    path = datafile(DEPNAME, TESTIMAGES, dir)
    images = Reader.readimages(path, args...)
    bytes_to_type(T, images)
end

function testtensor(args...; dir = nothing)
    testtensor(N0f8, args...; dir = dir)
end

"""
    trainlabels([indices]; [dir])

Returns the Fashion-MNIST **trainset** labels corresponding to
the given `indices` as an `Int` or `Vector{Int}`. The values of
the labels denote the zero-based class-index that they represent
(see [`FashionMNIST.classnames`](@ref) for the corresponding
names). If `indices` is omitted, all labels are returned.

```julia-repl
julia> FashionMNIST.trainlabels() # full training set
60000-element Array{Int64,1}:
 9
 0
 ⋮
 0
 5

julia> FashionMNIST.trainlabels(1:3) # first three labels
3-element Array{Int64,1}:
 9
 0
 0

julia> y = FashionMNIST.trainlabels(1) # first label
9

julia> FashionMNIST.classnames()[y + 1] # corresponding name
"Ankle boot"
```

$(download_docstring("FashionMNIST", DEPNAME))
"""
function trainlabels(args...; dir = nothing)
    path = datafile(DEPNAME, TRAINLABELS, dir)
    Vector{Int}(Reader.readlabels(path, args...))
end

function trainlabels(index::Integer; dir = nothing)
    path = datafile(DEPNAME, TRAINLABELS, dir)
    Int(Reader.readlabels(path, index))
end

"""
    testlabels([indices]; [dir])

Returns the Fashion-MNIST **testset** labels corresponding to the
given `indices` as an `Int` or `Vector{Int}`. The values of the
labels denote the class-index that they represent (see
[`FashionMNIST.classnames`](@ref) for the corresponding names).
If `indices` is omitted, all labels are returned.

```julia-repl
julia> FashionMNIST.testlabels() # full test set
10000-element Array{Int64,1}:
 9
 2
 ⋮
 1
 5

julia> FashionMNIST.testlabels(1:3) # first three labels
3-element Array{Int64,1}:
 9
 2
 1

julia> y = FashionMNIST.testlabels(1) # first label
9

julia> FashionMNIST.classnames()[y + 1] # corresponding name
"Ankle boot"
```

$(download_docstring("FashionMNIST", DEPNAME))
"""
function testlabels(args...; dir = nothing)
    path = datafile(DEPNAME, TESTLABELS, dir)
    Vector{Int}(Reader.readlabels(path, args...))
end

function testlabels(index::Integer; dir = nothing)
    path = datafile(DEPNAME, TESTLABELS, dir)
    Int(Reader.readlabels(path, index))
end

"""
    traindata([T = N0f8], [indices]; [dir]) -> images, labels

Returns the Fashion-MNIST **trainingset** corresponding to the
given `indices` as a two-element tuple. If `indices` is omitted
the full trainingset is returned. The first element of the
return values will be the images as a multi-dimensional array,
and the second element the corresponding labels as integers.

The image(s) is/are returned in the native horizontal-major
memory layout as a single numeric array of eltype `T`. If `T <:
Integer`, then all values will be within `0` and `255`, otherwise
the values are scaled to be between `0` and `1`. The integer
values of the labels correspond 1-to-1 the digit that they
represent.

```julia
train_x, train_y = FashionMNIST.traindata() # full datatset
train_x, train_y = FashionMNIST.traindata(2) # only second observation
train_x, train_y = FashionMNIST.traindata(dir="./FashionMNIST") # custom folder
```

$(download_docstring("FashionMNIST", DEPNAME))

Take a look at [`FashionMNIST.traintensor`](@ref) and
[`FashionMNIST.trainlabels`](@ref) for more information.
"""
function traindata(::Type{T}, args...; dir = nothing) where T
    (traintensor(T, args...; dir = dir),
     trainlabels(args...; dir = dir))
end

traindata(args...; dir = nothing) = traindata(N0f8, args...; dir = dir)

"""
    testdata([T = N0f8], [indices]; [dir]) -> images, labels

Returns the Fashion-MNIST **testset** corresponding to the given
`indices` as a two-element tuple. If `indices` is omitted the
full testset is returned. The first element of the return values
will be the images as a multi-dimensional array, and the second
element the corresponding labels as integers.

The image(s) is/are returned in the native horizontal-major
memory layout as a single numeric array of eltype `T`. If `T <:
Integer`, then all values will be within `0` and `255`, otherwise
the values are scaled to be between `0` and `1`. The integer
values of the labels correspond 1-to-1 the digit that they
represent.

```julia
test_x, test_y = FashionMNIST.testdata() # full datatset
test_x, test_y = FashionMNIST.testdata(2) # only second observation
test_x, test_y = FashionMNIST.testdata(dir="./FashionMNIST") # custom folder
```

$(download_docstring("FashionMNIST", DEPNAME))

Take a look at [`FashionMNIST.testtensor`](@ref) and
[`FashionMNIST.testlabels`](@ref) for more information.
"""
function testdata(::Type{T}, args...; dir = nothing) where T
    (testtensor(T, args...; dir = dir),
     testlabels(args...; dir = dir))
end

testdata(args...; dir = nothing) = testdata(N0f8, args...; dir = dir)


text:
using Pkg
Pkg.activate(@__DIR__)
using Soss
using ZigZagBoomerang

m = @model x begin
    α ~ SpikeMixture(Normal(), 0.3) # 0.2*Normal() + 0.8*Dirac(0)
    β ~ SpikeMixture(Normal(), 0.3)
    yhat = α .+ β .* x
    y ~ For(eachindex(x)) do j
        Normal(yhat[j], 2.0)
    end
end

vars = (α=1, β=2)

x = randn(20);
obs = -0.1 .+ 2x + 1randn(20);
T = 200.0
posterior = m(x=x) | (y=obs,)

bm_ = Soss.basemeasure(posterior, (β=0.0, α=0.0));
bm = getindex.(Ref(bm_.data), (keys(vars)));
@show kappa.(bm)

include("sparsezigzag.jl")
trace, final, (num, acc) = @time sparse_zigzag(posterior, T, c=50)

ts, xs = ZigZagBoomerang.sep(discretize(trace, 0.1))

xs = xform(posterior).(xs)

using Plots
p = plot(ts, getindex.(xs, 1))
plot!(ts, getindex.(xs, 2), color=:red)
savefig("tracesticky.png")
p2 = plot(first.(xs), last.(xs))
savefig("phasesticky.png")
p


text:
# module optimization
# TODO:
# (1) make a wrapper for Optim package

export optimize!, StochasticGradient, RProp

abstract Optimizer

@doc doc"Loop over parameters and call the optimizer for each parameter set with non null gradient" ->
function optimize!(s::Optimizer, m::Layer)
  for v in parameters(m)
    if !isnull(v.gradient)
        optimize!(s, v)
    end
  end
end

@doc doc"Stochatic gradient" ->
type StochasticGradient <: Optimizer
    learningRate::Float64
    # learningRateDecay::Float64

    # weithDecay::Float64
    # momentum::Float64

    # @doc doc"dampening for momentum" ->
    # dampening::Float64

    # @doc doc"Whether we use Nesterov momentum"
    # nesterov::Bool

    StochasticGradient(learningRate) = new(learningRate)
end

function optimize!(s::StochasticGradient, p::ArrayParameters)
    # TODO : Implement momentum
    axpy!(-s.learningRate, get(p.gradient), p.values)
end

#
# RProp optimizer
#

@doc doc"RProp optimizer
- stepsize    : initial step size, common to all components
- etaplus     : multiplicative increase factor, > 1 (default 1.2)
- etaminus    : multiplicative decrease factor, < 1 (default 0.5)
- stepsizemax : maximum stepsize allowed (default 50)
- stepsizemin : minimum stepsize allowed (default 1e-6)
- niter       : number of iterations (default 1)
" ->
immutable RProp <: Optimizer
    stepsize::Float64

    etaplus::Float64
    etaminus::Float64

    stepsizemax::Float64
    stepsizemin::Float64

    function RProp()
        self = new()
        self.stepsize = 0.1
        self.etaplus = 1.2
        self.etaminus = 0.5
        self.stepsizemax = 50.0
        self.stepsizemin = 1e-6
        self
    end

    RProp(stepsize) = new(stepsize, 1.2, 0.5, 50., 1e-6)
    RProp() = new(0.1)
end

type RPropState
    delta
    stepsize

    sign::BitArray{2}

    function RPropState(l::Int, stepsize::Float64)
        self = new()
        self.delta = zeros(l)
        self.stepsize = fill(stepsize, l)
        self.sign = BitArray(l, 2)
        self
    end
end

const RProp_POSITIVE = bitpack([false, false])
const RProp_NEGATIVE = bitpack([false, true])
const RProp_ZERO = bitpack([true, true])

function optimize!(s::RProp, p::ArrayParameters)
    # initialize auxiliary storage
    if isnull(p.optimization_state)
        p.optimization_state = Nullable(RPropState(length(p.values), s.stepsize))
    end

    state = get(p.optimization_state)::RPropState
    gradient = get(p.gradient)

    @inbounds for i in eachindex(gradient)
        # Compute the new step size
        sign = gradient[i] > 0 ? RProp_POSITIVE : (gradient[i] < 0 ? RProp_NEGATIVE : RProp_ZERO)

        if sign != RProp_ZERO
            if sign == state.sign[i]
                state.stepsize[i] = min(state.stepsize[i] * s.etaplus, s.stepsizemax)
            else
                state.stepsize[i] = max(state.stepsize[i] * s.etaminus, s.stepsizemin)
            end
        end

        # Update weight
        p.values[i] -= state.stepsize[i] * gradient[i]
    end

end


text:
# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.6.2
#     language: julia
#     name: julia-1.6
# ---

# %% [markdown]
# # Examples of marginal plots
#
# Extended version of https://discourse.julialang.org/t/multiple-plots-some-refinements/64800/5

# %%
using Plots
using Distributions
using KernelDensity
using LinearAlgebra

# true distribution of sample
dist_true = MvNormal([1.0, 2.0], [2.0 -1.0; -1.0 4.0])
mu_x, mu_y = mean(dist_true)
s_x, s_y = .√diag(cov(dist_true))
distx_true, disty_true = Normal(mu_x, s_x), Normal(mu_y, s_y)

# generate test sample
n = 2^10
sample = rand(dist_true, n)

# kernel density estimation
X, Y = sample[1, :], sample[2, :]
k, kx, ky = kde.(((X, Y), X, Y))
ik, ikx, iky = InterpKDE.((k, kx, ky));

# %% [markdown]
# ## Sample marginal plots

# %%
layout = @layout [
    a             _
    b{0.8w, 0.8h} c
]

xlim, ylim = extrema.((k.x, k.y))
legend = colorbar = false
a = plot(kx.x, kx.density; xlim, legend)
c = plot(ky.density, ky.x; ylim, legend, xrotation=90)
b = contour(k.x, k.y, k.density; xlim, ylim, legend)

plot(a, b, c; layout, link=:both, size=(500, 500))

# %%
layout = @layout [
    a             _
    b{0.8w, 0.8h} c
]

xlim, ylim = extrema.((k.x, k.y))
legend = colorbar = false
a = plot(kx.x, kx.density; xlim, legend)
c = plot(ky.density, ky.x; ylim, legend, xrotation=90)
b = contourf(k.x, k.y, k.density; xlim, ylim, legend)

plot(a, b, c; layout, link=:both, size=(500, 500))

# %%
layout = @layout [
    a             _
    b{0.8w, 0.8h} c
]

xlim, ylim = extrema.((k.x, k.y))
legend = colorbar = false
a = plot(kx.x, kx.density; xlim, legend)
c = plot(ky.density, ky.x; ylim, legend, xrotation=90)
b = heatmap(k.x, k.y, k.density; xlim, ylim, legend)

plot(a, b, c; layout, link=:both, size=(500, 500))

# %%
layout = @layout [
    a             _
    b{0.8w, 0.8h} c
]

xlim, ylim = extrema.((k.x, k.y))
legend = colorbar = false
a = plot(kx.x, kx.density; xlim, legend)
c = plot(ky.density, ky.x; ylim, legend, xrotation=90)
b = scatter(X, Y; xlim, ylim, legend, marker_z=pdf.(Ref(ik), X, Y), alpha=0.7, msw=0)

plot(a, b, c; layout, link=:both, size=(500, 500))

# %%
layout = @layout [
    a             _
    b{0.8w, 0.8h} c
]

xlim, ylim = extrema.((k.x, k.y))
legend = colorbar = false
a = plot(kx.x, kx.density; xlim, legend)
c = plot(ky.density, ky.x; ylim, legend, xrotation=90)
b = scatter(X, Y; xlim, ylim, legend, marker_z=pdf.(Ref(ik), X, Y), alpha=0.7, msw=0, color=:rainbow)

plot(a, b, c; layout, link=:both, size=(500, 500))

# %% [markdown]
# ## True distribution marginal plots

# %%
layout = @layout [
    a             _
    b{0.8w, 0.8h} c
]

xlim, ylim = extrema.((k.x, k.y))
legend = colorbar = false
a = plot(kx.x, pdf.(distx_true, kx.x); xlim, label="true")
c = plot(pdf.(disty_true, ky.x), ky.x; ylim, xrotation=90, label="true")
b = contour(k.x, k.y, (x, y)->pdf(dist_true, [x, y]); xlim, ylim, legend)

plot(a, b, c; layout, link=:both, size=(500, 500))

# %%
layout = @layout [
    a             _
    b{0.8w, 0.8h} c
]

xlim, ylim = extrema.((k.x, k.y))
legend = colorbar = false
a = plot(kx.x, pdf.(distx_true, kx.x); xlim, label="true")
c = plot(pdf.(disty_true, ky.x), ky.x; ylim, xrotation=90, label="true")
b = contourf(k.x, k.y, (x, y)->pdf(dist_true, [x, y]); xlim, ylim, legend)

plot(a, b, c; layout, link=:both, size=(500, 500))

# %%
layout = @layout [
    a             _
    b{0.8w, 0.8h} c
]

xlim, ylim = extrema.((k.x, k.y))
legend = colorbar = false
a = plot(kx.x, pdf.(distx_true, kx.x); xlim, label="true")
c = plot(pdf.(disty_true, ky.x), ky.x; ylim, xrotation=90, label="true")
b = heatmap(k.x, k.y, (x, y)->pdf(dist_true, [x, y]); xlim, ylim, legend)

plot(a, b, c; layout, link=:both, size=(500, 500))

# %% [markdown]
# ## Other examples

# %%
layout = @layout [
    a             _
    b{0.8w, 0.8h} c
]

xlim, ylim = extrema.((k.x, k.y))
legend = colorbar = false

a = plot(kx.x, kx.density; xlim, label="X")
plot!(kx.x, pdf.(distx_true, kx.x); label="true", ls=:dash)

c = plot(ky.density, ky.x; ylim, xrotation=90, label="Y")
plot!(pdf.(disty_true, ky.x), ky.x; label="true", ls=:dash)

b = contour(k.x, k.y, k.density; xlim, ylim, legend)

plot(a, b, c; layout, link=:both, size=(500, 500))

# %%
layout = @layout [
    a             _
    b{0.8w, 0.8h} c
]

xlim, ylim = extrema.((k.x, k.y))
legend = colorbar = false

a = plot(kx.x, kx.density; xlim, label="X")
plot!(kx.x, pdf.(distx_true, kx.x); label="true", ls=:dash)

c = plot(ky.density, ky.x; ylim, xrotation=90, label="Y")
plot!(pdf.(disty_true, ky.x), ky.x; label="true", ls=:dash)

b = contourf(k.x, k.y, k.density; xlim, ylim, legend)

plot(a, b, c; layout, link=:both, size=(500, 500))

# %%
layout = @layout [
    a             _
    b{0.8w, 0.8h} c
]

xlim, ylim = extrema.((k.x, k.y))
legend = colorbar = false

a = plot(kx.x, kx.density; xlim, label="X")
plot!(kx.x, pdf.(distx_true, kx.x); label="true", ls=:dash)

c = plot(ky.density, ky.x; ylim, xrotation=90, label="Y")
plot!(pdf.(disty_true, ky.x), ky.x; label="true", ls=:dash)

b = heatmap(k.x, k.y, k.density; xlim, ylim, legend)

plot(a, b, c; layout, link=:both, size=(500, 500))

# %%
layout = @layout [
    a             _
    b{0.8w, 0.8h} c
]

xlim, ylim = extrema.((k.x, k.y))
legend = colorbar = false

a = plot(kx.x, kx.density; xlim, label="X")
plot!(kx.x, pdf.(distx_true, kx.x); label="true", ls=:dash)

c = plot(ky.density, ky.x; ylim, xrotation=90, label="Y")
plot!(pdf.(disty_true, ky.x), ky.x; label="true", ls=:dash)

b = scatter(X, Y; xlim, ylim, colorbar, marker_z=pdf.(Ref(ik), X, Y), alpha=0.7, msw=0, label="(X, Y)")
contour!(k.x, k.y, (x, y)->pdf(dist_true, [x, y]); xlim, ylim, legend, ls=:dash, color=2)

plot(a, b, c; layout, link=:both, size=(500, 500))

# %%
layout = @layout [
    a             _
    b{0.8w, 0.8h} c
]

xlim, ylim = extrema.((k.x, k.y))
legend = colorbar = false

a = plot(kx.x, kx.density; xlim, label="X")
plot!(kx.x, pdf.(distx_true, kx.x); label="true", ls=:dash)

c = plot(ky.density, ky.x; ylim, xrotation=90, label="Y")
plot!(pdf.(disty_true, ky.x), ky.x; label="true", ls=:dash)

b = scatter(X, Y; xlim, ylim, colorbar, marker_z=pdf.(Ref(ik), X, Y), alpha=0.7, msw=0, color=:rainbow, label="(X, Y)")
contour!(k.x, k.y, (x, y)->pdf(dist_true, [x, y]); xlim, ylim, legend, ls=:dash, color=2)

plot(a, b, c; layout, link=:both, size=(500, 500))

# %%

max_stars_repo_path: 0012/marginal plots.jl

max_stars_repo_name: genkuroki/public

max_issues_repo_path: 0012/marginal plots.jl

max_issues_repo_name: genkuroki/public

max_forks_repo_path: 0012/marginal plots.jl

max_forks_repo_name: genkuroki/public


text:
function test_pareto_dominance_func(pareto_dominance_func)
  test("simple cases with float vector fitness values of length 1") do
    @check pareto_dominance_func([0.0], [1.0])
    @check pareto_dominance_func([-100.0], [100.0])
    @check !pareto_dominance_func([-2.5], [-2.5])
    @check !pareto_dominance_func([12.32], [12.32])
  end

  test("simple cases with int vector fitness values of length 2") do
    @check pareto_dominance_func([0, 1], [1, 2])
    @check pareto_dominance_func([-10, -5], [0, 0])
    @check pareto_dominance_func([0, 1], [0, 2])
    @check pareto_dominance_func([-1, 2], [0, 2])

    @check !pareto_dominance_func([0, 1], [0, 1])
  end
end

describe("pareto_dominance") do
  describe("pareto_dominates_clear") do
    test_pareto_dominance_func(pareto_dominates_clear)
  end

  describe("pareto_dominates_fast") do
    test_pareto_dominance_func(pareto_dominates_fast)
  end
end


text:
import ..Flux: Flux, relu
using CuArrays.CUDAnative

CuRNN{T} = Flux.RNNCell{<:Union{typeof(tanh),typeof(relu)},<:CuArray{T,2},<:CuArray{T,1}}
CuGRU{T} = Flux.GRUCell{<:CuArray{T,2},<:CuArray{T,1}}
CuLSTM{T} = Flux.LSTMCell{<:CuArray{T,2},<:CuArray{T,1}}
CuRNNs{T} = Union{CuRNN{T},CuGRU{T},CuLSTM{T}}

function CUDNN.RNNDesc(m::CuRNNs{T}) where T
  h, i = length(m.h), size(m.Wi, 2)
  mode = m isa CuRNN ?
    (m.σ == tanh ? CUDNN.CUDNN_RNN_TANH : CUDNN.CUDNN_RNN_RELU) :
    m isa CuGRU ? CUDNN.CUDNN_GRU : CUDNN.CUDNN_LSTM
  r = CUDNN.RNNDesc{T}(mode, i, h)
  return r
end

const descs = WeakKeyDict()

function desc(rnn)
  d = haskey(descs, rnn) ? descs[rnn] : (descs[rnn] = CUDNN.RNNDesc(rnn))
  CUDNN.setweights!(d, rnn.Wi, rnn.Wh, rnn.b)
  return d
end

import Zygote
using Zygote: @adjoint

function (m::CuRNN{T})(h::CuArray{T}, x::CuArray{T}) where T <: Union{Float32,Float64}
  y, h′ = CUDNN.forward(desc(m), x, h)
  return h′, y
end

function (m::CuGRU{T})(h::CuArray{T}, x::CuArray{T}) where T <: Union{Float32,Float64}
  y, h′ = CUDNN.forward(desc(m), x, h)
  return h′, y
end

function (m::CuLSTM{T})(h::NTuple{2,CuArray{T}}, x::CuArray{T}) where T <: Union{Float32,Float64}
  y, h′, c′ = CUDNN.forward(desc(m), x, h[1], h[2])
  return (h′, c′), y
end

(m::CuRNN{T})(h::CuArray{T}, x) where T <: Union{Float32,Float64} = m(h, CuArray{T}(x))
(m::CuGRU{T})(h::CuArray{T}, x) where T <: Union{Float32,Float64} = m(h, CuArray{T}(x))
(m::CuLSTM{T})(h::NTuple{2,CuArray{T}}, x) where T <: Union{Float32,Float64} = m(h, CuArray{T}(x))

trim(x, Δ) = reshape(Δ, ntuple(i -> size(Δ, i), Val(ndims(x))))

unbroadcast(x::AbstractArray, Δ) =
  size(x) == size(Δ) ? Δ :
  length(x) == length(Δ) ? trim(x, Δ) :
    trim(x, sum(Δ, dims = ntuple(i -> size(x, i) == 1 ? i : ndims(Δ)+1, Val(ndims(Δ)))))

coerce_cuda(x::Union{CuArray,Nothing}) = x
coerce_cuda(x::Tuple) = coerce_cuda.(x)

coerce_cuda(x::AbstractArray) = x .+ CuArrays.fill(0)

function struct_grad!(cx::Zygote.Context, x, x̄)
  for f in fieldnames(typeof(x))
    Zygote.accum_param(cx, getfield(x, f), getfield(x̄, f))
  end
  dx = Zygote.grad_mut(cx, x)
  dx[] = Zygote.accum(dx[], x̄)
  return dx
end

for RNN in (CuRNN, CuGRU)
  @eval @adjoint function (m::$RNN{T})(h::CuArray{T}, x::CuArray{T}) where T <: Union{Float32,Float64}
    (y, ho), back = CUDNN.pullback(desc(m), x, h)
    (ho, y), function (Δ)
      dho, dy = coerce_cuda(Δ) # Support FillArrays etc.
      m̄ = back(dy, dho)
      dm = struct_grad!(__context__, m, (σ=nothing,Wi=transpose(m̄.Wi),Wh=transpose(m̄.Wh),b=m̄.b,h=nothing))
      (dm, unbroadcast(h, m̄.h), m̄.x)
    end
  end
end

@adjoint function (m::CuLSTM)((h, c)::Tuple{CuArray{T},CuArray{T}}, x::CuArray{T}) where T <: Union{Float32,Float64}
  (y, ho, co), back = CUDNN.pullback(desc(m), x, h, c)
  ((ho, co), y), function (Δ)
    dhc, dy = coerce_cuda(Δ) # Support FillArrays etc.
    dho, dco = dhc === nothing ? (nothing, nothing) : dhc
    m̄ = back(dy, dho, dco)
    dm = struct_grad!(__context__, m, (σ=nothing,Wi=transpose(m̄.Wi),Wh=transpose(m̄.Wh),b=m̄.b,h=nothing,c=nothing))
    (dm, (unbroadcast(h, m̄.h), unbroadcast(c, m̄.c)), m̄.x)
  end
end


text:
using CongruenceClasses
using Test

@testset "CongruenceClasses.jl" begin
    # Write your tests here.
end


text:
# Julia to Knossos
using Zygote
using IRTools
using ArgParse
import Base.show
import IRTools: isconditional

### TODOS
# 1. Need to propagate types
#    (still some Anys in basic block defs)
# 2. Nice printing :)

function parse_commandline()
    s = ArgParseSettings()

    @add_arg_table! s begin
        "--input"
            help = "input .jl file, last expression used, absolute path"
            required = true
        "--output", "-o"
            help = "output Knossos .ks extension recommended"
            required = true
    end

    return parse_args(s)
end

parsed_args = parse_commandline()
inputfilename = parsed_args["input"]
outputfilename = parsed_args["output"]

mkpath(dirname(outputfilename))

io = open(outputfilename, "w")

# Newline constant for s-expr printing
nl = "\n"
tab = "\t"

# Simple s-expression structure
struct SExpr
    terms :: Array{Any}
end

# Return the head of an s-expression
head(e) = e.terms[1]
tail(e) = SExpr(e.terms[2:end])

# Print an s-expression
function show(io::IO, ex::SExpr)
    print(io, "(")
    terms = ex.terms
    if length(terms) > 0
        for a in terms[1:end-1]
            print(io, a)
            print(io, " ")
        end
        print(io, terms[end])
    end
    print(io, ")")
end

sexp(terms...) = SExpr([terms...])

## IR utilities
blockid(id::Int) = string("b",id)

## make_sexp converts IR to s-expressions
function make_sexp(x :: Any)
    sexp("[UNK", typeof(x), ":", x, "]")
end

function make_sexp(x :: Number)
    x
end

function make_sexp(x :: IRTools.Variable)
    string("_", x.id)
end

function make_sexp(x :: GlobalRef)
    string(x.mod, ".", x.name)
end

function make_sexp(x :: Expr)
    if x.head == :call
        sexp(map(make_sexp, x.args)...)
    else
        sexp("[EXPR<", x.head, ">:", x, "]")
    end
end

function make_let(x :: IRTools.Variable, st :: IRTools.Statement, body)
    if isnothing(st.expr)
        ex = sexp("!!!") # When does this happen?
    else
        ex = make_sexp(st.expr)
    end
    if isnothing(x)
        sexp("nil", ex)
    else
        sexp("let", sexp(make_sexp(x),ex), nl, body)
    end
end

function make_lets(let_args, body)
    n = length(let_args)
    if n == 0
        return body
    end

    x,st = let_args[1]
    body = make_lets(let_args[2:end], body)
    make_let(x, st, body)
end

function make_if(cond, texpr, fexpr)
    sexp("if", make_sexp(cond), nl,
         tab, texpr, nl,
         tab, fexpr)
end

# struct IRTools.Branch
#     condition :: Any
#     block     :: Int64
#     args      :: Array{Any,1}

function make_sexp_branch(id :: Int, br :: IRTools.Branch)
    if br.block == 0
        sexp(br.args...)
    else
        sexp(blockid(br.block), br.args...)
    end
end

function make_sexp_branches(ir :: IRTools.IR, id :: Int, br :: IRTools.Branch, tail...)
    # Must be conditional if tail nonempty
    @assert isconditional(br)
    make_if(br.condition,
            make_sexp_branch(id, br),
            make_sexp_branches(ir, id, tail...))
end

function make_sexp_branches(ir :: IRTools.IR, id :: Int, br :: IRTools.Branch)
    # The last branch in a list
    if isconditional(br)
        # if conditional, is either the br or a fallthrough
        make_if(br.condition,
                sexp(blockid(id+1), IRTools.arguments(ir.blocks[id])),
                make_sexp_branch(id, br))
    else
        # unconditional...
        make_sexp_branch(id, br)
    end
end

function block_args(b :: IRTools.Block)
    ats = zip(IRTools.arguments(b),IRTools.argtypes(b))
    [sexp(a,":",t) for (a,t) in ats]
end

# struct IRTools.Block
#    arguments   4: (%1, %2)
#    argtypes
#    statements   %1 = e
#    branches     br n if e     -- if e, call block n, else fallthrough
function make_def(is_toplevel :: Bool, name, b :: IRTools.Block)
    args = sexp(block_args(b)...)
    if is_toplevel
        # Strip the method name for top-level functions
        args = tail(args)
    end

    branches = make_sexp_branches(b.ir, b.id, IRTools.branches(b)...)

    let_args = [(x,t) for (x,t) in b if x != undef]
    body = make_lets(let_args, branches)

    sexp("def",name,
            "Any", nl, tab,
            args, nl,
            body)
end

# Make s-expression for method f
# Generate a top-level function for each basic block,
# Top level: ir is a list of blocks
function make_sexps(ir :: IRTools.IR)
    blocks = IRTools.blocks(ir)

    # Emit blocks n:2 in reverse
    for b in reverse(blocks[2:end])
        se = make_def(false, blockid(b.id), b)
        println(io, se,"\n")
    end
    # Emit block 1
    f = IRTools.argtypes(blocks[1])[1].val
    se = make_def(true, string(f), blocks[1])
    println(io, se)
end

###################
### Test make_sexp

function jl2ks(f, argtypes)
    meta = IRTools.typed_meta(Tuple{typeof(f),argtypes...})
    ir = IRTools.IR(meta)
    IRTools.expand!(ir)

    println(io, "#| -------- IR ----------")
    println(io, ir)
    println(io, "-------- |#")

    make_sexps(ir)
end

sample = evalfile(joinpath(pwd(), inputfilename))

jl2ks(sample, (Vector{Float64},Float64))

max_stars_repo_name: tomjaguarpaw/knossos-ksc

max_issues_repo_name: tomjaguarpaw/knossos-ksc

max_forks_repo_name: tomjaguarpaw/knossos-ksc


text:
include("../../testcases.jl")
using PyCall
using DFTK

atoms = [ElementPsp(:Fe, psp=load_psp("hgh/pbe/Fe-q16.hgh")) => iron_bcc.positions]
model = model_PBE(iron_bcc.lattice, atoms, temperature=iron_bcc.temperature,
                  smearing=Smearing.FermiDirac(), spin_polarization=:collinear)

abinitpseudos = [joinpath(@__DIR__, "Fe-q16-pbe.abinit.hgh")]
DFTK.run_abinit_scf(model, @__DIR__; abinitpseudos=abinitpseudos,
                    Ecut=50, kgrid=[11, 11, 11],
                    spinat=[0 0 4.0], n_bands=14, tol=1e-10, iscf=13, nstep=100)


text:
#!/usr/bin/env julia

"""
Simulateur de propagation d'épidémies
"""

module ICI_simulateur

    ###################### Importation des librairies ##########################
    using Statistics
    using Random
    using Distributions
    using DelimitedFiles
    using SpecialFunctions
    using MetaGraphs
    using LightGraphs
    using GraphMLReader
    using JSON
    using MPI
    using GraphIO
    using Dates
    using CSV
    using DataFrames
    using StatsBase
    ############################################################################

    ################# Include all functions and modules ########################
    include("Module_Parameters.jl")
    include("Module_Mesures.jl")
    include("functions_pretraitement.jl")
    include("struct_city.jl")
    include("struct_population.jl")
    include("struct_indicateur.jl")
    include("functions_deplacements.jl")
    include("struct_covid.jl")
    include("struct_simulateur.jl")
    include("functions_posttraitement.jl")
    ############################################################################

    export Principal
    """
    Function Principal
    Fonction principale du simulateur. Permet de lancer les simulations de propagation d'épidémies sur une région donnée.
    C'est cette fonction qu'il faut appeler dans le programme principal
    """
    function Principal()

        deb = time() ;
        ############ Initialisation de la parallélisation ###################
        comm,rank = parallelisation() ;
        #####################################################################

        ############### Définition de la ville ##############################
        Ville = City() ;
        #####################################################################

        ############# Définition de la population ###########################
        Individu = Population() ;
        #####################################################################

        ########################## Simulateur ###############################
        Propagation = Simulateur(Ville,Individu,rank) ;
        #####################################################################

        ################## Partie des résultats #############################
        save(Propagation,rank,comm) ;
        #####################################################################

        if(ParamSimu.vMPI)
            MPI.Barrier(comm) ;
        end

        if(rank == 0)
            elapsed = time() - deb ;
            println("Temps reel : ", elapsed, " secondes.") ;
        end

        if(ParamSimu.vMPI)
            MPI.Finalize() ;
        end

    end
end


text:
SRC_DIR = joinpath(Base.source_dir(), "..", "..", "..", "src")
OUT_DIR = joinpath(Base.source_dir(), "..", "..", "..", "output")
mkpath(OUT_DIR)

include(joinpath(SRC_DIR, "BridgeSDEInference.jl"))
using Main.BridgeSDEInference
#include(joinpath(SRC_DIR, "BridgeSDEInference_for_tests.jl"))

using StaticArrays, LinearAlgebra, GaussianDistributions
using Distributions
using Random
using Bridge
using CSV, DataFrames
DIR = "auxiliary"
include(joinpath(SRC_DIR, DIR, "data_simulation_fns.jl"))
include(joinpath(SRC_DIR, DIR, "utility_functions.jl"))
include(joinpath(SRC_DIR, DIR, "plotting_fns.jl"))

#=============================================================================
                Routines for setting up the MCMC sampler
=============================================================================#
function _prepare_setup(ρ=0.96, σ_RW_step=0.5, num_mcmc_steps=4*10^3,
                        do_blocking=false, b_step=2, num_steps_for_change_pt=10,
                        thin_path=10^0, save_path_every=1*10^3, update_σ=false)
    t_kernel = ( update_σ ? [RandomWalk([], []),
                             RandomWalk([2.0, 1.0, 0.64, σ_RW_step], 4)] :
                            [RandomWalk([], [])] )
    updt_type = ( update_σ ? (ConjugateUpdt(), MetropolisHastingsUpdt()) :
                             (ConjugateUpdt(),) )
    prior_list = ( update_σ ? (MvNormal([0.0,0.0,0.0], diagm(0=>[1000.0, 1000.0, 1000.0])),
                               ImproperPrior()) :
                              (MvNormal([0.0,0.0,0.0], diagm(0=>[1000.0, 1000.0, 1000.0])),) )
    updt_coord = (update_σ ? [[1,2,3],[4]] : [[1,2,3]])
    setup = MCMCSetup(Pˣ, P̃, PartObs())
    set_observations!(setup, [L for _ in P̃], [Σ for _ in P̃], obs_vals, obs_time) # uses default fpt
    set_imputation_grid!(setup, 1/2000)
    set_transition_kernels!(setup,
                            t_kernel,
                            ρ, true, updt_coord,
                            updt_type,                           # update types
                            Adaptation(x0,
                                       [0.7, 0.4, 0.2, 0.2, 0.2],
                                       [500, 500, 500, 500, 500],
                                       1)
                            )
    set_priors!(setup,
                Priors(prior_list),               # priors over parameters
                GsnStartingPt(x0, @SMatrix [400.0 0.0 0.0;
                                            0.0 20.0 0.0;
                                            0.0 0.0 20.0]), # prior over starting point
                x0
                )
    set_mcmc_params!(setup,
                     num_mcmc_steps,    # number of mcmc steps
                     save_path_every,   # save path every ... iteration
                     10^2,              # print progress message every ... iteration
                     thin_path,         # thin the path imputatation points for save
                     0                # number of first iterations without param update
                     )
    if do_blocking
        set_blocking!(setup, ChequeredBlocking(),
                      (collect(1:length(obs_vals)-2)[1:b_step:end], 10^(-10),
                       SimpleChangePt(num_steps_for_change_pt)))
    end
    set_solver!(setup, Vern7(), NoChangePt())
    initialise!(eltype(x0), setup)
    setup
end

#==============================================================================
                        Auxiliary routines for saving
==============================================================================#

function save_paths(tt, paths, filename)
    d = length(paths[1][1])
    xx = copy(tt)
    for i in 1:length(paths)
        for j in 1:d
            xx = hcat(xx, [p[j] for p in paths[i]])
        end
    end
    CSV.write(joinpath(OUT_DIR, filename), DataFrame(xx))
end

function save_marginals(tt, paths, filename, indices)
    d = length(paths[1][1])
    xx = copy(tt[indices])
    N = length(paths)
    stride = div(N, 100)
    for i in 1:length(paths)
        if i % stride == 0
            print(div(i, 100), "% done...\n")
        end
        for j in 1:d
            xx = hcat(xx, [p[j] for p in paths[i][indices]])
        end
    end
    CSV.write(joinpath(OUT_DIR, filename), DataFrame(xx'))
end

function save_param_chain(chain, filename)
    d = length(chain[1])
    xx = [c[1] for c in chain]
    for i in 2:d
        xx = hcat(xx, [c[i] for c in chain])
    end
    CSV.write(joinpath(OUT_DIR, filename), DataFrame(xx))
end

function save_obs(tt, paths, filename)
    d = length(paths[1][1])
    xx = copy(tt)
    for i in 1:length(paths)
        for j in 1:d
            xx = hcat(xx, [p[j] for p in paths[i]])
        end
    end
    CSV.write(joinpath(OUT_DIR, filename), DataFrame(xx))
end

#==============================================================================
                        Generate densely observed process
==============================================================================#

Random.seed!(4)
θˣ = [10.0, 28.0, 8.0/3.0, 3.0]
Pˣ = LorenzCV(θˣ...)

x0, dt, T = ℝ{3}(1.5, -1.5, 25.0), 1/5000, 2.0
tt = 0.0:dt:T
XX, _ = simulateSegment(ℝ{3}(0.0, 0.0, 0.0), x0, Pˣ, tt)

Σdiagel = 5*10^0
Σ = SMatrix{2,2}(1.0I)*Σdiagel
L = @SMatrix[0.0 1.0 0.0;
             0.0 0.0 1.0]

skip = 50
obs_time, obs_vals = XX.tt[1:skip:end], [rand(Gaussian(L*x, Σ)) for x in XX.yy[1:skip:end]]

θ_init = [5.0, 15.0, 6.0, 3.0]
aux_flag = Val{(false,true,true)}()
P̃ = [LorenzCVAux(θ_init..., t₀, u, T, v, aux_flag, x0[3]) for (t₀, T, u, v)
     in zip(obs_time[1:end-1], obs_time[2:end], obs_vals[1:end-1], obs_vals[2:end])]
Pˣ = LorenzCV(θ_init...)
print("mid-point time: ", XX.tt[div(length(XX.yy),2)+1], ", mid-point value: ",
      XX.yy[div(length(XX.yy),2)+1], "\n")
print("three-quarters time: ", XX.tt[3*div(length(XX.yy),4)+1], ", three-quarters value: ",
      XX.yy[3*div(length(XX.yy),4)+1], "\n")

save_paths(obs_time, [obs_vals], "many_obs_data.csv")

#==============================================================================
                    Run the expriment: very large number of blocks
==============================================================================#
setup = _prepare_setup(0.1, 0.3, 1*10^4, true, 1, 15, 1, 10^2)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup)
display(out.accpt_tracker)

plot_chains(out; truth=[10.0, 28.0, 8.0/3.0, 3.0],
            ylims=[nothing, (25,30), (2,5), (0,10)])
plot_paths(out; obs=(times=obs_time[2:end],
                     vals=[[v[1] for v in obs_vals[2:end]],
                           [v[2] for v in obs_vals[2:end]]], indices=[2,3]))

save_paths(out.time, out.paths, "many_obs_many_blocks_paths.csv")
save_param_chain(out.θ_chain.θ_chain, "many_obs_many_blocks_chain.csv")

# repeat, this time save marginals
setup = _prepare_setup(0.1, 0.3, 1*10^4, true, 1, 15, 1000, 1)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup)
display(out.accpt_tracker)
save_marginals(out.time, out.paths, "many_obs_many_blocks_marginals.csv", [1,50,100,150,200])

#==============================================================================
                    Run the expriment: longer blocks (less of them)
==============================================================================#
setup = _prepare_setup(0.5, 0.8, 1*10^4, true, 4, 15, 1, 10^2)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup)
display(out.accpt_tracker)

plot_chains(out; truth=[10.0, 28.0, 8.0/3.0, 3.0],
            ylims=[nothing, (25,30), (2,5), (0,10)])
plot_paths(out; obs=(times=obs_time[2:end],
                     vals=[[v[1] for v in obs_vals[2:end]],
                           [v[2] for v in obs_vals[2:end]]], indices=[2,3]))

save_paths(out.time, out.paths, "many_obs_medium_blocks_paths.csv")
save_param_chain(out.θ_chain.θ_chain, "many_obs_medium_blocks_chain.csv")

setup = _prepare_setup(0.5, 0.8, 1*10^4, true, 4, 15, 1000, 1)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup)
display(out.accpt_tracker)
save_marginals(out.time, out.paths, "many_obs_medium_blocks_marginals.csv", [1,50,100,150,200])

#==============================================================================
                    Run the expriment: very long blocks
==============================================================================#
setup = _prepare_setup(0.2, 0.8, 1*10^4, false, 40, 15, 1, 10^2)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup)
display(out.accpt_tracker)

plot_chains(out; truth=[10.0, 28.0, 8.0/3.0, 3.0],
            ylims=[nothing, (25,30), (2,5), (0,10)])
plot_paths(out; obs=(times=obs_time[2:end],
                     vals=[[v[1] for v in obs_vals[2:end]],
                           [v[2] for v in obs_vals[2:end]]], indices=[2,3]))

save_paths(out.time, out.paths, "many_obs_long_blocks_paths.csv")
save_param_chain(out.θ_chain.θ_chain, "many_obs_long_blocks_chain.csv")

setup = _prepare_setup(0.2, 0.8, 1*10^4, false, 4, 15, 1000, 1)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup)
display(out.accpt_tracker)
save_marginals(out.time, out.paths, "many_obs_long_blocks_marginals.csv", [1,50,100,150,200])

#==============================================================================
                        Generate sparsely observed process
==============================================================================#

Random.seed!(4)
θˣ = [10.0, 28.0, 8.0/3.0, 3.0]
Pˣ = LorenzCV(θˣ...)

x0, dt, T = ℝ{3}(1.5, -1.5, 25.0), 1/5000, 2.0
tt = 0.0:dt:T
XX, _ = simulateSegment(ℝ{3}(0.0, 0.0, 0.0), x0, Pˣ, tt)

Σdiagel = 5*10^(-1)
Σ = SMatrix{2,2}(1.0I)*Σdiagel
L = @SMatrix[0.0 1.0 0.0;
             0.0 0.0 1.0]

skip = 1000
obs_time, obs_vals = XX.tt[1:skip:end], [rand(Gaussian(L*x, Σ)) for x in XX.yy[1:skip:end]]

θ_init = [5.0, 15.0, 6.0, 3.0]
aux_flag = Val{(false,true,true)}()
P̃ = [LorenzCVAux(θ_init..., t₀, u, T, v, aux_flag, x0[3]) for (t₀, T, u, v)
     in zip(obs_time[1:end-1], obs_time[2:end], obs_vals[1:end-1], obs_vals[2:end])]
Pˣ = LorenzCV(θ_init...)
print("mid-point time: ", XX.tt[div(length(XX.yy),2)+1], ", mid-point value: ",
      XX.yy[div(length(XX.yy),2)+1], "\n")
print("three-quarters time: ", XX.tt[3*div(length(XX.yy),4)+1], ", three-quarters value: ",
      XX.yy[3*div(length(XX.yy),4)+1], "\n")
save_paths(obs_time, [obs_vals], "few_obs_data.csv")

#==============================================================================
                    Run the expriment: very long blocks
==============================================================================#
setup = _prepare_setup(0.96, 0.1, 1*10^4, false, 40, 15, 1, 10^2, false)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup)
display(out.accpt_tracker)

plot_chains(out; truth=[10.0, 28.0, 8.0/3.0, 3.0],
            ylims=[nothing, nothing, nothing, (0,10)])
plot_paths(out; obs=(times=obs_time[2:end],
                     vals=[[v[1] for v in obs_vals[2:end]],
                           [v[2] for v in obs_vals[2:end]]], indices=[2,3]))

save_paths(out.time, out.paths, "few_obs_long_blocks_paths.csv")
save_param_chain(out.θ_chain.θ_chain, "few_obs_long_blocks_chain.csv")

setup = _prepare_setup(0.96, 0.8, 1*10^3, false, 4, 15, 1000, 1)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup)
display(out.accpt_tracker)
save_marginals(out.time, out.paths, "few_obs_long_blocks_marginals.csv", [1,3,6,9,10])

#==============================================================================
                    Run the expriment: very large number of blocks
==============================================================================#
setup = _prepare_setup(0.7, 0.3, 1*10^4, true, 1, 100, 1, 10^2)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup)
display(out.accpt_tracker)

plot_chains(out; truth=[10.0, 28.0, 8.0/3.0, 3.0],
            ylims=[nothing, (25,30), (2,5), (0,10)])
plot_paths(out; obs=(times=obs_time[2:end],
                     vals=[[v[1] for v in obs_vals[2:end]],
                           [v[2] for v in obs_vals[2:end]]], indices=[2,3]))

save_paths(out.time, out.paths, "few_obs_many_blocks_paths.csv")
save_param_chain(out.θ_chain.θ_chain, "few_obs_many_blocks_chain.csv")

# repeat, this time save marginals
setup = _prepare_setup(0.7, 0.3, 1*10^4, true, 1, 100, 1000, 1)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup)
display(out.accpt_tracker)
save_marginals(out.time, out.paths, "few_obs_many_blocks_marginals.csv", [1,3,6,9,10])

#==============================================================================
            Run the expriment: very large number of blocks and adaptive
#NOTE because adaptation is currently implemented in a rather inefficient way
# the code responsible for it is by default **COMMENTED OUT**. Consequently,
# to run the code with adaptation one needs to manually remove `#` from the file
# `src/examples/lorenz_system_const_vola.jl` in line 104--105. Run the code
# below and comment it out again.
==============================================================================#
setup = _prepare_setup(0.7, 0.3, 1*10^4, true, 1, 100, 1, 10^2)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup)
display(out.accpt_tracker)

plot_chains(out; truth=[10.0, 28.0, 8.0/3.0, 3.0],
            ylims=[nothing, (25,30), (2,5), (0,10)])
plot_paths(out; obs=(times=obs_time[2:end],
                     vals=[[v[1] for v in obs_vals[2:end]],
                           [v[2] for v in obs_vals[2:end]]], indices=[2,3]))

save_paths(out.time, out.paths, "few_obs_many_blocks_and_adpt_paths.csv")
save_param_chain(out.θ_chain.θ_chain, "few_obs_many_blocks_and_adpt_chain.csv")

setup = _prepare_setup(0.7, 0.3, 1*10^4, true, 1, 100, 1000, 1)
Random.seed!(4)
out, elapsed = @timeit mcmc(setup)
display(out.accpt_tracker)
save_marginals(out.time, out.paths, "few_obs_many_blocks_and_adpt_marginals.csv", [1,3,6,9,10])


text:
"""
    PricingModel{C<:Currencies.Currency}

Represents a Pricing Model for a `Contract`.

# Requires

* `is_riskfree_discountfactor(model::PricingModel, kind::String, rf::DiscountFactor) :: Bool`

# Provides

* `get_functional_currency(model) :: Currencies.Currency`

"""
abstract type PricingModel{C<:Currencies.Currency} end

function is_riskfree_discountfactor end

@inline function get_functional_currency(m::PricingModel{C}) :: C where {C<:Currencies.Currency}
    return C()
end

"""
    StaticHedgingModel

Standard discounted cashflow model. Used to price instruments that can be hedged with a static strategy.
"""
struct StaticHedgingModel{C} <: PricingModel{C}
    functional_currency::C
    riskfree_curve_map::Dict{String, Dict{Symbol, Symbol}} # "onshore" -> [ :BRL -> :PRE ]
end

"""
    FutureValueModel

Wrapps a `StaticHedgingModel` and changes its behavior
to disable the discounted cashflow rule,
preserving any projection to future values.

Used to project cashflows.
"""
struct FutureValueModel{C} <: PricingModel{C}
    underlying_model::StaticHedgingModel{C}
end

# docstrings generated by docs.jl
struct BlackScholesModel{C} <: PricingModel{C}
    static_model::StaticHedgingModel{C}
end

struct BinomialModelDaily{C, U<:CashRiskFactor, T<:InterestRates.DayCountConvention} <: PricingModel{C}
    static_model::StaticHedgingModel{C}
    underlying::U # CashRiskFactor -> Stock or SpotCurrency
    numeraire_daycount_convention::T # the numeraire is static_model.functional_currency
end


text:
function gamma_epsilon_rate(binaries::Array{Int,2})
  n_ones = sum(binaries, dims=1)
  gamma_rate = n_ones .> length(binaries[:,1]) / 2
  epsilon_rate = .!gamma_rate

  return bool_to_int(gamma_rate), bool_to_int(epsilon_rate)
end

function power_consumption(binaries::Array{String,1})
  binaries = parse_input(binaries)
  gamma, epsilon = gamma_epsilon_rate(binaries)
  power_consumption = gamma * epsilon
  println("Power consumption is $power_consumption")
end

function oxygen_rating(arr::Array{Int,2})
  bin_considered = 1
  length_arr = length(arr[:,1])
  while length(arr[:,1]) > 1
    n_ones = sum(arr[:, bin_considered])
    bit_to_keep = n_ones >= length_arr / 2 ? 1 : 0
    to_be_removed = Int[]
    for i = 1:1:length_arr
      if arr[i, bin_considered] != bit_to_keep
        push!(to_be_removed, i)
      end
    end
    arr = arr[setdiff(1:length_arr, to_be_removed),:]
    length_arr = length(arr[:,1])
    bin_considered += 1
  end
  return arr_to_int(arr)
end

function carbon_dioxide_rating(arr::Array{Int,2})
  bin_considered = 1
  length_arr = length(arr[:,1])
  while length(arr[:,1]) > 1
    n_ones = sum(arr[:, bin_considered])
    bit_to_keep = n_ones >= length_arr / 2 ? 0 : 1
    to_be_removed = Int[]
    for i = 1:1:length_arr
      if arr[i, bin_considered] != bit_to_keep
        push!(to_be_removed, i)
      end
    end
    arr = arr[setdiff(1:length_arr, to_be_removed),:]
    length_arr = length(arr[:,1])
    bin_considered += 1
  end
  return arr_to_int(arr)
end

function life_support_rating(arr::Array{String,1})
  parsed_input = parse_input(arr)
  oxygen = oxygen_rating(parsed_input)
  carbon_dioxide = carbon_dioxide_rating(parsed_input)
  life_support = oxygen * carbon_dioxide

  println("Life support rating is $life_support")
end

#=====================================
# Formatter and parser functions below
=====================================#
function parse_input(arr::Array{String,1})
  n_bits = length(arr[1])
  n_numbers = length(arr)
  rval = Array{Int,2}(undef, n_numbers, n_bits)
  for i in 1:1:n_numbers
    for j in 1:1:n_bits
      rval[i, j] = parse(Int, arr[i][j])
    end
  end
  return rval
end

function bool_to_int(arr)
  str = ""
  for bool in arr
    if bool
      str *= "1"
    else
      str *= "0"
    end
  end
  return parse(Int, str, base=2)
end

function arr_to_int(arr)
  str = ""
  for bit in arr
    str *= string(bit)
  end
  return parse(Int, str, base=2)
end

#=====================================
# Main
=====================================#
test_input = ["00100", "11110", "10110", "10111", "10101", "01111", "00111",
              "11100", "10000", "11001", "00010", "01010"]
power_consumption(test_input)
life_support_rating(test_input)

# Application to acutal input.
println()
open("input.txt", "r") do f
  data = readlines(f)
  power_consumption(data)
  life_support_rating(data)
end


text:
# ============================================================
# Collision Term
# ============================================================

"""
Right-hand side of universal Boltzmann equation
`ube_dfdt(f, p, t)`
* f: particle distribution function in 1D formulation
* p: M, τ, (ann) (parameters)
* t: tspan

"""
function ube_dfdt(f, p, t)
    M, τ, ann = p

    if ann[1] isa FastChain
        df = (M - f) / τ + ann[1](M - f, ann[2])
    elseif ann[1] isa Chain
        df = (M - f) / τ + ann[1](M - f)
    end

    return df
end

"""
Right-hand side of universal Boltzmann equation
`ube_dfdt!(df, f, p, t)`

"""
function ube_dfdt!(df, f, p, t)
    M, τ, ann = p

    if ann[1] isa FastChain
        df .= (M - f) / τ + ann[1](M - f, ann[2])
    elseif ann[1] isa Chain
        df .= (M - f) / τ + ann[1](M - f)
    end
end


text:
using .Space: Box

mutable struct PendulumEnv <: AbstractEnv
    max_speed::Float32
    max_torque::Float32
    last_u
    dt::Float32
    viewer
    state
    action_space::Box
    observation_space::Box
end

include("vis/pendulum.jl")

function PendulumEnv()
    max_speed = 8f0
    max_torque = 2f0
    dt = 5f-2
    viewer = nothing
    # Observation space limits
    high = [1f0, 1f0, max_speed]
    action_space = Box(-max_torque, max_torque, (1,), Float32)
    observation_space = Box(-high, high, Float32)

    PendulumEnv(max_speed, max_torque, nothing, dt, nothing, nothing, action_space, observation_space)
end

function step!(env::PendulumEnv, u)
    # @assert u ∈ env.action_space "Invalid action ($(u)) issued"
    θ, θ̇  = env.state[1:1], env.state[2:2]
    g, m, l, dt = 10f0, 1f0, 1f0, env.dt

    v = clamp.(u, -env.max_torque, env.max_torque)
    env.last_u = v[1] # for rendering

    costs = angle_normalize.(θ).^2 .+ 1f-1(θ̇ .^2) .+ 1f-3(u.^2)

    tempθ̇_ = θ̇  .+ (-3g/(2l) * sin.(θ .+ pi) .+ 3/(m*l^2)*u) * dt
    θ_     = θ  .+ tempθ̇_*dt
    θ̇_     = clamp.(tempθ̇_, -env.max_speed, env.max_speed)

    env.state = vcat(θ_, θ̇_)
    r = -sum(costs)
    return _get_obs(env), r, false, Dict()
end

function reset!(env::PendulumEnv)
    high = Float32.([π, 1])
    env.state = 2rand(Float32, 2) .* high .- high

    if isdefined(Main, :CuArrays)
        env.state = env.state |> gpu
    end

    env.last_u = nothing
    return _get_obs(env)
end

function _get_obs(env::PendulumEnv)
    θ, θ̇ = env.state[1:1], env.state[2:2]
    return vcat(cos.(θ), sin.(θ), θ̇ )
end

pi = Float32(π)

angle_normalize(x) = ((x + pi) % 2pi + 2pi) % 2pi - pi


text:
"""
    pybool(x)

Convert `x` to a Python `bool`.
"""
pybool(x::Bool=false) = Py(x ? pybuiltins.True : pybuiltins.False)
pybool(x::Number) = pybool(!iszero(x))
pybool(x) = pybuiltins.bool(x)
export pybool

pyisTrue(x) = pyis(x, pybuiltins.True)
pyisFalse(x) = pyis(x, pybuiltins.False)
pyisbool(x) = pyisTrue(x) || pyisFalse(x)

pybool_asbool(x) =
    @autopy x if pyisTrue(x_)
        true
    elseif pyisFalse(x_)
        false
    else
        error("not a bool")
    end

function pyconvert_rule_bool(::Type{T}, x::Py) where {T<:Number}
    val = pybool_asbool(x)
    if T in (Bool, Int8, Int16, Int32, Int64, Int128, UInt8, UInt16, UInt32, UInt64, UInt128, BigInt)
        pyconvert_return(T(val))
    else
        pyconvert_tryconvert(T, val)
    end
end

pyconvert_rule_fast(::Type{Bool}, x::Py) =
    if pyisTrue(x)
        pyconvert_return(true)
    elseif pyisFalse(x)
        pyconvert_return(false)
    else
        pyconvert_unconverted()
    end


text:
using Test
include("../src/Threadings.jl")
@testset "Threadings tests" begin
function f(x)
  # this bit takes about 10 seconds
  local a = 0.0
  for i in 1:10
    a += sum(cbrt.(sqrt.(rand(10_00_000))))
  end
  return x .+ 0 * a
end
itr = [i*collect(rand(5)) for i in 1:10]
a = Threadings.tmapreduce(f, +, itr, init=zeros(5))
b = mapreduce(f, +, itr, init=zeros(5))
ta = @elapsed Threadings.tmapreduce(f, +, itr, init=zeros(5))
tb = @elapsed mapreduce(f, +, itr, init=zeros(5))
@show ta, tb
Base.Threads.nthreads() == 1 && @test ta < 1.1*tb
Base.Threads.nthreads() > 1 && @test ta < tb
for i in 1:5
  @test a[i] ≈ b[i] rtol=sqrt(eps()) atol=0.0
end

itr = 1:10
a = Threadings.tmap(f, itr)
b = map(f, itr)
ta = @elapsed Threadings.tmap(f, itr)
tb = @elapsed map(f, itr)
@show ta, tb
Base.Threads.nthreads() == 1 && @test ta < 1.1*tb
Base.Threads.nthreads() > 1 && @test ta < tb
@test a ≈ b rtol=sqrt(eps()) atol=0.0
end

max_stars_repo_name: jwscook/Threadings

max_issues_repo_name: jwscook/Threadings

max_forks_repo_name: jwscook/Threadings


text:
# definition of Agent structure

using Agents

# define household
mutable struct household_type
  id::Int
  members::Array{Int}
end

# define agents:
mutable struct agent_type <: AbstractAgent
  id::Int
  pos::Int
  # define extra attributes
  # category of agent:
  worker::Bool
  traveller::Bool
  #worker-specific properties
  workdays::Array{Int64}
  present_in_workplace::Bool
  # traveller-specific properties:
  hh_id::Int64
  p_compliance::Float64 # probability of being compliant on a given day
  compliant::Bool
  quarantine_extended::Bool
  isolation_extended::Bool
  # some clocks
  time_in_quar::Float64
  time_in_extended_quar::Float64
  time_in_isolation::Float64
  time_remaining_in_isolation::Float64
  #general properties
  vaccinated::Bool
  infected::Bool
  latent::Bool
  incubating::Bool
  recovering::Bool
  recovered::Bool
  asymptomatic::Bool
  expressing_symptoms::Bool
  # some clocks
  time_infected::Float64
  time_incubating::Float64
  time_recovering::Float64
  time_infectious_in_community::Float64
  test_sensitivity::Float64
  b1::Float64
  b2::Float64
  b3::Float64
  breakpoint::Float64
  beta_t::Float64
  beta_max::Float64
  beta_min::Float64
  k_inc::Float64
  k_r::Float64

  incubation_period::Float64
  recovery_period::Float64

  test_positive::Bool
  awaiting_test_result::Bool
  time_to_test_result::Float64

  discharged::Bool
  detected::Bool
  recorded::Bool

  FoI_in_community::Float64

  time_discharged::Float64

  time_until_test::Array{Float64}

  index_case::Bool

  agent_type(id, pos) = new(id, pos,
                            false, #worker
                            false, #traveller
                            [0], #work days
                            false, # present_in_workplace::Bool
                            0, #hh_id::Int64
                            1.0,  #p_compliance::Float64
                            true, #compliant::Bool
                            false, #quarantine_extended::Bool
                            false, #isolation_extended::Bool
                            0.0, #time_in_quar::Float64
                            0.0, #time_in_extended_quar::Float64
                            0.0, #time_in_isolation::Float64
                            0.0,  #time_remaining_in_isolation::Float64
                            false, #vaccinated::Bool
                            false, #infected::Bool
                            false, #latent::Bool
                            false, #incubating::Bool
                            false, #recovering::Bool
                            false, #recovered::Bool
                            false, #asymptomatic::Bool
                            false, #expressing_symptoms::Bool
                            0.0, #time_infected::Float64
                            0.0, #time_incubating::Float64
                            0.0, #time_recovering::Float64
                            0.0, #time_infectious_in_community::Float64
                            0.0, #test_sensitivity::Float64
                            0.0, #b1::Float64
                            0.0, #b2::Float64
                            0.0, #b3::Float64
                            0.0, #breakpoint::Float64
                            0.0, #beta_t::Float64
                            0.0, #beta_max::Float64
                            0.0, #beta_min::Float64
                            0.0, #k_inc::Float64
                            0.0, #k_r::Float64
                            0.0, #incubation_period::Float64
                            0.0, #recovery_period::Float64
                            false, #test_positive::Bool
                            false, #awaiting_test_result::Bool
                            0.0, #time_to_test_result::Float64
                            false, #discharged::Bool
                            false, #detected::Bool
                            false, #recorded::Bool
                            0.0,#FoI_in_community::Float64
                            0.0, #time_discharged::Float64
                            [0.0], #time_until_test::Array{Float64}
                            false #index_case::Bool
                            )

end


text:
"""
    model_input(year::Int, d::Dict{Symbol,DataFrame}, set::Dict, idx::Dict)

# Arguments
- `year::Int`: year for which to perform calibration
- `d::Dict{Symbol,DataFrame}` of DataFrames containing the model data.
- `set::Dict` of Arrays describing region, sector, final demand, etc.
- `idx::Dict` of parameter indices

# Returns
- `d::Dict{Symbol,Dict}` of model data with all zeros filled
- `set::Dict` of model indices, with added arrays of tuples indicating which (region,sector)
    or (region,good) combinations are relevant to the zero-profit and market-clearing
    conditions.
- `idx::Dict` of parameter indices.
"""
function model_input(year::Any, d::Dict, set::Dict, idx::Dict=Dict())
    return model_input(ensurearray(year), d, set, idx)
end

function model_input(year::Array{Int,1}, d::Dict, set::Dict, idx::Dict=Dict())
    @info("Preparing model data for $year.")
    if length(year) > 1
        d2 = Dict(k => filter_with(df, (yr = year,); extrapolate=true) for (k, df) in d)

        isempty(idx) && (idx = Dict(k => findindex(df) for (k, df) in d2))
        (set, idx) = _model_set!(d2, set, idx)

        d1 = Dict(k => filter_with(df, (yr = year,); drop=true) for (k, df) in d)
        d1 = Dict(k => convert_type(Dict, fill_zero(df; with=set)) for (k, df) in d1)
        return (d1, set, idx)
    else
        d = Dict(k => filter_with(df, (yr = year,); drop=true) for (k, df) in d)

        isempty(idx) && (idx = Dict(k => findindex(df) for (k, df) in d))
        (set, idx) = _model_set!(d, set, idx)

        d = Dict(k => convert_type(Dict, fill_zero(df; with=set)) for (k, df) in d)
        return (d, set, idx)
    end
end

"""
    _model_set!(d::Dict{Symbol,DataFrame}, set::Dict, idx::Dict)
This function returns subsets intended to limit the size of the model by including only
non-zero values when mapping zero-profit and market-clearing conditions.

# Arguments
- `d::Dict{Symbol,DataFrame}` of DataFrames containing the model data.
- `set::Dict` of Arrays describing region, sector, final demand, etc.
- `idx::Dict` of parameter indices

# Returns
- `d::Dict{Symbol,Dict}` of model data with all zeros filled
- `set::Dict` of model indices, with added arrays of tuples indicating which (region,sector)
    or (region,good) combinations are relevant to the zero-profit and market-clearing
    conditions.
- `idx::Dict` of indices, updated to include those used to define the newly-added sets.
"""
function _model_set!(d::Dict{Symbol,DataFrame}, set::Dict, idx::Dict)
    (set[:A], idx[:A]) = nonzero_subset(d[:a0] + d[:rx0])
    (set[:Y], idx[:Y]) = nonzero_subset(combine_over(d[:ys0], :g))
    (set[:X], idx[:X]) = nonzero_subset(d[:s0])
    (set[:PA], idx[:PA]) = nonzero_subset(d[:a0])
    (set[:PD], idx[:PD]) = nonzero_subset(d[:xd0])
    (set[:PK], idx[:PK]) = nonzero_subset(d[:kd0])
    (set[:PY], idx[:PY]) = nonzero_subset(d[:s0])

    if :yr in idx[:kd0]
        set[:PKT] = convert_type(Array{Tuple}, unique(d[:kd0][:,setdiff(idx[:kd0], [:yr])]))
    end

    return (set, idx)
end

"""
    nonzero_subset(df::DataFrame)

# Returns
- `x::Array{Tuple,1}` of all parameter indices corresponding with non-zero values
- `idx::Array{Symbol,1}` of parameter indices in `df`
"""
function nonzero_subset(df::DataFrame)
    idx = findindex(df)
    val = convert_type(Array{Tuple}, dropzero(df)[:,idx])
    return (val, idx)
end

max_stars_repo_name: NREL/SLiDE

max_issues_repo_name: NREL/SLiDE

max_forks_repo_name: NREL/SLiDE


text:
"""
    new_trace = map_optimize(trace, selection::Selection,
        max_step_size=0.1, tau=0.5, min_step_size=1e-16, verbose=false)

Perform backtracking gradient ascent to optimize the log probability of the trace over selected continuous choices.

Selected random choices must have support on the entire real line.
"""
function map_optimize(trace, selection::Selection;
                      max_step_size=0.1, tau=0.5, min_step_size=1e-16, verbose=false)
    args = get_args(trace)
    argdiffs = map((_) -> NoChange(), args)
    retval_grad = accepts_output_grad(get_gen_fn(trace)) ? zero(get_retval(trace)) : nothing

    (_, values, gradient) = choice_gradients(trace, selection, retval_grad)
    values_vec = to_array(values, Float64)
    gradient_vec = to_array(gradient, Float64)
    step_size = max_step_size
    score = get_score(trace)
    while true
        new_values_vec = values_vec + gradient_vec * step_size
        values = from_array(values, new_values_vec)
        # TODO discard and weight are not actually needed, there should be a more specialized variant
        (new_trace, _, _, discard) = update(trace, args, argdiffs, values)
        new_score = get_score(new_trace)
        change = new_score - score
        if verbose
            println("step_size: $step_size, prev score: $score, new score: $new_score, change: $change")
        end
        if change >= 0.
            # it got better, return it
            return new_trace
        elseif step_size < min_step_size
            # it got worse, but we ran out of attempts
            return trace
        end

        # try again with a smaller step size
        step_size = tau * step_size
    end
end

export map_optimize


text:
# ---
# title: 1144. Decrease Elements To Make Array Zigzag
# id: problem1144
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Array
# link: <https://leetcode.com/problems/decrease-elements-to-make-array-zigzag/description/>
# hidden: true
# ---
#
# Given an array `nums` of integers, a _move_  consists of choosing any element
# and **decreasing it by 1**.
#
# An array `A` is a  _zigzag array_  if either:
#
#   * Every even-indexed element is greater than adjacent elements, ie. `A[0] > A[1] < A[2] > A[3] < A[4] > ...`
#   * OR, every odd-indexed element is greater than adjacent elements, ie. `A[0] < A[1] > A[2] < A[3] > A[4] < ...`
#
# Return the minimum number of moves to transform the given array `nums` into a
# zigzag array.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [1,2,3]
#     Output: 2
#     Explanation: We can decrease 2 to 0 or 3 to 1.
#
#
# **Example 2:**
#
#
#
#     Input: nums = [9,6,1,6,2]
#     Output: 4
#
#
#
#
# **Constraints:**
#
#   * `1 <= nums.length <= 1000`
#   * `1 <= nums[i] <= 1000`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
export dimension

mutable struct ModAlgAss{S, T, V}
  base_ring::S
  action::Vector{T}
  dimension::Int
  M::AbstractAlgebra.FPModule{V}
  isirreducible::Int
  dimension_splitting_field::Int
  algebra::AlgAss{V}
  action_of_gens::Vector{T}
  action_of_basis::Vector{T}

  function ModAlgAss{S, T}(action::Vector{T}) where {S, T}
    z = new{S, T, elem_type(S)}()
    z.action = action
    z.dimension = ncols(action[1])
    z.base_ring = base_ring(action[1])
    if z.dimension == 1
      z.isirreducible = 1
      z.dimension_splitting_field = 1
    else
      z.isirreducible = 0
      z.dimension_splitting_field = 0
    end
    return z
  end
end

mutable struct Lat{S, T}
  mod::S
  order::AlgAssAbsOrd
  basis_matrix::T

  function Lat{S, T}(mod::S, order::AlgAssAbsOrd, basis_matrix::T) where {S, T}
    z = new{S, T}(mod, order, basis_matrix)
    return z
  end
end

function vector_space(M::ModAlgAss)
  return M.M
end

function algebra(M::ModAlgAss)
  if isdefined(M :algebra)
    return M.algebra
  else
    A = algebra(M.action)
    M.algebra = A
    M.action_of_gens = M.action
    return A
  end
end

function action(a::AlgAssElem, M::ModAlgAss)
  parent(a) == algebra(M)
end

function basis_matrix(L::Lat)
  return L.basis_matrix
end

function action(a::AlgAssAbsOrdElem, L::Lat)
end

function ModAlgAss(action::Vector{T}) where {T}
  @assert length(action) > 0
  S = typeof(base_ring(action[1]))
  return ModAlgAss{S, T}(action)
end

function ModAlgAss(M::AbstractAlgebra.FPModule{T}, action::Vector{<:Generic.ModuleHomomorphism{T}}) where {T <: FieldElem}
  all(x->domain(x) == codomain(x) == M, action) || error("module and maps do not match")
  V = ModAlgAss([x.matrix for x = action])
  V.M = M
  return V
end

function isirreducible_known(M::ModAlgAss)
  return M.isirreducible != 0
end

function isirreducible(M::ModAlgAss)
  if M.isirreducible != 0
    return M.isirreducible == 1
  else
    fl, N = meataxe(M)
    return fl
  end
end

function dimension(M::ModAlgAss)
  return M.dimension
end

function base_ring(M::ModAlgAss)
  return M.base_ring
end

# Compute the algebra generated by M and images
# of the elements in M

function algebra(M::Vector{T}) where {T <: MatElem}
  @assert length(M) > 0
  A = M[1]
  n = nrows(A)
  n2 = n^2
  @assert n == ncols(A)
  K = base_ring(A)
  Mprod = M
  Morig = copy(M)

  current_dim = -1

  B = zero_matrix(K, length(Mprod) + 1, n2)

  l = 0
  while true
    if l != 0
      B = zero_matrix(K, length(Mprod), n2)
    end
    for k in 1:length(Mprod)
      for i in 1:n
        for j in 1:n
          B[k, (i - 1)* n  + j] = Mprod[k][i, j]
        end
      end
    end
    # Add the identity
    if l == 0
      for i in 1:n
        B[length(M) + 1, (i - 1)*n + i] = one(K)
      end
    end
    new_dim = rref!(B)
    if new_dim == current_dim
      break
    end
    current_dim = new_dim
    M = [ matrix(K, n, n, [B[k, (i - 1)*n + j] for i in 1:n for j in 1:n]) for k in 1:new_dim]
    Mprod = [ M[i] * M[j] for i in 1:length(M) for j in 1:length(M) ]
    l = l + 1
  end

  dim = current_dim
  B = sub(B, 1:dim, 1:ncols(B))

  basis = [ matrix(K, n, n, [B[k, (i - 1)*n + j] for i in 1:n for j in 1:n]) for k in 1:dim]

  @assert isone(basis[1])

  v = zero_matrix(K, 1, n2)

  structure = Array{elem_type(K), 3}(dim, dim, dim)

  for k in 1:dim
    for l in 1:dim
      N = basis[k] * basis[l]
      for i in 1:n
        for j in 1:n
          v[1, (i - 1)* n  + j] = N[i, j]
        end
      end
      b, u = can_solve_with_solution(B, v, side = :left)
      error("NOT HERE!")
      @assert b
      @assert N == sum(u[i]*basis[i] for i in 1:dim)
      for m in 1:dim
        structure[k, l, m] = u[m, 1]
      end
    end
  end

  A = AlgAss(K, structure)

  gens = Vector{AlgAssElem{elem_type(K)}}(length(Morig))

  for l in 1:length(Morig)
    N = Morig[l]
    for i in 1:n
      for j in 1:n
        v[1, (i - 1)* n  + j] = N[i, j]
      end
    end
    b, u = can_solve_with_solution(B, v, side = :left)
    gens[l] =  A([u[1, m] for m in 1:dim])
  end

  A.gens = gens

  return A
end

#function gens(A::AlgAss{T}) where {T}
#  #return A.gens::Vector{AlgAssElem{T}}
#end


text:
# Load modules
using Plots, LaTeXStrings
pyplot()
using ThinFilmsTools

function main()
    # Define beam
    λ = LinRange(400,1000,600) # wavelength range [nm]
    θ = [0.] # angle of incidence [degrees]
    beam = PlaneWave(λ,θ)
    # Define layers with their parameters
    n = [RIdb.air(beam.λ) RIdb.silicon(beam.λ)]
    p1 = 0.89
    p2 = 0.70
    p3 = 0.41
    layers = [
        LayerTMMO(n[:,1]), # incident medium
        LayerTMMO(RI.looyenga([p1 1-p1],n); d=77.),
        LayerTMMO(RI.looyenga([p2 1-p2],n); d=56.),
        LayerTMMO(RI.looyenga([p3 1-p3],n), d=39.),
        LayerTMMO(n[:,2]),  # emergent (substrate) medium
    ]
    # solve
    return tmm_optics(beam, layers)
end

sol = main()

# plot the R, T and A spectra
plot(Spectrum1D(),
     sol.beam.λ,
     [sol.Spectra.Rp, sol.Spectra.Tp, 1.0 .- (sol.Spectra.Rp .+ sol.Spectra.Tp)],
     label=["Reflectance" "Transmittance" "Absorbance"],
     line=([:solid :dash :dashdot]),
)
gui()

# plot the refractive index profile
plot(RIprofile(), sol)
gui()


text:
@doc raw"""
```
Z,p,q = retrieve_surface(algorithm::Shah, img::AbstractArray, iterations::Int=200)
```
Attempts to produce a hieghtmap from a grayscale image using Shah's algorithm.

Under the assumption that the albedo is constant, and the surface is Lambertian,
the algorithm employs discrete approximations of p and q using finite
differences in order to linearize the reflectance map in terms of Z by taking
the first term of its Taylor series expansion.
# Output
Returns an M by N array (matching dimensions of original image) of Float `Z`
that represents the reconstructed height at the point and the gradients in
M by N arrays of Float 'p' and 'q'.
# Details
Given the reflectance map for the surface defined as bellow:
```math
R(x,y)=\dfrac{-i_xp-i_yq+i_z}{\sqrt{1+p^2+q^2}}=\dfrac{\cos\sigma+p\cos\tau\sin
\sigma+q\sin\tau\sin\sigma}{\sqrt{1+p^2+q^2}}
```
where ``i_x=\frac{I(1)}{I(3)}=\frac{\cos\tau\sin\sigma}{\cos\sigma}=\cos\tau
\tan\sigma`` and ``i_y=\frac{I(2)}{I(3)}=\frac{\sin\tau\sin\sigma}{\cos\sigma}
=\sin\tau\tan\sigma``.

and p and q are discretely approximated as:
```math
\begin{gathered}
p=Z(x,y)-Z(x-1,y)\\
q=Z(x,y)-Z(x,y-1)\\
\end{gathered}
```

Shah linearized the function ``f=E-R=0`` in terms of ``Z`` in the vicinity of
``Z^{k-1}`` by crating a system of linear equations which can be solved iteratively
using the Jacobi iterative scheme, simplifying the Taylor series expansion to the
first order to get the following:
```math
f(Z(x,y))=0\approx f(Z^{n-1}(x,y))+(Z(x,y)-Z^{n-1}(x,y))\dfrac{df(Z^{n-1}(x,y))}
{dZ(x,y)}
```
which by letting ``Z^n(x,y)=Z(x,y)`` gives:
```math
Z^n(x,y=Z^{n-1}(x,y)-\dfrac{f(Z^{n-1}(x,y))}{\dfrac{df(Z^{n-1}(x,y))}{dZ(x,y)}}
```
where,
```math
\dfrac{df(Z^{n-1}(x,y))}{dZ(x,y)}=\dfrac{(p+q)(pi_x+qi_y+1)}{\sqrt{(1+p^2+q^2)^3}
\sqrt{1+i_x+i_y}}-\dfrac{i_x+i_y}{\sqrt{1+p^2+q^2}\sqrt{1+i_x+i_y}}
```
which as ``Z^0(x,y)=0``, allows the algorithm to iteratively solve for ``Z(x,y)``.

The `slant` and `tilt` can be manually defined using the function signature:
```
Z,p,q = retrieve_surface(algorithm::Shah, img::AbstractArray, slant::Real, tilt::Real, iterations::Int=200)
```
# Arguments
The function arguments are described in more detail below.
##  `img`
An `AbstractArray` storing the grayscale value of each pixel within
the range [0,1].
## `iterations`
An `Int` that specifies the number of iterations the algorithm is to perform. If
left unspecified a default value of 200 is used.
## `slant`
A `Real` that specifies the slant value to be used by the algorithm. The `slant`
should be a value in the range [0,π/2]. If `slant` is specified to must the `tilt`.
## `tilt`
A `Real` that specifies the tilt value to be used by the algorithm. The `tilt`
should be a value in the range [0,2π]. If `tilt` is specified to must
the `slant`.
# Example
Compute the heightmap for a synthetic image generated by `generate_surface`.
```julia
using Images, Makie, ShapeFromShading

#generate synthetic image
img = generate_surface(SynthSphere(), 1, [0.2,0,0.9], radius = 5)

#calculate the heightmap
Z,p,q = retrieve_surface(Shah(), img)

#normalize to maximum of 1 (not necessary but makes displaying easier)
Z = Z./maximum(Z)

#display using Makie (Note: Makie can often take several minutes first time)
r = 0.0:0.1:2
surface(r, r, Z)
```
# Reference
1. T. Ping-Sing and M. Shah, "Shape from shading using linear approximation", Image and Vision Computing, vol. 12, no. 8, pp. 487-498, 1994. [doi:10.1016/0262-8856(94)90002-7](https://doi.org/10.1016/0262-8856(94)90002-7)
"""
function retrieve_surface(algorithm::Shah, img::AbstractArray, iterations::Int=200)
    ρ,I,σ,τ = estimate_img_properties(img)
    return retrieve_surface(Shah(), img, σ, τ, iterations)
end

function retrieve_surface(algorithm::Shah, img::AbstractArray, slant::Real, tilt::Real, iterations::Int=200)
    σ, τ = slant,tilt
    E = Float64.(img)
    E = E .* 255
    M, N = size(E)
    p = zeros(axes(E))
    q = zeros(axes(E))
    Z = zeros(axes(E))
    Zx = zeros(axes(E))
    Zy = zeros(axes(E))
    ix = cos(τ) * tan(σ)
    iy = sin(τ) * tan(σ)
    R = zeros(axes(E))
    δfδZ = zeros(axes(E))
    f = zeros(axes(E))
    @inbounds for i = 1:iterations
        #calculate reflectance map
        for i in CartesianIndices(R)
            R[i] = (cos(σ) + p[i] * cos(τ) * sin(σ) + q[i] * sin(τ) * sin(σ)) /
                sqrt(1 + p[i]^2 + q[i]^2)

            R[i] = max(0, R[i])
            f[i] = E[i] - R[i]
        end

        #calculate derivative of f in respect to Z and update Z
        for i in CartesianIndices(δfδZ)
            δfδZ[i] = (p[i] + q[i]) * (ix * p[i] + iy * q[i] + 1) / (sqrt((1
                + p[i]^2 + q[i]^2)^3) * sqrt(1 + ix^2 + iy^2)) - (ix + iy) /
                (sqrt(1 .+ p[i]^2 + q[i]^2) * sqrt(1 + ix^2 + iy^2))

            Z[i] = Z[i] - f[i] / (δfδZ[i] + eps())
        end

        #update surface normals
        for i = 2:M
            for j = 1:N
                Zx[i,j] = Z[i-1,j]
            end
        end
        for i = 2:N
            for j = 1:M
                Zy[j,i] = Z[j,i-1]
            end
        end
        for i in CartesianIndices(p)
            p[i] = Z[i] - Zx[i]
            q[i] = Z[i] - Zy[i]
        end
    end

    #smooth Z using a 21X21 median filter of the absolute values
    for i in CartesianIndices(Z)
        Z[i] = abs(Z[i])
    end
    Z = mapwindow(median!, Z, (21,21))
    return Z, p, q
end


text:
# M is the number of positions to keep
# 𝓟 is a collection of candidate positions.
function removepointssimple!(𝓟::Vector{T}, 𝑒::Function, M::Int, h::T, tol::T) where T
    @assert length(𝓟) > M

    abs_𝑒_𝓟 = abs.(𝑒.(𝓟))

    if length(𝓟) > M
        # remove the positions associated with the smallest errors.
        ind_sort = sortperm(abs_𝑒_𝓟,rev=true)

        out = 𝓟[ind_sort[1:M]]

        resize!(𝓟,length(out))
        𝓟[:] = out
    end

    @assert length(𝓟) == M

    return nothing
end

function lg7(a::Vector{T}, N::Int) where T
    𝓟 = Vector{T}(undef,0)
    removepointsAlg7!(𝓟,a,N)

    return 𝓟
end

# Algorithm 7 from Filip's thesis.
# N is the number of points to keep.
function removepointsAlg7!(𝓟::Vector{T}, 𝑒::Function, N) where T

    ### preparation steps before removing positions.
    sort!(𝓟)

    # for each interval, find the position that has the maximum absolute error.
    𝓐 = forcealternating(𝓟, 𝑒)

    # This should be true since 𝓟 should include the previous reference set, which
    #   has N alternating error signs.
    @assert length(𝓐) >= N

    if length(𝓐) == N
        return 𝓐
    end

    ##
    N_remove = length(𝓐) - N

    # remove the first or last position if N_remove cannot be fulfilled by
    #   removing only pairs of positions.
    if isodd(N_remove)
        if abs( 𝑒(𝓐[1]) ) > abs( 𝑒(𝓐[end]) )
            deleteat!(𝓐, length(𝓐))
        else
            deleteat!(𝓐, 1)
        end
        N_remove -= 1
    end
    @assert iseven(N_remove)

    # iteratively prune pairs of values.
    while N_remove > 0

        tmp = collect( max( abs(𝑒(𝓐[i])), abs(𝑒(𝓐[i+1])) ) for i = 1:length(𝓐)-1 )
        min_e, min_ind = findmin(tmp)

        # check if the wrap-around pair has the smallest absolute error.
        if max( abs(𝑒(𝓐[1])), abs(𝑒(𝓐[end])) ) < min_e
            # remove the abs error from wrap-around positions.
            deleteat!(𝓐, (1, length(𝓐)))
        else
            # remove smallest abs error.
            deleteat!(𝓐, (min_ind, min_ind+1))
        end

        N_remove -= 2
    end

    @assert length(𝓐) == N
    return 𝓐
end

function addpointsinwithinterval(a, N)
    𝓟 = Vector{T}(undef,0)
    addpointsinwithinterval!(𝓟, a, N)

    return 𝓟
end

# N is the number of points to add to a.
# assumes a is sorted in ascending order.
function addpointsinwithinterval!(𝓟::Vector{T}, a::Vector{T}, N)::Nothing where T
    M = length(a)

    ### add at least s number of points between every pair of extrema_positions
    s = div(N,M-1)
    r = N - s*(M-1)

    # reset container.
    resize!(𝓟, 0)

    # add points in the first interval.
    push!(𝓟, a[1])

    st = a[1]

    fin = a[2]
    tmp = collect(LinRange(st,fin,s+r+2)) # +2 since we want the end points, which are the extrema.

    push!(𝓟, tmp[2:end]...)

    st = fin

    # add points for the rest of the intervals.
    for i = 2:M-1
        # add points.
        fin = a[i+1]
        tmp = collect(LinRange(st,fin,s+2)) # +2 since we want the end points, which are the extrema.
        push!(𝓟, tmp[2:end]...)

        # update.
        st = fin
    end
    @assert length(𝓟) == N+M

    return nothing
end

function isbettercandidate(a::T, b::T, f::Function, η::T = one(T)) where T
    score_a = abs(Calculus.derivative(f,a)) + one(T)/abs(f(a))
    score_b = abs(Calculus.derivative(f,b)) + one(T)/abs(f(b))

    return score_a < score_b
end

max_stars_repo_name: RoyCCWang/FIRRemez

max_issues_repo_name: RoyCCWang/FIRRemez

max_forks_repo_name: RoyCCWang/FIRRemez