text:
"""
For `obj isa AbstractSearch{T,P}`, the method `search(obj)` finds the realization of obj.parameter
"""
@with_kw struct Search{T, PRM <: AbstractParameter{<:MaybeVariable{T}}, TGT <: AbstractTarget{T}}
    parameter::PRM
    target::TGT
end

struct SearchExecution
    search::Search
    result::Execution
end
function execute(search::Search{T,<:Simulation}) where T
    varying_simulation = search.parameter
    result = search(varying_simulation, target; kwargs...)
    result_simulation = param_from_p(varying_simulation, minimizer(result))
    SearchExecution(search, execute(result_simulation))
end

minimizer(res::BlackBoxOptim.OptimizationResults) = best_candidate(res)
minimizer(res::Optim.OptimizationResults) = Optim.minimizer(res)

result_simulation(search_execution::SearchExecution) = search_execution.result.simulation

# * Run the search

make_problem_generator() = error("undefined.")

function _build_loss_objective(initial_problem, solver::Solver{T,Euler}, space, loss_fn, prob_generator) where T
    build_loss_objective(initial_problem, Euler(), loss_fn;
        prob_generator=prob_generator, dt=solver.simulated_dt,
        save_at=saved_dt(solver), save_idxs=save_idxs(solver, space))
end

function _build_loss_objective(initial_problem, solver::Solver{T,Nothing}, space, loss_fn, prob_generator) where T
    initial_problem = problem_generator(nothing, initial_values)
    build_loss_objective(initial_problem, Tsit5(), loss_fn;
        prob_generator=prob_generator,
        save_at=saved_dt(solver),
        timeseries_steps=solver.time_save_every,
        save_idxs=save_idxs(solver, space),
        alg_hints=[solver.stiffness])
end

function search(varying_simulation::Simulation{MaybeVariable{T}}, target::AbstractTarget{T}; MaxSteps=11e3) where T
    varying_deconstruction = VariableDeconstruction(varying_simulation)
    problem_generator = make_problem_generator(varying_deconstruction)
    initial_problem = problem_generator(nothing, initial_values)
    loss_fn = target_loss(target, varying_simulation)
    loss_obj = _build_loss_objective(initial_problem, varying_simulation.solver, loss_fn, problem_generator)
    result = bboptimize(loss_obj; NumDimensions=length(varying_deconstruction.initial_values),
        MaxSteps=MaxSteps, SearchRange=varying_deconstruction.bounds)
    return result
end

# function run_search_optim(varying_model, variable_map, solver, target, initial_p, p_bounds::Array{<:Tuple}; MaxSteps=11e3)
#     @show MaxSteps
#     initial_model = model_from_p(varying_model, variable_map, initial_p)
#     problem_generator = make_problem_generator(initial_model, solver, variable_map)
#     initial_problem = problem_generator(nothing, initial_p)
#     loss_fn = target_loss(target, initial_model, solver)
#     loss_obj = _build_loss_objective(initial_problem, solver, initial_model.space, loss_fn,
#                                                 problem_generator)
#     lower = [bounds[1] for bounds in p_bounds]
#     upper = [bounds[2] for bounds in p_bounds]
#     result = optimize(loss_obj, lower, upper, initial_p, Fminbox(ParticleSwarm()),
#                         Optim.Options(
#                          iterations = 10,
#                          show_every = 1))
#     return result
# end

function run_search(jl_filename::AbstractString)
    include(jl_filename)
    filecopy(p_search.output, jl_filename, basename(jl_filename))
    return p_search
end

function make_problem_generator(deconstructed_simulation::VariableDeconstruction{SIM}) where {SIM <: Simulation}
    function problem_generator(prob, values_vec::AbstractArray)
        simulation = reconstruct_with_vec(deconstructed_simulation, values_vec)
        system_function = make_system_function(simulation)
        ode_fn = convert(ODEFunction{true}, system_function)
        return ODEProblem(ode_fn, initial_value(simulation), time_span(simulation), values_vec)
    end

    return problem_generator
end


text:
#=
Copyright (c) 2018-2022 Chris Coey, Lea Kapelevich, and contributors

This Julia package Hypatia.jl is released under the MIT license; see LICENSE
file in the root directory or at https://github.com/chriscoey/Hypatia.jl
=#

insts = OrderedDict()
insts["minimal"] = [
    ((2, 3, 4, 2, true, false, false, true, true, false),),
    ((2, 3, 4, 2, true, false, false, true, true, true),),
    ((2, 3, 4, 2, false, true, false, true, true, true),),
    ((2, 3, 4, 2, false, false, true, true, true, true),),
    ((2, 3, 4, 2, true, false, false, false, false, true),),
    ((2, 3, 4, 2, false, true, false, false, false, true),),
    ]
insts["fast"] = [
    ((3, 5, 7, 2, true, false, false, true, true, false),),
    ((3, 5, 7, 2, true, false, false, true, true, true),),
    ((3, 5, 7, 2, false, true, false, true, true, true),),
    ((3, 5, 7, 2, false, false, true, true, true, true),),
    ((3, 5, 7, 2, true, false, false, false, false, true),),
    ((3, 5, 7, 2, false, true, false, false, false, true),),
    ((8, 15, 25, 5, true, false, false, true, true, false),),
    ((8, 15, 25, 5, true, false, false, true, true, true),),
    ((8, 15, 25, 5, false, true, false, true, true, true),),
    ((8, 15, 25, 5, false, false, true, true, true, true),),
    ((8, 15, 25, 5, true, false, false, false, false, true),),
    ((8, 15, 25, 5, false, true, false, false, false, true),),
    ((25, 75, 125, 5, true, false, false, true, true, false),),
    ((25, 75, 125, 5, true, false, false, true, true, true),),
    ((25, 75, 125, 5, false, true, false, true, true, true),),
    ((25, 75, 125, 5, false, false, true, true, true, true),),
    ((25, 75, 125, 5, true, false, false, false, false, true),),
    ((25, 75, 125, 5, false, true, false, false, false, true),),
    ]
return (DOptimalDesignNative, insts)

ext: jl


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

struct ScaleShiftTrafo{T<:Union{Real,AbstractVector{<:Real}}} <: Function
    a::T
    b::T
end

@functor ScaleShiftTrafo

Base.:(==)(a::ScaleShiftTrafo, b::ScaleShiftTrafo) = a.a == b.a && a.b == b.b
Base.isequal(a::ScaleShiftTrafo, b::ScaleShiftTrafo) = isequal(a.a, b.a) && isequal(a.b, b.b)
Base.hash(x::ScaleShiftTrafo, h::UInt) = hash(x.a, hash(x.b, hash(:JohnsonTrafoInv, hash(:ScaleShiftTrafo, h))))

(f::ScaleShiftTrafo{<:Real})(x::Real) = muladd(x, f.a, f.b)
(f::ScaleShiftTrafo)(x) = muladd.(x, f.a, f.b)

function ChangesOfVariables.with_logabsdet_jacobian(
    f::ScaleShiftTrafo{<:AbstractVector{<:Real}},
    x::AbstractMatrix{<:Real}
)
    ladj = sum(log.(abs.(f.a)))
    f(x), similar_fill(ladj, x, (size(x, 2),))'
end

function InverseFunctions.inverse(f::ScaleShiftTrafo)
    a_inv = inv.(f.a)
    b_inv = - a_inv .* f.b
    ScaleShiftTrafo(a_inv, b_inv)
end


text:
module VsRecorderBase

using Reexport: @reexport

import Dates, TOML

import DataStructures
import YAML
@reexport using DataStructures: Queue, enqueue!, dequeue!,
    OrderedDict
using FileIO: Stream, @format_str

@reexport using LinearAlgebra, Statistics
@reexport using VideoIO, Images

import StringDistances
using Tesseract, ImageFiltering
@reexport using MacroTools: @forward

export Missable, Nullable, @missable, @nullable,
    Yes,
    SimpleTypeWrapper, @type_wrapper,
    ∞, ±
include("utils/misc.jl")

export AbstractVsStrategy, AbstractVsSource, AbstractVsScene, AbstractVsStream,
    VsStream, VsFrame,
    VsContextData,
    VsConfig,
    VsContext
export image, time
export vs_setup, vs_init!, vs_parse_frame!, vs_update!, vs_result, vs_tryparse_scene
include("types.jl")

export BoundedBinaryHeap
include("utils/data_structures.jl")

export find_closest_n, find_closest,
    data_search_n, data_search
include("utils/data_search.jl")

export remove_spaces, parse_int
include("utils/string_utils.jl")

export to_gray_image, blur,
    color_distance, floodfill, floodfill!,
    draw_outline, draw_outline!,
    cycled_translate,
    blend_color,
    bounding_rect, shrink,
    prepare_text_for_ocr
include("utils/image_transform.jl")

export Rect, subimage, topleft, is_gray
include("utils/image_data.jl")

export image_distance, block,
    template_match_all, template_match,
    table_search
export SpriteSheet
include("utils/template_match.jl")

export open_video, open_camera,
    read_frame, load_image,
    enum_to_string, enum_from_string,
    serialize, deserialize
include("io.jl")
include("yaml.jl")

export create_ocr_instance,
    download_ocr_language,
    default_ocr_instance,
    ocr, ocr_tsv,
    is_cjk,
    init_multiple_ocr!, ocr_multiple_lang
include("ocr.jl")

export DefaultStrategy
include("strategies/default.jl")
using .DefaultStrategyModule: DefaultStrategy

include("main.jl")
include("api.jl")

function __init__()
    mkpath(OCR_DATA_DIR)
    nothing
end

end


text:
using SplitApplyPlot
using Documenter
using Literate, Glob
using CairoMakie

CairoMakie.activate!()

ENV["DATADEPS_ALWAYS_ACCEPT"] = true

# generate examples
GENERATED = joinpath(@__DIR__, "src", "generated")
SOURCE_FILES = Glob.glob("*.jl", GENERATED)
foreach(fn -> Literate.markdown(fn, GENERATED), SOURCE_FILES)

DocMeta.setdocmeta!(SplitApplyPlot, :DocTestSetup, :(using SplitApplyPlot); recursive=true)

makedocs(;
    modules=[SplitApplyPlot],
    authors="piever <pietro.vertechi@protonmail.com> and contributors",
    repo="https://github.com/piever/SplitApplyPlot.jl/blob/{commit}{path}#{line}",
    sitename="SplitApplyPlot.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://piever.github.io/SplitApplyPlot.jl",
        assets=["assets/favicon.ico"],
    ),
    pages=Any[
        "index.md",
        "Getting Started" => [
            "generated/penguins.md",
            "generated/gallery.md",
        ],
        "Algebra of Layers" => [
            "layers/introduction.md",
            "layers/data.md",
            "layers/mappings.md",
            "Transformations" => [
                "generated/visualtransformations.md",
                "generated/datatransformations.md",
            ],
            "layers/operations.md",
            "layers/draw.md",
        ],
        "Internals" => [
            "generated/entries.md",
        ],
        "API.md",
    ],
    strict = true,
)

deploydocs(;
    repo="github.com/piever/SplitApplyPlot.jl",
    push_preview = true,
)


text:
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
# ==============================================================================
#
#   Functions related to the conversion from quaternion to Euler angle and axis.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

export quat_to_angleaxis

"""
    quat_to_angleaxis(q::Quaternion{T}) where T

Convert the quaternion `q` to a Euler angle and axis representation (see
[`EulerAngleAxis`](@ref)). By convention, the Euler angle will be kept between
`[0, π] rad`.

# Remarks

This function will not fail if the quaternion norm is not 1. However, the
meaning of the results will not be defined, because the input quaternion does
not represent a 3D rotation. The user must handle such situations.

# Examples

```jldoctest
julia> q = Quaternion(cosd(45/2), sind(45/2), 0, 0);

julia> quat_to_angleaxis(q)
EulerAngleAxis{Float64}:
  Euler angle : 0.785398 rad  (45.0°)
  Euler axis  : [1.0, 0.0, 0.0]
```
"""
function quat_to_angleaxis(q::Quaternion{T}) where T
    # If `q0` is 1 or -1, then we have an identity rotation.
    if abs(q.q0) >= 1 - eps()
        return EulerAngleAxis(T(0), SVector{3, T}(0, 0, 0))
    else
        # Compute sin(θ/2).
        sθo2 = √(q.q1 * q.q1 + q.q2 * q.q2 + q.q3 * q.q3)

        # Compute θ in range [0, 2π].
        θ = 2acos(q.q0)

        # Keep θ between [0, π].
        s = +1
        if θ > π
            θ = T(2π) - θ
            s = -1
        end

        return EulerAngleAxis(θ, s * SVector{3}(q.q1, q.q2, q.q3) / sθo2)
    end
end


text:
export CategoryWithProducts, ob, terminal, delete, product, proj1, proj2, pair,
  CategoryWithCoproducts, initial, create, coproduct, coproj1, coproj2, copair,
  CompleteCategory, equalizer, incl, factorize,
  CocompleteCategory, coequalizer, proj

# Limits
########

""" Theory of a *category with (finite) products*

Finite products are presented in biased style, via the nullary case (terminal
objects) and the binary case (binary products). The equational axioms are
standard, especially in type theory (Lambek & Scott, 1986, Section 0.5 or
Section I.3).

For a monoidal category axiomatization, see [`CartesianCategory`](@ref).
"""
@theory Category(Ob,Hom) => CategoryWithProducts(Ob,Hom,Terminal,Product) begin
  Terminal()::TYPE
  Product(foot1::Ob, foot2::Ob)::TYPE

  # Terminal object.
  terminal()::Terminal()
  ob(⊤::Terminal())::Ob
  delete(⊤::Terminal(), C::Ob)::(C → ob(⊤))

  # Binary products.
  product(A::Ob, B::Ob)::Product(A,B)
  ob(Π::Product(A,B))::Ob ⊣ (A::Ob, B::Ob)
  proj1(Π::Product(A,B))::(ob(Π) → A) ⊣ (A::Ob, B::Ob)
  proj2(Π::Product(A,B))::(ob(Π) → B) ⊣ (A::Ob, B::Ob)
  (pair(Π::Product(A,B), f::(C → A), g::(C → B))::(C → ob(Π))
    ⊣ (A::Ob, B::Ob, C::Ob))

  # Projection axioms.
  (pair(Π,f,g) ⋅ proj1(Π) == f
    ⊣ (A::Ob, B::Ob, C::Ob, Π::Product(A,B), f::(C → A), g::(C → B)))
  (pair(Π,f,g) ⋅ proj2(Π) == g
    ⊣ (A::Ob, B::Ob, C::Ob, Π::Product(A,B), f::(C → A), g::(C → B)))

  # Uniqueness axioms.
  f == g ⊣ (C::Ob, ⊤::Terminal(), f::(C → ob(⊤)), g::(C → ob(⊤)))
  (pair(h⋅proj1(Π), h⋅proj2(Π)) == h
    ⊣ (A::Ob, B::Ob, C::Ob, Π::Product(A,B), h::(C → ob(Π))))
end

""" Theory of a *(finitely) complete category*

Finite limits are presented in biased style, via finite products and equalizers.
The equational axioms for equalizers are obscure, but can found in (Lambek &
Scott, 1986, Section 0.5), apparently following "Burroni's pioneering ideas".
"""
@theory CategoryWithProducts(Ob,Hom,Terminal,Product) =>
    CompleteCategory(Ob,Hom,Terminal,Product,Equalizer) begin
  Equalizer(f::(A → B), g::(A → B))::TYPE ⊣ (A::Ob, B::Ob)

  # Equalizers.
  equalizer(f::(A → B), g::(A → B))::Equalizer(f,g) ⊣ (A::Ob, B::Ob)
  ob(eq::Equalizer(f,g))::Ob ⊣ (A::Ob, B::Ob, f::(A → B), g::(A → B))
  (incl(eq::Equalizer(f,g))::(ob(eq) → A)
    ⊣ (A::Ob, B::Ob, f::(A → B), g::(A → B)))
  (factorize(eq::Equalizer(f,g), h::(C → A),
             eq_h::Equalizer(h⋅f,h⋅g))::(ob(eq_h) → ob(eq))
    ⊣ (A::Ob, B::Ob, C::Ob, f::(A → B), g::(A → B)))

  # Equalizer axioms.
  (incl(eq)⋅f == incl(eq)⋅g
    ⊣ (A::Ob, B::Ob, f::(A → B), g::(A → B), eq::Equalizer(f,g)))
  (incl(eq) == id(A)
    ⊣ (A::Ob, B::Ob, f::(A → B), eq::Equalizer(f,f)))
  (factorize(eq,h,eq_h) ⋅ incl(eq) == incl(eq_h) ⋅ h
    ⊣ (A::Ob, B::Ob, C::Ob, f::(A → B), g::(A → B), h::(C → A),
       eq::Equalizer(f,g), eq_h::Equalizer(h⋅f, h⋅g)))
  (factorize(eq, k⋅incl(eq), eq_k) == k
    ⊣ (A::Ob, B::Ob, D::Ob, f::(A → B), g::(A → B), eq::Equalizer(f,g),
       k::(D → ob(eq)), eq_k::Equalizer(k⋅incl(eq)⋅f, k⋅incl(eq)⋅g)))
end

# Colimits
##########

""" Theory of a *category with (finite) coproducts*

Finite coproducts are presented in biased style, via the nullary case (initial
objects) and the binary case (binary coproducts). The axioms are dual to those
of [`CategoryWithProducts`](@ref).

For a monoidal category axiomatization, see [`CocartesianCategory`](@ref).
"""
@theory Category(Ob,Hom) => CategoryWithCoproducts(Ob,Hom,Initial,Coproduct) begin
  Initial()::TYPE
  Coproduct(foot1::Ob, foot2::Ob)::TYPE

  # Initial object.
  initial()::Initial()
  ob(⊥::Initial())::Ob
  create(⊥::Initial(), C::Ob)::(ob(⊥) → C)

  # Binary coproducts.
  coproduct(A::Ob, B::Ob)::Coproduct(A,B)
  ob(⨆::Coproduct(A,B))::Ob ⊣ (A::Ob, B::Ob)
  coproj1(⨆::Coproduct(A,B))::(A → ob(⨆)) ⊣ (A::Ob, B::Ob)
  coproj2(⨆::Coproduct(A,B))::(B → ob(⨆)) ⊣ (A::Ob, B::Ob)
  (copair(⨆::Coproduct(A,B), f::(A → C), g::(B → C))::(ob(⨆) → C)
    ⊣ (A::Ob, B::Ob, C::Ob))

  # Coprojection axioms.
  (coproj1(⨆) ⋅ copair(⨆,f,g) == f
    ⊣ (A::Ob, B::Ob, C::Ob, ⨆::Coproduct(A,B), f::(A → C), g::(B → C)))
  (coproj2(⨆) ⋅ copair(⨆,f,g) == g
    ⊣ (A::Ob, B::Ob, C::Ob, ⨆::Coproduct(A,B), f::(A → C), g::(B → C)))

  # Uniqueness axioms.
  f == g ⊣ (C::Ob, ⊥::Initial(), f::(ob(⊥) → C), g::(ob(⊥) → C))
  (copair(coproj1(⨆)⋅h, coproj2(⨆)⋅h) == h
    ⊣ (A::Ob, B::Ob, C::Ob, ⨆::Coproduct(A,B), h::(ob(⨆) → C)))
end

""" Theory of a *(finitely) cocomplete category*

Finite colimits are presented in biased style, via finite coproducts and
coequalizers. The axioms are dual to those of [`CompleteCategory`](@ref).
"""
@theory CategoryWithCoproducts(Ob,Hom,Initial,Coproduct) =>
    CocompleteCategory(Ob,Hom,Initial,Coproduct,Coequalizer) begin
  Coequalizer(f::(A → B), g::(A → B))::TYPE ⊣ (A::Ob, B::Ob)

  # Coequalizers.
  coequalizer(f::(A → B), g::(A → B))::Coequalizer(f,g) ⊣ (A::Ob, B::Ob)
  ob(eq::Coequalizer(f,g))::Ob ⊣ (A::Ob, B::Ob, f::(A → B), g::(A → B))
  (proj(eq::Coequalizer(f,g))::(B → ob(eq))
    ⊣ (A::Ob, B::Ob, f::(A → B), g::(A → B)))
  (factorize(coeq::Coequalizer(f,g), h::(B → C),
             coeq_h::Coequalizer(f⋅h,g⋅h))::(ob(coeq) → ob(coeq_h))
    ⊣ (A::Ob, B::Ob, C::Ob, f::(A → B), g::(A → B)))

  # Coequalizer axioms.
  (f⋅proj(coeq) == g⋅proj(coeq)
    ⊣ (A::Ob, B::Ob, f::(A → B), g::(A → B), coeq::Coequalizer(f,g)))
  (proj(coeq) == id(B)
    ⊣ (A::Ob, B::Ob, f::(A → B), coeq::Coequalizer(f,f)))
  (proj(coeq) ⋅ factorize(coeq,h,coeq_h) == h ⋅ proj(coeq_h)
    ⊣ (A::Ob, B::Ob, C::Ob, f::(A → B), g::(A → B), h::(B → C),
       coeq::Coequalizer(f,g), coeq_h::Coequalizer(f⋅h, g⋅h)))
  (factorize(coeq, proj(coeq)⋅k, coeq_k) == k
    ⊣ (A::Ob, B::Ob, D::Ob, f::(A → B), g::(A → B),
       coeq::Coequalizer(f,g), k::(ob(coeq) → D),
       coeq_k::Coequalizer(f⋅proj(coeq)⋅k, g⋅proj(coeq)⋅k)))
end


text:
using Random
using Plots
using DifferentialEquations,DiffEqFlux,LinearAlgebra
using Optim
using ForwardDiff
include("../src/GLOQ.jl")
pyplot()
#import GLOQ
############################################################
# Step 1: Generate Synthetic Data
############################################################
HK = [0.0 0.0 1.0;
      0.0 0.0 0.0;
      1.0 0.0 0.0]
HS = [0.0 1.0 0.0;
     -1.0 0.0 0.0;
      0.0 0.0 0.0]
L1 = [0.0 0.25 0.0;
      0.0 0.0  0.5;
      0.0 0.0  0.0]
L2 = [0.0 0.0 0.0;
      0.0 0.5 0.0;
      0.0 0.0 1.0]
L_list = (L1,L2)
LK,LS,LD = GLOQ.make_lindblad_operator(HK,HS,L_list)
println("Lindblad operator assembled")
# define the time interval
tspan_length = 21
t_span = range(0.0,pi,length=tspan_length)
# forward solves with different initial conditions,
# use exponential_solver
u1 = [1.0;0;0]
rho1_real,rho1_imag_adj = GLOQ.exponential_solver(real(u1),-imag(u1),LK,LS,LD,t_span;initial_type = "states")
u2 = [0;1.0;0]
rho2_real,rho2_imag_adj = GLOQ.exponential_solver(real(u2),-imag(u2),LK,LS,LD,t_span;initial_type = "states")
u3 = [1.0/sqrt(2.0);-im/sqrt(2.0);0]
rho3_real,rho3_imag_adj = GLOQ.exponential_solver(real(u3),-imag(u3),LK,LS,LD,t_span;initial_type = "states")
u4 = [0.0;1.0/sqrt(2.0);-im/sqrt(2.0)]
rho4_real,rho4_imag_adj = GLOQ.exponential_solver(real(u4),-imag(u4),LK,LS,LD,t_span;initial_type = "states")
u5 = [1.0/sqrt(2.0);0.0;-im/sqrt(2.0)]
rho5_real,rho5_imag_adj = GLOQ.exponential_solver(real(u5),-imag(u5),LK,LS,LD,t_span;initial_type = "states")

population1_history = GLOQ.get_population(rho1_real)
population2_history = GLOQ.get_population(rho2_real)
population3_history = GLOQ.get_population(rho3_real)
population4_history = GLOQ.get_population(rho4_real)
population5_history = GLOQ.get_population(rho5_real)

fig = plot(t_span,population1_history,lab=:false)
plot!(fig,t_span,population2_history,lab=:false)
plot!(fig,t_span,population3_history,lab=:false)
plot!(fig,t_span,population4_history,lab=:false)
plot!(fig,t_span,population5_history,lab=:false)
display(fig)

# add noise to the synthetic data
noise_amp = 0.025
noisy_data1 = population1_history+noise_amp.*rand(Float64,size(population1_history))
noisy_data2 = population2_history+noise_amp.*rand(Float64,size(population2_history))
noisy_data3 = population3_history+noise_amp.*rand(Float64,size(population3_history))
noisy_data4 = population4_history+noise_amp.*rand(Float64,size(population4_history))
noisy_data5 = population5_history+noise_amp.*rand(Float64,size(population5_history))
# present the noisy data

fig = plot(t_span,noisy_data1,lab=:false)
plot!(fig,t_span,noisy_data2,lab=:false)
plot!(fig,t_span,noisy_data3,lab=:false)
plot!(fig,t_span,noisy_data4,lab=:false)
plot!(fig,t_span,noisy_data5,lab=:false)
display(fig)

p_true = [1.0 1.0 0.25 0.5 0.5]
############################################################
# Step 2: Learn the operators
############################################################
rho1_initial_real,rho1_initial_adj_imag = GLOQ.convert_state_to_density(real(u1),-imag(u1))
rho2_initial_real,rho2_initial_adj_imag = GLOQ.convert_state_to_density(real(u2),-imag(u2))
rho3_initial_real,rho3_initial_adj_imag = GLOQ.convert_state_to_density(real(u3),-imag(u3))
rho4_initial_real,rho4_initial_adj_imag = GLOQ.convert_state_to_density(real(u4),-imag(u4))
rho5_initial_real,rho5_initial_adj_imag = GLOQ.convert_state_to_density(real(u5),-imag(u5))
rho1_initial = [rho1_initial_real;rho1_initial_adj_imag]
rho2_initial = [rho2_initial_real;rho2_initial_adj_imag]
rho3_initial = [rho3_initial_real;rho3_initial_adj_imag]
rho4_initial = [rho4_initial_real;rho4_initial_adj_imag]
rho5_initial = [rho5_initial_real;rho5_initial_adj_imag]
# p[1]-p[5] upper diagonal of the Hamiltonian part
# p[6]-p[9] parameters in the Lindblad term

function loss(p)
      _HK = [0  0  p[2];
             0  0  0;
             p[2] 0  0]
      _HS = [0  p[1]  0;
            -p[1]  0  0;
             0  0  0]
      _L1 = [0.0 p[3] 0.0;
            0.0 0.0  p[4];
            0.0 0.0  0.0]
      _L2 = [0.0 0.0 0.0;
            0.0 p[5] 0.0;
            0.0 0.0 2.0*p[5]]
      _L_list = (_L1,_L2)
      _LK,_LS,_LD = GLOQ.make_lindblad_operator(_HK,_HS,_L_list)

      _problem = GLOQ.LindbladODEProblem(rho1_initial_real,rho1_initial_adj_imag,
                                        _LK,_LS,_LD,
                                         t_span[end];initial_type="density")
      rho1_data = solve(_problem,saveat=t_span)
      rho2_data = solve(remake(_problem, u0=rho2_initial), saveat=t_span)
      rho3_data = solve(remake(_problem, u0=rho3_initial), saveat=t_span)
      rho4_data = solve(remake(_problem, u0=rho4_initial), saveat=t_span)
      rho5_data = solve(remake(_problem, u0=rho5_initial), saveat=t_span)

      pop_data1 = GLOQ.get_population(Array(rho1_data[1:Int64(end/2),:]))
      pop_data2 = GLOQ.get_population(Array(rho2_data[1:Int64(end/2),:]))
      pop_data3 = GLOQ.get_population(Array(rho3_data[1:Int64(end/2),:]))
      pop_data4 = GLOQ.get_population(Array(rho4_data[1:Int64(end/2),:]))
      pop_data5 = GLOQ.get_population(Array(rho5_data[1:Int64(end/2),:]))

      _loss = sum(abs2, pop_data1-noisy_data1)
      _loss += sum(abs2, pop_data2-noisy_data2)
      _loss += sum(abs2, pop_data3-noisy_data3)
      _loss += sum(abs2, pop_data4-noisy_data4)
      _loss += sum(abs2, pop_data5-noisy_data5)
      return _loss
end

p_initial = [1.0;1.0;0.5;0.5;0.5]
loss_initial = loss(p_initial)
println("Initial loss: ",loss_initial)

@time p_DEFlux = DiffEqFlux.sciml_train(loss, p_initial)
@time p_Optim = Optim.optimize(loss, p_initial)
println("True p: ",p_true)
println("DiffEqFlux p: ",p_DEFlux.u)
println("DiffEqFlux training history: ",p_DEFlux)
println("Optim p: ",p_Optim.minimizer)
println("Optim training history",p_Optim)

################################
function loss3(p)
      _HK = [0    0    p[2];
             0    p[3] p[4];
             p[2] p[4] p[5]]
      _HS = [0  p[1]  0;
            -p[1]  0  0;
             0  0  0]
      _L1 = [0.0 p[6] 0.0;
            0.0 0.0  p[7];
            0.0 0.0  0.0]
      _L2 = [0.0 0.0 0.0;
            0.0 p[8] 0.0;
            0.0 0.0 2.0*p[8]]
      _L_list = (_L1,_L2)
      _LK,_LS,_LD = GLOQ.make_lindblad_operator(_HK,_HS,_L_list)

      _problem = GLOQ.LindbladODEProblem(rho1_initial_real,rho1_initial_adj_imag,
                                              _LK,_LS,_LD,
                                               t_span[end];initial_type="density")
      rho1_data = solve(_problem,saveat=t_span)
      rho2_data = solve(remake(_problem, u0=rho2_initial), saveat=t_span)
      rho3_data = solve(remake(_problem, u0=rho3_initial), saveat=t_span)
      rho4_data = solve(remake(_problem, u0=rho4_initial), saveat=t_span)
      rho5_data = solve(remake(_problem, u0=rho5_initial), saveat=t_span)

      pop_data1 = GLOQ.get_population(Array(rho1_data[1:Int64(end/2),:]))
      pop_data2 = GLOQ.get_population(Array(rho2_data[1:Int64(end/2),:]))
      pop_data3 = GLOQ.get_population(Array(rho3_data[1:Int64(end/2),:]))
      pop_data4 = GLOQ.get_population(Array(rho4_data[1:Int64(end/2),:]))
      pop_data5 = GLOQ.get_population(Array(rho5_data[1:Int64(end/2),:]))

      _loss = sum(abs2, pop_data1-noisy_data1)
      _loss += sum(abs2, pop_data2-noisy_data2)
      _loss += sum(abs2, pop_data3-noisy_data3)
      _loss += sum(abs2, pop_data4-noisy_data4)
      _loss += sum(abs2, pop_data5-noisy_data5)
      return _loss
end

p3_initial = [1.0;1.0;0.25;0.1;0.1;0.5;0.5;0.5]
@time p3_DEFlux = DiffEqFlux.sciml_train(loss3, p3_initial)
@time p3_Optim = Optim.optimize(loss3, p3_initial)
println("True p: ",p_true)
println("DEFlux p: ",p3_DEFlux.u)
println("Optim p: ",p3_Optim.minimizer)
println(p3_DEFlux)
println(p3_Optim)

################################
function loss_exp(p)
      _HK = [0    0    p[2];
             0    p[3] p[4];
             p[2] p[4] p[5]]
      _HS = [0  p[1]  0;
            -p[1]  0  0;
             0  0  0]
      _L1 = [0.0 p[6] 0.0;
            0.0 0.0  p[7];
            0.0 0.0  0.0]
      _L2 = [0.0 0.0 0.0;
            0.0 p[8] 0.0;
            0.0 0.0 2.0*p[8]]
      _L_list = (_L1,_L2)
      _LK,_LS,_LD = GLOQ.make_lindblad_operator(_HK,_HS,_L_list)

      rho1_history_u,rho1_history_v = GLOQ.exponential_solver(real(u1),-imag(u1),_LK,_LS,_LD,t_span;initial_type = "states")
      rho2_history_u,rho2_history_v = GLOQ.exponential_solver(real(u2),-imag(u2),_LK,_LS,_LD,t_span;initial_type = "states")
      rho3_history_u,rho3_history_v = GLOQ.exponential_solver(real(u3),-imag(u3),_LK,_LS,_LD,t_span;initial_type = "states")
      rho4_history_u,rho4_history_v = GLOQ.exponential_solver(real(u4),-imag(u4),_LK,_LS,_LD,t_span;initial_type = "states")
      rho5_history_u,rho5_history_v = GLOQ.exponential_solver(real(u5),-imag(u5),_LK,_LS,_LD,t_span;initial_type = "states")

      pop_data1 = GLOQ.get_population(rho1_history_u)
      pop_data2 = GLOQ.get_population(rho2_history_u)
      pop_data3 = GLOQ.get_population(rho3_history_u)
      pop_data4 = GLOQ.get_population(rho4_history_u)
      pop_data5 = GLOQ.get_population(rho5_history_u)

      _loss = sum(abs2, pop_data1-noisy_data1)
      _loss += sum(abs2, pop_data2-noisy_data2)
      _loss += sum(abs2, pop_data3-noisy_data3)
      _loss += sum(abs2, pop_data4-noisy_data4)
      _loss += sum(abs2, pop_data5-noisy_data5)
      return _loss
end

p_exp_initial = [1.0;1.0;0.25;0.1;0.1;0.5;0.5;0.5]
@time p_exp_DEFlux = DiffEqFlux.sciml_train(loss_exp, p_exp_initial)
@time p_exp_Optim = Optim.optimize(loss_exp, p_exp_initial)
println("True p: ",p_true)
println("DEFlux p: ",p_exp_DEFlux.u)
println("Optim p: ",p_exp_Optim.minimizer)
println(p_exp_DEFlux)
println(p_exp_Optim)


text:
"""
    log2i(x::Integer) -> Integer

Return log2(x), this integer version of `log2` is fast but only valid for number equal to 2^n.
"""
function log2i end

for N in [8, 16, 32, 64, 128]
    T = Symbol(:Int, N)
    UT = Symbol(:UInt, N)
    @eval begin
        log2i(x::$T) =
            !signbit(x) ? ($(N - 1) - leading_zeros(x)) :
            throw(ErrorException("nonnegative expected ($x)"))
        log2i(x::$UT) = $(N - 1) - leading_zeros(x)
    end
end

log2dim(x::AbstractVector) = log2i(size(x, 1))
log2dim(x::AbstractMatrix) = log2i(size(x, 2))

function ctrl_offset(ctrl::CtrlLocations)
    locs = plain(ctrl)
    mask = 0
    for i in 1:length(ctrl)
        if ctrl.flags[i]
            mask += (1 << (locs[i] - 1))
        end
    end
    return mask
end


text:
# Design: DP

# Float64 version of value_to_go
function DP_2_finale_mean( number_of_allocations :: Int64 , success_probability_arm_1 :: Float64 , success_probability_arm_2 :: Float64 , value_to_go_evaluation :: Array{ Float64 , 1 } , value_to_go :: Array{ Float64 , 1 } , prior_success_arm_1 :: Int64 = Int64( 1 ) , prior_failure_arm_1 :: Int64 = Int64( 1 ) , prior_success_arm_2 :: Int64 = Int64( 1 ) , prior_failure_arm_2 :: Int64 = Int64( 1 ) )
  # This function implements the DP design for 2 arms
  # Uses linear indexing for value_to_go
  # Output is the expected value of the finale array (the sum of all finale elements multiplied by their probabilities under success_probability_arm_1 & 2 )

        # backwards recursion: the finale (i.e., number_of_observed_responses = number_of_allocations ) not needed as it is given by value_to_go

        # backwards recursion: t-th step
        for number_of_observed_responses = ( number_of_allocations - 1 ) : -1 : 0
            value_to_go_lin_index = 0
            for number_of_successes_arm_2 = 0 : number_of_observed_responses , number_of_failures_arm_1 = 0 : ( number_of_observed_responses - number_of_successes_arm_2 ) , number_of_successes_arm_1 = 0 : ( number_of_observed_responses - number_of_successes_arm_2 - number_of_failures_arm_1 )
                @inbounds begin

                value_to_go_lin_index += 1
                belief_of_success_arm_1 = ( ( prior_success_arm_1 + number_of_successes_arm_1 ) / ( prior_success_arm_1 + prior_failure_arm_1 + number_of_successes_arm_1 + number_of_failures_arm_1 ) )
                belief_of_success_arm_2 = ( ( prior_success_arm_2 + number_of_successes_arm_2 ) / ( prior_success_arm_2 + prior_failure_arm_2 + number_of_observed_responses - number_of_successes_arm_1 - number_of_failures_arm_1 ) )

                value_to_go_if_action_1 = belief_of_success_arm_1 * ( 1.0 + value_to_go[ value_to_go_lin_index + div( ( number_of_observed_responses + 2 ) * ( number_of_observed_responses + 3 ) - ( number_of_observed_responses - number_of_successes_arm_2 + 2 ) * ( number_of_observed_responses - number_of_successes_arm_2 + 3 ) , 2 ) + number_of_failures_arm_1 + 1 ] ) + ( 1.0 - belief_of_success_arm_1 ) * value_to_go[ value_to_go_lin_index + div( ( number_of_observed_responses + 2 ) * ( number_of_observed_responses + 3 ) - ( number_of_observed_responses - number_of_successes_arm_2 + 2 ) * ( number_of_observed_responses - number_of_successes_arm_2 + 3 ) , 2 ) + ( number_of_observed_responses - number_of_successes_arm_2 + 2 ) ]
                value_to_go_if_action_2 = belief_of_success_arm_2 * ( 1.0 + value_to_go[ value_to_go_lin_index + div( ( number_of_observed_responses + 2 ) * ( number_of_observed_responses + 3 ) , 2 ) ] ) + ( 1.0 - belief_of_success_arm_2 ) * value_to_go[ value_to_go_lin_index + number_of_failures_arm_1 + div( ( number_of_observed_responses + 2 ) * ( number_of_observed_responses + 3 ) - ( number_of_observed_responses - number_of_successes_arm_2 + 2 ) * ( number_of_observed_responses - number_of_successes_arm_2 + 3 ) , 2 ) ]

                if ( value_to_go_if_action_1 - value_to_go_if_action_2 ) > BB_numerical_precision_64 * ( value_to_go_if_action_1 + value_to_go_if_action_2 )

                    value_to_go[ value_to_go_lin_index ] = value_to_go_if_action_1

                    value_to_go_evaluation[ value_to_go_lin_index ] = success_probability_arm_1 * value_to_go_evaluation[ value_to_go_lin_index + div( ( number_of_observed_responses + 2 ) * ( number_of_observed_responses + 3 ) - ( number_of_observed_responses - number_of_successes_arm_2 + 2 ) * ( number_of_observed_responses - number_of_successes_arm_2 + 3 ) , 2 ) + number_of_failures_arm_1 + 1 ] + ( 1.0 - success_probability_arm_1 ) * value_to_go_evaluation[ value_to_go_lin_index + div( ( number_of_observed_responses + 2 ) * ( number_of_observed_responses + 3 ) - ( number_of_observed_responses - number_of_successes_arm_2 + 2 ) * ( number_of_observed_responses - number_of_successes_arm_2 + 3 ) , 2 ) + ( number_of_observed_responses - number_of_successes_arm_2 + 2 ) ]

                elseif ( value_to_go_if_action_2 - value_to_go_if_action_1 ) > BB_numerical_precision_64 * ( value_to_go_if_action_1 + value_to_go_if_action_2 )

                    value_to_go[ value_to_go_lin_index ] = value_to_go_if_action_2

                    value_to_go_evaluation[ value_to_go_lin_index ] = success_probability_arm_2 * ( value_to_go_evaluation[ value_to_go_lin_index + div( ( number_of_observed_responses + 2 ) * ( number_of_observed_responses + 3 ) , 2 ) ] ) + ( 1.0 - success_probability_arm_2 ) * value_to_go_evaluation[ value_to_go_lin_index + number_of_failures_arm_1 + div( ( number_of_observed_responses + 2 ) * ( number_of_observed_responses + 3 ) - ( number_of_observed_responses - number_of_successes_arm_2 + 2 ) * ( number_of_observed_responses - number_of_successes_arm_2 + 3 ) , 2 ) ]

                else #if value_to_go_if_action_1 approx== value_to_go_if_action_2

                    value_to_go[ value_to_go_lin_index ] = ( value_to_go_if_action_1 + value_to_go_if_action_2 ) / 2

                    value_to_go_evaluation[ value_to_go_lin_index ] = ( success_probability_arm_1 * value_to_go_evaluation[ value_to_go_lin_index + div( ( number_of_observed_responses + 2 ) * ( number_of_observed_responses + 3 ) - ( number_of_observed_responses - number_of_successes_arm_2 + 2 ) * ( number_of_observed_responses - number_of_successes_arm_2 + 3 ) , 2 ) + number_of_failures_arm_1 + 1 ] + ( 1.0 - success_probability_arm_1 ) * value_to_go_evaluation[ value_to_go_lin_index + div( ( number_of_observed_responses + 2 ) * ( number_of_observed_responses + 3 ) - ( number_of_observed_responses - number_of_successes_arm_2 + 2 ) * ( number_of_observed_responses - number_of_successes_arm_2 + 3 ) , 2 ) + ( number_of_observed_responses - number_of_successes_arm_2 + 2 ) ] + success_probability_arm_2 * ( value_to_go_evaluation[ value_to_go_lin_index + div( ( number_of_observed_responses + 2 ) * ( number_of_observed_responses + 3 ) , 2 ) ] ) + ( 1.0 - success_probability_arm_2 ) * value_to_go_evaluation[ value_to_go_lin_index + number_of_failures_arm_1 + div( ( number_of_observed_responses + 2 ) * ( number_of_observed_responses + 3 ) - ( number_of_observed_responses - number_of_successes_arm_2 + 2 ) * ( number_of_observed_responses - number_of_successes_arm_2 + 3 ) , 2 ) ] ) / 2

                end

            end # @inbounds
            end
        end

        return value_to_go_evaluation[ 1 ]

end

function DP_2_NS( number_of_allocations :: Int64 , success_probability_arm_1 :: Float64 , success_probability_arm_2 :: Float64 , float_version :: Int64 = Int64( 64 ) , prior_success_arm_1 :: Int64 = Int64( 1 ) , prior_failure_arm_1 :: Int64 = Int64( 1 ) , prior_success_arm_2 :: Int64 = Int64( 1 ) , prior_failure_arm_2 :: Int64 = Int64( 1 ) )
  # This function implements the DP design for 2 arms
  # Uses linear indexing for value_to_go
  # Output is the mean and the variance of the number of successes CHECK (as Float64 irrespectively of the value of "precision")

#        if float_version == 16
#
#            value_to_go_16 :: Array{ Float16 , 1 } = zeros( Float16 , div( ( number_of_allocations + 1 ) * ( number_of_allocations + 2 ) * ( number_of_allocations + 3 ) - 1 * 2 * 3 , 6 ) + 1 )
#            DP_2_finale_mean( number_of_allocations , value_to_go_16 , prior_success_arm_1 , prior_failure_arm_1 , prior_success_arm_2 , prior_failure_arm_2 )
#
#        elseif float_version == 32
#
#            value_to_go_32 :: Array{ Float32 , 1 } = zeros( Float32 , div( ( number_of_allocations + 1 ) * ( number_of_allocations + 2 ) * ( number_of_allocations + 3 ) - 1 * 2 * 3 , 6 ) + 1 )
#            DP_2_finale_mean( number_of_allocations , value_to_go_32 , prior_success_arm_1 , prior_failure_arm_1 , prior_success_arm_2 , prior_failure_arm_2 )
#
#        else

            value_to_go :: Array{ Float64 , 1 } = zeros( Float64 , div( ( number_of_allocations + 1 ) * ( number_of_allocations + 2 ) * ( number_of_allocations + 3 ) - 1 * 2 * 3 , 6 ) + 1 )

            value_to_go_evaluation :: Array{ Float64 , 1 } = zeros( Float64 , div( ( number_of_allocations + 1 ) * ( number_of_allocations + 2 ) * ( number_of_allocations + 3 ) - 1 * 2 * 3 , 6 ) + 1 )
            value_to_go_lin_index = 0
            for number_of_successes_arm_2 = 0 : number_of_allocations , number_of_failures_arm_1 = 0 : number_of_allocations - number_of_successes_arm_2 , number_of_successes_arm_1 = 0 : number_of_allocations - number_of_successes_arm_2 - number_of_failures_arm_1
                value_to_go_lin_index += 1
                @inbounds begin
                    value_to_go_evaluation[ value_to_go_lin_index ] = number_of_successes_arm_1 + number_of_successes_arm_2
                end # @inbounds
            end

            NS_mean = DP_2_finale_mean( number_of_allocations , success_probability_arm_1 , success_probability_arm_2 , value_to_go_evaluation , value_to_go , prior_success_arm_1 , prior_failure_arm_1 , prior_success_arm_2 , prior_failure_arm_2 )

            value_to_go = zeros( Float64 , div( ( number_of_allocations + 1 ) * ( number_of_allocations + 2 ) * ( number_of_allocations + 3 ) - 1 * 2 * 3 , 6 ) + 1 )

            value_to_go_lin_index = 0
            for number_of_successes_arm_2 = 0 : number_of_allocations , number_of_failures_arm_1 = 0 : number_of_allocations - number_of_successes_arm_2 , number_of_successes_arm_1 = 0 : number_of_allocations - number_of_successes_arm_2 - number_of_failures_arm_1
                value_to_go_lin_index += 1
                @inbounds begin
                    value_to_go_evaluation[ value_to_go_lin_index ] = ( number_of_successes_arm_1 + number_of_successes_arm_2 - NS_mean ) ^ 2
                end # @inbounds
            end

            NS_variance = DP_2_finale_mean( number_of_allocations , success_probability_arm_1 , success_probability_arm_2 , value_to_go_evaluation , value_to_go , prior_success_arm_1 , prior_failure_arm_1 , prior_success_arm_2 , prior_failure_arm_2 )

            return NS_mean , NS_variance

#        end

end

max_stars_repo_name: PeterJacko/BinaryBandit

max_issues_repo_name: PeterJacko/BinaryBandit

max_forks_repo_name: PeterJacko/BinaryBandit


text:
##################################
##################################
# Figures 5a, 7, A2, A3
# Tables 2, 3, A4
##################################
##################################

function solve_uncert(na = 1; expand = 1., quad = 11)

    inputs = Dict{Symbol,Any}(:scaling_factor => 1.0::Float64,
        :initial_approximation_level => [9 9 9 9 9]::Array{Int64,2},
        :expand_domain => expand::Float64,
        :increase_level => [-111]::Vector{Int64},
        :increase_dims => [1 0 0 0 0]::Array{Int64,2},
        :adapt => false::Bool,
        :num_adapts => 1::Int64,
        :num_runs_adapt => 1000::Int64,
        :sim_length_adapt => 300::Int64,
        :increase_adapt_iter => [1]::Vector{Int64},
        :increase_level_adapt => [[502]]::Vector{Vector{Int64}},
        :increase_dims_adapt => [[2 0 0 0 0]]::Vector{Array{Int64,2}},
        :new_bounds_weight_up => [.3 .4 .5 .0 .0]::Array{Float64,2},
        :new_bounds_weight_down => [.1 .3 .5 .0 .5]::Array{Float64,2},
        :load => false::Bool,
        :time_horizon => 600::Int64,
        :sim_length => 101::Int64,
        :perfect_info => false::Bool,
        :learn_e => false::Bool,
        :robust => false::Bool,
        :theta => 3.9::Float64,
        :damages => :base::Symbol,
        :num_runs => 1000::Int64,
        :nqnodes => quad::Int64,
        :smax => [0. 0.; 0. 0.]::Array{Float64,2},
        :smin => [0. 0.; 0. 0.]::Array{Float64,2})

    # Solve the model
    space, c_store, s, mparams, opts, smax, smin = @time compute_value_function(inputs);

    if inputs[:nqnodes] == 11 && expand == 1.0
        mparams[:num_runs] = 50000
    else
        mparams[:num_runs] = 10000
    end

    @time consumption, abatement, inv_rate, capital, co2,
      prior_loc, prior_scale, damage, sim_values, loc_quad, dam_quad,
      tax_channels, channels_out_samples, vars =
      simulate_paths(space, c_store, mparams, opts, smax, smin);

    # Save variables
    if inputs[:nqnodes] == 11 && expand == 1.0
        bson("results/uncert_results.bson", Dict{String, Any}("inputs" => inputs, "mparams" => mparams, "smax" => smax, "smin" => smin,
          "tax_channels" => tax_channels, "consumption" => consumption, "abatement" => abatement,
          "capital" => capital, "co2" => co2, "prior_loc" => prior_loc,
          "prior_scale" => prior_scale, "damage" => damage, "opts" => opts,
          "channels_out_samples" => channels_out_samples, "space" => space, "c_store" => c_store))
    elseif expand > 1.0
        bson("results/uncert_results_wide.bson",
            Dict{String, Any}("tax_channels" => tax_channels, "consumption" => consumption))
    else
        bson("results/uncert_results_quad.bson",
            Dict{String, Any}("tax_channels" => tax_channels, "consumption" => consumption))
    end

end

function solve_learn(na = 1; load = false, quad = 11, adapt_iters = 20)

    if !load

        inputs = Dict{Symbol,Any}(:scaling_factor => 1.0::Float64,
            :initial_approximation_level => [3 5 3 7 5]::Array{Int64,2},
            :expand_domain => 1.::Float64,
            :increase_level => [-111]::Vector{Int64},
            :increase_dims => [1 0 0 0 0]::Array{Int64,2},
            :adapt => true::Bool,
            :num_adapts => adapt_iters::Int64,
            :num_runs_adapt => 1000::Int64,
            :sim_length_adapt => 300::Int64,
            :increase_adapt_iter => [17]::Vector{Int64},
            :increase_level_adapt => [[502]]::Vector{Vector{Int64}},
            :increase_dims_adapt => [[2 0 0 0 0]]::Vector{Array{Int64,2}},
            :new_bounds_weight_up => [.3 .4 .5 .0 .0]::Array{Float64,2},
            :new_bounds_weight_down => [.1 .3 .5 .0 .5]::Array{Float64,2},
            :load => false::Bool,
            :time_horizon => 600::Int64,
            :sim_length => 101::Int64,
            :perfect_info => false::Bool,
            :learn_e => true::Bool,
            :robust => false::Bool,
            :theta => 3.9::Float64,
            :damages => :base::Symbol,
            :num_runs => 1000::Int64,
            :nqnodes => 11::Int64,
            :smax => [0. 0. ; 0. 0.]::Array{Float64,2},
            :smin => [0. 0. ; 0. 0.]::Array{Float64,2})

        # Solve the model to get the final collocation bounds
        space, c_store, s, mparams, opts, smax, smin = @time compute_value_function(inputs)

        # Save variables
        bson("results/learning_results_adapt_$adapt_iters.bson",
            Dict{String, Any}("smax" => smax, "smin" => smin))

    else

        learn = BSON.load("results/learning_results_adapt_20.bson")
        smax = learn["smax"]
        smin = learn["smin"]

    end

    inputs = Dict{Symbol,Any}(:scaling_factor => 1.0::Float64,
        :initial_approximation_level => [3 5 3 7 5]::Array{Int64,2},
        :expand_domain => 1.::Float64,
        :increase_level => [-111]::Vector{Int64},
        :increase_dims => [1 0 0 0 0]::Array{Int64,2},
        :adapt => true::Bool,
        :num_adapts => 1::Int64,
        :num_runs_adapt => 1000::Int64,
        :sim_length_adapt => 300::Int64,
        :increase_adapt_iter => [1]::Vector{Int64},
        :increase_level_adapt => [[502]]::Vector{Vector{Int64}},
        :increase_dims_adapt => [[1 0 0 0 0]]::Vector{Array{Int64,2}},
        :new_bounds_weight_up => [.3 .4 .5 .0 .0]::Array{Float64,2},
        :new_bounds_weight_down => [.1 .3 .5 .0 .5]::Array{Float64,2},
        :load => true::Bool,
        :time_horizon => 600::Int64,
        :sim_length => 101::Int64,
        :perfect_info => false::Bool,
        :learn_e => true::Bool,
        :robust => false::Bool,
        :theta => 3.9::Float64,
        :damages => :base::Symbol,
        :num_runs => 1000::Int64,
        :nqnodes => quad::Int64,
        :smax => smax::Array{Float64,2},
        :smin => smin::Array{Float64,2})

    # Solve the model
    space, c_store, s, mparams, opts, smax, smin = @time compute_value_function(inputs)

    if inputs[:nqnodes] == 11 && adapt_iters == 20
        mparams[:num_runs] = 50000
    else
        mparams[:num_runs] = 10000
    end

    @time consumption, abatement, inv_rate, capital, co2,
        prior_loc, prior_scale, damage, sim_values, loc_quad, dam_quad,
        tax_channels, channels_out_samples, vars =
        simulate_paths(space, c_store, mparams, opts, smax, smin);

    # Save variables
    if inputs[:nqnodes] == 11
        bson("results/learning_results_adapt_$adapt_iters.bson", Dict{String, Any}("inputs" => inputs, "mparams" => mparams, "smax" => smax, "smin" => smin,
          "tax_channels" => tax_channels, "consumption" => consumption, "abatement" => abatement,
          "capital" => capital, "co2" => co2, "prior_loc" => prior_loc,
          "prior_scale" => prior_scale, "damage" => damage, "opts" => opts,
          "channels_out_samples" => channels_out_samples, "space" => space, "c_store" => c_store))
    else
        bson("results/learning_results_quad.bson",
            Dict{String, Any}("tax_channels" => tax_channels, "consumption" => consumption))
    end

end

function solve_rc(theta; expand = 1., quad = 11)

    inputs = Dict{Symbol,Any}(:scaling_factor => 1.06::Float64,
        :initial_approximation_level => [9 9 9 9 9]::Array{Int64,2},
        :expand_domain => expand::Float64,
        :increase_level => [-111]::Vector{Int64},
        :increase_dims => [1 0 0 0 0]::Array{Int64,2},
        :adapt => false::Bool,
        :num_adapts => 1::Int64,
        :num_runs_adapt => 1000::Int64,
        :sim_length_adapt => 300::Int64,
        :increase_adapt_iter => [1]::Vector{Int64},
        :increase_level_adapt => [[-111]]::Vector{Vector{Int64}},
        :increase_dims_adapt => [[2 0 0 0 0]]::Vector{Array{Int64,2}},
        :new_bounds_weight_up => [.3 .4 .5 .0 .0]::Array{Float64,2},
        :new_bounds_weight_down => [.1 .3 .5 .0 .5]::Array{Float64,2},
        :load => false::Bool,
        :time_horizon => 600::Int64,
        :sim_length => 101::Int64,
        :perfect_info => false::Bool,
        :learn_e => false::Bool,
        :robust => true::Bool,
        :theta => theta::Float64,
        :damages => :base::Symbol,
        :num_runs => 1000::Int64,
        :nqnodes => quad::Int64,
        :smax => [0. 0.; 0. 0.]::Array{Float64,2},
        :smin => [0. 0.; 0. 0.]::Array{Float64,2})

    # Solve the model
    space, c_store, s, mparams, opts, smax, smin = @time compute_value_function(inputs);

    if inputs[:nqnodes] == 11 && theta < 4.0 && expand == 1.0
        mparams[:num_runs] = 50000
    elseif theta < 4.0
        mparams[:num_runs] = 10000
    else
        mparams[:num_runs] = 10
    end

    @time consumption, abatement, inv_rate, capital, co2,
      prior_loc, prior_scale, damage, sim_values, loc_quad, dam_quad,
      tax_channels, channels_out_samples, vars =
      simulate_paths(space, c_store, mparams, opts, smax, smin);

    # Save variables
    if inputs[:nqnodes] == 11 && expand == 1.0
        bson("results/rc_results_$theta.bson", Dict{String, Any}("inputs" => inputs, "mparams" => mparams, "smax" => smax, "smin" => smin,
          "tax_channels" => tax_channels, "consumption" => consumption, "abatement" => abatement,
          "capital" => capital, "co2" => co2, "prior_loc" => prior_loc,
          "prior_scale" => prior_scale, "damage" => damage, "opts" => opts,
          "channels_out_samples" => channels_out_samples, "space" => space, "c_store" => c_store))
    elseif expand > 1
        bson("results/rc_results_$(theta)_wide.bson",
            Dict{String, Any}("tax_channels" => tax_channels, "consumption" => consumption))
    else
        bson("results/rc_results_$(theta)_quad.bson",
            Dict{String, Any}("tax_channels" => tax_channels, "consumption" => consumption))
    end

end

function solve_rcl(theta; quad = 11, adapt_iters = 20)

    learn = BSON.load("results/learning_results_adapt_$(adapt_iters).bson")

    inputs = Dict{Symbol,Any}(:scaling_factor => quad == 11 ? 2.9 : 3.1::Float64,
        :initial_approximation_level => [3 5 3 7 5]::Array{Int64,2},
        :expand_domain => 1.::Float64,
        :increase_level => [-111]::Vector{Int64},
        :increase_dims => [1 0 0 0 0]::Array{Int64,2},
        :adapt => true::Bool,
        :num_adapts => 1::Int64,
        :num_runs_adapt => 1000::Int64,
        :sim_length_adapt => 300::Int64,
        :increase_adapt_iter => [1]::Vector{Int64},
        :increase_level_adapt => [[502]]::Vector{Vector{Int64}},
        :increase_dims_adapt => [[1 0 0 0 0]]::Vector{Array{Int64,2}},
        :new_bounds_weight_up => [.3 .4 .5 .0 .0]::Array{Float64,2},
        :new_bounds_weight_down => [.1 .3 .5 .0 .5]::Array{Float64,2},
        :load => true::Bool,
        :time_horizon => 600::Int64,
        :sim_length => 101::Int64,
        :perfect_info => false::Bool,
        :learn_e => true::Bool,
        :robust => true::Bool,
        :theta => theta::Float64,
        :damages => :base::Symbol,
        :num_runs => 1000::Int64,
        :nqnodes => quad::Int64,
        :smax => learn["smax"]::Array{Float64,2},
        :smin => learn["smin"]::Array{Float64,2})

    # Solve the model
    space, c_store, s, mparams, opts, smax, smin = @time compute_value_function(inputs)

    if inputs[:nqnodes] == 11 && theta < 4.0 && adapt_iters == 20
        mparams[:num_runs] = 50000
    elseif theta < 4.0
        mparams[:num_runs] = 10000
    else
        mparams[:num_runs] = 10
    end

    @time consumption, abatement, inv_rate, capital, co2,
        prior_loc, prior_scale, damage, sim_values, loc_quad, dam_quad,
        tax_channels, channels_out_samples, vars =
        simulate_paths(space, c_store, mparams, opts, smax, smin);

    # Save variables
    if inputs[:nqnodes] == 11
        bson("results/rcl_results_$(theta)_adapt_$adapt_iters.bson", Dict{String, Any}("inputs" => inputs, "mparams" => mparams, "smax" => smax, "smin" => smin,
          "tax_channels" => tax_channels, "consumption" => consumption, "abatement" => abatement,
          "capital" => capital, "co2" => co2, "prior_loc" => prior_loc,
          "prior_scale" => prior_scale, "damage" => damage, "opts" => opts,
          "channels_out_samples" => channels_out_samples, "space" => space, "c_store" => c_store))
    else
        bson("results/rcl_results_$(theta)_quad.bson",
            Dict{String, Any}("tax_channels" => tax_channels, "consumption" => consumption))
    end

end

##################################
##################################
# Figures 5b, 5c, 5d
##################################
##################################

function simulate_alt_damage_params(model_path, d1, d2)

    model = BSON.load(model_path)

    model["mparams"][:num_runs] = 5000

    if model["opts"][:robust] && model["opts"][:learn_e]
        type = "rcl"
    elseif model["opts"][:robust] && !model["opts"][:learn_e]
        type = "rc"
    elseif !model["opts"][:robust] && model["opts"][:learn_e]
        type = "learn"
    else
        type = "uncert"
    end

    @time consumption, abatement, inv_rate, capital, co2,
        prior_loc, prior_scale, damage, sim_values, loc_quad, dam_quad,
        tax_channels, channels_out_samples, vars = simulate_paths(model["space"], model["c_store"], model["mparams"],
        model["opts"], model["smax"], model["smin"],
        alt_damages = true, d1 = d1, d2 = d2);

    # Save variables
    bson("results/" * type * "_results_dam_d1_$(d1)_d2_$(d2).bson",
        Dict{String, Any}("mparams" => model["mparams"], "tax_channels" => tax_channels,
        "consumption" => consumption, "abatement" => abatement,
        "capital" => capital, "co2" => co2, "prior_loc" => prior_loc,
        "prior_scale" => prior_scale, "damage" => damage))

end

##################################
##################################
# Figure 8
##################################
##################################

function simulate_misspecified_damages(model_path, damage_type)

    model = BSON.load(model_path)

    model["mparams"][:num_runs] = 10000
    model["mparams"][:sim_length] = 201

    if model["opts"][:robust] && model["opts"][:learn_e]
        type = "rcl"
    elseif model["opts"][:robust] && !model["opts"][:learn_e]
        type = "rc"
    elseif !model["opts"][:robust] && model["opts"][:learn_e]
        type = "learn"
    else
        type = "uncert"
    end

    @time consumption, abatement, inv_rate, capital, co2,
        prior_loc, prior_scale, damage, sim_values, loc_quad, dam_quad,
        tax_channels, channels_out_samples, vars =
        simulate_paths(model["space"], model["c_store"], model["mparams"],
        model["opts"], model["smax"], model["smin"],
        misspec_damages = damage_type);

    # Save variables
    bson("results/" * type * "_results_dam_misspec_$(string(damage_type)).bson",
        Dict{String, Any}("mparams" => model["mparams"], "tax_channels" => tax_channels,
        "consumption" => consumption, "abatement" => abatement,
        "capital" => capital, "co2" => co2, "prior_loc" => prior_loc,
        "prior_scale" => prior_scale, "damage" => damage))

end

##################################
##################################
# Figures 6a and 6b
##################################
##################################

function solve_alt_penalty_parameter_rc()

    thetas = [6.0, 8.0, 10.0, sqrt(10.)*10., 100.0, sqrt(10.)*100., 1000.0]

    for theta in thetas

        solve_rc(theta)

    end

end

function solve_alt_penalty_parameter_rcl()

    thetas = [6.0, 8.0, 10.0, sqrt(10.)*10., 100.0, sqrt(10.)*10., 1000.0]

    for theta in thetas

        solve_rcl(theta)

    end

end

##################################
##################################
# Table A4: 1st panel
##################################
##################################

function solve_uncert_error_coll()

inc_dims_vec = [[1 0 0 0 0],
                [0 1 0 0 0],
                [0 0 0 0 1],
                [-1 0 0 0 0]]

    for (index, inc_dims) in enumerate(inc_dims_vec)

        inputs = Dict{Symbol,Any}(:scaling_factor => 1.0::Float64,
            :initial_approximation_level => [9 9 9 9 9]::Array{Int64,2},
            :expand_domain => 1.::Float64,
            :increase_level => [102]::Vector{Int64},
            :increase_dims => inc_dims::Array{Int64,2},
            :adapt => false::Bool,
            :num_adapts => 1::Int64,
            :num_runs_adapt => 1000::Int64,
            :sim_length_adapt => 300::Int64,
            :increase_adapt_iter => [1]::Vector{Int64},
            :increase_level_adapt => [[502]]::Vector{Vector{Int64}},
            :increase_dims_adapt => [[2 0 0 0 0]]::Vector{Array{Int64,2}},
            :new_bounds_weight_up => [.3 .4 .5 .0 .0]::Array{Float64,2},
            :new_bounds_weight_down => [.1 .3 .5 .0 .5]::Array{Float64,2},
            :load => false::Bool,
            :time_horizon => 600::Int64,
            :sim_length => 101::Int64,
            :perfect_info => false::Bool,
            :learn_e => false::Bool,
            :robust => false::Bool,
            :theta => 3.9::Float64,
            :damages => :base::Symbol,
            :num_runs => 1000::Int64,
            :nqnodes => 11::Int64,
            :smax => [0. 0.; 0. 0.]::Array{Float64,2},
            :smin => [0. 0.; 0. 0.]::Array{Float64,2})

        # Solve the model
        space, c_store, s, mparams, opts, smax, smin = @time compute_value_function(inputs);

        mparams[:num_runs] = 10000
        @time consumption, abatement, inv_rate, capital, co2,
          prior_loc, prior_scale, damage, sim_values, loc_quad, dam_quad,
          tax_channels, channels_out_samples, vars =
          simulate_paths(space, c_store, mparams, opts, smax, smin);

        # Save variables
        bson("results/uncert_results_coll_$index.bson",
            Dict{String, Any}("tax_channels" => tax_channels, "consumption" => consumption))

    end

end

function solve_learn_error_coll()

    inc_dims_vec = [[1 0 0 0 0; 1 0 0 0 0],
                    [1 0 0 0 0; 0 1 0 0 0],
                    [1 0 0 0 0; 0 0 1 0 0],
                    [1 0 0 0 0; 0 0 0 1 0],
                    [1 0 0 0 0; 0 0 0 0 1],
                    [1 0 0 0 0; -1 0 0 0 0]]

    for (index, inc_dims) in enumerate(inc_dims_vec)

        learn = BSON.load("results/learning_results_adapt_20.bson")

        inputs = Dict{Symbol,Any}(:scaling_factor => 1.0::Float64,
            :initial_approximation_level => [3 5 3 7 5]::Array{Int64,2},
            :expand_domain => 1.::Float64,
            :increase_level => [-111]::Vector{Int64},
            :increase_dims => [1 0 0 0 0]::Array{Int64,2},
            :adapt => true::Bool,
            :num_adapts => 1::Int64,
            :num_runs_adapt => 1000::Int64,
            :sim_length_adapt => 300::Int64,
            :increase_adapt_iter => [1]::Vector{Int64},
            :increase_level_adapt => [[502, 102]]::Vector{Vector{Int64}},
            :increase_dims_adapt => [inc_dims]::Vector{Array{Int64,2}},
            :new_bounds_weight_up => [.3 .4 .5 .0 .0]::Array{Float64,2},
            :new_bounds_weight_down => [.1 .3 .5 .0 .5]::Array{Float64,2},
            :load => true::Bool,
            :time_horizon => 600::Int64,
            :sim_length => 101::Int64,
            :perfect_info => false::Bool,
            :learn_e => true::Bool,
            :robust => false::Bool,
            :theta => 3.9::Float64,
            :damages => :base::Symbol,
            :num_runs => 1000::Int64,
            :nqnodes => 11::Int64,
            :smax => learn["smax"]::Array{Float64,2},
            :smin => learn["smin"]::Array{Float64,2})

        # Solve the model
        space, c_store, s, mparams, opts, smax, smin = @time compute_value_function(inputs)

        mparams[:num_runs] = 10000

        @time consumption, abatement, inv_rate, capital, co2,
            prior_loc, prior_scale, damage, sim_values, loc_quad, dam_quad,
            tax_channels, channels_out_samples, vars =
            simulate_paths(space, c_store, mparams, opts, smax, smin);

        # Save variables
        bson("results/learning_results_coll_$index.bson",
            Dict{String, Any}("tax_channels" => tax_channels, "consumption" => consumption))

    end

end

function solve_rc_error_coll()

inc_dims_vec = [[1 0 0 0 0],
                [0 1 0 0 0],
                [0 0 0 0 1],
                [-1 0 0 0 0]]

    for (index, inc_dims) in enumerate(inc_dims_vec)

        inputs = Dict{Symbol,Any}(:scaling_factor => 1.06::Float64,
            :initial_approximation_level => [9 9 9 9 9]::Array{Int64,2},
            :expand_domain => 1.0::Float64,
            :increase_level => [502]::Vector{Int64},
            :increase_dims => inc_dims::Array{Int64,2},
            :adapt => false::Bool,
            :num_adapts => 1::Int64,
            :num_runs_adapt => 1000::Int64,
            :sim_length_adapt => 300::Int64,
            :increase_adapt_iter => [1]::Vector{Int64},
            :increase_level_adapt => [[502]]::Vector{Vector{Int64}},
            :increase_dims_adapt => [[2 0 0 0 0]]::Vector{Array{Int64,2}},
            :new_bounds_weight_up => [.3 .4 .5 .0 .0]::Array{Float64,2},
            :new_bounds_weight_down => [.1 .3 .5 .0 .5]::Array{Float64,2},
            :load => false::Bool,
            :time_horizon => 600::Int64,
            :sim_length => 101::Int64,
            :perfect_info => false::Bool,
            :learn_e => false::Bool,
            :robust => true::Bool,
            :theta => 3.9::Float64,
            :damages => :base::Symbol,
            :num_runs => 1000::Int64,
            :nqnodes => 11::Int64,
            :smax => [0. 0.; 0. 0.]::Array{Float64,2},
            :smin => [0. 0.; 0. 0.]::Array{Float64,2})

        # Solve the model
        space, c_store, s, mparams, opts, smax, smin = @time compute_value_function(inputs);

        mparams[:num_runs] = 10000

        @time consumption, abatement, inv_rate, capital, co2,
          prior_loc, prior_scale, damage, sim_values, loc_quad, dam_quad,
          tax_channels, channels_out_samples, vars =
          simulate_paths(space, c_store, mparams, opts, smax, smin);

        # Save variables
        bson("results/rc_results_coll_$index.bson",
            Dict{String, Any}("tax_channels" => tax_channels, "consumption" => consumption))

    end

end

function solve_rcl_error_coll()

    inc_dims_vec = [[1 0 0 0 0; 1 0 0 0 0],
                    [1 0 0 0 0; 0 1 0 0 0],
                    [1 0 0 0 0; 0 0 1 0 0],
                    [1 0 0 0 0; 0 0 0 1 0],
                    [1 0 0 0 0; 0 0 0 0 1],
                    [1 0 0 0 0; -1 0 0 0 0]]

    for (index, inc_dims) in enumerate(inc_dims_vec)

        learn = BSON.load("results/learning_results_adapt_20.bson")

        inputs = Dict{Symbol,Any}(:scaling_factor => 6.4::Float64,
            :initial_approximation_level => [3 5 3 7 5]::Array{Int64,2},
            :expand_domain => 1.::Float64,
            :increase_level => [-111]::Vector{Int64},
            :increase_dims => [1 0 0 0 0]::Array{Int64,2},
            :adapt => true::Bool,
            :num_adapts => 1::Int64,
            :num_runs_adapt => 1000::Int64,
            :sim_length_adapt => 300::Int64,
            :increase_adapt_iter => [1]::Vector{Int64},
            :increase_level_adapt => [[502, 102]]::Vector{Vector{Int64}},
            :increase_dims_adapt => [inc_dims]::Vector{Array{Int64,2}},
            :new_bounds_weight_up => [.3 .4 .5 .0 .0]::Array{Float64,2},
            :new_bounds_weight_down => [.1 .3 .5 .0 .5]::Array{Float64,2},
            :load => true::Bool,
            :time_horizon => 600::Int64,
            :sim_length => 101::Int64,
            :perfect_info => false::Bool,
            :learn_e => true::Bool,
            :robust => true::Bool,
            :theta => 3.9::Float64,
            :damages => :base::Symbol,
            :num_runs => 1000::Int64,
            :nqnodes => 11::Int64,
            :smax => learn["smax"]::Array{Float64,2},
            :smin => learn["smin"]::Array{Float64,2})

        # Solve the model
        space, c_store, s, mparams, opts, smax, smin = @time compute_value_function(inputs)

        mparams[:num_runs] = 10000

        @time consumption, abatement, inv_rate, capital, co2,
            prior_loc, prior_scale, damage, sim_values, loc_quad, dam_quad,
            tax_channels, channels_out_samples, vars =
            simulate_paths(space, c_store, mparams, opts, smax, smin);

        # Save variables
        bson("results/rcl_results_coll_$index.bson",
            Dict{String, Any}("tax_channels" => tax_channels, "consumption" => consumption))

    end

end

##################################
##################################
# Figure A2
##################################
##################################

function solve_learn_previous_iter(quad = 11, adapt_iters = 19)

    inputs = Dict{Symbol,Any}(:scaling_factor => 1.0::Float64,
        :initial_approximation_level => [3 5 3 7 5]::Array{Int64,2},
        :expand_domain => 1.::Float64,
        :increase_level => [-111]::Vector{Int64},
        :increase_dims => [1 0 0 0 0]::Array{Int64,2},
        :adapt => true::Bool,
        :num_adapts => adapt_iters::Int64,
        :num_runs_adapt => 1000::Int64,
        :sim_length_adapt => 300::Int64,
        :increase_adapt_iter => [17]::Vector{Int64},
        :increase_level_adapt => [[502]]::Vector{Vector{Int64}},
        :increase_dims_adapt => [[2 0 0 0 0]]::Vector{Array{Int64,2}},
        :new_bounds_weight_up => [.3 .4 .5 .0 .0]::Array{Float64,2},
        :new_bounds_weight_down => [.1 .3 .5 .0 .5]::Array{Float64,2},
        :load => false::Bool,
        :time_horizon => 600::Int64,
        :sim_length => 101::Int64,
        :perfect_info => false::Bool,
        :learn_e => true::Bool,
        :robust => false::Bool,
        :theta => 3.9::Float64,
        :damages => :base::Symbol,
        :num_runs => 1000::Int64,
        :nqnodes => 11::Int64,
        :smax => [0. 0. ; 0. 0.]::Array{Float64,2},
        :smin => [0. 0. ; 0. 0.]::Array{Float64,2})

    # Solve the model
    space, c_store, s, mparams, opts, smax, smin = @time compute_value_function(inputs)

    # Save variables
    bson("results/learning_results_adapt_$adapt_iters.bson",
        Dict{String, Any}("smax" => smax, "smin" => smin))

end

max_stars_repo_name: irudik/optimal-climate-policy-aej

max_issues_repo_name: irudik/optimal-climate-policy-aej

max_forks_repo_name: irudik/optimal-climate-policy-aej


text:
module TestTLS

using Base.Test
using Distributions
using TimeData
using NLopt
using JuMP

import Econometrics

## reload(joinpath(Pkg.dir(), "Econometrics", "src", "Econometrics.jl"))

## constructors / display
##-----------------------

tls = Econometrics.TLSDist([3.2, 0.2, 1.2])
tls = Econometrics.TLSDist(3.2, 0.2, 1.2)

println("\nTesting display output:")
display(tls)

@test_throws Exception Econometrics.TLSDist(-0.2, 0.2, 1.2)
@test_throws Exception Econometrics.TLSDist(0.2, 0.2, -1.2)

## dof
##----

tls = Econometrics.TLSDist([3.2, 0.2, 1.2])
@test Econometrics.dof(tls) == 3.2

## getParams
##----------

@test Econometrics.getParams(tls) == (3.2, 0.2, 1.2)

## rand
##-----

nSim = 5000
tls = Econometrics.TLSDist([3.2, 0.2, 1.2])
simVals = rand(tls, nSim)

## pdf
##----

tls = Econometrics.TLSDist([3.2, 0.2, 1.2])
pdfVal = Econometrics.pdf(tls, 0.2)

@test isa(pdfVal, Float64)

pdfVals = Econometrics.pdf(tls, simVals)

@test isa(pdfVals, Array{Float64, 1})
@test all(pdfVals .>= 0)

## cdf
##----

tls = Econometrics.TLSDist([3.2, 0.2, 1.2])
cdfVal = Econometrics.cdf(tls, 0.2)

@test isa(cdfVal, Float64)

cdfVals = Econometrics.cdf(tls, simVals)

@test isa(cdfVals, Array{Float64, 1})
@test all(cdfVals .>= 0)
@test all(cdfVals .<= 1)

## quantile
##---------

tls = Econometrics.TLSDist([3.2, 0.2, 1.2])
pVal = Econometrics.quantile(tls, 0.2)

@test isa(pVal, Float64)

## test cdf(quantile(x)) and quantile(cdf(x))
##-------------------------------------------

nSim = 5000
tls = Econometrics.TLSDist([3.2, 0.2, 1.2])
simVals = rand(tls, nSim)
cdfVals = Econometrics.cdf(tls, simVals)
pVals = Econometrics.quantile(tls, cdfVals)
@test_approx_eq pVals simVals

grid = [0.01:0.01:0.99]
pVals = Econometrics.quantile(tls, grid)
cdfVals = Econometrics.cdf(tls, pVals)
@test_approx_eq cdfVals grid

## nllh
##-----

nSim = 5000
tls = Econometrics.TLSDist([3.2, 0.2, 1.2])
simVals = rand(tls, nSim)

nllhVal1 = Econometrics.nllh(tls, simVals)
nllhVal2 = Econometrics.nllh_tls([Econometrics.getParams(tls)...], simVals)

@test_approx_eq nllhVal1 nllhVal2

@test_throws Exception Econometrics.nllh_tls([3.2, 0.2, 1.2, 4.2], simVals)

## fit
##----

nSim = 1000
params = [3.2, 0.2, 1.2]
tls = Econometrics.TLSDist(params)

@time tls1 = Econometrics.fit(Econometrics.TLSDist, simVals)
@time tls2 = Econometrics.fit_jump(Econometrics.TLSDist, simVals)

@test isa(tls1, Econometrics.TLSDist)
@test isa(tls2, Econometrics.TLSDist)

## compare parameters
[params'
[Econometrics.getParams(tls1)...]'
[Econometrics.getParams(tls2)...]']

## compare nllhs
Econometrics.nllh(tls1, simVals)
Econometrics.nllh(tls2, simVals)

end


text:
import Pkg
Pkg.add("Documenter")
using Documenter
using MolecularMinimumDistances
push!(LOAD_PATH, "../src/")
makedocs(
    modules = [MolecularMinimumDistances],
    sitename = "MolecularMinimumDistances.jl",
    pages = [
        "Home" => "index.md",
        "Basic use" => "basic.md",
        "Advanced use" => "advanced.md",
        "Reference" => "reference.md",
    ],
)
deploydocs(
    repo = "github.com/m3g/MolecularMinimumDistances.jl.git",
    target = "build",
    branch = "gh-pages",
    versions = ["stable" => "v^", "v#.#"],
)


text:
######### Stan ODE example  ###########
#
# This is very experimental.
#
# Running this example with a single chain often returns a completed simulation.
# As more chains are added, you will find the ODE solver and/or Stan HMC might not finish,
# If this happens, in the REPL, ^C will terminate the run. Trying it a few more time
# will typically lead to a successful simulation. On my system completion takes < 1 minute.
#
# Bob Carpenter suggested to lower the stepsize and increase acceptance rate. This can
# be done as shown below (commented out lines 97 and 97). Limited testing has not shown
# guaranteed improvements in successful completion of the simulation.
#

using Stan, Distributions

ProjDir = dirname(@__FILE__)
cd(ProjDir) do

  odemodel = "
  functions {
         real[] sho(real t,
                    real[] y,
                    real[] theta,
                    real[] x_r,
                    int[] x_i) {
           real dydt[2];
           dydt[1] <- y[2];
           dydt[2] <- -y[1] - theta[1] * y[2];
           return dydt;
          }
  }

  data {
    int<lower=1> T;
    int<lower=1> M;
    real t0;
    #real y0[M];
    real ts[T];
    real y[T,M];
  }

  transformed data {
    real x_r[0];
    int x_i[0];
  }

  parameters {
    real y0[2];
    vector<lower=0>[2] sigma;
    real theta[1];
  }

  model {
    real y_hat[T,2];
    sigma ~ cauchy(0,2.5);
    theta ~ normal(0,1);
    y0 ~ normal(0,1);
    y_hat <- integrate_ode(sho, y0, t0, ts, theta, x_r, x_i);
    for (t in 1:T)
     y[t] ~ normal(y_hat[t], sigma);
  }
  "

  #Just store date in ho.csv, drop the header
  ho_data = readdlm("ho.csv.data", header=true)[1];

  #Extract initial values in row 1 of ho_data
  t0 = ho_data[1, 1]
  y0 = reshape(ho_data[1, 2:3], 2)

  # Now we can drop the first row of ho_data
  ts = ho_data[2:end, 1]
  ho_data = ho_data[2:end, 2:3]

  #Get the proper sizes to pass to CmdStan
  T = size(ho_data, 1)
  M = size(ho_data, 2)

  #Prepare to add the noise
  sigma = 0.15
  noise = [rand(Normal(0, sigma)) for i in 1:T, j in 1:M]
  y = ho_data + noise

  odedict =
    Dict(
      "T" => T,
      "M" => M,
      "t0" => t0,
      "y0" => y0,
      "ts" => ts,
      "y" => y)

  global stanmodel, rc, sim1
  stanmodel = Stanmodel(name="ode", model=odemodel, nchains=2);
  #stanmodel.method.algorithm.stepsize = 0.7 # default stepsize=1.0
  #stanmodel.method.adapt.delta = 0.9 # default delta=0.8

  println("\nStanmodel that will be used:")
  stanmodel |> display
  println()

  rc, sim1 = stan(stanmodel, [odedict], ProjDir, diagnostics=false,
    CmdStanDir=CMDSTAN_HOME);

end # cd


text:
using Roots # for fzero fction

"""
   vbmf_trial_parameters

Compound type for vbmf_trial computation. Contains the following fields:\n

    L::Int - number of rows of the original matrix
    M::Int - number of columns of the original matrix
    M0::Int - columns belonging to negative bags
    M1::Int - columns belonging to positive bags
    MH::Int - columns*internal factorization dimension
    H::Int - internal product dimension, H = H0 + H1
    H0::Int - size of base of noncontaminated class
    H1::Int - size of base of contaminated class

    AHat::Array{Float64, 2} - mean value of A, size (M, H), A = [A0 A1]
    ATVecHat::Array{Float64, 1} - mean value of vec(A^T), size (MH,1)
    SigmaATVec::Array{Float64, 2} - covariance of vec(A^T), size (MH, MH)
    diagSigmaATVec::Array{Float64,1} - diagonal of covariance of vec(A^T), size (MH)
    invSigmaATVec::Array{Float64, 2} - inverse of covariance of vec(A^T), size (MH, MH)
    SigmaA::Array{Float64,2} - covariance of A matrix, size (H, H)
    A1Hat::Array{Float64, 2} - mean value of A1, size (M, H0)
    A2Hat::Array{Float64, 2} - mean value of A2, size (M0, H1)
    A3Hat::Array{Float64, 2} - mean value of A3, size (M1, H1)

    BHat::Array{Float64, 2} - mean value of B, size (L, H)
    SigmaB::Array{Float64, 2} - covariance of B matrix, size (H, H)

    CA::Array{Float64, 1} - diagonal of the prior covariance of vec(A^T), size (MH)
    alpha::Array{Float64, 1} - shapes of CA gamma posterior, size (2,1)
    beta::Array{Float64,1} - scale of CA gamma posterior, size (MH, 1)

    CA1::Array{Float64, 1} - prior covariance of A1, size (MH0)
    alpha01::Float64 - shape of CA1 gamma prior
    beta01::Float64 - scale of CA1 gamma prior
    alpha1::Float64 - shape of CA1 gamma posterior
    beta1::::Array{Float64,1} - scale of CA1 gamma posterior, size (MH0, 1)
    CA2::Array{Float64, 1} - prior covariance of A2, size (M0H1)
    alpha02::Float64 - shape of CA2 gamma prior
    beta02::Float64 - scale of CA2 gamma prior
    alpha2::Float64 - shape of CA2 gamma posterior
    beta2::::Array{Float64,1} - scale of CA2 gamma posterior, size (M0H1, 1)
    CA3::Array{Float64, 1} - prior covariance of A3, size (M1H1)
    alpha03::Float64 - shape of CA3 gamma prior
    beta03::Float64 - scale of CA3 gamma prior
    alpha3::Float64 - shape of CA3 gamma posterior
    beta3::::Array{Float64,1} - scale of CA3 gamma posterior, size (M1H1, 1)

    CB::Array{Float64, 1} - diagonal of prior covariance of B, size (H)
    gamma0::Float64 - shape of CB gamma prior
    delta0::Float64 - scale of CB gamma prior
    gamma::Float64 - shape of CB posterior
    delta::Array{Float64,1} - scale of CB posterior, size (H)

    sigmaHat::Float64 - variance of data - homoscedastic case
    eta0::Float64 - shape of sigma gamma prior
    zeta0::Float64 - scale of sigma gamma prior
    eta::Float64 - shape of homoscedastic sigma posterior
    zeta::Float64 - scale of homoscedastic sigma posterior
    sigmaVecHat::Array{Float64,1} - variance of data (rows) - heteroscedastic case, size (L)
    etaVec::Array{Float64,1} - shape of heteroscedastic sigma posterior, size (L)
    zetaVec::Array{Float64,1} - scale of heteroscedastic sigma posterior, size (L)

    YHat::Array{Float64, 2} - estimate of Y, size (L, M)
    trYTY::Float64 - trace(Y^T*Y), saved so that it does not have to be recomputed
"""
type vbmf_trial_parameters
    L::Int
    M::Int
    M0::Int
    M1::Int
    MH::Int
    H::Int
    H0::Int
    H1::Int

    AHat::Array{Float64, 2}
    ATVecHat::Array{Float64, 1}
    SigmaATVec::Array{Float64, 2}
    diagSigmaATVec::Array{Float64,1}
    invSigmaATVec::Array{Float64, 2}
    SigmaA::Array{Float64,2}
    A1Hat::Array{Float64, 2}
    A2Hat::Array{Float64, 2}
    A3Hat::Array{Float64, 2}

    BHat::Array{Float64, 2}
    SigmaB::Array{Float64, 2}

    CA::Array{Float64, 1}
    alpha::Array{Float64, 1}
    beta::Array{Float64,1}
    CA1::Array{Float64, 1}
    alpha01::Float64
    beta01::Float64
    alpha1::Float64
    beta1::Array{Float64,1}
    CA2::Array{Float64, 1}
    alpha02::Float64
    beta02::Float64
    alpha2::Float64
    beta2::Array{Float64,1}
    CA3::Array{Float64, 1}
    alpha03::Float64
    beta03::Float64
    alpha3::Float64
    beta3::Array{Float64,1}

    CB::Array{Float64, 1}
    gamma0::Float64
    delta0::Float64
    gamma::Float64
    delta::Array{Float64,1}

    sigmaHat::Float64
    eta0::Float64
    zeta0::Float64
    eta::Float64
    zeta::Float64
    sigmaVecHat::Array{Float64,1}
    etaVec::Array{Float64,1}
    zetaVec::Array{Float64,1}

    YHat::Array{Float64, 2}
    trYTY::Float64

    vbmf_trial_parameters() = new()
end

"""
    vbmf_trial_init(Y::Array{Float64,2}, H::Int, H0::Int, M0::Int; ca::Float64 = 1.0,
    alpha0::Float64 = 1e-10, beta0::Float64 = 1e-10, cb::Float64 = 1.0,
    gamma0::Float64 = 1e-10, delta0::Float64 = 1e-10,
    sigma::Float64 = 1.0, eta0::Float64 = 1e-10, zeta0::Float64 = 1e-10)

Returns an initialized structure of type vbmf_trial_parameters.
"""
function vbmf_trial_init(Y::Array{Float64,2}, H::Int, H0::Int, M0::Int; ca::Float64 = 1.0,
    alpha0::Float64 = 1e-10, beta0::Float64 = 1e-10, cb::Float64 = 1.0,
    gamma0::Float64 = 1e-10, delta0::Float64 = 1e-10,
    sigma::Float64 = 1.0, eta0::Float64 = 1e-10, zeta0::Float64 = 1e-10)
    if H < H0
        error("H must be at least H0!")
    end

    params = vbmf_trial_parameters()
    L, M = size(Y)

    params.L, params.M = L, M
    params.H = H
    params.M0 = M0
    M1 = M - M0
    params.M1 = M1
    params.MH = M*H
    params.H0 = H0
    H1 = H - H0
    params.H1 = H1

    params.AHat = randn(M, H)
    params.ATVecHat = reshape(params.AHat', M*H)
    params.SigmaATVec = eye(M*H, M*H)
    params.diagSigmaATVec = ones(M*H)
    params.invSigmaATVec = eye(M*H, M*H)
    params.SigmaA = zeros(H, H)
    params.A1Hat = params.AHat[:,1:H0]
    params.A2Hat = params.AHat[1:M0,(params.H0+1):end]
    params.A3Hat = params.AHat[(M0+1):end,(params.H0+1):end]

    params.BHat = randn(L, H)
    params.SigmaB = zeros(H, H)

    params.CA1 = ca*ones(M*H0)
    params.CA2 = ca*ones(M0*H1)
    params.CA3 = ca*ones(M1*H1)
    params.CA = Array{Float64, 1}()
    for m in 1:params.M0
        params.CA = cat(1, params.CA, params.CA1[((m-1)*params.H0+1):m*params.H0])
        params.CA = cat(1, params.CA, params.CA2[((m-1)*params.H1+1):m*params.H1])
    end
    for m in (params.M0+1):params.M
        params.CA = cat(1, params.CA, params.CA1[((m-1)*params.H0+1):m*params.H0])
        params.CA = cat(1, params.CA, params.CA3[((m-1-params.M0)*params.H1+1):(m-params.M0)*params.H1])
    end
    params.alpha01 = alpha0
    params.alpha02 = alpha0
    params.alpha03 = alpha0
    params.beta01 = beta0
    params.beta02 = beta0
    params.beta03 = beta0
    params.alpha1 = alpha0 + 0.5
    params.alpha2 = alpha0 + 0.5
    params.alpha3 = alpha0 + 0.5
    params.beta1 = beta0*ones(M*H0)
    params.beta2 = beta0*ones(M0*H1)
    params.beta3 = beta0*ones(M1*H1)
    params.alpha = [params.alpha1, params.alpha2, params.alpha3]
    params.beta = Array{Float64, 1}()
    for m in 1:params.M0
        params.beta = cat(1, params.beta, params.beta1[((m-1)*params.H0+1):m*params.H0])
        params.beta = cat(1, params.beta, params.beta2[((m-1)*params.H1+1):m*params.H1])
    end
    for m in (params.M0+1):params.M
        params.beta = cat(1, params.beta, params.beta1[((m-1)*params.H0+1):m*params.H0])
        params.beta = cat(1, params.beta, params.beta3[((m-1-params.M0)*params.H1+1):(m-params.M0)*params.H1])
    end

    params.CB = cb*ones(H)
    params.gamma0 = gamma0
    params.delta0 = delta0
    params.gamma = gamma0 + L/2
    params.delta = delta0*ones(H)

    params.sigmaHat = sigma
    params.eta0 = eta0
    params.zeta0 = zeta0
    params.eta = eta0 + L*M/2
    params.zeta = zeta0
    params.sigmaVecHat = sigma*ones(L)
    params.etaVec = (eta0 + M/2)*ones(L)
    params.zetaVec = zeta0*ones(L)

    params.YHat = params.BHat*params.AHat'
    params.trYTY = traceXTY(Y, Y)

    return params
end

"""
   copy(params_in::vbmf_trial_parameters)

Copy function for vbmfa_trial_parameters.
"""
function copy(params_in::vbmf_trial_parameters)
    params = vbmf_trial_parameters()

    for field in fieldnames(params_in)
        setfield!(params, field, copy(getfield(params_in, field)))
    end

    return params
end

"""
    updateA!(Y::Array{Float64,2}, params::vbmf_trial_parameters; full_cov::Bool, diag_var::Bool = false)

Updates mean and covariance of vec(A^T) and also of the A matrix. If full_cov is true,
then inverse of full covariance matrix is computed, otherwise just the diagonal is estimated.
"""
function updateA!(Y::Array{Float64,2}, params::vbmf_trial_parameters; full_cov::Bool = false, diag_var::Bool = false)
    # either just the diagonal or the full covariance
    if full_cov
        # compute the inverse of covariance
        if diag_var
            params.invSigmaATVec = kron(eye(params.M),
                (params.BHat'*diagm(params.sigmaVecHat)*params.BHat + params.L*mean(params.sigmaVecHat)*params.SigmaB)) + diagm(params.CA)
        else
            params.invSigmaATVec = params.sigmaHat*kron(eye(params.M),
                (params.BHat'*params.BHat + params.L*params.SigmaB)) + diagm(params.CA)
        end

        params.SigmaATVec = inv(params.invSigmaATVec)
        params.diagSigmaATVec = diag(params.SigmaATVec)

        # now the mean
        if diag_var
            params.ATVecHat = params.SigmaATVec*reshape(params.BHat'*diagm(params.sigmaVecHat)*Y, params.H*params.M)
        else
            params.ATVecHat = params.sigmaHat*params.SigmaATVec*reshape(params.BHat'*Y, params.H*params.M)
        end

        # now compute the covariance of AHat
        params.SigmaA = zeros(params.H, params.H)
        for m in 1:params.M
            params.SigmaA += view(params.SigmaATVec, (m-1)*params.H+1:m*params.H, (m-1)*params.H+1:m*params.H)
        end

    else # this updates just the diagonal of the covariance of A
        # compute the inverse of covariance
        # only use for large problems, for small problems the full covariance is more precise
        if diag_var
            for h in 1:params.H
                # instead of the full matrix multiplication as above, we compute just
                # the diagonal
                params.diagSigmaATVec[h] = norm2(params.BHat[:,h].*params.sigmaVecHat) + params.L*mean(params.sigmaVecHat)*params.SigmaB[h,h]
            end
        else         # this is computed in case the variance of data is homoscedastic
            for h in 1:params.H
                # instead of the full matrix multiplication as above, we compute just
                # the diagonal
                params.diagSigmaATVec[h] = params.sigmaHat*norm2(params.BHat[:,h]) + params.L*params.SigmaB[h,h]
            end
        end
        # we filled just the first H elements, so we just copy it to fill the rest
        params.diagSigmaATVec[params.H+1:end] = repeat(params.diagSigmaATVec[1:params.H], inner = params.M-1)
        # also add the CA vector
        params.diagSigmaATVec += params.CA

        # finally, invert it
        params.diagSigmaATVec = 1./params.diagSigmaATVec

        # now the mean
        if diag_var
            params.ATVecHat = params.diagSigmaATVec.*reshape(params.BHat'*diagm(params.sigmaVecHat)*Y, params.H*params.M)
        else
            params.ATVecHat = params.sigmaHat*params.diagSigmaATVec.*reshape(params.BHat'*Y, params.H*params.M)
        end

        # now compute the covariance of AHat
        params.SigmaA = zeros(params.H, params.H)
        for m in 1:params.M
            params.SigmaA += diagm(view(params.diagSigmaATVec, (m-1)*params.H+1:m*params.H))
        end
    end

    params.AHat = reshape(params.ATVecHat, params.H, params.M)'
    params.A1Hat = params.AHat[:,1:params.H0]
    params.A2Hat = params.AHat[1:params.M0,(params.H0+1):end]
    params.A3Hat = params.AHat[(params.M0+1):end,(params.H0+1):end]
end

"""
    updateB!(Y::Array{Float64,2}, params::vbmf_trial_parameters, diag_var::Bool = false)

Updates mean and covariance of the B matrix.
"""
function updateB!(Y::Array{Float64,2}, params::vbmf_trial_parameters;  diag_var::Bool = false)
    # first, compute the inverse of the covariance
    if diag_var
        params.SigmaB = diagm(params.CB) + mean(params.sigmaVecHat)*(params.AHat'*params.AHat + params.SigmaA)
        #invert it
        params.SigmaB = inv(params.SigmaB)
        # now compute the mean
        params.BHat = diagm(params.sigmaVecHat)*Y*params.AHat*params.SigmaB
    else
        params.SigmaB = diagm(params.CB) + params.sigmaHat*(params.AHat'*params.AHat + params.SigmaA)
        #invert it
        params.SigmaB = inv(params.SigmaB)
        params.BHat = params.sigmaHat*Y*params.AHat*params.SigmaB
    end
end

"""
    updateYHat!(params::vbmf_trial_parameters)

Updates estimate of Y.
"""
function updateYHat!(params::vbmf_trial_parameters)
    params.YHat = params.BHat*params.AHat'
end

"""
    updateCA!(params::vbmf_trial_parameters)

Updates the estimate of CA.
"""
function updateCA!(params::vbmf_trial_parameters)
    # compute the updates
    params.alpha1 = params.alpha01 + 1/2
    params.alpha2 = params.alpha02 + 1/2
    params.alpha3 = params.alpha03 + 1/2

    diagSigmaA = reshape(params.diagSigmaATVec, params.H, params.M)
    diagSigmaA1Vec = diagSigmaA[1:params.H0,:]
    diagSigmaA2Vec = diagSigmaA[(params.H0+1):end,1:params.M0]
    diagSigmaA3Vec = diagSigmaA[(params.H0+1):end,(params.M0+1):end]

    params.beta1 = params.beta01*ones(params.M*params.H0)
    params.beta1 += 1/2*vec(params.A1Hat'.*params.A1Hat' + diagSigmaA1Vec)
    params.beta2 = params.beta02*ones(params.M0*params.H1)
    params.beta2 += 1/2*vec(params.A2Hat'.*params.A2Hat')
    params.beta2 += 1/2*vec(diagSigmaA2Vec)
    params.beta3 = params.beta03*ones(params.M1*params.H1)
    params.beta3 += 1/2*vec(params.A3Hat'.*params.A3Hat' + diagSigmaA3Vec)

    params.CA1 = params.alpha1*ones(params.M*params.H0)./params.beta1
    params.CA2 = params.alpha2*ones(params.M0*params.H1)./params.beta2
    params.CA3 = params.alpha3*ones(params.M1*params.H1)./params.beta3

    # now just recombine the results back to the original arrays
    params.CA = Array{Float64, 1}()
    for m in 1:params.M0
        params.CA = cat(1, params.CA, params.CA1[((m-1)*params.H0+1):m*params.H0])
        params.CA = cat(1, params.CA, params.CA2[((m-1)*params.H1+1):m*params.H1])
    end
    for m in (params.M0+1):params.M
        params.CA = cat(1, params.CA, params.CA1[((m-1)*params.H0+1):m*params.H0])
        params.CA = cat(1, params.CA, params.CA3[((m-1-params.M0)*params.H1+1):(m-params.M0)*params.H1])
    end
    params.alpha = [params.alpha1, params.alpha2, params.alpha3]
    params.beta = Array{Float64, 1}()
    for m in 1:params.M0
        params.beta = cat(1, params.beta, params.beta1[((m-1)*params.H0+1):m*params.H0])
        params.beta = cat(1, params.beta, params.beta2[((m-1)*params.H1+1):m*params.H1])
    end
    for m in (params.M0+1):params.M
        params.beta = cat(1, params.beta, params.beta1[((m-1)*params.H0+1):m*params.H0])
        params.beta = cat(1, params.beta, params.beta3[((m-1-params.M0)*params.H1+1):(m-params.M0)*params.H1])
    end
end

"""
    updateCB!(params::vbmf_trial_parameters)

Updates the estimate of CB.
"""
function updateCB!(params::vbmf_trial_parameters)
    for h in 1:params.H
        params.delta[h] = params.delta0 + 1/2*(params.BHat[:,h]'*params.BHat[:,h])[1] + 1/2*params.SigmaB[h,h]
        params.CB[h] = params.gamma/params.delta[h]
    end
end

"""
    updateSigma!(Y::Array{Float64,2}, params::vbmf_trial_parameters, diag_var::Bool = false)

Updates estimate of the measurement variance.
"""
function updateSigma!(Y::Array{Float64,2}, params::vbmf_trial_parameters; diag_var::Bool = false)
    if diag_var
        for l in 1:params.L
            params.zetaVec[l] = params.zeta0 + 1/2*norm2(Y[l,:]) - sum(Y[l,:].*(params.AHat*params.BHat[l,:])) +
                1/2*traceXTY(params.AHat'*params.AHat + params.SigmaA,
                    params.BHat[l,:]*params.BHat[l,:]' + params.SigmaB)

            params.sigmaVecHat[l] = params.etaVec[l]/params.zetaVec[l]
        end
    else
        params.zeta = params.zeta0 + 1/2*params.trYTY - traceXTY(params.BHat, Y*params.AHat) +
            1/2*traceXTY(params.AHat'*params.AHat + params.SigmaA,
                params.BHat'*params.BHat + params.L*params.SigmaB)

        params.sigmaHat = params.eta/params.zeta
    end
end

"""
    updateAlpha01!(params::vbmf_trial_parameters)

Estimate alpha01 by maximizing the lower bound.
"""
function updateAlpha01!(params::vbmf_trial_parameters)
    function fAlpha01(x)
        return params.M*params.H0*log(params.beta01) - params.M*params.H0*digamma(x) + sum(map(gammaELn, params.alpha1*ones(size(params.beta1)), params.beta1))
    end
    try
        root = fzero(fAlpha01, 1e-10, 1e10, ftol = 1e-5)
        params.alpha01 = root
    end
end

"""
    updateBeta01!(params::vbmf_trial_parameters)

Estimate beta01 by maximizing the lower bound.
"""
function updateBeta01!(params::vbmf_trial_parameters)
    params.beta01 = params.M*params.H0*params.alpha01/sum(params.CA1)
end

"""
    updateAlpha02!(params::vbmf_trial_parameters)

Estimate alpha02 by maximizing the lower bound.
"""
function updateAlpha02!(params::vbmf_trial_parameters)
    function fAlpha02(x)
        return params.M0*params.H1*log(params.beta02) - params.M0*params.H1*digamma(x) + sum(map(gammaELn, params.alpha2*ones(size(params.beta2)), params.beta2))
    end
    try
        root = fzero(fAlpha02, 1e-10, 1e10, ftol = 1e-5)
        params.alpha02 = root
    end
end

"""
    updateBeta02!(params::vbmf_trial_parameters)

Estimate beta02 by maximizing the lower bound.
"""
function updateBeta02!(params::vbmf_trial_parameters)
    params.beta02 = params.M0*params.H1*params.alpha02/sum(params.CA2)
end

"""
    updateAlpha03!(params::vbmf_trial_parameters)

Estimate alpha03 by maximizing the lower bound.
"""
function updateAlpha03!(params::vbmf_trial_parameters)
    function fAlpha03(x)
        return params.M1*params.H1*log(params.beta03) - params.M1*params.H1*digamma(x) + sum(map(gammaELn, params.alpha3*ones(size(params.beta3)), params.beta3))
    end
    try
        root = fzero(fAlpha03, 1e-10, 1e10, ftol = 1e-5)
        params.alpha03 = root
    end
end

"""
    updateBeta03!(params::vbmf_trial_parameters)

Estimate beta03 by maximizing the lower bound.
"""
function updateBeta03!(params::vbmf_trial_parameters)
    params.beta03 = params.M1*params.H1*params.alpha03/sum(params.CA3)
end

"""
    vbmf_trial!(Y::Array{Float64, 2}, params::vbmf_trial_parameters, niter::Int, eps::Float64 = 1e-6,
    est_var = false, full_cov::Bool = false, logdir = "", desc = "")

Computes variational bayes matrix factorization of Y = AB' + E. Independence of A and B is assumed.
Estimate of variance sigma can be either a scalar or a vector - for row variance estimation.
ARD property is imposed upon columns of B and also
upon vec(A) through gamma prior CB and CA and estimation of covariance.
The prior model is following:

    p(Y|A,B) = N(Y|BA^T, 1/sigma*I) or p(Y|A,B) = N(Y|BA^T, inv(diag(sigma)))
    p(vec(A^T)) = N(vec(A^T)|0, invCA), CA = diag(c_a)
    p(B_h) = MN(B|0, I, invCB), CB = diag(c1, ..., cH)
    p(CA) = G(C_A| alpha0, beta0)
    p(CB) = G(C_B| gamma0, delta0)

The params argument with initialized data is modified and contains the resulting estimates after the
algorithm stops.
"""
function vbmf_trial!(Y::Array{Float64, 2}, params::vbmf_trial_parameters, niter::Int; eps::Float64 = 1e-6,
    diag_var::Bool = false, full_cov::Bool = false, logdir = "", desc = "", verb = false, est_priors = true,
    est_cb::Bool = true)
    priors = Dict()

    # create the log dictionary
    loging = false
    if logdir !=""
        loging = true
        logVar = create_log(params)
    end

    # choice of convergence control variable
    convergence_var = :BHat
    #convergence_var = :dY

    if convergence_var == :dY
        old = Y - params.YHat
    else
        old = copy(getfield(params, convergence_var))
    end
    d = eps + 1.0 # delta
    i = 1

    # run the loop for a given number of iterations
    while ((i <= niter) && (d > eps))
        updateA!(Y, params, full_cov = full_cov, diag_var = diag_var)
        updateB!(Y, params, diag_var = diag_var)

        updateCA!(params)
        if est_cb
            updateCB!(params)
        end

        updateSigma!(Y, params, diag_var = diag_var)

        # estimate priors
        if est_priors
            updateAlpha01!(params)
            updateAlpha02!(params)
            updateAlpha03!(params)
            updateBeta01!(params)
            updateBeta02!(params)
            updateBeta03!(params)
        end
        if loging
            update_log!(logVar, params)
        end

        # check convergence
        if convergence_var == :dY
            updateYHat!(params)
            d = delta(Y - params.YHat, old)
            old = Y - params.YHat
        else
            d = delta(getfield(params, convergence_var), old)
            old = copy(getfield(params, convergence_var))
        end
        i += 1
    end

    # finally, compute the estimate of Y
    updateYHat!(params)

    # convergence info
    if verb
        print("Factorization finished after ", i-1, " iterations, eps = ", d, "\n")
    end

    # save inputs and outputs
    if loging
        println("Saving outputs and inputs under ", logdir, "/")
        save_log(logVar, Y, priors, logdir, desc = desc)
    end

    return d
end

"""
    vbmf_trial(Y::Array{Float64, 2}, params_in::vbmf_trial_parameters, niter::Int, eps::Float64 = 1e-6,
    est_var = false, full_cov::Bool = false, logdir = "", desc = "")

Calls vbmf_trial!() but copies the params_in argument so that it is not modified and can be reused.
"""
function vbmf_trial(Y::Array{Float64, 2}, params_in::vbmf_trial_parameters, niter::Int; eps::Float64 = 1e-6,
    diag_var::Bool = false, full_cov::Bool = false, logdir = "", desc = "", verb = false, est_priors = true,
    est_cb::Bool = true)
    # make a copy of input params
    params = copy(params_in)

    # run the algorithm
    d = vbmf_trial!(Y, params, niter, eps = eps, diag_var = diag_var, full_cov = full_cov,
        logdir = logdir, desc = desc, verb = verb, est_priors = est_priors, est_cb = est_cb)

    return params, d
end

"""
    lowerBound(Y::Array{Float64,2}, params::vbmf_trial_parameters)

Compute the lower bound for logarithm of data distribution.
"""
function lowerBound(Y::Array{Float64,2}, params::vbmf_trial_parameters)
    L = 0.0
    # E[lnp(Y|params)]
    L += - params.L*params.M/2*ln2pi + params.L*params.M/2*gammaELn(params.eta, params.zeta)
    L += - params.sigmaHat/2*(params.trYTY - 2*traceXTY(params.BHat, Y*params.AHat)
         + traceXTY(params.AHat'*params.AHat + params.SigmaA, params.BHat'*params.BHat + params.L*params.SigmaB))
    # E[lnp(vec(A'))]
    L += - params.MH/2*ln2pi + 1/2*sum(map(gammaELn, params.alpha1*ones(size(params.beta1)), params.beta1))
    L += 1/2*sum(map(gammaELn, params.alpha2*ones(size(params.beta2)), params.beta2))
    L += 1/2*sum(map(gammaELn, params.alpha3*ones(size(params.beta3)), params.beta3))
    L += - (1/2*params.CA'*(params.ATVecHat.^2 + params.diagSigmaATVec))[1]
    # E[lnp(B)]
    L += - params.L*params.H/2*ln2pi
    L += params.L/2*sum(map(gammaELn, params.gamma*ones(size(params.delta)), params.delta))
    L += - 1/2*traceXTY(diagm(params.CB), (params.BHat'*params.BHat + params.L*params.SigmaB))
    # E[lnp(sigma)]
    L += params.eta0*log(params.zeta0) - lgamma(params.eta0)
    L += (params.eta0 - 1)*gammaELn(params.eta, params.zeta) - params.zeta0*params.sigmaHat
    # E[lnp(CA1)]
    L += params.M*params.H0*(params.alpha01*log(params.beta01) - lgamma(params.alpha01))
    L += (params.alpha01 - 1)*sum(map(gammaELn, params.alpha1*ones(size(params.beta1)), params.beta1))
    L += - params.beta01*sum(params.CA1)
    # E[lnp(CA2)]
    L += params.M0*params.H1*(params.alpha02*log(params.beta02) - lgamma(params.alpha02))
    L += (params.alpha02 - 1)*sum(map(gammaELn, params.alpha2*ones(size(params.beta2)), params.beta2))
    L += - params.beta02*sum(params.CA2)
    # E[lnp(CA3)]
    L += params.M1*params.H1*(params.alpha03*log(params.beta03) - lgamma(params.alpha03))
    L += (params.alpha03 - 1)*sum(map(gammaELn, params.alpha3*ones(size(params.beta3)), params.beta3))
    L += - params.beta03*sum(params.CA3)
    # E[lnp(CB)]
    L += params.H*(params.gamma0*log(params.delta0) - lgamma(params.gamma0))
    L += (params.gamma0 - 1)*sum(map(gammaELn, params.gamma*ones(size(params.delta)), params.delta))
    L += - params.gamma0*sum(params.CB)

    # H(vec(A'))
    L += normalEntropy(params.diagSigmaATVec)
    # H(B)
    L += normalEntropy(kron(params.SigmaB,eye(params.L)))
    # H(sigma)
    L += gammaEntropy(params.eta, params.zeta)
    # H(CA1)
    L += sum(map(gammaEntropy, params.alpha1*ones(size(params.beta1)), params.beta1))
    # H(CA2)
    L += sum(map(gammaEntropy, params.alpha2*ones(size(params.beta2)), params.beta2))
    # H(CA3)
    L += sum(map(gammaEntropy, params.alpha3*ones(size(params.beta3)), params.beta3))
    # H(CB)
    L += sum(map(gammaEntropy, params.gamma*ones(size(params.delta)), params.delta))
    return L
end

"""
    lowerBoundTrimmed(Y::Array{Float64,2}, params_in::vbmf_trial_parameters, trim = 1e-1)

AHat with abs(AHat[i,j]) < trim and the corresponding elements of CA are not involved in the computation.
"""
function lowerBoundTrimmed(Y::Array{Float64,2}, params_in::vbmf_trial_parameters, trim = 1e-1)
    params = copy(params_in)

    trim_inds = (abs(params.ATVecHat) .> trim)
    params.ATVecHat = params.ATVecHat[trim_inds]
    params.MH = size(params.ATVecHat)[1]
    params.beta = params.beta[trim_inds]
    params.CA = params.CA[trim_inds]
    params.diagSigmaATVec = params.diagSigmaATVec[trim_inds]

    return lowerBound(Y, params)
end


text:
using BinDeps

@BinDeps.setup
netcdf = library_dependency("libnetcdf", aliases = ["libnetcdf4","libnetcdf-7"])

using Conda
provides(Conda.Manager, "libnetcdf", netcdf)

@osx_only begin
    using Homebrew
		provides(Homebrew.HB, "netcdf", netcdf, os = :Darwin )
end

provides(AptGet, Dict("libnetcdf-dev"=>netcdf), os = :Linux)
provides(Yum, Dict("netcdf-devel"=>netcdf), os = :Linux)

@BinDeps.install


text:
module Final

import Evolve

L = 2                       # width of the box
xmin = -1                   # leftmost point in the box
N = 20000                   # number of subintervals
ϵ = L/N                     # width of each subinterval
k = 20                      # wavenumbers
σ = 0.2                     # wavepacket width
x0 = 0.0                    # wavepacket starting position

V(xs) = map(x->0.0, xs)     # zero potential energy
AV = Evolve.A(V)            # evolution matrix constructor for square well
ϕV = Evolve.Φ(k, σ, x0)     # initial state vector constructor for Gaussian packet

x = (L/N .* [0:N] + xmin)   # x positions

end