text:
using Test
using JuMP
import DiffOpt
import MathOptInterface
const MOI = MathOptInterface
const MOIU = MOI.Utilities
import LinearAlgebra: dot, ⋅, I
import Ipopt
import OSQP
import GLPK
import SCS
import DelimitedFiles

@testset "Testing forward on trivial QP" begin
    # using example on https://osqp.org/docs/examples/setup-and-solve.html
    Q = [
        4.0 1.0
        1.0 2.0
    ]
    q = [1.0, 1.0]
    G = [
         1.0 1.0
         1.0 0.0
         0.0 1.0
        -1.0 -1.0
        -1.0 0.0
         0.0 -1.0
    ]
    h = [1, 0.7, 0.7, -1, 0, 0]

    model = JuMP.Model(() -> DiffOpt.diff_optimizer(Ipopt.Optimizer))
    MOI.set(model, MOI.Silent(), true)
    @variable(model, x[1:2])
    @objective(model, Min, dot(Q * x, x) + dot(q, x))
    @constraint(model,
        G * x .<= h,
    )
    optimize!(model)

    @test JuMP.value.(x) ≈ [0.3, 0.7] atol=ATOL rtol=RTOL
end

@testset "Differentiating trivial QP 1" begin
    Q = [
        4.0 1.0
        1.0 2.0
    ]
    q = [1.0, 1.0]
    G = [1.0 1.0]
    h = [-1.0]

    model = JuMP.direct_model(DiffOpt.diff_optimizer(Ipopt.Optimizer))
    MOI.set(model, MOI.Silent(), true)
    x = @variable(model, [1:2])
    @objective(model, Min, dot(Q * x, x) + dot(q, x))
    @constraint(model, ctr_le,
        G * x .<= h,
    )
    optimize!(model)

    @test JuMP.value.(x) ≈ [-0.25, -0.75] atol=ATOL rtol=RTOL

    # grad_wrt_h = backward(JuMP.backend(model), ["h"], ones(2))[1]

    MOI.set.(model, DiffOpt.BackwardInVariablePrimal(), x, 1.0)

    DiffOpt.backward(model)

    grad = JuMP.constant(MOI.get(model, DiffOpt.BackwardOutConstraint(), ctr_le[]))
    @test grad ≈ -1.0  atol=ATOL rtol=RTOL

    # TODO: this simple show fails
    # @show ctr_le
end

@testset "Differentiating QP with inequality and equality constraints" begin
    # refered from: https://www.mathworks.com/help/optim/ug/quadprog.html#d120e113424
    # Find equivalent qpth program here - https://github.com/AKS1996/jump-gsoc-2020/blob/master/DiffOpt_tests_4_py.ipynb

    Q = [
         1.0 -1.0 1.0;
        -1.0  2.0 -2.0;
         1.0 -2.0 4.0
    ]
    q = [2.0, -3.0, 1.0]
    G = [0.0 0.0 1.0
         0.0 1.0 0.0
         1.0 0.0 0.0
         0.0 0.0 -1.0
         0.0 -1.0 0.0
         -1.0 0.0 0.0
    ]
    h = [1.0, 1.0, 1.0, 0.0, 0.0, 0.0,]
    A = [1.0 1.0 1.0;]
    b = [0.5]

    model = Model(() -> DiffOpt.diff_optimizer(Ipopt.Optimizer))
    MOI.set(model, MOI.Silent(), true)
    @variable(model, x[1:3])
    @objective(model, Min, dot(Q * x, x) + dot(q, x))
    @constraint(model, ctr_le,
        G * x .<= h,
    )
    @constraint(model, ctr_eq,
        A * x .== b,
    )
    optimize!(model)
    # doptimizer = JuMP.backend(model).optimizer.model

    # nv = MOI.get(doptimizer, MOI.NumberOfVariables())
    # v = MOI.VariableIndex.(collect(1:nv))
    # z = MOI.get.(doptimizer, MOI.VariablePrimal(), v)

    @test JuMP.value.(x) ≈ [0.0, 0.5, 0.0] atol=ATOL rtol=RTOL

    MOI.set.(model, DiffOpt.BackwardInVariablePrimal(), x, 1.0)

    DiffOpt.backward(model)

    @test MOIU.isapprox_zero(moi_function(MOI.get(model, DiffOpt.BackwardOutObjective())), ATOL)

    db = [1.0]
    dA = [0.0 -0.5 0.0]
    for (j,jc) in enumerate(ctr_eq)
        grad = MOI.get(model, DiffOpt.BackwardOutConstraint(), jc)
        @test JuMP.constant(grad) ≈ -db[j] atol=ATOL rtol=RTOL
        for (i,iv) in enumerate(x)
            @test JuMP.coefficient(grad, iv) ≈ dA[j,i] atol=ATOL rtol=RTOL
        end
    end

    dh = zeros(6,1)
    dG = zeros(6,3)
    for (j,jc) in enumerate(ctr_le)
        grad = MOI.get(model, DiffOpt.BackwardOutConstraint(), jc)
        @test JuMP.constant(grad) ≈ -dh[j] atol=ATOL rtol=RTOL
        for (i,iv) in enumerate(x)
            @test JuMP.coefficient(grad, iv) ≈ dG[j,i] atol=ATOL rtol=RTOL
        end
    end
end

# refered from https://github.com/jump-dev/MathOptInterface.jl/blob/master/src/Test/contquadratic.jl#L3
# Find equivalent CVXPYLayers and QPTH code here:
#               https://github.com/AKS1996/jump-gsoc-2020/blob/master/DiffOpt_tests_1_py.ipynb
@testset "Differentiating MOI examples 1" begin
    # homogeneous quadratic objective
    # Min x^2 + xy + y^2 + yz + z^2
    # st  x + 2y + 3z >= 4 (c1)
    #     x +  y      >= 1 (c2)
    #     x, y, z \in R

    model = JuMP.direct_model(DiffOpt.diff_optimizer(OSQP.Optimizer))
    MOI.set(JuMP.backend(model).optimizer, MOI.RawOptimizerAttribute("eps_prim_inf"), 1e-7)
    MOI.set(JuMP.backend(model).optimizer, MOI.RawOptimizerAttribute("eps_dual_inf"), 1e-7)
    MOI.set(model, MOI.Silent(), true)
    @variables(model, begin
        x
        y
        z
    end)
    @constraint(model, c1,
        x + 2y + 3z >= 4
    )
    @constraint(model, c2,
        x +  y      >= 1
    )
    @objective(model, Min,
        x^2 + x*y + y^2 + y*z + z^2
    )

    optimize!(model)

    z = [x, y, z]

    @test JuMP.value.(z) ≈ [4/7, 3/7, 6/7] atol=ATOL rtol=RTOL

    MOI.set.(model, DiffOpt.BackwardInVariablePrimal(), z, 1.0)

    DiffOpt.backward(model)

    dl_dq = [-0.2142857;  0.21428567; -0.07142857]
    dl_dQ = [-0.12244895  0.01530609 -0.11224488;
              0.01530609  0.09183674  0.07653058;
             -0.11224488  0.07653058 -0.06122449]
    expected = dl_dq' * z + z' * (dl_dQ / 2.0) * z
    @test moi_function(MOI.get(model, DiffOpt.BackwardOutObjective())) ≈ moi_function(expected)  atol=ATOL rtol=RTOL

    dh = [0.35714284; 0.4285714]
    # This ws already broken
    # @test_broken dl_dG ≈ [0.05102035   0.30612245  0.255102;
    #                0.06122443   0.36734694  0.3061224] atol=ATOL rtol=RTOL
    dG = -[0.05102035   0.30612245  0.255102;
           0.06122443   0.36734694  0.3061224]
    for (j,jc) in enumerate([c1, c2])
        grad = MOI.get(model, DiffOpt.BackwardOutConstraint(), jc)
        @test JuMP.constant(grad) ≈ -dh[j] atol=ATOL rtol=RTOL
        for (i,iv) in enumerate(z)
            @test_broken JuMP.coefficient(grad, iv) ≈ dG[j,i] atol=ATOL rtol=RTOL
        end
    end
end

@testset "Differentiating MOI examples 2 - non trivial backward pass vector" begin
    # non-homogeneous quadratic objective
    #    minimize 2 x^2 + y^2 + xy + x + y
    #       s.t.  x, y >= 0
    #             x + y = 1 (c1)

    model = JuMP.Model(() -> DiffOpt.diff_optimizer(Ipopt.Optimizer))
    MOI.set(model, MOI.Silent(), true)
    @variable(model, x ≥ 0);
    @variable(model, y ≥ 0);
    @constraint(model, c1, x + y == 1)
    @objective(model, Min, 2 * x^2 + y^2 + x * y + x + y)

    optimize!(model)

    z = [x, y]

    @test JuMP.value.(z) ≈ [0.25, 0.75] atol=ATOL rtol=RTOL

    MOI.set.(model, DiffOpt.BackwardInVariablePrimal(), z, [1.3, 0.5])

    DiffOpt.backward(model)

    dl_dq = [-0.2; 0.2]
    dl_dQ = [-0.05   -0.05;
             -0.05    0.15]
    expected = dl_dq' * z + z' * (dl_dQ / 2.0) * z
    @test moi_function(MOI.get(model, DiffOpt.BackwardOutObjective())) ≈ moi_function(expected)  atol=ATOL rtol=RTOL

    func = MOI.get(model, DiffOpt.BackwardOutConstraint(), c1)
    @test -0.7 ≈ JuMP.constant(func) atol=ATOL rtol=RTOL

    # This ws already broken
    # @test_broken dl_dG ≈ [0.05102035   0.30612245  0.255102;
    #                0.06122443   0.36734694  0.3061224] atol=ATOL rtol=RTOL
    dl_dA = [0.375 -1.075]
    for (j, vi) in enumerate(z)
        @test dl_dA[j] ≈ JuMP.coefficient(func, vi) atol=ATOL rtol=RTOL
    end

    c_le = [LowerBoundRef(x), LowerBoundRef(y)]

    dl_dh = [1e-8; 1e-8]
    dl_dG = [1e-8  1e-8; 1e-8 1e-8]
    for (j,jc) in enumerate(c_le)
        grad = MOI.get(model, DiffOpt.BackwardOutConstraint(), jc)
        @test JuMP.constant(grad) ≈ dl_dh[j] atol=ATOL rtol=RTOL
        for (i,iv) in enumerate(z)
            @test JuMP.coefficient(grad, iv) ≈ dl_dG[j,i] atol=ATOL rtol=RTOL
        end
    end

end

@testset "Differentiating non trivial convex QP JuMP" begin
    nz = 10
    nineq_le = 25
    neq = 10

    # read matrices from files
    names = ["P", "q", "G", "h", "A", "b"]
    matrices = []

    for name in names
        push!(matrices, DelimitedFiles.readdlm(joinpath(dirname(dirname(pathof(DiffOpt))), "test", "data", name * ".txt"), ' ', Float64, '\n'))
    end

    Q, q, G, h, A, b = matrices
    q = vec(q)
    h = vec(h)
    b = vec(b)

    model = JuMP.Model(() -> DiffOpt.diff_optimizer(Ipopt.Optimizer))
    MOI.set(model, MOI.Silent(), true)

    @variable(model, x[1:nz])

    @objective(model, Min, dot(Q * x, x) + dot(q, x))
    @constraint(model, c_le, G * x .<= h)
    @constraint(model, c_eq, A * x .== b)

    optimize!(model)

    MOI.set.(model, DiffOpt.BackwardInVariablePrimal(), x, 1.0)

    DiffOpt.backward(model)

    # obtain gradients
    # grads = backward(doptimizer, ["Q", "q", "G", "h", "A", "b"], ones(nz))  # using dl_dz=[1,1,1,1,1,....]

    # read gradients from files
    param_names = ["dP", "dq", "dG", "dh", "dA", "db"]
    grads_actual = []

    for name in param_names
        push!(grads_actual,
            DelimitedFiles.readdlm(
                joinpath(@__DIR__, "data", "$(name).txt"),
                ' ', Float64, '\n',
            )
        )
    end

    dq = grads_actual[2] # = vec(grads_actual[2])
    dh = grads_actual[4] # = vec(grads_actual[4])
    db = grads_actual[6] # = vec(grads_actual[6])

    grad = MOI.get(model, DiffOpt.BackwardOutObjective())
    @test moi_function(grad) ≈ moi_function(dq ⋅ x) atol=1e-2 rtol=1e-2

    for (i, ci) in enumerate(c_le)
        @test -dh[i] ≈ JuMP.constant(MOI.get(model,
            DiffOpt.BackwardOutConstraint(), ci)) atol=1e-2 rtol=1e-2
    end

    for (i, ci) in enumerate(c_eq)
        @test -db[i] ≈ JuMP.constant(MOI.get(model,
            DiffOpt.BackwardOutConstraint(), ci)) atol=1e-2 rtol=1e-2
    end

    # # testing differences
    # for i in 1:size(grads)[1]
    #     @test grads[i] ≈  grads_actual[i] atol=1e-2 rtol=1e-2
    # end
end

@testset "Differentiating LP; checking gradients for non-active contraints" begin
    # Issue #40 from Gurobi.jl
    # min  x
    # s.t. x >= 0
    #      x >= 3

    model = direct_model(DiffOpt.diff_optimizer(GLPK.Optimizer))
    MOI.set(model, MOI.Silent(), true)

    @variable(model, x[1:1])

    @objective(model, Min, 1.1 * x[1])
    @constraint(model, c1, x[1] ≥ 0)
    @constraint(model, c2, x[1] ≥ 3)

    optimize!(model)

    # obtain gradients
    MOI.set.(model, DiffOpt.BackwardInVariablePrimal(), x, 1.0)

    DiffOpt.backward(model)

    c_le = [c1, c2]

    dh = [0.0, 1.0]
    dG = [0.0, -3.0]
    for (j,jc) in enumerate(c_le)
        grad = MOI.get(model, DiffOpt.BackwardOutConstraint(), jc)
        @test JuMP.constant(grad) ≈ -dh[j] atol=ATOL rtol=RTOL
        for (i,iv) in enumerate(x)
            @test JuMP.coefficient(grad, iv) ≈ dG[j,i] atol=ATOL rtol=RTOL
        end
    end

    model = direct_model(DiffOpt.diff_optimizer(GLPK.Optimizer))
    MOI.set(model, MOI.Silent(), true)

    @variable(model, x[1:1])

    @objective(model, Min, x[1])
    @constraint(model, c1, x[1] ≥ 0)
    @constraint(model, c2, x[1] ≥ 3)

    optimize!(model)

    MOI.set.(model, DiffOpt.BackwardInVariablePrimal(), x, 1.0)

    DiffOpt.backward(model)

    c_le = [c1, c2]

    dh = [0.0, 1.0]
    dG = [0.0, -3.0]
    for (j,jc) in enumerate(c_le)
        grad = MOI.get(model, DiffOpt.BackwardOutConstraint(), jc)
        @test JuMP.constant(grad) ≈ -dh[j] atol=ATOL rtol=RTOL
        for (i,iv) in enumerate(x)
            @test JuMP.coefficient(grad, iv) ≈ dG[j,i] atol=ATOL rtol=RTOL
        end
    end
end

@testset "Differentiating LP; checking gradients for non-active contraints" begin
    # refered from - https://en.wikipedia.org/wiki/Simplex_algorithm#Example

    # max 2x + 3y + 4z
    # s.t. 3x+2y+z <= 10
    #      2x+5y+3z <= 15
    #      x,y,z >= 0

    model = direct_model(DiffOpt.diff_optimizer(SCS.Optimizer))
    MOI.set(model, MOI.Silent(), true)
    @variable(model, v[1:3] ≥ 0)
    @objective(model, Min, dot([-2.0, -3.0, -4.0], v))
    (x, y, z) = v

    @constraint(model, c1, 3x+2y+z <= 10)
    @constraint(model, c2, 2x+5y+3z <= 15)
    optimize!(model)

    MOI.set.(model, DiffOpt.BackwardInVariablePrimal(), v, 1.0)

    DiffOpt.backward(model)

    dQ = zeros(3,3)
    dc = zeros(3)
    dG = [0.0 0.0 0.0;
          0.0 0.0 -5/3;
          0.0 0.0 -5/3;
          0.0 0.0 10/3;
          0.0 0.0 0.0]
    dh = [0.0; 1/3; 1/3; -2/3; 0.0]

    cb = LowerBoundRef.(v)
    cc = [c1, c2]
    ctrs = vcat(cc, cb)#, cc)

    expected = dc' * v + v' * dQ * v
    grad = MOI.get(model, DiffOpt.BackwardOutObjective())
    @test moi_function(grad) ≈ moi_function(expected) atol=ATOL rtol=RTOL

    for (j,jc) in enumerate(ctrs)
        grad = MOI.get(model, DiffOpt.BackwardOutConstraint(), jc)
        @test JuMP.constant(grad) ≈ -dh[j] atol=ATOL rtol=RTOL
        for (i,iv) in enumerate(v)
            @test JuMP.coefficient(grad, iv) ≈ dG[j,i] atol=ATOL rtol=RTOL
        end
    end
end

@testset "Differentiating LP with variable bounds" begin
    # max 2x + 3y + 4z
    # s.t. 3x+2y+z <= 10
    #      2x+5y+3z <= 15
    #      x ≤ 3
    #      0 ≤ y ≤ 2
    #      z ≥ 2
    #      x,y,z >= 0
    # variant of previous test with same solution

    model = direct_model(DiffOpt.diff_optimizer(GLPK.Optimizer))
    MOI.set(model, MOI.Silent(), true)
    @variable(model, v[1:3] ≥ 0)
    (x, y, z) = v
    @objective(model, Max, 2x + 3y + 4z)
    @constraint(model, c1, 3x+2y+z <= 10)
    @constraint(model, c2, 2x+5y+3z <= 15)
    @constraint(model, c3, x ≤ 3)
    @constraint(model, c4, 1.0 * y ≤ 2)
    @constraint(model, c5, z ≥ 2)

    optimize!(model)

    MOI.set.(model, DiffOpt.BackwardInVariablePrimal(), v, 1.0)

    DiffOpt.backward(model)

    dQ = zeros(3,3)
    dc = zeros(3)
    dG = [0.0 0.0 0.0
          0.0 0.0 -5/3
          0.0 0.0 0.0
          0.0 0.0 0.0
          0.0 0.0 0.0
          0.0 0.0 -5/3
          0.0 0.0 10/3
          0.0 0.0 0.0
          ]
    dh = [0.0, 1/3, 0.0, 0.0, 0.0, 1/3, -2/3, 0.0]

    cb = LowerBoundRef.(v)
    cc = [c1, c2, c3, c4, c5]
    ctrs = vcat(cc, cb)#, cc)

    expected = dc' * v + v' * dQ * v
    grad = MOI.get(model, DiffOpt.BackwardOutObjective())
    @test moi_function(grad) ≈ moi_function(expected) atol=ATOL rtol=RTOL

    for (j,jc) in enumerate(ctrs)
        grad = MOI.get(model, DiffOpt.BackwardOutConstraint(), jc)
        @test JuMP.constant(grad) ≈ -dh[j] atol=ATOL rtol=RTOL
        for (i,iv) in enumerate(v)
            @test JuMP.coefficient(grad, iv) ≈ dG[j,i] atol=ATOL rtol=RTOL
        end
    end
end

@testset "Differentiating simple SOCP" begin
    # referred from _soc2test, https://github.com/jump-dev/MathOptInterface.jl/blob/master/src/Test/contconic.jl#L1355
    # find equivalent diffcp python program here: https://github.com/AKS1996/jump-gsoc-2020/blob/master/diffcp_socp_1_py.ipynb

    # Problem SOC2
    # min  x
    # s.t. y ≥ 1/√2
    #      x² + y² ≤ 1
    # in conic form:
    # min  x
    # s.t.  -1/√2 + y ∈ R₊
    #        1 - t ∈ {0}
    #      (t,x,y) ∈ SOC₃

    model = JuMP.direct_model(DiffOpt.diff_optimizer(SCS.Optimizer))
    MOI.set(model, MOI.Silent(), true)
    @variable(model, x)
    @variable(model, y)
    @variable(model, t)

    vv = [x,y,t]

    @objective(model, Min, x)
    @constraint(model, c1, [-1/√2 + y] in MOI.Nonnegatives(1))
    @constraint(model, c2, [1 - t] in MOI.Zeros(1))
    @constraint(model, c3, [1.0 * t, x, y] in SecondOrderCone())

    optimize!(model)

    v = JuMP.value.(vv)
    # slack variables
    s = collect(Iterators.flatten(JuMP.value.([c1, c2, c3])))
    y = collect(Iterators.flatten(JuMP.dual.([c1, c2, c3])))

    # these matrices are benchmarked with the output generated by diffcp
    # refer the python file mentioned above to get equivalent python source code
    @test v ≈ [-1/√2; 1/√2; 1.0] atol=ATOL rtol=RTOL
    @test s ≈ [0.0, 0.0, 1.0, -1/√2, 1/√2] atol=ATOL rtol=RTOL
    @test y ≈ [1, √2, √2, 1, -1] atol=ATOL rtol=RTOL

    dA = zeros(5, 3)
    dA[1:3, :] .= Matrix(1.0I, 3, 3)
    db = zeros(5)
    dc = zeros(3)

    MOI.set.(model, DiffOpt.BackwardInVariablePrimal(), vv, 1.0)

    @test_broken DiffOpt.backward(model)

    # TODO add tests

end


text:
# author: Ziyi (Francis) Yin
# This script benchmarks the time and memory usage for spectral conv layer in FNO
# The PDE is in 2D while FNO requires 3D FFT
# Date: Aug 2021

using DrWatson
@quickactivate "FNO"
import Pkg; Pkg.instantiate()

using PyPlot
using BSON
using Flux, Random, FFTW, Zygote, NNlib
using MAT, Statistics, LinearAlgebra
using CUDA
using ProgressMeter
using BenchmarkTools
using JLD2

try
    CUDA.device()
    global gpu_flag=true
catch e
    println("CUDA.device() found no GPU device on this machine.")
    global gpu_flag=false
end

gpu_flag = false

CUDA.culiteral_pow(::typeof(^), a::Complex{Float32}, b::Val{2}) = real(conj(a)*a)
CUDA.sqrt(a::Complex) = cu(sqrt(a))
Base.broadcasted(::typeof(sqrt), a::Base.Broadcast.Broadcasted) = Base.broadcast(sqrt, Base.materialize(a))

include("../utils.jl")

Random.seed!(3)

mutable struct SpectralConv3d_fast{T,N}
    weights1::AbstractArray{T,N}
    weights2::AbstractArray{T,N}
    weights3::AbstractArray{T,N}
    weights4::AbstractArray{T,N}
end

@Flux.functor SpectralConv3d_fast

# Constructor
function SpectralConv3d_fast(in_channels::Integer, out_channels::Integer, modes1::Integer, modes2::Integer, modes3::Integer)
    scale = (1f0 / (in_channels * out_channels))
    if gpu_flag
        weights1 = scale*randn(Complex{Float32}, modes1, modes2, modes3, in_channels, out_channels) |> gpu
        weights2 = scale*randn(Complex{Float32}, modes1, modes2, modes3, in_channels, out_channels) |> gpu
        weights3 = scale*randn(Complex{Float32}, modes1, modes2, modes3, in_channels, out_channels) |> gpu
        weights4 = scale*randn(Complex{Float32}, modes1, modes2, modes3, in_channels, out_channels) |> gpu
    else
        weights1 = scale*rand(Complex{Float32}, modes1, modes2, modes3, in_channels, out_channels)
        weights2 = scale*rand(Complex{Float32}, modes1, modes2, modes3, in_channels, out_channels)
        weights3 = scale*rand(Complex{Float32}, modes1, modes2, modes3, in_channels, out_channels)
        weights4 = scale*rand(Complex{Float32}, modes1, modes2, modes3, in_channels, out_channels)
    end
    return SpectralConv3d_fast{Complex{Float32}, 5}(weights1, weights2, weights3, weights4)
end

function compl_mul3d(x::AbstractArray{Complex{Float32}}, y::AbstractArray{Complex{Float32}})
    # complex multiplication
    # x in (modes1, modes2, modes3, input channels, batchsize)
    # y in (modes1, modes2, modes3, input channels, output channels)
    # output in (modes1,modes2,modes3,output channels,batchsize)
    x_per = permutedims(x,[5,4,1,2,3]) # batchsize*in_channels*modes1*modes2*modes3
    y_per = permutedims(y,[4,5,1,2,3]) # in_channels*out_channels*modes1*modes2*modes3
    x_resh = reshape(x_per,size(x_per,1),size(x_per,2),:) # batchsize*in_channels*(modes1*modes2*modes3)
    y_resh = reshape(y_per,size(y_per,1),size(y_per,2),:) # in_channels*out_channels*(modes1*modes2*modes3)
    out_resh = batched_mul(x_resh,y_resh) # batchsize*out_channels*(modes1*modes2*modes3)
    out_per = reshape(out_resh,size(out_resh,1),size(out_resh,2),size(x,1),size(x,2),size(x,3)) # batchsize*out_channels*modes1*modes2*modes3
    out = permutedims(out_per,[3,4,5,2,1])
    return out
end

function (L::SpectralConv3d_fast)(x::AbstractArray{Float32})
    # x in (size_x, size_y, time, channels, batchsize)
    x_ft = rfft(x,[1,2,3])      ## full size FFT
    modes1 = size(L.weights1,1)
    modes2 = size(L.weights1,2)
    modes3 = size(L.weights1,3)
    ### only keep low frequency coefficients
    out_ft = cat(cat(cat(compl_mul3d(x_ft[1:modes1, 1:modes2, 1:modes3, :,:], L.weights1),
                zeros(Complex{Float32},modes1,modes2,size(x_ft,3)-2*modes3,size(x_ft,4),size(x_ft,5)),
                compl_mul3d(x_ft[1:modes1, 1:modes2, end-modes3+1:end,:,:], L.weights2),dims=3),
                zeros(Complex{Float32},modes1,size(x_ft, 2)-2*modes2,size(x_ft,3),size(x_ft,4),size(x_ft,5)),
                cat(compl_mul3d(x_ft[1:modes1, end-modes2+1:end, 1:modes3,:,:], L.weights3),
                zeros(Complex{Float32},modes1,modes2,size(x_ft,3)-2*modes3,size(x_ft,4),size(x_ft,5)),
                compl_mul3d(x_ft[1:modes1, end-modes2+1:end, end-modes3+1:end,:,:], L.weights4),dims=3),
                dims=2),
                zeros(Complex{Float32},size(x_ft,1)-modes1,size(x_ft,2),size(x_ft,3),size(x_ft,4),size(x_ft,5)),dims=1)
    out_ft = irfft(out_ft, size(x,1),[1,2,3])
end

sizes = [2^i for i=5:7]
modes = [2^i for i=1:4]
width = [2^i for i=2:4]

batchsize = 1

time_ = zeros(Float32, length(sizes), length(modes), length(width))
memory_ = zeros(Float32, length(sizes), length(modes), length(width))

for (i, m)=enumerate(modes)
    for (j, w)=enumerate(width)
        global SC = SpectralConv3d_fast(w, w, m, m, m)
        Flux.trainmode!(SC, true)
        global weight_ = Flux.params(SC)
        for (k, s)=enumerate(sizes)
            global x = randn(Float32, s, s, s, w, batchsize)
            global y = randn(Float32, s, s, s, w, batchsize)
            if gpu_flag
                x = x |> gpu
                y = y |> gpu
            end
            B = (@benchmark grads = gradient(() -> Flux.mse(SC(x),y;agg=sum), weight_))
            time_[k, i, j] = mean(B).time/1f9
            memory_[k, i, j] = B.memory
            println("SpectralConv3D test: $(m) modes, $(w) width, ($(s),$(s),$(s)) problem size in (x,y,t), $(batchsize) batchsize")
            println("Time = $(time_[k, i, j]) seconds")
            println("Memory = $(memory_[k, i, j]) Bytes")
            Base.flush(Base.stdout)
        end
    end
end

device_ = (gpu_flag ? "GPU" : "CPU")
JLD2.@save "BenchmarkSpectralConv3D_$(device_).jld2" time_ memory_ gpu_flag sizes modes width

ext: jl


text:
# ---
# title: 59. Spiral Matrix II
# id: problem59
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Array
# link: <https://leetcode.com/problems/spiral-matrix-ii/description/>
# hidden: true
# ---
#
# Given a positive integer `n`, generate an `n x n` `matrix` filled with
# elements from `1` to `n2` in spiral order.
#
#
#
# **Example 1:**
#
# ![](https://assets.leetcode.com/uploads/2020/11/13/spiraln.jpg)
#
#
#
#     Input: n = 3
#     Output: [[1,2,3],[8,9,4],[7,6,5]]
#
#
# **Example 2:**
#
#
#
#     Input: n = 1
#     Output: [[1]]
#
#
#
#
# **Constraints:**
#
#   * `1 <= n <= 20`
#
#
## @lc code=start
using LeetCode

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


text:
#
# routines for handling functions as processes
#

"""
    loop(p::SimProcess, start::Channel, cycles::Number)

Put a `SimProcess` in a loop, which can be broken by a `SimException`.

# Arguments
- `p::SimProcess`:
- `start::Channel`:
- `cycles=Inf`:
"""
function loop(p::SimProcess, start::Channel, cycles::Number)
    take!(start)
    while cycles > 0
        try
            p.func(p.arg...; p.kw...)
        catch exc
            if isa(exc, SimException)
                exc.ev == Stop() ? break : nothing
            end
            rethrow(exc)
        end
        cycles -= 1
    end
    p.sim.processes = delete!(p.sim.processes, p.id)
end

"""
    startup!(p::SimProcess)

Start a `SimProcess` as a task in a loop.
"""
function startup!(p::SimProcess, cycles::Number)
    start = Channel(0)
    p.task = @async loop(p, start, cycles)
    p.state = Idle()
    put!(start, 1) # let the process start
end

"""
```
process!(sim::Clock, p::SimProcess, cycles=Inf)
process!(p::SimProcess, cycles=Inf)
```
Register a `SimProcess` to a clock, start it as an asynchronous process and
return the `id` it was registered with. It can then be found under `sim.processes[id]`.

# Arguments
- `sim::Clock`: clock, if no clock is given, it runs under 𝐶,
- `p::SimProcess`
- `cycles::Number=Inf`: number of cycles, the process should run.
"""
function process!(sim::Clock, p::SimProcess, cycles::Number=Inf)
    id = p.id
    while haskey(sim.processes, id)
        if isa(id, Float64)
            id = nextfloat(id)
        elseif isa(id, Int)
            id += 1
        elseif isa(id, String)
            s = split(id, "#")
            id = length(s) > 1 ? chop(id)*string(s[end][end]+1) : id*"#1"
        else
            throw(ArgumentError("process id $id is duplicate, cannot convert!"))
        end
    end
    sim.processes[id] = p
    p.id = id
    p.sim = sim
    startup!(p, cycles)
    id
end
process!(p::SimProcess, cycles=Inf) = process!(𝐶, p, cycles)

"stop a SimProcess"
function step!(p::SimProcess, ::Idle, ::Stop)
    schedule(p.task, SimException(Stop()))
    yield()
    p.state = Halted()
end

"""
```
delay!(sim::Clock, t::Number)
delay!(t::Number)
```
Delay a process for a time interval `t` on the clock `sim`. Suspend the calling
process until being reactivated by the clock at the appropriate time.

# Arguments
- `sim::Clock`: clock, if no clock is given, the delay goes to `𝐶`.
- `t::Number`: the time interval for the delay.
"""
function delay!(sim::Clock, t::Number)
    c = Channel(0)
    event!(sim, SF(put!, c, t), after, t)
    take!(c)
end
delay!(t::Number) = delay!(𝐶, t)

"""
```
wait!(sim::Clock, cond::Union{SimExpr, Array, Tuple}; scope::Module=Main)
wait!(cond::Union{SimExpr, Array, Tuple}; scope::Module=Main)
```
Wait on a clock for a condition to become true. Suspend the calling process
until the given condition is true.

# Arguments
- `sim::Clock`: clock, if no clock is given, the delay goes to `𝐶`.
- `cond::Union{SimExpr, Array, Tuple}`: a condition is an expression or SimFunction
    or an array or tuple of them. It is true if all expressions or SimFunctions
    therein return true.
- `scope::Module=Main`: evaluation scope for given expressions
"""
function wait!(sim::Clock, cond::Union{SimExpr, Array, Tuple}; scope::Module=Main)
    if all(simExec(sconvert(cond)))   # all conditions met
        return         # return immediately
    else
        c = Channel(0)
        event!(sim, SF(put!, c, 1), cond, scope=scope)
        take!(c)
    end
end
wait!(cond::Union{SimExpr, Array, Tuple}; scope::Module=Main) = wait!(𝐶, cond, scope=scope)

"""
    interrupt!(p::SimProcess, ev::SEvent, value=nothing)

Interrupt a `SimProcess` by throwing a `SimException` to it.
"""
function interrupt!(p::SimProcess, ev::SEvent, value=nothing)
    schedule(p.task, SimException(ev, value), error=true)
    yield()
end

"Stop a SimProcess"
stop!(p::SimProcess, value=nothing) = interrupt!(p, Stop(), value)

"""
```
now!(sim::Clock, ex::Union{SimExpr, Array, Tuple})
now!(ex::Union{SimExpr, Array, Tuple})
```
Lock the clock, execute the given expression, then unlock the clock again.

# Arguments
- `sim::Clock`:
- `ex::Union{SimExpr, Array, Tuple}`:
"""
function now!(sim::Clock, ex::Union{SimExpr, Array, Tuple})
    lock(sim.lock)
    simExec(sconvert(ex))
    unlock(sim.lock)
end
now!(ex::Union{SimExpr, Array, Tuple}) = now!(𝐶, ex)


text:
# Algorithm for Optimization

diff_forward(f,x;h=sqrt(eps(Float64))) = (f(x+h)-f(x))/h
diff_central(f,x;h=sqrt(eps(Float64))) = (f(x+h/2)-f(x-h/2))/h
diff_backward(f,x;h=sqrt(eps(Float64))) = (f(x)-f(x-h))/2
diff_complex(f,x;h=1e-20) = imag(f(x+h*im))/h

f=x->sin(x^2);
v = f(π/2+0.001im);
@show real(v) #f(x)
@show imag(v)/0.001 #f'(x)
@show sin(pi/2*pi/2), π*cos(π^2/4)

max_stars_repo_name: terasakisatoshi/juliaExer

max_issues_repo_name: terasakisatoshi/juliaExer

max_forks_repo_name: terasakisatoshi/juliaExer


text:
module model

export size_au
export AbstractParameters, ParametersStandard, ParametersNuker, convert_vector, convert_dict, registered_params
export Upsilon_nu, S_nu, temperature, DeltaV2, get_grids

# The double dot is because we are now inside the model module, and we want to import the
# constants module, which is part of the enclosing DiskJockey package.
using ..constants

# Define an abstract Parameters type, then subset of parameters that can be used to dispatch specifics

# In this file, we need some way of specifying the RT constants necessary
# E.g., frequency (or rel velocity), transition, isotope
# E.g., 0.1 km/s, J=2-1, 12CO

abstract type AbstractParameters end

"Parameters for the standard model."
mutable struct ParametersStandard <: AbstractParameters
    M_star::Float64 # [M_sun] stellar mass
    r_c::Float64 # [AU] characteristic radius
    T_10::Float64 # [K] temperature at 10 AU
    q::Float64 # temperature gradient exponent
    gamma::Float64 # surface density gradient exponent
    Sigma_c::Float64 # [g/cm^2] surface density at characteristic radius
    ksi::Float64 # [cm s^{-1}] microturbulence
    dpc::Float64 # [pc] distance to system
    incl::Float64 # [degrees] inclination 0 deg = face on, 90 = edge on.
    PA::Float64 # [degrees] position angle (East of North)
    vel::Float64 # [km/s] systemic velocity (positive is redshift/receeding)
    mu_RA::Float64 # [arcsec] central offset in RA
    mu_DEC::Float64 # [arcsec] central offset in DEC
end

"Parameters for the NUKER model."
mutable struct ParametersNuker <: AbstractParameters
    M_star::Float64 # [M_sun] stellar mass
    r_c::Float64 # [AU] characteristic radius
    T_10::Float64 # [K] temperature at 10 AU
    q::Float64 # temperature gradient exponent
    gamma::Float64 # surface density gradient within r_c (negative values yield holes)
    alpha::Float64 # sharpness of transition (2 = smooth, 16 = sharp)
    beta::Float64 # gradient power law outside r_c (~7)
    Sigma_c::Float64 # [g/cm^2] surface density at characteristic radius
    ksi::Float64 # [cm s^{-1}] microturbulence
    dpc::Float64 # [pc] distance to system
    incl::Float64 # [degrees] inclination 0 deg = face on, 90 = edge on.
    PA::Float64 # [degrees] position angle (East of North)
    vel::Float64 # [km/s] systemic velocity (positive is redshift/receeding)
    mu_RA::Float64 # [arcsec] central offset in RA
    mu_DEC::Float64 # [arcsec] central offset in DEC
end

"A dictionary of parameter lists for conversion."
registered_params = Dict([("standard", ["M_star", "r_c", "T_10", "q", "gamma", "Sigma_c", "ksi", "dpc", "incl", "PA", "vel", "mu_RA", "mu_DEC"]),
("nuker", ["M_star", "r_c", "T_10", "q", "gamma", "alpha", "beta", "Sigma_c", "ksi", "dpc", "incl", "PA", "vel", "mu_RA", "mu_DEC"])])

registered_types = Dict([("standard", ParametersStandard), ("nuker", ParametersNuker)])

"Unroll a vector of parameter values into a parameter type."
function convert_vector(p::Vector{Float64}, model::AbstractString, fix_params::Vector; args...)
    args = Dict{Symbol}{Float64}(args)

    # The goal is to assemble a length-nparam vector that can be unrolled into a parameter type
    # e.g., ParametersStandard(M_star, r_c, T_10, q, gamma, Sigma_c, ksi, dpc, incl, PA, vel, mu_RA, mu_DEC)

    # Select the registerd parameters corresponding to this model
    # These are the names listed in this file (model.jl)
    reg_params = registered_params[model]

    # fix_params is a list of strings from a config file that list the names of parameters to be fixed to the config
    # value. These names correspond to the registered names.

    # fit_params are the ones in reg_params that are not in (∉) fix_params
    fit_params = filter(x->∉(x,fix_params), reg_params)

    nparams = length(reg_params)

    # Make an empty vector of this same length
    par_vec = Array{Float64}(nparams)

    # This requires assigning p to fit_params
    # Find the indexes that correspond to fit_params
    par_indexes = findin(reg_params, fit_params)
    # Stuff p directly into these
    par_vec[par_indexes] = p

    # Then reading fix_params from args
    # First, create an array of fixed values analogous to p
    p_fixed = Float64[args[convert(Symbol, par)] for par in fix_params]
    par_indexes = findin(reg_params, fix_params)
    par_vec[par_indexes] = p_fixed

    # Now that we are sampling for log10M_gas for the verticalEta model, this part gets tricky.

    # Find the location of logSigma_c and make it Sigma_c
    # Even if we are using verticalEta, this will still be in here because it is in reg_params
    # Only though it currently corresponds to log10M_gas instead of log10Sigma_c
    indSigma_c = findin(reg_params, ["Sigma_c"])
    @assert length(indSigma_c) == 1 "Could not find Sigma_c in order to convert from logSigma_c or logM_gas."

    if model == "standard"
      # Convert from log10M_gas to Sigma_c
      M_gas = 10.^par_vec[indSigma_c] * M_sun # [g]

      # Find gamma and r_c
      r_c = par_vec[2] * AU # [cm]
      gamma = par_vec[5]

      Sigma_c = M_gas * (2 - gamma) / (2 * pi * r_c^2)
      par_vec[indSigma_c] = Sigma_c

    elseif model == "nuker"
      indalpha = findin(reg_params, ["alpha"])
      # println("Converting logalpha at index: ", indalpha)
      par_vec[indalpha] = 10.^par_vec[indalpha]
      par_vec[indSigma_c] = 10.^par_vec[indSigma_c]
    else
      par_vec[indSigma_c] = 10.^par_vec[indSigma_c]
    end
    # Then assembling these in the same orignial order as registered_params, into the parameter
    # type corresponding to the model.
    return registered_types[model](par_vec...)

end

"Used to turn a dictionary of parameter values (from config.yaml) directly into a parameter type. Generally used for synthesis and plotting command line scripts."
function convert_dict(p::Dict, model::AbstractString)
    # Select the registerd parameters corresponding to this model
    reg_params = registered_params[model]
    nparams = length(reg_params)

    if model == "standard"

      M_gas = 10.^p["logM_gas"] * M_sun # [g]

      # Find gamma and r_c
      r_c = p["r_c"] * AU # [cm]
      gamma = p["gamma"]

      Sigma_c = M_gas * (2 - gamma) / (2 * pi * r_c^2)
      p["Sigma_c"] = Sigma_c

    elseif model == "nuker"
      # add a new field, which is the conversion of logSigma_c to Sigma_c
      p["alpha"] = 10^p["logalpha"]
      p["Sigma_c"] = 10^p["logSigma_c"]
    else
      # add a new field, which is the conversion of logSigma_c to Sigma_c
      p["Sigma_c"] = 10^p["logSigma_c"]
    end

    # Using this order of parameters, unpack the dictionary p into a vector
    # reg_params reads Sigma_c, not logSigma_c; alpha, not logalpha
    par_vec = Float64[p[par_name] for par_name in reg_params]

    return registered_types[model](par_vec...)

end

# Assume all inputs to these functions are in CGS units and in *cylindrical* coordinates.
# Parametric type T allows passing individual Float64 or Vectors.
# # Alternate functions accept pars passed around, where pars is in M_star, AU, etc...
"The Keplerian velocity assuming non-zero thickness to the disk."
function velocity(r::Real, z::Real, M_star::Real)
    sqrt.(G * M_star / (r^2 + z^2)^(3./2)) * r
end
velocity{T}(r::T, z::T, pars::AbstractParameters) = velocity(r, z, pars.M_star * M_sun)

"Calculate temperature in cylindrical coordinates."
function temperature{T}(r::T, T_10::Real, q::Real)
    T_10 * (r ./ (10. * AU)).^(-q)
end
temperature{T}(r::T, pars::AbstractParameters) = temperature(r, pars.T_10, pars.q)

"Scale height, calculate in cylindrical coordinates."
function Hp{T}(r::T, M_star::Real, T_10::Real, q::Real)
    temp = temperature(r, T_10, q)
    sqrt.(kB * temp .* r.^3./(mu_gas * m_H * G * M_star))
end
Hp{T}(r::T,  pars::AbstractParameters) = Hp(r, pars.M_star * M_sun, pars.T_10, pars.q)

"Calculate the gas surface density using cylindrical coordinates."
function Sigma(r::Real, pars::AbstractParameters)
    r_c = pars.r_c * AU

    gamma = pars.gamma
    Sigma_c = pars.Sigma_c

    S = Sigma_c * (r/r_c)^(-gamma) * exp(-(r/r_c)^(2 - gamma))

    return S
end

"
    Sigma(r::Float64, pars::ParametersNuker)

Calculate the gas surface density using the Nuker profile."
function Sigma(r::Real, pars::ParametersNuker)
    r_c = pars.r_c * AU

    gamma = pars.gamma
    Sigma_c = pars.Sigma_c

    alpha = pars.alpha
    beta = pars.beta

    S = Sigma_c * (r/r_c)^(-gamma) * (1 + (r/r_c)^alpha)^((gamma - beta)/alpha)

end

# Delivers a gas density in g/cm^3
function rho(r::Real, z::Real, pars::AbstractParameters)
    H = Hp(r, pars)
    S = Sigma(r, pars)

    # Calculate the density
    rho = S/(sqrt(2. * pi) * H) * exp(-0.5 * (z/H)^2)

    return rho
end

# Ksi is microturbulent broadining width in units of km/s. Output of this function
# is in cm/s for RADMC (RADMC manual, eqn 7.12)
function microturbulence(ksi::Real)
    return ksi * 1.e5 # convert from km/s to cm/s
end

microturbulence(pars::AbstractParameters) = microturbulence(pars.ksi)

# Calculate the partition function for the temperature
# uses Mangum and Shirley expansion
# assumes B0 in Hz, depends on molecule
function Z_partition(T::Real, mol::Molecule)
    nugget = (h * mol.B0) / (kB * T)
    return 1/nugget + 1/3. + 1/15. * nugget + 4 * nugget^2 / 315 + nugget^3 / 315
end

# Calculate the source function
"Calculate the source function. nu in Hz."
function S_nu(r::Real, z::Real, nu::Real, pars::ParametersStandard)
    T = temperature(r, pars)
    return (2 * h * nu^3)/cc^2 / (exp(h * nu / (kB * T)) - 1)
end

# Calculate Upsilon
"Calculate Upsilon, the pre-factor."
function Upsilon_nu(r::Real, z::Real, nu::Real, pars::ParametersStandard, mol::Molecule)

    T = temperature(r, pars)

    # Many of the terms in Upsilon could be pre-calculated for the molecule.
    g_l = 2 * mol.l + 1

    n_l = mol.X_mol * rho(r, z, pars) * g_l * exp(-mol.T_L / T) / Z_partition(T, mol)

    # sigma_0
    sigma_0 = 16 * pi^3 / (h^2 * cc) * (kB * mol.T_L) * (mol.l + 1) / (mol.l * (2 * mol.l + 1)) * mol.mu^2

    # calculate Delta V
    DeltaV = sqrt(2 * kB * T / mol.mol_weight + microturbulence(pars)^2)

    return n_l * sigma_0 * cc / (DeltaV * mol.nu_0 * sqrt(pi)) * (1 - exp(- h * nu / (kB * T)))
end

"Calculate (Delta V)^2. Returns in cm/s."
function DeltaV2(r::Real, z::Real, pars::ParametersStandard, mol::Molecule)
    T = temperature(r, pars)
    return 2 * kB * T/mol.mol_weight + (pars.ksi * 1e5)^2
end

# Functions to compute grid objects for fast interpolation.
# Because we will always be querying for the same rcyl, z values,
# Return the three quantities at once.

function get_grids(pars::AbstractParameters, mol::Molecule, nu::Number, nr::Int, rmax::Number, nz::Int, zmax::Number)
    # Fill out 2D array spaced in zmax, rmax
    rs = linspace(0.0, rmax, nr)
    zs = linspace(0.0, zmax, nz)

    DeltaV2_grid = Array{Float64}(nz, nr)
    S_nu_grid = Array{Float64}(nz, nr)
    Upsilon_nu_grid = Array{Float64}(nz, nr)

    # Don't fill out the first r column, since this blows up to infty due to singularities in the temperature and surface
    # density profiles at the origin (rcyl = 0).
    # Simply copy the second and third values.
    for i=1:nz
        for j=1:10
            DeltaV2_grid[i,j] = DeltaV2(rs[10], zs[i], pars, mol)
            S_nu_grid[i,j] = S_nu(rs[10], zs[i], nu, pars)
            Upsilon_nu_grid[i,j] = Upsilon_nu(rs[10], zs[i], nu, pars, mol)
        end
        for j=11:nr
            DeltaV2_grid[i,j] = DeltaV2(rs[j], zs[i], pars, mol)
            S_nu_grid[i,j] = S_nu(rs[j], zs[i], nu, pars)
            Upsilon_nu_grid[i,j] = Upsilon_nu(rs[j], zs[i], nu, pars, mol)
        end
    end

    println("Max upsilon ", maximum(Upsilon_nu_grid) * AU, " /AU")

    # Closure which works on this grid.
    function interp(rcyl::Number, z::Number)
        # Look up what index rcyl is in rs
        ind_r = searchsortedfirst(rs, rcyl)

        # Look up what index z is in zs (assume -z = z)
        ind_z = searchsortedfirst(zs, abs(z))

        # If either of these errored, return a large value for Delta V, and 0.0 for S_nu and Upsilon
        if (ind_r > nr) | (ind_z > nz)
            return 1.0e5, 0.0, 0.0
        else
            return DeltaV2_grid[ind_z, ind_r], S_nu_grid[ind_z, ind_r], Upsilon_nu_grid[ind_z, ind_r]
        end

    end

    return interp
end

# Calculate a grid/interpolation object
function get_S_nu_grid(pars, mol, nu, nr, rmax, nz, zmax)

    # Fill out 2D array spaced in zmax, rmax
    rs = linspace(0.0, rmax, nr)
    zs = linspace(0.0, zmax, nz)

    S_nu_grid = Array{Float64}(nz, nr)

    # Don't fill out the first r column, since this blows up to infty due to temperature.
    # Simply copy the second value.
    for i=1:nz
        S_nu_grid[i,1] = S_nu(rs[2], zs[i], nu, pars)
        for j=2:nr
            S_nu_grid[i,j] = S_nu(rs[j], zs[i], nu, pars)
        end
    end

    # Closure which works on this grid.
    function interp(rcyl, z)
        # Look up what index rcyl is in rs
        ind_r = searchsortedfirst(rs, rcyl)

        # Look up what index z is in zs (assume -z = z)
        ind_z = searchsortedfirst(zs, abs(z))

        # If either of these errored, return 0.0
        if (ind_r > nr) | (ind_z > nz)
            return 0.0
        else
            return S_nu_grid[ind_z, ind_r]
        end

    end

    return interp
end

function get_Upsilon_nu_grid(pars, mol, nu, nr, rmax, nz, zmax)

    # Fill out 2D array spaced in zmax, rmax
    rs = linspace(0.0, rmax, nr)
    zs = linspace(0.0, zmax, nz)

    Upsilon_nu_grid = Array{Float64}(nz, nr)

    # Don't fill out the first r column, since this blows up to infty due to temperature.
    # Simply copy the second value.
    for i=1:nz
        Upsilon_nu_grid[i,1] = Upsilon_nu(rs[2], zs[i], nu, pars, mol)
        for j=2:nr
            Upsilon_nu_grid[i,j] = Upsilon_nu(rs[j], zs[i], nu, pars, mol)
        end
    end

    # Closure which works on this grid.
    function interp(rcyl, z)
        # Look up what index rcyl is in rs
        ind_r = searchsortedfirst(rs, rcyl)

        # Look up what index z is in zs (assume -z = z)
        ind_z = searchsortedfirst(zs, abs(z))

        # If either of these errored, return 0.0
        if (ind_r > nr) | (ind_z > nz)
            return 0.0
        else
            return Upsilon_nu_grid[ind_z, ind_r]
        end

    end

    return interp
end

function get_DeltaV2_grid(pars, mol, nr, rmax, nz, zmax)

    # Fill out 2D array spaced in zmax, rmax
    rs = linspace(0.0, rmax, nr)
    zs = linspace(0.0, zmax, nz)

    DeltaV2_grid = Array{Float64}(nz, nr)

    # Don't fill out the first r column, since this blows up to infty due to temperature.
    # Simply copy the second value.
    for i=1:nz
        DeltaV2_grid[i,1] = DeltaV2(rs[2], zs[i], pars, mol)
        for j=2:nr
            DeltaV2_grid[i,j] = DeltaV2(rs[j], zs[i], pars, mol)
        end
    end

    # Closure which works on this grid.
    function interp(rcyl, z)
        # Look up what index rcyl is in rs
        ind_r = searchsortedfirst(rs, rcyl)

        # Look up what index z is in zs (assume -z = z)
        ind_z = searchsortedfirst(zs, abs(z))

        # If either of these errored, return 0.0
        if (ind_r > nr) | (ind_z > nz)
            return 0.0
        else
            return DeltaV2_grid[ind_z, ind_r]
        end

    end

    return interp
end

# Various projection matrices which may or may not be correct.

function P_x(var)
    mat =  Float64[[1, 0, 0]  [0, cos(var), -sin(var)] [0, sin(var), cos(var)]]
    return mat
end

function P_z(var)
    mat =   Float64[[cos(var), -sin(var), 0] [sin(var), cos(var),   0] [0,        0,          1]]
    return mat
end

# function to transform spherical to cartesian
function P_project(theta, phi)
    mat = Float64[ [sin(theta) * cos(phi), sin(theta) * sin(phi), cos(theta)] [cos(theta) * cos(phi), cos(theta) * sin(phi), -sin(theta)] [-sin(phi), cos(phi), 0]]
    return mat
end

function convert_position(pars, r, theta, phi)
    x = r * sin(theta) * cos(phi)
    y = r * sin(theta) * sin(phi)
    z = r * cos(theta)

    sphere = Float64[r, theta, phi]
    cart = Float64[x, y, z]

    Px = P_x(pars.incl_inner * pi/180)
    Pz = P_z(pars.Omega * pi/180)

    cart_prime = Px.' * Pz.' * cart

    x_prime, y_prime, z_prime = cart_prime

    # now convert to spherical coordinates in the primed frame
    r_prime = r
    theta_prime = acos(z_prime / r)
    phi_prime = atan2(y_prime, x_prime)

    if phi_prime < 0
        phi_prime += 2pi
    end

    sphere_prime = Float64[r_prime, theta_prime, phi_prime]

    # Assert that the magnitudes of all of these vectors are the same
    @assert isapprox(norm(cart), r) "Cartesian norm doesn't match "
    @assert isapprox(norm(cart_prime), r) "Rotated cartesion norm doesn't match"

    return (cart, cart_prime, sphere_prime)
end

# # project it to the primed cartesian unit vectors
# vel_cart_prime = Pproj * vel_sphere_prime
#
# # now rotate this from the primed cartesian frame to the RADMC-3D cartesian frame
# vel_cart = Pz * Px * vel_cart_prime
#
# # now project it on to the spherical unit vectors in the RADMC-3D frame
# vel_sphere = Pproj.' * vel_cart

end # module


text:
el_ug = Vector{Int64}[[3, 6], [3, 5], [1, 2, 4, 5], [3], [2, 3], [1]]
el_dg = Vector{Int64}[[3, 6], [3, 5], [4, 5], [], [], []]

@testset "Test adjlist" begin
    ug = SimpleGraph(6)
    add_edge!(ug, 1, 3); add_edge!(ug, 2, 3); add_edge!(ug, 1, 6)
    add_edge!(ug, 2, 5); add_edge!(ug, 3, 4); add_edge!(ug, 3, 5)
    @test adjlist(ug) == el_ug

    dg = SimpleDiGraph(6)
    add_edge!(dg, 1, 3); add_edge!(dg, 2, 3); add_edge!(dg, 1, 6)
    add_edge!(dg, 2, 5); add_edge!(dg, 3, 4); add_edge!(dg, 3, 5)
    @test adjlist(dg) == el_dg

    ug = SimpleWeightedGraph(6)
    add_edge!(ug, 1, 3, 2); add_edge!(ug, 2, 3, 2); add_edge!(ug, 1, 6, 1)
    add_edge!(ug, 2, 5, -2); add_edge!(ug, 3, 4, -2); add_edge!(ug, 3, 5, -1)
    @test adjlist(ug) == el_ug

    dg = SimpleWeightedDiGraph(6)
    add_edge!(dg, 1, 3, 2); add_edge!(dg, 2, 3, 2); add_edge!(dg, 1, 6, 1)
    add_edge!(dg, 2, 5, -2); add_edge!(dg, 3, 4, -2); add_edge!(dg, 3, 5, -1)
    @test adjlist(dg) == el_dg
end


text:
export BrickletIO16PortConfiguration
struct BrickletIO16PortConfiguration
    direction_mask::Integer
    value_mask::Integer
end

export BrickletIO16PortMonoflop
struct BrickletIO16PortMonoflop
    value::Integer
    time::Integer
    time_remaining::Integer
end

export BrickletIO16EdgeCountConfig
struct BrickletIO16EdgeCountConfig
    edge_type::Integer
    debounce::Integer
end

export BrickletIO16Identity
struct BrickletIO16Identity
    uid::String
    connected_uid::String
    position::Char
    hardware_version::Vector{Integer}
    firmware_version::Vector{Integer}
    device_identifier::Integer
end

export BrickletIO16
"""
16-channel digital input/output
"""
mutable struct BrickletIO16 <: TinkerforgeDevice
    ipcon::IPConnection
    deviceInternal::PyObject

    """
    Creates an object with the unique device ID *uid* and adds it to
    the IP Connection *ipcon*.
    """
    function BrickletIO16(uid::String, ipcon::IPConnection)
        package = pyimport("tinkerforge.bricklet_io16")
        deviceInternal = package.BrickletIO16(uid, ipcon.ipconInternal)

        return new(ipcon, deviceInternal)
    end
end

export set_port
"""
    $(SIGNATURES)

Sets the output value (high or low) for a port ("a" or "b") with a bitmask
(8bit). A 1 in the bitmask means high and a 0 in the bitmask means low.

For example: The value 15 or 0b00001111 will turn the pins 0-3 high and the
pins 4-7 low for the specified port.

All running monoflop timers of the given port will be aborted if this function
is called.

.. note::
 This function does nothing for pins that are configured as input.
 Pull-up resistors can be switched on with :func:`Set Port Configuration`.
"""
function set_port(device::BrickletIO16, port, value_mask)
    device.deviceInternal.set_port(port, value_mask)
end

export get_port
"""
    $(SIGNATURES)

Returns a bitmask of the values that are currently measured on the
specified port. This function works if the pin is configured to input
as well as if it is configured to output.
"""
function get_port(device::BrickletIO16, port)
    return device.deviceInternal.get_port(port)
end

export set_port_configuration
"""
    $(SIGNATURES)

Configures the value and direction of a specified port. Possible directions
are 'i' and 'o' for input and output.

If the direction is configured as output, the value is either high or low
(set as *true* or *false*).

If the direction is configured as input, the value is either pull-up or
default (set as *true* or *false*).

For example:

* ('a', 255, 'i', true) or ('a', 0b11111111, 'i', true) will set all pins of port A as input pull-up.
* ('a', 128, 'i', false) or ('a', 0b10000000, 'i', false) will set pin 7 of port A as input default (floating if nothing is connected).
* ('b', 3, 'o', false) or ('b', 0b00000011, 'o', false) will set pins 0 and 1 of port B as output low.
* ('b', 4, 'o', true) or ('b', 0b00000100, 'o', true) will set pin 2 of port B as output high.

Running monoflop timers for the selected pins will be aborted if this
function is called.
"""
function set_port_configuration(device::BrickletIO16, port, selection_mask, direction, value)
    device.deviceInternal.set_port_configuration(port, selection_mask, direction, value)
end

export get_port_configuration
"""
    $(SIGNATURES)

Returns a direction bitmask and a value bitmask for the specified port. A 1 in
the direction bitmask means input and a 0 in the bitmask means output.

For example: A return value of (15, 51) or (0b00001111, 0b00110011) for
direction and value means that:

* pins 0 and 1 are configured as input pull-up,
* pins 2 and 3 are configured as input default,
* pins 4 and 5 are configured as output high
* and pins 6 and 7 are configured as output low.
"""
function get_port_configuration(device::BrickletIO16, port)
    return device.deviceInternal.get_port_configuration(port)
end

export set_debounce_period
"""
    $(SIGNATURES)

Sets the debounce period of the :cb:`Interrupt` callback.

For example: If you set this value to 100, you will get the interrupt
maximal every 100ms. This is necessary if something that bounces is
connected to the IO-16 Bricklet, such as a button.
"""
function set_debounce_period(device::BrickletIO16, debounce)
    device.deviceInternal.set_debounce_period(debounce)
end

export get_debounce_period
"""
    $(SIGNATURES)

Returns the debounce period as set by :func:`Set Debounce Period`.
"""
function get_debounce_period(device::BrickletIO16)
    return device.deviceInternal.get_debounce_period()
end

export set_port_interrupt
"""
    $(SIGNATURES)

Sets the pins on which an interrupt is activated with a bitmask.
Interrupts are triggered on changes of the voltage level of the pin,
i.e. changes from high to low and low to high.

For example: ('a', 129) or ('a', 0b10000001) will enable the interrupt for
pins 0 and 7 of port a.

The interrupt is delivered with the :cb:`Interrupt` callback.
"""
function set_port_interrupt(device::BrickletIO16, port, interrupt_mask)
    device.deviceInternal.set_port_interrupt(port, interrupt_mask)
end

export get_port_interrupt
"""
    $(SIGNATURES)

Returns the interrupt bitmask for the specified port as set by
:func:`Set Port Interrupt`.
"""
function get_port_interrupt(device::BrickletIO16, port)
    return device.deviceInternal.get_port_interrupt(port)
end

export set_port_monoflop
"""
    $(SIGNATURES)

Configures a monoflop of the pins specified by the second parameter as 8 bit
long bitmask. The specified pins must be configured for output. Non-output
pins will be ignored.

The third parameter is a bitmask with the desired value of the specified
output pins. A 1 in the bitmask means high and a 0 in the bitmask means low.

The forth parameter indicates the time that the pins should hold
the value.

If this function is called with the parameters ('a', 9, 1, 1500) or
('a', 0b00001001, 0b00000001, 1500): Pin 0 will get high and pin 3 will get
low on port 'a'. In 1.5s pin 0 will get low and pin 3 will get high again.

A monoflop can be used as a fail-safe mechanism. For example: Lets assume you
have a RS485 bus and an IO-16 Bricklet connected to one of the slave
stacks. You can now call this function every second, with a time parameter
of two seconds and pin 0 set to high. Pin 0 will be high all the time. If now
the RS485 connection is lost, then pin 0 will get low in at most two seconds.
"""
function set_port_monoflop(device::BrickletIO16, port, selection_mask, value_mask, time)
    device.deviceInternal.set_port_monoflop(port, selection_mask, value_mask, time)
end

export get_port_monoflop
"""
    $(SIGNATURES)

Returns (for the given pin) the current value and the time as set by
:func:`Set Port Monoflop` as well as the remaining time until the value flips.

If the timer is not running currently, the remaining time will be returned
as 0.
"""
function get_port_monoflop(device::BrickletIO16, port, pin)
    return device.deviceInternal.get_port_monoflop(port, pin)
end

export set_selected_values
"""
    $(SIGNATURES)

Sets the output value (high or low) for a port ("a" or "b" with a bitmask,
according to the selection mask. The bitmask is 8 bit long and a 1 in the
bitmask means high and a 0 in the bitmask means low.

For example: The parameters ('a', 192, 128) or ('a', 0b11000000, 0b10000000)
will turn pin 7 high and pin 6 low on port A, pins 0-6 will remain untouched.

Running monoflop timers for the selected pins will be aborted if this
function is called.

.. note::
 This function does nothing for pins that are configured as input.
 Pull-up resistors can be switched on with :func:`Set Port Configuration`.
"""
function set_selected_values(device::BrickletIO16, port, selection_mask, value_mask)
    device.deviceInternal.set_selected_values(port, selection_mask, value_mask)
end

export get_edge_count
"""
    $(SIGNATURES)

Returns the current value of the edge counter for the selected pin on port A.
You can configure the edges that are counted with :func:`Set Edge Count Config`.

If you set the reset counter to *true*, the count is set back to 0
directly after it is read.

.. versionadded:: 2.0.3\$nbsp;(Plugin)
"""
function get_edge_count(device::BrickletIO16, pin, reset_counter)
    return device.deviceInternal.get_edge_count(pin, reset_counter)
end

export set_edge_count_config
"""
    $(SIGNATURES)

Configures the edge counter for the selected pin of port A. Pins 0 and 1
are available for edge counting.

The edge type parameter configures if rising edges, falling edges or
both are counted if the pin is configured for input. Possible edge types are:

* 0 = rising
* 1 = falling
* 2 = both

Configuring an edge counter resets its value to 0.

If you don't know what any of this means, just leave it at default. The
default configuration is very likely OK for you.

.. versionadded:: 2.0.3\$nbsp;(Plugin)
"""
function set_edge_count_config(device::BrickletIO16, pin, edge_type, debounce)
    device.deviceInternal.set_edge_count_config(pin, edge_type, debounce)
end

export get_edge_count_config
"""
    $(SIGNATURES)

Returns the edge type and debounce time for the selected pin of port A as set by
:func:`Set Edge Count Config`.

.. versionadded:: 2.0.3\$nbsp;(Plugin)
"""
function get_edge_count_config(device::BrickletIO16, pin)
    return device.deviceInternal.get_edge_count_config(pin)
end

export get_identity
"""
    $(SIGNATURES)

Returns the UID, the UID where the Bricklet is connected to,
the position, the hardware and firmware version as well as the
device identifier.

The position can be 'a', 'b', 'c', 'd', 'e', 'f', 'g' or 'h' (Bricklet Port).
A Bricklet connected to an :ref:`Isolator Bricklet <isolator_bricklet>` is always at
position 'z'.

The device identifier numbers can be found :ref:`here <device_identifier>`.
|device_identifier_constant|
"""
function get_identity(device::BrickletIO16)
    return device.deviceInternal.get_identity()
end

export register_callback
"""
Registers the given *function* with the given *callback_id*.
"""
function register_callback(device::BrickletIO16, callback_id, function_)
    if isnothing(function_)
        device.registered_callbacks.pop(callback_id, None)
    else
        device.registered_callbacks[callback_id] = function_
    end
end


text:
mutable struct SwitchSimulateState{T}
    score::Float64
    noise::Float64
    index::Int
    subtrace::Trace
    retval::T
    SwitchSimulateState{T}(score::Float64, noise::Float64) where T = new{T}(score, noise)
end

function process!(gen_fn::Switch{C, N, K, T},
                  index::Int,
                  args::Tuple,
                  state::SwitchSimulateState{T}) where {C, N, K, T}
    local retval::T
    subtrace = simulate(getindex(gen_fn.branches, index), args)
    state.index = index
    state.noise += project(subtrace, EmptySelection())
    state.subtrace = subtrace
    state.score += get_score(subtrace)
    state.retval = get_retval(subtrace)
end

@inline process!(gen_fn::Switch{C, N, K, T}, index::C, args::Tuple, state::SwitchSimulateState{T}) where {C, N, K, T} = process!(gen_fn, getindex(gen_fn.cases, index), args, state)

function simulate(gen_fn::Switch{C, N, K, T},
                  args::Tuple) where {C, N, K, T}

    index = args[1]
    state = SwitchSimulateState{T}(0.0, 0.0)
    process!(gen_fn, index, args[2 : end], state)
    return SwitchTrace(gen_fn, state.subtrace,
                       state.retval, args,
                       state.score, state.noise)
end


text:
const AbstractVec = SVector{3, T} where T <: AbstractFloat
const Vec = AbstractVec{Float64}

function normalize(v::AbstractVec)::AbstractVec
    v / norm(v)
end

@inline function x(v::AbstractVec)::AbstractFloat
    v[1]
end

@inline function y(v::AbstractVec)::AbstractFloat
    v[2]
end

@inline function z(v::AbstractVec)::AbstractFloat
    v[3]
end

function random_on_unit_disk()::AbstractVec
    v::AbstractVec = Vec(randn(), randn(), 0.0)
    normalize(v)
end

function random_in_unit_disk()::AbstractVec
    v::AbstractVec = random_on_unit_disk()
    r::AbstractFloat = rand()
    sqrt(r) * v
end

function random_on_unit_sphere()::AbstractVec
    v::AbstractVec = Vec(randn(), randn(), randn())
    normalize(v)
end

function random_in_unit_sphere()::AbstractVec
    v::AbstractVec = random_on_unit_disk()
    r::AbstractFloat = rand()
    cbrt(r) * v
end

function reflect(v::AbstractVec, n::AbstractVec)::AbstractVec
    v::AbstractVec = normalize(v)
    v - 2dot(v, n)n
end

function refract(v::AbstractVec, n::AbstractVec, index::AbstractFloat)::Tuple{Bool, Union{Nothing, AbstractVec}}
    v::AbstractVec = normalize(v)
    dt::AbstractFloat = dot(v, n)
    discriminant::AbstractFloat = 1.0 - index^2 * (1.0 - dt^2)
    if discriminant > 0
        refracted::AbstractVec = index * (v - n * dt) - n * sqrt(discriminant)
        true, refracted
    end
    false, nothing
end

max_stars_repo_name: goooxu/JuliaRayTracing

max_issues_repo_name: goooxu/JuliaRayTracing

max_forks_repo_name: goooxu/JuliaRayTracing


text:
using Gadfly

set_default_plot_size(14cm, 8cm)

# Issue #1344

datas = Gadfly.Data(color= ["D", "A", "C", "D", "A", "C", "D", "D", "A", "B"],
    shape=["E", "F", "F", "E", "G", "E", "E", "G", "H"], size=["I","I","K","J","L","L","I","J","K","L"])
scales = [Scale.color_discrete(levels=["A","B","C","D"]), Scale.shape_discrete(levels=["E","F","G","H"]),
    Scale.size_discrete2(levels=["I","J","K","L"])]

aes = Scale.apply_scales(scales, datas)
theme1 = Theme(key_max_columns=4, point_size=5pt, key_swatch_shape=Shape.circle)
guides = [render(g, theme1, aes[1])[1].ctxs  for g in (Guide.colorkey(), Guide.shapekey(), Guide.sizekey())]

gridstack([guides[i][j] for i in 1:3, j in 1:4])


text:
@testset "LPIP Interface" begin
    @testset "Production problem" begin
        A = [
            2. 1
            1 2
        ]
        b = [4.;4]
        c = [-4.; -3]
        d = 3.

        lpip_pb = LPIPLinearProblem{Float64}(A, b, c, d, 0)
        params = LPIP.Params()
        lp = LPIP.interior_points(lpip_pb, params)

        @test lp.obj_val ≈ -6.3333 atol = 1e-4 rtol = 1e-4
        @test lp.dual_obj_val ≈ -6.3333 atol = 1e-4 rtol = 1e-4
        @test lp.variables.x ≈ [1.3333; 1.3333] atol = 1e-4 rtol = 1e-4
        @test lp.variables.p ≈ [-1.6666; -0.6666] atol = 1e-4 rtol = 1e-4
        @test lp.variables.s ≈ [0.0; 0.0] atol = 1e-4 rtol = 1e-4
        @test lp.status == 1
    end
end


text:
"""
	tsls(exogenous::AbstractMatrix, instruments::AbstractMatrix,
		endogenous::AbstractMatrix)

	This functions creates the model matrix for 2SLS using a matrix of exogenous
	variables, additional instruments, and the endogenous variables. It verifies
	that the equation is identifiable (enough instruments). It returns the
	linearly independent version of the model matrix, the Gram matrix of it,
	and an indicator of which variables were linearly independent (which were
	not dropped).
"""
function tsls(exogenous::AbstractMatrix, instruments::AbstractMatrix, endogenous::AbstractMatrix)
	n_exogenous = size(exogenous, 2)
	n_endogenous = size(endogenous, 2)
	instrumented = hcat(exogenous, instruments)
	instrumented, LI = linearindependent(instrumented)
	@assert size(instrumented, 2) ≥ (n_exogenous + n_endogenous) "Not sufficient instruments."
	mm = hcat(exogenous, instrumented * inv(cholfact!(Hermitian(instrumented.'instrumented))) * instrumented.'endogenous)
	mm, LI = linearindependent(mm)
	return mm, Symmetric(mm.'mm), LI
end


text:
export ChebyshevT

"""
    ChebyshevT{<:Number}(coeffs::AbstractVector, var=:x)

Chebyshev polynomial of the first kind.

Construct a polynomial from its coefficients `a`, lowest order first, optionally in
terms of the given variable `x`. `x` can be a character, symbol, or string.

# Examples

```jldoctest ChebyshevT
julia> using Polynomials

julia> ChebyshevT([1, 0, 3, 4])
ChebyshevT(1⋅T_0(x) + 3⋅T_2(x) + 4⋅T_3(x))

julia> ChebyshevT([1, 2, 3, 0], :s)
ChebyshevT(1⋅T_0(s) + 2⋅T_1(s) + 3⋅T_2(s))

julia> one(ChebyshevT)
ChebyshevT(1.0⋅T_0(x))
```
"""
struct ChebyshevT{T <: Number} <: AbstractPolynomial{T}
    coeffs::Vector{T}
    var::Symbol
    function ChebyshevT{T}(coeffs::AbstractVector{T}, var::Symbol) where {T <: Number}
        length(coeffs) == 0 && return new{T}(zeros(T, 1), var)
        last_nz = findlast(!iszero, coeffs)
        last = max(1, last_nz === nothing ? 0 : last_nz)
        return new{T}(coeffs[1:last], var)
    end
end

@register ChebyshevT

function Base.convert(P::Type{<:Polynomial}, ch::ChebyshevT)
    if length(ch) < 3
        return P(ch.coeffs, ch.var)
    end
    c0 = P(ch[end - 1], ch.var)
    c1 = P(ch[end], ch.var)
    @inbounds for i in degree(ch):-1:2
        tmp = c0
        c0 = P(ch[i - 2], ch.var) - c1
        c1 = tmp + c1 * variable(P) * 2
    end
    return c0 + c1 * variable(P)
end

function Base.convert(C::Type{<:ChebyshevT}, p::Polynomial)
    res = zero(C)
    @inbounds for i in degree(p):-1:0
        res = variable(C) * res + p[i]
    end
    return res
end

domain(::Type{<:ChebyshevT}) = Interval(-1, 1)
variable(P::Type{<:ChebyshevT}, var::SymbolLike=:x ) = P([0,1], var)
"""
    (::ChebyshevT)(x)

Evaluate the Chebyshev polynomial at `x`. If `x` is outside of the domain of [-1, 1], an error will be thrown. The evaluation uses Clenshaw Recursion.

# Examples
```jldoctest ChebyshevT
julia> using Polynomials

julia> c = ChebyshevT([2.5, 1.5, 1.0])
ChebyshevT(2.5⋅T_0(x) + 1.5⋅T_1(x) + 1.0⋅T_2(x))

julia> c(0)
1.5

julia> c.(-1:0.5:1)
5-element Array{Float64,1}:
 2.0
 1.25
 1.5
 2.75
 5.0
```
"""
function (ch::ChebyshevT{T})(x::S) where {T,S}
    x ∉ domain(ch) && error("$x outside of domain")
    R = promote_type(T, S)
    length(ch) == 0 && return zero(R)
    length(ch) == 1 && return R(ch[0])
    c0 = ch[end - 1]
    c1 = ch[end]
    @inbounds for i in lastindex(ch) - 2:-1:0
        c0, c1 = ch[i] - c1, c0 + c1 * 2x
    end
    return R(c0 + c1 * x)
end

function vander(P::Type{<:ChebyshevT}, x::AbstractVector{T}, n::Integer) where {T <: Number}
    A = Matrix{T}(undef, length(x), n + 1)
    A[:, 1] .= one(T)
    if n > 0
        A[:, 2] .= x
        @inbounds for i in 3:n + 1
            A[:, i] .= A[:, i - 1] .* 2x .- A[:, i - 2]
        end
    end
    return A
end

function integrate(p::ChebyshevT{T}, C::S) where {T,S <: Number}
    R = promote_type(eltype(one(T) / 1), S)
    if hasnan(p) || isnan(C)
        return ChebyshevT([NaN])
    end
    n = length(p)
    if n == 1
        return ChebyshevT{R}([C, p[0]])
    end
    a2 = Vector{R}(undef, n + 1)
    a2[1] = zero(R)
    a2[2] = p[0]
    a2[3] = p[1] / 4
    @inbounds for i in 2:n - 1
        a2[i + 2] = p[i] / (2 * (i + 1))
        a2[i] -= p[i] / (2 * (i - 1))
    end
    a2[1] += R(C) - ChebyshevT(a2)(0)
    return ChebyshevT(a2, p.var)
end

function derivative(p::ChebyshevT{T}, order::Integer = 1) where {T}
    order < 0 && throw(ArgumentError("Order of derivative must be non-negative"))
    R  = eltype(one(T)/1)
    order == 0 && return convert(ChebyshevT{R}, p)
    hasnan(p) && return ChebyshevT(R[NaN], p.var)
    order > length(p) && return zero(ChebyshevT{R})

    q =  convert(ChebyshevT{R}, copy(p))
    n = length(p)
    der = Vector{R}(undef, n)

    for j in n:-1:2
        der[j] = 2j * q[j]
        q[j - 2] += j * q[j] / (j - 2)
    end
    if n > 1
        der[2] = 4q[2]
    end
    der[1] = q[1]

    pp = ChebyshevT(der, p.var)
    return order > 1 ?  derivative(pp, order - 1) : pp

end

function companion(p::ChebyshevT{T}) where T
    d = length(p) - 1
    d < 1 && error("Series must have degree greater than 1")
    d == 1 && return diagm(0 => [-p[0] / p[1]])
    R = eltype(one(T) / one(T))

    scl = vcat(1.0, fill(R(√0.5), d - 1))

    diag = vcat(√0.5, fill(R(0.5), d - 2))
    comp = diagm(1 => diag,
                 -1 => diag)
    monics = p.coeffs ./ p.coeffs[end]
    comp[:, end] .-= monics[1:d] .* scl ./ scl[end] ./ 2
    return R.(comp)
end

function Base.:+(p1::ChebyshevT, p2::ChebyshevT)
    p1.var != p2.var && error("Polynomials must have same variable")
    n = max(length(p1), length(p2))
    c = [p1[i] + p2[i] for i = 0:n]
    return ChebyshevT(c, p1.var)
end

function Base.:*(p1::ChebyshevT{T}, p2::ChebyshevT{S}) where {T,S}
    p1.var != p2.var && error("Polynomials must have same variable")
    z1 = _c_to_z(p1.coeffs)
    z2 = _c_to_z(p2.coeffs)
    prod = fastconv(z1, z2)
    ret = ChebyshevT(_z_to_c(prod), p1.var)
    return truncate!(ret)
end

function Base.divrem(num::ChebyshevT{T}, den::ChebyshevT{S}) where {T,S}
    num.var != den.var && error("Polynomials must have same variable")
    n = length(num) - 1
    m = length(den) - 1

    R = typeof(one(T) / one(S))
    P = ChebyshevT{R}

    if n < m
        return zero(P), convert(P, num)
    elseif m == 0
        den[0] ≈ 0 && throw(DivideError())
        return num ./ den[end], zero(P)
    end

    znum = _c_to_z(num.coeffs)
    zden = _c_to_z(den.coeffs)
    quo, rem = _z_division(znum, zden)
    q_coeff = _z_to_c(quo)
    r_coeff = _z_to_c(rem)
    return P(q_coeff, num.var), P(r_coeff, num.var)
end

function showterm(io::IO, ::Type{ChebyshevT{T}}, pj::T, var, j, first::Bool, mimetype) where {N, T}
    iszero(pj) && return false
    !first &&  print(io, " ")
    print(io, hasneg(T)  && isneg(pj) ? "- " :  (!first ? "+ " : ""))
    print(io, "$(abs(pj))⋅T_$j($var)")
    return true
end

#=
zseries =#

function _c_to_z(cs::AbstractVector{T}) where {T}
    n = length(cs)
    U = typeof(one(T) / 2)
    zs = zeros(U, 2n - 1)
    zs[n:end] = cs ./ 2
    return zs .+ reverse(zs)
end

function _z_to_c(z::AbstractVector{T}) where {T}
    n = (length(z) + 1) ÷ 2
    cs = z[n:end]
    cs[2:n] *= 2
    return cs
end

function _z_division(z1::AbstractVector{T}, z2::AbstractVector{S}) where {T,S}
    R = eltype(one(T) / one(S))
    length(z1)
    length(z2)
    if length(z2) == 1
        z1 ./= z2
        return z1, zero(R)
    elseif length(z1) < length(z2)
        return zero(R), R.(z1)
    end
    dlen = length(z1) - length(z2)
    scl = z2[1]
    z2 ./= scl
    quo = Vector{R}(undef, dlen + 1)
    i = 1
    j = dlen + 1
    while i < j
        r = z1[i]
        quo[i] = z1[i]
        quo[end - i + 1] = r
        tmp = r .* z2
        z1[i:i + length(z2) - 1] .-= tmp
        z1[j:j + length(z2) - 1] .-= tmp
        i += 1
        j -= 1
    end

    r = z1[i]
    quo[i] = r
    tmp = r * z2
    z1[i:i + length(z2) - 1] .-= tmp
    quo ./= scl
    rem = z1[i + 1:i - 2 + length(z2)]
    return quo, rem
end


text:
using Test;

using AdalmPluto;

@testset "libIIO/context.jl" begin
    # disable errors on null pointers
    toggleNoAssertions(true);

    # Need iiod running for those
    # C_iio_create_default_context
    @test_skip C_iio_create_default_context() != C_NULL;
    # C_iio_create_local_context
    @test_skip C_iio_create_local_context() != C_NULL;
    # C_iio_create_network_context
    @test_skip C_iio_create_network_context("127.0.0.1") != C_NULL;

    try
        global uri = AdalmPluto.scan("usb", 1, false);
    catch
        @error "Could not get a valid uri, skipping the remaining tests"; return;
    end

    # C_iio_create_context_from_uri
    @test (global context = C_iio_create_context_from_uri(uri)) != C_NULL;

    if context == C_NULL
        @error "Could not create a context, skipping the remaining tests"; return;
    end

    @testset "Context description" begin
        # C_iio_context_get_name
        @test C_iio_context_get_name(context) != "";
        # C_iio_context_get_description
        @test C_iio_context_get_description(context) != "";
        # C_iio_context_get_version
        ctx =  C_iio_context_get_version(context)
        @test ctx[3] > 17
        # @test (C_iio_context_get_version(context) == (0, 0, 17, "v0.17  ") ||    C_iio_context_get_version(context) == (0, 0, 18, "v0.18  "));
    end

    @testset "Context devices" begin
        # C_iio_context_get_devices_count
        @test (global nb_devices = C_iio_context_get_devices_count(context)) > 0;

        if nb_devices < 1
            @warn "Could not find devices in the context, skipping related tests";
        else
            # C_iio_context_find_device
            @test C_iio_context_find_device(context, "dummy") == C_NULL;
            @test C_iio_context_find_device(context, "cf-ad9361-lpc") != C_NULL;

            # C_iio_context_get_device
            @test C_iio_context_get_device(context, UInt32(666)) == C_NULL;
            @test C_iio_context_get_device(context, UInt32(0)) != C_NULL;
        end

    end

    @testset "Context attrs" begin
        # C_iio_context_get_attrs_count
        @test (global nb_attrs = C_iio_context_get_attrs_count(context)) > 0

        if nb_attrs < 1
            @warn "Could not find context specific attributes, skipping related tests";
        else
            # NOTE : Specific the radio used for the test, this test failing can be normal.
            # C_iio_context_get_attr
            flag1 = C_iio_context_get_attr(context, UInt32(0)) ==  (0, "hw_model", "Analog Devices PlutoSDR Rev.B (Z7010-AD9364)")
            flag2 = C_iio_context_get_attr(context, UInt32(0)) ==  (0, "hw_model", "Analog Devices PlutoSDR Rev.B (Z7010-AD9363A)")
            @test flag1 || flag2
            # C_iio_context_get_attr_value

            flag1 =  C_iio_context_get_attr_value(context, "hw_model") ==   (0, "hw_model", "Analog Devices PlutoSDR Rev.B (Z7010-AD9364)")
            flag3 =  C_iio_context_get_attr_value(context, "hw_model") ==   ("Analog Devices PlutoSDR Rev.B (Z7010-AD9364)")
            @test flag1 || flag2 || flag3
        end

    end

    # C_iio_context_set_timeout
    @test C_iio_context_set_timeout(context, UInt32(0)) == 0;

    @testset "Context xml & clone" begin
        # C_iio_context_get_xml
        @test (global xml = C_iio_context_get_xml(context)) != "";
        # C_iio_context_destroy
        @test C_iio_context_destroy(context) === nothing;
        # C_iio_create_xml_context
        # C_iio_create_xml_context_mem
        @test (global context = C_iio_create_xml_context_mem(xml, UInt64(length(xml)))) != C_NULL;
        # C_iio_context_clone
        @test (global clone = C_iio_context_clone(context)) != C_NULL;
        C_iio_context_destroy(clone);
    end

    C_iio_context_destroy(context);
end


text:
using Turing
using LinearAlgebra

# Packages for visualization
using VegaLite, DataFrames, StatsPlots

# Import Fisher's iris example data set
using RDatasets

# Set a seed for reproducibility.
using Random
Random.seed!(1789);

n_cells = 60
n_genes = 9
mu_1 = 10.0 * ones(n_genes ÷ 3)
mu_0 = zeros(n_genes ÷ 3)
S = I(n_genes ÷ 3)
mvn_0 = MvNormal(mu_0, S)
mvn_1 = MvNormal(mu_1, S)

# create a diagonal block like expression matrix, with some non-informative genes;
# not all features/genes are informative, some might just not differ very much between cells)
expression_matrix = transpose(
    vcat(
        hcat(rand(mvn_1, n_cells ÷ 2), rand(mvn_0, n_cells ÷ 2)),
        hcat(rand(mvn_0, n_cells ÷ 2), rand(mvn_0, n_cells ÷ 2)),
        hcat(rand(mvn_0, n_cells ÷ 2), rand(mvn_1, n_cells ÷ 2)),
    ),
)

df_exp = DataFrame(expression_matrix, :auto)
df_exp[!, :cell] = 1:n_cells

@vlplot(
    :rect,
    x = "cell:o",
    color = :value,
    encoding = {
        y = {field = "variable", type = "nominal", sort = "-x", axis = {title = "gene"}}
    }
)(
    DataFrames.stack(df_exp, 1:n_genes)
)

@model function pPCA(x, ::Type{TV}=Array{Float64}) where {TV}
    # Dimensionality of the problem.
    N, D = size(x)

    # latent variable z
    z ~ filldist(Normal(), D, N)

    # side note for the curious
    # we use the more concise filldist syntax partly for compactness, but also for compatibility with other AD
    # backends, see the [Turing Performance Tipps](https://turing.ml/dev/docs/using-turing/performancetips)
    # w = TV{2}(undef, D, D)
    # for d in 1:D
    #  w[d, :] ~ MvNormal(ones(D))
    # end

    # weights/loadings W
    w ~ filldist(Normal(), D, D)

    # mean offset
    m ~ MvNormal(ones(D))
    mu = (w * z .+ m)'
    for d in 1:D
        x[:, d] ~ MvNormal(mu[:, d], ones(N))
    end
end;

ppca = pPCA(expression_matrix)
chain_ppca = sample(ppca, NUTS(), 500);

# Extract parameter estimates for plotting - mean of posterior
w = reshape(mean(group(chain_ppca, :w))[:, 2], (n_genes, n_genes))
z = permutedims(reshape(mean(group(chain_ppca, :z))[:, 2], (n_genes, n_cells)))'
mu = mean(group(chain_ppca, :m))[:, 2]

X = w * z

df_rec = DataFrame(X', :auto)
df_rec[!, :cell] = 1:n_cells

@vlplot(
    :rect,
    x = "cell:o",
    color = :value,
    encoding = {
        y = {field = "variable", type = "nominal", sort = "-x", axis = {title = "gene"}}
    }
)(
    DataFrames.stack(df_rec, 1:n_genes)
)

let
    diff = X' - expression_matrix
    @assert mean(diff[:, 4]) < 0.5
    @assert mean(diff[:, 5]) < 0.5
    @assert mean(diff[:, 6]) < 0.5
end

df_pca = DataFrame(z', :auto)
rename!(df_pca, Symbol.(["z" * string(i) for i in collect(1:n_genes)]))
df_pca[!, :cell] = 1:n_cells

@vlplot(:rect, "cell:o", "variable:o", color = :value)(DataFrames.stack(df_pca, 1:n_genes))

df_pca[!, :type] = repeat([1, 2]; inner=n_cells ÷ 2)
@vlplot(:point, x = :z1, y = :z2, color = "type:n")(df_pca)

@model function pPCA_ARD(x, ::Type{TV}=Array{Float64}) where {TV}
    # Dimensionality of the problem.
    N, D = size(x)

    # latent variable z
    z ~ filldist(Normal(), D, N)

    # weights/loadings w with Automatic Relevance Determination part
    alpha ~ filldist(Gamma(1.0, 1.0), D)
    w ~ filldist(MvNormal(zeros(D), 1.0 ./ sqrt.(alpha)), D)

    mu = (w' * z)'

    tau ~ Gamma(1.0, 1.0)
    for d in 1:D
        x[:, d] ~ MvNormal(mu[:, d], 1.0 / sqrt(tau))
    end
end;

ppca_ARD = pPCA_ARD(expression_matrix)
chain_pccaARD = sample(ppca_ARD, NUTS(), 500)

StatsPlots.plot(group(chain_pccaARD, :alpha))

# Extract parameter estimates for plotting - mean of posterior
w = permutedims(reshape(mean(group(chain_pccaARD, :w))[:, 2], (n_genes, n_genes)))
z = permutedims(reshape(mean(group(chain_pccaARD, :z))[:, 2], (n_genes, n_cells)))'
α = mean(group(chain_pccaARD, :alpha))[:, 2]
α

alpha_indices = sortperm(α)[1:2]
X = w[alpha_indices, alpha_indices] * z[alpha_indices, :]

df_rec = DataFrame(X', :auto)
df_rec[!, :cell] = 1:n_cells
@vlplot(:rect, "cell:o", "variable:o", color = :value)(DataFrames.stack(df_rec, 1:2))

df_pre = DataFrame(z', :auto)
rename!(df_pre, Symbol.(["z" * string(i) for i in collect(1:n_genes)]))
df_pre[!, :cell] = 1:n_cells

@vlplot(:rect, "cell:o", "variable:o", color = :value)(DataFrames.stack(df_pre, 1:n_genes))

df_pre[!, :type] = repeat([1, 2]; inner=n_cells ÷ 2)
df_pre[!, :ard1] = df_pre[:, alpha_indices[1]]
df_pre[!, :ard2] = df_pre[:, alpha_indices[2]]
@vlplot(:point, x = :ard1, y = :ard2, color = "type:n")(df_pre)

# Example data set - generate synthetic gene expression data

# dataset available in RDatasets
data = dataset("datasets", "iris")
species = data[!, "Species"]

# we extract the four measured quantities
d = 4
dat = data[!, 1:d]
# and the number of measurements
n = size(dat)[1];

ppca = pPCA(dat)

# Here we use a different sampler, we don't always have to use NUTS:
# Hamiltonian Monte Carlo (HMC) sampler parameters
ϵ = 0.05
τ = 10
chain_ppca2 = sample(ppca, HMC(ϵ, τ), 1000)

# Extract parameter estimates for plotting - mean of posterior
w = permutedims(reshape(mean(group(chain_ppca2, :w))[:, 2], (d, d)))
z = permutedims(reshape(mean(group(chain_ppca2, :z))[:, 2], (d, n)))'
mu = mean(group(chain_ppca2, :m))[:, 2]

X = w * z
# X = w * z .+ mu

df_rec = DataFrame(X', :auto)
df_rec[!, :species] = species
@vlplot(:rect, "species:o", "variable:o", color = :value)(DataFrames.stack(df_rec, 1:d))

df_iris = DataFrame(z', :auto)
rename!(df_iris, Symbol.(["z" * string(i) for i in collect(1:d)]))
df_iris[!, :sample] = 1:n
df_iris[!, :species] = species

@vlplot(:point, x = :z1, y = :z2, color = "species:n")(df_iris)

## Introduce batch effect
batch = rand(Binomial(1, 0.5), 150)
effect = rand(Normal(2.4, 0.6), 150)
batch_dat = dat .+ batch .* effect

ppca_batch = pPCA(batch_dat)
chain_ppcaBatch = sample(ppca_batch, HMC(ϵ, τ), 1000)
describe(chain_ppcaBatch)[1]

z = permutedims(reshape(mean(group(chain_ppcaBatch, :z))[:, 2], (d, n)))'
df_pre = DataFrame(z', :auto)
rename!(df_pre, Symbol.(["z" * string(i) for i in collect(1:d)]))
df_pre[!, :sample] = 1:n
df_pre[!, :species] = species
df_pre[!, :batch] = batch

@vlplot(:point, x = :z1, y = :z2, color = "species:n", shape = :batch)(df_pre)

@model function pPCA_residual(x, batch, ::Type{TV}=Array{Float64}) where {TV}

    # Dimensionality of the problem.
    N, D = size(x)

    # latent variable z
    z ~ filldist(Normal(), D, N)

    # weights/loadings w
    w ~ filldist(Normal(), D, D)

    # covariate vector
    w_batch = TV{1}(undef, D)
    w_batch ~ MvNormal(ones(D))

    # mean offset
    m = TV{1}(undef, D)
    m ~ MvNormal(ones(D))
    mu = m .+ w * z + w_batch .* batch'

    for d in 1:D
        x[:, d] ~ MvNormal(mu'[:, d], ones(N))
    end
end;

ppca_residual = pPCA_residual(batch_dat, convert(Vector{Float64}, batch))
chain_ppcaResidual = sample(ppca_residual, HMC(ϵ, τ), 1000);

z = permutedims(reshape(mean(group(chain_ppcaResidual, :z))[:, 2], (d, n)))'
df_post = DataFrame(z', :auto)
rename!(df_post, Symbol.(["z" * string(i) for i in collect(1:d)]))
df_post[!, :sample] = 1:n
df_post[!, :species] = species
df_post[!, :batch] = batch

@vlplot(:point, x = :z1, y = :z2, color = "species:n", shape = :batch)(df_post)

## helper functions for Householder transform
function V_low_tri_plus_diag(Q::Int, V)
    for q in 1:Q
        V[:, q] = V[:, q] ./ sqrt(sum(V[:, q] .^ 2))
    end
    return (V)
end

function Householder(k::Int, V)
    v = V[:, k]
    sgn = sign(v[k])

    v[k] += sgn
    H = LinearAlgebra.I - (2.0 / dot(v, v) * (v * v'))
    H[k:end, k:end] = -1.0 * sgn .* H[k:end, k:end]

    return (H)
end

function H_prod_right(V)
    D, Q = size(V)

    H_prod = zeros(Real, D, D, Q + 1)
    H_prod[:, :, 1] = Diagonal(repeat([1.0], D))

    for q in 1:Q
        H_prod[:, :, q + 1] = Householder(Q - q + 1, V) * H_prod[:, :, q]
    end

    return (H_prod)
end

function orthogonal_matrix(D::Int64, Q::Int64, V)
    V = V_low_tri_plus_diag(Q, V)
    H_prod = H_prod_right(V)

    return (H_prod[:, 1:Q, Q + 1])
end

@model function pPCA_householder(x, K::Int, ::Type{T}=Float64) where {T}

    # Dimensionality of the problem.
    D, N = size(x)
    @assert K <= D

    # parameters
    sigma_noise ~ LogNormal(0.0, 0.5)
    v ~ filldist(Normal(0.0, 1.0), Int(D * K - K * (K - 1) / 2))
    sigma ~ Bijectors.ordered(MvLogNormal(MvNormal(ones(K))))

    v_mat = zeros(T, D, K)
    v_mat[tril!(trues(size(v_mat)))] .= v
    U = orthogonal_matrix(D, Q, v_mat)

    W = zeros(T, D, K)
    W += U * Diagonal(sigma)

    Kmat = zeros(T, D, D)
    Kmat += W * W'
    for d in 1:D
        Kmat[d, d] = Kmat[d, d] + sigma_noise^2 + 1e-12
    end
    L = LinearAlgebra.cholesky(Kmat).L

    for q in 1:Q
        r = sqrt.(sum(dot(v_mat[:, q], v_mat[:, q])))
        Turing.@addlogprob! (-log(r) * (D - q))
    end

    Turing.@addlogprob! -0.5 * sum(sigma .^ 2) + (D - Q - 1) * sum(log.(sigma))
    for qi in 1:Q
        for qj in (qi + 1):Q
            Turing.@addlogprob! log(sigma[Q - qi + 1]^2) - sigma[Q - qj + 1]^2
        end
    end
    Turing.@addlogprob! sum(log.(2.0 * sigma))

    L_full = zeros(T, D, D)
    L_full += L * transpose(L)
    # fix numerical instability (non-posdef matrix)
    for d in 1:D
        for k in (d + 1):D
            L_full[d, k] = L_full[k, d]
        end
    end

    return x ~ filldist(MvNormal(L_full), N)
end;

# Dimensionality of latent space
Random.seed!(1789);
Q = 2
n_samples = 700
ppca_householder = pPCA_householder(Matrix(dat)', Q)
chain_ppcaHouseholder = sample(ppca_householder, NUTS(), n_samples);

# Extract mean of v from chain
N, D = size(dat)
vv = mean(group(chain_ppcaHouseholder, :v))[:, 2]
v_mat = zeros(Real, D, Q)
v_mat[tril!(trues(size(v_mat)))] .= vv
sigma = mean(group(chain_ppcaHouseholder, :sigma))[:, 2]
U_n = orthogonal_matrix(D, Q, v_mat)
W_n = U_n * (LinearAlgebra.I(Q) .* sigma)

# Create array with projected values
z = W_n' * transpose(Matrix(dat))
df_post = DataFrame(convert(Array{Float64}, z)', :auto)
rename!(df_post, Symbol.(["z" * string(i) for i in collect(1:Q)]))
df_post[!, :sample] = 1:n
df_post[!, :species] = species

@vlplot(:point, x = :z1, y = :z2, color = "species:n")(df_post)

## Create data projections for each step of chain
vv = collect(get(chain_ppcaHouseholder, [:v]).v)
v_mat = zeros(Real, D, Q)
vv_mat = zeros(Float64, n_samples, D, Q)
for i in 1:n_samples
    index = BitArray(zeros(n_samples, D, Q))
    index[i, :, :] = tril!(trues(size(v_mat)))
    tmp = zeros(size(vv)[1])
    for j in 1:(size(vv)[1])
        tmp[j] = vv[j][i]
    end
    vv_mat[index] .= tmp
end

ss = collect(get(chain_ppcaHouseholder, [:sigma]).sigma)
sigma = zeros(Q, n_samples)
for d in 1:Q
    sigma[d, :] = Array(ss[d])
end

samples_raw = Array{Float64}(undef, Q, N, n_samples)
for i in 1:n_samples
    U_ni = orthogonal_matrix(D, Q, vv_mat[i, :, :])
    W_ni = U_ni * (LinearAlgebra.I(Q) .* sigma[:, i])
    z_n = W_ni' * transpose(Matrix(dat))
    samples_raw[:, :, i] = z_n
end

# initialize a 3D plot with 1 empty series
plt = plot(
    [100, 200, 300];
    xlim=(-4.00, 7.00),
    ylim=(-100.00, 0.00),
    group=["Setosa", "Versicolor", "Virginica"],
    markercolor=["red", "blue", "black"],
    title="Visualization",
    seriestype=:scatter,
)

anim = @animate for i in 1:n_samples
    scatter!(
        plt,
        samples_raw[1, 1:50, i],
        samples_raw[2, 1:50, i];
        color="red",
        seriesalpha=0.1,
        label="",
    )
    scatter!(
        plt,
        samples_raw[1, 51:100, i],
        samples_raw[2, 51:100, i];
        color="blue",
        seriesalpha=0.1,
        label="",
    )
    scatter!(
        plt,
        samples_raw[1, 101:150, i],
        samples_raw[2, 101:150, i];
        color="black",
        seriesalpha=0.1,
        label="",
    )
end
gif(anim, "anim_fps.gif"; fps=5)

let
    m1 = mean(samples_raw[1, 1:50, :])
    m2 = mean(samples_raw[1, 51:100, :])
    m3 = mean(samples_raw[1, 101:150, :])
    @assert m1 - m2 > 3
    @assert m1 - m3 > 3
    @assert m2 - m3 < 2
end

# kernel density estimate
using KernelDensity
dens = kde((vec(samples_raw[1, :, :]), vec(samples_raw[2, :, :])))
StatsPlots.plot(dens)

max_stars_repo_name: TuringLang/TuringTutorialsOutput

max_issues_repo_name: TuringLang/TuringTutorialsOutput

max_forks_repo_name: TuringLang/TuringTutorialsOutput


text:
#!/usr/bin/julia

using Pkg
include("gcc.jl")

file_prefix = "processed"
num_executions = 15

out_file = open("output.csv", "w")

function output_csv_header()
	header = "filename; functions; expected_insns; inlined_functions; parallel; "

	for i in 1:(num_executions-1)
		header = header * "execution_" * string(i) * "; "
	end

	header = header * "execution_" * string(num_executions)

	println(out_file, header)
end

function output_csv_line(file, num_functions, expected_insn, inlined_functions, parallel, data)
	str = file * "; " * string(num_functions) * "; " * string(expected_insn) * "; " *
	string(inlined_functions) * "; " * string(parallel)

	for d in data
		str = str * "; " * string(d)
	end

	println(out_file, str)
end

function y(x)
	runs = Float64[]

	for i in 1:num_executions
		push!(runs, run_gcc(x))
	end

	return runs
end

function generate_database()
	prefix = file_prefix

	files = readdir(prefix)
	output_csv_header()

	num_files = length(files)

	i = 0
	for file in files
		if file == "README.txt"
			continue
		end

		a = Int64(floor((i * 100)/num_files))
		percent = "[" * string(a) * "%]"

		println(percent, " Testing file: ", file)

		splitted = split(file, ".")
		base = splitted[1]
		ext = splitted[2]

		base_split = split(base, "-")
		num_func_str = base_split[1]
		expected_insn_str = base_split[2]
		inlined_funcs_str = base_split[3]
		num_funcs = parse(Int64, num_func_str)
		expected_insn = parse(Int64, expected_insn_str)
		inlined_funcs = parse(Int64, inlined_funcs_str)

		for threads in (1, 2, 4, 8)
			result = y([file, 1, threads])
			output_csv_line(file, num_funcs, expected_insn, inlined_funcs, threads, result)
		end

		i = i + 1
	end
end

generate_database()
close(out_file)

max_stars_repo_name: giulianobelinassi/LTO-Timing-Analysis

max_issues_repo_name: giulianobelinassi/LTO-Timing-Analysis

max_forks_repo_name: giulianobelinassi/LTO-Timing-Analysis


text:
using BioSequences

# something later

max_stars_repo_name: KorfLab/julie

max_issues_repo_name: KorfLab/julie

max_forks_repo_name: KorfLab/julie