text:
# Default configuration of operators and functions

import .Parsers: Op, Un, OpPolicy
import Base: show

function prelude(ops::OpPolicy)

    prec_levels = [
        Op.([:.]),
        Op.([:^]),
        Un.([:+, :-, :!, :~]),
        Op.([:*, :/, :%, :÷]),
        Op.([:-]),
        Op.([:+]),
        Op.([:(:)]),
        Op.([:<<, :>>, :>>>]),
        Op.([:<, :<=, :≤, :>, :>=, :≥]),
        Op.([:(==), :!=, :≠]),
        Op.([:&]),
        Op.([:⊻]),
        Op.([:|]),
        Op.([:&&]),
        Op.([:||]),
        Op.([:(=), :(:=)]),
    ]

    for (lvl1, lvl2) in zip(prec_levels[1:end-1], prec_levels[2:end])
        for op1 in lvl1, op2 in lvl2
            push!(ops, op2 => op1)
        end
    end

    return ops

end

function prelude(env::Env)

    # Named literals

    define(env, Symbol("true"),  true)   # :true and :false don't work
    define(env, Symbol("false"), false)
    define(env, :nil,            nil)

    # Operators

    for op in [:^, :+, :-, :!, :~, :*, :/, :%, :÷, :-, :+, :(:), :<<, :>>,
               :>>>, :<, :<=, :≤, :>, :>=, :≥, :(==), :!=, :≠, :&, :⊻, :|]
        define(env, op, getfield(Base, op))
    end

    # Basic functions

    for func in [
                 # Rounding functions
                 :round, :floor, :ceil, :trunc,

                 # Division functions
                 :div, :fld, :cld, :rem, :mod, :mod1, :mod2pi, :divrem,
                 :fldmod, :gcd, :lcm,

                 # Sign and absolute value functions
                 :abs, :abs2, :sign, :signbit, :copysign, :flipsign,

                 # Powers, logs and roots
                 :sqrt, :cbrt, :hypot, :exp, :expm1, :ldexp, :log, :log2,
                 :log10, :log1p, :exponent, :significand,

                 # Trigonometric and hyperbolic functions
                 :sin, :cos, :tan, :cot, :sec, :csc, :sinh, :cosh, :tanh,
                 :coth, :sech, :csch, :asin, :acos, :atan, :acot, :asec,
                 :acsc, :asinh, :acosh, :atanh, :acoth, :asech, :acsch, :sinc,
                 :cosc,

                 # Trigonometric functions using degrees
                 :sind, :cosd, :tand, :cotd, :secd, :cscd, :asind, :acosd,
                 :atand, :acotd, :asecd, :acscd,
                ]
        define(env, func, getfield(Base, func))
    end

    # Output

    define(env, :stdout,    IOBuffer())

    # Special forms

    define(env, :do,        Do())
    define(env, :.,         Dot())
    define(env, :at,        At())
    define(env, :print,     Print())
    define(env, :&&,        And())
    define(env, :||,        Or())
    define(env, :if,        If())
    define(env, :while,     While())
    define(env, :for,       For())
    define(env, :(:=),      Define())
    define(env, :(=),       Assign())
    define(env, :procedure, ProcDef())
    define(env, :return,    Return())
    define(env, :break,     Break())
    define(env, :continue,  Continue())

    return env

end

# Special forms, which require different evaluation than functions

abstract type Form end
abstract type Sequential <: Form end
abstract type Loop <: Form end

show(io::IO, form::Form) = print(io::IO, nameof(typeof(form)), "()")

apply!(ev::Eval, ::Sequential, args) = return!(ev, last(ev))

# The assumption here is that the special form is accessed via a symbol, so
# that the evaluation of the head itself is of the form
# eval!(ev, ::Symbol, ex). Thus the calls to eval! here are assumed to only
# apply to the arguments.

# apply! should return the Eval object for the next execution step, if any.

## do: really just a function, but grouped with other "statements" ###########

struct Do       <: Sequential end

## print: needs to access stdout #############################################

struct Print    <: Form       end

pfmt(x)         = ppfmt(x)
pfmt(x::String) = x # To override ppfmt

apply!(ev::Eval, ::Print, args) =
    (print(lookup(ev, :stdout)::IO, pfmt.(args)...); return!(ev, nil))

## ||, &&: short-circuit evaluation ##########################################

struct And      <: Sequential end
struct Or       <: Sequential end

conj(ev, ex, halt) = halt ? return!(ev, last(ev)) :
                            Eval(ex, par=ev, env=ev.env)

eval!(ev::Eval, ::Or, ex)  = conj(ev, ex, argn(ev.tgt) ≥ 1 &&  last(ev))
eval!(ev::Eval, ::And, ex) = conj(ev, ex, argn(ev.tgt) ≥ 1 && !last(ev))

## if: evaluates only relevant clauses #######################################

struct If       <: Sequential end

function eval!(ev::Eval, ::If, ex)

    # As arguments to if, both if and else are treated as "decorative noise
    # words", and are ignored.
    (ex == :if || ex == :else) && return skip!(ev)

    vals = args(ev.tgt)
    pos  = length(vals) + 1

    if isodd(pos) # A condition or an unconditional final else clause

        # A final, unconditional else clause
        pos == argn(ev) && !vals[pos-2] && return eval!(ev, ex)

        # No condition succeeded yet, so evaluate this:
        (pos == 1 || !vals[pos-2]) && return eval!(ev, ex)

        # The previous condition succeeded:
        return apply!(ev)

    else # A clause

        # Our condition succeeded:
        vals[pos-1] && return eval!(ev, ex)

        # Skipping this clause:
        return eval!(ev, nil)

    end

end

## while loops: evaluates a variable number of times #########################

struct While    <: Loop end

function eval!(ev::Eval, ::While, ex)

    ev.idx == 1 && return eval!(ev, ex)

    @assert ev.idx == 2

    if !first(args(ev.tgt))
        return!(ev, nil)
    else
        Eval(ex, par=ev, env=Env(ev.env))
    end

end

apply!(ev::Eval, ::While, args) = rewind!(ev)

## Assignment: rebinds one, evaluates the other ##############################

abstract type Assignment <: Form end
struct Define   <: Assignment end
struct Assign   <: Assignment end

struct LHS{T} <: Expression
    ref::T
end

lhseval!(ev::Eval, ex::Symbol) = noeval!(ev, ex)
lhseval!(ev::Eval, ex::Call) = LHS(eval!(ev::Eval, ex))

# Left-hand-side calls are evaluated as normal, except for the application:
step!(lhs::LHS) = done(lhs.ref) ? noapply!(lhs.ref) : LHS(eval!(lhs.ref))

eval!(ev::Eval, ::Assignment, ex) =
    ev.idx == 1 ? lhseval!(ev, ex) : eval!(ev, ex)

apply!(ev::Eval, ::Define, args) =
    ((x, y) = args; define(ev, x, y); return!(ev, y))

function apply!(ev::Eval, ::Assign, (x, y))
    if x isa Symbol
        assign(ev, x, y)
    elseif head(x) isa Dot
        a = collect(args(x))
        lhs = foldl(getfield, a[1:end-1])
        setfield!(lhs, a[end], y)
    else
        @assert head(x) isa At
        a = collect(args(x))
        obj, idx = a[1], int.(a[2:end])
        obj[idx...] = y
    end
    return!(ev, y)
end

## Procedure objects #########################################################

struct Proc     <: Sequential
    name    :: Symbol
    pars    :: Vector{Symbol}
    body    :: Expression
    env     :: Env
end

show(io::IO, proc::Proc) = print(io, "Proc(:$(proc.name))")

function apply!(ev::Eval, proc::Proc, args)

    env = Env(proc.env)
    for (par, arg) in zip(proc.pars, args)
        define(env, par, arg)
    end

    Eval(proc.body, env=env, par=ev.par, frame=true)

end

## Procedure definitions #####################################################

struct ProcDef  <: Form end

eval!(ev::Eval, ::ProcDef, ex) = noeval!(ev, ex)

function apply!(ev::Eval, ::ProcDef, (sign, body))

    name = head(sign)
    pars = args(sign)

    proc = Proc(name, pars, body, ev.env)

    define(ev, name, proc)

    push!(parent(ev), nil)
    parent(ev)

end

## return ####################################################################

struct Return   <: Form end

function apply!(ev::Eval, ::Return, args)

    v = collect(args)
    @assert length(v) ≤ 1
    res = isempty(v) ? nil : first(args)

    ret = parent(ev)
    while !isframe(ret)
        ret = parent(ret)
    end

    return!(ret, res)

end

## break #####################################################################

struct Break    <: Form end

function apply!(ev::Eval, ::Break, ())

    ret = parent(ev)
    while !(head(ret) isa Loop)
        ret = parent(ret)
    end

    return!(ret, nil)

end

## continue ##################################################################

struct Continue <: Form end

function apply!(ev::Eval, ::Continue, ())

    ret = parent(ev)
    while !(head(ret) isa Loop)
        ret = parent(ret)
    end

    rewind!(ret)
    ret

end

## for #######################################################################

# for i in vals { dosomething }
#
# Can be rendered differently if vals is an OrdinalRange (e.g., "for i = 1 to
# n" or "for i = n downto 1 by 2" or the like) than if it's any other iterable
# object.
#
# Note that unless/until negative numbers are handled differently from
# unary operators, things like n:-1:1 won't work. Use n:(-1):1 or n : -1 : 1.

struct For      <: Loop end

function eval!(ev::Eval, ::For, ex)

    ev.idx == 2 && @assert ex == :in
    ev.idx <= 2 && return skip!(ev)
    ev.idx == 3 && return eval!(ev, ex)

    @assert ev.idx == 4

    i_name = first(args(ev.src))

    # Defined in the environment ev has inherited, so it might clash with
    # other loops, i.e., be reused by them, but that doesn't matter.
    i_state_name = Symbol("$i_name iteration state")

    range = first(args(ev.tgt))

    env = Env(ev.env)

    if haskey(ev.env.vars, i_state_name)
        i_state = ev.env.vars[i_state_name]
        next = iterate(range, i_state)
    else
        next = iterate(range)
    end

    if next ≡ nil
        # Clean up after us:
        haskey(ev.env.vars, i_state_name) && delete(ev.env, i_state_name)
        return return!(ev, nil)
    end

    i, i_state = next

    define(env, i_name, i)
    define(ev, i_state_name, i_state)

    Eval(ex, par=ev, env=env)

end

apply!(ev::Eval, ::For, args) = rewind!(ev, 1)

## attribute access ##########################################################

struct Dot      <: Form end

eval!(ev::Eval, ::Dot, ex) = ev.idx == 1 ? eval!(ev, ex) : noeval!(ev, ex)
apply!(ev::Eval, ::Dot, args) = return!(ev, foldl(getfield,  args))

## indexing ##################################################################

struct At       <: Form end

int(x)         = x
int(x::Number) = Int(x)

apply!(ev::Eval, ::At, args) = return!(ev, getindex(int.(args)...))


text:
struct Placeholder{T <: Number}
    r::T
    c::T
end

struct scaleDataType
    newDim::Placeholder{Int}
    nTiles::Placeholder{Int}
    paddedDim::Placeholder{Int}
    offset::Placeholder{Int}
    scale::Float64
    alpha::Placeholder{Float64}
    isFFT::Bool
    transform::Function
end

scale_chirpz(eb::Eigenbrot, xfactor::Real, yfactor::Real) =
    scale_y_chirpz(scale_x_chirpz(eb,xfactor), yfactor)

function scale_x_chirpz(eb::Eigenbrot, factor::Real)
    factor == 1 && return copy(eb)
    h, w = size(eb)
    newDim = Placeholder(h, round(Int, w * factor))
    sc = scaleDataType(
        newDim,                                  # newDim
        Placeholder(1, ceil(Int, newDim.c / w)), # nTiles
        Placeholder(h, 2 * w),                   # paddedDim
        Placeholder(0, w >> 1),                  # offset
        sqrt(2.0),                               # scale
        Placeholder(0.0, π / (w * factor)),      # alpha
        eb.fft,                                  # isFFT
        fftx                                     # transform
    )
    scale_1d(eb, sc)
end

function scale_y_chirpz(eb::Eigenbrot, factor::Real)
    factor == 1 && return copy(eb)
    h, w = size(eb)
    newDim = Placeholder(round(Int, h * factor), w)
    sc = scaleDataType(
        newDim,                                  # newDim
        Placeholder(ceil(Int, newDim.r / h), 1), # nTiles
        Placeholder(2 * h, w),                   # paddedDim
        Placeholder(h >> 1, 0),                  # offset
        sqrt(2.0),                               # scale
        Placeholder(π / (h * factor), 0.0),      # alpha
        eb.fft,                                  # isFFT
        ffty                                     # transform
    )
    scale_1d(eb, sc)
end

function scale_1d(input::Eigenbrot, sc::scaleDataType)
    h, w = size(input)
    result = Eigenbrot(sc.newDim.r, sc.newDim.c, sc.isFFT)
    quadProd1 = Eigenbrot(sc.paddedDim.r, sc.paddedDim.c, !sc.isFFT)
    fill!(quadProd1.vals, 0.0)
    inputFFT = sc.transform(input, true)
    inputFFT.fft = sc.isFFT

    for cc in 1:w
        r2x = (div(w, 2) - cc + 1) ^ 2 * sc.alpha.c
        for rr in 1:h
            r2y = (div(h, 2) - rr + 1) ^ 2 * sc.alpha.r
            quadProd1[rr + sc.offset.r, cc + sc.offset.c] =
                inputFFT[rr, cc] * (cos(r2x + r2y) - im * sin(r2x + r2y))
        end
    end
    quadFFT1 = sc.transform(quadProd1, true)
    quadProd2 = Eigenbrot(sc.paddedDim.r, sc.paddedDim.c, !sc.isFFT)
    fftProd = Eigenbrot(sc.paddedDim.r, sc.paddedDim.c, sc.isFFT)
    for deltaHor in 1:sc.nTiles.c
        firstXpixel = (deltaHor == sc.nTiles.c && mod(sc.newDim.c, w) != 0) ?
            w - mod(sc.newDim.c, w) + 1 : 1
        for deltaVer in 1:sc.nTiles.r
            firstYpixel = (deltaVer == sc.nTiles.r && mod(sc.newDim.r, h) != 0) ?
                h - mod(sc.newDim.r, h) + 1 : 1

            for cc in 1:sc.paddedDim.c
                r2x = (div(sc.paddedDim.c, 2) - cc + 1 + (deltaHor - 1) * w) ^ 2 * sc.alpha.c
                for rr in 1:sc.paddedDim.r
                    r2y = (div(sc.paddedDim.r, 2) - rr + 1 + (deltaVer - 1) * h) ^ 2 * sc.alpha.r
                    quadProd2[rr, cc] = cos(r2x + r2y) + im * sin(r2x + r2y)
                end
            end
            quadFFT2 = sc.transform(quadProd2, true)
            fftProd = quadFFT2 * quadFFT1
            conv = sc.transform(fftProd, true)
            for cc in firstXpixel:w
                r2x = (div(w, 2) - cc + 1 + (deltaHor - 1) * w) ^ 2 * sc.alpha.c
                col = mod1(cc + div(sc.newDim.c - w, 2) - (deltaHor - 1) * w + sc.newDim.c,
                           sc.newDim.c)
                for rr in firstYpixel:h
                    r2y = (div(h, 2) - rr + 1 + (deltaVer - 1) * h) ^ 2 * sc.alpha.r
                    row = mod1(rr + div(sc.newDim.r - h, 2) - (deltaVer - 1) * h + sc.newDim.r,
                               sc.newDim.r)
                    newval = sc.scale * conv[rr + sc.offset.r, cc + sc.offset.c] *
                        (cos(r2x + r2y) - im * sin(r2x + r2y))
                    result[row,col] =
                        sc.scale * conv[rr + sc.offset.r, cc + sc.offset.c] *
                        (cos(r2x + r2y) - im * sin(r2x + r2y))
                end
            end
        end
    end
    return result
end

ext: jl


text:
# This file is part of IntegerSequences.
# Copyright Peter Luschny. License is MIT.

(@__DIR__) ∉ LOAD_PATH && push!(LOAD_PATH, (@__DIR__))

module Polynomials

using Nemo
# import AbstractAlgebra.lead

export ModulePolynomials
export Poly, AltPoly, EgfPoly, AltEgfPoly
export Coeffs, CoeffSum, CoeffAltSum, CoeffConst, CoeffLeading, AltCoeffs
export Diagonal, Central, EgfCoeffs, AltEgfCoeffs, ReflectPoly

"""

Mostly convenient functions to deal with polynomials as often used in connection with ordinary and exponential generating functions. The naming scheme used is roughly described by:

```
   Poly       <-> Coeffs
   AltPoly    <-> Poly(AltCoeffs)
   EgfPoly    <-> Poly(EgfCoeffs)
   OgfPoly    <-> Poly(OgfCoeffs)
   AltEgfPoly <-> Poly(AltEgfCoeffs)
```

Here 'Alt' stands for alternating, 'Egf' for exponential generating function, 'Ogf' for ordinary generating function.

* Coeffs, CoeffSum, CoeffAltSum, CoeffConst, CoeffLeading, AltCoeffs, Diagonal, Central, EgfCoeffs, AltEgfCoeffs, Poly, AltPoly, EgfPoly, AltEgfPoly, ReflectPoly.
"""
const ModulePolynomials = ""

"""

Return the coefficients of the polynomial ``p``.
"""
Coeffs(p::fmpz_poly) = [coeff(p, k) for k ∈ 0:degree(p)]

"""

Return the coefficients of the polynomial ``p`` with alternating signs.
"""
AltCoeffs(p::fmpz_poly) = [(-1)^k * coeff(p, k) for k ∈ 0:degree(p)]

"""

Return the coefficients of the polynomial ``p`` multiplied by ``k!``.
"""
OgfCoeffs(p::fmpz_poly) = [factorial(k)*coeff(p, k) for k ∈ 0:degree(p)]

"""

Return the coefficients of the polynomial ``p`` divided by ``k!``. Note that integer division is used.
"""
EgfCoeffs(p::fmpz_poly) = [div(coeff(p, k), factorial(k)) for k ∈ 0:degree(p)]

"""

Return the coefficients of the polynomial ``p`` divided by ``(-1)^k k!``.
"""
AltEgfCoeffs(p::fmpz_poly) = [(-1)^k * div(coeff(p, k), factorial(k)) for k ∈ 0:degree(p)]

"""

Return the polynomial ``p`` with the coefficients C.
"""
function Poly(C)
    T, x = PolynomialRing(ZZ, "x")
    sum(c * x^k for (k, c) ∈ enumerate(C))
end

"""

Return the polynomial ``p`` with the coefficients C and alternating signs (i.e. with ``(-1)^k c[k] x^k)``.
"""
function AltPoly(C)
    T, x = PolynomialRing(ZZ, "x")
    sum((-1)^k * c * x^k for (k, c) ∈ enumerate(C))
end

"""

Return the polynomial ``p`` with the coefficients C used in the form ``c[k] x^k/k!``. Note that integer division is used.
"""
function EgfPoly(C)
    T, x = PolynomialRing(ZZ, "x")
    sum(div(c, factorial(k)) * x^k for (k, c) ∈ enumerate(C))
end

"""

Return the polynomial ``p`` with the coefficients C used in the form ``c[k] k! x^k``.
"""
function OgfPoly(C)
    T, x = PolynomialRing(ZZ, "x")
    sum(c * factorial(k) * x^k for (k, c) ∈ enumerate(C))
end

"""

Return the polynomial ``p`` with alternating signs attached to the coefficients .
"""
AltPoly(p::fmpz_poly) = Poly(AltCoeffs(p))
#    T, x = PolynomialRing(ZZ, "x")
#    sum((-1)^k*coeff(p, k)*x^k for k ∈ 0:degree(p))

"""

Return the polynomial ``p`` with coefficients in exponential form (i.e. with ``c[k] x^k/k!``).
"""
EgfPoly(p::fmpz_poly) = Poly(EgfCoeffs(p))
#    T, x = PolynomialRing(ZZ, "x")
#    sum(div(coeff(p, k), factorial(k))*x^k for k ∈ 0:degree(p))

"""

Return the polynomial ``p`` with coefficients in exponential form and alternating signs (i.e. with ``(-1)^k c[k] x^k/k!``).
"""
AltEgfPoly(p::fmpz_poly) = Poly(AltEgfCoeffs(p))
#    T, x = PolynomialRing(ZZ, "x")
#    sum((-1)^k*div(coeff(p, k), factorial(k))*x^k for k ∈ 0:degree(p))

"""

Return the list of the coefficients of the first ``len`` polynomials of the sequence of polynomials ``P`` as a triangle.
"""
Coeffs(P, len) = [[coeff(P(n), k) for k ∈ 0:degree(P(n))] for n ∈ 0:len-1]

"""

Return the sum of the coefficients of the polynomial ``p``.
"""
CoeffSum(p) = evaluate(p, 1)

"""

Return the sequence of the sum of coefficients of the sequence of polynomials ``P``.
"""
CoeffSum(P, len) = [CoeffSum(P(n)) for n ∈ 0:len-1]

"""

Return the alternating sum of the coefficients of the polynomial ``p``.
"""
CoeffAltSum(p) = evaluate(p, -1)

"""

Return the sequence of the alternating sums of the coefficients of the sequence of polynomials ``P``.
"""
CoeffAltSum(P, len) = [CoeffAltSum(P(n)) for n ∈ 0:len-1]

"""

Return the leading coefficient of the polynomial ``p``.
"""
CoeffLeading(p) = coeff(p, degree(p)) # lead(p)

"""

Return the sequence of the leading coefficient of the sequence of polynomials ``P``.
"""
Diagonal(P, len) = [CoeffLeading(P(n)) for n ∈ 0:len-1]

"""

Return the constant coefficient of the polynomial ``p``.
"""
CoeffConst(p) = coeff(p, 0)

"""

Return the sequence of the constant coefficients of the sequence of polynomials ``P``.
"""
CoeffConst(P, len) = [CoeffConst(P(n)) for n ∈ 0:len-1]

"""

Return the central column of the coefficients of the sequence of polynomials ``P``.
"""
Central(P, len) = [Coeffs(P(2n))[n+1] for n ∈ 0:len-1]

"""

Return the reflected polynomial of ``p``.
"""
function ReflectPoly(p::fmpz_poly)
    T, x = PolynomialRing(ZZ, "x")
    p(0) != 1 && throw(ValueError("Constant coefficient must be 1."))
    d = degree(p)
    x^d + sum(coeff(p, k) * x^(d - k) for k ∈ 1:d)
end

#START-TEST-########################################################

using Test, SeqUtils

function test()

    @testset "Polynomials" begin
        T, x = PolynomialRing(ZZ, "x")
        p = 63063000 * x^4 + 2702700 * x^3 + 16510 * x^2 + x
        @test Coeffs(p) == [0, 1, 16510, 2702700, 63063000]
        @test CoeffSum(p) == 65782211
        @test CoeffAltSum(p) == 60376809
        @test CoeffConst(p) == T(0)
        @test CoeffLeading(p) == 63063000

        q = x^2 - x - 1
        p = 1 - x - x^2
        r = ReflectPoly(p)
        @test r == q
    end

end

function demo()

    function P(m, n)
        R, x = PolynomialRing(ZZ, "x")
        function recP(m, n)
            n == 0 && return R(1)
            sum(binomial(m * n, m * k) * recP(m, n - k) * x for k ∈ 1:n)
        end
        recP(m, n)
    end

    # ... given a polynomial p:
    for m ∈ 0:4, n ∈ 0:5
        println("---> m: $m, n: $n")
        q = P(m, n)

        typeof(q) |> println
        q |> println
        AltPoly(q) |> println
        EgfPoly(q) |> println
        AltEgfPoly(q) |> println

        Coeffs(q) |> Println
        EgfCoeffs(q) |> Println
        AltEgfCoeffs(q) |> Println

        CoeffSum(q) |> println
        CoeffAltSum(q) |> println
        CoeffConst(q) |> println
        CoeffLeading(q) |> println
    end

    # ... given a sequence of polynomials P:
    for m ∈ 0:4
        println("---> m ", m)
        println("\nPolynomial:")
        Q(n) = P(m, n)
        println(Q(4))
        println("\nTriangle of coefficients:")
        Coeffs(Q, 7) |> Println
        println("\nSum of coefficients:")
        CoeffSum(Q, 7) |> Println
        println("\nAlternating sum of coefficients:")
        CoeffAltSum(Q, 7) |> Println
        println("\nLeading coefficients:")
        Diagonal(Q, 7) |> Println
        println("\nCentral column of coefficient triangle:")
        Central(Q, 7) |> Println
    end
end

function perf() end

function main()
    test()
    demo()
    perf()
end

main()

end # module

max_stars_repo_name: OpenLibMathSeq/Sequences

max_issues_repo_name: OpenLibMathSeq/Sequences


text:
module InteratomicPotentials

using StaticArrays
using LinearAlgebra
using AtomsBase
using Unitful
using UnitfulAtomic

# Import files
include("NeighborList/nnlist.jl")
include("Parameters/param.jl")
include("PotentialTypes/types.jl")

# Export energies, forces, virial evaluations
export energy_and_force, potential_energy, force, virial, virial_stress
export grad_potential_energy, grad_force, grad_virial, grad_virial_stress

end


text:
Inputs = InputSettings(

    # 1. T
    [6],
    # 2. n_items
    Int(250),
    # 3. n_groups
    Int(1),
    # 4. groups
    [""],

    ####################################################################
    ########################## IRT #####################################
    ####################################################################

    # 5. IRT_model
    "2PL",
    # 6. IRT_parameters
    ["a","b"],
    # 7. IRT_parametrization
    "at-ab",
    # 8. IRT_D
    1,

    ####################################################################
    ######################FRIENDS AND ENEMIes###########################
    ####################################################################

    # 9. enemy_sets_var
    String[],
    # 10. friend_sets_var
    String[],

    ####################################################################
    ########################### ITEM USE  ##############################
    ####################################################################

    # 11. item_use_min
    fill(0, 250),
    # 12. item_use_max
    fill(6, 250),

    ####################################################################
    ########################## TesT LENGHT #############################
    ####################################################################

    # 13. lenght_min
    [50, 50, 50],
    # 14. lenght_max
    [50, 50, 50],
    # 15. lenght_weight
    [1.0, 1.0, 1.0],

    ####################################################################
    ####################### EXPECTED SCORE #############################
    ####################################################################

    # 16. expected_score_var
    String[],
    # 17. expected_score_pts
    Vector{Vector{Float64}}(undef, 0),
    # 18. expected_score_min
    Vector{Vector{Float64}}(undef, 0),
    # 19. expected_score_max
    Vector{Vector{Float64}}(undef, 0),

    ####################################################################
    ###################### GENERIC CONSTRAINTS#################
    ####################################################################

    # 20. mean_vars
    Vector{Vector{String}}(undef, 0),  # (future)
    # 21. mean_vars_min
    Vector{Vector{Float64}}(undef, 0), # (future)
    # 22. mean_vars_max
    Vector{Vector{Float64}}(undef, 0), # (future)
    # 23. sum_vars
    Vector{Vector{String}}(undef, 0),
    # 24. sum_vars_min
    Vector{Vector{Float64}}(undef, 0),
    # 25. sum_vars_max
    Vector{Vector{Float64}}(undef, 0),

    ####################################################################
    ######################### OBJECTIVE ################################
    ####################################################################

    # 26. obj_type
    "", # "MAXIMIN", "CCMAXIMIN", "", "custom", "MINIMAX"
    # 27. obj_points (required in MAXIMIN, CC, MINIMAX)
    Vector{Vector{Float64}}(undef, 0),
    # 28. obj_targets (required in MINIMAX)
    Vector{Vector{Float64}}(),
    # 29. obj_aux_int
    0.0,
    # 30. obj_aux_float
    0.0,

    ####################################################################
    ######################### OUTPUT ###################################
    ####################################################################
    # 31. categories
    ["guessed_difficulty"],
);


text:
export QFT, qft
using BitBasis
using LinearAlgebra
using FFTW

"""
    QFT{N} <: PrimitiveBlock{N}

Quantum Fourier Transform node. See also [`qft_circuit`](@ref).
"""
struct QFT{N} <: PrimitiveBlock{N} end

"""
    qft(n)

Create a Quantum Fourier Transform block. See also [`qft_circuit`](@ref).
"""
qft(n) = QFT{n}()

Yao.mat(::Type{T}, q::QFT{N}) where {T, N} = T.(applymatrix(q))

function _apply!(r::ArrayReg, x::QFT)
    α = sqrt(1<<nactive(r))
    st = state(invorder!(r))
    ifft!(st, 1)
    lmul!(α, st)
    return r
end

function _apply!(r::ArrayReg, ::Daggered{<:QFT})
    # (reg.state = invorder_firstdim(fft!(reg|>state, 1)/sqrt(1<<nactive(reg))); reg)
    α = sqrt(1 << nactive(r))
    st = state(r)
    fft!(st, 1)
    st ./= α
    invorder!(r)
    return r
end

# traits
Yao.ishermitian(q::QFT{N}) where N = N==1
Yao.isreflexive(q::QFT{N}) where N = N==1
Yao.isunitary(q::QFT{N}) where N = true

function Yao.print_block(io::IO, pb::QFT{N}) where N
    printstyled(io, "qft(1-$N)"; bold=true, color=:blue)
end

function Yao.print_block(io::IO, pb::Daggered{N,<:QFT}) where {N, T}
    printstyled(io, "iqft(1-$N)"; bold=true, color=:blue)
end

max_stars_repo_name: QuantumBFS/YaoExtensions

max_issues_repo_name: QuantumBFS/YaoExtensions

max_forks_repo_name: QuantumBFS/YaoExtensions


text:
export ShiftedPotentialCutoff, ShiftedForceCutoff, NoCutoff

struct NoCutoff <: AbstractCutoff end

cutoff_points(::Type{NoCutoff}) = 0

struct ShiftedPotentialCutoff{F, T} <: AbstractCutoff
    max_force::F
    cutoff_dist::T
    sqdist_cutoff::T
    inv_sqdist_cutoff::T
end

ShiftedPotentialCutoff(cutoff_dist) = ShiftedPotentialCutoff(
    100,
    cutoff_dist,
    cutoff_dist^2,
    inv(cutoff_dist),
)

cutoff_points(::Type{ShiftedPotentialCutoff{F, T}}) where {F, T} = 1

struct ShiftedForceCutoff{F, T} <: AbstractCutoff
    max_force::F
    cutoff_dist::T
    sqdist_cutoff::T
    inv_sqdist_cutoff::T
end

ShiftedForceCutoff(cutoff_dist) = ShiftedForceCutoff(
    100,
    cutoff_dist,
    cutoff_dist^2,
    inv(cutoff_dist),
)

cutoff_points(::Type{ShiftedForceCutoff{F, T}}) where {F, T} = 1

force_cutoff(::NoCutoff, r2, inter, params) = force_nocutoff(inter, r2, inv(r2), params)
potential_cutoff(::NoCutoff, r2, inter, params) = potential(inter, r2, inv(r2), params)

function force_cutoff(cutoff::ShiftedPotentialCutoff{F}, r2, inter, params) where F
    f = force_nocutoff(inter, r2, inv(r2), params)

    if @generated
        quote
            if !(F === Nothing)
                f = min(f, cutoff.max_force)
            end
        end
    else
        if !(F === Nothing)
            f = min(f, cutoff.max_force)
        end
    end

    return f
end

function potential_cutoff(cutoff::ShiftedPotentialCutoff, r2, inter, params)
    potential(inter, r2, inv(r2), params) - potential(inter, cutoff.sqdist_cutoff, cutoff.inv_sqdist_cutoff, params)
end

function force_cutoff(cutoff::ShiftedForceCutoff{F}, r2, inter, params) where F
    invr2 = inv(r2)
    f = force_nocutoff(inter, r2, invr2, params) - force_nocutoff(
                                inter, cutoff.sqdist_cutoff, cutoff.inv_sqdist_cutoff, params)

    if @generated
        quote
            if !(F === Nothing)
                f = min(f, cutoff.max_force)
            end
        end
    else
        if !(F === Nothing)
            f = min(f, cutoff.max_force)
        end
    end

    return f
end

@fastmath function potential_cutoff(cutoff::ShiftedForceCutoff, r2, inter, params)
    invr2 = inv(r2)
    r = √r2
    rc = cutoff.cutoff_dist
    fc = force_nocutoff(inter, cutoff.sqdist_cutoff, cutoff.inv_sqdist_cutoff, params)

    potential(inter, r2, invr2, params) - (r - rc) * fc -
        potential(inter, cutoff.sqdist_cutoff, cutoff.inv_sqdist_cutoff, params)
end


text:
# imports
using Pkg; Pkg.activate(".")
using GRASS
using Statistics

# plotting imports
using LaTeXStrings
import PyPlot; plt = PyPlot; mpl = plt.matplotlib; plt.ioff()
using PyCall; animation = pyimport("matplotlib.animation");
mpl.style.use(GRASS.moddir * "figures/fig.mplstyle")

# define some functions
include(GRASS.moddir * "figures/fig_functions.jl")

# get command line args and output directories
run, plot = parse_args(ARGS)
grassdir, plotdir, datadir = check_plot_dirs()

# figure 1b -- input bisectors w/ variability
function main()
    # get input data
    bisinfo = GRASS.SolarData()

    # initialize plot objects
    fig, ax1 = plt.subplots()

    # loop and plot curves
    keyz = [(:c, :mu10), (:w, :mu06), (:w, :mu03)]
    labels = [L"\mu = 1.0", L"\mu = 0.6", L"\mu = 0.3"]
    linestyles = ["-", "--", ":"]
    for (i, key) in enumerate(keyz)
        # find average and std
        avg_bis = mean(bisinfo.bis[key], dims=2)
        avg_wav = mean(bisinfo.wav[key], dims=2)
        std_bis = std(bisinfo.bis[key], dims=2)
        std_wav = std(bisinfo.wav[key], dims=2)

        # convert to doppler velocity
        restwav = 5434.5232                         # angstroms
        avg_wav = avg_wav ./ restwav .* GRASS.c_ms
        std_wav = std_wav ./ restwav .* GRASS.c_ms

        # cut off top portion, where uncertainty is large
        ind = findfirst(avg_bis .> 0.86)[1]
        avg_bis = avg_bis[1:ind]
        avg_wav = avg_wav[1:ind]
        std_bis = std_bis[1:ind]
        std_wav = std_wav[1:ind]

        # fix dimensions
        y = reshape(avg_bis, length(avg_bis))
        x1 = reshape(avg_wav .+ std_wav, length(avg_bis))
        x2 = reshape(avg_wav .- std_wav, length(avg_bis))

        # plot the curve
        ax1.fill_betweenx(y, x1, x2, color="C"*string(i-1), alpha=0.5)
        ax1.plot(avg_wav, avg_bis, ls=linestyles[i], color="C"*string(i-1), label=labels[i])
    end

    # set axes labels and save the figure
    ax1.legend(loc="upper right")
    ax1.set_ylim()
    ax1.set_xlabel(L"{\rm Doppler\ Velocity\ (ms}^{-1} {\rm )}")
    ax1.set_ylabel(L"{\rm Normalized\ Intensity}")
    fig.savefig(plotdir * "fig1b.pdf")
    plt.clf(); plt.close()
    println(">>> Figure written to: " * plotdir * "fig1b.pdf")
    return nothing
end

if (run | plot)
    main()
end

max_stars_repo_name: palumbom/GRASS

max_issues_repo_name: palumbom/GRASS

max_forks_repo_name: palumbom/GRASS


text:
"""
    BSplineKit.BoundaryConditions

Contains some boundary condition definitions.
"""
module BoundaryConditions

export Natural

abstract type BoundaryCondition end

"""
    Natural

Generalised natural boundary condition.

This boundary condition is convenient for spline interpolations, as it
provides extra constraints enabling to equate the number of unique B-spline
knots to the number of data points.

For cubic splines (order ``k = 4``), this corresponds to [natural cubic
splines](https://en.wikipedia.org/wiki/Spline_(mathematics)#Examples), imposing
the second derivatives to be zero at the boundaries (``S''(a) = S''(b) = 0``).

For higher-order splines, this boundary condition generalises the standard
natural cubic splines, by setting derivatives of order ``2, 3, …, k/2`` to
be zero at the boundaries.
For instance, for ``k = 6`` (quintic splines), this imposes ``S'' = S''' = 0``.
In practice, BSplineKit.jl achieves this by using [basis recombination](@ref
basis-recombination-api).

Note that, for symmetry reasons, only even-order splines are supported by this BC.
"""
struct Natural <: BoundaryCondition end

end

max_stars_repo_name: jipolanco/BasisSplines

max_issues_repo_name: jipolanco/BasisSplines

max_forks_repo_name: jipolanco/BasisSplines


text:
## PARTE 1
cd(@__DIR__)

##
open("input.txt") do f
    sum(strip(read(f, String))) do x
        2 * (')' - x) - 1
    end
end

## PARTE 2
input = open("input.txt") do f
    strip(read(f, String))
end |> collect;

##
findfirst(==(-1), accumulate((f, x) -> f + 2 * (')' - x) - 1, input; init=0))


text:
# This file is a part of Julia. License is MIT: https://julialang.org/license

using Test
using Base.Threads

# This test has not been added to CI as there can be unexpected delays
# which cause timing-dependent actions to fail.

#=
Test to ensure that the interactive threadpool works as designed.

Task A is a standard task that does a lot of work (~2 seconds) without
yielding. This would prevent ordinarily prevent other tasks from running.

Task B is an interactive task that does a little work (~0.02 seconds) and
yields.

With an interactive threadpool, multiple task Bs should not see notable
delays in execution even when multiple task As are occupying Julia's
default threads.

This test should fail in the absence of an interactive thread.
=#
const N = 263000000 # busywork(N) takes ~1 sec on an i7-9750H @ 2.6GHz
function busywork(n::Int)
    acc = 0
    for i = 1:n
        x = rand(2:10)
        acc += i * x
    end
    return acc
end

function itask()
    h = N ÷ 50
    for i = 1:100
        t1 = time()
        busywork(h)
        yield()
        t2 = time()
        @test t2 - t1 < 0.15
    end
end

it1 = @spawn :interactive itask()
ti1 = @spawn busywork(N * 2);
it2 = @spawn :interactive itask()
ti2 = @spawn busywork(N * 2);
wait(it1)
wait(it2)

max_stars_repo_name: SmiteWindows/julia

max_issues_repo_name: SmiteWindows/julia

max_forks_repo_name: SmiteWindows/julia


text:
abstract type AbstractNamedBijector <: AbstractBijector end

forward(b::AbstractNamedBijector, x) = (rv = b(x), logabsdetjac = logabsdetjac(b, x))

#######################
### `NamedBijector` ###
#######################
"""
    NamedBijector <: AbstractNamedBijector

Wraps a `NamedTuple` of key -> `Bijector` pairs, implementing evaluation, inversion, etc.

# Examples
```julia-repl
julia> using Bijectors: NamedBijector, Scale, Exp

julia> b = NamedBijector((a = Scale(2.0), b = Exp()));

julia> x = (a = 1., b = 0., c = 42.);

julia> b(x)
(a = 2.0, b = 1.0, c = 42.0)

julia> (a = 2 * x.a, b = exp(x.b), c = x.c)
(a = 2.0, b = 1.0, c = 42.0)
```
"""
struct NamedBijector{names, Bs<:NamedTuple{names}} <: AbstractNamedBijector
    bs::Bs
end

# fields contain nested numerical parameters
function Functors.functor(::Type{<:NamedBijector{names}}, x) where names
    function reconstruct_namedbijector(xs)
        return NamedBijector{names,typeof(xs.bs)}(xs.bs)
    end
    return (bs = x.bs,), reconstruct_namedbijector
end

names_to_bijectors(b::NamedBijector) = b.bs

@generated function (b::NamedBijector{names1})(
    x::NamedTuple{names2}
) where {names1, names2}
    exprs = []
    for n in names2
        if n in names1
            # Use processed value
            push!(exprs, :($n = b.bs.$n(x.$n)))
        else
            # Use existing value
            push!(exprs, :($n = x.$n))
        end
    end
    return :($(exprs...), )
end

@generated function Base.inv(b::NamedBijector{names}) where {names}
    return :(NamedBijector(($([:($n = inv(b.bs.$n)) for n in names]...), )))
end

@generated function logabsdetjac(b::NamedBijector{names}, x::NamedTuple) where {names}
    exprs = [:(logabsdetjac(b.bs.$n, x.$n)) for n in names]
    return :(+($(exprs...)))
end

######################
### `NamedInverse` ###
######################
"""
    NamedInverse <: AbstractNamedBijector

Represents the inverse of a `AbstractNamedBijector`, similarily to `Inverse` for `Bijector`.

See also: [`Inverse`](@ref)
"""
struct NamedInverse{B<:AbstractNamedBijector} <: AbstractNamedBijector
    orig::B
end
Base.inv(nb::AbstractNamedBijector) = NamedInverse(nb)
Base.inv(ni::NamedInverse) = ni.orig

logabsdetjac(ni::NamedInverse, y::NamedTuple) = -logabsdetjac(inv(ni), ni(y))

##########################
### `NamedComposition` ###
##########################
"""
    NamedComposition <: AbstractNamedBijector

Wraps a tuple of array of `AbstractNamedBijector` and implements their composition.

This is very similar to `Composed` for `Bijector`, with the exception that we do not require
the inputs to have the same "dimension", which in this case refers to the *symbols* for the
`NamedTuple` that this takes as input.

See also: [`Composed`](@ref)
"""
struct NamedComposition{Bs} <: AbstractNamedBijector
    bs::Bs
end

# Essentially just copy-paste from impl of composition for 'standard' bijectors,
# with minor changes here and there.
composel(bs::AbstractNamedBijector...) = NamedComposition(bs)
composer(bs::AbstractNamedBijector...) = NamedComposition(reverse(bs))
∘(b1::AbstractNamedBijector, b2::AbstractNamedBijector) = composel(b2, b1)

inv(ct::NamedComposition) = NamedComposition(reverse(map(inv, ct.bs)))

function (cb::NamedComposition{<:AbstractArray{<:AbstractNamedBijector}})(x)
    @assert length(cb.bs) > 0
    res = cb.bs[1](x)
    for b ∈ Base.Iterators.drop(cb.bs, 1)
        res = b(res)
    end

    return res
end

(cb::NamedComposition{<:Tuple})(x) = foldl(|>, cb.bs; init=x)

function logabsdetjac(cb::NamedComposition, x)
    y, logjac = forward(cb.bs[1], x)
    for i = 2:length(cb.bs)
        res = forward(cb.bs[i], y)
        y = res.rv
        logjac += res.logabsdetjac
    end

    return logjac
end

@generated function logabsdetjac(cb::NamedComposition{T}, x) where {T<:Tuple}
    N = length(T.parameters)

    expr = Expr(:block)
    push!(expr.args, :((y, logjac) = forward(cb.bs[1], x)))

    for i = 2:N - 1
        temp = gensym(:res)
        push!(expr.args, :($temp = forward(cb.bs[$i], y)))
        push!(expr.args, :(y = $temp.rv))
        push!(expr.args, :(logjac += $temp.logabsdetjac))
    end
    # don't need to evaluate the last bijector, only it's `logabsdetjac`
    push!(expr.args, :(logjac += logabsdetjac(cb.bs[$N], y)))

    push!(expr.args, :(return logjac))

    return expr
end

function forward(cb::NamedComposition, x)
    rv, logjac = forward(cb.bs[1], x)

    for t in cb.bs[2:end]
        res = forward(t, rv)
        rv = res.rv
        logjac = res.logabsdetjac + logjac
    end
    return (rv=rv, logabsdetjac=logjac)
end

@generated function forward(cb::NamedComposition{T}, x) where {T<:Tuple}
    expr = Expr(:block)
    push!(expr.args, :((y, logjac) = forward(cb.bs[1], x)))
    for i = 2:length(T.parameters)
        temp = gensym(:temp)
        push!(expr.args, :($temp = forward(cb.bs[$i], y)))
        push!(expr.args, :(y = $temp.rv))
        push!(expr.args, :(logjac += $temp.logabsdetjac))
    end
    push!(expr.args, :(return (rv = y, logabsdetjac = logjac)))

    return expr
end

############################
### `NamedCouplingLayer` ###
############################
# TODO: Add ref to `Coupling` or `CouplingLayer` once that's merged.
"""
    NamedCoupling{target, deps, F} <: AbstractNamedBijector

Implements a coupling layer for named bijectors.

# Examples
```julia-repl
julia> using Bijectors: NamedCoupling, Scale

julia> b = NamedCoupling(:b, (:a, :c), (a, c) -> Scale(a + c))
NamedCoupling{:b,(:a, :c),var"#3#4"}(var"#3#4"())

julia> x = (a = 1., b = 2., c = 3.);

julia> b(x)
(a = 1.0, b = 8.0, c = 3.0)

julia> (a = x.a, b = (x.a + x.c) * x.b, c = x.c)
(a = 1.0, b = 8.0, c = 3.0)
```
"""
struct NamedCoupling{target, deps, F} <: AbstractNamedBijector where {F, target}
    f::F
end

NamedCoupling(target, deps, f::F) where {F} = NamedCoupling{target, deps, F}(f)
function NamedCoupling(::Val{target}, ::Val{deps}, f::F) where {target, deps, F}
    return NamedCoupling{target, deps, F}(f)
end

coupling(b::NamedCoupling) = b.f
# For some reason trying to use the parameteric types doesn't always work
# so we have to do this weird approach of extracting type and then index `parameters`.
target(b::NamedCoupling{Target}) where {Target} = Target
deps(b::NamedCoupling{<:Any, Deps}) where {Deps} = Deps

@generated function (nc::NamedCoupling{target, deps, F})(x::NamedTuple) where {target, deps, F}
    return quote
        b = nc.f($([:(x.$d) for d in deps]...))
        return merge(x, ($target = b(x.$target), ))
    end
end

@generated function (ni::NamedInverse{<:NamedCoupling{target, deps, F}})(
    x::NamedTuple
) where {target, deps, F}
    return quote
        b = ni.orig.f($([:(x.$d) for d in deps]...))
        return merge(x, ($target = inv(b)(x.$target), ))
    end
end

@generated function logabsdetjac(nc::NamedCoupling{target, deps, F}, x::NamedTuple) where {target, deps, F}
    return quote
        b = nc.f($([:(x.$d) for d in deps]...))
        return logabsdetjac(b, x.$target)
    end
end


text:
# Test relative speed of Deleglise Rivat and tabsieve
function tprimepi(n)
    tdr = @elapsed primepi(n, alg = :dr)
    primesieve_num_threads(8)
    tts = @elapsed primepi(n, alg = :tabsieve)
    return tdr/tts
end

# Test function above m times with random numbers
# between n0 and n0+n1. Report how many times one or
# the other is faster.
# On
function tprimepia(n0,n1,m)
    c1 = 0
    c2 = 0
    for i in 1:m
        nr = int(n1 * rand()) + n0
        res = tprimepi(nr)
        res > 1 ? c1 += 1 : c2 += 1
    end
    return (c1/m,c2/m)
end

tprimepia(n,m) = tprimepia(n,n,m)

function print_tprimepia()
    for (n,m) in ((1,10^5), (2, 10^5),(3, 10^4),
                  (4, 10^4),(5, 10^4),(6, 10^4),
                  (7, 10^4),(8, 10^4),(9, 10^3),
                  (10, 10^3),(11, 10^3),(12, 5*10^2),
                  (13, 10^2))
        println("10^$n: ", tprimepia(10^n,m))
    end
end

function print_tprimepia1()
    m = 10^2
    for n in (1,2,3,4,5,6,7,8,9,10)
        println("$n * 10^11: ", tprimepia(n*10^11,m))
    end
end

function print_sieve(n0)
    x = 10^6
    primesieve_num_threads(8)
    while x <= n0/(10^4)
        t = @elapsed ntcountprimes(n0,n0+x)
        println("$(x+n0)  $t  $(log10(x))")
        x *= 2
    end
end

function print_dr()
    x = 10^5
    while true
        t = @elapsed primepi(x, alg = :dr)
        println("$x  $t")
        x *= 2
    end
end

###############################

function ts0(n,n1)
    for i in 1:n1
        nextprimea(n)
    end
end

function ts1(n,n1)
    for i in 1:n1
        PrimeSieve.nextprime(n)
    end
end

function ps0(n,n1)
    for i in 1:n1
        prevprimea(n)
    end
end

function ps1(n,n1)
    for i in 1:n1
        PrimeSieve.prevprime(n)
    end
end

function ts0b(n,n1)
    for i in 1:n1
        nextprimeb(n)
    end
end

# function ts2(n,n1)
#     for i in 1:n1
#         PrimeSieve.nextprime2(n)
#     end
# end

function trat(n,n1)
    t0 = @elapsed ts0(n,n1)
    t1 = @elapsed ts1(n,n1)
    t0/t1
end

function prat(n,n1)
    t0 = @elapsed ps0(n,n1)
    t1 = @elapsed ps1(n,n1)
    t0/t1
end

function trata(n,n1)
    t1 = @elapsed ts1(n,n1)
    t0 = @elapsed ts0(n,n1)
    t0/t1
end

function tratb(n,n1)
    t1 = @elapsed ts1(n,n1)
    t0 = @elapsed ts0(n,n1)
    t0/t1
end


text:
function klein(m::AbstractModel)

    #################
    # Linearization:
    #################
    Jac1 = Matrix{Float64}(jacobian(m))
    ##################################################################################
    # Klein Solution Method---apply generalized Schur decomposition a la Klein (2000)
    ##################################################################################

    # NK is number of predetermined variables
    NK = get_setting(m, :n_predetermined_variables)
    # n is number of variables (predet + non-predet)
	n = size(Jac1, 1)

    # A and B matrices
	A::Matrix{Float64} = Jac1[:, 1:n]
	B = -Jac1[:, n+1:2*n]

    # Apply generlaized Schur decomposition
    # A ≈ QZ[:Q]*QZ[:S]*QZ[:Z]'
    # B ≈ QZ[:Q]*QZ[:T]*QZ[:Z]'
	QZ = schur(A,B)

    # Reorder so that stable comes first
    alpha::Vector{complex(promote_type(eltype(A), eltype(B)))} = QZ.α
    beta::Vector{complex(promote_type(eltype(A), eltype(B)))} = QZ.β
	eigs = QZ.β ./ QZ.α #real(QZ.β ./ QZ.α)
	eigselect::AbstractArray{Bool} = abs.(eigs) .< 1 # returns false for NaN gen. eigenvalue which is correct here bc they are > 1
	ordschur!(QZ, eigselect)

    # Check that number of stable eigenvalues equals the number of predetermined state variables
	nk = sum(eigselect)
	if nk>NK
	    @warn "Equilibrium is locally indeterminate"
	elseif nk<NK
	    @warn "No local equilibrium exists"
	end

	U::Matrix{Float64} = QZ.Z'
	T::Matrix{Float64} = QZ.T
	S::Matrix{Float64} = QZ.S

    U11 = Matrix{Float64}(undef, NK, NK)
    U12 = Matrix{Float64}(undef, NK, NK-2)
    U21 = Matrix{Float64}(undef, NK-2, NK)
    U22 = Matrix{Float64}(undef, NK-2, NK-2)
    S11 = Matrix{Float64}(undef, NK, NK)
    T11 = Matrix{Float64}(undef, NK, NK)

    U11 = U[1:NK,1:NK]
	U12 = U[1:NK,NK+1:end]
	U21 = U[NK+1:end,1:NK]
	U22 = U[NK+1:end,NK+1:end]

	S11 = S[1:NK,1:NK]
	T11 = T[1:NK,1:NK]

    # Find minimum norm solution to U₂₁ + U₂₂*g_x = 0 (more numerically stable than -U₂₂⁻¹*U₂₁)
    gx_coef = Matrix{Float64}(undef, n-NK, NK)
	gx_coef = try
        -U22'*pinv(U22*U22')*U21
    catch ex
        if isa(ex, LinearAlgebra.LAPACKException)
            #@info "LAPACK exception thrown while computing pseudo inverse of U22*U22'"
            return gx_coef, Array{Float64, 2}(undef, NK, NK), -1
        else
            rethrow(ex)
        end
    end

    # Solve for h_x (in a more numerically stable way)
	S11invT11 = S11\T11;
	Ustuff = (U11 + U12*gx_coef);
	invterm = try
        pinv(eye(NK) + gx_coef' * gx_coef)
    catch ex
        if isa(ex, LinearAlgebra.LAPACKException)
            #@info "LAPACK exception thrown while computing pseudo inverse of eye(NK) + gx_coef'*gx_+coef"
            return gx_coef, Array{Float64, 2}(undef, NK, NK), -1
        else
            rethrow(ex)
        end
    end
    # hx_coef = Array{Float64, 2}(NK, NK)
	hx_coef = invterm*Ustuff'*S11invT11*Ustuff;

	# Ensure that hx and S11invT11 should have same eigenvalues
	# (eigst,valst) = eig(S11invT11);
	# (eighx,valhx) = eig(hx_coef);
	eigst = eigvals(S11invT11)
    eighx = eigvals(hx_coef)
    if abs(norm(eighx, Inf) - norm(eigst, Inf)) > 1e-4
		@warn "max abs eigenvalue of S11invT11 and hx are different!"
	end

	# next, want to represent policy functions in terms of meaningful things
	# gx_fval = Qy'*gx_coef*Qx
	# hx_fval = Qx'*hx_coef*Qx
    return gx_coef, hx_coef, 0
end

# Need an additional transition_equation function to properly stack the
# individual state and jump transition matrices/shock mapping matrices to
# a single state space for all of the model_states
function klein_transition_matrices(m::AbstractModel,
                                   TTT_state::Matrix{Float64}, TTT_jump::Matrix{Float64})
    TTT = zeros(n_model_states(m), n_model_states(m))

    # Loading mapping time t states to time t+1 states
    TTT[1:n_backward_looking_states(m), 1:n_backward_looking_states(m)] = TTT_state

    # Loading mapping time t jumps to time t+1 states
    TTT[1:n_backward_looking_states(m), n_backward_looking_states(m)+1:end] .= 0.

    # Loading mapping time t states to time t+1 jumps
    TTT[n_backward_looking_states(m)+1:end, 1:n_backward_looking_states(m)] = TTT_jump*TTT_state

    # Loading mapping time t jumps to time t+1 jumps
    TTT[n_backward_looking_states(m)+1:end, n_backward_looking_states(m)+1:end] .= 0.

    RRR = shock_loading(m, TTT_jump)

    return TTT, RRR
end


text:
psize = 5

srand(1)

# TODO: Links gets set up properly on instantiation.

context("population.jl") do
  landscapes = [
    NK.NKLandscape(4, 1),
    NK.NKqLandscape(4, 1, 2),
    NK.NKpLandscape(4, 1, 0.90),
  ]

  for ls = landscapes
    pr = rand(NK.Population, ls, psize)
    pz = zeros(NK.Population, ls, psize)

    context("$(ls)") do
      context("popsize(...)") do
        @fact NK.popsize(pr) --> psize
        @fact NK.popsize(pz) --> psize
      end

      context("popfits(...)") do
        @fact length(NK.popfits(pr)) --> psize
        @fact length(NK.popfits(pz)) --> psize
      end
    end
  end
end


text:
"""
$(TYPEDEF)

Implements the Mass Weighted Urn Design (MWUD) of [Zhao (2015)](https://doi.org/10.1016/j.cct.2015.06.008).

The MWU has conditional allocation probability

```math
P_{im} = \\frac{\\max{\\alpha w_m - n_{i-1,m} + (i - 1)w_m, 0}}{\\sum \\max{\\alpha w_m - n_{i-1,m} + (i - 1)w_m, 0}}
```

# Fields
$(TYPEDFIELDS)
"""
struct MassWeightedUrn{Tv<:AbstractVector,T<:Real} <: MultiArmRandomisationModel
    target::Tv
    α::T
    mass::Tv
    sequence::AbstractVector{Int}
    dist::AbstractVector{Int}
    MassWeightedUrn{Tv,T}(
        target::Tv,
        α::T,
        mass::Tv,
        sequence::AbstractVector{Int},
        dist::AbstractVector{Int},
    ) where {Tv<:AbstractVector,T<:Real} = new{Tv,T}(target, α, mass, sequence, dist)
end

"""
$(TYPEDSIGNATURES)
"""
function MassWeightedUrn(target::AbstractVector{T}, α::Real) where {T<:Real}
    if !isposvec(target)
        error("All values must be ≥ 0")
    end
    if !isprobvec(target)
        target = target ./ sum(target)
    end
    if α ≤ zero(typeof(α))
        error("α must be > 0")
    end
    target = collect(target)
    mass = α .* target
    MassWeightedUrn{typeof(target),typeof(α)}(
        target,
        α,
        mass,
        zeros(Int, 0),
        zeros(Int, length(target)),
    )
end

"""
$(TYPEDSIGNATURES)
"""
function MassWeightedUrn(target::AbstractVector{T}) where {T<:Real}
    if !isprobvec(target)
        # error("target must sum to 1")
        target = target ./ sum(target)
    end
    α = 2 * length(target)
    MassWeightedUrn(target, α)
end

Base.deepcopy(m::MassWeightedUrn) = MassWeightedUrn(deepcopy(m.target), deepcopy(m.α))
function show(io::IO, MWU::MassWeightedUrn)
    k = length(MWU.target)
    n = length(MWU.sequence)
    print(io, "MassWeightedUrn(balls = $k, mass = $(MWU.α), n = $n)")
end

mass(MWU::MassWeightedUrn) = MWU.mass
prob(MWU::MassWeightedUrn) =
    max.(MWU.α .* mass(MWU), 0.0) ./ sum(max.(MWU.α .* mass(MWU), 0.0))

"""
$(TYPEDSIGNATURES)
"""
function update!(MWU::MassWeightedUrn, y::Int)
    MWU.mass[y] += target(MWU)[y] - 1.0
    MWU.mass[1:end.!=y] += target(MWU)[1:end.!=y]
    push!(MWU.sequence, y)
    MWU.dist[y] += 1
    if !(sum(mass(MWU)) ≈ MWU.α)
        error("Mass no longer equal to α, something went wrong")
    end
    return nothing
end


text:
# clone a given ChemicalHyperGraph{Num} n times
# the symbolic variables will be renamed with a susbscript going from 1 to n
# note: set rename_weights to true to also rename parameters in the same way
function clone(chx::ChemicalHyperGraph{Num}, n::Int; rename_weights = false)

    # initialise an empty vector that will hold the cloned hypergraphs
    n == 1 ? (return deepcopy(chx)) : (hypergraphs = Vector{ChemicalHyperGraph{Num}}(undef, n))

    # iteration over clones
    for i in 1:n

        # define new vertices
        new_vs = subscript_vars(vertices(chx), i)
        v_d = Dict(old_v => new_v for (old_v, new_v) in zip(vertices(chx), new_vs)) # dictionary to keep track of changes

        # rename weights, if needed
        if rename_weights
            old_ws = rate.(hyperedges(chx))
            new_ws = Vector{Any}(undef, length(old_ws))
            iswsym = [(w isa Num && get_value(w) isa Symbolic) for w in old_ws] # check which weights are symbolic variables
            [new_ws[k] = (iswsym[k] ? subscript_var(old_ws[k], i) : old_ws[k]) for k in 1:length(old_ws)]
            w_d = Dict(old_w => new_w for (old_w, new_w) in zip(old_ws, new_ws)) # dictionary to keep track of changes
        end

        # define new hyperedges
        new_es = Vector{ChemicalHyperEdge{Num}}(undef, nhe(chx))
        for (j, e) in enumerate(hyperedges(chx))
            new_e = deepcopy(e)
            replace!(v -> isa(get_value(v), Symbolic) ? v_d[v] : v, src(new_e))
            replace!(v -> isa(get_value(v), Symbolic) ? v_d[v] : v, tgt(new_e))
            rename_weights && (new_e.weight = w_d[rate(new_e)])
            new_es[j] = new_e
        end

        hypergraphs[i] = ChemicalHyperGraph(new_vs, new_es)

    end

    hypergraphs

end


text:
module Types

using RBNF
using MLStyle
using RBNF: Token

export MainProgram, IfStmt, Opaque, Barrier, RegDecl, Include, GateDecl, Gate, Reset, Measure,
    Instruction, UGate, CXGate, Bit, Call, Neg, Add, Sub, Mul, Div, ASTNode

abstract type ASTNode end

struct MainProgram <: ASTNode
    version::VersionNumber
    prog::Vector{Any}

    MainProgram(version, prog) = new(version, Vector{Any}(prog))
end

struct IfStmt <: ASTNode
    left
    right
    body
end

struct Opaque <: ASTNode
    name
    cargs::Vector{Any}
    qargs::Vector{Any}

    function Opaque(name, cargs, qargs)
        new(name, _force_any(cargs), _force_any(qargs))
    end
end

struct Barrier <: ASTNode
    qargs::Vector{Any}

    function Barrier(qargs)
        new(_force_any(qargs))
    end
end

struct RegDecl <: ASTNode
    type
    name
    size
end

struct Include <: ASTNode
    file
end

struct GateDecl <: ASTNode
    name
    # we remove type annotations for now
    # due to JuliaLang/julia/issues/38091
    cargs::Vector{Any}
    qargs::Vector{Any}

    function GateDecl(name, cargs, qargs)
        new(name, _force_any(cargs), _force_any(qargs))
    end
end

struct Gate <: ASTNode
    decl::GateDecl
    body::Vector{Any}
end

struct Reset <: ASTNode
    qarg
end

struct Measure <: ASTNode
    qarg
    carg
end

struct Instruction <: ASTNode
    name::String
    cargs::Vector{Any}
    qargs::Vector{Any}

    function Instruction(name, cargs, qargs)
        new(name, _force_any(cargs), _force_any(qargs))
    end
end

struct UGate <: ASTNode
    z1
    y
    z2
    qarg
end

struct CXGate <: ASTNode
    ctrl
    qarg
end

struct Bit <: ASTNode
    name
    address
end

function Bit(name::String, address::Int)
    Bit(Token{:id}(name), Token{:int}(string(address)))
end

Bit(name::String) = Bit(Token{:id}(name), nothing)

struct Call <: ASTNode
    name::Symbol
    args
end

struct Neg <: ASTNode
    val
end

@deprecate Negative(val) Neg(val)
@deprecate FnExp(fn, arg) Call(fn, arg)

Base.show(io::IO, x::ASTNode) = print_qasm(io, x)

print_kw(io::IO, xs...) = printstyled(io, xs...; color=:light_blue)

function print_list(io::IO, list::Vector)
    for k in eachindex(list)
        print_qasm(io, list[k])

        if k != lastindex(list)
            print(io, ", ")
        end
    end
end

print_list(io::IO, x) = print_qasm(io, x)

print_qasm(ast) = print_qasm(stdout, ast)
print_qasm(io::IO) = x->print_qasm(io, x)
print_qasm(io::IO, ::Nothing) = nothing

print_qasm(io::IO, t::Token) = print(io, t.str)

function print_qasm(io::IO, t::Token{:reserved})
    print_kw(io, t.str)
end

function print_qasm(io::IO, stmt::RBNF.Token{:id})
    printstyled(io, stmt.str; color=:light_cyan)
end

function print_qasm(io::IO, stmt::RBNF.Token{:float64})
    printstyled(io, stmt.str; color=:green)
end

function print_qasm(io::IO, stmt::RBNF.Token{:int})
    printstyled(io, stmt.str; color=:green)
end

# NOTE:
# In order to preserve some line number
# we usually don't annote types to AST

# work around JuliaLang/julia/issues/38091
function _force_any(x)
    if x === nothing
        return Any[]
    else
        return Vector{Any}(x)
    end
end

function print_qasm(io::IO, x::MainProgram)
    printstyled(io, "OPENQASM "; bold=true)
    printstyled(io, x.version.major, ".", x.version.minor; color=:yellow)
    print(io, ";")
    println(io)

    for k in 1:length(x.prog)
        stmt = x.prog[k]
        print_qasm(io, stmt)

        # print extra line
        # when there is a gate decl
        if stmt isa Gate
            println(io)
        end

        if k != length(x.prog)
            println(io)
        end
    end
end
# nested inst list
function print_qasm(io::IO, stmts::Vector{Any})
    for each in stmts
        print_qasm(io, each)
    end
end

function print_qasm(io::IO, stmt::IfStmt)
    print_kw(io, "if ")
    print(io, "(")
    print_qasm(io, stmt.left)
    print(io, " == ")
    print_qasm(io, stmt.right)
    print(io, ") ")
    print_qasm(io, stmt.body)
end

function print_qasm(io::IO, stmt::Opaque)
    print_kw(io, "opaque ")
    if !isempty(stmt.cargs)
        print(io, "(")
        print_list(io, stmt.cargs)
        print(io, ") ")
    end
    print_list(io, stmt.qargs)
    print(io, ";")
end

function print_qasm(io::IO, stmt::Barrier)
    print_kw(io, "barrier ")
    print_list(io, stmt.qargs)
    print(io, ";")
end

function print_qasm(io::IO, stmt::RegDecl)
    print_qasm(io, stmt.type)
    print(io, " ")
    print_qasm(io, stmt.name)
    print(io, "[")
    print_qasm(io, stmt.size)
    print(io, "];")
end

function print_qasm(io::IO, stmt::Include)
    print_kw(io, "include ")
    print_qasm(io, stmt.file)
    print(io, ";")
end

function print_qasm(io::IO, stmt::GateDecl)
    print_kw(io, "gate ")
    print_qasm(io, stmt.name)

    if !isempty(stmt.cargs)
        print(io, "(")
        print_list(io, stmt.cargs)
        print(io, ")")
    end

    print(io, " ")
    print_list(io, stmt.qargs)
    print(io, " {")
end

function print_qasm(io::IO, stmt::Gate)
    print_qasm(io, stmt.decl)
    println(io)
    for k in 1:length(stmt.body)
        print(io, " "^2)
        print_qasm(io, stmt.body[k])
        println(io)
    end
    print(io, "}")
end

function print_qasm(io::IO, stmt::Reset)
    print_kw(io, "reset ")
    print_qasm(io, stmt.qarg)
end

function print_qasm(io::IO, stmt::Measure)
    print_kw(io, "measure ")
    print_qasm(io, stmt.qarg)
    print_kw(io, " -> ")
    print_qasm(io, stmt.carg)
    print(io, ";")
end

function print_qasm(io::IO, stmt::Instruction)
    printstyled(io, stmt.name; color=:light_magenta)

    if !isempty(stmt.cargs)
        print(io, "(")
        print_list(io, stmt.cargs)
        print(io, ")")
    end
    print(io, " ")
    print_list(io, stmt.qargs)
    print(io, ";")
end

function print_qasm(io::IO, stmt::UGate)
    print_kw(io, "U")
    print(io, "(")
    print_qasm(io, stmt.z1)
    print(io, ", ")
    print_qasm(io, stmt.y)
    print(io, ", ")
    print_qasm(io, stmt.z2)
    print(io, ") ")
    print_qasm(io, stmt.qarg)
    print(io, ";")
end

function print_qasm(io::IO, stmt::CXGate)
    print_kw(io, "CX ")
    print_qasm(io, stmt.ctrl)
    print(io, ", ")
    print_qasm(io, stmt.qarg)
    print(io, ";")
end

function print_qasm(io::IO, stmt::Bit)
    print_qasm(io, stmt.name)
    if stmt.address !== nothing
        print(io, "[")
        print_qasm(io, stmt.address)
        print(io, "]")
    end
end

function print_qasm(io::IO, stmt::Call)
    print_kw(io, stmt.name)
    print(io, "(")
    print_qasm(io, stmt.args)
    print(io, ")")
end

function print_qasm(io::IO, stmt::Neg)
    print(io, "-")
    print(io, "(")
    print_qasm(io, stmt.val)
    print(io, ")")
end

# exp
function print_qasm(io::IO, stmt::Tuple)
    if get(io, :parathesis, false)
        print(io, "(")
    end

    foreach(print_qasm(IOContext(io, :parathesis=>true)), stmt)

    if get(io, :parathesis, false)
        print(io, ")")
    end
end

@as_record MainProgram
@as_record IfStmt
@as_record Opaque
@as_record Barrier
@as_record RegDecl
@as_record Include
@as_record GateDecl
@as_record Gate
@as_record Reset
@as_record Measure
@as_record Instruction
@as_record UGate
@as_record CXGate
@as_record Bit
@as_record Call
@as_record Neg

end


text:
using MortIntl
using Test

@testset "MortIntl.jl" begin
    # Write your tests here.
end


text:
using Base.Test

include("raindrops.jl")

@testset "detect numbers" begin
    @testset "the sound for 1 is 1" begin
      @test raindrops(1) == "1"
    end
    @testset "2 to the power 3 does not make a raindrop sound as 3 is the exponent not the base" begin
      @test raindrops(8) == "8"
    end
    @testset "the sound for 52 is 52" begin
      @test raindrops(52) == "52"
    end
end

@testset "detect pling" begin
    @testset "the sound for 3 is Pling" begin
      @test raindrops(3) == "Pling"
    end
    @testset "the sound for 6 is Pling as it has a factor 3" begin
      @test raindrops(6) == "Pling"
    end
    @testset "the sound for 9 is Pling as it has a factor 3" begin
      @test raindrops(9) == "Pling"
    end
    @testset "the sound for 27 is Pling as it has a factor 3" begin
      @test raindrops(27) == "Pling"
    end
end

@testset "detect plang" begin
    @testset "the sound for 5 is Plang" begin
      @test raindrops(5) == "Plang"
    end
    @testset "the sound for 10 is Plang as it has a factor 5" begin
      @test raindrops(10) == "Plang"
    end
    @testset "the sound for 25 is Plang as it has a factor 5" begin
      @test raindrops(25) == "Plang"
    end
    @testset "the sound for 3125 is Plang as it has a factor 5" begin
      @test raindrops(3125) == "Plang"
    end
end

@testset "detect plong" begin
    @testset "the sound for 7 is Plong" begin
      @test raindrops(7) == "Plong"
    end
    @testset "the sound for 14 is Plong as it has a factor of 7" begin
      @test raindrops(14) == "Plong"
    end
    @testset "the sound for 49 is Plong as it has a factor 7" begin
      @test raindrops(49) == "Plong"
    end
end

@testset "detect multiple sounds" begin
    @testset "the sound for 15 is PlingPlang as it has factors 3 and 5" begin
      @test raindrops(15) == "PlingPlang"
    end
    @testset "the sound for 21 is PlingPlong as it has factors 3 and 7" begin
      @test raindrops(21) == "PlingPlong"
    end
    @testset "the sound for 35 is PlangPlong as it has factors 5 and 7" begin
      @test raindrops(35) == "PlangPlong"
    end
    @testset "the sound for 105 is PlingPlangPlong as it has factors 3, 5 and 7" begin
      @test raindrops(105) == "PlingPlangPlong"
    end
end


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct VirtualMachineScaleSetPublicIPAddressConfigurationProperties <: SwaggerModel
    idleTimeoutInMinutes::Any # spec type: Union{ Nothing, Int32 } # spec name: idleTimeoutInMinutes
    dnsSettings::Any # spec type: Union{ Nothing, VirtualMachineScaleSetPublicIPAddressConfigurationDnsSettings } # spec name: dnsSettings
    ipTags::Any # spec type: Union{ Nothing, Vector{VirtualMachineScaleSetIpTag} } # spec name: ipTags
    publicIPPrefix::Any # spec type: Union{ Nothing, SubResource } # spec name: publicIPPrefix
    publicIPAddressVersion::Any # spec type: Union{ Nothing, String } # spec name: publicIPAddressVersion

    function VirtualMachineScaleSetPublicIPAddressConfigurationProperties(;idleTimeoutInMinutes=nothing, dnsSettings=nothing, ipTags=nothing, publicIPPrefix=nothing, publicIPAddressVersion=nothing)
        o = new()
        validate_property(VirtualMachineScaleSetPublicIPAddressConfigurationProperties, Symbol("idleTimeoutInMinutes"), idleTimeoutInMinutes)
        setfield!(o, Symbol("idleTimeoutInMinutes"), idleTimeoutInMinutes)
        validate_property(VirtualMachineScaleSetPublicIPAddressConfigurationProperties, Symbol("dnsSettings"), dnsSettings)
        setfield!(o, Symbol("dnsSettings"), dnsSettings)
        validate_property(VirtualMachineScaleSetPublicIPAddressConfigurationProperties, Symbol("ipTags"), ipTags)
        setfield!(o, Symbol("ipTags"), ipTags)
        validate_property(VirtualMachineScaleSetPublicIPAddressConfigurationProperties, Symbol("publicIPPrefix"), publicIPPrefix)
        setfield!(o, Symbol("publicIPPrefix"), publicIPPrefix)
        validate_property(VirtualMachineScaleSetPublicIPAddressConfigurationProperties, Symbol("publicIPAddressVersion"), publicIPAddressVersion)
        setfield!(o, Symbol("publicIPAddressVersion"), publicIPAddressVersion)
        o
    end
end # type VirtualMachineScaleSetPublicIPAddressConfigurationProperties

const _property_map_VirtualMachineScaleSetPublicIPAddressConfigurationProperties = Dict{Symbol,Symbol}(Symbol("idleTimeoutInMinutes")=>Symbol("idleTimeoutInMinutes"), Symbol("dnsSettings")=>Symbol("dnsSettings"), Symbol("ipTags")=>Symbol("ipTags"), Symbol("publicIPPrefix")=>Symbol("publicIPPrefix"), Symbol("publicIPAddressVersion")=>Symbol("publicIPAddressVersion"))
const _property_types_VirtualMachineScaleSetPublicIPAddressConfigurationProperties = Dict{Symbol,String}(Symbol("idleTimeoutInMinutes")=>"Int32", Symbol("dnsSettings")=>"VirtualMachineScaleSetPublicIPAddressConfigurationDnsSettings", Symbol("ipTags")=>"Vector{VirtualMachineScaleSetIpTag}", Symbol("publicIPPrefix")=>"SubResource", Symbol("publicIPAddressVersion")=>"String")
Base.propertynames(::Type{ VirtualMachineScaleSetPublicIPAddressConfigurationProperties }) = collect(keys(_property_map_VirtualMachineScaleSetPublicIPAddressConfigurationProperties))
Swagger.property_type(::Type{ VirtualMachineScaleSetPublicIPAddressConfigurationProperties }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_VirtualMachineScaleSetPublicIPAddressConfigurationProperties[name]))}
Swagger.field_name(::Type{ VirtualMachineScaleSetPublicIPAddressConfigurationProperties }, property_name::Symbol) =  _property_map_VirtualMachineScaleSetPublicIPAddressConfigurationProperties[property_name]

const _allowed_VirtualMachineScaleSetPublicIPAddressConfigurationProperties_publicIPAddressVersion = ["IPv4", "IPv6"]

function check_required(o::VirtualMachineScaleSetPublicIPAddressConfigurationProperties)
    true
end

function validate_property(::Type{ VirtualMachineScaleSetPublicIPAddressConfigurationProperties }, name::Symbol, val)
    if name === Symbol("publicIPAddressVersion")
        Swagger.validate_param(name, "VirtualMachineScaleSetPublicIPAddressConfigurationProperties", :enum, val, _allowed_VirtualMachineScaleSetPublicIPAddressConfigurationProperties_publicIPAddressVersion)
    end
end


text:
abstract type AbstractRate end

# TODO: Add more methods here such as MultiRate, Explicit [can't reuse word]
Base.@kwdef struct IMEX{ℱ}
    method::ℱ
end

function IMEX()
    return IMEX(ARK2GiraldoKellyConstantinescu)
end

Base.@kwdef struct Explicit{ℳ, ℛ} <: AbstractRate
    model::ℳ
    rate::ℛ = 1
end

function Explicit(model::BalanceLaw; rate = 1)
    return Explicit(model = model, rate = rate)
end

function Explicit(model::SpaceDiscretization; rate = 1)
    return Explicit(model = model, rate = rate)
end

# perhaps "solver_method" instead of "method"?
Base.@kwdef struct Implicit{ℳ, ℛ, 𝒜} <: AbstractRate
    model::ℳ
    rate::ℛ = 1
    method::𝒜 = LinearBackwardEulerSolver(ManyColumnLU(); isadjustable = false)
end

function Implicit(
    model::BalanceLaw;
    rate = 1,
    adjustable = false,
    method = LinearBackwardEulerSolver(ManyColumnLU(), isadjustable = false),
)
    return Implicit(model = model, rate = rate, method = method)
end

function Implicit(
    model::SpaceDiscretization;
    rate = 1,
    adjustable = false,
    method = LinearBackwardEulerSolver(ManyColumnLU(), isadjustable = false),
)
    return Implicit(model = model, rate = rate, method = method)
end

# helper functions
function explicit(models::Tuple)
    explicit_models = []
    for model in models
        if model isa Explicit
            push!(explicit_models, model)
        end
    end
    return explicit_models
end
# extension
explicit(model) = explicit((model,))
explicit(model::BalanceLaw) = explicit(Explicit(model))

function implicit(models::Tuple)
    explicit_models = []
    for model in models
        if model isa Implicit
            push!(explicit_models, model)
        end
    end
    return explicit_models
end
# extension
implicit(model) = implicit(())

# Default to Explicit
# TODO: Extend whatever struct an ODE solver is
function construct_odesolver(method, rhs, state, Δt; t0 = 0, split_explicit_implicit = false)
    # put error checking here
    explicit_rhs = explicit(rhs)
    implicit_rhs = implicit(rhs)
    number_implicit = length(implicit_rhs)
    number_explicit = length(explicit_rhs)
    if (number_implicit != 0) | (number_explicit != 1)
        error_string = "Explicit methods require one (and only one) explicit model"
        error_string *= "\n for example, a tuple of the form  (Explicit(model),) or just (model,)"
        @error(error_string)
    end

    explicit_model = explicit_rhs[1]

    # Instantiate time stepping method
    odesolver = method(explicit_rhs.model, state; dt = Δt, t0 = t0, split_explicit_implicit = split_explicit_implicit)
    return odesolver
end

# IMEX
function construct_odesolver(timestepper::IMEX, rhs, state, Δt; t0 = 0, split_explicit_implicit = false)
    # put error checking here
    implicit_rhs = implicit(rhs)
    explicit_rhs = explicit(rhs)
    number_implicit = length(implicit_rhs)
    number_explicit = length(explicit_rhs)
    if (number_implicit != 1) | (number_explicit != 1)
        error_string = "IMEX methods require 1 implicit model and one explicit model"
        error_string *= "\n for example, a tuple of the form  (Explicit(model), Implicit(linear_model),)"
        @error(error_string)
    end
    explicit_rhs = explicit_rhs[1]
    implicit_rhs = implicit_rhs[1]

    # Instantiate time stepping method
    odesolver = timestepper.method(
        explicit_rhs.model,
        implicit_rhs.model,
        implicit_rhs.method,
        state;
        dt = Δt,
        t0 = t0,
        split_explicit_implicit = split_explicit_implicit,
    )
    return odesolver
end

max_stars_repo_name: CliMA/CouplerMachine

max_issues_repo_name: CliMA/CouplerMachine


text:
export iNALU

"""
    iNALU(in::Int, out::Int; initNAC=glorot_uniform, initG=glorot_uniform, ϵ=1e-7, ω=20)

Improved NALU that can process negative numbers by recovering the multiplication
sign. Implemented as suggested in: https://arxiv.org/abs/2003.07629
"""
struct iNALU{Tg<:AbstractMatrix,T<:Real}
    a_nac::NAC
    m_nac::NAC
    G::Tg
    ϵ::T
    ω::T
end

Flux.@functor iNALU

iNALU(a_nac::NAC, m_nac::NAC, G::AbstractMatrix) =
  iNALU(a_nac, m_nac, G{T}, T(1e-7), T(20)) where T

function iNALU(in::Int, out::Int;
              initNAC=glorot_uniform, initG=glorot_uniform, ϵ=1e-7, ω=20)
    a_nac = NAC(in, out, initW=initNAC, initM=initNAC)
    m_nac = NAC(in, out, initW=initNAC, initM=initNAC)
    G = initG(out, in)
    T = eltype(G)
    iNALU(a_nac,m_nac,G,T(ϵ),T(ω))
end

function Base.sign(nalu::iNALU, x::AbstractVector)
    W  = abs.(weights(nalu.m_nac))
    sm = reshape(sign.(x),1,:) .* W .+ 1 .- W
    s  = vec(prod(sm, dims=2))
end

function Base.sign(nalu::iNALU, x::AbstractMatrix)
    buf = Zygote.Buffer(x, size(nalu.G,1), size(x,2))
    for i in 1:size(x,2)
        buf[:,i] = sign(nalu, x[:,i])
    end
    copy(buf)
end

add(nalu::iNALU, x) = nalu.a_nac(x)
mult(nalu::iNALU, x) = exp.(min.(nalu.m_nac(log.(max.(abs.(x),nalu.ϵ))),nalu.ω))
gate(nalu::iNALU, x) = σ.(nalu.G*x)

function (nalu::iNALU)(x)
    a = add(nalu, x)
    m = mult(nalu, x)
    s = sign(nalu, x)
    g = gate(nalu, x)
    g .* a .+ (1.0 .- g) .* m .* s
end

Base.show(io::IO, l::iNALU) =
  print(io, "iNALU(in=$(size(l.G,2)), out=$(size(l.G,1)))")


text:
# """
#     ≪(μ,ν)

# # Absolute continuity

# A measure μ is _absolutely continuous_ with respect to ν, written μ ≪ ν, if
# ν(A)==0 implies μ(A)==0 for every ν-measurable set A.

# Less formally, suppose we have a set A with ν(A)==0. If μ(A)≠0, then there can
# be no way to "reweight" ν to get to μ. We can't make something from nothing.

# This "reweighting" is really a density function. If μ≪ν, then there is some
# function f that makes `μ == ∫(f,ν)` (see the help section for `∫`).

# We can get this f directly via the Radon-Nikodym derivative, `f == 𝒹(μ,ν)` (see
# the help section for `𝒹`).

# Note that `≪` is not a partial order, because it is not antisymmetric. That is
# to say, it's possible (in fact, common) to have two different measures `μ` and
# `ν` with `μ ≪ ν` and `ν ≪ μ`. A simple example of this is
# ```
# μ = Normal()
# ν = Lebesgue(ℝ)
# ```

# When `≪` holds in both directions, the measures μ and ν are _equivalent_,
# written `μ ≃ ν`. See the help section for `≃` for more information.
# """
# function ≪ end

# export ≃

# """
#     ≃(μ,ν)

# # Equivalence of Measure

# Measures μ and ν on the same space X are equivalent, written `μ ≃ ν`, if `μ ≪ ν`
# and `ν ≪ μ`. Note that this is often written `~` in the literature, but this is
# overloaded in probabilistic programming, so we use this alternate notation.

# Also note that equivalence is very different from equality. For two equivalent
# measures, the sets of non-zero measure will be identical, but what that measure
# is in each case can be very different.
# """
# function ≃(μ,ν)
#     return (μ≪ν && ν≪μ)
# end

# function ≪(μ, ν)
#     μ == ν && return true
#     representative(μ) ≪ representative(ν) && return true
#     return false
# end


text:
# This script shows stansummary ESS vakues vs. MCMCChains ESS values

using CmdStan, StatsPlots

ProjDir = @__DIR__
cd(ProjDir)

berstanmodel = "
data {
  int<lower=0> N;
  int<lower=0,upper=1> y[N];
}
parameters {
  real<lower=0,upper=1> theta;
}
model {
  theta ~ beta(1,1);
  for (n in 1:N)
    y[n] ~ bernoulli(theta);
}
"

data = Dict("y" => [0,1,0,0,0,0,0,0,0,1],
  "N" => 10)

stanmodel = Stanmodel(
   name = "berstanmodel", model = berstanmodel, nchains = 4,
   printsummary=false, tmpdir=mktempdir());

ess_array_mcmcchains = []
ess_array_cmdstan = []

for i in 1:100
  rc, chns, cnames = stan(stanmodel,
     data, summary=true, ProjDir);
  dfc = describe(chns)[1]
  append!(ess_array_mcmcchains, dfc[:theta, :ess])
  rs = read_summary(stanmodel)
  append!(ess_array_cmdstan, rs[:theta, :ess])
end

p = Array{Plots.Plot{Plots.GRBackend}}(undef, 1);
p[1] = plot(ess_array_cmdstan, lab="CmdStan", xlim=(0, 170))
p[1] = plot!(ess_array_mcmcchains, lab="MCMCChains")
plot(p..., layout=(1,1))
savefig("ess_comparison_plot.pdf")


text:
module TermInterface

using ..Util
using Base.Meta

# TODO document this interface for extending for custom types
# Interface for Expr. Implement these methods for your own type to
# Use it instead of Expr in egraphs!
gethead(e::Expr) = isexpr(e, :call) ? e.args[1] : e.head
getargs(e::Expr) = @view e.args[(isexpr(e, :call) ? 2 : 1):end]
istree(e::Expr) = true
getmetadata(e::Expr) = (iscall=isexpr(e, :call),)
arity(e::Expr) = length(getargs(e)) # optional
preprocess(e::Expr) = cleanast(e)

# Fallback implementation for all other types
gethead(e) = e
getargs(e) = []
istree(a) = false
getmetadata(e) = nothing
arity(e) = 0 # optional
preprocess(e) = e

export gethead
export getargs
export istree
export getmetadata
export preprocess
export arity
end


text:
# # RadialBasisFunctionModels
#
# [![Stable](https://img.shields.io/badge/docs-stable-blue.svg)](https://manuelbb-upb.github.io/RadialBasisFunctionModels.jl/stable)
# [![Dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://manuelbb-upb.github.io/RadialBasisFunctionModels.jl/dev)
# [![Build Status](https://github.com/manuelbb-upb/RadialBasisFunctionModels.jl/workflows/CI/badge.svg)](https://github.com/manuelbb-upb/RadialBasisFunctionModels.jl/actions)
# [![Coverage](https://codecov.io/gh/manuelbb-upb/RadialBasisFunctionModels.jl/branch/master/graph/badge.svg)](https://codecov.io/gh/manuelbb-upb/RadialBasisFunctionModels.jl)
#

# # Description
# This package provides Radial Basis Function (RBF) models with polynomial tails.
# RBF models are a special case of kernel machines that can interpolate high-dimensional
# and nonlinear data.

# # Usage Examples

# First load the `RadialBasisFunctionModels` package.
using RadialBasisFunctionModels

# We also use `Test` to validate the results and `BenchmarkTools` for comparisons. #jl
using Test #jl

# ## Interpolating RBF Model

# ### One dimensional data
# The main type `RBFModel` uses vectors internally, but we can easily
# interpolate 1-dimensional data.
# Assume, e.g., we want to interpolate ``f:ℝ → ℝ, f(x) = x^2``:
f = x -> x^2
# Define 5 training sites `X` and evaluate to get `Y`
X = collect( LinRange(-4,4,5) )
Y = f.(X)

# Initialize the `RadialFunction` to use for the RBF model:
φ = Multiquadric()

# Construct an interpolating model with linear polynomial tail:
rbf = RBFModel( X, Y, φ, 1)

# We can evaluate `rbf` at the data points;
# By default, vectors are returned.
Z = rbf.(X)
# Now
Z isa Vector #md
# and
Z[1] isa AbstractVector{<:Real} #md

@test typeof(Z[1]) <: AbstractVector{<:Real} #jl
@test length(Z[1]) == 1     #jl
# The results should be close to the data labels `Y`, i.e., `Z[1] ≈ Y[1]` etc.
@test all( isapprox(Z[i][1], Y[i]; atol = 1e-10) for i = 1 : length(Z) ) #jl

# `X` contains Floats, but we can pass them to `rbf`.
# Usually you have feature vectors and they are always supported:
@test rbf( [ X[1], ] ) == Z[1]

# For 1 dimensional labels we can actually disable the vector output:
rbf_scalar = RBFInterpolationModel( X, Y, φ, 1; vector_output = false)
Z_scalar = rbf_scalar.( X )

@test(#jl
Z_scalar isa Vector{Float64}
)#jl
@test( #jl
all( Z_scalar[i] == Z[i][1] for i = 1 : length(Z) )
) #jl

# The data precision of the training data is preserved when determining the model
# coefficients.
# Accordingly, the return type precision is also at least that of the training data.
X_f0 = Float32.(X)
Y_f0 = f.(X_f0)
rbf_f0 = RBFInterpolationModel( X_f0, Y_f0, φ, 1)
@test(#jl
rbf_f0.(X_f0) isa Vector{Vector{Float32}}
)#jl

# If you are using statically sized arrays, they work too!
# You can provide a vector of statically sized arrays or, if
# you have only few centers ( number of variables × number of centers <= 100),
# provide a statically sized vector of statically sized vectors to maybe profit
# from faster matrix multiplications when evaluating:
using StaticArrays
features = [ @SVector(rand(3)) for i = 1 : 5 ]
labels = [ @SVector(rand(3)) for i = 1 : 5 ]
centers = SVector{5}(features)
rbf_sized = RBFModel( features, labels; centers )

# Now, the model uses sized matrices internally.
# For most input vectors a `SizedVector` would be returned.
# But there is a "type guarding" function for static arrays so that output has the same
# array type (by conversion, if necessary):
x_vec = rand(3)
x_s = SVector{3}(x_vec)
x_m = MVector{3}(x_vec)
x_sized = SizedVector{3}(x_vec)

@test(#jl
rbf_sized( x_vec ) isa Vector
)#jl
@test(#jl
rbf_sized( x_s ) isa SVector
)#jl
@test(#jl
rbf_sized( x_m ) isa MVector
)#jl
@test(#jl
rbf_sized( x_sized ) isa SizedVector
)#jl

# ### Using Kernel Names

# Instead of initializing `RadialFunction`s beforehand,
# their names can be used. Currently supported are:
# ```
# :gaussian, :multiquadric, :inv_multiquadric, :cubic, :thin_plate_spline
# ```

# You can do
RBFModel(features, labels, :gaussian)
# or, to specify kernel arguments:
RBFModel(features, labels, :multiquadric, [1.0, 1//2])

# ## Machines

# There is an MLJ wrapper for the RBFInterpolationModel, exported as `RBFInterpolator`.
# It can be used like other regressors and takes the kernel name as a symbol (and kernel arguments as a vector).

using MLJBase
X,y = @load_boston

r = RBFInterpolator(; kernel_name = :multiquadric )
R = machine(r, X, y)

MLJBase.fit!(R)
MLJBase.predict(R, X)

# You can do similar things (for vector valued data) with the `RBFMachineWithKernel`:
X = [ rand(2) for i = 1 : 10 ]
Y = [ rand(2) for i = 1 : 10 ]

R = RBFMachine(;features = X, labels = Y, kernel_name = :gaussian )
R isa RBFMachineWithKernel
RadialBasisFunctionModels.fit!(R)
@test(#jl
R( X[1] ) ≈ Y[1]
)#jl

# Such a machine can be initialized empty and data can be added:
R = RBFMachine()
add_data!(R, X, Y)


text:
__precompile__()
module BarycentricLagrangeInterpolation

using Measurements, Unitful, PyCall
import Measurements: value, uncertainty
import Unitful: AbstractQuantity
using PyCall
@pyimport numpy as np
@pyimport import matplotlib.pyplot as plt

#const nNodes = 14

# TODO
#const n =

end  # module BarycentricLagrangeInterpolation

max_stars_repo_name: LaGuer/barycentricLagrangeInterpolation

max_issues_repo_name: LaGuer/barycentricLagrangeInterpolation

max_forks_repo_name: LaGuer/barycentricLagrangeInterpolation


text:
@unit € "€" EuroSign 1.0EUR true
@unit USdollar "US\$" USDollarSign 1.0USD true
@unit CAdollar "CA\$" CanadianDollarSign 1.0CAD true
@unit Real "R\$" BrazilianRealSign 1.0BRL true
@unit £ "£" PoundSterlingSign 1.0GBP true


text:
#=
  From http://mc-stan.org/workshops/vanderbilt2016/carp-1.pdf
  Slide 5ff
  """
  Laplace's data on live births in Paris 1745-1770:
     sex        live births
     ----------------------
     female      241 945
     male        251 527

  * Question 1 (Estimation):
    What is the birth rate of boys vs. girls?
  * Question 2 (Event Probability)
    Is a boy more likely to be born than a girl?

    Bayes (1763) set up the 'Bayesian' model  Laplace (1781, 1786) solved for the posterior
  ...
  (Answers:)
  * Q1: θis 99% certain to lie in(0.508,0.512)
  * Q2: Laplace 'morally certain' boys more prevalent
  """

  See Stan model ~/stan/laplace_births.stan (see slide 10)
  """
                    mean se_mean     sd   2.5%    25%    50%    75%  97.5%  n_eff   Rhat
  theta           0.51  2.4e-5 6.8e-4   0.51   0.51   0.51   0.51   0.51  793.0    1.0
  theta_gt_half    1.0     nan    0.0    1.0    1.0    1.0    1.0    1.0    nan    nan
  """

  Summary Statistics
     parameters          mean        std   naive_se      mcse        ess      rhat   ess_per_sec
         Symbol       Float64    Float64    Float64   Float64    Float64   Float64       Float64

          theta        0.5096     0.0007     0.0000    0.0001    80.1135    1.0098       42.9333
  theta_gt_half        1.0000     0.0000     0.0000    0.0000        NaN       NaN           NaN
      male_post   251491.8328   492.9858     4.9299   28.3048   162.6618    1.0060       87.1714

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

          theta        0.5083        0.5092        0.5095        0.5101        0.5111
  theta_gt_half        1.0000        1.0000        1.0000        1.0000        1.0000
      male_post   250533.9750   251162.0000   251485.0000   251822.2500   252475.0750

  Cf ~/webppl/laplace_births.jl

=#

using Turing
include("jl_utils.jl")

@model function laplace_births(male=251527)
    female = 241945
    theta ~ Beta(2,2)
    # Note that this works i.e. male is both
    # in LHS and RHS.
    male ~ Binomial(male + female, theta)

    theta_gt_half ~ Dirac(theta > 0.5)

    # Posterior predictive of male
    male_post ~ Binomial(male + female, theta)

end

model = laplace_births()

num_chns = 4
# chns = sample(model, Prior(), 100_000)
# chns = sample(model, MH(), 100_000)
# chns = sample(model, PG(15), 10_000)
chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 100_000)

display(chns)

max_stars_repo_name: tias/hakank

max_issues_repo_name: tias/hakank

max_forks_repo_name: tias/hakank


text:
using Documenter, MonteCarloMeasurements

makedocs(sitename="MonteCarloMeasurements Documentation", doctest = false, modules=[MonteCarloMeasurements]) # Due to lots of plots, this will just have to be run on my local machine

deploydocs(
    deps   = Deps.pip("pygments", "mkdocs", "python-markdown-math", "mkdocs-cinder"),
    repo = "github.com/baggepinnen/MonteCarloMeasurements.jl.git"
)


text:
"""
Fitness objects are used to provide an empirical means of assessing the
relative quality of one potential solution to another.
"""
module fitness
importall common
using utility, core, distance
export assign, uses, maximise, Fitness, FitnessScheme

"""
Base type used by all fitness values.
"""
abstract Fitness

"""
Base type used by all fitness schemes.
"""
abstract FitnessScheme

# Load each of the fitness schemes.
include("fitness/_scalar.jl")
include("fitness/_pareto.jl")
include("fitness/_aggregate.jl")
include("fitness/_belegundu.jl")
include("fitness/_goldberg.jl")
include("fitness/_moga.jl")

"""
Responsible for generating a fitness value according to a given scheme.
"""
assign(s::FitnessScheme) =
  error("No 'fitness' function defined for this fitness scheme.")

"""
Returns the type of object by this fitness scheme to store fitness values.
"""
uses(s::FitnessScheme) =
  error("No 'uses' function defined for this fitness scheme.")

"""
Returns a flag indicating whether or not fitness values produced by a given
scheme should be maximised or minimised.
"""
maximise(::FitnessScheme) =
  error("No 'maximise' function defined for this fitness scheme.")
end


text:
using ThreeBodyDecay
using Test
using Parameters

@testset "Infererred type" begin
    tbs = ThreeBodySystem( 2.0, 1.0, 1.5; m0=6.0,
        two_js = ThreeBodySpins(1, 0, 0; two_h0=1))  # 1/2+ 0- 0- 1/2+

    dc = DecayChainLS(3, σ->1/(4.1^2-σ-0.1im); two_s=3, parity='-', Ps=['+','-','-','+'], tbs=tbs)

    @unpack σs, two_λs = randomPoint(tbs)

    @inferred Complex{Float64} amplitude(σs, two_λs, dc)

    # using InteractiveUtils
    # @code_warntype amplitude(σs, two_λs, dc)
end

max_stars_repo_name: mmikhasenko/ThreeBodyDecay

max_issues_repo_name: mmikhasenko/ThreeBodyDecay

max_forks_repo_name: mmikhasenko/ThreeBodyDecay


text:
include("../apm.jl")

y = apm_solve("2nd_order",4)

print(y)


text:
import PlutoSliderServer
PlutoSliderServer.run_directory("."; config_toml_path="./PlutoDeployment.toml")


text:
module GeneOntology

export
    # term
    Term,
    cellular_component, biological_process, molecular_function,
    isobsolete, @go_str,

    # typedef
    Typedef,

    # annotation
    AnnotationRecord,

    # parser
    OBOParser, eachterm, eachtypedef,
    AnnotaionParser, eachannot,

    # graph
    GOGraph, is_a, parents, children

include("term.jl")
include("typedef.jl")
include("annotation.jl")
include("parser.jl")
include("graph.jl")

end # GeneOntology module


text:
struct SparseValueIterationSolver <: Solver
    max_iterations::Int64
    belres::Float64 # the Bellman Residual
    include_Q::Bool
    verbose::Bool
    init_util::Vector{Float64}
end

function SparseValueIterationSolver(;max_iterations=500,
                                    belres::Float64=1e-3,
                                    include_Q::Bool=true,
                                    verbose::Bool=false,
                                    init_util::Vector{Float64}=Vector{Float64}(undef, 0))
    return SparseValueIterationSolver(max_iterations, belres, include_Q, verbose, init_util)
end

@POMDP_require solve(solver::SparseValueIterationSolver, mdp::MDP) begin
    P = typeof(mdp)
    S = statetype(P)
    A = actiontype(P)
    @req discount(::P)
    @subreq ordered_states(mdp)
    @subreq ordered_actions(mdp)
    @req transition(::P,::S,::A)
    @req reward(::P,::S,::A,::S)
    @req stateindex(::P,::S)
    @req actionindex(::P, ::A)
    @req actions(::P, ::S)
    as = actions(mdp)
    ss = states(mdp)
    @req length(::typeof(ss))
    @req length(::typeof(as))
    a = first(as)
    s = first(ss)
    dist = transition(mdp, s, a)
    D = typeof(dist)
    @req support(::D)
    @req pdf(::D,::S)
    @subreq SparseTabularMDP(mdp)
end

function qvalue!(m::Union{MDP,POMDP}, transition_A_S_S2, reward_S_A::AbstractMatrix{F}, value_S::AbstractVector{F}, out_qvals_S_A) where {F}
    @assert size(out_qvals_S_A) == (length(states(m)), length(actions(m)))
    for a in 1:length(actions(m))
        out_qvals_S_A[:, a] = view(reward_S_A, :, a) + discount(m) * transition_A_S_S2[a] * value_S
    end
end

function solve(solver::SparseValueIterationSolver, mdp::SparseTabularMDP)
    nS = length(states(mdp))
    nA = length(actions(mdp))
    if isempty(solver.init_util)
        v_S = zeros(nS)
    else
        @assert length(solver.init_util) == nS "Input utility dimension mismatch"
        v_S = solver.init_util
    end

    transition_A_S_S2 = transition_matrices(mdp)
    reward_S_A = reward_matrix(mdp)
    qvals_S_A = zeros(nS, nA)
    maxchanges_T = zeros(solver.max_iterations)

    total_time = 0.0
    for i in 1:solver.max_iterations
        iter_time = @elapsed begin
            qvalue!(mdp, transition_A_S_S2, reward_S_A, v_S, qvals_S_A)
            new_v_S = dropdims(maximum(qvals_S_A, dims=2), dims=2)
            @assert size(v_S) == size(new_v_S)
            maxchanges_T[i] = maximum(abs.(new_v_S .- v_S))
            v_S = new_v_S
        end
        total_time += iter_time
        if solver.verbose
            @info "residual: $(maxchanges_T[i]), time: $(iter_time), total time: $(total_time) " i
        end
        maxchanges_T[i] < solver.belres ? break : nothing
    end
    qvalue!(mdp, transition_A_S_S2, reward_S_A, v_S, qvals_S_A)
    # Rounding to avoid floating point error noise
    policy_S = dropdims(getindex.(argmax(round.(qvals_S_A, digits=20), dims=2), 2), dims=2)
    if solver.include_Q
        policy = ValueIterationPolicy(mdp, qvals_S_A, v_S, policy_S)
    else
        policy = ValueIterationPolicy(mdp, utility=v_S, policy=policy_S, include_Q=false)
    end
    return policy
end

function solve(solver::SparseValueIterationSolver, mdp::MDP)
    p = solve(solver, SparseTabularMDP(mdp))
    return ValueIterationPolicy(p.qmat, p.util, p.policy, ordered_actions(mdp), p.include_Q, mdp)
end

function solve(::SparseValueIterationSolver, ::POMDP)
    throw("""
           ValueIterationError: `solve(::SparseValueIterationSolver, ::POMDP)` is not supported,
          `SparseValueIterationSolver` supports MDP models only, look at QMDP.jl for a POMDP solver that assumes full observability.
           If you still wish to use the transition and reward from your POMDP model you can use the `UnderlyingMDP` wrapper from POMDPModelTools.jl as follows:
           ```
           solver = ValueIterationSolver()
           mdp = UnderlyingMDP(pomdp)
           solve(solver, mdp)
           ```
           """)
end


text:
using Revise
using ADCME
using ADSeismic
using PyPlot
using DelimitedFiles
matplotlib.use("agg")

param = AcousticPropagatorParams(NX=150, NY=150, NSTEP=1000, DELTAT=1e-4,  DELTAX=1.0, DELTAY=1.0,
                                 vp_ref = 3000.0, Rcoef=0.001,
                                 USE_PML_XMIN=true, USE_PML_XMAX=true, USE_PML_YMIN=true, USE_PML_YMAX=true)

rc = Ricker(param, 15.0, 100.0, 1e10)
srci = [param.NX ÷ 2]
srcj = [param.NY ÷ 2]
srcv = reshape(rc, :, 1)
src = AcousticSource(srci, srcj, srcv)

c = 3000.0*ones(param.NX+2, param.NY+2)
model = AcousticPropagatorSolver(param, src, c)

sess = Session(); init(sess)
u = run(sess, model.u)

p = visualize_wavefield(u, param)
saveanim(p, "acoustic-wavefield.gif")


text:
# Load Packages
using MPI
using CLIMA
using CLIMA.Mesh.Topologies
using CLIMA.Mesh.Grids
using CLIMA.DGmethods
using CLIMA.DGmethods.NumericalFluxes
using CLIMA.MPIStateArrays
using CLIMA.LowStorageRungeKuttaMethod
using CLIMA.ODESolvers
using CLIMA.GenericCallbacks
using CLIMA.Atmos
using CLIMA.VariableTemplates
using CLIMA.MoistThermodynamics
using CLIMA.PlanetParameters
using LinearAlgebra
using StaticArrays
using Logging, Printf, Dates
using CLIMA.VTK
using Random
using CLIMA.Atmos: vars_state, vars_aux

const ArrayType = CLIMA.array_type()

if !@isdefined integration_testing
  const integration_testing =
    parse(Bool, lowercase(get(ENV,"JULIA_CLIMA_INTEGRATION_TESTING","false")))
end

# -------------- Problem constants ------------------- #
const xmin      = 0
const ymin      = 0
const zmin      = 0
const xmax      = 1000
const ymax      = 400
const zmax      = 1000
const Ne        = (30,5,30)
const polynomialorder = 4
const dim       = 3
const Δx        = (xmax-xmin)/(Ne[1]*polynomialorder+1)
const Δy        = (ymax-ymin)/(Ne[2]*polynomialorder+1)
const Δz        = (zmax-zmin)/(Ne[3]*polynomialorder+1)
const Δ         = cbrt(Δx * Δy * Δz)
const dt        = 0.005
const timeend   = 100
const T_bot     = 299
const T_lapse   = -0.01
const T_top     = T_bot + T_lapse*zmax
const C_smag    = 0.18
# ------------- Initial condition function ----------- #
function initialise_rayleigh_benard!(state::Vars, aux::Vars, (x1,x2,x3), t)
  FT            = eltype(state)
  R_gas::FT     = R_d
  c_p::FT       = cp_d
  c_v::FT       = cv_d
  γ::FT         = c_p / c_v
  p0::FT        = MSLP
  δT            = sinpi(6*x3/(zmax-zmin)) * cospi(6*x3/(zmax-zmin))
  δw            = sinpi(6*x3/(zmax-zmin)) * cospi(6*x3/(zmax-zmin))
  ΔT            = T_lapse * x3 + δT
  T             = T_bot + ΔT
  P             = p0*(T/T_bot)^(grav/R_gas/T_lapse)
  ρ             = P / (R_gas * T)
  ρu, ρv, ρw    = FT(0) , FT(0) , ρ * δw
  E_int         = ρ * c_v * (T-T_0)
  E_pot         = ρ * grav * x3
  E_kin         = ρ * FT(1/2) * δw^2
  ρe_tot        = E_int + E_pot + E_kin
  state.ρ       = ρ
  state.ρu      = SVector(ρu, ρv, ρw)
  state.ρe      = ρe_tot
  state.moisture.ρq_tot = FT(0)
end
# --------------- Driver definition ------------------ #
function run(mpicomm,
             topl, dim, Ne, polynomialorder,
             timeend, FT, dt, model)
  # -------------- Define grid ----------------------------------- #
  grid = DiscontinuousSpectralElementGrid(topl,
                                          FloatType = FT,
                                          DeviceArray = ArrayType,
                                          polynomialorder = polynomialorder
                                           )
  # -------------- Define model ---------------------------------- #
  model = model
  # -------------- Define dgbalancelaw --------------------------- #
  dg = DGModel(model,
               grid,
               Rusanov(),
               CentralNumericalFluxDiffusive(),
               CentralGradPenalty())

  Q = init_ode_state(dg, FT(0))

  lsrk = LSRK54CarpenterKennedy(dg, Q; dt = dt, t0 = 0)

  eng0 = norm(Q)
  @info @sprintf """Starting
  norm(Q₀) = %.16e
  ArrayType = %s
  FloatType = %s""" eng0 ArrayType FT

  # Set up the information callback (output field dump is via vtk callback: see cbinfo)
  starttime = Ref(now())
  cbinfo = GenericCallbacks.EveryXWallTimeSeconds(60, mpicomm) do (s=false)
    if s
      starttime[] = now()
    else
      energy = norm(Q)
      @info @sprintf("""Update
                     simtime = %.16e
                     runtime = %s
                     norm(Q) = %.16e""", ODESolvers.gettime(lsrk),
                     Dates.format(convert(Dates.DateTime,
                                          Dates.now()-starttime[]),
                                  Dates.dateformat"HH:MM:SS"),
                     energy)
    end
  end

  step = [0]
  cbvtk = GenericCallbacks.EveryXSimulationSteps(3000)  do (init=false)
    mkpath("./vtk-rb/")
      outprefix = @sprintf("./vtk-rb/RB_%dD_mpirank%04d_step%04d", dim,
                           MPI.Comm_rank(mpicomm), step[1])
      @debug "doing VTK output" outprefix
      writevtk(outprefix, Q, dg, flattenednames(vars_state(model,FT)), dg.auxstate, flattenednames(vars_aux(model,FT)))
      step[1] += 1
      nothing
  end

  solve!(Q, lsrk; timeend=timeend, callbacks=(cbinfo,cbvtk))
  # End of the simulation information
  engf = norm(Q)
  Qe = init_ode_state(dg, FT(timeend))
  engfe = norm(Qe)
  errf = euclidean_distance(Q, Qe)
  @info @sprintf """Finished
  norm(Q)                 = %.16e
  norm(Q) / norm(Q₀)      = %.16e
  norm(Q) - norm(Q₀)      = %.16e
  norm(Q - Qe)            = %.16e
  norm(Q - Qe) / norm(Qe) = %.16e
  """ engf engf/eng0 engf-eng0 errf errf / engfe
engf/eng0
end
# --------------- Test block / Loggers ------------------ #
using Test
let
  CLIMA.init()
  mpicomm = MPI.COMM_WORLD
  ll = uppercase(get(ENV, "JULIA_LOG_LEVEL", "INFO"))
  loglevel = ll == "DEBUG" ? Logging.Debug :
    ll == "WARN"  ? Logging.Warn  :
    ll == "ERROR" ? Logging.Error : Logging.Info
  logger_stream = MPI.Comm_rank(mpicomm) == 0 ? stderr : devnull
  global_logger(ConsoleLogger(logger_stream, loglevel))

  for FT in (Float32,)
    SGSmodels = (AnisoMinDiss{FT}(1), Vreman{FT}(C_smag), SmagorinskyLilly{FT}(C_smag))
    Expected = (FT(9.9859344959259033e-01),FT(1.0038942098617554e+00),FT(1.0027571916580200e+00))
    for ii=1:length(SGSmodels)
      model = AtmosModel(FlatOrientation(),
                         NoReferenceState(),
                         SGSmodels[ii],
                         EquilMoist(),
                         NoRadiation(),
                         NoSubsidence{FT}(),
                         Gravity(),
                         RayleighBenardBC{FT}(T_bot,T_top),
                         initialise_rayleigh_benard!)
      brickrange = (range(FT(xmin); length=Ne[1]+1, stop=xmax),
                    range(FT(ymin); length=Ne[2]+1, stop=ymax),
                    range(FT(zmin); length=Ne[3]+1, stop=zmax))
      topl = StackedBrickTopology(mpicomm, brickrange, periodicity = (true, true, false), boundary=((0,0),(0,0),(1,2)))
      engf_eng0 = run(mpicomm,
                      topl, dim, Ne, polynomialorder,
                      timeend, FT, dt, model)
      @test engf_eng0 ≈ Expected[ii]
    end
  end
end

#nothing

max_stars_repo_name: emassoud/CLIMA

max_issues_repo_name: emassoud/CLIMA

max_forks_repo_name: emassoud/CLIMA