text:
using WignerD, Test
using WignerD: NorthPole, SouthPole, Equator
using Aqua
using HalfIntegers
using OffsetArrays
using LinearAlgebra
using Rotations
using StaticArrays

@testset "project quality" begin
    if VERSION >= v"1.6.0"
        Aqua.test_all(WignerD)
    else
        Aqua.test_all(WignerD, ambiguities = false)
    end
end

@testset "trigonometric functions for special points" begin
	@testset "Equator" begin
        @test promote_rule(Equator, Float64) == Float64
	    for α = -10:10
	    	@test (@inferred WignerD._sincos(α, Equator()))[1] ≈ sincos(α*π/2)[1] atol=1e-14 rtol=1e-8
	    	@test (@inferred WignerD._sincos(α, Equator()))[2] ≈ sincos(α*π/2)[2] atol=1e-14 rtol=1e-8
            @test (@inferred WignerD._sincos(half(α), Equator()))[1] ≈ sincos(half(α)*π/2)[1] atol=1e-14 rtol=1e-8
            @test (@inferred WignerD._sincos(half(α), Equator()))[2] ≈ sincos(half(α)*π/2)[2] atol=1e-14 rtol=1e-8
            @test sincos(α*Equator()) == sincos(α*π/2)
	    end
    	@test @inferred cos(Equator()) == 0
    	@test @inferred sin(Equator()) == 1
    	@test @inferred float(Equator()) == π/2
    	@test @inferred AbstractFloat(Equator()) == π/2
    	@test @inferred Float64(Equator()) == π/2
	end
    @testset "NorthPole" begin
        @test @inferred cos(NorthPole()) == 1
        @test @inferred sin(NorthPole()) == 0
        @test @inferred float(NorthPole()) == 0
        @test @inferred AbstractFloat(NorthPole()) == 0
        @test @inferred Float64(NorthPole()) == 0
    end
    @testset "SouthPole" begin
        @test @inferred cos(SouthPole()) == -1
        @test @inferred sin(SouthPole()) == 0
        @test @inferred float(SouthPole()) == float(pi)
        @test @inferred AbstractFloat(SouthPole()) == float(pi)
        @test @inferred Float64(SouthPole()) == float(pi)
    end
end

@testset "wignerdjmn" begin

    function testapprox(m, n, dj_m_n, dj_m_n2)
    	@test begin
    		res = isapprox(dj_m_n, dj_m_n2, atol=1e-14, rtol=sqrt(eps(Float64)))
    		if !res
    			@show m n dj_m_n dj_m_n2
    		end
    		res
    	end
    end

	function test_special(j, Jy...)
        @testset "j = $j" begin
            @testset "Equator" begin
                for m in -j:j, n in -j:j
                    dj_m_n = @inferred WignerD.wignerdjmn(j, m, n, π/2, Jy...)
                    dj_m_n2 = @inferred WignerD.wignerdjmn(j, m, n, Equator(), Jy...)
                    testapprox(m, n, dj_m_n, dj_m_n2)
                end
            end
            @testset "NorthPole" begin
                for m in -j:j, n in -j:j
                    dj_m_n = @inferred WignerD.wignerdjmn(j, m, n, 0, Jy...)
                    dj_m_n2 = @inferred WignerD.wignerdjmn(j, m, n, NorthPole(), Jy...)
                    testapprox(m, n, dj_m_n, dj_m_n2)
                end
            end
            @testset "SouthPole" begin
                for m in -j:j, n in -j:j
                    dj_m_n = @inferred WignerD.wignerdjmn(j, m, n, π, Jy...)
                    dj_m_n2 = @inferred WignerD.wignerdjmn(j, m, n, SouthPole(), Jy...)
                    testapprox(m, n, dj_m_n, dj_m_n2)
                end
            end
        end
    end

    for j in 0:3
        test_special(j)
        Jy = zeros(ComplexF64, 2j+1, 2j+1)
        test_special(j, Jy)
    end
    for j in half(1):1:half(7)
        test_special(j)
        Jy = zeros(ComplexF64, Int(2j+1), Int(2j+1))
        test_special(j, Jy)
    end

    @testset "wignerDjmn" begin
        for j in 0:3
            Jy = zeros(ComplexF64, 2j+1, 2j+1)
            for β in LinRange(0, pi, 10), m in -j:j, n in -j:j
                Djmn = WignerD.wignerDjmn(j, m, n, 0, β, 0)
                Djmn2 = WignerD.wignerDjmn(j, m, n, 0, β, 0, Jy)
                @test Djmn == Djmn2
                djmn = WignerD.wignerdjmn(j, m, n, β)
                @test isapprox(Djmn, djmn, atol = 1e-14, rtol = sqrt(eps(Float64)))
            end
        end
        for j in half(1):1:half(7)
            Jy = zeros(ComplexF64, Int(2j+1), Int(2j+1))
            for β in LinRange(0, pi, 10), m in -j:j, n in -j:j
                Djmn = WignerD.wignerDjmn(j, m, n, 0, β, 0)
                Djmn2 = WignerD.wignerDjmn(j, m, n, 0, β, 0, Jy)
                @test Djmn == Djmn2
                djmn = WignerD.wignerdjmn(j, m, n, β)
                @test isapprox(Djmn, djmn, atol = 1e-14, rtol = sqrt(eps(Float64)))
            end
        end
    end
end

function testwignerd(testelements, j, θ)
    d = wignerd(j, θ)
    testelements(WignerD._offsetmatrix(j, d), θ)
    @test wignerd(HalfInt(j), θ) ≈ d atol=1e-14 rtol=1e-8
    @test isapprox(tr(d), sum(cos(m*θ) for m in -j:j), atol = 1e-14, rtol = sqrt(eps(Float64)))
    @test det(d) ≈ 1
end

@testset "d0_mn(θ)" begin
    function test(d, θ)
        @test isapprox(d[0, 0], 1, atol=1e-14, rtol=1e-8)
    end

    n = 100
    j = 0

    for θ in LinRange(0, π, n)
        testwignerd(test, j, θ)
    end

    @testset "Equator" begin
        testwignerd(test, j, Equator())
    end

    @testset "NorthPole" begin
        testwignerd(test, j, NorthPole())
    end

    @testset "SouthPole" begin
        testwignerd(test, j, SouthPole())
    end
end

@testset "d1/2_mn(θ)" begin
    function test(d, θ)
        @test isapprox(d[-1/2, -1/2], cos(θ/2), atol=1e-14, rtol=1e-8)
        @test isapprox(d[1/2, -1/2], -sin(θ/2), atol=1e-14, rtol=1e-8)

        @test isapprox(d[-1/2, 1/2], sin(θ/2), atol=1e-14, rtol=1e-8)
        @test isapprox(d[1/2, 1/2], cos(θ/2), atol=1e-14, rtol=1e-8)
    end

    n = 100
    j = half(1)

    for θ in LinRange(0, π, n)
        testwignerd(test, j, θ)
    end

    @testset "Equator" begin
        testwignerd(test, j, Equator())
    end

    @testset "NorthPole" begin
        testwignerd(test, j, NorthPole())
    end

    @testset "SouthPole" begin
        testwignerd(test, j, SouthPole())
    end
end

@testset "d1_mn(θ)" begin
	function test(d, θ)
		@test isapprox(d[-1, -1], (1+cos(θ))/2, atol=1e-14, rtol=1e-8)
		@test isapprox(d[ 0, -1], -sin(θ)/√2, atol=1e-14, rtol=1e-8)
		@test isapprox(d[ 1, -1], (1-cos(θ))/2, atol=1e-14, rtol=1e-8)

		@test isapprox(d[-1, 0],  sin(θ)/√2, atol=1e-14, rtol=1e-8)
		@test isapprox(d[ 0, 0],  cos(θ), atol=1e-14, rtol=1e-8)
		@test isapprox(d[ 1, 0], -sin(θ)/√2, atol=1e-14, rtol=1e-8)

		@test isapprox(d[-1, 1], (1-cos(θ))/2, atol=1e-14, rtol=1e-8)
		@test isapprox(d[ 0, 1],  sin(θ)/√2, atol=1e-14, rtol=1e-8)
        @test isapprox(d[ 1, 1], (1+cos(θ))/2, atol=1e-14, rtol=1e-8)
	end

	n = 100
    j = 1

	for θ in LinRange(0, π, n)
		testwignerd(test, j, θ)
	end

	@testset "Equator" begin
        testwignerd(test, j, Equator())
    end

    @testset "NorthPole" begin
        testwignerd(test, j, NorthPole())
    end

    @testset "SouthPole" begin
        testwignerd(test, j, SouthPole())
    end
end

@testset "d3/2_mn(θ)" begin
    function test(d, θ)
        @test isapprox(d[-3/2, -3/2],  cos(θ/2)^3, atol=1e-14, rtol=1e-8)
        @test isapprox(d[-1/2, -3/2], -√3*sin(θ/2)cos(θ/2)^2, atol=1e-14, rtol=1e-8)
        @test isapprox(d[ 1/2, -3/2],  √3*sin(θ/2)^2*cos(θ/2), atol=1e-14, rtol=1e-8)
        @test isapprox(d[ 3/2, -3/2],  -sin(θ/2)^3, atol=1e-14, rtol=1e-8)

        @test isapprox(d[-3/2, -1/2],  √3*sin(θ/2)cos(θ/2)^2, atol=1e-14, rtol=1e-8)
        @test isapprox(d[-1/2, -1/2],  cos(θ/2)*(3cos(θ/2)^2 - 2), atol=1e-14, rtol=1e-8)
        @test isapprox(d[ 1/2, -1/2],  sin(θ/2)*(3sin(θ/2)^2 - 2), atol=1e-14, rtol=1e-8)
        @test isapprox(d[ 3/2, -1/2],  √3*sin(θ/2)^2*cos(θ/2), atol=1e-14, rtol=1e-8)

        @test isapprox(d[-3/2,  1/2],  √3*sin(θ/2)^2*cos(θ/2), atol=1e-14, rtol=1e-8)
        @test isapprox(d[-1/2,  1/2],  -sin(θ/2)*(3sin(θ/2)^2 - 2), atol=1e-14, rtol=1e-8)
        @test isapprox(d[ 1/2,  1/2],  cos(θ/2)*(3cos(θ/2)^2 - 2), atol=1e-14, rtol=1e-8)
        @test isapprox(d[ 3/2,  1/2], -√3*sin(θ/2)cos(θ/2)^2, atol=1e-14, rtol=1e-8)

        @test isapprox(d[-3/2,  3/2],  sin(θ/2)^3, atol=1e-14, rtol=1e-8)
        @test isapprox(d[-1/2,  3/2],  √3*sin(θ/2)^2*cos(θ/2), atol=1e-14, rtol=1e-8)
        @test isapprox(d[ 1/2,  3/2],  √3*sin(θ/2)cos(θ/2)^2, atol=1e-14, rtol=1e-8)
        @test isapprox(d[ 3/2,  3/2],  cos(θ/2)^3, atol=1e-14, rtol=1e-8)
    end

    n = 100
    j = 3/2

    for θ in LinRange(0, π, n)
        testwignerd(test, j, θ)
    end

    @testset "Equator" begin
        testwignerd(test, j, Equator())
    end

    @testset "NorthPole" begin
        testwignerd(test, j, NorthPole())
    end

    @testset "SouthPole" begin
        testwignerd(test, j, SouthPole())
    end
end

@testset "d2_mn(θ)" begin
	function test(d, θ)
		@test isapprox(d[-2, -2],  (1+cos(θ))^2/4, atol=1e-14, rtol=1e-8)
		@test isapprox(d[-1, -2], -sin(θ)*(1+cos(θ))/2, atol=1e-14, rtol=1e-8)
		@test isapprox(d[ 0, -2],  1/2*√(3/2)*sin(θ)^2, atol=1e-14, rtol=1e-8)
		@test isapprox(d[ 1, -2], -sin(θ)*(1-cos(θ))/2, atol=1e-14, rtol=1e-8)
		@test isapprox(d[ 2, -2],  (1-cos(θ))^2/4, atol=1e-14, rtol=1e-8)

        @test isapprox(d[-2, -1],  sin(θ)*(1+cos(θ))/2, atol=1e-14, rtol=1e-8)
        @test isapprox(d[-1, -1],  (2cos(θ)^2+cos(θ)-1)/2, atol=1e-14, rtol=1e-8)
        @test isapprox(d[ 0, -1], -√(3/2)*sin(θ)*cos(θ), atol=1e-14, rtol=1e-8)
        @test isapprox(d[ 1, -1], -(2cos(θ)^2-cos(θ)-1)/2, atol=1e-14, rtol=1e-8)
        @test isapprox(d[ 2, -1], -sin(θ)*(1-cos(θ))/2, atol=1e-14, rtol=1e-8)

        @test isapprox(d[-2, 0],  1/2*√(3/2)*sin(θ)^2, atol=1e-14, rtol=1e-8)
        @test isapprox(d[-1, 0],  √(3/2)*sin(θ)*cos(θ), atol=1e-14, rtol=1e-8)
        @test isapprox(d[ 0, 0],  1/2*(3cos(θ)^2-1), atol=1e-14, rtol=1e-8)
        @test isapprox(d[ 1, 0], -√(3/2)*sin(θ)*cos(θ), atol=1e-14, rtol=1e-8)
        @test isapprox(d[ 2, 0],  1/2*√(3/2)*sin(θ)^2, atol=1e-14, rtol=1e-8)

        @test isapprox(d[-2, 1],  sin(θ)*(1-cos(θ))/2, atol=1e-14, rtol=1e-8)
        @test isapprox(d[-1, 1], -(2cos(θ)^2-cos(θ)-1)/2, atol=1e-14, rtol=1e-8)
        @test isapprox(d[ 0, 1],  √(3/2)*sin(θ)*cos(θ), atol=1e-14, rtol=1e-8)
        @test isapprox(d[ 1, 1],  (2cos(θ)^2+cos(θ)-1)/2, atol=1e-14, rtol=1e-8)
		@test isapprox(d[ 2, 1], -sin(θ)*(1+cos(θ))/2, atol=1e-14, rtol=1e-8)

		@test isapprox(d[-2, 2],  (1-cos(θ))^2/4, atol=1e-14, rtol=1e-8)
		@test isapprox(d[-1, 2],  sin(θ)*(1-cos(θ))/2, atol=1e-14, rtol=1e-8)
		@test isapprox(d[ 0, 2],  1/2*√(3/2)*sin(θ)^2, atol=1e-14, rtol=1e-8)
		@test isapprox(d[ 1, 2],  sin(θ)*(1+cos(θ))/2, atol=1e-14, rtol=1e-8)
        @test isapprox(d[ 2, 2],  (1+cos(θ))^2/4, atol=1e-14, rtol=1e-8)
	end

    n = 100
	j = 2

    for θ in LinRange(0, π, n)
        testwignerd(test, j, θ)
    end

    @testset "Equator" begin
        testwignerd(test, j, Equator())
    end

    @testset "NorthPole" begin
        testwignerd(test, j, NorthPole())
    end

    @testset "SouthPole" begin
        testwignerd(test, j, SouthPole())
    end
end

@testset "symmetry" begin
    for β in LinRange(0, pi, 10), j in 0:half(1):5
        djβ = wignerd(j, β)
        djβ_offset = WignerD._offsetmatrix(j, djβ)
        @test wignerd(j, -β)' ≈ djβ
        for k in 1:4
            @test wignerd(j, β + 2pi * k) ≈ (-1)^(2j*k) * djβ
            @test wignerd(j, β - 2pi * k) ≈ (-1)^(2j*k) * djβ

            dj2np1pi = WignerD._offsetmatrix(j, wignerd(j, β + (2k + 1)*pi))
            for m in -j:j, n in -j:j
                @test isapprox(dj2np1pi[m, n], (-1)^((2k+1)*j - n)*djβ_offset[m, -n], atol =1e-14, rtol = 1e-8)
            end
            dj2np1pi = WignerD._offsetmatrix(j, wignerd(j, β - (2k + 1)*pi))
            for m in -j:j, n in -j:j
                @test isapprox(dj2np1pi[m, n], (-1)^(-(2k+1)*j - n)*djβ_offset[m, -n], atol =1e-14, rtol = 1e-8)
            end
        end
        djpimβ = WignerD._offsetmatrix(j, wignerd(j, π - β))
        for m in -j:j, n in -j:j
            @test isapprox(djpimβ[m, n], (-1)^(j + m) * djβ_offset[m, -n], atol = 1e-14, rtol = 1e-8)
            @test isapprox(djpimβ[m, n], (-1)^(j - n) * djβ_offset[-m, n], atol = 1e-14, rtol = 1e-8)
        end
    end
end

@testset "WignerD" begin
    for β in LinRange(0, pi, 10), j in 0:half(1):3
        d = wignerd(j, β)
        D = wignerD(j, 0, β, 0)
        @test d == D
        for α in LinRange(0, 2pi, 10), γ in LinRange(0, 2pi, 10)
            D = wignerD(j, α, β, γ)
            D_w = WignerD._offsetmatrix(j, D)
            for n in -j:j, m in -j:j
                @test isapprox(D_w[m, n], WignerD.wignerDjmn(j, m, n, α, β, γ), atol=1e-14, rtol = sqrt(eps(Float64)))
            end
            @test det(D) ≈ 1
            @test D' * D ≈ Diagonal(ones(Int(2j+1)))

            Dinv = wignerD(j, -γ, -β, -α)
            @test inv(D) ≈ Dinv
            @test D' ≈ Dinv
        end
    end
end

@testset "WignerMatrix" begin
    w = WignerD._offsetmatrix(0.5, wignerd(0.5, 0))
    @test IndexStyle(w) == IndexStyle(parent(w))
    @test size(w) == (2,2)
end

@testset "comparison with Cartesian rotations" begin
    # The spherical covariant basis transforms through D¹*
    # The Cartesian basis transforms through RotZYZ
    # The two may be related through U * RotZYZ(α,β,γ) * U' == D¹*(α,β,γ)
    # where U is the matrix that converts from Cartesian to spherical covariant bases
    # U * [ex, ey, ez] == [χ₋₁, χ₀, χ₊₁]
    U = SMatrix{3,3}(
        [
            1/√2   -im/√2   0
            0       0       1
            -1/√2  -im/√2   0
        ]
    )
    j = 1
    Jy = zeros(ComplexF64, 3, 3)
    D = zeros(ComplexF64, 3, 3)
    for α in LinRange(0, 2pi, 10), β in LinRange(0, 2pi, 10), γ in LinRange(0, 2pi, 10)
        wignerD!(D, j, α, β, γ)
        conj!(D)
        @test D ≈ U * RotZYZ(α, β, γ) * U'
    end
end


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

using Markdown
using InteractiveUtils

# ╔═╡ b53dd37e-0dc5-11eb-37ee-555eed0b2c9f
using CSV, OffsetArrays, Plots, Plots.Measures, PyCall, SplitApplyCombine, Statistics

# ╔═╡ 4592edbe-2edd-11eb-2ed0-679711b2d714
# Code to load CSC matrix from NPZ file
# based on description here: https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.csc_matrix.html
# see also:
# + https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.save_npz.html
# + https://docs.julialang.org/en/v1/stdlib/SparseArrays/#man-csc

module NPZ_CSC
using PyCall: pyimport, pyrepr
using SparseArrays: SparseMatrixCSC

np = pyimport("numpy")

function load(filename)
	contents = np.load(filename, allow_pickle=true)
	format = get(contents, "format")
	pyrepr(format) == "array(b'csc', dtype='|S3')" ||
		throw("Array format seems wrong. Got $(pyrepr(format))")

	shape = get(contents, "shape")
	data = get(contents, "data")
	indptr = get(contents, "indptr") .+ 1
	indices = get(contents, "indices") .+ 1

	SparseMatrixCSC(shape[1], shape[2], indptr, indices, data)
end

end

# ╔═╡ eaffbf72-7bed-11eb-2b72-b759220943ef
md"""
# Resource analysis figure

Un/comment the relevant line for each choice of days.
"""

# ╔═╡ e88985de-2f4e-11eb-32cf-1f7a851792bb
# prevalence ≈ 0.033808% - 2.459128%
t0, t1, q1pos, q2pos, q3pos, pbestpos = 40, 90, 8.5, 7, 4, missing

# prevalence ≈ 0.103552%
# t0, t1, q1pos, q2pos, q3pos, pbestpos = 53, 53, 8, 5.5, 3, missing

# prevalence ≈ 1.055872%
# t0, t1, q1pos, q2pos, q3pos, pbestpos = 80, 80, 8.5, 7, 4, missing

# prevalence ≈ 1.362224%
# t0, t1, q1pos, q2pos, q3pos, pbestpos = 83, 83, 8.5, 1, 4, missing

# prevalence ≈ 1.484104%
# t0, t1, q1pos, q2pos, q3pos, pbestpos = 84, 84, 8.5, 1.5, missing, 5

# prevalence ≈ 2.459128%
# t0, t1, q1pos, q2pos, q3pos, pbestpos = 90, 90, 9, 1.5, missing, 5

# prevalence ≈ 3.149712%
# t0, t1, q1pos, q2pos, q3pos, pbestpos = 93, 93, 9, 1.5, 5, missing

# ╔═╡ a17ac18e-7bee-11eb-1624-fde146fd6293
swabsweep = [12,24,48,96,192,384,768,1536,3072,6144]

# ╔═╡ a45a36b4-7bee-11eb-18d7-9f83f59c66fa
testsweep = [12,24,48,96,192,384,768,1536,3072,6144]

# ╔═╡ 6e2887ba-0e95-11eb-1475-e7ea7e5cf740
budgetsweep = collect(Iterators.product(swabsweep, testsweep));

# ╔═╡ 08baf128-2958-11eb-294d-35b2a6eeef5d
selectedcells = [(9,1),(9,2),(9,3),(9,7),(4,2),(6,3)]

# ╔═╡ 9b6802e0-1558-11eb-1f6e-b5f67e4cafd3
methodlist = ["individual","hypergraph","array_comb","pbest"]

# ╔═╡ f7be67a4-7bed-11eb-032c-e9f443225026
md"""
## Utility functions
"""

# ╔═╡ 213f310c-155e-11eb-0720-b748da5a582d
module FindMin
import Base.findmin, Base.findmax, Base.argmin, Base.argmax
isgreater(x, y) = _is_reflexive(x) && _is_reflexive(y) ? isless(y, x) : isless(x, y)
_is_reflexive(x) = let eq = x == x
	isa(eq, Bool) && eq
end
findmax(f, domain) = mapfoldl(x -> (f(x), x), _rf_findmax, domain)
_rf_findmax((fm, m), (fx, x)) = isless(fm, fx) ? (fx, x) : (fm, m)
findmin(f, domain) = mapfoldl(x -> (f(x), x), _rf_findmin, domain)
_rf_findmin((fm, m), (fx, x)) = isgreater(fm, fx) ? (fx, x) : (fm, m)
argmax(f, domain) = findmax(f, domain)[2]
argmin(f, domain) = findmin(f, domain)[2]
end

# ╔═╡ 2e6ae32c-7bee-11eb-12b2-a146668beb80
md"""
## Load results
"""

# ╔═╡ 685aa44c-2edd-11eb-237f-6715b380cb39
ViralLoad = NPZ_CSC.load("../simulation/Simulated_populations/seir_viral_loads_swab.viral_loads.npz")

# ╔═╡ 6ec783d4-2edd-11eb-2be6-7f0631417b11
prevalence = OffsetArray(vec(mean(>(0), ViralLoad, dims=1)), -1)

# ╔═╡ b37b43e4-2edd-11eb-23a1-e1c242eb0029
prev0 = round(prevalence[t0]*100,digits=2)

# ╔═╡ bf7652c6-2edd-11eb-05df-33027bb8ab4e
prev1 = round(prevalence[t1]*100,digits=2)

# ╔═╡ 0faadc44-0dc6-11eb-0c3f-db78ba861a01
results = CSV.File("summary.resource_t0-$(t0)_t1-$(t1).csv")

# ╔═╡ d33a506c-0e2b-11eb-24a6-41ddae38afa5
results_methods = group(res -> res[Symbol("Design method")], results)

# ╔═╡ d1b9df5a-0e2b-11eb-1767-e7762414041c
begin
	results_methods_budget = map(results_methods) do results_method
		map(budgetsweep) do (swabs, tests)
			filter(results_method) do res
				res[Symbol("Sample budget")] == swabs && res[Symbol("Test budget")] == tests
			end
		end
	end
	array_comb = map(results_methods_budget["array_8x12"],results_methods_budget["array_16x24"]) do a1,a2
		a = [a1;a2]
		length(a) <= 1 ? a : [FindMin.argmax(r->r[Symbol("Design effectiveness")],a)]
	end
	insert!(results_methods_budget,"array_comb",array_comb)
end

# ╔═╡ ccdb2506-0e95-11eb-34b4-8db7c3d43ead
efftests = map(results_methods_budget) do res
	map(r -> length(r) == 0 ? 0.0 : only(r)[Symbol("Design effectiveness")], res)
end

# ╔═╡ c3cc7816-2943-11eb-2807-c742e3307662
bestresults = map(Tuple.(CartesianIndices(budgetsweep))) do (swabidx,testidx)
	bestmethod = FindMin.argmax(method->efftests[method][swabidx,testidx], methodlist)
	only(results_methods_budget[bestmethod][swabidx,testidx])
end

# ╔═╡ e3aa4e4a-7be8-11eb-1a29-239b9fbce803
md"""
## Figs a-f
"""

# ╔═╡ 3d2aa366-295f-11eb-24be-f93fa0a0a9de
methodlabels = Dict(
	"individual"=>"Indiv.",
	"hypergraph"=>"HYPER",
	"array_comb"=>"Array",
	"pbest"=>"P-BEST"
)

# ╔═╡ eb95975e-2959-11eb-22af-8dc3c3bd4a3c
selectedplts = map(enumerate(selectedcells)) do (t,(swabidx,testidx))
	result = bestresults[swabidx,testidx]
	method = result[Symbol("Design method")]
	if method == "individual"
		bestq = 0
	elseif method == "hypergraph"
		bestq = result[Symbol("Design sample split")]
	else
		bestq = missing
	end
	ymax = 1.2*result[Symbol("Design effectiveness")]

	methodeffs = [efftests[method][swabidx,testidx] for method in methodlist]
	effticks = [0,round(result[Symbol("Design effectiveness")],digits=1)]
	effticklabels = effticks
	bar([methodlabels[method] for method in methodlist],methodeffs,
		bordercolor=:black,
		framestyle=:box,xgrid=nothing,label="",size=(325,250),
		bar_width=0.99,
		ylims=(0,ymax),xlims=(-0.01,4.01),
		ylabel="Effective screening capacity",
		yticks=(effticks,string.(" ",effticklabels)),tick_direction=:out)

	swabs, tests = result[Symbol("Sample budget")], result[Symbol("Test budget")]
	title!("$swabs samples, $tests tests")

	for (idx,method) in enumerate(methodlist)
		iszero(methodeffs[idx]) && continue
		annotate!([(idx-0.5,methodeffs[idx]+0.05*ymax,Plots.text("$(round(methodeffs[idx],digits=1))",6,:black,:center))])

		method == "individual" && continue
		res = only(results_methods_budget[method][swabidx,testidx])
		n = res[Symbol("Design samples")]
		m = res[Symbol("Design pools")]
		q = res[Symbol("Design sample split")]
		b = res[Symbol("Design runs per day")]

		if method == "hypergraph"
			if q == 1
				n = convert(Int,n / m)
				b = b*m
				m = 1
			end
			annotate!([(idx-0.5,methodeffs[idx]-0.07ymax,Plots.text("n:$n",5,:white,:center))])
			annotate!([(idx-0.5,methodeffs[idx]-0.15ymax,Plots.text("m:$m",5,:white,:center))])
			annotate!([(idx-0.5,methodeffs[idx]-0.23ymax,Plots.text("q:$q",5,:white,:center))])
		elseif method == "array_comb"
			arraylabel = split(only(results_methods_budget["array_comb"][swabidx,testidx])[Symbol("Design method")],'_')[end]
			annotate!([(idx-0.5,methodeffs[idx]-0.07*ymax,Plots.text("$arraylabel",5,:white,:center))])
		end
		annotate!([(idx-0.5,0.07*ymax,Plots.text("b:$(isinteger(b) ? convert(Int,b) : round(b,digits=1))",5,:black,:center))])
	end

	annotate!(-1.05,1.07*ymax,text("ABCDEFGHIJKLMNOPQRSTUVWXYZ"[t] |> lowercase,12))

	plot!()
end |> plts -> plot(plts...,layout=(3,2),link=:x,xtickfontsize=6,ytickfontsize=6,titlefontsize=7,guidefontsize=6,dpi=150,size=(375,500),top_margin=-2mm,bottom_margin=0.5mm,left_margin=0.0mm)

# ╔═╡ cc143562-2967-11eb-06fe-d7f4cd31664d
savefig(selectedplts,"res,analysis,selected,t0-$(t0),t1-$(t1).png")

# ╔═╡ e2fff794-7be8-11eb-0bfb-a17b3f02641e
md"""
## Fig g
"""

# ╔═╡ 8b5a9e38-7be9-11eb-1d71-f5529db9d4e8
HEATMAP_PALETTE = palette(:OrRd_4,4)

# ╔═╡ 422424ee-3032-11eb-2454-972aa1343745
pbestcolor=colorant"forestgreen"

# ╔═╡ fc03d208-5f2e-11eb-366e-45ca469d44ed
arraycolor=colorant"dodgerblue"

# ╔═╡ 17caf5de-2951-11eb-10b4-89609425f75a
bestfig = let cpal = palette([collect(HEATMAP_PALETTE); pbestcolor])
	bestq = map(bestresults) do result
		method = result[Symbol("Design method")]
		if method == "individual"
			return 0
		elseif method == "hypergraph" || method == "dorfman"
			return result[Symbol("Design sample split")]
		elseif method == "pbestdecoder"
			return 4
		else
			return missing
		end
	end

	heatmap(permutedims(bestq),colorbar=nothing,color=cpal,clims=(0,4),aspect_ratio=1,
		size=(550,500),dpi=150,right_margin=15mm,grid=nothing,
		xlims=(0.5,length(swabsweep)+0.5),ylims=(0.5,length(testsweep)+0.5),
		xlabel="Daily sample collection capacity (average)",ylabel="Daily test capacity (average)",guidefontsize=10,
		xticks=(axes(swabsweep,1),swabsweep),yticks=(axes(testsweep,1),testsweep),
		framestyle=:box)

	if !ismissing(pbestpos)
		annotate!([(10.75,pbestpos, Plots.text("P-BEST",10,:left,cpal[5]))])
	end

	if !ismissing(q1pos)
		annotate!([(10.75,q1pos+0.2, Plots.text("HYPER",10,:left,cpal[2]))])
		annotate!([(10.75,q1pos-0.2, Plots.text("(q=1)",10,:left,cpal[2]))])
	end

	if !ismissing(q2pos)
		annotate!([(10.75,q2pos+0.2, Plots.text("HYPER",10,:left,cpal[3]))])
		annotate!([(10.75,q2pos-0.2, Plots.text("(q=2)",10,:left,cpal[3]))])
	end

	if !ismissing(q3pos)
		annotate!([(10.75,q3pos+0.2, Plots.text("HYPER",10,:left,cpal[4]))])
		annotate!([(10.75,q3pos-0.2, Plots.text("(q=3)",10,:left,cpal[4]))])
	end

	bestefftests = getindex.(bestresults,Symbol("Design effectiveness"))
	for idx in CartesianIndices(bestefftests)
		plot!(Shape([(idx[1]-0.5,idx[2]-0.5),(idx[1]+0.5,idx[2]-0.5),(idx[1]+0.5,idx[2]+0.5),(idx[1]-0.5,idx[2]+0.5)]),
			label="",fillcolor=nothing,linewidth=0.5,linealpha=0.2)
	end

	for (t,idx) in enumerate(selectedcells)
		plot!(Shape([(idx[1]-0.5,idx[2]-0.5),(idx[1]+0.5,idx[2]-0.5),(idx[1]+0.5,idx[2]+0.5),(idx[1]-0.5,idx[2]+0.5)]),label="",
			fillcolor=nothing,linecolor=:black,linewidth=2)
	end

	for idx in CartesianIndices(bestefftests)
		idx[1] >= idx[2] && annotate!([(idx[1],idx[2]+0.2,Plots.text(round(bestefftests[idx],digits=1),7,:black,:center))])
		if idx[1] > idx[2]
			n, m, q = bestresults[idx][Symbol("Design samples")], bestresults[idx][Symbol("Design pools")], bestresults[idx][Symbol("Design sample split")]
			method = bestresults[idx][Symbol("Design method")]
			if method == "pbestdecoder"
				annotate!([(idx[1],idx[2]-0.20,Plots.text("P-BEST",5,:white,:center))])
				continue
			end
			if q == 1
				n = convert(Int,n / m)
				m = 1
			end
			annotate!([(idx[1],idx[2]-0.11,Plots.text("n=$n",5,ismissing(bestq[idx]) ? :black : :white,:center))])
			annotate!([(idx[1],idx[2]-0.29,Plots.text("m=$m",5,ismissing(bestq[idx]) ? :black : :white,:center))])
		end
	end
	for idx in axes(bestresults,1)[begin:end-1]
		plot!([(idx,idx+0.35),(idx,idx+1.25)],arrow=1,color=:black,label="")
	end

	annotate!([(3.5    ,7.5    ,Plots.text("Testing-rich regime:",     12,:black,:center,rotation=45))])
	annotate!([(3.5+0.4,7.5-0.4,Plots.text("individual testing optimal",12,:black,:center,rotation=45))])

	annotate!(-0.5,10.95,text("ABCDEFGHIJKLMNOPQRSTUVWXYZ"[length(selectedcells)+1] |> lowercase,12))

	if t0 == t1
		plot!(title="Day $t0: prevalence of $prev0%",titlefontsize=8)
	else
		plot!(title="Days $t0-$t1: prevalence grows exponentially from $prev0% to $prev1%",titlefontsize=8)
	end
	plot!(size=(550,500),bottom_margin=-8mm,left_margin=-2mm,top_margin=-8mm)
end

# ╔═╡ 20d151f0-2951-11eb-124b-b124fe93a1c8
savefig(bestfig,"res,analysis,bestmethods,t0-$(t0),t1-$(t1).png")

# ╔═╡ 4526a4f6-7bef-11eb-17e6-61b5bd0a2f00
md"""
## Fig h
"""

# ╔═╡ 85b7b0d6-1962-11eb-33c2-b759b672eac8
resfig = map(Tuple.(CartesianIndices(budgetsweep))) do (swabidx,testidx)
	swabs, tests = budgetsweep[swabidx,testidx]
	bestmethod = FindMin.argmax(method->efftests[method][swabidx,testidx], methodlist)
	bestresult = only(results_methods_budget[bestmethod][swabidx,testidx])
	n, m, q = bestresult[Symbol("Design samples")], bestresult[Symbol("Design pools")], bestresult[Symbol("Design sample split")]
	if q == 1
		n = convert(Int,n / m)
		m = 1
	end
	besteff = bestresult[Symbol("Design effectiveness")]
	hyperq = only(results_methods_budget["hypergraph"][swabidx,testidx])[Symbol("Design sample split")]
	bar([efftests[method][swabidx,testidx] for method in methodlist],
		color=[HEATMAP_PALETTE[1],HEATMAP_PALETTE[hyperq+1],arraycolor,pbestcolor],
		bar_width=0.99,xlims=(0.47,4.53),
		size=(100,100), ylims=(0,1.3*besteff), label="", framestyle=:box,grid=nothing,
		xticks=(testidx == 1 ? (2.5,swabs) : nothing),yticks=(swabidx == 1 ? (0.65*besteff,tests) : nothing),
	)
	hline!([besteff],color=:black,opacity=0.2,label="")
	annotate!([(2.5,1.16*besteff,Plots.text(round(besteff,digits=1),7,:black,:center))])
	annotate!([(2,efftests["hypergraph"][swabidx,testidx]/2, text("n:$n, m:$m", 5, :white, :center, rotation=90))])
	efftests["individual"][swabidx,testidx] > 0.25*besteff && annotate!([(1,efftests["individual"][swabidx,testidx]/2, text("Ind.", 5, :black, :center, rotation=90))])
	efftests["pbest"][swabidx,testidx] > 0 && annotate!([(4,efftests["pbest"][swabidx,testidx]/2, text("P-BEST", 5, :white, :center, rotation=90))])
	if efftests["array_comb"][swabidx,testidx] > 0
		arraylabel = split(only(results_methods_budget["array_comb"][swabidx,testidx])[Symbol("Design method")],'_')[end]
		annotate!([(3,efftests["array_comb"][swabidx,testidx]/2, text(arraylabel, 5, :white, :center, rotation=90))])
	end
	xlabel!(swabidx == 5 && testidx == 1 ? "Daily sample collection capacity (average)" : "",labelfontsize=10)
	ylabel!(testidx == 5 && swabidx == 1 ? "Daily test capacity (average)" : "",labelfontsize=10)

	swabidx == 1 && testidx == 10 && !(t0 == 40 && t1 == 90) && annotate!(-2.4,1.6*besteff,text("ABCDEFGHIJKLMNOPQRSTUVWXYZ"[length(selectedcells)+2] |> lowercase,12))

	if swabidx == 5 && testidx == 10
		if t0 == t1
			plot!(title="Day $t0: prevalence of $prev0%",titlefontsize=10)
		else
			plot!(title="Days $t0-$t1: prevalence grows exponentially from $prev0% to $prev1%",titlefontsize=10)
		end
	end

	plot!(leftmargin=(swabidx == 1 ? 0mm : -3.25mm),rightmargin=0mm,topmargin=-1mm,bottommargin=(testidx == 1 ? 0mm : -2.25mm),guidefontsize=10)
end |> plts->plot(plts[:,end:-1:begin]...,layout=size(budgetsweep),size=(800,700),legend=nothing,dpi=150)

# ╔═╡ e425af4e-1ad2-11eb-3317-e7dcb97dc924
savefig(resfig,"res,analysis,detail,t0-$(t0),t1-$(t1).png")

# ╔═╡ Cell order:
# ╟─eaffbf72-7bed-11eb-2b72-b759220943ef
# ╠═b53dd37e-0dc5-11eb-37ee-555eed0b2c9f
# ╠═e88985de-2f4e-11eb-32cf-1f7a851792bb
# ╠═a17ac18e-7bee-11eb-1624-fde146fd6293
# ╠═a45a36b4-7bee-11eb-18d7-9f83f59c66fa
# ╠═6e2887ba-0e95-11eb-1475-e7ea7e5cf740
# ╠═08baf128-2958-11eb-294d-35b2a6eeef5d
# ╠═9b6802e0-1558-11eb-1f6e-b5f67e4cafd3
# ╟─f7be67a4-7bed-11eb-032c-e9f443225026
# ╠═213f310c-155e-11eb-0720-b748da5a582d
# ╠═4592edbe-2edd-11eb-2ed0-679711b2d714
# ╟─2e6ae32c-7bee-11eb-12b2-a146668beb80
# ╠═685aa44c-2edd-11eb-237f-6715b380cb39
# ╠═6ec783d4-2edd-11eb-2be6-7f0631417b11
# ╠═b37b43e4-2edd-11eb-23a1-e1c242eb0029
# ╠═bf7652c6-2edd-11eb-05df-33027bb8ab4e
# ╠═0faadc44-0dc6-11eb-0c3f-db78ba861a01
# ╠═d33a506c-0e2b-11eb-24a6-41ddae38afa5
# ╠═d1b9df5a-0e2b-11eb-1767-e7762414041c
# ╠═ccdb2506-0e95-11eb-34b4-8db7c3d43ead
# ╠═c3cc7816-2943-11eb-2807-c742e3307662
# ╟─e3aa4e4a-7be8-11eb-1a29-239b9fbce803
# ╟─3d2aa366-295f-11eb-24be-f93fa0a0a9de
# ╠═eb95975e-2959-11eb-22af-8dc3c3bd4a3c
# ╠═cc143562-2967-11eb-06fe-d7f4cd31664d
# ╟─e2fff794-7be8-11eb-0bfb-a17b3f02641e
# ╠═8b5a9e38-7be9-11eb-1d71-f5529db9d4e8
# ╠═422424ee-3032-11eb-2454-972aa1343745
# ╠═fc03d208-5f2e-11eb-366e-45ca469d44ed
# ╠═17caf5de-2951-11eb-10b4-89609425f75a
# ╠═20d151f0-2951-11eb-124b-b124fe93a1c8
# ╟─4526a4f6-7bef-11eb-17e6-61b5bd0a2f00
# ╠═85b7b0d6-1962-11eb-33c2-b759b672eac8
# ╠═e425af4e-1ad2-11eb-3317-e7dcb97dc924

ext: jl


text:
module StaticCompile

using Mixtape
using MacroTools

module Factorial

f(x::Int64) = x <= 1 ? 1 : x * f(x - 1)

end

@ctx (false, false, true) struct MyMix end
allow(ctx::MyMix, m::Module) = m == Factorial

swap(e) = e
function swap(e::Expr)
    new = MacroTools.postwalk(e) do s
        isexpr(s, :call) || return s
        s.args[1] == Base.:(*) || return s
        return Expr(:call, Base.:(+), e.args[2:end]...)
    end
    return new
end

function transform(::MyMix, b)
    for (v, st) in b
        replace!(b, v, swap(st))
    end
    return b
end

optimize!(::MyMix, ir) = ir

p = "libf.o"
Mixtape.aot(MyMix(), Factorial.f, Tuple{Int}; path = p)
rm(p)

end # module


text:
module NorahPhd
include("./Algorithm1.jl")
# include("../Examples/Example1.jl")
# c1 = [1; 2]
# c = [1; 1]
# z1 = [1; 0]
# z2 = [0; 1]
# C = [c1 c]
# Z = [z1 z2]
# Xbar, Tbar, V, p = primal(Γ, P1, P2, D2, C, Z, ϵ = 0.05)
# println(Vertex2D.(V))

# include("../Examples/Example2.jl")
# c1 = [1; 0]
# c = [1; 1]
# z1 = [1; 0]
# z2 = [0; 1]
# C = [c1 c]
# Z = [z1 z2]
# # v0 = [0; 0]
# # # w, u = D2(v0)
# # println(P1(z1))
# Xbar, Tbar, V, p = primal(Γ, P1, P2, D2, C, Z, ϵ = 0.01)
end


text:
#
# MmapIO
#
# An IO built on top of mmap to avoid the overhead of ordinary disk IO
const MMAP_GROW_SIZE = 2^24
const FILE_GROW_SIZE = 2^18

const Plain = Union{Int16,Int32,Int64,Int128,UInt16,UInt32,UInt64,UInt128,Float16,Float32,
                    Float64}
const PlainType = Union{Type{Int16},Type{Int32},Type{Int64},Type{Int128},Type{UInt16},
                        Type{UInt32},Type{UInt64},Type{UInt128},Type{Float16},
                        Type{Float32},Type{Float64}}

mutable struct MmapIO <: IO
    f::IOStream
    arr::Vector{UInt8}
    curptr::Ptr{Void}
    endptr::Ptr{Void}
end

function MmapIO(fname::AbstractString, write::Bool, create::Bool, truncate::Bool)
    truncate && !write && throw(ArgumentError("cannot truncate file that is not writable"))

    f = open(fname, true, write, create, truncate, false)
    initialsz = truncate ? 0 : filesize(fname)
    arr = Mmap.mmap(f, Vector{UInt8}, (initialsz + MMAP_GROW_SIZE,); grow=false)
    ptr = Ptr{Void}(pointer(arr))
    io = MmapIO(f, arr, ptr, ptr + initialsz)
end

Base.show(io::IO, ::MmapIO) = print(io, "MmapIO")

if is_linux()
    # This is substantially faster than truncate on Linux, but slower on OS X.
    # TODO: Benchmark on Windows
    grow(io::IOStream, sz::Integer) =
        systemerror("pwrite", ccall(:jl_pwrite, Cssize_t,
                                    (Cint, Ptr{UInt8}, UInt, Int64),
                                    fd(io), &UInt8(0), 1, sz - 1) < 1)
else
    grow(io::IOStream, sz::Integer) = truncate(io, sz)
end

function Base.resize!(io::MmapIO, newend::Ptr{Void})
    # Resize file
    ptr = pointer(io.arr)
    newsz = Int(max(newend - ptr, io.curptr - ptr + FILE_GROW_SIZE))
    grow(io.f, newsz)

    if newsz > length(io.arr)
        # If we have not mapped enough memory, map more
        io.arr = Mmap.mmap(io.f, Vector{UInt8}, (newsz + MMAP_GROW_SIZE,); grow=false)
        newptr = pointer(io.arr)
        io.curptr += newptr - ptr
        ptr = newptr
    end

    # Set new end
    io.endptr = ptr + newsz
    io
end

@inline function ensureroom(io::MmapIO, n::Int)
    ep = io.curptr + n
    if ep > io.endptr
        resize!(io, ep)
    end
    nothing
end

Base.truncate(io::MmapIO, pos) = truncate(io.f, pos)

function Base.close(io::MmapIO)
    Mmap.sync!(io.arr)
    close(io.f)
end

@inline function _write(io::MmapIO, x)
    cp = io.curptr
    ep = cp + sizeof(x)
    if ep > io.endptr
        resize!(io, ep)
        cp = io.curptr
        ep = cp + sizeof(x)
    end
    unsafe_store!(Ptr{typeof(x)}(cp), x)
    io.curptr = ep
    return sizeof(x)
end
@inline Base.write(io::MmapIO, x::UInt8) = _write(io, x)
@inline Base.write(io::MmapIO, x::Int8) = _write(io, x)
@inline Base.write(io::MmapIO, x::Plain)  = _write(io, x)

function Base.unsafe_write(io::MmapIO, x::Ptr{UInt8}, n::UInt64)
    cp = io.curptr
    ep = cp + n
    if ep > io.endptr
        resize!(io, ep)
        cp = io.curptr
        ep = cp + n
    end
    unsafe_copy!(Ptr{UInt8}(cp), x, n)
    io.curptr = ep
    return n
end

@inline function _read(io::MmapIO, T::DataType)
    cp = io.curptr
    ep = cp + sizeof(T)
    ep > io.endptr && throw(EOFError())
    v = unsafe_load(Ptr{T}(cp))
    io.curptr = ep
    v
end
@inline Base.read(io::MmapIO, T::Type{UInt8}) = _read(io, T)
@inline Base.read(io::MmapIO, T::Type{Int8}) = _read(io, T)
@inline Base.read(io::MmapIO, T::PlainType) = _read(io, T)

function Base.read{T}(io::MmapIO, ::Type{T}, n::Int)
    cp = io.curptr
    ep = cp + sizeof(T)*n
    ep > io.endptr && throw(EOFError())
    arr = Vector{T}(n)
    unsafe_copy!(pointer(arr), Ptr{T}(cp), n)
    io.curptr = ep
    arr
end
Base.read{T}(io::MmapIO, ::Type{T}, n::Integer) =
    read(io, T, Int(n))

# Read a null-terminated string
function read_bytestring(io::MmapIO)
    # TODO do not try to read outside the buffer
    cp = io.curptr
    str = unsafe_string(convert(Ptr{UInt8}, cp))
    io.curptr = cp + sizeof(str) + 1
    str
end

function Base.seek(io::MmapIO, offset::Integer)
    io.curptr = pointer(io.arr) + offset
    nothing
end

function Base.skip(io::MmapIO, offset::Integer)
    io.curptr += offset
    nothing
end

Base.position(io::MmapIO) = Int64(io.curptr - pointer(io.arr))

# We sometimes need to compute checksums. We do this by first calling
# begin_checksum when starting to handle whatever needs checksumming,
# and calling end_checksum afterwards. Note that we never compute
# nested checksums.
# XXX not thread-safe!

const CHECKSUM_POS = Int64[]
const NCHECKSUM = Ref{Int}(0)
function begin_checksum(io::MmapIO)
    idx = NCHECKSUM[] += 1
    if idx > length(CHECKSUM_POS)
        push!(CHECKSUM_POS, position(io))
    else
        @inbounds CHECKSUM_POS[idx] = position(io)
    end
    io
end
function begin_checksum(io::MmapIO, sz::Integer)
    ensureroom(io, sz)
    begin_checksum(io)
end
function end_checksum(io::MmapIO)
    @inbounds v = CHECKSUM_POS[NCHECKSUM[]]
    NCHECKSUM[] -= 1
    Lookup3.hash(io.arr, v+1, position(io) - v)
end


text:
include("../../../../std/util/src/util_module.jl") # for binary search
# include("../../datastructs/src/datastructs_module.jl") # for stack

module numint
using ..util
# using ..datastructs

struct Interval
    x_lo::Float64
    x_hi::Float64
    function Interval(x_lo, x_hi)
        (x_lo, x_hi) = (x_lo < x_hi) ? (x_lo, x_hi) : (x_hi, x_lo)
        return new(x_lo, x_hi)
    end
end
function range(interval::Interval)
    return interval.x_hi - interval.x_lo
end

function trapezoidal(fcn, interval::Interval, N::Int)

    int_range = range(interval)
    dx = int_range / N
    I  = 0.0
    x1 = interval.x_lo
    f1 = fcn(x1)
    @inbounds for kk = 1:N
        x2 = x1 + dx
        f2 = fcn(x2)
        I += 0.5 * (f1 + f2) * dx
        x1 = x2
        f1 = f2
    end
    return I
end

function trapezoidal(fcn, interval::Interval, xs::Vector)
    I  = 0.0
    x1 = interval.x_lo
    f1 = fcn(x1)
    @inbounds for kk = 2:length(xs)
        x2 = xs[kk]
        dx = x2 - x1
        f2 = fcn(x2)
        I += 0.5 * (f1 + f2) * dx
        x1 = x2
        f1 = f2
    end
    return I
end

struct RombergOptions
    N::Int
    abs_tol::Float64
    rel_tol::Float64
end
function RombergOptions()
    return RombergOptions(128, 1e-8, 1e-8)
end
function RombergOptions(;N = 128, abs_tol = 1e-8, rel_tol = 1e-8)
    return RombergOptions(N, abs_tol, rel_tol)
end

@enum StoppingCrit begin
    max_iter_reached
    rel_tol
    abs_tol
    unknown
end

struct RombergResult
    integral::Float64
    n_fcn_eval::Int
    stopping_crit::StoppingCrit
    I::Vector{Float64}
end

function romberg(
    fcn,
    interval::Interval,
    options::RombergOptions = RombergOptions()
    )

    # allocations
    int_range   = range(interval)
    Nmax        = ceil(Int, log2(options.N)) + 1
    I           = zeros(Nmax)

    # N == 2
    # stat
    n_fcn_eval    = 2
    stopping_crit = unknown::StoppingCrit

    idxi    = 1
    fi_a    = fcn(interval.x_lo)
    fi_b    = fcn(interval.x_hi)
    I[idxi] = 0.5 * (fi_a + fi_b) * int_range

    # R = datastructs.Stack{Float64}()
    R    = zeros(Nmax)
    R[1] = I[1]

    # auxilary arrays
    r   = zeros(Nmax)
    tmp = zeros(Nmax)

    # main loop
    @inbounds for N = 2:Nmax

        # fill up the trapezoidal table first column
        hn = int_range / 2^(N-1)

        # evaluate the function or steal it from the previous iteration
        I_partial = 0.0
        @inbounds for kk = 1:2^(N-2)
            x = interval.x_lo + hn * (2 * kk - 1)
            n_fcn_eval += 1
            I_partial  += fcn(x)
        end
        R[N] = 0.5 * R[N-1] + hn * I_partial

        copyto!(r, R)
        @inbounds for jj = 2:N
            @inbounds for ii = jj:N
                R_n_m1  = r[ii]
                R_n1_m1 = r[ii-1]
                tmp[ii] = R_n_m1 + (R_n_m1 - R_n1_m1) / (4^(jj-1)-1)
            end
            # swap the arrays such that the previous array is always in r
            r, tmp = tmp, r
        end

        idxi    += 1
        I[idxi] = r[N]
        if idxi >= 2
            dI = abs(I[idxi-1] - I[idxi])
            if dI <= options.abs_tol
                stopping_crit = abs_tol::StoppingCrit
                break
            end
            if  dI / abs(I[idxi]) <= options.rel_tol
                stopping_crit = rel_tol::StoppingCrit
                break
            end
        end

    end

    if idxi == Nmax
        stopping_crit = max_iter_reached::StoppingCrit
    end

    stat = RombergResult(I[idxi], n_fcn_eval, stopping_crit, I)
    return stat
end

struct GaussKronrodOptions
    abs_tol::Float64
    rel_tol::Float64
    min_interval::Float64
end
function GaussKronrodOptions(;
    abs_tol = 1e-8,
    rel_tol = 1e-8,
    min_interval = 1e-10
    )
    return GaussKronrodOptions(abs_tol, rel_tol, min_interval)
end

struct GaussKronrodResult
    integral::Float64
    n_fcn_eval::Int
end

function gauss_kronrod(
    fcn,
    interval::Interval,
    options::GaussKronrodOptions = GaussKronrodOptions()
    )

    (gi_idx,
    kronrod_weights,
    kronrod_nodes,
    gauss_weights) = __gk_nodes_weights()

    n_fcn_eval = 0
    x_lo = interval.x_lo
    x_hi = interval.x_hi

    (K15, n_fcn_eval) = __rec_gk(
        fcn,
        x_lo,
        x_hi,
        options,
        gi_idx,
        kronrod_weights,
        kronrod_nodes,
        gauss_weights,
        n_fcn_eval)

    stat = GaussKronrodResult(K15, n_fcn_eval)
    return stat
end

function __rec_gk(
    fcn,
    x_lo,
    x_hi,
    options,
    gi_idx,
    kronrod_weights,
    kronrod_nodes,
    gauss_weights,
    n_fcn_eval)

    # constants
    Nk = 15
    Ng = 7

    int_range = x_hi - x_lo
    mid_point = (x_lo + x_hi) / 2

    if int_range < options.min_interval
        return (0.0, n_fcn_eval)
    end

    fi = zeros(Nk)
    @inbounds for ii = 1:Nk
        xi = int_range/2 * kronrod_nodes[ii] + mid_point
        fi[ii] = fcn(xi)
    end
    n_fcn_eval += 15

    gi = fi[gi_idx]

    K15 = 0.0
    @inbounds for ii = 1:Nk
        K15 += kronrod_weights[ii] * fi[ii]
    end
    K15 *= int_range/2

    G7 = 0.0
    @inbounds for ii = 1:Ng
        G7 += gauss_weights[ii] * gi[ii]
    end
    G7 *= int_range/2

    error_est = abs(G7 - K15)
    if error_est <= options.abs_tol
        return (K15, n_fcn_eval)
    end
    if error_est / abs(K15) <= options.rel_tol
        return (K15, n_fcn_eval)
    end

    # subdivide the interval
    x_mid = mid_point

    (K15_left, n_fcn_eval) = __rec_gk(
        fcn,
        x_lo,
        x_mid,
        options,
        gi_idx,
        kronrod_weights,
        kronrod_nodes,
        gauss_weights,
        n_fcn_eval)

    (K15_right, n_fcn_eval) = __rec_gk(
        fcn,
        x_mid,
        x_hi,
        options,
        gi_idx,
        kronrod_weights,
        kronrod_nodes,
        gauss_weights,
        n_fcn_eval)

    return (K15_left + K15_right, n_fcn_eval)
end

function __gk_nodes_weights()
# gauss_nodes = [
    # -0.949107912342759  # 1
    # -0.741531185599394  # 2
    # -0.405845151377397  # 3
    # 0.0                 # 4
    # +0.405845151377397
    # +0.741531185599394
    # +0.949107912342759
    # ]

    gauss_weights = [
        0.129484966168870
        0.279705391489277
        0.381830050505119

        0.417959183673469

        0.381830050505119
        0.279705391489277
        0.129484966168870
    ]

    kronrod_nodes = [
        -0.991455371120813
        -0.949107912342759  # 1
        -0.864864423359769
        -0.741531185599394  # 2
        -0.586087235467691
        -0.405845151377397  # 3
        -0.207784955007898

        0.0                 # 4

        0.207784955007898
        0.405845151377397
        0.586087235467691
        0.741531185599394
        0.864864423359769
        0.949107912342759
        0.991455371120813
    ]

    kronrod_weights = [
        0.022935322010529
        0.063092092629979
        0.104790010322250
        0.140653259715525
        0.169004726639267
        0.190350578064785
        0.204432940075298

        0.209482141084728

        0.204432940075298
        0.190350578064785
        0.169004726639267
        0.140653259715525
        0.104790010322250
        0.063092092629979
        0.022935322010529
    ]

    gi_idx = [  2, 4, 6,
                8,
                10, 12, 14
            ]

    return (gi_idx, kronrod_weights, kronrod_nodes, gauss_weights)
end

end # module

max_stars_repo_name: HomoModelicus/julia

max_issues_repo_name: HomoModelicus/julia

max_forks_repo_name: HomoModelicus/julia


text:
function walk_entry_file(d::DataEntry, file_name)
    res = match_entry_file(d, file_name)
    isempty(res) ? nothing : to_dataframe(res)
end

# function walk_entry(d::DataEntry, file_name)
#     res = match_entry_file(d, file_name)
#     for (val_dict, file_path) in res
#     end
# end

function match_entry_file(d::DataEntry, file_name)
    res = []
    for (root, dirs, files) in walkdir(d.root)
        if !isempty(files)
            file_entries = splitpath(relpath(root, d.root))
            length(file_entries) != length(d.folder_entries) && break
            miss_match_flag = false
            split_entries = []
            for (file_entry, folder_entry) in zip(file_entries, d.folder_entries)
                split_entry = [split(x, "=") for x in split(file_entry, "_")]
                if !any((i) -> split_entry[i][1] in folder_entry, 1:length(split_entry))
                    miss_match_flag = true
                    break
                end
                split_entries = vcat(split_entries, [x[1]=>x[2] for x in split_entry])
            end
            val_dict = Dict{String, String}(split_entries)
            for f in files
                if f == file_name
                    push!(res, (val_dict, joinpath(root, f)))
                end
            end
        end
    end
    res
end

function to_dataframe(file_entries)
    entry_names = [x for x in keys(file_entries[1][1])]
    l = length(entry_names)
    res = [Vector{String}() for i in 1:l+1]
    for entry in file_entries
        for i in 1:l
            push!(res[i], entry[1][entry_names[i]])
        end
        push!(res[end], entry[2])
    end
    entry_symbols = [Symbol(x) for x in entry_names]
    entry_symbols = [entry_symbols; :file]
    DataFrame(res, entry_symbols)
end


text:
using Test
using NLsolve
using Distributions
# using Statistics

using IncrementalInference

import IncrementalInference: getSample

##

@testset "Ensure lambda's work with NLsolve" begin

function f!(F, x, dummy)
    F[1] = (x[1]+3)*(x[2]^3-7)+18
    F[2] = sin(x[2]*exp(x[1])-1)
end

data = randn(2,3)
r = nlsolve( (res, x) -> f!(res, x, data), zeros(2), inplace=true )

@test norm(r.zero - [0, 1.0]) < 1e-10

end

#


text:
"""
    unwrap_nested_exception(err::Exception) -> Bool

Extracts the "core" exception from a nested exception."
"""
unwrap_nested_exception(err::CapturedException) =
    unwrap_nested_exception(err.ex)
unwrap_nested_exception(err::RemoteException) =
    unwrap_nested_exception(err.captured)
unwrap_nested_exception(err) = err

"Prepares the scheduler to schedule `thunk`."
function reschedule_inputs!(state, thunk, seen=Dict{Thunk,Bool}())
    haskey(seen, thunk) && return seen[thunk]
    haskey(state.cache, thunk) && return false
    w = get!(()->Set{Thunk}(), state.waiting, thunk)
    scheduled = false
    for input in unwrap_weak.(thunk.inputs)
        if istask(input) || (input isa Chunk)
            push!(get!(()->Set{Thunk}(), state.waiting_data, input), thunk)
        end
        istask(input) || continue
        if get(state.errored, input, false)
            set_failed!(state, input, thunk)
            scheduled = true
        end
        haskey(state.cache, input) && continue
        if (input in state.running) ||
           (input in state.ready) ||
           reschedule_inputs!(state, input, seen)
            push!(w, input)
            scheduled = true
        else
            error("Failed to reschedule $(input.id) for $(thunk.id)")
        end
    end
    if isempty(w)
        # Inputs are ready
        delete!(state.waiting, thunk)
        if !get(state.errored, thunk, false)
            push!(state.ready, thunk)
        end
        return true
    else
        return scheduled
    end
end

"Marks `thunk` and all dependent thunks as failed."
function set_failed!(state, origin, thunk=origin)
    filter!(x->x!==thunk, state.ready)
    state.cache[thunk] = ThunkFailedException(thunk, origin, state.cache[origin])
    state.errored[thunk] = true
    if haskey(state.futures, thunk)
        for future in state.futures[thunk]
            put!(future, state.cache[thunk]; error=true)
        end
        delete!(state.futures, thunk)
    end
    if haskey(state.waiting_data, thunk)
        for dep in state.waiting_data[thunk]
            haskey(state.waiting, dep) &&
                delete!(state.waiting, dep)
            haskey(state.errored, dep) &&
                continue
            set_failed!(state, origin, dep)
        end
        delete!(state.waiting_data, thunk)
    end
    if haskey(state.waiting, thunk)
        delete!(state.waiting, thunk)
    end
end

"Internal utility, useful for debugging scheduler state."
function print_sch_status(state, thunk; kwargs...)
    iob = IOBuffer()
    print_sch_status(iob, state, thunk; kwargs...)
    seek(iob, 0)
    write(stderr, iob)
end
function print_sch_status(io::IO, state, thunk; offset=0, limit=5, max_inputs=3)
    function status_string(node)
        status = ""
        if get(state.errored, node, false)
            status *= "E"
        end
        if node in state.ready
            status *= "r"
        elseif node in state.running
            status *= "R"
        elseif haskey(state.cache, node)
            status *= "C"
        else
            status *= "?"
        end
        status
    end
    if offset == 0
        println(io, "Ready ($(length(state.ready))): $(join(map(t->t.id, state.ready), ','))")
        println(io, "Running: ($(length(state.running))): $(join(map(t->t.id, collect(state.running)), ','))")
        print(io, "($(status_string(thunk))) ")
    end
    println(io, "$(thunk.id): $(thunk.f)")
    for (idx,input) in enumerate(thunk.inputs)
        input isa Thunk || continue
        if idx > max_inputs
            println(io, repeat(' ', offset+2), "…")
            break
        end
        status = status_string(input)
        if haskey(state.waiting, thunk) && input in state.waiting[thunk]
            status *= "W"
        end
        if haskey(state.waiting_data, input) && thunk in state.waiting_data[input]
            status *= "w"
        end
        if haskey(state.futures, input)
            status *= "f($(length(state.futures[input])))"
        end
        print(io, repeat(' ', offset+2), "($status) ")
        if limit > 0
            print_sch_status(io, state, input; offset=offset+2, limit=limit-1)
        else
            println(io, "…")
        end
    end
end

function fetch_report(task)
    try
        fetch(task)
    catch err
        @static if VERSION >= v"1.1"
            stk = Base.catch_stack(task)
            err, frames = stk[1]
            rethrow(CapturedException(err, frames))
        else
            rethrow(task.result)
        end
    end
end

function signature(task::Thunk, state)
    inputs = map(x->istask(x) ? state.cache[x] : x, unwrap_weak.(task.inputs))
    Tuple{typeof(task.f), map(x->x isa Chunk ? x.chunktype : typeof(x), inputs)...}
end

function report_catch_error(err, desc=nothing)
    iob = IOContext(IOBuffer(), :color=>true)
    if desc !== nothing
        println(iob, desc)
    end
    Base.showerror(iob, err)
    Base.show_backtrace(iob, catch_backtrace())
    println(iob)
    seek(iob.io, 0)
    write(stderr, iob)
end


text:
using Survival
using Distributions
using Base.Test
using DataFrames
using JLD

# Check aft
include("aft_test.jl")

#test cox
include("cox_rossi.jl")


text:
module FluxUtils

using Flux
using Random

function kaiming_uniform(dims...; gain=sqrt(2))
   fan_in = length(dims) <= 2 ? dims[end] : div(*(dims...), dims[end])
   bound = sqrt(3.0) * gain / sqrt(fan_in)
   return Float32.(rand(Uniform(-bound, bound), dims...))
 end

 function kaiming_normal(dims...; gain=sqrt(2))
   fan_in = length(dims) <= 2 ? dims[end] : div(*(dims...), dims[end])
   std = gain / sqrt(fan_in)
   return Float32.(rand(Normal(0.0, std), dims...))
 end

# glorot_uniform(rng::Random.AbstractRNG, T::Type, dims...) = (rand(rng, T, dims...) .- 0.5f0) .* sqrt(24.0f0/sum(dims))
# glorot_normal(rng::Random.AbstractRNG, T::Type, dims...) = randn(rng, T, dims...) .* sqrt(2.0f0/sum(dims))

# glorot_uniform(rng::Random.AbstractRNG, dims::Vararg{Int64}) = (rand(rng, Float32, dims...) .- 0.5f0) .* sqrt(24.0f0/sum(dims))
# glorot_normal(rng::Random.AbstractRNG, dims...) = randn(rng, Float32, dims...) .* sqrt(2.0f0/sum(dims))

function get_activations(model::M, data) where {M<:Flux.Chain}
    # Assume it is a chain of things
    activations = [data]
    for (idx, layer) in enumerate(model)
        push!(activations, layer(activations[idx]))
    end
    return activations
end

struct ParallelStreams{T<:Tuple}
    l::T
end

ParallelStreams(args...) = ParallelStreams((args))

Flux.@functor ParallelStreams
(l::ParallelStreams)(x) = map((mdl)->mdl(x), l.l)

function Base.show(io::IO, l::ParallelStreams)
  print(io, "ParallelStreams(", (string(layer)*", " for layer in l.l)..., ")")
end

struct DualStreams{M1, M2}
    m1::M1
    m2::M2
end

Flux.@functor DualStreams
(l::DualStreams)(x) = (l.m1(x), l.m2(x))

struct ConcatStreams{M1, M2}
    m1::M1
    m2::M2
end

Flux.@functor ConcatStreams
(l::ConcatStreams)(x) = vcat(l.m1(x), l.m2(x))

function Base.show(io::IO, l::ConcatStreams)
  print(io, "ConcatStreams(", string(l.m1), ", ", string(l.m2), ")")
end

end


text:
using KernelFunctions
using AxisArrays
using Distances
using Documenter
using Kronecker: Kronecker
using LinearAlgebra
using LogExpFunctions
using PDMats
using Random
using SpecialFunctions
using Test
using Flux
using Zygote: Zygote
using ForwardDiff: ForwardDiff
using ReverseDiff: ReverseDiff
using FiniteDifferences: FiniteDifferences

using KernelFunctions: SimpleKernel, metric, kappa, ColVecs, RowVecs, TestUtils

using KernelFunctions.TestUtils: test_interface

# Writing tests:
# 1. The file structure of the test should match precisely the file structure of src.
#   Amongst other things, this means that there should be exactly 1 test file per src file.
#   This makes it trivially easy for someone to find the tests associated with a particular
#   src file.
# 2. A consequence of 1 is that there should be exactly 1 test file per src file.
# 3. A test file called foo.jl should have the structure:
#   @testset "foo" begin
#       code
#   end
#
#   Note that the testset is called `foo`, not `foo.jl`. Use whatever testset structure
#   seems appropriate within a given file. eg. if multiple types / functions are defined in
#   a particular source file, you might want multiple testsets in the test file.
# 4. Each directory should have its own testset, in which each test file is `include`d.
# 5. Each test file should create its own state, and shouldn't rely on state defined in
#   other test files. e.g. don't define a matrix used by all of the files in kernels. If
#   two test files are similar enough to share state, perhaps the corresponding source code
#   should be in the same file.
# 6. If you write a src file without any tests, create a corresponding test file with the
#   usual structure, but without any tests.
# 7. Explicitly create a new random number generate for _at_ _least_ each new test file, and
#   use it whenever generating randomness. This ensures complete control over random number
#   generation and makes it clear what randomness depends on other randomness.
# 8. All `using` statements should appear in runtests.jl.
# 9. List out all test files explicitly (eg. don't loop over them). This makes it easy to
#   disable tests by simply commenting them out, and makes it very clear which tests are not
#   currently being run.
# 10. If utility functionality is required, it should be placed in `src/test_utils.jl` so
#   that other packages can benefit from it when implementing new kernels.
@info "Packages Loaded"

include("test_utils.jl")

@testset "KernelFunctions" begin
    include("utils.jl")

    @testset "distances" begin
        include(joinpath("distances", "pairwise.jl"))
        include(joinpath("distances", "dotproduct.jl"))
        include(joinpath("distances", "delta.jl"))
        include(joinpath("distances", "sinus.jl"))
    end
    @info "Ran tests on Distances"

    @testset "transform" begin
        include(joinpath("transform", "transform.jl"))
        print(" ")
        include(joinpath("transform", "scaletransform.jl"))
        print(" ")
        include(joinpath("transform", "ardtransform.jl"))
        print(" ")
        include(joinpath("transform", "lineartransform.jl"))
        print(" ")
        include(joinpath("transform", "functiontransform.jl"))
        print(" ")
        include(joinpath("transform", "selecttransform.jl"))
        print(" ")
        include(joinpath("transform", "chaintransform.jl"))
        print(" ")
        include(joinpath("transform", "periodic_transform.jl"))
        print(" ")
    end
    @info "Ran tests on Transform"

    @testset "basekernels" begin
        include(joinpath("basekernels", "constant.jl"))
        print(" ")
        include(joinpath("basekernels", "cosine.jl"))
        print(" ")
        include(joinpath("basekernels", "exponential.jl"))
        print(" ")
        include(joinpath("basekernels", "exponentiated.jl"))
        print(" ")
        include(joinpath("basekernels", "fbm.jl"))
        print(" ")
        include(joinpath("basekernels", "gabor.jl"))
        print(" ")
        include(joinpath("basekernels", "matern.jl"))
        print(" ")
        include(joinpath("basekernels", "nn.jl"))
        print(" ")
        include(joinpath("basekernels", "periodic.jl"))
        print(" ")
        include(joinpath("basekernels", "piecewisepolynomial.jl"))
        print(" ")
        include(joinpath("basekernels", "polynomial.jl"))
        print(" ")
        include(joinpath("basekernels", "rational.jl"))
        print(" ")
        include(joinpath("basekernels", "sm.jl"))
        print(" ")
        include(joinpath("basekernels", "wiener.jl"))
        print(" ")
    end
    @info "Ran tests on BaseKernel"

    @testset "kernels" begin
        include(joinpath("kernels", "kernelproduct.jl"))
        include(joinpath("kernels", "kernelsum.jl"))
        include(joinpath("kernels", "kerneltensorproduct.jl"))
        include(joinpath("kernels", "overloads.jl"))
        include(joinpath("kernels", "scaledkernel.jl"))
        include(joinpath("kernels", "transformedkernel.jl"))
        include(joinpath("kernels", "normalizedkernel.jl"))
        include(joinpath("kernels", "neuralkernelnetwork.jl"))
    end
    @info "Ran tests on Kernel"

    @testset "matrix" begin
        include(joinpath("matrix", "kernelmatrix.jl"))
        include(joinpath("matrix", "kernelkroneckermat.jl"))
        include(joinpath("matrix", "kernelpdmat.jl"))
    end
    @info "Ran tests on matrix"

    @testset "multi_output" begin
        include(joinpath("mokernels", "moinput.jl"))
        include(joinpath("mokernels", "independent.jl"))
        include(joinpath("mokernels", "slfm.jl"))
        include(joinpath("mokernels", "intrinsiccoregion.jl"))
        include(joinpath("mokernels", "lmm.jl"))
    end
    @info "Ran tests on Multi-Output Kernels"

    @testset "approximations" begin
        include(joinpath("approximations", "nystrom.jl"))
    end

    include("generic.jl")
    include("chainrules.jl")
    include("zygoterules.jl")

    @testset "doctests" begin
        DocMeta.setdocmeta!(
            KernelFunctions,
            :DocTestSetup,
            quote
                using KernelFunctions
                using LinearAlgebra
                using Random
                using PDMats: PDMats
            end;
            recursive=true,
        )
        doctest(
            KernelFunctions;
            doctestfilters=[
                r"{([a-zA-Z0-9]+,\s?)+[a-zA-Z0-9]+}",
                r"(Array{[a-zA-Z0-9]+,\s?1}|Vector{[a-zA-Z0-9]+})",
                r"(Array{[a-zA-Z0-9]+,\s?2}|Matrix{[a-zA-Z0-9]+})",
            ],
        )
    end
end


text:
#derivative logic

#ForwardDiff compiles one method per function, so separating the
#differentiation logic from the property logic allows the differentials
#to be compiled only once

"""
    ∂f∂t(model,v,t,z=SA[1.0])

returns `f` and `∂f/∂T` at constant total volume and composition, where f is the total helmholtz energy, given by `eos(model,v,T,z)`

"""
function ∂f∂t(model,v,t,z=SA[1.0])
    return ForwardDiff.derivative(∂t -> eos(model,v,∂t,z),t)
end

"""
    ∂f∂v(model,v,t,z=SA[1.0])

returns `f` and `∂f/∂v` at constant temperature and composition, where f is the total helmholtz energy, given by `eos(model,v,T,z)`, and v is the total volume

"""
function ∂f∂v(model,v,t,z)
    return ForwardDiff.derivative(∂v -> eos(model,∂v,t,z),v)
end

#returns a tuple of the form ([∂f∂v,∂f∂t],f),using the least amount of computation
"""
    ∂f(model,v,t,z)

returns zeroth order (value) and first order derivative information of the total helmholtz energy (given by `eos(model,v,T,z)`).
the result is given in two values:

```
grad_f,fval = ∂2f(model,v,t,z)
```

where:
```
fval   = f(v,T) = eos(model,v,T,z)

grad_f = [ ∂f/∂v; ∂f/∂T]

 ```

Where `v` is the total volume, `T` is the temperature and `f` is the total helmholtz energy.
"""
function ∂f(model,v,t,z)
    f(w) = eos(model,first(w),last(w),z)
    v,t = promote(v,t)
    vt_vec = SVector(v,t)
    ∂result = DiffResults.GradientResult(vt_vec)
    res_∂f =  ForwardDiff.gradient!(∂result, f,vt_vec)
    _f =  DiffResults.value(res_∂f)
    _∂f = DiffResults.gradient(res_∂f)
    return (_∂f,_f)
end

#returns p and ∂p∂v at constant t
#it doesnt do a pass over temperature, so its
#faster that d2f when only requiring d2fdv2

"""
    p∂p∂v(model,v,t,z=SA[1.0])

returns `p` and `∂p/∂v` at constant temperature, where p is the pressure = `pressure(model,v,T,z)` and `v` is the total volume.

"""
function p∂p∂v(model,v,t,z=SA[1.0])
    v_vec =   SVector(v)
    f(∂v) = pressure(model,only(∂v),t,z)
    ∂result = DiffResults.GradientResult(v_vec)
    res_∂f =  ForwardDiff.gradient!(∂result, f,v_vec)
    _p =  DiffResults.value(res_∂f)
    _∂p∂v = only(DiffResults.gradient(res_∂f))
    return _p,_∂p∂v
end

"""
    ∂2f(model,v,t,z)

returns zeroth order (value), first order and second order derivative information of the total helmholtz energy (given by `eos(model,v,T,z)`).
the result is given in three values:

```
hess_f,grad_f,fval = ∂2f(model,v,t,z)
```

where:
```
fval   = f(v,T) = eos(model,v,T,z)

grad_f = [ ∂f/∂v; ∂f/∂T]

hess_f = [ ∂²f/∂v²; ∂²f/∂v∂T
          ∂²f/∂v∂T; ∂²f/∂v²]
 ```

Where `v` is the total volume, `T` is the temperature and `f` is the total helmholtz energy.
"""
function ∂2f(model,v,t,z)
    f(w) = eos(model,first(w),last(w),z)
    v,t = promote(v,t)
    vt_vec =   SVector(v,t)
    ∂result = DiffResults.HessianResult(vt_vec)
    res_∂f =  ForwardDiff.hessian!(∂result, f,vt_vec)
    _f =  DiffResults.value(res_∂f)
    _∂f = DiffResults.gradient(res_∂f)
    _∂2f = DiffResults.hessian(res_∂f)
    return (_∂2f,_∂f,_f)
end

"""
    ∂2p(model,v,t,z)

returns zeroth order (value), first order and second order derivative information of the pressure.
the result is given in three values:

```
hess_p,grad_p,pval = ∂2p(model,v,t,z)
```

where:
```
pval   = p(v,T) = pressure(model,v,T,z)

grad_p = [ ∂p/∂v; ∂p/∂T]

hess_p = [ ∂²p/∂v²; ∂²p/∂v∂T
          ∂²p/∂v∂T; ∂²p/∂v²]
 ```

Where `v` is the total volume, `T` is the temperature and `p` is the pressure.
"""
function ∂2p(model,v,t,z)
    f(w) = pressure(model,first(w),last(w),z)
    v,t = promote(v,t)
    vt_vec =   SVector(v,t)
    ∂result = DiffResults.HessianResult(vt_vec)
    res_∂f =  ForwardDiff.hessian!(∂result, f,vt_vec)
    _f =  DiffResults.value(res_∂f)
    _∂f = DiffResults.gradient(res_∂f)
    _∂2f = DiffResults.hessian(res_∂f)
    return (_∂2f,_∂f,_f)
end

"""
    f_hess(model,v,t,z)

returns the second order volume (`v`) and temperature (`T`) derivatives of the total helmholtz energy (given by `eos(model,v,T,z)`). the result is given in a 2x2 `SMatrix`, in the form:

```
[ ∂²f/∂v²  ∂²f/∂v∂T
 ∂²f/∂v∂T  ∂²f/∂v²]
 ```

use this instead of the ∂2f if you only need second order information. ∂2f also gives zeroth and first order derivative information, but due to a bug in the used AD, it allocates more than necessary.
"""
function f_hess(model,v,t,z)
    f(w) = eos(model,first(w),last(w),z)
    v,t = promote(v,t)
    vt_vec = SVector(v,t)
    return ForwardDiff.hessian(f,vt_vec)
end

"""
    ∂p2∂p3(model,v,t,z=SA[1.0])

returns `∂²p/∂v²` and `∂³p/∂v³`, in a single ForwardDiff pass. used mainly in `crit_pure` objective function

"""
function ∂p2∂p3(model,v,t,z=SA[1.0])
    v_vec =   SVector(v)
    f(∂v) = ForwardDiff.derivative(∂2v -> pressure(model,only(∂2v),t,z),only(∂v))
    ∂result = DiffResults.GradientResult(v_vec)
    res_∂f =  ForwardDiff.gradient!(∂result, f,v_vec)
    _p =  DiffResults.value(res_∂f)
    _∂p∂v = only(DiffResults.gradient(res_∂f))
    return _p,_∂p∂v
end


text:
using Statistics

mae(output, label; agg=mean) = agg(abs.(output .- label))

mse(output, label; agg=mean) = agg((output .- label) .^ 2)

function cross_entropy(probs, label; dims=1, agg=mean)
    ce_summands = label .* log.(probs .+ eps(eltype(probs)))
    agg(.-sum(ce_summands, dims=dims))
end

function class_error(probs, label; dims=1, agg=mean)
    class_predicted = argmax(probs, dims=dims)
    class_actual = argmax(label, dims=dims)
    return agg(1 .- (class_predicted .== class_actual))
end


text:
function moose_body_mass(latitude)
    mass = 0

    # Your code goes here!

    return mass
end

max_stars_repo_name: basimr/lovelace

max_issues_repo_name: basimr/project-lovelace

max_forks_repo_name: basimr/project-lovelace


text:
#== # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#   This function reads the embedded gravity models.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
#
#   [1] http://icgem.gfz-potsdam.de/home
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # ==#

export load_gravity_model
export EGM96, JGM2, JGM3

################################################################################
#                               Public Functions
################################################################################

"""
    function load_gravity_model(T)

Load an embedded gravity model coefficients `T` and return an instance of the
structure `GravityModel_Coefs` with the parsed values.

The current supported values for `T` are:

| `T`       | Model Name                     | Maximum Degree |
|:---------:|:-------------------------------|:---------------|
| `EGM96()` | Earth Gravitational Model 1996 | 360            |
| `JGM2()`  | Joint Earth Gravity Model 2    | 70             |
| `JGM3()`  | Joint Earth Gravity Model 3    | 70             |
|-----------|--------------------------------|----------------|

For other models, you can downlad the `gfc` file at

    http://icgem.gfz-potsdam.de/home

and load it using the functions `parse_icgem` and `create_gravity_model_coefs`.

"""
function load_gravity_model(::Type{Val{:egm96}})
    dir      = @__DIR__
    filename = dir * "/coefficients/EGM96.gfc"

    return create_gravity_model_coefs(parse_icgem(filename))
end

function load_gravity_model(::Type{Val{:jgm2}})
    dir      = @__DIR__
    filename = dir * "/coefficients/JGM2.gfc"

    return create_gravity_model_coefs(parse_icgem(filename))
end

function load_gravity_model(::Type{Val{:jgm3}})
    dir      = @__DIR__
    filename = dir * "/coefficients/JGM3.gfc"

    return create_gravity_model_coefs(parse_icgem(filename))
end

################################################################################
#                             Auxiliary Functions
################################################################################

@inline EGM96() = Val{:egm96}
@inline JGM2()  = Val{:jgm2}
@inline JGM3()  = Val{:jgm3}


text:
"""
    function create_main_index(metadata_folder="../metadata/"::String,
        main_file="../metadata/main_idx.tsv"::String
    )

Create main index TSV file by combining all metadata files

Parameters
* metadata_folder: Folder where metadata TSVs are stored
* main_file: TSV file name for combined metadata

"""
function create_main_index(metadata_folder="../metadata/"::String,
                            main_file="../metadata/main_idx.tsv"::String)

    # Import all csv files into a dataframe
    metadata_files = [i for i in readdir(metadata_folder; join=true) if (i!=main_file) & occursin(".tsv", i)]
    df = reduce(vcat, [DataFrame(CSV.File(i, delim="|")) for i in metadata_files])

    # remove duplicates
    df = df[.!nonunique(df), :]

    # export df
    CSV.write(main_file, df, delim="|")

    return

end


text:
function CRT_GIA_kin_fb(m, Nx, ϕ, u, D, k_f, k_b, nu_A, nu_B, nu_C, a, b, c, C_bc, C_0, N_comp, dt, T, tol)
  BC = Array{Any}(N_comp)
  for i in 1:N_comp
      BC[i] = createBC(m)
      # Dirichlet BC on the left side
      BC[i].left.a[:] = 0.0
      BC[i].left.b[:] = 1.0
      BC[i].left.c[:] = C_bc[i]
  end

  # initial condition
  C_old = Array{Any}(N_comp)
  for i in 1:N_comp
      C_old[i] = createCellVariable(m, C_0[i], BC[i]) # initial value of concentration
  end

  # Advective term
  u_vec = createFaceVariable(m, u)
  M_adv = convectionUpwindTerm(u_vec) # matrix of coef for the advection term

  # Diffusive term
  M_diff = Array{Any}(N_comp)
  for i in 1:N_comp
    D_face = createFaceVariable(m, D[i]*ϕ)
    M_diff[i] = diffusionTerm(D_face)      # matrix of coef for the diffusion term
  end

  # Boundary terms
  M_bc = Array{Any}(N_comp)
  RHS_bc  = Array{Any}(N_comp)
  for i in 1:N_comp
      M_bc[i], RHS_bc[i] = boundaryConditionTerm(BC[i]) # matrix and RHS for the BC
  end

  # Constructing stationary part of each system matrix
  M_stat = Array{Any}(N_comp)
  for i in 1:N_comp
    M_stat[i] = M_adv - M_diff[i] + M_bc[i]
  end

  # Preparing transient terms
  M_t = Array{Any}(N_comp)
  RHS_t = Array{Any}(N_comp)

  # Preparing Jacobian terms
  J_eps = Array{Any}(N_comp)
  RHS_eps = Array{Any}(N_comp)

  # Preparing system matrices and RHS
  M_system = Array{Any}(N_comp)
  RHS_system = Array{Any}(N_comp)
  C_new = Array{Any}(N_comp)
  C_temp = Array{Any}(N_comp)

  # Defining time-stepping parameters
  maxit = 5000 # maximum number of iterations

  # Time-stepping
  for t in dt:dt:T
    iter = 0
    error = 1000 # forcing at least one Newton-iteration

    #Assigning C_old as the new start guess for the Newton-Iteration and creating the transport related system matrix
    for i in 1:N_comp
      C_temp[i] = copyCell(C_old[i])
      M_t[i], RHS_t[i] = transientTerm(C_old[i], dt, ϕ)
      M_system[i] = M_t[i] + M_stat[i]
    end

    #Newton iterations loop
    while (error > tol && iter < maxit)

      # Constructing Jacobian terms
      J_eps[1] = linearSourceTerm(createCellVariable(m, a*k_f.*(internalCells(C_temp[1]).^(a-1)).*(internalCells(C_temp[2]).^b), BC[1]))
      J_eps[2] = linearSourceTerm(createCellVariable(m, b*k_f.*(internalCells(C_temp[1]).^a).*(internalCells(C_temp[2]).^(b-1)), BC[2]))
      J_eps[3] = linearSourceTerm(createCellVariable(m, -c*k_b.*(internalCells(C_temp[3]).^(c-1)), BC[3]))

      RHS_reacA = constantSourceTerm(createCellVariable(m, -k_f.*(internalCells(C_temp[1]).^a).*(internalCells(C_temp[2]).^b).*(1-a-b)+k_b*(internalCells(C_temp[3]).^c)*(1-c), BC[1]))
      RHS_reacB = constantSourceTerm(createCellVariable(m, -k_f.*(internalCells(C_temp[1]).^a).*(internalCells(C_temp[2]).^b).*(1-a-b)+k_b*(internalCells(C_temp[3]).^c)*(1-c), BC[2]))
      RHS_reacC = constantSourceTerm(createCellVariable(m, -k_f.*(internalCells(C_temp[1]).^a).*(internalCells(C_temp[2]).^b).*(1-a-b)+k_b*(internalCells(C_temp[3]).^c)*(1-c), BC[3]))

      # Constructing the RHS's
      RHS_system[1] = RHS_t[1] + RHS_bc[1] - nu_A*RHS_reacA
      RHS_system[2] = RHS_t[2] + RHS_bc[2] - nu_B*RHS_reacB
      RHS_system[3] = RHS_t[3] + RHS_bc[3] + nu_C*RHS_reacC

      RHS = [RHS_system[1];
            RHS_system[2];
            RHS_system[3]]

      # Constructing block system matrix
      M = [M_system[1]-(nu_A*J_eps[1]) -(nu_A*J_eps[2]) -(nu_A*J_eps[3]);
          -(nu_B*J_eps[1]) M_system[2]-(nu_B*J_eps[2]) -(nu_B*J_eps[3]);
          (nu_C*J_eps[1]) (nu_C*J_eps[2]) M_system[3]+(nu_C*J_eps[3])]

      # Solving the system
      sol = M\RHS

      # Putting solution into JFVM structure
      for i in 1:N_comp
        C_new[i] = copyCell(C_temp[i])
        C_new[i].value[:] = sol[((i-1)*(Nx+2)+1):(i*(Nx+2))]
      end

      # Calculating convergence criteria
      error = 0
      for i in 1:N_comp
        error = error + sum(abs(C_new[i].value-C_temp[i].value))
        C_temp[i] = copyCell(C_new[i])
      end

      iter = iter + 1
    end

    #if (t == dt) break end

    for i in 1:N_comp
      C_old[i] = copyCell(C_new[i])
    end

    #print(iter, ", ", t, ", ", error, "\n")
    #print(sum(internalCells(C_old[3])-K_eq*internalCells(C_old[1]).*internalCells(C_old[2])), "\n") #mass balance
  end

  return C_old
end

max_stars_repo_name: FrederikLehn/coupled-solvers

max_issues_repo_name: FrederikLehn/coupled-solvers

max_forks_repo_name: FrederikLehn/coupled-solvers


text:
# debug functionality

"""
    enable_timings()

Enable the recording of debug timings.
"""
enable_timings() = (TimerOutputs.enable_debug_timings(CUDA); return)
disable_timings() = (TimerOutputs.disable_debug_timings(CUDA); return)

isdebug(group, mod=CUDA) =
    Base.CoreLogging.current_logger_for_env(Base.CoreLogging.Debug, group, mod) !== nothing


text:
module TestScripts

using Test

@testset "scripts" begin
    empty!(ARGS)
    push!(ARGS, "arg", "--opt1=2", "--opt2", "3", "-f")
    @test Base.include(Main, "scripts/hello.jl") == 0

    empty!(ARGS)
    push!(ARGS, "activate", "-h")
    @test Base.include(Main, "scripts/pkg.jl") == 0

    empty!(ARGS)
    push!(ARGS, "activate", "path", "--shared")
    @test Base.include(Main, "scripts/pkg.jl") == 0

    empty!(ARGS)
    push!(ARGS, "Author - Year.pdf")
    @test Base.include(Main, "scripts/searchpdf.jl") == 0
end

end


text:
using SoftSquishyMatter # load simulation package
using Random # useful if you want to shuffle arrays
cd(dirname(@__FILE__)) # set directory to this file's folder

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# Initialize empty simulation with a description for future reference
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
simulation = Simulation()
simulation.descriptor = "Run-and-tumble particles with funnels"

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# Initialize active particles and boundary particles (walls and funnels)
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
R_swimmer = 1e-6 # radius of active particles
γ_trans_swimmer = 6 * pi * 8.9e-4 * R_swimmer # translational friction coefficient
D_trans_swimmer = 1.38e-23 * 300 / γ_trans_swimmer # translational diffusivity

R_boundary = 1e-6 # radius of boundary particles

# initialize positions of active particles on a triangular lattice
# use lattice dimensions for the simulation dimensions (L_x, L_y)
swimmer_positions, simulation.L_x, simulation.L_y = triangular_lattice(s = 4e-6, M_x = 50, M_y = 15)

# initialize boundary particles
boundary_positions = Array{Array{Float64, 1}, 1}()
# add two walls of particles on the leftmost and rightmost edges of simulation
append!(boundary_positions, flat_edge(from = [R_boundary, R_boundary], to = [R_boundary, simulation.L_y - R_boundary], spacing = R_boundary))
append!(boundary_positions, flat_edge(from = [simulation.L_x - R_boundary, R_boundary], to = [simulation.L_x - R_boundary, simulation.L_y - R_boundary], spacing = R_boundary))
# add lanes of wedges to form funnels
angle = pi / 2 # wedge angle
edge_length = 12e-6 # length of wedge arms
num_wedges_per_lane = 4 # number of wedges per lane
num_lanes = 3 # number of lanes of funnels
wedge_to_wedge = simulation.L_y / num_wedges_per_lane # distance between wedges
lane_to_lane = simulation.L_x / num_lanes # distance between lanes
for i = 0 : num_lanes - 1, j = 0 : num_wedges_per_lane - 1
    tip_position = [lane_to_lane / 2 + i * lane_to_lane - edge_length / (2 * sqrt(2)), wedge_to_wedge / 2 + j * wedge_to_wedge]
    append!(boundary_positions, wedge(angle = angle, edge_length = edge_length, spacing = R_boundary, tip = tip_position))
end

# remove active particles that overlap with boundary particles
remove_overlaps!(swimmer_positions; fixed_positions = boundary_positions, minimum_distance = 2^(1.0 / 6.0) * (R_swimmer + R_boundary), period_x = simulation.L_x, period_y = simulation.L_y)

# create the actual particles with properties
for (x, y) in swimmer_positions
    active_force = RunAndTumble(γv = γ_trans_swimmer * 15e-6, α = 1.0)
    push!(simulation.particles, Particle(ptype = :swimmer, x = x, y = y, R = R_swimmer, γ_trans = γ_trans_swimmer, D_trans = D_trans_swimmer, active_force = active_force))
end
for (x, y) in boundary_positions
    push!(simulation.particles, Particle(ptype = :boundary, x = x, y = y, R = R_boundary, D_trans = 0.0))
end
# particle groups are useful for interactions and integrators
pgroup_swimmer = group_by_type(simulation.particles; ptype = :swimmer)
pgroup_all = simulation.particles

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# Initialize cell list and Lennard-Jones interactions for swimmer-swimmer and swimmer-boundary
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
cell_list = CellList(particles = pgroup_all, L_x = simulation.L_x, L_y = simulation.L_y, cutoff = 2^(1.0 / 6.0) * (R_swimmer + R_boundary))
lj = LennardJones(particles = pgroup_swimmer, cell_list = cell_list, ϵ = 1.38e-23 * 300, multithreaded = true)
push!(simulation.cell_lists, cell_list)
push!(simulation.interactions, lj)

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# Initialize Brownian integrator for overdamped dynamics
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
simulation.dt = 0.0001
# only swimmers will be integrated since boundary doesn't move
brownian = Brownian(particles = pgroup_swimmer, dt = simulation.dt, multithreaded = true)
push!(simulation.integrators, brownian)

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# Set the number of steps and which particles to save periodically
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
simulation.num_steps = trunc(Int64, 10 / simulation.dt)
simulation.save_interval = trunc(Int64, 1 / simulation.dt)
simulation.particles_to_save = pgroup_all

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# Run simulation and save data
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
run_simulation!(simulation; save_as = "out/Example_RunAndTumbleFunnels_data.out")

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# OPTIONAL: Load simulation data and animate the simulation.
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
simulation = load_simulation(file = "out/Example_RunAndTumbleFunnels_data.out")
visualize!(simulation; save_as = (:gif, "frames/Example_RunAndTumbleFunnels.gif"), fps = 1, particle_colors = Dict(:boundary => "black", :swimmer => "green"))


text:
# In this tutorial, we will learn
# - How to implement the method of manufactured solutions
# - How to perform a convergence test
# - How to define the discretization error
# - How to integrate error norms
# - How to generate Cartesian meshes in arbitrary dimensions
#
# ## Problem statement
#
# In this tutorial, we show how to validate a code using the well known *method of manufactured solutions*. For the sake of simplicity, we consider the Poisson equation in the unit square $\Omega\doteq (0,1)^2$ as a model problem,
#
#
# ```math
# \left\lbrace
# \begin{aligned}
# -\Delta u = f  \ \text{in} \ \Omega\\
# u = g \ \text{on}\ \partial\Omega.\\
# \end{aligned}
# \right.
# ```
#
# We are going to consider two different manufactured solutions. On the one hand, we consider function $u(x)=x_1+x_2$, which can be exactly represented by the FE interpolation that we construct below. Thus, one expects that the obtained approximation error is near the machine precision. We are going to check that this is true in the code. On the other hand, we consider a function that cannot be captured exactly by the interpolation, namely $u(x)=x_2 \sin(2 \pi\ x_1)$. Here, our goal is to confirm that the convergence order of the discretization error is the optimal one.
#
#
# ## Manufactured solution
#
# We start by defining the manufactured solution $u(x) = x_1+x_2$ and the source term $f$ associated with it, namely $f\doteq-\Delta(x_1+x_2)=0$.

using Gridap

u(x) = x[1] + x[2]
f(x) = 0

# We also need to define the gradient of $u$ since we will compute the $H^1$ error norm later. In that case, the gradient is simply defined as
#

∇u(x) = VectorValue(1,1)

# Note that we have used the constructor `VectorValue` to build the vector that represents the gradient. However, we still need a final trick. We need to tell the Gridap library that the gradient of the function `u` is available in the function `∇u` (at this moment `u` and `∇u` are two standard Julia functions without any connection between them). This is done by adding an extra method to the function `gradient` (aka `∇`) defined in Gridap:

import Gridap: ∇
∇(::typeof(u)) = ∇u

# Now, it is possible to recover function `∇u` from function `u` as `∇(u)`. You can check that the following expression evaluates to `true`.

∇(u) === ∇u

# Note: the definition of the gradient is optional. If not provided, the gradient will be computed with automatic differentiation.

# ## Cartesian mesh generation
#
# In order to discretize the geometry of the unit square, we use the Cartesian mesh generator available in Gridap:

domain = (0,1,0,1)
partition = (4,4)
model = CartesianDiscreteModel(domain,partition)

# The type `CartesianDiscreteModel` is a concrete type that inherits from `DiscreteModel`, which is specifically designed for building Cartesian meshes. The `CartesianDiscreteModel` constructor takes a tuple containing limits of the box we want to discretize  plus a tuple with the number of cells to be generated in each direction (here 4 by 4 cells). Note that the `CaresianDiscreteModel` is implemented for arbitrary dimensions. For instance, the following lines build a `CartesianDiscreteModel`  for the unit cube $(0,1)^3$ with 4 cells per direction

domain3d = (0,1,0,1,0,1)
partition3d = (4,4,4)
model3d = CartesianDiscreteModel(domain3d,partition3d)

# You could also generate a mesh for the unit tesseract $(0,1)^4$ (i.e., the unit cube in 4D). Look how the 2D and 3D models are built and just follow the sequence.
#

# Let us return to the 2D `CartesianDiscreteModel` that we have already constructed. You can inspect it by writing it into vtk format. Note that you can also print a 3D model, but not a 4D one. In the future, it would be cool to generate a movie from a 4D model, but this functionality is not yet implemented.

writevtk(model,"model")

# If you open the generated files, you will see that the boundary vertices and facets are identified with the name "boundary". This is just what we need to impose the Dirichlet boundary conditions in this example.
#
# These are the vertices in the model
#
# ![](../assets/validation/model_0.png)
#
# and these the facets
#
# ![](../assets/validation/model_1.png)
#
# As you can see in the pictures, the objects on the boundary are correctly tagged with the name "boundary".

# ## FE approximation
#
# We compute a FE approximation of the Poisson problem above by following the steps detailed in the previous tutorial:

order = 1
reffe = ReferenceFE(lagrangian,Float64,order)
V0 = TestFESpace(model,reffe,conformity=:H1,dirichlet_tags="boundary")
U = TrialFESpace(V0,u)

degree = 2
Ω = Triangulation(model)
dΩ = Measure(Ω,degree)

a(u,v) = ∫( ∇(v)⊙∇(u) )*dΩ
b(v) = ∫( v*f )*dΩ

op = AffineFEOperator(a,b,U,V0)

uh = solve(op)

# Note that we are imposing Dirichlet boundary conditions on the objects tagged as "boundary" and that we are using the manufactured solution `u` to construct the trial FE space. Not also that we are not explicitly constructing an `Assembler` object nor a `FESolver`. We are relying on default values.
#
#
# ## Measuring the discretization error
#
# Our goal is to check that the discratization error associated with the computed approximation `uh` is close to machine precision. To this end, the first step is to compute the discretization error, which is done as you would expect:

e = u - uh

# Once the error is defined, you can, e.g., visualize it.

writevtk(Ω,"error",cellfields=["e" => e])

# This generates a file called `error.vtu`. Open it with Paraview to check that the error is of the order of the machine precision.
#
# ![](../assets/validation/error.png)
#
# A more rigorous way of quantifying the error is to measure it with a norm. Here, we use the $L^2$ and $H^1$ norms, which are defined as
#
# ```math
# \| w \|_{L^2}^2 \doteq \int_{\Omega} w^2 \ \text{d}\Omega, \quad
# \| w \|_{H^1}^2 \doteq \int_{\Omega} w^2 + \nabla w \cdot \nabla w \ \text{d}\Omega.
# ```
#
# In order to compute these norms, we use again the `∫` function and the integration measure `dΩ`, namely

el2 = sqrt(sum( ∫( e*e )*dΩ ))
eh1 = sqrt(sum( ∫( e*e + ∇(e)⋅∇(e) )*dΩ ))

# The expression `∫( fun )*dΩ` returns an object storing the cell contributions of the integral of the given function `fun`.
#  To end up with the desired error norms, one has to sum these contributions and take the square root. You can check that the computed error norms are close to machine precision (as one would expect).

tol = 1.e-10
@assert el2 < tol
@assert eh1 < tol

# ## Convergence test
#
# We end up this tutorial by performing a convergence test, where we are going to use all the new concepts we have learned.  We will consider a manufactured solution that does not belong to the FE interpolation space. In this test, we expect to see the optimal convergence order of the FE discretization.

# Here, we define the manufactured functions
p = 3
u(x) = x[1]^p+x[2]^p
∇u(x) = VectorValue(p*x[1]^(p-1),p*x[2]^(p-1))
f(x) = -p*(p-1)*(x[1]^(p-2)+x[2]^(p-2))

# Since we have redefined the valiables `u`, `∇u`, and `f`, we need to execute these lines again

∇(::typeof(u)) = ∇u
b(v) = ∫( v*f )*dΩ

# In order to perform the convergence test, we write in a function all the code needed to perform a single computation and measure its error. The input of this function is the number of cells in each direction and the interpolation order. The output is the computed $L^2$ and $H^1$ error norms.

function run(n,k)

  domain = (0,1,0,1)
  partition = (n,n)
  model = CartesianDiscreteModel(domain,partition)

  reffe = ReferenceFE(lagrangian,Float64,k)
  V0 = TestFESpace(model,reffe,conformity=:H1,dirichlet_tags="boundary")
  U = TrialFESpace(V0,u)

  degree = 2*p
  Ω = Triangulation(model)
  dΩ = Measure(Ω,degree)

  a(u,v) = ∫( ∇(v)⊙∇(u) )*dΩ
  b(v) = ∫( v*f )*dΩ

  op = AffineFEOperator(a,b,U,V0)

  uh = solve(op)

  e = u - uh

  el2 = sqrt(sum( ∫( e*e )*dΩ ))
  eh1 = sqrt(sum( ∫( e*e + ∇(e)⋅∇(e) )*dΩ ))

  (el2, eh1)

end

# The following function does the convergence test. It takes a vector of integers (representing the number of cells per direction in each computation) plus the interpolation order. It returns the $L^2$ and $H^1$ error norms for each computation as well as the corresponding cell size.

function conv_test(ns,k)

  el2s = Float64[]
  eh1s = Float64[]
  hs = Float64[]

  for n in ns

    el2, eh1 = run(n,k)
    h = 1.0/n

    push!(el2s,el2)
    push!(eh1s,eh1)
    push!(hs,h)

  end

  (el2s, eh1s, hs)

end

# We are ready to perform the test! We consider several mesh sizes and interpolation order $k=1$ and $k=2$ (for $k=3$ the error will be close to machine precision, as before, except if you change the value of $p$ above).

el2s1, eh1s1, hs = conv_test([8,16,32,64,128],1);
el2s2, eh1s2, hs = conv_test([8,16,32,64,128],2);

# With the generated data, we do the classical convergence plot.

using Plots

plot(hs,[el2s1 eh1s1 el2s2 eh1s2],
    xaxis=:log, yaxis=:log,
    label=["L2 k=1" "H1 k=1" "L2 k=2" "H1 k=2"],
    shape=:auto,
    xlabel="h",ylabel="error norm")

#src savefig("conv.png")

#md # If you run the code in a notebook, you will see a figure like this one:
#md # ![](../assets/validation/conv.png)
#
#
# The generated curves make sense. For a given interpolation order it is observed that the convergence of the $H^1$ error is slower that $L^2$ one whereas increasing the order makes convergence faster both in $L^2$ and in $H^1$. However, in order to be more conclusive, we need to compute the slope of these lines. It can be done with this little function that internally uses a linear regression.

function slope(hs,errors)
  x = log10.(hs)
  y = log10.(errors)
  linreg = hcat(fill!(similar(x), 1), x) \ y
  linreg[2]
end

# The slopes for the $L^2$ error norm is computed as

slope(hs,el2s1)
slope(hs,el2s2)

# and for the $H^1$ error norm as

slope(hs,eh1s1)
slope(hs,eh1s2)

#md # If your run these lines in a notebook, you will see that
#nb # As you can see,
# the slopes for the $L^2$ and $H^1$ error norms are 2 and 1 for $k=1$ and 3 and 2 for $k=2$ respectively.
#
# Congrats, another tutorial done!

max_stars_repo_name: principejavier/Tutorials

max_issues_repo_name: principejavier/Tutorials

max_forks_repo_name: principejavier/Tutorials


text:
x = range(0, stop=1, length=25)
y = @. sin(x)
z = @. cos(x)
t = y .+ z
scatter(x, hcat(y, z), t)


text:
export AbstractPolynomial

const SymbolLike = Union{AbstractString,Char,Symbol}

"""
    AbstractPolynomial{T,X}

An abstract container for various polynomials.

# Properties
- `coeffs` - The coefficients of the polynomial
"""
abstract type AbstractPolynomial{T,X} end

# We want  ⟒(P{α…,T}) = P{α…}; this default
# works for most cases
⟒(P::Type{<:AbstractPolynomial}) = constructorof(P)

# convert `as` into polynomial of type P based on instance, inheriting variable
# (and for LaurentPolynomial the offset)
_convert(p::P, as) where {T,X,P <: AbstractPolynomial{T,X}} = ⟒(P)(as, X)

"""
    Polynomials.@register(name)

Given a polynomial with `name`, creates some common convenience constructors and conversions to minimize code required for implementation of a new polynomial type.

# Example
```julia
struct MyPolynomial{T,X} <: AbstractPolynomial{T,X} end

Polynomials.@register MyPolynomial
```

# Implementations
This will implement simple self-conversions like `convert(::Type{MyPoly}, p::MyPoly) = p` and creates two promote rules. The first allows promotion between two types (e.g. `promote(Polynomial, ChebyshevT)`) and the second allows promotion between parametrized types (e.g. `promote(Polynomial{T}, Polynomial{S})`).

For constructors, it implements the shortcut for `MyPoly(...) = MyPoly{T}(...)`, singleton constructor `MyPoly(x::Number, ...)`,  conversion constructor `MyPoly{T}(n::S, ...)`, and `variable` alternative  `MyPoly(var=:x)`.
"""
macro register(name)
    poly = esc(name)
    quote
        Base.convert(::Type{P}, p::P) where {P<:$poly} = p
        function Base.convert(P::Type{<:$poly}, p::$poly{T,X}) where {T,X}
            isconstant(p) && return constructorof(P){eltype(P),indeterminate(P)}(constantterm(p))
            constructorof(P){eltype(P), indeterminate(P,p)}(coeffs(p))
        end
        Base.promote(p::P, q::Q) where {X, T, P <:$poly{T,X}, Q <: $poly{T,X}} = p,q
        Base.promote_rule(::Type{<:$poly{T,X}}, ::Type{<:$poly{S,X}}) where {T,S,X} =  $poly{promote_type(T, S),X}
        Base.promote_rule(::Type{<:$poly{T,X}}, ::Type{S}) where {T,S<:Number,X} =
            $poly{promote_type(T, S),X}
        $poly(coeffs::AbstractVector{T}, var::SymbolLike = :x) where {T} =
            $poly{T, Symbol(var)}(coeffs)
        $poly{T}(x::AbstractVector{S}, var::SymbolLike = :x) where {T,S<:Number} =
            $poly{T,Symbol(var)}(T.(x))
        function $poly(coeffs::G, var::SymbolLike=:x) where {G}
            !Base.isiterable(G) && throw(ArgumentError("coeffs is not iterable"))
            cs = collect(coeffs)
            $poly{eltype(cs), Symbol(var)}(cs)
        end
        $poly{T,X}(n::S) where {T, X, S<:Number} =
            T(n) *  one($poly{T, X})
        $poly{T}(n::S, var::SymbolLike = :x) where {T, S<:Number} =
            T(n) *  one($poly{T, Symbol(var)})
        $poly(n::S, var::SymbolLike = :x)  where {S  <: Number} = n * one($poly{S, Symbol(var)})
        $poly{T}(var::SymbolLike=:x) where {T} = variable($poly{T, Symbol(var)})
        $poly(var::SymbolLike=:x) = variable($poly, Symbol(var))
        (p::$poly)(x) = evalpoly(x, p)
    end
end

macro registerN(name, params...)
    poly = esc(name)
    αs = tuple(esc.(params)...)
    quote
        Base.convert(::Type{P}, q::Q) where {$(αs...),T, P<:$poly{$(αs...),T}, Q <: $poly{$(αs...),T}} = q
        Base.convert(::Type{$poly{$(αs...)}}, q::Q) where {$(αs...),T, Q <: $poly{$(αs...),T}} = q
        Base.promote(p::P, q::Q) where {$(αs...),T, X, P <:$poly{$(αs...),T,X}, Q <: $poly{$(αs...),T,X}} = p,q
        Base.promote_rule(::Type{<:$poly{$(αs...),T,X}}, ::Type{<:$poly{$(αs...),S,X}}) where {$(αs...),T,S,X} =
            $poly{$(αs...),promote_type(T, S),X}
        Base.promote_rule(::Type{<:$poly{$(αs...),T,X}}, ::Type{S}) where {$(αs...),T,X,S<:Number} =
            $poly{$(αs...),promote_type(T,S),X}

        function $poly{$(αs...),T}(x::AbstractVector{S}, var::SymbolLike = :x) where {$(αs...),T,S}
            $poly{$(αs...),T, Symbol(var)}(T.(x))
        end
        $poly{$(αs...)}(coeffs::AbstractVector{T}, var::SymbolLike=:x) where {$(αs...),T} =
            $poly{$(αs...),T,Symbol(var)}(coeffs)

        $poly{$(αs...),T,X}(n::Number) where {$(αs...),T,X} = T(n)*one($poly{$(αs...),T,X})
        $poly{$(αs...),T}(n::Number, var::SymbolLike = :x) where {$(αs...),T} = T(n)*one($poly{$(αs...),T,Symbol(var)})
        $poly{$(αs...)}(n::S, var::SymbolLike = :x) where {$(αs...), S<:Number} =
            n*one($poly{$(αs...),S,Symbol(var)})
        $poly{$(αs...),T}(var::SymbolLike=:x) where {$(αs...), T} = variable($poly{$(αs...),T,Symbol(var)})
        $poly{$(αs...)}(var::SymbolLike=:x) where {$(αs...)} = variable($poly{$(αs...)},Symbol(var))
        (p::$poly)(x) = evalpoly(x, p)
    end
end


text:
# template log-on and fetch graph

using Caesar
using IncrementalInference

include(joinpath(dirname(@__FILE__),"blandauthremote.jl"))
addrdict["session"] = "SESS??"
addrdict["robot"] = "robot"
addrdict["user"] = "user"
cloudGraph, addrdict = standardcloudgraphsetup(addrdict=addrdict)

fg = initfg(cloudgraph=cloudGraph, sessionname=addrdict["session"])
fullLocalGraphCopy!(fg)
savejld(fg, file="somefg.jld")

plotKDE(fg, :x1, dims=[1;2], api=localapi)
fncvar = getfnctype(getVert(fg,fg.fIDs[:x1710l200050]))
# TODO: Please fix with correct session for example.
ret = whosNear2D(cloudGraph, "SESS??", robot, user, x=21.7, y=-54.0)
plotLocalProduct(fg, :l200050, dims=[1;2], api=localapi)

#


text:
module GPUBenchmarks

using GPUArrays, CUDAdrv, FileIO, BenchmarkTools
import CuArrays, CLArrays, CUDAnative, CUDAdrv, OpenCL

function devices()
    (:cuarrays, :clarrays_cpu, :clarrays_gpu, :juliabase, #=:clarrayfire, :cuarrayfire=#)
end

function normit(x, mini, maxi)
    width = (maxi - mini)
    width == 0 && return x
    (x .- mini) ./ width
end

function meandifference(juliabaseline, gpuarray)
    mini, maxi = extrema(juliabaseline)
    normed1 = normit(juliabaseline, mini, maxi)
    normed2 = normit(Array(gpuarray), mini, maxi)
    mean(abs.(normed1 .- normed2))
end

function init(device::OpenCL.cl.Device)
    CLArrays.init(device)
    GPUArrays.name(device), CLArrays.CLArray
end

function init(device::Symbol)
    if device == :cuarrays
        # cuarrays uses the default device 0
        CUDAdrv.name(CUDAnative.CuDevice(0)), CuArrays.CuArray
    elseif device == :juliabase
        FFTW.set_num_threads(1)
        Sys.cpu_info()[1].model, Array
    elseif device == :clarrayfire
        ArrayFire.set_backend(ArrayFire.AF_BACKEND_OPENCL)
        gpu_devices = first(OpenCL.cl.devices(:gpu))
        replace(gpu_devices[:name], r"\s+", " "), ArrayFire.AFArray
    elseif device == :cuarrayfire
        ArrayFire.set_backend(ArrayFire.AF_BACKEND_CUDA)
        CUDAdrv.name(CUDAnative.CuDevice(ArrayFire.get_device())), ArrayFire.AFArray
    elseif device == :clarrays_cpu
        device = first(CLArrays.devices(CLArrays.is_cpu))
        ctx = CLArrays.init(device)
        GPUArrays.name(device), CLArrays.CLArray
    elseif device == :clarrays_gpu
        device = first(CLArrays.devices(CLArrays.is_gpu))
        ctx = CLArrays.init(device)
        GPUArrays.name(device), CLArrays.CLArray
    end
end

is_cudanative(x) = x == :cuarrays
is_arrayfire(x) = x in (:clarrayfire, :cuarrayfire)
is_gpuarrays(x) = x in (:clarrays_cpu, :cuarrays, :clarrays_gpu)

synchronize(x) = GPUArrays.synchronize(x)

function synchronize(x::CuArrays.CuArray)
    ctx = CUDAdrv.CuContext(CUDAdrv.CuPrimaryContext(CUDAdrv.CuDevice(0)))
    CUDAdrv.synchronize(ctx)
    x
end

# function synchronize(x::ArrayFire.AFArray)
#     ArrayFire.afgc()
#     ArrayFire.sync(x)
#     return x
# end

free(x) = finalize(x)

dir(paths...) = normpath(joinpath(@__DIR__, "..", paths...))

"""
Runs a script in a new julia process.
Usage:
```
@run_julia (JULIA_NUM_THREADS = 8, "-O3",  "--math-mode=fast") begin
    println(Threads.nthreads())
end
```
"""
macro run_julia(args, expr)
    envs = String[]
    julia_args = []
    variables = []
    if isa(args, Expr) && args.head == :tuple
        for elem in args.args
            if isa(elem, String)
                push!(julia_args, elem)
            elseif isa(elem, Expr) && elem.head == :(=)
                key, value = elem.args
                push!(envs, string(key, "=", value))
            elseif isa(elem, Symbol)
                push!(variables, elem)
            else
                error("Unsupported argument: $elem")
            end
        end
    else
        error("Unsupported expression for args $args")
    end

    new_args = []
    for arg in expr.args
        if isa(arg, Expr) && arg.head == :toplevel
            for toplevel in arg.args
                push!(new_args, toplevel)
            end
        else
            push!(new_args, arg)
        end
    end
    expr.args = new_args
    str = string(expr)
    command_1 = ["julia", julia_args...]
    variable_expr = map(variables) do var
        quote
            src_str *= sprint() do io
                val = $(esc(var))
                print(io, $(string(var)), " = ")
                show(io, val)
                println(io)
            end
        end
    end
    expr = quote
        src_str = ""
        $(variable_expr...)
        src_str *= $(string(expr))
        command_2 = Cmd([$(command_1)..., "-e", src_str])
        command_3 = Cmd(command_2, env = $(envs))
        run(command_3)
    end
    expr
end

include("database.jl")

export devices, init, is_cudanative, free, synchronize, is_arrayfire, is_gpuarrays
export @run_julia, BenchResult, meandifference

end # module