text:
module vobs_to_sqlite

import Base
import SQLite
import DataFrames
import Dates

    const parameters = Base.Dict([("ID", "ID"),
                    ("LAT", "Latitude"),
                    ("LON", "Longitude"),
                    ("NN", "CloudCover2D"),
                    ("FI", "GeopotentialHeight"),
                    ("DD", "WindDirection"),
                    ("FF", "WindSpeed"),
                    ("GG", "WindGust"),
                    ("TT", "T2m"),
                    ("RH", "RH2m"),
                    ("PS", "MSLP"),
                    ("PSS", "PSurface"),
                    ("PE", "Precip12H"),
                    ("PE1", "Precip1H"),
                    ("PE3", "Precip3H"),
                    ("PE6", "Precip6H"),
                    ("PE24", "Precip24H"),
                    ("QQ", "Q2m"),
                    ("VI", "Visibility"),
                    ("TD", "TD2m"),
                    ("TX", "T2m_MaximaPast6H"),
                    ("TM", "T2m_MinimaPast6H"),
                    ("GM", "WindSpeed_Maxima1H"),
                    ("GX", "Max_Windgust1H"),
                    ("WX", "Unknown1"),
                    ("GW", "Unknown2"),
                    ("TIME", "Time")])

    function make_sqlite(cmd_message)
        """Converts VFLD files to a single SQLite file
        Assumes that the cmd_message comes in the following format:
        ("vobs_to_sqlite", "<starttime>", "<endtime>", "<indir>", "<sqlfile>")
        """
        @info "Starting make_sqlite"
        starttime   = convert_string_to_datetime(cmd_message[2])
        endtime     = convert_string_to_datetime(cmd_message[3])
        indir       = cmd_message[4]
        sqlfile     = cmd_message[5]

        vobs_files = find_vobs_files(indir, starttime, endtime)

        vobs2sqlite(vobs_files, sqlfile)

        @info "Finished"
    end

    function convert_string_to_datetime(time)
        df = Dates.DateFormat("y-m-d-H")
        return Dates.DateTime(time, df)
    end

    function find_vobs_files(indir::String, starttime::Dates.DateTime, endtime::Dates.DateTime)

        files = readdir(indir, join=false)
        vobs_files = [x for x in files if startswith(x, "vobs")]

        vobs_files_within_range = []
        for f in vobs_files
            println("Going through file $f")
            #println(f[length(f)-9:length(f)])
            try
                dl = Dates.DateTime(f[length(f)-9:length(f)],"yyyymmddHH")
                if dl >= starttime && dl<endtime
                    append!(vobs_files_within_range, [indir*f])
                end
            catch e
                println("Error in getting vobs: ",e)
            end

        end
        return vobs_files_within_range
    end

    function vobs2sqlite(vobs_files, sqlfile::String)

        target_column_names = [k for (k,v) in parameters]

        db = SQLite.DB(sqlfile)
        make_missing_table(db)

        for f in vobs_files
            @info "Reading $f"
            i=1
            record = 1
            parameter_count = 1
            processed_header = false

            no_records = 0
            header_lines = 0
            column_names = []
            for l in eachline(f)
                if i == 1
                    no_records = parse(Int,Base.split(l)[1])
                elseif i == 2
                    header_lines = parse(Int,Base.split(l)[1])
                    column_names = ["" for k in 1:header_lines]
                elseif i > 2 && i <= header_lines::Int+2 #Parameters in file
                    column_parameter = Base.split(l)[1]
                    column_names[parameter_count] = column_parameter
                    parameter_count+=1
                elseif i > header_lines::Int+2 && i <= no_records::Int+header_lines::Int+2
                    #Adding header_lines::Int+2 to simulate resetting counter i
                    if !processed_header
                        # First we need to check if "FI" is present. In VOBS FI is always there,
                        # even though it is not listed.
                        is_FI_there = is_string_present("FI", column_names)
                        if !is_FI_there
                            column_names = append!(["FI"], column_names)
                        end
                        column_names = append!(["ID", "LAT", "LON"], column_names)
                        #println(column_names)
                        no_columns = length(column_names)
                        println("Number of columns: $no_columns")
                        processed_header = true
                        global data = zeros(Float64, no_records, no_columns)
                    end

                    dataline = parse.(Float64, Base.split(l))

                    data[record,:] = dataline

                    record+=1

                end
                i+=1
            end

            df = DataFrames.DataFrame(Base.zeros(Float64, no_records, length(target_column_names)))

            DataFrames.rename!(df, target_column_names)

            df = get_and_put_time(f, df)
            df = set_and_reorder_columns(df, column_names)

            inject_data(db, df)

        end

    end

    function is_string_present(str, str_array)
        b = filter(x->occursin(str,x), str_array)
        if length(b) == 0
            return false
        elseif length(b) > 0
            return true
        end
    end

    function get_and_put_time(f, df)
        """Inserting the Time into DataFrame"""

        try
            current_time = Dates.DateTime(f[length(f)-9:length(f)],"yyyymmddHH")
        catch e
            print("Error adding times ",e)
        end
        current_time = Int(Dates.datetime2unix(current_time))
        println("Current time $current_time")
        println(first(df))
        println(names(df))
        #colname = "TIME"
        # this only works if TIME not defined before
        nrows = DataFrames.nrow(df)
        this_time = repeat([current_time],nrows)
        df[!, "TIME"] = this_time
        #DataFrames.insertcols!(df,1,:TIME => this_time)
        #df.TIME = this_time #current_time
        #df["TIME"] = current_time #This not working!!!
        #println(names(df))
        #df[!,:TIME] = current_time #[current_time]
        #df[!,colname] = current_time #[current_time]
        #df[:"TIME"] = current_time #[current_time]
        return df
    end

    function set_and_reorder_columns(df, column_names)
        """Sets data is present columns and reorder to match SQL table"""
        k_itr = 1
        for k in column_names
            #println("column ",k)
            #df[k] = data[:,k_itr] # does not work anymore
            df[!,k] = data[:,k_itr]
            k_itr+=1
        end

        # Reordering columns
        # Order of columns gets important when inserted into sql table
        DataFrames.select!(df,[:ID, :TIME, :LAT, :LON, :FI, :NN, :DD, :FF, :GG, :TT, :RH, :PS, :PSS, :PE, :PE1, :PE3, :PE6, :PE24, :QQ, :VI, :TD, :TX, :TM, :GM, :GX, :WX, :GW])
        return df
    end

    function inject_data(db, dataTable)
        """Inject data into SQLite Table"""
       SQLite.load!(dataTable, db, "vobs")
    end

    function make_missing_table(db)
        """Makes the SQL Table if it does not exist"""

        sqliteCreateTable   = """CREATE TABLE IF NOT EXISTS vobs
                                    (ID INT DEFAULT NULL,
                                    TIME INT DEFAULT NULL,
                                    LAT REAL DEFAULT NULL,
                                    LON REAL DEFAULT NULL,
                                    FI REAL DEFAULT NULL,
                                    NN INT DEFAULT NULL,
                                    DD REAL DEFAULT NULL,
                                    FF REAL DEFAULT NULL,
                                    GG REAL DEFAULT NULL,
                                    TT REAL DEFAULT NULL,
                                    RH REAL DEFAULT NULL,
                                    PS REAL DEFAULT NULL,
                                    PSS REAL DEFAULT NULL,
                                    PE REAL DEFAULT NULL,
                                    PE1 REAL DEFAULT NULL,
                                    PE3 REAL DEFAULT NULL,
                                    PE6 REAL DEFAULT NULL,
                                    PE24 REAL DEFAULT NULL,
                                    QQ REAL DEFAULT NULL,
                                    VI REAL DEFAULT NULL,
                                    TD REAL DEFAULT NULL,
                                    TX REAL DEFAULT NULL,
                                    TM REAL DEFAULT NULL,
                                    GM REAL DEFAULT NULL,
                                    GX REAL DEFAULT NULL,
                                    WX REAL DEFAULT NULL,
                                    GW REAL DEFAULT NULL,
                                    PRIMARY KEY (ID, TIME));"""

        SQLite.execute(db, sqliteCreateTable)
    end

end


text:
# Composite vector spaces
#-------------------------
# Obtained as (some subspace of) the tensor product elementary vector spaces
abstract CompositeSpace{S<:ElementarySpace} <: VectorSpace

Base.eltype{S<:ElementarySpace}(V::CompositeSpace{S}) = eltype(S)
Base.eltype{S<:ElementarySpace}(::Type{CompositeSpace{S}}) = eltype(S)
Base.eltype{S<:CompositeSpace}(::Type{S}) = eltype(super(S))

# Composite spaces with finite number of elements N, defining spaces of tensors
abstract TensorSpace{S<:ElementarySpace,N} <: CompositeSpace{S}
Base.isfinite{S<:ElementarySpace}(::TensorSpace{S,:Inf}) = false
Base.isfinite(::TensorSpace) = true

# Explicit realizations:
#------------------------
# ProductSpace: type and methods for tensor products of ElementarySpace objects
include("spaces/productspace.jl")

# InvariantSpace: invariant subspace of tensor product of UnitaryRepresentationSpace objects
include("spaces/invariantspace.jl")

# Other examples might include:
# braidedspace and fermionspace
# symmetric and antisymmetric subspace of a tensor product of identical vector spaces
# ...

ext: jl


text:
import LinearAlgebra: norm, dot, cross

mutable struct Vec <: NonAbelian
	v
end
Vec(a...)=Vec([a...])
@delegate Vec.v [#=getindex, setindex!,=# iterate, length]
function Base.broadcast(::typeof(*),v::Vec,a::Factor)
	nv=Any[]
	for val in v.v
		push!(nv,val*a)
	end
	return Vec(nv)
end
mutable struct Cross <: Component
	x
	y
end
function getvec(c::Cross)
	ind=indsin(c,Vec)
	args=deepcopy(getargs(c))
	ex=args[ind[1][1]][ind[1][2][1][1]]
	vec=ex[ind[1][2][1][2][1]]
	deleteat!(ex,ind[1][2][1][2][1])
	return (ex,vec)
end
function getvec(ex::Expression)
	@assert length(ex)==1
	ind=indsin(ex,Vec)
	ex=deepcopy(ex)
				#implement support for multiple vectors/terms
	vec=ex[1][ind[1][2][1]]
	deleteat!(ex[1],ind[1][2][1])
	return (extract(ex),vec)
end
function cross(v1::Vec,v2::Vec)
	if isa(v1.v,Vector)&&isa(v2.v,Vector)
		return Vec(cross(v1.v,v2.v))
	end
	return Cross(v1,v2)
end
function simplify(c::Cross)
	if (isa(c.x,Vec)&&isa(c.y,Vec)&&isa(c.x.v,Vector)&&isa(c.y.v,Vector))
		return Vec(cross(c.x.v,c.y.v))
	end
	if isa(c.x,Expression)&&isa(c.y,Expression)
		ind=indsin(c,Vec)
		args=deepcopy(getargs(c))
		ex1=args[ind[1][1]][ind[1][2][1][1]]
		v1=ex1[ind[1][2][1][2][1]]
		deleteat!(ex1,ind[1][2][1][2][1])
		ex2=args[ind[2][1]][ind[2][2][1][1]]
		v2=ex2[ind[2][2][1][2][1]]
		deleteat!(ex2,ind[2][2][1][2][1])
		return extract(expression(ex1))*extract(expression(ex2))*Vec(cross(v1.v,v2.v))
	end
	return c
end
print(io::IO,c::Cross)=print(io,c.x,'×',c.y)
mutable struct Norm <: Component
	x
end
norm(ex::Ex)=Norm(ex)
function norm(vec::Vec)
	if isa(vec.v,Vector)
		return norm(vec.v)
	elseif isa(vec.v,Expression)
		ex,v=getvec(vec.v)
		return ex*norm(v)
	else
		return Norm(vec)
	end
#	has(vec.v,Ex)?Norm(vec):norm(vec.v)
end
function simplify(n::Norm)
	if isa(n.x,Expression)&&has(n.x,Vec)
		ex,v=getvec(n.x)
		return ex*norm(v)
	end
	norm(n.x)
end
mutable struct Dot <: Component
	x
	y
end
print(io::IO,c::Dot)=print(io,c.x,'⋅',c.y)
dot(x::EX,y::EX)=Dot(x,y)
function dot(v1::Vec,v2::Vec)
	if isa(v1.v,Array)&&isa(v2.v,Array)&&length(v1.v)==length(v2.v)
		ex=Expression(Term[Factor[v1.v[1]*v2.v[1]]])
		for t in 2:length(v1.v)
			push!(ex,Factor[v1.v[t],v2.v[t]])
		end
		return simplify(ex)
	end
	return Dot(v1,v2)
end
function simplify(d::Dot)
	if !isempty(methods(dot,(typeof(d.x),typeof(d.y))))
		return dot(d.x,d.y)
	end
	return d
end


text: GLOQ_MICRO_SEC = 1000.0


text:
"""
    ReplitDatabase

The `ReplitDatabase` module allows for easy access to the Repl.it database.

## Exports

- `ReplitDB` - an `AbstractDict{String,String}` interface to the Repl.it database.
- `ReplitDatabaseCore` - the core functions which `ReplitDB` uses to communicate with the database.
  - `get(key; url)`
  - `set!(key, value; url)`
  - `list([prefix]; url)`
  - `delete!(key; url)`
"""
module ReplitDatabase
export db, ReplitDB, ReplitDatabaseCore

include(joinpath(@__DIR__, "ReplitDatabaseShell.jl"))

struct ReplitDB <: AbstractDict{String,String}
    url::String
end

"""
    ReplitDB(url = ENV["REPLIT_DB_URL"])

An `AbstractDict{String,String}` interface representing the Repl.it database.

## Examples

```julia
julia> db = ReplitDB()
ReplitDB()

julia> db["foo"] = "bar"
"bar"

julia> db["hello"] = "world"
"bar"

julia> collect(keys(db))
2-element Array{String,1}:
 "foo"
 "hello"

julia> collect(keys(db, "fo"))
1-element Array{String,1}:
 "foo"

julia> delete!(db, "foo")
ReplitDB()

julia> isempty(db)
false

julia> empty!(db)
ReplitDB()

julia> isempty(db)
true
```
"""
function ReplitDB()
    if haskey(ENV, "REPLIT_DB_URL")
        ReplitDB(ENV["REPLIT_DB_URL"])
    else
        throw(ArgumentError("Could not find the ReplitDB URL in the environment"))
    end
end

function Base.delete!(db::ReplitDB, key::AbstractString)
    if haskey(db, key)
        ReplitDatabaseCore.delete!(key, url=db.url)
        db
    else
        throw(KeyError(key))
    end
end

Base.eltype(::ReplitDB) = Pair{String,String}

function Base.empty!(db::ReplitDB)
    for key in keys(db)
        delete!(db, key)
    end
    db
end

Base.get(db::ReplitDB, key::AbstractString, default) = get(() -> default, db, key)

Base.get(default::Function, db::ReplitDB, key::AbstractString) = haskey(db, key) ? ReplitDatabaseCore.get(key, url=db.url) : default()

Base.getkey(db::ReplitDB, key::AbstractString, default) = haskey(db, key) ? key : default

Base.haskey(db::ReplitDB, key::AbstractString) = key in keys(db, key)

function Base.iterate(db::ReplitDB, state=keys(db))
    if !isempty(state)
        key, rest = Iterators.peel(state)
        (key => db[key], rest)
    end
end

Base.isempty(db::ReplitDB) = isempty(keys(db))

Base.IteratorEltype(::ReplitDB) = Base.HasEltype()

Base.IteratorSize(::ReplitDB) = Base.SizeUnknown()

Base.keys(db::ReplitDB, prefix::AbstractString="") = ReplitDatabaseCore.list(prefix, url=db.url)

Base.length(db::ReplitDB) = length(keys(db))

function Base.pop!(db::ReplitDB, key::AbstractString)
    value = get(db, key) do
        throw(KeyError(key))
    end
    delete!(db, key)
    value
end

function Base.pop!(db::ReplitDB, key::AbstractString, default)
    if haskey(db, key)
        value = get(db, key, nothing)
        delete!(db, key)
        value
    else
        default
    end
end

function Base.setindex!(db::ReplitDB, value::AbstractString, key::AbstractString)
    ReplitDatabaseCore.set!(key, value, url=db.url)
    value
end

Base.setindex!(db::ReplitDB, ::Nothing, key::AbstractString) = delete!(db, key)

Base.setindex!(db::ReplitDB, value, key::AbstractString) = db[key] = string(value)

Base.show(io::IO, ::MIME"text/plain", db::ReplitDB) = print(io, "ReplitDB()")

"""
    db::AbstractDict{String,String}

The default Repl.it database.

## Examples
```julia
julia> db["foo"] = "bar"  # set a key/value pair
"bar"

julia> db["foo"]  # get a value
"bar"

julia> keys(db, "f")  # list all keys starting with an optional prefix
1-element Array{String,1}:
 "foo"

julia> db["foo"] = nothing  # delete a key/value pair
```
"""
db = ReplitDB()

end # module


text:
using Test

using MathOptInterface
const MOI = MathOptInterface
const MOIT = MathOptInterface.Test
const MOIU = MathOptInterface.Utilities
const MOIB = MathOptInterface.Bridges

include("../utilities.jl")

mock = MOIU.MockOptimizer(MOIU.Model{Float64}())
config = MOIT.TestConfig()

bridged_mock = MOIB.Objective.Functionize{Float64}(mock)

@testset "solve_singlevariable_obj" begin
    MOIU.set_mock_optimize!(mock,
        (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(mock, [1.0], MOI.FEASIBLE_POINT)
    )
    MOIT.solve_singlevariable_obj(bridged_mock, config)
    @test MOI.get(mock, MOI.ObjectiveFunctionType()) == MOI.ScalarAffineFunction{Float64}
    @test MOI.get(bridged_mock, MOI.ObjectiveFunctionType()) == MOI.SingleVariable
    @test MOI.get(mock, MOI.ObjectiveSense()) == MOI.MIN_SENSE
    @test MOI.get(bridged_mock, MOI.ObjectiveSense()) == MOI.MIN_SENSE
    vis = MOI.get(bridged_mock, MOI.ListOfVariableIndices())
    func = MOI.SingleVariable(vis[1])
    @test MOI.get(mock, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}()) ≈
        convert(MOI.ScalarAffineFunction{Float64}, func)
    @test MOI.get(bridged_mock, MOI.ObjectiveFunction{MOI.SingleVariable}()) == func
    MOI.set(bridged_mock, MOI.ObjectiveSense(), MOI.MAX_SENSE)
    @test MOI.get(mock, MOI.ObjectiveSense()) == MOI.MAX_SENSE
    @test MOI.get(bridged_mock, MOI.ObjectiveSense()) == MOI.MAX_SENSE
    test_delete_objective(bridged_mock, 1, tuple())
end

# Tests that the `ObjectiveValue` attribute passed has the correct
# `result_index`.
@testset "solve_result_index" begin
    MOIU.set_mock_optimize!(mock,
        (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(
            mock,
            MOI.OPTIMAL,
            (MOI.FEASIBLE_POINT, [1.0]),
            MOI.FEASIBLE_POINT,
            (MOI.SingleVariable, MOI.GreaterThan{Float64}) => [1.0],
        )
    )
    MOIT.solve_result_index(bridged_mock, config)
end


text:
"""
    fit_echo_linear(
        phas::AbstractArray{<:AbstractFloat, 4},
        W::AbstractArray{<:AbstractFloat, 4},
        TEs::NTuple{N, Real};
        phase_offset::Bool = true
    ) -> NTuple{1 or 2, typeof(similar(phas)){3}}

Weighted least squares for multi-echo data.

### Arguments
- `phas::AbstractArray{<:AbstractFloat, 4}`: unwrapped multi-echo phase
- `W::AbstractArray{<:AbstractFloat, 4}`: reciprocal of error variance of voxel
- `TEs::NTuple{N, Real}`: echo times

### Keywords
- `phase_offset::Bool = true`: model phase offset (`true`)

### Returns
- `typeof(similar(phas)){3}`: weighted least-squares estimate for phase
- [`typeof(similar(phas)){3}`]: weighted least-squares estimate for phase offset
    if `phase_offset = true`
"""
function fit_echo_linear(
    phas::AbstractArray{<:AbstractFloat, 4},
    W::AbstractArray{<:AbstractFloat, 4},
    TEs::NTuple{N, Real};
    phase_offset::Bool = true
) where {N}
    p = tzero(phas, size(phas)[1:3])
    if !phase_offset
        return fit_echo_linear!(p, phas, W, TEs)
    else
        return fit_echo_linear!(p, tzero(p), phas, W, TEs)
    end
end

"""
    fit_echo_linear!(
        p::AbstractArray{<:AbstractFloat, 3},
        phas::AbstractArray{<:AbstractFloat, 4},
        W::AbstractArray{<:AbstractFloat, 4},
        TEs::NTuple{N, Real}
    ) -> p

Weighted least squares for multi-echo data (phase offset = 0).

### Arguments
- `p::AbstractArray{<:AbstractFloat, 3}`: weighted least-squares estimate for phase
- `phas::AbstractArray{<:AbstractFloat, 4}`: unwrapped multi-echo phase
- `W::AbstractArray{<:AbstractFloat, 4}`: reciprocal of error variance of voxel
- `TEs::NTuple{N, Real}`: echo times

### Returns
- `p`: weighted least-squares estimate for phase
"""
function fit_echo_linear!(
    p::AbstractArray{Tp, 3},
    phas::AbstractArray{Tphas, 4},
    W::AbstractArray{TW, 4},
    TEs::NTuple{NT, Real}
) where {Tp<:AbstractFloat, Tphas<:AbstractFloat, TW<:Real, NT}
    nx, ny, nz, nt = size(phas)

    nt == NT || throw(DimensionMismatch())
    size(p) == (nx, ny, nz) || throw(DimensionMismatch())
    size(W) == size(phas)   || throw(DimensionMismatch())

    T = foldl(promote_type, (Tphas, TW), init=Tp)
    tes = convert.(T, TEs)

    _zeroT = zero(T)
    _zeroTp = zero(Tp)

    @threads for k in 1:nz
        @inbounds for j in 1:ny
            for i in 1:nx
                num = _zeroT
                den = _zeroT
                for t in Base.OneTo(NT)
                    w = W[i,j,k,t] * W[i,j,k,t] * tes[t]
                    num = muladd(w, phas[i,j,k,t], num)
                    den = muladd(w, tes[t], den)
                end
                p[i,j,k] = iszero(den) ? _zeroTp : num * inv(den)
            end
        end
    end

    return p
end

"""
    fit_echo_linear!(
        p::AbstractArray{<:AbstractFloat, 3},
        p0::AbstractArray{<:AbstractFloat, 3},
        phas::AbstractArray{<:AbstractFloat, 4},
        W::AbstractArray{<:AbstractFloat, 4},
        TEs::NTuple{N, Real}
    ) -> (p, p0)

Weighted least squares for multi-echo data (estimate phase offset).

### Arguments
- `p::AbstractArray{<:AbstractFloat, 3}`: weighted least-squares estimate for phase
- `p0::AbstractArray{<:AbstractFloat, 3}`: weighted least-squares estimate for phase offset
- `phas::AbstractArray{<:AbstractFloat, 4}`: unwrapped multi-echo phase
- `W::AbstractArray{<:AbstractFloat, 4}`: reciprocal of error variance of voxel
- `TEs::NTuple{N, Real}`: echo times

### Returns
- `p`: weighted least-squares estimate for phase
- `p0`: weighted least-squares estimate for phase offset
"""
function fit_echo_linear!(
    p::AbstractArray{Tp, 3},
    p0::AbstractArray{Tp0, 3},
    phas::AbstractArray{Tphas, 4},
    W::AbstractArray{TW, 4},
    TEs::NTuple{NT, Real}
) where {Tp<:AbstractFloat, Tp0<:AbstractFloat, Tphas<:AbstractFloat, TW<:Real, NT}
    nx, ny, nz, nt = size(phas)

    nt == NT || throw(DimensionMismatch())
    size(p)  == (nx, ny, nz) || throw(DimensionMismatch())
    size(p0) == (nx, ny, nz) || throw(DimensionMismatch())
    size(W)  == size(phas)   || throw(DimensionMismatch())

    T = foldl(promote_type, (Tp0, Tphas, TW), init=Tp)
    tes = convert.(T, TEs)

    _zeroT = zero(T)
    _zeroTp = zero(Tp)

    @threads for k in 1:nz
        @inbounds for j in 1:ny
            for i in 1:nx
                x = _zeroT
                y = _zeroT
                w = _zeroT
                for t in Base.OneTo(NT)
                    w += W[i,j,k,t]
                    x = muladd(W[i,j,k,t], tes[t], x)
                    y = muladd(W[i,j,k,t], phas[i,j,k,t], y)
                end

                w = inv(w)
                x *= w
                y *= w

                num = _zeroT
                den = _zeroT
                for t in Base.OneTo(NT)
                    xx = tes[t] - x
                    yy = phas[i,j,k,t] - y
                    ww = W[i,j,k,t] * xx
                    num = muladd(ww, yy, num)
                    den = muladd(ww, xx, den)
                end

                p[i,j,k] = iszero(den) ? _zeroTp : num * inv(den)
                p0[i,j,k] = y - p[i,j,k] * x
            end
        end
    end

    return p, p0
end


text:
# Tomography Tests
using Schrodinger, Distributions
using Schrodinger: mle_state_tomo, build_density_matrix, pgd_process_tomo,
    apply_process, operator_to_choi, trace_norm, gate_fidelity_choi,
    state_likelihood_model, process_likelihood_model

using Schrodinger: kraus_to_natural, natural_to_kraus, natural_to_choi,
    choi_to_natural, kraus_to_choi, choi_to_kraus, choi_to_chi, chi_to_choi

using Compat.Test, Compat.LinearAlgebra, Compat.Random
if VERSION > v"0.7.0-"
    using Compat.Statistics: mean
end
println("Testing Tomography...")

@testset "Tomography" begin
# Build a few different variable for testing
d = 2; N = 10^5
g = basis(2,0); e1 = basis(2,1)
ONEQUBIT = normalize!.([(g+e1),(g-e1),(g+1im*e1),(g-1im*e1),g,e1])

@testset "State Tomography" begin
# Input states and POVM operators for state tomo
ψ = coherent(d,0.32)
E_m = Operator.(ONEQUBIT)./3
# Create the A matrix
A = @inferred state_likelihood_model(E_m)
# Simulate tomography measurement by drawing from a binomial distribution
P = [real(expect(ψ,E)) for E ∈ E_m]
M = [rand(Binomial(N,p))/N for p in P]
# Reconstruct state
ρ = Operator(build_density_matrix(mle_state_tomo(M,A).minimizer))
# Fidelity
F1 = @inferred fidelity2(ρ,ψ)
@test 1-F1 < 0.001
end

@testset "Process Tomography" begin
# Input states and POVM operators for process tomo
ρ_in = Operator.(ONEQUBIT)
E_m  = ρ_in./3
@test sum(E_m) ≈ qeye(d)

# Create the likelihood model matrix ("A-matrix")
A = @inferred process_likelihood_model(ρ_in,E_m)

# Test reconstruction with a couple of notable gates
X½ = Operator([1 -im; -im 1]/√2)
Y½ = Operator([1 -1; 1 1]/√2)
T = Operator([1 0; 0 cis(π/4)])
H = Operator([1 1; 1 -1]/√2)
for G ∈ [σ0, σx, σy, σz, X½, Y½, T, H]
    # Simulate tomography measurement by drawing from a binomial distribution
    P = [real(expect(G*ρ*G',E)) for E ∈ E_m, ρ ∈ ρ_in]
    M = [rand(Binomial(N,p))/N for p in P]./size(P,2)
    # Reconstruct the Choi matrix from the measurements
    Crec = Operator(@inferred(pgd_process_tomo(M,A,info=false)),(d,d))
    # Compare using the J distance, aka the trace norm of the difference
    Ctrue = @inferred operator_to_choi(G) # ideal Choi matrix
    @test @inferred(trace_norm(Ctrue-Crec))/2d < 0.003
    # Compare with gate fidelity
    F1 = @inferred gate_fidelity_choi(Crec,G)
    F2 = mean(ψ->fidelity2(G*ψ,apply_process(Crec,ψ)), ONEQUBIT)
    @test 1-F1 < 0.002
    @test F1 ≈ F2

    # also test superoperator conversions (TODO: move somewhere else)
    @test Crec ≈ (Crec |> choi_to_kraus |> kraus_to_natural |> natural_to_choi)
    @test Crec ≈ (Crec |> choi_to_natural |> natural_to_kraus |> kraus_to_choi)
    @test Crec ≈ (Crec |> choi_to_chi |> chi_to_choi)
end

# Also test with some 2Q gates
d = 4
ρ_in = Operator.(vec(ONEQUBIT .⊗ permutedims(ONEQUBIT)))
E_m = ρ_in/9 # POVM
@test sum(E_m) ≈ qeye((2,2))
A = process_likelihood_model(ρ_in,E_m)
include(joinpath(@__DIR__,"twoqubit.jl"))

CNOT = Operator(
  [1 0 0 0
   0 1 0 0
   0 0 0 1
   0 0 1 0], (2,2))
CZ = Operator(
  [1 0 0 0
   0 1 0 0
   0 0 1 0
   0 0 0 -1], (2,2))
SWAP = Operator(
  [1 0 0 0
   0 0 1 0
   0 1 0 0
   0 0 0 1], (2,2))
iSWAP = Operator(
  [1    0  0    0
   0    0  1im  0
   0  1im  0    0
   0    0  0    1], (2,2))
sqrtSWAP = Operator(
  [1         0         0  0
   0  (1+im)/2  (1-im)/2  0
   0  (1-im)/2  (1+im)/2  0
   0         0         0  1], (2,2))
sqrtiSWAP = Operator(
  [1        0          0  0
   0    1/√(2)  1im/√(2)  0
   0  1im/√(2)    1/√(2)  0
   0         0         0  1], (2,2))
for G ∈ [CNOT, CZ, SWAP, sqrtSWAP, iSWAP, sqrtiSWAP]
   # Simulate tomography measurement by drawing from a binomial distribution
   P = [real(expect(G*ρ*G',E)) for E ∈ E_m,ρ ∈ ρ_in]
   M = [rand(Binomial(N,p))/N for p in P]./size(P,2)
   # Reconstruct the Choi matrix from the measurements
   Crec = Operator(@inferred(pgd_process_tomo(M,A,info=false)),(2,2,2,2))
   # Compare using the J distance, aka the trace norm of the difference
   Ctrue = @inferred operator_to_choi(G) # ideal Choi matrix
   @test @inferred(trace_norm(Ctrue-Crec))/2d < 0.003
   # Compare with gate fidelity
   F1 = @inferred gate_fidelity_choi(Crec,G)
   F2 = mean(ψ->fidelity2(G*ψ,apply_process(Crec,ψ)), TWOQUBIT)
   @test 1-F1 < 0.002
   @test F1 ≈ F2

   # also test superoperator conversions (TODO: move somewhere else)
   #@test Crec ≈ (Crec |> choi_to_kraus |> kraus_to_natural |> natural_to_choi)
   #@test Crec ≈ (Crec |> choi_to_natural |> natural_to_kraus |> kraus_to_choi)
   #@test Crec ≈ (Crec |> choi_to_chi |> chi_to_choi)
end

end

end

# the below is for historical testing purpose
# the measurements come from a poorly implemented Z gate
M = [ 47   736   395   358   421   383
     710   107   323   468   423   315
     338   428   128   646   429   338
     440   359   731    46   407   408
     403   340   349   400   753    36
     400   391   384   408    10   755]./(6*3*786)
# benchmark with CPTP tol of 1E-4
# BenchmarkTools.Trial:
#   memory estimate:  965.63 KiB
#   allocs estimate:  2349
#   --------------
#   minimum time:     873.043 μs (0.00% GC)
#   median time:      902.563 μs (0.00% GC)
#   mean time:        974.332 μs (6.05% GC)
#   maximum time:     2.970 ms (50.09% GC)
#   --------------
#   samples:          5116
#   evals/sample:     1


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

using Markdown
using InteractiveUtils

# ╔═╡ e077ab16-22e5-11eb-15ed-5f0654bf7328
using Pkg, DrWatson

# ╔═╡ e75a9a06-22e5-11eb-2e36-8d4f62b830ed
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using CSV, DataFrames
end

# ╔═╡ ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
md"## Listing4.04"

# ╔═╡ fa68607e-22e5-11eb-0558-c9a4d9f77426
data = CSV.read(datadir("purchaseData.csv"), copycols = true)

# ╔═╡ 700ec498-243d-11eb-30a7-0bd31057f6d6
("Grade of person 1: ", data[1, 3],
        ", ", data[1,:Grade],
        ", ", data.Grade[1], "\n")

# ╔═╡ 700ef812-243d-11eb-2b07-89169e32b9d9
data[[1,2,4], :]

# ╔═╡ 700f7f12-243d-11eb-3f40-d1a2c7403417
data[13:15, :Name]

# ╔═╡ 7019d3ce-243d-11eb-0703-a386528626b7
data.Name[13:15]

# ╔═╡ 701a53d6-243d-11eb-3b69-a54dfe2b3626
data[13:15, [:Name]]

# ╔═╡ 475ff888-22e6-11eb-2354-09f8f40a8e12
md"## End of listing4.04"

# ╔═╡ Cell order:
# ╟─ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
# ╠═e077ab16-22e5-11eb-15ed-5f0654bf7328
# ╠═e75a9a06-22e5-11eb-2e36-8d4f62b830ed
# ╠═fa68607e-22e5-11eb-0558-c9a4d9f77426
# ╠═700ec498-243d-11eb-30a7-0bd31057f6d6
# ╠═700ef812-243d-11eb-2b07-89169e32b9d9
# ╠═700f7f12-243d-11eb-3f40-d1a2c7403417
# ╠═7019d3ce-243d-11eb-0703-a386528626b7
# ╠═701a53d6-243d-11eb-3b69-a54dfe2b3626
# ╟─475ff888-22e6-11eb-2354-09f8f40a8e12


text:
module DanceDocker

import Dance.Configuration

export setup

"""
    setup(path::String="."; version::String=string(VERSION))

- Julia version number can be overwritten, else defaults to that of system
- Check if project contains `Project.toml`, `dance.jl` and `settings` dir, else means project is not properly configured with DanceJL
- Obtain server port (specified under `Dockerfile`) from `Dance.Configuration.Settings` dict
- Generate `Dockerfile` using template from package `files` dir
"""
function setup(path::String="."; version::String=string(VERSION)) :: Nothing
    if !isfile("Project.toml")
        @error "Please add Project.toml to project root"
    elseif !isfile("dance.jl")
        @error "Please add dance.jl to project root"
    elseif !isdir("settings")
        @error "Please add settings dir to project root"
    else
        if Configuration.populate(path)
            settings::Dict{Symbol, Any} = Configuration.Settings
            port::Int64 = settings[:server_port]

            dest_dockerfile_path::String = joinpath(abspath(path), "Dockerfile")
            cp(joinpath(@__DIR__, "../files/Dockerfile"), dest_dockerfile_path)
            if Sys.iswindows()
                run(`icacls.exe $dest_dockerfile_path /reset /T /Q`)
                username::String = read(run(`whoami`), String)
                run(`icacls.exe $dest_dockerfile_path /grant $username:F /T /Q`)
            else
                run(`chmod 755 $dest_dockerfile_path`)
            end

            file_string::String = ""
            open("Dockerfile", "r") do io
                file_string = read(io, String)
            end
            file_string = replace(file_string, "\${version}" => version)
            file_string = replace(file_string, "\${port}" => string(port))
            open("Dockerfile", "w") do io
                write(io, file_string)
            end

            @info "Dockerfile added to project root"
        else
            @error "Error setting up project settings. Please see error above."
        end
    end
end

end


text:
#-------------genetic.jl-------------------------------------------------------#
#
# Purpose: To implement a simple genetic algorithm for function minimization
#
#------------------------------------------------------------------------------#

using PyPlot

# For now, we will evaluate (x-1)^2
function evaluate(x::Float64)
    return (x-1)*(x-1)
end

# Initialize num positions centered around origin
function init(num::Int64, box_length::Float64)
    return rand(num) * box_length - box_length * 0.5
end

# Function for actual tournament
function tournament(population::Vector{Float64}, size::Int64)
    # Selecting first individual
    best = rand(population)

    # Choosing best from random selection of others
    for j = 2:size
        candidate = rand(population)
        if (evaluate(candidate) < evaluate(best))
            best = candidate
        end
    end

    return best
end

# Tournament selection, size is tournament size, and prob is prob of selection
# num_parents must be even
# The num_children = num_parents, make sure that
#     the children size < population size - elite number
function selection(population::Vector{Float64}, size::Int64,
                   num_parents::Int64)

    if (num_parents % 2 != 0)
        num_parents += 1
    end

    # result is the final selected population
    parents = fill(0.0, num_parents)

    # selection loop
    for i = 1:num_parents
        parents[i] = tournament(population, size)
    end

    return parents
end

# Crossover
function crossover(parents::Vector{Float64}, res::Int64, cross_rate::Float64)

    # vector of offspring to work from
    offspring = fill(0.0, length(parents))
    child1::Float64 = 0.0
    child2::Float64 = 0.0

    # We are selecting two parents and creating two children
    for i = 1:2:length(parents)
        child1 = 0
        child2 = 0
        # splitting each parent into segments based on location
        for j = 1:res
           index = rand()
           if (index > cross_rate)
               child1 += parents[i] / res
               child2 += parents[i+1] / res
           else
               child1 += parents[i+1] / res
               child2 += parents[i] / res
           end
        end
        offspring[i] = child1
        offspring[i+1] = child2
    end

    return offspring
end

# Section of individuals for repopulation
function repopulate(population::Vector{Float64}, parents::Vector{Float64},
                    offspring::Vector{Float64}, elite::Int64, box_size::Float64)

    # grabbing the elite from the previous population
    sort!(population; lt = (a, b) -> evaluate(a) < evaluate(b))

    # Randomizing all elements that are not elite or children
    leftover_num = length(population) - length(offspring) - elite + 1
    if leftover_num < 0
        println("More offspring / elite than in population")
    end
    leftovers = init(abs(leftover_num), box_size)

    println(length(parents), '\t', length(offspring), '\t',
            length(population), '\t', elite, '\t', leftover_num)

    j = 1
    for i = elite:length(population)
        if j <= length(offspring)
            population[i] = offspring[j]
        else
            population[i] = leftovers[j - length(offspring)]
        end
        j += 1
    end

    return population
end

# Function for mutation of an individual
function mutation(offspring::Vector{Float64}, mutation_rate::Float64,
                  box_length::Float64)

    # randomly mutating offspring
    for i = 1:length(offspring)
       if rand() < mutation_rate
           offspring[i] = rand() * box_length - box_length * 0.5
       end
    end

    return offspring

end

# main function
function main(timesteps::Int64)
    num = 100
    box_length = 5.0
    population::Vector{Float64} = init(num, box_length)
    plot(population)

    for i =1:timesteps
        parents::Vector{Float64} = selection(population, 10, num)
        offspring::Vector{Float64} = crossover(parents, 10, .5)
        #offspring = mutation(offspring, 0.1, box_length)
        population = repopulate(population, parents, offspring, 10, box_length)
        plot(population)
        sleep(0.5)
    end
end

main(10)

max_issues_repo_name: shiffman/simuleios

max_forks_repo_name: shiffman/simuleios


text:
struct Features{T}
    """
    A type for holding all the fixed quantities during a simulation. This involves
    simulation constants for detection probabiliites, constraints for consistency
    with data and some data based measures like the local cases population size,
    duration of sims and the length of the observation perod.
    """
    # constraints for the simulation
    max_forecast_cases::T
    cases_pre_forecast::T
    N::T
    T_observed::T
    T_end::T
    omicron_dominant_day::T
end

struct Realisations{T1, T2}
    """
    A type for holding the realisations of the simulations. This is a cleaner way of
    holding the information for the three different matrices used. Z is for infections
    D is for observed cases and U is for undetected cases.
    """
    Z::T1
    D::T2
    U::T2
end

struct Constants{T}
    """
    A type for the dynamical / simulation constants
    """
    k::T
    p_symp::T
    p_detect_given_symp::T
    p_detect_given_asymp::T
    p_detect::T
    p_symp_given_detect::T
    γ::T
    α_s::T
    α_a::T
    qi::T
    # import model parameters
    prior_alpha::T
    prior_beta::T
    ϕ::T
    # multiplier for the consistency between observed cases
    consistency_multiplier::T
end

struct IndividualTypeMap{T}
    S::T
    A::T
    I::T
end

struct Forecast{T1,T2,T3,T4,T5}
    """
    This type wraps all the information relating to the forecast in a single accessible
    object. This drastically simplifies the layout of the code at no performance hit.
    """
    sim_features::Features{T1}
    sim_realisations::Realisations{T2,T3}
    sim_constants::Constants{T4}
    individual_type_map::IndividualTypeMap{T5}
end


text:
using JuMP, Mosek, MosekTools, LinearAlgebra, Suppressor, StatsBase#, SCS#, Ipopt

function getMatrixNuclearNorm(coords, entries, M_noisy, n, numEntriesSampled, mu=sqrt(2*numEntriesSampled/n), imposeConstraints=false) #mu picked in line with Candes and Plan
    m=Model(Mosek.Optimizer)
    @variable(m, X[1:n, 1:n])
    @variable(m, U[1:n, 1:n])
    @variable(m, V[1:n, 1:n])
    if imposeConstraints
        for t=1:Int(numEntriesSampled)
            @constraint(m, X[coords[entries[t],1], coords[entries[t],2]]==M_noisy[coords[entries[t],1], coords[entries[t],2]])
        end
        @objective(m, Min, sum(U[i,i] for i=1:n)+sum(V[i,i] for i=1:n))
    else
        @objective(m, Min, mu*sum(U[i,i] for i=1:n)+mu*sum(V[i,i] for i=1:n)+sum((X[coords[entries[t],1], coords[entries[t],2]]-M_noisy[coords[entries[t],1], coords[entries[t],2]])^2 for t in 1:Int(numEntriesSampled)));
    end
    @constraint(m, Symmetric([U X; X' V]) in PSDCone());
    @suppress optimize!(m)
   return value.(X)
end

function getMatrixFrobeniusNorm_reg2ndstage(coords, entries, M, n, numEntriesSampled, k=sqrt(n), gamma=200*sqrt(n)/(numEntriesSampled), imposeConstraints=false)
    m=Model(Mosek.Optimizer)

    @variable(m, X[1:n, 1:n])
    @variable(m, Y[1:n, 1:n], Symmetric)
    @variable(m, theta[1:n, 1:n])

    if imposeConstraints
        for t=1:Int(numEntriesSampled)
            @constraint(m, X[coords[entries[t],1], coords[entries[t],2]]==M[coords[entries[t],1], coords[entries[t],2]])
        end
        @objective(m, Min, 1.0/(2.0*gamma)*sum(theta[i,i] for i=1:n))
    else
        @objective(m, Min, 1.0/(2.0*gamma)*sum(theta[i,i] for i=1:n)+sum((X[coords[entries[t],1], coords[entries[t],2]]-M[coords[entries[t],1], coords[entries[t],2]])^2 for t in 1:Int(numEntriesSampled)))
    end

    @constraint(m, Symmetric(Matrix(1.0I, n, n)-Y) in PSDCone());

    @constraint(m, Symmetric([Y X; X' theta]) in PSDCone());
    @constraint(m, sum(Y[i,i] for i=1:n)<=k)

    @suppress optimize!(m)
    u,v,w=LinearAlgebra.svd(value.(Y))

    # do 100 iterations of randomized rounding and keep best solution in terms of in-sample objective.
    X_best=zeros(n,n)
    ofv_best=1e10
    numIters=100
    for i=1:numIters
        v_rounded=(rand(n).<v)
        if sum(v_rounded)<=k
            Y_rounded=u*Diagonal(v_rounded)*u'
            m2=Model(with_optimizer(Mosek.Optimizer))
            @variable(m2, X2[1:n, 1:n])
            @objective(m2, Min, 1.0/(2.0*gamma)*sum(X2[i,j]^2 for i=1:n for j=1:n)+sum((X2[coords[entries[t],1], coords[entries[t],2]]-M[coords[entries[t],1], coords[entries[t],2]])^2 for t in 1:Int(numEntriesSampled)))
            @constraint(m2, X2.==Y_rounded*X2)
            @suppress optimize!(m2)
            if objective_value(m2) < ofv_best
               ofv_best = objective_value(m2)
               X_best=value.(X2)
            end
        end
    end
   return X_best
end

function getMatrixFrobeniusNorm_greedy(coords, entries, M, n, numEntriesSampled, k=sqrt(n), gamma=200*sqrt(n)/(numEntriesSampled), imposeConstraints=false, useRotation=false)
    m=Model(Mosek.Optimizer)

    @variable(m, X[1:n, 1:n])
    @variable(m, Y[1:n, 1:n], Symmetric)
    @variable(m, theta[1:n, 1:n])
    @variable(m, sigma>=0.0)

    @constraint(m, sigma>=1.0/(2.0*gamma)*sum(theta[i,i] for i=1:n))
    @SDconstraint(m, Matrix(1.0I, n, n)>=Y);

    if imposeConstraints
        for t=1:Int(numEntriesSampled)
            @constraint(m, X[coords[entries[t],1], coords[entries[t],2]]==M[coords[entries[t],1], coords[entries[t],2]])
        end

        @objective(m, Min, sigma)
    else
        @objective(m, Min, sigma+0.5*sum((X[coords[entries[t],1], coords[entries[t],2]]-M[coords[entries[t],1], coords[entries[t],2]])^2 for t in 1:Int(numEntriesSampled)))
    end
    @constraint(m, Symmetric([Y X; X' theta]) in PSDCone());
    @constraint(m, sum(Y[i,i] for i=1:n)<=k)

    @suppress optimize!(m)
    @show objective_value(m)
    u,v,w=LinearAlgebra.svd(value.(Y))

    diag_v=zeros(n)
    if useRotation
        # Compute alpha, use to compute weights (leaving out -1/gamma term as scales out from squaring)
        @show "Using rotation"
        alpha_star=pinv(value.(Y))*value.(X) #strictly speaking only proportional to this
        w_full=alpha_star*alpha_star'
        diag_v[sortperm(diag(w_full), rev=true)[1:k]].=1.0
    else
        diag_v[sortperm(v, rev=true)[1:k]].=1.0
        # diag_v2[sortperm(v2, rev=true)[1:k]].=1.0
    end
    @show sum(diag_v)
    Y_rounded=u*Diagonal(diag_v)*u'
    # Y_rounded2=u2*Diagonal(diag_v2)*u2'
    @show sum(abs.(Y_rounded))

    m2=Model(with_optimizer(Mosek.Optimizer))
    @variable(m2, X2[1:n, 1:n])
    @objective(m2, Min, 1.0/(2.0*gamma)*sum(X2[i,j]^2 for i=1:n for j=1:n)+0.5*sum((X2[coords[entries[t],1], coords[entries[t],2]]-M[coords[entries[t],1], coords[entries[t],2]])^2 for t in 1:Int(numEntriesSampled)))
    @constraint(m2, X2.==Y_rounded*X2)
    # @constraint(m2, X2.==X2*Y_rounded2)
    @suppress optimize!(m2)
    X_final=value.(X2)
    @show sum(abs.(value.(X2)))
    #@show 1.0/(2.0*gamma)*sum(value(X2[i,j])^2 for i=1:n for j=1:n)+0.5*sum((value(X2[coords[entries[t],1], coords[entries[t],2]])-M[coords[entries[t],1], coords[entries[t],2]])^2 for t in 1:Int(numEntriesSampled))

   return X_final, objective_value(m), objective_value(m2), value.(Y)
end

function getMatrix_BM(coords, entries, M, n, numEntriesSampled, k=sqrt(n), gamma=200*sqrt(n)/(numEntriesSampled)) #can't impose constraints here because then we need to warm-start with a low-rank feasible solution, and if we have a method for doing that then we don't need to solve this problem.

    X_best=zeros(n,n)
    ofv_best=1e10
    numIters=1 #keeping it at 1 for now, since IPopt is rather slow with even a moderate number of constraints.

    m=Model(with_optimizer(Ipopt.Optimizer))

    @variable(m, X[1:n, 1:n])
    @variable(m, U[1:n, 1:k])
    @variable(m, V[1:k, 1:n])

    @objective(m, Min, 1.0/(2.0*gamma)*sum(X[i,j]^2 for i=1:n for j=1:n)+sum((X[coords[entries[t],1], coords[entries[t],2]]-M[coords[entries[t],1], coords[entries[t],2]])^2 for t in 1:Int(numEntriesSampled)))
    @constraint(m, X.>=U*V.-1e-6*ones(n,n))
    @constraint(m, X.<=U*V.+1e-6*ones(n,n))

    #@constraint(m, X.==U*V) #non-convex constraint encoding low-rank relation

    for iter=1:numIters
      U_init=randn(n,k) #same order of magnitude as actual solution
      V_init=randn(k,n)
      JuMP.set_start_value.(U, U_init)
      JuMP.set_start_value.(V, V_init)

      @suppress optimize!(m)
      if objective_value(m) < ofv_best
         @show ofv_best = objective_value(m)
         @show iter
         X_best=value.(X)
      end
    end

   return X_best
end

function getMatrixFrobeniusNorm_rademacher(coords, entries, M, n, numEntriesSampled, k=sqrt(n), gamma=200*sqrt(n)/(numEntriesSampled), imposeConstraints=false)
    m=Model(with_optimizer(Mosek.Optimizer))

    @variable(m, X[1:n, 1:n])
    @variable(m, Y[1:n, 1:n], Symmetric) #Need symmetry on Y because of the SVD step.
    @variable(m, theta[1:n, 1:n])

    @constraint(m, Symmetric(Matrix(1.0I, n, n)-Y) in PSDCone());

    if imposeConstraints
        for t=1:Int(numEntriesSampled)
            @constraint(m, X[coords[entries[t],1], coords[entries[t],2]]==M[coords[entries[t],1], coords[entries[t],2]])
        end
        @objective(m, Min, 1.0/(2.0*gamma)*sum(theta[i,i] for i=1:n))
    else
        @objective(m, Min, 1.0/(2.0*gamma)*sum(theta[i,i] for i=1:n)+sum((X[coords[entries[t],1], coords[entries[t],2]]-M[coords[entries[t],1], coords[entries[t],2]])^2 for t in 1:Int(numEntriesSampled)))
    end
    @constraint(m, Symmetric([Y X; X' theta]) in PSDCone());
    @constraint(m, sum(Y[i,i] for i=1:n)<=k)

    @suppress optimize!(m)
    u,v,w=LinearAlgebra.svd(value.(Y))

    # do 100 iterations of randomized rounding and keep best solution in terms of in-sample objective.
    X_best=zeros(n,n)
    ofv_best=1e10
    numIters=100
    for i=1:numIters
            Ξ=(2.0)*(rand(n,k).<=0.5).-1.0 # using distributions package would be more efficient, but introduces weird dependency issues
            Y_intermediate=(1.0/n)*u*Diagonal(sqrt.(v))*Ξ*Ξ'*Diagonal(sqrt.(v))*u'
            u2,v2,w2=LinearAlgebra.svd(Y_intermediate)
            v_rounded=zeros(n)
            v_rounded[sortperm(v2, rev=true)[1:k]].=1.0

            Y_rounded=u2*Diagonal(v_rounded)*u2'
            m2=Model(with_optimizer(Mosek.Optimizer))
            @variable(m2, X2[1:n, 1:n])
            @objective(m2, Min, 1.0/(2.0*gamma)*sum(X2[i,j]^2 for i=1:n for j=1:n)+sum((X2[coords[entries[t],1], coords[entries[t],2]]-M[coords[entries[t],1], coords[entries[t],2]])^2 for t in 1:Int(numEntriesSampled)))
            @constraint(m2, X2.==Y_rounded*X2)
            @suppress optimize!(m2)
            if objective_value(m2) < ofv_best
               ofv_best = objective_value(m2)
               X_best=value.(X2)
            end
    end
    return X_best
end

function getMatrixFrobeniusNorm_restrictedcoorddescent(coords, entries, M, n, numEntriesSampled, k=sqrt(n), gamma=200*sqrt(n)/(numEntriesSampled), imposeConstraints=false)
    m=Model(with_optimizer(Mosek.Optimizer))

    @variable(m, X[1:n, 1:n])
    @variable(m, Y[1:n, 1:n], Symmetric) #Need symmetry on Y because of the SVD step.
    @variable(m, theta[1:n, 1:n])

    @constraint(m, Symmetric(Matrix(1.0I, n, n)-Y) in PSDCone());

    if imposeConstraints
        for t=1:Int(numEntriesSampled)
            @constraint(m, X[coords[entries[t],1], coords[entries[t],2]]==M[coords[entries[t],1], coords[entries[t],2]])
        end
        @objective(m, Min, 1.0/(2.0*gamma)*sum(theta[i,i] for i=1:n))
    else
        @objective(m, Min, 1.0/(2.0*gamma)*sum(theta[i,i] for i=1:n)+sum((X[coords[entries[t],1], coords[entries[t],2]]-M[coords[entries[t],1], coords[entries[t],2]])^2 for t in 1:Int(numEntriesSampled)))
    end
    @constraint(m, Symmetric([Y X; X' theta]) in PSDCone());
    @constraint(m, sum(Y[i,i] for i=1:n)<=k)

    @suppress optimize!(m)
    u,v,w=LinearAlgebra.svd(value.(Y))

    # This gives us our u, v
    S=sortperm(v, rev=true)[1:min(2*k, n)]
    Y_search=zeros(n,n)
    for t=1:k
            best_ofv=1e10
            j=0
        for i in S
            # Evaluate f(Y+u_i u_i')
            m2=Model(with_optimizer(Mosek.Optimizer))
            @variable(m2, X2[1:n, 1:n])
            @objective(m2, Min, 1.0/(2.0*gamma)*sum(X2[i,j]^2 for i=1:n for j=1:n)+sum((X2[coords[entries[t],1], coords[entries[t],2]]-M[coords[entries[t],1], coords[entries[t],2]])^2 for t in 1:Int(numEntriesSampled)))
            @constraint(m2, X2.==(Y_search.+u[:,i]*u[:,i]')*X2)
            @suppress optimize!(m2)

            if objective_value(m2)<best_ofv
                best_ofv=objective_value(m2)
                j=i
            end
        end
        filter!(x->x≠j, S)
        Y_search=Y_search+u[:,j]*u[:,j]'
    end

# Finally, evaluate X and return it (could keep track of X as we go, but the idea of the paper is to work with Y wherever possible so we should really be doing that)
m2=Model(with_optimizer(Mosek.Optimizer))
@variable(m2, X2[1:n, 1:n])
@objective(m2, Min, 1.0/(2.0*gamma)*sum(X2[i,j]^2 for i=1:n for j=1:n)+sum((X2[coords[entries[t],1], coords[entries[t],2]]-M[coords[entries[t],1], coords[entries[t],2]])^2 for t in 1:Int(numEntriesSampled)))
@constraint(m2, X2.==Y_search*X2)
@suppress optimize!(m2)

return value.(X2)
end

# Note: this function is currently not used in the paper
function getMatrixFrobeniusNorm_greedystrengthened(coords, entries, M, n, numEntriesSampled, k=sqrt(n), gamma=200*sqrt(n)/(numEntriesSampled), imposeConstraints=false, useRotation=false)
    m=Model(with_optimizer(Mosek.Optimizer))
    U_bound=0.3

    @variable(m, X[1:n, 1:n])

    @variable(m, Y[1:n, 1:n], Symmetric)
    @variable(m, Y_sub[1:n, 1:n, 1:k])
    @variable(m, Y_sub2[1:n, 1:n, 1:k, 1:k])
    @variable(m, W[1:n, 1:k, 1:k])

    @variable(m, theta[1:n, 1:n])
    @variable(m, sigma>=0.0)
    @variable(m, U[1:n, 1:k])

    @constraint(m, Y.==sum(Y_sub[:,:,t] for t=1:k))

    @constraint(m, traceoneconstraint[t=1:k], sum(Y_sub[i,i,t] for i=1:n).==1.0)

    @constraint(m, imposeSDPSlice[t=1:k], ([1.0 U[:,t]' ;U[:,t] Y_sub[:,:,t] ]) in PSDCone());

    @constraint(m, imposeSDPSlice2[t=1:k, t2=(t+1):k], ([1.0 U[:,t]' U[:,t2]' ;U[:,t] Y_sub2[:,:,t, t] Y_sub2[:,:,t, t2]; U[:,t2] Y_sub2[:,:,t,t2]' Y_sub2[:,:,t2,t2]]) in PSDCone());

    @constraint(m, tracezeroconstraint[t=1:k, t2=(t+1):k], sum(Y_sub2[i,i,t, t2] for i=1:n).==0.0)

    @constraint(m, definediag[t=1:k], Y_sub[:,:,t].==Y_sub2[:,:,t,t])

    @constraint(m, strengthen4[i=1:n, t=1:k, t2=1:k], W[i,t,t2]==Y_sub2[i,i,t,t2])

    @constraint(m, McCormick5[i=1:n, t1=1:k, t2=1:k], W[i,t1,t2]>=-U_bound*(U[i,t1]+U[i,t2])-U_bound^2)
    @constraint(m, McCormick6[i=1:n, t1=1:k, t2=1:k], W[i,t1,t2]>=U_bound*(U[i,t1]+U[i,t2])-U_bound^2)
    @constraint(m, McCormick7[i=1:n, t1=1:k, t2=1:k], W[i,t1,t2]<=U_bound*(U[i,t1]-U[i,t2])+U_bound^2)
    @constraint(m, McCormick8[i=1:n, t1=1:k, t2=1:k], W[i,t1,t2]<=U_bound*(-U[i,t1]+U[i,t2])+U_bound^2)
    # @constraint(m, McCormick9[i=1:n, j=(i+1):n, t=1:k], Y_sub[i,j,t]<=U_bound^2)
    # @constraint(m, McCormick10[i=1:n, j=(i+1):n, t=1:k], Y_sub[i,j,t]>=-U_bound^2) #Product of U_it, U_jt but ||U||_2 \leq 1

    @constraint(m, orthogonality[t1=1:k, t2=1:k], sum(W[i,t1,t2] for i=1:n).==1.0*(t1==t2))
    @constraint(m, strengthen3[i=1:n, t=1:k], Y_sub[i,i,t].==W[i,t,t])

    @constraint(m, Symmetric(Matrix(1.0I, n, n)-Y) in PSDCone());
    @constraint(m, sigma>=1.0/(2.0*gamma)*sum(theta[i,i] for i=1:n))

    if imposeConstraints
        for t=1:Int(numEntriesSampled)
            @constraint(m, X[coords[entries[t],1], coords[entries[t],2]]==M[coords[entries[t],1], coords[entries[t],2]])
        end

        @objective(m, Min, sigma)
    else
        @objective(m, Min, sigma+0.5*sum((X[coords[entries[t],1], coords[entries[t],2]]-M[coords[entries[t],1], coords[entries[t],2]])^2 for t in 1:Int(numEntriesSampled)))
    end
    @constraint(m, Symmetric([Y X; X' theta]) in PSDCone());

    @suppress optimize!(m)
    @show objective_value(m)

    u,v,w=LinearAlgebra.svd(value.(Y))

    diag_v=zeros(n)
    diag_v[sortperm(v, rev=true)[1:k]].=1.0
    Y_rounded=u*Diagonal(diag_v)*u'
    @show sum(abs.(Y_rounded))

    m2=Model(with_optimizer(Mosek.Optimizer))
    @variable(m2, X2[1:n, 1:n])
    @objective(m2, Min, 1.0/(2.0*gamma)*sum(X2[i,j]^2 for i=1:n for j=1:n)+0.5*sum((X2[coords[entries[t],1], coords[entries[t],2]]-M[coords[entries[t],1], coords[entries[t],2]])^2 for t in 1:Int(numEntriesSampled)))
    @constraint(m2, X2.==Y_rounded*X2)
    @suppress optimize!(m2)
    X_best=value.(X2)
    @show sum(abs.(value.(X2)))

   return X_best, objective_value(m), objective_value(m2), value.(Y)
end

function getMatrixFrobeniusNorm_cuttingplanes_kelley(coords, entries, M, n, numEntriesSampled, k=sqrt(n), gamma=200*sqrt(n)/(numEntriesSampled), imposeConstraints=false)
    m=Model(Mosek.Optimizer)

    theIndices=zeros(n,n)
    for t=1:Int(numEntriesSampled)
        theIndices[coords[entries[t],1], coords[entries[t],2]]=1.0
    end

    @variable(m, Y[1:n, 1:n], PSD)
    @constraint(m, Symmetric(Matrix(1.0I, n, n)-Y) in PSDCone());
    @variable(m, theta>=-1e5)
    @objective(m, Min, theta)
    @constraint(m, sum(Y[i,i] for i=1:n)<=k)

    UB=1e12
    LB=-1e12
    max_iter=200
    Y_best=zeros(n,n)
    for iter=1:max_iter

        @suppress optimize!(m)

        LB=objective_value(m)

        # Solve subproblem for fixed Y, in alpha
        m2=Model(Ipopt.Optimizer)
        #set_start_value.(all_variables(m2), value.(all_variables(m2)))

        @variable(m2, alpha[1:n, 1:n])
        @constraint(m2, alpha.*(ones(n,n)-theIndices).==0.0) #if not in set of indices then eliminate variable
        @objective(m2, Max, -0.5*sum((alpha[i,j]+M[i,j])^2*theIndices[i,j] for i=1:n for j=1:n)-(gamma/2.0)*Compat.dot(alpha*alpha', value.(Y)))
        set_start_value.(alpha, rand(n,n))
        @suppress optimize!(m2)
        alpha_t=value.(alpha)
        # Get UB and cut from second-stage optimization problem
        f_Y=objective_value(m2)+0.5*sum((M.*theIndices).^2)
        # Impose constraint
        @constraint(m, theta>=f_Y-gamma/2*Compat.dot(alpha_t*alpha_t', Y-value.(Y)))
        #set_start_value.(all_variables(m), value.(all_variables(m)))
        if UB>f_Y
            Y_best=value.(Y)
            UB=f_Y
        end
        @show UB, LB, iter
        if abs.(UB-LB)<1e-3
            iter=max_iter
            break
        end
    end

    u,v,w=LinearAlgebra.svd(Y_best)

    diag_v=zeros(n)
    diag_v[sortperm(v, rev=true)[1:k]].=1.0
    Y_rounded=u*Diagonal(diag_v)*u'

    # greedy rounding

    m2=Model(with_optimizer(Mosek.Optimizer))
    @variable(m2, X2[1:n, 1:n])
    @objective(m2, Min, 1.0/(2.0*gamma)*sum(X2[i,j]^2 for i=1:n for j=1:n)+sum((X2[coords[entries[t],1], coords[entries[t],2]]-M[coords[entries[t],1], coords[entries[t],2]])^2 for t in 1:Int(numEntriesSampled)))
    @constraint(m2, X2.==Y_rounded*X2)
    @suppress optimize!(m2)
    X_best=value.(X2)

   return X_best, LB, objective_value(m2)
end

function getMatrixFrobeniusNorm_cuttingplanes_inout(coords, entries, M, n, numEntriesSampled, k=sqrt(n), gamma=200*sqrt(n)/(numEntriesSampled), stabilizationPoint=k/n*Diagonal(ones(n)))
    m=Model(Mosek.Optimizer)
    set_optimizer_attribute(m, "MSK_DPAR_INTPNT_CO_TOL_PFEAS", 1e-10)
    set_optimizer_attribute(m, "MSK_DPAR_INTPNT_CO_TOL_DFEAS", 1e-10)

    theIndices=zeros(n,n)
    for t=1:Int(numEntriesSampled)
        theIndices[coords[entries[t],1], coords[entries[t],2]]=1.0
    end

    @variable(m, Y[1:n, 1:n], PSD)
    @constraint(m, Symmetric(Matrix(1.0I, n, n)-Y) in PSDCone());
    @variable(m, theta>=-1e4)
    @objective(m, Min, theta)
    @constraint(m, sum(Y[i,i] for i=1:n)<=k)

    UB=1e12
    LB=-1e12
    rootStabilizationTrick = :inOut
     #Use the solution to the SDP relaxation instead, where appropriate
    rootCutsSense=1
    ε= 1e-10
    λ = (rootStabilizationTrick == :inOut) ? .1 : 1.
    δ = (rootStabilizationTrick == :inOut || rootStabilizationTrick == :twoEps) ? 2*ε : 0.
    rootCutsLim=1000
    rootCutCount = 0
    oaRootCutCount=0
    consecutiveNonImprov_1 = 0
    consecutiveNonImprov_2 = 0
    # theCutPool=Cut[]

    Y_best=zeros(n,n)
    for epoch in 1:1000
      #@show epoch, LB, UB, oaRootCutCount
      @suppress optimize!(m)
      Ystar=value.(Y)
      stabilizationPoint += Ystar; stabilizationPoint /= 2

      if LB >= objective_value(m) - eps()
        if consecutiveNonImprov_1 == 5
          consecutiveNonImprov_2 += 1
        else
          consecutiveNonImprov_1 += 1
        end
      else
        if consecutiveNonImprov_1 < 5
          consecutiveNonImprov_1 = 0
        end
        consecutiveNonImprov_2 = 0
      end
      LB = max(LB, objective_value(m))

      if consecutiveNonImprov_1 == 5
        λ = 1
      elseif consecutiveNonImprov_2 == 5
        δ = 0.
      end

      Y0 = λ*Ystar + (1-λ)*stabilizationPoint .+ δ*Diagonal(ones(n))

      # Perform an SVD step here to ensure still in convex hull
      u,v,=svd(Y0)
      v=v*k/sum(v)
      v.=min.(ones(n), max.(v, zeros(n)))

      Y0=u*Diagonal(v)*u'

      # Get a cut by solving the dual subproblem
      # Solve subproblem for fixed Y, in alpha
      m2=Model(Mosek.Optimizer)
      @variable(m2, alpha[1:n, 1:n])
      @constraint(m2, alpha.*(ones(n,n)-theIndices).==0.0) #if not in set of indices then eliminate variable
      @objective(m2, Max, -0.5*sum((alpha[i,j]+M[i,j])^2*theIndices[i,j] for i=1:n for j=1:n)-(gamma/2.0)*Compat.dot(alpha*alpha', Y0))
      set_start_value.(alpha, rand(n,n))
      @suppress optimize!(m2)
      alpha_t=value.(alpha)
      # Get UB and cut from second-stage optimization problem
      f_Y=objective_value(m2)+0.5*sum((M.*theIndices).^2)
      @show f_Y
      @show sum(diag(Y0))
      if f_Y<UB
          Y_best=Y0
          UB=f_Y
      end
      LB=max(LB, objective_value(m))
      # Impose constraint
      @constraint(m, theta>=f_Y-gamma/2*Compat.dot(alpha_t*alpha_t', Y-Y0))
      # Add cut to pool (add in later if it works)
      @show LB, UB, epoch
      if abs(UB-LB) <= 1e-2 || consecutiveNonImprov_2 >= 10
          break
      end
    end

    u,v,w=LinearAlgebra.svd(Y_best)

    diag_v=zeros(n)
    diag_v[sortperm(v, rev=true)[1:k]].=1.0
    Y_rounded=u*Diagonal(diag_v)*u'

    # greedy rounding

    m2=Model(with_optimizer(Mosek.Optimizer))
    @variable(m2, X2[1:n, 1:n])
    @objective(m2, Min, 1.0/(2.0*gamma)*sum(X2[i,j]^2 for i=1:n for j=1:n)+sum((X2[coords[entries[t],1], coords[entries[t],2]]-M[coords[entries[t],1], coords[entries[t],2]])^2 for t in 1:Int(numEntriesSampled)))
    @constraint(m2, X2.==Y_rounded*X2)
    @suppress optimize!(m2)
    X_best=value.(X2)

   return X_best, LB, objective_value(m2)
end

max_stars_repo_path: Matrix Completion/convexpenalties.jl

max_stars_repo_name: ryancorywright/MixedProjectionSoftware

max_issues_repo_path: Matrix Completion/convexpenalties.jl

max_issues_repo_name: ryancorywright/MixedProjectionSoftware

max_forks_repo_path: Matrix Completion/convexpenalties.jl

max_forks_repo_name: ryancorywright/MixedProjectionSoftware


text:
# Albrecher, H., Mayer, P., & Schoutens, W. T. J. (2007)
# The Little Heston Trap
# Wilmott Magazine, January issue, 83-92.

function analyticinterval(m::Heston)
  # an approximation of the analytic interval is the
  # area between the roots of p(x) = ax^2+bx+c
  a = (1.0 - m.ρ^2) * m.σ^2
  b = 2.0 * m.ρ * m.σ * m.κ - m.σ^2
  c = -m.κ^2
  A = -0.5 * b / a
  B = sqrt(0.25 * b^2 / a^2 - c / a)
  A - B, A + B
end

function logmgf(m::Heston, x0, T, u)
  v0 = x0[2]
  σ2 = m.σ * m.σ
  a0 = m.ρ * m.σ * u - m.κ
  a1 = sqrt(a0 * a0 + σ2 * u * (1 - u))
  a2 = m.ρ * m.σ * u - m.κ
  a3 = -(a2 + a1)
  a4 = a3 / (a1 - a2)
  a5 = exp(-a1 * T)
  a6 = 1 - a4 * a5
  ((a3 * T - 2 * log(a6 / (1 - a4))) * m.θ * m.κ + v0 * a3 * (1 - a5) / a6) / σ2
end


text:
export Result, Field, FieldTM, FieldTE, FieldAll
export AxisX, AxisY, AxisComponent

const AxisX = Axis{:x}
const AxisY = Axis{:y}
const AxisComponent = Axis{:component, Array{Symbol,1}}

abstract type Result{T,N} <: AbstractArray{T,N} end
abstract type Field{T,N} <: Result{T,N} end

Base.size(A::Result) = size(A.data)
Base.IndexStyle(A::Result) = IndexStyle(A.data)

AxisArrays.axisdim(A::Result, ax) = axisdim(A.data, ax)
AxisArrays.axes(A::Result, i...) = AxisArrays.axes(A.data, i...)
AxisArrays.axisnames(A::Result) = axisnames(A.data)
AxisArrays.axisvalues(A::Result) = axisvalues(A.data)

# function Base.similar(x::Result, ::Type{S}, dims::NTuple{N,Int}) where {S,N}
#     if size(x) == dims
#         data = similar(AxisArray(x),S)
#         similar_helper(x,data,x.params)
#     else
#         similar(AxisArray(x),S,dims)
#     end
# end
#
# Base.similar(x::Result, ax1::Axis, axs::Axis...) = similar(x, eltype(x), ax1, axs...)
#
# function Base.similar(x::Result, ::Type{S}, ax1::Axis, axs::Axis...) where S
#     if ndims(x) == 1+length(axs)
#         data = similar(AxisArray(x), ax1, axs...)
#         similar_helper(x, data, x.params)
#     else
#         similar(AxisArray(x), S, ax1, axs...)
#     end
# end

@inline @Base.propagate_inbounds Base.getindex(A::Result, i...) =
  getindex(A.data, i...)
@inline @Base.propagate_inbounds Base.getindex(A::Result, i::Int...) =
  getindex(A.data, i...)
@inline @Base.propagate_inbounds Base.setindex!(A::Result, val, i...) =
  setindex!(A.data, val, i...)
@inline @Base.propagate_inbounds Base.setindex!(A::Result{T}, val::T, i::Int...) where T =
  setindex!(A.data, val, i...)

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

struct FieldTM <: Field{Complex,3}
    grid::Grid{2}
    ω::Complex
    data::AxisArray
end

function FieldTM(grid::Grid{2}, ω::Number, data::Array{<:Complex,3})
    return FieldTM(grid, Complex(ω), AxisArray(data, AxisX(xc(grid)), AxisY(yc(grid)), AxisComponent([:Ez, :Hx, :Hy])))
end

function FieldTM(grid::Grid{2}, ω::Number, Ez::Array{<:Complex,1}, Hx::Array{<:Complex,1}, Hy::Array{<:Complex,1})
  	sz = size(grid);
    return FieldTM(grid, ω, cat(reshape(Ez, sz), reshape(Hx, sz), reshape(Hy, sz), dims=3))
end

struct FieldTE <: Field{Complex,3}
    grid::Grid{2}
    ω::Complex
    data::AxisArray
end

function FieldTE(grid::Grid{2}, ω::Number, data::Array{<:Complex,3})
    return FieldTE(grid, Complex(ω), AxisArray(data, AxisX(xc(grid)), AxisY(yc(grid)), AxisComponent([:Hz, :Ex, :Ey])))
end

function FieldTE(grid::Grid{2}, ω::Number, Hz::Array{<:Complex,1}, Ex::Array{<:Complex,1}, Ey::Array{<:Complex,1})
  	sz = size(grid);
    return FieldTE(grid, ω, cat(reshape(Hz, sz), reshape(Ex, sz), reshape(Ey, sz), dims=3))
end

struct FieldAll <: Field{Complex,3}
    grid::Grid{2}
    ω::Complex
    data::AxisArray
end

function FieldAll(grid::Grid{2}, ω::Number, data::Array{<:Complex,3})
    return FieldAll(grid, Complex(ω), AxisArray(data, AxisX(xc(grid)), AxisY(yc(grid)), AxisComponent([:Ex, :Ey, :Ez, :Hx, :Hy, :Hz])))
end

function FieldAll(grid::Grid{2}, ω::Number, Ex::Array{<:Complex,1}, Ey::Array{<:Complex,1}, Ez::Array{<:Complex,1}, Hx::Array{<:Complex,1}, Hy::Array{<:Complex,1}, Hz::Array{<:Complex,1})
  	sz = size(grid);
    return FieldAll(grid, ω, cat(reshape(Ex, sz), reshape(Ey, sz), reshape(Ez, sz), reshape(Hx, sz), reshape(Hy, sz), reshape(Hz, sz), dims=3))
end

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

struct Flux <: Field{Float,3}
    grid::Grid{2}
    ω::Complex
    data::AxisArray
end

function Flux(grid::Grid{2}, ω::Number, data::Array{Float,3})
    return Flux(grid, Complex(ω), AxisArray(data, AxisX(xc(grid)), AxisY(yc(grid)), AxisComponent([:Sx, :Sy])))
end

function Flux(grid::Grid{2}, ω::Number, Sx::Array{Float,2}, Sy::Array{Float,2})
    return Flux(grid, ω, cat(3, Sx, Sy))
end

# poynt = poynting(field_dipole);
# ax = plot_field(field_dipole, funcz=real)
# xinterval = -2..2
# yinterval = 1..3
# sub_poynt = poynt[xinterval,yinterval,:]
# x_vals = AxisArrays.axisvalues(sub_poynt)[1]
# y_vals = AxisArrays.axisvalues(sub_poynt)[2]
# skip = 9
# quiver(x_vals[1:skip:end], y_vals[1:skip:end], sub_poynt[1:skip:end,1:skip:end,:Sx]', sub_poynt[1:skip:end,1:skip:end,:Sy]')


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

using Markdown
using InteractiveUtils

# ╔═╡ 6fd616f6-1ea1-11eb-3814-8bfb4a096c49
using Pkg, DrWatson

# ╔═╡ 6ffe2628-1ea1-11eb-24ea-57f985146a72
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Distributions, StatsPlots, LaTeXStrings
end;

# ╔═╡ ed174bc4-1ea0-11eb-1e2f-a32874cec549
md"## Listing 3.9"

# ╔═╡ 70179b62-1ea1-11eb-13ce-7fb4d54bad4b
begin
	dist = TriangularDist(0,2,1)
	xGrid = 0:0.01:2
	uGrid = 0:0.01:1
end

# ╔═╡ a715f0ac-2034-11eb-029c-9be9db193aed
begin
	p1 = plot( xGrid, pdf.(dist,xGrid), c=:blue,
			xlims=(0,2), ylims=(0,1.1),
			xlabel="x", ylabel="f(x)")
	p2 = plot( xGrid, cdf.(dist,xGrid), c=:blue,
			xlims=(0,2), ylims=(0,1),
			xlabel="x", ylabel="F(x)")
	p3 = plot( uGrid,quantile.(dist,uGrid), c=:blue,
			xlims=(0,1), ylims=(0,2),
			xlabel="u", ylabel=(L"F^{-1}(u)"))

	plot(p1, p2, p3, legend=false, layout=(1,3), size=(1200, 400))
end

# ╔═╡ 70182cd8-1ea1-11eb-094d-8d0d49cf15f3
md"## End of listing 3.9"

# ╔═╡ Cell order:
# ╟─ed174bc4-1ea0-11eb-1e2f-a32874cec549
# ╠═6fd616f6-1ea1-11eb-3814-8bfb4a096c49
# ╠═6ffe2628-1ea1-11eb-24ea-57f985146a72
# ╠═70179b62-1ea1-11eb-13ce-7fb4d54bad4b
# ╠═a715f0ac-2034-11eb-029c-9be9db193aed
# ╟─70182cd8-1ea1-11eb-094d-8d0d49cf15f3


text:
"""
    ess_rhat(chains::Chains; duration=compute_duration, kwargs...)

Estimate the effective sample size and the potential scale reduction.

ESS per second options include `duration=MCMCChains.compute_duration` (the default)
and `duration=MCMCChains.wall_duration`.
"""
function MCMCDiagnosticTools.ess_rhat(
    chains::Chains;
    sections = _default_sections(chains), duration = compute_duration, kwargs...
)
    # Subset the chain
    _chains = Chains(chains, _clean_sections(chains, sections))

    # Estimate the effective sample size and rhat
    ess, rhat = MCMCDiagnosticTools.ess_rhat(_chains.value.data; kwargs...)

    # Calculate ESS/minute if available
    dur = duration(chains)

    # Convert to NamedTuple
    nt = if dur === missing
        merge((parameters = names(_chains),), (ess = ess, rhat = rhat))
    else
        merge((parameters = names(_chains),), (ess = ess, rhat = rhat, ess_per_sec=ess/dur))
    end

    return ChainDataFrame("ESS", nt)
end


text:
struct Sprite
    x::Int
    y::Int

    width::Int
    height::Int

    offsetX::Int
    offsetY::Int
    realWidth::Int
    realHeight::Int

    surface::Cairo.CairoSurface
    filename::String
end

mutable struct SpriteHolder
    sprite::Sprite

    readtime::Float64
    path::String
end

function loadSprites(metaFn::String, spritesFn::String)
    surface = open(Cairo.read_from_png, spritesFn)
    fh = open(metaFn)

    res = Dict{String, SpriteHolder}()

    # Mostly useless information
    read(fh, Int32)
    Maple.readString(fh)
    read(fh, Int32)

    count = read(fh, Int16)
    for i in 1:count
        dataFile = Maple.readString(fh)
        sprites = read(fh, Int16)

        for j in 1:sprites
            path = replace(Maple.readString(fh), "\\" => "/")

            sprite = Sprite(
                read(fh, Int16),
                read(fh, Int16),

                read(fh, Int16),
                read(fh, Int16),

                read(fh, Int16),
                read(fh, Int16),
                read(fh, Int16),
                read(fh, Int16),

                surface,
                spritesFn
            )

            res[path] = SpriteHolder(
                sprite,

                time(),
                spritesFn
            )
        end
    end

    close(fh)

    return res
end

function findTextureAnimations(texture::String, sprites::Dict{String, SpriteHolder}; maxPadding=7)::Array{String, 1}
    textures = String[]

    for i in 1:maxPadding
        testFrame = texture * "0"^i

        if haskey(sprites, testFrame) && sprites[testFrame].sprite.width != 0 && sprites[testFrame].sprite.height != 0
            for j in 0:10^i - 1
                frame = texture * lpad(j, i, "0")

                if haskey(sprites, frame) && sprites[frame].sprite.width != 0 && sprites[frame].sprite.height != 0
                    push!(textures, frame)

                else
                    break
                end
            end
        end
    end

    return textures
end

max_stars_repo_name: asdfsdfadfasfasd/Ahorn

max_issues_repo_name: asdfsdfadfasfasd/Ahorn

max_forks_repo_name: asdfsdfadfasfasd/Ahorn


text:
## Example:
## DGP - Bivariate VAR(1) Model from Kilian, RESTAT, 1998
# B11 set to 0.01
using VectorAutoregressions, Plots
plotly()

const T,K = 1000,2
const H = 24
const nrep = 1000
const p = 1
const intercept = false
const burnin = 100

srand(1234)
mR  = [0.9 0.01; 0.5 0.5]             #  Coefficients of the 1st lag B11 set to 0.01
mΣ_true = [1 0.3; 0.3 1]              #  vcv of u
mP = cholfact(mΣ_true)[:L]            #  Cholesky PP'= Σ
mI = eye(K)
mIRF_true = At_mul_Bt(mP,mI)[:]       # Initialize structural IRF
# True IRF
for h=1:H
    vIRF = At_mul_Bt(mP,mR^h)[:]
    mIRF_true = hcat(mIRF_true,vIRF)
end

# Generate data from the VAR(1) DGP
y = zeros(T+burnin,K)
gen_var1_data!(y,mR,mP,burnin)

# Fit VAR(1) to data
V = VAR(y,p,intercept)
mIRFb = IRFs_b(V,H,nrep,intercept)
mIRFa = IRFs_a(V,H,intercept)

# Plot true irf vs model irf + asy ci
pIRF_asy = plot(layout = grid(K,K));
[plot!(pIRF_asy, [mIRF_true[i,:] mIRFa.IRF[i,:] mIRFa.CI.CIl[i,:] mIRFa.CI.CIh[i,:]], color = ["blue" "red" "red" "red"],
line = [:solid :solid :dash :dash], legend = false, label = ["true" "model"], subplot = i) for i in 1:K^2]
gui(pIRF_asy)

# Plot true irf vs model irf + bootstraped ci
pIRF_boot = plot(layout = grid(K,K));
[plot!(pIRF_boot, [mIRF_true[i,:] mIRFb.CI.CIl[i,:] mIRFb.IRF[i,:] mIRFb.CI.CIh[i,:]], color = ["blue" "black" "black" "black"],
line = [:solid :solid :dash :dash], legend = false, label = ["true" "model"], subplot = i) for i in 1:K^2]
gui(pIRF_boot)


text:
# Test parameter function methods
@testset "Parameter Function Methods" begin
    # setup the needed info
    m = InfiniteModel()
    @infinite_parameter(m, t in [0, 1])
    @infinite_parameter(m, x[1:2] in [-1, 1])
    func = ParameterFunction(sin, IC.VectorTuple(t), [1], [1])
    object = ParameterFunctionData(func, "test")
    idx = ParameterFunctionIndex(1)
    fref = ParameterFunctionRef(m, idx)
    gvref = GeneralVariableRef(m, 1, ParameterFunctionIndex)
    bad_vref = ParameterFunctionRef(m, ParameterFunctionIndex(-1))
    # JuMP.owner_model
    @testset "JuMP.owner_model" begin
        @test owner_model(fref) === m
        @test owner_model(gvref) === m
    end
    # JuMP.index
    @testset "JuMP.index" begin
        @test index(fref) == idx
        @test index(gvref) == idx
    end
    # dispatch_variable_ref
    @testset "dispatch_variable_ref" begin
        @test dispatch_variable_ref(m, idx) == fref
        @test dispatch_variable_ref(gvref) == fref
    end
    # _add_data_object
    @testset "_add_data_object" begin
        @test InfiniteOpt._add_data_object(m, object) == idx
    end
    # _data_dictionary
    @testset "_data_dictionary" begin
        @test InfiniteOpt._data_dictionary(fref) === m.param_functions
        @test InfiniteOpt._data_dictionary(gvref) === m.param_functions
    end
    # JuMP.is_valid
    @testset "JuMP.is_valid" begin
        @test is_valid(m, fref)
        @test is_valid(m, gvref)
    end
    # _data_object
    @testset "_data_object" begin
        @test InfiniteOpt._data_object(fref) === object
        @test InfiniteOpt._data_object(gvref) === object
        @test_throws ErrorException InfiniteOpt._data_object(bad_vref)
    end
    # _core_variable_object
    @testset "_core_variable_object" begin
        @test InfiniteOpt._core_variable_object(fref) === func
        @test InfiniteOpt._core_variable_object(gvref) === func
    end
    # _object_numbers
    @testset "_object_numbers" begin
        @test InfiniteOpt._object_numbers(fref) == [1]
    end
    # _parameter_numbers
    @testset "_parameter_numbers" begin
        @test InfiniteOpt._parameter_numbers(fref) == [1]
    end
    # JuMP.name
    @testset "JuMP.name" begin
        @test name(fref) == "test"
        @test name(gvref) == "test"
        @test name(bad_vref) == ""
    end
    # JuMP.set_name
    @testset "JuMP.set_name" begin
        @test isa(set_name(fref, "new"), Nothing)
        @test name(fref) == "new"
        @test isa(set_name(gvref, "new2"), Nothing)
        @test name(fref) == "new2"
        @test_throws ErrorException set_name(bad_vref, "")
    end
    # _measure_dependencies
    @testset "_measure_dependencies" begin
        @test InfiniteOpt._measure_dependencies(fref) == MeasureIndex[]
        @test InfiniteOpt._measure_dependencies(gvref) == MeasureIndex[]
    end
    # _constraint_dependencies
    @testset "_constraint_dependencies" begin
        @test InfiniteOpt._constraint_dependencies(fref) == InfOptConstraintIndex[]
        @test InfiniteOpt._constraint_dependencies(gvref) == InfOptConstraintIndex[]
    end
    # _semi_infinite_variable_dependencies
    @testset "_semi_infinite_variable_dependencies" begin
        @test InfiniteOpt._semi_infinite_variable_dependencies(fref) == SemiInfiniteVariableIndex[]
        @test InfiniteOpt._semi_infinite_variable_dependencies(gvref) == SemiInfiniteVariableIndex[]
    end
    # _derivative_dependencies
    @testset "_derivative_dependencies" begin
        @test InfiniteOpt._derivative_dependencies(fref) == DerivativeIndex[]
        @test InfiniteOpt._derivative_dependencies(gvref) == DerivativeIndex[]
    end
    # raw_parameter_refs
    @testset "raw_parameter_refs" begin
        @test raw_parameter_refs(fref) == IC.VectorTuple(t)
        @test raw_parameter_refs(gvref) == IC.VectorTuple(t)
    end
    # parameter_refs
    @testset "parameter_refs" begin
        @test parameter_refs(fref) == (t, )
        @test parameter_refs(gvref) == (t, )
    end
    # parameter_list
    @testset "parameter_list" begin
        @test parameter_list(fref) == [t]
        @test parameter_list(gvref) == [t]
    end
    # raw_function
    @testset "raw_function" begin
        @test raw_function(fref) == sin
        @test raw_function(gvref) == sin
    end
    # call_function
    @testset "call_function" begin
        @test call_function(fref, 0) == 0
        @test call_function(gvref, 0) == 0
    end
    # test used_by_semi_infinite_variable
    @testset "used_by_semi_infinite_variable" begin
        @test !used_by_semi_infinite_variable(fref)
        push!(InfiniteOpt._semi_infinite_variable_dependencies(fref),
              SemiInfiniteVariableIndex(1))
        @test used_by_semi_infinite_variable(fref)
        empty!(InfiniteOpt._semi_infinite_variable_dependencies(fref))
    end
    # test used_by_derivative
    @testset "used_by_derivative" begin
        @test !used_by_derivative(fref)
        push!(InfiniteOpt._derivative_dependencies(fref), DerivativeIndex(1))
        @test used_by_derivative(fref)
        empty!(InfiniteOpt._derivative_dependencies(fref))
    end
    # test used_by_measure
    @testset "used_by_measure" begin
        @test !used_by_measure(fref)
        push!(InfiniteOpt._measure_dependencies(fref), MeasureIndex(1))
        @test used_by_measure(fref)
        empty!(InfiniteOpt._measure_dependencies(fref))
    end
    # test used_by_constraint
    @testset "used_by_constraint" begin
        @test !used_by_constraint(fref)
        push!(InfiniteOpt._constraint_dependencies(fref), InfOptConstraintIndex(1))
        @test used_by_constraint(fref)
        empty!(InfiniteOpt._constraint_dependencies(fref))
    end
    # test is_used
    @testset "is_used" begin
        @test !is_used(fref)
        push!(InfiniteOpt._constraint_dependencies(fref), InfOptConstraintIndex(1))
        @test is_used(fref)
        empty!(InfiniteOpt._constraint_dependencies(fref))
    end
    # test _update_param_var_mapping
    @testset "_update_param_var_mapping" begin
        @test InfiniteOpt._update_param_var_mapping(fref, IC.VectorTuple(t)) isa Nothing
        @test used_by_parameter_function(t)
        empty!(InfiniteOpt._parameter_function_dependencies(t))
    end
    # test _delete_data_object
    @testset "_delete_data_object" begin
        @test InfiniteOpt._delete_data_object(fref) isa Nothing
        @test length(InfiniteOpt._data_dictionary(fref)) == 0
        @test !is_valid(m, fref)
    end
    # test build_parameter_function
    @testset "build_parameter_function" begin
        f3(ts, xs, a...; b...) = 42
        # test errors
        @test_throws ErrorException build_parameter_function(error, sin, t, bob = 1)
        @test_throws ErrorException build_parameter_function(error, sin, x[1])
        @test_throws ErrorException build_parameter_function(error, sin, (t, x))
        @test_throws ErrorException build_parameter_function(error, sin, 2)
        # test normal
        @test build_parameter_function(error, (a, b) -> 2, (t, x)) isa ParameterFunction
        @test build_parameter_function(error, f3, (t, x)) isa ParameterFunction
        @test build_parameter_function(error, (ts, xs) -> f3(ts, xs, 1), (t, x)) isa ParameterFunction
        @test build_parameter_function(error, (ts, xs) -> f3(ts, xs, d = 1), (t, x)) isa ParameterFunction
        @test build_parameter_function(error, sin, t) isa ParameterFunction
    end
    # test add_parameter_function
    @testset "add_parameter_function" begin
        # test does not belong
        @infinite_parameter(InfiniteModel(), t2 in [0, 1])
        func = build_parameter_function(error, (a, b) -> 2, (t2, x))
        @test_throws VariableNotOwned{GeneralVariableRef} add_parameter_function(m, func)
        # test normal
        func = build_parameter_function(error, (a, b) -> 2, (t, x))
        fref = GeneralVariableRef(m, 2, ParameterFunctionIndex)
        @test add_parameter_function(m, func, "test") == fref
        @test name(fref) == "test"
        @test parameter_refs(fref) == (t, x)
        @test used_by_parameter_function(t)
        # test default name
        func = build_parameter_function(error, cos, t)
        fref = GeneralVariableRef(m, 3, ParameterFunctionIndex)
        @test add_parameter_function(m, func) == fref
        @test name(fref) == "cos"
    end
    # test parameter_function
    @testset "parameter_function" begin
        f4(ts, xs, a...; b...) = 42
        # test normal
        @test parameter_function(sin, t) isa GeneralVariableRef
        @test parameter_function(sin, t, name = "name") isa GeneralVariableRef
        @test parameter_function((ts, xs) -> f4(ts, xs, 1, d = 1), (t, x)) isa GeneralVariableRef
        # test errors
        @test_throws ErrorException parameter_function(sin, (t, x))
    end
    # test making other objects
    @testset "Other Objects" begin
        f = parameter_function((a,b) -> 2, (t, x))
        # test making semi_infinite variable
        d = Dict{Int, Float64}(1 => 0)
        @test add_variable(m, build_variable(error, f, d)) isa GeneralVariableRef
        # test making derivative
        @test deriv(f, t) isa GeneralVariableRef
    end
end

# Test the parameter function methods for macro definition
@testset "Parameter Function Macro" begin
    # setup
    m = InfiniteModel()
    @infinite_parameter(m, t in [0, 1])
    @infinite_parameter(m, x[1:2] in [0, 1])
    f5(t, x, a...; b...) = 42
    # test _process_func_expr
    @testset "_process_func_expr" begin
        # test normal
        @test InfiniteOpt._process_func_expr(error, :(f(t, x))) == (esc(:f), esc(:(t, x)))
        anon = :((t, x) -> f(t, x, 1, d = 1))
        @test InfiniteOpt._process_func_expr(error, anon) == (esc(anon), esc(:(t,x)))
        # test errors
        @test_throws ErrorException InfiniteOpt._process_func_expr(error, :(f(t, d = 2)))
        @test_throws ErrorException InfiniteOpt._process_func_expr(error, :(f(x, t; d = 2)))
        @test_throws ErrorException InfiniteOpt._process_func_expr(error, :(f >= 2))
        @test_throws ErrorException InfiniteOpt._process_func_expr(error, :(f[t, x]))
    end
    # test @parameter_function
    @testset "@parameter_function" begin
        # test errors
        @test_macro_throws ErrorException @parameter_function(m)
        @test_macro_throws ErrorException @parameter_function(m, func = f5)
        @test_macro_throws ErrorException @parameter_function(m, y == sin(t), Int)
        @test_macro_throws ErrorException @parameter_function(m, [1:2])
        @test_macro_throws ErrorException @parameter_function(m, "a$(1)" == f5(t, x))
        @test_macro_throws ErrorException @parameter_function(m, a[m = 1:2] == f5(t, x))
        @test_macro_throws ErrorException @parameter_function(Model(), sin(t))
        @test_macro_throws ErrorException @parameter_function(m, [a...] == sin(t))
        # test anonymous singular
        idx = 1
        ref = GeneralVariableRef(m, idx, ParameterFunctionIndex)
        @test @parameter_function(m, f5(t, x), base_name = "a") == ref
        @test name(ref) == "a"
        @test raw_function(ref) == f5
        @test parameter_refs(ref) == (t, x)
        idx += 1
        ref = GeneralVariableRef(m, idx, ParameterFunctionIndex)
        @test @parameter_function(m, f5(t, x)) == ref
        @test name(ref) == "f5"
        @test raw_function(ref) == f5
        @test parameter_refs(ref) == (t, x)
        idx += 1
        # test anonymous single argument multi-dim
        refs = [GeneralVariableRef(m, idx + i, ParameterFunctionIndex) for i in 0:1]
        @test @parameter_function(m, [1:2] == f5(t, x)) == refs
        @test parameter_refs(refs[1]) == (t, x)
        @test raw_function(refs[2]) == f5
        @test name.(refs) == ["f5", "f5"]
        idx += 2
        refs = [GeneralVariableRef(m, idx + i, ParameterFunctionIndex) for i in 0:1]
        @test @parameter_function(m, [i = 1:2; i >= 1] == (sin, cos)[i](t)) isa JuMPC.SparseAxisArray
        @test parameter_refs(refs[1]) == (t,)
        @test raw_function(refs[2]) == cos
        @test raw_function(refs[1]) == sin
        @test name(refs[1]) == "sin"
        @test name(refs[2]) == "cos"
        idx += 2
        # test explicit single
        ref = GeneralVariableRef(m, idx, ParameterFunctionIndex)
        @test @parameter_function(m, a == f5(t, x), base_name = "bob") == ref
        @test name(ref) == "bob"
        @test raw_function(ref) == f5
        @test parameter_refs(ref) == (t, x)
        idx += 1
        ref = GeneralVariableRef(m, idx, ParameterFunctionIndex)
        @test @parameter_function(m, c == (t, x) -> f5(t, x, 2, d = 1)) == ref
        @test name(ref) == "c"
        @test raw_function(ref) != f5
        @test raw_function(ref)(1, [1, 1]) == 42
        @test parameter_refs(ref) == (t, x)
        idx += 1
        # test explicit multi-dim
        refs = [GeneralVariableRef(m, idx + i, ParameterFunctionIndex) for i in 0:1]
        @test @parameter_function(m, d[1:2] == f5(t, x)) == refs
        @test parameter_refs(refs[1]) == (t, x)
        @test raw_function(refs[2]) == f5
        @test name.(refs) == ["d[1]", "d[2]"]
        idx += 2
        refs = [GeneralVariableRef(m, idx + i, ParameterFunctionIndex) for i in 0:1]
        @test @parameter_function(m, e[1:2] == (t, x) -> f5(t, x, 2; s = 1)) == refs
        @test parameter_refs(refs[1]) == (t, x)
        @test raw_function(refs[2]) != f5
        @test name.(refs) == ["e[1]", "e[2]"]
        idx += 2
    end
end

# Test _all_function_variables
@testset "_all_function_variables" begin
    # initialize model and references
    m = InfiniteModel()
    @infinite_parameter(m, par in [0, 1])
    @variable(m, inf, Infinite(par))
    @variable(m, pt, Point(inf, 0))
    @variable(m, finite)
    data = TestData(par, 0, 5)
    meas = Measure(finite, data, Int[], Int[], true)
    object = MeasureData(meas, "test")
    mindex = MeasureIndex(1)
    @test InfiniteOpt._add_data_object(m, object) == mindex
    meas = InfiniteOpt._make_variable_ref(m, mindex)
    dinf = @deriv(inf, par)
    # test for variable reference
    @testset "Variable" begin
        @test InfiniteOpt._all_function_variables(par) == [par]
        @test InfiniteOpt._all_function_variables(inf) == [inf]
        @test InfiniteOpt._all_function_variables(meas) == [meas]
        @test InfiniteOpt._all_function_variables(dinf) == [dinf]
    end
    # test for GenericAffExpr
    @testset "AffExpr" begin
        # make expressions
        aff1 = meas + 2par + finite - dinf
        aff2 = zero(GenericAffExpr{Float64, GeneralVariableRef})
        # test expressions
        @test isempty(setdiff(InfiniteOpt._all_function_variables(aff1),
                              [meas, par, finite, dinf]))
        @test InfiniteOpt._all_function_variables(aff2) == GeneralVariableRef[]
    end
    # test for GenericQuadExpr
    @testset "QuadExpr" begin
        # make expressions
        quad1 = pt^2 + inf * pt - meas + 2par + finite - dinf
        quad2 = pt^2 + inf * pt
        quad3 = zero(GenericQuadExpr{Float64, GeneralVariableRef})
        # test expressions
        @test isempty(setdiff(InfiniteOpt._all_function_variables(quad1),
                      [meas, par, finite, dinf, pt, inf]))
        @test isempty(setdiff(InfiniteOpt._all_function_variables(quad2),
                              [pt, inf]))
        @test InfiniteOpt._all_function_variables(quad3) == GeneralVariableRef[]
    end
    # test for Array of expressions
    @testset "AbstractArray" begin
        # make expressions
        ex1 = [inf, pt]
        ex2 = [inf + pt, meas + pt]
        # test expressions
        @test isempty(setdiff(InfiniteOpt._all_function_variables(ex1),
                      [pt, inf]))
        @test isempty(setdiff(InfiniteOpt._all_function_variables(ex2),
                              [pt, inf, meas]))
    end
    # test backup
    @testset "Fallback" begin
        @variable(Model(), x)
        @test_throws ErrorException InfiniteOpt._all_function_variables(x)
    end
end

# Test comparisons
@testset "Comparisons" begin
    # initialize model and references
    m = InfiniteModel()
    @infinite_parameter(m, par in [0, 1])
    @infinite_parameter(m, par2 in [0, 1])
    @variable(m, inf, Infinite(par))
    @variable(m, inf2, Infinite(par, par2))
    @variable(m, pt, Point(inf, 0))
    @variable(m, finite)
    # test AffExpr comparison
    @testset "Base.:(==) AffExpr" begin
        @test par + par2 + inf - 2 == par + (par2 + inf) - 2
        @test 0.25par- inf == 0.25par - inf
        @test inf + 3 - inf != par + inf
        @test inf + 3 - inf != inf + 3 - finite
    end
    # test QuadExpr comparison
    @testset "Base.:(==) QuadExpr" begin
        @test par * inf + par2 + inf - 2 == par * inf + (par2 + inf) - 2
        @test inf * inf - inf == inf * inf - inf
        @test inf * inf + 3 - inf != par * inf2 + inf
        @test par * par2 + inf * inf2 == par * par2 + inf * inf2
        @test par * par2 + inf * inf2 != par * par2
        @test par * par2 + 2 * inf * inf2 !=  par * par2 + inf * inf2
        @test par * inf + par2 + inf - 2 != par * inf + (par2 + inf) - 3
    end
end

# Test _object_numbers
@testset "_object_numbers" begin
    # initialize model and references
    m = InfiniteModel()
    @infinite_parameter(m, par in [0, 1])
    @infinite_parameter(m, pars[1:2] in [0, 1])
    @variable(m, inf, Infinite(par))
    @variable(m, inf2, Infinite(par, pars))
    @variable(m, pt, Point(inf, 0))
    @variable(m, finite)
    var = build_variable(error, inf2, Dict{Int, Float64}(1 => 0.5), check = false)
    red = add_variable(m, var)
    # test for finite variable reference
    @testset "FiniteVariable" begin
        @test InfiniteOpt._object_numbers(pt) == []
        @test InfiniteOpt._object_numbers(finite) == []
    end
    # test for infinite variable reference
    @testset "InfiniteVariable" begin
        @test InfiniteOpt._object_numbers(inf) == [1]
    end
    # test for parameter reference
    @testset "Parameter" begin
        @test InfiniteOpt._object_numbers(par) == [1]
        @test InfiniteOpt._object_numbers(pars[1]) == [2]
    end
    # test for semi-infinite variable reference
    @testset "SemiInfiniteInfinite" begin
        @test InfiniteOpt._object_numbers(red) == [2]
    end
    # test for GenericAffExpr
    @testset "AffExpr" begin
        # make expressions
        aff1 = inf + inf2 + pt - 3
        aff2 = pt + finite - 2
        # test expressions
        @test sort!(InfiniteOpt._object_numbers(aff1)) == [1, 2]
        @test InfiniteOpt._object_numbers(aff2) == []
    end
    # test for GenericQuadExpr
    @testset "QuadExpr" begin
        # make expressions
        quad1 = inf * inf2 + inf + inf2 + pt - 3 - par
        quad2 = pt * pt + pt + finite - 2
        # test expressions
        @test sort!(InfiniteOpt._object_numbers(quad1)) == [1, 2]
        @test InfiniteOpt._object_numbers(quad2) == []
    end
end

# Test _parameter_numbers
@testset "_parameter_numbers" begin
    # initialize model and references
    m = InfiniteModel()
    @infinite_parameter(m, par in [0, 1])
    @infinite_parameter(m, pars[1:2] in [0, 1])
    @variable(m, inf, Infinite(par))
    @variable(m, inf2, Infinite(par, pars))
    @variable(m, pt, Point(inf, 0))
    @variable(m, finite)
    var = build_variable(error, inf2, Dict{Int, Float64}(1 => 0.5), check = false)
    red = add_variable(m, var)
    # test for finite variable reference
    @testset "FiniteVariable" begin
        @test InfiniteOpt._parameter_numbers(pt) == []
        @test InfiniteOpt._parameter_numbers(finite) == []
    end
    # test for infinite variable reference
    @testset "InfiniteVariable" begin
        @test InfiniteOpt._parameter_numbers(inf) == [1]
        @test InfiniteOpt._parameter_numbers(inf2) == [1, 2, 3]
    end
    # test for parameter reference
    @testset "Parameter" begin
        @test InfiniteOpt._parameter_numbers(par) == [1]
        @test InfiniteOpt._parameter_numbers(pars[2]) == [3]
    end
    # test for semi-infinite variable reference
    @testset "SemiInfiniteInfinite" begin
        @test InfiniteOpt._parameter_numbers(red) == [2, 3]
    end
    # test for GenericAffExpr
    @testset "AffExpr" begin
        # make expressions
        aff1 = inf + inf2 + pt - 3
        aff2 = pt + finite - 2
        # test expressions
        @test sort!(InfiniteOpt._parameter_numbers(aff1)) == [1, 2, 3]
        @test InfiniteOpt._parameter_numbers(aff2) == []
    end
    # test for GenericQuadExpr
    @testset "QuadExpr" begin
        # make expressions
        quad1 = inf * inf2 + inf + inf2 + pt - 3 - par
        quad2 = pt * pt + pt + finite - 2
        # test expressions
        @test sort!(InfiniteOpt._parameter_numbers(quad1)) == [1, 2, 3]
        @test InfiniteOpt._parameter_numbers(quad2) == []
    end
end

# Test _model_from_expr
@testset "_model_from_expr" begin
    # initialize model and references
    m = InfiniteModel()
    @variable(m, hd)
    # test for variable reference
    @testset "Variable" begin
        @test InfiniteOpt._model_from_expr(hd) === m
    end
    # test for GenericAffExpr
    @testset "AffExpr" begin
        # make expressions
        aff1 = hd + 2
        aff2 = zero(JuMP.GenericAffExpr{Float64, GeneralVariableRef})
        # test expressions
        @test InfiniteOpt._model_from_expr(aff1) === m
        @test InfiniteOpt._model_from_expr(aff2) isa Nothing
    end
    # test for GenericQuadExpr
    @testset "QuadExpr" begin
        # make expressions
        quad1 = hd * hd + hd + 1
        quad2 = zero(JuMP.GenericQuadExpr{Float64, GeneralVariableRef})
        quad3 = hd * hd + 1
        # test expressions
        @test InfiniteOpt._model_from_expr(quad1) === m
        @test InfiniteOpt._model_from_expr(quad2) isa Nothing
        @test InfiniteOpt._model_from_expr(quad3) === m
    end
    # test for Vector{GeneralVariableRef}
    @testset "Vector{GeneralVariableRef}" begin
        vrefs1 = GeneralVariableRef[]
        vrefs2 = [hd]
        # test expressions
        @test InfiniteOpt._model_from_expr(vrefs1) isa Nothing
        @test InfiniteOpt._model_from_expr(vrefs2) === m
    end
    # test Fallback
    @testset "Fallback" begin
        @test_throws ErrorException InfiniteOpt._model_from_expr(42)
    end
end

# Test _remove_variable
@testset "_remove_variable" begin
    # initialize model and references
    m = InfiniteModel()
    @infinite_parameter(m, par in [0, 1])
    @variable(m, inf, Infinite(par))
    @variable(m, pt, Point(inf, 0))
    @variable(m, finite)
    # test for GenericAffExpr
    @testset "AffExpr" begin
        # make expressions
        aff1 = pt + 2par + finite
        aff2 = zero(GenericAffExpr{Float64, GeneralVariableRef})
        # test expressions
        @test isa(InfiniteOpt._remove_variable(aff1, finite), Nothing)
        @test !haskey(aff1.terms, finite)
        @test isa(InfiniteOpt._remove_variable(aff2, inf), Nothing)
    end
    # test for GenericQuadExpr
    @testset "QuadExpr" begin
        # make expressions
        quad1 = pt^2 + inf * pt + 2par + finite
        quad2 = zero(GenericQuadExpr{Float64, GeneralVariableRef})
        # test expressions
        quad = copy(quad1)
        @test isa(InfiniteOpt._remove_variable(quad, inf), Nothing)
        @test !haskey(quad.aff.terms, inf)
        @test !haskey(quad.terms, UnorderedPair{GeneralVariableRef}(inf, pt))
        quad = copy(quad1)
        @test isa(InfiniteOpt._remove_variable(quad, pt), Nothing)
        @test !haskey(quad.terms, UnorderedPair{GeneralVariableRef}(inf, pt))
        @test !haskey(quad.terms, UnorderedPair{GeneralVariableRef}(pt, pt))
        @test isa(InfiniteOpt._remove_variable(quad2, inf), Nothing)
    end
    # test for AbstractArray
    @testset "AbstractArray" begin
        # make expressions
        ex = [pt + 2par + finite, 2inf + pt]
        # test expressions
        @test isa(InfiniteOpt._remove_variable(ex, finite), Nothing)
        @test !haskey(ex[1].terms, finite)
    end
end

# Test _set_variable_coefficient!
@testset "_set_variable_coefficient!" begin
    # initialize model and references
    m = InfiniteModel()
    @infinite_parameter(m, par in [0, 1])
    @infinite_parameter(m, par2 in [0, 1])
    @variable(m, x, Infinite(par))
    @variable(m, y, Infinite(par, par2))
    @variable(m, z)
    # test with GeneralVariableRef
    @testset "GeneralVariableRef" begin
        @test InfiniteOpt._set_variable_coefficient!(x, x, 2) == 2 * x
        @test InfiniteOpt._set_variable_coefficient!(z, x, 2) == z + 2x
    end
    # test with GenericAffExpr
    @testset "AffExpr" begin
        @test InfiniteOpt._set_variable_coefficient!(x + z, x, 2) == 2x + z
        @test InfiniteOpt._set_variable_coefficient!(x + z, y, 2) == x + z + 2y
    end
    # test with GenericQuadExpr
    @testset "QuadExpr" begin
        @test InfiniteOpt._set_variable_coefficient!(y ^2 + x, x, 2) == y^2 + 2x
        @test InfiniteOpt._set_variable_coefficient!(y^2 + x, y, 2) == y^2 + x + 2y
    end
    # test fallabck
    @testset "Fallback" begin
        @test_throws ErrorException InfiniteOpt._set_variable_coefficient!(2, x, 2)
    end
end

# Test parameter reference methods
@testset "Parameter References" begin
    m = InfiniteModel()
    @infinite_parameter(m, t in [0, 1])
    @infinite_parameter(m, y in [0, 1])
    @infinite_parameter(m, x[1:3] in [0, 1])
    @variable(m, z)
    @expression(m, c1, 2z)
    @expression(m, c2, z + t + x[1])
    # test _make_param_tuple_element (IndependentParameterIndex)
    @testset "_make_param_tuple_element (Independent)" begin
        @test InfiniteOpt._make_param_tuple_element(m, index(t)) == t
        @test InfiniteOpt._make_param_tuple_element(m, index(y)) == y
    end
    # test _make_param_tuple_element (DependentParametersIndex)
    @testset "_make_param_tuple_element (Dependent)" begin
        obj_idx = index(first(x)).object_index
        @test InfiniteOpt._make_param_tuple_element(m, obj_idx) == x
    end
    # test parameter_refs
    @testset "parameter_refs" begin
        @test parameter_refs(c1) == ()
        @test parameter_refs(c2) == (t, x)
        @test parameter_refs(zero(AffExpr)) == ()
    end
end