text:
include("log.jl")
include("sort.jl")
include("print.jl")
include("material.jl")
include("select_type.jl")
include("tree.jl")

max_stars_repo_name: khurrumsaleem/MOCNeutronTransport

max_issues_repo_name: khurrumsaleem/MOCNeutronTransport

max_forks_repo_name: khurrumsaleem/MOCNeutronTransport


text:
const subscription_tracker = Ref{Dict}(Dict())

"""
    open_subscription(fn::Function,
                      [client::Client],
                      subscription_name::Union{Alias, AbstractString},
                      output_type::Type=Any;
                      sub_args=Dict(),
                      output_fields=String[],
                      initfn=nothing,
                      retry=true,
                      subtimeout=0,
                      stopfn=nothing,
                      throw_on_execution_error=false)

Subscribe to `subscription_name`, running `fn` on each received result and ending the
subcription when `fn` returns `true`.

By default `fn` receives a `GQLReponse{Any}`, where the data for an individual result
object can be found by `gql_response.data[subscription_name]`.

If used, `initfn` is called once the subscription is open.

The subscription uses the `ws_endpoint` field of the `client.`

This function is designed to be used with the `do` keyword.

# Arguments
- `fn::Function`: function to be run on each result, recieves the response from the
    subscription`. Must return a boolean to indicate whether or not to close the subscription,
    with `true` closing the subscription.
- `client::Client`: GraphQL client (optional). If not supplied, [`global_graphql_client`](@ref) is used.
- `subscription_name::Union{Alias, AbstractString}`: name of subscription in server.
- `output_type::Type=Any`: output data type for subscription response object. An object
    of type `GQLResponse{output_type}` will be returned.For further information, see
    documentation for `GQLResponse`.

# Keyword Arguments
- `sub_args=Dict()`: dictionary of subscription argument key value pairs - can be
    nested with dictionaries and vectors.
- `output_fields=String[]`: output fields to be returned. Can be a string, or
    composed of dictionaries and vectors.
- `initfn=nothing`: optional function to be run once subscription is itialised.
- `retry=true`: retry if subscription fails to open.
- `subtimeout=0`: if `stopfn` supplied, this is the period that it is called at.
    If `stopfn` is not supplied, this is the timeout for waiting for data. The timer
    is reset after every subscription result is received.
- `stopfn=nothing`: a function to be called every `subtimeout` that stops the
    subscription if it returns positive. The timer is reset after every subscription
    result is received.
- `throw_on_execution_error=false`: set to `true` to stop an error being thrown if the GraphQL server
    response contains errors that occurred during execution.
- `verbose=0`: set to 1, 2 for extra logging.

# Examples
```julia
julia> open_subscription("subSaveUser", sub_args=Dict("role" => "SYSTEM_ADMIN")) do result
           fn(result)
       end
```

See also: [`GQLResponse`](@ref)
"""
function open_subscription(fn::Function,
                           subscription_name::Union{Alias, AbstractString},
                           output_type::Type=Any;
                           kwargs...)
    return open_subscription(fn, global_graphql_client(), subscription_name, output_type; kwargs...)
end
function open_subscription(fn::Function,
                           client::Client,
                           subscription_name::Union{Alias, AbstractString},
                           output_type::Type=Any;
                           sub_args=Dict(),
                           output_fields=String[],
                           initfn=nothing,
                           retry=true,
                           subtimeout=0,
                           stopfn=nothing,
                           throw_on_execution_error=false,
                           verbose=0)

    !in(get_name(subscription_name), get_subscriptions(client)) && throw(GraphQLError("$(get_name(subscription_name)) is not an existing subscription"))

    output_str = get_output_str(output_fields)
    payload = get_generic_query_payload(client, "subscription", subscription_name, sub_args, output_str)

    sub_id = string(length(keys(subscription_tracker[])) + 1)
    sub_id *= "-" * string(Threads.threadid())
    message = Dict(
        "id" => string(sub_id),
        "type" => "start",
        "payload" => payload
    )
    message_str = JSON3.write(message)

    HTTP.WebSockets.open(client.ws_endpoint; retry=retry, headers=client.headers) do ws
        # Start sub
        output_info(verbose) && println("Starting $(get_name(subscription_name)) subscription with ID $sub_id")
        write(ws, message_str)
        subscription_tracker[][sub_id] = "open"

        # Init function
        if !isnothing(initfn)
            output_debug(verbose) && println("Running subscription initialisation function")
            initfn()
        end

        # Get listening
        output_debug(verbose) && println("Listening to $(get_name(subscription_name)) with ID $sub_id...")

        # Run function
        finish = false
        while !finish
            data = readfromwebsocket(ws, stopfn, subtimeout)
            if data === :timeout
                output_info(verbose) && println("Subscription $sub_id timed out")
                break
            elseif data === :stopfn
                output_info(verbose) && println("Subscription $sub_id stopped by the stop function supplied")
                break
            end
            response = JSON3.read(data::Vector{UInt8}, GQLSubscriptionResponse{output_type})
            payload = response.payload
            if !isnothing(payload.errors) && !isempty(payload.errors) && throw_on_execution_error
                subscription_tracker[][sub_id] = "errored"
                throw(GraphQLError("Error during subscription.", payload))
            end
            # Handle multiple subs, do we need this?
            if response.id == string(sub_id)
                output_debug(verbose) && println("Result recieved on subscription with ID $sub_id")
                finish = fn(payload)
                if !isa(finish, Bool)
                    subscription_tracker[][sub_id] = "errored"
                    error("Subscription function must return a boolean")
                end
            end
        end
    end
    output_debug(verbose) && println("Finished. Closing subscription")
    subscription_tracker[][sub_id] = "closed"
    return
end

"""
    clear_subscriptions()

Removes all subscriptions from the `subscription_tracker`, throwing an error if any are still open.
"""
function clear_subscriptions()
    for (sub_id, val) in subscription_tracker[]
        val == "open" ?
            throw(GraphQLError("Subscription $sub_id is still open and cannot be cleared")) :
            delete!(subscription_tracker[], sub_id)
    end
end

function async_reader_with_timeout(io::IO, subtimeout)::Channel
    ch = Channel(1)
    task = @async begin
        reader_task = current_task()
        function timeout_cb(timer)
            put!(ch, :timeout)
            Base.throwto(reader_task, InterruptException())
        end
        timeout = Timer(timeout_cb, subtimeout)
        data = readavailable(io)
        subtimeout > 0 && close(timeout) # Cancel the timeout
        put!(ch, data)
    end
    bind(ch, task)
    return ch
end

function async_reader_with_stopfn(io::IO, stopfn, checktime)::Channel
    ch = Channel(1) # Could we make this channel concretely typed?
    task = @async begin
        reader_task = current_task()
        function timeout_cb(timer)
            if stopfn()
                put!(ch, :stopfn)
                Base.throwto(reader_task, InterruptException())
            else
                timeout = Timer(timeout_cb, checktime)
            end
        end
        timeout = Timer(timeout_cb, checktime)
        data = readavailable(io)
        close(timeout) # Cancel the timeout
        put!(ch, data)
    end
    bind(ch, task)
    return ch
end

"""
    readfromwebsocket(ws::IO, stopfn, subtimeout)

Read from the websocket with the following logic:
- If `stopfn` is nothing and `subtimeout` is 0, use `readavailable`
    which blocks data is written to the stream.
- If `stopfn` is not nothing, check the value of `stopfn` periodically.
    If it returns true, the websocket is closed. The period is set to
    `subtimeout` if greater than 0, otherwise 2 seconds is used.
- If `stopfn` is nothing but `subtimeout` > 0, stop listening after
    `subtimeout` seconds if no data has been received.

A channel is returned with the data. If `stopfn` stops the websocket,
the data will be `:stopfn`. If the timeout stops the websocket,
the data will be `:timeout`
"""
function readfromwebsocket(ws::IO, stopfn, subtimeout)
    if isnothing(stopfn) && subtimeout > 0
        ch_out = async_reader_with_timeout(ws, subtimeout)
        data = take!(ch_out)
    elseif !isnothing(stopfn)
        checktime = subtimeout > 0 ? subtimeout : 2
        ch_out = async_reader_with_stopfn(ws, stopfn, checktime)
        data = take!(ch_out)
    else
        data = readavailable(ws)
    end
    return data
end

ext: jl


text:
#=
Motivation: I want to model psychometric data to some kind of sigmoid function, but the nature
of psychometric data is that it is forced choice.  This means that the output is binary.
GLM doesn't support forced choice models yet——however, with a little bit of work and some
type piracy, the results are attractive.
=#

using GLM, DataFrames, CSV, FreqTables, StatsPlots, Distributions, Statistics
const Dist = Distributions
# imports for overloading
import GLM: Link, Link01, linkfun, linkinv, mueta, inverselink

# obtain dataset
df_raw = DataFrame(CSV.File(download("https://gist.githubusercontent.com/jakewilliami/7a84208f33de8076c7b17abb52a9b242/raw/b27d63700f4072830251eb4063e790c27c3154d0/mcd.csv")))

# statistical notations
Φ(x::Real)   = Dist.cdf(Dist.Normal(), x)
Φ⁻¹(x::Real) = Dist.quantile(Dist.Normal(), x)
φ(x::Real)   = Dist.pdf(Dist.Normal(), x)

# define a new link function for 2-alternative forced choice
struct Probit2AFCLink <: Link end

# overloading
linkfun(::Probit2AFCLink, μ::Real) = Φ⁻¹(2 * max(μ, nextfloat(0.5)) - 1)
linkinv(::Probit2AFCLink, η::Real) = (1 + Φ(η)) / 2
mueta(::Probit2AFCLink,   η::Real) = φ(η) / 2
function inverselink(::Probit2AFCLink, η::Real)
    μ = (1 + Φ(η)) / 2
    d = φ(η) / 2
    return μ, d, oftype(μ, NaN)
end

# define models
logit(p) = log(p / (1 - p))
logit⁻¹(α) = 1 / (1 + exp(-α))
logit⁻¹(α) = logit⁻¹(α) * 0.5 + 0.5 # shift and squish for 2-AFC
logit2afc⁻¹(α) = 0.5 + 0.5 / (1 + exp(-α))
probit(p) = Dist.quantile(Dist.Normal(), p)
probit⁻¹(x) = Dist.cdf(Dist.Normal(), x)
Φ⁻¹(z) = probit(z) # statistical notation
Φ(α) = probit⁻¹(α) # statistical notation
probit2afc⁻¹(x) = probit⁻¹(x) * 0.5 + 0.5 # shift and squish for 2-AFC

df_pivot = combine(groupby(df_raw, :condition1)) do df
    μ = mean(df.correct)
    n = nrow(df)
    (
    correct_mean = μ,
    n = n,
    k = μ * n
    )
end

model = glm(@formula(correct ~ condition1) , df_raw, Binomial(), Probit2AFCLink())
a, b = coef(model)

theme(:solarized)
plot = @df df_pivot scatter(
        :condition1,
        :correct_mean,
        title = "Psychometric Curve of Motion Coherence",
        label = false,
        xaxis = "Coherence",
        yaxis = "Accuracy",
        fontfamily = font("Times")
    )

# add models to graph
plot!(plot, x -> probit2afc⁻¹(a + b*x), 0, 0.32, label = "Probit Link")
plot!(plot, x -> logit2afc⁻¹(a + b*x), 0, 0.32, label = "Logit Link")

max_stars_repo_name: Humans-of-Julia/Challenges

max_issues_repo_name: Humans-of-Julia/WeeklyContest

max_forks_repo_name: Humans-of-Julia/WeeklyContest


text:
"""
    EventTracker.event_table()

Export tracked events as a table.
"""
EventTracker.event_table

"""
    EventTracker.location_table()

Export tracked locations as a table.
"""
EventTracker.location_table

"""
    EventTracker.event_dataframe(; owntime = false)

Export tracked events as a `DataFrame`. If `owntime = true` is passed, include
`owntime` column by analyzing the nesting of the intervals and compute the
during excluding the time spent in the sub-intervals.
"""
EventTracker.event_dataframe

"""
    EventTracker.location_dataframe()

Export tracked locations as a `DataFrame`.
"""
EventTracker.location_dataframe

EventTracker.event_table() = append_events!(StructVector{Event}(undef, 0))
EventTracker.location_table() = append_locations!(StructVector{Location}(undef, 0))

EventTracker.event_dataframe(; owntime::Bool = false) =
    _event_dataframe(EventTracker.event_table(), owntime)
EventTracker.location_dataframe() =
    DataFrame(EventTracker.location_table(); copycols = true)

function _event_dataframe(etbl, include_owntime::Bool, ltbl = nothing)
    edf = DataFrame(etbl; copycols = true)
    edf[!, :time] = (something.(edf.stop, missing) .- edf.start) ./ 1e9
    if include_owntime
        if ltbl === nothing
            ltbl = EventTracker.location_table()
        end
        stks = EventTracker.stacks(etbl, ltbl)
        event_to_owntime = owntime_dict(stks)
        edf[!, :owntime] = map(edf.eventid) do i
            v = get(event_to_owntime, i, nothing)
            v === nothing && return missing
            return v.owntime / 1e9
        end
    end
    return edf
end

function move_to_end!(cols, name)
    i = findfirst(==(name), cols)
    @assert i !== nothing
    deleteat!(cols, i)
    push!(cols, name)
end

"""
    EventTracker.summary_dataframe(; owntime = false)

Compute the summary of the events.

It computes the summary statistics of the events for each location and join the
location information.
"""
EventTracker.summary_dataframe
function EventTracker.summary_dataframe(; owntime::Bool = false, sortby = nothing)
    etbl = EventTracker.event_table()
    ltbl = EventTracker.location_table()
    edf = _event_dataframe(etbl, owntime, ltbl)
    ldf = DataFrame(ltbl; copycols = true)

    combine_spec = Pair[]
    append!(combine_spec, [
        :time => length => :ncalls,
        # average duration, including sub-calls:
        :time => mean => :time,
    ])
    if owntime
        push!(combine_spec, :owntime => mean => :owntime)
    end
    append!(
        combine_spec,
        [
            # sort by this to approximate call tree?
            :start => minimum => :firstcall,
            :eventkind => first => :eventkind,
        ],
    )
    summary_df = combine(groupby(edf, :locationid), combine_spec)
    summary_df = leftjoin(summary_df, ldf, on = :locationid)

    # Reorder columns so that important information comes first:
    cols = propertynames(summary_df)
    move_to_end!(cols, :firstcall)
    move_to_end!(cols, :eventkind)
    move_to_end!(cols, :locationid)
    summary_df = select!(summary_df, cols)

    if sortby !== nothing
        sort!(summary_df, sortby)
    end

    return summary_df
end

"""
    EventTracker.summary()

Compute and print the summary. Currently just pints
[`EventTracker.summary_dataframe`](@ref).
"""
EventTracker.summary
function EventTracker.summary(; kwargs...)
    summary_df = EventTracker.summary_dataframe(; kwargs...)
    display(summary_df)
end


text:
module FunctionsTest

using Test
using LinearAlgebra
using StaticArrays: @SVector, SVector
using Parametron
using Parametron.Functions

@testset "canonicalize, prune_zero" begin
    x = Variable(1)
    y = Variable(2)

    @test canonicalize(3 * x * y) === canonicalize(3 * y * x) == QuadraticTerm(3, x, y)
    @test canonicalize(y + x - 2 * y + 3) == x - y + 3
    @test canonicalize(x * y + y * x + y + y + x - y + 4) == 2 * x * y + x + y + 4
    @test canonicalize(0 * y + x + 1) == x + 0 * y + 1
    @test prune_zero(canonicalize(0 * y + x + 1)) == x + 1
    @test canonicalize(0 * y^2 + x^2 + 0 * x + y + 1) == x^2 + 0 * y^2 + 0 * x + y + 1
    @test prune_zero(canonicalize(0 * y^2 + x^2 + 0 * x + y + 1)) == x^2 + y + 1

    f = x * y + y * x + y + y + x - y + 4
    global allocs
    for i = 1 : 2
        fcopy = QuadraticFunction(f)
        allocs = @allocated canonicalize!(fcopy)
    end
    @test allocs == 0
end

@testset "LinearTerm" begin
    x = Variable(1)
    @test LinearTerm(4.5, x) === 4.5 * x
    @test LinearTerm(4.5, x) === x * 4.5
    @test convert(LinearTerm{Float64}, x) === 1.0 * x
    @test promote(LinearTerm(4.5, x), x) === (4.5 * x, 1.0 * x)
    @test +(2 * x) === 2 * x
    @test -(2 * x) === -2 * x
    @test 2 * (3 * x) === x * 6
    @test (2 * x) * 3 === 6 * x
    buf = IOBuffer()
    show(buf, LinearTerm(4.5, x))
    @test String(take!(buf)) == "4.5 * x1"
end

@testset "QuadraticTerm" begin
    x = Variable(1)
    y = Variable(2)
    @test x * y === QuadraticTerm(1, x, y)
    @test (2 * x) * y === QuadraticTerm(2, x, y)
    @test x * (2 * y) === QuadraticTerm(2, x, y)

    @test promote(QuadraticTerm(4.5, x, y), QuadraticTerm(3, x, y)) === (QuadraticTerm(4.5, x, y), QuadraticTerm(3.0, x, y))
    @test +(2 * x * y) === 2 * x * y
    @test -(2 * x * y) === -2 * x * y
    @test 2 * (3 * x * y) === x * 6 * y
    @test (y * 2 * x) * 3 === 6 * y * x
    buf = IOBuffer()
    show(buf, QuadraticTerm(1, x, y))
    @test String(take!(buf)) == "1 * x1 * x2"
end

@testset "AffineFunction" begin
    x = Variable(1)
    y = Variable(2)

    vals1 = Dict(x => 1.0, y => 2.0)
    f1 = 2 * x + 3 * y + 5
    @test f1(vals1) == 2.0 + 3.0 * 2.0 + 5
    buf = IOBuffer()
    show(buf, f1)
    @test String(take!(buf)) == "2 * x1 + 3 * x2 + 5"

    vals2 = Dict(x => 2, y => -1)
    f2 = 2.5 * x + 4 * y + 1
    @test f2(vals2) == 2.5 * 2 - 4 + 1

    f3 = zero(typeof(f2))
    @test typeof(f3) == typeof(f2)
    @test f3(vals2) === 0.0

    f4 = f1 + f2
    @test f4(vals1) == f1(vals1) + f2(vals1)
    @test f4(vals2) == f1(vals2) + f2(vals2)

    f5 = f1 + 4.0
    @test f5(vals1) == f1(vals1) + 4.0

    f6 = 1 + f5
    @test f6(vals1) == 1 + f5(vals1)

    f7 = f6 - f4
    @test f7(vals1) == f6(vals1) - f4(vals1)

    f8 = f7 - LinearTerm(4, y)
    @test f8(vals1) == f7(vals1) - 4 * vals1[y]

    @test AffineFunction{Int}(x) == x + 0
end

@testset "QuadraticFunction" begin
    x = Variable.(1 : 2)
    vals = Dict(zip(x, [1.0, 2.0]))
    intvals = Dict(zip(x, [1, 2]))
    xvals = map(var -> vals[var], x)
    y = zero(QuadraticFunction{Int})

    @test (x ⋅ x)(vals) == xvals ⋅ xvals
    @test Functions.vecdot!(y, x, x)(vals) == xvals ⋅ xvals

    a = [4.0, 5.0]
    @test (x ⋅ (a .* x))(vals) == xvals ⋅ (a .* xvals)
    @test Functions.vecdot!(y, x , (a .* x))(vals) == xvals ⋅ (a .* xvals)
    @test ((a .* x) ⋅ x)(vals) == (a .* xvals) ⋅ xvals
    @test Functions.vecdot!(y, (a .* x), x)(vals) == (a .* xvals) ⋅ xvals
    @test ((a .* x) ⋅ (a .* x))(vals) == (a .* xvals) ⋅ (a .* xvals)
    @test Functions.vecdot!(y, (a .* x), (a .* x))(vals) == (a .* xvals) ⋅ (a .* xvals)

    b = [6.0, 7.0]
    @test (x ⋅ (a .* x .+ b))(vals) == xvals ⋅ (a .* xvals .+ b)
    @test Functions.vecdot!(y, x, (a .* x .+ b))(vals) == xvals ⋅ (a .* xvals .+ b)
    @test ((a .* x .+ b) ⋅ x)(vals) == xvals ⋅ (a .* xvals .+ b)
    @test Functions.vecdot!(y, (a .* x .+ b), x)(vals) == xvals ⋅ (a .* xvals .+ b)
    @test ((b .* x .+ a) ⋅ (a .* x .+ b))(vals) == (b .* xvals .+ a) ⋅ (a .* xvals .+ b)
    @test Functions.vecdot!(y, b .* x .+ a, a .* x .+ b)(vals) == (b .* xvals .+ a) ⋅ (a .* xvals .+ b)

    f1 = x[1]^2 + 2 * x[1] * x[2] + 3 * x[2] + 4
    f2 = x[2]^2 - 2 * x[1] + 3 * x[2] - 1
    @test (f1 - f2)(vals) == f1(vals) - f2(vals)

    f3 = (3 * x[1]^2) - x[2]
    @test f3(vals) == 3 * xvals[1]^2 - xvals[2]

    f4 = (3 * x[1] * x[2]) + (0.5 * x[2])
    @test f4(vals) == 3 * xvals[1] * xvals[2] + 0.5 * xvals[2]

    f5 = (0.5 * x[2]) + (3 * x[1] * x[2])
    @test f5(vals) == 3 * xvals[1] * xvals[2] + 0.5 * xvals[2]

    @test QuadraticFunction{Int}(x[1]) == QuadraticFunction(x[1] + 0)

    @test convert(QuadraticFunction{Float64}, 1)(intvals) === 1.0
    @test convert(QuadraticFunction{Float64}, x[1])(intvals) === vals[x[1]]
    @test convert(QuadraticFunction{Float64}, 3 * x[1])(intvals) === 3 * vals[x[1]]
    @test convert(QuadraticFunction{Float64}, x[1] + x[2] + 2)(intvals) === vals[x[1]] + vals[x[2]] + 2
    @test convert(QuadraticFunction{Float64}, x[1]^2 + x[2]^2)(intvals) === vals[x[1]]^2 + vals[x[2]]^2
end

@testset "matvecmul!" begin
    A = ones(Int, 3, 4)
    x = Variable.(1 : size(A, 2))
    y = Vector{AffineFunction{Int}}(undef, size(A, 1))
    Functions.matvecmul!(y, A, x)
    @test y == fill(sum(x), size(A, 1))
    allocs = @allocated Functions.matvecmul!(y, A, x)
    @test allocs == 0

    Bx = rand(4, 4) * x
    y2 = Vector{AffineFunction{Float64}}(undef, size(A, 1))
    Functions.matvecmul!(y2, A, Bx)
    @test y2 == fill(sum(Bx), size(A, 1))
    allocs = @allocated Functions.matvecmul!(y2, A, Bx)
    @test allocs == 0
end

@testset "mul!" begin
    x = map(Variable, 1 : 3)
    aff = [1, 2, 3]' * x + 4
    quad = x[1]^2 + 2 * x[1] * x[3] + 3 * x[2] + 4

    @testset "affine" begin
        dest = zero(aff)
        for i = 1 : 2
            Functions.mul!(dest, aff, 2)
            @test dest == 2 * x[1] + 4 * x[2] + 6 * x[3] + 8

            Functions.mul!(dest, 3, aff)
            @test dest == 3 * x[1] + 6 * x[2] + 9 * x[3] + 12
        end
    end
    @testset "quadratic" begin
        dest = zero(quad)
        for i = 1 : 2
            Functions.mul!(dest, quad, 2)
            @test dest == 2 * x[1]^2 + 4 * x[1] * x[3] + 6 * x[2] + 8

            Functions.mul!(dest, 2, quad)
            @test dest == 2 * x[1]^2 + 4 * x[1] * x[3] + 6 * x[2] + 8

            Functions.mul!(dest, aff, x[1])
            @test dest == [1, 2, 3]' * (x .* x[1]) + 4 * x[1]

            Functions.mul!(dest, x[1], aff)
            @test dest == [1, 2, 3]' * (x .* x[1]) + 4 * x[1]

            Functions.mul!(dest, aff, aff)
            @test dest == ([1, 2, 3]' * x + 4)^2
        end
    end
end

@testset "dot with matrix variables" begin
    inds = LinearIndices((1:2, 1:2))
    x = [Variable(inds[i]) for i in CartesianIndices(inds)]
    w = [0.1 0.2; 0.3 0.4]
    @test x ⋅ w == w ⋅ x == 0.1 * x[1] + 0.3 * x[2] + 0.2 * x[3] + 0.4 * x[4] + 0.0
end

@testset "Matrix operations" begin
    x = Variable.(1 : 2)
    A1 = [1.0 2.0; 3.0 4.0]
    fs = A1 * x
    @test fs isa Vector{AffineFunction{Float64}}
    vals = Dict(zip(x, [2.0, 5.0]))
    fvals = [f(vals) for f in fs]
    xvals = map(var -> vals[var], x)
    @test fvals == A1 * xvals

    a = [1, 2]
    @test a ⋅ x == x ⋅ a == a[1] * x[1] + a[2] * x[2]

    gs = fs .+ a
    gvals = [g(vals) for g in gs]

    @test (gs ⋅ gs)(vals) == gvals ⋅ gvals
    @test (gs ⋅ x)(vals) == gvals ⋅ xvals
    @test (x ⋅ gs)(vals) == xvals ⋅ gvals

    h = x ⋅ x
    @test h(vals) == xvals ⋅ xvals
    buf = IOBuffer()
    show(buf, h)
    @test String(take!(buf)) == "1 * x1 * x1 + 1 * x2 * x2 + 0"
end

@testset "concatenation" begin
    x = Variable(1)
    y = Variable(2)
    f1 = 3x + 10
    f2 = 0.1 * y - 0.5
    f3 = x + y
    v1 = @SVector [f1, f2]
    v2 = @SVector [f2, f3]
    v3 = [f3, f2, f1]
    @test Functions.vcat!(deepcopy(v1), v1) == v1
    @test Functions.vcat!(deepcopy(vcat(v1, v2)), v1, v2) == vcat(v1, v2)
    @test Functions.vcat!(deepcopy(vcat(v1, v2, v3)), v1, v2, v3) == vcat(v1, v2, v3)
    @test_throws DimensionMismatch Functions.vcat!(zeros(AffineFunction{Float64}, 1), v1)
    @test_throws DimensionMismatch Functions.vcat!(zeros(AffineFunction{Float64}, 3), v1)
    @test_throws DimensionMismatch Functions.vcat!(zeros(AffineFunction{Float64}, 3), v1, v2)
    @test_throws DimensionMismatch Functions.vcat!(zeros(AffineFunction{Float64}, 5), v1, v2)
    @test_throws DimensionMismatch Functions.vcat!(zeros(AffineFunction{Float64}, 6), v1, v2, v3)
    @test_throws DimensionMismatch Functions.vcat!(zeros(AffineFunction{Float64}, 8), v1, v2, v3)
end

@testset "scalar times vector" begin
    x = Variable(1)
    y = Variable(2)
    @test 1 * SVector(x, y) === SVector(x, y) * 1 === SVector(1 * x, 1 * y)
    @test x * SVector(x, y) === SVector(x * x, x * y)
    @test SVector(x, y) * y === SVector(x * y, y * y)
    @test x^2 * SVector(1, 2) === SVector(1, 2) * x^2 === SVector(x^2, 2 * x^2)
end

@testset "zero" begin
    x = Variable(1)
    @test zero(x) === LinearTerm(0, x)
    @test zero(3.0 * x) === LinearTerm(0.0, x)
    @test zero(4 * x * x) === QuadraticTerm(0, x, x)
end

@testset "SVector dot" begin
    x = Variable(1)
    y = Variable(2)
    @test canonicalize(dot(SVector(x, y), SVector(x, y))) == x^2 + y^2
end

end # module


text:
export PerformanceOneHotEncoding
import ..encode_event, ..decode_event

using ..NoteSequences: MIN_MIDI_PITCH, MAX_MIDI_PITCH

"""
    PerformanceOneHotEncoding <: Any

One-Hot encoder for [`PerformanceEvent`](@ref)s.

## Fields
* `event_ranges::Vector{NTuple{3, Int}}`: Stores the min and max values of each event type.
* `num_classes::Int`: Total number of event classes. It is the sum of all `NOTE_ON`,
`NOTE_OFF`,`TIME_SHIFT` and `VELOCITY` events.
"""
struct PerformanceOneHotEncoding
    event_ranges::Vector{NTuple{3, Int}}
    num_classes::Int

    function PerformanceOneHotEncoding(;num_velocitybins::Int=0,
                                        max_shift_steps::Int=DEFAULT_MAX_SHIFT_STEPS,
                                        minpitch::Int=MIN_MIDI_PITCH,
                                        maxpitch::Int=MAX_MIDI_PITCH)
        event_ranges = [
            (NOTE_ON, minpitch, maxpitch),
            (NOTE_OFF, minpitch, maxpitch),
            (TIME_SHIFT, 1, max_shift_steps)
        ]
        num_velocitybins > 0 && push!(event_ranges, (VELOCITY, 1, num_velocitybins))
        num_classes = sum([max - min + 1 for (type, min, max) in event_ranges])
        new(event_ranges, num_classes)
    end
end

function Base.getproperty(encoder::PerformanceOneHotEncoding, sym::Symbol)
    if sym === :labels
        return 1:encoder.num_classes
    elseif sym === :defaultevent
        return PerformanceEvent(TIME_SHIFT, DEFAULT_MAX_SHIFT_STEPS)
    else
        getfield(encoder, sym)
    end
end

"""
    encode_event(event::PerformanceEvent, encoder::PerformanceOneHotEncoding)

Encodes a `PerformanceEvent` to its corresponding one hot index.
"""
function encode_event(event::PerformanceEvent, encoder::PerformanceOneHotEncoding)
    # Start at one to account for 1-based indexing
    offset = 1
    for (type, min, max) in encoder.event_ranges
        if event.event_type == type
            return offset + event.event_value - min
        end
        offset += (max - min + 1)
    end

    throw(error("Unknown type for PerformanceEvent"))
end

"""
    decode_event(index::Int, encoder::PerformanceOneHotEncoding)

Decodes a one hot index to its corresponding `PerformanceEvent`.
"""
function decode_event(index::Int, encoder::PerformanceOneHotEncoding)
    # Start at one to account for 1-based indexing
    offset = 1
    for (type, min, max) in encoder.event_ranges
        if index < offset + (max - min + 1)
            return PerformanceEvent(type, min + index - offset)
        end
        offset += (max - min + 1)
    end

    throw(error("Unknown index"))
end


text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/Materials.jl/blob/master/LICENSE
#
# Low-level definitions for one_elem_disp_chaboche.jl.

mutable struct Continuum3D <: FieldProblem
    material_model :: Symbol
end

Continuum3D() = Continuum3D(:PerfectPlastic)
FEMBase.get_unknown_field_name(::Continuum3D) = "displacement"

function FEMBase.assemble_elements!(problem::Problem{Continuum3D},
                                    assembly::Assembly,
                                    elements::Vector{Element{Hex8}},
                                    time::Float64)

    for element in elements
        for ip in get_integration_points(element)
            material = ip("material", time)
            preprocess_increment!(material, element, ip, time)
        end
    end
    bi = BasisInfo(Hex8)

    dim = 3
    nnodes = 8
    ndofs = dim*nnodes

    BL = zeros(6, ndofs)
    Km = zeros(ndofs, ndofs)
    f_int = zeros(ndofs)
    f_ext = zeros(ndofs)
    D = zeros(6, 6)
    S = zeros(6)

    dtime = 0.05
    # super dirty hack
    # data = first(elements).fields["displacement"].data
    # if length(data) > 1
    #     time0 = data[end-1].first
    #     dtime = time - time0
    # end

    for element in elements

        u = element("displacement", time)
        fill!(Km, 0.0)
        fill!(f_int, 0.0)
        fill!(f_ext, 0.0)

        for ip in get_integration_points(element)
            J, detJ, N, dN = element_info!(bi, element, ip, time)
            material = ip("material", time)
            w = ip.weight*detJ

            # Kinematic matrix, linear part

            fill!(BL, 0.0)
            for i=1:nnodes
                BL[1, 3*(i-1)+1] = dN[1,i]
                BL[2, 3*(i-1)+2] = dN[2,i]
                BL[3, 3*(i-1)+3] = dN[3,i]
                BL[4, 3*(i-1)+1] = dN[2,i]
                BL[4, 3*(i-1)+2] = dN[1,i]
                BL[5, 3*(i-1)+2] = dN[3,i]
                BL[5, 3*(i-1)+3] = dN[2,i]
                BL[6, 3*(i-1)+1] = dN[3,i]
                BL[6, 3*(i-1)+3] = dN[1,i]
            end

            # Calculate stress response
            integrate_material!(material)
            D = material.jacobian
            S = material.stress + material.dstress
            #@info("material matrix", D)

            # Material stiffness matrix
            Km += w*BL'*D*BL

            # Internal force vector
            f_int += w*BL'*S

            # External force vector
            for i=1:dim
                haskey(element, "displacement load $i") || continue
                b = element("displacement load $i", ip, time)
                f_ext[i:dim:end] += w*B*vec(N)
            end

        end

        # add contributions to K, Kg, f
        gdofs = get_gdofs(problem, element)
        add!(assembly.K, gdofs, gdofs, Km)
        add!(assembly.f, gdofs, f_ext - f_int)

    end

    return nothing
end

function FEMBase.assemble_elements!(problem::Problem{Continuum3D},
                                    assembly::Assembly,
                                    elements::Vector{Element{Quad4}},
                                    time::Float64)

    nnodes = 4
    ndofs = 3
    f = zeros(nnodes*ndofs)
    bi = BasisInfo(Quad4)

    for element in elements

        fill!(f, 0.0)

        for ip in get_integration_points(element)

            J, detJ, N, dN = element_info!(bi, element, ip, time)
            w = ip.weight*detJ

            if haskey(element, "surface pressure")
                J = element(ip, time, Val{:Jacobian})'
                n = cross(J[:,1], J[:,2])
                n /= norm(n)
                # sign convention, positive pressure is towards surface
                p = element("surface pressure", ip, time)
                f += w*p*vec(n*N)
            end
        end

        gdofs = get_gdofs(problem, element)
        add!(assembly.f, gdofs, f)

    end

    return nothing

end


text:
export Multiplication

abstract type Multiplication{D,S,T} <:Operator{T} end

struct ConcreteMultiplication{D<:Space,S<:Space,T} <: Multiplication{D,S,T}
    f::VFun{D,T}
    space::S

    ConcreteMultiplication{D,S,T}(f::Fun{D,T},sp::S) where {D,S,T} = new{D,S,T}(f,sp)
end

function ConcreteMultiplication(::Type{V},f::Fun{D,T},sp::Space) where {V,D,T}
    if !domainscompatible(space(f),sp)
        error("Domain mismatch: cannot multiply function on $(domain(f)) to function on $(domain(sp))")
    end
    ConcreteMultiplication{D,typeof(sp),V}(
        convert(Fun{D,V},chop(f,40*eps(cfstype(f)))),sp)
end

function ConcreteMultiplication(f::Fun{D,T},sp::Space) where {D,T}
    if !domainscompatible(space(f),sp)
        error("Domain mismatch: cannot multiply function on $(domain(f)) to function on $(domain(sp))")
    end
    V = promote_type(T,rangetype(sp))
    ConcreteMultiplication{D,typeof(sp),V}(convert(Fun{D,V},chop(f,40*eps(cfstype(f)))),sp)
end

# We do this in two stages to support Modifier spaces
# without ambiguity errors
function defaultMultiplication(f::Fun,sp::Space)
    csp=space(f)
    if csp==sp || !hasconversion(sp,csp)
        error("Implement Multiplication(::Fun{$(typeof(space(f)))},::$(typeof(sp)))")
    end
    MultiplicationWrapper(f,Multiplication(f,csp)*Conversion(sp,csp))
end

Multiplication(f::Fun,sp::Space) = defaultMultiplication(f,sp)

Multiplication(f::Fun,sp::UnsetSpace) = ConcreteMultiplication(f,sp)
Multiplication(f::Fun) = Multiplication(f,UnsetSpace())

Multiplication(c::Number,sp::Space) = Multiplication(Fun(c),sp)
Multiplication(sp::Space,c::Number) = Multiplication(sp,Fun(c))
Multiplication(c::Number) = Multiplication(Fun(c) )

# This covers right multiplication unless otherwise specified.
Multiplication(S::Space,f::Fun) = Multiplication(f,S)

function convert(::Type{Operator{T}},C::ConcreteMultiplication{S,V}) where {S,V,T}
    if T==eltype(C)
        C
    else
        ConcreteMultiplication{S,V,T}(Fun{S,T}(C.f),C.space)
    end
end

domainspace(M::ConcreteMultiplication{D,S,T}) where {D,S,T} = M.space
domain(T::ConcreteMultiplication) = domain(T.f)

## Default implementation: try converting to space of M.f

# avoid ambiguity
rangespace(D::ConcreteMultiplication{F,UnsetSpace,T}) where {F,T} = UnsetSpace()
getindex(D::ConcreteMultiplication{F,UnsetSpace,T},k::Integer,j::Integer) where {F,T} =
    error("No range space attached to Multiplication")

##multiplication can always be promoted, range space is allowed to change
promotedomainspace(D::Multiplication,sp::UnsetSpace) = D
promotedomainspace(D::Multiplication,sp::Space) = Multiplication(D.f,sp)
promoterangespace(D::ConcreteMultiplication{P,UnsetSpace},sp::UnsetSpace) where {P} = D
promoterangespace(D::ConcreteMultiplication{P,UnsetSpace},sp::Space) where {P} =
    promoterangespace(Multiplication(D.f,ConstantSpace(domain(sp))), sp)

choosedomainspace(M::ConcreteMultiplication{D,UnsetSpace},::UnsetSpace) where {D} = space(M.f)
# we assume multiplication maps spaces to themselves
choosedomainspace(M::ConcreteMultiplication{D,UnsetSpace},sp::Space) where {D} = sp

diagm(a::Fun) = Multiplication(a)

struct MultiplicationWrapper{D<:Space,S<:Space,O<:Operator,T} <: Multiplication{D,S,T}
    f::VFun{D,T}
    op::O
end

MultiplicationWrapper(T::Type,f::Fun{D,V},op::Operator) where {D<:Space,V} = MultiplicationWrapper{D,typeof(domainspace(op)),typeof(op),T}(f,op)
MultiplicationWrapper(f::Fun{D,V},op::Operator) where {D<:Space,V} = MultiplicationWrapper(eltype(op),f,op)

@wrapper MultiplicationWrapper

function convert(::Type{Operator{TT}},C::MultiplicationWrapper{S,V,O,T}) where {TT,S,V,O,T}
    if TT==T
        C
    else
        MultiplicationWrapper(Fun{S,TT}(C.f),Operator{TT}(C.op))::Operator{TT}
    end
end

## Multiplication operators allow us to multiply two spaces

hasfasttransform(_) = false
hasfasttransform(f::Fun) = hasfasttransform(space(f))
hasfasttransformtimes(f,g) = pointscompatible(f,g) && hasfasttransform(f) && hasfasttransform(g)

# This should be overriden whenever the multiplication space is different
function default_mult(f::Fun,g::Fun)
    # When the spaces differ we promote and multiply
    if domainscompatible(space(f),space(g))
        m,n = ncoefficients(f),ncoefficients(g)
        # Heuristic division of parameter space between value-space and coefficient-space multiplication.
        if hasfasttransformtimes(f,g) && log10(m)*log10(n)>4
            transformtimes(f,g)
        elseif m≤n
            coefficienttimes(f,g)
        else
            coefficienttimes(g,f)
        end
    else
        sp=union(space(f),space(g))
        Fun(f,sp)*Fun(g,sp)
    end
end

*(f::Fun,g::Fun) = default_mult(f,g)

coefficienttimes(f::Fun,g::Fun) = Multiplication(f,space(g))*g

function transformtimes(f::Fun,g::Fun,n)
    @assert pointscompatible(space(f),space(g))
    isempty(f.coefficients) && return f
    isempty(g.coefficients) && return g
    f2,g2,sp = pad(f,n),pad(g,n),space(f)
    hc = transform(sp,values(f2).*values(g2))
    chop!(Fun(sp,hc),10eps(eltype(hc)))
end
transformtimes(f::Fun,g::Fun) = transformtimes(f,g,ncoefficients(f) + ncoefficients(g) - 1)

*(a::Fun,L::UniformScaling) = Multiplication(a*L.λ,UnsetSpace())
*(L::UniformScaling,a::Fun) = L.λ*a

## docs

"""
`Multiplication(f::Fun,sp::Space)` is the operator representing multiplication by
`f` on functions in the space `sp`.
"""
Multiplication(::Fun,::Space)

"""
`Multiplication(f::Fun)` is the operator representing multiplication by
`f` on an unset space of functions.  Spaces will be inferred when applying or
manipulating the operator.
"""
Multiplication(::Fun)


text:
include("config_syn.jl")

using JuMP, Gurobi, SVDD, OneClassSampling

solver = with_optimizer(Gurobi.Optimizer; OutputFlag=0, Threads=1)
init_strategy = Dict{Symbol, Any}(
    :svdd => SVDD.SimpleCombinedStrategy(SVDD.ModifiedMeanCriterion(), SVDD.BoundedTaxErrorEstimate(0.05, 0.02, 0.98)),
    :pwc => Dict(:gamma => SVDD.modified_mean_criterion, :threshold_strat => OutlierPercentageThresholdStrategy())
)

max_stars_repo_name: englhardt/ocs-evaluation

max_issues_repo_name: englhardt/ocs-evaluation

max_forks_repo_name: englhardt/ocs-evaluation


text:
module Belos
  using CxxWrap, MPI
  import ..libjltrilinos
  import ..Teuchos

  @wrapmodule(libjltrilinos, :register_belos)
end


text:
uv = UncertainValue(rand(1000))

#######################
# TruncateLowerQuantile
#######################
uvc = constrain(uv, TruncateLowerQuantile(0.2))
@test minimum(resample(uvc, 1000)) >= quantile(uv, 0.2) - tol

#######################
# TruncateUpperQuantile
#######################
uvc = constrain(uv, TruncateUpperQuantile(0.8))
@test maximum(resample(uvc, 1000)) <= quantile(uv, 0.8) + tol

#######################
# TruncateQuantiles
#######################
uvc = constrain(uv, TruncateQuantiles(0.2, 0.8))
@test all(quantile(uv, 0.2) - tol .<= resample(uvc, 1000) .<= quantile(uv, 0.8) + tol)

#######################
# TruncateMinimum
#######################
uvc = constrain(uv, TruncateMinimum(0.5))
@test minimum(resample(uvc, 1000)) >= 0.5 - tol

#######################
# TruncateMaximum
#######################
uvc = constrain(uv, TruncateMaximum(1.5))
@test maximum(resample(uvc, 1000)) <= 1.5 + tol

#######################
# TruncateRange
#######################
uvc = constrain(uv, TruncateRange(0.5, 1.5))
@test all(0.5 - tol .<= resample(uvc, 1000) .<= 1.5 + tol)


text:
__precompile__()

module TrajOpt

using mjWrap
using LearningStrategies
import LearningStrategies: setup!, update!, hook, cleanup!
export setup!, update!, hook, cleanup!
using ExpFunctions
using Common

export TrajOptModel
export MPPIStrategy
export ControlExp

using Random
using Statistics

struct TrajOptModel
   theta::Matrix{Float64} # parameters to improve
   mjsys::mjWrap.mjSet    # environment
   s0::Vector{Float64}    # initial / current state
   terminalvalue::Function

   MPC::Bool

   # scratch stuff
   samples::TrajSamples
   traj::TrajSamples

   trace::Dict{Symbol,Vector{Float64}}

   delay::Int
   niter::Int

   function TrajOptModel(mjsys, H, T, K;
                         theta=zeros(mjsys.nu, H),
                         s0=zeros(mjsys.nq + mjsys.nv)*NaN, # set NaN init to signal Rollout functions
                         valuefunction=(x...)->nothing,
                         delay=0,
                         niter=T)
      ns = mjsys.ns

      if H >= T # not MPC mode
         theta=zeros(mjsys.nu, T)
         samples = mjw.allocateTrajSamples(mjsys, T, K, ns)
         traj    = mjw.allocateTrajSamples(mjsys, T, 1, ns)
         mpc     = false
         @warn("MPPI in FULL TRAJECTORY Mode")
      else
         samples = mjw.allocateTrajSamples(mjsys, H, K, ns)
         traj    = mjw.allocateTrajSamples(mjsys, T, 1, ns)
         mpc     = true
         @warn("MPPI in Model Predictive Mode")
      end

      return new(theta, mjsys, s0, valuefunction, mpc,
                 samples, traj,
                 Dict(:stocR => Vector{Float64}(),
                      :meanR => Vector{Float64}(),
                      :evalscore => Vector{Float64}() ),
                 delay,
                 niter)
   end
end

include("mppi.jl")
#include("ilqg.jl")
#include("udp.jl")

##### Control Experiment Functions

using Common: makesavedir, save

mutable struct ControlExp <: LearningStrategy
   prefix::String
   expfile::String
   expname::String
   meshfile::String
   dir::String
   statefile::String
   expresults::String

   function ControlExp(prefix::String, expfile::String,
                       expname::String, meshfile::String)
      return new(prefix, expfile, expname, meshfile, "", "", "")
   end
end

function setup!(s::ControlExp, model)
   @info("Setting up Trajectory Optimization Experiment")

   dir          = "$(s.prefix)/$(s.expname)_opt"
   expdir       = makesavedir(dir, s.expfile, model.mjsys.name, s.meshfile)
   s.dir        = "$(expdir)"
   s.statefile  = "$(expdir)/data.jld2"
   s.expresults = "$(expdir)/expmt.jld2"
end

#function hook(s::ControlExp, model, i)
#   if i==1 || model.stocR[i] >= maximum(model.stocR[1:i-1])
#      @info("\tSaving traj for experiment ", s.dir)
#      save_traj(model.mjsys, model.traj, s.statefile)
#      save(s.expresults, Dict(:stocR => model.stocR[1:i],
#                              :meanR => model.meanR[1:i],
#                              :evals => model.evalscore[1:i]))
#   end
#end

function cleanup!(s::ControlExp, model)
   mjWrap.save_traj(model.mjsys, model.traj, s.statefile) # save result of MPC anyway
   save(s.expresults, model.trace)
   @info("Saved Experiment to ", s.dir)
end

end


text:
using Base.Threads
using Distributed
@everywhere include("Trap.jl")

function getInput()
    println("Enter start, end, n")
    collect([parse(Int64, x) for x in split(readline())])
end

function main(a=0, b=131072, n=32768)
    # a, b, n = getInput()
    h = (b - a) / n

    n_threads = nprocs()
    if n_threads == 1
        println("
        Number of threads detected to be =1\n
        environment variables may not be properly set up!\n
        exiting..."
        )
        return
    end
    workers = Array{Future}(undef, n_threads - 1)
    total = 0
    local_n::Int64 = div(n, n_threads - 1)
    for rank in 2:n_threads
        local_a = a + (rank + 1) * local_n * h
        local_b = local_a + local_n * h
        # println("at $rank: local_a = $local_a, local_b = $local_b")
        workers[rank - 1] = @spawnat rank Trap.trap(local_a, local_b, local_n, h)
    end

    for rank in 2:n_threads
        local_ans = fetch(workers[rank - 1])
        # println("at $rank: ans = $local_ans")
        total += local_ans
    end

    println("the answer for a = $a, b = $b, n = $n: $(total[])")
end

for i=1:10000
    main()
end


text:
# pyplot()
using Plots
default(grid=false) #src
plot(randn(10))
# closeall();gui()
using Parameters, LTVModelsBase, ValueHistories, DSP, ForwardDiff#, JLD, ReverseDiff#, JacProp
const Diff = ForwardDiff
# const RDiff = ReverseDiff
@with_kw struct LinearSys
    A
    B
    N     = 1000
    nx    = size(A,1)
    nu    = size(B,2)
    h     = 0.02
    σ0    = 0
    sind  = 1:nx
    uind  = nx+1:(nx+nu)
    s1ind = (nx+nu+1):(nx+nu+nx)
end

function LinearSys(seed; nx = 10, nu = nx, h=0.02, kwargs...)
    srand(seed)
    A = randn(nx,nx)
    A = A-A'        # skew-symmetric = pure imaginary eigenvalues
    A = A - h*I     # Make 'slightly' stable
    A = expm(h*A)   # discrete time
    B = h*randn(nx,nu)
    LinearSys(;A=A, B=B, nx=nx, nu=nu, h=h, kwargs...)
end

function generate_data(sys::LinearSys, seed, validation=false)
    Parameters.@unpack A,B,N, nx, nu, h, σ0 = sys
    srand(seed)
    u      = DSP.filt(ones(5),[5], 10randn(N+2,nu))'
    t      = h:h:N*h+h
    x0     = randn(nx)
    x      = zeros(nx,N+1)
    x[:,1] = x0
    for i = 1:N-1
        x[:,i+1] = A*x[:,i] + B*u[:,i]
    end

    validation || (x .+= σ0 * randn(size(x)))
    u = u[:,1:N]
    @assert all(isfinite, u)
    x,u
end

function true_jacobian(sys::LinearSys, x, u)
    [sys.A sys.B]
end

num_params = 30
wdecay     = 0
stepsize   = 0.02
const sys  = LinearSys(1, nx=5, N=100, h=0.02, σ0 = 0.01)
true_jacobian(x,u) = true_jacobian(sys,x,u)
nu         = sys.nu
nx         = sys.nx

# Generate validation data
function valdata()
    vx,vu,vy = Vector{Float64}[],Vector{Float64}[],Vector{Float64}[]
    for i = 20:60
        x,u = generate_data(sys,i, true)
        for j in 10:5:(sys.N-1)
            push!(vx, x[:,j])
            push!(vy, x[:,j+1])
            push!(vu, u[:,j])
        end
    end
    hcat(vx...),hcat(vu...),hcat(vy...)
end
vx,vu,vy = valdata()
const vt = Trajectory(vx,vu,vy)

# Generate training trajectories
const trajs = [Trajectory(generate_data(sys, i)...) for i = 1:3]
# TODO: take care of separate trajs
function todata(trajs::Vector{Trajectory})
    xu = hcat(getfield.(trajs,:xu)...)
    y = hcat(getfield.(trajs,:y)...)
    xu, y
end

function i2m(w,i,sizes)
    s = [1; cumsum([prod.(sizes)...]) .+ 1]
    reshape(view(w,s[i]:(s[i+1]-1)), sizes[i])
end

function pred(w,x,sizes)
    state = copy(x)
    for i=1:2:length(sizes)-2
        state = tanh.(i2m(w,i,sizes)*state .+ i2m(w,i+1,sizes))
    end
    return i2m(w,length(sizes)-1,sizes)*state .+ i2m(w,length(sizes),sizes) .+ x[1:nx,:]
end

cost(w,x,y)  = sum(abs2, pred(w,x,sizes) .- y)/size(y,2)
# cost(w,data) = sum(cost(w, d...) for d in data)/length(data)
cost(w,data) = cost(w, data...)

function loss(w,x,y)
    chunk = Diff.Chunk(x[:,1])
    function lf(w)
        # println("Entering loss function, typeof(w):", typeof(w))
        model(x)    = pred(w,x,sizes)
        jcfg        = Diff.JacobianConfig(model, x[:,1], chunk)
        jacobian(x) = Diff.jacobian(model, x, jcfg)
        l = cost(w,x,y)
        J1 = jacobian(x[:,1])
        for t = 2:size(x,2)
            J2 = jacobian(x[:,t])
            l += 2sum(abs2.(J1.-J2))
            J1 = J2
        end
        l
    end
end

##

function runtrain(w, loss; epochs = 500)
    dtrn        = todata(trajs)
    dtst        = todata([vt])
    x,y         = dtrn
    lossfun     = loss(w,x,y)
    gcfg        = RDiff.GradientConfig(w)
    @time tape  = RDiff.compile(RDiff.GradientTape(lossfun, w, gcfg))
    # gradient(w) = Diff.gradient(lossfun, w, gcfg)
    trace       = History(Float64)
    tracev      = History(Float64)
    push!(trace, 0, lossfun(w))
    push!(tracev, 0, cost(w,dtst))
    g = similar(w)
    m = zeros(size(g))
    plot(reuse=false)
    for epoch=1:epochs
        lossfun     = loss(w,x,y)
        # gcfg        = Diff.GradientConfig(lossfun, w)
        # gradient(w) = RDiff.gradient!(g,tape, w)
        # RDiff.gradient!(g, lossfun, w, gcfg)
        γ = 0.85
        RDiff.gradient!(g,tape,w.-γ.*m)
        @. m = 0.002g + γ*m
        @. w -= m
        if epoch % 10 == 0
            push!(trace, epoch, lossfun(w))
            push!(tracev, epoch, cost(w,dtst))
            plot(trace, reuse=true)
            plot!(tracev)
            gui()
            println(last(trace), last(tracev))
        end
    end
    trace,tracev
end

const sizes = ((num_params,nx+nu), (num_params,1), (nx,num_params), (nx,1))
tovec(w) = vcat([vec(w) for w in w]...)

srand(1)
w1 = [ 0.1randn(Float64,sizes[1]), zeros(Float64,sizes[2]),0.1randn(Float64,sizes[3]),  zeros(Float64,sizes[4]) ]
w1 = tovec(w1)
w2 = deepcopy(w1)

# res = runtrain(w1, (w,x,y)-> (w-> cost(w,x,y)), epochs=300)
resj = runtrain(w2, loss, epochs=1000)
# @save "res1.jld" res w1
# @save "resj.jld" resj w2
##

# Eval jac
model(x)    = pred(w2,x,sizes)
jcfg        = Diff.JacobianConfig(model, vt.xu[:,1])
jacobian(x) = Diff.jacobian(model, x, jcfg)
jacs = vcat([eigvals(jacobian(vt.xu[:,i])[1:nx,1:nx]) for i = 1:length(vt)]...)
scatter(real.(jacs), imag.(jacs))
scatter!(real.(eigvals(sys.A)), imag.(eigvals(sys.A)))
phi = linspace(-π/5,π/5,200)
plot!(real.(exp.(phi.*im)), imag.(exp.(phi.*im)), l=(:dash, :black), yaxis=[-0.5,0.5], xaxis=[0,1.2])

@load "resj.jld"
model(x)    = pred(w2,x,sizes)
jfg         = ForwardDiff.JacobianConfig(model, zeros(nx+nu))
jacobian(x) = ForwardDiff.jacobian(model, x)
m = JacProp.ADDiffSystem(w2,sizes,nx,nu, 1, jacobian)
JacProp.eigvalplot(m,vt, true_jacobian)

# Simple test
using BenchmarkTools
using ForwardDiff
const Diff = ForwardDiff
w = randn(20,20)
# x = [randn(20) for i = 1:100]
x = randn(20,100)
f(w::AbstractMatrix,x) = w*x
f(w::AbstractVector,x) = f(reshape(w, 20,20),x)
# const jcfg  = Diff.JacobianConfig(model, x)
function get_loss(w,x)
    modelw(w)    = [f(w,x) for x = x]
    function loss(w)
        modelx(x::AbstractVector)    = f(w,x)
        jacobian(x) = Diff.jacobian(modelx, x)
        # @assert jacobian(x) == w
        # sum(sum(modelw(w))) + sum(sum(jacobian(xi) for xi in x))
        sum(sum(modelw(w))) + sum(sum(jacobian(x[:,i]) for i in 1:size(x,2)))
        # sum(modelw(w)) + sum(jacobian(x))
    end
    loss
end
loss = get_loss(w,x)
grad(w) = Diff.gradient(loss,vec(w))
grad(w)

@btime grad($w)

# f(w,x) = NeuralNetwork(w,x)
# function loss(w,x)
#     ∇xf(x) = jacobian(x->f(w,x), x)
#     w -> norm(∇xf(x))
# end
#
# l = loss(w,x)
# gradient(loss,w) # Good luch with this


text:
mutable struct ReplacementFunction
    db::DB
    replacement_func::Function
    extra_data::Any
    uuid::UUID
end

struct ReplacementFunctionInfo
    handle::duckdb_replacement_scan_info
    main_function::ReplacementFunction
    table_name::String

    function ReplacementFunctionInfo(
        handle::duckdb_replacement_scan_info,
        main_function::ReplacementFunction,
        table_name::String
    )
        result = new(handle, main_function, table_name)
        return result
    end
end

function _replacement_scan_function(handle::duckdb_replacement_scan_info, table_name::Ptr{UInt8}, data::Ptr{Cvoid})
    func::ReplacementFunction = unsafe_pointer_to_objref(data)
    tname = unsafe_string(table_name)
    info = ReplacementFunctionInfo(handle, func, tname)
    func.replacement_func(info)
    return
end

function getdb(info::ReplacementFunctionInfo)
    return info.main_function.db
end

function get_extra_data(info::ReplacementFunctionInfo)
    return info.main_function.extra_data
end

function get_table_name(info::ReplacementFunctionInfo)
    return info.table_name
end

function set_function_name(info::ReplacementFunctionInfo, function_name::String)
    return duckdb_replacement_scan_set_function_name(info.handle, function_name)
end

function add_function_parameter(info::ReplacementFunctionInfo, parameter::Value)
    return duckdb_replacement_scan_add_parameter(info.handle, parameter.handle)
end

function _replacement_func_cleanup(data::Ptr{Cvoid})
    info::ReplacementFunction = unsafe_pointer_to_objref(data)
    delete!(info.db.handle.registered_objects, info.uuid)
    return
end

function add_replacement_scan!(db::DB, replacement_func::Function, extra_data::Any)
    func = ReplacementFunction(db, replacement_func, extra_data, uuid4())
    db.handle.registered_objects[func.uuid] = func
    return duckdb_add_replacement_scan(
        db.handle.handle,
        @cfunction(_replacement_scan_function, Cvoid, (duckdb_replacement_scan_info, Ptr{UInt8}, Ptr{Cvoid})),
        pointer_from_objref(func),
        @cfunction(_replacement_func_cleanup, Cvoid, (Ptr{Cvoid},))
    )
end

max_stars_repo_name: lokax/duckdb

max_issues_repo_name: lokax/duckdb

max_forks_repo_name: lokax/duckdb


text:
using MCMCBenchmarks, Test, Random
Random.seed!(4504)
tests = [
    "Gaussian_Tests",
    "Plot_Tests",
]

res = map(tests) do t
    @eval module $(Symbol("Test_", t))
        include($t * ".jl")
    end
end


text:
#TODO move into Leaf class?
@system Nitrogen begin
    pheno: phenology ~ hold
    PD: planting_density ~ hold

    initial_seed_mass ~ hold
    shoot_mass ~ hold
    green_leaf_area ~ hold

    # SK: get N fraction allocated to leaves, this code is just moved from the end of the procedure, this may be taken out to become a separate fn

    # def setup(self):
    #     # assume nitrogen concentration at the beginning is 3.4% of the total weight of the seed
    #     # need to check with Yang. This doesn't look correct
    #     self.set_pool(self.initial_pool)
    #
    #     #TODO set up interface
    #     self.ratio = 0
    #     self.hourly_soil_uptake = 0
    #     self.hourly_demand = 0
    #     self.cumulative_demand = 0
    #     self.cumulative_soil_uptake = 0

    nitrogen_pool_from_shoot(shoot_mass, PD, frac=0.063) => begin
        if shoot_mass * PD <= 100u"g/m^2"
            # when shoot biomass is lower than 100 g/m2, the maximum [N] allowed is 6.3%
            # shoot biomass and Nitrogen are in g
            # need to adjust demand or else there will be mass balance problems
            #FIXME self.initial_pool or just pool?
            #pool = min(0.063 * shoot_mass, self.initial_pool)
            #pool = min(0.063 * shoot_mass, pool)
            frac * shoot_mass
        else
            #FIXME what about other case? should be no cyclic dependency
            #return pool?
            @error "pool_from_shoot"
        end
    end ~ track(u"g") # Nitrogen

    initial_nitrogen_pool(initial_seed_mass, frac=0.034) => begin
        # assume nitrogen concentration at the beginning is 3.4% of the total weight of the seed
        frac * initial_seed_mass # 0.275
    end ~ track(u"g") # Nitrogen

    #FIXME: how to use `initial_pool` when track has no init?
    nitrogen_pool(nitrogen_pool_from_shoot, nitrogen_uptake_from_soil) => begin
        nitrogen_pool_from_shoot + nitrogen_uptake_from_soil
    end ~ track(u"g") # Nitrogen

    #TODO for 2DSOIL interface
    nitrogen_uptake_from_soil => 0 ~ track(u"g") # Nitrogen

    #TODO currently not implemented in the original code
    # def nitrogen_remobilize(self):
    #     pass
        #droppedLfArea = (1-greenLeafArea/potentialLeafArea)*potentialLeafArea; //calculated dropped leaf area YY
        #SK 8/20/10: Changed it to get all non-green leaf area
        #currentDroppedLfArea = droppedLfArea - previousDroppedlfArea; //leaf dropped at this time step
        #this->set_N((this->get_N()-(leaf_N/leafArea)*currentDroppedLfArea)); //calculated the total amount of nitrogen after the dropped leaves take some nitrogen out
        #no nitrogen remobilization from old leaf to young leaf is considered for now YY

    #TODO rename to `leaf_to_plant_ratio`? or just keep it?
    leaf_nitrogen_fraction(tt=pheno.gdd_after_emergence) => begin
        # Calculate faction of nitrogen in leaves (leaf NFraction) as a function of thermal time from emergence
        # Equation from Lindquist et al. 2007 YY
        #SK 08/20/10: TotalNitrogen doesn't seem to be updated at all anywhere else since initialized from the seed content
        #SK: I see this is set in crop.cpp ln 253 from NUptake from 2dsoil
        # but this appears to be the amount gained from the soil for the time step; so how does it represent totalNitrogen of a plant?
        # tt: record thermal time from emergency YY
        # Calculate faction of nitrogen in leaves (leaf NFraction) as a function of thermal time from emergence
        # Equation from Lindquist et al. 2007 YY
        0.79688 - 0.00023747tt - 0.000000086145tt^2
        # fraction of leaf n in total shoot n can't be smaller than zero. YY
    end ~ track(min=0)

    #TODO rename to `leaves`?
    leaf_nitrogen(leaf_nitrogen_fraction, nitrogen_pool) => begin
        # calculate total nitrogen amount in the leaves YY units are grams N in all the leaves
        leaf_nitrogen_fraction * nitrogen_pool
    end ~ track(u"g") # Nitrogen

    #TODO rename to `unit_leaf`?
    # Calculate leaf nitrogen content of per unit area
    leaf_nitrogen_content(leaf_nitrogen, green_leaf_area) => begin
        # defining leaf nitrogen content this way, we did not consider the difference in leaf nitrogen content
        # of sunlit and shaded leaf yet YY
        #SK 8/22/10: set avg greenleaf N content before update in g/m2
        iszero(green_leaf_area) ? 0. : (leaf_nitrogen / green_leaf_area)
    end ~ track(u"g/cm^2") # Nitrogen
end


text:
# FASTQ Reader
# ============

struct Reader{S <: TranscodingStream} <: BioGenerics.IO.AbstractReader
    state::State{S}
    seq_transform::Union{Function, Nothing}
end

"""
    FASTQ.Reader(input::IO; fill_ambiguous=nothing)

Create a data reader of the FASTQ file format.

# Arguments
* `input`: data source
* `fill_ambiguous=nothing`: fill ambiguous symbols with the given symbol

# Extended help
`Reader`s take ownership of the underlying IO. That means the underlying IO may
not be directly modified such as writing or reading from it, or seeking in it.

`Reader`s carry their own buffer. This means that the underlying IO may be `eof`
before the `Reader` itself. To avoid issues, do not interact with the underlying
IO while a `Reader` is using it.

`Reader`s are iterable and yield `Record` objects. For improved reading speed,
instantiate a single `Record`, then `read!` into the record repeatedly until
`eof(reader)`. Note that this approach overwrite the same mutable `Record`.
"""
function Reader(input::IO; fill_ambiguous = nothing)
    if fill_ambiguous === nothing
        seq_transform = nothing
    else
        seq_transform = generate_fill_ambiguous(fill_ambiguous)
    end
    if !(input isa TranscodingStream)
        stream = TranscodingStreams.NoopStream(input)
        return Reader(State(stream, 1, 1, false), seq_transform)
    else
        return Reader(State(input, 1, 1, false), seq_transform)
    end
end

function Base.eltype(::Type{<:Reader})
    return Record
end

function BioGenerics.IO.stream(reader::Reader)
    return reader.state.stream
end

function Base.read!(rdr::Reader, rec::Record)
    cs, ln, f = readrecord!(rdr.state.stream, rec, (rdr.state.state, rdr.state.linenum), rdr.seq_transform)
    rdr.state.state = cs
    rdr.state.linenum = ln
    rdr.state.filled = f
    if !f
        cs == 0 && throw(EOFError())
        throw(ArgumentError("malformed FASTQ file"))
    end
    return rec
end

function Base.close(reader::Reader)
    if reader.state.stream isa IO
        close(reader.state.stream)
    end
    return nothing
end

function generate_fill_ambiguous(symbol::BioSymbols.DNA)
    certain = map(UInt8, ('A', 'C', 'G', 'T', 'a', 'c', 'g', 't'))
    # return transform function
    return function (data, range)
        fill = convert(UInt8, convert(Char, symbol))
        for i in range
            if data[i] ∉ certain
                data[i] = fill
            end
        end
        return data
    end
end

function index!(record::Record)
    stream = TranscodingStreams.NoopStream(IOBuffer(record.data))
    cs, linenum, found = readrecord!(stream, record, (1, 1), nothing)
    if !found || !allspace(stream)
        throw(ArgumentError("invalid FASTQ record"))
    end
    return record
end

function allspace(stream)
    while !eof(stream)
        if !isspace(read(stream, Char))
            return false
        end
    end
    return true
end

#=

# NOTE: This does not support line-wraps within sequence and quality.
isinteractive() && @info "Compiling FASTQ parser..."
const record_machine, file_machine = (function ()
    cat = Automa.RegExp.cat
    rep = Automa.RegExp.rep
    rep1 = Automa.RegExp.rep1
    alt = Automa.RegExp.alt
    opt = Automa.RegExp.opt
    any = Automa.RegExp.any
    space = Automa.RegExp.space

    hspace = re"[ \t\v]"

    header1 = let
        identifier = rep(any() \ space())
        identifier.actions[:enter] = [:mark]
        identifier.actions[:exit]  = [:header1_identifier]

        description = cat(any() \ hspace, re"[^\r\n]*")
        description.actions[:enter] = [:mark]
        description.actions[:exit]  = [:header1_description]

        cat('@', identifier, opt(cat(rep1(hspace), description)))
    end

    sequence = re"[A-z]*"
    sequence.actions[:enter] = [:mark]
    sequence.actions[:exit]  = [:sequence]

    header2 = let
        identifier = rep1(any() \ space())
        identifier.actions[:enter] = [:mark]
        identifier.actions[:exit]  = [:header2_identifier]

        description = cat(any() \ hspace, re"[^\r\n]*")
        description.actions[:enter] = [:mark]
        description.actions[:exit]  = [:header2_description]

        cat('+', opt(cat(identifier, opt(cat(rep1(hspace), description)))))
    end

    quality = re"[!-~]*"
    quality.actions[:enter] = [:mark]
    quality.actions[:exit]  = [:quality]

    newline = let
        lf = re"\n"
        lf.actions[:enter] = [:countline]

        cat(opt('\r'), lf)
    end

    record′ = cat(header1, newline, sequence, newline, header2, newline, quality)
    record′.actions[:enter] = [:anchor]
    record′.actions[:exit]  = [:record]
    record = cat(record′, newline)

    file = rep(record)

    return map(Automa.compile, (record, file))
end)()

#=
write("fastq.dot", Automa.machine2dot(file_machine))
run(`dot -Tsvg -o fastq.svg fastq.dot`)
=#

function check_identical(data1, range1, data2, range2)
    if length(range1) != length(range2) ||
       memcmp(pointer(data1, first(range1)), pointer(data2, first(range2)), length(range1)) != 0
       error("sequence and quality have non-matching header")
    end
end

function memcmp(p1::Ptr, p2::Ptr, len::Integer)
    return ccall(:memcmp, Cint, (Ptr{Cvoid}, Ptr{Cvoid}, Csize_t), p1, p2, len) % Int
end

const record_actions = Dict(
    :header1_identifier  => :(record.identifier  = (mark:p-1)),
    :header1_description => :(record.description = (mark:p-1)),
    :header2_identifier  => :(check_identical(record.data, mark:p-1, record.data, record.identifier)),
    :header2_description => :(check_identical(record.data, mark:p-1, record.data, record.description)),
    :sequence => :(record.sequence = (mark:p-1)),
    :quality  => :(record.quality  = (mark:p-1)),
    :record   => :(record.filled   = 1:p-1),
    :anchor => :(),
    :mark   => :(mark = p),
    :countline => :())
eval(
    BioCore.ReaderHelper.generate_index_function(
        Record,
        record_machine,
        :(mark = 0),
        record_actions))
eval(
    BioCore.ReaderHelper.generate_read_function(
        Reader,
        file_machine,
        :(mark = offset = 0),
        merge(record_actions, Dict(
            :header1_identifier  => :(record.identifier  = (mark:p-1) .- stream.anchor .+ 1),
            :header1_description => :(record.description = (mark:p-1) .- stream.anchor .+ 1),
            :header2_identifier  => :(check_identical(data, mark:p-1, data, (record.identifier) .+ stream.anchor .- 1)),
            :header2_description => :(check_identical(data, mark:p-1, data, (record.description) .+ stream.anchor .- 1)),
            :sequence            => :(record.sequence    = (mark:p-1) .- stream.anchor .+ 1),
            :quality             => :(record.quality     = (mark:p-1) .- stream.anchor .+ 1),
            :record => quote
                if length(record.sequence) != length(record.quality)
                    error("the length of sequence does not match the length of quality")
                end
                BioCore.ReaderHelper.resize_and_copy!(record.data, data, BioCore.ReaderHelper.upanchor!(stream):p-1)
                record.filled = (offset+1:p-1) .- offset
                if reader.seq_transform != nothing
                    reader.seq_transform(record.data, record.sequence)
                end
                found_record = true
                @escape
            end,
            :countline => :(linenum += 1),
            :anchor => :(BioCore.ReaderHelper.anchor!(stream, p); offset = p - 1)))))

=#

max_forks_repo_name: BioJulia/FASTX


text:
#!/usr/bin/env julia

using Luxor

using Test

using Random
Random.seed!(42)

function run_line_intersection_test(fname)
    Drawing(2000, 2000, fname)
    origin()
    sethue("magenta")
    setline(.2)
    fontsize(4)
    tiles = Tiler(1000, 1000, 10, 10)
    for (pos, n) in tiles
        gsave()
        randomhue()
        translate(pos)
        topleft = Point(-tiles.tilewidth/2, -tiles.tileheight/2)
        bottomright = Point(tiles.tilewidth/2, tiles.tileheight/2)
        a = randompoint(topleft, bottomright)
        b = randompoint(topleft, bottomright)
        c = randompoint(topleft, bottomright)
        d = randompoint(topleft, bottomright)
        line(a, b, :stroke)
        line(c, d, :stroke)
        (flag, ip) = intersectionlines(a, b, c, d, crossingonly=true)
        text("the lines $(flag ? "do" : "don't") cross", O)
        if flag
            gsave()
            setline(.5)
            setdash("dot")
            if distance(a, ip) < distance(b, ip)
                arrow(a, ip, arrowheadlength=1)
            else
                arrow(b, ip, arrowheadlength=1)
            end
            if distance(c, ip) < distance(d, ip)
                arrow(c, ip, arrowheadlength=1)
            else
                arrow(d, ip, arrowheadlength=1)
            end
            circle(ip, 2, :fill)
            grestore()
        else
            if ip != O
                box(ip, 2, 2, :fill)
            end
            dist = distance(O, ip)
            if dist > 500
                text("intersection point is $(dist) units away", O + (0, 10))
            end
        end
        grestore()
    end
    @test finish() == true

end

fname = "line-intersection-options.pdf"
run_line_intersection_test(fname)
println("...finished test: output in $(fname)")


text:
using Oceananigans.Operators: Δy_uᶠᶜᵃ, Δx_vᶜᶠᵃ, Δx_uᶠᶜᵃ, Δy_vᶜᶠᵃ, ℑxyᶠᶠᵃ, ℑxzᶠᵃᶠ, ℑyzᵃᶠᶠ

#####
##### Viscous flux divergences
#####

@inline function ∂ⱼ_τ₁ⱼ(i, j, k, grid, closure::AbstractTurbulenceClosure, args...)
    disc = time_discretization(closure)
    return 1/Vᶠᶜᶜ(i, j, k, grid) * (δxᶠᵃᵃ(i, j, k, grid, Ax_cᶜᶜᶜ, viscous_flux_ux, disc, closure, args...) +
                                    δyᵃᶜᵃ(i, j, k, grid, Ay_ζᶠᶠᶜ, viscous_flux_uy, disc, closure, args...) +
                                    δzᵃᵃᶜ(i, j, k, grid, Az_ηᶠᶜᵃ, viscous_flux_uz, disc, closure, args...))
end

@inline function ∂ⱼ_τ₂ⱼ(i, j, k, grid, closure::AbstractTurbulenceClosure, args...)
    disc = time_discretization(closure)
    return 1/Vᶜᶠᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, Ax_ζᶠᶠᶜ, viscous_flux_vx, disc, closure, args...) +
                                    δyᵃᶠᵃ(i, j, k, grid, Ay_cᶜᶜᶜ, viscous_flux_vy, disc, closure, args...) +
                                    δzᵃᵃᶜ(i, j, k, grid, Az_ξᶜᶠᵃ, viscous_flux_vz, disc, closure, args...))
end

@inline function ∂ⱼ_τ₃ⱼ(i, j, k, grid, closure::AbstractTurbulenceClosure, args...)
    disc = time_discretization(closure)
    return 1/Vᶜᶜᶠ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, Ax_ηᶠᶜᶠ, viscous_flux_wx, disc, closure, args...) +
                                    δyᵃᶜᵃ(i, j, k, grid, Ay_ξᶜᶠᶠ, viscous_flux_wy, disc, closure, args...) +
                                    δzᵃᵃᶠ(i, j, k, grid, Az_cᶜᶜᵃ, viscous_flux_wz, disc, closure, args...))
end

#####
##### Viscosities at different locations
#####

@inline νᶜᶜᶜ(i, j, k, grid, clock, ν::Number) = ν
@inline νᶠᶠᶜ(i, j, k, grid, clock, ν::Number) = ν
@inline νᶠᶜᶠ(i, j, k, grid, clock, ν::Number) = ν
@inline νᶜᶠᶠ(i, j, k, grid, clock, ν::Number) = ν

@inline νᶜᶜᶜ(i, j, k, grid, clock, ν::AbstractArray) = @inbounds ν[i, j, k]
@inline νᶠᶠᶜ(i, j, k, grid, clock, ν::AbstractArray) = @inbounds ℑxyᶠᶠᵃ(i, j, k, grid, ν)
@inline νᶠᶜᶠ(i, j, k, grid, clock, ν::AbstractArray) = @inbounds ℑxzᶠᵃᶠ(i, j, k, grid, ν)
@inline νᶜᶠᶠ(i, j, k, grid, clock, ν::AbstractArray) = @inbounds ℑyzᵃᶠᶠ(i, j, k, grid, ν)

@inline νᶜᶜᶜ(i, j, k, grid, clock, ν::F) where F<:Function = ν(xnode(Center(), i, grid), ynode(Center(), j, grid), znode(Center(), k, grid), clock.time)
@inline νᶠᶠᶜ(i, j, k, grid, clock, ν::F) where F<:Function = ν(xnode(Face(),   i, grid), ynode(Face(),   j, grid), znode(Center(), k, grid), clock.time)
@inline νᶠᶜᶠ(i, j, k, grid, clock, ν::F) where F<:Function = ν(xnode(Face(),   i, grid), ynode(Center(), j, grid), znode(Face(),   k, grid), clock.time)
@inline νᶜᶠᶠ(i, j, k, grid, clock, ν::F) where F<:Function = ν(xnode(Center(), i, grid), ynode(Face(),   j, grid), znode(Face(),   k, grid), clock.time)

#####
##### Convenience "vanilla" viscous flux functions
#####

@inline viscous_flux_ux(i, j, k, grid, clock, ν, u) = - νᶜᶜᶜ(i, j, k, grid, clock, ν) * ∂xᶜᶜᵃ(i, j, k, grid, u)
@inline viscous_flux_uy(i, j, k, grid, clock, ν, u) = - νᶠᶠᶜ(i, j, k, grid, clock, ν) * ∂yᶠᶠᵃ(i, j, k, grid, u)
@inline viscous_flux_uz(i, j, k, grid, clock, ν, u) = - νᶠᶜᶠ(i, j, k, grid, clock, ν) * ∂zᵃᵃᶠ(i, j, k, grid, u)

@inline viscous_flux_vx(i, j, k, grid, clock, ν, v) = - νᶠᶠᶜ(i, j, k, grid, clock, ν) * ∂xᶠᶠᵃ(i, j, k, grid, v)
@inline viscous_flux_vy(i, j, k, grid, clock, ν, v) = - νᶜᶜᶜ(i, j, k, grid, clock, ν) * ∂yᶜᶜᵃ(i, j, k, grid, v)
@inline viscous_flux_vz(i, j, k, grid, clock, ν, v) = - νᶜᶠᶠ(i, j, k, grid, clock, ν) * ∂zᵃᵃᶠ(i, j, k, grid, v)

@inline viscous_flux_wx(i, j, k, grid, clock, ν, w) = - νᶠᶜᶠ(i, j, k, grid, clock, ν) * ∂xᶠᶜᵃ(i, j, k, grid, w)
@inline viscous_flux_wy(i, j, k, grid, clock, ν, w) = - νᶜᶠᶠ(i, j, k, grid, clock, ν) * ∂yᶜᶠᵃ(i, j, k, grid, w)
@inline viscous_flux_wz(i, j, k, grid, clock, ν, w) = - νᶜᶜᶜ(i, j, k, grid, clock, ν) * ∂zᵃᵃᶜ(i, j, k, grid, w)

#####
##### Products of viscosity and divergence, vorticity, vertical momentum gradients
#####

@inline ν_δᶜᶜᶜ(i, j, k, grid, clock, ν, u, v) = νᶜᶜᶜ(i, j, k, grid, clock, ν) * div_xyᶜᶜᵃ(i, j, k, grid, u, v)
@inline ν_ζᶠᶠᶜ(i, j, k, grid, clock, ν, u, v) = νᶠᶠᶜ(i, j, k, grid, clock, ν) * ζ₃ᶠᶠᵃ(i, j, k, grid, u, v)

@inline ν_uzᶠᶜᶠ(i, j, k, grid, clock, ν, u) = νᶠᶜᶠ(i, j, k, grid, clock, ν) * ∂zᵃᵃᶠ(i, j, k, grid, u)
@inline ν_vzᶜᶠᶠ(i, j, k, grid, clock, ν, v) = νᶜᶠᶠ(i, j, k, grid, clock, ν) * ∂zᵃᵃᶠ(i, j, k, grid, v)

@inline ν_uzzzᶠᶜᶠ(i, j, k, grid, clock, ν, u) = νᶠᶜᶠ(i, j, k, grid, clock, ν) * ∂³zᵃᵃᶠ(i, j, k, grid, u)
@inline ν_vzzzᶜᶠᶠ(i, j, k, grid, clock, ν, v) = νᶜᶠᶠ(i, j, k, grid, clock, ν) * ∂³zᵃᵃᶠ(i, j, k, grid, v)

# See https://mitgcm.readthedocs.io/en/latest/algorithm/algorithm.html#horizontal-dissipation
@inline function δ★ᶜᶜᶜ(i, j, k, grid, u, v)

    # These closures seem to be needed to help the compiler infer types
    # (either of u and v or of the function arguments)
    @inline Δy_∇²u(i, j, k, grid, u) = Δy_uᶠᶜᵃ(i, j, k, grid, ∇²hᶠᶜᶜ, u)
    @inline Δx_∇²v(i, j, k, grid, v) = Δx_vᶜᶠᵃ(i, j, k, grid, ∇²hᶜᶠᶜ, v)

    return 1 / Azᶜᶜᵃ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, Δy_∇²u, u) +
                                       δyᵃᶜᵃ(i, j, k, Δx_∇²v, v))
end

@inline function ζ★ᶠᶠᶜ(i, j, k, grid, u, v)

    # These closures seem to be needed to help the compiler infer types
    # (either of u and v or of the function arguments)
    @inline Δy_∇²v(i, j, k, grid, v) = Δy_vᶜᶠᵃ(i, j, k, grid, ∇²hᶜᶠᶜ, v)
    @inline Δx_∇²u(i, j, k, grid, u) = Δx_uᶠᶜᵃ(i, j, k, grid, ∇²hᶠᶜᶜ, u)

    return 1 / Azᶠᶠᵃ(i, j, k, grid) * (δxᶠᵃᵃ(i, j, k, Δy_∇²v, v) -
                                       δyᵃᶠᵃ(i, j, k, Δx_∇²u, u))
end

@inline ν_δ★ᶜᶜᶜ(i, j, k, grid, clock, ν, u, v) = νᶜᶜᶜ(i, j, k, grid, clock, ν) * δ★ᶜᶜᶜ(i, j, k, grid, u, v)
@inline ν_ζ★ᶠᶠᶜ(i, j, k, grid, clock, ν, u, v) = νᶠᶠᶜ(i, j, k, grid, clock, ν) * ζ★ᶠᶠᶜ(i, j, k, grid, u, v)