text:
# This file is part of the TaylorIntegration.jl package; MIT licensed

using TaylorIntegration, DiffEqBase
using Test
using LinearAlgebra: norm
using InteractiveUtils: methodswith
using Elliptic
using Base.Threads

@testset "Testing `taylorize.jl`" begin

    # Constants for the integrations
    local _order = 20
    local _abstol = 1.0e-20
    local t0 = 0.0
    local tf = 1000.0

    # Scalar integration
    @testset "Scalar case: xdot(x, p, t) = b-x^2" begin
        b1 = 3.0
        @taylorize xdot1(x, p, t) = b1-x^2
        @test (@isdefined xdot1)

        x0 = 1.0
        tv1, xv1 = taylorinteg( xdot1, x0, t0, tf, _order, _abstol, maxsteps=1000,
            parse_eqs=false)
        tv1p, xv1p = taylorinteg( xdot1, x0, t0, tf, _order, _abstol, maxsteps=1000)

        @test length(tv1) == length(tv1p)
        @test iszero( norm(tv1-tv1p, Inf) )
        @test iszero( norm(xv1-xv1p, Inf) )

        # Now using `local` constants
        @taylorize xdot2(x, p, t) = (local b2 = 3; b2-x^2)
        @test (@isdefined xdot2)

        tv2, xv2 = taylorinteg( xdot2, x0, t0, tf, _order, _abstol, maxsteps=1000,
            parse_eqs=false)
        tv2p, xv2p = taylorinteg( xdot2, x0, t0, tf, _order, _abstol, maxsteps=1000)

        @test length(tv2) == length(tv2p)
        @test iszero( norm(tv2-tv2p, Inf) )
        @test iszero( norm(xv2-xv2p, Inf) )

        # Passing a parameter
        @taylorize xdot3(x, p, t) = p-x^2
        @test (@isdefined xdot3)

        tv3, xv3 = taylorinteg( xdot3, x0, t0, tf, _order, _abstol, b1, maxsteps=1000,
            parse_eqs=false)
        tv3p, xv3p = taylorinteg( xdot3, x0, t0, tf, _order, _abstol, b1, maxsteps=1000)

        @test length(tv3) == length(tv3p)
        @test iszero( norm(tv3-tv3p, Inf) )
        @test iszero( norm(xv3-xv3p, Inf) )

        # Comparing integrations
        @test length(tv1p) == length(tv2p) == length(tv3p)
        @test iszero( norm(tv1p-tv2p, Inf) )
        @test iszero( norm(tv1p-tv3p, Inf) )
        @test iszero( norm(xv1p-xv2p, Inf) )
        @test iszero( norm(xv1p-xv3p, Inf) )

        xv4 = taylorinteg( xdot2, x0, t0:0.5:tf, _order, _abstol, maxsteps=1000,
            parse_eqs=false)
        xv4p = taylorinteg( xdot2, x0, t0:0.5:tf, _order, _abstol, maxsteps=1000)
        @test iszero( norm(xv4-xv4p, Inf) )

        # Compare to exact solution
        exact_sol(t, b, x0) = sqrt(b)*((sqrt(b)+x0)-(sqrt(b)-x0)*exp(-2sqrt(b)*t)) /
            ((sqrt(b)+x0)+(sqrt(b)-x0)*exp(-2sqrt(b)*t))
        @test norm(xv1p[end] - exact_sol(tv1p[end], b1, x0), Inf) < 1.0e-15

        # Check that the parsed `jetcoeffs` produces the correct series in `x`
        tT = t0 + Taylor1(_order)
        xT = x0 + zero(tT)
        TaylorIntegration.__jetcoeffs!(Val(true), xdot1, tT, xT, nothing)
        @test xT ≈ exact_sol(tT, b1, x0)

        xT = x0 + zero(tT)
        TaylorIntegration.__jetcoeffs!(Val(true), xdot2, tT, xT, nothing)
        @test xT ≈ exact_sol(tT, 3.0, x0)

        xT = x0 + zero(tT)
        TaylorIntegration.__jetcoeffs!(Val(true), xdot2, tT, xT, b1)
        @test xT ≈ exact_sol(tT, b1, x0)

        # The macro returns a (parsed) jetcoeffs! function which yields an error and
        # therefore it runs with the default jetcoeffs! method. A warning is issued.
        @taylorize function xdot2_err(x, p, t)
            b2 = 3
            return b2-x^2
        end

        @test (@isdefined xdot2_err)
        @test !isempty(methodswith(Val{xdot2_err}, TaylorIntegration.jetcoeffs!))
        tv2e, xv2e = taylorinteg( xdot2_err, x0, t0, tf, _order, _abstol, maxsteps=1000)
        @test length(tv2) == length(tv2e)
        @test iszero( norm(tv2-tv2e, Inf) )
        @test iszero( norm(xv2-xv2e, Inf) )
        @test_throws UndefVarError TaylorIntegration.__jetcoeffs!(Val(true), xdot2_err, tT, qT, similar(qT), similar(qT), [1.0])
    end

    @testset "Scalar case: xdot(x, p, t) = -10" begin
        xdot1(x, p, t) = -10 + zero(t) # `zero(t)` is needed; cf #20
        @taylorize xdot1_parsed(x, p, t) = -10# `zero(t)` can be avoided here

        @test (@isdefined xdot1_parsed)
        tv1, xv1   = taylorinteg( xdot1, 10, 1, 20.0, _order, _abstol)
        tv1p, xv1p = taylorinteg( xdot1_parsed, 10.0, 1.0, 20.0, _order, _abstol)

        @test length(tv1) == length(tv1p)
        @test iszero( norm(tv1-tv1p, Inf) )
        @test iszero( norm(xv1-xv1p, Inf) )

        # Now using `local`
        @taylorize function xdot2(x, p, t)
            local ggrav = -10 + zero(t)  # zero(t) is needed for the parse_eqs=false case
            tmp = ggrav  # needed to avoid an error when parsing
        end

        @test (@isdefined xdot2)

        tv2, xv2   = taylorinteg( xdot2, 10.0, 1.0, 20.0, _order, _abstol, parse_eqs=false)
        tv2p, xv2p = taylorinteg( xdot2, 10.0, 1, 20.0, _order, _abstol)

        @test length(tv2) == length(tv2p)
        @test iszero( norm(tv2-tv2p, Inf) )
        @test iszero( norm(xv2-xv2p, Inf) )

        # Passing a parameter
        @taylorize xdot3(x, p, t) = p + zero(t) # `zero(t)` can be avoided here

        @test (@isdefined xdot3)

        tv3, xv3   = taylorinteg( xdot3, 10.0, 1.0, 20.0, _order, _abstol, -10,
            parse_eqs=false)
        tv3p, xv3p = taylorinteg( xdot3, 10, 1.0, 20.0, _order, _abstol, -10)

        @test length(tv3) == length(tv3p)
        @test iszero( norm(tv3-tv3p, Inf) )
        @test iszero( norm(xv3-xv3p, Inf) )

        # Comparing integrations
        @test length(tv1p) == length(tv2p) == length(tv3p)
        @test iszero( norm(tv1p-tv2p, Inf) )
        @test iszero( norm(tv1p-tv3p, Inf) )
        @test iszero( norm(xv1p-xv2p, Inf) )
        @test iszero( norm(xv1p-xv3p, Inf) )

        # Compare to exact solution
        exact_sol(t, g, x0) = x0 + g*(t-1.0)
        @test norm(xv1p[end] - exact_sol(20, -10, 10), Inf) < 10eps(exact_sol(20, -10, 10))

        # Check that the parsed `jetcoeffs` produces the correct series in `x`
        tT = 1.0 + Taylor1(_order)
        xT = 10.0 + zero(tT)
        TaylorIntegration.__jetcoeffs!(Val(true), xdot1_parsed, tT, xT, nothing)
        @test xT ≈ exact_sol(tT, -10, 10)

        xT = 10.0 + zero(tT)
        TaylorIntegration.__jetcoeffs!(Val(true), xdot2, tT, xT, nothing)
        @test xT ≈ exact_sol(tT, -10, 10)

        xT = 10.0 + zero(tT)
        TaylorIntegration.__jetcoeffs!(Val(true), xdot3, tT, xT, -10)
        @test xT ≈ exact_sol(tT, -10, 10)
    end

    # Pendulum integrtf = 100.0
    @testset "Integration of the pendulum and DiffEqs interface" begin
        @taylorize function pendulum!(dx::Array{T,1}, x::Array{T,1}, p, t) where {T}
            dx[1] = x[2]
            dx[2] = -sin( x[1] )
            nothing
        end

        @test (@isdefined pendulum!)
        q0 = [pi-0.001, 0.0]
        tv2, xv2 = taylorinteg(pendulum!, q0, t0, tf, _order, _abstol, parse_eqs=false,
            maxsteps=5000)
        tv2p, xv2p = taylorinteg(pendulum!, q0, t0, tf, _order, _abstol, maxsteps=5000)

        @test length(tv2) == length(tv2p)
        @test iszero( norm(tv2-tv2p, Inf) )
        @test iszero( norm(xv2-xv2p, Inf) )

        prob = ODEProblem(pendulum!, q0, (t0, tf), nothing) # no parameters
        sol1 = solve(prob, TaylorMethod(_order), abstol=_abstol, parse_eqs=true)
        sol2 = solve(prob, TaylorMethod(_order), abstol=_abstol, parse_eqs=false)
        sol3 = solve(prob, TaylorMethod(_order), abstol=_abstol)

        @test sol1.t == sol2.t == sol3.t == tv2p
        @test sol1.u[end] == sol2.u[end] == sol3.u[end] == xv2p[end,1:2]

        # Check that the parsed `jetcoeffs` produces the correct series in `x`;
        # we check that it does not throws an error
        tT = 0.0 + Taylor1(_order)
        qT = q0 .+ zero(tT)
        TaylorIntegration.__jetcoeffs!(Val(true), pendulum!, tT, qT, similar(qT), similar(qT), nothing)
    end

    # Complex dependent variables
    @testset "Complex dependent variable" begin
        cc = complex(0.0,1.0)
        @taylorize eqscmplx(x, p, t) = cc*x

        @test (@isdefined eqscmplx)
        cx0 = complex(1.0, 0.0)
        tv1, xv1 = taylorinteg(eqscmplx, cx0, t0, tf, _order, _abstol, maxsteps=1500,
            parse_eqs=false)
        tv1p, xv1p = taylorinteg(eqscmplx, cx0, t0, tf, _order, _abstol, maxsteps=1500)

        @test length(tv1) == length(tv1p)
        @test iszero( norm(tv1-tv1p, Inf) )
        @test iszero( norm(xv1-xv1p, Inf) )

        # Using `local` for the constant value
        @taylorize eqscmplx2(x, p, t) = (local cc1 = Complex(0.0,1.0); cc1*x)

        @test (@isdefined eqscmplx2)
        tv2, xv2 = taylorinteg(eqscmplx2, cx0, t0, tf, _order, _abstol, maxsteps=1500,
            parse_eqs=false)
        tv2p, xv2p = taylorinteg(eqscmplx2, cx0, t0, tf, _order, _abstol, maxsteps=1500)

        @test length(tv2) == length(tv2p)
        @test iszero( norm(tv2-tv2p, Inf) )
        @test iszero( norm(xv2-xv2p, Inf) )

        # Passing a parameter
        @taylorize eqscmplx3(x, p, t) = p*x
        @test (@isdefined eqscmplx3)
        tv3, xv3 = taylorinteg(eqscmplx3, cx0, t0, tf, _order, _abstol, cc, maxsteps=1500,
            parse_eqs=false)
        tv3p, xv3p = taylorinteg(eqscmplx3, cx0, t0, tf, _order, _abstol, cc, maxsteps=1500)

        @test length(tv3) == length(tv3p)
        @test iszero( norm(tv3-tv3p, Inf) )
        @test iszero( norm(xv3-xv3p, Inf) )

        # Comparing integrations
        @test length(tv1p) == length(tv2p) == length(tv3p)
        @test iszero( norm(tv1p-tv2p, Inf) )
        @test iszero( norm(tv1p-tv3p, Inf) )
        @test iszero( norm(xv1p-xv2p, Inf) )
        @test iszero( norm(xv1p-xv3p, Inf) )

        # Compare to exact solution
        exact_sol(t, p, x0) = x0*exp(p*t)
        nn = norm.(abs2.(xv1p)-abs2.(exact_sol.(tv1p, cc, cx0)), Inf)
        @test maximum(nn) < 1.0e-14

        eqs1(z, p, t) = -z
        eqs2(z, p, t) = im*z
        @taylorize function eqs3!(Dz, z, p, t)
            local _eltype_z_ = eltype(z)
            tmp = Array{_eltype_z_}(undef, 2)
            tmp[1] = eqs1(z[1], p, t)
            tmp[2] = eqs2(z[2], p, t)
            Dz[1] = tmp[1]
            Dz[2] = tmp[2]
            nothing
        end
        @test (@isdefined eqs3!)
        z0 = complex(0.0, 1.0)
        zz0 = [z0, z0]
        ts = 0.0:pi:2pi
        zsol = taylorinteg(eqs3!, zz0, ts, _order, _abstol, parse_eqs=false, maxsteps=10)
        zsolp = taylorinteg(eqs3!, zz0, ts, _order, _abstol, maxsteps=10)
        @test length(tv3) == length(tv3p)
        @test iszero( norm(tv3-tv3p, Inf) )
        @test iszero( norm(xv3-xv3p, Inf) )

        tT = t0 + Taylor1(_order)
        zT = zz0 .+ zero(tT)
        TaylorIntegration.__jetcoeffs!(Val(true), eqs3!, tT, zT, similar(zT), similar(zT), nothing)
        @test zT[1] == exact_sol(tT, -1, z0)
        @test zT[2] == exact_sol(tT, z0, z0)
    end

    @testset "Time-dependent integration (with and without `local` vars)" begin
        @taylorize function integ_cos1(x, p, t)
            y = cos(t)
            return y
        end
        @taylorize function integ_cos2(x, p, t)
            local y = cos(t)  # allows to calculate directly `cos(t)` *once*
            yy = y            # needed to avoid an error
            return yy
        end

        @test (@isdefined integ_cos1)
        @test (@isdefined integ_cos2)

        tv11, xv11 = taylorinteg(integ_cos1, 0.0, 0.0, pi, _order, _abstol, parse_eqs=false)
        tv12, xv12 = taylorinteg(integ_cos1, 0.0, 0.0, pi, _order, _abstol)
        @test length(tv11) == length(tv12)
        @test iszero( norm(tv11-tv12, Inf) )
        @test iszero( norm(xv11-xv12, Inf) )

        tv21, xv21 = taylorinteg(integ_cos2, 0.0, 0.0, pi, _order, _abstol, parse_eqs=false)
        tv22, xv22 = taylorinteg(integ_cos2, 0.0, 0.0, pi, _order, _abstol)
        @test length(tv21) == length(tv22)
        @test iszero( norm(tv21-tv22, Inf) )
        @test iszero( norm(xv21-xv22, Inf) )

        @test iszero( norm(tv12-tv22, Inf) )
        @test iszero( norm(xv12-xv22, Inf) )

        # Compare to exact solution
        @test norm(xv11[end] - sin(tv11[end]), Inf) < 1.0e-15
    end

    @testset "Time-dependent integration (vectorial case)" begin
        @taylorize function integ_vec(dx, x, p, t)
            dx[1] = cos(t)
            dx[2] = -sin(t)
            return dx
        end

        @test (@isdefined integ_vec)
        x0 = [0.0, 1.0]
        tv11, xv11 = taylorinteg(integ_vec, x0, 0.0, pi, _order, _abstol, parse_eqs=false)
        tv12, xv12 = taylorinteg(integ_vec, x0, 0.0, pi, _order, _abstol)
        @test length(tv11) == length(tv12)
        @test iszero( norm(tv11-tv12, Inf) )
        @test iszero( norm(xv11-xv12, Inf) )

        # Compare to exact solution
        @test norm(xv12[end, 1] - sin(tv12[end]), Inf) < 1.0e-15
        @test norm(xv12[end, 2] - cos(tv12[end]), Inf) < 1.0e-15
    end

    # Simple harmonic oscillator
    @testset "Simple harmonic oscillator" begin
        @taylorize function harm_osc!(dx, x, p, t)
            local ω = p[1]
            local ω2 = ω^2
            dx[1] = x[2]
            dx[2] = - (ω2 * x[1])
            return nothing
        end

        @test (@isdefined harm_osc!)
        q0 = [1.0, 0.0]
        p = [2.0]
        tv1, xv1 = taylorinteg(harm_osc!, q0, t0, tf, _order, _abstol,
            p, maxsteps=1000, parse_eqs=false)
        tv1p, xv1p = taylorinteg(harm_osc!, q0, t0, tf, _order, _abstol,
            p, maxsteps=1000)

        @test length(tv1) == length(tv1p)
        @test iszero( norm(tv1-tv1p, Inf) )
        @test iszero( norm(xv1-xv1p, Inf) )

        # Comparing to exact solution
        @test norm(xv1p[end, 1] - cos(p[1]*tv1p[end]), Inf) < 1.0e-12
        @test norm(xv1p[end, 2] + p[1]*sin(p[1]*tv1p[end]), Inf) < 3.0e-12

        # Check that the parsed `jetcoeffs` produces the correct series in `x`
        tT = t0 + Taylor1(_order)
        qT = q0 .+ zero(tT)
        TaylorIntegration.__jetcoeffs!(Val(true), harm_osc!, tT, qT, similar(qT), similar(qT), [1.0])
        @test qT[1] ≈ cos(tT)
        @test qT[2] ≈ -sin(tT)

        # The macro returns a (parsed) jetcoeffs! function which yields an error and
        # therefore it runs with the default jetcoeffs! methdo. A warining is issued.
        @taylorize function harm_osc_error!(dx, x, p, t)
            local ω = p[1]
            dx[1] = x[2]
            dx[2] = - (ω^2 * x[1])
            return nothing
        end

        @test !isempty(methodswith(Val{harm_osc_error!}, TaylorIntegration.jetcoeffs!))
        tv2e, xv2e = taylorinteg(harm_osc_error!, q0, t0, tf, _order, _abstol,
            p, maxsteps=1000, parse_eqs=true)
        @test length(tv1) == length(tv1p)
        @test iszero( norm(tv1-tv2e, Inf) )
        @test iszero( norm(xv1-xv2e, Inf) )
        @test_throws MethodError TaylorIntegration.__jetcoeffs!(Val(true), harm_osc_error!, tT, qT, similar(qT), similar(qT), [1.0])
    end

    @testset "Multiple pendula" begin
        NN = 3
        nnrange = 1:3
        @taylorize function multpendula1!(dx, x, p, t)
            for i in p[2]
                dx[i] = x[p[1]+i]
                dx[i+p[1]] = -sin( x[i] )
            end
            return nothing
        end

        @test (@isdefined multpendula1!)
        q0 = [pi-0.001, 0.0, pi-0.001, 0.0,  pi-0.001, 0.0]
        tv1, xv1 = taylorinteg(multpendula1!, q0, t0, tf, _order, _abstol,
            [NN, nnrange], maxsteps=1000, parse_eqs=false)
        tv1p, xv1p = taylorinteg(multpendula1!, q0, t0, tf, _order, _abstol,
            [NN, nnrange], maxsteps=1000)

        @test length(tv1) == length(tv1p)
        @test iszero( norm(tv1-tv1p, Inf) )
        @test iszero( norm(xv1-xv1p, Inf) )

        @taylorize function multpendula2!(dx, x, p, t)
            local NN = 3
            local nnrange = 1:NN
            for i in nnrange
                dx[i] = x[NN+i]
                dx[i+NN] = -sin( x[i] )
            end
            return nothing
        end

        @test (@isdefined multpendula2!)
        tv2, xv2 = taylorinteg(multpendula2!, q0, t0, tf, _order, _abstol,
            maxsteps=1000, parse_eqs=false)
        tv2p, xv2p = taylorinteg(multpendula2!, q0, t0, tf, _order, _abstol,
            maxsteps=1000)

        @test length(tv2) == length(tv2p)
        @test iszero( norm(tv2-tv2p, Inf) )
        @test iszero( norm(xv2-xv2p, Inf) )

        @taylorize function multpendula3!(dx, x, p, t)
            nn, nnrng = p    # `local` is not needed to reassign `p`; internally is treated as local)
            for i in nnrng
                dx[i] = x[nn+i]
                dx[i+nn] = -sin( x[i] )
            end
            return nothing
        end

        @test (@isdefined multpendula3!)
        tv3, xv3 = taylorinteg(multpendula3!, q0, t0, tf, _order, _abstol,
            [NN, nnrange], maxsteps=1000, parse_eqs=false)
        tv3p, xv3p = taylorinteg(multpendula3!, q0, t0, tf, _order, _abstol,
            [NN, nnrange], maxsteps=1000)

        # Comparing integrations
        @test length(tv1) == length(tv2) == length(tv3)
        @test iszero( norm(tv1-tv2, Inf) )
        @test iszero( norm(tv1-tv3, Inf) )
        @test iszero( norm(xv1-xv2, Inf) )
        @test iszero( norm(xv1-xv3, Inf) )

        # Check that the parsed `jetcoeffs` produces the correct series in `x`;
        # we check that it does not throws an error
        tT = 0.0 + Taylor1(_order)
        qT = q0 .+ zero(tT)
        TaylorIntegration.__jetcoeffs!(Val(true), multpendula1!, tT, qT, similar(qT), similar(qT), (NN, nnrange))

        qT = q0 .+ zero(tT)
        TaylorIntegration.__jetcoeffs!(Val(true), multpendula2!, tT, qT, similar(qT), similar(qT), nothing)

        qT = q0 .+ zero(tT)
        TaylorIntegration.__jetcoeffs!(Val(true), multpendula3!, tT, qT, similar(qT), similar(qT), [NN, nnrange])
    end

    # Kepler problem
    # Redefining some constants
    local _order = 28
    local tf = 2π*100.0
    @testset "Kepler problem (using `^`)" begin
        @taylorize function kepler1!(dq, q, p, t)
            local μ = p
            r_p3d2 = (q[1]^2+q[2]^2)^1.5

            dq[1] = q[3]
            dq[2] = q[4]
            dq[3] = μ * q[1]/r_p3d2
            dq[4] = μ * q[2]/r_p3d2

            return nothing
        end

        @taylorize function kepler2!(dq, q, p, t)
            nn, μ = p
            r2 = zero(q[1])
            for i = 1:nn
                r2_aux = r2 + q[i]^2
                r2 = r2_aux
            end
            r_p3d2 = r2^(3/2)
            for j = 1:nn
                dq[j] = q[nn+j]
                dq[nn+j] = μ*q[j]/r_p3d2
            end

            nothing
        end

        @test (@isdefined kepler1!)
        @test (@isdefined kepler2!)
        pars = (2, -1.0)
        q0 = [0.2, 0.0, 0.0, 3.0]
        tv1, xv1 = taylorinteg(kepler1!, q0, t0, tf, _order, _abstol, -1.0,
            maxsteps=500000, parse_eqs=false)
        tv1p, xv1p = taylorinteg(kepler1!, q0, t0, tf, _order, _abstol, -1.0,
            maxsteps=500000)

        tv6p, xv6p = taylorinteg(kepler2!, q0, t0, tf, _order, _abstol, pars,
            maxsteps=500000, parse_eqs=false)
        tv7p, xv7p = taylorinteg(kepler2!, q0, t0, tf, _order, _abstol, pars,
            maxsteps=500000)

        @test length(tv1) == length(tv1p) == length(tv6p)
        @test iszero( norm(tv1-tv1p, Inf) )
        @test iszero( norm(xv1-xv1p, Inf) )
        @test iszero( norm(tv1-tv6p, Inf) )
        @test iszero( norm(xv1-xv6p, Inf) )
        @test iszero( norm(tv7p-tv6p, Inf) )
        @test iszero( norm(xv7p-xv6p, Inf) )

        @taylorize function kepler3!(dq, q, p, t)
            local mμ = -1.0
            r_p3d2 = (q[1]^2+q[2]^2)^1.5

            dq[1] = q[3]
            dq[2] = q[4]
            dq[3] = mμ*q[1]/r_p3d2
            dq[4] = mμ*q[2]/r_p3d2

            return nothing
        end

        @taylorize function kepler4!(dq, q, p, t)
            local mμ = -1.0
            local NN = 2
            r2 = zero(q[1])
            for i = 1:NN
                r2_aux = r2 + q[i]^2
                r2 = r2_aux
            end
            r_p3d2 = r2^(3/2)
            for j = 1:NN
                dq[j] = q[NN+j]
                dq[NN+j] = mμ*q[j]/r_p3d2
            end

            nothing
        end

        @test (@isdefined kepler3!)
        @test (@isdefined kepler4!)
        tv3, xv3 = taylorinteg(kepler3!, q0, t0, tf, _order, _abstol,
            maxsteps=500000, parse_eqs=false)
        tv3p, xv3p = taylorinteg(kepler3!, q0, t0, tf, _order, _abstol,
            maxsteps=500000)

        tv4, xv4 = taylorinteg(kepler4!, q0, t0, tf, _order, _abstol,
            maxsteps=500000, parse_eqs=false)
        tv4p, xv4p = taylorinteg(kepler4!, q0, t0, tf, _order, _abstol,
            maxsteps=500000)

        @test length(tv3) == length(tv3p) == length(tv4)
        @test iszero( norm(tv3-tv3p, Inf) )
        @test iszero( norm(xv3-xv3p, Inf) )
        @test iszero( norm(tv3-tv4, Inf) )
        @test iszero( norm(xv3-xv4, Inf) )
        @test iszero( norm(tv4p-tv4, Inf) )
        @test iszero( norm(xv4p-xv4, Inf) )

        # Comparing both integrations
        @test iszero( norm(tv1p-tv3p, Inf) )
        @test iszero( norm(xv1p-xv3p, Inf) )

        # Check that the parsed `jetcoeffs` produces the correct series in `x`;
        # we check that it does not throws an error
        tT = 0.0 + Taylor1(_order)
        qT = q0 .+ zero(tT)
        TaylorIntegration.__jetcoeffs!(Val(true), kepler1!, tT, qT, similar(qT), similar(qT), -1.0)

        qT = q0 .+ zero(tT)
        TaylorIntegration.__jetcoeffs!(Val(true), kepler2!, tT, qT, similar(qT), similar(qT), pars)

        qT = q0 .+ zero(tT)
        TaylorIntegration.__jetcoeffs!(Val(true), kepler3!, tT, qT, similar(qT), similar(qT), nothing)

        qT = q0 .+ zero(tT)
        TaylorIntegration.__jetcoeffs!(Val(true), kepler4!, tT, qT, similar(qT), similar(qT), nothing)
    end

    @testset "Kepler problem (using `sqrt`)" begin
        @taylorize function kepler1!(dq, q, p, t)
            local μ = -1.0
            r = sqrt(q[1]^2+q[2]^2)
            r_p3d2 = r^3

            dq[1] = q[3]
            dq[2] = q[4]
            dq[3] = μ * q[1] / r_p3d2
            dq[4] = μ * q[2] / r_p3d2

            return nothing
        end

        @taylorize function kepler2!(dq, q, p, t)
            local NN = 2
            local μ = p
            r2 = zero(q[1])
            for i = 1:NN
                r2_aux = r2 + q[i]^2
                r2 = r2_aux
            end
            r = sqrt(r2)
            r_p3d2 = r^3
            for j = 1:NN
                dq[j] = q[NN+j]
                dq[NN+j] =  μ * q[j] / r_p3d2
            end

            nothing
        end

        @test (@isdefined kepler1!)
        @test (@isdefined kepler2!)
        q0 = [0.2, 0.0, 0.0, 3.0]
        tv1, xv1 = taylorinteg(kepler1!, q0, t0, tf, _order, _abstol,
            maxsteps=500000, parse_eqs=false)
        tv1p, xv1p = taylorinteg(kepler1!, q0, t0, tf, _order, _abstol,
            maxsteps=500000)

        tv2, xv2 = taylorinteg(kepler2!, q0, t0, tf, _order, _abstol, -1.0,
            maxsteps=500000, parse_eqs=false)
        tv2p, xv2p = taylorinteg(kepler2!, q0, t0, tf, _order, _abstol, -1.0,
            maxsteps=500000)

        @test length(tv1) == length(tv1p) == length(tv2)
        @test iszero( norm(tv1-tv1p, Inf) )
        @test iszero( norm(xv1-xv1p, Inf) )
        @test iszero( norm(tv1-tv2, Inf) )
        @test iszero( norm(xv1-xv2, Inf) )
        @test iszero( norm(tv2p-tv2, Inf) )
        @test iszero( norm(xv2p-xv2, Inf) )

        @taylorize function kepler3!(dq, q, p, t)
            local μ = p
            x, y, px, py = q
            r = sqrt(x^2+y^2)
            r_p3d2 = r^3

            dq[1] = px
            dq[2] = py
            dq[3] = μ * x / r_p3d2
            dq[4] = μ * y / r_p3d2

            return nothing
        end

        @taylorize function kepler4!(dq, q, p, t)
            local NN = 2
            local μ = -1.0
            r2 = zero(q[1])
            for i = 1:NN
                r2_aux = r2 + q[i]^2
                r2 = r2_aux
            end
            r = sqrt(r2)
            r_p3d2 = r^3
            for j = 1:NN
                dq[j] = q[NN+j]
                dq[NN+j] = μ * q[j] / r_p3d2
            end

            nothing
        end

        @test (@isdefined kepler3!)
        @test (@isdefined kepler4!)
        tv3, xv3 = taylorinteg(kepler3!, q0, t0, tf, _order, _abstol, -1.0,
            maxsteps=500000, parse_eqs=false)
        tv3p, xv3p = taylorinteg(kepler3!, q0, t0, tf, _order, _abstol, -1.0,
            maxsteps=500000)

        tv4, xv4 = taylorinteg(kepler4!, q0, t0, tf, _order, _abstol,
            maxsteps=500000, parse_eqs=false)
        tv4p, xv4p = taylorinteg(kepler4!, q0, t0, tf, _order, _abstol,
            maxsteps=500000)

        @test length(tv3) == length(tv3p) == length(tv4)
        @test iszero( norm(tv3-tv3p, Inf) )
        @test iszero( norm(xv3-xv3p, Inf) )
        @test iszero( norm(tv3-tv4, Inf) )
        @test iszero( norm(xv3-xv4, Inf) )
        @test iszero( norm(tv4p-tv4, Inf) )
        @test iszero( norm(xv4p-xv4, Inf) )

        # Comparing both integrations
        @test iszero( norm(tv1p-tv3p, Inf) )
        @test iszero( norm(xv1p-xv3p, Inf) )

        # Check that the parsed `jetcoeffs` produces the correct series in `x`;
        # we check that it does not throws an error
        tT = 0.0 + Taylor1(_order)
        qT = q0 .+ zero(tT)
        TaylorIntegration.__jetcoeffs!(Val(true), kepler1!, tT, qT, similar(qT), similar(qT), nothing)

        qT = q0 .+ zero(tT)
        TaylorIntegration.__jetcoeffs!(Val(true), kepler2!, tT, qT, similar(qT), similar(qT), -1.0)

        qT = q0 .+ zero(tT)
        TaylorIntegration.__jetcoeffs!(Val(true), kepler3!, tT, qT, similar(qT), similar(qT), -1.0)

        qT = q0 .+ zero(tT)
        TaylorIntegration.__jetcoeffs!(Val(true), kepler4!, tT, qT, similar(qT), similar(qT), nothing)
    end

    local tf = 20.0
    @testset "Lyapunov spectrum and `@taylorize`" begin
        #Lorenz system parameters
        params = [16.0, 45.92, 4.0]

        #Lorenz system ODE:
        @taylorize function lorenz1!(dq, q, p, t)
            x, y, z = q
            local σ, ρ, β = p
            dq[1] = σ*(y-x)
            dq[2] = x*(ρ-z) - y
            dq[3] = x*y - β*z
            nothing
        end

        #Lorenz system Jacobian (in-place):
        function lorenz1_jac!(jac, q, p, t)
            x, y, z = q
            σ, ρ, β = p
            jac[1,1] = -σ+zero(x)
            jac[2,1] = ρ-z
            jac[3,1] = y
            jac[1,2] = σ+zero(x)
            jac[2,2] = -one(y)
            jac[3,2] = x
            jac[1,3] = zero(z)
            jac[2,3] = -x
            jac[3,3] = -β+zero(x)
            nothing
        end

        q0 = [19.0, 20.0, 50.0] #the initial condition
        xi = set_variables("δ", order=1, numvars=length(q0))

        tv1, lv1, xv1 = lyap_taylorinteg(lorenz1!, q0, t0, tf, _order, _abstol,
            params, maxsteps=2000, parse_eqs=false);

        tv1p, lv1p, xv1p = lyap_taylorinteg(lorenz1!, q0, t0, tf, _order, _abstol,
            params, maxsteps=2000);

        @test tv1 == tv1p
        @test lv1 == lv1p
        @test xv1 == xv1p

        tv2, lv2, xv2 = lyap_taylorinteg(lorenz1!, q0, t0, tf, _order, _abstol,
            params, lorenz1_jac!, maxsteps=2000, parse_eqs=false);

        tv2p, lv2p, xv2p = lyap_taylorinteg(lorenz1!, q0, t0, tf, _order, _abstol,
            params, lorenz1_jac!, maxsteps=2000,  parse_eqs=true);

        @test tv2 == tv2p
        @test lv2 == lv2p
        @test xv2 == xv2p

        # Comparing both integrations (lorenz1)
        @test tv1 == tv2
        @test lv1 == lv2
        @test xv1 == xv2

        #Lorenz system ODE:
        @taylorize function lorenz2!(dx, x, p, t)
            #Lorenz system parameters
            local σ = 16.0
            local β = 4.0
            local ρ = 45.92

            dx[1] = σ*(x[2]-x[1])
            dx[2] = x[1]*(ρ-x[3])-x[2]
            dx[3] = x[1]*x[2]-β*x[3]
            nothing
        end

        #Lorenz system Jacobian (in-place):
        function lorenz2_jac!(jac, x, p, t)
            #Lorenz system parameters
            local σ = 16.0
            local β = 4.0
            local ρ = 45.92

            jac[1,1] = -σ+zero(x[1])
            jac[2,1] = ρ-x[3]
            jac[3,1] = x[2]
            jac[1,2] = σ+zero(x[1])
            jac[2,2] = -1.0+zero(x[1])
            jac[3,2] = x[1]
            jac[1,3] = zero(x[1])
            jac[2,3] = -x[1]
            jac[3,3] = -β+zero(x[1])
            nothing
        end

        tv3, lv3, xv3 = lyap_taylorinteg(lorenz2!, q0, t0, tf, _order, _abstol,
            maxsteps=2000, parse_eqs=false);

        tv3p, lv3p, xv3p = lyap_taylorinteg(lorenz2!, q0, t0, tf, _order, _abstol,
            maxsteps=2000);

        @test tv3 == tv3p
        @test lv3 == lv3p
        @test xv3 == xv3p

        tv4, lv4, xv4 = lyap_taylorinteg(lorenz2!, q0, t0, tf, _order, _abstol,
            lorenz2_jac!, maxsteps=2000, parse_eqs=false);

        tv4p, lv4p, xv4p = lyap_taylorinteg(lorenz2!, q0, t0, tf, _order, _abstol,
            lorenz2_jac!, maxsteps=2000,  parse_eqs=true);

        @test tv4 == tv4p
        @test lv4 == lv4p
        @test xv4 == xv4p

        # Comparing both integrations (lorenz2)
        @test tv3 == tv4
        @test lv3 == lv4
        @test xv3 == xv4

        # Comparing both integrations (lorenz1! vs lorenz2!)
        @test tv1 == tv3
        @test lv1 == lv3
        @test xv1 == xv3

        # Using ranges
        lv5, xv5 = lyap_taylorinteg(lorenz2!, q0, t0:0.125:tf, _order, _abstol,
            maxsteps=2000, parse_eqs=false);

        lv5p, xv5p = lyap_taylorinteg(lorenz2!, q0, t0:0.125:tf, _order, _abstol,
            maxsteps=2000, parse_eqs=true);

        @test lv5 == lv5p
        @test xv5 == xv5p

        lv6, xv6 = lyap_taylorinteg(lorenz2!, q0, t0:0.125:tf, _order, _abstol,
            lorenz2_jac!, maxsteps=2000, parse_eqs=false);

        lv6p, xv6p = lyap_taylorinteg(lorenz2!, q0, t0:0.125:tf, _order, _abstol,
            lorenz2_jac!, maxsteps=2000,  parse_eqs=true);

        @test lv6 == lv6p
        @test xv6 == xv6p

        tT = 0.0 + Taylor1(_order)
        qT = q0 .+ zero(tT)
        TaylorIntegration.__jetcoeffs!(Val(true), lorenz1!, tT, qT, similar(qT), similar(qT), params)

        qT = q0 .+ zero(tT)
        TaylorIntegration.__jetcoeffs!(Val(true), lorenz2!, tT, qT, similar(qT), similar(qT), nothing)
    end

    @testset "Tests for throwing errors" begin
        # Wrong number of arguments
        ex = :(function f_p!(dx, x, p, t, y)
            dx[1] = x[2]
            dx[2] = -sin( x[1] )
        end)
        @test_throws ArgumentError TaylorIntegration._make_parsed_jetcoeffs(ex)

        # `&&` is not yet implemented
        ex = :(function f_p!(x, p, t)
            true && x
        end)
        @test_throws ArgumentError TaylorIntegration._make_parsed_jetcoeffs(ex)

        # a is not an Expr; String
        ex = :(function f_p!(x, p, t)
            "a"
        end)
        @test_throws ArgumentError TaylorIntegration._make_parsed_jetcoeffs(ex)

        # KeyError: key :fname not found
        ex = :(begin
            x=1
            x+x
        end)
        @test_throws KeyError TaylorIntegration._make_parsed_jetcoeffs(ex)

        # BoundsError; no return variable defined
        ex = :(function f_p!(x, p, t)
            local cos(t)
        end)
        @test_throws BoundsError TaylorIntegration._make_parsed_jetcoeffs(ex)

        # ArgumentError; .= is not yet implemented
        ex = :(function err_bbroadcasting!(Dz, z, p, t)
            Dz .= z
            nothing
        end)
        @test_throws ArgumentError TaylorIntegration._make_parsed_jetcoeffs(ex)

        # The macro works fine, but the `jetcoeffs!` method does not work
        # (so the integration would silently run using `parse_eqs=false`)
        @taylorize function harm_osc!(dx, x, p, t)
            local ω = p[1]
            # local ω2 = ω^2    # Needed this to avoid an error
            dx[1] = x[2]
            dx[2] = - (ω^2 * x[1])  # ω^2 -> ω2
            return nothing
        end
        tT = t0 + Taylor1(_order)
        qT = [1.0, 0.0] .+ zero(tT)
        @test_throws MethodError TaylorIntegration.__jetcoeffs!(Val(true), harm_osc!, tT, qT, similar(qT), similar(qT), [2.0])

        @taylorize function kepler1!(dq, q, p, t)
            μ = p
            r_p3d2 = (q[1]^2+q[2]^2)^1.5

            dq[1] = q[3]
            dq[2] = q[4]
            dq[3] = μ * q[1]/r_p3d2
            dq[4] = μ * q[2]/r_p3d2

            return nothing
        end
        tT = t0 + Taylor1(_order)
        qT = [0.2, 0.0, 0.0, 3.0] .+ zero(tT)
        @test_throws MethodError TaylorIntegration.__jetcoeffs!(Val(true), kepler1!, tT, qT, similar(qT), similar(qT), -1.0)
    end

    @testset "Jet transport with @taylorize macro" begin
        @taylorize function pendulum!(dx, x, p, t)
            dx[1] = x[2]
            dx[2] = -sin( x[1] )
            nothing
        end

        varorder = 2 #the order of the variational expansion
        p = set_variables("ξ", numvars=2, order=varorder) #TaylorN steup
        q0 = [1.3, 0.0] #the initial conditions
        q0TN = q0 + p #parametrization of a small neighbourhood around the initial conditions
        # T is the librational period == 4Elliptic.K(sin(q0[1]/2)^2)
        T = 4Elliptic.K(sin(q0[1]/2)^2) # equals 7.019250311844546
        integstep = 0.25*T #the time interval between successive evaluations of the solution vector

        #the time range
        tr = t0:integstep:T;
        #note that as called below, taylorinteg uses the parsed jetcoeffs! method by default
        xvp = taylorinteg(pendulum!, q0, tr, _order, _abstol, maxsteps=100)

        # "warmup" for jet transport integration
        xvTN = taylorinteg(pendulum!, q0TN, tr, _order, _abstol, maxsteps=1, parse_eqs=false)
        xvTN = taylorinteg(pendulum!, q0TN, tr, _order, _abstol, maxsteps=1)
        @test size(xvTN) == (5,2)
        #jet transport integration with parsed jetcoeffs!
        xvTNp = taylorinteg(pendulum!, q0TN, tr, _order, _abstol, maxsteps=100)
        #jet transport integration with non-parsed jetcoeffs!
        xvTN = taylorinteg(pendulum!, q0TN, tr, _order, _abstol, maxsteps=100, parse_eqs=false)
        @test xvTN == xvTNp
        @test norm(xvTNp[:,:]() - xvp, Inf) < 1e-15

        dq = 0.0001rand(2)
        q1 = q0 + dq
        yv = taylorinteg(pendulum!, q1, tr, _order, _abstol, maxsteps=100)
        yv_jt = xvTNp[:,:](dq)
        @test norm(yv-yv_jt, Inf) < 1e-11

        dq = 0.001
        t = Taylor1([0.0, 1.0], 10)
        x0T1 = q0+[0t,t]
        q1 = q0+[0.0,dq]
        tv, xv = taylorinteg(pendulum!, q1, t0, 2T, _order, _abstol)
        tvT1, xvT1 = taylorinteg(pendulum!, x0T1, t0, 2T, _order, _abstol, parse_eqs=false)
        tvT1p, xvT1p = taylorinteg(pendulum!, x0T1, t0, 2T, _order, _abstol)
        @test tvT1 == tvT1p
        @test xvT1 == xvT1p
        xv_jt = xvT1p[:,:](dq)
        @test norm(xv_jt[end,:]-xv[end,:]) < 20eps(norm(xv[end,:]))
    end

    @testset "Poincare maps with the pendulum" begin
        @taylorize function pendulum!(dx, x, p, t)
            dx[1] = x[2]
            dx[2] = -sin( x[1] )
            nothing
        end

        # Function defining the crossing events: the
        # surface of section is x[2]==0, x[1]>0
        function g(dx, x, params, t)
            if constant_term(x[1]) > 0
                return (true, x[2])
            else
                return (false, x[2])
            end
        end

        t0 = 0.0
        x0 = [1.3, 0.0]
        Tend = 7.019250311844546

        p = set_variables("ξ", numvars=length(x0), order=2)
        x0N = x0 + p
        ξ = Taylor1(2)
        x01 = x0 + [ξ, ξ]

        #warm-up lap and preliminary tests
        taylorinteg(pendulum!, g, x0, t0, Tend, _order, _abstol, maxsteps=1)
        @test_throws AssertionError taylorinteg(pendulum!, g, x0, t0, Tend,
            _order, _abstol, maxsteps=1, eventorder=_order+1)

        #testing 0-th order root-finding
        tv, xv, tvS, xvS, gvS = taylorinteg(pendulum!, g, x0, t0, 3Tend,
            _order, _abstol, maxsteps=1000)
        @test tv[1] == t0
        @test xv[1,:] == x0
        @test size(tvS) == (2,)
        @test norm(tvS-[Tend, 2Tend], Inf) < 1e-13
        @test norm(gvS, Inf) < eps()

        #testing 0-th order root-finding with time ranges/vectors
        tvr = [t0, Tend/2, Tend, 3Tend/2, 2Tend, 5Tend/2, 3Tend]
        taylorinteg(pendulum!, g, x0, view(tvr, :), _order, _abstol, maxsteps=1)
        @test_throws AssertionError taylorinteg(pendulum!, g, x0, view(tvr, :),
            _order, _abstol, maxsteps=1, eventorder=_order+1)
        xvr, tvSr, xvSr, gvSr = taylorinteg(pendulum!, g, x0, view(tvr, :),
            _order, _abstol, maxsteps=1000)
        @test xvr[1,:] == x0
        @test size(tvSr) == (2,)
        @test norm(tvSr-tvr[3:2:end-1], Inf) < 1e-13
        @test norm(tvr[3:2:end-1]-tvSr, Inf) < 1e-14
        @test norm(xvr[3:2:end-1,:]-xvSr, Inf) < 1e-14
        @test norm(gvSr[:]) < eps()
        @test norm(tvS-tvSr, Inf) < 5e-15
    end

    @testset "Tests parsing specific aspects of the expression" begin
        ex = :(
            function f!(dq::Array{T,1}, q::Array{T,1}, p, t) where {T}
                aa = my_simple_function(q, p, t)
                for i in 1:length(q)
                    if i == 1
                        dq[i] = 2q[i]
                    elseif i == 2
                        dq[i] = q[i]
                    elseif i == 3
                        dq[i] = aa
                        continue
                    else
                        dq[i] = my_complicate_function(q)
                        break
                    end
                end
                nothing
            end
        )
        newex = TaylorIntegration._make_parsed_jetcoeffs(ex)

        # Ignore declarations in the function
        @test newex.args[1] == :(
            TaylorIntegration.jetcoeffs!(::Val{f!}, t::Taylor1{_T},
                q::AbstractVector{Taylor1{_S}},
                dq::AbstractVector{Taylor1{_S}}, p) where
                    {_T <: Real, _S <: Number})

        # Include not recognized functions as they appear
        @test newex.args[2].args[2] == :(aa = my_simple_function(q, p, t))
        @test newex.args[2].args[5].args[2].args[2] ==
            :(aa = my_simple_function(q, p, t))

        # Issue 96: deal with `elseif`s, `continue` and `break`
        @test newex.args[2].args[5].args[2].args[3] ==
            :(for i = 1:length(q)
                  if i == 1
                      TaylorSeries.mul!(dq[i], 2, q[i], ord)
                  else
                      if i == 2
                          TaylorSeries.identity!(dq[i], q[i], ord)
                      else
                          if i == 3
                              TaylorSeries.identity!(dq[i], aa, ord)
                              continue
                          else
                              dq[i] = my_complicate_function(q)
                              break
                          end
                      end
                  end
              end
              )
    end

    @testset "Test @taylorize with @threads" begin
        @taylorize function f1!(dq, q, params, t)
            for i in 1:10
                dq[i] = q[i]
            end
            nothing
        end

        @taylorize function f1_parsed!(dq, q, params, t)
            for i in 1:10
                dq[i] = q[i]
            end
            nothing
        end

        @taylorize function f2!(dq, q, params, t)
            Threads.@threads for i in 1:10
                dq[i] = q[i]
            end
            nothing
        end

        @taylorize function f3!(dq, q, params, t)
            @threads for i in 1:10
                dq[i] = q[i]
            end
            nothing
        end

        x01 = Taylor1.(rand(10), _order)
        dx01 = similar(x01)
        t1 = Taylor1(_order)
        xaux1 = similar(x01)

        x01p = deepcopy(x01)
        dx01p = similar(x01)
        t1p = deepcopy(t1)

        x02 = deepcopy(x01)
        dx02 = similar(x01)
        t2 = deepcopy(t1)

        x03 = deepcopy(x01)
        dx03 = similar(x01)
        t3 = deepcopy(t1)

        TaylorIntegration.jetcoeffs!(f1!, t1, x01, dx01, xaux1, nothing)
        TaylorIntegration.jetcoeffs!(Val(f1_parsed!), t1p, x01p, dx01p, nothing)
        TaylorIntegration.jetcoeffs!(Val(f2!), t2, x02, dx02, nothing)
        TaylorIntegration.jetcoeffs!(Val(f3!), t3, x03, dx03, nothing)

        @test x01 == x01p
        @test x01p == x02
        @test x02 == x03

        # An example involving coupled dofs: a 1D chain of harmonic oscillators
        @taylorize function harmosc1dchain!(dq, q, params, t)
            local N = Int(length(q)/2)
            local _eltype_q_ = eltype(q)
            local μ = params
            X = Array{_eltype_q_}(undef, N, N)
            accX = Array{_eltype_q_}(undef, N) #acceleration
            for j in 1:N
                accX[j] = zero(q[1])
                dq[j] = q[N+j]
            end
            #compute accelerations
            for j in 1:N
                for i in 1:N
                    if i == j
                    else
                        X[i,j] = q[i]-q[j]
                        temp_001 = accX[j] + (μ[i]*X[i,j])
                        accX[j] = temp_001
                    end #if i != j
                end #for, i
            end #for, j
            for i in 1:N
                dq[N+i] = accX[i]
            end
            nothing
        end

        # Same as harmosc1dchain!, but adding a `@threads for`
        @taylorize function harmosc1dchain_threads!(dq, q, params, t)
            local N = Int(length(q)/2)
            local _eltype_q_ = eltype(q)
            local μ = params
            X = Array{_eltype_q_}(undef, N, N)
            accX = Array{_eltype_q_}(undef, N) #acceleration
            for j in 1:N
                accX[j] = zero(q[1])
                dq[j] = q[N+j]
            end
            #compute accelerations
            Threads.@threads for j in 1:N
                for i in 1:N
                    if i == j
                    else
                        X[i,j] = q[i]-q[j]
                        temp_001 = accX[j] + (μ[i]*X[i,j])
                        accX[j] = temp_001
                    end #if i != j
                end #for, i
            end #for, j
            for i in 1:N
                dq[N+i] = accX[i]
            end
            nothing
        end

        N = 200
        x0 = 10randn(2N)
        t = Taylor1(_order)
        μ = 1e-7rand(N)
        x = Taylor1.(x0, t.order)
        dx = similar(x)
        t_ = deepcopy(t)
        x_ = Taylor1.(x0, t.order)
        dx_ = similar(x_)

        @show Threads.nthreads()

        TaylorIntegration.jetcoeffs!(Val(harmosc1dchain!), t, x, dx, μ)
        @time TaylorIntegration.jetcoeffs!(Val(harmosc1dchain!), t, x, dx, μ)

        TaylorIntegration.jetcoeffs!(Val(harmosc1dchain_threads!), t_, x_, dx_, μ)
        @time TaylorIntegration.jetcoeffs!(Val(harmosc1dchain_threads!), t_, x_, dx_, μ)

        @test x == x_
        @test dx == dx_

        tv, xv = taylorinteg(harmosc1dchain!, x0, t0, 10000.0, _order, _abstol, μ)
        tv_, xv_ = taylorinteg(harmosc1dchain_threads!, x0, t0, 10000.0, _order, _abstol, μ)

        @test tv == tv_
        @test xv == xv_
    end

end


text:
module DoubletSource

## Package imports
#==========================================================================================#

using LinearAlgebra
import Base.Iterators: product
using StaticArrays
import SplitApplyCombine: combinedimsview

import ..MathTools: rotation, inverse_rotation, midpair_map

import ..Laplace: Uniform2D, magnitude, angle, velocity

import ..NonDimensional: pressure_coefficient

import ..PanelGeometry: AbstractPanel2D, Panel2D, WakePanel2D, collocation_point, p1, p2, transform_panel, affine_2D, panel_length, panel_angle, panel_tangent, panel_normal, distance, wake_panel, wake_panels, panel_points, panel_vector

import ..AeroMDAO: solve_system, surface_velocities, surface_coefficients

## Doublet-source Dirichlet boundary condition
#===========================================================================#

source_potential(str, x, z, x1, x2) = str / 4π * ((x - x1) * log((x - x1)^2 + z^2) - (x - x2) * log((x - x2)^2 + z^2) + 2z * (atan(z, x - x2) - atan(z, x - x1)))

source_velocity(str, x, z, x1, x2) = SVector(str / (4π) * log(((x - x1)^2 + z^2) / ((x - x2)^2 + z^2)), doublet_potential(str, x, z, x1, x2))

doublet_potential(str, x, z, x1, x2) = str / (2π) * (atan(z, x - x1) - atan(z, x - x2))

doublet_velocity(str, x, z, x1, x2) = SVector(str / (2π) * - (z / ((x - x1)^2 + z^2) - z / ((x - x2)^2 + z^2) ), str / (2π) * ( (x - x1) / ((x - x1)^2 + z^2) - (x - x2) / ((x - x2)^2 + z^2)))

## Matrix helpers
#===========================================================================#

function doublet_influence(panel_j :: AbstractPanel2D, panel_i :: AbstractPanel2D)
    xp, yp = transform_panel(panel_j, panel_i)
    ifelse(panel_i == panel_j, 0.5, doublet_potential(1., xp, yp, 0., panel_length(panel_j)))
end

function source_influence(panel_j :: AbstractPanel2D, panel_i :: AbstractPanel2D)
    xp, yp = transform_panel(panel_j, panel_i)
    source_potential(1., xp, yp, 0., panel_length(panel_j))
end

boundary_condition(panel_j :: AbstractPanel2D, panel_i :: AbstractPanel2D, u) = -source_influence(panel_j, panel_i) * dot(u, panel_normal(panel_j))

## Aerodynamic coefficients
#===========================================================================#

surface_velocity(dφ, dr, u, α) = dφ / dr + dot(u, α)

# """
#     aerodynamic_coefficients(vels, Δrs, panel_angles, speed, α)

# Compute the lift, moment, and pressure coefficients given associated arrays of edge speeds, adjacent collocation point distances, panel angles, the freestream speed, and angle of attack ``α``.
# """
# function evaluate_coefficients(vels, Δrs, xjs, panel_angles, speed, α)
#     cps   = @. pressure_coefficient(speed, vels)
#     cls   = @. lift_coefficient(cps, Δrs, panel_angles)
#     cms   = @. -cls * xjs * cos(α)

#     cls, cms, cps
# end

## Matrix assembly
#===========================================================================#

include("matrix_func.jl")

struct DoubletSourceSystem{T <: Real, M <: AbstractMatrix{T}, N <: AbstractVector{T}, O <: AbstractVector{<: AbstractPanel2D}, R <: WakePanel2D, P <: Uniform2D}
    influence_matrix   :: M
    boundary_condition :: N
    singularities      :: N
    surface_panels     :: O
    wake_panels        :: R
    freestream         :: P
end

function Base.show(io :: IO, sys :: DoubletSourceSystem)
    println(io, "DoubletSourceSystem —")
    println(io, length(sys.surface_panels), " ", eltype(sys.surface_panels), " Elements")
end

function solve_system(panels, uni :: Uniform2D, sources :: Bool, wake_length)
    # Freestream conditions
    u, α  = velocity(uni), uni.angle

    # Build wake
    wake_pan = wake_panel(panels, wake_length, α)

    # speed           = norm(u)
    # xs              = getindex.(panel_points(panels)[2:end-1], 1)

    # Blunt trailing edge tests
    # te_panel        = Panel2D((p2 ∘ last)(panels), (p1 ∘ first)(panels))
    # r_te            = panel_vector(te_panel)
    # φ_TE            = dot(u, r_te)

    # Solve for doublet strengths
    φs, AIC, boco   = solve_linear(panels, u, α, wakes; bound = wake_length)

    DoubletSourceSystem(AIC, boco, φs, panels, wakes, uni)

    # # Evaluate inviscid edge velocities
    # u_es, Δrs       = tangential_velocities(panels, φs, u, sources)

    # # Compute coefficients
    # cls, cms, cps   = evaluate_coefficients(u_es, Δrs, xs, panel_angle.(panels[2:end]), speed, α)

    # # Evaluate lift coefficient from wake doublet strength
    # cl_wake         = lift_coefficient(φs[end] - φs[1] + φ_TE, speed)

    # cls, cms, cps, cl_wake
end

function solve_system(panels, uni :: Uniform2D, num_wake :: Integer, wake_length)
    u, α  = velocity(uni), uni.angle

    wake_pan = wake_panel(panels, wake_length, α)
    # wakes = wake_panels(panels, wake_length, num_wake)

    # Solve for doublet strengths
    φs, AIC, boco   = solve_linear(panels, u, wake_pan) # ; bound = wake_length)

    DoubletSourceSystem(AIC, boco, φs, panels, wake_pan, uni)
end

function surface_velocities(prob :: DoubletSourceSystem)
    # Panel properties
    ps   = prob.surface_panels
    Δrs  = @views @. distance(ps[2:end], ps[1:end-1])
    αs   = @views panel_tangent.(ps[2:end])

    @views surface_velocities(prob.singularities[1:end-1], Δrs, αs, velocity(prob.freestream), false)
end

function surface_coefficients(prob :: DoubletSourceSystem)
    # Panel properties
    ps   = prob.surface_panels
    Δrs  = @views @. distance(ps[2:end], ps[1:end-1])
    xs   = @views combinedimsview(panel_points(ps)[2:end-1])[1,:]
    θs   = @views panel_angle.(ps[2:end])

    # Inviscid edge velocities
    u_es = @views surface_velocities(prob)

    # Aerodynamic coefficients
    cps  = @. pressure_coefficient(prob.freestream.magnitude, u_es)
    cls  = @. -cps * Δrs * cos(θs)
    cms  = @. -cls * xs * cos(prob.freestream.angle)

    cls, cms, cps
end

lift_coefficient(prob :: DoubletSourceSystem) = 2 * last(prob.singularities) / prob.freestream.magnitude

end

ext: jl


text:
using ITensors
using ITensorVisualizationBase
using Graphs
using LayeredLayouts
using PastaQ: randomcircuit

Nx, Ny = 3, 3
N = Nx * Ny
# TODO: change to (Nx, Ny) with PastaQ v0.0.16
gates = randomcircuit(
  Nx, Ny; depth=4, twoqubitgates="CX", onequbitgates="Rn", layered=false, rotated=false
)

s = siteinds("Qubit", N)

U, s̃ = circuit_network(gates, s)
ψ = MPS(s)
ψ̃ = MPS(s̃)
tn = [prod(ψ), U..., prod(ψ̃)]

edge_labels = (; plevs=true)
layout(g) = layered_layout(solve_positions(Zarate(), g))
@visualize fig tn arrow_show = true edge_labels = edge_labels layout = layout edge_textsize =
  20

fig


text:
@testset "active pattern" begin
    @testset "regular case" begin
        @active LessThan0(x) begin
            if x > 0
                nothing
            else
                x
            end
        end

        @test (@match 15 begin
            LessThan0(_) => :a
            _ => :b
        end) === :b

        @test (@match -15 begin
            LessThan0(a) => a
            _ => 0
        end) == -15
    end
    @testset "parametric case" begin

        @active Re{r :: Regex}(x) begin
             match(r, x)
        end

        @test (@match "123" begin
            Re{r"\d+"}(x) => x.match
            _ => @error ""
        end) == "123"

        @test_skip @match "abc" begin
            Re{r"\d+"}(x) => x
            _ => @error ""
        end
    end
end


text:
"""
A LatticeBond is a struct which represents
a single bond in a geometrical lattice or
else on interaction graph defining a physical
model such as a quantum Hamiltonian.

LatticeBond has the following data fields:
* s1::Int -- number of site 1
* s2::Int -- number of site 2
* x1::Int -- x coordinate of site 1
* y1::Int -- y coordinate of site 1
* x2::Int -- x coordinate of site 2
* y2::Int -- y coordinate of site 2
* type::String -- optional description of bond type
"""
struct LatticeBond
  s1::Int
  s2::Int
  x1::Float64
  y1::Float64
  x2::Float64
  y2::Float64
  type::String
end

"""
    LatticeBond(s1::Int,s2::Int)

    LatticeBond(s1::Int,s2::Int,
                x1::Real,y1::Real,
                x2::Real,y2::Real,
                type::String="")

Construct a LatticeBond struct by
specifying just the numbers of sites
1 and 2, or additional details including
the (x,y) coordinates of the two sites and
an optional type string.
"""
function LatticeBond(s1::Int,s2::Int)
  return LatticeBond(s1,s2,0.0,0.0,0.0,0.0,"")
end

function LatticeBond(s1::Int,s2::Int,
                     x1::Real,y1::Real,
                     x2::Real,y2::Real,
                     bondtype::String="")
  cf(x) = convert(Float64,x)
  return LatticeBond(s1,s2,cf(x1),cf(y1),cf(x2),cf(y2),bondtype)
end

"""
Lattice is an alias for Vector{LatticeBond}
"""
const Lattice = Vector{LatticeBond}

"""
    square_lattice(Nx::Int,
                   Ny::Int;
                   kwargs...)::Lattice

Return a Lattice (array of LatticeBond
objects) corresponding to the two-dimensional
square lattice of dimensions (Nx,Ny).
By default the lattice has open boundaries,
but can be made periodic in the y direction
by specifying the keyword argument
`yperiodic=true`.
"""
function square_lattice(Nx::Int,
                        Ny::Int;
                        kwargs...)::Lattice
  yperiodic = get(kwargs,:yperiodic,false)
  yperiodic = yperiodic && (Ny > 2)
  N = Nx*Ny
  Nbond = 2N-Ny + (yperiodic ? 0 : -Nx)
  latt = Lattice(undef,Nbond)
  b = 0
  for n=1:N
    x = div(n-1,Ny)+1
    y = mod(n-1,Ny)+1
    if x < Nx
      latt[b+=1] = LatticeBond(n,n+Ny,x,y,x+1,y)
    end
    if Ny > 1
      if y < Ny
        latt[b+=1] = LatticeBond(n,n+1,x,y,x,y+1);
      end
      if yperiodic && y==1
        latt[b+=1] = LatticeBond(n,n+Ny-1,x,y,x,y+Ny)
      end
    end
  end
  return latt
end

"""
    triangular_lattice(Nx::Int,
                       Ny::Int;
                       kwargs...)::Lattice

Return a Lattice (array of LatticeBond
objects) corresponding to the two-dimensional
triangular lattice of dimensions (Nx,Ny).
By default the lattice has open boundaries,
but can be made periodic in the y direction
by specifying the keyword argument
`yperiodic=true`.
"""
function triangular_lattice(Nx::Int,
                            Ny::Int;
                            kwargs...)::Lattice
  yperiodic = get(kwargs,:yperiodic,false)
  yperiodic = yperiodic && (Ny > 2)
  N = Nx*Ny
  Nbond = 3N-2Ny + (yperiodic ? 0 : -2Nx+1)
  latt = Lattice(undef,Nbond)
  b = 0
  for n=1:N
    x = div(n-1,Ny)+1
    y = mod(n-1,Ny)+1

    # x-direction bonds
    if x < Nx
      latt[b+=1] = LatticeBond(n,n+Ny)
    end

    # 2d bonds
    if Ny > 1
      # vertical / y-periodic diagonal bond
      if (n+1 <= N) && ((y < Ny) || yperiodic)
        latt[b+=1] = LatticeBond(n,n+1);
      end
      # periodic vertical bond
      if yperiodic && y==1
        latt[b+=1] = LatticeBond(n,n+Ny-1)
      end
      # diagonal bonds
      if x < Nx && y < Ny
        latt[b+=1] = LatticeBond(n,n+Ny+1)
      end
    end
  end
  return latt
end


text:
# include(joinpath(WATERFALL_DIR,"src","includes.jl"))
import Base
import Combinatorics
import CSV
import DataFrames
import LinearAlgebra, LinearAlgebra.I
import Printf
import SparseArrays
import Statistics
import StatsPlots

# DATA_DIR = joinpath(WATERFALL_DIR,"data","pvrd")
# df = CSV.read(joinpath(DATA_DIR,"pvrd2-investment-metrics.csv"), DataFrames.DataFrame)

DATA_DIR = joinpath(WATERFALL_DIR,"data")
df = CSV.read(joinpath(DATA_DIR,"viz","investment.csv"), DataFrames.DataFrame)
[df[!,ii] .= coalesce.(df[:,ii],missing,"") for ii in 1:size(df,2)]
# df = df[[1:11;15],:]

# df[!,:Process] .= replace.(df[:,:Process], "interconnection"=>"intercon")
# df[!,:Label] .= string.(titlecase.(df[:,:Process])," (",df[:,:Step],")")
# df[!,:Units] .= "Efficiency (%)"

# Define global arguments.
distribution=:normal
fuzziness=(0.01,0.3)
interactivity=(0.2,0.4)
sample=:Sample
value=:Value
label=:Step
sublabel=:Process
nsample=5

kwargs = (
    label=label,
    sublabel=sublabel,
    distribution=distribution,
    fuzziness=fuzziness,
    interactivity=interactivity,
    # vmin=12.0,
    # ncor=ncor,
    # nsample=nsample,
)


text:
# Operations on 2D and 3D vectors, and colors

const Vec3{T<:AbstractFloat} = SVector{3, T}

# Adapted from @Christ_Foster's: https://discourse.julialang.org/t/ray-tracing-in-a-week-end-julia-vs-simd-optimized-c/72958/40
# Christ added this feature directly in StaticArrays.jl on 2021-12-17, so
# this will soon be no longer needed.
@inline @inbounds function Base.getproperty(v::SVector{3}, name::Symbol)
    name === :x || name === :r ? getfield(v, :data)[1] :
    name === :y || name === :g ? getfield(v, :data)[2] :
    name === :z || name === :b ? getfield(v, :data)[3] : getfield(v, name)
end

@inline @inbounds function Base.getproperty(v::SVector{2}, name::Symbol)
    name === :x || name === :r ? getfield(v, :data)[1] :
    name === :y || name === :g ? getfield(v, :data)[2] : getfield(v, name)
end

@inline squared_length(v::Vec3) = v ⋅ v
@inline near_zero(v::Vec3) = squared_length(v) < 1e-5
@inline rgb(v::Vec3) = RGB(v...)
@inline rgb_gamma2(v::Vec3) = RGB(sqrt.(v)...)


text:
"""
    type Layer
		name::Symbol
		neurons::Neuron
		n_of::Int64
		inputconnections::Dict{Symbol, Dict{Symbol, Connection}}

"""
type Layer
	name::Symbol
	neurons::Neuron
	n_of::Int64
	inputconnections::Dict{Symbol, Dict{Symbol, Connection}}
end
export Layer

"""
    Layer(name::Symbol, neuron::DataType, neuronparams::NeuronParameters, n_of::Int64)

Create an unconnected `Layer` of `n_of` neurons of type `neuron` with parameters
`neuronparams`.
"""
function Layer(name::Symbol, neuron::DataType,
			   neuronparams::NeuronParameters, n_of::Int64)
	Layer(name, neuron(n_of, neuronparams), n_of,
	      Dict{Symbol, Dict{Symbol, Connection}}())
end

"""
    connect!(from::Layer, to::Layer, con::Connection; label = :default)

Add to inputconnections of layer `to` the connection `con` with key
[`label`][`from.name`].
"""
function connect!(from::Layer, to::Layer, con::Connection; label = :default)
	if haskey(to.inputconnections, label)
		to.inputconnections[label][from.name] = con
	else
		to.inputconnections[label] = Dict(from.name => con)
	end
	string(typeof(con))
end

"""
    connect!(from::Layer, to::Layer, kind::Union{DataType, Function}; label = :default)

Create connection `kind(from, to)` and add it.
"""
function connect!(from::Layer, to::Layer, kind::Union{DataType, Function};
				  label = :default)
	connect!(from, to, kind(from, to), label = label)
end

export connect!

max_stars_repo_name: jbrea/NetSim

max_issues_repo_name: jbrea/NetSim

max_forks_repo_name: jbrea/NetSim


text:
# The SteinInverseMultiquadricKernel for R^d is given by
#
# k(r) = (c^2 + r^2)^{-beta}
#
# where r = |x-y|_2.

type SteinInverseMultiquadricKernel <: SteinKernel
    beta::Float64       # -beta is exponent
    c2::Float64         # equal to c^2
    SteinInverseMultiquadricKernel(beta, c2) = (
        @assert (0 < c2);
        @assert (0 < beta);
        new(beta, c2)
    )
end

# default will be 1.0 for c^2
SteinInverseMultiquadricKernel(beta::Float64) = SteinInverseMultiquadricKernel(beta, 1.0)
# default will be 0.5 for beta
SteinInverseMultiquadricKernel() = SteinInverseMultiquadricKernel(0.5)

function k(ker::SteinInverseMultiquadricKernel, x::Array{Float64,1}, y::Array{Float64,1})
    r = norm(x - y)
    c2 = ker.c2
    beta = ker.beta
    (c2 + r^2)^(-beta)
end

function gradxk(ker::SteinInverseMultiquadricKernel, x::Array{Float64,1}, y::Array{Float64,1})
    r = norm(x - y)
    c2 = ker.c2
    beta = ker.beta
    -2.0 * beta * (x - y) * (c2 + r^2)^(-beta-1.0)
end

function gradxyk(ker::SteinInverseMultiquadricKernel, x::Array{Float64,1}, y::Array{Float64,1})
    r = norm(x - y)
    c2 = ker.c2
    beta = ker.beta
    d = length(x)
    2 * d * beta * (c2 + r^2)^(-beta-1.0) - 4 * beta * (beta + 1) * r^2 * (c2 + r^2)^(-beta-2.0)
end

# This is reimplemented here, optimized for speed!
function k0(ker::SteinInverseMultiquadricKernel,
            x::Array{Float64, 1},
            y::Array{Float64, 1},
            gradlogpx::Array{Float64, 1},
            gradlogpy::Array{Float64, 1})
    # compute all the core values only once and store them
    c2 = ker.c2
    beta = ker.beta
    d = length(x)

    z = x - y
    r2 = sum(abs2, z)
    base = c2 + r2
    base_beta = base^(-beta)
    base_beta1 = base_beta / base

    coeffk = dot(gradlogpx, gradlogpy)
    coeffgrad = -2.0 * beta * base_beta1

    kterm = coeffk * base_beta
    gradandgradgradterms = coeffgrad * (
        (dot(gradlogpy, z) - dot(gradlogpx, z)) +
        (-d + 2 * (beta + 1) * r2 / base)
    )
    kterm + gradandgradgradterms
end


text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

using JuliaFEM
using JuliaFEM.Testing

#= TODO: Fix test

facts("test interpolation of fields") do

    # interpolation of field in spatial domain
    N = Basis((xi) -> [0.5*(1.0-xi[1]), 0.5*(1.0+xi[1])])
    u = Field(0.0, [0.0, 1.0])
    @fact interpolate(N, u, [0.0]) --> 0.5

    # interpolation of fieldset in time domain
    u1 = Field(0.0, [0.0, 1.0])
    u2 = Field(1.0, [1.0, 2.0])
    u = FieldSet("displacement", [u1, u2])
    @fact interpolate(u, 0.5).values --> [0.5, 1.5]
    @fact interpolate(u, 0.5).time --> 0.5

    # interpolation of fieldset is defined for every time value:
    u1 = Field(0.0, [0.0, 0.0])
    u2 = Field(1.0, [1.0, 2.0])
    u3 = Field(2.0, [0.5, 1.5])
    u = FieldSet("displacement", [u1, u2, u3])
    @fact interpolate(u, -1.0).values --> [0.0, 0.0]  # "out of range -" -> first known value
    @fact interpolate(u, 0.0).values --> [0.0, 0.0]
    @fact interpolate(u, 1.0).values --> [1.0, 2.0]
    @fact interpolate(u, 2.0).values --> [0.5, 1.5]
    @fact interpolate(u, 3.0).values --> [0.5, 1.5]  # "out of range +" -> last known value
    @fact interpolate(u, 0.5).values --> [0.5, 1.0]
    @fact interpolate(u, 1.5).values --> [0.75, 1.75]
    # use Inf to get very first or last value of field
    @fact interpolate(u, -Inf).values --> [0.0, 0.0]
    @fact interpolate(u, +Inf).values --> [0.5, 1.5]

    # multidimensional interpolation with and without derivatives
    h(xi) = [
        (1-xi[1])*(1-xi[2])/4
        (1+xi[1])*(1-xi[2])/4
        (1+xi[1])*(1+xi[2])/4
        (1-xi[1])*(1+xi[2])/4]
    dh(xi) = [
        -(1-xi[2])/4.0   -(1-xi[1])/4.0
         (1-xi[2])/4.0   -(1+xi[1])/4.0
         (1+xi[2])/4.0    (1+xi[1])/4.0
        -(1+xi[2])/4.0    (1-xi[1])/4.0]
    N = Basis(h, dh)

    X = Field(0.0, Vector[[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0]])

    # get midpoint of field in spatial domain
    @fact interpolate(N, X, [0.0, 0.0]) --> [0.5, 0.5]

    # interpolate of scalar field -> scalar
    H = Field(0.0, 6.0)
    @fact interpolate(N, H, [0.0, 0.0]) --> 6.0

    # multiplying scalar field with a vector -> vector
    # this is actually not so good idea...
    #h(xi) = [1/2*(1-xi[1]), 1/2*(1+xi[1])]
    #dh(xi) = [-1/2 1/2]'
    #N = Basis(h, dh)
    #f = Field(0.0, 100.0)
    #@fact interpolate(N, f, [0.0]) --> [50.0, 50.0]

end

=#


text:
@testset "propagators" begin
    p = OpenTelemetryAPI.CompositePropagator()
    push!(p, TraceContextTextMapPropagator())

    test_trace_id = TraceIdType(12345678901234567890123456789012)
    test_span_id = SpanIdType(1234567890123456)
    tracestate_value = "foo=1,bar=2,baz=3"
    header = [
        "traceparent" => "00-$test_trace_id-$test_span_id-00",
        "tracestate" => tracestate_value,
    ]

    with_context(; x = 123) do
        ctx = extract(header, p)
        @test ctx[:x] == 123
        @test ctx |> current_span |> span_status == SpanStatus(SPAN_STATUS_UNSET)

        sc = span_context(ctx)
        @test sc.trace_id == test_trace_id
        @test sc.span_id == test_span_id
        @test string(sc.trace_state) ==
              string(TraceState("foo" => "1", "bar" => "2", "baz" => "3"))
        @test sc.trace_flag == TraceFlag(false)
    end

    header = Dict("Content-Type" => "text/plain")
    with_span("test") do
        inject!(header, p)
        @test header["traceparent"] ==
              "00-00000000000000000000000000000000-0000000000000000-00"  # invalid trace parent
        @test !haskey(header, "tracestate")
    end

    inject!(nothing, p)  # Should not throw error
    extract(nothing, p)  # Should not throw error
end


text:
# This file is a part of EmpiricalDistributions.jl, licensed under the MIT License (MIT).

"""
    MvBinnedDist <: StatsBase.Distribution{Multivariate,Continuous}

A binned multivariate distribution, usually derived from a histogram.

Constructors:

```julia
    MvBinnedDist(h::StatsBase.Histogram{<:Real,N})
    MvBinnedDist{T<:Real}(h::StatsBase.Histogram{<:Real,N})
```

You can convert a `MvBinnedDist` back to a histogram via

```julia
convert(StatsBase.Histogram, dist::MvBinnedDist)
StatsBase.Histogram(dist::MvBinnedDist)
```
"""
struct MvBinnedDist{
    T <: Real,
    N,
    U <: Real,
    ET <: NTuple{N,AbstractVector{<:Real}},
    VT <: AbstractVector{T},
    MT <: AbstractMatrix{T},
    VU <: AbstractVector{U},
    AU <: AbstractArray{U,N}
} <: Distributions.Distribution{Multivariate,Continuous}
    _edges::ET
    _bin_pdf::AU
    _bin_linidx_cdf::VU
    _closed_left::Bool
    _mean::VT
    _mode::VT
    _var::VT
    _cov::MT
end

export MvBinnedDist

function MvBinnedDist{T}(h::Histogram{<:Real}) where {T<:Real}
    nh = normalize(h)

    edges = nh.edges
    bin_pdf = nh.weights

    closed_left = nh.closed == :left

    Y = nh.weights
    X = _bin_centers.(nh.edges)
    W = _bin_widths.(nh.edges)

    bin_linidx_cdf = cumsum(broadcast(idx -> Y[idx] .* prod(map(getindex, W, idx.I)), vec(CartesianIndices(Y))))
    @assert last(bin_linidx_cdf) ≈ 1
    bin_linidx_cdf[end] = 1

    mean_est_tpl = _mean(nh)
    mean_est = [T.(mean_est_tpl)...]
    mode_est = [T.(_mode(nh))...]
    var_est = [T.(_var(nh, mean_est_tpl))...]
    cov_est = T.(_cov(nh, mean_est_tpl))

    return MvBinnedDist(
        edges, bin_pdf, bin_linidx_cdf, closed_left,
        mean_est, mode_est, var_est, cov_est
    )
end

MvBinnedDist(h::Histogram{<:Real}) = MvBinnedDist{float(promote_type(map(eltype, h.edges)...))}(h)

function Adapt.adapt_structure(to, d::MvBinnedDist)
    MvBinnedDist(
        map(e -> adapt(to, e), d._edges), adapt(to, d._bin_pdf), adapt(to, d._bin_linidx_cdf),
        adapt(to, d._closed_left), adapt(to, d._mean), adapt(to, d._mode), adapt(to, d._var), adapt(to, d._cov)
    )
end

Histogram(d::MvBinnedDist) = Histogram(map(Array, d._edges), Array(d._bin_pdf), (d._closed_left ? :left : :right), true)
Base.convert(::Type{Histogram}, d::MvBinnedDist) = Histogram(d)

Base.length(d::MvBinnedDist{T,N}) where {T,N} = N
Base.size(d::MvBinnedDist{T,N}) where {T,N} = (N,)
Base.eltype(d::MvBinnedDist{T,N}) where {T,N} = T

Statistics.mean(d::MvBinnedDist) = d._mean
StatsBase.mode(d::MvBinnedDist) = d._mode
Statistics.var(d::MvBinnedDist) = d._var
Statistics.cov(d::MvBinnedDist) = d._cov

function Distributions._rand!(rng::AbstractRNG, d::MvBinnedDist{T,N}, A::AbstractVector{<:Real}) where {T,N}
    @assert length(eachindex(A)) == N
    u = rand(rng)
    i = searchsortedfirst(d._bin_linidx_cdf, u)
    idx_lo = CartesianIndices(d._bin_pdf)[i]
    idx_hi = idx_lo + CartesianIndex(ntuple(i -> 1, Val(N))...)
    x_lo = map(getindex, d._edges, idx_lo.I)
    x_hi = map(getindex, d._edges, idx_hi.I)
    for i in 1:N
        A[i] = _rand_uniform(rng, T, x_lo[i], x_hi[i])
    end
    return A
end

function Distributions._rand!(rng::AbstractRNG, d::MvBinnedDist{T,N}, A::AbstractMatrix{<:Real}) where {T,N}
    Distributions._rand!.(Ref(rng), (d,), nestedview(A))
    return A
end

# Similar to unroll_tuple in StaticArrays.jl:
@generated function _unsafe_unroll_tuple(A::AbstractArray, ::Val{L}) where {L}
    exprs = [:(A[idx0 + $j]) for j = 0:(L-1)]
    quote
        idx0 = firstindex(A)
        Base.@_inline_meta
        @inbounds return $(Expr(:tuple, exprs...))
    end
end

function Distributions.pdf(d::MvBinnedDist{T,N,U}, x::AbstractVector{<:Real}) where {T,N,U}
    length(eachindex(x)) == N || throw(ArgumentError("Length of variate doesn't match dimensionality of distribution"))
    xs = _unsafe_unroll_tuple(x, Val(N))

    idxs = _find_bin(d._edges, d._closed_left, xs)
    if checkbounds(Bool, d._bin_pdf, idxs...)
        @inbounds r = d._bin_pdf[idxs...]
        convert(U, r)
    else
        zero(U)
    end
end

function Distributions.logpdf(d::MvBinnedDist{T,N}, x::AbstractArray{<:Real, 1}) where {T,N}
    return log(pdf(d, x))
end

function Distributions._logpdf(d::MvBinnedDist{T,N}, x::AbstractArray{<:Real, 1}) where {T,N}
    return logpdf(d, x)
end


text:
all_Fs = [20, 25, 50]

SAVE_OUTPUT = true
PSI_HOVMOLLER = false

#-------------------
include("onehalflay.jl")

max_stars_repo_name: apaloczy/RoughOceanQGEnergetics

max_issues_repo_name: apaloczy/RoughOceanQGEnergetics

max_forks_repo_name: apaloczy/RoughOceanQGEnergetics


text:
# Writing a test for a new Hamiltonian formulation:

n = 4
x = zeros(3,n)
v = zeros(3,n)
vcom = zeros(3)
m = zeros(n)
h_cartesian = 0.0

# Compute Hamiltonian in Cartesian coordinates:
for i=1:n
  x[:,i] = randn(3)
  v[:,i] = randn(3)
  m[i] = exp(rand())
  h_cartesian += 0.5*m[i]*dot(v[:,i],v[:,i])
  vcom .+= m[i]*v[:,i]
end

# Center of mass velocity:
mtot = sum(m)
vcom ./= mtot

# Kinetic energy of the system:
h_sakura = h_cartesian
h_keplerian = 0.5*mtot*dot(vcom,vcom)

for i=1:n-1
  for j=i+1:n
    xij = x[:,j].-x[:,i]
    vij = v[:,j].-v[:,i]
    rij = norm(xij)
    h_cartesian -= m[i]*m[j]/rij
    mij = m[i]+m[j]
    mred = m[i]*m[j]/mij
    v2 = dot(vij,vij)
    h_sakura += mred*(0.5*v2-mij/rij) - 0.5*mred*v2
    h_keplerian += mred*(0.5*v2-mij/rij) - 0.5*(1.0-mij/mtot)*mred*v2
  end
end

println("h_cartesian: ",h_cartesian)
println("h_sakura:    ",h_sakura)
println("h_keplerian:    ",h_keplerian," ",h_keplerian/h_cartesian-1.)


text:
using FunctionalCollections
using FluxionArrays

function rearrange_dims(arr, arr_dims, target_dims)
    """
    Contracts:
    - elements of `arr_dims` and `target_dims` support the `==` operator
    - all elements of `arr_dims` are different
    - each element of `arr_dims` is present in `target_dims` exactly once
    - `arr` needs to support `reshape()`, `permutedims()` and `size()` methods.
    """
    if length(target_dims) == 0 || length(arr_dims) == 0
        return arr
    end

    old_shape = [size(arr, dim_num) for dim_num in 1:length(arr_dims)]

    # transpose the array so that the order of its dimensions
    # is the same as in `target_dims`
    dims_correspondence = [
        (findfirst(target_dims .== dim), old_dim_num)
        for (old_dim_num, dim) in enumerate(arr_dims)]
    new_dims_order = [old_dim_num for (new_dim_num, old_dim_num) in sort(dims_correspondence)]

    # Don't permute if the required order is already achieved
    if new_dims_order != collect(1:length(new_dims_order))
        arr = permutedims(arr, new_dims_order)
        new_shape = old_shape[new_dims_order]
        new_dims = arr_dims[new_dims_order]
    else
        new_shape = old_shape
        new_dims = arr_dims
    end

    # We may need to reshape the array adding new array dimensions (of size 1)
    # in place of unused array dimensions.
    # Currently just checking for the difference in shape lengths;
    # we probably don't need to add dummy dimensions if broadcasting works for us,
    # but that's harder to check
    if length(new_shape) != length(target_dims)

        new_dim_num = 1
        target_shape = Array{Any, 1}(length(target_dims))
        for target_dim_num = 1:length(target_dims)
            if (new_dim_num <= length(new_dims)
                    && target_dims[target_dim_num] == new_dims[new_dim_num])
                target_shape[target_dim_num] = new_shape[new_dim_num]
                new_dim_num += 1
            else
                target_shape[target_dim_num] = 1
            end
        end

        arr = reshape(arr, target_shape...)
    end

    return arr
end

function ensure_shape(arr, arr_shape, target_shape)

    if arr_shape == target_shape
        return arr
    end

    if length(arr_shape) < length(target_shape)
        extend_dims = length(target_shape) - length(arr_shape)
        extend_shape = [1 for i in 1:extend_dims]
        arr_shape = [arr_shape..., extend_shape...]
        arr = reshape(arr, arr_shape...)
    end

    outer = [
        (arr_dim == 1 && target_dim > 1) ? target_dim : 1
        for (arr_dim, target_dim) in zip(arr_shape, target_shape)
    ]

    repeat(arr, outer=outer)
end

function map_accum{T}(func, state, vals::Array{T, 1}, args...)
    new_vals = Array{Any, 1}(length(vals))
    for (i, val) in enumerate(vals)
        state, val = func(state, val, args...)
        new_vals[i] = val
    end
    state, new_vals
end

immutable ArrayExprBuildingState
    unknowns :: PersistentHashMap{FieldExpr, Unknown}
    random_states :: Nullable{Pair{Unknown, ArrayExpr}}
    random_variables :: PersistentHashMap{RandomVariable, ArrayExpr}
end

ArrayExprBuildingState() = ArrayExprBuildingState(
    PersistentHashMap{FieldExpr, Unknown}(),
    nothing,
    PersistentHashMap{RandomVariable, ArrayExpr}())

function add_if_missing(gen_val, state::ArrayExprBuildingState, key::FieldExpr)
    if haskey(state.unknowns, key)
        state, state.unknowns[key]
    else
        val = gen_val(key)
        new_unknowns = assoc(state.unknowns, key, val)
        new_state = ArrayExprBuildingState(
            new_unknowns, state.random_states, state.random_variables)
        new_state, val
    end
end

function set_random_state(state::ArrayExprBuildingState, rs::Unknown)
    ArrayExprBuildingState(
        state.unknowns, Pair{Unknown, ArrayExpr}(rs, rs), state.random_variables)
end

function update_random_state(state::ArrayExprBuildingState, rs::ArrayExpr)
    input_rs, output_rs = get(state.random_states)
    ArrayExprBuildingState(
        state.unknowns, Pair{Unknown, ArrayExpr}(input_rs, rs), state.random_variables)
end

function add_random_variable(state::ArrayExprBuildingState, fexpr::RandomVariable, aexpr::ArrayExpr)
    ArrayExprBuildingState(
        state.unknowns, state.random_states,
        assoc(state.random_variables, fexpr, aexpr))
end

function clear_random_variables(state::ArrayExprBuildingState)
    ArrayExprBuildingState(
        state.unknowns, state.random_states,
        PersistentHashMap{RandomVariable, ArrayExpr}())
end

function set_random_variables(state::ArrayExprBuildingState, random_variables)
    ArrayExprBuildingState(
        state.unknowns, state.random_states,
        random_variables)
end

random_states_set(state::ArrayExprBuildingState) = !isnull(state.random_states)

random_states(state::ArrayExprBuildingState) = get(state.random_states)
output_random_state(state::ArrayExprBuildingState) = random_states(state)[2]

function unknown_array_like(x::UnknownDimension)
    return unknown_array(shape=[], eltype=Float64)
end

function unknown_array_like(x::UnknownField)
    return unknown_array(
        shape=[Nullable{Int64}(length(vertices(grid(dim)))) for dim in known_dimensions(x)],
        eltype=Float64)
end

_as_array_expr(state, x::Number, target_dimensions) = state, lazy_array(x)
_as_array_expr(state, x::KnownDimension, target_dimensions) =
    state, rearrange_dims(lazy_array(vertices(grid(x))), [x], target_dimensions)
_as_array_expr(state, x::TransformedDimension, target_dimensions) =
    state, rearrange_dims(lazy_array(vertices(grid(x))), [x], target_dimensions)
_as_array_expr(state, x::KnownField, target_dimensions) =
    state, rearrange_dims(x.data, known_dimensions(x), target_dimensions)
_as_array_expr(state, x::UnknownDimension, target_dimensions) =
    add_if_missing(unknown_array_like, state, x)

function _as_array_expr(state, x::UnknownField, target_dimensions)
    state, arr_var = add_if_missing(unknown_array_like, state, x)
    state, rearrange_dims(arr_var, known_dimensions(x), target_dimensions)
end

function _as_array_expr(state, x::Fourier, target_dimensions)

    dims_dict = map(reverse, x.dimensions)

    fft_axes = []
    target_dimensions_inner = []
    for (i, dim) in enumerate(target_dimensions)
        if haskey(dims_dict, dim)
            push!(fft_axes, i)
            push!(target_dimensions_inner, dims_dict[dim])
        else
            push!(target_dimensions_inner, dim)
        end
    end

    state, expr = _as_array_expr(state, x.fexpr, target_dimensions_inner)

    if x.inverse
        result = ifft(expr, fft_axes)
    else
        result = fft(expr, fft_axes)
    end
    state, result
end

function _as_array_expr(state, fexpr::Sample, target_dimensions)
    target_dimensions = [
        x === STOCHASTIC_DIMENSION ? KnownStochasticDimension(fexpr.trajectories) : x
        for x in target_dimensions]

    # Random variables inside a sample() call are independent from the rest
    outer_rvs = state.random_variables
    inner_state = clear_random_variables(state)
    state, aexpr = _as_array_expr(inner_state, fexpr.fexpr, target_dimensions)
    state = set_random_variables(state, outer_rvs)
    state, aexpr
end

function _as_array_expr(state, fexpr::StochasticMean, target_dimensions)
    trs = trajectories_number(fexpr.fexpr)
    target_dimensions = [target_dimensions..., KnownStochasticDimension(trs)]
    state, aexpr = _as_array_expr(state, fexpr.fexpr, target_dimensions)
    state, sum(aexpr, length(target_dimensions)) ./ trs
end

immutable KnownStochasticDimension <: KnownDimension
    id :: Symbol
    trajectories :: Int64
    KnownStochasticDimension(trajectories) = new(gensym(), trajectories)
end

function _as_array_expr(state, fexpr::RandomVariable, target_dimensions)

    if haskey(state.random_variables, fexpr)
        return state, state.random_variables[fexpr]
    end

    if fexpr.distribution == :random_normal
        sdims = find(dim -> typeof(dim) == KnownStochasticDimension, target_dimensions)
        @assert length(sdims) == 1
        random_dimensions = [fexpr.dimensions..., target_dimensions[sdims[1]]]

        shape = [
            typeof(x) == KnownStochasticDimension ? x.trajectories : length(vertices(grid(x)))
            for x in random_dimensions]

        if random_states_set(state)
            rs = output_random_state(state)
        else
            rs = unknown_random_state()
            state = set_random_state(state, rs)
        end

        # TODO: we need to either use the known mean/std here,
        # or expose the basic random_normal(rs, shape...) from Arrays.
        res, new_rs = FluxionArrays.random_normal(rs, 0, 1, shape...)

        res = rearrange_dims(res, random_dimensions, target_dimensions)

        state = update_random_state(state, new_rs)
        state = add_random_variable(state, fexpr, res)

        return state, res
    end

end

function _as_array_expr(state, fexpr::FieldExprNode, target_dimensions)
    const broadcasted_operators = Set([:+, :-, :*, :/, :^])

    op = fexpr.head
    state, args = map_accum(_as_array_expr, state, fexpr.args, target_dimensions)
    if in(op, broadcasted_operators)
        res = broadcast(getfield(Base, op), args...)
    else
        res = getfield(Base, op)(args...)
    end
    state, res
end

function as_array_expr(state, fexpr::FieldExpr, target_dimensions)

    state, aexpr = _as_array_expr(state, fexpr, target_dimensions)

    if is_stochastic(fexpr)
        tr_num = trajectories_number(fexpr)
    else
        tr_num = nothing
    end
    target_shape = [
        dim === STOCHASTIC_DIMENSION ? tr_num : length(vertices(grid(dim)))
        for dim in target_dimensions]

    shape = size(aexpr)
    @assert !isa(shape, ArrayExpr)
    shape = collect(shape)

    aexpr = ensure_shape(aexpr, shape, target_shape)

    state, aexpr
end

immutable WrappedArrayFunction
    array_function
    uses_random_state :: Bool
end

function wrapped_array_function(returns, parameters)

    array_returns = []

    state = ArrayExprBuildingState()

    for (fexpr, result_dimensions) in returns
        state, aexpr = as_array_expr(state, fexpr, result_dimensions)
        push!(array_returns, aexpr)
    end

    array_parameters = [
        get(state.unknowns, parameter, unknown_array())
        for parameter in parameters]

    uses_random_state = !isnull(state.random_states)

    if uses_random_state
        input_rs, output_rs = get(state.random_states)
        array_parameters = [input_rs, array_parameters...]
        array_returns = [output_rs, array_returns...]
    end

    afunc = array_function(array_returns, array_parameters)

    WrappedArrayFunction(afunc, uses_random_state)
end

immutable WrappedBackendArrayFunction
    backend_array_function
    uses_random_state :: Bool
end

function compile_wrapped_array_function(backend::ArrayBackend, afunc::WrappedArrayFunction)
    WrappedBackendArrayFunction(
        compile_array_function(backend, afunc.array_function), afunc.uses_random_state)
end

import FluxionArrays: array_backend

array_backend(wbafunc::WrappedBackendArrayFunction) = array_backend(wbafunc.backend_array_function)

function (wcafunc::WrappedBackendArrayFunction)(args...; seed=nothing)
    if wcafunc.uses_random_state
        backend = array_backend(wcafunc.backend_array_function)
        rs = random_state(backend, seed=seed)
        rs, result = wcafunc.backend_array_function(rs, args...)
        return result
    else
        wcafunc.backend_array_function(args...)
    end
end

immutable FieldFunction
    wrapped_array_function
    output_dimensions
end

function field_function(returns, parameters)
    wafunc = wrapped_array_function(returns, parameters)
    output_dimensions = [rd for (_, rd) in returns]
    FieldFunction(wafunc, output_dimensions)
end

immutable BackendFieldFunction
    wrapped_backend_array_function :: WrappedBackendArrayFunction
    output_dimensions
end

function compile_field_function(backend::ArrayBackend, ffunc::FieldFunction)
    BackendFieldFunction(
        compile_wrapped_array_function(backend, ffunc.wrapped_array_function),
        ffunc.output_dimensions)
end

unwrap_field(x::Number) = x
unwrap_field(x::KnownField) = x.data

function (bffunc::BackendFieldFunction)(args...; seed=nothing)

    args = map(unwrap_field, args)

    if bffunc.wrapped_backend_array_function.uses_random_state
        backend = array_backend(bffunc.wrapped_backend_array_function)
        rs = random_state(backend, seed=seed)
        results = bcfunc.wrapped_backend_array_function(rs, args...)
        results = results[2:end]
    else
        results = bffunc.wrapped_backend_array_function(args...)
        if length(bffunc.output_dimensions) == 1
            results = [results]
        end
    end

    wrapped_results = []
    for (result, odims) in zip(results, bffunc.output_dimensions)
        push!(wrapped_results, known_field(result, odims))
    end

    if length(wrapped_results) == 1
        return wrapped_results[1]
    else
        return tuple(wrapped_results)
    end
end


text:
using Revise
using Juno
using ShaleDrillingLikelihood
using ShaleDrillingLikelihood.SDLParameters

# detect if using SLURM
if "SLURM_JOBID" in keys(ENV)
    SLURM_JOBID = ENV["SLURM_JOBID"]
    using ClusterManagers
else
    SLURM_JOBID = ""
end

using CountPlus, Distributed, JLD2
using ShaleDrillingLikelihood: value_function, EVobj, cost
using SparseArrays: nonzeros
using Formatting: generate_formatter
using Dates: today

num_nonzeros(x) = length(nonzeros(x))
intstr = generate_formatter("%'d")

using ShaleDrillingLikelihood: learn, royalty, constr, tech, tax,
    PerfectInfo, MaxLearning, drill,
    simulationPrimitives_information,
    doSimulations, average_cost_df,
    Theta_NoTech, theta_revenue,
    time_idx, DateQuarter

# ------------------- number of simulations ----------------------

pargs = parse_commandline_counterfactuals()
print_parsed_args_counterfactuals(pargs)

DO_PAR         = !pargs["noPar"]
JLD2FILE       = pargs["jld2"]
DATE_STOP      = pargs["dateStop"]
TECH_YEAR_ZERO = pargs["techYearZero"]
RFILEDIR       = pargs["rFileDir"]

DO_PAR = false

# --------------- create data ---------------

if "SLURM_JOBID" in keys(ENV)
    DATADIR = "/home/magerton/haynesville/intermediate_data"
else
    DATADIR = "E:/projects/haynesville/intermediate_data"
end

println_time_flush("Loading results from $JLD2FILE")
file = jldopen(JLD2FILE, "r")
    DATAPATH = file["DATAPATH"]
    M        = file["M"]
    ddmnovf  = file["ddm_novf"]
    theta    = file["theta1"]
close(file)
println_time_flush("Done!")

REWARD = reward(ddmnovf)
ddm = DynamicDrillModel(ddmnovf, REWARD)

# --------------- create data ---------------

# load in data from disk
rdatapath = joinpath(DATADIR, DATAPATH)
println_time_flush("loading $rdatapath")

data_royalty    = DataRoyalty(REWARD, rdatapath)
data_produce    = DataProduce(REWARD, rdatapath)
data_drill_prim = DataDrillPrimitive(REWARD, rdatapath)
data_drill      = DataDrill(data_drill_prim, ddm)

evdims = size(EVobj(value_function(ddm)))
println_time_flush("EV dimension is $evdims, implying $(intstr(prod(evdims))) states")

# full dataset
dataset_full = DataSetofSets(data_drill, data_royalty, data_produce, CoefLinks(data_drill))

println_time_flush("Data created")

let thts = split_thetas(dataset_full, theta)
    println("theta_drill = $(round.(thts[1]; digits=3))")
    println("theta_roy   = $(round.(thts[2]; digits=3))")
    println("theta_pdxn  = $(round.(thts[3]; digits=3))")
end

# ------------------- simulations -----------------------

wp = statespace(ddm)
PP = PerpetualProblem(wp)

# ------------------- average costs -----------------------

cost_df = average_cost_df(dataset_full, theta)

# ------------------- simulation info -----------------------
R = ShaleDrillingLikelihood.revenue(REWARD)
C = ShaleDrillingLikelihood.cost(REWARD)
E = ShaleDrillingLikelihood.extend(REWARD)

DrillRev(lrn, roy) = DrillingRevenue(constr(R), tech(R), tax(R), lrn, roy)
DrillRwrd(lrn, roy) = DrillReward(DrillRev(lrn,roy), C, E)

theta_notech = Theta_NoTech(dataset_full, theta, TECH_YEAR_ZERO)
TSTOP = time_idx(zchars(drill(dataset_full)), DateQuarter(DATE_STOP))

simlist = [
    # (DrillRwrd( learn(R),      royalty(R) ),   wp, theta, ),           # "Baseline"
    (reward(ddm), wp, theta),
    # (DrillRwrd( learn(R),      royalty(R) ),   wp, theta_notech, ),    # "Where to drill"
    # (DrillRwrd( NoLearn(),     royalty(R) ),   wp, theta, ),           # "No update ($\\psi_{it} = \\psi_i^0$)"
    # (DrillRwrd( PerfectInfo(), royalty(R) ),   wp, theta, ),           # "Perfect information ($\\psi_{it} = \\psi_i^1$)"
    # (DrillRwrd( MaxLearning(), royalty(R) ),   wp, theta, ),           # "Uninformative signals ($\\rho = 0$)"
    # (DrillRwrd( learn(R),      NoRoyalty()),   wp, theta, ),           # "No royalty"
    # (DrillRwrd( learn(R),      royalty(R) ),   PP, theta, ),           # "No expiration"
    # (DrillRwrd( learn(R),      NoRoyalty()),   PP, theta, ),           # "Ownership"
    # (DrillRwrd(NoLearn(),      NoRoyalty()  ), PP, theta_notech, ),    # "Price only"
    # (DrillRwrd(NoLearn(),      NoRoyalty()  ), PP, theta, ),           # "Price + tech (no distortions)"
    # (DrillRwrd(NoLearn(),      WithRoyalty()), wp, theta_notech, ),    # "Price + leases"
    # (DrillRwrd(Learn(),        NoRoyalty()  ), PP, theta_notech, ),    # "Price + learning"
    # (DrillRwrd(NoLearn(),      WithRoyalty()), PP, theta, ),           # "How to drill"
]

simlist_meta = [(a,b,copy(theta_revenue(dataset_full, c))) for (a,b,c) in simlist]

simulations_meta_info = simulationPrimitives_information(simlist_meta)

# ------------------- simulations -----------------------

if DO_PAR
    pids = start_up_workers(ENV)
    @everywhere using ShaleDrillingLikelihood
    println_time_flush("Library loaded on workers")
end

# df_d, df_D = doSimulations(dataset_full, simlist, TSTOP, M)

# generate smaller DataDrill to transfer to workers
using ShaleDrillingLikelihood: SharedSimulations, dataFramesForSimulations, SimulationTmp
using ShaleDrillingLikelihood: revenue, theta_drill, vw_revenue, tech, learn, royalty, SimulationPrimitives
rwrd = reward(ddm)
k=1
datafull = dataset_full
Tstop = TSTOP

datadrill = drill(datafull)
ddm_novf = DDM_NoVF(_model(datadrill))
datadrill_dso = DataDrillStartOnly(datadrill)
datadrill_bare= DataDrill(datadrill_dso, ddm_novf)

# DataFull set for transfer
data_bare = DataSetofSets(datafull, datadrill_bare)

# generate common variables
sharesim = SharedSimulations(datadrill_bare)
simMat = SimulationDraws(data_bare, M)

# DataFrames for results
df_t, df_Tstop = dataFramesForSimulations(simlist, datadrill_bare, Tstop)

fill!(sharesim, 0)
rev = revenue(rwrd)
thet_d = theta_drill(datafull, theta)
thet_r = vw_revenue(rwrd, thet_d)
println_time_flush("Simulation $k of $(length(simlist))")
println("\tTech = $(tech(rev))")
println("\tLearn = $(learn(rev))")
println("\tRoyalty = $(royalty(rev))")
println("\tStatespace = $(wp)")
println("\tTheta_rev = $(thet_r)")

# @eval @everywhere set_g_SimulationPrimitives($rwrd, $wp, $Tstop, $theta)
# simprim = get_g_SimulationPrimitives()
simprim = SimulationPrimitives(data_bare, simMat, rwrd, wp, Tstop, theta, sharesim)
simtmp = SimulationTmp(simprim)
simulate_unit!(simprim, 9, true)

using Test
@test sum(simtmp.sa) ≈ 1
@test sum(simtmp.sb) ≈ 1
for x in sum(simtmp.Pprime; dims=1)
    @test x ≈ 1 || x == 0
end

simtmp.sa .- simtmp.sb
@test simtmp.sa == simtmp.Pprime * simtmp.sb


text:
using Random

struct VDPInitDist end
sampletype(::Type{VDPInitDist}) = TagState
function rand(rng::AbstractRNG, d::VDPInitDist)
    # return TagState([0.0, 0.0], 8.0*rand(rng, 2)-4.0)
    return TagState([0.0, 0.0], 8.0*rand(rng, 2) .- 4.0)

end

initialstate_distribution(::VDPTagProblem) = VDPInitDist()


text:
struct StorePutKey{T} <: ResourceKey
  priority :: Int
  id :: UInt
  item :: T
  StorePutKey{T}(priority, id, item) where T = new(priority, id, item)
end

struct StoreGetKey <: ResourceKey
  priority :: Int
  id :: UInt
  filter :: Function
end

mutable struct Store{T} <: AbstractResource
  env :: Environment
  capacity :: UInt
  items :: Set{T}
  seid :: UInt
  put_queue :: DataStructures.PriorityQueue{Put, StorePutKey{T}}
  get_queue :: DataStructures.PriorityQueue{Get, StoreGetKey}
  function Store{T}(env::Environment; capacity::UInt=typemax(UInt)) where {T}
    new(env, capacity, Set{T}(), zero(UInt), DataStructures.PriorityQueue{Put, StorePutKey{T}}(), DataStructures.PriorityQueue{Get, StoreGetKey}())
  end
end

function put(sto::Store{T}, item::T; priority::Int=0) where T
  put_ev = Put(sto.env)
  sto.put_queue[put_ev] = StorePutKey{T}(priority, sto.seid+=one(UInt), item)
  @callback trigger_get(put_ev, sto)
  trigger_put(put_ev, sto)
  put_ev
end

get_any_item(::T) where T = true

function get(sto::Store{T}, filter::Function=get_any_item; priority::Int=0) where T
  get_ev = Get(sto.env)
  sto.get_queue[get_ev] = StoreGetKey(priority, sto.seid+=one(UInt), filter)
  @callback trigger_put(get_ev, sto)
  trigger_get(get_ev, sto)
  get_ev
end

function do_put(sto::Store{T}, put_ev::Put, key::StorePutKey{T}) where {T}
  if length(sto.items) < sto.capacity
    push!(sto.items, key.item)
    schedule(put_ev)
  end
  false
end

function do_get(sto::Store{T}, get_ev::Get, key::StoreGetKey) where {T}
  for item in sto.items
    if key.filter(item)
      delete!(sto.items, item)
      schedule(get_ev; value=item)
      break
    end
  end
  true
end


text:
using Plots, DataFrames, Flux, Plots, Random, ReverseDiff
using Turing # v0.20

# create half-moon data:
function make_moons(n_samples::Int = 100)

    n = n_samples

    n_samples_out = trunc(Int, round(n_samples / 2))
    n_samples_in = n_samples - n_samples_out

    outer_circ_x = cos.(LinRange(0, π, n_samples_out))
    outer_circ_y = sin.(LinRange(0, π, n_samples_out))
    inner_circ_x = 1 .- cos.(LinRange(0, π, n_samples_in))
    inner_circ_y = 1 .- sin.(LinRange(0, π, n_samples_in)) .- 0.5

    X = hcat(vcat(outer_circ_x, inner_circ_x), vcat(outer_circ_y, inner_circ_y))
    Y = hcat(zeros(Int, 1, 2 .* n_samples_out), ones(Int, 1, 2 .* n_samples_in))

    return X, Y, n

end

X, Y, n = make_moons()

# rearrange half-moon data:
n_2 = trunc(Int, n/2)
xscatter_0 = X[1:n_2]
xscatter_1 = X[1+n_2:n]
yscatter_0 = X[1+n:n+n_2]
yscatter_1 = X[1+n+n_2:end]

# plot half-moon data:
plot(
    xscatter_0, yscatter_0,
    seriestype = :scatter,
    legend = true,
    label = "Class 0",
    title = "half-moons problem"
)
plot!(
    xscatter_1, yscatter_1,
    seriestype = :scatter,
    legend = true,
    label = "Class 1"
)

# optional: half-moon data in a dataframe:
half_moon_df = DataFrame(
    X = vcat(xscatter_0, xscatter_1),
    Y = vcat(yscatter_0, yscatter_1),
    Class = vcat(zeros(Int, n_2), ones(Int, n_2))
)

# Turing settings:
Turing.setprogress!(true)
Turing.setadbackend(:reversediff)

# architecture:
network_shape = [
    (5,2, :tanh),
    (5,5, :tanh),
    (2,5, :σ)]

# Regularization, parameter variance, and total number of parameters.
alpha = 0.09
sig = sqrt(1.0 / alpha)
num_params = sum([i * o + i for (i, o, _) in network_shape])

# This modification of the unpack function generates a series of vectors given a network shape.
function unpack(θ::AbstractVector, network_shape::AbstractVector)
    index = 1
    weights = []
    biases = []
    for layer in network_shape
        rows, cols, _ = layer
        size = rows * cols
        last_index_w = size + index - 1
        last_index_b = last_index_w + rows
        push!(weights, reshape(θ[index:last_index_w], rows, cols))
        push!(biases, reshape(θ[last_index_w+1:last_index_b], rows))
        index = last_index_b + 1
    end
    return weights, biases
end

# Generate an abstract neural network given a shape, and return a prediction.
function nn_forward(x, θ::AbstractVector, network_shape::AbstractVector)
    weights, biases = unpack(θ, network_shape)
    layers = []
    for i in eachindex(network_shape)
        push!(layers, Dense(weights[i],
            biases[i],
            eval(network_shape[i][3])))
    end
    nn = Chain(layers...)
    return nn(x)
end

# rearrange half-moon data yet again:
xscatter = vcat(xscatter_0, xscatter_1)
yscatter = vcat(yscatter_0, yscatter_1)
xs = Array([[xscatter[i]; yscatter[i]] for i = 1:n])
ts = vcat(zeros(Int, n_2), ones(Int, n_2))

# General Turing specification for a BNN model.
@model bayes_nn_general(xs, ts, network_shape, num_params) = begin
    θ ~ MvNormal(zeros(num_params), sig .* ones(num_params))
    preds = nn_forward(xs, θ, network_shape)
    for i = 1:length(ts)
        ts[i] ~ Bernoulli(preds[i])
    end
end

# Perform inference.
num_samples = 2000
ch2 = sample(bayes_nn_general(hcat(xs...), ts, network_shape, num_params), NUTS(0.65), num_samples);

# This function makes predictions based on network shape.
function nn_predict(x, theta, num, network_shape)
    mean([nn_forward(x, theta[i,:], network_shape)[1] for i in 1:10:num])
end;

# Extract the θ parameters from the sampled chain.
params2 = MCMCChains.group(ch2, :θ).value

# classification plot
plot(
    xscatter_0, yscatter_0,
    seriestype = :scatter,
    label = "Class 0",
    title = "half-moons problem"
)
plot!(
    xscatter_1, yscatter_1,
    seriestype = :scatter,
    legend = :bottomright,
    label = "Class 1"
)
x_range = collect(range(-2.5,stop=3.5,length=200))
y_range = collect(range(-1.2,stop=1.7,length=200))
Z = [nn_predict([x, y], params2, length(ch2), network_shape)[1] for x=x_range, y=y_range]
contour!(x_range, y_range, Z)

max_stars_repo_name: cb-g/machine-learning-in-finance

max_issues_repo_name: cb-g/machine-learning-in-finance

max_forks_repo_name: cb-g/machine-learning-in-finance