text:
# Helper functions to simplify ERFA usage

using ERFA: DAS2R, DJM0, apco13, atciq, atco13, atioq
using RadioInterferometry: dms2rad
using EarthOrientation: getΔUT1, getxp, getyp

const MJD0 = DJM0

"""
    radec2hadec(α, δ, jdutc;
                longitude, latitude, altitude,
                jdutc2=0,
                dut1=getΔUT1(jdutc+jdutc2),
                xp=getxp(jdutc+jdutc2)*DAS2R,
                yp=getxp(jdutc+jdutc2)*DAS2R
               ) -> (hob, δob)

    radec2hadec(αδ::AbstractMatrix, jdutc; kwargs...) -> Matrix[hobs; δobs]

Convert ICRS right ascension `α` and declination `δ` to observed hour angle
`hob` and observed declination `δob` as seen from the location specified by
`longitude`, `latitude`, and `altitude` at the Julian date `jdutc+jdutc2`.
Earth orientation parameters `dut1`, `xp`, and `yp` may be specified if desired,
otherwise values from EarthOrientation for `jdutc+jdutc2` will be used.

If `α` and `δ` are Vectors and `jdutc+jdutc2` is a scalar, the
position-independent astrometric parameters for `jdutc+jdutc2` will be computed
only once and `hob` and `δob` will also be Vectors.  Multiple right ascensions
and declinations given as a 2xN matrix `αδ` consisting of [`α`, `δ`] columns
will also compute the position-independent astrometric parameters only once for
all positions.

The `longitude` and `latitude` parameters should be given in degrees.  They may
be numeric values or Strings in decimal ("dd.ddd") or sexagesimal
("dd:mm:ss.sss") format.

!!! note
    Be sure to use the right units for the various parameters!

| Parameter | Description            | Units   |
|-----------|:-----------------------|:--------|
|         α | Right ascention (ICRS) | radians |
|         δ | Declination (ICRS)     | radians |
|     jdutc | Julian Date (UTC)      | days    |
| longitude | Longitude              | degrees |
|  latitude | Geodetic latitude      | degrees |
|  altitude | Altitude               | meters  |
|    jdutc2 | Julian Date (UTC)      | days    |
|      dut1 | UTC-UT1                | seconds |
|        xp | Polar motion X         | radians |
|        yp | Polar motion Y         | radians |
"""
function radec2hadec(α, δ, jdutc;
                     longitude, latitude, altitude,
                     jdutc2=0,
                     dut1=getΔUT1(jdutc+jdutc2),
                     xp=getxp(jdutc+jdutc2)*DAS2R,
                     yp=getxp(jdutc+jdutc2)*DAS2R,
                     _ignored_kwargs...
                    )
    _aob, _zob, hob, δob, _αob, _eo = atco13(α, δ,
        0, 0, 0, 0,
        jdutc, jdutc2, dut1,
        dms2rad(longitude), dms2rad(latitude), altitude,
        xp, yp,
        0, 0, 0, 0
    )

    (hob, δob)
end

function radec2hadec(α::AbstractVector, δ::AbstractVector, jdutc::Real;
                     longitude, latitude, altitude,
                     jdutc2=0,
                     dut1=getΔUT1(jdutc+jdutc2),
                     xp=getxp(jdutc+jdutc2)*DAS2R,
                     yp=getxp(jdutc+jdutc2)*DAS2R,
                     _ignored_kwargs...
                    )

    # Star-independent astrometry parameters
    astrom, _eo = apco13(jdutc, jdutc2, dut1,
                         dms2rad(longitude), dms2rad(latitude), altitude,
                         xp, yp,
                         0, 0, 0, 0
                        )

    # Transform ICRS to CIRS
    αδis = atciq.(α, δ, 0, 0, 0, 0, Ref(astrom))

    αis = getindex.(αδis, 1)
    δis = getindex.(αδis, 2)

    # Transform CIRS to observed
    azhδαobs = atioq.(αis, δis, Ref(astrom))

    hobs = getindex.(azhδαobs, 3)
    δobs = getindex.(azhδαobs, 4)
    (hobs, δobs)
end

function radec2hadec(αδ::AbstractMatrix, jdutc::Real;
                     longitude, latitude, altitude,
                     jdutc2=0,
                     dut1=getΔUT1(jdutc+jdutc2),
                     xp=getxp(jdutc+jdutc2)*DAS2R,
                     yp=getxp(jdutc+jdutc2)*DAS2R,
                     _ignored_kwargs...
                    )
    @assert size(αδ, 1) == 2 "expected 2xN Matrix"

    hobs, δobs = radec2hadec(αδ[1,:], αδ[2,:], jdutc;
                             longitude, latitude, altitude,
                             jdutc2, dut1, xp, yp
                            )

    [permutedims(hobs); permutedims(δobs)]
end


text:
module LabelPropagation

using Random: randperm
using StatsBase: countmap
using LightGraphs

export labelpropagation

function labelpropagation(graph::AbstractGraph, max_iter=10000)
    labels = collect(1:nv(graph))
    stable_labels = false
    for _ in 1:max_iter
        if reassign_labels!(graph, labels)
            break
        end
    end
    return labels
end

function reassign_labels!(graph::AbstractGraph{T}, labels::Vector{T}) where T<:Integer
    permu = randperm(nv(graph))
    c = 0
    stable = true
    # for i in permutation
    for (i, node) in zip(permu, vertices(graph)[permu])
        label = labels[i]
        most_frequent_label = find_most_frequent_label(graph, node, labels)
        if most_frequent_label != label
            labels[i] = most_frequent_label
            stable = false
            c += 1
        end
    end
    println(c)
    return stable
end

function find_most_frequent_label(graph::AbstractGraph{T}, node::T,
                                  labels::Array{T}) where T<:Integer
    neighbours_labels = countmap(labels[neighbors(graph, node)])
    most_frequent_label, _ = findmax(neighbours_labels)
    return most_frequent_label
end

end

ext: jl

max_stars_repo_name: quentin-burthier/GrandsGraphes

max_issues_repo_name: quentin-burthier/GrandsGraphes

max_forks_repo_name: quentin-burthier/GrandsGraphes


text:
"""
    assemble_flux_matrix(facets, nodes[, weights])

Assemble 3D flux matrix using P1 finite elements.

This function is based on the Matlab function `flux_matrixP1_3D.m` from the Matlab
nodal matrix assembly toolbox by Jan Valdman and Talal Rahman.

https://uk.mathworks.com/matlabcentral/fileexchange/27826-fast-fem-assembly-nodal-elements

Talal Rahman and Jan Valdman: Fast MATLAB assembly of FEM matrices in 2D and 3D: nodal
elements, Applied Mathematics and Computation 219, 7151–7158 (2013).
"""
function assemble_flux_matrix(facets, nodes, weights = 1)
    nnode = size(nodes, 1)
    areas = compute_areas(facets, nodes)

    x = kron(ones(1, 3), facets)
    y = kron(facets, ones(1, 3))
    if length(weights) == 1 || length(weights) == size(facets, 1)
        # P0 weights
        z = kron(weights .* areas, reshape((ones(3, 3) + I(3)) / 12, 1, 9))
    else
        # P1 weights
        M1 = [6 2 2; 2 2 1; 2 1 2] / 60
        M2 = M1[[3, 1, 2], [3, 1, 2]]
        M3 = M2[[3, 1, 2], [3, 1, 2]]

        z = (
            kron(areas .* weights[facets[:, 1]], reshape(M1, 1, 9)) +
            kron(areas .* weights[facets[:, 2]], reshape(M2, 1, 9)) +
            kron(areas .* weights[facets[:, 3]], reshape(M3, 1, 9))
        )
    end

    # Enforce symmetry
    sym(x) = (x + x') / 2

    sym(sparse(x[:], y[:], z[:], nnode, nnode))
end

"""
    compute_areas(facets, nodes)

Compute facet areas.
"""
function compute_areas(facets, nodes)
    nfacet = size(facets, 1)

    v1 = nodes[facets[:, 2], :] - nodes[facets[:, 1], :]
    v2 = nodes[facets[:, 3], :] - nodes[facets[:, 1], :]

    matrix_3D = zeros(2, 2, size(facets, 1))
    matrix_3D[1, 1, :] = sum(v1 .* v1, dims = 2)
    matrix_3D[1, 2, :] = sum(v1 .* v2, dims = 2)
    matrix_3D[2, 1, :] = sum(v1 .* v2, dims = 2)
    matrix_3D[2, 2, :] = sum(v2 .* v2, dims = 2)

    [√abs(det(matrix_3D[:, :, i])) / 2 for i = 1:nfacet]
end


text:
using DifferentialEquations
using Plots; gr()
using Formatting

function popsize(n0,r,t)
    n0*exp(r*t)
end

function graph(dname;u0vals=[],rvals=[],tmax=10,deltat=0.1)
    if !isdir(dname)
     mkdir(dname)
end
fno = 1
tspan = 0.0:deltat:tmax
for u in u0vals
    for r in rvals
        psize = [popsize(u,r,t) for t in tspan]
            fname = format(dname*"/frame{:05d}.png",fno)
            fno = fno + 1
            plot(tspan,psize, linewidth=5,title="Exponential Growth",
                 xaxis="Time",yaxis="Population size",legend=false)
            savefig(fname)
        end
    end
end

graph("test";u0vals=[0:0.01:1], rvals=[1.1], tmax=10.0, deltat=0.1)


text:
using Markdown

Base.@kwdef struct MarkdownBlock
    block::Block
    text::Markdown.MD
end

function draw(d::MarkdownBlock, area::Rect, buf::Buffer)
    draw(d.block, area, buf);
    text_area = inner(d.block, area)

    title_area = Rect(left(area), top(area), width(area), height(area) ÷ 5)
    # Assume first block is always the header
    draw(d.text.content[1], title_area, buf)

    body_area = Rect(left(area), bottom(title_area), width(area), height(area) - height(title_area))
    body_area = Rect(body_area, Margin(5, 5))
    items = d.text.content[2:end]
    inner_area = Rect(left(body_area), top(body_area) - height(body_area) ÷ length(items), width(body_area), height(body_area) ÷ length(items))
    for (i, c) in enumerate(d.text.content[2:end])
        inner_area = Rect(left(inner_area), bottom(inner_area), width(inner_area), height(inner_area))
        draw(c, inner_area, buf)
    end
end

function draw(h::Markdown.Header, area::Rect, buf::Buffer)
    b = Block(border_type = BorderTypeLight)
    draw(b, area, buf);

    title = h.text[1]

    row = top(area) + height(area) ÷ 2
    col = width(area) ÷ 2 - ( length(title) ÷ 2 )
    set(buf, col, row, title, Crayon(bold = true))
end

function draw(l::Markdown.List, area::Rect, buf::Buffer)

    _s = String[]
    for item in l.items
        iob = IOBuffer()
        print(iob, item[1])
        push!(_s, String(take!(iob)))
    end
    for (i, item) in enumerate(l.items)
        item_area = Rect(
            left(area) + maximum(length.(_s)) ÷ 2,
            top(area) + (i - 1) * (height(area) ÷ length(l.items)),
            width(area),
            height(area) ÷ length(l.items)
        )
        row = top(item_area)
        col = left(item_area)
        set(buf, col, row, " ● ", Crayon(bold = true, foreground = :red))
        inner_area = Rect(item_area, Margin(0, 5))
        draw(item[1], inner_area, buf)
    end

end

function draw(p::Markdown.Paragraph, area::Rect, buf::Buffer)

    row = top(area)
    col = left(area)
    for c in p.content
        if typeof(c) == String
            l = length(c)
            set(buf, col, row, c)
            col += l
        elseif typeof(c) == Markdown.Bold
            l = length(c.text[1])
            set(buf, col, row, c.text[1], Crayon(bold = true))
            col += l
        elseif typeof(c) == Markdown.Italic
            l = length(c.text[1])
            set(buf, col, row, c.text[1], Crayon(italics = true))
            col += l
        elseif typeof(c) == Markdown.Image
            draw(c, area, buf)
        end
    end

end


text:
using Test, ReinforcementLearningEnvironmentDiscrete, POMDPModels,
ReinforcementLearningBase

for env in [MDP(), MDPEnv(LegacyGridWorld()), POMDPEnv(TigerPOMDP()), DiscreteMaze()]
    test_envinterface(env)
end

env = CliffWalking()
reset!(env)
@test env.mdp.state == 1


text:
using Oceananigans.Diagnostics

function nan_checker_aborts_simulation(arch)
    grid = RegularCartesianGrid(size=(4, 2, 1), extent=(1, 1, 1))
    model = IncompressibleModel(grid=grid, architecture=arch)
    simulation = Simulation(model, Δt=1, stop_iteration=1)

    model.velocities.u[1, 1, 1] = NaN

    run!(simulation)

    return nothing
end

TestModel(::GPU, FT, ν=1.0, Δx=0.5) =
    IncompressibleModel(
          grid = RegularCartesianGrid(FT, size=(3, 3, 3), extent=(3Δx, 3Δx, 3Δx)),
       closure = IsotropicDiffusivity(FT, ν=ν, κ=ν),
  architecture = GPU(),
    float_type = FT
)

TestModel(::CPU, FT, ν=1.0, Δx=0.5) =
    IncompressibleModel(
          grid = RegularCartesianGrid(FT, size=(3, 3, 3), extent=(3Δx, 3Δx, 3Δx)),
       closure = IsotropicDiffusivity(FT, ν=ν, κ=ν),
  architecture = CPU(),
    float_type = FT
)

function max_abs_field_diagnostic_is_correct(arch, FT)
    model = TestModel(arch, FT)
    set!(model.velocities.u, rand(size(model.grid)))
    u_max = FieldMaximum(abs, model.velocities.u)
    return u_max(model) == maximum(abs, model.velocities.u.data.parent)
end

function advective_cfl_diagnostic_is_correct(arch, FT)
    model = TestModel(arch, FT)

    Δt = FT(1.3e-6)
    Δx = FT(model.grid.Δx)
    u₀ = FT(1.2)
    CFL_by_hand = Δt * u₀ / Δx

    model.velocities.u.data.parent .= u₀
    cfl = AdvectiveCFL(FT(Δt))

    return cfl(model) ≈ CFL_by_hand
end

function diffusive_cfl_diagnostic_is_correct(arch, FT)
    Δt = FT(1.3e-6)
    Δx = FT(0.5)
    ν = FT(1.2)
    CFL_by_hand = Δt * ν / Δx^2

    model = TestModel(arch, FT, ν, Δx)
    cfl = DiffusiveCFL(FT(Δt))

    return cfl(model) ≈ CFL_by_hand
end

get_iteration(model) = model.clock.iteration
get_time(model) = model.clock.time

function diagnostics_getindex(arch, FT)
    model = TestModel(arch, FT)
    simulation = Simulation(model, Δt=0, stop_iteration=0)
    nc = NaNChecker(model, schedule=IterationInterval(1), fields=model.velocities)
    simulation.diagnostics[:nc] = nc

    # The first diagnostic is the NaN checker.
    return simulation.diagnostics[2] == nc
end

function diagnostics_setindex(arch, FT)
    model = TestModel(arch, FT)
    simulation = Simulation(model, Δt=0, stop_iteration=0)

    nc1 = NaNChecker(model, schedule=IterationInterval(1), fields=model.velocities)
    nc2 = NaNChecker(model, schedule=IterationInterval(2), fields=model.velocities)
    nc3 = NaNChecker(model, schedule=IterationInterval(3), fields=model.velocities)

    push!(simulation.diagnostics, nc1, nc2)
    simulation.diagnostics[2] = nc3

    return simulation.diagnostics[:diag2] == nc3
end

@testset "Diagnostics" begin
    @info "Testing diagnostics..."

    for arch in archs
        @testset "NaN Checker [$(typeof(arch))]" begin
            @info "  Testing NaN Checker [$(typeof(arch))]"
            @test_throws ErrorException nan_checker_aborts_simulation(arch)
        end
    end

    for arch in archs
        @testset "Miscellaneous timeseries diagnostics [$(typeof(arch))]" begin
            @info "  Testing miscellaneous timeseries diagnostics [$(typeof(arch))]"
            for FT in float_types
                @test diffusive_cfl_diagnostic_is_correct(arch, FT)
                @test advective_cfl_diagnostic_is_correct(arch, FT)
                @test max_abs_field_diagnostic_is_correct(arch, FT)
                @test diagnostics_getindex(arch, FT)
                @test diagnostics_setindex(arch, FT)
            end
        end
    end
end


text:
"""
Contains information about a subproblem in the branch and bound scheme.
`I` is a set of schools that are "in" the portfolio and `N` is the set that
are "negotiable" and used to generate an LP upper bound and child nodes.
"""
mutable struct Node
    I::Vector{Int}
    N::Vector{Int}
    t_bar::Vector{Float64}
    H_bar::Int
    v_I::Float64
    v_LP::Float64

    "Generate a new node with its LP relaxation value and empty child set."
    function Node(
        I::Vector{Int},
        N::Vector{Int},
        t_bar::Vector{Float64},
        H_bar::Int,
        v_I::Float64,
        mkt::VariedCostsMarket
    )
        # Leaf node
        isempty(N) && new(I, N, t_bar, H_bar, v_I, v_I)

        sort!(
            N,
            by=function (j)
                mkt.f[j] * t_bar[j] / mkt.g[j]
            end,
            rev=true,
            alg=InsertionSort # Typically just one item is out of place
        )

        v_LP = v_I
        H_left = H_bar
        @inbounds for j in N
            if mkt.g[j] ≤ H_left
                H_left -= mkt.g[j]
                v_LP += mkt.f[j] * t_bar[j]
            else
                v_LP += mkt.f[j] * t_bar[j] * H_left / mkt.g[j]
                break
            end
        end

        return new(I, N, t_bar, H_bar, v_I, v_LP)
    end
end

isless(n::Node, m::Node) = isless(n.v_LP, m.v_LP)

"""
    generatechildren(nd, mkt)

With respect to the data in `mkt`, generates the child(ren) of node `nd`.
"""
function generatechildren(nd::Node, mkt::VariedCostsMarket)::Vector{Node}
    newN = filter(j -> mkt.g[j] ≤ nd.H_bar, nd.N)

    # No way to add any school to this: just skip to the leaf node
    if isempty(newN)
        child = Node(nd.I, Int[], nd.t_bar, nd.H_bar, nd.v_I, mkt)
        return [child]
    end

    # School we will branch on. In principle it can by any school in newN.
    i = popfirst!(newN) # = argmax(j -> mkt.f[j] * nd.t_bar[j], newN) by sortation policy

    # Node without i
    child2 = Node(nd.I, newN, nd.t_bar, nd.H_bar, nd.v_I, mkt)

    if mkt.g[i] < nd.H_bar
        # Node with i
        t_bar1 = copy(nd.t_bar)
        @inbounds for j in newN
            if t_bar1[j] ≤ nd.t_bar[i]
                t_bar1[j] *= mkt.omf[i]
            else
                t_bar1[j] -= mkt.f[i] * nd.t_bar[i]
            end
        end

        child1 = Node(
            vcat(nd.I, i),
            newN,
            t_bar1,
            nd.H_bar - mkt.g[i],
            nd.v_I + mkt.f[i] * nd.t_bar[i],
            mkt
        )

        return [child1, child2]
    else
        child1 = Node(
            vcat(nd.I, i),
            Int[],
            Float64[], # nd.t_bar
            0,
            nd.v_I + mkt.f[i] * nd.t_bar[i],
            mkt
        )

        return [child1, child2]
    end
end

"""
    optimalportfolio_branchbound(mkt::VariedCostsMarket; maxit=10000, verbose=false) -> X, v

Use the branch-and-bound algorithm to produce the optimal portfolio for the `VariedCostsMarket`
defined by `mkt`. Intractable for large markets.

```julia-repl
julia> mkt = VariedCostsMarket([0.2, 0.5, 0.1, 0.6, 0.1], [1, 4, 9, 1, 8], [2, 4, 2, 5, 1], 8);

julia> optimalportfolio_branchbound(mkt)
([2, 3, 5], 3.24)
```
"""
function optimalportfolio_branchbound(mkt::VariedCostsMarket; maxit::Int=100000, verbose::Bool=false)::Tuple{Vector{Int},Float64}
    mkt.m ≥ 33 && @warn "Branch and bound is slow for large markets"

    C = collect(1:mkt.m)

    rootnode::Node = Node(Int[], C, Float64.(mkt.t), mkt.H, 0.0, mkt)
    LB::Float64 = 0.0
    LB_node::Node = rootnode
    tree = BinaryMaxHeap{Node}([rootnode])

    for k in 1:maxit
        verbose && @show k, LB, length(tree)

        # All branches have either reached leaves or fathomed: done
        isempty(tree) && return mkt.perm[collect(LB_node.I)], LB

        # Inspect the node with the highest UB
        children::Vector{Node} = generatechildren(pop!(tree), mkt)

        for child in children
            if child.v_LP > LB
                if child.v_I > LB
                    LB_node = child
                    LB = child.v_I
                end
                isempty(child.N) || push!(tree, child)
            end
        end
    end

    @warn "Unable to find optimum in $maxit iterations; returning best so far.\n" *
          "         Worst-case optimality ratio: $(LB/rootnode.v_LP)"
    return mkt.perm[collect(LB_node.I)], LB
end

max_stars_repo_name: maxkapur/OptimalApplication

max_issues_repo_name: maxkapur/OptimalApplication

max_forks_repo_name: maxkapur/OptimalApplication


text:
# example forms
#   epsilon = < {1,1/2,1/4,...} | {} >
#   omega = < {0,1,2,3,...} | {} >
#
#   could do non-dyadic rationals, e.g., 1/3 by truncating binary approximations
#

#
# the idea here is to build the infinite bits of these using lazy evaluation
#       sim. to lazy indexing into a sequence
#           https://github.com/MikeInnes/Lazy.jl
#           http://mikeinnes.github.io/2017/06/05/lazy.html
#
# but we'd still have to be able to compare sup and inf such sequences to make it all work
# so maybe there is a better mathematical way to specify the sequences such that this is easy
#

struct SurrealTrans <: Surreal
    shorthand::String
    X_L::SurrealSequence
    X_R::SurrealSequence
    # constructor should check that X_L <= X_R
    function SurrealTrans(name::String, X_L::SurrealSequence, X_R::SurrealSequence)
        if X_L < X_R
            return new(name, X_L, X_R)
        else
            error("Must have X_L <= X_R")
        end
    end
end

struct SurrealSequence <: Function
#  provide a map from an index to the surreal number
#     but its a bit restricted
#        i.e., for inf sequence make it a map from Naturals -> SurrealFinite
#              but for finite or empty sequences do something else?
#  seem similar to indexing
end

#
# ω = SurrealTrans("ω", x->x,   x->ϕ)
# ε = SurrealTrans("ε", x->1/x, x->ϕ)
#

# do + using
#  x+y = < X^L + Y^L | X^R + Y^R >
#
# where
#    X + Y = {  x_i + y_j | for all x_i in X, and y_j in Y }
# so
#
#    [X+Y]_I = x_i + y_j
# where we take indexes (i,j) from table
#              j
#      1 3 6 10
#      2 5 9
#    i 4 8
#      7
#
# and so in (the entry in the table is I)
#


text:
@testset "Embeddings" begin
  Qx, x = QQ["x"]
  L, b = NumberField([x^2 - 2, x^2 - 3, x^3 + 2], "a")
  r, s = @inferred signature(L)
  @test (r, s) == (4, 4)
  @test complex_embeddings(L) === complex_embeddings(L)

  emb = @inferred complex_embeddings(L)
  @test (@inferred number_field(emb[1])) === L
  @test emb[1] == emb[1]
  @test emb[1] != emb[2]

  # isreal
  @test length(real_embeddings(L)) == 4
  @test @inferred isreal(emb[1])
  @test count(isreal, emb) == 4
  @test count(!isreal, emb) == 8
  @test count(isimaginary, emb) == 8

  for e in emb
    if isreal(e)
      @test (@inferred conj(e)) == e
    else
      @test conj(e) in emb && conj(e) != e
    end
  end

  # print
  sprint(show, emb[1])
  sprint(show, MIME"text/plain"(), emb[1])

  # compare with absolute simple field
  Labs, LabstoL = @inferred absolute_simple_field(L)
  embabs = complex_embeddings(Labs)
  binLabs = [ LabstoL\bb for bb in b]
  @inferred emb[1](b[1])

  for e in emb
    imageofb = e.(b)
    @test all(overlaps.(conj(e).(b), conj.(e.(b))))
    @test count(ee -> all(overlaps.(ee.(binLabs), imageofb)), embabs) == 1
    @test all(Hecke.radiuslttwopower.(e.(b, 128), -128))
  end

  # restrict along a morphism
  LtoLabs = inv(LabstoL)
  for e in embabs
    ee = @inferred restrict(e, LtoLabs)
    @test overlaps(ee(LabstoL(gen(Labs))), e(gen(Labs)))
    @test ee == LtoLabs * e
    @test_throws ArgumentError restrict(e, LabstoL)
  end

  for e in embabs
    ee = @inferred extend(e, LabstoL)
    @test length(ee) == 1
    @test LabstoL * ee[1] == e
  end

  for e in emb
    ee = @inferred extend(e, LtoLabs)
    @test length(ee) == 1
    @test LtoLabs * ee[1] == e
  end

  z = rand(L, -10:10)^100
  for p in [32, 128, 256, 1024]
    for e in emb
      fn = evaluation_function(e, p)
      @test Hecke.radiuslttwopower(fn(z), -p)
      c = e(z, p)
      @test Hecke.radiuslttwopower(c, -p)
    end
  end
end


text:
#TODO: transverse_mode_1: give transverse mode definite sign so it doesn't change from call to call
#TODO: incident_pc, pc_transverse_field: higher gap-order modes
#TODO: higher order gaps in build_channels!
#TODO: fix total kluge in waveguide_dispersion which eliminates bleeding of bands into gap

################################################################################
### SOURCE SYNTHESIS
################################################################################
"""
    j, φ₊₋, φ₊, φ₋ = synthesize_source(sim, k, a)

`a[n]` complex amplitude associated with channel `n`
`j` source for scattering calculation
`φ₊[:,m]` is incident wave on waveguide `m`
`φ₋[:,m]` is outgoing wave on waveguide `m`
`φ₊₋` is outgoing + incident wave at disordered sites
"""
function synthesize_source(sim::Simulation, k, a)

    N = prod(sim.dis.N)
    W = length(sim.sct.waveguides_used)

    φ₊₋ = zeros(ComplexF64,N,W)
    φ₊ = zeros(ComplexF64,N,W)  # incident
    φ₋ = zeros(ComplexF64,N,W)  # outgoing

    φt₊₋ = zeros(ComplexF64,N) # incident
    φt₊ = zeros(ComplexF64,N) # incident
    φt₋ = zeros(ComplexF64,N) # outgoing

    for m ∈ 1:length(sim.sct.channels)
        ind = findall(sim.sct.channels[m].waveguide .== sim.sct.waveguides_used)
        φt₊₋, φt₊, φt₋ = ScalarFDFD.incident_mode(sim, k, m)
        φ₊₋[:,ind] += a[m]*φt₊₋
        φ₊[:,ind] += a[m]*φt₊
        φ₋[:,ind] += a[m]*φt₋
    end

    k² = k^2
    k²δ = Array{ComplexF64}(undef,N)
    j = zeros(ComplexF64,N)
    for w ∈ 1:W
        φ₊₋[:,w] = (sim.sys.ε[:] .!== sim.sct.ε₀[w][:]).*φ₊₋[:,w]
        φ₊[:,w] = (sim.sys.ε[:] .== sim.sct.ε₀[w][:]).*φ₊[:,w]
        φ₋[:,w] = (sim.sys.ε[:] .== sim.sct.ε₀[w][:]).*φ₋[:,w]
        k²δ[:] = k²*(sim.sys.ε[:]-sim.sct.ε₀[w][:])
        j += -k²δ.*φ₊₋[:,w]
    end

    return j, sum(φ₊₋, dims=2), sum(φ₊, dims=2), sum(φ₋, dims=2)
end

"""
    φ₊₋, φ₊, φ₋ = incident_mode(sim, k, m)
"""
function incident_mode(sim::Simulation, k, m)
    if isempty(sim.sys.domains[ScalarFDFD.get_waveguide_domains(sim, sim.sct.channels[m].waveguide)])
        φ₊₋, φ₊, φ₋ = incident_free_space_mode(sim, k, m, :in)
    elseif ScalarFDFD.ispc(sim.sys.domains[ScalarFDFD.get_waveguide_domains(sim, sim.sct.channels[m].waveguide)][end])
        φ₊₋, φ₊, φ₋ = incident_pc_mode(sim, k, m, :in)
    elseif ishalfspace(sim.sys.domains[ScalarFDFD.get_waveguide_domains(sim, sim.sct.channels[m].waveguide)][1])
        φ₊₋, φ₊, φ₋ = ScalarFDFD.incident_halfspace_mode(sim, k, m, :in)
    elseif  ScalarFDFD.isplanar(sim.sys.domains[ScalarFDFD.get_waveguide_domains(sim, sim.sct.channels[m].waveguide)][1])
        φ₊₋, φ₊, φ₋ = ScalarFDFD.incident_planar_mode(sim, k, m, :in)
    else
        throw(ArgumentError("unrecognized waveguide type"))
    end

    return φ₊₋, φ₊, φ₋
end

################################################################################
### FREE SPACE MODES (ANGULAR MOMENTUM)
################################################################################
"""
    φ₊₋, φ₊, φ₋ = incident_free_space_mode(sim, k, m)
"""
function incident_free_space_mode(sim::Simulation, k, m, direction)

    x = sim.dis.x[1]
    y = sim.dis.x[2]

    R = hypot.(x,y)
    θ = atan.(x,y)

    φ₊ = exp(1im*m*θ).*hankelh2(m,k*R)/2
    φ₋ = exp(1im*m*θ).*hankelh1(m,k*R)/2
    φ₊₋ = exp(1im*m*θ).*besselj(m,k*R)/4

    return φ₊₋, φ₊, φ₋
end

################################################################################
### PHOTONIC CRYSTAL WAVEGUIDE MODES
################################################################################
"""
    φ₊₋, φ₊, φ₋, φ_interpolation = incident_pc_mode(sim, k, m)
"""
function incident_pc_mode(sim::Simulation, k, m, direction)

    β, ψ, wg_sim = ScalarFDFD.pc_transverse_field(sim, k, m, direction)

    ψ = ψ.*exp.(-complex.(0,wg_sim.dis.x[1])*β[1] .- complex.(0,wg_sim.dis.x[2])*β[2])[:]
    xs = wg_sim.dis.x[1][1] .+ (0:wg_sim.dis.N[1]-1)*wg_sim.dis.dx[1]
    ys = wg_sim.dis.x[2][1] .+ (0:wg_sim.dis.N[2]-1)*wg_sim.dis.dx[2]
    utp = CubicSplineInterpolation((xs,ys),reshape(ψ[:,1],wg_sim.dis.N[1],wg_sim.dis.N[2]), bc=Periodic(OnCell()), extrapolation_bc=Periodic(OnCell()))

    X = repeat(sim.dis.x[1],1,sim.dis.N[2])
    Y = repeat(sim.dis.x[2],sim.dis.N[1],1)

    v1 = wg_sim.lat.v1
    v2 = wg_sim.lat.v2

    φ = utp.(X,Y)
    φ₊ = φ.*exp.(1im*(β[1]*(v1[1]*X + v1[2]*Y) + β[2]*(v2[1]*X + v2[2]*Y)))
    φ₋ = zeros(ComplexF64, size(φ₊))
    φ₊₋ = φ₊ + φ₋

    return φ₊₋[:], φ₊[:], φ₋[:]
end

"""
    pc_transverse_field(sim, k, m)
"""
function pc_transverse_field(sim::Simulation, k, m, direction)

    waveguide = sim.sct.channels[m].waveguide
    wg_sim = extract_waveguide_simulation(sim, waveguide)
    β = [0.,0.]
    if sim.sys.domains[ScalarFDFD.get_waveguide_domains(sim, waveguide)][1].which_asymptote == :bottom
        if direction == :in
            sol = optimize(x->(sim.sct.channels[m].dispersion[1](x)-float(k))^2, 0, +π/wg_sim.lat.b)
            β[2] = (sol.minimizer[1])::Float64
        else
            sol = optimize(x->(sim.sct.channels[m].dispersion[1](x)-float(k))^2, -π/wg_sim.lat.b, 0)
            β[2] = (sol.minimizer[1])::Float64
        end
    elseif sim.sys.domains[ScalarFDFD.get_waveguide_domains(sim, waveguide)][1].which_asymptote == :top
        if direction == :in
            sol = optimize(x->(sim.sct.channels[m].dispersion[1](x)-float(k))^2, -π/wg_sim.lat.b, 0)
            β[2] = (sol.minimizer[1])::Float64
        else
            sol = optimize(x->(sim.sct.channels[m].dispersion[1](x)-float(k))^2, 0, +π/wg_sim.lat.b)
            β[2] = (sol.minimizer[1])::Float64
        end
    elseif sim.sys.domains[ScalarFDFD.get_waveguide_domains(sim, waveguide)][1].which_asymptote == :left
        if direction == :in
            sol = optimize(x->(sim.sct.channels[m].dispersion[1](x)-float(k))^2, 0, +π/wg_sim.lat.a)
            β[1] = (sol.minimizer[1])::Float64
        else
            sol = optimize(x->(sim.sct.channels[m].dispersion[1](x)-float(k))^2, -π/wg_sim.lat.a, 0)
            β[1] = (sol.minimizer[1])::Float64
        end
    else
        if direction == :in
            sol = optimize(x->(sim.sct.channels[m].dispersion[1](x)-float(k))^2, -π/wg_sim.lat.a, 0)
            β[1] = (sol.minimizer[1])::Float64
        else
            sol = optimize(x->(sim.sct.channels[m].dispersion[1](x)-float(k))^2, 0, +π/wg_sim.lat.a)
            β[1] = (sol.minimizer[1])::Float64
        end
    end
    if !Optim.converged(sol)
        return NaN, zeros(ComplexF64, prod(wg_sim.dis.N), 1), wg_sim
    end
    if sim.sys.domains[ScalarFDFD.get_waveguide_domains(sim, waveguide)][1].which_asymptote ∈ [:bottom, :top]
        f(x) = abs2.(eig_k(wg_sim, k, 1; kb=x[1])[1][1]-k)
        a = optimize(f,[β[2]],BFGS(), Optim.Options(g_tol=1e-10, f_tol=1e-10))
        β[2] = a.minimizer[1]::Float64
        k_new, ψ = eig_k(wg_sim, k, 1; kb=β[2])
        if sim.sys.domains[ScalarFDFD.get_waveguide_domains(sim, waveguide)][1].which_asymptote == :top
            𝒩² = surface_flux(wg_sim, ψ)[2][4][1]
        else
            𝒩² = surface_flux(wg_sim, ψ)[2][3][1]
        end
    else
        g(x) = abs2.(eig_k(wg_sim, k, 1; ka=x[1])[1][1]-k)
        a = optimize(g,[β[1]],BFGS(), Optim.Options(g_tol=1e-10, f_tol=1e-10))
        β[1] = a.minimizer[1]::Float64
        k_new, ψ = eig_k(wg_sim, k, 1; ka=β[1])
        if sim.sys.domains[ScalarFDFD.get_waveguide_domains(sim, waveguide)][1].which_asymptote == :right
            𝒩² = surface_flux(wg_sim, ψ)[2][2][1]
        else
            𝒩² = surface_flux(wg_sim, ψ)[2][1][1]
        end
    end
    if !isapprox(k_new[1],k; atol=1e-2)
        @warn "computed k $(k_new[1]) not consistent with precomputed dispersion $k."
    end
    gauge_fixing = exp(-complex(0,angle(ψ[findmax(abs.(ψ))[2]])))
    ψ = ψ*gauge_fixing/sqrt(abs(𝒩²))
    return β, ψ, wg_sim
end

################################################################################
### HALFSPACE WAVEGUIDES MODES
################################################################################
"""
    φ₊₋, φ₊, φ₋ = incident_halfspace_mode(sim, k, m)
"""
function incident_halfspace_mode(sim::Simulation, k, m, direction::Symbol)

    c, prop_const = ScalarFDFD.halfspace_transverse_field(sim, k, m)

    if sim.dis.N[1]==1
        transverse_dim = 1
        longitudinal_dim = 2
    elseif sim.dis.N[2]==1
        transverse_dim = 2
        longitudinal_dim = 1
    else
        throw(ArgumentError("halfspace incident mode will not work for 2d"))
    end

    xs = sim.dis.x[longitudinal_dim][1] .+ (1:sim.dis.N[longitudinal_dim])*sim.dis.dx[longitudinal_dim]

    if ScalarFDFD.get_asymptote(sim, sim.sct.channels[m].waveguide) ∈ [:left]
        β = +prop_const
    elseif ScalarFDFD.get_asymptote(sim, sim.sct.channels[m].waveguide) ∈ [:right]
        β = -prop_const
    elseif ScalarFDFD.get_asymptote(sim, sim.sct.channels[m].waveguide) ∈ [:bottom]
        β = +prop_const
    elseif ScalarFDFD.get_asymptote(sim, sim.sct.channels[m].waveguide) ∈ [:top]
        β = -prop_const
    end

    if issubset(sim.bnd.bl[:,longitudinal_dim],[:pml_out, :abs_out])
        φ₊ = exp.(1im*β*xs)/sqrt(abs(real(β)))
        φ₋ = zeros(ComplexF64, size(φ₊))
    elseif issubset([:none],sim.bnd.bl[:,longitudinal_dim])
        x0 = sim.bnd.bc[2,longitudinal_dim]
        if sim.bnd.bc[2,longitudinal_dim] == :d
            φ₊ = exp.(1im*β*(xs-x0))/sqrt(abs(real(β)))
            φ₋ = -1im*exp.(-1im*β*(xs-x0))/sqrt(abs(real(β)))/2
        elseif sim.bnd.bc[2,longitudinal_dim] == :n
            φ₊ = exp.(1im*β*(xs-x0))/sqrt(abs(real(β)))
            φ₋ = +1im*exp.(-1im*β*(xs-x0))/sqrt(abs(real(β)))/2
        else
            throw(ErrorException("for now one-sided scattering can only be done from the left or bottom"))
        end
    end
    φ₊₋ = φ₊ + φ₋

    return φ₊₋[:], φ₊[:], φ₋[:]
end

"""
    1, β = halfspace_transverse_field(sim, k, m)
"""
function halfspace_transverse_field(sim::Simulation, k, m)
    prop_const = (optimize(x->(sim.sct.channels[m].dispersion[1](x[1])-float(k))^2, [float(k)], BFGS()).minimizer[1])::Float64
    return 1, prop_const
end

################################################################################
### PLANAR WAVEGUIDE MODES
################################################################################
"""
    incident_planar_mode(sim, k, m)
"""
function incident_planar_mode(sim::Simulation, k, m)

    β, u_itp, longitudinal_dim, transverse_dim = ScalarFDFD.planar_transverse_field(sim, k, m)

    if k > β
        @warn "mode $(sim.sct.channels[m].quantum_number) on waveguide $(sim.sct.channels[m].waveguide) is not guided at frequency $k"
    end

    waveguide = sim.sct.channels[m].waveguide
    asymptote = ScalarFDFD.get_asymptote(sim, waveguide)

    if issubset(sim.bnd.bl[:,longitudinal_dim],[:pml_out, :abs_out])
        if asymptote ∈ [:left, :bottom]
            φ₊ = u_itp.(sim.dis.x[transverse_dim]).*exp.(complex(0,+β)*sim.dis.x[longitudinal_dim])/sqrt(β)
        elseif asymptote ∈ [:right, :top]
            φ₊ = u_itp.(sim.dis.x[transverse_dim]).*exp.(complex(0,-β)*sim.dis.x[longitudinal_dim])/sqrt(β)
        else
            throw(ErrorException("invalid aymptote $(asymptote)"))
        end
        φ₋ = zeros(ComplexF64, size(φ₊))
        φ₊₋ = copy(φ₊)
    elseif issubset([:none],sim.bnd.bl[:,longitudinal_dim])
        if asymptote ∈ [:left, :bottom]
            if sim.bnd.bc[2,longitudinal_dim] == :d
                φ₊ = +u_itp.(sim.dis.x[transverse_dim]).*exp.(complex(0,+β)*(sim.dis.x[longitudinal_dim] .- sim.dis.x[longitudinal_dim][end]))/sqrt(β)
                φ₋ = -u_itp.(sim.dis.x[transverse_dim]).*exp.(complex(0,-β)*(sim.dis.x[longitudinal_dim] .- sim.dis.x[longitudinal_dim][end]))/sqrt(β)
            elseif sim.bnd.bc[2,longitudinal_dim] == :n
                φ₊ = +u_itp.(sim.dis.x[transverse_dim]).*exp.(complex(0,+β)*(sim.dis.x[longitudinal_dim] .- sim.dis.x[longitudinal_dim][end]))/sqrt(β)
                φ₋ = +u_itp.(sim.dis.x[transverse_dim]).*exp.(complex(0,-β)*(sim.dis.x[longitudinal_dim] .- sim.dis.x[longitudinal_dim][end]))/sqrt(β)
            end
        else
            throw(ErrorException("for now one-sided scattering can only be done from the left or bottom"))
        end
    else
        throw(ArgumentError("invalid boundary layer $(sim.bnd.bl)"))
    end
    φ₊₋ = φ₊ + φ₋

    return φ₊₋[:], φ₊[:], φ₋[:]
end

#@code_warntype checked
"""
    β, u_itp, longitudinal_dim, transverse_dim = planar_transverse_field(sim, k, m; η_init=-2, num_modes_multiplier=3)
"""
function planar_transverse_field(sim::Simulation, k, m)

    β = (optimize(x->(sim.sct.channels[m].dispersion[1](x[1])-float(k))^2, [float(k)], BFGS()).minimizer[1])::Float64

    wg_sim = extract_waveguide_simulation(sim, sim.sct.channels[m].waveguide)

    k_new, ψ = ScalarFDFD.planar_kl(wg_sim, β, k, 1)
    𝒩² = quadrature(wg_sim, abs2.(ψ[:]))
    gauge_fixing = exp(-complex(0,angle(ψ[findmax(abs.(ψ))[2]])))
    ψ = ψ*gauge_fixing/sqrt(𝒩²)
    if !isapprox(k_new[1],k; atol=1e-2)
        @warn "computed k $(k_new[1]) not consistent with precomputed dispersion $k."
    end

    if wg_sim.dis.N[1]==1
        transverse_dim = 2
        longitudinal_dim=1
    else
        transverse_dim=1
        longitudinal_dim=2
    end
    xs = sim.dis.x[transverse_dim][1] .+ (1:sim.dis.N[transverse_dim])*sim.dis.dx[transverse_dim]
    utp = CubicSplineInterpolation(xs,real(ψ[:,1]), extrapolation_bc=Line())

    return β, utp, longitudinal_dim, transverse_dim
end

################################################################################
### AUXILLIARIES
################################################################################
"""
    waveguide_domains = get_waveguide_domains(sim::Simulation, waveguide)
"""
function get_waveguide_domains(sim::Simulation, waveguide)
    return findall([sim.sys.domains[i].which_waveguide for i ∈ eachindex(sim.sys.domains)] .== waveguide)
end

"""
    which_asymptote = get_asymptote(sim, waveguide, wg_inds)
"""
function get_asymptote(sim::Simulation, waveguide, wg_inds=get_waveguide_domains(sim, waveguide))
    return sim.sys.domains[wg_inds[1]].which_asymptote
end

"""
    wg_sim = extract_waveguide_simulation(sim, waveguide)
"""
function extract_waveguide_simulation(sim::Simulation, waveguide::Int)

    wg_inds = ScalarFDFD.get_waveguide_domains(sim, waveguide)

    if isempty(sim.sys.domains[wg_inds])
        throw(ErrorException("haven't done free space waveguide extraction yet"))
    elseif all(isplanar.(sim.sys.domains[wg_inds]))
        wg_sim = extract_planar_waveguide_simulation(sim, waveguide)
    elseif all(ispc.(sim.sys.domains[wg_inds]))
        wg_sim = extract_pc_waveguide_simulation(sim, waveguide)
    elseif all(ishalfspace.(sim.sys.domains[wg_inds]))
        throw(ErrorException("haven't done halfplane waveguide extraction yet"))
    else
        throw(ErrorException("invalid waveguide type"))
    end

    return deepcopy(wg_sim)
end

"""
    wg_sim = extract_planar_waveguide_simulation(sim, waveguide)
"""
function extract_planar_waveguide_simulation(sim::Simulation, waveguide::Int)

        wg_inds = ScalarFDFD.get_waveguide_domains(sim, waveguide)
        which_asymptote = ScalarFDFD.get_asymptote(sim, waveguide)

        domains = Array{Domain}(undef,length(wg_inds))
        for i ∈ eachindex(wg_inds)
              domains[i] = Domain(sim.sys.domains[wg_inds[i]]; :which_asymptote => :none, :which_waveguide => 0)
        end
        waveguide = System(domains)

        sim = deepcopy(sim)

        if which_asymptote == :right
            sim.bnd.∂Ω[1,1] = sim.bnd.∂Ω_tr[2,1]
            sim.dis.dx[1] = Inf
        elseif which_asymptote == :left
            sim.bnd.∂Ω[1,1] = sim.bnd.∂Ω_tr[1,1]
            sim.dis.dx[1] = Inf
        elseif which_asymptote == :bottom
            sim.bnd.∂Ω[1,2] = sim.bnd.∂Ω_tr[1,2]
            sim.dis.dx[2] = Inf
        elseif which_asymptote == :top
            sim.bnd.∂Ω[1,2] = sim.bnd.∂Ω_tr[2,2]
            sim.dis.dx[2] = Inf
        else
            throw(ArgumentError("Invalid asymptotic region $(which_asymptote)."))
        end
        sim.bnd.bl[:] .= :none
        sim.bnd.bl_depth[:] .= 0
        sim.bnd.bc[:] .= :p

        wg_sim = Simulation(sys=waveguide, dis=sim.dis, bnd=Boundary(sim.bnd))

        return wg_sim
end

"""
    wg_sim = extract_pc_waveguide_simulation(sim, waveguide)
"""
function extract_pc_waveguide_simulation(sim::Simulation, waveguide::Int)

    sim = deepcopy(sim)

    bnd = sim.bnd
    sys = sim.sys

    wg_inds = ScalarFDFD.get_waveguide_domains(sim, waveguide)
    which_asymptote = ScalarFDFD.get_asymptote(sim, waveguide, wg_inds)

    domains = Array{Domain}(undef,length(wg_inds))
    for i ∈ eachindex(wg_inds)
          domains[i] = Domain(sys.domains[wg_inds[i]]; :which_asymptote => :none, :which_waveguide => 0)
    end
    waveguide = System(domains)

    wg_inds = findall([waveguide.domains[i].domain_type for i ∈ eachindex(waveguide.domains)] .== :pc_waveguide)
    pc_inds = findall([waveguide.domains[i].domain_type for i ∈ eachindex(waveguide.domains)] .== :pc_waveguide_background)
    if isempty(wg_inds)
        ind = 1
    else
        ind = wg_inds[1]
    end
    if which_asymptote == :right
        waveguide_lattice = Bravais(waveguide.domains[ind].lattice; :x0 => bnd.∂Ω_tr[2,1], :y0 => bnd.∂Ω[1,2], :b => (bnd.∂Ω[2,2]-bnd.∂Ω[1,2])/sin(waveguide.domains[ind].lattice.β))
        bnd.∂Ω[1,1] = bnd.∂Ω_tr[2,1]
        bnd.∂Ω[2,1] = bnd.∂Ω[1,1] + waveguide_lattice.a*waveguide_lattice.v1[1]
    elseif which_asymptote == :left
        waveguide_lattice = Bravais(waveguide.domains[ind].lattice; :x0 => bnd.∂Ω[1,1], :y0 => bnd.∂Ω[1,2], :b => (bnd.∂Ω[2,2]-bnd.∂Ω[1,2])/sin(waveguide.domains[ind].lattice.β))
        bnd.∂Ω[2,1] = bnd.∂Ω_tr[1,1]
        bnd.∂Ω[1,1] = bnd.∂Ω[2,1] - waveguide_lattice.a*waveguide_lattice.v1[1]
    elseif which_asymptote == :bottom
        waveguide_lattice = Bravais(waveguide.domains[ind].lattice; :x0 => bnd.∂Ω[1,1], :y0 => bnd.∂Ω[1,2], :a => (bnd.∂Ω[2,1]-bnd.∂Ω[1,1])/cos(waveguide.domains[ind].lattice.α))
        bnd.∂Ω[2,2] = bnd.∂Ω_tr[1,2]
        bnd.∂Ω[1,2] = bnd.∂Ω[2,2] - waveguide_lattice.b*waveguide_lattice.v2[2]
    elseif which_asymptote == :top
        waveguide_lattice = Bravais(waveguide.domains[ind].lattice; :x0 => bnd.∂Ω[1,1], :y0 => bnd.∂Ω[2,2], :a => (bnd.∂Ω[2,1]-bnd.∂Ω[1,1])/cos(waveguide.domains[ind].lattice.α))
        bnd.∂Ω[1,2] = bnd.∂Ω_tr[2,2]
        bnd.∂Ω[2,2] = bnd.∂Ω[1,2] + waveguide_lattice.b*waveguide_lattice.v2[2]
    end
    if which_asymptote ∈ [:left, :right]
        bnd.bc[:,1] .= :p
        bnd.bl[:,1] .= :none
        bnd.bl_depth[:,1] .= 0
    elseif which_asymptote ∈ [:bottom, :top]
        bnd.bc[:,2] .= :p
        bnd.bl[:,2] .= :none
        bnd.bl_depth[:,2] .= 0
    else
        throw(ArgumentError("Invalid asymptotic region $(which_asymptote)."))
    end

    bnd.bc[:] .= :p
    bnd.bl[:] .= :none
    bnd.bl_depth[:] .= 0
    wg_sim = Simulation(sys=waveguide, dis=sim.dis, bnd=Boundary(bnd), lat=waveguide_lattice)

    return wg_sim
end


text:
using Compat
using Compat.Test: @test, @test_broken, @testset

makepath = joinpath(dirname(@__FILE__), "..", "deps")

if !isfile(joinpath(makepath, "Makefile"))
    start = pwd()
    cd(makepath)
    if Compat.Sys.iswindows()
        run(`cmake -DCMAKE_IGNORE_PATH="C:/Program Files/Git/usr/bin" -G "MinGW Makefiles" .`)
    elseif Compat.Sys.isapple() || Compat.Sys.islinux()
        run(`cmake .`)
    else
        error("This system is not supported by AndersonMoore")
    end

    run(`make`)
    cd(start)
end

const AndersonMooreModule = joinpath(dirname(@__FILE__), "..", "src", "AndersonMoore.jl")
include(AndersonMooreModule)

# test_path = joinpath(dirname(@__FILE__), "defineShiftRightTestFuncs.jl")
include(joinpath(dirname(@__FILE__), "defineShiftRightTestFuncs.jl"))
@testset "outer"  begin# an outer so that it does't quit on first fail
@testset "test shiftRight" begin
@test ShiftRightTests.firmvalue()
@test ShiftRightTests.firmvalue3Leads2Lags()
@test ShiftRightTests.example7()
@test ShiftRightTests.oneEquationNoLead()
@test ShiftRightTests.reliablePaperExmpl()
@test ShiftRightTests.athan()
@test ShiftRightTests.habitmod()
end

# test_path = joinpath(dirname(@__FILE__), "defineExactShiftTestFuncs.jl")
include(joinpath(dirname(@__FILE__), "defineExactShiftTestFuncs.jl"))
@testset "test exactShift" begin
@test ExactShiftTests.firmvalue()
@test ExactShiftTests.firmvalue3Leads2Lags()
@test ExactShiftTests.example7()
@test ExactShiftTests.oneEquationNoLead()
@test ExactShiftTests.reliablePaperExmpl()
@test ExactShiftTests.athan()
@test ExactShiftTests.habitmod()
end

# test_path = joinpath(dirname(@__FILE__), "defineNumericShiftTestFuncs.jl")
include(joinpath(dirname(@__FILE__), "defineNumericShiftTestFuncs.jl"))
@testset "test numericShift" begin
@test NumericShiftTests.firmvalueTrue()
@test NumericShiftTests.firmvalue3Leads2LagsTrue()
@test NumericShiftTests.example7True()
@test NumericShiftTests.oneEquationNoLeadTrue()
@test NumericShiftTests.reliablePaperExmplTrue()
@test NumericShiftTests.athanTrue()
@test NumericShiftTests.habitmodTrue()

@test NumericShiftTests.firmvalueFalse()
@test NumericShiftTests.firmvalue3Leads2LagsFalse()
@test NumericShiftTests.example7False()
@test NumericShiftTests.oneEquationNoLeadFalse()
@test NumericShiftTests.reliablePaperExmplFalse()
@test NumericShiftTests.athanFalse()
@test NumericShiftTests.habitmodFalse()
end

# test_path = joinpath(dirname(@__FILE__), "defineBuildATestFuncs.jl")
include(joinpath(dirname(@__FILE__), "defineBuildATestFuncs.jl"))
@testset "test buildA" begin
@test BuildATests.firmvalueFalse()
@test BuildATests.firmvalue3Leads2LagsFalse()
@test BuildATests.example7False()
@test BuildATests.oneEquationNoLeadFalse()
@test BuildATests.reliablePaperExmplFalse()
@test BuildATests.athanFalse()
@test BuildATests.habitmodFalse()

@test BuildATests.firmvalueTrue()
@test BuildATests.firmvalue3Leads2LagsTrue()
@test BuildATests.example7True()
@test BuildATests.oneEquationNoLeadTrue()
@test BuildATests.reliablePaperExmplTrue()
@test BuildATests.athanTrue()
@test BuildATests.habitmodTrue()
end

# test_path = joinpath(dirname(@__FILE__), "defineEigenSysTestFuncs.jl")
include(joinpath(dirname(@__FILE__), "defineEigenSysTestFuncs.jl"))
@testset "test eigenSys" begin
@test EigenSysTests.firmvalue()
@test EigenSysTests.firmvalue3Leads2Lags()
@test EigenSysTests.example7()
@test EigenSysTests.oneEquationNoLead()
@test EigenSysTests.reliablePaperExmpl()
@test EigenSysTests.athan()
@test EigenSysTests.habitmod()
end

# test_path = joinpath(dirname(@__FILE__), "defineAugmentQTestFuncs.jl")
include(joinpath(dirname(@__FILE__), "defineAugmentQTestFuncs.jl"))
@testset "test augmentQ" begin
@test AugmentQTests.firmvalue()
@test AugmentQTests.firmvalue3Leads2Lags()
@test AugmentQTests.example7()
@test AugmentQTests.oneEquationNoLead()
@test AugmentQTests.reliablePaperExmpl()
@test AugmentQTests.athan()
@test AugmentQTests.habitmod()
end

# test_path = joinpath(dirname(@__FILE__), "defineReducedFormTestFuncs.jl")
include(joinpath(dirname(@__FILE__), "defineReducedFormTestFuncs.jl"))
@testset "test reducedForm" begin
@test ReducedFormTests.firmvalueFalse()
@test ReducedFormTests.firmvalue3Leads2LagsFalse()
@test ReducedFormTests.example7False()
@test ReducedFormTests.oneEquationNoLeadFalse()
@test ReducedFormTests.reliablePaperExmplFalse()
@test ReducedFormTests.athanFalse()
@test ReducedFormTests.habitmodFalse()

@test ReducedFormTests.firmvalueTrue()
@test ReducedFormTests.firmvalue3Leads2LagsTrue()
@test ReducedFormTests.example7True()
@test ReducedFormTests.oneEquationNoLeadTrue()
@test ReducedFormTests.reliablePaperExmplTrue()
@test ReducedFormTests.athanTrue()
@test ReducedFormTests.habitmodTrue()
end

# test_path = joinpath(dirname(@__FILE__), "defineAndersonMooreAlgTestFuncs.jl")
include(joinpath(dirname(@__FILE__), "defineAndersonMooreAlgTestFuncs.jl"))
@testset "test AndersonMooreAlg" begin

    @test AndersonMooreAlgTests.firmvalueFalse()
    @test AndersonMooreAlgTests.firmvalue3Leads2LagsFalse()
    @test AndersonMooreAlgTests.example7False()
    @test AndersonMooreAlgTests.oneEquationNoLeadFalse()
    @test AndersonMooreAlgTests.reliablePaperExmplFalse()
    @test AndersonMooreAlgTests.athanFalse()
    @test AndersonMooreAlgTests.habitmodFalse()

    @test AndersonMooreAlgTests.firmvalueTrue()
    @test AndersonMooreAlgTests.firmvalue3Leads2LagsTrue()
    @test AndersonMooreAlgTests.example7True()
    @test AndersonMooreAlgTests.oneEquationNoLeadTrue()
    @test AndersonMooreAlgTests.reliablePaperExmplTrue()
    @test AndersonMooreAlgTests.athanTrue()
    @test AndersonMooreAlgTests.habitmodTrue()

end

# test_path = joinpath(dirname(@__FILE__), "defineCcallTestFuncs.jl")
include(joinpath(dirname(@__FILE__), "defineCcallTestFuncs.jl"))
@testset "test ccall" begin

    @test CcallTests.firmvalueFalse()
    @test CcallTests.firmvalue3Leads2LagsFalse()
    @test CcallTests.example7False()
    @test CcallTests.oneEquationNoLeadFalse()
    @test CcallTests.reliablePaperExmplFalse()
    @test CcallTests.athanFalse()
    @test CcallTests.habitmodFalse()

    @test CcallTests.firmvalueTrue()
    @test CcallTests.firmvalue3Leads2LagsTrue()
    @test CcallTests.example7True()
    @test CcallTests.oneEquationNoLeadTrue()
    @test CcallTests.reliablePaperExmplTrue()
    @test CcallTests.athanTrue()
    @test CcallTests.habitmodTrue()

end

# test_path = joinpath(dirname(@__FILE__), "defineErrTestFuncs.jl")
include(joinpath(dirname(@__FILE__), "defineErrTestFuncs.jl"))
@testset "test err" begin
    @test ErrTests.noErrors()
    @test ErrTests.tooManyRoots()
    @test ErrTests.tooFewRoots()
    #@test ErrTests.tooManyExactShifts()
    #@test ErrTests.tooManyNumericShifts()
    @test ErrTests.spurious()
end

# test_path = joinpath(dirname(@__FILE__), "defineGensysTestFuncs.jl")
include(joinpath(dirname(@__FILE__), "defineGensysTestFuncs.jl"))
@testset "test gensysToAMA" begin
    @test GensysTests.example0()
    @test GensysTests.example1()
    @test GensysTests.example2()
    @test GensysTests.example3()
    @test GensysTests.example4()
    @test GensysTests.example5()
    #@test GensysTests.example6()
    @test GensysTests.example7()
end

end #outer


text:
using CSV
using DelimitedFiles
using DataFrames
using Dates
using Plots
using TimerOutputs
using BenchmarkTools
using Revise

include("branching_process.jl")
include("processing_sim.jl")

function simulate_single_state(
    file_date,
    state,
    nsims,
    run_simulation;
    truncation_days = 7,
    adjust_TP = false,
)
    """
    This runs the branching process for a single state.
    """

    # set seed for consistent plots (NOTE: this is not useful when multithreading
    # enabled as we use separate seeds but the simulation pool should handle that)
    rng = Random.seed!(2022)

    jurisdiction_assumptions = JurisdictionAssumptions()

    # get the latest onset date
    latest_start_date = Dates.Date(
        maximum(v for v in values(jurisdiction_assumptions.simulation_start_dates))
    )

    (dates, local_case_dict, import_case_dict) = read_in_cases(file_date, rng)
    last_date_in_data = dates[end]
    forecast_end_date = last_date_in_data + Dates.Day(35)

    # create vector for dates
    onset_dates = latest_start_date:Dates.Day(1):forecast_end_date

    if run_simulation
        forecast_start_date = Dates.Date(
            jurisdiction_assumptions.simulation_start_dates[state]
        )
        local_cases = local_case_dict[state]
        import_cases = import_case_dict[state]

        # named tuple for initial condi
        D0 = jurisdiction_assumptions.initial_conditions[state]
        N = jurisdiction_assumptions.pop_sizes[state]

        # get the observed cases
        local_cases = local_case_dict[state][dates .>= forecast_start_date]
        import_cases = import_case_dict[state]
        # cutoff the last bit of the local cases
        local_cases = local_cases[begin:end - truncation_days + 1]
        import_cases = import_cases[begin:end - truncation_days + 1]

        (D, U, Z) = simulate_branching_process(
            D0,
            N,
            nsims,
            local_cases,
            import_cases,
            forecast_start_date,
            file_date,
            state,
            adjust_TP = adjust_TP,
        )

        save_simulations(
            D,
            state,
            file_date,
            onset_dates,
            rng,
        )
    end

    return nothing

end

function simulate_all_states(
    file_date,
    states,
    nsims,
    run_simulation;
    truncation_days = 7,
    p_detect_omicron = 0.5,
    adjust_TP = false,
)
    """
    This runs the branching process for each of the states in states and then
    merges and saves the files. If running a single state, this merging of files is
    done in a post-processing step.
    """

    # set seed for consistent plots (NOTE: this is not useful when multithreading
    # enabled as we use separate seeds but the simulation pool should handle that)
    rng = Random.seed!(2022)

    jurisdiction_assumptions = JurisdictionAssumptions()

    # get the latest onset date
    latest_start_date = Dates.Date(
        maximum(v for v in values(jurisdiction_assumptions.simulation_start_dates))
    )

    (local_case_dict, import_case_dict) = read_in_cases(file_date, rng)
    dates = local_case_dict["date"]
    last_date_in_data = dates[end]
    forecast_end_date = last_date_in_data + Dates.Day(35)

    # create vector for dates
    onset_dates = latest_start_date:Dates.Day(1):forecast_end_date

    if run_simulation
        for state in states

            println("======================")
            println("Simulating ", state)
            println("======================")

            forecast_start_date = Dates.Date(
                jurisdiction_assumptions.simulation_start_dates[state]
            )
            local_cases = local_case_dict[state]
            import_cases = import_case_dict[state]

            # named tuple for initial condi
            D0 = jurisdiction_assumptions.initial_conditions[state]
            N = jurisdiction_assumptions.pop_sizes[state]

            # get the observed cases
            cases_pre_forecast = sum(local_case_dict[state][dates .< forecast_start_date])
            local_cases = local_case_dict[state][dates .>= forecast_start_date]
            # cutoff the last bit of the local cases
            import_cases = import_case_dict[state]
            local_cases = local_cases[begin:end - truncation_days + 1]
            import_cases = import_cases[begin:end - truncation_days + 1]

            (D, U, TP_local) = simulate_branching_process(
                D0,
                N,
                nsims[state],
                local_cases,
                import_cases,
                cases_pre_forecast,
                forecast_start_date,
                file_date,
                jurisdiction_assumptions.omicron_dominant_date,
                state,
                p_detect_omicron = p_detect_omicron,
                adjust_TP = adjust_TP,
                parallel = true
            )

            save_simulations(
                D,
                TP_local,
                state,
                file_date,
                onset_dates,
                rng,
            )

        end
    end

    # merge all the simulation and TP files for states into single CSV
    merge_simulation_files(file_date)
    merge_TP_files(file_date)

    return nothing

end


text:
"""
*Bayesian Neural Networks for Julia.*

# Introduction

Formulates and fits Bayesian Neural Network with Variational Inference.

# Examples

```julia
using Distributions, BayesNN, StatsFuns

# data size
N = 1000  # number of observations
D = 3     # number of covariates including intercept

# generate data
srand(1)                  # insure reproducibility
X = rand(Normal(), N, D)  # generate covariates
X[:, 1] = 1               # make first column intercept
b = rand(Normal(), D)     # coefficients
θ = logistic.(X * b)      # simulate dependent variables
y = rand.(Bernoulli.(θ))  # simulate dependent variables

net = Nnet(y, X)
```

"""
module BayesNN

using StatsFuns
using Distributions
using GradDescent
using ForwardDiff

export
    Nnet,
    fit,
    predict

include("helper.jl")
include("propagate.jl")
include("Nnet.jl")

end # module


text:
# Does not work
#=

using BinDeps

@BinDeps.setup

klb = library_dependency("klb", runtime=true, aliases=["libklb.so", "klb.dll"])

deps = [ klb ]

provides(Sources,
         [ URI("https://bitbucket.org/fernandoamat/keller-lab-block-filetype/get/1796d6334bd3.zip") ],
         klb
         )

@BinDeps.install begin

prefix = joinpath(BinDeps.depsdir(klb), "usr")
uprefix = replace(replace(prefix,"\\" => "/"), "C:/" => "/c/")
klbsrcdir   = joinpath(BinDeps.depsdir(klb), "src",    "klb")
klbbuilddir = joinpath(BinDeps.depsdir(klb), "builds", "klb")

println(prefix)
println(uprefix)
println(klbsrcdir)
println(klbbuilddir)
println(joinpath(prefix, "lib", "libklb.dll"))

end

=#


text: calc_H(reactor, R_0, B_0) = K_H(reactor) * B_0 ^ 1.19 * R_0 ^ 1.02


text:
# Indicates how coefficients should be discretised for the budgeted implementation.
abstract type ESCB2BudgetedDiscretisationScheme end

struct ESCB2BudgetedDiscretisationSchemeConstant <: ESCB2BudgetedDiscretisationScheme
  v::Float64
end

struct ESCB2BudgetedDiscretisationSchemeLambda <: ESCB2BudgetedDiscretisationScheme
  f::Function # Argument: the current round t.
end

function _discretisation_provably_converging(instance::CombinatorialInstance, round::Int)
  logm = max(1.0, log(instance.m))
  δ = instance.n_arms * (logm ^ 2) / round
  return δ / instance.m
end

ESCB2BudgetedDiscretisationSchemeConservative(instance::CombinatorialInstance, max_rounds::Int) =
  ESCB2BudgetedDiscretisationSchemeConstant(_discretisation_provably_converging(instance, max_rounds))
ESCB2BudgetedDiscretisationSchemeAdaptive(instance::CombinatorialInstance) =
  ESCB2BudgetedDiscretisationSchemeLambda(t -> _discretisation_provably_converging(instance, t))

_get_ξ(::ESCB2BudgetedDiscretisationScheme, ::Int) = error("Not implemented")
_get_ξ(s::ESCB2BudgetedDiscretisationSchemeConstant, ::Int) = s.v
_get_ξ(s::ESCB2BudgetedDiscretisationSchemeLambda, t::Int) = s.f(t)

# Parameters for the budgeted algorithm.
struct ESCB2Budgeted <: ESCB2OptimisationAlgorithm
  discretisation_scheme::ESCB2BudgetedDiscretisationScheme
  solve_all_budgets_at_once::Union{Bool, Nothing} # Some optimisation algorithms can solve the budgeted problems for all values of the budget at once more efficiently.
end

ESCB2Budgeted(i::CombinatorialInstance, s::Union{Bool, Nothing}=nothing) =
  ESCB2Budgeted(ESCB2BudgetedDiscretisationSchemeAdaptive(i), s)
ESCB2Budgeted(ε::Float64, s::Union{Bool, Nothing}=nothing) =
  ESCB2Budgeted(ESCB2BudgetedDiscretisationSchemeConstant(ε), s)

# Optimisation of ESCB2's objective function.
function optimise_linear_sqrtlinear(instance::CombinatorialInstance{T}, algo::ESCB2Budgeted,
                                    linear::Dict{T, Float64}, sqrtlinear::Dict{T, Float64},
                                    sqrtlinear_weight::Float64, bandit_round::Int;
                                    with_trace::Bool=false) where T
  # Transform the linear term in a budget constraint, the nonlinear term
  # becoming the objective function (in which case the concave function
  # can be dropped).
  # For this to work, the linear part of the objective function must
  # take only integer values, i.e. the linear coefficients are actually
  # only integers.
  ξ = _get_ξ(algo.discretisation_scheme, bandit_round)
  linear_discrete = Dict(k => round(Int, v / ξ, RoundUp) for (k, v) in linear)

  # Fill automatic values.
  if algo.solve_all_budgets_at_once === nothing
    algo = copy(algo) # Don't modify the user's object.
    algo.solve_all_budgets_at_once = supports_solve_all_budgeted_linear(instance)
  end

  # Start timing.
  t0 = time_ns()

  # Maximum value of the linear term?
  max_budget = instance.n_arms * maximum(values(linear_discrete))

  # Solve the family of problems with an increasing budget.
  best_solution, best_objective = _maximise_nonlinear_through_budgeted(instance, sqrtlinear, linear_discrete, max_budget,
                                                                       (sol, _) -> escb2_index(linear, sqrtlinear, sqrtlinear_weight, sol),
                                                                       algo.solve_all_budgets_at_once)
  t1 = time_ns()

  # Return the nice data structures.
  if with_trace
    run_details = ESCB2Details()
    run_details.n_iterations = 1
    run_details.best_objective = best_objective
    run_details.solver_time = (t1 - t0) / 1_000_000_000
    return best_solution, run_details
  else
    return best_solution
  end
end


text:
using Polynomials
export char_poly
"""
`char_poly(A::Matrix)` returns the characteristic polynomial
of the matrix `A`.
"""
function char_poly(A::Matrix{T}) where T
    r,c = size(A)
    @assert r==c "Matrix must be square"

    L = zeros(Poly{T},r,r)
    x = Poly{T}([0,1])
    for i=1:r
        L[i,i] = x
    end

    cofactor_det(L-A)
end


text:
"""
This file is used to run the
"""

include("simulate_states.jl")

# parameters to pass to the main function
const file_date = ARGS[1]
const nsims = parse(Int, ARGS[2])

# states to simulate
const states_to_run = [
    "NSW",
]

# run main
simulate_all_states(file_date, states_to_run, nsims)


text:
using ImagineWorker
using Test


text:
using Documenter

makedocs(
    sitename="PersonalPlaygound.jl",
    pages = [
        "index.md",
    ]
)

deploydocs(
    repo = "github.com/jjj999/PersonalPlayground.jl.git",
    devbranch = "develop",
    push_preview = true,
)


text:
using ConstrainedDynamics
using GPR

function max2mincoordinates(cstate::CState, mechanism::Mechanism)
    oldstates = getStates(mechanism)
    states = toStates(cstate)
    resetMechanism!(mechanism, states)
    cstate_min = Vector{Float64}()
    for eqc in mechanism.eqconstraints
        append!(cstate_min, ConstrainedDynamics.minimalCoordinates(mechanism, eqc))
        append!(cstate_min, ConstrainedDynamics.minimalVelocities(mechanism, eqc))
    end
    for (id, state) in enumerate(oldstates)
        mechanism.bodies[id].state = state  # Reset mechanism to default values
    end
    return cstate_min
end

"""
    Fourbar doesn't work well with the generic max2mincoordinates.
"""
function max2mincoordinates_fb(cstate::CState{T,4}) where T
    minstate = Vector{T}(undef, 4)
    q1 = UnitQuaternion(cstate[4], cstate[5:7])
    q2 = UnitQuaternion(cstate[30], cstate[31:33])  # angle 2 is body 3
    minstate[1] = Rotations.rotation_angle(q1)*sign(q1.x)*sign(q1.w)
    minstate[2] = cstate[11]
    minstate[3] = Rotations.rotation_angle(q2)*sign(q2.x)*sign(q2.w)
    minstate[4] = cstate[37]
    return minstate
end

"""
    Noise for v is calculated by doing a forward integration of ω in minimal coordinates, finding x2 and building a v that is consistent
    with the variational integration approach.
"""
function applynoise!(df, Σ, etype, Δtsim, varargs...)
    etype == "P1" && return _applynoise_p1!(df, Σ, Δtsim, varargs...)
    etype == "P2" && return _applynoise_p2!(df, Σ, Δtsim, varargs...)
    etype == "CP" && return _applynoise_cp!(df, Σ, Δtsim, varargs...)
    etype == "FB" && return _applynoise_fb!(df, Σ, Δtsim, varargs...)
    throw(ArgumentError("Experiment $etype is not supported!"))
end

function _applynoise_p1!(df, Σ, Δtsim, l)
    for col in eachcol(df)
        for t in 1:length(col)
            col[t][1].qc = UnitQuaternion(RotX(Σ["q"]*randn())) * col[t][1].qc  # Small error around θ
            col[t][1].ωc += Σ["ω"]*[randn(), 0, 0]  # Zero noise in fixed ωy, ωz
            θ = Rotations.rotation_angle(col[t][1].qc)*sign(col[t][1].qc.x)*sign(col[t][1].qc.w)  # Signum for axis direction
            ω = col[t][1].ωc[1]
            col[t][1].xc = [0, l/2*sin(θ), -l/2*cos(θ)]  # Noise is consequence of θ and ω
            θnext = ω*Δtsim + θ
            xnext = [0, l/2*sin(θnext), -l/2*cos(θnext)]
            v = (xnext - col[t][1].xc)/Δtsim
            col[t][1].vc = v
        end
    end
end

function _applynoise_p2!(df, Σ, Δtsim, l1, l2)
    for col in eachcol(df)
        for t in 1:length(col)
            col[t][1].qc = UnitQuaternion(RotX(Σ["q"]*randn())) * col[t][1].qc
            col[t][1].ωc += Σ["ω"]*[randn(), 0, 0]  # Zero noise in fixed ωy, ωz
            col[t][2].qc = UnitQuaternion(RotX(Σ["q"]*randn())) * col[t][2].qc
            col[t][2].ωc += Σ["ω"]*[randn(), 0, 0]  # Zero noise in fixed ωy, ωz
            θ1 = Rotations.rotation_angle(col[t][1].qc)*sign(col[t][1].qc.x)*sign(col[t][1].qc.w)  # Signum for axis direction
            θ2 = Rotations.rotation_angle(col[t][2].qc)*sign(col[t][2].qc.x)*sign(col[t][2].qc.w) - θ1
            ω1, ω2 = col[t][1].ωc[1], col[t][2].ωc[1] - col[t][1].ωc[1]
            col[t][1].xc = [0, l1/2*sin(θ1), -l1/2*cos(θ1)]  # Noise is consequence of θ and ω
            col[t][2].xc = [0, l1*sin(θ1) + l2/2*sin(θ1+θ2), -l1*cos(θ1) - l2/2*cos(θ1+θ2)]  # Noise is consequence of θ and ω
            θ1next = θ1 + ω1*Δtsim
            θ2next = θ2 + ω2*Δtsim
            x1next = [0, l1/2*sin(θ1next), -l1/2*cos(θ1next)]
            x2next = [0, l1*sin(θ1next) + l2/2*sin(θ1next+θ2next), -l1*cos(θ1next) - l2/2*cos(θ1next+θ2next)]
            v1 = (x1next - col[t][1].xc) / Δtsim
            v2 = (x2next - col[t][2].xc) / Δtsim
            col[t][1].vc = v1
            col[t][2].vc = v2
        end
    end
end

function _applynoise_cp!(df, Σ, Δtsim, l)
    for col in eachcol(df)
        for t in 1:length(col)
            col[t][1].xc += Σ["x"]*[0, randn(), 0]  # Cart pos noise only y, no orientation noise
            col[t][1].vc += Σ["v"]*[0, randn(), 0]  # Same for v
            col[t][2].qc = UnitQuaternion(RotX(Σ["q"]*randn())) * col[t][2].qc
            col[t][2].ωc += Σ["ω"]*[randn(), 0, 0]
            θ = Rotations.rotation_angle(col[t][2].qc)*sign(col[t][2].qc.x)*sign(col[t][2].qc.w)  # Signum for axis direction
            ω = col[t][2].ωc[1]
            col[t][2].xc = col[t][1].xc + [0, l/2*sin(θ), -l/2*cos(θ)]
            θnext = ω*Δtsim + θ
            xnext = col[t][1].xc + col[t][1].vc*Δtsim + [0, l/2*sin(θnext), -l/2*cos(θnext)]
            v = (xnext - col[t][2].xc)/Δtsim
            col[t][2].vc = v
        end
    end
end

function _applynoise_fb!(df, Σ, Δtsim, l)
    for col in eachcol(df)
        for t in 1:length(col)
            col[t][1].qc = UnitQuaternion(RotX(Σ["q"]*randn())) * col[t][1].qc
            col[t][1].ωc += Σ["ω"]*[randn(), 0, 0]  # Zero noise in fixed ωy, ωz
            col[t][3].qc = UnitQuaternion(RotX(Σ["q"]*randn())) * col[t][3].qc
            col[t][3].ωc += Σ["ω"]*[randn(), 0, 0]  # Zero noise in fixed ωy, ωz
            θ1 = Rotations.rotation_angle(col[t][1].qc)*sign(col[t][1].qc.x)*sign(col[t][1].qc.w)  # Signum for axis direction
            θ2 = Rotations.rotation_angle(col[t][3].qc)*sign(col[t][3].qc.x)*sign(col[t][3].qc.w)
            ω1, ω2 = col[t][1].ωc[1], col[t][3].ωc[1]
            col[t][2].qc = UnitQuaternion(RotX(θ2))
            col[t][4].qc = UnitQuaternion(RotX(θ1))
            col[t][1].xc = [0, 0.5sin(θ1)l, -0.5cos(θ1)l]
            col[t][2].xc = [0, sin(θ1)l + 0.5sin(θ2)l, -cos(θ1)l - 0.5cos(θ2)l]
            col[t][3].xc = [0, 0.5sin(θ2)l, -0.5cos(θ2)l]
            col[t][4].xc = [0, sin(θ2)l + 0.5sin(θ1)l, -cos(θ2)l - 0.5cos(θ1)l]
            θ1next = ω1*Δtsim + θ1
            θ2next = ω2*Δtsim + θ2
            x1next = [0, 0.5sin(θ1next)l, -0.5cos(θ1next)l]
            x2next = [0, sin(θ1next)l + 0.5sin(θ2next)l, -cos(θ1next)l - 0.5cos(θ2next)l]
            x3next = [0, 0.5sin(θ2next)l, -0.5cos(θ2next)l]
            x4next = [0, sin(θ2next)l + 0.5sin(θ1next)l, -cos(θ2next)l - 0.5cos(θ1next)l]
            v1 = (x1next - col[t][1].xc)/Δtsim
            v2 = (x2next - col[t][2].xc)/Δtsim
            v3 = (x3next - col[t][3].xc)/Δtsim
            v4 = (x4next - col[t][4].xc)/Δtsim
            col[t][1].vc = v1
            col[t][2].vc = v2
            col[t][3].vc = v3
            col[t][4].vc = v4
            col[t][2].ωc = col[t][3].ωc
            col[t][4].ωc = col[t][1].ωc
        end
    end
end


text:
using Devectorize
using Gadfly

"""
Linear Regression Cost function with Regularisation

Parameters:
	params      = vector containing rolled up X and Θ
		X       = array of features for each movie
		Θ       = array of weights for each user
	Y           = array of ratings for each movie
	R           = array of the existance of ratings for each movie
	numUsers    = number of users
	numMovies   = number of movies
	numFeatures = number of features
	λ           = regularisation constant

Returns:
	J = linear regression cost of the given weights with the given regularisation constant
"""
function cost(params, Y::Array{Int}, numUsers, numMovies, numFeatures, λ)
	# Unfold params -> X, Θ
	X, Θ = unroll(params, numUsers, numMovies, numFeatures)

	# Calculate cost
	J = 0
 	@simd for i = 1:size(Y)[2]
		@inbounds J += (dot(X[Y[2, i], :][:], Θ[Y[1, i], :][:]) - Y[3, i])^2
	end
	J = 1/2 * J
	# add regularisation
	J += λ/2 * (sum(Θ.^2) + sum(X.^2))

	return J
end

"""
Unroll parameter vector into X and Θ

Parameters:
	params = vector containing rolled up X and Θ
	numUsers    = number of users
	numMovies   = number of movies
	numFeatures = number of features

Returns:
	X = array of features for each movie
	Θ = array of weights for each user
"""
function unroll(params, numUsers, numMovies, numFeatures)
	X = reshape(params[1 : numMovies * numFeatures], numMovies, numFeatures)
	Θ = reshape(params[numMovies * numFeatures + 1 : end], numUsers, numFeatures)
	return X, Θ
end

"""
Roll up X and Θ into single vector

Parameters:
	X = array of features for each movie
	Θ = array of weights for each user

Returns:
	params = vector containing rolled up X and Θ
"""
function rollup(X, Θ)
	return params = vcat(X..., Θ...)
end

function predict(movie, me, others, Y, users)
	meIndices = find(Y[1, :] .== me)
	meMovies = Y[2, meIndices]
	weightSum = 0.0
	summation = 0.0
	otherIndices = Int[]
	tempMeIndices = Int[]
	for other in others
		otherIndices = users[other]
		otherMovies = Y[2, otherIndices]
		jointMovies = intersect(meMovies, otherMovies)
		if length(jointMovies) == 0
			continue
		end
		tempMeIndices = filter(x -> in(Y[2, x], jointMovies), meIndices)
		tempOtherIndices = filter(x -> in(Y[2, x], jointMovies), otherIndices)

		meRatings = Y[3, tempMeIndices]
		meAverage = mean(meRatings)
		otherRating = Y[3, find(Y[2, otherIndices] .== movie)][1]
		otherRatings = Y[3, tempOtherIndices]
		otherAverage = mean(otherRatings)

		temp = sqrt(sum((meRatings - meAverage).^2)) * sqrt(sum((otherRatings - otherAverage).^2))
		if temp == 0
			continue
		end

		@devec weight = sum((meRatings - meAverage).*(otherRatings - otherAverage)) ./ temp

		summation += (otherRating - otherAverage) * weight
		weightSum += weight
	end

	meRatings = Y[3, meIndices]
	meAverage = mean(meRatings)

	return meAverage + summation / weightSum
end

"""
Entry point

Params:
	dataPath = path to() CSVs
"""
function main()
	dataPath = ""
	if ARGS[1] == "small"
		dataPath = "../ml-latest-small/fixed/"
	elseif ARGS[1] == "full"
		dataPath = "../ml-latest/fixed/"
	end

	# Load data from CSV - using wc to speed up loading by preallocating
	# N.B. requires cat and wc to be installed - a shell should not be nessecary though
	#lines = parse(readall(pipeline(`cat $[dataPath]ratings.csv`, `wc -l`)))
	#@time ratings = readdlm(dataPath * "ratings.csv", ',', header=true, dims=(lines, 4))[1]

	#lines = parse(readall(pipeline(`cat $[dataPath]train.csv`, `wc -l`)))
	#@time train = readdlm(dataPath * "train.csv", ',', header=true, dims=(lines, 4))[1]

	lines = parse(readall(pipeline(`cat $[dataPath]cross.csv`, `wc -l`)))
	@time cross = readdlm(dataPath * "cross.csv", ',', header=true, dims=(lines, 4))[1]

	lines = parse(readall(pipeline(`cat $[dataPath]test.csv`, `wc -l`)))
	@time test = readdlm(dataPath * "test.csv", ',', header=true, dims=(lines, 4))[1]

	lines = parse(readall(pipeline(`cat $[dataPath]movies.csv`, `wc -l`)))
	@time movies = readdlm(dataPath * "movies.csv", ',', header=true, dims=(lines, 3))
	moviesHeader = movies[2]
	movies = movies[1]

	lines = parse(readall(pipeline(`cat X-results.csv`, `wc -l`)))
	@time X = readdlm("X-results.csv", ',', header=false, dims=(lines, 18))

	lines = parse(readall(pipeline(`cat Theta-results.csv`, `wc -l`)))
	@time Θ = readdlm("Theta-results.csv", ',', header=false, dims=(lines, 18))

	testUsersData = readcsv("test-users-data.csv")
	users = testUsersData[:, 2]
	usersNumMovies = testUsersData[:, 3]
	usersAvgRatings = testUsersData[:, 4]
	usersRatingsDifference = testUsersData[:, 5]

	#=@printf "users\n"
	usersIndices = Vector{Vector{Int}}()
	f = open("usersIndices.csv")
	i = 0
	@time while (l = readline(f)) != nothing && strip(l) != ""
		i += 1
		if i % 10_000 == 0
			@printf "%i/%i\t%s\n" i 240_000 now()
		end
		l = replace(l, "[", "")
		l = replace(l, "]", "")
		temp = map(parse, split(l, ","))
		if typeof(temp) != Array{Int, 1}
			temp = []
		end
		push!(usersIndices, temp)
	end
	close(f)=#

	#data = [train; cross]
	#data[:, 3] *= 2
	#data = Array{Int}(data[:, 1:3]')
	test[:, 3] *= 2
	test = Array{Int}(test[:, 1:3]')

	#nonEmptyUserIndices = filter(x -> length(usersIndices[x]) > 0, collect(1:size(usersIndices)[1]))
	#realUsers = map(x -> data[1, nonEmptyUserIndices[x][1]], collect(1:size(nonEmptyUserIndices)[1]))
	#reallyNonEmptyUserIndices = usersIndices[nonEmptyUserIndices]

	function userCostModel(user, indices)
		cost = 0
		for i in indices
			cost += (dot(X[test[2, i], :][:], Θ[test[1, i], :][:]) - test[3, i])^2
		end
		return cost / length(indices)
	end

	#=function userCostMemory(user, indices)
		cost = 0
		for i in indices
			movie = test[2, i]
			rating = test[3, i]

			others = find(data[2, :] .== movie)
			others = data[1, others]
			others = intersect(others, realUsers)
			cost += (predict(movie, user, others, data, reallyNonEmptyUserIndices) - rating)^2
		end
		return cost / length(indices)
	end=#

	#memoryCost = zeros(length(realUsers))
	modelCost = zeros(length(users))
	idx = 0
	for user in users
		if idx % 1000 == 0
			@show idx
		end
		idx += 1
		indices = find(test[1, :] .== user)
		try
			modelCost[idx] = userCostModel(user, indices)
			#memoryCost[idx] = userCostMemory(user, indices)
		catch e
			@show e
			continue
		end
	end

	#writecsv("memoryCost.csv", memoryCost)
	writecsv("modelCost.csv", modelCost)
	#writecsv("final-out.csv", hcat(users, usersNumMovies, usersAvgRatings, usersRatingsDifference, memoryCost, modelCost))

	#=@printf "movies\n"
	moviesIndices = Vector{Vector{Int}}()
	f = open("moviesIndices.csv")
	i = 0
	@time while (l = readline(f)) != nothing && strip(l) != ""
		i += 1
		if i % 1_000 == 0
			@printf "%i/%i\t%s\n" i 33_000 now()
		end
		l = replace(l, "[", "")
		l = replace(l, "]", "")
		temp = map(parse, split(l, ","))
		if typeof(temp) != Array{Int, 1}
			temp = []
		end
		push!(moviesIndices, temp)
	end
	close(f)

	@printf "Finished loading, begin processing\n"

	data = [train; cross]
	data[:, 3] *= 2
	data = Array{Int}(data[:, 1:3]')
	test[:, 3] *= 2
	test = Array{Int}(test[:, 1:3]')

	users = union(test[1, :])
	@time usersAvgRatings = pmap(x -> mean(data[2, usersIndices[x]]), users)
	@show size(usersAvgRatings)
	@show size(usersIndices)
	@show size(usersIndices[1])
	@show usersIndices[1]
	@show data[2, usersIndices[1]]
	@show mean(data[2, usersIndices[1]])
	@show map(x -> mean(data[2, usersIndices[x]]), collect(1:5))
	@show size(data[2, usersIndices[1]])
	# Count users, movies and features
	numUsers = length(union(ratings[:,1]))
	numMovies = size(movies)[1]

	numFeatures = 18
	λ = 0.05

	function costTemp(x)
		cost(x, train, numUsers, numMovies, numFeatures, λ)
	end

	#@time moviesIndices = pmap(x -> find(data[2, :] .== x), movies[:, 1])
	@time moviesRatingsCount = pmap(x -> length(moviesIndices[x]), movies[:, 1])
	@time moviesAvgRating = pmap(x -> mean(data[3, moviesIndices[x]]), movies[:, 1])

	users = union(test[1, :])
	@time usersNumMovies = pmap(x -> length(usersIndices[x]), users)
	@time usersAvgRatings = pmap(x -> mean(moviesRatingsCount[data[2, usersIndices[x]]]), users)
	@time usersRatingsDifference = pmap(x -> mean(map(y -> moviesAvgRating[data[2, y]] - data[3, y], usersIndices[x])), users)

	@printf "Done calculating, saving\n"

	writecsv("test-users-data.csv", hcat(users, usersNumMovies, usersAvgRatings, usersRatingsDifference))

	@printf "Done :D\n"=#
end

main()


text:
function solve!(pn::ProjectedNewtonSolver)
    copyto!(pn.Z̄data, pn.Zdata)
    projection_solve!(pn)
end

function projection_solve!(pn::ProjectedNewtonSolver)
    ϵ_feas = pn.opts.constraint_tolerance
    evaluate_constraints!(pn)
    viol = max_violation(pn, nothing)
    max_projection_iters = pn.opts.n_steps
    count = 0
    J_prev = cost(pn, pn.Z̄)
    while count <= max_projection_iters && viol > ϵ_feas
        _qdldl_solve!(pn)
        J = cost(pn, pn.Z̄)
        viol = max_violation(pn)  # calculate again, updating the active set
        if (pn.opts.multiplier_projection)
            # TODO: (#35) Re-implement this
            # res = multiplier_projection!(pn)
            res = Inf
        else
            res = Inf
        end
        count += 1
        dJ = J_prev - J
        record_iteration!(pn, viol, res, J, dJ)
        J_prev = J
    end
    return viol
end

function record_iteration!(pn::ProjectedNewtonSolver, viol, res, J, dJ)
    # TODO: (#35) compute the gradient
    record_iteration!(pn.stats, cost=J, c_max=viol, is_pn=true, dJ=dJ, # gradient=res,
        penalty_max=NaN)
end

function update_b!(pn::ProjectedNewtonSolver)
    Np = num_primals(pn)
    Nd = num_duals(pn)
    Na = sum(pn.active) - Np
    b = pn.b
    b .= 0
    resize!(pn.b, Na + Np)
    nzcount = 1
    for i = 1:Nd
        if pn.active[Np + i]
            b[Np + nzcount] = -pn.d[i]
            nzcount += 1
        end
    end
    return b
end

function getKKTMatrix(pn::ProjectedNewtonSolver)
    Np = num_primals(pn)
    Nd = num_duals(pn)
    Na = sum(pn.active) - Np

    nnz0 = nnz(pn.Atop) + Nd
    resize!(pn.colptr, Np + Nd + 1)
    resize!(pn.rowval, nnz0)
    resize!(pn.nzval, nnz0)
    nnz_new = triukkt!(pn.Atop, pn.active, pn.colptr, pn.rowval, pn.nzval, reg=1e-8)
    colptr = resize!(pn.colptr, Np + Na + 1)
    rowval = resize!(pn.rowval, nnz_new)
    nzval = resize!(pn.nzval, nnz_new)
    A = SparseMatrixCSC(Np + Na, Np + Na, colptr, rowval, nzval)
end

function _qdldl_solve!(pn::ProjectedNewtonSolver)
    Np = num_primals(pn)
    Nd = num_duals(pn)

    max_refinements = 10
    convergence_rate_threshold = pn.opts.r_threshold

    # Update everything
    evaluate_constraints!(pn)
    constraint_jacobians!(pn)
    cost_hessian!(pn)
    update_active_set!(pn)
    primalregularization!(pn)
    Na = sum(pn.active) - Np

    # Get A,b
    A = getKKTMatrix(pn)

    # Copy the active constraints to the b vector and resize
    b = update_b!(pn)
    resize!(pn.dY, Np + Na)

    # Factorize the matrix
    resize!(pn.qdldl, Np + Na)
    Cqdldl.eliminationtree!(pn.qdldl, A)
    Cqdldl.factor!(pn.qdldl)
    F = Cqdldl.QDLDLFactorization(pn.qdldl)

    # Line search
    viol_prev = max_violation(pn, nothing)
    count = 1
    while count < max_refinements
        viol = _qdldl_linesearch(pn, F, b)
        convergence_rate = log10(viol) / log10(viol_prev)
        # println("  r iter = $(count + 1), v = $viol, r = $convergence_rate")
        viol_prev = viol
        count += 1
        if viol < pn.opts.constraint_tolerance
            break
        elseif convergence_rate < convergence_rate_threshold
            break
        end
    end
    return
end

function _qdldl_linesearch(pn, F, b)
    Np = num_primals(pn)
    Na = length(b) - Np
    ia = Np .+ (1:Na)
    viol0 = max_violation(pn, nothing)
    dY = pn.dY
    # update_b!(pn)
    p = view(dY, 1:Np)
    α = 1.0
    for i = 1:10
        dY .= b
        ldiv!(F, dY)
        # QDLDL.solve!(F, dY)
        pn.Z̄data .= pn.Zdata .+ α .* p

        evaluate_constraints!(pn, pn.Z̄)
        update_b!(pn)
        v = max_violation(pn, nothing)   # don't update active set
        if v < viol0
            copyto!(pn.Zdata, pn.Z̄data)
            return v
        else
            # b[ia] .= pn.d[pn.active]
            α /= 2
        end
    end
    return NaN
end

# function multiplier_projection!(pn::ProjectedNewtonSolver)
#     λ = pn.Ydata[pn.active]
#     D,d = active_constraints(pn)
#     g = pn.g
#     res0 = g + D'λ
#     A = D*D'
#     Areg = A + I*pn.opts.ρ_primal
#     b = D*res0
#     δλ = -reg_solve(A, b, Areg)
#     λ += δλ
#     res = g + D'λ  # dual feasibility
#     pn.Ydata[pn.active] .= λ
#     return norm(res)
# end


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

using Markdown
using InteractiveUtils

# ╔═╡ 8595f030-5e9b-11eb-3bd7-ef8250cd1fab
begin
	# import the entire package
	using Optim
	using Distributions
	using DataFrames

	# import just a few functions
	using Random: seed!
	using StatsFuns: logsumexp, softmax
	using LinearAlgebra: diag
	using StatsBase: countmap

	# autodiff instead of finite diff?
	using FiniteDiff: finite_difference_gradient
end

# ╔═╡ 0ece0030-5e9d-11eb-0092-254fc6e72e83
md"
# Logit

Suppose that agents can choose an action $y \in \{0,1\}$. The payoffs to each are

```math
\begin{align*}
u_0(X,\epsilon) &= \epsilon_{i0} \\
u_1(X,\epsilon) &= x_i^\top \beta + \epsilon_{i1}
\end{align*}
```

The shocks $\epsilon_{i0},\epsilon_{i1}$ are distributed as iid Type-I extreme value with mean 0 and scale parameter 1. Agents choose $y_i=1$ if $u_1 \geq u_0$ and $y_0$ if $u_1 < u_0$
"

# ╔═╡ b7910ac0-5ffe-11eb-3925-13fbd88d3f67
md"
Denote $cdf(z) = F(z)$. Let choice $y_i \in \{0,1\}$. For a symmetric distribution and a binary discrete choice model, we can use this shortcut (trick is in Greene's Econometrics tome, Greene 6th ed p. 779):
```math
\log L(y|X) = \sum_i \log F\left( 2(y_i-1) x_i^\top \beta \right)
```
Score is a vector
```math
\nabla_\beta \log L(y_i|x_i) = \left[y_i - F\left(x_i^\top\beta\right)\right]x_i
```
Information matrix is
```math
\left[\sum_i \nabla \log L_i \nabla \log L_i^\top\right] \to Var(\beta)
```
"

# ╔═╡ e2248c2e-5ffe-11eb-189c-85dbf7703168
md"
We can vectorize stuff to make it simpler. The $.$ means element-by-element operations a la MATLAB/Julia. Define
```math
\boldsymbol q \equiv 2.\boldsymbol y .- 1
```
Then
```math
\log L(y|X) = \sum_i \log. F.\left(\boldsymbol q .* X\beta \right)
```
Matrix of scores
```math
\frac{\partial \log L_i}{\partial \beta} = \left(y .- F.(X\beta)\right) .* X
```
Information matrix
```math
\left(\frac{\partial \log L_i}{\partial \beta}\right)^\top \frac{\partial \log L_i}{\partial \beta} \to Var(\beta)
```

"

# ╔═╡ 1508d920-6000-11eb-34bf-81e7d3ceb9bc
md"
In the version below, we use the `Distributions.jl` package, which means we could actually change to a binary probit just by swapping out the distribution from `Logistic` to `Normal`.

Alternatively, for lower-level control, we can use the following functions from [`StatsFuns.jl`](https://github.com/JuliaStats/StatsFuns.jl). This is useful for more computationally intensive work with multinomial discrete choice.

```julia
logsumexp      # log(exp(x) + exp(y)) or log(sum(exp(x)))
softmax        # exp(x_i) / sum(exp(x)), for i
```
"

# ╔═╡ b9c02010-5e9b-11eb-1e5d-9358c778cbba
# set seed for random # generator
seed!(1234)

# ╔═╡ b9a85250-5e9b-11eb-1787-6b671b235984
begin
	nobs = 1_000
	β = [1.0, -2.0, 1.0, 0.5]
	k = length(β)
	X = randn(nobs, k);

	# choice utilities
	u0 = zeros(nobs)
	u1 = X*β
	u = hcat(u0, u1)
end

# ╔═╡ b976e210-5e9b-11eb-0b18-e95537e1d64f
# multinomial logit probabilities
prob_actions = mapslices(softmax, u; dims=2)

# ╔═╡ dc8270d0-5e9b-11eb-02cb-1faa58ce2656
cum_prob = cumsum(prob_actions; dims=2)

# ╔═╡ dc82bef0-5e9b-11eb-20f5-0d1ed55e9bb5
# will throw an error if we goof
@assert all(cum_prob[:,2] .≈ 1)

# ╔═╡ dc87c800-5e9b-11eb-194c-4f7f0e83a5c5
# instead of simulating random type-1 extreme values, we just
# use a uniform variable and the CDF
y = [searchsortedfirst(row, rand()) for row in eachrow(cum_prob) ] .-1

# ╔═╡ 694cc4c0-5e9c-11eb-3d00-458ca8c8fa28
countmap(y)

# ╔═╡ 90b963b0-5e9c-11eb-3e39-bdf160b8655e
function loglik(y, X, theta)
    n,k = size(X)
    ff(z) = logcdf(Logistic(), z)

    # see footnote 6 on p. 778 in Greene 6th ed for this shortcut
    q = 2 .* y .- 1
    u1 = X*theta
    LL = sum(ff.(q.*u1))

    return -LL  # I *think* you'll need to flip sign to maximize
end

# ╔═╡ 90b963b0-5e9c-11eb-1d75-25729dfb57ee
# note that the `!` means we're updating the first argument(s)
function dloglik!(grad, y, X, theta)

	n,k = size(X)
    u1 = X*theta

	# create function
	ff(z) = cdf(Logistic(), z)

	# all the broadcasting fuses operations into a single
	# loop instead of allocating temp vectors
	# this can help w/ speed + memory
    grad .= -vec(sum( (y .- ff.(u1)) .* X; dims=1))
    return grad
end

# ╔═╡ 90b963b0-5e9c-11eb-0dea-81d683f42601
# wrapper to allocate gradient vector
dloglik(y, X, theta) = dloglik!(similar(theta), y, X, theta)

# ╔═╡ 90c65c02-5e9c-11eb-032a-7d600b60b9c4
function informationmatrix(y, X, theta)

	n,k = size(X)
    infomatrix = zeros(k,k)

    u1 = X*theta
    ff(z) = cdf(Logistic(), z)
    g = y .- ff.(u1)   # as per Greene 6th ed p. 779

    infomatrix = (g .* X)' * (g .* X)

    return infomatrix # maybe flip signs?
end

# ╔═╡ 640b3fa2-5e9c-11eb-19d3-97ebb410851a
# closures wrap likelihood & gradient
begin
	f(thet) = loglik(y,X,thet)
	g!(grad,thet) = dloglik!(grad,y,X,thet)
end

# ╔═╡ ba5f02b0-5e9c-11eb-30cc-fbb3e2d4b6e2
# initial guess
theta0 = zeros(k)

# ╔═╡ ba5f02b0-5e9c-11eb-2221-4b0d5e525673
# Check gradient against finite difference
begin
	fdgrad = finite_difference_gradient(f, theta0, Val{:central})
	@assert  fdgrad ≈ dloglik(y,X,theta0)
	fdgrad .- dloglik(y,X,theta0)
end

# ╔═╡ ba614ca0-5e9c-11eb-3619-0507168a18d6
res = optimize(f, g!, theta0, BFGS(), Optim.Options(;show_trace=true))

# ╔═╡ 61327d20-5e9c-11eb-3704-bf5055d67bb0
begin
	theta1 = res.minimizer  # should be about β
	vcov = informationmatrix(y, X, theta1)
	vcovinv = inv(vcov)
	stderr = sqrt.(diag(vcovinv))
	tstats = theta1 ./ stderr
	pvals = map(z -> 2 .* cdf(Normal(), -abs(z)), tstats)

	DataFrame(beta = β, betahat = theta1, tstat=tstats, se = stderr, pval=pvals)
end

# ╔═╡ Cell order:
# ╠═0ece0030-5e9d-11eb-0092-254fc6e72e83
# ╠═b7910ac0-5ffe-11eb-3925-13fbd88d3f67
# ╠═e2248c2e-5ffe-11eb-189c-85dbf7703168
# ╠═1508d920-6000-11eb-34bf-81e7d3ceb9bc
# ╠═8595f030-5e9b-11eb-3bd7-ef8250cd1fab
# ╠═b9c02010-5e9b-11eb-1e5d-9358c778cbba
# ╠═b9a85250-5e9b-11eb-1787-6b671b235984
# ╠═b976e210-5e9b-11eb-0b18-e95537e1d64f
# ╠═dc8270d0-5e9b-11eb-02cb-1faa58ce2656
# ╠═dc82bef0-5e9b-11eb-20f5-0d1ed55e9bb5
# ╠═dc87c800-5e9b-11eb-194c-4f7f0e83a5c5
# ╠═694cc4c0-5e9c-11eb-3d00-458ca8c8fa28
# ╠═90b963b0-5e9c-11eb-3e39-bdf160b8655e
# ╠═90b963b0-5e9c-11eb-1d75-25729dfb57ee
# ╠═90b963b0-5e9c-11eb-0dea-81d683f42601
# ╠═90c65c02-5e9c-11eb-032a-7d600b60b9c4
# ╠═640b3fa2-5e9c-11eb-19d3-97ebb410851a
# ╠═ba5f02b0-5e9c-11eb-30cc-fbb3e2d4b6e2
# ╠═ba5f02b0-5e9c-11eb-2221-4b0d5e525673
# ╠═ba614ca0-5e9c-11eb-3619-0507168a18d6
# ╠═61327d20-5e9c-11eb-3704-bf5055d67bb0


text:
# Import modules
using Distributions  # probability distributions and associated functions
using LinearAlgebra
using StatsPlots

using Random
using JLD2

using Plots
# CES
using EnsembleKalmanProcesses.EnsembleKalmanProcessModule
using EnsembleKalmanProcesses.Observations
using EnsembleKalmanProcesses.ParameterDistributionStorage

rng_seed = 4137
Random.seed!(rng_seed)

# Output figure save directory
homedir = pwd()
println(homedir)
figure_save_directory = homedir * "/output/"
data_save_directory = homedir * "/output/"
if ~isdir(figure_save_directory)
    mkdir(figure_save_directory)
end
if ~isdir(data_save_directory)
    mkdir(data_save_directory)
end

###
###  model settings
###

###  Generate (artificial) truth samples
###  Note: The observables y are related to the parameters θ by:
###        y = G(θ) + η
###
# G(u) = [u[1] ; u[2]]
G(u) = [sqrt(u[1] - u[2]^2) ; u[2]]
θ1_true, θ2_true = 2.0, 1.0
prior_mean = [θ1_true + 1.0; θ2_true + 0.5]
prior_cov = Array(Diagonal(fill(0.1^2,  length(prior_mean))))
# prior_cov = Array(Diagonal(fill(1.0^2,  length(prior_mean))))

params_true = [θ1_true, θ2_true]
param_names = ["θ1", "θ2"]
n_param = length(param_names)
params_true = reshape(params_true, (n_param, 1))

println(n_param)
println(params_true)

# prior_stds = [2.0, 0.5*A_true]

prior_distributions = [Parameterized(Normal(prior_mean[1], sqrt(prior_cov[1,1]))), Parameterized(Normal(prior_mean[2], sqrt(prior_cov[2,2])))]
constraints = [[no_constraint()], [no_constraint()]]

priors = ParameterDistribution(prior_distributions, constraints, param_names)

###
###  Define the data from which we want to learn the parameters
###
data_names = ["y0", "y1"]

function run_G_ensemble(params)
    N_ensemble = size(params, 2)
    return hcat([G(params[:, i]) for i in 1:N_ensemble]...)
end
# Lorenz forward
# Input: params: [N_params, N_ens]
# Output: gt: [N_data, N_ens]
# Dropdims of the output since the forward model is only being run with N_ens=1
# corresponding to the truth construction
n_samples = 100
G_t = G(params_true)
y_t = zeros(length(G_t), n_samples)
Γy = Array(Diagonal(fill(0.1^2,  length(G_t))))

# Add noise
for i in 1:n_samples
    y_t[:, i] = G_t .+ rand(MvNormal(zeros(size(G_t)), Γy))
end

# Construct observation object
truth = Observations.Obs(y_t, Γy, data_names)
truth_sample = truth.mean


text:
# -*- coding: utf-8 -*-
# # Multivariate Nonlinear Regression
#
# This example demonstrates how to use a product group to manipulate a multivariate regression model in an algebraically precise way. It builds on the example `monomial_regression.jl`.
#

# Let `m` be a model like
# ```julia
# f(x,y,z) = a*x + b*y
# x,y = rand(Normal(0,1), (n,2))
# target[i] = f(x[i],y[i]) + rand(Normal(0,1), 1)
# min_{a,b} sum_i (f(x[i],y[i],z[i]) - target[i])^2
# ```
#
# So it is a bivariate least squares regression model.
# Define the set of transformations $T$ as the free monoid generated by
#
# ```julia
# x^i->x^i+1
# x^i->x^i-1
# y^i->y^i+1
# y^i->y^i-1
# ```
#
# The free monoid over this set of transformations is actually the product of cyclic groups $Z^2 = Z x Z$ with the
# group operation of $+:Z^2->Z^2$. so you get the set of models:
#
# ```julia
# f(x,y) = a*x^i + b*y^j
# x,y = rand(Normal(0,1), (n,2))
# target[i] = f(x[i],y[i]) + rand(Normal(0,1), 1)
# min_{a,b} sum_i (f(x[i],y[i]) - target[i])^2
# ```
#
# for all combinations of integers $i,j$ which is a pretty cool structure. If you do the transformations with `i mod N` then you get $Z_N^2$ which is a finite abelian group. It doesn't get more computationally tractable than component-wise arithmetic mod N.

# We could then talk about the orbits of a given model under this group. For this group, there is only one orbit, but if you take subgroups like the group generated by $(1,1)$ aka $\langle 1,1\rangle$ you get multiple orbits. For example if you start with
#
# ```julia
# f(x,y) = a*x^1 + b*y^2
# x,y = rand(Normal(0,1), (n,2))
# target[i] = f(x[i],y[i]) + rand(Normal(0,1), 1)
# min_{a,b} sum_i (f(x[i],y[i]) - target[i])^2
# ```
# the orbit is all of the models of the form:
# ```julia
# f(x,y) = a*x^1+i + b*y^2+i
# x,y = rand(Normal(0,1), (n,2))
# target[i] = f(x[i],y[i]) + rand(Normal(0,1), 1)
# min_{a,b} sum_i (f(x[i],y[i]) - target[i])^2
# ```
# where $i$ is in $Z_N$. You could reason about finding the best fitting model from a given orbit. and then comparing
# the orbits to see which one was contained the best fitting model.
#
# Here is an example implementation of this theory

# ## Example Model

#Our working example of a multivariate nonlinear regression model
expr = quote
    module Regression
    using Random
    using LsqFit
    using LinearAlgebra

    function f(X, β)
        a = β[1]
        b = β[2]
        x = X[:, 1]
        y = X[:, 2]
        return a.*x.^0 .+ b.*y.^0
    end

    function sample(g::Function, n)
        x = randn(Float64, n)
        target = g(x) .+ randn(Float64, n[1])./1600
        return x, target
    end

    function describe(fit)
        if !fit.converged
            error("Did not converge")
        end
        return (β = fit.param, r=norm(fit.resid,2))
    end
    #setup

    # Random.seed!(42)
    β = (1/2, 1/2)
    n = (1000,2)
    g(x) = β[1].*x[:,1].^2 + β[2].*x[:,2].^3
    X, target = sample(g, n)
    # x, y = X[:,1], X[:,2]
    # @show size(X), size(target)
    # loss(a) = sum((f.(a, x, y) .- target).^2)
    # @show loss.([-1,-1/2,-1/4, 0, 1/4,1/3,1/2,2/3, 1])

    a₀ = [1.5, 0.1]
    try
        ŷ₀ = f(X, a₀)
    catch
        error("Could not execute f on the initial data")
    end

    #solving
    fit = curve_fit(f, X, target, a₀; autodiff=:forwarddiff)
    result = describe(fit)
    end
end;

# ## Implementation Details
# The following code is the implementation details for representing the models as an `AbstractProblem` and representing the transformations as `Product{Tuple{Pow{Int}, Pow{Int}}}` and applying the transformations onto the models.

using SemanticModels
using SemanticModels.ModelTools
using SemanticModels.ModelTools.Transformations

import SemanticModels.ModelTools: model, AbstractModel
import SemanticModels.Parsers: findfunc, findassign
import Base: show

# +
""""    MultivariateLsq

A program that solves min_β || f(X,β) - y ||_2

Example:

`f(X, β) = β[1]*X[:,1]^p + β[2]*X[:,2]^q`
"""
struct MultivariateLsq <: AbstractModel
    expr
    f
    coefficient
    p₀
end
# -

function show(io::IO, m::MultivariateLsq)
    write(io, "MultivariateLsq(\n  f=$(repr(m.f)),\n  coefficient=$(repr(m.coefficient)),\n  p₀=$(repr(m.p₀))\n)")
end

function model(::Type{MultivariateLsq}, expr::Expr)
    if expr.head == :block
        return model(MultivariateLsq, expr.args[2])
    end
    objective = :l2norm
    f = callsites(expr, :curve_fit)[end].args[3]
    coeff = callsites(expr, f)[1].args[end]
    p₀ = callsites(expr, :curve_fit)[end].args[end]
    return MultivariateLsq(expr, f, coeff, p₀)
end

"""    poly(m::MultivariateLsq)::Expr

find the part of the model that implements the polynomial model for regression.
"""
function poly(m::MultivariateLsq)
    func = findfunc(m.expr, m.f)[1]
    poly = func.args[2].args[end].args[1]
    return poly
end

"""    (t::Pow)(m::MultivariateLsq, i::Int)

apply the power transformation to the i-th coordinate of m.

Example:

If `m` is a program implementing `f(X, β) = β[1]*X[:,1]^p + β[2]*X[:,2]^q`

a) and `t = Pow(2)` then `t(m, 1)` is the model implementing
`f(X, β) = β[1]*X[:,1]^p+2 + β[2]*X[:,2]^q`.

or

b) and `t = Pow(2)` then `t(m, 2)` is the model implementing
`f(X, β) = β[1]*X[:,1]^p + β[2]*X[:,2]^q+2`.

See also [`(t::Pow)(m::MultivariateLsq, i::Int)`](@ref)
"""
function (t::Pow)(m::MultivariateLsq, i::Int)
    p = poly(m)
    comp = p.args[i+1]
    pow = callsites(comp, :(.^))
    pow[end].args[3] += t.inc
    return m
end

"""    (t::Product)(m::MultivariateLsq)

apply the product transformation to the model m.
For a MultivariateLsq model a product transformation works
by applying the transformation component-wise where each component represents a
column of the data.

Example:

If `m` is a program implementing `f(X, β) = β[1]*X[:,1]^p + β[2]*X[:,2]^q`
and `t = Pow(2,-1)` then `t(m)` is the model implementing
`f(X, β) = β[1]*X[:,1]^p+2 + β[2]*X[:,2]^q-1`.
"""
function (t::Product)(m::MultivariateLsq)
    m = m
    for (i, tᵢ) in enumerate(t.dims)
        m = tᵢ(m, i)
    end
    return m
end

# ## Model Augmentation with Group Actions
# the Group `Product{Tuple{Pow{Int},Pow{Int}}}` which is isomorphic to $(Z, +, 0)\times (Z,+,0)$ can act on our class of models
#
# Let's build an instance of the model object from the code snippet `expr`

m = model(MultivariateLsq, expr)
@show m
poly(m)

# We can get the poly from the model and show how the Pow group elements can act on it componentwise.

@show poly((one(Pow)^4)(deepcopy(m), 1))
@show poly((one(Pow)^4)(deepcopy(m), 2))

# Some *basis vectors* will come in handy for building elements of the product group.
# $T_1,T_2$ are generators for $Z\times Z$

@show T₁ = Product((one(Pow), zero(Pow)))
@show T₂ = Product((zero(Pow), one(Pow)))

# Some Checks that we get what is expected.

m′ = @show (T₁^2 ∘ T₂^3)(deepcopy(m))
p = poly(m)
p′ = poly(m′)
@assert p′.args[2].args[3].args[3] == p.args[2].args[3].args[3] + 2
@assert p′.args[3].args[3].args[3] == p.args[3].args[3].args[3]+3
p′

# ## Exploring the Orbits
#
# We use the Levenberg Marquardt algorithm for general least squares problems in data analysis. The goal is to fix a model class and find the best coefficients for that class. Our algebraic representation allows us to have a similar treatment of model class (in this case the exponents $i,j$ in our formula). We can sample from the orbits of the group when applied to the model and solve for the best coefficients in order to find the best model class.
#
# In this case we have reduced the code search space to a model search space to an algebraic formulation.

m = model(MultivariateLsq, expr)
results = []
for i in 1:3
    for j in 1:3
        m₀ = deepcopy(m)
        T = (T₁^i) ∘ (T₂^j)
        mᵢⱼ = T(m₀)
        p = poly(mᵢⱼ)
        M = eval(mᵢⱼ.expr)
        push!(results, (i=i, j=j, M.result..., p=p))
    end
end

# It turns out that this method recovers the *true* model order $2,3$
#
#
#
#

sort(results, by=x->x.r)

# ## Conclusions
#
# We have seen how abstract algebra can be applied to the category of models to build a systematic treatment of model augmentation. This proves that the model transformations can be arranged into a simple algebraic structure that can act on a model to build new models. The structure of the transformations are easier to analyze than the changes to the models themselves.


text:
using DifferentialEquations, Flux, DiffEqFlux
using Plots

u0 = [1.0,1.0]
tspan = (0.0,10.0)
tsteps = 0.0:0.1:10.0

my_nn = Chain(
  Dense(2, 32, relu),
  Dense(32, 2)) |> f64

p, re = Flux.destructure(my_nn)

function right_hand_side(du,u,p,t)
  m = re(p)
  nn_output = m(u)
  du[1] = nn_output[1]
  du[2] = nn_output[2]
end

prob = ODEProblem(right_hand_side,u0,tspan,p)
sol = solve(prob, Tsit5())

# Plot the solution
plot(sol)
savefig("Initial_NN_ODE.png")

function predict_ode(p)
  tmp_prob = remake(prob, p = p)
  Array(solve(tmp_prob, Tsit5(), saveat = tsteps))
end

function loss_ode(p)
    pred = predict_ode(p)
    loss = sum(abs2, 1.0 .- pred)
    return loss, pred
end

callback = function (p, l, pred)
  display(l)
  plt = plot(pred, ylim = (0, 6))
  display(plt)
  # Tell sciml_train to not halt the optimization. If return true, then
  # optimization stops.
  return false
end

result_ode = DiffEqFlux.sciml_train(loss_ode, p,
                                    cb = callback,
                                    maxiters = 100)

prob_final = ODEProblem(right_hand_side, u0, tspan, result_ode)
sol_final = solve(prob_final,Tsit5())
plot(sol_final)
savefig("NN_ode_optimized.png") # Both state components close to 1

max_stars_repo_name: Romit-Maulik/Tutorials-Demos-Practice

max_issues_repo_name: Romit-Maulik/Tutorials-Demos-Practice

max_forks_repo_name: Romit-Maulik/Tutorials-Demos-Practice


text:
f(x)=x^2              # 定义了函数 f 的一个方法，其需要一个不特定类型的参数 (i)
f(x::Float64)=x^2+0.1 # 定义了函数 f 的 第二个方法，需要一个 Float64 类型的参数 (ii)
                      # REPL 会告诉你 f 现在有两个方法
f(1)                  # 1 的类型是 Int，只能调用方法 (i)，返回 1
f("1")                # "1" 的类型是 String，同样调用方法 (i)
                      # 这里看到，执行 x^2 的时候
                      # 实际上已经去调用特定为字符串定义的 ^ 方法了
                      # 会返回 "11"
f(1.0)                # 1.0 的类型是 Float64，调用方法 (ii)，返回 1.1
                      # 这里展现了一个不好的编程实践
                      # 实际工作时，如无特殊情况，要保证 f(1) 和 f(1.0) 的返回值相等

max_stars_repo_name: zhaiyusci/Julia-Notes

max_issues_repo_name: zhaiyusci/Julia-Notes

max_forks_repo_name: zhaiyusci/Julia-Notes


text:
include("abstract_approximators.jl")
include("q_approximators/q_approximators.jl")
include("v_approximators.jl")