text:
function chol!(A)
    C, info = LinearAlgebra._chol!(A, LowerTriangular)
    Cholesky(C.data, 'L', info)
end

function foo(C,A,B)
    matmul!(C,A,B,1.0, 2.0)
    C .+= 1
    C
end

"""
    backwardpass!(::iLQRSolver)

Compute the backwardpass of iLQR, which forms a local quadratic approximation of the
cost-to-go and constructs a local affine control policy about the current trajectory
using time-varying LQR. The expansions of the objective function and dynamics
should already be computed before calling this function.
"""
function backwardpass!(solver::iLQRSolver)
    # Extract Variables
    N = solver.N
    D = solver.D     # dynamics expansion
    E = solver.Eerr  # cost expansion
    S = solver.S     # quadratic cost-to-go
    Q = solver.Q     # action-value expansion
    K = solver.K
    d = solver.d
    Quu_reg = solver.Quu_reg
    Qux_reg = solver.Qux_reg
    ΔV = solver.ΔV
    Qtmp = solver.Qtmp

    # Terminal cost-to-go
    S[N].xx .= E[N].xx
    S[N].x .= E[N].x

    k = N-1
    ΔV .= 0
    while k > 0
        A = D[k].fx
        B = D[k].fu
        m = RD.control_dim(solver.model[k])

        # Action-value expansion
        matmul!(Q[k].x, A', S[k+1].x)
        Q[k].x .+= E[k].x

        matmul!(Q[k].u, B', S[k+1].x)
        Q[k].u .+= E[k].u

        matmul!(Qtmp.xx, S[k+1].xx, A)
        matmul!(Q[k].xx, A', Qtmp.xx)
        Q[k].xx .+= E[k].xx

        matmul!(Qtmp.ux, B', S[k+1].xx)
        matmul!(Q[k].uu, Qtmp.ux, B)
        Q[k].uu .+= E[k].uu

        matmul!(Qtmp.xx, S[k+1].xx, A)
        matmul!(Q[k].ux, B', Qtmp.xx)
        Q[k].ux .+= E[k].ux

        # Regularization
        ρ = solver.reg.ρ
        if solver.opts.bp_reg_type == :state
            Quu_reg .= Q[k].uu
            matmul!(Quu_reg, B', B, ρ, 1.0)
            Qux_reg .= Qux
            matmul!(Qux_reg, A', A, ρ, 1.0)
        elseif solver.opts.bp_reg_type == :control
            Quu_reg .= Q[k].uu
            for i = 1:m
                Quu_reg[i,i] += ρ
            end
            Qux_reg .= Q[k].ux
        end

        # Solve for gains
        K[k] .= Qux_reg
        d[k] .= Q[k].u
        # LAPACK.potrf!('L',Quu_reg)
        Quu_fact = chol!(Quu_reg)::Cholesky
        if !isposdef(Quu_fact)  # this is a super cheap check
            # TODO: add log message
            @warn "Backwardpass cholesky failed at time step $k"
            increaseregularization!(solver)
            k = N-1
            ΔV .= 0
            continue
        end
        # Save time by solving for K and d at the same time (1 BLAS call)
        ldiv!(Quu_fact, solver.gains[k])
        # LAPACK.potrs!('L', Quu_reg, K[k])
        # LAPACK.potrs!('L', Quu_reg, d[k])
        # K[k] .*= -1
        # d[k] .*= -1
        solver.gains[k] .*= -1

        # k == N-5 && break

        # Update Cost-to-go
	    # S.x[k]  =  Qx + K[k]'*Quu*d[k] + K[k]'* Qu + Qux'd[k]
        S[k].x .= Q[k].x
        matmul!(Qtmp.u, Q[k].uu, d[k])
        matmul!(S[k].x, K[k]', Qtmp.u, 1.0, 1.0)
        matmul!(S[k].x, K[k]', Q[k].u, 1.0, 1.0)
        matmul!(S[k].x, Q[k].ux', d[k], 1.0, 1.0)

	    # S.xx[k] = Qxx + K[k]'*Quu*K[k] + K[k]'*Qux + Qux'K[k]
        S[k].xx .= Q[k].xx
        matmul!(Qtmp.ux, Q[k].uu, K[k])
        matmul!(S[k].xx, K[k]', Qtmp.ux, 1.0, 1.0)
        matmul!(S[k].xx, K[k]', Q[k].ux, 1.0, 1.0)
        matmul!(S[k].xx, Q[k].ux', K[k], 1.0, 1.0)

        # Symmeterize Cost-to-go Hessian
        transpose!(Qtmp.xx, S[k].xx)
        S[k].xx .+= Qtmp.xx
        S[k].xx ./= 2

        ΔV[1] += dot(d[k], Q[k].u)
        ΔV[2] += 0.5 * dot(d[k], Q[k].uu, d[k])
        # println("k = $k, ΔV = ", ΔV)

        k -= 1
    end
    decreaseregularization!(solver)
    return ΔV
end


text:
module EmacsVterm

import REPL
import Base64
import Markdown

struct Display <: AbstractDisplay
    io::IO
end
Display() = Display(stdout)

const user = get(ENV, "USER", "")
vterm_cmd(str) = "\e]$str\e\\"
prompt_suffix() = vterm_cmd("51;A$(user)@$(gethostname()):$(pwd())")
eval_elisp(elisp) = vterm_cmd("51;E$(elisp)")

# Show rendered Markdown in Emacs *julia-doc* buffer.
function Base.display(d::Display, md::Markdown.MD)
    write(d.io, eval_elisp("julia-repl--show documentation text/html \"$(md |> Markdown.html |> Base64.base64encode)\""))
end

function __init__()
    if !(isinteractive() && isdefined(Base, :active_repl))
        return
    end
    begin
        if get(ENV, "INSIDE_EMACS", "") == "vterm"
            @info "Emacs vterm detected"
            repl = Base.active_repl

            if !isdefined(repl,:interface)
                repl.interface = REPL.setup_interface(repl)
            end

            suffix = repl.interface.modes[1].prompt_suffix
            repl.interface.modes[1].prompt_suffix = function ()
                ((isa(suffix,Function) ? suffix() : suffix) * prompt_suffix())
            end

            pushdisplay(Display())
        end
    end
end

end

ext: jl


text:
function setupsim()
    ρ = 0.78
    kT = 0.85
    particles = 500
    cutoff = 2.5
    s = System(ρ, kT, particles, cutoff; lcell=1)
    return Simulation(s, NVT(0.25, 0.4), LennardJones())
end

@testset "Parallel Simulations" begin
    sim = setupsim()
    @test_nowarn simulate!(sim; steps=10_000, parallel=true)
end

@testset "Serial Simulations" begin
    sim = setupsim()
    @test_nowarn simulate!(sim; steps=10_000, parallel=false)
end

@testset "N-square" begin
    sim = setupsim()
    @test_nowarn simulate!(sim; pairwise=:squared, steps=10_000, parallel=false)
end


text:
# RELOG: Reverse Logistics Optimization
# Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
# Released under the modified BSD license. See COPYING.md for more details.

using DataStructures
using JSON
using JSONSchema
using Printf
using Statistics

function parsefile(path::String)::Instance
    return RELOG.parse(JSON.parsefile(path))
end

function parse(json)::Instance
    basedir = dirname(@__FILE__)
    json_schema = JSON.parsefile("$basedir/../schemas/input.json")
    validate(json, Schema(json_schema))

    T = json["parameters"]["time horizon (years)"]
    json_schema["definitions"]["TimeSeries"]["minItems"] = T
    json_schema["definitions"]["TimeSeries"]["maxItems"] = T
    validate(json, Schema(json_schema))

    building_period = [1]
    if "building period (years)" in keys(json)
        building_period = json["building period (years)"]
    end

    plants = Plant[]
    products = Product[]
    collection_centers = CollectionCenter[]
    prod_name_to_product = Dict{String,Product}()

    # Create products
    for (product_name, product_dict) in json["products"]
        cost = product_dict["transportation cost (\$/km/tonne)"]
        energy = zeros(T)
        emissions = Dict()

        if "transportation energy (J/km/tonne)" in keys(product_dict)
            energy = product_dict["transportation energy (J/km/tonne)"]
        end

        if "transportation emissions (tonne/km/tonne)" in keys(product_dict)
            emissions = product_dict["transportation emissions (tonne/km/tonne)"]
        end

        product = Product(product_name, cost, energy, emissions)
        push!(products, product)
        prod_name_to_product[product_name] = product

        # Create collection centers
        if "initial amounts" in keys(product_dict)
            for (center_name, center_dict) in product_dict["initial amounts"]
                if "location" in keys(center_dict)
                    region = geodb_query(center_dict["location"])
                    center_dict["latitude (deg)"] = region.centroid.lat
                    center_dict["longitude (deg)"] = region.centroid.lon
                end
                center = CollectionCenter(
                    length(collection_centers) + 1,
                    center_name,
                    center_dict["latitude (deg)"],
                    center_dict["longitude (deg)"],
                    product,
                    center_dict["amount (tonne)"],
                )
                push!(collection_centers, center)
            end
        end
    end

    # Create plants
    for (plant_name, plant_dict) in json["plants"]
        input = prod_name_to_product[plant_dict["input"]]
        output = Dict()

        # Plant outputs
        if "outputs (tonne/tonne)" in keys(plant_dict)
            output = Dict(
                prod_name_to_product[key] => value for
                (key, value) in plant_dict["outputs (tonne/tonne)"] if value > 0
            )
        end

        energy = zeros(T)
        emissions = Dict()

        if "energy (GJ/tonne)" in keys(plant_dict)
            energy = plant_dict["energy (GJ/tonne)"]
        end

        if "emissions (tonne/tonne)" in keys(plant_dict)
            emissions = plant_dict["emissions (tonne/tonne)"]
        end

        for (location_name, location_dict) in plant_dict["locations"]
            sizes = PlantSize[]
            disposal_limit = Dict(p => [0.0 for t = 1:T] for p in keys(output))
            disposal_cost = Dict(p => [0.0 for t = 1:T] for p in keys(output))

            # GeoDB
            if "location" in keys(location_dict)
                region = geodb_query(location_dict["location"])
                location_dict["latitude (deg)"] = region.centroid.lat
                location_dict["longitude (deg)"] = region.centroid.lon
            end

            # Disposal
            if "disposal" in keys(location_dict)
                for (product_name, disposal_dict) in location_dict["disposal"]
                    limit = [1e8 for t = 1:T]
                    if "limit (tonne)" in keys(disposal_dict)
                        limit = disposal_dict["limit (tonne)"]
                    end
                    disposal_limit[prod_name_to_product[product_name]] = limit
                    disposal_cost[prod_name_to_product[product_name]] =
                        disposal_dict["cost (\$/tonne)"]
                end
            end

            # Capacities
            for (capacity_name, capacity_dict) in location_dict["capacities (tonne)"]
                push!(
                    sizes,
                    PlantSize(
                        Base.parse(Float64, capacity_name),
                        capacity_dict["variable operating cost (\$/tonne)"],
                        capacity_dict["fixed operating cost (\$)"],
                        capacity_dict["opening cost (\$)"],
                    ),
                )
            end
            length(sizes) > 1 || push!(sizes, sizes[1])
            sort!(sizes, by = x -> x.capacity)

            # Storage
            storage_limit = 0
            storage_cost = zeros(T)
            if "storage" in keys(location_dict)
                storage_dict = location_dict["storage"]
                storage_limit = storage_dict["limit (tonne)"]
                storage_cost = storage_dict["cost (\$/tonne)"]
            end

            # Validation: Capacities
            if length(sizes) != 2
                throw("At most two capacities are supported")
            end
            if sizes[1].variable_operating_cost != sizes[2].variable_operating_cost
                throw("Variable operating costs must be the same for all capacities")
            end

            plant = Plant(
                length(plants) + 1,
                plant_name,
                location_name,
                input,
                output,
                location_dict["latitude (deg)"],
                location_dict["longitude (deg)"],
                disposal_limit,
                disposal_cost,
                sizes,
                energy,
                emissions,
                storage_limit,
                storage_cost,
            )

            push!(plants, plant)
        end
    end

    @info @sprintf("%12d collection centers", length(collection_centers))
    @info @sprintf("%12d candidate plant locations", length(plants))

    return Instance(T, products, collection_centers, plants, building_period)
end

max_stars_repo_name: ANL-CEEESA/RELOG

max_issues_repo_name: ANL-CEEESA/RELOG

max_forks_repo_name: ANL-CEEESA/RELOG


text:
entropy(C::Array{T, 2} where T<:Real) = 0.5*logdet(2*π*ℯ*C*C')

entropy(C::AbstractArray{T, 1} where T<:Real) = 0.5*sum(log.(C.^2)) + 0.5*log(2*π*ℯ) * length(C)

entropy(σ::Real) = 0.5*log(2*π*ℯ*σ^2)


text:
# diffs.jl

using MIRT: diff_map, diff_forw, diff_adj
import MIRT: diff2d_forw, diff2d_adj, diff_length, diff2d_map # deprecated

using LinearMapsAA: LinearMapAA
using Test: @test, @test_throws

@test_throws ArgumentError diff_forw(ones(3), dims=2)
@test_throws ArgumentError diff_forw(ones(3), dims=[1,2])
@test_throws ArgumentError diff_forw(ones(1,2,1), dims=(1,2))
@test diff_forw(ones(2,4,6) ; dims=2) == zeros(2*3*6)

for N in [(3,), (3,4), (2,3,4), (4,3,2,2)]
	T = diff_map(N)
	@test Matrix(T)' == Matrix(T')
	@test T.name == "diff_map"
	T = diff_map(N, dims=[1])
	@test Matrix(T)' == Matrix(T')
	if length(N) >= 2
		for dims in [2, (2,), [1,2], (1,2)]
			T = diff_map(N, dims=dims)
			@test Matrix(T)' == Matrix(T')
		end
	end
	if length(N) >= 3
		for dims in [3, (1,3), [2,3], (1,2,3)]
			T = diff_map(N, dims=dims)
			@test Matrix(T)' == Matrix(T')
		end
	end
end

if true # test old 2D versions
	N = (3,5)
	x = rand(N...)
	@test diff2d_forw(x) == diff_forw(x)
	d = rand(sum(diff_length(N,dim) for dim=1:2))
	@test diff2d_adj(d, N... ; out2d=false) == vec(diff_adj(d, N))
	@test diff2d_adj(d, N... ; out2d=true) == diff_adj(d, N)
	@test Matrix(diff2d_map(N...)) == Matrix(diff_map(N))
end

N = (1,2)
@test_throws ArgumentError T = diff_map(N)
@test_throws ArgumentError T = diff_map(N, dims=1)
T = diff_map(N, dims=2)
@test Matrix(T)' == Matrix(T')


text:
module InteractiveErrors

# TODO: de-vendor once changes are upstreamed.
include("vendor/FoldingTrees/src/FoldingTrees.jl")
using .FoldingTrees

using REPL, REPL.TerminalMenus, InteractiveUtils, IterTools, Requires

export toggle, current_theme, set_theme!, reset_theme!, adjust_theme!

#
# Themes.
#

const DEFAULT_THEME = (
    function_name   = (bold = true,),
    directory       = (color = :light_black,),
    filename        = (color = :magenta, bold = true,),
    line_number     = (color = :green, bold = true),
    user_stack      = (color = :green, bold = true),
    system_stack    = (color = :red, bold = true),
    stdlib_module   = (color = :yellow,),
    base_module     = (color = :blue,),
    core_module     = (color = :light_black,),
    package_module  = (color = :cyan, bold = true,),
    unknown_module  = (color = :red,),
    inlined_frames  = (color = :light_black,),
    toplevel_frames = (color = :light_black,),
    repeated_frames = (color = :red,),
    file_contents   = (color = :light_black,),
    signature       = (color = :light_black, format = true, highlight = true),
    source          = (color = :normal, bold = true, highlight = true),
    line_range      = (before = 0, after = 5,),
    charset         = :unicode,
)
const THEME = Ref{Any}(DEFAULT_THEME)

current_theme() = THEME[]
set_theme!(nt::NamedTuple) = THEME[] = nt
set_theme!(; kws...) = set_theme!(_nt(kws))

_nt(kws) = NamedTuple{Tuple(keys(kws))}(values(kws))
reset_theme!() = set_theme!(DEFAULT_THEME)

adjust_theme!(nt::NamedTuple) = set_theme!(merge(current_theme(), nt))
adjust_theme!(; kws...) = adjust_theme!(_nt(kws))

get_theme(key) = get(NamedTuple, current_theme(), key)
get_theme(key, default) = get(current_theme(), key, default)

function style(str; kws...)
    sprint(; context = :color => true) do io
        printstyled(
            io, str;
            bold = get(kws, :bold, false),
            color = get(kws, :color, :normal),
        )
    end
end
style(str, key::Symbol) = style(str; get_theme(key)...)

#
# Stackframe Wrapping.
#

struct StackFrameWrapper
    sf::StackTraces.StackFrame
    n::Int
    StackFrameWrapper(tuple) = new(tuple...)
end

function Base.show(io::IO, s::StackFrameWrapper)
    func = style(s.sf.func, :function_name)
    file = rewrite_path(s.sf.file)
    dir, file = dirname(file), basename(file)
    file = style(file, :filename)
    dir = style(joinpath(dir, ""), :directory)
    line = style(s.sf.line, :line_number)
    repeated = s.n > 1 ? style("x $(s.n)", :repeated_frames) : ""
    print(io, strip("$func $dir$file:$line $repeated"))
end

function rewrite_path(path)
    fn(path, replacer) = replace(String(path), replacer; count = 1)
    path = fn(path, normpath(Sys.BUILD_STDLIB_PATH) => "@stdlib")
    path = fn(path, normpath(Sys.STDLIB) => "@stdlib")
    path = fn(path, homedir() => "~")
    return path
end

function find_source(file)
    # Binary versions of Julia have the wrong stdlib path, fix it.
    file = replace(string(file), normpath(Sys.BUILD_STDLIB_PATH) => Sys.STDLIB; count = 1)
    return Base.find_source_file(file)
end

#
# Explorer.
#

struct CapturedError
    err
    bt
end

Base.show(io::IO, ce::CapturedError) = showerror(io, ce.err, ce.bt)

explore(err::CapturedError) = explore(stdout, err)

function explore(io::IO, err::CapturedError; interactive = true)
    # Give a printout of the actual error message prior to launching tree
    # explorer since it's probably useful to have.
    println(io, sprint(showerror, err.err, context = :color => true))

    # Use the default cleaning functionality from Base. No need to reinvent.
    clean = Base.process_backtrace(err.bt)
    wrapped = StackFrameWrapper.(clean)

    toplevel = findfirst(s -> StackTraces.is_top_level_frame(s.sf), wrapped)
    toplevel = toplevel === nothing ? length(wrapped) : toplevel
    user_frames = wrapped[1:toplevel]
    system_frames = wrapped[toplevel+1:end]

    root = Node{Any}("(stacktrace)")

    function make_nodes(root_node, frames; fold = false)
        for (nth, frame_group) in enumerate(aggregate_modules(frames))
            m = module_of(first(frame_group))
            if m === :unknown
                for frame in frame_group
                    fold!(Node{Any}(frame, root_node))
                end
            else
                name =
                    m === :inlined    ? style("[inlined]", :inlined_frames) :
                    m === :toplevel   ? style("[top-level]", :toplevel_frames) :
                    is_from_stdlib(m) ? style("$(m)", :stdlib_module) :
                    is_from_base(m)   ? style("$(m)", :base_module) :
                    is_from_core(m)   ? style("$(m)", :core_module) :
                    is_from_package(m) ? style("$(m)", :package_module) :
                        style("$(m)", :unknown_module)

                node = Node{Any}(name, root_node)
                for frame in frame_group
                    current = Node{Any}(frame, node)
                    fold!(current)
                    # Formatted signature for the frame:
                    if !StackTraces.is_top_level_frame(frame.sf)
                        let lines = _formatted_signature(frame)
                            if !isempty(lines)
                                sig = Node{Any}(style("signature", :signature), current)
                                fold!(sig)
                                for line in lines
                                    Node{Any}(line, sig)
                                end
                            end
                        end
                    end
                    # Source code for the frame:
                    let lines = _lines_around(frame)
                        if !isempty(lines)
                            src = Node{Any}(style("source", :source), current)
                            for line in lines
                                Node{Any}(line, src)
                            end
                        end
                    end
                end
                # Hide any of the following by default:
                if m in (:inlined, :toplevel) || is_from_stdlib(m) || is_from_base(m) || is_from_core(m) || fold
                    fold!(node)
                end
                # Always open up the very first node, unless it's a toplevel.
                if nth === 1 && m !== :toplevel
                    unfold!(node)
                end
            end
        end
    end

    user_nodes = Node{Any}(style("(user)", :user_stack), root)
    make_nodes(user_nodes, user_frames)

    system_nodes = Node{Any}(style("(system)", :system_stack), root)
    make_nodes(system_nodes, system_frames; fold = true)
    fold!(system_nodes)

    menu = TreeMenu(root; dynamic = true, maxsize = 30)
    result = interactive ? TerminalMenus.request(menu; cursor = 3) : user_nodes
    result === nothing && return

    actions = [
        "clipboard"  => () -> (maybe_clipboard(sprint(showerror, err.err, err.bt[1:toplevel])); nothing),
        "print"      => () -> (showerror(io, err.err, err.bt[1:toplevel]); nothing),
        "stacktrace" => () -> clean,
        "exception"  => () -> err.err,
        "backtrace"  => () -> err.bt,
    ]

    data = result.data
    extras = []
    if isa(data, StackFrameWrapper)
        file, line = data.sf.file, data.sf.line
        file = find_source(file)
        if file !== nothing && isfile(file)
            file, line
            extras = [
                "edit" => () -> (edit(file, line); nothing),
                "retry" => () -> true,
            ]
            has_debugger() && push!(extras, "breakpoint" => () -> breakpoint(file, line))
            push!(extras, "less" => () -> (less(file, line); nothing))
            actions = vcat(extras, actions)
        end
        if isdefined(data.sf, :linfo)
            mi = data.sf.linfo
            if isa(mi, Core.MethodInstance)
                extras = []
                if has_cthulhu()
                    push!(extras, "ascend" => () -> ascend(mi))
                    push!(extras, "descend" => () -> descend(mi))
                end
                if has_jet()
                    push!(extras, "JET" => () -> report_call(mi))
                end
                actions = vcat(extras, actions)
            end
        end
    end

    result = interactive ? request(MultiSelectMenu(first.(actions); charset = get_theme(:charset, :unicode))) : collect(1:length(actions))
    choice = sort(collect(result))
    if !isempty(choice)
        output = []
        for (name, func) in actions[choice]
            out = func()
            out === nothing || push!(output, Symbol(name) => out)
        end
        isempty(output) || return NamedTuple{Tuple(first.(output))}(last.(output))
    end
    return nothing
end

function _lines_around(s::StackFrameWrapper)
    file, line = s.sf.file, s.sf.line
    file = find_source(file)
    if file !== nothing && isfile(file)
        lines = readlines(file)
        range = get_theme(:line_range)
        above = max(1, line - get(range, :before, 0))
        below = min(line + get(range, :after, 5), length(lines))
        highlighter = get(get_theme(:source), :highlight, true) === true ? highlight : s -> style(s, :file_contents)
        return highlighter.(lines[above:below])
    else
        return String[]
    end
end

function _formatted_signature(s::StackFrameWrapper)
    str = String(rsplit(string(s.sf), " at "; limit = 2)[1])
    str = replace(str, "#unused#" => "")
    formatter = get(get_theme(:signature), :format, true) === true ? format_julia_source : identity
    highlighter = get(get_theme(:signature), :highlight, true) === true ? highlight : s -> style(s, :file_contents)
    fmt = highlighter(formatter(str))
    return collect(eachline(IOBuffer(fmt)))
end

# Just give up when there is no clipboard available.
function maybe_clipboard(str)
    try
        clipboard(str)
    catch err
        @warn "Could not find a clipboard."
    end
end

rootmodule(m::Module) = m === Base ? m : m === parentmodule(m) ? m : rootmodule(parentmodule(m))
rootmodule(::Any) = nothing
modulepath(m::Module) = string(pkgdir(m))
modulepath(other) = ""

is_from_stdlib(m) = startswith(modulepath(rootmodule(m)), Sys.STDLIB)
is_from_base(m) = rootmodule(m) === Base
is_from_core(m) = rootmodule(m) === Core
is_from_package(m) = (r = rootmodule(m); !is_from_core(r) && !is_from_base(r) && !is_from_stdlib(r))

module_of(sf) =
    sf.sf.inlined ? :inlined :
    sf.sf.func === Symbol("top-level scope") ? :toplevel :
    isa(sf.sf.linfo, Core.MethodInstance) ? sf.sf.linfo.def.module :
        :unknown

aggregate_modules(stacktrace) = IterTools.groupby(module_of, stacktrace)

#
# REPL hook.
#

const ENABLED = Ref(true)

"""
Turn interactive errors on or off.
"""
toggle() = ENABLED[] = !ENABLED[]

is_toggle_expr(expr) = Meta.isexpr(expr, :call, 1) && expr.args[1] === :toggle

is_retry(::Nothing) = false
is_retry(nt::NamedTuple) = haskey(nt, :retry) && nt.retry === true

maybe_retry(out, expr) = is_retry(out) ? Core.eval(Main, _ast_transforms(expr)) : out

function _ast_transforms(ast)
    if isdefined(Base, :active_repl_backend)
        for xf in Base.active_repl_backend.ast_transforms
            ast = Base.invokelatest(xf, ast)
        end
    end
    return ast
end

function wrap_errors(expr)
    if ENABLED[] && !is_toggle_expr(expr)
        quote
            try
                $(Expr(:toplevel, expr))
            catch e
                $(maybe_retry)($(explore)(($CapturedError)(e, catch_backtrace())), $(Expr(:quote, expr)))
            end
        end
    else
        expr
    end
end

function setup_repl()
    @async begin
        done = false
        for _ in 1:10
            if isdefined(Base, :active_repl_backend)
                backend = Base.active_repl_backend
                if isdefined(backend, :ast_transforms)
                    pushfirst!(backend.ast_transforms, wrap_errors)
                    done = true
                    break
                end
            end
            sleep(0.5)
        end
        done || @warn "Could not start `InteractiveErrors` REPL hook."
    end
end

#
# Requires.
#

has_cthulhu(args...) = false
ascend(args...) = @warn "`import Cthulhu` to enable `ascend` action."
descend(args...) = @warn "`import Cthulhu` to enable `descend` action."

has_debugger(args...) = false
breakpoint(args...) = @warn "`import Debugger` to enable `breakpoint` action."

has_jet(args...) = false
report_call(args...) = @warn "`import JET` to enable `report_call` action."

has_juliaformatter(args...) = false
format_julia_source(source) = source

has_ohmyrepl(args...) = false
highlight(source) = style(source, :file_contents)

function requires()
    @require Cthulhu = "f68482b8-f384-11e8-15f7-abe071a5a75f" begin
        has_cthulhu() = true
        ascend(mi::Core.MethodInstance) = Cthulhu.ascend(mi)
        descend(mi::Core.MethodInstance) = Cthulhu.descend(mi)
    end
    @require Debugger = "31a5f54b-26ea-5ae9-a837-f05ce5417438" begin
        has_debugger() = true
        breakpoint(file::AbstractString, line::Integer) = Debugger.breakpoint(file, line)
    end
    @require JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" begin
        has_jet() = true
        function report_call(mi::Core.MethodInstance)
            func = Base.tuple_type_head(mi.specTypes).instance
            sig = Base.tuple_type_tail(mi.specTypes)
            result = JET.report_call(func, sig)
            @info "Press return to continue."
            readline()
            return result
        end
    end
    @require JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899" begin
        has_juliaformatter() = true
        format_julia_source(source::String) = try JuliaFormatter.format_text(source); catch err; source; end
    end
    @require OhMyREPL = "5fb14364-9ced-5910-84b2-373655c76a03" begin
        has_ohmyrepl() = true
        function highlight(source::String)
            O = OhMyREPL
            tokens = collect(O.tokenize(source))
            crayons = fill(O.Crayon(), length(tokens))
            O.Passes.SyntaxHighlighter.SYNTAX_HIGHLIGHTER_SETTINGS(crayons, tokens, 0)
            io = IOBuffer()
            for (token, crayon) in zip(tokens, crayons)
                print(io, crayon)
                print(io, O.untokenize(token))
                print(io, O.Crayon(reset = true))
            end
            return String(take!(io))
        end
    end
end

#
# Module Initialisation.
#

function __init__()
    setup_repl()
    requires()
end

end # module


text:
using Plots

Δ = .1
N = 100
τ = N*Δ/5
ran = 0.0:2pi/N:2pi
function fillxs(Ts = 1,N = 100, ini=0.0, fin=100)
    xs = []
    Δ = fin/N
    τ = N*Δ/5

    for j in 1:Ts
        for i in ini:Δ:τ-Δ
            push!(xs,1)
        end

        for i in τ:Δ:fin
            push!(xs,0)
        end

    end

    return xs
end

function Xforward(x,n)
    N = length(x)
    j = im
    b = 2pi/N
    sum = 0.0
    for k in 1:N
        sum += x[k] * exp(-j * k * b * n)
    end

    return abs(1/N * sum)
end

xs = Int.(fillxs())

ys =  [Xforward(xs,i) for i in 0:1:N]

scatter(xs, title = "Train signal")
scatter(ys, title = "DFT")

function xback(x,n)
    N = length(x)
    j = im
    b = 2pi/N
    sum = 0.0
    for k in 1:N
        sum += x[k] * exp(j * k * b * n)
    end

    return abs(sum)
end

ys2 = [xback(xs,i) for i in 0:1:N]

scatter(ys2, title = "DFT back")


text:
import jInv.ForwardShare.getSensMat
function getSensMat(m::Vector,pFor::DipParam)
	if pFor.method == MATFree || pFor.method == MATBased
		error("getSensMat: Not supported or too expensive");
	end
	return pFor.Jacobian';
end

export getSensMatVec
function getSensMatVec(v::Vector,m::Vector,pFor::DipParam)
if pFor.method == MATFree
	return getData(v,pFor)[1];
else
	S = pFor.Jacobian;
	d = S'*v[:];
	ndips = size(pFor.theta_phi_rad,1);
	d = reshape(d,div(length(d),ndips),ndips);
end
end

export getSensTMatVec
function getSensTMatVec(v::Vector,m::Vector,pFor::DipParam)
JtV = 0;
if pFor.method == MATFree
	Mesh = pFor.Mesh;
	n = Mesh.n;
	theta_phi = pFor.theta_phi_rad;
	b = pFor.b;
	ndips = size(theta_phi,1);
	sampleBinning = pFor.samplingBinning;
	traceLength = div(n[3],sampleBinning);
	d = reshape(v,traceLength,ndips);
	u = zeros(tuple(n...));
	JtV = zeros(size(m));
	doTranspose = true;
	for ii = 1:ndips
		sampleSlicesT(d[:,ii],prod(Mesh.h),u);
		#JtV .+= rotateAndMove3DTranspose(u,theta_phi[ii,:],(b[ii,:]./Mesh.h))[:];
		(mr,XT,XTT) = rotateAndMove3D(u,theta_phi[ii,:],(b[ii,:]./Mesh.h),doTranspose);
		JtV .+= mr[:];
	end
else
	S = pFor.Jacobian;
	JtV = S*v[:];
end
return JtV;
end

max_stars_repo_name: BGUCompSci/ShapeReconstructionPaLS

max_issues_repo_name: BGUCompSci/ShapeReconstructionPaLS

max_forks_repo_name: BGUCompSci/ShapeReconstructionPaLS


text:
module MarketObservables

using Observables, Dates

include("types.jl")
include("Ticker.jl")
include("orderbook/base.jl")

end # module


text:
@testset "Ridge" begin
    a2 = rand(2,2)
    C = rand(5,2)
    θ = rand(2)
    resid = rand(5)
    m = Ridge(rand(5), C, a2, rand(2), a2, 1.0, a2, nothing, 1.0, θ, resid, a2, 1.0, 1)
    @test modelmatrix(m) === C
    @test coef(m) === θ
    @test residuals(m) === resid
    @test sprint(show, m) == "Ridge regression"
end

@testset "SearchCriterion" begin
    c = LOOCV()
    @test sprint(show, c) == "LOOCV"
    @test sprint(show, MIME("text/plain"), c) == "Leave-one-out cross validation"

    c = GCV()
    @test sprint(show, c) == "GCV"
    @test sprint(show, MIME("text/plain"), c) == "Generalized cross validation"

    c = AIC()
    @test sprint(show, c) == "AIC"
    @test sprint(show, MIME("text/plain"), c) == "Akaike information criterion"
end

@testset "SmoothAlgorithm" begin
    a = DemmlerReinsch()
    @test sprint(show, a) == "DemmlerReinsch"
    @test sprint(show, MIME("text/plain"), a) == "Demmler-Reinsch orthogonalization"

    a = DirectSolve()
    @test sprint(show, a) == "DirectSolve"
end

@testset "GridSearch" begin
    @test grid(1) == GridSearch([1.0])
    @test grid(1:3) == GridSearch([1.0, 2.0, 3.0])
    @test grid([2.0, 1]) == GridSearch([2.0, 1.0])
    @test grid().v == exp.(LinRange(-6, 12, 50))
    @test grid([2 1], algo=DirectSolve()) == GridSearch([2.0, 1.0], algo=DirectSolve())
    @test_throws ArgumentError grid([2.0, -1.0])

    s = grid()
    @test sprint(show, s) == "GridSearch"
    @test sprint(show, MIME("text/plain"), s)[1:80] == """
        GridSearch{DemmlerReinsch, Float64} across 50 candidate values:
          [0.00247875, 0"""
end

@testset "GridSearchResult" begin
    iopt = Dict{SearchCriterion,Int}()
    iopt[LOOCV()] = 5
    iopt[GCV()] = 1
    iopt[AIC()] = 3
    v = rand(2)
    sr = GridSearchResult(iopt, rand(2,2), v, v, v, v, v, v)
    @test sprint(show, sr) == "GridSearchResult"
    @test sprint(show, MIME("text/plain"), sr) == """
        GridSearchResult across 2 candidate values:
          LOOCV => 5
          GCV   => 1
          AIC   => 3"""
end

@testset "SmoothLP" begin
    e = SmoothLP(:v)
    @test sprint(show, e) == "SmoothLP"
    @test sprint(show, MIME("text/plain"), e) == """
        Smooth Local Projection:
          smoothed regressor: v
          polynomial order: 3
          finite difference order: 3"""
    e = SmoothLP([:v1, 2], 2, 2, search=grid(), criterion=AIC(), fullX=true)
    @test sprint(show, MIME("text/plain"), e) == """
        Smooth Local Projection:
          smoothed regressors: v1 2
          polynomial order: 2
          finite difference order: 2"""
end

@testset "_basismatrix" begin
    b3 = [1/6 2/3 1/6]
    @test _basismatrix(3, 0, 0) ≈ b3
    bm = _basismatrix(3, 0, 2)
    @test size(bm) == (3,5)
    m = zeros(3, 5)
    for r in 1:3
        m[r,r:r+2] = b3
    end
    @test bm ≈ m
    @test _basismatrix(2, 0, 0) ≈ [0.5 0.5]
end

@testset "_makeYSr" begin
    T = 100
    ys, ss, xs = [randn(T)], [randn(T)], [randn(T)]
    ws = ys
    res, X, T1, e1, e2 = _makeYSr(ys, ss, xs, ws, nothing, 3, 5, nothing)
    @test T1 == T-8
    @test size(res) == (T1, 2)
    @test all(e1)
    @test all(e2)

    xs = ()
    res, X, T1, e1, e2 = _makeYSr(ys, ss, xs, ws, nothing, 1, 0, nothing)
    @test T1 == 99
    @test size(res) == (T1, 2)
    @test all(e1)
    @test all(e2)

    ys[1][2], ys[1][3] = NaN, Inf
    xs = (convert(Vector{Union{Float64, Missing}}, randn(T)),)
    xs[1][3] = missing
    res, X, T1, e1, e2 = _makeYSr(ys, ss, xs, ws, nothing, 1, 0, nothing)
    @test T1 == 96
    @test size(res) == (T1, 2)
    @test e1 == ((1:99).>3)
    @test all(e2)

    res, X, T1, e1, e2 = _makeYSr(ys, ss, xs, ws, nothing, 1, 1, ((1:100).<60).|((1:100).>=70))
    @test T1 == 83
    @test e1 == ((1:98).>3) .& (((1:98).<58) .| ((1:98).>=70))
    @test all(e2)

    ys, xs = [randn(T)], [randn(T)]
    ss[1][2], ss[1][4] = NaN, Inf
    ws = ys
    res, X, T1, e1, e2 = _makeYSr(ys, ss, xs, ws, nothing, 2, 0, nothing)
    @test T1 == 97
    @test all(e1)
    @test e2 == ((1:98).!=2)
end

@testset "slp" begin
    # Check estimates are close to OLS with low penalty
    df = exampledata(:gk)
    ns = (:logcpi, :logip, :ff, :ebp, :ff4_tc)
    r0 = lp(df, :ebp, xnames=:ff4_tc, wnames=ns, nlag=12, nhorz=48, vce=HRVCE())
    f0 = irf(r0, :ebp, :ff4_tc)
    est = SmoothLP(:ff4_tc, search=grid(1e-8), criterion=GCV())
    r1 = lp(est, df, :ebp, xnames=:ff4_tc, wnames=ns, nlag=12, nhorz=48, vce=HRVCE())
    f1 = irf(r1, :ebp, :ff4_tc)
    @test f1.B ≈ f0.B atol=1e-4
    # Check HR confidence interval
    ci1 = confint(f1)
    @test ci1[1][1] ≈ -0.04758915684832821 atol=1e-8
    @test ci1[2][1] ≈ 0.8758551271581806 atol=1e-8
    @test ci1[1][10] ≈ -0.35800938178031794 atol=1e-8
    @test ci1[2][10] ≈ 2.265871824705655 atol=1e-8

    @test sprint(show, MIME("text/plain"), r1) == """
        LocalProjectionResult with 12 lags over 48 horizons:
        ──────────────────────────────────────────────────────────────────────────────
        Variable Specifications
        ──────────────────────────────────────────────────────────────────────────────
        Outcome variable:                 ebp    Minimum horizon:                    0
        Regressors:           ff4_tc constant
        Lagged controls:                                    logcpi logip ff ebp ff4_tc
        ──────────────────────────────────────────────────────────────────────────────
        Smooth Local Projection
        ──────────────────────────────────────────────────────────────────────────────
        Smoothing parameter:             0.00    Smoothed regressor:            ff4_tc
        Polynomial order:                   3    Finite difference order:            3
        Selection criterion:              GCV    Selection algorithm:   DemmlerReinsch
        Leave-one-out CV:             2212.02    Generalized CV:               2212.25
        Akaike information:              7.70    Residual sum of squares:      2193.42
        ──────────────────────────────────────────────────────────────────────────────"""

    # Compare estimates with Matlab results from Barnichon and Brownlees (2019)
    df = exampledata(:bb)
    ns = (:ir, :pi, :yg)
    # 194 is their optimal smoothing parameter
    est = SmoothLP(:ir, 3, 2, search=grid(194), criterion=LOOCV())
    r1 = lp(est, df, :yg, xnames=ns, wnames=ns, nlag=4, nhorz=20, minhorz=1)
    # δ is the normalization factor they use taken from Matlab
    δ = 0.802327024988408
    f1 = irf(r1, :yg, :ir)
    B1 = f1.B.*δ
    @test B1[1] ≈ -0.189377969344491 atol=1e-10
    @test B1[2] ≈ -0.429648467594006 atol=1e-10
    @test B1[3] ≈ -0.533856512472860 atol=1e-10
    @test B1[9] ≈ -0.022406955605936 atol=1e-10
    @test B1[10] ≈ 0.146703133067349 atol=1e-10
    @test B1[11] ≈ 0.303602580252881 atol=1e-10
    @test B1[18] ≈ 0.222464247242188 atol=1e-10
    @test B1[19] ≈ 0.096090022750258 atol=1e-10
    @test B1[20] ≈ -0.065709407174569 atol=1e-10

    # Check EWC confidence interval
    ci1 = confint(f1)
    @test ci1[1][1] ≈ -0.8571544691033884 atol=1e-8
    @test ci1[2][1] ≈ 0.3850826992477342 atol=1e-8
    @test ci1[1][10] ≈ 0.020165100998330715 atol=1e-8
    @test ci1[2][10] ≈ 0.34552900750920656 atol=1e-8

    # Compare results based on DemmlerReinsch and DirectSolve
    gbb = 194.0.*(1:0.5:10)
    iopt = Dict{SearchCriterion,Int}(c=>1 for c in (LOOCV(), GCV(), AIC()))
    est1 = SmoothLP(:ir, 3, 2, search=grid(gbb), criterion=LOOCV())
    r1 = lp(est1, df, :yg, xnames=ns, wnames=ns, nlag=4, nhorz=20, minhorz=1)
    @test r1.estres.λ == 194
    @test r1.estres.search.iopt == iopt
    @test r1.estres.rss ≈ sum(r1.estres.m.resid.^2)
    est2 = SmoothLP(:ir, 3, 2, search=grid(gbb, algo=DirectSolve()), criterion=GCV())
    r2 = lp(est2, df, :yg, xnames=ns, wnames=ns, nlag=4, nhorz=20, minhorz=1)
    @test r2.estres.λ == 194
    @test r2.estres.search.iopt == iopt
    @test r2.estres.rss ≈ sum(r2.estres.m.resid.^2)
    @test r1.estres.search.θs ≈ r2.estres.search.θs atol=1e-2
    @test r1.estres.search.loocv ≈ r2.estres.search.loocv atol=1e-1
    @test r1.estres.search.rss ≈ r2.estres.search.rss atol=1e-1
    @test r1.estres.search.gcv ≈ r2.estres.search.gcv atol=1e-1
    @test r1.estres.search.aic ≈ r2.estres.search.aic atol=1e-5
    @test r1.estres.search.dof_fit ≈ r2.estres.search.dof_fit atol=1e-2
    # Check the recalculated final estimates are the same
    @test r1.estres.θ ≈ r2.estres.θ
    @test r1.estres.Σ ≈ r2.estres.Σ
    @test r1.estres.loocv ≈ r2.estres.loocv
    @test r1.estres.rss ≈ r2.estres.rss
    @test r1.estres.gcv ≈ r2.estres.gcv
    @test r1.estres.aic ≈ r2.estres.aic
    @test r1.estres.dof_fit ≈ r2.estres.dof_fit
    @test r1.estres.dof_res ≈ r2.estres.dof_res

    # Compare estimates from LeastSquaresLP
    r0 = lp(df, :yg, xnames=ns, wnames=ns, nlag=4, nhorz=20, minhorz=1)
    f0 = irf(r0, :yg, :ir)
    est3 = SmoothLP(:ir, 3, 2, search=grid(1e-8), criterion=AIC())
    r3 = lp(est3, df, :yg, xnames=ns, wnames=ns, nlag=4, nhorz=20, minhorz=1)
    f3 = irf(r3, :yg, :ir)
    @test f3.B ≈ f0.B atol=1e-8

    # Make a benchmark for comparisons
    est0 = SmoothLP(:ir, 3, 2, search=grid(194.0))
    r0 = lp(est0, df, :yg, xnames=ns, wnames=ns, nlag=4, nhorz=20, minhorz=1, vce=HRVCE())
    r4 = lp(r1, HRVCE())
    @test r4.B ≈ r1.B
    @test r4.V ≈ r0.V
    f4 = irf(r4, :yg, :ir)
    ci4 = confint(f4)
    @test ci4[1][1] ≈ -0.6790290214946426 atol=1e-8
    @test ci4[2][1] ≈ 0.20695725163898837 atol=1e-8
    @test ci4[1][10] ≈ -0.06529955673073387 atol=1e-8
    @test ci4[2][10] ≈ 0.4309936652382712 atol=1e-8

    r5 = lp(r1, 194.0, vce=HRVCE())
    @test r5.B ≈ r1.B
    @test r5.V ≈ r0.V
    est0 = SmoothLP(:ir, 3, 2, search=grid(30))
    r0 = lp(est0, df, :yg, xnames=ns, wnames=ns, nlag=4, nhorz=20, minhorz=1, vce=HRVCE())
    r6 = lp(r1, 30, vce=HRVCE())
    @test r6.B ≈ r0.B
    @test r6.V ≈ r0.V

    @test sprint(show, MIME("text/plain"), r1) == """
        LocalProjectionResult with 4 lags over 20 horizons:
        ──────────────────────────────────────────────────────────────────────────────
        Variable Specifications
        ──────────────────────────────────────────────────────────────────────────────
        Outcome variable:                  yg    Minimum horizon:                    1
        Regressors:         ir pi yg constant    Lagged controls:             ir pi yg
        ──────────────────────────────────────────────────────────────────────────────
        Smooth Local Projection
        ──────────────────────────────────────────────────────────────────────────────
        Smoothing parameter:           194.00    Smoothed regressor:                ir
        Polynomial order:                   3    Finite difference order:            2
        Selection criterion:            LOOCV    Selection algorithm:   DemmlerReinsch
        Leave-one-out CV:            34530.21    Generalized CV:              34511.09
        Akaike information:             10.45    Residual sum of squares:     34379.97
        ──────────────────────────────────────────────────────────────────────────────"""

    df = exampledata(:rz)
    est = SmoothLP(Cum(:g), 3, 3, search=grid([1, 1e3, 1e5]))
    r1 = lp(est, df, Cum(:y), xnames=Cum(:g), wnames=(:newsy, :y, :g), iv=Cum(:g)=>:newsy,
        nlag=4, nhorz=17, addylag=false, firststagebyhorz=true, subset=df.wwii.==0)
    @test all(i->i==3, values(r1.estres.search.iopt))
    f1 = irf(r1, Cum(:y), Cum(:g))
    @test coef(f1)[1] ≈ 0.6969380791400279 atol=1e-8
    @test coef(f1)[9] ≈ 0.8444504165548141 atol=1e-8
    @test coef(f1)[17] ≈ 0.7343077490303411 atol=1e-8
    @test stderror(f1)[1] ≈ 0.46635397544693047 atol=1e-8
    @test stderror(f1)[9] ≈ 0.08475156240143418 atol=1e-8
    @test stderror(f1)[17] ≈ 0.12534835370652056 atol=1e-8

    r2 = lp(est, df, Cum(:y), xnames=Cum(:g), wnames=(:newsy, :y, :g),
        iv=Cum(:g)=>(:newsy, :g), nlag=4, nhorz=16, minhorz=1, addylag=false,
        firststagebyhorz=true, subset=df.wwii.==0)
    @test all(i->i==3, values(r2.estres.search.iopt))
    f2 = irf(r2, Cum(:y), Cum(:g))
    @test coef(f2)[1] ≈ -0.178375389687782 atol=1e-8
    @test coef(f2)[8] ≈ 0.2573602225267132 atol=1e-8
    @test coef(f2)[16] ≈ 0.23087128211782626 atol=1e-8
    @test stderror(f2)[1] ≈ 0.09017317177271794 atol=1e-8
    @test stderror(f2)[8] ≈ 0.03145777721639394 atol=1e-8
    @test stderror(f2)[16] ≈ 0.06985430057347854 atol=1e-8

    @test sprint(show, MIME("text/plain"), r2) == """
        LocalProjectionResult with 4 lags over 16 horizons:
        ──────────────────────────────────────────────────────────────────────────────
        Variable Specifications
        ──────────────────────────────────────────────────────────────────────────────
        Outcome variable:              Cum(y)    Minimum horizon:                    1
        Regressors:           Cum(g) constant    Lagged controls:            newsy y g
        Endogenous variable:           Cum(g)    Instruments:                  newsy g
        ──────────────────────────────────────────────────────────────────────────────
        Smooth Local Projection
        ──────────────────────────────────────────────────────────────────────────────
        Smoothing parameter:        100000.00    Smoothed regressor:            Cum(g)
        Polynomial order:                   3    Finite difference order:            3
        Selection criterion:            LOOCV    Selection algorithm:   DemmlerReinsch
        Leave-one-out CV:             1595.11    Generalized CV:               1594.28
        Akaike information:              7.37    Residual sum of squares:      1592.98
        ──────────────────────────────────────────────────────────────────────────────"""
end


text:
"""
    Metaplex{T} <: AbstractEpidemicModel

A type representing a metaplex, i.e. a metapopulation model with an embedded individual network, with arbitrary dynamics.
"""
struct Metaplex{T} <: AbstractEpidemicModel where T
    g::AbstractSimpleGraph # individual relation network
    h::AbstractSimpleGraph # population network
    D::Vector{Float64}
    dynamics::T
end

Metaplex(g, h, D::Real, dynamics) = Metaplex(
    g,
    h,
    fill(float(D), num_states(dynamics)),
    dynamics
    )

Metapopulation(mpx::Metaplex) = Metapopulation(mpx.h, mpx.D, mpx.dynamics)

ContactProcess(mpx::Metaplex) = ContactProcess(mpx.g, mpx.dynamics)

function init_rates(mpx::Metaplex, state)
    N = nv(mpx.g)
    a = zeros(N)
    init_rates!(a, state, mpx)
    return a
end

output_type(::Metaplex) = Int

progressbar(::Metaplex) = true


text:
function metropolis_hastings end

check_is_kernel(::typeof(metropolis_hastings)) = true
is_custom_primitive_kernel(::typeof(metropolis_hastings)) = false
reversal(::typeof(metropolis_hastings)) = metropolis_hastings

"""
    (new_trace, accepted) = metropolis_hastings(
        trace, selection::Selection;
        check=false, observations=EmptyChoiceMap())

Perform a Metropolis-Hastings update that proposes new values for the selected addresses from the internal proposal (often using ancestral sampling), returning the new trace (which is equal to the previous trace if the move was not accepted) and a Bool indicating whether the move was accepted or not.
"""
function metropolis_hastings(
        trace, selection::Selection;
        check=false, observations=EmptyChoiceMap())
    args = get_args(trace)
    argdiffs = map((_) -> NoChange(), args)
    (new_trace, weight) = regenerate(trace, args, argdiffs, selection)
    check && check_observations(get_choices(new_trace), observations)
    if log(rand()) < weight
        # accept
        return (new_trace, true)
    else
        # reject
        return (trace, false)
    end
end

"""
    (new_trace, accepted) = metropolis_hastings(
        trace, proposal::GenerativeFunction, proposal_args::Tuple;
        check=false, observations=EmptyChoiceMap())

Perform a Metropolis-Hastings update that proposes new values for some subset of random choices in the given trace using the given proposal generative function, returning the new trace (which is equal to the previous trace if the move was not accepted) and a Bool indicating whether the move was accepted or not.

The proposal generative function should take as its first argument the current trace of the model, and remaining arguments `proposal_args`.
If the proposal modifies addresses that determine the control flow in the model, values must be provided by the proposal for any addresses that are newly sampled by the model.
"""
function metropolis_hastings(
        trace, proposal::GenerativeFunction, proposal_args::Tuple;
        check=false, observations=EmptyChoiceMap())
    # TODO add a round trip check
    model_args = get_args(trace)
    argdiffs = map((_) -> NoChange(), model_args)
    proposal_args_forward = (trace, proposal_args...,)
    (fwd_choices, fwd_weight, _) = propose(proposal, proposal_args_forward)
    (new_trace, weight, _, discard) = update(trace,
        model_args, argdiffs, fwd_choices)
    proposal_args_backward = (new_trace, proposal_args...,)
    (bwd_weight, _) = assess(proposal, proposal_args_backward, discard)
    alpha = weight - fwd_weight + bwd_weight
    check && check_observations(get_choices(new_trace), observations)
    if log(rand()) < alpha
        # accept
        return (new_trace, true)
    else
        # reject
        return (trace, false)
    end
end

"""
    (new_trace, accepted) = metropolis_hastings(
        trace, proposal::GenerativeFunction, proposal_args::Tuple,
        involution::Union{TraceTransformDSLProgram,Function};
        check=false, observations=EmptyChoiceMap())

Perform a generalized (reversible jump) Metropolis-Hastings update based on an involution (bijection that is its own inverse) on a space of choice maps, returning the new trace (which is equal to the previous trace if the move was not accepted) and a Bool indicating whether the move was accepted or not.

Most users will want to construct `involution` using the [Trace Transform DSL](@ref) with the [`@transform`](@ref) macro, but for more user control it is also possible to provide a Julia function for `involution`, that has the following signature:

    (new_trace, bwd_choices::ChoiceMap, weight) = involution(trace::Trace, fwd_choices::ChoiceMap, fwd_retval, fwd_args::Tuple)

The generative function `proposal` is executed on arguments `(trace, proposal_args...)`, producing a choice map `fwd_choices` and return value `fwd_ret`.
For each value of model arguments (contained in `trace`) and `proposal_args`, the `involution` function applies an involution that maps the tuple `(get_choices(trace), fwd_choices)` to the tuple `(get_choices(new_trace), bwd_choices)`.
Note that `fwd_ret` is a deterministic function of `fwd_choices` and `proposal_args`.
When only discrete random choices are used, the `weight` must be equal to `get_score(new_trace) - get_score(trace)`.

When continuous random choices are used, the `weight` returned by the involution must include an additive correction term that is the determinant of the the Jacobian of the bijection on the continuous random choices that is obtained by currying the involution on the discrete random choices (this correction term is automatically computed if the involution is constructed using the [Trace Transform DSL](@ref)).
NOTE: The Jacobian matrix of the bijection on the continuous random choices must be full-rank (i.e. nonzero determinant).
The `check` keyword argument to the involution can be used to enable or disable any dynamic correctness checks that the involution performs; for successful executions, `check` does not alter the return value.
"""
function metropolis_hastings(
        trace, proposal::GenerativeFunction,
        proposal_args::Tuple, involution::Union{TraceTransformDSLProgram,Function};
        check=false, observations=EmptyChoiceMap())
    trace_translator = SymmetricTraceTranslator(proposal, proposal_args, involution)
    (new_trace, log_weight) = trace_translator(trace; check=check, observations=observations)
    if log(rand()) < log_weight
        # accept
        (new_trace, true)
    else
        # reject
        (trace, false)
    end
end

"""
    (new_trace, accepted) = mh(trace, selection::Selection; ..)
    (new_trace, accepted) = mh(trace, proposal::GenerativeFunction, proposal_args::Tuple; ..)
    (new_trace, accepted) = mh(trace, proposal::GenerativeFunction, proposal_args::Tuple, involution; ..)

Alias for [`metropolis_hastings`](@ref). Perform a Metropolis-Hastings update on the given trace.
"""
const mh = metropolis_hastings

"""
    (new_trace, accepted) = involutive_mcmc(trace, proposal::GenerativeFunction, proposal_args::Tuple, involution; ..)

Alias for the involutive form of [`metropolis_hastings`](@ref).
"""
function involutive_mcmc(trace, proposal, proposal_args, involution; check=false, observations=EmptyChoiceMap())
    return metropolis_hastings(
        trace, proposal, proposal_args, involution;
        check=check, observations=observations)
end

export metropolis_hastings, mh, involutive_mcmc

max_stars_repo_name: georgematheos/Gen

max_issues_repo_name: georgematheos/Gen

max_forks_repo_name: georgematheos/Gen


text:
"""
Exported module that provides build and deploy dependencies and related functions.

Currently only [`pip`](@ref) is implemented.
"""
module Deps

export pip

using DocStringExtensions

"""
$(SIGNATURES)

Installs (as non-root user) all python packages listed in `deps`.

# Examples

```julia
using Documenter

makedocs(
    # ...
)

deploydocs(
    deps = Deps.pip("pygments", "mkdocs", "mkdocs-material"),
    # ...
)
```
"""
pip(deps...) = () -> map(dep -> run(`pip install --user $(dep)`), deps)

function localbin()
    Sys.islinux() ? joinpath(homedir(), ".local", "bin") :
    Sys.isapple() ? joinpath(homedir(), "Library", "Python", "2.7", "bin") : ""
end

function updatepath!(p = localbin())
    if occursin(p, ENV["PATH"])
        ENV["PATH"]
    else
        ENV["PATH"] = "$p:$(ENV["PATH"])"
    end
end

end


text:
# Plot the loss function on a random 1D line
# Authors: Ali Siahkoohi, alisk@gatech.edu
# Date: December 2020

using DrWatson
@quickactivate :ExtendedConv

using Flux
using PyPlot
using Seaborn
using Printf
using LinearAlgebra
using Optim
using JLD2
using Random
using ProgressMeter
seed1 = 21
seed2 = 12

Random.seed!(seed1)

matplotlib.use("Agg")
set_style("whitegrid")
rc("font", family="serif", size=10)
font_prop = matplotlib.font_manager.FontProperties(
    family="serif",
    style="normal",
    size=10
)
sfmt=matplotlib.ticker.ScalarFormatter(useMathText=true)
sfmt.set_powerlimits((0, 0))

sim_name = "2D-visualization"

nx = 16
ny = 16
nc = 1
batchsize = 1
k = 3
nlayers = 4
lambda = .04f0
f_tol = 1f-6
iterations = 100

G = Chain([Conv((k, k), nc => nc, sigmoid, pad=(1, 1), stride=(1, 1))  for i=1:nlayers]...)

w = deepcopy(Flux.params(G))
w_true = cat([w[i] for i=1:2:2*nlayers]..., dims=3)[:, :, :, 1]

C = extended_conv_op(nx, ny; nc=nc, batchsize=1, k=k, nlayers=nlayers, w=w_true)

z_std = 1f-2
z = z_std*randn(Float32, nx, ny, nc, batchsize)
y = G(z)

G_ext(x) = reshape(C(reshape(x, :)), nx, ny, nc, 1)[end:-1:1, end:-1:1]
perm_y(y, idx) = reshape(y[end:-1:1, end:-1:1, :, idx], :)

αs = range(-4f0, 4f0, step=4f-1)
βs = range(-4f0, 4f0, step=4f-1)
n_grid = length(αs)

Random.seed!(seed2)

opt_var_proj = Optim.Options(
    iterations=iterations,
    f_tol=f_tol,
    show_trace=false,
    store_trace=false,
    show_every=1
)

num_figs = 4
p = Progress(num_figs*n_grid^2)

for fig_idx = 1:num_figs
    v1 = filter_normalization(w)
    v2 = filter_normalization(w)

    loss = zeros(Float32, length(αs), length(βs))
    loss_ext = zeros(Float32, length(αs), length(βs))

    for (i, α) in enumerate(αs)
        for (j, β) in enumerate(βs)
            Base.flush(Base.stdout)

            idx = rand(1:batchsize)

            ŵ = w + α*v1 + β*v2
            Flux.loadparams!(G, ŵ)

            ŵ0 = cat([ŵ[i] for i=1:2:2*nlayers]..., dims=3)[:, :, :, 1]
            C = extended_conv_op(nx, ny; nc=nc, batchsize=1, k=k, nlayers=nlayers, w=ŵ0)
            w_ex_0 = cat(
                [randn(Float32, size(C[i].mask)).*C[i].mask for i=1:length(C)]...,
                dims=3
            )

            loss_ext[i, j] = var_proj_optim(
                reshape(z[:, :, :, idx:idx], :), perm_y(y, idx), C, w_ex_0, ŵ0, lambda,
                opt_var_proj
            )[1]

            loss[i, j] = 5f-1*norm(G(z[:, :, :, idx:idx]) - y[:, :, :, idx:idx])^2

            ProgressMeter.next!(
                p;
                showvalues = [
                    ("Figure number", fig_idx),
                    (:α, α),
                    (:β, β),
                    ("Objectibe", loss[i, j]),
                    ("Extended objective", loss_ext[i, j])
                ]
            )

        end
    end

    save_dict = @strdict loss loss_ext lambda n_grid f_tol iterations batchsize z_std nlayers nx ny k v1 v2 seed1 seed2 w z y αs βs
    save_path = plotsdir(sim_name, savename(save_dict; digits=12))

    @tagsave(
        datadir(sim_name, savename(save_dict, "jld2"; digits=12)),
        save_dict;
        safe=true
    )

    loss = (loss .- minimum(loss))./maximum(loss)
    loss_ext = (loss_ext .- minimum(loss_ext))./maximum(loss_ext)

    extent = [minimum(αs), maximum(αs), maximum(βs), minimum(βs)]

    fig = figure("conventional", figsize=(4, 4))
    contour(αs, βs, loss, 25, cmap=Seaborn.ColorMap("mako"), linewidths=1.5,
        vmin=0.0, vmax=1f0)
    title("Loss landscape")
    colorbar(fraction=0.057, pad=0.01, format=sfmt)
    PyPlot.scatter(
        βs[argmin(loss)[2]], αs[argmin(loss)[1]], s=30.0,
        color="#b80000", marker="*"
    )
    xlabel(L"$\beta$")
    ylabel(L"$\alpha$")
    safesave(joinpath(save_path, "loss.png"), fig)
    close(fig)

    fig = figure("conventional", figsize=(4, 4))
    contour(αs, βs, loss_ext, 25, cmap=Seaborn.ColorMap("mako"), linewidths=1.5,
        vmin=0.0, vmax=1f0)
    title("Extended loss landscape")
    colorbar(fraction=0.057, pad=0.01, format=sfmt)
    PyPlot.scatter(
        βs[argmin(loss_ext)[2]], αs[argmin(loss_ext)[1]], s=30.0,
        color="#b80000", marker="*"
    )
    xlabel(L"$\beta$")
    ylabel(L"$\alpha$")
    safesave(joinpath(save_path, "loss_ext.png"), fig)
    close(fig)

    fig = figure("conventional", figsize=(4, 4))
    contourf(αs, βs, loss, 25, cmap=Seaborn.ColorMap("mako"), linewidths=1.5,
        vmin=-1f-1, vmax=1f0)
    title("Loss landscape")
    cb = colorbar(fraction=0.057, pad=0.01, format=sfmt)
    PyPlot.scatter(
        βs[argmin(loss)[2]], αs[argmin(loss)[1]], s=30.0,
        color="#b80000", marker="*"
    )
    xlabel(L"$\beta$")
    ylabel(L"$\alpha$")
    safesave(joinpath(save_path, "loss_contourf.png"), fig)
    close(fig)

    fig = figure("conventional", figsize=(4, 4))
    contourf(αs, βs, loss_ext, 25, cmap=Seaborn.ColorMap("mako"), linewidths=1.5,
        vmin=-1f-1, vmax=1f0)
    title("Extended loss landscape")
    colorbar(fraction=0.057, pad=0.01, format=sfmt)
    PyPlot.scatter(
        βs[argmin(loss_ext)[2]], αs[argmin(loss_ext)[1]], s=30.0,
        color="#b80000", marker="*"
    )
    xlabel(L"$\beta$")
    ylabel(L"$\alpha$")
    safesave(joinpath(save_path, "loss_ext_contourf.png"), fig)
    close(fig)

    fig = figure("trace", figsize=(3*1.5, 7.5*1.5))
    plt.plot(αs, loss[:, argmin(loss)[2]], color="b", ls="--")
    plt.plot(αs, loss_ext[:, argmin(loss_ext)[2]], color="b")

    plt.plot(βs, loss[argmin(loss)[1], :], color="k", ls="--")
    plt.plot(βs, loss_ext[argmin(loss_ext)[1], :], color="k")

    plt.plot(αs, diag(loss), color="r", ls="--")
    plt.plot(αs, diag(loss_ext), color="r")
    safesave(joinpath(save_path, "trace.png"), fig)
    close(fig)

end


text:
using NetCDF, Statistics, Plots, Distributed, StatGeochem, Polynomials

"""
Name:Nestor Orozco-Llamas
Class: EARS13
Final Project

!!!!!!!!!!!!!!!!! IMPORTANT!!!!!!!!!!!!!!!!

Notes about the datasets (Reanalysis data is a few gigabytes):

You can download the dataset using an ftp client accessing: ftp2.psl.noaa.gov
The datasets are all in the /Datasets/20thC_ReanV3/Dailies/accumsMO/ directory
In the directory you need to download:

apcp.1981.nc .... apcp.1981.nc              #Precip amounts
air.sig995.1981.nc .... air.sig995.2015.nc  #Temps
snowd.1981.nc .... snowd.2015.nc            #Snowd amounts

You can also download them online visiting:

http://psl.noaa.gov/cgi-bin/db_search/DBListFiles.pl?did=210&tid=76965&vid=5202
http://psl.noaa.gov/cgi-bin/db_search/DBListFiles.pl?did=210&tid=76964&vid=5246
http://psl.noaa.gov/cgi-bin/db_search/DBListFiles.pl?did=210&tid=76964&vid=5205

The station data was self-processed and made into a csv

"""

#-------------------  Data Processing Variables -----------------------------
#Data years
dataYears = 1981:2015
#Latitudes
lats = [[41,41,41,40,40],[39,38,38],[37,37,36]]
#Longitudes
longs = [[123,122,121,122,121],[121,121,120],[120,119,119]]

#Equation for index adjustment in reanalysis dataset
f(x) = 181 .+ (-1 .*x .+ 180)
g(x) = x .+ 91

#Mapping them to their index
lats = g.(lats)
longs = f.(longs)
#-------------------------Analysis Variables ----------------------------

#decade separation
decades = [[1981,1989],[1990,1999],[2000,2009],[2010,2015]]
#decade names for title purposes
decadeNames = ["1980s", "1990s", "2000s", "2010s"]
#for labeling
years = [1985,1995,2005,2015]

#switch for each analysis
pack = 1
#for naming purposes
packNames = ["Northern","Central","Southern"]
packName = packNames[pack]

#for resampling
randomSamples = 10_000

#------------------- Reanalysis Data Processing -----------------------------

#Data stored in dictionary
#Format:
#       String --> List
# ex.   y2000pack1 --> List of avg value each day of 2000 in pack 1
# snowd: snow depth (m)
# airtemp:  air temperature (F)
# precip: rainfall (mm)
# pack1: northern sierra snow pack
# pack2: central sierra snow pack
# pack3: southern sierra snow pack
snowd = Dict{String, Array{Float32}}()
airtemp = Dict{String, Array{Float32}}()
precip = Dict{String, Array{Float32}}()

#Storing data using nested for loops
#first loop is to iterate through each year
for i=dataYears
    #this iterates through each cluster coordinates
    for j=1:length(lats)
        #key for dictionary
        key = "y$(i)pack$(j)"
        #filling with zeros
        snowd[key] = zeros(Float32, 365)
        airtemp[key] = zeros(Float32, 365)
        precip[key] = zeros(Float32, 365)

        #iterating for each square grid for each cluster
        for z=1:length(lats[j])
            #adding the snow depth data
            curr = ncread("./Data/snod.$(i).nc", "snod")[longs[j][z],:,:][lats[j][z],:][1:365]
            curr[curr.<=0] .= 0 #removes any noisy value below 0
            snowd[key] = snowd[key] .+ curr
            #adding the temperature data in fahrenheit
            curr = ncread("./Data/air.sig995.$(i).nc", "air")[longs[j][z],:,:][lats[j][z],:][1:365]
            curr[curr.<=0] .= 0 #removes any noisy value below 0
            curr = (curr .- 273.15) .* (9/5) .+ 32 #kelvin to fahrenheit conversion
            airtemp[key] = airtemp[key] .+ curr
            #adding the precip data in mm each day
            curr = ncread("./Data/apcp.$(i).nc", "apcp")[longs[j][z],:,:][lats[j][z],:][1:365]
            curr[curr.<=0] .= 0 #removes any noisy value below 0
            curr = curr .* 8    #data is in 3 hour mean so mulply by 8 for 24 hour amount
            precip[key] = precip[key] .+ curr

        end
        #averaging them out by dividing by area of each cluster
        snowd[key] = snowd[key] ./ length(lats[j])
        airtemp[key] = airtemp[key] ./ length(lats[j])
        precip[key] = precip[key] ./ length(lats[j])
    end
    #Loading bar to keep track of progress
    print("\rProcessing Data: [")
    for r=1981:i  print("#") end
    for r=i:2015 print(" ") end
    print("] $(round((i-1980)/0.35;digits=2))% complete  ")
end

#------------------- Station Data Processing and Comparison to Reanalysis-----------------------------

#Importing stationData.csv as a named tuple
#Contains year, mean temperature, precip in inches, snowfall in inches
data = importdataset("stationData.csv",'\t',importas=:Tuple)

#annual total precipitation in Central Sierra in inches
annualPrecipPack2 = Array{Float32}(undef, 35)
for i=dataYears
    annualPrecipPack2[i-1980] = getPrecip(i, i, 1, 365, 2)
end

#regression line slope
slopeReanalysis = Polynomials.fit(dataYears,annualPrecipPack2,1)[1]
slopeStation = Polynomials.fit(data.year,data.precip,1)[1]

#plotting the annual precip in reanalysis vs station data
p1 = plot(dataYears,annualPrecipPack2,smooth=:true, xlabel = "Year",framestyle=:box, ylabel="Precipitation (in)",label="Reanalysis Data - Slope = $(round(slopeReanalysis;digits=3))", title = "Precipitation from 1981-2015")
plot!(data.year,data.precip,smooth=:true, label="Station Data - Slope = $(round(slopeStation;digits=3))")
savefig(p1, "./Figures/StationReanalysisPrecip.png")

#average annual temperature in Central Sierra
annualTempPack2 = Array{Float32}(undef, 35)
for i=dataYears
    annualTempPack2[i-1980] = getMeanTemp(i, i, 1, 365, 2)
end

#regression line slope
slopeReanalysis = Polynomials.fit(dataYears,annualTempPack2,1)[1]
slopeStation = Polynomials.fit(data.year,data.Tmean,1)[1]

#plotting the annual temperature in reanalysis vs station data
p2 = plot(dataYears,annualTempPack2,smooth=:true, framestyle=:box,legend=:topleft,xlabel = "Year", ylabel="Temperature (ºF)",label="Reanalysis Data - Slope = $(round(slopeReanalysis;digits=3))", title = "Average Temperature from 1981-2015")
plot!(data.year,data.Tmean,smooth=:true, label="Station Data - Slope = $(round(slopeStation;digits=3))")
savefig(p2, "./Figures/StationReanalysisTemp.png")

#average snow depth from january 1 to june 30
annualSnowPack2 = Array{Float32}(undef, 35)
for i=dataYears
    annualSnowPack2[i-1980] = getSnowDepth(i, i, 1, 151, 2)
end
#plotting the snow depth in reanalysis vs snowfall in station data
p3 = plot(dataYears,annualSnowPack2, smooth=:true,framestyle=:box,legend=:topleft,xlabel = "Year",label="Central Sierra - Reanalysis Data", title = "Snow Depth in Reanalysis vs Snowfall in Station Data")
plot!(data.year,data.snow,label="Tahoe City, CA - Station Data",smooth=:true,yticks=:false)
savefig(p3, "./Figures/StationReanalysisSnow.png")

#----------------------------- Data Analysis ------------------------------

#----------------------------- Snow Depth ------------------------------

#integral of snow depth for each season
sumSnowDepthWinter = Array{Float32}(undef, 35)
sumSnowDepthSpring = Array{Float32}(undef, 35)

#getting the area under the curve for each season
for i=dataYears
    curr = snowd["y$(i)pack$(pack)"]
    sumSnowDepthWinter[i-1980] = sum(vcat(curr[1:59],curr[335:365]))
    sumSnowDepthSpring[i-1980] = sum(curr[60:151])
end
#regression line slope
slopeWinter = Polynomials.fit(dataYears,sumSnowDepthWinter,1)[1]
slopeSpring = Polynomials.fit(dataYears,sumSnowDepthSpring,1)[1]

#plotting the figures
p4 = plot(dataYears,sumSnowDepthWinter,framestyle=:box,title="$(packName) Pack Snow Depth", ylabel="Snow Depth (m)", xlabel="Year",smooth=:true, label="Winter - Slope = $(round(slopeWinter;digits=3))")
plot!(dataYears,sumSnowDepthSpring, smooth=:true, label="Spring - Slope = $(round(slopeSpring;digits=3))")
savefig(p4, "./Figures/$(packName)-SD-Season.png")

#getting snow depth by decade for each day of the year
snowDepthByDecade = zeros(Float32, length(decades), 365)
p5 = plot(title="$(packName) Pack Snow Depth",framestyle=:box,ylabel="Snow Depth (m)", xlabel="Year")
#iterating through decades
for j=1:length(decades)
    #extracting the decade data
    for i=(decades[j][1]:decades[j][2])
        key = "y$(i)pack$(pack)"
        snowDepthByDecade[j,:] = snowDepthByDecade[j,:] .+ snowd[key]
    end
    #averaging them out and plotting
    snowDepthByDecade[j,:] = snowDepthByDecade[j,:]./(decades[j][2]-decades[j][1]+1)
    plot!(1:365, snowDepthByDecade[j,:], label="$(decadeNames[j])" )
end
savefig(p5, "./Figures/$(packName)SD-Decade-Daily.png")

#snow depth by decade with monte carlo resampling to get mean of means and std of means
snowd_decade = [getAllDataPoints(1981,1989,1,365,pack,snowd),getAllDataPoints(1990,1999,1,365,pack,snowd),getAllDataPoints(2000,2009,1,365,pack,snowd),getAllDataPoints(2010,2015,1,365,pack,snowd)]
snowd_means = Array{Float32}(undef,4,randomSamples)
#applying it for a large amount of times
for i=1:4
    for j=1:randomSamples
        snowd_means[i,j] = mean(rand(snowd_decade[i],randomSamples))
    end
end
#getting the mean of means and std of means and plotting the distrbution data
meansofmeans = [mean(snowd_means[1,:]),mean(snowd_means[2,:]),mean(snowd_means[3,:]),mean(snowd_means[4,:])]
stdsofmeans = [std(snowd_means[1,:]),std(snowd_means[2,:]),std(snowd_means[3,:]),std(snowd_means[4,:])]
p6 = plot(years, meansofmeans,yerr=2*stdsofmeans,seriestype=:scatter,framestyle=:box,title="$(packName) Sierra Resampled Mean Snow Depths",ylabel="Snow Depth (m)",xlabel="year",legend=:false)
savefig(p6, "./Figures/$(packName)SD-Decade-Resampled.png")

#----------------------------- Precip ------------------------------

#integral of precip for each season
sumPrecipFall = Array{Float32}(undef, 35)
sumPrecipWinter = Array{Float32}(undef, 35)
sumPrecipSpring = Array{Float32}(undef, 35)

#getting the area under the curve for each season and converting to inches
for i=dataYears
    curr = precip["y$(i)pack$(pack)"]

    sumPrecipFall[i-1980] = sum(curr[244:334])/25.4
    sumPrecipWinter[i-1980] = sum(vcat(curr[1:59],curr[335:365]))/25.4
    sumPrecipSpring[i-1980] = sum(curr[60:151])/25.4
end
#getting the slope for each season
slopeWinter = Polynomials.fit(dataYears,sumPrecipWinter,1)[1]
slopeSpring = Polynomials.fit(dataYears,sumPrecipSpring,1)[1]
slopeFall = Polynomials.fit(dataYears,sumPrecipFall,1)[1]
#plotting the data
p8 = plot(dataYears,sumPrecipWinter,framestyle=:box,title="$(packName) Sierra Precipitation", ylabel="Precipitation (in)", xlabel="Year",smooth=:true, label="Winter - Slope = $(round(slopeWinter;digits=3))")
plot!(dataYears,sumPrecipSpring, smooth=:true, label="Spring - Slope = $(round(slopeSpring;digits=3))")
plot!(dataYears,sumPrecipFall, smooth=:true, label="Fall - Slope = $(round(slopeFall;digits=3))")
savefig(p8, "./Figures/$(packName)-Precip-Season.png")

#Precip by decade with monte carlo resampling to get mean of means and std of means
precip_decade = [getAllDataPoints(1981,1989,1,365,pack,precip),getAllDataPoints(1990,1999,1,365,pack,precip),getAllDataPoints(2000,2009,1,365,pack,precip),getAllDataPoints(2010,2015,1,365,pack,precip)]
precip_means = Array{Float32}(undef,4,randomSamples)
#applying it for a large amount of times
for i=1:4
    for j=1:randomSamples
        precip_means[i,j] = mean(rand(precip_decade[i],randomSamples))/25.4
    end
end
#getting the mean of means and std of means and plotting the distrbution data
meansofmeans = [mean(precip_means[1,:]),mean(precip_means[2,:]),mean(precip_means[3,:]),mean(precip_means[4,:])]
stdsofmeans = [std(precip_means[1,:]),std(precip_means[2,:]),std(precip_means[3,:]),std(precip_means[4,:])]
p9=plot(years, meansofmeans,yerr=2*stdsofmeans,seriestype=:scatter,framestyle=:box,title="$(packName) Sierra Resampled Mean Precip",ylabel="Precip (in)",xlabel="year",legend=:false)
savefig(p9, "./Figures/$(packName)Precip-Decade-Resampled.png")

#----------------------------- Air Temperature ------------------------------

#integral of snow depth for each season
meanTempFall = Array{Float32}(undef, 35)
meanTempWinter = Array{Float32}(undef, 35)
meanTempSpring = Array{Float32}(undef, 35)
meanTempSummer = Array{Float32}(undef, 35)

#getting the area under the curve for each season and dividing by days for mean temp
for i=dataYears
    curr = airtemp["y$(i)pack$(pack)"]

    meanTempFall[i-1980] = sum(curr[244:334])/91
    meanTempWinter[i-1980] = sum(vcat(curr[1:59],curr[335:365]))/90
    meanTempSpring[i-1980] = sum(curr[60:151])/92
    meanTempSummer[i-1980] = sum(curr[152:243])/92
end

#getting the regression line slope
slopeWinter = Polynomials.fit(dataYears,meanTempWinter,1)[1]
slopeSpring = Polynomials.fit(dataYears,meanTempSpring,1)[1]
slopeFall = Polynomials.fit(dataYears,meanTempFall,1)[1]
slopeSummer= Polynomials.fit(dataYears,meanTempSummer,1)[1]
#plotting the figures
p11= plot(dataYears,meanTempWinter,legend=:topleft,framestyle=:box,title="$(packName) Sierra Temperatures", ylabel="Temperature (ºF)", xlabel="Year",smooth=:true, label="Winter - Slope = $(round(slopeWinter;digits=3))")
plot!(dataYears,meanTempSpring, smooth=:true, label="Spring - Slope = $(round(slopeSpring;digits=3))")
plot!(dataYears,meanTempFall, smooth=:true, label="Fall - Slope = $(round(slopeFall;digits=3))")
plot!(dataYears,meanTempSummer, smooth=:true, label="Summer - Slope = $(round(slopeSummer;digits=3))")
savefig(p11, "./Figures/$(packName)-Temp-Season.png")

#Getting the annual mean snow depth
sumSnowAnnual = Array{Float32}(undef, 4,35)
#Getting the annual mean temperature
meanTempAnnual = Array{Float32}(undef, 4,35)

for j=1:length(lats)
    for i=dataYears
        meanTempAnnual[j,i-1980] = sum(airtemp["y$(i)pack$(j)"])/365
        sumSnowAnnual[j,i-1980] = sum(snowd["y$(i)pack$(j)"])
    end
end

#Plotting the results of snow depth vs temperature for each pack
p13 = plot(seriestype=:scatter)
plot!(meanTempAnnual[1,:],sumSnowAnnual[1,:], xlabel = "Temperature (ºF)",ylabel="Snow Depth (m)", framestyle=:box,label="$(packNames[1]) Sierra",smooth=:true,seriestype=:scatter)
savefig(p13,"./Figures/$(packNames[1])TempvsSD.png" )
p14 = plot(seriestype=:scatter)
plot!(meanTempAnnual[2,:],sumSnowAnnual[2,:],xlabel = "Temperature (ºF)",ylabel="Snow Depth (m)", framestyle=:box, label="$(packNames[2]) Sierra",smooth=:true,seriestype=:scatter)
savefig(p14,"./Figures/$(packNames[2])TempvsSD.png" )
p15 = plot(seriestype=:scatter)
plot!(meanTempAnnual[3,:],sumSnowAnnual[3,:],xlabel = "Temperature (ºF)",ylabel="Snow Depth (m)", framestyle=:box, label="$(packNames[3]) Sierra",smooth=:true,seriestype=:scatter)
savefig(p15,"./Figures/$(packNames[3])TempvsSD.png" )

#-------------------------------- Pack Comparisons -----------------------------------
#Shows the average snow depth of each pack from day 1 to 365
snowdByPack = zeros(Float32, length(lats), 365)
p7 = plot(title="Pack Snow Depth Comparison",framestyle=:box,ylabel="Snow Depth (m)", xlabel="Day of Year")
for j=1:length(lats)
    for i=dataYears
        key = "y$(i)pack$(j)"
        snowdByPack[j,:] = snowdByPack[j,:] .+ snowd[key]
    end
    snowdByPack[j,:] = snowdByPack[j,:]./35
    plot!(1:365,snowdByPack[j,:], label="$(packNames[j]) Sierra")
end
savefig(p7, "./Figures/SD-Pack-Comp.png")

#Shows the total precip of each pack from day 1 to 365
p10 = plot(title="Pack Precipitation Comparison",framestyle=:box,ylabel="Precipitation (in)", xlabel="Day of Year")
for j=1:length(lats)
    precipPack = Array{Float32}(undef,365)
    for i=1:365
        precipPack[i] = getPrecip(1981,2015,i,i,j)
    end
    plot!(1:365,precipPack, label="$(packNames[j]) Sierra")
end
savefig(p10, "./Figures/Precip-Pack-Comp.png")

#Shows the average temperature of each pack from day 1 to 365
tempByPack = zeros(Float32, length(lats), 365)
p13 = plot(title="Pack Temperature Comparison",framestyle=:box,ylabel="Temperature (ºF)", xlabel="Day of Year")

for j=1:length(lats)
    for i=1981:2015
        key = "y$(i)pack$(j)"
        tempByPack[j,:] = tempByPack[j,:] .+ airtemp[key]
    end
    tempByPack[j,:] = tempByPack[j,:]./35
    plot!(1:365,tempByPack[j,:], label="$(packName) Sierra")
end
savefig(p13, "./Figures/Temp-Pack-Comp.png")

#--------------------------- functions -------------------------------

#Takes in a start year and end year, and start day and end day
# and returns all the data points for that data and pack in that period
# pack = 1 ---> northern
# pack = 2 ---> central
# pack = 3 ---> southern
function getAllDataPoints(startYear, endYear, startDay, endDay, pack, data)
    dataPoints = [data["y$(i)pack$(pack)"][j] for i in startYear:endYear for j in startDay:endDay]
    return dataPoints
end

#Takes in a start year and end year, and start day and end day
# and returns the total precip during that period in that pack
# pack = 1 ---> northern
# pack = 2 ---> central
# pack = 3 ---> southern
function getPrecip(startYear, endYear, startDay, endDay, pack)
    #mm
    totalPrecip = 0
    #sums the precip for all the days
    for i=startYear:endYear
        key = "y$(i)pack$(pack)"
        totalPrecip += sum(precip[key][startDay:endDay])
    end

    #Conversion to inches
    totalPrecip /= 25.4

    return totalPrecip
end

#Takes in a start year and endy year, and start day and end day
# and returns the mean temp during that period in that pack
# pack = 1 ---> northern
# pack = 2 ---> central
# pack = 3 ---> southern
function getMeanTemp(startYear, endYear, startDay, endDay, pack)
    #Mean
    meanTemp = 0
    #iterating through each year and summing the counts for those days
    for i=startYear:endYear
        key = "y$(i)pack$(pack)"
        meanTemp += sum(airtemp[key][startDay:endDay])
    end

    #Averaging out based on days
    meanTemp /= ((endDay-startDay+1)*(endYear-startYear+1))

    return meanTemp
end

#Takes in a start year and endy year, and start day and end day
# and returns the mean snow depth during that period in that pack
# pack = 1 ---> northern
# pack = 2 ---> central
# pack = 3 ---> southern
function getSnowDepth(startYear, endYear, startDay, endDay, pack)
    #mean
    meanDepth = 0

    #iterating through each year and summing the counts for those days
    for i=startYear:endYear
        key = "y$(i)pack$(pack)"
        meanDepth += sum(snowd[key][startDay:endDay])
    end

    #Averaging out based on days summed
    meanDepth /= ((endDay-startDay+1)*(endYear-startYear+1))
    #turns to millimeter
    meanDepth *= 1000
    return meanDepth
end

max_stars_repo_name: nestoreo/SierraNevadaCA-WaterAnalysis

max_issues_repo_name: nestoreo/SierraNevadaCA-WaterAnalysis

max_forks_repo_name: nestoreo/SierraNevadaCA-WaterAnalysis


text:
# ============================================================
# AbstractLazyReachSet: reach-set that behaves like a LazySet
# ============================================================

"""
    AbstractLazyReachSet{N} <: AbstractReachSet{N}

Abstract type for all reach-set types that use a `LazySet` for the underlying
set representation.

### Notes

An `AbstractLazyReachSet` is the interface for reach-sets such that the geometric
set is represented by any subtype of `LazySet`.

This types implements the `LazySets` interface, namely support function (`ρ`),
support vector (`σ`) and ambient dimension (`dim`) functions. Hence, these
functions directly apply to concrete subtypes of an `AbstractLazyReachSet`. The
set wrapped by this type is obtained through `set(R)`.

The following functions should be implemented by any concrete subtype:

- `reconstruct` -- create a new instance of the given reach-set with a different
                   set representation but sharing the other fields, i.e. the same
                   time span (and the same for other fields, if applicable)

In addition to the functions inherited from `AbstractReachSet`, the following
are available:

- `linear_map`  -- concrete linear map of a reach-set
- `project`     -- projection of a reach-set
- `shift`       -- time-shift of a reach-set
- `vars`        -- tuple of integers associated to the variables of the given reach-set
"""
abstract type AbstractLazyReachSet{N} <: AbstractReachSet{N} end

# Implement LazySets interface
LazySets.dim(R::AbstractLazyReachSet) = dim(set(R))

LazySets.ρ(d::AbstractVector, R::AbstractLazyReachSet) = ρ(d, set(R))
LazySets.ρ(d::AbstractVector, R::Vector{<:AbstractLazyReachSet}) = map(Ri -> ρ(d, set(Ri)), R) |> maximum
LazySets.ρ(d::AbstractVector, R::SubArray{<:AbstractLazyReachSet}) = map(Ri -> ρ(d, set(Ri)), R) |> maximum

LazySets.σ(d::AbstractVector, R::AbstractLazyReachSet) = σ(d, set(R))
LazySets.σ(d::AbstractVector, R::Vector{<:AbstractLazyReachSet}) = _σ_vec(d, Rvec)
LazySets.σ(d::AbstractVector, R::SubArray{<:AbstractLazyReachSet}) = _σ_vec(d, Rvec)

function _σ_vec(d, Rvec)
    σarray = map(Ri -> σ(d, set(Ri)), Rvec)
    ρarray = map(vi -> dot(d, vi), σarray)
    m = argmax(ρarray)
    return σarray[m]
end

# Expose common LazySets operations
LazySets.constraints_list(R::AbstractLazyReachSet) = constraints_list(set(R))
LazySets.vertices_list(R::AbstractLazyReachSet) = vertices_list(set(R))
Base.:⊆(R::AbstractLazyReachSet, X::LazySet) = ⊆(set(R), X)
Base.:⊆(X::LazySet, R::AbstractLazyReachSet) = ⊆(X, set(R))
Base.:⊆(R::AbstractLazyReachSet, S::AbstractLazyReachSet) = ⊆(set(R), set(S))
LazySets.area(R::AbstractLazyReachSet) = area(set(R))
LazySets.volume(R::AbstractLazyReachSet) = volume(set(R))
Base.convert(::Type{ST}, R::AbstractLazyReachSet) where {ST<:LazySet} = convert(ST, set(R))
Base.convert(::Type{<:IntervalBox}, R::AbstractLazyReachSet) = convert(IntervalBox, set(R))
LazySets.complement(R::AbstractLazyReachSet) = reconstruct(R, complement(set(R)))
Base.isempty(R::AbstractLazyReachSet) = isempty(set(R))

function LinearMap(M::Union{AbstractMatrix, Number}, R::AbstractLazyReachSet)
    return reconstruct(R, LinearMap(M, set(R)))
end

Base.:(*)(M::AbstractMatrix, R::AbstractLazyReachSet) = LinearMap(M, R)
Base.:(*)(α::AbstractFloat, R::AbstractLazyReachSet) = reconstruct(R, α * set(R))

function linear_map(M::AbstractMatrix, R::AbstractLazyReachSet)
    return reconstruct(R, linear_map(M, set(R)))
end

function linear_map(M::AbstractMatrix, R::Vector{<:AbstractLazyReachSet})
    return map(Ri -> reconstruct(Ri, linear_map(M, set(Ri))), R)
end

function overapproximate(R::AbstractLazyReachSet, func)
    return reconstruct(R, overapproximate(set(R), func))
end

# ------------------------------------
# Concrete projection of a reach-set
# ------------------------------------

"""
    project(R::AbstractLazyReachSet, variables::NTuple{D, M};
            check_vars::Bool=true) where {D, M<:Integer}

Projects a reach-set onto the subspace spanned by the given variables.

### Input

- `R`          -- reach-set
- `vars`       -- tuple of variables for the projection
- `check_vars` -- (optional, default: `true`) if `true`, check that the given variable
                  indices `vars` are a subset of the variables of `R`

### Output

A `SparseReachSet` whose variable indices are given by `vars`.

The type of the new reach-set depends on the type of the reach-set `R`:

- If `R` contains a hyperrectangular set, the output is a hyperrectangle.
- If `R` contains a zonotopic set, the output is a zonotope.
- Otherwise, the return type is a polytope either in constraint representation
  or in vertex representation, depending on the dimension and the properties of `M`.
  For details, see `LazySets.project`.

### Notes

This function can be used to project a reach-set onto a lower-dimensional
sub-space. The projection is concrete, and it consists of mapping the reach-set
`X = set(R)` to a new reach-set through to `MX`, where `M` is the projection matrix
associated with the given variables `vars`.

To project onto the time variable, use the index `0`. For instance, `(0, 1)` projects
onto the time variable and the first variable in `R`.
"""
function project(R::AbstractLazyReachSet, variables::NTuple{D, M};
                 check_vars::Bool=true) where {D, M<:Integer}
     vR = vars(R)
     vRvec = collect(vR)

     if check_vars && !(setdiff(variables, 0) ⊆ vR)
         throw(ArgumentError("the variables $vars do not belong to the variables " *
                             " of this reach-set, $(vR)"))
     end

     if 0 ∈ variables  # the projection involves "time"
         vars_idx = _get_vars_idx(variables, vcat(0, vRvec))
         Δt = convert(Interval, tspan(R))
         proj =  project(Δt × set(R), vars_idx)
     else
         vars_idx = _get_vars_idx(variables, vRvec)
         proj = project(set(R), vars_idx)
     end

     return SparseReachSet(proj, tspan(R), variables)
end

# assumes that variables is a subset of all_variables
@inline function _get_vars_idx(variables, all_variables)
    vars_idx = indexin(variables, all_variables) |> Vector{Int}
end

# concrete projection onto a single variable
function project(R::AbstractLazyReachSet, variable::Int; check_vars::Bool=true)
    return project(R, (variable,), check_vars=check_vars)
end

# concrete projection overloads
project(R::AbstractLazyReachSet, vars::AbstractVector{M}) where {M<:Integer} = project(R, Tuple(vars))
project(R::AbstractLazyReachSet; vars) = project(R, Tuple(vars))

# projection of an array of reach-sets
_project_vec(R, vars) = map(Ri -> project(Ri, Tuple(vars)), R)
project(R::Vector{<:AbstractLazyReachSet}, vars::VecOrTupleOrInt) = _project_vec(R, vars)
project(R::Vector{<:AbstractLazyReachSet}; vars) = _project_vec(R, vars)
project(R::SubArray{<:AbstractLazyReachSet}, vars) = _project_vec(R, vars)
project(R::SubArray{<:AbstractLazyReachSet}; vars) = _project_vec(R, vars)

# -------------------------------
# Lazy projection of a reach-set
# -------------------------------

function Projection(R::AbstractLazyReachSet, variables::NTuple{D, M},
                    check_vars::Bool=true) where {D, M<:Integer}

    vR = vars(R)
    vRvec = collect(vR)

    # TODO: make vars check faster, specific for ReachSets and number of vars D
    if check_vars && !(setdiff(variables, 0) ⊆ vR)
        throw(ArgumentError("the variables $variables do not belong to the variables " *
                            " of this reach-set, $(vR)"))
    end

    if 0 ∈ variables  # the projection involves "time"
        vars_idx = _get_vars_idx(variables, vcat(0, vRvec))
        Δt = convert(Interval, tspan(R))
        proj =  Projection(Δt × set(R), vars_idx)
    else
        vars_idx = _get_vars_idx(variables, vRvec)
        proj = Projection(set(R), vars_idx)
    end

    return SparseReachSet(proj, tspan(R), variables)
end

# handle generic vars vector
function Projection(R::AbstractLazyReachSet, vars::AbstractVector{M}) where {M<:Integer}
    return Projection(R, Tuple(vars))
end

# handle generic kwargs vars
function Projection(R::AbstractLazyReachSet; vars)
    return Projection(R, Tuple(vars))
end

# ----------------------
# Additional extensions
# ----------------------

# membership test
function ∈(x::AbstractVector{N}, R::AbstractLazyReachSet{N}) where {N}
    return ∈(x, set(R))
end

# splitting a reach-set according to a given partition; the partition should be
# a vector of integers
function LazySets.split(R::AbstractLazyReachSet, partition)
    Y = split(set(R), partition)
    [reconstruct(R, y) for y in Y]
end


text:
encode(obj::String)="\"$(esc_str2(obj))\""
encode(obj::Bool)=obj ? "T" : "F"
encode(obj::Missing)="mis"
encode(obj::Nothing)="n"
encode(obj::Number)=string(obj)
encode(obj::Symbol)="&\"$(escape_string(obj))\""
encode(obj::Char)="(Char,$(Int(obj)))"
encode(obj::RGB)="(RGB_N0f8,$(obj.r.i),$(obj.g.i),$(obj.b.i))"
RGB_N0f8(r::Int,g::Int,b::Int)=RGB{N0f8}(reinterpret(N0f8,UInt8(r)),reinterpret(N0f8,UInt8(g)),reinterpret(N0f8,UInt8(b)))
function encode(obj::Vector)
	l=length(obj)
	if isempty(l)
		return "[]"
	end
	s="[$(encode(obj[1]))"
	for i in 2:l
		s*=",$(encode(obj[i]))"
	end
	return s*']'
end
function encode(obj::Dict)
	s="{"
	for i in obj
		s*="$(i.first):$(i.second),"
	end
	# s=s[1:end-1]
	return s*'}'
end
function encode(obj) # default
	ty=typeof(obj) # ty::DataType
	s="($(ty.name.name)" # ty.name::Core.TypeName
	li=fieldnames(ty)
	for i in li
		s*=",$(getfield(obj,i))"
	end
	return s*')'
end

function decode_get_elements(s::String)
	l=length(s);st=zero(UInt32);instr=false;a=[];last=1
	for i in 1:l
		c=s[i]
		if instr
			if c=='"'
				instr=false
			end
			continue
		end
		if c=='('||c=='['||c=='{'
			st+=1
		elseif c==')'||c==']'||c=='}'
			st-=1
		elseif c=='"'
			instr=true
		elseif c==','
			if st==0
				push!(a,decode(s[last:i-1]))
				last=i+1
			end
		end
	end
	if last!=l+1
		push!(a,decode(s[last:l]))
	end
	return a
end
function decode_get_pairs(s::String)
	l=length(s);st=zero(UInt32);instr=false;a=Dict();last=1;tkey=nothing
	for i in 1:l
		c=s[i]
		if instr
			if c=='"'
				instr=false
			end
			continue
		end
		if c=='('||c=='['||c=='{'
			st+=1
		elseif c==')'||c==']'||c=='}'
			st-=1
		elseif c=='"'
			instr=true
		elseif c==':'
			if st==0
				tkey=decode(s[last:i-1])
				last=i+1
			end
		elseif c==','
			if st==0
				a[tkey]=decode(s[last:i-1])
				last=i+1
			end
		end
	end
	if last!=l+1
		a[tkey]=decode(s[last:l])
	end
	return a
end
function decode_get_parts(s::String)
	l=length(s);st=zero(UInt32);instr=false;a=Vector{String}();last=1
	for i in 1:l
		c=s[i]
		if instr
			if c=='"'
				instr=false
			end
			continue
		end
		if c=='('||c=='['||c=='{'
			st+=1
		elseif c==')'||c==']'||c=='}'
			st-=1
		elseif c=='"'
			instr=true
		elseif c==','
			if st==0
				push!(a,s[last:i-1])
				last=i+1
			end
		end
	end
	if last!=l+1
		push!(a,s[last:l])
	end
	return a
end
function decode(s::String)
	if s=="T"
		return true
	elseif s=="F"
		return false
	elseif s=="mis"
		return missing
	elseif s=="n"
		return nothing
	elseif s=="Inf"
		return Inf
	elseif s=="-Inf"
		return -Inf
	elseif s=="NaN"
		return NaN
	end
	s=clean_str(s)
	c=s[1]
	if c=='"' return unesc_str2(s[2:end-1])
	elseif c=='&'
		if s[2]=='"' return Symbol(unesc_str2(s[3:end-1])) end
		return Symbol(s[2:end])
	elseif c=='['
		return decode_get_elements(s[2:end-1])
	elseif c=='{'
		return decode_get_pairs(s[2:end-1])
	elseif c=='('
		a=decode_get_parts(s[2:end-1])
		name=popfirst!(a)
		l=length(a)
		b=Vector{Any}(undef,l)
		for i in 1:l
			b[i]=decode(a[i])
		end
		f=eval(Symbol(name))
		return f(b...)
	elseif '0'<=c<='9'
		fl=true
		for i in 1:length(s)
			c=s[i]
			if c<'0'||c>'9'
				fl=false
				break
			end
		end
		if fl
			return parse(Int,s)
		else
			return parse(Float64,s)
		end
	end
end

function read_code(io::IO)
	s=""
	for i in eachline(io)
		s*=i
	end
	return decode(s)
end
function write_code(io::IO,obj)
	println(io,encode(obj))
end
function read_code(p::String)
	io=open(p,"r")
	s=""
	for i in eachline(io)
		s*=i
	end
	close(io)
	return decode(s)
end
function write_code(p::String)
	io=open(p,"w")
	println(io,encode(obj))
	close(io)
end


text:
using LightXML

# create an empty XML document
xdoc = XMLDocument()

# create & attach a root node
xroot = create_root(xdoc, "States")

# create the first child
xs1 = new_child(xroot, "State")

# add the inner content
add_text(xs1, "Massachusetts")

# set attribute
set_attribute(xs1, "tag", "MA")

# likewise for the second child
xs2 = new_child(xroot, "State")
add_text(xs2, "Illinois")
# set multiple attributes using a dict
set_attributes(xs2, Dict("tag"=>"IL", "cap"=>"Springfield"))

testdoc = XMLDocument()
g = create_root(testdoc, "graphml")
c1 = new_child(g, "key")
set_attribute(c1, "attr.name", "FA_std")
set_attribute(c1, "attr.type", "double")
set_attribute(c1, "for", "edge")
set_attribute(c1, "id", "d13")

function add_keys!(root)
    for i in 1:13
        c = new_child(root, "key")
        set_attribute(c, "attr.name", name_dict[i])
        set_attribute(c, "attr.type", type_dict[i])
        set_attribute(c, "for", for_dict[i])
        set_attribute(c, "id", "d$i")
    end
end

for_dict = Dict(1 => "node",
                2 => "node",
                3 => "node",
                4 => "node",
                5 => "node",
                6 => "node",
                7 => "node",
                8 => "node",
                9 => "edge",
                10 => "edge",
                11 => "edge",
                12 => "edge",
                13 => "edge"
)

type_dict = Dict(1 => "double",
                 2 => "double",
                 3 => "double",
                 4 => "string",
                 5 => "string",
                 6 => "string",
                 7 => "string",
                 8 => "string",
                 9 => "int",
                 10 => "double",
                 11 => "double",
                 12 => "double",
                 13 => "double"
)

name_dict = Dict(1 => "dn_position_x",
                 2 => "dn_position_y",
                 3 => "dn_position_z",
                 4 => "dn_correspondence_id",
                 5 => "dn_region",
                 6 => "fn_fsname",
                 7 => "dn_name",
                 8 => "dn_hemisphere",
                 9 => "number_of_fibers",
                 10 => "FA_mean",
                 11 => "fiber_length_std",
                 12 => "fiber_length_mean",
                 13 => "FA_std"
)

function make_xml()
    xdoc = XMLDocument()
    root = create_root(xdoc, "graphml")
    add_keys!(root)

    xdoc
end


text:
# %% Surrogate application (From Lancaster)
using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot

using TimeseriesSurrogates, Statistics, Random

ro = Systems.roessler(ones(3); a = 0.165, b = 0.2, c = 10.0)

tr = trajectory(ro, 500; Δt = 0.1, Ttr = 500)
x = tr[:, 1]
x ./= std(x)
x .+= randn(length(x))*std(x)*0.1

# generate autoregressive process
Random.seed!(77163)
η = randn(5000)
s = ones(5000)
for n in 4:5000
    s[n] = 1.625s[n-1] - 0.284s[n-2] - 0.355s[n-3] + η[n] - 0.96η[n-1]
end
s ./= std(s)

# Do the surrogate calculation
εro, εma = std.((x, s))./4
algs = [RandomFourier(), AAFT()]
names = ["FT", "AAFT"]
sgx = [surrogenerator(x, m) for m in algs]
sgs = [surrogenerator(s, m) for m in algs]
τx = estimate_delay(x, "ac_zero")
τs = estimate_delay(s, "ac_zero")
Cx = grassberger(embed(x, 3, τx))
Cs = grassberger(embed(s, 4, τs))
A = length(algs)

xboxes = []
sboxes = []
for i in 1:A
    sx, ss = sgx[i], sgs[i]
    bx, bs = [], []
    for j in 1:100
        X = embed(sx(), 3, τx)
        S = embed(ss(), 4, τs)
        push!(bx, grassberger(X))
        push!(bs, grassberger(S))
    end
    push!(xboxes, bx)
    push!(sboxes, bs)
end

# %% Plot everything
fig = figure()
ax1 = fig.add_subplot(1,2,1)
ax2 = fig.add_subplot(1,4,3)
ax3 = fig.add_subplot(1,4,4)
axs = [ax1, ax2, ax3]

axs[1].plot(x .+ 2,  label = "Rössler", lw = 1.5)
axs[1].plot(s .- 2,label = "ARMA", lw = 1.5, color = "C2")
axs[1].set_xlim(0, 1000)
axs[1].set_ylim(-5, 5)
axs[1].set_yticks(-4:2:4)
leg = axs[1].legend(bbox_to_anchor=(0., 1.02, 1., .102), loc="lower left",
           ncol=2, mode="expand", borderaxespad=0., fontsize = 26)
for line in leg.get_lines()
    line.set_linewidth(4.0)
end

for (j, boxes) in enumerate((xboxes, sboxes))
    c = j == 1 ? "C0" : "C2"
    ax = axs[j+1]
    for (i, b) in enumerate(boxes)
        ax.boxplot([b]; positions = [i],
        patch_artist=true, boxprops=Dict(:facecolor=>c, :color=>c),
        medianprops=Dict(:color=>"w"), flierprops=Dict(:markeredgecolor=>c))
    end
end

l = [1-0.2, A+0.2]
axs[2].plot(l, fill(Cx, 2), color = "C0", ls = "dashed")
# axs[2].plot(l, fill(2.75, 2), color = "C0", ls = "dashed")
axs[3].plot(l, fill(Cs, 2), color = "C2", ls = "dashed")
# axs[2].set_ylim(2.7, 3)
axs[2].set_yticks(2.6:0.2:3.0)
axs[2].set_title("\$\\Delta^{(C)}\$, Rössler")
axs[3].set_title("\$\\Delta^{(C)}\$, ARMA")
axs[3].set_yticks(3.4:0.3:4.1)

for ax in axs[2:3]
    ax.grid(false; axis = "x")
    ax.set_xticks(1:A)
    ax.set_xticklabels(names, rotation = 0, size = 20)
end

fig.tight_layout(pad=0.3)
fig.subplots_adjust(top = 0.85, left = 0.05, bottom = 0.1, right = 0.98, wspace = 0.3)
wsave(plotsdir("7", "surrogates"), fig)

max_stars_repo_name: JuliaDynamics/NonlinearDynamicsTextbook

max_issues_repo_name: JuliaDynamics/NonlinearDynamicsTextbook

max_forks_repo_name: JuliaDynamics/NonlinearDynamicsTextbook