text:
function βdf(data::NCCSData; minPointsForβ::Int = MIN_POINTS_FOR_BETA,
  path::String=WORKING_PATH, charitytypes::Symbol = :all)
  local βdf::DataFrame = DataFrame(ein = unique(data.df[:ein]))
  local Nβ::Int = size(βdf,1)
  local reg::FMLM

  local plotNames::Vector{String} = Vector{String}()
  local plots::Vector{PlotContainer} = Vector{PlotContainer}()

  #NOTE: DELETE THE BELOW

  (charitytypes ≠ :all) && (βdf = βdf[βdf[:charitytype] .== charitytypes, :])

  #initialize the rest of the dataframe
  βdf[:b0return] = Vector{MFloat64}(missing,Nβ)
  βdf[:b1return] = similar(βdf[:b0return])

  βdf[:b0contributions] = similar(βdf[:b0return])
  βdf[:b1contributions] = similar(βdf[:b0return])

  βdf[:b0contribonret] = similar(βdf[:b0return])
  βdf[:b1contribonret] = similar(βdf[:b0return])

  βdf[:b0expenses] = similar(βdf[:b0return])
  βdf[:b1expenses] = similar(βdf[:b0return])

  βdf[:avgwealth] = similar(βdf[:b0return])
  βdf[:N] = Vector{MInt}(missing, Nβ)
  βdf[:charitytype] = Vector{MSymbol}(missing, Nβ)

  #index the dataframe
  βdfindex::Dict = Dict(βdf[i,:ein] => βdf[i,:] for i::Int  ∈ 1:Nβ)

  local YSpecs::Vector{Symbol} = Vector{Symbol}() # holds the LHS
  local XSpecs::Vector{FMExpr} = Vector{FMExpr}() #holds the RHS
  local XNames::Vector{Vector{Symbol}} = Vector{Vector{Symbol}}() # holds the names of the x variables

  #regress excess returns on excess s&p
  push!(XSpecs, Meta.parse("sp500ret"))
  push!(XNames, [:intercept; :sp500ret])
  push!(YSpecs, :excess)

  #regress contributions on excess s&p
  push!(XSpecs, Meta.parse("sp500ret"))
  push!(XNames, [:intercept; :lagsp500ret])
  push!(YSpecs, :pcontributions)

  #regress contributions on excess returns
  push!(XSpecs, Meta.parse("excess"))
  push!(XNames, [:intercept; :excess])
  push!(YSpecs, :pcontributions)

  #regress contributions on excess returns
  push!(XSpecs, Meta.parse("sp500ret"))
  push!(XNames, [:intercept; :sp500ret])
  push!(YSpecs, :pexpenses)

  println("beginning beta computations")
  #print(data.df[50000:50100, [:ein, :datefiscal, :return, :sp500ret]])
  for subdf ∈ groupby(data.df, :ein) #compute the beta for each firm
    N::Int = sum(completecases(subdf[[:excess, :sp500ret, :pcontributions, :pexpenses]]))
    βrow = βdfindex[subdf[1,:ein]]
    βrow[:N] = N
    βrow[:charitytype] = subdf[1,:charitytype]

    if N ≥ minPointsForβ
      local models::Vector{FMLM} = Vector{FMLM}() #holds the regression models

      #####run the regressions
      try
        for i ∈ 1:length(XSpecs)
          push!(models, FMLM(subdf, XSpecs[i], YSpecs[i],
            XNames = XNames[i], YName=YSpecs[i]))
        end
      catch err
        println(subdf[[:ein, :datefiscal, :excess, :sp500ret]])
        error("Regression failed. $err")
      end

      #store the values
      βrow[:b0return] = models[1].β[1]
      βrow[:b1return] = models[1].β[2]
      βrow[:b0contributions] = models[2].β[1]
      βrow[:b1contributions] = models[2].β[2]
      βrow[:b0contribonret] = models[3].β[1]
      βrow[:b1contribonret] = models[3].β[2]
      βrow[:b0expenses] = models[4].β[1]
      βrow[:b1expenses] = models[4].β[2]
      βrow[:avgwealth] =
        mean(skipmissing((ismissing).(subdf[:excess]) .* subdf[:adjnetassets]))
    end

  end

  #histogram for betas of returns
  βdf = βdf[completecases(βdf[[:b1return, :b0return]]),:]
  push!(plotNames, "return_beta_hist_$charitytypes")
  push!(plots, plot(βdf[(f->abs(f)<1.1).(βdf[:b1return]),:], x=:b1return,
    Guide.title("histogram of return betas"), Geom.histogram(bincount=25)))

  #histogram for betas of the contributions
  push!(plotNames, "return_beta_contrib_$charitytypes")
  βdf = βdf[completecases(βdf[[:b1contributions, :b0contributions]]),:]
  push!(plots, plot(βdf[(f->abs(f)<1.1).(βdf[:b1contributions]),:], x=:b1contributions,
    Guide.title("histogram of contribution betas"), Geom.histogram(bincount=25)))

  #histogram for betas of the contribonret
  βdf = βdf[completecases(βdf[[:b1contribonret, :b0contribonret]]),:]
  push!(plotNames, "return_beta_ret_on_contrib_$charitytypes")
  push!(plots,  plot(βdf[(f->abs(f)<1.1).(βdf[:b1contribonret]),:], x=:b1contribonret,
    Guide.title("histogram of return on contribution betas"), Geom.histogram(bincount=25)))

    #histogram for betas of expenses
  βdf = βdf[completecases(βdf[[:b1expenses, :b0contribonret]]),:]
  push!(plotNames, "return_beta_expense_$charitytypes")
  push!(plots, plot(βdf[(f->abs(f)<1.1).(βdf[:b1expenses]),:], x=:b1expenses,
      Guide.title("histogram of expense betas"), Geom.histogram(bincount=25)))

  println("return beta computations complete with median $(median(βdf[:b1return])),
    mean $(mean(βdf[:b1return])),
    weighted mean $(mean(βdf[:b1return], Weights(Vector{Float64}(βdf[:avgwealth])))),
    sd $(std(βdf[:b1return])), and se $(std(βdf[:b1return])/size(βdf,1)^0.5)")

  println("contribution on return beta computations complete with median $(median(βdf[:b1contributions])),
    mean $(mean(βdf[:b1contributions])),
    weighted mean $(mean(βdf[:b1contributions], Weights(Vector{Float64}(βdf[:avgwealth])))),
    sd $(std(βdf[:b1contributions])), and se $(std(βdf[:b1contributions])/size(βdf,1)^0.5)")

  println("contribution on return beta computations complete with median $(median(βdf[:b1contribonret])),
    mean $(mean(βdf[:b1contribonret])),
    weighted mean $(mean(βdf[:b1contribonret], Weights(Vector{Float64}(βdf[:avgwealth])))),
    sd $(std(βdf[:b1contribonret])), and se $(std(βdf[:b1contribonret])/size(βdf,1)^0.5)")

  println("expense beta computations complete with median $(median(βdf[:b1expenses])),
    mean $(mean(βdf[:b1expenses])),
    weighted mean $(mean(βdf[:b1expenses], Weights(Vector{Float64}(βdf[:avgwealth])))),
    sd $(std(βdf[:b1expenses])), and se $(std(βdf[:b1expenses])/size(βdf,1)^0.5)")

  bronbc = FMLM(βdf, Meta.parse("b1return"), :b1contributions,
    XNames = [:b1return], YName=:b1contributions)
  println("beta contributions on beta returns is $(bronbc.β)")
  #NOTE: THinking of a table which has something like the above stats with a pooled beta stat too

  for i ∈ 1:length(plots) #write the graphs
    try
      draw(SVG("$path\\$(plotNames[i]).svg", 9inch, 7inch), plots[i])
    catch err
      @error "ERROR writing $(plotNames[i]). Error message $err"
      println("Size: $(size(subdf))")
    end
    #println("Graph $(plotNames[i]) written.")
  end

  return βdf
end

function processβ(data::NCCSData)::NCCSData
  βdf(data)

  return data
end

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
## Preloads
using Statistics
using FFTW
using Plots
using BenchmarkTools
using Profile
using LinearAlgebra

# %% Code Cell
function fink_filter_bank_fast(J, L)

    # -------- set parameters
    dθ   = π/8        # 8 angular bins hardwired
    nx   = 256
    dx   = nx/2-1

    # -------- allocate output array of zeros
    filt = zeros(256, 256, J, L)

    # -------- allocate theta and logr arrays
    logr = zeros(nx, nx)
    θ    = zeros(nx, nx)

    for l = 0:L-1
        θ_l = dθ*l

    # -------- allocate anggood BitArray
        anggood = falses(nx, nx)

    # -------- loop over pixels
        for x = 1:nx
            sx = mod(x+dx,nx)-dx -1    # define sx,sy so that no fftshift() needed
            for y = 1:nx
                sy = mod(y+dx,nx)-dx -1
                θ_pix  = mod(atan(sy, sx)+π -θ_l, 2*π)
                θ_good = abs(θ_pix-π) <= dθ

            # If this is a pixel we might use, calculate log2(r)
                if θ_good
                    anggood[y, x] = θ_good
                    θ[y, x]       = θ_pix
                    r = sqrt(sx^2 + sy^2)
                    logr[y, x] = log2(max(1,r))
                end
            end
        end
        angmask = findall(anggood)
    # -------- compute the wavelet in the Fourier domain
    # -------- the angular factor is the same for all j
        F_angular = cos.((θ[angmask].-π).*4)

    # -------- loop over j for the radial part
        for j = 0:J-1
            jrad  = 7-j
            Δj    = abs.(logr[angmask].-jrad)
            rmask = (Δj .<= 1)

    # -------- radial part
            F_radial = cos.(Δj[rmask] .* (π/2))
            ind      = angmask[rmask]
            filt[ind,j+1,l+1] = F_radial .* F_angular[rmask]
        end
    end
    return filt
end

fink_filter_set = fink_filter_bank_fast(8,8)
test_img = ones(256,256)
copyto!(test_img,fink_filter_set[:,:,1,3])

# %% Code Cell

function speedy_DHC(image::Array{Float64,2}, filter_set::Array{Float64,4})
    FFTW.set_num_threads(2)

    # array sizes
    (Nx, Ny)  = size(image)
    (_,_,J,L) = size(fink_filter_set)

    out_coeff = []
    # allocate coeff arrays
    S0  = zeros(2)
    S1  = zeros(J, L)
    S20 = zeros(Float64, J*L, J*L)
    S12 = zeros(Float64, J*L, J*L)
    im_rd_0_1  = zeros(Float64, Nx, Ny, J, L)
    im_fd_0_1 = zeros(Float64, Nx, Ny, J, L)

    ## 0th Order
    S0[1]   = mean(image)
    norm_im = image.-S0[1]
    S0[2]   = BLAS.dot(norm_im,norm_im)/(Nx*Ny)
    norm_im ./= sqrt(Nx*Ny*S0[2])

    append!(out_coeff,S0[:])

    ## 1st Order
    im_fd_0 = fft(norm_im)

    Atmp = zeros(ComplexF64,Nx,Ny)
    Btmp = zeros(Float64,Nx,Ny)
    ## Main 1st Order and Precompute 2nd Order
    @views for l = 1:L
        for j = 1:J
            @inbounds copyto!(Atmp, im_fd_0 .* filter_set[:,:,j,l])
            Btmp .= abs.(Atmp)
            @inbounds im_fd_0_1[:,:,j,l] .= Btmp
            @inbounds S1[j,l] +=BLAS.dot(Btmp,Btmp)
            ifft!(Atmp)
            @inbounds im_rd_0_1[:,:,j,l] .= abs.(Atmp)
        end
    end
    append!(out_coeff,S1[:])

    #DPF version of 2nd Order
    Amat    = reshape(im_fd_0_1, Nx*Ny, J*L)
    mul!(S12,Amat',Amat)

    Amat    = reshape(im_rd_0_1, Nx*Ny, J*L)
    mul!(S20,Amat',Amat)

    append!(out_coeff,S20[:])
    append!(out_coeff,S12[:])

    return out_coeff
end

temp = speedy_DHC(test_img,fink_filter_set)

@time speedy_DHC(test_img,fink_filter_set)

Profile.clear()
@profile speedy_DHC(test_img,fink_filter_set)

Juno.profiler()

BenchmarkTools.DEFAULT_PARAMETERS.seconds = 30
@benchmark speedy_DHC(test_img,fink_filter_set)

function speedy_DHC_old(image::Array{Float64,2}, filter_set::Array{Float64,4})
    FFTW.set_num_threads(2)

    # array sizes
    (Nx, Ny)  = size(image)
    (_,_,J,L) = size(fink_filter_set)

    out_coeff = []
    # allocate coeff arrays
    S0  = zeros(2)
    S1  = zeros(J, L)
    S20 = zeros(Float64, J, L, J, L)
    S12 = zeros(Float64, J, L, J, L)
    im_rd_0_1  = zeros(Float64,    Nx, Ny, J, L)
    im_fdf_0_1 = zeros(Float64,    Nx, Ny, J, L)
    #im_fd_0_1  = zeros(ComplexF64, Nx, Ny, J, L)

    ## 0th Order
    S0[1]   = mean(image)
    norm_im = image.-S0[1]
    S0[2]   = sum(norm_im .* norm_im)/(Nx*Ny)
    norm_im ./= sqrt(Nx*Ny*S0[2])

    # Was this intentional to set this back to image?
    norm_im = image

    append!(out_coeff,S0[:])

    ## 1st Order
    im_fd_0 = fft(norm_im)

    foo = zeros(ComplexF64,Nx,Ny)
    Btmp = zeros(Float64,Nx,Ny)
    ## Main 1st Order and Precompute 2nd Order
    for l = 1:L
        for j = 1:J
            #@inbounds had!(im_fd_0_1[:,:,j,l],filter_set[:,:,j,l]) #wavelet already in fft domain not shifted
            # @inbounds Btmp .= abs.(im_fd_0_1[:,:,j,l])
            # We don't need an fftshift()
            # @inbounds im_fdf_0_1[:,:,j,l] .= fftshift(Btmp)
            copyto!(foo, im_fd_0 .* filter_set[:,:,j,l])
            Btmp .= abs.(foo)
            @inbounds im_fdf_0_1[:,:,j,l] .= Btmp
            # I think you are doing a 256x256 matrix multiplication here
            #@inbounds S1[j,l]+=BLAS.dot(Btmp,Btmp) #normalization choice arb to make order unity
            # You mean this:  (and BLAS won't speed this up much)
            # S1[j,l] = sum(Btmp.*Btmp)
            S1[j,l] = sum([x^2 for x in Btmp])  #slightly faster? (no malloc)
            #@inbounds Atmp .= ifft(foo)
            @inbounds im_rd_0_1[:,:,j,l] .= abs.(ifft(foo))

        end
    end
    append!(out_coeff,S1[:])

    #DPF version of 2nd Order
    Amat    = reshape(im_fdf_0_1, Nx*Ny, J*L)
    S12 = reshape(Amat' * Amat, J, L, J, L)

    Amat    = reshape(im_rd_0_1, Nx*Ny, J*L)
    S20 = reshape(Amat' * Amat, J, L, J, L)

    append!(out_coeff,S20)
    append!(out_coeff,S12)

    return out_coeff
end

temp = speedy_DHC_old(test_img,fink_filter_set)

@time speedy_DHC_old(test_img,fink_filter_set)

Profile.clear()
@profile speedy_DHC_old(test_img,fink_filter_set)

Juno.profiler()

BenchmarkTools.DEFAULT_PARAMETERS.seconds = 30
@benchmark speedy_DHC_old(test_img,fink_filter_set)

ext: jl

max_stars_repo_name: andrew-saydjari/DHC

max_issues_repo_name: andrew-saydjari/DHC

max_forks_repo_name: andrew-saydjari/DHC


text:
#-------------julia-fractal.jl-------------------------------------------------#
#
# Purpose: This file tests the creation of the julia fractal in julia for fun
#
#------------------------------------------------------------------------------#

# We will need a few libraries for visualization
using PyPlot

# Function to calculate and output julia fractal to image
function create_fractal(cutoff::Float64, step::Float64, c_val::Float64,
                        res::Int64, range::Float64)
    # Setting up the array to be used
    carr = Array{Float64}(res, res)

    # For now, let's just assume that the imaginary and real ranges are from
    #     -range < x < range
    for i = 1:res
        for j = 1:res
            z = complex((i/res)*range - range/2.0,(j/res)*range - range/2.0)
            c = complex(0,c_val)
            n = 255
            while (abs(z) < cutoff && n > 5)
                z = z*z + c
                n -= 5
            end
            carr[i,j] = n/255.0;
        end
    end

    return carr

end

# Fuction to scan through C values
function c_scan(cutoff::Float64, step::Float64, res::Int64,
                range::Float64, max_c::Float64, c_step::Float64)

    id = 0
    for i = 0:c_step:max_c
        println(i)
        carr = create_fractal(cutoff, step, i, res, range)
        figure(figsize = (10,10))
        imshow(carr, cmap=ColorMap("Blues"))
        subplots_adjust(top = 1, bottom = 0, right = 1, left = 0,
                        hspace = 0, wspace = 0)
        savefig(string("c_scan", lpad(id, 5, 0), ".png"))
        id += 1
    end
end

# Function to zoom in on fractal
function fractal_zoom(cutoff::Float64, step::Float64, c_val::Float64,
                      res::Int64, max_range::Float64,
                      min_range::Float64)
    id = 0
    range = max_range
    while range > min_range
        println(range)
        carr = create_fractal(cutoff, step, c_val, res, range)
        figure(figsize = (10,10))
        imshow(carr, cmap=ColorMap("Blues"))
        subplots_adjust(top = 1, bottom = 0, right = 1, left = 0,
                        hspace = 0, wspace = 0)
        savefig(string("fractal_zoom", lpad(id, 5, 0), ".png"))
        id += 1
        range -= range*0.05
    end
end

fractal_zoom(10.0, 0.1, 1.0, 1024, 1.0, 0.0001)

#c_scan(10.,5., 1024, 3.0, 1.0, 0.025)

max_stars_repo_name: vchuravy/skillpill-julia

max_issues_repo_name: oist/skillpill-julia

max_forks_repo_name: oist/skillpill-julia


text:
module NumericalMethods

include("bisection_method.jl")
include("divided_difference.jl")
include("hermite_interpolation.jl")
include("newtons_method.jl")
include("secant_method.jl")

export bisect
export newton
export secant
export divid
export hermit

end


text:
# Julia wrapper for header: System.h
# Automatically generated using Clang.jl

# Julia wrapper for header: Clock.h
# Automatically generated using Clang.jl

function sfClock_create()
    ccall((:sfClock_create, libcsfml_system), Ptr{sfClock}, ())
end

function sfClock_copy(clock)
    ccall((:sfClock_copy, libcsfml_system), Ptr{sfClock}, (Ptr{sfClock},), clock)
end

function sfClock_destroy(clock)
    ccall((:sfClock_destroy, libcsfml_system), Cvoid, (Ptr{sfClock},), clock)
end

function sfClock_getElapsedTime(clock)
    ccall((:sfClock_getElapsedTime, libcsfml_system), sfTime, (Ptr{sfClock},), clock)
end

function sfClock_restart(clock)
    ccall((:sfClock_restart, libcsfml_system), sfTime, (Ptr{sfClock},), clock)
end
# Julia wrapper for header: Export.h
# Automatically generated using Clang.jl

# Julia wrapper for header: InputStream.h
# Automatically generated using Clang.jl

# Julia wrapper for header: Mutex.h
# Automatically generated using Clang.jl

function sfMutex_create()
    ccall((:sfMutex_create, libcsfml_system), Ptr{sfMutex}, ())
end

function sfMutex_destroy(mutex)
    ccall((:sfMutex_destroy, libcsfml_system), Cvoid, (Ptr{sfMutex},), mutex)
end

function sfMutex_lock(mutex)
    ccall((:sfMutex_lock, libcsfml_system), Cvoid, (Ptr{sfMutex},), mutex)
end

function sfMutex_unlock(mutex)
    ccall((:sfMutex_unlock, libcsfml_system), Cvoid, (Ptr{sfMutex},), mutex)
end
# Julia wrapper for header: Sleep.h
# Automatically generated using Clang.jl

function sfSleep(duration)
    ccall((:sfSleep, libcsfml_system), Cvoid, (sfTime,), duration)
end
# Julia wrapper for header: Thread.h
# Automatically generated using Clang.jl

function sfThread_create(_function, userData)
    ccall((:sfThread_create, libcsfml_system), Ptr{sfThread}, (Ptr{Cvoid}, Ptr{Cvoid}), _function, userData)
end

function sfThread_destroy(thread)
    ccall((:sfThread_destroy, libcsfml_system), Cvoid, (Ptr{sfThread},), thread)
end

function sfThread_launch(thread)
    ccall((:sfThread_launch, libcsfml_system), Cvoid, (Ptr{sfThread},), thread)
end

function sfThread_wait(thread)
    ccall((:sfThread_wait, libcsfml_system), Cvoid, (Ptr{sfThread},), thread)
end

function sfThread_terminate(thread)
    ccall((:sfThread_terminate, libcsfml_system), Cvoid, (Ptr{sfThread},), thread)
end
# Julia wrapper for header: Time.h
# Automatically generated using Clang.jl

function sfTime_asSeconds(time)
    ccall((:sfTime_asSeconds, libcsfml_system), Cfloat, (sfTime,), time)
end

function sfTime_asMilliseconds(time)
    ccall((:sfTime_asMilliseconds, libcsfml_system), sfInt32, (sfTime,), time)
end

function sfTime_asMicroseconds(time)
    ccall((:sfTime_asMicroseconds, libcsfml_system), sfInt64, (sfTime,), time)
end

function sfSeconds(amount)
    ccall((:sfSeconds, libcsfml_system), sfTime, (Cfloat,), amount)
end

function sfMilliseconds(amount)
    ccall((:sfMilliseconds, libcsfml_system), sfTime, (sfInt32,), amount)
end

function sfMicroseconds(amount)
    ccall((:sfMicroseconds, libcsfml_system), sfTime, (sfInt64,), amount)
end
# Julia wrapper for header: Types.h
# Automatically generated using Clang.jl

# Julia wrapper for header: Vector2.h
# Automatically generated using Clang.jl

# Julia wrapper for header: Vector3.h
# Automatically generated using Clang.jl


text:
module PairVelocities

include("pairvels.jl")
include("read.jl")

end


text:
using Etcetera
using Documenter

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

makedocs(;
    modules=[Etcetera],
    authors="René Olsthoorn <39967488+rene-olsthoorn@users.noreply.github.com> and contributors",
    repo="https://github.com/rene-olsthoorn/Etcetera.jl/blob/{commit}{path}#{line}",
    sitename="Etcetera.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://rene-olsthoorn.github.io/Etcetera.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/rene-olsthoorn/Etcetera.jl",
    devbranch="main",
)


text:
function dummy_data(n, move_type, param_gen, D_obj)
    t = TraceLogger(String, Int)
    moves = [move_type for i in 1:n]
    ϕ = [param_gen() for i in 1:n]
    for i in 1:n
        v = rand(D_obj)
        log!(t, moves[i] * "$(ϕ[i])", v, Δv ≤ 0, time_ns())
    end
    return t
end

max_stars_repo_name: Michiel-VL/PolyaViz

max_issues_repo_name: Michiel-VL/PolyaViz


text:
__precompile__(false)

module ReviseTest

square(x) = x^2

cube(x) = x^3

fourth(x) = x^4  # this is an addition to the file

module Internal

mult2(x) = 2*x
mult3(x) = 3*x

"""
This has a docstring
"""
unchanged(x) = x

unchanged2(@nospecialize(x)) = x

end  # Internal

end


text:
# Create line data with NaN as polyline separator
x = [randn(10000); NaN; randn(10000) .+ 5 ]
y = [randn(10000); NaN; randn(10000) .+ 5]
# Draw shade
shade(x, y, xform="loglog")


text:
module OMEinsumContractionOrders

using KaHyPar
using SparseArrays, Suppressor
using OMEinsum
using OMEinsum: NestedEinsum

include("kahypar.jl")

end


text:
@testset "Solve: car" begin
    # ## horizon
    T = 51

    # ## car
    num_state = 3
    num_action = 2
    num_parameter = 0

    function car(x, u, w)
        [u[1] * cos(x[3]); u[1] * sin(x[3]); u[2]]
    end

    function midpoint_explicit(x, u, w)
        h = 0.1 # timestep
        x + h * car(x + 0.5 * h * car(x, u, w), u, w)
    end

    # ## model
    dynamics = Dynamics(midpoint_explicit, num_state, num_action, num_parameter=num_parameter)
    model = [dynamics for t = 1:T-1]

    # ## initialization
    x1 = [0.0; 0.0; 0.0]
    xT = [1.0; 1.0; 0.0]

    # ## rollout
    ū = [1.0e-2 * [1.0; 0.1] for t = 1:T-1]
    w = [zeros(num_parameter) for t = 1:T]
    x̄ = rollout(model, x1, ū, w)

    # ## objective
    ot = (x, u, w) -> 1.0 * dot(x - xT, x - xT) + 1.0e-2 * dot(u, u)
    oT = (x, u, w) -> 0.0 * dot(x - xT, x - xT)
    ct = Cost(ot, num_state, num_action, num_parameter=num_parameter)
    cT = Cost(oT, num_state, 0, num_parameter=0)
    objective = [[ct for t = 1:T-1]..., cT]

    # ## constraints
    ul = -5.0 * ones(num_action)
    uu = 5.0 * ones(num_action)

    p_obs = [0.5; 0.5]
    r_obs = 0.1

    function stage_con(x, u, w)
        e = x[1:2] - p_obs
        [
        ul - u; # control limit (lower)
        u - uu; # control limit (upper)
        r_obs^2.0 - dot(e, e); # obstacle
        ]
    end

    function terminal_con(x, u, w)
        e = x[1:2] - p_obs
        [
        x - xT; # goal
        r_obs^2.0 - dot(e, e); # obstacle
        ]
    end

    cont = Constraint(stage_con, num_state, num_action, indices_inequality=collect(1:5))
    conT = Constraint(terminal_con, num_state, 0, indices_inequality=collect(3 .+ (1:1)))
    constraints = [[cont for t = 1:T-1]..., conT]

    # ## problem
    options = Options(verbose=false,
        line_search=:armijo,
        min_step_size=1.0e-5,
        objective_tolerance=1.0e-3,
        lagrangian_gradient_tolerance=1.0e-3,
        max_iterations=100,
        max_dual_updates=10,
        initial_constraint_penalty=1.0,
        scaling_penalty=10.0)
    s = Solver(model, objective, constraints,
        options=options)
    initialize_controls!(s, ū)
    initialize_states!(s, x̄)

    # ## solve
    solve!(s)

    # ## solution
    x_sol, u_sol = get_trajectory(s)

    @test all([all(stage_con(x_sol[t], u_sol[t], w[t]) .<= options.constraint_tolerance) for t = 1:T-1])
    @test all(abs.(terminal_con(x_sol[T], zeros(0), zeros(0))[1:3]) .<= options.constraint_tolerance)

    # ## allocations
    # info = @benchmark solve!($prob, a, b) setup=(a=deepcopy(x̄), b=deepcopy(ū))
    # @test info.allocs == 0
end


text:
module GraphMakie

using NetworkLayout
using LightGraphs
using Makie
using LinearAlgebra

import Makie: DocThemer, ATTRIBUTES, project, automatic

include("recipes.jl")
include("interaction.jl")

end


text:
using Test

@test 1 ≡ 1


text:
ENV["GKSwstype"] = "100"  # set 'GR environment' to 'no output' (for Travis CI)
using Documenter, StructuralDynamicsODESolvers

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

# generate Literate documentation
include("generate.jl")

makedocs(
    format = Documenter.HTML(prettyurls = haskey(ENV, "GITHUB_ACTIONS"),  # disable for local builds
                             collapselevel = 1),
    sitename = "StructuralDynamicsODESolvers.jl",
    doctest = false,
    strict = false,
    pages = [
        "Home" => "index.md",
        "Algorithms" => Any["First-order problems" => "lib/first_order.md",
                            "Second-order problems" => "lib/second_order.md"],
        "Examples" => Any["Example (Ch.9 Bathe)" => "models/example_9_1_Bathe.md",
                          "Spring-mass" => "models/massDashpotSpring.md",
                          "Dynamic Von Mises" => "models/dynamic_von_mises_truss.md"],
        "API Reference" => "lib/api.md",
        "References" => "references.md",
        "About" => "about.md"
    ]
)

# Deploy built documentation from Travis.
deploydocs(
    repo = "github.com/ONSAS/StructuralDynamicsODESolvers.jl.git",
    push_preview = true,
)


text:
using Test
using IntPrint

@test intprint(1) == "1"
@test intprint(-1) == "-1"
@test intprint(123456) == "123,456"
@test intprint(-123456) == "-123,456"
@test intprint(1234,except4=true) == "1234"
@test intprint(1234,sep=".") == "1.234"


text:
include("../src/PowerFlowNetworks.jl")
using .PowerFlowNetworks
using SQLite
using ArgParse

function parse_commandline()
    s = ArgParseSettings()
    @add_arg_table s begin
        "--dbpath", "-p"
            help = "Path of the sqlite database file."
            default = "data/PowerFlowNetworks.sqlite"
        "--recompute"
            help = "Wether to recompute the basic features if they're already not NULL in the database."
            default = false
            action = :store_true
    end
    return parse_args(s)
end

function main()
    parsed_args = parse_commandline()
    db = SQLite.DB(parsed_args["dbpath"])
    save_basic_features_instances!(db; recompute=parsed_args["recompute"])
end

main()

max_stars_repo_name: charlyalizadeh/PowerFlowNetworks

max_issues_repo_name: charlyalizadeh/PowerFlowNetworks

max_forks_repo_name: charlyalizadeh/PowerFlowNetworks


text:
using Pkg, DrWatson
#@quickactivate "StatisticalRethinkingStan"
using StanSample
using StatisticalRethinking

function sim_train_test(N::Int, K::Int, rho = [0.15, -0.4])

    n_dim = 1 + length(rho)
    n_dim = n_dim < K ? K : n_dim

    Rho = Matrix{Float64}(I, n_dim, n_dim)
    for i in 1:length(rho)
        Rho[i+1, 1] = Rho[1, i + 1] = rho[i]
    end

    x_train = Matrix(rand(MvNormal(zeros(n_dim), Rho), N)')
    x_test = Matrix(rand(MvNormal(zeros(n_dim), Rho), N)')
    y = x_train[:, 1]
    x_train = x_train[:, 2:n_dim]
    (y, x_train, x_test[:, 2:n_dim])
end

stan_01 = "
data {
    int<lower=1> K;
    int<lower=0> N;
    matrix[N, K] x;
    vector[N] y;

    int<lower=0> N_new;
    matrix[N_new, K] x_new;
}
parameters {
    real a;
    vector[K] b;
    real<lower=0> sigma;

    vector[N_new] y_new;                  // predictions
}
transformed parameters {
    vector[N] mu;
    vector[N_new] mu_new;

    mu = a + x * b;
    mu_new = a + x_new * b;
}
model {
    y ~ normal(mu, sigma);          // observed model
    y_new ~ normal(mu_new, sigma);  // prediction model
}
generated quantities {
    vector[N] log_lik;
    vector[N_new] log_lik_new;

    for (i in 1:N)
        log_lik[i] = normal_lpdf(y[i] | mu[i], sigma);

    for (j in 1:N_new)
        log_lik_new[j] = normal_lpdf(y_new[j] | mu_new[j], sigma);

}
";

sm = SampleModel("stan_01", stan_01)
N = 100
K = 6
rho = [0.15, -0.4]
y, x_train, x_test = sim_train_test(N, K, rho)
println(length(y))
println(size(x_train))
println(size(x_test))
k = K < 2 ? K : K - 1
data = (N = size(x_train, 1), K = size(x_train, 2), y = y,
    x = x_train, N_new = size(x_test, 1), x_new = x_test)
rc = stan_sample(sm; data)
if success(rc)
    sm_df = read_samples(sm, :dataframe)
    precis(sm_df[:, vcat([:a], [Symbol("b.$i") for i in 1:k])])
end

nt = read_samples(sm)
log_lik = nt.log_lik'
log_lik_new = nt.log_lik_new'

loo, loos, pk = psisloo(log_lik)
loo_new, loos_new, pk_new = psisloo(log_lik_new)

[K, loo, loo_new] |> display

#plot(y)
#plot!(nt.y_new')


text:
export id, X, Y, Z, H, S, Sdg, T, Tdg, U1, U2, U3, CU, CNOT, CCX, SWAP

"""
    id(n=1)

Identity operator for n qubits system.
"""
function id(n::Int=1)
    return identityoperator(SpinBasis(1//2)^n)
end
const X = sigmax
const Y = sigmay
const Z = sigmaz
const H = hadamard
const S = sgate
const Sdg = sdggate
const T = tgate
const Tdg = tdggate
const U1 = phaseshift
"""
    U2(ϕ::Real, λ::Real)

Equivalent to `singleunitary(π/2, ϕ, λ)`.

```math
U(\\phi, \\lambda) = \\frac{1}{\\sqrt{2}} \\begin{bmatrix}
       1 & -e^{i \\lambda} \\\\
       e^{i \\phi} & e^{i (\\lambda + \\phi)}
\\end{bmatrix}
```
"""
U2(ϕ::Real, λ::Real) = singleunitary(π/2, ϕ, λ)
const U3 = singleunitary
const CU = controlsgate
const CNOT = cnot
const CCX = toffoli
const SWAP = swap


text:
using Hecke, ArgParse

include(joinpath(Hecke.pkgdir,"examples/FieldEnumeration/FieldEnumeration.jl"))

function ArgParse.parse_item(::Type{fmpz}, x::AbstractString)
  if in('^', x)
    l = split(x, '^')
    if length(l) != 2
      throw(error("Could not parse $x as fmpz"))
    end
    l = string.(l)
    return (parse(fmpz, l[1]))^parse(Int, l[2])
  else
    return parse(fmpz, string(x))
  end
end

function ArgParse.parse_item(::Type{Vector{Int}}, x::AbstractString)
  return Meta.eval(Meta.parse(x))
end

function parse_commandline()
  s = ArgParseSettings()

  @add_arg_table s begin
    "--type", "-t"
      help = "Type of the abelian group"
      arg_type = Vector{Int}
      required = true
    "--conductors", "-c"
      help = "File containing the conductors"
      arg_type = String
      required = true
    "--only-real"
      help = "Only totally real fields"
      action = :store_true
    "--only-cm"
      help = "Only CM fields"
      action = :store_true
    "--non-simple"
      help = "Compute also non-simple presentation"
      action = :store_true
  end

  return parse_args(s)
end

function main()
  parsed_args = parse_commandline()

  local type::Vector{Int}
  local conds::String
  local only_real::Bool
  local only_cm::Bool
  local non_simple::Bool

  for (arg, val) in parsed_args
    println("$arg => $val")
    if arg == "type"
      type = val
    elseif arg == "conductors"
      conds = val
    elseif arg == "only-real"
      only_real = val
    elseif arg == "only-cm"
      only_cm = val
    elseif arg == "non-simple"
      non_simple = val
    end
  end

  sprint_formatted(fmt, args...) = @eval @sprintf($fmt, $(args...))

  @show type
  @show conds
  @show only_real
  @show only_cm
  @show non_simple

  file = "fields_" * string(type) * "_from_" * conds

  println("Output file $file")

  if isfile(file)
    throw(error("File $file does already exist"))
  end

  l_conductors = Vector{Int}()

  print("Loading conductors ... ")
  open(conds) do f
    for (i, lline) in enumerate(eachline(f))
      push!(l_conductors, parse(Int, lline))
    end
  end

  gtype = type

  Qx, x = PolynomialRing(QQ, "x")
  K, a = NumberField(x - 1, "a")
  O = maximal_order(K)

  n=prod(gtype)
  expo=lcm(gtype)

  class_fields = Vector{Hecke.ClassField{Hecke.MapRayClassGrp, GrpAbFinGenMap}}()
  for (i, k) in enumerate(l_conductors)
    r, mr = Hecke.ray_class_groupQQ(O, k, true, expo)
    if !Hecke.has_quotient(r, gtype)
      continue
    end
    ls = subgroups(r, quotype = gtype, fun = (x, y) -> quo(x, y, false)[2])
    for s in ls
      C = ray_class_field(mr, s)::Hecke.ClassField{Hecke.MapRayClassGrp, GrpAbFinGenMap}
      if Hecke._is_conductor_minQQ(C, n)
        push!(class_fields, C)
      end
    end
  end
  fields = Vector()
  for i = 1:length(class_fields)
    println("Computing class field $(i) /$(length(class_fields))")
    C = class_fields[i]
    r, s = signature(C)
    if (!only_cm && !only_real) || (only_cm && r == 0) || (only_real && s == 0)
      NC = number_field(C)
      sfield = Hecke._from_relative_to_absQQ(NC, Hecke.automorphism_groupQQ(C))[1]
      if non_simple
        push!(fields, (sfield, NC, discriminant(maximal_order(sfield))))
      else
        push!(fields, (sfield, discriminant(maximal_order(sfield))))
      end
    end
  end

  #=
  #The non simple absolute extension and the absolute value of its discriminant
  fields = Vector{Tuple{NfAbsNS, fmpz}}()
  for i = 1:length(class_fields)
    println("Computing class field $(i) /$(length(class_fields))")
    C = class_fields[i]
    r, s = signature(C)
    if (!only_cm && !only_real) || (only_cm && r == 0) || (only_real && s == 0)
      L = number_field(C)
      polys = Vector{fmpq_poly}(undef, length(L.pol))
      for t = 1:length(L.pol)
        fK = Hecke.isunivariate(L.pol[t])[2]
        f = Qx(fmpq[coeff(coeff(fK, j), 0) for j = 0:degree(fK)])
        polys[t] = f
      end
      NS, gNS = number_field(polys, check = false, cached = false)
      push!(fields, (NS, Hecke.discriminantQQ(O, C, Int(minimum(defining_modulus(C)[1])))))
    end
  end
  =#

  ll = fields

  # Discriminant always at the end
  sort!(ll, lt = (x, y) -> abs(x[end]) <= abs(y[end]))

  @show length(ll)

  println("Writing to $file...")

  ll = identity.(ll)
  _write_fields(ll, file)
end

main()


text:
# MadNLP.jl
# Created by Sungho Shin (sungho.shin@wisc.edu)

include("MOI_interface.jl")

# Thin wrapper for NLPModels.jl
function madnlp(model::AbstractNLPModel;buffered=true, kwargs...)
    ips = InteriorPointSolver(model;kwargs...)
    initialize!(ips.kkt)
    return optimize!(ips)
end


text:
module Shapley

# Write your package code here.
function shapley()
    return "this is the Shapley"
end

end


text:
__precompile__()

module SkyCoords
using StaticArrays

export AbstractSkyCoords,
       ICRSCoords,
       GalCoords,
       FK5Coords,
       separation,
       position_angle,
       offset

include("types.jl")

# -----------------------------------------------------------------------------
# Helper functions: Create rotation matrix about a given axis (x, y, z)

function xrotmat(angle)
    s, c = sincos(angle)
    SMatrix{3,3}(1, 0, 0, 0, c, -s, 0, s, c)
end

function yrotmat(angle)
    s, c = sincos(angle)
    SMatrix{3,3}(c, 0, s, 0, 1, 0, -s, 0, c)
end

function zrotmat(angle)
    s, c = sincos(angle)
    SMatrix{3,3}(c, -s, 0, s, c, 0, 0, 0, 1)
end

# (lon, lat) -> [x, y, z] unit vector
function coords2cart(lon, lat)
    sinlon, coslon = sincos(lon)
    sinlat, coslat = sincos(lat)
    SVector{3}(coslat * coslon, coslat * sinlon, sinlat)
end

# [x, y, z] unit vector -> (lon, lat)
cart2coords(v) = atan(v[2], v[1]), atan(v[3], sqrt(v[1] * v[1] + v[2] * v[2]))

# -----------------------------------------------------------------------------
# Constant rotation matricies and precession matrix function

# ICRS --> FK5 at J2000 (See USNO Circular 179, section 3.5)
eta0 = deg2rad(-19.9 / 3600000.0)
xi0 = deg2rad(9.1 / 3600000.0)
da0 = deg2rad(-22.9 / 3600000.0)
const ICRS_TO_FK5J2000 = xrotmat(-eta0) * yrotmat(xi0) * zrotmat(da0)
const FK5J2000_TO_ICRS = ICRS_TO_FK5J2000'

# FK5J2000 --> Gal
# Note that galactic pole and zeropoint of l are somewhat arbitrary
# and not officially defined (as far as I know). The values below are
# from astropy.coordinates, which includes the following comment:
# | "This gives better consistency with other codes than using the values
# |  from Reid & Brunthaler 2004 and the best self-consistency between FK5
# |  -> Galactic and FK5 -> FK4 -> Galactic. The lon0 angle was found by
# |  optimizing the self-consistency."
ngp_fk5j2000_ra = deg2rad(192.8594812065348)
ngp_fk5j2000_dec = deg2rad(27.12825118085622)
lon0_fk5j2000 = deg2rad(122.9319185680026)
const FK5J2000_TO_GAL = (zrotmat(pi - lon0_fk5j2000) *
                         yrotmat(pi / 2.0 - ngp_fk5j2000_dec) * zrotmat(ngp_fk5j2000_ra))
const GAL_TO_FK5J2000 = FK5J2000_TO_GAL'

# Gal --> ICRS: simply chain through FK5J2000
const GAL_TO_ICRS = FK5J2000_TO_ICRS * GAL_TO_FK5J2000
const ICRS_TO_GAL = GAL_TO_ICRS'

# FK5J2000 --> FK5{epoch}
# Computes the precession matrix from J2000 to the given Julian equinox.
# Expression from from Capitaine et al. 2003 as expressed in the USNO
# Circular 179.  This should match the IAU 2006 standard from SOFA.
const pzeta = [2.650545, 2306.083227, 0.2988499, 0.01801828, -0.000005971, -0.0000003173]
const pz = [-2.650545, 2306.077181, 1.0927348, 0.01826837, -0.000028596, -0.0000002904]
const ptheta = [0.0, 2004.191903, -0.4294934, -0.04182264, -0.000007089, -0.0000001274]
function precess_from_j2000(equinox)
    t = (equinox - 2000.0) / 100.0
    tn = 1.0
    zeta = pzeta[1]
    z = pz[1]
    theta = ptheta[1]
    for i = 2:6
        tn *= t
        zeta += pzeta[i] * tn
        z += pz[i] * tn
        theta += ptheta[i] * tn
    end
    zeta = deg2rad(zeta / 3600.0)
    z = deg2rad(z / 3600.0)
    theta = deg2rad(theta / 3600.0)
    zrotmat(-z) * yrotmat(theta) * zrotmat(-zeta)
end

# -----------------------------------------------------------------------------
# Type-dependent methods

lon(c::GalCoords) = c.l
lat(c::GalCoords) = c.b
lon(c::AbstractSkyCoords) = c.ra
lat(c::AbstractSkyCoords) = c.dec

# Abstract away specific field names (ra, dec vs l, b)
coords2cart(c::AbstractSkyCoords) = coords2cart(lon(c), lat(c))

# Rotation matrix between coordinate systems: `rotmat(to, from)`
# Note that all of these return SMatrix{3,3}{Float64}, regardless of
# element type of input coordinates.
rotmat(::Type{T1}, ::Type{T2}) where {T1<:GalCoords,T2<:ICRSCoords} = ICRS_TO_GAL
rotmat(::Type{T1}, ::Type{T2}) where {T1<:ICRSCoords,T2<:GalCoords} = GAL_TO_ICRS

# Define both these so that `convert(FK5Coords{e}, ...)` and
# `convert(FK5Coords{e,T}, ...)` both work. Similar with other
# FK5Coords rotmat methods below.
@generated rotmat(::Type{FK5Coords{e1}}, ::Type{T2}) where {e1,T2<:ICRSCoords} =
    precess_from_j2000(e1) * ICRS_TO_FK5J2000
@generated rotmat(::Type{FK5Coords{e1,T1}}, ::Type{T2}) where {e1,T1,T2<:ICRSCoords} =
    precess_from_j2000(e1) * ICRS_TO_FK5J2000

@generated rotmat(::Type{FK5Coords{e1}}, ::Type{T2}) where {e1,T2<:GalCoords} =
    precess_from_j2000(e1) * GAL_TO_FK5J2000
@generated rotmat(::Type{FK5Coords{e1,T1}}, ::Type{T2}) where {e1,T1,T2<:GalCoords} =
    precess_from_j2000(e1) * GAL_TO_FK5J2000

@generated rotmat(::Type{T1}, ::Type{FK5Coords{e2,T2}}) where {T1<:ICRSCoords,e2,T2} =
    FK5J2000_TO_ICRS * precess_from_j2000(e2)'

@generated rotmat(::Type{T1}, ::Type{FK5Coords{e2,T2}}) where {T1<:GalCoords,e2,T2} =
    FK5J2000_TO_GAL * precess_from_j2000(e2)'

@generated rotmat(::Type{FK5Coords{e1}}, ::Type{FK5Coords{e2,T2}}) where {e1,e2,T2} =
    precess_from_j2000(e1) * precess_from_j2000(e2)'
@generated rotmat(::Type{FK5Coords{e1,T1}}, ::Type{FK5Coords{e2,T2}}) where {e1,T1,e2,T2} =
    precess_from_j2000(e1) * precess_from_j2000(e2)'

# ------------------------------------------------------------------------------
# Distance between coordinates

function _separation(λ_1, ϕ_1, λ_2, ϕ_2)
    Δλ = λ_2 - λ_1
    sin_Δλ, cos_Δλ = sincos(Δλ)
    sin_ϕ1, cos_ϕ1 = sincos(ϕ_1)
    sin_ϕ2, cos_ϕ2 = sincos(ϕ_2)
    return atan(
        hypot(cos_ϕ2 * sin_Δλ, cos_ϕ1 * sin_ϕ2 - sin_ϕ1 * cos_ϕ2 * cos_Δλ),
        sin_ϕ1 * sin_ϕ2 + cos_ϕ1 * cos_ϕ2 * cos_Δλ,
    )
end

"""
    separation(c1::AbstractSkyCoords, c2::AbstractSkyCoords) -> distance

Return angular separation between two sky coordinates, in radians.

The angular separation is calculated using the [Vincenty formula](http://en.wikipedia.org/wiki/Great-circle_distance), which is slightly more
complex and computationally expensive than some alternatives, but is stable at
at all distances, including the poles and antipodes.
"""
separation(c1::T, c2::T) where {T<:AbstractSkyCoords} =
    _separation(lon(c1), lat(c1), lon(c2), lat(c2))

separation(c1::T1, c2::T2) where {T1<:AbstractSkyCoords,T2<:AbstractSkyCoords} =
    separation(c1, convert(T1, c2))

"""
    position_angle(c1::AbstractSkyCoords, c2::AbstractSkyCoords) -> angle

Return position angle between two sky coordinates, in positive radians.

# Examples
```jldoctest
julia> c1 = ICRSCoords(0, 0); c2 = ICRSCoords(deg2rad(1), 0);

julia> position_angle(c1, c2) |> rad2deg
90.0
```
"""
position_angle(c1::T, c2::T) where {T <: AbstractSkyCoords} = _position_angle(lon(c1), lat(c1), lon(c2), lat(c2))
position_angle(c1::T1, c2::T2) where {T1 <: AbstractSkyCoords,T2 <: AbstractSkyCoords} = position_angle(c1, convert(T1, c2))

function _position_angle(λ1, ϕ1, λ2, ϕ2)
    sin_Δλ, cos_Δλ = sincos(λ2 - λ1)
    sin_ϕ1, cos_ϕ1 = sincos(ϕ1)
    sin_ϕ2, cos_ϕ2 = sincos(ϕ2)

    x = sin_ϕ2 * cos_ϕ1 - cos_ϕ2 * sin_ϕ1 * cos_Δλ
    y = sin_Δλ * cos_ϕ2

    return mod2pi(atan(y, x))
end

"""
    offset(::AbstractSkyCoords, separation, pa) -> coordinate

Offset a coordinate by a given angular separation, `separation`, in radians and position angle, `pa`, in radians.

Uses the sine and cosine rules in spherical coordinates with corrections for the antipodes. Returns a sky coordinate of the same type as input.

# Examples
```jldoctest
julia> c1 = ICRSCoords(0, 0);

julia> c2 = offset(c1, deg2rad(1), deg2rad(90))
ICRSCoords{Float64}(0.017453292519943295, 1.0686516840418957e-18)

julia> offset(c1, c2) .|> rad2deg
(1.0, 90.0)
```

# See Also
* [`separation`](@ref), [`position_angle`](@ref)
"""
offset(c::T, sep, pa) where T <: AbstractSkyCoords = T(_offset(lon(c), lat(c), sep, pa)...)

"""
    offset(::AbstractSkyCoords, AbstractSkyCoords) -> angle, angle

Return the separation and position angle in radians between two sky coordinates.

# Examples
```jldoctest
julia> c1 = ICRSCoords(0, 0); c2 = ICRSCoords(deg2rad(1), 0);

julia> offset(c1, c2) .|> rad2deg
(1.0, 90.0)
```

# See Also
* [`separation`](@ref), [`position_angle`](@ref)
"""
offset(c1::AbstractSkyCoords, c2::AbstractSkyCoords) = separation(c1, c2), position_angle(c1, c2)

#= use the cosine rule in spherical geometry with three points, the north pole, the starting point,
and the final point.
angles: (change in lon), (position angle), (-1/position angle)
sides: (separation), (final co-latitude), (starting co-latitude)
=#
function _offset(λ, ϕ, separation, pa)
    sin_a, cos_a = sincos(separation)
    cos_c, sin_c = sincos(ϕ)
    sin_B, cos_B = sincos(pa)

    # solving cosine rule
    cos_b = cos_c * cos_a + sin_c * sin_a * cos_B

    # solving sine rule
    xsin_A = sin_a * sin_B * sin_c
    xcos_A = cos_a - cos_b * cos_c

    # correction for antipodes, otherwise atan2
    ang = sin_c < 1e-12 ? π/2 + cos_c * (π/2 - pa) : atan(xsin_A, xcos_A)

    return mod2pi(λ + ang), asin(cos_b)
end

end # module


text:
# These functions should be overloaded from ModelKit, so we need to make sure
# that they are already defined in the HomotopyContinuation module.
function evaluate end
function evaluate! end
function evaluate_and_jacobian end
function evaluate_and_jacobian! end
function jacobian end
function jacobian! end
function dt end
function dt! end
function jacobian_and_dt end
function jacobian_and_dt! end

module ModelKit

using OrderedCollections: OrderedDict
using StaticArrays: @SVector, @SMatrix

import LinearAlgebra: det, dot
import Latexify
# Overload these functions vom MP to be able to export this without any name clashes
import MultivariatePolynomials: variables, nvariables, differentiate, subs, monomials

export Expression, Constant, Variable, Operation
import ..HomotopyContinuation:
    evaluate,
    evaluate!,
    evaluate_and_jacobian,
    evaluate_and_jacobian!,
    jacobian,
    jacobian!,
    dt,
    dt!,
    jacobian_and_dt,
    jacobian_and_dt!

export @var,
    @unique_var,
    subs,
    variables,
    nvariables,
    differentiate,
    monomials,
    Compiled,
    CompiledSystem,
    CompiledHomotopy,
    compile,
    interpreted,
    System,
    Homotopy,
       # reexport
    evaluate,
    evaluate!,
    evaluate_gradient,
    evaluate_and_jacobian,
    evaluate_and_jacobian!,
    jacobian,
    jacobian!,
    dt,
    dt!,
    jacobian_and_dt,
    jacobian_and_dt!

include("model_kit/expression.jl")
include("model_kit/codegen.jl")

end


text:
##### structs for individuals
mutable struct plankton
    data::AbstractArray
    p::NamedTuple
end

struct individuals
    phytos::NamedTuple
    # zoos::NamedTuple
end

##### find indices (halo points included)
@kernel function find_inds_kernel!(plank, g::AbstractGrid)
    i = @index(Global)
    @inbounds plank.xi[i] = unsafe_trunc(Int, get_xf_index(plank.x[i]) * plank.ac[i]) + g.Hx
    @inbounds plank.yi[i] = unsafe_trunc(Int, get_yf_index(plank.y[i]) * plank.ac[i]) + g.Hy
    @inbounds plank.zi[i] = unsafe_trunc(Int, get_zf_index(plank.z[i]) * plank.ac[i]) + g.Hz
end
function find_inds!(plank, g::AbstractGrid, arch::Architecture)
    kernel! = find_inds_kernel!(device(arch), 256, (size(plank.ac,1)))
    event = kernel!(plank, g)
    wait(device(arch), event)
    return nothing
end

@kernel function find_NPT_kernel!(nuts, x, y, z, ac, NH4, NO3, PO4, DOC, par, temp, pop)
    i = @index(Global)
    @inbounds nuts.NH4[i] = max(1.0e-10, NH4[x[i], y[i], z[i]]) * ac[i]
    @inbounds nuts.NO3[i] = max(1.0e-10, NO3[x[i], y[i], z[i]]) * ac[i]
    @inbounds nuts.PO4[i] = max(1.0e-10, PO4[x[i], y[i], z[i]]) * ac[i]
    @inbounds nuts.DOC[i] = max(1.0e-10, DOC[x[i], y[i], z[i]]) * ac[i]
    @inbounds nuts.par[i] = par[x[i], y[i], z[i]] * ac[i]
    @inbounds nuts.T[i]   =temp[x[i], y[i], z[i]] * ac[i]
    @inbounds nuts.pop[i] = pop[x[i], y[i], z[i]] * ac[i]
end
function find_NPT!(nuts, x, y, z, ac, NH4, NO3, PO4, DOC, par, temp, pop, arch::Architecture)
    kernel! = find_NPT_kernel!(device(arch), 256, (size(ac,1)))
    event = kernel!(nuts, x, y, z, ac, NH4, NO3, PO4, DOC, par, temp, pop)
    wait(device(arch), event)
    return nothing
end

##### calculate Chla and individual counts based on the status of plankton individuals
@inline function gpu_acc_counts_kernel!(ctsChl, ctspop, Chl, ac, x, y, z)
    index = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    stride = blockDim().x * gridDim().x
    for i = index:stride:size(ac,1)
        @inbounds CUDA.@atomic ctsChl[x[i], y[i], z[i]] += Chl[i] * ac[i]
        @inbounds CUDA.@atomic ctspop[x[i], y[i], z[i]] += ac[i]
    end
    return nothing
end
function acc_counts!(ctsChl, ctspop, Chl, ac, x, y, z, ::GPU)
    @cuda threads=256 blocks=ceil(Int, size(ac,1)/256) gpu_acc_counts_kernel!(ctsChl, ctspop, Chl, ac, x, y, z)
    return nothing
end
function acc_counts!(ctsChl, ctspop, Chl, ac, x, y, z, ::CPU)
    for i in 1:size(ac,1)
        @inbounds ctsChl[x[i], y[i], z[i]] += Chl[i] * ac[i]
        @inbounds ctspop[x[i], y[i], z[i]] += ac[i]
    end
    return nothing
end

##### calculate PAR field based on Chla field and depth
@kernel function calc_par_kernel!(par, Chl, PARF, g::AbstractGrid, kc, kw)
    i, j = @index(Global, NTuple)
    @unroll for k in 1:g.Nz
        ii = i + g.Hx
        jj = j + g.Hy
        kk = k + g.Hz
        atten = (Chl[ii,jj,kk]/volume(ii, jj, kk, g) * kc + kw) * ΔzF(ii, jj, kk, g)
        par[ii,jj,kk] = PARF[i,j] * (1.0 - exp(-atten)) / atten
        PARF[i,j] = PARF[i,j] * exp(-atten)
    end
end
function calc_par!(par, arch::Architecture, Chl, PARF, g::AbstractGrid, kc, kw)
    kernel! = calc_par_kernel!(device(arch), (16,16), (g.Nx, g.Ny))
    event = kernel!(par, Chl, PARF, g, kc, kw)
    wait(device(arch), event)
    return nothing
end

@kernel function mask_individuals_kernel!(plank, mask)
    i = @index(Global)
    xi = unsafe_trunc(Int, plank.x[i]) + 1 # 0-based index to 1-based index
    yi = unsafe_trunc(Int, plank.y[i]) + 1
    zi = unsafe_trunc(Int, plank.z[i]) + 1
    plank.ac[i] = mask[xi, yi, zi] * plank.ac[i]
end
function mask_individuals!(plank, mask, N, arch)
    kernel! = mask_individuals_kernel!(device(arch), 256, (N,))
    event = kernel!(plank, mask)
    wait(device(arch), event)
    return nothing
end


text:
#############################################################
# abstractions for constructing custom generative functions #
#############################################################

"""
    CustomDetermGFTrace{T,S} <: Trace

Trace type for custom deterministic generative function.
"""
struct CustomDetermGFTrace{T,S} <: Trace
    retval::T
    state::S
    args::Tuple
    gen_fn::Any
end

get_args(trace::CustomDetermGFTrace) = trace.args

get_retval(trace::CustomDetermGFTrace) = trace.retval

get_choices(trace::CustomDetermGFTrace) = EmptyChoiceMap()

get_score(trace::CustomDetermGFTrace) = 0.

project(trace::CustomDetermGFTrace, selection::Selection) = 0.

get_gen_fn(trace::CustomDetermGFTrace) = trace.gen_fn

"""
    CustomDetermGF{T,S} <: GenerativeFunction{T,CustomDetermGFTrace{T,S}}

Abstract type for a custom deterministic generative function.
"""
abstract type CustomDetermGF{T,S} <: GenerativeFunction{T,CustomDetermGFTrace{T,S}} end

"""
    retval, state = execute_determ(gen_fn::CustomDetermGF, args)

Execute the generative function and return the return value and the state.
"""
function execute_determ end

"""
    state, retval, retdiff = update_determ(gen_fn::CustomDetermGF, state, args, argdiffs)

Update the arguments to the generative function and return new return value and state.
"""
function update_determ end

"""
    arg_grads = gradient_determ(gen_fn::CustomDetermGF, state, retgrad)

Return the gradient tuple with respect to the arguments.
"""
function gradient_determ end

"""
    arg_grads = accumulate_param_gradients_get!(trace::CustomDefGF, retgrad, scaler)

Increment gradient accumulators for parameters the gradient with respect to the
arguments, optionally scaled, and return the gradient with respect to the
arguments (not scaled).

Given the previous state and a gradient with respect to the return value \$∇_y
J\$ (`retgrad`), return the following gradient (`arg_grads`) with respect to
the arguments \$x\$:
```math
∇_x J
```
Also increment the gradient accumulators for the trainable parameters \$Θ\$ of
the function by:
```math
s * ∇_Θ J
```
where \$s\$ is `scaler`.
"""
function accumulate_param_gradients_determ! end

function simulate(gen_fn::CustomDetermGF{T,S}, args::Tuple) where {T,S}
    retval, state = execute_determ(gen_fn, args)
    CustomDetermGFTrace{T,S}(retval, state, args, gen_fn)
end

function generate(gen_fn::CustomDetermGF{T,S}, args::Tuple, choices::ChoiceMap) where {T,S}
    if !isempty(choices)
        error("Deterministic generative function makes no random choices")
    end
    retval, state = execute_determ(gen_fn, args)
    trace = CustomDetermGFTrace{T,S}(retval, state, args, gen_fn)
    trace, 0.
end

function update(trace::CustomDetermGFTrace{T,S}, args::Tuple, argdiffs::Tuple, choices::ChoiceMap) where {T,S}
    if !isempty(choices)
        error("Deterministic generative function makes no random choices")
    end
    state, retval, retdiff = update_determ(trace.gen_fn, trace.state, args, argdiffs)
    new_trace = CustomDetermGFTrace{T,S}(retval, state, args, trace.gen_fn)
    (new_trace, 0., retdiff)
end

function regenerate(trace::CustomDetermGFTrace, args::Tuple, argdiffs::Tuple, selection::Selection)
    update(trace, args, argdiffs, EmptyChoiceMap())
end

function choice_gradients(trace::CustomDetermGFTrace, selection::Selection, retgrad)
    arg_grads = gradient_determ(trace.gen_fn, trace.state, retgrad)
    (arg_grads, EmptyChoiceMap(), EmptyChoiceMap())
end

function accumulate_param_gradients!(trace::CustomDetermGFTrace, retgrad, scaler)
    accumulate_param_gradients_determ!(trace.gen_fn, trace.state, retgrad, scaler)
end

export CustomDetermGF, CustomDetermGFTrace, execute_determ, update_determ, gradient_determ, accumulate_param_gradients_determ!

max_stars_repo_name: cameronraysmith/Gen

max_issues_repo_name: cameronraysmith/Gen

max_forks_repo_name: cameronraysmith/Gen


text:
module FDA

using BSplines

struct basisobj
    basis_type
    basis
end

struct fdobj
    basis_type
    basis
    coefs
end

function second_order_diff(f, x, h)
    n = f(x+h) - 2*f(x) + f(x-h)
    return n / (h^2)
end

function estimate_R_mat(basis_type, basis, Psi_matrix, argvals, nderiv)
    dpsi = zeros(length(basis), length(argvals))
    if basis_type == "BSpline"
        for i in 1:length(argvals)
            tpsi = Psi_matrix[i,:]
            tid = [j for j in 1:length(tpsi) if tpsi[j]!=0]
            t_time = argvals[i]
            temp_deriv = bsplines(basis, t_time, BSplines.Derivative(nderiv))
            for k in tid
                dpsi[k, i] = temp_deriv[:,1][k]
            end
        end
    elseif basis_type == "Fourier"
        h = (argvals[2] - argvals[1])/10
        for i in 1:length(basis)
            tf = basis[i]
            temp_deriv = [second_order_diff(tf, t, h) for t in argvals]
            dpsi[i,:] = temp_deriv
        end
    end

    R_matrix = dpsi * transpose(dpsi)
    R_matrix = R_matrix ./ length(argvals)
    return R_matrix
end

function gen_bspline_basis(range_eval, nbasis, order=4)
    nbreaks = nbasis - order + 2
    break_points = range(range_eval[1], range_eval[2]; length = nbreaks)
    basis = BSplineBasis(order, break_points)
    r = basisobj("BSpline", basis)
    return r
end

function gen_fourier_basis(range_eval, nbasis)
    normalizer = (range_eval[2] - range_eval[1])
    basis = []

    function tf1(t)
        return 1/normalizer
    end

    push!(basis, tf1)

    n = Int(ceil((nbasis-1)/2))
    for i in 1:n
        function tfa(t)
            return sin(2*pi*i*t/normalizer)
        end
        push!(basis, tfa)
        function tfb(t)
            return cos(2*pi*i*t/normalizer)
        end
        push!(basis, tfb)
    end
    basis = basis[1:nbasis]

    r = basisobj("Fourier", basis)
end

function get_coef_mat(argvals, basisobj, lambda, nderiv)
    basis = basisobj.basis
    basis_type = basisobj.basis_type
    if basis_type == "BSpline"
        Psi_matrix = basismatrix(basis, argvals)
        R_matrix = estimate_R_mat(basis_type, basis, Psi_matrix, argvals, nderiv)
    elseif basis_type == "Fourier"
        Psi_matrix = zeros(length(argvals), length(basis))
        for i in 1:length(basis)
            tf = basis[i]
            Psi_matrix[:,i] = tf.(argvals)
            R_matrix = estimate_R_mat(basis_type, basis, Psi_matrix, argvals, nderiv)
        end
    end

    coef_mat = inv(transpose(Psi_matrix) * Psi_matrix .+ R_matrix .* lambda) *
            transpose(Psi_matrix)
    return coef_mat

end

function Data2fd(argvals, y, basisobj;
    lambda=3e-8/(maximum(argvals) - minimum(argvals)),
    nderiv=2)

    basis_type = basisobj.basis_type
    basis = basisobj.basis

    coefs = zeros(length(basis), ncol(y))
    coef_mat = get_coef_mat(argvals, basisobj,lambda, nderiv)
    for i in 1:ncol(y)
        coefs[:,i] = coef_mat * y[:,i]
    end

    res = fdobj(basis_type, basis, coefs)
    return res
end

function eval_fd(eval_grid, fdobj)
    basis = fdobj.basis
    coefs = fdobj.coefs
    if fdobj.basis_type == "BSpline"
        Psi_matrix = basismatrix(basis, eval_grid)
    elseif fdobj.basis_type == "Fourier"
        Psi_matrix = zeros(length(eval_grid), length(basis))
        for i in 1:length(basis)
            tf = basis[i]
            Psi_matrix[:,i] = tf.(eval_grid)
        end
    end
    obs = Psi_matrix * coefs
    return obs
end

end


text:
"""
Contains the results of
"""
struct RefinementQuery
    # The points (vertices) in the original simplex space
    points::AbstractArray{Float64, 2}

    # The points (vertices) in the image simplex space
    impoints::AbstractArray{Float64, 2}

    # Indices corresponding to the simplices of the triangulation. The indices in a single
    # row of this matrix indicates from which rows of `points` and `impoints` the
    # corresponding simplices should be constructed.
    simplex_inds::AbstractArray{Int64, 2}

    # The maximal allowed radius in the resulting triangulation  splitting simplices
    # that are too big.
    target_radius::Float64

    # The number of new simplices created by splitting to the desired `target_radius`
    n_new_simplices::Int

    ks::AbstractVector{Int}

    # Indices of the simplices that will be split
    inds_toolarge::AbstractVector{Int}

    # The resulting maximum subsimplex radius for each simplex simplex that is being
    # split due to being larger than the allowed target size. Should contain
    # max(radius(simplex), radius(image simplex)).
    resulting_radii::AbstractVector{Float64}

    # The indices of the simplices that will not be split
    inds_untouched::AbstractVector{Int}
end

"""
   query_refinement(t::Triangulation, target_radius::Float64)

Calculates information about the triangulation resulting from calling `refine_variable_k`
with identical arguments.

`refine_variable_k` works as follows:

Refines a triangulation so that the maximum simplex radius is less than `target_radius`.

The splitting is done in a single iteration. For each simplex that we want to split (given
by the integer vector `inds_to_split`, referencing rows of the `simplex_inds` array),
we indentify the necessary splitting factor `k`

Each row in the integer array `simplex_inds` represents a simplex in both the original
space and in the image space. The E+1 elements of each row vector in `simplex_inds`
references a set of corresponding rows in `points` and `impoints`, such that

points[simplex_inds[k, :], :] gives the k-th simplex in the original space, and
impoints[simplex_inds[k, :], :] gives the k-th simplex in the image space.
"""
function query_refinement(t::Triangulation,
                        target_radius::Float64)
   dim = size(t.points, 2)
   # For each simplex-imagesimplex pair, determine if either of them has a radius
   # larger than the allowed `target_radius`. Simultaneously, calculate the necessary
   # splitting factors k needed to reduce their radius to below `target_radius`
   inds_toolarge = Int[]
   ks = Int[]
   resulting_radii = Float64[] # store the resulting radii for  splitting
   n_new_simplices = 0

   # Indices of the simplices that are not split
   untouched_indices = setdiff(1:size(t.simplex_inds, 1), inds_toolarge)

   for i = 1:length(t.radii)
      if t.radii[i] >= target_radius || t.radii_im[i] >= target_radius
         largest_radius = max(t.radii[i], t.radii_im[i])
         k = find_k(largest_radius, target_radius)
         n_new_simplices += k^dim
         append!(resulting_radii, largest_radius/k)
         append!(ks, k)
         append!(inds_toolarge, i)
      end
   end

   # Indices of the simplices that are not split
   untouched_indices = setdiff(1:size(t.simplex_inds, 1), inds_toolarge)

   return RefinementQuery(t.points, t.impoints, t.simplex_inds, target_radius,
                           n_new_simplices, ks, inds_toolarge, resulting_radii,
                           untouched_indices)
end

"""
    find_k(radius::Float64, target_radius::Float64)

Finds the splitting factor `k` needed to split a simplex with a given `radius` into
subsimplices each having radii less than `target_radius`.
"""
function find_k(radius::Float64, target_radius::Float64; tol = 0.0001)
    ceil(Int, radius/(target_radius - tol))
end

"""
    SplitInfo(inds_toolarge::Vector{Int}, ks::Vector{Int})

Stores information about a potential splitting of a triangulation.
"""
struct SplitInfo
    # Indices of the simplices that are too large
    inds_toolarge::Vector{Int}

    # The corresponding splitting factors needed to reduce the radii of the simplices
    # in `inds_toolarge` to below the desired target radius
    ks::Vector{Int}
end

"""
    get_split_info(t::Triangulation,
                target_radius::Float64)

Given a `Triangulation`, returns the indices of simplices of the triangulation that must be
split in order for the maximum simplex radius to be below a `target_radius`.

The corresponding splitting factors are also returned.
"""
function get_split_info(t::Triangulation,
                        target_radius::Float64)

    # Find the radii and centroids of the simplices we want to split
    dim = size(t.points, 2)
    # For each simplex-imagesimplex pair, determine if either of them has a radius
    # larger than the allowed `target_radius`. Simultaneously, calculate the necessary
    # splitting factors k needed to reduce their radius to below `target_radius`
    inds_toolarge = Int[]
    ks = Int[]
    resulting_radii = Float64[] # store the resulting radii for  splitting
    n_new_simplices = 0

    for i = 1:length(t.radii)
        if t.radii[i] >= target_radius || t.radii_im[i] >= target_radius
            largest_radius = max(t.radii[i], t.radii_im[i])
            k = find_k(largest_radius, target_radius)
            n_new_simplices += k^dim
            append!(resulting_radii, largest_radius/k)
            append!(ks, k)
            append!(inds_toolarge, i)
        end
    end

   return SplitInfo(inds_toolarge, ks)
end

"""
    SplitRules(SplitInfo)

Stores the simplex splitting rules (a result of a `simplicial_subdivision(dim, k)` call)
for a given dimension and splitting factor.
"""
struct SplitRules
    rules::AbstractArray{Int, 2}
    indices::AbstractArray{Int, 2}
end

"""
Refines a triangulation `t`using shape-preserving simplex subdivision. Uses
a variable splitting factor, so that all simplices' radii are rendered less
than `target_radius`. Operates in-place on `t`.
"""
function refine_variable_k!(t::Triangulation, target_radius::Float64)
   #=
   # Get:
   # (1) The indices of simplices that must be split in order for all the radii
   #     of the provided simplices to be below a `target_radius`.
   # (2) The corresponding splitting factors (`ks`).
   =#
   split_info = get_split_info(t, target_radius)
   query = query_refinement(t, target_radius)
   @assert split_info.ks == query.ks

   # If any simplices does not need splitting, keep track of their indices
   n_simplices = size(t.simplex_inds, 1)
   untouched_inds = setdiff(1:n_simplices, split_info.inds_toolarge)

   if length(untouched_inds) > 0
       simplex_inds = vcat(t.simplex_inds, t.simplex_inds[untouched_inds, :])
   end

   # Loop over splitting factors
   for i in 1:length(unique(split_info.ks))
      #=
      # Split all simplices that needs to be split with this splitting factor
      # in order to reduce their radius to below `target_radius`
      =#
      k = unique(split_info.ks)[i]
      split_inds = split_info.inds_toolarge[find(split_info.ks .== k)]

      if length(split_inds) > 0
         refine_t!(t, split_inds, k)
      else
         warn("Target radius $target_radius is too large. All simplex radii are
               already smaller than this.")
      end
   end
end

"""
Refines a triangulation `t`using shape-preserving simplex subdivision. Uses
a variable splitting factor, so that all simplices' radii are rendered less
than `target_radius`.
"""
function refine_variable_k(t::Triangulation, target_radius::Float64)
   #=
   # Get:
   # (1) The indices of simplices that must be split in order for all the radii
   #     of the provided simplices to be below a `target_radius`.
   # (2) The corresponding splitting factors (`ks`).
   =#
   split_info = get_split_info(t, target_radius)
   query = query_refinement(t, target_radius)
   @assert split_info.ks == query.ks

   # If any simplices does not need splitting, keep track of their indices
   dim = size(t.points, 2)
   n_simplices = size(t.simplex_inds, 1)
   untouched_inds = setdiff(1:n_simplices, split_info.inds_toolarge)

   if length(untouched_inds) > 0
       simplex_inds = vcat(t.simplex_inds, t.simplex_inds[untouched_inds, :])
   end

   # Loop over splitting factors
   for i in 1:length(unique(split_info.ks))
      #=
      # Split all simplices that needs to be split with this splitting factor
      # in order to reduce their radius to below `target_radius`
      =#
      k = unique(split_info.ks)[i]
      split_inds = split_info.inds_toolarge[find(split_info.ks .== k)]

      if length(split_inds) > 0
         refine_t!(t, split_inds, k)
      else
         warn("Target radius $target_radius is too large. All simplex radii are
               already smaller than this.")
      end
   end
end


text:
#
# Prepare cost, proximal maps and differentials
M = PowerManifold(pixelM, NestedPowerRepresentation(), size(f)...)
N = TangentBundle(M)
fidelity(M, x) = 1 / 2 * distance(M, x, f)^2
Λ(M, x) = ProductRepr(x, forward_logs(M, x))
prior(M, x) = norm(norm.(Ref(pixelM), x, submanifold_component(N, Λ(M, x), 2)), 1)
cost(M, x) = (1 / α) * fidelity(M, x) + prior(M, x)

prox_F(M, λ, x) = prox_distance(M, λ / α, f, x, 2)
function prox_G_dual(N, n, λ, ξ)
    return ProductRepr(
        submanifold_component(N, ξ, 1),
        project_collaborative_TV(
            base_manifold(N),
            λ,
            submanifold_component(N, n, 1),
            submanifold_component(N, ξ, 2),
            Inf,
            Inf,
            1.0,
        ),
    )
end

DΛ(M, m, X) = ProductRepr(zero_vector(M, m), differential_forward_logs(M, m, X))
adjoint_DΛ(N, m, n, ξ) = adjoint_differential_forward_logs(N.manifold, m, ξ[N, :vector])


text:
using Gridap
using GridapGmsh
using Gridap.Geometry

mutable struct ModelInfo
    "Reference velocity"
    Uref::Float64
    "Velocity on different surfaces"
    U::Vector{Float64}
    "Different material regions"
    regions::Vector{String}
    "Surface regions where convection occurs"
    hsurf::Vector{String}
    "Surface regions where heat flux is prescribed"
    qsurf::Vector{String}
    "Surface regions where temperature is prescribed"
    tsurf::Vector{String}
    "Regions to which the prescribed temperature belongs"
    surfreg::Dict{String,String}
    "Heat conductivity on different materials"
    k::Dict{String,Float64}
    "Convection coefficient on different surfaces"
    h::Dict{String,Float64}
    "Prescribed temperature on different surfaces"
    T::Dict{String,Float64}
    "Prescribed heat flux on different surfaces"
    q::Dict{String,Float64}
    "Heat conductivity of fluif"
    kₐ::Float64
    "Prandtl number of fluid"
    Prₐ::Float64
    "Kinematic viscosity of the fluid"
    νₐ::Float64
    "Fluid temperature"
    Tref::Float64
end

"""
`ModelInfo(u=10.0; k=Dict("glass"=>1.0, "steel"=>50.0), hsurf=["hglass", "hsteel"],
                   T=Dict("temp_glass"=>200.0, "temp_steel"=>200.0),
                   q=Dict{String,Float64}(), Lmm = [2.0, 0.5],
                   qsurf=Dict{String,Float64}(),
                   Prₐ=0.7, kₐ=25e-3, νₐ=1.6e-5, Tref=20.0)`

Specifies material parameters and boundary conditions for FEM analysis.
The heat conduction model assumes that several regions with differing materials
can exist. There are 3 possible boundary conditions

 1. Prescribed temperature, Dirichlet BC
 2. Prescribed heat flux, Neumann BC
 3. Prescribed convection coefficient, mixed BC

Boundaries where no BC is prescribed automatically uses a prescribed
0 heat flux. This BC can be used for symmetry.

In the case of convection, different surfaces can have different velocities.

For now, to calculate the convection coefficient, a correlation for circular
cylinders is used:

      ``Nu = 0.51⋅Re^(½)⋅Pr^(0.37)``

"""
function ModelInfo(u=10.0; k=Dict("glass"=>1.0, "steel"=>50.0), hsurf=["hglass", "hsteel"],
                   T=Dict("temp_glass"=>200.0, "temp_steel"=>200.0),
                   surfreg=Dict("temp_glass"=>"glass", "temp_steel"=>"steel",
                                "hglass"=>"glass", "hsteel"=>"steel"),
                   q=Dict{String,Float64}(), Lmm = [2.0, 0.5],
                   qsurf=Dict{String,Float64}(),
                   Prₐ=0.707, kₐ=25.3e-3, νₐ=1.6e-5, Tref=20.0)
    nh = length(hsurf)
    nu = length(u)

    Uref = u[1]

    if nu==1
        U = fill(u[1], nh)
    elseif nu != nh
        error("Incompatible number of convection surfaces with velocities!")
    else
        U = [uu for uu in u]
    end

    if length(Lmm) != nh
        error("The number of dimensions should be the same as the number of convection regions!")
    end

    L = Lmm .* 1e-3
    Re = (U ./ νₐ) .* L
    Nu = (0.51*Prₐ^0.37) .* Re .^ 0.35 * 1.35 # .^ 0.5
    h1 = kₐ .* Nu ./ L

    h = Dict{String,Float64}()

    for i in 1:nh
        h[ hsurf[i] ] = h1[i]
    end
    return  ModelInfo(Uref, U,  collect(keys(k)), hsurf, collect(keys(q)), collect(keys(T)),
                      surfreg, k, h, T, q, kₐ, Prₐ, νₐ, Tref)

end

r(x) = x[2]

function fem_model(th::ModelInfo, model; order=2, dimension=2)

    degree = 2*order
    reffe = ReferenceFE(lagrangian, Float64, order)

    # Vamos montar os espaçoes de funções

    tsurf = th.tsurf
    uw = [th.T[s]-th.Tref for s in tsurf]
    V = TestFESpace(model, reffe, dirichlet_tags=tsurf)
    U = TrialFESpace(V, uw)

    Ω = Triangulation(model)
    dΩ = Measure(Ω, degree)

    hsurf = th.hsurf  # Boundaries with convection
    Γh = [BoundaryTriangulation(model, tags=hs) for hs in hsurf]
    dΓh = [Measure(Γ,degree) for Γ in Γh]

    qsurf = th.qsurf
    #Γq = [BoundaryTriangulation(model, tags=hs) for qs in qsurf]
    #dΓq = [Measure(Γ,degree) for Γ in Γq]

    Γt = [BoundaryTriangulation(model, tags=ts) for ts in tsurf]
    dΓt = [Measure(Γ,degree) for Γ in Γt]

    labels = get_face_labeling(model)
    tags = get_face_tag(labels, dimension)
    kdict = Dict{Int,Float64}()
    for r in th.regions
        tagidx = get_tag_from_name(labels, r)
        kdict[tagidx] = th.k[r]
    end

    function heatflux(∇u, tag)
        k = kdict[tag]
        return k * ∇u
    end
    f(x) = 0.0
    nh = length(hsurf)
    hh = zeros(nh)
    for i in 1:nh
        hh[i] = th.h[ hsurf[i] ]
    end

    function a(u,v)
        z = ∫( r*∇(v)⋅(heatflux∘(∇(u), tags) ) )*dΩ
        for i in 1:nh
            z = z + ∫(r*v*u*hh[i])*dΓh[i]
        end
        return z
    end
    #return
    b(v) = ∫(r*v*f)*dΩ

    op = AffineFEOperator(a,b,U,V)

    ls = LUSolver()
    solver = LinearFESolver(ls)

    uh = solve(solver, op)

    return ( u=uh, Ω=(Ω,dΩ), Γh=(Γh,dΓh), Γt=(Γt,dΓt) )

end

function process_thermistor(th, fem)

    (;u, Ω, Γh, Γt) = fem
    # Calculate surface areas (m²)
    Ah = 2π .* [sum(∫(r)*dΓ) for dΓ in Γh[2]]
    At = 2π .* [sum(∫(r)*dΓ) for dΓ in Γt[2]]

    # Calculate the mean temperature on the surfaces
    Th = 2π .* [sum(∫( r*u ) *dΓ ) for dΓ in Γh[2]] ./ Ah

    # Heat flux from surfaces:
    # Convection surfaces:
    Qh = Float64[]
    nh = length(th.hsurf)
    for i = 1:nh
        hs = th.hsurf[i]
        h = th.h[hs]
        dΓ = Γh[2][i]
        push!(Qh, 4π*h * sum(∫( r*u ) * dΓ))
    end
    # Heat flux from prescribed temperature regions
    Qt = Float64[]
    nt = length(th.tsurf)
    for i in 1:nt
        s = th.tsurf[i]
        k = th.k[th.surfreg[s]]
        Γ  = Γt[1][i]
        dΓ = Γt[2][i]
        n⃗ = get_normal_vector(Γ)
        push!(Qt, -4π*k * sum( ∫( r*(∇(u)⋅n⃗ ) ) * dΓ ))
    end

    # Heat flux from convection surfaces calculated from heat conduction
    Qh1 = Float64[]
    for i in 1:nh
        s = th.tsurf[i]
        k = th.k[th.surfreg[s]]
        Γ  = Γh[1][i]
        dΓ = Γh[2][i]
        n⃗ = get_normal_vector(Γ)
        push!(Qh1, -4π*k * sum( ∫( r*(∇(u)⋅n⃗ ) ) * dΓ ))
    end

    return Th, Qh, Qt, Qh1, Ah, At

end


text:
##### This file contain routines related to evaluating polynomials or interpolators.

# the second (more stable) Barycentric interpolation formula.
function Barycentric2nditp(x::T, 𝓧, 𝓨::Vector{T})::T where T
    w = get𝑤(𝓧)
    return Barycentric2nditp(x,w,𝓧,𝓨)
end

function Barycentric2nditp(x::T, 𝑤::Vector{T}, 𝓧, 𝓨::Vector{T})::T where T

    tmp_numerator = zero(T)
    tmp_denominator = zero(T)
    for k = 1:length(𝓧)
        tmp = x-𝓧[k]
        if isnumericallyclose(tmp,zero(T))
            return 𝓨[k]
        end

        tmp_numerator += 𝓨[k]*𝑤[k]/(x-𝓧[k])
        tmp_denominator += 𝑤[k]/(x-𝓧[k])
    end

    @assert tmp_denominator != zero(T)
    return tmp_numerator/tmp_denominator
end

##### other evaluations.
function evalunivariateLagrangepolynomial(  x::T,
                                            𝓧::Vector{T},
                                            a::Vector{T})::T where T
    out = zero(T)
    for k = 1:length(a)
        out += a[k]*getkthLangrangebasis(x,𝓧,k)
    end

    return out
end

function getkthLangrangebasis(x::T, 𝓧::Vector{T}, k::Int)::T where T
    Lp1 = length(𝓧)

    numerator = one(T)
    denominator = one(T)
    for i = 1:Lp1
        if i != k
            numerator *= x-𝓧[i]
            denominator *= 𝓧[k]-𝓧[i]
        end
    end

    return numerator/denominator
end

##### routines related to the evaluation of Chebyshev polynomials at nodes.

function evalunivariateChebyshevpolynomial(c::Vector{T},
                                            X::Vector{T}) where T

    return collect( evalunivariateChebyshevpolynomial(c,x) for x in X )
end

# # naive implementation.
function evalunivariateChebyshevpolynomial( c::Vector{T},
                                            x::T) where T
    L = length(c) - 1

    out = zero(T)
    for k = 0:L
        out += c[k+1]*Chebyshev1st(x,k)
    end

    return out
end

function evalunivariateChebyshevpolynomialparallel( c::Vector{T},
                                                    x::T) where T
    L = length(c) - 1

    out = @distributed (+) for k = 0:L
        c[k+1]*Chebyshev1st(x,k)
    end

    return out
end

# Clenshaw implementation. does evalunivariateChebyshevpolynomial()
function clenshaweval(c::Vector{T},x::T) where T
    two = one(T) + one(T)

    L = length(c) - 1

    b_kp2 = zero(T)
    b_kp1 = c[end]

    for k = L-1:-1:1
        b_k = two*x*b_kp1 - b_kp2 + c[k+1]

        # update
        b_kp2 = b_kp1
        b_kp1 = b_k
    end

    return c[1] +x*b_kp1 - b_kp2
end

function getChebyshevinterpolator(  f::Function,
                                    L::Int)
    Lp1 = L + 1
    ν = collect( Chebyshev2ndnode(k,L) for k = 0:L )

    return getChebyshevinterpolator(f,L,ν,ν)
end

function getChebyshevinterpolator(  f::Function,
                                    L::Int, a::T, b::T) where T<:Real
    Lp1 = L + 1
    ν = collect( convert(T,Chebyshev2ndnode(k,L)) for k = 0:L )

    # rescale nodes to [a,b].
    νf = collect( scalefromChebyinterval(ν[i],a,b) for i = 1:length(ν) )

    return getChebyshevinterpolator(f,L,νf,ν)
end

# uses Float64 SharedArrays.
function getChebyshevinterpolatorfast(  f::Function,
                                    L::Int)
    Lp1 = L + 1
    ν = SharedArray{Float64}(L+1)
    @sync @distributed for k = 0:L
        ν[k+1] = convert(Float64, Chebyshev2ndnode(k,L))
    end

    return getChebyshevinterpolatorfast(f,L,ν,ν)
end

function getChebyshevinterpolatorfast(  f::Function,
                                    L::Int, a_inp, b_inp)
    Lp1 = L + 1
    ν = SharedArray{Float64}(L+1)
    @sync @distributed for k = 0:L
        ν[k+1] = convert(Float64, Chebyshev2ndnode(k,L))
    end

    # rescale nodes to [a,b].
    a = convert(T,a_inp)
    b = convert(T,b_inp)

    νf = SharedArray{Float64}(L+1)
    @sync @distributed for i = 1:length(ν)
        νf[i] = convert(Float64, scalefromChebyinterval(ν[i],a,b))
    end

    return getChebyshevinterpolatorfast(f,L,νf,ν)
end

# # naive implementation.
function evalunivariateChebyshevpolynomial( c::Vector{T},
                                            x::T,
                                            a::T,
                                            b::T) where T
    L = length(c) - 1

    v = scaletoChebyinterval(x,a,b)

    out = zero(T)
    for k = 0:L
        out += c[k+1]*Chebyshev1st(v,k)
    end

    return out
end

function evalunivariateChebyshevpolynomialderivative(a::Vector{T}, x::T) where T
    L = length(a) - 1
    c = collect( k*a[k+1] for k = 1:L)

    out = zero(T)
    for i = 0:L-1
        out += c[i+1]*Chebyshev2nd(x, i)
    end

    return out
end

function evalunivariateChebyshevpolynomialderivative(c::Vector{T}, X::Vector{T}) where T
    return collect( evalunivariateChebyshevpolynomialderivative(c,x) for x in X )
end

function evalChebyshevpolynomialinterval(   x::T,
                                            d_set::Vector{Vector{T}},
                                            X_set::Vector{Vector{T}}) where T

    for i = 1:length(X_set)
        X = X_set[i]
        if X[1] <= x <= X[end]
            return evalunivariateChebyshevpolynomial(d_set[i], x, X[1], X[end])
        end
    end

    return convert(T,-Inf) # when x isn't in the intervals specified by X_set
end

##### others.

# integer exponents.
# c has the format: [c_0; c_1; ...; c_{L}]
function evalunivariatepolynomial(c::Vector{T}, x::T)::T where T
    L = length(c) - 1

    out = c[1]
    for l = 1:L
        out += c[l+1]x^l
    end

    return out
end

max_stars_repo_name: RoyCCWang/FIRRemez

max_issues_repo_name: RoyCCWang/FIRRemez

max_forks_repo_name: RoyCCWang/FIRRemez


text:
module TestTreeLasCxx

using Test
import Libdl
import TreeLas

const cxx_jl = joinpath(dirname(pathof(TreeLas)), "cxx.jl")

@static if !Sys.iswindows()

@testset "TreeLas.Cxx                    " begin
    @assert isfile(cxx_jl)
    include(cxx_jl)
    @assert isfile(Cxx.lib) "Call `Pkg.build(\"TreeLas\")`"

    syms = Set((String(m.captures[1])
                for m in eachmatch(r"(_Z\w+)", read(open(cxx_jl), String))))
    @test syms == Set()
    syms = Set([
        "tree_dp_f64_i32",
        "tree_dual_f64_i32",
    ])
    Libdl.dlopen(Cxx.lib) do tl
        @testset "Check DLL symbols" begin
            for s in syms
                @testset "Check symbol $s" begin
                    @test Libdl.dlsym_e(tl, s) != Ptr{Nothing}()
                end
            end
        end
    end
end
end

end

max_stars_repo_name: EQt/treelas

max_issues_repo_name: EQt/treelas

max_forks_repo_name: EQt/treelas


text:
defaultfigure(;kwargs...) = Figure(
    ;resolution=(800, 800),
    background=RGBA(0, 0, 0, 0),
    kwargs...)

# ## Plotting interface definition

"""
    plotsample(method, sample)

Plot a `sample` of a `method`.

Learning methods should implement [`plotsample!`](#) to make this work.
"""
function plotsample(method, sample; figkwargs...)
    f = defaultfigure(; figkwargs...)
    plotsample!(f, method, sample)
    return f
end

"""
    plotsample!(f, method, sample)

Plot a `sample` of a `method` type on a Makie.jl figure or
axis `f`. See also [`plotsample`](#).

Part of the plotting interface for learning methods.
"""
function plotsample! end

"""
    plotxy(method, x, y)

Plot a single pair of encoded input `x` and encoded target `y`.
Requires [`plotxy!`](#) to be implemented.
"""
function plotxy(method, x, y; figkwargs...)
    f = defaultfigure(; figkwargs...)
    plotxy!(f, method, x, y)
    return f
end
"""
    plotxy!(f, method, x, y)
"""
function plotxy! end

function plotprediction(method, x, ŷ, y; figkwargs...)
    f = defaultfigure(; figkwargs...)
    plotprediction!(f, method, x, ŷ, y)
    return f
end

"""
    plotprediction(method, x, ŷ, y)
    plotprediction!(f, method, x, ŷ, y)

Plot a comparison of model output `ŷ` with the ground truth `y` on
Makie.jl figure or axis `f`. `x` is the model input.
"""
function plotprediction! end

"""
    plotpredictions(method, xs, ŷs, ys)

Plot a comparison of batches of model outputs `ŷs` with the ground truths
`ys` on Makie.jl figure or axis `f`. `xs` is a batch of model inputs.

[`plotprediction!`](#) needs to be implemented for `method`.
"""
plotpredictions(method, xs, ŷs, ys) = plotpredictions!(defaultfigure(), method, xs, ŷs, ys)

function plotpredictions!(f, method, xs, ŷs, ys)
    n = size(xs)[end]
    nrows = Int(ceil(sqrt(n)))
    is = Iterators.product(1:nrows, 1:nrows)
    for (i, (x, ŷ, y)) in zip(is, DataLoaders.obsslices((xs, ŷs, ys)))
        plotprediction!(f[i...], method, x, ŷ, y)
    end
    return f
end

"""
    plotbatch(method, xs, ys)

Plot an encoded batch of data in a grid.

[`plotxy!`](#) needs to be implemented for `method`.
"""
plotbatch(method, xs, ys) = plotbatch!(defaultfigure(), method, xs, ys)

function plotbatch!(f, method, xs, ys)
    n = DataLoaders._batchsize(xs, DataLoaders.BatchDimLast())
    nrows = Int(ceil(sqrt(n)))
    is = Iterators.product(1:nrows, 1:nrows)
    for (i, (x, y)) in zip(is, DataLoaders.obsslices((xs, ys)))
        plotxy!(f[i...], method, x, y)
    end
    return f
end
# TODO: implement `plotbatch(method, dataloader; n)`
"""
    plotsamples(method::LearningMethod, samples)

Plot a vector of samples for a `LearningMethod` in a grid. Use
[`plotbatch`](#) instead if you want to plot an encoded batch of data.

[`plotsample!`](#) needs to be implemented for `method`.
"""
function plotsamples(method::LearningMethod, samples)
    n = length(samples)
    nrows = Int(ceil(sqrt(n)))
    f = Figure()
    is = Iterators.product(1:nrows, 1:nrows)
    for (i, sample) in zip(is, samples)
        plotsample!(f[i...], method, sample)
    end
    return f
end

# ## Utilities

"""
    imageaxis(f)

Create a `Makie.Axis` with no interactivity, decorations and aspect distortion
suitable for showing images on.
"""
function imageaxis(f; kwargs...)
    ax = Makie.Axis(f; kwargs...)
    ax.aspect = DataAspect()
    ax.xzoomlock = true
    ax.yzoomlock = true
    ax.xrectzoom = false
    ax.yrectzoom = false
    ax.panbutton = nothing
    ax.xpanlock = true
    ax.ypanlock = true
    ax.bottomspinevisible = false
    ax.leftspinevisible = false
    ax.rightspinevisible = false
    ax.topspinevisible = false
    MakieLayout.tightlimits!(ax)
    hidedecorations!(ax)

    return ax
end

#imageaxis(f::Makie.GridPosition; kwargs...) = imageaxis(f.fig; kwargs...)

# ## Plot recipes

@recipe(PlotImage, image) do scene
    Attributes(
        alpha = 1
    )
end

function Makie.plot!(plot::PlotImage)
    im = plot[:image]
    rim = @lift alphacolor.(copy(rotr90($im)), $(plot.attributes[:alpha]))
    image!(plot, rim; plot.attributes...)
    return plot
end

@recipe(PlotMask, mask, classes) do scene
    Attributes()
end

function Makie.plot!(plot::PlotMask; kwargs...)
    mask = plot[:mask]
    classes = try
        classes = plot[:classes]
    catch
        classes = @lift unique($mask)
    end
    im = @lift maskimage($mask, $classes)
    plotimage!(plot, im; alpha = 1, plot.attributes...)
    return plot
end

function maskimage(mask, classes)
    classtoidx = Dict(class => i for (i, class) in enumerate(classes))
    colors = distinguishable_colors(length(classes), transform=deuteranopic)
    return map(x -> colors[classtoidx[x]], mask)
end

maskimage(mask::AbstractArray{<:Gray{T}}, args...) where T =
    maskimage(reinterpret(T, mask), args...)
maskimage(mask::AbstractArray{<:Normed{T}}, args...) where T =
    maskimage(reinterpret(T, mask), args...)

_toindex(v) = CartesianIndex(Tuple(round.(Int, v)))
_boxIs(I, r) = I-(r*CartesianIndex(1, 1)):I+(r*CartesianIndex(1, 1))
function _drawkeypoint!(img, v; r = 2, c = RGB(0, 0, 1))
    Is = _boxIs(_toindex(v), r)
    for I in Is
        checkbounds(Bool, img, I) && (img[I] = c)
    end
    return img
end


text:
module BlockArraysCooTests

using Test
using Gridap.Arrays
using Gridap.Algebra
using BlockArrays
using LinearAlgebra
using Gridap.Arrays: BlockArrayCooMap

# Check MultiLevelBlockedUnitRange

r1 = Base.OneTo(2)
r2 = Base.OneTo(3)

r = append_ranges([r1,r2])
@test length(r) == 5
@test num_blocks_equal(r,r)

s1 = Base.OneTo(1)
s2 = Base.OneTo(2)
s3 = Base.OneTo(0)

s = append_ranges([s1,s2,s3])
@test length(s) == 3
@test !num_blocks_equal(r,s)

o = append_ranges([r,s])
@test length(o) == 8
@test blocks_equal(o,o)
@test num_blocks_equal(o,o)
@test o.local_ranges[1] === r
@test o.local_ranges[2] === s

a = append_ranges([r,s])
@test blocks_equal(o,a)
@test num_blocks_equal(o,a)

#using BenchmarkTools
#@btime blocks_equal($o,$a)
#@btime num_blocks_equal($o,$a)

b = append_ranges([s,s])
@test !blocks_equal(a,b)
@test BlockArrays.blocksize(a) == BlockArrays.blocksize(b) # This only refers to the top level blocks
@test !num_blocks_equal(a,b) # This takes into account all levels

r1_axes = (r1,)
r2_axes = (r2,)
s1_axes = (s1,)
s2_axes = (s2,)

r_axes = (append_ranges(map(first,[r1_axes,r2_axes])),)
s_axes = (append_ranges(map(first,[s1_axes,s2_axes])),)
o_axes = (append_ranges(map(first,[r_axes,s_axes])),)
@test blocks_equal(o_axes,o_axes)
@test num_blocks_equal(o_axes,o_axes)

#using BenchmarkTools
#@btime blocks_equal($o_axes,$o_axes)
#@btime num_blocks_equal($o_axes,$o_axes)

# BlockArrayCoo

# Minimal constructor

blocks = Matrix{Float64}[ [1 2; 3 4], [5 6; 7 8; 9 10] ]
blockids = [(1,1),(2,1)]
ax = (blockedrange([2,3]), blockedrange([2,4]))
a = BlockArrayCoo(ax,blockids,blocks)
a = BlockArrayCooMap((2,2),blockids)(ax,blocks...)

k = BlockArrayCooMap((2,2),blockids)
@test blocksize(k) == (2,2)
@test is_zero_block(k,1,1) == false
@test is_zero_block(k,2,1) == false
@test is_zero_block(k,1,2) == true
@test is_zero_block(k,2,2) == true
@test is_nonzero_block(k,1,1) == true
@test is_nonzero_block(k,first(eachblockid(k)))

# Specific API

@test is_zero_block(rand(4,5),1,1) == false
@test is_zero_block(a,2,2) == true
@test is_zero_block(a,2,1) == false
@test is_nonzero_block(a,2,1) == true
@test is_zero_block(a,Block(1,2)) == true
@test is_zero_block(a,Block(1),Block(2)) == true

# AbstractBlockArray interface

@test a[Block(1),Block(1)] === blocks[1]
@test a[Block(1,1)] === blocks[1]
@test a[Block(1,2)] === a.zero_blocks[1]

c = copy(a[Block(1,1)])
fill!(c,0)

getblock!(c,a,1,1)
@test c == a[Block(1,1)]
@test a[Block(1,1)] === blocks[1]

c = copy(a)
@test isa(c,BlockArrayCoo)
@test c == a
c[Block(1,1)] = a[Block(1,1)]
@test c[Block(1,1)] === a[Block(1,1)]
@test a[Block(1,1)] === blocks[1]
c[1,2] = 3
@test c[1,2] == 3

#using BenchmarkTools
#@btime getindex($c,1,2)
#bi = BlockIndex((1,1),(1,2))
#@btime getindex($c,$bi)

@test a[BlockIndex((1,1),(2,1))] === blocks[1][2,1]
@test a[BlockIndex(1,2),BlockIndex(1,1)] === blocks[1][2,1]
@test a[2,1] === blocks[1][2,1]
@test a[3,2] === blocks[2][1,2]
@test axes(a) === ax
@test size(a) == (5,6)

for (i,b) in enumerateblocks(a)
  @test a[i] === b
end

b21 = zeros(3,2)
getblock!(b21,a,Block(2,1))
@test b21 == a[Block(2,1)]

b12 = ones(2,4)
getblock!(b12,a,Block(1,2))
@test b12 == a[Block(1,2)]

blocks = Vector{Float64}[ [1,2,3] ]
blockids = [(2,)]
ax = (blockedrange([2,3,4]),)
a = BlockArrayCoo(ax,blockids,blocks)

@test a[Block(1)] === a.zero_blocks[1]
@test a[Block(2)] === blocks[1]
@test a[Block(3)] === a.zero_blocks[2]
@test a[BlockIndex(2,3)] === blocks[1][3]

blocks = Matrix{Float64}[ [1 2; 3 4], [5 6 7 8; 9 10 11 12; 13 14 15 16], [1 2 3 4; 5 6 7 8], [1 2 3; 4 5 6; 7 8 9] ]
blockids = [(1,1),(2,2),(1,2),(3,3)]
ax = (blockedrange([2,3,3]), blockedrange([2,4,3]))
a = BlockArrayCoo(ax,blockids,blocks)

blocks = Vector{Float64}[ 10*[1,2], 20*[1,2,3] ]
blockids = [(1,),(3,)]
ax = (blockedrange([2,4,3]),)
b = BlockArrayCoo(ax,blockids,blocks)

# similar

c = similar(a)
@test typeof(c) == typeof(a)
@test axes(a) == axes(c)
@test a.blockids == c.blockids

c = similar(a,Int)
@test axes(a) == axes(c)
@test a.blockids == c.blockids

ax = (blockedrange([2,3,3]), blockedrange([2,4,3]))
c = similar(a,eltype(a),ax)
@test typeof(c) == typeof(a)
@test axes(a) == axes(c)
@test a.blockids == c.blockids

# In this case, we can still preserve the zero block structure
ax = (blockedrange([2,5,3]), blockedrange([2,4,1]))
c = similar(a,eltype(a),ax)
@test typeof(c) == typeof(a)
@test axes(a) != axes(c)
@test axes(c) === ax
@test a.blockids == c.blockids
fill_entries!(c,2)

# In this case, we cannot preserve the zero block structure
ax = (blockedrange([2,5,3,5]), blockedrange([2,4,1]))
c = similar(a,eltype(a),ax)
@test typeof(c) == typeof(a)
@test axes(a) != axes(c)
@test axes(c) === ax
@test a.blockids != c.blockids
@test length(c.blockids) == 12

# zero

c = similar(a)
fill_entries!(c,1)

#using BenchmarkTools
#@btime fill_entries!($c,1)

z = zero(c)
@test isa(z,BlockArrayCoo)
@test length(z.blocks) == 0

# + and -

c = z + a
@test isa(c,BlockArrayCoo)
@test c == a
@test Array(c) == Array(a)

c = z - a
@test isa(c,BlockArrayCoo)
@test c == -a
@test Array(c) == Array(-a)

# mat vec *

c = a*b
@test isa(c,BlockArrayCoo)
@test axes(c,1) === axes(a,1)
@test blocksize(c) == (3,)
@test Array(a)*Array(b) == c

mul!(c,a,b)
@test axes(c,1) === axes(a,1)
@test blocksize(c) == (3,)
@test Array(a)*Array(b) == c

d = copy(c)
mul!(d,a,b,2,3)
@test axes(d,1) === axes(a,1)
@test blocksize(d) == (3,)
@test 2*Array(a)*Array(b) + 3*Array(c) == d

#using BenchmarkTools
#@btime mul!($c,$a,$b)
#@btime mul!($d,$a,$b,2,3)

# mat mat *

b = transpose(a)
@test isa(b,BlockArrayCoo)
c = a*b
@test isa(c,BlockArrayCoo)
@test axes(c,1) === axes(a,1)
@test axes(c,2) === axes(a,1)
@test blocksize(c) == (3,3)
@test Array(a)*Array(b) == c

mul!(c,a,b)
@test axes(c,1) === axes(a,1)
@test axes(c,2) === axes(a,1)
@test blocksize(c) == (3,3)
@test Array(a)*Array(b) == c

#using BenchmarkTools
#@btime mul!($c,$a,$b)

# Cached array
cc = CachedArray(c)

# Compatible axes, we do not modify the cache
axs = (blockedrange([2,3,3]), blockedrange([2,3,3]))
@test blocks_equal(axs,axes(c))
setaxes!(cc,axs)
@test map(blocklasts,axes(cc.array)) == map(blocklasts,axs)
@test cc.array === c

#using BenchmarkTools
#@btime setaxes!($cc,$axs)

# Not compatible axes but same number of blocks
# we create a new cache, but we can preserve the zero block structure
axs = (blockedrange([4,5,3]), blockedrange([2,4,3]))
@test !blocks_equal(axs,axes(c))
@test num_blocks_equal(axs,axes(c))
setaxes!(cc,axs)
@test map(blocklasts,axes(cc.array)) == map(blocklasts,axs)
@test cc.array != c
@test cc.array.blockids == c.blockids
scale_entries!(cc.array,2)

# Not compatible axes and different number of blocks
# we create a new cache without preserving zero block structure
axs = (blockedrange([5,4,3,4]), blockedrange([4,2,3]))
@test !blocks_equal(axs,axes(cc.array))
@test !num_blocks_equal(axs,axes(cc.array))
setaxes!(cc,axs)
@test map(blocklasts,axes(cc.array)) == map(blocklasts,axs)
@test cc.array != c
@test cc.array.blockids != c.blockids

axs1 = (blockedrange([5,4,3]), blockedrange([4,2,3]))
axs2 = (blockedrange([4,5,3]), blockedrange([2,4,3]))
axs3 = (blockedrange([4,5,3]), blockedrange([2,4,3]))
@test blocks_equal(axs1,axs2) == false
@test blocks_equal(axs2,axs2)
@test blocks_equal(axs2,axs3)

blocks = [ 10*[1,2], 20*[1,2,3] ]
blockids = [(1,),(3,)]
axs = (blockedrange([2,4,3]),)
b = BlockArrayCoo(axs,blockids,blocks)

cb = CachedArray(b)

setaxes!(cb,axs)
@test map(blocklasts,axes(cb.array)) == map(blocklasts,axs)
@test cb.array === b

axs = (blockedrange([3,2,3]),)
setaxes!(cb,axs)
@test map(blocklasts,axes(cb.array)) == map(blocklasts,axs)
@test size(cb) == (8,)

c = copy(a)
@test isa(c,BlockArrayCoo)
@test c == a
@test axes(c) == axes(a)

fill!(c,0)
copyto!(c,a)
@test c == a
@test axes(c) == axes(a)

c = 2*a
@test isa(c,BlockArrayCoo)
@test 2*Array(a) == c

c = a*2
@test isa(c,BlockArrayCoo)
@test 2*Array(a) == c

c = a + a
@test isa(c,BlockArrayCoo)
@test 2*Array(a) == c

blocks = [ [1 2; 3 4], [5 6 7 8; 9 10 11 12; 13 14 15 16], [1 2 3 4; 5 6 7 8], [1 2 3; 4 5 6; 7 8 9] ]
blockids = [(1,1),(2,2),(1,2),(3,3)]
ax = (blockedrange([2,3,3]), blockedrange([2,4,3]))
a = BlockArrayCoo(ax,blockids,blocks)

blocks = [ [1 2; 3 4], [1 2; 4 5; 8 9], [5 6 7 8; 9 10 11 12; 13 14 15 16] ]
blockids = [(1,1),(2,1),(3,2)]
a = BlockArrayCoo(ax,blockids,blocks)

c = a - 2*a
@test isa(c,BlockArrayCoo)
@test Array(a)-2*Array(a) == c

@test a[1,2,1] == a[1,2]

blocks = [ 10*[1,2], 20*[1,2,3] ]
blockids = [(1,),(3,)]
axs = (blockedrange([2,4,3]),)
b = BlockArrayCoo(axs,blockids,blocks)

@test b[2] == b[2,1]

# Blocks of Blocks

axs = (blockedrange([2,4,3]),)
bL = BlockArrayCoo(axs,[(1,)],[[1,2]])
bR = BlockArrayCoo(axs,[(1,),(3,)],[[2,3],[4,5,6]])

axsS = (append_ranges(map(first,[axs,axs])),)

bS = BlockArrayCoo(axsS,[(1,),(2,)],[bL,bR])

zS = zero(bS)
@test isa(zS,BlockArrayCoo)
@test isa(zS[Block(1)],BlockArrayCoo)
@test isa(zS[Block(2)],BlockArrayCoo)
@test length(zS.blocks) == 0
@test length(zS[Block(1)].blocks) == 0
@test length(zS[Block(2)].blocks) == 0

cS = 2*bS
@test isa(cS,BlockArrayCoo)
@test isa(cS[Block(1)],BlockArrayCoo)

dS = BlockArrayCoo(axsS,[(2,)],[bR])

r = cS - dS
@test isa(r,BlockArrayCoo)
@test isa(r[Block(1)],BlockArrayCoo)

r = cS + dS
@test isa(r,BlockArrayCoo)
@test isa(r[Block(1)],BlockArrayCoo)

# Block of block matrices

ax = (blockedrange([2,3]), blockedrange([2,4]))
aLL = BlockArrayCoo(ax,[(1,1),(2,1)],[[1 2; 3 4],[5 6; 7 8; 9 10]])
aLR = BlockArrayCoo(ax,[(2,2),(2,1)],[[1 2 5 6; 3 4 1 2; 1 2 3 4],[5 6; 7 8; 9 10]])
aRR = BlockArrayCoo(ax,[(1,1),(2,1)],[[1 2; 3 4],[5 6; 7 8; 9 10]])

axsS = ( append_ranges([ax[1],ax[1]]) , append_ranges([ax[2],ax[2]]))
aS = BlockArrayCoo(axsS,[(1,1),(1,2),(2,2)],[aLL,aLR,aRR])
@test is_zero_block(aS,2,1)
@test is_nonzero_block(aS,2,2)

aS2 = copy(aS)
@test isa(aS2,BlockArrayCoo)
@test isa(aS2[Block(2),Block(2)],BlockArrayCoo)
@test is_zero_block(aS2,2,1)
@test is_nonzero_block(aS2,2,2)

aSt = transpose(aS)
@test Array(aSt) == transpose(Array(aS))
@test isa(aSt,BlockArrayCoo)
@test isa(aSt[Block(2),Block(2)],BlockArrayCoo)
@test is_zero_block(aSt,1,2)
@test is_nonzero_block(aSt,2,2)

rS = aS*aSt
@test rS == Array(aS)*Array(aSt)
@test isa(rS,BlockArrayCoo)
@test isa(rS[Block(2),Block(2)],BlockArrayCoo)

mul!(rS,aS,aSt)
@test rS == Array(aS)*Array(aSt)
@test isa(rS,BlockArrayCoo)
@test isa(rS[Block(2),Block(2)],BlockArrayCoo)

#using BenchmarkTools
#@btime mul!($rS,$aS,$aSt)

axs =(blockedrange([2,4]),)
bL = BlockArrayCoo(axs,[(1,)],[[1,2]])
bR = BlockArrayCoo(axs,[(1,),(2,)],[[2,3],[4,5,6,8]])
axsS = (append_ranges(map(first,[axs,axs])),)
bS = BlockArrayCoo(axsS,[(1,),(2,)],[bL,bR])

rS = aS*bS
@test rS == Array(aS)*Array(bS)
@test isa(rS,BlockArrayCoo)
@test isa(rS[Block(2)],BlockArrayCoo)

mul!(rS,aS,bS)
@test rS == Array(aS)*Array(bS)
@test isa(rS,BlockArrayCoo)
@test isa(rS[Block(2)],BlockArrayCoo)

#using BenchmarkTools
#@btime mul!($rS,$aS,$bS)

cS = copy(rS)
mul!(rS,aS,bS,3,2)
@test rS == 3*Array(aS)*Array(bS) + 2*cS
@test isa(rS,BlockArrayCoo)
@test isa(rS[Block(2)],BlockArrayCoo)

#using BenchmarkTools
#@btime mul!($rS,$aS,$bS,3,2)

#using BenchmarkTools
#@btime Arrays._same_axes($axsA,$axsA)

end # module


text:
struct ActiveState
    weighted_label::Float64
    weight::Float64
    lower::Int64
    upper::Int64

end

function merge_state(l::ActiveState, r::ActiveState)
    return ActiveState(l.weighted_label + r.weighted_label,
                       l.weight + r.weight,
                       l.lower,
                       r.upper)
end

function below(l::ActiveState, r::ActiveState)
    return l.weighted_label * r.weight <= l.weight * r.weighted_label
end

function active_set_isotonic_regression!(y::Vector{Float64}, weights::Vector{Float64})

    n = length(y)
    if n <= 1
        return y
    end
    if n != length(weights)
        throw(DimensionMismatch("Lengths of values and weights mismatch"))
    end

    @inbounds begin
        active_set = [ActiveState(weights[i] * y[i], weights[i], i, i) for i in 1 : n]
        current = 1
        while current < size(active_set, 1)
            while current < size(active_set, 1) && below(active_set[current], active_set[current+1])
                current += 1
            end
            if current == size(active_set, 1)
                break
            end

            merged = merge_state(active_set[current], active_set[current+1])
            splice!(active_set, current:current+1, [merged])
            while current > 1 && !below(active_set[current-1], active_set[current])
                current -= 1
                merged = merge_state(active_set[current], active_set[current+1])
                splice!(active_set, current:current+1, [merged])
            end
        end

        for as in active_set
            y[as.lower:as.upper] .= as.weighted_label / as.weight
        end
    end
    return y
end

active_set_isotonic_regression!(y::Vector{Float64}) = active_set_isotonic_regression!(y, fill!(similar(y), one(eltype(y))))

# non-mutating versions
active_set_isotonic_regression(y::Vector{Float64}, weights::Vector{Float64}) = active_set_isotonic_regression!(copy(y), weights)
active_set_isotonic_regression(y::Vector{Float64}) = active_set_isotonic_regression(y, fill!(similar(y), one(eltype(y))))