text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

using JuliaFEM
using JuliaFEM.Preprocess
using JuliaFEM.Postprocess
using JuliaFEM.Testing

#=
- solve 2d plane stress problem with known solution:
    surface traction force in 2d
    volume load in 2d
    reaction force
- test postprocessing of nodal fields:
    geometry
    displacement
    reaction force
    concentrated force
=#
@testset "test 2d linear elasticity with surface + volume load" begin
    meshfile = "/geometry/2d_block/BLOCK_1elem.med"
    mesh = aster_read_mesh(Pkg.dir("JuliaFEM")*meshfile)

    # field problem
    block = Problem(Elasticity, "BLOCK", 2)
    block.properties.formulation = :plane_stress
    block.properties.finite_strain = false
    block.properties.geometric_stiffness = false
    block.elements = create_elements(mesh, "BLOCK")
    update!(block.elements, "youngs modulus", 288.0)
    update!(block.elements, "poissons ratio", 1/3)
    update!(block.elements, "displacement load 2", 576.0)

    # traction
    traction = Problem(Elasticity, "TRACTION", 2)
    traction.properties.formulation = :plane_stress
    traction.properties.finite_strain = false
    traction.properties.geometric_stiffness = false
    traction.elements = create_elements(mesh, "TOP")
    update!(traction, "displacement traction force 2", 288.0)

    # boundary conditions
    bc_sym_23 = Problem(Dirichlet, "symmetry bc 23", 2, "displacement")
    bc_sym_23.elements = create_elements(mesh, "LEFT")
    update!(bc_sym_23, "displacement 1", 0.0)
    bc_sym_13 = Problem(Dirichlet, "symmetry bc 13", 2, "displacement")
    bc_sym_13.elements = create_elements(mesh, "BOTTOM")
    update!(bc_sym_13, "displacement 2", 0.0)

    solver = LinearSolver(block, traction, bc_sym_23, bc_sym_13)
    solver()

    info("u = ", block.assembly.u)
    info("λ = ", block.assembly.la)

    f = 288.0
    g = 576.0
    E = 288.0
    nu = 1/3
    u3_expected = f/E*[-nu, 1] + g/(2*E)*[-nu, 1]

    # fetch nodal results X + u and join them into one table using DataFrames
    X = solver(DataFrame, "geometry", :COOR)
    u = solver(DataFrame, "displacement", :U)
    la = solver(DataFrame, "reaction force", :RF)
    f = solver(DataFrame, "concentrated force", :CF)
    results = join(X, u, on=:NODE, kind=:outer)
    results = join(results, la, on=:NODE, kind=:outer)
    length(f) != 0 && (results = join(results, f, on=:NODE, kind=:outer))
    sort!(results, cols=[:NODE])
    println(results)

    u3 = extract(results, NODE=:N3, :U1, :U2)
    @test isapprox(u3, u3_expected)

    # element details
    el = first(block.elements)
    S1 = block(el, [0.0, 0.0], 0.0, Val{:S})
    S1 = S1[[1,4,2]]
    E1= block(el, [0.0, 0.0], 0.0, Val{:E})
    E1 = E1[[1,4,2]]
    C1 = block(el, [0.0, 0.0], 0.0, Val{:COORD})
    info("strain = $E1, stress = $S1, at $C1")
    @test isapprox(E1, [-2/3, 2.0, 0.0])
    @test isapprox(S1, [0.0, 576.0, 0.0])
    @test isapprox(C1, [0.5, 0.5])

    S1 = block(DataFrame, 0.0, Val{:S})
    E1 = block(DataFrame, 0.0, Val{:E})
    C1 = block(DataFrame, 0.0, Val{:COORD})

    println(S1)
    println(E1)
    println(C1)

    S = solver(DataFrame, 0.0, Val{:S})
    println(S)

    info(solver("displacement", 0.0))
    solver()
    info(solver("displacement", 0.0))
    u = solver("displacement", 0.0)[3]
    info("u3 = $u")
    @test isapprox(u, u3_expected)

    info("calling nonlinear solver")
    solver2 = NonlinearSolver(block, traction, bc_sym_23, bc_sym_13)
    solver2()
    u = solver2("displacement", 0.0)[3]
    info("nlsolver u3 = $u, expected = $u3_expected")
    @test isapprox(u, u3_expected; rtol=1.0e-5)
end

#= TODO: to other file
@testset "test dump model to disk and read back before and after solution" begin
    solver = get_model("test 2d linear elasticity with surface + volume load")
    save("/tmp/model.jld", "linear_model", solver)
    solver2 = load("/tmp/model.jld")["linear_model"]
    solver2()
    save("/tmp/model.jld", "results", solver2)
    solver3 = load("/tmp/model.jld")["results"]
    block = solver3["BLOCK"]
    u3 = reshape(block.assembly.u, 2, 4)[:,3]
    f = 288.0
    g = 576.0
    E = 288.0
    nu = 1/3
    u3_expected = f/E*[-nu, 1] + g/(2*E)*[-nu, 1]
    @test isapprox(u3, u3_expected)
end
=#


text:
# compatiblestreamingmodetest.jl
using PicoScope
using Base.Test

println("testing compatable streaming mode...")
handle = ps3000_open_unit()
println("handle = ",handle)
enabled = true
dc_coupling = true
ps3000_set_channel(handle, PS3000Channel.A, enabled, dc_coupling, PS3000Range.V1)
sample_interval_ms = 10
max_samples = 1000 # size of buffer for ps3000_run_streaming
windowed = 0 # only return values since last call
number_of_samples = 100
vbuffer = ValueBuffer(number_of_samples)
ps3000_run_streaming(handle, sample_interval_ms, max_samples, windowed)
wait(Timer(1))
(values_per_channel_returned, overflow) =
  ps3000_get_values!(handle,vbuffer,number_of_samples)
# @test values_per_channel_returned<100
ps3000_stop(handle)
# ps3000_close_unit(handle)  # does not return from close when in compatable streaming mode
println("done")

ext: jl


text:
# ---
# title: 811. Subdomain Visit Count
# id: problem811
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Hash Table
# link: <https://leetcode.com/problems/subdomain-visit-count/description/>
# hidden: true
# ---
#
# A website domain like "discuss.leetcode.com" consists of various subdomains.
# At the top level, we have "com", at the next level, we have "leetcode.com",
# and at the lowest level, "discuss.leetcode.com". When we visit a domain like
# "discuss.leetcode.com", we will also visit the parent domains "leetcode.com"
# and "com" implicitly.
#
# Now, call a "count-paired domain" to be a count (representing the number of
# visits this domain received), followed by a space, followed by the address. An
# example of a count-paired domain might be "9001 discuss.leetcode.com".
#
# We are given a list `cpdomains` of count-paired domains. We would like a list
# of count-paired domains, (in the same format as the input, and in any order),
# that explicitly counts the number of visits to each subdomain.
#
#
#
#     **Example 1:**
#     Input:
#     ["9001 discuss.leetcode.com"]
#     Output:
#     ["9001 discuss.leetcode.com", "9001 leetcode.com", "9001 com"]
#     Explanation:
#     We only have one website domain: "discuss.leetcode.com". As discussed above, the subdomain "leetcode.com" and "com" will also be visited. So they will all be visited 9001 times.
#
#
#
#
#     **Example 2:**
#     Input:
#     ["900 google.mail.com", "50 yahoo.com", "1 intel.mail.com", "5 wiki.org"]
#     Output:
#     ["901 mail.com","50 yahoo.com","900 google.mail.com","5 wiki.org","5 org","1 intel.mail.com","951 com"]
#     Explanation:
#     We will visit "google.mail.com" 900 times, "yahoo.com" 50 times, "intel.mail.com" once and "wiki.org" 5 times. For the subdomains, we will visit "mail.com" 900 + 1 = 901 times, "com" 900 + 50 + 1 = 951 times, and "org" 5 times.
#
#
#
# **Notes:**
#
#   * The length of `cpdomains` will not exceed `100`.
#   * The length of each domain name will not exceed `100`.
#   * Each address will have either 1 or 2 "." characters.
#   * The input count in any count-paired domain will not exceed `10000`.
#   * The answer output can be returned in any order.
#
#
## @lc code=start
using LeetCode

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


text:
module GeometricFlux

using Statistics: mean
using LinearAlgebra: Adjoint, norm, Transpose
using Reexport

using CUDA
using ChainRulesCore: @non_differentiable
using FillArrays: Fill
using Flux
using Flux: glorot_uniform, leakyrelu, GRUCell, @functor
using NNlib, NNlibCUDA
using GraphSignals
using Graphs
using Zygote

export
    # layers/graphlayers
    AbstractGraphLayer,

    # layers/gn
    GraphNet,

    # layers/msgpass
    MessagePassing,

    # layers/conv
    GCNConv,
    ChebConv,
    GraphConv,
    GATConv,
    GatedGraphConv,
    EdgeConv,
    GINConv,
    CGConv,

    # layer/pool
    GlobalPool,
    LocalPool,
    TopKPool,
    topk_index,

    # models
    GAE,
    VGAE,
    InnerProductDecoder,
    VariationalEncoder,
    summarize,
    sample,

    # layer/selector
    bypass_graph,

    # utils
    generate_cluster

include("datasets.jl")

include("utils.jl")

include("layers/graphlayers.jl")
include("layers/gn.jl")
include("layers/msgpass.jl")

include("layers/conv.jl")
include("layers/pool.jl")
include("models.jl")
include("layers/misc.jl")

include("cuda/conv.jl")

using .Datasets

end


text:
atol = 1e-8
rtol = 1e-16
max_iter = 500
r = PH.ScalarPenaltyParameter(25.0)

@testset "Residuals" begin
    (n, err, rerr, obj, soln, phd) = PH.solve(build_scen_tree(),
                                              create_model,
                                              r,
                                              opt=Ipopt.Optimizer,
                                              opt_args=(print_level=0,tol=1e-12),
                                              atol=atol,
                                              rtol=rtol,
                                              max_iter=max_iter,
                                              report=0,
                                              save_residuals=1,
                                              timing=false,
                                              warm_start=false)

    res_df = residuals(phd)
    @test size(res_df,1) == n + 1
    rid = n + 1
    @test res_df[rid, :iteration] == n
    @test res_df[rid, :absolute] == err
    @test res_df[rid, :relative] == rerr
end

@testset "Iterates" begin
    (n, err, rerr, obj, soln, phd) = PH.solve(build_scen_tree(),
                                              create_model,
                                              r,
                                              opt=Ipopt.Optimizer,
                                              opt_args=(print_level=0,tol=1e-12),
                                              atol=atol,
                                              rtol=rtol,
                                              max_iter=max_iter,
                                              report=0,
                                              save_iterates=1,
                                              timing=false,
                                              warm_start=false)

    rid = n + 1

    xhat_df = retrieve_xhat_history(phd)
    @test size(xhat_df,1) == n + 1
    @test xhat_df[rid, :iteration] == n
    for (k,var) in enumerate(soln[:,:variable])
        if soln[k,:stage] < 3
            vname = var * "_" * soln[k,:scenarios]
            @test xhat_df[rid, vname] == soln[k, :value]
        end
    end

    x_hist_df = retrieve_no_hat_history(phd)
    x_df = retrieve_no_hats(phd)
    @test size(x_hist_df,1) == n + 1
    @test x_hist_df[rid, :iteration] == n
    for (k,var) in enumerate(x_df[:,:variable])
        if x_df[k,:stage] < 3
            vname = var * "_" * string(x_df[k,:scenario])
            @test x_hist_df[rid, vname] == x_df[k,:value]
        end
    end

    w_hist_df = retrieve_w_history(phd)
    w_df = retrieve_w(phd)
    @test size(w_hist_df,1) == n + 1
    @test w_hist_df[rid, :iteration] == n
    for (k,var) in enumerate(w_df[:,:variable])
        vname = var * "_" * string(w_df[k,:scenario])
        @test w_hist_df[rid, vname] == w_df[k,:value]
    end
end

@testset "Reporting" begin
    global n, err, rerr, obj, soln, phd

    reports = @capture_out((n, err, rerr, obj, soln, phd) = PH.solve(build_scen_tree(),
                                                                     create_model,
                                                                     r,
                                                                     opt=Ipopt.Optimizer,
                                                                     opt_args=(print_level=0,
                                                                               tol=1e-12),
                                                                     atol=atol,
                                                                     rtol=rtol,
                                                                     max_iter=max_iter,
                                                                     report=1,
                                                                     lower_bound=1,
                                                                     save_residuals=1,
                                                                     timing=false,
                                                                     warm_start=false)
                           )

    report_rex = r"Iter:\s*(\d+)\s*AbsR:\s*(\d\.\d+e[+-]\d+)\s*RelR:\s*(\d\.\d+e[+-]\d+)\s*Xhat:\s*(\d\.\d+e[+-]\d+)\s*X:\s*(\d\.\d+e[+-]\d+)"
    bound_rex = r"Iter:\s*(\d+)\s*Bound:\s*(-?\d\.\d+e[+-]\d+)\s*Abs Gap:\s*(\d\.\d+e[+-]\d+)\s*Rel Gap:\s*((\d)+\.\d+(e[+-]\d+)?)"

    res_df = residuals(phd)
    lb_df = lower_bounds(phd)
    N = 5 # number of consensus variables

    nreps = 0
    nbds = 0

    for line in split(reports,"\n")

        if occursin(report_rex, line)

            nreps += 1
            if nreps == Int(floor(n/2))
                m = match(report_rex, line)
                iter = parse(Int, m.captures[1])
                abs = parse(Float64, m.captures[2])
                @test isapprox(res_df[iter + 1, :absolute], abs, rtol=1e-6)
                rel = parse(Float64, m.captures[3])
                @test isapprox(res_df[iter + 1, :relative], rel, rtol=1e-6)
                xhat = parse(Float64, m.captures[4])
                @test isapprox(res_df[iter + 1, :xhat_sq], xhat^2 * N, rtol=1e-6)
                x = parse(Float64, m.captures[5])
                @test isapprox(res_df[iter + 1, :x_sq], x^2 * N, rtol=1e-6)
            end

        elseif occursin(bound_rex, line)

            nbds += 1
            if nbds == Int(floor(n/2))
                m = match(bound_rex, line)
                iter = parse(Int, m.captures[1])
                bound = parse(Float64, m.captures[2])
                @test isapprox(lb_df[iter + 1, "bound"], bound, rtol=1e-3)
                abs = parse(Float64, m.captures[3])
                @test isapprox(lb_df[iter + 1, "absolute gap"], abs, rtol=1e-3)
                rel = parse(Float64, m.captures[4])
                @test isapprox(lb_df[iter + 1, "relative gap"], rel, rtol=1e-3)
            end

        end
    end

    @test nreps == n + 1
    @test nbds == n + 1
end

@testset "Timing" begin
    global n, err, rerr, obj, soln, phd
    reports = @capture_out((n, err, rerr, obj, soln, phd) = PH.solve(build_scen_tree(),
                                                                     create_model,
                                                                     r,
                                                                     opt=Ipopt.Optimizer,
                                                                     opt_args=(print_level=0,
                                                                               tol=1e-12),
                                                                     atol=atol,
                                                                     rtol=rtol,
                                                                     max_iter=max_iter,
                                                                     report=1,
                                                                     save_residuals=1,
                                                                     timing=false,
                                                                     warm_start=false)
                           )
    @test length(reports) > 1

end


text:
srand(1)
a = mapweight(grid2(20),x->1/(rand(1)[1]));
la = lap(a)
n = size(la)[1]
b = randn(n)
b = b - mean(b);

d = diag(la)
pre = (x)->(x ./ d)

x = pcg(la,b,pre,maxits=2000,maxtime=1,verbose=true);
x = Laplacians.pcgBLAS(la,b,pre,maxits=2000,maxtime=1,verbose=true);
x = Laplacians.pcgSlow(la,b,pre,maxits=2000,maxtime=1,verbose=true);
x = pcg(la,b,pre,maxits=2000,verbose=true);
x = Laplacians.pcgBLAS(la,b,pre,maxits=2000,verbose=true);
x = Laplacians.pcgSlow(la,b,pre,maxits=2000,verbose=true);
x = Laplacians.pcgBLAS(la,b,pre,tol=1e-2);
x = Laplacians.pcgSlow(la,b,pre,tol=1e-2);
x = pcg(la,b,pre,tol=1e-2);

x = cg(la,b,verbose=true,maxtime=1,maxits=1000);
x = Laplacians.cgBLAS(la,b,verbose=true,maxtime=1,maxits=1000);
x = Laplacians.cgSlow(la,b,verbose=true,maxtime=1,maxits=1000);
x = cg(la,b,verbose=true,maxits=2000);
x = Laplacians.cgBLAS(la,b,verbose=true,maxits=2000);
x = Laplacians.cgSlow(la,b,verbose=true,maxits=2000);
x = cg(la,b,verbose=true,tol=0.5);
x = Laplacians.cgBLAS(la,b,verbose=true,tol=0.5);
x = Laplacians.cgSlow(la,b,verbose=true,tol=0.5);

n = 100
a = wtedChimera(n,3)
la = lap(a)

f = cgSolver(la,verbose=false)
b = randn(n)
b = b - mean(b)
x = f(b)
norm(la*x-b)

x = f(b,maxits=15,verbose=true, maxtime=1)

d = diagm(1./diag(la))
pre = x -> d*x

x = pcg(la, b, diagm(diag(la)),maxits=10)
x = pcg(la, b, pre,maxits=10,verbose=true)
f = pcgSolver(la, diagm(diag(la)) ,maxits=10,verbose=true)
x = f(b)
f = pcgSolver(la,pre,maxits=10,verbose=false)
x = f(b,verbose=true, maxits=1000, maxtime = 2)


text:
module MapInterfacesTests

using Test
using Gridap.Arrays
using Gridap.TensorValues

# using BenchmarkTools
using FillArrays

# Map Interfaces

a = [3,2]
b = [2,1]
test_map(a+b,+,a,b)

bm = Broadcasting(+)
cache = return_cache(bm,a,b)
# @btime evaluate!($cache,$bm,$a,$b)

# m = rand(2,2)
# test_map(m,(a,b),m)
# testitem(m,a,b)

# m = rand(2,2)
# test_map(m,(a,b),m)
# testitem(m,a,b)

# cs = map(op -> return_cache(op,a,b),(+,m))
# cs = return_caches((+,m),a,b)
# evaluate!(cs,(+,m),a,b) == (a+b,m)
# evaluate((+,m),a,b) == (a+b,m)

# @test map(op -> return_type(op,a,b),(+,m)) == (Array{Int,1}, Array{Float64,2})

z = evaluate(+,a,b)
c = return_cache(+,a,b)
evaluate!(c,+,a,b) == a+b
typeof(z) == typeof(a+b)
return_type(+,a,b)

f = Broadcasting(+)
a = rand(3,2)
b = 3
c = a .+ b
test_map(c,f,a,b)

k = Broadcasting(-)
test_map(-1,k,1)
test_map([-1,-2],k,[1,2])
test_map(-1,k,1,2)
test_map(-1.0,k,1.0,2)
test_map(-1.0,k,1,2.0)
test_map([-1,0],k,[1,2],2)
test_map([1,0],k,2,[1,2])
test_map([2,2],k,[3,4],[1,2])

f = Broadcasting(⋅)
a = fill(TensorValue(2,0,0,0,2,0,0,0,2),2)
b = VectorValue(1,2,3)
c = zeros(VectorValue{3,Int},2)
broadcast!(⋅,c,a,b)
test_map(c,f,a,b)

cache = return_cache(f,a,b)
# @btime evaluate!($cache,$f,$a,$b)

# x = [1,2]
x = rand(3,3)

fa(x) = 2*x
test_map(2*x,fa,x)

fb(x) = sqrt.(x)
test_map(sqrt.(x),fb,x)

op = Broadcasting(*)
cache = return_cache(op,2,x)
# @btime evaluate!($cache,$op,$2,$x)
test_map(2*x,Broadcasting(*),2,x)

fab = Operation(fa)(fb)
test_map(2*(sqrt.(x)),fab,x)

bm = Broadcasting(*)
cache = return_cache(bm,x,2)
# @btime evaluate!($cache,$bm,$x,$2)

bs = Broadcasting(sqrt)
cache = return_cache(bs,x)
# @btime evaluate!($cache,$bs,$x)

h = Operation(bs)(bm)
cache = return_cache(h,x,2)
# @btime evaluate!($cache,$h,$x,$2)

# More exotic functions as Maps

myf(::Type{T}) where T = 1.0*one(T)
@test return_type(myf,Int) == Float64

# k = -

# v = 3.0
# d = 2
# f = MockField{d}(v)
# g = composition(k,f)

# p = 4
# p1 = Point(1,2)
# p2 = Point(2,1)
# p3 = Point(4,3)
# p4 = Point(6,1)
# x = [p1,p2,p3,p4]

# fx = evaluate(f,x)
# ∇fx = evaluate(∇(f),x)
# gx = evaluate(k,fx)
# ∇gx = evaluate(k,∇fx)
# test_map(g,(x,),gx) #,grad=∇gx)

# # @santiagobadia : Create the gradient method for composition
# # using chain rule
# ∇g = composition(k,∇(f))
# test_map(∇g,(x,),∇gx)

# g = field_composition(k,f)
# @test g isa NewField
# gx = evaluate(g,x)
# cache = return_cache(g,x)
# test_map(g,(x,),gx) #grad=∇gx)

# fi = 3.0
# gi = 4.5
# d = 2
# f = MockField{d}(fi)
# g = MockField{d}(gi)

# # @santiagobadia : I like the tuple, k ∘ (f,g)
# h = composition(k,f,g)

# hx = evaluate(k,evaluate(f,x),evaluate(g,x))
# ∇hx = evaluate(k,evaluate(∇(f),x),evaluate(∇(g),x))
# test_map(h,(x,),hx) #,grad=∇hx)

# ∇h = composition(k,∇(f),∇(g))
# test_map(∇h,(x,),∇hx)

# fi = VectorValue(3.0,0.0)
# gi = VectorValue(4,5)
# d = 2
# f = MockField{d}(fi)
# g = ConstantField(gi)
# h = composition(k,f,g)

# hx = evaluate(k,evaluate(f,x),evaluate(g,x))
# ∇hx = evaluate(k,evaluate(∇(f),x),evaluate(∇(g),x))
# test_map(h,(x,),hx)

# ∇h = composition(k,∇(f),∇(g))
# test_map(∇h,(x,),∇hx)

# k = Broadcasting(-)
# fi = VectorValue(3.0,0.0)
# d = 2
# f = MockField{d}(fi)
# g = broadcast(ConstantField,x)
# h = composition(k,f,g)

# # @santiagobadia : A field can be paralelised wrt array of points
# # A mapping on the contrary can take multiple arguments...
# # We must create two evaluates, one for arrays of tuples and
# # one for one tuple
# # @santiagobadia : The previous version of gradient is wrong
# # @santiagobadia : We could also have derivative (transpose of gradient),
# # i.e., a Jacobian for f : Rn -> Rm

# hx = evaluate(k,evaluate(f,x),evaluate(g,x))
# test_map(h,(x,),hx) #,grad=∇hx)

# G = ∇(g)

# v = 3.0
# v = VectorValue(3.0,0.0)
# f = MockField{d}(v)

# f = ConstantField(v)

# fx = evaluate(f,x)
# test_field(f,(x,),fx)

end # module


text:
using SnpArrays, Test, Statistics, Random, TraitSimulation
using DataFrames, LinearAlgebra, OrdinalMultinomialModels, GLM
Random.seed!(1234)

# demo genetic data can be found in Snparrays package directory
filename = "EUR_subset"
EUR = SnpArray(SnpArrays.datadir(filename * ".bed"));
GRM = grm(EUR, minmaf = 0.05);
EUR_data = SnpData(SnpArrays.datadir(filename));
bimfile = EUR_data.snp_info # store the snp_info with the snp names
snpid  = bimfile[:, :snpid] # store the snp names in the snpid vector
causal_snp_index = findfirst(x -> x ==  "rs62057050" , snpid) # find the index of the snp of interest by snpid
locus = convert(Vector{Float64}, @view EUR[:, causal_snp_index])
famfile = EUR_data.person_info
sex = map(x -> strip(x) == "F" ? 0.0 : 1.0, famfile[!, :sex])
X = [sex locus]
β_full =
 [0.52
0.05]

# first ordinal power for the single snp
θ = [1.0, 1.2, 1.4]
trait_ordered = OrderedMultinomialTrait(X, β_full, θ, LogitLink())

nsim = 10
γs = collect(0.0:0.05:0.5)

randomseed = 12345
ordinal_pvals = power_simulation(nsim, γs, trait_ordered, randomseed)

alpha = 0.005
powers = power(ordinal_pvals, alpha)
@test powers[1] < powers[end]

# next glm power
dist = Poisson()
link = LogLink()
pdf_age = Normal(45, 8)
# simulate age under the specified pdf_age and standardize to be ~ N(0, 1)
age = rand(pdf_age, length(sex))
X = [sex age locus]
β_full =
[2.52
1.0
0.05]
glm_trait = GLMTrait(X, β_full, dist, link)

nsim = 10
γs = collect(0.0:0.5:2.0)

randomseed = 12345
glm_pvals = power_simulation(nsim, γs, glm_trait, randomseed)

alpha = 0.20
powers = power(glm_pvals, alpha)
@test powers[1] < powers[end]


text:
@testset "noncentral beta cdf" begin
    @test ncbeta(Float32(5.0),Float32(5.0),Float32(54.0),Float32(0.8640)) ≈ Float32(0.456302)
    @test ncbeta(Float32(5.0),Float32(5.0),Float32(140.0),Float32(0.90)) ≈ Float32(0.104134)
    @test ncbeta(Float32(5.0),Float32(5.0),Float32(170.0),Float32(0.9560)) ≈ Float32(0.602235)
    @test ncbeta(Float32(10.0),Float32(10.0),Float32(140.0),Float32(0.90)) ≈ Float32(0.600811)
    @test ncbeta(Float32(10.0),Float32(10.0),Float32(250.0),Float32(0.9)) ≈ Float32(0.902850E-01)
    @test ncbeta(Float32(20.0),Float32(20.0),Float32(54.0),Float32(0.8787)) ≈ Float32(0.999865)
    @test ncbeta(Float32(20.0),Float32(20.0),Float32(140.0),Float32(0.9)) ≈ Float32(0.992600)
    @test ncbeta(Float32(20.0),Float32(20.0),Float32(250.0),Float32(0.922)) ≈ Float32(0.964111)
    @test ncbeta(Float32(10.0),Float32(20.0),Float32(150.0),Float32(0.868)) ≈ Float32(0.937663)
    @test ncbeta(Float32(10.0),Float32(10.0),Float32(120.0),Float32(0.9)) ≈ Float32(0.730682)
    @test ncbeta(Float32(15.0),Float32(5.0),Float32(80.0),Float32(0.88)) ≈ Float32(0.160426)
    @test ncbeta(Float32(20.0),Float32(10.0),Float32(110.0),Float32(0.85)) ≈ Float32(0.186749)
    @test ncbeta(Float32(20.0),Float32(30.0),Float32(65.0),Float32(0.66)) ≈ Float32(0.655939)
    @test ncbeta(Float32(20.0),Float32(50.0),Float32(130.0),Float32(0.720)) ≈ Float32(0.979688)
    @test ncbeta(Float32(30.0),Float32(20.0),Float32(80.0),Float32(0.720)) ≈ Float32(0.116239)
    @test ncbeta(Float32(30.0),Float32(40.0),Float32(130.0),Float32(0.8)) ≈ Float32(0.993043)
    @test ncbeta(Float32(10.0),Float32(5.0),Float32(20.0),Float32(0.644)) ≈ Float32(0.506899E-01)
    @test ncbeta(Float32(10.0),Float32(10.0),Float32(54.0),Float32(0.7)) ≈ Float32(0.103096)
    @test ncbeta(Float32(10.0),Float32(30.0),Float32(80.0),Float32(0.78)) ≈ Float32(0.997842)
    @test ncbeta(Float32(15.0),Float32(20.0),Float32(120.0),Float32(0.76)) ≈ Float32(0.255555)
    @test ncbeta(Float32(10.0),Float32(5.0),Float32(55.0),Float32(0.795)) ≈ Float32(0.668307E-01)
    @test ncbeta(Float32(12.0),Float32(17.0),Float32(64.0),Float32(0.56)) ≈ Float32(0.113601E-01)
    @test ncbeta(Float32(30.0),Float32(30.0),Float32(140.0),Float32(0.80)) ≈ Float32(0.781337)
    @test ncbeta(Float32(35.0),Float32(30.0),Float32(20.0),Float32(0.670)) ≈ Float32(0.886713)
    @test ncF(Float32(2.0), Float32(3.0), Float32(4.0), Float32(5.0)) ≈ Float32(0.3761448105)
end


text:
"""
Buckley Leverett equation
dependent variables: water saturation
Prepared for educational purposes by Ali A Eftekhari
"""
function forced_imb_impes(mu_water, mu_oil, ut, phi0,
  perm_val, swc, sor, kro0, no, krw0, nw, sw0, sw_inj, L, pv_inj; Nx = 100)
  m = createMesh1D(Nx, L)
  # m = createMeshCylindrical2D(Nx, Ny, W, H) # creates a 2D mesh
  ## define the physical parametrs

  p0 = 1e5 # [bar] pressure
  pin = 150e5 # [bar] injection pressure at the left boundary
  u_in= ut # [m/s] equal to 1 m/day
  sw_in = sw_inj

  eps1=1e-6

  k=createCellVariable(m, perm_val)
  phi=createCellVariable(m, phi0)

  lw = geometricMean(k)/mu_water
  lo = geometricMean(k)/mu_oil

  ## Define the boundaries: all fixed Sw=1, fixed pressure everywhere(?)
  BCp = createBC(m) # Neumann BC for pressure
  BCs = createBC(m) # Neumann BC for saturation
  # left boundary pressure gradient
  # change the right boandary to constant pressure (Dirichlet)
  BCp.left.a[:]=perm_val/mu_water
  BCp.left.b[:]=0
  BCp.left.c[:]=-u_in
  BCp.right.a[:]=0.0
  BCp.right.b[:]=1.0
  BCp.right.c[:]=p0

  BCs.left.a[:]=0.0
  BCs.left.b[:]=1.0
  BCs.left.c[:]=sw_in

  t_end = pv_inj*phi0*L/u_in # [s] final time
  dt=L/Nx/u_in
  eps_p = 1e-7 # pressure accuracy
  eps_sw = 1e-7 # saturation accuracy
  ## define the variables
  sw_old = createCellVariable(m, sw0, BCs)
  p_old = createCellVariable(m, p0, BCp)
  sw = copyCell(sw_old)
  oil_init=domainInt(1-sw_old)
  p = copyCell(p_old)
  uw = -gradientTerm(p_old) # an estimation of the water velocity
  ## start the main loop
  # generate intial pressure profile (necessary to initialize the fully
  # implicit solver)
  rec_fact=zeros(1)
  t_day=zeros(1)
  dp_hist = zeros(1)
  t = 0.0
  dt0=dt
  dsw_alwd= 0.05
  dp_alwd= 100.0 # Pa
  labdaw=createFaceVariable(m,1.0) # initialize
  # prog_1=ProgressThresh(t_end)
  while (t<t_end)
    print("progress is $((t/t_end*100)) %", "\u1b[1G")
    # for i=1:5
      error_p = 1e5
      error_sw = 1e5
      # Implicit loop
  #     while ((error_p>eps_p) || (error_sw>eps_sw))
      while(true)
          # calculate parameters
          pgrad = gradientTerm(p)
          sw_face = upwindMean(sw, -labdaw.*pgrad) # average value of water saturation
          sw_ave=arithmeticMean(sw)
          # solve for pressure at known Sw
          labdao = lo.*faceEval(s -> kro.(s, kro0, sor, swc, no), sw_face)
          labdaw = lw.*faceEval(s -> krw.(s, krw0, sor, swc, nw), sw_face)
          labda = labdao+labdaw
          # compute [Jacobian] matrices
          Mdiffp1 = diffusionTerm(-labda)
          Mbcp, RHSbcp = boundaryConditionTerm(BCp)
          RHS1 = RHSbcp # with capillary
          p_new=solvePDE(m, Mdiffp1+Mbcp, RHS1)

          # solve for Sw
          pgrad = gradientTerm(p_new)
          uw=-labdaw.*pgrad
          Mbcsw, RHSbcsw = boundaryConditionTerm(BCs)
          RHS_sw=-divergenceTerm(uw)
          sw_new=solveExplicitPDE(sw_old, dt, RHS_sw, BCs, phi)

          error_p = maximum(abs((internalCells(p_new)-internalCells(p))./internalCells(p_new)))
          error_sw = maximum(abs(internalCells(sw_new)-internalCells(sw)))
          dt_new=dt*min(dp_alwd/error_p, dsw_alwd/error_sw)

          # assign new values of p and sw
          if error_sw>dsw_alwd
              dt=dt*(dsw_alwd/error_sw)
              # print("new time step = $dt \n")
          else
              t=t+dt
              p = copyCell(p_new)
              sw = copyCell(sw_new)
              p_old = copyCell(p)
              sw_old = copyCell(sw)
              dt=min(dt*(dsw_alwd/error_sw), 15*dt, t_end-t-eps())
              break
          end
      end

      push!(rec_fact, (oil_init-domainInt(1-sw))/oil_init)
      push!(t_day, t)
      push!(dp_hist, 0.5*sum(p.value[1:2])-p0)
      # print(t)
      # GR.plot(sw.value[2:end-1])
      # GR.imshow(sw.value[2:end-1,2:end-1])
      # visualizeCells(1-sw)
      # GR.plot(t_day/3600/24, rec_fact)
      # xlabel('time [day]')
      # ylabel('recovery factor')
      # title([num2str(t/3600/24) ' day']) drawnow
  end
  dp_hist[1] = dp_hist[2]
  t_day, rec_fact, dp_hist, sw
end # imb_impes

max_stars_repo_name: simulkade/peteng

max_issues_repo_name: simulkade/peteng

max_forks_repo_name: simulkade/peteng


text:
__precompile__()

module Common

using UnicodePlots
using JLD2, FileIO
using Flux

const WIDTH=50
const HEIGHT=8

function addmachine(name, count, tunnel::Bool=false)
   #topo=:all_to_all
   topo=:master_slave
   println("ADDING $name to cluster")
   if tunnel
      addprocs([(name, count)]; max_parallel=1, topology=topo, tunnel=true)
   else
      addprocs([(name, count)]; max_parallel=1, topology=topo)
   end
   println("WORKERS: ", nworkers())
end

#samples.ctrl[:,1,:]   .= 0.0 # x0 for ou process
#samples.ctrl[:,end,:] .= 0.0 # μ for ou process
#for k=1:K
#   Common.ou_process!(@view(samples.ctrl[:,:,k]),s.Σ[1],2.0,1/T)
#end

# ornstein-uhlenbeck process for ctrl noise for stochastic search
# μ  = target value of x[end]; get some random target?
# σ  = noise. sure.
# θ  = how quickly to go to μ; fights against σ
# dt = dt of model * skip; dt of environment
# x0 = initial value of x[1]
function ou_process!(x::AbstractMatrix{Float64},
                     σ::Float64=1.0, θ::Float64=2.0,
                     dt::Float64=1/size(x,2))
   nu = size(x,1)
   T  = size(x,2)-1
   W  = zeros(nu)
   x0 = copy(x[:,1])
   μ  = copy(x[:,end])
   for (i,t) in enumerate(0.0:dt:T*dt)
      ex = exp(-θ*t)
      for j=1:nu
         x[j,i] = x0[j]*ex + μ[j]*(1.0-ex) + σ*ex*W[j]
         W[j] += exp(θ*t)*sqrt(dt)*randn()
      end
   end
end

function ou_process!(x::AbstractVector{Float64}, x0::Float64,
                     μ::Float64=(2rand()-1),
                     σ::Float64=1.0, θ::Float64=2.0,
                     dt::Float64=1/length(x))
   T = length(x)-1
   W = 0.0
   for (i,t) in enumerate(0.0:dt:T*dt)
      ex = exp(-θ*t)
      x[i] = x0*ex + μ*(1.0-ex) + σ*ex*W
      W += exp(θ*t)*sqrt(dt)*randn()
   end
end

function Rtau(x::AbstractVector{Float64}, gamma::Float64)::Float64
   T = length(x)
   R = 0.0
   for t=1:T
      R += gamma^(t-1) * x[t]
   end
   return R
end

# X: TxK
function Rtau(x::AbstractMatrix{Float64}, gamma::Float64)::Vector{Float64}
   T, K = size(x)
   R = zeros(K)
   for t=1:T
      g = gamma^(t-1)
      for k=1:K
         R[k] += g * x[t,k]
      end
   end
   return R
end

function plotlines(iter::Integer,title::String,lines::AbstractMatrix)
   yin = minimum(lines)
   yax = maximum(lines)
   p = lineplot(lines[1,:],
                xlim=[0, size(lines,2)],
                ylim=[yin, yax],
                title=title,
                width=WIDTH, height=HEIGHT)
   for l=2:size(lines, 1)
      lineplot!(p, lines[l,:])
   end

   display(p)
end

function plotlines(iter::Integer,title::String,lines::AbstractVector)
   yin = minimum(lines)
   yax = maximum(lines)
   p = lineplot(lines,
                xlim=[0, size(lines,2)],
                ylim=[yin, yax],
                title=title,
                width=WIDTH, height=HEIGHT)

   display(p)
end

function plotlines(iter::Integer,title::String,lines...)
   if length(lines) < 1
      error("plotlines needs data and title")
   end

   n,m = extrema(lines[1][1][:])
   numP = length(lines[1][1][:])
   #r = collect(Int, round.(linspace(1, numP, min(100, numP))))
   r = (numP-min(200, numP)+1):numP
   for l in lines[2:end]
      n2, m2 = extrema(l[1][r])
      if n2 < n n = n2 end
      if m2 > m m = m2 end
   end

   p = lineplot(lines[1][1][r],
                title=title,
                ylim=[n, m],
                width=WIDTH, height=HEIGHT, name=lines[1][2])
   for l in lines[2:end]
      lineplot!(p, l[1][r], name=l[2])
   end

   display(p)
end

function plotexpmt(file)
   d = load(file)
   plotlines(length(d["stocR"]),"Reward",
             (d["stocR"],"Stoc"),
             (d["meanR"],"Mean"))
   return d
end

function makesavedir(dir_name::String, exp_file::String,
                     m_file::String, mesh_file::String="",
                     overwrite::Bool=false)
   if isdir(dir_name) == false
      mkdir(dir_name)
   else
      if overwrite == false
         val = 1
         while isdir("$(dir_name)_$(val)") == true
            val += 1
         end
         dir_name = "$(dir_name)_$(val)"
         mkdir(dir_name)
      end
   end
   # copy file to save directory
   filename = basename(exp_file)
   cp(exp_file, "$(dir_name)/$(filename)", force=overwrite)
   @info("Saving experiment $exp_file")
   open(exp_file) do f
      for l in eachline(f)
         if startswith(l, "include")
            includefile = basename(split(l, "\"")[2])
            @info("Including $includefile")
            cp(dirname(exp_file)*"/"*includefile, "$(dir_name)/$(includefile)", force=overwrite)
         end
      end
   end
   filename = basename(m_file)
   cp(m_file, "$(dir_name)/$(filename)", force=overwrite)
   @info("Model file $m_file")
   if (mesh_file != "")
      if isdir(mesh_file)
         symlink(mesh_file, dir_name*"/meshes")
      else
         symlink(dirname(mesh_file), dir_name*"/meshes")
      end
   end
   return dir_name
end

function save(expresults::String, data::Dict{Symbol, Vector{Float64}})
    jldopen(expresults, "w") do file
        for d in data
           write(file, "$(d[1])", d[2])
        end
    end
end

function copymodel(src::Flux.Chain)
   i = 1
   DT = eltype(params(m)[1].data)
   dst = Vector{DT}( sum(length.(params(src))) )
   copymodel!(dst, src)
   return dst
end

function copymodel!(dst::AbstractVector, src::Flux.Chain)
   i = 1
   for l in Tracker.data.(Flux.params(src))
      r = i:(i+length(l)-1)
      dst[r] = l
      i += length(l)
   end
end

function copymodel!(dst::Flux.Chain, src::AbstractVector)
   i = 1
   for l in Tracker.data.(Flux.params(dst))
      r = i:(i+length(l)-1)
      l[:] = src[r]
      i += length(l)
   end
end

function minibatch(nsamples::Int, szbatch::Int,
                   nbatch = div(nsamples, szbatch) )
   return [ randperm(nsamples)[1:szbatch] for i=1:nbatch ]
end

function vecvectomat(vv)
   if length(vv) > 0
      return reshape(vcat(vv[:]...), length(vv[1]), length(vv))
   else
      return Matrix{eltype(eltype(vv))}(0,0)
   end
end

function vecvectomat!(m, vv)
   r, c = size(m)
   @assert c == length(vv)
   @assert r == length(vv[1]) # assumes all elements of vv are same size
   for i=1:length(vv)
      m[:,i] = vv[i]
   end
end

end


text:
using Caesar

robotId = "Munster"
sessionId = "Test"

fg = initfg(sessionname=sessionId, robotname=robotId)

# also add a PriorPose2 to pin the first pose at a fixed location
addVariable!(fg, :x0, Pose2) # , tags=[:POSE]
addFactor!(fg, [:x0], PriorPose2(MvNormal(zeros(3), 0.01*Matrix(LinearAlgebra.I, 3,3))) )

# Drive around in a hexagon
for i in 0:5
  psym = Symbol("x$i")
  nsym = Symbol("x$(i+1)")
  addVariable!(fg, nsym, Pose2) # , tags=[:VARIABLE;:POSE]
  addFactor!(fg, [psym;nsym], Pose2Pose2( MvNormal([10.0;0;pi/3], 0.01*Matrix(LinearAlgebra.I, 3,3)) ), autoinit=true )
  # Pose2Pose2_NEW(MvNormal([10.0;0;pi/3], diagm([0.1;0.1;0.1].^2)))
end

response = addVariable!(fg, :l1, Point2("LANDMARK"))


text:
using SixDOF
import ForwardDiff
import ReverseDiff
# using FLOWMath: centraldiff

function wrapper(x)
    num = 2
    type = COCOUNTER

    mp = MassProp(x[1:5]...)
    ref = Reference(x[6:8]...)
    sd = StabilityDeriv(x[9:45]...)
    propulsion = MotorPropBatteryDataFit(x[46], x[47], x[48], x[49], x[50], x[51], x[52], num, type, x[53:56]...)
    atm = ConstantAtmosphere(x[57:59], x[60:62], x[63], x[64], x[65])
    controller = ConstantController(x[66:70]...)
    inertial = UniformGravitationalField()
    Vinf = 10.0
    alpha = 3.0*pi/180
    s0 = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, Vinf*cos(alpha), 0.0, Vinf*sin(alpha), 0.0, 0.0, 0.0]
    time = 1.0
    params = mp, ref, sd, propulsion, inertial, atm, controller
    ds = zeros(eltype(x[1]), 12)
    sixdof!(ds, s0, params, time)
    return ds
end

x = zeros(71)

x[1] = 1.56
x[2] = 0.1147
x[3] = 0.0576
x[4] = 0.1712
x[5] = 0.0015
x[6] = 0.2589
x[7] = 1.4224
x[8] = 0.3302
x[9] = 0.09167 # Zero-alpha lift
x[10] = 3.5016  # lift curve slope
x[11] = 2.8932 # Pitch rate derivative
x[12] = 0.0 # Mach derivative
x[13] = 0.0  # flaps derivative
x[14] = 0.2724  # elevator derivative
x[15] = 20*pi/180
x[16] = 0.01631  # zero-lift drag coerff
x[17] = 10.0  # velocity corresponding to Reynolds number of CD0
x[18] = -0.2  # exponent in Reynolds number calc
x[19] = 0.8  # Oswald efficiency
x[20] = 0.7  # crest critcal Mach number
x[21] = 0.0  # flaps
x[22] = 0.3045  # elevator
x[23] = 0.0  # aileron
x[24] = 0.0  # rudder
x[25] = -0.07359 # Sideslip derivative
x[26] = 0.0  # Roll rate derivative
x[27] = 0.0 # Yaw rate derivative
x[28] = 0.0 # Roll control (aileron) derivative
x[29] = 0.0 # Yaw control (rudder) derivative
x[30] = -0.02854  # Sideslip derivative
x[31] = -0.3209  # Roll rate derivative
x[32] = 0.03066  # Yaw rate derivative
x[33] = 0.1682  # Roll (aileron) control derivative
x[34] = 0.0  #Yaw (rudder) control derivative
x[35] = -0.02338 # Zero-alpha pitch
x[36] = -0.5675 # Alpha derivative
x[37] = -1.3990 # Pitch rate derivative
x[38] = 0.0
x[39] = 0.0
x[40] = -0.3254 # Pitch control derivative
x[41] = -0.00040  # Slideslip derivative
x[42] = -0.01297  # Roll rate derivative
x[43] = -0.00434  # Yaw rate derivative
x[44] = -0.00328  # Roll (aileron) control derivative
x[45] = 0.0  # Yaw (rudder) control derivative
x[46] = 0.11221
x[47] = -0.13803
x[48] = -0.047394
x[49] = 0.0062
x[50] = 0.00314
x[51] = -0.015729
x[52] = 10*0.0254
x[53] = 0.5
x[54] = 2500.0 * pi/30
x[55] = 0.3
x[56] = 8.0
x[57:59] .= [0.0, 0.0, 0.0]
x[60:62] .= [0.0, 0.0, 0.0]
x[63] = 1.2682
x[64] = 300.0
x[65] = 9.81
x[66:70] .= [-3.0*pi/180, 0.0, 0.0, 0.0, 0.8]

f = wrapper(x)

J = ForwardDiff.jacobian(wrapper, x)
J2 = ReverseDiff.jacobian(wrapper, x)
# J3 = centraldiff(wrapper, x)

# diff = J - J3
# idx = (diff .!= 0.0)
# rerr = diff[idx]./J3[idx]

# @test maximum(abs.(rerr)) < 1e-5

# diff = J2 - J3
# idx = (diff .!= 0.0)
# rerr = diff[idx]./J3[idx]

# @test maximum(abs.(rerr)) < 1e-5

diff = J - J2
idx = (diff .!= 0.0)
rerr = diff[idx]./J[idx]
# println(maximum(abs.(rerr)))
@test maximum(abs.(rerr)) < 1e-7

# @code_warntype wrapper(x)


text:
import NaiveNASlib
import NaiveNASlib:OutputsVertex
import InteractiveUtils:subtypes

@testset "Mutation vertices" begin

    @testset "OutputsVertex" begin
        iv = OutputsVertex(InputVertex(1))

        @test inputs(iv) == []
        @test outputs(iv) == []

        @test issame(iv, clone(iv))

        cv = OutputsVertex(CompVertex(x -> 2x, iv))
        NaiveNASlib.init!(cv, base(cv))

        @test inputs(cv) == [iv]
        @test outputs(iv) == [base(cv)]

        ivc = clone(iv)
        cvc = clone(cv, ivc)
        NaiveNASlib.init!(cvc,base(cvc))

        @test issame(iv, ivc)
    end

    @testset "InputSizeVertex" begin
        iv = InputSizeVertex(InputVertex(1), 3)

        @test nout(iv) == 3
        @test nin(iv) == []
        @test inputs(iv) == []
        @test outputs(iv) == []

        @test issame(iv, clone(iv))

        cv =  OutputsVertex(CompVertex(x -> 2x, iv))
        NaiveNASlib.init!(cv, base(cv))

        @test inputs(cv) == [iv]
        @test outputs(iv) == [base(cv)]

        ivc = clone(iv)
        cvc = clone(cv, ivc)
        NaiveNASlib.init!(cvc,base(cvc))

        @test issame(iv, ivc)
    end

    @testset "Pretty printing" begin
        import NaiveNASlib: show_less

        @testset "OutputsVertex" begin
            iv = OutputsVertex(InputVertex(1))
            cv = OutputsVertex(CompVertex(identity, iv))
            NaiveNASlib.init!(cv, base(cv))
            sv = OutputsVertex(CompVertex(hcat, iv, cv))
            NaiveNASlib.init!(sv, base(sv))

            @test showstr(show, iv) == "InputVertex(1), outputs=[CompVertex(identity), CompVertex(hcat)]"
            @test showstr(show, cv) == "CompVertex(identity), inputs=[InputVertex(1)], outputs=[CompVertex(hcat)]"
            @test showstr(show, sv) == "CompVertex(hcat), inputs=[InputVertex(1), CompVertex(identity)], outputs=[]"
        end

        @testset "InputSizeVertex" begin
            iv = InputSizeVertex("iv", 3)
            @test showstr(show_less, iv) == "iv"
        end

        @testset "NamedTrait" begin

            v1 = MutationVertex(CompVertex(identity, InputSizeVertex("input1", 3)),  NoOp(), NamedTrait(SizeInvariant(), "mv"))

            @test showstr(show_less, v1) == "mv"

            v2 = MutationVertex(CompVertex(+, v1, InputSizeVertex("input2", 3)), NoOp(), NamedTrait(SizeInvariant(), "sv"))

            @test showstr(show, v2) == "MutationVertex(CompVertex(+), inputs=[mv, input2], outputs=[], NoOp(), NamedTrait(SizeInvariant(), \"sv\"))"
        end

        @testset "Info strings" begin
            v1 = InputSizeVertex("v1", 3)
            v2 = MutationVertex(CompVertex(identity, v1), IoSize(nout(v1), 5), NamedTrait(SizeAbsorb(), "v2"))
            v3 = MutationVertex(CompVertex(identity, v1,v2), IoSize([3, 5], 8), NamedTrait(SizeStack(), "v3"))
            v4 = MutationVertex(CompVertex(identity, v3), IoSize(nout(v3), nout(v3)), NamedTrait(SizeInvariant(), "v4"))

            @test infostr(NameInfoStr(), v1) == "v1"
            @test infostr(NameInfoStr(), v2) == "v2"
            @test infostr(NameInfoStr(), v3) == "v3"
            @test infostr(NameInfoStr(), v4) == "v4"

            @test infostr(NinInfoStr(), v1) == "N/A"
            @test infostr(NinInfoStr(), v2) == "3"
            @test infostr(NinInfoStr(), v3) == "3, 5"
            @test infostr(NinInfoStr(), v4) == "8"

            @test infostr(SizeInfoStr(), v1) == "nin=[N/A], nout=[3]"
            @test infostr(SizeInfoStr(), v2) == "nin=[3], nout=[5]"
            @test infostr(SizeInfoStr(), v3) == "nin=[3, 5], nout=[8]"
            @test infostr(SizeInfoStr(), v4) == "nin=[8], nout=[8]"

            @test infostr(OutputsInfoStr(), v1) == "[v2, v3]"
            @test infostr(OutputsInfoStr(), v2) == "[v3]"
            @test infostr(OutputsInfoStr(), v3) == "[v4]"
            @test infostr(OutputsInfoStr(), v4) == "[]"

            @test infostr(NameAndIOInfoStr(), v1) == "v1, inputs=[], outputs=[v2, v3]"
            @test infostr(NameAndIOInfoStr(), v2) == "v2, inputs=[v1], outputs=[v3]"
            @test infostr(NameAndIOInfoStr(), v3) == "v3, inputs=[v1, v2], outputs=[v4]"
            @test infostr(NameAndIOInfoStr(), v4) == "v4, inputs=[v3], outputs=[]"

            @test infostr(MutationTraitInfoStr(), v1) == "Immutable()"
            @test infostr(MutationTraitInfoStr(), v2) == "NamedTrait(SizeAbsorb(), v2)"
            @test infostr(MutationTraitInfoStr(), v3) == "NamedTrait(SizeStack(), v3)"
            @test infostr(MutationTraitInfoStr(), v4) == "NamedTrait(SizeInvariant(), v4)"

            @test infostr(MutationSizeTraitInfoStr(), v1) == "Immutable()"
            @test infostr(MutationSizeTraitInfoStr(), v2) == "SizeAbsorb()"
            @test infostr(MutationSizeTraitInfoStr(), v3) == "SizeStack()"
            @test infostr(MutationSizeTraitInfoStr(), v4) == "SizeInvariant()"

            @test infostr(FullInfoStr(), v1) == "v1, inputs=[], outputs=[v2, v3], nin=[N/A], nout=[3], Immutable()"
            @test infostr(FullInfoStr(), v2) == "v2, inputs=[v1], outputs=[v3], nin=[3], nout=[5], SizeAbsorb()"
            @test infostr(FullInfoStr(), v3) == "v3, inputs=[v1, v2], outputs=[v4], nin=[3, 5], nout=[8], SizeStack()"
            @test infostr(FullInfoStr(), v4) == "v4, inputs=[v3], outputs=[], nin=[8], nout=[8], SizeInvariant()"
        end

    end
end


text:
# Compact Hilbert Indices by Chris Hamilton. Technical Report CS-2006-07.
# 6059 University Ave., Halifax, Nova Scotia, B3H 1W5, Canada.
#
# In this document, the notation is:
# ∧ = & for and
# ∨ = | for or
# ⊻ = for xor
# ▷ = >> for shift right
# ◁ = << for shift left
# ⌈x⌉ = ceil(x)
# ⌊x⌋ = floor(x)

# hi_g(i) is the direction in which brgc changes.
# It is also tells us the axis along which the exit of one
# subcube meets the entrance of the next subcube.
hi_g(i) = trailing_set_bits(i)
# This is the definition of g, to compare against for testing.
hi_g_orig(i) = floor(typeof(i), log2(brgc(i) ⊻ brgc(i + 1)))

# e is the entry vertex of the ith sub-hypercube in a Gray code
# ordering of sub-hypercubes. If e is 0b011, then this is z-y-x
# within the subcube of 0 in the z, 1 in the y, 1 in the x.
# e = gc(2⌊(i-1)/2⌋) = gc(2*floor((i-1)/2)).
# domain is i=0 to i=2^n - 1.
function hi_e(i::Integer)
    i == zero(i) && return zero(i)
    brgc((i - one(i)) & ~one(i))
end

# This is verbatim from the article, for comparison.
function hi_e_orig(i::Base.BitInteger)
    i == zero(i) && return zero(i)
    brgc(2 * floor(typeof(i), (i - 1) / 2))
end

"""
f is the exit vertex of the ith sub-hypercube in a Gray code
ordering of the sub-hypercubes.
Corllary 2.7 on pg. 12 says:
    hi_f(i, n) = hi_e(i) ⊻ hi_d(i, n)
"""
hi_f(i, n) = hi_e(one(i)<<n - one(i) - i) ⊻ (one(i)<<(n-1))
# hi_f(i, n) = hi_e(i) ⊻ hi_d(i, n)
#hi_f(i, n) = hi_e(i) ⊻ (one(i)<<hi_d(i, n))

"""
page 12. directions
d(i) = 0 if i=0
    = g(i-1) mod n if i=0 mod 2
    = g(i) mod n if i=1 mod 2.
Domain is 0 <= i <= 2^n - 1.
"""
function hi_d(i, n)
    i == zero(i) && return zero(i)
    ((i & 1) == 0) && return hi_g(i - one(i)) % n
    hi_g(i) % n
end

# bitrotate is a left shift, so negate d+1.
# T_{(e,d)}(b), so read right-to-left.
# The paper means to bit rotate over the n bits that are in use,
# not all n bits. This is not the usual bitrotate!
# This had a +1 in the paper.
hi_T(b, d, e, n) = rotateright(b ⊻ e, d, n)

# The author's code differs with his paper. It doesn't add one.
# https://github.com/pdebuyl/libhilbert/blob/master/include/Hilbert/Algorithm.hpp
hi_T_inv(b, d, e, n) = rotateleft(b, d, n) ⊻ e

# Lemma 2.12, page 15.
# Is it -2 or -1?
# function hi_T_inv(b, d, e, n)
#     hi_T(b, n - d - one(b) - one(b), bitrotate(e, -(d + one(b)), n), n)
# end

"""
    ith_bit_of_indices(n, p, i)

Given `n` indices in `p`, take the `i`-th bit of each one
and place it so that the first vector's value is at the
0-th place, the second vector's value at the 1st place, and so-on.
`i` is zero-indexed.
"""
function ith_bit_of_indices(n, p, i)
    l = zero(eltype(p))
    for j = 0:(n - 1)
        b = (p[j + 1] & (one(eltype(p))<<i)) >>i
        l |= b << j
    end
    l
end

# hamilton version of getting the ith bit, zero-indexed i.
function get_location(p::Vector{T}, i) where {T}
    l = zero(T)
    for j = eachindex(p)
        if (p[j] & (one(T) << i)) != 0
            l |= (one(T) << (j - 1))
        end
    end
    l
end

"""
    set_indices_bits(p, l, m, i)

Given bits in `l`, set each index in the array `p` according
to the bit in `l`. Set the i-th bit of each index in `p`.
"""
function set_indices_bits!(p, l, m, i)
    for j in 0:(m - 1)
        v = (l & (1<<j)) >>j
        p[j + 1] = p[j + 1] | (v<<i)
    end
end

"""
Return `m` bits from the `i`-th set of `m` bits
in the integer `h`. `i` is zero-based.
"""
function bitrange(h, n, i)
    v = zero(h)
    for j in 0:(n - 1)
        v = v | (h & (1<<(i*n + j)))
    end
    v
end

function update1(l, t, w, n, e, d)
    e = l ⊻ (one(l) << d)
    d += one(d) + first_set_bit(t)
    while d >= n
        d -= n
    end
    if (w & one(w)) == zero(w)
        if d == zero(d)
            e ⊻= one(e) << (n - 1)
        else
            e ⊻= one(e) << (d - 1)
        end
    end
    e, d
end

function update2(l, t, w, n, e, d)
    e = l
    e ⊻= (one(e) << d)
    d += one(d) + first_set_bit(t)
    while d >= n
        d -= n
    end
    e, d
end

"""
Algorithm.hpp: _coordsToIndex
"""
function hilbert_index_paper!(::Type{T}, n, m, p, ds) where {T <: Integer}
    h = zero(T)  # hilbert index
    e = zero(T)  # entry point
    d = one(T)  # direction
    nmask = fbvn1s(T, n)
    for i = (m - 1):-1:0  # i is an index. Can be any type.
        ds[i + 1] = d
        l = T(ith_bit_of_indices(n, p, i))
        t = hi_T(l, d, e, n)
        w = t
        if i < m - 1
            w ⊻= (one(w) << (n - 1))
        end
        # Concatenate to the index
        h |= (w & nmask) << (i * n)
        e, d = update2(l, t, w, n, e, d)
    end
    brgc_inv(h)
end

"""
    hilbert_index(n, m, p)

Hilbert index for an `n`-dimensional vector `p`, with each
component of extent less than 2^m. Algorithm 1 of Hamilton and
Rau-Chaplin.
"""
function hilbert_index_paper(::Type{T}, n, m, p) where {T <: Integer}
    ds = zeros(Int, m)
    hilbert_index_paper!(T, n, m, p, ds)
end

function hilbert_index_inv_paper!(::Type{T}, n, m, h, p) where {T <: Integer}
    e = zero(T)
    d = one(T)
    l = zero(T)
    p .= zero(eltype(p))
    nmask = fbvn1s(T, n)
    for i = (m - 1):-1:0
        w = (h >> (i * n)) & nmask
        t = brgc(w)
        l = hi_T_inv(t, d, e, n)
        set_indices_bits!(p, l, n, i)
        e, d = update1(l, t, w, n, e, d)
    end
end

# Make d an argument because it determines the initial direction and maybe
# it isn't initialized correctly. Maybe d and e need different initial values.
function hilbert_index(n, m, p, d = zero(eltype(p)))
    h = zero(eltype(p))  # hilbert index
    e = zero(eltype(p))  # entry point
    for i = (m - 1):-1:0  # i is an index. Can be any type.
        l = ith_bit_of_indices(n, p, i)
        # @show l
        l = hi_T(l, d, e, n)  # n or m?
        w = brgc_inv(l)
        h = (h << n) | w
        e = e ⊻ rotateleft(hi_e(w), d + one(d), n)
        d = (d + hi_d(w, n) + one(d)) % n  # n or m for hi_d?
    end
    h
end

function hilbert_index_inv(n, m, h)
    e = zero(h)  # entry point
    d = zero(h)  # direction
    p = zeros(typeof(h), m)
    for i = (m - 1):-1:0  # i is an index. Can be any type.
        w = bitrange(h, n, i)
        l = brgc(w)
        l = hi_T_inv(l, d, e, n)
        set_indices_bits!(p, l, m, i)
        e = e ⊻ rotateleft(hi_e(w), d + one(d), n)
        d = (d + hi_d(w, n) + one(d)) % n  # n or m for hi_d?
    end
    p
end

"""
Calculates a bit-mask for excluding Gray code values when the level
is below the resolution of the dimension.

Bits free at iteration `i`. Vector of resolutions, in powers of two,
`m`. Dimensions `n` so that `length(m)==n`. `i` is the level down
in the Hilbert curve. `d` is the direction, called `hi_d` above.

Returns both the mask and the number of bits set in the mask.
"""
function extract_mask(m::Vector, n, d, i)
    T = UInt64
    mask = zero(T)
    b = 0
    jm = one(T)
    j = d
    while true
        if m[j + 1] > i
            mask |= jm
            b += 1
        end
        jm <<= one(T)
        if jm == zero(T)
            jm = one(T)
        end
        j += 1
        if j == n
            j = 0
        end
        if j == d
            break
        end
    end
    mask, b
end

function extract_mask_paper(m::Vector, n, d, i)
    T = UInt64
    mask = zero(T)
    b = 0
    for j = (n-1):-1:0
        mask <<= one(T)
        jn = (j + d) % n
        if m[jn + 1] > i
            mask |= one(T)
            b += 1
        end
    end
    mask, b
end

"""
From GrayCodeRank.hpp: compactIndex.
"""
function compact_index(ms::Vector, ds::Vector, n, m, h::T) where {T}
    hc = zero(T)
    hr = 0
    hcr = 0
    hm = one(T)
    hcm = one(T)
    for i = 0:(m-1)
        j = ds[i + 1]
        while true
            if ms[j + 1] > i
                if hr > 0
                    error("hr on next rack")
                end
                if ((h & hm) != 0)
                    if hcr == 0
                        hc |= hcm
                    else
                        error("should only be one rack")
                    end
                end
                hcm <<= 1
                if hcm == 0
                    hcm = one(T)
                    hcr += 1
                end
            end
            j += 1
            if j == n
                j = 0
            end
            hm <<= 1
            if hm == zero(T)
                hm = one(T)
                hr += 1
            end
            if j == ds[i + 1]
                break
            end
        end
    end
    hc
end

function coords_to_compact_index(::Type{T}, p::Vector{A}, ms::Vector, n) where {A,T}
    m = maximum(ms)
    M = sum(ms)
    mn = m * n
    ds = zeros(Int, m)
    h = hilbert_index_paper!(T, n, m, p, ds)
    compact_index(ms, ds, n, m, h)
end

function compact_index_to_coords!(p::Vector{A}, ms, n, hc::T) where {A, T}
    m = maximum(ms)
    M = sum(ms)
    bit_cnt = 8 * sizeof(T)

    e = zero(T)
    d = one(T)
    l = zero(T)
    p .= zero(A)
    # work from most significant bit to least significant bit
    for i = (m - 1):-1:0
        mask, b = extract_mask(ms, n, d, i)
        # rotateright(val, shift_cnt, total_bits)
        ptrn = rotateright(e, d, n)

        # Get the Hilbert index bits.
        M -= b
        # b bits from hc at index M, into r
        r = get_bits(hc, b, M)

        t, w = brgc_rank_inv(mask, ptrn, r, n, b)

        l = hi_T_inv(t, d, e, n)
        set_indices_bits!(p, l, n, i)
        e, d = update1(l, t, w, n, e, d)
    end
end

"""
    SpaceGray(b, n)
    SpaceGray(::Type{T}, b, n)

This is an n-dimensional Hilbert curve where all `n` dimensions
must have `b` bits in size. It was described in the same paper
and examples as the `Compact` algorithm.
"""
struct SpaceGray{T} <: HilbertAlgorithm{T}
    b::Int
    n::Int
end

axis_type(gg::SpaceGray) = large_enough_unsigned(gg.b)

function SpaceGray(b, n)
    ttype = large_enough_unsigned(b * n)
    SpaceGray{ttype}(b, n)
end

function SpaceGray(::Type{T}, b, n) where {T}
    SpaceGray{T}(b, n)
end

function encode_hilbert_zero(g::SpaceGray{T}, X::Vector)::T where {T}
    hilbert_index_paper(T, g.n, g.b, X)
end

function decode_hilbert_zero!(g::SpaceGray{T}, X::Vector, h::T) where {T}
    hilbert_index_inv_paper!(T, g.n, g.b, h, X)
end

"""
    Compact(ms::Vector{Int})
    Compact(::Type{T}, ms::Vector{Int})

This algorithm is n-dimensional and permits dimensions to use different
numbers of bits, specified in the `ms` vector. The type `T` is an
optional data type for the Hilbert index. It should be greater than or
equal to the sum of the bits.

This algorithm comes from three sources:

* A technical report, "Compact Hilbert Indices" by Chris Hamilton.
  Technical Report CS-2006-07. 6059 University Ave., Halifax, Nova
  Scotia, B3H 1W5, Canada. This report is informative but has many
  errors.

* A paper by Hamilton and Rau-Chaplin, "Compact Hilbert Indices for
  Multi-Dimensional Data," 2007. Nice paper. Also wrong.

* The [libhilbert source code](https://github.com/pdebuyl/libhilbert)
  is a copy of Hamilton's work and has many corrections. This, ultimately,
  lead to the working code.
"""
struct Compact{T} <: HilbertAlgorithm{T}
    ms::Vector{Int}
    n::Int
end

axis_type(gg::Compact) = large_enough_unsigned(maximum(gg.ms))

function Compact(ms::Vector{Int})
    n = length(ms)
    b = maximum(ms)
    ttype = large_enough_unsigned(b * n)
    Compact{ttype}(ms, n)
end

function Compact(::Type{T}, ms::Vector{Int}) where {T}
    Compact{T}(ms, length(ms))
end

function encode_hilbert_zero(g::Compact{T}, X::Vector)::T where {T}
    coords_to_compact_index(index_type(g), X, g.ms, g.n)
end

function decode_hilbert_zero!(g::Compact{T}, X::Vector, h::T) where {T}
    compact_index_to_coords!(X, g.ms, g.n, h)
end


text:
module FooTests

import Foo
import Example

Foo.greet()
Example.hello("human")

end


text:
#----------------------------------------------------------------------
# Boids algorithm visualization
# Based on pseudocode at http://www.kfish.org/boids/pseudocode.html
# Simulate flocking birds using simple rules.
# Initial version developed and contributed by Iain Dunning (@IainNZ)
# as part of the JuliaCon 2015 workshop
# TODO:
# - Add interactivity for parameters that govern behavior.
# - Add a force that attracts to mouse pointer.
#----------------------------------------------------------------------
using Reactive, GeometryTypes, GLAbstraction, GLPlot, GLVisualize, Colors, ColorTypes
GLPlot.init()
#----------------------------------------------------------------------
typealias Position Point{2,Float32}
typealias Velocity Vec{2,Float32}

immutable Boids
    position::Vector{Position}
    velocity::Vector{Velocity}
end

# Create new boid at random location and velocity
Boids(n=200) = Boids([rand(Position) for i=1:n], [rand(Velocity)/1000 for i=1:n])

#----------------------------------------------------------------------

function simulate!(t, boids)
    len = Float32(length(boids.position))
    boid_center = sum(boids.position)/len
    @inbounds for (i, vel) in enumerate(boids.velocity)
        boids.velocity[i] = vel + Vec(boid_center - boids.position[i])/600
    end

    # Force 2: Avoid others
    @inbounds for (i, boidpos) in enumerate(boids.position)
        avoidance = zero(Velocity)
        for (j, other_boidpos) in enumerate(boids.position)
            i == j && continue
            if norm(boidpos - other_boidpos) <= 0.2
                avoidance -= Vec(other_boidpos - boidpos)
            end
        end
        bvel = boids.velocity[i]
        boids.velocity[i] = bvel + (avoidance/1000f0)
    end
    @inbounds for (i, bvel) in enumerate(boids.velocity)
        perceived_vel = zero(Velocity)
        for (j, other_boidvel) in enumerate(boids.velocity)
            i == j && continue
            perceived_vel += other_boidvel
        end
        perceived_vel /= len-1f0
        boids.velocity[i] = bvel + (perceived_vel - bvel)/900f0
    end
    # Limit max velocity
    MAXVEL = 0.5f0
    @inbounds for (i, vel) in boids.velocity
        absvel = norm(vel)
        if absvel >= MAXVEL
            boids.velocity[i] = vel / absvel*MAXVEL
        end
    end
    # Update positions
    @inbounds for (i, vel) in enumerate(boids.velocity)
        boids.position[i] = boids.position[i] + Point(vel)
    end
    boids.position, boids.velocity
end
Base.clamp{T}(x::T) = clamp(x, zero(T), one(T))

to_color(velocities, mul) = RGBA{U8}[RGBA{U8}(clamp(velocity[1]*mul), clamp(velocity[2]*mul), 1.0, 0.8) for velocity in velocities]
function main()
    # Create population of boids
    boids = Boids()
    pv = map(simulate!, bounce(0:1000), Signal(boids))
    positions = map(first, pv)
    velocity = map(last, pv)
    glplot(
       (Circle, positions), scale=Vec2f0(0.05),
       intensity=map(norm, velocity), color_map=GLVisualize.default(Vector{RGBA}),
       color_norm=Vec2f0(0, 2)
    )

end

main()


text:
const c = 299792458.
const mu0 = 4*pi*1e-7
const eps0 = 1/(mu0*c^2)

function Hertz_dipole (r, p, R, phi, f, t=0, epsr=1.)
    #=
    Calculate E field and B field strength of hertzian dipole(s)
    p: array of dipole moments
    R: array of dipole positions
    r: observation point
    f: frequency
    t: time
    phi: array with dipole phase angles (0..2pi)
    return: field strength at observation point r at time t (3-tuple: Ex, Ey, Ez)
    =#
    rprime=r-R  # r'=r-R
    magrprime=sqrt(sum(rprime.^2)) # |r-R|
    w=2*pi*f  # omega
    k=w/c     # wave number
    krp=k*magrprime  # k*|r-R|
    rprime_cross_p = cross(rprime, p) # (r-R) x p
    rp_c_p_c_rp = cross(rprime_cross_p, rprime) # ((r-R) x p) x (r-R)
    expfac=exp(1im*(-w*t+krp-phi))/(4*pi*eps0*epsr)
    E=expfac*(w^2/(c^2*magrprime^3)*rp_c_p_c_rp+(1/magrprime^3-w*im/(c*magrprime^2))*(3*rprime*dot(rprime,p)/magrprime^2-p))
	B=expfac/(magrprime*c^3)*(w^2*rprime_cross_p)/magrprime*(1-c/(im*w*magrprime))
    return E,B
  end

function Hertz_dipole_ff (r, p, R, phi, f, t=0, epsr=1.)
  #=
  Calculate E field and B field strength in far field of hertzian dipole(s)
  p: array of dipole moments
  R: array of dipole positions
  r: observation point
  f: frequency
  t: time
  phi: array with dipole phase angles (0..2pi)
  return: field strength at observation point r at time t (3-tuple: Ex, Ey, Ez)
  =#
  rprime=r-R  # r'=r-R
  magrprime=sqrt(sum(rprime.^2)) # |r-R|
  w=2*pi*f  # omega
  k=w/c     # wave number
  krp=k*magrprime  # k*|r-R|
  rprime_cross_p = cross(rprime, p) # (r-R) x p
  rp_c_p_c_rp = cross(rprime_cross_p, rprime) # ((r-R) x p) x (r-R)
  expfac=exp(1im*(-w*t+krp-phi))/(4*pi*eps0*epsr)
  E=expfac*(w^2/(c^2*magrprime^3)*rp_c_p_c_rp)
  B=expfac/(magrprime^2*c^3)*(w^2*rprime_cross_p)
  return E,B
end

function Hertz_dipole_nf (r, p, R, phi, f, t=0, epsr=1.)
    #=
  Calculate E field and B field strength in near field of hertzian dipole(s)
  p: array of dipole moments
  R: array of dipole positions
  r: observation point
  f: frequency
  t: time
  phi: array with dipole phase angles (0..2pi)
  return: field strength at observation point r at time t (3-tuple: Ex, Ey, Ez)
  =#
    rprime=r-R  # r'=r-R
    magrprime=sqrt(sum(rprime.^2)) # |r-R|
    w=2*pi*f  # omega
    k=w/c     # wave number
    krp=k*magrprime  # k*|r-R|
	rprime_cross_p = cross(rprime, p) # (r-R) x p
    expfac=exp(1im*(-w*t+krp-phi))/(4*pi*eps0*epsr)
    E=expfac*((1/magrprime^3-w*im/(c*magrprime^2))*(3*rprime*dot(rprime,p)/magrprime^2-p))
	B=expfac/(magrprime*c^2)*(w^2*rprime_cross_p)/magrprime*(-1/(im*w*magrprime))
    return E,B
	end

max_stars_repo_name: manuamador/Z

max_issues_repo_name: manuamador/Z

max_forks_repo_name: manuamador/Z


text:
using LinearAlgebra: checksquare
using LinearAlgebra.BLAS: gemv, gemv!, gemm!, trsm!, axpy!, ger!

#####
##### `lu`
#####

# These rules are necessary because the primals call LAPACK functions

# frule for square matrix was introduced in Eq. 3.6 of
# de Hoog, F.R., Anderssen, R.S. and Lukas, M.A. (2011)
# Differentiation of matrix functionals using triangular factorization.
# Mathematics of Computation, 80 (275). p. 1585.
# doi: http://doi.org/10.1090/S0025-5718-2011-02451-8
# for derivations for wide and tall matrices, see
# https://sethaxen.com/blog/2021/02/differentiating-the-lu-decomposition/

function frule(
    (_, ΔA), ::typeof(lu!), A::StridedMatrix, pivot::Union{Val{false},Val{true}}; kwargs...
)
    F = lu!(A, pivot; kwargs...)
    ∂factors = pivot === Val(true) ? ΔA[F.p, :] : ΔA
    m, n = size(∂factors)
    q = min(m, n)
    if m == n  # square A
        # minimal allocation computation of
        # ∂L = L * tril(L \ (P * ΔA) / U, -1)
        # ∂U = triu(L \ (P * ΔA) / U) * U
        # ∂factors = ∂L + ∂U
        L = UnitLowerTriangular(F.factors)
        U = UpperTriangular(F.factors)
        rdiv!(∂factors, U)
        ldiv!(L, ∂factors)
        ∂L = lmul!(L, tril(∂factors, -1))
        ∂U = rmul!(triu(∂factors), U)
        ∂factors .= ∂L .+ ∂U
    elseif m < n  # wide A, system is [P*A1 P*A2] = [L*U1 L*U2]
        L = UnitLowerTriangular(F.L)
        U = F.U
        ldiv!(L, ∂factors)
        @views begin
            ∂factors1 = ∂factors[:, 1:q]
            ∂factors2 = ∂factors[:, (q + 1):end]
            U1 = UpperTriangular(U[:, 1:q])
            U2 = U[:, (q + 1):end]
        end
        rdiv!(∂factors1, U1)
        ∂L = tril(∂factors1, -1)
        mul!(∂factors2, ∂L, U2, -1, 1)
        lmul!(L, ∂L)
        rmul!(triu!(∂factors1), U1)
        ∂factors1 .+= ∂L
    else  # tall A, system is [P1*A; P2*A] = [L1*U; L2*U]
        L = F.L
        U = UpperTriangular(F.U)
        rdiv!(∂factors, U)
        @views begin
            ∂factors1 = ∂factors[1:q, :]
            ∂factors2 = ∂factors[(q + 1):end, :]
            L1 = UnitLowerTriangular(L[1:q, :])
            L2 = L[(q + 1):end, :]
        end
        ldiv!(L1, ∂factors1)
        ∂U = triu(∂factors1)
        mul!(∂factors2, L2, ∂U, -1, 1)
        rmul!(∂U, U)
        lmul!(L1, tril!(∂factors1, -1))
        ∂factors1 .+= ∂U
    end
    ∂F = Composite{typeof(F)}(; factors=∂factors)
    return F, ∂F
end

function rrule(
    ::typeof(lu), A::StridedMatrix, pivot::Union{Val{false},Val{true}}; kwargs...
)
    F = lu(A, pivot; kwargs...)
    function lu_pullback(ΔF::Composite)
        Δfactors = ΔF.factors
        Δfactors isa AbstractZero && return (NO_FIELDS, Δfactors, DoesNotExist())
        factors = F.factors
        ∂factors = eltype(A) <: Real ? real(Δfactors) : Δfactors
        ∂A = similar(factors)
        m, n = size(A)
        q = min(m, n)
        if m == n  # square A
            # ∂A = P' * (L' \ (tril(L' * ∂L, -1) + triu(∂U * U')) / U')
            L = UnitLowerTriangular(factors)
            U = UpperTriangular(factors)
            ∂U = UpperTriangular(∂factors)
            tril!(copyto!(∂A, ∂factors), -1)
            lmul!(L', ∂A)
            copyto!(UpperTriangular(∂A), UpperTriangular(∂U * U'))
            rdiv!(∂A, U')
            ldiv!(L', ∂A)
        elseif m < n  # wide A, system is [P*A1 P*A2] = [L*U1 L*U2]
            triu!(copyto!(∂A, ∂factors))
            @views begin
                factors1 = factors[:, 1:q]
                U2 = factors[:, (q + 1):end]
                ∂A1 = ∂A[:, 1:q]
                ∂A2 = ∂A[:, (q + 1):end]
                ∂L = tril(∂factors[:, 1:q], -1)
            end
            L = UnitLowerTriangular(factors1)
            U1 = UpperTriangular(factors1)
            triu!(rmul!(∂A1, U1'))
            ∂A1 .+= tril!(mul!(lmul!(L', ∂L), ∂A2, U2', -1, 1), -1)
            rdiv!(∂A1, U1')
            ldiv!(L', ∂A)
        else  # tall A, system is [P1*A; P2*A] = [L1*U; L2*U]
            tril!(copyto!(∂A, ∂factors), -1)
            @views begin
                factors1 = factors[1:q, :]
                L2 = factors[(q + 1):end, :]
                ∂A1 = ∂A[1:q, :]
                ∂A2 = ∂A[(q + 1):end, :]
                ∂U = triu(∂factors[1:q, :])
            end
            U = UpperTriangular(factors1)
            L1 = UnitLowerTriangular(factors1)
            tril!(lmul!(L1', ∂A1), -1)
            ∂A1 .+= triu!(mul!(rmul!(∂U, U'), L2', ∂A2, -1, 1))
            ldiv!(L1', ∂A1)
            rdiv!(∂A, U')
        end
        if pivot === Val(true)
            ∂A = ∂A[invperm(F.p), :]
        end
        return NO_FIELDS, ∂A, DoesNotExist()
    end
    return F, lu_pullback
end

#####
##### functions of `LU`
#####

# this rrule is necessary because the primal mutates

function rrule(::typeof(getproperty), F::TF, x::Symbol) where {T,TF<:LU{T,<:StridedMatrix{T}}}
    function getproperty_LU_pullback(ΔY)
        ∂factors = if x === :L
            m, n = size(F.factors)
            S = eltype(ΔY)
            tril!([ΔY zeros(S, m, max(0, n - m))], -1)
        elseif x === :U
            m, n = size(F.factors)
            S = eltype(ΔY)
            triu!([ΔY; zeros(S, max(0, m - n), n)])
        elseif x === :factors
            Matrix(ΔY)
        else
            return (NO_FIELDS, DoesNotExist(), DoesNotExist())
        end
        ∂F = Composite{TF}(; factors=∂factors)
        return NO_FIELDS, ∂F, DoesNotExist()
    end
    return getproperty(F, x), getproperty_LU_pullback
end

# these rules are needed because the primal calls a LAPACK function

function frule((_, ΔF), ::typeof(LinearAlgebra.inv!), F::LU{<:Any,<:StridedMatrix})
    # factors must be square if the primal did not error
    L = UnitLowerTriangular(F.factors)
    U = UpperTriangular(F.factors)
    # compute ∂Y = -(U \ (L \ ∂L + ∂U / U) / L) * P while minimizing allocations
    m, n = size(F.factors)
    q = min(m, n)
    ∂L = tril(m ≥ n ? ΔF.factors : view(ΔF.factors, :, 1:q), -1)
    ∂U = triu(m ≤ n ? ΔF.factors : view(ΔF.factors, 1:q, :))
    ∂Y = ldiv!(L, ∂L)
    ∂Y .+= rdiv!(∂U, U)
    ldiv!(U, ∂Y)
    rdiv!(∂Y, L)
    rmul!(∂Y, -1)
    return LinearAlgebra.inv!(F), ∂Y[:, invperm(F.p)]
end

function rrule(::typeof(inv), F::LU{<:Any,<:StridedMatrix})
    function inv_LU_pullback(ΔY)
        # factors must be square if the primal did not error
        L = UnitLowerTriangular(F.factors)
        U = UpperTriangular(F.factors)
        # compute the following while minimizing allocations
        # ∂U = - triu((U' \ ∂Y * P' / L') / U')
        # ∂L = - tril(L' \ (U' \ ∂Y * P' / L'), -1)
        ∂factors = ΔY[:, F.p]
        ldiv!(U', ∂factors)
        rdiv!(∂factors, L')
        rmul!(∂factors, -1)
        ∂L = tril!(L' \ ∂factors, -1)
        triu!(rdiv!(∂factors, U'))
        ∂factors .+= ∂L
        ∂F = Composite{typeof(F)}(; factors=∂factors)
        return NO_FIELDS, ∂F
    end
    return inv(F), inv_LU_pullback
end

#####
##### `svd`
#####

function rrule(::typeof(svd), X::AbstractMatrix{<:Union{Real, Complex}})
    F = svd(X)
    function svd_pullback(Ȳ::Composite)
        # `getproperty` on `Composite`s ensures we have no thunks.
        ∂X = svd_rev(F, Ȳ.U, Ȳ.S, Ȳ.V)
        return (NO_FIELDS, ∂X)
    end
    return F, svd_pullback
end

function rrule(::typeof(getproperty), F::T, x::Symbol) where T <: SVD
    function getproperty_svd_pullback(Ȳ)
        C = Composite{T}
        ∂F = if x === :U
            C(U=Ȳ,)
        elseif x === :S
            C(S=Ȳ,)
        elseif x === :V
            C(V=Ȳ,)
        elseif x === :Vt
            # TODO: https://github.com/JuliaDiff/ChainRules.jl/issues/106
            throw(ArgumentError("Vt is unsupported; use V and transpose the result"))
        end
        return NO_FIELDS, ∂F, DoesNotExist()
    end
    return getproperty(F, x), getproperty_svd_pullback
end

@inline function safe_inv(x::T) where T<:Real
    if x == 0
        return inv(eps(T))
    elseif abs(x) < eps(T)
        return inv(sign(x)*eps(x))
    else
        return inv(x)
    end
end

# When not `Zero`s expect `Ū::AbstractMatrix, s̄::AbstractVector, V̄::AbstractMatrix`
function svd_rev(USV::SVD{T}, Ū, s̄, V̄) where T<:Real
    # Note: assuming a thin factorization, i.e. svd(A, full=false), which is the default
    U = USV.U
    s = USV.S
    V = USV.V
    Vt = USV.Vt

    k = length(s)
    F = T[i == j ? 1 : safe_inv(@inbounds s[j]^2 - s[i]^2) for i = 1:k, j = 1:k]

    # We do a lot of matrix operations here, so we'll try to be memory-friendly and do
    # as many of the computations in-place as possible. Benchmarking shows that the in-
    # place functions here are significantly faster than their out-of-place, naively
    # implemented counterparts, and allocate no additional memory.
    Ut = U'
    FUᵀŪ = _mulsubtrans!!(Ut*Ū, F)  # F .* (UᵀŪ - ŪᵀU)
    FVᵀV̄ = _mulsubtrans!!(Vt*V̄, F)  # F .* (VᵀV̄ - V̄ᵀV)
    ImUUᵀ = _eyesubx!(U*Ut)  # I - UUᵀ
    ImVVᵀ = _eyesubx!(V*Vt)  # I - VVᵀ

    S = Diagonal(s)
    S̄ = s̄ isa AbstractZero ? s̄ : Diagonal(s̄)

    # TODO: consider using MuladdMacro here
    Ā = add!!(U * FUᵀŪ * S, ImUUᵀ * (Ū / S)) * Vt
    Ā = add!!(Ā, U * S̄ * Vt)
    Ā = add!!(Ā, U * add!!(S * FVᵀV̄ * Vt, (S \ V̄') * ImVVᵀ))

    return Ā
end

#References:
#   https://j-towns.github.io/papers/svd-derivative.pdf
#   https://giggleliu.github.io/2019/04/02/einsumbp.html
#   https://arxiv.org/abs/1909.02659
function svd_rev(USV::SVD{T}, dU, dS, dV; η::Real=1e-40) where {FT,T<:Complex{FT}}
    U, S, V = USV
    all(x -> x isa AbstractZero, (dU, dS, dV)) && return Zero()
    η = FT(η)
    NS = length(S)
    S2 = S .^ 2
    Sinv = @. S/(S2+η)
    F = S2' .- S2
    F ./= (F .^ 2 .+ η)

    res = Zero()
    if !(dU isa AbstractZero)
        UdU = U'*dU
        J = F.*(UdU)
        res += (J+J')*LinearAlgebra.Diagonal(S) + LinearAlgebra.Diagonal(1im*imag(LinearAlgebra.diag(UdU)) .* Sinv)
    end
    if !(dV isa AbstractZero)
        VdV = V'*dV
        K = F.*(VdV)
        res += LinearAlgebra.Diagonal(S) * (K+K')
    end
    if !(dS isa AbstractZero)
        res += LinearAlgebra.Diagonal(dS)
    end

    res = U * res * V'

    if !(dU isa AbstractZero) && size(U, 1) != size(U, 2)
        res += (dU - U* (U'*dU)) * LinearAlgebra.Diagonal(Sinv) * V'
    end

    if !(dV isa AbstractZero) && size(V, 1) != size(V, 2)
        res = res + U * LinearAlgebra.Diagonal(Sinv) * (dV' - (dV'*V)*V')
    end
    res
end

#####
##### `eigen`
#####

# TODO:
# - support correct differential of phase convention when A is hermitian
# - simplify when A is diagonal
# - support degenerate matrices (see #144)

function frule((_, ΔA), ::typeof(eigen!), A::StridedMatrix{T}; kwargs...) where {T<:BlasFloat}
    ΔA isa AbstractZero && return (eigen!(A; kwargs...), ΔA)
    if ishermitian(A)
        sortby = get(kwargs, :sortby, VERSION ≥ v"1.2.0" ? LinearAlgebra.eigsortby : nothing)
        return if sortby === nothing
            frule((Zero(), Hermitian(ΔA)), eigen!, Hermitian(A))
        else
            frule((Zero(), Hermitian(ΔA)), eigen!, Hermitian(A); sortby=sortby)
        end
    end
    F = eigen!(A; kwargs...)
    λ, V = F.values, F.vectors
    tmp = V \ ΔA
    ∂K = tmp * V
    ∂Kdiag = @view ∂K[diagind(∂K)]
    ∂λ = eltype(λ) <: Real ? real.(∂Kdiag) : copy(∂Kdiag)
    ∂K ./= transpose(λ) .- λ
    fill!(∂Kdiag, 0)
    ∂V = mul!(tmp, V, ∂K)
    _eigen_norm_phase_fwd!(∂V, A, V)
    ∂F = Composite{typeof(F)}(values = ∂λ, vectors = ∂V)
    return F, ∂F
end

function rrule(::typeof(eigen), A::StridedMatrix{T}; kwargs...) where {T<:Union{Real,Complex}}
    F = eigen(A; kwargs...)
    function eigen_pullback(ΔF::Composite)
        λ, V = F.values, F.vectors
        Δλ, ΔV = ΔF.values, ΔF.vectors
        ΔV isa AbstractZero && Δλ isa AbstractZero && return (NO_FIELDS, Δλ + ΔV)
        if eltype(λ) <: Real && ishermitian(A)
            hermA = Hermitian(A)
            ∂V = ΔV isa AbstractZero ? ΔV : copyto!(similar(ΔV), ΔV)
            ∂hermA = eigen_rev!(hermA, λ, V, Δλ, ∂V)
            ∂Atriu = _symherm_back(typeof(hermA), ∂hermA, Symbol(hermA.uplo))
            ∂A = ∂Atriu isa AbstractTriangular ? triu!(∂Atriu.data) : ∂Atriu
        elseif ΔV isa AbstractZero
            ∂K = Diagonal(Δλ)
            ∂A = V' \ ∂K * V'
        else
            ∂V = copyto!(similar(ΔV), ΔV)
            _eigen_norm_phase_rev!(∂V, A, V)
            ∂K = V' * ∂V
            ∂K ./= λ' .- conj.(λ)
            ∂K[diagind(∂K)] .= Δλ
            ∂A = mul!(∂K, V' \ ∂K, V')
        end
        return NO_FIELDS, T <: Real ? real(∂A) : ∂A
    end
    eigen_pullback(ΔF::AbstractZero) = (NO_FIELDS, ΔF)
    return F, eigen_pullback
end

# mutate ∂V to account for the (arbitrary but consistent) normalization and phase condition
# applied to the eigenvectors.
# these implementations assume the convention used by eigen in LinearAlgebra (i.e. that of
# LAPACK.geevx!; eigenvectors have unit norm, and the element with the largest absolute
# value is real), but they can be specialized for `A`

function _eigen_norm_phase_fwd!(∂V, A, V)
    @inbounds for i in axes(V, 2)
        v, ∂v = @views V[:, i], ∂V[:, i]
        # account for unit normalization
        ∂c_norm = -real(dot(v, ∂v))
        if eltype(V) <: Real
            ∂c = ∂c_norm
        else
            # account for rotation of largest element to real
            k = _findrealmaxabs2(v)
            ∂c_phase = -imag(∂v[k]) / real(v[k])
            ∂c = complex(∂c_norm, ∂c_phase)
        end
        ∂v .+= v .* ∂c
    end
    return ∂V
end

function _eigen_norm_phase_rev!(∂V, A, V)
    @inbounds for i in axes(V, 2)
        v, ∂v = @views V[:, i], ∂V[:, i]
        ∂c = dot(v, ∂v)
        # account for unit normalization
        ∂v .-= real(∂c) .* v
        if !(eltype(V) <: Real)
            # account for rotation of largest element to real
            k = _findrealmaxabs2(v)
            @inbounds ∂v[k] -= im * (imag(∂c) / real(v[k]))
        end
    end
    return ∂V
end

# workaround for findmax not taking a mapped function
function _findrealmaxabs2(x)
    amax = abs2(first(x))
    imax = 1
    @inbounds for i in 2:length(x)
        xi = x[i]
        !isreal(xi) && continue
        a = abs2(xi)
        a < amax && continue
        amax, imax = a, i
    end
    return imax
end

#####
##### `eigvals`
#####

function frule((_, ΔA), ::typeof(eigvals!), A::StridedMatrix{T}; kwargs...) where {T<:BlasFloat}
    ΔA isa AbstractZero && return eigvals!(A; kwargs...), ΔA
    if ishermitian(A)
        λ, ∂λ = frule((Zero(), Hermitian(ΔA)), eigvals!, Hermitian(A))
        sortby = get(kwargs, :sortby, VERSION ≥ v"1.2.0" ? LinearAlgebra.eigsortby : nothing)
        _sorteig!_fwd(∂λ, λ, sortby)
    else
        F = eigen!(A; kwargs...)
        λ, V = F.values, F.vectors
        tmp = V \ ΔA
        ∂λ = similar(λ)
        # diag(tmp * V) without computing full matrix product
        if eltype(∂λ) <: Real
            broadcast!((a, b) -> sum(real ∘ prod, zip(a, b)), ∂λ, eachrow(tmp), eachcol(V))
        else
            broadcast!((a, b) -> sum(prod, zip(a, b)), ∂λ, eachrow(tmp), eachcol(V))
        end
    end
    return λ, ∂λ
end

function rrule(::typeof(eigvals), A::StridedMatrix{T}; kwargs...) where {T<:Union{Real,Complex}}
    F, eigen_back = rrule(eigen, A; kwargs...)
    λ = F.values
    function eigvals_pullback(Δλ)
        ∂F = Composite{typeof(F)}(values = Δλ)
        _, ∂A = eigen_back(∂F)
        return NO_FIELDS, ∂A
    end
    return λ, eigvals_pullback
end

# adapted from LinearAlgebra.sorteig!
function _sorteig!_fwd(Δλ, λ, sortby)
    Δλ isa AbstractZero && return (sort!(λ; by=sortby), Δλ)
    if sortby !== nothing
        p = sortperm(λ; alg=QuickSort, by=sortby)
        permute!(λ, p)
        permute!(Δλ, p)
    end
    return (λ, Δλ)
end

#####
##### `cholesky`
#####

function rrule(::typeof(cholesky), A::Real, uplo::Symbol=:U)
    C = cholesky(A, uplo)
    function cholesky_pullback(ΔC::Composite)
        return NO_FIELDS, ΔC.factors[1, 1] / (2 * C.U[1, 1]), DoesNotExist()
    end
    return C, cholesky_pullback
end

function rrule(::typeof(cholesky), A::Diagonal{<:Real}, ::Val{false}; check::Bool=true)
    C = cholesky(A, Val(false); check=check)
    function cholesky_pullback(ΔC::Composite)
        Ā = Diagonal(diag(ΔC.factors) .* inv.(2 .* C.factors.diag))
        return NO_FIELDS, Ā, DoesNotExist()
    end
    return C, cholesky_pullback
end

# The appropriate cotangent is different depending upon whether A is Symmetric / Hermitian,
# or just a StridedMatrix.
# Implementation due to Seeger, Matthias, et al. "Auto-differentiating linear algebra."
function rrule(
    ::typeof(cholesky),
    A::LinearAlgebra.HermOrSym{<:LinearAlgebra.BlasReal, <:StridedMatrix},
    ::Val{false};
    check::Bool=true,
)
    C = cholesky(A, Val(false); check=check)
    function cholesky_pullback(ΔC::Composite)
        Ā, U = _cholesky_pullback_shared_code(C, ΔC)
        Ā = BLAS.trsm!('R', 'U', 'C', 'N', one(eltype(Ā)) / 2, U.data, Ā)
        return NO_FIELDS, _symhermtype(A)(Ā), DoesNotExist()
    end
    return C, cholesky_pullback
end

function rrule(
    ::typeof(cholesky),
    A::StridedMatrix{<:LinearAlgebra.BlasReal},
    ::Val{false};
    check::Bool=true,
)
    C = cholesky(A, Val(false); check=check)
    function cholesky_pullback(ΔC::Composite)
        Ā, U = _cholesky_pullback_shared_code(C, ΔC)
        Ā = BLAS.trsm!('R', 'U', 'C', 'N', one(eltype(Ā)), U.data, Ā)
        idx = diagind(Ā)
        @views Ā[idx] .= real.(Ā[idx]) ./ 2
        return (NO_FIELDS, UpperTriangular(Ā), DoesNotExist())
    end
    return C, cholesky_pullback
end

function _cholesky_pullback_shared_code(C, ΔC)
    U = C.U
    Ū = ΔC.U
    Ā = similar(U.data)
    Ā = mul!(Ā, Ū, U')
    Ā = LinearAlgebra.copytri!(Ā, 'U', true)
    Ā = ldiv!(U, Ā)
    return Ā, U
end

function rrule(::typeof(getproperty), F::T, x::Symbol) where {T <: Cholesky}
    function getproperty_cholesky_pullback(Ȳ)
        C = Composite{T}
        ∂F = if x === :U
            if F.uplo === 'U'
                C(U=UpperTriangular(Ȳ),)
            else
                C(L=LowerTriangular(Ȳ'),)
            end
        elseif x === :L
            if F.uplo === 'L'
                C(L=LowerTriangular(Ȳ),)
            else
                C(U=UpperTriangular(Ȳ'),)
            end
        end
        return NO_FIELDS, ∂F, DoesNotExist()
    end
    return getproperty(F, x), getproperty_cholesky_pullback
end


text:
using Foundation
using Test

@test [1,2,3].count == 3

a = [1,2,3]
a.append(5)

@test a == [1,2,3,5]


text:
using SeisNoiseAWS
using Test

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


text:
# We demonstrate the full power of the sparse grid basis
# by successfully evolving a 4+1-dimensional PDE corresponding
# to the phase evolution of some matter distribution.
# This is the Vlassov-Poisson or "Collisionless Boltzman" equation
using GalerkinSparseGrids
using LinearAlgebra
using ODE

# The spatial dimension is 2, meaning a 4D phase space
# We will use mode order 5 at 5th level resolution along each axis
# Higher resolution at the moment requires a supercomputer
D = 2; k = 5; n = 5
# D = 2; k = 5; n = 2

# We will evolve from an initial t = 0 to a final t1 = 0.54
t0 = 0.0; t1 = 0.54; nout = 101
# t0 = 0.0; t1 = 0.0054; nout = 2

# Initial distribution function in phase space:
@inline fx(x) = exp(-2 * pi^2 * norm(x.-1/2)^2)
@inline fv(v) = exp(-2 * pi^2 * norm(v)^2)
fx_modal = coeffs_DG(1, k, n, fx)
fv_modal = coeffs_DG(1, k, n, fv)
# We use the tensor_construct method as before because of its
# increased accuracy
f0_modal_array = [i <= D ? fx_modal : fv_modal for i in 1:2*D]
f0_modal = 2*pi*D2V(2*D,k,n,tensor_construct(2*D, k, n, f0_modal_array))

# Constructing the transformation matrices takes the longest time
# and most memory by far.
# Once the matrices have been constructed, however, the same
# matrices can forever be used at a given (D, k, n)
# It would therefore be advantageous to load these in from some
# external file rather than computing them (likely multiple times over)
m2n, n2p = make_modal2point_matrices(2*D, k, n)
p2n, n2m = make_point2modal_matrices(2*D, k, n)

# the force as a function of r^2:
# (We use r^2) here because |r| is dicontinuous at the origin and gives
# instabilities for interpolation
F_radial = x -> x == 0 ? 10/9 : 10/3 * (sqrt(x) - atan(sqrt(x)))/(sqrt(x)^2)
fr2 = broadcast(F_radial, get_r2_point(2*D, k, n, m2n, n2p))
# I have defined F_radial so that multiplying it by x_i
# gives the ith component of the force vector
F_point = [get_xi_point(2*D, i, k, n, m2n, n2p) .* fr2 for i in 1:D]

# Finally, we do the full evolution of the Vlasov equation:
vlasov_evolve(D, k, n,
                m2n, n2p, p2n, n2m,
                f0_modal, F_point,
                t0, t1, nout;
                order="45", points=:specified)
# points=:specified only saves the first and last set of coefficients
# in the evolution. This is to save memory.
# If one desires the full set of coefficients, simply remove this keyword arg


text:
module JuliaPackageWithRustDep

const deps_file = joinpath(dirname(@__FILE__), "..", "deps", "deps.jl")
if !isfile(deps_file)
    error("JuliaPackageWithRustDep.jl is not installed properly, run Pkg.build(\"JuliaPackageWithRustDep\") and restart Julia.")
end
include(deps_file)

function __init__()
    check_deps()
end

include("api.jl")

function read_rust_owned_string() :: String
    cstring = rustdylib_generate_rust_owned_string()
    result = unsafe_string(cstring) # copies the contents of the string
    rustdylib_free_rust_owned_string(cstring) # ask Rust to free the memory
    return result
end

end # module


text:
module Sha256
export digest

#######################################################################

# Helper functions
ch(x::Uint32, y::Uint32, z::Uint32) = (x & y) $ (~x & z)
maj(x::Uint32, y::Uint32, z::Uint32) = (x & y) $ (x & z) $ (y & z)
rot(x::Uint32, n) = x >>> n | x << (32 - n)
shift(x::Uint32, n) = x >>> n
sigma0(x::Uint32) = rot(x, 2) $ rot(x, 13) $ rot(x, 22)
sigma1(x::Uint32) = rot(x, 6) $ rot(x, 11) $ rot(x, 25)
sgm0(x::Uint32) = rot(x, 7) $ rot(x, 18) $ shift(x, 3)
sgm1(x::Uint32) = rot(x, 17) $ rot(x, 19) $ shift(x, 10)

#######################################################################

#######################################################################

# S-box
const k = [0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
           0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,
           0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
           0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967,
           0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13, 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85,
           0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
           0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
           0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, 0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2]

#######################################################################

#######################################################################

# Message expansion
# msg is an array of Uint32
function expand_msg(msg)
    w = Array(Uint32, 64)
    w[1:16] = msg

    for j = 17:64
        w[j] = sgm1(w[j - 2]) + w[j - 7] + sgm0(w[j - 15]) + w[j - 16]
    end

    return w
end

#######################################################################

#######################################################################

# Message padding
# msg is an array of Uint8, representing the original message in bytes
function pad(msg)
    function construct(slc)
        local _::Uint32 = 0
        for i in [int(i) for i in slc]
            _ <<= 8
            _ |= i
        end
        return _
    end

    function to_byte_array(number)
        reverse(reinterpret(Uint8, [number]))
    end

    len, = size(msg)
    l = len * 8
    k = mod(448 - 1 - l, 512)
    k_byte = div(k + 1, 8) - 1
    msg = vcat(msg, [0x80], zeros(k_byte), to_byte_array(l))

    len, = size(msg)
    n = div(len * 8, 512)

    const rounds::Int32 = 64 / 4

    ret = zeros(Uint32, rounds * n)
    for i = 1:n
        for j = 1:rounds
            ret[16(i - 1)+ j] = construct(msg[64(i - 1) + 4(j - 1) + 1:64(i - 1) + 4j])
        end
    end
    return ret, n
end

#######################################################################

#######################################################################
# Compression
# init is an array of Uint32 (the initial hash values)
# w is an array of Uint32 (the expanded message)
function compress(init, w)
    a, b, c, d, e, f, g, h = init

    for j = 1:64
        t1::Uint32 = h + sigma1(e) + ch(e, f, g) + k[j] + w[j]
        t2::Uint32 = sigma0(a) + maj(a, b, c)
        h = g
        g = f
        f = e
        e::Uint32 = d + t1
        d = c
        c = b
        b = a
        a::Uint32 = t1 + t2
    end
    init[1] += a
    init[2] += b
    init[3] += c
    init[4] += d
    init[5] += e
    init[6] += f
    init[7] += g
    init[8] += h

    return init
end

#######################################################################

#######################################################################
# Message digestion
# msg is an array of Uint8, representing the original message in bytes
function digest(msg)
    blocks, n = pad(msg)
    m = [0x6a09e667, 0xbb67ae85,
         0x3c6ef372, 0xa54ff53a,
         0x510e527f, 0x9b05688c,
         0x1f83d9ab, 0x5be0cd19]
    init = copy(m)
    for i = 1:n
        w = expand_msg(blocks[16(i - 1) + 1:16i])
        compress(m, w)
    end

    return m
end

#######################################################################

end # module


text:
#######################################################################
# Lattice presets
#######################################################################

module LatticePresets

using Quantica
using Quantica: NameType

linear(; a0 = 1, kw...) =
    lattice(sublat((0.,)); bravais = a0 .* (1,), kw...)

square(; a0 = 1, kw...) =
    lattice(sublat((0., 0.)); bravais = a0 * SA[1. 0.; 0. 1.]', kw...)

triangular(; a0 = 1, kw...) =
    lattice(sublat((0., 0.)); bravais = a0 * SA[cos(pi/3) sin(pi/3); -cos(pi/3) sin(pi/3)]', kw...)

honeycomb(; a0 = 1, kw...) =
    lattice(sublat((0.0, -0.5*a0/sqrt(3.0)), name = :A),
            sublat((0.0,  0.5*a0/sqrt(3.0)), name = :B);
            bravais = a0 * SA[cos(pi/3) sin(pi/3); -cos(pi/3) sin(pi/3)]', kw...)

cubic(; a0 = 1, kw...) =
    lattice(sublat((0., 0., 0.)); bravais = a0 * SA[1. 0. 0.; 0. 1. 0.; 0. 0. 1.]', kw...)

fcc(; a0 = 1, kw...) =
    lattice(sublat((0., 0., 0.)); bravais = (a0/sqrt(2.)) * SA[-1. -1. 0.; 1. -1. 0.; 0. 1. -1.]', kw...)

bcc(; a0 = 1, kw...) =
    lattice(sublat((0., 0., 0.)); bravais = a0 * SA[1. 0. 0.; 0. 1. 0.; 0.5 0.5 0.5]', kw...)

hcp(; a0 = 1, kw...) =
    lattice(sublat((0., 0., 0.), (a0*0.5, a0*0.5/sqrt(3.), a0*0.5)); bravais = a0 * SA[1. 0. 0.; -0.5 0.5*sqrt(3.) 0.; 0. 0. 1.]', kw...)

end # module

#######################################################################
# Hamiltonian presets
#######################################################################

module HamiltonianPresets

using Quantica, LinearAlgebra

function graphene(; a0 = 0.246, range = a0/sqrt(3), t0 = 2.7, β = 3, kw...)
    lat = LatticePresets.honeycomb(a0 = a0)
    h = hamiltonian(lat, hopping((r, dr) -> t0 * exp(-β*(norm(dr)/a0 - 1)) * I, range = range); kw...)
    return h
end

function twisted_bilayer_graphene(;
    twistindex = 1, twistindices = (twistindex, 1), a0 = 0.246, interlayerdistance = 1.36a0,
    rangeintralayer = a0/sqrt(3), rangeinterlayer = 4a0/sqrt(3), hopintra = 2.70,
    hopinter = 0.48, modelintra = hopping(hopintra, range = rangeintralayer), kw...)

    (m, r) = twistindices
    θ = acos((3m^2 + 3m*r +r^2/2)/(3m^2 + 3m*r + r^2))
    sAbot = sublat((0.0, -0.5a0/sqrt(3.0), - interlayerdistance / 2); name = :Ab)
    sBbot = sublat((0.0,  0.5a0/sqrt(3.0), - interlayerdistance / 2); name = :Bb)
    sAtop = sublat((0.0, -0.5a0/sqrt(3.0),   interlayerdistance / 2); name = :At)
    sBtop = sublat((0.0,  0.5a0/sqrt(3.0),   interlayerdistance / 2); name = :Bt)
    brbot = a0 * SA[ cos(pi/3) sin(pi/3) 0; -cos(pi/3) sin(pi/3) 0]'
    brtop = a0 * SA[ cos(pi/3) sin(pi/3) 0; -cos(pi/3) sin(pi/3) 0]'
    # Supercell matrices sc.
    # The one here is a [1 0; -1 1] rotation of the one in Phys. Rev. B 86, 155449 (2012)
    if gcd(r, 3) == 1
        scbot = SA[m -(m+r); (m+r) 2m+r] * SA[1 0; -1 1]
        sctop = SA[m+r -m; m 2m+r] * SA[1 0; -1 1]
    else
        scbot = SA[m+r÷3 -r÷3; r÷3 m+2r÷3] * SA[1 0; -1 1]
        sctop = SA[m+2r÷3 r÷3; -r÷3 m+r÷3] * SA[1 0; -1 1]
    end
    latbot = lattice(sAbot, sBbot; bravais = brbot)
    lattop = lattice(sAtop, sBtop; bravais = brtop)
    htop = hamiltonian(lattop, modelintra; kw...) |> unitcell(sctop)
    hbot = hamiltonian(latbot, modelintra; kw...) |> unitcell(scbot)
    let R = SA[cos(θ/2) -sin(θ/2) 0; sin(θ/2) cos(θ/2) 0; 0 0 1]
        transform!(r -> R * r, htop)
    end
    let R = SA[cos(θ/2) sin(θ/2) 0; -sin(θ/2) cos(θ/2) 0; 0 0 1]
        transform!(r -> R * r, hbot)
    end
    modelinter = hopping((r,dr) -> (
        hopintra * exp(-3*(norm(dr)/a0 - 1))  *  dot(dr, SVector(1,1,0))^2/sum(abs2, dr) -
        hopinter * exp(-3*(norm(dr)/a0 - interlayerdistance/a0)) * dr[3]^2/sum(abs2, dr)),
        range = rangeinterlayer)
    return combine(hbot, htop; coupling = modelinter)
end

end # module

#######################################################################
# Region presets
#######################################################################

module RegionPresets

using StaticArrays

struct Region{E,F}
    f::F
end

Region{E}(f::F) where {E,F<:Function} = Region{E,F}(f)

(region::Region{E})(r::SVector{E2}) where {E,E2} = region.f(r)

Base.show(io::IO, ::Region{E}) where {E} =
    print(io, "Region{$E} : region in $(E)D space")

Base.:&(r1::Region{E}, r2::Region{E}) where {E} = Region{E}(r -> r1.f(r) && r2.f(r))
Base.:|(r1::Region{E}, r2::Region{E}) where {E}  = Region{E}(r -> r1.f(r) || r2.f(r))
Base.xor(r1::Region{E}, r2::Region{E}) where {E}  = Region{E}(r -> xor(r1.f(r), r2.f(r)))
Base.:!(r1::Region{E}) where {E}  = Region{E}(r -> !r1.f(r))

extended_eps(T = Float64) = sqrt(eps(T))

circle(radius = 10.0, c...) = Region{2}(_region_ellipse((radius, radius), c...))

ellipse(radii = (10.0, 15.0), c...) = Region{2}(_region_ellipse(radii, c...))

square(side = 10.0, c...) = Region{2}(_region_rectangle((side, side), c...))

rectangle(sides = (10.0, 15.0), c...) = Region{2}(_region_rectangle(sides, c...))

sphere(radius = 10.0, c...) = Region{3}(_region_ellipsoid((radius, radius, radius), c...))

spheroid(radii = (10.0, 15.0, 20.0), c...) = Region{3}(_region_ellipsoid(radii, c...))

cube(side = 10.0, c...) = Region{3}(_region_cuboid((side, side, side), c...))

cuboid(sides = (10.0, 15.0, 20.0), c...) = Region{3}(_region_cuboid(sides, c...))

function _region_ellipse((rx, ry), (cx, cy) = (0, 0))
    return r -> ((r[1]-cx)/rx)^2 + ((r[2]-cy)/ry)^2 <= 1 + extended_eps(Float64)
end

function _region_rectangle((lx, ly), (cx, cy) = (0, 0))
    return r -> abs(2*(r[1]-cx)) <= lx * (1 + extended_eps()) &&
                abs(2*(r[2]-cy)) <= ly * (1 + extended_eps())
end

function _region_ellipsoid((rx, ry, rz), (cx, cy, cz) = (0, 0, 0))
    return r -> ((r[1]-cx)/rx)^2 + ((r[2]-cy)/ry)^2 + ((r[3]-cz)/rz)^2 <= 1 + eps()
end

function _region_cuboid((lx, ly, lz), (cx, cy, cz) = (0, 0, 0))
    return r -> abs(2*(r[1]-cx)) <= lx * (1 + extended_eps()) &&
                abs(2*(r[2]-cy)) <= ly * (1 + extended_eps()) &&
                abs(2*(r[3]-cy)) <= lz * (1 + extended_eps())
end

end # module

#######################################################################
# Aliases
#######################################################################

const RP = RegionPresets
const LP = LatticePresets
const HP = HamiltonianPresets

max_stars_repo_name: pablosanjose/Quantica

max_issues_repo_name: pablosanjose/Quantica

max_forks_repo_name: pablosanjose/Quantica