text:
# Solving the baseline version of the model with estimated parameters
using AAABK2018

# parameters
λ =  1.131874350361868; ψ =  0.037175089899834; ν =  0.205622086535123
α =  0.925520992611854; ϕ =  0.216311889844990; θˡ=  1.390807772162332
θʰ=  1.751156379867318; θᴱ=  0.023600456924424; ε =  2.900000000000000
ρ =  0.020000000000000; γ =  0.500000000000000; γᴱ=  0.500000000000000
σ =  2.000000000000000; Lˢ=  0.165500000000000; ω =  1.0

p = Params(λ,ψ,ν, α, ϕ, θˡ, θʰ, θᴱ, ε, ρ, γ, γᴱ, σ, Lˢ, ω)

# initial guess for equilibirum objects
eqInit = [ 2.00, 0.55, 0.06,
		   0.72, 0.83, 1.73]

# solve
eq,res = solveBGP(p,eqInit);

# solve with 5% incumbent R&D subsidy
pIncSub = Params(p, (sⁱ = 0.05,))   # Params(pOld, modifiedParameters)
eqIncSub,resIncSub = solveBGP(pIncSub, eqInit);

# solve with 5% operation cost subsidy
pFixSub = Params(p, (sᶠ = 0.05,))
eqFixSub,resFixSub = solveBGP(pFixSub,eqInit);

# solve with 5% entrant R&D subsidy
pEntSub = Params(p, (sᴱ = 0.05,))
eqEntSub,resEntSub = solveBGP(pEntSub,eqInit);

#social planner problem
eqSoc,resSoc = solveSocialPlanner(p)

# optimal policy
# one-tool policies
optPolInc = policy_opt(p, sⁱ = true)
optPolFix = policy_opt(p, sᶠ = true)
optPolEnt = policy_opt(p, sᴱ = true)

# two-tool policy
optPolInc_Fix = policy_opt(p, sⁱ = true, sᶠ = true)


text:
using KernelTools
using Base.Test

include("macros.jl")

ext: jl


text:
## multi.jl - multiprocessing
##
## julia starts with one process, and processors can be added using:
##   addprocs(n)                         using exec
##   addprocs({"host1","host2",...})     using remote execution
##
## remotecall(w, func, args...) -
##     tell a worker to call a function on the given arguments.
##     returns a RemoteRef to the result.
##
## remote_do(w, f, args...) - remote function call with no result
##
## wait(rr) - wait for a RemoteRef to be finished computing
##
## fetch(rr) - wait for and get the value of a RemoteRef
##
## remotecall_fetch(w, func, args...) - faster fetch(remotecall(...))
##
## pmap(func, lst) -
##     call a function on each element of lst (some 1-d thing), in
##     parallel.
##
## RemoteRef() - create an uninitialized RemoteRef on the local processor
##
## RemoteRef(p) - ...or on a particular processor
##
## put!(r, val) - store a value to an uninitialized RemoteRef
##
## @spawn expr -
##     evaluate expr somewhere. returns a RemoteRef. all variables in expr
##     are copied to the remote processor.
##
## @spawnat p expr - @spawn specifying where to run
##
## @async expr -
##     run expr as an asynchronous task on the local processor
##
## @parallel (r) for i=1:n ... end -
##     parallel loop. the results from each iteration are reduced using (r).
##
## @everywhere expr - run expr everywhere.

# todo:
# - more indexing
# * take!() to empty a Ref (full/empty variables)
# * have put!() wait on non-empty Refs
# - removing nodes
# - more dynamic scheduling
# * fetch/wait latency seems to be excessive
# * message aggregation
# - timer events
# - send pings at some interval to detect failed/hung machines
# - integrate event loop with other kinds of i/o (non-messages)
# ? method_missing for waiting (getindex/setindex!/localdata seems to cover a lot)
# * serializing closures
# * recover from i/o errors
# * handle remote execution errors
# * all-to-all communication
# * distributed GC
# * call&wait and call&fetch combined messages
# * aggregate GC messages
# * dynamically adding nodes (then always start with 1 and grow)
# * add readline to event loop
# * GOs/darrays on a subset of nodes

## workers and message i/o ##

function send_msg_unknown(s::IO, kind, args)
    error("attempt to send to unknown socket")
end

function send_msg(s::IO, kind, args...)
    id = worker_id_from_socket(s)
    if id > -1
        return send_msg(worker_from_id(id), kind, args...)
    end
    send_msg_unknown(s, kind, args)
end

function send_msg_now(s::IO, kind, args...)
    id = worker_id_from_socket(s)
    if id > -1
        return send_msg_now(worker_from_id(id), kind, args...)
    end
    send_msg_unknown(s, kind, args)
end

abstract ClusterManager

type Worker
    host::ByteString
    port::Uint16
    socket::TcpSocket
    sendbuf::IOBuffer
    del_msgs::Array{Any,1}
    add_msgs::Array{Any,1}
    id::Int
    gcflag::Bool
    bind_addr::IpAddr
    manage::Function
    config::Dict

    Worker(host::String, port::Integer, sock::TcpSocket, id::Int) =
        new(bytestring(host), uint16(port), sock, IOBuffer(), {}, {}, id, false)
end
Worker(host::String, port::Integer, sock::TcpSocket) =
    Worker(host, port, sock, 0)
function Worker(host::String, port::Integer)
    # Connect to the loopback port if requested host has the same ipaddress as self.
    if host == string(LPROC.bind_addr)
        w = Worker(host, port, connect("127.0.0.1", uint16(port)))
    else
        w = Worker(host, port, connect(host, uint16(port)))
    end
    # Avoid calling getaddrinfo if possible - involves a DNS lookup
    # host may be a stringified ipv4 / ipv6 address or a dns name
    if host == "localhost"
        w.bind_addr = parseip("127.0.0.1")
    else
        try
            w.bind_addr = parseip(host)
        catch
            w.bind_addr = getaddrinfo(host)
        end
    end
    w
end
function Worker(host::String, bind_addr::String, port::Integer, tunnel_user::String, sshflags)
    w = Worker(host, port, connect("localhost",
                               ssh_tunnel(tunnel_user, host, bind_addr, uint16(port), sshflags)))
    w.bind_addr = parseip(bind_addr)
    w
end

function send_msg_now(w::Worker, kind, args...)
    send_msg_(w, kind, args, true)
end

function send_msg(w::Worker, kind, args...)
    send_msg_(w, kind, args, false)
end

function flush_gc_msgs(w::Worker)
    w.gcflag = false
    msgs = copy(w.add_msgs)
    if !isempty(msgs)
        empty!(w.add_msgs)
        remote_do(w, add_clients, msgs...)
    end

    msgs = copy(w.del_msgs)
    if !isempty(msgs)
        empty!(w.del_msgs)
        #print("sending delete of $msgs\n")
        remote_do(w, del_clients, msgs...)
    end
end

#TODO: Move to different Thread
function enq_send_req(sock::TcpSocket,buf,now::Bool)
    arr=takebuf_array(buf)
    write(sock,arr)
    #TODO implement "now"
end

function send_msg_(w::Worker, kind, args, now::Bool)
    #println("Sending msg $kind")
    buf = w.sendbuf
    serialize(buf, kind)
    for arg in args
        serialize(buf, arg)
    end

    if !now && w.gcflag
        flush_gc_msgs(w)
    else
        enq_send_req(w.socket,buf,now)
    end
end

function flush_gc_msgs()
    for w in (PGRP::ProcessGroup).workers
        if isa(w,Worker)
            k = w::Worker
            if k.gcflag
                flush_gc_msgs(k)
            end
        end
    end
end

## process group creation ##

type LocalProcess
    id::Int
    bind_addr::IpAddr
    LocalProcess() = new(1)
end

const LPROC = LocalProcess()

const map_pid_wrkr = Dict{Int, Union(Worker, LocalProcess)}()
const map_sock_wrkr = ObjectIdDict()
const map_del_wrkr = Set{Int}()

let next_pid = 2    # 1 is reserved for the client (always)
    global get_next_pid
    function get_next_pid()
        retval = next_pid
        next_pid += 1
        retval
    end
end

type ProcessGroup
    name::String
    workers::Array{Any,1}

    # global references
    refs::Dict

    ProcessGroup(w::Array{Any,1}) = new("pg-default", w, Dict())
end
const PGRP = ProcessGroup({})

get_bind_addr(pid::Integer) = get_bind_addr(worker_from_id(pid))
function get_bind_addr(w::Union(Worker, LocalProcess))
    if !isdefined(w, :bind_addr)
        if w.id != myid()
            w.bind_addr = remotecall_fetch(w.id, get_bind_addr, w.id)
        else
            error("LPROC.bind_addr not defined") # Should never happend since LPROC.bind_addr
                                           # is defined early on during process init.
        end
    end
    w.bind_addr
end

function add_workers(pg::ProcessGroup, ws::Array{Any,1})
    # NOTE: currently only node 1 can add new nodes, since nobody else
    # has the full list of address:port
    assert(LPROC.id == 1)
    for w in ws
        w.id = get_next_pid()
        register_worker(w)
        create_message_handler_loop(w.socket)
    end

    all_locs = map(x -> isa(x, Worker) ? (string(x.bind_addr), x.port, x.id, x.manage == manage_local_worker) : ("", 0, x.id, true), pg.workers)

    for w in ws
        send_msg_now(w, :join_pgrp, w.id, all_locs, w.manage == manage_local_worker)
    end
    for w in ws
        @schedule begin
            w.manage(w.id, w.config, :register)
        end
    end
    [w.id for w in ws]
end

myid() = LPROC.id

nprocs() = length(PGRP.workers)
function nworkers()
    n = nprocs()
    n == 1 ? 1 : n-1
end

procs() = Int[x.id for x in PGRP.workers]
function procs(pid::Integer)
    if myid() == 1
        if (pid == 1) || (map_pid_wrkr[pid].manage == manage_local_worker)
            Int[x.id for x in filter(w -> (w.id==1) || (w.manage == manage_local_worker), PGRP.workers)]
        else
            ipatpid = get_bind_addr(pid)
            Int[x.id for x in filter(w -> get_bind_addr(w) == ipatpid, PGRP.workers)]
        end
    else
        remotecall_fetch(1, procs, pid)
    end
end

function workers()
    allp = procs()
    if nprocs() == 1
       allp
    else
       filter(x -> x != 1, allp)
    end
end

rmprocset = Set()
function rmprocs(args...; waitfor = 0.0)
    # Only pid 1 can add and remove processes
    if myid() != 1
        error("only process 1 can add and remove processes")
    end

    global rmprocset
    empty!(rmprocset)

    for i in [args...]
        if i == 1
            warn("rmprocs: process 1 not removed")
        else
            if haskey(map_pid_wrkr, i)
                push!(rmprocset, i)
                remote_do(i, exit)
            end
        end
    end

    start = time()
    while (time() - start) < waitfor
        if length(rmprocset) == 0
            break;
        else
            sleep(0.1)
        end
    end

    ((waitfor > 0) && (length(rmprocset) > 0)) ? :timed_out : :ok
end

type ProcessExitedException <: Exception end

worker_from_id(i) = worker_from_id(PGRP, i)
function worker_from_id(pg::ProcessGroup, i)
#   Processes with pids > ours, have to connect to us. May not have happened. Wait for some time.
    if in(i, map_del_wrkr)
        throw(ProcessExitedException())
    end
    if myid()==1 && !haskey(map_pid_wrkr,i)
        error("no process with id $i exists")
    end
    start = time()
    while (!haskey(map_pid_wrkr, i) && ((time() - start) < 60.0))
        sleep(0.1)
        yield()
    end
    map_pid_wrkr[i]
end

function worker_id_from_socket(s)
    w = get(map_sock_wrkr, s, nothing)
    if isa(w,Worker)
        if is(s, w.socket) || is(s, w.sendbuf)
            return w.id
        end
    end
    if isa(s,IOStream) && fd(s)==-1
        # serializing to a local buffer
        return myid()
    end
    return -1
end

register_worker(w) = register_worker(PGRP, w)
function register_worker(pg, w)
    push!(pg.workers, w)
    map_pid_wrkr[w.id] = w
    if isa(w, Worker)
        map_sock_wrkr[w.socket] = w
        map_sock_wrkr[w.sendbuf] = w
    end
end

deregister_worker(pid) = deregister_worker(PGRP, pid)
function deregister_worker(pg, pid)
    pg.workers = filter(x -> !(x.id == pid), pg.workers)
    w = pop!(map_pid_wrkr, pid, nothing)
    if isa(w, Worker)
        pop!(map_sock_wrkr, w.socket)
        pop!(map_sock_wrkr, w.sendbuf)

        # Notify the cluster manager of this workers death
        if myid() == 1
            w.manage(w.id, w.config, :deregister)
        end
    end
    push!(map_del_wrkr, pid)

    # delete this worker from our RemoteRef client sets
    ids = {}
    tonotify = {}
    for (id,rv) in pg.refs
        if in(pid,rv.clientset)
            push!(ids, id)
        end
        if rv.waitingfor == pid
            push!(tonotify, (id,rv))
        end
    end
    for id in ids
        del_client(pg, id, pid)
    end

    # throw exception to tasks waiting for this pid
    for (id,rv) in tonotify
        notify_error(rv.full, ProcessExitedException())
        delete!(pg.refs, id)
    end
end

## remote refs ##

const client_refs = WeakKeyDict()

type RemoteRef
    where::Int
    whence::Int
    id::Int
    # TODO: cache value if it's fetched, but don't serialize the cached value

    function RemoteRef(w, wh, id)
        r = new(w,wh,id)
        found = getkey(client_refs, r, false)
        if !is(found,false)
            return found
        end
        client_refs[r] = true
        finalizer(r, send_del_client)
        r
    end

    REQ_ID::Int = 0
    function RemoteRef(pid::Integer)
        rr = RemoteRef(pid, myid(), REQ_ID)
        REQ_ID += 1
        if mod(REQ_ID,200) == 0
            # force gc after making a lot of refs since they take up
            # space on the machine where they're stored, yet the client
            # is responsible for freeing them.
            gc()
        end
        rr
    end

    RemoteRef(w::LocalProcess) = RemoteRef(w.id)
    RemoteRef(w::Worker) = RemoteRef(w.id)
    RemoteRef() = RemoteRef(myid())

    global next_id
    next_id() = (id=(myid(),REQ_ID); REQ_ID+=1; id)
end

hash(r::RemoteRef, h::Uint) = hash(r.whence, hash(r.id, h))
==(r::RemoteRef, s::RemoteRef) = (r.whence==s.whence && r.id==s.id)

rr2id(r::RemoteRef) = (r.whence, r.id)

lookup_ref(id) = lookup_ref(PGRP, id)
function lookup_ref(pg, id)
    rv = get(pg.refs, id, false)
    if rv === false
        # first we've heard of this ref
        rv = RemoteValue()
        pg.refs[id] = rv
        push!(rv.clientset, id[1])
    end
    rv
end

# is a ref uninitialized? (for locally-owned refs only)
#function ref_uninitialized(id)
#    wi = lookup_ref(id)
#    !wi.done && is(wi.thunk,bottom_func)
#end
#ref_uninitialized(r::RemoteRef) = (assert(r.where==myid());
#                                   ref_uninitialized(rr2id(r)))

function isready(rr::RemoteRef)
    rid = rr2id(rr)
    if rr.where == myid()
        lookup_ref(rid).done
    else
        remotecall_fetch(rr.where, id->lookup_ref(id).done, rid)
    end
end

del_client(id, client) = del_client(PGRP, id, client)
function del_client(pg, id, client)
    rv = lookup_ref(id)
    delete!(rv.clientset, client)
    if isempty(rv.clientset)
        delete!(pg.refs, id)
        #print("$(myid()) collected $id\n")
    end
    nothing
end

function del_clients(pairs::(Any,Any)...)
    for p in pairs
        del_client(p[1], p[2])
    end
end

any_gc_flag = Condition()
function start_gc_msgs_task()
    @schedule while true
        wait(any_gc_flag)
        flush_gc_msgs()
    end
end

function send_del_client(rr::RemoteRef)
    if rr.where == myid()
        del_client(rr2id(rr), myid())
    else
        if in(rr.where, map_del_wrkr)
            # for a removed worker, don't bother
            return
        end
        w = worker_from_id(rr.where)
        push!(w.del_msgs, (rr2id(rr), myid()))
        w.gcflag = true
        notify(any_gc_flag)
    end
end

function add_client(id, client)
    #println("$(myid()) adding client $client to $id")
    rv = lookup_ref(id)
    push!(rv.clientset, client)
    nothing
end

function add_clients(pairs::(Any,Any)...)
    for p in pairs
        add_client(p[1], p[2])
    end
end

function send_add_client(rr::RemoteRef, i)
    if rr.where == myid()
        add_client(rr2id(rr), i)
    elseif i != rr.where
        # don't need to send add_client if the message is already going
        # to the processor that owns the remote ref. it will add_client
        # itself inside deserialize().
        w = worker_from_id(rr.where)
        #println("$(myid()) adding $((rr2id(rr), i)) for $(rr.where)")
        push!(w.add_msgs, (rr2id(rr), i))
        w.gcflag = true
        notify(any_gc_flag)
    end
end

function serialize(s, rr::RemoteRef)
    i = worker_id_from_socket(s)
    #println("$(myid()) serializing $rr to $i")
    if i != -1
        #println("send add $rr to $i")
        send_add_client(rr, i)
    end
    invoke(serialize, (Any, Any), s, rr)
end

function deserialize(s, t::Type{RemoteRef})
    rr = invoke(deserialize, (Any, DataType), s, t)
    where = rr.where
    if where == myid()
        add_client(rr2id(rr), myid())
    end
    # call ctor to make sure this rr gets added to the client_refs table
    RemoteRef(where, rr.whence, rr.id)
end

# data stored by the owner of a RemoteRef
type RemoteValue
    done::Bool
    result
    full::Condition   # waiting for a value
    empty::Condition  # waiting for value to be removed
    clientset::IntSet
    waitingfor::Int   # processor we need to hear from to fill this, or 0

    RemoteValue() = new(false, nothing, Condition(), Condition(), IntSet(), 0)
end

function work_result(rv::RemoteValue)
    v = rv.result
    if isa(v,WeakRef)
        v = v.value
    end
    v
end

function wait_full(rv::RemoteValue)
    while !rv.done
        wait(rv.full)
    end
    return work_result(rv)
end

function wait_empty(rv::RemoteValue)
    while rv.done
        wait(rv.empty)
    end
    return nothing
end

## core messages: do, call, fetch, wait, ref, put! ##

function run_work_thunk(thunk)
    local result
    try
        result = thunk()
    catch err
        print(STDERR, "exception on ", myid(), ": ")
        display_error(err,catch_backtrace())
        result = err
    end
    result
end
function run_work_thunk(rv::RemoteValue, thunk)
    put!(rv, run_work_thunk(thunk))
    nothing
end

function schedule_call(rid, thunk)
    rv = RemoteValue()
    (PGRP::ProcessGroup).refs[rid] = rv
    push!(rv.clientset, rid[1])
    schedule(@task(run_work_thunk(rv,thunk)))
    rv
end

#localize_ref(b::Box) = Box(localize_ref(b.contents))

#function localize_ref(r::RemoteRef)
#    if r.where == myid()
#        fetch(r)
#    else
#        r
#    end
#end

#localize_ref(x) = x

# make a thunk to call f on args in a way that simulates what would happen if
# the function were sent elsewhere
function local_remotecall_thunk(f, args)
    if isempty(args)
        return f
    end
    return ()->f(args...)

    # TODO: this seems to be capable of causing deadlocks by waiting on
    # Refs buried inside the closure that we don't want to wait on yet.
    # linfo = ccall(:jl_closure_linfo, Any, (Any,), f)
    # if isa(linfo,LambdaStaticData)
    #     env = ccall(:jl_closure_env, Any, (Any,), f)
    #     buf = memio()
    #     serialize(buf, env)
    #     seek(buf, 0)
    #     env = deserialize(buf)
    #     f = ccall(:jl_new_closure, Any, (Ptr{Void}, Any, Any),
    #               C_NULL, env, linfo)::Function
    # end
    # f(map(localize_ref,args)...)
end

function remotecall(w::LocalProcess, f, args...)
    rr = RemoteRef(w)
    schedule_call(rr2id(rr), local_remotecall_thunk(f,args))
    rr
end

function remotecall(w::Worker, f, args...)
    rr = RemoteRef(w)
    #println("$(myid()) asking for $rr")
    send_msg(w, :call, rr2id(rr), f, args)
    rr
end

remotecall(id::Integer, f, args...) = remotecall(worker_from_id(id), f, args...)

# faster version of fetch(remotecall(...))
function remotecall_fetch(w::LocalProcess, f, args...)
    run_work_thunk(local_remotecall_thunk(f,args))
end

function remotecall_fetch(w::Worker, f, args...)
    # can be weak, because the program will have no way to refer to the Ref
    # itself, it only gets the result.
    oid = next_id()
    rv = lookup_ref(oid)
    rv.waitingfor = w.id
    send_msg(w, :call_fetch, oid, f, args)
    v = wait_full(rv)
    delete!(PGRP.refs, oid)
    v
end

remotecall_fetch(id::Integer, f, args...) =
    remotecall_fetch(worker_from_id(id), f, args...)

# faster version of wait(remotecall(...))
remotecall_wait(w::LocalProcess, f, args...) = wait(remotecall(w,f,args...))

function remotecall_wait(w::Worker, f, args...)
    prid = next_id()
    rv = lookup_ref(prid)
    rv.waitingfor = w.id
    rr = RemoteRef(w)
    send_msg(w, :call_wait, rr2id(rr), prid, f, args)
    wait_full(rv)
    delete!(PGRP.refs, prid)
    rr
end

remotecall_wait(id::Integer, f, args...) =
    remotecall_wait(worker_from_id(id), f, args...)

function remote_do(w::LocalProcess, f, args...)
    # the LocalProcess version just performs in local memory what a worker
    # does when it gets a :do message.
    # same for other messages on LocalProcess.
    thk = local_remotecall_thunk(f, args)
    schedule(Task(thk))
    nothing
end

function remote_do(w::Worker, f, args...)
    send_msg(w, :do, f, args)
    nothing
end

remote_do(id::Integer, f, args...) = remote_do(worker_from_id(id), f, args...)

# have the owner of rr call f on it
function call_on_owner(f, rr::RemoteRef, args...)
    rid = rr2id(rr)
    if rr.where == myid()
        f(rid, args...)
    else
        remotecall_fetch(rr.where, f, rid, args...)
    end
end

wait_ref(rid) = (wait_full(lookup_ref(rid)); nothing)
wait(r::RemoteRef) = (call_on_owner(wait_ref, r); r)

fetch_ref(rid) = wait_full(lookup_ref(rid))
fetch(r::RemoteRef) = call_on_owner(fetch_ref, r)
fetch(x::ANY) = x

# storing a value to a Ref
function put!(rv::RemoteValue, val::ANY)
    wait_empty(rv)
    rv.result = val
    rv.done = true
    notify_full(rv)
    rv
end

put_ref(rid, v) = put!(lookup_ref(rid), v)
put!(rr::RemoteRef, val::ANY) = (call_on_owner(put_ref, rr, val); rr)

function take!(rv::RemoteValue)
    wait_full(rv)
    val = rv.result
    rv.done = false
    rv.result = nothing
    notify_empty(rv)
    val
end

take_ref(rid) = take!(lookup_ref(rid))
take!(rr::RemoteRef) = call_on_owner(take_ref, rr)

function deliver_result(sock::IO, msg, oid, value)
    #print("$(myid()) sending result $oid\n")
    if is(msg,:call_fetch)
        val = value
    else
        val = oid
    end
    try
        send_msg_now(sock, :result, oid, val)
    catch e
        # terminate connection in case of serialization error
        # otherwise the reading end would hang
        print(STDERR, "fatal error on ", myid(), ": ")
        display_error(e, catch_backtrace())
        wid = worker_id_from_socket(sock)
        close(sock)
        if myid()==1
            rmprocs(wid)
        elseif wid == 1
            exit(1)
        else
            remote_do(1, rmprocs, wid)
        end
    end
end

# notify waiters that a certain job has finished or Ref has been emptied
notify_full (rv::RemoteValue) = notify(rv.full, work_result(rv))
notify_empty(rv::RemoteValue) = notify(rv.empty)

## message event handlers ##

# activity on accept fd
function accept_handler(server::TcpServer, status::Int32)
    if status == -1
        error("an error occured during the creation of the server")
    end
    client = accept_nonblock(server)
    create_message_handler_loop(client)
end

function create_message_handler_loop(sock::AsyncStream) #returns immediately
    schedule(@task begin
        global PGRP
        #println("message_handler_loop")
        disable_nagle(sock)
        start_reading(sock)
        wait_connected(sock)

        try
            while true
                msg = deserialize(sock)
                # println("got msg: ",msg)
                # handle message
                if is(msg, :call)
                    id = deserialize(sock)
                    #print("$(myid()) got id $id\n")
                    f0 = deserialize(sock)
                    #print("$(myid()) got call $f0\n")
                    args0 = deserialize(sock)
                    #print("$(myid()) got args $args0\n")
                    let f=f0, args=args0
                        schedule_call(id, ()->f(args...))
                    end
                elseif is(msg, :call_fetch)
                    id = deserialize(sock)
                    f = deserialize(sock)
                    args = deserialize(sock)
                    @schedule begin
                        v = run_work_thunk(()->f(args...))
                        deliver_result(sock, msg, id, v)
                        v
                    end
                elseif is(msg, :call_wait)
                    id = deserialize(sock)
                    notify_id = deserialize(sock)
                    f = deserialize(sock)
                    args = deserialize(sock)
                    @schedule begin
                        rv = schedule_call(id, ()->f(args...))
                        deliver_result(sock, msg, notify_id, wait_full(rv))
                    end
                elseif is(msg, :do)
                    f = deserialize(sock)
                    args = deserialize(sock)
                    #print("got args: $args\n")
                    @schedule begin
                        run_work_thunk(RemoteValue(), ()->f(args...))
                    end
                elseif is(msg, :result)
                    # used to deliver result of wait or fetch
                    oid = deserialize(sock)
                    #print("$(myid()) got $msg $oid\n")
                    val = deserialize(sock)
                    put!(lookup_ref(oid), val)
                elseif is(msg, :identify_socket)
                    otherid = deserialize(sock)
                    register_worker(Worker("", 0, sock, otherid))
                elseif is(msg, :join_pgrp)
                    # first connection; get process group info from client
                    self_pid = LPROC.id = deserialize(sock)
                    locs = deserialize(sock)
                    self_is_local = deserialize(sock)
                    #print("\nLocation: ",locs,"\nId:",myid(),"\n")
                    # joining existing process group

                    register_worker(Worker("", 0, sock, 1))
                    register_worker(LPROC)

                    for (rhost, rport, rpid, r_is_local) in locs
                        if (rpid < self_pid) && (!(rpid == 1))
                            # Connect to them
                            if self_is_local && r_is_local
                                # If on localhost, use the loopback address - this addresses
                                # the special case of system suspend wherein the local ip
                                # may be changed upon system awake.
                                w = Worker("127.0.0.1", rport)
                            else
                                w = Worker(rhost, rport)
                            end
                            w.id = rpid
                            register_worker(w)
                            create_message_handler_loop(w.socket)
                            send_msg_now(w, :identify_socket, self_pid)
                        else
                            # Others will connect to us. Don't do anything just yet
                            continue
                        end
                    end
                end

            end # end of while
        catch e
            iderr = worker_id_from_socket(sock)
            # If error occured talking to pid 1, commit harakiri
            if iderr == 1
                if isopen(sock)
                    print(STDERR, "fatal error on ", myid(), ": ")
                    display_error(e, catch_backtrace())
                end
                exit(1)
            end

            # Will treat any exception as death of node and cleanup
            # since currently we do not have a mechanism for workers to reconnect
            # to each other on unhandled errors
            deregister_worker(iderr)

            if isopen(sock) close(sock) end

            if (myid() == 1)
                global rmprocset
                if in(iderr, rmprocset)
                    delete!(rmprocset, iderr)
                else
                    println("Worker $iderr terminated.")
                    rethrow(e)
                end
            end

            return nothing
        end
    end)
end

function disable_threaded_libs()
    blas_set_num_threads(1)
end

## worker creation and setup ##

# the entry point for julia worker processes. does not return.
# argument is descriptor to write listening port # to.
start_worker() = start_worker(STDOUT)
function start_worker(out::IO)
    # we only explicitly monitor worker STDOUT on the console, so redirect
    # stderr to stdout so we can see the output.
    # at some point we might want some or all worker output to go to log
    # files instead.
    # Currently disabled since this caused processes to spin instead of
    # exit when process 1 shut down. Don't yet know why.
    #redirect_stderr(STDOUT)

    (actual_port,sock) = listenany(uint16(9009))
    sock.ccb = accept_handler
    print(out, "julia_worker:")  # print header
    print(out, "$(dec(actual_port))#") # print port
    print(out, LPROC.bind_addr)
    print(out, '\n')
    flush(out)
    # close STDIN; workers will not use it
    #close(STDIN)

    disable_threaded_libs()
    disable_nagle(sock)

    try
        check_master_connect(60.0)
        while true; wait(); end
    catch err
        print(STDERR, "unhandled exception on $(myid()): $(err)\nexiting.\n")
    end

    close(sock)
    exit(0)
end

function start_cluster_workers(np::Integer, config::Dict, cman::ClusterManager)
    ws = cell(np)

    # Get the cluster manager to launch the instance
    (insttype, instances) = cman.launch(cman, np, config)

    if insttype == :io_only
        read_cb_response(inst) =
        begin
            (host, port) = read_worker_host_port(inst[1])
            inst[1], host, port, host, inst[2]
        end
    elseif  insttype == :io_host
        read_cb_response(inst) =
        begin
            (bind_addr, port) = read_worker_host_port(inst[1])
            inst[1], bind_addr, port, inst[2], inst[3]
        end
    elseif insttype == :io_host_port
        read_cb_response(inst) = (inst[1], inst[2], inst[3], inst[2], inst[4])
    elseif insttype == :host_port
        read_cb_response(inst) = (nothing, inst[1], inst[2], inst[1], inst[3])
    else
        error("unsupported format from ClusterManager callback")
    end

    for i=1:np
        (io, bind_addr, port, pubhost, wconfig) = read_cb_response(instances[i])
        ws[i] = create_worker(bind_addr, port, pubhost, io, wconfig, cman.manage)
    end
    ws
end

function read_worker_host_port(io::IO)
    io.line_buffered = true
    while true
        conninfo = readline(io)
        bind_addr, port = parse_connection_info(conninfo)
        if bind_addr != ""
            return bind_addr, port
        end
    end
end

function create_worker(bind_addr, port, pubhost, stream, config, manage)
    tunnel = config[:tunnel]

    s = split(pubhost,'@')
    user = ""
    if length(s) > 1
        user = s[1]
        pubhost = s[2]
    else
        if haskey(ENV, "USER")
            user = ENV["USER"]
        elseif tunnel
            error("USER must be specified either in the environment or as part of the hostname when tunnel option is used")
        end
    end

    if tunnel
        sshflags = config[:sshflags]
        w = Worker(pubhost, bind_addr, port, user, sshflags)
    else
        w = Worker(bind_addr, port)
    end

    w.config = config
    w.manage = manage

    if isa(stream, AsyncStream)
        let wrker = w
            # redirect console output from workers to the client's stdout:
            @async begin
                while !eof(stream)
                    line = readline(stream)
                    print("\tFrom worker $(wrker.id):\t$line")
                end
            end
        end
    end

    # install a finalizer to perform cleanup if necessary
    finalizer(w, (w)->if myid() == 1 w.manage(w.id, w.config, :finalize) end)

    w
end

function parse_connection_info(str)
    m = match(r"^julia_worker:(\d+)#(.*)", str)
    if m != nothing
        (m.captures[2], parseint(Int16, m.captures[1]))
    else
        ("", int16(-1))
    end
end

tunnel_port = 9201
# establish an SSH tunnel to a remote worker
# returns P such that localhost:P connects to host:port
function ssh_tunnel(user, host, bind_addr, port, sshflags)
    global tunnel_port
    localp = tunnel_port::Int
    while !success(detach(`ssh -f -o ExitOnForwardFailure=yes $sshflags $(user)@$host -L $localp:$bind_addr:$(int(port)) sleep 60`)) && localp < 10000
        localp += 1
    end

    if localp >= 10000
        error("unable to assign a local tunnel port between 9201 and 10000")
    end

    tunnel_port = localp+1
    localp
end

immutable LocalManager <: ClusterManager
    launch::Function
    manage::Function

    LocalManager() = new(launch_local_workers, manage_local_worker)
end

show(io::IO, cman::LocalManager) = println("LocalManager()")

function launch_local_workers(cman::LocalManager, np::Integer, config::Dict)
    dir = config[:dir]
    exename = config[:exename]
    exeflags = config[:exeflags]

    io_objs = cell(np)
    configs = cell(np)

    # start the processes first...
    for i in 1:np
        io, pobj = open(detach(`$(dir)/$(exename) $exeflags --bind-to $(LPROC.bind_addr)`), "r")
        io_objs[i] = io
        configs[i] = merge(config, {:process => pobj})
    end

    # ...and then read the host:port info. This optimizes overall start times.
    return (:io_only, collect(zip(io_objs, configs)))
end

function manage_local_worker(id::Integer, config::Dict, op::Symbol)
    if op == :interrupt
        kill(config[:process], 2)
    end
end

immutable SSHManager <: ClusterManager
    launch::Function
    manage::Function
    machines::AbstractVector

    SSHManager(; machines=[]) = new(launch_ssh_workers, manage_ssh_worker, machines)
end

show(io::IO, cman::SSHManager) = println("SSHManager(machines=", cman.machines, ")")

function launch_ssh_workers(cman::SSHManager, np::Integer, config::Dict)
    dir = config[:dir]
    exename = config[:exename]
    exeflags_base = config[:exeflags]

    io_objs = cell(np)
    configs = cell(np)

    # start the processes first...

    for i in 1:np
        # machine could be of the format [user@]host[:port] bind_addr
        machine_bind = split(cman.machines[i])
        if length(machine_bind) > 1
            exeflags = `--bind-to $(machine_bind[2]) $exeflags_base`
        else
            exeflags = exeflags_base
        end
        machine_def = machine_bind[1]

        machine_def = split(machine_def, ':')
        portopt = length(machine_def) == 2 ? ` -p $(machine_def[2]) ` : ``
        config[:sshflags] = `$(config[:sshflags]) $portopt`

        sshflags = config[:sshflags]
        host = cman.machines[i] = machine_def[1]

        # Build up the ssh command
        cmd = `cd $dir && $exename $exeflags`          # launch julia
        cmd = `sh -l -c $(shell_escape(cmd))`                   # shell to launch under
        cmd = `ssh -n $sshflags $host $(shell_escape(cmd))`     # use ssh to remote launch

        io, pobj = open(detach(cmd), "r")
        io_objs[i] = io
        configs[i] = merge(config, {:machine => cman.machines[i]})
    end

    # ...and then read the host:port info. This optimizes overall start times.
    # For ssh, the tunnel connection, if any, has to be with the specified machine name.
    # but the port needs to be forwarded to the bound hostname/ip-address
    return (:io_host, collect(zip(io_objs, cman.machines, configs)))
end

function manage_ssh_worker(id::Integer, config::Dict, op::Symbol)
    if op == :interrupt
        if haskey(config, :ospid)
            machine = config[:machine]
            if !success(`ssh -n $(config[:sshflags]) $machine "kill -2 $(config[:ospid])"`)
                println("Error sending a Ctrl-C to julia worker $id on $machine")
            end
        else
            # This state can happen immediately after an addprocs
            println("Worker $id cannot be presently interrupted.")
        end
    elseif op == :register
        config[:ospid] = remotecall_fetch(id, getpid)
    end
end

# start and connect to processes via SSH.
# optionally through an SSH tunnel.
# the tunnel is only used from the head (process 1); the nodes are assumed
# to be mutually reachable without a tunnel, as is often the case in a cluster.
function addprocs_internal(np::Integer;
                  tunnel=false, dir=JULIA_HOME,
                  exename=(ccall(:jl_is_debugbuild,Cint,())==0?"./julia":"./julia-debug"),
                  sshflags::Cmd=``, cman=LocalManager(), exeflags=``)

    config={:dir=>dir, :exename=>exename, :exeflags=>`$exeflags --worker`, :tunnel=>tunnel, :sshflags=>sshflags}
    disable_threaded_libs()
    add_workers(PGRP, start_cluster_workers(np, config, cman))
end

addprocs(np::Integer; kwargs...) = addprocs_internal(np; kwargs...)

function addprocs(machines::AbstractVector; kwargs...)
    cman_defined = any(x -> begin k,v = x; k==:cman end, kwargs)
    if cman_defined
        error("custom cluster managers unsupported on the ssh interface")
    else
        addprocs_internal(length(machines); cman=SSHManager(machines=machines), kwargs...)
    end
end

## higher-level functions: spawn, pmap, pfor, etc. ##

let nextidx = 1
    global chooseproc
    function chooseproc(thunk::Function)
        p = -1
        env = thunk.env
        if isa(env,Tuple)
            for v in env
                if isa(v,Box)
                    v = v.contents
                end
                if isa(v,RemoteRef)
                    p = v.where; break
                end
            end
        end
        if p == -1
            p = workers()[(nextidx % nworkers()) + 1]
            nextidx += 1
        end
        p
    end
end

spawnat(p, thunk) = sync_add(remotecall(p, thunk))

spawn_somewhere(thunk) = spawnat(chooseproc(thunk),thunk)

macro spawn(expr)
    expr = localize_vars(:(()->($expr)), false)
    :(spawn_somewhere($(esc(expr))))
end

macro spawnat(p, expr)
    expr = localize_vars(:(()->($expr)), false)
    :(spawnat($(esc(p)), $(esc(expr))))
end

macro fetch(expr)
    expr = localize_vars(:(()->($expr)), false)
    quote
        thunk = $(esc(expr))
        remotecall_fetch(chooseproc(thunk), thunk)
    end
end

macro fetchfrom(p, expr)
    expr = localize_vars(:(()->($expr)), false)
    :(remotecall_fetch($(esc(p)), $(esc(expr))))
end

function at_each(f, args...)
    for w in PGRP.workers
        sync_add(remotecall(w.id, f, args...))
    end
end

macro everywhere(ex)
    quote
        @sync begin
            at_each(()->eval(Main,$(Expr(:quote,ex))))
        end
    end
end

function pmap_static(f, lsts...)
    np = nprocs()
    n = length(lsts[1])
    { remotecall(PGRP.workers[(i-1)%np+1].id, f, map(L->L[i], lsts)...) for i = 1:n }
end

pmap(f) = f()

# dynamic scheduling by creating a local task to feed work to each processor
# as it finishes.
# example unbalanced workload:
# rsym(n) = (a=rand(n,n);a*a')
# L = {rsym(200),rsym(1000),rsym(200),rsym(1000),rsym(200),rsym(1000),rsym(200),rsym(1000)};
# pmap(eig, L);
function pmap(f, lsts...; err_retry=true, err_stop=false)
    len = length(lsts)

    results = Dict{Int,Any}()

    retryqueue = {}
    task_in_err = false
    is_task_in_error() = task_in_err
    set_task_in_error() = (task_in_err = true)

    nextidx = 0
    getnextidx() = (nextidx += 1)

    states = [start(lsts[idx]) for idx in 1:len]
    function getnext_tasklet()
        if is_task_in_error() && err_stop
            return nothing
        elseif !any(idx->done(lsts[idx],states[idx]), 1:len)
            nxts = [next(lsts[idx],states[idx]) for idx in 1:len]
            for idx in 1:len; states[idx] = nxts[idx][2]; end
            nxtvals = [x[1] for x in nxts]
            return (getnextidx(), nxtvals)

        elseif !isempty(retryqueue)
            return shift!(retryqueue)
        else
            return nothing
        end
    end

    @sync begin
        for wpid in workers()
            @async begin
                tasklet = getnext_tasklet()
                while (tasklet != nothing)
                    (idx, fvals) = tasklet
                    try
                        result = remotecall_fetch(wpid, f, fvals...)
                        if isa(result, Exception)
                            ((wpid == myid()) ? rethrow(result) : throw(result))
                        else
                            results[idx] = result
                        end
                    catch ex
                        if err_retry
                            push!(retryqueue, (idx,fvals, ex))
                        else
                            results[idx] = ex
                        end
                        set_task_in_error()
                        break # remove this worker from accepting any more tasks
                    end

                    tasklet = getnext_tasklet()
                end
            end
        end
    end

    for failure in retryqueue
        results[failure[1]] = failure[3]
    end
    [results[x] for x in 1:nextidx]
end

# Statically split range [1,N] into equal sized chunks for np processors
function splitrange(N::Int, np::Int)
    each = div(N,np)
    extras = rem(N,np)
    nchunks = each > 0 ? np : extras
    chunks = Array(UnitRange{Int}, nchunks)
    lo = 1
    for i in 1:nchunks
        hi = lo + each - 1
        if extras > 0
            hi += 1
            extras -= 1
        end
        chunks[i] = lo:hi
        lo = hi+1
    end
    return chunks
end

function preduce(reducer, f, N::Int)
    chunks = splitrange(N, nworkers())
    results = cell(length(chunks))
    for i in 1:length(chunks)
        results[i] = @spawn f(first(chunks[i]), last(chunks[i]))
    end
    mapreduce(fetch, reducer, results)
end

function pfor(f, N::Int)
    for c in splitrange(N, nworkers())
        @spawn f(first(c), last(c))
    end
    nothing
end

function make_preduce_body(reducer, var, body, ran)
    localize_vars(
    quote
        function (lo::Int, hi::Int)
            R = $(esc(ran))
            $(esc(var)) = R[lo]
            ac = $(esc(body))
            if lo != hi
                for $(esc(var)) in R[(lo+1):hi]
                    ac = ($(esc(reducer)))(ac, $(esc(body)))
                end
            end
            ac
        end
    end
                  )
end

function make_pfor_body(var, body, ran)
    localize_vars(
    quote
        function (lo::Int, hi::Int)
            for $(esc(var)) in ($(esc(ran)))[lo:hi]
                $(esc(body))
            end
        end
    end
                  )
end

macro parallel(args...)
    na = length(args)
    if na==1
        loop = args[1]
        if isa(loop,Expr) && loop.head === :comprehension
            ex = loop.args[1]
            loop.args[1] = esc(ex)
            nd = length(loop.args)-1
            ranges = map(e->esc(e.args[2]), loop.args[2:end])
            for i=1:nd
                var = loop.args[1+i].args[1]
                loop.args[1+i] = :( $(esc(var)) = ($(ranges[i]))[I[$i]] )
            end
            return :( DArray((I::(UnitRange{Int}...))->($loop),
                             tuple($(map(r->:(length($r)),ranges)...))) )
        end
    elseif na==2
        reducer = args[1]
        loop = args[2]
    else
        throw(ArgumentError("wrong number of arguments to @parallel"))
    end
    if !isa(loop,Expr) || !is(loop.head,:for)
        error("malformed @parallel loop")
    end
    var = loop.args[1].args[1]
    r = loop.args[1].args[2]
    body = loop.args[2]
    if na==1
        quote
            pfor($(make_pfor_body(var, body, r)), length($(esc(r))))
        end
    else
        quote
            preduce($(esc(reducer)),
                    $(make_preduce_body(reducer, var, body, r)), length($(esc(r))))
        end
    end
end

## demos ##

# fv(a)=eig(a)[2][2]
# A=randn(800,800);A=A*A';
# pmap(fv, {A,A,A})

#all2all() = at_each(hello_from, myid())

#hello_from(i) = print("message from $i to $(myid())\n")

# monte carlo estimate of pi
# function buffon(niter)
#     nc =
#     @parallel (+) for i=1:niter
#         rand() <= sin(rand()*pi/2) ? 1 : 0
#     end
#     2/(nc/niter)
# end

function check_master_connect(timeout)
    # If we do not have at least process 1 connect to us within timeout
    # we log an error and exit
    @schedule begin
        start = time()
        while !haskey(map_pid_wrkr, 1) && (time() - start) < timeout
            sleep(1.0)
        end

        if !haskey(map_pid_wrkr, 1)
            print(STDERR, "Master process (id 1) could not connect within $timeout seconds.\nexiting.\n")
            exit(1)
        end
    end
end

function timedwait(testcb::Function, secs::Float64; pollint::Float64=0.1)
    start = time()
    done = RemoteRef()
    timercb(aw) = begin
        try
            if testcb()
                put!(done, :ok)
            elseif (time() - start) > secs
                put!(done, :timed_out)
            end
        catch e
            put!(done, :error)
        finally
            isready(done) && stop_timer(aw)
        end
    end

    if !testcb()
        t = Timer(timercb)
        start_timer(t, pollint, pollint)
        ret = fetch(done)
        stop_timer(t)
    else
        ret = :ok
    end
    ret
end

function interrupt(pid::Integer)
    assert(myid() == 1)
    w = map_pid_wrkr[pid]
    if isa(w, Worker)
        w.manage(w.id, w.config, :interrupt)
    end
end
interrupt(pids::Integer...) = interrupt([pids...])

function interrupt(pids::AbstractVector=workers())
    assert(myid() == 1)
    @sync begin
        for pid in pids
            @async interrupt(pid)
        end
    end
end

function disable_nagle(sock)
    # disable nagle on all OSes
    ccall(:uv_tcp_nodelay, Cint, (Ptr{Void}, Cint), sock.handle, 1)
    @linux_only begin
        # tcp_quickack is a linux only option
        if ccall(:jl_tcp_quickack, Cint, (Ptr{Void}, Cint), sock.handle, 1) < 0
            warn_once("Parallel networking unoptimized ( Error enabling TCP_QUICKACK : ", strerror(errno()), " )")
        end
    end
end

function check_same_host(pids)
    if myid() != 1
        return remotecall_fetch(1, check_same_host, pids)
    else
        # We checkfirst if all test pids have been started using the local manager,
        # else we check for the same bind_to addr. This handles the special case
        # where the local ip address may change - as during a system sleep/awake
        if all(p -> (p==1) || (map_pid_wrkr[p].manage == manage_local_worker), pids)
            return true
        else
            first_bind_addr = map_pid_wrkr[pids[1]].bind_addr
            return all(p -> (p != 1) && (map_pid_wrkr[p].bind_addr == first_bind_addr), pids[2:end])
        end
    end
end

function terminate_all_workers()
    if myid() != 1
        return
    end

    if nprocs() > 1
        ret = rmprocs(workers(); waitfor=0.5)
        if ret != :ok
            warn("Forcibly interrupting busy workers")
            # Might be computation bound, interrupt them and try again
            interrupt(workers())
            ret = rmprocs(workers(); waitfor=0.5)
            if ret != :ok
                warn("Unable to terminate all workers")
            end
        end
     end
end

max_stars_repo_name: Aerlinger/julia

max_issues_repo_name: Aerlinger/julia

max_forks_repo_name: Aerlinger/julia


text:
# ============================================================
# Continuum Theory
# ============================================================

"""
    prim_conserve(prim::AbstractVector{T}, γ) where {T<:Real}
    prim_conserve(ρ, U, λ, γ)
    prim_conserve(ρ, U, V, λ, γ)
    prim_conserve(ρ, U, V, W, λ, γ)

Transform primitive -> conservative variables
"""
function prim_conserve(prim::AbstractVector{T}, γ) where {T<:Real}
    if eltype(prim) <: Integer
        W = similar(prim, Float64)
    else
        W = similar(prim)
    end

    if length(prim) == 3 # 1D
        W[1] = prim[1]
        W[2] = prim[1] * prim[2]
        W[3] = 0.5 * prim[1] / prim[3] / (γ - 1.0) + 0.5 * prim[1] * prim[2]^2
    elseif length(prim) == 4 # 2D
        W[1] = prim[1]
        W[2] = prim[1] * prim[2]
        W[3] = prim[1] * prim[3]
        W[4] = 0.5 * prim[1] / prim[4] / (γ - 1.0) + 0.5 * prim[1] * (prim[2]^2 + prim[3]^2)
    elseif length(prim) == 5 # 3D
        W[1] = prim[1]
        W[2] = prim[1] * prim[2]
        W[3] = prim[1] * prim[3]
        W[4] = prim[1] * prim[4]
        W[5] =
            0.5 * prim[1] / prim[5] / (γ - 1.0) +
            0.5 * prim[1] * (prim[2]^2 + prim[3]^2 + prim[4]^2)
    else
        throw("prim -> w : dimension error")
    end

    return W
end

prim_conserve(ρ, U, λ, γ) = prim_conserve([ρ, U, λ], γ)

prim_conserve(ρ, U, V, λ, γ) = prim_conserve([ρ, U, V, λ], γ)

prim_conserve(ρ, U, V, W, λ, γ) = prim_conserve([ρ, U, V, W, λ], γ)

#--- Rykov ---#
function prim_conserve(prim::AbstractVector{T}, γ, Kr) where {T<:Real}
    if eltype(prim) <: Integer
        W = similar(prim, Float64, length(prim) - 1)
    else
        W = similar(prim, length(prim) - 1)
    end

    if length(prim) == 5 # 1D
        W[1] = prim[1]
        W[2] = prim[1] * prim[2]
        W[3] = 0.5 * prim[1] / prim[3] / (γ - 1.0) + 0.5 * prim[1] * prim[2]^2
        W[4] = 0.25 * prim[1] * Kr / prim[5]
    elseif length(prim) == 6 # 2D
        W[1] = prim[1]
        W[2] = prim[1] * prim[2]
        W[3] = prim[1] * prim[3]
        W[4] = 0.5 * prim[1] / prim[4] / (γ - 1.0) + 0.5 * prim[1] * (prim[2]^2 + prim[3]^2)
        W[5] = 0.25 * prim[1] * Kr / prim[6]
    else
        throw("prim -> w : dimension error")
    end

    return W
end

"""
    mixture_prim_conserve(prim::AbstractMatrix{T}, γ) where {T<:Real}

Transform multi-component primitive -> conservative variables
"""
function mixture_prim_conserve(prim::AbstractMatrix{T}, γ) where {T<:Real}
    if eltype(prim) <: Integer
        w = similar(prim, Float64)
    else
        w = similar(prim)
    end

    for j in axes(w, 2)
        w[:, j] .= prim_conserve(prim[:, j], γ)
    end

    return w
end

"""
* scalar: pseudo primitive vector for scalar conservation laws

    conserve_prim(u)

    conserve_prim(u, a)

* vector: primitive vector for Euler, Navier-Stokes and extended equations

    conserve_prim(W::T, γ) where {T<:AbstractArray{<:Real,1}}

    conserve_prim(ρ, M, E, γ)

    conserve_prim(ρ, MX, MY, E, γ)

Transform conservative -> primitive variables
"""
conserve_prim(u) = [u, 0.5 * u, 1.0]

conserve_prim(u, a) = [u, a, 1.0]

function conserve_prim(W::AbstractVector{T}, γ) where {T<:Real}
    if eltype(W) <: Integer
        prim = similar(W, Float64)
    else
        prim = similar(W)
    end

    if length(W) == 3 # 1D
        prim[1] = W[1]
        prim[2] = W[2] / W[1]
        prim[3] = 0.5 * W[1] / (γ - 1.0) / (W[3] - 0.5 * W[2]^2 / W[1])
    elseif length(W) == 4 # 2D
        prim[1] = W[1]
        prim[2] = W[2] / W[1]
        prim[3] = W[3] / W[1]
        prim[4] = 0.5 * W[1] / (γ - 1.0) / (W[4] - 0.5 * (W[2]^2 + W[3]^2) / W[1])
    elseif length(W) == 5 # 3D
        prim[1] = W[1]
        prim[2] = W[2] / W[1]
        prim[3] = W[3] / W[1]
        prim[4] = W[4] / W[1]
        prim[5] = 0.5 * W[1] / (γ - 1.0) / (W[5] - 0.5 * (W[2]^2 + W[3]^2 + W[4]^2) / W[1])
    else
        throw("w -> prim : dimension dismatch")
    end

    return prim
end

conserve_prim(ρ, M, E, γ) = conserve_prim([ρ, M, E], γ)

conserve_prim(ρ, MX, MY, E, γ) = conserve_prim([ρ, MX, MY, E], γ)

conserve_prim(ρ, MX, MY, MZ, E, γ) = conserve_prim([ρ, MX, MY, MZ, E], γ)

#--- Rykov ---#
function conserve_prim(w::AbstractVector{T}, K, Kr) where {T<:Real}
    if eltype(w) <: Integer
        prim = similar(w, Float64, length(w) + 1)
    else
        prim = similar(w, length(w) + 1)
    end

    if length(w) == 4 # 1D
        prim[1] = w[1]
        prim[2] = w[2] / w[1]
        prim[3] = 0.25 * w[1] * (K + Kr + 1.0) / (w[3] - 0.5 * w[2]^2 / w[1])
        prim[4] = 0.25 * w[1] * (K + 1.0) / (w[3] - w[4] - 0.5 * w[2]^2 / w[1])
        prim[5] = 0.25 * w[1] * Kr / w[4]
    elseif length(w) == 5 # 2D
        prim[1] = w[1]
        prim[2] = w[2] / w[1]
        prim[3] = w[3] / w[1]
        prim[4] = 0.25 * w[1] * (K + Kr + 2.0) / (w[4] - 0.5 * (w[2]^2 + w[3]^2) / w[1])
        prim[5] = 0.25 * w[1] * (K + 2.0) / (w[4] - w[5] - 0.5 * (w[2]^2 + w[3]^2) / w[1])
        prim[6] = 0.25 * w[1] * Kr / w[5]
    else
        throw("w -> prim : dimension dismatch")
    end

    return prim
end

"""
    mixture_conserve_prim(W::T, γ) where {T<:AbstractArray{<:Real,2}}

Transform multi-component conservative -> primitive variables
"""
function mixture_conserve_prim(W::AbstractMatrix{T}, γ) where {T<:Real}
    if eltype(W) <: Integer
        prim = similar(W, Float64)
    else
        prim = similar(W)
    end

    for j in axes(prim, 2)
        prim[:, j] .= conserve_prim(W[:, j], γ)
    end

    return prim
end

"""
    aap_hs_prim(
        prim::AbstractArray{<:Real,2},
        tau::AbstractArray{<:Real,1},
        mi::Real,
        ni::Real,
        me::Real,
        ne::Real,
        kn::Real,
    )

Calculate mixture primitive variables from AAP model
"""
function aap_hs_prim(
    prim::AbstractMatrix{T1},
    tau::AbstractVector{T2},
    mi,
    ni,
    me,
    ne,
    kn,
) where {T1<:Real,T2<:Real}

    mixprim = similar(prim)

    if size(prim, 1) == 3

        mixprim[1, :] = deepcopy(prim[1, :])
        mixprim[2, 1] =
            prim[2, 1] +
            tau[1] / kn *
            (
                4.0 * sqrt(2.0) / (3.0 * sqrt(π)) * prim[1, 2] / (ni + ne) / (mi + me) *
                sqrt(1.0 / prim[end, 1] + 1.0 / prim[end, 2])
            ) *
            (prim[2, 2] - prim[2, 1])
        mixprim[2, 2] =
            prim[2, 2] +
            tau[2] / kn *
            (
                4.0 * sqrt(2.0) / (3.0 * sqrt(π)) * prim[1, 1] / (ni + ne) / (mi + me) *
                sqrt(1.0 / prim[end, 1] + 1.0 / prim[end, 2])
            ) *
            (prim[2, 1] - prim[2, 2])
        mixprim[3, 1] =
            1.0 / (
                1.0 / prim[end, 1] - 2.0 / 3.0 * (mixprim[2, 1] - prim[2, 1])^2 +
                tau[1] / kn * 2.0 * mi / (mi + me) *
                (
                    4.0 * sqrt(2.0) / (3.0 * sqrt(π)) * prim[1, 2] / (ni + ne) / (mi + me) *
                    sqrt(1.0 / prim[end, 1] + 1.0 / prim[end, 2])
                ) *
                (
                    1.0 / prim[end, 2] * me / mi - 1.0 / prim[end, 1] +
                    2.0 / 3.0 * me / mi * (prim[2, 2] - prim[2, 1])^2
                )
            )
        mixprim[3, 2] =
            1.0 / (
                1.0 / prim[end, 2] - 2.0 / 3.0 * (mixprim[2, 2] - prim[2, 2])^2 +
                tau[2] / kn * 2.0 * me / (mi + me) *
                (
                    4.0 * sqrt(2.0) / (3.0 * sqrt(π)) * prim[1, 1] / (ni + ne) / (mi + me) *
                    sqrt(1.0 / prim[end, 1] + 1.0 / prim[end, 2])
                ) *
                (
                    1.0 / prim[end, 1] * mi / me - 1.0 / prim[end, 2] +
                    2.0 / 3.0 * mi / me * (prim[2, 1] - prim[2, 2])^2
                )
            )

    elseif size(prim, 1) == 4

        mixprim[1, :] = deepcopy(prim[1, :])
        mixprim[2, 1] =
            prim[2, 1] +
            tau[1] / kn *
            (
                4.0 * sqrt(2.0) / (3.0 * sqrt(π)) * prim[1, 2] / (ni + ne) / (mi + me) *
                sqrt(1.0 / prim[end, 1] + 1.0 / prim[end, 2])
            ) *
            (prim[2, 2] - prim[2, 1])
        mixprim[2, 2] =
            prim[2, 2] +
            tau[2] / kn *
            (
                4.0 * sqrt(2.0) / (3.0 * sqrt(π)) * prim[1, 1] / (ni + ne) / (mi + me) *
                sqrt(1.0 / prim[end, 1] + 1.0 / prim[end, 2])
            ) *
            (prim[2, 1] - prim[2, 2])
        mixprim[3, 1] =
            prim[3, 1] +
            tau[1] / kn *
            (
                4.0 * sqrt(2.0) / (3.0 * sqrt(π)) * prim[1, 2] / (ni + ne) / (mi + me) *
                sqrt(1.0 / prim[end, 1] + 1.0 / prim[end, 2])
            ) *
            (prim[3, 2] - prim[3, 1])
        mixprim[3, 2] =
            prim[3, 2] +
            tau[2] / kn *
            (
                4.0 * sqrt(2.0) / (3.0 * sqrt(π)) * prim[1, 1] / (ni + ne) / (mi + me) *
                sqrt(1.0 / prim[end, 1] + 1.0 / prim[end, 2])
            ) *
            (prim[3, 1] - prim[3, 2])
        mixprim[4, 1] =
            1.0 / (
                1.0 / prim[end, 1] - 2.0 / 3.0 * (mixprim[2, 1] - prim[2, 1])^2 -
                2.0 / 3.0 * (mixprim[3, 1] - prim[3, 1])^2 +
                tau[1] / kn * 2.0 * mi / (mi + me) *
                (
                    4.0 * sqrt(2.0) / (3.0 * sqrt(π)) * prim[1, 2] / (ni + ne) / (mi + me) *
                    sqrt(1.0 / prim[end, 1] + 1.0 / prim[end, 2])
                ) *
                (
                    1.0 / prim[end, 2] * me / mi - 1.0 / prim[end, 1] +
                    2.0 / 3.0 * me / mi * (prim[2, 2] - prim[2, 1])^2 +
                    2.0 / 3.0 * me / mi * (prim[3, 2] - prim[3, 1])^2
                )
            )
        mixprim[4, 2] =
            1.0 / (
                1.0 / prim[end, 2] - 2.0 / 3.0 * (mixprim[2, 2] - prim[2, 2])^2 -
                2.0 / 3.0 * (mixprim[3, 2] - prim[3, 2])^2 +
                tau[2] / kn * 2.0 * me / (mi + me) *
                (
                    4.0 * sqrt(2.0) / (3.0 * sqrt(π)) * prim[1, 1] / (ni + ne) / (mi + me) *
                    sqrt(1.0 / prim[end, 1] + 1.0 / prim[end, 2])
                ) *
                (
                    1.0 / prim[end, 1] * mi / me - 1.0 / prim[end, 2] +
                    2.0 / 3.0 * mi / me * (prim[2, 1] - prim[2, 2])^2 +
                    2.0 / 3.0 * mi / me * (prim[3, 1] - prim[3, 2])^2
                )
            )

    elseif size(prim, 1) == 5

        mixprim[1, :] = deepcopy(prim[1, :])
        mixprim[2, 1] =
            prim[2, 1] +
            tau[1] / kn *
            (
                4.0 * sqrt(2.0) / (3.0 * sqrt(π)) * prim[1, 2] / (ni + ne) / (mi + me) *
                sqrt(1.0 / prim[end, 1] + 1.0 / prim[end, 2])
            ) *
            (prim[2, 2] - prim[2, 1])
        mixprim[2, 2] =
            prim[2, 2] +
            tau[2] / kn *
            (
                4.0 * sqrt(2.0) / (3.0 * sqrt(π)) * prim[1, 1] / (ni + ne) / (mi + me) *
                sqrt(1.0 / prim[end, 1] + 1.0 / prim[end, 2])
            ) *
            (prim[2, 1] - prim[2, 2])
        mixprim[3, 1] =
            prim[3, 1] +
            tau[1] / kn *
            (
                4.0 * sqrt(2.0) / (3.0 * sqrt(π)) * prim[1, 2] / (ni + ne) / (mi + me) *
                sqrt(1.0 / prim[end, 1] + 1.0 / prim[end, 2])
            ) *
            (prim[3, 2] - prim[3, 1])
        mixprim[3, 2] =
            prim[3, 2] +
            tau[2] / kn *
            (
                4.0 * sqrt(2.0) / (3.0 * sqrt(π)) * prim[1, 1] / (ni + ne) / (mi + me) *
                sqrt(1.0 / prim[end, 1] + 1.0 / prim[end, 2])
            ) *
            (prim[3, 1] - prim[3, 2])
        mixprim[4, 1] =
            prim[4, 1] +
            tau[1] / kn *
            (
                4.0 * sqrt(2.0) / (3.0 * sqrt(π)) * prim[1, 2] / (ni + ne) / (mi + me) *
                sqrt(1.0 / prim[end, 1] + 1.0 / prim[end, 2])
            ) *
            (prim[4, 2] - prim[4, 1])
        mixprim[4, 2] =
            prim[4, 2] +
            tau[2] / kn *
            (
                4.0 * sqrt(2.0) / (3.0 * sqrt(π)) * prim[1, 1] / (ni + ne) / (mi + me) *
                sqrt(1.0 / prim[end, 1] + 1.0 / prim[end, 2])
            ) *
            (prim[4, 1] - prim[4, 2])
        mixprim[5, 1] =
            1.0 / (
                1.0 / prim[end, 1] - 2.0 / 3.0 * (mixprim[2, 1] - prim[2, 1])^2 -
                2.0 / 3.0 * (mixprim[3, 1] - prim[3, 1])^2 -
                2.0 / 3.0 * (mixprim[4, 1] - prim[4, 1])^2 +
                tau[1] / kn * 2.0 * mi / (mi + me) *
                (
                    4.0 * sqrt(2.0) / (3.0 * sqrt(π)) * prim[1, 2] / (ni + ne) / (mi + me) *
                    sqrt(1.0 / prim[end, 1] + 1.0 / prim[end, 2])
                ) *
                (
                    1.0 / prim[end, 2] * me / mi - 1.0 / prim[end, 1] +
                    2.0 / 3.0 * me / mi * (prim[2, 2] - prim[2, 1])^2 +
                    2.0 / 3.0 * me / mi * (prim[3, 2] - prim[3, 1])^2 +
                    2.0 / 3.0 * me / mi * (prim[4, 2] - prim[4, 1])^2
                )
            )
        mixprim[5, 2] =
            1.0 / (
                1.0 / prim[end, 2] - 2.0 / 3.0 * (mixprim[2, 2] - prim[2, 2])^2 -
                2.0 / 3.0 * (mixprim[3, 2] - prim[3, 2])^2 -
                2.0 / 3.0 * (mixprim[4, 2] - prim[4, 2])^2 +
                tau[2] / kn * 2.0 * me / (mi + me) *
                (
                    4.0 * sqrt(2.0) / (3.0 * sqrt(π)) * prim[1, 1] / (ni + ne) / (mi + me) *
                    sqrt(1.0 / prim[end, 1] + 1.0 / prim[end, 2])
                ) *
                (
                    1.0 / prim[end, 1] * mi / me - 1.0 / prim[end, 2] +
                    2.0 / 3.0 * mi / me * (prim[2, 1] - prim[2, 2])^2 +
                    2.0 / 3.0 * mi / me * (prim[3, 1] - prim[3, 2])^2 +
                    2.0 / 3.0 * mi / me * (prim[4, 1] - prim[4, 2])^2
                )
            )

    else

        throw("AAP mixture: dimension dismatch")

    end

    return mixprim

end

"""
    advection_flux(u, a)

Theoretical flux of linear advection equation
"""
advection_flux(u, a) = a * u

"""
    burgers_flux(u)

Theoretical flux of Burgers' equation
"""
burgers_flux(u) = 0.5 * u^2

"""
    euler_flux(w::AbstractVector{T}, γ; frame = :cartesian::Symbol) where {T<:Real}

Theoretical fluxes of Euler Equations

* @return: flux tuple

"""
function euler_flux(w::AbstractVector{T}, γ; frame = :cartesian::Symbol) where {T<:Real}
    prim = conserve_prim(w, γ)
    p = 0.5 * prim[1] / prim[end]

    if eltype(w) <: Integer
        F = similar(w, Float64)
        G = similar(w, Float64)
        H = similar(w, Float64)
    else
        F = similar(w)
        G = similar(w)
        H = similar(w)
    end

    if length(w) == 3
        F[1] = w[2]
        F[2] = w[2]^2 / w[1] + p
        F[3] = (w[3] + p) * w[2] / w[1]

        return (F,)
    elseif length(w) == 4
        F[1] = w[2]
        F[2] = w[2]^2 / w[1] + p
        F[3] = w[2] * w[3] / w[1]
        F[4] = (w[end] + p) * w[2] / w[1]

        G[1] = w[3]
        G[2] = w[3] * w[2] / w[1]
        G[3] = w[3]^2 / w[1] + p
        G[4] = (w[end] + p) * w[3] / w[1]

        return F, G
    elseif length(w) == 5
        F[1] = w[2]
        F[2] = w[2]^2 / w[1] + p
        F[3] = w[2] * w[3] / w[1]
        F[4] = w[2] * w[4] / w[1]
        F[5] = (w[end] + p) * w[2] / w[1]

        G[1] = w[3]
        G[2] = w[3] * w[2] / w[1]
        G[3] = w[3]^2 / w[1] + p
        G[4] = w[3] * w[4] / w[1]
        G[5] = (w[end] + p) * w[3] / w[1]

        H[1] = w[4]
        H[2] = w[4] * w[2] / w[1]
        H[3] = w[4] * w[3] / w[1]
        H[4] = w[4]^2 / w[1] + p
        H[5] = (w[end] + p) * w[4] / w[1]

        return F, G, H
    end
end

"""
    euler_jacobi(w::AbstractVector{T}, γ) where {T<:Real}

Flux Jacobian of Euler Equations

* @return: Jacobian matrix A

"""
function euler_jacobi(w::AbstractVector{T}, γ) where {T<:Real}
    if eltype(w) <: Integer
        A = similar(w, Float64, 3, 3)
    else
        A = similar(w, 3, 3)
    end

    A[1, 1] = 0.0
    A[1, 2] = 1.0
    A[1, 3] = 0.0

    A[2, 1] = 0.5 * (γ - 3.0) * (w[2] / w[1])^2
    A[2, 2] = (3.0 - γ) * w[2] / w[1]
    A[2, 3] = γ - 1.0

    A[3, 1] = -γ * w[3] * w[2] / w[1]^2 + (γ - 1.0) * (w[2] / w[1])^3
    A[3, 2] = γ * w[3] / w[1] - 1.5 * (γ - 1.0) * (w[2] / w[1])^2
    A[3, 3] = γ * w[2] / w[1]

    return A
end


text:
"""
# TarFile

TarFile provides methods to extract files in a tar archive into memory.

The main function is `extract`, which returns an array of UInt8 arrays,
each corresponding to each file in the archive, and a file name for each.

## Dependencies
This module relies on the python module `tarfile` to do all the work and simply
provides a convenient wrapper.  Hence you need:

- PyCall
- A python installation with `tarfile`.

## Using
```
julia> using TarFile

julia> content, names = extract("/path/to/tarfile.tar")
```

## Limitations
- TarFile only operates on tar files on disk at present.
"""
module TarFile

import Glob: FilenameMatch, @fn_str
using PyCall

export extract

@pyimport tarfile as TF

"""
    extract(file, pattern="*", verbose=false) -> file_contents::Array{Array{Uint8}}, file_names::Array{String}

Read a tape archive ('tar') file and return the contents as `file_contents` and
file paths as `file_names` to all the members of the archive.  If `pattern` is
specified, then only return files matching the pattern.  This can be a Regex
(specified by `r"..."`) or a file globbing pattern.

Archives with gzip or bzip2 compression are supported without any further action
required by the user.
"""
function extract(file, pattern::Union{Regex,FilenameMatch}=r""; verbose=false)
    # If this gets too slow, look at:
    #   https://blogs.it.ox.ac.uk/inapickle/2011/06/20/high-memory-usage-when-using-pythons-tarfile-module/
    verbose && info("Opening tar file '$file'")
    contents = Array{Array{UInt8,1},1}()
    names = Array{String,1}()
    f = TF.open(file)
    try
        for member in f
            filename = member[:path]
            if ismatch(pattern, filename) && member[:isfile]()
                handle = f[:extractfile](member)
                handle == nothing && continue
                verbose && info("  Extracting file $filename")
                push!(names, filename)
                push!(contents, handle[:read]())
            end
        end
    catch err
        verbose && info("  Cleaning up after encountering error:\n$err")
    finally
        f[:close]()
    end
    contents, names
end
extract(file, pattern::AbstractString; kw...) = extract(file, FilenameMatch(pattern); kw...)

end # module

max_stars_repo_name: anowacki/assorted-julia-modules

max_issues_repo_name: anowacki/assorted-julia-modules

max_forks_repo_name: anowacki/assorted-julia-modules


text:
using DFTK
using Printf
using HDF5
using Dates
import FFTW
import JSON
using LinearAlgebra
import DFTK: HybridMixing
import DFTK: ScfDiagtol, ScfDefaultCallback

function write_scfres(h5gr, scfres; store_wfn=false)
    for symbol in (:ρin, :ρ)
        hasproperty(scfres, symbol) || continue
        object = getproperty(scfres, symbol)
        h5gr["$(symbol)_real", "shuffle", (), "compress", 3] = object.real
        h5gr["$(symbol)_fourier", "shuffle", (), "compress", 3] = object.fourier
    end

    for symbol in (:εF, )
        hasproperty(scfres, symbol) || continue
        object = getproperty(scfres, symbol)
        object === nothing && continue
        h5gr[string(symbol)] = object
    end

    for symbol in (:eigenvalues, :occupation)
        hasproperty(scfres, symbol) || continue
        h5gr[string(symbol)] = hcat(getproperty(scfres, symbol)...)
    end

    for key in keys(scfres.energies)
        h5gr["energy_$(key)"] = scfres.energies[key]
    end

    if store_wfn
        for symbol in (:ψ, )
            hasproperty(scfres, symbol) || continue
            object = getproperty(scfres, symbol)
            object === nothing && continue

            for ik in 1:length(object)
                h5gr["ψ_$(ik)_fourier", "shuffle", (), "compress", 3] = object[ik]
            end
        end
    end
end

function load_scfres(basis, fileprefix)
    has_results(fileprefix) || error("Need result in $fileprefix")

    h5open(fileprefix * ".hdf5", "r") do h5f
        res = h5f["result"]
        n_kpoints = length(basis.kpoints)
        n_bands = size(res["ψ_1_fourier"], 2)

        # Checks for consistency
        @assert size(res["ρ_real"]) == basis.fft_size
        @assert size(res["occupation"]) == (n_bands, n_kpoints)
        @assert size(res["eigenvalues"]) == (n_bands, n_kpoints)

        # Read state
        εF = read(res, "εF")
        ρ = from_real(basis, res["ρ_real"][:, :, :])
        ψ = Matrix{ComplexF64}[]
        occupation = Vector{Float64}[]
        eigenvalues = Vector{Float64}[]
        for ik = 1:n_kpoints
            n_fft = length(G_vectors(basis.kpoints[ik]))
            @assert size(res["ψ_$(ik)_fourier"]) == (n_fft, n_bands)
            push!(ψ, res["ψ_$(ik)_fourier"][:, :])

            push!(occupation, res["occupation"][:, ik])
            push!(eigenvalues, res["eigenvalues"][:, ik])
        end

        # Construct Hamiltonian and compute energies
        energies, ham = energy_hamiltonian(basis, ψ, occupation;
                                           ρ=ρ, eigenvalues=eigenvalues, εF=εF)

        (ham=ham, basis=basis, energies=energies, converged=true, ρ=ρ,
         eigenvalues=eigenvalues, occupation=occupation, εF=εF, n_iter=nothing, ψ=ψ)
    end
end

function has_results(fileprefix)
    !isfile(fileprefix * ".hdf5") && return false
    try
        return h5open(fileprefix * ".hdf5", "r") do h5f
            exists(h5f, "result")
        end
    catch
    end
    false
end

function run_dftk(lattice, atoms; Ecut=20, tol=1e-11, mixing=DFTK.SimpleMixing(),
                  maxiter=50, kspacing=0.3, smearing=Smearing.None(),
                  temperature=0, fileprefix="dftk",
                  determine_bands=nelec -> nelec // 2 + sqrt(nelec // 2),
                  stepsize=10, randomise_guess=false, kgrid=nothing,
                  supersampling=1.5, ensure_odd_fft_size=false,
                  find_occupation=find_occupation, random_ampl=1e-2,
                  diagtol_max=1e-3,
                 )
    model = model_DFT(lattice, atoms, [:gga_x_pbe, :gga_c_pbe],
                      smearing=smearing, temperature=temperature)
    isnothing(kgrid) && (kgrid = kgrid_size_from_minimal_spacing(lattice, kspacing))

    fft_size = determine_grid_size(model, Ecut, supersampling=supersampling)
    ensure_odd_fft_size && (fft_size = Vec3(fft_size .+ iseven.(fft_size)))
    basis = PlaneWaveBasis(model, Ecut, kgrid=kgrid, kshift=[0, 0, 0], fft_size=fft_size)
    n_bands = ceil(Int, determine_bands(model.n_electrons))

    if has_results(fileprefix)
        @info "Loading existing results instead of SCF"
        return load_scfres(basis, fileprefix)
    end
    n_blas = Int(ccall((BLAS.@blasfunc(openblas_get_num_threads), BLAS.libblas), Cint, ()))

    println()
    println("hostname       = $(read(`hostname`, String))")
    println("started on     = $(Dates.now())")
    println("julia threads  = $(Threads.nthreads())")
    println("BLAS threads   = $(n_blas)")
    println()
    println("temperature    = $(basis.model.temperature)")
    println("smearing       = $(basis.model.smearing)")
    if isdiag(lattice)
        println("diag(lattice)  = $(round.(diag(lattice), sigdigits=4))")
    else
        println("lattice        = $(round.(lattice, sigdigits=4))")
    end
    println("Ecut           = $Ecut")
    println("fft_size       = $(basis.fft_size)")
    println("kgrid          = $kgrid")
    println("irreducible k  = $(length(basis.kpoints))")
    println("n_bands        = $(n_bands)")
    println("n_electr       = $(basis.model.n_electrons)")
    println("mixing         = $mixing")
    flush(stdout)

    # Collection of data to dump later
    save_dict = Dict(
        "temperature"  => basis.model.temperature,
        "smearing"     => string(basis.model.smearing),
        "diag_lattice" => diag(lattice),
        "lattice"      => lattice,
        "Ecut"         => Ecut,
        "fft_size"     => basis.fft_size,
        "kgrid"        => kgrid,
        "n_kirred"     => length(basis.kpoints),
        "n_bands"      => n_bands,
        "n_electr"     => basis.model.n_electrons,
        "mixing"       => string(mixing),
        "α"            => mixing.α,
        "energies"     => Float64[],
    )

    # Empty the existing HDF5 file
    isdir(dirname(fileprefix)) || mkdir(dirname(fileprefix))
    close(h5open(fileprefix * ".hdf5", "w"))

    default_callback = ScfDefaultCallback()
    function callback(info)
        default_callback(info)
        flush(stdout)
        key = info.stage == :finalize ? "result" : "neval_$(info.n_iter)"
        store_wfn = info.stage == :finalize
        h5open(fileprefix * ".hdf5", "r+") do h5f
            write_scfres(g_create(h5f, key), info, store_wfn=store_wfn)
        end
        push!(save_dict["energies"], info.energies.total)
    end

    ρ0 = guess_density(basis)
    if randomise_guess
        # Add randomness and adjust electron count
        ρ0.real .= abs.(ρ0.real .+ random_ampl * rand(size(ρ0.real)))
        ρ0.real .*= basis.model.n_electrons ./ (sum(ρ0.real) ./ length(ρ0.real) .* basis.model.unit_cell_volume)
    end
    scfres = self_consistent_field(basis, ρ=ρ0, tol=tol, mixing=mixing,
                                   callback=callback,
                                   determine_diagtol=ScfDiagtol(ratio_ρdiff=0.04,
                                                                diagtol_max=diagtol_max,
                                                                diagtol_min=1e-12),
                                   solver=scf_nlsolve_solver(stepsize),
                                   maxiter=maxiter, n_bands=n_bands,
                                   find_occupation=find_occupation,
                                   enforce_symmetry=true)

    println(scfres.energies)
    flush(stdout)

    open(fp -> JSON.print(fp, save_dict), fileprefix * ".json", "w")
    scfres
end


text:
mutable struct Segment
   P0::SVector{3,Float64}    # Point 0 (support vector of line segment)
   P1::SVector{3,Float64}    # Point 1
   function Segment(P0::SVector{3,Float64},P1::SVector{3,Float64})
     new(P0,P1)
   end
end

mutable struct Plane
   V0::SVector{3,Float64}    # Point 0 of plane (support vector)
   V1::SVector{3,Float64}    # Point 1
   V2::SVector{3,Float64}    # Point 1
   n::SVector{3,Float64}     # normal of plane
   function Plane(V0::SVector{3,Float64}, V1::SVector{3,Float64}, V2::SVector{3,Float64})
      n = cross(V1-V0, V2-V0)
      new(V0,V1,V2,n)
   end
end

function intersect3DSegmentPlane(seg::Segment, plane::Plane, neps)
   u = seg.P1 - seg.P0
   w = seg.P0 - plane.V0

   dividend = -dot(plane.n, w)
   divisor = dot(plane.n, u)

   # checks if segment and plane are parallel
   if abs(divisor) < neps  # segment is parallel to plane
      if (abs(dividend) < neps)  # segemnt is parallel to plane AND lies in plane
         error("Segment is parallel to plane and lies in plane")
         return (2, nothing)
      else
         error("Segment is parallel to plane")
         return (0, nothing)
      end
   end

   # computes intersection point of plane and segment
   SI = dividend/divisor
   I = seg.P0 + SI*u
   return (1, I)
end

function sameSideTriangle(p1, p2, a, b)::Bool
   cp1 = cross(b-a, p1-a)
   cp2 = cross(b-a, p2-a)
   return (dot(cp1, cp2) >= 0.0)
end
pointInTriangle(p,a,b,c) = (sameSideTriangle(p,a,b,c) && sameSideTriangle(p,b,a,c) && sameSideTriangle(p,c,a,b))

function doesRayIntersectPortal(r1,r2,r3,point,neps)
   plane = Plane(r1,r2,r3)
   segment = Segment(point, SVector(0.0, 0.0, 0.0))

   (value, intersectionPoint) = intersect3DSegmentPlane(segment, plane, neps)
   if value == 1
      if !pointInTriangle(intersectionPoint, r1, r2, r3)
         error("Ray = ", point ," does not intersect portal (r1,r2,r3) = ", r1, r2, r3)
         return false
      end
   end

   return true
end

function analyzeFinalPortal(r1, r2, r3, r4, neps)
   plane = Plane(r1,r2,r3)
   segment = Segment(r4, SVector(0.0, 0.0, 0.0))
   (value, intersectionPoint) = intersect3DSegmentPlane(segment, plane, neps)
   if value == 1
      if !pointInTriangle(intersectionPoint, r1, r2, r3)
         println("Ray r4 = ", r4 ," is not intersecting last portal (r1,r2,r3) = ", r1, r2, r3)
         return false
      end
   end

   areaPlane = norm(plane.n)*1/2
   println("areaPlane = ", areaPlane)
    # det([r2-r1, r3-r1, r4-r1])*1/6
   volumeTetraeder1 = abs(dot(r2-r1,cross(r3-r1,r4-r1)))*1/6
   #volumeTetraeder2 = abs(det(hcat(r2-r1,r3-r1,r4-r1))) *1/6
   println("volumeTetraeder = ", volumeTetraeder1)
   #println("volumeTetraeder2 = ", volumeTetraeder2)
   return true
end

function sameSideTetrahedron(v1,v2,v3,v4,p)
   normal = cross(v2-v1,v3-v1)
   dotV4  = dot(normal, v4-v1)
   dotP   = dot(normal, p-v1)
   return ( sign(dotV4) == sign(dotP) )
end
pointInTetrahedron(v1,v2,v3,v4,p) = (sameSideTetrahedron(v1,v2,v3,v4,p) && sameSideTetrahedron(v2,v3,v4,v1,p) && sameSideTetrahedron(v3,v4,v1,v2,p) && sameSideTetrahedron(v4,v1,v2,v3,p) )


text:
@testset "Agent" begin
    action_space = DiscreteSpace(3)
    agent = Agent(;
        policy = RandomPolicy(action_space),
        trajectory = VectorialCompactSARTSATrajectory(),
    )

    obs1 = (state = 1,)
    agent(PRE_EPISODE_STAGE, obs1)
    a1 = agent(PRE_ACT_STAGE, obs1)
    @test a1 ∈ action_space

    obs2 = (reward = 1.0, terminal = true, state = 2)
    agent(POST_ACT_STAGE, obs2)
    dummy_action = agent(POST_EPISODE_STAGE, obs2)

    @test length(agent.trajectory) == 1
    @test get_trace(agent.trajectory, :state) == [1]
    @test get_trace(agent.trajectory, :action) == [a1]
    @test get_trace(agent.trajectory, :reward) == [get_reward(obs2)]
    @test get_trace(agent.trajectory, :terminal) == [get_terminal(obs2)]
    @test get_trace(agent.trajectory, :next_state) == [get_state(obs2)]
    @test get_trace(agent.trajectory, :next_action) == [dummy_action]
end


text:
### A Pluto.jl notebook ###
# v0.18.1

using Markdown
using InteractiveUtils

# ╔═╡ d6cfceae-779d-474d-a609-e1fe6fb1d276
begin
	import Pkg
	Pkg.activate(".")
end

# ╔═╡ ea6b61fa-51d1-499a-884a-b7fba187d005
begin
	using FinSight
	using StatsPlots
	using Distributions
	using PlutoUI
	using DataFrames
	using MonteCarloMeasurements
end

# ╔═╡ c2f43cb6-e226-4be6-b54c-96417ae953fa
# Illistrative "pro forma" DCF valuation of a rental property

# ╔═╡ d7e8f31d-86e9-4832-878b-be2d925767dc
begin
	# initial purchase price of property $ per unit area
	initial_investment = 1_000

	# rate
	r = 0.07

	# Capital Expenditures 10% of gross income
	ci = 0.10  # Could be zero

	#Growth Rate 2%
	g = 0.02

	# Expenses 35% of gross income
	expense_ratio = 0.35

	# Exit Capitalisation Rate 5%
	y = 0.05

	# Time horizon now (year 0) plus 11 years
	year = [0,0,0, 0,0,0, 0,0,0, 0, 0,0]

	# Potential Gross Income first year $100
	pgi_fy = 100.0

	nothing
end

# ╔═╡ 454d71be-0ca1-4234-a736-8d1033a5d68d
begin
	# Vacancy 5.0% of gross income
	# v_distr = TriangularDist(0.0, 0.8, 0.05)
     v_distr =Truncated(InverseGamma(2,0.1), 0, 1)
	# v_distr = LogNormal(0,1)
	# Try various types
	# v_distr = GeneralizedPareto(0.04, 0.8, 0.01)
	v = Particles(2000, v_distr)

	plot(v_distr, title="Vacancy distribution (non-gaussian)", leg=false, )
end

# ╔═╡ 6d70da11-1f37-46c1-bae2-c9635674f4e2
# Potential Gross Income nominal cash flow
begin
	pgi = zeros(typeof(pgi_fy), 12)
	pgi[2] = pgi_fy   # first year gross income is the second value
	growth_factor = 1.0 + g

	for i = 3:12
			prior_year_income = pgi[i-1]
	 		pgi[i] =   prior_year_income * growth_factor
	end

	pgi
end

# ╔═╡ 06a77d3e-e321-446a-9e18-8a53e3ad2214
pgi  # Potential gross income per square metre of office space

# ╔═╡ 46f48d40-e8d6-4b4a-8176-06ae91298209
begin
	df = DataFrame();
	df.years = 0:length(year)-1
	df.pgi = pgi
	df.vacancy_allowance = df.pgi .* v
	df.effective_gross_income = df.pgi - df.vacancy_allowance
	df.operating_expenses = expense_ratio * df.pgi
	df.net_operating_income = df.effective_gross_income - df.operating_expenses
	df.capital_expenditures = ci * df.pgi
	df.net_annual_cash_flows = df.net_operating_income - df.capital_expenditures
	df.reversion_cash_flows = [-initial_investment, 0,0,0,0,0, 0,0,0,0, df.net_annual_cash_flows[12] / y ,0]

	df.net_cash_flows = df.net_annual_cash_flows + df.reversion_cash_flows

	df
end;

# ╔═╡ 25391fc1-5a43-4dcd-b69b-ff77628dbb83
# transpose the dataframe for presentation with columns as time periods
nominal_cash_flows = DataFrame([[names(df)]; collect.(eachrow(df))], [:Period; Symbol.(axes(df, 1))])

# ╔═╡ faa662b0-d0d7-4a8d-8328-5f324a0d7f3f
net_present_value = npv(r, df.net_cash_flows)

# ╔═╡ 931cdff5-fe89-4f8c-bb5c-033ed228fdc5

	density(net_present_value, ylab="relative prob", title="Net present value distribution", leg=false)

# ╔═╡ Cell order:
# ╟─d6cfceae-779d-474d-a609-e1fe6fb1d276
# ╠═ea6b61fa-51d1-499a-884a-b7fba187d005
# ╠═c2f43cb6-e226-4be6-b54c-96417ae953fa
# ╠═d7e8f31d-86e9-4832-878b-be2d925767dc
# ╠═454d71be-0ca1-4234-a736-8d1033a5d68d
# ╠═6d70da11-1f37-46c1-bae2-c9635674f4e2
# ╠═06a77d3e-e321-446a-9e18-8a53e3ad2214
# ╠═46f48d40-e8d6-4b4a-8176-06ae91298209
# ╠═25391fc1-5a43-4dcd-b69b-ff77628dbb83
# ╠═faa662b0-d0d7-4a8d-8328-5f324a0d7f3f
# ╠═931cdff5-fe89-4f8c-bb5c-033ed228fdc5

max_stars_repo_name: davidtaylor-fastlanebi/FinSight

max_issues_repo_name: davidtaylor-fastlanebi/FinSight

max_forks_repo_name: davidtaylor-fastlanebi/FinSight


text:
using ArcadeLearningEnvironment
using ArgParse
using Cambrian
using CartesianGeneticProgramming
using Dates
using IICGP
using Random

# function extension
import Cambrian.mutate

```
Playing Atari games using DualCGP on screen input values

This uses a single Game seed, meaning an unfair deterministic Atari
environment. To evolve using a different game seed per generation, add in
reset_expert=true and seed=evo.gen below.
```

s = ArgParseSettings()
@add_arg_table! s begin
    "--cfg"
    help = "configuration script"
    default = "cfg/test_dual.yaml"
    "--game"
    help = "game rom name"
    default = "assault"
    "--seed"
    help = "random seed for evolution"
    arg_type = Int
    default = 0
    "--out"
    help = "output directory"
    arg_type = String
    default = dirname(@__DIR__)
    "--ind"
    help = "individual for evaluation"
    arg_type = String
    default = ""
end
args = parse_args(ARGS, s)
const game = args["game"]
const seed = args["seed"]
const resdir = args["out"]
Random.seed!(seed)

main_cfg, enco_cfg, cont_cfg, reducer, bootstrap = IICGP.dualcgp_config(
    args["cfg"], game
)
const max_frames = main_cfg["max_frames"]
const stickiness = main_cfg["stickiness"]
const grayscale = main_cfg["grayscale"]
const downscale = main_cfg["downscale"]
const logid = enco_cfg.id
const state_ref = get_state_ref(game, seed)

function play_atari(
    encoder::CGPInd,
    reducer::Reducer,
    controller::CGPInd,
    lck::ReentrantLock;
    rom=game,
    seed=seed,
    rom_state_ref=state_ref,
    max_frames=max_frames,
    grayscale=grayscale,
    downscale=downscale,
    stickiness=stickiness
)
    # Random.seed!(seed)
    mt = MersenneTwister(seed)
    game = Game(rom, seed, lck=lck, state_ref=rom_state_ref)
    IICGP.reset!(reducer) # zero the buffers
    reward = 0.0
    frames = 0
    prev_action = Int32(0)
    while ~game_over(game.ale)
        if rand(mt) > stickiness || frames == 0
            s = get_state(game, grayscale, downscale)
            output = IICGP.process(encoder, reducer, controller, s)
            action = game.actions[argmax(output)]
        else
            action = prev_action
        end
        reward += act(game.ale, action)
        frames += 1
        prev_action = action
        if frames > max_frames
            break
        end
    end
    close!(game)
    [reward]
end

if length(args["ind"]) > 0
    ind = CGPInd(cfg, read(args["ind"], String))
    ftn = fitness(ind, inps, outs)
    println("Fitness: ", ftn)
else
    # Define mutate and fit functions
    function mutate(ind::CGPInd, ind_type::String)
        if ind_type == "encoder"
            return goldman_mutate(enco_cfg, ind, init_function=IPCGPInd)
        elseif ind_type == "controller"
            return goldman_mutate(cont_cfg, ind)
        end
    end
    lck = ReentrantLock()
    fit(encoder::CGPInd, controller::CGPInd) = play_atari(encoder, reducer,
                                                          controller, lck)
    # Create an evolution framework
    e = IICGP.DualCGPEvolution(enco_cfg, cont_cfg, fit, resdir,
                               encoder_init_function=IPCGPInd, logid=logid,
                               bootstrap=bootstrap, game=game)
    # Run evolution
    init_backup(logid, resdir, args["cfg"])
    run!(e)
end


text:
using LinearAlgebra
using DynamicPolynomials
using SwitchOnSafety
using Combinatorics
using SparseArrays
using JuMP, MosekTools
using SpecialFunctions

include("../src/RandomTrajectories.jl")
include("../src/AlgebraicLift.jl")
include("../src/ScenarioOpti.jl")
include("../src/ProbabilisticCertificates.jl")
include("../src/WhiteBoxAnalysis.jl")

dim = 2; numMode = 3; dimIn = 1#Integer(floor(dim/2)+1)

numScen_budget = 2000

A = [[0.7 0.16; 1.1 -1.1], [0.4 -0.84; 0.83 0.35], [0.37 0.96; 0.34 -1.2]]
B = reshape([-0.9; -1.2],2,1)

jsrboundopen= white_box_jsr(A)
println("JSR open loop: $jsrboundopen")
gaWB,KWB = white_box_stabilization_quad(A,B)
println("White-box stabilization: $gaWB")
jsrboundcloseWB = white_box_jsr([Ai+B*KWB for Ai in A])
println("White-box jsr: $jsrboundcloseWB")

(state0_budget,state_budget) = generate_trajectories(1,A,numScen_budget)
K0 = zeros(dimIn,dim)

K,jsr_bound,ga,P = probabilistc_stability_certificate(state0_budget,state_budget;B=B,numMode=numMode,d=1,batchsize=numScen_budget,K0=K0,beta=0.01,tol=1e-3)

println("JSR bound: $jsr_bound")

Aclose = [Ai+B*K for Ai in A]
jsrboundclose = white_box_jsr(Aclose)
println("JSR closed: $jsrboundclose")
println("gamma: $ga")
println("P: $P")

max_stars_repo_name: zhemingwang/DataDrivenSwitchControl

max_issues_repo_name: zhemingwang/DataDrivenSwitchControl

max_forks_repo_name: zhemingwang/DataDrivenSwitchControl


text:
using SkyRoom
using Test

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


text:
using Pkg.Artifacts
import ONNXNaiveNASflux.BaseOnnx: readproto, TensorProto
import ONNXNaiveNASflux:  CompGraphBuilder, extract
using Downloads
using Downloads: download

const last_dl_time = Dict()

function prepare_node_test(name, ninputs, noutputs)
    ahash = get_node_artifact(name, ninputs=ninputs, noutputs=noutputs)
    apath = artifact_path(ahash)
    model = extract(joinpath(apath, "model.onnx"))

    inputs = readinput.(apath, 0:ninputs-1)
    outputs = readoutput.(apath, 0:noutputs-1)

    graph = model.graph
    for (vi, val) in zip(graph.input, inputs)
        val.name = vi.name
        push!(graph.initializer, val)
    end

    gb = CompGraphBuilder(graph)
    return model, gb, array.(inputs), array.(outputs)
end

# Browsable link: https://github.com/onnx/onnx/tree/v1.6.0/onnx/backend/test/data/node
function get_node_artifact(name, node=name; ninputs, noutputs, location = "https://raw.githubusercontent.com/onnx/onnx/v1.6.0/onnx/backend/test/data/node/")
    return get_artifact(name, joinpath(location, node), ninputs, noutputs)
end

# Copy paste from documentation basically
function get_artifact(name, location, ninputs, noutputs, rcnt=5)
    artifacts_toml = joinpath(@__DIR__, "Artifacts.toml")
    ahash = artifact_hash(name, artifacts_toml)
    if isnothing(ahash) || !artifact_exists(ahash)

        throttle_download(location)
        try
            ahash = create_artifact() do artifact_dir
                mkpath(joinpath(artifact_dir, "test_data_set_0"))

                download(location * "/model.onnx", joinpath(artifact_dir, "model.onnx"))
                for i in 0:ninputs-1
                    download(inputfile(location, i), inputfile(artifact_dir, i))
                end

                for i in 0:noutputs-1
                    download(outputfile(location, i), outputfile(artifact_dir, i))
                end
            end
        catch e
            rcnt == 0 && rethrow(e)
            @warn sprint(showerror, e) * " Retry nr $rcnt"
            get_artifact(name, location, ninputs, noutputs, rcnt-1)
        end

        bind_artifact!(artifacts_toml, name, ahash; force=true)
    end
    return ahash
end
    # Limit minimum time between downloads so the server doesn't block us
function throttle_download(location, waittime=0.5)
    ctime = time()
    baseurl = splitpath(location)[2]
    ptime = get(last_dl_time, baseurl, ctime-waittime)
    sleep(max(0, ptime-ctime + waittime))
    last_dl_time[baseurl] = time()
end

inputfile(apath, i) = apath * "/test_data_set_0" * "/" * join(["input_", i, ".pb"])
outputfile(apath, i) = apath * "/test_data_set_0" * "/" * join(["output_", i, ".pb"])

readinput(apath, i) = readdata(inputfile(apath, i))
readoutput(apath, i) = readdata(outputfile(apath, i))
readdata(filename) = readproto(open(filename), TensorProto())


text:
module IterativeFactorizations

using LinearAlgebra

using IterativeSolvers
import IterativeSolvers: GMRESIterable, gmres_iterable!
using AlgebraicMultigrid
using SparseArrays

using BlockBandedMatrices

const ConjugateGradient = Union{CGIterable,PCGIterable}

LinearAlgebra.isposdef(T::Union{Tridiagonal,SymTridiagonal}) = isposdef(Matrix(T))

mutable struct IterativeFactorization{Iterator,X,B,T}
    iterator::Iterator
    x::X # Solution vector
    b::B # Right-hand side vector
    tol::T
    verbosity::Int
end

Base.size(A::IterativeFactorization) = (length(A.b), length(A.x))
Base.size(A::IterativeFactorization, i::Int) = size(A)[i]
Base.eltype(A::IterativeFactorization) = eltype(A.x)

getscalarresidual(it::ConjugateGradient) = it.residual
getscalarresidual(it::GMRESIterable) = it.residual.current

preconditioner(Pl::P) where {P<:AbstractMatrix} =
    aspreconditioner(ruge_stuben(sparse(Pl)))
preconditioner(Pl::Union{Diagonal,SymTridiagonal}) = factorize(Pl)

function IterativeFactorization(A::M,
                                x::X=zeros(eltype(A), size(A, 2)),
                                b::B=similar(x);
                                tol=√(eps(real(eltype(b)))),
                                prec=preconditioner(A),
                                isposdefA=isposdef(A),
                                verbosity=0,
                                kwargs...) where {M,X,B}
    iterator,b = if isposdefA
        verbosity > 0 && println("Positive-definite -> Conjugate-gradient")
        iterator = cg_iterator!(x, A, b, prec;
                                tol=tol, initially_zero=iszero(x), kwargs...)
        iterator,iterator.r
    else
        verbosity > 0 && println("General -> GMRES")
        iterator = gmres_iterable!(x, A, b; Pl=prec, tol=tol,
                                   initially_zero=iszero(x), kwargs...)
        iterator,iterator.b
    end
    verbosity > 0 && println("Tolerance: $(tol), max #iteration: $(iterator.maxiter)")

    IterativeFactorization(iterator, x, b, tol, verbosity)
end

update_preconditioner!(A::IterativeFactorization, Pl::P) where {P<:AbstractMatrix} =
    A.iterator.Pl = preconditioner(Pl)

function reset_iterator!(it::It, tol::T) where {It<:ConjugateGradient,T}
    # It is assumed that the new RHS is already copied to it.r

    initially_zero = iszero(it.x)
    it.mv_products = if initially_zero
        0
    else
        mul!(it.c, it.A, it.x)
        it.r .-= it.c
        1
    end

    it.residual = norm(it.r)
    it.reltol = tol*it.residual
    it.u .= zero(eltype(it.x))
end

function reset_iterator!(it::It, tol::T) where {It<:GMRESIterable, T}
    # It is assumed that the new RHS is already copied to it.b
    initially_zero=iszero(it.x)
    it.mv_products = initially_zero ? 0 : 1

    β = IterativeSolvers.init!(it.arnoldi, it.x, it.b, it.Pl, it.Ax,
                               initially_zero=initially_zero)
    it.residual.current = β
    IterativeSolvers.init_residual!(it.residual, β)
    it.reltol = tol * β
    it.β = β
end

function _ldiv!(A::IterativeFactorization)
    iterator = A.iterator
    reset_iterator!(iterator, A.tol)

    A.verbosity > 1 && println("Initial residual: $(getscalarresidual(iterator))")

    ii = 0
    for (iteration,item) in enumerate(iterator)
        iterator isa ConjugateGradient && (iterator.mv_products += 1)
        R = getscalarresidual(iterator)
        A.verbosity > 1 && println("#$iteration: $(R)")
        ii += 1
        (isnan(R) || !isfinite(R)) && break
    end
    A.verbosity > 0 && println("Solution converged: ",
                              IterativeSolvers.converged(iterator) ? "yes" : "no",
                              ", #iterations: ", ii, "/", iterator.maxiter,
                              ", residual: ", getscalarresidual(iterator))
end

function LinearAlgebra.ldiv!(A::IterativeFactorization, b)
    copyto!(A.x, b)
    copyto!(A.b, b)

    _ldiv!(A)

    copyto!(b, A.iterator.x)
    b
end

function LinearAlgebra.ldiv!(x, A::IterativeFactorization, b)
    copyto!(A.x, x)
    copyto!(A.b, b)

    _ldiv!(A)

    copyto!(x, A.iterator.x)
    x
end

to_block_banded_matrix(A::BlockBandedMatrix) = A
function to_block_banded_matrix(A::BlockSkylineMatrix)
    m,n = size(A)
    bs = A.block_sizes
    l,u = maximum(bs.l),maximum(bs.u)
    axs = bs.axes
    rows = diff(vcat(0,axs[1].lasts))
    cols = diff(vcat(0,axs[2].lasts))
    BlockBandedMatrix(A, rows, cols, (l,u))
end

factorization(A::AbstractMatrix;kwargs...) = IterativeFactorization(A; kwargs...)
factorization(A::SymTridiagonal;kwargs...) = factorize(A)
factorization(A::BlockSkylineMatrix) = qr(to_block_banded_matrix(A))
LinearAlgebra.factorize(A::BlockSkylineMatrix) = factorization(A)

export IterativeFactorization, factorization

end # module