text:
# using ..CovidSurvey
#include("../data.jl")

const regions = ["Sweden","France","Hungary","UK","USA","Italy","Germany","Denmark"]
const nregions = length(regions)
const populations = fill(1e10, nregions)

function regions!(data)
    data["regions"] = regions
    data["num_regions"] = nregions
    data["populations"] = populations
    return nothing
end

function observables!(data)
    epilong = data["data"] = @_ projectdir("data/epidata/", "epidata_longformat.csv") |>
        CSV.read( __, DataFrame ) |>
        preprocess

    data["cases"] = @_ DataFrames.select(epilong, :date, :cases, :region) |>
        unstack(__, :region, :cases) |>
        disallowmissing(__)

    data["hospit"] = nothing

    data["deaths"] = @_ DataFrames.select(epilong, :date, :deaths, :region) |>
        unstack(__, :region, :deaths) |>
        disallowmissing(__)

    return nothing
end

function preprocess(df::DataFrame)
    startdate = Date("2020-01-05") #Date("2020-01-05") #avoid missing value at very early date
    enddate = Date("2021-04-01")
    @_ df |>
        filter(:date => >=(startdate), __) |>
        filter(:date => <=(enddate), __) |>
        disallowmissing
end

# max_vacc_rate = 0.05
# min_case_rate = 20
# epidemic = CSV.read(projectdir("data/epidata/epidata_region=$(region).csv"), DataFrame)
# is = findfirst(>(22), epidemic.cases_per_100000)
# ds = epidemic.date[is]
# de1 = epidemic.date[is:end][findfirst(<(10), epidemic.cases_per_100000[is:end])]
# de2 = epidemic.date[findfirst(>(0.05), epidemic.fraction_vaccinated)]
# de  = min(de1, de2)

# ============================================================================
# observation model defaults
# ============================================================================
@kwdef struct CaseInit <: CovidSurvey.ObservationInit
    obs_start        ::String  = ""
    obs_stop         ::String  = ""
    delay_dist                 = NegativeBinomial2
    delay_length     ::Int64   = 40
    delay_dispersion ::Float64 = 5.41
    populations      ::Vector{Int64} = International.populations
end

@kwdef struct DeathInit <: CovidSurvey.ObservationInit
    obs_start        ::String  = ""
    obs_stop         ::String  = ""
    delay_dist                 = NegativeBinomial2
    delay_length     ::Int64   = 40
    delay_dispersion ::Float64 = 14.26
    populations      ::Vector{Int64} = International.populations
end

ObservParams!(data, o::CaseInit)   = data["casemodel"]   = CovidSurvey.ObservParams(data, o)
ObservParams!(data, o::DeathInit)  = data["deathmodel"]  = CovidSurvey.ObservParams(data, o)

# ============================================================================

function lockdown!(data, lockdowns, include_early_dynamic)
    @unpack dates_turing, num_regions = data
    include_early_dynamic && (data["lockdown_indices"] = fill(0, num_regions); return)
    li = Vector{Int64}()
    for (d, lockdown) in zip(dates_turing, lockdowns)
        i = findfirst(==(Date(lockdown)), d)
        push!(li, i)
    end
    data["lockdown_indices"] = li
    return nothing
end

# ===============================================================================
#                           main function
# ===============================================================================
function load_data(;
    observationsend  = nothing,
    predictors       = nothing,
    rwstep           = 7,
    epidemicstart    = Date("2020-06-01"),
    epidemic_threshold = 30,
    numimpute        = 6,
    casemodel         = CaseInit(),
    deathmodel        = DeathInit(),
    link              = KLogistic(3.),
    invlink           = KLogit(3.),
    lockdowns         = nothing,
    predictive        = false,
    include_early_dynamic = true,
    covariates_kwargs = Dict(
        :semiparametric => false,
        :fname => projectdir( homedir(), "data/covidsurvey/smoothed_contacts.csv" ),
        :shift => 0,
        :startdate => nothing,
        :enddate => nothing,
        :aggregation => 1,
        :mobility => nothing,
        :standardize => false
    )
)

    data = Dict{String, Any}()
    regions!(data)
    observables!(data)
    CovidSurvey.consistent!(data; startdates = epidemicstart)
    CovidSurvey.turingformat!(data, observationsend, predictive)
    lockdown!(data, lockdowns, include_early_dynamic)
    # cases_start_date!(data, cases_start)
    CovidSurvey.covariates!(data, predictors; covariates_kwargs...)
    CovidSurvey.randomwalk!(data, rwstep, covariates_kwargs)
    ObservParams!(data, casemodel)
    ObservParams!(data, deathmodel)

    # iar_idx = get_iar_idx(num_tot, num_observations, cases_start_idx, iar_step)
    # num_iar_steps = length(unique(iar_idx))

    turing_data = (
        num_impute                  = numimpute,
        num_total_days              = data["num_tot"],
        dates                       = data["dates_turing"],
        cases                       = data["cases_turing"],
        casemodel                   = data["casemodel"],
        deaths                      = data["deaths_turing"],
        deathmodel                  = data["deathmodel"],
        epidemic_start              = epidemicstart,
        serial_interval             = CovidSurvey.serialinterval(15),#padzeros(serialinterval(30), 0, num_tot), #data_usa.turing_data.serial_intervals[1:50],
        num_si                      = 15,
        lockdown_indices            = data["lockdown_indices"],
        # cases_start_indices         = data["cases_start_idx"],
        # num_case_observations       = data["num_case_obs"],
        num_rt_steps                = data["num_rt_steps"],
        rt_step_indices             = data["rt_step_index"],
        num_regions                 = length(regions),
        num_observations            = data["num_observations"],
        link                        = link,
        invlink                     = invlink,
        num_covariates              = data["num_covariates"],
        covariates                  = data["covariates"],
        covariates_starts           = data["covariates_starts"],
        semiparametric              = covariates_kwargs[:semiparametric],
        include_early_dynamic       = include_early_dynamic,
        rwscale                     = sqrt(rwstep),
        populations                 = populations,
        init_infected               = first.(data["cases_turing"])
        # num_i2h                     = 40,
        # ϕ_i2h                       = 5.41,
        # num_i2d                     = 60,
        # ϕ_i2d                       = 14.26,
        # num_i2c                     = 40,
        # ϕ_i2c                       = 5.41,
    )

    dates = [ collect(d:Day(1):d+Day(n-1)) for (d, n) in zip(data["startdates"], data["num_tot"])]
    startdates = [d[data["epidemicstart"]] for d in data["dates_turing"]]

    CovidSurvey.Data(
        turing_data,
        data["cases"],
        data["deaths"],
        data["hospit"],
        data["num_observations"], #num_observations
        data["num_tot"],
        dates,
        populations,
        predictors,
        nothing,
        regions,
        nothing,
        data["observationsend"],
        nothing,
        startdates
    )
end

max_stars_repo_name: andreaskoher/CovidSurvey

max_issues_repo_name: andreaskoher/CovidSurvey

max_forks_repo_name: andreaskoher/CovidSurvey


text:
################################################################################
### FUNCTIONS FOR RUNNING COMPLETE WORKFLOWS

function RunSSCase(CaseDataDirectory=pwd(); PrintModel=false, WriteReport=true, PrintOutput=true, ModelOutputFileName="_Model.txt", SolutionOutputFileName="_SolutionData.txt", PrintSpacer="*"^50, TestMode=false)

    # Load data
    A,N,P,D,S,T,L,Sets,Pars = LoadSSCaseData(CaseDataDirectory);

    # Run steady state model
    Output, Stats = OptimizeSSCase(A,N,P,D,S,T,L,Sets,Pars,
        CaseDataDirectory=CaseDataDirectory,
        PrintModel=PrintModel,
        WriteReport=WriteReport,
        PrintOutput=PrintOutput,
        ModelOutputFileName=ModelOutputFileName,
        SolutionOutputFileName=SolutionOutputFileName,
        PrintSpacer=PrintSpacer);

    # Record results
    SSRecordMaker(A,N,P,D,S,T,L,Output,Stats,
        CaseDataDirectory=CaseDataDirectory,
        PrintSpacer=PrintSpacer);

    # Generate network plot
    SSNetworkPlot(A,N,P,D,S,L,Output.f,
        CaseDataDirectory=CaseDataDirectory,
        PrintSpacer=PrintSpacer,
        TestMode=TestMode);

    # Update User
    println(PrintSpacer*"\nSteady State Case Done!\n"*PrintSpacer)
end

ext: jl


text:
"""
	ComputeAngles(angx,angy,dip_flag=false,vel="vel",wav="wav",sz=0.,nhx=100,ohx=0,dhx=10,nhy=1,ohy=0,dhy=10,pade_flag=false,fmin=0,fmax=80,padt=2,padx=2,verbose=false,sx=[0],sy=[0])

Compute angles for shot gathers, for example, to map migrated shots into angle gathers during shot profile migration.

# Arguments
- `angx::AbstractString` : filename for incidence angles in the x direction for each shot
- `angy::AbstractString` : filename for incidence angles in the y direction for each shot
- `dip_flag = false` : flag to subtract reflector dip from the computed angles to make them with reference to reflector normal
- `vel = "vel"` : seis file containing the velocity (should have same x and z dimensions as the desired image)
- `wav = "wav"` : seis file containing the source wavelet (in time domain)
- `sz = 0.` : source depth (Dev: read this from source wavelet file for variable source depth)
- `nhx = 101` : number of offset bins
- `ohx = 1000.` : min offset (surface offset in the data)
- `dhx = 10.` : offset increment
- `nhy = 101` : number of offset bins
- `ohy = 1000.` : min offset (surface offset in the data)
- `dhy = 10.` : offset increment
- `pade_flag = false` : flag for Pade Fourier correction
- `fmin = 0.` : min frequency to process (Hz)
- `fmax = 80.` : max frequency to process (Hz)
- `padt = 2` : pad factor for the time axis
- `padx = 2` : pad factor for the spatial axes
- `verbose = false` : flag for error / debugging messages
- `sx = [0.]` : array of source X positions (meters)
- `sy = [0.]` : array of source Y positions (meters)

*Credits: AS, 2015*
"""
function ComputeAngles(angx::AbstractString,angy::AbstractString,dip_flag=false,vel="vel",wav="wav",sz=0.,nhx=100,ohx=0,dhx=10,nhy=1,ohy=0,dhy=10,pade_flag=false,fmin=0,fmax=80,padt=2,padx=2,verbose=false,sx=[0],sy=[0])

	nshot = length(sx)
	if (dip_flag=="y")
		dipx,h = SeisRead(dipx_name)
		dipy,h = SeisRead(dipy_name)
	end

	v,h = SeisRead(vel)
	min_imx = h[1].imx
	max_imx = h[end].imx
	dmx = dhx
	nmx = max_imx - min_imx + 1
	min_imy = h[1].imy
	max_imy = h[end].imy
	dmy = dhy
	nmy = max_imy - min_imy + 1
	nz = h[1].n1
	dz = h[1].d1
	oz = h[1].o1
	min_gx = h[1].mx
	max_gx = h[end].mx
	min_gy = h[1].my
	max_gy = h[end].my

	shot_list = Array(ShotAngles,nshot)
	for ishot = 1 : nshot
		shot_list[ishot] = ShotAngles(0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0)
		shot_list[ishot].angx = join(["angx_shot_" int(floor(sx[ishot])) "_" int(floor(sy[ishot]))])
		shot_list[ishot].angy = join(["angy_shot_" int(floor(sx[ishot])) "_" int(floor(sy[ishot]))])
		shot_list[ishot].vel = vel
		shot_list[ishot].dipx = dipx_name
		shot_list[ishot].dipy = dipy_name
		shot_list[ishot].wav = wav
		shot_list[ishot].sx = sx[ishot]
		shot_list[ishot].sy = sy[ishot]
		shot_list[ishot].sz = sz
		shot_list[ishot].ohx = ohx
		shot_list[ishot].dhx = dhx
		shot_list[ishot].nhx = nhx
		shot_list[ishot].ohy = ohy
		shot_list[ishot].dhy = dhy
		shot_list[ishot].nhy = nhy
		shot_list[ishot].fmin = fmin
		shot_list[ishot].fmax = fmax
		shot_list[ishot].omp = omp
		shot_list[ishot].pade_flag = pade_flag
		shot_list[ishot].dip_flag = (dip_flag == true) ? "y" : "n"
		shot_list[ishot].verbose = (verbose == true) ? "y" : "n"
	end

	a = pmap(compute_angles,shot_list)

	j = 1
	for ishot = 1 : nshot
		angx_shot,h_shot = SeisRead(shot_list[ishot].angx)
		angy_shot,h_shot = SeisRead(shot_list[ishot].angy)
		SeisWrite(angx,angx_shot,h_shot,itrace=j)
		SeisWrite(angy,angy_shot,h_shot,itrace=j)
		j += size(angx_shot,2)
		SeisRemove(shot_list[ishot].angx)
		SeisRemove(shot_list[ishot].angy)
	end

end

struct ShotAngles
	angx
	angy
	vel
	dipx
	dipy
	wav
	sx
	sy
	sz
	ohx
	dhx
	nhx
	ohy
	dhy
	nhy
	fmin
	fmax
	pade_flag
	dip_flag
	omp
	verbose
end

function compute_angles(shot)
	ENV["OMP_NUM_THREADS"] = shot.omp
	if (find_library(["compute_angles"]) == "")
		error("Couldn't find shared library compute_angles.so, make sure it is installed and check LD_LIBRARY_PATH environmental variable.")
	end
	argv = ["0",
	join(["angx=",shot.angx]), join(["angy=",shot.angy]), join(["vel=",shot.vel]), join(["dipx=",shot.dipx]), join(["dipy=",shot.dipy]), join(["wav=",shot.wav]),
	join(["sx=",shot.sx]), join(["sy=",shot.sy]),  join(["sz=",shot.sz]),
	join(["ohx=",shot.ohx]),  join(["dhx=",shot.dhx]),  join(["nhx=",shot.nhx]),
	join(["ohy=",shot.ohy]),  join(["dhy=",shot.dhy]),  join(["nhy=",shot.nhy]),
	join(["fmin=",shot.fmin]),  join(["fmax=",shot.fmax]),
	join(["pade_flag=",shot.pade_flag]),  join(["dip_flag=",shot.dip_flag]),  join(["verbose=",shot.verbose]) ]
	a = ccall((:main, "compute_angles"), Int32, (Int32, Ptr{Ptr{UInt8}}), length(argv), argv)
	return(a)

end


text:
struct BlockMap{T,As<:Tuple{Vararg{LinearMap}},Rs<:Tuple{Vararg{Int}},Rranges<:Tuple{Vararg{UnitRange{Int}}},Cranges<:Tuple{Vararg{UnitRange{Int}}}} <: LinearMap{T}
    maps::As
    rows::Rs
    rowranges::Rranges
    colranges::Cranges
    function BlockMap{T,R,S}(maps::R, rows::S) where {T, R<:Tuple{Vararg{LinearMap}}, S<:Tuple{Vararg{Int}}}
        for n in eachindex(maps)
            A = maps[n]
            @assert promote_type(T, eltype(A)) == T "eltype $(eltype(A)) cannot be promoted to $T in BlockMap constructor"
        end
        rowranges, colranges = rowcolranges(maps, rows)
        return new{T,R,S,typeof(rowranges),typeof(colranges)}(maps, rows, rowranges, colranges)
    end
end

BlockMap{T}(maps::As, rows::S) where {T,As<:Tuple{Vararg{LinearMap}},S} = BlockMap{T,As,S}(maps, rows)

MulStyle(A::BlockMap) = MulStyle(A.maps...)

Base.parent(A::BlockMap) = A.maps

"""
    rowcolranges(maps, rows)

Determines the range of rows for each block row and the range of columns for each
map in `maps`, according to its position in a virtual matrix representation of the
block linear map obtained from `hvcat(rows, maps...)`.
"""
function rowcolranges(maps, rows)
    rowranges = ()
    colranges = ()
    mapind = 0
    rowstart = 1
    for row in rows
        xinds = vcat(1, map(a -> size(a, 2), maps[mapind+1:mapind+row])...)
        cumsum!(xinds, xinds)
        mapind += 1
        rowend = rowstart + size(maps[mapind], 1) - 1
        rowranges = (rowranges..., rowstart:rowend)
        colranges = (colranges..., xinds[1]:xinds[2]-1)
        for colind in 2:row
            mapind +=1
            colranges = (colranges..., xinds[colind]:xinds[colind+1]-1)
        end
        rowstart = rowend + 1
    end
    return rowranges::NTuple{length(rows), UnitRange{Int}}, colranges::NTuple{length(maps), UnitRange{Int}}
end

Base.size(A::BlockMap) = (last(last(A.rowranges)), last(last(A.colranges)))

############
# hcat
############
"""
    hcat(As::Union{LinearMap,UniformScaling,AbstractVecOrMat}...)::BlockMap

Construct a (lazy) representation of the horizontal concatenation of the arguments.
All arguments are promoted to `LinearMap`s automatically.

# Examples
```jldoctest; setup=(using LinearMaps)
julia> CS = LinearMap{Int}(cumsum, 3)::LinearMaps.FunctionMap;

julia> L = [CS LinearMap(ones(Int, 3, 3))]::LinearMaps.BlockMap;

julia> L * ones(Int, 6)
3-element Array{Int64,1}:
 4
 5
 6
```
"""
function Base.hcat(As::Union{LinearMap,UniformScaling,AbstractVecOrMat}...)
    T = promote_type(map(eltype, As)...)
    nbc = length(As)

    nrows = -1
    # find first non-UniformScaling to detect number of rows
    for A in As
        if !(A isa UniformScaling)
            nrows = size(A, 1)
            break
        end
    end
    nrows == -1 && throw(ArgumentError("hcat of only UniformScaling objects cannot determine the linear map size"))
    return BlockMap{T}(promote_to_lmaps(fill(nrows, nbc), 1, 1, As...), (nbc,))
end

############
# vcat
############
"""
    vcat(As::Union{LinearMap,UniformScaling,AbstractVecOrMat}...)::BlockMap

Construct a (lazy) representation of the vertical concatenation of the arguments.
All arguments are promoted to `LinearMap`s automatically.

# Examples
```jldoctest; setup=(using LinearMaps)
julia> CS = LinearMap{Int}(cumsum, 3)::LinearMaps.FunctionMap;

julia> L = [CS; LinearMap(ones(Int, 3, 3))]::LinearMaps.BlockMap;

julia> L * ones(Int, 3)
6-element Array{Int64,1}:
 1
 2
 3
 3
 3
 3
```
"""
function Base.vcat(As::Union{LinearMap,UniformScaling,AbstractVecOrMat}...)
    T = promote_type(map(eltype, As)...)
    nbr = length(As)

    ncols = -1
    # find first non-UniformScaling to detect number of columns
    for A in As
        if !(A isa UniformScaling)
            ncols = size(A, 2)
            break
        end
    end
    ncols == -1 && throw(ArgumentError("vcat of only UniformScaling objects cannot determine the linear map size"))

    return BlockMap{T}(promote_to_lmaps(fill(ncols, nbr), 1, 2, As...), ntuple(i->1, nbr))
end

############
# hvcat
############
"""
    hvcat(rows::Tuple{Vararg{Int}}, As::Union{LinearMap,UniformScaling,AbstractVecOrMat}...)::BlockMap

Construct a (lazy) representation of the horizontal-vertical concatenation of the arguments.
The first argument specifies the number of arguments to concatenate in each block row.
All arguments are promoted to `LinearMap`s automatically.

# Examples
```jldoctest; setup=(using LinearMaps)
julia> CS = LinearMap{Int}(cumsum, 3)::LinearMaps.FunctionMap;

julia> L = [CS CS; CS CS]::LinearMaps.BlockMap;

julia> L.rows
(2, 2)

julia> L * ones(Int, 6)
6-element Array{Int64,1}:
 2
 4
 6
 2
 4
 6
```
"""
Base.hvcat

function Base.hvcat(rows::Tuple{Vararg{Int}}, As::Union{LinearMap,UniformScaling,AbstractVecOrMat}...)
    nr = length(rows)
    T = promote_type(map(eltype, As)...)
    sum(rows) == length(As) || throw(ArgumentError("mismatch between row sizes and number of arguments"))
    n = fill(-1, length(As))
    j = 0
    for i in 1:nr # infer UniformScaling sizes from row counts, if possible:
        ni = -1 # number of rows in this block-row, -1 indicates unknown
        for k in 1:rows[i]
            if !isa(As[j+k], UniformScaling)
                na = size(As[j+k], 1)
                ni >= 0 && ni != na &&
                    throw(DimensionMismatch("mismatch in number of rows"))
                ni = na
            end
        end
        if ni >= 0
            for k = 1:rows[i]
                n[j+k] = ni
            end
        end
        j += rows[i]
    end
    # check for consistent total column number
    nc = -1
    j = 0
    for i in 1:nr
        nci = 0
        rows[i] > 0 && n[j+1] == -1 && (j += rows[i]; continue)
        for k = 1:rows[i]
            nci += isa(As[j+k], UniformScaling) ? n[j+k] : size(As[j+k], 2)
        end
        nc >= 0 && nc != nci && throw(DimensionMismatch("mismatch in number of columns"))
        nc = nci
        j += rows[i]
    end
    nc == -1 && throw(ArgumentError("sizes of UniformScalings could not be inferred"))
    j = 0
    for i in 1:nr
        if rows[i] > 0 && n[j+1] == -1 # this row consists entirely of UniformScalings
            nci, r = divrem(nc, rows[i])
            r != 0 && throw(DimensionMismatch("indivisible UniformScaling sizes"))
            for k = 1:rows[i]
                n[j+k] = nci
            end
        end
        j += rows[i]
    end

    return BlockMap{T}(promote_to_lmaps(n, 1, 1, As...), rows)
end

function check_dim(A, dim, n)
    n == size(A, dim) || throw(DimensionMismatch("Expected $n, got $(size(A, dim))"))
    return nothing
end

promote_to_lmaps_(n::Int, dim, A::AbstractMatrix) = (check_dim(A, dim, n); LinearMap(A))
promote_to_lmaps_(n::Int, dim, A::AbstractVector) = (check_dim(A, dim, n); LinearMap(reshape(A, length(A), 1)))
promote_to_lmaps_(n::Int, dim, J::UniformScaling) = UniformScalingMap(J.λ, n)
promote_to_lmaps_(n::Int, dim, A::LinearMap) = (check_dim(A, dim, n); A)
promote_to_lmaps(n, k, dim) = ()
promote_to_lmaps(n, k, dim, A) = (promote_to_lmaps_(n[k], dim, A),)
@inline promote_to_lmaps(n, k, dim, A, B, Cs...) =
    (promote_to_lmaps_(n[k], dim, A), promote_to_lmaps_(n[k+1], dim, B), promote_to_lmaps(n, k+2, dim, Cs...)...)

############
# basic methods
############

function isblocksquare(A::BlockMap)
    rows = A.rows
    N = length(rows)
    return all(==(N), rows)
end

# the following rules are sufficient but not necessary
function LinearAlgebra.issymmetric(A::BlockMap)
    isblocksquare(A) || return false
    N = length(A.rows)
    maps = A.maps
    symindex = vec(permutedims(reshape(collect(1:N*N), N, N)))
    for i in 1:N*N
        if (i == symindex[i] && !issymmetric(maps[i]))
            return false
        elseif (maps[i] != transpose(maps[symindex[i]]))
            return false
        end
    end
    return true
end

LinearAlgebra.ishermitian(A::BlockMap{<:Real}) = issymmetric(A)
function LinearAlgebra.ishermitian(A::BlockMap)
    isblocksquare(A) || return false
    N = length(A.rows)
    maps = A.maps
    symindex = vec(permutedims(reshape(collect(1:N*N), N, N)))
    for i in 1:N*N
        if (i == symindex[i] && !ishermitian(maps[i]))
            return false
        elseif (maps[i] != adjoint(maps[symindex[i]]))
            return false
        end
    end
    return true
end

############
# comparison of BlockMap objects, sufficient but not necessary
############

Base.:(==)(A::BlockMap, B::BlockMap) = (eltype(A) == eltype(B) && A.maps == B.maps && A.rows == B.rows)

############
# multiplication helper functions
############

function _blockmul!(y, A::BlockMap, x, α, β)
    if iszero(α)
        iszero(β) && return fill!(y, zero(eltype(y)))
        isone(β) && return y
        return rmul!(y, β)
    end
    return __blockmul!(MulStyle(A), y, A, x, α, β)
end

# provide one global intermediate storage vector if necessary
__blockmul!(::FiveArg, y, A, x, α, β)  = ___blockmul!(y, A, x, α, β, nothing)
__blockmul!(::ThreeArg, y, A, x, α, β) = ___blockmul!(y, A, x, α, β, similar(y))

function ___blockmul!(y, A, x, α, β, ::Nothing)
    maps, rows, yinds, xinds = A.maps, A.rows, A.rowranges, A.colranges
    mapind = 0
    for (row, yi) in zip(rows, yinds)
        yrow = selectdim(y, 1, yi)
        mapind += 1
        _unsafe_mul!(yrow, maps[mapind], selectdim(x, 1, xinds[mapind]), α, β)
        for _ in 2:row
            mapind +=1
            _unsafe_mul!(yrow, maps[mapind], selectdim(x, 1, xinds[mapind]), α, true)
        end
    end
    return y
end
function ___blockmul!(y, A, x, α, β, z)
    maps, rows, yinds, xinds = A.maps, A.rows, A.rowranges, A.colranges
    mapind = 0
    for (row, yi) in zip(rows, yinds)
        yrow = selectdim(y, 1, yi)
        zrow = selectdim(z, 1, yi)
        mapind += 1
        if MulStyle(maps[mapind]) === ThreeArg() && !iszero(β)
            !isone(β) && rmul!(yrow, β)
            muladd!(ThreeArg(), yrow, maps[mapind], selectdim(x, 1, xinds[mapind]), α, zrow)
        else
            _unsafe_mul!(yrow, maps[mapind], selectdim(x, 1, xinds[mapind]), α, β)
        end
        for _ in 2:row
            mapind +=1
            muladd!(MulStyle(maps[mapind]), yrow, maps[mapind], selectdim(x, 1, xinds[mapind]), α, zrow)
        end
    end
    return y
end

function _transblockmul!(y, A::BlockMap, x, α, β, transform)
    maps, rows, xinds, yinds = A.maps, A.rows, A.rowranges, A.colranges
    if iszero(α)
        iszero(β) && return fill!(y, zero(eltype(y)))
        isone(β) && return y
        return rmul!(y, β)
    else
        # first block row (rowind = 1) of A, meaning first block column of A', fill all of y
        xcol = selectdim(x, 1, first(xinds))
        for rowind in 1:first(rows)
            _unsafe_mul!(selectdim(y, 1, yinds[rowind]), transform(maps[rowind]), xcol, α, β)
        end
        mapind = first(rows)
        # subsequent block rows of A (block columns of A'),
        # add results to corresponding parts of y
        # TODO: think about multithreading
        for (row, xi) in zip(Base.tail(rows), Base.tail(xinds))
            xcol = selectdim(x, 1, xi)
            for _ in 1:row
                mapind +=1
                _unsafe_mul!(selectdim(y, 1, yinds[mapind]), transform(maps[mapind]), xcol, α, true)
            end
        end
    end
    return y
end

############
# multiplication with vectors & matrices
############

for (intype, outtype) in ((AbstractVector, AbstractVecOrMat), (AbstractMatrix, AbstractMatrix))
    @eval begin
        function _unsafe_mul!(y::$outtype, A::BlockMap, x::$intype)
            require_one_based_indexing(y, x)
            return _blockmul!(y, A, x, true, false)
        end
        function _unsafe_mul!(y::$outtype, A::BlockMap, x::$intype,
                            α::Number, β::Number)
            require_one_based_indexing(y, x)
            return _blockmul!(y, A, x, α, β)
        end
    end

    for (maptype, transform) in ((:(TransposeMap{<:Any,<:BlockMap}), :transpose), (:(AdjointMap{<:Any,<:BlockMap}), :adjoint))
        @eval begin
            function _unsafe_mul!(y::$outtype, wrapA::$maptype, x::$intype)
                require_one_based_indexing(y, x)
                return _transblockmul!(y, wrapA.lmap, x, true, false, $transform)
            end
            function _unsafe_mul!(y::$outtype, wrapA::$maptype, x::$intype,
                            α::Number, β::Number)
                require_one_based_indexing(y, x)
                return _transblockmul!(y, wrapA.lmap, x, α, β, $transform)
            end
        end
    end
end

############
# BlockDiagonalMap
############

struct BlockDiagonalMap{T,As<:Tuple{Vararg{LinearMap}},Ranges<:Tuple{Vararg{UnitRange{Int}}}} <: LinearMap{T}
    maps::As
    rowranges::Ranges
    colranges::Ranges
    function BlockDiagonalMap{T,As}(maps::As) where {T, As<:Tuple{Vararg{LinearMap}}}
        for n in eachindex(maps)
            A = maps[n]
            @assert promote_type(T, eltype(A)) == T "eltype $(eltype(A)) cannot be promoted to $T in BlockDiagonalMap constructor"
        end
        # row ranges
        inds = vcat(1, size.(maps, 1)...)
        cumsum!(inds, inds)
        rowranges = ntuple(i -> inds[i]:inds[i+1]-1, Val(length(maps)))
        # column ranges
        inds[2:end] .= size.(maps, 2)
        cumsum!(inds, inds)
        colranges = ntuple(i -> inds[i]:inds[i+1]-1, Val(length(maps)))
        return new{T,As,typeof(rowranges)}(maps, rowranges, colranges)
    end
end

BlockDiagonalMap{T}(maps::As) where {T,As<:Tuple{Vararg{LinearMap}}} =
    BlockDiagonalMap{T,As}(maps)
BlockDiagonalMap(maps::LinearMap...) =
    BlockDiagonalMap{promote_type(map(eltype, maps)...)}(maps)

# since the below methods are more specific than the Base method,
# they would redefine Base/SparseArrays behavior
for k in 1:8 # is 8 sufficient?
    Is = ntuple(n->:($(Symbol(:A,n))::AbstractVecOrMat), Val(k-1))
    # yields (:A1, :A2, :A3, ..., :A(k-1))
    L = :($(Symbol(:A,k))::LinearMap)
    # yields :Ak
    mapargs = ntuple(n -> :($(Symbol(:A,n))), Val(k-1))
    # yields (:LinearMap(A1), :LinearMap(A2), ..., :LinearMap(A(k-1)))

    @eval begin
        function SparseArrays.blockdiag($(Is...), $L, As::Union{LinearMap,AbstractVecOrMat}...)
            return BlockDiagonalMap(convert_to_lmaps($(mapargs...))..., $(Symbol(:A,k)), convert_to_lmaps(As...)...)
        end

        function Base.cat($(Is...), $L, As::Union{LinearMap,AbstractVecOrMat}...; dims::Dims{2})
            if dims == (1,2)
                return BlockDiagonalMap(convert_to_lmaps($(mapargs...))..., $(Symbol(:A,k)), convert_to_lmaps(As...)...)
            else
                throw(ArgumentError("dims keyword in cat of LinearMaps must be (1,2)"))
            end
        end
    end
end

"""
    blockdiag(As::Union{LinearMap,AbstractVecOrMat}...)::BlockDiagonalMap

Construct a (lazy) representation of the diagonal concatenation of the arguments.
To avoid fallback to the generic `SparseArrays.blockdiag`, there must be a `LinearMap`
object among the first 8 arguments.
"""
SparseArrays.blockdiag

"""
    cat(As::Union{LinearMap,AbstractVecOrMat}...; dims=(1,2))::BlockDiagonalMap

Construct a (lazy) representation of the diagonal concatenation of the arguments.
To avoid fallback to the generic `Base.cat`, there must be a `LinearMap`
object among the first 8 arguments.
"""
Base.cat

Base.size(A::BlockDiagonalMap) = (last(A.rowranges[end]), last(A.colranges[end]))

MulStyle(A::BlockDiagonalMap) = MulStyle(A.maps...)

Base.parent(A::BlockDiagonalMap) = A.maps

LinearAlgebra.issymmetric(A::BlockDiagonalMap) = all(issymmetric, A.maps)
LinearAlgebra.ishermitian(A::BlockDiagonalMap{<:Real}) = all(issymmetric, A.maps)
LinearAlgebra.ishermitian(A::BlockDiagonalMap) = all(ishermitian, A.maps)

LinearAlgebra.adjoint(A::BlockDiagonalMap{T}) where {T} = BlockDiagonalMap{T}(map(adjoint, A.maps))
LinearAlgebra.transpose(A::BlockDiagonalMap{T}) where {T} = BlockDiagonalMap{T}(map(transpose, A.maps))

Base.:(==)(A::BlockDiagonalMap, B::BlockDiagonalMap) = (eltype(A) == eltype(B) && A.maps == B.maps)

for (intype, outtype) in ((AbstractVector, AbstractVecOrMat), (AbstractMatrix, AbstractMatrix))
    @eval begin
        function _unsafe_mul!(y::$outtype, A::BlockDiagonalMap, x::$intype)
            require_one_based_indexing(y, x)
            return _blockscaling!(y, A, x, true, false)
        end
        function _unsafe_mul!(y::$outtype, A::BlockDiagonalMap, x::$intype,
                            α::Number, β::Number)
            require_one_based_indexing(y, x)
            return _blockscaling!(y, A, x, α, β)
        end
    end
end

function _blockscaling!(y, A::BlockDiagonalMap, x, α, β)
    maps, yinds, xinds = A.maps, A.rowranges, A.colranges
    # TODO: think about multi-threading here
    @views for i in eachindex(yinds, maps, xinds)
        _unsafe_mul!(selectdim(y, 1, yinds[i]), maps[i], selectdim(x, 1, xinds[i]), α, β)
    end
    return y
end


text:
include("../raylib.jl")
const s_w = 800
const s_h = 450
InitWindow(s_w,s_h,"raylib [shapes] - raylib logo")
SetTargetFPS(60)
while !WindowShouldClose()
	BeginDrawing()
		ClearBackground(Color(245,245,245,255))
		DrawRectangle(s_w÷2 - 128, s_h÷2 - 128,256,256,Color(77,100,174,255))
		DrawRectangle(s_w÷2 - 112, s_h÷2 - 112,224,224,Color(245,245,245,255))
		DrawCircle(s_w÷2 + 70, s_h÷2 , 12.0, Color(56,152,38,255))
		DrawCircle(s_w÷2 + 55, s_h÷2 + 25 , 12.0, Color(203,60,51,255))
		DrawCircle(s_w÷2 + 85, s_h÷2 + 25 , 12.0, Color(149,88,178,255))
		DrawText("raylib.jl",s_w÷2 - 44, s_h÷2+48,40,Color(0,0,0,255))
	EndDrawing()
end


text:
@testset "Neighborhoods" begin
  @testset "MetricBall" begin
    # Euclidean metric
    b = MetricBall(T(1/2))
    r = radius(b)
    m = metric(b)
    @test evaluate(m, T[0], T[0]) ≤ r
    @test evaluate(m, T[0], T[1]) > r

    b = MetricBall(T(1))
    r = radius(b)
    m = metric(b)
    @test evaluate(m, T[0,0], T[0,0]) ≤ r
    @test evaluate(m, T[0,0], T[1,0]) ≤ r
    @test evaluate(m, T[0,0], T[0,1]) ≤ r

    # Chebyshev metric
    b = MetricBall(T(1/2), Chebyshev())
    r = radius(b)
    m = metric(b)
    @test evaluate(m, T[0], T[0]) ≤ r
    @test evaluate(m, T[0], T[1]) > r

    for r in [1.,2.,3.,4.,5.]
      b = MetricBall(r, Chebyshev())
      r = radius(b)
      m = metric(b)
      for i in 0.0:1.0:r, j in 0.0:1.0:r
        @test evaluate(m, T[0,0], T[i,j]) ≤ r
      end
    end

    # 2D simple test of default convention
    m = metric(MetricBall(T.((1,1))))
    @test evaluate(m, T[1,0], T[0,0]) == evaluate(m, T[0,1], T[0,0])

    m = metric(MetricBall(T.((1,2))))
    @test evaluate(m, T[1,0], T[0,0]) != evaluate(m, T[0,1], T[0,0])

    # 3D simple test of default convention
    m = metric(MetricBall(T.((1.0,0.5,0.5)), TaitBryanAngles(T(π/4),T(0),T(0))))
    @test evaluate(m, [1.,1.,0.], [0.,0.,0.]) ≈ √T(2)
    @test evaluate(m, [-1.,1.,0.], [0.,0.,0.]) ≈ √T(8)

    # make sure the correct constructor is called
    m = metric(MetricBall(T[1.0,0.5,0.2], EulerAngles(T(0), T(0), T(0))))
    @test m isa Mahalanobis
  end
end


text:
type ScalarFitnessScheme{T <: Number} <: FitnessScheme
  maximise::Bool

  ScalarFitnessScheme(m::Bool) = new(m)
end

"""
The scalar fitness scheme represents the fitness of an individual using a
single scalar value, which is subject to either maximisation or minimisation.

**Properties:**

* `of::Float`, the type used to represent fitness values. Defaults to `Float`.
* `maximise::Bool`, a flag specifying whether the fitness value should be
maximised. Defaults to `true`.
"""
scalar(s::Dict{Any, Any}) =
  scalar(eval(Base.parse(Base.get(s, "of", "Float"))), Base.get(s, "maximise", true))
scalar(maximise::Bool) = scalar(Float, maximise)
scalar() = scalar(Float, true)
scalar(t::Type) = scalar(t, true)
scalar(t::Type, maximise::Bool) =
  ScalarFitnessScheme{t}(maximise)

uses{T}(s::ScalarFitnessScheme{T}) = T
maximise(s::ScalarFitnessScheme) = s.maximise
assign{T}(s::ScalarFitnessScheme{T}, v::T) = v

compare{T <: Number}(s::ScalarFitnessScheme{T}, x::T, y::T) =
  if x > y
    s.maximise ? -1 : 1
  elseif x < y
    s.maximise ? 1 : -1
  else
    0
  end


text:
#=
# spin10.jl - This module demonstrates grape on the spin system in julia.
=#
module Foo

# using BenchmarkTools
# using CUDAnative
# using CuArrays
# using DiffEqSensitivity
# using DifferentialEquations
# using ForwardDiff
using HDF5
using LinearAlgebra
# using Polynomials
using Printf
using StaticArrays
using Zygote

EXPERIMENT_META = "spin"
EXPERIMENT_NAME = "spin0"
WDIR = ENV["QC_PATH"]
SAVE_PATH = joinpath(WDIR, "out/$EXPERIMENT_META/$EXPERIMENT_NAME")

function generate_save_file_path(save_file_name, save_path)
    # Ensure the path exists.
    mkpath(save_path)

    # Create a save file name based on the one given; ensure it will
    # not conflict with others in the directory.
    max_numeric_prefix = -1
    for (_, _, files) in walkdir(save_path)
        for file_name in files
            if occursin("_$save_file_name.h5", file_name)
                max_numeric_prefix = max(parse(Int, split(file_name, "_")[1]))
            end
        end
    end

    save_file_name = "_$save_file_name.h5"
    save_file_name = @sprintf("%05d%s", max_numeric_prefix + 1, save_file_name)

    return joinpath(save_path, save_file_name)
end

mutable struct Reporter
    cost :: Float64
    costs :: Array{Float64, 1}
end

mutable struct PState
    control_count :: Int64
    control_eval_count :: Int64
    control_eval_times :: AbstractArray{Float64}
    cost_multipliers :: Array{Float64, 1}
    dt :: Float64
    evolution_time :: Float64
    initial_astate :: Any
    lagrange_multipliers :: Array{Float64, 1}
end

# We need an isomorphism from complex numbers to real numbers.
# TODO: It may be more memory efficient to store the augmented
# matrix as [mat_r ; mat_i] and define special multiplication
# operations.
"""
Promote a complex vector on the Hilbert space to
a real operator on the augmented Hilbert space.
"""
function augment_vec(vec)
    return [real(vec); imag(vec)]
end

"""
Complex conjugate an augmented vector.
"""
function conj_vec(vec)
    (len,) = size(vec)
    imag_indices = Int(len / 2) + 1:len
    vec[imag_indices] = -vec[imag_indices]
    return vec
end

"""
Promote a complex operator on the Hilbert space to a real
operator on the augmented Hilbert space.
"""
function augment_mat(mat_r, mat_i)
    return [mat_r -mat_i ; mat_i mat_r]
end

"""
Complex conjugate an augmented matrix.
"""
function conj_mat(mat)
    (len, len) = size(mat)
    len_by_2 = Int(len / 2)
    lo = 1:len_by_2
    hi = len_by_2 + 1:len
    mat[lo, hi] = -mat[lo, hi]
    mat[hi, lo] = -mat[hi, lo]
    return mat
end

"""
Multiply an augmented matrix by the scalar negative imaginary unit (-i).

This is faster than multiplying by a matrix but
cannot be used in the differentiable code
because Zygote does not like the indexing.
"""
function neg_i_mat(mat)
    (len, len) = size(mat)
    len_by_2 = Int(len / 2)
    lo = 1:len_by_2
    hi = len_by_2 + 1:len
    mat_r = mat[lo, lo]
    mat_i = mat[hi, lo]
    return [mat_i mat_r ; -mat_r mat_i]
end

"""
Construct an augmented matrix that represents the conjugate transpose of
a vector.
"""
function conjugate_transpose_vec(vec)
    vec_r_t = real(vec)'
    vec_i_t = imag(vec)'
    return [vec_r_t vec_i_t ; -vec_i_t vec_r_t]
end

# Define experimental constants.
OMEGA = 2 * pi * 1e-2
MAX_CONTROL_NORM_0 = 2 * pi * 3e-1
HAMILTONIAN_ARGS = [
    OMEGA,
]
MAX_CONTROL_NORMS = [
    MAX_CONTROL_NORM_0,
]

# Define the system.
NEG_I = SA_F64[0   0  1  0 ;
               0   0  0  1 ;
               -1  0  0  0 ;
               0  -1  0  0 ;]
SIGMA_X = SA_F64[0 1 0 0;
                 1 0 0 0;
                 0 0 0 1;
                 0 0 1 0]
SIGMA_Z = SA_F64[1   0  0   0;
                 0  -1  0   0;
                 0   0  1   0;
                 0   0  0  -1]
H_S = SIGMA_Z / 2
NEG_I_H_S = NEG_I * H_S
OMEGA_NEG_I_H_S = OMEGA * NEG_I_H_S
H_C1 = SIGMA_X / 2
NEG_I_H_C1 = NEG_I * H_C1

# Define the optimization.
EVOLUTION_TIME = 120
COMPLEX_CONTROLS = false
CONTROL_COUNT = 1
CONTROL_EVAL_COUNT = SYSTEM_EVAL_COUNT = Int(EVOLUTION_TIME) + 1
CONTROL_EVAL_TIMES = SVector{CONTROL_EVAL_COUNT, Float64}(Array(range(0., stop=EVOLUTION_TIME, length=CONTROL_EVAL_COUNT)))
DT = 1e-2
GRAB_CONTROLS = false
if GRAB_CONTROLS
    initial_controls_ = h5open(joinpath(SAVE_PATH, "00007_spin10.h5")) do save_file
        save_file["controls"][200, :, :]
    end
    initial_controls_ = SMatrix{CONTROL_EVAL_COUNT, CONTROL_COUNT, Float64}(initial_controls_)
else
    initial_controls_ = @SMatrix ones(CONTROL_EVAL_COUNT, CONTROL_COUNT)
end
INITIAL_CONTROLS = initial_controls_

# Define the problem.
INITIAL_STATE = augment_vec(SA[1.; 0])
ZEROS_INITIAL_STATE = augment_vec(SA[0; 0])
TARGET_STATE_DAGGER = conjugate_transpose_vec(SA[0; 1.])
INITIAL_ASTATE = [
    INITIAL_STATE; # state
    # ZEROS_INITIAL_STATE; # dstate_dw
    # ZEROS_INITIAL_STATE; # d2state_dw2
]
(STATE_SIZE,) = size(INITIAL_STATE)
(ASTATE_SIZE,) = size(INITIAL_ASTATE)
STATE_INDICES = 1:STATE_SIZE
DSTATE_DW_INDICES = STATE_SIZE + 1:STATE_SIZE * 2
D2STATE_DW2_INDICES = STATE_SIZE * 2 + 1:STATE_SIZE * 3

HILBERT_DIM = Int(STATE_SIZE / 2)

# Define Misc
LEARNING_RATE = 1e-3
ITERATION_COUNT = 200
INITIAL_COSTS = [0., 0]
INITIAL_COST_MULTIPLIERS = [1., 1.]
INITIAL_LAGRANGE_MULTIPLIERS = [0., 0]
COST_MULTIPLIER_STEP = 5.
SAVE = false

"""
"""
function rk3_step(rhs_, x, u, t, dt, us)
    k1 = rhs_(x,             u, t       , us) * dt
    k2 = rhs_(x + k1/2,      u, t + dt/2, us) * dt
    k3 = rhs_(x - k1 + 2*k2, u, t + dt  , us) * dt
    return x + (k1 + 4 * k2 + k3) / 6
end

# Attempt at in-place was slower than static array version.
# function rk3_step(rhs_, x, u, t, dt, dx, ks)
#     rhs_(dx, x,             u, t       )
#     ks[1,:] = dx * dt
#     rhs_(dx, x + ks[1,:]/2,      u, t + dt/2)
#     ks[2,:] = dx * dt
#     rhs_(dx, x - ks[1,:] + 2*ks[2,:], u, t + dt  )
#     ks[3, :] = dx * dt
#     return x + (ks[1,:] + 4 * ks[2,:] + ks[3,:]) / 6
# end

"""
"""
function rk4_step(rhs_, x, u, t, dt, us)
    k1 = rhs_(x,        u, t       , us) * dt
	k2 = rhs_(x + k1/2, u, t + dt/2, us) * dt
	k3 = rhs_(x + k2/2, u, t + dt/2, us) * dt
	k4 = rhs_(x + k3,   u, t + dt  , us) * dt
	return x + (k1 + 2 * k2 + 2 * k3 + k4) / 6
end

"""
Evaluate the polynomial using horner's method.
"""
function horner(coeffs, x)
    (len,) = size(coeffs)
    if len == 1
        val = coeffs[1]
    else
        val = coeffs[len]
        for i in len:-1:2
            val = coeffs[i - 1] + x * val
        end
    end
    return val
end

"""
Interpolate at x3 given (x1, y1), (x2, y2)
"""
function interpolate_linear_points(x1, x2, x3, y1, y2)
    return y1 + (((y2 - y1) / (x2 - x1)) * (x3 - x1))
end

"""
Interpolate the controls at a given time.
"""
function interpolate(controls, control_eval_times, time)
    (len,) = size(control_eval_times)
    if time < control_eval_times[1]
        t1 = control_eval_times[1]
        t2 = control_eval_times[2]
        c1 = controls[1]
        c2 = controls[2]
    elseif time >= control_eval_times[len]
        t1 = control_eval_times[len]
        t2 = control_eval_times[len - 1]
        c1 = controls[len]
        c2 = controls[len - 1]
    else
        hi = Int(findall(t -> t > time, control_eval_times)[1][1])
        t1 = control_eval_times[hi]
        t2 = control_eval_times[hi - 1]
        c1 = controls[hi]
        c2 = controls[hi - 1]
    end

    controls_ = interpolate_linear_points(t1, t2, time, c1, c2)
    return controls_
end

"""
The right-hand-side of the Schroedinger equation for the dynamics
we consider.
"""
function rhs(astate, controls, time, control_step)
    # Unpack augmented state.
    state = astate[STATE_INDICES]
    # dstate_dw = astate[DSTATE_DW_INDICES]
    # d2state_dw2 = astate[D2STATE_DW2_INDICES]
    # less gc time but more allocations = slower in the no gradient regime
    # state = SVector{4, Float64}(astate[STATE_INDICES])
    # dstate_dw = SVector{4, Float64}(astate[DSTATE_DW_INDICES])
    # d2state_dw2 = SVector{4, Float64}(astate[D2STATE_DW2_INDICES])

    # Compute
    # linear interpolation
    # controls_ = interpolate_linear_points(CONTROL_EVAL_TIMES[control_step],
    #                                       CONTROL_EVAL_TIMES[control_step + 1],
    #                                       time,
    #                                       controls[control_step],
    #                                       controls[control_step + 1])

    # zero-order-hold
    controls_ = controls[control_step]

    neg_i_hamiltonian = OMEGA * NEG_I_H_S + controls_[1] * NEG_I_H_C1

    delta_state = neg_i_hamiltonian * state
    # delta_dstate_dw = neg_i_H_S * state + neg_i_hamiltonian * dstate_dw
    # delta_d2state_dw2 = neg_i_H_S * dstate_dw + neg_i_hamiltonian * d2state_dw2

    delta_astate = [
        delta_state;
        # delta_dstate_dw;
        # delta_d2state_dw2
    ]

    return delta_astate
end

"""
Infidelity captures the difference between
the final state and the target final state.
"""
function infidelity(final_astate)
    final_state = final_astate[STATE_INDICES]
    inner_product = TARGET_STATE_DAGGER * final_state
    fidelity = inner_product[1]^2 + inner_product[2]^2
    infidelity_ = 1 - fidelity
    return infidelity_
end

"""
Infidelity Robustness captures the sensitivity of the difference between
the final state and the target final state to the parameter OMEGA.
"""
function infidelity_robustness(final_astate)
    d2final_state_dw2 = final_astate[D2STATE_DW2_INDICES]
    inner_product = TARGET_STATE_DAGGER * d2final_state_dw2
    d2j_dw2 = 2 * sqrt(inner_product[1]^2 + inner_product[2]^2)
    return d2j_dw2 * 1e-3
end

function augment_cost(cost, multiplier, lagrange_multiplier)
    # return cost * lagrange_multiplier + (multiplier / 2) * cost ^ 2
    return cost
end

"""
Evolve a state and compute associated costs.
"""
function evolve(controls, pstate, reporter)
    # Evolve the state.
    t = 0
    dt = pstate.dt
    N = pstate.evolution_time / dt
    final_astate = pstate.initial_astate
    control_step = 1
    for i=1:N
        final_astate = rk3_step(rhs, final_astate, controls, t, dt, control_step)
        t += dt
        if t > pstate.control_eval_times[control_step + 1]
            control_step += 1
        end
    end
    final_state = final_astate[STATE_INDICES]
    # println(final_state)

    # Compute costs.
    infidelity_cost = infidelity(final_astate)
    # infidelity_robustness_cost = infidelity_robustness(final_astate)
    augmented_cost = (
        augment_cost(infidelity_cost, pstate.cost_multipliers[1], pstate.lagrange_multipliers[1])
        # + augment_cost(infidelity_robustness_cost, pstate.cost_multipliers[2], pstate.lagrange_multipliers[2])
    )

    # Report.
    reporter.costs = [
        infidelity_cost
        # infidelity_robustness_cost
    ]
    reporter.cost = augmented_cost

    return augmented_cost
end

function main()
    # Grab relevant information.
    control_count = CONTROL_COUNT
    control_eval_count = CONTROL_EVAL_COUNT
    controls = INITIAL_CONTROLS
    cost_multipliers = INITIAL_COST_MULTIPLIERS
    dt = DT
    evolution_time = EVOLUTION_TIME
    control_eval_times = Array(range(0., stop=evolution_time, length=control_eval_count))
    initial_astate = INITIAL_ASTATE
    iteration_count = ITERATION_COUNT
    lagrange_multipliers = INITIAL_LAGRANGE_MULTIPLIERS
    learning_rate = LEARNING_RATE
    save = SAVE
    # Construct the program state.
    pstate = PState(
        control_count, control_eval_count, control_eval_times,
        cost_multipliers, dt, evolution_time,
        initial_astate, lagrange_multipliers,
    )
    reporter = Reporter(0, INITIAL_COSTS)
    # Create a wrapper around the cost function. Zygote takes gradients with repsect to all arguments.
    evolve_(controls_) = evolve(controls_, pstate, reporter)
    # Generate save file path
    if save
        save_file_path = generate_save_file_path(EXPERIMENT_NAME, SAVE_PATH)
        println("saving this optimization to $save_file_path")
        h5write(save_file_path, "controls", zeros(iteration_count, control_eval_count, control_count))
    end

    # AL loop
    # for i = 1:5
    #     # iLQR loop
    #     println("costs")
    #     println(reporter.costs)
    #     println("cost_multipliers")
    #     println(pstate.cost_multipliers)
    #     println("lagrange_multipliers")
    #     println(pstate.lagrange_multipliers)

    for iteration = 1:iteration_count
        (dcontrols,) = Zygote.gradient(evolve_, controls)
        dcontrols_norm = norm(dcontrols)
        controls = controls - dcontrols * learning_rate
        @printf("%05d %f %f\n", iteration, reporter.cost, dcontrols_norm,)
        if save
            h5open(save_file_path, "r+") do save_file
                save_file["controls"][iteration, :, :] = controls
            end
        end
    end

    #     pstate.cost_multipliers = pstate.cost_multipliers * COST_MULTIPLIER_STEP
    #     pstate.lagrange_multipliers = (
    #         pstate.lagrange_multipliers + pstate.cost_multipliers .* reporter.costs
    #     )
    # end

    # evolve_(controls)
    # cfg = ForwardDiff.GradientConfig(evolve_, controls)
    # out = zeros(11, control_eval_count, control_count)
    # ForwardDiff.gradient!(out[11, :, :], evolve_, controls, cfg)
    # Zygote.gradient(evolve_, controls)
    # function test()
    #     for i=1:10
    #         # evolve_(controls)
    #         # ForwardDiff.gradient!(out[i, :, :], evolve_, controls, cfg)
    #         Zygote.gradient(evolve_, controls)
    #     end
    #     return
    # end

    # @time test()
end

end

max_stars_repo_name: SchusterLab/rbqoc

max_issues_repo_name: SchusterLab/rbqoc

max_forks_repo_name: SchusterLab/rbqoc


text:
#
const BASE_UUID = UUID("985c312f-129b-4acd-9e40-cb629d184183")
const BASE_PARAMS = load_system()

function initialise_settings()::Settings
    settings = Settings()
	settings.uuid = BASE_UUID
	settings.means_tested_routing = modelled_phase_in
    settings.run_name="run-$(date_string())"
	settings.income_data_source = ds_frs
	settings.dump_frames = false
	settings.do_marginal_rates = true
	settings.requested_threads = 4
	settings.dump_frames = true
	return settings
end

const BASE_SETTINGS = initialise_settings()

## FIXME POVERTY LINE!
function make_base_results()
    obs = Observable(
		Progress(BASE_SETTINGS.uuid, "",0,0,0,0))
	tot = 0
	of = on(obs) do p
		tot += p.step
		PROGRESS[p.uuid] = (progress=p,total=tot)
	end
    settings = deepcopy(BASE_SETTINGS)
	results = do_one_run( settings, [BASE_PARAMS], obs )
    println( BASE_PARAMS)
    println( settings )
    settings.poverty_line = make_poverty_line( results.hh[1], settings )
    outf = summarise_frames( results, settings )
	gl = make_gain_lose( results.hh[1], results.hh[1], settings )
	exres = calc_examples( BASE_PARAMS, BASE_PARAMS, settings )
    println( settings.poverty_line)
	aout = AllOutput( BASE_SETTINGS.uuid, "DEFAULT", results, outf, gl, exres )
	return aout;
end

const BASE_RESULTS = make_base_results()
const BASE_TEXT_OUTPUT = results_to_html( BASE_UUID, BASE_RESULTS, BASE_RESULTS )


text:
module Lasem

    import Cairo: CairoContext  # for rendering

    const liblasem = "liblasem-0.6.5.dylib"
    const libgobject = "libgobject-2.0.0.dylib"

################################################################################
    # low-level wrapper for lasem functions

    immutable LsmBox
        x::Cdouble
        y::Cdouble
        width::Cdouble
        height::Cdouble
    end

    # general
    destroy(ptr::Ptr{Void}) = ccall((:g_object_unref, libgobject), Void, (Ptr{Void},), ptr)

    # document
    function lsm_dom_document_new_from_path(filename)
        ccall((:lsm_dom_document_new_from_path, liblasem), Ptr{Void}, (Ptr{Uint8}, Ptr{Void}), filename, C_NULL)
    end

    function lsm_dom_document_new_from_memory(buffer::Ptr{Uint8}, bufsize=length(buffer))
        ccall((:lsm_dom_document_new_from_memory, liblasem), Ptr{Void}, (Ptr{Uint8}, Uint32, Ptr{Void}), buffer, bufsize, C_NULL)
    end

    function lsm_mathml_document_new_from_itex(buffer, bufsize=length(buffer))
        ccall((:lsm_mathml_document_new_from_itex, liblasem), Ptr{Void}, (Ptr{Uint8}, Uint32, Ptr{Void}), buffer, bufsize, C_NULL)
    end

    function lsm_dom_document_create_view(doc::Ptr{Void})
        ccall((:lsm_dom_document_create_view, liblasem), Ptr{Void}, (Ptr{Void}, ), doc)
    end

    # view
    function lsm_dom_view_set_resolution(view::Ptr{Void}, ppi::Float64)
        ccall((:lsm_dom_view_set_resolution, liblasem), Void, (Ptr{Void}, Cdouble, ), view, ppi)
    end

    function lsm_dom_view_set_viewport_pixels(view::Ptr{Void}, viewport::LsmBox)
        ccall((:lsm_dom_view_set_viewport_pixels, liblasem), Void, (Ptr{Void}, Ptr{LsmBox} ), view, &viewport)
    end

    function lsm_dom_view_get_size_pixels(view::Ptr{Void})
        w_and_h = Uint32[0., 0.]
        ccall((:lsm_dom_view_get_size_pixels, liblasem), Void, (Ptr{Void}, Ptr{Uint32}, Ptr{Uint32}, Ptr{Void} ), view, pointer(w_and_h, 1), pointer(w_and_h, 2), C_NULL)

        return w_and_h[1], w_and_h[2]
    end

    function lsm_dom_view_get_size(view::Ptr{Void})
        w_and_h = Float64[0., 0.]
        ccall((:lsm_dom_view_get_size, liblasem), Void, (Ptr{Void}, Ptr{Float64}, Ptr{Float64}, Ptr{Void} ), view, pointer(w_and_h, 1), pointer(w_and_h, 2), C_NULL)

        return w_and_h[1], w_and_h[2]
    end

    function lsm_dom_view_render(view::Ptr{Void}, cairo::CairoContext, x::Float64=0., y::Float64=0.)
        ccall((:lsm_dom_view_render, liblasem), Void, (Ptr{Void}, Ptr{Void}, Cdouble, Cdouble ), view, cairo.ptr, x, y)
    end

################################################################################
    # TODO
    # high level interface
    type LsmDoc
        ptr::Ptr{Void}
    end

    type LsmView
        ptr::Ptr{Void}

        function LsmView(doc::LsmDoc)
            view = lsm_dom_document_create_view(doc.ptr)
            view != C_NULL || error("Creating lasem view failed.")

            new(view)
        end
    end

    function render(view::LsmView, ctx::CairoContext)
        lsm_dom_view_render(view.ptr, ctx)
        return nothing
    end

end # module


text:
using ..ShallowWater: TurbulenceClosure, LinearDrag, ConstantViscosity

function null_init_state!(
    ::HomogeneousBox,
    ::TurbulenceClosure,
    state,
    aux,
    local_geometry,
    t,
)
    T = eltype(state.U)
    state.U = @SVector zeros(T, 2)
    state.η = 0
    return nothing
end

η_lsw(x, y, t) = cos(π * x) * cos(π * y) * cos(√2 * π * t)
u_lsw(x, y, t) = 2^(-0.5) * sin(π * x) * cos(π * y) * sin(√2 * π * t)
v_lsw(x, y, t) = 2^(-0.5) * cos(π * x) * sin(π * y) * sin(√2 * π * t)

function lsw_init_state!(
    m::ShallowWaterModel,
    p::HomogeneousBox,
    state,
    aux,
    local_geometry,
    t,
)

    coords = local_geometry.coord

    state.U = @SVector [
        u_lsw(coords[1], coords[2], t),
        v_lsw(coords[1], coords[2], t),
    ]

    state.η = η_lsw(coords[1], coords[2], t)

    return nothing
end

v_lkw(x, y, t) = 0
u_lkw(x, y, t) = exp(-0.5 * y^2) * exp(-0.5 * (x - t + 5)^2)
η_lkw(x, y, t) = 1 + u_lkw(x, y, t)

function lkw_init_state!(
    m::ShallowWaterModel,
    p::HomogeneousBox,
    state,
    aux,
    local_geometry,
    t,
)

    coords = local_geometry.coord

    state.U = @SVector [
        u_lkw(coords[1], coords[2], t),
        v_lkw(coords[1], coords[2], t),
    ]

    state.η = η_lkw(coords[1], coords[2], t)

    return nothing
end

R₋(ϵ) = (-1 - sqrt(1 + (2 * π * ϵ)^2)) / (2ϵ)
R₊(ϵ) = (-1 + sqrt(1 + (2 * π * ϵ)^2)) / (2ϵ)
D(ϵ) =
    (R₊(ϵ) * (exp(R₋(ϵ)) - 1) + R₋(ϵ) * (1 - exp(R₊(ϵ)))) /
    (exp(R₊(ϵ)) - exp(R₋(ϵ)))
R₂(x₁, ϵ) =
    (1 / D(ϵ)) * (
        (
            (R₊(ϵ) * (exp(R₋(ϵ)) - 1)) * exp(R₊(ϵ) * x₁) +
            (R₋(ϵ) * (1 - exp(R₊(ϵ)))) * exp(R₋(ϵ) * x₁)
        ) / (exp(R₊(ϵ)) - exp(R₋(ϵ)))
    )
R₁(x₁, ϵ) =
    (π / D(ϵ)) * (
        1 .+
        (
            (exp(R₋(ϵ)) - 1) * exp(R₊(ϵ) * x₁) .+
            (1 - exp(R₊(ϵ))) * exp(R₋(ϵ) * x₁)
        ) / (exp(R₊(ϵ)) - exp(R₋(ϵ)))
    )

𝒱(x₁, y₁, ϵ) = R₂(x₁, ϵ) * sin.(π * y₁)
𝒰(x₁, y₁, ϵ) = -R₁(x₁, ϵ) * cos.(π * y₁)
ℋ(x₁, y₁, ϵ, βᵖ, fₒ, γ) =
    (R₂(x₁, ϵ) / (π * fₒ)) * γ * cos(π * y₁) +
    (R₁(x₁, ϵ) / π) *
    (sin(π * y₁) * (1.0 + βᵖ * (y₁ - 0.5)) + (βᵖ / π) * cos(π * y₁))

function gyre_init_state!(
    m::SWModel,
    p::HomogeneousBox,
    T::LinearDrag,
    state,
    aux,
    local_geometry,
    t,
)

    coords = local_geometry.coord

    FT = eltype(state)
    τₒ = p.τₒ
    fₒ = m.fₒ
    β = m.β
    Lˣ = p.Lˣ
    Lʸ = p.Lʸ
    H = p.H

    γ = T.λ

    βᵖ = β * Lʸ / fₒ
    ϵ = γ / (Lˣ * β)

    _grav::FT = grav(m.param_set)

    uˢ(ϵ) = (τₒ * D(ϵ)) / (H * γ * π)
    hˢ(ϵ) = (fₒ * Lˣ * uˢ(ϵ)) / _grav

    u = uˢ(ϵ) * 𝒰(coords[1] / Lˣ, coords[2] / Lʸ, ϵ)
    v = uˢ(ϵ) * 𝒱(coords[1] / Lˣ, coords[2] / Lʸ, ϵ)
    h = hˢ(ϵ) * ℋ(coords[1] / Lˣ, coords[2] / Lʸ, ϵ, βᵖ, fₒ, γ)

    state.U = @SVector [H * u, H * v]

    state.η = h

    return nothing
end

t1(x, δᵐ) = cos((√3 * x) / (2 * δᵐ)) + (√3^-1) * sin((√3 * x) / (2 * δᵐ))
t2(x, δᵐ) = 1 - exp((-x) / (2 * δᵐ)) * t1(x, δᵐ)
t3(y, Lʸ) = π * sin(π * y / Lʸ)
t4(x, Lˣ, C) = C * (1 - x / Lˣ)

η_munk(x, y, Lˣ, Lʸ, δᵐ, C) = t4(x, Lˣ, C) * t3(y, Lʸ) * t2(x, δᵐ)

function gyre_init_state!(
    m::SWModel,
    p::HomogeneousBox,
    V::ConstantViscosity,
    state,
    aux,
    local_geometry,
    t,
)

    coords = local_geometry.coord

    FT = eltype(state.U)
    _grav::FT = grav(m.param_set)

    τₒ = p.τₒ
    fₒ = m.fₒ
    β = m.β
    Lˣ = p.Lˣ
    Lʸ = p.Lʸ
    H = p.H

    ν = V.ν

    δᵐ = (ν / β)^(1 / 3)
    C = τₒ / (_grav * H) * (fₒ / β)

    state.η = η_munk(coords[1], coords[2], Lˣ, Lʸ, δᵐ, C)
    state.U = @SVector zeros(FT, 2)

    return nothing
end

max_forks_repo_name: jm-c/CLIMA


text:
using Logging, Random
using JuMP, Gurobi
using DecisionProgramming

Random.seed!(42)

const dᴾ = 1    # Decision node: range for number of patents
const cᵀ = 2    # Chance node:   technical competitiveness
const dᴬ = 3    # Decision node: range for number of applications
const cᴹ = 4    # Chance node:   market share
const DP_states = ["0-3 patents", "3-6 patents", "6-9 patents"]
const CT_states = ["low", "medium", "high"]
const DA_states = ["0-5 applications", "5-10 applications", "10-15 applications"]
const CM_states = ["low", "medium", "high"]

S = States([
    (length(DP_states), [dᴾ]),
    (length(CT_states), [cᵀ]),
    (length(DA_states), [dᴬ]),
    (length(CM_states), [cᴹ]),
])
C = Vector{ChanceNode}()
D = Vector{DecisionNode}()
V = Vector{ValueNode}()
X = Vector{Probabilities}()
Y = Vector{Consequences}()

I_DP = Vector{Node}()
push!(D, DecisionNode(dᴾ, I_DP))

I_CT = [dᴾ]
X_CT = zeros(S[dᴾ], S[cᵀ])
X_CT[1, :] = [1/2, 1/3, 1/6]
X_CT[2, :] = [1/3, 1/3, 1/3]
X_CT[3, :] = [1/6, 1/3, 1/2]
push!(C, ChanceNode(cᵀ, I_CT))
push!(X, Probabilities(cᵀ, X_CT))

I_DA = [dᴾ, cᵀ]
push!(D, DecisionNode(dᴬ, I_DA))

I_CM = [cᵀ, dᴬ]
X_CM = zeros(S[cᵀ], S[dᴬ], S[cᴹ])
X_CM[1, 1, :] = [2/3, 1/4, 1/12]
X_CM[1, 2, :] = [1/2, 1/3, 1/6]
X_CM[1, 3, :] = [1/3, 1/3, 1/3]
X_CM[2, 1, :] = [1/2, 1/3, 1/6]
X_CM[2, 2, :] = [1/3, 1/3, 1/3]
X_CM[2, 3, :] = [1/6, 1/3, 1/2]
X_CM[3, 1, :] = [1/3, 1/3, 1/3]
X_CM[3, 2, :] = [1/6, 1/3, 1/2]
X_CM[3, 3, :] = [1/12, 1/4, 2/3]
push!(C, ChanceNode(cᴹ, I_CM))
push!(X, Probabilities(cᴹ, X_CM))

# Dummy value node
push!(V, ValueNode(5, [cᴹ]))
push!(Y, Consequences(5, zeros(S[cᴹ])))

@info("Validate influence diagram.")
validate_influence_diagram(S, C, D, V)
sort!.((C, D, V, X, Y), by = x -> x.j)

@info("Creating path probability.")
P = DefaultPathProbability(C, X)

@info("Defining DecisionModel")
model = Model()
z = DecisionVariables(model, S, D)

@info("Creating problem specific constraints and expressions")

function variables(model::Model, dims::AbstractVector{Int}; binary::Bool=false)
    v = Array{VariableRef}(undef, dims...)
    for i in eachindex(v)
        v[i] = @variable(model, binary=binary)
    end
    return v
end

n_T = 5                     # number of technology projects
n_A = 5                     # number of application projects
I_t = rand(n_T)*0.1         # costs of technology projects
O_t = rand(1:3,n_T)         # number of patents for each tech project
I_a = rand(n_T)*2           # costs of application projects
O_a = rand(2:4,n_T)         # number of applications for each appl. project

V_A = rand(S[cᴹ], n_A).+0.5 # Value of an application
V_A[1, :] .+= -0.5          # Low market share: less value
V_A[3, :] .+= 0.5           # High market share: more value

x_T = variables(model, [S[dᴾ]...,n_T]; binary=true)
x_A = variables(model, [S[dᴾ]...,S[cᵀ]...,S[dᴬ]..., n_A]; binary=true)

M = 20                      # a large constant
ε = 0.5*minimum([O_t O_a])  # a helper variable, allows using ≤ instead of < in constraints (28b) and (29b)

q_P = [0, 3, 6, 9]          # limits of the technology intervals
q_A = [0, 5, 10, 15]        # limits of the application intervals
z_dP = z.z[1]
z_dA = z.z[2]

@constraint(model, [i=1:3],
    sum(x_T[i,t] for t in 1:n_T) <= z_dP[i]*n_T)            #(25)
@constraint(model, [i=1:3, j=1:3, k=1:3],
    sum(x_A[i,j,k,a] for a in 1:n_A) <= z_dP[i]*n_A)        #(26)
@constraint(model, [i=1:3, j=1:3, k=1:3],
    sum(x_A[i,j,k,a] for a in 1:n_A) <= z_dA[i,j,k]*n_A)    #(27)

#(28a)
@constraint(model, [i=1:3],
    q_P[i] - (1 - z_dP[i])*M <= sum(x_T[i,t]*O_t[t] for t in 1:n_T))
#(28b)
@constraint(model, [i=1:3],
    sum(x_T[i,t]*O_t[t] for t in 1:n_T) <= q_P[i+1] + (1 - z_dP[i])*M - ε)
#(29a)
@constraint(model, [i=1:3, j=1:3, k=1:3],
    q_A[k] - (1 - z_dA[i,j,k])*M <= sum(x_A[i,j,k,a]*O_a[a] for a in 1:n_A))
#(29b)
@constraint(model, [i=1:3, j=1:3, k=1:3],
    sum(x_A[i,j,k,a]*O_a[a] for a in 1:n_A) <= q_A[k+1] + (1 - z_dA[i,j,k])*M - ε)

@constraint(model, [i=1:3, j=1:3, k=1:3], x_A[i,j,k,1] <= x_T[i,1])
@constraint(model, [i=1:3, j=1:3, k=1:3], x_A[i,j,k,2] <= x_T[i,1])
@constraint(model, [i=1:3, j=1:3, k=1:3], x_A[i,j,k,2] <= x_T[i,2])

@info("Creating model objective.")
patent_investment_cost = @expression(model, [i=1:S[1]], sum(x_T[i, t] * I_t[t] for t in 1:n_T))
application_investment_cost = @expression(model, [i=1:S[1], j=1:S[2], k=1:S[3]], sum(x_A[i, j, k, a] * I_a[a] for a in 1:n_A))
application_value = @expression(model, [i=1:S[1], j=1:S[2], k=1:S[3], l=1:S[4]], sum(x_A[i, j, k, a] * V_A[l, a] for a in 1:n_A))
@objective(model, Max, sum( sum( P((i,j,k,l)) * (application_value[i,j,k,l] - application_investment_cost[i,j,k]) for j in 1:S[2], k in 1:S[3], l in 1:S[4] ) - patent_investment_cost[i] for i in 1:S[1] ))

@info("Starting the optimization process.")
optimizer = optimizer_with_attributes(
    () -> Gurobi.Optimizer(Gurobi.Env()),
    "IntFeasTol"      => 1e-9,
    "LazyConstraints" => 1,
)
set_optimizer(model, optimizer)
optimize!(model)

@info("Extracting results.")
Z = DecisionStrategy(z)

@info("Printing decision strategy:")
print_decision_strategy(S, Z)

@info("Extracting path utilities")
struct PathUtility <: AbstractPathUtility
    data::Array{AffExpr}
end
Base.getindex(U::PathUtility, i::Int) = getindex(U.data, i)
Base.getindex(U::PathUtility, I::Vararg{Int,N}) where N = getindex(U.data, I...)
(U::PathUtility)(s::Path) = value.(U[s...])

path_utility = [@expression(model,
    sum(x_A[s[1:3]..., a] * (V_A[s[4], a] - I_a[a]) for a in 1:n_A) -
    sum(x_T[s[1], t] * I_t[t] for t in 1:n_T)) for s in paths(S)]
U = PathUtility(path_utility)

@info("Computing utility distribution.")
udist = UtilityDistribution(S, P, U, Z)

@info("Printing utility distribution.")
print_utility_distribution(udist)

@info("Printing statistics")
print_statistics(udist)


text:
using Base: Float64
import SparseArrays.SparseVector
import NearestNeighbors: KDTree, nn, inrange, knn
import LinearAlgebra.Vector
import DataStructures: MutableBinaryMaxHeap, top_with_handle, pop!, update!

function construct_multicolor_ordering(input_array::AbstractVector, ρh::Real, tree_function)
    RT = real_type(eltype(input_array))
    # Vector (colors) of Vectors of supernodes
    out = Vector{eltype(input_array)}[]

    assigned = falses(length(input_array))
    # While not all nodes are assigned to a color
    tree = tree_function(center.(input_array))
    heap = MutableBinaryMaxHeap(fill(typemax(RT), length(input_array)))
    while !isempty(heap)
        # add new color
        push!(out, typeof(input_array)(undef, 0))
        # reset heap values
        # Presently a little inefficient since we revisit every entry of the multicolor
        # ordering. Still doesn't change the asymptotic complexity
        for k = 1 : length(input_array)
            # only assigns a value to points that have not already been assigned and thus
            # removed from the heap
            !assigned[k] && (heap[k] = typemax(RT))
        end
        # do furthest point sampling, finish if either the heap is empty, or there are no
        # sufficiently distant points left
        while !isempty(heap) && first(heap) > 2 * ρh
            # get the id of the new pivot
            ~, top_id = top_with_handle(heap)
            # remove the new pivot from the heap
            pop!(heap); assigned[top_id] = true; push!(out[end], input_array[top_id])
            # add the new pivot to the
            # obtain the affected nodes and their distance to the pivot
            number_affected_nodes = length(inrange(tree, center(input_array[top_id]), 2 * ρh))
            affected_nodes, distances = knn(tree, center(input_array[top_id]), number_affected_nodes)
            # Update the distances of nodes that could be affedcted.
            for (node, dist) in zip(affected_nodes, distances)
                # update distance if node is still in heap
                !assigned[node] && update!(heap, node, dist)
            end
        end
    end
    return out
end

# function that can directly take an array of arrays (corresponding to different scales) as supernodes
function construct_multicolor_ordering(input_arrays::AbstractVector{<:AbstractVector}, ρh::AbstractVector{<:Real}, tree_function)
    # Both lengths should be equal to total number of scales
    @assert length(ρh) == length(input_arrays)
    q = length(ρh)
    out = Vector{typeof(input_arrays)}(undef, q)
    for k = 1 : q
        out[k] = construct_multicolor_ordering(input_arrays[k], ρh[k], tree_function)
    end
    return vcat(out...)
end


text:
push!(LOAD_PATH,"../src/")

using Documenter
using EcologicalNetworks
using Random

makedocs(
    sitename = "EcologicalNetworks",
    authors = "Timothée Poisot",
    modules = [EcologicalNetworks],
    pages = [
        "Index" => "index.md",
        "Interface" => [
            "Types" => "interface/types.md",
            "Conversions" => "interface/conversions.md",
            "Core functions" => "interface/highlevel.md",
            "AbstractArray" => "interface/abstractarray.md",
            "Iterate" => "interface/iterate.md"
        ],
        "Examples" => [
            "Modularity" => "examples/modularity.md",
            "Centrality" => "examples/centrality.md",
            "Integration with Mangal" => "examples/mangal.md",
            "Network beta-diversity" => "examples/beta-diversity.md",
            "Extinctions" => "examples/extinctions.md"
        ],
        "Basic measures" => [
            "Links" => "properties/links.md",
            "Modularity" => "properties/modularity.md",
            "Nestedness" => "properties/nestedness.md",
            "Motifs" => "properties/motifs.md",
            "Centrality and paths" => "properties/paths.md",
            "Overlap and similarity" => "properties/overlap.md"
        ],
        "Advanced information" => [
            "Beta-diversity" => "properties/betadiversity.md",
            "Resilience" => "properties/resilience.md",
            "Information theory" => "properties/information.md"
        ],
        "Generating networks" => [
            "Null models" => "random/null.md",
            "Structural models" => "random/structure.md"
        ]
    ]
)

deploydocs(
deps   = Deps.pip("pygments", "python-markdown-math"),
repo   = "github.com/EcoJulia/EcologicalNetworks.jl.git",
devbranch = "main",
push_preview = true
)


text:
"""
    Model

Abstract type represnting a BoundedFluidQueue or DiscretisedFluidQueue
"""
abstract type Model end

"""
    Phase

Represents a phase of the fluid queue.

# Arguments:
- `c::Float64`: the rate of change of the fluid X(t) associated with the phase
- `m::Int64`: either -1 or 1, membership of the phase with either the set of positive or negative phases.
    Default is positive, i.e. 1. Only relevant for FRAPApproximation.
- `lpm::Bool`: whether the phase has a point mass at the left/lower boundary
- `rpm::Bool`: whether the phase has a point mass at the right/upper boundary
"""
struct Phase
    c::Float64
    m::Int64 # membership to +1 or -1
    lpm::Bool # left point mass
    rpm::Bool # right point mass
    function Phase(c::Float64,m::Int,lpm::Bool,rpm::Bool)
        !((m===1)||(m===-1))&&throw(DomainError("m is 1, -1 only"))
        !((sign(c)==m)||(sign(c)==0.0))&&throw(DomainError("sign(c) must be m or 0"))
        !((c<0.0)==lpm||(c==0.0))&&throw(DomainError("negative phases must have lpm=true"))
        !((c>0.0)==rpm||(c==0.0))&&throw(DomainError("positive phases must have rpm=true"))
        return new(c,m,lpm,rpm)
    end
end
Phase(c::Float64) = Phase(c,-1+2*Int(_strictly_pos(c)),c<=0,c>=0)
# Phase(c::Int) = Phase(Float64(c),-1+2*Int(_strictly_pos(c)),c<=0,c>=0)
Phase(c::Float64,m::Int) = Phase(c,m,c<=0,c>=0)

"""
    const PhaseSet = Array{Phase,1}

A container for phases.
"""
const PhaseSet = Array{Phase,1}
PhaseSet(c::Array{Float64,1}) = [Phase(c[i]) for i in 1:length(c)]
PhaseSet(c::Array{Float64,1},m::Array{Int,1}) = (length(m)==length(c))&&[Phase(c[i],m[i]) for i in 1:length(c)]
"""
    PhaseSet(c::Array{Float64, 1}[, m::Array{Int, 1}, lpm::BitArray, rpm::BitArray])

Construct an Array of phases.

# Arguments:
- `c`: Array of rates
- `m`: (optional) array of memberships, see Phase()
- `lpm`: (optional) array of left point mass indicators, see Phase()
- `rpm`: (optional) array of right point mass indicators, see Phase()
"""
PhaseSet(c::Array{Float64,1},m::Array{Int,1},lpm::BitArray,rpm::BitArray) =
    (length(m)==length(c)==length(lpm)==length(rpm))&&[Phase(c[i],m[i],lpm[i],rpm[i]) for i in 1:length(c)]

"""
    n_phases(S::PhaseSet)

Return the number of phases from a PhaseSet, BoundedFluidQueue, or DiscretisedFluidQueue
"""
n_phases(S::PhaseSet) = length(S)
"""
    rates(S::PhaseSet, i::Int)

Return the rate of the phase from a PhaseSet, BoundedFluidQueue, or DiscretisedFluidQueue
"""
rates(S::PhaseSet,i::Int) = S[i].c
rates(S::PhaseSet) = [S[i].c for i in 1:n_phases(S)]

negative_phases(S::PhaseSet,i::Int) = S[i].c.<0.0
negative_phases(S::PhaseSet) = [S[i].c.<0.0 for i in 1:n_phases(S)]
positive_phases(S::PhaseSet,i::Int) = S[i].c.>0.0
positive_phases(S::PhaseSet) = [S[i].c.>0.0 for i in 1:n_phases(S)]

"""
    membership(S::PhaseSet, i::Int)

Return the membership of phases from a PhaseSet
"""
membership(S::PhaseSet,i::Int) = S[i].m
membership(S::PhaseSet) = [S[i].m for i in 1:n_phases(S)]
"""
    phases(S::PhaseSet)

Return the iterator 1:n_phases(S), for a PhaseSet, BoundedFluidQueue, or DiscretisedFluidQueue
"""
phases(S::PhaseSet) = 1:n_phases(S)

_strictly_neg(x::Float64) = (x<0.0) || (x.===-0.0)
_strictly_pos(x::Float64) = !_strictly_neg(x)
_strictly_neg(x::Int) = (x===0) ? true : throw(DomainError("invalid membership detected"))
_strictly_pos(x::Int) = (x===0) ? true : throw(DomainError("invalid membership detected"))

_has_left_boundary(i::Phase) = i.lpm
_has_right_boundary(i::Phase) = i.rpm
_has_left_boundary(S::PhaseSet,i::Int) = S[i].lpm
_has_right_boundary(S::PhaseSet,i::Int) = S[i].rpm
_has_left_boundary(S::PhaseSet) = _has_left_boundary.(S)
_has_right_boundary(S::PhaseSet) = _has_right_boundary.(S)

"""
    N₋(S::PhaseSet)

Return the number of left/lower point masses of a PhaseSet, BoundedFluidQueue or DiscretisedFluidQueue
"""
N₋(S::PhaseSet) = sum(_has_left_boundary(S))
"""
    N₊(S::PhaseSet) = begin

Return the number of right/upper point masses of a PhaseSet, BoundedFluidQueue or DiscretisedFluidQueue
"""
N₊(S::PhaseSet) = sum(_has_right_boundary(S))

checksquare(A::AbstractArray{<:Any,2}) = !(size(A,1)==size(A,2)) ? throw(DomainError(A," must be square")) : nothing

function _fluid_queue_checks(T,S)
    checksquare(T)
    !(size(T,1)==length(S))&&throw(DomainError("PhaseSet must have length the size(T,1)"))
    return
end

"""
    BoundedFluidQueue <: Model

Constructor for a fluid queue model.

# Arguments:
- `T::Array{Float64, 2}`: Generator of the phase process
- `S::PhaseSet`: An array of phases describing the evolution of the fluid level in each phase.
- `P_lwr::Array{Float64,2}`: Probabilities of phase change upon hitting the lower boundary.
    A |S₋| by |S| array.
- `P_upr::Array{Float64,2}`: Probabilities of phase change upon hitting the upper boundary.
    A |S₊| by |S| array.
- `b::Float64`: the upper boundary of the fluid level (the lower boundary is 0 by default).
"""
struct BoundedFluidQueue <: Model
    T::Array{Float64,2}
    S::PhaseSet
    P_lwr::Matrix{Float64}
    P_upr::Matrix{Float64}
    b::Float64
    function BoundedFluidQueue(T::Array{Float64,2}, S::PhaseSet, P_lwr::Matrix{Float64}, P_upr::Matrix{Float64},b::Float64)
        _fluid_queue_checks(T,S)
        !(all(sum(T,dims=2).≈0.0))&&@warn "row sums of T should be 0"
        !all(sum(P_lwr,dims=2).≈1.0)&&throw(DomainError("P_lwr sum should be 1"))#@warn "row sums of P_lwr should be 1.0"
        !all(sum(P_upr,dims=2).≈1.0)&&throw(DomainError("P_lwr sum should be 1"))#@warn "row sums of P_upr should be 1.0"
        !((sum(rates(S).<0.0)==size(P_lwr,1))&&(size(T,2)==size(P_lwr,2)))&&throw(DomainError("P_lwr must be |S₋| by |S|"))
        !((sum(rates(S).>0.0)==size(P_upr,1))&&(size(T,2)==size(P_upr,2)))&&throw(DomainError("P_upr must be |S₊| by |S|"))
        (b<0)&&throw(DomainError("upperboundary, b, must be positive"))
        return new(T,S,P_lwr,P_upr,b)
    end
end
"""
    BoundedFluidQueue(T::Array{Float64,2}, S::PhaseSet, b::Float64)

Constructor for a fluid queue model with regulated boundaries by default

# Arguments:
- `T::Array{Float64, 2}`: Generator of the phase process
- `S::PhaseSet`: An array of phases describing the evolution of the fluid level in each phase
- `b::Float64`: the upper boundary of the fluid level (lower boundary is 0 by default).
"""
function BoundedFluidQueue(T::Array{Float64,2}, S::PhaseSet, b::Float64)
    _fluid_queue_checks(T,S)
    P_lwr = zeros(sum(negative_phases(S)),size(T,2))
    P_upr = zeros(sum(positive_phases(S)),size(T,2))
    i_lwr = 0
    i_upr = 0
    C = rates(S)
    for j in 1:size(T,2)
        if C[j] < 0.0
            i_lwr += 1
            P_lwr[i_lwr,j] = 1.0
        elseif C[j] > 0.0
            i_upr += 1
            P_upr[i_upr,j] = 1.0
        end
    end
    return BoundedFluidQueue(T,S,P_lwr,P_upr,b)
end
"""
    BoundedFluidQueue(T::Array{Float64,2},c::Array{Float64,1},b::Float64)

Alias to `BoundedFluidQueue(T,PhaseSet(c),b)`.
"""
BoundedFluidQueue(T::Array{Float64,2},c::Array{Float64,1},b::Float64) = BoundedFluidQueue(T,PhaseSet(c),b)
BoundedFluidQueue(T::Array{Float64,2}, c::Array{Float64,1}, P_lwr::Matrix{Float64}, P_upr::Matrix{Float64}, b::Float64) =
    BoundedFluidQueue(T,PhaseSet(c),P_lwr,P_upr,b)

rates(m::Model) = rates(m.S)
rates(m::Model,i::Int) = rates(m.S,i)

negative_phases(m::Model,i::Int) = negative_phases(m.S,i)
negative_phases(m::Model) = negative_phases(m.S)
positive_phases(m::Model,i::Int) = positive_phases(m.S,i)
positive_phases(m::Model) = positive_phases(m.S)

n_phases(m::Model) = n_phases(m.S)
phases(m::Model) = 1:n_phases(m.S)
N₋(m::Model) = N₋(m.S)
N₊(m::Model) = N₊(m.S)

function _duplicate_zero_states_C(T::Array{Float64,2},C::Array{Float64,1})
    n_0 = sum(C.==0)

    # assign rates for the augmented model
    C_aug = Array{Float64,1}(undef,length(C)+n_0)
    c_zero = 0
    plus_idx = falses(length(C_aug)) # zero states associated with +
    neg_idx = falses(length(C_aug)) # zero states associated with -
    for i in 1:length(C)
        C_aug[i+c_zero] = C[i]
        if C[i] == 0
            C_aug[i+c_zero+1] = C[i]
            plus_idx[i+c_zero] = true
            neg_idx[i+c_zero+1] = true
            c_zero += 1
        end
    end
    return C_aug, plus_idx, neg_idx
end
function _duplicate_zero_states_membership(C_aug,neg_idx,plus_idx)
    m_aug = Int.(sign.(C_aug))
    m_aug[plus_idx] .= +1
    m_aug[neg_idx] .= -1
    lpm_aug = (C_aug.<0.0) .| neg_idx
    rpm_aug = (C_aug.>0.0) .| plus_idx
    return m_aug, lpm_aug, rpm_aug
end
function _duplicate_zero_states_T(T,C,C_aug,neg_idx,plus_idx)
    # assign augmented generator
    c_zero = 0
    T_aug = zeros(length(C_aug),length(C_aug))
    for i in 1:length(C)
        if C[i] == 0
            # duplicate
            T_aug[i+c_zero,(C_aug.!=0).|(plus_idx)] = T[i,:]
            T_aug[i+c_zero+1,(C_aug.!=0).|(neg_idx)] = T[i,:]
            c_zero += 1
        elseif C[i] < 0
            T_aug[i+c_zero,(C_aug.!=0).|(neg_idx)] = T[i,:]
        elseif C[i] > 0
            T_aug[i+c_zero,(C_aug.!=0).|(plus_idx)] = T[i,:]
        end
    end
    return T_aug
end
function _duplicate_zero_states_P(neg_idx,plus_idx,C_aug,P_lwr,P_upr)
    P_lwr_aug = zeros(sum(C_aug.<0.0),length(neg_idx))
    P_upr_aug = zeros(sum(C_aug.>0.0),length(plus_idx))
    c_zero_lwr = 0
    c_zero_upr = 0
    for i in 1:length(neg_idx)
        if !((plus_idx[i])&&(C_aug[i]==0.0))
            c_zero_lwr += 1
            P_lwr_aug[:,i] = P_lwr[:,c_zero_lwr]
        end
        if !((neg_idx[i])&&(C_aug[i]==0.0))
            c_zero_upr += 1
            P_upr_aug[:,i] = P_upr[:,c_zero_upr]
        end
    end
    return P_lwr_aug, P_upr_aug
end

function _duplicate_zero_states(T::Array{Float64,2},C::Array{Float64,1},P_lwr::Array{Float64,2},P_upr::Array{Float64,2})
    C_aug, plus_idx, neg_idx = _duplicate_zero_states_C(T,C)
    m_aug, lpm_aug, rpm_aug = _duplicate_zero_states_membership(C_aug,neg_idx,plus_idx)
    T_aug = _duplicate_zero_states_T(T,C,C_aug,neg_idx,plus_idx)
    P_lwr_aug, P_upr_aug = _duplicate_zero_states_P(neg_idx,plus_idx,C_aug,P_lwr,P_upr)

    return T_aug, C_aug, m_aug, lpm_aug, rpm_aug, P_lwr_aug, P_upr_aug
end

"""
    augment_model(model::Model)

Given a BoundedFluidQueue, return a BoundedFluidQueue with twice as many phases with rate 0, one set associated
with m=1 phases and one associated with m=-1 phases.
"""
function augment_model(model::Model)
    if (any(rates(model).==0))
        T_aug, C_aug, m_aug, lpm_aug, rpm_aug, P_lwr_aug, P_upr_aug =
            _duplicate_zero_states(model.T,rates(model),model.P_lwr,model.P_upr)
        S_aug = PhaseSet(C_aug,m_aug,lpm_aug,rpm_aug)
        return BoundedFluidQueue(T_aug,S_aug,P_lwr_aug,P_upr_aug,model.b)
    else # no zero states, no augmentation needed
       return model
    end
end


text:
# This file is a part of RunStatistics.jl, licensed under the MIT License (MIT).

using RunStatistics
using Test

# Before you make a release, you should copy over all the tests I implemented in https://github.com/fredRos/runs/tree/master/test and make sure the numbers agree. There may some discrepancy because a different numerical integration is used but otherwise you should get the same numbers

@testset "squares" begin
    T_obs = 3.4
    N = 30

    @test squares_pvalue(T_obs, N) ≈ 0.818208032939994
    @test squares_cdf(T_obs, N) ≈ 0.18179196706000597
end


text:
# This file is a part of Julia. License is MIT: https://julialang.org/license

# RUN: julia --startup-file=no %s | opt -load libjulia-internal%shlibext -LateLowerGCFrame -FinalLowerGC -S - | FileCheck %s

println("""
declare {} addrspace(10)* @alloc()
declare void @one_arg_boxed({} addrspace(10)*)
declare {}*** @julia.ptls_states()

define void @stress(i64 %a, i64 %b) {
    %ptls = call {}*** @julia.ptls_states()
""")

# CHECK: %gcframe = alloca {} addrspace(10)*, i32 10002
for i = 1:10000
    println("\t%arg$i = call {} addrspace(10)* @alloc()")
end

for i = 1:10000
    println("\tcall void @one_arg_boxed({} addrspace(10)* %arg$i)")
end

println("""
    ret void
}
""")

max_stars_repo_name: udesou/julia

max_issues_repo_name: udesou/julia

max_forks_repo_name: udesou/julia


text:
# Test storage model using a multiperiod optimization

# - Model: `_FP.BFARadPowerModel` is used to be able to test storage in both active and
#   reactive power, while keeping the model linear. It requires a distribution network.
# - Problem: `flex_tnep` is the simplest problem that implements both storage and flexible
#   loads.

## Settings

file = normpath(@__DIR__,"..","test","data","case2","case2_d_strg.m") # Input case. Here 2-bus distribution network having 1 generator, 1 load, 1 storage device, and 1 candidate storage device, all on bus 1 (bus 2 is empty).
number_of_hours = 4 # Number of time periods

## Plot function

# Uncomment this part and the commented lines further down to display a nice plot when manually editing a testset
#=
using StatsPlots
function plot_storage(mn_data, result; candidate::Bool, id::Int=1)
    id = string(id)
    storage_type = candidate ? "ne_storage" : "storage"
    res_storage(key) = [result["solution"]["nw"]["$n"][storage_type][id][key] for n in 1:number_of_hours]
    data_storage(key) = [mn_data["nw"]["$n"][storage_type][id][key] for n in 1:number_of_hours]
    repeatfirst(x) = vcat(x[1:1,:], x)
    p = plot(0:number_of_hours, repeatfirst(res_storage(candidate ? "ps_ne" : "ps"));
        yguide = "Power [p.u.]",
        xformatter = _ -> "",
        legend = :none,
        framestyle = :zerolines,
        seriestype = :steppre,
        fillrange = 0,
        linewidth = 2,
        seriescolor = HSLA(203,1,0.49,0.5),
    )
    plot!(p, 0:number_of_hours, hcat(repeatfirst(data_storage("charge_rating")), -repeatfirst(data_storage("discharge_rating")));
        seriestype=:steppre, linewidth=2, linestyle=:dot, seriescolor=HSLA(203,1,0.49,1)
    )
    e = plot(0:number_of_hours, vcat(mn_data["nw"]["1"][storage_type][id]["energy"]*get(result["solution"]["nw"]["1"][storage_type][id],"isbuilt",1), res_storage(candidate ? "se_ne" : "se"));
        yguide = "Energy [p.u.]",
        xguide = "Time [h]",
        legend = :none,
        framestyle = :zerolines,
        fillrange = 0,
        linewidth = 2,
        seriescolor = HSLA(15,0.73,0.58,0.5),
    )
    plot!(e, 0:number_of_hours, repeatfirst(data_storage("energy_rating"));
        seriestype=:steppre, linewidth=2, linestyle=:dot, seriescolor=HSLA(15,0.73,0.58,1)
    )
    plt = plot(p, e; layout = (2,1), plot_title = (candidate ? "Candidate storage" : "Storage") * " $id" )
    display(plt)
end
=#

## Test results

@testset "Storage model" begin

    @testset "Common features" begin

        # Case with a storage device and a candidate storage device. As demand exceeds by
        # far the available generation in the second half of the time horizon, both storage
        # devices are charged in the first half of the time horizon and discharged in the
        # second half. Involuntary curtailment of the load covers the remaining excess of
        # demand.
        data = _FP.parse_file(file)
        _FP.add_dimension!(data, :hour, number_of_hours)
        _FP.add_dimension!(data, :year, 1; metadata = Dict{String,Any}("scale_factor"=>1))
        _FP.scale_data!(data; cost_scale_factor=1e-6)
        loadprofile = collect(reshape(range(0,2;length=number_of_hours),:,1)) # Create a load profile: ramp from 0 to 2 times the rated value of load
        time_series = _FP.make_time_series(data; loadprofile) # Compute time series by multiplying the rated value by the profile
        mn_data = _FP.make_multinetwork(data, time_series)
        result = _FP.flex_tnep(mn_data, _FP.BFARadPowerModel, cbc)

        @testset "Existing storage" begin
            #plot_storage(mn_data, result; candidate=false)
            @test result["solution"]["nw"]["1"]["storage"]["1"]["sc"]               ≈  1.0    rtol=1e-3
            @test result["solution"]["nw"]["2"]["storage"]["1"]["sc"]               ≈  1.0    rtol=1e-3
            @test result["solution"]["nw"]["3"]["storage"]["1"]["sc"]               ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["4"]["storage"]["1"]["sc"]               ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["1"]["storage"]["1"]["sd"]               ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["2"]["storage"]["1"]["sd"]               ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["3"]["storage"]["1"]["sd"]               ≈  1.0    rtol=1e-3
            @test result["solution"]["nw"]["4"]["storage"]["1"]["sd"]               ≈  0.6098 rtol=1e-3 # Less than 1.0 because of "charge_efficiency", "discharge_efficiency" and "self_discharge_rate"
            @test result["solution"]["nw"]["1"]["storage"]["1"]["ps"]               ≈  1.0    rtol=1e-3 # Storage model uses load convention: positive power when absorbed from the grid
            @test result["solution"]["nw"]["2"]["storage"]["1"]["ps"]               ≈  1.0    rtol=1e-3
            @test result["solution"]["nw"]["3"]["storage"]["1"]["ps"]               ≈ -1.0    rtol=1e-3 # Storage model uses load convention: negative power when injected into the grid
            @test result["solution"]["nw"]["4"]["storage"]["1"]["ps"]               ≈ -0.6098 rtol=1e-3
            @test result["solution"]["nw"]["1"]["storage"]["1"]["se"]               ≈  2.898  rtol=1e-3 # Greater than 2.0 because refers to the end of period
            @test result["solution"]["nw"]["2"]["storage"]["1"]["se"]               ≈  3.795  rtol=1e-3
            @test result["solution"]["nw"]["3"]["storage"]["1"]["se"]               ≈  2.680  rtol=1e-3
            @test result["solution"]["nw"]["4"]["storage"]["1"]["se"]               ≈  2.0    rtol=1e-3 # Must match "energy" parameter in data model
            @test haskey(result["solution"]["nw"]["4"]["storage"]["1"], "e_abs")    == false # Absorbed energy is not computed if is not bounded
        end

        @testset "Candidate storage" begin
            #plot_storage(mn_data, result; candidate=true)
            @test result["solution"]["nw"]["1"]["ne_storage"]["1"]["investment"]    ≈  1.0    atol=1e-3 # Invested in candidate storage device because it costs less than load curtailment
            @test result["solution"]["nw"]["1"]["ne_storage"]["1"]["isbuilt"]       ≈  1.0    atol=1e-3 # Candidate storage device is built accordingly to investment decision
            @test result["solution"]["nw"]["1"]["ne_storage"]["1"]["sc_ne"]         ≈  1.0    rtol=1e-3
            @test result["solution"]["nw"]["2"]["ne_storage"]["1"]["sc_ne"]         ≈  1.0    rtol=1e-3
            @test result["solution"]["nw"]["3"]["ne_storage"]["1"]["sc_ne"]         ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["4"]["ne_storage"]["1"]["sc_ne"]         ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["1"]["ne_storage"]["1"]["sd_ne"]         ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["2"]["ne_storage"]["1"]["sd_ne"]         ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["3"]["ne_storage"]["1"]["sd_ne"]         ≈  1.0    rtol=1e-3
            @test result["solution"]["nw"]["4"]["ne_storage"]["1"]["sd_ne"]         ≈  0.6098 rtol=1e-3 # Less than 1.0 because of "charge_efficiency", "discharge_efficiency" and "self_discharge_rate"
            @test result["solution"]["nw"]["1"]["ne_storage"]["1"]["ps_ne"]         ≈  1.0    rtol=1e-3 # Storage model uses load convention: positive power when absorbed from the grid
            @test result["solution"]["nw"]["2"]["ne_storage"]["1"]["ps_ne"]         ≈  1.0    rtol=1e-3
            @test result["solution"]["nw"]["3"]["ne_storage"]["1"]["ps_ne"]         ≈ -1.0    rtol=1e-3 # Storage model uses load convention: negative power when injected into the grid
            @test result["solution"]["nw"]["4"]["ne_storage"]["1"]["ps_ne"]         ≈ -0.6098 rtol=1e-3
            @test result["solution"]["nw"]["1"]["ne_storage"]["1"]["se_ne"]         ≈  2.898  rtol=1e-3 # Greater than 2.0 because refers to the end of period
            @test result["solution"]["nw"]["2"]["ne_storage"]["1"]["se_ne"]         ≈  3.795  rtol=1e-3
            @test result["solution"]["nw"]["3"]["ne_storage"]["1"]["se_ne"]         ≈  2.680  rtol=1e-3
            @test result["solution"]["nw"]["4"]["ne_storage"]["1"]["se_ne"]         ≈  2.0    rtol=1e-3 # Must match "energy" parameter in data model
            @test haskey(result["solution"]["nw"]["4"]["ne_storage"]["1"], "e_abs") == false # Absorbed energy is not computed if is not bounded
        end

    end

    @testset "Bounded absorption" begin

        # Same as base case, but the two storage devices have bounded absorption.
        data = _FP.parse_file(file)
        data["storage"]["1"]["max_energy_absorption"] = 1.0 # Limit the maximum energy absorption of existing storage device
        data["ne_storage"]["1"]["max_energy_absorption"] = 1.0 # Limit the maximum energy absorption of candidate storage device
        _FP.add_dimension!(data, :hour, number_of_hours)
        _FP.add_dimension!(data, :year, 1; metadata = Dict{String,Any}("scale_factor"=>1))
        _FP.scale_data!(data; cost_scale_factor=1e-6)
        loadprofile = collect(reshape(range(0,2;length=number_of_hours),:,1)) # Create a load profile: ramp from 0 to 2 times the rated value of load
        time_series = _FP.make_time_series(data; loadprofile) # Compute time series by multiplying the rated value by the profile
        mn_data = _FP.make_multinetwork(data, time_series)
        result = _FP.flex_tnep(mn_data, _FP.BFARadPowerModel, cbc)

        @testset "Existing storage" begin
            #plot_storage(mn_data, result; candidate=false)
            @test result["solution"]["nw"]["1"]["storage"]["1"]["ps"]               ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["2"]["storage"]["1"]["ps"]               ≈  1.0    rtol=1e-3
            @test result["solution"]["nw"]["3"]["storage"]["1"]["ps"]               ≈ -0.8020 rtol=1e-3
            @test result["solution"]["nw"]["4"]["storage"]["1"]["ps"]               ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["4"]["storage"]["1"]["e_abs"]            ≈  1.0    rtol=1e-3 # Must match "max_energy_absorption" parameter
        end

        @testset "Candidate storage" begin
            #plot_storage(mn_data, result; candidate=true)
            @test result["solution"]["nw"]["1"]["ne_storage"]["1"]["ps_ne"]         ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["2"]["ne_storage"]["1"]["ps_ne"]         ≈  1.0    rtol=1e-3
            @test result["solution"]["nw"]["3"]["ne_storage"]["1"]["ps_ne"]         ≈ -0.8020 rtol=1e-3
            @test result["solution"]["nw"]["4"]["ne_storage"]["1"]["ps_ne"]         ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["4"]["ne_storage"]["1"]["e_abs_ne"]      ≈  1.0    rtol=1e-3 # Must match "max_energy_absorption" parameter
        end

    end

    @testset "Candidate storage only" begin

        # Case with a storage device and a candidate storage device. The high demand in
        # period 4 requires using existing storage at full power and some load curtailment.
        # The candidate storage device is not built even though it would avoid load
        # curtailment because its construction costs more than load curtailment.
        data = _FP.parse_file(file)
        _FP.add_dimension!(data, :hour, number_of_hours)
        _FP.add_dimension!(data, :year, 1; metadata = Dict{String,Any}("scale_factor"=>1))
        _FP.scale_data!(data; cost_scale_factor=1e-6)
        loadprofile = collect(reshape(range(0,1.1005;length=number_of_hours),:,1)) # Create a load profile: ramp from 0 to 1.1005 times the rated value of load
        time_series = _FP.make_time_series(data; loadprofile) # Compute time series by multiplying the rated value by the profile
        mn_data = _FP.make_multinetwork(data, time_series)
        result = _FP.flex_tnep(mn_data, _FP.BFARadPowerModel, cbc)

        @testset "Not built if not needed" begin
            #plot_storage(mn_data, result; candidate=true)
            @test result["solution"]["nw"]["1"]["ne_storage"]["1"]["investment"]    ≈  0.0    atol=1e-3 # Not invested in candidate storage device because it costs more than the small amount of load curtailment needed to satisfy all power bounds in the last period
            @test result["solution"]["nw"]["1"]["ne_storage"]["1"]["isbuilt"]       ≈  0.0    atol=1e-3 # Candidate storage device is not built, accordingly to investment decision
            @test result["solution"]["nw"]["1"]["ne_storage"]["1"]["sc_ne"]         ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["2"]["ne_storage"]["1"]["sc_ne"]         ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["3"]["ne_storage"]["1"]["sc_ne"]         ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["4"]["ne_storage"]["1"]["sc_ne"]         ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["1"]["ne_storage"]["1"]["sd_ne"]         ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["2"]["ne_storage"]["1"]["sd_ne"]         ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["3"]["ne_storage"]["1"]["sd_ne"]         ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["4"]["ne_storage"]["1"]["sd_ne"]         ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["1"]["ne_storage"]["1"]["ps_ne"]         ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["2"]["ne_storage"]["1"]["ps_ne"]         ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["3"]["ne_storage"]["1"]["ps_ne"]         ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["4"]["ne_storage"]["1"]["ps_ne"]         ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["1"]["ne_storage"]["1"]["se_ne"]         ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["2"]["ne_storage"]["1"]["se_ne"]         ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["3"]["ne_storage"]["1"]["se_ne"]         ≈  0.0    atol=1e-3
            @test result["solution"]["nw"]["4"]["ne_storage"]["1"]["se_ne"]         ≈  0.0    atol=1e-3 # Even if "energy" parameter in data model is positive, this must be zero because the candidate storage device is not built
        end

    end

end;


text:
module psf

using FFTW
using ..params
import SpecialFunctions.besselj
import FastGaussQuadrature.gausslegendre
import LinearAlgebra.Symmetric

export integratePSF, fresnelH, originPSFproj

# NOTE: Number of quadrature nodes can be tuned for speed or accuracy
const (nodes, weights) = gausslegendre(100)

function intervalchange(x::Float64,a::Float64,b::Float64)
    newx = x*(b-a)/2 + (a+b)/2
    return newx
end

function lfmPSF(θ::Float64, α::Float64, v::Float64, u::Float64, a₀::Float64)
    if a₀ > zero(Float64)
        # PSF as given by Li et al. 2018. Used when MLA is offset from NIP by distance a₀
        I = sqrt(cos(θ)) * besselj(0, sin(θ)/sin(α)*v) * (1+cos(θ)) * exp((im*u*sin(θ/2)^2)/(2*sin(α/2)^2)) * sin(θ)
    else
        # PSF for classic LFM
        I = sqrt(cos(θ)) * besselj(0, sin(θ)/sin(α)*v) * (1+cos(θ)) * exp((-im*u*sin(θ/2)^2)/(2*sin(α/2)^2)) * sin(θ)
    end
    return I
end

function integratePSF(v::Float64, u::Float64, a₀::Float64, α::Float64)
    θ = intervalchange.(nodes, 0.0, α)
    integral = (α/2) * sum(lfmPSF.(θ, α, v, u, a₀) .* weights)
    return integral
end

function fresnelH(f0::Array{Complex{Float64},N}, par::ParameterSet, z::Float64) where N

    Nx = size(f0,1)
    f0length = par.sim.subpixelpitch*Nx

    # Generates spatial frequency range at 256-bit precision to avoid inexact error
    truestep = BigFloat(one(BigFloat)/f0length)
    endpt = BigFloat(one(BigFloat)/(2*par.sim.subpixelpitch))
    spfreq = range(-endpt, stop=endpt-truestep, length=Nx)

    # Setup frequency axes
    fx = spfreq'
    fy = reverse(spfreq, dims=1)
    if N == 2
        fxfy = fx.^2 .+ fy.^2
    elseif N == 1
        fxfy = fy.^2
    else
        throw("something wrong")
    end

    # Fresnel propagation func. (see: Computational Fourier Optics p.54-55,63)
    H::Array{ComplexF64,N} = exp(im*par.con.k0*z).*exp.((-im*pi*par.opt.lambda*z) .* fxfy)
    # Shift H for later fourier space calc + convert back to 64-bit precision
    return fftshift(H)
end

struct OriginPSF
    patternstack::Array{Array{Complex{Float64},2},1}
    originimgs::Array{Complex{Float64},3}

    function OriginPSF(lf::LightFieldSimulation)
        zmax = maximum(lf.obj.z)
        patternstack = Array{Array{Complex{Float64},2},1}(undef, lf.obj.zlen)
        originimgs = Threads.@spawn zeros(ComplexF64, lf.img.xlen, lf.img.ylen, lf.obj.zlen)

        Threads.@threads for p in 1:lf.obj.zlen
            imgrefsize = cld((lf.img.xlen * abs(lf.obj.z[p])), zmax)
            halfwidth = max(imgrefsize * lf.par.sim.subvpix, 2 * lf.par.sim.subvpix)
            centerarea = max((lf.img.center - halfwidth + 1), 1):1:min((lf.img.center + halfwidth - 1), lf.img.xlen)
            patternstack[p] = Array{Complex{Float64},2}(undef, length(centerarea), length(centerarea))
        end

        new(patternstack, fetch(originimgs))
    end
end

function originPSFproj(lf::LightFieldSimulation)

    psf = OriginPSF(lf)
    originPSFproj(psf, lf)

    #This is what happens when Shu Jia lab super resolution is applied
    if lf.par.opt.a0 > 0.0
        steps::Int64 = 10
        stepz = lf.par.opt.a0/steps
        Ha0 = fresnelH(psf.originimgs[:,:,1], lf.par, stepz)
        itrfresnelconv!(psf.originimgs, Ha0, steps, lf.obj)
    end

    return psf.originimgs
end

function unfold(integral::Array{Complex{Float64},2}, pattern::Array{Complex{Float64},2})
    middle::Int64 = cld(size(pattern,1),2)
    pattern[1:middle, 1:middle] .= integral[:,:]
    pattern[1:middle, middle:end] .= reverse(integral, dims=2)
    pattern[middle:end, 1:end] .= reverse(pattern[1:middle, 1:end], dims=1)
    return pattern
end

function originPSFproj(psf::OriginPSF, lf::LightFieldSimulation)

    img = lf.img
    obj = lf.obj
    con = lf.par.con
    sim = lf.par.sim
    opt = lf.par.opt

    zmax = maximum(obj.z)
    vscalar = con.k * sin(con.alpha)
    uscalar = 4 * con.k * (sin(con.alpha / 2)^2)

    Threads.@threads for j in 1:obj.zlen

        sizeref = cld((img.xlen * abs(obj.z[j])),zmax)
        halfwidth::Int64 =  max(sizeref * sim.subvpix, 2 * sim.subvpix)
        centerarea::Array{Int64,1} = max((img.center - halfwidth + 1), 1):1:min((img.center + halfwidth - 1), img.xlen)
        triangle = falses(length(centerarea[1]:img.center), length(centerarea[1]:img.center))

        for x in 1:size(triangle,1)
            triangle[1:x, x] .= true
        end

        xl2 = ((img.x[centerarea[1]:img.center]'.^2  .+ img.y[centerarea[1]:img.center].^2) .^0.5) ./ opt.M
        v = xl2 .* vscalar
        u = obj.z[j] * uscalar
        ki = opt.M / ((opt.fobj* opt.lambda)^2) * exp(-im * u / (4 * (sin(con.alpha / 2)^2)))

        integral = zeros(ComplexF64,size(v))
        integral[triangle] .= integratePSF.(v[triangle], u, opt.a0, con.alpha)
        integral[triangle] .= integral[triangle] .* ki
        copyto!(integral, Symmetric(integral))

        psf.patternstack[j] .= unfold(integral, psf.patternstack[j])
        psf.originimgs[centerarea, centerarea, j] .= psf.patternstack[j][:,:]

    end
    return
end

function itrfresnelconv!(originimgs::Array{Complex{Float64},3}, Ha0::Array{Complex{Float64},2}, steps::Int64, obj::Space)
    f0 = originimgs[:,:,1]
    p = plan_fft!(f0, [1,2])

    for h in 1:length(obj.z)
        f0 .= originimgs[:,:,h]
        p*f0                   # Applies fft in place
        f0 .= f0.*(Ha0.^steps) # Multiply by transfer function each multiplication is an incremental proj
        p\f0                   # Applies ifft in place
        originimgs[:,:,h] .= f0
    end
    return
end

end # psf module end


text:
# GPU memory allocation is very expensive.  So we create an
# application specific memory manager.  Typically same type, size, and
# number of arrays are needed in every iteration of training and
# testing.  During training these arrays should not be overwritten
# until the backward pass.  During testing they should be recycled as
# much as possible.

# We will have a Dict(length=>arrays) for each array type.

using CUDArt

type TmpStack; arr::Vector; idx::Int; TmpStack()=new([],0); end
typealias TmpDict Dict{Int,TmpStack}
for F in (Float16,Float32,Float64)
    A = CudaArray{F}
    let d=TmpDict()
        global tmpdict{T<:A}(::Type{T})=d
    end
end

tmpdict(a)=tmpdict(typeof(a))
tmpdict(T::Type)=error("TmpDict for $T not defined")

tmpfree()=(for T in (Float16,Float32,Float64), s in values(tmpdict(CudaArray{T})); s.idx=0; end)

function tmplike(a, dims::Dims=size(a))
    s = get!(TmpStack, tmpdict(a), prod(dims))
    s.idx += 1
    if s.idx > length(s.arr)
        push!(s.arr, similar(a,dims))
        s.idx > length(s.arr) && error("short stack")
    end
    if size(s.arr[s.idx]) != dims
        s.arr[s.idx] = reshape(s.arr[s.idx], dims)
    end
    return s.arr[s.idx]
end

function tmpmem()
    a = []
    for T in (Float16,Float32,Float64)
        D = tmpdict(CudaArray{T})
        isempty(D) && continue
        push!(a, (T,:sizes,length(D),:arrays,sum(s->length(s.arr),values(D)),:nelem,sum(d->d[1]*length(d[2].arr),D)))
    end
    return a
end