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

module DFT

# DFT plan where the inputs are an array of eltype T
abstract type Plan{T} end

import Base: show, summary, size, ndims, length, eltype,
             *, A_mul_B!, inv, \, A_ldiv_B!

eltype(::Type{Plan{T}}) where {T} = T

# size(p) should return the size of the input array for p
size(p::Plan, d) = size(p)[d]
ndims(p::Plan) = length(size(p))
length(p::Plan) = prod(size(p))::Int

##############################################################################
export fft, ifft, bfft, fft!, ifft!, bfft!,
       plan_fft, plan_ifft, plan_bfft, plan_fft!, plan_ifft!, plan_bfft!,
       rfft, irfft, brfft, plan_rfft, plan_irfft, plan_brfft

const FFTWFloat = Union{Float32,Float64}
fftwfloat(x) = _fftwfloat(float(x))
_fftwfloat(::Type{T}) where {T<:FFTWFloat} = T
_fftwfloat(::Type{Float16}) = Float32
_fftwfloat(::Type{Complex{T}}) where {T} = Complex{_fftwfloat(T)}
_fftwfloat(::Type{T}) where {T} = error("type $T not supported")
_fftwfloat(x::T) where {T} = _fftwfloat(T)(x)

complexfloat(x::StridedArray{Complex{<:FFTWFloat}}) = x
realfloat(x::StridedArray{<:FFTWFloat}) = x

# return an Array, rather than similar(x), to avoid an extra copy for FFTW
# (which only works on StridedArray types).
complexfloat(x::AbstractArray{T}) where {T<:Complex} = copy1(typeof(fftwfloat(zero(T))), x)
complexfloat(x::AbstractArray{T}) where {T<:Real} = copy1(typeof(complex(fftwfloat(zero(T)))), x)

realfloat(x::AbstractArray{T}) where {T<:Real} = copy1(typeof(fftwfloat(zero(T))), x)

# copy to a 1-based array, using circular permutation
function copy1(::Type{T}, x) where T
    y = Array{T}(map(length, indices(x)))
    Base.circcopy!(y, x)
end

to1(x::AbstractArray) = _to1(indices(x), x)
_to1(::Tuple{Base.OneTo,Vararg{Base.OneTo}}, x) = x
_to1(::Tuple, x) = copy1(eltype(x), x)

# implementations only need to provide plan_X(x, region)
# for X in (:fft, :bfft, ...):
for f in (:fft, :bfft, :ifft, :fft!, :bfft!, :ifft!, :rfft)
    pf = Symbol("plan_", f)
    @eval begin
        $f(x::AbstractArray) = (y = to1(x); $pf(y) * y)
        $f(x::AbstractArray, region) = (y = to1(x); $pf(y, region) * y)
        $pf(x::AbstractArray; kws...) = (y = to1(x); $pf(y, 1:ndims(y); kws...))
    end
end

"""
    plan_ifft(A [, dims]; flags=FFTW.ESTIMATE;  timelimit=Inf)

Same as [`plan_fft`](@ref), but produces a plan that performs inverse transforms
[`ifft`](@ref).
"""
plan_ifft

"""
    plan_ifft!(A [, dims]; flags=FFTW.ESTIMATE;  timelimit=Inf)

Same as [`plan_ifft`](@ref), but operates in-place on `A`.
"""
plan_ifft!

"""
    plan_bfft!(A [, dims]; flags=FFTW.ESTIMATE;  timelimit=Inf)

Same as [`plan_bfft`](@ref), but operates in-place on `A`.
"""
plan_bfft!

"""
    plan_bfft(A [, dims]; flags=FFTW.ESTIMATE;  timelimit=Inf)

Same as [`plan_fft`](@ref), but produces a plan that performs an unnormalized
backwards transform [`bfft`](@ref).
"""
plan_bfft

"""
    plan_fft(A [, dims]; flags=FFTW.ESTIMATE;  timelimit=Inf)

Pre-plan an optimized FFT along given dimensions (`dims`) of arrays matching the shape and
type of `A`.  (The first two arguments have the same meaning as for [`fft`](@ref).)
Returns an object `P` which represents the linear operator computed by the FFT, and which
contains all of the information needed to compute `fft(A, dims)` quickly.

To apply `P` to an array `A`, use `P * A`; in general, the syntax for applying plans is much
like that of matrices.  (A plan can only be applied to arrays of the same size as the `A`
for which the plan was created.)  You can also apply a plan with a preallocated output array `Â`
by calling `A_mul_B!(Â, plan, A)`.  (For `A_mul_B!`, however, the input array `A` must
be a complex floating-point array like the output `Â`.) You can compute the inverse-transform plan by `inv(P)`
and apply the inverse plan with `P \\ Â` (the inverse plan is cached and reused for
subsequent calls to `inv` or `\\`), and apply the inverse plan to a pre-allocated output
array `A` with `A_ldiv_B!(A, P, Â)`.

The `flags` argument is a bitwise-or of FFTW planner flags, defaulting to `FFTW.ESTIMATE`.
e.g. passing `FFTW.MEASURE` or `FFTW.PATIENT` will instead spend several seconds (or more)
benchmarking different possible FFT algorithms and picking the fastest one; see the FFTW
manual for more information on planner flags.  The optional `timelimit` argument specifies a
rough upper bound on the allowed planning time, in seconds. Passing `FFTW.MEASURE` or
`FFTW.PATIENT` may cause the input array `A` to be overwritten with zeros during plan
creation.

[`plan_fft!`](@ref) is the same as [`plan_fft`](@ref) but creates a
plan that operates in-place on its argument (which must be an array of complex
floating-point numbers). [`plan_ifft`](@ref) and so on are similar but produce
plans that perform the equivalent of the inverse transforms [`ifft`](@ref) and so on.
"""
plan_fft

"""
    plan_fft!(A [, dims]; flags=FFTW.ESTIMATE;  timelimit=Inf)

Same as [`plan_fft`](@ref), but operates in-place on `A`.
"""
plan_fft!

"""
    rfft(A [, dims])

Multidimensional FFT of a real array `A`, exploiting the fact that the transform has
conjugate symmetry in order to save roughly half the computational time and storage costs
compared with [`fft`](@ref). If `A` has size `(n_1, ..., n_d)`, the result has size
`(div(n_1,2)+1, ..., n_d)`.

The optional `dims` argument specifies an iterable subset of one or more dimensions of `A`
to transform, similar to [`fft`](@ref). Instead of (roughly) halving the first
dimension of `A` in the result, the `dims[1]` dimension is (roughly) halved in the same way.
"""
rfft

"""
    ifft!(A [, dims])

Same as [`ifft`](@ref), but operates in-place on `A`.
"""
ifft!

"""
    ifft(A [, dims])

Multidimensional inverse FFT.

A one-dimensional inverse FFT computes

```math
\\operatorname{IDFT}(A)[k] = \\frac{1}{\\operatorname{length}(A)}
\\sum_{n=1}^{\\operatorname{length}(A)} \\exp\\left(+i\\frac{2\\pi (n-1)(k-1)}
{\\operatorname{length}(A)} \\right) A[n].
```

A multidimensional inverse FFT simply performs this operation along each transformed dimension of `A`.
"""
ifft

"""
    fft!(A [, dims])

Same as [`fft`](@ref), but operates in-place on `A`, which must be an array of
complex floating-point numbers.
"""
fft!

"""
    bfft(A [, dims])

Similar to [`ifft`](@ref), but computes an unnormalized inverse (backward)
transform, which must be divided by the product of the sizes of the transformed dimensions
in order to obtain the inverse. (This is slightly more efficient than [`ifft`](@ref)
because it omits a scaling step, which in some applications can be combined with other
computational steps elsewhere.)

```math
\\operatorname{BDFT}(A)[k] = \\operatorname{length}(A) \\operatorname{IDFT}(A)[k]
```
"""
bfft

"""
    bfft!(A [, dims])

Same as [`bfft`](@ref), but operates in-place on `A`.
"""
bfft!

# promote to a complex floating-point type (out-of-place only),
# so implementations only need Complex{Float} methods
for f in (:fft, :bfft, :ifft)
    pf = Symbol("plan_", f)
    @eval begin
        $f(x::AbstractArray{<:Real}, region=1:ndims(x)) = $f(complexfloat(x), region)
        $pf(x::AbstractArray{<:Real}, region; kws...) = $pf(complexfloat(x), region; kws...)
        $f(x::AbstractArray{<:Complex{<:Union{Integer,Rational}}}, region=1:ndims(x)) = $f(complexfloat(x), region)
        $pf(x::AbstractArray{<:Complex{<:Union{Integer,Rational}}}, region; kws...) = $pf(complexfloat(x), region; kws...)
    end
end
rfft(x::AbstractArray{<:Union{Integer,Rational}}, region=1:ndims(x)) = rfft(realfloat(x), region)
plan_rfft(x::AbstractArray, region; kws...) = plan_rfft(realfloat(x), region; kws...)

# only require implementation to provide *(::Plan{T}, ::Array{T})
*(p::Plan{T}, x::AbstractArray) where {T} = p * copy1(T, x)

# Implementations should also implement A_mul_B!(Y, plan, X) so as to support
# pre-allocated output arrays.  We don't define * in terms of A_mul_B!
# generically here, however, because of subtleties for in-place and rfft plans.

##############################################################################
# To support inv, \, and A_ldiv_B!(y, p, x), we require Plan subtypes
# to have a pinv::Plan field, which caches the inverse plan, and which
# should be initially undefined.  They should also implement
# plan_inv(p) to construct the inverse of a plan p.

# hack from @simonster (in #6193) to compute the return type of plan_inv
# without actually calling it or even constructing the empty arrays.
_pinv_type(p::Plan) = typeof([plan_inv(x) for x in typeof(p)[]])
pinv_type(p::Plan) = eltype(_pinv_type(p))

inv(p::Plan) =
    isdefined(p, :pinv) ? p.pinv::pinv_type(p) : (p.pinv = plan_inv(p))
\(p::Plan, x::AbstractArray) = inv(p) * x
A_ldiv_B!(y::AbstractArray, p::Plan, x::AbstractArray) = A_mul_B!(y, inv(p), x)

##############################################################################
# implementations only need to provide the unnormalized backwards FFT,
# similar to FFTW, and we do the scaling generically to get the ifft:

mutable struct ScaledPlan{T,P,N} <: Plan{T}
    p::P
    scale::N # not T, to avoid unnecessary promotion to Complex
    pinv::Plan
    ScaledPlan{T,P,N}(p, scale) where {T,P,N} = new(p, scale)
end
ScaledPlan{T}(p::P, scale::N) where {T,P,N} = ScaledPlan{T,P,N}(p, scale)
ScaledPlan(p::Plan{T}, scale::Number) where {T} = ScaledPlan{T}(p, scale)
ScaledPlan(p::ScaledPlan, α::Number) = ScaledPlan(p.p, p.scale * α)

size(p::ScaledPlan) = size(p.p)

show(io::IO, p::ScaledPlan) = print(io, p.scale, " * ", p.p)
summary(p::ScaledPlan) = string(p.scale, " * ", summary(p.p))

*(p::ScaledPlan, x::AbstractArray) = scale!(p.p * x, p.scale)

*(α::Number, p::Plan) = ScaledPlan(p, α)
*(p::Plan, α::Number) = ScaledPlan(p, α)
*(I::UniformScaling, p::ScaledPlan) = ScaledPlan(p, I.λ)
*(p::ScaledPlan, I::UniformScaling) = ScaledPlan(p, I.λ)
*(I::UniformScaling, p::Plan) = ScaledPlan(p, I.λ)
*(p::Plan, I::UniformScaling) = ScaledPlan(p, I.λ)

# Normalization for ifft, given unscaled bfft, is 1/prod(dimensions)
normalization(T, sz, region) = one(T) / Int(prod([sz...][[region...]]))
normalization(X, region) = normalization(real(eltype(X)), size(X), region)

plan_ifft(x::AbstractArray, region; kws...) =
    ScaledPlan(plan_bfft(x, region; kws...), normalization(x, region))
plan_ifft!(x::AbstractArray, region; kws...) =
    ScaledPlan(plan_bfft!(x, region; kws...), normalization(x, region))

plan_inv(p::ScaledPlan) = ScaledPlan(plan_inv(p.p), inv(p.scale))

A_mul_B!(y::AbstractArray, p::ScaledPlan, x::AbstractArray) =
    scale!(p.scale, A_mul_B!(y, p.p, x))

##############################################################################
# Real-input DFTs are annoying because the output has a different size
# than the input if we want to gain the full factor-of-two(ish) savings
# For backward real-data transforms, we must specify the original length
# of the first dimension, since there is no reliable way to detect this
# from the data (we can't detect whether the dimension was originally even
# or odd).

for f in (:brfft, :irfft)
    pf = Symbol("plan_", f)
    @eval begin
        $f(x::AbstractArray, d::Integer) = $pf(x, d) * x
        $f(x::AbstractArray, d::Integer, region) = $pf(x, d, region) * x
        $pf(x::AbstractArray, d::Integer;kws...) = $pf(x, d, 1:ndims(x);kws...)
    end
end

for f in (:brfft, :irfft)
    @eval begin
        $f(x::AbstractArray{<:Real}, d::Integer, region=1:ndims(x)) = $f(complexfloat(x), d, region)
        $f(x::AbstractArray{<:Complex{<:Union{Integer,Rational}}}, d::Integer, region=1:ndims(x)) = $f(complexfloat(x), d, region)
    end
end

"""
    irfft(A, d [, dims])

Inverse of [`rfft`](@ref): for a complex array `A`, gives the corresponding real
array whose FFT yields `A` in the first half. As for [`rfft`](@ref), `dims` is an
optional subset of dimensions to transform, defaulting to `1:ndims(A)`.

`d` is the length of the transformed real array along the `dims[1]` dimension, which must
satisfy `div(d,2)+1 == size(A,dims[1])`. (This parameter cannot be inferred from `size(A)`
since both `2*size(A,dims[1])-2` as well as `2*size(A,dims[1])-1` are valid sizes for the
transformed real array.)
"""
irfft

"""
    brfft(A, d [, dims])

Similar to [`irfft`](@ref) but computes an unnormalized inverse transform (similar
to [`bfft`](@ref)), which must be divided by the product of the sizes of the
transformed dimensions (of the real output array) in order to obtain the inverse transform.
"""
brfft

function rfft_output_size(x::AbstractArray, region)
    d1 = first(region)
    osize = [size(x)...]
    osize[d1] = osize[d1]>>1 + 1
    return osize
end

function brfft_output_size(x::AbstractArray, d::Integer, region)
    d1 = first(region)
    osize = [size(x)...]
    @assert osize[d1] == d>>1 + 1
    osize[d1] = d
    return osize
end

plan_irfft(x::AbstractArray{Complex{T}}, d::Integer, region; kws...) where {T} =
    ScaledPlan(plan_brfft(x, d, region; kws...),
               normalization(T, brfft_output_size(x, d, region), region))

"""
    plan_irfft(A, d [, dims]; flags=FFTW.ESTIMATE;  timelimit=Inf)

Pre-plan an optimized inverse real-input FFT, similar to [`plan_rfft`](@ref)
except for [`irfft`](@ref) and [`brfft`](@ref), respectively. The first
three arguments have the same meaning as for [`irfft`](@ref).
"""
plan_irfft

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

export fftshift, ifftshift

fftshift(x) = circshift(x, div.([size(x)...],2))

"""
    fftshift(x)

Swap the first and second halves of each dimension of `x`.
"""
fftshift(x)

function fftshift(x,dim)
    s = zeros(Int,ndims(x))
    for i in dim
        s[i] = div(size(x,i),2)
    end
    circshift(x, s)
end

"""
    fftshift(x,dim)

Swap the first and second halves of the given dimension or iterable of dimensions of array `x`.
"""
fftshift(x,dim)

ifftshift(x) = circshift(x, div.([size(x)...],-2))

"""
    ifftshift(x, [dim])

Undoes the effect of `fftshift`.
"""
ifftshift

function ifftshift(x,dim)
    s = zeros(Int,ndims(x))
    for i in dim
        s[i] = -div(size(x,i),2)
    end
    circshift(x, s)
end

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

# FFTW module (may move to an external package at some point):
"""
    fft(A [, dims])

Performs a multidimensional FFT of the array `A`. The optional `dims` argument specifies an
iterable subset of dimensions (e.g. an integer, range, tuple, or array) to transform along.
Most efficient if the size of `A` along the transformed dimensions is a product of small
primes; see `nextprod()`. See also `plan_fft()` for even greater efficiency.

A one-dimensional FFT computes the one-dimensional discrete Fourier transform (DFT) as
defined by

```math
\\operatorname{DFT}(A)[k] =
  \\sum_{n=1}^{\\operatorname{length}(A)}
  \\exp\\left(-i\\frac{2\\pi
  (n-1)(k-1)}{\\operatorname{length}(A)} \\right) A[n].
```

A multidimensional FFT simply performs this operation along each transformed dimension of `A`.

!!! note
    * Julia starts FFTW up with 1 thread by default. Higher performance is usually possible by
      increasing number of threads. Use `FFTW.set_num_threads(Sys.CPU_CORES)` to use as many
      threads as cores on your system.

    * This performs a multidimensional FFT by default. FFT libraries in other languages such as
      Python and Octave perform a one-dimensional FFT along the first non-singleton dimension
      of the array. This is worth noting while performing comparisons. For more details,
      refer to the [Noteworthy Differences from other Languages](@ref)
      section of the manual.
"""
fft

"""
    plan_rfft(A [, dims]; flags=FFTW.ESTIMATE;  timelimit=Inf)

Pre-plan an optimized real-input FFT, similar to [`plan_fft`](@ref) except for
[`rfft`](@ref) instead of [`fft`](@ref). The first two arguments, and the
size of the transformed result, are the same as for [`rfft`](@ref).
"""
plan_rfft

"""
    plan_brfft(A, d [, dims]; flags=FFTW.ESTIMATE;  timelimit=Inf)

Pre-plan an optimized real-input unnormalized transform, similar to
[`plan_rfft`](@ref) except for [`brfft`](@ref) instead of
[`rfft`](@ref). The first two arguments and the size of the transformed result, are
the same as for [`brfft`](@ref).
"""
plan_brfft

module FFTW
    import ..DFT: fft, bfft, ifft, rfft, brfft, irfft, plan_fft, plan_bfft, plan_ifft,
        plan_rfft, plan_brfft, plan_irfft, fft!, bfft!, ifft!, plan_fft!, plan_bfft!, plan_ifft!,
        Plan, rfft_output_size, brfft_output_size, plan_inv, normalization, ScaledPlan

    export r2r, r2r!, plan_r2r, plan_r2r!

    """
        plan_dct!(A [, dims [, flags [, timelimit]]])

    Same as [`plan_dct`](@ref), but operates in-place on `A`.
    """
    function plan_dct! end

    """
        plan_idct(A [, dims [, flags [, timelimit]]])

    Pre-plan an optimized inverse discrete cosine transform (DCT), similar to
    [`plan_fft`](@ref) except producing a function that computes
    [`idct`](@ref). The first two arguments have the same meaning as for
    [`idct`](@ref).
    """
    function plan_idct end

    """
        plan_dct(A [, dims [, flags [, timelimit]]])

    Pre-plan an optimized discrete cosine transform (DCT), similar to
    [`plan_fft`](@ref) except producing a function that computes
    [`dct`](@ref). The first two arguments have the same meaning as for
    [`dct`](@ref).
    """
    function plan_dct end

    """
        plan_idct!(A [, dims [, flags [, timelimit]]])

    Same as [`plan_idct`](@ref), but operates in-place on `A`.
    """
    function plan_idct! end

    """
        dct(A [, dims])

    Performs a multidimensional type-II discrete cosine transform (DCT) of the array `A`, using
    the unitary normalization of the DCT. The optional `dims` argument specifies an iterable
    subset of dimensions (e.g. an integer, range, tuple, or array) to transform along.  Most
    efficient if the size of `A` along the transformed dimensions is a product of small primes;
    see [`nextprod`](@ref). See also [`plan_dct`](@ref) for even greater
    efficiency.
    """
    function dct end

    """
        idct(A [, dims])

    Computes the multidimensional inverse discrete cosine transform (DCT) of the array `A`
    (technically, a type-III DCT with the unitary normalization). The optional `dims` argument
    specifies an iterable subset of dimensions (e.g. an integer, range, tuple, or array) to
    transform along.  Most efficient if the size of `A` along the transformed dimensions is a
    product of small primes; see [`nextprod`](@ref).  See also
    [`plan_idct`](@ref) for even greater efficiency.
    """
    function idct end

    """
        dct!(A [, dims])

    Same as [`dct!`](@ref), except that it operates in-place on `A`, which must be an
    array of real or complex floating-point values.
    """
    function dct! end

    """
        idct!(A [, dims])

    Same as [`idct!`](@ref), but operates in-place on `A`.
    """
    function idct! end

    """
        r2r(A, kind [, dims])

    Performs a multidimensional real-input/real-output (r2r) transform
    of type `kind` of the array `A`, as defined in the FFTW manual.
    `kind` specifies either a discrete cosine transform of various types
    (`FFTW.REDFT00`, `FFTW.REDFT01`, `FFTW.REDFT10`, or
    `FFTW.REDFT11`), a discrete sine transform of various types
    (`FFTW.RODFT00`, `FFTW.RODFT01`, `FFTW.RODFT10`, or
    `FFTW.RODFT11`), a real-input DFT with halfcomplex-format output
    (`FFTW.R2HC` and its inverse `FFTW.HC2R`), or a discrete
    Hartley transform (`FFTW.DHT`).  The `kind` argument may be
    an array or tuple in order to specify different transform types
    along the different dimensions of `A`; `kind[end]` is used
    for any unspecified dimensions.  See the FFTW manual for precise
    definitions of these transform types, at http://www.fftw.org/doc.

    The optional `dims` argument specifies an iterable subset of
    dimensions (e.g. an integer, range, tuple, or array) to transform
    along. `kind[i]` is then the transform type for `dims[i]`,
    with `kind[end]` being used for `i > length(kind)`.

    See also [`plan_r2r`](@ref) to pre-plan optimized r2r transforms.
    """
    function r2r end

    """
        r2r!(A, kind [, dims])

    Same as [`r2r`](@ref), but operates in-place on `A`, which must be
    an array of real or complex floating-point numbers.
    """
    function r2r! end

    """
        plan_r2r!(A, kind [, dims [, flags [, timelimit]]])

    Similar to [`plan_fft`](@ref), but corresponds to [`r2r!`](@ref).
    """
    function plan_r2r! end

    """
        plan_r2r(A, kind [, dims [, flags [, timelimit]]])

    Pre-plan an optimized r2r transform, similar to [`plan_fft`](@ref)
    except that the transforms (and the first three arguments)
    correspond to [`r2r`](@ref) and [`r2r!`](@ref), respectively.
    """
    function plan_r2r end

    (Base.USE_GPL_LIBS || Base.fftw_vendor() == :mkl) && include(joinpath("fft", "FFTW.jl"))
end

importall .FFTW
export FFTW, dct, idct, dct!, idct!, plan_dct, plan_idct, plan_dct!, plan_idct!

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

end

max_stars_repo_name: caodi/julia

max_issues_repo_name: caodi/julia

max_forks_repo_name: caodi/julia


text:
using Test;
# Learnings
# - structs must be mutable to update them
# - be careful with variable scopings
# - `@test` works similarly to `@assert` but provides feedback on the failure

mutable struct Command
    Name::String
    Value::Int
end

function parse_line(s)
    parts = split(s)
    name = parts[1]
    value = parse(Int, parts[2])
    return Command(name, value)
end

function swap!(cmd)
    if cmd.Name == "jmp"
        cmd.Name = "nop"
    elseif cmd.Name == "nop"
        cmd.Name = "jmp"
    end
end

function run(fname)
    # read file
    f = open(fname, "r")
    lines = readlines(f)
    close(f)

    result = -1
    instructions = map(parse_line, lines)
    inf_loop = true
    for idx_to_swap in 1:length(instructions)
        # try swapping an instruction
        if idx_to_swap > 1
            swap!(instructions[idx_to_swap - 1])
        end
        swap!(instructions[idx_to_swap])

        idx = 1
        accum = 0
        visited = Set()
        while true
            push!(visited, idx)
            if idx > length(instructions)
                inf_loop = false
                break
            end

            cmd = instructions[idx]
            if cmd.Name == "nop"
                idx += 1
            elseif cmd.Name == "acc"
                idx += 1
                accum += cmd.Value
            elseif cmd.Name == "jmp"
                idx += cmd.Value
            end

            if in(idx, visited)
                # println("Infinite Loop, at idx = ", idx)
                break
            end
        end

        if ! inf_loop
            result = accum
            break
        end
    end

    println("Ran ", fname, " got ", result)
    return result
end

@test run("8ex.txt") == 8

run("8.txt")

ext: jl


text:
# ---
# title: 264. Ugly Number II
# id: problem264
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Math, Dynamic Programming, Heap
# link: <https://leetcode.com/problems/ugly-number-ii/description/>
# hidden: true
# ---
#
# Write a program to find the `n`-th ugly number.
#
# Ugly numbers are **positive numbers** whose prime factors only include `2, 3,
# 5`.
#
# **Example:**
#
#
#
#     Input: n = 10
#     Output: 12
#     Explanation:1, 2, 3, 4, 5, 6, 8, 9, 10, 12 is the sequence of the first 10 ugly numbers.
#
# **Note:**
#
#   1. `1` is typically treated as an ugly number.
#   2. `n` **does not exceed 1690**.
#
#
## @lc code=start
using LeetCode

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


text:
## Allows for determining the step size for first difference.
function gaps(obj::DataFrame,
              Step::Base.Dates.DatePeriod;
              PID::Symbol = names(obj)[1],
              TID::Symbol = names(obj)[2])
    obj = obj[[PID, TID]]
    T = typeof(Step)
    if typeof(obj[TID]) <: AbstractCategoricalVector
        obj[TID] = get.(obj[TID])
    end
    if (T <: Dates.Day) | (T <: Dates.Week)
        output = by(obj, PID) do subdf
            if size(subdf, 1) > 1
                output = vcat(false, diff(subdf[TID]) .== Step)
            else
                output = DataFrame(x1 = false)
            end
            return output
        end
    else
        output = by(obj, PID) do subdf
            if size(subdf, 1) > 1
                output = reduce(vcat, diff(map(ym -> [12 * ym[1] + ym[2]], Dates.yearmonth.(subdf[TID]))))
                output = vcat(false, Dates.Month.(output) .== Step)
            else
                output = DataFrame(x1 = false)
            end
            return output
        end
    end
    names!(output, vcat(names(output)[1], :Valid))
    return output
end
function gaps(obj::DataFrame,
              Step::Real;
              PID::Symbol = names(obj)[1],
              TID::Symbol = names(obj)[2])
    obj = obj[[PID, TID]]
    T = typeof(Step)
    if typeof(obj[TID]) <: AbstractCategoricalVector
        obj[TID] = get.(obj[TID])
    end
    output = by(obj, PID) do subdf
        if size(subdf, 1) > 1
            output = vcat(false, diff(subdf[TID]) .== Step)
        else
            output = DataFrame(x1 = false)
        end
        return output
    end
    names!(output, vcat(names(output)[1], :Valid))
    return output
end
function frequency(obj::AbstractVector)
    if length(obj) < 2
        return -1
    end
    output = minimum(diff(obj))
    return output
end
frequency(obj::AbstractCategoricalVector) = frequency(get.(obj))
function frequency(obj::AbstractVector{T}) where T <: Dates.Date
    if length(obj) < 2
        return -1
    end
    FirstDifference = diff(obj)
    FormalFirstDifference = Dates.canonicalize.(Dates.CompoundPeriod.(FirstDifference))
    if all(length.(getfield.(FormalFirstDifference, :periods)) .== 1)
        output = Dates.CompoundPeriod(minimum(FirstDifference))
        output = Dates.canonicalize(output)
    else
        output = diff(map(ymd -> [12 * ymd[1] + ymd[2], ymd[3]], Dates.yearmonthday.(obj)))
        if all(getindex.(output, 2) .== 0)
            output = minimum(reduce(vcat, first.(output)))
            output = Dates.canonicalize(Dates.CompoundPeriod(Dates.Month(output)))
        else
            output = Dates.CompoundPeriod(minimum(FirstDifference))
        end
    end
    return Dates.Period(first(output.periods))
end

## Drop support for missing data
dropsupportformissing(obj::AbstractVector) = obj
dropsupportformissing(obj::AbstractVector{T}) where T <: Union = Vector{eltype(obj).b}(obj)
function dropsupportformissing(obj::AbstractCategoricalVector)
	T = eltype(obj)
	if isa(T, Union)
		T = T.b
	end
	return CategoricalVector{T}(obj)
end
function dropsupportformissing!(obj::AbstractDataFrame)
	for col ∈ eachcol(obj)
		obj[col[1]] = dropsupportformissing(col[2])
	end
	categorical!(obj, find(col -> col <: AbstractString, (eltype.(obj.columns))))
	return
end

## Linear Independent
function linearindependent(obj::AbstractMatrix{T}) where T <: Real
    cf = cholfact!(Symmetric(obj.'obj, :U), Val{true}, tol = -one(eltype(obj)))
    r = cf.rank
    p = size(obj, 2)
    if r < p
        LI = sort!(cf.piv[1:r])
        obj = obj[:, LI]
    else
        LI = eachindex(cf.piv)
    end
    return obj, LI
end

## Make groups
groups(formula::Formula, data::AbstractDataFrame) =
       groups.(interact.(StatsModels.Terms(formula).terms, data))
groups(obj::AbstractVector) = map(val -> find(equalto(val), obj), unique(obj))
interact(ex::Expr, data::AbstractDataFrame) = interact(getindex.(data, ex.args[2:end]))
interact(name::Symbol, data::AbstractDataFrame) = getindex(data, name)
function interact(obj::AbstractVector{T}) where T <: AbstractCategoricalVector
    n = length(obj[1])
    obj = groups.(obj)
    mapper = Dict{Any,Int64}()
    output = Vector{Int64}(n)
    for idx ∈ eachindex(output)
        output[idx] = get!(mapper, map(each -> findfirst(idx .∈ each), obj), length(mapper))
    end
    return output
end


text:
import LinearAlgebra.LAPACK: syev!, @blasfunc, BlasInt, chkstride1, checksquare, chklapackerror, liblapack
import LinearAlgebra.BLAS: libblas
import LinearAlgebra: eigen, eigen!
import Base: @propagate_inbounds

abstract type Spin end
struct Up <: Spin end
struct Down <: Spin end

"Represents a magnetic Hamiltonian matrix with the block structure [up updown;downup down]"
abstract type AbstractMagneticMatrix{T} <: AbstractMatrix{T} end

data(m::AbstractMatrix) = m
data(m::AbstractMagneticMatrix) = m.data

Base.similar(::Type{M}, i::NTuple{2,Int}) where {M <: AbstractMagneticMatrix} =
    M(Matrix{M.parameters[1]}(undef, i))

for f in (:length, :size, :setindex!, :elsize)
    @eval @inline @propagate_inbounds Base.$f(c::AbstractMagneticMatrix, args...) =
    Base.$f(c.data, args...)
end

Base.pointer(c::AbstractMagneticMatrix, i::Integer) = pointer(c.data, i)

"Magnetic block dimensions"
blockdim(c::AbstractMatrix) = div(size(data(c), 2), 2)

up(c::AbstractMatrix) =   (d = blockdim(c); view(data(c), 1:d, 1:d))
down(c::AbstractMatrix) = (d = blockdim(c); r = d + 1:2 * d; view(data(c), r, r))

# Standard getindex behavior

for f in (:view, :getindex)
    @eval @inline @propagate_inbounds Base.$f(c::AbstractMagneticMatrix, args...) =
        getindex(c.data, args...)
    @eval @inline @propagate_inbounds Base.$f(c::AbstractMagneticMatrix, r::Union{Colon, AbstractUnitRange}, i::Int) =
        MagneticVector(Base.$f(c.data, r, i))
    @eval @inline @propagate_inbounds Base.$f(c::AbstractMagneticMatrix, i::Int, r::Union{Colon, AbstractUnitRange}) =
        MagneticVector(Base.$f(c.data, i, r))
end

Base.similar(c::M, args::AbstractUnitRange...) where {M <: AbstractMagneticMatrix} =
    M(similar(c.data), args...)

"""
    ColinMatrix{T, M <: AbstractMatrix{T}} <: AbstractMagneticMatrix{T}

Defines a Hamiltonian Matrix with [up zeros; zeros down] structure.
It is internally only storing the up and down block.
"""
struct ColinMatrix{T,M <: AbstractMatrix{T}} <: AbstractMagneticMatrix{T}
    data::M
end

function ColinMatrix(up::AbstractMatrix, down::AbstractMatrix)
    @assert size(up) == size(down)
    return ColinMatrix([up down])
end

Base.Array(c::ColinMatrix{T}) where T =
    (d = blockdim(c); [c[Up()] zeros(T, d, d); zeros(T, d, d) c[Down()]])

down(c::ColinMatrix) = (d = blockdim(c); view(c.data, 1:d, d + 1:2 * d))
blockdim(c::ColinMatrix) = size(c.data, 1)

function LinearAlgebra.diag(c::ColinMatrix)
    d = blockdim(c)
    r = LinearAlgebra.diagind(d, d)
    [c[r];c[r.+last(r)]]
end

"""
    NonColinMatrix{T, M <: AbstractMatrix{T}} <: AbstractMagneticMatrix{T}

Defines a Hamiltonian Matrix with [up updown;downup down] structure.
Since atomic projections w.r.t spins are defined rather awkwardly in Wannier90 for exchange calculations,
i.e. storing the up-down parts of an atom sequentially,
a NonColinMatrix reshuffles the entries of a matrix such that it follows the above structure.
"""
struct NonColinMatrix{T,M <: AbstractMatrix{T}} <: AbstractMagneticMatrix{T}
    data::M
end

"Reshuffles standard Wannier90 up-down indices to the ones for the structure of a NonColinMatrix."
function Base.convert(::Type{NonColinMatrix}, m::M) where {M <: AbstractMatrix}
    @assert iseven(size(m, 1)) "Error, dimension of the supplied matrix is odd, i.e. it does not contain both spin components."
    data = similar(m)
    d    = blockdim(m)
    for i in 1:2:size(m, 1), j in 1:2:size(m, 2)
        up_id1 = div1(i, 2)
        up_id2 = div1(j, 2)
        data[up_id1, up_id2] = m[i, j]
        data[up_id1 + d, up_id2] = m[i + 1, j]
        data[up_id1, up_id2 + d] = m[i, j + 1]
        data[up_id1 + d, up_id2 + d] = m[i + 1, j + 1]
    end
    return NonColinMatrix(data)
end

function NonColinMatrix(up::AbstractMatrix{T}, down::AbstractMatrix{T}) where {T}
    @assert size(up) == size(down)
    return NonColinMatrix([up zeros(T, size(up));zeros(T, size(up)) down])
end

Base.Array(c::NonColinMatrix) = copy(c.data)

function uprange(a::DFC.Structures.Projection)
    projrange = range(a)
    if length(projrange) > a.orbital.size
        return range(div1(first(projrange), 2), length = div(length(projrange), 2))
    else
        return projrange
    end
end
uprange(a::DFC.Structures.Atom) = vcat(uprange.(a.projections)...)

## Indexing ##
for f in (:view, :getindex)
    @eval function Base.$f(c::ColinMatrix, a1::T, a2::T) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}}
        projrange1 = range(a1)
        projrange2 = range(a2)

        return ColinMatrix($f(c, projrange1, projrange2), $f(c, projrange1, projrange2 .+ blockdim(c)))
    end
    @eval function Base.$f(c::NonColinMatrix, a1::T, a2::T) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}}
        up_range1 = uprange(a1)
        up_range2 = uprange(a2)
        d = blockdim(c)
        dn_range1 = up_range1 .+ d
        dn_range2 = up_range2 .+ d
        return NonColinMatrix([$f(c, up_range1, up_range2) $f(c, up_range1, dn_range2)
                               $f(c, dn_range1, up_range2) $f(c, dn_range1, dn_range2)])
    end

    @eval Base.$f(c::AbstractMagneticMatrix, a1::T) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}} =
        $f(c, a1, a1)

    @eval Base.$f(c::ColinMatrix, a1::T, a2::T, ::Up) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}} =
        $f(c, range(a1), range(a2))

    @eval Base.$f(c::NonColinMatrix, a1::T, a2::T, ::Up) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}} =
        $f(c, uprange(a1), uprange(a2))

    @eval Base.$f(c::ColinMatrix, a1::T, a2::T, ::Down) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}} =
        $f(c, range(a1), range(a2) .+ blockdim(c))

    @eval Base.$f(c::NonColinMatrix, a1::T, a2::T, ::Down) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}} =
        $f(c, uprange(a1) .+ blockdim(c), uprange(a2) .+ blockdim(c))

    @eval Base.$f(c::NonColinMatrix, a1::T, a2::T, ::Up, ::Down) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}} =
        $f(c, uprange(a1), uprange(a2) .+ blockdim(c))

    @eval Base.$f(c::NonColinMatrix, a1::T, a2::T, ::Down, ::Up) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}} =
        $f(c, uprange(a1) .+ blockdim(c), uprange(a2))

    @eval Base.$f(c::NonColinMatrix, ::Up, ::Down) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}} =
        (s = size(c,1); $f(c, 1:div(s, 2), div(s, 2)+1:s))

    @eval Base.$f(c::NonColinMatrix, ::Down, ::Up) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}} =
        (s = size(c,1); $f(c, div(s, 2)+1:s, 1:div(s, 2)))

    @eval Base.$f(c::AbstractMatrix, a1::T, a2::T, ::Up) where {T<:Union{DFC.Structures.Projection, DFC.Structures.Atom}} =
        $f(c, range(a1), range(a2))

    @eval Base.$f(c::AbstractMatrix, a1::T, a2::T, ::Down) where {T<:Union{DFC.Structures.Projection, DFC.Structures.Atom}} =
        $f(c, range(a1), range(a2) .+ blockdim(c))

    @eval Base.$f(c::AbstractMagneticMatrix, ::Up) =
        (r = 1:blockdim(c); $f(c, r, r))

    @eval Base.$f(c::ColinMatrix, ::Down) =
        (d = blockdim(c); r = 1:d; $f(c, r, r .+ d))

    @eval Base.$f(c::NonColinMatrix, ::Down) =
        (d = blockdim(c); r = d+1 : 2*d; $f(c, r, r))

    @eval Base.$f(c::AbstractMatrix, ::Up) =
        (r = 1:blockdim(c); $f(c, r, r))

    @eval Base.$f(c::AbstractMatrix, ::Down) =
        (d = blockdim(c); r = d + 1:2 * d; $f(c, r, r))

end

for op in (:*, :-, :+, :/)
    @eval @inline Base.$op(c1::ColinMatrix, c2::ColinMatrix) =
        ColinMatrix($op(c1[Up()], c2[Up()]), $op(c1[Down()], c2[Down()]))
    @eval @inline Base.$op(c1::NonColinMatrix, c2::NonColinMatrix) =
        NonColinMatrix($op(c1.data, c2.data))
end

    # BROADCASTING
Base.BroadcastStyle(::Type{<:AbstractMagneticMatrix}) =
    Broadcast.ArrayStyle{AbstractMagneticMatrix}()

Base.ndims(::Type{<:AbstractMagneticMatrix}) =
    2

Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{AbstractMagneticMatrix}}, ::Type{ElType}) where {ElType} =
    Base.similar(AbstractMagneticMatrix{ElType}, axes(bc))

Base.axes(c::AbstractMagneticMatrix) =
    Base.axes(c.data)

@inline @propagate_inbounds Base.broadcastable(c::AbstractMagneticMatrix) =
    c

@inline @propagate_inbounds Base.unsafe_convert(::Type{Ptr{T}}, c::AbstractMagneticMatrix{T}) where {T} =
    Base.unsafe_convert(Ptr{T}, c.data)

for (elty, cfunc) in zip((:ComplexF32, :ComplexF64), (:cgemm_, :zgemm_))
    @eval @inline function LinearAlgebra.mul!(C::ColinMatrix{$elty}, A::ColinMatrix{$elty}, B::ColinMatrix{$elty})
        dim = blockdim(C)
        dim2 = dim * dim
        ccall((@blasfunc($(cfunc)), libblas), Cvoid,
                        (Ref{UInt8}, Ref{UInt8}, Ref{BlasInt}, Ref{BlasInt},
                         Ref{BlasInt}, Ref{$elty}, Ptr{$elty}, Ref{BlasInt},
                         Ptr{$elty}, Ref{BlasInt}, Ref{$elty}, Ptr{$elty},
                         Ref{BlasInt}),
                         'N', 'N', dim, dim,
                         dim, one($elty), A, dim,
                         B, dim, zero($elty), C, dim)
        ccall((@blasfunc($(cfunc)), libblas), Cvoid,
                        (Ref{UInt8}, Ref{UInt8}, Ref{BlasInt}, Ref{BlasInt},
                         Ref{BlasInt}, Ref{$elty}, Ptr{$elty}, Ref{BlasInt},
                         Ptr{$elty}, Ref{BlasInt}, Ref{$elty}, Ptr{$elty},
                         Ref{BlasInt}),
                         'N', 'N', dim, dim,
                         dim, one($elty), pointer(A, dim2 + 1), dim,
                         pointer(B, dim2 + 1), dim, zero($elty), pointer(C, dim2 + 1), dim)

        return C
    end
end

@inline function LinearAlgebra.adjoint(c::AbstractMagneticMatrix)
    out = similar(c)
    adjoint!(out, c)
end

@inline @inbounds function LinearAlgebra.adjoint!(out::ColinMatrix, in1::ColinMatrix)
    dim = blockdim(out)
    for i in 1:dim, j in 1:dim
        out[j, i] = in1[i, j]'
        out[j, i + dim] = in1[i, j + dim]'
    end
    return out
end

@inline LinearAlgebra.adjoint!(out::NonColinMatrix, in1::NonColinMatrix) =
    adjoint!(out.data, in1.data)

for (elty, relty, cfunc) in zip((:ComplexF32, :ComplexF64), (:Float32, :Float64), (:cheev_, :zheev_))
# We use Upper Triangular blas for everything! And Eigvals are always all calculated
    @eval @inline function blas_eig_ccall(A::AbstractMatrix{$elty},
                                          W::AbstractVector{$relty},
                                          work::Vector{$elty},
                                          lwork::BlasInt,
                                          rwork::Vector{$relty},
                                          n::Int,
                                          info::Ref{BlasInt})

    ccall((@blasfunc($cfunc), liblapack), Cvoid,
           (Ref{UInt8}, Ref{UInt8}, Ref{BlasInt}, Ptr{$elty}, Ref{BlasInt},
           Ptr{$relty}, Ptr{$elty}, Ref{BlasInt}, Ptr{$relty}, Ptr{BlasInt}),
           'V', 'U', n, A, n, W, work, lwork, rwork, info)
    chklapackerror(info[])
    end
end

"Vector following the same convention as the in AbstractMagneticMatrix, i.e. first half of the indices contain the up part, second the down part"
struct MagneticVector{T} <: AbstractVector{T}
    data::Vector{T}
end

for f in (:length, :size, :setindex!, :elsize)
    @eval @inline @propagate_inbounds Base.$f(c::MagneticVector, args...) =
    Base.$f(c.data, args...)
end

up(c::MagneticVector) = view(c.data, 1:div(length(c), 2))
down(c::MagneticVector) = (lc = length(c); view(c.data, div(lc, 2):lc))

# Standard getindex behavior
@inline @propagate_inbounds Base.getindex(c::MagneticVector, args...) =
    getindex(c.data, args...)

for f in (:view, :getindex)
    @eval @inline @propagate_inbounds Base.$f(c::MagneticVector, args::AbstractUnitRange...) =
    Base.$f(c.data, args...)
end

Base.similar(v::MagneticVector) = MagneticVector(similar(v.data))

"Reshuffles standard Wannier90 up-down indices to the ones for the structure of a MagneticVector."
function Base.convert(::Type{MagneticVector}, v::V) where {V <: AbstractVector}
    @assert iseven(length(v)) "Error, dimension of the supplied matrix is odd, i.e. it does not contain both spin components."
    data = similar(v)
    vl = length(v)
    d    =div(vl, 2)

    for i in 1:2:vl
        up_id1 = div1(i, 2)
        data[up_id1] = v[i]
        data[up_id1 + d] = v[i + 1]
    end
    return NonColinMatrix(data)
end

## Indexing with atoms and spins ##
for f in (:view, :getindex)
    @eval function Base.$f(c::MagneticVector, a1::T) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}}
        projrange1 = uprange(a1)
        return MagneticVector([$f(c, projrange1); $f(c, projrange1 .+ div(length(c), 2))])
    end
    @eval Base.$f(c::MagneticVector, a1::T, ::Up) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}} =
        $f(c, uprange(a1))

    @eval Base.$f(c::MagneticVector, a1::T, ::Down) where {T <: Union{DFC.Structures.Projection,DFC.Structures.Atom}} =
        $f(c, range(a1), uprange(a2) + div(length(c), 2))

    @eval Base.$f(c::MagneticVector, ::Up) =
        $f(c, 1:div(length(c), 2))

    @eval Base.$f(c::MagneticVector, ::Down) =
        (lc = length(c); $f(c, div(lc, 2) + 1:lc))
end

for op in (:*, :-, :+, :/)
    @eval @inline Base.$op(c1::MagneticVector, c2::MagneticVector) =
        MagneticVector($op(c1.data, c2.data))
end

struct EigCache{T <: AbstractFloat}
    work::Vector{Complex{T}}
    lwork::BlasInt
    rwork::Vector{T}
    n::Int
    info::Ref{BlasInt}
    function EigCache(A::AbstractMatrix{Complex{relty}}) where {relty}
        chkstride1(A)
        elty = Complex{relty}
        n = checksquare(A)
        W     = similar(A, relty, n)
        work  = Vector{elty}(undef, 1)
        lwork = BlasInt(-1)
        rwork = Vector{relty}(undef, max(1, 3n - 2))
        info  = Ref{BlasInt}()

        blas_eig_ccall(A, W, work, lwork, rwork, n, info)
        lwork = BlasInt(real(work[1]))
        resize!(work, lwork)
        return new{relty}(work, lwork, rwork, n, info)
    end
end
EigCache(A::ColinMatrix) = EigCache(up(A))
EigCache(A::NonColinMatrix) = EigCache(A.data)

@inline function eigen!(vals::AbstractVector{T}, vecs::AbstractMatrix{Complex{T}}, c::EigCache{T}) where {T}
    blas_eig_ccall(vecs, vals, c.work, c.lwork, c.rwork, c.n, c.info)
    return Eigen(vals, vecs)
end

@inline function eigen!(vals::AbstractVector{T}, vecs::ColinMatrix{Complex{T}}, c::EigCache{T}) where {T}
    eigen!(view(vals.data, 1:c.n), up(vecs), c)
    eigen!(view(vals.data, c.n + 1:2 * c.n), down(vecs), c)
    return Eigen(vals, vecs)
end

@inline function eigen!(vals::AbstractVector{T}, vecs::NonColinMatrix{Complex{T}}, c::EigCache{T}) where {T}
    eigen!(vals.data, vecs.data, c)
    return Eigen(vals, vecs)
end

@inline function eigen(vecs::AbstractMatrix{Complex{T}}, c::EigCache{T}) where {T}
    out  = copy(vecs)
    vals = similar(out, T, size(out, 2))
    return eigen!(vals, out, c)
end

@inline function eigen(vecs::AbstractMagneticMatrix{Complex{T}}, c::EigCache{T}) where {T}
    out  = copy(vecs)
    vals = MagneticVector(similar(out, T, size(out, 2)))
    return eigen!(vals, out, c)
end

# @inline function eigen!(orig::M) where {M<:AbstractMagneticMatrix}
#     vals, vecs = eigen!(orig.data)
#     return Eigen(vals, M(vecs))
# end

@inline function eigen(h::AbstractMagneticMatrix)
    return eigen(h, EigCache(h))
end
    # Very unoptimized
function Base.Array(e::Eigen{Complex{T},T,<:ColinMatrix{Complex{T}}}) where {T}
    d = size(e.vectors, 1)
    return [e.vectors[1:d, 1:d] * diagm(0 => e.values[1:d]) * e.vectors[1:d, 1:d]' e.vectors[1:d, d + 1:2d] * diagm(0 => e.values[d + 1:2d]) * e.vectors[1:d, d + 1:2d]']
end


text:
#basis functions
begin
	P = [0 1 3 5 6 8 9;0 2 4 2 0 0 3]
	U = [0 0 0 0 2/5 3/5 3/5 1 1 1 1]
	n = size(P,2) - 1
	m = length(U) - 1
	p = m-n-1
	u = range(0, stop=1, length=100)
	funs = zeros(length(u),n+1)
	for i = 1:length(u)
	    funs[i,:] = BasisFuns(n,p,u[Int(i)],U)
	end

	plot(u,funs)
	plot!([2/5, 2/5],[0, 1], line=(:dash,2))
	plot!([3/5, 3/5],[0, 1], line=(:dash,2))
	png("Basisfuns")

	C = funs*P'
	plot(C[:,1],C[:,2], line=(:dot,2))
	plot!(P[1,:],P[2,:], line=(:dash,2))
	png("Example_curve")
end
#Bspline_example
begin
	P1 = [0 1 3 5 5 8 9;0 2 4 2 0 0 3]
	U1 = [0 0 0 0 2/5 3/5 3/5 1 1 1 1]
	n1 = size(P1,2) - 1
	m1 = length(U1) - 1
	p1 = m1-n1-1
	x=4
	u1 = range(0,1,length=1*10^x)
	C1 = zeros(size(P1,1),length(u1))
	for i = 1:length(u1)
    	C1[1,i] = bsplevalc(u1[i],U1,p1,P1[1,:])
    	C1[2,i] = bsplevalc(u1[i],U1,p1,P1[2,:])
	end

	plot(P1[1,:],P1[2,:],line=(:dot,2),show=false)
	plot!(C1[1,:],C1[2,:],line=(:dot,2),show=false)
	png("Bspline_example")
end
#Bspline_surface
begin
	Px2 = [0 3 6 9;0 3 6 9;0 3 6 9]'
	Py2 = [0 0 0 0;2 2 2 2;4 4 4 4]'
	Pz2 = [0 3 3 0;2 5 5 2;0 3 3 0]'
	U2 = [0 0 0 1/2 1 1 1]
	V2 = [0 0 0 1 1 1]
	n2 = size(Px2,2) - 1
	m2 = size(Px2,1) - 1
	h2 = length(U2) - 1
	k2 = length(V2) - 1

	q2 = k2-n2-1
	p2 = h2-m2-1

	x2 = 2
	u2 = range(0,1,length=30)
	v2 = range(0,1,length=30)

	lu2 = length(u2)
	lv2 = length(v2)
	Cx2 = zeros(lu2,lv2)
	Cy2 = zeros(lu2,lv2)
	Cz2 = zeros(lu2,lv2)

	for i = 1:length(u2)
	    for j = 1:length(v2)
	        Cx2[i,j] = bsplevals(u2[i],U2,v2[j],V2,p2,q2,Px2)
	        Cy2[i,j] = bsplevals(u2[i],U2,v2[j],V2,p2,q2,Py2)
	        Cz2[i,j] = bsplevals(u2[i],U2,v2[j],V2,p2,q2,Pz2)
	    end
	end
	plot(Cx2[1,:],Cy2[1,:],Cz2[1,:],show=false)
	for i = 1:length(u2)
	    plot!(Cx2[i,:],Cy2[i,:],Cz2[i,:])
	    plot!(Cx2[:,i],Cy2[:,i],Cz2[:,i])
	end
	for i = 1:size(Px2,1)
	   	plot!(Px2[i,:],Py2[i,:],Pz2[i,:])
	end
	for i = 1:size(Px2,2)
	    plot!(Px2[:,i],Py2[:,i],Pz2[:,i])
	end
	png("Bspline_surface")
end

#NURBS_spiral
begin
	P3 = [0            0
      0.55221      0.77547
     -0.62274       1.7993
      -2.7403      0.80463
      -2.9933      -2.3539
      0.22649      -4.7546
       4.8052      -3.0881
       6.1866       2.4768
       1.7955       7.4013
      -5.6108       6.4752
      -9.4769     -0.90493
       -5.236       -9.069
       4.7457      -10.392
       12.152      -2.3422
       9.6459       9.1974
      -2.0323       14.135
      -13.539       6.9797
      -14.385      -7.4159
      -2.4387      -16.962
       13.091      -12.482
       ]'
	U3 = [0 0 0 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 17 17 17]./17
	W3 = ones(1,size(P3,2))

	npts3 = 150
	nrbplot(P3,W3,U3,npts3)

	u3=0.3 # % for derivative
	n3 = size(P3,2) - 1
	m3 = length(U3) - 1
	p3 = m3-n3-1
	T3 = zeros(length(u3),2)
	C3 = T3
	for i = 1:length(u3)
		T3[i,:] = nrbder(P3,U3,W3,u3[i])
		C3[i,:],~ = nrbevalc(u3[i],U3,p3,P3,W3)
	end
	plot!(C3[:,1],C3[:,2],show=false)
	png("Image_spiral_nrbs")
end

#NURBSsurface
begin
	Px4 = [0 3 6 9;	0 3 6 9;0 3 6 9]'
	Py4 = [0 0 0 0;	2 2 2 2;4 4 4 4]'
	Pz4 = [0 3 3 0;2 5 5 2;0 3 3 0]'
	W4 = ones(4,3)
	U4 = [0 0 0 1/2 1 1 1]
	V4 = [0 0 0 1 1 1]

	nptsu4 = 20
	nptsv4 = 20
	nrbplot3(Px4,Py4,Pz4,W4,U4,V4,nptsu4,nptsv4,1)
	png("NURBSsurface")
end

max_stars_repo_name: barche/ShipstabNURBS

max_issues_repo_name: barche/ShipstabNURBS

max_forks_repo_name: barche/ShipstabNURBS


text:
module Maker

using JLD

export directory,
       file,
       make,
       phony,
       task,
       tasks,
       variable,
       @desc

const CACHEFILE = ".maker-cache.jld"
const CACHEVERSION = v"0.1"

"""
`Maker.AbstractTarget`

`AbstractTarget` is an abstract type covering various "targets" or "tasks".
Each of these targets normally has an action and zero or more dependencies.
"""
abstract AbstractTarget

abstract AbstractCached

include("targets.jl")
include("abstracttarget.jl")
include("task.jl")
include("phony.jl")
include("directory.jl")
include("file.jl")
include("variable.jl")
include("utils.jl")

end # module


text:
using Test: @test, @testset
using LightGraphs: nv, ne
using MetaGraphs: get_prop
using LogicCircuits
using ProbabilisticCircuits

@testset "Chow-Liu Tree learner tests" begin
    data = dataset(twenty_datasets("nltcs"); do_shuffle=false, batch_size=-1)
    train_x = train(data)
    clt = learn_chow_liu_tree(train_x; α=1.0, clt_root="graph_center")
    pv = parent_vector(clt)

    @test clt isa CLT
    @test nv(clt) == num_features(train_x)
    @test ne(clt) == num_features(train_x) - 1

    root = findall(x -> x == 0, pv)[1]
    cpt = get_prop(clt, root, :cpt)
    @test cpt[0] ≈ 0.7828236021007106
    @test cpt[0] + cpt[1] ≈ 1

    cpt1 = get_prop(clt, 1, :cpt)
    @test cpt1[(0,0)] + cpt1[(1, 0)] ≈ 1
    @test cpt1[(0,1)] + cpt1[(1, 1)] ≈ 1
    @test get_prop(clt, 1, :parent) == pv[1]

end


text:
struct HamiltonianProblem{iip} <: DiffEqBase.AbstractDynamicalODEProblem end

function HamiltonianProblem(H, p0, q0, tspan, param=nothing; kwargs...)
  iip = (typeof(q0) <: AbstractArray) && !(typeof(q0) <: SArray)
  HamiltonianProblem{iip}(H, p0, q0, tspan, param; kwargs...)
end

struct PhysicsTag end

function generic_derivative(q0, hami, x)
    ForwardDiff.gradient(hami, x)
end

function generic_derivative(q0::Number, hami, x)
    ForwardDiff.derivative(hami, x)
end

function HamiltonianProblem{false}(H, p0, q0, tspan, param=nothing; kwargs...)
    dp = function (p, q, param, t)
        generic_derivative(q0, q -> -H(p, q, param), q)
    end
    dq = function (p, q, param, t)
        generic_derivative(q0, p -> H(p, q, param), p)
    end
    return ODEProblem(DynamicalODEFunction{false}(dp, dq), ArrayPartition(p0, q0), tspan, param; kwargs...)
end

function HamiltonianProblem{true}(H, p0, q0, tspan, param=nothing; kwargs...)
    let cfg = ForwardDiff.GradientConfig(PhysicsTag(), p0), cfg2 = ForwardDiff.GradientConfig(PhysicsTag(), q0)
        dp = function (dv, p, q, param, t)
            fun2 = q -> -H(p, q, param)
            ForwardDiff.gradient!(dv, fun2, q, cfg2, Val{false}())
        end
        dq = function (dx, p, q, param, t)
            fun1 = p-> H(p, q, param)
            ForwardDiff.gradient!(dx, fun1, p, cfg, Val{false}())
        end
        return ODEProblem(DynamicalODEFunction{true}(dp, dq), ArrayPartition(p0, q0), tspan, param; kwargs...)
    end
end


text:
@cdef :PyDict_New PyPtr ()
@cdef :PyDict_GetItem PyPtr (PyPtr, PyPtr)
@cdef :PyDict_GetItemString PyPtr (PyPtr, Cstring)
@cdef :PyDict_SetItem Cint (PyPtr, PyPtr, PyPtr)
@cdef :PyDict_SetItemString Cint (PyPtr, Cstring, PyPtr)
@cdef :PyDict_DelItemString Cint (PyPtr, Cstring)

const PyDict_Type__ref = Ref(PyNULL)
PyDict_Type() = pyglobal(PyDict_Type__ref, :PyDict_Type)

PyDict_Check(o) = Py_TypeCheckFast(o, Py_TPFLAGS_DICT_SUBCLASS)

PyDict_CheckExact(o) = Py_TypeCheckExact(o, PyDict_Type())

PyDict_FromPairs(kvs) = begin
    r = PyDict_New()
    isnull(r) && return PyNULL
    try
        for (k, v) in kvs
            ko = PyObject_From(k)
            isnull(ko) && (Py_DecRef(r); return PyNULL)
            vo = PyObject_From(v)
            isnull(vo) && (Py_DecRef(r); Py_DecRef(ko); return PyNULL)
            err = PyDict_SetItem(r, ko, vo)
            Py_DecRef(ko)
            Py_DecRef(vo)
            ism1(err) && (Py_DecRef(r); return PyNULL)
        end
        return r
    catch err
        Py_DecRef(r)
        PyErr_SetString(PyExc_Exception(), "Julia error: $err")
        return PyNULL
    end
end

PyDict_FromStringPairs(kvs) = begin
    r = PyDict_New()
    isnull(r) && return PyNULL
    try
        for (k, v) in kvs
            vo = PyObject_From(v)
            isnull(vo) && (Py_DecRef(r); return PyNULL)
            err = PyDict_SetItemString(r, string(k), vo)
            Py_DecRef(vo)
            ism1(err) && (Py_DecRef(r); return PyNULL)
        end
        return r
    catch err
        Py_DecRef(r)
        PyErr_SetString(PyExc_Exception(), "Julia error: $err")
        return PyNULL
    end
end

PyDict_From(x::AbstractDict) = PyDict_FromPairs(x)
PyDict_From(x::AbstractDict{String}) = PyDict_FromStringPairs(x)
PyDict_From(x::AbstractDict{Symbol}) = PyDict_FromStringPairs(x)
PyDict_From(x::NamedTuple) = PyDict_FromStringPairs(pairs(x))
PyDict_From(x::Base.Iterators.Pairs) = PyDict_FromPairs(x)
PyDict_From(x::Base.Iterators.Pairs{String}) = PyDict_FromStringPairs(x)
PyDict_From(x::Base.Iterators.Pairs{Symbol}) = PyDict_FromStringPairs(x)


text:
module Benchmarks

export soundof, soundof_switch, soundof_dict, soundof_little, soundof_dynamic

using ValSplit
using OrderedCollections: LittleDict

const SOUND_OF = Dict{Symbol,String}()
const SOUND_OF_LITTLE = LittleDict{Symbol,String}()

N_VALS = 10
animals = [QuoteNode(Symbol(:animal, i)) for i in 1:N_VALS]
sounds = ["sound $i" for i in 1:N_VALS]

# Generate method definitions and dictionary assignments
for (animal, sound) in zip(animals, sounds)
    @eval begin
        soundof(animal::Val{$animal}) = $sound
        SOUND_OF[$animal] = $sound
        SOUND_OF_LITTLE[$animal] = $sound
    end
end

# Generate manual switch statement
cond_exprs = [:(animal == $a) for a in animals]
branch_exprs = [:(soundof(Val{$a}())) for a in animals]
switch_expr = ValSplit.generate_switch_stmt(cond_exprs, branch_exprs)
@eval function soundof_switch(animal::Symbol)
    $switch_expr
end

soundof_dynamic(animal::Symbol) = soundof(Val(animal))
soundof_dict(animal::Symbol) = SOUND_OF[animal]
soundof_little(animal::Symbol) = SOUND_OF_LITTLE[animal]
@valsplit soundof(Val(animal::Symbol)) = nothing

end

using BenchmarkTools
using .Benchmarks

idxs = rand(1:10, 10000)
animals = [Symbol(:animal, i) for i in idxs]
function test(f, animals)
    for a in animals
        f(a)
    end
end

println("Manual switch statement:")
@btime test(soundof_switch, animals)
println("Global Dict{Symbol,String}:")
@btime test(soundof_dict, animals)
println("Global LittleDict{Symbol,String}:")
@btime test(soundof_little, animals)
println("Dynamic dispatch:")
@btime test(soundof_dynamic, animals)
println("Val-splitting with @valsplit:")
@btime test(soundof, animals)


text:
function normalize(x)
    mag = norm(x)
    if mag > 0.0
        return x ./ mag
    else
        return ones(length(x)) ./ length(x)
    end
end

function ∂normalize∂x(x)
    mag = norm(x)
    n = length(x)
    if mag > 0.0
        n = length(x)
        return 1.0 * I(n) ./ mag - x * transpose(x) ./ mag^3
    else
        return 1.0 * I(n) # TODO: confirm this is good choice
    end
end

function ∂norm∂x(x)
    mag = norm(x)
    if mag > 0.0
        return x' ./ mag
    else
        return ones(1, length(x))
    end
end


text:
# automatically generated by the FlatBuffers compiler, do not modify

MyGame.Example.eval(quote

import ..Example2

struct Any_Monster end
Any_Monster(args...; kwargs...) = Monster(args...; kwargs...)

struct Any_TestSimpleTableWithEnum end
Any_TestSimpleTableWithEnum(args...; kwargs...) = TestSimpleTableWithEnum(args...; kwargs...)

struct Any_MyGame_Example2_Monster end
Any_MyGame_Example2_Monster(args...; kwargs...) = Example2.Monster(args...; kwargs...)

FlatBuffers.@UNION(Any_, (
    Nothing,
    Any_Monster,
    Any_TestSimpleTableWithEnum,
    Any_MyGame_Example2_Monster,
))

end)

max_stars_repo_name: rjkat/flatbuffers-julia

max_issues_repo_name: rjkat/flatbuffers-julia

max_forks_repo_name: rjkat/flatbuffers-julia


text:
"""
This function calculates the operating market profits.
"""
function calculate_realized_operating_profit(prices::AxisArrays.AxisArray{Float64, 2},
                                            marginal_cost::Float64,
                                            output::Array{Float64, 2},
                                            realized_hour_weight::Vector{Float64})

    profit = sum(realized_hour_weight[i] *
                ((prices[1, i] - marginal_cost) * output[1, i])
                 for i in 1:length(realized_hour_weight))

    return profit

end

"""
This function does nothing if the method is not specified for a product.
"""
function calculate_realized_profit(project_name::String,
                                   size::Float64,
                                   zone::String,
                                   product::T,
                                   market_prices::MarketPrices,
                                   capacity_factors::Dict{String, Array{Float64, 2}},
                                   reserve_up_perc::Dict{String, Array{Float64, 2}},
                                   reserve_down_perc::Dict{String, Array{Float64, 2}},
                                   capacity_accepted_bids::Dict{String, Float64},
                                   rec_accepted_bids::Dict{String, Float64},
                                   realized_hour_weight::Vector{Float64},
                                   iteration_year::Int64,
                                   capacity_forward_years::Int64) where T <: Product
    return nothing
end

"""
This function calculates the realized energy market profits.
"""
function calculate_realized_profit(project_name::String,
                                   size::Float64,
                                   zone::String,
                                   product::Energy,
                                   market_prices::MarketPrices,
                                   capacity_factors::Dict{String, Array{Float64, 2}},
                                   reserve_up_perc::Dict{String, Array{Float64, 2}},
                                   reserve_down_perc::Dict{String, Array{Float64, 2}},
                                   capacity_accepted_bids::Dict{String, Float64},
                                   rec_accepted_bids::Dict{String, Float64},
                                   realized_hour_weight::Vector{Float64},
                                   iteration_year::Int64,
                                   capacity_forward_years::Int64)

    update_year = iteration_year

    if in(project_name, keys(capacity_factors))
        output = size * capacity_factors[project_name]

        profit = calculate_realized_operating_profit(get_prices(market_prices, product)["realized"][zone, :, :],
                                                get_marginal_cost(product),
                                                output,
                                                realized_hour_weight)

        return profit, update_year
    else
        return nothing, update_year
    end
end

"""
This function calculates the realized reserve up market profits.
"""
function calculate_realized_profit(project_name::String,
                                  size::Float64,
                                  zone::String,
                                  product::OperatingReserve{ReserveUpEMIS},
                                  market_prices::MarketPrices,
                                  capacity_factors::Dict{String, Array{Float64, 2}},
                                  reserve_up_perc::Dict{String, Array{Float64, 2}},
                                  reserve_down_perc::Dict{String, Array{Float64, 2}},
                                  capacity_accepted_bids::Dict{String, Float64},
                                  rec_accepted_bids::Dict{String, Float64},
                                  realized_hour_weight::Vector{Float64},
                                  iteration_year::Int64,
                                  capacity_forward_years::Int64)

    update_year = iteration_year

    if in(project_name, keys(reserve_up_perc))
        output = size * reserve_up_perc[project_name]

        profit =calculate_realized_operating_profit(get_prices(market_prices, product)["realized"][zone, :, :],
                                                get_marginal_cost(product),
                                                output,
                                                realized_hour_weight)

        return profit, update_year
    else
        return nothing, update_year
    end
end

"""
This function calculates the realized reserve down market profits.
"""
function calculate_realized_profit(project_name::String,
                                  size::Float64,
                                  zone::String,
                                  product::OperatingReserve{ReserveDownEMIS},
                                  market_prices::MarketPrices,
                                  capacity_factors::Dict{String, Array{Float64, 2}},
                                  reserve_up_perc::Dict{String, Array{Float64, 2}},
                                  reserve_down_perc::Dict{String, Array{Float64, 2}},
                                  capacity_accepted_bids::Dict{String, Float64},
                                  rec_accepted_bids::Dict{String, Float64},
                                  realized_hour_weight::Vector{Float64},
                                  iteration_year::Int64,
                                  capacity_forward_years::Int64)

    update_year = iteration_year

    if in(project_name, keys(reserve_down_perc))
        output = size * reserve_down_perc[project_name]

        profit = calculate_realized_operating_profit(get_prices(market_prices, product)["realized"][zone, :, :],
                                                get_marginal_cost(product),
                                                output,
                                                realized_hour_weight)

        return profit, update_year
    else
        return nothing, update_year
    end
end

"""
This function calculates the realized capacity market profits.
"""
function calculate_realized_profit(project_name::String,
                                  size::Float64,
                                  zone::String,
                                  product::Capacity,
                                  market_prices::MarketPrices,
                                  capacity_factors::Dict{String, Array{Float64, 2}},
                                  reserve_up_perc::Dict{String, Array{Float64, 2}},
                                  reserve_down_perc::Dict{String, Array{Float64, 2}},
                                  capacity_accepted_bids::Dict{String, Float64},
                                  rec_accepted_bids::Dict{String, Float64},
                                  realized_hour_weight::Vector{Float64},
                                  iteration_year::Int64,
                                  capacity_forward_years::Int64)

    update_year = iteration_year + capacity_forward_years - 1

    if in(project_name, keys(capacity_accepted_bids))
        profit = size *
                get_derating(product) *
                get_prices(market_prices, product)["realized"][1] *
                capacity_accepted_bids[project_name]

    return profit, update_year
    else
        return nothing, update_year
    end
end

"""
This function calculates the realized REC market profits.
"""
function calculate_realized_profit(project_name::String,
                                  size::Float64,
                                  zone::String,
                                  product::REC,
                                  market_prices::MarketPrices,
                                  capacity_factors::Dict{String, Array{Float64, 2}},
                                  reserve_up_perc::Dict{String, Array{Float64, 2}},
                                  reserve_down_perc::Dict{String, Array{Float64, 2}},
                                  capacity_accepted_bids::Dict{String, Float64},
                                  rec_accepted_bids::Dict{String, Float64},
                                  realized_hour_weight::Vector{Float64},
                                  iteration_year::Int64,
                                  capacity_forward_years::Int64)

    update_year = iteration_year

    if in(project_name, keys(rec_accepted_bids))
        profit = get_prices(market_prices, product)["realized"][1] *
                rec_accepted_bids[project_name]

    return profit, update_year
    else
        return nothing, update_year
    end
end


text:
#### Types and constructors ####

# Convenience constructors
export grayim
function grayim{T<:Union{UFixed,Bool}}(A::AbstractArray{T})
    Base.depwarn("grayim is deprecated, please use ColorView{Gray}(A), possibly in conjunction with ufixedview", :grayim)
    ColorView{Gray}(A)
end
grayim(A::Array{UInt8,2})  = grayim(reinterpret(UFixed8, A))
grayim(A::Array{UInt16,2}) = grayim(reinterpret(UFixed16, A))
grayim(A::Array{UInt8,3})  = grayim(reinterpret(UFixed8, A))
grayim(A::Array{UInt16,3}) = grayim(reinterpret(UFixed16, A))
grayim(A::AbstractArray{UInt8,2})  = grayim(ufixedview(A))
grayim(A::AbstractArray{UInt8,3})  = grayim(ufixedview(A))
grayim(A::AbstractArray{UInt16,2}) = grayim(ufixedview(U16, A))
grayim(A::AbstractArray{UInt16,3}) = grayim(ufixedview(U16, A))

export colorim
function colorim{T}(A::AbstractArray{T,3})
    if size(A, 1) == 4 || size(A, 3) == 4
        error("The array looks like a 4-channel color image. Please specify the colorspace explicitly (e.g. \"ARGB\" or \"RGBA\".)")
    end

    colorim(A, "RGB")
end
function colorim{T<:Fractional}(A::AbstractArray{T,3}, colorspace)
    Base.depwarn("colorim(A, colorspace) is deprecated, use ColorView{C}(A) instead, possibly in conjunction with permuteddimsview and/or ufixedview", :colorim)
    CT = getcolortype(colorspace, eltype(A))
    if 3 <= size(A, 1) <= 4 && 3 <= size(A, 3) <= 4
        error("Both first and last dimensions are of size 3 or 4; impossible to guess which is for color. Use the Image constructor directly.")
    elseif 3 <= size(A, 1) <= 4  # Image as returned by imread for regular 2D RGB images
        ColorView{CT}(A)
    elseif 3 <= size(A, 3) <= 4  # "Matlab"-style image, as returned by convert(Array, im).
        ColorView{CT}(permuteddimsview(A, (3,1,2)))
    else
        error("Neither the first nor the last dimension is of size 3. This doesn't look like an RGB image.")
    end
end
colorim(A::Array{UInt8,3},  colorspace) = colorim(reinterpret(UFixed8,  A), colorspace)
colorim(A::Array{UInt16,3}, colorspace) = colorim(reinterpret(UFixed16, A), colorspace)
colorim(A::AbstractArray{UInt8,3},  colorspace) = colorim(ufixedview(A), colorspace)
colorim(A::AbstractArray{UInt16,3}, colorspace) = colorim(ufixedview(U16, A), colorspace)

colorspacedict = Dict{String,Any}()
for ACV in (Color, AbstractRGB)
    for CV in subtypes(ACV)
        (length(CV.parameters) == 1 && !(CV.abstract)) || continue
        str = string(CV.name.name)
        colorspacedict[str] = CV
    end
end
function getcolortype{T}(str::String, ::Type{T})
    if haskey(colorspacedict, str)
        CV = colorspacedict[str]
        return CV{T}
    else
        if endswith(str, "A")
            CV = colorspacedict[str[1:end-1]]
            return coloralpha(CV){T}
        elseif startswith(str, "A")
            CV = colorspacedict[str[2:end]]
            return alphacolor(CV){T}
        else
            error("colorspace $str not recognized")
        end
    end
end

#### Data and traits ####

# NOTE: when this is deleted, modify ImagesMeta so it doesn't import this
export data
function data(img::AbstractArray)
    Base.depwarn("""

data(A) is deprecated for arrays that are not an ImageMeta. To avoid
using `data`, structure your code like this:

    function myfunction(img::AbstractArray, args...)
        # "real" algorithm goes here
    end
    myfunction(img::ImageMeta, args...) = myfunction(data(img), args...)

""", :data)
    img
end

# Using plain arrays, we used to have to make all sorts of guesses
# about colorspace and storage order. This was a big problem for
# three-dimensional images, image sequences, cameras with more than
# 16-bits, etc.

# Here are the two most important assumptions (see also colorspace below):
defaultarraycolordim = 3
# defaults for plain arrays ("vertical-major")
const yx = [:y, :x]
# order used in Cairo & most image file formats (with color as the very first dimension)
const xy = [:x, :y]
export spatialorder
function spatialorder(img::AbstractArray)
    Base.depwarn("spatialorder is deprecated for general AbstractArrays, please switch to ImagesAxes instead", :spatialorder)
    _spatialorder(img)
end
_spatialorder{M<:Matrix}(::Type{M}) = yx
_spatialorder(img::AbstractArray) = (sdims(img) == 2) ? _spatialorder(Matrix) : error("cannot guess spatial order for images with ", sdims(img), " spatial dimensions")

@deprecate isdirect(img::AbstractArray) true

export colorspace
function colorspace(img)
    Base.depwarn("""
colorspace(img) is deprecated, use eltype(img) instead, possibly in conjunction
with colorview(img)""", :colorspace)
    _colorspace(img)
end
_colorspace{C<:Colorant}(img::AbstractVector{C}) = ColorTypes.colorant_string(C)
_colorspace{C<:Colorant}(img::AbstractMatrix{C}) = ColorTypes.colorant_string(C)
_colorspace{C<:Colorant}(img::AbstractArray{C,3}) = ColorTypes.colorant_string(C)
_colorspace{C<:Colorant}(img::AbstractArray{C}) = ColorTypes.colorant_string(C)
@noinline _colorspace(img::AbstractVector{Bool}) = "Binary"  # noinlines are just for test coverage (julia bug)
@noinline _colorspace(img::AbstractMatrix{Bool}) = "Binary"
@noinline _colorspace(img::AbstractArray{Bool}) = "Binary"
@noinline _colorspace(img::AbstractArray{Bool,3}) = "Binary"
@noinline _colorspace(img::AbstractMatrix{UInt32}) = "RGB24"
@noinline _colorspace(img::AbstractVector) = "Gray"
@noinline _colorspace(img::AbstractMatrix) = "Gray"
_colorspace{T}(img::AbstractArray{T,3}) = (size(img, defaultarraycolordim) == 3) ? "RGB" : error("Cannot infer colorspace of Array, use a color eltype (e.g., colorview)")

export colordim
function colordim(img)
    Base.depwarn("colordim(img) is deprecated, use colorview(img) to represent as a color image", :colordim)
    _colordim(img)
end
_colordim{C<:Colorant}(img::AbstractArray{C}) = 0
_colordim(img::AbstractArray) = 0

export timedim
function timedim(img)
    Base.depwarn("timedim(img) is deprecated for general AbstractArrays, please switch to ImagesAxes", :timedim)
    return 0
end

export limits
function limits(img)
    Base.depwarn("limits(img) is deprecated, use (zero(T),one(T)) where T is the eltype", :limits)
    _limits(img)
end
_limits(img::AbstractArray{Bool}) = 0,1
_limits{T}(img::AbstractArray{T}) = zero(T), one(T)

export storageorder
function storageorder(img::AbstractArray)
    Base.depwarn("storageorder is deprecated, please switch to ImagesAxes and use `axisnames`", :storageorder)
    so = Array(Symbol, ndims(img))
    so[[coords_spatial(img)...]] = spatialorder(img)
    td = timedim(img)
    if td != 0
        so[td] = :t
    end
    so
end

# number of array elements used for each pixel/voxel
export ncolorelem
function ncolorelem{T}(img::AbstractArray{T})
    Base.depwarn("ncolorelem is deprecated, please encode as a color array (possibly with `colorview`) and use `length(eltype(img))`.\nNumeric arrays are assumed to be grayscale and will return 1.", :ncolorelem)
    T <: Colorant ? length(T) : 1
end

#### Utilities for writing "simple algorithms" safely ####
# If you don't feel like supporting multiple representations, call these

# Two-dimensional images
export assert2d
function assert2d(img::AbstractArray)
    Base.depwarn("assert2d is deprecated, write your algorithm as `myfunc{T}(img::AbstractArray{T,2}) instead", :assert2d)
    if ndims(img) != 2
        error("Only two-dimensional images are supported")
    end
end

# "Scalar color", either grayscale, RGB24, or an immutable type
export assert_scalar_color
function assert_scalar_color(img::AbstractArray)
    Base.depwarn("assert_scalar_color is deprecated and can be removed", :assert_scalar_color)
    nothing
end

# Spatial storage order
export isyfirst, isxfirst, assert_yfirst, assert_xfirst
"""
    isyfirst(img)

Return true if the first spatial dimension is `:y`. Supported only
if you use ImagesAxes.

See also: `isxfirst`, `assert_yfirst`.
"""
function isyfirst(img::AbstractArray)
    Base.depwarn("isyfirst is deprecated, please use ImagesAxes and test axisnames(img)[1] directly", :isyfirst)
    spatialorder(img)[1] == :y
end

"""
    assert_yfirst(img)

Throw an error if the first spatial dimension is not `:y`.
"""
function assert_yfirst(img)
    Base.depwarn("assert_yfirst is deprecated, please use ImagesAxes and test axisnames(img)[1] directly", :assert_yfirst)
    if !isyfirst(img)
        error("Image must have y as its first dimension")
    end
end

"""
    isxfirst(img)

Return true if the first spatial dimension is `:x`. Supported only
if you use ImagesAxes.

See also: `isyfirst`, `assert_xfirst`.
"""
function isxfirst(img::AbstractArray)
    Base.depwarn("isxfirst is deprecated, please use ImagesAxes and test axisnames(img)[1] directly", :isxfirst)
    spatialorder(img)[1] == :x
end

"""
    assert_xfirst(img)

Throw an error if the first spatial dimension is not `:x`.
"""
function assert_xfirst(img::AbstractArray)
    Base.depwarn("assert_xfirst is deprecated, please use ImagesAxes and test axisnames(img)[1] directly", :asset_xfirst)
    if !isxfirst(img)
        error("Image must have x as its first dimension")
    end
end

#### Permutations over dimensions ####

export spatialproperties
function spatialproperties(img::AbstractArray)
    Base.depwarn("spatialproperties is deprecated for any arrays other than ImageMeta", :spatialproperties)
    String[]
end

@deprecate spatialpermutation permutation

# Permute the dimensions of an image, also permuting the relevant properties. If you have non-default properties that are vectors or matrices relative to spatial dimensions, include their names in the list of spatialprops.
import Base: permutedims
@deprecate permutedims{S<:AbstractString}(img::StridedArray, pstr::Union{Vector{S}, Tuple{S,Vararg{S}}}, spatialprops::Vector=spatialproperties(img)) permutedims(img, map(Symbol, pstr), spatialprops)
@deprecate permutedims{S<:AbstractString}(img::AbstractArray, pstr::Union{Vector{S}, Tuple{S,Vararg{S}}}, spatialprops::Vector=spatialproperties(img)) permutedims(img, map(Symbol, pstr), spatialprops)

Base.permutedims{S<:Symbol}(img::StridedArray, pstr::Union{Vector{S}, Tuple{Vararg{S}}}, spatialprops::Vector = spatialproperties(img)) = error("not supported, please switch to ImagesAxes")
Base.permutedims{S<:Symbol}(img::AbstractArray, pstr::Union{Vector{S}, Tuple{Vararg{S}}}, spatialprops::Vector = spatialproperties(img)) = error("not supported, please switch to ImagesAxes")

### Functions ###
@deprecate raw rawview
@deprecate raw{C<:Colorant}(A::AbstractArray{C}) rawview(channelview(A))
@deprecate separate{C<:Colorant,N}(img::AbstractArray{C,N}) permuteddimsview(channelview(img), (ntuple(n->n+1, Val{N})..., 1))
if squeeze1
    @deprecate separate{C<:Color1,N}(img::AbstractArray{C,N}) channelview(img)
end


text:
using Plots.PlotMeasures
using LaTeXStrings
using Plots

import Plots: plot,plot!
function plot(steady_states::Vector{<:Branch},data::StateSpace; kwargs...)

	layout = @layout [a;b{1.0w,0.5h}]
	default(); default( ; grid=false,label="",margin=1mm,linewidth=2,lims=:round, kwargs...)

	figure = plot(layout = layout, link = :x, xlim=extrema(data.parameter), size=(300,500), ylabel=L"\mathrm{steady\,\,states}\quad |\,u\,\,|" )
	hline!( [0],subplot=2,linewidth=1,color=:black, xlabel=L"\mathrm{control\,\,condition,}p",
		xmirror=true, topmargin=-5mm, ylabel=L"\mathrm{spectrum}\,\quad \rho(\lambda)")

	vline!(data.targets,subplot=1,linewidth=1,color=:gold)
	vline!(data.targets,subplot=2,linewidth=1,color=:gold)

	for branch ∈ steady_states

		alpha = map(s->window_function(data.parameter,s.z),branch)
		parameter = map( s -> s.z.p, branch)
		states = map(s-> norm(s.z.u), branch)

		realExtrema = map( s->extrema(real(s.λ)), branch)
		imagExtrema = map( s->extrema(imag(s.λ)), branch)

		plot!( parameter, states, subplot=1, alpha=alpha, color=:lightgreen,
			linewidth=map(r->any(λ->λ≠0,r) ? 5 : 0, imagExtrema) )

		plot!( parameter, states, subplot=1, alpha=alpha,
			color=map(r->last(r)<0 ? :darkblue : :lightblue, realExtrema) )

		plot!( parameter, map(r->asinh(first(r)),imagExtrema), fillrange=map(r->asinh(last(r)),imagExtrema),
			subplot=2, alpha=alpha, color=:lightgreen, fillalpha=0.25*alpha, linewidth=0 )

		plot!( parameter, map(r->asinh(last(r)),realExtrema), ribbon=( map(r->asinh(last(r))-asinh(first(r)),realExtrema), zeros(length(parameter))),
			color=map(r->last(r)<0 ? :darkblue : :lightblue, realExtrema),
			subplot=2, alpha=alpha, fillalpha=0.2*alpha )
	end

	plot!( [NaN], [NaN], fillrange=[NaN], color=:lightgreen, fillalpha=0.25, linewidth=0, subplot=2, label=L"\Im\mathrm{m}\lambda" )
	plot!( [NaN], [NaN], fillrange=[NaN], color=:darkblue, fillalpha=0.2, subplot=2, label=L"\Re\mathrm{e}\lambda" )
	xticks!([NaN],subplot=1)
	return figure
end

function plot!(steady_states::Vector{<:Branch},data::StateSpace; determinant=false, kwargs...)

	hline!([0];linewidth=0,color=:black, ylabel=L"\mathrm{steady\,\,states}\quad F_{\theta}(u,p)=0", xlabel=L"\mathrm{control\,\,condition,}p", kwargs...)
	vline!(data.targets;linewidth=1,color=:gold, kwargs...)
	for branch ∈ steady_states

		stability = map( s -> all(real(s.λ).<0), branch)
		parameter = map( s -> s.z.p, branch)

		for idx ∈ 1:dim(branch)

			plot!( parameter, map(s->s.z.u[idx],branch); grid=false,label="",margin=1mm,linewidth=2,
				color=map( stable -> stable ? :darkblue : :lightblue, stability ), kwargs... )
		end
	end
end

function plot(F::Function, θ::AbstractVector, data::StateSpace; kwargs...)

	parameters = (θ=θ,p=minimum(data.parameter))
	hyperparameters = getParameters(data)

	steady_states = deflationContinuation(F,data.roots,parameters,hyperparameters)
	println(steady_states)

	return plot(steady_states,data;kwargs...)
end

function plot!(F::Function, θ::AbstractVector, data::StateSpace; kwargs...)

	parameters = (θ=θ,p=minimum(data.parameter))
	hyperparameters = getParameters(data)

	steady_states = deflationContinuation(F,data.roots,parameters,hyperparameters)
	println(steady_states)

	return plot!(steady_states,data;kwargs...)
end


text:
######################################################################
# pto_linear.jl
#
# like pto, but for linear systems. Should be faster
######################################################################

export pto_linear

function pto_linear(em::ErgodicManager, tm::TrajectoryManager; verbose::Bool=true, logging::Bool=false, max_iters::Int=100, es_crit::Float64=0.0, dd_crit::Float64=1e-6)
	xd0, ud0 = initialize(tm.initializer, em, tm)
	pto_linear(em, tm, xd0, ud0; verbose=verbose, logging=logging, max_iters=max_iters, es_crit=es_crit, dd_crit = dd_crit)
end

function pto_linear(em::ErgodicManager, tm::TrajectoryManager, xd0::VVF, ud0::VVF; verbose::Bool=true, logging::Bool=false, max_iters::Int=100, es_crit::Float64=0.0, dd_crit::Float64=1e-6)

	# let's not overwrite the initial trajectories
	xd = deepcopy(xd0)
	ud = deepcopy(ud0)
	N = tm.N

	# matrices for gradients
	ad = zeros(tm.dynamics.n, N+1)
	bd = zeros(tm.dynamics.m, N)

	# prepare for logging if need be
	if logging
		outfile = open("temp.csv", "w")
		save(outfile, xd)
	end

	if verbose; print_header(); end
	i = 1
	not_finished = true
	es = 0.; cs = 0.; ts = 0.; dd = 0.; step_size = 0.; es_diff=0.0
	es_prev = 0.0
	es_count = 1
	A, B = linearize(tm.dynamics, xd, ud, tm.h)
	K, Ginv = LQ2(A, B, tm.Qn, tm.Rn, tm.N)
	while not_finished

		# determine gradients used in optimization
		gradients!(ad, bd, em, tm, xd, ud)

		# find gains K and descent direction (zd, vd) using LQ
		#A, B = linearize(tm.dynamics, xd, ud, tm.h)
		C = LQ(K, Ginv, A, B, ad, bd, tm.Qn, tm.Rn, tm.N)
		zd, vd = apply_LQ_gains(A, B, K, C)

		# determine step size and descend
		step_size = get_step_size2(tm.descender, em, tm, xd, ud, zd, vd, ad, bd, K, i)

		# descend and project
		#xd, ud = project(em, tm, K, xd, ud, zd, vd, step_size)
		xd, ud = project2(em, tm, K, xd, ud, zd, vd, step_size)

		# compute statistics and report
		es, cs, ts = all_scores(em, tm, xd, ud)
		dd = directional_derivative(ad, bd, zd, vd)
		if verbose; step_report(i, es, cs, ts, dd, step_size); end
		if logging; save(outfile, xd); end

		# check convergence
		i += 1
		es_count = abs(es - es_prev) < 1e-7 ? es_count + 1 : 0
		es_prev = es
		not_finished = check_convergence(es,es_crit,i,max_iters,dd,dd_crit,verbose, es_count)
	end

	# now that we are done, print a special finished report
	if verbose
		print_header()
		if verbose; step_report(i-1, es, cs, ts, dd, step_size); end
	end

	if logging; close(outfile); end

	return xd, ud
end


text:
include("runtests_complex.jl")
include("runtests_complex_all_indices.jl")
include("runtests_autodiff_separate_real_imag.jl")


text:
using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBottom
using Oceananigans.Architectures: arch_array
using Oceananigans.TurbulenceClosures
using Oceananigans.Models.HydrostaticFreeSurfaceModels: compute_vertically_integrated_volume_flux!
using Oceananigans.Models.HydrostaticFreeSurfaceModels: compute_implicit_free_surface_right_hand_side!
using Oceananigans.Models.HydrostaticFreeSurfaceModels: pressure_correct_velocities!

@testset "Immersed boundaries test divergent flow solve with hydrostatic free surface models" begin
    @info "Testing immersed boundaries divergent flow solve"

    for arch in archs
        Nx = 11
        Ny = 11
        Nz = 1

        underlying_grid = RectilinearGrid(arch,
                                          size = (Nx, Ny, Nz),
                                          extent = (Nx, Ny, 1),
                                          halo = (3, 3, 3),
                                          topology = (Periodic, Periodic, Bounded))

        imm1=Int( floor((Nx+1)/2)   )
        imp1=Int( floor((Nx+1)/2)+1 )
        jmm1=Int( floor((Ny+1)/2)   )
        jmp1=Int( floor((Ny+1)/2)+1 )

        bottom = [-1. for i=1:Nx, j=1:Ny ]
        bottom[imm1-1:imp1+1, jmm1-1:jmp1+1] .= 0

        B = arch_array(arch, bottom)
        grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(B))

        free_surfaces = [ImplicitFreeSurface(solver_method=:HeptadiagonalIterativeSolver, gravitational_acceleration=1.0),
                         ImplicitFreeSurface(solver_method=:PreconditionedConjugateGradient, gravitational_acceleration=1.0),
                         ImplicitFreeSurface(gravitational_acceleration=1.0)]

        sol = ()
        f = ()

        for free_surface in free_surfaces

            model = HydrostaticFreeSurfaceModel(grid = grid,
                                                free_surface = free_surface,
                                                buoyancy = nothing,
                                                tracers = nothing,
                                                closure = nothing)

            # Now create a divergent flow field and solve for
            # pressure correction
            u, v, w     = model.velocities
            u[imm1, jmm1, 1:Nz ] .=  1
            u[imp1, jmm1, 1:Nz ] .= -1
            v[imm1, jmm1, 1:Nz ] .=  1
            v[imm1, jmp1, 1:Nz ] .= -1

            events = ((device_event(arch), device_event(arch)), (device_event(arch), device_event(arch)))
            implicit_free_surface_step!(model.free_surface, model, 1.0, 1.5, events)

            sol = (sol..., model.free_surface.η)
            f  = (f..., model.free_surface)
        end

        @test all(interior(sol[1]) .≈ interior(sol[2]) .≈ interior(sol[3]))
    end
end


text:
using Test
using LinearAlgebra
using Distributions
using OceanTurbulenceParameterEstimation
using Oceananigans
using Oceananigans.Units
using Oceananigans.Models.HydrostaticFreeSurfaceModels: ColumnEnsembleSize
using Oceananigans.TurbulenceClosures: ConvectiveAdjustmentVerticalDiffusivity

using OceanTurbulenceParameterEstimation.Observations: FieldTimeSeriesCollector, initialize_simulation!, observation_times
using OceanTurbulenceParameterEstimation.InverseProblems: transpose_model_output, forward_run!

@testset "Forward map tests" begin

    # Generate synthetic observations
    Nz = 8
    N² = 1e-5
    stop_iteration = 3*60*10
    experiment_name = "forward_map_test"

    default_closure() = ConvectiveAdjustmentVerticalDiffusivity(; convective_κz=1.0,
                                                                  background_κz=1e-5,
                                                                  convective_νz=0.9,
                                                                  background_νz=1e-4)

    model_kwargs = (
        tracers = :b,
        buoyancy = BuoyancyTracer(),
        coriolis = FPlane(f=1e-4),
        boundary_conditions = (u = FieldBoundaryConditions(top = FluxBoundaryCondition(-1e-4)),
                               b = FieldBoundaryConditions(top = FluxBoundaryCondition(1e-7), bottom = GradientBoundaryCondition(N²)))
    )

    function build_simulation(size=Nz)
        grid = RectilinearGrid(size=size, z=(-128, 0), topology=(Flat, Flat, Bounded))
        model = HydrostaticFreeSurfaceModel(; grid,
                                              closure = default_closure(),
                                              model_kwargs...)

        set!(model, b = (x, y, z) -> N² * z)

        simulation = Simulation(model; Δt=20.0, stop_iteration)

        return simulation
    end

    # Make observations
    truth_simulation = build_simulation()
    model = truth_simulation.model
    truth_simulation.output_writers[:fields] = JLD2OutputWriter(model, merge(model.velocities, model.tracers),
                                                                schedule = IterationInterval(round(Int, stop_iteration / 10)),
                                                                prefix = experiment_name,
                                                                array_type = Array{Float64},
                                                                field_slicer = nothing,
                                                                force = true)

    run!(truth_simulation)

    data_path = experiment_name * ".jld2"
    observations = SyntheticObservations(data_path, field_names=(:u, :b), normalize=ZScore)

    #####
    ##### Make model data
    #####

    @testset "Single-member forward map" begin
        # First test transpose_model_output
        test_simulation = build_simulation()
        collected_fields = (u = test_simulation.model.velocities.u, b = test_simulation.model.tracers.b)

        time_series_collector = FieldTimeSeriesCollector(collected_fields, observation_times(observations))

        # Test initialize_simulation!
        @info "  Testing initialize_simulation!..."
        random_initial_condition(x, y, z) = rand()

        for field in fields(test_simulation.model)
            set!(field, random_initial_condition)
        end

        test_u = test_simulation.model.velocities.u
        test_v = test_simulation.model.velocities.v
        test_b = test_simulation.model.tracers.b
        @test !all(test_u .== 0)
        @test !all(test_v .== 0)
        @test !all(test_b .== 0)

        test_simulation.stop_iteration = Inf
        initialize_simulation!(test_simulation, observations, time_series_collector)

        @test all(test_v .== 0)

        run!(test_simulation)

        @info "  Testing transpose_model_output..."
        output = transpose_model_output(time_series_collector, observations)[1]

        test_b = output.field_time_serieses.b
        truth_b = observations.field_time_serieses.b
        test_u = output.field_time_serieses.u
        truth_u = observations.field_time_serieses.u

        @test interior(test_b) == interior(truth_b)
        @test interior(test_u) == interior(truth_u)

        @info "  Testing forward_map and output_map..."
        closure = default_closure()

        priors = (
            convective_κz = Normal(closure.convective_κz, 0.05),
            background_κz = Normal(closure.background_κz, 1e-5),
            convective_νz = Normal(closure.convective_νz, 0.01),
            background_νz = Normal(closure.background_νz, 1e-5),
        )

        free_parameters = FreeParameters(priors)

        ensemble_size = 1
        column_ensemble_size = ColumnEnsembleSize(Nz=Nz, ensemble=(ensemble_size, 1), Hz=1)
        ensemble_grid = RectilinearGrid(size=column_ensemble_size, z = (-128, 0), topology = (Flat, Flat, Bounded))
        closure_ensemble = [default_closure() for i = 1:ensemble_grid.Nx, j = 1:ensemble_grid.Ny]

        ensemble_model = HydrostaticFreeSurfaceModel(; grid = ensemble_grid,
                                                       closure = closure_ensemble,
                                                       model_kwargs...)

        ensemble_simulation = Simulation(ensemble_model; Δt=20.0)

        calibration = InverseProblem(observations, ensemble_simulation, free_parameters)
        optimal_parameters = NamedTuple(name => getproperty(closure, name) for name in keys(priors))

        forward_run!(calibration, optimal_parameters)
        truth_u = observations.field_time_serieses.u
        test_u = calibration.time_series_collector.field_time_serieses.u
        @test interior(test_b) == interior(truth_b)
        @test interior(test_u) == interior(truth_u)

        forward_run!(calibration, optimal_parameters)
        truth_u = observations.field_time_serieses.u
        test_u = calibration.time_series_collector.field_time_serieses.u
        @test interior(test_b) == interior(truth_b)
        @test interior(test_u) == interior(truth_u)

        x₁ = forward_map(calibration, optimal_parameters)
        x₂ = forward_map(calibration, optimal_parameters)

        y = observation_map(calibration)

        @test x₁[:, 1:1] == y
        @test x₂[:, 1:1] == y
    end

    @testset "Two-member (2x1) transposition of model output" begin
        ensemble_size = 2
        column_ensemble_size = ColumnEnsembleSize(Nz=Nz, ensemble=(ensemble_size, 1), Hz=1)
        test_simulation = build_simulation(column_ensemble_size)
        collected_fields = (u = test_simulation.model.velocities.u, b = test_simulation.model.tracers.b)
        time_series_collector = FieldTimeSeriesCollector(collected_fields, observation_times(observations))
        initialize_simulation!(test_simulation, observations, time_series_collector)
        run!(test_simulation)
        output = transpose_model_output(time_series_collector, observations)

        truth_b = observations.field_time_serieses.b
        truth_u = observations.field_time_serieses.u

        test_b = output[1].field_time_serieses.b
        test_u = output[1].field_time_serieses.u

        for n in 1:ensemble_size
            @test test_b[n, :, :, :] == truth_b[1, :, : ,:]
            @test test_u[n, :, :, :] == truth_u[1, :, : ,:]
        end
    end

    @testset "Six-member (2x3) transposition of model output" begin
        ensemble_size = 2
        batch_size = 3
        observations_batch = [observations, observations, observations]
        column_ensemble_size = ColumnEnsembleSize(Nz=Nz, ensemble=(ensemble_size, batch_size), Hz=1)
        test_simulation = build_simulation(column_ensemble_size)
        collected_fields = (u = test_simulation.model.velocities.u, b = test_simulation.model.tracers.b)
        time_series_collector = FieldTimeSeriesCollector(collected_fields, observation_times(observations))
        initialize_simulation!(test_simulation, observations_batch, time_series_collector)
        run!(test_simulation)
        output = transpose_model_output(time_series_collector, observations_batch)

        for j in 1:batch_size
            truth_b = observations_batch[j].field_time_serieses.b
            truth_u = observations_batch[j].field_time_serieses.u

            test_b = output[j].field_time_serieses.b
            test_u = output[j].field_time_serieses.u

            for n in 1:ensemble_size
                @test test_b[n, :, :, :] == truth_b[1, :, : ,:]
                @test test_u[n, :, :, :] == truth_u[1, :, : ,:]
            end
        end
    end
end