text:
function bottom_drag_quad!(Bu,Bv,KEu,KEv,sqrtKE,sqrtKE_u,sqrtKE_v,u,v,η,H,h,u²,v²,h_u,h_v)
    # quadratic bottom drag Bu,Bv = c_D/h * | u⃗ | * u⃗

    thickness!(h,η,H)
    Ix!(h_u,h)
    Iy!(h_v,h)

    speed!(u²,v²,u,v)
    Ix!(KEu,u²)
    Iy!(KEv,v²)

    # sqrt of KE, which is actually the kinetic energy without the 0.5 factor
    m,n = size(sqrtKE)
    @boundscheck (m+ep,n+2) == size(KEu) || throw(BoundsError())
    @boundscheck (m+2,n) == size(KEv) || throw(BoundsError())

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            sqrtKE[i,j] = sqrt(KEu[i+ep,j+1] + KEv[i+1,j])
        end
    end

    Ix!(sqrtKE_u,sqrtKE)
    Iy!(sqrtKE_v,sqrtKE)

    m,n = size(Bu)
    @boundscheck (m,n) == size(sqrtKE_u) || throw(BoundsError())
    @boundscheck (m,n) == size(h_u) || throw(BoundsError())
    @boundscheck (m+2+ep,n+2) == size(u) || throw(BoundsError())

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            Bu[i,j] = c_D*sqrtKE_u[i,j] * u[i+1+ep,j+1] / h_u[i,j]
        end
    end

    m,n = size(Bv)
    @boundscheck (m,n) == size(sqrtKE_v) || throw(BoundsError())
    @boundscheck (m,n) == size(h_v) || throw(BoundsError())
    @boundscheck (m+2,n+2) == size(v) || throw(BoundsError())

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            Bv[i,j] = c_D*sqrtKE_v[i,j] * v[i+1,j+1] / h_v[i,j]
        end
    end
end

"""Linear bottom drag computed as r_B*(u,v). r_B is negative and contains the
grid spacing Δ as gradient operators are dimensionless."""
function bottom_drag_lin!(Bu,Bv,KEu,KEv,sqrtKE,sqrtKE_u,sqrtKE_v,u,v,η,H,h,u²,v²,h_u,h_v)
    m,n = size(Bu)
    @boundscheck (m+2+ep,n+2) == size(u) || throw(BoundsError())

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            Bu[i,j] = r_B * u[i+1+ep,j+1]
        end
    end

    m,n = size(Bv)
    @boundscheck (m+2,n+2) == size(v) || throw(BoundsError())

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            Bv[i,j] = r_B * v[i+1,j+1]
        end
    end
end

function no_bottom_drag!(Bu,Bv,KEu,KEv,sqrtKE,sqrtKE_u,sqrtKE_v,u,v,η,H,h,u²,v²,h_u,h_v)
    Bu .= zeero
    Bv .= zeero
end

if bottom_friction == "linear"
    bottom_drag! = bottom_drag_lin!
elseif bottom_friction == "quadratic"
    bottom_drag! = bottom_drag_quad!
elseif bottom_friction == "none"
    bottom_drag! = no_bottom_drag!
else
    throw(error("Bottom friction not correctly specified. Only linear, quadratic or none allowed."))
end

max_stars_repo_name: eschnett/juls

max_issues_repo_name: eschnett/juls

max_forks_repo_name: eschnett/juls


text:
using Klimakoffer

Klimakoffer.upscale_albedo(joinpath(@__DIR__,"..","input","albedo"), "albedo128x65.dat",150)

upscaled = Klimakoffer.read_albedo(joinpath(@__DIR__,"..","input","albedo","albedo150x76.dat"), 150, 76)

ref = Klimakoffer.read_albedo(joinpath(@__DIR__,"..","examples","reference_data","albedo","ref_albedo150x76.dat"), 150, 76)

diff = ref-upscaled

res = sum(abs.(diff))

rm(joinpath(@__DIR__,"..","input","albedo","albedo150x76.dat"))

ext: jl


text:
# ---
# title: 1353. Maximum Number of Events That Can Be Attended
# id: problem1353
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Greedy, Sort, Segment Tree
# link: <https://leetcode.com/problems/maximum-number-of-events-that-can-be-attended/description/>
# hidden: true
# ---
#
# Given an array of `events` where `events[i] = [startDayi, endDayi]`. Every
# event `i` starts at `startDayi` and ends at `endDayi`.
#
# You can attend an event `i` at any day `d` where `startTimei <= d <=
# endTimei`. Notice that you can only attend one event at any time `d`.
#
# Return _the maximum number of events  _you can attend.
#
#
#
# **Example 1:**
#
# ![](https://assets.leetcode.com/uploads/2020/02/05/e1.png)
#
#
#
#     Input: events = [[1,2],[2,3],[3,4]]
#     Output: 3
#     Explanation: You can attend all the three events.
#     One way to attend them all is as shown.
#     Attend the first event on day 1.
#     Attend the second event on day 2.
#     Attend the third event on day 3.
#
#
# **Example 2:**
#
#
#
#     Input: events= [[1,2],[2,3],[3,4],[1,2]]
#     Output: 4
#
#
# **Example 3:**
#
#
#
#     Input: events = [[1,4],[4,4],[2,2],[3,4],[1,1]]
#     Output: 4
#
#
# **Example 4:**
#
#
#
#     Input: events = [[1,100000]]
#     Output: 1
#
#
# **Example 5:**
#
#
#
#     Input: events = [[1,1],[1,2],[1,3],[1,4],[1,5],[1,6],[1,7]]
#     Output: 7
#
#
#
#
# **Constraints:**
#
#   * `1 <= events.length <= 105`
#   * `events[i].length == 2`
#   * `1 <= startDayi <= endDayi <= 105`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
"""
Period at which we should reevaluate the source period (seconds)
"""
const SOURCE_PERIOD_REEVAL_PERIOD = 1
const SOURCE_PERIOD_REEVAL_ALPHA = 0.1

struct Frame
    data::Any
    timestamp::Int64
end

struct Packet
    """
    The data contained in the packet
    """
    data::Any
    """
    The associated frame (that was initially written to the graph)
    """
    frame::Frame
    """
    Whether this is a DONE packet or not
    """
    done::Bool

    # FIXME: ugly
    Packet(frame::Frame) = new(nothing, frame, true)
    Packet(data, frame::Frame) = new(data, frame, false)
end

get_frame_timestamp(p::Packet) = p.frame.timestamp
get_frame(p::Packet) = p.frame
get_data(p::Packet) = p.data
is_done_packet(p::Packet) = p.done

struct CalculatorWrapper
    calculator::Any
    """
    All Channels which connect the wrapped calculator with the rest of the graph.

    Invariant: the order of the channels corresponds to the order of the arguments
    of the corresponding `process()`.
    """
    input_channels::Vector{Channel}

    """
    All channels corresponding to the one output stream (one channel per subscriber)
    """
    output_channels::Vector{Channel}

    "Whether the wrapped calculator is a sink calculator"
    is_sink::Bool

    "Whether the resolved process method has a `graph` keyword parameter"
    process_has_graph_kw::Bool

    """
    Statistic about the (exponentially smoothed) execution time of the calculator's
    `process()`.
    """
    exec_time::Threads.Atomic{Float64}

    CalculatorWrapper(calc, input_channels, output_channels, is_sink, has_graph_kw) = new(
        calc,
        input_channels,
        output_channels,
        is_sink,
        has_graph_kw,
        Threads.Atomic{Float64}(-1.0),
    )
end

get_calculator(cw::CalculatorWrapper) = cw.calculator
get_input_channels(cw::CalculatorWrapper)::Vector{Channel} = cw.input_channels
get_output_channels(cw::CalculatorWrapper)::Vector{Channel} = cw.output_channels
is_sink(cw::CalculatorWrapper) = cw.is_sink
process_has_graph_kw(cw::CalculatorWrapper) = cw.process_has_graph_kw
get_exec_time(cw::CalculatorWrapper) = cw.exec_time[]

function add_new_exec_time(cw::CalculatorWrapper, new_time::Float64)
    current_time = get_exec_time(cw)
    if current_time > 0.0
        cw.exec_time[] =
            SOURCE_PERIOD_REEVAL_ALPHA * current_time +
            (1 - SOURCE_PERIOD_REEVAL_ALPHA) * new_time
    else
        cw.exec_time[] = new_time
    end
end

"""
Blocks on all channels until data is available on all channels.

Discards all old packets in the case where all channels don't refer to the same frame.
This could happen if an upstream calculator threw away a packet.

Contract
+ All returned packets refer to the same frame
"""
function fetch_input_streams(cw::CalculatorWrapper)::Vector{Packet}
    packets = map(take!, cw.input_channels)

    # Drop old packets
    newest_timestamp = max(map(get_frame_timestamp, packets)...)
    for (idx, packet) in enumerate(packets)
        while get_frame_timestamp(packet) < newest_timestamp
            packet = take!(cw.input_channels[idx])
        end

        packets[idx] = packet
    end

    return packets
end

"""
Monitors input streams, calls on the calculator to process the data, and forwards
the return value in the output channel.
"""
function run_calculator(graph::Graph, cw::CalculatorWrapper)
    # TODO: Refactor to make more readable

    done = false
    while !done
        in_packets::Vector{Packet} = fetch_input_streams(cw)

        if all(is_done_packet, in_packets)
            # Note: calculators will be `close()`d in the main thread
            done_packet = in_packets[1]
            for out_channel in get_output_channels(cw)
                put!(out_channel, done_packet)
            end

            done = true
        else
            out_value = nothing
            try
                stats = nothing
                if process_has_graph_kw(cw)
                    stats = @timed process(
                        get_calculator(cw),
                        map(get_data, in_packets)...,
                        graph = graph,
                    )
                else
                    stats = @timed process(get_calculator(cw), map(get_data, in_packets)...)
                end
                add_new_exec_time(cw, stats.time)
                out_value = stats.value
            catch e
                println("Error caught: $e :")
                display(stacktrace(catch_backtrace()))
                println()
                stop(graph)
                done = true
            end

            if !done && out_value !== nothing
                frame = get_frame(in_packets[1])
                out_packet = Packet(out_value, frame)

                for out_channel in get_output_channels(cw)
                    put!(out_channel, out_packet)
                end
            end
        end

        if is_sink(cw) && is_last_sink_not_init(graph)
            # Source period bootstrap just ended
            lk = get_flow_limiter_bootstrap_lock(graph)

            lock(lk)
            notify(get_flow_limiter_bootstrap_cond(graph))
            unlock(lk)
        end

    end
end

"""
Evaluates the flow limiter period (i.e. the amount of time to sleep in between fetching new packets from the source calculator)
"""
function evaluate_packet_period(graph)
    non_gen_cws = filter(
        cw -> cw != get_source_calculator_wrapper(graph),
        get_calculator_wrappers(graph),
    )

    exec_times = map(get_exec_time, non_gen_cws)
    if all(time -> time > 0.0, exec_times)
        new_gen_period = max(exec_times...)
        set_source_period(graph, new_gen_period)
    end
end

"""
Generates a packet using the source calculator and writes it to the graph.

Returns true if generate and write were successful, false if not, in which case the graph is stopped.

"""
function generate_packet_and_write(graph::Graph)::Bool
    packet_data = nothing
    try
        source_cw = get_source_calculator_wrapper(graph)
        packet_data = process(get_calculator(source_cw))
    catch e
        println("Error caught in source's process():\n$e")
    end

    if packet_data === nothing
        stop(graph)
        return false
    end

    write(graph, packet_data)
    return true
end

"""
    start(graph)

Start all the calculators and begin pulling data out of the source calculator.

All calculators run in their own [task](https://docs.julialang.org/en/v1/base/parallel/#Tasks). A call to `start(graph)` does not block; it returns after spawning all the tasks.
"""
function start(graph::Graph)
    # TODO: If already started, error

    # Start all non-source calculators
    for calculator_wrapper in get_calculator_wrappers(graph)
        if calculator_wrapper != get_source_calculator_wrapper(graph)
            t = @spawn run_calculator(graph, calculator_wrapper)
            push!(get_cw_tasks(graph), t)
        end
    end

    # Generate the first packet
    success = generate_packet_and_write(graph)
    if !success
        return
    end

    # Start bootstrap
    @spawn begin
        lk = get_flow_limiter_bootstrap_lock(graph)
        cond = get_flow_limiter_bootstrap_cond(graph)

        lock(lk)
        try
            while get_num_sinks_not_init(graph) > 0
                wait(cond)
            end
        finally
            unlock(lk)
        end

        evaluate_packet_period(graph)

        # Start source function
        @spawn begin
            while !is_done(graph)
                if generate_packet_and_write(graph)
                    sleep(get_source_period(graph))
                else
                    break
                end
            end

            write_done(graph)
        end

        # Start flow limiter period evaluation
        @spawn begin
            sleep(SOURCE_PERIOD_REEVAL_PERIOD)
            while !is_done(graph)
                evaluate_packet_period(graph)
                sleep(SOURCE_PERIOD_REEVAL_PERIOD)
            end
        end
    end

end

"""
Return and increments timestamp.
"""
function inc_next_timestamp(graph::Graph)
    next_ts = get_next_timestamp(graph)
    set_next_timestamp(graph, next_ts + 1)

    next_ts
end

"""
Writes a value to the graph's input stream.
"""
function write(graph::Graph, data::Any)
    timestamp = inc_next_timestamp(graph)
    packet_frame = Frame(data, timestamp)
    packet = Packet(data, packet_frame)

    for channel in get_output_channels(get_source_calculator_wrapper(graph))
        put!(channel, packet)
    end
end

"""
Writes a DONE packet to the graph's input stream.
"""
function write_done(graph::Graph)
    timestamp = inc_next_timestamp(graph)
    packet_frame = Frame(nothing, timestamp)
    done_packet = Packet(packet_frame)

    for channel in get_output_channels(get_source_calculator_wrapper(graph))
        put!(channel, done_packet)
    end
end


text:
module JointMoments

    using StatsBase

    import Base: std

    export
        normalize,
        coskew,
        cokurt,
        coskewness,
        cokurtosis,
        center,
        collapse,
        std,
        flip,
        _cov,
        outer,
        replicate,
        recombine

    include("center.jl")
    include("tensors.jl")
    include("replicate.jl")
    include("collapse.jl")
    include("statistics.jl")

end # module


text:
""" Array utilities.

- Dense and sparse arrays with a uniform interface
- Lazy operations on arrays
- Wapper structs for arrays to enable "typed" array computing
"""
module ArrayUtils
export @parts_array_struct, @vector_struct, applydiag, applynz, fromnz,
  enumeratenz, lazy

using LazyArrays: ApplyArray
using SparseArrays

# Dense and sparse arrays
#########################

abstract type AbstractArrayBuilder{T,N} end

struct ArrayBuilder{T,N,Arr<:AbstractArray{T,N}} <: AbstractArrayBuilder{T,N}
  array::Arr
end
add!(builder::ArrayBuilder, v, inds...) = builder.array[inds...] += v
take(builder::ArrayBuilder) = builder.array

struct SparseVectorBuilder{Tv,Ti<:Integer} <: AbstractArrayBuilder{Tv,1}
  n::Int
  I::Vector{Ti}
  V::Vector{Tv}
end
function add!(builder::SparseVectorBuilder, v, i::Integer)
  push!(builder.I, i); push!(builder.V, v)
end
function add!(builder::SparseVectorBuilder, v, i::AbstractVector)
  append!(builder.I, i); append!(builder.V, v)
end
take(builder::SparseVectorBuilder) = sparsevec(builder.I, builder.V, builder.n)

struct SparseMatrixBuilder{Tv,Ti<:Integer} <: AbstractArrayBuilder{Tv,2}
  m::Int
  n::Int
  I::Vector{Ti}
  J::Vector{Ti}
  V::Vector{Tv}
end
function add!(builder::SparseMatrixBuilder, v, i::Integer, j::Integer)
  push!(builder.I, i); push!(builder.J, j); push!(builder.V, v)
end
function add!(builder::SparseMatrixBuilder, v, i::AbstractVector, j::AbstractVector)
  append!(builder.I, i); append!(builder.J, j); append!(builder.V, v)
end
take(builder::SparseMatrixBuilder) = sparse(builder.I, builder.J, builder.V,
                                            builder.m, builder.n)

""" Object for efficient incremental construction of dense or sparse vector.
"""
nzbuilder(::Type{Arr}, dims::Integer...) where Arr =
  ArrayBuilder(zeros(Arr, dims))
nzbuilder(::Type{<:SparseVector{Tv}}, n::Integer) where Tv =
  SparseVectorBuilder(n, Int[], Tv[])
nzbuilder(::Type{<:SparseMatrixCSC{Tv}}, m::Integer, n::Integer) where {Tv,Ti} =
  SparseMatrixBuilder(m, n, Int[], Int[], Tv[])

""" Apply diagonal operator to dense or sparse vector.
"""
applydiag(f, x::AbstractVector) = map(f, eachindex(x), x)

function applydiag(f, x::SparseVector)
  I, V = findnz(x)
  sparsevec(I, map(f, I, V), length(x))
end

""" Apply linear map defined in terms of structural nonzero values.
"""
function applynz(f, x::Vec, m::Integer, n::Integer) where Vec <: AbstractVector
  length(x) == n ||
    error("Vector length does not match domain: $(length(x)) != $n")
  y = nzbuilder(Vec, m)
  for (i, a) in enumeratenz(x)
    if !iszero(a)
      j, b = f(i)
      add!(y, a*b, j)
    end
  end
  take(y)
end

""" Construct array, not necessarily sparse, from structural nonzero values.
"""
function fromnz(f, ::Type{Mat}, m::Integer, n::Integer) where Mat <: AbstractMatrix
  A = nzbuilder(Mat, m, n)
  for j in 1:n
    for (i, a) in zip(f(j)...)
      add!(A, a, i, j)
    end
  end
  take(A)
end

function fromnz(::Type{Vec}, I::AbstractVector, V::AbstractVector,
                n::Integer) where Vec <: AbstractVector
  vec = zeros(Vec, n)
  vec[I] = V
  vec
end
fromnz(::Type{<:SparseVector}, I::AbstractVector,
       V::AbstractVector, n::Integer) = sparsevec(I, V, n)

""" Enumerate structural nonzeros of vector.
"""
enumeratenz(v::AbstractVector) = enumerate(v)
enumeratenz(v::SparseVector) = zip(findnz(v)...)

""" Alternative to `Base.zeros` supporting dense and sparse arrays.
"""
zeros(::Type{Arr}, dims::Integer...) where Arr = zeros(Arr, dims)
zeros(::Type{Array{T,N}}, dims::Tuple{Vararg{Integer,N}}) where {T,N} =
  Base.zeros(T, dims)
zeros(::Type{<:SparseVector{T}}, dims::Tuple{Vararg{Integer,1}}) where T =
  spzeros(T, dims...)
zeros(::Type{<:SparseMatrixCSC{T}}, dims::Tuple{Vararg{Integer,2}}) where T =
  spzeros(T, dims...)

# Lazy arrays
#############

""" Perform operation lazily on arrays.
"""
lazy(::typeof(vcat), args...) = ApplyArray(vcat, args...)

# Wrapped arrays
################

""" Generate struct for a named vector struct.
"""
macro vector_struct(struct_sig)
  name, params = parse_struct_signature(struct_sig)
  :T ∉ params || error("Type parameter `T` is reserved")
  expr = quote
    Core.@__doc__ struct $(name){$(params...),T,V<:AbstractVector{T}} <: AbstractVector{T}
      data::V
    end
    $(if !isempty(params); quote
        $(name){$(params...)}(v::V) where {$(params...),T,V<:AbstractVector{T}} =
          $(name){$(params...),T,V}(v)
      end end)
    Base.size(v::$name) = size(v.data)
    Base.getindex(v::$name, i::Int) = getindex(v.data, i)
    Base.setindex!(v::$name, x, i::Int) = setindex!(v.data, x, i)
    Base.IndexStyle(::Type{<:$name}) = IndexLinear()
  end
  esc(expr)
end

""" Generate struct for C-set part or parts of a certain type.

The struct wraps an array and is mainly useful for dispatching on part type.
Example usage:

```julia
@parts_struct V
@parts_struct E

E(1)     # Edge #1 (zero-dimensional array)
V([1,4]) # Vertices #2 and #4 (one-dimensional array)
```
"""
macro parts_array_struct(struct_sig)
  name, params = parse_struct_signature(struct_sig)
  :N ∉ params || error("Type parameter `N` is reserved")
  expr = quote
    Core.@__doc__ struct $(name){$(params...),N,Data} <: AbstractArray{Int,N}
      data::Data
      $(name){$(params...)}(x::Int) where {$(params...)} =
        new{$(params...),0,Int}(x)
      $(name){$(params...)}(xs::Arr) where {$(params...),N,Arr<:AbstractArray{Int,N}} =
        new{$(params...),N,Arr}(xs)
    end
    Base.size(A::$name) = size(A.data)
    Base.getindex(A::$name, args...) = getindex(A.data, args...)
  end
  esc(expr)
end

function parse_struct_signature(expr)
  if expr isa Expr
    expr.head == :curly || error("Invalid struct expression: $expr")
    (expr.args[1]::Symbol, collect(Symbol, expr.args[2:end]))
  else
    (expr::Symbol, Symbol[])
  end
end

end


text:
module FrontendBenchmarks

using BenchmarkTools

const SUITE = BenchmarkGroup(["nestedscopes"])

nested_lets = :(:(()->$(Expr(:let, Expr(:block, Any[ Expr(:(=), Symbol("x$i"), :x1) for i = 1:2000 ]...), Expr(:block)))))

SUITE["nestedscopes"] = @benchmarkable Meta.lower(@__MODULE__, $nested_lets)

end


text:
# Autogenerated wrapper script for ICU_jll for x86_64-apple-darwin
export libicudata, libicui18n, libicuio, libicutest, libicutu, libicuuc

JLLWrappers.@generate_wrapper_header("ICU")
JLLWrappers.@declare_library_product(libicudata, "@rpath/libicudata.69.dylib")
JLLWrappers.@declare_library_product(libicui18n, "@rpath/libicui18n.69.dylib")
JLLWrappers.@declare_library_product(libicuio, "@rpath/libicuio.69.dylib")
JLLWrappers.@declare_library_product(libicutest, "@rpath/libicutest.69.dylib")
JLLWrappers.@declare_library_product(libicutu, "@rpath/libicutu.69.dylib")
JLLWrappers.@declare_library_product(libicuuc, "@rpath/libicuuc.69.dylib")
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_library_product(
        libicudata,
        "lib/libicudata.69.1.dylib",
        nothing,
    )

    JLLWrappers.@init_library_product(
        libicui18n,
        "lib/libicui18n.69.1.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libicuio,
        "lib/libicuio.69.1.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libicutest,
        "lib/libicutest.69.1.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libicutu,
        "lib/libicutu.69.1.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libicuuc,
        "lib/libicuuc.69.1.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
try
    using ThreadedScansTests
    true
catch
    false
end || begin
    let path = joinpath(@__DIR__, "ThreadedScansTests")
        path in LOAD_PATH || push!(LOAD_PATH, path)
    end
    using ThreadedScansTests
end


text:
export denseblocks

type DenseBlockIterator
    readers::Array{Any}
    blockSize::Int64
    blockWidth::Int64
    block::Array{Float64,2}
    offset::Int64
    done::Bool
    constantColumn::Bool
    loop::Bool
end
function denseblocks(readers, blockSize::Int64; constantColumn=false, loop=false)
    blockWidth = constantColumn ? length(readers) + 1 : length(readers)
    block = ones(Float64, blockSize, blockWidth)
    if constantColumn
        block[:,end] = 1.0
    end
    DenseBlockIterator(readers, blockSize, blockWidth, block, 0, false, constantColumn, loop)
end
Base.start(it::DenseBlockIterator) = 0
Base.done(it::DenseBlockIterator, nil) = it.done && !it.loop # never done with a constant column
function Base.next(it::DenseBlockIterator, nil)

    if it.constantColumn
        it.block[:,1:end-1] = 0.0
    else
        it.block[:,:] = 0.0
    end

    # Fill in the block
    if !it.done

        foundRead = false
        for i in 1:length(it.readers)
            reader = it.readers[i]

            while !eof(reader) && position(reader) <= it.offset + it.blockSize
                it.block[position(reader) - it.offset, i] += value(reader)
                advance!(reader)
                foundRead = true
            end
        end

        # See if we are really done or just found a blank block
        if !foundRead
            it.done = true
            for i in 1:length(it.readers)
                it.done = it.done && eof(it.readers[i])
            end
        end
    end

    # update the offset
    it.offset += it.blockSize

    it.block,nothing
end


text:
@testset "StrictlyLessThan2StrictlyGreaterThan: $(T)" for T in [Int, Float64]
    mock = MOIU.MockOptimizer(StrictlyGreaterThanMILPModel{T}())
    model = COIB.StrictlyLessThan2StrictlyGreaterThan{T}(mock)

    if T == Int
        @test MOI.supports_constraint(model, MOI.VariableIndex, MOI.Integer)
    end
    @test MOI.supports_constraint(
        model,
        MOI.ScalarAffineFunction{T},
        CP.Strictly{MOI.GreaterThan{T}, T},
    )
    @test MOIB.supports_bridging_constraint(
        model,
        MOI.ScalarAffineFunction{T},
        CP.Strictly{MOI.LessThan{T}, T},
    )

    if T == Int
        x, _ = MOI.add_constrained_variable(model, MOI.Integer())
    elseif T == Float64
        x = MOI.add_variable(model)
    end

    fct = one(T) * x
    c = MOI.add_constraint(model, fct, CP.Strictly(MOI.LessThan(one(T))))

    @test MOI.is_valid(model, x)
    @test MOI.is_valid(model, c)

    bridge = first(MOIBC.bridges(model))[2]

    @testset "Bridge properties" begin
        @test MOIBC.concrete_bridge_type(typeof(bridge), MOI.ScalarAffineFunction{T}, CP.Strictly{MOI.LessThan{T}, T}) == typeof(bridge)
        @test MOIB.added_constrained_variable_types(typeof(bridge)) == Tuple{Type}[]
        @test MOIB.added_constraint_types(typeof(bridge)) == [
            (MOI.ScalarAffineFunction{T}, CP.Strictly{MOI.GreaterThan{T}, T}),
        ]

        @test MOI.get(bridge, MOI.NumberOfVariables()) == 0
        @test MOI.get(bridge, MOI.NumberOfConstraints{MOI.ScalarAffineFunction{T}, CP.Strictly{MOI.GreaterThan{T}, T}}()) == 1

        @test MOI.get(bridge, MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{T}, CP.Strictly{MOI.GreaterThan{T}, T}}()) == [bridge.con]
    end

    @testset "Constraint" begin
        @test MOI.is_valid(model, bridge.con)
        f = MOI.get(model, MOI.ConstraintFunction(), bridge.con)
        @test length(f.terms) == 1
        @test MOI.get(model, MOI.ConstraintSet(), bridge.con) == CP.Strictly(MOI.GreaterThan(-one(T)))

        t1 = f.terms[1]
        @test t1.coefficient === -one(T)
        @test t1.variable == x
    end
end


text:
# attr_list = []
# for func in (atomics..., contour)
#     Typ = to_type(func)
#     attr = keys(default_theme(nothing, Typ))
#     push!(attr_list, attr...)
# end
# attr_list = string.(sort!(unique(attr_list)))
# # filter out fxaa attribute
# attr_list = filter!(x -> x ≠ "fxaa", attr_list)

const plot_attr_desc = Dict(
    :absorption => "Float32. Sets the absorption value for `volume` plots.",
    :algorithm => "Algorithm to be used for `volume` plots. Can be one of `:iso`, `:absorption`, `:mip`, `:absorptionrgba`, or `:indexedabsorption`.",
    :align => "`(:pos, :pos)`. Specify the text alignment, where `:pos` can be `:left`, `:center`, or `:right`.",
    :alpha => "Float in [0,1]. The alpha value (transparency).",
    :color => "Symbol or Colorant. The color of the main plot element (markers, lines, etc.). Can be a color symbol/string like :red, or a Colorant.",
    :colormap => "The color map of the main plot. Call available_gradients() to see what gradients are available. Can also be used with any Vector{<: Colorant}, or e.g. [:red, :black]",
    :colorrange => "A tuple `(min, max)` where `min` and `max` specify the data range to be used for indexing the colormap. E.g. color = [-2, 4] with colorrange = (-2, 4) will map to the lowest and highest color value of the colormap.",
    :fillrange => "Bool. Toggles range filling in `contour` plots.",
    :font => "String. Specifies the font, and can choose any font available on the system.",
    :glowcolor => "Color Type. Color of the marker glow (outside the border) in `scatter` plots.",
    :glowwidth => "Number. Width of the marker glow in `scatter` plots.",
    :image => "The image to be plotted on the plot.",
    :interpolate => "Bool. For `heatmap` and `images`. Toggles color interpolation between nearby pixels.",
    :isorange => "Float32. Sets the isorange for `volume` plots.",
    :isovalue => "Float32. Sets the isovalue for `volume` plots.",
    :levels => "Integer. Number of levels for a `contour`-type plot.",
    :linestyle => "Symbol. Style of the line (for `line` and `linesegments` plots). Available styles are `:dash`, `:dot`, `:dashdot`, and `:dashdotdot`. You can also supply an array describing the length of each gap/fill.",
    :linewidth => "Number. Width of the line in `line` and `linesegments` plots.",
    :marker => "Symbol, Shape, or AbstractVector.",
    :marker_offset => "Array of `GeometryTypes.Point`'s. Specifies the offset coordinates for the markers. See the [Marker offset](@ref) example.",
    :markersize => "Number or AbstractVector. Specifies size (radius pixels) of the markers.",
    :position => "NTuple{2,Float}, `(x, y)`. Specify the coordinates to position text at.",
    :rotation => "Float32. Specifies the rotation in radians.",
    :rotations => "AbstractVector{Float32}. Similar to `:rotation`, except it specifies the rotations for each element in the plot.",
    :shading => "Bool. Specifies if shading should be on or not (for meshes).",
    :strokecolor => "Color Type. Color of the marker stroke (border).",
    :strokewidth => "Number. Width of the marker stroke (in pixels).",
    :textsize => "Integer. Font pointsize for text.",
    :transformation => "`(:plane, location)`. Transforms the `:plane` to the specified location. Possible `:plane`'s are `:xy`, `:yz`, and `:xz`.",
    :visible => "Bool. Toggle visibility of plot."
)


text:
# Plot a 1D simplicial subcomplex (graph) of `cplx`
@recipe function f(cplx::T, args...) where {T<:SimplicialComplex}
    data = length(args) > 0 ? args[1] :
                              hcat(circlepoints(size(cplx, 0)+1, 1.0)...)
    for (i,c) in enumerate(cells(cplx, 1))
        pts = data[values(c),:]
        @series begin
            seriestype := :path
            linewidth --> 2
            linecolor --> :black
            label --> "S$i"
            pts[:,1], pts[:,2]
        end
    end
end


text:
# Define and test Monoid instance for Int
# ---------------------------------------

TypeClasses.combine(a::Int, b::Int) = a + b

@test reduce_monoid([1,2,3,1]) == 7
@test foldl_monoid([1,2,3,4]) == 10
@test foldr_monoid([1,2,3,100], init=3000) == 3106

# Test default Semigroup instance for String
# ------------------------------------------

@test reduce_monoid(["hi"," ","du"], init="!") == "!hi du"

# Test `neutral` singleton value
# ------------------------------

@test combine(neutral, :anything) == :anything
@test combine("anything", neutral) == "anything"
@test combine(neutral, neutral) == neutral


text:
module SparseND

type Banded3D{T} <: AbstractArray{T,3}
    k0::Array{Int,2}
    k1::Array{Int,2}
    data::Array{T,3}
    function Banded3D{T}(k0,k1,data) where T
        @assert size(k0) == size(k1)
        @assert size(k0) == size(data,1,2)
        return new(k0,k1,data)
    end
end

Banded3D{T}(k0,k1,data::Array{T}) = Banded3D{T}(k0,k1,data)

import Base: size, getindex, setindex!, linearindexing

size(A::Banded3D) = size(A.data)
#linearindexing{T<:Banded3D}(::Type{T}) = Base.LinearSlow()
function getindex(A::Banded3D, m::Int, n::Int, k::Int)
    (A.k0[m,n] <= k <= A.k1[m,n]) || return zero(eltype(A.data))
    A.data[m,n,k-A.k0[m,n]+1]
end
function setindex!(A::Banded3D, v, m, n, k)
    @assert A.k0[m,n] <= k <= A.k1[m,n]
    A.data[m,n,k-A.k0[m,n]+1] = v
end

"""
    convolve{T}(A::Banded3D, x::Array{T,2})

Compute the *space-time* convolution of A and x.

Returns an array `y` of size `(size(A,1),size(x,2))` such that
`y[m,i] = sum([A[]])``
"""
function convolve{T}(A::Banded3D, x::Array{T,2})
    V = promote_type(eltype(A),eltype(x))
    y = zeros(V,size(A,1),size(x,2))
    for i in 1:size(A,1)
        for j in 1:size(A,2)
            for k in A.k0[i,j] : A.k1[i,j]
                for l in 1 : size(y,2)
                    l-k+1 < 1 && continue
                    y[i,l] += A[i,j,k] * x[j,l-k+1]
                end
            end
        end
    end
    return y
end

end # module


text:
using JLD2
using FileIO

function test_merge(fname_base="output/simpleDemoCompareCost_results_", files=100)
    if files < 1
        return
    end
    o = load(string(fname_base, 1, ".jld2"))
    AS_cost_ = o["AS_cost"]
    BS_cost_ = o["BS_cost"]
    ts_ = o["ts"]
    Ns_ = o["Ns"]
    maxIter_ = o["maxIter"]
    for i in 2:files
        o = load(string(fname_base, i, ".jld2"))
        AS_cost_ = cat(AS_cost_,o["AS_cost"]; dims=3)
        BS_cost_ = cat(BS_cost_,o["BS_cost"]; dims=3)
    end
    ts_, AS_cost_, BS_cost_, Ns_, maxIter_
end

ts, AS_iter, BS_iter, Ns, maxIter = test_merge()
@save "simpleDemoCompareCostAll_results.jld2" ts AS_iter BS_iter Ns maxIter


text:
"""
    NCEI

The `NCEI` module provides an interface to the National Centers of Environmental
Information API. The Climate Data Online API is fully supported.

Please see the documentation for more information.
"""
module NCEI

using Dates: Day, Date, DateTime, TimeType, today, Year
using DataFrames: DataFrame
using HTTP: request
using LazyJSON: value
using LazyJSON.PropertyDicts: get

import Base: parse

foreach(include, ["Endpoint.jl", "API.jl"])

export
    Date,
    DateTime,
    cdo_datasets,
    cdo_datacategories,
    cdo_datatypes,
    cdo_locationcategories,
    cdo_locations,
    cdo_stations,
    cdo_data
end


text:
export run_acdctnepopf_bf

""
function run_acdctnepopf_bf(file::String, model_type::Type, solver; kwargs...)
    data = _PM.parse_file(file)
    process_additional_data!(data)
    return run_acdctnepopf_bf(data, model_type, solver; ref_extensions = [add_ref_dcgrid!, add_candidate_dcgrid!, _PM.ref_add_on_off_va_bounds!, _PM.ref_add_ne_branch!], kwargs...)
end

""
function run_acdctnepopf_bf(data::Dict{String,Any}, model_type::Type, solver; kwargs...)
    return _PM.run_model(data, model_type, solver, post_acdctnepopf_bf; ref_extensions = [add_ref_dcgrid!, add_candidate_dcgrid!, _PM.ref_add_on_off_va_bounds!, _PM.ref_add_ne_branch!], kwargs...)
end

""
function post_acdctnepopf_bf(pm::_PM.AbstractPowerModel)
    # PowerModelsACDC.add_ref_dcgrid!(pm)
    # add_candidate_dcgrid!(pm)
    _PM.variable_bus_voltage(pm)
    _PM.variable_gen_power(pm)
    _PM.variable_branch_power(pm)

    _PM.variable_branch_current(pm)
    variable_voltage_slack(pm)
    _PM.constraint_model_current(pm)

    variable_active_dcbranch_flow(pm)
    variable_dc_converter(pm)
    variable_dcbranch_current(pm)
    variable_dcgrid_voltage_magnitude(pm)

    # new variables for TNEP problem
    _PM.variable_ne_branch_indicator(pm)
    _PM.variable_ne_branch_power(pm)
    _PM.variable_ne_branch_current(pm)
    variable_active_dcbranch_flow_ne(pm)
    variable_branch_ne(pm)
    variable_dc_converter_ne(pm)
    variable_dcbranch_current_ne(pm)
    variable_dcgrid_voltage_magnitude_ne(pm)
    objective_min_cost_acdc(pm)

    _PM.constraint_model_voltage(pm)
    _PM.constraint_ne_model_voltage(pm)
    constraint_voltage_dc(pm)
    constraint_voltage_dc_ne(pm)
    for i in _PM.ids(pm, :ref_buses)
        _PM.constraint_theta_ref(pm, i)
    end
    for i in _PM.ids(pm, :bus)
        constraint_power_balance_acne_dcne(pm, i)
    end
    for i in _PM.ids(pm, :branch)
        _PM.constraint_ohms_yt_from(pm, i)
        _PM.constraint_ohms_yt_to(pm, i)
        _PM.constraint_voltage_angle_difference(pm, i)
        _PM.constraint_thermal_limit_from(pm, i)
        _PM.constraint_thermal_limit_to(pm, i)
    end
    for i in _PM.ids(pm, :ne_branch)
        _PM.constraint_ne_ohms_yt_from(pm, i)
        _PM.constraint_ne_ohms_yt_to(pm, i)
        _PM.constraint_ne_voltage_angle_difference(pm, i)
        _PM.constraint_ne_thermal_limit_from(pm, i)
        _PM.constraint_ne_thermal_limit_to(pm, i)
    end

    for i in _PM.ids(pm, :busdc)
        constraint_power_balance_dc_dcne(pm, i)
    end
    for i in _PM.ids(pm, :busdc_ne)
        constraint_power_balance_dcne_dcne(pm, i)
    end

    for i in _PM.ids(pm, :branchdc)
        constraint_ohms_dc_branch(pm, i)
    end
    for i in _PM.ids(pm, :branchdc_ne)
        constraint_ohms_dc_branch_ne(pm, i)
        constraint_branch_limit_on_off(pm, i)
    end

    for i in _PM.ids(pm, :convdc)
        constraint_converter_losses(pm, i)
        constraint_converter_current(pm, i)
        constraint_conv_transformer(pm, i)
        constraint_conv_reactor(pm, i)
        constraint_conv_filter(pm, i)
        if pm.ref[:it][:pm][:nw][_PM.nw_id_default][:convdc][i]["islcc"] == 1
            constraint_conv_firing_angle(pm, i)
        end
    end

    for i in _PM.ids(pm, :convdc_ne)
        constraint_converter_losses_ne(pm, i)
        constraint_converter_current_ne(pm, i)
        constraint_converter_limit_on_off(pm, i)
        constraint_conv_transformer_ne(pm, i)
        constraint_conv_reactor_ne(pm, i)
        constraint_conv_filter_ne(pm, i)
        if pm.ref[:it][:pm][:nw][_PM.nw_id_default][:convdc_ne][i]["islcc"] == 1
            constraint_conv_firing_angle_ne(pm, i)
        end
    end
end


text:
abstract type ClosedForm end

mutable struct CFiniteClosedForm <: ClosedForm
    f::SymFunction
    n::Sym
    exp::Array{Sym}
    coeff::Array{Sym}
    expvars::Array{Sym}
end

struct HyperClosedForm <: ClosedForm
    exp::Array{Sym} # TODO: do exponentials contain loop counter?
    fact::Array{Sym}
    coeff::Array{Sym}
end

struct ClosedFormSystem
    cforms::Array{ClosedForm,1}
    lc::Sym
    vars::Array{Sym,1}
end

#-------------------------------------------------------------------------------

CFiniteClosedForm(f::SymFunction, n::Sym, exp::Array{Sym}, coeff::Array{Sym}) = CFiniteClosedForm(f, n, exp, coeff, [])

function shift(cf::CFiniteClosedForm, sh::Sym)
    coeff = copy(cf.coeff)
    for i in 1:length(cf.exp)
        coeff[i] = coeff[i] |> subs(cf.n, cf.n+sh)
        coeff[i] *= cf.exp[i]^sh
    end
    CFiniteClosedForm(cf.f, cf.n, cf.exp, coeff, cf.expvars)
end

function poly(cf::CFiniteClosedForm)
    return sum([ex^cf.n * c for (ex, c) in zip(cf.exp, cf.coeff)])
end

function rhs(cf::CFiniteClosedForm)
    poly(cf)
end

function subs!(cf::CFiniteClosedForm, x::Pair...)
    if !isempty(x)
        cf.coeff = [c |> subs(x...) for c in cf.coeff]
    end
    cf
end

function polynomial(cf::CFiniteClosedForm)
    if isempty(cf.expvars)
        return sum(cf.coeff)
    else
        return sum(cf.expvars .* cf.coeff)
    end
    # error("No expvars given. This should not happen.")
end

exponentials(cf::ClosedForm) = cf.exp

function expvars!(cf::ClosedForm, d::Dict{Sym,Sym})
    cf.expvars = replace(cf.exp, d)
end

function Base.show(io::IO, cf::ClosedForm)
    print(io, "$(cf.f(cf.n)) = $(poly(cf))")
end

function Base.show(io::IO, cfs::Aligator.ClosedFormSystem)
    if get(io, :compact, false)
        show(io, cfs.cforms)
    else
        println(io, "$(length(cfs.cforms))-element $(typeof(cfs)):")
        for cf in cfs.cforms
            print(io, " ")
            show(io, cf)
            print(io, "\n")
        end
    end
end


text:
# See [Chapter 9, in Finite Element Procedures, K-J Bathe].
function forced_oscillator(; X0=(zeros(2), zeros(2)), X=Universe(2))
    # Mx'' + Kx = R
    M = [2 0; 0 1.]
    K = [6 -2; -2 4.]
    C = zeros(2, 2)
    R = [0, 10.]
    example_9_1_Bathe = SecondOrderAffineContinuousSystem(M, C, K, R)
    prob = InitialValueProblem(example_9_1_Bathe, X0)
    tspan = (0.0, 20.0)
    return prob, tspan
end

# analytic solution (see Example 9.7 in Bathe)
function forced_oscillator_solution()
    A = [1/√3  (1/2)*√(2/3);
         1/√3      -√(2/3)]
    x₁(t) = (5 / √3) * (1 - cos(t*√2))
    x₂(t) = (2 * √(2/3)) * (-1 + cos(t*√5))
    U(t) = A * [x₁(t), x₂(t)]
end


text:
function trim_arr(arr; trim_at_begin, trim_at_end)
    if (trim_at_begin + trim_at_end) > 1
        error("you are trying to remove more than you have")
    elseif (trim_at_begin < 0) | (trim_at_end < 0)
        error("you are trying to remove less than zero")
    end
    l = length(arr)
    trim_b = round(Int, float(l*trim_at_begin), RoundDown)
    trim_e = round(Int, float(l*trim_at_end), RoundDown)
    return arr[begin+trim_b:end-trim_e]
end

rms(ar) = sqrt(mean(ar.^2))

######

function arr_amplitude(;testarr, trim_at_begin=0, trim_at_end=0)
    ar = trim_arr(testarr; trim_at_begin, trim_at_end)
    return (;y=(maximum(ar)-minimum(ar)))
end

function arr_length(;testarr, trim_at_begin=0, trim_at_end=0)
    ar = trim_arr(testarr; trim_at_begin, trim_at_end)
    return (;y=length(ar))
end

function arr_mean(;testarr, trim_at_begin=0, trim_at_end=0)
    ar = trim_arr(testarr; trim_at_begin, trim_at_end)
    return (;y=mean(ar))
end

function arr_median(;testarr, trim_at_begin=0, trim_at_end=0)
    ar = trim_arr(testarr; trim_at_begin, trim_at_end)
    return (;y=median(ar))
end

function arr_std(;testarr, trim_at_begin=0, trim_at_end=0)
    ar = trim_arr(testarr; trim_at_begin, trim_at_end)
    return (;y=std(ar))
end

function arr_rms(;testarr, trim_at_begin=0, trim_at_end=0)
    ar = trim_arr(testarr; trim_at_begin, trim_at_end)
    return (;y=rms(ar))
end

function vec_amplify(;testarr, trim_at_begin=0, trim_at_end=0, p=1)
    vec = trim_arr(testarr; trim_at_begin, trim_at_end).*p
    return (;y=rms(vec), bigarrs=(;vec))
end

function vec_attenuate(;testarr, trim_at_begin=0, trim_at_end=0, p=1)
    vec = trim_arr(testarr; trim_at_begin, trim_at_end)./p
    return (;y=rms(vec), bigarrs=(;vec))
end

function vec_power(;testarr, trim_at_begin=0, trim_at_end=0, p=1)
    vec = abs.(trim_arr(testarr; trim_at_begin, trim_at_end)).^p
    return (;y=rms(vec), bigarrs=(;vec))
end


text:
# Goals: minimize calls to ordered_states (allocates memory)

# needs pomdp for state_index in pdf(b, s)
# needs list of ordered_states for rand(b)

"""
    DiscreteBelief

A belief specified by a probability vector.

Normalization of `b` is NOT enforced at all times, but the `DiscreteBeleif(pomdp, b)` constructor will warn, and `update(...)` always returns a belief with normalized `b`.
"""
struct DiscreteBelief{P<:POMDP, S}
    pomdp::P
    state_list::Vector{S}       # vector of ordered states
    b::Vector{Float64}
end

function DiscreteBelief(pomdp::POMDP, state_list::AbstractVector, b::AbstractVector{Float64}, check::Bool=true)
    if check
        if !isapprox(sum(b), 1.0, atol=0.001)
            warn("""
                 b in DiscreteBelief(pomdp, b) does not sum to 1.

                 To suppress this warning use `DiscreteBelief(pomdp, b, check=false)`
                 """)
        end
        if !all(0.0 <= p <= 1.0 for p in b)
            warn("""
                 b in DiscreteBelief(pomdp, b) contains entries outside [0,1].

                 To suppress this warning use `DiscreteBelief(pomdp, b, check=false)`
                 """)
        end
    end
    return DiscreteBelief(pomdp, state_list, b)
end

function DiscreteBelief(pomdp::POMDP, b::Vector{Float64}; check::Bool=true)
    return DiscreteBelief(pomdp, ordered_states(pomdp), b, check)
end

function DiscreteBelief(pomdp::POMDP, b; check::Bool=true)
    # convert b to a vector representation
    state_list = ordered_states(pomdp)
    bv = Vector{Float64}(n_states(pomdp))
    for (i, s) in enumerate(state_list)
        bv[i] = pdf(b, s)
    end
    return DiscreteBelief(pomdp, state_list, bv, check)
end

"""
Return a DiscreteBelief with equal probability for each state.
"""
function uniform_belief(pomdp)
    state_list = ordered_states(pomdp)
    ns = length(state_list)
    return DiscreteBelief(pomdp, state_list, ones(ns) / ns)
end

pdf(b::DiscreteBelief, s) = b.b[state_index(b.pomdp, s)]

function rand(rng::AbstractRNG, b::DiscreteBelief)
    i = sample(rng, Weights(b.b))
    return b.state_list[i]
end

function Base.fill!(b::DiscreteBelief, x::Float64)
    fill!(b.b, x)
    return b
end

Base.length(b::DiscreteBelief) = length(b.b)

iterator(b::DiscreteBelief) = b.state_list

==(b1::DiscreteBelief, b2::DiscreteBelief) = b1.state_list == b2.state_list && b1.b == b2.b

Base.hash(b::DiscreteBelief, h::UInt) = hash(b.b, hash(b.state_list, h))

mutable struct DiscreteUpdater{P<:POMDP} <: Updater
    pomdp::P
end

uniform_belief(up::DiscreteUpdater) = uniform_belief(up.pomdp)

function initialize_belief(bu::DiscreteUpdater, dist::Any)
    state_list = ordered_states(bu.pomdp)
    ns = length(state_list)
    b = zeros(ns)
    belief = DiscreteBelief(bu.pomdp, state_list, b)
    for s in iterator(dist)
        sidx = state_index(bu.pomdp, s)
        belief.b[sidx] = pdf(dist, s)
    end
    return belief
end

function update(bu::DiscreteUpdater, b::DiscreteBelief, a, o)
    pomdp = b.pomdp
    state_space = b.state_list
    bp = zeros(length(state_space))

    bp_sum = 0.0   # to normalize the distribution

    for (spi, sp) in enumerate(state_space)

        # po = O(a, sp, o)
        od = observation(pomdp, a, sp)
        po = pdf(od, o)

        if po == 0.0
            continue
        end

        b_sum = 0.0
        for (si, s) in enumerate(state_space)
            td = transition(pomdp, s, a)
            pp = pdf(td, sp)
            b_sum += pp * b.b[si]
        end

        bp[spi] = po * b_sum
        bp_sum += bp[spi]
    end

    if bp_sum == 0.0
        error("""
              Failed discrete belief update: new probabilities sum to zero.

              b = $b
              a = $a
              o = $o

              Failed discrete belief update: new probabilities sum to zero.
              """)
    else
        bp ./= bp_sum
    end

    return DiscreteBelief(pomdp, b.state_list, bp)
end

update(bu::DiscreteUpdater, b::Any, a, o) = update(bu, initialize_belief(bu, b), a, o)

# DEPRECATED
@generated function create_belief(bu::DiscreteUpdater)
    Core.println("WARNING: create_belief(up::DiscreteUpdater) is deprecated. Use uniform_belief(up) instead.")
    return :(uniform_belief(bu))
end

# alphas are |A|x|S|
# computes dot product of alpha vectors and belief
# util is array with utility of each alpha vecotr for belief b
@generated function product(alphas::Matrix{Float64}, b::DiscreteBelief)
    Core.println("WARNING: product(alphas, b::DiscreteBelief) is deprecated.")
    quote
        @assert size(alphas, 1) == length(b) "Alpha and belief sizes not equal"
        n = size(alphas, 2)
        util = zeros(n)
        for i = 1:n
            s = 0.0
            for j = 1:length(b)
                s += alphas[j,i]*b.b[j]
            end
            util[i] = s
        end
        return util
    end
end


text:
###################################################
## viz.jl
##      Visualizes Student/Schools/Stops data (without network)
## Authors: Arthur Delarue, Sébastien Martin, 2019
###################################################

struct RouteSFML
    stopids::Vector{Int}
    arcshapes::Vector{SFML.Line}
    clickpoints::Vector{Vector2f}
    bus::Int
end

PALETTE = Plots.get_color_palette(:auto, Plots.plot_color(:white), 17)
SFMLPALETTE = repeat([SFML.Color(round.(Int,[c.r,c.g,c.b]*255)...) for c in PALETTE], outer=20)
function visualize(data::SchoolBusData)
    nodeRadius = 100.

    # Defines the window, an event listener, and view
    window_w, window_h = 1200,1200
    window = SFML.RenderWindow("School Bus Problem", window_w, window_h)
    SFML.set_vsync_enabled(window, true)
    event = SFML.Event()

    # positions
    schoolpositions = [(school.position.x, school.position.y) for school in data.schools]
    yardpositions = [(yard.position.x, yard.position.y) for yard in data.yards]
    tmp = Dict(stop.uniqueId => (stop.position.x, stop.position.y) for stop in vcat(data.stops...))
    stoppositions = [tmp[i] for i=1:maximum(keys(tmp))]
    allpositions = vcat(schoolpositions,yardpositions,stoppositions)

    #Set the window size
    minX, maxX = extrema(x for (x,y) in allpositions)
    minY, maxY = extrema(y for (x,y) in allpositions)

    # Do the Y-axis transformation
    minY, maxY = -maxY, -minY
    networkLength = max(maxX-minX, maxY-minY)
    viewWidth = max(maxX-minX, (maxY-minY)*window_w/window_h)
    viewHeigth = max(maxY-minY, (maxX-minX)*window_h/window_w)
    view = SFML.View(SFML.Vector2f((minX+maxX)/2,(minY+maxY)/2), SFML.Vector2f(viewWidth, viewHeigth))
    zoomLevel = 1.0

    # init visualizer
    hideStops = false
    routemode = false
    selectedSchool = 0
    selectedbus = 0
    schoolShapes = [SFML.CircleShape() for i in eachindex(data.schools)]
    for s in schoolShapes
        SFML.set_pointcount(s,5)
    end

    stopShapes = [SFML.CircleShape() for i in eachindex(stoppositions)]
    for s in stopShapes
        SFML.set_pointcount(s,4)
        SFML.set_fillcolor(s, SFML.Color(255,0,0))
    end

    function drawroute(sfmlroute)
        for cornerstopid in sfmlroute.stopids
            SFML.draw(window, stopShapes[cornerstopid])
        end
        for routeedge in sfmlroute.arcshapes
            SFML.draw(window, routeedge)
        end
    end

    if data.withFinalBuses
        routetobus = Dict{Tuple{Int,Int},Int}() # (schoolid, routeid)->busid
        for bus in data.buses
            for (schoolid,routeid) in zip(bus.schools,bus.routes)
                routetobus[schoolid,routeid] = bus.id
            end
        end
        sfmlroutes = RouteSFML[]
        schoolroutes = Vector{Int}[Int[] for school in data.schools]
        busroutes = Vector{Int}[Int[] for bus in data.buses]
        for school in data.schools
            chosenscenario = data.scenarios[school.id][data.usedScenario[school.id]]
            routes = data.routes[school.id][chosenscenario.routeIDs]

            for route in routes
                routesegments = SFML.Line[]
                cornerstops = Int[]
                clickpoints = SFML.Vector2f[]

                busstopids = route.stops

                originstop = data.stops[school.id][busstopids[1]]
                originpos = (originstop.position.x, originstop.position.y)
                push!(cornerstops, originstop.uniqueId)

                for i in 2:length(busstopids)
                    deststop = data.stops[school.id][busstopids[i]]
                    destpos = (deststop.position.x, deststop.position.y)
                    l = SFML.Line(SFML.Vector2f(originpos[1], -originpos[2]), SFML.Vector2f(destpos[1], -destpos[2]))
                    SFML.set_thickness(l, nodeRadius/1.5)
                    SFML.set_fillcolor(l, SFML.Color(100,100,100))
                    push!(routesegments,l)
                    push!(cornerstops, deststop.uniqueId)

                    dist = sqrt((originpos[1]-destpos[1])^2 + (originpos[2]-destpos[2])^2)
                    xs = linspace(originpos[1], destpos[1], 2+floor(Int, dist/50.))
                    ys = linspace(originpos[2], destpos[2], 2+floor(Int, dist/50.))
                    for (x,y) in zip(xs,ys)
                        push!(clickpoints, SFML.Vector2f(x,y))
                    end
                    originpos = destpos
                end
                # add last move to school
                destpos = schoolpositions[school.id]
                l = SFML.Line(SFML.Vector2f(originpos[1], -originpos[2]), SFML.Vector2f(destpos[1], -destpos[2]))
                SFML.set_thickness(l, nodeRadius/1.5)
                SFML.set_fillcolor(l, SFML.Color(100,100,100))
                dist = sqrt((originpos[1]-destpos[1])^2 + (originpos[2]-destpos[2])^2)
                xs = linspace(originpos[1], destpos[1], 2+floor(Int, dist/50.))
                ys = linspace(originpos[2], destpos[2], 2+floor(Int, dist/50.))
                for (x,y) in zip(xs,ys)
                    push!(clickpoints, SFML.Vector2f(x,y))
                end
                push!(routesegments,l)

                push!(sfmlroutes, RouteSFML(cornerstops,routesegments, clickpoints, routetobus[school.id,route.id]))
                push!(schoolroutes[school.id], length(sfmlroutes))
                push!(busroutes[routetobus[school.id,route.id]], length(sfmlroutes))
            end
        end
        connectingsegments = Vector{SFML.Line}[SFML.Line[] for bus in data.buses]
        for bus in data.buses
            if length(bus.schools) > 1
                for i in 2:length(bus.schools)
                    schoolpos = schoolpositions[bus.schools[i-1]]
                    route = data.routes[bus.schools[i]][bus.routes[i]]
                    firststoppos = (data.stops[bus.schools[i]][route.stops[1]].position.x,
                                    data.stops[bus.schools[i]][route.stops[1]].position.y)
                    l = SFML.Line(SFML.Vector2f(schoolpos[1], -schoolpos[2]), SFML.Vector2f(firststoppos[1], -firststoppos[2]))
                    SFML.set_thickness(l, nodeRadius/1.5)
                    SFML.set_fillcolor(l, SFML.Color(100,100,100))
                    push!(connectingsegments[bus.id], l)
                end
            end
        end
    end

    schoolstops = Int[] # school stops to be drawn
    walkingshapes = SFML.Line[]

    function selectschool(schoolId)
        selectedSchool = schoolId
        schoolstops = Int[]
        for stop in data.stops[schoolId]
            cornerstopid = stop.uniqueId
            stopposition = stoppositions[cornerstopid]
            push!(schoolstops, cornerstopid)
        end
    end

    function redraw()
        schoolRadius = selectedSchool > 0 ? nodeRadius*1.5 : nodeRadius*4
        stopRadius = nodeRadius*2
        for (i,s) in enumerate(schoolShapes)
            SFML.set_radius(s, schoolRadius)
            SFML.set_fillcolor(s, SFMLPALETTE[i])

            SFML.set_position(s, SFML.Vector2f(schoolpositions[i][1] - schoolRadius,-schoolpositions[i][2] - schoolRadius))
        end
        for (i,s) in enumerate(stopShapes)
            SFML.set_radius(s, stopRadius)
            SFML.set_position(s, SFML.Vector2f(stoppositions[i][1] - stopRadius,-stoppositions[i][2] - stopRadius))
        end
        if selectedSchool > 0
            for (i,s) in enumerate(schoolShapes)
                SFML.set_fillcolor(s, SFML.Color(255,255,255))
            end
            s = schoolShapes[selectedSchool]
#             SFML.set_fillcolor(s, SFML.Color(0,255,0))
            SFML.set_radius(s, 4*schoolRadius)
            SFML.set_fillcolor(s, SFMLPALETTE[selectedSchool])

            SFML.set_position(s, SFML.Vector2f(schoolpositions[selectedSchool][1] - 4*schoolRadius,-schoolpositions[selectedSchool][2] - 4*schoolRadius))
        end
        if data.withFinalBuses
            for route in sfmlroutes, (arcid, arc) in enumerate(route.arcshapes)
                SFML.set_thickness(arc, nodeRadius/1.5)
            end
            for busconnectors in connectingsegments, line in busconnectors
                SFML.set_thickness(line, nodeRadius/1.5)
            end
        end
    end
    redraw()

    # Constructing school position tree
    schoolPos = Array{Float64}(2,length(data.schools))
    for (i,school) in enumerate(data.schools)
       schoolPos[1,i] = schoolpositions[i][1]
       schoolPos[2,i] = schoolpositions[i][2]
    end
    schoolTree = NearestNeighbors.KDTree(schoolPos)

    if data.withFinalBuses
    # Constructing route position tree
        positions = Array{Float64}(2,sum(length(route.clickpoints) for route in sfmlroutes))
        clickpoint2route = Int[]
        i = 1
        for (routeid,route) in enumerate(sfmlroutes), clickpoint in route.clickpoints
            positions[1,i] = clickpoint[1]
            positions[2,i] = clickpoint[2]
            push!(clickpoint2route, routeid)
            i+=1
        end
        routeTree = NearestNeighbors.KDTree(positions)
    end

    clock = SFML.Clock()
    # gc_enable(false)
    while SFML.isopen(window)
        frameTime = Float64(SFML.as_seconds(SFML.restart(clock)))
        while SFML.pollevent(window, event)
            if SFML.get_type(event) == SFML.EventType.CLOSED
                SFML.close(window)
            end
            if SFML.get_type(event) == SFML.EventType.RESIZED
                size = SFML.get_size(event)
                window_w, window_h = size.width, size.height
                viewWidth = max(maxX-minX, (maxY-minY)*window_w/window_h)
                viewHeigth = max(maxY-minY, (maxX-minX)*window_h/window_w)
                SFML.set_size(view, SFML.Vector2f(viewWidth, viewHeigth))
                SFML.zoom(view, zoomLevel)
            end
            if SFML.get_type(event) == SFML.EventType.KEY_PRESSED
                k = SFML.get_key(event).key_code
                if k == SFML.KeyCode.ESCAPE || k == SFML.KeyCode.Q
                    SFML.close(window)
                    break;
                elseif k == SFML.KeyCode.A
                    nodeRadius *= 1.3
                    redraw()
                elseif k == SFML.KeyCode.S
                    nodeRadius /= 1.3
                    redraw()
                elseif k == SFML.KeyCode.G
                    if selectedSchool == 0
                        hideStops = !hideStops
                    end
                elseif k == SFML.KeyCode.R
                    if data.withFinalBuses
                        routemode = !routemode
                    end
                end
            end
            if SFML.get_type(event) == SFML.EventType.MOUSE_BUTTON_PRESSED && SFML.get_mousebutton(event).button == SFML.MouseButton.LEFT
                x,y = SFML.get_mousebutton(event).x, SFML.get_mousebutton(event).y
                coord = SFML.pixel2coords(window,SFML.Vector2i(x,y))
                if routemode
                    if selectedbus == 0
                        clickpoint = NearestNeighbors.knn(routeTree, [Float64(coord.x),-Float64(coord.y)],1)[1][1]
                        routeid = clickpoint2route[clickpoint]
                        selectedbus = sfmlroutes[routeid].bus
                        selectedSchool = 0
                    else
                        selectedbus = 0
                    end
                else
                    schoolId = NearestNeighbors.knn(schoolTree,[Float64(coord.x),-Float64(coord.y)],1)[1][1]
                    if schoolId == selectedSchool
                        selectedSchool = 0
                    else
                        selectschool(schoolId)
                    end
                end
                redraw()
            end
        end
        if SFML.is_key_pressed(SFML.KeyCode.LEFT)
            SFML.move(view, SFML.Vector2f(-networkLength/2*frameTime*zoomLevel,0.))
        end
        if SFML.is_key_pressed(SFML.KeyCode.RIGHT)
            SFML.move(view, SFML.Vector2f(networkLength/2*frameTime*zoomLevel,0.))
        end
        if SFML.is_key_pressed(SFML.KeyCode.UP)
            SFML.move(view, SFML.Vector2f(0.,-networkLength/2*frameTime*zoomLevel))
        end
        if SFML.is_key_pressed(SFML.KeyCode.DOWN)
            SFML.move(view, SFML.Vector2f(0.,networkLength/2*frameTime*zoomLevel))
        end
        if SFML.is_key_pressed(SFML.KeyCode.Z)
            SFML.zoom(view, 0.6^frameTime)
            zoomLevel = SFML.get_size(view).x/viewWidth
        end
        if SFML.is_key_pressed(SFML.KeyCode.X)
            SFML.zoom(view, 1/(0.6^frameTime))
            zoomLevel = SFML.get_size(view).x/viewWidth
        end
        SFML.set_view(window,view)
        SFML.clear(window, SFML.Color(255,255,255))

        # Draw students
        if selectedSchool != 0
            if routemode
                for sfmlrouteid in schoolroutes[selectedSchool]
                    drawroute(sfmlroutes[sfmlrouteid])
                end
            else
                for walkingline in walkingshapes
                    SFML.draw(window, walkingline)
                end
                # only draw corner stops that are used
                for s in stopShapes[schoolstops]
                    SFML.draw(window, s)
                end
            end
        else
            if routemode
                if selectedbus > 0
                    for sfmlrouteid in busroutes[selectedbus]
                        drawroute(sfmlroutes[sfmlrouteid])
                    end
                    for line in connectingsegments[selectedbus]
                        SFML.draw(window, line)
                    end
                else
                    for route in sfmlroutes
                        drawroute(route)
                    end
                end
            else
                # Draw stops
                if !hideStops
                    for s in stopShapes
                        SFML.draw(window, s)
                    end
                end
            end
        end
        # Draw schools
        for s in schoolShapes
            SFML.draw(window, s)
        end
        SFML.display(window)
    end
end

max_stars_repo_name: jihoyeo/SchoolBusRouting

max_issues_repo_name: jihoyeo/SchoolBusRouting

max_forks_repo_name: jihoyeo/SchoolBusRouting


text:
using PythonCallChainRules.Torch: TorchModuleWrapper, ispysetup

using Test

using PythonCall
using DLPack

if !ispysetup[]
    return
end

torch = pyimport("torch")
dlpack = pyimport("torch.utils.dlpack")

device = torch.device("cpu")

@testset "dlpack" begin
    for dims in ((10,), (1, 10), (2, 3, 5), (2, 3, 4, 5))
        xto = torch.randn(dims..., device=device)
        xjl = DLPack.wrap(xto, dlpack.to_dlpack)
        @test Tuple(xto.size()) == reverse(size(xjl))
        @test isapprox(pyconvert(Float32, xto.cpu().sum().item()), sum(xjl))
    end
end

function compare_grad_wrt_params(modelwrap, inputs...)
    params = map(x -> DLPack.share(x, Py, dlpack.from_dlpack).to(device = device, dtype = modelwrap.dtype).requires_grad_(true), (modelwrap.params))
    torch_out = modelwrap.torch_stateless_module(params, modelwrap.buffers, map(z->DLPack.share(z, Py, dlpack.from_dlpack).to(dtype=modelwrap.dtype, device=device), inputs)...).sum()
    torchgrad = map(x-> (x.cpu().numpy()), torch.autograd.grad(torch_out, params))
    grad,  = Zygote.gradient(m->sum(m(inputs...)), modelwrap)
    @test length(torchgrad) == length(grad.params)
    for i in 1:length(grad.params)
        @test isapprox(sum(torchgrad[i]), sum(grad.params[i]))
    end
    @test length(grad.params) == length(modelwrap.params)
    @test grad.params[1] !== nothing
    @test grad.params[2] !== nothing
    @test size(grad.params[1]) == size(modelwrap.params[1])
    @test size(grad.params[2]) == size(modelwrap.params[2])
end

function compare_grad_wrt_inputs(modelwrap, x)
    params = map(z -> DLPack.share(z, Py, dlpack.from_dlpack).to(device = device, dtype = modelwrap.dtype).requires_grad_(true), (modelwrap.params))
    xtorch = DLPack.share(copy(x), Py, dlpack.from_dlpack).to(dtype=modelwrap.dtype, device=device).requires_grad_(true)
    torch_out = modelwrap.torch_stateless_module(params, modelwrap.buffers, xtorch).sum()
    torchgrad = map(z-> (copy(z.cpu().numpy())), torch.autograd.grad(torch_out, xtorch))[1]
    grad, = Zygote.gradient(z->sum(modelwrap(z)), x)
    @test size(grad) == size(x)
    @test length(torchgrad) == length(grad)
    @test isapprox(sum(torchgrad), sum(grad))
end

batchsize = 1
indim = 3
outdim = 2
hiddendim = 4

#@testset "linear" begin
    lin = torch.nn.Linear(indim, outdim).to(device=device)
    linwrap = TorchModuleWrapper(lin)
    # if CUDA.functional()
    #     linwrap = fmap(CUDA.cu, linwrap)
    # end
    x = randn(Float32, indim, batchsize)
    # if CUDA.functional()
    #     x = cu(x)
    # end
    y = linwrap(x)
    @test size(y) == (outdim, batchsize)
    compare_grad_wrt_params(linwrap, x)
    compare_grad_wrt_inputs(linwrap, x)

#end

@testset "mlp" begin
    mlp = torch.nn.Sequential(torch.nn.Linear(indim, hiddendim), torch.nn.ReLU(), torch.nn.Linear(hiddendim, outdim)).to(device=device)
    mlpwrap = TorchModuleWrapper(mlp)
    # if CUDA.functional()
    #     mlpwrap = fmap(CUDA.cu, mlpwrap)
    # end

    x = randn(Float32, indim, batchsize)
    # if CUDA.functional()
    #     x = cu(x)
    # end
    y = mlpwrap(x)
    @test size(y) == (outdim, batchsize)
    compare_grad_wrt_params(mlpwrap, x)
    compare_grad_wrt_inputs(mlpwrap, x)
end


text:
@init @require Stan="682df890-35be-576f-97d0-3d8c8b33a550" begin

  using Turing, Stan

  @model gdemo(x) = begin
    s ~ InverseGamma(2,3)
    m ~ Normal(0,sqrt.(s))
    x[1] ~ Normal(m, sqrt.(s))
    x[2] ~ Normal(m, sqrt.(s))
    return s, m
  end

  model_f = gdemo([1.5, 2.0])

  chn = sample(model_f, 2000, 1000, false, 1, Stan.Adapt(), Stan.Hmc())

end


text:
@static if !isdefined(Base, :IdentityUnitRange)
    const IdentityUnitRange = Base.Slice
else
    using Base: IdentityUnitRange
end

export ModifiedOffsetArray, ModifiedOffsetVector

"""
An `OffsetMod` results in an array for which the indices start at a value
different from `1`. The modifier is functionally equivalent to an `OffsetArray`
"""
struct OffsetMod{N} <: ArrayModifier
    offsets ::  NTuple{N,Int}
end

ModifiedOffsetArray{T,N,AA} = ModifiedArray{T,N,AA,OffsetMod{N}}
ModifiedOffsetVector{T,AA} = ModifiedOffsetArray{T,1,AA}

ModifiedOffsetArray(A::AbstractArray{T,N}, offsets::NTuple{N,Int}) where {T,N} =
    ModifiedArray{T,N}(A, OffsetMod(offsets))
ModifiedOffsetArray(A::AbstractArray{T,N}, offsets::Vararg{Int,N}) where {T,N} =
    ModifiedOffsetArray(A, offsets)

ModStyle(::OffsetMod, ::IF_eltype) = ModNothing()
ModStyle(::OffsetMod, ::IF_ndims) = ModNothing()
ModStyle(::OffsetMod, ::IF_size) = ModNothing()

ModStyle(::OffsetMod, ::IF_getindex) = ModRecursive()
ModStyle(::OffsetMod, ::IF_setindex!) = ModRecursive()
ModStyle(::Type{<:OffsetMod}, ::IF_IndexStyle) = ModNothing()
ModStyle(::OffsetMod, ::IF_axes) = ModRecursive()
ModStyle(::OffsetMod, ::IF_similar) = ModNothing()

mod_getindex_pre(mod::OffsetMod{N}, I::Vararg{Int,N}) where {N} = offset(mod.offsets, I)
mod_getindex_pre(mod::OffsetMod, i::Int) = (i,)

mod_setindex!_pre(mod::OffsetMod{N}, val, I::Vararg{Int,N}) where {N} = (val,offset(mod.offsets, I)...)
mod_setindex!_pre(mod::OffsetMod, val, i::Int) = (val, i)

mod_axes_post(mod::OffsetMod, Z) = _axes(Z, mod.offsets)
_axes(inds, offsets) =
    (_slice(inds[1], offsets[1]), _axes(tail(inds), tail(offsets))...)
_axes(::Tuple{}, ::Tuple{}) = ()

# Avoid the kw-arg on the range(r+x, length=length(r)) call in r .+ x
@inline _slice(r, x) = IdentityUnitRange(Base._range(first(r) + x, nothing, nothing, length(r)))

## Take from OffsetArrays source code:
# Computing a shifted index (subtracting the offset)
@inline offset(offsets::NTuple{N,Int}, inds::NTuple{N,Int}) where {N} =
    (inds[1]-offsets[1], offset(Base.tail(offsets), Base.tail(inds))...)
offset(::Tuple{}, ::Tuple{}) = ()

# Support trailing 1s
@inline offset(offsets::Tuple{Vararg{Int}}, inds::Tuple{Vararg{Int}}) =
    (offset(offsets, Base.front(inds))..., inds[end])
offset(offsets::Tuple{Vararg{Int}}, inds::Tuple{}) = error("inds cannot be shorter than offsets")

indexoffset(r::AbstractRange) = first(r) - 1
indexoffset(i::Integer) = 0
indexlength(r::AbstractRange) = length(r)
indexlength(i::Integer) = i

# Everything below here is copied from OffsetArrays and translated to the current setting

using Base: Indices, Dims

const OffsetAxis = Union{Integer, UnitRange, Base.Slice, Base.OneTo, IdentityUnitRange}

## similar

function Base.similar(A::AbstractArray, ::Type{T}, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}}) where {T}
    B = similar(A, T, map(indexlength, inds))
    ModifiedOffsetArray(B, map(indexoffset, inds))
end

Base.similar(::Type{T}, shape::Tuple{OffsetAxis,Vararg{OffsetAxis}}) where {T<:AbstractArray} =
    ModifiedOffsetArray(T(undef, map(indexlength, shape)), map(indexoffset, shape))

## reshape

Base.reshape(A::AbstractArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}}) =
    ModifiedOffsetArray(reshape(A, map(indexlength, inds)), map(indexoffset, inds))

# Reshaping OffsetArrays can "pop" the original OffsetArray wrapper and return
# an OffsetArray(reshape(...)) instead of an OffsetArray(reshape(OffsetArray(...)))
Base.reshape(A::ModifiedOffsetArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}}) =
    ModifiedOffsetArray(reshape(parent(A), map(indexlength, inds)), map(indexoffset, inds))
# And for non-offset axes, we can just return a reshape of the parent directly
Base.reshape(A::ModifiedOffsetArray, inds::Tuple{Union{Integer,Base.OneTo},Vararg{Union{Integer,Base.OneTo}}}) = reshape(parent(A), inds)
Base.reshape(A::ModifiedOffsetArray, inds::Dims) = reshape(parent(A), inds)

## fill

Base.fill(v, inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {N} =
    fill!(ModifiedOffsetArray(Array{typeof(v), N}(undef, map(indexlength, inds)), map(indexoffset, inds)), v)
Base.zeros(::Type{T}, inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {T, N} =
    fill!(ModifiedOffsetArray(Array{T, N}(undef, map(indexlength, inds)), map(indexoffset, inds)), zero(T))
Base.ones(::Type{T}, inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {T, N} =
    fill!(ModifiedOffsetArray(Array{T, N}(undef, map(indexlength, inds)), map(indexoffset, inds)), one(T))
Base.trues(inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {N} =
    fill!(ModifiedOffsetArray(BitArray{N}(undef, map(indexlength, inds)), map(indexoffset, inds)), true)
Base.falses(inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {N} =
    fill!(ModifiedOffsetArray(BitArray{N}(undef, map(indexlength, inds)), map(indexoffset, inds)), false)

# Compatibility with OffsetArray constructors

const ArrayInitializer = Union{UndefInitializer, Missing, Nothing}
ModifiedOffsetArray{T,N}(init::ArrayInitializer, inds::Indices{N}) where {T,N} =
    ModifiedOffsetArray{T,N,Array{T,N}}(Array{T,N}(init, map(indexlength, inds)), map(indexoffset, inds))
ModifiedOffsetArray{T}(init::ArrayInitializer, inds::Indices{N}) where {T,N} = ModifiedOffsetArray{T,N}(init, inds)
ModifiedOffsetArray{T,N}(init::ArrayInitializer, inds::Vararg{AbstractUnitRange,N}) where {T,N} = ModifiedOffsetArray{T,N}(init, inds)
ModifiedOffsetArray{T}(init::ArrayInitializer, inds::Vararg{AbstractUnitRange,N}) where {T,N} = ModifiedOffsetArray{T,N}(init, inds)
ModifiedOffsetArray(A::AbstractArray{T,0}) where {T} = ModifiedOffsetArray{T,0,typeof(A)}(A, ())

# ModifiedOffsetVector constructors
ModifiedOffsetVector(A::AbstractVector, offset) = ModifiedOffsetArray(A, offset)
ModifiedOffsetVector{T}(init::ArrayInitializer, inds::AbstractUnitRange) where {T} = ModifiedOffsetArray{T}(init, inds)

ModifiedOffsetArray{T,N,AA}(A::AbstractArray{T,N}, offsets::NTuple{N,Int}) where {T,N,AA} =
    modify(A, OffsetMod(offsets))

# The next two are necessary for ambiguity resolution. Really, the
# second method should not be necessary.
ModifiedOffsetArray(A::AbstractArray{T,0}, inds::Tuple{}) where {T} = ModifiedOffsetArray{T,0,typeof(A)}(A, ())
ModifiedOffsetArray(A::AbstractArray{T,N}, inds::Tuple{}) where {T,N} = error("this should never be called")
function ModifiedOffsetArray(A::AbstractArray{T,N}, inds::NTuple{N,AbstractUnitRange}) where {T,N}
    lA = map(indexlength, axes(A))
    lI = map(indexlength, inds)
    lA == lI || throw(DimensionMismatch("supplied axes do not agree with the size of the array (got size $lA for the array and $lI for the indices"))
    ModifiedOffsetArray(A, map(indexoffset, inds))
end
ModifiedOffsetArray(A::AbstractArray{T,N}, inds::Vararg{AbstractUnitRange,N}) where {T,N} =
    ModifiedOffsetArray(A, inds)

# avoid a level of indirection when nesting ModifiedOffsetArrays
function ModifiedOffsetArray(A::ModifiedOffsetArray, inds::NTuple{N,AbstractUnitRange}) where {N}
    ModifiedOffsetArray(parent(A), inds)
end
ModifiedOffsetArray(A::ModifiedOffsetArray{T,0}, inds::Tuple{}) where {T} = ModifiedOffsetArray{T,0,typeof(A)}(parent(A), ())
ModifiedOffsetArray(A::ModifiedOffsetArray{T,N}, inds::Tuple{}) where {T,N} = error("this should never be called")

### Special handling for AbstractRange

const OffsetRange{T} = ModifiedOffsetArray{T,1,<:AbstractRange{T}}
const IIUR = IdentityUnitRange{S} where S<:AbstractUnitRange{T} where T<:Integer

Base.step(a::OffsetRange) = step(parent(a))

Base.getindex(a::OffsetRange, r::OffsetRange) = ModifiedOffsetArray(a[parent(r)], r.modifier.offsets)
Base.getindex(a::OffsetRange, r::AbstractRange) = a.parent[r .- a.modifier.offsets[1]]
Base.getindex(a::AbstractRange, r::OffsetRange) = ModifiedOffsetArray(a[parent(r)], r.modifier.offsets)

Base.getindex(r::UnitRange, s::IIUR) =
    ModifiedOffsetArray(r[s.indices], s)

Base.getindex(r::StepRange, s::IIUR) =
    ModifiedOffsetArray(r[s.indices], s)

Base.getindex(r::StepRangeLen{T,<:Base.TwicePrecision,<:Base.TwicePrecision}, s::IIUR) where T =
    ModifiedOffsetArray(r[s.indices], s)
Base.getindex(r::StepRangeLen{T}, s::IIUR) where {T} =
    ModifiedOffsetArray(r[s.indices], s)

Base.getindex(r::LinRange, s::IIUR) =
    ModifiedOffsetArray(r[s.indices], s)

function Base.show(io::IO, r::OffsetRange)
    show(io, r.parent)
    o = r.modifier.offsets[1]
    print(io, " with indices ", o+1:o+length(r))
end
Base.show(io::IO, ::MIME"text/plain", r::OffsetRange) = show(io, r)

### Convenience functions ###

Base.fill(x, inds::Tuple{UnitRange,Vararg{UnitRange}}) =
    fill!(ModifiedOffsetArray{typeof(x)}(undef, inds), x)
@inline Base.fill(x, ind1::UnitRange, inds::UnitRange...) = fill(x, (ind1, inds...))

## IO routines

function Base.showarg(io::IO, a::ModifiedOffsetArray, toplevel)
    print(io, "ModifiedOffsetArray(")
    Base.showarg(io, parent(a), false)
    if ndims(a) > 0
        print(io, ", ")
        printindices(io, axes(a)...)
    end
    print(io, ')')
    toplevel && print(io, " with eltype ", eltype(a))
end
printindices(io::IO, ind1, inds...) =
    (print(io, _unslice(ind1), ", "); printindices(io, inds...))
printindices(io::IO, ind1) = print(io, _unslice(ind1))
_unslice(x) = x
_unslice(x::IdentityUnitRange) = x.indices

### Some mutating functions defined only for OffsetVector ###

Base.resize!(A::ModifiedOffsetVector, nl::Integer) = (resize!(A.parent, nl); A)
Base.push!(A::ModifiedOffsetVector, x...) = (push!(A.parent, x...); A)
Base.pop!(A::ModifiedOffsetVector) = pop!(A.parent)
Base.empty!(A::ModifiedOffsetVector) = (empty!(A.parent); A)

function no_offset_view(A::AbstractArray)
    if Base.has_offset_axes(A)
        ModifiedOffsetArray(A, map(r->1-first(r), axes(A)))
    else
        A
    end
end

no_offset_view(A::ModifiedOffsetArray) = no_offset_view(parent(A))


text:
module YAAD

using DiffRules

abstract type AbstractNode end
abstract type Operator end

# core
include("comput_graph.jl")

# operators
include("operators/broadcast.jl")
include("operators/math.jl")
include("operators/linalg.jl")
include("operators/array.jl")
include("operators/cat.jl")
include("operators/reduce.jl")
# include("operators/batched.jl")

# printings
include("show.jl")

# utils
include("utils.jl")
# test utils
include("test_utils.jl")

end # module


text:
#=
This first example comes almost straight from the documentation for DiffEqFlux

https://github.com/SciML/DiffEqFlux.jl#optimizing-parameters-of-an-ode-for-an-optimal-control-problem
https://julialang.org/blog/2019/01/fluxdiffeq/

=#

module lotka_volterra

using DifferentialEquations, Flux, Optim, DiffEqFlux, Plots

u0 = [1.0,1.0]
tstart=0.0
tend=10.0
sampling=0.1

model_params= [1.5,1.0,3.0,1.0]

function model(du,u,p,t)
  x, y = u
  α, β, δ, γ = p
  du[1] = dx = α*x - β*x*y
  du[2] = dy = -δ*y + γ*x*y
end

function InitPlot(;vars=(0,1))
    #Simple function to plot the dynamics
    tspan=(tstart,tend)
    prob = ODEProblem(model,u0,tspan,model_params)
    sol = solve(prob,Tsit5())
    plot(sol, lw=2, legend=false, vars=(0,1))
end

function predict_adjoint(param) # Our 1-layer neural network
    prob=ODEProblem(model,[1.0,0.0],(tstart,tend), model_params)
    Array(concrete_solve(prob,Tsit5(),param[1:2],param[3:end],saveat=tstart:sampling:tend, abstol=1e-8,reltol=1e-6))
end

# Generate some data to fit, and add some noise to it
data=predict_adjoint([1.0,1.0,1.5,1.0,3.0,1.0])
σN=0.1
data+=σN*randn(size(data))
data=abs.(data) #Keep measurements positive

# Returning more than just the loss function breaks the Flux optim
function loss_adjoint(param)
  prediction = predict_adjoint(param)
  loss = sum(abs2,prediction - data)
  loss
end;

#Test

function train_model(;pguess=[0.8,1.2,1.2,1.0,2.9,1.1])
    println("The initial loss function is $(loss_adjoint(pguess)[1])")
    #Train the ODE
    resinit=DiffEqFlux.sciml_train(loss_adjoint,pguess,ADAM(), maxiters=3000)
    res = DiffEqFlux.sciml_train(loss_adjoint,resinit.minimizer,BFGS(initial_stepnorm = 1e-5))
    println("The parameters are $(res.minimizer) with final loss value $(res.minimum)")
    return(res)
end

function plotFit(param)

    tspan=(tstart,tend)
    sol_fit=solve(ODEProblem(model,param[1:2],tspan,param[3:end]), Tsit5())

    tgrid=tstart:sampling:tend
    pl=plot(sol_fit, lw=2, legend=false)
    scatter!(pl,tgrid, data[1,:], color=:blue)
    scatter!(pl,tgrid, data[2,:], color=:red)
    xlabel!(pl,"Time")
    ylabel!(pl,"Population")
    title!(pl,"Model Parameter Fits")
    savefig("Lotka_Volterra_ParamFit.png")
    display(pl)
end

function validationPlot(param, ic)
    tspan=(tstart,tend)
    sol_fit=solve(ODEProblem(model,ic,tspan,param[3:end]), Tsit5())
    sol_actual=solve(ODEProblem(model,ic,tspan,model_params), Tsit5(), saveat=0.0:0.1:10.0)

    pl=scatter(sol_actual, lw=2.0, color=:blue, vars=(0,1))
    scatter!(sol_actual, color=:red, vars=(0,2))
    plot!(pl, sol_fit, lw=2, legend=false, color=:blue, vars=(0,1))
    plot!(pl, sol_fit, lw=2, color=:red, vars=(0,2))
    xlabel!(pl,"Rabbits")
    ylabel!(pl,"Lynx")
    title!(pl,"Validation Plot")
    savefig("Lotka_Volterra_Validation_Plot.png")
    display(pl)
end

end #module

lotka_volterra.InitPlot()

resL=lotka_volterra.train_model()

lotka_volterra.plotFit(resL.minimizer)

lotka_volterra.validationPlot(resL.minimizer, [6.,6.])

max_stars_repo_name: khannay/FittingParamsDiffEqFlux

max_issues_repo_name: khannay/FittingParamsDiffEqFlux

max_forks_repo_name: khannay/FittingParamsDiffEqFlux


text:
function stepsize_controller!(integrator::SDDEIntegrator, controller::PIController, alg)
  integrator.q11 = DiffEqBase.value(DiffEqBase.fastpow(integrator.EEst,controller.beta1))
  integrator.q = DiffEqBase.value(integrator.q11/DiffEqBase.fastpow(integrator.qold,controller.beta2))
  @fastmath integrator.q = DiffEqBase.value(max(inv(integrator.opts.qmax),min(inv(integrator.opts.qmin),integrator.q/integrator.opts.gamma)))
end

@inline function step_accept_controller!(integrator::SDDEIntegrator, alg)
  step_accept_controller!(integrator, integrator.opts.controller, alg)
end

function step_accept_controller!(integrator::SDDEIntegrator, controller::PIController, alg)
  integrator.dtnew = DiffEqBase.value(integrator.dt/integrator.q) * oneunit(integrator.dt)
end

function step_reject_controller!(integrator::SDDEIntegrator, controller::PIController, alg)
  integrator.dtnew = integrator.dt/min(inv(integrator.opts.qmin),integrator.q11/integrator.opts.gamma)
end


text:
## AbstractSparseTensor type and functions/methods

## Overloads for specializing outputs

@testset "show(io, A)" begin
    @testset "nstored=$nstored, N=$N, Ti=$Ti, Tv=$Tv" for nstored in 0:3, N in 1:3, Ti in [Int, UInt8], Tv in [Float64, BigFloat, UInt8]
        dims = (10, 3, 2)[1:N]
        inds = (Ti[2, 1, 4], Ti[1, 3, 2], Ti[1, 2, 1])[1:N]
        inds = tuple.(inds...)
        vals = Tv[0, 1, 10]

        # Take subset of entries
        inds = inds[1:nstored]
        vals = vals[1:nstored]

        # SparseTensorCOO
        perm = sortperm(inds; by = CartesianIndex)
        sinds, svals = inds[perm], vals[perm]
        C = SparseTensorCOO(dims, inds, vals)
        @test sprint(show, C) == "SparseTensorCOO{$Tv, $Ti, $N}($dims, $sinds, $svals)"

        # SparseTensorDOK
        dict = Dict(inds .=> vals)
        D = SparseTensorDOK(dims, dict)
        @test sprint(show, D) == "SparseTensorDOK{$Tv, $Ti, $N}($dims, $dict)"
    end
end

@testset "show(io, ::MIME\"text/plain\", A)" begin
    @testset "nstored=$nstored, N=$N, Ti=$Ti, Tv=$Tv" for nstored in 0:3, N in 1:3, Ti in [Int, UInt8], Tv in [Float64, BigFloat, UInt8]
        dims = (10, 3, 2)[1:N]
        inds = (Ti[2, 1, 4], Ti[1, 3, 2], Ti[1, 2, 1])[1:N]
        inds = tuple.(inds...)
        vals = Tv[0, 1, 10]

        # Take subset of entries
        inds = inds[1:nstored]
        vals = vals[1:nstored]

        # SparseTensorCOO
        perm = sortperm(inds; by = CartesianIndex)
        sinds, svals = inds[perm], vals[perm]
        C = SparseTensorCOO(dims, inds, vals)
        entrystrs = map(sinds, svals) do ind, val
            indstr = join(lpad.(Int.(ind), ndigits.(dims)), ", ")
            return "  [$indstr]  =  $val"
        end
        showstr = length(vals) == 0 ? summary(C) : "$(summary(C)):\n$(join(entrystrs,'\n'))"
        @test sprint(show, MIME("text/plain"), C) == showstr

        # SparseTensorDOK
        dict = Dict(inds .=> vals)
        D = SparseTensorDOK(dims, dict)
        entrystrs = map(collect(dict)) do (ind, val)
            indstr = join(lpad.(Int.(ind), ndigits.(dims)), ", ")
            return "  [$indstr]  =  $val"
        end
        showstr = length(vals) == 0 ? summary(D) : "$(summary(D)):\n$(join(entrystrs,'\n'))"
        @test sprint(show, MIME("text/plain"), D) == showstr
    end

    @testset "displayheight=$displayheight" for displayheight in 0:11
        iocontext = IOContext(IOBuffer(), :displaysize => (displayheight, 80), :limit => true)
        dims = (5, 10, 2)
        inds = tuple.(
            UInt8[4, 1, 2, 5, 5, 3, 2],
            UInt8[1, 2, 3, 4, 10, 7, 6],
            UInt8[1, 1, 2, 1, 2, 1, 2],
        )
        vals = Float32[0.0, 0.5, 0.4, 0.2, 0.3, 0.8, 0.9]

        # SparseTensorCOO
        showstr = Dict(
            0 => "5×10×2 SparseTensorCOO{Float32, UInt8, 3} with 7 stored entries: \u2026",
            1 => "5×10×2 SparseTensorCOO{Float32, UInt8, 3} with 7 stored entries: \u2026",
            2 => "5×10×2 SparseTensorCOO{Float32, UInt8, 3} with 7 stored entries: \u2026",
            3 => "5×10×2 SparseTensorCOO{Float32, UInt8, 3} with 7 stored entries: \u2026",
            4 => "5×10×2 SparseTensorCOO{Float32, UInt8, 3} with 7 stored entries: \u2026",
            5 => """
            5×10×2 SparseTensorCOO{Float32, UInt8, 3} with 7 stored entries:
             \u22ee""",
            6 => """
            5×10×2 SparseTensorCOO{Float32, UInt8, 3} with 7 stored entries:
              [4,  1, 1]  =  0.0
                          \u22ee""",
            7 => """
            5×10×2 SparseTensorCOO{Float32, UInt8, 3} with 7 stored entries:
              [4,  1, 1]  =  0.0
                          \u22ee
              [5, 10, 2]  =  0.3""",
            8 => """
            5×10×2 SparseTensorCOO{Float32, UInt8, 3} with 7 stored entries:
              [4,  1, 1]  =  0.0
              [1,  2, 1]  =  0.5
                          \u22ee
              [5, 10, 2]  =  0.3""",
            9 => """
            5×10×2 SparseTensorCOO{Float32, UInt8, 3} with 7 stored entries:
              [4,  1, 1]  =  0.0
              [1,  2, 1]  =  0.5
                          \u22ee
              [2,  6, 2]  =  0.9
              [5, 10, 2]  =  0.3""",
            10 => """
            5×10×2 SparseTensorCOO{Float32, UInt8, 3} with 7 stored entries:
              [4,  1, 1]  =  0.0
              [1,  2, 1]  =  0.5
              [5,  4, 1]  =  0.2
                          \u22ee
              [2,  6, 2]  =  0.9
              [5, 10, 2]  =  0.3""",
            11 => """
            5×10×2 SparseTensorCOO{Float32, UInt8, 3} with 7 stored entries:
              [4,  1, 1]  =  0.0
              [1,  2, 1]  =  0.5
              [5,  4, 1]  =  0.2
              [3,  7, 1]  =  0.8
              [2,  3, 2]  =  0.4
              [2,  6, 2]  =  0.9
              [5, 10, 2]  =  0.3""",
        )[displayheight]
        C = SparseTensorCOO(dims, inds, vals)
        @test sprint(show, MIME("text/plain"), C; context = iocontext) == showstr

        # SparseTensorDOK
        showstr = Dict(
            0 => "5×10×2 SparseTensorDOK{Float32, UInt8, 3} with 7 stored entries: \u2026",
            1 => "5×10×2 SparseTensorDOK{Float32, UInt8, 3} with 7 stored entries: \u2026",
            2 => "5×10×2 SparseTensorDOK{Float32, UInt8, 3} with 7 stored entries: \u2026",
            3 => "5×10×2 SparseTensorDOK{Float32, UInt8, 3} with 7 stored entries: \u2026",
            4 => "5×10×2 SparseTensorDOK{Float32, UInt8, 3} with 7 stored entries: \u2026",
            5 => """
            5×10×2 SparseTensorDOK{Float32, UInt8, 3} with 7 stored entries:
             \u22ee""",
            6 => """
            5×10×2 SparseTensorDOK{Float32, UInt8, 3} with 7 stored entries:
              [1,  2, 1]  =  0.5
                          \u22ee""",
            7 => """
            5×10×2 SparseTensorDOK{Float32, UInt8, 3} with 7 stored entries:
              [1,  2, 1]  =  0.5
                          \u22ee
              [5,  4, 1]  =  0.2""",
            8 => """
            5×10×2 SparseTensorDOK{Float32, UInt8, 3} with 7 stored entries:
              [1,  2, 1]  =  0.5
              [2,  3, 2]  =  0.4
                          \u22ee
              [5,  4, 1]  =  0.2""",
            9 => """
            5×10×2 SparseTensorDOK{Float32, UInt8, 3} with 7 stored entries:
              [1,  2, 1]  =  0.5
              [2,  3, 2]  =  0.4
                          \u22ee
              [5, 10, 2]  =  0.3
              [5,  4, 1]  =  0.2""",
            10 => """
            5×10×2 SparseTensorDOK{Float32, UInt8, 3} with 7 stored entries:
              [1,  2, 1]  =  0.5
              [2,  3, 2]  =  0.4
              [3,  7, 1]  =  0.8
                          \u22ee
              [5, 10, 2]  =  0.3
              [5,  4, 1]  =  0.2""",
            11 => """
            5×10×2 SparseTensorDOK{Float32, UInt8, 3} with 7 stored entries:
              [1,  2, 1]  =  0.5
              [2,  3, 2]  =  0.4
              [3,  7, 1]  =  0.8
              [2,  6, 2]  =  0.9
              [4,  1, 1]  =  0.0
              [5, 10, 2]  =  0.3
              [5,  4, 1]  =  0.2""",
        )[displayheight]
        D = SparseTensorDOK(dims, Dict(inds .=> vals))
        @test sprint(show, MIME("text/plain"), D; context = iocontext) == showstr
    end
end

@testset "summary(io, A)" begin
    @testset "nstored=$nstored, N=$N, Ti=$Ti, Tv=$Tv" for nstored in 0:3, N in 1:3, Ti in [Int, UInt8], Tv in [Float64, BigFloat, UInt8]
        dims = (10, 3, 2)[1:N]
        inds = (Ti[2, 1, 4], Ti[1, 3, 2], Ti[1, 2, 1])[1:N]
        inds = tuple.(inds...)
        vals = Tv[0, 1, 10]

        # Take subset of entries
        inds = inds[1:nstored]
        vals = vals[1:nstored]

        # References
        dimstr = ["10-element", "10×3", "10×3×2"]
        valstr = "with $(length(vals)) stored " * (length(vals) == 1 ? "entry" : "entries")

        # SparseTensorCOO
        C = SparseTensorCOO(dims, inds, vals)
        @test sprint(summary, C) == "$(dimstr[N]) SparseTensorCOO{$Tv, $Ti, $N} $valstr"

        # SparseTensorDOK
        D = SparseTensorDOK(dims, Dict(inds .=> vals))
        @test sprint(summary, D) == "$(dimstr[N]) SparseTensorDOK{$Tv, $Ti, $N} $valstr"
    end
end

## Overloads for improving efficiency

@testset "findall(A)" begin
    @testset "N=$N, Ti=$Ti" for N in 1:3, Ti in [Int, UInt8]
        dims = (10, 3, 2)[1:N]
        inds = (Ti[2, 5, 4], Ti[1, 3, 2], Ti[1, 2, 1])[1:N]
        inds = tuple.(inds...)
        vals = Bool[1, 0, 1]

        C = SparseTensorCOO(dims, inds, vals)
        D = SparseTensorDOK(dims, Dict(inds .=> vals))
        A = collect(C)

        @test typeof(findall(C)) == typeof(findall(D)) == typeof(findall(A))
        @test findall(C) == findall(D) == findall(A)
    end
end

@testset "findall(f, A)" begin
    @testset "N=$N, Ti=$Ti, Tv=$Tv" for N in 1:3, Ti in [Int, UInt8], Tv in [Float64, BigFloat, UInt8]
        dims = (10, 3, 2)[1:N]
        inds = (Ti[2, 5, 4], Ti[1, 3, 2], Ti[1, 2, 1])[1:N]
        inds = tuple.(inds...)
        vals = Tv[0, 1, 10]

        # Form tensors
        C = SparseTensorCOO(dims, inds, vals)
        D = SparseTensorDOK(dims, Dict(inds .=> vals))
        A = collect(C)

        # findall(!iszero, A)
        @test typeof(findall(!iszero, C)) == typeof(findall(!iszero, D)) == typeof(findall(!iszero, A))
        @test findall(!iszero, C) == findall(!iszero, D) == findall(!iszero, A)

        # findall(in(...), A) - f(0) == false
        @test typeof(findall(in(Tv[1, 2]), C)) == typeof(findall(in(Tv[1, 2]), D)) == typeof(findall(in(Tv[1, 2]), A))
        @test findall(in(Tv[1, 2]), C) == findall(in(Tv[1, 2]), D) == findall(in(Tv[1, 2]), A)

        # findall(in(...), A) - f(0) == true
        @test typeof(findall(in(Tv[0, 1]), C)) == typeof(findall(in(Tv[0, 1]), D)) == typeof(findall(in(Tv[0, 1]), A))
        @test findall(in(Tv[0, 1]), C) == findall(in(Tv[0, 1]), D) == findall(in(Tv[0, 1]), A)
    end
end

## Generic methods

@testset "indtype" begin
    # indtype(T::Type{<:AbstractSparseTensor})
    @testset "T::Type{<:AbstractSparseTensor}" begin
        AT = AbstractSparseTensor
        @test indtype(AT) == Integer
        for T in subtypes(AT)
            @test indtype(T) == Integer
        end
    end
    @testset "T::Type{<:AbstractSparseTensor{$Tv}}" for Tv in [Float64, BigFloat, Int8]
        AT = AbstractSparseTensor{Tv}
        @test indtype(AT) == Integer
        for T in subtypes(AT)
            @test indtype(T) == Integer
        end
    end
    @testset "T::Type{<:AbstractSparseTensor{$Tv,$Ti}}" for Tv in [Float64, BigFloat, Int8], Ti in [Int, UInt8]
        AT = AbstractSparseTensor{Tv,Ti}
        @test indtype(AT) == Ti
        for T in subtypes(AT)
            @test indtype(T) == Ti
        end
    end
    @testset "T::Type{<:AbstractSparseTensor{$Tv,$Ti,$N}}" for Tv in [Float64, BigFloat, Int8], Ti in [Int, UInt8], N in 1:3
        AT = AbstractSparseTensor{Tv,Ti,N}
        @test indtype(AT) == Ti
        for T in subtypes(AT)
            @test indtype(T) == Ti
        end
    end

    # indtype(A::AbstractSparseTensor)
    @testset "A::AbstractSparseTensor{$Tv, $Ti, $N}" for Tv in [Float64, BigFloat, Int8], Ti in [Int, UInt8], N in 1:3
        dims = (5, 3, 2)[1:N]
        inds = (Ti[2, 1, 4], Ti[1, 3, 2], Ti[1, 2, 1])[1:N]
        vals = Tv[1, 0, 10]

        @test indtype(SparseTensorCOO(dims, tuple.(inds...), vals)) == Ti
        @test indtype(SparseTensorDOK(dims, Dict(tuple.(inds...) .=> vals))) == Ti
    end
end


text:
using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBoundary
using Oceananigans.TurbulenceClosures: VerticallyImplicitTimeDiscretization

@testset "Immersed boundaries with hydrostatic free surface models" begin
    @info "Testing immersed boundaries with hydrostatic free surface models..."

    for arch in archs
        underlying_grid = RectilinearGrid(arch, size=(8, 8, 8), x = (-5, 5), y = (-5, 5), z = (0, 2))

        bump(x, y, z) = z < exp(-x^2 - y^2)

        grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBoundary(bump))

        for closure in (IsotropicDiffusivity(ν=1, κ=0.5),
                        IsotropicDiffusivity(ν=1, κ=0.5, time_discretization=VerticallyImplicitTimeDiscretization()))

            model = HydrostaticFreeSurfaceModel(grid = grid,
                                                tracers = :b,
                                                buoyancy = BuoyancyTracer(),
                                                closure = closure)

            u = model.velocities.u
            b = model.tracers.b

            # Linear stratification
            set!(model, u = 1, b = (x, y, z) -> 4 * z)

            # Inside the bump
            @test b[4, 4, 2] == 0
            @test u[4, 4, 2] == 0

            simulation = Simulation(model, Δt = 1e-3, stop_iteration=2)

            run!(simulation)

            # Inside the bump
            @test b[4, 4, 2] == 0
            @test u[4, 4, 2] == 0
        end
    end
end


text:
###
### Sampler states
###

struct Emcee{space, E<:AMH.Ensemble} <: InferenceAlgorithm
    ensemble::E
end

function Emcee(n_walkers::Int, stretch_length=2.0)
    # Note that the proposal distribution here is just a Normal(0,1)
    # because we do not need AdvancedMH to know the proposal for
    # ensemble sampling.
    prop = AMH.StretchProposal(nothing, stretch_length)
    ensemble = AMH.Ensemble(n_walkers, prop)
    return Emcee{(), typeof(ensemble)}(ensemble)
end

function Sampler(
    alg::Emcee,
    model::Model,
    s::Selector=Selector()
)
    # Set up info dict.
    info = Dict{Symbol, Any}()

    # Set up state struct.
    state = SamplerState(VarInfo(model))

    # Generate a sampler.
    return Sampler(alg, info, s, state)
end

function AbstractMCMC.sample_init!(
    rng::AbstractRNG,
    model::Model,
    spl::Sampler{<:Emcee},
    N::Integer;
    verbose::Bool=true,
    resume_from=nothing,
    kwargs...
)
    # Resume the sampler.
    set_resume!(spl; resume_from=resume_from, kwargs...)

    # Get `init_theta`
    initialize_parameters!(spl; verbose=verbose, kwargs...)

    # Link everything before sampling.
    link!(spl.state.vi, spl)
end

function AbstractMCMC.sample_end!(
    rng::AbstractRNG,
    model::Model,
    spl::Sampler{<:Emcee},
    N::Integer,
    transitions;
    kwargs...
)
    # Invlink everything after sampling.
    invlink!(spl.state.vi, spl)
end

function _typed_draw(model, vi::VarInfo)
    spl = SampleFromPrior()
    empty!(vi)
    model(vi, spl)
    DynamicPPL.link!(vi, spl)
    params = vi[spl]
    return AMH.Transition(params, getlogp(vi))
end

function AbstractMCMC.step!(
    rng::AbstractRNG,
    model::Model,
    spl::Sampler{<:Emcee},
    N::Integer,
    transition::Nothing;
    kwargs...
)
    # Generate a log joint function.
    densitymodel = AMH.DensityModel(gen_logπ(spl.state.vi, DynamicPPL.SampleFromPrior(), model))

    # Make the first transition.
    walkers = map(x -> _typed_draw(model, spl.state.vi), 1:spl.alg.ensemble.n_walkers)

    return walkers
end

function AbstractMCMC.step!(
    rng::AbstractRNG,
    model::Model,
    spl::Sampler{<:Emcee},
    N::Integer,
    transition;
    kwargs...
)
    # Generate a log joint function.
    densitymodel = AMH.DensityModel(gen_logπ(spl.state.vi, DynamicPPL.SampleFromPrior(), model))

    # Make the first transition.
    new_transitions = AbstractMCMC.step!(rng, densitymodel, spl.alg.ensemble, 1, transition)
    return new_transitions
end

function transform_transition(spl::Sampler{<:Emcee}, ts::Vector{<:Vector}, w::Int, i::Int; linked=true)
    trans = ts[i][w]
    linked && DynamicPPL.link!(spl.state.vi, spl)
    spl.state.vi[spl] = trans.params
    linked && DynamicPPL.invlink!(spl.state.vi, spl)
    setlogp!(spl.state.vi, trans.lp)

    return Transition(spl)
end

function AbstractMCMC.bundle_samples(
    rng::AbstractRNG,
    model::AbstractModel,
    spl::Sampler{<:Emcee},
    N::Integer,
    ts::Vector{<:Vector},
    chain_type::Type{MCMCChains.Chains};
    save_state = false,
    kwargs...
)
    # Transform the transitions by linking them to the constrained space.
    ts_transform = [map(i -> transform_transition(spl, ts, w, i), 1:N) for w in 1:spl.alg.ensemble.n_walkers]

     # Convert transitions to array format.
    # Also retrieve the variable names.
    params_vec = map(_params_to_array, ts_transform)

    # Extract names and values separately.
    nms = params_vec[1][1]
    vals_vec = [p[2] for p in params_vec]

    # Get the values of the extra parameters in each transition.
    extra_vec = map(get_transition_extras, ts_transform)

    # Get the extra parameter names & values.
    extra_params = extra_vec[1][1]
    extra_values_vec = [e[2] for e in extra_vec]

    # Extract names & construct param array.
    nms = [nms; extra_params]
    parray = map(x -> hcat(x[1], x[2]), zip(vals_vec, extra_values_vec))
    parray = cat(parray..., dims=3)

    # Get the average or final log evidence, if it exists.
    le = getlogevidence(spl)

    # Set up the info tuple.
    if save_state
        info = (range = rng, model = model, spl = spl)
    else
        info = NamedTuple()
    end

    # Concretize the array before giving it to MCMCChains.
    parray = MCMCChains.concretize(parray)

    # Chain construction.
    return MCMCChains.Chains(
        parray,
        nms,
        deepcopy(TURING_INTERNAL_VARS);
        evidence=le,
        info=info,
    ) |> sort
end


text:
using SmartTransitionSim

#Creating MapData object
mapfile = "reno_east3.osm"
datapath = "./example/";
RoadSet = 5
map_data = get_map_data(datapath, mapfile, use_cache=false; road_levels = Set(1:RoadSet));

#Defining starting and ending area
Start = [Rect((39.50,-119.70),(39.55,-119.74))]
End = [Rect((39.50,-119.80),(39.55,-119.76))]

#Dictionary for memorizing used k-shortest paths
KDict = Dict{Tuple{Int,Int},Array{Vector{Int}}}()

#Declaring simulation input parameters
density_factor = 5.0
smart_perc = 0.5
agents_num = 100
update_period = 150
T = 0.1
k = 3
AvgDelay = 1.0

#Generating agents
Agents = generate_agents(map_data, agents_num, Start, End, smart_perc, k, T , AvgDelay, KDict)
#Running base simulation - no V2I system
BaseOutput = simulation_run(:base, map_data, Agents)
SmartOutput = simulation_run(:smart, map_data, Agents, KDict, update_period, T, k, density_factor)

simulation_statistics = gather_statistics(getfield.(Agents,:smart),
                                    BaseOutput.TravelTimes,
                                    SmartOutput.TravelTimes)
#Gathered simulation statistics - percentage time reduction
println(simulation_statistics)