text:
#=
This file is part of the replication code for: Hasenzagl, T., Pellegrino, F., Reichlin, L., & Ricco, G. (2020). A Model of the Fed's View on Inflation.
Please cite the paper if you are using any part of the code for academic work (including, but not limited to, conference and peer-reviewed papers).
=#

function get_progress(chain::Array{Float64, 2}, istart::Int64, iend::Int64, print_progress=1::Int64, acc_str=""::String)

# ----------------------------------------------------------------------------------------------------------------------
# Get acceptance rate and print current status of the estimation (if enabled)
# ----------------------------------------------------------------------------------------------------------------------

     chain_iter = chain[:, istart:iend];
     acc_rate         = 100*mean(chain_iter[1, 2:end] .!= chain_iter[1, 1:end-1]);
     acc_rate         = round(acc_rate, digits=2);

     if print_progress == 1
          if acc_str != ""
               print("- $acc_str: $acc_rate%\n");
          else
               print("- Acceptance rate: $acc_rate%\n");
          end
     end

     return acc_rate;
end


text:
module Kurpel2018

using JuMP

"""
    build(model, LWH, XYZ)
    build(model, L, W, H, X, Y, Z)

Add the variables and constraints of Kurpel 2018 3D-packing (minimize wasted
container volume, weakly heterogeneous containers, heterogeneous boxes,
free rotation allowed) to the model. The report presents formulations for
many variants using about the same main variables and constraints. The model
implemented here is not the same as any of the presented, but a slight
variation of the model in Section 2.1 (instead homogeneous containers and
minimizing the number of containers, here we have heterogeneous containers
and minimize the space wasted by the containers, or the total volume of the
containers used). This was done to allow comparison with the Chen 1995 model.
The parameter `b` (amount of boxes of some type) was removed, making the
formulation more BPP-like (instead of CSP-like).

For more information see: Kurpel, Deidson Vitorio, Cleder Marcos Schenekemberg,
Cassius Tadeu Scarpin, José Eduardo Pécora Junior, and Leandro C. Coelho.
“The Exact Solutions of Several Classes of Container Loading Problems,” n.d.

# Arguments
- `model`: JuMP model to which the variables and contraints will be added to.
- `lwh′`: three sequences of numbers (with the same length) denoting
  respectively the length/width/height of the boxes/cartons
- `XYZ`: three sequences of numbers (with the same length) denoting
  respectively the length/width/height of the containers
- `no_rotation`: if true the boxes will not be allowed to rotate. Default:
  false.
- `max_packed_volume`: if true the model will not fail if not all boxes cannot
  be packed, but it will maximize the box volume that can be packed, however
  there is no guarantee of minimizing wasted container space anymore (i.e., if
  all boxes can be packed in less than all containers, the solver can put some
  boxes in each container anyway, it will not worry about trying to pack the
  boxes in the most dense way).
"""
function build(model, L, W, H, X, Y, Z;
  no_rotation = false,
  max_packed_volume = false
)
  build(model, (L, W, H), (X, Y, Z);
    no_rotation = no_rotation,
    max_packed_volume = max_packed_volume
  )
end
function build(model, LWH, XYZ;
  no_rotation = false,
  max_packed_volume = false
)
  @assert length(LWH) == 3
  @assert (isone ∘ length ∘ unique ∘ map)(length, LWH)
  L, W, H = LWH # temporary names while the l/w/h is not yet created
  m = length(L)

  if !no_rotation
    # the true l/h/w parameters take orientation in account
    l = Matrix(undef, m, 6)
    w = Matrix(undef, m, 6)
    h = Matrix(undef, m, 6)
    for i = 1:m
      a, b, c = L[i], W[i], H[i]
      l[i,1] = a; l[i,2] = a; l[i,3] = b; l[i,4] = b; l[i,5] = c; l[i,6] = c
      w[i,1] = b; w[i,2] = c; w[i,3] = a; w[i,4] = c; w[i,5] = a; w[i,6] = b
      h[i,1] = c; h[i,2] = b; h[i,3] = c; h[i,4] = a; h[i,5] = b; h[i,6] = a
    end
  end
  @assert length(XYZ) == 3
  @assert (isone ∘ length ∘ unique ∘ map)(length, XYZ)
  X, Y, Z = XYZ

  C = length(X)

  if max_packed_volume
    v = map((a,b,c) -> a*b*c, L, W, H) # box volumes
  else
    v = map((a,b,c) -> a*b*c, X, Y, Z) # container volumes
  end

  # TODO: for each container dimension, x uses the largest value of that
  # dimension among all containers. This way, many unnecessary variables
  # are created. One way to reduce this is to put all containers in a
  # canonical form (shortest side, middle side, longest side). Nevertheless,
  # if the variables are created but not used in constraints, I think they
  # do not cause major performance problems (to be verified).
  if no_rotation
    @variable(model,
      x[1:m, 1:C, 1:maximum(X), 1:maximum(Y), 1:maximum(Z)],
    Bin)
  else
    @variable(model,
      x[1:m, 1:6, 1:C, 1:maximum(X), 1:maximum(Y), 1:maximum(Z)],
    Bin)
  end
  @variable(model, e[1:C], Bin)

  if max_packed_volume
    if no_rotation
      @objective(model, Max, sum(v[i]*x[i,j,p,q,r] for i = 1:m, j = 1:C, p = 1:(X[j]-L[i]+1), q = 1:(Y[j]-W[i]+1), r = 1:(Z[j]-H[i]+1)))
    else
      @objective(model, Max, sum(v[i]*x[i,g,j,p,q,r] for i = 1:m, j = 1:C, g = 1:6, p = 1:(X[j]-l[i,g]+1), q = 1:(Y[j]-w[i,g]+1), r = 1:(Z[j]-h[i,g]+1)))
    end
  else
    @objective(model, Min, sum(e[j] * v[j] for j = 1:C))
  end

  # Constraint (13) -- no overlap.
  for j = 1:C, s = 1:X[j], t = 1:Y[j], u = 1:Z[j]
    if no_rotation
      @constraint(model, sum(
        x[i,j,p,q,r] for i = 1:m, p = max(1,s-L[i]+1):s,
          q = max(1,t-W[i]+1):t, r = max(1,u-H[i]+1):u
      ) <= e[j])
    else
      @constraint(model, sum(
        x[i,g,j,p,q,r] for i = 1:m, g = 1:6, p = max(1,s-l[i,g]+1):s,
          q = max(1,t-w[i,g]+1):t, r = max(1,u-h[i,g]+1):u
      ) <= e[j])
    end
  end

  # Constraint (14) -- every box needs to be packed (and it needs to be in a
  # valid non-protuding-outside-the-container position).
  if !max_packed_volume
    for i = 1:m
      if no_rotation
        @constraint(model, sum(
          x[i,j,p,q,r] for j = 1:C, p = 1:(X[j]-L[i]+1),
            q = 1:(Y[j]-W[i]+1), r = 1:(Z[j]-H[i]+1)
        ) == 1)
      else
        @constraint(model, sum(
          x[i,g,j,p,q,r] for g = 1:6, j = 1:C, p = 1:(X[j]-l[i,g]+1),
            q = 1:(Y[j]-w[i,g]+1), r = 1:(Z[j]-h[i,g]+1)
        ) == 1)
      end
    end
  end

  model
end

"""
    extract_attributions(model; no_rotation = false)
      :: Dict{Int64, Vector{Int64}}

Returns a dictionary with the index of the containers as keys and the index
of the boxes inside that container as values.
"""
function extract_attributions(model;
  no_rotation = false
) :: Dict{Int64, Vector{Int64}}
  # TODO: this is slower than it could be. If the dimensions of every
  # box and each container were passed to this function, then the loops
  # could iterate a lot less unused variables. Nevertheless, 99% is
  # spent on solving the model, retrieving the solution is not
  # performance critical. More than that, is possible that getvalue
  # already copy the Matrix, so the damage is already done.
  x = getvalue(model[:x])
  ix = 0
  m = size(x, ix += 1) # number of boxes
  # size(x, 2) will be always six if we allow rotation
  if !no_rotation
    ix += 1
  end
  C = size(x, ix += 1) # number of containers
  X = size(x, ix += 1)
  Y = size(x, ix += 1)
  Z = size(x, ix += 1)

  d = Dict{Int64, Vector{Int64}}()
  # builds the dictionary from container index to box indexes
  for i = 1:m
    for j = 1:C, p = 1:X, q = 1:Y, r = 1:Z
      if no_rotation
        if isone(round(x[i,j,p,q,r]))
          if haskey(d, j)
            push!(d[j], i)
          else
            d[j] = [i]
          end
          break # already found where the box is so we can go to the next box
        end
      else
        for g = 1:6
          if isone(round(x[i,g,j,p,q,r]))
            if haskey(d, j)
              push!(d[j], i)
            else
              d[j] = [i]
            end
            break # already found where the box is so we can go to the next box
          end
        end
      end
    end
  end
  d
end

"""
    extract_positions(model) :: NTuple{3, Vector{Int64}}

Returns three arrays (x, y, z) of coordinates where the boxes were placed.
"""
function extract_positions(model) :: NTuple{3, Vector{Int64}}
  # TODO: see extract_attributions comment about the performance.
  # TODO: if attributions works then we need to change extract_positions to
  # also take a 'no_rotation' parameter.
  x = getvalue(model[:x])
  m = size(x, 1) # number of boxes
  # size(x, 2) will be always six
  C = size(x, 3) # number of containers
  X = size(x, 4)
  Y = size(x, 5)
  Z = size(x, 6)

  s = Vector{Int64}(undef, m)
  t = Vector{Int64}(undef, m)
  u = Vector{Int64}(undef, m)
  # find and save the position of every box
  for i = 1:m
    for g = 1:6, j = 1:C, p = 1:X, q = 1:Y, r = 1:Z
      if isone(round(x[i,g,j,p,q,r]))
        # we change the coordinate system from base one to base zero here
        s[i], t[i], u[i] = p-1, q-1, r-1
        break # we already found where the box is so we can go to the next box
      end
    end
  end

  (s, t, u)
end

function orientate(g, l, w, h)
  if     g == 1; (l, w, h)
  elseif g == 2; (l, h, w)
  elseif g == 3; (w, l, h)
  elseif g == 4; (w, h, l)
  elseif g == 5; (h, l, w)
  elseif g == 6; (h, w, l)
  else
    error("The g parameter of Kurpel2018.orientate should stay between " *
          "1 and 6 (both inclusive) but it was " * string(g) * ".")
  end
end

"""
    extract_oriented_sizes(model, pqr) :: NTuple{3, Vector{Int64}}

Returns the oriented sizes of the boxes in the solved model.
"""
function extract_oriented_sizes(model, LWH) :: NTuple{3, Vector{Int64}}
  # TODO: see extract_attributions comment about the performance.
  x = getvalue(model[:x])
  m = size(x, 1) # number of boxes
  # size(x, 2) will be always six
  C = size(x, 3) # number of containers
  X = size(x, 4)
  Y = size(x, 5)
  Z = size(x, 6)

  L, W, H = LWH

  l = Vector{Int64}(undef, m)
  w = Vector{Int64}(undef, m)
  h = Vector{Int64}(undef, m)
  # discover which orientation every box is and give the oriented dimensions
  for i = 1:m
    for g = 1:6, j = 1:C, p = 1:X, q = 1:Y, r = 1:Z
      if isone(round(x[i,g,j,p,q,r]))
        l[i], w[i], h[i] = orientate(g, L[i], W[i], H[i])
        break # we already found where the box is so we can go to the next box
      end
    end
  end

  (l, w, h)
end

"""
    extract_solution(model, LWH)

Returns the triple of the results of extract\\_{attributions, positions,
oriented\\_sizes} over the solved model.

This is faster than executing the three functions separately.
"""
function extract_solution(model, LWH)
  # TODO: see extract_attributions comment about the performance.
  x = getvalue(model[:x])
  m = size(x, 1) # number of boxes
  # size(x, 2) will be always six
  C = size(x, 3) # number of containers
  X = size(x, 4)
  Y = size(x, 5)
  Z = size(x, 6)

  L, W, H = LWH

  # for some container, which boxes are inside it
  d = Dict{Int64, Vector{Int64}}()
  # coordinate of the box inside the container
  s = Vector{Int64}(undef, m)
  t = Vector{Int64}(undef, m)
  u = Vector{Int64}(undef, m)
  # the dimensions of the items when already rotated
  l = Vector{Int64}(undef, m)
  w = Vector{Int64}(undef, m)
  h = Vector{Int64}(undef, m)

  for i = 1:m
    for g = 1:6, j = 1:C, p = 1:X, q = 1:Y, r = 1:Z
      if isone(round(x[i,g,j,p,q,r]))
        if haskey(d, j)
          push!(d[j], i)
        else
          d[j] = [i]
        end
        # we change the coordinate system from base one to base zero here
        s[i], t[i], u[i] = p-1, q-1, r-1
        l[i], w[i], h[i] = orientate(g, L[i], W[i], H[i])
        break # we already found where the box is so we can go to the next box
      end
    end
  end

  (d, (s, t, u), (l, w, h))
end

end # module Kurpel2018

ext: jl


text:
"""
For observed covariance matrices and means.

# Constructor

    SemObservedCovariance(;
        specification,
        obs_cov,
        obs_colnames = nothing,
        meanstructure = false,
        obs_mean = nothing,
        n_obs = nothing,
        kwargs...)

# Arguments
- `specification`: either a `RAMMatrices` or `ParameterTable` object (1)
- `obs_cov`: observed covariance matrix
- `obs_colnames::Vector{Symbol}`: column names of the covariance matrix
- `meanstructure::Bool`: does the model have a meanstructure?
- `obs_mean`: observed mean vector
- `n_obs::Number`: number of observed data points (necessary for fit statistics)

# Extended help
## Interfaces
- `n_obs(::SemObservedCovariance)` -> number of observed data points
- `n_man(::SemObservedCovariance)` -> number of manifest variables

- `obs_cov(::SemObservedCovariance)` -> observed covariance matrix
- `obs_mean(::SemObservedCovariance)` -> observed means

## Implementation
Subtype of `SemObserved`

## Remarks
(1) the `specification` argument can also be `nothing`, but this turns of checking whether
the observed data/covariance columns are in the correct order! As a result, you should only
use this if you are shure your covariance matrix is in the right format.

## Additional keyword arguments:
- `spec_colnames::Vector{Symbol} = nothing`: overwrites column names of the specification object
"""
struct SemObservedCovariance{B, C, D, O} <: SemObserved
    obs_cov::B
    obs_mean::C
    n_man::D
    n_obs::O
end
using StructuralEquationModels
function SemObservedCovariance(;
        specification,
        obs_cov,

        obs_colnames = nothing,
        spec_colnames = nothing,

        obs_mean = nothing,
        meanstructure = false,

        n_obs = nothing,

        kwargs...)

    if !meanstructure & !isnothing(obs_mean)

        throw(ArgumentError("observed means were passed, but `meanstructure = false`"))

    elseif meanstructure & isnothing(obs_mean)

        throw(ArgumentError("`meanstructure = true`, but no observed means were passed"))

    end

    if isnothing(spec_colnames) spec_colnames = get_colnames(specification) end

    if !isnothing(spec_colnames) & isnothing(obs_colnames)

        throw(ArgumentError("no `obs_colnames` were specified"))

    elseif !isnothing(spec_colnames) & !(eltype(obs_colnames) <: Symbol)

        throw(ArgumentError("please specify `obs_colnames` as a vector of Symbols"))

    end

    if !isnothing(spec_colnames)
        obs_cov = reorder_obs_cov(obs_cov, spec_colnames, obs_colnames)
        obs_mean = reorder_obs_mean(obs_mean, spec_colnames, obs_colnames)
    end

    n_man = Float64(size(obs_cov, 1))

    return SemObservedCovariance(obs_cov, obs_mean, n_man, n_obs)
end

############################################################################################
### Recommended methods
############################################################################################

n_obs(observed::SemObservedCovariance) = observed.n_obs
n_man(observed::SemObservedCovariance) = observed.n_man

############################################################################################
### additional methods
############################################################################################

obs_cov(observed::SemObservedCovariance) = observed.obs_cov
obs_mean(observed::SemObservedCovariance) = observed.obs_mean

############################################################################################
### Additional functions
############################################################################################

# reorder covariance matrices --------------------------------------------------------------
function reorder_obs_cov(obs_cov, spec_colnames, obs_colnames)
    if spec_colnames == obs_colnames
        return obs_cov
    else
        new_position = [findall(x .== obs_colnames)[1] for x in spec_colnames]
        indices = reshape([CartesianIndex(i, j) for j in new_position for i in new_position], size(obs_cov, 1), size(obs_cov, 1))
        obs_cov = obs_cov[indices]
        return obs_cov
    end
end

# reorder means ----------------------------------------------------------------------------
reorder_obs_mean(obs_mean::Nothing, spec_colnames, obs_colnames) = nothing

function reorder_obs_mean(obs_mean, spec_colnames, obs_colnames)
    if spec_colnames == obs_colnames
        return obs_mean
    else
        new_position = [findall(x .== obs_colnames)[1] for x in spec_colnames]
        obs_mean = obs_mean[new_position]
        return obs_mean
    end
end


text:
asvector(x::AbstractVector) = x
asvector(x) = collect(x)

"""
    pytable(src, format=:pandas; ...)

Construct a Python table from the Tables.jl-compatible table `src`.

The `format` controls the type of the resulting table, and is one of:
- `:pandas`: A `pandas.DataFrame`. Keyword arguments are passed to the `DataFrame` constructor.
- `:columns`: A `dict` mapping column names to columns.
- `:rows`: A `list` of rows, which are `namedtuple`s.
- `:rowdicts`: A `list` of rows, which are `dict`s.
"""
function pytable(src, format=:pandas; opts...)
    format = Symbol(format)
    if format == :pandas
        _pytable_pandas(src; opts...)
    elseif format == :columns
        _pytable_columns(src; opts...)
    elseif format == :rows
        _pytable_rows(src; opts...)
    elseif format == :rowdicts
        _pytable_rowdicts(src; opts...)
    else
        error("invalid format")
    end
end
export pytable

function _pytable_columns(src, cols=Tables.columns(src))
    pydict(
        pystr(String(n)) => asvector(Tables.getcolumn(cols, n))
        for n in Tables.columnnames(cols)
    )
end

function _pytable_rows(src, rows=Tables.rows(src))
    names = Tables.columnnames(rows)
    t = pyimport("collections"=>"namedtuple")("Row", pylist(pystr(string(n)) for n in names))
    pylist(
        t(map(n->Tables.getcolumn(row, n), names)...)
        for row in rows
    )
end

function _pytable_rowdicts(src, rows=Tables.rows(src))
    names = Tables.columnnames(rows)
    pynames = [pystr(string(n)) for n in names]
    pylist(
        pydict(
            p => Tables.getcolumn(row, n)
            for (n,p) in zip(names, pynames)
        )
        for row in rows
    )
end

aspandasvector(x) = asvector(x)

function _pytable_pandas(src, cols=Tables.columns(src); opts...)
    pyimport("pandas").DataFrame(
        pydict(
            pystr(string(n)) => aspandasvector(Tables.getcolumn(cols, n))
            for n in Tables.columnnames(cols)
        );
        opts...
    )
end

function init_tables()
    @require CategoricalArrays="324d7699-5711-5eae-9e2f-1d82baa6b597" @eval begin
        aspandasvector(x::CategoricalArrays.CategoricalArray) = begin
            codes = map(x -> x===missing ? -1 : Int(CategoricalArrays.levelcode(x))-1, x)
            cats = CategoricalArrays.levels(x)
            ordered = x.pool.ordered
            pyimport("pandas").Categorical.from_codes(codes, cats, ordered=ordered)
        end
    end
end


text:
# using LibSerialPort, COBS,
using Statistics
using DataStructures
using AbstractPlotting, WGLMakie, JSServe, Markdown
using AbstractPlotting.MakieLayout
using JSServe: Slider
markdown_css = JSServe.Asset(JSServe.dependency_path("markdown.css"))

top_rpm = 11_500 + 1_150
t4 = 60e6/4
shortest_t = t4/1.1top_rpm
longest_t = 100_000
fps = 30
baudrate = 115200
history = 50
t₀ = time()

gettime() = time() - t₀

function toint(msg)
    y = zero(UInt32)
    for c in msg
        y <<= 8
        y += c
    end
    return y
end

function getrpm(t)
    if iszero(t)
        0.0
    elseif 0 < t ≤ shortest_t
        nothing
    elseif shortest_t < t ≤ longest_t
        t4/t
    else
        0.0
    end
end

last2 = last ∘ last

gety(line, ::Nothing) = last2(line)
gety(_, y) = y

struct Arduino
    c::ReentrantLock
    sp#::SerialPort
    line::Observable{CircularBuffer{Point2f0}}
    function Arduino(sp)
        line = Node(CircularBuffer{Point2f0}(history))
        for i in 1:history
            push!(line[], Point2f0(i/fps, 0.0))
        end
        c = ReentrantLock()
        new(c, sp, line)
    end
end

function sample!(a::Arduino)
    ts = lock(a.c) do
        # sp_flush(a.sp, SP_BUF_INPUT)
        # encode(a.sp, UInt8(0))
        rand(UInt32)
    end
    t = toint(ts)
    rpm = getrpm(t)
    y = gety(a.line[], rpm)
    xy = Point2f0(gettime(), y)
    push!(a.line[], xy)
    a.line[] = a.line[]
end

# port = only(get_port_list())
# sp = LibSerialPort.open(port, baudrate)

a = Arduino(nothing)

reading = @async while true#isopen(sp)
    cond = Condition()
    Timer(x->notify(cond), 1/fps)
    t = @async sample!(a)
    wait(cond)
    wait(t)
end

pwm = Node(1)
on(pwm) do i
    lock(a.c) do
        println(i)
    end
end

function handler(session, request)
    scene, layout = layoutscene(0)
    ax = layout[1,1] = LAxis(scene, xlabel = "Time (s)", ylabel = "RPM", title = "Fan 1")
    lines!(ax, a.line, color = :blue)
    on(a.line) do xys
        xmin = first(xys[1])
        xmax = first(xys[end]) + 1
        ax.targetlimits[] = FRect2D(xmin, 0, xmax - xmin, top_rpm)
    end
    rpm = lift(a.line) do xys
        round(Int, last(last(xys)))
    end
    slider_s = Slider(1:255, pwm)
    dom = md"""
    $scene
    Speed setting: $pwm $slider_s RPM: $rpm
    """
    return JSServe.DOM.div(markdown_css, dom)
end

app = JSServe.Application(handler, "0.0.0.0", 8000)


text:
## R code 6.3
m6.1 <- quap(
    alist(
        height ~ dnorm( mu , sigma ) ,
        mu <- a + bl*leg_left + br*leg_right ,
        a ~ dnorm( 10 , 100 ) ,
        bl ~ dnorm( 2 , 10 ) ,
        br ~ dnorm( 2 , 10 ) ,
        sigma ~ dexp( 1 )
    ) , data=d )

## R code 6.7
m6.2 <- quap(
    alist(
        height ~ dnorm( mu , sigma ) ,
        mu <- a + bl*leg_left,
        a ~ dnorm( 10 , 100 ) ,
        bl ~ dnorm( 2 , 10 ) ,
        sigma ~ dexp( 1 )
    ) , data=d )

## R code 6.9
# kcal.per.g regressed on perc.fat
m6.3 <- quap(
    alist(
        K ~ dnorm( mu , sigma ) ,
        mu <- a + bF*F ,
        a ~ dnorm( 0 , 0.2 ) ,
        bF ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) , data=d )

# kcal.per.g regressed on perc.lactose
m6.4 <- quap(
    alist(
        K ~ dnorm( mu , sigma ) ,
        mu <- a + bL*L ,
        a ~ dnorm( 0 , 0.2 ) ,
        bL ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) , data=d )

## R code 6.10
m6.5 <- quap(
    alist(
        K ~ dnorm( mu , sigma ) ,
        mu <- a + bF*F + bL*L ,
        a ~ dnorm( 0 , 0.2 ) ,
        bF ~ dnorm( 0 , 0.5 ) ,
        bL ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) ,
    data=d )

## R code 6.15
m6.6 <- quap(
    alist(
        h1 ~ dnorm( mu , sigma ),
        mu <- h0*p,
        p ~ dlnorm( 0 , 0.25 ),
        sigma ~ dexp( 1 )
    ), data=d )

## R code 6.16
m6.7 <- quap(
    alist(
        h1 ~ dnorm( mu , sigma ),
        mu <- h0 * p,
        p <- a + bt*treatment + bf*fungus,
        a ~ dlnorm( 0 , 0.2 ) ,
        bt ~ dnorm( 0 , 0.5 ),
        bf ~ dnorm( 0 , 0.5 ),
        sigma ~ dexp( 1 )
    ), data=d )

## R code 6.17
m6.8 <- quap(
    alist(
        h1 ~ dnorm( mu , sigma ),
        mu <- h0 * p,
        p <- a + bt*treatment,
        a ~ dlnorm( 0 , 0.2 ),
        bt ~ dnorm( 0 , 0.5 ),
        sigma ~ dexp( 1 )
    ), data=d )

## R code 6.23
d2$mid <- d2$married + 1
m6.9 <- quap(
    alist(
        happiness ~ dnorm( mu , sigma ),
        mu <- a[mid] + bA*A,
        a[mid] ~ dnorm( 0 , 1 ),
        bA ~ dnorm( 0 , 2 ),
        sigma ~ dexp(1)
    ) , data=d2 )
precis(m6.9,depth=2)

## R code 6.24
m6.10 <- quap(
    alist(
        happiness ~ dnorm( mu , sigma ),
        mu <- a + bA*A,
        a ~ dnorm( 0 , 1 ),
        bA ~ dnorm( 0 , 2 ),
        sigma ~ dexp(1)
    ) , data=d2 )

## R code 6.27
m6.11 <- quap(
    alist(
        C ~ dnorm( mu , sigma ),
        mu <- a + b_PC*P + b_GC*G,
        a ~ dnorm( 0 , 1 ),
        c(b_PC,b_GC) ~ dnorm( 0 , 1 ),
        sigma ~ dexp( 1 )
    ), data=d )

## R code 6.28
m6.12 <- quap(
    alist(
        C ~ dnorm( mu , sigma ),
        mu <- a + b_PC*P + b_GC*G + b_U*U,
        a ~ dnorm( 0 , 1 ),
        c(b_PC,b_GC,b_U) ~ dnorm( 0 , 1 ),
        sigma ~ dexp( 1 )
    ), data=d )

## R code 6.29
library(dagitty)
dag_6.1 <- dagitty( "dag {
    U [unobserved]
    X -> Y
    X <- U <- A -> C -> Y
    U -> B <- C
}")
adjustmentSets( dag_6.1 , exposure="X" , outcome="Y" )

## R code 6.30
library(dagitty)
dag_6.2 <- dagitty( "dag {
    A -> D
    A -> M -> D
    A <- S -> M
    S -> W -> D
}")
adjustmentSets( dag_6.2 , exposure="W" , outcome="D" )


text:
using CorrelationTrackers
using CorrelationFunctions
using XUnit

@testset "Test correlation tracker"  begin include("test-tracker.jl") end


text:
"""
    PrimalDualPoint

# Attributes
- `x::AbstractVector{T}`: vector of primal variables
- `w::AbstractVector{T}`: vector of primal slacks with respect to upper bounds
- `y::AbstractVector{T}`: vector of dual variables
- `s::AbstractVector{T}`: vector of dual variables (reduced costs of `x`)
- `z::AbstractVector{T}`: vector of dual variables (reduced costs of `w`)
"""
mutable struct PrimalDualPoint{T<:Real}
    x::AbstractVector{T}
    w::AbstractVector{T}

    y::AbstractVector{T}
    s::AbstractVector{T}
    z::AbstractVector{T}
end

# extend Base operation on primal-dual points
import Base:
    copy, +

copy(p::PrimalDualPoint) = PrimalDualPoint(copy(p.x), copy(p.w), copy(p.y), copy(p.s), copy(p.z))
+(p1::PrimalDualPoint, p2::PrimalDualPoint) = PrimalDualPoint(
    p1.x + p2.x,
    p1.w + p2.w,
    p1.y + p2.y,
    p1.s + p2.s,
    p1.z + p2.z
)

"""
    Model
    Data structure for a model

# Attributes
-`nconstr::Integer`: Number of constraints
-`nvars::Integer`: Number of variables
-`A::AbstractMatrix{T1<:Real}`: Constraint matrix
-`b::AbstractVector{T2<:Real}`: Right-hand side of the equality constraints
-`c::AbstractVector{T3<:Real}`: Objective coefficient
-`uind::AbstractVector{Ti<:Integer}`: Indices of upper-bounded variables
-`uval::AbstractVector{T4<:Real}`: Upper bounds on the variables. Only finite
    upper bounds are stored.
-`sol::PrimalDualPoint`: Current solution to the problem (may be infeasible at
    the beginning of the optimization)
-`status::Symbol`: Optimization status
"""
mutable struct Model{T1<:Real, T2<:Real, T3<:Real, T4<:Real, Ti<:Integer}

    nconstr::Integer
    nvars::Integer

    A::AbstractMatrix{T1}
    b::AbstractVector{T2}
    c::AbstractVector{T3}
    uind::AbstractVector{Ti}
    uval::AbstractVector{T4}

    sol::PrimalDualPoint
    status::Symbol

end

function Model(
    A::AbstractMatrix{T1},
    b::AbstractVector{T2},
    c::AbstractVector{T3},
    uind::AbstractVector{Ti},
    uval::AbstractVector{T4}
    ) where{T1<:Real, T2<:Real, T3<:Real, T4<:Real, Ti<:Integer}

    (m, n) = size(A)
    p = size(uind, 1)
    n == size(c, 1) || throw(DimensionMismatch("Expected c to have size $(n) but got $(size(c, 1))"))
    m == size(b, 1) || throw(DimensionMismatch("Expected b to have size $(m) but got $(size(b, 1))"))
    if p > 0
        uind[end] <= n  || throw(DimensionMismatch("Model has $(n) vars but upper bound given for var $(uind[end])"))
    end
    s = PrimalDualPoint(ones(n), ones(p), zeros(m), ones(n), ones(p))

    model = Model(m, n, A, b, c, uind, uval, s, :Built)
    return model

end


text:
using TriangularSolve, LinearAlgebra
using Test

function test_solve(::Type{T}) where {T}
  for n ∈ 1:(T === Float32 ? 100 : 200)
    @show n
    for m ∈ max(1,n-10):n+10
      A = rand(T, m, n); res = similar(A);
      B = rand(T, n, n) + I
      @test TriangularSolve.rdiv!(res, A, UpperTriangular(B)) * UpperTriangular(B) ≈ A
      @test TriangularSolve.rdiv!(res, A, UnitUpperTriangular(B)) * UnitUpperTriangular(B) ≈ A
      @test TriangularSolve.rdiv!(res, A, UpperTriangular(B), Val(false)) * UpperTriangular(B) ≈ A
      @test TriangularSolve.rdiv!(res, A, UnitUpperTriangular(B), Val(false)) * UnitUpperTriangular(B) ≈ A
      A = rand(T, n, m); res = similar(A);
      @test LowerTriangular(B) * TriangularSolve.ldiv!(res, LowerTriangular(B), A) ≈ A
      @test UnitLowerTriangular(B) * TriangularSolve.ldiv!(res, UnitLowerTriangular(B), A) ≈ A
      @test LowerTriangular(B) * TriangularSolve.ldiv!(res, LowerTriangular(B), A, Val(false)) ≈ A
      @test UnitLowerTriangular(B) * TriangularSolve.ldiv!(res, UnitLowerTriangular(B), A, Val(false)) ≈ A
    end
  end
end

@testset "TriangularSolve.jl" begin
  @time @testset "Float64" begin
    test_solve(Float64)
  end
  @time @testset "Float32" begin
    test_solve(Float32)
  end
end

using Aqua
Aqua.test_all(TriangularSolve)


text:
using Pkg
using GZip
using CSV
using FileIO
using DelimitedFiles
using Feather
using MatrixMarket
using CSVFiles
using SparseArrays
using NamedArrays
using Arpack
using LinearAlgebra: svd, lu!, qr!
using DataFrames
using BenchmarkTools
using StatsBase
using ThreadsX
using TSne
using UMAP
using Statistics: std, var
using Pipe
using Setfield
using ProgressMeter
using JuliaDB
using SortingLab
using Match
using ProgressBars
using Loess

max_stars_repo_name: Cortalak/Pissaro

max_issues_repo_name: Cortalak/Pissaro

max_forks_repo_name: Cortalak/Pissaro


text:
#define XCB_PACKED __attribute__((__packed__))
#define X_PROTOCOL 11
#define X_PROTOCOL_REVISION 0
#define X_TCP_PORT 6000
#define XCB_CONN_ERROR 1
#define XCB_CONN_CLOSED_EXT_NOTSUPPORTED 2
#define XCB_CONN_CLOSED_MEM_INSUFFICIENT 3
#define XCB_CONN_CLOSED_REQ_LEN_EXCEED 4
#define XCB_CONN_CLOSED_PARSE_ERR 5
#define XCB_CONN_CLOSED_INVALID_SCREEN 6
#define XCB_CONN_CLOSED_FDPASSING_FAILED 7
#define XCB_TYPE_PAD(T,I) (-(I) & (sizeof(T) > 4 ? 3 : sizeof(T) - 1))
#define XCB_NONE 0L
#define XCB_COPY_FROM_PARENT 0L
#define XCB_CURRENT_TIME 0L
#define XCB_NO_SYMBOL 0L

const XCB_CONN_ERROR = (UInt32)(0)

const XCB_CONN_CLOSED_EXT_NOTSUPPORTED = (UInt32)(2)
const XCB_CONN_CLOSED_MEM_INSUFFICIENT = (UInt32)(3)
const XCB_CONN_CLOSED_REQ_LEN_EXCEED = (UInt32)(4)
const XCB_CONN_CLOSED_PARSE_ERR = (UInt32)(5)
const XCB_CONN_CLOSED_INVALID_SCREEN = (UInt32)(6)
const XCB_CONN_CLOSED_FDPASSING_FAILED = (UInt32)(7)

const XCB_NONE = (UInt32)(0)
const XCB_COPY_FROM_PARENT = (UInt8)(0)
const XCB_CURRENT_TIME = (UInt32)(0)
const XCB_NO_SYMBOL = (UInt32)(0)


text:
import Base: rand,
             ∈

export Ball1

"""
    Ball1{N, VN<:AbstractVector{N}} <: AbstractCentrallySymmetricPolytope{N}

Type that represents a ball in the 1-norm (also known as the Manhattan norm).
The ball is also known as a
[cross-polytope](https://en.wikipedia.org/wiki/Cross-polytope).

It is defined as the set

```math
\\mathcal{B}_1^n(c, r) = \\{ x ∈ \\mathbb{R}^n : ∑_{i=1}^n |c_i - x_i| ≤ r \\},
```
where ``c ∈ \\mathbb{R}^n`` is its center and ``r ∈ \\mathbb{R}_+`` its radius.

### Fields

- `center` -- center of the ball as a real vector
- `radius` -- radius of the ball as a scalar (``≥ 0``)

### Examples

Unit ball in the 1-norm in the plane:

```jldoctest ball1_constructor
julia> B = Ball1(zeros(2), 1.)
Ball1{Float64, Vector{Float64}}([0.0, 0.0], 1.0)
julia> dim(B)
2
```

We evaluate the support vector in the East direction:

```jldoctest ball1_constructor
julia> σ([0.,1], B)
2-element Vector{Float64}:
 0.0
 1.0
```
"""
struct Ball1{N, VN<:AbstractVector{N}} <: AbstractCentrallySymmetricPolytope{N}
    center::VN
    radius::N

    # default constructor with domain constraint for radius
    function Ball1(center::VN, radius::N) where {N, VN<:AbstractVector{N}}
        @assert radius >= zero(N) "radius must not be negative"
        return new{N, VN}(center, radius)
    end
end

isoperationtype(::Type{<:Ball1}) = false
isconvextype(::Type{<:Ball1}) = true

# --- AbstractCentrallySymmetric interface functions ---

"""
    center(B::Ball1)

Return the center of a ball in the 1-norm.

### Input

- `B` -- ball in the 1-norm

### Output

The center of the ball in the 1-norm.
"""
function center(B::Ball1)
    return B.center
end

# --- AbstractPolytope interface functions ---

"""
    vertices_list(B::Ball1{N, VN}) where {N, VN<:AbstractVector}

Return the list of vertices of a ball in the 1-norm.

### Input

- `B` -- ball in the 1-norm

### Output

A list containing the vertices of the ball in the 1-norm.
"""
function vertices_list(B::Ball1{N, VN}) where {N, VN<:AbstractVector}
    # fast evaluation if B has radius 0
    if iszero(B.radius)
        return [B.center]
    end
    vertices = Vector{VN}(undef, 2 * dim(B))
    j = 0
    v = copy(B.center)
    @inbounds for i in 1:dim(B)
        ci = v[i]
        v[i] += B.radius
        j += 1
        vertices[j] = copy(v)
        v[i] = ci - B.radius
        j += 1
        vertices[j] = copy(v)
        v[i] = ci  # restore old value
    end
    return vertices
end

# --- LazySet interface functions ---

"""
    σ(d::AbstractVector, B::Ball1)

Return the support vector of a ball in the 1-norm in a given direction.

### Input

- `d` -- direction
- `B` -- ball in the 1-norm

### Output

Support vector in the given direction.
"""
function σ(d::AbstractVector, B::Ball1)
    res = copy(B.center)
    imax = argmax(abs.(d))
    res[imax] += sign(d[imax]) * B.radius
    return res
end

"""
    ∈(x::AbstractVector, B::Ball1, [failfast]::Bool=false)

Check whether a given point is contained in a ball in the 1-norm.

### Input

- `x` -- point/vector
- `B` -- ball in the 1-norm
- `failfast` -- (optional, default: `false`) optimization for negative answer

### Output

`true` iff ``x ∈ B``.

### Notes

The default behavior (`failfast == false`) is worst-case optimized, i.e., the
implementation is optimistic and first computes (see below) the whole sum before
comparing to the radius.
In applications where the point is typically far away from the ball, the option
`failfast == true` terminates faster.

### Algorithm

Let ``B`` be an ``n``-dimensional ball in the 1-norm with radius ``r`` and let
``c_i`` and ``x_i`` be the ball's center and the vector ``x`` in dimension
``i``, respectively.
Then ``x ∈ B`` iff ``∑_{i=1}^n |c_i - x_i| ≤ r``.

### Examples

```jldoctest
julia> B = Ball1([1., 1.], 1.);

julia> [.5, -.5] ∈ B
false
julia> [.5, 1.5] ∈ B
true
```
"""
function ∈(x::AbstractVector, B::Ball1, failfast::Bool=false)
    @assert length(x) == dim(B) "a $(length(x))-dimensional vector is " *
        "incompatible with a $(dim(B))-dimensional set"
    N = promote_type(eltype(x), eltype(B))
    sum = zero(N)
    @inbounds for (i, xi) in enumerate(x)
        sum += abs(B.center[i] - xi)
        if failfast && sum > B.radius
            return false
        end
    end
    return sum <= B.radius
end

"""
    rand(::Type{Ball1}; [N]::Type{<:Real}=Float64, [dim]::Int=2,
         [rng]::AbstractRNG=GLOBAL_RNG, [seed]::Union{Int, Nothing}=nothing
        )

Create a random ball in the 1-norm.

### Input

- `Ball1` -- type for dispatch
- `N`     -- (optional, default: `Float64`) numeric type
- `dim`   -- (optional, default: 2) dimension
- `rng`   -- (optional, default: `GLOBAL_RNG`) random number generator
- `seed`  -- (optional, default: `nothing`) seed for reseeding

### Output

A random ball in the 1-norm.

### Algorithm

All numbers are normally distributed with mean 0 and standard deviation 1.
Additionally, the radius is nonnegative.
"""
function rand(::Type{Ball1};
              N::Type{<:Real}=Float64,
              dim::Int=2,
              rng::AbstractRNG=GLOBAL_RNG,
              seed::Union{Int, Nothing}=nothing
             )
    rng = reseed(rng, seed)
    center = randn(rng, N, dim)
    radius = abs(randn(rng, N))
    return Ball1(center, radius)
end

"""
    constraints_list(P::Ball1{N}) where {N}

Return the list of constraints defining a ball in the 1-norm.

### Input

- `B` -- ball in the 1-norm

### Output

The list of constraints of the ball.

### Algorithm

The constraints can be defined as ``d_i^T (x-c) ≤ r`` for all ``d_i``, where
``d_i`` is a vector with elements ``1`` or ``-1`` in ``n`` dimensions. To span
all possible ``d_i``, the function `Iterators.product` is used.
"""
function constraints_list(B::Ball1{N}) where {N}
    n = dim(B)
    c, r = B.center, B.radius
    clist = Vector{LinearConstraint{N, Vector{N}}}(undef, 2^n)
    for (i, di) in enumerate(Iterators.product([[one(N), -one(N)] for i = 1:n]...))
        d = collect(di) # tuple -> vector
        clist[i] = LinearConstraint(d, dot(d, c) + r)
    end
    return clist
end

"""
    translate(B::Ball1, v::AbstractVector)

Translate (i.e., shift) a ball in the 1-norm by a given vector.

### Input

- `B` -- ball in the 1-norm
- `v` -- translation vector

### Output

A translated ball in the 1-norm.

### Algorithm

We add the vector to the center of the ball.

### Notes

See also [`translate!(::Ball1, ::AbstractVector)`](@ref) for the in-place version.
"""
function translate(B::Ball1, v::AbstractVector)
    return translate!(copy(B), v)
end

"""
    translate!(B::Ball1, v::AbstractVector)

Translate (i.e., shift) a ball in the 1-norm by a given vector, in-place.

### Input

- `B` -- ball in the 1-norm
- `v` -- translation vector

### Output

The ball `B` translated by `v`.

### Notes

See also [`translate(::Ball1, ::AbstractVector)`](@ref) for the out-of-place version.
"""
function translate!(B::Ball1, v::AbstractVector)
    @assert length(v) == dim(B) "cannot translate a $(dim(B))-dimensional " *
                                "set by a $(length(v))-dimensional vector"
    c = center(B)
    c .+= v
    return B
end

function project(B::Ball1, block::AbstractVector{Int}; kwargs...)
    return Ball1(B.center[block], B.radius)
end


text:
"""
    minmaxflip(x, y)

Returns `x, y` sorted lowest to highest and a bool that indicates if a swap
was needed.
"""
minmaxflip(x, y) = y < x ? (y, x, true) : (x, y, false)

permutationtuples(n) = permutationtuples(Val(n))
permutationtuples(::Val{1}) = ((1,),)
permutationtuples(::Val{2}) = ((1,2),(2,1))
permutationtuples(::Val{3}) = ((1,2,3),(1,3,2),(2,1,3),(2,3,1),(3,1,2),(3,2,1))
permutationtuples(::Val{4}) = ((1,2,3,4), (1,2,4,3), (1,3,2,4), (1,3,4,2),
                               (1,4,2,3), (1,4,3,2), (2,1,3,4), (2,1,4,3),
                               (2,3,1,4), (2,3,4,1), (2,4,1,3), (2,4,3,1),
                               (3,1,2,4), (3,1,4,2), (3,2,1,4), (3,2,4,1),
                               (3,4,1,2), (3,4,2,1), (4,1,2,3), (4,1,3,2),
                               (4,2,1,3), (4,2,3,1), (4,3,1,2), (4,3,2,1))
"""
    tuplesort(a)

Returns the tuple `a` as a sorted tuple.
"""
@inline tuplesort(a::Tuple) = tuplesort(a...)

"""
    tuplesortpermnum(a)

Returns the lexicographic permutation number `p` to sort the tuple `a`.
"""
@inline tuplesortpermnum(a::Tuple) = tuplesortpermnum(a...)

"""
    tuplesortperm(a)

Returns the permutation to sort `a`.
"""
function tuplesortperm(a::Tuple)
    return @inbounds permutationtuples(length(a))[tuplesortpermnum(a)]
end

tuplesortpermutation(x::Tuple) = Permutation{length(x)}(tuplesortpermnum(x...))
tuplesortpermutation(args...) = Permutation{length(args)}(tuplesortpermnum(args...))

tuplesort(a) = Tuple(a)
tuplesortpermnum(a) = 1

tuplesort(a, b) = minmax(a, b)
function tuplesortpermnum(a, b)
    a, b, s1 = minmaxflip(a, b)
    p = s1 ? 2 : 1
    return p
end

function tuplesort(a, b, c)
    # Use a (Bose-Nelson Algorithm based) sorting network from
    # <http://pages.ripco.net/~jgamble/nw.html> to sort the vertices.
    b, c = minmax(b, c)
    a, c = minmax(a, c)
    a, b = minmax(a, b)

    return (a, b, c)
end

function tuplesortpermnum(a, b, c)
    # Use a (Bose-Nelson Algorithm based) sorting network from
    # <http://pages.ripco.net/~jgamble/nw.html> to sort the vertices.
    b, c, s1 = minmaxflip(b, c)
    a, c, s2 = minmaxflip(a, c)
    a, b, s3 = minmaxflip(a, b)

    if     !s1 && !s2 && !s3; return 1
    elseif  s1 && !s2 && !s3; return 2
    elseif !s1 &&  s2 && !s3; return 4
    elseif  s1 &&  s2 && !s3
        @assert c < b && b ≤ c
        throw(AssertionError())
    elseif !s1 && !s2 &&  s3; return 3
    elseif  s1 && !s2 &&  s3; return 5
    elseif !s1 &&  s2 &&  s3; return 4
    elseif  s1 &&  s2 &&  s3; return 6
    else
        # Throw an error as this should never be reached.
        throw(AssertionError())
    end
end

function tuplesort(a, b, c, d)
    # Use a (Bose-Nelson Algorithm based) sorting network from
    # <http://pages.ripco.net/~jgamble/nw.html> to sort the vertices.
    a, b = minmax(a, b)
    c, d = minmax(c, d)
    a, c = minmax(a, c)
    b, d = minmax(b, d)
    b, c = minmax(b, c)

    return (a, b, c, d)
end

function tuplesortpermnum(a, b, c, d)
    # Use a (Bose-Nelson Algorithm based) sorting network from
    # <http://pages.ripco.net/~jgamble/nw.html> to sort the vertices.
    a, b, s1 = minmaxflip(a, b)
    c, d, s2 = minmaxflip(c, d)
    a, c, s3 = minmaxflip(a, c)
    b, d, s4 = minmaxflip(b, d)
    b, c, s5 = minmaxflip(b, c)

    if     !s1 && !s2 && !s3 && !s4 && !s5; return 1
    elseif  s1 && !s2 && !s3 && !s4 && !s5; return 7
    elseif !s1 &&  s2 && !s3 && !s4 && !s5; return 2
    elseif  s1 &&  s2 && !s3 && !s4 && !s5; return 8
    elseif !s1 && !s2 &&  s3 && !s4 && !s5; return 15
    elseif  s1 && !s2 &&  s3 && !s4 && !s5; return 13
    elseif !s1 &&  s2 &&  s3 && !s4 && !s5; return 19
    elseif  s1 &&  s2 &&  s3 && !s4 && !s5
        @assert b < a && d < c && d < b && a ≤ c && a ≤ b
        throw(AssertionError())
    elseif !s1 && !s2 && !s3 &&  s4 && !s5; return 4
    elseif  s1 && !s2 && !s3 &&  s4 && !s5; return 10
    elseif !s1 &&  s2 && !s3 &&  s4 && !s5
        @assert a ≤ b && d < c && a ≤ d && c < b && c ≤ d
        throw(AssertionError())
    elseif  s1 &&  s2 && !s3 &&  s4 && !s5
        @assert b < a && d < c && b ≤ d && c < a && c ≤ d
        throw(AssertionError())
    elseif !s1 && !s2 &&  s3 &&  s4 && !s5; return 17
    elseif  s1 && !s2 &&  s3 &&  s4 && !s5; return 18
    elseif !s1 &&  s2 &&  s3 &&  s4 && !s5; return 23
    elseif  s1 &&  s2 &&  s3 &&  s4 && !s5; return 24
    elseif !s1 && !s2 && !s3 && !s4 &&  s5; return 3
    elseif  s1 && !s2 && !s3 && !s4 &&  s5; return 9
    elseif !s1 &&  s2 && !s3 && !s4 &&  s5; return 5
    elseif  s1 &&  s2 && !s3 && !s4 &&  s5; return 11
    elseif !s1 && !s2 &&  s3 && !s4 &&  s5; return 13
    elseif  s1 && !s2 &&  s3 && !s4 &&  s5; return 15
    elseif !s1 &&  s2 &&  s3 && !s4 &&  s5; return 19
    elseif  s1 &&  s2 &&  s3 && !s4 &&  s5; return 21
    elseif !s1 && !s2 && !s3 &&  s4 &&  s5; return 4
    elseif  s1 && !s2 && !s3 &&  s4 &&  s5; return 10
    elseif !s1 &&  s2 && !s3 &&  s4 &&  s5; return 6
    elseif  s1 &&  s2 && !s3 &&  s4 &&  s5; return 12
    elseif !s1 && !s2 &&  s3 &&  s4 &&  s5; return 14
    elseif  s1 && !s2 &&  s3 &&  s4 &&  s5; return 16
    elseif !s1 &&  s2 &&  s3 &&  s4 &&  s5; return 20
    elseif  s1 &&  s2 &&  s3 &&  s4 &&  s5; return 22
    else
        # Throw an error as this should never be reached.
        throw(AssertionError())
    end
end


text:
using DocOpt

doc="""Align two feature vector sequences and create parallel data in batch.
Note that filename of source and target feature is assumed to be same.

Usage:
    align.jl [options] <src_dir> <tgt_dir> <dst_dir>
    align.jl --version
    align.jl -h | --help

Options:
    -h --help       show this message
    --threshold=TH  threshold that is used to remove silence [default: -14.0]
    --max=MAX       Maximum number that will be processed [default: 200]
"""

using VoiceConversion
using HDF5, JLD

using Logging
@Logging.configure(level=DEBUG, output=STDOUT)

searchdir(path, key) = filter(x -> contains(x, key), readdir(path))

function mkdir_if_not_exist(dir)
    if !isdir(dir)
        println("Create $(dir)")
        run(`mkdir -p $dir`)
    end
end

function align!(src::Dict,
                tgt::Dict;
                threshold::AbstractFloat=-14.0,
                )
    src_fm = src["feature_matrix"]
    tgt_fm = tgt["feature_matrix"]

    max_num_frames = max(size(src_fm, 2), size(tgt_fm, 2))

    if src["type"] == "MelCepstrum"
        src_fm, tgt_fm = align_mcep(src_fm, tgt_fm,
                                    src["alpha"],
                                    src["fftlen"];
                                    threshold=threshold)
    else
        src_fm, tgt_fm = _align(src_fm, tgt_fm)
    end

    @assert size(src_fm) == size(tgt_fm)

    @info("The number of aligned frames: $(size(src_fm, 2))/$(max_num_frames)")
    if size(src_fm, 2) ==  0
        @warn("No frame found in aligned data. Probably threshold is too high.")
    end

    @assert !any(isnan(src_fm))
    @assert !any(isnan(tgt_fm))

    src["feature_matrix"] = src_fm
    tgt["feature_matrix"] = tgt_fm

    # type Dict{Union(UTF8String, ASCIIString), Any} is saved as
    # Dict{UTF8String, Any} and cause error in reading JLD file.
    # remove off Union and then save do the trick (but why? bug in HDF5?)
    @assert isa(src, Dict{String, Any})
    @assert isa(tgt, Dict{String, Any})

    Dict{String, Any}(src), Dict{String, Any}(tgt)
end

let
    args = docopt(doc, version=v"0.0.1")

    srcdir = args["<src_dir>"]
    tgtdir = args["<tgt_dir>"]
    dstdir = args["<dst_dir>"]
    mkdir_if_not_exist(dstdir)

    threshold = parse(Float64, args["--threshold"])
    nmax = parse(Int, args["--max"])

    files = searchdir(srcdir, ".jld")
    @info("$(length(files)) data found.")

    count = 0
    for filename in files
        # filename is assumed to be same between src and tgt
        srcpath = joinpath(srcdir, filename)
        tgtpath = joinpath(tgtdir, filename)
        savepath = joinpath(dstdir, string(splitext(basename(srcpath))[1],
                                      "_parallel.jld"))

        @info("Start processing $(srcpath) and $(tgtpath)")
        elapsed = @elapsed begin
            src = load(srcpath)
            tgt = load(tgtpath)

            src, tgt = align!(src, tgt, threshold=threshold)
            save(savepath,
                 "description", "Parallel data",
                 "src", src,
                 "tgt", tgt
                 )
        end
        @info("Elapsed time in alignment is $(elapsed) sec.")
        @info("Dumped to $(savepath)")

        count += 1
        count >= nmax && break
    end

    @info("Finished")
end


text:
using Documenter, AlgebraOfGraphics, Literate, Glob
# avoid font caching warning in docs
using CairoMakie
CairoMakie.activate!()
scatter(rand(10), rand(10))

# generate examples
GENERATED = joinpath(@__DIR__, "src", "generated")
SOURCE_FILES = Glob.glob("*.jl", GENERATED)
foreach(fn -> Literate.markdown(fn, GENERATED), SOURCE_FILES)

makedocs(
         sitename="Algebra of Graphics",
         pages = Any[
                     "index.md",
                     "generated/tutorial.md",
                     "generated/statistics.md",
                     "analyses.md",
                     "generated/internals.md",
                    ]
        )

deploydocs(
    repo = "github.com/JuliaPlots/AlgebraOfGraphics.jl.git",
    push_preview = true
)


text:
#module CUDNN

include("cudnn_retry.jl")       # TODO: export this to Ops20_gpu, auto-add unsafe_cudnn?
include("common.jl")
include("activation.jl")
include("dropout.jl")
#include("multiheadattn.jl")

#end

## grep cudnnHandle_t
# add8: cudnnBackendExecute
# add8: cudnnBuildRNNDynamic
# add8: cudnnCTCLoss_v8
# add8: cudnnGetFoldedConvBackwardDataDescriptors
# add8: cudnnGetNormalizationBackwardWorkspaceSize
# add8: cudnnGetNormalizationForwardTrainingWorkspaceSize
# add8: cudnnGetNormalizationTrainingReserveSpaceSize
# add8: cudnnGetRNNTempSpaceSizes
# add8: cudnnGetRNNWeightParams
# add8: cudnnGetRNNWeightSpaceSize
# add8: cudnnNormalizationBackward
# add8: cudnnNormalizationForwardInference
# add8: cudnnNormalizationForwardTraining
# add8: cudnnRNNBackwardData_v8
# add8: cudnnRNNBackwardWeights_v8
# add8: cudnnRNNForward
# cudnnActivationBackward
# cudnnActivationForward
# cudnnAddTensor
# cudnnBatchNormalizationBackward
# cudnnBatchNormalizationBackwardEx
# cudnnBatchNormalizationForwardInference
# cudnnBatchNormalizationForwardTraining
# cudnnBatchNormalizationForwardTrainingEx
# cudnnConvolutionBackwardBias
# cudnnConvolutionBackwardData
# cudnnConvolutionBackwardFilter
# cudnnConvolutionBiasActivationForward
# cudnnConvolutionForward
# cudnnCreate
# cudnnCTCLoss
# cudnnDestroy
# cudnnDivisiveNormalizationBackward
# cudnnDivisiveNormalizationForward
# cudnnDropoutBackward
# cudnnDropoutForward
# cudnnDropoutGetStatesSize
# cudnnFindConvolutionBackwardDataAlgorithm
# cudnnFindConvolutionBackwardDataAlgorithmEx
# cudnnFindConvolutionBackwardFilterAlgorithm
# cudnnFindConvolutionBackwardFilterAlgorithmEx
# cudnnFindConvolutionForwardAlgorithm
# cudnnFindConvolutionForwardAlgorithmEx
# cudnnFusedOpsExecute
# cudnnGetBatchNormalizationBackwardExWorkspaceSize
# cudnnGetBatchNormalizationForwardTrainingExWorkspaceSize
# cudnnGetBatchNormalizationTrainingExReserveSpaceSize
# cudnnGetConvolutionBackwardDataAlgorithm_v7
# cudnnGetConvolutionBackwardDataAlgorithmMaxCount
# cudnnGetConvolutionBackwardDataWorkspaceSize
# cudnnGetConvolutionBackwardFilterAlgorithm_v7
# cudnnGetConvolutionBackwardFilterAlgorithmMaxCount
# cudnnGetConvolutionBackwardFilterWorkspaceSize
# cudnnGetConvolutionForwardAlgorithm_v7
# cudnnGetConvolutionForwardAlgorithmMaxCount
# cudnnGetConvolutionForwardWorkspaceSize
# cudnnGetCTCLossWorkspaceSize
# cudnnGetCTCLossWorkspaceSize_v8
# cudnnGetDropoutDescriptor
# cudnnGetMultiHeadAttnBuffers
# cudnnGetMultiHeadAttnWeights
# cudnnGetReductionIndicesSize
# cudnnGetReductionWorkspaceSize
# cudnnGetStream
# cudnnIm2Col
# cudnnLRNCrossChannelBackward
# cudnnLRNCrossChannelForward
# cudnnMakeFusedOpsPlan
# cudnnMultiHeadAttnBackwardData
# cudnnMultiHeadAttnBackwardWeights
# cudnnMultiHeadAttnForward
# cudnnOpTensor
# cudnnPoolingBackward
# cudnnPoolingForward
# cudnnQueryRuntimeError
# cudnnReduceTensor
# cudnnReorderFilterAndBias
# cudnnRestoreDropoutDescriptor
# cudnnScaleTensor
# cudnnSetDropoutDescriptor
# cudnnSetStream
# cudnnSetTensor
# cudnnSoftmaxBackward
# cudnnSoftmaxForward
# cudnnSpatialTfGridGeneratorBackward
# cudnnSpatialTfGridGeneratorForward
# cudnnSpatialTfSamplerBackward
# cudnnSpatialTfSamplerForward
# cudnnTransformFilter
# cudnnTransformTensor
# cudnnTransformTensorEx
# dep8: cudnnFindRNNBackwardDataAlgorithmEx
# dep8: cudnnFindRNNBackwardWeightsAlgorithmEx
# dep8: cudnnFindRNNForwardInferenceAlgorithmEx
# dep8: cudnnFindRNNForwardTrainingAlgorithmEx
# dep8: cudnnGetAlgorithmSpaceSize
# dep8: cudnnGetRNNBackwardDataAlgorithmMaxCount
# dep8: cudnnGetRNNBackwardWeightsAlgorithmMaxCount
# dep8: cudnnGetRNNDescriptor_v6
# dep8: cudnnGetRNNForwardInferenceAlgorithmMaxCount
# dep8: cudnnGetRNNForwardTrainingAlgorithmMaxCount
# dep8: cudnnGetRNNLinLayerBiasParams
# dep8: cudnnGetRNNLinLayerMatrixParams
# dep8: cudnnGetRNNParamsSize
# dep8: cudnnGetRNNProjectionLayers
# dep8: cudnnGetRNNTrainingReserveSize
# dep8: cudnnGetRNNWorkspaceSize
# dep8: cudnnRestoreAlgorithm
# dep8: cudnnRNNBackwardData
# dep8: cudnnRNNBackwardDataEx
# dep8: cudnnRNNBackwardWeights
# dep8: cudnnRNNBackwardWeightsEx
# dep8: cudnnRNNForwardInference
# dep8: cudnnRNNForwardInferenceEx
# dep8: cudnnRNNForwardTraining
# dep8: cudnnRNNForwardTrainingEx
# dep8: cudnnRNNGetClip
# dep8: cudnnRNNSetClip
# dep8: cudnnSaveAlgorithm
# dep8: cudnnSetRNNAlgorithmDescriptor
# dep8: cudnnSetRNNDescriptor_v6
# dep8: cudnnSetRNNProjectionLayers


text:
function barycenter_matrices(d, models, normalize=true; allow_shortcut=true)
    r = roots.(SpectralDistances.Continuous(), models)
    w = [Float64.(d.weight(r)) for r in r]
    realpoles = any(any(iszero ∘ imag, r) for r in r)

    realpoles = realpoles || !allow_shortcut # if we don't allow the shortcut it's the  same as if there are real poles.

    if !realpoles
        r = [r[1:end÷2] for r in r]
        w = [2w[1:end÷2] for w in w]
    end
    X = [[Float64.(real(r)'); Float64.(imag(r)')] for r in r]
    if !all(sum.(w) .≈ 1)
        if normalize
            w = s1.(w)
        else
            @warn "sum.(w) ≠ 1" sum.(w)
        end
    end

    X, w, realpoles
end

function embedding(model::AR,args...)
    embedding(roots(SpectralDistances.Continuous(), model), args...)
end

"""
    embedding([::Type{Vector}], m, [full=true])

Returns a `Vector/Matrix` containing the roots of `m`.
`full` indicates whether or not to use all poles or only one half-plane.
"""
embedding(::Type{Vector}, model, args...) = embedding(model, args...)
embedding(::Type{Matrix}, model, args...) = reshape(embedding(model, args...), :, 2)

function embedding(r::AbstractRoots, full::Bool=true)
    if full
        [real.(r); imag.(r)]
    else
        realpoles = any(any(iszero ∘ imag, r) for r in r)
        realpoles && @error("Model contained real poles, this produces weird results when used in an embedding")
        @views [real.(r[1:end÷2]); imag.(r[1:end÷2])]
    end
end

"""
    barycenter(d::OptimalTransportRootDistance, models; normalize = true, kwargs...)
    $(SIGNATURES)

**Approximately** calculate the barycenter supported on the same number of atoms as the number of poles in the models.

The solver can be selected by providing a keword argument, example: `solver=IPOT`.

Uses the algorithms from ["Fast Computation of Wasserstein Barycenters"](https://arxiv.org/pdf/1310.4375.pdf)

# Example:
```julia
models = examplemodels(10)

d = OptimalTransportRootDistance(domain=SpectralDistances.Continuous(),p=2, weight=residueweight, β=0.01)
Xe = barycenter(d, models, solver=sinkhorn_log!, uniform=true)

plot()
pzmap!.(models)
pzmap!(tf(Xe), m=:c, title="Barycenter OptimalTransportRootDistance", lab="BC")
```

# Arguments:
- `models`: vector of AR models
- `normalize`: make sure weights sum to 1
- `kwargs`: are sent to the solver
"""
function barycenter(d::OptimalTransportRootDistance,models, λ=s1(ones(length(models))); normalize=true, uniform=true, solver=sinkhorn_log!, kwargs...)
    d.p == 2 || throw(ArgumentError("p must be 2"))
    X, w, realpoles = barycenter_matrices(d, models, normalize)
    if uniform
        bc = barycenter(X, λ; uniform=true, solver=solver, β=d.β, kwargs...)
    else
        bc,w = barycenter(X,w, λ; uniform=false, solver=solver, β=d.β, kwargs...)
    end
    bc2model(bc, realpoles)

end

function bc2model(bc, realpoles)
    r1 = complex.(bc[1,:],bc[2,:])
    if realpoles
        bcr = ContinuousRoots(r1)
    else
        @assert !any(iszero ∘ imag, r1) "A real root was found in barycenter even though inputs had no real roots"
        bcr = ContinuousRoots([r1; conj.(r1)])
    end
    AR(bcr)
end

"""
    barycenter(d::EuclideanRootDistance, models::AbstractVector, [λ])

# Example:
```julia
models = examplemodels(10)

Xe = barycenter(EuclideanRootDistance(domain=SpectralDistances.Continuous(),p=2), models)

G = tf.(models)
plot()
pzmap!.(G)
pzmap!(tf(Xe), m=:c, title="Barycenter EuclideanRootDistance")
current()
```
"""
function barycenter(d::EuclideanRootDistance,models::AbstractVector;kwargs...)
    r = roots.(SpectralDistances.Continuous(), models)
    w = map(d.weight, r)
    bc = map(1:length(r[1])) do pi
        sum(w[pi]*r[pi] for (w,r) in zip(w,r))/sum(w[pi] for w in w)
    end
    AR(ContinuousRoots(bc))
end

function barycenter(d::EuclideanRootDistance,models::AbstractVector, λ;kwargs...)
    N = length(models)
    r = roots.(SpectralDistances.Continuous(), models)
    w = map(d.weight, r)
    bc = map(1:length(r[1])) do pi
        sum(λ[i]*w[pi]*r[pi] for (i,w,r) in zip(1:N,w,r))/sum(λ[i]*w[pi] for (i,w) in enumerate(w))
    end
    AR(ContinuousRoots(bc))
end

function distmat_euclidean(X::AbstractMatrix,Y::AbstractMatrix)
    C = [mean(abs2, c1-c2) for c1 in eachcol(X), c2 in eachcol(Y)]
end

function distmat_euclidean!(C,X::AbstractMatrix,Y::AbstractMatrix)
    for (j,c2) in enumerate(eachcol(Y))
        for (i,c1) in enumerate(eachcol(X))
            C[i,j] = mean(((c1,c2),) -> abs2(c1-c2), zip(c1,c2))
        end
    end
    C
end

"""
    barycenter(X::Vector{<:AbstractArray}, λ)

Calculate the weighted barycenter for point clouds in `X`.
Each `X[i]` has the shame `n_dims × n_atoms`
`λ` is the weight vector that should sum to 1.
"""
function barycenter(X::Vector{<:AbstractArray}, λ; uniform=true, solver=sinkhorn_log, kwargs...)
    N = length(X)
    n = size(X[1],2)
    w = s1(ones(n))

    # X0 = mean(X)
    ind = rand(1:length(X))
    X0 = X[ind] .- mean(X[ind],dims=2) .+ mean(mean.(X, dims=2))
    perturb!(X0,X)
    bc = alg2(X0,X,w,fill(w,N); solver=solver, weights=λ, uniform=uniform, kwargs...)[1]
end

function perturb!(X0,X)
    m = 0.1minimum(std, X)
    X0 .+= m .* randn.()
end

function barycenter(X::Vector{<:AbstractArray}, p, λ; uniform=true, solver=sinkhorn_log!, kwargs...)
    N = length(X)
    n = size(X[1],2)
    w = s1(ones(n))
    m = 0.1minimum(std, X)
    # X0 = mean(X)
    ind = rand(1:length(X))
    X0 = X[ind] .- mean(X[ind],dims=2) .+ mean(mean.(X, dims=2))
    X0 .+= m .* randn.()
    alg2(X0,X,w,p; solver=solver, weights=λ, uniform=uniform, kwargs...)
end

barycentric_weighting(X,λ,sw) = sum(λ[i].*X[i][:,sw[i]] for i in eachindex(sw))

function softmax(x)
    e = exp.(x)
    e ./= sum(e)
end

using Random

"""
    λ = barycentric_coordinates(pl, ql, p, q; options, kwargs...)

Compute the barycentric coordinates `λ` such that
sum(λᵢ W(pᵢ,q) for i in eachindex(p)) is minimized.

This function works best with the `sinkhorn_log!` solver, a large β (around 1) and small tolerance. These are set using `kwargs...`.

# Arguments:
- `pl`: Atoms in measures `p`, vector, length `n_measures`, of matrices of size `n_dims × n_atoms`
- `ql`: Atoms in measure `q`
- `p`: Measures `p`, a matrix of weight vectors, size `n_atoms × n_measures` that sums to 1
- `q`: the veight vector for measure `q`, length is `n_atoms`
- `options`: For the Optim solver. Defaults are `options = Optim.Options(store_trace=false, show_trace=false, show_every=0, iterations=20, allow_f_increases=true, time_limit=100, x_tol=1e-5, f_tol=1e-6, g_tol=1e-6, f_calls_limit=0, g_calls_limit=0)`
- `solver`: = [`sinkhorn_log!`](@ref) solver
- `tol`:    = 1e-7 tolerance
- `β`:      = 0.1 entropy regularization. This function works best with rather large regularization, hence the large default value.
- `kwargs`: these are sent to the solver algorithm.

# Example:
```julia
using SpectralDistances, ControlSystems, Optim
models = examplemodels(10)

d = OptimalTransportRootDistance(
    domain = SpectralDistances.Continuous(),
    p      = 2,
    weight = residueweight,
    β      = 0.01,
)
Xe = barycenter(d, models, solver=sinkhorn_log!)

G = tf.(models)
plot()
pzmap!.(G)
pzmap!(tf(Xe), m=:c, title="Barycenter OptimalTransportRootDistance", lab="BC")

options = Optim.Options(store_trace       = true,
                        show_trace        = false,
                        show_every        = 1,
                        iterations        = 50,
                        allow_f_increases = true,
                        time_limit        = 100,
                        x_tol             = 1e-7,
                        f_tol             = 1e-7,
                        g_tol             = 1e-7,
                        f_calls_limit     = 0,
                        g_calls_limit     = 0)

method = LBFGS()
λ = barycentric_coordinates(d, models, Xe, method,
    options = options,
    solver  = sinkhorn_log!,
    robust  = true,
    uniform = true,
    tol     = 1e-6,
)
bar(λ, title="Barycentric coorinates")

G = tf.(models)
plot()
pzmap!.(G, lab="")
pzmap!(tf(Xe), m = :c, title = "Barycenter OptimalTransportRootDistance", lab = "BC")
# It's okay if the last system dot does not match the barycenter exactly, there are limited models to choose from.
pzmap!(G[argmax(λ)], m = :c, lab = "Largest bc coord", legend = true)
```
"""
function barycentric_coordinates(pl, ql, p, q::AbstractVector{T}, method=LBFGS();
    options = Optim.Options(store_trace       = false,
                            show_trace        = false,
                            show_every        = 1,
                            iterations        = 20,
                            allow_f_increases = true,
                            time_limit        = 100,
                            x_tol             = 1e-5,
                            f_tol             = 1e-6,
                            g_tol             = 1e-6,
                            f_calls_limit     = 0,
                            g_calls_limit     = 0),
    robust = true,
    solver = sinkhorn_log!,
    tol    = 1e-7,
    β      = 0.1,
    plot   = nothing,
    kwargs...) where T

    # C = [[mean(abs2, x1-x2) for x1 in eachcol(Xi), x2 in eachcol(ql)] for Xi in pl]

    S = length(pl)
    k = length(p[1])
    # λl = 1e-8randn(S)
    # dists = map(i->sum(C[i].*IPOT(C[i],p[i],q)[1]), eachindex(p))
    # λl = -v1(dists) # Initial guess based on distances between anchors and query point
    # λl .*= 0sqrt(length(λl)) # scale so that softmax(λ) is reasonably sparse.
    # randn!(λl)
    # return softmax(λl)
    # function fg!(F,G,λl)
    #     λ = softmax(λl) # optimization done in log domain
    #     local cost = 0.0
    #     if G !== nothing
    #         cost,P,∇ℰ = sinkhorn_diff(pl,ql,p,q,C,λ; L=L, kwargs...)
    #         G .= ∇ℰ
    #     end
    #     if F !== nothing
    #         if G === nothing
    #             cost = sinkhorn_cost(pl,ql,p,q,C,λ; L=L, kwargs...)
    #         end
    #         return cost
    #     end
    # end
    # if robust
    #     res = Optim.optimize(Optim.only_fg!(fg!), λl, NelderMead(), Optim.Options(iterations=60, store_trace=false))
    #     λl = res.minimizer
    # end
    # res = Optim.optimize(Optim.only_fg!(fg!), λl, method, options)

    λl = zeros(T,S)
    C = zeros(T,k,k)
    costfun = λ -> sinkhorn_cost(pl, ql, p, q, softmax(λ);
        solver = solver,
        tol    = tol,
        β      = β,
        kwargs...)

    if plot !== nothing
        @assert S == 2 "Can only plot for two anchor measures"
        cf = x -> costfun(log.([x,1-x]))
        plot(LinRange(1e-3, 1-1e-3, 100), cf) |> display
    end
    if robust
        res = Optim.optimize(costfun, λl, ParticleSwarm(), Optim.Options(iterations=100, store_trace=false))
        λl = res.minimizer
    end
    local λh
    try
        res = Optim.optimize(costfun, λl, method, options, autodiff=:forward)
        λh = softmax(res.minimizer)
    catch err
        @error("Barycentric coordinates: optimization failed, retrying with robust options: ", err)
        res = Optim.optimize(costfun, λl, NelderMead(), options)
        λh = softmax(res.minimizer)
    end

    # cost,P,∇ℰ = sinkhorn_diff(pl,ql,p,q,C,λh; kwargs...)
    λh
end

function barycentric_coordinates(d::OptimalTransportRootDistance,models, qmodel, method=BFGS(); kwargs...)

    d.p == 2 || throw(ArgumentError("p must be 2"))
    pl, p, realpolesp = barycenter_matrices(d, models)
    ql, q, realpolesq = barycenter_matrices(d, [qmodel])
    if realpolesp != realpolesq
        pl, p, realpolesp = barycenter_matrices(d, models, allow_shortcut=false)
        ql, q, realpolesq = barycenter_matrices(d, [qmodel], allow_shortcut=false)
    end

    for pi in p
        pi ./= sum(pi)
    end
    q = q[1] ./= sum(q[1])
    @assert sum(q) ≈ 1
    @assert all(sum(p) ≈ 1 for p in p)

    if d.divergence !== nothing
        if :solver ∈ keys(kwargs)
            throw(ArgumentError("Solver choice not available for unbalanced distances, solver `sinkhorn_unbalanced` will be used."))
        end
        solver, kwargs = unbalanced_solver_closure(d, kwargs...)
    else
        solver = get(kwargs, :solver, IPOT)
    end
    λ = barycentric_coordinates(pl, ql[1], p, vec(q), method; solver=solver, β=d.β, kwargs...)
    return λ
end

function barycentric_coordinates(d::EuclideanRootDistance, models, qmodel, args...; kwargs...)
    d.p == 2 || throw(ArgumentError("p must be 2"))
    N = length(models)
    pl = embedding.(models)
    ql = embedding(qmodel)
    if d.weight != unitweight
        return _wrd_barycentric_coordinates(d,models,pl,ql)
    end
    ⅅ  = reduce(hcat, pl) # Dictionary matrix
    TotalLeastSquares.sls(Float64.(ⅅ), Float64.(ql); kwargs...)
end

function _wrd_barycentric_coordinates(d,models,pl,ql)
    N = length(models)

    R = roots.(domain(d), models)
    w = map(d.weight, R)
    w2 = [Float64.(getindex.(w, i)) for i in 1:length(w[1])]
    w2 = [w2;w2]

    P = [Float64.(getindex.(pl, i)) for i in 1:length(pl[1])]
    n = length(ql)

    function scostfun(λ)
        λ = softmax(λ)
        sum(abs2(dot(λ,Diagonal(P[j]),w2[j])/dot(λ,w2[j]) - ql[j]) for j in 1:n)
    end

    res = Optim.optimize(scostfun, zeros(N), LBFGS(m=20), Optim.Options(store_trace=false, show_trace=false, show_every=1, iterations=100, allow_f_increases=false, time_limit=100, x_tol=0, f_tol=0, g_tol=1e-8, f_calls_limit=0, g_calls_limit=0), autodiff=:forward)

    softmax(res.minimizer)
end

##
"Sum over j≠i. Internal function."
function ∑jni(X,i,S,k)
    s = zero(X[1][:,1])
    d = size(X[1],1)
    @inbounds for j in eachindex(X)
        j == i && continue
        for l in 1:d
            s[l] += X[j][l,S[j][k]]
        end
    end
    s
end

"""
    ISA(X, w = nothing; iters = 100, printerval = typemax(Int))

Iterative swapping algorithm from "On the Computation of Wasserstein barycenters", Giovanni Puccetti et al.

# Arguments:
- `X`: vector of d×k matrices where d is dimension and k number of atoms
- `w`: weights. See the files `test_barycenter.jl` for different uses.
- `iters`: maximum number of iterations
- `printerval`: print this often
"""
function ISA(X, w=nothing; iters=100, printerval = typemax(Int))
    n = length(X)
    d,k = size(X[1])

    if w !== nothing
        # X = deepcopy(X)
        X = map(eachindex(X)) do i # This does both copy and handles weird input types
            X[i] .* w[i] # This should work for both w[i] scalar and vector
        end
    end

    fudgefactor = 1.0

    σ = [collect(1:k) for _ in 1:n] # let σᵢ = Id, 1 ≤ i ≤ n.
    σ′ = deepcopy(σ)
    @inbounds for iter = 1:iters
        swaps = 0
        for i = 1:n
            σᵢ = σ[i]
            σᵢ′ = σ′[i]
            for k₁ = 1:k-1, k₂ = k₁+1:k
                Xik1 = @view X[i][:,σᵢ[k₁]]
                Xik2 = @view X[i][:,σᵢ[k₂]]
                if dot(Xik1, ∑jni(X,i,σ,k₁)) + dot(Xik2, ∑jni(X,i,σ,k₂)) < (dot(Xik2, ∑jni(X,i,σ,k₁)) + dot(Xik1, ∑jni(X,i,σ,k₂)))*fudgefactor
                    σᵢ′[k₁],σᵢ′[k₂] = σᵢ[k₂],σᵢ[k₁]
                    swaps += 1
                end
            end
        end
        iter % printerval == 0 && @show iter, swaps
        swaps == 0 && (return σ)
        σ = deepcopy(σ′) # Update assignment
        fudgefactor *= 1-1/iters^3
    end
    σ
end

## ======================
"""
    alg1(X, Y, â, b; β = 1, printerval = typemax(Int), tol = 1.0e-5, iters = 10000, solver = IPOT)

Algorithm 1 from ["Fast Computation of Wasserstein Barycenters"](https://arxiv.org/pdf/1310.4375.pdf) Notation is the same as in the paper.

# Arguments:
- `X`: Initial guess for barycenter support points
- `Y`: Support points for measures to calc barycenter of
- `a`: initial guess of barycenter weights
- `b`: Weigts of measures in Y
- `β`: Reg param, higher is more reg (inverse of λ in paper)
- `printerval`: DESCRIPTION
- `tol`: DESCRIPTION
- `iters`: DESCRIPTION
- `solver`: any of [`IPOT`](@ref) (default), [`sinkhorn`](@ref), [`sinkhorn_log`](@ref)
"""
function alg1(X,Y,â::AbstractVector{T},b;β=1, printerval=typemax(Int), tol=1e-5, iters=10000, solver=IPOT, weights=nothing)::Vector{T} where T
    N = length(Y)
    â = copy(â)
    a = copy(â)
    ã = copy(â)
    t0 = 1
    t = 0
    𝛂 = similar(a, length(a), N)
    Mth = [distmat_euclidean(X,Y[1]) for i in 1:Threads.nthreads()]
    for outer t = 1:iters
        B = (t0+t)/2
        a .= (1-inv(B)).*â .+ inv(B).*ã
        @sync for i in 1:N
            Threads.@spawn begin
                M = distmat_euclidean!(Mth[Threads.threadid()], X,Y[i])
                ai = solver(M,a,b[i]; iters=50000, β=β, tol=tol)[2]
                if !all(isfinite, a)
                    @warn "Got nan in inner sinkhorn alg 1, increasing precision"
                    ai = solver(M,big.(â),big.(b[i]); iters=50000, β=β, tol=tol)[2]
                    ai = eltype(â).(ai)
                end
                scale!(ai, i, weights)
                𝛂[:,i] .= ai
            end
        end

        ã .= ã .* exp.((-t0*B).*vec(mean(𝛂, dims=2)))
        ã ./= sum(ã)
        aerr = sum(abs2,â-ã)
        t % printerval == 0 && @info "Sinkhorn alg1:  iter: $t, aerr: $aerr"
        â .= (1-inv(B)).*â .+ inv(B).*ã
        â ./ sum(â)
        if aerr < tol
            t > printerval && @info "Sinkhorn alg1 done at iter $t"
            return â
        end
    end
    t > printerval && @info "Sinkhorn alg1 maximum number of iterations reached: $iters"
    â
end

"""
    alg2(X, Y, a, b;
            β          = 1/10,
            θ          = 0.5,
            printerval = typemax(Int),
            tol        = 1.0e-6,
            innertol   = 1.0e-5,
            iters      = 500,
            inneriters = 1000,
            atol       = 1.0e-32,
            solver     = IPOT,
            γ          = 1,
        )

Algorithm 2 from ["Fast Computation of Wasserstein Barycenters"](https://arxiv.org/pdf/1310.4375.pdf) Notation is the same as in the paper.

# Arguments:
- `X`: Initial guess for barycenter support points
- `Y`: Support points for measures to calc barycenter of
- `a`: initial guess of barycenter weights
- `b`: Weigts of measures in Y
- `β`: Reg param, higher is more reg
- `θ`: step size ∈ [0,1]
- `printerval`: print this often
- `tol`: outer tolerance
- `innertol`: inner tolerance
- `solver`: any of [`IPOT`](@ref) (default), [`sinkhorn`](@ref), [`sinkhorn_log`](@ref)
- `γ`: Sparsity parameter, if <1, encourage a uniform weight vector, if >1, do the opposite. Kind of like the inverse of α in the Dirichlet distribution.
"""
function alg2(X,Y,a,b; β = 1/10, θ = 0.5, printerval=typemax(Int), tol=1e-6, innertol=1e-4, iters=500, inneriters=10000, atol=1e-32, solver=IPOT, γ=0.0, weights=nothing, uniform=false)
    uniform || @warn("This function is known to be buggy when not enforcing uniform weights", maxlog=10)
    N  = length(Y)
    a  = copy(a)
    ao = copy(a)
    X  = copy(X)
    if weights !== nothing && eltype(weights) != eltype(X)
        X = convert(Matrix{eltype(weights)}, X)
    end
    Xo = copy(X)
    weights === nothing && (weights = fill(1/N, N))
    fill!(ao, 1/length(ao))
    for iter = 1:iters
        uniform || (a = alg1(X,Y,ao,b,β=β, printerval=printerval, tol=innertol, iters=inneriters, solver=solver, weights=weights))
        if γ > 0 && γ != 1
            a .= softmax(γ.*log.(a))
        else
            a = s1(a)
        end
        YT = sum(1:N) do i
            M = distmat_euclidean(X,Y[i])
            T,_ = solver(M,a,b[i]; iters=inneriters, β=β, tol=innertol)
            @assert !any(isnan, T) "Got nan in sinkhorn alg 2"
            scale!(Y[i]*T', i, weights)
        end

        X .= (1-θ).*X .+ θ.*(YT ./ (a' .+ eps()))
        aerr = mean(abs2, a-ao)
        xerr = mean(abs2, X-Xo)
        iter % printerval == 0 && @info "Sinkhorn alg2:  iter: $iter, aerr: $aerr, xerr: $xerr"
        if xerr < tol
            iter > printerval && @info "Sinkhorn alg2 done at iter $iter"
            return X,(a./=sum(a))
        end
        copyto!(ao,a)
        copyto!(Xo,X)
        ao ./= sum(ao)
        # θ *= 0.999
    end
    iters > printerval && @info "Sinkhorn alg2 maximum number of iterations reached: $iters"
    X,a
end

scale!(x,_,::Nothing) = x
scale!(x::AbstractArray{T},i,w::AbstractArray{T}) where T = (x .*= w[i])
scale!(x,i,w::AbstractArray) = (x * w[i]) # for dual numbers etc.

struct BCWorkspace{T,TK}
    K::TK
    KV::Array{T,3}
    U::Array{T,3}
    b::Array{T,2}
    bold::Array{T,2}
    S::Array{T,2}
end

function BCWorkspace(A, β)
    N    = length(A)
    m, n = size(A[1])
    T    = eltype(eltype(A))
    β    = T(β)

    b    = zeros(T, m, n)
    bold = zeros(T, m, n)
    U    = ones(T, m, n, N)
    KV   = ones(T, m, n, N)
    S1   = Matrix{T}(undef, m, n)
    S2   = Matrix{T}(undef, m, n)

    xi1 = Matrix{T}(undef, m, m)
    xi2 = Matrix{T}(undef, n, n)
    _initialize_conv_op!(xi1, xi2, β)

    function K(u,x)
        # xi1 * x * xi2
        mul!(S2,x,xi2)
        mul!(u,xi1,S2)
    end
    BCWorkspace(K, KV, U, b, bold, S1)
end

"""
    barycenter_convolutional(A, [λ]; β = 0.01, iters = 1000, tol = 1e-9, ϵ = 1e-90, verbose = false)

Convolutional barycenters.

´β` is the regularization and `λ` (optional) are the weights (barycentric coordinates). To reuse allocated space between successive calls, use the "workspace" method from the example below. `ϵ` is a truncation parameter for numerical stability.
```julia
a1      = zeros(10, 10)
a1[2,2] = 1
a2      = zeros(10, 10)
a2[6,6] = 1
A       = [a1,a2]
β       = 0.01
λ       = [0.5, 0.5] # Barycentric coordinates, must sum to 1
w       = BCWorkspace(A, β)
b       = barycenter_convolutional(w,A,λ)
plot(heatmap(a1), heatmap(a2), heatmap(b))
```

Ref: J. Solomon, F. de Goes, G. Peyré, M. Cuturi, A. Butscher, A. Nguyen, T. Du, L. Guibas. Convolutional Wasserstein Distances: Efficient Optimal Transportation on Geometric Domains.  2015 https://people.csail.mit.edu/jsolomon/assets/convolutional_w2.compressed.pdf
"""
function barycenter_convolutional(
    w::BCWorkspace{T,TK},
    A::AbstractVector{<:AbstractMatrix},
    λ::Union{Vector{<:Union{Float64,Float32}}, <: Fill} = Fill(1 / length(models), length(models));
    iters   = 1000,
    tol     = 1e-6,
    ϵ       = 1e-90,
    verbose = false,
) where {T,TK}

    length(A) < 2 && return A[]
    @fastmath sum(A[1]) ≈ sum(A[2]) || @warn "Input matrices do not appear to have the same mass (sum)"
    0.99 < sum(λ) < 1.01 || throw(ArgumentError("sum of barycentric coordinates λ was $(sum(λ)) but should be 1"))
    N = length(A)
    @unpack K, KV, U, b, bold, S = w
    iter = 0
    err = one(T)

    @views while err > tol && iter < iters
        copyto!(bold, b)
        iter = iter + 1
        b .= 0
        for r = 1:N # TODO: if S storage is expanded, this loop can be split into several and @avx or @tullio can be put on some. Maybe even K can be manually inlined here to make everythin avx-able. Can also be run on GPU easily.
            K(S, U[:, :, r])
            @avx S .= A[r] ./ max.(ϵ, S)
            K(KV[:, :, r], S)
            @avx @. b += λ[r] * log(max(ϵ, U[:, :, r] * KV[:, :, r]))
        end
        @avx b .= exp.(b)
        for r = 1:N
            @avx @. U[:, :, r] = b / max(ϵ, KV[:, :, r])
        end
        @avx b ./= sum(b)

        if iter % 10 == 1
            @fastmath err = sum(abs(bold - b) for (bold, b) in zip(bold, b))
            verbose && @info "Sinkhorn conv barycenters: iter = $iter, error = $err"
        end
    end
    b#, cost
end

# function barycenter_convolutional_diff(
#     A::AbstractVector{<:AbstractMatrix},
#     λ::AbstractVector{T};
#     β       = 0.01,
#     iters   = 1000,
#     tol     = 1e-6,
#     ϵ       = 1e-90,
#     verbose = false,
# ) where {T}
#
#     length(A) < 2 && return A[]
#     @fastmath sum(A[1]) ≈ sum(A[2]) ||
#               @warn "Input matrices do not appear to have the same mass (sum)"
#     0.99 < sum(λ) < 1.01 ||
#         throw(ArgumentError("sum of barycentric coordinates λ was $(sum(λ)) but should be 1"))
#
#     N    = length(A)
#     err  = one(T)
#     m, n = size(A[1])
#     b    = similar(λ, m, n) .= 0
#     # S  = similar(b)
#     U    = similar(λ, m, n, N) .= 1
#     KV   = similar(λ, m, n, N) .= 1
#
#     xi1  = Matrix{T}(undef, m, m)
#     xi2  = Matrix{T}(undef, n, n)
#     _initialize_conv_op!(xi1, xi2, β)
#
#
#
#     for i = 1:iters
#         b .= 0
#         for r = 1:N
#             S = xi1 * U[:, :, r] * xi2
#             # mul!(S,U,xi2)
#             # mul!(V,xi1,S) # TODO: this is by far the most expensive operation
#             @. S = A[r] / max(ϵ, S)
#             KV[:, :, r] = xi1 * S * xi2
#             # mul!(S,V,xi2)
#             # mul!(U,xi1,S)
#             @. b += λ[r] * log(max(ϵ, U[:, :, r] * KV[:, :, r]))
#         end
#         @. b = exp(b)
#         for r = 1:N
#             @. U[:, :, r] = b / max(ϵ, KV[:, :, r])
#         end
#
#     end
#     b
# end

function barycenter_convolutional(
    A::AbstractVector{<:AbstractMatrix},
    λ = Fill(1 / length(A), length(A));
    β = 0.001,
    kwargs...,
)
    w = BCWorkspace(A, β)
    barycenter_convolutional(w, A, λ; kwargs...)
end

"""
    barycenter_convolutional(models::Vector{<:DSP.Periodograms.TFR}, λ = Fill(1 / length(models), length(models)); dynamic_floor = default_dynamic_floor(models), kwargs...)

Covenience function for the calculation of spectrograms. This function transforms the spectrograms to log-power and adjusts the floor to `dynamic_floor`, followed by a normalization to sum to 1.

This function will be called if [`barycenter`](@ref) is called with [`ConvOptimalTransportDistance`](@ref) as first argument.

# Arguments:
- `dynamic_floor`: Sets the floor of the spectrogram in log-domain, i.e., all values below this will be truncated. The default value is based on a quantile of the spectrogram powers. If your spectrograms are mostly low entropy, you can try to increase this number to get sharper results.
- `kwargs`: Same as for the base method

# Example:
```julia
using SpectralDistances, DSP, Plots
N     = 24_000
t     = 1:N
f     = range(0.8, stop=1.2, length=N)
y1    = sin.(t .* f) .+ 0.1 .* randn.()
y2    = sin.(t .* reverse(f .+ 0.5)) .+ 0.1 .* randn.()
S1,S2 = spectrogram.((y1,y2), 1024)

A = [S1,S2]
β = 0.0001     # Regularization parameter (higher implies more smoothing and a faster, more stable solution)
λ = [0.5, 0.5] # Barycentric coordinates (must sum to 1)
B = barycenter_convolutional(A, β=β, tol=1e-6, iters=200, ϵ=1e-100, dynamic_floor=-2)
plot(plot(S1, title="S1"), plot(B, title="Barycenter"), plot(S2, title="S2"), layout=(1,3), colorbar=false)
```
"""
function barycenter_convolutional(
    models::Vector{<:DSP.Periodograms.TFR},
    λ = Fill(1 / length(models), length(models));
    dynamic_floor = default_dynamic_floor(models),
    kwargs...,
)

    A  = normalize_spectrogram.(models, dynamic_floor)
    ms = mean(sum, A)
    sumnorm = x -> x .* (ms/sum(x))
    A .= sumnorm.(A)
    b  = barycenter_convolutional(A, λ; kwargs...)
    B  = deepcopy(models[1])
    Bp = power(B)
    Bp .= exp.(b .+ dynamic_floor)
    B
end

function barycenter(d::ConvOptimalTransportDistance, args...; kwargs...)
    @assert d.invariant_axis == 0 "Invariant axis not supported"
    barycenter_convolutional( args...; β = d.β, kwargs...)
end

function barycenter(d::ConvOptimalTransportDistance, A::Vector{<:DSP.Periodograms.TFR}, args...; kwargs...)
    @assert d.invariant_axis == 0 "Invariant axis not supported"
    barycenter_convolutional(
        A, args...;
        dynamic_floor = isnan(d.dynamic_floor) ? default_dynamic_floor(A) : d.dynamic_floor,
        β = d.β,
        kwargs...,
    )
end

# function barycenter(d::OptimalTransportRootDistance,models; normalize=true, kwargs...)
#     X, w, realpoles = barycenter_matrices(d, models, normalize)
#     S = ISA(X; kwargs...)
#     bc = barycentric_weighting(X,w,S)
#
#     r1 = complex.(bc[1,:],bc[2,:])
#     if realpoles
#         bcr = ContinuousRoots(r1)
#     else
#         @assert !any(iszero ∘ imag, r1) "A real root was found in barycenter even though inputs had no real roots"
#         bcr = ContinuousRoots([r1; conj.(r1)])
#     end
#     AR(bcr)
# end
# function barycenter(d::OptimalTransportRootDistance,models,w; kwargs...)
#     r = roots.(SpectralDistances.Continuous(), models)
#     realpoles = any(any(iszero ∘ imag, r) for r in r)
#
#     if !realpoles
#         r = [r[1:end÷2] for r in r]
#     end
#     X = [[real(r)'; imag(r)'] for r in r]
#
#     S = ISA(X, w; kwargs...)
#     X̂ = [X[i][:,S[i]] for i in eachindex(X)]
#     bc = sum(w.*X̂)
#     r1 = complex.(bc[1,:],bc[2,:])
#     if realpoles
#         bcr = ContinuousRoots(r1)
#     else
#         @assert !any(iszero ∘ imag, r1) "A real root was found in barycenter even though inputs had no real roots"
#         bcr = ContinuousRoots([r1; conj.(r1)])
#     end
#
#     AR(bcr)
# end
# the method below does not work very well due to likely bug in ISA
# function barycenter(X::Vector{<:AbstractArray}, λ; iters=100, kwargs...)
#     sw = ISA(X, λ; iters=iters, kwargs...)
#     barycentric_weighting(X,λ,sw)
# end
# function barycentric_weighting2(X,λ,sw)
#     sum(λ[sw[i]]'.*X[i][:,sw[i]] for i in eachindex(sw))
# end # This one was definitely bad
## ==========================================================
# function simplex_ls(ⅅ, ql; iters=2000, verbose=false, kwargs...)
#
#     N = length(ql)
#     α = 1.0
#     ε = 0.1/N
#     @show λ = max.(ⅅ\ql .+ 0.01 .* randn.(), 0)
#     λ ./= sum(λ)
#     λo = copy(λ)
#     g = similar(λ)
#     local ng, err
#     err = 0.0
#     for iter = 1:iters # Do a few projected gradient iterations with soft thresholding
#         g .= ⅅ'*(ql-ⅅ*λ) # This is the negative gradient so + below
#         ng = norm(g)
#         ng < 1e-10 && break
#         λ .+= α.*g
#         λ .-= minimum(λ)
#         λ .= soft_th.(λ, ε)
#         λ ./= sum(λ)
#         α *= 0.999
#         err = norm(λ-λo)
#         verbose && @info "Iter $iter norm(g): $ng norm(λ-λo): $err"
#         err < 1e-5 && break
#         λo .= λ
#     end
#     verbose &&  @info "Converged norm(g): $ng norm(λ-λo): $err"
#     λ
# end

struct BCCWorkspace{T}
    w::Vector{T}
    b::Vector{Array{T,3}}
    r::Array{T,3}
    φ::Vector{Array{T,3}}
    C::Matrix{T}
    C2::Matrix{T}
    C3::Matrix{T}
    C4::Matrix{T}
    C5::Matrix{T}
    S2::Matrix{T}
    xi1::Matrix{T}
    xi2::Matrix{T}
    scw::SCWorkspace{T}
end

"""
    BCCWorkspace(X::Vector{<:AbstractMatrix{T}}, L, β) where T

Create a workspace cache for [`barycentric_coordinates`](@ref) with the convolutional distance.

# Arguments:
- `X`: Input matrices
- `L`: Number of iterations
- `β`: Regularization factor
"""
function BCCWorkspace(p::Vector{<:AbstractMatrix{T}}, L, β) where T
    N   = length(p)
    S   = length(p)
    m,n = size(p[1])
    w   = zeros(T,S)
    b   = [fill(1/N, m, n, S) for _ in 0:L+1]
    r   = zeros(T,m,n,S)
    φ   = [Array{T}(undef,m,n,S) for _ in 1:L]
    C   = zeros(T,m,n)
    C2  = zeros(T,m,n)
    C3  = zeros(T,m,n)
    C4  = zeros(T,m,n)
    C5  = zeros(T,m,n)
    S2  = zeros(T,m,n)
    scw = SCWorkspace(p[1], p[1], β)

    xi1 = Matrix{T}(undef, m, m)
    xi2 = Matrix{T}(undef, n, n)
    _initialize_conv_op!(xi1, xi2, β)
    BCCWorkspace{T}(w,b,r,φ,C,C2,C3,C4,C5,S2,xi1,xi2,scw)
end

"""
    cost, B, grad = sinkhorn_convolutional_diff(workspace::BCCWorkspace{T}, X, q::AbstractMatrix{T}, λ::AbstractVector; β = 0.01) where T

Returns cost, barycenter and gradient of `λ`. Called from within [`barycentric_coordinates`](@ref)`(
    d::ConvOptimalTransportDistance,
    X::Vector{<:AbstractMatrix},
    q::AbstractMatrix)`

# Arguments:
- `workspace`: See [`BCCWorkspace`](@ref)
- `X`: Input matrices (anchors)
- `q`: Query matrix
- `λ`: barycentric coordinates
- `β`: reg parameter
"""
function sinkhorn_convolutional_diff(workspace::BCCWorkspace{T}, p, q::AbstractMatrix{T}, λ::AbstractVector; β=0.01) where T

    @unpack w,b,r,φ,C,C2,C3,C4,C5,S2,xi1,xi2,scw = workspace
    w .= 0
    r .= 0

    L = length(φ)
    N = length(p)
    S = length(p)
    m,n = size(p[1])

    function K(u,x)
        # xi1 * x * xi2
        mul!(S2,x,xi2)
        mul!(u,xi1,S2)
    end

    local P
    @views for l = 1:L
        for s in 1:S
            K(C,b[l][:,:,s])
            @avx @. C = p[s] / C
            K(φ[l][:,:,s], C)
        end
        P = dropdims(prod(φ[l].^reshape(λ,1,1,:), dims=3), dims=3)
        @avx b[l+1] .= P ./ φ[l]
    end

    cost,a,_ = sinkhorn_convolutional(scw, P, q; β=β)

    @avx ∇W = @. (a = β * a)
    # ∇W = bb
    @avx g = (∇W .= ∇W .* P)

    for l = L:-1:1
        @views for s in 1:S
            @avx S2 .= log.(φ[l][:,:,s])
            w[s] += dot(S2, g)
            K(C,b[l][:,:,s])
            @avx @. C = abs2(C)
            @avx @. C4 = (λ[s]* g - r[:,:,s]) / φ[l][:,:,s]
            K(C2,C4)
            @avx @. C2 = C2 * p[s] / C
            K(C3,C2)
            @avx @. r[:,:,s] = -C3 * b[l][:,:,s]
        end
        g = dropdims(sum(r, dims=3), dims=3)
    end
    cost, P,w
end

"""
    sinkhorn_convolutional_diff(p::Vector, q::AbstractMatrix, λ::AbstractVector; β = 0.01, L = 32, kwargs...)
"""
function sinkhorn_convolutional_diff(p::Vector, q::AbstractMatrix, λ::AbstractVector; β=0.01, L = 32, kwargs...)
    w = BCCWorkspace(p, L, β)
    sinkhorn_convolutional_diff(w, p, q, λ; β=β, kwargs...)
end

"""
    barycentric_coordinates(d::ConvOptimalTransportDistance, X::Vector{<:AbstractMatrix}, q::AbstractMatrix; method = LBFGS(), kwargs...)

Calculate the barycentric coordinates of a vector of matrices `X` using the convolutional method.

# Arguments:
- `q`: Query matrix
- `method`: The optimizer from Optim
- `kwargs`: Are sent to [`sinkhorn_convolutional`](@ref)

## Optim options
The default options are
```
options = Optim.Options(
        store_trace       = true,
        show_trace        = false,
        show_every        = 1,
        iterations        = 10,
        allow_f_increases = false,
        time_limit        = 150,
        x_tol             = 1e-3,
        f_tol             = 1e-3,
        g_tol             = 1e-4,
    )
```
"""
function barycentric_coordinates(
    d::ConvOptimalTransportDistance,
    X::Vector{<:AbstractMatrix},
    q::AbstractMatrix;
    L = 40,
    method  = LBFGS(),
    options = Optim.Options(
        store_trace       = true,
        show_trace        = false,
        show_every        = 1,
        iterations        = 10,
        allow_f_increases = false,
        time_limit        = 150,
        x_tol             = 1e-3,
        f_tol             = 1e-3,
        g_tol             = 1e-4,
    ),
    kwargs...,
)
    workspace = BCCWorkspace(X, L, d.β)
    @assert d.invariant_axis == 0 "Invariant axis not supported"

    function fg!(F, G, λ)
        λ = softmax(λ)
        cost, B, g = sinkhorn_convolutional_diff(workspace, X, q, λ; β=d.β)

        if G != nothing
            G .= λ .* (g .- dot(g, λ)) # Chain rule for softmax
            # G .= g
        end
        return cost
    end

    λ = zeros(length(X))

    res = Optim.optimize(Optim.only_fg!(fg!), λ, method, options)
    softmax(res.minimizer), res
end


text:
using Measurements
using Printf

printstyled("\nExperimento 1: Capacidade térmica\n", color =:blue);

# em gramas
Δm = 0.01;
mcopo = 51.82;
maq = 122.07;
maf = 125.02;

# em graus celsius
ΔT = 0.1;
taf = 12.4;
taq = 53.2;
tf = 31.4;

# calor especifico da água
ca = 1;
@printf("%.3f · %.3f ( %.3f - %.3f) / (%.3f - %.3f) - %.3f · %.3f \n", maq,ca,taq,tf,tf,taf,maf,ca);
C = maq * ca * ( (taq - tf) / (tf - taf) ) - maf * ca;

ΔC = ((Δm * taq + ΔT * maq) * (tf - taf) + (2 * ΔT) * (maq * taq) ) / (tf - taf)^2

println("\nCapacidade térmica do calorimetro");
@printf("Ccal = %.2f ± %.2f \n", C, ΔC);

# em grama
mcopo = 51.82 ±  Δm;
cAluminio = 0.218;
Ccopo =  mcopo * cAluminio;

println("\nCapacidade térmica do copo");
@printf("Ccopo = %.3f ± %0.3f \n", Ccopo.val, Ccopo.err);

println("\nRelação de equivalência");
@printf("| x1 - x2 | < 2 · (σ1 + σ2) \n");
@printf("| %.3f - %.3f| < 2 · (%.3f + %.3f) \n", C,Ccopo.val,ΔC,Ccopo.err);
println(abs(C - Ccopo.val) < 2 * (ΔC + Ccopo.err));

printstyled("\nExperimento 2: Calor específico de um metal  ", color =:blue);

# em gramas
mobjeto = 205.75 ± Δm;
ma = 216.07 ± Δm;

# em graus celsius
ta = 13.1 ± ΔT;
to = 97.1 ± ΔT;
tf = 19.2 ± ΔT;

co = ((ma * ca + C ) * (ta - tf)) / (mobjeto * (tf - to));

println("\nCalor específico de um metal");
@printf("Co = %.4f ± %.4f \n", co.val, co.err);

printstyled("\nExperimento 3: Calor latente de condensação da água\n",color=:blue);

ta = 8.8  ± ΔT;
tf = 71.7  ± ΔT;
tc = 100  ± ΔT;

ma = 167.479  ± Δm;
mc = 187.02 - ma  ± Δm;

@printf("Massa condensada = %.3f", mc.val);

L = (((ma + C) * (ta - tf)) / mc) + (tc - tf);

println("\nCalor latente");
@printf("Lc = %.3f ± %.3f",L.val, L.err);

max_stars_repo_path: Fisica experimental/02/05/calcs.jl

max_issues_repo_path: Fisica experimental/02/05/calcs.jl

max_forks_repo_path: Fisica experimental/02/05/calcs.jl


text:
using Flux, BSON, Random, Plots; pyplot()
Random.seed!(0)

latentDim = 100
outputX, outputY = 6, 3

gen =  Chain(Dense(latentDim,7*7*256),BatchNorm(7*7*256,relu),
        x->reshape(x,7,7,256,:),ConvTranspose((5,5),256=>128;stride=1,pad=2),
        BatchNorm(128,relu),ConvTranspose((4,4),128=>64;stride=2,pad=1),
        BatchNorm(64,relu),ConvTranspose((4,4),64=>1,tanh;stride=2,pad=1))

cd(@__DIR__)
BSON.@load "../data/mnistGAN40.bson" genParams
Flux.loadparams!(gen, genParams)

fixedNoise = [randn(latentDim, 1) for _ in 1:outputX*outputY]
fakeImages = @. gen(fixedNoise)
imageArray = permutedims(dropdims(reduce(vcat,
            reduce.(hcat, Iterators.partition(fakeImages, outputY)));
             dims=(3, 4)), (2, 1))

heatmap(imageArray, yflip = true, color = :Greys,
        size = (300,150), legend=false, ticks=false)

max_stars_repo_name: Yoshinobu-Ishizaki/StatsWithJuliaBook

max_issues_repo_name: Yoshinobu-Ishizaki/StatsWithJuliaBook

max_forks_repo_name: Yoshinobu-Ishizaki/StatsWithJuliaBook