text:
module Testoperators

using Test
import TreeParzen

CONFIG = TreeParzen.Config()

@testset "Delayed binary operators -- addition" begin

    choicevals = [1, 2, 3]

    a = TreeParzen.HP.Choice(:a, choicevals)
    b = TreeParzen.HP.Choice(:b, choicevals)

    expr = a + b

    trials = TreeParzen.Trials.Trial[]

    for i in 1:100
        trial = TreeParzen.ask(expr)
        TreeParzen.tell!(trials, trial, 1.)
        @test 2 <= trial.hyperparams <= 6
    end

    for i in 1:1000
        trial = TreeParzen.ask(expr, trials, CONFIG)
        @test 2 <= trial.hyperparams <= 6
    end

    # test scalar expressions real quick using identity
    trial = TreeParzen.ask(expr + 0)
    @test 2 <= trial.hyperparams <= 6
    trial = TreeParzen.ask(expr + 0, trials, CONFIG)
    @test 2 <= trial.hyperparams <= 6

end

@testset "Delayed binary operators -- subtraction" begin

    choicevals = [1, 2, 3]

    a = TreeParzen.HP.Choice(:a, choicevals)
    b = TreeParzen.HP.Choice(:b, choicevals)

    expr = a - b

    trials = TreeParzen.Trials.Trial[]

    for i in 1:100
        trial = TreeParzen.ask(expr)
        TreeParzen.tell!(trials, trial, 1.)
        @test -2 <= trial.hyperparams <= 2
    end

    for i in 1:1000
        trial = TreeParzen.ask(expr, trials, CONFIG)
        @test -2 <= trial.hyperparams <= 2
    end

    # test scalar expressions real quick using identity
    trial = TreeParzen.ask(expr - 0)
    @test -2 <= trial.hyperparams <= 2
    trial = TreeParzen.ask(expr - 0, trials, CONFIG)
    @test -2 <= trial.hyperparams <= 2

end

@testset "Delayed binary operators -- multiplication" begin

    choicevals = [1, 2, 3]

    a = TreeParzen.HP.Choice(:a, choicevals)
    b = TreeParzen.HP.Choice(:b, choicevals)

    expr = a * b

    trials = TreeParzen.Trials.Trial[]

    for i in 1:100
        trial = TreeParzen.ask(expr)
        TreeParzen.tell!(trials, trial, 1.)
        @test 1 <= trial.hyperparams <= 9
    end

    for i in 1:1000
        trial = TreeParzen.ask(expr, trials, CONFIG)
        @test 1 <= trial.hyperparams <= 9
    end

    # test scalar expressions real quick using identity
    trial = TreeParzen.ask(expr * 1)
    @test 1 <= trial.hyperparams <= 9
    trial = TreeParzen.ask(expr * 1, trials, CONFIG)
    @test 1 <= trial.hyperparams <= 9

end

@testset "Delayed binary operators -- division" begin

    choicevals = [1, 2, 3]

    a = TreeParzen.HP.Choice(:a, choicevals)
    b = TreeParzen.HP.Choice(:b, choicevals)

    expr = a / b

    trials = TreeParzen.Trials.Trial[]

    for i in 1:100
        trial = TreeParzen.ask(expr)
        TreeParzen.tell!(trials, trial, 1.)
        @test 0.333 <= trial.hyperparams <= 3
    end

    for i in 1:1000
        trial = TreeParzen.ask(expr, trials, CONFIG)
        @test 0.333 <= trial.hyperparams <= 3
    end

    # test scalar expressions real quick using identity
    trial = TreeParzen.ask(expr / 1)
    @test 0.333 <= trial.hyperparams <= 3
    trial = TreeParzen.ask(expr / 1, trials, CONFIG)
    @test 0.333 <= trial.hyperparams <= 3

end

@testset "Delayed binary operators -- modulus" begin

    choicevals = [1, 2, 3]

    a = TreeParzen.HP.Choice(:a, choicevals)
    b = TreeParzen.HP.Choice(:b, choicevals)

    expr = a % b

    trials = TreeParzen.Trials.Trial[]

    for i in 1:100
        trial = TreeParzen.ask(expr)
        TreeParzen.tell!(trials, trial, 1.)
        @test 0 <= trial.hyperparams <= 2
    end

    for i in 1:1000
        trial = TreeParzen.ask(expr, trials, CONFIG)
        @test 0 <= trial.hyperparams <= 2
    end

    # test scalar expressions real quick using identity (large modulus)
    trial = TreeParzen.ask(expr % 500)
    @test 0 <= trial.hyperparams <= 2
    trial = TreeParzen.ask(expr % 500, trials, CONFIG)
    @test 0 <= trial.hyperparams <= 2

end

@testset "Delayed binary operators -- power" begin

    choicevals = [1, 2, 3]

    a = TreeParzen.HP.Choice(:a, choicevals)
    b = TreeParzen.HP.Choice(:b, choicevals)

    expr = a ^ b

    trials = TreeParzen.Trials.Trial[]

    for i in 1:100
        trial = TreeParzen.ask(expr)
        TreeParzen.tell!(trials, trial, 1.)
        @test 1 <= trial.hyperparams <= 27
    end

    for i in 1:1000
        trial = TreeParzen.ask(expr, trials, CONFIG)
        @test 1 <= trial.hyperparams <= 27
    end

    # test scalar expressions real quick using identity
    trial = TreeParzen.ask(expr ^ 1)
    @test 1 <= trial.hyperparams <= 27
    trial = TreeParzen.ask(expr ^ 1, trials, CONFIG)
    @test 1 <= trial.hyperparams <= 27

end

@testset "Delayed unary operators -- subtraction" begin

    choicevals = [1, 2, 3]

    a = TreeParzen.HP.Choice(:a, choicevals)

    expr = -a

    trials = TreeParzen.Trials.Trial[]

    for i in 1:100
        trial = TreeParzen.ask(expr)
        TreeParzen.tell!(trials, trial, 1.)
        @test -3 <= trial.hyperparams <= -1
    end

    for i in 1:1000
        trial = TreeParzen.ask(expr, trials, CONFIG)
        @test -3 <= trial.hyperparams <= -1
    end

end

@testset "Delayed unary operators -- user-specified (ceil)" begin

    a = TreeParzen.HP.Uniform(:a, 0., 1.)
    expr = TreeParzen.Delayed.UnaryOperator(a, ceil)

    trials = TreeParzen.Trials.Trial[]

    for i in 1:100
        trial = TreeParzen.ask(expr)
        TreeParzen.tell!(trials, trial, 1.)
        @test trial.hyperparams == 1
    end

    for i in 1:1000
        trial = TreeParzen.ask(expr, trials, CONFIG)
        @test trial.hyperparams == 1
    end

end
@testset "Delayed unary operators -- round" begin

    a = TreeParzen.HP.Uniform(:a, 0.1, 1.9)
    expr = round(a)

    trials = TreeParzen.Trials.Trial[]

    for i in 1:100
        trial = TreeParzen.ask(expr)
        TreeParzen.tell!(trials, trial, 1.)
        @test isinteger(trial.hyperparams)  == true
    end

    for i in 1:1000
        trial = TreeParzen.ask(expr, trials, CONFIG)
        @test isinteger(trial.hyperparams)  == true
    end

end

end # module
true


text:
function show(io::IO, x::DoubleFloat{T}) where {T<:IEEEFloat}
    str = string(x)
    print(io, str)
end

show(io::IO, ::MIME"text/plain", x::DoubleFloat{T}) where {T<:IEEEFloat} = show(io, x)

function show(io::IO, x::DoubleFloat{DoubleFloat{T}}) where {T<:IEEEFloat}
    str = string(x)
    print(io, str)
end

function showtyped(io::IO, x::DoubleFloat{T}) where {T<:IEEEFloat}
    str = stringtyped(x)
    print(io, str)
end

showtyped(io::IO, ::MIME"text/plain", x::DoubleFloat{T}) where {T<:IEEEFloat} = showtyped(io, x)

function showtyped(io::IO, x::DoubleFloat{DoubleFloat{T}}) where {T<:IEEEFloat}
    str = stringtyped(x)
    print(io, str)
end

function showtyped(x::DoubleFloat{T}) where {T<:IEEEFloat}
    str = stringtyped(x)
    print(str)
end

function showtyped(x::DoubleFloat{DoubleFloat{T}}) where {T<:IEEEFloat}
    str = stringtyped(x)
    print(str)
end

ext: jl


text:
abstract type AbstractLoadBalancer end

# ------ worker helper functions
"""
    w_idx(wid::Integer)
    w_idx()

Return the index of a worker in the workers() array, given its process id.

If no wid is specified, then this returns the index of the calling process.
Note that this should only be called from the controller if the controller
is also a worker.
"""
w_idx(wid::Integer) = nprocs() > 1 ? wid - 1 : wid
w_idx() = w_idx(myid())

"""
    get_msg_chl(balancer<:AbstractLoadBalancer, wid::Integer)
    get_msg_chl(balancer<:AbstractLoadBalancer)

Return a reference to the message channel associated with worker wid.
"""
get_msg_chl(balancer::T, wid::Integer) where {T<:AbstractLoadBalancer} = balancer.msg_chls[w_idx(wid)]
get_msg_chl(balancer::T) where {T<:AbstractLoadBalancer} = get_msg_chl(balancer, myid())

"""
    create_msg_chls(capacity::Integer=64)

Create nworkers() remote channels for message passing.

# Arguments
- `capacity`: the capacity of the channels

"""
function create_msg_chls(capacity::Integer=64)
    msg_chls = [RemoteChannel(()->Channel{Message}(capacity), pid) for pid in workers()]
    msg_chls
end

# -- Simulate work

mutable struct WorkUnit{T<:Integer, S<:Real}
    units::T
    unitcost::S
    WorkUnit{T, S}(u, c) where {T<:Integer, S<:Real} = u < 0 || c ≤ 0 ? error("negative units or cost less than 0") : new(u, c)
end
WorkUnit(u::T, c::S) where {T<:Integer, S<:Real} = WorkUnit{T, S}(u, c)

"""
    split_work(work::WorkUnit, p<:AbstractFloat)
    split_work(work::WorkUnit)

Split a work into two WorkUnits, u and v, according to percentage p.

work must have at least 2 units. work will always be split in such a way to guarantee
that both u and v have at least 1 unit each. If `p` is not given, then a random value
0 < p < 1 will be used.

# Arguments
- `work`: An object of type WorkUnit to be split
- `p`: The percentage of units to give to the first split, u.

# Examples

```jldoctest
julia> w = WorkUnit(10, 3.)
WorkUnit{Int64,Float64}(10, 3.0)
julia> u, v = split_work(w, .6)
(WorkUnit{Int64,Float64}(6, 3.0), WorkUnit{Int64,Float64}(4, 3.0))
```
"""
function split_work(work::WorkUnit, p::T where T<:AbstractFloat)

    # check that valid percentage
    !(0 < p < 1) && (@printf("p must be: 0 < p < 1, got %g\n", p); throw(DomainError()))
    # check that there are at least 2 units
    work.units < 2 && (@printf("work must have more than 2 units\n"); throw(DomainError()))

    # set the rounding mode to up or down depending on p
    mode = p ≤ .5 ? RoundUp : RoundDown

    u_units = round(typeof(work.units), work.units * p, mode)
    v_units = work.units - u_units

    return WorkUnit(u_units, work.unitcost), WorkUnit(v_units, work.unitcost)

end
split_work(work::WorkUnit) = (p = rand(1:999); split_work(work, p/1000))

# ------ Message type for passing messages between processes and tasks
struct Message
    kind::Symbol
    data::Int64
    _data2 # TODO: rename data and _data2 to more sensible names
end

Message(kind, data) = Message(kind, data, -1)

"""
    recv_results(balancer<:AbstractLoadBalancer)

Receive and handle the results from workers passed to res_chl.
"""
function recv_results(balancer::T) where {T<:AbstractLoadBalancer}
    n_ended = 0
    res_count = 0
    total_work_done = zeros(Int64, nworkers())
    while n_ended < nworkers()
        work = take!(balancer.res_chl)
        if work.kind == :work
            res_count += 1
            w_idx = nprocs() > 1 ? work.data - 1 : work.data
            # add the work this worker did
            total_work_done[w_idx] += work._data2
            #info("Receiving results #", string(res_count), " from worker ", string(work.data))
        elseif work.kind == :end
            n_ended += 1
        end
    end

    for w_idx in 1:nworkers()
        info("Worker ", string(workers()[w_idx]), " did ", string(total_work_done[w_idx]), "s of work.")
    end
end

max_stars_repo_name: dillondaudert/jlance

max_issues_repo_name: dillondaudert/jlance

max_forks_repo_name: dillondaudert/jlance


text:
module VanDerPol
export van_der_pol

using Parameters: @with_kw, @unpack
using ..ExampleBase: LEDemo, ContinuousExample

@with_kw struct VanDerPolParam
    d::Float64 = 5.0
    a::Float64 = 5.0
    ω::Float64 = 2.466
end

@inline function phase_dynamics!(du, u, p, t)
    @unpack d, a, ω = p
    du[1] = u[2]
    du[2] = -u[1] - d * (u[1]^2 - 1) * u[2] + a * cos(ω * t)
end

@inline @views function tangent_dynamics!(du, u, p, t)
    @unpack d, a, ω = p
    phase_dynamics!(du[:, 1], u[:, 1], p, t)
    du[1, 2:end] .= u[2, 2:end]
    du[2, 2:end] .= -u[1, 2:end] .-
        d * (u[1, 1]^2 - 1) .* u[2, 2:end] .-
        d * 2.0 * u[1, 1] * u[2, 1] .* u[1, 2:end]
end

"""
Return a [`LEDemo`](@ref) for the van der Pol oscillator with
periodic forcing.

`.known_exponents` are extracted from Figure 6 of Geist, Parlitz &
Lauterborn (1990).

* <http://scholarpedia.org/article/Van_der_Pol_oscillator>
* <https://en.wikipedia.org/wiki/Van_der_Pol_oscillator>
* van der Pol and van der Mark. “Frequency Demultiplication.”
  Nature 120, no. 3019 (September 1927): 363.
  <https://doi.org/10.1038/120363a0>.
* Parlitz, Ulrich, and Werner Lauterborn. “Period-Doubling Cascades and
  Devil’s Staircases of the Driven van Der Pol Oscillator.”
  Physical Review A 36, no. 3 (August 1, 1987): 1428–34.
  <https://doi.org/10.1103/PhysRevA.36.1428>.
  (Figure 10a)
* Geist, K., Parlitz, U., & Lauterborn, W. (1990).
  Comparison of Different Methods for Computing Lyapunov Exponents.
  Progress of Theoretical Physics, 83, 875–893.
  <https://doi.org/10.1143/PTP.83.875>.
  (Figure 6)
"""
function van_der_pol(;
        params = [],
        u0=[0.1, 0.1],
        t_attr = 800,
        atol=0, rtol=2e-1,
        terminator_options = [:atol => atol, :rtol => rtol],
        kwargs...)
    # Note that with larger t_attr (e.g., 10000), the last Lyapunov
    # exponents negatively overshoots what Geist, Parlitz & Lauterborn
    # (1990) reported.  t_attr=100 is required for test to pass.
    if isempty(params)
        param = VanDerPolParam()
        known_exponents = [0.085, -6.7]
    else
        param = VanDerPolParam(; params...)
        known_exponents = Float64[]
    end
    t_renorm = 2 * π / param.ω
    LEDemo(ContinuousExample(
        "van der Pol & van der Mark (1927)",
        phase_dynamics!, u0, t_renorm, param,
        tangent_dynamics!,
        t_attr * t_renorm,
        known_exponents,
        atol, rtol,
        terminator_options,
    ); kwargs...)
end

end


text:
using .NumericalFluxes: CentralHyperDiffusiveFlux, CentralDivPenalty

struct DGModel{BL, G, NFND, NFD, GNF, AS, DS, HDS, D, DD, MD}
    balancelaw::BL
    grid::G
    numfluxnondiff::NFND
    numfluxdiff::NFD
    gradnumflux::GNF
    auxstate::AS
    diffstate::DS
    hyperdiffstate::HDS
    direction::D
    diffusion_direction::DD
    modeldata::MD
end
function DGModel(
    balancelaw,
    grid,
    numfluxnondiff,
    numfluxdiff,
    gradnumflux;
    auxstate = create_auxstate(balancelaw, grid),
    diffstate = create_diffstate(balancelaw, grid),
    hyperdiffstate = create_hyperdiffstate(balancelaw, grid),
    direction = EveryDirection(),
    diffusion_direction = direction,
    modeldata = nothing,
)
    DGModel(
        balancelaw,
        grid,
        numfluxnondiff,
        numfluxdiff,
        gradnumflux,
        auxstate,
        diffstate,
        hyperdiffstate,
        direction,
        diffusion_direction,
        modeldata,
    )
end

function (dg::DGModel)(dQdt, Q, ::Nothing, t; increment = false)

    bl = dg.balancelaw
    device = typeof(Q.data) <: Array ? CPU() : CUDA()

    grid = dg.grid
    topology = grid.topology

    dim = dimensionality(grid)
    N = polynomialorder(grid)
    Nq = N + 1
    Nqk = dim == 2 ? 1 : Nq
    Nfp = Nq * Nqk
    nrealelem = length(topology.realelems)

    Qvisc = dg.diffstate
    Qhypervisc_grad, Qhypervisc_div = dg.hyperdiffstate
    auxstate = dg.auxstate

    FT = eltype(Q)
    nviscstate = num_diffusive(bl, FT)
    nhyperviscstate = num_hyperdiffusive(bl, FT)

    Np = dofs_per_element(grid)

    workgroups_volume = (Nq, Nq, Nqk)
    ndrange_volume = (nrealelem * Nq, Nq, Nqk)
    workgroups_surface = Nfp
    ndrange_surface = Nfp * nrealelem

    communicate =
        !(isstacked(topology) && typeof(dg.direction) <: VerticalDirection)

    aux_comm = update_aux!(dg, bl, Q, t)
    @assert typeof(aux_comm) == Bool

    if nhyperviscstate > 0
        hypervisc_indexmap = create_hypervisc_indexmap(bl)
    else
        hypervisc_indexmap = nothing
    end

    ########################
    # Gradient Computation #
    ########################
    if communicate
        MPIStateArrays.start_ghost_exchange!(Q)
        if aux_comm
            MPIStateArrays.start_ghost_exchange!(auxstate)
        end
    end

    if nviscstate > 0 || nhyperviscstate > 0

        event = Event(device)
        event = volumeviscterms!(device, workgroups_volume)(
            bl,
            Val(dim),
            Val(N),
            dg.diffusion_direction,
            Q.data,
            Qvisc.data,
            Qhypervisc_grad.data,
            auxstate.data,
            grid.vgeo,
            t,
            grid.D,
            hypervisc_indexmap,
            topology.realelems,
            ndrange = ndrange_volume,
            dependencies = (event,),
        )
        wait(device, event)

        if communicate
            MPIStateArrays.finish_ghost_recv!(Q)
            if aux_comm
                MPIStateArrays.finish_ghost_recv!(auxstate)
            end
        end

        event = Event(device)
        event = faceviscterms!(device, workgroups_surface)(
            bl,
            Val(dim),
            Val(N),
            dg.diffusion_direction,
            dg.gradnumflux,
            Q.data,
            Qvisc.data,
            Qhypervisc_grad.data,
            auxstate.data,
            grid.vgeo,
            grid.sgeo,
            t,
            grid.vmap⁻,
            grid.vmap⁺,
            grid.elemtobndy,
            hypervisc_indexmap,
            topology.realelems;
            ndrange = ndrange_surface,
            dependencies = (event,),
        )
        wait(device, event)

        if communicate
            nviscstate > 0 && MPIStateArrays.start_ghost_exchange!(Qvisc)
            nhyperviscstate > 0 &&
            MPIStateArrays.start_ghost_exchange!(Qhypervisc_grad)
        end

        if nviscstate > 0
            aux_comm = update_aux_diffusive!(dg, bl, Q, t)
            @assert typeof(aux_comm) == Bool
        end

        if aux_comm
            MPIStateArrays.start_ghost_exchange!(auxstate)
        end
    end

    if nhyperviscstate > 0
        #########################
        # Laplacian Computation #
        #########################

        event = Event(device)
        event = volumedivgrad!(device, workgroups_volume)(
            bl,
            Val(dim),
            Val(N),
            dg.diffusion_direction,
            Qhypervisc_grad.data,
            Qhypervisc_div.data,
            grid.vgeo,
            grid.D,
            topology.realelems;
            ndrange = ndrange_volume,
            dependencies = (event,),
        )
        wait(device, event)

        communicate && MPIStateArrays.finish_ghost_recv!(Qhypervisc_grad)

        event = Event(device)
        event = facedivgrad!(device, workgroups_surface)(
            bl,
            Val(dim),
            Val(N),
            dg.diffusion_direction,
            CentralDivPenalty(),
            Qhypervisc_grad.data,
            Qhypervisc_div.data,
            grid.vgeo,
            grid.sgeo,
            grid.vmap⁻,
            grid.vmap⁺,
            grid.elemtobndy,
            topology.realelems;
            ndrange = ndrange_surface,
            dependencies = (event,),
        )
        wait(device, event)

        communicate && MPIStateArrays.start_ghost_exchange!(Qhypervisc_div)

        ####################################
        # Hyperdiffusive terms computation #
        ####################################

        event = Event(device)
        event = volumehyperviscterms!(device, workgroups_volume)(
            bl,
            Val(dim),
            Val(N),
            dg.diffusion_direction,
            Qhypervisc_grad.data,
            Qhypervisc_div.data,
            Q.data,
            auxstate.data,
            grid.vgeo,
            grid.ω,
            grid.D,
            topology.realelems,
            t;
            ndrange = ndrange_volume,
            dependencies = (event,),
        )
        wait(device, event)

        communicate && MPIStateArrays.finish_ghost_recv!(Qhypervisc_div)

        event = Event(device)
        event = facehyperviscterms!(device, workgroups_surface)(
            bl,
            Val(dim),
            Val(N),
            dg.diffusion_direction,
            CentralHyperDiffusiveFlux(),
            Qhypervisc_grad.data,
            Qhypervisc_div.data,
            Q.data,
            auxstate.data,
            grid.vgeo,
            grid.sgeo,
            grid.vmap⁻,
            grid.vmap⁺,
            grid.elemtobndy,
            topology.realelems,
            t;
            ndrange = ndrange_surface,
            dependencies = (event,),
        )
        wait(device, event)

        communicate && MPIStateArrays.start_ghost_exchange!(Qhypervisc_grad)
    end

    ###################
    # RHS Computation #
    ###################
    event = Event(device)
    event = volumerhs!(device, workgroups_volume)(
        bl,
        Val(dim),
        Val(N),
        dg.direction,
        dQdt.data,
        Q.data,
        Qvisc.data,
        Qhypervisc_grad.data,
        auxstate.data,
        grid.vgeo,
        t,
        grid.ω,
        grid.D,
        topology.realelems,
        increment;
        ndrange = ndrange_volume,
        dependencies = (event,),
    )
    wait(device, event)

    if communicate
        if nviscstate > 0 || nhyperviscstate > 0
            if nviscstate > 0
                MPIStateArrays.finish_ghost_recv!(Qvisc)
                if aux_comm
                    MPIStateArrays.finish_ghost_recv!(auxstate)
                end
            end
            nhyperviscstate > 0 &&
            MPIStateArrays.finish_ghost_recv!(Qhypervisc_grad)
        else
            MPIStateArrays.finish_ghost_recv!(Q)
            if aux_comm
                MPIStateArrays.finish_ghost_recv!(auxstate)
            end
        end
    end

    event = Event(device)
    event = facerhs!(device, workgroups_surface)(
        bl,
        Val(dim),
        Val(N),
        dg.direction,
        dg.numfluxnondiff,
        dg.numfluxdiff,
        dQdt.data,
        Q.data,
        Qvisc.data,
        Qhypervisc_grad.data,
        auxstate.data,
        grid.vgeo,
        grid.sgeo,
        t,
        grid.vmap⁻,
        grid.vmap⁺,
        grid.elemtobndy,
        topology.realelems;
        ndrange = ndrange_surface,
        dependencies = (event,),
    )
    wait(device, event)

    # Just to be safe, we wait on the sends we started.
    if communicate
        MPIStateArrays.finish_ghost_send!(Qhypervisc_div)
        MPIStateArrays.finish_ghost_send!(Qvisc)
        MPIStateArrays.finish_ghost_send!(Qhypervisc_grad)
        MPIStateArrays.finish_ghost_send!(Q)
    end
end

function init_ode_state(
    dg::DGModel,
    args...;
    init_on_cpu = false,
    commtag = 888,
)
    device = arraytype(dg.grid) <: Array ? CPU() : CUDA()

    bl = dg.balancelaw
    grid = dg.grid

    state = create_state(bl, grid, commtag)

    topology = grid.topology
    Np = dofs_per_element(grid)

    auxstate = dg.auxstate
    dim = dimensionality(grid)
    N = polynomialorder(grid)
    nrealelem = length(topology.realelems)

    if !init_on_cpu
        event = Event(device)
        event = initstate!(device, Np)(
            bl,
            Val(dim),
            Val(N),
            state.data,
            auxstate.data,
            grid.vgeo,
            topology.realelems,
            args...;
            ndrange = Np * nrealelem,
            dependencies = (event,),
        )
        wait(device, event)
    else
        h_state = similar(state, Array)
        h_auxstate = similar(auxstate, Array)
        h_auxstate .= auxstate
        event = initstate!(CPU(), Np)(
            bl,
            Val(dim),
            Val(N),
            h_state.data,
            h_auxstate.data,
            Array(grid.vgeo),
            topology.realelems,
            args...;
            ndrange = Np * nrealelem,
        )
        wait(event) # XXX: This could be `wait(device, event)` once KA supports that.
        state .= h_state
    end

    MPIStateArrays.start_ghost_exchange!(state)
    MPIStateArrays.finish_ghost_exchange!(state)

    return state
end

# fallback
function update_aux!(dg::DGModel, bl::BalanceLaw, Q::MPIStateArray, t::Real)
    return false
end

function update_aux_diffusive!(
    dg::DGModel,
    bl::BalanceLaw,
    Q::MPIStateArray,
    t::Real,
)
    return false
end

function indefinite_stack_integral!(
    dg::DGModel,
    m::BalanceLaw,
    Q::MPIStateArray,
    auxstate::MPIStateArray,
    t::Real,
)

    device = typeof(Q.data) <: Array ? CPU() : CUDA()

    grid = dg.grid
    topology = grid.topology

    dim = dimensionality(grid)
    N = polynomialorder(grid)
    Nq = N + 1
    Nqk = dim == 2 ? 1 : Nq

    FT = eltype(Q)

    # do integrals
    nelem = length(topology.elems)
    nvertelem = topology.stacksize
    nhorzelem = div(nelem, nvertelem)

    event = Event(device)
    event = knl_indefinite_stack_integral!(device, (Nq, Nqk))(
        m,
        Val(dim),
        Val(N),
        Val(nvertelem),
        Q.data,
        auxstate.data,
        grid.vgeo,
        grid.Imat,
        1:nhorzelem;
        ndrange = (nhorzelem * Nq, Nqk),
        dependencies = (event,),
    )
    wait(device, event)
end

function reverse_indefinite_stack_integral!(
    dg::DGModel,
    m::BalanceLaw,
    Q::MPIStateArray,
    auxstate::MPIStateArray,
    t::Real,
)

    device = typeof(auxstate.data) <: Array ? CPU() : CUDA()

    grid = dg.grid
    topology = grid.topology

    dim = dimensionality(grid)
    N = polynomialorder(grid)
    Nq = N + 1
    Nqk = dim == 2 ? 1 : Nq

    FT = eltype(auxstate)

    # do integrals
    nelem = length(topology.elems)
    nvertelem = topology.stacksize
    nhorzelem = div(nelem, nvertelem)

    event = Event(device)
    event = knl_reverse_indefinite_stack_integral!(device, (Nq, Nqk))(
        m,
        Val(dim),
        Val(N),
        Val(nvertelem),
        Q.data,
        auxstate.data,
        1:nhorzelem;
        ndrange = (nhorzelem * Nq, Nqk),
        dependencies = (event,),
    )
    wait(device, event)
end

function nodal_update_aux!(
    f!,
    dg::DGModel,
    m::BalanceLaw,
    Q::MPIStateArray,
    t::Real;
    diffusive = false,
)
    device = typeof(Q.data) <: Array ? CPU() : CUDA()

    grid = dg.grid
    topology = grid.topology

    dim = dimensionality(grid)
    N = polynomialorder(grid)
    Nq = N + 1
    nrealelem = length(topology.realelems)

    Np = dofs_per_element(grid)

    nodal_update_aux! = knl_nodal_update_aux!(device, Np)
    ### update aux variables
    event = Event(device)
    if diffusive
        event = nodal_update_aux!(
            m,
            Val(dim),
            Val(N),
            f!,
            Q.data,
            dg.auxstate.data,
            dg.diffstate.data,
            t,
            topology.realelems;
            ndrange = Np * nrealelem,
            dependencies = (event,),
        )
    else
        event = nodal_update_aux!(
            m,
            Val(dim),
            Val(N),
            f!,
            Q.data,
            dg.auxstate.data,
            t,
            topology.realelems;
            ndrange = Np * nrealelem,
            dependencies = (event,),
        )
    end
    wait(device, event)
end

"""
    courant(local_courant::Function, dg::DGModel, m::BalanceLaw,
            Q::MPIStateArray, direction=EveryDirection())
Returns the maximum of the evaluation of the function `local_courant`
pointwise throughout the domain.  The function `local_courant` is given an
approximation of the local node distance `Δx`.  The `direction` controls which
reference directions are considered when computing the minimum node distance
`Δx`.
An example `local_courant` function is
    function local_courant(m::AtmosModel, state::Vars, aux::Vars,
                           diffusive::Vars, Δx)
      return Δt * cmax / Δx
    end
where `Δt` is the time step size and `cmax` is the maximum flow speed in the
model.
"""
function courant(
    local_courant::Function,
    dg::DGModel,
    m::BalanceLaw,
    Q::MPIStateArray,
    Δt,
    simtime,
    direction = EveryDirection(),
)
    grid = dg.grid
    topology = grid.topology
    nrealelem = length(topology.realelems)

    if nrealelem > 0
        N = polynomialorder(grid)
        dim = dimensionality(grid)
        Nq = N + 1
        Nqk = dim == 2 ? 1 : Nq
        device = grid.vgeo isa Array ? CPU() : CUDA()
        pointwise_courant = similar(grid.vgeo, Nq^dim, nrealelem)
        event = Event(device)
        event = Grids.knl_min_neighbor_distance!(device, (Nq, Nq, Nqk))(
            Val(N),
            Val(dim),
            direction,
            pointwise_courant,
            grid.vgeo,
            topology.realelems;
            ndrange = (nrealelem * Nq, Nq, Nqk),
            dependencies = (event,),
        )
        event = knl_local_courant!(device, Nq * Nq * Nqk)(
            m,
            Val(dim),
            Val(N),
            pointwise_courant,
            local_courant,
            Q.data,
            dg.auxstate.data,
            dg.diffstate.data,
            topology.realelems,
            Δt,
            simtime,
            direction;
            ndrange = nrealelem * Nq * Nq * Nqk,
            dependencies = (event,),
        )
        wait(device, event)
        rank_courant_max = maximum(pointwise_courant)
    else
        rank_courant_max = typemin(eltype(Q))
    end

    MPI.Allreduce(rank_courant_max, max, topology.mpicomm)
end

function copy_stack_field_down!(
    dg::DGModel,
    m::BalanceLaw,
    auxstate::MPIStateArray,
    fldin,
    fldout,
)

    device = typeof(auxstate.data) <: Array ? CPU() : CUDA()

    grid = dg.grid
    topology = grid.topology

    dim = dimensionality(grid)
    N = polynomialorder(grid)
    Nq = N + 1
    Nqk = dim == 2 ? 1 : Nq

    # do integrals
    nelem = length(topology.elems)
    nvertelem = topology.stacksize
    nhorzelem = div(nelem, nvertelem)

    event = Event(device)
    event = knl_copy_stack_field_down!(device, (Nq, Nqk))(
        Val(dim),
        Val(N),
        Val(nvertelem),
        auxstate.data,
        1:nhorzelem,
        Val(fldin),
        Val(fldout);
        ndrange = (nhorzelem * Nq, Nqk),
        dependencies = (event,),
    )
    wait(device, event)
end

function MPIStateArrays.MPIStateArray(dg::DGModel, commtag = 888)
    bl = dg.balancelaw
    grid = dg.grid

    state = create_state(bl, grid, commtag)

    return state
end

function create_hypervisc_indexmap(bl::BalanceLaw)
    # helper function
    _getvars(v, ::Type) = v
    function _getvars(v::Vars, ::Type{T}) where {T <: NamedTuple}
        fields = getproperty.(Ref(v), fieldnames(T))
        collect(Iterators.Flatten(_getvars.(fields, fieldtypes(T))))
    end

    gradvars = vars_gradient(bl, Int)
    gradlapvars = vars_gradient_laplacian(bl, Int)
    indices = Vars{gradvars}(1:varsize(gradvars))
    SVector{varsize(gradlapvars)}(_getvars(indices, gradlapvars))
end

max_stars_repo_name: trontrytel/CLIMA

max_issues_repo_name: trontrytel/CLIMA

max_forks_repo_name: trontrytel/CLIMA


text:
@contract begin
    require(b > 0)
    ensure(toReturn >= 1)
    function flawedReturnDiv(a::Int64, b::Int64)
        return a / b
    end
end

@contract begin
    require(b > 0)
    ensure(toReturn >= 1)
    returnName = toReturn
    function fixedDiv(a::Int64, b::Int64)
        return a / b
    end
end

@testset "Return name change" begin
    @test_throws ContractBreachException flawedReturnDiv(1, 0)
    @test_throws UndefVarError flawedReturnDiv(1, 2)
    @test_throws UndefVarError flawedReturnDiv(10, 2)

    @test_throws ContractBreachException fixedDiv(1, 0)
    @test_throws ContractBreachException fixedDiv(1, 2)
    @test fixedDiv(10, 2) == 5
end

wrongReturnNameInt = quote
    @contract begin
        returnName = 2
        function stub()
            return 2
        end
    end
end

wrongReturnNameBool = quote
    @contract begin
        returnName = true
        function stub()
            return 2
        end
    end
end

wrongReturnNameExpr = quote
    @contract begin
        returnName = x -> x + 2
        function stub()
            return 2
        end
    end
end

@testset "new ReturnName Type" begin
    try
        eval(wrongReturnNameInt)
    catch e
        @test typeof(e.error) == ArgumentError
    end
    try
        eval(wrongReturnNameBool)
    catch e
        @test typeof(e.error) == ArgumentError
    end
    try
        eval(wrongReturnNameExpr)
    catch e
        @test typeof(e.error) == ArgumentError
    end
end

@contract begin
    ensure(result >= 0)
    function wrongNestedValue(; x = 0, y = 0)
        if x == 0
            return 0
        elseif x == 1
            if y == 0
                return 0
            elseif y == 1
                return -1
            end
        end
        return -2
    end
end

@testset "Nested returns" begin
    @test_throws ContractBreachException wrongNestedValue(x = 1, y = 1)
    @test_throws ContractBreachException wrongNestedValue(x = 2)
    @test wrongNestedValue() == 0
    @test wrongNestedValue(x = 1) == 0
end


text:
# Original implementations by @nantonel
# https://github.com/maxmouchet/HMMBase.jl/pull/6

function viterbilog!(
    T1::AbstractMatrix,
    T2::AbstractMatrix,
    z::AbstractVector,
    a::AbstractVector,
    A::AbstractMatrix,
    LL::AbstractMatrix,
)
    T, K = size(LL)
    (T == 0) && return

    fill!(T1, 0.0)
    fill!(T2, 0)

    al = log.(a)
    Al = log.(A)

    for i in OneTo(K)
        T1[1, i] = al[i] + LL[1, i]
    end

    @inbounds for t = 2:T
        for j in OneTo(K)
            amax = 0
            vmax = -Inf

            for i in OneTo(K)
                v = T1[t-1, i] + Al[i, j]
                if v > vmax
                    amax = i
                    vmax = v
                end
            end

            T1[t, j] = vmax + LL[t, j]
            T2[t, j] = amax
        end
    end

    z[T] = argmax(T1[T, :])
    for t = T-1:-1:1
        z[t] = T2[t+1, z[t+1]]
    end
end

"""
    viterbi(a, A, LL) -> Vector

Find the most likely hidden state sequence, see [Viterbi algorithm](https://en.wikipedia.org/wiki/Viterbi_algorithm).
"""
function viterbi(a::AbstractVector, A::AbstractMatrix, LL::AbstractMatrix; logl = nothing)
    ## < v1.1 compatibility
    (logl !== nothing) && deprecate_kwargs("logl")
    (logl == false) && (LL = log.(LL))
    ## --------------------
    T1 = Matrix{Float64}(undef, size(LL))
    T2 = Matrix{Int}(undef, size(LL))
    z = Vector{Int}(undef, size(LL, 1))
    viterbilog!(T1, T2, z, a, A, LL)
    z
end

"""
    viterbi(hmm, observations; robust) -> Vector

**Example**
```julia
using Distributions, HMMBase
hmm = HMM([0.9 0.1; 0.1 0.9], [Normal(0,1), Normal(10,1)])
y = rand(hmm, 1000)
zv = viterbi(hmm, y)
```
"""
function viterbi(hmm::AbstractHMM, observations; logl = nothing, robust = false)
    (logl !== nothing) && deprecate_kwargs("logl")
    LL = loglikelihoods(hmm, observations; robust = robust)
    viterbi(hmm.a, hmm.A, LL)
end


text:
# basic_spaces.jl

##############################
# Definitions of basic spaces
##############################

"A geometric space with integer type."
const IntegerSpace{T <: Integer} = GeometricSpace{T}

"A geometric space with rational types."
const RationalSpace{T} = GeometricSpace{Rational{T}}

subeltype(::Type{RationalSpace{T}}) where {T} = T
similar_space(::Type{RationalSpace{T}}, ::Type{S}) where {T,S} = RationalSpace{S}

"A geometric space with floating point type."
const RealSpace{T <: AbstractFloat} = GeometricSpace{T}

"A geometric space with complex type."
const ComplexSpace{T} = GeometricSpace{Complex{T}}

subeltype(::Type{ComplexSpace{T}}) where {T} = T
similar_space(::Type{ComplexSpace{T}}, ::Type{S}) where {T,S} = ComplexSpace{S}

"""
A Euclidean space with static vector types. It is the space of all vectors
of fixed length `N`, with entries of type `T`.
"""
const VectorSpace{N,T} = GeometricSpace{SVector{N,T}}
const EuclideanSpace{N,T} = VectorSpace{N,T}

subeltype(::Type{VectorSpace{N,T}}) where {N,T} = T
similar_space(::Type{VectorSpace{N,T}}, ::Type{S}) where {N,T,S} = VectorSpace{N,S}

"""
A general Array space, with arrays of dimension `N` and element type `T`.
Note that the arrays can have any size. Thus, `ArraySpace{2,Float64}` contains
all possible matrices of size `m × n` for any combination of `m,n ∈ ℕ`.
"""
const ArraySpace{N,T} = GeometricSpace{Array{T,N}}

subeltype(::Type{ArraySpace{N,T}}) where {N,T} = T
similar_space(::Type{ArraySpace{N,T}}, ::Type{S}) where {N,T,S} = ArraySpace{N,S}

## Convenience functions

^(::Type{GeometricSpace{T}}, ::Type{Val{N}}) where {T <: Number,N} = VectorSpace{N,T}

## Some standard spaces

"The set of integers of type Int (ℤ = \\BbbZ)."
const ℤ = IntegerSpace{Int}
"The set of rational numbers of type Rational{Int} (ℚ = \\BbbQ)."
const ℚ = RationalSpace{Int}
"The set of reals of type Float64 (ℝ = \\BbbR)."
const ℝ = RealSpace{Float64}
"The complex plane with Float64 real and imaginar parts (ℂ = \\BbbC)."
const ℂ = ComplexSpace{Float64}

"The space ℝ^1"
const ℝ1 = ℝ^Val{1}
"The space ℝ^2"
const ℝ2 = ℝ^Val{2}
"The space ℝ^3"
const ℝ3 = ℝ^Val{3}
"The space ℝ^4"
const ℝ4 = ℝ^Val{4}

##############################
# Embeddings and conversions
##############################

## Isomorphism between scalars and 1D vectors

isomorphism_reduction(::Type{GeometricSpace{T}}, ::Type{VectorSpace{1,S}}) where {T <: Number,S} =
    (GeometricSpace{T},GeometricSpace{S})

convert_space(::Type{GeometricSpace{T}}, x::SVector{1,T}) where {T} = x[1]
convert_space(::Type{VectorSpace{1,T}}, x::T) where {T} = SVector(x)

## Isomorphism between the complex plane and ℝ^2

isomorphism_reduction(::Type{VectorSpace{2,T}}, ::Type{ComplexSpace{S}}) where {T,S} =
    (GeometricSpace{T},GeometricSpace{S})

convert_space(::Type{ComplexSpace{T}}, x::SVector{2,T}) where {T} = x[1] + im*x[2]
convert_space(::Type{VectorSpace{2,T}}, x::Complex{T}) where {T} = SVector(real(x), imag(x))

## Isomorphism between vector spaces of the same dimension

# Normally, this line never gets called when T equals S, because that case is
# caught higher up in the algorithms.
isomorphism_reduction(::Type{VectorSpace{N,T}}, ::Type{VectorSpace{N,S}}) where {N,T,S} =
    (GeometricSpace{T},GeometricSpace{S})

# Just in case, we make sure if it ever happens that we notice
isomorphism_reduction(::Type{VectorSpace{N,T}}, ::Type{VectorSpace{N,T}}) where {N,T} =
    error("isomorphism_reduction was called with identical types")

convert_space(A::Type{VectorSpace{N,T}}, x::SVector{N,T}) where {N,T} = x
convert_space(A::Type{VectorSpace{N,T}}, x::SVector{N,S}) where {N,T,S} =
    SVector{N,T}(convert_space.(GeometricSpace{T}, x))

restrict_space(A::Type{VectorSpace{N,T}}, x::SVector{N,T}) where {N,T} = x
restrict_space(A::Type{VectorSpace{N,T}}, x::SVector{N,S}) where {N,T,S} =
    SVector{N,T}(restrict_space.(GeometricSpace{T}, x))

## For completeness, we have an isomorphism rule for ComplexSpace and RationalSpace
## to themselves as well. These are probably not useful.

isomorphism_reduction(::Type{ComplexSpace{T}}, ::Type{ComplexSpace{S}}) where {T,S} =
    (GeometricSpace{T},GeometricSpace{S})

convert_space(A::Type{ComplexSpace{T}}, x::Complex{T}) where {T} = x
convert_space(A::Type{ComplexSpace{T}}, x::Complex{S}) where {T,S} =
    convert_space(GeometricSpace{T}, real(x)) + im * convert_space(GeometricSpace{T}, imag(x))

restrict_space(A::Type{ComplexSpace{T}}, x::Complex{T}) where {T} = x
restrict_space(A::Type{ComplexSpace{T}}, x::Complex{S}) where {T,S} =
    restrict_space(GeometricSpace{T}, real(x)) + im * restrict_space(GeometricSpace{T}, imag(x))

isomorphism_reduction(::Type{RationalSpace{T}}, ::Type{RationalSpace{S}}) where {T,S} =
    (GeometricSpace{T},GeometricSpace{S})

convert_space(A::Type{RationalSpace{T}}, x::Rational{T}) where {T} = x
convert_space(A::Type{RationalSpace{T}}, x::Rational{S}) where {T,S} =
    convert_space(GeometricSpace{T}, numerator(x)) // convert_space(GeometricSpace{T}, denominator(x))

restrict_space(A::Type{RationalSpace{T}}, x::Rational{T}) where {T} = x
restrict_space(A::Type{RationalSpace{T}}, x::Rational{S}) where {T,S} =
    restrict_space(GeometricSpace{T}, numerator(x)) // restrict_space(GeometricSpace{T}, denominator(x))

## Embedding of R^N in R^(N+1)

# We only go up to dimension four
embedding_reduction(::Type{VectorSpace{1,T}}, ::Type{VectorSpace{2,S}}) where {T,S} =
    (GeometricSpace{T},GeometricSpace{S})
embedding_reduction(::Type{VectorSpace{2,T}}, ::Type{VectorSpace{3,S}}) where {T,S} =
    (GeometricSpace{T},GeometricSpace{S})
embedding_reduction(::Type{VectorSpace{3,T}}, ::Type{VectorSpace{4,S}}) where {T,S} =
    (GeometricSpace{T},GeometricSpace{S})

convert_space(::Type{VectorSpace{2,T}}, x::SVector{1,T}) where {T} = SVector(x[1], 0)
convert_space(::Type{VectorSpace{3,T}}, x::SVector{2,T}) where {T} = SVector(x[1], x[2], 0)
convert_space(::Type{VectorSpace{4,T}}, x::SVector{3,T}) where {T} = SVector(x[1], x[2], x[3], 0)

# Since these are true embeddings, we also provide a left inverse
restrict_space(::Type{VectorSpace{1,T}}, x::SVector{2,T}) where {T} = SVector(x[1])
restrict_space(::Type{VectorSpace{2,T}}, x::SVector{3,T}) where {T} = SVector(x[1], x[2])
restrict_space(::Type{VectorSpace{3,T}}, x::SVector{4,T}) where {T} = SVector(x[1], x[2], x[3])

# Superspaces of a vector space have one higher dimension
superspace(::Type{VectorSpace{1,T}}) where {T} = VectorSpace{2,T}
superspace(::Type{VectorSpace{2,T}}) where {T} = VectorSpace{3,T}
superspace(::Type{VectorSpace{3,T}}) where {T} = VectorSpace{4,T}


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule LoopTools_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("LoopTools")
JLLWrappers.@generate_main_file("LoopTools", UUID("e275ce21-2bb6-5a50-939e-e917944b1aed"))
end  # module LoopTools_jll


text:
@testset "UI" begin
@testset "ArrayOp" begin
    ctx = Context()
    log = Dagger.LocalEventLog()
    ctx.log_sink = log

    sz = 2^4
    bsz = 2^2
    X = rand(Float32, sz, sz)
    XD = Distribute(Blocks(bsz, bsz), X)

    dag = XD*XD
    collect(ctx, dag)
    logs = Dagger.get_logs!(log)
    @test logs isa Vector{Dagger.Timespan}
    @test length(logs) > 0
    plan = Dagger.show_plan(logs)
    @test plan isa String
    @test occursin("digraph {", plan)
    @test occursin("Move:", plan)
    @test endswith(plan, "}\n")
end

@testset "Argument Merging" begin
    ctx = Context()
    log = Dagger.LocalEventLog()
    ctx.log_sink = log

    X = rand(Float32, 4, 3)
    a = delayed(sum)(X)
    b = delayed(sum)(X)
    c = delayed(+)(a,b)

    X1 = rand(Float32, 5, 10)
    X2 = rand(Float32, 5, 10)
    f = delayed(sum)(X1)
    g = delayed(sum)(X2)
    h = delayed(+)(f,g)

    j = delayed(+)(c,h)
    collect(ctx, j)
    logs = Dagger.get_logs!(log)
    plan = Dagger.show_plan(logs, j)
end

@testset "Automatic Plan Rendering" begin
    x = compute(rand(Blocks(2,2),4,4))
    mktemp() do path, io
        ctx = Context(;log_sink=Dagger.LocalEventLog(),log_file=path)
        compute(ctx, x * x)
        plan = String(read(io))
        @test occursin("digraph {", plan)
        @test occursin("Move:", plan)
        @test endswith(plan, "}\n")
    end
end
end


text:
# FIXME: this is grossly inadequate

@testset "refine" begin
    Thigh = Double64
    Tlow = Float64
    # make a moderately unfriendly matrix
    n = 128
    B = rand(Tlow,n,n)
    SF = svd(B)
    acond = 1/sqrt(eps(Tlow))
    A = (1/acond) * SF.U * diagm(0 => exp.((0:n-1)*log(acond)/(n-1))) * SF.Vt

    Ad = Thigh.(A)
    x = rand(Thigh,n)
    bd = Ad * x
    b = Tlow.(bd)
    Flow = lu(A)
    xlow = Flow \ b
    # make sure we didn't sneak through
    # println("low-prec error: ",norm(x-xlow) / norm(x))
    @assert norm(x-xlow) / norm(x) > 1e4 * n * acond * eps(Thigh)
    xhigh,cvg = refinedldiv(Ad, Flow, bd)
    # println("high-prec error: ",norm(x-xhigh) / norm(x))
    @test norm(x-xhigh) / norm(x) < n * acond * eps(Thigh)
end

max_stars_repo_name: RalphAS/DoubleBLAS


text:
"""
    gen_checkboard(side)

Generate cube array of ones and zeros with checkboard-like
pattern. The size of checkboard cell is 2.
"""
function gen_checkboard(side)
    checkboard = Array{Int, 3}(undef, side, side, side)
    for i in 1:side
        for j in 1:side
            for k in 1:side
                x = ((i-1) & 2 != 0) ? 1 : 0
                y = ((j-1) & 2 != 0) ? 1 : 0
                z = ((k-1) & 2 != 0) ? 1 : 0
                checkboard[k,j,i] = x ⊻ y ⊻ z
            end
        end
    end
    return checkboard
end

cb = gen_checkboard(60)

@testset "L2" begin
    for phase in 0:1
        l = mean(Directional.l2(cb, phase; periodic = true))
        @test l[1] ≈ 1/2
        @test l[2] ≈ 1/4
        @test all(x -> x == 0, l[3:end])
    end
end

# C2 behaves exactly as L2 with segmentation algorithm chosen which
# does not join segments in diagonal directions.
@testset "C2" begin
    for phase in 0:1
        c = mean(Directional.c2(cb, phase; periodic = true))
        @test c[1] ≈ 1/2
        @test c[2] ≈ 1/4
        @test all(x -> x == 0, c[3:end])
    end
end

@testset "S2" begin
    for phase in 0:1
        s = mean(Directional.s2(cb, phase; periodic = true))
        @test all(x -> x ≈ 1/4, s[2:2:end])
        @test all(x -> isapprox(x, 0; atol = 1e-2), s[3:4:end])
        @test all(x -> x ≈ 1/2, s[1:4:end])
    end
end


text:
import Statistics
import StatsBase

function _sortperm(A)
          tmp = [(@inbounds A[i],i) for i=1:length(A)]
          sort!(tmp; alg=QuickSort)
          [t[2] for t in tmp]
end

"""
compute_possible_splits

"""
function compute_possible_splits(
    x                   :: Array{Float64,1},
    min_leaf_size       :: Int
    )

    x = sort(x)
    n = length(x)
    splits = Vector{Float64}(undef,(n-1))
    @inbounds for i in min_leaf_size:(n-min_leaf_size)
        splits[i] = (x[i] + x[i+1])/2.0
    end
    return(splits)
end

"""
evaluate_possible_splits

"""
function evaluate_possible_splits(
    x                   :: Array{Float64,1},
    y                   :: Array{Float64,1},
    min_leaf_size       :: Int
    )

    oj = sortperm(x)
    yj = y[oj]
    n = length(x)
    # first split
    split = min_leaf_size  # actually (x[1] + x[2])/2  x["1.5"]
    yl_ = Statistics.mean(yj[1:min_leaf_size])
    yr_ = Statistics.mean(yj[(min_leaf_size+1):end])
    best_loss = min_leaf_size*yl_^2 + (n-min_leaf_size)*yr_^2 # update to function calling
    @inbounds for i in (min_leaf_size+1):(n-min_leaf_size+1)
        yl_ = update_average(yl_,yj[i],i)
        yr_ = update_average(yr_,yj[i],n-i,true)
        loss_i = i*yl_^2 + (n-i)*yr_^2
        if loss_i > best_loss
            best_loss = loss_i
            split = i
        end
    end
    # compute threshold
    threshold = (x[oj][split] + x[oj][split+1])/2.0
    # return
    return(split,threshold,best_loss)
end

"""
Update an average based on a newly observed value

"""
function update_average(
    current             :: Float64,
    val                 :: Float64,
    n                   :: Int64,
    remove              :: Bool=false
    )

    if remove == false
        update = (n/(n+1))*current + (1/(n+1))*val
    elseif remove == true
        update = (n/(n-1))*current - (1/(n-1))*val
    end
    return(update)
end

"""
learn a split given two datasets X and y

Arguments

node: the node to split
X: feature matrix
y: response
loss: loss function
min_leaf_size: minimum number of unique values in a leaf

Note: aim is to maximise the loss function
"""
function learn_split!(
    node                :: Node,
    X                   :: Array{Float64,2},
    y                   :: Array{Float64,1},
    loss                :: Function,
    min_leaf_size       :: Int64,
    min_loss_increase   :: Float64,
    max_depth           :: Int64
    )

    region = node.region
    n_samples = length(region)
    n_features = size(X)[2]

    # split parameters
    feature = 0
    threshold = NaN

    # X and y values in partition region
    Xp = X[region,:]
    yp = y[region]

    # prediction and loss in partition region
    mu_ = Statistics.mean(yp)
    current_loss = n_samples*mu_^2

    #println(n_samples," ",min_leaf_size, node.depth < max_depth)

    # test whether maximum depth reached
    if ((node.depth < max_depth) & (n_samples > min_leaf_size))
        # for each column evaluate all possible split points
        # outer loop: loop over columns
        @inbounds for j in 1:n_features
            xpj = Xp[:,j]
            #spj = compute_possible_splits(xpj,min_leaf_size)  # no!
            split,prop_threshold,prop_loss = evaluate_possible_splits(xpj,yp,min_leaf_size)

            if prop_loss > current_loss
                feature = j
                current_loss = prop_loss
                threshold = prop_threshold
                global region_l = region[xpj .< threshold]
                global region_r = region[xpj .> threshold]
            end
        end
    end

    if feature != 0
        # modify the node
        node.value = mu_  # placeholder for pruning
        node.is_leaf = false
        node.feature = feature
        node.threshold = threshold
        node.l = Node(region_l,node.depth+1)
        node.r = Node(region_r,node.depth+1)
    else
        # modify the node
        node.value = mu_
        node.is_leaf = true
    end
end

"""
learn_tree :
learn a decision tree

X: feature matrix
y: response
loss: loss function
min_leaf_size: minimum number of unique values in a leaf
max_depth: maximum tree depth

Note: aim is to maximise the loss function
"""
function learn_tree(
    X                   :: Array{Float64,2},
    y                   :: Array{Float64,1},
    loss                :: Function,
    min_leaf_size       :: Int64,
    min_loss_increase   :: Float64,
    max_depth           :: Int64
    )

    n_samples, n_features = size(X)
    indX = collect(1:n_samples)

    # root node
    root = Node(indX,0)
    stack = Node[root] # here tree.Node indicates array type
    ids = Array{Int64,1}()

    while length(stack) > 0
        node = pop!(stack)
        node.id = node_id(ids)
        append!(ids,node.id)
        learn_split!(node,X,y,loss,min_leaf_size,min_loss_increase,max_depth)  # in place modification of node.l, node.r
        if !node.is_leaf
            push!(stack,node.r)
            push!(stack,node.l)
        end
    end
    return(root)
end

"""
estimate_leaf_value :

"""
function learn_leaf_value(
    tree                :: Node,
    X                   :: Array{Float64,2},
    y                   :: Array{Float64,1}
    )

    n_samples, n_features = size(X)
    id_vec = match_to_id(tree,X)

    id_to_val = Dict()  # add types

    for id in unique(id_vec)
        yp = y[id .== id_vec]
        Xp = X[id .== id_vec,:]
        X_id = Statistics.mean(Xp,dims=1)
        id_to_val[id] = [Statistics.mean(yp),X_id]
    end

    for id in unique(id_vec)
        node = tree
        while node.is_leaf == false
            if id_to_val[id][2][node.feature] < node.threshold
                node = node.l
            else
                node = node.r
            end
        end
        node.value = id_to_val[id][1]
    end
    return(tree)
end

"""
fit a decision tree

Arguments

structure_p : the proportion of the data used in estimating the tree structure (i.e. nodes)
"""
function decision_tree(
    X                   :: Array{Float64,2},
    y                   :: Array{Float64,1},
    loss                :: Function;
    min_leaf_size       :: Int64=10,
    min_loss_increase   :: Float64=0.0,
    max_depth           :: Int64=3,
    honesty             :: Bool=false,
    structure_p         :: Float64=0.5
    )

    # dataset dimensions
    n_samples, n_features = size(X)

    if (honesty == true)
        indX = collect(1:n_samples)

        # indexes for the structural and leaf parts of the tree
        n_struct = convert(Int64,floor(structure_p*n_samples))
        indS = StatsBase.sample(indX,n_struct,replace=false)
        indH = setdiff(indX,indS)

        # structure stage
        XS = X[indS,:]
        yS = y[indS]
        tree = learn_tree(XS,yS,loss,min_leaf_size,
                          min_loss_increase,max_depth)

        # estimate leaf values
        XH = X[indH,:]
        yH = y[indH]
        tree = learn_leaf_value(tree,XH,yH)
    elseif (honesty == false)
        tree = learn_tree(X,y,loss,min_leaf_size,
                          min_loss_increase,max_depth)
    end

    # remove region from tree
    recursive_action(tree,remove_region!)
    # return as type Tree
    return(Tree(tree,n_features))

end

"""
causal_forest : learn a causal forest

"""
function random_forest(
    X                   :: Array{Float64,2},
    y                   :: Array{Float64,1},
    loss                :: Function;
    n_trees             :: Int64=100,
    col_subsamp_p       :: Float64=1.0,
    row_subsamp_p       :: Float64=1.0,
    row_resample        :: Bool=true,
    min_leaf_size       :: Int64=10,
    min_loss_increase   :: Float64=0.0,
    max_depth           :: Int64=3,
    honesty             :: Bool=false,
    structure_p         :: Float64=0.5
    )

    # dataset dimensions
    n_samples, n_features = size(X)
    ind_r = collect(1:n_samples)
    ind_c = collect(1:n_features)

    # column and row subsampling number
    col_subsamp_n = convert(Int64,floor(col_subsamp_p*n_features))
    row_subsamp_n = convert(Int64,floor(row_subsamp_p*n_samples))

    # initialise empty array for trees and column indices
    trees = Array{CausalTrees.Tree,1}(undef,n_trees)
    feature_index = Array{Int64,2}(undef,(n_trees,col_subsamp_n))

    @inbounds for t in 1:n_trees
        # subsamples
        ind_r_t = StatsBase.sample(1:n_samples,row_subsamp_n,replace=row_resample)
        ind_c_t = StatsBase.sample(1:n_features,col_subsamp_n,replace=false)
        Xt = X[ind_r_t,ind_c_t]
        yt = y[ind_r_t]
        # train forest
        trees[t] = decision_tree(Xt,yt,loss,
            min_leaf_size=min_leaf_size,
            min_loss_increase=min_loss_increase,
            max_depth=max_depth,
            honesty=honesty,
            structure_p=structure_p)
        feature_index[t,:] = ind_c_t
    end
    return(Forest(trees,feature_index,n_trees,n_features))
end


text:
module PyCallUtils

using SparseArrays, Serialization, Random
using PyCall, Requires, BSON
using PyCall: jlWrapIteratorType, pyjlwrap_type!, pyincref_, pyjlwrap_iternext, pyjlwrap_new

export @imports, @from, pytypename, pymodulename, pymat

macro imports(lib, as, abbrev)
    esc(:($abbrev = pyimport($(string(lib)))))
end

macro imports(lib)
    exs = []
    if isa(lib, Expr) && lib.head == :tuple
        for arg in lib.args
            var, mdl = arg, string(arg)
            push!(exs, :($var = pyimport($mdl)))
        end
    elseif isa(lib, Expr) && lib.head == :.
        var, mdl = gensym(lib.args[2].value), string(lib)
        mdl = replace(mdl, r"[\(\)]" => "")
        push!(exs, :($var = pyimport($mdl)))
    else
        var, mdl = lib, string(lib)
        push!(exs, :($var = pyimport($mdl)))
    end
    esc(Expr(:block, exs...))
end

macro from(lib, imports, fs)
    fs = isa(fs, Expr) ? fs.args : [fs]
    mdl = replace(string(lib), r"[\(\)]" => "")
    exs = [:($f = pyimport($mdl).$(string(f))) for f in fs]
    esc(Expr(:block, exs...))
end

PyCall.PyObject(x::Colon) = pybuiltin("slice")(nothing, nothing, nothing)

Base.getindex(o::PyObject, s) = o.__getitem__(s)
Base.getindex(o::PyObject, is...) = getindex(o, is)
Base.setindex!(o::PyObject, v, s) = o.__setitem__(s, v)
Base.setindex!(o::PyObject, v, is...) = setindex!(o, v, is)

Base.getindex(o::PyObject, s::Union{Symbol, AbstractString}) = invoke(getindex, Tuple{PyObject, Any}, o, s)
Base.setindex!(o::PyObject, v, s::Union{Symbol, AbstractString}) = invoke(setindex!, Tuple{PyObject, Any, Any}, o, v, s)

Base.getindex(o::PyObject, i::Integer) = invoke(getindex, Tuple{PyObject, Any}, o, i)
Base.getindex(o::PyObject, i1::Integer, i2::Integer) = getindex(o, (i1, i2))
Base.getindex(o::PyObject, is::Integer...) = getindex(o, is)
Base.setindex!(o::PyObject, v, i::Integer) = invoke(setindex!, Tuple{PyObject, Any, Any}, o, v, i)
Base.setindex!(o::PyObject, v, i1::Integer, i2::Integer) = setindex!(o, (i1, i2), v)
Base.setindex!(o::PyObject, v, is::Integer...) = setindex!(o, is, v)

function PyCall.jlwrap_iterator(o::Any)
    if jlWrapIteratorType.tp_name == C_NULL # lazily initialize
        pyjlwrap_type!(jlWrapIteratorType, "PyCall.jlwrap_iterator") do t
            t.tp_iter = @cfunction(pyincref_, PyPtr, (PyPtr,)) # new reference to same object
            t.tp_iternext = @cfunction(pyjlwrap_iternext, PyPtr, (PyPtr,))
        end
    end
    iter_result = iterate(o)
    return pyjlwrap_new(jlWrapIteratorType, (o, Ref{Union{Nothing,typeof(iter_result)}}(iter_result)))
end

function PyCall.PyObject(S::SparseMatrixCSC)
    @from scipy.sparse imports csc_matrix
    csc_matrix((S.nzval, S.rowval .- 1, S.colptr .- 1), shape = size(S))
end

function Base.convert(::Type{SparseMatrixCSC}, o::PyObject)
    o = o.tocsc()
    SparseMatrixCSC(o.shape..., o.indptr .+ 1, o.indices .+ 1, o.data)
end

function Base.:∘(f::PyObject, g::PyObject)
    sf, sg, sh = ntuple(i -> randstring(['a':'z'; 'A':'Z'], 24), 3)
py"""
$$sf = $f; $$sg = $g; $$sh = lambda *args, **kwargs: $$sf($$sg(*args, **kwargs))
"""
    h = py"$$sh"
end

pytypename(obj) = obj.__class__.__name__
pymodulename(obj) = obj.__class__.__module__

pymat(x) = x
pymat(x::AbstractArray) = PyReverseDims(reshape(x, size(x, 1), :))

function BSON.lower(x::PyObject)
    io = IOBuffer()
    serialize(io, x)
    BSON.BSONDict(:tag => "pyobject", :data => take!(io))
end

@init BSON.tags[:pyobject] = d -> deserialize(IOBuffer(d[:data]))

end # module


text:
function wrapconsts(headerpath::String, outfile::String)
    lines = readlines(headerpath)
    filter!(x -> startswith(x, "const"), lines)

    juliaconsts = String[]
    lasttype = ""
    for line in lines
        parts = split(line, r" +")
        if occursin("//", line)
            comment = string(" #",split(line, "//")[2])
        else
            comment = ""
        end
        if length(parts) >= 4
            if parts[4] != "="
                val = strip(parts[4], ';')
                name = strip(parts[3], '=')
                type = parts[2]
                juliaconst = "const $(name) = $(type)($(val))$comment"
            else
                val = rstrip(parts[5], ';')
                name = parts[3]
                type = parts[2]
                juliaconst = "const $(name) = $(type)($(val))$comment"
            end
            !(type == lasttype) && push!(juliaconsts, "")
            push!(juliaconsts, juliaconst)
            lasttype = type
        else
            @warn "Couldn't parse $line"
        end
    end
    @info "Found and wrapped $(length(juliaconsts)) consts"
    isfile(outfile) && rm(outfile)
    @info "Writing outfile"
    open(outfile, "w") do io
        write(io, "# Consts wrapped by wrapconsts.jl\n")
        for juliaconst in juliaconsts
            write(io, juliaconst)
            write(io, "\n")
        end
    end
end


text:
using InteractiveUtils
using BenchmarkTools
using PyCall

versioninfo()

a = rand(10^7) # 1D vector of random numbers, uniform on [0,1)

function mysum(arr)
    s = 0.0
    for a in arr
        s+=a
    end
    s
end

function mysum_simd(arr)
    s = 0.0
    @simd for a in arr
        s += a
    end
    s
end

j_bench=@benchmark sum($a)
println("Julia embedded sum: Fastest time was $(minimum(j_bench.times) / 1e6) msec")

npsum=pyimport("numpy")["sum"]
np_bench=@benchmark npsum($a)
println("np.sum: Fastest time was $(minimum(np_bench.times) / 1e6) msec")

j_bench_hand = @benchmark mysum($a)
println("mysum: Fastest time was $(minimum(j_bench_hand.times)/1e6) msec")

elapsed=@elapsed mysum(a)
@show elapsed

elapsed=@elapsed mysum_simd(a)
@show elapsed

# このコードは動かない・・・
#j_bench_hand_simd @benchmark $(mysum_simd(a))
#println("mysum_simd: Fastest time was $(minimum(j_bench_hand_simd.times)/1e6) msec")

#= output
Julia Version 0.7.0
Commit a4cb80f3ed (2018-08-08 06:46 UTC)
Platform Info:
  OS: macOS (x86_64-apple-darwin14.5.0)
  CPU: Intel(R) Core(TM) M-5Y51 CPU @ 1.10GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-6.0.0 (ORCJIT, broadwell)
Julia embedded sum: Fastest time was 5.339243 msec
np.sum: Fastest time was 5.530426 msec
mysum: Fastest time was 11.793079 msec
elapsed = 0.025699647
elapsed = 0.026391872
=#

max_stars_repo_name: terasakisatoshi/juliaExer

max_issues_repo_name: terasakisatoshi/juliaExer

max_forks_repo_name: terasakisatoshi/juliaExer


text:
function check_dimensionality(c::CrystalNet{T}) where T
    edgs = [c.pos[dst(e)] .+ PeriodicGraphs.ofs(e) .- c.pos[src(e)] for e in edges(c.graph)]
    sort!(edgs)
    unique!(edgs)
    mat = reduce(hcat, edgs)
    if !isrank3(mat)
        throw(ArgumentError("The input net does not have dimensionality 3, so it cannot be analyzed."))
    end
    nothing
end

function check_valid_translation(c::CrystalNet{T}, t::SVector{3,<:Rational{<:Integer}}, r=nothing) where T
    U = soft_widen(T)
    n = length(c.pos)
    vmap = Int[]
    offsets = SVector{3,Int}[]
    for k in 1:n
        curr_pos = c.pos[k]
        transl = U.(isnothing(r) ? curr_pos : (r * curr_pos)) .+ t
        ofs = floor.(Int, transl)
        x = transl .- ofs
        (i, j) = x < curr_pos ? (1, k) : (k, length(c.types)+1)
        while j - i > 1
            m = (j+i)>>1
            if cmp(x, c.pos[m]) < 0
                j = m
            else
                i = m
            end
        end
        (c.pos[i] == x && c.types[i] == c.types[k]) || return nothing
        push!(vmap, i)
        push!(offsets, ofs)
    end
    for e in edges(c.graph)
        src = vmap[e.src]
        dst = vmap[e.dst.v]
        newofs = (isnothing(r) ? e.dst.ofs : r * e.dst.ofs) .+ offsets[e.dst.v] .- offsets[e.src]
        has_edge(c.graph, PeriodicGraphs.unsafe_edge{3}(src, dst, newofs)) || return nothing
    end
    return vmap
end

function possible_translations(c::CrystalNet{T}) where T
    n = length(c.pos)
    ts = Tuple{Int, Int, Int, SVector{3,T}}[]
    sortedpos = copy(c.pos)
    origin = popfirst!(sortedpos)
    @assert iszero(origin)
    sort!(SVector{3,widen(soft_widen(widen(T)))}.(sortedpos), by=norm)
    for t in sortedpos
        @assert t == back_to_unit.(t)
        max_den, i_max_den = findmax(denominator.(t))
        numerator(t[i_max_den]) == 1 || continue
        nz = count(iszero, t)
        push!(ts, (nz, i_max_den, max_den, t))
    end
    return sort!(ts; by=(x->(x[1], x[2], x[3])))
end

# TODO remove?
function find_first_valid_translations(c::CrystalNet)
    for (nz, i_max_den, max_den, t) in possible_translations(c)
        !isnothing(check_valid_translation(c, t)) && return (nz, i_max_den, max_den, t)
    end
    return nothing
end

function find_all_valid_translations(c::CrystalNet{T}) where T
    ret = NTuple{3, Vector{Tuple{Int, Int, SVector{3,T}}}}(([], [], []))
    for (nz, i_max_den, max_den, t) in possible_translations(c)
        !isnothing(check_valid_translation(c, t)) && push!(ret[nz+1], (i_max_den, max_den, t))
    end
    return ret
end

function minimal_volume_matrix(translations::Tuple{T,T,T}) where T
    nz0, nz1, nz2 = translations

    denmax = [1//1, 1//1, 1//1]
    imax = [0, 0, 0]
    for j in 1:length(nz2)
        i, den, _ = nz2[j]
        if den > denmax[i]
            imax[i] = j
            denmax[i] = den
        end
    end
    for j in 1:3
        if imax[j] == 0
            push!(nz2, (j, 1, [j==1, j==2, j==3]))
            imax[j] = length(nz2)
        end
    end
    _nz2 = [nz2[i] for i in imax]
    empty!(nz2)
    append!(nz2, _nz2)

    # TODO optimize
    all = vcat(nz0, nz1, nz2)
    n = length(all)
    detmax = 1//1
    best = (0, 0, 0)
    @inbounds for i in 1:n-2
        for j in i+1:n-1
            for k in j+1:n
                d = abs(det(hcat(all[i][3], all[j][3], all[k][3])))
                d == 0 && continue
                if d < detmax
                    detmax = d
                    best = (i, j, k)
                end
            end
        end
    end
    @assert !any(iszero.(best))
    i, j, k = best
    ret = hcat(all[i][3], all[j][3], all[k][3])
    if det(ret) < 0
        ret = hcat(all[i][3], all[j][3], .-all[k][3])
    end
    return ret
end

function reduce_with_matrix(c::CrystalNet{<:Rational{T}}, mat) where T
    U = soft_widen(widen(T))
    lengths = degree(c.graph)
    cell = Cell(c.cell, c.cell.mat * mat)
    imat = soft_widen(T).(inv(mat)) # The inverse should only have integer coefficients
    poscol = (U.(imat),) .* c.pos

    offset = [floor.(Int, x) for x in poscol]
    for i in 1:length(poscol)
        poscol[i] = poscol[i] .- offset[i]
    end
    I_sort = sort(1:length(poscol); by=i->(poscol[i], hash_position(offset[i])))
    i = popfirst!(I_sort)
    @assert iszero(offset[i])
    I_kept = Int[i]
    sortedcol = SVector{3,Rational{U}}[poscol[i]]
    for i in I_sort
        x = poscol[i]
        if x != sortedcol[end]
            push!(I_kept, i)
            push!(sortedcol, x)
        end
    end
    edges = PeriodicEdge3D[]
    for i in 1:length(I_kept)
        ofs_i = offset[I_kept[i]]
        for neigh in neighbors(c.graph, I_kept[i])
            x = poscol[neigh.v]
            j = searchsortedfirst(sortedcol, x)
            @assert j <= length(sortedcol) && sortedcol[j] == x
            ofs_x = offset[neigh.v]
            push!(edges, (i, j, ofs_x - ofs_i .+ imat*neigh.ofs))
        end
    end
    graph = PeriodicGraph3D(edges)
    @assert degree(graph) == lengths[I_kept]
    return CrystalNet(cell, c.types[I_kept], sortedcol, graph)
end

function minimize(net::CrystalNet)
    translations = find_all_valid_translations(net)
    while !all(isempty.(translations))
        mat = minimal_volume_matrix(translations)
        net = reduce_with_matrix(net, mat)
        translations = find_all_valid_translations(net) # TODO don't recompute
    end
    return net
end

function findfirstbasis(offsets)
    newbasis = nf3D(offsets)
    invbasis = inv(newbasis)
    intcoords = [Int.(invbasis * x) for x in offsets]
    return newbasis, intcoords
end

function findbasis(edges::Vector{Tuple{Int,Int,SVector{3,T}}}) where T
    m = length(edges)
    positivetrans = SVector{3,T}[]
    map_to_ofs = Vector{Int}(undef, m)
    tmp_map = Int[]
    for i in 1:m
        trans = last(edges[i])
        c = cmp(trans, zero(SVector{3,Int}))
        if iszero(c)
            map_to_ofs[i] = 0
        elseif c < 0
            push!(tmp_map, -i)
            push!(positivetrans, .-trans)
        else
            push!(tmp_map, i)
            push!(positivetrans, trans)
        end
    end

    I_sort = sortperm(positivetrans)
    uniques = SVector{3,T}[]
    last_trans = zero(SVector{3,T})

    for j in 1:length(positivetrans)
        i = tmp_map[I_sort[j]]
        trans = positivetrans[I_sort[j]]
        if trans != last_trans
            push!(uniques, trans)
            last_trans = trans
        end
        map_to_ofs[abs(i)] = sign(i) * length(uniques)
    end

    basis, unique_coords = findfirstbasis(uniques)

    newedges = Vector{PeriodicEdge3D}(undef, m)
    for j in 1:m
        i = map_to_ofs[j]
        src, dst, _ = edges[j]
        if i == 0
            newedges[j] = PeriodicEdge3D(src, dst, zero(SVector{3,Int}))
        else
            newedges[j] = PeriodicEdge3D(src, dst, sign(i) .* unique_coords[abs(i)])
        end
    end

    @assert all(z-> begin x, y = z; x == (y.src, y.dst.v, basis*y.dst.ofs) end, zip(edges, newedges))

    return basis, newedges
end

function candidate_key(net::CrystalNet{T}, u, basis, minimal_edgs) where T
    V = widen(widen(T))
    n = nv(net.graph)
    h = 2 # Next node to assign
    origin = net.pos[u]
    newpos = Vector{SVector{3,V}}(undef, n) # Positions of the kept representatives
    newpos[1] = zero(SVector{3,V})
    offsets = Vector{SVector{3,Int}}(undef, n)
    # Offsets of the new representatives with respect to the original one, in the original basis
    offsets[1] = zero(SVector{3,Int})
    vmap = Vector{Int}(undef, n) # Bijection from the old to the new node number
    vmap[1] = u
    rev_vmap = zeros(Int, n) # inverse of vmap
    rev_vmap[u] = 1
    flag_bestedgs = false # Marks whether the current list of edges is lexicographically
    # below the best known one. If not by the end of the algorithm, return false
    edgs = Tuple{Int,Int,SVector{3,V}}[]
    mat = V.(inv(widen(V).(basis)))
    for t in 1:n # t is the node being processed
        neighs = neighbors(net.graph, vmap[t])
        ofst = offsets[t]
        pairs = Vector{Tuple{SVector{3,V},Int}}(undef, length(neighs))
        for (i,x) in enumerate(neighs)
            pairs[i] = (V.(mat*(net.pos[x.v] .+ x.ofs .- origin .+ ofst)), x.v)
        end
        # (x,i) ∈ pairs means that vertex i (in the old numerotation) has position x in the new basis
        order = unique(last.(sort(pairs)))
        inv_order = Vector{Int}(undef, n)
        for i in 1:length(order)
            inv_order[order[i]] = i
        end
        sort!(pairs, lt = ((x, a), (y, b)) -> begin inv_order[a] < inv_order[b] ||
                                                    (inv_order[a] == inv_order[b] && x < y) end)
        # pairs is sorted such that different nodes first appear in increasing order of their position
        # but different representatives of the same node are contiguous and also sorted by position.
        bigbasis = widen(V).(basis)
        for (coordinate, v) in pairs
            idx = rev_vmap[v]
            if idx == 0 # New node to which h is assigned
                @assert t < h
                vmap[h] = v
                rev_vmap[v] = h
                newpos[h] = coordinate
                offsets[h] = SVector{3,Int}(bigbasis * coordinate .+ origin .- net.pos[v])
                push!(edgs, (t, h, zero(SVector{3,T})))
                h += 1
            else
                realofs = coordinate .- newpos[idx]
                # offset between this representative of the node and that which was first encountered
                push!(edgs, (t, idx, realofs))
            end
            if !flag_bestedgs
                j = length(edgs)
                c = cmp(minimal_edgs[j], edgs[j])
                c < 0 && return (Int[], Tuple{Int,Int,SVector{3,V}}[])
                c > 0 && (flag_bestedgs = true)
            end
        end
    end
    @assert allunique(edgs)
    if !flag_bestedgs
        @assert minimal_edgs == edgs
        return (Int[], Tuple{Int,Int,SVector{3,V}}[])
    end
    return vmap, edgs
end

"""
Partition the vertices of the graph into disjoint categories, one for each
coordination sequence. The partition is then sorted by order of coordination sequence.
This partition does not depend on the representation of the graph.
The optional argument vmaps is a set of permutations of the vertices that leave the
graph unchanged. In other words, vmaps is a set of symmetry operations of the graph.

Also returns the map from vertices to an identifier such that all vertices with
the same identifier are symmetric images of one another, as well as a list of
unique representative for each symmetry class.
"""
function partition_by_coordination_sequence(graph, vmaps=Vector{Int}[])
    # First, union-find on the symmetries to avoid computing useless coordination sequences
    n = nv(graph)
    unionfind = collect(1:n)
    for vmap in vmaps
        for (i,j) in enumerate(vmap)
            i == j && continue
            if i > j
                i, j = j, i
            end
            repri = i
            while repri != unionfind[repri]
                repri = unionfind[repri]
            end
            reprj = j
            while reprj != unionfind[reprj]
                reprj = unionfind[reprj]
            end
            repri == reprj && continue
            if repri < reprj
                unionfind[j] = repri
            else
                unionfind[i] = reprj
            end
        end
    end

    cat_map = zeros(Int, n)
    unique_reprs = Vector{Int}[]
    categories = Vector{Int}[]
    for i in 1:n
        repri = i
        if repri != unionfind[repri]
            repri = unionfind[repri]
            while repri != unionfind[repri]
                repri = unionfind[repri]
            end
            descent = i
            while descent != unionfind[descent]
                tmp = unionfind[descent]
                unionfind[descent] = repri
                descent = tmp
            end
            @assert descent == repri
        end
        cat = cat_map[repri]
        if iszero(cat)
            push!(unique_reprs, [i])
            push!(categories, [i])
            cat = length(categories)
            cat_map[repri] = cat
        else
            push!(categories[cat], i)
        end
        cat_map[i] = cat
    end

    PeriodicGraphs.graph_width!(graph) # setup for the computation of coordination sequences
    csequences = [coordination_sequence(graph, first(cat), 10) for cat in categories]
    I = sortperm(csequences)
    @assert csequences[I[1]][1] >= 2 # vertices of degree <= 1 should have been removed at input creation

    todelete = falses(length(categories))
    last_i = I[1]
    for j in 2:length(I)
        i = I[j]
        if csequences[i] == csequences[last_i]
            todelete[i] = true
            append!(categories[last_i], categories[i])
            push!(unique_reprs[last_i], pop!(unique_reprs[i]))
            empty!(categories[i])
        else
            last_i = i
        end
    end
    deleteat!(categories, todelete)
    deleteat!(unique_reprs, todelete)
    deleteat!(csequences, todelete)

    num = length(categories)
    numoutgoingedges = Vector{Tuple{Int,Vector{Int}}}(undef, num)
    for i in 1:num
        seq = csequences[i]
        numoutgoingedges[i] = (length(categories[i]) * seq[1], seq)
    end
    sortorder = sortperm(numoutgoingedges)
    # categories are sorted by the total number of outgoing directed edges from each category
    # and by the coordination sequence in case of ex-aequo.
    # categories are thus uniquely determined and ordered independently of the representation of the net

    @assert allunique(csequences)
    for i in 1:num
        @assert all(coordination_sequence(graph, x, 10) == csequences[i] for x in categories[i])
    end # TODO comment out these costly asserts
    return categories[sortorder], cat_map, unique_reprs[sortorder]
end

"""
Find a set of pairs (v, basis) that are candidates for finding the key with
candidate_key(net, v, basis).
The returned set is independent of the representation of the graph used in net.
"""
function find_candidates(net::CrystalNet{T}) where T
    rotations, vmaps, hasmirror = find_symmetries(net)
    @assert length(rotations) == length(vmaps)
    categories, symmetry_map, unique_reprs = partition_by_coordination_sequence(net.graph, vmaps)
    category_map = Vector{Int}(undef, nv(net.graph))
    for (i, cat) in enumerate(categories)
        for j in cat
            category_map[j] = i
        end
    end
    @assert sort.(categories) == sort.(first(partition_by_coordination_sequence(net.graph, [])))
    candidates = Dict{Int,Vector{SMatrix{3,3,soft_widen(T),9}}}()
    for reprs in unique_reprs
        # First, we try to look for triplet of edges all starting from the same vertex within a category
        degree(net.graph, first(reprs)) <= 3 && continue
        candidates = find_candidates_onlyneighbors(net, reprs, category_map)
        !isempty(candidates) && break
    end
    if isempty(candidates)
        # If we arrive at this point, it means that all vertices only have coplanar neighbors.
        # Then we look for all triplets of edges two of which start from a vertex
        # and one does not, stopping and the first pair of categories for which this
        # set of candidates is non-empty.
        if isempty(candidates)
            for reprs in unique_reprs
                candidates = find_candidates_fallback(net, reprs, categories, category_map)
                !isempty(candidates) && break
            end
        end
    end
    if isempty(candidates)
        if !allunique(net.pos)
            throw(ArgumentError("""The net is unstable, hence it cannot be analyzed."""))
        else
            check_dimensionality(net)
            error("Internal error: no candidate found.")
        end
    end
    @assert !isempty(candidates)
    return extract_through_symmetry(candidates, vmaps, rotations), category_map
end

function extract_through_symmetry(candidates::Dict{Int,Vector{SMatrix{3,3,T,9}}}, vmaps, rotations) where T
    unique_candidates = Pair{Int,SMatrix{3,3,T,9}}[]
    for (i, mats) in candidates
        @assert i == minimum(vmap[i] for vmap in vmaps)
        indices = [j for j in 1:length(vmaps) if vmaps[j][i] == i]
        min_mats = Set{SVector{9,T}}()
        for mat in mats
            flattenedmat = SVector{9,T}(mat)
            min_mat = flattenedmat
            for j in indices
                new_mat = SVector{9,T}(rotations[j] * mat)
                if new_mat < min_mat
                    min_mat = new_mat
                end
            end
            push!(min_mats, min_mat)
        end
        for x in min_mats
            push!(unique_candidates, i => SMatrix{3,3,T,9}(x))
        end
    end
    return unique_candidates
end

function find_candidates_onlyneighbors(net::CrystalNet{T}, candidates_v, category_map) where T
    U = soft_widen(T)
    deg = degree(net.graph, first(candidates_v)) # The degree is the same for all vertices of the same category
    initial_candidates = Pair{Int,Tuple{Matrix{U},Vector{Int}}}[]
    fastlock = SpinLock()
    @threads for v in candidates_v
        a = Matrix{U}(undef, 3, deg)
        cats = Vector{Int}(undef, deg)
        posi = net.pos[v]
        for (j, x) in enumerate(neighbors(net.graph, v))
            a[:,j] .= net.pos[x.v] .+ x.ofs .- posi
            cats[j] = category_map[x.v]
        end
        if isrank3(a)
            pair = v => (a, cats)
            lock(fastlock) do
                push!(initial_candidates, pair)
            end
        end
    end

    candidates = Dict{Int,Vector{SMatrix{3,3,U,9}}}()
    current_cats::SizedVector{3,Int} = SizedVector{3,Int}(fill(length(category_map), 3))
    current_ordertype::Int = 1
    isempty(initial_candidates) && return candidates

    @threads for (v, (mat, cats)) in initial_candidates
        _, n = size(mat)
        matv = SMatrix{3,3,U,9}[]
        lock(fastlock)
        ordertype = current_ordertype
        mincats = copy(current_cats)
        unlock(fastlock)
        for _i in 1:n, _j in (_i+1):n, _k in (_j+1):n
            m = SMatrix{3,3,widen(T),9}(mat[:,[_i,_j,_k]])
            issingular(m) && continue
            orders = SVector{3,Int}[]
            subcats = cats[SVector{3,Int}(_i, _j, _k)]
            reorder = begin
                cmp23 = subcats[2] >= subcats[3]
                cmp13 = subcats[1] >= subcats[3]
                if subcats[1] >= subcats[2]
                    cmp23 ? (3,2,1) : (cmp13 ? (2,3,1) : (2,1,3))
                else
                    cmp23 ? (cmp13 ? (3,1,2) : (1,3,2)) : (1,2,3)
                end
            end
            subcats = subcats[SVector{3,Int}(reorder)]
            if subcats[1] == subcats[3]
                (ordertype != 1 || mincats[1] < subcats[1]) && continue
                if mincats[1] > subcats[1]
                    empty!(matv)
                    mincats = subcats
                end
                orders = SVector{3,Int}[(1,2,3), (1,3,2), (2,1,3), (2,3,1), (3,1,2), (3,2,1)]
            elseif subcats[2] == subcats[3]
                (ordertype > 2 || (ordertype == 2 && mincats < subcats)) && continue
                if ordertype == 1 || mincats > subcats
                    empty!(matv)
                    ordertype = 2
                    mincats = subcats
                end
                orders = SVector{3,Int}[reorder, (reorder[1], reorder[3], reorder[2])]
            elseif subcats[1] == subcats[2]
                (ordertype == 4 || (ordertype == 3 && mincats < subcats)) && continue
                if ordertype <= 2 || mincats > subcats
                    empty!(matv)
                    ordertype = 3
                    mincats = subcats
                end
                orders = SVector{3,Int}[reorder, (reorder[2], reorder[1], reorder[3])]
            else
                ordertype == 4 && mincats < subcats && continue
                if ordertype != 4 || mincats > subcats
                    empty!(matv)
                    ordertype = 4
                    mincats = subcats
                end
                orders = SVector{3,Int}[reorder]
            end
            mats = [m[:,o] for o in orders]
            append!(matv, mats)
        end
        if ordertype < current_ordertype || (ordertype == current_ordertype
                                             && mincats > current_cats)
            continue # fast-path to avoid unecessary locking
        end

        lock(fastlock)
        if ordertype < current_ordertype || (ordertype == current_ordertype
                                             && mincats > current_cats)
            unlock(fastlock)
            continue
        end
        if ordertype > current_ordertype || mincats < current_cats
            empty!(candidates)
            current_ordertype = ordertype
            current_cats = mincats
        end
        candidates[v] = matv
        unlock(fastlock)
    end

    return candidates
end

function find_candidates_fallback(net::CrystalNet{T}, reprs, othercats, category_map) where T
    U = soft_widen(T)
    candidates = Dict{Int,Vector{SMatrix{3,3,U,9}}}(u => [] for u in reprs)
    n = length(category_map)
    mincats = SizedVector{3,Int}(fill(n, 3))
    current_cats = SizedVector{3,Int}(fill(0, 3))
    for cat in othercats
        for u in reprs
            mats = SMatrix{3,3,T,9}[]
            posu = net.pos[u]
            neighu = neighbors(net.graph, u)
            for (j1, x1) in enumerate(neighu)
                category_map[x1.v] > mincats[1] && category_map[x1.v] > mincats[2] && continue
                for j2 in j1+1:length(neighu)
                    x2 = neighu[j2]
                    if category_map[x1.v] > category_map[x2.v]
                        x1, x2 = x2, x1
                    end
                    current_cats[1] = category_map[x1.v]
                    current_cats[1] > mincats[1] && continue
                    current_cats[2] = category_map[x2.v]
                    current_cats[1] == mincats[1] && current_cats[2] > mincats[2] && continue
                    vec1 = net.pos[x1.v] .+ x1.ofs .- posu
                    vec2 = net.pos[x2.v] .+ x2.ofs .- posu
                    j = iszero(vec1[1]) ? iszero(vec1[2]) ? 3 : 2 : 1
                    vec1[j] * vec2 == vec2[j] .* vec1 && continue # colinear
                    for v in cat
                        posv = net.pos[v]
                        for x3 in neighbors(net.graph, v)
                            vec3 = net.pos[x3.v] .+ x3.ofs .- posv
                            mat = SMatrix{3,3,widen(T),9}(hcat(vec1, vec2, vec3))
                            issingular(mat) && continue
                            current_cats[3] = category_map[x3.v]
                            current_cats > mincats && continue
                            if current_cats < mincats
                                mincats .= current_cats
                                foreach(empty!, values(candidates))
                                empty!(mats)
                            end
                            # if d < 0
                            #     mat = hcat(vec2, vec1, vec3)
                            # end
                            push!(mats, mat)
                            if current_cats[1] == current_cats[2]
                                push!(mats, hcat(vec2, vec1, vec3))
                            end
                        end
                    end
                end
            end
            append!(candidates[u], mats)
        end
        # If at this point candidates is not empty, then we have found the first
        # other category cat such that there is a candidate with two edges from
        # maincat and one from cat. This is uniquely determined independently of
        # the representation of the net, so we can simply return the found candidates
        # without needing to check for the remaining categories.
        if !all(isempty, values(candidates))
            for (u, mats) in candidates
                if isempty(mats)
                    delete!(candidates, u)
                end
            end
            return candidates
         end
    end
    @assert all(isempty, values(candidates))
    return Dict{Int,Vector{SMatrix{3,3,U,9}}}()
end

# function parallel_topological_key(net::CrystalNet{T}) where T
#     candidates = find_candidates(net)
#     # @show length(candidates)
#     numthreads = min(nthreads(), length(candidates))
#     minimal_graph = Vector{PeriodicGraph3D}(undef, numthreads)
#     minimal_vmap = Vector{Vector{Int}}(undef, numthreads)
#     minimal_basis = Vector{SMatrix{3,3,T,9}}(undef, numthreads)
#     @threads for i in 1:numthreads
#         v, basis = candidates[end+1-i]
#         minimal_basis[i], minimal_vmap[i], minimal_graph[i] = candidate_key(net, v, basis)
#     end
#     resize!(candidates, length(candidates) - numthreads)
#     @threads for (v, basis) in candidates
#         id = threadid() # necessarily inbounds because otherwise the candidates is empty
#         newbasis, vmap, graph = candidate_key(net, v, basis)
#         if edges(graph) < edges(minimal_graph[id])
#             minimal_graph[id] = graph
#             minimal_vmap[id] = vmap
#             minimal_basis[id] = newbasis
#         end
#     end
#     _, j = findmin(edges.(minimal_graph))
#     return minimal_basis[j], minimal_vmap[j], minimal_graph[j]
# end

function topological_key(net::CrystalNet{T}) where T
    if isempty(net.pos)
        throw(ArgumentError("The net is empty. This probably indicates that the input is not 3-dimensional."))
    end
    # @assert allunique(net.pos) # FIXME make a more precise check for net stability. Currently fails for sxt
    candidates, category_map = find_candidates(net)
    v, minimal_basis = popfirst!(candidates)
    n = nv(net.graph)
    _edgs = [(n + 1, 0, zero(SVector{3,T}))]
    minimal_vmap, minimal_edgs = candidate_key(net, v, minimal_basis, _edgs)
    for (v, basis) in candidates
        vmap, edgs = candidate_key(net, v, basis, minimal_edgs)
        isempty(vmap) && continue
        @assert edgs < minimal_edgs
        if edgs < minimal_edgs
            minimal_edgs = edgs
            minimal_vmap = vmap
            minimal_basis = basis
        end
    end

    newbasis, coords = findbasis(minimal_edgs)
    # @show abs(det(newbasis))
    if abs(det(newbasis)) != 1
        # @warn "The input structure does not appear to be connected. Analyzing the first component only."
        # see for instance "3 1 2 2 0 0 1 2 0 2 0 1 2 0 0 2 1 2 0 0 0"
        # TODO check that only a disconnected input can yield this.
    end

    return minimal_basis * newbasis, minimal_vmap, PeriodicGraph3D(n, coords)
end

function find_new_representation(pos, basis, vmap, graph)
    return collect(eachcol(CIFTypes.equilibrium(graph))) # TODO optimize
end

function topological_genome(g::PeriodicGraph3D, skipminimize=true)
    net = CrystalNet(g)
    if !allunique(net.pos)
        throw(ArgumentError("""The net is unstable, hence it cannot be analyzed."""))
    end
    if !skipminimize
        net = minimize(net)
    end
    return last(topological_key(net))
end

function topological_genome(net::CrystalNet, skipminimize=false)
    if !allunique(net.pos)
        throw(ArgumentError("""The net is unstable, hence it cannot be analyzed."""))
    end
    if !skipminimize
        net = minimize(net)
    end
    basis, vmap, graph = topological_key(net)

    return graph

    # return CrystalNet(graph) # TODO remove this temporary bandaid
    # # FIXME CrystalNet(graph).graph != graph so this is invalid !
    #
    # positions = find_new_representation(net.pos, basis, vmap, graph)
    #
    # ret = CrystalNet(Cell(net.cell, net.cell.mat*basis), net.types[vmap], positions, graph)
    # means = mean(ret.pos)
    # for i in 1:length(ret.pos)
    #     ret.pos[i] = ret.pos[i] .+ [1//2, 1//2, 1//2] .- means
    # end
    # @assert all(ret.pos[i] == mean(ret.pos[x.v] .+ x.ofs for x in neighbors(ret.graph, i)) for i in 1:length(net.pos))
    #
    # return ret
end

function reckognize_topology(genome::AbstractString, arc=CRYSTAL_NETS_ARCHIVE)
    get(arc, genome, "UNKNOWN")
end
function reckognize_topology(graph::PeriodicGraph3D, arc=CRYSTAL_NETS_ARCHIVE)
    reckognize_topology(string(graph), arc)
end

function reckognize_topologies(path; ignore_atoms=())
    ret = Dict{String,String}()
    failed = Dict{String, Tuple{Exception, Vector{Union{Ptr{Nothing}, Base.InterpreterIP}}}}()
    newarc = copy(CRYSTAL_NETS_ARCHIVE)
    @threads for f in readdir(path)
        name = first(splitext(f))
        genome::String = try
            string(topological_genome(CrystalNet(parse_chemfile(path*f; ignore_atoms))))
        catch e
            if e isa InterruptException ||
              (e isa TaskFailedException && e.task.result isa InterruptException)
                rethrow()
            end
            failed[name] = (e, catch_backtrace())
            ""
        end
        if !isempty(genome)
            id = reckognize_topology(genome, newarc)
            if id == "UNKNOWN"
                newarc[genome] = name
                ret[name] = genome
            else
                ret[name] = id
            end
        end
    end
    return ret, failed
end


text:
"Returns the last b digits of a number in little-endian binary form"
function lastbbits(n::Integer,b=5)
    map(d->Bool(rem(n>>(d-1),2)), 1:b)
end

using StaticArrays: @SVector
function secret_handshake(code::Integer)
    actions = @SVector ["wink", "double blink", "close your eyes", "jump"]
    # special "reverse" action isn't in actions, so we account for it here
    n_actions = length(actions) + 1
    instructions = lastbbits(code,n_actions)
    decoded = actions[instructions[1:length(actions)]]
    instructions[5] ? decoded |> reverse : decoded
end

# 1.3x speedup on average
function secret_handshake_quick(code::Integer)
    instructions = lastbbits(code)
    decoded = Array{String,1}()

    # reverse is quite an expensive operation, so we want to avoid it
    if instructions[5]
        instructions[4] && push!(decoded,"jump")
        instructions[3] && push!(decoded,"close your eyes")
        instructions[2] && push!(decoded,"double blink")
        instructions[1] && push!(decoded,"wink")
    else
        instructions[1] && push!(decoded,"wink")
        instructions[2] && push!(decoded,"double blink")
        instructions[3] && push!(decoded,"close your eyes")
        instructions[4] && push!(decoded,"jump")
    end
    decoded
end