text:
#holds the parameters for forming the weights
#not meant to hold any meaningful data- just parameters
struct WeightSpec
  preconditioner::Function
  crosssectionweights::Function
  weightprefix::Symbol

  gross::Float64
  net::Float64
  long::Float64
  short::Float64

  specname::Symbol
end

#const WeightSpec = NamedTuple

#helper functions
@inline computelong(gross::Real, net::Real) = (gross+net)/2
@inline computeshort(gross::Real, net::Real) = (gross-net)/2
@inline maxx0(x::Real) = max(x, 0.0)
@inline minx0(x::Real) = min(x, 0.0)

function assignweights!(panel::DataFrame, ξ::AbstractCharacteristic)

  (ξ.X[:Fw] ∈ names(panel)) && error("$(ξ.X[:Fw]) already exists in dataframe")
  panel[!,ξ.X[:Fw]] = Vector{MFloat64}(undef, size(panel,1))
  @info "assigning weight for $(ξ.X[:Fw])"

  spanels::GroupedDataFrame = groupby(panel, :date)

  ####### NOTE can delete this after testing
  thresholds=PARAM[:sorteddefaultthresholds]
  Fw = ξ.X[:Fw]
  Fξ = ξ.Fξ
  ws = ξ.ws
  ######

  #=Threads.@threads =#for i ∈ 1:length(spanels)
    spanel = spanels[i]

    ####WARNING FOR TESTING ONLY!!!!
    #=ξv::Vector{Float64} = spanel[:,Fξ]#ws.preconditioner(spanel[:,Fξ])
    ranks::Vector{Float64} = ξv  |> quantiles

    Fwcol = zeros(length(ranks))
    Fwcol[ranks .≥ thresholds.upper] .=  1.0
    Fwcol[ranks .< thresholds.lower] .=  -1.0
    #Fwcol .= ((ranks .≥ thresholds.upper) .- (ranks .< thresholds.lower))
    if !(minimum(Fwcol) ≈ -1.0 * maximum(Fwcol))
      testpanel = spanel |> DataFrame
      testpanel.raw = ξv
      testpanel.unconditioned = (ranks .≥ thresholds.upper) .- (ranks .< thresholds.lower)
      testpanel.ranks = ranks
      testpanel |> CSV.write("$(PARAM[:testpath])\\error_subpanel_$(panel[1,:date])_$(Fw).csv")

      error("Ranking Failed!!
        for $Fξ -> $Fw, $i, $(spanel.date[1])
        sort category: $((ranks .≥ thresholds.upper) .- (ranks .< thresholds.lower))
        ranks: \$ranks
        spanel[:,Fξ]: \$(spanel[:,Fξ])
        typeof:$(typeof((ranks .≥ thresholds.upper) .- (ranks .< thresholds.lower)))
        $(spanel[1, Fw] |> typeof)")
    end=#
    ###############################
    ξ.ws.crosssectionweights(spanels[i], ξ.ws, ξ.Fξ, ξ.X[:Fw])

    ####sanity checks
    if ξ.ws.specname === :fixedweight #this is a long only strategy
      @assert 1.0 ≈ sum(spanels[i][!, ξ.X[:Fw]]) ≈ sum(abs.(spanels[i][!, ξ.X[:Fw]]))
      continue
    end
    @assert 1.0 ≈ 1.0 + sum(spanels[i][!, ξ.X[:Fw]])
    @assert 2.0 ≈ sum(abs.(spanels[i][!, ξ.X[:Fw]]))
  end

  #we want the weight fields, and the lagged product of the weights with the returns
  #NOTE: the accounting identity requires contemporaneous return
  #the characteristic however can be defined as either lag, contemporaneous or soemthing else
  # its all definitional
  lagwithin2!(panel, [ξ.X[:Fw], :ret], :permno, date=:date)
  panel[!, ξ.X[:FRLw]] = panel[!,ξ.X[:FLw]] .* (1 .+ panel[!, :ret])

  lagwithin2!(panel, [ξ.X[:FLw]], :permno, date=:date)
  panel[!, ξ.X[:FRLLw]] = panel[!,ξ.X[:FLLw]] .* (1 .+ panel[!, :Lret])
  lagwithin2!(panel, [ξ.X[:FLLw]], :permno, date=:date)

  #compute the forward return weights
  transform!(groupby(panel, :date),
    [ξ.X[:FRLw], ξ.X[:Fw]] => ((RLw,w) -> sum(skipmissing(RLw)) .* w) => ξ.X[:FRtw])

#  @assert ismissing.(panel[!, ξ.X[:FRLw]]) .== (ismissing.(panel[!,ξ.X[:FLw]]) .& ismissing.(panel[!, :ret])) |> all

  #check this was done correctly
  for spanel ∈ groupby(view(panel, completecases(panel, ξ.X[:FLw]), :), :date)
    #@assert (!ismissing).(spanel[:, ξ.X[:FRtw]]) |> all

    @assert (spanel[:, ξ.X[:FRtw]] ≈ (
      sum(spanel[:,ξ.X[:FLw]] .* (1.0 .+ spanel.ret)) .* spanel[:, ξ.X[:Fw]]))
  end

    return nothing
end

###########################################CARA Specs
#returns a vector of weights according to the CARA methodology
function caraweights!(panel::AbstractDataFrame, ws::WeightSpec, Fξ::Symbol, Fw::Symbol;
  Fiivol::Symbol = PARAM[:portfolioiivol], minweightmagnitude=PARAM[:minweightmagnitude])

  local longraw::Float64
  local shortraw::Float64

  ξ = ws.preconditioner(panel[:,Fξ])

  panel[:, Fw] .= ξ .* panel[!, Fiivol]
  panel[:, Fw] .-= mean(panel[:, Fw])

  longraw = sum((maxx0).(panel[!, Fw]))
  shortraw = sum((minx0).(panel[!, Fw])) .* -1.0

  #correct for small values
  for (i,wₙ) ∈ enumerate(panel[!, Fw])
    if wₙ ≥ 0 && wₙ/longraw*ws.long < minweightmagnitude
      panel[i, Fw] = 0.0
    elseif wₙ < 0 && abs(wₙ)/shortraw*ws.short < minweightmagnitude
      panel[i, Fw] = 0.0
    end
  end

  longraw = sum((maxx0).(panel[!, Fw]))
  shortraw = sum((minx0).(panel[!, Fw])) .* -1.0

  if (!isfinite(longraw)) || (longraw==0.0)
    println("ERROR!")
    println(panel[!,[:date, :permno, Fiivol, Fw]])
    error("Longraw ($longraw) is not non-zero finite.")
  elseif (!isfinite(shortraw)) || (shortraw==0.0)
    println(panel[!,[:date, :permno, Fiivol, Fw]])
    error("Shortraw ($shortraw) is not non-zero finite.")
  end

  #assigns the weight by first checking if the position is long or short
  for (i,wₙ) ∈ enumerate(panel[!, Fw])
    panel[i, Fw] = ifelse(wₙ ≥ 0, wₙ/longraw*ws.long, wₙ/shortraw*ws.short)
  end

  return nothing
end

#creates a specification for cara based weights
caraweightspec(;
  preconditioner = error("preconditioning function is required"),
  crosssectionweights = caraweights!,
  weightprefix::Symbol = PARAM[:caraweightprefix],
  gross=PARAM[:caragross],
  net=PARAM[:caranet],
  specificationname=PARAM[:caraname]) = WeightSpec(
    preconditioner,
    crosssectionweights,
    weightprefix,
    gross,
    net,
    computelong(gross,net),
    computeshort(gross,net),
    specificationname)

################################Z-weight spec (Equal conditional weights)

function zweights!(panel::AbstractDataFrame, ws::WeightSpec, Fξ::Symbol, Fw::Symbol;
  minweightmagnitude=PARAM[:minweightmagnitude])

  ξ = ws.preconditioner(panel[:,Fξ])

  panel[:, Fw] .= ξ
  panel[:, Fw] .-= mean(panel[:, Fw])

  longraw = sum((maxx0).(panel[!, Fw]))
  shortraw = sum((minx0).(panel[!, Fw])) .* -1.0

  #correct for small values
  for (i,wₙ) ∈ enumerate(panel[!, Fw])
    if wₙ ≥ 0 && wₙ/longraw*ws.long < minweightmagnitude
      panel[i, Fw] = 0.0
    elseif wₙ < 0 && abs(wₙ)/shortraw*ws.short < minweightmagnitude
      panel[i, Fw] = 0.0
    end
  end

  longraw = sum((maxx0).(panel[!, Fw]))
  shortraw = sum((minx0).(panel[!, Fw])) .* -1.0

  if (!isfinite(longraw)) || (longraw==0.0)
    println("ERROR!")
    println(panel[!,[:date, :permno, Fw]])
    error("Longraw ($longraw) is not non-zero finite.")
  elseif (!isfinite(shortraw)) || (shortraw==0.0)
    println(panel[!,[:date, :permno, Fw]])
    error("Shortraw ($shortraw) is not non-zero finite.")
  end

  #assigns the weight by first checking if the position is long or short
  for (i,wₙ) ∈ enumerate(panel[!, Fw])
    panel[i, Fw] = ifelse(wₙ ≥ 0, wₙ/longraw*ws.long, wₙ/shortraw*ws.short)
  end

  #=panel[!, Fw] .= panel[!, Fξ] .- mean(panel[:, Fξ])
  panel[!, Fw] .= panel[!, Fξ] ./ sum((abs).(panel[!, Fw])) .* 2

  (sum((x->ismissing(x)||(!isfinite(x))).(panel[!,Fw])) == 0) || error(
    "INVALID values detected in weights for Fξ")=#

  @assert sum(panel[!, Fw]) + 1.0 ≈ 1.0
  @assert sum(abs.(panel[!, Fw])) ≈ 2.0

  return nothing
end

#creates a specification given equal conditional weights
zweightspec(;
  preconditioner = error("preconditioning function is required"),
  crosssectionweights = zweights!,
  weightprefix::Symbol = PARAM[:zweightprefix],
  gross=PARAM[:zgross],
  net=PARAM[:znet],
  specificationname=PARAM[:zname]) = WeightSpec(
    preconditioner,
    crosssectionweights,
    weightprefix,
    gross,
    net,
    computelong(gross,net),
    computeshort(gross,net),
    specificationname)

#####################sortspec
#Applies preconditioner to the characteristic -usually either demeaning, quantilizaiton, or both,
#The preconditioned characteristic is then multiplied by a conditioner (such as market cap)
#The conditioner weights the characteristic within a particular sort threshold
#Then the portfolios are formed within quantile thresholds, both long and short
#These are then finally normalized to add to a value (usually 1) per side
#note the conditioning vector can be overridden if customization is needed
function _sortedweights!(panel::AbstractDataFrame, ws::WeightSpec, Fξ::Symbol, Fw::Symbol;
  thresholds::NamedTuple=error("thresholds is required"),
  conditioning::AbstractVector=error("conditioning is required"),
  minweightmagnitude=PARAM[:minweightmagnitude])

  ξ::Vector{Float64} = ws.preconditioner(panel[:,Fξ])
  ranks::Vector{Float64} = ξ  |> quantiles
  terciles = ((ranks .≥ thresholds.upper) .- (ranks .< thresholds.lower))
  minimum(terciles) ≈ -1.0 * maximum(terciles) || error("
    (minimum(ξ.X[:Fw]) /≈ -1.0 * maximum(ξ.X[:Fw])): This probably
    means the sort failed to differentiate between the quantiles. Check winsorization
    and sorting thresholds to try again.")

  #conditioning is usually marketcap or a vector of ones for equal wieghting
  #see other sortwedweights! function
  panel[:, Fw] .= terciles .* conditioning

  longraw = sum((maxx0).(panel[!, Fw]))
  shortraw = sum((minx0).(panel[!, Fw])) .* -1.0

  #correct for small values
  for (i,wₙ) ∈ enumerate(panel[!, Fw])
    if wₙ ≥ 0 && wₙ/longraw*ws.long < minweightmagnitude
      panel[i, Fw] = 0.0
    elseif wₙ < 0 && abs(wₙ)/shortraw*ws.short < minweightmagnitude
      panel[i, Fw] = 0.0
    end
  end

  longraw = sum((maxx0).(panel[!, Fw]))
  shortraw = sum((minx0).(panel[!, Fw])) .* -1.0

  if (!isfinite(longraw)) || (longraw==0.0)
    #println("ERROR!")
    #println(panel[1:10,[:date, :permno, Fw]])
    error("Longraw ($longraw) is not non-zero finite for date/field $(panel[1,[:date, Fw]]).")
  elseif (!isfinite(shortraw)) || (shortraw==0.0)
    #println(panel[1:10,[:date, :permno, Fw]])
    #=testpanel = panel |> DataFrame
    testpanel.conditioning = conditioning
    testpanel.unconditioned = (ranks .≥ thresholds.upper) .- (ranks .< thresholds.lower)
    testpanel.ranks = ranks
    testpanel |> CSV.write("$(PARAM[:testpath])\\error_subpanel_$(panel[1,:date])_$(Fw).csv")=#

    error("Shortraw ($shortraw) is not non-zero finite for date/field $(panel[1,[:date, Fw]]).")
  end

  #assigns the weight by first checking if the position is long or short
  for (i,wₙ) ∈ enumerate(panel[!, Fw])
    panel[i, Fw] = ifelse(wₙ ≥ 0, wₙ/longraw*ws.long, wₙ/shortraw*ws.short)
  end

  #=panel[!, Fw] .= panel[!, Fξ] .- mean(panel[:, Fξ])
  panel[!, Fw] .= panel[!, Fξ] ./ sum((abs).(panel[!, Fw])) .* 2

  (sum((x->ismissing(x)||(!isfinite(x))).(panel[!,Fw])) == 0) || error(
    "INVALID values detected in weights for Fξ")=#

  @assert sum(panel[!, Fw]) + 1.0 ≈ 1.0
  @assert sum(abs.(panel[!, Fw])) ≈ 2.0

  return nothing
end

#override the kw arguments to get characteristic specific threshholds and conditioning
function sortedweights!(panel::AbstractDataFrame, ws::WeightSpec, Fξ::Symbol, Fw::Symbol;
  thresholds=PARAM[:sorteddefaultthresholds],
  Fconditioning::Union{Nothing, Symbol} = PARAM[:sorteddefaultconditioning],)

  #extract the conditioning vector
  conditioning =  (Fconditioning === nothing ?
    ones(size(panel, 1)) |> Vector{MFloat64} : panel[!, Fconditioning])

  return _sortedweights!(panel, ws, Fξ, Fw; thresholds, conditioning)
end

#creates a specification given equal conditional weights
#WARNING- switch this back to zweights!
sortedweightspec(;
  preconditioner = error("preconditioning function is required"),
  crosssectionweights = sortedweights!,
  weightprefix::Symbol = PARAM[:sortedweightprefix],
  gross=PARAM[:sortedgross],
  net=PARAM[:sortednet],
  specificationname=PARAM[:sortedname]) = WeightSpec(
    preconditioner,
    crosssectionweights,
    weightprefix,
    gross,
    net,
    computelong(gross,net),
    computeshort(gross,net),
    specificationname)

#####################doublesortspec
#Applies preconditioner to the characteristic -usually either demeaning, quantilizaiton, or both,
#The preconditioned characteristic is then multiplied by a conditioner (such as market cap)
#The conditioner weights the characteristic within a particular sort threshold
#Then the portfolios are formed within quantile thresholds, both long and short
#These are then finally normalized to add to a value (usually 1) per side
#note the conditioning vector can be overridden if customization is needed
function _doublesortedweights!(panel::AbstractDataFrame, ws::WeightSpec, Fξ::Symbol, Fw::Symbol;
  thresholds::NamedTuple=error("thresholds is required"),
  conditioning::AbstractVector=error("conditioning is required"),
  secondary::AbstractVector = error("secondary sort variable \"secondary\" is required"),
  minweightmagnitude=PARAM[:minweightmagnitude])

  ξ::Vector{Float64} = ws.preconditioner(panel[:,Fξ])
  primaryranks::Vector{Float64} = ξ  |> quantiles
  primaryterciles::Vector{Int} = ((primaryranks .≥ thresholds.primaryupper) .-
    (primaryranks .< thresholds.primarylower))
  minimum(primaryterciles) ≈ -1.0 * maximum(primaryterciles) || error("
    (minimum(ξ.X[:Fw]) /≈ -1.0 * maximum(ξ.X[:Fw])): This probably
    means the primary sort failed to differentiate between the quantiles. Check winsorization
    and sorting thresholds to try again.")

  secondaryranks::Vector{Float64} = ws.preconditioner(secondary) |> quantiles
  secondaryterciles::Vector{Int} = ((secondaryranks .≥ thresholds.secondaryupper) .-
    (secondaryranks .< thresholds.secondarylower))
  minimum(secondaryterciles) ≈ -1.0 * maximum(secondaryterciles) || error("
    (secondary) /≈ -1.0 * maximum(secondary)): This probably
    means the secondary sort failed to differentiate between the quantiles. Check winsorization
    and sorting thresholds to try again.")

  #conditioning is usually marketcap or a vector of ones for equal wieghting
  #see other sortwedweights! function
  w = (
    longupper = ((primaryterciles .== 1) .& (secondaryterciles .== 1)) .* conditioning,
    longlower = ((primaryterciles .== 1) .& (secondaryterciles .== -1)) .* conditioning,
    shortupper = -1 .* ((primaryterciles .== -1) .& (secondaryterciles .== 1)) .* conditioning,
    shortlower = -1 .* ((primaryterciles .== -1) .& (secondaryterciles .== -1)) .* conditioning)

  #below checks to make sure there are no collisions
  wcheck = reduce(+, w |> values)
  @assert (wcheck .+ 1.0 .≈ 1.0) .| (abs.(wcheck) .≈ conditioning) |> all
  @assert all(vcat(w.longupper, w.longlower, -1 .* w.shortupper, -1 .* w.shortlower) .≥ 0.0)
  longupperraw = sum(w.longupper)
  longlowerraw = sum(w.longlower)
  shortupperraw = sum(w.shortupper)  .* -1.0
  shortlowerraw = sum(w.shortlower)  .* -1.0

  #zero the small weights- note we divide the target leg leverage by two since we weight
  #the secondary upper and secondary lower portfolio components equally
  zerosmallweights(wₙ, raw, legleverage) = ifelse(
    abs(wₙ)/raw*legleverage<minweightmagnitude, 0.0, wₙ)
  w.longupper .= zerosmallweights.(w.longupper, longupperraw, ws.long/2)
  w.longlower .= zerosmallweights.(w.longlower, longlowerraw, ws.long/2)
  w.shortupper .= zerosmallweights.(w.shortupper, shortupperraw, ws.short/2)
  w.shortlower .= zerosmallweights.(w.shortlower, shortlowerraw, ws.short/2)

  #update weight totals
  longupperraw = sum(w.longupper)
  longlowerraw = sum(w.longlower)
  shortupperraw = sum(w.shortupper)  .* -1.0
  shortlowerraw = sum(w.shortlower)  .* -1.0

  @assert all(isfinite.(w |> values .|> sum))
  @assert !all((legtot -> abs(legtot) ≈ 1.0).(w |> values .|> sum))

  #assigns the weight by first checking if the position is long or short
  panel[:, Fw] .= 0.0
  panel[:, Fw] .+= w.longupper/longupperraw*ws.long/2
  panel[:, Fw] .+= w.longlower/longlowerraw*ws.long/2
  panel[:, Fw] .+= w.shortupper/shortupperraw*ws.short/2
  panel[:, Fw] .+= w.shortlower/shortlowerraw*ws.short/2

  #sanity check on total exposure
  @assert sum(panel[panel[!,Fw] .> 0,Fw]) ≈ 1.0
  @assert sum(panel[panel[!,Fw] .< 0,Fw]) ≈ -1.0

  return nothing
end

#override the kw arguments to get characteristic specific threshholds and conditioning
function doublesortedweights!(panel::AbstractDataFrame, ws::WeightSpec, Fξ::Symbol, Fw::Symbol;
  thresholds=PARAM[:doublesorteddefaultthresholds],
  Fconditioning::Union{Nothing, Symbol} = PARAM[:doublesorteddefaultconditioning],
  Fsecondary = PARAM[:doublesorteddefaultsecondary],)

  #extract the conditioning vector
  conditioning =  (Fconditioning === nothing ?
    ones(size(panel, 1)) |> Vector{MFloat64} : panel[!, Fconditioning])

  secondary = panel[!, Fsecondary]

  return _doublesortedweights!(panel, ws, Fξ, Fw; thresholds, conditioning, secondary)
end

#creates a specification given equal conditional weights
#WARNING- switch this back to zweights!
doublesortedweightspec(;
  preconditioner = error("preconditioning function is required"),
  crosssectionweights = doublesortedweights!,
  weightprefix::Symbol = PARAM[:doublesortedweightprefix],
  gross=PARAM[:doublesortedgross],
  net=PARAM[:doublesortednet],
  specificationname=PARAM[:doublesortedname]) = WeightSpec(
    preconditioner,
    crosssectionweights,
    weightprefix,
    gross,
    net,
    computelong(gross,net),
    computeshort(gross,net),
    specificationname)

#############fixedweightspec

function fixedweights!(panel::AbstractDataFrame, ws::WeightSpec, Fξ::Symbol, Fw::Symbol;
  minweightmagnitude=PARAM[:minweightmagnitude])

  ξ = ws.preconditioner(panel[:,Fξ])

  panel[:, Fw] .= ξ
  longraw = sum((maxx0).(panel[!, Fw]))

  #correct for small values
  for (i,wₙ) ∈ enumerate(panel[!, Fw])
    if wₙ ≥ 0 && wₙ/longraw*ws.long < minweightmagnitude
      panel[i, Fw] = 0.0
    elseif wₙ < 0 && abs(wₙ)/shortraw*ws.short < minweightmagnitude
      error("negative weights not allowed for fixed weight spec")
      #panel[i, Fw] = 0.0
    end
  end

  longraw = sum((maxx0).(panel[!, Fw]))

  if (!isfinite(longraw)) || (longraw==0.0)
    println("ERROR!")
    println(panel[!,[:date, :permno, Fw]])
    error("Longraw ($longraw) is not non-zero finite.")
  end

  #assigns the weight by first checking if the position is long or short
  for (i,wₙ) ∈ enumerate(panel[!, Fw])
    panel[i, Fw] = wₙ/longraw*ws.long
  end

  @assert sum(panel[!, Fw]) ≈ sum(abs.(panel[!, Fw])) ≈ 1.0

  return nothing
end

#creates a specification given equal conditional weights
fixedweightspec(;
  preconditioner = error("preconditioning function is required"),
  crosssectionweights = fixedweights!,
  weightprefix::Symbol = PARAM[:fixedweightprefix],
  gross=1.0,
  net=1.0,
  specificationname=PARAM[:fixedname]) = WeightSpec(
    preconditioner,
    crosssectionweights,
    weightprefix,
    gross,
    net,
    computelong(gross,net),
    computeshort(gross,net),
    specificationname)

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
# Interesting MPS TensorMap types
const MPOTensor{S} = AbstractTensorMap{S,2,2} where {S<:EuclideanSpace}
const MPSBondTensor{S} = AbstractTensorMap{S,1,1} where {S<:EuclideanSpace}
const GenericMPSTensor{S,N} = AbstractTensorMap{S,N,1} where {S<:EuclideanSpace,N} #some functions are also defined for "general mps tensors" (used in peps code)
const MPSTensor{S} = GenericMPSTensor{S,2} where {S<:EuclideanSpace} #the usual mps tensors on which we work
#const ExMPSTensor{S,N,A,G,F1,F2}=GenericMPSTensor{S,3,A,G,F1,F2} #and mps tensor with an extra excitation - utility leg

abstract type AbstractMPS end

Base.eltype(psi::AbstractMPS) = eltype(typeof(psi))

TensorKit.leftorth(psi::AbstractMPS, args...; kwargs...) =
    leftorth!(copy(psi), args...; kwargs...)
TensorKit.rightorth(psi::AbstractMPS, args...; kwargs...) =
    rightorth!(copy(psi), args...; kwargs...)

ext: jl


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct DiffDiskSettings <: SwaggerModel
    option::Any # spec type: Union{ Nothing, DiffDiskOption } # spec name: option
    placement::Any # spec type: Union{ Nothing, DiffDiskPlacement } # spec name: placement

    function DiffDiskSettings(;option=nothing, placement=nothing)
        o = new()
        validate_property(DiffDiskSettings, Symbol("option"), option)
        setfield!(o, Symbol("option"), option)
        validate_property(DiffDiskSettings, Symbol("placement"), placement)
        setfield!(o, Symbol("placement"), placement)
        o
    end
end # type DiffDiskSettings

const _property_map_DiffDiskSettings = Dict{Symbol,Symbol}(Symbol("option")=>Symbol("option"), Symbol("placement")=>Symbol("placement"))
const _property_types_DiffDiskSettings = Dict{Symbol,String}(Symbol("option")=>"DiffDiskOption", Symbol("placement")=>"DiffDiskPlacement")
Base.propertynames(::Type{ DiffDiskSettings }) = collect(keys(_property_map_DiffDiskSettings))
Swagger.property_type(::Type{ DiffDiskSettings }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_DiffDiskSettings[name]))}
Swagger.field_name(::Type{ DiffDiskSettings }, property_name::Symbol) =  _property_map_DiffDiskSettings[property_name]

function check_required(o::DiffDiskSettings)
    true
end

function validate_property(::Type{ DiffDiskSettings }, name::Symbol, val)
end


text:
using Pkg
using Test

Pkg.activate("../..")

using ColonyRuntime
using ColonyRuntime.Crypto

function main(args)
  server_prvkey = args[1]
  println("server_prvkey: ", server_prvkey)

  colony_prvkey = Crypto.prvkey()
  colonyid = Crypto.id(colony_prvkey)

  println("colony prvkey: ", colony_prvkey)
  println("colonyid: ", colonyid)

  server = ColonyRuntime.ColoniesServer("localhost", 8080)

  colony = ColonyRuntime.Colony(colonyid, "my_colony")
  ColonyRuntime.addcolony(server, colony, server_prvkey)
end

main(ARGS)


text:
using Plots.PlotMeasures
using Plots

export
    plot_ϵ,
    plot_e_field

function plot_ϵ(s::Simulator; figsize=(600, 750), left_margin=-100px)
    ϵ = s.permittivity.ϵ

    return heatmap(
		axes(s.grid, 1), axes(s.grid, 2), ϵ',
		color=:algae,
		size=figsize, left_margin=left_margin, aspect_ratio=:equal
	)
end

function plot_e_field(s::Simulator; figsize=(600, 750), left_margin=-100px)
    ez = s.ez
    ϵ = s.permittivity.ϵ

    lim = maximum(abs.(ez))
    lim_ϵ = maximum(abs.(ϵ))

    p = plot(
        clim=(-lim, lim),  colorbar=false,
		size=figsize, left_margin=left_margin, aspect_ratio=:equal
    )
    p = heatmap!(
        p,
		axes(s.grid, 1), axes(s.grid, 2), ez',
		color=:coolwarm
	)
    p = contour!(
        p,
        axes(s.grid, 1), axes(s.grid, 2), lim .* ϵ' ./ lim_ϵ,
        color=:algae
    )

    return p
end


text:
# This file was generated, do not modify it. # hide
B_range = range(ensemble_model, :bagging_fraction,
                lower=0.5, upper=1.0)
K_range = range(ensemble_model, :(atom.K),
                lower=1, upper=20);


text:
@testset "elastic net" begin
    a = 0.5 * L2Penalty() + 0.3 * L1Penalty()
    b = 0.3 * L1Penalty() + 0.5 * L2Penalty()
    c = 0.3 * L1Penalty() + 0.2 * L2Penalty() + 0.3 * L1Penalty()
    @test R.is_elnet(a)
    @test R.is_elnet(b)
    @test R.getscale_l1(a) == 0.3
    @test R.getscale_l2(c) == 0.2
    @test R.getscale_l1(c) == 0.6

    @test R.getscale_l2(2*L2Penalty()) == 2
    @test R.getscale_l1(3.4*L1Penalty()) == 3.4
end


text:
using Statistics

include("common.jl")

example = """
[({(<(())[]>[[{[]{<()<>>
[(()[<>])]({[<{<<[]>>(
{([(<{}[<>[]}>{[]{[(<()>
(((({<>}<{<{<>}{[]{[]{}
[[<[([]))<([[{}[[()]]]
[{[{({}]{}}([{[{{{}}([]
{<[[]]>}<{[{[{[]{()[[[]
[<(<(<(<{}))><([]([]()
<{([([[(<>()){}]>(<<{{
<{([{{}}[<[[[<>{}]]]>[]]
"""

abstract type Status end
struct Complete <: Status end
struct Incomplete <: Status
    stack::Vector{Char}
end
struct Corrupted <: Status
    char::Char
end

isopen(char::Char) = char in ('(', '[', '{', '<')

function toopen(char::Char)
    return if char == ')'
        '('
    elseif char == ']'
        '['
    elseif char == '}'
        '{'
    elseif char == '>'
        '<'
    else
        throw(ArgumentError("Unsupported char $char"))
    end
end

function get_status(line::AbstractString)
    stack = Char[]
    for char in line
        if isopen(char)
            push!(stack, char)
            continue
        else
            # The close char should match what we pop off the stack.
            last = pop!(stack)
            last == toopen(char) || return Corrupted(char)
        end
    end

    return isempty(stack) ? Complete() : Incomplete(stack)
end

function score(s::Corrupted)
    return if s.char == ')'
        3
    elseif s.char == ']'
        57
    elseif s.char == '}'
        1197
    elseif s.char == '>'
        25137
    else
        throw(ArgumentError("Unsupported char $char"))
    end
end

function get_ans(input::AbstractString)
    total = 0
    for line in split(input, '\n')
        isempty(line) && continue
        status = get_status(line)
        isa(status, Corrupted) || continue
        total += score(status)
    end
    return total
end

@assert get_ans(example) == 26397
get_ans(load_input(10))

# Part 2

function score(s::Incomplete)
    total = 0
    for char in Iterators.reverse(s.stack)
        # Unwind the stack.
        total *= 5
        total += if char == '('
            1
        elseif char == '['
            2
        elseif char == '{'
            3
        elseif char == '<'
            4
        end
    end
    return total
end

function get_ans2(input::AbstractString)
    scores = Int64[]
    for line in split(input, '\n')
        isempty(line) && continue
        status = get_status(line)
        isa(status, Incomplete) || continue
        push!(scores, score(status))
    end
    @assert length(scores) % 2 == 1
    return Int64(median!(scores))
end

@assert get_ans2(example) == 288957
get_ans2(load_input(10))


text:
abstract type Algorithm end
abstract type ObjectivePenalty end

has_regularizer(algorithm::Algorithm) = false
get_regularizer(algorithm::Algorithm) = error("Algorithm does not have regularizer.")

mutable struct Learner{A <: Algorithm}
    train_data::DataSet
    test_data::Union{Nothing, DataSet}
    validation_data::Union{Nothing, DataSet}
    algorithm::A
    param_lb::Vector{Float64}
    param_ub::Vector{Float64}
    offset_lb::Float64
    offset_ub::Float64
    objective_penalties::Vector{ObjectivePenalty}
end
function Learner(train_data::DataSet, algorithm::A; test_data=nothing, validation_data=nothing, objective_penalties=ObjectivePenalty[], bounding_diameter=10.0) where {A <: Algorithm}
    max_abs_val = _max_abs_val(train_data)
    param_lb = -bounding_diameter * max_abs_val
    param_ub = bounding_diameter * max_abs_val
    max_bid_price = maximum(dp.bid_price for dp in train_data.samples)
    offset_lb = -bounding_diameter * max_bid_price
    offset_ub = bounding_diameter * max_bid_price
    if has_regularizer(algorithm)
        push!(objective_penalties, get_regularizer(algorithm))
    end
    return Learner{A}(train_data, test_data, validation_data, algorithm, param_lb, param_ub, offset_lb, offset_ub, objective_penalties)
end

struct ProximalPenalty <: ObjectivePenalty
    β::Vector{Float64}
    λ::Float64
end

struct L1Regularizer <: ObjectivePenalty
    coefficient::Float64
end
L1Regularizer() = L1Regularizer(0.0)

struct L2Regularizer <: ObjectivePenalty
    coefficient::Float64
end
L2Regularizer() = L2Regularizer(0.0)

function _add_objective_penalty!(learner::Learner, objective_penalty::ObjectivePenalty)
    push!(learner.objective_penalties, objective_penalty)
    return nothing
end

struct LinearModel{T <: Union{JuMP.VariableRef, Float64}}
    weights::Vector{T}
    offset::T
end
LinearModel(weights::Vector{Float64}) = LinearModel(weights, 0.0)

evaluate(model::LinearModel, x::Vector{Float64}) = LinearAlgebra.dot(model.weights, x) + model.offset

struct Evaluation
    reward::Float64
    proportion_sold::Float64
    reward_bound::Union{Nothing, Float64}
end

_evaluate(::Nothing, model::LinearModel) = nothing
function _evaluate(dataset::DataSet, model::LinearModel, reward_bound=nothing)
    reward = _compute_expected_reward(dataset, model)
    proportion_sold = _proportion_sold(dataset, model)
    Evaluation(reward, proportion_sold, reward_bound)
end

struct IterationResult
    is_eval::Evaluation
    os_eval::Union{Nothing,Evaluation}
    batch_eval::Union{Nothing,Evaluation}
    meta
end
function IterationResult(is_eval::Evaluation, os_eval::Union{Nothing,Evaluation}, batch_eval::Union{Nothing,Evaluation})
    return IterationResult(is_eval, os_eval, batch_eval, nothing)
end

mutable struct TrainingResults
    model::LinearModel
    final_is_eval::Evaluation
    final_os_eval::Union{Nothing,Evaluation}
    final_validation_eval::Union{Nothing,Evaluation}
    iteration_evals::Vector{IterationResult}
end
function TrainingResults(learner::Learner, model::LinearModel; reward_bound=nothing, iteration_results=IterationResult[])
    final_is_eval = _evaluate(learner.train_data, model, reward_bound)
    final_os_eval = _evaluate(learner.test_data, model)
    final_validation_eval = _evaluate(learner.validation_data, model)
    return TrainingResults(model, final_is_eval, final_os_eval, final_validation_eval, iteration_results)
end

max_stars_repo_name: joehuchette/reserve-price-optimization

max_issues_repo_name: joehuchette/reserve-price-optimization

max_forks_repo_name: joehuchette/reserve-price-optimization


text:
module ParameterHandling

using Compat: only
using ChainRulesCore
using InverseFunctions: inverse
using LogExpFunctions: logit, logistic
using LinearAlgebra
using SparseArrays

export flatten,
    value_flatten, positive, bounded, fixed, deferred, orthogonal, positive_definite

include("flatten.jl")
include("parameters_base.jl")
include("parameters_meta.jl")
include("parameters_scalar.jl")
include("parameters_matrix.jl")

include("test_utils.jl")

end # module


text:
__precompile__()
##===----------------------------------------------------------------------===##
##                                   ACCELERATE.JL                            ##
##                                                                            ##
##===----------------------------------------------------------------------===##
##                                                                            ##
##  This file checks for the available frameworks and imports only the        ##
##  frameworks available; it also coordinates access to three frameworks      ##
##  through the unified Vectorize.f() syntax                                  ##
##                                                                            ##
##===----------------------------------------------------------------------===##
module Vectorize

export @vectorize

# Cross-version compatibility
if VERSION < v"0.5.0-dev" # Julia v0.4
    OS = OS_NAME
else
    OS = Sys.KERNEL
end

## FUNCTION DICTIONARY
functions = Dict()

## Function for adding to function dictionary
function addfunction(d::Dict, k, v)
    if k in keys(d)
        append!(d[k], [v])
    else
        d[k] = []
        append!(d[k], [v])
    end
end

## START OSX
## On OS X, check if Accelerate can be found - future proofing
OS == :Darwin ? (ACCELERATE = true) : (ACCELERATE = false)
if ACCELERATE == true
    if isfile("/System/Library/Frameworks/Accelerate.framework/Accelerate")
        eval(:(include("Accelerate.jl")))
    else
        warn("Unable to locate the Apple Accelerate framework")
    end
end
## END OSX

## Find Yeppp library
libyeppp_ = Libdl.find_library(["libyeppp"])
if libyeppp_ != ""
    @eval const global libyeppp = libyeppp_
else
    currdir = @__FILE__
    @static if is_windows()
        bindir = currdir[1:end-16]*"deps\\src\\yeppp\\binaries\\"
    else
        bindir = currdir[1:end-16]*"deps/src/yeppp/binaries/"
    end

    if OS == :Darwin
        @eval const global libyeppp = bindir*"macosx/x86_64/libyeppp.dylib"
    elseif OS == :Linux
        @eval const global libyeppp = bindir*"linux/x86_64/libyeppp.so"
    elseif OS == :NT
        @eval const global libyeppp = bindir*"windows\\amd64\\yeppp.dll"
    end
end

if isfile(libyeppp)
    include("Yeppp.jl") # include Yeppp if present
end

## Find VML
const global libvml = Libdl.find_library(["libmkl_rt"], ["/opt/intel/mkl/lib"])
if libvml != ""
    include("VML.jl")
end

# vectorize macro
macro vectorize(ex)
    len = length(ex.args)
    if len == 2
        f = :(Vectorize.$(ex.args[1]))
        arg = ex.args[2]
        return esc(:($f($arg)))
    elseif len == 3
        f = :(Vectorize.$(ex.args[1]))
        arg1 = ex.args[2]
        arg2 = ex.args[3]
        return esc(:($f($arg1, $arg2)))
    else
        f = :(Vectorize.$(ex.args[1]))
        arg1 = ex.args[2]
        arg2 = ex.args[3]
        arg3 = ex.args[4]
        return esc(:($f($arg1, $arg2, $arg3)))
    end
end

# Include optimized functions
include("Functions.jl")

end # module


text:
using PyPlot, MAT, Printf
using LinearAlgebra, Random, Statistics
using NNNCVXF #(https://github.com/PetersBas/NNNCVXF)
using Flux
using InvertibleNetworks #(https://github.com/slimgroup/InvertibleNetworks.jl)
using SetIntersectionProjection  #(https://github.com/slimgroup/SetIntersectionProjection.jl)

# Download data at:
# https://rslab.ut.ac.ir/documents/81960329/82034892/Hyperspectral_Change_Datasets.zip
# from https://rslab.ut.ac.ir/data ("the USA Dataset")

use_gpu = true

cd("/home/bpete23/Results/") #working directory
data_dir = "/home/bpete23/Data"

file = matopen(joinpath(data_dir,"USA_Change_Dataset.mat"));
data = zeros(Float32,307,241,154,2,1)
labels = zeros(Float32,307,241,154,2,1)

data[:,:,:,1,1] = read(file,"T1")
data[:,:,:,2,1] = read(file,"T2")

labels[:,:,66,1,1] = read(file,"Binary")
labels[:,:,66,2,1] = 1f0 .- labels[:,:,66,1,1]

#cut data so we can divide by 2 sufficiently many times
data = data[1:304,1:240,1:152,:,:]
labels = labels[1:304,1:240,1:152,:,:]

data = data[1:2:end,1:2:end,2:2:end,:,:]
labels = labels[1:2:end,1:2:end,2:2:end,:,:]

data = data[1:100,:,:,:,:]
labels = labels[1:100,:,:,:,:]

#is this normalized properly?>
#looks like incorrect, check initial vlaues
for i=1:size(data,4)
    data[:,:,:,i,1] .= data[:,:,:,i,1] .- mean(data[:,:,:,i,1]);
    #data[:,:,:,i,1] .= data[:,:,:,i,1] .- minimum(data[:,:,:,i,1]);
    data[:,:,:,i,1] .= data[:,:,:,i,1] ./ maximum(data[:,:,:,i,1]);
end

data = repeat(data,outer=(1,1,1,8,1))

n                  = size(labels)
active_channels    = [1,2]
#n_samples_per_iter = 40

#change data to a vector of examples
dataL    = Vector{Array{Float32,5}}(undef,1)
dataL[1] = data
data     = deepcopy(dataL)
dataL    = []

labelsL   = Vector{Any}(undef,1)
labelsL[1] = labels
labels     = deepcopy(labelsL)
labelsL    = []

#labels_V = Vector{Any}(undef,1)

architecture = ((0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12))
#architecture = ((0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12),(0,12))
k = 3   # kernel size
s = 1   # stride
p = 1   # padding
n_chan_in = size(data[1],4)
#(nx, ny)  = size(data[1])[1:2]
α = 0.2^2 #artificial time-step in the nonlinear telegraph equation discretization that is the neural network
#HN = H = NetworkHyperbolic(nx,ny,n_chan_in,1,architecture; α)
HN = H = NetworkHyperbolic3D(n_chan_in,architecture; α)

#set up one projector per channel to enforce constraints
#input comes as a tensor (matrix in 2D, 3D array for 3d data)
#output should be same size as input
pos_inds = findall(labels[1][:,:,33,1,1].>0)
pos_inds_select = shuffle(pos_inds)[1:20]

function Proj_bounds_card_ch1(input,pos_inds_select,max_card_pos)
  proj=deepcopy(input)
  proj[pos_inds_select] .= 1.1f0 #needs to be >1 otherwise the next line may remove the point annotation
  proj = reshape(project_cardinality!(vec(proj),round(Int,max_card_pos*prod(size(proj)[1:2]))),size(proj)[1:2])
  proj = reshape(project_bounds!(vec(proj),0.0f0,1.0f0),size(proj)[1:2])
  return proj
end

function Proj_bounds_card_ch2(input,pos_inds_select,max_card_neg)
  proj=deepcopy(input)
  proj[pos_inds_select] .= 0.0f0
  proj = reshape(project_cardinality!(vec(proj),round(Int,max_card_neg*prod(size(proj)[1:2]))),size(proj)[1:2])
  #proj = reshape(project_bounds!(vec(proj),0.0,1.0),size(x)[1:2])
  return proj
end

P    = Vector{Vector{Any}}(undef,1)
P[1] = Vector{Any}(undef,2)
P[1][1] = x -> Proj_bounds_card_ch1(x,pos_inds_select,0.45)
P[1][2] = x -> Proj_bounds_card_ch2(x,pos_inds_select,0.75)

#set up the training options
TrainOptions = TrOpts()
TrainOptions.eval_every      = 5
TrainOptions.batchsize       = 1
TrainOptions.use_gpu         = use_gpu
TrainOptions.lossf           = []
TrainOptions.lossg           = []
TrainOptions.active_channels = [1,2]
TrainOptions.flip_dims       = []
TrainOptions.permute_dims    = []

TrainOptions.alpha   = 1f0
TrainOptions.maxiter = 300
TrainOptions. opt    = Flux.Momentum(1f-4,0.9)

active_z_slice = 33

if use_gpu==true
  HN   = H |> gpu  #move network to use_gpu
  data = data|>gpu #move data to gpu
end

logs = Log() #create structure of logs of misfits

#fill unused data/labels with empty arrays
val_data      = [[]]
train_labels  = [[]]
val_labels    = [[]]
output_samples_train = [[]]
output_samples_val   = [[]]

#train the network; returns the loss, network parameters are updated inplace so HN is updated after training
logs = Train(HN,logs,TrainOptions,data,val_data,train_labels,val_labels,P,output_samples_train,output_samples_val,active_z_slice)

if isdir("point_annotations_1_class_plus_constraints")==false
  mkdir("point_annotations_1_class_plus_constraints")
end
cd("point_annotations_1_class_plus_constraints")
#plot the distance function per iteration
PlotHyperspectralLossConstrained(logs,TrainOptions.eval_every)

#plot data and results
PlotDataLabelPredictionHyperspectral(1,data,labels,HN,TrainOptions.active_channels,active_z_slice,pos_inds_select,[],"train")
cd("..")

######### Compare to point annotations for both classes + constraints ##########
######### Compare to point annotations for both classes + constraints ##########

#HN = H = NetworkHyperbolic3D(n_chan_in,architecture; α)

#set up one projector per channel to enforce constraints
#input comes as a tensor (matrix in 2D, 3D array for 3d data)
#output should be same size as input
#pos_inds = findall(labels[1][:,:,33,1,1].>0)
#pos_inds_select = shuffle(pos_inds)[1:20]
neg_inds = findall(labels[1][:,:,33,1,1].==0)
neg_inds_select = shuffle(neg_inds)[1:20]

function Proj_bounds_card_ch1(input,pos_inds_select,neg_inds_select,max_card_pos)
  proj=deepcopy(input)
  proj[pos_inds_select] .= 1.1f0
  proj[neg_inds_select] .= 0.0f0
  proj = reshape(project_cardinality!(vec(proj),round(Int,max_card_pos*prod(size(proj)[1:2]))),size(proj)[1:2])
  proj = reshape(project_bounds!(vec(proj),0.0f0,1.0f0),size(proj)[1:2])
  return proj
end

function Proj_bounds_card_ch2(input,pos_inds_select,neg_inds_select,max_card_neg)
  proj=deepcopy(input)
  proj[pos_inds_select] .= 0.0f0
  proj[neg_inds_select] .= 1.1f0
  proj = reshape(project_cardinality!(vec(proj),round(Int,max_card_neg*prod(size(proj)[1:2]))),size(proj)[1:2])
  proj = reshape(project_bounds!(vec(proj),0.0f0,1.0f0),size(input)[1:2])
  return proj
end

P    = Vector{Vector{Any}}(undef,1)
P[1] = Vector{Any}(undef,2)
P[1][1] = x -> Proj_bounds_card_ch1(x,pos_inds_select,0.45)
P[1][2] = x -> Proj_bounds_card_ch2(x,pos_inds_select,0.75)

if use_gpu==true
  HN   = H |> gpu  #move network to use_gpu
  data = data|>gpu #move data to gpu
end

logs = Log() #create structure of logs of misfits

#fill unused data/labels with empty arrays
val_data      = [[]]
train_labels  = [[]]
val_labels    = [[]]
output_samples_train = [[]]
output_samples_val   = [[]]

#train the network; returns the loss, network parameters are updated inplace so HN is updated after training
logs = Train(HN,logs,TrainOptions,data,val_data,train_labels,val_labels,P,output_samples_train,output_samples_val,active_z_slice)

if isdir("point_annotations_2_class_plus_constraints")==false
  mkdir("point_annotations_2_class_plus_constraints")
end
cd("point_annotations_2_class_plus_constraints")
#plot the distance function per iteration
PlotHyperspectralLossConstrained(logs,TrainOptions.eval_every)

#plot data and results
PlotDataLabelPredictionHyperspectral(1,data,labels,HN,TrainOptions.active_channels,active_z_slice,pos_inds_select,neg_inds_select,"train")
cd("..")

######### Compare to point annotations for both classes WITHOUT other constraints ##########
######### Compare to point annotations for both classes WITHOUT other constraints ##########

#HN = H = NetworkHyperbolic3D(n_chan_in,architecture; α)

#set up one projector per channel to enforce constraints
#input comes as a tensor (matrix in 2D, 3D array for 3d data)
#output should be same size as input

# pos_inds = findall(labels[1][:,:,33,1,1].>0)
# pos_inds_select = shuffle(pos_inds)[1:20]
# neg_inds = findall(labels[1][:,:,33,1,1].==0)
# neg_inds_select = shuffle(neg_inds)[1:20]

function Proj_bounds_ch1(input,pos_inds_select,neg_inds_select)
  proj=deepcopy(input)
  proj[pos_inds_select] .= 1.0f0
  proj[neg_inds_select] .= 0.0f0
  #proj = reshape(project_cardinality!(vec(proj),round(Int,max_card_pos*prod(size(proj)[1:2]))),size(proj)[1:2])
  proj = reshape(project_bounds!(vec(proj),0.0f0,1.0f0),size(input)[1:2])
  return proj
end

function Proj_bounds_ch2(input,pos_inds_select,neg_inds_select)
  proj=deepcopy(input)
  proj[pos_inds_select] .= 0.0f0
  proj[neg_inds_select] .= 1.0f0
  #proj = reshape(project_cardinality!(vec(proj),round(Int,max_card_neg*prod(size(proj)[1:2]))),size(proj)[1:2])
  proj = reshape(project_bounds!(vec(proj),0.0f0,1.0f0),size(input)[1:2])
  return proj
end

P    = Vector{Vector{Any}}(undef,1)
P[1] = Vector{Any}(undef,2)
P[1][1] = x -> Proj_bounds_ch1(x,pos_inds_select,neg_inds_select)
P[1][2] = x -> Proj_bounds_ch2(x,pos_inds_select,neg_inds_select)

if use_gpu==true
  HN   = H |> gpu  #move network to use_gpu
  data = data|>gpu #move data to gpu
end

logs = Log() #create structure of logs of misfits

#fill unused data/labels with empty arrays
val_data      = [[]]
train_labels  = [[]]
val_labels    = [[]]
output_samples_train = [[]]
output_samples_val   = [[]]

#train the network; returns the loss, network parameters are updated inplace so HN is updated after training
logs = Train(HN,logs,TrainOptions,data,val_data,train_labels,val_labels,P,output_samples_train,output_samples_val,active_z_slice)

if isdir("point_annotations_2_class_no_constraints")==false
  mkdir("point_annotations_2_class_no_constraints")
end
cd("point_annotations_2_class_no_constraints")
#plot the distance function per iteration
PlotHyperspectralLossConstrained(logs,TrainOptions.eval_every)

#plot data and results
PlotDataLabelPredictionHyperspectral(1,data,labels,HN,TrainOptions.active_channels,active_z_slice,pos_inds_select,neg_inds_select,"train")
cd("..")


text:
# A lot of this module is adapted from Tracker.jl and ReverseDiff.jl
# ReverseDiff.jl is not actively developed but it would be nice to move the code in this
# module to ReverseDiff at some point.

##########
## fill ##
##########

function Base.fill(
    value::TrackedReal,
    dims::Vararg{Union{Integer, AbstractUnitRange}},
)
    return track(fill, value, dims...)
end
@grad function fill(v::Real, dims...)
    return fill(value(v), dims...), function(Δ)
        size(Δ) ≢  dims && error("Dimension mismatch")
        return (sum(Δ), map(_->nothing, dims)...)
    end
end

###############
## any & all ##
###############

Base.any(f::Function, x::TrackedArray; dims=:) = any(f, value(x), dims = dims)
Base.all(f::Function, x::TrackedArray; dims=:) = all(f, value(x), dims = dims)

#########
## cat ##
#########

function combinations(xs, n)
    n < 1 && return [[]]
    cs = combinations(xs, n-1)
    [[x, c...] for x in xs, c in cs]
end

for f in [:hcat, :vcat]
    for i = 0:2, c = combinations([:AbstractArray, :TrackedArray, :Number, :TrackedReal], i)
        cnames = map(_ -> gensym(), c)
        @eval Base.$f($([:($x::$c) for (x, c) in zip(cnames, c)]...), x::Union{TrackedArray,TrackedReal}, xs::Union{AbstractArray,Number}...) = track($f, $(cnames...), x, xs...)
    end
    @eval begin
        Base.$f(xs::TrackedVector{T}...) where T = track($f, xs...)
        Base.$f(xs::TrackedMatrix{T}...) where T = track($f, xs...)
        Base.$f(x::TrackedVecOrMat{T}, xs::AbstractVecOrMat{T}...) where T = track($f, x, xs...)
        Base.$f(x1::TrackedVecOrMat{T}, x2::TrackedVecOrMat{T}, xs::AbstractVecOrMat{T}...) where T = track($f, x1, x2, xs...)
        Base.$f(x::TrackedVector{T}, xs::AbstractVector{T}...) where T = track($f, x, xs...)
        Base.$f(x1::TrackedVector{T}, x2::TrackedVector{T}, xs::AbstractVector{T}...) where T = track($f, x1, x2, xs...)

        @grad function $f(x::Real)
            $f(value(x)), (Δ) -> (Δ[1],)
        end
        @grad function $f(x1::Real, x2::Real)
            $f(value(x1), value(x2)), (Δ) -> (Δ[1], Δ[2])
        end
        @grad function $f(x1::AbstractVector, x2::Real)
            $f(value(x1), value(x2)), (Δ) -> (Δ[1:length(x1)], Δ[length(x1)+1])
        end
    end
end

@grad function vcat(xs::AbstractVecOrMat...)
    xs_value = value.(xs)
    out_value = reduce(vcat,xs_value)
    function back(Δ)
        start = 0
        Δs = map(xs) do xsi
          x = map(_ -> :, size(xsi))
          i = isempty(x) ? x : Base.tail(x)
          d = Δ[start+1:start+size(xsi,1), i...]
          start += size(xsi, 1)
          d
        end
        return (Δs...,)
    end
    return out_value, back
end

@grad function hcat(xs::AbstractVecOrMat...)
    xs_value = value.(xs)
    out_value = reduce(hcat,xs_value)
    function back(Δ)
        start = 0
        Δs = map(xs) do xsi
          d = if ndims(xsi) == 1
            Δ[:, start+1]
          else
            i = map(_ -> :, size(xsi)) |> Base.tail |> Base.tail
            Δ[:, start+1:start+size(xsi,2), i...]
          end
          start += size(xsi, 2)
          d
        end
        return (Δs...,)
    end
    return out_value, back
end

Base.cat(Xs::TrackedArray...; dims) = track(cat, Xs...; dims = dims)
@grad function cat(Xs::AbstractArray...; dims)
    Xs_value = value.(Xs)
    return cat(Xs_value...; dims = dims), Δ -> begin
        start = ntuple(i -> 0, Val(ndims(Δ)))
        Δs = map(Xs) do xs
          dim_xs = 1:ndims(xs)
          till_xs = ntuple((i -> i in dims ? (i in dim_xs ? size(xs,i) : 1) : 0), Val(ndims(Δ)))
          xs_in_Δ = ntuple(i -> till_xs[i] > 0 ? (start[i]+1:start[i]+till_xs[i]) : Colon(), Val(ndims(Δ)))
          d = reshape(Δ[xs_in_Δ...],size(xs))
          start = start .+ till_xs
          d
        end
        return (Δs...,)
    end
end

###############
## logsumexp ##
###############

logsumexp(x::TrackedArray; dims=:) = track(logsumexp, x, dims = dims)
@grad function logsumexp(x::AbstractArray; dims)
    lse = logsumexp(value(x), dims = dims)
    return lse, Δ -> (Δ .* exp.(x .- lse),)
end

############
## linalg ##
############

Base.copy(A::Adjoint{<:TrackedReal, <:TrackedVecOrMat}) = copyadjoint(parent(A))
copyadjoint(A) = copy(A')
copyadjoint(A::TrackedVecOrMat) = track(copyadjoint, A)
@grad function copyadjoint(A::AbstractVecOrMat)
    return copy(value(A)'), ∇ -> (copy(∇'),)
end

Base.:*(A::Adjoint{<:Real, <:TrackedVector{<:Real}}, B::AbstractVector{<:Real}) = dot(A, B)
Base.:*(A::Adjoint{<:Real, <:TrackedVector{<:Real}}, B::TrackedVector{<:Real}) = dot(A, B)
Base.:*(A::AbstractVector{<:Real}, B::Adjoint{<:Real, <:TrackedVector{<:Real}}) = dot(A, B)
Base.:*(A::TrackedVector{<:Real}, B::Adjoint{<:Real, <:TrackedVector{<:Real}}) = dot(A, B)

function LinearAlgebra.cholesky(A::Symmetric{<:Any, <:TrackedMatrix}; check=true)
    uplo = A.uplo == 'U' ? (:U) : (:L)
    factors, info = symm_turing_chol(parent(A), check, uplo)
    return Cholesky{eltype(factors), typeof(factors)}(factors, 'U', info)
end
function LinearAlgebra.cholesky(A::TrackedMatrix; check=true)
    factors, info = turing_chol(A, check)
    return Cholesky{eltype(factors), typeof(factors)}(factors, 'U', info)
end

function symm_turing_chol(x::TrackedArray{V,D}, check, uplo) where {V,D}
    tp = tape(x)
    x_value = value(x)
    (factors,info), back = DistributionsAD.symm_turing_chol_back(x_value, check, uplo)
    C = Cholesky{eltype(factors), typeof(factors)}(factors, 'U', info)
    out = track(C.factors, D, tp)
    record!(tp, SpecialInstruction, symm_turing_chol, (x, check, uplo), out, (back, issuccess(C)))
    return out, C.info
end
function turing_chol(x::TrackedArray{V,D}, check) where {V,D}
    tp = tape(x)
    x_value = value(x)
    (factors,info), back = DistributionsAD.turing_chol_back(x_value, check)
    C = Cholesky{eltype(factors), typeof(factors)}(factors, 'U', info)
    out = track(C.factors, D, tp)
    record!(tp, SpecialInstruction, turing_chol, (x, check), out, (back, issuccess(C)))
    return out, C.info
end

for f in (:turing_chol, :symm_turing_chol)
    @eval begin
        @noinline function ReverseDiff.special_reverse_exec!(
            instruction::SpecialInstruction{typeof($f)},
        )
            output = instruction.output
            instruction.cache[2] || throw(PosDefException(C.info))
            input = instruction.input
            input_deriv = deriv(input[1])
            P = instruction.cache[1]
            input_deriv .+= P((factors = deriv(output),))[1]
            unseed!(output)
            return nothing
        end
    end
end

@noinline function ReverseDiff.special_forward_exec!(
    instruction::SpecialInstruction{typeof(turing_chol)},
)
    output, input = instruction.output, instruction.input
    factors = turing_chol(value.(input)...)[1]
    value!(output, factors)
    return nothing
end

@noinline function ReverseDiff.special_forward_exec!(
    instruction::SpecialInstruction{typeof(symm_turing_chol)},
)
    output, input = instruction.output, instruction.input
    factors = symm_turing_chol(value.(input)...)[1]
    value!(output, factors)
    return nothing
end

##############
## getindex ##
##############

function Base.getindex(t::TrackedArray, inds::AbstractArray{<:CartesianIndex})
    tp = tape(t)
    out = TrackedArray(value(t)[inds], deriv(t)[inds], tp)
    record!(tp, SpecialInstruction, getindex, (t, inds), out)
    return out
end
function Base.getindex(t::TrackedArray, i::Int...)
    ind = LinearIndices(t)[i...]
    return TrackedReal(value(t)[i...], deriv(t)[i...], tape(t), ind, t)
end
function Base.getindex(t::TrackedArray, _inds::Union{Integer, Colon, AbstractArray{<:Integer}}...)
    inds = ntuple(Val(length(_inds))) do i
        _inds[i] isa Colon && return firstindex(t,i):lastindex(t,i)
        return _inds[i]
    end
    tp = tape(t)
    out = TrackedArray(value(t)[inds...], deriv(t)[inds...], tp)
    record!(tp, SpecialInstruction, (getindex, Val(:generic)), (t, inds), out)
    return out
end
@noinline function ReverseDiff.special_reverse_exec!(instruction::SpecialInstruction{<:Tuple{typeof(getindex), Val{:generic}}})
    input, inds = instruction.input
    output = instruction.output
    cinds = CartesianIndices(map(i -> 1:length(i), inds))
    input_deriv = deriv(input)
    output_deriv = deriv(output)
    i = 0
    for _idx in cinds
        idx = CartesianIndex(map(getindex, inds, Tuple(_idx)))
        input_deriv[idx] += output_deriv[i += 1]
    end
    unseed!(output)
    return nothing
end
@noinline function ReverseDiff.special_forward_exec!(instruction::SpecialInstruction{<:Tuple{typeof(getindex), Val{:generic}}})
    input, inds = instruction.input
    input_value = value(input)
    output_value = value(instruction.output)
    cinds = CartesianIndices(map(i -> 1:length(i), inds))
    i = 0
    for cind in cinds
        idx = CartesianIndex(map(getindex, inds, Tuple(cind)))
        output_value[i += 1] = input_value[idx]
    end
    return nothing
end

#############
## reshape ##
#############

function Base.reshape(t::TrackedArray, dims::Colon...)
	return TrackedArray(reshape(value(t), dims), reshape(deriv(t), dims), tape(t))
end

##################
## Broadcasting ##
##################

"""
    NotTracked(f::Function)

A struct that can be used to wrap around closures, structs and arrays of structs declaring that they do not contain tracked variables. This enables a more efficient broadcasting of such functions and structs when doing automatic differentiation with `ReverseDiff` producing a `TrackedArray` instead of an `Array{<:TrackedReal}`.
"""
struct NotTracked{F} <: Function
    f::F
end
(f::NotTracked{<:Union{Function, Type}})(args...; kwargs...) = f.f(args...; kwargs...)

istypeorclosure(::F) where {F} = _istypeorclosure(F)
istypeorclosure(::AbstractArray{F}) where {F} = _istypeorclosure(F)
istypeorclosure(::Base.RefValue{F}) where {F} = _istypeorclosure(F)
istypeorclosure(::AbstractArray{<:Real}) = false
istypeorclosure(::TrackedArray) = false
istypeorclosure(::AbstractArray{<:TrackedReal}) = true
istypeorclosure(::Real) = false
@generated _istypeorclosure(::Type{F}) where {F} = :($(fieldcount(F) > 0))

mayhavetracked(b) = istypeorclosure(b)
mayhavetracked(b::Type) = false
mayhavetracked(b::NotTracked) = false
mayhavetracked(b::Base.RefValue{<:NotTracked}) = false
mayhavetracked(b::Broadcasted) = mayhavetracked(b.f) || any(mayhavetracked, b.args)

struct TrackedStyle <: BroadcastStyle end

Broadcast.BroadcastStyle(::Type{<:Union{TrackedArray, TrackedReal}}) = TrackedStyle()
Broadcast.BroadcastStyle(::TrackedStyle, b::BroadcastStyle) = TrackedStyle()

# We have to re-build the original broadcast struct to get the appropriate array
# style. We need this primarily to support CuArrays' broadcasting fixes.
broadcast_rebuild(xs) = value(xs)
function broadcast_rebuild(bc::Broadcasted)
    broadcasted(bc.f, broadcast_rebuild.(bc.args)...)
end

getstyle(::Broadcasted{Style}) where {Style} = Style
remove_not_tracked(f) = f
remove_not_tracked(f::NotTracked) = f.f
remove_not_tracked(f::Base.RefValue{<:NotTracked}) = Ref(remove_not_tracked(f[]))
remove_not_tracked(f::Base.RefValue{<:NotTracked{<:AbstractArray}}) = remove_not_tracked(f[])
function remove_not_tracked(b::Broadcasted{style}) where {style}
    return Broadcasted{style}(remove_not_tracked(b.f), remove_not_tracked.(b.args), b.axes)
end

onlyrealarrays(args::Tuple) = onlyrealarray(first(args)) && onlyrealarrays(Base.tail(args))
onlyrealarrays(::Tuple{}) = true
onlyrealarray(::AbstractArray{<:Real}) = true
onlyrealarray(::AbstractArray) = false
onlyrealarray(::Any) = true

anyreals(args::Tuple) = first(args) isa Real || anyreals(Base.tail(args))
anyreals(args::Tuple{}) = false

function get_implementation(bc, f, T, args)
    outputisreal = (T <: AbstractArray{<:Real}) && (T !== Union{})
    # Each arg is either a real number, an array of untraked reals, a tracked array of reals or an array of untracked non-reals,
    # Output is real, and
    # No tracked closure or arguments, except TrackedReal and TrackedArray.
    if !mayhavetracked(bc) && outputisreal && (anyreals(args) || !onlyrealarrays(args))
        return Val(:tracker)
    # No arg is a real number and array args must be arrays of untracked reals or tracked arrays of reals,
    # Output is real, and
    # No tracked closure or arguments, except TrackedReal and TrackedArray.
    elseif !mayhavetracked(bc) && outputisreal
        return Val(:reversediff)
    # Function or any arg is possibly a tracked non-real or an array of tracked reals/non-reals,
    # Or output is not an array of reals
    else
        return Val(:fallback)
    end
end
function Base.copy(_bc::Broadcasted{TrackedStyle})
    bc = remove_not_tracked(_bc)
    flattened_bc = Broadcast.flatten(bc)
    untracked_bc = broadcast_rebuild(bc)
    flattened_untracked_bc = Broadcast.flatten(untracked_bc)
    T = Core.Compiler.return_type(copy, Tuple{typeof(untracked_bc)})
    f, args = flattened_untracked_bc.f, flattened_bc.args
    implementation = get_implementation(_bc, f, T, args)
    if implementation isa Val{:reversediff}
        return ∇broadcast(f, args...)
    elseif implementation isa Val{:tracker}
        return tracker_∇broadcast(f, args...)
    else
        style, axes = getstyle(flattened_untracked_bc), flattened_bc.axes
        return copy(Broadcasted{style, typeof(axes), typeof(f), typeof(args)}(f, args, axes))
    end
end

# https://github.com/FluxML/Flux.jl/issues/353
if VERSION < v"1.1.0-DEV.548"
    @eval Base.Broadcast begin
        function flatten(bc::Broadcasted{Style}) where {Style}
            isflat(bc) && return bc
            args = cat_nested(bc)
            let makeargs = make_makeargs(bc), f = bc.f
                newf = @inline function(args::Vararg{Any,N}) where N
                f(makeargs(args...)...)
                end
                return Broadcasted{Style}(newf, args, bc.axes)
            end
        end
        @inline function make_makeargs(makeargs, t::Tuple{<:Broadcasted,Vararg{Any}})
            bc = t[1]
            let makeargs = make_makeargs(makeargs, tail(t)), f = bc.f
                let makeargs = make_makeargs(makeargs, bc.args)
                    headargs, tailargs = make_headargs(bc.args), make_tailargs(bc.args)
                    return @inline function(args::Vararg{Any,N}) where N
                        args1 = makeargs(args...)
                        a, b = headargs(args1...), tailargs(args1...)
                        (f(a...), b...)
                    end
                end
            end
        end
    end
end

getouttype(::TrackedReal{<:Any, D}) where {D} = D
getouttype(::TrackedArray{<:Any, D}) where {D} = D
getouttype(::Any) = Union{}

deref(x) = x
deref(x::Base.RefValue) = x[]

@generated function splatcall(f, x::SVector{N}, utargs::T, ::Val{tinds}) where {N, T <: Tuple, tinds}
    args = []
    ti = 1
    uti = 1
    for i in 1:(N + length(T.types))
        if i in tinds
            push!(args, :(deref(x[$ti])))
            ti += 1
        else
            push!(args, :(deref(utargs[$uti])))
            uti += 1
        end
    end
    return quote
        $(Expr(:meta, :inline))
        $(Expr(:call, :f, args...))
    end
end

@generated function splitargs(args::T) where {T <: Tuple}
    N = length(T.types)
    RealOrArray = Union{Real, AbstractArray}
    inds = [i for i in 1:N if T.types[i] <: RealOrArray]
    indsval = :(Val{$(Expr(:tuple, [:($i) for i in inds]...))}())
    maybetracked = Expr(:tuple, [:(args[$i]) for i in inds]...)
    untracked = Expr(:tuple, [:(args[$i]) for i in 1:N if !(i in inds)]...)
    return :($indsval, $maybetracked, $untracked)
end

## A generalization of the broadcasting approach in ReverseDiff for general functions

@inline function ∇broadcast(f::F, args::Vararg{<:Any}) where {F}
    inds, targs, untracked = splitargs(args)
    N = length(targs)
    D = promote_type(getouttype.(targs)...)
    result = DiffResults.GradientResult(zero(SVector{N, D}))
    function df(x...)
        return ForwardDiff.gradient!(
            result,
            s -> splatcall(f, s, untracked, inds),
            SVector(x),
        )
    end
    results = broadcast(df, value.(targs)...)
    tp = tape(targs...)
    out_value = DiffResults.value.(results)
    eltype(out_value) == Bool && return out_value
    out = track(out_value, D, tp)
	cache = (results, df, ReverseDiff.index_bound.(targs, (out,)))
	record!(tp, SpecialInstruction, ∇broadcast, targs, out, cache)
    return out
end
@noinline function ReverseDiff.special_reverse_exec!(instruction::SpecialInstruction{typeof(∇broadcast)})
    input = instruction.input
    output = instruction.output
    output_deriv = deriv(output)
    results, _, bounds = instruction.cache
    N = length(input)
    if N == 1 || all(isequal(size(input[1])), size.(Base.tail(input)))
        add_to_deriv!(input, output_deriv, results)
    else
        add_to_deriv!(input, output_deriv, results, bounds)
    end
    unseed!(output)
    return nothing
end

@generated function add_to_deriv!(xs::T, o, r) where {T <: Tuple}
    N = length(T.types)
    return Expr(:block, [:(_add_to_deriv!(xs[$i], o, r, Val($i))) for i in 1:N]...)
end
_add_to_deriv!(_, _, _, _) = nothing
function _add_to_deriv!(x::Union{TrackedReal, TrackedArray}, out_deriv, results, ::Val{i}) where {i}
    return ReverseDiff.istracked(x) && ReverseDiff.diffresult_increment_deriv!(x, out_deriv, results, i)
end

@generated function add_to_deriv!(xs::T, o, r, bounds) where {T <: Tuple}
    N = length(T.types)
    return Expr(:block, [:(_add_to_deriv!(xs[$i], o, r, Val($i), bounds[$i])) for i in 1:N]...)
end
_add_to_deriv!(_, _, _, _, _) = nothing
function _add_to_deriv!(x::Union{TrackedReal,TrackedArray}, out_deriv, results, ::Val{i}, bound) where {i}
    return ReverseDiff.istracked(x) && ReverseDiff.diffresult_increment_deriv!(x, out_deriv, results, i, bound)
end

add_to_deriv!(d1, d2) = nothing
function add_to_deriv!(d1::Union{TrackedReal,TrackedArray}, d2)
    ReverseDiff.increment_deriv!(d1, d2)
end

@noinline function ReverseDiff.special_forward_exec!(instruction::SpecialInstruction{typeof(∇broadcast)})
    input, output = instruction.input, instruction.output
    results, df, _ = instruction.cache
    broadcast!(df, results, value.(input)...)
    output_value = value(output)
    output_value .= DiffResults.value.(results)
    return nothing
end

## Tracker style broadcasting
## Good for broadcasting real numbers or arrays of non-tracked structs

trim(x, Δ) = reshape(Δ, ntuple(i -> size(Δ, i), Val(ndims(x))))

unbroadcast(x::AbstractArray, Δ) =
  size(x) == size(Δ) ? Δ :
  length(x) == length(Δ) ? trim(x, Δ) :
    trim(x, sum(Δ, dims = ntuple(i -> size(x, i) == 1 ? i : ndims(Δ)+1, Val(ndims(Δ)))))

unbroadcast(x::Number, Δ) = sum(Δ)
unbroadcast(x::Base.RefValue, _) = nothing

dual(x, p) = x
dual(x::Real, p) = Dual(x, p)

function _deriv(f, G, ::Val{i}, args::Vararg{Any, N}) where {N, i}
    dargs = ntuple(j -> dual(args[j], i==j), Val(N))
    return f(dargs...).partials[1] * G
end
@generated function _derivs(f, G, args::Vararg{Any, N}) where {N}
    return Expr(:tuple, [:(_deriv.(f, G, Val($i), args...)) for i in 1:N]...)
end
@inline function tracker_∇broadcast(f, args::Vararg{Any, N}) where {N}
    args_values = map(value, args)
    out_value = broadcast(f, args_values...)
    tp = tape(args...)
    eltype(out_value) == Bool && return out_value
	out = track(out_value, tp)
    cache = (f,)
	record!(tp, SpecialInstruction, tracker_∇broadcast, args, out, cache)
    return out
end

@noinline function ReverseDiff.special_forward_exec!(instruction::SpecialInstruction{typeof(tracker_∇broadcast)})
    input, output = instruction.input, instruction.output
    f = instruction.cache[1]
    output_value = value(output)
    broadcast!(f, output_value, value.(input)...)
    return nothing
end

@noinline function ReverseDiff.special_reverse_exec!(instruction::SpecialInstruction{typeof(tracker_∇broadcast)})
    input = instruction.input
    output = instruction.output
    f = instruction.cache[1]
    output_deriv = deriv(output)
    N = length(input)
    Δargs = _derivs(f, output_deriv, value.(input)...)
    dxs = map(unbroadcast, input, Δargs)
    map(add_to_deriv!, input, dxs)
    unseed!(output)
    return nothing
end

## Limited ReverseDiff broadcasting
## Efficient broadcasting for specific functions, e.g. +, -

@inline _materialize(f, args) = broadcast(f, args...)

for (M, f, arity) in ReverseDiff.DiffRules.diffrules()
    isdefined(ReverseDiffX, M) || continue
    if arity == 1
        @eval @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{TrackedArray}}) = _materialize(bc.f, bc.args)
    elseif arity == 2
        @eval begin
            @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{TrackedArray,TrackedArray}}) = _materialize(bc.f, bc.args)
            @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{TrackedArray,TrackedReal}}) = _materialize(bc.f, bc.args)
            @noinline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{TrackedReal,TrackedArray}}) = _materialize(bc.f, bc.args)
        end
        for A in ReverseDiff.ARRAY_TYPES
            @eval begin
                @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{$A,TrackedArray}}) = _materialize(bc.f, bc.args)
                @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{TrackedArray, $A}}) = _materialize(bc.f, bc.args)
                @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{$A, TrackedReal}}) = _materialize(bc.f, bc.args)
                @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{TrackedReal,$A}}) = _materialize(bc.f, bc.args)
            end
        end
        for R in ReverseDiff.REAL_TYPES
            @eval begin
                @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{$R,TrackedArray}}) = _materialize(bc.f, bc.args)
                @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{TrackedArray,$R}}) = _materialize(bc.f, bc.args)
            end
        end
    end
end


text:
using FoodCachingExperiments
using Test

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


text:
module ISRLutils

# -----------------------------------------------------------------------------------
# This file contains functions which are needed by other packages developed at ISRLab
# -----------------------------------------------------------------------------------

# Convert linear index to multi-dimensional index
# ===============================================
function sub2ind(i,N)
    d = length(N);
    if d == 1
        return i;
    else
        x = i-1;
        ii = zeros(Integer,d);
        for j in d:-1:1
            A =  prod(N[1:(j-1)]);
            ii[j] = div(x,A) + 1;
            x = x % A;
        end
        return ii
    end
end

# Generate ND grid
# ================
"""
G = GenerateNDGrid(lb,ub,N)
lb, ub, N are d-dimensional vectors specifying the limits and number of points along each dimension.
G is d x prod(N) matrix with grid points.
"""
function GenerateNDGrid(lb,ub,N)
    d = length(N);

    if length(N) != length(lb) || length(N) != length(ub) || length(lb) != length(ub)
        error("Dimension mismatch.")
    end

    g = [range(lb[i],ub[i],length=N[i]) for i in 1:d];
    nmax = prod(N);
    G = zeros(d,nmax);

    for i in 1:nmax
        ii = sub2ind(i,N);
        for j = 1:d
            G[j,i] = g[j][ii[j]];
        end
    end
    return G
end

# Hypercube corners
# =================
function cornerPoints(x1,x2)
	x1 = x1[:]; x2 = x2[:]; # Vectorize
	d = length(x1);
	if d!=length(x2)
		error("dimension mismatch in x1, x2")
	end

	X = [x1 x2]; # d x 2 matrix
	p = []; # Not the best way to do it.

	ii = zeros(Int,d);
	for i in 0:(2^d-1)
		digits!(ii,i,base=2);
		xx = zeros(Float64,d);
		for j in 1:d
			xx[j] = X[j,ii[j]+1];
		end
		push!(p,xx');
	end
	return vcat(p...);
end

end # module

max_stars_repo_name: isrlab/ISRLutils

max_issues_repo_name: isrlab/ISRLutils

max_forks_repo_name: isrlab/ISRLutils


text:
using MechanismGeometries
using RigidBodyDynamics
const rbd = RigidBodyDynamics
using GeometryTypes
using StaticArrays: SVector, SDiagonal
using CoordinateTransformations: AffineMap, IdentityTransformation, LinearMap, Translation
using ColorTypes: RGBA
using MeshCat, RigidBodyDynamics, MeshCatMechanisms

function homog(::IdentityTransformation)
  eye(4)
end

function homog(m::AffineMap)
  CoordinateTransformations.Transformation(m)
end

homog(t::Translation) = homog(AffineMap(LinearMap(eye(3)),t.v))
homog(r::LinearMap,t::Translation) = homog(AffineMap(r.m,t.v))

vis = Visualizer()
urdf = joinpath(pwd(), "resources", "quadrotor.urdf")
robot = parse_urdf(Float64, urdf)
delete!(vis)
mvis = MechanismVisualizer(robot, URDFVisuals(urdf), vis)
body = bodies(robot)[end]
body_frame = default_frame(body)
setelement!(mvis,body_frame)
at = homog(Translation(0,0,-1))
settransform!(vis,homog(LinearMap(eye(3)),Translation(0,0,0)))
open(vis)
for i=1:1000
  settransform!(vis,homog(LinearMap(CoordinateTransformations.RotX(i/10)),Translation(0,0,0)))
  sleep(0.05)
end

max_stars_repo_name: umd-agrc/SimpleControlSim

max_issues_repo_name: umd-agrc/SimpleControlSim

max_forks_repo_name: umd-agrc/SimpleControlSim


text:
using MLJ
using PyPlot
using PrettyPrinting
using StableRNGs
import DataFrames: DataFrame, describe

X, y = @load_boston
sch = schema(X)
p = length(sch.names)
describe(y)  # From DataFrames

DecisionTreeRegressor = @load DecisionTreeRegressor pkg=DecisionTree

tree = machine(DecisionTreeRegressor(), X, y)
e = evaluate!(tree, resampling=Holdout(fraction_train=0.8),
              measure=[rms, rmslp1])
e

forest = EnsembleModel(model=DecisionTreeRegressor())
forest.model.n_subfeatures = 3

rng = StableRNG(5123) # for reproducibility
m = machine(forest, X, y)
r = range(forest, :n, lower=10, upper=1000)
curves = learning_curve!(m, resampling=Holdout(fraction_train=0.8, rng=rng),
                         range=r, measure=rms);

figure(figsize=(8,6))
plot(curves.parameter_values, curves.measurements)
ylabel("Root Mean Squared error", fontsize=16)
xlabel("Number of trees", fontsize=16)
xticks([10, 100, 250, 500, 750, 1000], fontsize=14)
yticks(fontsize=14)

forest.n = 150;

params(forest) |> pprint

r_sf = range(forest, :(model.n_subfeatures), lower=1, upper=12)
r_bf = range(forest, :bagging_fraction, lower=0.4, upper=1.0);

tuned_forest = TunedModel(model=forest,
                          tuning=Grid(resolution=3),
                          resampling=CV(nfolds=6, rng=StableRNG(32)),
                          ranges=[r_sf, r_bf],
                          measure=rms)
m = machine(tuned_forest, X, y)
e = evaluate!(m, resampling=Holdout(fraction_train=0.8),
              measure=[rms, rmslp1])
e

r = report(m)

figure(figsize=(8,6))

res = r.plotting

vals_sf = res.parameter_values[:, 1]
vals_bf = res.parameter_values[:, 2]

tricontourf(vals_sf, vals_bf, res.measurements)
xticks(1:3:12, fontsize=12)
xlabel("Number of sub-features", fontsize=14)
yticks(0.4:0.2:1, fontsize=12)
ylabel("Bagging fraction", fontsize=14)

ŷ = predict(m, X)
@show rms(ŷ, y)

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl

max_issues_repo_name: alan-turing-institute/MLJTutorials

max_forks_repo_name: alan-turing-institute/MLJTutorials


text:
# Try to approximate squared-exponential kernel:
using PyPlot

ntime = 1000
time = collect(linspace(0,10,ntime))

# If we want to approximate Gaussian, then we just get a cosine... that doesn't work well!

# How about sum of two?  This I solved numerically with Mathematically, and unfortunately I find
# a growing solution in each case... darn.

#nexp = 3
nexp = 4
npar = 2*nexp
abest=zeros(npar)
#abest[1:6]=[-9.,1.,10.,0.9,0.1,1.1]
abest[1:6]=[-3.,1.,4.,0.75,0.1,1.1]
# Here are best parameters so far:
#New minimum: 7063 0.0009743003566026671
#abest=[-4.081065227181365,1.0304513099118429,0.7384172905709343,1.6404128688104038,-0.6973216477808734,0.47273293933522775,5.0399695843913035,0.659463641544771]
#New minimum: 4091 0.0007528492057342187
#abest=[-3.983277629878132,1.0735443595362018,0.9310909822220776,1.609723984179903,-0.728531695566164,0.4833054951112925,4.780718343222218,0.6546148284917492]
#New minimum: 3509 0.0006919748190206947 [-4.012577986446492,1.0702556640310126,0.9246366188173811,1.6094033319207477,-0.756585292781486,0.48539249997921724,4.844526660410597,0.655092253875969]
#abest=[-4.012577986446492,1.0702556640310126,0.9246366188173811,1.6094033319207477,-0.756585292781486,0.48539249997921724,4.844526660410597,0.655092253875969]
#New minimum: 99883 0.0004297805010976177 [-4.707752923722873,1.0471511183944253,1.0707583835332415,1.5530678690438415,-1.023610702073631,0.522600717254369,5.660605242263262,0.6716080388250383]
#abest=[-4.707752923722873,1.0471511183944253,1.0707583835332415,1.5530678690438415,-1.023610702073631,0.522600717254369,5.660605242263262,0.6716080388250383]
#New minimum: 99770 0.00034643674752663026 [-5.071392864979707,1.041877623915747,1.1886472038873641,1.521389181893939,-1.1341375315197493,0.5358293183936108,6.016883192612092,0.6786038154302075]
abest=[-5.071392864979707,1.041877623915747,1.1886472038873641,1.521389181893939,-1.1341375315197493,0.5358293183936108,6.016883192612092,0.6786038154302075]
#New minimum: 999639 0.00016786630230513215 [-6.52117872187123,1.0277101206108639,1.6861024710175314,1.4296430478894617,-1.5767703086820388,0.5777664250974277,7.4118465595357375,0.7018981839031324]
#abest=[-6.52117872187123,1.0277101206108639,1.6861024710175314,1.4296430478894617,-1.5767703086820388,0.5777664250974277,7.4118465595357375,0.7018981839031324]

#New minimum: 950075 0.0001478308089395661 [-6.80168281621302,1.0282509070629977,1.8226885846422511,1.4133527852506547,-1.7003897369362933,0.5856387342941094,7.679383968507063,0.7049456900610803]

#abest[1:4]=[-9.,1.,10.1,0.9]
abest[7]=1.0-abest[1]-abest[3]-abest[5]
abest[8]=-(abest[1]*abest[2]+abest[3]*abest[4]+abest[5]*abest[6])/abest[7]
#abest[5]=1.0-abest[1]-abest[3]
#abest[6]=-(abest[1]*abest[2]+abest[3]*abest[4])/abest[5]
atrial = zeros(npar)
for k=1:npar
  atrial[k] = abest[k]
end
model = zeros(ntime)
for k=1:nexp
  model += atrial[1+(k-1)*2].*exp(-atrial[2+(k-1)*2].*time)
end
clf()
plot(time,exp(-time.^2./2))
plot(time,model)
#chibest=sum(((1.0+time+time.^2./3.0).*exp(-time)-model).^2)
chibest=sum((exp(-time.^2/2)-model).^2)
println("Initial chi-square: ",chibest)
read(STDIN,Char)
atrial = zeros(npar)
for i=1:1000000
  for j=1:npar-2
    atrial[j] = abest[j] + 0.01*randn()
    atrial[7]= 1.0-atrial[1]-atrial[3]-atrial[5]
    atrial[8]=-(atrial[1]*atrial[2]+atrial[3]*atrial[4]+atrial[5]*atrial[6])/atrial[7]
#    atrial[5]= 1.0-atrial[1]-atrial[3]
#    atrial[6]=-(atrial[1]*atrial[2]+atrial[3]*atrial[4])/atrial[5]
    model = zeros(ntime)
    for k=1:nexp
      model += atrial[1+(k-1)*2].*exp(-atrial[2+(k-1)*2].*time)
    end
    chi=sum((exp(-time.^2/2)-model).^2)
    if chi < chibest
      chibest = chi
      for k=1:npar
        abest[k] = atrial[k]
      end
      println("New minimum: ",i," ",chibest," ",abest)
      clf()
      plot(time,exp(-time.^2/2))
      plot(time,model)
      plot(time,exp(-time.^2/2)-model)
    end
  end
end

max_stars_repo_name: dfm/ess

max_issues_repo_name: dfm/ess

max_forks_repo_name: dfm/ess


text:
using StatsFuns: normcdf, normccdf, normlogcdf, normlogccdf, normlogpdf, normpdf, logsumexp

@rule Probit(:in, Marginalisation) (m_out::Union{PointMass, Bernoulli}, ) = begin

    # extract parameters
    p = mean(m_out)

    # calculate outward message
    f = (z) -> log(1 - p + (2*p - 1)*normcdf(z))

    # return message
    return ContinuousUnivariateLogPdf(f)

end

@rule Probit(:in, MomentMatching) (m_out::Union{PointMass, Bernoulli}, m_in::UnivariateNormalDistributionsFamily) = begin

    # extract parameters
    mz, vz = mean_cov(m_in)
    p = mean(m_out)
    @assert p >= zero(p) && p <= one(p) "The Probit node only accepts messages on its output with values between 0 and 1."

    # calculate auxiliary variables
    γ = mz/sqrt(1+vz)

    # calculate moments of g
    if γ > 0 && p > 0.5
        log_mom0_pz = logsumexp((log(1-p), log(2*p-1) + normlogccdf(-γ)))
    elseif γ <= 0 && p > 0.5
        log_mom0_pz = logsumexp((log(1-p), log(2*p-1) + normlogcdf(γ)))
    elseif γ > 0 && p <= 0.5
        log_mom0_pz = logsumexp((log(1-p) + normlogcdf(-γ), log(p) + normlogcdf(γ)))
    else
        log_mom0_pz = logsumexp((log(1-p) + normlogccdf(γ), log(p) + normlogcdf(γ)))
    end
    tmp = log(vz) + normlogpdf(γ) - log(1+vz)/2 - log_mom0_pz
    mom1_pz = mz + (2*p-1)*exp(tmp)
    mom2_pz = vz + mz^2 + (2*p-1)*2*mz*exp(tmp) - (2p-1)*γ*exp(log(vz) - log(1 + vz)/2 + tmp)

    # calculate parameters of posterior
    mpz = mom1_pz
    vpz = mom2_pz - mom1_pz^2
    vpz = clamp(vpz, tiny, vz)# ensure variance of marginal is not larger than the variance of the cavity distribution.

    # calculate parameters of outgoing message
    wz_out = clamp(1/vpz - 1/vz, tiny, huge) # Ensure precision isn't too small or too large
    ξz_out = mpz/vpz - mz/vz

    # return message
    return NormalWeightedMeanPrecision(ξz_out, wz_out)

end


text:
export SumOfSmallestIntervals
export SmallestInterval
export HighestDensityRegion

abstract type AbstractRankingCriterion end

"""
    struct SmallestInterval <: AbstractRankingCriterion

Type for specifying the ranking criterion to be the total width of the
one-dimensional marginalized smallest interval containing a probability mass `p` for a certain parameter.

Fields:
* `key::Symbol`: Name of the parameter that is used for the marginalized intervals.
* `p::Float64 = 0.9`: Probability mass to be enclosed in the smallest interval.

Constructors:

```julia
SmallestInterval(key::Symbol, p=0.9, bins=200)
```
"""
@with_kw struct SmallestInterval{T} <: AbstractRankingCriterion
    key::Symbol
    p::Float64 = 0.9
    bins::T = 200
end

"""
    struct SumOfSmallestIntervals <: AbstractRankingCriterion

Type for specifying the ranking criterion to be the summed width of all
one-dimensional marginalized smallest intervals containing a probability mass `p`.

Fields:
* `p::Float64 = 0.9`: Probability mass to be enclosed in the smallest intervals.
* `bins::T = 200`: Number of bins for the histograms to calculate interval widths.

Constructors:

```julia
SumOfSmallestIntervals(p=0.9, bins=200)
```
"""
@with_kw struct SumOfSmallestIntervals{T} <: AbstractRankingCriterion
    p::Float64 = 0.9
    bins::T = 200
end

"""
    struct HighestDensityRegion <: AbstractRankingCriterion

Type for specifying the ranking criterion to be the area of the two-dimensional
marginal posterior region containing a probability mass `p` for two certain parameters.

Fields:
* `keys::NTuple{2, Symbol}`: Names of the parameters that are used for the marginalized regions.
* `p::Float64 = 0.9`: Probability mass to be enclosed in the smallest interval.
* `bins::T = 200`: Number of bins for the histograms to calculate interval widths.

Constructors:

```julia
HighestDensityRegion(keys::NTuple{2, Symbol}, p=0.9, bins=200)
```
"""
@with_kw struct HighestDensityRegion{N, T} <: AbstractRankingCriterion
    keys::NTuple{N, Symbol}
    p::Float64 = 0.9
    bins::T = 200
end

"""
    apply_criterion(criterion<:AbstractRankingCriterion, samples::DensitySampleVector)

Applies the `criterion` to the `samples` and returns the corresponding value.

Available ranking criteria:

* `SumOfSmallestIntervals(p=0.9)`: summed width of all one-dimensional marginalized smallest intervals containing p=90% of the posterior probability
* `SmallestInterval(key=:C1, p=0.9)`: width of the one-dimensional marginalized smallest interval of parameter `:C1` containing p=90% of the posterior probability.
* `HighestDensityRegion(keys::NTuple{2, Symbol}, p=0.9, bins=200): volume of the highest density regione
"""
function apply_criterion(criterion::SumOfSmallestIntervals, samples::DensitySampleVector)
    return get_all_intervals(samples, criterion.p, bins=criterion.bins)
end

function apply_criterion(criterion::SmallestInterval, samples::DensitySampleVector)
    return get_volume(samples, (criterion.key, ), criterion.p, bins=criterion.bins)
end

function apply_criterion(criterion::HighestDensityRegion, samples::DensitySampleVector)
    return get_volume(samples, criterion.keys, criterion.p, bins=criterion.bins)
end

#------------------------------------------------------------------
# helper functions for calculating the intervals

# total width of 1d marginal smallest intervals
function get_interval(samples, key, p; bins=200)
    marg = BAT.get_marginal_dist(samples, key, bins=bins).result

    marghist = convert(Histogram, marg.dist)
    hist_p = BAT.get_smallest_intervals(marghist, [p])[1][1]

    lower, upper = BAT.get_interval_edges(hist_p)
    return sum(upper .- lower)
end

#TODO: export & docu
# interval edges of 1d marginal smallest intervals
function get_interval_edges(
    samples::DensitySampleVector,
    key::Union{Symbol, Real},
    p::Real;
    bins=200,
    atol=0.0)

    marg = BAT.get_marginal_dist(samples, key, bins=bins).result

    marghist = convert(Histogram, marg.dist)
    hist_p = BAT.get_smallest_intervals(marghist, [p])[1][1]

    lower, upper = BAT.get_interval_edges(hist_p, atol=atol)
    return (lower=lower, upper=upper)
end

function get_volume(samples, keys, p; bins=200)
    marghist = get_marginal_hist(samples, keys, bins=bins)

    hist_p = BAT.get_smallest_intervals(marghist, [p])[1][1]

    vol = bin_volume(hist_p)
    n = count(x->x>0, hist_p.weights)
    return vol*n
end

function bin_volume(hist::StatsBase.Histogram)
    edges = hist.edges
    vol = prod([abs(edges[i][1]-edges[i][2]) for i in 1:length(edges)])
    return vol
end

function get_all_intervals(samples, p; bins=200)
    sample_keys = BAT.all_active_names(BAT.varshape(samples))

    intervals = [get_volume(samples, (Meta.parse(k),), p, bins=bins) for k in sample_keys]
    return intervals
end

function get_marginal_hist(
    maybe_shaped_samples::BAT.DensitySampleVector,
    key::Union{NTuple{n,Integer}, NTuple{n,Union{Symbol, Expr}}} where n;
    bins = 200,
    closed::Symbol = :left,
    filter::Bool = false
)
    samples = unshaped.(maybe_shaped_samples)

    if filter
        samples = BAT.drop_low_weight_samples(samples)
    end

    idxs = BAT.asindex.(Ref(maybe_shaped_samples), key)
    s = Tuple(BAT.flatview(samples.v)[i, :] for i in idxs)

    edges = if isa(bins, Integer)
        BAT._get_edges(s, (bins,), closed)
    else
        Tuple(BAT._get_edges(s[i], bins[i], closed) for i in 1:length(bins))
    end

    return fit(Histogram, s, FrequencyWeights(samples.weight), edges, closed = closed)
end


text:
function addNoise(Img::Array{Float32,3}, nSig::Int64)

    (N, M) = size(Img)

    Img_N = Img .+ nSig * randn(N, M)

end


text:
"""
Decoding functions based on contracting a tensor-network associated with a
[`TensorNetworkCode`](@ref), see [`tn_decode`](@ref).
"""
module TNDecode

#imports
using ..TensorNetworkCodes: TensorNetworkCode
using ..TensorNetworkCodes: pauli_product, pauli_product_pow
using ..TensorNetworkCodes: num_qubits, find_pure_error
using ..TensorNetworkCodes: physical_neighbours, edge_indices
using ..TensorNetworkCodes: create_virtual_tensor, physical_tensor
using ITensors: ITensor, MPO, MPS, array, contract, hastags

#exports
export basic_contract, mps_contract, tn_decode

"""
    basic_contract(bond_dim=10) -> function

Return an exact contract function for use with [`tn_decode`](@ref) that contracts a matrix
of `ITensor` working along row-by-row.

See also: [`tn_decode`](@ref).
"""
function basic_contract()
    return function _basic_contract(imatrix)
        output = ITensor(1)
        for i in 1:size(imatrix)[1], j in 1:size(imatrix)[2]
            output = output * imatrix[i,j]
        end

        return output
    end
end

"""
    mps_contract(bond_dim=10) -> function

Return a contract function for use with [`tn_decode`](@ref) that contracts a matrix of
`ITensor`, treating rows as MPS/MPO and contracting rows element-wise, retaining a maximum
bond dimension of `bond_dim`, before contracting along the final row.

See also: [`tn_decode`](@ref).
"""
function mps_contract(bond_dim=10)
    # let trick for captured variable performance
    # https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-captured
    fn = let bond_dim = bond_dim  #
        function _mps_contract(imatrix)
            ψ = MPS(imatrix[1,:])
            for j in 2:size(imatrix)[1]-1
                Σ = MPO(imatrix[j,:])
                ψ = contract(Σ, ψ; maxdim = bond_dim)
            end

            output = imatrix[end,1] * ψ[1]
            for j in 2:size(imatrix)[2]
                output = output * imatrix[end,j] * ψ[j]
            end

            return output
        end
    end
    return fn
end

"""
    tn_decode(
        code::TensorNetworkCode, syndrome::AbstractVector{Int}, error_probability::Real;
        contract_fn=basic_contract()
    ) -> (Vector{Int}, Float64)

Return a recovery operator for the code, that is consistent with the syndrome and is from
the most-likely logical coset. Additionally, the predicted success probability, defined as
the ratio of the probability of the mostly-likely logical coset to the sum of the
probabilities of all logical cosets, is returned.

The method of tensor-network contraction defaults to `basic_contract`, which corresponds to
exact contraction. Alternatively, `mps_contract` may be specified for efficient but
approximate contraction.

See also: [`basic_contract`](@ref), [`mps_contract`](@ref)

!!! note
    This function currently assumes a depolarizing noise model and only supports codes that
    can be laid out on a square lattice.

# Examples
```jldoctest
julia> using TensorNetworkCodes.TNDecode

julia> code = rotated_surface_code(3);

julia> error = [0, 0, 0, 1, 0, 0, 3, 0, 0];  # IIIXIIZII

julia> syndrome = find_syndrome(code, error);

julia> p = 0.2;  # error_probability

julia> recovery, success_prob = tn_decode(code, syndrome, p; contract_fn=mps_contract(8))
([0, 0, 0, 3, 3, 0, 1, 3, 0], 0.8250539267483997)

julia> find_syndrome(code, recovery) == syndrome
true
```
"""
function tn_decode(
    code::TensorNetworkCode, syndrome::AbstractVector{Int}, error_probability;
    contract_fn=basic_contract()
)
    pure_error = find_pure_error(code, syndrome)

    coset_probabilities = array(
        _contract(code, pure_error, error_probability, contract_fn)
    )

    if length(coset_probabilities) != 4
        error("decoder not set up for multiple logicals yet!")
    end

    # find best coset index
    coset_probabilities = abs.(coset_probabilities)
    best_coset_idx = argmax(coset_probabilities)

    # evaluate correction
    best_logical = best_coset_idx - 1  # logical in format 0=I, 1=X, 2=Y, 3=Z
    powers = Int[best_logical in (1, 2), best_logical in (2, 3)]  # convert to bsf
    correction = pauli_product_pow(code.logicals, powers)
    correction = pauli_product.(correction, pure_error)

    # evaluate success probability
    predicted_success_rate = coset_probabilities[best_coset_idx] / sum(coset_probabilities)

    return correction, predicted_success_rate
end

"""
    _contract(
        code, pure_error::AbstractVector{Int}, error_prob::Float64, contract_fn
    )

Contracts the tensor network described by `code_graph`.  Uses any specified `contract_fn`.
"""
function _contract(code, pure_error::AbstractVector{Int}, error_prob::Float64, contract_fn)
    L = Int64(sqrt(num_qubits(code)))
    imatrix = [ITensor(1) for i in 1:L, j in 1:L]

    for i in 1:L, j in 1:L
        node =  -(L*(i-1) + j)
        tensor = create_virtual_tensor(code, node)

        neighbours = physical_neighbours(code, node)
        for neighbour in neighbours
            indices = edge_indices(code, Set([node, neighbour]))
            α = findfirst(x->hastags(x, "physical"), indices)
            index = indices[α]
            tensor = tensor *
            physical_tensor(index, error_prob, pure_error[neighbour])
        end
        imatrix[i,j] = tensor
    end

    return contract_fn(imatrix)
end

end


text:
module DsstWrapper

using JavaCall

const AbsoluteDate = @jimport org.orekit.time.AbsoluteDate
const AbstractGradientConverter = @jimport org.orekit.propagation.integration.AbstractGradientConverter
const AbstractIntegratedPropagator = @jimport org.orekit.propagation.integration.AbstractIntegratedPropagator
const AbstractJacobiansMapper = @jimport org.orekit.propagation.integration.AbstractJacobiansMapper
const AbstractPropagator = @jimport org.orekit.propagation.AbstractPropagator
const AdditionalEquations = @jimport org.orekit.propagation.integration.AdditionalEquations
const AdditionalStateProvider = @jimport org.orekit.propagation.AdditionalStateProvider
const AttitudeProvider = @jimport org.orekit.attitudes.AttitudeProvider
const BoundedPropagator = @jimport org.orekit.propagation.BoundedPropagator
const Class = @jimport java.lang.Class
const Collection = @jimport java.util.Collection
const DSSTForceModel = @jimport org.orekit.propagation.semianalytical.dsst.forces.DSSTForceModel
const DSSTGradientConverter = @jimport org.orekit.propagation.semianalytical.dsst.DSSTGradientConverter
const DSSTJacobiansMapper = @jimport org.orekit.propagation.semianalytical.dsst.DSSTJacobiansMapper
const DSSTPartialDerivativesEquations = @jimport org.orekit.propagation.semianalytical.dsst.DSSTPartialDerivativesEquations
const DSSTPropagator = @jimport org.orekit.propagation.semianalytical.dsst.DSSTPropagator
const EventDetector = @jimport org.orekit.propagation.events.EventDetector
const Field = @jimport org.hipparchus.Field
const FieldAbsoluteDate = @jimport org.orekit.time.FieldAbsoluteDate
const FieldAbstractIntegratedPropagator = @jimport org.orekit.propagation.integration.FieldAbstractIntegratedPropagator
const FieldAbstractPropagator = @jimport org.orekit.propagation.FieldAbstractPropagator
const FieldAdditionalEquations = @jimport org.orekit.propagation.integration.FieldAdditionalEquations
const FieldAdditionalStateProvider = @jimport org.orekit.propagation.FieldAdditionalStateProvider
const FieldBoundedPropagator = @jimport org.orekit.propagation.FieldBoundedPropagator
const FieldDSSTPropagator = @jimport org.orekit.propagation.semianalytical.dsst.FieldDSSTPropagator
const FieldEventDetector = @jimport org.orekit.propagation.events.FieldEventDetector
const FieldODEIntegrator = @jimport org.hipparchus.ode.FieldODEIntegrator
const FieldOrbit = @jimport org.orekit.orbits.FieldOrbit
const FieldOrekitFixedStepHandler = @jimport org.orekit.propagation.sampling.FieldOrekitFixedStepHandler
const FieldOrekitStepHandler = @jimport org.orekit.propagation.sampling.FieldOrekitStepHandler
const FieldSpacecraftState = @jimport org.orekit.propagation.FieldSpacecraftState
const Frame = @jimport org.orekit.frames.Frame
const Gradient = @jimport org.hipparchus.analysis.differentiation.Gradient
const JString = @jimport java.lang.String
const List = @jimport java.util.List
const ODEIntegrator = @jimport org.hipparchus.ode.ODEIntegrator
const Object = @jimport java.lang.Object
const Orbit = @jimport org.orekit.orbits.Orbit
const OrbitType = @jimport org.orekit.orbits.OrbitType
const OrekitFixedStepHandler = @jimport org.orekit.propagation.sampling.OrekitFixedStepHandler
const OrekitStepHandler = @jimport org.orekit.propagation.sampling.OrekitStepHandler
const ParameterDriversList = @jimport org.orekit.utils.ParameterDriversList
const PositionAngle = @jimport org.orekit.orbits.PositionAngle
const PropagationType = @jimport org.orekit.propagation.PropagationType
const RealFieldElement = @jimport org.hipparchus.RealFieldElement
const Set = @jimport java.util.Set
const SpacecraftState = @jimport org.orekit.propagation.SpacecraftState
const TimeStampedFieldPVCoordinates = @jimport org.orekit.utils.TimeStampedFieldPVCoordinates
const TimeStampedPVCoordinates = @jimport org.orekit.utils.TimeStampedPVCoordinates

export DSSTGradientConverter
export DSSTJacobiansMapper
export DSSTPartialDerivativesEquations
export DSSTPropagator
export FieldDSSTPropagator
export add_additional_equations
export add_additional_state_provider
export add_event_detector
export add_force_model
export clear_events_detectors
export compute_derivatives
export compute_mean_state
export compute_osculating_state
export equals
export get_additional_state_dimension
export get_additional_state_providers
export get_all_force_models
export get_attitude_provider
export get_basic_dimension
export get_calls
export get_class
export get_events_detectors
export get_field
export get_frame
export get_free_state_parameters
export get_generated_ephemeris
export get_initial_state
export get_managed_additional_states
export get_mapper
export get_mode
export get_mu
export get_name
export get_orbit_type
export get_parameters
export get_parameters_jacobian
export get_position_angle_type
export get_pv_coordinates
export get_satellite_revolution
export get_selected_coefficients
export get_selected_parameters
export get_state
export get_state_jacobian
export hash_code
export init
export initial_is_osculating
export is_additional_state_managed
export notify
export notify_all
export propagate
export remove_force_models
export reset_initial_state
export set_attitude_provider
export set_ephemeris_mode
export set_initial_jacobians
export set_initial_state
export set_interpolation_grid_to_fixed_number_of_points
export set_interpolation_grid_to_max_time_gap
export set_master_mode
export set_mu
export set_reset_at_end
export set_satellite_revolution
export set_selected_coefficients
export set_short_period_jacobians
export set_slave_mode
export to_string
export tolerances
export wait

include("ForcesWrapper/ForcesWrapper.jl")
include("UtilitiesWrapper/UtilitiesWrapper.jl")
include("dsst_gradient_converter.jl")
include("dsst_jacobians_mapper.jl")
include("dsst_partial_derivatives_equations.jl")
include("dsst_propagator.jl")
include("field_dsst_propagator.jl")

end


text:
using AstroTime: Period, Epoch, TimeScale
using Roots: find_zero

import AstroTime: unit
import ..Interfaces: AbstractTrajectory

export Trajectory, initial, final, state, events, times,
    LogEntry, count_id, id, epoch, detector,
    Apocenter, Pericenter, Eclipse

function transpose_aoa(input::Vector{Vector{T}}) where T
    n = length(input[1])
    output = Vector{Vector{T}}(undef, n)
    matrix = hcat(input...)
    for i = 1:n
        output[i] = matrix[i, :]
    end
    return output
end

struct Trajectory{S, F, B, T, TType, Unit, PType, NT} <: AbstractTrajectory{S, F, B, T}
    series::NT
    events::Vector{Event{S, TType}}
    data::Vector{Vector{T}}
    time::Vector{Period{Unit, PType}}
    function Trajectory(epoch::Epoch{S, TType},
                        time::Vector{Period{Unit, PType}},
                        data::Vector{Vector{T}};
                        frame::F=icrf,
                        body::B=earth,
                        names::Vector{Symbol}=Symbol[]) where {S, F, B, TType, Unit, PType, T}
        columns = transpose_aoa(data)
        if isempty(names)
            n = length(columns)
            if n == 3
                names = [:x, :y, :z]
            elseif n == 6
                names = [:x, :y, :z, :vx, :vy, :vz]
            else
                throw(ArgumentError("Explicit `names` need to provided if `data` does not have either 3 or 6 columns."))
            end
        end
        series = (; zip(names, [TimeSeries(epoch, time, c) for c in columns])...)
        events = Vector{Event{S, TType}}[]
        new{S::TimeScale, frame::AbstractFrame, body::CelestialBody,
            T, TType, Unit, PType, typeof(series)}(series, events, data, time)
    end
end

unit(tra::Trajectory) = unit(tra.time[1])

function find_events!(tra::Trajectory, detectors::Vector{<:Detector})
    for detector in detectors
        for (i, t) in enumerate(tra.time[1:end-1])
            t1 = tra.time[i+1]
            s0 = sign(detect(detector, t, tra))
            s1 = sign(detect(detector, t1, tra))
            detected, typ = detect(s0, s1, detector)
            detected || continue

            t_event = find_zero(t -> detect(detector, t * unit(tra), tra), (value(t), value(t1)))
            evt = Event(typ, epoch(tra) + t_event * unit(tra), detector)
            evt in tra.events || push!(tra.events, evt)
        end
    end
    sort!(tra.events, lt=(a, b)->isless(a.epoch, b.epoch))
end

function detect(s0, s1, detector::DiscreteDetector)
    return s1 > s0, :discrete
end

function detect(s0, s1, detector::IntervalDetector)
    s1 > s0 && return true, :start
    s1 < s0 && return true, :stop
    return false, :none
end

function (tra::Trajectory)(t)
    typ = promote_type(eltype.(values(tra.series))...)
    out = Vector{typ}(undef, length(tra.series))
    for (i, ts) in enumerate(tra.series)
        out[i] = ts(t)
    end
    return out
end

epoch(tra::Trajectory) = tra.series[1].epoch

function keplerian(tra::Trajectory, t)
    rv = tra(t)[1:6]
    return keplerian(rv[1:3], rv[4:6], grav_param(centralbody(tra)))
end

function State(tra::Trajectory, t::Period)
    rv = tra(t)[1:6]
    return State(epoch(tra) + t, rv[1:3], rv[4:6],
                 scale=timescale(tra),
                 frame=refframe(tra),
                 body=centralbody(tra))
end

function State(tra::Trajectory, ep::Epoch)
    rv = tra(t)[1:6]
    return State(ep, rv[1:3], rv[4:6],
                 scale=timescale(tra),
                 frame=refframe(tra),
                 body=centralbody(tra))
end

function KeplerianState(tra::Trajectory, t::Period)
    ele = keplerian(tra, t)
    return KeplerianState(epoch(tra) + t, ele...,
                          scale=timescale(tra),
                          frame=refframe(tra),
                          body=centralbody(tra))
end

function KeplerianState(tra::Trajectory, ep::Epoch)
    ele = keplerian(tra, ep)
    return KeplerianState(ep, ele...,
                          scale=timescale(tra),
                          frame=refframe(tra),
                          body=centralbody(tra))
end


text:
export HJBOneDim, HJBTwoDim

abstract AbstractHJB

type HJBOneDim{T1<:Real} <: AbstractHJB
    b::Function # Drift term
    σ::Function # Volatility term
    f::Function # "Gain" function
    g::Function # Terminal value
    T::T1     # Terminal time
    amin::T1  # Control minimum value
    amax::T1  # Control maximum value
    xmin::T1  # Domain minimum value
    xmax::T1  # Domain maximum value
    # bcond = (true, false) means Dirichlet condition on x=x_{min},
    #         but no condition on x=x_{max}
    bcond::Tuple{Bool,Bool}
    Dfun::Tuple{Function,Function} # Dirichlet condition on boundaries
end

type HJBTwoDim{T1<:Real} <: AbstractHJB
    b::Function      # Drift term
    σ::Function      # Volatility term
    f::Function      # "Gain" function
    g::Function      # Terminal value
    T::T1            # Terminal time
    amin::Vector{T1} # Control minimum value
    amax::Vector{T1} # Control maximum value
    xmin::Vector{T1} # Domain minimum value
    xmax::Vector{T1} # Domain maximum value
    Dbound::Function # Dirichlet condition on x-boundary
    # TODO: assert T>0, amin < amax, xmin < xmax
    # TODO: assert that the vectors are of size 2
    # TODO: make sure b and sigma return a length-2 vector
end

max_stars_repo_name: anriseth/HJBSolver

max_issues_repo_name: anriseth/HJBSolver

max_forks_repo_name: anriseth/HJBSolver