text:
using Statistics: mean
function stepColumn!(
    mb :: ModelBlock,
    Δt :: Float64,
)

    fi    = mb.fi
    tmpfi = mb.tmpfi
    co    = mb.co
    ev    = mb.ev
    cfg   = ev.config
    wksp  = co.wksp

#    println("avg of HMXL: ", mean(fi.HMXL[isfinite.(fi.HMXL)]))

    _INTMT_ = view(tmpfi._INTMX_, :, 1)
    _INTMS_ = view(tmpfi._INTMX_, :, 2)

    # Surface fluxes of temperature
    rad = ( co.mtx[:T_swflxConv_sT] * view(fi.SWFLX, :) + co.mtx[:T_sfcflxConv_sT] * view(fi.NSWFLX, :)) / ρcp_sw
    @. _INTMT_ += Δt * rad

    # Surface fluxes of salinity
    fflx = co.mtx[:T_sfcflxConv_sT] * view(fi.VSFLX, :)
    @. _INTMS_ += Δt * fflx

    # Surface fluxes and QFLX

    if cfg["Qflx"] == "on"

        qflxt = co.amo.T_mask_T * view( fi._QFLXX_, :, 1)
        qflxs = co.amo.T_mask_T * view( fi._QFLXX_, :, 2)
        @. _INTMT_ += Δt * qflxt
        @. _INTMS_ += Δt * qflxs

    end

    # *** Diffusion and restoring ***
    # (b_t+1 - b_t) / dt =  OP1 * b_t+1 + OP2 * (b_t+1 - b_target) + const
    # b_t+1 - b_t =  dt * OP1 * b_t+1 + dt * OP2 * (b_t+1 - b_target) + dt * const
    # (I - OP1 * dt - OP2 * dt) b_t+1 = b_t - dt * OP2 * b_target + dt * const
    # b_t+1 = (I - OP1 * dt - OP2 * dt) \ (b_t - dt * OP2 * b_target + dt * const)
    #
    # Change of b
    #
    # Δb = b_t+1 - b_t = Δt * OP1 * b_t+1 + Δt * OP2 * (b_t+1 - b_target) + Δt * const
    #
    # Notice that I carefully use the buoyancy before
    # stepAdvection which means operators should be at
    # t = t0. I empirically know if we use newton method
    # to find a steady state, then this scheme seem to
    # ensure that physical stepping will be steady state
    # too.
    #

    @. fi._b = TS2b(fi.sv[:_TEMP], fi.sv[:_SALT])
    op_vdiff = calOp_vdiff(co.vd, fi._b, fi.HMXL)

    op_TEMP = op_vdiff
    op_SALT = op_vdiff

    # Save this operator for diagnostic purpose
    mb.tmpfi.op_vdiff = op_vdiff

    RHS_TEMP = getSpace!(wksp, :T, true; o=0.0)
    RHS_SALT = getSpace!(wksp, :T, true; o=0.0)

    RHS_TEMP[:] = _INTMT_
    RHS_SALT[:] = _INTMS_

    if cfg["weak_restoring"] == "on"
        op_TEMP   += co.mtx[:T_invτwk_TEMP_T]
        op_SALT   += co.mtx[:T_invτwk_SALT_T]

        idx = isnan.(tmpfi.datastream["TEMP"])
        tmpfi.datastream["TEMP"][idx] .= 0.0
        tmpfi.datastream["SALT"][idx] .= 0.0

        RHS_TEMP .-= Δt * co.amo.T_mask_T * co.mtx[:T_invτwk_TEMP_T] * reshape( tmpfi.datastream["TEMP"] , :)
        RHS_SALT .-= Δt * co.amo.T_mask_T * co.mtx[:T_invτwk_SALT_T] * reshape( tmpfi.datastream["SALT"] , :)
    end

    F_EBM_TEMP = lu( I - Δt * op_TEMP )
    F_EBM_SALT = lu( I - Δt * op_SALT )

    tmpfi.sv[:NEWTEMP][:] = F_EBM_TEMP \ RHS_TEMP
    tmpfi.sv[:NEWSALT][:] = F_EBM_SALT \ RHS_SALT
    #tmpfi.sv[:NEWTEMP][:] = RHS_TEMP
    #tmpfi.sv[:NEWSALT][:] = RHS_SALT

    # Compute the change of tracers attributed to vertical diffusion
   fi.sv[:VDIFFT][:] = op_vdiff * reshape(tmpfi.sv[:NEWTEMP], :, 1)
   fi.sv[:VDIFFS][:] = op_vdiff * reshape(tmpfi.sv[:NEWSALT], :, 1)

    #=
    beg_idx = 0
    jmp = ev.Nz*ev.Nx
    for i=1:ev.Ny
        rng = beg_idx+1:beg_idx+jmp
        sub_op_TEMP  = view(op_TEMP, rng, rng)
        sub_op_SALT  = view(op_SALT, rng, rng)

        sub_RHS_TEMP = view(RHS_TEMP, rng)
        sub_RHS_SALT = view(RHS_SALT, rng)

        F_EBM_TEMP = lu( I - Δt * sub_op_TEMP )
        F_EBM_SALT = lu( I - Δt * sub_op_SALT )

        tmpfi.sv[:NEWTEMP][rng] = F_EBM_TEMP \ sub_RHS_TEMP
        tmpfi.sv[:NEWSALT][rng] = F_EBM_SALT \ sub_RHS_SALT

        beg_idx += jmp
    end
    =#

    # Compute source and sink of tracers due to weak restoring
    if cfg["weak_restoring"] == "on"
        tmpfi._WKRSTΔX_[:, 1] = tmpfi._NEWX_[:, 1] - reshape(tmpfi.datastream["TEMP"], :)
        tmpfi._WKRSTΔX_[:, 2] = tmpfi._NEWX_[:, 2] - reshape(tmpfi.datastream["SALT"], :)
        fi._WKRSTX_[:, 1] = co.mtx[:T_invτwk_TEMP_T] * view(tmpfi._WKRSTΔX_, :, 1)
        fi._WKRSTX_[:, 2] = co.mtx[:T_invτwk_SALT_T] * view(tmpfi._WKRSTΔX_, :, 2)
    else
        fi._WKRSTX_ .= 0.0
    end

    #
    # Freezing Potential
    #
    # Please read:
    # "CICE: the Los Alamos Sea Ice Model Documentation and Software User’s Manual"
    # https://csdms.colorado.edu/w/images/CICE_documentation_and_software_user's_manual.pdf
    #
    # Section 2.2 "Ocean":
    #
    # New sea ice forms when the ocean temperature drops below its freezing temperature, Tf = −µS, where S
    # is the seawater salinity and µ = 0.054 psu is the ratio of the freezing temperature of brine to its salinity.
    # The ocean model performs this calculation; if the freezing/melting potential Ffrzmlt is positive, its value
    # represents a certain amount of frazil ice that has formed in one or more layers of the ocean and floated to the
    # surface. (The ocean model assumes that the amount of new ice implied by the freezing potential actually
    # forms.) In general, this ice is added to the thinnest ice category. The new ice is grown in the open water area
    # of the grid cell to a specified minimum thickness; if the open water area is nearly zero or if there is more
    # new ice than will fit into the thinnest ice category, then the new ice is spread over the entire cell.
    # If Ffrzmlt is negative, it is used to heat already existing ice from below. In particular, the sea surface
    # temperature and salinity are used to compute an oceanic heat flux Fw (|Fw| ≤ |Ffrzmlt|) which is applied at
    # the bottom of the ice. The portion of the melting potential actually used to melt ice is returned to the coupler
    # in Fhocn.
    #
    # By reading the code of CICE, I know that Q_FRZMLTPOT in my model is
    # the variable "frzmlt" in CICE.
    #
    # If frzmlt < 0:
    # In ice_therm_vertical.F90 of lines 703-740, CICE model uses
    # frzmlt only when it is negative (T_ocean > T_sw_frz) to
    # compute the melting of sea ice.
    #
    # If frzmlt > 0:
    # In ice_therm_itd.F90 of lines 1000-1022, CIMICE uses positive frzmlt
    # to compute the amount of newly formed sea ice.
    #
    # So, frzmlt > 0 has to be consistent with my energy gain due to freezing.
    # I can safely adapt the original "instant response" scheme to Newtonion
    # relaxation.
    #
    # If frzmlt < 0, I guess this melting potential is used by CICE model and
    # recompute the actual heat that would be taken away by sea ice, and this
    # amount of energy is then passed to "Fioi_melth". So it looks like I cannot
    # modify its definition because frzmlt positive and negative do not share
    # the same concep. frzmlt > 0 represents the actual energy transfer while
    # frzmlt < 0 represents temperature differences.
    #
    # *** This is potentially a CESM design flaw because freezing and
    # melting can happen at the same time. For example, ocean
    # surface is above freezing point and the bottom of ocean is
    # below freezing point. So, a way to imagine this still works is to modify
    # the convective adjustment scheme:
    #
    # If freezing happens somewhere below the ocean surface (top grid), the
    # formed sea ice will float upward due to less density. This indirectly
    # produce turbulence, and hence increase the vertical diffusivity above
    # this deep ocean grid.
    #
    # Also, I surprisingly find that sensible heat flux between atm and ocn
    # is not affected by the presence of sea ice. At least I cannot find anything
    # considering the sea ice when the coupler computes sensible heat flux in
    # "csm_share/shr/shr_flux_mod.F90" "Faox_sen, Foxx_sen"
    #
    # Maybe it is worth asking this question on CESM forum.
    #

    # *** Restoring ***
    # (b_t+1 - b_t) / dt =  OP1 * b_t+1 + OP2 * (b_t+1 - b_target) + const
    # b_t+1 - b_t =  dt * OP1 * b_t+1 + dt * OP2 * (b_t+1 - b_target) + dt * const
    # (I - OP1 * dt - OP2 * dt) b_t+1 = b_t - dt * OP2 * b_target + dt * const
    # b_t+1 = (I - OP1 * dt - OP2 * dt) \ (b_t - dt * OP2 * b_target + dt * const)
    # Δb = b_t+1 - b_t = Δt * OP1 * b_t+1 + Δt * OP2 * (b_t+1 - b_target) + Δt * const

    # Tackle freeze / melt potential
    ΔT_sT  = getSpace!(wksp, :sT, true)
    tmp_sT = getSpace!(wksp, :sT, true)

    NEWSST = reshape(view(tmpfi.sv[:NEWTEMP], 1:1, :, :), :)
    sfcΔz_sT = reshape(view(co.amo.gd.Δz_T, 1:1, :, :), :)

    Q_FRZHEAT        = reshape(fi.Q_FRZHEAT,       :)
    Q_FRZMLTPOT      = reshape(fi.Q_FRZMLTPOT,     :)
    Q_FRZMLTPOT_NEG  = reshape(fi.Q_FRZMLTPOT_NEG, :)
    Q_LOST           = reshape(fi.Q_LOST,          :)

    Q_FRZHEAT       .= 0.0
    Q_FRZMLTPOT     .= 0.0
    Q_FRZMLTPOT_NEG .= 0.0
    Q_LOST          .= 0.0

    @. ΔT_sT = NEWSST - T_sw_frz
    below_frz_mask_sT = ΔT_sT .< 0.0
    above_frz_mask_sT = ΔT_sT .>= 0.0

    tmp_sT .= 0.0
    tmp_sT[below_frz_mask_sT] .= - 1.0 / cfg["τfrz"]
    op_frz = co.amo_slab.T_mask_T * spdiagm(0 => tmp_sT)

    tmp_sT .= T_sw_frz
    NEWSST[:] = lu( I - Δt * op_frz ) \ ( NEWSST - Δt * op_frz * tmp_sT )

    Q_FRZHEAT[:]   = ρcp_sw * sfcΔz_sT .* (op_frz * (NEWSST .- T_sw_frz))

    # Capping freezing potential at 10 W/m^2 to avoid instability
    # in cice model due to advection noise in Ekman flow. However,
    # this breaks energy conservation so we record this with Q_LOST.
    for (i, q_frzheat) in enumerate(Q_FRZHEAT)
        if q_frzheat > 10.0
            Q_FRZHEAT[i] = 10.0
            Q_LOST[i] = q_frzheat - 10.0
        end
    end
    #Q_FRZHEAT[Q_FRZHEAT .> 10.0] .= 10.0

    @. tmp_sT = - ρcp_sw * sfcΔz_sT * ΔT_sT / Δt

    Q_FRZMLTPOT_NEG[above_frz_mask_sT] = tmp_sT[above_frz_mask_sT]

    Q_FRZMLTPOT[above_frz_mask_sT]     = tmp_sT[above_frz_mask_sT]
    Q_FRZMLTPOT[below_frz_mask_sT]     = Q_FRZHEAT[below_frz_mask_sT]

    #=
    if any(fi.Q_FRZMLTPOT_NEG .> 0)
        throw(ErrorException("[1] Something is wrong when computing freeze melt potential."))
    end

    if any(fi.Q_FRZHEAT .< 0)
        throw(ErrorException("[2] Something is wrong when computing freeze melt potential."))
    end
    =#
end

max_stars_repo_name: meteorologytoday/IOM

max_issues_repo_name: meteorologytoday/IOM

max_forks_repo_name: meteorologytoday/IOM


text:
push!(LOAD_PATH,"../src/")

using Documenter
using Literate
using Plots
using VlasovSolvers

ENV["GKSwstype"] = "100"

makedocs(modules=[VlasovSolvers],
         sitename = "VlasovSolvers.jl",
         authors="Julia Vlasov",
         format=Documenter.HTML(;
         prettyurls=get(ENV, "CI", "false") == "true",
         canonical="https://juliavlasov.github.io/VlasovSolvers.jl",
         assets=String[],
         ),
         doctest = false,
         pages = ["Home"           => "index.md",
                  "Problem type"   => "problem.md",
                  "Vlasov-Poisson" => "vlasov-poisson.md",
                  "Vlasov-Ampere"  => "vlasov-ampere.md",
                  "Bump On Tail"   => "bump_on_tail.md",
                  "Rotation 2D"    => "rotation2d.md",
                  "Vlasov-HMF"     => "vlasov-hmf.md",
                  "Two-stream instability" => "tsi.md",
                  "Contents"       => "contents.md"])

deploydocs(;
    branch = "gh-pages",
    devbranch = "master",
    repo   = "github.com/JuliaVlasov/VlasovSolvers.jl"
)

ext: jl

max_stars_repo_name: JuliaVlasov/VlasovSolvers

max_issues_repo_name: JuliaVlasov/VlasovSolvers

max_forks_repo_name: JuliaVlasov/VlasovSolvers


text:
function pass_names(dest, src, map)
    for vi in MOI.get(src, MOI.ListOfVariableIndices())
        name = MOI.get(src, MOI.VariableName(), vi)
        if name != ""
            MOI.set(dest, MOI.VariableName(), map[vi], name)
        end
    end
    for (F,S) in MOI.get(src, MOI.ListOfConstraints())
        for con in MOI.get(src, MOI.ListOfConstraintIndices{F,S}())
            name = MOI.get(src, MOI.ConstraintName(), con)
            if name != ""
                MOI.set(dest, MOI.ConstraintName(), map[con], name)
            end
        end
    end
end

function append_to(dest::MOI.ModelLike, src::MOI.ModelLike, idxmap, copy_names::Bool,
    filter_constraints::Union{Nothing, Function}=nothing; allow_single_bounds::Bool = true)

    #=
        This function follows closely the function `default_copy_to` defined in
        MathOptInterface.Utilities
        due to some caveats of this function we keep the commented functions
        from the original function to highlight the differences and
        easen the burden of updating when `default_copy_to` is updated.
    =#

    # MOI.empty!(dest)

    # idxmap = MOIU.IndexMap()

    vis_src = MOI.get(src, MOI.ListOfVariableIndices())
    # index_map_for_variable_indices only initializes the data structure
    # idxmap = index_map_for_variable_indices(vis_src)

    # The `NLPBlock` assumes that the order of variables does not change (#849)
    if MOI.NLPBlock() in MOI.get(src, MOI.ListOfModelAttributesSet())
        error("NLP models are not supported.")
        # constraint_types = MOI.get(src, MOI.ListOfConstraints())
        # single_variable_types = [S for (F, S) in constraint_types
        #                          if F == MOI.SingleVariable]
        # vector_of_variables_types = [S for (F, S) in constraint_types
        #                              if F == MOI.VectorOfVariables]
        # vector_of_variables_not_added = [
        #     MOI.get(src, MOI.ListOfConstraintIndices{MOI.VectorOfVariables, S}())
        #     for S in vector_of_variables_types
        # ]
        # single_variable_not_added = [
        #     MOI.get(src, MOI.ListOfConstraintIndices{MOI.SingleVariable, S}())
        #     for S in single_variable_types
        # ]
    else
        # the key asusmption here is that MOI keeps the following behaviour
        # "The copy is only done when
        # the variables to be copied are not already keys of `idxmap`. It returns a list
        # of the constraints copied and not copied."
        # from copy_single_variable and copy_vector_of_variables.
        # this is very importante because variables are shered between
        # upper, lower and lower dual levels
        vector_of_variables_types, _, vector_of_variables_not_added,
        single_variable_types, _, single_variable_not_added = MOIU.try_constrain_variables_on_creation(
            dest, src, idxmap, MOI.add_constrained_variables, MOI.add_constrained_variable
        )
    end

    # MOIU.copy_free_variables(dest, idxmap, vis_src, MOI.add_variables)
    # copy variables has a size check that dows not generalizes here
    # because we have previously added variables
    for vi in vis_src
        if !haskey(idxmap.varmap, vi)
            var = MOI.add_variable(dest)
            idxmap.varmap[vi] = var
        end
    end

    # Copy variable attributes
    MOIU.pass_attributes(dest, src, copy_names, idxmap, vis_src)

    # Copy model attributes
    # attention HERE to no pass objective functions!
    # pass_attributes(dest, src, copy_names, idxmap)

    # Copy constraints
    MOIU.pass_constraints(dest, src, copy_names, idxmap,
                     single_variable_types, single_variable_not_added,
                     vector_of_variables_types, vector_of_variables_not_added,
                     filter_constraints=filter_constraints)

    return idxmap
end

# scalar
function MOIU.promote_operation(::typeof(LinearAlgebra.dot), ::Type{T},
    ::Type{<:Union{MOI.SingleVariable, MOI.ScalarAffineFunction{T}}},
    ::Type{T}
    ) where T
    MOI.ScalarAffineFunction{T}
end
function MOIU.promote_operation(::typeof(LinearAlgebra.dot), ::Type{T},
    ::Type{T},
    ::Type{<:Union{MOI.SingleVariable, MOI.ScalarAffineFunction{T}}}
    ) where T
    MOI.ScalarAffineFunction{T}
end
function MOIU.promote_operation(::typeof(LinearAlgebra.dot), ::Type{T},
    ::Type{<:Union{MOI.SingleVariable, MOI.ScalarAffineFunction{T}}},
    ::Type{<:Union{MOI.SingleVariable, MOI.ScalarAffineFunction{T}}}
    ) where T
    MOI.ScalarQuadraticFunction{T}
end
function MOIU.promote_operation(::typeof(LinearAlgebra.dot), ::Type{T},
    ::Type{MOI.ScalarQuadraticFunction{T}},
    ::Type{T}
    ) where T
    MOI.ScalarQuadraticFunction{T}
end
function MOIU.promote_operation(::typeof(LinearAlgebra.dot), ::Type{T},
    ::Type{T},
    ::Type{MOI.ScalarQuadraticFunction{T}}
    ) where T
    MOI.ScalarQuadraticFunction{T}
end
# flip
function MOIU.operate(::typeof(LinearAlgebra.dot), ::Type{T},
    f::Union{
        MOI.SingleVariable,
        MOI.ScalarAffineFunction{T},
        MOI.ScalarQuadraticFunction{T}
        },
    α::T) where T
    return MOIU.operate(LinearAlgebra.dot, T, α, f)
end
# pass to *
function MOIU.operate(::typeof(LinearAlgebra.dot), ::Type{T},
    f::Union{
        T,
        MOI.SingleVariable,
        MOI.ScalarAffineFunction{T}
        },
    g::Union{
        MOI.SingleVariable,
        MOI.ScalarAffineFunction{T}
        }
    ) where T
    return MOIU.operate(*, T, f, g)
end
function MOIU.operate(::typeof(LinearAlgebra.dot), ::Type{T},
    α::T,
    f::MOI.ScalarQuadraticFunction{T}
    ) where T
    return MOIU.operate(*, T, f, α)
end

# vector
function MOIU.promote_operation(::typeof(LinearAlgebra.dot), ::Type{T},
    ::Type{<:Union{MOI.VectorOfVariables, MOI.VectorAffineFunction{T}}},
    ::Type{Vector{T}}
    ) where T
    MOI.VectorAffineFunction{T}
end
function MOIU.promote_operation(::typeof(LinearAlgebra.dot), ::Type{T},
    ::Type{Vector{T}},
    ::Type{<:Union{MOI.VectorOfVariables, MOI.VectorAffineFunction{T}}}
    ) where T
    MOI.VectorAffineFunction{T}
end
function MOIU.promote_operation(::typeof(LinearAlgebra.dot), ::Type{T},
    ::Type{<:Union{MOI.VectorOfVariables, MOI.VectorAffineFunction{T}}},
    ::Type{<:Union{MOI.VectorOfVariables, MOI.VectorAffineFunction{T}}}
    ) where T
    MOI.VectorQuadraticFunction{T}
end
function MOIU.promote_operation(::typeof(LinearAlgebra.dot), ::Type{T},
    ::Type{MOI.VectorQuadraticFunction{T}},
    ::Type{Vector{T}}
    ) where T
    MOI.VectorQuadraticFunction{T}
end
function MOIU.promote_operation(::typeof(LinearAlgebra.dot), ::Type{T},
    ::Type{Vector{T}},
    ::Type{MOI.VectorQuadraticFunction{T}}
    ) where T
    MOI.VectorQuadraticFunction{T}
end
# flip
function MOIU.operate(::typeof(LinearAlgebra.dot), ::Type{T},
    f::Union{
        MOI.VectorOfVariables,
        MOI.VectorAffineFunction{T},
        MOI.VectorQuadraticFunction{T}
        },
    α::Vector{T}) where T
    return MOIU.operate(LinearAlgebra.dot, T, α, f)
end
# pass to _operate(LinearAlgebra.dot, ...)
function MOIU.operate(::typeof(LinearAlgebra.dot), ::Type{T},
    f::Union{
        Vector{T},
        MOI.VectorOfVariables,
        MOI.VectorAffineFunction{T}
        },
    g::Union{
        MOI.VectorOfVariables,
        MOI.VectorAffineFunction{T}
        }
    ) where T
    return _operate(LinearAlgebra.dot, T, f, g)
end
function MOIU.operate(::typeof(LinearAlgebra.dot), ::Type{T},
    α::T,
    f::MOI.VectorQuadraticFunction{T}
    ) where T
    return _operate(LinearAlgebra.dot, T, f, α)
end
function _operate(::typeof(LinearAlgebra.dot), ::Type{T},
    f::Union{
        Vector{T},
        MOI.VectorOfVariables,
        MOI.VectorAffineFunction{T},
        MOI.VectorQuadraticFunction{T}
        },
    g::Union{
        MOI.VectorOfVariables,
        MOI.VectorAffineFunction{T},
        MOI.VectorQuadraticFunction{T}
    }) where T

    dim = MOI.output_dimension(g)
    if MOI.output_dimension(f) != dim
        throw(DimensionMismatch("f and g are of different MOI.output_dimension's!"))
    end

    fs = MOIU.scalarize(f)
    gs = MOIU.scalarize(g)

    out = MOIU.operate(*, T, fs[1], gs[1])
    for i in 2:dim
        MOIU.operate!(+, T, out, MOIU.operate(*, T, fs[i], gs[i]))
    end

    return out
end
MOIU.scalarize(v::Vector{T}) where T<:Number = v
MOI.output_dimension(v::Vector{T}) where T<:Number = length(v)#


text:
##
using MRIgeneralizedBloch
using DifferentialEquations
using StaticArrays
using Test

## choose random parameters
α = 0
TFP = 100e-6 + 900e-6 * rand()
ω1 = α
B1 = 1
ω0 = 1000 * rand() - 500
m0s = 0.2 * rand()
m0f = 1 - m0s
R1f = 0.5 + 0.2 * rand()
R1s = 2 + rand()
R2f = 1 / (40e-3 + 100e-3 * rand())
T2s = 5e-6 + 10e-6 * rand()
R2s = 1 / T2s
Rx = 100 * rand()

ms = m0s * rand()
mf = (1 - m0s) * rand()
ϑ = rand() * π / 2
φ = rand() * 2π

## Solve generalized Bloch-McConnell with Lorentzian lineshape
p = (ω0, m0s, R1f, R2f, Rx, R1s)
alg = Tsit5()
m0 = [mf * sin(ϑ) * cos(φ), mf * sin(ϑ) * sin(φ), mf * cos(ϑ), ms, 1]

sol = solve(
    ODEProblem(MRIgeneralizedBloch.apply_hamiltonian_freeprecession!, m0, (0.0, TFP), p),
    alg,
)

u_gBloch = sol[end]

## Solve original Bloch-McConnell with Lorentzian lineshape
H = @SMatrix [
    -R2f  -ω0       B1*ω1      0           0       0
      ω0 -R2f           0      0           0       0
  -B1*ω1    0 -R1f-Rx*m0s      0      Rx*m0f R1f*m0f
       0    0           0   -R2s       B1*ω1       0
       0    0      Rx*m0s -B1*ω1 -R1s-Rx*m0f R1s*m0s
       0    0           0      0           0       0
]

m0 = [mf * sin(ϑ) * cos(φ), mf * sin(ϑ) * sin(φ), mf * cos(ϑ), 0, ms, 1]
u_Bloch = exp(H * TFP) * m0

## Test!
max_error = 1e-3
@test u_gBloch[1] ≈ u_Bloch[1] atol = max_error
@test u_gBloch[2] ≈ u_Bloch[2] atol = max_error
@test u_gBloch[3] ≈ u_Bloch[3] atol = max_error
@test u_gBloch[4] ≈ u_Bloch[5] atol = max_error
@test u_gBloch[5] ≈ u_Bloch[6] atol = max_error


text:
# Ensure that the `@benchmark` user-facing API supports the following cases:
using Benchmarks
# no-op function calls:
f() = nothing
@benchmark f()

# infix operators
@benchmark (3.0+5im)^3.2

# indexing
A = rand(2,2)
@benchmark A[end,end]

# Keyword arguments
@benchmark svds(A, nsv=1)

# local scopes
x = 1
let B = copy(A), y = 2
    @benchmark B[1]
    @benchmark x+y
end


text:
module Stripe82Score

import JLD
import FITSIO
using DataFrames

import ..SDSSIO
import ..SDSSIO: RunCamcolField
import ..Model: CatalogEntry, ids

immutable MatchException <: Exception
    msg::String
end

"""
mag_to_flux(m)

convert SDSS mags to SDSS flux
"""
mag_to_flux(m::AbstractFloat) = 10.^(0.4 * (22.5 - m))
@vectorize_1arg AbstractFloat mag_to_flux

flux_to_mag(nm::AbstractFloat) = nm > 0 ? 22.5 - 2.5 * log10(nm) : NaN
@vectorize_1arg AbstractFloat flux_to_mag

"""
where(condition, x, y)

Construct a new Array containing elements from `x` where `condition` is true
otherwise elements from `y`.
"""
function where(condition, x, y)
    @assert length(condition) == length(x) == length(y)
    out = similar(x)
    for i=1:length(condition)
        out[i] = condition[i]? x[i]: y[i]
    end
    return out
end

"""
Return distance in pixels using small-distance approximation. Falls
apart at poles and RA boundary.
"""
dist(ra1, dec1, ra2, dec2) = (3600 / 0.396) * (sqrt((dec2 - dec1).^2 +
                                  (cos(dec1) .* (ra2 - ra1)).^2))

"""
match_position(ras, decs, ra, dec, dist)

Return index of first position in `ras`, `decs` that is within a distance
`maxdist` (in degrees) of the target position `ra`, `dec`. If none found,
an exception is raised.
"""
function match_position(ras, decs, ra, dec, maxdist)
    @assert length(ras) == length(decs)
    for i in 1:length(ras)
        dist(ra, dec, ras[i], decs[i]) < maxdist && return i
    end
    throw(MatchException(@sprintf("No source found at %f  %f", ra, dec)))
end

"""
load_s82(fname)

Load Stripe 82 objects into a DataFrame. `fname` should be a FITS file
created by running a CasJobs (skyserver.sdss.org/casjobs/) query
on the Stripe82 database. Run the following query in the \"Stripe82\"
context, then download the table as a FITS file.

```
select
  objid, rerun, run, camcol, field, flags,
  ra, dec, probpsf,
  psfmag_u, psfmag_g, psfmag_r, psfmag_i, psfmag_z,
  devmag_u, devmag_g, devmag_r, devmag_i, devmag_z,
  expmag_u, expmag_g, expmag_r, expmag_i, expmag_z,
  fracdev_r,
  devab_r, expab_r,
  devphi_r, expphi_r,
  devrad_r, exprad_r
into mydb.s82_0_1_0_1
from stripe82.photoobj
where
  run in (106, 206) and
  ra between 0. and 1. and
  dec between 0. and 1.
```
"""
function load_s82(fname)

    # First, simply read the FITS table into a dictionary of arrays.
    f = FITSIO.FITS(fname)
    keys = [:objid, :rerun, :run, :camcol, :field, :flags,
            :ra, :dec, :probpsf,
            :psfmag_u, :psfmag_g, :psfmag_r, :psfmag_i, :psfmag_z,
            :devmag_u, :devmag_g, :devmag_r, :devmag_i, :devmag_z,
            :expmag_u, :expmag_g, :expmag_r, :expmag_i, :expmag_z,
            :fracdev_r,
            :devab_r, :expab_r,
            :devphi_r, :expphi_r,
            :devrad_r, :exprad_r,
            :flags]
    objs = Dict(key=>read(f[2], string(key)) for key in keys)
    close(f)

    usedev = objs[:fracdev_r] .> 0.5  # true=> use dev, false=> use exp

    # Convert to "celeste" style results.
    gal_mag_u = where(usedev, objs[:devmag_u], objs[:expmag_u])
    gal_mag_g = where(usedev, objs[:devmag_g], objs[:expmag_g])
    gal_mag_r = where(usedev, objs[:devmag_r], objs[:expmag_r])
    gal_mag_i = where(usedev, objs[:devmag_i], objs[:expmag_i])
    gal_mag_z = where(usedev, objs[:devmag_z], objs[:expmag_z])

    result = DataFrame()
    result[:objid] = objs[:objid]
    result[:ra] = objs[:ra]
    result[:dec] = objs[:dec]
    result[:is_star] = [x != 0 for x in objs[:probpsf]]
    result[:star_mag_r] = objs[:psfmag_r]
    result[:gal_mag_r] = gal_mag_r

    # star colors
    result[:star_color_ug] = objs[:psfmag_u] .- objs[:psfmag_g]
    result[:star_color_gr] = objs[:psfmag_g] .- objs[:psfmag_r]
    result[:star_color_ri] = objs[:psfmag_r] .- objs[:psfmag_i]
    result[:star_color_iz] = objs[:psfmag_i] .- objs[:psfmag_z]

    # gal colors
    result[:gal_color_ug] = gal_mag_u .- gal_mag_g
    result[:gal_color_gr] = gal_mag_g .- gal_mag_r
    result[:gal_color_ri] = gal_mag_r .- gal_mag_i
    result[:gal_color_iz] = gal_mag_i .- gal_mag_z

    # gal shape -- fracdev
    result[:gal_fracdev] = objs[:fracdev_r]

    # Note that the SDSS photo pipeline doesn't constrain the de Vaucouleur
    # profile parameters and exponential disk parameters (A/B, angle, scale)
    # to be the same, whereas Celeste does. Here, we pick one or the other
    # from SDSS, based on fracdev - we'll get the parameters corresponding
    # to the dominant component. Later, we limit comparison to objects with
    # fracdev close to 0 or 1 to ensure that we're comparing apples to apples.

    result[:gal_ab] = where(usedev, objs[:devab_r], objs[:expab_r])

    # gal effective radius (re)
    re_arcsec = where(usedev, objs[:devrad_r], objs[:exprad_r])
    re_pixel = re_arcsec / 0.396
    result[:gal_scale] = re_pixel

    # gal angle (degrees)
    raw_phi = where(usedev, objs[:devphi_r], objs[:expphi_r])
    result[:gal_angle] = raw_phi - floor(raw_phi / 180) * 180

    return result
end

"""
Convert two fluxes to a color: mag(f1) - mag(f2), assuming the same zeropoint.
Returns NaN if either flux is nonpositive.
"""
function fluxes_to_color(f1::Real, f2::Real)
    (f1 <= 0. || f2 <= 0.) && return NaN
    return -2.5 * log10(f1 / f2)
end
@vectorize_2arg Real fluxes_to_color

"""
Load the SDSS photoObj catalog used to initialize celeste, and reformat column
names to match what the rest of the scoring code expects.
"""
function load_primary(rcf::RunCamcolField, stagedir::String)
    dir = @sprintf "%s/%d/%d/%d" stagedir rcf.run rcf.camcol rcf.field
    fname = @sprintf "%s/photoObj-%06d-%d-%04d.fits" dir rcf.run rcf.camcol rcf.field
    println(fname)
    objs = SDSSIO.read_photoobj(fname)

    usedev = objs["frac_dev"] .> 0.5  # true=> use dev, false=> use exp

    gal_flux_u = where(usedev, objs["devflux_u"], objs["expflux_u"])
    gal_flux_g = where(usedev, objs["devflux_g"], objs["expflux_g"])
    gal_flux_r = where(usedev, objs["devflux_r"], objs["expflux_r"])
    gal_flux_i = where(usedev, objs["devflux_i"], objs["expflux_i"])
    gal_flux_z = where(usedev, objs["devflux_z"], objs["expflux_z"])

    result = DataFrame()
    result[:objid] = objs["objid"]
    result[:ra] = objs["ra"]
    result[:dec] = objs["dec"]
    result[:is_star] = objs["is_star"]
    result[:star_mag_r] = flux_to_mag(objs["psfflux_r"])
    result[:gal_mag_r] = flux_to_mag(gal_flux_r)

    # star colors
    result[:star_color_ug] = fluxes_to_color(objs["psfflux_u"], objs["psfflux_g"])
    result[:star_color_gr] = fluxes_to_color(objs["psfflux_g"], objs["psfflux_r"])
    result[:star_color_ri] = fluxes_to_color(objs["psfflux_r"], objs["psfflux_i"])
    result[:star_color_iz] = fluxes_to_color(objs["psfflux_i"], objs["psfflux_z"])

    # gal colors
    result[:gal_color_ug] = fluxes_to_color(gal_flux_u, gal_flux_g)
    result[:gal_color_gr] = fluxes_to_color(gal_flux_g, gal_flux_r)
    result[:gal_color_ri] = fluxes_to_color(gal_flux_r, gal_flux_i)
    result[:gal_color_iz] = fluxes_to_color(gal_flux_i, gal_flux_z)

    # gal shape -- fracdev
    result[:gal_fracdev] = objs["frac_dev"]

    # gal shape -- axis ratio
    #TODO: filter when 0.5 < frac_dev < .95
    result[:gal_ab] = where(usedev, objs["ab_dev"], objs["ab_exp"])

    # gal effective radius (re)
    re_arcsec = where(usedev, objs["theta_dev"], objs["theta_exp"])
    result[:gal_scale] = re_arcsec / 0.396 #pixel scale

    # gal angle (degrees)
    raw_phi = where(usedev, objs["phi_dev"], objs["phi_exp"])
    result[:gal_angle] = raw_phi - floor(raw_phi / 180) * 180

    return result#[!objs["is_saturated"], :]
end

const color_names = ["ug", "gr", "ri", "iz"]

"""
This function loads one catalog entry into row of i of df, a results data
frame.
ce = Catalog Entry, a row of an astronomical catalog
"""
function load_ce!(i::Int, ce::CatalogEntry, df::DataFrame)
    df[i, :ra] = ce.pos[1]
    df[i, :dec] = ce.pos[2]
    df[i, :is_star] = ce.is_star ? 1. : 0.

    for j in 1:2
        s_type = ["star", "gal"][j]
        fluxes = j == 1 ? ce.star_fluxes : ce.gal_fluxes
        df[i, Symbol("$(s_type)_mag_r")] = flux_to_mag(fluxes[3])
        for c in 1:4
            cc = Symbol("$(s_type)_color_$(color_names[c])")
            cc_sd = Symbol("$(s_type)_color_$(color_names[c])_sd")
            if fluxes[c] > 0 && fluxes[c + 1] > 0  # leave as NA otherwise
                df[i, cc] = -2.5log10(fluxes[c] / fluxes[c + 1])
            else
                df[i, cc] = NaN
            end
        end
    end

    df[i, :gal_fracdev] = ce.gal_frac_dev
    df[i, :gal_ab] = ce.gal_ab
    df[i, :gal_angle] = (180/pi)ce.gal_angle
    df[i, :gal_scale] = ce.gal_scale
    df[i, :objid] = ce.objid
end

"""
Convert Celeste results to a dataframe.
"""
function celeste_to_df(results::Vector{Dict})
    # Initialize dataframe
    N = length(results)
    color_col_names = ["color_$cn" for cn in color_names]
    color_sd_col_names = ["color_$(cn)_sd" for cn in color_names]
    col_names = vcat(["objid", "ra", "dec", "is_star", "star_mag_r",
                      "star_mag_r_sd", "gal_mag_r", "gal_mag_r_sd"],
                     ["star_$c" for c in color_col_names],
                     ["star_$c" for c in color_sd_col_names],
                     ["gal_$c" for c in color_col_names],
                     ["gal_$c" for c in color_sd_col_names],
                     ["gal_fracdev", "gal_ab", "gal_angle", "gal_scale"])
    col_Symbols = Symbol[Symbol(cn) for cn in col_names]
    col_types = Array(DataType, length(col_names))
    fill!(col_types, Float64)
    col_types[1] = String
    df = DataFrame(col_types, N)
    names!(df, col_Symbols)

    # Fill dataframe row-by-row.
    i = 0
    for result in results
        i += 1
        vs = result["vs"]

        function get_fluxes(i::Int)
            ret = Array(Float64, 5)
            ret[3] = exp(vs[ids.r1[i]] + 0.5 * vs[ids.r2[i]])
            ret[4] = ret[3] * exp(vs[ids.c1[3, i]])
            ret[5] = ret[4] * exp(vs[ids.c1[4, i]])
            ret[2] = ret[3] / exp(vs[ids.c1[2, i]])
            ret[1] = ret[2] / exp(vs[ids.c1[1, i]])
            ret
        end

        ce = CatalogEntry(
            vs[ids.u],
            vs[ids.a[1, 1]] > 0.5,
            get_fluxes(1),
            get_fluxes(2),
            vs[ids.e_dev],
            vs[ids.e_axis],
            vs[ids.e_angle],
            vs[ids.e_scale],
            result["objid"],
            result["thing_id"])
        load_ce!(i, ce, df)

        df[i, :is_star] = vs[ids.a[1, 1]]

        #TODO: update UQ to mag units not flux. Also, log-normal now, not gamma.
#        for j in 1:2
#            s_type = ["star", "gal"][j]
#            df[i, Symbol("$(s_type)_flux_r_sd")] =
#                sqrt(df[i, Symbol("$(s_type)_flux_r")]) * vs[ids.r2[j]]
#            for c in 1:4
#                cc_sd = Symbol("$(s_type)_color_$(color_names[c])_sd")
#                df[i, cc_sd] = 2.5 * log10(e) * vs[ids.c2[c, j]]
#            end
#        end
    end

    return df
end

"""
Given two results data frame, one containing ground truth (i.e Coadd)
and one containing predictions (i.e., either Primary of Celeste),
compute an a data frame containing each prediction's error.
(It's not an average of the errors, it's each error.)
Let's call the return type of this function an \"error data frame\".
"""
function get_err_df(truth::DataFrame, predicted::DataFrame)
    color_cols = [Symbol("color_$cn") for cn in color_names]
    abs_err_cols = [:gal_fracdev, :gal_ab, :gal_scale]
    col_Symbols = vcat([:objid, :position, :missed_stars,
                        :missed_gals, :mag_r],
                       color_cols,
                       abs_err_cols,
                       :gal_angle)

    col_types = fill(Float64, length(col_Symbols))
    col_types[1] = String
    col_types[3] = col_types[4] = Bool
    ret = DataFrame(col_types, size(truth, 1))
    names!(ret, col_Symbols)
    ret[:objid] = truth[:objid]

    predicted_gal = predicted[:is_star] .< .5
    true_gal = truth[:is_star] .< .5
    ret[:missed_stars] =  predicted_gal & !(true_gal)
    ret[:missed_gals] =  !predicted_gal & true_gal

    ret[:position] = dist(truth[:ra], truth[:dec],
                          predicted[:ra], predicted[:dec])

    ret[true_gal, :mag_r] =
        abs(truth[true_gal, :gal_mag_r] - predicted[true_gal, :gal_mag_r])
    ret[!true_gal, :mag_r] =
        abs(truth[!true_gal, :star_mag_r] - predicted[!true_gal, :star_mag_r])

    for cn in color_names
        ret[true_gal, Symbol("color_$cn")] =
            abs(truth[true_gal, Symbol("gal_color_$cn")] -
                predicted[true_gal, Symbol("gal_color_$cn")])
        ret[!true_gal, Symbol("color_$cn")] =
            abs(truth[!true_gal, Symbol("star_color_$cn")] -
                predicted[!true_gal, Symbol("star_color_$cn")])
    end

    for n in abs_err_cols
        ret[n] = abs(predicted[n] - truth[n])
    end

    function degrees_to_diff(a, b)
        angle_between = abs(a - b) % 180
        min(angle_between, 180 - angle_between)
    end

    ret[:gal_angle] = degrees_to_diff(truth[:gal_angle], predicted[:gal_angle])

    ret
end

function match_catalogs(rcf::RunCamcolField,
               results, truthfile, stagedir)
    # convert Celeste results to a DataFrame.
    celeste_full_df = celeste_to_df(results)
    println("celeste: $(size(celeste_full_df, 1)) objects")

    # load coadd catalog
    coadd_full_df = load_s82(truthfile)
    println("coadd catalog: $(size(coadd_full_df, 1)) objects")

    # find matches in coadd catalog by position
    disttol = 1.0 / 0.396  # 1 arcsec
    good_coadd_indexes = Int[]
    good_celeste_indexes = Int[]
    for i in 1:size(celeste_full_df, 1)
        try
            j = match_position(coadd_full_df[:ra], coadd_full_df[:dec],
                         celeste_full_df[i, :ra], celeste_full_df[i, :dec],
                         disttol)
            push!(good_celeste_indexes, i)
            push!(good_coadd_indexes, j)
        catch y
            isa(y, MatchException) || throw(y)
        end
    end

    celeste_df = celeste_full_df[good_celeste_indexes, :]
    coadd_df = coadd_full_df[good_coadd_indexes, :]

    # load "primary" catalog (the SDSS photoObj catalog used to initialize
    # celeste).
    primary_full_df = load_primary(rcf, stagedir)
    println("primary catalog: $(size(primary_full_df, 1)) objects")

    # match Primary to Celeste by object id
    pc_matches = Int[findfirst(primary_full_df[:objid], objid)
                   for objid in celeste_df[:objid]]
    pc_matches = filter(x->x!=0, pc_matches)
    primary_df = primary_full_df[pc_matches, :]

    # match Celeste to Primary by object id
    cp_matches = Int[findfirst(celeste_df[:objid], objid)
                   for objid in primary_df[:objid]]
    cp_matches = filter(x->x!=0, cp_matches)
    celeste_df = celeste_df[cp_matches, :]
    coadd_df = coadd_df[cp_matches, :]

    # show that all catalogs have same size, and (hopefully)
    # that not too many sources were filtered
    println("matched celeste catalog: $(size(celeste_df, 1)) objects")
    println("matched coadd catalog: $(size(coadd_df, 1)) objects")
    println("matched primary catalog: $(size(primary_df, 1)) objects")

    # ensure that all objects are matched
    if size(primary_df, 1) != size(celeste_df, 1)
        error("catalog mismatch between celeste and primary")
    end

    (celeste_df, primary_df, coadd_df)
end

function get_scores_df(celeste_err, primary_err, coadd_df)
    ttypes = [Symbol, Float64, Float64, Float64, Float64, Int]
    scores_df = DataFrame(ttypes, size(celeste_err, 2) - 1)
    names!(scores_df, [:field, :primary, :celeste, :diff, :diff_sd, :N])

    for i in 1:(size(celeste_err, 2) - 1)
        nm = names(celeste_err)[i + 1]
        nm != :objid || continue

        pe_good = Bool[!isnan(x) for x in primary_err[:, nm]]
        ce_good = Bool[!isnan(x) for x in celeste_err[:, nm]]
        good_row = pe_good & ce_good
        if string(nm)[1:5] == "star_"
            good_row &= (coadd_df[:is_star] .> 0.5)
        elseif string(nm)[1:4] == "gal_"
            good_row &= (coadd_df[:is_star] .< 0.5)
            if nm in [:gal_ab, :gal_scale, :gal_angle, :gal_fracdev]
                good_row &= !(0.05 .< coadd_df[:gal_fracdev] .< 0.95)
            end
            if nm == :gal_angle
                good_row &= coadd_df[:gal_ab] .< .6
            end
        end

        scores_df[i, :field] = nm
        N_good = sum(good_row)
        scores_df[i, :N] = N_good
        N_good > 0 || continue

        scores_df[i, :primary] = mean(primary_err[good_row, nm])
        scores_df[i, :celeste] = mean(celeste_err[good_row, nm])

        diffs = primary_err[good_row, nm] .- celeste_err[good_row, nm]
        scores_df[i, :diff] = mean(diffs)

        # compute the difference in error rates between celeste and primary
        # if we have enough data to get confidence intervals
        N_good > 1 || continue
        abs_errs = Float64[abs(x) for x in diffs]
        scores_df[i, :diff_sd] = std(abs_errs) / sqrt(N_good)
    end

    scores_df
end

function score_field(rcf::RunCamcolField, results, truthfile, stagedir)
    (celeste_df, primary_df, coadd_df) = match_catalogs(rcf,
                                results, truthfile, stagedir)

    suffix = @sprintf "%06d-%d-%04d.csv" rcf.run rcf.camcol rcf.field
    #writetable("celeste_results_"suffix, celeste_df)
    #writetable("primary_results_"suffix, primary_df)
    #writetable("coadd_results_"suffix, coadd_df)

    # difference between celeste and coadd
    celeste_err = get_err_df(coadd_df, celeste_df)
    primary_err = get_err_df(coadd_df, primary_df)

#=
    JLD.save("results_and_errors.jld",
             "celeste_df", celeste_df,
             "primary_df", primary_df,
             "coadd_df", coadd_df,
             "celeste_err", celeste_err,
             "primary_err", primary_err)
=#

    # create scores
    get_scores_df(celeste_err, primary_err, coadd_df)
end

"""
Score the celeste results for a particular field
"""
function score_field_disk(rcf::RunCamcolField, resultdir, truthfile, stagedir)
    fname = @sprintf "%s/celeste-%06d-%d-%04d.jld" resultdir rcf.run rcf.camcol rcf.field
    results = JLD.load(fname, "results")

    println( score_field(rcf, results, truthfile, stagedir) )
end

"""
Display results from Celeste, Primary, and Coadd for a particular object
"""
function score_object_disk(rcf::RunCamcolField, objid, resultdir, truthfile, stagedir)
    fname = @sprintf "%s/celeste-objid-%s.jld" resultdir objid
    results = JLD.load(fname, "results")

    (celeste_df, primary_df, coadd_df) = match_catalogs(rcf,
                                results, truthfile, stagedir)

    println("\n\nceleste results:\n")
    println(celeste_df)
    println("\n\nprimary results:\n")
    println(primary_df)
    println("\n\ncoadd results:\n")
    println(coadd_df)
end

end


text:
"""
    tabf_Klnnp

Table of the true anomalies, f, used for the discretisation of the integrals in the geometric factor Klnnp.
Set as `SVector{nbK,Float64}([(pi/nbK)*(k-0.5) for k=1:nbK])` by default.
Using `StaticArrays` for better performances.
"""
const tabf_Klnnp = SVector{nbK,Float64}([(pi/nbK)*(k-0.5) for k=1:nbK])

"""
    tabcosf_Klnnp

Table of cos(f) used to compute the radius in the geometric factor Klnnp, where f are the true anomalies.
Set as `SVector{nbK,Float64}([cos(tabf_Klnnp[k]) for k=1:nbK])` by default.
Using `StaticArrays` for better performances.
"""
const tabcosf_Klnnp = SVector{nbK,Float64}([cos(tabf_Klnnp[k]) for k=1:nbK])

"""
    tabcosnf_Klnnp

Table of cos(nf) used to compute the integrand in Klnnp the geometric factor Klnnp, where f are the true anomalies.
Set as `SMatrix{lmax,nbK,Float64}([cos(n*tabf_Klnnp[k]) for n=1:lmax, k=1:nbK])` by default.
Using `StaticArrays` for better performances.
"""
const tabcosnf_Klnnp = SMatrix{lmax,nbK,Float64}([cos(n*tabf_Klnnp[k]) for n=1:lmax, k=1:nbK])

"""
    IntTable_K

`IntTable` structure used in the computation of the geometric factor Klnnp.
Use `Table_create!()` to initialize.
"""
const IntTable_K_serial = IntTable_create!()

##################################################
# Computing the interaction coefficients Klnnp
##################################################
"""
    Klnnp(l,n,np,a,jt,ap,jp,IntTable_K=IntTable_K_serial)

Geometric factor K_nn'(a,j,a',j') involved in the computation of the scalar resonant diffusion coefficient DRR_jj.

# Remarks
    In practise this function does not exactly satisfy the symmetry Klnnp(a,j,ap,jp)=Klnnp(ap,jp,a,j).

# Arguments
- `l   ::Int64`  : mode of the multipole expansion.
- `n   ::Int64`  : first resonance number of the orbit.
- `np  ::Int64`  : second resonance number of the orbit.
- `a   ::Float64`: semi-major axis of the first orbit.
- `jt  ::Float64`: reduced angular momentum of the first orbit.
- `ap  ::Float64`: semi-major axis of the second orbit.
- `jp  ::Float64`: reduced angular momentum of the second orbit.
- `IntTable_K=IntTable_K_serial:: IntTable`: Structure which contains the tables needed for the computation of this integral.
"""
function Klnnp(l::Int64,n::Int64,np::Int64,a::Float64,j::Float64,ap::Float64,jp::Float64,IntTable_K::IntTable=IntTable_K_serial)

    tabw  = IntTable_K.tabw
    tabg  = IntTable_K.tabg
    tabgp = IntTable_K.tabgp
    tabR  = IntTable_K.tabR
    tabRp = IntTable_K.tabRp
    tabdeltaP = IntTable_K.tabdeltaP
    tabdeltaQ = IntTable_K.tabdeltaQ

    #####
    # Computing tabdeltaP and tabdeltaQ
    for j=1:nbK
        #####
        # Computing tabdeltaP
        x = 0.0
        for i=(tabw[j]+1):(tabw[j+1]) # ATTENTION, indices are shifted by one
            x += tabg[i]*(tabR[i]/tabRp[j])^(l)/(tabRp[j]) # Adding the contribution
        end
        tabdeltaP[j] = x # Updating tabdeltaP
        #####
        # Computing tabdeltaQ
        x = 0.0
        for i=(tabw[j+1]+1):(tabw[j+2]) # ATTENTION, indices are shifted by one
            x += tabg[i]*(tabRp[j]/tabR[i])^(l)/(tabR[i]) # Adding the contribution
        end
        tabdeltaQ[j] = x # Updating tabdeltaQ
    end
    #####
    # Computing the contributions from the sum P
    x = tabdeltaP[1] # Initial value of tabP[1]
    sumP = tabgp[1]*x # Initialising the contribution from P with the value of tabP[1]
    for j=1:(nbK-1) # Loop over the samples
        x = (tabRp[j]/tabRp[j+1])^(l+1)*x + tabdeltaP[j+1] # Computing the value of P[j+1]
        sumP += tabgp[j+1]*x # Updating the total contribution from P with P[j+1]
    end
    #####
    # Computing the contributions from the sum Q
    x = tabdeltaQ[nbK] # Initial value of tabQ[1]
    sumQ = tabgp[nbK]*x # Initialising the contribution from Q with the value of tabQ[nbK]
    for j=nbK:-1:2 # Loop over the samples. ATTENTION, this is a decreasing recurrence
        x = (tabRp[j-1]/tabRp[j])^(l)*x + tabdeltaQ[j-1] # Computing the value of Q[j-1]
        sumQ += tabgp[j-1]*x # Updating the total contribution from Q with Q[j-1]
    end
    #####
    return (sumP+sumQ)/(nbK^(2)) # Computing the total contribution from both P and Q
end

##################################################
# Computing the wrapped interaction coefficients
##################################################
"""
    AnnpSQ(l,n,np,a,jt,ap,jp,IntTable_K=IntTable_K_serial)

Coefficient |A_nnp|^2 at (a,jt,a',j') involved in the computation of the scalar resonant diffusion coefficient DRR_jj.

# Remark
    Do not forget that this function returns the SQUARE of the coefficient!

# Arguments
- `l   ::Int64`  : mode of the multipole expansion.
- `n   ::Int64`  : first resonance number of the orbit.
- `np  ::Int64`  : second resonance number of the orbit.
- `a   ::Float64`: semi-major axis of the first orbit.
- `jt  ::Float64`: reduced angular momentum of the first orbit.
- `ap  ::Float64`: semi-major axis of the second orbit.
- `jp  ::Float64`: reduced angular momentum of the second orbit.
- `IntTable_K=IntTable_K_serial:: IntTable`: Structure which contains the tables needed for the computation of this integral.
"""
function AnnpSQ(n::Int64,np::Int64,a::Float64,j::Float64,ap::Float64,jp::Float64,IntTable_K::IntTable=IntTable_K_serial)
    # Getting rid of the cases where the coefficients is necessarily zero
    if (abs(n) > lmax || abs(np) > lmax || mod(n-np,2) != 0)
        return 0.0
    else # In that case the result is non-zero
        res = 0.0 # Initialising the result
        #####
        IntTable_init!(IntTable_K,n,np,a,j,ap,jp) # Updating tabR, tabRp, tabg, tabgp, and tabw for the respective orders of the particles
        #####
        for l = max(abs(n),abs(np)):2:lmax # Loop over the l that have non-vanishing contributions !! ATTENTION, the step of the loop is 2
            res += (ylnSQ(l,n)*ylnSQ(l,np))/((2.0*l+1.0)^(3))*(Klnnp(l,n,np,a,j,ap,jp,IntTable_K))^(2) # Adding the contribution from (l,n,np) !! ATTENTION, not to forget the squared
        end
        #####
        res *= 16.0*pi^(2) # Putting the prefactor
        return res # Output
    end
end


text:
using Binning, SafeTestsets, Test

# in the root of the directory run:
# pkg> test BinningLayer
# to run all of the tests.

@time begin

@time @safetestset "layer tests" begin include("layer_tests.jl") end

end


text:
#=

  Assignment 1:
  https://edu.swi-prolog.org/mod/assign/view.php?id=242
  """
  http://cplint.eu/p/urns_and_balls.swinb

  Urns and balls

  Suppose you have two urns: urn1 contains 40 blue balls and 20 red balls and urn2 contains 25
  blue balls and 30 red balls.

  You throw an unbiased coin and, if it turns out head, you draw a ball from the first urn,
  it it turns out tails you draw a ball from the second urn.

  Write a program modeling this process and a query for answering the question
  "What is the probability of drawing a blue ball?"

  Write the program and the query in the cells below. Then save the notebook and submit the url.

  """
  The exact answer of drawing a blue call is
     0.5*40/60 + 0.5*25/55 = 0.56060606060606060606

  Cf ~/cplint/course_urns_and_balls.pl
     ~/blog/urns_and_balls.blog
     ~/psi/urns_and_balls.blog
     ~/webppl/urns_and_balls.wppl

=#

using Turing, StatsPlots, DataFrames
using ReverseDiff, Zygote, Tracker
# Turing.setadbackend(:reversediff)
# Turing.setadbackend(:zygote)
# Turing.setadbackend(:tracker)
include("jl_utils.jl")

#=
Summary Statistics
  parameters      mean       std   naive_se      mcse          ess      rhat
      Symbol   Float64   Float64    Float64   Float64      Float64   Float64

        coin    1.5008    0.5000     0.0025    0.0025   39377.7877    1.0000
        draw    1.4408    0.4965     0.0025    0.0025   39019.4117    1.0000

Quantiles
  parameters      2.5%     25.0%     50.0%     75.0%     97.5%
      Symbol   Float64   Float64   Float64   Float64   Float64

        coin    1.0000    1.0000    2.0000    2.0000    2.0000
        draw    1.0000    1.0000    1.0000    2.0000    2.0000

length(data): 40001
Dict{Any,Any} with 4 entries:
  ["tail", "blue"] => 8991
  ["head", "red"]  => 6654
  ["tail", "red"]  => 10979
  ["head", "blue"] => 13377
p(tail):0.49923751906202346
p(blue):0.5591860203494913
  2.282226 seconds (8.63 M allocations: 532.511 MiB, 7.70% gc time)

=#

data = []
@model function urns_and_balls()
    tail = 1
    head = 2
    coin ~ Categorical(simplex([0.5,0.5]))

    blue = 1
    red = 2
    draw ~ coin == head ? Categorical(simplex([40,20])) : Categorical(simplex([25,30]))
    push!(data,[
                coin == tail ? "tail" : "head"
                draw == blue ? "blue" : "red"
    ])
end

model = urns_and_balls()

num_chains = 4
# chains = sample(model, Prior(), MCMCThreads(), 1000, num_chains)

# chains = sample(model, MH(), 40_000)
# chains = sample(model, MH(), MCMCThreads(), 40_000, num_chains)

# chains = sample(model, PG(20), MCMCThreads(), 10_000, num_chains)
# chains = sample(model, PG(20), 40_000)
# chains = sample(model, IS(), MCMCThreads(), 10_000, num_chains)
chains = sample(model, IS(), 40_000)

# chains = sample(model, NUTS(1000,0.65), MCMCThreads(), 40_000, num_chains)
# chains = sample(model, Gibbs(MH(:zlabels),NUTS(1000,0.65,:m,:b,:sigma)), MCMCThreads(), 40_000, num_chains)

display(chains)
println("length(data): ", length(data))

h = make_hash(data)
display(h)
# println("data:", make_hash(data))
# println("p(blue): ", (df[!:,2] .== 2)/length(data))

println("p(tail):", sum([d[1] == "tail" for d in data])/length(data))
println("p(blue):", sum([d[2] == "blue" for d in data])/length(data))

# df = DataFrame(chains)

max_stars_repo_name: Wikunia/hakank

max_issues_repo_name: Wikunia/hakank

max_forks_repo_name: Wikunia/hakank


text:
push!(LOAD_PATH,"../src/")

using Documenter, FeatureEng

makedocs(
    modules = [FeatureEng],
    sitename="FeatureEng.jl",
    authors="Austin Poor",
    clean=true,
    # debug=true,
    pages=[
        "index.md",
        "Manual" => "manual.md",
        "Examples" => "examples.md",
    ],
    repo   = "github.com/a-poor/FeatureEng.jl.git",
)

deploydocs(
    # root   = dirname(abspath(PROGRAM_FILE)),
    # target = "build",
    repo   = "github.com/a-poor/FeatureEng.jl.git",
    # branch = "gh-pages",
    # deps   = nothing | <Function>,
    # make   = nothing | <Function>,
    devbranch = "main",
    # devurl = "",
    # versions = ["stable" => "v^", "v#.#", devurl => devurl],
    versions = ["stable" => "v^", "v#.#", "latest" => "v^"],
    # push_preview    = true,
    # repo_previews   = repo,
    # branch_previews = branch,
)


text:
origin(prim::HyperRectangle) = prim.origin
maximum(prim::HyperRectangle) = origin(prim) + widths(prim)
minimum(prim::HyperRectangle) = origin(prim)
length{T, N}(prim::HyperRectangle{N, T}) = N
widths(prim::HyperRectangle) = prim.widths

"""
Splits an HyperRectangle into two along an axis at a given location.
"""
split{H<:HyperRectangle}(b::H, axis, value::Integer) = _split(b,axis,value)
split{H<:HyperRectangle}(b::H, axis, value::Number) = _split(b,axis,value)
function _split{H<:HyperRectangle}(b::H, axis, value)
    bmin = minimum(b)
    bmax = maximum(b)
    b1max = setindex(bmax, value, axis)
    b2min = setindex(bmin, value, axis)

    return H(bmin, b1max-bmin),
           H(b2min, bmax-b2min)
end

# empty constructor such that update will always include the first point
@compat function (HR::Type{HyperRectangle{N,T}}){T,N}()
    HR(Vec{N,T}(typemax(T)), Vec{N,T}(typemin(T)))
end

# conversion from other HyperRectangles
@compat function (HR::Type{HyperRectangle{N,T1}}){N,T1,T2}(a::HyperRectangle{N,T2})
    HR(Vec{N, T1}(minimum(a)), Vec{N, T1}(widths(a)))
end

function HyperRectangle{N,T1,T2}(v1::Vec{N,T1}, v2::Vec{N,T2})
    T = promote_type(T1,T2)
    HyperRectangle{N,T}(Vec{N,T}(v1), Vec{N,T}(v2))
end

@compat function (HR::Type{HyperRectangle{N,T}}){N,T}(a::GeometryPrimitive)
    HR(Vec{N, T}(minimum(a)), Vec{N, T}(widths(a)))
end
"""
```
HyperRectangle(vals::Number...)
```
HyperRectangle constructor for indidually specified intervals.
e.g. HyperRectangle(0,0,1,2) has origin == Vec(0,0) and
width == Vec(1,2)
"""
@generated function HyperRectangle(vals::Number...)
    # Generated so we get goodish codegen on each signature
    n = length(vals)
    @assert iseven(n)
    mid = div(n,2)
    v1 = Expr(:call, :Vec)
    v2 = Expr(:call, :Vec)
    # TODO this can be inbounds
    append!(v1.args, [:(vals[$i]) for i = 1:mid])
    append!(v2.args, [:(vals[$i]) for i = mid+1:length(vals)])
    Expr(:call, :HyperRectangle, v1, v2)
end

function HyperRectangle{T}(r::SimpleRectangle{T})
    HyperRectangle{2,T}(r)
end

@compat function (::Type{HyperRectangle{N,T}}){N,T}(r::SimpleRectangle)
    if N > 2
        return HyperRectangle(Vec{N, T}(T(r.x), T(r.y), Vec{N-2,T}(zero(T))...),
                              Vec{N, T}(T(r.w), T(r.h), Vec{N-2,T}(zero(T))...))
    else
        return HyperRectangle(Vec{N, T}(T(r.x), T(r.y)),
                              Vec{N, T}(T(r.w), T(r.h)))
    end
end

"""
Transform a `HyperRectangle` using a matrix. Maintains axis-align properties
so a significantly larger HyperRectangle may be generated.
"""
function *{N1,N2,T1,T2}(m::Mat{N1,N1,T1}, h::HyperRectangle{N2,T2})

    # TypeVar constants
    T = promote_type(T1,T2)
    D = N1 - N2

    # get all points on the HyperRectangle
    d   = decompose(Point, h)
    pts = (Vec{N1,T}[Vec{N1,T}(pt, Vec{D,T}(one(T))...) for pt in d]...)::NTuple{2^N2,Vec{N1,T}}

    # make sure our points are sized for the tranform
    vmin = Vec{N1, T}(typemax(T))
    vmax = Vec{N1, T}(typemin(T))

    # tranform all points, tracking min and max points
    for pt in pts
        pn = m*pt
        vmin = min(pn,vmin)
        vmax = max(pn,vmax)
    end
    HyperRectangle{N2,T}(Vec{N2,T}(vmin), Vec{N2,T}(vmax-vmin))
end

function *{N,T1,T2}(m::Mat{N,N,T1}, h::HyperRectangle{N,T2})
    # equal dimension case

    # TypeVar constants
    T = promote_type(T1,T2)

    # get all points on the HyperRectangle
    pts = decompose(Point, h)

    # make sure our points are sized for the tranform
    vmin = Vec{N, T}(typemax(T))
    vmax = Vec{N, T}(typemin(T))

    # tranform all points, tracking min and max points
    for pt in pts
        pn = m*Vec(pt)
        vmin = min(pn,vmin)
        vmax = max(pn,vmax)
    end
    HyperRectangle{N,T}(vmin, vmax-vmin)
end

function HyperRectangle{N,T}(geometry::Array{Point{N, T}})
    HyperRectangle{N,T}(geometry)
end

"""
Construct a HyperRectangle enclosing all points.
"""
@compat function (t::Type{HyperRectangle{N1, T1}}){N1, T1, PT<:Point}(
        geometry::AbstractArray{PT}
    )
    N2, T2 = length(PT), eltype(PT)
    @assert N1 >= N2
    vmin = Point{N2, T2}(typemax(T2))
    vmax = Point{N2, T2}(typemin(T2))
    for p in geometry
         vmin = min(p, vmin)
         vmax = max(p, vmax)
    end
    o = vmin
    w = vmax - vmin
    if N1 > N2
        z = Vec{N1-N2,T1}(zero(T1))
        return HyperRectangle{N1,T1}(Vec{N1,T1}(o, z...),
                                     Vec{N1,T1}(w, z...))
    else
        return HyperRectangle{N1,T1}(Vec{N1,T1}(o),
                                     Vec{N1,T1}(w))
   end
end

xwidth(a::SimpleRectangle)  = a.w + a.x
width(a::SimpleRectangle)  = a.w
yheight(a::SimpleRectangle) = a.h + a.y
height(a::SimpleRectangle) = a.h
area(a::SimpleRectangle) = a.w*a.h
widths{T}(a::SimpleRectangle{T}) = Point{2,T}(a.w, a.h)
maximum{T}(a::SimpleRectangle{T}) = Point{2, T}(a.x + widths(a)[1], a.y +widths(a)[2])
minimum{T}(a::SimpleRectangle{T}) = Point{2, T}(a.x, a.y)
origin{T}(a::SimpleRectangle{T}) = Point{2, T}(a.x, a.y)

@compat (::Type{SimpleRectangle}){T}(val::Vec{2, T}) = SimpleRectangle{T}(0, 0, val...)
function SimpleRectangle{T}(position::Vec{2,T}, width::Vec{2,T})
    SimpleRectangle{T}(position..., width...)
end

function getindex{T}(a::Array{T,2}, rect::SimpleRectangle)
    a[rect.x+1:rect.w, rect.y+1:rect.h]
end

function setindex!{T}(a::Array{T,2}, b::Array{T,2}, rect::SimpleRectangle)
    a[rect.x+1:rect.w, rect.y+1:rect.h] = b
end


text:
"""
    SOCRBridge{T, F, G}

We simply do the inverse transformation of [`RSOCBridge`](@ref). In fact, as the
transformation is an involution, we do the same transformation.
"""
struct SOCRBridge{T, F, G} <: AbstractBridge
    rsoc::CI{F, MOI.RotatedSecondOrderCone}
end
function bridge_constraint(::Type{SOCRBridge{T, F, G}}, model,
                           f::MOI.AbstractVectorFunction,
                           s::MOI.SecondOrderCone) where {T, F, G}
    soc = MOI.add_constraint(model, rotate_function(f, T),
                             MOI.RotatedSecondOrderCone(MOI.dimension(s)))
    return SOCRBridge{T, F, G}(soc)
end

function MOI.supports_constraint(::Type{SOCRBridge{T}},
                                ::Type{<:MOI.AbstractVectorFunction},
                                ::Type{MOI.SecondOrderCone}) where T
    return true
end
MOIB.added_constrained_variable_types(::Type{<:SOCRBridge}) = Tuple{DataType}[]
function MOIB.added_constraint_types(::Type{<:SOCRBridge{T, F}}) where {T, F}
    return [(F, MOI.RotatedSecondOrderCone)]
end
function concrete_bridge_type(::Type{<:SOCRBridge{T}},
                              G::Type{<:MOI.AbstractVectorFunction},
                              ::Type{MOI.SecondOrderCone}) where T
    S = MOIU.promote_operation(/, T, MOIU.scalar_type(G), T)
    Y = MOIU.promote_operation(-, T, S, S)
    Z = MOIU.promote_operation(+, T, S, S)
    F = MOIU.promote_operation(vcat, T, Z, Y, G)
    return SOCRBridge{T, F, G}
end

# Attributes, Bridge acting as an model
function MOI.get(b::SOCRBridge{T, F},
                 ::MOI.NumberOfConstraints{F, MOI.RotatedSecondOrderCone}) where {T, F}
    return 1
end
function MOI.get(b::SOCRBridge{T, F},
                 ::MOI.ListOfConstraintIndices{F, MOI.RotatedSecondOrderCone}) where {T, F}
    return [b.rsoc]
end

# References
function MOI.delete(model::MOI.ModelLike, bridge::SOCRBridge)
    MOI.delete(model, bridge.rsoc)
end

# Attributes, Bridge acting as a constraint
function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintFunction,
                 bridge::SOCRBridge{T, F, G}) where {T, F, G}
    # As it is an involution, we can just reapply the same transformation
    func = MOI.get(model, attr, bridge.rsoc)
    return MOIU.convert_approx(G, rotate_function(func, T))
end
function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintSet, bridge::SOCRBridge)
    set = MOI.get(model, attr, bridge.rsoc)
    return MOI.SecondOrderCone(MOI.dimension(set))
end
# As the linear transformation is a symmetric involution,
# the constraint primal and dual both need to be processed by reapplying the same transformation
function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintPrimal,
                 bridge::SOCRBridge)
    return rotate_result(model, attr, bridge.rsoc)
end
function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintDual,
                 bridge::SOCRBridge)
    return rotate_result(model, attr, bridge.rsoc)
end


text:
@testset "single-column kets" begin
    h = LatticePresets.honeycomb() |> hamiltonian(hopping(I), orbitals = (Val(1), Val(3)))
    k = ket(h)
        @test iszero(k)
        @test eltype(k) <: SVector{3}
        @test size(k) == (size(h, 1), 1)
    @test_throws ArgumentError ket(ketmodel(1), h)
    @test_throws ArgumentError ket(ketmodel(r -> SA[1,2,3]), h)
    @test_throws ArgumentError ket(ketmodel(2, sublats = :B), h)
    @test_throws ArgumentError ket(ketmodel(SA[1,2,3], maporbitals = true), h)
    k = ket(ketmodel(SA[1], maporbitals = true), h)
        @test parent(k) == SVector.(hcat([(0.5,0,0); (0.5,0.5,0.5)]))
    k = ket(ketmodel(fill(1, 1, 1), maporbitals = true), h)
        @test parent(k) == SVector.(hcat([(0.5,0,0); (0.5,0.5,0.5)]))
    k = ket(ketmodel(r -> [1], maporbitals = true), h)
        @test parent(k) == SVector.(hcat([(0.5,0,0); (0.5,0.5,0.5)]))
    k = ket(ketmodel(r -> SA[3], maporbitals = true, sublats = :A), h)
        @test parent(k) == SVector.(hcat([(1,0,0); (0,0,0)]))
    k = ket(ketmodel(2, sublats = :A), h)
        @test parent(k) == SVector.(hcat([(1,0,0); (0,0,0)]))
    k = ket(ketmodel(SA[1, 2, 2], sublats = :B), h)
        @test parent(k) == SVector.(hcat([(0,0,0); (1/3,2/3,2/3)]))
    k = ket(ketmodel(r -> SMatrix{3,1}((1, 2, 2)), sublats = :B), h)
        @test parent(k) == SVector.(hcat([(0,0,0); (1/3,2/3,2/3)]))
    k = ket(ketmodel(r -> 2r[2], sublats = :A), h)
        @test parent(k) == SVector.(hcat([(-1,0,0); (0,0,0)]))
    k = ket(ketmodel(2, sublats = :A, normalization = missing), h)
        @test parent(k) == SVector.(hcat([(2,0,0); (0,0,0)]))
    k = ket(ketmodel(1; maporbitals = true), h)
        @test parent(k) == SVector.(hcat([ (0.5,0,0); (0.5,0.5,0.5)]))
    k = ket(ketmodel(1; maporbitals = true), h)
        @test parent(k) == SVector.(hcat([(0.5,0,0); (0.5,0.5,0.5)]))
end

@testset "multi-column kets" begin
    h = LatticePresets.honeycomb() |> hamiltonian(hopping(I), orbitals = (Val(1), Val(3)))
    k = ket(ketmodel(SA[1 1]; maporbitals = true), h)
        @test parent(k) == SVector.([(0.5,0,0) (0.5,0,0); (0.5,0.5,0.5) (0.5,0.5,0.5)])
    k´ = ket(ketmodel([1 1]; maporbitals = true), h)
        @test k == k´
    @test_throws ArgumentError ket(ketmodel(SA[1 0; 0 1]; maporbitals = true), h)
    @test_throws ArgumentError ket(ketmodel([1 0; 0 2; 0 1]; sublats = :A), h)
    @test_throws ArgumentError ket(ketmodel(SA[1 0; 0 1]), h)
    @test_throws ArgumentError ket(ketmodel(SA[1 0 0; 0 2 0; 0 0 3]), h)
    k = ket(ketmodel(SA[1 0; 0 2; 0 1], sublats = :B, normalization = missing), h)
        @test parent(k) == SVector.([(0,0,0) (0,0,0); (1,0,0) (0,2,1)])
    k´ = ket(ketmodel([1 0 0; 0 2 1]', sublats = :B, normalization = missing), h)
    k´´= ket(ketmodel(r -> SA[1 0 0; 0 2 1]', sublats = :B, normalization = missing), h)
        @test k == k´´
    @test_throws Exception ket(ketmodel(2I), h)
    k = ket(ketmodel([1 0; 0 2; 0 1], sublats = :B, normalization = missing), h)
        @test parent(k) == SVector.([(0,0,0) (0,0,0); (1,0,0) (0,2,1)])
end

@testset "randomkets" begin
    h = LatticePresets.honeycomb() |> hamiltonian(hopping(I), orbitals = Val(2))
    ks = ket.(randomkets(3, r-> 2*SA[1 0; 0 1], normalization = missing), Ref(h))
        @test length(ks) == 3
        @test ks isa Vector{<:Quantica.Ket}
        @test ks[1] == ks[2] == ks[3]
    k = first(ks)
        @test parent(k) == SVector.([(2,0) (0,2); (2,0) (0,2)])
    k1, k2 = ket.(randomkets(2, r-> randn()*SA[1 0; 0 1], normalization = missing), Ref(h))
        r, s = k1[1,1][1], k1[2,1][1]
        @test k1 ≈ SVector.([(r,0) (0,r); (s,0) (0,s)])
        r, s = k2[1,1][1], k2[2,1][1]
        @test k2 ≈ SVector.([(r,0) (0,r); (s,0) (0,s)])
end

@testset "flatten/unflatten kets" begin
    h = LatticePresets.honeycomb() |> hamiltonian(hopping(I), orbitals = (Val(1), Val(3)))
    k = ket(ketmodel(SA[1 1]; maporbitals = true), h)
    @test unflatten(flatten(k), orbitalstructure(k)) == k
    kf = ket(ketmodel(SA[1 1]), flatten(h))
    Quantica.ket!(kf, ketmodel(SA[1 1]; maporbitals = true), h)
    @test unflatten(kf, orbitalstructure(k)) == k
end

@testset "ket algebra" begin
    h = LatticePresets.honeycomb() |> hamiltonian(hopping(I), orbitals = (Val(1), Val(3))) |> unitcell(2)
    km = 2*ketmodel(SA[1 1], sublats = :A) + ketmodel(SA[1 1; 0 0; 1 1], sublats = :B)
    k = ket(km, h)
    @test size(k) == (8,2)
    @test all(k[1:4, 1:2] .== Ref(k[1,1]))
    @test all(k[5:8, 1:2] .== Ref(k[5,1]))
    km = ketmodel(SA[1 1], sublats = :A) - 2*ketmodel(SA[1; 0; 1], sublats = :B)
    @test_throws ArgumentError ket(km, h)
    km = ketmodel(1, sublats = :A) - 2*ketmodel(SA[1 2; 0 1; 1 1], sublats = :B, singlesitekets = true)
    k = ket(km, h)
    @test size(k) == (8, 12)
    km = ketmodel(I, sublats = :A) + 2*ketmodel(I, sublats = :B, singlesitekets = true)
    k = ket(km, h)
    @test size(k) == (8, 16)
    @test parent(flatten(k)) == I
    km = basiskets()
    k´ = ket(km, h)
    @test k == k´
    km = ketmodel(I, sublats = :A) + ketmodel(SA[1 0 0; 0 1 0; 0 0 -1], sublats = :B, singlesitekets = true)
    @test km.singlesitekets == true
    k = ket(km, h)
    @test size(k) == (8, 16)
    @test parent(flatten(k)) != I && abs.(parent(flatten(k))) == I
end

max_stars_repo_name: pablosanjose/Quantica

max_issues_repo_name: pablosanjose/Quantica

max_forks_repo_name: pablosanjose/Quantica


text:
#
# Developed for:
#
#  C O N S T R U C T I V I S T__ __
#   _____ ____   ____   /  |/  // /
#  / ___// __ \ / __ \ / /|_/ // /
# / /__ / /_/ // / / // /  / // /___
# \___/ \____//_/ /_//_/  /_//_____/
#  M A C H I N E   L E A R N I N G
#
#
# A Project by
# Thomas Schmid | UNIVERSITÄT LEIPZIG
# www.constructivist.ml
#
# Code Author: Florian Große
# Licence: MIT
#

"""
# Krippendorff.jl

Krippendorff's alpha in Julia

This package aims to provide easy access to the afforementioned inter-rater reliability measure.
The main entry point (and sole exported function) is [`krippendorffs_alpha`](@ref), which provides an
easy-to-use interface for most requirements. `alpha` is the unexported alias for `krippendorffs_alpha`,
allowing the usage of `Krippendorff.alpha`.

Alpha is used to quantify the amount of agreement or disagreement of a number of raters that assign
values (responses) to different units (subjects/tasks) according to a common code. It is conceptually
similar to other measures like Cohen's κ, Scott's π, Fleiss' κ or intra-class correlation, but is more
generally applicable. Alpha can be computer for any number of raters and units, allows for missing values,
is corrected for small sample sizes and works with various levels of measurement, which govern the
distance metric that is used in the computation. Thus, alpha works with nominal, ordinal, interval,
(bi)polar, circular and ratio variables. (and possibly more)

The high-level function will *just work* with most data. Inputs are required to either be iterable
julia objects or satisfy the [`Tables.jl`](https://github.com/JuliaData/Tables.jl) interface (as
determined by `Tables.istable`), in which case they will be accessed via `Tables.rows` or `Tables.columns`.
Missing values are handled and the algorithm tries to determine the set of possible responses
automatically. Depending on these, one of currently two computational backends is choosen to compute
alpha. Thin wrappers around these backends with proper documentation are [`compute_alpha_generical`](@ref Krippendorff.compute_alpha_generical)
and [`compute_alpha_with_coincidences`](@ref Krippendorff.compute_alpha_with_coincidences).

# References

- Krippendorff, Klaus. (2011). Computing Krippendorff's Alpha-Reliability. Retrieved from https://repository.upenn.edu/asc_papers/43
- Hayes, Andrew & Krippendorff, Klaus. (2007). Answering the Call for a Standard Reliability Measure for Coding Data. Communication Methods and Measures. 1. 77-89. 10.1080/19312450709336664.
"""
module Krippendorff

using Tables

const UNITSDEFAULT = :rows
const KNOWNMETRICS = [:nominal, :interval]

## common entry point

# warn that NamedTuples/Dicts of Vectors and Vectors of NamedTuples/Dicts
# can satisfy the tables interface. See Krippendorff.istable
"""
    Krippendorff.alpha(...)
    krippendorffs_alpha(input; units = $(UNITSDEFAULT), metric = :nominal, R = :discrete, silent = false)

Compute the Krippendorff's-α inter-rater reliability measure from the supplied input. The input will
be checked to determine how to iterate over it, `missing` values will be handled automatically if present
and the most efficient algorithm to compute α will be determined heuristically. By default, `Tables.jl`
tables and table-like inputs are assumed to have columns representing raters and rows representing units.
See [`prepare_iterator`](@ref Krippendorff.prepare_iterator) for more information about the input requirements.

# Arguments
- `units::Union{Symbol,AbstractString}`: `rows` or `col(umn)s`, see [`prepare_iterator`](@ref Krippendorff.prepare_iterator) for explanation
- `metric`: a metric computing the squared distance between any pair of responses. Any of $(KNOWNMETRICS) or a custom function. Should satisfy `f(x,y) = [0 if x==y], [>0 otherwise]` but this is not enforced. See README for explanation of the default metrics.
- `R`: The space of possible responses. Either `:discrete` (relatively few possible responses, uses fast computation via coincidence matrix), `:continuous` (many possible responses up to continuous range, uses a slower but generically applicable algorithm with minimal allocation) or a precomputed value (implies discrete, but avoids searching for all unique values). If precomputed, it should be supplied as a tuple or vector of possible values (e.g. the output of unique(yourdata)), or as an appropriate range object where possible (slightly more efficient).
- `silent::Bool`: set to disable all optional output (`@info` and `@warn`, doesn't affect error messages)

See also: [`compute_alpha_generical`](@ref Krippendorff.compute_alpha_generical), [`compute_alpha_with_coincidences`](@ref Krippendorff.compute_alpha_with_coincidences)
"""
function krippendorffs_alpha(input; units::Union{Symbol,AbstractString} = UNITSDEFAULT, metric = :nominal, R = :discrete, silent::Bool = false)
    units_iterator, elementtype, inputsize = _dispatch_units_iterator(input, Symbol(units); silent)
    implementation, possible_responses = _dispatch_responses(R, units_iterator, elementtype, inputsize)
    distance_metric_squared = _dispatch_metric(metric, elementtype; silent)
    return _compute_alpha(implementation, possible_responses, distance_metric_squared, units_iterator)
end

const alpha = krippendorffs_alpha # enables Krippendorff.alpha
export krippendorffs_alpha

## input type dispatch + find a meaningful elementtype + size of the whole thing

function _dispatch_units_iterator(input, units::Symbol; silent::Bool = false)
    units ∈ (:rows, :columns, :cols) || error("\"units\" must be one of :rows or :col(umn)s, got: $(units)")
    userows = units === :rows # use rows as units or not (use columns then)
    units, elementtype, sizeestimate = _units_eltype_and_size(input, userows; silent)
    # if missings are allowed, wrap every unit in a skipmissings
    prepared_iter, prepared_eltype = Missing<:elementtype ? ((skipmissing(u) for u in units), nonmissingtype(elementtype)) : (units, elementtype)
    return (prepared_iter, prepared_eltype, sizeestimate)
end

# use multiple dispatch to find an appropriate units_iterator, eltype and size for different types

function _units_eltype_and_size(input, userows::Bool; silent::Bool = false)
    # fallback, check Tables.jl interface via istable, else assume generic iterability
    if Tables.istable(input)
        if userows
            rows = Tables.rows(input)
            schema = Tables.schema(rows)
            if isnothing(schema)
                # infer from first row
                if !silent
                    @info "Schema-less table detected, infering layout from first row only. This may lead to unexpected results."
                end
                firstrow, _ = iterate(Iterators.take(rows,1))
                elementtype = reduce((x,y)->Union{x,y}, (typeof(entry) for entry in firstrow), init=Base.Bottom)
                colnames = Tables.columnnames(firstrow)
                sizeestimate = length(rows)*length(colnames)
                let colnames = colnames
                    return (( [Tables.getcolumn(unit, name) for name in colnames] for unit in rows ), elementtype, sizeestimate)
                end
            else
                elementtype = reduce((x,y)->Union{x,y}, schema.types, init=Base.Bottom)
                sizeestimate = length(rows)*length(schema.names)
                return (rows, elementtype, sizeestimate)
            end
        else # usecols
            cols = Tables.columns(input)
            schema = Tables.schema(cols)
            if isnothing(schema)
                # infer from first row
                if !silent
                    @info "Schema-less table detected, infering types from first column only. This may lead to unexpected results."
                end
                colnames = Tables.columnnames(columns)
                firstcol = Tables.getcolumn(columns, 1)
                elementtype = reduce((x,y)->Union{x,y}, (typeof(entry) for entry in firstcol), init=Base.Bottom)
                sizeestimate = length(firstcol)*length(colnames)
                let colnames = colnames
                    return ((Tables.getcolumn(columns, name) for name in colnames), elementtype, sizeestimate)
                end
            else
                elementtype = reduce((x,y)->Union{x,y}, schema.types, init=Base.Bottom)
                sizeestimate = length(Tables.getcolumn(cols,1))*length(schema.names)
                let colnames = schema.names
                    return ((Tables.getcolumn(cols, name) for name in colnames), elementtype, sizeestimate)
                end
            end
        end
    else
        # last fallback, wrap everything in values(...) to accomodate dicts
        units = (values(unit) for unit in values(input))
        sizeestimate = sum(length(unit) for unit in units)
        elementtype = reduce((x,y)->Union{x,y}, (eltype(unit) for unit in units), init=Base.Bottom)
        if !silent
            @info "Input has no notion of rows and columns, iterating it is expected to yield units."
        end
        return (units, elementtype, sizeestimate)
    end
end

function _units_eltype_and_size(input::AbstractMatrix, userows::Bool; silent::Bool = false)
    iter = userows ? eachrow(input) : eachcol(input)
    return (iter, eltype(input), *(length(axes(input,1)), length(axes(input,2))))
end

## dispatch possible responses

function _dispatch_responses(R, _,_,_)
    # if not a Symbol: assume R is precomputed and reasonable
    return (Coincidence(), R)
end

function _dispatch_responses(R::Union{Symbol,AbstractString}, itr, elementtype, inputsize)
    Rsym = Symbol(R)
    if Rsym === :continuous
        return (Generic(), nothing)
    elseif Rsym === :discrete
        # get unique values of input, if there are lots of unique values
        # in comparison to the input size, use generic version
        # also don't bother with coincidence etc if the input is tiny
        uniques = unique(Base.Iterators.flatten(itr))
        # TODO determine cutoffs
        if inputsize / length(uniques) > 0.5
            return (Generic(), nothing)
        else
            return (Coincidence(), optimize_possible_responses(uniques, elementtype))
        end
    else
        error("Unknown input for R: $(R). Use either :discrete, :continuous or a precomputed set of responses.")
    end
end

# try to find a more efficient mapping strategy than dict(unique_values => indices)
# i.e. if the input looks like a range and is sufficiently dense, then make it a range
function _optimize_possible_reponses(uniquevalues, elementtype)
    # if extended to a range, how many empty indices may appear at worst
    allowedsparsity::Float64 = 0.5
    if elementtype <: Integer # relax to number?
        minR, maxR = extrema(uniquevalues)
        # TODO check for range with a constant non-one stepsize
        if length(uniquevalues)/(maxR-minR) >= allowedsparsity # assumes stepsize == 1
            if minR === oneunit(elementtype) # one-based is especially easy
                return Base.OneTo(maxR)
            else
                return UnitRange(minR, maxR)
            end
        else # making a range would introduce too many unused responses
            return uniquevalues
        end
    else # all non-numerics can not be directly mapped to indices anyways
        return uniquevalues
    end
end

# return a function that maps all possible values to indices into a coincidence matrix axis
_map_responses_to_indices(::Nothing) = nothing
_map_responses_to_indices(::Base.OneTo) = identity
_map_responses_to_indices(possible_responses::UnitRange) = Base.Fix2(-, possible_responses.start - 1)
function _map_responses_to_indices(possible_responses)
    let mapping = Dict(reverse.(collect(pairs(possible_responses))))
        return x -> mapping[x]
    end
end

## dispatch inputs for metric (and assert that it fits the elementtype)

function _dispatch_metric(metric::Union{Symbol, AbstractString}, elementtype; silent::Bool = false)
    m = Symbol(metric)
    if m === :nominal
        return (!=)
    elseif m === :interval
        elementtype <: Number || (!silent && @warn "Interval metric is unlikely to work with non-numeric responses.")
        return (x,y)->abs2(x-y)
    else
        throw(ArgumentError("Metric $(m) is currently not known. Try to use one of: $(KNOWNMETRICS)"))
    end
end

function _dispatch_metric(metric, elt; _...) # takes everything that's not a symbol or string
    # assumes that it's a valid function defined on every pair of possible responses
    return metric
end

# from a distance metric and the set of possible responses mapping to the coincidence matrix indices,
# compute a symmetric distance matrix once for later evaluation

function _compute_distance_matrix(possible_responses, mapping_function::F, distance_metric_squared) where {F}
    l = length(possible_responses)
    D = zeros(l,l)
    for rᵢ in possible_responses, rⱼ in possible_responses
        i = mapping_function(rᵢ)
        j = mapping_function(rⱼ)
        D[i,j] = distance_metric_squared(rᵢ, rⱼ)
    end
    return D
end

## computation backends

abstract type KrippendorffComputationStrategy end
struct Coincidence<:KrippendorffComputationStrategy end
struct Generic<:KrippendorffComputationStrategy end

# TODO docstrings
# using coincidence matrix
function _compute_alpha(::Coincidence, possible_responses, distance_metric_squared, units_iterator)
    mapping_function = _map_responses_to_indices(possible_responses)
    distances = _compute_distance_matrix(possible_responses, mapping_function, distance_metric_squared)
    l = length(possible_responses)
    C = zeros(l, l) # initialize coincidence matrix
    _fill_coincidence_matrix!(C, mapping_function, units_iterator)
    return _evaluate_coincidence_matrix(C, distances)
end

function _fill_coincidence_matrix!(C, mapping_function::F, units_iterator) where {F} # should force specialization
    for unit in units_iterator
        increment::Float64 = 1 / (count(x->true,unit) - 1) # count(...) because skipmissings have no length
        for i in eachindex(unit), j in eachindex(unit)
            if i<j # this assumes ordered indices, won't work for e.g. NamedTuples but I think this is rarely a problem
                cᵢ = mapping_function(unit[i])
                cⱼ = mapping_function(unit[j])
                C[ cᵢ, cⱼ ] += increment
                C[ cⱼ, cᵢ ] += increment
            end
        end
    end
end

function _evaluate_coincidence_matrix(C, distance_matrix)
    l = size(C,1)
    marginals = sum(C, dims=1)
    n = sum(marginals)
    disagreement_observed = sum( C[i,j] * distance_matrix[i,j] for i in 1:l for j in i+1:l)
    disagreement_expected = sum( marginals[i] * marginals[j] * distance_matrix[i,j] for i in 1:l for j in i+1:l)
    # it would be sufficient to evaluate just the upper or lower triangle of each, but I don't think this would be faster
    return 1. - (n-1)*(disagreement_observed / disagreement_expected)
end

# generic without coincidence matrix
function _compute_alpha(::Generic, _, distance_metric, units_iterator)
    disagreement_observed = 0.
    disagreement_expected = 0.
    n_all = 0
    # one huge pass over the data: always add pairs to the expected term and
    # to the observed term only when within one unit
    # to avoid duplications, enumerate both units iterator and skip all redundant
    for (i_s, unit_s) in enumerate(units_iterator)
        n_inunit = count(x->true,unit_s) # count(...) because skipmissings have no length
        n_inunit > 1 || continue
        n_all += n_inunit
        for (i_t, unit_t) in enumerate(units_iterator)
            count(x->true,unit_t) > 1 || continue
            i_t < i_s && continue # skip backwards pairs of units
            if i_s == i_t
                # both iters in same unit, iter forward pairs only
                unitsum = sum( distance_metric(i,j)
                                for (i_pos, i) in enumerate(unit_s)
                                for (j_pos, j) in enumerate(unit_t)
                                if i_pos ≥ j_pos) # count same position pairs once for consistency, should be equal to i_pos > j_pos unless the metric violates d(x,x)==0
                disagreement_observed += unitsum / (n_inunit-1)
                disagreement_expected += unitsum
            else
                # iters in different units, iter both full, add only to expected disagreement term
                disagreement_expected += sum( distance_metric(i,j) for i in unit_s for j in unit_t )
            end
        end
    end
    return 1. - ((n_all-1)*(disagreement_observed / disagreement_expected))
end

## utility stuff

# thin wrappers for backends with explanation

"""
    compute_alpha_generical(units_iterator, squared_distance_metric)

Generic computation backend for Krippendorffs alpha, bypassing the creation of coincidence tables
entirely. The tradeof in this case is the necessity to iterate over all pairs of units. Thus, scaling
is typically much worse (in O(U²*R) for (U)nits and (R)aters) than when using coicidence matrices.
Nonetheless, this backend will be used by default if the number of possible responses is large
compared to the size of the input, since a coincidence matrix can get huge in this case.
It is also preferable when all possible responses are not known beforehand or span a continuous
spectrum of values.

# Arguments

- `units_iterator`: An iterable object which is assumed to yield units with no `missings`. If this is not the case, you can call [`prepare_iterator`](@ref Krippendorff.prepare_iterator) on it.
- `squared_distance_metric`: An object callable on any pair of possible responses in the supplied iterator. Should satisfy `f(x,y) = [0 if x==y], [>0 otherwise]` but this is not checked explicitly.
"""
compute_alpha_generical(units, metric) = _compute_alpha(Generic(), nothing, metric, units)

"""
    compute_alpha_with_coincidences(units, metric, possible_responses)

The default fast computation backend for Krippendorffs alpha. This will iterate over all units only once
and thus scales preferably if the number of differing possible responses is bounded. A coincidence matrix
is generated to keep track of the observed disagreement, while the expected disagreement will be computed
from the marginals of the coincidence matrix.
Be careful, if the number of possible responses is large, this backend may allocate a lot!

# Arguments

- `units_iterator`: An iterable object which is assumed to yield units with no `missings`. If this is not the case, you can call [`prepare_iterator`](@ref Krippendorff.prepare_iterator) on it.
- `squared_distance_metric`: An object callable on any pair of possible responses in the supplied iterator. Should satisfy `f(x,y) = [0 if x==y], [>0 otherwise]` but this is not checked explicitly.
- `R`: Space of possible responses. This is necessary to generate the coincidence matrix efficiently. It should be supplied as a tuple or vector of possible values (e.g. the output of unique(yourdata)), or as an appropriate range object where possible.
"""
compute_alpha_with_coincidences(units, metric, R) = _compute_alpha(Coincidence(), R, metric, units)

"""
    prepare_iterator(input; units = $(UNITSDEFAULT))

Prepare an object for iteration by one of the `compute_alpha_...` functions.
This involves determining how to iterate over units in the object and probing for the elementtype.
If appropriate, the `units` argument is used to determine the direction of iteration. This is however
not always possible. If no sense of direction is found, the heuristic will assume the input is already
a suitable iterator over units. Furthermore, if the input is found to contain `missing` values
(or has them in it's eltype), all units will be wrapped in `skipmissing` automatically.

Since some seemingly unstructured iterables can satisfy the `Tables.jl` interface somewhat surprisingly
(Dict{Symbol,Vector} does, but not Dict{String,Vector} for example) and this may change the order of
iteration implied, you can call the helper function [`Krippendorff.istable`](@ref) to see whether your
input looks like a `table` and if yes, how many rows and columns it appears to have.

# Arguments

- `input`: The input to be prepared. Should support generic iteration via `iterate`, `eachrow` or `eachcol` or satisfy the [`Tables.jl`](https://github.com/JuliaData/Tables.jl) interface as determined by `Tables.istable`.
- `units::Union{Symbol,AbstractString}`: either `rows` or `col(umn)s`. This is used to determine how to iterate units in the input. For example, if the input iterator was a `Matrix`, `:rows` would make the function call `eachrow(input)` (and a little bit more).
"""
function prepare_iterator(input, units::Union{Symbol,AbstractString} = UNITSDEFAULT)
    units_iter, _, _ = _dispatch_units_iterator(input, Symbol(units))
    return units_iter
end

"""
    istable(input; IO = stdout)

A thin wrapper around `Tables.istable` that additionally prints how man rows and columns the input appears
to have when iterated through the `Tables.jl` interface. (`Tables.columns` specifically)
`IO` can be used to redirect the written output. Pass `IO=devnull` to supress output (making it equivalent to calling `Tables.istable`

# Examples

The `Tables.jl` interface assumes named columns and unnamed rows, which may lead to confusion
if one wanted to pass a dictionary of rows for examples:

```jldoctest; setup = :(using Krippendorff, Tables)
julia> testmatrix = reshape(1:15, (3,5))
3×5 reshape(::UnitRange{Int64}, 3, 5) with eltype Int64:
 1  4  7  10  13
 2  5  8  11  14
 3  6  9  12  15

julia> Krippendorff.istable(testmatrix)
false

julia> Krippendorff.istable(Tables.table(testmatrix));
Input satisfies the Tables.jl table interface and appears to have 3 rows and 5 columns.

julia> testvectordict = Dict([k=>v for (k,v) in zip([:row1, :row2, :row3], eachrow(testmatrix))]); [println(entry) for entry in testvectordict];
:row1 => [1, 4, 7, 10, 13]
:row2 => [2, 5, 8, 11, 14]
:row3 => [3, 6, 9, 12, 15]

julia> Krippendorff.istable(testvectordict)
Input satisfies the Tables.jl table interface and appears to have 5 rows and 3 columns.
true
```
"""
function istable(input; IO = stdout)
    table = Tables.istable(input)
    if table
        itr = Tables.columns(input)
        cols = length(Tables.columnnames(itr))
        rows = cols==0 ? 0 : length(Tables.getcolumn(itr,1))
        println(IO, "Input satisfies the Tables.jl table interface and appears to have $(rows) rows and $(cols) columns.")
    end
    return table
end

end # module

max_stars_repo_name: FPGro/Krippendorff

max_issues_repo_name: FPGro/Krippendorff

max_forks_repo_name: FPGro/Krippendorff


text:
int_rules_3_1_1 = @theory begin

    #= ::Subsection::Closed:: =#
    #= 3.1.1*(a+b*log(c*x^n))^p =#
    @apply_utils Antiderivative(log(~(c') * (~x) ^ ~(n')), ~x) => ~x * log(~c * (~x) ^ ~n) - ~n * ~x <-- FreeQ([~c, ~n], ~x)
    @apply_utils Antiderivative((~(a') + ~(b') * log(~(c') * (~x) ^ ~(n'))) ^ ~(p'), ~x) => ~x * (~a + ~b * log(~c * (~x) ^ ~n)) ^ ~p - ~b * ~n * ~p * Antiderivative((~a + ~b * log(~c * (~x) ^ ~n)) ^ (~p - 1), ~x) <-- FreeQ([~a, ~b, ~c, ~n], ~x) && (GtQ(~p, 0) && IntegerQ(2 * ~p))
    @apply_utils Antiderivative((~(a') + ~(b') * log(~(c') * (~x) ^ ~(n'))) ^ ~p, ~x) => (~x * (~a + ~b * log(~c * (~x) ^ ~n)) ^ (~p + 1)) / (~b * ~n * (~p + 1)) - (1 / (~b * ~n * (~p + 1))) * Antiderivative((~a + ~b * log(~c * (~x) ^ ~n)) ^ (~p + 1), ~x) <-- FreeQ([~a, ~b, ~c, ~n], ~x) && (LtQ(~p, -1) && IntegerQ(2 * ~p))
    @apply_utils Antiderivative(1 / log(~(c') * ~x), ~x) => LogIntegral(~c * ~x) / ~c <-- FreeQ(~c, ~x)
    @apply_utils Antiderivative((~(a') + ~(b') * log(~(c') * (~x) ^ ~(n'))) ^ ~p, ~x) => (1 / (~n * (~c) ^ (1 / ~n))) * Subst(Antiderivative(ℯ ^ (~x / ~n) * (~a + ~b * ~x) ^ ~p, ~x), ~x, log(~c * (~x) ^ ~n)) <-- FreeQ([~a, ~b, ~c, ~p], ~x) && IntegerQ(1 / ~n)
    @apply_utils Antiderivative((~(a') + ~(b') * log(~(c') * (~x) ^ ~(n'))) ^ ~p, ~x) => (~x / (~n * (~c * (~x) ^ ~n) ^ (1 / ~n))) * Subst(Antiderivative(ℯ ^ (~x / ~n) * (~a + ~b * ~x) ^ ~p, ~x), ~x, log(~c * (~x) ^ ~n)) <-- FreeQ([~a, ~b, ~c, ~n, ~p], ~x)
end

max_stars_repo_path: src/rules/3 Logarithms/3.1.1 (a+b log(c x^n))^p.jl


text:
using LinearAlgebra

function scalarproduct(a::AbstractVector{T}, b::AbstractVector{T}, c::AbstractVector{T}) where {T<:Number}
    return dot(a, cross(b, c))
end

function vectorproduct(a::AbstractVector{T}, b::AbstractVector{T}, c::AbstractVector{T}) where {T<:Number}
    return cross(a, cross(b, c))
end

const a = [3, 4, 5]
const b = [4, 3, 5]
const c = [-5, -12, -13]

println("Test Vectors:")
@show a b c

println("\nVector Products:")
@show dot(a, b)
@show cross(a, b)
@show scalarproduct(a, b, c)
@show vectorproduct(a, b, c)

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
#=
ndgrid-unitr.jl
ndgrid for a set of "UnitRange{Int}" inputs.
=#

export ndgrid

"""
    GridUR{T,d,D} <: AbstractGrid{T,d,D}
The `d`th component of `D`-dimensional `ndgrid(a:b, c:d, ...)`
where `1 ≤ d ≤ D`.
"""
struct GridUR{T,d,D} <: AbstractGrid{T,d,D}
    dims::Dims{D}
    first0::T # first-1

    function GridUR(dims::Dims{D}, v::AbstractRange{T}, d::Int) where {T,D}
        1 ≤ d ≤ D || throw(ArgumentError("$d for $dims"))
        T == Int || throw(ArgumentError("only Int supported"))
        new{T,d,D}(dims, first(v) - one(T))
    end
end

Base.size(a::GridUR) = a.dims
Base.eltype(::GridUR{T}) where T = T

@inline Base.@propagate_inbounds function Base.getindex(
    a::GridUR{T,d,D},
    i::Vararg{Int,D},
) where {T <: Int, d, D}
    @boundscheck checkbounds(a, i...)
#   @boundscheck checkbounds(i, d)
    return a.first0 + @inbounds i[d] # i don't want to deal with one(T) here
end

#=
"""
    (xg, yg, ...) = ndgrid(a:b, c:d, ...)
Construct `ndgrid` tuple for `UnitRange{Int}` inputs.

# Example

```jldoctest
julia> ndgrid(1:2, 3:5)
([1 1 1; 2 2 2], [3 4 5; 3 4 5])
```
"""
function ndgrid(vs::Vararg{UnitRange{Int}})
    D = length(vs)
    dims = ntuple(i -> length(vs[i]), D)
    return ntuple(i -> GridUR(dims, vs[i], i), D)
end
=#


text:
using AugmentedGaussianProcesses
using Distributions,LinearAlgebra
using Random: seed!; seed!(42)
using Plots; pyplot()

N = 200;#Number of training points
N_test = 1000 # Number of points for predictions
X = sort(rand(N)) # Training points
shift = 0.1*(maximum(X)-minimum(X))
X_test = collect(range(minimum(X)-shift,length=N_test,stop=maximum(X)+shift))
## Noise parameters
σ = 5e-1 #Gaussian noise
ν = 4.0; #Degrees of freedom for Student T
β = 5e-1; # Shape for Laplace
mu0 = 0.3; # Constant for heteroscedastic gp
logit(x) =  1.0./(1.0.+exp.(-x))
## Kernel
k = RBFKernel(0.1) # Squared exponential kernel
K = Symmetric(kernelmatrix(reshape(X,:,1),k)+1e-9*Diagonal(I,N)) #Kernel matrix
f = rand(MvNormal(zeros(N),K)) #Random GP sampled from the prior
g = rand(MvNormal(mu0*ones(N),K)) #Random GP sampled from the prior for Heteroscedasticity
y_gaussian = f .+ rand(Normal(0,σ),N)
y_studentt = f .+ rand(TDist(ν),N)
y_laplace = f .+ rand(Laplace(0,β),N)
y_hetero = f .+ randn(N).*sqrt.(inv.(logit.(g)))

gpmodel= GP(X,y_gaussian,k,noise=σ^2);
train!(gpmodel,iterations=50)#,callback=intplot)
gppred,gppred_cov = proba_y(gpmodel,X_test)

stumodel = VGP(X,y_studentt,k,StudentTLikelihood(ν),AnalyticVI());
train!(stumodel,iterations=50)#,callback=intplot)
stupred,stupred_cov = proba_y(stumodel,X_test)

lapmodel = VGP(X,y_laplace,k,LaplaceLikelihood(β),AnalyticVI());
train!(lapmodel,iterations=50)#,callback=intplot)
lappred,lappred_cov = proba_y(lapmodel,X_test)

hetmodel = VGP(X,y_hetero,k,HeteroscedasticLikelihood(k,convert(PriorMean,mu0)),AnalyticVI())
train!(hetmodel,iterations=50)
hetpred,hetpred_cov = proba_y(hetmodel,X_test)

## Plotting
lw = 3.0 #Width of the lines
nsig = 2 #Number of sigmas for the uncertainty
falpha = 0.3

pstudent = plot(X,y_studentt,t=:scatter,lab="")
plot!(X_test,stupred,title="StudentT Regression",lw=lw,color=1,lab="")
plot!(X,f,color=:black,lw=1.0,lab="")
plot!(X,f.+nsig*sqrt(ν/(ν-2)),lab="",color=:black,lw=1.0,linestyle=:dash)
plot!(X,f.-nsig*sqrt(ν/(ν-2)),lab="",color=:black,lw=1.0,linestyle=:dash)
maxlims = ylims(pstudent)
plot!(X_test,stupred.+ nsig  .* sqrt.(stupred_cov),linewidth=0.0,
    fillrange=stupred .- nsig .* sqrt.(stupred_cov),
    fillalpha=falpha,
    fillcolor=1,
    label="")
ylims!(maxlims)

pgauss = plot(X,y_gaussian,t=:scatter,lab="")
plot!(X_test,gppred,title="Gaussian Regression",lw=lw,color=1,lab="")
plot!(X,f,color=:black,lw=1.0,lab="")
plot!(X,f.+nsig*σ,lab="",color=:black,lw=1.0,linestyle=:dash)
plot!(X,f.-nsig*σ,lab="",color=:black,lw=1.0,linestyle=:dash)
plot!(X_test,gppred .+ nsig .* sqrt.(gppred_cov),linewidth=0.0,
    fillrange=gppred .- nsig  .* sqrt.(gppred_cov),
    fillalpha=falpha,
    fillcolor=1,
    label="")
ylims!(maxlims)

plaplace = plot(X,y_laplace,t=:scatter,lab="")
plot!(X_test,lappred,title="Laplace Regression",lw=lw,color=1,lab="")
plot!(X,f,color=:black,lw=1.0,lab="")
plot!(X,f.+nsig*β*sqrt(2),lab="",color=:black,lw=1.0,linestyle=:dash)
plot!(X,f.-nsig*β*sqrt(2),lab="",color=:black,lw=1.0,linestyle=:dash)
plot!(X_test,lappred.+ nsig  .* sqrt.(lappred_cov),linewidth=0.0,
    fillrange=lappred .- nsig .* sqrt.(lappred_cov),
    fillalpha=falpha,
    fillcolor=1,
    label="")
ylims!(maxlims)

phetero = plot(X,y_hetero,t=:scatter,lab="")
plot!(X_test,lappred,title="Heteroscedastic Regression",lw=lw,color=1,lab="")
plot!(X,f,color=:black,lw=1.0,lab="")
plot!(X,f.+nsig*sqrt.(inv.(logit.(g))),lab="",color=:black,lw=1.0,linestyle=:dash)
plot!(X,f.-nsig*sqrt.(inv.(logit.(g))),lab="",color=:black,lw=1.0,linestyle=:dash)
plot!(X_test,hetpred.+ nsig  .* sqrt.(hetpred_cov),linewidth=0.0,
    fillrange=hetpred .- nsig .* sqrt.(hetpred_cov),
    fillalpha=falpha,
    fillcolor=1,
    label="")
ylims!(maxlims)

default(legendfontsize=10.0,xtickfontsize=10.0,ytickfontsize=10.0)
p=plot(pgauss,pstudent,plaplace,phetero)
display(p)
savefig(p,joinpath(@__DIR__,"Regression.png"))


text:
# This file includes blocks of the communication systems

export SymbolGenerator, Modulator, AWGNChannel, Detector

# ------------------------------ Symbol Generator --------------------------------- #

"""
    $TYPEDEF

Bit generator

# FIELDS

    $TYPEDFIELDS
"""
struct SymbolGenerator
    "Generated symbols"
    symbols::Vector{Int}
    SymbolGenerator(ns::Int, M::Int) = new(rand(1:M, ns))
end

# ------------------------------ CSK Modulator --------------------------------- #

abstract type AbstractScheme end

"""
    $TYPEDEF

Chaos Shift Keying scheme
"""
struct CSK <: AbstractScheme end

"""
    $TYPEDEF

Differential Chaos Shift Keying
"""
struct DCSK <: AbstractScheme end

"""
    $TYPEDEF

Chaos-Shif-Keying Modulator

# FIELDS

    $TYPEDFIELDS
"""
struct Modulator{T<:AbstractScheme, S<:AbstractOscillator}
    "Modulation scheme"
    scheme::T
    "Chaotic generator"
    gens::Vector{S}
    "Symbol duration"
    tsymbol::Float64
    "Sampling period"
    tsample::Float64
end

(mdltr::Modulator)(symbols) = iscontinuous(mdltr) ? contmodulate!(mdltr, symbols) : discmodulate!(mdltr, symbols)

function discmodulate!(modulator::Modulator, symbols)
    K = numsamples(modulator)
    refs = map(gen -> trajectory!(gen, K - 1), modulator.gens)
    refs, refs[symbols]
end

function contmodulate!(modulator::Modulator, symbols)
    tf = modulator.tsymbol
    ts = modulator.tsample
    refs = map(gen -> trajectory!(gen, tf - ts, ts), modulator.gens)
    refs, refs[symbols]
end

"""
    $SIGNATURES

Returns true if `modulator` operates in continuous time
"""
iscontinuous(modulator::Modulator) = typeof(modulator.scheme) <: AbstractContinuousOscillator

"""
    $SIGNATURES

Returns the number of samples per symbol of `modulator`
"""
numsamples(modulator::Modulator) = Int(modulator.tsymbol / modulator.tsample)

# ------------------------------ Channel  --------------------------------- #

"""
    $SIGNATURES

Additive white Gaussian noise channel

# Fields
    $TYPEDFIELDS
"""
mutable struct AWGNChannel
    "Energy per symbol to noise power spectral density ratio (in dB)"
    esno::Float64
    "Symbol duration"
    tsymbol::Float64
    "Sampling period"
    tsample::Float64
end

function (channel::AWGNChannel)(tx)
    numsymbols = length(tx)
    numsamples = length(tx[1])
    ts = channel.tsample
    Es = sum(energy.(tx)) / numsymbols
    No = Es / (2 * dbtoval(channel.esno)) * ts
    n = collect(eachrow(sqrt(σ) * randn(numsymbols, numsamples)))
    tx + n
end

"""
    $SIGNATURES

Converts `dB` to numerical value.
"""
dbtoval(db) = 10^(db / 10)

"""
    $SIGNATURES

Converts `dB` to numerical value.
"""
valtodb(db) = 10 * log10(db)

"""
    $SIGNATURES

Returns the energy of `x` sampled with a sampling frequency `ts`.
"""
energy(x, ts=1) = sum(abs.(x).^2) * ts

# ---------------------------------- Detector --------------------------------- #
"""
    $TYPEDEF

Correlation detector
"""
struct Detector end

function (detector)(refs, rx)
    map(enumerate(rx)) do (i, rxi)
        argmax(map(ref -> map(rxi ⋅ ref[i]), refs))
    end
end