text:
@testset "GoldState" begin
  tree = build_gold_tree()
  config = build_configuration()
  system = ArcEagerSystem()
  state = GoldState(tree, config, system)

  head_in_stack = DependencyParser.DependencyParsing.ArcEager.HEAD_IN_STACK
  head_in_buffer = DependencyParser.DependencyParsing.ArcEager.HEAD_IN_BUFFER

  correct_root_dependents_in_stack = 1
  correct_root_dependents_in_buffer = 0
  correct_dependents_in_stack = [0, 1, 0, 0, 0, 0]
  correct_dependents_in_buffer = [0 ,2, 0, 0, 1, 0]
  correct_heads_states = [
    head_in_stack,
    head_in_stack,
    head_in_stack,
    head_in_buffer,
    head_in_stack,
    head_in_stack
  ]

  # Expect it correct calculate dependents arcs count
  @test state.root_dependents_in_buffer_count == correct_root_dependents_in_buffer
  @test state.root_dependents_in_stack_count == correct_root_dependents_in_stack
  @test state.dependents_in_buffer_count == correct_dependents_in_buffer
  @test state.dependents_in_stack_count == correct_dependents_in_stack
  @test state.heads_states == correct_heads_states
end

max_stars_repo_name: LarsWl/DependencyParser

max_issues_repo_name: LarsWl/DependencyParser

max_forks_repo_name: LarsWl/DependencyParser


text:
using Trilinos
using Test
using MPI

if Test.get_testset_depth() == 0
  MPI.Init()
end

comm = Teuchos.MpiComm(MPI.CComm(MPI.COMM_WORLD))

const n = 20

# tri-diagonal nxn test matrix
A = spdiagm(-1 => 2*ones(n-1), 0 => ones(n), 1=> 3*ones(n-1))

rowmap = Tpetra.Map(n, 0, comm)
A_crs = Tpetra.CrsMatrix(A,rowmap)
Tpetra.describe(A_crs, Teuchos.VERB_LOW)

@test sqrt(sum(nonzeros(A).^2)) ≈ Tpetra.getFrobeniusNorm(A_crs)

function loop_gid(rowmap)
  n_my_elms = Tpetra.getNodeNumElements(rowmap)
  i_sum  = 0
  for i in 1:n_my_elms
    i_sum += Tpetra.getGlobalElement(rowmap,i)
  end
  return i_sum
end

@show methods(Tpetra.getGlobalElement)

@time loop_gid(rowmap[])
@time loop_gid(rowmap[])
@time loop_gid(rowmap[])

if Test.get_testset_depth() == 0
  MPI.Finalize()
end

ext: jl


text:
import Base: start, next, done

immutable Edge
    interior::Bool
    inner::UnitRange{Int}
    outer::UnitRange{Int}
end

start(e::Edge) = (k = first(e.outer); in(k, e.inner) ? last(e.inner)+1 : k)
next(e::Edge, i) = (i, e.interior && in(i+1, e.inner) ? last(e.inner)+1 : i+1)
done(e::Edge, i) = i == last(e.outer)+1

#=
@nedge prototype usage:
for N = 1:3
    @eval begin
        function foo{T}(A::AbstractArray{T,$N}, region)
            interior = [in(d, region) ? 2:size(A,d)-1 : 1:size(A,d) for d = 1:$N]
            @nexprs $N d->(interior_d = interior[d])
            @nloops $N i d->interior[d] begin
                # do something, not worrying about the boundaries
            end
            @nedge $N i d->1:size(A,d) d->interior_d begin
                # do the same thing, this time worrying about the boundaries
            end
            retval
        end
    end
end

=#

macro nedge(N, itersym, outerexpr, innerexpr, args...)
    _nedge(N, itersym, outerexpr, innerexpr, args...)
end

_nedge(N::Int, itersym::Symbol, outersym::Symbol, innerexpr::Expr, args::Expr...) = _nedge(N, itersym, :(d->1:size($outersym,d)), innerexpr, args...)

function _nedge(N::Int, itersym::Symbol, outerexpr::Expr, innerexpr::Expr, args::Expr...)
    outerexpr.head == :-> || error("outerexpr must be an anonymous-function expression")
    innerexpr.head == :-> || error("innerexpr must be an anonymous-function expression")
    1 <= length(args) <= 3 || error("Wrong number of arguments")
    body = args[end]
    ex = Expr(:escape, body)
    firstsym = gensym()           # Boolean scalar; initialization variable
    fastest_edge_sym = gensym()   # Boolean ntuple; true only for the fastest interior dimension
    interior_sym = gensym()       # Boolean ntuple; true if all slower indexes are in the interior
    fastestexpr = Array(Expr, N)  # used during initialization
    for dim = 1:N
        itervar = Cartesian.inlineanonymous(itersym, dim)
        outer = Cartesian.inlineanonymous(outerexpr, dim)
        inner = Cartesian.inlineanonymous(innerexpr, dim)
        interior_cur  = Cartesian.inlineanonymous(interior_sym, dim-1)
        interior_prev = Cartesian.inlineanonymous(interior_sym, dim)
        fastvar = Cartesian.inlineanonymous(fastest_edge_sym, dim)
        # Set up initialization of fastest_edge_sym
        fastestexpr[dim] = Expr(:escape, :($fastvar = $firstsym & ($outer != $inner); $firstsym &= !$fastvar))
        # Set up updating of interior_sym
        interiorexpr = dim > 1 ? (:($interior_cur = $interior_prev & in($itervar, $inner))) : (:(nothing))
        preexpr = length(args) > 1 ? Cartesian.inlineanonymous(args[1], dim) : (:(nothing))
        postexpr = length(args) > 2 ? Cartesian.inlineanonymous(args[2], dim) : (:(nothing))
        ex = quote
            for $(esc(itervar)) in KernelTools.Edge($(esc(fastvar)) & $(esc(interior_prev)), $(esc(inner)), $(esc(outer)))
                $(esc(interiorexpr))
                $(esc(preexpr))
                $ex
                $(esc(postexpr))
            end
        end
    end
    lastinterior = Cartesian.inlineanonymous(interior_sym, N)
    # Initialize fastest_edge_sym and set the last interior_sym flag to be true before we do the loops
    quote
        $(esc(firstsym)) = true
        $(fastestexpr...)
        $(esc(lastinterior)) = true
        $ex
    end
end


text:
module ComplexBigMatrices

"""
This module calculates exp(), log() and eigen() of a complex BigFloat Matrix.
"""

export exp, log, eigen
import .Base: exp, log

using LinearAlgebra
using GenericLinearAlgebra
using GenericSchur

"""
    exp(a::Matrix{Complex{BigFloat}})
Extends exp() to work with Matrix{Complex{BigFloat}}
"""
function exp(a::Matrix{Complex{BigFloat}})
    eig = eigen(a)
    eig.vectors*Diagonal(exp.(eig.values))*eig.vectors^(-1)
end

"""
    log(a::Matrix{Complex{BigFloat}})
Extends log() to work with Matrix{Complex{BigFloat}}
"""
function log(a::Matrix{Complex{BigFloat}})
    eig = eigen(a)
    eig.vectors*Diagonal(log.(eig.values))*eig.vectors^(-1)
end

"""
    eigen(a::Matrix{Complex{BigFloat}})
Extends eigen() to work with Matrix{Complex{BigFloat}}
"""
function eigen(a::Matrix{Complex{BigFloat}})
    LinearAlgebra.eigen(a)
end

end


text:
# Files sourced from: https://www.wavsource.com
sounds_folder = joinpath(dirname(pathof(Beep)), "sounds")
readdir(sounds_folder)

beep_sounds = Dict{Symbol, Tuple{Array{Float64, 2}, Float64}}()

for fname in readdir(sounds_folder)

    wav_file = wavread(joinpath(sounds_folder, fname))

    sound_symbol = Symbol(replace(fname, ".wav" => ""))
    beep_sounds[sound_symbol] = (wav_file[1], wav_file[2])
end


text:
# ============================================================
# Solution Algorithm for Particle Simulations
# ============================================================

"""
    sample_particle!(ptc::Particle1D, m, x, v, e, idx, flag, tc)
    sample_particle!(ptc::Particle1D, m, prim::T, x, dx, idx, μᵣ, ω, flag = 0) where {T<:AbstractArray{<:Real,1}}
    sample_particle!(ptc::Particle1D, m, prim::T, umin, umax, x, dx, idx, μᵣ, ω, flag = 0) where {T<:AbstractArray{<:Real,1}}
    sample_particle!(ptc::Particle1D, KS::SolverSet, ctr, idx)
    sample_particle!(ptc::Particle, ip, p)
    sample_particle!(ptc::Particle, ip, KS::SolverSet, ctr, idx)

Sample particles from local flow conditions

"""
function sample_particle!(ptc::Particle1D, m, x, v, e, idx, flag, tc)
    ptc.m = m
    ptc.x = x
    ptc.v .= v
    ptc.e = e
    ptc.idx = idx
    ptc.flag = flag
    ptc.tc = tc

    return nothing
end

function sample_particle!(
    ptc::Particle1D,
    m,
    x,
    dx,
    prim::T,
    idx,
    μᵣ,
    ω,
    flag = 0,
) where {T<:AbstractArray{<:Real,1}}
    ptc.m = m
    ptc.x = x + (rand() - 0.5) * dx
    ptc.v .= sample_maxwell(prim)
    ptc.e = 0.5 / prim[end]
    ptc.idx = idx
    ptc.flag = flag
    τ = vhs_collision_time(prim, μᵣ, ω)
    ptc.tc = next_collision_time(τ)

    return nothing
end

function sample_particle!(
    ptc::Particle1D,
    m,
    x,
    dx,
    prim::T,
    umin,
    umax,
    idx,
    μᵣ,
    ω,
    flag = 0,
) where {T<:AbstractArray{<:Real,1}}
    ptc.m = m
    ptc.x = x + (rand() - 0.5) * dx
    ptc.v .= sample_maxwell(prim, umin, umax)
    ptc.e = 0.5 / prim[end]
    ptc.idx = idx
    ptc.flag = flag
    τ = vhs_collision_time(prim, μᵣ, ω)
    ptc.tc = next_collision_time(τ)

    return nothing
end

function sample_particle!(ptc::Particle1D, KS::SolverSet, ctr, idx, x, Δx)
    ptc.m = KS.gas.m
    ptc.x = x + (rand() - 0.5) * Δx
    ptc.v .= sample_maxwell(ctr.prim)
    ptc.e = 0.5 / ctr.prim[end]
    ptc.idx = idx
    ptc.flag = 0
    τ = vhs_collision_time(ctr.prim, KS.gas.μᵣ, KS.gas.ω)
    ptc.tc = next_collision_time(τ)

    return nothing
end

function sample_particle!(ptc::Particle, ip, p)
    m, x, v, e, idx, flag, tc = p

    ptc.m[ip] = m
    ptc.x[ip] = x
    ptc.v[ip] .= v
    ptc.e[ip] = e
    ptc.idx[ip] = idx
    ptc.flag[ip] = flag
    ptc.tc[ip] = tc

    return nothing
end

function sample_particle!(ptc::Particle, ip, KS::SolverSet, ctr, idx)
    ptc.m[ip] = KS.gas.m
    ptc.x[ip] = ctr.x + (rand() - 0.5) * ctr.dx
    ptc.v[ip] .= sample_maxwell(ctr.prim)
    ptc.e[ip] = 0.5 / ctr.prim[end]
    ptc.idx[ip] = idx
    ptc.flag[ip] = 0
    τ = vhs_collision_time(ctr.prim, KS.gas.μᵣ, KS.gas.ω)
    ptc.tc[ip] = next_collision_time(τ)

    return nothing
end

"""
    free_transport!(KS::SolverSet, x, v, flag, dt, np = length(x))

Flight particles along trajectories

"""
function free_transport!(KS::SolverSet, x, v, flag, dt, np = length(x))

    @inbounds for i = 1:np
        x[i] += v[i, 1] * dt

        flag[i] = 0
        if x[i] < KS.pSpace.x0
            flag[i] = 1
        elseif x[i] > KS.pSpace.x1
            flag[i] = 2
        end

        #if flag[i] != 0
        #    x0 = x[i] - v[i, 1]*dt
        #    vi = @view v[i, :]
        #    x[i] = boundary!(x[i], vi, (KS.ib, [KS.pSpace.x0, KS.pSpace.x1], x0, flag[i], dt))
        #end
    end

    return nothing

end

function bgk_transport!(KS::SolverSet, ctr, ptc, ptc_new, dt, np = KS.gas.np)

    @inbounds for i = 1:KS.pSpace.nx
        ctr[i].wf .= ctr[i].w
    end

    npt = 0
    @inbounds for i = 1:np
        cellid = ptc.idx[i]
        ptc.tc[i] = -ctr[cellid].τ * log(rand())

        if ptc.tc[i] > dt # class 1: free transport particles

            ptc.x[i] += dt * ptc.v[i, 1]

            ctr[cellid].wf[1] -= ptc.m[i] / ctr[cellid].dx
            ctr[cellid].wf[2] -= ptc.m[i] * ptc.v[i, 1] / ctr[cellid].dx
            ctr[cellid].wf[3] -= 0.5 * ptc.m[i] * sum(ptc.v[i, :] .^ 2) / ctr[cellid].dx

            if ptc.x[i] >= KS.pSpace.x0 && ptc.x[i] <= KS.pSpace.x1
                npt += 1
                ptc.idx[i] = find_idx(KS.pSpace.x[1:end], ptc.x[i], mode = :uniform)

                ptc_new.m[npt] = ptc.m[i]
                ptc_new.x[npt] = ptc.x[i]
                ptc_new.v[npt, :] .= ptc.v[i, :]
                ptc_new.e[npt] = ptc.e[i]
                ptc_new.idx[npt] = ptc.idx[i]
                ptc_new.flag[npt] = 0
                ptc_new.tc[npt] = ptc.tc[i]
            end

        elseif ptc.tc[i] >
               boundary_time(ptc.x[i], ptc.v[i, 1], ctr[cellid].x, ctr[cellid].dx)   # class 2: jumping collision particles

            ptc.x[i] += ptc.tc[i] * ptc.v[i, 1]
            #ptc.x[i] += dt * ptc.v[i, 1]

            if ptc.x[i] > ctr[cellid].x + 0.5 * ctr[cellid].dx &&
               ptc.x[i] < ctr[cellid].x - 0.5 * ctr[cellid].dx

                ctr[cellid].wf[1] -= ptc.m[i] / ctr[cellid].dx
                ctr[cellid].wf[2] -= ptc.m[i] * ptc.v[i, 1] / ctr[cellid].dx
                ctr[cellid].wf[3] -= 0.5 * ptc.m[i] * sum(ptc.v[i, :] .^ 2) / ctr[cellid].dx

                if ptc[i].x >= KS.pSpace.x0 && ptc[i].x <= KS.pSpace.x1
                    cellid_tmp = find_idx(KS.pSpace.x[1:end], ptc.x[i], mode = :uniform)

                    ctr[cellid_tmp].wf[1] += ptc.m[i] / ctr[cellid_tmp].dx
                    ctr[cellid_tmp].wf[2] += ptc.m[i] * ptc.v[i, 1] / ctr[cellid_tmp].dx
                    ctr[cellid_tmp].wf[3] +=
                        0.5 * ptc.m[i] * sum(ptc.v[i] .^ 2) / ctr[cellid_tmp].dx
                end
            end

        end

    end

    @inbounds for i = 1:KS.pSpace.nx
        ctr[i].wp .= 0.0
    end
    @inbounds for i = 1:npt
        cellid = ptc_new.idx[i]

        ctr[cellid].wp[1] += ptc_new.m[i] / ctr[cellid].dx
        ctr[cellid].wp[2] += ptc_new.m[i] * ptc_new.v[i, 1] / ctr[cellid].dx
        ctr[cellid].wp[3] += 0.5 * ptc_new.m[i] * sum(ptc_new.v[i, :] .^ 2) / ctr[cellid].dx
    end

    KS.gas.np = npt
    return nothing

end

function boundary!(KS, ctr, ptc, face, dt, bc = :maxwell::Symbol)

    np = KS.gas.np

    if bc == :fix

        ng = 1 - firstindex(KS.pSpace.x)

        # left boundary
        @inbounds for i = 1:ng
            idx = 1 - i
            npc = ceil(ctr[idx].w[1] * ctr[idx].dx / KS.gas.m) |> Int
            for j = 1:npc
                np += 1

                ptc.m[np] = KS.gas.m
                ptc.x[np] = ctr[idx].x + (rand() - 0.5) * ctr[idx].dx
                ptc.v[np, :] .= sample_maxwell(ctr[idx].prim)
                ptc.e[np] = 0.5 / ctr[idx].prim[end]
                ptc.idx[np] = idx
                ptc.flag[np] = 0
                τ = vhs_collision_time(ctr[idx].prim, KS.gas.μᵣ, KS.gas.ω)
                ptc.tc[np] = next_collision_time(τ)

                #sample_particle!(ptc, np, KS, ctr[idx], idx)
            end
        end

        # right boundary
        @inbounds for i = 1:ng
            idx = KS.pSpace.nx + i
            npc = ceil(ctr[idx].w[1] * ctr[idx].dx / KS.gas.m) |> Int
            for j = 1:npc
                np += 1

                ptc.m[np] = KS.gas.m
                ptc.x[np] = ctr[idx].x + (rand() - 0.5) * ctr[idx].dx
                ptc.v[np, :] .= sample_maxwell(ctr[idx].prim)
                ptc.e[np] = 0.5 / ctr[idx].prim[end]
                ptc.idx[np] = idx
                ptc.flag[np] = 0
                τ = vhs_collision_time(ctr[idx].prim, KS.gas.μᵣ, KS.gas.ω)
                ptc.tc[np] = next_collision_time(τ)

                #sample_particle!(ptc, np, KS, ctr[idx], idx)
            end
        end

        KS.gas.np = np

    elseif bc == :maxwell

        @inbounds for i = 1:np
            if ptc.flag[i] != 0
                x0 = ptc.x[i] - ptc.v[i, 1] * dt
                vi = @view ptc.v[i, :]
                ptc.x[i] = maxwell_boundary!(
                    ptc.x[i],
                    vi,
                    (KS.ib, [KS.pSpace.x0, KS.pSpace.x1], x0, ptc.flag[i], dt),
                )
            end
        end

    end

    return nothing

end

"""
    particle_boundary_maxwell!(x, v, p)

Maxwell diffusive boundary for particle reflection

"""
function maxwell_boundary!(x, v, p)

    ib, xwall, x0, flag, dt = p

    if flag == 1
        xw = xwall[1]
        primB = ib.bc(xw)
        bound = [0.0, Inf]
        vw = [primB[2]; zeros(2)]#ib.vL
    elseif flag == 2
        xw = xwall[2]
        primB = ib.bc(xw)
        bound = [-Inf, 0.0]
        vw = [primB[2]; zeros(2)]#ib.vR
    end

    #v[1] = sqrt(-log(1.0-rand()))
    v[1] = sample_maxwell(primB[end], bound[1], bound[2])
    v[2] = sample_maxwell(primB[end], vw[2])
    v[3] = sample_maxwell(primB[end], vw[3])

    dtr = dt * (x - xw) / (x - x0)
    x = xw + v[1] * dtr

    return x

end

function bgk_collision!(
    KS::SolverSet,
    ctr::AbstractArray{ControlVolumeParticle1D,1},
    ptc::Particle,
    face::AbstractArray{Interface1D,1},
    res::AbstractArray{<:AbstractFloat,1},
)

    sum_res = zeros(3)
    sum_avg = zeros(3)

    npt = KS.gas.np
    @inbounds for i = 1:KS.pSpace.nx
        # fluid variables
        wold = deepcopy(ctr[i].w)
        @. ctr[i].wf += (face[i].fw - face[i+1].fw) / ctr[i].dx

        if ctr[i].wf[1] < 0.0
            ctr[i].wf .= zero(ctr[i].wf)
        else
            primG = conserve_prim(ctr[i].wf, KS.gas.γ)
            if primG[end] < 0.0
                primG[end] = 1e8
                ctr[i].wf .= prim_conserve(primG, KS.gas.γ)
            end
        end

        @. ctr[i].w = ctr[i].wf + ctr[i].wp
        if ctr[i].w[1] < 0.0
            ctr[i].w .= zero(ctr[i].w) .+ 1e-8
        else
            ctr[i].prim = conserve_prim(ctr[i].w, KS.gas.γ)
            if ctr[i].prim[end] < 0.0
                ctr[i].prim[end] = eltype(ctr[i].prim)(1e8)
                ctr[i].w .= prim_conserve(ctr[i].prim, KS.gas.γ)
            end
            ctr[i].τ = vhs_collision_time(ctr[i].prim, KS.gas.μᵣ, KS.gas.ω)
        end
        @. sum_res += (ctr[i].w - wold)^2
        @. sum_avg += abs.(ctr[i].w)

        # particles
        no_cellid = Int(round(ctr[i].wf[1] * ctr[i].dx / KS.gas.m))
        for j = 1:no_cellid
            npt += 1

            ptc.m[npt] = KS.gas.m
            ptc.x[npt] = ctr[i].x + (rand() - 0.5) * ctr[i].dx
            ptc.v[npt, :] .= sample_maxwell(ctr[i].prim)
            ptc.e[npt] = 0.5 / ctr[i].prim[end]
            ptc.idx[npt] = i
            ptc.flag[npt] = 0
            τ = vhs_collision_time(ctr[i].prim, KS.gas.μᵣ, KS.gas.ω)
            ptc.tc[npt] = next_collision_time(τ)
        end
    end

    @. res = sqrt(KS.pSpace.nx * sum_res) / (sum_avg + 1e-8)
    KS.gas.np = npt

    return nothing

end

function sort!(
    KS,
    ctr::T,
    x,
    idx,
    ref,
    np = length(idx);
    mode = :uniform,
) where {T<:AbstractArray{<:AbstractControlVolume1D,1}}

    # calculate cell indices of particles
    @inbounds for i = 1:np
        idx[i] = find_idx(KS.pSpace.x[1:end], x[i], mode = mode)
    end

    # count the number of particles in each cell
    @inbounds for i in eachindex(ctr)
        ctr[i].np = 0
    end
    for i = 1:np
        ctr[idx[i]].np += 1
    end

    # build index list as cumulative sum of the number of particles in each cell
    m = 1
    @inbounds for jcell in eachindex(ctr)
        ctr[jcell].ip = m
        m += ctr[jcell].np
    end

    # build cross-reference list
    temp = zeros(Int, axes(ctr))
    @inbounds for i = 1:np
        jcell = idx[i]
        k = ctr[jcell].ip + temp[jcell]
        ref[k] = i
        temp[jcell] += 1
    end

    return nothing

end

function dsmc!(KS, ctr, ref, v, dt, ne = 1)

    vcm = zeros(3)
    vrel = zeros(3)
    col = 0

    for jcell in eachindex(ctr)
        number = ctr[jcell].np
        if number > 1

            #select = coeff*number**2*vrmax(jcell) + remainder(jcell)
            select =
                ne * 0.5 / √2 * dt / ctr[jcell].dx / KS.gas.Kn *
                number^2 *
                ctr[jcell].vrmax + ctr[jcell].remainder

            nsel = Int(floor(select))
            ctr[jcell].remainder = select - nsel
            crm = ctr[jcell].vrmax

            for isel = 1:nsel

                # pick two particles at random out of this cell
                k = floor(rand() * number) |> Int
                kk = mod(Int(floor(k + rand() * (number - 1)) + 1), number)
                ip1 = ref[k+ctr[jcell].ip]  # first particle
                ip2 = ref[kk+ctr[jcell].ip] # second particle

                # calculate pair's relative speed
                cr = sqrt(
                    (v[ip1, 1] - v[ip2, 1])^2 +
                    (v[ip1, 2] - v[ip2, 2])^2 +
                    (v[ip1, 3] - v[ip2, 3])^2,
                )
                if cr > crm    # If relative speed larger than crm,
                    crm = cr                # then reset crm to larger value
                end

                # accept or reject candidate pair according to relative speed
                if cr / ctr[jcell].vrmax > rand()
                    # If pair accepted, select post-collision velocities
                    col += 1                     # Collision counter
                    for k = 1:3
                        vcm[k] = 0.5 * (v[ip1, k] + v[ip2, k])       # Center of mass velocity
                    end
                    cos_th = 1.0 - 2.0 * rand()       # Cosine and sine of
                    sin_th = sqrt(1.0 - cos_th^2)      # collision angle theta
                    phi = 2.0 * pi * rand()             # Collision angle phi
                    vrel[1] = cr * cos_th                 # Compute post-collision
                    vrel[2] = cr * sin_th * cos(phi)        # relative velocity
                    vrel[3] = cr * sin_th * sin(phi)
                    for k = 1:3
                        v[ip1, k] = vcm[k] + 0.5 * vrel[k]   # Update post-collision
                        v[ip2, k] = vcm[k] - 0.5 * vrel[k]   # velocities
                    end

                end
                ctr[jcell].vrmax = crm     # Update max relative speed

            end
        end
    end

    return col

end

function stat!(KS, ctr, ptc)
    @inbounds Threads.@threads for i = 1:KS.pSpace.nx
        ctr[i].w .= 0.0
    end

    @inbounds Threads.@threads for i in eachindex(ptc.x)
        ctr[ptc.idx[i]].w[1] += ptc.m[i] / ctr[ptc.idx[i]].dx
        ctr[ptc.idx[i]].w[2] += ptc.m[i] * ptc.v[i, 1] / ctr[ptc.idx[i]].dx
        ctr[ptc.idx[i]].w[3] += ptc.m[i] * ptc.v[i, 2] / ctr[ptc.idx[i]].dx
        ctr[ptc.idx[i]].w[4] += 0.5 * ptc.m[i] * sum(ptc.v[i, :] .^ 2) / ctr[ptc.idx[i]].dx
    end

    @inbounds Threads.@threads for i = 1:KS.pSpace.nx
        ctr[i].prim .= KitBase.conserve_prim(ctr[i].w, KS.gas.γ)
    end
end

function update!(
    KS::AbstractSolverSet,
    ctr::AbstractArray{ControlVolumeParticle1D,1},
    ptc::AbstractArray{Particle1D,1},
    ptc_temp::AbstractArray{Particle1D,1},
    face::AbstractArray{Interface1D,1},
    dt,
    res;
    coll = :bgk::Symbol,
    bc = :fix::Symbol,
)

    res .= 0.0

    particle_transport!(KS, ctr, ptc, ptc_temp, dt)
    particle_collision!(KS, ctr, ptc_temp, face, res, coll)
    particle_boundary!(KS, ctr, ptc, ptc_temp, face, dt, coll, bc)
    duplicate!(ptc, ptc_temp, KS.gas.np)

    return nothing

end

"""
Update algorithm for particle transports

    update_transport!(KS::AbstractSolverSet, ctr::ControlVolumeParticle1D, ptc::Particle1D, ptc_temp::Particle1D, dt)

"""
function particle_transport!(
    KS::T1,
    ctr::T2,
    ptc::T3,
    ptc_tmp::T3,
    dt,
) where {
    T1<:AbstractSolverSet,
    T2<:AbstractArray{ControlVolumeParticle1D,1},
    T3<:AbstractArray{Particle1D,1},
}

    @inbounds Threads.@threads for i = 1:KS.pSpace.nx
        ctr[i].wf .= ctr[i].w
    end

    npt = 0
    @inbounds for i = 1:KS.gas.np
        cellid = ptc[i].idx
        ptc[i].tc = -vhs_collision_time(ctr[cellid].prim, KS.gas.μᵣ, KS.gas.ω) * log(rand())

        if ptc[i].tc > dt # class 1: free transport particles
            ptc[i].x += dt * ptc[i].v[1]

            ctr[cellid].wf[1] -= ptc[i].m / ctr[cellid].dx
            ctr[cellid].wf[2] -= ptc[i].m * ptc[i].v[1] / ctr[cellid].dx
            ctr[cellid].wf[3] -= 0.5 * ptc[i].m * sum(ptc[i].v .^ 2) / ctr[cellid].dx

            if ptc[i].x >= KS.pSpace.x0 && ptc[i].x <= KS.pSpace.x1
                npt += 1

                cellid_tmp = Int(ceil((ptc[i].x - KS.pSpace.x0) / ctr[1].dx))
                #cellid_tmp = find_idx(KS.pSpace.x[1:KS.pSpace.nx], ptc[i].x, mode=:uniform)

                sample_particle!(
                    ptc_tmp[npt],
                    ptc[i].m,
                    ptc[i].x,
                    ptc[i].v,
                    0.5 / ctr[cellid_tmp].prim[end],
                    cellid_tmp,
                    ptc[i].tc,
                )
            end
        elseif ptc[i].tc > boundary_time(ptc[i].x, ptc[i].v, ctr[cellid].x, ctr[cellid].dx)   # class 2: jumping collision particles
            #ptc[i].x += ptc[i].tc * ptc[i].v[1]
            ptc[i].x += dt * ptc[i].v[1]

            if ptc[i].x > ctr[cellid].x + 0.5 * ctr[cellid].dx &&
               ptc[i].x < ctr[cellid].x - 0.5 * ctr[cellid].dx

                ctr[cellid].wf[1] -= ptc[i].m / ctr[cellid].dx
                ctr[cellid].wf[2] -= ptc[i].m * ptc[i].v[1] / ctr[cellid].dx
                ctr[cellid].wf[3] -= 0.5 * ptc[i].m * sum(ptc[i].v .^ 2) / ctr[cellid].dx

                if ptc[i].x >= KS.pSpace.x0 && ptc[i].x <= KS.pSpace.x1
                    cellid_tmp = Int(ceil((ptc[i].x - KS.pSpace.x0) / ctr[1].dx))
                    #cellid_tmp = find_idx(KS.pSpace.x[1:KS.pSpace.nx], ptc[i].x, mode=:uniform)

                    ctr[cellid_tmp].wf[1] += ptc[i].m / ctr[cellid_tmp].dx
                    ctr[cellid_tmp].wf[2] += ptc[i].m * ptc[i].v[1] / ctr[cellid_tmp].dx
                    ctr[cellid_tmp].wf[3] +=
                        0.5 * ptc[i].m * sum(ptc[i].v .^ 2) / ctr[cellid_tmp].dx
                end
            end
        end
    end

    @inbounds Threads.@threads for i = 1:KS.pSpace.nx
        ctr[i].wp .= 0.0
    end
    @inbounds for i = 1:npt
        cellid = ptc_tmp[i].idx

        ctr[cellid].wp[1] += ptc_tmp[i].m / ctr[cellid].dx
        ctr[cellid].wp[2] += ptc_tmp[i].m * ptc_tmp[i].v[1] / ctr[cellid].dx
        ctr[cellid].wp[3] += 0.5 * ptc_tmp[i].m * sum(ptc_tmp[i].v .^ 2) / ctr[cellid].dx
    end

    KS.gas.np = npt
    return nothing

end

"""
Update algorithm for particle collisions

    update_collision!(KS::AbstractSolverSet, ctr::ControlVolumeParticle1D, ptc::Particle1D, ptc_temp::Particle1D, dt)

"""
function particle_collision!(
    KS::T,
    ctr::AbstractArray{ControlVolumeParticle1D,1},
    ptc_temp::AbstractArray{Particle1D,1},
    face::AbstractArray{Interface1D,1},
    res::AbstractArray{<:AbstractFloat,1},
    coll = :bgk::Symbol,
) where {T<:AbstractSolverSet}

    sum_res = zeros(3)
    sum_avg = zeros(3)

    npt = KS.gas.np
    @inbounds for i = 1:KS.pSpace.nx
        # fluid variables
        wold = deepcopy(ctr[i].w)
        @. ctr[i].wf += (face[i].fw - face[i+1].fw) / ctr[i].dx

        if ctr[i].wf[1] < 0.0
            ctr[i].wf .= zero(ctr[i].wf)
            #end
        else
            primG = conserve_prim(ctr[i].wf, KS.gas.γ)
            if primG[end] < 0.0
                primG[end] = 1e8
                ctr[i].wf .= prim_conserve(primG, KS.gas.γ)
            end
        end

        @. ctr[i].w = ctr[i].wf + ctr[i].wp
        if ctr[i].w[1] < 0.0
            ctr[i].w .= zero(ctr[i].w) .+ 1e-8
        else
            ctr[i].prim = conserve_prim(ctr[i].w, KS.gas.γ)
            if ctr[i].prim[end] < 0.0
                ctr[i].prim[end] = eltype(ctr[i].prim)(1e8)
                ctr[i].w .= prim_conserve(ctr[i].prim, KS.gas.γ)
            end
        end
        @. sum_res += (ctr[i].w - wold)^2
        @. sum_avg += abs.(ctr[i].w)

        ctr[i].τ = vhs_collision_time(ctr[i].prim, KS.gas.μᵣ, KS.gas.ω)

        # particles
        no_cellid = Int(round(ctr[i].wf[1] * ctr[i].dx / KS.gas.m))
        for j = 1:no_cellid
            npt += 1
            sample_particle!(
                ptc_temp[npt],
                KS.gas.m,
                primG,
                ctr[i].x,
                ctr[i].dx,
                i,
                KS.gas.μᵣ,
                KS.gas.ω,
            )
            #sample_particle!(ptc_temp[npt], KS.gas.m, primG, KS.vSpace.u0, KS.vSpace.u1, ctr[i].x, ctr[i].dx, i, KS.gas.μᵣ, KS.gas.ω)
        end
    end

    @. res = sqrt(KS.pSpace.nx * sum_res) / (sum_avg + 1e-8)
    KS.gas.np = npt
    return nothing

end

function particle_boundary!(
    KS::T,
    ctr::AbstractArray{ControlVolumeParticle1D,1},
    ptc::AbstractArray{Particle1D,1},
    ptc_temp::AbstractArray{Particle1D,1},
    face::AbstractArray{Interface1D,1},
    dt,
    coll = :bgk::Symbol,
    bc = :fix::Symbol,
) where {T<:AbstractSolverSet}

    np = KS.gas.np

    if bc == :fix
        # ghost cell
        ng = 1 - first(eachindex(KS.pSpace.x))

        # left boundary
        @inbounds for i = 1:ng
            idx = 1 - i
            npc = ceil(ctr[idx].w[1] * ctr[idx].dx / KS.gas.m) |> Int
            for j = 1:npc
                np += 1
                sample_particle!(ptc_temp[np], KS, ctr[idx], idx)
            end
        end

        # right boundary
        @inbounds for i = 1:ng
            idx = KS.pSpace.nx + i
            npc = ceil(ctr[idx].w[1] * ctr[idx].dx / KS.gas.m) |> Int
            for j = 1:npc
                np += 1
                sample_particle!(ptc_temp[np], KS, ctr[idx], idx)
            end
        end
    end

    KS.gas.np = np
    return nothing

end

"""
Duplicate particles: ptc_new -> ptc

"""
function duplicate!(ptc, ptc_new, n = length(ptc))
    @inbounds Threads.@threads for i = 1:n
        ptc[i].m = ptc_new[i].m
        ptc[i].x = ptc_new[i].x
        ptc[i].v .= ptc_new[i].v
        ptc[i].e = ptc_new[i].e
        ptc[i].idx = ptc_new[i].idx
        ptc[i].tc = ptc_new[i].tc
    end

    return nothing
end

function duplicate!(ptc::Particle, ptc_tmp, n = length(ptc.x))
    @inbounds Threads.@threads for i = 1:n
        ptc.m[i] = ptc_tmp.m[i]
        ptc.x[i] = ptc_tmp.x[i]
        ptc.v[i, :] .= ptc_tmp.v[i, :]
        ptc.e[i] = ptc_tmp.e[i]
        ptc.idx[i] = ptc_tmp.idx[i]
        ptc.tc[i] = ptc_tmp.tc[i]
    end

    return nothing
end


text:
# Alpao.jl -
#
# Julia interface to Alpao's library for managing their deformable mirrors.
#
#------------------------------------------------------------------------------
#
# This file is part of the `Alpao.jl` package which is licensed under the MIT
# "Expat" License.
#
# Copyright (C) 2016-2021, Éric Thiébaut & Jonathan Léger.
#

module Alpao

export
    send!,
    send,
    stop

import Base: reset
import Sockets: send

@static if isdefined(Base, :stop)
    import Base: stop
end

isfile(joinpath(@__DIR__,"..","deps","deps.jl")) ||
    error("Alpao not properly installed.  Please run Pkg.build(\"Alpao\")")
include("../deps/deps.jl")

# Notes:
#
# C-types are defined in "asdkType.h" and function prototypes in
# "asdkWrapper.h" of the Alpao SDK.
#
# All functions but 2 return status with type `COMPL_STAT` which is an
# enumeration (hence a `Cint`) with 2 possible values: 0 for success and -1
# for failure.  The other functions are `asdkInit` which returns a pointer
# to a structure (considered as anonymous, hence a `Ptr{Cvoid}` here) and
# `asdkPrintLastError` which returns nothing.
#
# The other C types used in the API are `Int` and `UInt` which are 32-bit
# integers, `CString` and `CStrConst` which are C-strings, `Scalar` which
# is C `double` and `Size_T` which is `size_t`.
#
# The following table summarizes the types:
#
# -------------------------------------------------
# Alpao API      C Type        Julia Type
# -------------------------------------------------
# Char           char          Cchar
# UChar          uint8_t       UInt8
# Short          int16_t       Int16
# UShort         uint16_t      UInt16
# Int            int32_t       Int32
# UInt           uint32_t      UInt32
# Long           int64_t       Int64
# ULong          uint64_t      UInt64
# Size_T         size_t        Csize_t
# Scalar         double        Cdouble
# asdkDM*        struct DM*    Ptr{Cvoid}
# CString        char*         Cstring / Ptr{UInt8}
# CStrConst      char const*   Cstring / Ptr{UInt8}
# -------------------------------------------------
#

const Keyword = AbstractString
const Scalar  = Cdouble
struct Status; val::Cint; end

const SUCCESS = Status(0)
const FAILURE = Status(-1)
const CMDMAX = Scalar(1.0)
const CMDMIN = -CMDMAX

"""
    dm = Alpao.DeformableMirror(name)

creates an instance `dm` to manage Alpao's deformable mirror identifed by the
`name` of its configuration file.

The deformable mirror instance `dm` can be used as follows:

```julia
length(dm)        # yields number of actuators
eltype(dm)        # yields bit type for an actuator command
dm[key]           # yields value of keyword `key`
dm[key] = val     # sets value of keyword `key`
keys(dm)          # yields a list of possible keywords
dm[]              # yields the last commands sent to the mirror
dm[:]             # yields a copy of the last commands
dm[i]             # yields the value of i-th actuator in the last commands
send(dm, cmd)     # sets the shape of the deformable mirror
send!(dm, cmd)    # idem but, on return, `cmd` contains actual commands
reset(dm)         # resets the deformable mirror values
stop(dm)          # stops asynchronous commands sent to the deformable mirror
close(dm)         # release the deformable mirror resources
minimum(dm)       # yields the minimum value of a command
maximum(dm)       # yields the maximum value of a command
extrema(dm)       # yields the minimum and maximum values of a command
VersionNumber(dm) # yields the SDK version number
```

List of parameter keywords

```
---------------------------------------------------------------------------
Keyword         Get  Set  Value  Description
---------------------------------------------------------------------------
"AckTimeout"     x    x    > 0   For Ethernet / USB interface only, set the
                                 time-out (ms); can be set in synchronous
                                 mode only (see SyncMode).
"DacReset"            x    1     Reset all digital to analog converters of
                                 drive electronics.
"ItfState"       x         0:1   Return 1 if PCI interface is busy or 0
                                 otherwise.
"LogDump"             x    1     Dump the log stack on the standard output.
"LogPrintLevel"  x    x    0:4   Changes the output level of the logger to
                                 the standard output.
"NbOfActuator"   x         ≥ 1   Get the numbers of actuator for that
                                 mirror.
"SyncMode"            x    0:1   0: Synchronous mode, will return when send
                                    is done.
                                 1: Asynchronous mode, return immediately
                                    after safety checks.
"TriggerMode"         x    0:1   Set mode of the (optional) electronics
                                 trigger output. 0: long pulse width or 1:
                                 short pulse width on each command.
"TriggerIn"           x    0:2   Set mode of the (optional) input trigger.
                                 0: disabled, 1: trig on rising edge or
                                 2: trig on falling edge.
"UseException"   x    x    0:1   Enables or disables the throwing of an
                                 exception on error.
"VersionInfo"    x         > 0   Alpao SDK core version, e.g. 3040500612 is
                                 SDK v3.04.05.0612 where 0612 is build
                                 number.
---------------------------------------------------------------------------
```

"""
mutable struct DeformableMirror
    ptr::Ptr{DeformableMirror} # handle to device
    num::Int                   # number of actuators
    cmd::Vector{Scalar}        # last command
    function DeformableMirror(ident::AbstractString)
        local ptr::Ptr{Cvoid}, num::Int
        if '/' in ident
            odir = pwd()
            try
                cd(dirname(ident))
                ptr = ccall((:asdkInit, libasdk), Ptr{Cvoid}, (Cstring,),
                            basename(ident))
            finally
                cd(odir)
            end
        else
            ptr = ccall((:asdkInit, libasdk), Ptr{Cvoid}, (Cstring,), ident)
        end
        if ptr == C_NULL
            code, mesg = lasterror()
            error("failed to open $ident ($mesg)")
        end
        # Create object first to be able to release resources in case of errors.
        dm = new(ptr, 0, Scalar[])
        try
            num = Int(dm["NbOfActuator"])
            dm.num = num
            fill!(resize!(dm.cmd, num), zero(Scalar))
        catch err
            close(dm)
            rethrow(err)
        end
        return finalizer(close, dm)
    end
end

function Base.close(dm::DeformableMirror)
    if dm.ptr != C_NULL
        status = ccall((:asdkRelease, libasdk), Status,
                       (Ptr{DeformableMirror},), dm.ptr)
        dm.ptr = C_NULL # to never release more than once
        _check(status)
    end
end

Base.length(dm::DeformableMirror) = dm.num
Base.eltype(::DeformableMirror) = Scalar
Base.unsafe_convert(::Type{Ptr{DeformableMirror}}, dm::DeformableMirror) =
    (dm.ptr != C_NULL ? dm.ptr : error("device has been closed"))
Base.extrema(dm::DeformableMirror) = (minimum(dm), maximum(dm))
Base.minimum(dm::DeformableMirror) = CMDMIN
Base.maximum(dm::DeformableMirror) = CMDMAX
Base.keys(dm::DeformableMirror) = (
    "AckTimeout", "DacReset", "ItfState", "LogDump", "LogPrintLevel",
    "NbOfActuator", "SyncMode", "TriggerMode", "TriggerIn",
    "UseException", "VersionInfo")

function Base.VersionNumber(dm::DeformableMirror)
    value = Int(dm["VersionInfo"])
    value, build = divrem(value, 10_000)
    value, patch = divrem(value, 100)
    major, minor = divrem(value, 100)
    return VersionNumber("$major.$minor.$patch-$build")
end

"""
    send(dm, cmd) -> actcmd

sends actuator commands `cmd` to deformable mirror `dm` and returns a vector of
actual commands sent to the mirror.  The actual commands may be different from
`cmd` due to bound constraints.  Actual commands are stored in an internal
buffer specific to the deformable mirror instance and allocated at construction
time.

See also: [`send!`](@ref), [`Alpao.lastcommand`](@ref).

"""
function send(dm::DeformableMirror, cmd::AbstractVector{<:AbstractFloat})
    num = length(dm) # FIXME: there may be several mirrors
    axes(cmd,1) == Base.OneTo(num) ||
        throw(DimensionMismatch("invalid indices for commands"))
    length(dm.cmd) == num || resize!(dm.cmd, num)
    @inbounds for i in 1:num
        dm.cmd[i] = clamp(Scalar(cmd[i]), CMDMIN, CMDMAX)
    end
    _check(ccall((:asdkSend, libasdk), Status,
                 (Ptr{DeformableMirror}, Ptr{Scalar}),
                 dm, dm.cmd))
    return dm.cmd
end

"""
    send!(dm, cmd) -> cmd

sends actuator commands `cmd` to deformable mirror `dm` and return, possibly
modified, commands.  The command values in `cmd` may be modified due to bound
constraints.  Thus, on return, `cmd` contains the actual commands sent to the
mirror.

See also: [`send`](@ref), [`Alpao.lastcommand`](@ref).

"""
send!(dm::DeformableMirror, cmd::AbstractVector{<:AbstractFloat}) =
    copy!(cmd, send(dm, cmd))

"""
    Alpao.lastcommand(dm) -> cmd

yields the last commands actually sent to the deformable mirror `dm`.

See also: [`send(::DeformableMirror)`](@ref).

"""
lastcommand(dm::DeformableMirror) = dm.cmd

"""
    Alpao.lasterror() -> (code, mesg)

pops the last error from the stack and returns error code and corresponding
message.

"""
function lasterror()
    code = Ref{UInt32}(0)
    mesg = zeros(UInt8, 512) # message buffer filled with zeroes
    if ccall((:asdkGetLastError, libasdk), Status,
             (Ptr{UInt32}, Ptr{UInt8}, Csize_t),
             code, mesg, sizeof(mesg)) != SUCCESS
        error("failed to retrieve last error message")
    end
    mesg[end] = 0
    (code[], String(mesg))
end

printlasterror() = ccall((:asdkPrintLastError, libasdk), Cvoid, ())

function _check(status::Status)
    if status != SUCCESS
        code, mesg = lasterror()
        error(mesg)
    end
end

stop(dm::DeformableMirror) =
    _check(ccall((:asdkStop, libasdk), Status, (Ptr{DeformableMirror},), dm))

function reset(dm::DeformableMirror)
    fill!(dm.cmd, Scalar(0))
    _check(ccall((:asdkReset, libasdk), Status, (Ptr{DeformableMirror},), dm))
end

Base.getindex(dm::DeformableMirror) = dm.cmd
Base.getindex(dm::DeformableMirror, ::Colon) = dm.cmd[:]
Base.getindex(dm::DeformableMirror, i::Integer) = dm.cmd[i]
Base.getindex(dm::DeformableMirror, i::AbstractUnitRange{<:Integer}) = dm.cmd[i]
function Base.getindex(dm::DeformableMirror, key::Keyword)
    val = Ref{Scalar}(0)
    _check(ccall((:asdkGet, libasdk), Status,
                 (Ptr{DeformableMirror}, Cstring, Ptr{Scalar}),
                 dm, key, val))
    return val[]
end

Base.setindex!(dm::DeformableMirror, val::Real, key::Keyword) = begin
    _check(ccall((:asdkSet, libasdk), Status,
                 (Ptr{DeformableMirror}, Cstring, Scalar),
                 dm, key, val))
    return dm
end

Base.setindex!(dm::DeformableMirror, val::AbstractString, key::Keyword) = begin
    _check(ccall((:asdkSetString, libasdk), Status,
                 (Ptr{DeformableMirror}, Cstring, Cstring),
                 dm, key, val))
    return dm
end

"""

`runtests(name)` run simple tests for deformable mirror identified by `name`.

"""
function runtests(name::String="BOL143")
    dm = DeformableMirror(name)
    num = length(dm)
    println("Number of actuators ", round(Int, dm["NbOfActuator"]))
    println("Number of actuators ", num)
    image_count = 0
    tot_image_count = 0
    cmd = zeros(Scalar, num)
    while tot_image_count < 10000
        for i in 1:num
            cmd[i] = 0.12;
            dm.send( cmd );
            cmd[i] = 0.0;
            if time() - t >= 1.0
                print(image_count," FPS\r")
                image_count = 0
                t = time()
            end
            image_count += 1
            tot_image_count += 1
        end
    end
    println()
end

end # module Alpao


text:
export len, siz, siz3, sizem, sizen, sizeo, rangem, rangen, rangeo, range

len(a) = length(a)
len(a::AbstractArray{T,N}) where {T,N} = size(a,N)

siz(a) = ones(Int, 2, 1)

function siz(a::Tuple)
    r = ones(Int, 2, 1)
    r[1] = 1
    r[2] = length(a)
    r
end

function siz(a::Union{UnitRange, Char, AbstractString})
    r = ones(Int, 2,1)
    r[2] = length(a)
    r
end

function siz(a::AbstractArray{T,N}) where {T,N}
    r = ones(Int, max(2,N), 1)
    for i = 1:N
        r[i] = size(a,i)
    end
    r
end

function siz3(a::Array{T,N}) where {T,N}
    if N > 3
        error("siz3 can only be used for arrays with ndims <= 3")
    end
    r = ones(Int, 3, 1)
    for i = 1:N
        r[i] = size(a,i)
    end
    r
end
siz3(a) = siz(a)

sizem(a) = size(a,1)
sizen(a) = size(a,2)
sizeo(a) = size(a,3)
rangem(a) = 1:size(a,1)
rangen(a) = 1:size(a,2)
rangeo(a) = 1:size(a,3)

import Base.range
range(a::AbstractArray) = 1:len(a)

import Base.ndims
ndims(a::Tuple) = 1


text:
#
#  OutputData.jl
#  CollagenPolymer
#
#  Created by Christopher Revell on 30/03/2020.
#
#

module OutputData

# Import Julia packages
using DelimitedFiles
using LinearAlgebra

# Import local program modules

function outputData(pos,outfile)#,t,tMax)#,nMonomers,nDomains,σ)

    writedlm(outfile,pos,", ")
    flush(outfile)

    # Measure monomer lengths
    # for i=0:nMonomers-1
    #     l = 0
    #     for j=1:nDomains-1
    #         dx = pos[i*nDomains+j+1,:]-pos[i*nDomains+j,:]
    #         l += sqrt(dot(dx,dx))
    #     end
    #     println("$(i+1) => $(l+σ), nDomains=$nDomains")
    # end

    return nothing

end

export outputData

end

max_stars_repo_name: chris-revell/collagen-model-jl

max_issues_repo_name: chris-revell/collagen-model-jl

max_forks_repo_name: chris-revell/collagen-model-jl


text:
# This file was generated, do not modify it. # hide
using Gnuplot, ColorSchemes, RDatasets, Colors
function gp_palette(colormap=:viridis)
    cmap = get(colorschemes[colormap], LinRange(0,1,256))
    ctmp = "0 '#$(hex(cmap[1]))',"
    for i in 2:256; ctmp = ctmp*"$(i-1) '#$(hex(cmap[i]))'," end;
    "set palette defined("*ctmp[1:end-1]*")"
end

volcano = Matrix{Float64}(dataset("datasets", "volcano"))
@gsp(volcano, "w image", "set view map", "set auto fix",
    gp_palette(:inferno), title = "Auckland s Maunga Whau Volcano")
save(term="pngcairo size 900,600", output="plt3d_ex6.png")


text:
function Interval(arg0::jdouble, arg1::jdouble)
    return Interval((jdouble, jdouble), arg0, arg1)
end

function check_point(obj::Interval, arg0::jdouble, arg1::jdouble)
    return jcall(obj, "checkPoint", Region_Location, (jdouble, jdouble), arg0, arg1)
end

function get_barycenter(obj::Interval)
    return jcall(obj, "getBarycenter", jdouble, ())
end

function get_inf(obj::Interval)
    return jcall(obj, "getInf", jdouble, ())
end

function get_size(obj::Interval)
    return jcall(obj, "getSize", jdouble, ())
end

function get_sup(obj::Interval)
    return jcall(obj, "getSup", jdouble, ())
end


text:
###############################################################################
#
# The Stopping structure eases the implementation of algorithms and the
# stopping criterion.
#
# The following examples illustrate a specialized Stopping for the
# LinearOperators type.
# https://github.com/JuliaSmoothOptimizers/LinearOperators.jl
#
using LinearOperators #v.1.0.1
using LinearAlgebra, SparseArrays, Main.Stopping, Test
#
# This tutorial illustrates the different step in preparing the resolution of a
# new problem.
# - This Stopping handles linear system
#  Ax = b with A an m x n matrix.
#  or linear Least Square problem
#  min_{x ∈ ℜⁿ} || Ax - b ||_p^2
# - evaluations are measured by the LinearOperator.
# - we will use the GenericState (a specialized State would keep track of evals)
#
###############################################################################
mutable struct LinearSystem
    A :: Union{AbstractLinearOperator, AbstractMatrix}
    b :: AbstractVector
end

function linear_system_check!(pb    :: LinearSystem,
                              state :: AbstractState;
                              pnorm :: Float64 = Inf,
                              kwargs...)
 return norm(pb.A * state.x - pb.b, pnorm)
end

"""
Type: LinearOperatorStopping (specialization of GenericStopping)
Methods: start!, stop!, update_and_start!, update_and_stop!, fill_in!, reinit!, status

Stopping structure for non-linear programming problems using NLPModels.
    Input :
       - pb         : a LinearSystem
       - optimality_check : a stopping criterion through an admissibility function
       - state      : The information relative to the problem, see GenericState
       - max_cntrs  : Dict contains the max number of evaluations
       - (opt) meta : Metadata relative to stopping criterion.
       - (opt) main_stp : Stopping of the main loop in case we consider a Stopping
                          of a subproblem.
                          If not a subproblem, then nothing.

 Note:
 * optimality_check : takes two inputs (LinearSystem, AbstractState)
 and returns a Float64 to be compared at 0.
 * Not specific State targeted
 * State don't necessarily keep track of evals
 * Evals are checked only for pb.A being a LinearOperator

 Warning:
 * optimality_check does not necessarily fill in the State.
 """
mutable struct LinearOperatorStopping <: AbstractStopping

    # problem
    pb :: LinearSystem
    # stopping criterion
    optimality_check :: Function
    # Common parameters
    meta      :: AbstractStoppingMeta
    # Contains the max number of evaluations
    max_cntrs :: Dict
    # current state of the problem
    current_state :: AbstractState
    # Stopping of the main problem, or nothing
    main_stp :: Union{AbstractStopping, Nothing}

    function LinearOperatorStopping(pb             :: LinearSystem,
                                    admissible     :: Function,
                                    current_state  :: AbstractState;
                                    meta           :: AbstractStoppingMeta = StoppingMeta(),
                                    max_cntrs      :: Dict = _init_max_counters_linear_operators(),
                                    main_stp       :: Union{AbstractStopping, Nothing} = nothing,
                                    kwargs...)

        if !(isempty(kwargs))
           meta = StoppingMeta(;kwargs...)
        end

        return new(pb, admissible, meta, max_cntrs, current_state, main_stp)
    end
end

"""
_init_max_counters_linear_operators():
"""
function _init_max_counters_linear_operators(;nprod   :: Int64 = 20000,
                                              ntprod  :: Int64 = 20000,
                                              nctprod :: Int64 = 20000,
                                              sum     :: Int64 = 20000*11)

  cntrs = Dict([(:nprod,   nprod),   (:ntprod, ntprod),
                (:nctprod, nctprod), (:neval_sum,    sum)])

 return cntrs
end

import Main.Stopping: _optimality_check, _resources_check!, start!, stop!, update_and_start!, update_and_stop!, fill_in!, reinit!, status

function _optimality_check(stp  :: LinearOperatorStopping; kwargs...)

 optimality = stp.optimality_check(stp.pb, stp.current_state; kwargs...)
 stp.current_state.current_score = optimality

 return optimality
end

"""
_resources_check!: check if the optimization algorithm has exhausted the resources.
                   This is the LinearOperator specialized version that takes
                   into account the evaluation of the functions.

Note:
* function does _not_ keep track of the evals in the State
* check :nprod, :ntprod, :nctprod in the LinearOperator entries
"""
function _resources_check!(stp    :: LinearOperatorStopping,
                           x      :: Union{Vector, Nothing})

  max_cntrs = stp.max_cntrs

  # check all the entries in the counter
  max_f = false
  sum   = 0
 if typeof(stp.pb.A) <: AbstractLinearOperator
  for f in [:nprod, :ntprod, :nctprod]
   max_f = max_f || (getfield(stp.pb.A, f) > max_cntrs[f])
   sum  += getfield(stp.pb.A, f)
  end
 end

 # Maximum number of function and derivative(s) computation
 max_evals = sum > max_cntrs[:neval_sum]

 # global user limit diagnostic
 stp.meta.resources = max_evals || max_f

 return stp
end

##############################################################
#TEST

n = 5
A = rand(n,n)
xref = 100 * rand(n)
x0 = zeros(n)
b    = A * xref
sA  = sparse(A)
opA = LinearOperator(A)

mLO = LinearSystem(A, b)
sLO = LinearSystem(sA, b)
opLO = LinearSystem(opA, b)
mLOstp = LinearOperatorStopping(mLO, linear_system_check!, GenericState(x0))
sLOstp = LinearOperatorStopping(sLO, linear_system_check!, GenericState(x0))
maxcn = _init_max_counters_linear_operators(nprod = 1)
opLOstp = LinearOperatorStopping(opLO, linear_system_check!, GenericState(x0), max_cntrs = maxcn)
@test start!(mLOstp) == false
@test start!(sLOstp) == false
@test start!(opLOstp) == false

update_and_stop!(mLOstp, x = xref)
@test status(mLOstp) == :Optimal
@test mLOstp.meta.resources == false
update_and_stop!(sLOstp, x = xref)
@test status(sLOstp) == :Optimal
@test sLOstp.meta.resources == false
update_and_stop!(opLOstp, x = xref)
@test status(opLOstp) == :Optimal
@test opLOstp.meta.resources == true

"""
Randomized block Kaczmarz
"""
function RandomizedBlockKaczmarz(stp :: AbstractStopping; kwargs...)

    #A,b = stp.current_state.Jx, stp.current_state.cx
    A,b = stp.pb.A, stp.pb.b
    x0  = stp.current_state.x

    m,n = size(A)
    xk  = x0

    OK = start!(stp)

    while !OK

     i  = Int(floor(rand() * m)+1) #rand a number between 1 and m
     Ai = A[i,:]
     xk  = Ai == 0 ? x0 : x0 - (dot(Ai,x0)-b[i])/dot(Ai,Ai) * Ai
     #xk  = Ai == 0 ? x0 : x0 - (Ai' * x0-b[i])/(Ai' * Ai) * Ai

     OK = update_and_stop!(stp, x = xk)
     x0  = xk

    end

 return stp
end

m, n = 400, 200 #size of A: m x n
A    = 100 * rand(m, n)
xref = 100 * rand(n)
b    = A * xref

#Our initial guess
x0 = zeros(n)

la_stop = LinearOperatorStopping(LinearSystem(A, b), linear_system_check!, GenericState(x0), max_iter = 150000, rtol = 1e-6)
#Be careful using GenericState(x0) would not work here without forcing convert = true
#in the update function. As the iterate will be an SparseVector to the contrary of initial guess.
#Tangi: maybe start! should send a Warning for such problem !?
sa_stop = LinearOperatorStopping(LinearSystem(sparse(A), b), linear_system_check!, GenericState(sparse(x0)), max_iter = 150000, rtol = 1e-6)
op_stop = LinearOperatorStopping(LinearSystem(LinearOperator(A), b), linear_system_check!, GenericState(x0), max_iter = 150000, rtol = 1e-6)

@time RandomizedBlockKaczmarz(la_stop)
@test status(la_stop) == :Optimal
@time RandomizedBlockKaczmarz(sa_stop)
@test status(sa_stop) == :Optimal
#No compatible dot product with LinearOperators.
#@time RandomizedBlockKaczmarz(op_stop)
#@test status(op_stop) == :Optimal


text:
module TestComponentsOrdering

using Mimi
using Test

my_model = Model()

#Testing that you cannot add two components of the same name
@defcomp testcomp1 begin
    var1 = Variable(index=[time])
    par1 = Parameter(index=[time])

    function run_timestep(p, v, d, t)
        v.var1[t] = p.par1[t]
    end
end

@defcomp testcomp2 begin
    var1 = Variable(index=[time])
    par1 = Parameter(index=[time])

    function run_timestep(p, v, d, t)
        v.var1[t] = p.par1[t]
    end
end

@defcomp testcomp3 begin
    var1 = Variable(index=[time])
    par1 = Parameter(index=[time])

    function run_timestep(p, v, d, t)
        v.var1[t] = p.par1[t]
    end
end

set_dimension!(my_model, :time, 2015:5:2110)
addcomponent(my_model, testcomp1)

@test_throws ErrorException addcomponent(my_model, testcomp1)

# Testing to catch adding component twice
@test_throws ErrorException addcomponent(my_model, testcomp1)

# Testing to catch if before or after does not exist
@test_throws ErrorException addcomponent(my_model, testcomp2, before=:testcomp3)

@test_throws ErrorException addcomponent(my_model, testcomp2, after=:testcomp3)

end #module


text:
module TestCovariance
    using StreamStats
    using Distributions
    using Base.Test

    # Means of uniform draws
    for n in rand(1:1_000_000, 10)
        xs = randn(n)
        ys = rand() * xs + randn(n)
        stat = StreamStats.Covariance()
        for (x, y) in zip(xs, ys)
            update!(stat, x, y)
        end
        online_cov = cov(stat)
        online_cor = cor(stat)
        online_n = nobs(stat)
        batch_cov = cov(xs, ys)
        batch_cor = cor(xs, ys)

        @test_approx_eq(online_cov, batch_cov)
        @test_approx_eq(online_cor, batch_cor)
        @test online_n == n
    end
end


text:
"""
$(TYPEDEF)
"""
struct StandardODEProblem end

# Mu' = f
"""
$(TYPEDEF)

Defines an ODE problem.

# Fields

$(FIELDS)
"""
struct ODEProblem{uType,tType,isinplace,P,F,K,PT} <:
               AbstractODEProblem{uType,tType,isinplace}
  """The ODE is `du/dt = f(u,p,t)`."""
  f::F
  """The initial condition is `u(tspan[1]) = u0`."""
  u0::uType
  """The solution `u(t)` will be computed for `tspan[1] ≤ t ≤ tspan[2]`."""
  tspan::tType
  """Constant parameters to be supplied as the second argument of `f`."""
  p::P
  """A callback to be applied to every solver which uses the problem."""
  kwargs::K
  # TODO
  problem_type::PT
  @add_kwonly function ODEProblem{iip}(f::AbstractODEFunction{iip},
                                       u0,tspan,p=NullParameters(),
                                       problem_type=StandardODEProblem();
                                       kwargs...) where {iip}
    _tspan = promote_tspan(tspan)
    new{typeof(u0),typeof(_tspan),
       isinplace(f),typeof(p),typeof(f),
       typeof(kwargs),
       typeof(problem_type)}(
       f,u0,_tspan,p,kwargs,problem_type)
  end

  """
      ODEProblem{isinplace}(f,u0,tspan,p=NullParameters(),callback=CallbackSet())

  Define an ODE problem with the specified function.
  `isinplace` optionally sets whether the function is inplace or not.
  This is determined automatically, but not inferred.
  """
  function ODEProblem{iip}(f,u0,tspan,p=NullParameters();kwargs...) where {iip}
    ODEProblem(convert(ODEFunction{iip},f),u0,tspan,p;kwargs...)
  end

  @add_kwonly function ODEProblem{iip,recompile}(f,u0,tspan,p=NullParameters();kwargs...) where {iip,recompile}
    if !recompile
      if iip
        ODEProblem{iip}(wrapfun_iip(f,(u0,u0,p,tspan[1])),u0,tspan,p;kwargs...)
      else
        ODEProblem{iip}(wrapfun_oop(f,(u0,p,tspan[1])),u0,tspan,p;kwargs...)
      end
    else
      ODEProblem{iip}(f,u0,tspan,p;kwargs...)
    end
  end
end

"""
    ODEProblem(f::ODEFunction,u0,tspan,p=NullParameters(),callback=CallbackSet())

Define an ODE problem from an [`ODEFunction`](@ref).
"""
function ODEProblem(f::AbstractODEFunction,u0,tspan,args...;kwargs...)
  ODEProblem{isinplace(f)}(f,u0,tspan,args...;kwargs...)
end

function ODEProblem(f,u0,tspan,p=NullParameters();kwargs...)
  ODEProblem(convert(ODEFunction,f),u0,tspan,p;kwargs...)
end

"""
$(TYPEDEF)
"""
abstract type AbstractDynamicalODEProblem end

"""
$(TYPEDEF)
"""
struct DynamicalODEProblem{iip} <: AbstractDynamicalODEProblem end
# u' = f2(v)
# v' = f1(t,u)
"""
    DynamicalODEProblem(f::DynamicalODEFunction,v0,u0,tspan,p=NullParameters(),callback=CallbackSet())

Define a dynamical ODE function from a [`DynamicalODEFunction`](@ref).
"""
function DynamicalODEProblem(f::DynamicalODEFunction,du0,u0,tspan,p=NullParameters();kwargs...)
  ODEProblem(f,ArrayPartition(du0,u0),tspan,p;kwargs...)
end
function DynamicalODEProblem(f1,f2,du0,u0,tspan,p=NullParameters();kwargs...)
  ODEProblem(DynamicalODEFunction(f1,f2),ArrayPartition(du0,u0),tspan,p;kwargs...)
end

"""
    DynamicalODEProblem{isinplace}(f1,f2,v0,u0,tspan,p=NullParameters(),callback=CallbackSet())

Define a dynamical ODE problem from the two functions `f1` and `f2`.

# Arguments
* `f1` and `f2`: The functions in the ODE.
* `v0` and `u0`: The initial conditions.
* `tspan`: The timespan for the problem.
* `p`: Parameter values for `f1` and `f2`.
* `callback`: A callback to be applied to every solver which uses the problem. Defaults to nothing.

`isinplace` optionally sets whether the function is inplace or not.
This is determined automatically, but not inferred.
"""
function DynamicalODEProblem{iip}(f1,f2,du0,u0,tspan,p=NullParameters();kwargs...) where iip
  ODEProblem(DynamicalODEFunction{iip}(f1,f2),ArrayPartition(du0,u0),tspan,p,DynamicalODEProblem{iip}();kwargs...)
end

# u'' = f(t,u,du,ddu)
"""
$(TYPEDEF)
"""
struct SecondOrderODEProblem{iip} <: AbstractDynamicalODEProblem end
function SecondOrderODEProblem(f,du0,u0,tspan,p=NullParameters();kwargs...)
  iip = isinplace(f,5)
  SecondOrderODEProblem{iip}(f,du0,u0,tspan,p;kwargs...)
end

"""
    SecondOrderODEProblem{isinplace}(f,du0,u0,tspan,p=NullParameters(),callback=CallbackSet())

Define a second order ODE problem with the specified function.

# Arguments
* `f`: The function for the second derivative.
* `du0`: The initial derivative.
* `u0`: The initial condition.
* `tspan`: The timespan for the problem.
* `callback`: A callback to be applied to every solver which uses the problem. Defaults to nothing.
"""
function SecondOrderODEProblem{iip}(f,du0,u0,tspan,p=NullParameters();kwargs...) where iip
  if iip
    f2 = function (du,v,u,p,t)
      du .= v
    end
  else
    f2 = function (v,u,p,t)
      v
    end
  end
  _u0 = ArrayPartition((du0,u0))
  ODEProblem(DynamicalODEFunction{iip}(f,f2),_u0,tspan,p,
                  SecondOrderODEProblem{iip}();kwargs...)
end
function SecondOrderODEProblem(f::DynamicalODEFunction,du0,u0,tspan,p=NullParameters();kwargs...)
  iip = isinplace(f.f1, 5)
  _u0 = ArrayPartition((du0,u0))
  if f.f2.f === nothing
    if iip
      f2 = function (du,v,u,p,t)
        du .= v
      end
    else
      f2 = function (v,u,p,t)
        v
      end
    end
    return ODEProblem(DynamicalODEFunction{iip}(f.f1,f2;mass_matrix=f.mass_matrix,analytic=f.analytic),_u0,tspan,p,
                  SecondOrderODEProblem{iip}();kwargs...)
  else
    return ODEProblem(DynamicalODEFunction{iip}(f.f1,f.f2;mass_matrix=f.mass_matrix,analytic=f.analytic),_u0,tspan,p,
                  SecondOrderODEProblem{iip}();kwargs...)
  end
end

"""
$(TYPEDEF)
"""
abstract type AbstractSplitODEProblem end

"""
$(TYPEDEF)
"""
struct SplitODEProblem{iip} <: AbstractSplitODEProblem end
# u' = Au + f
function SplitODEProblem(f1,f2,u0,tspan,p=NullParameters();kwargs...)
  f = SplitFunction(f1,f2)
  SplitODEProblem(f,u0,tspan,p;kwargs...)
end
"""
    SplitODEProblem{isinplace}(f1,f2,u0,tspan,p=NullParameters();kwargs...)

Define a split ODE problem from separate functions `f1` and `f2`.

# Arguments
* `f1`, `f2`: The functions in the ODE.
* `u0`: The initial condition.
* `tspan`: The timespan for the problem.
* `p`: The parameters for the problem.
* `callback`: A callback to be applied to every solver which uses the problem. Defaults to nothing.

The `isinplace parameter` can be omitted and will be determined using the
signature of `f2`. Note that both `f1` and `f2` should support the in-place style if
`isinplace` is true or they should both support the out-of-place style if
`isinplace` is false. You cannot mix up the two styles.
"""
function SplitODEProblem{iip}(f1,f2,u0,tspan,p=NullParameters();kwargs...) where iip
  f = SplitFunction{iip}(f1,f2)
  SplitODEProblem(f,u0,tspan,p;kwargs...)
end

"""
$(SIGNATURES)

Define a split ODE problem from a [`SplitFunction`](@ref).
"""
SplitODEProblem(f::SplitFunction,u0,tspan,p=NullParameters();kwargs...) =
  SplitODEProblem{isinplace(f)}(f,u0,tspan,p;kwargs...)
function SplitODEProblem{iip}(f::SplitFunction,u0,tspan,p=NullParameters();kwargs...) where iip
  if f.cache === nothing && iip
    cache = similar(u0)
    f = SplitFunction{iip}(f.f1, f.f2; mass_matrix=f.mass_matrix,
                     _func_cache=cache, analytic=f.analytic)
  end
  ODEProblem(f,u0,tspan,p,SplitODEProblem{iip}();kwargs...)
end

abstract type AbstractIncrementingODEProblem end

struct IncrementingODEProblem{iip} <: AbstractIncrementingODEProblem end

function IncrementingODEProblem(f,u0,tspan,p=NullParameters();kwargs...)
  f = IncrementingODEFunction(f)
  IncrementingODEProblem(f,u0,tspan,p;kwargs...)
end

function IncrementingODEProblem{iip}(f,u0,tspan,p=NullParameters();kwargs...) where iip
  f = IncrementingODEFunction{iip}(f)
  IncrementingODEProblem(f,u0,tspan,p;kwargs...)
end

IncrementingODEProblem(f::IncrementingODEFunction,u0,tspan,p=NullParameters();kwargs...) =
  IncrementingODEProblem{isinplace(f)}(f,u0,tspan,p;kwargs...)

function IncrementingODEProblem{iip}(f::IncrementingODEFunction,u0,tspan,p=NullParameters();kwargs...) where iip
  ODEProblem(f,u0,tspan,p,IncrementingODEProblem{iip}();kwargs...)
end


text:
#=
Tools for solving the standard optimal savings / income fluctuation
problem for an infinitely lived consumer facing an exogenous income
process that evolves according to a Markov chain.

@author : Spencer Lyon <spencer.lyon@nyu.edu>

@date: 2014-08-18

References
----------

http://quant-econ.net/jl/ifp.html

=#
# using PyCall
# @pyimport scipy.optimize as opt
# brentq = opt.brentq

"""
Income fluctuation problem

##### Fields

- `u::Function` : Utility `function`
- `du::Function` : Marginal utility `function`
- `r::Real` : Strictly positive interest rate
- `R::Real` : The interest rate plus 1 (strictly greater than 1)
- `bet::Real` : Discount rate in (0, 1)
- `b::Real` :  The borrowing constraint
- `Pi::Matrix` : Transition matrix for `z`
- `z_vals::Vector` : Levels of productivity
- `asset_grid::AbstractVector` : Grid of asset values
"""
type ConsumerProblem
    u::Function
    du::Function
    r::Real
    R::Real
    bet::Real
    b::Real
    Pi::Matrix
    z_vals::Vector
    asset_grid::AbstractVector
end

"Marginal utility for log utility function"
default_du{T <: Real}(x::T) = 1.0 / x

"""
Constructor with default values for `ConsumerProblem`

##### Arguments

- `r::Real(0.01)` : Strictly positive interest rate
- `bet::Real(0.96)` : Discount rate in (0, 1)
- `Pi::Matrix{Float64}([0.6 0.4; 0.05 0.95])` : Transition matrix for `z`
- `z_vals::Vector{Float64}([0.5, 1.0])` : Levels of productivity
- `b::Real(0.0)` : Borrowing constraint
- `grid_max::Real(16)` : Maximum in grid for asset holdings
- `grid_size::Int(50)` : Number of points in grid for asset holdings
- `u::Function(log)` : Utility `function`
- `du::Function(x->1/x)` : Marginal utility `function`

##### Notes

$(____kwarg_note)

"""
function ConsumerProblem(r=0.01, bet=0.96, Pi=[0.6 0.4; 0.05 0.95],
                         z_vals=[0.5, 1.0], b=0.0, grid_max=16, grid_size=50,
                         u=log, du=default_du)
    R = 1 + r
    asset_grid = linspace_range(-b, grid_max, grid_size)

    ConsumerProblem(u, du, r, R, bet, b, Pi, z_vals, asset_grid)
end

# make kwarg version
function ConsumerProblem(;r=0.01, beta=0.96, Pi=[0.6 0.4; 0.05 0.95],
                         z_vals=[0.5, 1.0], b=0.0, grid_max=16, grid_size=50,
                         u=log, du=x -> 1./x)
    ConsumerProblem(r, beta, Pi, z_vals, b, grid_max, grid_size, u, du)
end

"""
$(____bellman_main_docstring).

##### Arguments

- `cp::ConsumerProblem` : Instance of `ConsumerProblem`
- `v::Matrix`: Current guess for the value function
- `out::Matrix` : Storage for output
- `;ret_policy::Bool(false)`: Toggles return of value or policy functions

##### Returns

None, `out` is updated in place. If `ret_policy == true` out is filled with the
policy function, otherwise the value function is stored in `out`.

"""
function bellman_operator!(cp::ConsumerProblem, V::Matrix, out::Matrix;
                           ret_policy::Bool=false)
    # simplify names, set up arrays
    R, Pi, bet, u, b = cp.R, cp.Pi, cp.bet, cp.u, cp.b
    asset_grid, z_vals = cp.asset_grid, cp.z_vals

    new_V = similar(V)
    new_c = similar(V)

    z_idx = 1:length(z_vals)

    # Linear interpolation of V along the asset grid
    vf(a, i_z) = CoordInterpGrid(asset_grid, V[:, i_z], BCnearest,
                                 InterpLinear)[a]

    # compute lower_bound for optimization
    opt_lb = minimum(z_vals) - 1e-5

    # solve for RHS of Bellman equation
    for (i_z, z) in enumerate(z_vals)
        for (i_a, a) in enumerate(asset_grid)

            function obj(c)
                y = sum([vf(R*a+z-c, j) * Pi[i_z, j] for j=z_idx])
                return -u(c)  - bet * y
            end
            res = optimize(obj, opt_lb, R.*a.+z.+b)
            c_star = res.minimum
            if ret_policy
                out[i_a, i_z] = c_star
            else
               out[i_a, i_z] = - obj(c_star)
            end
        end
    end
end

function bellman_operator(cp::ConsumerProblem, V::Matrix; ret_policy=false)
    out = similar(V)
    bellman_operator!(cp, V, out, ret_policy=ret_policy)
    return out
end

"""
$(____greedy_main_docstring).

##### Arguments

- `cp::CareerWorkerProblem` : Instance of `CareerWorkerProblem`
- `v::Matrix`: Current guess for the value function
- `out::Matrix` : Storage for output

##### Returns

None, `out` is updated in place to hold the policy function

"""
function get_greedy!(cp::ConsumerProblem, V::Matrix, out::Matrix)
    bellman_operator!(cp, V, out, ret_policy=true)
end

function get_greedy(cp::ConsumerProblem, V::Matrix)
    bellman_operator(cp, V, ret_policy=true)
end

"""
The approximate Coleman operator.

Iteration with this operator corresponds to policy function
iteration. Computes and returns the updated consumption policy
c.  The array c is replaced with a function cf that implements
univariate linear interpolation over the asset grid for each
possible value of z.

##### Arguments

- `cp::CareerWorkerProblem` : Instance of `CareerWorkerProblem`
- `c::Matrix`: Current guess for the policy function
- `out::Matrix` : Storage for output

##### Returns

None, `out` is updated in place to hold the policy function

"""
function coleman_operator!(cp::ConsumerProblem, c::Matrix, out::Matrix)
    # simplify names, set up arrays
    R, Pi, bet, du, b = cp.R, cp.Pi, cp.bet, cp.du, cp.b
    asset_grid, z_vals = cp.asset_grid, cp.z_vals
    z_size = length(z_vals)
    gam = R * bet
    vals = Array(Float64, z_size)

    # linear interpolation to get consumption function. Updates vals inplace
    function cf!(a, vals)
        for i=1:z_size
            vals[i] = CoordInterpGrid(asset_grid, c[:, i], BCnearest,
                                      InterpLinear)[a]
        end
        Void
    end

    # compute lower_bound for optimization
    opt_lb = minimum(z_vals) - 1e-2
    for (i_z, z) in enumerate(z_vals)
        for (i_a, a) in enumerate(asset_grid)
            function h(t)
                cf!(R*a+z-t, vals)  # update vals
                expectation = dot(du(vals), vec(Pi[i_z, :]))
                return abs(du(t) - max(gam * expectation, du(R*a+z+b)))
            end
            opt_ub = R*a + z + b  # addresses issue #8 on github
            res = optimize(h, min(opt_lb, opt_ub - 1e-2), opt_ub, method=:brent)
            out[i_a, i_z] = res.minimum
        end
    end
    return out
end

"""
Apply the Coleman operator for a given model and initial value

See the specific methods of the mutating version of this function for more
details on arguments
"""
function coleman_operator(cp::ConsumerProblem, c::Matrix)
    out = similar(c)
    coleman_operator!(cp, c, out)
    return out
end

function init_values(cp::ConsumerProblem)
    # simplify names, set up arrays
    R, bet, u, b = cp.R, cp.bet, cp.u, cp.b
    asset_grid, z_vals = cp.asset_grid, cp.z_vals
    shape = length(asset_grid), length(z_vals)
    V, c = Array(Float64, shape...), Array(Float64, shape...)

    # Populate V and c
    for (i_z, z) in enumerate(z_vals)
        for (i_a, a) in enumerate(asset_grid)
            c_max = R*a + z + b
            c[i_a, i_z] = c_max
            V[i_a, i_z] = u(c_max) ./ (1 - bet)
        end
    end

    return V, c
end

max_stars_repo_name: mwhchen/quantecon

max_issues_repo_name: mwhchen/quantecon

max_forks_repo_name: mwhchen/quantecon


text:
using .Turing, Random, ExprTools, Distributions, Test
using DynamicPPL: DynamicPPL, vsym, vinds, @varname, VarInfo, VarName, build_model_info, namedtuple, build_model_info

dir = splitdir(splitdir(pathof(DynamicPPL))[1])[1]
include(dir*"/test/test_utils/AllUtils.jl")

Random.seed!(129)

priors = 0 # See "new grammar" test.

macro custom(expr)
    (Meta.isexpr(expr, :call, 3) && expr.args[1] === :~) ||
        error("incorrect macro usage")
    quote
        $(esc(expr.args[2])) = 0.0
    end
end

@testset "build_model_info" begin
    @testset "single arg" begin
        expr = :(model(obs) = begin
            obs
        end)

        result = build_model_info(expr)
        @test result[:args] == [:obs]
        @test result[:arg_syms] == [:obs]
        @test result[:name] == :model
        result
    end

    @testset "default arg missing" begin
        expr = :(model(x = missing) = begin
            x
        end)

        result = build_model_info(expr)

        @test result[:args] == [:($(Expr(:kw, :x, :missing)))]
        @test result[:arg_syms] == [:x]
        args_nt = result[:args_nt]
        @test args_nt.args[2] ==  :(NamedTuple{(:x,), Tuple{Core.Typeof(x)}})
        @test args_nt.args[3] ==  :((x,))
        defaults_nt = result[:defaults_nt]
        @test defaults_nt.args[2] ==  :(NamedTuple{(:x,), Tuple{Core.Typeof(missing)}})
        @test defaults_nt.args[3] ==  :((missing,))
    end

    @testset "default arg with type annotation" begin
        expr1 = :(model(x::Int = 2) = begin
            x
        end)

        result1 = build_model_info(expr1)

        @test result1[:args] == [:($(Expr(:kw, :(x::Int), 2)))]
        @test result1[:arg_syms] == [:x]
        args_nt = result1[:args_nt]
        @test args_nt.args[2] ==  :(NamedTuple{(:x,), Tuple{Core.Typeof(x)}})
        @test args_nt.args[3] ==  :((x,))
        defaults_nt = result1[:defaults_nt]
        @test defaults_nt.args[2] ==  :(NamedTuple{(:x,), Tuple{Core.Typeof(2)}})
        @test defaults_nt.args[3] ==  :((2,))

        expr2 = :(model(x::Vector{Float64} = [1,2,3]) = begin
            x
        end)

        result2 = build_model_info(expr2)
        @test result2[:arg_syms] == [:x]
    end

    @testset "default arg type" begin
        expr1 = :(model(::Type{T}=Float64) where {T <: Float64} = begin
            T
        end)

        result1 = build_model_info(expr1)
        @test result1[:args] == [:($(Expr(:kw, :(T::Type{<:Float64}), :Float64)))]
        @test result1[:arg_syms] == [:T]

        expr2 = :(model(::Type{T}=Float64) where {T} = begin
            T
        end)

        result2 = build_model_info(expr2)
        @test result2[:args] == [:($(Expr(:kw, :(T::Type{<:Any}), :Float64)))]
        @test result2[:arg_syms] == [:T]

        #TODO support t::Type{T}
        expr3 = :(model(t::Type{T}=Float64) where {T <: Float64} = begin
            T
        end)

        @test_throws ArgumentError build_model_info(expr3)

        #Invalid expression
        expr4 = :(model(::Type{T}=Float64) where {X} = begin
            T
        end)
        @test_throws ArgumentError build_model_info(expr4)
    end
end

@testset "compiler.jl" begin
    @testset "assume" begin
        @model test_assume() = begin
            x ~ Bernoulli(1)
            y ~ Bernoulli(x / 2)
            x, y
        end

        smc = SMC()
        pg = PG(10)

        res1 = sample(test_assume(), smc, 1000)
        res2 = sample(test_assume(), pg, 1000)

        check_numerical(res1, [:y], [0.5], atol=0.1)
        check_numerical(res2, [:y], [0.5], atol=0.1)

        # Check that all xs are 1.
        @test all(isone, res1[:x].value)
        @test all(isone, res2[:x].value)
    end
    @testset "beta binomial" begin
        prior = Beta(2,2)
        obs = [0,1,0,1,1,1,1,1,1,1]
        exact = Beta(prior.α + sum(obs), prior.β + length(obs) - sum(obs))
        meanp = exact.α / (exact.α + exact.β)

        @model testbb(obs) = begin
            p ~ Beta(2,2)
            x ~ Bernoulli(p)
            for i = 1:length(obs)
                obs[i] ~ Bernoulli(p)
            end
            p, x
        end

        smc = SMC()
        pg = PG(10)
        gibbs = Gibbs(HMC(0.2, 3, :p), PG(10, :x))

        chn_s = sample(testbb(obs), smc, 1000)
        chn_p = sample(testbb(obs), pg, 2000)
        chn_g = sample(testbb(obs), gibbs, 1500)

        check_numerical(chn_s, [:p], [meanp], atol=0.05)
        check_numerical(chn_p, [:x], [meanp], atol=0.1)
        check_numerical(chn_g, [:x], [meanp], atol=0.1)
    end
    @testset "forbid global" begin
        xs = [1.5 2.0]
        # xx = 1

        @model fggibbstest(xs) = begin
            s ~ InverseGamma(2,3)
            m ~ Normal(0,sqrt(s))
            # xx ~ Normal(m, sqrt(s)) # this is illegal

            for i = 1:length(xs)
                xs[i] ~ Normal(m, sqrt(s))
                # for xx in xs
                # xx ~ Normal(m, sqrt(s))
            end
            s, m
        end

        gibbs = Gibbs(PG(10, :s), HMC(0.4, 8, :m))
        chain = sample(fggibbstest(xs), gibbs, 2);
    end
    @testset "model macro" begin
        @model testmodel_comp(x, y) = begin
            s ~ InverseGamma(2,3)
            m ~ Normal(0,sqrt(s))

            x ~ Normal(m, sqrt(s))
            y ~ Normal(m, sqrt(s))

            return x, y
        end
        testmodel_comp(1.0, 1.2)

        # check if drawing from the prior works
        @model testmodel0(x = missing) = begin
            x ~ Normal()
            return x
        end
        f0_mm = testmodel0()
        @test mean(f0_mm() for _ in 1:1000) ≈ 0. atol=0.1

        # Test #544
        @model testmodel0(x = missing) = begin
            if x === missing
                x = Vector{Float64}(undef, 2)
            end
            x[1] ~ Normal()
            x[2] ~ Normal()
            return x
        end
        f0_mm = testmodel0()
        @test all(x -> isapprox(x, 0; atol = 0.1), mean(f0_mm() for _ in 1:1000))

        @model testmodel01(x = missing) = begin
            x ~ Bernoulli(0.5)
            return x
        end
        f01_mm = testmodel01()
        @test mean(f01_mm() for _ in 1:1000) ≈ 0.5 atol=0.1

        # test if we get the correct return values
        @model testmodel1(x1, x2) = begin
            s ~ InverseGamma(2,3)
            m ~ Normal(0,sqrt(s))

            x1 ~ Normal(m, sqrt(s))
            x2 ~ Normal(m, sqrt(s))

            return x1, x2
        end
        f1_mm = testmodel1(1., 10.)
        @test f1_mm() == (1, 10)
        f1_mm = testmodel1(x1=1., x2=10.)
        @test f1_mm() == (1, 10)

        @info "Testing the compiler's ability to catch bad models..."

        # Test for assertions in observe statements.
        @model brokentestmodel_observe1(x1, x2) = begin
            s ~ InverseGamma(2,3)
            m ~ Normal(0,sqrt(s))

            x1 ~ Normal(m, sqrt(s))
            x2 ~ x1 + 2

            return x1, x2
        end

        btest = brokentestmodel_observe1(1., 2.)
        @test_throws ArgumentError btest()

        @model brokentestmodel_observe2(x) = begin
            s ~ InverseGamma(2,3)
            m ~ Normal(0,sqrt(s))

            x = Vector{Float64}(undef, 2)
            x ~ [Normal(m, sqrt(s)), 2.0]

            return x
        end

        btest = brokentestmodel_observe2([1., 2.])
        @test_throws ArgumentError btest()

        # Test for assertions in assume statements.
        @model brokentestmodel_assume1() = begin
            s ~ InverseGamma(2,3)
            m ~ Normal(0,sqrt(s))

            x1 ~ Normal(m, sqrt(s))
            x2 ~ x1 + 2

            return x1, x2
        end

        btest = brokentestmodel_assume1()
        @test_throws ArgumentError btest()

        @model brokentestmodel_assume2() = begin
            s ~ InverseGamma(2,3)
            m ~ Normal(0,sqrt(s))

            x = Vector{Float64}(undef, 2)
            x ~ [Normal(m, sqrt(s)), 2.0]

            return x
        end

        btest = brokentestmodel_assume2()
        @test_throws ArgumentError btest()

        # Test missing input arguments
        @model testmodel(x) = begin
            x ~ Bernoulli(0.5)
            return x
        end
        @test_throws UndefKeywordError testmodel()

        # Test missing initialization for vector observation turned parameter
        @model testmodel(x) = begin
            x[1] ~ Bernoulli(0.5)
            return x
        end
        @test_throws MethodError testmodel(missing)()

        # Test use of internal names
        @model testmodel(x) = begin
            x[1] ~  Bernoulli(0.5)
            global varinfo_ = _varinfo
            global sampler_ = _sampler
            global model_ = _model
            global context_ = _context
            global lp = getlogp(_varinfo)
            return x
        end
        model = testmodel([1.0])
        varinfo = DynamicPPL.VarInfo(model)
        model(varinfo)
        @test getlogp(varinfo) == lp
        @test varinfo_ isa AbstractVarInfo
        @test model_ === model
        @test sampler_ === SampleFromPrior()
        @test context_ === DefaultContext()

        # disable warnings
        @model testmodel(x) = begin
            x[1] ~  Bernoulli(0.5)
            global varinfo_ = _varinfo
            global sampler_ = _sampler
            global model_ = _model
            global context_ = _context
            global lp = getlogp(_varinfo)
            return x
        end false
        lpold = lp
        model = testmodel([1.0])
        varinfo = DynamicPPL.VarInfo(model)
        model(varinfo)
        @test getlogp(varinfo) == lp == lpold

        # test DPPL#61
        @model testmodel(z) = begin
            m ~ Normal()
            z[1:end] ~ MvNormal(fill(m, length(z)), 1.0)
            return m
        end
        model = testmodel(rand(10))
        @test all(z -> isapprox(z, 0; atol = 0.2), mean(model() for _ in 1:1000))
    end
    @testset "nested model" begin
        function makemodel(p)
            @model testmodel(x) = begin
                x[1] ~ Bernoulli(p)
                global lp = getlogp(_varinfo)
                return x
            end
            return testmodel
        end
        model = makemodel(0.5)([1.0])
        varinfo = DynamicPPL.VarInfo(model)
        model(varinfo)
        @test getlogp(varinfo) == lp
    end
    @testset "new grammar" begin
        x = Float64[1 2]

        @model gauss(x) = begin
            priors = TArray{Float64}(2)
            priors[1] ~ InverseGamma(2,3)         # s
            priors[2] ~ Normal(0, sqrt(priors[1])) # m
            for i in 1:length(x)
                x[i] ~ Normal(priors[2], sqrt(priors[1]))
            end
            priors
        end

        chain = sample(gauss(x), PG(10), 10)
        chain = sample(gauss(x), SMC(), 10)

        @model gauss2(x, ::Type{TV}=Vector{Float64}) where {TV} = begin
            priors = TV(undef, 2)
            priors[1] ~ InverseGamma(2,3)         # s
            priors[2] ~ Normal(0, sqrt(priors[1])) # m
            for i in 1:length(x)
                x[i] ~ Normal(priors[2], sqrt(priors[1]))
            end
            priors
        end

        chain = sample(gauss2(x), PG(10), 10)
        chain = sample(gauss2(x=x, TV=Vector{Float64}), PG(10), 10)
        chain = sample(gauss2(x), SMC(), 10)
        chain = sample(gauss2(x=x, TV=Vector{Float64}), SMC(), 10)
    end
    @testset "new interface" begin
        obs = [0, 1, 0, 1, 1, 1, 1, 1, 1, 1]

        @model newinterface(obs) = begin
          p ~ Beta(2,2)
          for i = 1:length(obs)
            obs[i] ~ Bernoulli(p)
          end
          p
        end

        chain = sample(
            newinterface(obs),
            HMC{Turing.ForwardDiffAD{2}}(0.75, 3, :p, :x),
            100)
    end
    @testset "no return" begin
        @model noreturn(x) = begin
            s ~ InverseGamma(2,3)
            m ~ Normal(0, sqrt(s))
            for i in 1:length(x)
                x[i] ~ Normal(m, sqrt(s))
            end
        end

        chain = sample(noreturn([1.5 2.0]), HMC(0.15, 6), 1000)
        check_numerical(chain, [:s, :m], [49/24, 7/6])
    end
    @testset "observe" begin
        @model test() = begin
          z ~ Normal(0,1)
          x ~ Bernoulli(1)
          1 ~ Bernoulli(x / 2)
          0 ~ Bernoulli(x / 2)
          x
        end

        is  = IS()
        smc = SMC()
        pg  = PG(10)

        res_is = sample(test(), is, 10000)
        res_smc = sample(test(), smc, 1000)
        res_pg = sample(test(), pg, 100)

        @test all(isone, res_is[:x].value)
        @test res_is.logevidence ≈ 2 * log(0.5)

        @test all(isone, res_smc[:x].value)
        @test res_smc.logevidence ≈ 2 * log(0.5)

        @test all(isone, res_pg[:x].value)
    end
    @testset "sample" begin
        alg = Gibbs(HMC(0.2, 3, :m), PG(10, :s))
        chn = sample(gdemo_default, alg, 1000);
    end
    @testset "vectorization @." begin
        @model vdemo1(x) = begin
            s ~ InverseGamma(2,3)
            m ~ Normal(0, sqrt(s))
            @. x ~ Normal(m, sqrt(s))
            return s, m
        end

        alg = HMC(0.01, 5)
        x = randn(100)
        res = sample(vdemo1(x), alg, 250)

        @model vdemo1b(x) = begin
            s ~ InverseGamma(2,3)
            m ~ Normal(0, sqrt(s))
            @. x ~ Normal(m, $(sqrt(s)))
            return s, m
        end

        res = sample(vdemo1b(x), alg, 250)

        D = 2
        @model vdemo2(x) = begin
            μ ~ MvNormal(zeros(D), ones(D))
            @. x ~ $(MvNormal(μ, ones(D)))
        end

        alg = HMC(0.01, 5)
        res = sample(vdemo2(randn(D, 100)), alg, 250)

        # Vector assumptions
        N = 10
        setchunksize(N)
        alg = HMC(0.2, 4)

        @model vdemo3() = begin
            x = Vector{Real}(undef, N)
            for i = 1:N
                x[i] ~ Normal(0, sqrt(4))
            end
        end

        t_loop = @elapsed res = sample(vdemo3(), alg, 1000)

        # Test for vectorize UnivariateDistribution
        @model vdemo4() = begin
          x = Vector{Real}(undef, N)
          @. x ~ Normal(0, 2)
        end

        t_vec = @elapsed res = sample(vdemo4(), alg, 1000)

        @model vdemo5() = begin
            x ~ MvNormal(zeros(N), 2 * ones(N))
        end

        t_mv = @elapsed res = sample(vdemo5(), alg, 1000)

        println("Time for")
        println("  Loop : ", t_loop)
        println("  Vec  : ", t_vec)
        println("  Mv   : ", t_mv)

        # Transformed test
        @model vdemo6() = begin
            x = Vector{Real}(undef, N)
            @. x ~ InverseGamma(2, 3)
        end

        sample(vdemo6(), alg, 1000)

        N = 3
        @model vdemo7() = begin
            x = Array{Real}(undef, N, N)
            @. x ~ [InverseGamma(2, 3) for i in 1:N]
        end

        sample(vdemo7(), alg, 1000)
    end

    # Notation is ugly since `x .~ Normal(μ, σ)` cannot be parsed in Julia 1.0
    @testset "vectorization .~" begin
        @model vdemo1(x) = begin
            s ~ InverseGamma(2,3)
            m ~ Normal(0, sqrt(s))
            (.~)(x, Normal(m, sqrt(s)))
            return s, m
        end

        alg = HMC(0.01, 5)
        x = randn(100)
        res = sample(vdemo1(x), alg, 250)

        D = 2
        @model vdemo2(x) = begin
            μ ~ MvNormal(zeros(D), ones(D))
            (.~)(x, MvNormal(μ, ones(D)))
        end

        alg = HMC(0.01, 5)
        res = sample(vdemo2(randn(D,100)), alg, 250)

        # Vector assumptions
        N = 10
        setchunksize(N)
        alg = HMC(0.2, 4)

        @model vdemo3() = begin
            x = Vector{Real}(undef, N)
            for i = 1:N
                x[i] ~ Normal(0, sqrt(4))
            end
        end

        t_loop = @elapsed res = sample(vdemo3(), alg, 1000)

        # Test for vectorize UnivariateDistribution
        @model vdemo4() = begin
            x = Vector{Real}(undef, N)
            (.~)(x, Normal(0, 2))
        end

        t_vec = @elapsed res = sample(vdemo4(), alg, 1000)

        @model vdemo5() = begin
            x ~ MvNormal(zeros(N), 2 * ones(N))
        end

        t_mv = @elapsed res = sample(vdemo5(), alg, 1000)

        println("Time for")
        println("  Loop : ", t_loop)
        println("  Vec  : ", t_vec)
        println("  Mv   : ", t_mv)

        # Transformed test
        @model vdemo6() = begin
            x = Vector{Real}(undef, N)
            (.~)(x, InverseGamma(2, 3))
        end

        sample(vdemo6(), alg, 1000)

        @model vdemo7() = begin
            x = Array{Real}(undef, N, N)
            (.~)(x, [InverseGamma(2, 3) for i in 1:N])
        end

        sample(vdemo7(), alg, 1000)
    end

    @testset "Type parameters" begin
        N = 10
        setchunksize(N)
        alg = HMC(0.01, 5)
        x = randn(1000)
        @model vdemo1(::Type{T}=Float64) where {T} = begin
            x = Vector{T}(undef, N)
            for i = 1:N
                x[i] ~ Normal(0, sqrt(4))
            end
        end

        t_loop = @elapsed res = sample(vdemo1(), alg, 250)
        t_loop = @elapsed res = sample(vdemo1(Float64), alg, 250)
        t_loop = @elapsed res = sample(vdemo1(T=Float64), alg, 250)

        @model vdemo2(::Type{T}=Float64) where {T <: Real} = begin
            x = Vector{T}(undef, N)
            @. x ~ Normal(0, 2)
        end

        t_vec = @elapsed res = sample(vdemo2(), alg, 250)
        t_vec = @elapsed res = sample(vdemo2(Float64), alg, 250)
        t_vec = @elapsed res = sample(vdemo2(T=Float64), alg, 250)

        @model vdemo3(::Type{TV}=Vector{Float64}) where {TV <: AbstractVector} = begin
            x = TV(undef, N)
            @. x ~ InverseGamma(2, 3)
        end

        sample(vdemo3(), alg, 250)
        sample(vdemo3(Vector{Float64}), alg, 250)
        sample(vdemo3(TV=Vector{Float64}), alg, 250)
    end
    @testset "var name splitting" begin
        var_expr = :(x)
        @test vsym(var_expr) == :x
        @test vinds(var_expr) == :(())

        var_expr = :(x[1,1][2,3])
        @test vsym(var_expr) == :x
        @test vinds(var_expr) == :(((1, 1), (2, 3)))

        var_expr = :(x[:,1][2,:])
        @test vsym(var_expr) == :x
        @test vinds(var_expr) == :(((:, 1), (2, :)))

        var_expr = :(x[2:3,1][2,1:2])
        @test vsym(var_expr) == :x
        @test vinds(var_expr) == :(((2:3, 1), (2, 1:2)))

        var_expr = :(x[2:3,2:3][[1,2],[1,2]])
        @test vsym(var_expr) == :x
        @test vinds(var_expr) == :(((2:3, 2:3), ([1, 2], [1, 2])))

        var_expr = :(x[end])
        @test vsym(var_expr) == :x
        @test vinds(var_expr) == :((($lastindex(x),),))

        var_expr = :(x[1, end])
        @test vsym(var_expr) == :x
        @test vinds(var_expr) == :(((1, $lastindex(x, 2)),))
    end
    @testset "user-defined variable name" begin
        @model f1() = begin
            x ~ NamedDist(Normal(), :y)
        end
        @model f2() = begin
            x ~ NamedDist(Normal(), @varname(y[2][:,1]))
        end
        @model f3() = begin
            x ~ NamedDist(Normal(), @varname(y[1]))
        end
        vi1 = VarInfo(f1())
        vi2 = VarInfo(f2())
        vi3 = VarInfo(f3())
        @test haskey(vi1.metadata, :y)
        @test vi1.metadata.y.vns[1] == VarName(:y)
        @test haskey(vi2.metadata, :y)
        @test vi2.metadata.y.vns[1] == VarName(:y, ((2,), (Colon(), 1)))
        @test haskey(vi3.metadata, :y)
        @test vi3.metadata.y.vns[1] == VarName(:y, ((1,),))
    end
    @testset "custom tilde" begin
        @model demo() = begin
            $(@custom m ~ Normal())
            return m
        end
        model = demo()
        @test all(iszero(model()) for _ in 1:1000)
    end
    @testset "docstring" begin
        "This is a test"
        @model function demo(x)
            m ~ Normal()
            x ~ Normal(m, 1)
        end

        s = @doc(demo)
        @test string(s) == "This is a test\n"

        # Verify that adding docstring didn't completely break execution of model
        m = demo(0.)
        @test m() isa Float64
    end
end


text:
module GeneralizedSampling

using Compat
import Compat: view

using Base.Cartesian

using NFFT

using IntervalWavelets
import IntervalWavelets: isuniform, van_moment

import VoronoiCells: density, voronoiarea
import Wavelets: wavelet, WT
import ArrayViews: flatten_view, reshape_view
import Base: A_mul_B!, Ac_mul_B!

export
	# Types
	CoB,
	Freq2Wave,
	Freq2BoundaryWave,
	Freq2NoBoundaryWave,

	# Fourier transforms
	FourScalingFunc,
	FourHaarScaling,
	FourDaubScaling,

	# Linear equation solvers
	cgnr,

	# Special CoB functions
	collect,

	# misc
	had!,
	hadc!,
	yphad!,
	isuniform,
	grid,
	weights,
	density,
	van_moment,
	isdaubechies,
	ishaar,
	hasboundary,
	dim,
	split,
	wscale,
	wsize,
	frac,
	frac!

const sqrt2 = sqrt(2)
const SMALL_EPS = eps()
const LARGE_EPS= sqrt(eps())
const ComplexOne = one(Complex{Float64})
const ComplexZero = zero(Complex{Float64})
const twoπ = 2.0*pi

include("Freq2Wave/Types.jl")
include("Misc.jl")
include("FourierTransforms.jl")
include("NFFT.jl")
include("CGNR.jl")

end # module


text:
function plot_surface!(vis::Visualizer, env::Environment{R3};  col=zeros(3), α=0.4,
		n::Int=50, xlims = [-1.0, 5.0], ylims = [-2.0, 2.0])
    f(x) = x[3] - env.surf(x[1:2])[1]
    plot_surface!(vis, f, xlims=xlims, ylims=ylims, col=col, α=α, n=n)
    return nothing
end

function plot_surface!(vis::Visualizer, env::Environment{R2};  col=zeros(3), α=0.4,
		n::Int=50, xlims = [-1.0, 5.0], ylims = [-0.1, 0.1])
    f(x) = x[3] - env.surf(x[1:1])[1]
    plot_surface!(vis, f, xlims=xlims, ylims=ylims, col=col, α=α, n=n)
    return nothing
end

function plot_surface!(vis::Visualizer, f::Any; xlims = [-1.0, 5.0],
        ylims = [-0.1, 0.1], col=zeros(3), α=0.4, n::Int=200)
    mesh = GeometryBasics.Mesh(f,
    	HyperRectangle(Vec(xlims[1], ylims[1], -2.0), Vec(xlims[2]-xlims[1], ylims[2]-ylims[1], 4.0)),
        Meshing.MarchingCubes(), samples=(n, n, Int(floor(n/8))))
    setobject!(vis["surface"], mesh,
    		   MeshPhongMaterial(color=RGBA{Float32}(col..., α)))
    return nothing
end

function animate_robot!(vis::Visualizer, anim::MeshCat.Animation, model::ContactModel,
		q::AbstractVector; name::Symbol=model_name(model))
	for t in 1:length(q)
		MeshCat.atframe(anim, t) do
			set_robot!(vis, model, q[t], name=name)
		end
	end
	setanimation!(vis, anim)
	return nothing
end

function visualize_robot!(vis::Visualizer, model::ContactModel, q::AbstractVector;
		h=0.01, α=1.0,
		anim::MeshCat.Animation=MeshCat.Animation(Int(floor(1/h))),
		name::Symbol=model_name(model))

	build_robot!(vis, model, name=name, α=α)
	animate_robot!(vis, anim, model, q, name=name)
	return anim
end

function visualize_robot!(vis::Visualizer, model::ContactModel, traj::ContactTraj;
		sample=max(1, Int(floor(traj.H / 100))), h=traj.h*sample,  α=1.0,
		anim::MeshCat.Animation=MeshCat.Animation(Int(floor(1/h))),
		name::Symbol=model_name(model))

	visualize_robot!(vis, model, traj.q[3:sample:end]; anim=anim, name=name, h=h, α=α)
	return anim
end

function build_meshrobot!(vis::Visualizer, model::ContactModel, urdf::String,
		package_path::String; name::Symbol=model_name(model), α=1.0)
	default_background!(vis)

	mechanism = MeshCatMechanisms.parse_urdf(urdf)
    visuals = URDFVisuals(urdf, package_path=[package_path])
    state = MeshCatMechanisms.MechanismState(mechanism)
    vis_el = MeshCatMechanisms.visual_elements(mechanism, visuals)
    set_alpha!(vis_el,α)

    mvis = MechanismVisualizer(state, vis[name, :world])
    MeshCatMechanisms._set_mechanism!(mvis, vis_el)
    MeshCatMechanisms._render_state!(mvis)
	return mvis
end

function set_alpha!(visuals::Vector, α)
    for el in visuals
        c = el.color
        c_new = RGBA(red(c),green(c),blue(c),α)
        el.color = c_new
    end
end

function set_meshrobot!(vis::Visualizer, mvis::MechanismVisualizer, model::ContactModel,
		q::AbstractVector; name::Symbol=model_name(model))

	q_mesh, tform = convert_config(model, q)
	set_configuration!(mvis, q_mesh)
	settransform!(vis[name, :world], tform)

	return nothing
end

function animate_meshrobot!(vis::Visualizer, mvis::MechanismVisualizer, anim::MeshCat.Animation,
		model::ContactModel, q::AbstractVector; name::Symbol=model_name(model))
	for t in 1:length(q)
		MeshCat.atframe(anim, t) do
			set_meshrobot!(vis, mvis, model, q[t], name=name)
		end
	end
	setanimation!(vis, anim)
	return nothing
end

function visualize_meshrobot!(vis::Visualizer, model::ContactModel, q::AbstractVector;
		h=0.01,
		anim::MeshCat.Animation=MeshCat.Animation(Int(floor(1/h))),
		name::Symbol=model_name(model),
		α=α)

	mvis = build_meshrobot!(vis, model, name=name, α=α)
	animate_meshrobot!(vis, mvis, anim, model, q, name=name)

	return anim
end

function visualize_meshrobot!(vis::Visualizer, model::ContactModel, traj::ContactTraj;
		sample=max(1, Int(floor(traj.H / 100))), h=traj.h*sample,
		anim::MeshCat.Animation=MeshCat.Animation(Int(floor(1/h))),
		name::Symbol=model_name(model),
		α=1.0)

	anim = visualize_meshrobot!(vis, model, traj.q[3:sample:end]; anim=anim, name=name, h=h, α=α)
	return anim
end

function build_force!(vis::Visualizer, model::ContactModel; name::Symbol=model_name(model))
	default_background!(vis)
	orange_mat, blue_mat, black_mat = get_material()
	nc = model.dim.c

	for i = 1:nc
		fri_vis = ArrowVisualizer(vis[name][:force][:friction]["$i"])
		imp_vis = ArrowVisualizer(vis[name][:force][:impact]["$i"])
		con_vis = ArrowVisualizer(vis[name][:force][:contact]["$i"])
		setobject!(imp_vis, orange_mat)
		setobject!(fri_vis, black_mat)
		setobject!(con_vis, blue_mat)
	end
	return nothing
end

function contact_force(model::ContactModel, env::Environment, pc::AbstractVector,
		γ::AbstractVector, b::AbstractVector;
		E::AbstractMatrix=Matrix(friction_mapping(env)))
	nc = model.dim.c
	nb = nc * friction_dim(env)
	nf = Int(nb/nc)

	r = [Matrix(rotation(env, pc[i])') for i=1:nc]
	bc = [r[i] * [E*b[(i-1) * nf .+ (1:nf)];  0.0] for i = 1:nc] # TODO: make efficient
	λc = [r[i] * [E*b[(i-1) * nf .+ (1:nf)]; γ[i]] for i = 1:nc] # TODO: make efficient
	return bc, λc
end

function set_force!(vis::Visualizer, model::ContactModel, env::Environment, q::AbstractVector,
		γ::AbstractVector, b::AbstractVector; name::Symbol=model_name(model), shift=-0.14,
		E::AbstractMatrix=Matrix(friction_mapping(env)))

	nc = model.dim.c
	nb = nc * friction_dim(env)
	nf = Int(nb/nc)
	p_shift = [0, shift, 0]

	pc = contact_point(model, q)
	bc, λc = contact_force(model, env, pc, γ, b, E=E)

	imp_vis = [ArrowVisualizer(vis[name][:force][:impact]["$i"]) for i=1:nc]
	fri_vis = [ArrowVisualizer(vis[name][:force][:friction]["$i"]) for i=1:nc]
	con_vis = [ArrowVisualizer(vis[name][:force][:contact]["$i"]) for i=1:nc]

	for i = 1:nc
		bi = cast3d(bc[i])
		λi = cast3d(λc[i])

		settransform!(imp_vis[i],
			Point(0.0, 0.0, 0.0),
			Vec(γ[i], 0.0, 0.0),
			shaft_radius=0.008,
			max_head_radius=0.020)
		settransform!(fri_vis[i],
			Point(0.0, 0.0, 0.0),
			Vec(norm(bi), 0.0, 0.0),
			shaft_radius=0.008,
			max_head_radius=0.020)
		settransform!(con_vis[i],
			Point(0.0, 0.0, 0.0),
			Vec(norm(λi), 0.0, 0.0),
			shaft_radius=0.008,
			max_head_radius=0.020)

		trans = Translation(pc[i] + p_shift)
		rot = LinearMap(rotation_3d(env, pc[i])')
		transform = compose(trans, rot)
		settransform!(vis[name][:force][:impact]["$i"], transform)

		rot_fri = LinearMap(bivector_rotation([0.0, 0.0, 1.0], bi))
		transform_fri = compose(trans, rot_fri)
		settransform!(vis[name][:force][:friction]["$i"], transform_fri)

		rot_con = LinearMap(bivector_rotation([0.0, 0.0, 1.0], λi))
		transform_con = compose(trans, rot_con)
		settransform!(vis[name][:force][:contact]["$i"], transform_con)
	end
	return nothing
end

function animate_force!(vis::Visualizer, anim::MeshCat.Animation, model::ContactModel, env::Environment,
	q::AbstractVector, γ::AbstractVector, b::AbstractVector;
	name::Symbol=model_name(model), shift=-0.14)
	E = Matrix(friction_mapping(env))
	for t in 1:length(q)
		MeshCat.atframe(anim, t) do
			set_force!(vis, model, env, q[t], γ[t], b[t], name=name, E=E, shift=shift)
		end
	end
	setanimation!(vis, anim)
	return nothing
end

function visualize_force!(vis::Visualizer, model::ContactModel, env::Environment,
		q::AbstractVector, γ::AbstractVector, b::AbstractVector; h=0.01,
		anim::MeshCat.Animation=MeshCat.Animation(Int(floor(1/h))),
		name::Symbol=model_name(model), shift=-0.14)

	build_force!(vis, model, name=name)
	animate_force!(vis, anim, model, env, q, γ, b, name=name, shift=shift)
	return anim
end

function visualize_force!(vis::Visualizer, model::ContactModel, env::Environment,
		traj::ContactTraj; sample=max(1,Int(floor(traj.H/100))), h=traj.h*sample,
		anim::MeshCat.Animation=MeshCat.Animation(Int(floor(1/h))),
		name::Symbol=model_name(model), shift=-0.14)

	visualize_force!(vis, model, env, traj.q[3:sample:end], traj.γ[1:sample:end], traj.b[1:sample:end];
			anim=anim, name=name, h=h, shift=shift)
	return anim
end