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

using Logging
import Logging: Info,
    shouldlog, handle_message, min_enabled_level, catch_exceptions
import Base: occursin

#-------------------------------------------------------------------------------
# Log records
struct LogRecord
    level
    message
    _module
    group
    id
    file
    line
    kwargs
end
LogRecord(args...; kwargs...) = LogRecord(args..., kwargs)

struct Ignored ; end

#-------------------------------------------------------------------------------
# Logger with extra test-related state
mutable struct TestLogger <: AbstractLogger
    logs::Vector{LogRecord}
    min_level::LogLevel
    catch_exceptions::Bool
    shouldlog_args
end

TestLogger(; min_level=Info, catch_exceptions=false) = TestLogger(LogRecord[], min_level, catch_exceptions, nothing)
min_enabled_level(logger::TestLogger) = logger.min_level

function shouldlog(logger::TestLogger, level, _module, group, id)
    logger.shouldlog_args = (level, _module, group, id)
    true
end

function handle_message(logger::TestLogger, level, msg, _module,
                        group, id, file, line; kwargs...)
    push!(logger.logs, LogRecord(level, msg, _module, group, id, file, line, kwargs))
end

# Catch exceptions for the test logger only if specified
catch_exceptions(logger::TestLogger) = logger.catch_exceptions

function collect_test_logs(f; kwargs...)
    logger = TestLogger(; kwargs...)
    value = with_logger(f, logger)
    logger.logs, value
end

#-------------------------------------------------------------------------------
# Log testing tools

# Failure result type for log testing
mutable struct LogTestFailure <: Result
    orig_expr
    source::Union{Nothing,LineNumberNode}
    patterns
    logs
end
function Base.show(io::IO, t::LogTestFailure)
    printstyled(io, "Log Test Failed"; bold=true, color=Base.error_color())
    print(io, " at ")
    printstyled(io, t.source.file, ":", t.source.line, "\n"; bold=true, color=:default)
    println(io, "  Expression: ", t.orig_expr)
    println(io, "  Log Pattern: ", join(t.patterns, " "))
    println(io, "  Captured Logs: ")
    for l in t.logs
        println(io, "    ", l)
    end
end

# Patch support for LogTestFailure into Base.Test test set types
# TODO: Would be better if `Test` itself allowed us to handle this more neatly.
function record(::FallbackTestSet, t::LogTestFailure)
    println(t)
    throw(FallbackTestSetException("There was an error during testing"))
end

function record(ts::DefaultTestSet, t::LogTestFailure)
    if myid() == 1
        printstyled(ts.description, ": ", color=:white)
        print(t)
        Base.show_backtrace(stdout, scrub_backtrace(backtrace()))
        println()
    end
    # Hack: convert to `Fail` so that test summarization works correctly
    push!(ts.results, Fail(:test, t.orig_expr, t.logs, nothing, t.source))
    t
end

"""
    @test_logs [log_patterns...] [keywords] expression

Collect a list of log records generated by `expression` using
`collect_test_logs`, check that they match the sequence `log_patterns`, and
return the value of `expression`.  The `keywords` provide some simple filtering
of log records: the `min_level` keyword controls the minimum log level which
will be collected for the test, the `match_mode` keyword defines how matching
will be performed (the default `:all` checks that all logs and patterns match
pairwise; use `:any` to check that the pattern matches at least once somewhere
in the sequence.)

The most useful log pattern is a simple tuple of the form `(level,message)`.
A different number of tuple elements may be used to match other log metadata,
corresponding to the arguments to passed to `AbstractLogger` via the
`handle_message` function: `(level,message,module,group,id,file,line)`.
Elements which are present will be matched pairwise with the log record fields
using `==` by default, with the special cases that `Symbol`s may be used for
the standard log levels, and `Regex`s in the pattern will match string or
Symbol fields using `occursin`.

# Examples

Consider a function which logs a warning, and several debug messages:

    function foo(n)
        @info "Doing foo with n=\$n"
        for i=1:n
            @debug "Iteration \$i"
        end
        42
    end

We can test the info message using

    @test_logs (:info,"Doing foo with n=2") foo(2)

If we also wanted to test the debug messages, these need to be enabled with the
`min_level` keyword:

    @test_logs (:info,"Doing foo with n=2") (:debug,"Iteration 1") (:debug,"Iteration 2") min_level=Debug foo(2)

The macro may be chained with `@test` to also test the returned value:

    @test (@test_logs (:info,"Doing foo with n=2") foo(2)) == 42

"""
macro test_logs(exs...)
    length(exs) >= 1 || throw(ArgumentError("""`@test_logs` needs at least one arguments.
                               Usage: `@test_logs [msgs...] expr_to_run`"""))
    patterns = Any[]
    kwargs = Any[]
    for e in exs[1:end-1]
        if e isa Expr && e.head == :(=)
            push!(kwargs, esc(Expr(:kw, e.args...)))
        else
            push!(patterns, esc(e))
        end
    end
    expression = exs[end]
    orig_expr = QuoteNode(expression)
    sourceloc = QuoteNode(__source__)
    Base.remove_linenums!(quote
        let testres=nothing, value=nothing
            try
                didmatch,logs,value = match_logs($(patterns...); $(kwargs...)) do
                    $(esc(expression))
                end
                if didmatch
                    testres = Pass(:test, nothing, nothing, value)
                else
                    testres = LogTestFailure($orig_expr, $sourceloc,
                                             $(QuoteNode(exs[1:end-1])), logs)
                end
            catch e
                testres = Error(:test_error, $orig_expr, e, catch_backtrace(), $sourceloc)
            end
            Test.record(Test.get_testset(), testres)
            value
        end
    end)
end

function match_logs(f, patterns...; match_mode::Symbol=:all, kwargs...)
    logs,value = collect_test_logs(f; kwargs...)
    if match_mode == :all
        didmatch = length(logs) == length(patterns) &&
            all(occursin(p, l) for (p,l) in zip(patterns, logs))
    elseif match_mode == :any
        didmatch = all(any(occursin(p, l) for l in logs) for p in patterns)
    end
    didmatch,logs,value
end

# TODO: Use a version of parse_level from stdlib/Logging, when it exists.
function parse_level(level::Symbol)
    if      level == :belowminlevel  return  Logging.BelowMinLevel
    elseif  level == :debug          return  Logging.Debug
    elseif  level == :info           return  Logging.Info
    elseif  level == :warn           return  Logging.Warn
    elseif  level == :error          return  Logging.Error
    elseif  level == :abovemaxlevel  return  Logging.AboveMaxLevel
    else
        throw(ArgumentError("Unknown log level $level"))
    end
end

logfield_contains(a, b) = a == b
logfield_contains(a, r::Regex) = occursin(r, a)
logfield_contains(a::Symbol, r::Regex) = occursin(r, String(a))
logfield_contains(a::LogLevel, b::Symbol) = a == parse_level(b)
logfield_contains(a, b::Ignored) = true

function occursin(pattern::Tuple, r::LogRecord)
    stdfields = (r.level, r.message, r._module, r.group, r.id, r.file, r.line)
    all(logfield_contains(f, p) for (f, p) in zip(stdfields[1:length(pattern)], pattern))
end

"""
    @test_deprecated [pattern] expression

When `--depwarn=yes`, test that `expression` emits a deprecation warning and
return the value of `expression`.  The log message string will be matched
against `pattern` which defaults to `r"deprecated"i`.

When `--depwarn=no`, simply return the result of executing `expression`.  When
`--depwarn=error`, check that an ErrorException is thrown.

# Examples

```
# Deprecated in julia 0.7
@test_deprecated num2hex(1)

# The returned value can be tested by chaining with @test:
@test (@test_deprecated num2hex(1)) == "0000000000000001"
```
"""
macro test_deprecated(exs...)
    1 <= length(exs) <= 2 || throw(ArgumentError("""`@test_deprecated` expects one or two arguments.
                               Usage: `@test_deprecated [pattern] expr_to_run`"""))
    pattern = length(exs) == 1 ? r"deprecated"i : esc(exs[1])
    expression = esc(exs[end])
    res = quote
        dw = Base.JLOptions().depwarn
        if dw == 2
            # TODO: Remove --depwarn=error if possible and replace with a more
            # flexible mechanism so we don't have to do this.
            @test_throws ErrorException $expression
        elseif dw == 1
            @test_logs (:warn, $pattern, Ignored(), :depwarn) match_mode=:any $expression
        else
            $expression
        end
    end
    # Propagate source code location of @test_logs to @test macro
    # FIXME: Use rewrite_sourceloc!() for this - see #22623
    res.args[4].args[2].args[2].args[2] = __source__
    res.args[4].args[3].args[2].args[2].args[2] = __source__
    res
end

max_stars_repo_name: djsegal/julia-fork

max_issues_repo_name: djsegal/julia-fork

max_forks_repo_name: djsegal/julia-fork


text:
"""
    scalar(t::AbstractTensorMap{S}) where {S<:IndexSpace}

Change the tensor map which is from 1D space to 1D space to its scalar value.
"""
scalar(t::AbstractTensorMap{S}) where {S<:IndexSpace} =
    dim(codomain(t)) == dim(domain(t)) == 1 ?
        first(blocks(t))[2][1, 1] : throw(SpaceMismatch())

"""
    _add_trivial_kernel!(α, tsrc::AbstractTensorMap, β, tdst::AbstractTensorMap,
                            p1::IndexTuple, p2::IndexTuple, fusiontreemap)

`tsrc` and `tdst` are TrivialTensorMap, thus `tsrc[]` and `tdst[]` gives the data of
the tensor map in a form of multidimensional array. `pdata = (p1..., p2...)` gives the
permutation of the indices of the tensor map `tsrc`. After permutation `tsrc` changes
to `tsrc2` and has the same spaces with `tdst`.

After running, `tdst` will be replaced in-place by the tensor map `tsrc2*α + tdst*β`.

Note that the parameter `fusiontreemap` is useless here.
"""
function _add_trivial_kernel!(α, tsrc::AbstractTensorMap, β, tdst::AbstractTensorMap,
                                p1::IndexTuple, p2::IndexTuple, fusiontreemap)

    pdata = (p1..., p2...)
    axpby!(α, permutedims(tsrc[], pdata), β, tdst[])
    return nothing
end

"""
    _addabelianblock!(α, tsrc::AbstractTensorMap, β, tdst::AbstractTensorMap,
        p1::IndexTuple, p2::IndexTuple, f1::FusionTree, f2::FusionTree, fusiontreemap)

Only works for tensor maps with UniqueFusion.

The fusion trees of the `tdst` may not be the same with `tsrc`, and the fusion trees of
`tsrc` can be transformed to that of `tdst` by `fusiontreemap`.

Consider the fusion tree `(f1,f2)` in `tsrc` with data `tsrc[f1, f2]`.

Applying the `fusiontreemap` on `(f1,f2)`, we obtain `(f1′, f2′) => coeff`, where
`(f1′, f2′)` is also a fusion tree of `tdst`.

At the same time the data `tsrc[f1, f2]` should be permuted by `pdata = (p1..., p2...)` to
be consistent with the structure of the new fusion tree `(f1′, f2′)`.

After running, the block `tdst[f1′, f2′]` is replaced in-place by
`α*coeff*permutedims(tsrc[f1, f2], pdata)+β*tdst[f1′, f2′]`.
"""
function _addabelianblock!(α, tsrc::AbstractTensorMap, β, tdst::AbstractTensorMap,
        p1::IndexTuple, p2::IndexTuple, f1::FusionTree, f2::FusionTree, fusiontreemap)

    (f1′, f2′), coeff = first(fusiontreemap(f1, f2))
    pdata = (p1..., p2...)
    @inbounds axpby!(α*coeff, permutedims(tsrc[f1, f2], pdata), β, tdst[f1′, f2′])
end

"""
    _add_abelian_kernel!(α, tsrc::AbstractTensorMap, β, tdst::AbstractTensorMap,
                            p1::IndexTuple, p2::IndexTuple, fusiontreemap)

Only works for tensor maps with UniqueFusion.

For every fusion tree in `tsrc`, apply the operation `_addabelianblock!`.
"""
function _add_abelian_kernel!(α, tsrc::AbstractTensorMap, β, tdst::AbstractTensorMap,
                                p1::IndexTuple, p2::IndexTuple, fusiontreemap)
    if Threads.nthreads() > 1
        nstridedthreads = Strided.get_num_threads()
        Strided.set_num_threads(1)
        Threads.@sync for (f1, f2) in fusiontrees(tsrc)
            Threads.@spawn _addabelianblock!(α, tsrc, β, tdst, p1, p2, f1, f2,
                                                fusiontreemap)
        end
        Strided.set_num_threads(nstridedthreads)
    else # debugging is easier this way
        for (f1, f2) in fusiontrees(tsrc)
            _addabelianblock!(α, tsrc, β, tdst, p1, p2, f1, f2, fusiontreemap)
        end
    end
    return nothing
end

"""
    _add_general_kernel!(α, tsrc::AbstractTensorMap, β, tdst::AbstractTensorMap,
                            p1::IndexTuple, p2::IndexTuple, fusiontreemap)

Works for general tensor maps.

Same operations on each fusion trees.
"""
function _add_general_kernel!(α, tsrc::AbstractTensorMap, β, tdst::AbstractTensorMap,
                                p1::IndexTuple, p2::IndexTuple, fusiontreemap)

    pdata = (p1..., p2...)
    if iszero(β)
        fill!(tdst, β)
    elseif β != 1
        mul!(tdst, β, tdst)
    end
    for (f1, f2) in fusiontrees(tsrc)
        for ((f1′, f2′), coeff) in fusiontreemap(f1, f2)
            @inbounds axpy!(α*coeff, permutedims(tsrc[f1, f2], pdata), tdst[f1′, f2′])
        end
    end
    return nothing
end

const _add_kernels = (_add_trivial_kernel!, _add_abelian_kernel!, _add_general_kernel!)

"""
    _add!(α, tsrc::AbstractTensorMap{S}, β, tdst::AbstractTensorMap{S, N₁, N₂},
            p1::IndexTuple{N₁}, p2::IndexTuple{N₂}, fusiontreemap) where {S, N₁, N₂}

Make the general add operation between two tensor maps.

The fusion trees of the `tdst` may not be the same with `tsrc`, and the fusion trees of
`tsrc` can be transformed to that of `tdst` by `fusiontreemap`.

For each fusion tree `(f1,f2)` of `tsrc`, apply the manipulation `fusiontreemap`, which
give new fusion trees `(f1′, f2′)`, which are also fusion trees of `tdst`, and the
corresponding coefficients `coeff`.

At the same time apply the permutation on the data of the old fusion trees based on
`(p1..., p2...)` to make the data structure consistent with the new fusion trees.

Return `tdst`, whose data for each fusion tree is replaced in-place by
`α*coeff*permutedims(tsrc[f1, f2], pdata)+β*tdst[f1′, f2′]`.
"""
function _add!(α, tsrc::AbstractTensorMap{S}, β, tdst::AbstractTensorMap{S, N₁, N₂},
                p1::IndexTuple{N₁}, p2::IndexTuple{N₂}, fusiontreemap) where {S, N₁, N₂}
    @boundscheck begin
        all(i->space(tsrc, p1[i]) == space(tdst, i), 1:N₁) ||
            throw(SpaceMismatch("tsrc = $(codomain(tsrc))←$(domain(tsrc)),
            tdst = $(codomain(tdst))←$(domain(tdst)), p1 = $(p1), p2 = $(p2)"))
        all(i->space(tsrc, p2[i]) == space(tdst, N₁+i), 1:N₂) ||
            throw(SpaceMismatch("tsrc = $(codomain(tsrc))←$(domain(tsrc)),
            tdst = $(codomain(tdst))←$(domain(tdst)), p1 = $(p1), p2 = $(p2)"))
    end

    # do some kind of dispatch which is compiled away if S is known at compile time,
    # and makes the compiler give up quickly if S is unknown
    I = sectortype(S)
    i = I === Trivial ? 1 : (FusionStyle(I) isa UniqueFusion ? 2 : 3)
    if p1 == codomainind(tsrc) && p2 == domainind(tsrc)
        axpby!(α, tsrc, β, tdst)
    else
        _add_kernel! = _add_kernels[i]
        _add_kernel!(α, tsrc, β, tdst, p1, p2, fusiontreemap)
    end
    return tdst
end

"""
    add_permute!(α, tsrc::AbstractTensorMap{S}, β, tdst::AbstractTensorMap{S, N₁, N₂},
                    p1::IndexTuple{N₁}, p2::IndexTuple{N₂}) where {S, N₁, N₂}

Work for the case that `tsrc` are related to `tdst` by permutation `(p1...,p2...)`.

Replace `tdst` with `permute(tsrc)*α + tdst*β`.
"""
@propagate_inbounds function add_permute!(α, tsrc::AbstractTensorMap{S},
                                            β, tdst::AbstractTensorMap{S, N₁, N₂},
                                            p1::IndexTuple{N₁},
                                            p2::IndexTuple{N₂}) where {S, N₁, N₂}

    _add!(α, tsrc, β, tdst, p1, p2, (f1, f2)->permute(f1, f2, p1, p2))
end

"""
    add_braid!(α, tsrc::AbstractTensorMap{S}, β, tdst::AbstractTensorMap{S, N₁, N₂},
            p1::IndexTuple{N₁}, p2::IndexTuple{N₂}, levels::IndexTuple) where {S, N₁, N₂}

Work for the case that `tsrc` are related to `tdst` by braiding according to `(p1...,p2...)`
and `levels`.

Replace `tdst` with `braid(tsrc)*α + tdst*β`.
"""
@propagate_inbounds function add_braid!(α, tsrc::AbstractTensorMap{S},
                                        β, tdst::AbstractTensorMap{S, N₁, N₂},
                                        p1::IndexTuple{N₁},
                                        p2::IndexTuple{N₂},
                                        levels::IndexTuple) where {S, N₁, N₂}

    length(levels) == numind(tsrc) ||
        throw(ArgumentError("incorrect levels $levels for tensor map
                                $(codomain(tsrc)) ← $(domain(tsrc))"))

    levels1 = TupleLabXD.getindices(levels, codomainind(tsrc))
    levels2 = TupleLabXD.getindices(levels, domainind(tsrc))
    _add!(α, tsrc, β, tdst, p1, p2, (f1, f2)->braid(f1, f2, levels1, levels2, p1, p2))
end

"""
    add_transpose!(α, tsrc::AbstractTensorMap{S}, β, tdst::AbstractTensorMap{S, N₁, N₂},
                    p1::IndexTuple{N₁}, p2::IndexTuple{N₂}) where {S, N₁, N₂}

Work for the case that `tsrc` are related to `tdst` by transpose, which can be realized
by repartition and cyclic permutation.

Replace `tdst` with `transpose(tsrc)*α + tdst*β`.
"""
@propagate_inbounds function add_transpose!(α, tsrc::AbstractTensorMap{S},
                                            β, tdst::AbstractTensorMap{S, N₁, N₂},
                                            p1::IndexTuple{N₁},
                                            p2::IndexTuple{N₂}) where {S, N₁, N₂}

    _add!(α, tsrc, β, tdst, p1, p2, (f1, f2)->transpose(f1, f2, p1, p2))
end

"""
    add!(α, tsrc::AbstractTensorMap{S}, β, tdst::AbstractTensorMap{S},
            p1::IndexTuple, p2::IndexTuple) where {S}

Work only for the case with `SymmetricBraiding`, and equivalent to `add_permute!`.

Replace `tdst` with `permute(tsrc)*α + tdst*β`.
"""
@propagate_inbounds function add!(α, tsrc::AbstractTensorMap{S},
                                    β, tdst::AbstractTensorMap{S},
                                    p1::IndexTuple, p2::IndexTuple) where {S}
    I = sectortype(S)
    if BraidingStyle(I) isa SymmetricBraiding
        add_permute!(α, tsrc, β, tdst, p1, p2)
    else
        throw(ArgumentError("add! without levels is defined only if
                                `BraidingStyle(sectortype(...)) isa SymmetricBraiding`"))
    end
end

"""
    add!(α, tsrc::AbstractTensorMap{S}, β, tdst::AbstractTensorMap{S},
            p1::IndexTuple, p2::IndexTuple, levels::IndexTuple) where {S}

Equivalent to `add_braid!`.

Replace `tdst` with `braid(tsrc)*α + tdst*β`.
"""
@propagate_inbounds function add!(α, tsrc::AbstractTensorMap{S},
                                    β, tdst::AbstractTensorMap{S},
                                    p1::IndexTuple, p2::IndexTuple,
                                    levels::IndexTuple) where {S}
    add_braid!(α, tsrc, β, tdst, p1, p2, levels)
end

"""
    trace!(α, tsrc::AbstractTensorMap{S}, β, tdst::AbstractTensorMap{S, N₁, N₂},
            p1::IndexTuple{N₁}, p2::IndexTuple{N₂},
            q1::IndexTuple{N₃}, q2::IndexTuple{N₃}) where {S, N₁, N₂, N₃}

Work only for tensors with symmetric braiding.

Implements `tdst = β*tdst+α*partialtrace(tsrc)` where `tsrc` is permuted and partially
traced, such that the codomain (domain) of `tdst` correspond to the spaces `p1` (`p2`) of
`tsrc`, and indices `q1[i]` are contracted with indices `q2[i]`.
"""
function trace!(α, tsrc::AbstractTensorMap{S}, β, tdst::AbstractTensorMap{S, N₁, N₂},
                p1::IndexTuple{N₁}, p2::IndexTuple{N₂},
                q1::IndexTuple{N₃}, q2::IndexTuple{N₃}) where {S, N₁, N₂, N₃}

    if !(BraidingStyle(sectortype(S)) isa SymmetricBraiding)
        throw(SectorMismatch("only tensors with symmetric braiding rules can be contracted;
                                try `@planar` instead"))
    end
    @boundscheck begin
        all(i->space(tsrc, p1[i]) == space(tdst, i), 1:N₁) ||
            throw(SpaceMismatch("trace: tsrc = $(codomain(tsrc))←$(domain(tsrc)),
                    tdst = $(codomain(tdst))←$(domain(tdst)), p1 = $(p1), p2 = $(p2)"))
        all(i->space(tsrc, p2[i]) == space(tdst, N₁+i), 1:N₂) ||
            throw(SpaceMismatch("trace: tsrc = $(codomain(tsrc))←$(domain(tsrc)),
                    tdst = $(codomain(tdst))←$(domain(tdst)), p1 = $(p1), p2 = $(p2)"))
        all(i->space(tsrc, q1[i]) == dual(space(tsrc, q2[i])), 1:N₃) ||
            throw(SpaceMismatch("trace: tsrc = $(codomain(tsrc))←$(domain(tsrc)),
                    q1 = $(q1), q2 = $(q2)"))
    end

    I = sectortype(S)
    if I === Trivial
        pdata = (p1..., p2...)
        TO._trace!(α, tsrc[], β, tdst[], pdata, q1, q2)
    # elseif FusionStyle(I) isa UniqueFusion
    # TODO: is it worth multithreading UniqueFusion case for traces?
    else
        cod = codomain(tsrc)
        dom = domain(tsrc)
        n = length(cod)
        pdata = (p1..., p2...)
        if iszero(β)
            fill!(tdst, β)
        elseif β != 1
            mul!(tdst, β, tdst)
        end
        r1 = (p1..., q1...)
        r2 = (p2..., q2...)
        for (f1, f2) in fusiontrees(tsrc)
            for ((f1′, f2′), coeff) in permute(f1, f2, r1, r2)
                f1′′, g1 = split(f1′, N₁)
                f2′′, g2 = split(f2′, N₂)
                if g1 == g2
                    coeff *= dim(g1.coupled)/dim(g1.uncoupled[1])
                    for i = 2:length(g1.uncoupled)
                        if !(g1.isdual[i])
                            coeff *= twist(g1.uncoupled[i])
                        end
                    end
                    TO._trace!(α*coeff, tsrc[f1, f2], true, tdst[f1′′, f2′′], pdata, q1, q2)
                end
            end
        end
    end
    return tdst
end

# TODO: contraction with either A or B a rank (1, 1) tensor does not require to
# permute the fusion tree and should therefore be special cased. This will speed
# up MPS algorithms

"""
    contract!(α, A::AbstractTensorMap{S}, B::AbstractTensorMap{S},
                β, C::AbstractTensorMap{S},
                oindA::IndexTuple{N₁}, cindA::IndexTuple,
                oindB::IndexTuple{N₂}, cindB::IndexTuple,
                p1::IndexTuple, p2::IndexTuple,
                syms::Union{Nothing, NTuple{3, Symbol}} = nothing) where {S, N₁, N₂}

Implements `C = β*C+α*contract(A,B)` where `A` and `B` are contracted, such that
the indices `cindA` of `A` are contracted with indices `cindB` of `B`. The open indices
`oindA` of `A` and `oindB` of `B` are permuted such that `C` has codomain (domain)
corresponding to indices `p1` (`p2`) out of `(oindA..., oindB...)`.

Together, `(oindA..., cindA...)` is a permutation of 1 to the number of indices of `A` and
`(oindB..., cindB...)` is a permutation of 1 to the number of indices of `B`.

`length(cindA) == length(cindB)`, and `length(oindA)+length(oindB)` equals the number of
indices of `C` and `(p1..., p2...)` is a permutation of `1` ot the number of indices of `C`.

The final argument `syms` is optional and can be either `nothing`, or a tuple of three
symbols, which are used to identify temporary objects in the cache to be used for permuting
`A`, `B` and `C` so as to perform the contraction as a matrix multiplication.
"""
function contract!(α, A::AbstractTensorMap{S}, B::AbstractTensorMap{S},
                    β, C::AbstractTensorMap{S},
                    oindA::IndexTuple{N₁}, cindA::IndexTuple,
                    oindB::IndexTuple{N₂}, cindB::IndexTuple,
                    p1::IndexTuple, p2::IndexTuple,
                    syms::Union{Nothing, NTuple{3, Symbol}} = nothing) where {S, N₁, N₂}
    # find optimal contraction scheme
    hsp = has_shared_permute
    ipC = TupleLabXD.invperm((p1..., p2...))
    oindAinC = TupleLabXD.getindices(ipC, ntuple(n->n, N₁))
    oindBinC = TupleLabXD.getindices(ipC, ntuple(n->n+N₁, N₂))

    qA = TupleLabXD.sortperm(cindA)
    cindA′ = TupleLabXD.getindices(cindA, qA)
    cindB′ = TupleLabXD.getindices(cindB, qA)

    qB = TupleLabXD.sortperm(cindB)
    cindA′′ = TupleLabXD.getindices(cindA, qB)
    cindB′′ = TupleLabXD.getindices(cindB, qB)

    dA, dB, dC = dim(A), dim(B), dim(C)

    # keep order A en B, check possibilities for cind
    memcost1 = memcost2 = dC*(!hsp(C, oindAinC, oindBinC))
    memcost1 += dA*(!hsp(A, oindA, cindA′)) +
                dB*(!hsp(B, cindB′, oindB))
    memcost2 += dA*(!hsp(A, oindA, cindA′′)) +
                dB*(!hsp(B, cindB′′, oindB))

    # reverse order A en B, check possibilities for cind
    memcost3 = memcost4 = dC*(!hsp(C, oindBinC, oindAinC))
    memcost3 += dB*(!hsp(B, oindB, cindB′)) +
                dA*(!hsp(A, cindA′, oindA))
    memcost4 += dB*(!hsp(B, oindB, cindB′′)) +
                dA*(!hsp(A, cindA′′, oindA))

    if min(memcost1, memcost2) <= min(memcost3, memcost4)
        if memcost1 <= memcost2
            return _contract!(α, A, B, β, C, oindA, cindA′, oindB, cindB′, p1, p2, syms)
        else
            return _contract!(α, A, B, β, C, oindA, cindA′′, oindB, cindB′′, p1, p2, syms)
        end
    else
        p1′ = map(n->ifelse(n>N₁, n-N₁, n+N₂), p1)
        p2′ = map(n->ifelse(n>N₁, n-N₁, n+N₂), p2)
        if memcost3 <= memcost4
            return _contract!(α, B, A, β, C, oindB, cindB′, oindA, cindA′, p1′, p2′, syms)
        else
            return _contract!(α, B, A, β, C, oindB, cindB′′, oindA, cindA′′, p1′, p2′, syms)
        end
    end
end

function _contract!(α, A::AbstractTensorMap{S}, B::AbstractTensorMap{S},
                    β, C::AbstractTensorMap{S},
                    oindA::IndexTuple{N₁}, cindA::IndexTuple,
                    oindB::IndexTuple{N₂}, cindB::IndexTuple,
                    p1::IndexTuple, p2::IndexTuple,
                    syms::Union{Nothing, NTuple{3, Symbol}} = nothing) where {S, N₁, N₂}

    if !(BraidingStyle(sectortype(S)) isa SymmetricBraiding)
        throw(SectorMismatch("only tensors with symmetric braiding rules can be contracted;
                                try `@planar` instead"))
    end
    copyA = false
    if BraidingStyle(sectortype(S)) isa Fermionic
        for i in cindA
            if !isdual(space(A, i))
                copyA = true
            end
        end
    end
    if syms === nothing
        A′ = permute(A, oindA, cindA; copy = copyA)
        B′ = permute(B, cindB, oindB)
    else
        A′ = cached_permute(syms[1], A, oindA, cindA; copy = copyA)
        B′ = cached_permute(syms[2], B, cindB, oindB)
    end
    if BraidingStyle(sectortype(S)) isa Fermionic
        for i in domainind(A′)
            if !isdual(space(A′, i))
                A′ = twist!(A′, i)
            end
        end
    end
    ipC = TupleLabXD.invperm((p1..., p2...))
    oindAinC = TupleLabXD.getindices(ipC, ntuple(n->n, N₁))
    oindBinC = TupleLabXD.getindices(ipC, ntuple(n->n+N₁, N₂))
    if has_shared_permute(C, oindAinC, oindBinC)
        C′ = permute(C, oindAinC, oindBinC)
        mul!(C′, A′, B′, α, β)
    else
        if syms === nothing
            C′ = A′*B′
        else
            p1′ = ntuple(identity, N₁)
            p2′ = N₁ .+ ntuple(identity, N₂)
            TC = eltype(C)
            C′ = TO.cached_similar_from_indices(syms[3], TC, oindA, oindB, p1′, p2′, A, B,
                                                    :N, :N)
            mul!(C′, A′, B′)
        end
        add!(α, C′, β, C, p1, p2)
    end
    return C
end

# Add support for cache and API (`@tensor` and other macros) from TensorContractionsXD.jl:
"""
    TO.memsize(t::TensorMap)

Return total size of the data for the tensor map `t` in bytes, i.e., the memory size that
needed to save the tensor map.
"""
function TO.memsize(t::TensorMap)
    s = 0
    for (c, b) in blocks(t)
        s += sizeof(b)
    end
    return s
end
TO.memsize(t::AdjointTensorMap) = TensorContractionsXD.memsize(t')

"""
    _similarstructure_from_indices(::Type{T}, p1::IndexTuple{N₁}, p2::IndexTuple{N₂},
                                t::AbstractTensorMap{S}) where {T, S<:IndexSpace, N₁, N₂}

Return the structure of an object similar to the tensor map `t`, i.e., the HomSpace from
the domain to the codomain, where the codomain is the tensor product of spaces selected by
`p1` from the tensor map `t`, and the domain is that by `p2`.
"""
function _similarstructure_from_indices(::Type{T}, p1::IndexTuple{N₁}, p2::IndexTuple{N₂},
                                t::AbstractTensorMap{S}) where {T, S<:IndexSpace, N₁, N₂}

    cod = ProductSpace{S, N₁}(space.(Ref(t), p1))
    dom = ProductSpace{S, N₂}(dual.(space.(Ref(t), p2)))
    return dom→cod
end

"""
    TO.similarstructure_from_indices(T::Type, p1::IndexTuple, p2::IndexTuple,
                                        A::AbstractTensorMap, CA::Symbol = :N)

Returns the structure of an object similar to the tensor map `A`, i.e. the HomSpace from
the domain to the codomain, where the codomain is the tensor product of spaces selected by
`p1` from the tensor map `op(A)`, and the domain is that by `p2`. Th `op` is `conj` if
`CA == :C` or does nothing if `CA == :N` (default).

The `eltype` of the returned structure is supposed given by `T`, but this is not actually
implemented.
"""
function TO.similarstructure_from_indices(T::Type, p1::IndexTuple, p2::IndexTuple,
        A::AbstractTensorMap, CA::Symbol = :N)
    if CA == :N
        _similarstructure_from_indices(T, p1, p2, A)
    else
        p1 = adjointtensorindices(A, p1)
        p2 = adjointtensorindices(A, p2)
        _similarstructure_from_indices(T, p1, p2, adjoint(A))
    end
end

"""
    _similarstructure_from_indices(::Type{T}, oindA::IndexTuple, oindB::IndexTuple,
            p1::IndexTuple{N₁}, p2::IndexTuple{N₂}, tA::AbstractTensorMap{S},
            tB::AbstractTensorMap{S}) where {T, S<:IndexSpace, N₁, N₂}

Return the structure of an object similar to the tensor map `tA`, i.e., the HomSpace from
the domain to the codomain, where the codomain is the tensor product of spaces selected by
`p1` from the list of spaces that are selected by `oindA` from `tA` and `oindB` from `tB` in
sequence, and the domain is that by `p2`.
"""
function _similarstructure_from_indices(::Type{T}, oindA::IndexTuple, oindB::IndexTuple,
        p1::IndexTuple{N₁}, p2::IndexTuple{N₂}, tA::AbstractTensorMap{S},
        tB::AbstractTensorMap{S}) where {T, S<:IndexSpace, N₁, N₂}

    spaces = (space.(Ref(tA), oindA)..., space.(Ref(tB), oindB)...)
    cod = ProductSpace{S, N₁}(getindex.(Ref(spaces), p1))
    dom = ProductSpace{S, N₂}(dual.(getindex.(Ref(spaces), p2)))
    return dom→cod
end

"""
    TO.similarstructure_from_indices(T::Type, poA::IndexTuple, poB::IndexTuple,
            p1::IndexTuple, p2::IndexTuple, A::AbstractTensorMap, B::AbstractTensorMap,
            CA::Symbol = :N, CB::Symbol = :N)

Return the structure of an object similar to the tensor map `A`, i.e., the HomSpace from
the domain to the codomain, where the codomain is the tensor product of spaces selected by
`p1` from the list of spaces that are selected by `poA` from `op(A)` and `poB` from `op(B)`
in sequence, and the domain is that by `p2`. Th `opA` is `conj` if `CA == :C` or does
nothing if `CA == :N` (default), and similarly for `opB`.

The `eltype` of the returned structure is supposed given by `T`, but this is not actually
implemented.
"""
function TO.similarstructure_from_indices(T::Type, poA::IndexTuple, poB::IndexTuple,
        p1::IndexTuple, p2::IndexTuple, A::AbstractTensorMap, B::AbstractTensorMap,
        CA::Symbol = :N, CB::Symbol = :N)

    if CA == :N && CB == :N
        _similarstructure_from_indices(T, poA, poB, p1, p2, A, B)
    elseif CA == :C && CB == :N
        poA = adjointtensorindices(A, poA)
        _similarstructure_from_indices(T, poA, poB, p1, p2, adjoint(A), B)
    elseif CA == :N && CB == :C
        poB = adjointtensorindices(B, poB)
        _similarstructure_from_indices(T, poA, poB, p1, p2, A, adjoint(B))
    else
        poA = adjointtensorindices(A, poA)
        poB = adjointtensorindices(B, poB)
        _similarstructure_from_indices(T, poA, poB, p1, p2, adjoint(A), adjoint(B))
    end
end

TO.scalar(t::AbstractTensorMap) = scalar(t)

function TO.add!(α, tsrc::AbstractTensorMap{S}, CA::Symbol, β,
    tdst::AbstractTensorMap{S, N₁, N₂}, p1::IndexTuple, p2::IndexTuple) where {S, N₁, N₂}

    if CA == :N
        p = (p1..., p2...)
        pl = TupleLabXD.getindices(p, codomainind(tdst))
        pr = TupleLabXD.getindices(p, domainind(tdst))
        add!(α, tsrc, β, tdst, pl, pr)
    else
        p = adjointtensorindices(tsrc, (p1..., p2...))
        pl = TupleLabXD.getindices(p, codomainind(tdst))
        pr = TupleLabXD.getindices(p, domainind(tdst))
        add!(α, adjoint(tsrc), β, tdst, pl, pr)
    end
    return tdst
end

function TO.trace!(α, tsrc::AbstractTensorMap{S}, CA::Symbol, β,
    tdst::AbstractTensorMap{S, N₁, N₂}, p1::IndexTuple, p2::IndexTuple,
    q1::IndexTuple, q2::IndexTuple) where {S, N₁, N₂}

    if CA == :N
        p = (p1..., p2...)
        pl = TupleLabXD.getindices(p, codomainind(tdst))
        pr = TupleLabXD.getindices(p, domainind(tdst))
        trace!(α, tsrc, β, tdst, pl, pr, q1, q2)
    else
        p = adjointtensorindices(tsrc, (p1..., p2...))
        pl = TupleLabXD.getindices(p, codomainind(tdst))
        pr = TupleLabXD.getindices(p, domainind(tdst))
        q1 = adjointtensorindices(tsrc, q1)
        q2 = adjointtensorindices(tsrc, q2)
        trace!(α, adjoint(tsrc), β, tdst, pl, pr, q1, q2)
    end
    return tdst
end

function TO.contract!(α,
    tA::AbstractTensorMap{S}, CA::Symbol,
    tB::AbstractTensorMap{S}, CB::Symbol,
    β, tC::AbstractTensorMap{S, N₁, N₂},
    oindA::IndexTuple, cindA::IndexTuple,
    oindB::IndexTuple, cindB::IndexTuple,
    p1::IndexTuple, p2::IndexTuple,
    syms::Union{Nothing, NTuple{3, Symbol}} = nothing) where {S, N₁, N₂}

    p = (p1..., p2...)
    pl = ntuple(n->p[n], N₁)
    pr = ntuple(n->p[N₁+n], N₂)
    if CA == :N && CB == :N
        contract!(α, tA, tB, β, tC, oindA, cindA, oindB, cindB, pl, pr, syms)
    elseif CA == :N && CB == :C
        oindB = adjointtensorindices(tB, oindB)
        cindB = adjointtensorindices(tB, cindB)
        contract!(α, tA, tB', β, tC, oindA, cindA, oindB, cindB, pl, pr, syms)
    elseif CA == :C && CB == :N
        oindA = adjointtensorindices(tA, oindA)
        cindA = adjointtensorindices(tA, cindA)
        contract!(α, tA', tB, β, tC, oindA, cindA, oindB, cindB, pl, pr, syms)
    elseif CA == :C && CB == :C
        oindA = adjointtensorindices(tA, oindA)
        cindA = adjointtensorindices(tA, cindA)
        oindB = adjointtensorindices(tB, oindB)
        cindB = adjointtensorindices(tB, cindB)
        contract!(α, tA', tB', β, tC, oindA, cindA, oindB, cindB, pl, pr, syms)
    else
        error("unknown conjugation flags: $CA and $CB")
    end
    return tC
end

function cached_permute(sym::Symbol, t::TensorMap{S},
                            p1::IndexTuple{N₁},  p2::IndexTuple{N₂}=();
                            copy::Bool = false) where {S, N₁, N₂}
    cod = ProductSpace{S, N₁}(map(n->space(t, n), p1))
    dom = ProductSpace{S, N₂}(map(n->dual(space(t, n)), p2))
    # share data if possible
    if !copy
        if p1 === codomainind(t) && p2 === domainind(t)
            return t
        elseif has_shared_permute(t, p1, p2)
            return TensorMap(reshape(t.data, dim(cod), dim(dom)), cod, dom)
        end
    end
    # general case
    @inbounds begin
        tp = TO.cached_similar_from_indices(sym, eltype(t), p1, p2, t, :N)
        return add!(true, t, false, tp, p1, p2)
    end
end

function cached_permute(sym::Symbol, t::AdjointTensorMap{S},
                            p1::IndexTuple,  p2::IndexTuple=();
                            copy::Bool = false) where {S, N₁, N₂}
    p1′ = adjointtensorindices(t, p2)
    p2′ = adjointtensorindices(t, p1)
    adjoint(cached_permute(sym, adjoint(t), p1′, p2′; copy = copy))
end

ext: jl


text:
l = [1,2,3,4,5]

# 1 based index
println(l[1])

max_stars_repo_name: mkanenobu/trashbox

max_issues_repo_name: mkanenobu/trashbox

max_forks_repo_name: mkanenobu/trashbox


text:
using Pkg
Pkg.activate(joinpath(@__DIR__, "../.."))
Pkg.instantiate()

# ## visualizer
vis = Visualizer()
open(vis)

# ## setup
using Dojo
using IterativeLQR
using LinearAlgebra
using FiniteDiff
using DojoEnvironments
using JLD2

# ## scripts
include(joinpath(module_dir(), "examples/policy/methods/continuation.jl"))
include(joinpath(module_dir(), "examples/policy/methods/tvlqr.jl"))
include(joinpath(module_dir(), "DojoEnvironments/src",
    "quadruped/methods/template.jl"))

################################################################################
# ## system
################################################################################
gravity = -9.81
timestep = 0.02
friction_coefficient = 0.8
damper = 0.5
spring = 1.0
env = get_environment(:quadruped,
    representation=:minimal,
    timestep=timestep,
    contact_body=false,
    gravity=gravity,
    friction_coefficient=friction_coefficient,
    damper=damper,
    spring=spring,
    infeasible_control=true,
    vis=vis)

# ## dimensions
nx = env.num_states
nu = env.num_inputs
nu_infeasible = 6

################################################################################
# Load trajectory and gains
################################################################################
file = JLD2.jldopen(joinpath(@__DIR__, "../data/trotting_forward.jld2"))
x_ref = file["x"][1:36]
u_ref = file["u"][1:35]
JLD2.close(file)

file = JLD2.jldopen(joinpath(@__DIR__, "../data/tvlqr_gains.jld2"))
K_tvlqr = file["K"]
JLD2.close(file)

file = JLD2.jldopen(joinpath(@__DIR__, "../data/tuned_gains.jld2"))
K_tuned = file["K"]
JLD2.close(file)

T = length(x_ref)

################################################################################
# Simulate accurately the TVLQR policy
################################################################################
# divide timestep
S = 1
mech_sim = get_mechanism(:quadruped,
    contact_body=true,
    timestep=timestep/S,
    gravity=gravity,
    friction_coefficient=friction_coefficient,
    damper=damper,
    spring=spring)

# controller
function ctrl!(m, k; x_ref=x_ref, u_ref=u_ref, K_ref=K_ref)
    N = length(u_ref)
    nu = input_dimension(m)
    @show Int(floor((k-1)/S))
    ind = (Int(floor((k-1)/S))) % N + 1

    x = get_minimal_state(m)
    # x[1] -= 0.15
    u = u_ref[ind] - K_ref[ind] * (x - x_ref[ind])
    u = clamp.(u, -0.2, +0.2)
    u = [zeros(6); u[7:end]] / S
    set_input!(m, SVector{nu}(u))
end

controller_tvlqr!(m, k) = ctrl!(m, k; K_ref=K_tvlqr)
controller_tuned!(m, k) = ctrl!(m, k; K_ref=K_tuned)

x1_roll = deepcopy(xref[1])
x1_roll[3] += 0.25
x1_roll[4] += 0.0
set_minimal_state!(mech_sim, x1_roll)
Main.@elapsed storage_tvlqr = simulate!(mech_sim, 4.0, controller_tvlqr!,
    record=true,
    opts=SolverOptions(rtol=1e-5, btol=1e-4, undercut=5.0),
    )

x1_roll = deepcopy(xref[1])
x1_roll[3] += 0.25
x1_roll[4] += 0.0
set_minimal_state!(mech_sim, x1_roll)
Main.@elapsed storage_tuned = simulate!(mech_sim, 4.0, controller_tuned!,
    record=true,
    opts=SolverOptions(rtol=1e-5, btol=1e-4, undercut=5.0),
    )

vis, anim = Dojo.visualize(mech_sim, storage_tvlqr, vis=vis, build=true, name=:tvlqr)
vis, anim = Dojo.visualize(mech_sim, storage_tuned, vis=vis, animation=anim, build=true, name=:tuned)

settransform!(vis[:tuned], Translation(0,0.5,0))
# convert_frames_to_video_and_gif("quadruped_tuned_vs_tvlqr_facing")


text:
# Code derived from metabolicEP (https://github.com/anna-pa-m/Metabolic-EP)

abstract type AbstractEPMat end


text:
struct CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace}
    groupindex::I
    ndrange::NDRange
    iterspace::Iterspace

    # CPU variant
    function CompilerMetadata{NDRange, CB}(idx, ndrange, iterspace) where {NDRange, CB}
        if ndrange !== nothing
            ndrange = CartesianIndices(ndrange)
        end
        return new{NDRange, CB, typeof(idx), typeof(ndrange), typeof(iterspace)}(idx, ndrange, iterspace)
    end

    # GPU variante: index is given implicit
    function CompilerMetadata{NDRange, CB}(ndrange, iterspace) where {NDRange, CB}
        if ndrange !== nothing
            ndrange = CartesianIndices(ndrange)
        end
        return new{NDRange, CB, Nothing, typeof(ndrange), typeof(iterspace)}(nothing, ndrange, iterspace)
    end
end

@inline __iterspace(cm::CompilerMetadata)  = cm.iterspace
@inline __groupindex(cm::CompilerMetadata) = cm.groupindex
@inline __groupsize(cm::CompilerMetadata) = size(workitems(__iterspace(cm)))
@inline __dynamic_checkbounds(::CompilerMetadata{NDRange, CB}) where {NDRange, CB} = CB <: DynamicCheck
@inline __ndrange(cm::CompilerMetadata{NDRange}) where {NDRange<:StaticSize}  = CartesianIndices(get(NDRange))
@inline __ndrange(cm::CompilerMetadata{NDRange}) where {NDRange<:DynamicSize} = cm.ndrange

include("compiler/contract.jl")
include("compiler/pass.jl")

function generate_overdubs(Ctx)
   @eval begin
        @inline Cassette.overdub(ctx::$Ctx, ::typeof(groupsize)) = __groupsize(ctx.metadata)
        @inline Cassette.overdub(ctx::$Ctx, ::typeof(__workitems_iterspace)) = workitems(__iterspace(ctx.metadata))

        ###
        # Cassette fixes
        ###
        @inline Cassette.overdub(::$Ctx, ::typeof(Core.kwfunc), f) = return Core.kwfunc(f)
        @inline Cassette.overdub(::$Ctx, ::typeof(Core.apply_type), args...) = return Core.apply_type(args...)
        @inline Cassette.overdub(::$Ctx, ::typeof(StaticArrays.Size), x::Type{<:AbstractArray{<:Any, N}}) where {N} = return StaticArrays.Size(x)

        @inline Cassette.overdub(::$Ctx, ::typeof(+), a::T, b::T) where T<:Union{Float32, Float64} = add_float_contract(a, b)
        @inline Cassette.overdub(::$Ctx, ::typeof(-), a::T, b::T) where T<:Union{Float32, Float64} = sub_float_contract(a, b)
        @inline Cassette.overdub(::$Ctx, ::typeof(*), a::T, b::T) where T<:Union{Float32, Float64} = mul_float_contract(a, b)

        function Cassette.overdub(::$Ctx, ::typeof(:), start::T, step::T, stop::T) where T<:Union{Float16,Float32,Float64}
            lf = (stop-start)/step
            if lf < 0
                len = 0
            elseif lf == 0
                len = 1
            else
                len = round(Int, lf) + 1
                stop′ = start + (len-1)*step
                # if we've overshot the end, subtract one:
                len -= (start < stop < stop′) + (start > stop > stop′)
            end
            Base.steprangelen_hp(T, start, step, 0, len, 1)
        end
    end
end


text:
@windows_only begin
    using WinRPM
    WinRPM.install("gcc",yes=true)
end


text:
module Api

# Shared api trough the modules, i.e. names used by more than one module
# Modules are free to use other functions but these are defined here to avoid name conflicts
# and allows instead Multiple Dispatch to handle them

export predict, fit, fit!, partition

predict()   = nothing
fit()       = nothing
fit!()      = nothing
partition() = nothing

end


text:
function movetoside!(r,side)
    count=0
    while isborder(r,side)==0
        move!(r,side)
        count=count+1
    end
    return count
end

function corner!(r)
home=[]
    while isborder(r,Nord)==0 || isborder(r,West)==0
        x=movetoside!(r,Nord)
        y=movetoside!(r,West)
        push!(home,x)
        push!(home,y)
    end
    return home
end

function countmove!(r,side,count)
    for i in 1:count
        move!(r,side)
    end
end

function gohome!(r,home)
    reverse!(home)
    for i = 1:length(home)
        if i%2==0
            countmove!(r,Sud,home[i])
        else
            countmove!(r,Ost,home[i])
        end
    end
end

function dotsborder!(r)
    putmarker!(r)
    move!(r,Sud)
    if ismarker(r)==0
        for side in [Sud,Ost,Nord,West]
            while isborder(r,side)==0
                putmarker!(r)
                move!(r,side)
            end
        putmarker!(r)
        end
    end

end
function task15!(r)
home=[]
home=corner!(r)
dotsborder!(r)
gohome!(r,home)
end
task5!(r)


text:
using Documenter, WignerDMatrices

DocMeta.setdocmeta!(WignerDMatrices, :DocTestSetup, :(using WignerDMatrices); recursive=true)

makedocs(;sitename="WignerDMatrices",
	modules=[WignerDMatrices],
    authors="Jishnu Bhattacharya",
    repo="https://github.com/jishnub/WignerDMatrices.jl/blob/{commit}{path}#L{line}",
    format = Documenter.HTML(
        prettyurls = get(ENV, "CI", nothing) == "true"
    ),
    pages=[
        "Reference" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/jishnub/WignerDMatrices.jl",
)


text:
using SeisIO, SeisNoise, Dates, Glob, BenchmarkTools, StatsPlots, Statistics
using CuArrays, PyCall, Plots.PlotMeasures, Printf, HTTP, DelimitedFiles
## This is a test of single core Julia CPU performace vs GPU performance vs python
# for one year of cross-correlation between two LHZ stations
# numpy, scipy and obspy must be installed via Conda for this benchmark

fs = 1.
cc_len, cc_step = 2^15, 26768
freqmin = 0.003
freqmax = 0.005
maxlag = 12000.
file1 = expanduser("~/SeisNoise_paper/DATA/YEARSAC/2019.001.00.00.00.000.BK.BDM.00.LHZ.R.SAC")
file2 = expanduser("~/SeisNoise_paper/DATA/YEARSAC/2019.001.00.00.00.000.BK.BKS.00.LHZ.R.SAC")

# check if these YEARSAC directory exist, if not download from zenodo
if !isdir(expanduser("~/SeisNoise_paper/DATA/YEARSAC"))
    mkpath(expanduser("~/SeisNoise_paper/DATA/YEARSAC"))
end

# download file
if !isfile(file1) | !isfile(file2)
    zenodourl = "https://zenodo.org/record/3823283/files/"

    dlpath1 = joinpath(zenodourl,basename(file1) * "?download=1")
    dlpath2 = joinpath(zenodourl,basename(file2) * "?download=1")
    outpath1 = joinpath(expanduser("~/SeisNoise_paper/DATA/YEARSAC"), file1)
    outpath2 = joinpath(expanduser("~/SeisNoise_paper/DATA/YEARSAC"), file2)
    println("Downloading file $file1, 1 of 2")
    HTTP.download(dlpath1,outpath1)

    println("Downloading file $file1, 2 of 2")
    HTTP.download(dlpath2,outpath2)
end

# function for cross-correlation on the CPU
function CPU_CC(file1,file2,cc_len,cc_step,fs,freqmin,freqmax,maxlag)
    S1 = read_data("sac",file1)
    S2 = read_data("sac",file2)
    R1 = RawData(S1[1],cc_len,cc_step)
    R2 = RawData(S2[1],cc_len,cc_step)
    detrend!(R1)
    detrend!(R2)
    taper!(R1)
    taper!(R2)
    highpass!(R1,0.001)
    highpass!(R2,0.001)
    onebit!(R1)
    onebit!(R2)
    F1 = rfft(R1)
    F2 = rfft(R2)
    C = compute_cc(F1,F2,maxlag)
    stack!(C,allstack=true)
    clean_up!(C,freqmin,freqmax)
    return C
end

# function for cross-correlation on the CPU
function GPU_CC(file1,file2,cc_len,cc_step,fs,freqmin,freqmax,maxlag)
    S1 = read_data("sac",file1)
    S2 = read_data("sac",file2)
    R1 = RawData(S1[1],cc_len,cc_step) |> gpu
    R2 = RawData(S2[1],cc_len,cc_step) |> gpu
    detrend!(R1)
    detrend!(R2)
    taper!(R1)
    taper!(R2)
    highpass!(R1,0.001)
    highpass!(R2,0.001)
    onebit!(R1)
    onebit!(R2)
    F1 = rfft(R1)
    F2 = rfft(R2)
    C = compute_cc(F1,F2,maxlag)
    stack!(C,allstack=true)
    clean_up!(C,freqmin,freqmax)
    return C |> cpu
end

#### TEST SINGLE CORE SINGLE CC ####
CPU_CC(file1,file2,cc_len,cc_step,fs,freqmin,freqmax,maxlag)
GPU_CC(file1,file2,cc_len,cc_step,fs,freqmin,freqmax,maxlag)

# benchmark total compute time
println("Benchmarking Total CPU Time")
println("---------------------------")
a = @benchmarkable CPU_CC(file1,file2,cc_len,cc_step,fs,freqmin,freqmax,maxlag) seconds=60 time_tolerance=0.01
TOTALCPU = run(a)
display("text/plain", TOTALCPU)

println("\n\nBenchmarking Total GPU Time")
println("---------------------------")
b = @benchmarkable begin CuArrays.@sync GPU_CC(file1,file2,cc_len,cc_step,fs,freqmin,freqmax,maxlag) end seconds=60 time_tolerance=0.01
TOTALGPU = run(b)
display("text/plain", TOTALGPU)

# benchmark file read speed
println("\n\nBenchmarking SeisIO Read Speed")
println("------------------------------")
SEISIO = @benchmark read_data("sac",file1)
display("text/plain", SEISIO)
S1 = read_data("sac",file1);
S2 = read_data("sac",file2);

# benchmark file read speed
println("\n\nBenchmarking Obspy Read Speed")
println("-----------------------------")
obspy = pyimport("obspy")
OBSPYIO = @benchmark obspy.read(file1)
display("text/plain", OBSPYIO)

# benchmark RawData creation
println("\n\nBenchmarking RawData CPU")
println("------------------------")
RAWCPU = @benchmark RawData(S1[1],cc_len,cc_step)
display("text/plain", RAWCPU)
R1 = RawData(S1[1],cc_len,cc_step);
R2 = RawData(S2[1],cc_len,cc_step);

# benchmark detrending
println("\n\nBenchmarking detrending CPU")
println("---------------------------")
DETRENDCPU = @benchmark detrend!(R1)
display("text/plain", DETRENDCPU)

# benchmark tapering
println("\n\nBenchmarking tapering CPU")
println("-------------------------")
TAPERCPU = @benchmark taper!(R1,max_length=400.)
display("text/plain", TAPERCPU)

# benchmark filtering
println("\n\nBenchmarking filtering CPU")
println("-------------------------")
FILTERCPU = @benchmark highpass!(R1,0.001)
display("text/plain", FILTERCPU)

# benchmark onebit
println("\n\nBenchmarking time-normalization CPU")
println("-----------------------------------")
ONEBITCPU = @benchmark onebit!(R1)
display("text/plain", ONEBITCPU)

# benchmark Fourier transform CPU
println("\n\nBenchmarking rfft CPU")
println("---------------------")
RFFTCPU = @benchmark rfft(R1)
display("text/plain", RFFTCPU)
F1 = rfft(R1);
F2 = rfft(R2);

# benchmark cross-correlation CPU
println("\n\nBenchmarking cross-correlation CPU")
println("----------------------------------")
CCCPU = @benchmark compute_cc(F1,F2,maxlag)
display("text/plain", CCCPU)
C = compute_cc(F1,F2,maxlag)

# benchmark stacking CPU
println("\n\nBenchmarking stacking CPU")
println("-------------------------")
# this times just the stack and not the reallocation
STACKCPU = @benchmark mean(C.corr,dims=2)
display("text/plain", STACKCPU)
stack!(C,allstack=true)

# benchmark final clean-up
println("\n\nBenchmarking post-processing CPU")
println("--------------------------------")
# this times just the stack and not the reallocation
POSTCPU = @benchmark clean_up!(C,freqmin,freqmax)
display("text/plain", POSTCPU)

# benchmark RawData creation
println("\n\nBenchmarking RawData GPU")
println("------------------------")
c = @benchmarkable begin CuArrays.@sync RawData(S1[1],cc_len,cc_step) |> gpu end
RAWGPU = run(c)
display("text/plain", RAWGPU)
R1 = RawData(S1[1],cc_len,cc_step) |> gpu;
R2 = RawData(S2[1],cc_len,cc_step) |> gpu;

println("\n\nBenchmarking detrending GPU")
println("---------------------------")
d = @benchmarkable begin CuArrays.@sync detrend!(R1) end
DETRENDGPU = run(d)
display("text/plain", DETRENDGPU)

# benchmark tapering
println("\n\nBenchmarking tapering GPU")
println("-------------------------")
e = @benchmarkable begin CuArrays.@sync taper!(R1,max_length=400.) end
TAPERGPU = run(e)
display("text/plain", TAPERGPU)

# benchmark filtering
println("\n\nBenchmarking filtering GPU")
println("-------------------------")
f = @benchmarkable begin CuArrays.@sync highpass!(R1,0.001) end
FILTERGPU = run(f)
display("text/plain", FILTERGPU)

# benchmark onebit
println("\n\nBenchmarking time-normalization GPU")
println("-----------------------------------")
g = @benchmarkable begin CuArrays.@sync onebit!(R1) end
ONEBITGPU = run(g)
display("text/plain", ONEBITGPU)

# benchmark Fourier transform CPU
println("\n\nBenchmarking rfft GPU")
println("---------------------")
h = @benchmarkable begin CuArrays.@sync rfft(R1) end
RFFTGPU = run(h)
display("text/plain", RFFTGPU)
F1 = rfft(R1);
F2 = rfft(R2);

# benchmark cross-correlation CPU
println("\n\nBenchmarking cross-correlation GPU")
println("----------------------------------")
i = @benchmarkable begin CuArrays.@sync compute_cc(F1,F2,maxlag) |> cpu end
CCGPU = run(i)
display("text/plain", CCGPU)
C = compute_cc(F1,F2,maxlag)

# benchmark stacking CPU
println("\n\nBenchmarking stacking GPU")
println("-------------------------")
# this times just the stack and not the reallocation
j = @benchmarkable begin CuArrays.@sync mean(C.corr,dims=2) end
STACKGPU = run(j)
display("text/plain", STACKGPU)
stack!(C,allstack=true)

# benchmark final clean-up
println("\n\nBenchmarking post-processing GPU")
println("--------------------------------")
k = @benchmarkable begin CuArrays.@sync clean_up!(cpu(C),freqmin,freqmax) end
POSTGPU = run(k)
display("text/plain", POSTGPU)

## python benchmarks ##
# reload RawData
R1 = RawData(S1[1],cc_len,cc_step);
R2 = RawData(S2[1],cc_len,cc_step);
# send to numpy
Rnp1 = PyReverseDims(R1.x)
Rnp2 = PyReverseDims(R2.x)
# import scipy modules
np = pyimport("numpy")
signal = pyimport("scipy.signal")
spf = pyimport("scipy.fft")

println("\n\nBenchmarking detrending SciPy")
println("-----------------------------")
DETRENDSCI = @benchmark @pycall signal.detrend(Rnp1)::PyArray
display("text/plain", DETRENDSCI)

# benchmark tapering
println("\n\nBenchmarking tapering SciPy")
println("---------------------------")
py"""
from scipy import signal
import numpy as np
def taper(A, fs,max_length=400):
    window = np.int(fs * max_length)
    tape = signal.hann(max_length * 2)
    if A.ndim == 2:
        A[:,:window] *= tape[:window]
        A[:,-window:] *= tape[-window:]
    elif A.ndim == 1:
        A[:window] *= tape[:window]
        A[-window:] *= tape[-window:]

    return A
"""
TAPERSCI = @benchmark @pycall py"taper"(Rnp1,fs)::PyArray
display("text/plain", TAPERCPU)

# benchmark filtering
println("\n\nBenchmarking filtering SciPy")
println("----------------------------")
sos = @pycall signal.butter(4,0.001,btype="highpass",fs=fs,output="sos")::PyArray
FILTERSCI = @benchmark @pycall signal.sosfiltfilt(sos,Rnp1)::PyArray
display("text/plain", FILTERSCI)

# benchmark onebit
println("\n\nBenchmarking time-normalization SciPy")
println("-------------------------------------")
ONEBITSCI = @benchmark @pycall np.sign(Rnp1)::PyArray
display("text/plain", ONEBITSCI)

# benchmark Fourier transform CPU
println("\n\nBenchmarking rfft SciPy")
println("-----------------------")
RFFTSCI = @benchmark @pycall spf.rfft(Rnp1)::PyArray
display("text/plain", RFFTSCI)
Fsc1 = @pycall spf.rfft(Rnp1)::PyArray
Fsc2 = @pycall spf.rfft(Rnp2)::PyArray

# benchmark cross-correlation
println("\n\nBenchmarking cross-correlation SciPy")
println("------------------------------------")
py"""
from scipy.fft import rfft, irfft, fftshift
import numpy as np
# cross-correlation function
def xcorr(A,B,N,maxlag):
    corr = irfft(np.conj(A) * B)
    t =  np.hstack((np.arange(0,N/2),np.arange(-N/2,0)))
    ind = np.where(np.abs(t) <= maxlag)[0]
    newind = fftshift(ind)
    return corr[:,newind]
"""
N = np.shape(Rnp1)[2]
CCSCI = @benchmark @pycall py"xcorr"(Fsc1,Fsc2,N,maxlag)::PyArray
display("text/plain", CCSCI)
Xsci = @pycall py"xcorr"(Fsc1,Fsc2,N,maxlag)::PyArray

# benchmark stacking CPU
println("\n\nBenchmarking stacking SciPy")
println("---------------------------")
# this times just the stack and not the reallocation
STACKSCI = @benchmark @pycall np.mean(Xsci,axis=0)::PyArray
display("text/plain", STACKSCI)
Xstack = @pycall np.mean(Xsci,axis=0)::PyArray

# benchmark final clean-up
println("\n\nBenchmarking post-processing SciPy")
println("----------------------------------")
py"""
def clean_up(A,freqmin,freqmax,fs):
    A = signal.detrend(A, type='linear')
    A = taper(A,fs)
    sos = signal.butter(4,[freqmin,freqmax],btype="bandpass",fs=fs,output="sos")
    A = signal.sosfiltfilt(sos,A)
    return A
"""
POSTSCI = @benchmark @pycall py"clean_up"(Xstack,freqmin,freqmax,fs)::PyArray
display("text/plain", POSTSCI)

processes = [SEISIO,RAWGPU,DETRENDGPU,TAPERGPU,FILTERGPU,ONEBITGPU,RFFTGPU,CCGPU,
             STACKGPU,POSTGPU,SEISIO, RAWCPU,DETRENDCPU,TAPERCPU,FILTERCPU,ONEBITCPU,
             RFFTCPU,CCCPU,STACKCPU,POSTCPU]
SCIPY = [DETRENDSCI,TAPERSCI,FILTERSCI,ONEBITSCI,RFFTSCI,CCSCI,STACKSCI,POSTSCI]
ctg = repeat(["CUDA", "Julia","Python"], inner = 10)
colors = repeat(["white", "black","grey"], inner = 10)
nam = repeat(["I/O","RawData","Detrend","Taper","Filter","Time\n Normalization","RFFT","Xcorr","Stack","Post-\nProcess"], outer = 3)

times = [median(p).time * 1e-6 for p in processes]
scitimes = vcat([median(OBSPYIO).time,0], [median(p).time for p in SCIPY]) .* 1e-6
scitotal = sum(scitimes .* [1,1,2,2,2,2,2,1,1,1])
scitotal += (median(OBSPYIO).time + median(RAWCPU).time) * 1e-6 * 2
cputotal = median(TOTALCPU).time * 1e-6
gputotal = median(TOTALGPU).time * 1e-6
times = vcat(times,scitimes)

# final timing results
println("\n\nFinal Timing Results")
println("--------------------")
println("GPU    Total Time: $(@sprintf("%.2f",gputotal)) ms")
println("CPU    Total Time: $(@sprintf("%.2f",cputotal)) ms")
println("Python Total Time: $(@sprintf("%.2f",scitotal)) ms")
println("GPU vs CPU    Speedup: $(@sprintf("%.2f",cputotal/gputotal))x")
println("GPU vs Python Speedup: $(@sprintf("%.2f",scitotal/gputotal))x")
println("CPU vs Python Speedup: $(@sprintf("%.2f",scitotal/cputotal))x")

println("\n\nI/O Timing Results")
println("------------------")
seisiogpu = median(SEISIO).time * 1e-6
seisiocpu = median(SEISIO).time * 1e-6
obspyio = median(OBSPYIO).time * 1e-6
println("GPU    Total Time: $(@sprintf("%.2f",seisiogpu)) ms")
println("CPU    Total Time: $(@sprintf("%.2f",seisiocpu)) ms")
println("Python Total Time: $(@sprintf("%.2f",obspyio)) ms")
println("GPU vs CPU    Speedup: $(@sprintf("%.2f",seisiocpu/seisiogpu))x")
println("GPU vs Python Speedup: $(@sprintf("%.2f",obspyio/seisiogpu))x")
println("CPU vs Python Speedup: $(@sprintf("%.2f",obspyio/seisiocpu))x")

println("\n\nDetrend Results")
println("---------------")
detrendgpu = median(DETRENDGPU).time * 1e-6
detrendcpu = median(DETRENDCPU).time * 1e-6
detrendsci = median(DETRENDSCI).time * 1e-6
println("GPU    Total Time: $(@sprintf("%.2f",detrendgpu)) ms")
println("CPU    Total Time: $(@sprintf("%.2f",detrendcpu)) ms")
println("Python Total Time: $(@sprintf("%.2f",detrendsci)) ms")
println("GPU vs CPU    Speedup: $(@sprintf("%.2f",detrendcpu/detrendgpu))x")
println("GPU vs Python Speedup: $(@sprintf("%.2f",detrendsci/detrendgpu))x")
println("CPU vs Python Speedup: $(@sprintf("%.2f",detrendsci/detrendcpu))x")

println("\n\nTaper Results")
println("-------------")
tapergpu = median(TAPERGPU).time * 1e-6
tapercpu = median(TAPERCPU).time * 1e-6
tapersci = median(TAPERSCI).time * 1e-6
println("GPU    Total Time: $(@sprintf("%.2f",tapergpu)) ms")
println("CPU    Total Time: $(@sprintf("%.2f",tapercpu)) ms")
println("Python Total Time: $(@sprintf("%.2f",tapersci)) ms")
println("GPU vs CPU    Speedup: $(@sprintf("%.2f",tapercpu/tapergpu))x")
println("GPU vs Python Speedup: $(@sprintf("%.2f",tapersci/tapergpu))x")
println("CPU vs Python Speedup: $(@sprintf("%.2f",tapersci/tapercpu))x")

println("\n\nFilter Results")
println("--------------")
filtergpu = median(FILTERGPU).time * 1e-6
filtercpu = median(FILTERCPU).time * 1e-6
filtersci = median(FILTERSCI).time * 1e-6
println("GPU    Total Time: $(@sprintf("%.2f",filtergpu)) ms")
println("CPU    Total Time: $(@sprintf("%.2f",filtercpu)) ms")
println("Python Total Time: $(@sprintf("%.2f",filtersci)) ms")
println("GPU vs CPU    Speedup: $(@sprintf("%.2f",filtercpu/filtergpu))x")
println("GPU vs Python Speedup: $(@sprintf("%.2f",filtersci/filtergpu))x")
println("CPU vs Python Speedup: $(@sprintf("%.2f",filtersci/filtercpu))x")

println("\n\nAmplitude Normalization Results")
println("-------------------------------")
onebitgpu = median(ONEBITGPU).time * 1e-6
onebitcpu = median(ONEBITCPU).time * 1e-6
onebitsci = median(ONEBITSCI).time * 1e-6
println("GPU    Total Time: $(@sprintf("%.2f",onebitgpu)) ms")
println("CPU    Total Time: $(@sprintf("%.2f",onebitcpu)) ms")
println("Python Total Time: $(@sprintf("%.2f",onebitsci)) ms")
println("GPU vs CPU    Speedup: $(@sprintf("%.2f",onebitcpu/onebitgpu))x")
println("GPU vs Python Speedup: $(@sprintf("%.2f",onebitsci/onebitgpu))x")
println("CPU vs Python Speedup: $(@sprintf("%.2f",onebitsci/onebitcpu))x")

println("\n\nrFFT Results")
println("------------")
rfftgpu = median(RFFTGPU).time * 1e-6
rfftcpu = median(RFFTCPU).time * 1e-6
rfftsci = median(RFFTSCI).time * 1e-6
println("GPU    Total Time: $(@sprintf("%.2f",rfftgpu)) ms")
println("CPU    Total Time: $(@sprintf("%.2f",rfftcpu)) ms")
println("Python Total Time: $(@sprintf("%.2f",rfftsci)) ms")
println("GPU vs CPU    Speedup: $(@sprintf("%.2f",rfftcpu/rfftgpu))x")
println("GPU vs Python Speedup: $(@sprintf("%.2f",rfftsci/rfftgpu))x")
println("CPU vs Python Speedup: $(@sprintf("%.2f",rfftsci/rfftcpu))x")

println("\n\nCross-Correlation Results")
println("-------------------------")
ccgpu = median(CCGPU).time * 1e-6
cccpu = median(CCCPU).time * 1e-6
ccsci = median(CCSCI).time * 1e-6
println("GPU    Total Time: $(@sprintf("%.2f",ccgpu)) ms")
println("CPU    Total Time: $(@sprintf("%.2f",cccpu)) ms")
println("Python Total Time: $(@sprintf("%.2f",ccsci)) ms")
println("GPU vs CPU    Speedup: $(@sprintf("%.2f",cccpu/ccgpu))x")
println("GPU vs Python Speedup: $(@sprintf("%.2f",ccsci/ccgpu))x")
println("CPU vs Python Speedup: $(@sprintf("%.2f",ccsci/cccpu))x")

println("\n\nStacking Results")
println("----------------")
stackgpu = median(STACKGPU).time * 1e-6
stackcpu = median(STACKCPU).time * 1e-6
stacksci = median(STACKSCI).time * 1e-6
println("GPU    Total Time: $(@sprintf("%.2f",stackgpu)) ms")
println("CPU    Total Time: $(@sprintf("%.2f",stackcpu)) ms")
println("Python Total Time: $(@sprintf("%.2f",stacksci)) ms")
println("GPU vs CPU    Speedup: $(@sprintf("%.2f",stackcpu/stackgpu))x")
println("GPU vs Python Speedup: $(@sprintf("%.2f",stacksci/stackgpu))x")
println("CPU vs Python Speedup: $(@sprintf("%.2f",stacksci/stackcpu))x")

println("\n\nPost-Processing Results")
println("-----------------------")
postgpu = median(POSTGPU).time * 1e-6
postcpu = median(POSTCPU).time * 1e-6
postsci = median(POSTSCI).time * 1e-6
println("GPU    Total Time: $(@sprintf("%.2f",postgpu)) ms")
println("CPU    Total Time: $(@sprintf("%.2f",postcpu)) ms")
println("Python Total Time: $(@sprintf("%.2f",postsci)) ms")
println("GPU vs CPU    Speedup: $(@sprintf("%.2f",postcpu/postgpu))x")
println("GPU vs Python Speedup: $(@sprintf("%.2f",postsci/postgpu))x")
println("CPU vs Python Speedup: $(@sprintf("%.2f",postsci/postcpu))x")

groupedbar(nam, times , group = ctg,xlabel = "Processing", ylabel = "Time [ms]",
        bar_width = 0.67, lw = 0, framestyle = :box,color=colors,legend=:best,
        size=(800,500),left_margin=10mm,bottom_margin=10mm,dpi=1000)
png(expanduser("~/SeisNoise_paper/CPU-GPU-BENCHMARK"))


text:
__precompile__(true)

module UnitfulPlasmas

  import Unitful
  using Unitful: @unit

  import Unitful.DefaultSymbols

  @unit n20 "n20" PlasmaDensity 1e20/Unitful.m^3 false

  # Some gymnastics required here because if we precompile, we cannot add to
  # Unitful.basefactors at compile time and expect the changes to persist to runtime.

  const localunits = Unitful.basefactors

  function __init__()
      merge!(Unitful.basefactors, localunits)
      Unitful.register(UnitfulPlasmas)
  end

end


text:
"""
# Gives the number of occurrences (counts) of the vector for each
#  unique value in the vector
"""
function countUniqueValues{T<:Number}(inValues::Array{T,1})
  uValues=unique(inValues)
  uCounts=zeros(Int64,(length(uValues),))
  for j=1:length(inValues)
    for i=1:length(uValues)
      (inValues[j]==uValues[i]) ? uCounts[i]+=1 : 0
    end
  end
  return uCounts,uValues
end

max_stars_repo_name: jadamc/MetoceanTools

max_issues_repo_name: jadamc/MetoceanTools

max_forks_repo_name: jadamc/MetoceanTools


text:
using DataGeneratorTranslators
using Test

#
# temporarily stub out BaseTestAuto features
#
macro mcheck_values_include(params...)
  quote warn("Skipping mcheck_values_include") end
end
macro mcheck_values_are(params...)
  quote warn("Skipping mcheck_values_are") end
end
macro mcheck_values_vary(params...)
  quote warn("Skipping mcheck_values_are") end
end
macro mcheck_that_sometimes(params...)
  quote warn("Skipping mcheck_that_sometimes") end
end
#

@testset "DataGeneratorTranslators" begin
	include(joinpath("regex", "runtests.jl"))
	include(joinpath("type", "runtests.jl"))
end


text:
@testset "gibbskernel" begin
    x = randn()
    y = randn()

    # this is the gibbs lengthscale function.
    ell(x) = exp(sum(sin, x))
    # create a gibbs kernel with our specific lengthscale function
    k_gibbs = GibbsKernel(ell)

    @test k_gibbs(x, y) ≈
          sqrt((2 * ell(x) * ell(y)) / (ell(x)^2 + ell(y)^2)) *
          exp(-(x - y)^2 / (ell(x)^2 + ell(y)^2))
end


text:
"""
---------------------------------------------------
Types and Constructors
---------------------------------------------------
"""

mutable struct CURVED_APPROACH_SAMPLER <: SAMPLER
	z_init::Sampleable
	vx_init::Sampleable
	vz_init::Sampleable
	radius::Sampleable
	t_extra::Sampleable
end

function curved_approach_sampler(;z_init = Uniform(300,600),
						vx_init = Uniform(30,45),
						vz_init = Uniform(-500fpm2fps,-300fpm2fps),
						radius = Uniform(200,400),
						t_extra = Normal(20,4))
	return CURVED_APPROACH_SAMPLER(z_init, vx_init, vz_init, radius, t_extra)
end

mutable struct CURVED_APPROACH <: UAM_TRAJECTORY
	z_init::Float64
	vx_init::Float64
	vz_init::Float64
	radius::Float64
	t_extra::Float64
	sampler::SAMPLER
	dt::Float64
	vz_max::Float64
	az_max::Float64
	ax_max::Float64
	p::Array{Float64,2}
	v::Array{Float64,2}
end

function curved_approach(;z_init = 0.0,
						vx_init = 0.0,
						vz_init = 0.0,
						radius = 0.0,
						t_extra = 0.0,
						sampler = curved_approach_sampler(),
						dt = 1.0,
						vz_max = 550fpm2fps, # this doesn't actually get used right now
						az_max = 0.3g_ft,
						ax_max = 0.1g, # this doesn't actually get used right now
						p = zeros(1,3),
						v = zeros(1,3))
	return CURVED_APPROACH(z_init, vx_init, vz_init, radius, t_extra, sampler, dt,
							vz_max, az_max, ax_max, p, v)
end

"""
---------------------------------------------------
Functions
---------------------------------------------------
"""

function solve_trajectory!(τ::CURVED_APPROACH)
	# Get some needed info
	ttot = abs(τ.z_init/τ.vz_init) + τ.t_extra
	N = convert(Int64, floor(ttot/τ.dt))
	# Difference matrix
	D = zeros(N-1,N)
	for i = 1:(N-1)
	    D[i,i] = -1
	    D[i,i+1] = 1
	end
	D = D./τ.dt

	# Create and solve the optimization problem
	px = Variable(N)
	py = Variable(N)
	pz = Variable(N)
	vx = Variable(N)
	vy = Variable(N)
	vz = Variable(N)

	objec = 100*sumsquares(D*vx) + 100*sumsquares(D*vy) + sumsquares(D*vz*ft2m)

	constraints = [[px[i+1] == px[i] + vx[i]*τ.dt for i = 1:N-1];
				   [py[i+1] == py[i] + vy[i]*τ.dt for i = 1:N-1];
				   [pz[i+1] == pz[i] + vz[i]*τ.dt for i = 1:N-1];
				   px[1] == 0; py[1] == 2*τ.radius; pz[1] == τ.z_init;
				   vx[1] == -τ.vx_init; vy[1] == 0; vz[1] == τ.vz_init;
				   px[end] == 0; py[end] == 0; pz[end] == 0;
				   vx[end] == 0.001; vy[end] == 0; vz[end] == 0; # Giving small x speed because I want it to constrain heading
				   D*vz*ft2m ≥ -τ.az_max; D*vz*ft2m ≤ τ.az_max]

	prob = minimize(objec, constraints)
	solve!(prob, ECOS.Optimizer(verbose=0))

	# Update the trajectory
	p = zeros(N,3)
	p[:,1] = px.value[:,1]
	p[:,2] = py.value[:,1]
	p[:,3] = pz.value[:,1]
	v = zeros(N,3)
	v[:,1] = vx.value[:,1]
	v[:,2] = vy.value[:,1]
	v[:,3] = vz.value[:,1]

	τ.p = copy(p)
	τ.v = copy(v)

	return Int(prob.status) == 1
end

max_stars_repo_name: Airspace-Encounter-Models/em-model-uam

max_issues_repo_name: Airspace-Encounter-Models/em-model-uam

max_forks_repo_name: Airspace-Encounter-Models/em-model-uam


text:
module COESA

export atmosphere, density, temperature, pressure, speed_of_sound, mean_molecular_weight

const r0 = 6356766.0 # (m), effective Earth radius at 45 deg N latitude
const g0 = 9.80665 # (m / s²) or (m² / s² m')
const M0 = 28.9644 # (kg / kmol)
const Rstar = 8.31432e3 # (N m / kmol K)
const γ = 1.4

geopotential_altitude(Z) = r0 * Z / (r0 + Z)

const Hb = [0, 11, 20, 32, 47, 51, 71] * 1000 # (m')
const Lmb = [-6.5, 0.0, 1.0, 2.8, 0.0, -2.8, -2.0] / 1000 # (K / m')

function findb(H)
    i = 1
    while i < length(Hb) && H > Hb[i + 1]
        i += 1
    end
    b = i - 1
    b
end

const Tmb = let
    local Tmb
    Tmb = [288.15] # (K)
    for i in 1:(length(Hb) - 1)
        push!(Tmb, Tmb[i] + Lmb[i] * (Hb[i + 1] - Hb[i]))
    end
    Tmb
end
function Tm(H)
    b = findb(H)
    i = b + 1
    Tmb[i] + Lmb[i] * (H - Hb[i])
end
temperature_lower(H, M) = Tm(H) / M0 * M

const Pb = let
    local Pb
    Pb = [101325.0]
    for i in 1:(length(Hb) - 1)
        if Lmb[i] == 0
            push!(Pb, Pb[i] * exp(-g0 * M0 * (Hb[i + 1] - Hb[i]) / (Rstar * Tmb[i])))
        else
            push!(Pb, Pb[i] * (Tmb[i] / (Tmb[i] + Lmb[i] * (Hb[i + 1] - Hb[i]))) ^ (g0 * M0 / (Rstar * Lmb[i])))
        end
    end
    Pb
end
function pressure_lower(H)
    b = findb(H)
    i = b + 1
    if Lmb[i] == 0
        return Pb[i] * exp(-g0 * M0 * (H - Hb[i]) / (Rstar * Tmb[i]))
    else
        return Pb[i] * (Tmb[i] / (Tmb[i] + Lmb[i] * (H - Hb[i]))) ^ (g0 * M0 / (Rstar * Lmb[i]))
    end
end

speed_of_sound_lower(T, M) = sqrt(γ * Rstar * T / M)

const Ztable = [80.0, 80.5, 81.0, 81.5, 82.0, 82.5, 83.0, 83.5, 84.0, 84.5, 85.0, 85.5, 86.0] * 1000 # (m)
const Mratiotable = [1.0, 0.999996, 0.999989, 0.999971, 0.999941, 0.999909, 0.999870, 0.999829, 0.999786, 0.999741, 0.999694, 0.999641, 0.999579]
function mean_molecular_weight_ratio_lower(Z)
    Z < Ztable[1] && return 1.0
    Z > Ztable[end] && error("altitude above maximum value in table")
    i = 1
    while i < length(Ztable) && Z > Ztable[i + 1]
        i += 1
    end
    Mratiotable[i] + (Mratiotable[i + 1] - Mratiotable[i]) / (Ztable[i + 1] - Ztable[i]) * (Z - Ztable[i])
end
mean_molecular_weight_lower(Z) = M0 * mean_molecular_weight_ratio_lower(Z)

function temperature_upper(Z)
    if Z <= 91_000
        return 186.8673 # (K)
    elseif Z <= 110_000
        Tc = 263.1905 # (K)
        A = -76.3232 # (K)
        a = -19.9429 * 1000 # (m)
        return Tc + A * sqrt(1 - ((Z - 91_000) / a) ^ 2)
    elseif Z <= 120_000
        T9 = 240 # (K)
        LK9 = 12 / 1000 # (K / m)
        Z9 = 110_000 # (m)
        return T9 + LK9 * (Z - Z9)
    elseif Z <= 1_000_000
        T10 = 360 # (K)
        Z10 = 120_000 # (m)
        Tinf = 1000 # (K)
        λ = 0.01875 / 1000 # (1 / m)
        ξ = (Z - Z10) * (r0 + Z10) / (r0 + Z)
        return Tinf - (Tinf - T10) * exp(-λ * ξ)
    end
end

# based on David's code, which was based on Regan
# M and log(P) are interpolated quadratically
const Ztableupper = [86000.0, 87000.0, 88000.0, 89000.0, 90000.0, 91000.0,
    93000.0, 95000.0, 97000.0, 99000.0, 101000.0, 103000.0, 105000.0, 107000.0,
    109000.0, 110000.0, 111000.0, 112000.0, 113000.0, 114000.0, 115000.0,
    116000.0, 117000.0, 118000.0, 119000.0, 120000.0, 125000.0, 130000.0,
    135000.0, 140000.0, 145000.0, 150000.0, 160000.0, 170000.0, 180000.0,
    190000.0, 200000.0, 210000.0, 220000.0, 230000.0, 240000.0, 250000.0,
    260000.0, 270000.0, 280000.0, 290000.0, 300000.0, 310000.0, 320000.0,
    330000.0, 340000.0, 350000.0, 360000.0, 370000.0, 380000.0, 390000.0,
    400000.0, 410000.0, 420000.0, 430000.0, 440000.0, 450000.0, 460000.0,
    470000.0, 480000.0, 490000.0, 500000.0, 525000.0, 550000.0, 575000.0,
    600000.0, 625000.0, 650000.0, 675000.0, 700000.0, 725000.0, 750000.0,
    775000.0, 800000.0, 825000.0, 850000.0, 875000.0, 900000.0, 925000.0,
    950000.0, 975000.0, 1000000.0] # (m)
const Ptableupper = [3.7338E-1, 3.1259E-1, 2.6173E-1, 2.1919E-1, 1.8359E-1,
    1.5381E-1, 1.0801E-1, 7.5966E-2, 5.3571E-2, 3.7948E-2, 2.7192E-2,
    1.9742E-2, 1.4477E-2, 1.0751E-2, 8.1142E-3, 7.1042E-3, 6.2614E-3,
    5.5547E-3, 4.9570E-3, 4.4473E-3, 4.0096E-3, 3.6312E-3, 3.3022E-3,
    3.0144E-3, 2.7615E-3, 2.5382E-3, 1.7354E-3, 1.2505E-3, 9.3568E-4,
    7.2028E-4, 5.6691E-4, 4.5422E-4, 3.0395E-4, 2.1210E-4, 1.5271E-4,
    1.1266E-4, 8.4736E-5, 6.4756E-5, 5.0149E-5, 3.9276E-5, 3.1059E-5,
    2.4767E-5, 1.9894E-5, 1.6083E-5, 1.3076E-5, 1.0683E-5, 8.7704E-6,
    7.2285E-6, 5.9796E-6, 4.9630E-6, 4.1320E-6, 3.4498E-6, 2.8878E-6,
    2.4234E-6, 2.0384E-6, 1.7184E-6, 1.4518E-6, 1.2291E-6, 1.0427E-6,
    8.8645E-7, 7.5517E-7, 6.4468E-7, 5.5155E-7, 4.7292E-7, 4.0642E-7,
    3.5011E-7, 3.0236E-7, 2.1200E-7, 1.5137E-7, 1.1028E-7, 8.2130E-8,
    6.2601E-8, 4.8865E-8, 3.9048E-8, 3.1908E-8, 2.6611E-8, 2.2599E-8,
    1.9493E-8, 1.7036E-8, 1.5051E-8, 1.3415E-8, 1.2043E-8, 1.0873E-8,
    9.8635E-9, 8.9816E-9, 8.2043E-9, 7.5138E-9] # (Pa)
const logPtableupper = log.(Ptableupper)
const Mtableupper = [28.95, 28.95, 28.94, 28.93, 28.91, 28.89, 28.82, 28.73,
    28.62, 28.48, 28.30, 28.10, 27.88, 27.64, 27.39, 27.27, 27.14, 27.02,
    26.90, 26.79, 26.68, 26.58, 26.48, 26.38, 26.29, 26.20, 25.80, 25.44,
    25.09, 24.75, 24.42, 24.10, 23.49, 22.90, 22.34, 21.81, 21.30, 20.83,
    20.37, 19.95, 19.56, 19.19, 18.85, 18.53, 18.24, 17.97, 17.73, 17.50,
    17.29, 17.09, 16.91, 16.74, 16.57, 16.42, 16.27, 16.13, 15.98, 15.84,
    15.70, 15.55, 15.40, 15.25, 15.08, 14.91, 14.73, 14.54, 14.33, 13.76,
    13.09, 12.34, 11.51, 10.62, 9.72, 8.83, 8.00, 7.24, 6.58, 6.01, 5.54, 5.16,
    4.85, 4.60, 4.40, 4.25, 4.12, 4.02, 3.94] # (kg / kmol)

function interpolation_index(Z)
    # Find the index for the lower side of the altitude interval
    i = 1
    while i < length(Ztableupper) && Z > Ztableupper[i+1]
        i += 1
    end

    # We are going to reference all elements from i - 1 to i + 1, so we need to
    # adjust the index away from the boundaries
    i == 1 && (i = 2)

    i
end

function interpolation_scale_factors(i, Z)
    Z0 = Ztableupper[i - 1]
    Z1 = Ztableupper[i]
    Z2 = Ztableupper[i + 1]

    scale0 = (Z - Z1) * (Z - Z2) / ((Z0 - Z1) * (Z0 - Z2))
    scale1 = (Z - Z0) * (Z - Z2) / ((Z1 - Z0) * (Z1 - Z2))
    scale2 = (Z - Z0) * (Z - Z1) / ((Z2 - Z0) * (Z2 - Z1))

    scale0, scale1, scale2
end

function pressure_upper(Z)
    i = interpolation_index(Z)
    scale0, scale1, scale2 = interpolation_scale_factors(i, Z)
    logP0 = logPtableupper[i - 1]
    logP1 = logPtableupper[i]
    logP2 = logPtableupper[i + 1]
    logP = logP0 * scale0 + logP1 * scale1 + logP2 * scale2
    exp(logP)
end

function mean_molecular_weight_upper(Z)
    i = interpolation_index(Z)
    scale0, scale1, scale2 = interpolation_scale_factors(i, Z)
    M0 = Mtableupper[i - 1]
    M1 = Mtableupper[i]
    M2 = Mtableupper[i + 1]
    M0 * scale0 + M1 * scale1 + M2 * scale2
end

const speed_of_sound_86km = let
    Z = 86_000.0 # (m)
    H = geopotential_altitude(Z)
    M = mean_molecular_weight_lower(Z)
    T = temperature_lower(H, M)
    speed_of_sound_lower(T, M)
end

struct State
    mean_molecular_weight::Float64
    temperature::Float64
    pressure::Float64
    speed_of_sound::Float64
end

mean_molecular_weight(s::State) = s.mean_molecular_weight
temperature(s::State) = s.temperature
pressure(s::State) = s.pressure
density(s::State) = pressure(s) * mean_molecular_weight(s) / (Rstar * temperature(s))
speed_of_sound(s::State) = s.speed_of_sound

function check_altitude(Z)
    Z < -5000 && error("altitude below lower bound of -5000 m")
    Z > 1_000_000 && error("altitude above upper bound of 1000000 m")
end

function atmosphere(Z)
    check_altitude(Z)
    if Z < 86_000
        H = geopotential_altitude(Z)
        M = mean_molecular_weight_lower(Z)
        T = temperature_lower(H, M)
        P = pressure_lower(H)
        c = speed_of_sound_lower(T, M)
    else
        T = temperature_upper(Z)
        P = pressure_upper(Z)
        M = mean_molecular_weight_upper(Z)
        c = speed_of_sound_86km
    end
    State(M, T, P, c)
end

end


text:
# This function converts GPS week and day to T_TAI (MJD referenced to International Atomic Time)
function gps2MJD(week, day)
    # GPS Epoch in Mod. Julian Date
    gpsMJDEpoch = 2444244.5 - 2400000.5

    # Add GPS week and days
    MJD = gpsMJDEpoch + week*7.0 + day + 19.0 / 86400.0

    return MJD
end

max_stars_repo_name: GrantHecht/OptimalEstimationProject


text:
using Spider
using Test

@testset "Basic site generation" begin
  args = (source_dir="basic_src",
          output_dir="basic_output",
          clear_output_dir=true,
          header_file="basic_header.html",
          footer_file="basic_footer.html")

  run_spider(;args...)
  # Check html generated from Markdown:
  @test isfile("basic_output/index.html")
  # Check non-Markdown files copied:
  @test isfile("basic_output/style.css")
end

run(`rm -fr basic_output`)

nothing


text:
function LHop!(LH,vbasis,sbasis,quad,linemap::LineMap,NED,scale,nhdofs)

    detjac = determinant_jacobian(linemap)
    for nk in NED
        for (p,w) in quad
            vals = vbasis(linemap(p...))
            svals = sbasis(p)

            M = make_row_matrix(vals,nk)
            N = interpolation_matrix(svals,nhdofs)

            LH .+= M'*N*detjac*scale*w
        end
    end
end

function LHop!(LH,vbasis,sbasis,quad,imap::InterpolatingPolynomial,normals,
    ED,scale,nhdofs)

    @assert length(scale) == length(quad)
    @assert size(normals)[2] == length(quad)

    for (idx,(p,w)) in enumerate(quad)
        n = view(normals,:,idx)
        vals = vbasis(imap(p))
        svals = sbasis(p)
        detjac = determinant_jacobian(imap,p)

        for k in 1:length(ED)
            NED = n[k]*ED[k]

            M = make_row_matrix(vals,NED)
            N = interpolation_matrix(svals,nhdofs)

            LH .+= M'*N*detjac*scale[idx]*w

        end
    end
end

function LHop(vbasis,sbasis,facequads,facemaps,normals,Dhalf,cellmap)

    nfaces = length(facemaps)
    @assert length(normals) == nfaces
    @assert length(facequads) == nfaces

    dim = dimension(vbasis)
    sdim = symmetric_tensor_dimension(dim)
    Ek = vec_to_symm_mat_converter(dim)

    NLF = number_of_basis_functions(vbasis)
    NHF = number_of_basis_functions(sbasis)

    isactiveface = active_faces(facequads)

    LH = [zeros(sdim*NLF,dim*NHF) for i in 1:count(isactiveface)]
    facescale = face_determinant_jacobian(cellmap)

    counter = 1
    for faceid in 1:nfaces
        if isactiveface[faceid]
            n = normals[faceid]
            NED = [n[k]*Ek[k]'*Dhalf for k in 1:length(Ek)]
            LHop!(LH[counter],vbasis,sbasis,facequads[faceid],facemaps[faceid],
                NED,facescale[faceid],dim)
            counter += 1
        end
    end
    return LH
end

function LHop_on_interface(vbasis,sbasis,iquad,imap,normals,Dhalf,cellmap)

    dim = dimension(vbasis)
    sdim = symmetric_tensor_dimension(dim)
    nq = length(iquad)

    @assert size(Dhalf) == (sdim,sdim)
    @assert size(normals) == (dim,nq)

    Ek = vec_to_symm_mat_converter(dim)
    NLF = number_of_basis_functions(vbasis)
    NHF = number_of_basis_functions(sbasis)

    LH = zeros(sdim*NLF,dim*NHF)
    scale = scale_area(cellmap,normals)

    ED = [E'*Dhalf for E in Ek]
    LHop!(LH,vbasis,sbasis,iquad,imap,normals,ED,scale,dim)

    return LH
end

function UHop!(UH,vbasis,sbasis,quad,linemap::LineMap,
    scale,nudofs)

    detjac = determinant_jacobian(linemap)
    for (p,w) in quad
        vals = vbasis(linemap(p...))
        svals = sbasis(p)

        Nv = interpolation_matrix(vals,nudofs)
        Ns = interpolation_matrix(svals,nudofs)

        UH .+= Nv'*Ns*detjac*scale*w
    end
end

function UHop!(UH,vbasis,sbasis,quad,imap::InterpolatingPolynomial,
    scale,nudofs)

    @assert length(scale) == length(quad)

    for (idx,(p,w)) in enumerate(quad)
        vals = vbasis(imap(p))
        svals = sbasis(p)
        detjac = determinant_jacobian(imap,p)

        Nv = interpolation_matrix(vals,nudofs)
        Ns = interpolation_matrix(svals,nudofs)

        UH .+= Nv'*Ns*detjac*scale[idx]*w
    end
end

function UHop(vbasis,sbasis,facequads,facemaps,stabilization,cellmap)

    nfaces = length(facemaps)
    @assert length(facequads) == nfaces

    dim = dimension(vbasis)
    sdim = dimension(sbasis)
    @assert sdim == dim-1

    NUF = number_of_basis_functions(vbasis)
    NHF = number_of_basis_functions(sbasis)

    isactiveface = [length(fq) > 0 ? true : false for fq in facequads]

    facejac = face_determinant_jacobian(cellmap)
    UH = [zeros(dim*NUF,dim*NHF) for i = 1:count(isactiveface)]

    counter = 1
    for faceid in 1:nfaces
        if isactiveface[faceid]
            UHop!(UH[counter],vbasis,sbasis,facequads[faceid],facemaps[faceid],
                stabilization*facejac[faceid],dim)
            counter += 1
        end
    end
    return UH
end

function UHop_on_interface(vbasis,sbasis,squad,imap,normals,
    stabilization,cellmap)

    dim = dimension(vbasis)
    sdim = dimension(sbasis)
    @assert sdim == dim-1

    NUF = number_of_basis_functions(vbasis)
    NHF = number_of_basis_functions(sbasis)

    UH = zeros(dim*NUF,dim*NHF)
    scale = stabilization*scale_area(cellmap,normals)

    UHop!(UH,vbasis,sbasis,squad,imap,scale,dim)

    return UH
end

function local_hybrid_operator(vbasis,sbasis,facequads,facemaps,normals,
    Dhalf,stabilization,cellmap)

    LH = LHop(vbasis,sbasis,facequads,facemaps,normals,Dhalf,cellmap)
    UH = UHop(vbasis,sbasis,facequads,facemaps,stabilization,cellmap)

    @assert length(LH) == length(UH)

    return [[LH[i];UH[i]] for i in 1:length(LH)]

end

function local_hybrid_operator_on_interface(vbasis,sbasis,iquad,imap,inormals,
    Dhalf,stabilization,cellmap)

    LH = LHop_on_interface(vbasis,sbasis,iquad,imap,inormals,Dhalf,cellmap)
    UH = UHop_on_interface(vbasis,sbasis,iquad,imap,inormals,
        stabilization,cellmap)

    return [LH;UH]
end


text:
import .AxisKeys
const AK = AxisKeys

function from_namedtuple(::Type{T}, nt::NamedTuple) where {T<:AK.KeyedArray}
    AK.KeyedArray(nt.values; nt.axes...)::T
end

namedaxes(ak::AK.KeyedArray) = AK.named_axiskeys(ak)
values(ak::AK.KeyedArray) = parent(parent(ak))

roottype(::Type{T}) where {T<:AK.KeyedArray} = AK.KeyedArray


text:
#=

    statistical.jl

    Contains code to simulate the statistical model.

    Author:     Alexander P. Browning
                ======================
                School of Mathematical Sciences
                Queensland University of Technology
                ======================
                ap.browning@icloud.com
                alexbrowning.me

=#
"""
    simulate_model(θ,t,d,E;N=1000)

Simulate the full statistical model for 2N cells, returning `[[Q̄,Ū],[Q,U]]` where
`[Q̄,Ū]` are paired measurements of the quenched sample, and `[Q,U]` of the unquenched
sample.

Parameter values are given as a vector `θ` where
    `α₁,α₂,σ₁,σ₂ = θ[1:4]`
are the observation process parameters,
    `d(θ[5:end-1])`
provides a multivariate distribution parameterised by hyperparameters `θ[5:end-1]`, and
    `p = θ[end]`
is the proportion recycled.
"""
function simulate_model(θ::Vector,t::Number,d::Function,E;N::Int=1000,η::Number=1.0)

    # Parameters
    α₁,α₂,σ₁,σ₂ = θ[1:4]        # Observation process parameters
    ξᵢ = rand(d(θ[5:end-1]),2N) # Sample from parameter distribution
    p  = θ[end]                 # Proportion recycled
    Rᵢ = @view ξᵢ[1,:]
    λᵢ = @view ξᵢ[2,:]
    βᵢ = @view ξᵢ[3,:]

    # Sample Aᵢ and Iᵢ
    Aᵢ = similar(λᵢ)
    Iᵢ = similar(λᵢ)
    for i = 1:2N
        Aᵢ[i],Iᵢ[i] = solve_ode_model(t,[λᵢ[i],βᵢ[i],p])
    end

    # Sample errors
    idx1 = sample(1:length(E[1]),N)
    idx2 = sample(1:length(E[2]),N)

    V1 = max.(0.0,Aᵢ[1:N] .* Rᵢ[1:N])    # Total antibody
    V2 = max.(0.0,Aᵢ[N+1:end] .* Rᵢ[N+1:end])    # Total antibody
    V̄2 = max.(0.0,Iᵢ[N+1:end] .* Rᵢ[N+1:end] .+ (1 - η) * (Aᵢ[N+1:end] - Iᵢ[N+1:end]))    # Total internal antibody

    # Unquenched
    Q = α₁ * V1 + σ₁ * sqrt.(α₁ * V1) .* randn(N) + E[1][idx1]
    U = α₂ * V1 + σ₂ * sqrt.(α₂ * V1) .* randn(N) + E[2][idx1]

    # Quenched
    Q̄ = α₁ * V̄2 + σ₁ * sqrt.(α₁ * V̄2) .* randn(N) + E[1][idx2]
    Ū = α₂ * V2 + σ₂ * sqrt.(α₂ * V2) .* randn(N) + E[2][idx2]

    return [[Q̄,Ū],[Q,U]]

end
simulate_model(θ::Vector,T::Vector,d::Function,E;kwargs...) = [simulate_model(θ,t,d,E;kwargs...) for t in T]

function simulate_model_noiseless(θ::Vector,T::Vector,d::Function,E=nothing;N::Int=1000,η::Number=1.0)

    # Parameters
    ξᵢ = rand(d(θ[5:end-1]),N) # Sample from parameter distribution
    p  = θ[end]                 # Proportion recycled
    Rᵢ = @view ξᵢ[1,:]
    λᵢ = @view ξᵢ[2,:]
    βᵢ = @view ξᵢ[3,:]

    function sample_quenched_quenchable_signal(t)
        Aᵢ = similar(λᵢ)
        Iᵢ = similar(λᵢ)
        for i = 1:N
            Aᵢ[i],Iᵢ[i] = solve_ode_model(t,[λᵢ[i],βᵢ[i],p])
        end
        return Iᵢ .* Rᵢ .+ (1 - η) * (Aᵢ - Iᵢ)
    end

    return [sample_quenched_quenchable_signal(t) for t in T]

end

function parameters_to_antibody(t::Number,λ::Vector,β::Vector,p::Number)
    Aᵢ = similar(λ)
    Iᵢ = similar(λ)
    @simd for i = 1:length(Aᵢ)
        Aᵢ[i],Iᵢ[i] = solve_ode_model(t,[λ[i],β[i],p])
    end
    return Aᵢ,Iᵢ
end
parameters_to_antibody(T::Vector,λ::Vector,β::Vector,p::Number) = [parameters_to_antibody(t,λ,β,p) for t in T]

max_stars_repo_name: ap-browning/internalisation

max_issues_repo_name: ap-browning/internalisation

max_forks_repo_name: ap-browning/internalisation


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

mutable struct ShareInfoElement <: SwaggerModel
    vmUri::Any # spec type: Union{ Nothing, String } # spec name: vmUri

    function ShareInfoElement(;vmUri=nothing)
        o = new()
        validate_property(ShareInfoElement, Symbol("vmUri"), vmUri)
        setfield!(o, Symbol("vmUri"), vmUri)
        o
    end
end # type ShareInfoElement

const _property_map_ShareInfoElement = Dict{Symbol,Symbol}(Symbol("vmUri")=>Symbol("vmUri"))
const _property_types_ShareInfoElement = Dict{Symbol,String}(Symbol("vmUri")=>"String")
Base.propertynames(::Type{ ShareInfoElement }) = collect(keys(_property_map_ShareInfoElement))
Swagger.property_type(::Type{ ShareInfoElement }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_ShareInfoElement[name]))}
Swagger.field_name(::Type{ ShareInfoElement }, property_name::Symbol) =  _property_map_ShareInfoElement[property_name]

function check_required(o::ShareInfoElement)
    true
end

function validate_property(::Type{ ShareInfoElement }, name::Symbol, val)
end


text:
"Abstract type for ASPRS LAS point data record formats 0 - 3"
abstract type LasPoint end

"""
Custom PointVector struct for memory mapped LasPoints.
Inspiration taken from UnalignedVector.jl
and extended it with custom indexing and packing.
"""
struct PointVector{T} <: AbstractArray{T,1}
    io::IOBuffer
    n::Int
    pointsize::Int

    function PointVector{T}(data::Vector{UInt8}, pointsize::Integer) where {T}
        n = length(data) ÷ pointsize

        # IOBuffer takes (data, readable, writable)
        new{T}(IOBuffer(data; read=true, write=false), n, pointsize)
    end
end

Base.IndexStyle(::Type{PointVector{T}}) where {T} = IndexLinear()
Base.length(pv::PointVector) = pv.n
Base.size(pv::PointVector) = (length(pv),)

function Base.getindex(pv::PointVector{T}, i::Int) where {T}
    offset = (i-1) * pv.pointsize  # seeking uses 0 based indexing!
    position(pv.io) != offset && seek(pv.io, offset)
    read(pv.io, T)
end

function Base.setindex!(pv::PointVector{T}, val::T, i::Int) where {T}
    # offset = (i-1) * pv.pointsize  # seeking uses 0 based indexing!
    # position(pv.io) != offset && seek(pv.io, offset)
    # write(pv.io, val)
    error("Can't write to read only memory mapped file.")
end

function Base.show(io::IO, pointdata::AbstractVector{<:LasPoint})
    n = size(pointdata, 1)
    println(io, "$(typeof(pointdata)) with $n points.")
end

"ASPRS LAS point data record format 0"
@gen_io struct LasPoint0 <: LasPoint
    x::Int32
    y::Int32
    z::Int32
    intensity::UInt16
    flag_byte::UInt8
    raw_classification::UInt8
    scan_angle::Int8
    user_data::UInt8
    pt_src_id::UInt16
end

"ASPRS LAS point data record format 1"
@gen_io struct LasPoint1 <: LasPoint
    x::Int32
    y::Int32
    z::Int32
    intensity::UInt16
    flag_byte::UInt8
    raw_classification::UInt8
    scan_angle::Int8
    user_data::UInt8
    pt_src_id::UInt16
    gps_time::Float64
end

"ASPRS LAS point data record format 2"
@gen_io struct LasPoint2 <: LasPoint
    x::Int32
    y::Int32
    z::Int32
    intensity::UInt16
    flag_byte::UInt8
    raw_classification::UInt8
    scan_angle::Int8
    user_data::UInt8
    pt_src_id::UInt16
    red::N0f16
    green::N0f16
    blue::N0f16
end

"ASPRS LAS point data record format 3"
@gen_io struct LasPoint3 <: LasPoint
    x::Int32
    y::Int32
    z::Int32
    intensity::UInt16
    flag_byte::UInt8
    raw_classification::UInt8
    scan_angle::Int8
    user_data::UInt8
    pt_src_id::UInt16
    gps_time::Float64
    red::N0f16
    green::N0f16
    blue::N0f16
end

# for convenience in function signatures
const LasPointColor = Union{LasPoint2,LasPoint3}
const LasPointTime = Union{LasPoint1,LasPoint3}

function Base.show(io::IO, p::LasPoint)
    x = Int(p.x)
    y = Int(p.y)
    z = Int(p.z)
    cl = Int(classification(p))
    println(io, "LasPoint(x=$x, y=$y, z=$z, classification=$cl)")
end

# Extend base by enabling reading/writing relevant FixedPointNumbers from IO.
Base.read(io::IO, ::Type{N0f16}) = reinterpret(N0f16, read(io, UInt16))
Base.write(io::IO, t::N0f16) = write(io, reinterpret(UInt16, t))

# functions for IO on points

"X coordinate (Float64), apply scale and offset according to the header"
xcoord(p::LasPoint, h::LasHeader) = muladd(p.x, h.x_scale, h.x_offset)
"Y coordinate (Float64), apply scale and offset according to the header"
ycoord(p::LasPoint, h::LasHeader) = muladd(p.y, h.y_scale, h.y_offset)
"Z coordinate (Float64), apply scale and offset according to the header"
zcoord(p::LasPoint, h::LasHeader) = muladd(p.z, h.z_scale, h.z_offset)

# inverse functions of the above
"X value (Int32), as represented in the point data, reversing the offset and scale from the header"
xcoord(x::Real, h::LasHeader) = round(Int32, (x - h.x_offset) / h.x_scale)
"Y value (Int32), as represented in the point data, reversing the offset and scale from the header"
ycoord(y::Real, h::LasHeader) = round(Int32, (y - h.y_offset) / h.y_scale)
"Z value (Int32), as represented in the point data, reversing the offset and scale from the header"
zcoord(z::Real, h::LasHeader) = round(Int32, (z - h.z_offset) / h.z_scale)

# functions to access common LasPoint fields
"Integer representation of the pulse return magnitude."
intensity(p::LasPoint) = p.intensity
"Angle at which the laser point was output, including the roll of the aircraft."
scan_angle(p::LasPoint) = p.scan_angle
"This field may be used at the user’s discretion."
user_data(p::LasPoint) = p.user_data
"This value indicates the file from which this point originated."
pt_src_id(p::LasPoint) = p.pt_src_id

# time
# 1.3159648e9 = 315964800.0 + 1.0e9
# 1.0e9: see LAS spec (for float accuracy)
# 315964800: seconds between 1970 (Unix epoch) and 1980 (GPS epoch), from http://stackoverflow.com/a/20528332
"""Get the DateTime that the point was collected.
Assumes time is recorded in Adjusted Standard GPS Time; see `is_standard_gps`.
Note that DateTime has millisecond precision, any higher precision is lost."""
Dates.DateTime(p::LasPointTime) = Dates.unix2datetime(p.gps_time) + Dates.Second(1315964800)
# A conversion of the GPS Week Time format to DateTime is not yet implemented

"""Time tag value at which the point was aquired,
a Float64; see `is_standard_gps` for what it represents"""
gps_time(p::LasPointTime) = p.gps_time
"""Convert DateTime to GPS Float64, as represented in the point data,
assumes time is recorded in Adjusted Standard GPS Time; see `is_standard_gps`"""
gps_time(d::DateTime) = Dates.datetime2unix(d - Dates.Second(1315964800))

# color
"The red image channel value associated with this point"
ColorTypes.red(p::LasPointColor) = p.red
"The green image channel value associated with this point"
ColorTypes.green(p::LasPointColor) = p.green
"The blue image channel value associated with this point"
ColorTypes.blue(p::LasPointColor) = p.blue
"The RGB color associated with this point"
ColorTypes.RGB(p::LasPointColor) = RGB(red(p), green(p), blue(p))

# functions to extract sub-byte items from a LasPoint's flag_byte
"The pulse return number for a given output pulse, starting at one."
return_number(p::LasPoint) = (p.flag_byte & 0b00000111)
"The total number of returns for a given pulse."
number_of_returns(p::LasPoint) = (p.flag_byte & 0b00111000) >> 3
"If true, the scanner mirror was traveling from left to right at the time of the output pulse."
scan_direction(p::LasPoint) = Bool((p.flag_byte & 0b01000000) >> 6)
"If true, it is the last point before the scanner changes direction."
edge_of_flight_line(p::LasPoint) = Bool((p.flag_byte & 0b10000000) >> 7)

"Flag byte, contains return number, number of returns, scan direction flag and edge of flight line"
flag_byte(p::LasPoint) = p.flag_byte
"Flag byte, as represented in the point data, built up from components"
function flag_byte(return_number::UInt8, number_of_returns::UInt8,
                   scan_direction::Bool, edge_of_flight_line::Bool)
    # Bool to UInt8 conversion because a bit shift on a Bool produces an Int
    UInt8(edge_of_flight_line) << 7 | UInt8(scan_direction) << 6 | number_of_returns << 3 | return_number
end

# functions to extract sub-byte items from a LasPoint's raw_classification
"Classification value as defined in the ASPRS classification table."
classification(p::LasPoint) = (p.raw_classification & 0b00011111)
"If true, the point was not created from lidar collection"
synthetic(p::LasPoint) = Bool((p.raw_classification & 0b00100000) >> 5)
"If true, this point is considered to be a model key-point."
key_point(p::LasPoint) = Bool((p.raw_classification & 0b01000000) >> 6)
"If true, this point should not be included in processing"
withheld(p::LasPoint) = Bool((p.raw_classification & 0b10000000) >> 7)

"Raw classification byte, contains classification, synthetic, key point and withheld"
raw_classification(p::LasPoint) = p.raw_classification
"Raw classification byte, as represented in the point data, built up from components"
function raw_classification(classification::UInt8, synthetic::Bool,
                            key_point::Bool, withheld::Bool)
    UInt8(withheld) << 7 | UInt8(key_point) << 6 | UInt8(synthetic) << 5 | classification
end

function Base.convert(::Type{Point{3, Float64}}, p::LasPoint, h::LasHeader)
    Point{3, Float64}(xcoord(p, h), ycoord(p, h), zcoord(p, h))
end

# beware of the limited precision, for instance with UTM coordinates
function Base.convert(::Type{Point{3, Float32}}, p::LasPoint, h::LasHeader)
    Point{3, Float32}(xcoord(p, h), ycoord(p, h), zcoord(p, h))
end


text:
module ITAE

export itae

"""
    ```julia
    itae(error, time)
    ```
    ITAE (Integral of Time-Weighted Absolute Error) from error vector.
    Inputs:
    error -> error vector.
    time  -> time vector.
"""
function itae(error::Vector{T}, time::Vector{T}) where T
    sum(time .* abs.(error))
end

"""
    ```julia
    itae(error, time)
    ```
    ITAE (Integral of Time-Weighted Absolute Error) from vector of error vector.
    Inputs:
    error -> vector of error vector.
    time  -> time vector.
"""
function itae(error::Vector{Vector{T}}, time::Vector{T}) where T
    map(x->sum(time .* abs.(x)), error)
end

"""
    ```julia
    itae(x, t, desired)
    ```
    ITAE (Integral of Time-Weighted Absolute Error) from error betwwen desired value and controler output.
    Inputs:
    x -> controler output vector.
    t -> time vector.
    desired -> desired value function.
"""
function itae(x::Vector{T}, t::Vector{T}, desired::Function) where T
    x = x .- map(desired, t)
    sum(t .* abs.(x))
end

end