text:
mutable struct HMM
    dict::AhoCorasickAutomaton{Ti}
    words::Vector{String}
    usr_words::Int
    tags::Vector{String}
    hpr::Vector{Tv}
    h2h::Matrix{Tv}
    h2v::Vector{Dict{Int, Tv}}
    INF::Vector{Tv}
end

function HMM(corpus, dictpath, major)
        ispath(dictpath) || @warn "dictpath is not a valid path"
        @assert 0. < major <= 1.

        tags, words = KongYiji.postags(corpus), KongYiji.words(corpus)
        np = length(tags)

        hpr, h2h, h2v, INF = fill(Tv(0), np), fill(Tv(0), (np, np)), [Dict{Int, Tv}() for _ in 1:np], fill(Tv(0), np)
        h2v2 = [Dict{Int, Tv}() for _ in 1:np]

        for sent in sents(corpus)
                pp = 0
                for (ip, iw) in sent
                        if pp == 0 hpr[ip] += 1 else h2h[pp,ip] += 1 end
                        pp = ip
                        h2v[ip][iw] = get(h2v[ip], iw, 0) + 1
                end
        end

        if major < 1.
                  h2v2 = [Dict{Int, Tv}() for _ in 1:np]
                  for ip = 1:np
                          pct, tot = 0., sum(values(h2v[ip]))
                          for (iw, cnt) in sort!(collect(h2v[ip]); by=last, rev=true)
                                  h2v2[ip][iw] = cnt
                                  pct += cnt / tot
                                  if pct >= major; break; end
                          end
                  end
                  h2v = h2v2
        end

        for ip = 1:np
                if !startswith(tags[ip], "N"); continue; end
                for (iw, cnt) in h2v[ip]; h2v[ip][iw] = 1; end
        end

        wmp = KongYiji.dict(words)

        if isdir(dictpath)
                cd(() -> begin
                                for tag in tags
                                        if !ispath(tag); continue; end
                                        pmp = KongYiji.dict(tags)
                                        ip = get(pmp, tag, 0)
                                        if ip == 0; continue; end
                                        maxcnt = maximum(values(h2v[ip]))
                                        for file in readdir(tag)
                                                for word in eachline(joinpath(tag, file))
                                                        word = mask(word) #todo forgotten
                                                        iw = getid!(wmp, word)
                                                        if ip > 0
                                                                if startswith(tag, "N"); h2v[ip][iw] = 1;
                                                                else h2v[ip][iw] = maxip;
                                                                end
                                                        end
                                                end
                                        end
                                end
                        end, dictpath)
        end

        dict = AhoCorasickAutomaton{Ti}(wmp)
        usr_words = length(wmp) - length(words)
        words = KongYiji.vec(wmp)
        normalize!.((hpr, h2h))
        normalize!(h2v, INF)
        log!.((hpr, h2h, h2v, INF))

        HMM(dict, words, usr_words, tags, hpr, h2h, h2v, INF)
end

HMM() = load(KongYiji.dir("hmm.jld2"))["m"]

function (m::HMM)(xs::Vector{String}; recover=true, withpos=false)
        nc_max = mapreduce(ncodeunits, max, xs)
        np = length(m.hpr)
        @assert np > 0
        dp = fill(Tv(-Inf), (nc_max + 1, np))
        pr = fill((1, 0), (nc_max + 1, np))
        return [m(x, dp, pr, recover, withpos) for x in xs]
end

function (m::HMM)(x::String; recover=true, withpos=false)
        return m([x]; recover=recover, withpos=withpos)[1]
end

function (m::HMM)(x::String, dp::Matrix{Tv}, pr::Matrix{Tuple{Int, Int}}, recover, withpos)::Vector{String}
        origin = x
        x = mask(x)
        chrs = codeunits(x) #todo slow?
        vtxs = collect(eachmatch(m.dict, chrs))
        sort!(vtxs)
        nv, nc, np = length.((vtxs, chrs, m.hpr))
        for i = 1:nc + 1, j in 1:np dp[i,j] = -Inf end
        pv = 1
        dp[1, :] = m.hpr
        pre_i = 1
        for i in 1:nc + 1
                if dp[i,1] != -Inf pre_i = i end
                while pv <= nv && vtxs[pv].s < i pv = pv + 1 end
                if !(pv <= nv && vtxs[pv].s == i || i == nc + 1) continue end
                if dp[i, 1] == -Inf
                        #@show i, pre_i
                        for pi = 1:np, pj = 1:np
                                maybe = dp[pre_i, pj] + m.h2h[pj, pi] + logph2v(m, pj, 0)
                                if maybe > dp[i, pi] dp[i, pi] = maybe; pr[i, pi] = (pre_i, pj) end
                        end
                        #pre_i = i
                end
                while pv <= nv && vtxs[pv].s == i
                        vtx = vtxs[pv]
                        j = i + length(vtx)
                        for pi = 1:np
                                for pj = 1:np
                                        maybe = dp[i, pi] + m.h2h[pi, pj] + logph2v(m, pi, vtx.i)
                                        if maybe > dp[j,pj] dp[j,pj] = maybe; pr[j,pj] = (i, pi) end
                                end
                        end
                        pv = pv + 1
                end
        end
        #@show m
        #@show maximum(dp[nc + 1, :])
        #@show dp'
        #@show pr'
        #@show vtxs
        v = (nc + 1, argmax(dp[nc + 1,:]))
        words = fill("", 0)
        while v[1] != 1
                pv = pr[v[1],v[2]]
                word = x[pv[1]:prevind(x, v[1])]
                push!(words, word)
                v = pv
        end
        reverse!(words)
        if recover; words = demask(words, origin); end
        ret = words

        if withpos
                v = (nc + 1, argmax(dp[nc + 1,:]))
                postags = fill("", 0)
                while v[1] != 1
                        pv = pr[v[1],v[2]]
                        postag = m.tags[pv[2]]
                        push!(postags, postag)
                        v = pv
                end
                reverse!(postags)
                ret = [join(p, "/") for p in zip(words, postags)]
        end

        return ret
end

logph2v(m::HMM, p, v) = haskey(m.h2v[p], v) ? m.h2v[p][v] : m.INF[p]

function ==(a::HMM, b::HMM)
        return all(fname -> getfield(a, fname) == getfield(b, fname), fieldnames(HMM))
end

# tagger
function (m::HMM)(xs::Vector{Vector{T}}; recover=true, withword=true) where {T}
	maxnw = mapreduce(length, max, xs)
	np = length(m.tags)
	dp = fill(-oo, (np, maxnw + 1))
	pr = fill((0, 1), (np, maxnw + 1))
	[m(x, dp, pr, recover, withword) for x in xs]
end

function (m::HMM)(x::Vector{T}; recover=true, withword=false) where {T}
        return m([x]; recover=recover, withword=withword)[1]
end

function (m::HMM)(x::Vector{T}, dp::Matrix{Tv}, pr::Matrix{Tuple{Int, Int}}, recover, withword)::Vector{String} where {T}
	nw, np = length.((x, m.hpr))
        for i = 1:nw + 1, j in 1:np dp[j, i] = -oo end
        dp[:, 1] = m.hpr

        for i = 1:nw
        	iw = get(m.dict, mask(x[i]), Ti(0))
        	j = i + 1
        	for ip = 1:np, jp = 1:np
        		maybe = dp[ip, i] + get(m.h2v[ip], iw, m.INF[ip]) + m.h2h[ip, jp]
        		if dp[jp, j] < maybe; dp[jp, j] = maybe; pr[jp, j] = (ip, i); end
		end
	end

	v = (argmax(dp[:, nw + 1]), nw + 1)
	r = fill("", 0)
        while v[2] != 1
                pv = pr[v[1],v[2]]
                postag = m.tags[pv[1]]
                word = x[pv[2]]
                if !recover; word = mask(word); end
                push!(r, string(word, '/', postag))
                v = pv
        end
        reverse!(r)
        r
end


text:
####
#### Defines list of tutorials given `generated_dir`
####

generate_tutorials = true

tutorials = Any[]

# Allow flag to skip generated
# tutorials since this is by
# far the slowest part of the
# docs build.
if generate_tutorials

    tutorials_dir = joinpath(@__DIR__, "..", "tutorials")      # julia src files

    non_tutorial_files =
        ["KinematicModel.jl", "clean_build_folder.jl", "plothelpers.jl"]
    skip_execute = [
        "heldsuarez.jl",                 # broken
        "risingbubble.jl",               # broken
        "topo.jl",                       # broken
        "dry_rayleigh_benard.jl",        # takes too long
        "nonnegative.jl",                # takes too long
        "ex_2_Kessler.jl",               # takes too long
        "ex_1_saturation_adjustment.jl", # takes too long
    ]

    # generate tutorials
    import Literate

    tutorials_jl = [
        joinpath(root, f)
        for (root, dirs, files) in Base.Filesystem.walkdir(tutorials_dir)
        for f in files
    ]
    filter!(x -> endswith(x, ".jl"), tutorials_jl) # only grab .jl files
    filter!(x -> !any(occursin.(non_tutorial_files, Ref(x))), tutorials_jl)

    println("Building literate tutorials:")
    for tutorial in tutorials_jl
        println("    $(tutorial)")
    end

    for tutorial in tutorials_jl
        gen_dir =
            joinpath(generated_dir, relpath(dirname(tutorial), tutorials_dir))
        input = abspath(tutorial)
        script = Literate.script(input, gen_dir)
        code = strip(read(script, String))
        mdpost(str) = replace(str, "@__CODE__" => code)
        Literate.markdown(input, gen_dir, postprocess = mdpost)
        if !any(occursin.(skip_execute, Ref(input)))
            Literate.notebook(input, gen_dir, execute = true)
        end
    end

    # TODO: Should we use AutoPages.jl?

    # These files mirror the .jl files in `ClimateMachine.jl/tutorials/`:
    tutorials = Any[
        "Atmos" => Any[
            "Dry Idealized GCM" => "generated/Atmos/heldsuarez.md",
            "Rising Thermal Bubble" => "generated/Atmos/risingbubble.md",
        ],
        "Ocean" => Any[],
        "Land" => Any["Heat" => Any["Heat Equation" => "generated/Land/Heat/heat_equation.md"],],
        "Numerics" => Any["LinearSolvers" => Any[
            "Conjugate Gradient" => "generated/Numerics/LinearSolvers/cg.md",
            "Batched Generalized Minimal Residual" => "generated/Numerics/LinearSolvers/bgmres.md",
        ],],
        "Contributing" => Any["Notes on Literate" => "generated/literate_markdown.md",],
    ]

end

ext: jl


text:
using MISAlgorithms

# mask-copy version, 1.5 second
using Random
Random.seed!(2)
eg = rand_eg(60, 0.1)
using BenchmarkTools
#@benchmark neighborcover($eg, 2) seconds=1
@time mis1(eg)
@code_warntype mis1(eg)
@benchmark mis1($eg) seconds=1

using Profile
Profile.clear()
Profile.init(delay=1e-3)
@profile neighborcover(eg, 2)

using Random
Random.seed!(2)
eg = rand_eg(120, 0.1)
@time mis2(eg)

using Profile
eg = rand_eg(80, 0.1)
Profile.clear()
Profile.init(delay=1e-3)
@profile for i=1:30 mis2(eg) end


text:
using IntervalWavelets
using Test

@testset "Scaling function basis" begin
    @testset "Non-default reconstruction intervals" begin
        J = 3
        R = 3
        l = DyadicRational(-1, 1)
        r = DyadicRational(1, 1)
        B = interval_scaling_function_basis(2, J, R; left_boundary = l, right_boundary = r)

        x = reconstruct(B, ones(2^J))[1]

        # This only works because the interval has length 1
        @test x == all_dyadic_rationals(l, r, R)
    end

    @testset "Support of basis function" begin
        l = DyadicRational(-1, 1)
        r = DyadicRational(1, 0)
        B = interval_scaling_function_basis(2, 3, 4; left_boundary = l, right_boundary = r)

        # Left boundary functions
        @test B[1][1][1] == l
        @test B[2][1][1] == l

        # Right boundary functions
        @test B[7][1][end] == r
        @test B[8][1][end] == r
    end

    @testset "Reconstruct from unit vectors" begin
        p = 2
        J = 3
        R = 4
        B = interval_scaling_function_basis(p, J, R)

        x = all_dyadic_rationals(B.left_boundary, B.right_boundary, R)
        L = boundary_scaling_functions(LEFT, p, R)
        for k in 1:2
            coef = zeros(2^J)
            coef[k] = 1

            y = 2.0^(J/2) * L[k - 1].(2^J * x)
            xx, z = reconstruct(B, coef)

            @test x == xx
            @test y ≈ z atol = 1e-8
        end

        phi = interior_scaling_function(interior_filter(p), R)
        for k in 3:6
            coef = zeros(2^J)
            coef[k] = 1

            translation = k - 1
            y = 2.0^(J/2) * phi.(2^J * x .- translation)
            xx, z = reconstruct(B, coef)

            @test x == xx
            @test y ≈ z atol = 1e-8
        end

        right = boundary_scaling_functions(RIGHT, p, R)
        for k in 7:8
            coef = zeros(2^J)
            coef[k] = 1

            y = 2.0^(J/2) * right[8 - k].(2^J * (x .- B.right_boundary))
            xx, z = reconstruct(B, coef)

            @test x == xx
            @test y ≈ z atol = 1e-8
        end
    end

    @testset "Reconstruct constant function 1D" begin
        B = interval_scaling_function_basis(2, 3, 4)

        # Coefficients found by solving [ B[j](x_i) ] * coef = ones()
        coef = [
            0.128004734,
            0.354064418,
            0.353553390,
            0.353553390,
            0.353553390,
            0.353553390,
            0.385317706,
            0.458021222
        ]

        y = reconstruct(B, coef)[2]

        @test y ≈ ones(17) atol = 1e-8
    end

    @testset "Reconstruct constant function 2D" begin
        B = interval_scaling_function_basis(2, 3, 4)

        # Coefficients reused from 1D
        coef1 = [
            0.128004734,
            0.354064418,
            0.353553390,
            0.353553390,
            0.353553390,
            0.353553390,
            0.385317706,
            0.458021222
        ]
        coef = coef1 * coef1'

        y = reconstruct(B, coef)

        @test y ≈ ones(17, 17) atol = 1e-7
    end
end


text:
export Engine

module Engine

include("sdl.jl")
include("io.jl")

export initialize, run

using Printf

using ..Nodes
using ..Ranger
using ..Rendering
using ..Events
using ..Geometry

manager = nothing

# billion ns in a second
const SECOND = 1000000000
# Maximum updates per second
const UPDATES_PER_SECOND = 30

const UPDATE_PERIOD = 1000000000.0 / Float64(UPDATES_PER_SECOND)

function initialize(title::String, build::Function)
    println("Initializing Ranger engine...")
    ranger_path = joinpath(@__DIR__, "../")
    println("Ranger path: ", ranger_path)

    println("Loading fonts...")
    raster_font = Rendering.RasterFont()
    loaded = Rendering.load_font!(raster_font, ranger_path * "assets/raster_font.data")

    if !loaded
        println("Engine failed to initialize")
        return nothing
    end

    vector_font = Rendering.VectorFont()
    loaded = Rendering.load_font!(vector_font, ranger_path * "assets/vector_font.data")

    if !loaded
        println("Engine failed to initialize")
        return nothing
    end

    println("Fonts loaded.")

    world = Ranger.World(title)

    SDL(world)

    global manager = Nodes.NodeManager(world, raster_font, vector_font)

    # Allow client to build game
    build(world)

    world
end

function run(world::Ranger.World)
    println("running...")

    running = true

    ns_per_update = UInt64(round(UPDATE_PERIOD))
    # println("ns_per_update: ", ns_per_update, "ns")

    frame_dt = Float64(ns_per_update) / 1000000.0
    println("Frame period: ", frame_dt, "ms")

    lag = 0

    # ***************************
    # Debugging only
    # ***************************
    fps_cnt = fps = 0
    ups_cnt = ups = 0
    previous_t = time_ns()
    second_cnt = 0
    avg_render = 0.0
    render_elapsed_cnt = 0
    display_coords = true
    display_on = true
    step_enabled = false

    keyboard = Events.KeyboardEvent()
    mouse = Events.MouseMotionEvent()
    mouse_wheel = Events.MouseWheelEvent()
    mouse_button = Events.MouseButtonEvent()

    # Main game loop
    while running
        current_t = time_ns()

        # ~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--
        # Handle Events
        # ~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--
        haveEvents = true

        # Process any events that have been queued.
        while haveEvents
            event, haveEvents = Events.poll_event!()

            if haveEvents
                ev_type = Events.get_event_type(event)

                if (ev_type == SDL2.KEYDOWN || ev_type == SDL2.KEYUP)
                    keySym = Events.get_key_code_sym(event)

                    if (keySym == SDL2.SDLK_ESCAPE)
                        running = false
                        continue
                    end

                    Events.set!(keyboard, event)
                    # Events.print_keyboard_event(event)

                    if keyboard.keycode == SDL2.SDLK_m
                        display_coords = !display_coords
                    end

                    if keyboard.keycode == SDL2.SDLK_n
                        display_on = !display_on
                    end

                    # Route event to registered Nodes
                    Nodes.route_events(manager, keyboard)
                elseif ev_type == SDL2.MOUSEMOTION
                    Events.set!(mouse, event)

                    Nodes.route_events(manager, mouse)
                elseif ev_type == SDL2.MOUSEWHEEL
                    Events.set!(mouse_wheel, event)

                    Nodes.route_events(manager, mouse_wheel)
                elseif ev_type == SDL2.MOUSEBUTTONDOWN || ev_type == SDL2.MOUSEBUTTONUP
                    Events.set!(mouse_button, event)

                    Nodes.route_events(manager, mouse_button)
                end
            end
        end

        # ~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--
        # Update
        # ~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--
        elapsed_t = current_t - previous_t
        previous_t = current_t

        # Note: This update loop is based on:
        # https://gameprogrammingpatterns.com/game-loop.html

        if !step_enabled
            lag += elapsed_t
            lagging = true
            while lagging
                if lag >= ns_per_update
                    Nodes.update(manager, frame_dt)
                    lag -= ns_per_update
                    ups_cnt += 1
                else
                    lagging = false
                end
            end
        end

        # sleep(0.01666)

        # ~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--
        # Render
        # ~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--
        Nodes.pre_visit(manager)
        # **** Any rendering must occur AFTER this point ****

        # Capture time AFTER pre_visit otherwise if vsync is enabled
        # then time includes the vertical refresh which ~16.667ms
        render_t = time_ns()

        interpolation = Float64(lag) / Float64(ns_per_update)

        more_scenes = Nodes.visit(manager, interpolation)

        if !more_scenes
            running = false
            continue
        end

        second_cnt += elapsed_t

        # ***************************
        # Debugging only
        # ***************************
        if display_on
            if display_coords
                draw_space_coords(mouse.x, mouse.y, world)
            end
            draw_stats(fps, ups, avg_render, world)
        end

        render_elapsed_t =  time_ns() - render_t
        render_elapsed_cnt += render_elapsed_t

        if second_cnt >= SECOND
            avg_render = (Float64(render_elapsed_cnt) /  Float64(fps_cnt)) / 1000000.0
            fps = fps_cnt
            ups = ups_cnt
            # text = @sprintf("fps (%2d), ups(%2d) %2.4f", fps, ups, avg_render)
            # println(text)
            # println("fps (", fps, ") ups (", ups, "), ren ", avg_render)

            fps_cnt = 0
            ups_cnt = 0
            second_cnt = 0
            render_elapsed_cnt = 0
        end

        fps_cnt += 1
        # ***************************

        # ~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--
        # Present
        # ~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--~--
        Nodes.post_visit(manager)

    end

    exit();
end

function push(scene::Ranger.AbstractScene)
    Nodes.push_node(manager, scene)
end

function exit()
    SDL2.Quit()
    println("Game exited");
end

# ***************************
# Debugging only
# ***************************
orange = Rendering.Orange()
white = Rendering.White()
green = Rendering.Lime()
view_point = Geometry.Point{Float64}()

function draw_stats(fps::Integer, ups::Integer, avg_render::Float64, world::Ranger.World)
    Rendering.set_draw_color(manager.context, white)
    text = @sprintf("(%2d)(%2d)(%2.4f)", fps, ups, avg_render)
    x = 10
    y = world.window_height - 24
    Rendering.draw_text(manager.context, x, y, text, 2, 2, false)
end

function draw_space_coords(mx::Int32, my::Int32, world::Ranger.World)
    Rendering.set_draw_color(manager.context, green)
    text = @sprintf("M: %2d, %2d", mx, my)
    x = 10
    y = 30
    Rendering.draw_text(manager.context, x, y, text, 2, 2, false)

    Nodes.map_device_to_view(world, mx, my, view_point)
    text = @sprintf("V: %2d, %2d", view_point.x, view_point.y)
    y += 20
    Rendering.draw_text(manager.context, x, y, text, 2, 2, false)
end

end # Module -----------------------------------------------------------

max_stars_repo_name: wdevore/Ranger-Julia-SDL

max_issues_repo_name: wdevore/Ranger-Julia-SDL

max_forks_repo_name: wdevore/Ranger-Julia-SDL


text:
using BPM
using Test

@testset "Rosiere Validation" begin
##################################################################################
##################################################################################
##################################################################################

# Rosiere Valdiation (243.84 m, 800 ft should be 47 dB)
# http://www.mge.com/environment/green-power/wind/rosiere.htm
x_test = [0.] # x-locations of turbines (m)
y_test = [0.] # y-locations of turbines (m)
obs_test = [0., 243.84, 0.] # x-, y-, and z-location of the observer (m)
winddir_test = 0. # wind direction (deg)
rpm_test = [28.5] # rotation rate of the tubrines (rpm)
windvel_test = [15.] # wind velocity (m/s)
B_test = 3 # number of blades
h_test = 25. # height of the turbine hub (m)
noise_corr = 0.8697933840957954 # correction factor for noise

rad =  [1.069324603174603, 2.088888888888889, 3.1084531746031745, 4.382936507936508, 5.912301587301587, 7.441666666666666, 8.971031746031747, 10.500396825396825, 12.029761904761905, 13.559126984126985, 15.088492063492065, 16.617857142857144, 18.147222222222222, 19.6765873015873, 20.951070634920637, 21.97063492063492, 22.990199206349207, 23.5] # radial positions (m)
c =  [2.8941253867777776, 3.1490568155396828, 3.404805332214286, 3.7234696181666673, 3.8010929698730163, 3.6425779148095243, 3.4718065410555554, 3.2740712661825397, 3.062445496793651, 2.861441870269841, 2.660438243746032, 2.459434617222222, 2.258430990698413, 2.057427364174603, 1.889924342071429, 1.7044453858888888, 1.1594477481190477] # chord lengths (m)
alpha =  [13.308, 13.308, 13.308, 13.308, 11.48, 10.162, 9.011, 7.795, 6.544, 5.361, 4.188, 3.125, 2.319, 1.526, 0.863, 0.37, 0.106] # angles of attack (deg)
c1 = c*0.25 # pitch axis (m)
AR = 17. # blade aspect ratio

nu = 1.78e-5 # kinematic viscosity (m^2/s)
c0 = 343.2 # speed of sound (m/s)
psi = 14.0 # solid angle (deg)
# f    (nturb,  nseg,     nobs,       x,            y,          obs,    winddir,windvel,rpm, B, Hub, rad,  c,  c1,alpha,nu,    c0,      psi, AR, noise_corr)
db_test_ros, db_test_ros_splf, db_test_ros_splfA = BPM.turbinepos(x_test, y_test, obs_test, winddir_test, windvel_test, rpm_test, B_test, h_test, rad, c, c1, alpha, nu, c0, psi, AR, noise_corr)

println("Test Cases:")
println("Rosiere Validation (47 dB): $db_test_ros")

@test isapprox(db_test_ros[1], 47.0; atol=1e-6)

end #Rosiere Test

@testset "Eric's Example" begin
##################################################################################
##################################################################################
##################################################################################

# SPL Test Paremeters (changed to whatever desired)
x_test = [0.,5.] # x-locations of turbines (m)
y_test = [0.,5.] # y-locations of turbines (m)
obs_test = [0., 200., 0.] # x-, y-, and z-location of the observer (m)
winddir_test = 0. # wind direction (deg)
rpm_test = [28.5,28.5] # rotation rate of the tubrines (rpm)
windvel_test = [10.0,10.0] # wind velocity (m/s)
B_test = 3 # number of blades
h_test = 25. # height of the turbine hub (m)
noise_corr = 0.8697933840957954 # correction factor for noise

rad =  [1.069324603174603, 2.088888888888889, 3.1084531746031745, 4.382936507936508, 5.912301587301587, 7.441666666666666, 8.971031746031747, 10.500396825396825, 12.029761904761905, 13.559126984126985, 15.088492063492065, 16.617857142857144, 18.147222222222222, 19.6765873015873, 20.951070634920637, 21.97063492063492, 22.990199206349207, 23.5] # radial positions (m)
c =  [2.8941253867777776, 3.1490568155396828, 3.404805332214286, 3.7234696181666673, 3.8010929698730163, 3.6425779148095243, 3.4718065410555554, 3.2740712661825397, 3.062445496793651, 2.861441870269841, 2.660438243746032, 2.459434617222222, 2.258430990698413, 2.057427364174603, 1.889924342071429, 1.7044453858888888, 1.1594477481190477] # chord lengths (m)
alpha =  [13.308, 13.308, 13.308, 13.308, 11.48, 10.162, 9.011, 7.795, 6.544, 5.361, 4.188, 3.125, 2.319, 1.526, 0.863, 0.37, 0.106] # angles of attack (deg)
c1 = c*0.25 # pitch axis (m)
AR = 17. # blade aspect ratio

nu = 1.78e-5 # kinematic viscosity (m^2/s)
c0 = 343.2 # speed of sound (m/s)
psi = 14.0 # solid angle (deg)

db_test, db_test_splf, db_test_splfA = BPM.turbinepos(x_test, y_test, obs_test, winddir_test, windvel_test, rpm_test, B_test, h_test, rad, c, c1, alpha, nu, c0, psi, AR, noise_corr)

println("Test SPL (50.8446366094): $db_test")
@test isapprox(db_test[1], 50.8446366094; atol=1e-6)

end #Eric's other example

@testset "Mariah Windspire 5 MW" begin
##################################################################################
##################################################################################
##################################################################################

div=5

nu = 1.78e-5 # kinematic viscosity (m^2/s)
c0 = 343.2 # speed of sound (m/s)
psi = 14.0 # solid angle (deg)

turbx = [0.] # turbine x-positions (m)
turby = [0.] # turbine y-positions (m)
Vinf = 8. # free stream wind speed (m/s)
dia = 1.2 # turbine diameter (m)
rad = dia/2. # turbine radius (m)
tsrd = 2.625 # tip-speed ratio

rot = [1]*tsrd*Vinf/rad # turbine rotation rates (rad/s)

winddir = 180. # wind direction (deg)

# Aerodynamic properties
twist = 0.0
delta = 0.0

B = 3 # number of blades
chord = 0.128 # chord length (m)
c = ones(div)*chord # chord lengths over height positions (m)
c1 = c*0.5 # pitch axis location (m)
alpha = ones(div)*0.0 # angles of attack (deg)
Hub = 2. # hub height (m)
H = 6.1 # blade height (m)
high = range(0,H, length=div+1) # height positions of the blade (m)

AR = 5. # aspect ratio

# Wake velocities from surronding turbines
ntheta = 4
wakex = zeros(ntheta)
wakey = zeros(ntheta)

noise_corr = 1.0 # correction factor for noise

#Test Point
X=5
Y=5
Z=0
OASPL = BPM.turbinepos_VAWT(ntheta, turbx, turby, [X,Y,Z], winddir, B, Hub, high, rad, c, c1, alpha, nu, c0, psi, AR, noise_corr, rot, Vinf, wakex, wakey)

println("Test OASPL (70.263597969): $OASPL")
@test isapprox(OASPL, 70.263597969; atol=1e-6)

end


text:
" Implements a Diffusion (Laplacian) operator "
type Diffusion <: Kernel
    u::Variable
end

@inline function computeQpResidual(kernel::Diffusion, qp::Integer, i::Integer)
    u = kernel.u

    return u.grad[qp] ⋅ u.grad_phi[qp][i]
end

@inline function computeQpJacobian(kernel::Diffusion, v::Variable, qp::Integer, i::Integer, j::Integer)::Float64
    u = kernel.u

    if u.id == v.id
        return v.grad_phi[qp][j] ⋅ u.grad_phi[qp][i]
    end

    return 0
end

@inline function coupledVars(kernel::Diffusion)::Array{Variable}
    return [kernel.u]
end


text:
using NamedArrays

taskNr = 3

####### Code used by several tasks #######
function trafoCSV(filename)
    # transforms a given CSV file with 3 columns into a matrix and the into
    # a named array.
    # The first two columns of the CSV file need to be categorical
    # The third column needs to be numerical

    open(filename) do myFile
        lines = readlines(myFile)

        # Header information
        header = split(lines[1],",");
        rowName, colName = header[1:2];

        colNr = length(header)
        rowNr = length(lines) - 1

        # Empty data matrix
        data = Array{String}(undef, rowNr, colNr)

        # Split and put each line into the data matrix:
        for (idx, line) in enumerate(lines[2:end])
            data[idx, 1:colNr] = split(line,",");
        end

        # Get the categorical values:
        hairCols = unique(data[:,1]);
        eyeCols = unique(data[:,2]);
        # Alternatively:
    #        hairCols = Set(data[:,1]);
    #        eyeCols = Set(data[:,2]);

        # Make a named array:
        namedData = NamedArray(zeros(length(eyeCols),length(hairCols)),
                            (eyeCols,hairCols), (colName, rowName))

        # Populate it with absolute numbers:
        for i in 1:length(data[:,1])
            namedData[data[i,2],data[i,1]] = parse(Float64,data[i,3]);
        end

        # Tranfsorm to relative proportions:
        namedData = namedData ./ sum(namedData);

        return namedData
    end
end
###################################

if taskNr == 1
    thisFile = "HairEyeColor.csv";
    println(trafoCSV(thisFile))

elseif taskNr == 2
    thisFile = "HairEyeColor.csv";
    M = trafoCSV(thisFile);

    # a)
    println("\nMarginal of the hair colors:")
    println(sum(M, dims=1))
    # b)
    println("\nMarginal of the eye colors:")
    println(sum(M, dims=2))
    # c )
    println("\nTotal sum:")
    println(sum(M, dims=:))

elseif taskNr == 3
    thisFile = "HairEyeColor.csv";
    M = trafoCSV(thisFile);

    # a)
    print("\nP(Blue eyes, Blond hair) = ")
    println(M["Blue","Blond"])

    # b)
    print("\nP(Brown eyes) = ")
    println(sum(M["Brown",:]))

    # c)
    # Conditional prob -> reallocate on the summed "universe"
    print("\nP(Brown eyes | Red hair) = ")
    println(M["Brown","Red"] / sum(M["Brown",:]))

    # Allow several categories to survive by using a vector as input
    # d) Make intersection of subspaces
    print("\nP((Red OR Blond) AND (Brown OR Blue) = ")
    overlap = sum(M[["Brown","Blue"],["Red","Blond"]]);
    println(overlap)

    # e) Sum the two subspaces (remove the intersection to avoid double counting)
    print("\nP((Red OR Blond) OR (Brown OR Blue) = ")
    combine = sum(M[["Brown","Blue"],:]) + sum(M[:,["Red","Blond"]]) - overlap;
    println(combine)

    # f) Show that P(x,y) != P(x)P(y)
    println("\nP(Blue eyes, Blond hair) != P(Blue eyes) P(Blond hair)")
    print(M["Blue","Blond"])
    print(" != ")
    println(sum(M["Blue",:]) * sum(M[:,"Blond"]))

end


text:
using StatisticalRethinking, GLM
#gr(size=(600, 600));

howell1 = CSV.read(rel_path("..", "data", "Howell1.csv"), delim=';')
df = convert(DataFrame, howell1);
df2 = filter(row -> row[:age] >= 18, df);
first(df2, 5)

m = lm(@formula(height ~ weight), df2)

coef(m)

scatter( df2[!, :height], residuals(m), xlab="Height",
ylab="Model residual values", lab="Model residuals")

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


text:
# grids.jl : introduces the Grid type for easy use of multigrid algorithms

# coarsen the problem A into a sequence of coarser grids
"""
coarsen(A, sz, R_op, P_op, ngrids)

Coarsen the matrix `A` using `R_op` as restriction operator and `P_op` as interpolation operator. The matrix `A` is the discrete version of a PDE defined on an `sz`-point mesh, and `ngrids` is the number of levels. For MSG, `ngrids` is a tuple with the number of grids in each direction. A Galerkin approach is used to compose the coarse matrices, unless a function `f` is provided for direct discretization.
"""
function coarsen(A::SparseMatrixCSC, sz::NTuple, R_op::TransferKind, P_op::TransferKind, ngrids::NTuple)
    @assert length(sz) == length(ngrids)
    d = length(ngrids)
    grids = Array{Grid{typeof(A),Vector{eltype(A)},NTuple{d,typeof(A)},typeof(sz)},length(sz)}(undef, ngrids...)
    grids[1] = Grid(A,zero_x(A),zero_x(A),tuple(R̃.(1:d,R_op,sz...)...),ntuple(i->spzeros(0,0),d),sz)
    for (idx,i) in enumerate(Base.Iterators.drop(Base.product(UnitRange.(1,ngrids)...),1))
        sz_c = sz.>>(i.-1)
        R_mats = tuple(R̃.(1:d,R_op,sz_c...)...)
        P_mats = tuple(P̃.(1:d,P_op,sz_c...)...)
        grid_num = i.-δ(argmax(i),d)
        A_c = grids[grid_num...].R[argmax(i)]*grids[grid_num...].A*P_mats[argmax(i)]
        grids[i...] = Grid(A_c, zero_x(A_c), zero_x(A_c), R_mats, P_mats, sz_c)
    end
    grids
end

function coarsen(f::Function, sz::NTuple, R_op::TransferKind, P_op::TransferKind, ngrids::NTuple)
    @assert length(sz) == length(ngrids)
    d = length(ngrids)
    A0 = f(sz...)
    grids = Array{Grid{typeof(A0),Vector{eltype(A0)},NTuple{d,typeof(A0)},typeof(sz)},length(sz)}(undef, ngrids...)
    grids[1] = Grid(A0,zero_x(A0),zero_x(A0),tuple(R̃.(1:d,R_op,sz...)...),ntuple(i->spzeros(0,0),d),sz)
    for (idx,i) in enumerate(Base.Iterators.drop(Base.product(UnitRange.(1,ngrids)...),1))
        sz_c = sz.>>(i.-1)
        R_mats = tuple(R̃.(1:d,R_op,sz_c...)...)
        P_mats = tuple(P̃.(1:d,P_op,sz_c...)...)
        grid_num = i.-δ(argmax(i),d)
        A_c = f(sz_c...)
        grids[i...] = Grid(A_c, zero_x(A_c), zero_x(A_c), R_mats, P_mats, sz_c)
    end
    grids
end

δ(i,n) = ntuple(j->j==i,n)

norm_of_residu(grid::Grid) = SimpleMultigrid.norm_of_residu(grid)


text:
"""
    SingletonCommsContext()

A singleton communications context, used for single-process runs.
"""
struct SingletonCommsContext <: AbstractCommsContext end

init(::SingletonCommsContext) = (1, 1)

mypid(::SingletonCommsContext) = 1
iamroot(::SingletonCommsContext) = true
nprocs(::SingletonCommsContext) = 1
barrier(::SingletonCommsContext) = nothing
reduce(::SingletonCommsContext, val, op) = val
gather(::SingletonCommsContext, array) = array

struct SingletonGraphContext <: AbstractGraphContext
    context::SingletonCommsContext
end

function graph_context(
    ctx::SingletonCommsContext,
    send_array,
    send_lengths,
    send_pids,
    recv_array,
    recv_lengths,
    recv_pids,
)

    @assert isempty(send_array)
    @assert isempty(send_lengths)
    @assert isempty(send_pids)
    @assert isempty(recv_array)
    @assert isempty(recv_lengths)
    @assert isempty(recv_pids)
    return SingletonGraphContext(ctx)
end

start(gctx::SingletonGraphContext) = nothing
progress(gctx::SingletonGraphContext) = nothing
finish(gctx::SingletonGraphContext) = nothing


text:
module GemmKernels

include("tiling.jl")

include("config.jl")
include("epilogue.jl")
include("kernel.jl")
include("layout.jl")
include("operator.jl")
include("transform.jl")

include("launch.jl")

include("blas.jl")

end


text:
# Utilities to convert between UTM+zone and Mercator projections

# bound theta in degs, so on output -pi <= theta < pi
bound_thetad(theta) = theta - floor((theta+180) / (360)) * 360

"""
    (zone, isnorth) = utm_zone(lat, lon)
    (zone, isnorth) = utm_zone(ll::LatLon)
    (zone, isnorth) = utm_zone(lla::LLA)
    (zone, isnorth) = utm_zone(ecef::ECEF, datum)

Find the UTM zone and hemisphere (`isnorth = true` or `false`) for the given
latitude and longitude (or world point), including the special rules for Norway
and Svalbard. Zone 0 corresponds to the poles, using the UPS regions.
"""
function utm_zone{T}(lat::T, lon::T)
    if lat > 84
        return (0, true)
    elseif lat < -80
        return (0, false)
    end

    # int versions
    ilat = floor(Int64, bound_thetad(lat))
    ilon = floor(Int64, bound_thetad(lon))

    # get the latitude band
    band = max(-10, min(9,  fld((ilat + 80), 8) - 10))

    # and check for weird ones
    zone = fld((ilon + 186), 6)
    if ((band == 7) && (zone == 31) && (ilon >= 3)) # Norway
        zone = 32
    elseif ((band == 9) && (ilon >= 0) && (ilon < 42)) # Svalbard
        zone = 2 * fld((ilon + 183), 12) + 1
    end

    return (zone, lat >= 0)
end

utm_zone(ll::LatLon) = utm_zone(ll.lat, ll.lon)
utm_zone(lla::LLA) = utm_zone(lla.lat, lla.lon)
utm_zone(ecef::ECEF, datum) = utm_zone(LLAfromECEF(datum)(ecef))

"""
    utm_meridian(zone)

Central meridian of the given UTM `zone` (note - does not include the
conventiional 500km false easting offset).
"""
function utm_meridian(zone::Integer)
    if zone < 1 || zone > 60
        error("UTM zone $zone not defined")
    end
    return 6 * zone - 183
end


text:
function ev_sim(lambda,mu,gamma,Tfinal,C,policy,snapshots=[Inf])

    #barra de progreso
    prog=Progress(101, dt=0.5, desc="Simulando... ");

    #guardo parametros
    params = Dict(
        "ArrivalRate" => lambda,
        "AvgEnergy" => 1.0/mu,
        "AvgDeadline" => 1.0/mu + 1.0/gamma,
        "SimTime" => Tfinal,
        "Capacity" => C,
        "Policy" => get_policy_name(policy),
        "Snapshots" => length(snapshots)
    )

    #variables aleatorias de los clientes
    arr_rng=Exponential(1.0/lambda);
    work_rng=Exponential(1.0/mu);
    deadline_rng=Exponential(1.0/gamma);

    #valores iniciales
    num=convert(Integer,round(3.3*lambda*Tfinal)+length(snapshots));
    T=zeros(num);
    X=zeros(UInt16,length(T));   #charging vehicles
    Y=zeros(UInt16,length(T));   #already charged
    P=zeros(Float64,length(T));   #used power

    snaps = Array{Snapshot}(undef,0);

    t=0.0;
    x=0;
    y=0;
    p=0.0;

    T[1]=t;
    X[1]=x;
    Y[1]=y;
    P[1]=p;

    i=1;    #event counter
    m=1;    #snapshot counter

    charging = Array{EVinstance}(undef,0);
    alreadyCharged = Array{EVinstance}(undef,0);
    finished = Array{EVinstance}(undef,0);

    powerAllocation = Array{Float64}(undef,0);

    arrivals=0;
    expired=0;

    nextArr = rand(arr_rng);
    nextCharge = Inf;
    nextDepON = Inf;
    nextDepOFF = Inf;
    nextSnapshot = snapshots[m];

    dt,caso = findmin([nextArr;nextCharge;nextDepON;nextDepOFF;nextSnapshot])

    while t<Tfinal

        t=t+dt;
        nextArr=nextArr-dt;
        nextSnapshot = nextSnapshot - dt;

        map(v->update_vehicle(v,dt),charging);
        map(v->update_vehicle(v,dt),alreadyCharged);

        if caso==1          #arribo
            arrivals=arrivals+1;
            nextArr = rand(arr_rng);
            x=x+1;

            #sorteo trabajo y deadline
            w=rand(work_rng);
            dep=t+w+rand(deadline_rng);
            push!(charging,EVinstance(t,dep,w,1.0););

        elseif caso==2      #charge completed
            x=x-1;
            y=y+1;

            #guardo el auto que termina
            aux,k = findmin([ev.currentWorkload for ev in charging]); ##TODO agregar un assert de que aux==0?
            ev = charging[k];

            push!(finished,ev);
            push!(alreadyCharged,ev);
            deleteat!(charging,k)

            ev.currentPower=0.0;
            ev.departureWorkload=0.0;
            ev.completionTime=t;

        elseif caso==3      #departure without full charge
            expired=expired+1;
            x=x-1;

            #guardo el auto que termina
            aux,k = findmin([ev.currentDeadline for ev in charging]);
            ev=charging[k];

            push!(finished,ev);
            deleteat!(charging,k)

            ev.currentPower=0.0;
            ev.departureWorkload = ev.currentWorkload;
            ev.completionTime=t;

        elseif caso==4      #departure after full charge
            y=y-1;
            aux,k = findmin([ev.currentDeadline for ev in alreadyCharged]);
            deleteat!(alreadyCharged,k);

        elseif caso==5      #take snapshot

            snapshot=Snapshot(t,deepcopy(charging),deepcopy(alreadyCharged));
            push!(snaps,snapshot);

            if (m<length(snapshots))
                m=m+1;
                nextSnapshot = snapshots[m]-t;
            else
                nextSnapshot = Inf;
            end

        end

        if x>0
            #Apply the policy to all vehicles. Return total user power
            p = policy(charging,C);

            if minimum([ev.currentWorkload for ev in charging])==0
                nextCharge=0;
            else
                nextCharge = minimum([ev.currentWorkload/ev.currentPower for ev in charging]);
            end

            nextDepON = minimum([ev.currentDeadline for ev in charging]);
        else
            nextCharge = Inf;
            nextDepON = Inf;
            p=0.0;
        end

        if y>0
            nextDepOFF = minimum([ev.currentDeadline for ev in alreadyCharged]);
        else
            nextDepOFF = Inf;
        end

        i=i+1;
        T[i]=t;
        X[i]=x;
        Y[i]=y;
        P[i]=p;

        dt,caso = findmin([nextArr;nextCharge;nextDepON;nextDepOFF;nextSnapshot])

        progreso = ceil(Integer,t/Tfinal*100);
        ProgressMeter.update!(prog,progreso);

    end

    T=T[1:i];
    X=X[1:i];
    Y=Y[1:i];
    P=P[1:i];

    pD=expired/arrivals;
    next!(prog);

    trace = TimeTrace(T,X,Y,P);
    stats = SimStatistics([],[],[],[],NaN,NaN,NaN,NaN);
    return EVSim(params,trace,finished,snaps,stats)
end


text:
using SubstitutionModels
using BioSymbols
using Test
using LinearAlgebra
using StaticArrays

import SubstitutionModels._π,
       SubstitutionModels._scale,
       SubstitutionModels.scale_generic,
       SubstitutionModels.P_generic

function test_mod_fun(mod::Type{T}, n_params::Int64, equal_base_freqs::Bool, closed_form_p::Bool) where T <: NASM
  _dummy_params = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]
  _dummy_freqs = [0.21, 0.29, 0.23, 0.27]
  if equal_base_freqs
    @test_throws ErrorException convert(mod, _dummy_params[1:n_params+1], safe=true)
    if n_params > 0
      @test_throws BoundsError mod(_dummy_params[1:n_params-1], false) # safe = false
    end
    @test_nowarn mod(_dummy_params[1:n_params+1], false) # safe = false
    for i in 1:n_params
      flip = fill(1.0, n_params)
      flip[i] *= -1
      @test_throws ErrorException convert(mod, _dummy_params[1:n_params] .* flip, safe=true)
      @test_nowarn mod(_dummy_params[1:n_params] .* flip, false) # safe = false
    end
    @test_throws MethodError convert(mod, _dummy_params[1:n_params], _dummy_freqs)
    @test_nowarn convert(mod, _dummy_params[1:n_params], safe=true)
    x = mod(_dummy_params[1:n_params])
    @test x == supertype(mod)(_dummy_params[1:n_params]) # Convenience constructor
  else
    @test_throws ErrorException convert(mod, _dummy_params[1:n_params+1], _dummy_freqs)
    if n_params > 0
      @test_throws BoundsError mod(_dummy_params[1:n_params-1], _dummy_freqs, false) # safe = false
    end
    for i in 1:n_params
      flip = fill(1.0, n_params)
      flip[i] *= -1
      @test_throws ErrorException convert(mod, _dummy_params[1:n_params] .* flip, _dummy_freqs)
      @test_nowarn mod(_dummy_params[1:n_params] .* flip, _dummy_freqs, false) # safe = false
    end
    @test_throws ErrorException convert(mod, _dummy_params[1:n_params], _dummy_freqs .+ 0.1)
    @test_throws ErrorException convert(mod, _dummy_params[1:n_params], _dummy_freqs[1:3])
    @test_nowarn mod(_dummy_params[1:n_params], [_dummy_freqs; 0.1], false) # safe = false
    @test_throws BoundsError mod(_dummy_params[1:n_params], _dummy_freqs[1:3], false) # safe = false
    @test_throws MethodError convert(mod, _dummy_params[1:n_params])
    @test_nowarn convert(mod, _dummy_params[1:n_params], _dummy_freqs, safe=true)
    x = mod(_dummy_params[1:n_params], _dummy_freqs)
    @test x == supertype(mod)(_dummy_params[1:n_params], _dummy_freqs) # Convenience constructor
  end
  @test_nowarn Q(x)
  @test_nowarn Q(x, true) # Scaled q matrix
  q1 = Q(x)
  q2 = Q(x, true) # Scaled q matrix
  @test all(.≈(sum(q1, dims=2), 0.0, atol=1e-13)) # Q matrix col sums
  @test all(.≈(sum(q2, dims=2), 0.0, atol=1e-13)) # Scaled Q matrix col sums
  @test _scale(x) ≈ scale_generic(x) # Test specific vs. generic scale method
  @test _π(x) ⋅ -diag(q2) ≈ 1.0 # Consistency of π with scaled Q matrix
  @test_throws ErrorException P(x, -1e3)
  @test P(x, [1e3]) ≈ P_generic(x, [1e3]) # Test P generic function
  @test P(x, [1e3], true) ≈ P_generic(x, [1e3], true) # Test P generic function with scaling
  if closed_form_p # If closed form solution to P matrix calculation
    @test P(x, Inf) ≈ _π(x)' .* [1, 1, 1, 1] # P matrix asymptotics
  end
end

for testmod in [(JC69abs, 1, true, true)
                (JC69rel, 0, true, true)
                (K80abs, 2, true, true)
                (K80rel, 1, true, true)
                (F81abs, 1, false, true)
                (F81rel, 0, false, true)
                (F84abs, 2, false, true)
                (F84rel, 1, false, true)
                (HKY85abs, 2, false, true)
                (HKY85rel, 1, false, true)
                (TN93abs, 3, false, true)
                (TN93rel, 2, false, true)
                (GTRabs, 6, false, false)
                (GTRrel, 5, false, false)]
  @testset "$(testmod[1])" begin
    test_mod_fun(testmod...)
  end
end

@testset "Indexing" begin
  testmat = MMatrix{4, 4,Int64}(1:16)
  @test testmat[DNA_A] == testmat[DNA_A, DNA_A]
  testmat[DNA_G] = 100
  @test testmat[DNA_G] == 100
end


text:
using Xyce
using Plots

res = Xyce.simulate("test/VRC.sp")

plot(res["TIME"], [res["V(NET2)"], res["V(NET1)"]])
savefig("xyce.png")


text:
export
    ProjectedNewtonSolverOptions,
    ProjectedNewtonSolver

@with_kw mutable struct ProjectedNewtonStats{T}
    iterations::Int = 0
    c_max::Vector{T} = zeros(5)
    cost::Vector{T} = zeros(5)
end

"""$(TYPEDEF)
Solver options for the Projected Newton solver.
$(FIELDS)
"""
@with_kw mutable struct ProjectedNewtonSolverOptions{T} <: DirectSolverOptions{T}
    verbose::Bool = true
    n_steps::Int = 2
    solve_type::Symbol = :feasible
    active_set_tolerance::T = 1e-3
    constraint_tolerance::T = 1e-6
    ρ::T = 1e-2
    r_threshold::T = 1.1
end

function ProjectedNewtonSolverOptions(opts::SolverOptions)
    ProjectedNewtonSolverOptions(
        constraint_tolerance=opts.constraint_tolerance,
        active_set_tolerance=opts.active_set_tolerance,
        verbose=opts.verbose,
    )
end

struct ProblemInfo{T,N}
    model::AbstractModel
    obj::Objective
    conSet::ConstraintSet{T}
    x0::SVector{N,T}
    xf::SVector{N,T}
end

function ProblemInfo(prob::Problem)
    n = size(prob)[1]
    ProblemInfo(prob.model, prob.obj, prob.constraints, SVector{n}(prob.x0), SVector{n}(prob.xf))
end

"""
$(TYPEDEF)
Projected Newton Solver
Direct method developed by the REx Lab at Stanford University
Achieves machine-level constraint satisfaction by projecting onto the feasible subspace.
    It can also take a full Newton step by solving the KKT system.
This solver is to be used exlusively for solutions that are close to the optimal solution.
    It is intended to be used as a "solution polishing" method for augmented Lagrangian methods.
"""
struct ProjectedNewtonSolver{T,N,M,NM} <: DirectSolver{T}
    # Problem Info
    prob::ProblemInfo{T,N}
    Z::Vector{KnotPoint{T,N,M,NM}}
    Z̄::Vector{KnotPoint{T,N,M,NM}}

    opts::ProjectedNewtonSolverOptions{T}
    stats::ProjectedNewtonStats{T}
    P::Primals{T,N,M}
    P̄::Primals{T,N,M}

    H::SparseMatrixCSC{T,Int}
    g::Vector{T}
    E::Vector{SizedCostExpansion{T,N,N,M}}

    D::SparseMatrixCSC{T,Int}
    d::Vector{T}

    dyn_vals::DynamicsVals{T}
    active_set::Vector{Bool}

    dyn_inds::Vector{SVector{N,Int}}
    con_inds::Vector{<:Vector}
end

function ProjectedNewtonSolver(prob::Problem, opts=SolverOptions())
    Z = prob.Z  # grab trajectory before copy to keep associativity
    prob = copy(prob)  # don't modify original problem

    n,m,N = size(prob)
    NN = n*N + m*(N-1)
    stats = ProjectedNewtonStats()

    # Add dynamics constraints
    add_dynamics_constraints!(prob)
    conSet = prob.constraints
    NP = sum(num_constraints(conSet))

    # Trajectory
    prob_info = ProblemInfo(prob)
    Z̄ = copy(prob.Z)

    # Create concatenated primal vars
    P = Primals(n,m,N)
    P̄ = Primals(n,m,N)

    # Allocate Cost Hessian & Gradient
    H = spzeros(NN,NN)
    g = zeros(NN)
    E = [SizedCostExpansion{Float64}(n,m) for k = 1:N]

    D = spzeros(NP,NN)
    d = zeros(NP)

    fVal = [@SVector zeros(n) for k = 1:N]
    xMid = [@SVector zeros(n) for k = 1:N-1]
    ∇F = [@SMatrix zeros(n,n+m+1) for k = 1:N]
    dyn_vals = DynamicsVals(fVal, xMid, ∇F)
    active_set = zeros(Bool,NP)

    con_inds = gen_con_inds(conSet)

    # Set constant pieces of the Jacobian
    xinds,uinds = P.xinds, P.uinds

    opts_pn = ProjectedNewtonSolverOptions(opts)
    dyn_inds = SVector{n,Int}[]
    ProjectedNewtonSolver(prob_info, Z, Z̄, opts_pn, stats,
        P, P̄, H, g, E, D, d, dyn_vals, active_set, dyn_inds, con_inds)
end

primals(solver::ProjectedNewtonSolver) = solver.P.Z
primal_partition(solver::ProjectedNewtonSolver) = solver.P.xinds, solver.P.uinds

# AbstractSolver interface
Base.size(solver::ProjectedNewtonSolver{T,n,m}) where {T,n,m} = n,m,length(solver.Z)
get_model(solver::ProjectedNewtonSolver) = solver.prob.model
get_constraints(solver::ProjectedNewtonSolver) = solver.prob.conSet
get_trajectory(solver::ProjectedNewtonSolver) = solver.Z
get_objective(solver::ProjectedNewtonSolver) = solver.prob.obj
get_active_set(solver::ProjectedNewtonSolver) = solver.active_set


text:
"""
    Sequential quadratic programming with line search
"""
mutable struct SqpLS{T,Tv,Tt} <: AbstractSqpOptimizer
    @sqp_fields

    soc::Tv # second-order correction direction

    phi::T # merit function value
    μ::Tv # penalty parameters for the merit function
    directional_derivative::T # directional derivative
    alpha::T # stepsize

    function SqpLS(problem::Model{T,Tv,Tt}) where {T,Tv<:AbstractArray{T},Tt}
        sqp = new{T,Tv,Tt}()
        sqp.problem = problem
        sqp.x = Tv(undef, problem.n)
        sqp.p = zeros(problem.n)
        sqp.p_slack = Dict()
        sqp.lambda = zeros(problem.m)
        sqp.mult_x_L = zeros(problem.n)
        sqp.mult_x_U = zeros(problem.n)
        sqp.df = Tv(undef, problem.n)
        sqp.E = Tv(undef, problem.m)
        sqp.dE = Tv(undef, length(problem.j_str))

        sqp.j_row = Vector{Int}(undef, length(problem.j_str))
        sqp.j_col = Vector{Int}(undef, length(problem.j_str))
        for i=1:length(problem.j_str)
            sqp.j_row[i] = Int(problem.j_str[i][1])
            sqp.j_col[i] = Int(problem.j_str[i][2])
        end
        sqp.Jacobian = sparse(sqp.j_row, sqp.j_col, ones(length(sqp.j_row)), problem.m, problem.n)
        sqp.h_row = Vector{Int}(undef, length(problem.h_str))
        sqp.h_col = Vector{Int}(undef, length(problem.h_str))
        for i=1:length(problem.h_str)
            sqp.h_row[i] = Int(problem.h_str[i][1])
            sqp.h_col[i] = Int(problem.h_str[i][2])
        end
        sqp.h_val = Tv(undef, length(problem.h_str))
        sqp.Hessian = sparse(sqp.h_row, sqp.h_col, ones(length(sqp.h_row)), problem.n, problem.n)

        sqp.soc = zeros(problem.n)

        sqp.phi = Inf
        sqp.μ = Tv(undef, problem.m)
        fill!(sqp.μ, 10.0)

        sqp.alpha = 1.0

        sqp.prim_infeas = Inf
        sqp.dual_infeas = Inf
        sqp.compl = Inf

        sqp.options = problem.parameters
        sqp.optimizer = nothing

        sqp.feasibility_restoration = false
        sqp.iter = 1
        sqp.ret = -5
        sqp.start_time = 0.0
        sqp.start_iter_time = 0.0

        return sqp
    end
end

"""
    run!

Run the line-search SQP algorithm
"""
function run!(sqp::SqpLS)

    sqp.start_time = time()

    if sqp.options.OutputFlag == 1
        sparsity_val = ifelse(
            sqp.problem.m > 0,
            length(sqp.problem.j_str) / (sqp.problem.m * sqp.problem.n),
            0.0,
        )
        @printf("Constraint sparsity: %e\n", sparsity_val)
        add_statistics(sqp.problem, "sparsity", sparsity_val)
    else
        Logging.disable_logging(Logging.Info)
    end

    # Set initial point from MOI
    @assert length(sqp.x) == length(sqp.problem.x)
    sqp.x .= sqp.problem.x
    # Adjust the initial point to satisfy the column bounds
    for i = 1:sqp.problem.n
        if sqp.problem.x_L[i] > -Inf
            sqp.x[i] = max(sqp.x[i], sqp.problem.x_L[i])
        end
        if sqp.problem.x_U[i] > -Inf
            sqp.x[i] = min(sqp.x[i], sqp.problem.x_U[i])
        end
    end

    sqp.iter = 1
    is_valid_step = true
    while true

        # Iteration counter limit
        if sqp.iter > sqp.options.max_iter
            sqp.ret = -1
            if sqp.prim_infeas <= sqp.options.tol_infeas
                sqp.ret = 6
            end
            break
        end

        sqp.start_iter_time = time()

        # evaluate function, constraints, gradient, Jacobian
        eval_functions!(sqp)
        sqp.alpha = 0.0
        sqp.prim_infeas = norm_violations(sqp, Inf)
        sqp.dual_infeas = KT_residuals(sqp)
        sqp.compl = norm_complementarity(sqp)

        LP_time_start = time()
        # solve QP subproblem (to initialize dual multipliers)
        # sqp.p, lambda, mult_x_U, mult_x_L, sqp.p_slack, status =
        sqp.p, sqp.lambda, sqp.mult_x_U, sqp.mult_x_L, sqp.p_slack, status =
            sub_optimize!(sqp)

        # directions for dual multipliers
        # p_lambda = lambda - sqp.lambda
        # p_x_U = mult_x_U - sqp.mult_x_U
        # p_x_L = mult_x_L - sqp.mult_x_L

        add_statistics(sqp.problem, "QP_time", time() - LP_time_start)

        if status ∈ [MOI.OPTIMAL, MOI.ALMOST_LOCALLY_SOLVED, MOI.LOCALLY_SOLVED]
            # do nothing
        elseif status ∈ [MOI.INFEASIBLE, MOI.DUAL_INFEASIBLE, MOI.NORM_LIMIT]
            if sqp.feasibility_restoration == true
                @info "Failed to find a feasible direction"
                if sqp.prim_infeas <= sqp.options.tol_infeas
                    sqp.ret = 6
                else
                    sqp.ret = 2
                end
                break
            else
                # println("Feasibility restoration ($(status), |p| = $(norm(sqp.p, Inf))) begins.")
                sqp.feasibility_restoration = true
                continue
            end
        else
            @warn("Unexpected QP subproblem solution status ($status)")
            sqp.ret == -3
            if sqp.prim_infeas <= sqp.options.tol_infeas
                sqp.ret = 6
            end
            break
        end

        compute_mu!(sqp)
        sqp.phi = compute_phi(sqp, sqp.x, 0.0, sqp.p)
        sqp.directional_derivative = compute_derivative(sqp)

        # step size computation
        is_valid_step = compute_alpha(sqp)

        print(sqp)
        collect_statistics(sqp)

        if norm(sqp.p, Inf) <= sqp.options.tol_direction
            if sqp.feasibility_restoration
                sqp.feasibility_restoration = false
                sqp.iter += 1
                continue
            else
                sqp.ret = 0
                break
            end
        end

        if sqp.prim_infeas <= sqp.options.tol_infeas && sqp.compl <= sqp.options.tol_residual
            if sqp.feasibility_restoration
                sqp.feasibility_restoration = false
                sqp.iter += 1
                continue
            elseif sqp.dual_infeas <= sqp.options.tol_residual
                sqp.ret = 0
                break
            end
        end

        # Failed to find a step size
        if !is_valid_step
            @info "Failed to find a step size"
            # if sqp.feasibility_restoration
            #     if sqp.prim_infeas <= sqp.options.tol_infeas
            #         sqp.ret = 6
            #     else
            #         sqp.ret = 2
            #     end
            #     break
            # else
            #     sqp.feasibility_restoration = true
            # end
            # sqp.iter += 1
            # continue

            ## Second-order correction step
            sqp.alpha = 1.0
            sqp.soc, _, _, _, _, status = sub_optimize_soc!(sqp)

            if status ∈ [MOI.OPTIMAL, MOI.ALMOST_LOCALLY_SOLVED, MOI.LOCALLY_SOLVED]
                # TODO: Do we need a line search on this correction
                f_k = sqp.problem.eval_f(sqp.x)
                f_kk = sqp.problem.eval_f(sqp.x + sqp.p)
                f_soc = sqp.problem.eval_f(sqp.x + sqp.p + sqp.soc)
                @info "Second-order correction"  f_k f_kk f_soc
            else
                @warn "Unexpected status ($status) from second-order correction subproblem"
            end
        end

        # @info "solution at k  " sqp.x sqp.problem.eval_f(sqp.x)

        # update primal points
        sqp.x += sqp.alpha .* sqp.p + sqp.soc
        fill!(sqp.soc, 0.0)
        # sqp.lambda += sqp.alpha .* p_lambda
        # sqp.mult_x_U += sqp.alpha .* p_x_U
        # sqp.mult_x_L += sqp.alpha .* p_x_L
        # sqp.lambda += p_lambda
        # sqp.mult_x_U += p_x_U
        # sqp.mult_x_L += p_x_L
        # @info "solution at k+1" sqp.x sqp.problem.eval_f(sqp.x)

        sqp.iter += 1
    end
    sqp.problem.obj_val = sqp.problem.eval_f(sqp.x)
    sqp.problem.status = Int(sqp.ret)
    sqp.problem.x .= sqp.x
    sqp.problem.g .= sqp.E
    sqp.problem.mult_g .= sqp.lambda
    sqp.problem.mult_x_U .= sqp.mult_x_U
    sqp.problem.mult_x_L .= sqp.mult_x_L
    add_statistic(sqp.problem, "iter", sqp.iter)
end

"""
    sub_optimize!

Solve QP subproblems by using JuMP
"""
sub_optimize!(sqp::SqpLS) = sub_optimize!(sqp, JuMP.Model(sqp.options.external_optimizer), 1000.0)

"""
    sub_optimize_soc!

Solve second-order correction subproblem
"""
sub_optimize_soc!(sqp::SqpLS) = sub_optimize_soc!(sqp, JuMP.Model(sqp.options.external_optimizer), 1000.0)

"""
    compute_mu!

Compute the penalty parameter for the merit fucntion
"""
compute_mu!(sqp::AbstractSqpOptimizer) = compute_mu_rule2!(sqp)

function compute_mu_rule1!(sqp::AbstractSqpOptimizer)
    denom = max((1-sqp.options.rho)*norm_violations(sqp, 1), 1.0e-8)
    Hess_part = max(0.5 * sqp.p' * sqp.Hessian * sqp.p, 0.0)
    for i = 1:sqp.problem.m
        sqp.μ[i] = max(sqp.μ[i], (sqp.df' * sqp.p + Hess_part) / denom)
        sqp.μ[i] = max(sqp.μ[i], abs(sqp.lambda[i]))
    end
end
function compute_mu_rule2!(sqp::AbstractSqpOptimizer)
    if sqp.iter == 1
        denom = max((1-sqp.options.rho)*norm_violations(sqp, 1), 1.0e-8)
        Hess_part = max(0.5 * sqp.p' * sqp.Hessian * sqp.p, 0.0)
        for i = 1:sqp.problem.m
            sqp.μ[i] = (sqp.df' * sqp.p + Hess_part) / denom
        end
    else
        for i = 1:sqp.problem.m
            sqp.μ[i] = max(sqp.μ[i], abs(sqp.lambda[i]))
        end
    end
end
function compute_mu_rule3!(sqp::AbstractSqpOptimizer)
    for i = 1:sqp.problem.m
        sqp.μ[i] = max(sqp.μ[i], abs(sqp.lambda[i]))
    end
end

"""
    compute_alpha

Compute step size for line search
"""
function compute_alpha(sqp::AbstractSqpOptimizer)::Bool
    is_valid = true
    sqp.alpha = 1.0

    if norm(sqp.p, Inf) <= sqp.options.tol_direction
        return true
    end

    phi_x_p = compute_phi(sqp, sqp.x, sqp.alpha, sqp.p)
    eta = sqp.options.eta

    # if phi_x_p > sqp.phi
    #     @info "Increasing ϕ" phi_x_p sqp.phi sqp.f sqp.problem.eval_f(sqp.x + sqp.alpha * sqp.p)
    # end
    # E_k = norm_violations(sqp, sqp.x)
    # f_k = sqp.problem.eval_f(sqp.x)

    while phi_x_p > sqp.phi + eta * sqp.alpha * sqp.directional_derivative
        # The step size can become too small.
        if sqp.alpha < sqp.options.min_alpha
            is_valid = false
            break
        end
        sqp.alpha *= sqp.options.tau
        phi_x_p = compute_phi(sqp, sqp.x, sqp.alpha, sqp.p)
        # E_k_p = norm_violations(sqp, sqp.x + sqp.alpha * sqp.p, 1)
        # f_k_p = sqp.problem.eval_f(sqp.x + sqp.alpha * sqp.p)
        # @info "step" sqp.alpha sqp.phi phi_x_p f_k norm(E_k, 1) f_k_p norm(E_k_p, 1)
    end
    # @show phi_x_p, sqp.phi, sqp.alpha, sqp.directional_derivative, is_valid
    return is_valid
end

"""
    print

Print iteration information.
"""
function print(sqp::SqpLS)
    if sqp.options.OutputFlag == 0
        return
    end
    if (sqp.iter - 1) % 25 == 0
        @printf("  %6s", "iter")
        @printf("  %15s", "f(x_k)")
        @printf("  %15s", "ϕ(x_k)")
        @printf("  %15s", "|μ|")
        # @printf("  %15s", "D(ϕ,p)")
        @printf("  %14s", "α")
        @printf("  %14s", "|p|")
        # @printf("  %14s", "α|p|")
        @printf("  %14s", "inf_pr")
        @printf("  %14s", "inf_du")
        @printf("  %14s", "compl")
        @printf("  %10s", "time")
        @printf("\n")
    end
    st = ifelse(sqp.feasibility_restoration, "FR", "  ")
    @printf("%2s%6d", st, sqp.iter)
    @printf("  %+6.8e", sqp.f)
    @printf("  %+6.8e", sqp.phi)
    @printf("  %+6.8e", norm(sqp.μ,Inf))
    # @printf("  %+.8e", sqp.directional_derivative)
    @printf("  %6.8e", sqp.alpha)
    @printf("  %6.8e", norm(sqp.p, Inf))
    # @printf("  %6.8e", sqp.alpha * norm(sqp.p, Inf))
    @printf("  %6.8e", sqp.prim_infeas)
    @printf("  %.8e", sqp.dual_infeas)
    @printf("  %6.8e", sqp.compl)
    @printf("  %10.2f", time() - sqp.start_time)
    @printf("\n")
end

"""
    collect_statistics

Collect iteration information.
"""
function collect_statistics(sqp::SqpLS)
    if sqp.options.StatisticsFlag == 0
        return
    end
    add_statistics(sqp.problem, "f(x)", sqp.f)
    add_statistics(sqp.problem, "ϕ(x_k))", sqp.phi)
    add_statistics(sqp.problem, "D(ϕ,p)", sqp.directional_derivative)
    add_statistics(sqp.problem, "|p|", norm(sqp.p, Inf))
    add_statistics(sqp.problem, "|J|2", norm(sqp.dE, 2))
    add_statistics(sqp.problem, "|J|inf", norm(sqp.dE, Inf))
    add_statistics(sqp.problem, "inf_pr", sqp.prim_infeas)
    # add_statistics(sqp.problem, "inf_du", dual_infeas)
    add_statistics(sqp.problem, "compl", sqp.compl)
    add_statistics(sqp.problem, "alpha", sqp.alpha)
    add_statistics(sqp.problem, "iter_time", time() - sqp.start_iter_time)
    add_statistics(sqp.problem, "time_elapsed", time() - sqp.start_time)
end


text:
##################################################
# Component types
abstract type UCComponentType end
abstract type RequiredConstraints <: UCComponentType end
abstract type SystemConstraints <: UCComponentType end
abstract type GenerationLimits <: UCComponentType end
abstract type PiecewiseProduction <: UCComponentType end
abstract type UpDownTime <: UCComponentType end
abstract type ReserveConstraints <: UCComponentType end
abstract type RampLimits <: UCComponentType end
abstract type StartupCosts <: UCComponentType end
abstract type ShutdownCosts <: UCComponentType end

##################################################
# Components
"""
Generic component of the unit commitment problem.

Elements
===
 * `name`: name of the component
 * `description`: gives a brief summary of what the component adds
 * `type`: reference back to the UCComponentType being modeled
 * `vars`: required variables
 * `constrs`: constraints that are created by this function
 * `add_component`: function to add constraints and update the objective to capture this component
 * `params`: extra parameters the component might use
"""
mutable struct UCComponent
  "Name of the component."
  name::String
  "Description of what the component adds."
  description::String
  "Which part of the unit commitment problem is modeled by this component."
  type::Type{<:UCComponentType}
  "Variables that are needed for the component (subset of `var_list`)."
  vars::Union{Array{Symbol},Nothing}
  "Equations that are modified for the component (subset of `constr_list`)."
  constrs::Union{Array{Symbol},Nothing}
  "Function to add constraints and objective coefficients needed for this component to the model. Signature should be (component, mip, model)."
  add_component::Union{Function,Nothing}
  "Extra parameters for the component."
  params::Any
end # struct UCComponent

export UCComponent


text:
abstract type AbstractEmitter{T} end

@inline get_doa(emitter::AbstractEmitter) = get_doa(emitter.doa)
@inline get_existence(emitter::AbstractEmitter) = get_existence(emitter.exists)

# Can be removed once https://github.com/JuliaLang/julia/pull/32968 is merged
# Julia v1.4
filteremitters(f, xs::Tuple) = Base.afoldl((ys, x) -> f(x) ? (ys..., x) : ys, (), xs...)

function get_steer_vec(manifold::AbstractManifold, emitter::AbstractEmitter, attitude)
    get_steer_vec(manifold, get_doa(emitter), attitude)
end

function get_steer_vec(manifold::IdealManifold{1}, emitter::AbstractEmitter, attitude)
    1.0
end


text:
   #---FD main (fdmain.jl)----
export  init_fields_Por,
        update_velocity_1st_Por!,
        update_stress_1st_Por!,
        ApplyBCLeft_vel!,
        ApplyBCLeft_stress!,
        ApplyBCRight_vel!,
        ApplyBCRight_stress!,
        ApplyBC_stress_AcousticMedia!,
        ApplyBC_stress_ElasticMedia_Ou!,
   #---Additional consideration at borehole wall (boreholewall.jl)
        update_vr_vfr_1st_vertical,
   #---PML (pml.jl)---------
        init_PML_Por,
        init_PML_profile,
        init_memory_variables_Por,
        PML_save_vel!,
        PML_save_stress!,
        PML_update_vel!,
        PML_update_stress!,
        PML_update_memRS!,
        PML_update_memPQ!,
       #---snapshots (snapshots.jl)----
        init_snapshots,
        get_snapshots!,
      #---receivers (receivers.jl)---
        init_receiver,
        init_receiver_hydrophone,
        init_receiver_geophone,
        getRecData_from_index!,
      #---other functions (misc.jl)---------
        mycopy_mat,
        get_Flag_vf_zero,
        check_stability01,
      #---sources (sources.jl)-----
        myricker2,
        srcapply!,
        get_srcindex_monopole,
        get_srcindex_pGauss

  include("./fdmain.jl")
  include("./pml.jl")
  include("./boreholewall.jl")
  include("./receivers.jl")
  include("./sources.jl")
  include("./snapshots.jl")
  include("./misc.jl")


text:
using Pkg; Pkg.activate("HeatEqCalc")
using HeatEqCalc

ENV["GKSwstype"]=100
using Plots
using Printf

sol = HeatEqCalc.load_sol("sol")
t, x, u, m = sol
u0 = u[:, begin]

P = plot(; legend=false)
for k in 1:m
    plot!(x, (p -> p[k]).(u0); lw=0.1, color=:blue)
end
plot()
savefig(P, "heateq0.png")

ylim = (minimum(minimum.(u0)), maximum(maximum.(u0)))
anim = @animate for i in 1:length(t)
    title = @sprintf("t = %4.2f", t[i])
    plot(; legend=false)
    for k in 1:m
        plot!(x, (p -> p[k]).(u[:, i]); ylim, lw=0.1, color=:blue)
    end
    title!(title)
end
gif(anim, "heateq.gif")

max_stars_repo_name: genkuroki/public

max_issues_repo_name: genkuroki/public

max_forks_repo_name: genkuroki/public


text:
"calculates the entropy of a given state"
entropy(state::InfiniteMPS) = [-tr(c*c'*log(c*c')) for c in state.CR]

infinite_temperature(ham::MPOHamiltonian) = [permute(isomorphism(Matrix{eltype(ham[1,1,1])},oneunit(sp)*sp,oneunit(sp)*sp),(1,2,4),(3,)) for sp in ham.pspaces]

"calculates the galerkin error"
calc_galerkin(state::Union{InfiniteMPS,FiniteMPS,MPSComoving},loc,envs)::Float64 =
    norm(leftnull(state.AC[loc])'*ac_prime(state.AC[loc], loc,state,envs))
calc_galerkin(state::Union{InfiniteMPS,FiniteMPS,MPSComoving}, envs)::Float64 =
    maximum([calc_galerkin(state,loc,envs) for loc in 1:length(state)])
calc_galerkin(state::MPSMultiline, envs::PerMPOInfEnv)::Float64 =
    maximum([norm(leftnull(state.AC[row+1,col])'*ac_prime(state.AC[row,col], row,col,state,envs)) for (row,col) in product(1:size(state,1),1:size(state,2))][:])
calc_galerkin(state::MPSMultiline, envs::MixPerMPOInfEnv)::Float64 =
    maximum([norm(leftnull(state.AC[row+1,col])'*ac_prime(envs.above.AC[row,col], row,col,state,envs)) for (row,col) in product(1:size(state,1),1:size(state,2))][:])

"
Calculates the (partial) transfer spectrum
"
function transfer_spectrum(above::InfiniteMPS;below=above,tol=Defaults.tol,num_vals = 20,sector=first(sectors(oneunit(virtualspace(above,1)))))
    init = TensorMap(rand, eltype(eltype(above)), virtualspace(below,0),ℂ[typeof(sector)](sector => 1)'*virtualspace(above,0))

    transferspace = fuse(virtualspace(above,0)*virtualspace(below,0))
    num_vals = min(dim(transferspace, sector), num_vals); # we can ask at most this many values

    eigenvals, eigenvecs,convhist = eigsolve(x->transfer_left(x, above.AL, below.AL) , init, num_vals, :LM, tol=tol)
    convhist.converged < num_vals && @warn "correlation length failed to converge $(convhist.normres)"

    return eigenvals
end

"
Returns the (full) entanglement spectrum at site I
"
function entanglement_spectrum(st::Union{InfiniteMPS,FiniteMPS,MPSComoving},site::Int=0)
    @assert site<=length(st)

    (_,S,_) = tsvd(st.CR[site]);
    diag(convert(Array,S))
end

"""
Find the closest fractions of π, differing at most ```tol_angle```
"""
function approx_angles(spectrum; tol_angle=0.1)
    angles = angle.(spectrum) ./ π                          # ∈ ]-1, 1]
    angles_approx = rationalize.(angles, tol=tol_angle)     # ∈ [-1, 1]

    # Remove the effects of the branchcut.
    angles_approx[findall(angles_approx .== -1)] .= 1       # ∈ ]-1, 1]

    return angles_approx .* π                               # ∈ ]-π, π]
end

"""
Given an InfiniteMPS, compute the gap ```ϵ``` for the asymptotics of the transfer matrix, as well as the Marek gap ```δ``` as a scaling measure of the bond dimension.
"""
function marek_gap(above::InfiniteMPS; tol_angle=0.1, kwargs...)
    spectrum = transfer_spectrum(above; kwargs...)
    return marek_gap(spectrum; tol_angle)
end

function marek_gap(spectrum; tol_angle=0.1)
    # Remove 1s from the spectrum
    inds = findall(abs.(spectrum) .< 1 - 1e-12)
    length(spectrum) - length(inds) < 2 || @warn "Non-injective mps?"

    spectrum = spectrum[inds]

    angles = approx_angles(spectrum; tol_angle=tol_angle)
    θ = first(angles)

    spectrum_at_angle = spectrum[findall(angles .== θ)]

    lambdas = -log.(abs.(spectrum_at_angle));

    ϵ = first(lambdas);

    δ = Inf;
    if length(lambdas) > 2
        δ = lambdas[2]-lambdas[1]
    end

    return ϵ, δ, θ

end

"""
Compute the correlation length of a given InfiniteMPS.
"""
function correlation_length(above::InfiniteMPS; kwargs...)
    ϵ, = marek_gap(above; kwargs...)
    return 1/ϵ
end

function correlation_length(spectrum; kwargs...)
    ϵ, = marek_gap(spectrum; kwargs...)
    return 1/ϵ
end

function variance(state::InfiniteMPS,ham::MPOHamiltonian,envs = environments(state,ham))
    rescaled_ham = ham-expectation_value(state,ham,envs);
    real(sum(expectation_value(state,rescaled_ham*rescaled_ham)))
end

function variance(state::FiniteMPS,ham::MPOHamiltonian,envs = environments(state,ham))
    ham2 = ham*ham;
    real(sum(expectation_value(state,ham2)) - sum(expectation_value(state,ham,envs))^2)
end

function variance(state::MPSComoving,ham::MPOHamiltonian,envs = environments(state,ham))
    #tricky to define
    (ham2,nenvs) = squaredenvs(state,ham,envs);
    real(expectation_value(state,ham2,nenvs)[2] - expectation_value(state,ham,envs)[2]^2)
end

variance(state::FiniteQP,ham::MPOHamiltonian,args...) = variance(convert(FiniteMPS,state),ham);

function variance(state::InfiniteQP,ham::MPOHamiltonian,envs=environments(state,ham))
    # I remember there being an issue here @gertian?
    state.trivial || throw(ArgumentError("variance of domain wall excitations is not implemented"));

    rescaled_ham = ham - expectation_value(state.left_gs,ham);

    #I don't remember where the formula came from
    E_ex = dot(state,effective_excitation_hamiltonian(ham,state,envs));
    E_f = expectation_value(state.left_gs,rescaled_ham,0);

    ham2 = rescaled_ham*rescaled_ham

    real(dot(state,effective_excitation_hamiltonian(ham2,state))-2*(E_f+E_ex)*E_ex+E_ex^2)
end

"""
You can impose periodic boundary conditions on an mpo-hamiltonian (for a given size)
That creates a new mpo-hamiltonian with larger bond dimension
The interaction never wraps around multiple times
"""
function periodic_boundary_conditions(ham::MPOHamiltonian{S,T,E},len = ham.period) where {S,T,E}
    sanitycheck(ham) || throw(ArgumentError("invalid ham"))
    mod(len,ham.period) == 0 || throw(ArgumentError("$(len) is not a multiple of unitcell"))

    fusers = PeriodicArray(map(1:len) do loc
        map(Iterators.product(ham.domspaces[loc,:],ham.domspaces[len+1,:],ham.domspaces[loc,:])) do (v1,v2,v3)
            isomorphism(fuse(v1*v2'*v3),v1*v2'*v3)
        end
    end)

    #a -> what progress have I made in the upper layer?
    #b -> what virtual space did I "lend" in the beginning?
    #c -> what progress have I made in the lower layer?
    χ = ham.odim;
    χ´ = Int((χ-1)*χ*(χ+1)/2+1);

    indmap(a,b,c) = Int((χ-b)*(χ-b+1)/2+(a-1)*χ*(χ+1)/2+(c-b)+1)

    #do the bulk
    bulk = PeriodicArray(convert(Array{Union{T,E},3},fill(zero(E),ham.period,χ´,χ´)));

    for loc in 1:ham.period,
        (j,k) in keys(ham,loc)

        #apply (j,k) above
        l = ham.odim
        for i in 2:ham.odim

            k <= i && i<=l || continue

            f1 = fusers[loc][j,i,l]
            f2 = fusers[loc+1][k,i,l]

            @plansor bulk[loc,indmap(j,i,l),indmap(k,i,l)][-1 -2;-3 -4]:=
                ham[loc,j,k][1 2;-3 6]*f1[-1;1 3 5]*conj(f2[-4;6 7 8])*τ[2 3;7 4]*τ[4 5;8 -2]

        end

        #apply (j,k) below
        i = 1
        for l in 2:(ham.odim-1)

            l > 1 && l >= i && l<=j || continue

            f1 = fusers[loc][i,l,j];
            f2 = fusers[loc+1][i,l,k];

            @plansor bulk[loc,indmap(i,l,j),indmap(i,l,k)][-1 -2;-3 -4] :=
                ham[loc,j,k][1 -2;3 6]*f1[-1;4 2 1]*conj(f2[-4;8 7 6])*τ[5 2;7 3]*τ[-3 4;8 5]
        end
    end

    # make the starter
    starter = convert(Array{Union{T,E},2},fill(zero(E),χ´,χ´));
    for (j,k) in keys(ham,1)

        #apply (j,k) above
        if j == 1
            f1 = fusers[1][1,end,end];
            f2 = fusers[2][k,end,end];

            @plansor starter[1,indmap(k,ham.odim,ham.odim)][-1 -2;-3 -4]:=
                ham[1,j,k][1 2;-3 6]*f1[-1;1 3 5]*conj(f2[-4;6 7 8])*τ[2 3;7 4]*τ[4 5;8 -2]
        end

        #apply (j,k) below
        if j > 1 && j < ham.odim
            f1 = fusers[1][1,j,j];
            f2 = fusers[2][1,j,k];

            @plansor starter[1,indmap(1,j,k)][-1 -2;-3 -4]:=
                ham[1,j,k][4 -2;3 1]*conj(f2[-4;6 2 1])*τ[5 4;2 3]*τ[-3 -1;6 5]

        end

    end
    starter[1,1] = one(E);
    starter[end,end] = one(E);

    # make the ender
    ender = convert(Array{Union{T,E},2},fill(zero(E),χ´,χ´));
    for (j,k) in keys(ham,ham.period)

        if k > 1
            f1 = fusers[end][j,k,ham.odim]
            @plansor ender[indmap(j,k,ham.odim),end][-1 -2;-3 -4]:=
                f1[-1;1 2 6]*ham[ham.period,j,k][1 3;-3 4]*τ[3 2;4 5]*τ[5 6;-4 -2]
        end
    end
    ender[1,1] = one(E);
    ender[end,end] = one(E);

    # fill in the entire ham
    nos = convert(Array{Union{T,E},3},fill(zero(E),len,χ´,χ´))
    nos[1,:,:] = starter[:,:];
    nos[end,:,:] = ender[:,:];

    for i in 2:len-1
        nos[i,:,:] = bulk[i,:,:];
    end

    return MPOHamiltonian(nos)
end

#impose periodic boundary conditions on a normal mpo
function periodic_boundary_conditions(mpo::InfiniteMPO{O},len = length(mpo)) where O
    mod(len,length(mpo)) == 0 || throw(ArgumentError("len not a multiple of unitcell"))

    output = PeriodicArray{O,1}(undef,len);

    sp = _firstspace(mpo[1])';
    utleg = Tensor(ones,oneunit(sp));

    #do the bulk
    for j in 2:len-1
        f1 = isomorphism(fuse(sp*_firstspace(mpo[j])),sp*_firstspace(mpo[j]))
        f2 = isomorphism(fuse(sp*_lastspace(mpo[j])'),sp*_lastspace(mpo[j])')

        @plansor output[j][-1 -2;-3 -4] := mpo[j][2 -2;3 5]*f1[-1;1 2]*conj(f2[-4;4 5])*τ[-3 1;4 3]
    end

    #do the left
    f2 = isomorphism(fuse(sp*_lastspace(mpo[1])'),sp*_lastspace(mpo[1])')
    @plansor output[1][-1 -2;-3 -4] := mpo[1][1 -2;3 5]*conj(f2[-4;4 5])*τ[-3 1;4 3]*utleg[-1]

    #do the right
    f2 = isomorphism(fuse(sp*_firstspace(mpo[len])),sp*_firstspace(mpo[len]))
    @plansor output[end][-1 -2;-3 -4] := mpo[len][2 -2;3 4]*f2[-1;1 2]*τ[-3 1;4 3]*conj(utleg[-4])

    InfiniteMPO(output)
end


text:
include(joinpath(@__DIR__, "../../../PlotUtils/PlotUtils.jl"))
include(joinpath(@__DIR__, "../../../PlotUtils/imports.jl"))

using Distributed

rmprocs(filter(w -> w > 1, workers()))
addprocs(27)

@everywhere include(joinpath(@__DIR__, "../../../PlotUtils/PlotUtils.jl"))
@everywhere include(joinpath(@__DIR__, "../../../PlotUtils/imports.jl"))

# Directory where results are
simname = basename(strip(@__DIR__, '/'))
results_dir = "/scratchdata/alui2/cytof/results/repfam/$(simname)"

### Create metrics ###

# Name of output file
@everywhere OUTPUT_FILE = "output.bson"

println("Producing metrics ...")

# directories to loop through
results_dirs = ["$(results_dir)/dataseed_$(dataseed)/mcmcseed_$(mcmcseed)/scale_$(scale)"
                for dataseed in 1:3
                for mcmcseed in 1:3
                for scale in (0, 1, 10)]

# Number of directories
num_dirs = length(results_dirs)

parsekey(T, key, path) = parse(T, match(Regex("(?<=$(key)_)\\d+"), path).match)

# Construct a data frame for R
function fill_metrics_df!(path, metrics, metrics_df)
  # Parse dataseed
  dataseed = parsekey(Int, "dataseed", path)

  # Parse mcmcseed
  mcmcseed = parsekey(Int, "mcmcseed", path)

  # Parse scale
  scale = parsekey(Int, "scale", path)

  I = length(metrics[:Rmean])
  for i in 1:I
    append!(metrics_df,
            DataFrames.DataFrame(i=i,
                                 dataseed=dataseed,
                                 mcmcseed=mcmcseed,
                                 scale=scale,
                                 DIC=metrics[:DIC],
                                 LPML=metrics[:LPML],
                                 Rmean=metrics[:Rmean][i],
                                 R_975=metrics[:R_975][i],
                                 R_025=metrics[:R_025][i]))
  end
end

# Create make metrics function which only takes directory as input
@everywhere function getmetrics(d, OUTPUT_FILE)
  output = BSON.load("$(d)/$(OUTPUT_FILE)")
  metrics = Dict{Symbol, Any}()
  metrics[:LPML] = output[:LPML]
  metrics[:DIC] = output[:DIC]

  # Get posterior samples of W
  Ws = [o[:theta__W] for o in output[:samples][1]]

  # Posterior samples of R (I x K) -- number of active features / sample
  # Rs (I x num_mcmc_samples)
  Rs = hcat([vec(sum(W .> 0, dims=2)) for W in Ws]...)

  # Mean of R, by sample (vector of length I)
  Rmean = vec(mean(Rs, dims=2))

  # Mean of R, by sample (vector of length I)
  R_025 = PlotUtils.quantiles(Rs, .025, dims=2, drop=true)

  # Mean of R, by sample (vector of length I)
  R_975 = PlotUtils.quantiles(Rs, .975, dims=2, drop=true)

  # Append Rmean to metrics dictionary
  metrics[:Rmean] = Rmean

  # Append Rmean to metrics dictionary
  metrics[:R_025] = R_025

  # Append Rmean to metrics dictionary
  metrics[:R_975] = R_975

  return metrics
end

@everywhere make_metrics(d) = getmetrics(d, OUTPUT_FILE)

# Get all metrics in parallel
allmetrics = pmap(make_metrics, results_dirs)

# Put results in a data frame.
metrics_df = let
  df = DataFrames.DataFrame()
  for (d, metrics) in zip(results_dirs, allmetrics)
    fill_metrics_df!(d, metrics, df)
  end
  df
end

CSV.write("metrics/metrics.csv", metrics_df)

# Remove extra processors
rmprocs(filter(w -> w > 1, workers()))

println("DONE!")


text:
export
    MMTFDict,
    generatechainid,
    writemmtf

"""
    MMTFDict(filepath)
    MMTFDict(io)
    MMTFDict()

A Macromolecular Transmission Format (MMTF) dictionary.

Can be accessed using similar functions to a standard `Dict`.
Keys are field names as a `String` and values are various types.
To directly access the underlying dictionary of `MMTFDict` `d`, use
`d.dict`.
Call `MMTFDict` with a filepath or stream to read the dictionary from that
source.
The keyword argument `gzip` (default `false`) determines if the file is gzipped.
"""
struct MMTFDict
    dict::Dict{String, Any}
end

# Create an empty MMTF dictionary
# Matches the decoded form of a MMTF file using MMTF.jl
# Encoding and decoding this Dict gives an identical Dict
MMTFDict() = MMTFDict(Dict{String, Any}(
        "altLocList"         => Char[],
        "atomIdList"         => Int32[],
        "bFactorList"        => Float32[],
        "bioAssemblyList"    => Any[],
        "bondAtomList"       => Int32[],
        "bondOrderList"      => Int8[],
        "chainIdList"        => String[],
        "chainNameList"      => String[],
        "chainsPerModel"     => Any[],
        "depositionDate"     => "",
        "entityList"         => Any[],
        "experimentalMethods"=> Any[],
        "groupIdList"        => Int32[],
        "groupList"          => Any[],
        "groupsPerChain"     => Any[],
        "groupTypeList"      => Int32[],
        "insCodeList"        => Char[],
        "mmtfProducer"       => "",
        "mmtfVersion"        => "",
        "ncsOperatorList"    => Any[],
        "numAtoms"           => 0,
        "numBonds"           => 0,
        "numChains"          => 0,
        "numGroups"          => 0,
        "numModels"          => 0,
        "occupancyList"      => Float32[],
        "releaseDate"        => "",
        "resolution"         => 0.0,
        "rFree"              => "",
        "rWork"              => "",
        "secStructList"      => Int8[],
        "sequenceIndexList"  => Int32[],
        "spaceGroup"         => "",
        "structureId"        => "",
        "title"              => "",
        "unitCell"           => Any[],
        "xCoordList"         => Float32[],
        "yCoordList"         => Float32[],
        "zCoordList"         => Float32[]))

MMTFDict(filepath::AbstractString; gzip::Bool=false) = MMTFDict(parsemmtf(filepath; gzip=gzip))
MMTFDict(io::IO; gzip::Bool=false) = MMTFDict(parsemmtf(io; gzip=gzip))

Base.getindex(mmtf_dict::MMTFDict, field::AbstractString) = mmtf_dict.dict[field]

function Base.setindex!(mmtf_dict::MMTFDict,
                    val,
                    field::AbstractString)
    mmtf_dict.dict[field] = val
    return mmtf_dict
end

Base.keys(mmtf_dict::MMTFDict) = keys(mmtf_dict.dict)
Base.values(mmtf_dict::MMTFDict) = values(mmtf_dict.dict)
Base.haskey(mmtf_dict::MMTFDict, key) = haskey(mmtf_dict.dict, key)

function Base.show(io::IO, mmtf_dict::MMTFDict)
    print(io, "MMTF dictionary with $(length(keys(mmtf_dict))) fields")
end

function Base.read(input::IO,
            ::Type{MMTF};
            structure_name::AbstractString="",
            remove_disorder::Bool=false,
            read_std_atoms::Bool=true,
            read_het_atoms::Bool=true,
            gzip::Bool=false)
    d = parsemmtf(input; gzip=gzip)
    struc = ProteinStructure(structure_name)
    # Extract hetero atom information from entity list
    hets = trues(length(d["chainIdList"]))
    for e in d["entityList"]
        if e["type"] == "polymer"
            for i in e["chainIndexList"]
                # 0-based indexing in MMTF
                hets[i + 1] = false
            end
        end
    end
    model_i = 0
    chain_i = 0
    group_i = 0
    atom_i = 0
    for modelchaincount in d["chainsPerModel"]
        model_i += 1
        struc[model_i] = Model(model_i, struc)
        for ci in 1:modelchaincount
            chain_i += 1
            for gi in 1:d["groupsPerChain"][chain_i]
                group_i += 1
                # 0-based indexing in MMTF
                group = d["groupList"][d["groupTypeList"][group_i] + 1]
                for ai in 1:length(group["atomNameList"])
                    atom_i += 1
                    if (read_std_atoms || hets[chain_i]) && (read_het_atoms || !hets[chain_i])
                        unsafe_addatomtomodel!(
                            struc[model_i],
                                AtomRecord(
                                    hets[chain_i],
                                    d["atomIdList"][atom_i],
                                    group["atomNameList"][ai],
                                    d["altLocList"][atom_i] == '\0' ? ' ' : d["altLocList"][atom_i],
                                    group["groupName"],
                                    d["chainNameList"][chain_i],
                                    d["groupIdList"][group_i],
                                    d["insCodeList"][group_i] == '\0' ? ' ' : d["insCodeList"][group_i],
                                    [
                                        d["xCoordList"][atom_i],
                                        d["yCoordList"][atom_i],
                                        d["zCoordList"][atom_i],
                                    ],
                                    d["occupancyList"][atom_i],
                                    d["bFactorList"][atom_i],
                                    group["elementList"][ai],
                                    # Add + to positive charges to match PDB convention
                                    group["formalChargeList"][ai] > 0 ? "+$(group["formalChargeList"][ai])" :
                                                string(group["formalChargeList"][ai])
                                );
                                remove_disorder=remove_disorder)
                    end
                end
            end
        end
    end
    if length(d["atomIdList"]) != atom_i
        throw(ArgumentError("Discrepancy between atom count ($(length(d["atomIdList"]))) and atoms read in ($atom_i)"))
    end
    # Remove any models that were not added to
    for model in struc
        if countchains(model) == 0
            delete!(models(struc), modelnumber(model))
        end
    end
    fixlists!(struc)
    return struc
end

"""
    generatechainid(entity_id)

Convert a positive `Integer` into a chain ID.
Goes A to Z, then AA to ZA, AB to ZB etc.
This is in line with PDB conventions.
"""
function generatechainid(entity_id::Integer)
    entity_id > 0 || throw(ArgumentError("Entity ID $entity_id is not positive"))
    divisor = entity_id
    out_string = ""
    while divisor > 0
        mod = (divisor - 1) % 26
        out_string *= Char(65 + mod)
        divisor = Int(floor((divisor - mod) / 26))
    end
    return out_string
end

"""
    writemmtf(output, element, atom_selectors...)
    writemmtf(output, mmtf_dict)

Write a `StructuralElementOrList` or a `MMTFDict` to a MMTF file or output
stream.

Atom selector functions can be given as additional arguments - only atoms
that return `true` from all the functions are retained.
The keyword argument `expand_disordered` (default `true`) determines whether to
return all copies of disordered residues and atoms.
The keyword argument `gzip` (default `false`) determines if the file should be
gzipped.
"""
function writemmtf(output::Union{AbstractString, IO},
                d::MMTFDict;
                gzip::Bool=false)
    writemmtf(d.dict, output; gzip=gzip)
    return
end

function writemmtf(output::Union{AbstractString, IO},
                el::StructuralElementOrList,
                atom_selectors::Function...;
                expand_disordered::Bool=true,
                gzip::Bool=false)
    loop_el = isa(el, ProteinStructure) ? collectmodels(el) : el
    ats = sort(sort(sort(collectatoms(loop_el, atom_selectors...;
                            expand_disordered=expand_disordered)),
                            by=residue), by=model)
    d = MMTFDict()
    if length(ats) == 0
        return writemmtf(output, d; gzip=gzip)
    end
    last_model = model(first(ats))
    last_chain = chain(first(ats))
    last_res = residue(first(ats))
    chain_i = 0
    group_count = 0
    sequence = ""
    prev_resname = ""
    prev_het = true
    for (i, at) in enumerate(ats)
        if model(at) != last_model
            push!(d["chainsPerModel"], chain_i)
            chain_i = 0
            last_model = model(at)
        end

        if chain(at) != last_chain || i == 1
            # MMTF splits chains up by molecules so we determine chain splits
            #   at the residue level
            # The below just ensures we have a new entity for each chain
            prev_resname = ""
            prev_het = true
            last_chain = chain(at)
        end

        res = residue(at)
        if residue(at) != last_res || i == 1
            # Determine whether we have changed entity
            # ATOM blocks, and hetero molecules with the same name, are
            #   treated as the same entity
            if ishetero(res) != prev_het || (prev_het && resname(res) != prev_resname) || group_count == 0
                chain_i += 1
                push!(d["chainIdList"], generatechainid(chain_i))
                push!(d["chainNameList"], chainid(res))
                # Add the groupsPerChain and sequence for the previous chain
                if group_count > 0
                    push!(d["groupsPerChain"], group_count)
                    d["entityList"][end]["sequence"] = sequence
                    group_count = 0
                    sequence = ""
                end
                # Checking for similar entities is non-trivial so we treat
                #   each molecule as a separate entity
                push!(d["entityList"], Dict{Any, Any}(
                    "chainIndexList"=> Any[length(d["chainIdList"]) - 1],
                    "description"   => "",
                    "sequence"      => "", # This is changed later
                    "type"          => ishetero(res) ? "non-polymer" : "polymer"
                ))
            end
            if !ishetero(res)
                sequence *= string(AminoAcidSequence(res; gaps=false))
            end
            group_count += 1

            # Look for an existing group with the correct residue name and
            #   atom names present
            # Look forwards in the loop to get atom names
            ats_res = AbstractAtom[at]
            for j in (i + 1):length(ats)
                if residue(ats[j]) == res
                    push!(ats_res, ats[j])
                else
                    break
                end
            end
            at_names = atomname.(ats_res)
            group_i = 0
            for (gi, group) in enumerate(d["groupList"])
                if group["groupName"] == resname(res) && group["atomNameList"] == at_names
                    group_i = gi
                    break
                end
            end

            if group_i == 0
                push!(d["groupList"], Dict{Any, Any}(
                    "groupName"       => resname(res),
                    "bondAtomList"    => Any[],
                    "elementList"     => Any[element(at) for at in ats_res],
                    # MMTF specifies missing charges as zero
                    "formalChargeList"=> Any[charge(at) == "" ? 0 : parse(Int64, charge(at)) for at in ats_res],
                    "singleLetterCode"=> "",
                    "chemCompType"    => "",
                    "atomNameList"    => Any[at_names...],
                    "bondOrderList"   => Any[]
                ))
            end

            push!(d["groupIdList"], resnumber(res))
            push!(d["groupTypeList"], group_i == 0 ? length(d["groupList"]) - 1 : group_i - 1)
            push!(d["insCodeList"], inscode(res) == ' ' ? '\0' : inscode(res))
            push!(d["secStructList"], -1)
            push!(d["sequenceIndexList"], ishetero(res) ? -1 : length(sequence) - 1)

            prev_resname = resname(res)
            prev_het = ishetero(res)
            last_res = res
        end

        push!(d["altLocList"], altlocid(at) == ' ' ? '\0' : altlocid(at))
        push!(d["atomIdList"], serial(at))
        push!(d["bFactorList"], tempfactor(at))
        push!(d["occupancyList"], occupancy(at))
        push!(d["xCoordList"], x(at))
        push!(d["yCoordList"], y(at))
        push!(d["zCoordList"], z(at))
    end

    push!(d["chainsPerModel"], chain_i)
    push!(d["groupsPerChain"], group_count)
    d["entityList"][end]["sequence"] = sequence

    d["numModels"] = length(d["chainsPerModel"])
    d["numChains"] = length(d["chainIdList"])
    d["numGroups"] = length(d["groupIdList"])
    d["numAtoms"] = length(d["atomIdList"])
    d["structureId"] = structurename(first(el))
    d["mmtfVersion"] = "1.0.0"
    d["mmtfProducer"] = "BioStructures.jl"

    return writemmtf(output, d; gzip=gzip)
end


text:
using MPI
function main()
    MPI.Init()
    comm = MPI.COMM_WORLD
    rank = MPI.Comm_rank(comm)
    if rank == 1
        error("Goodbye")
    end
    MPI.Barrier(comm)
end
main()


text:
using MXNet

mlp = @mx.chain mx.Variable(:data)             =>
  mx.FullyConnected(name=:fc1, num_hidden=128) =>
  mx.Activation(name=:relu1, act_type=:relu)   =>
  mx.FullyConnected(name=:fc2, num_hidden=64)  =>
  mx.Activation(name=:relu2, act_type=:relu)   =>
  mx.FullyConnected(name=:fc3, num_hidden=10)  =>
  mx.SoftmaxOutput(name=:softmax)

# data provider
batch_size = 100
include(Pkg.dir("MXNet", "examples", "mnist", "mnist-data.jl"))
train_provider, eval_provider = get_mnist_providers(batch_size)

# setup model
model = mx.FeedForward(mlp, context=mx.gpu())

# optimization algorithm
optimizer = mx.SGD(lr=0.1, momentum=0.9)

# fit parameters
mx.fit(model, optimizer, train_provider, n_epoch=20, eval_data=eval_provider)

You can also predict using the model in the following way:

probs = mx.predict(model, eval_provider)

# collect all labels from eval data
labels = Array[]
for batch in eval_provider
    push!(labels, copy(mx.get(eval_provider, batch, :softmax_label)))
end
labels = cat(1, labels...)

# Now we use compute the accuracy
correct = 0
for i = 1:length(labels)
    # labels are 0...9
    if indmax(probs[:,i]) == labels[i]+1
        correct += 1
    end
end
accuracy = 100correct/length(labels)
println(mx.format("Accuracy on eval set: {1:.2f}%", accuracy))


text:
using Test
using QCMaterial

@testset "hamiltonian.generate_ham_1d_hubbard" begin
    nsite = 2
    n_qubit = 2*nsite
    U = 1.0
    t = -0.01
    μ = 0.0

    ham = generate_ham_1d_hubbard(t, U, nsite, μ)

    ham_ref = FermionOperator()
    up1 = up_index(1)
    down1 = down_index(1)
    up2 = up_index(2)
    down2 = down_index(2)

    #クーロン斥力
    ham_ref += FermionOperator("$(up1)^ $(down1)^ $(up1) $(down1)", -U)
    ham_ref += FermionOperator("$(up2)^ $(down2)^ $(up2) $(down2)", -U)

    #ホッピング積分
    ham_ref += FermionOperator("$(up2)^ $(up1)", t)
    ham_ref += FermionOperator("$(up1)^ $(up2)", t)
    ham_ref += FermionOperator("$(down2)^ $(down1)", t)
    ham_ref += FermionOperator("$(down1)^ $(down2)", t)

    #ケミカルポテンシャルの項
    ham_ref += FermionOperator("$(up1)^  $(up1) ", -μ)
    ham_ref += FermionOperator("$(down1)^ $(down1)", -μ)
    ham_ref += FermionOperator("$(up2)^  $(up2) ", -μ)
    ham_ref += FermionOperator("$(down2)^ $(down2)", -μ)

    @test ham == ham_ref
end

@testset "hamiltonian.generate_impurity_dimer" begin
    nsite = 2
    U = 1.0
    V = 1.0
    μ = 1.0
    ε = 1.0

    ham = generate_impurity_ham_with_1imp_multibath(U, V, μ, ε, nsite)

    ham_ref = FermionOperator()
    up1 = up_index(1)
    down1 = down_index(1)
    up2 = up_index(2)
    down2 = down_index(2)

    #Coulomb
    ham_ref += FermionOperator("$(up1)^ $(down1)^ $(up1) $(down1)", -U)

    #hybridization
    ham_ref += FermionOperator("$(up2)^ $(up1)", -V)
    ham_ref += FermionOperator("$(up1)^ $(up2)", -V)
    ham_ref += FermionOperator("$(down2)^ $(down1)", -V)
    ham_ref += FermionOperator("$(down1)^ $(down2)", -V)

    #chemical potential
    ham_ref += FermionOperator("$(up1)^  $(up1) ", -μ)
    ham_ref += FermionOperator("$(down1)^ $(down1)", -μ)

    #bath energy level
    ham_ref += FermionOperator("$(up2)^  $(up2) ", ε)
    ham_ref += FermionOperator("$(down2)^ $(down2)", ε)

    @test ham == ham_ref
end

@testset "hamiltonian.generate_impurity_ham_with_1imp_3bath" begin
    nsite = 4
    U = 1.0
    V = 1.0
    μ = 1.0
    ε = 1.0

    ham = generate_impurity_ham_with_1imp_multibath(U, V, μ, ε, nsite)

    ham_ref = FermionOperator()
    up1 = up_index(1)
    down1 = down_index(1)
    up2 = up_index(2)
    down2 = down_index(2)
    up3 = up_index(3)
    down3 = down_index(3)
    up4 = up_index(4)
    down4 = down_index(4)

    #Coulomb
    ham_ref += FermionOperator("$(up1)^ $(down1)^ $(up1) $(down1)", -U)

    #hybridization
    ham_ref += FermionOperator("$(up2)^ $(up1)", -V)
    ham_ref += FermionOperator("$(up1)^ $(up2)", -V)
    ham_ref += FermionOperator("$(down2)^ $(down1)", -V)
    ham_ref += FermionOperator("$(down1)^ $(down2)", -V)

    ham_ref += FermionOperator("$(up3)^ $(up1)", -V)
    ham_ref += FermionOperator("$(up1)^ $(up3)", -V)
    ham_ref += FermionOperator("$(down3)^ $(down1)", -V)
    ham_ref += FermionOperator("$(down1)^ $(down3)", -V)

    ham_ref += FermionOperator("$(up4)^ $(up1)", -V)
    ham_ref += FermionOperator("$(up1)^ $(up4)", -V)
    ham_ref += FermionOperator("$(down4)^ $(down1)", -V)
    ham_ref += FermionOperator("$(down1)^ $(down4)", -V)

    #chemical potential
    ham_ref += FermionOperator("$(up1)^  $(up1) ", -μ)
    ham_ref += FermionOperator("$(down1)^ $(down1)", -μ)

    #bath energy level
    ham_ref += FermionOperator("$(up2)^  $(up2) ", ε)
    ham_ref += FermionOperator("$(down2)^ $(down2)", ε)

    ham_ref += FermionOperator("$(up3)^  $(up3) ", ε)
    ham_ref += FermionOperator("$(down3)^ $(down3)", ε)

    ham_ref += FermionOperator("$(up4)^  $(up4) ", ε)
    ham_ref += FermionOperator("$(down4)^ $(down4)", ε)

    @test ham == ham_ref
end

max_stars_repo_name: Sakurairihito/QCMaterialNew

max_issues_repo_name: sakurairihito/QCMaterialNew

max_forks_repo_name: sakurairihito/QCMaterialNew