text:
using TravelingSalesmanHeuristics
using Test
using Random
using LinearAlgebra

###
# helpers
###

# generate a Euclidean distance matrix for n points in the unit square
function generate_planar_distmat(n)
	pts = rand(2, n)
	dm = [norm(pts[:,i] - pts[:,j]) for i in 1:n, j in 1:n]
end

# test that a path is acceptable:
# - at most one city appears twice, and if so must be first and last
# - all values 1:n are present where n is the maximum city index, and no other values
function testpathvalidity(path, iscycle)
	if iscycle
		@test path[1] == path[end]
	end
	n = iscycle ? length(path) - 1 : length(path)
	@test sort(unique(path)) == collect(1:n)
end

# tests that a path-cost pair is consistent with the distance matrix
# basically a reimplementation of `pathcost`, but hey...less likely to double a mistake?
function testpathcost(path, cost, dm)
	c = zero(eltype(dm))
	for i in 1:(length(path) - 1)
		c += dm[path[i], path[i+1]]
	end
	@test isapprox(c, cost)
end

###
# main tests
###

function test_nearest_neighbor()
	distmats = [generate_planar_distmat(10), generate_planar_distmat(2)]

	for dm in distmats
		n = size(dm, 1)
		randstartcity = rand(1:n)
		# standard
		path, cost = nearest_neighbor(dm)
		@test cost > 0
		testpathvalidity(path, true)
		# repetitive
		path, cost = repetitive_heuristic(dm, nearest_neighbor)
		@test cost > 0
		testpathvalidity(path, true)
		# no loop, 2 opt
		path, cost = nearest_neighbor(dm, closepath = false)
		@test cost > 0
		testpathvalidity(path, false)
		# no loop, no 2 opt
		path, cost = nearest_neighbor(dm, closepath = false, do2opt = false)
		@test cost > 0
		testpathvalidity(path, false)
		# fixed start, no loop
		path, cost = nearest_neighbor(dm, closepath = false, firstcity = randstartcity)
		@test cost > 0
		testpathvalidity(path, false)
	end

	dm_bad = rand(3,2)
	@test_throws ErrorException nearest_neighbor(dm_bad)
end

function test_cheapest_insertion()
	# default random start
	dm = generate_planar_distmat(8)
	path, cost = cheapest_insertion(dm)
	@test cost > 0
	testpathvalidity(path, true) # should be a closed path
	# repetitive start
	path, cost = repetitive_heuristic(dm, cheapest_insertion)
	@test cost > 0
	testpathvalidity(path, true)

	# bad input
	dmbad = rand(2,3)
	@test_throws ErrorException cheapest_insertion(dmbad)
end

function test_farthest_insertion()
	# standard symmetric case
	dm = generate_planar_distmat(8)
	path, cost = farthest_insertion(dm)
	testpathvalidity(path, true)
	testpathcost(path, cost, dm)
	# invalid argument
	@test_throws ErrorException farthest_insertion(dm; firstcity = 0)
	# asymmetric matrix
	dm = rand(20, 20)
	path, cost = farthest_insertion(dm; firstcity = 1, do2opt = true)
	testpathvalidity(path, true)
	testpathcost(path, cost, dm)
end

function test_simulated_annealing()
	dm = generate_planar_distmat(14)
	# single start
	path, cost = simulated_annealing(dm)
	@test cost > 0
	testpathvalidity(path, true) # closed path
	# multi-start
	dm = generate_planar_distmat(8)
	path, cost = simulated_annealing(dm, num_starts = 3)
	@test cost > 0
	testpathvalidity(path, true) # also closed
	# given init path
	init_path = collect(1:8)
	push!(init_path, 1)
	reverse!(init_path, 2, 6)
	path, cost = simulated_annealing(dm; init_path = init_path)
	@test cost > lowerbound(dm)
	testpathvalidity(path, true) # still closed
end

function test_bounds()
	dm = generate_planar_distmat(2)
	path, cost = solve_tsp(dm)
	lb = lowerbound(dm)
	@test lb <= cost
	dm = generate_planar_distmat(30)
	path, cost = solve_tsp(dm)
	lb = lowerbound(dm)
	@test lb <= cost
end

function test_path_costs()
	dm = [1 2 3; 4 5 6; 7 8 9]
	p1 = [1, 2, 3, 1]
	@test TravelingSalesmanHeuristics.pathcost(dm, p1) == 2 + 6 + 7
	# various reverse indices to test
	revs = [(1,4), (2,3), (2,4), (3,3)]
	for rev in revs
		path_rev = reverse(p1, rev[1], rev[2])
		@test TravelingSalesmanHeuristics.pathcost(dm, path_rev) ==
			  TravelingSalesmanHeuristics.pathcost_rev(dm, p1, rev[1], rev[2])
	end
end

function test_solve_tsp()
	dm = rand(10,10)
	quality_factors = [-1, 0, 1.1, 35, 101]
	for qf in quality_factors
		path, cost = solve_tsp(dm; quality_factor = qf)
		testpathvalidity(path, true)
	end
end

function test_two_opt()
	n = 50
	dm = generate_planar_distmat(n)
	path = collect(1:n)
	path_2opt, _ = two_opt(dm, path)
	testpathvalidity(path_2opt, false)
	# shouldn't affect endpoints
	@test path_2opt[1] == 1
	@test path_2opt[end] == n
	cycle = vcat(path, 1)
	cycle_2opt, _ = two_opt(dm, cycle)
	testpathvalidity(cycle_2opt, true)
	# same endpoints
	@test cycle_2opt[1] == 1 == cycle_2opt[n+1]
end

function test_atypical_types()
	# rational distance matrix
	dm = [rationalize(rand(); tol = .01) for i in 1:10, j in 1:10]
	initpath = 1:2:9 # not a Vector, but <:AbstractVector{<:Int}
	p, c = cheapest_insertion(dm, initpath)
	testpathvalidity(p, false)
end

###
# run
###
println("Hello tester")
Random.seed!(47)
test_nearest_neighbor()
test_cheapest_insertion()
test_farthest_insertion()
test_simulated_annealing()
test_bounds()
test_path_costs()
test_solve_tsp()
test_two_opt()
test_atypical_types()
println("Done testing.")


text:
struct LShapedSolver <: AbstractStructuredSolver
    variant::Symbol
    lpsolver::AbstractMathProgSolver
    subsolver::AbstractMathProgSolver
    crash::Crash.CrashMethod
    parameters

    function (::Type{LShapedSolver})(variant::Symbol, lpsolver::AbstractMathProgSolver; crash::Crash.CrashMethod = Crash.None(), subsolver = lpsolver, kwargs...)
        return new(variant,lpsolver,subsolver,crash,kwargs)
    end
end
LShapedSolver(lpsolver::AbstractMathProgSolver; kwargs...) = LShapedSolver(:ls, lpsolver, kwargs...)

function StructuredModel(solver::LShapedSolver,stochasticprogram::JuMP.Model)
    x₀ = solver.crash(stochasticprogram,solver.lpsolver)
    if solver.variant == :ls
        return LShaped(stochasticprogram,x₀,solver.lpsolver,solver.subsolver; solver.parameters...)
    elseif solver.variant == :dls
        return DLShaped(stochasticprogram,x₀,solver.lpsolver,solver.subsolver; solver.parameters...)
    elseif solver.variant == :rd
        return Regularized(stochasticprogram,x₀,solver.lpsolver,solver.subsolver; solver.parameters...)
    elseif solver.variant == :drd
        return DRegularized(stochasticprogram,x₀,solver.lpsolver,solver.subsolver; solver.parameters...)
    elseif solver.variant == :tr
        return TrustRegion(stochasticprogram,x₀,solver.lpsolver,solver.subsolver; solver.parameters...)
    elseif solver.variant == :dtr
        return DTrustRegion(stochasticprogram,x₀,solver.lpsolver,solver.subsolver; solver.parameters...)
    elseif solver.variant == :lv
        return LevelSet(stochasticprogram,x₀,solver.lpsolver,solver.subsolver; solver.parameters...)
    elseif solver.variant == :dlv
        return DLevelSet(stochasticprogram,x₀,solver.lpsolver,solver.subsolver; solver.parameters...)
    else
        error("Unknown L-Shaped variant: ", solver.variant)
    end
end

function optimsolver(solver::LShapedSolver)
    return solver.lpsolver
end

function optimize_structured!(lshaped::AbstractLShapedSolver)
    return lshaped()
end

function fill_solution!(lshaped::AbstractLShapedSolver,stochasticprogram::JuMP.Model)
    # First stage
    nrows, ncols = length(stochasticprogram.linconstr), stochasticprogram.numCols
    stochasticprogram.objVal = lshaped.solverdata.Q
    stochasticprogram.colVal = copy(lshaped.x)
    stochasticprogram.redCosts = try
        getreducedcosts(lshaped.mastersolver.lqmodel)[1:ncols]
    catch
        fill(NaN, ncols)
    end
    stochasticprogram.linconstrDuals = try
        getconstrduals(lshaped.mastersolver.lqmodel)[1:nrows]
    catch
        fill(NaN, nrows)
    end

    # Second stage
    fill_submodels!(lshaped,scenarioproblems(stochasticprogram))
end

ext: jl


text:
# Vicenty's inverse (two points -> ellipsoidal distance and azimuths)
const ANTIPODAL_ERROR =
    ArgumentError("Points too near antipodal for distance algorithm to converge")

function vicentys_inverse{T <: LL}(a::T, b::T)
    d = ellipsoid(T)
    f = 1 - d.b/d.a

    # Reduced latitude (latitude on the auxiliary sphere)
    U₁ = atan(d.b / d.a * tand(a.lat))
    U₂ = atan(d.b / d.a * tand(b.lat))

    sinU₁, cosU₁ = sin(U₁), cos(U₁)
    sinU₂, cosU₂ = sin(U₂), cos(U₂)

    cosU₁cosU₂ = cosU₁ * cosU₂
    sinU₁sinU₂ = sinU₁ * sinU₂
    sinU₁cosU₂ = sinU₁ * cosU₂
    cosU₁sinU₂ = cosU₁ * sinU₂

    ΔL = deg2rad(b.lon - a.lon)

    # Needed in scope after the loop
    oldλ = sinλ = cosλ = sinσ = cosσ = σ = cos²α = cos2σm = cos²2σm = Inf
    itr = 0

    # Initialize λ to ΔL for first iteration
    λ = ΔL

    for itr=1:200
        oldλ = λ

        sinλ, cosλ = sin(λ), cos(λ)

        x = cosU₂ * sinλ
        y = cosU₁sinU₂ - sinU₁cosU₂ * cosλ
        sinσ = sqrt(x*x + y*y)
        cosσ = sinU₁sinU₂ + cosU₁cosU₂ * cosλ
        σ = atan2(sinσ, cosσ)

        sinα = cosU₁cosU₂ * sinλ / sinσ
        cos²α = 1 - sinα*sinα
        cos2σm = cosσ - 2sinU₁sinU₂ / cos²α
        cos²2σm = cos2σm * cos2σm
        C = f / 16 * cos²α * (4 + f*(4-3cos²α))

        λ = ΔL + (1-C) * f * sinα *
            (σ + C * sinσ * (cos2σm + C * cosσ * (-1 + 2cos²2σm)))

        abs(λ - oldλ) < 1e-12 && break
    end

    # Failure to converge happens when abs(λ) > π after the first iteration
    itr == 200 && throw(ANTIPODAL_ERROR)

    a², b² = d.a * d.a, d.b * d.b
    u² = cos²α * (a² - b²) / b²
    A = 1 + u² / 16384 * (4096 + u² * (-768 + u² * (320 - 175u²)))
    B = u² / 1024 * (256 + u² * (-128 + u² * (74 - 47u²)))
    Δσ = B * sinσ * (cos2σm + B / 4 *
        (cosσ * (-1 + 2cos²2σm) -
         B / 6 * cos2σm * (-3 + 4sinσ*sinσ) * (-3 + 4cos²2σm)))

    s = d.b * A * (σ - Δσ) # ellipsoidal distance
    α₁ = atan2(cosU₂ * sinλ,  cosU₁sinU₂ - sinU₁cosU₂ * cosλ) # start azimuth
    α₂ = atan2(cosU₁ * sinλ, -sinU₁cosU₂ + cosU₁sinU₂ * cosλ) # end azimuth

    return s, rad2deg(α₁), rad2deg(α₂)
end


text:
paths_file = "$(Pkg.dir("BSR"))/src/paths.jl"

if !isfile(paths_file)
    bsr = try
        ENV["BSR"]
    catch
        error("Did not find BSR environment variables. Please define it and rebuild BSR.")
    end

    isfile("$bsr/bsr_prep") || error("Could not find bsr_prep exe in $bsr")

    println("Found BSR at $bsr")

    open(paths_file, "w") do file
        write(file, """bsr = \"$bsr\"
    """)
    end
end


text:
using WindMixing
using OceanParameterizations
using OrdinaryDiffEq
using DiffEqSensitivity
using Flux
using GalacticOptim
using LinearAlgebra
using Statistics
using Plots
using Oceananigans.Grids
using JLD2
using FileIO

BLAS.set_num_threads(1)

T_fraction = parse(Float32, ARGS[1])
N_sims = parse(Int, ARGS[2])
train_gradient = parse(Bool, ARGS[3])
optimizer_type = ARGS[4]
# optimizer_type = "ADAM"
# rate_str = ARGS[4]
# rate = parse(Float64, rate_str)

train_files_all = [
    "wind_-5e-4_cooling_3e-8_new",
    "wind_-5e-4_cooling_1e-8_new",
    "wind_-2e-4_cooling_3e-8_new",
    "wind_-2e-4_cooling_1e-8_new",
    "wind_-5e-4_heating_-3e-8_new",
    "wind_-2e-4_heating_-1e-8_new",
    "wind_-2e-4_heating_-3e-8_new",
    "wind_-5e-4_heating_-1e-8_new",

    "wind_-3.5e-4_cooling_2e-8_new",
    "wind_-3.5e-4_heating_-2e-8_new",

    "wind_-5e-4_cooling_2e-8_new",
    "wind_-3.5e-4_cooling_3e-8_new",
    "wind_-3.5e-4_cooling_1e-8_new",
    "wind_-2e-4_cooling_2e-8_new",
    "wind_-3.5e-4_heating_-3e-8_new",
    "wind_-3.5e-4_heating_-1e-8_new",
    "wind_-2e-4_heating_-2e-8_new",
    "wind_-5e-4_heating_-2e-8_new",
]

train_files = train_files_all[1:N_sims]

PATH = pwd()
# PATH = "D:\\University Matters\\MIT\\CLiMA Project\\OceanParameterizations.jl"

if train_gradient
    FILE_NAME = "parameter_optimisation_$(N_sims)sim_windcooling_windheating_5params_$(optimizer_type)_T$(T_fraction)_grad_new"
else
    FILE_NAME = "parameter_optimisation_$(N_sims)sim_windcooling_windheating_5params_$(optimizer_type)_T$(T_fraction)_nograd_new"
end

OUTPUT_PATH = joinpath(PATH, "training_output", "$(FILE_NAME).jld2")
@assert !isfile(OUTPUT_PATH)

EXTRACTED_OUTPUT_PATH = joinpath(PATH, "extracted_training_output", "$(FILE_NAME)_extracted.jld2")

timestepper = ROCK4()

if optimizer_type == "LBFGS"
    optimizers = [LBFGS()]
else
    optimizers = [BFGS()]
end

# optimizers = [ADAM(rate)]

tsteps = 1:20:1153
maxiters = 200

training_fractions = (T=T_fraction, profile=0.5f0, ∂T∂z=T_fraction)

optimise_modified_pacanowski_philander(train_files, tsteps, timestepper, optimizers, maxiters, OUTPUT_PATH, n_simulations=length(train_files),
                                       train_gradient=train_gradient, training_fractions=training_fractions)

extract_parameters_modified_pacanowski_philander_optimisation(OUTPUT_PATH, EXTRACTED_OUTPUT_PATH)

PATH = joinpath(pwd(), "extracted_training_output")

DATA_NAME = FILE_NAME

DATA_PATH = joinpath(PATH, "$(DATA_NAME)_extracted.jld2")
ispath(DATA_PATH)

FILE_PATH = joinpath(pwd(), "Output", DATA_NAME)

if !ispath(FILE_PATH)
    mkdir(FILE_PATH)
end

file = jldopen(DATA_PATH, "r")

train_files = file["training_info/train_files"]
train_parameters = file["training_info/parameters"]
loss_scalings = file["training_info/loss_scalings"]
mpp_parameters = file["parameters"]

𝒟train = WindMixing.data(train_files, scale_type=ZeroMeanUnitVarianceScaling, enforce_surface_fluxes=false)
close(file)

ν₀, ν₋, ΔRi, Riᶜ, Pr = mpp_parameters

N_inputs = 96
hidden_units = 400
N_outputs = 31

weights, re = Flux.destructure(Chain(Dense(N_inputs, hidden_units, relu), Dense(hidden_units, N_outputs)))

uw_NN = re(zeros(Float32, length(weights)))
vw_NN = re(zeros(Float32, length(weights)))
wT_NN = re(zeros(Float32, length(weights)))

to_run = [
    "wind_-5e-4_cooling_3e-8_new",
    "wind_-5e-4_cooling_1e-8_new",
    "wind_-2e-4_cooling_3e-8_new",
    "wind_-2e-4_cooling_1e-8_new",
    "wind_-5e-4_heating_-3e-8_new",
    "wind_-2e-4_heating_-1e-8_new",
    "wind_-2e-4_heating_-3e-8_new",
    "wind_-5e-4_heating_-1e-8_new",

    "wind_-3.5e-4_cooling_2e-8_new",
    "wind_-3.5e-4_heating_-2e-8_new",

    "wind_-5e-4_cooling_2e-8_new",
    "wind_-3.5e-4_cooling_3e-8_new",
    "wind_-3.5e-4_cooling_1e-8_new",
    "wind_-2e-4_cooling_2e-8_new",
    "wind_-3.5e-4_heating_-3e-8_new",
    "wind_-3.5e-4_heating_-1e-8_new",
    "wind_-2e-4_heating_-2e-8_new",
    "wind_-5e-4_heating_-2e-8_new",
  ]

for test_file in to_run
    @info "running $test_file"
    test_files = [test_file]
    𝒟test = WindMixing.data(test_files, scale_type=ZeroMeanUnitVarianceScaling, enforce_surface_fluxes=false)
    trange = 1:1:1153
    plot_data = NDE_profile(uw_NN, vw_NN, wT_NN, test_file, 𝒟test, 𝒟train, trange,
                            modified_pacanowski_philander=true,
                            ν₀=ν₀, ν₋=ν₋, ΔRi=ΔRi, Riᶜ=Riᶜ, Pr=Pr,
                            convective_adjustment=false,
                            smooth_NN=false, smooth_Ri=false,
                            zero_weights=true,
                            loss_scalings=loss_scalings)

    animation_type = "Pre-Training"
    n_trainings = length(train_files)
    training_types = "Modified Pacanowski-Philander"
    VIDEO_NAME = "$(test_file)_mpp"
    animate_profiles_fluxes_comparison(plot_data, plot_data, plot_data, joinpath(FILE_PATH, VIDEO_NAME), fps=30,
                                                    animation_type=animation_type, n_trainings=n_trainings, training_types=training_types)
end


text:
"""
    partition(d::Dict, set::Dict; kwargs...)

# Arguments
- `d::Dict` of DataFrames containing the model data.
- `set::Dict` of Arrays describing region, sector, final demand, etc.

# Keywords
- `save_build = true`
- `overwrite = false`
See [`SLiDE.build`](@ref) for keyword argument descriptions.

# Returns
- `d::Dict` of DataFrames containing the model data at the
"""
function partition(
    dataset::String,
    d::Dict,
    set::Dict;
    version::String=DEFAULT_VERSION,
    save_build::Bool=DEFAULT_SAVE_BUILD,
    overwrite::Bool=DEFAULT_OVERWRITE,
    map_fdcat::Bool=false,
)
    # !!!! different for "detailed" sector as not to overwrite.
    CURR_STEP = "partition"

    # If there is already partition data, read it and return.
    d_read = read_build(dataset, CURR_STEP; overwrite = overwrite)
    !(isempty(d_read)) && (return d_read)

    x = Deselect([:units],"==")
    [d[k] = edit_with(filter_with(d[k], (yr=set[:yr],)), x) for k in [:supply,:use]]

    map_fdcat && _filter_use!(d,set)

    _partition_io!(d, set)
    _partition_fd!(d, set)

    # _partition_fd0!(d, set)
    _partition_ts0!(d, set)
    _partition_va0!(d, set)
    _partition_x0!(d, set)

    _partition_cif0!(d, set)
    _partition_m0!(d, set)   # cif0
    _partition_trn0!(d, set) # cif0

    _partition_mrg0!(d, set)
    _partition_md0!(d, set)  # mrg0, trn0
    _partition_ms0!(d)       # mrg0, trn0

    # _partition_fs0!(d)       # fd0
    _partition_s0!(d)        # ys0

    _partition_y0!(d, set)   # ms0, fs0, ys0
    _partition_a0!(d, set)   # fd0, id0

    _partition_ta0!(d, set)  # a0, sbd0, tax0
    _partition_tm0!(d, set)  # duty0, m0

    write_build!(dataset, CURR_STEP, d; save_build = save_build)
    return d
end

"""
    _remove_imrg(df::DataFrame, x::Pair{Symbol,Array{String,1}})
Removes commodity taxes and subsidies on the goods which are produced solely
for supplying retail sales margin.
"""
function _remove_imrg(df::DataFrame, x::Pair{Symbol,Array{String,1}})
    df[findall(in(x.second), df[:,x.first]), :value] .= 0.0
    return df
end

"""
This function combines fd values into fdcat values upfront. Doing so changes the results of
the calibration routine slightly, but isn't too bad. This is the only mapping
still performed in the build stream. If the difference in output is alright, we can
move this mapping to the data stream.
"""
function _filter_use!(d::Dict, set::Dict)
    df = copy(d[:use])
    s0 = append(:s,:fd)

    x = [
        Rename(:s,s0),
        Map("crosswalk/fd.csv",[:fd],[:fdcat],[s0],[:s],:left),
        Replace(:s,missing,"$s0 value"),
        Combine("",propertynames(df)),
        Order(propertynames(df), eltype.(eachcol(df))),
    ]

    set[:fd] = unique(read_file(x[2])[:,:fdcat])
    d[:use] = edit_with(df,x)
end

"""
`ys0(yr,s,g)`, sectoral supply (with byproducts), and
`id0(yr,g,s)`, intermediate input demand

Filter from supply/use data:

```math
\\begin{aligned}
\\tilde{id}_{yr,s,} = \\left\\{{use}\\left(yr,i,j\\right)
\\;\\vert\\; yr,\\, g \\in i,\\, s \\in j \\right\\}
\\\\
\\tilde{ys}_{yr,s,g} = \\left\\{{supply}\\left(yr,j,i\\right)
\\;\\vert\\; yr,\\, s \\in j,\\, g \\in i \\right\\}
\\end{aligned}
```

Treat negative inputs as outputs:

```math
\\begin{aligned}
\\tilde{ys}_{yr,s,g} &= \\tilde{ys}_{yr,s,g} - \\min\\left\\{0, \\tilde{id}_{yr,g,s} \\right\\} \\\\
\\tilde{id}_{yr,g,s} &= \\max\\left\\{ 0, \\tilde{id}_{yr,s,g} \\right\\}
\\end{aligned}
```
"""
function _partition_io!(d::Dict, set::Dict)
    println("  Partitioning id0 and ys0, supply/demand data.")
    # (!!!!) filtering here assumes sector/good are the same.
    d[:id0] = filter_with(d[:use], set)
    d[:ys0] = filter_with(d[:supply], set)

    # In sectordisagg, the good/sector column names are switched...
    if d[:sector]==:detail
        x = Rename.([:g,:s,:g_temp],[:g_temp,:g,:s])
        d[:ys0] = edit_with(d[:ys0], x)
    end

    (d[:id0], d[:ys0]) = fill_zero(d[:id0], d[:ys0])

    # Treat negative inputs as outputs.
    d[:ys0][!,:value] = d[:ys0][:,:value] - min.(0, d[:id0][:,:value])
    d[:id0][!,:value] = max.(0, d[:id0][:,:value])

    [dropzero!(d[k]) for k in [:ys0,:id0]]
    return d
end

"""
`a(yr,g)`, Armington supply

```math
\\tilde{a}_{yr,g} = \\sum_{fd}\\tilde{fd}_{yr,g,fd} + \\sum_{s}\\tilde{id}_{yr,g,s}
```
"""
function _partition_a0!(d::Dict, set::Dict)
    println("  Partitioning a0, Armington supply")
    d[:a0] = combine_over(d[:fd0], :fd) + combine_over(d[:id0], :s)
    d[:a0] = _remove_imrg(d[:a0], :g => set[:imrg])
    return d[:a0]
end

"""
`bopdef(yr)`, balance of payments

```math
\\tilde{bop}_{yr} = 0
\\;\\forall\\; yr
```
"""
function _partition_bop!(d::Dict, set::Dict)
    println("  Partitioning bopdef0, balance of payments deficit")
    d[:bopdef] = fill_zero((yr = set[:yr], ))
    return d[:bopdef]
end

"""
`cif(yr,g)`, CIF/FOB Adjustments on Imports

```math
\\tilde{cif}_{yr,g} = \\left\\{{supply}\\left(yr,i,j\\right)
\\;\\vert\\; yr,\\, g \\in i,\\, j = ciffob \\right\\}
```
"""
function _partition_cif0!(d::Dict, set::Dict)
    println("  Partitioning CIF/FOB adjustments on imports")
    d[:cif0] = filter_with(d[:supply], (g = set[:g], s = "ciffob"); drop = true)
    return d[:cif0]
end

"""
`duty(yr,g)`, import duties

```math
\\tilde{duty}_{yr,g} = \\left\\{{supply}\\left(yr,i,j\\right)
\\;\\vert\\; yr,\\, g \\in i,\\, j = duties \\right\\}
```
"""
function _partition_duty0!(d::Dict, set::Dict)
    println("  Partitioning duty0, import duties")
    d[:duty0] = filter_with(d[:supply], (g = set[:g], s = "duties"); drop = true)
    d[:duty0] = _remove_imrg(d[:duty0], :g => set[:imrg])
    return d[:duty0]
end

"""
`fd(yr,g,fd)`, final demand, and
`fs(yr,g)`, household supply

```math
\\begin{aligned}
\\tilde{fd}_{yr,g,fd} &= \\left\\{{use}\\left(yr,i,j\\right) \\;\\vert\\; yr,\\, g \\in i,\\, fd \\in j \\right\\}
\\\\
\\tilde{fs}_{yr,g} &= \\left\\{\\tilde{fd}_{yr,g,fd} \\;\\vert\\; yr,\\, g \\in i,\\, fd = pce \\right\\}
\\end{aligned}
```

```math
\\begin{aligned}
\\tilde{fs}_{yr,g} &= - \\min\\left\\{0, \\tilde{fs}_{yr,g} \\right\\}
\\\\
\\tilde{fd}_{yr,g,fd} &= \\max\\left\\{0, \\tilde{fd}_{yr,g,fd} \\right\\} \\;\\vert\\; yr,\\, g,\\, fd = pce \\right\\}
\\end{aligned}
```
"""
function _partition_fd!(d::Dict, set::Dict)
    x = Rename(:s, :fd)
    d[:fd0] = filter_with(edit_with(d[:use],x), set)
    d[:fs0] = filter_with(d[:fd0], (fd=["pce","C"],); drop = true)

    d[:fs0][!,:value] .= - min.(0, d[:fs0][:,:value])

    # For the sectoral disaggregation,
    if d[:sector]==:detail
        d[:fd0][.&(d[:fd0][:,:fd].=="pce", d[:fd0][:,:value].<0),:value] .= 0.0
        d[:fd0][.&(d[:fd0][:,:fd].=="C", d[:fd0][:,:value].<0),:value] .= 0.0
    end

    [dropzero!(d[k]) for k in [:fd0,:fs0]]
    return d
end

"""
`fd(yr,g,fd)`, final demand

```math
\\tilde{fd}_{yr,g,fd} = \\left\\{{use}\\left(yr,i,j\\right)
\\;\\vert\\; yr,\\, g \\in i,\\, fd \\in j \\right\\}
\\\\
\\tilde{fd}_{yr,g,fd} = \\max\\left\\{0, \\tilde{fd}_{yr,g,fd} \\right\\} \\;\\vert\\; yr,\\, g,\\, fd = pce \\right\\}
```
"""
function _partition_fd0!(d::Dict, set::Dict)
    println("  Partitioning fd0, final demand")
    d[:fd0] = filter_with(d[:use], (g = set[:g],  s = set[:fd]))
    d[:fd0] = edit_with(d[:fd0], Rename(:s, :fd))
    return d[:fd0]
end

"""
`fd(yr,g,fd)`, final demand, and
`fs(yr,g)`, household supply

```math
\\tilde{fs}_{yr,g} = \\left\\{\\tilde{fd}_{yr,g,fd}
\\;\\vert\\; yr,\\, g \\in i,\\, fd = pce \\right\\}
\\\\
\\tilde{fs}_{yr,g} = - \\min\\left\\{0, \\tilde{fs}_{yr,g} \\right\\}
```
"""
function _partition_fs0!(d::Dict)
    println("  Partitioning fs0, household supply")
    d[:fs0] = filter_with(d[:fd0], (fd=["pce","C"],); drop=true)
    d[:fs0][!,:value] .= - min.(d[:fs0][:,:value], 0)
    return dropzero!(d[:fs0])
end

"""
`m(yr,g)`, imports

```math
\\tilde{m}_{yr,g} = \\left\\{{supply}\\left(yr,i,j\\right)
\\;\\vert\\; yr,\\, g \\in i,\\, j = imports \\right\\}
```

Adjust transport margins according to CIF/FOB adjustments:

```math
\\tilde{m}_{yr,g} = \\tilde{m}_{yr,g} + \\tilde{cif}_{yr,g}
\\;\\forall\\; g = ins
```
"""
function _partition_m0!(d::Dict, set::Dict)
    println("  Partitioning m0, imports")
    d[:m0]  = filter_with(d[:supply], (g=set[:g], s="imports"); drop = true)

    # Adjust transport margins for transport sectors according to CIF/FOB adjustments.
    # Insurance imports are specified as net of adjustments.
    if d[:sector]==:summary
        d[:m0] += filter_with(d[:cif0], (g = "ins",))
        d[:m0] = _remove_imrg(d[:m0], :g => set[:imrg])
    end
    return d[:m0]
end

"""
`md(yr,m,g)`, margin demand

```math
\\begin{aligned}
\\tilde{md}_{yr,m,g} &=
\\begin{cases}
\\tilde{mrg}_{yr,g}  & m = trd   \\\\
\\tilde{trn}_{yr,g}  & m = trn
\\end{cases}
\\\\
\\tilde{md}_{yr,m,g} &= \\max\\left\\{0, \\tilde{md}_{yr,m,g} \\right\\}
\\end{aligned}
```
"""
function _partition_md0!(d::Dict, set::Dict)
    println("  Partitioning md0, margin demand")
    d[:md0] = [edit_with(d[:mrg0], Add(:m, "trd")); edit_with(d[:trn0], Add(:m, "trn"))]
    d[:md0] = sort(d[:md0][:,[:yr,:m,:g,:value]])
    d[:md0] = _remove_imrg(d[:md0], :g => set[:imrg])

    d[:md0][!,:value] .= max.(d[:md0][:,:value], 0)
    return d[:md0]
end

"""
`ms(yr,g,m)`, margin supply

```math
\\begin{aligned}
\\tilde{ms}_{yr,g,m} &=
\\begin{cases}
\\tilde{mrg}_{yr,g}  & m = trd   \\\\
\\tilde{trn}_{yr,g}  & m = trn
\\end{cases}
\\\\
\\tilde{ms}_{yr,g,m} &= \\max\\left\\{0, -\\tilde{ms}_{yr,g,m} \\right\\}
\\end{aligned}
```
"""
function _partition_ms0!(d::Dict)
    println("  Partitioning ms0, margin supply")
    d[:ms0] = [edit_with(d[:mrg0], Add(:m, "trd")); edit_with(d[:trn0], Add(:m, "trn"))]
    d[:ms0] = sort(d[:ms0][:,[:yr,:g,:m,:value]])

    d[:ms0][!,:value] .= max.(-d[:ms0][:,:value], 0)
    return d[:ms0]
end

"""
`mrg(yr,g)`, trade margins

```math
\\tilde{mrg}_{yr,g} = \\left\\{{supply}\\left(yr,i,j\\right)
\\;\\vert\\; yr,\\, g \\in i,\\, j = margins \\right\\}
```
"""
function _partition_mrg0!(d::Dict, set::Dict)
    println("  Partitioning mrg0, trade margins")
    d[:mrg0] = filter_with(d[:supply], (g = set[:g], s = "margins"); drop = true)
    return d[:mrg0]
end

"""
`s(yr,s)`, aggregate supply

```math
\\tilde{s}_{yr,s} = \\sum_{g}\\tilde{ys}_{yr,s,g}
```
"""
function _partition_s0!(d::Dict)
    println("  Partitioning s0, aggregate supply")
    d[:s0] = edit_with(combine_over(d[:ys0], :g), Rename(:s,:g))
    return d[:s0]
end

"""
`sbd(yr,g)`, subsidies on products

```math
\\tilde{sbd}_{yr,g} = \\left\\{{supply}\\left(yr,i,j\\right)
\\;\\vert\\; yr,\\, g \\in i,\\, j = subsidies \\right\\}
```

Treat negative inputs as outputs:

```math
\\tilde{sbd}_{yr,g} = - \\tilde{sbd}_{yr,g}
```
"""
function _partition_sbd0!(d::Dict, set::Dict)
    println("  Partitioning sbd0, subsidies on products")
    d[:sbd0] = filter_with(d[:supply], (g = set[:g], s = "subsidies"); drop = true)
    d[:sbd0] = _remove_imrg(d[:sbd0], :g => set[:imrg])
    d[:sbd0][!,:value] *= -1
    return d[:sbd0]
end

"""
`ta(yr,g)`, import tariff

```math
\\tilde{ta}_{yr,g} = \\frac{\\tilde{tax}_{yr,g} - \\tilde{sbd}_{yr,g}}{\\tilde{a}_{yr,g}}
```
"""
function _partition_ta0!(d::Dict, set::Dict)
    println("  Partitioning ta0(yr,g), import tariffs")
    !(:tax0 in keys(d)) && _partition_tax0!(d, set)
    !(:sbd0 in keys(d)) && _partition_sbd0!(d, set)

    d[:ta0] = dropnan((d[:tax0] - d[:sbd0]) / d[:a0])
    # d[:ta0] = edit_with(d[:ta0], Replace(:units, "billions of us dollars (USD)", "USD/USD"))
    return d[:ta0]
end

"""
`tax(yr,g)`, taxes on products

```math
\\tilde{tax}_{yr,g} = \\left\\{{supply}\\left(yr,i,j\\right)
\\;\\vert\\; yr,\\, g \\in i,\\, j = tax \\right\\}
```
"""
function _partition_tax0!(d::Dict, set::Dict)
    println("  Partitioning tax0, taxes on products")
    d[:tax0] = filter_with(d[:supply], (g = set[:g], s = "tax"); drop = true)
    d[:tax0] = _remove_imrg(d[:tax0], :g => set[:imrg])
    return d[:tax0]
end

"""
`tm(yr,g)`, tax net subsidy rate on intermediate demand

```math
\\tilde{tm}_{yr,g} = \\frac{\\tilde{duty}_{yr,g}}{\\tilde{m}_{yr,g}}
```
"""
function _partition_tm0!(d::Dict, set::Dict)
    println("  Partitioning tm0, tax net subsidy rate on intermediate demand")
    !(:duty0 in keys(d)) && _partition_duty0!(d, set)
    d[:tm0] = dropnan(d[:duty0] / d[:m0])
    # d[:tm0] = edit_with(d[:tm0], Replace(:units, "billions of us dollars (USD)", "USD/USD"))
    return d[:tm0]
end

"""
`trn(yr,g)`, transportation costs

```math
\\tilde{trn}_{yr,g} = \\left\\{{supply}\\left(yr,i,j\\right)
\\;\\vert\\; yr,\\, g \\in i,\\, j = trncost \\right\\}
```

```math
\\tilde{trn}_{yr,g} = \\tilde{m}_{yr,g} + \\tilde{cif}_{yr,g}
\\;\\forall\\; g \\neq ins
```
"""
function _partition_trn0!(d::Dict, set::Dict)
    println("  Partitioning trn0, transportation costs")
    d[:trn0]  = filter_with(d[:supply], (g = set[:g], s = "trncost"); drop = true)

    # Adjust transport margins for transport sectors according to CIF/FOB adjustments.
    # Insurance imports are specified as net of adjustments.
    d[:trn0] += edit_with(d[:cif0], Drop(:g,"ins","=="))
    return d[:trn0]
end

"""
`ts(yr,ts,s)`, taxes and subsidies

```math
\\tilde{ts}_{yr,ts,s} = \\left\\{{use}\\left(yr,i,j\\right)
\\;\\vert\\; yr,\\, ts \\in i,\\, s \\in j \\right\\}
```
Treat negative inputs as outputs:

```math
\\tilde{ts}_{yr,ts,s} = - \\tilde{ts}_{yr,ts,s}
\\;\\forall\\; ts = subsidies
```
"""
function _partition_ts0!(d::Dict, set::Dict)
    println("  Partitioning ts0, taxes and subsidies")
    d[:ts0] = filter_with(d[:use], (g = set[:ts], s = set[:s]))
    d[:ts0] = edit_with(d[:ts0], Rename(:g, :ts))
    d[:ts0][d[:ts0][:,:ts] .== "subsidies", :value] *= -1  # treat negative inputs as outputs
    return d[:ts0]
end

"""
`va(yr,va,s)`, value added

```math
\\tilde{va}_{yr,va,s} = \\left\\{{use}\\left(yr,i,j\\right)
\\;\\vert\\; yr,\\, va \\in i,\\, s \\in j \\right\\}
```
"""
function _partition_va0!(d::Dict, set::Dict)
    println("  Partitioning va0, value added")
    d[:va0] = filter_with(d[:use], (g = set[:va], s = set[:s]))
    d[:va0] = edit_with(d[:va0], Rename(:g, :va))
    return d[:va0]
end

"""
`x(yr,g)`, exports of goods and services

```math
\\tilde{x}_{yr,g} = \\left\\{{use}\\left(yr,i,j\\right)
\\;\\vert\\; yr,\\, g \\in i,\\, exports \\in j \\right\\}
```
"""
function _partition_x0!(d::Dict, set::Dict)
    println("  Partitioning x0, exports of goods and services")
    d[:x0] = filter_with(d[:use], (g = set[:g], s = "exports"); drop = true)
    d[:x0] = _remove_imrg(d[:x0], :g => set[:imrg])
    return d[:x0]
end

"""
`y(yr,g)`, gross output
"Move household supply of recycled goods into the domestic output market
from which some may be exported. Net out margin supply from output."

```math
\\tilde{y}_{yr,g} = \\sum_{s}\\tilde{ys}_{yr,s,g} + \\tilde{fd}_{yr,g} - \\sum_{m}\\tilde{ms}_{yr,g,m}
```
"""
function _partition_y0!(d::Dict, set::Dict)
    println("  Partitioning y0, gross output")

    d[:y0] = if d[:sector]==:summary
        combine_over(d[:ys0], :s) + d[:fs0] - combine_over(d[:ms0], :m)
    elseif d[:sector]==:detail
        combine_over(d[:ys0], :s) + d[:fs0]
    end

    d[:y0] = _remove_imrg(d[:y0], :g => set[:imrg])
    return d[:y0]
end

max_stars_repo_name: NREL/SLiDE

max_issues_repo_name: NREL/SLiDE

max_forks_repo_name: NREL/SLiDE


text:
let
    using UnsteadyFlowSolvers

    # Set up a single vortex
    start_vc = 0.02 # Core radius
    start_vort = 1.
    re = 1e3
    ref_len = 1
    nu = start_vort / re
    vort = TwoDVVort(0., 0., start_vort, start_vc, 0., 0., 0.)
    field = TwoDVFlowField(nu)
    push!(field.tev, vort)

    # Solver settings
    dt = 0.015
    t0 = 0.
    tmax = 1.5
    t = t0
    nsteps = Int(ceil((tmax-t)/dt)) + 1

    # Comparison to the (Gaussian) Lamb-Oseen vortex
    rc = t->sqrt(4 * nu * (t-t0) + start_vc^2)

    for i = 1:nsteps
        t += dt
        UnsteadyFlowSolvers.wakeroll(field, dt)
        println("Step\t",i,":\tCore radius is \t",field.tev[1].vc, "\tLamb-Oseen radius is\t", rc(t))
    end
end


text:
using Base.Test

using BackgroundSegmenter

A = [1 0 1 1; 1 0 1 0; 0 1 1 0; 1 1 0 0]
B = [1 0 1 1; 0 0 1 0; 0 0 0 0; 0 0 1 1]
C = zeros(Int64, 2, 4, 4)
C[1, :, :] = A
C[2, :, :] = B
out = filter_components(A, 3)
@test out == [0 0 1 1 ; 0 0 1 0 ; 0 1 1 0; 1 1 0 0]
out = filter_components(A, 2)
@test out == [1 0 2 2 ; 1 0 2 0 ; 0 2 2 0; 2 2 0 0]
out = filter_components(A, 8)
@test out == zeros(A)
out = filter_components(C, 2)
@test sum(out) == 29
out = filter_components(C, 3)
@test sum(out) == 23
out = filter_components(C, 4)
@test sum(out) == sum(C) - 5
out = filter_components(C, 12)
@test out == zeros(C)

fgbg = filter_components(A, 3)
out = components(fgbg)
@test length(out) == 1
@test size(out[1])[1] == 7


text:
module AntColony

export aco

using Statistics

"""
    aco(dist_mat; start_node=nothing, end_node=nothing; is_tour=false, beta=2, rho=0.1,
        q=0.1, Q=1.0, tau_min=1.0, tau_max=10.0, max_iter=20, reset_iter=10,
        top_perc_ants=0.05, verbose=false)

Ant Colony Optimization for directed graphs based on Max-Min Ant System, Elitist Ant System and Ant Colony System.

# Arguments
- `dist_mat`: must be a square distance matrix (NxN) and edges are accessed like this `dist_mat[to, from]`
- `start_node`: must be an integer <= N or `nothing` (both nodes or none must be specified)
- `end_node`: must be an integer <= N or `nothing` (both nodes or none must be specified)
- `is_tour`: whether path should start and end with same node (will set `end_node = start_node`)
- `beta`: influences how important distance is for the ants decision making
- `rho`: pheromones evaporation percentage for each iteration (`0.0 <= rho <= 1.0`)
- `q`: percentage of roulette wheel style decisions of ants
- `Q`: pheromone deposit factor
- `tau_min`: lower bound for pheromone levels for any edge
- `tau_max`: upper bound for pheremone levels for any edge (initial level of pheromones)
- `max_iter`: number of iterations run for the algorithm
- `reset_iter`: number of iterations after which the pheromone levels are reset to `tau_max`
- `top_perc_ants`: the percentage of top perfoming ants which deposit pheromones on their trails
- `verbose`: print to console whenever a better path was found

# Examples
## Find a tour
Find any path with the same start and end node i.e. a tour. Note that the start node only appears once!

```julia-repl
julia> distances = rand(5, 5);
julia> aco(distances, is_tour = true)
5-element Array{Int64,1}:
 4
 3
 5
 1
 2
```

## Find a path
Finda a specific path from node 1 to node 5

```julia-repl
julia> distances = rand(5, 5);
julia> aco(distances, start_node = 1, end_node = 5)
5-element Array{Int64,1}:
 1
 2
 4
 3
 5
```
"""
function aco(
        dist_mat::AbstractMatrix{<:Number};
        start_node::Union{Nothing, Int} = nothing,
        end_node::Union{Nothing, Int} = nothing,
        is_tour = false,
        beta = 2,
        rho = 0.1,
        q = 0.1,
        Q = 1.0,
        tau_min = 1.0,
        tau_max = 10.0,
        max_iter = 20,
        reset_iter = 10,
        top_perc_ants = 0.05,
        verbose = false
    )

    @assert isnothing(start_node) && isnothing(end_node) || !isnothing(start_node) && !isnothing(end_node)
    @assert if is_tour && (!isnothing(start_node) || !isnothing(end_node)); start_node != end_node else true end

    _, n_nodes = size(dist_mat)
    ants = trunc(Int, n_nodes)
    top_k_ants = max(1, trunc(Int, top_perc_ants * ants))

    # how many factors an edge is better than the expected edge (median) to that node
    η = (median(dist_mat, dims=1)' ./ dist_mat) .^ beta
    τ = fill(tau_max, n_nodes, n_nodes)

    no_improv = 0
    best_path = nothing
    best_cost = Inf

    for i in 1:max_iter
        # probability matrix for the ants decision making
        P = τ .* η
        solutions = []
        _start_node, _end_node = start_node, end_node

        for _ in 1:ants
            # select random start/end nodes
            if isnothing(start_node)
                _start_node = rand(1:n_nodes)
                if is_tour
                    _end_node = _start_node
                end
            end
            path = travel(P, q, _start_node, _end_node)
            cost = sum(edge_distances(dist_mat, path, is_tour))
            push!(solutions, (cost, path))
        end

        # update if better solution found
        sort!(solutions)
        best_local_cost, best_local_path = solutions[1]
        if best_local_cost < best_cost
            if verbose
                println("Better solution found with cost $(best_local_cost) at iteration $(i)")
            end
            best_path = best_local_path
            best_cost = best_local_cost
            no_improv = 0
        else
            no_improv += 1
        end

        # deposit pheromones
        for (cost, path) in solutions[1:top_k_ants]
            Δτ = Q / cost
            for (from, to) in edges(path, is_tour)
                τ[to, from] += Δτ
            end
        end

        # reset when stagnation
        if no_improv > reset_iter
            no_improv = 0
            τ = fill(tau_max, n_nodes, n_nodes)
        else
            # evaporate pheromones
            τ = max.(tau_min, min.(tau_max, τ * (1 - rho)))
        end

    end

    best_path
end

"""
    travel(P, q, start_node, end_node)

Travel through graph based on probability `P`
"""
function travel(
        P::Matrix{<:Number},
        q::Float64,
        start_node::Int,
        end_node::Union{Nothing, Int} = nothing
    )

    not_visited = trues(size(P)[1])
    n_nodes = count(not_visited)
    not_visited[start_node] = false
    if !isnothing(end_node)
        not_visited[end_node] = false
    end
    path = Array{Int}(undef, n_nodes)
    path[1] = start_node

    node = start_node

    for i in 2:n_nodes
        if i == n_nodes && !isnothing(end_node) && start_node != end_node
            not_visited[end_node] = true
        end

        neighbor_idx = if rand() < q
            sample(1:count(not_visited), P[not_visited, node])
        else
            argmax(P[not_visited, node])
        end
        next_node = findall(not_visited)[neighbor_idx]

        not_visited[next_node] = false
        path[i] = next_node
        node = next_node
    end
    path
end

"""
    edges(path, is_tour)

Generate node pairs representing edges in `path`
"""
function edges(path::AbstractArray{Int}, is_tour::Bool)
    path_len = length(path)
    n_edges = if is_tour path_len else path_len - 1 end
    (
        (path[i], path[mod1(i + 1, path_len)])
        for i in 1:n_edges
    )
end

"""
    edge_distances(dist_mat, path, is_tour)

Generate a distances for each edge in `path`
"""
function edge_distances(
        dist_mat::AbstractMatrix{<:Number},
        path::AbstractArray{Int},
        is_tour::Bool
    )
    (dist_mat[to, from] for (from, to) in edges(path, is_tour))
end

"""
    sample(x, weights)

Sample from `x` with a custom distribution defined by `weights`
"""
function sample(x, weights)
    ecdf = cumsum(weights)
    p = rand() * ecdf[end]
    for i in 1:length(ecdf)
        if p <= ecdf[i]
            return x[i]
        end
    end
end

end


text:
using JSON

module Spark

include("Types.jl")
include("test_functions.jl")
include("RPC.jl")
include("Worker.jl")
include("Master.jl")
include("Partitioner.jl")
include("Transformations.jl")
include("Actions.jl")

end


text:
"This defines kernel in addition to BaytesMCMC, BaytesFilters and BaytesPMCMC."
################################################################################
#Import
include("SMC2/SMC2.jl")

################################################################################
#Export
#export


text: using Test, EngTherm


text:
export leaky_relu

doc"""
    leaky_relu(x::Var, alpha=0.1)

Leaky Rectified Linear Unit.

```math
f(x) =
\begin{cases}
x & x > 0 \\
\alpha x & x \leq 0
\end{cases}
```

# References
* Maas et al., ["Rectifier Nonlinearities Improve Neural Network Acoustic Models"](http://web.stanford.edu/~awni/papers/relu_hybrid_icml2013_final.pdf), ICML 2013.
"""
leaky_relu(x::Var, alpha=0.1) = Var(leaky_relu.(x.data,eltype(x)(alpha)), (leaky_relu,x,alpha))
leaky_relu(x::T, alpha::T) where T = x >= zero(T) ? x : x*alpha
leaky_relu(x::Node, alpha=0.1) = Node(leaky_relu, x, alpha)

function addgrad!(y::Var, ::typeof(leaky_relu), x::Var, alpha::Float64)
    isvoid(x.grad) && return
    ∇leaky_relu!(y.grad, x.data, x.grad, eltype(x)(alpha))
end

function ∇leaky_relu!(gy::Array{T}, x::Array{T}, gx::Array{T}, alpha::T) where T
    @inbounds for i = 1:length(x)
        gx[i] += x[i] >= zero(T) ? gy[i] : gy[i]*alpha
    end
end


text:
using Random

R = Float64
ntests = 100
nx = 100

for i in 1:ntests
    a = rand(R)
    b = 2.0 * rand(R)
    x = randn(R, nx)
    y = similar(x)

    nnzy = proxl0nonneg!( x, a, y )
    @assert all(y .>= 0.0)
    @assert sum(y .> 0.0) == nnzy

    nnzy = proxl0simplex!( x, a, b, y )
    @assert isapprox(sum(y), b)
    @assert all(y .>= 0.0)
    @assert sum(y .> 0.0) == nnzy

    projsimplex!(x, b, y)
    @assert isapprox(sum(y), b)
    @assert all(y .>= 0.0)

end

@printf "Passed"


text:
_REIF_LT_FLOAT_EPSILON = 1.0e-5

"""
Bridges `CP.Reification{MOI.LessThan}` to MILP constraints.
"""
struct ReificationLessThan2MILPBridge{T <: Real} <: MOIBC.AbstractBridge
    con_bigm::MOI.ConstraintIndex{MOI.ScalarAffineFunction{T}, MOI.LessThan{T}}
    con_smallm::MOI.ConstraintIndex{MOI.ScalarAffineFunction{T}, MOI.GreaterThan{T}}
end

function MOIBC.bridge_constraint(
    ::Type{ReificationLessThan2MILPBridge{T}},
    model,
    f::MOI.VectorOfVariables,
    s::CP.Reification{MOI.LessThan{T}},
) where {T}
    return MOIBC.bridge_constraint(
        ReificationLessThan2MILPBridge{T},
        model,
        MOI.VectorAffineFunction{T}(f),
        s,
    )
end

function MOIBC.bridge_constraint(
    ::Type{ReificationLessThan2MILPBridge{T}},
    model,
    f::MOI.VectorAffineFunction{T},
    s::CP.Reification{MOI.LessThan{T}},
) where {T <: Real}
    f_scalars = MOIU.scalarize(f)

    # For this formulation work, both lower and upper bounds are required on
    # the constrained expression (but obviously not on the binary reified
    # variable).
    @assert CP.is_binary(model, f_scalars[1])
    @assert CP.has_lower_bound(model, f_scalars[2])
    @assert CP.has_upper_bound(model, f_scalars[2])

    # If the reified expression is true/false, then the LessThan constraint
    # must/cannot be satisfied. (If the constraint is satisfied, the reified
    # expression is unconstrained.)
    bigm = T(max(
        abs(CP.get_upper_bound(model, f_scalars[2])),
        abs(CP.get_lower_bound(model, f_scalars[2]))
    ))
    con_bigm = MOI.add_constraint(
        model,
        f_scalars[2] - bigm * (one(T) - f_scalars[1]),
        MOI.LessThan(s.set.upper)
    )

    # If the constraint is satisfied, constrain the reified.
    smallm = if T <: Int
        one(T)
    else
        T(_REIF_LT_FLOAT_EPSILON)
    end

    con_smallm = MOI.add_constraint(
        model,
        f_scalars[2] - bigm * f_scalars[1],
        MOI.GreaterThan(s.set.upper + smallm)
    )

    return ReificationLessThan2MILPBridge{T}(con_bigm, con_smallm)
end

function MOI.supports_constraint(
    ::Type{ReificationLessThan2MILPBridge{T}},
    ::Union{Type{MOI.VectorOfVariables}, Type{MOI.VectorAffineFunction{T}}},
    ::Type{CP.Reification{MOI.LessThan{T}}},
) where {T <: Real}
    return true
end

function MOIB.added_constrained_variable_types(::Type{ReificationLessThan2MILPBridge{T}}) where {T}
    return Tuple{Type}[]
end

function MOIB.added_constraint_types(::Type{ReificationLessThan2MILPBridge{T}}) where {T}
    return [
        (MOI.VectorAffineFunction{T}, MOI.GreaterThan{T}),
        (MOI.VectorAffineFunction{T}, MOI.LessThan{T}),
    ]
end

function MOI.get(
    ::ReificationLessThan2MILPBridge{T},
    ::MOI.NumberOfConstraints{
        MOI.ScalarAffineFunction{T}, MOI.LessThan{T},
    },
) where {T <: Real}
    return 1
end

function MOI.get(
    ::ReificationLessThan2MILPBridge{T},
    ::MOI.NumberOfConstraints{
        MOI.ScalarAffineFunction{T}, MOI.GreaterThan{T},
    },
) where {T <: Real}
    return 1
end

function MOI.get(
    b::ReificationLessThan2MILPBridge{T},
    ::MOI.ListOfConstraintIndices{
        MOI.ScalarAffineFunction{T}, MOI.LessThan{T},
    },
) where {T <: Real}
    return [b.con_bigm]
end

function MOI.get(
    b::ReificationLessThan2MILPBridge{T},
    ::MOI.ListOfConstraintIndices{
        MOI.ScalarAffineFunction{T}, MOI.GreaterThan{T},
    },
) where {T <: Real}
    return [b.con_smallm]
end


text:
""""""
function corpus(db::SQLite.DB, tablename::AbstractString,
                colname::AbstractString)
    table = DBInterface.execute(db, """
        SELECT $colname FROM $tablename;
    """) |> Tables.columntable
    rows = table[Symbol(colname)]
    return Corpus(StringDocument.(rows))
end

""""""
function readdb(filename::AbstractString, tablename::AbstractString)
    db = SQLite.DB(filename)
    df = DBInterface.execute(db, """
        SELECT * FROM $tablename;
    """) |> DataFrame
    return df
end

""""""
function save(M::AbstractMatrix, filename::AbstractString)
    result = Dict(enumerate(eachcol(M)))

    open(filename, "w") do io
        JSON3.write(io, result)
    end
    open(filename * ".pretty", "w") do io
        JSON3.pretty(io, result)
    end
end

""""""
function save(clustering::DbscanResult, filename::AbstractString)
    y = assignments(clustering) .+ 1
    k = nclusters(clustering) + 1 # unassigned instances / noise
    ns = counts(clustering)

    enumeratedy = enumerate(y)
    assigns = map(collect, enumeratedy)

    invassigns = Vector{Vector{Union{Int, Vector{Int}}}}(undef, k)
    for i = 1:k
        invassigns[i] = [i, Int[]]
    end
    for (i, yi) in enumeratedy
        push!(invassigns[yi][2], i)
    end

    clusters = Vector{Dict}(undef, k)
    proto = length(invassigns[1][2]) == 0 ? 0 : invassigns[1][2][end]
    clusters[1] = Dict("id" => 1, "n" => length(y) - sum(ns), "J" => 0.0, "proto" => proto)
    for i = 2:k
        proto = length(invassigns[i][2]) == 0 ? 0 : invassigns[i][2][end]
        clusters[i] = Dict("id" => i, "n"  => ns[i - 1] , "J" => 0.0, "proto" => proto)
    end

    result = Dict("clusters" => clusters, "assignments" => assigns, "invAssignments" => invassigns)

    open(filename, "w") do io
        JSON3.write(io, result)
    end
    open(filename * ".pretty", "w") do io
        JSON3.pretty(io, result)
    end
end


text:
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

#=doc
SVMLight / LibSVM is a popular data format for sparse features. Some preprocessed
datasets in this format could be found at http://www.csie.ntu.edu.tw/~cjlin/libsvmtools/datasets/
=#
using MXNet
using SVMLightLoader

mutable struct SVMLightProvider <: mx.AbstractDataProvider
  filename   :: AbstractString
  batch_size :: Int
  fea_dim    :: Int
  data_name  :: Symbol
  label_name :: Symbol
end

function SVMLightProvider(filename::AbstractString, batch_size::Int; fea_dim::Int=-1,
                          data_name::Symbol=:data, label_name::Symbol=:label)
  if fea_dim == -1
    info("SVMLightProvider: going over file to get feature dimension of $filename")
    f = SVMLightFile(filename)
    for (data, label) in f
      fea_dim = max(fea_dim, length(data))
    end
  end

  return SVMLightProvider(filename, batch_size, fea_dim, data_name, label_name)
end

mx.get_batch_size(provider :: SVMLightProvider) = provider.batch_size
function mx.provide_data(provider :: SVMLightProvider)
  [(provider.data_name, (provider.fea_dim, provider.batch_size))]
end
function mx.provide_label(provider :: SVMLightProvider)
  [(provider.label_name, (provider.batch_size,))]
end

function mx.eachbatch(provider :: SVMLightProvider)
  data_jl  = zeros(mx.MX_float, (provider.fea_dim, provider.batch_size))
  data_nd  = mx.empty(size(data_jl))
  label_jl = zeros(mx.MX_float, (provider.batch_size,))
  label_nd = mx.empty(size(label_jl))

  batch = mx.DataBatch([data_nd], [label_nd], provider.batch_size)
  function _svmlight_iter()
    f = SVMLightFile(provider.filename)
    while true
      error("This is actually buggy and needs fixing")
      raw = collect(take(f, provider.batch_size))
      cnt = length(raw)
      if cnt == 0
        # end of file, no more data to see
        return
      end

      data_jl[:] = 0
      for i = 1:provider.batch_size
        vec, gnd = raw[min(i,cnt)]
        data_jl[1:length(vec),i] = vec
        label_jl[i]  = gnd
      end
      mx.copy!(data_nd, data_jl)
      mx.copy!(label_nd, label_jl)
      batch.count = cnt
      produce(batch)
    end
  end

  return Task(_svmlight_iter)
end

max_stars_repo_name: Vikas-kum/incubator-mxnet

max_issues_repo_name: Vikas-kum/incubator-mxnet

max_forks_repo_name: Vikas-kum/incubator-mxnet


text:
"""
    rand(::Type{MC}, m::Model)

Draw random configuration.
"""
Base.rand(::Type{MC}, m::Model) = MethodError(rand, (MC, m))

"""
    propose_local(mc::MC, m::Model, i::Int, conf) -> delta_E, passthrough

Propose a local move for lattice site `i` of current configuration `conf`.
Returns local energy difference `delta_E = E_new - E_old` alongside some
additional information `passthrough` which is passed to `accept_local!`.

See also [`accept_local!`](@ref).
"""
propose_local(mc::MC, m::Model, i::Int, conf) = MethodError(propose_local, (mc, m, i, conf))

"""
    accept_local(mc::MC, m::Model, i::Int, conf, delta_E::Float64, passthrough)

Accept a local move for site `i` of current configuration `conf`.
Arguments `delta_E` and `passthrough` correspond to output of `propose_local()`
for that local move.

See also [`propose_local`](@ref).
"""
accept_local!(mc::MC, m::Model, i::Int, conf, delta_E::Float64, passthrough) =
    MethodError(accept_local!, (mc, m, i, conf, delta_E, passthrough))


text:
# conversion routines between different sparse and dense storage formats

"""
    sparse(x::DenseCuMatrix; fmt=:csc)
    sparse(I::CuVector, J::CuVector, V::CuVector, [m, n]; fmt=:csc)

Return a sparse cuda matrix, with type determined by `fmt`.
Possible formats are :csc, :csr, :bsr, and :coo.
"""
function SparseArrays.sparse(x::DenseCuMatrix; fmt=:csc)
    if fmt == :csc
        return CuSparseMatrixCSC(x)
    elseif fmt == :csr
        return CuSparseMatrixCSR(x)
    elseif fmt == :bsr
        return CuSparseMatrixBSR(x)
    elseif fmt == :coo
        return CuSparseMatrixCOO(x)
    else
        error("Format :$fmt not available, use :csc, :csr, :bsr or :coo.")
    end
end

SparseArrays.sparse(I::CuVector, J::CuVector, V::CuVector; kws...) =
    sparse(I, J, V, maximum(I), maximum(J); kws...)

SparseArrays.sparse(I::CuVector, J::CuVector, V::CuVector, m, n; kws...) =
    sparse(Cint.(I), Cint.(J), V, m, n; kws...)

function SparseArrays.sparse(I::CuVector{Cint}, J::CuVector{Cint}, V::CuVector{Tv}, m, n;
                             fmt=:csc, sorted::Bool=false) where Tv
    coo = CuSparseMatrixCOO{Tv}(I, J, V, (m, n))

    # The COO format is assumed to be sorted by row.
    if !sorted
        coo = sort_rows(coo)
    end

    if fmt == :csc
        return CuSparseMatrixCSC(coo)
    elseif fmt == :csr
        return CuSparseMatrixCSR(coo)
    elseif fmt == :coo
        return coo
    else
        error("Format :$fmt not available, use :csc, :csr, or :coo.")
    end
end

function sort_rows(coo::CuSparseMatrixCOO{Tv,Ti}) where {Tv <: BlasFloat,Ti}
    m,n = size(coo)

    perm = CuArray{Ti}(undef, nnz(coo))
    cusparseCreateIdentityPermutation(handle(), nnz(coo), perm)

    sorted_rowInd = copy(coo.rowInd)
    sorted_colInd = copy(coo.colInd)
    function bufferSize()
        out = Ref{Csize_t}()
        cusparseXcoosort_bufferSizeExt(handle(), m, n, nnz(coo), coo.rowInd,
            coo.colInd, out)
        return out[]
    end
    with_workspace(bufferSize) do buffer
        cusparseXcoosortByRow(handle(), m, n, nnz(coo), sorted_rowInd, sorted_colInd, perm, buffer)
    end

    sorted_nzVal = similar(coo.nzVal)
    let spvec = CuSparseVector(perm, sorted_nzVal, nnz(coo))
        if version() >= v"11.1.1"
            gather!(spvec, nonzeros(coo), 'Z')
        else
            gthr!(spvec, nonzeros(coo), 'Z')
        end
    end

    CUDA.unsafe_free!(perm)
    CuSparseMatrixCOO{Tv}(sorted_rowInd, sorted_colInd, sorted_nzVal, size(coo))
end
function sort_rows(coo::CuSparseMatrixCOO{Tv,Ti}) where {Tv,Ti}
    perm = sortperm(coo.rowInd)
    sorted_rowInd = coo.rowInd[perm]
    sorted_colInd = coo.colInd[perm]
    sorted_nzVal = coo.nzVal[perm]
    CUDA.unsafe_free!(perm)

    CuSparseMatrixCOO{Tv}(sorted_rowInd, sorted_colInd, sorted_nzVal, size(coo))
end

## CSR to CSC

function CuSparseMatrixCSC{T}(S::Transpose{T, <:CuSparseMatrixCSR{T}}) where T
    csr = parent(S)
    return CuSparseMatrixCSC{T}(csr.rowPtr, csr.colVal, csr.nzVal, size(csr))
end

function CuSparseMatrixCSR{T}(S::Transpose{T, <:CuSparseMatrixCSC{T}}) where T
    csc = parent(S)
    return CuSparseMatrixCSR{T}(csc.colPtr, csc.rowVal, csc.nzVal, size(csc))
end

function CuSparseMatrixCSC{T}(S::Adjoint{T, <:CuSparseMatrixCSR{T}}) where {T <: Real}
    csr = parent(S)
    return CuSparseMatrixCSC{T}(csr.rowPtr, csr.colVal, csr.nzVal, size(csr))
end

function CuSparseMatrixCSR{T}(S::Adjoint{T, <:CuSparseMatrixCSC{T}}) where {T <: Real}
    csc = parent(S)
    return CuSparseMatrixCSR{T}(csc.colPtr, csc.rowVal, csc.nzVal, size(csc))
end

function CuSparseMatrixCSC{T}(S::Adjoint{T, <:CuSparseMatrixCSR{T}}) where {T <: Complex}
    csr = parent(S)
    return CuSparseMatrixCSC{T}(csr.rowPtr, csr.colVal, conj.(csr.nzVal), size(csr))
end

function CuSparseMatrixCSR{T}(S::Adjoint{T, <:CuSparseMatrixCSC{T}}) where {T <: Complex}
    csc = parent(S)
    return CuSparseMatrixCSR{T}(csc.colPtr, csc.rowVal, conj.(csc.nzVal), size(csc))
end

for SparseMatrixType in [:CuSparseMatrixCSC, :CuSparseMatrixCSR]
    @eval begin
        $SparseMatrixType(S::Diagonal) = $SparseMatrixType(cu(S))
        $SparseMatrixType(S::Diagonal{T, <:CuArray}) where T = $SparseMatrixType{T}(S)
        $SparseMatrixType{Tv}(S::Diagonal{T, <:CuArray}) where {Tv, T} = $SparseMatrixType{Tv, Cint}(S)
        function $SparseMatrixType{Tv, Ti}(S::Diagonal{T, <:CuArray}) where {Tv, Ti, T}
            m = size(S, 1)
            return $SparseMatrixType{Tv, Ti}(CuVector(1:(m+1)), CuVector(1:m), Tv.(S.diag), (m, m))
        end
    end
end

# by flipping rows and columns, we can use that to get CSC to CSR too
for elty in (Float32, Float64, ComplexF32, ComplexF64)
    @eval begin
        function CuSparseMatrixCSC{$elty}(csr::CuSparseMatrixCSR{$elty}; inda::SparseChar='O')
            m,n = size(csr)
            colPtr = CUDA.zeros(Cint, n+1)
            rowVal = CUDA.zeros(Cint, nnz(csr))
            nzVal = CUDA.zeros($elty, nnz(csr))
            # TODO: algorithm configuratibility?
            function bufferSize()
                out = Ref{Csize_t}(1)
                cusparseCsr2cscEx2_bufferSize(handle(), m, n, nnz(csr), nonzeros(csr),
                    csr.rowPtr, csr.colVal, nzVal, colPtr, rowVal,
                    $elty, CUSPARSE_ACTION_NUMERIC, inda,
                    CUSPARSE_CSR2CSC_ALG1, out)
                return out[]
            end
            with_workspace(bufferSize) do buffer
                cusparseCsr2cscEx2(handle(), m, n, nnz(csr), nonzeros(csr),
                    csr.rowPtr, csr.colVal, nzVal, colPtr, rowVal,
                    $elty, CUSPARSE_ACTION_NUMERIC, inda,
                    CUSPARSE_CSR2CSC_ALG1, buffer)
            end
            CuSparseMatrixCSC(colPtr,rowVal,nzVal,size(csr))
        end

        function CuSparseMatrixCSR{$elty}(csc::CuSparseMatrixCSC{$elty}; inda::SparseChar='O')
            m,n    = size(csc)
            rowPtr = CUDA.zeros(Cint,m+1)
            colVal = CUDA.zeros(Cint,nnz(csc))
            nzVal  = CUDA.zeros($elty,nnz(csc))
            # TODO: algorithm configuratibility?
            function bufferSize()
                out = Ref{Csize_t}(1)
                cusparseCsr2cscEx2_bufferSize(handle(), n, m, nnz(csc), nonzeros(csc),
                    csc.colPtr, rowvals(csc), nzVal, rowPtr, colVal,
                    $elty, CUSPARSE_ACTION_NUMERIC, inda,
                    CUSPARSE_CSR2CSC_ALG1, out)
                return out[]
            end
            with_workspace(bufferSize) do buffer
                cusparseCsr2cscEx2(handle(), n, m, nnz(csc), nonzeros(csc),
                    csc.colPtr, rowvals(csc), nzVal, rowPtr, colVal,
                    $elty, CUSPARSE_ACTION_NUMERIC, inda,
                    CUSPARSE_CSR2CSC_ALG1, buffer)
            end
            CuSparseMatrixCSR(rowPtr,colVal,nzVal,size(csc))
        end
    end
end

# implement Float16 conversions using wider types
# TODO: Float16 is sometimes natively supported
for (elty, welty) in ((:Float16, :Float32),
                      (:ComplexF16, :ComplexF32))
    @eval begin
        function CuSparseMatrixCSC{$elty}(csr::CuSparseMatrixCSR{$elty}; inda::SparseChar='O')
            m,n = size(csr)
            colPtr = CUDA.zeros(Cint, n+1)
            rowVal = CUDA.zeros(Cint, nnz(csr))
            nzVal = CUDA.zeros($elty, nnz(csr))
            # TODO: algorithm configuratibility?
            if $elty == Float16 #broken for ComplexF16?
                function bufferSize()
                    out = Ref{Csize_t}(1)
                    cusparseCsr2cscEx2_bufferSize(handle(), m, n, nnz(csr), nonzeros(csr),
                        csr.rowPtr, csr.colVal, nzVal, colPtr, rowVal,
                        $elty, CUSPARSE_ACTION_NUMERIC, inda,
                        CUSPARSE_CSR2CSC_ALG1, out)
                    return out[]
                end
                with_workspace(bufferSize) do buffer
                    cusparseCsr2cscEx2(handle(), m, n, nnz(csr), nonzeros(csr),
                        csr.rowPtr, csr.colVal, nzVal, colPtr, rowVal,
                        $elty, CUSPARSE_ACTION_NUMERIC, inda,
                        CUSPARSE_CSR2CSC_ALG1, buffer)
                end
                return CuSparseMatrixCSC(colPtr,rowVal,nzVal,size(csr))
            else
                wide_csr = CuSparseMatrixCSR(csr.rowPtr, csr.colVal, convert(CuVector{$welty}, nonzeros(csr)), size(csr))
                wide_csc = CuSparseMatrixCSC(wide_csr)
                return CuSparseMatrixCSC(wide_csc.colPtr, wide_csc.rowVal, convert(CuVector{$elty}, nonzeros(wide_csc)), size(wide_csc))
            end
        end

        function CuSparseMatrixCSR{$elty}(csc::CuSparseMatrixCSC{$elty}; inda::SparseChar='O')
            m,n    = size(csc)
            rowPtr = CUDA.zeros(Cint,m+1)
            colVal = CUDA.zeros(Cint,nnz(csc))
            nzVal  = CUDA.zeros($elty,nnz(csc))
            if $elty == Float16 #broken for ComplexF16?
                # TODO: algorithm configuratibility?
                function bufferSize()
                    out = Ref{Csize_t}(1)
                    cusparseCsr2cscEx2_bufferSize(handle(), n, m, nnz(csc), nonzeros(csc),
                        csc.colPtr, rowvals(csc), nzVal, rowPtr, colVal,
                        $elty, CUSPARSE_ACTION_NUMERIC, inda,
                        CUSPARSE_CSR2CSC_ALG1, out)
                    return out[]
                end
                with_workspace(bufferSize) do buffer
                    cusparseCsr2cscEx2(handle(), n, m, nnz(csc), nonzeros(csc),
                        csc.colPtr, rowvals(csc), nzVal, rowPtr, colVal,
                        $elty, CUSPARSE_ACTION_NUMERIC, inda,
                        CUSPARSE_CSR2CSC_ALG1, buffer)
                end
                return CuSparseMatrixCSR(rowPtr,colVal,nzVal,size(csc))
            else
                wide_csc = CuSparseMatrixCSC(csc.colPtr, csc.rowVal, convert(CuVector{$welty}, nonzeros(csc)), size(csc))
                wide_csr = CuSparseMatrixCSR(wide_csc)
                return CuSparseMatrixCSR(wide_csr.rowPtr, wide_csr.colVal, convert(CuVector{$elty}, nonzeros(wide_csr)), size(wide_csr))
            end
        end
    end
end

# implement Int conversions using reinterpreted Float
for (elty, felty) in ((:Int16, :Float16),
                      (:Int32, :Float32),
                      (:Int64, :Float64),
                      (:Int128, :ComplexF64))
    @eval begin
        function CuSparseMatrixCSR{$elty}(csc::CuSparseMatrixCSC{$elty})
            csc_compat = CuSparseMatrixCSC(
                csc.colPtr,
                csc.rowVal,
                reinterpret($felty, csc.nzVal),
                size(csc)
            )
            csr_compat = CuSparseMatrixCSR(csc_compat)
            CuSparseMatrixCSR(
                csr_compat.rowPtr,
                csr_compat.colVal,
                reinterpret($elty, csr_compat.nzVal),
                size(csr_compat)
            )
        end

        function CuSparseMatrixCSC{$elty}(csr::CuSparseMatrixCSR{$elty})
            csr_compat = CuSparseMatrixCSR(
                csr.rowPtr,
                csr.colVal,
                reinterpret($felty, csr.nzVal),
                size(csr)
            )
            csc_compat = CuSparseMatrixCSC(csr_compat)
            CuSparseMatrixCSC(
                csc_compat.colPtr,
                csc_compat.rowVal,
                reinterpret($elty, csc_compat.nzVal),
                size(csc_compat)
            )
        end
    end
end

## CSR to BSR and vice-versa

for (fname,elty) in ((:cusparseScsr2bsr, :Float32),
                     (:cusparseDcsr2bsr, :Float64),
                     (:cusparseCcsr2bsr, :ComplexF32),
                     (:cusparseZcsr2bsr, :ComplexF64))
    @eval begin
        function CuSparseMatrixBSR{$elty}(csr::CuSparseMatrixCSR{$elty}, blockDim::Integer;
                                          dir::SparseChar='R', inda::SparseChar='O',
                                          indc::SparseChar='O')
            m,n = size(csr)
            nnz_ref = Ref{Cint}(1)
            mb = cld(m, blockDim)
            nb = cld(n, blockDim)
            bsrRowPtr = CUDA.zeros(Cint,mb + 1)
            cudesca = CuMatrixDescriptor('G', 'L', 'N', inda)
            cudescc = CuMatrixDescriptor('G', 'L', 'N', indc)
            cusparseXcsr2bsrNnz(handle(), dir, m, n, cudesca, csr.rowPtr,
                                csr.colVal, blockDim, cudescc, bsrRowPtr, nnz_ref)
            bsrNzVal = CUDA.zeros($elty, nnz_ref[] * blockDim * blockDim )
            bsrColInd = CUDA.zeros(Cint, nnz_ref[])
            $fname(handle(), dir, m, n,
                   cudesca, nonzeros(csr), csr.rowPtr, csr.colVal,
                   blockDim, cudescc, bsrNzVal, bsrRowPtr,
                   bsrColInd)
            CuSparseMatrixBSR{$elty}(bsrRowPtr, bsrColInd, bsrNzVal, size(csr), blockDim, dir, nnz_ref[])
        end
    end
end

for (fname,elty) in ((:cusparseSbsr2csr, :Float32),
                     (:cusparseDbsr2csr, :Float64),
                     (:cusparseCbsr2csr, :ComplexF32),
                     (:cusparseZbsr2csr, :ComplexF64))
    @eval begin
        function CuSparseMatrixCSR{$elty}(bsr::CuSparseMatrixBSR{$elty};
                                          inda::SparseChar='O', indc::SparseChar='O')
            m,n = size(bsr)
            mb = cld(m, bsr.blockDim)
            nb = cld(n, bsr.blockDim)
            cudesca = CuMatrixDescriptor('G', 'L', 'N', inda)
            cudescc = CuMatrixDescriptor('G', 'L', 'N', indc)
            csrRowPtr = CUDA.zeros(Cint, m + 1)
            csrColInd = CUDA.zeros(Cint, nnz(bsr))
            csrNzVal  = CUDA.zeros($elty, nnz(bsr))
            $fname(handle(), bsr.dir, mb, nb,
                   cudesca, nonzeros(bsr), bsr.rowPtr, bsr.colVal,
                   bsr.blockDim, cudescc, csrNzVal, csrRowPtr,
                   csrColInd)
            # XXX: the size here may not match the expected size, when the matrix dimension
            #      is not a multiple of the block dimension!
            CuSparseMatrixCSR(csrRowPtr, csrColInd, csrNzVal, (mb*bsr.blockDim, nb*bsr.blockDim))
        end
    end
end

# implement Int conversions using reinterpreted Float
for (elty, felty) in ((:Int16, :Float16),
                      (:Int32, :Float32),
                      (:Int64, :Float64),
                      (:Int128, :ComplexF64))
    @eval begin
        function CuSparseMatrixCSR{$elty}(bsr::CuSparseMatrixBSR{$elty})
            bsr_compat = CuSparseMatrixBSR(
                bsr.rowPtr,
                bsr.colVal,
                reinterpret($felty, bsr.nzVal),
                bsr.blockDim,
                bsr.dir,
                bsr.nnzb,
                size(bsr)
            )
            csr_compat = CuSparseMatrixCSR(bsr_compat)
            CuSparseMatrixCSR(
                csr_compat.rowPtr,
                csr_compat.colVal,
                reinterpret($elty, csr_compat.nzVal),
                size(csr_compat)
            )
        end

        function CuSparseMatrixBSR{$elty}(csr::CuSparseMatrixCSR{$elty}, blockDim)
            csr_compat = CuSparseMatrixCSR(
                csr.rowPtr,
                csr.colVal,
                reinterpret($felty, csr.nzVal),
                size(csr)
            )
            bsr_compat = CuSparseMatrixBSR(csr_compat, blockDim)
            CuSparseMatrixBSR(
                bsr_compat.rowPtr,
                bsr_compat.colVal,
                reinterpret($elty, bsr_compat.nzVal),
                bsr_compat.blockDim,
                bsr_compat.dir,
                bsr_compat.nnzb,
                size(bsr_compat)
            )
        end
    end
end

## CSR to COO and vice-versa

# TODO: we can do similar for CSC conversions, but that requires the columns to be sorted

function CuSparseMatrixCSR(coo::CuSparseMatrixCOO{Tv}, ind::SparseChar='O') where {Tv}
    m,n = size(coo)
    csrRowPtr = CuVector{Cint}(undef, m+1)
    cusparseXcoo2csr(handle(), coo.rowInd, nnz(coo), m, csrRowPtr, ind)
    CuSparseMatrixCSR{Tv}(csrRowPtr, coo.colInd, nonzeros(coo), size(coo))
end

function CuSparseMatrixCOO(csr::CuSparseMatrixCSR{Tv}, ind::SparseChar='O') where {Tv}
    m,n = size(csr)
    cooRowInd = CuVector{Cint}(undef, nnz(csr))
    cusparseXcsr2coo(handle(), csr.rowPtr, nnz(csr), m, cooRowInd, ind)
    CuSparseMatrixCOO{Tv}(cooRowInd, csr.colVal, nonzeros(csr), size(csr), nnz(csr))
end

### CSC/BSR to COO and viceversa

CuSparseMatrixCSC(coo::CuSparseMatrixCOO) = CuSparseMatrixCSC(CuSparseMatrixCSR(coo)) # no direct conversion
CuSparseMatrixCOO(csc::CuSparseMatrixCSC) = CuSparseMatrixCOO(CuSparseMatrixCSR(csc)) # no direct conversion
CuSparseMatrixBSR(coo::CuSparseMatrixCOO, blockdim) = CuSparseMatrixBSR(CuSparseMatrixCSR(coo), blockdim) # no direct conversion
CuSparseMatrixCOO(bsr::CuSparseMatrixBSR) = CuSparseMatrixCOO(CuSparseMatrixCSR(bsr)) # no direct conversion

## sparse to dense, and vice-versa

for (cname,rname,elty) in ((:cusparseScsc2dense, :cusparseScsr2dense, :Float32),
                           (:cusparseDcsc2dense, :cusparseDcsr2dense, :Float64),
                           (:cusparseCcsc2dense, :cusparseCcsr2dense, :ComplexF32),
                           (:cusparseZcsc2dense, :cusparseZcsr2dense, :ComplexF64))
    @eval begin
        function CUDA.CuMatrix{$elty}(csr::CuSparseMatrixCSR{$elty}; ind::SparseChar='O')
            m,n = size(csr)
            denseA = CUDA.zeros($elty,m,n)
            if version() >= v"11.3.0" # CUSPARSE version from CUDA release notes
                desc_csr   = CuSparseMatrixDescriptor(csr, ind)
                desc_dense = CuDenseMatrixDescriptor(denseA)

                function bufferSize()
                    out = Ref{Csize_t}(1)
                    cusparseSparseToDense_bufferSize(handle(), desc_csr, desc_dense,
                        CUSPARSE_SPARSETODENSE_ALG_DEFAULT, out)
                    return out[]
                end
                with_workspace(bufferSize) do buffer
                    cusparseSparseToDense(handle(), desc_csr, desc_dense,
                        CUSPARSE_SPARSETODENSE_ALG_DEFAULT, buffer)
                end
                return denseA
            else
                cudesc = CuMatrixDescriptor('G', 'L', 'N', ind)
                lda = max(1,stride(denseA,2))
                $rname(handle(), m, n, cudesc, nonzeros(csr),
                       csr.rowPtr, csr.colVal, denseA, lda)
                return denseA
            end
        end
        function CUDA.CuMatrix{$elty}(csc::CuSparseMatrixCSC{$elty}; ind::SparseChar='O')
            m,n = size(csc)
            denseA = CUDA.zeros($elty,m,n)
            if version() >= v"11.3.0" # CUSPARSE version from CUDA release notes
                desc_csc   = CuSparseMatrixDescriptor(csc, ind; convert=false)
                desc_dense = CuDenseMatrixDescriptor(denseA)

                function bufferSize()
                    out = Ref{Csize_t}(1)
                    cusparseSparseToDense_bufferSize(handle(), desc_csc, desc_dense,
                        CUSPARSE_SPARSETODENSE_ALG_DEFAULT, out)
                    return out[]
                end
                with_workspace(bufferSize) do buffer
                    cusparseSparseToDense(handle(), desc_csc, desc_dense,
                        CUSPARSE_SPARSETODENSE_ALG_DEFAULT, buffer)
                end
                return denseA
            else
                lda = max(1,stride(denseA,2))
                cudesc = CuMatrixDescriptor('G', 'L', 'N', ind)
                $cname(handle(), m, n, cudesc, nonzeros(csc),
                       rowvals(csc), csc.colPtr, denseA, lda)
                return denseA
            end
        end
    end
end

for (elty, welty) in ((:Float16, :Float32),
                      (:ComplexF16, :ComplexF32))
    @eval begin
        function CUDA.CuMatrix{$elty}(csr::CuSparseMatrixCSR{$elty}; ind::SparseChar='O')
            m,n = size(csr)
            denseA = CUDA.zeros($elty,m,n)
            if version() >= v"11.3.0" # CUSPARSE version from CUDA release notes
                desc_csr   = CuSparseMatrixDescriptor(csr, ind)
                desc_dense = CuDenseMatrixDescriptor(denseA)

                function bufferSize()
                    out = Ref{Csize_t}(1)
                    cusparseSparseToDense_bufferSize(handle(), desc_csr, desc_dense,
                        CUSPARSE_SPARSETODENSE_ALG_DEFAULT, out)
                    return out[]
                end
                with_workspace(bufferSize) do buffer
                    cusparseSparseToDense(handle(), desc_csr, desc_dense,
                        CUSPARSE_SPARSETODENSE_ALG_DEFAULT, buffer)
                end
                return denseA
            else
                wide_csr = CuSparseMatrixCSR(csr.rowPtr, csr.colVal, convert(CuVector{$welty}, nonzeros(csr)), size(csr))
                wide_dense = CuArray{$welty}(wide_csr)
                denseA = convert(CuArray{$elty}, wide_dense)
                return denseA
            end
        end
        function CUDA.CuMatrix{$elty}(csc::CuSparseMatrixCSC{$elty}; ind::SparseChar='O')
            m,n = size(csc)
            denseA = CUDA.zeros($elty,m,n)
            if version() >= v"11.3.0" # CUSPARSE version from CUDA release notes
                desc_csc   = CuSparseMatrixDescriptor(csc, ind; convert=false)
                desc_dense = CuDenseMatrixDescriptor(denseA)

                function bufferSize()
                    out = Ref{Csize_t}(1)
                    cusparseSparseToDense_bufferSize(handle(), desc_csc, desc_dense,
                        CUSPARSE_SPARSETODENSE_ALG_DEFAULT, out)
                    return out[]
                end
                with_workspace(bufferSize) do buffer
                    cusparseSparseToDense(handle(), desc_csc, desc_dense,
                        CUSPARSE_SPARSETODENSE_ALG_DEFAULT, buffer)
                end
                return denseA
            else
                wide_csc = CuSparseMatrixCSC(csc.colPtr, csc.rowVal, convert(CuVector{$welty}, nonzeros(csc)), size(csc))
                wide_dense = CuArray{$welty}(wide_csc)
                denseA = convert(CuArray{$elty}, wide_dense)
                return denseA
            end
        end
    end
end

Base.copyto!(dest::Array{T, 2}, src::AbstractCuSparseMatrix{T}) where T = copyto!(dest, CuMatrix{T}(src))

for (nname,cname,rname,elty) in ((:cusparseSnnz, :cusparseSdense2csc, :cusparseSdense2csr, :Float32),
                                 (:cusparseDnnz, :cusparseDdense2csc, :cusparseDdense2csr, :Float64),
                                 (:cusparseCnnz, :cusparseCdense2csc, :cusparseCdense2csr, :ComplexF32),
                                 (:cusparseZnnz, :cusparseZdense2csc, :cusparseZdense2csr, :ComplexF64))
    @eval begin
        function CuSparseMatrixCSR(A::CuMatrix{$elty}; ind::SparseChar='O')
            m,n = size(A)
            lda = max(1, stride(A,2))
            cudesc = CuMatrixDescriptor('G',
                                        'L',
                                        'N', ind)
            nnzRowCol = CUDA.zeros(Cint, m)
            nnzTotal = Ref{Cint}(1)
            $nname(handle(),
                   'R', m, n, cudesc, A, lda, nnzRowCol,
                   nnzTotal)
            nzVal = CUDA.zeros($elty,nnzTotal[])

            rowPtr = CUDA.zeros(Cint,m+1)
            colInd = CUDA.zeros(Cint,nnzTotal[])
            $rname(handle(), m, n, cudesc, A,
                    lda, nnzRowCol, nzVal, rowPtr, colInd)
            return CuSparseMatrixCSR(rowPtr,colInd,nzVal,size(A))
        end

        function CuSparseMatrixCSC(A::CuMatrix{$elty}; ind::SparseChar='O')
            m,n = size(A)
            lda = max(1, stride(A,2))
            cudesc = CuMatrixDescriptor('G',
                                        'L',
                                        'N', ind)
            nnzRowCol = CUDA.zeros(Cint, n)
            nnzTotal = Ref{Cint}(1)
            $nname(handle(),
                   'C', m, n, cudesc, A, lda, nnzRowCol,
                   nnzTotal)
            nzVal = CUDA.zeros($elty,nnzTotal[])

            colPtr = CUDA.zeros(Cint,n+1)
            rowInd = CUDA.zeros(Cint,nnzTotal[])
            $cname(handle(), m, n, cudesc, A,
                    lda, nnzRowCol, nzVal, rowInd, colPtr)
            return CuSparseMatrixCSC(colPtr,rowInd,nzVal,size(A))
        end
    end
end

# to do: use cusparseDenseToSparse_convert here
for (elty, welty) in ((:Float16, :Float32),
                      (:ComplexF16, :ComplexF32))
    @eval begin
        function CuSparseMatrixCSR(A::CuMatrix{$elty}; ind::SparseChar='O')
            wide_csr = CuSparseMatrixCSR(convert(CuMatrix{$welty}, A))
            return CuSparseMatrixCSR(wide_csr.rowPtr, wide_csr.colVal, convert(CuVector{$elty}, nonzeros(wide_csr)), size(wide_csr))
        end
        function CuSparseMatrixCSC(A::CuMatrix{$elty}; ind::SparseChar='O')
            wide_csc = CuSparseMatrixCSC(convert(CuMatrix{$welty}, A))
            return CuSparseMatrixCSC(wide_csc.colPtr, wide_csc.rowVal, convert(CuVector{$elty}, nonzeros(wide_csc)), size(wide_csc))
        end
    end
end

function CUDA.CuMatrix{T}(bsr::CuSparseMatrixBSR{T}; inda::SparseChar='O',
                          indc::SparseChar='O') where {T}
    CuMatrix{T}(CuSparseMatrixCSR{T}(bsr; inda, indc))
end

function CuSparseMatrixBSR(A::CuMatrix; ind::SparseChar='O')
    m,n = size(A)   # TODO: always let the user choose, or provide defaults for other methods too
    CuSparseMatrixBSR(CuSparseMatrixCSR(A; ind), gcd(m,n))
end


text:
abstract type RandomisationModel end
abstract type TwoArmRandomisationModel <: RandomisationModel end
abstract type MultiArmRandomisationModel <: RandomisationModel end

target(RM::RandomisationModel) = RM.target
narms(RM::RandomisationModel) = length(RM.target)
sequence(RM::RandomisationModel) = RM.sequence
dist(RM::RandomisationModel) = RM.dist
imbalance(RM::RandomisationModel) = euclidean(dist(RM), sum(dist(RM)) * target(RM))
predictability(RM::RandomisationModel) = euclidean(prob(RM), target(RM))

function ==(a::V, b::V) where V <: RandomisationModel
    for name in fieldnames(V)
        if !(getfield(a, name) ≈ getfield(b, name))
            return false
        end
    end
    return true
end

"""
    randomise!(rng::AbstractRNG, RM::RandomisationModel)

Generate a single random allocation from a randomisation model, `RM`,
using `prob(RM)` and the random number generator provided.
"""
function randomise!(rng::AbstractRNG, RM::RandomisationModel)
    y = sample(rng, 1:narms(RM), Weights(prob(RM)))
    update!(RM, y)
    return y
end
randomize!(rng::AbstractRNG, RM::RandomisationModel) = randomise!(rng::AbstractRNG, RM::RandomisationModel)

"""
    randomise!(RM::RandomisationModel)

Generate a single random allocation from a randomisation model using prob(RM) and the global RNG.
"""
function randomise!(RM::RandomisationModel)
    y = sample(1:narms(RM), Weights(prob(RM)))
    update!(RM, y)
    return y
end
randomize!(RM::RandomisationModel) = randomise!(RM::RandomisationModel)

"""
    randomise!(rng::AbstractRNG, RM::RandomisationModel, n::Int)

Generate a n random allocation's from a randomisation model using prob(RM) and
the random number generator provided.
"""
function randomise!(rng::AbstractRNG, RM::RandomisationModel, n::Int)
    # u = rand(rng, n)
    y = zeros(Int, n)
    for i in 1:n
        y[i] = sample(rng, 1:narms(RM), Weights(prob(RM)))
        update!(RM, y[i])
    end
    return y
end
randomize!(rng::AbstractRNG, RM::RandomisationModel, n::Int) = randomise!(rng::AbstractRNG, RM::RandomisationModel, n::Int)

"""
    randomise!(RM::RandomisationModel, n::Int)

Generate a n random allocation's from a randomisation model using prob(RM) and
the global RNG.
"""
function randomise!(RM::RandomisationModel, n::Int)
    # u = rand(rng, n)
    y = zeros(Int, n)
    for i in 1:n
        # y[i] = random_sample(prob(MWU), u[i])
        y[i] = sample(1:narms(RM), Weights(prob(RM)))
        update!(RM, y[i])
    end
    return y
end
randomize!(RM::RandomisationModel, n::Int) = randomise!(RM::RandomisationModel, n::Int)

"""
    randomise(rng::AbstractRNG, RM::RandomisationModel)

Generates allocations from the randomisation model without updating the models.
"""
randomise(rng::AbstractRNG, RM::RandomisationModel) = sample(rng, 1:narms(RM), Weights(prob(RM)))

"""
    randomize(rng::AbstractRNG, RM::RandomisationModel)

Generates allocations from the randomisation model without updating the models.
"""
randomize(rng::AbstractRNG, RM::RandomisationModel) = randomise(rng::AbstractRNG, RM::RandomisationModel)
randomise(RM::RandomisationModel) = sample(1:narms(RM), Weights(prob(RM)))
randomize(RM::RandomisationModel) = randomise(RM::RandomisationModel)

"""
    randomise(rng::AbstractRNG, RM::RandomisationModel, n::Int)

Generates n allocations
"""
function randomise(rng::AbstractRNG, RM::RandomisationModel, n::Int)
    y = zeros(Int, n)
    tmpRM = deepcopy(RM)
    for i in 1:n
        y[i] = sample(rng, 1:narms(tmpRM), Weights(prob(tmpRM)))
        update!(tmpRM, y[i])
    end
    return y
end
randomize(rng::AbstractRNG, RM::RandomisationModel, n::Int) = randomise(rng::AbstractRNG, RM::RandomisationModel, n::Int)

function randomise(RM::RandomisationModel, n::Int)
    y = zeros(Int, n)
    tmpRM = deepcopy(RM)
    for i in 1:n
        y[i] = sample(1:narms(tmpRM), Weights(prob(tmpRM)))
        update!(tmpRM, y[i])
    end
    return y
end
randomize(RM::RandomisationModel, n::Int) = randomise(RM::RandomisationModel, n::Int)


text:
# The time-dependent GPE

"""
    integrate(s, d, ψ₀, ts; μ, dt)

Return a solution to the GPE

The state at t=0 is ψ₀, ts is a set of times to return states.
"""
function integrate(s, d, ψ₀, ts; μ = nothing, dt = default(:dt))
    # TODO amend DifferentialEquations so we can use saveat for this
    L = operators(s, d, :L) |> only
    isnothing(μ) && (μ = real(dot(L(ψ₀), ψ₀)))
    qs = [ψ₀]
    tt = ts[1] ≤ sqrt(eps()) ? ts : [0; ts]

    for j = 2:length(tt)
        P = DifferentialEquations.ODEProblem(
            (ψ, _, _) -> -1im * (L(ψ) - μ * ψ),
            qs[j-1],
            tt[j-1:j],
        )
        S = DifferentialEquations.solve(
            P,
            DifferentialEquations.RK4(),
            adaptive = false;
            dt,
            saveat = tt[j] - tt[j-1],
        )
        push!(qs, S[end])
    end

    length(tt) == length(ts) ? qs : qs[2:end]
end


text:
#
# Sensitivity analysis related types and functions
#

struct Sensitivity <: Wrapper
    pyo::PyObject

    function Sensitivity(pyo::PyObject)
        new(pyo)
    end
end

struct SAResult  <: Wrapper
    pyo::PyObject

    function SAResult(pyo::PyObject)
        new(pyo)
    end
end

function sa(m::Model, response; policy=Dict(), method="sobol", nsamples=1000, kwargs...)
    my_sa = py"my_sa"
    pyo = my_sa(m.pyo, response; policy=policy, method=method, nsamples=nsamples, kwargs...)
    return SAResult(pyo."sa_result")   # access without conversion
end

function plot(sar::SAResult; kwargs...)
    return rhodium.SAResult.plot(sar.pyo; kwargs...)
end

function plot_sobol(sar::SAResult;
                    radSc=2.0, scaling=1, widthSc=0.5, STthick=1, varNameMult=1.3,
                    colors=nothing, groups=nothing, gpNameMult=1.5, threshold="conf")

    fig = rhodium.SAResult.plot_sobol(sar.pyo; radSc=radSc, scaling=scaling, widthSc=widthSc,
                               STthick=STthick, varNameMult=varNameMult,
                               colors=colors, groups=groups, gpNameMult=gpNameMult,
                               threshold=threshold)
    return fig
end

function oat(m::Model, response; policy=Dict(), nsamples::Int=100, kwargs...)
    fig = rhodium.oat(m.pyo, response; policy=policy, nsamples=nsamples, kwargs...)
    return fig
end


text:
"""
A simple factor summary.
"""
mutable struct FactorSummary
  id::Int64
  label::String
  createdTimestamp::String
  lastUpdatedTimestamp::String
end


text:
# Classifies MNIST digits with a convolutional network.
# Writes out saved model to the file "mnist_conv.bson".
# Demonstrates basic model construction, training, saving,
# conditional early-exit, and learning rate scheduling.
#
# This model, while simple, should hit around 99% test
# accuracy after training for approximately 20 epochs.

using Flux, Flux.Data.MNIST, Statistics
using Flux: onehotbatch, onecold, logitcrossentropy
using Base.Iterators: partition
using Printf, BSON
using Parameters: @with_kw
using CUDAapi
if has_cuda()
    @info "CUDA is on"
    import CuArrays
    CuArrays.allowscalar(false)
end

@with_kw mutable struct Args
    lr::Float64 = 3e-3
    epochs::Int = 20
    batch_size = 128
    savepath::String = "./"
end

# Bundle images together with labels and group into minibatchess
function make_minibatch(X, Y, idxs)
    X_batch = Array{Float32}(undef, size(X[1])..., 1, length(idxs))
    for i in 1:length(idxs)
        X_batch[:, :, :, i] = Float32.(X[idxs[i]])
    end
    Y_batch = onehotbatch(Y[idxs], 0:9)
    return (X_batch, Y_batch)
end

function get_processed_data(args)
    # Load labels and images from Flux.Data.MNIST
    train_labels = MNIST.labels()
    train_imgs = MNIST.images()
    mb_idxs = partition(1:length(train_imgs), args.batch_size)
    train_set = [make_minibatch(train_imgs, train_labels, i) for i in mb_idxs]

    # Prepare test set as one giant minibatch:
    test_imgs = MNIST.images(:test)
    test_labels = MNIST.labels(:test)
    test_set = make_minibatch(test_imgs, test_labels, 1:length(test_imgs))

    return train_set, test_set

end

# Build model
function build_model(args; imgsize = (28,28,1), nclasses = 10)
    cnn_output_size = Int.(floor.([imgsize[1]/8,imgsize[2]/8,32]))

    return Chain(
    # First convolution, operating upon a 28x28 image
    Conv((3, 3), imgsize[3]=>16, pad=(1,1), relu),
    MaxPool((2,2)),

    # Second convolution, operating upon a 14x14 image
    Conv((3, 3), 16=>32, pad=(1,1), relu),
    MaxPool((2,2)),

    # Third convolution, operating upon a 7x7 image
    Conv((3, 3), 32=>32, pad=(1,1), relu),
    MaxPool((2,2)),

    # Reshape 3d tensor into a 2d one using `Flux.flatten`, at this point it should be (3, 3, 32, N)
    flatten,
    Dense(prod(cnn_output_size), 10))
end

# We augment `x` a little bit here, adding in random noise.
augment(x) = x .+ gpu(0.1f0*randn(eltype(x), size(x)))

# Returns a vector of all parameters used in model
paramvec(m) = vcat(map(p->reshape(p, :), params(m))...)

# Function to check if any element is NaN or not
anynan(x) = any(isnan.(x))

accuracy(x, y, model) = mean(onecold(cpu(model(x))) .== onecold(cpu(y)))

function train(; kws...)
    args = Args(; kws...)

    @info("Loading data set")
    train_set, test_set = get_processed_data(args)

    # Define our model.  We will use a simple convolutional architecture with
    # three iterations of Conv -> ReLU -> MaxPool, followed by a final Dense layer.
    @info("Building model...")
    model = build_model(args)

    # Load model and datasets onto GPU, if enabled
    train_set = gpu.(train_set)
    test_set = gpu.(test_set)
    model = gpu(model)

    # Make sure our model is nicely precompiled before starting our training loop
    model(train_set[1][1])

    # `loss()` calculates the crossentropy loss between our prediction `y_hat`
    # (calculated from `model(x)`) and the ground truth `y`.  We augment the data
    # a bit, adding gaussian random noise to our image to make it more robust.
    function loss(x, y)
        x̂ = augment(x)
        ŷ = model(x̂)
        return logitcrossentropy(ŷ, y)
    end

    # Train our model with the given training set using the ADAM optimizer and
    # printing out performance against the test set as we go.
    opt = ADAM(args.lr)

    @info("Beginning training loop...")
    best_acc = 0.0
    last_improvement = 0
    for epoch_idx in 1:args.epochs
        # Train for a single epoch
        Flux.train!(loss, params(model), train_set, opt)

        # Terminate on NaN
        if anynan(paramvec(model))
            @error "NaN params"
            break
        end

        # Calculate accuracy:
        acc = accuracy(test_set..., model)

        @info(@sprintf("[%d]: Test accuracy: %.4f", epoch_idx, acc))
        # If our accuracy is good enough, quit out.
        if acc >= 0.999
            @info(" -> Early-exiting: We reached our target accuracy of 99.9%")
            break
        end

        # If this is the best accuracy we've seen so far, save the model out
        if acc >= best_acc
            @info(" -> New best accuracy! Saving model out to mnist_conv.bson")
            BSON.@save joinpath(args.savepath, "mnist_conv.bson") params=cpu.(params(model)) epoch_idx acc
            best_acc = acc
            last_improvement = epoch_idx
        end

        # If we haven't seen improvement in 5 epochs, drop our learning rate:
        if epoch_idx - last_improvement >= 5 && opt.eta > 1e-6
            opt.eta /= 10.0
            @warn(" -> Haven't improved in a while, dropping learning rate to $(opt.eta)!")

            # After dropping learning rate, give it a few epochs to improve
            last_improvement = epoch_idx
        end

        if epoch_idx - last_improvement >= 10
            @warn(" -> We're calling this converged.")
            break
        end
    end
end

# Testing the model, from saved model
function test(; kws...)
    args = Args(; kws...)

    # Loading the test data
    _,test_set = get_processed_data(args)

    # Re-constructing the model with random initial weights
    model = build_model(args)

    # Loading the saved parameters
    BSON.@load joinpath(args.savepath, "mnist_conv.bson") params

    # Loading parameters onto the model
    Flux.loadparams!(model, params)

    test_set = gpu.(test_set)
    model = gpu(model)
    @show accuracy(test_set...,model)
end

cd(@__DIR__)
train()
test()

max_stars_repo_name: Moelf/model-zoo

max_issues_repo_name: Moelf/model-zoo

max_forks_repo_name: Moelf/model-zoo


text:
using Test
const cf = ChemistryFeaturization

@testset "binning" begin
    # test the which_bin function, can use a fake feature for numerical cases
    bins = [0,2,4]
    block_feat = AtomFeat(:Block, ["s", "p", "d", "f"])
    dummy_num_feat = AtomFeat(:Dummy, false, 2, false, [0,2,4])

    # test onehot_bins function
    @test cf.onehot_bins(block_feat, "s")==[1.0, 0., 0., 0.]
    @test cf.onehot_bins(dummy_num_feat, 1)==[1.0, 0.0]

    # and onecold_bins
    @test cf.onecold_bins(block_feat, [1,0,0,0])=="s"
    @test cf.onecold_bins(dummy_num_feat, [1, 0])==(0,2)

    # get_logspaced_vec
    @test cf.get_logspaced_vec(true, 3)==[true, true, true]
    @test cf.get_logspaced_vec(false, 3)==[false, false, false]
    @test cf.get_logspaced_vec([true,false,true], 3)==[true,false,true]
    @test cf.get_logspaced_vec([true,false,true], 2)==[true,false]
end

@testset "AtomFeat" begin
    # test that both constructors for categorical variables give the same result
    cat1 = AtomFeat(:feat, true, 3, false, ['a','b','c'])
    cat2 = AtomFeat(:feat, ['a','b','c'])
    for name in propertynames(cat1)
        @test getproperty(cat1,name)==getproperty(cat2,name)
    end

    # and now with categorical features that are numbers
    cat3 = AtomFeat(:feat, true, 3, 1, 3)
    cat3b = AtomFeat(:feat, true, 3, 1.0, 3)
    cat4 = AtomFeat(:feat, true, 3, false, [1,2,3])
    cat5 = AtomFeat(:feat, [1,2,3])

    for name in [:name, :categorical, :num_bins, :logspaced]
        results = getproperty.([cat3, cat3b, cat4, cat5], name)
        @test all(y->y==results[1], results)
    end

    @test cat3.vals==cat4.vals==cat5.vals==[1,2,3]
    @test cat3b.vals==[1.0,2.0,3.0]

    # and for numerical features, first linearly spaced...
    lin1 = AtomFeat(:feat, false, 3, false, [1.0,2,3,4])
    lin2 = AtomFeat(:feat, false, 3, 1.0, 4)
    @test lin1.vals==lin2.vals

    # now log...
    log1 = AtomFeat(:feat, false, 3, true, Float32.([0.1, 1, 10, 100]))
    log2 = AtomFeat(:feat, false, 3, 0.1, 100, true)
    @test log1.vals==log2.vals

    # check that arbitrary spacing does work, though I don't see a need for it right now...
    arb = AtomFeat(:feat, false, 3, false, [1,2,4,5])
    @test arb.vals==Float32.([1,2,4,5])

    # test build_atom_feats function
    X, MP, block = build_atom_feats([:X, Symbol("Melting point"), :Block]; logspaced=[false, true])
    @test X.name==:X
    @test X.categorical==false
    @test X.logspaced==false
    @test X.vals[1]==Float32(0.7)
    @test X.vals[end]==Float32(3.98)

    @test MP.categorical==false
    @test MP.logspaced==true

    @test block.categorical==true
    @test block.logspaced==false
    @test block.vals==["s", "p", "d", "f"]
end

@testset "encode/decode" begin
    # make_feature_vectors...pick some representative properties
    feature_names = [:X, Symbol("Atomic mass"), :Block]
    vecs, features = make_feature_vectors(feature_names)
    # test that they are what we think they should be for a couple elements
    H_feat = decode_feature_vector(vecs["H"], features)
    @test H_feat[:X][1] <= 2.2 <= H_feat[:X][2]
    @test H_feat[:Block] == "s"
    @test H_feat[Symbol("Atomic mass")][1] <= Float32(1.00794) <= H_feat[Symbol("Atomic mass")][2]
    Si_feat = decode_feature_vector(vecs["Si"], features)
    @test Si_feat[:X][1] <= 1.9 <= Si_feat[:X][2]
    @test Si_feat[:Block] == "p"
    @test Si_feat[Symbol("Atomic mass")][1] <= 28.0855 <= Si_feat[Symbol("Atomic mass")][2]

    # chunk_vec
    vec = [1,0,0,1,0]
    @test_throws AssertionError cf.chunk_vec(vec, [3,3])
    @test cf.chunk_vec(vec, [3,2])==[[1,0,0],[1,0]]

    # vec_valid
    @test cf.vec_valid(vec, [3,2])
    @test !cf.vec_valid(vec, [3,3])
    @test !cf.vec_valid([1,0,1,1,0], [3,2])
end


text:
using Decimals
using Compat.Test

@testset "Conversions" begin

@testset "String/Number to Decimal" begin
    @testset "Direct" begin
        @test parse(Decimal, "0.01") == Decimal(0.01) == Decimal(0, 1, -2)
        @test parse(Decimal, ".001") == Decimal(.001) == Decimal(0, 1, -3)
        @test parse(Decimal, "15.23") == Decimal(15.23) == Decimal(0, 1523, -2)
        @test parse(Decimal, "543") == Decimal(543) == Decimal(0, 543, 0)
        @test parse(Decimal, "-345") == Decimal(-345) == Decimal(1, 345, 0)
        @test parse(Decimal, "000123") == Decimal(000123) == Decimal(0, 123, 0)
        @test parse(Decimal, "-00032") == Decimal(-00032) == Decimal(1, 32, 0)
        @test parse(Decimal, "200100") == Decimal(200100) == Decimal(0, 2001, 2)
        @test parse(Decimal, "-.123") == Decimal(-.123) == Decimal(1, 123, -3)
        @test parse(Decimal, "1.23000") == Decimal(1.23000) == Decimal(0, 123, -2)
        @test parse(Decimal, "4734.612") == Decimal(4734.612) == Decimal(0, 4734612, -3)
        @test parse(Decimal, "541724.2") == Decimal(541724.2) == Decimal(0,5417242,-1)
        @test parse(Decimal, "2.5e6") == Decimal(2.5e6) == Decimal(0, 25, 5)
        @test parse(Decimal, "2.385350e8") == Decimal(2.385350e8) == Decimal(0, 238535, 3)
        @test parse(Decimal, "12.3e-4") == Decimal(12.3e-4) == Decimal(0, 123, -5)

        @test parse(Decimal, "-12.3e4") == Decimal(-12.3e4) == Decimal(1, 123, 3)

        @test parse(Decimal, "-12.3e-4") == Decimal(-12.3e-4) == Decimal(1, 123, -5)

        @test parse(Decimal, "0.1234567891") == Decimal(0.1234567891) == Decimal(0,1234567891, -10)
        @test parse(Decimal, "0.12345678912") == Decimal(0.12345678912) == Decimal(0,12345678912, -11)
    end

    @testset "Using `decimal`" begin
        @test decimal("1.0") == Decimal(0, 1, 0)
        @test decimal(8.1) == Decimal(0, 81, -1)
        @test decimal.(Float64.(d)) == d
    end
end

@testset "Array{<:Number} to Array{Decimal}" begin
    @test Decimal.([0.1 0.2 0.3]) == [Decimal(0.1) Decimal(0.2) Decimal(0.3)]
end

@testset "Decimal to String" begin
    @test string(Decimal(0, 1, -2)) == "0.01"
    @test string(Decimal(0, 1, -3)) == "0.001"
    @test string(Decimal(0, 1523, -2)) == "15.23"
    @test string(Decimal(0, 543, 0)) == "543"
    @test string(Decimal(1, 345, 0)) == "-345"
    @test string(Decimal(0, 123, 0)) == "123"
    @test string(Decimal(1, 32, 0)) == "-32"
    @test string(Decimal(0, 2001, 2)) == "200100"
    @test string(Decimal(1, 123, -3)) == "-0.123"
    @test string(Decimal(0, 123, -2)) == "1.23"
end

@testset "Decimal to Number" begin
    @test Float32(Decimal(0, 1, -2)) == 0.01f0
    @test Float64(Decimal(0, 1, -3)) == 0.001
    @test Float64(Decimal(0, 1523, -2)) == 15.23
    @test UInt(Decimal(0, 543, 0)) == 543
    @test Int(Decimal(1, 345, 0)) == -345
    @test Int32(Decimal(0, 123, 0)) == 123
    @test Int8(Decimal(1, 32, 0)) == -32
    @test BigInt(Decimal(0, 2001, 2)) == 200100
    @test BigFloat(Decimal(1, 123, -3)) == big"-0.123"
    @test Float64(Decimal(0, 123, -2)) == 1.23
    @test number(Decimal(0, 1, -2)) == 0.01
    @test number(Decimal(0, 1, -3)) == 0.001
    @test number(Decimal(0, 1523, -2)) == 15.23
    @test number(Decimal(0, 543, 0)) == 543
    @test number(Decimal(1, 345, 0)) == -345
    @test number(Decimal(0, 123, 0)) == 123
    @test number(Decimal(1, 32, 0)) == -32
    @test number(Decimal(0, 2001, 2)) == 200100
    @test number(Decimal(1, 123, -3)) == -0.123
    @test number(Decimal(0, 123, -2)) == 1.23
    @test string(Float64(Decimal(0, 543, 0))) == "543.0"
    @test string(number(Decimal(0, 543, 0))) == "543"
    @test string(number(Decimal(0, 543, -1))) == "54.3"
end

end