text:
# This code calls the first-order eccentric TTV code
# compute_ttv.jl.  Please cite Agol & Deck (2015) if
# you make use of this in your research.

include("compute_ttv.jl")
#include("laplace_coefficients_initialize.jl")

function test_ttv(jmax::Integer,n1::Integer,n2::Integer,data::Vector; WriteOutput::Bool = true, num_evals::Integer = 1, profile::Bool = false)
@assert(jmax>=1)  # Should there be a larger minimum?
@assert(n1>2)
@assert(n2>2)
@assert(length(data)==10)
# Performs a test of the transit_times.jl routine
# Set up planets planar-planet types for the inner and outer planets:
p1=TTVFaster.Planet_plane_hk(data[1],data[2],data[3],data[4],data[ 5])
p2=TTVFaster.Planet_plane_hk(data[6],data[7],data[8],data[9],data[10])
time1 = collect(p1.trans0 + linspace(0,n1-1,n1) * p1.period)
time2 = collect(p2.trans0 + linspace(0,n2-1,n2) * p2.period)
alpha0=(p1.period/p2.period)^(2//3)
# Initialize the computation of the Laplace coefficients:
b0=TTVFaster.LaplaceCoefficients.initialize(jmax+1,alpha0)
# Define arrays to hold the TTVs:
#ttv1=Array(Float64,n1)
#ttv2=Array(Float64,n2)
ttv_el_type = eltype(data) == Float64 ? Float64 : Number
ttv1=Array(ttv_el_type,n1)
ttv2=Array(ttv_el_type,n2)
# Define arrays to hold the TTV coefficients and Laplace coefficients:
f1=Array(Float64,jmax+2,5)
f2=Array(Float64,jmax+2,5)
b=Array(Float64,jmax+2,3)
hashsum = 0
for i in 1:num_evals
   # Call the compute_ttv code which implements equation (33)
   TTVFaster.compute_ttv!(jmax,p1,p2,time1,time2,ttv1,ttv2,f1,f2,b,alpha0,b0)
   if profile
     hashsum += hash(ttv1)+hash(ttv2)
   end
end
if profile
  println("# Ignore this: ",hashsum) # This just makes sure optimizer doesn't optimize away important calculations.
end

if WriteOutput
   # Write the mean ephemeris and TTV results to two files:
   writedlm("inner_ttv.txt",[time1 ttv1])
   writedlm("outer_ttv.txt",[time2 ttv2])
end
   return ttv1,ttv2
end

max_stars_repo_name: ericagol/TTVFaster

max_issues_repo_name: ericagol/TTVFaster

max_forks_repo_name: ericagol/TTVFaster


text:
using .DualNumbers

function blkimpv_dual(num1, num2, num3, num4, num5)
    @eval function blkimpv(S0::$num1, K::$num2, r::$num3, T::$num4, Price::$num5, FlagIsCall::Bool = true, xtol::Real = 1e-14, ytol::Real = 1e-15)
        blscheck(S0, K, r, T, 0.1, r)

        return blsimpv(S0, K, r, T, Price, r, FlagIsCall)
    end
end

type_blkimpv_dual_1 = [Dual, Union{Real, Dual}, Union{Real, Dual}, Union{Real, Dual}, Union{Real, Dual}]
type_blkimpv_dual_2 = [Real, Dual, Union{Real, Dual}, Union{Real, Dual}, Union{Real, Dual}]
type_blkimpv_dual_3 = [Real, Real, Union{Real, Dual}, Union{Real, Dual}, Union{Real, Dual}]
blkimpv_dual(type_blkimpv_dual_1[1], type_blkimpv_dual_1[2], type_blkimpv_dual_1[3], type_blkimpv_dual_1[4], type_blkimpv_dual_1[5])
blkimpv_dual(type_blkimpv_dual_2[1], type_blkimpv_dual_2[2], type_blkimpv_dual_2[3], type_blkimpv_dual_2[4], type_blkimpv_dual_2[5])
blkimpv_dual(type_blkimpv_dual_3[1], type_blkimpv_dual_3[2], type_blkimpv_dual_3[3], type_blkimpv_dual_3[4], type_blkimpv_dual_3[5])

function blsimpv_dual(num1, num2, num3, num4, num5, num6)
    @eval function blsimpv(S0::$num1, K::$num2, r::$num3, T::$num4, Price::$num5, d::$num6 = 0.0, FlagIsCall::Bool = true, xtol::Real = 1e-14, ytol::Real = 1e-15)
        if (Price < $num5(0))
            throw(ErrorException("Option Price Cannot Be Negative"))
        end
        FinancialToolbox.blscheck(S0, K, r, T, 0.1, d)
        value__(x) = x.value
        value__(x::Real) = x
        f(x) = (blsprice(value__(S0), value__(K), value__(r), value__(T), x, value__(d), FlagIsCall) - value__(Price))
        σ = FinancialToolbox.brentMethod(f, 0.001, 1.2, xtol, ytol)
        out = dual(0.0, 0.0)

        if (!isdual(Price))
            der_ = -(blsprice(S0, K, r, T, σ, d, FlagIsCall) / blsvega(value__(S0), value__(K), value__(r), value__(T), σ, value__(d), FlagIsCall)).epsilon
            out = dual(σ, der_)
        else
            der_ = 1 / blsvega(value__(S0), value__(K), value__(r), value__(T), σ, value__(d), FlagIsCall)
            out = dual(σ, der_)
        end

        return out
    end
end

type_blsimpv_dual_1 = [Dual, Union{Real, Dual}, Union{Real, Dual}, Union{Real, Dual}, Union{Real, Dual}, Union{Real, Dual}]
type_blsimpv_dual_2 = [Real, Dual, Union{Real, Dual}, Union{Real, Dual}, Union{Real, Dual}, Union{Real, Dual}]
type_blsimpv_dual_3 = [Real, Real, Union{Real, Dual}, Union{Real, Dual}, Union{Real, Dual}, Union{Real, Dual}]
blsimpv_dual(type_blsimpv_dual_1[1], type_blsimpv_dual_1[2], type_blsimpv_dual_1[3], type_blsimpv_dual_1[4], type_blsimpv_dual_1[5], type_blsimpv_dual_1[6])
blsimpv_dual(type_blsimpv_dual_2[1], type_blsimpv_dual_2[2], type_blsimpv_dual_2[3], type_blsimpv_dual_2[4], type_blsimpv_dual_2[5], type_blsimpv_dual_2[6])
blsimpv_dual(type_blsimpv_dual_3[1], type_blsimpv_dual_3[2], type_blsimpv_dual_3[3], type_blsimpv_dual_3[4], type_blsimpv_dual_3[5], type_blsimpv_dual_3[6])

ext: jl


text:
export chebygrid, ischebygrid, chebycoef, cheby, chebyderiv
export ChebyshevInterpolator, BichebyshevInterpolator

ξ2x(ξ, a, b) = (ξ + 1)*((b - a)/2) + a

x2ξ(x, a, b) = 2*(x - a)/(b - a) - 1

x2θ(x, a, b) = acos(x2ξ(x, a, b))

"""
    chebygrid(n)

Create an array of `n` chebyshev nodes in [-1,1]
"""
@memoize chebygrid(n::Int) = cos.(π*(n-1:-1:0)/(n-1))

"""
    chebygrid(xa, xb, n)

Create an array of `n` chebyshev nodes in [`xa`,`xb`]
"""
chebygrid(xa, xb, n::Int) = ξ2x.(chebygrid(n), xa, xb)

"""
    chebygrid(xa, xb, nx, ya, yb, ny)

Create a two-dimensional grid of chebyshev nodes using `nx` points along the first axis, in [`xa`,`xb`], and `ny` points along the second axis, in [`ya`,`yb`].
"""
function chebygrid(xa, xb, nx::Int, ya, yb, ny::Int)
    X = chebygrid(xa, xb, nx) .* ones(ny)'
    Y = chebygrid(ya, yb, ny)' .* ones(nx)
    return X, Y
end

function ischebygrid(x)::Bool
    n = length(x)
    c = chebygrid(n)
    xa, xb = minimum(x), maximum(x)
    @inbounds for i ∈ eachindex(x)
        ξ = x2ξ(x[i], xa, xb)
        #shift from [-1,1] to [1,3] to avoid issues comparing to zero
        !isapprox(ξ+2, c[i]+2) && return false
    end
    return true
end

@memoize function chebymatrix(n::Int)
    @assert n > 1 "can't construct cheby matrix smaller than 2 x 2"
    A = zeros(n,n)
    ξ = chebygrid(n)
    for j ∈ 1:n, k ∈ 1:n
        @inbounds A[j,k] = cos((k-1)*acos(ξ[j]))
    end
    return A
end

function chebyrecurrance!(T, ξ, L::Int)::Nothing
    Tₖ₋₂ = one(ξ)
    Tₖ₋₁ = ξ
    @inbounds T[2] = Tₖ₋₁
    for k = 3:L
        #compute next value
        Tₖ = 2ξ*Tₖ₋₁ - Tₖ₋₂
        #set array value
        @inbounds T[k] = Tₖ
        #swaps
        Tₖ₋₂ = Tₖ₋₁
        Tₖ₋₁ = Tₖ
    end
    nothing
end

function chebyrecurrance(ξ::U, L::Int) where {U}
    T = Vector{U}(undef, L)
    @inbounds T[1] = one(ξ)
    chebyrecurrance!(T, ξ, L)
    return T
end

"""
    chebycoef(y)

Compute the Chebyshev expansion coefficients for a set of points `y`, which are assumed to be located on the Chebyshev points for some interval.
"""
function chebycoef(y)
    n = length(y)
    @assert n > 1 "must have at least 2 points to form chebyshev coefficients"
    Aᴵ = invertedchebymatrix(n)
    return Aᴵ*y
end

"""
    cheby(coef, x, xa, xb)

Evaluates the Chebyshev expansion represented by the coefficients in `coef` and defined on the interval [`xa`,`xb`] at the point `x`.
"""
function cheby(coef, x, xa, xb)
    N = length(coef)
    @assert xa <= x <= xb "x must be inside [xa,xb]"
    #get coordinate in ξ space
    ξ = x2ξ(x, xa, xb)
    #first two elements of cheby recursion
    Tₖ₋₂ = one(ξ)
    Tₖ₋₁ = ξ
    #first two terms of dot product
    @inbounds y = Tₖ₋₂*coef[1] + Tₖ₋₁*coef[2]
    #cheby recursion and rest of terms in dot product, all at once
    for k = 3:N
        #next value in recursion
        Tₖ = 2*ξ*Tₖ₋₁ - Tₖ₋₂
        #next term in dot product
        @inbounds y += Tₖ*coef[k]
        #swaps
        Tₖ₋₂ = Tₖ₋₁
        Tₖ₋₁ = Tₖ
    end
    return y
end

"""
    chebyderiv(coef, xa, xb)

Generates the expansion coefficents for the derivative of a preexisting Chebyshev expansion defined on the interval [`xa`,`xb`].
"""
function chebyderiv(coef, xa, xb)
    n = length(coef)
    T = eltype(coef)
    @assert xa < xb "xa must be less than xb"
    #length 1 case is a trivial zero slope
    n == 1 && return zeros(T, 1)
    #start the array and check for length 2 case
    d = zeros(T, n-1)
    d[n-1] = 2*(n-2)*coef[n]
    n == 2 && return d
    #recurrance for all the rest
    d[n-2] = 2*(n-3)*coef[n-1]
    @inbounds for k ∈ n-2:-1:2
        d[k-1] = d[k+1] + 2*(k-1)*coef[k]
    end
    #interval width factor
    d .*= 2/(xb - xa)
    #extra DC scaling
    d[1] /= 2
    return d
end

#-------------------------------------------------------------------------------
# caching function for inverted cheby matrices, needed for interpolator setup

@memoize invertedchebymatrix(n::Int64)::Matrix{Float64} = inv(chebymatrix(n))

#-------------------------------------------------------------------------------
# one-dimensional interpolation

struct ChebyshevInterpolator{N,T}
    #lowest value in range
    xa::T
    #highest value in range
    xb::T
    #interpolation coefficents
    a::NTuple{N,T}
    #must always have strict boundaries
    boundaries::StrictBoundaries
end

function Base.show(io::IO, ϕ::ChebyshevInterpolator)
    print(io, "$(typeof(ϕ)) ∈ [$(ϕ.xa), $(ϕ.xb)]")
end

"""
    ChebyshevInterpolator(x, y)

Construct a `ChebyshevInterpolator` for the points defined by coordinates `x` and values `y`. The `x` coordinates *must* be arranged on a chebyshev grid, which can be generated using the [`chebygrid`](@ref) function.
"""
function ChebyshevInterpolator(x, y)
    #same types
    T = promote_type(eltype(x), eltype(y))
    x = collect(T, x)
    y = collect(T, y)
    #check for basic issues
    rangecheck(x, y, 3)
    #demand that the input points have chebyshev spacing
    @assert ischebygrid(x) "points must be on a chebyshev grid"
    #generate expansion coefficients
    a = Tuple(chebycoef(y))
    #construct
    ChebyshevInterpolator(minimum(x), maximum(x), a, StrictBoundaries())
end

"""
    ChebyshevInterpolator(f, xa, xb, n)

Construct a `ChebyshevInterpolator` for the function `f` using `n` function evaluations in the range [`xa`,`xb`]. The function evaluations will occur on the chebyshev nodes.
"""
function ChebyshevInterpolator(f::F, xa, xb, n::Int) where {F}
    #set up the range coordinates
    x = chebygrid(xa, xb, n)
    #evaluate the function at those coordinates
    y = f.(x)
    #call the other constructor
    ChebyshevInterpolator(x, y)
end

function (ϕ::ChebyshevInterpolator{N,U})(x) where {N,U}
    #always enforce boundaries
    ϕ.boundaries(x, ϕ.xa, ϕ.xb)
    #evaluate the expansion
    return cheby(ϕ.a, x, ϕ.xa, ϕ.xb)
end

"""
    chebyderiv(ϕ::ChebyshevInterpolator)

Construct a ChebyshevInterpolator representing the derivative of a preexisting interpolator.
"""
function chebyderiv(ϕ::ChebyshevInterpolator)
    #derivative expansion's coefficients
    coef = Tuple(chebyderiv(ϕ.a, ϕ.xa, ϕ.xb))
    #new interpolator
    ChebyshevInterpolator(ϕ.xa, ϕ.xb, coef, StrictBoundaries())
end

#-------------------------------------------------------------------------------
# bichebyshev interpolation, a little trickier now!

struct BichebyshevInterpolator{M,N,U}
    #lowest value on axis 1
    xa::U
    #highest value on axis 1
    xb::U
    #lowest value on axis 2
    ya::U
    #highest value on axis 2
    yb::U
    #matrix and vectors for doing the interpolation
    A::Matrix{U} # size (ny by nx) or (M by N)
    a::Vector{U} # length ny for cosine expansion in θy
    b::Vector{U} # length nx for cosine expansion in θx
    c::Vector{U} # length ny for doing M*b in place
    #must always use strict boundaries
    boundaries::StrictBoundaries
end

function Base.show(io::IO, ϕ::BichebyshevInterpolator)
    print(io, "$(typeof(ϕ)) ∈ [$(ϕ.xa), $(ϕ.xb)], [$(ϕ.ya), $(ϕ.yb)]")
end

"""
    BichebyshevInterpolator(x, y, Z)

Construct a `BichebyshevInterpolator` for the grid of points defined by coordinates (`x`,`y`) and values `Z`. The given points must lie on a chebyshev grid in each direction. These can be generated with the [`chebygrid`](@ref) function or the interpolator can be constructed directly from a function using the method below.

!!! warning

    The Bichebyshev interpolator is *not thread-safe*. It computes a cosine expansion and does some linear algebra in-place using arrays stored with the object. A single `BichebyshevInterpolator` should never be called by multiple threads at once.
"""
function BichebyshevInterpolator(x, y, Z)
    #same types
    T = promote_type(eltype(x), eltype(y), eltype(Z))
    x = collect(T, x)
    y = collect(T, y)
    Z = collect(T, Z)
    #check for basic grid problems
    gridcheck(x, y, Z, 3)
    #grid properties
    nx, ny = length(x), length(y)
    xa, xb = T(minimum(x)), T(maximum(x))
    ya, yb = T(minimum(y)), T(maximum(y))
    #reject any non-cheby grid spacing
    @assert ischebygrid(x) "axis 1 coordinates must be on a chebyshev grid"
    @assert ischebygrid(y) "axis 2 coordinates must be on a chebyshev grid"
    #get inverted matrices from cache or generate them
    B = invertedchebymatrix(nx)
    #generate interpolation coefficients along axis 1 for each value of axis 2
    α = zeros(T, nx, ny)
    for j = 1:ny
        mul!(view(α,:,j), B, view(Z,:,j))
    end
    #then combine α and A
    A = invertedchebymatrix(ny)*α'
    #other vectors we need for doing the actual interpolation
    a = ones(T, ny)
    b = ones(T, nx)
    c = zeros(T, ny)
    #done
    BichebyshevInterpolator{ny,nx,T}(xa, xb, ya, yb, A, a, b, c, StrictBoundaries())
end

"""
    BichebyshevInterpolator(f, xa, xb, nx, ya, yb, ny)

Construct a `BichebyshevInterpolator` for the function `f` using a grid of `nx` points on the first axis in [`xa`,`xb`] and `ny` points on the second axis in [`ya`,`yb`].
"""
function BichebyshevInterpolator(f::F, xa, xb, nx::Int, ya, yb, ny::Int) where {F}
    #set up the grid
    X, Y = chebygrid(xa, xb, nx, ya, yb, ny)
    #evaluate the function at chebyshev grid points
    Z = f.(X, Y)
    #call the other constructor
    BichebyshevInterpolator(X[:,1], Y[1,:], Z)
end

#=====
This is the fast implementation. It's executed when the types of the
input coordinates match the type of the stored coefficients in the
interpolator. When the types match, the Chebyshev expansions can be
evaluated in-place, using vectors pre-allocated in the interpolator.
See the a, b, and c fields of the struct. This method also guarantees
that the interpolator type and the coordinate types are <: AbstractFloat,
making the low-level linear algebra functions safe. Without such a
guarantee, there can be issues with, for example, the dual numbers in
FowardDiff routines.
=====#
function (Φ::BichebyshevInterpolator{M,N,U})(x::U, y::U) where {M,N,U<:AbstractFloat}
    #always enforce boundaries for Chebyshev
    Φ.boundaries(x, Φ.xa, Φ.xb, y, Φ.ya, Φ.yb)
    #evaluate Chebyshev polys at the coordinates recursively and in-place
    ξy = x2ξ(y, Φ.ya, Φ.yb)
    chebyrecurrance!(Φ.a, ξy, M)
    ξx = x2ξ(x, Φ.xa, Φ.xb)
    chebyrecurrance!(Φ.b, ξx, N)
    #perform M*b, which interpolates along the first axis, also in-place
    mul!(Φ.c, Φ.A, Φ.b)
    #then a'*c interpolates along the second axis
    return dot(Φ.a, Φ.c)
end

#=====
This is the slow, but type-flexible, implementation. It's
executed whenever the type of the interpolator's coefficients or
the coordinate types don't match OR are not AbstractFloats. The
price of this flexibility is allocations for the expansions and
loss of the low-level linear algebra routines 😭.
=====#
function (Φ::BichebyshevInterpolator{M,N,U})(x, y) where {M,N,U}
    #always enforce boundaries for Chebyshev
    Φ.boundaries(x, Φ.xa, Φ.xb, y, Φ.ya, Φ.yb)
    #coordinates in ξ space
    ξx, ξy = promote(x2ξ(x, Φ.xa, Φ.xb), x2ξ(y, Φ.ya, Φ.yb))
    #allocating expansion in the x direction
    b = chebyrecurrance(ξx, N)
    c = Φ.A * b
    #then perform the y-axis expansion and dot product simulataneously
    Tₖ₋₂ = one(ξy)
    Tₖ₋₁ = ξy
    @inbounds z = c[1] + ξy*c[2]
    for i = 3:M
        #compute next value
        Tₖ = 2ξy*Tₖ₋₁ - Tₖ₋₂
        #update running dot product
        @inbounds z += c[i]*Tₖ
        #swaps
        Tₖ₋₂ = Tₖ₋₁
        Tₖ₋₁ = Tₖ
    end
    return z
end


text:
######################################################################################################
# macros from jpeglib.h

jpeg_create_compress(cinfo) =
    jpeg_CreateCompress(cinfo, JPEG_LIB_VERSION, sizeof(jpeg_compress_struct))
jpeg_create_decompress(cinfo) =
    jpeg_CreateDecompress(cinfo, JPEG_LIB_VERSION, sizeof(jpeg_decompress_struct))


text:
function asin!(X::IntervalBox)
    x, y = X

    h = half_pi.lo
    y_new = y ∩ Interval(-h, h)  # range of asin
    x_new = sin(y_new)

    return IntervalBox(x_new, y_new)
end

"""
Reverse `asin`.
"""
function asin_rev(y::Interval, x::Interval)  # y = asin(x)

    h = half_pi.lo
    y_new = y ∩ Interval(-h, h)  # range of asin

    x_new = sin(y_new)

    return y_new, x_new  # return in order y, x
end

"""
Reverse `acos`.
"""
function acos_rev(y::Interval, x::Interval)
        y_new = y ∩ Interval(0.0,two_pi.hi)
        x_new = x ∩ cos(y_new)

        return y_new, x_new
end

"""
    atan_rev(y::Interval, x::Interval)

Inverse of `y = atan(x)`.
Returns the new `y` and `x`.
"""
function atan_rev(y::Interval, x::Interval)
        y_new = y ∩ Interval(-half_pi.hi, half_pi.hi)
        x_new = x ∩ tan(y_new)

        return y_new, x_new
end


text:
struct Callback{F, S}
    func :: F
    schedule :: S
end

(callback::Callback)(sim) = callback.func(sim)

"""
    Callback(func, schedule)

Return `Callback` that executes `func(sim::Simulation)` on `schedule`.

`schedule = IterationInterval(1)` by default.
"""
Callback(func) = Callback(func, IterationInterval(1))


text:
# Wrap PlotPolaron (and it's dependency on Plots) within it's own (sub)module.
module PlotPolaron

export plotpolaron # Should unify case here c.f. Module

using PolaronMobility
using Plots

function plotpolaron(fileprefix, p::Polaron; extension="png")
    println("Plotting polaron to $fileprefix...")

    #####
    ## Mass vs. Temperature plot
    plot(p.T,p.M,label="Phonon effective-mass",
         markersize=3,marker=:rect,
         xlab="Temperature (K)",ylab="Phonon effective-mass",
         ylim=(0,1.2))

    savefig("$fileprefix-mass.$extension")

    #####
    ## Relaxationtime vs. Temperature plot
    plot(p.T,p.Tau,label="Kadanoff relaxation time (ps)",markersize=3,marker=:rect,xlab="Temperature (K)",ylab="Relaxation time (ps)",ylim=(0,1.2))

    savefig("$fileprefix-tau.$extension")

    ## Mass + relaxation time vs. Temperature plot
    plot(p.T,p.M,label="Phonon effective-mass (m\$_b\$)",markersize=3,marker=:rect,
        xlab="Temperature (K)",ylab="Effective-mass / relaxation time",ylim=(0,1.2))
    plot!(p.T,p.Tau,label="Kadanoff relaxation time (ps)",markersize=3,marker=:diamond,
        xlab="Temperature (K)",ylab="Relaxation time (ps)",ylim=(0,1.2))

    savefig("$fileprefix-mass-tau.$extension")

    ####
    ## Variational parameters, v and w vs. Temperature plot
    plot(p.T,p.v,label="v",markersize=3, marker=:rect, xlab="Temperature (K)",ylab="hbar-omega")
    plot!(p.T,p.w,label="w",markersize=3, marker=:diamond)

    savefig("$fileprefix-vw.$extension")

    #####
    ## Spring Constants vs. Temperature plot
    plot(p.T,p.k,label="Polaron spring-constant",markersize=3, marker=:uptriangle, xlab="Temperature (K)",ylab="Spring-constant",)

    savefig("$fileprefix-spring.$extension")

    #####
    ## Variation Energy vs. Temperature plots
    plot( p.T,p.A,label="A",markersize=3,marker=:downtriangle, xlab="Temperature (K)",ylab="Polaron free-energy")
    plot!(p.T,p.B,label="B",markersize=3,marker=:diamond)
    plot!(p.T,p.C,label="C",markersize=3,marker=:uptriangle)
    plot!(p.T,p.F,label="F",markersize=3,marker=:rect)
    #plot!(Ts,Fs,label="F=-(A+B+C)",markersize=3,marker=:rect)

    savefig("$fileprefix-variational.$extension")

    #####
    ## Polaron radius vs. Temperature
    plot(p.T,p.rfsi.*10^10, markersize=3,marker=:rect,
        label="Polaron radius",xlab="Temperature (K)",ylab="Polaron Radius (Angstrom)",ylims=(0,Inf))
#    plot!(p.T,p.rfsmallalpha.*10^10,label="T=0 Schultz small alpha polaron radius") # obsolete
    savefig("$fileprefix-radius.$extension")

    #####
    ## Calculated mobility comparison plot
    plot(p.T,p.Kμ,label="Kadanoff",markersize=3,marker=:rect,xlab="Temperature (K)",ylab="Mobility (cm\$^2\$/Vs)",ylims=(0,1000))
    plot!(p.T,p.FHIPμ,label="FHIP",markersize=3,marker=:diamond)
    plot!(p.T,p.Hμ,label="Hellwarth1999",markersize=3,marker=:uptriangle)

    savefig("$fileprefix-mobility-calculated.$extension")
end

end


text:
module DictIO

export writedict, readdict

"Save the type and data of the passed dictionary as text."
function writedict(f::IOStream, dict::Dict)
    println(f, join(match(r"[A-z]+{([^,]+),([^}]+)}", string(typeof(dict))).captures, '\t'))
    for (k,v) in dict
        println(f, string(k)*"\t"*string(v))
    end
end
function writedict(fileName::AbstractString, dict::Dict)
    f = open(fileName, "w")
    writedict(f, dict)
    close(f)
end

"""
Read the type and data of the passed dictionary.

Numeric types are converted automatically, and string type are
preserved, but other types may fail to load.
"""
function readdict(f::IOStream; reversed=false)
    parts = split(strip(readline(f)), '\t')
    type1 = eval(parse(parts[1]))
    type2 = eval(parse(parts[2]))
    d = reversed ? Dict{type2,type1}() : Dict{type1,type2}()
    for line in eachline(f)
        k,v = split(strip(line), '\t', limit=2)
        if type1 <: Number
            k = parse(type1, k)
        end
        if type2 <: Number
            v = parse(type2, v)
        end

        reversed ? d[v] = k : d[k] = v
    end
    d
end
function readdict(fileName::AbstractString; reversed=false)
    f = open(fileName)
    d = readdict(f, reversed=reversed)
    close(f)
    d
end

end # module


text:
"""
    abstract type AbstractCompiler <: Any
Supertype for Clang compilers.
"""
abstract type AbstractCompiler end

"""
    abstract type AbstractIRGenerator <: AbstractCompiler
Supertype for LLVM IR generators.
"""
abstract type AbstractIRGenerator <: AbstractCompiler end

# interface
get_instance(x::AbstractIRGenerator) = x.instance
get_context(x::AbstractIRGenerator) = x.ts_ctx

"""
    struct IRGenerator <: AbstractIRGenerator
LLVM IR Generator via Clang's `LLVMOnlyAction`.
"""
struct IRGenerator <: AbstractIRGenerator
    ts_ctx::ThreadSafeContext
    instance::CompilerInstance
    act::LLVMOnlyAction
end

get_module(x::IRGenerator) = takeModule(x.act)

function IRGenerator(src::String, args::Vector{String}; diag_show_colors=true)
    ts_ctx = ThreadSafeContext()
    ctx = context(ts_ctx)
    instance = CompilerInstance()
    # diagnostics
    setShowPresumedLoc(instance, true)
    setShowColors(instance, diag_show_colors)
    createDiagnostics(instance)
    diag = getDiagnostics(instance)
    # invocation
    if Sys.isapple()
        # FIXME: Duplicate definition of symbol '___cxa_atexit'
        # icxxabi_inc = joinpath(@__DIR__, "..", "abi") |> normpath
        # insert!(args, length(args), "-isystem$icxxabi_inc")
        # insert!(args, length(args), "-includeicxxabi.h")
    elseif Sys.islinux()
        hack_inc = joinpath(@__DIR__, "..", "hack") |> normpath
        insert!(args, length(args), "-isystem$hack_inc")
        insert!(args, length(args), "-include"*"hack_linux.h")
    end
    invok = create_compiler_invocation_from_cmd(src, args, diag)
    setInvocation(instance, invok)
    # codegen action
    act = LLVMOnlyAction(ctx)
    ExecuteAction(instance, act)
    return IRGenerator(ts_ctx, instance, act)
end

function dispose(x::IRGenerator)
    dispose(x.instance)
    dispose(x.act)
    dispose(x.ts_ctx)
end

"""
    struct CXCompiler <: AbstractCompiler
[`IRGenerator`](@ref) + [`LLJIT`](@ref).
"""
struct CXCompiler <: AbstractCompiler
    irgen::IRGenerator
    jit::LLJIT
end

get_instance(x::CXCompiler) = get_instance(irgen)
get_context(x::CXCompiler) = get_context(x.irgen)
get_module(x::CXCompiler) = get_module(x.irgen)
get_jit(x::CXCompiler) = x.jit
get_dylib(x::CXCompiler) = JITDylib(x.jit)
get_codegen(x::CXCompiler) = x.irgen

function link_process_symbols(cc::CXCompiler)
    jd = get_dylib(cc)
    jit = get_jit(cc)
    prefix = LLVM.get_prefix(jit)
    dg = LLVM.CreateDynamicLibrarySearchGeneratorForProcess(prefix)
    add!(jd, dg)
end

function compile(cc::CXCompiler)
    ts_mod = ThreadSafeModule(get_module(cc); ctx=get_context(cc))
    jd = get_dylib(cc)
    jit = get_jit(cc)
    add!(jit, jd, ts_mod)
end

function dispose(x::CXCompiler)
    dispose(x.irgen)
    dispose(x.jit)
end

"""
    mutable struct IncrementalIRGenerator <: AbstractIRGenerator
Incremental LLVM IR Generator.
"""
mutable struct IncrementalIRGenerator <: AbstractIRGenerator
    ts_ctx::ThreadSafeContext
    instance::CompilerInstance
    parser::Parser
    modules::Vector{LLVM.Module}
    current_module::Int
    src_counter::Int
end

getParser(x::IncrementalIRGenerator) = x.parser
get_modules(x::IncrementalIRGenerator) = x.modules
get_current_module(x::IncrementalIRGenerator) = x.modules[x.current_module]

function IncrementalIRGenerator(src::String, args::Vector{String}; diag_show_colors=true)
    ts_ctx = ThreadSafeContext()
    ctx = context(ts_ctx)
    instance = CompilerInstance()
    # diagnostics
    setShowPresumedLoc(instance, true)
    setShowColors(instance, diag_show_colors)
    createDiagnostics(instance)
    diag = getDiagnostics(instance)
    # invocation
    invok = create_compiler_invocation_from_cmd(src, args, diag)
    setInvocation(instance, invok)
    setTargetAndLangOpts(instance)
    # source
    createFileManager(instance)
    createSourceManager(instance)
    setMainFileID(instance, src)
    # preprocessor & AST & sema
    createPreprocessor(instance, CXTranslationUnitKind_TU_Prefix)
    createASTContext(instance)
    codegen = CreateLLVMCodeGen(instance, ctx, "JLCC_Incremental_1")
    setASTConsumer(instance, codegen)
    createSema(instance, CXTranslationUnitKind_TU_Prefix)
    # parser
    preprocessor = getPreprocessor(instance)
    enable_incremental(preprocessor)
    sema = getSema(instance)
    parser = Parser(preprocessor, sema)
    begin_diag(instance)
    try
        InitializeBuiltins(preprocessor)
        EnterMainSourceFile(preprocessor)
        Initialize(parser)

        ast_ctx = getASTContext(instance)

        if tryParseAndSkipInvalidOrParsedDecl(parser, codegen)
            processWeakTopLevelDecls(sema, codegen)
            HandleTranslationUnit(codegen, ast_ctx)
        end
    finally
        end_diag(instance)
    end

    # generate llvm modules
    m_cur = release_llvm_module(codegen)
    m_cur.ref == C_NULL && error("failed to generate IR.")

    m_next = start_llvm_module(codegen, context(m_cur), "JLCC_Incremental_2")

    return IncrementalIRGenerator(ts_ctx, instance, parser, [m_cur, m_next], 2, 1)
end

function incremental_parse(irgen::IncrementalIRGenerator, code::String)
    ci = get_instance(irgen)
    parser = getParser(irgen)
    pp = getPreprocessor(ci)
    src_mgr = getSourceManager(ci)
    codegen = getASTConsumer(ci)
    buffer = get_buffer(code)
    fid = FileID(src_mgr, buffer)
    begin_diag(ci)
    loc = get_loc_with_offset(get_main_file_begin_loc(src_mgr), irgen.src_counter)
    EnterSourceFile(pp, fid, loc)
    try
        if tryParseAndSkipInvalidOrParsedDecl(parser, codegen)
            processWeakTopLevelDecls(getSema(ci), codegen)
            HandleTranslationUnit(codegen, getASTContext(ci))
        end
    finally
        dispose(fid)
        EndSourceFile(pp)
        end_diag(ci)
        irgen.src_counter += 1
    end
    # generate llvm modules
    m_cur = release_llvm_module(codegen)
    m_cur.ref == C_NULL && error("failed to generate IR.")
    irgen.current_module += 1
    m_next = start_llvm_module(codegen, context(m_cur), "JLCC_Incremental_$(irgen.current_module)")
    push!(irgen.modules, m_next)

    return m_cur
end

function dispose(x::IncrementalIRGenerator)
    dispose(x.parser)
    dispose(x.instance)
    pop!(x.modules)  # we don't have the ownership of the latest module
    dispose.(x.modules)
    dispose(x.ts_ctx)
end

function parse_cxx_scope_spec(irgen::IncrementalIRGenerator, str::String, spec::CXXScopeSpec)
    ci = get_instance(irgen)
    parser = getParser(irgen)
    parse_cxx_scope_spec(ci, parser, str, spec)
end

function (x::DeclFinder)(irgen::IncrementalIRGenerator, decl::String, scope::String="")
    ci = get_instance(irgen)
    parser = getParser(irgen)
    return x(ci, parser, decl, scope)
end


text:
#InspectDR: Event handlers for user input of Gtk layer
#-------------------------------------------------------------------------------

import Gtk: GConstants.GdkModifierType, Gtk.GConstants.GdkScrollDirection

#==Constants
===============================================================================#
const MODIFIER_ALT = GdkModifierType.GDK_MOD1_MASK #Is this bad? How do I query?
const MODIFIER_CTRL = GdkModifierType.GDK_CONTROL_MASK
const MODIFIER_SHIFT = GdkModifierType.GDK_SHIFT_MASK

#Other modifiers are ignored... ex: numlock (Mod2)
const MODIFIERS_SUPPORTED = MODIFIER_ALT|MODIFIER_CTRL|MODIFIER_SHIFT

#Symbols not currently supported by Gtk package:
const GdkKeySyms_Plus = 0xffab
const GdkKeySyms_Minus = 0xffad

#==Types
===============================================================================#
typealias KeyMap Dict{Int, Function}

#User input states
immutable ISNormal <: InputState; end
immutable ISSelectingArea <: InputState; end
immutable ISPanningData <: InputState; end #Typically with mouse

type KeyBindings
	nomod::KeyMap
	shiftmod::KeyMap
	ctrlmod::KeyMap
	altmod::KeyMap
end
KeyBindings() = KeyBindings(KeyMap(), KeyMap(), KeyMap(), KeyMap())

#==Module-level definitions
===============================================================================#
const keybindings = KeyBindings()

#==Helper functions
===============================================================================#
modifiers_pressed(eventstate, modmask) = (modmask == (MODIFIERS_SUPPORTED & eventstate))

function userinput_setstate_normal(pwidget::PlotWidget)
	pwidget.state = ISNormal()
end
userinput_donothing(pwidget::PlotWidget) = nothing

#==Wrapper functions
===============================================================================#
function raiseevent_plothover(pwidget::PlotWidget, event::Gtk.GdkEventMotion)
	ext = getextents_xfrm(pwidget.src)
	xf = Transform2D(ext, pwidget.graphbb)
	pt = ptmap_rev(xf, Point2D(event.x, event.y))
	raiseevent(pwidget.eh_plothover, pwidget, DReal(pt.x), DReal(pt.y))
	nothing
end

#==Default event handlers
===============================================================================#
handleevent_mousepress(::InputState, pwidget::PlotWidget, event::Gtk.GdkEventButton) =
	nothing
handleevent_mouserelease(::InputState, pwidget::PlotWidget, event::Gtk.GdkEventButton) =
	nothing
function handleevent_mousemove(::InputState, pwidget::PlotWidget, event::Gtk.GdkEventMotion)
	raiseevent_plothover(pwidget, event)
	nothing
end

#Event handler for plot widget keypress
#-------------------------------------------------------------------------------
function handleevent_keypress(::InputState, pwidget::PlotWidget, event::Gtk.GdkEventKey)
#event.event_type: GDK_KEY_PRESS/GDK_KEY_RELEASE
#@show event.event_type, event.state, event.keyval, event.length, event.string, event.hardware_keycode, event.group, event.flags
	global keybindings

	#Find bound functions & execute:
	if modifiers_pressed(event.state, 0)
		fn = get(keybindings.nomod, event.keyval, userinput_donothing)
		fn(pwidget)
	elseif modifiers_pressed(event.state, MODIFIER_SHIFT)
		fn = get(keybindings.shiftmod, event.keyval, userinput_donothing)
		fn(pwidget)
	elseif modifiers_pressed(event.state, MODIFIER_CTRL)
		fn = get(keybindings.ctrlmod, event.keyval, userinput_donothing)
		fn(pwidget)
	elseif modifiers_pressed(event.state, MODIFIER_ALT)
		fn = get(keybindings.altmod, event.keyval, userinput_donothing)
		fn(pwidget)
	end
#	@show event.keyval
end

function handleevent_mousescroll(::InputState, pwidget::PlotWidget, event::Gtk.GdkEventScroll)
	if modifiers_pressed(event.state, 0)
		if GdkScrollDirection.GDK_SCROLL_UP == event.direction
			pan_up(pwidget)
		elseif GdkScrollDirection.GDK_SCROLL_DOWN == event.direction
			pan_down(pwidget)
		end
	elseif modifiers_pressed(event.state, MODIFIER_SHIFT)
		if GdkScrollDirection.GDK_SCROLL_UP == event.direction
			pan_left(pwidget)
		elseif GdkScrollDirection.GDK_SCROLL_DOWN == event.direction
			pan_right(pwidget)
		end
	elseif modifiers_pressed(event.state, MODIFIER_CTRL)
		if GdkScrollDirection.GDK_SCROLL_UP == event.direction
			zoom_in(pwidget, event.x, event.y)
		elseif GdkScrollDirection.GDK_SCROLL_DOWN == event.direction
			zoom_out(pwidget, event.x, event.y)
		end
	end
end
function handleevent_scalechanged(::InputState, pwidget::PlotWidget)
	if !scalectrl_enabled(pwidget); return; end #Should not really happen
	scalectrl_apply(pwidget)
end

#==State: Normal input state
===============================================================================#
function handleevent_mousepress(::ISNormal, pwidget::PlotWidget, event::Gtk.GdkEventButton)
#	@show event.state, event.button, event.event_type
	if 3==event.button
		boxzoom_setstart(pwidget, event.x, event.y)
		pwidget.state = ISSelectingArea()
	elseif 1==event.button
		if modifiers_pressed(event.state, MODIFIER_SHIFT)
			mousepan_setstart(pwidget, event.x, event.y)
			pwidget.state = ISPanningData()
		end
	end

	focus(pwidget.widget) #In case not in focus
end

#==State: Selecting plot area
===============================================================================#
function handleevent_keypress(::ISSelectingArea, pwidget::PlotWidget, event::Gtk.GdkEventKey)
	if GdkKeySyms.Escape == event.keyval
		boxzoom_cancel(pwidget)
		pwidget.state = ISNormal()
	elseif 'h' == event.keyval
		locdir_h(pwidget)
	elseif 'v' == event.keyval
		locdir_v(pwidget)
	elseif 'b' == event.keyval
		locdir_any(pwidget)
	end
end
function handleevent_mouserelease(::ISSelectingArea, pwidget::PlotWidget, event::Gtk.GdkEventButton)
	if 3==event.button
		boxzoom_complete(pwidget, event.x, event.y)
		pwidget.state = ISNormal()
	end
end
function handleevent_mousemove(::ISSelectingArea, pwidget::PlotWidget, event::Gtk.GdkEventMotion)
	raiseevent_plothover(pwidget, event)
	boxzoom_setend(pwidget, event.x, event.y)
end

#==State: Mouse pan
===============================================================================#
function handleevent_keypress(::ISPanningData, pwidget::PlotWidget, event::Gtk.GdkEventKey)
	if GdkKeySyms.Escape == event.keyval
		mousepan_cancel(pwidget)
		pwidget.state = ISNormal()
	elseif 'h' == event.keyval || 'H' == event.keyval
		locdir_h(pwidget)
	elseif 'v' == event.keyval || 'V' == event.keyval
		locdir_v(pwidget)
	elseif 'b' == event.keyval || 'B' == event.keyval
		locdir_any(pwidget)
	end
end
function handleevent_mouserelease(::ISPanningData, pwidget::PlotWidget, event::Gtk.GdkEventButton)
	if 1==event.button
		mousepan_complete(pwidget, event.x, event.y)
		pwidget.state = ISNormal()
	end
end
function handleevent_mousemove(::ISPanningData, pwidget::PlotWidget, event::Gtk.GdkEventMotion)
	raiseevent_plothover(pwidget, event)
	mousepan_move(pwidget, event.x, event.y)
end

#==Key bindings
===============================================================================#
function keybindings_setdefaults(bnd::KeyBindings)
	bnd.nomod = KeyMap(
		GdkKeySyms.Escape => userinput_setstate_normal,
		GdkKeySyms.Up => pan_up,
		GdkKeySyms.Down => pan_down,
		GdkKeySyms.Left => pan_left,
		GdkKeySyms.Right => pan_right,
		GdkKeySyms_Plus => zoom_in,
		GdkKeySyms_Minus => zoom_out,
	)
	bnd.shiftmod = KeyMap(
	)
	bnd.ctrlmod = KeyMap(
		'f' => zoom_full,
		'h' => zoom_hfull,
		'v' => zoom_vfull,
	)
	bnd.altmod = KeyMap(
	)
end

#Last line


text:
struct MultiHeadAttentionLayer{L,I,B} <: MessagePassing
    query_layer::L
    key_layer::L
    value_layer::L
    edge_projection_layer::L
    node_output_layer::L
    in_dim::I
    out_dim::I
    num_heads::I
    use_bias::B
end

function MultiHeadAttentionLayer(
    in_dim::Int,
    out_dim::Int,
    num_heads::Int;
    bias=true,
    init=glorot_uniform
    )
    query_layer = Flux.Dense(in_dim, out_dim * num_heads, bias=bias, init=init)
    key_layer = Flux.Dense(in_dim, out_dim * num_heads, bias=bias, init=init)
    value_layer = Flux.Dense(in_dim, out_dim * num_heads, bias=bias, init=init)
    edge_projection_layer = Flux.Dense(in_dim, out_dim * num_heads, bias=bias, init=init)
    # The final projection layer for the output of attention module
    node_output_layer = Flux.Dense(out_dim * num_heads, out_dim * num_heads, bias=bias, init=init)
    return MultiHeadAttentionLayer(query_layer,
                                    key_layer,
                                    value_layer,
                                    edge_projection_layer,
                                    node_output_layer,
                                    in_dim, out_dim, num_heads, bias)
end

@functor MultiHeadAttentionLayer

# The message function
"""
Arguments:
- `x_i`: Feature vector of i'th node
- `x_j`: Feature vector of j'th node, where j is in the neighbourhood of i
- `e_ij`: Feature vector for the edge between i and j
"""
function message(mha::MultiHeadAttentionLayer, i, j, edge_index_map, X, E)
    x_i = _view(X, i)
    x_j = _view(X, j)

    key = mha.key_layer(x_i)
    query = mha.query_layer(x_j)
    value = mha.value_layer(x_j)

    if i != j
        edge_index = edge_index_map[(i, j)]
        edge_feat = _view(E, edge_index)
    else
        # If it is self edge then just use ones
        edge_feat = ones(size(value))
    end

    # Reshape - [out_dim, num_heads]
    key = reshape(key, :, mha.num_heads)
    query = reshape(query, :, mha.num_heads)
    value = reshape(value, :, mha.num_heads)
    edge_feat = reshape(edge_feat, :, mha.num_heads)

    # Add edge features
    scores = sum(key .* query .* edge_feat, dims=1) ./ sqrt(mha.out_dim) # shape - [1, num_heads]
    # scores = scores .* edge_feats
    return vcat(scores, value)
end

function edge_message(mha::MultiHeadAttentionLayer, i, j, edge_index_map, E)
    if i != j
        edge_index = edge_index_map[(i, j)]
        edge_feat = _view(E, edge_index)
    else
        # If it is self edge then just use ones
        edge_feat = zeros(mha.out_dim * mha.num_heads)
    end
    return edge_feat
end

function apply_batch_message(mha::MultiHeadAttentionLayer, i, js, edge_index_map, X, E)
    scores = mapreduce(j -> message(mha, i, j, edge_index_map, X, E), hcat, js)
    nodes = size(scores, 1)
    # score shape - [num_heads, num_nodes_in_neighbourhood]
    reshaped_scores = reshape(view(scores, 1, :), mha.num_heads, :)
    alphas = Flux.softmax(reshaped_scores, dims=2)
    # alpha shape - [1, num_heads, num_nodes_in_neighbourhood]
    alphas = reshape(alphas, 1, mha.num_heads, :)

    values = view(scores, 2:nodes, :)
    values = reshape(values, mha.out_dim, mha.num_heads, :) .* alphas
    # Reshape values to [out_dim * num_heads, num_nodes_in_neighbourhood]
    values = reshape(values, mha.out_dim * mha.num_heads, :)
    # Put the tensor with heads concatenated through a dense layer
    node_features = mha.node_output_layer(values)
    # Aggregate the values of all neighbour nodes
    node_features = dropdims(sum(node_features, dims=ndims(node_features)), dims=ndims(node_features))

    # Edge message passing
    edge_feats = mapreduce(j -> edge_message(mha, i, j, edge_index_map, E), hcat, js)
    edge_feats = reshape(edge_feats, mha.out_dim, mha.num_heads, :)
    edge_feats = edge_feats .* reshape(reshaped_scores, 1, mha.num_heads, :)
    edge_feats = reshape(edge_feats, mha.out_dim * mha.num_heads, :)
    # Edge projection layer
    edge_features = mha.edge_projection_layer(edge_feats)
    # edge_features = dropdims(sum(edge_feats, dims=ndims(edge_feats)), dims=ndims(edge_feats))

    return node_features, edge_features
end

function aggregate_neighbors(mha::MultiHeadAttentionLayer, adj, edge_index_map, E, V)
    n = size(adj, 1)
    # # a vertex must always receive a message from itself
    # Zygote.ignore() do
    #     GraphLaplacians.add_self_loop!(adj, n)
    # end
    multi_hcat = (args...) -> hcat.(args...)
    return mapreduce(i -> apply_batch_message(mha, i, adj[i], edge_index_map, V, E), multi_hcat, 1:n)
end

function propagate(mha::MultiHeadAttentionLayer, fg::FeaturedGraph)
    adj = adjacency_list(fg)

    # Get the edge map to retrieve edges
    edge_index_map = edge_index_table(adj, fg.directed)
    edge_feats, node_feats = fg.ef, fg.nf

    # Attention propagation
    node_feats, edge_feats = aggregate_neighbors(mha, adj, edge_index_map, edge_feats, node_feats)

    return node_feats, edge_feats
end

function (mha::MultiHeadAttentionLayer)(graph::FeaturedGraph)
    # Multi Headed Attention over neighbours
    return propagate(mha, graph)
end

max_stars_repo_name: DwaraknathT/Flux-Graph-Transformer

max_issues_repo_name: DwaraknathT/Flux-Graph-Transformer

max_forks_repo_name: DwaraknathT/Flux-Graph-Transformer


text:
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    intersecttype(b1, b2)

Compute the intersection type of two boxes `b1` and `b2`.

The intersection type can be one of four types:

1. overlap with non-zero measure
2. intersect at one of the boundaries
3. intersect at corner point
4. do not overlap nor intersect
"""
function intersecttype(b1::Box{Dim,T}, b2::Box{Dim,T}) where {Dim,T}
  m1, M1 = coordinates.(extrema(b1))
  m2, M2 = coordinates.(extrema(b2))

  # relevant vertices
  u = Point(max.(m1, m2))
  v = Point(min.(M1, M2))

  # branch on possible configurations
  if u ≺ v
    OverlappingBoxes(Box(u, v))
  elseif u ≻ v
    NoIntersection()
  elseif isapprox(u, v, atol=atol(T))
    CornerTouchingBoxes(u)
  else
    FaceTouchingBoxes(Box(u, v))
  end
end


text:
abstract FixedEffectProblem

##############################################################################
##
## A FixedEffectProblem must define three methods:
##
## solve_residuals!
## solve_coefficients!
## get_fes (accessor)
##
##############################################################################

##############################################################################
##
## get residuals
##
##############################################################################

function residualize!(r::AbstractVector{Float64}, fep::FixedEffectProblem,
                      iterationsv::Vector{Int}, convergedv::Vector{Bool};
                      kwargs...)
    r, iterations, converged = solve_residuals!(fep, r; kwargs...)
    push!(iterationsv, iterations)
    push!(convergedv, converged)
end

function residualize!(X::Matrix{Float64}, fep::FixedEffectProblem,
                      iterationsv::Vector{Int}, convergedv::Vector{Bool};
                      kwargs...)
    for j in 1:size(X, 2)
        residualize!(view(X, :, j), fep, iterationsv, convergedv; kwargs...)
    end
end

function residualize!(::Array, ::Void,
                      ::Vector{Int}, ::Vector{Bool};
                      kwargs...)
    nothing
end

##############################################################################
##
## Get fixed effects
##
## Fixed effects are generally not identified
## We standardize the solution in the following way :
## Mean within connected component of all fixed effects except the first
## is zero
##
## Unique solution with two components, not really with more
##
## Connected component : Breadth-first search
## components is an array of component
## A component is an array of set (length is number of values taken)
##
##############################################################################

function getfe!(fep::FixedEffectProblem, b::Vector{Float64}; kwargs...)

    # solve Ax = b
    x, iterations, converged = solve_coefficients!(fep, b; kwargs...)
    if !converged
       warn("getfe did not converge")
    end

    # The solution is generally not unique. Find connected components and scale accordingly
    findintercept = find(fe -> isa(fe.interaction, Ones), get_fes(fep))
    if length(findintercept) >= 2
        components = connectedcomponent(view(get_fes(fep), findintercept))
        rescale!(x, fep, findintercept, components)
    end

    return x
end

# Convert estimates to dataframes
function DataFrame(fev::Vector{Vector{Float64}}, fep::FixedEffectProblem, esample::BitVector)
    fes = get_fes(fep)
    newdf = DataFrame()
    len = length(esample)
    for j in 1:length(fes)
        name = fes[j].id
        T = eltype(fes[j].refs)
        refs = fill(zero(T), len)
        refs[esample] = fes[j].refs
        newdf[fes[j].id] = PooledDataArray(RefArray(refs), fev[j])
    end
    return newdf
end

function getfe!(fep::FixedEffectProblem, b::Vector{Float64},esample::BitVector;
                tol::Real = 1e-8, maxiter::Integer = 100_000)
    fev = getfe!(fep, b; tol = tol, maxiter = maxiter)
    return DataFrame(fev, fep, esample)
end

function connectedcomponent(fes::AbstractVector{FixedEffect})
    # initialize
    where = initialize_where(fes)
    refs = initialize_refs(fes)
    nobs = size(refs, 2)
    visited = fill(false, nobs)
    components = Vector{Set{Int}}[]
    # start
    @inbounds for i in 1:nobs
        if !visited[i]
            component = Set{Int}[Set{Int}() for fe in fes]
            connectedcomponent!(component, visited, i, refs, where)
            push!(components, component)
        end
    end
    return components
end

function initialize_where(fes::AbstractVector{FixedEffect})
    where = Vector{Set{Int}}[]
    for j in 1:length(fes)
        fe = fes[j]
        wherej = Set{Int}[Set{Int}() for fe in fe.scale]
        @inbounds for i in 1:length(fe.refs)
            push!(wherej[fe.refs[i]], i)
        end
        push!(where, wherej)
    end
    return where
end

function initialize_refs(fes::AbstractVector{FixedEffect})
    nobs = length(fes[1].refs)
    refs = fill(zero(Int), length(fes), nobs)
    @inbounds for j in 1:length(fes)
        ref = fes[j].refs
        for i in 1:length(ref)
            refs[j, i] = ref[i]
        end
    end
    return refs
end

# Breadth-first search
function connectedcomponent!(component::Vector{Set{Int}},
    visited::Vector{Bool}, i::Integer, refs::Matrix{Int},
    where::Vector{Vector{Set{Int}}})
    visited[i] = true
    tovisit = Set{Int}()
    # for each fixed effect
    @inbounds for j in 1:size(refs, 1)
        ref = refs[j, i]
        # if category has not been encountered
        if !(ref in component[j])
            # mark category as encountered
            push!(component[j], ref)
            # add other observations with same component in list to visit
            for k in where[j][ref]
                push!(tovisit, k)
            end
        end
    end
    for k in tovisit
        if k != i
            connectedcomponent!(component, visited, k, refs, where)
        end
    end
end

function rescale!(fev::Vector{Vector{Float64}}, fep::FixedEffectProblem,
                  findintercept,
                  components::Vector{Vector{Set{Int}}})
    fes = get_fes(fep)
    adj1 = zero(Float64)
    i1 = findintercept[1]
    @inbounds for component in components
        for i in reverse(findintercept)
            # demean all fixed effects except the first
            if i != 1
                adji = zero(Float64)
                for j in component[i]
                    adji += fev[i][j]
                end
                adji = adji / length(component[i])
                for j in component[i]
                    fev[i][j] -= adji
                end
                adj1 += adji
            else
                # rescale the first fixed effects
                for j in component[i1]
                    fev[i1][j] += adj1
                end
            end
        end
    end
end


text:
module Summation
export sum_naive, sum_kbn, sum_oro, sum_mixed
export dot_naive, dot_oro, dot_mixed

import VectorizationBase

using VectorizationBase: Vec, vload, vsum, vzero, AbstractSIMD, MM
fptype(::Type{V}) where {W, T, V <: AbstractSIMD{W, T}} = T
@inline f64(x::Number)     = convert(Float64, x)

using ..EFT: two_sum, fast_two_sum, two_prod

include("accumulators/sum.jl")
include("accumulators/dot.jl")
include("accumulators/compSum.jl")
include("accumulators/compDot.jl")
include("accumulators/mixedSum.jl")
include("accumulators/mixedDot.jl")

# T. Ogita, S. Rump and S. Oishi, "Accurate sum and dot product",
# SIAM Journal on Scientific Computing, 6(26), 2005.
# DOI: 10.1137/030601818
@generated function accumulate(x::NTuple{A, AbstractArray{T}},
                               accType::F,
                               rem_handling = Val(:scalar),
                               ::Val{Ushift} = Val(2),
                               ::Val{Prefetch} = Val(0),
                               )  where {F, A, T <: Union{Float32,Float64}, Ushift, Prefetch}
    @assert 0 ≤ Ushift < 6
    U = 1 << Ushift

    W, shift = VectorizationBase.pick_vector_width_shift(T)
    WU = W * U
    V = Vec{Int(W),T}

    quote
        $(Expr(:meta,:inline))
        px = VectorizationBase.zstridedpointer.(x)
        N = length(first(x))
        Base.Cartesian.@nexprs $U u -> begin
            acc_u = accType($V)
        end

        Nshift = N >> $(shift + Ushift)
        offset = MM{$(Int(W))}(0)
        for n in 1:Nshift
            if $Prefetch > 0
                # VectorizationBase.prefetch.(px.+offset.+$(Prefetch*WT), Val(3), Val(0))
                VectorizationBase.prefetch0.(px, offset + $(Prefetch*W))
            end

            Base.Cartesian.@nexprs $U u -> begin
                xi = vload.(px, tuple(tuple(offset))) # vload expects the index to be wrapped in a tuple, and broadcasting strips one layer
                add!(acc_u, xi...)
                offset += $W
            end
        end

        rem = N & $(WU-1)
        for n in 1:(rem >> $shift)
            xi = vload.(px, tuple(tuple(offset)))
            add!(acc_1, xi...)
            offset += $W
        end

        if $rem_handling <: Val{:mask}
            rem &= $(W-1)
            if rem > 0
                mask = VectorizationBase.mask($W, rem)
                xi = vload.(px, tuple(tuple(offset)), mask)
                add!(acc_1, xi...)
            end
        end

        Base.Cartesian.@nexprs $(U-1) u -> begin
            add!(acc_1, acc_{u+1})
        end

        acc = sum(acc_1)

        if $rem_handling <: Val{:scalar}
            _offset = VectorizationBase.data(offset)
            while _offset < N
                _offset += 1
                Base.Cartesian.@nexprs $A a -> begin
                    @inbounds xi_a = getindex(x[a], _offset)
                end
                Base.Cartesian.@ncall $A add! acc xi
            end
        end

        sum(acc)
    end
end

# Default values for unrolling
@inline default_ushift(::SumAcc)      = Val(3)
@inline default_ushift(::CompSumAcc)  = Val(2)
@inline default_ushift(::MixedSumAcc) = Val(3)
@inline default_ushift(::DotAcc)      = Val(3)
@inline default_ushift(::CompDotAcc)  = Val(2)
@inline default_ushift(::MixedDotAcc) = Val(3)
# dispatch
#   either default_ushift(x,    acc)
#   or     default_ushift((x,), acc)
@inline default_ushift(x::AbstractArray, acc) = default_ushift(acc(eltype(x)))
@inline default_ushift(x::NTuple, acc)        = default_ushift(first(x), acc)

# Default values for cache prefetching
@inline default_prefetch(::SumAcc)      = Val(0)
@inline default_prefetch(::CompSumAcc)  = Val(35)
@inline default_prefetch(::MixedSumAcc) = Val(0)
@inline default_prefetch(::DotAcc)      = Val(0)
@inline default_prefetch(::CompDotAcc)  = Val(20)
@inline default_prefetch(::MixedDotAcc) = Val(0)
# dispatch
#   either default_prefetch(x,    acc)
#   or     default_prefetch((x,), acc)
@inline default_prefetch(x::AbstractArray, acc) = default_prefetch(acc(eltype(x)))
@inline default_prefetch(x::NTuple, acc)        = default_prefetch(first(x), acc)

@inline _sum(x, acc) = if length(x) < 500
    # no cache prefetching for small vectors
    accumulate((x,), acc, Val(:scalar), default_ushift(x, acc), Val(0))
else
    accumulate((x,), acc, Val(:scalar), default_ushift(x, acc), default_prefetch(x, acc))
end

sum_naive(x) = _sum(x, sumAcc)
sum_kbn(x)   = _sum(x, compSumAcc(fast_two_sum))
sum_oro(x)   = _sum(x, compSumAcc(two_sum))
sum_mixed(x) = _sum(x, mixedSumAcc)

@inline _dot(x, y, acc) = if length(x) < 500
    # no cache prefetching for small vectors
    accumulate((x,y), acc, Val(:scalar), default_ushift(x, acc), Val(0))
else
    accumulate((x,y), acc, Val(:scalar), default_ushift(x, acc), default_prefetch(x, acc))
end

dot_naive(x, y) = _dot(x, y, dotAcc)
dot_oro(x, y)   = _dot(x, y, compDotAcc)
dot_mixed(x, y) = _dot(x, y, mixedDotAcc)

end


text:
using Documenter, LombScargle

# # Generate all images
# include(joinpath(@__DIR__, "src", "images.jl"))
cp(joinpath(@__DIR__, "..", "perf", "benchmarks.png"),
   joinpath(@__DIR__, "src", "benchmarks.png"))

makedocs(
    modules = [LombScargle],
    sitename = "LombScargle",
)

deploydocs(
    repo = "github.com/JuliaAstro/LombScargle.jl.git",
    target = "build",
    deps = nothing,
    make = nothing,
)


text:
include("../../fancy_toys.jl")

version = v"8.1.0"#.77

name = "CUDNN_CUDA$(cuda_version.major)$(cuda_version.minor)"

script = raw"""
mkdir -p ${libdir} ${prefix}/include

cd ${WORKSPACE}/srcdir
if [[ ${target} == x86_64-linux-gnu ]]; then
    cd cuda
    find .

    install_license NVIDIA_SLA_cuDNN_Support.txt

    mv lib64/libcudnn*.so* ${libdir}
    mv include/* ${prefix}/include
elif [[ ${target} == powerpc64le-linux-gnu ]]; then
    cd cuda/targets/ppc64le-linux
    find .

    install_license NVIDIA_SLA_cuDNN_Support.txt

    mv lib/libcudnn*.so* ${libdir}
    mv include/* ${prefix}/include
elif [[ ${target} == x86_64-w64-mingw32 ]]; then
    cd cuda
    find .

    install_license NVIDIA_SLA_cuDNN_Support.txt

    mv bin/cudnn*64_*.dll ${libdir}
    mv include/* ${prefix}/include
fi
"""

products = [
    LibraryProduct(["libcudnn", "cudnn64_$(version.major)"], :libcudnn),
]

dependencies = [Dependency(PackageSpec(name="CUDA_jll", version=cuda_version))]

non_reg_ARGS = filter(arg -> arg != "--register", ARGS)

if @isdefined(sources_linux_x64) && should_build_platform("x86_64-linux-gnu")
    build_tarballs(non_reg_ARGS, name, version, sources_linux_x64, script,
                   [Platform("x86_64", "linux")], products, dependencies)
end

if @isdefined(sources_linux_ppc64le) && should_build_platform("powerpc64le-linux-gnu")
    build_tarballs(non_reg_ARGS, name, version, sources_linux_ppc64le, script,
                   [Platform("powerpc64le", "linux")], products, dependencies)
end

if @isdefined(sources_windows) && should_build_platform("x86_64-w64-mingw32")
    build_tarballs(ARGS, name, version, sources_windows, script,
                   [Platform("x86_64", "windows")], products, dependencies)
end

max_stars_repo_name: blenessy/Yggdrasil

max_issues_repo_name: blenessy/Yggdrasil

max_forks_repo_name: blenessy/Yggdrasil


text:
"""
Simplest example of model with 1 variable, 1 shock and 1 transition equation.
It shows the boilerplate code for creating models.
"""
module E1

using ModelBaseEcon

model = Model()

@parameters model begin
    α = 0.5
    β = 0.5
end

@variables model y

@shocks model y_shk

@autoexogenize model y = y_shk

@equations model begin
    y[t] = α * y[t - 1] + β * y[t + 1] + y_shk[t]
end

@initialize model

end


text:
"""
	$(SIGNATURES)

Plots a set of regressions as a series of bar graph. Each bar is a regressor so that different scales are not a visual problem.
Standard errors are shown as error bars.

# Arguments

- coeffNameV:  Regressor names. Common across regressions. Symbols or Strings.
- regrLabelV:  Descriptive labels for the regressions. Become x labels under each bar. Symbols or Strings.
- coeffM: Coefficients by [regressor, regression].
- seM: Standard errors in the same format.
- `onePlot`: if `true`, all regressors are plotted in one grouped bar graph. Only works when scales are similar across regressors.
- `interceptAsText`: don't show a bar for intercept; show it as text in the graph (to avoid vastly different scales).
"""
function plot_regressions(coeffNameV, regrLabelV,
    coeffM :: AbstractMatrix{F1},
    seM :: AbstractMatrix{F1};
    onePlot :: Bool = false,
    interceptAsText = interceptAsText, interceptName = "cons",
    figTitle = nothing) where F1 <: AbstractFloat

    n, nRegr = size(coeffM);
    @assert length(coeffNameV) == n
    @assert length(regrLabelV) == nRegr

    p = blank_plot();
    if onePlot
        plot_regressions_together(coeffNameV, regrLabelV, coeffM, seM;
            interceptAsText = interceptAsText, interceptName = interceptName,
            fig = p, pos = (1,1), figTitle = figTitle);
    else
        # One plot per regressor
        nr, nc = subplot_layout(n);
        for iReg = 1 : n
            pos = subplot_pos(iReg, nr, nc);
            # No point having intercept as text here. It gets its own plot.
            plot_regression(string.(regrLabelV), coeffNameV[iReg],
                coeffM[iReg,:], seM[iReg,:]; fig = p, pos = pos,
                forSubPlot = true);
        end
        add_title!(p, figTitle);
    end
    return p
end

function plot_regression(xStrV, yStr, coeffV :: AbstractVector{F1},
    seV :: AbstractVector{F1};
    fig = blank_plot(), pos = (1,1),
    interceptAsText :: Bool = false,
    interceptName = :cons,
    kwargs...) where F1 <: AbstractFloat

    n = length(xStrV);
    @assert n == length(coeffV) == length(seV)  "Size mismatch"

    if interceptAsText
        idxV, constIdx = idx_drop_regressor(xStrV, interceptName);
    else
        idxV = 1 : length(xStrV);
        constIdx = nothing;
    end
    if !isnothing(seV)
        seV = seV[idxV];
    end

    _, ax = bar_graph(xStrV[idxV], coeffV[idxV]; fig = fig, pos = pos,
        ylabel = string(yStr), yerror = seV, kwargs...);

    if interceptAsText  &&  !isnothing(constIdx)
        show_intercept_text(ax, coeffV, seV, constIdx);
    end
    return fig, ax
end

# Drop a regressor by name. Return remaining indices and matching index.
function idx_drop_regressor(xStrV, regrName)
    idxV = findall(x -> !isequal(x, regrName), xStrV);
    matchIdx = findfirst(x -> isequal(x, regrName), xStrV);
    return idxV, matchIdx
end

function show_intercept_text(ax, coeffV :: AbstractVector{F1}, seV, constIdx) where F1
    xPos = 1.0;
    idxV = [j  for j = 1 : length(coeffV)  if j != constIdx];
    yPos = maximum(coeffV[idxV]) + 0.1;
    intercept = round(coeffV[constIdx], digits = 2);
    se = round(seV[constIdx], digits = 2);
    interStr = "Intercept: $intercept ($se)";
    text!(ax, interStr; position = (xPos, yPos));
end

# Plot regressions in one plot. Each group of bars is a regressor.
function plot_regressions_together(coeffNameV, regrLabelV, coeffM, seM;
    fig = blank_plot(), pos = (1,1),
    interceptAsText :: Bool = false,
    interceptName = :cons,
    figTitle = nothing, kwargs...)

    if interceptAsText
        idxV, constIdx = idx_drop_regressor(coeffNameV, interceptName);
    else
        idxV = 1 : length(coeffNameV);
        constIdx = nothing;
    end
    if !isnothing(seM)
        seM = seM[idxV, :];
    end

    _, ax = grouped_bar_graph(string.(coeffNameV[idxV]), coeffM[idxV, :];
        fig = fig, pos = pos, yerror = seM, kwargs...);
    isnothing(figTitle)  ||  (ax.title = figTitle);
    if interceptAsText  &&  !isnothing(constIdx)
        show_intercept_text(ax, coeffM, seM, constIdx);
    end
    return fig, ax
end

# Rows are coefficients.
function show_intercept_text(ax, coeffM :: AbstractMatrix{F1}, seM, constIdx) where F1
    nCoeff, nRegr = size(coeffM);
    if isnothing(seM)
        seV = nothing;
    else
        seV = seM[constIdx, :];
    end
    idxV = [j  for j = 1 : nCoeff  if j != constIdx];
    yPos = maximum(coeffM[idxV, :]) + 0.1;
    text!(ax, make_intercept_text(coeffM[constIdx,:], seV); position = (1.0, yPos));
end

function make_intercept_text(interceptV, seV)
    nRegr = length(interceptV);
    interStr = "Intercepts: ";
    for j = 1 : nRegr
        sepStr = (j == nRegr)  ?  ""  :  ", ";
        if isnothing(seV)
            seStr = "";
        else
            se = round(seV[j], digits = 2);
            seStr = " ($se)";
        end
        inter = round(interceptV[j], digits = 2);
        interStr = interStr * " $inter" * seStr * sepStr;
    end
    return interStr
end

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


text:
"""A specification for a particular instance of the environment. Used
to register the parameters for official evaluations.

Args:
    id (str): The official environment ID
    entry_point (Optional[str]): The Python entrypoint of the environment class (e.g. module.name:Class)
    trials (int): The number of trials to average reward over
    reward_threshold (Optional[int]): The reward threshold before the task is considered solved
    kwargs (dict): The kwargs to pass to the environment class
    nondeterministic (bool): Whether this environment is non-deterministic even after seeding
    tags (dict[str:any]): A set of arbitrary key-value tags on this environment, including simple property=True tags

Attributes:
    id (str): The official environment ID
    trials (int): The number of trials run in official evaluation
"""

struct EnvSpec
    id::Symbol
    entry_point::String

    trials::Int
    reward_threshold::Union{Int, Nothing}
    nondeterministic::Bool
    tags::Dict{String, Any}
    max_episode_steps::Union{Int, Nothing}
    max_episode_seconds::Union{Int, Nothing}
    kwargs::Dict
end

function EnvSpec(id, entry_point; trials=100, reward_threshold=nothing,
                nondeterministic=false, tags=Dict(), max_episode_steps=nothing,
                max_episode_seconds=nothing, kwargs=Dict())

    # Still need to figure out how to create a wrapper that stops the environment
    # after a certain number of episodes.
    tags["wrapper_config.TimeLimit.max_episode_steps"] = max_episode_steps
    EnvSpec(id, entry_point, trials, reward_threshold, nondeterministic, tags,
            max_episode_steps, max_episode_seconds, kwargs)
end

function _make(spec::EnvSpec, render_mode::Symbol; kwargs...)
    _kwargs = deepcopy(spec.kwargs)
    merge!(_kwargs, Dict(kwargs))

    env_var = load(spec.entry_point, spec.id)
    ctx_var = load(spec.entry_point, :Ctx)
    env = Base.invokelatest(env_var)
    ctx = Base.invokelatest(ctx_var, env, render_mode)
    env, ctx, spec.reward_threshold, spec.max_episode_steps
end

function load(path, id)
    #mod_name, attr_name = split(name, ':')
    #println(mod_name, attr_name)
    home_dir = @__DIR__
    fullpath = home_dir * path
    include(fullpath)
    eval(id)
end

mutable struct EnvRegistry
    env_specs::Dict{String, EnvSpec}

    EnvRegistry() = new(Dict{String, EnvSpec}())
end

function _register(reg::EnvRegistry, id_string, id, entry_point; kwargs...)
    !isnothing(get(reg.env_specs, id_string, nothing)) &&
                (throw(ErrorException("Cannot re-register id: $(id_string)")))

    reg.env_specs[id_string] = EnvSpec(id, entry_point; kwargs...)
end

function _make(reg::EnvRegistry, id_string, render_mode::Symbol; kwargs...)
    spec = get(reg.env_specs, id_string, nothing)

    isnothing(spec) &&
            (throw(ErrorException("Environment $(id_string) not found in the registry. Please ensure that you've spelled the name correctly.")))

    _make(spec, render_mode; kwargs...)
end

registry = EnvRegistry()

"""
Registers an environment to the global registry.
"""
register(id_string, id, entry_point; kwargs...) = _register(registry, id_string,
                                                            id, entry_point; kwargs...)

"""
Instantiates an instance of the environment with appropriate kwargs.

Optional keyword arguments:
    trials (Int): The number of trials to average reward over
    reward_threshold (Int): The reward threshold before the task is considered solved
    kwargs (Dict): The kwargs to pass to the environment class
    nondeterministic ( Bool): Whether this environment is non-deterministic even after seeding
    tags (Dict{String, Any}): A set of arbitrary key-value tags on this environment, including simple property=True tags
"""

include("env_wrapper.jl")

function make(id_string, mode=:human_pane, train=true; kwargs...)
	env, ctx, rt, max_ep_steps = _make(registry, id_string, mode; kwargs...)
	EnvWrapper(env, ctx, train; reward_threshold=rt, max_episode_steps=max_ep_steps)
end