text:
using ArrayLayouts, LinearAlgebra, FillArrays, Test
import ArrayLayouts: ApplyBroadcastStyle, QRCompactWYQLayout, QRCompactWYLayout, QRPackedQLayout, QRPackedLayout

@testset "Ldiv" begin
    @testset "Float64 \\ *" begin
        A = randn(5,5)
        b = randn(5)
        M = Ldiv(A,b)

        @test ndims(M) == 1
        @test size(M) == (5,)
        @test similar(M) isa Vector{Float64}
        @test materialize(M) isa Vector{Float64}
        @test all(materialize(M) .=== (A\b) .=== ldiv(A,b))

        @test Base.BroadcastStyle(typeof(Ldiv(A,b))) isa ApplyBroadcastStyle

        @test all(copyto!(similar(b), Ldiv(A,b)) .===
                    (similar(b) .= Ldiv(A,b)) .=== materialize(Ldiv(A,b)) .===
                  (A\b) .=== (b̃ =  copy(b); LAPACK.gesv!(copy(A), b̃); b̃))

        @test copyto!(similar(b), Ldiv(UpperTriangular(A) , b)) ≈ UpperTriangular(A) \ b
        @test all(copyto!(similar(b), Ldiv(UpperTriangular(A) , b)) .===
                    (similar(b) .= Ldiv(UpperTriangular(A),b)) .===
                    BLAS.trsv('U', 'N', 'N', A, b) )

        @test copyto!(similar(b), Ldiv(UpperTriangular(A)' , b)) ≈ UpperTriangular(A)' \ b
        @test all(copyto!(similar(b), Ldiv(UpperTriangular(A)' , b)) .===
                    (similar(b) .= Ldiv(UpperTriangular(A)',b)) .===
                    copyto!(similar(b), Ldiv(transpose(UpperTriangular(A)) , b)) .===
                            (similar(b) .= Ldiv(transpose(UpperTriangular(A)),b)) .===
                    BLAS.trsv('U', 'T', 'N', A, b))
    end

    @testset "ComplexF64 \\ *" begin
        T = ComplexF64
        A = randn(T,5,5)
        b = randn(T,5)
        @test all(copyto!(similar(b), Ldiv(A,b)) .===
                    (similar(b) .= Ldiv(A,b)) .===
                  (A\b) .=== (b̃ =  copy(b); LAPACK.gesv!(copy(A), b̃); b̃))

        @test copyto!(similar(b), Ldiv(UpperTriangular(A) , b)) ≈ UpperTriangular(A) \ b
        @test all(copyto!(similar(b), Ldiv(UpperTriangular(A) , b)) .===
                    (similar(b) .= Ldiv(UpperTriangular(A),b)) .===
                    BLAS.trsv('U', 'N', 'N', A, b) )

        @test copyto!(similar(b), Ldiv(UpperTriangular(A)' , b)) ≈ UpperTriangular(A)' \ b
        @test all(copyto!(similar(b), Ldiv(UpperTriangular(A)' , b)) .===
                    (similar(b) .= Ldiv(UpperTriangular(A)',b)) .===
                    BLAS.trsv('U', 'C', 'N', A, b))

        @test copyto!(similar(b), Ldiv(transpose(UpperTriangular(A)) , b)) ≈ transpose(UpperTriangular(A)) \ b
        @test all(copyto!(similar(b), Ldiv(transpose(UpperTriangular(A)) , b)) .===
                            (similar(b) .= Ldiv(transpose(UpperTriangular(A)),b)) .===
                    BLAS.trsv('U', 'T', 'N', A, b))
    end

    @testset "BigFloat Triangular \\" begin
        for T in (Float64, ComplexF64)
            A = big.(randn(T,5,5))
            b = big.(randn(T,5))
            @test ArrayLayouts.ldiv!(UpperTriangular(A),copy(b)) ≈ UpperTriangular(A)\b
            @test ArrayLayouts.ldiv!(UnitUpperTriangular(A),copy(b)) ≈ UnitUpperTriangular(A)\b
            @test ArrayLayouts.ldiv!(LowerTriangular(A),copy(b)) ≈ LowerTriangular(A)\b
            @test ArrayLayouts.ldiv!(UnitLowerTriangular(A),copy(b)) ≈ UnitLowerTriangular(A)\b

            @test ArrayLayouts.ldiv!(UpperTriangular(A)',copy(b)) ≈ UpperTriangular(A)'\b
            @test ArrayLayouts.ldiv!(UnitUpperTriangular(A)',copy(b)) ≈ UnitUpperTriangular(A)'\b
            @test ArrayLayouts.ldiv!(LowerTriangular(A)',copy(b)) ≈ LowerTriangular(A)'\b
            @test ArrayLayouts.ldiv!(UnitLowerTriangular(A)',copy(b)) ≈ UnitLowerTriangular(A)'\b

            @test ArrayLayouts.ldiv!(transpose(UpperTriangular(A)),copy(b)) ≈ transpose(UpperTriangular(A))\b
            @test ArrayLayouts.ldiv!(transpose(UnitUpperTriangular(A)),copy(b)) ≈ transpose(UnitUpperTriangular(A))\b
            @test ArrayLayouts.ldiv!(transpose(LowerTriangular(A)),copy(b)) ≈ transpose(LowerTriangular(A))\b
            @test ArrayLayouts.ldiv!(transpose(UnitLowerTriangular(A)),copy(b)) ≈ transpose(UnitLowerTriangular(A))\b

            B = big.(randn(T,5,5))
            @test ArrayLayouts.ldiv!(UpperTriangular(A),copy(B)) ≈ UpperTriangular(A)\B
        end
    end

    @testset "Triangular \\ matrix" begin
        A = randn(5,5)
        b = randn(5,5)
        M =  Ldiv(UpperTriangular(A), b)
        @test Base.Broadcast.broadcastable(M) === M
        @test UpperTriangular(A) \ b ≈ copyto!(similar(b) , Ldiv(UpperTriangular(A), b)) ≈
            (b .= Ldiv(UpperTriangular(A), b))
    end

    @testset "Int" begin
        A = [1 2 ; 3 4]; b = [5,6];
        @test eltype(Ldiv(A, b)) == Float64
    end

    @testset "Rdiv" begin
        @testset "Float64 \\ *" begin
            A = randn(3,5)
            B = randn(5,5)
            M = Rdiv(A,B)

            @test eltype(M) == Float64
            @test ndims(M) == 2
            @test size(M) == (3,5)
            @test axes(M) == (Base.OneTo(3),Base.OneTo(5))
            @test similar(M) isa Matrix{Float64}
        end
    end

    @testset "Diagonal" begin
        D = Diagonal(randn(5))
        F = Eye(5)
        A = randn(5,5)
        @test copy(Ldiv(D,F)) == ldiv(D,F) == D \ F
        @test ldiv(F,D) == F \ D
        @test ldiv(D,A) ≈ D \ A
        @test ldiv(A,D) == A \ D
        @test ldiv(F,A) ≈ F \ A
        @test ldiv(A,F) == A \ F

        @test copy(Rdiv(D,F)) == rdiv(D,F) == D / F
        @test rdiv(F,D) == F / D
        @test rdiv(A,D) ≈ A / D
        @test rdiv(A,F) == A / F
    end

    @testset "QR" begin
        @testset "QRCompactWYQ" begin
            for T in (Float64, ComplexF64)
                A = randn(T,10,10)
                b = randn(T,10)
                B = randn(T,10,10)
                F = qr(A)
                Q = F.Q
                @test MemoryLayout(F) isa QRCompactWYLayout
                @test MemoryLayout(Q) isa QRCompactWYQLayout
                @test all(ArrayLayouts.lmul!(Q,copy(b)) .=== lmul!(Q,copy(b)))
                @test all(ArrayLayouts.lmul!(Q',copy(b)) .=== ArrayLayouts.ldiv!(Q,copy(b)) .=== lmul!(Q',copy(b)))
                @test all(ArrayLayouts.lmul!(Q,copy(B)) .=== lmul!(Q,copy(B)))
                @test all(ArrayLayouts.lmul!(Q',copy(B)) .=== ArrayLayouts.ldiv!(Q,copy(B)) .=== lmul!(Q',copy(B)))
                @test all(ArrayLayouts.rmul!(copy(B),Q) .=== rmul!(copy(B),Q))
                @test all(ArrayLayouts.rmul!(copy(B),Q') .=== ArrayLayouts.rdiv!(copy(B),Q) .=== rmul!(copy(B),Q'))
                @test ArrayLayouts.ldiv!(F,copy(b)) ≈ ldiv!(F,copy(b)) # only approx since we use BLAS.trsv!
                @test ArrayLayouts.ldiv!(F,copy(B)) ≈ ldiv!(F,copy(B)) # only approx since we use BLAS.trsv!

                @test copyto!(similar(b), Lmul(Q,b)) == Q*b
                @test copyto!(similar(B), Lmul(Q,B)) == Q*B
            end
        end

        @testset "QRPacked" begin
            @testset "Blas square" begin
                for T in (Float64, ComplexF64)
                    A = randn(T,10,10)
                    b = randn(T,10)
                    B = randn(T,10,10)
                    F = LinearAlgebra.qrfactUnblocked!(copy(A))
                    Q = F.Q
                    @test MemoryLayout(F) isa QRPackedLayout
                    @test MemoryLayout(Q) isa QRPackedQLayout
                    @test all(ArrayLayouts.lmul!(Q,copy(b)) .=== lmul!(Q,copy(b)))
                    @test all(ArrayLayouts.lmul!(Q',copy(b)) .=== ArrayLayouts.ldiv!(Q,copy(b)) .=== lmul!(Q',copy(b)))
                    @test all(ArrayLayouts.lmul!(Q,copy(B)) .=== lmul!(Q,copy(B)))
                    @test all(ArrayLayouts.lmul!(Q',copy(B)) .=== ArrayLayouts.ldiv!(Q,copy(B)) .=== lmul!(Q',copy(B)))
                    @test all(ArrayLayouts.rmul!(copy(B),Q) .=== rmul!(copy(B),Q))
                    @test ArrayLayouts.rmul!(copy(B),Q') ≈ ArrayLayouts.rdiv!(copy(B),Q) ≈ rmul!(copy(B),Q')
                    @test ArrayLayouts.ldiv!(F,copy(b)) ≈ ldiv!(F,copy(b)) # only approx since we use BLAS.trsv!
                    @test ArrayLayouts.ldiv!(F,copy(B)) ≈ ldiv!(F,copy(B)) # only approx since we use BLAS.trsv!

                    @test copyto!(similar(b), Lmul(Q,b)) == Q*b
                    @test copyto!(similar(B), Lmul(Q,B)) == Q*B
                end
            end
            @testset "BigFloat" begin
                A = BigFloat.(randn(10,10))
                b = BigFloat.(randn(10))
                B = BigFloat.(randn(10,10))
                F = LinearAlgebra.qrfactUnblocked!(copy(A))
                Q = F.Q
                @test MemoryLayout(F) isa QRPackedLayout
                @test MemoryLayout(Q) isa QRPackedQLayout
                @test ArrayLayouts.lmul!(Q,copy(b)) == lmul!(Q,copy(b))
                @test ArrayLayouts.lmul!(Q',copy(b)) == ArrayLayouts.ldiv!(Q,copy(b)) == lmul!(Q',copy(b))
                @test ArrayLayouts.lmul!(Q,copy(B)) == lmul!(Q,copy(B))
                @test ArrayLayouts.lmul!(Q',copy(B)) == ArrayLayouts.ldiv!(Q,copy(B)) == lmul!(Q',copy(B))
                @test ArrayLayouts.rmul!(copy(B),Q) == rmul!(copy(B),Q)
                @test ArrayLayouts.rmul!(copy(B),Q') ≈ ArrayLayouts.rdiv!(copy(B),Q) ≈ rmul!(copy(B),Q')
                @test ArrayLayouts.ldiv!(F,copy(b)) ≈ ldiv!(F,copy(b)) # only approx since we use BLAS.trsv!
                @test ArrayLayouts.ldiv!(F,copy(B)) ≈ ldiv!(F,copy(B)) # only approx since we use BLAS.trsv!

                @test copyto!(similar(b), Lmul(Q,b)) == Q*b
                @test copyto!(similar(B), Lmul(Q,B)) == Q*B
            end
            @testset "rectangular" begin
                for T in (Float64,BigFloat)
                    A = T.(randn(12,10))
                    b = T.(randn(12))
                    B = T.(randn(12,12))
                    F = LinearAlgebra.qrfactUnblocked!(copy(A))
                    Q = F.Q
                    @test all(ArrayLayouts.lmul!(Q,copy(b)) == lmul!(Q,copy(b)))
                    @test all(ArrayLayouts.lmul!(Q',copy(b)) == ArrayLayouts.ldiv!(Q,copy(b)) == lmul!(Q',copy(b)))
                    @test all(ArrayLayouts.lmul!(Q,copy(B)) == lmul!(Q,copy(B)))
                    @test all(ArrayLayouts.lmul!(Q',copy(B)) == ArrayLayouts.ldiv!(Q,copy(B)) == lmul!(Q',copy(B)))
                    @test all(ArrayLayouts.rmul!(copy(B),Q) == rmul!(copy(B),Q))
                    @test ArrayLayouts.rmul!(copy(B),Q') ≈ ArrayLayouts.rdiv!(copy(B),Q) ≈ rmul!(copy(B),Q')
                    @test ArrayLayouts.ldiv!(F,copy(b)) ≈ ldiv!(F,copy(b)) # only approx since we use BLAS.trsv!
                    @test ArrayLayouts.ldiv!(F,copy(B)) ≈ ldiv!(F,copy(B)) # only approx since we use BLAS.trsv!

                    @test_throws BoundsError ArrayLayouts.ldiv!(F,randn(10))

                    A = T.(randn(10,12))
                    b = T.(randn(10))
                    B = T.(randn(10,10))
                    F = LinearAlgebra.qrfactUnblocked!(copy(A))
                    Q = F.Q
                    @test all(ArrayLayouts.lmul!(Q,copy(b)) == lmul!(Q,copy(b)))
                    @test all(ArrayLayouts.lmul!(Q',copy(b)) == ArrayLayouts.ldiv!(Q,copy(b)) == lmul!(Q',copy(b)))
                    @test all(ArrayLayouts.lmul!(Q,copy(B)) == lmul!(Q,copy(B)))
                    @test all(ArrayLayouts.lmul!(Q',copy(B)) == ArrayLayouts.ldiv!(Q,copy(B)) == lmul!(Q',copy(B)))
                    @test all(ArrayLayouts.rmul!(copy(B),Q) == rmul!(copy(B),Q))
                    @test ArrayLayouts.rmul!(copy(B),Q') ≈ ArrayLayouts.rdiv!(copy(B),Q) ≈ rmul!(copy(B),Q')

                    @test_throws DimensionMismatch ArrayLayouts.ldiv!(F,copy(b))
                    @test_throws DimensionMismatch ArrayLayouts.ldiv!(F,copy(B))

                    b = T.(randn(12))
                    B = T.(randn(12,12))
                    @test ArrayLayouts.ldiv!(F,copy(b)) ≈ A\b[1:10]
                    @test ArrayLayouts.ldiv!(F,copy(B)) ≈ A\B[1:10,:]
                end
            end
        end
    end
end


text:
# Summary statistics

"""
    genvar(X)

Compute the generalized sample variance of `X`. If `X` is a vector, one-column matrix,
or other one-dimensional iterable, this is equivalent to the sample variance.
Otherwise if `X` is a matrix, this is equivalent to the determinant of the covariance
matrix of `X`.

!!! note
    The generalized sample variance will be 0 if the columns of the matrix of deviations
    are linearly dependent.
"""
genvar(X::AbstractMatrix) = size(X, 2) == 1 ? var(vec(X)) : det(cov(X))
genvar(itr) = var(itr)

"""
    totalvar(X)

Compute the total sample variance of `X`. If `X` is a vector, one-column matrix,
or other one-dimensional iterable, this is equivalent to the sample variance.
Otherwise if `X` is a matrix, this is equivalent to the sum of the diagonal elements
of the covariance matrix of `X`.
"""
totalvar(X::AbstractMatrix) = sum(var(X, 1))
totalvar(itr) = var(itr)

ext: jl


text:
"""
    is_trivial_block_wrapper(ex)

Returns `true` if `ex` is a "trivial" `:block` expression.  This means
that it has one or two `args`; if two, the first must represent line
number information or be `nothing`.
"""
function is_trivial_block_wrapper(ex::Expr)
    if ex.head == :block
        return length(ex.args) == 1 ||
            (length(ex.args) == 2 && (is_linenumber(ex.args[1]) || ex.args[1]===nothing))
    end
    false
end
is_trivial_block_wrapper(arg::ANY) = false

function is_linenumber(stmt::ANY)
    (isa(stmt, Expr) && (stmt::Expr).head == :line) || isa(stmt, LineNumberNode)
end


text:
# Raytracer.jl
# Raytracing for the generation of photorealistic images in Julia
# Copyright (c) 2021 Samuele Colombo, Paolo Galli

# Unit test file for world.jl

@testset "World" begin
    @testset "RayIntersection" begin
        world = World()

        sphere1 = Sphere(transformation=translation(VEC_X * 2f0))
        sphere2 = Sphere(transformation=translation(VEC_X * 8f0))
        push!(world, sphere1)
        push!(world, sphere2)

        intersection1 = ray_intersection(Ray(Point(0f0, 0f0, 0f0), VEC_X), world)
        @test intersection1 !== nothing
        @test intersection1.world_point ≈ Point(1f0, 0f0, 0f0)

        intersection2 = ray_intersection(Ray(Point(10f0, 0f0, 0f0), -VEC_X), world)

        @test intersection2 !== nothing
        @test intersection2.world_point ≈ Point(9f0, 0f0, 0f0)
    end

    @testset "QuickRayIntersection" begin
        world = World()

        sphere1 = Sphere(transformation=translation(VEC_X * 2))
        sphere2 = Sphere(transformation=translation(VEC_X * 8))
        push!(world, sphere1)
        push!(world, sphere2)

        @test !is_point_visible(world, Point(10.0, 0.0, 0.0), ORIGIN)
        @test !is_point_visible(world, Point(5.0, 0.0, 0.0), ORIGIN)
        @test is_point_visible(world, Point(5.0, 0.0, 0.0), Point(4.0, 0.0, 0.0))
        @test is_point_visible(world, Point(0.5, 0.0, 0.0), ORIGIN)
        @test is_point_visible(world, Point(0.0, 10.0, 0.0), ORIGIN)
        @test is_point_visible(world, Point(0.0, 0.0, 10.0), ORIGIN)

    end
end


text:
import SloanDigitalSkySurvey: SDSS, PSF, WCSUtils
using Base.Test
using DataFrames

# PyPlot is only a dependency for this example, so it's not in the REQUIRE file.
using PyPlot

const field_dir =
  joinpath(Pkg.dir("SloanDigitalSkySurvey"), "dat", "sample_field")
const run_num = "003900"
const camcol_num = "6"
const field_num = "0269"

# Load the catalog and select an object.
cat_df = SDSS.load_catalog_df(field_dir, run_num, camcol_num, field_num);
objid = cat_df[:objid][1]

# Gain and dark variance for all the bands are in a single file.
band_gain, band_dark_variance =
  SDSS.load_photo_field(field_dir, run_num, camcol_num, field_num);

# Load a band's image.  b can be in 1:5.
b = 3
nelec, calib_col, sky_grid, sky_x, sky_y, sky_image, wcs =
  SDSS.load_raw_field(field_dir, run_num, camcol_num, field_num,
                      b, band_gain[b]);

# Mask the image.
nelec_original = deepcopy(nelec);
SDSS.mask_image!(nelec, field_dir, run_num, camcol_num, field_num, b);

# Load the point spread function.
raw_psf_comp =
  SDSS.load_psf_data(field_dir, run_num, camcol_num, field_num, b);

# Display the image with a dot at the object location.
pixel_graph = deepcopy(nelec_original)
pixel_graph_masked = deepcopy(nelec)
clip = 8000 # To see dimmer objects, clip the display at 8000 electrons
pixel_graph[pixel_graph .>= clip] = clip
pixel_graph_masked[pixel_graph_masked .>= clip] = clip

# Get the object location in pixel coordinates.
obj_loc  = Float64[cat_df[cat_df[:objid] .== objid, :ra][1],
                   cat_df[cat_df[:objid] .== objid, :dec][1]]
obj_px = WCSUtils.world_to_pix(wcs, obj_loc)

PyPlot.figure()
PyPlot.plt.subplot(1, 2, 1)
PyPlot.title("Band $b image without masking\nObj $objid")
PyPlot.imshow(pixel_graph', cmap=PyPlot.ColorMap("gray"), interpolation = "nearest")
# PyPlot uses zero indexing when plotting points.
PyPlot.scatter(obj_px[1] - 1, obj_px[2] - 1, marker="o", c="r", s=25)

PyPlot.plt.subplot(1, 2, 2)
PyPlot.title("Band $b image without masking\nObj $objid")
PyPlot.imshow(pixel_graph_masked', cmap=PyPlot.ColorMap("gray"), interpolation = "nearest")
PyPlot.scatter(obj_px[1] - 1, obj_px[2] - 1, marker="o", c="r", s=25)

# Display the point spread function at that object.
raw_psf = PSF.get_psf_at_point(obj_px[1], obj_px[2], raw_psf_comp);
PyPlot.figure()
PyPlot.title("Band $b PSF at pixel $(obj_px)\nObj $objid")
PyPlot.imshow(raw_psf', cmap=PyPlot.ColorMap("gray"), interpolation = "nearest")


text:
using DSP, ImageFiltering, Images, SparseArrays, FFTW, RecipesBase

import DSP.Periodograms: compute_window

function istft(S::AbstractMatrix{Complex{T}}, wlen::Int, overlap::Int=wlen÷2; onesided=true, nfft=nextfastfft(wlen), window=hanning) where T
    winc = wlen-overlap
    win, norm2 = compute_window(window, wlen)
    win² = win.^2
    nframes = size(S,2)-1
    outlen = wlen + nframes*winc
    out = zeros(T, outlen)
    tmp1 = Array{eltype(S)}(undef, size(S, 1))

    if onesided
        tmp2 = zeros(T, nfft)
        p = FFTW.plan_irfft(tmp1, nfft)
    else
        tmp2 = zeros(Complex{T}, nfft)
        p = FFTW.plan_ifft(tmp1)
    end
    wsum = zeros(outlen)
    for k = 1:size(S,2)
        copyto!(tmp1, 1, S, 1+(k-1)*size(S,1), length(tmp1))
        mul!(tmp2, p, tmp1)
        # tmp2 ./= AbstractFFTs.normalization(tmp2, size(tmp2)) # QUESTION: Not sure how to make this work

        ix = (k-1)*winc
        @inbounds for n=1:wlen
            out[ix+n] += real(tmp2[n])*win[n]
            wsum[ix+n] += win²[n]
        end

    end
    mw = median(wsum)
    @inbounds for i = eachindex(wsum,out)
        if wsum[i] >= 1e-3*mw
            out[i] /= wsum[i]
        end
    end
    out
end

function find_extrema(A, n)
    # q = quantile(vec(A), 0.95)
    # A = copy(A)
    # A .= max.(A, q)
    inds = findlocalmaxima(A)
    isempty(inds) && error("Did not find any peaks in the delay/amplitude histogram, try modifying the `kernel_size` or any of the other parameters.")
    amps = A[inds]
    perm = sortperm(amps, rev=true)
    inds[perm[1:min(n,length(perm))]]
end

function separate(αₚ::Vector{<:Real}, δₚ::Vector{<:Real}, S1::Matrix{Complex{T}}, S2, freqs, n; kwargs...) where T
    N_sources = length(αₚ)
    αₚ = @. (αₚ + sqrt(αₚ^2 + 4)) / 2 # Do not inplace this

    cisδf = Vector{Complex{T}}(undef, length(freqs))
    bestsofar = fill(T(Inf), size(S1,1)-1, size(S1,2))
    bestind = zeros(Int, size(S1,1)-1, size(S1,2))
    mask = BitArray{2}(undef, size(bestsofar))
    for i = 1:N_sources
        den = 1 + αₚ[i]^2
        @. cisδf = cis(-freqs * δₚ[i])
        @views score = @. abs2(αₚ[i] * cisδf * S1[2:end,:] - S2[2:end,:]) /
           den
        mask .= score .< bestsofar
        bestind[mask] .= i
        bestsofar[mask] .= score[mask]
    end
    masks = map(1:N_sources) do i
        mask = bestind .== i
    end
    spectrograms = map(1:N_sources) do i
        # mask = imfilter(masks[i], Kernel.gaussian(1))
        mask = masks[i]
        den = 1 + αₚ[i]^2
        # S = [
        #     zeros(eltype(S1), 1, size(S1, 2)) # Add in zero DC component
        #     @. ((S1 + αₚ[i] * cis(δₚ[i] * freqs) * S2) / (1 + αₚ[i]^2)) * mask
        # ]
        S = similar(S1) .= 0
        @. cisδf = cis(freqs * δₚ[i])
        @views @. S[2:end,:] = ((S1[2:end,:] + αₚ[i] * cisδf * S2[2:end,:]) / den) * mask
        S
    end
    est = map(spectrograms) do S
        istft(S, n, n÷2; kwargs...)
    end
    reduce(hcat, est), masks, spectrograms
end

struct DUET
    "histogram"
    A
    "ranges for amp and delay values"
    ar
    "ranges for amp and delay values"
    dr
    "peak locations"
    av
    "peak locations"
    dv
    "amp map"
    α
    "delay map"
    δ
    masks
    spectrograms
end

@recipe function plot(h::DUET)
    xguide --> "δ"
    yguide --> "A"
    title --> "Almplitude-delay histogram"
    @series begin
        seriestype --> :heatmap
        h.dr, h.ar, h.A
    end
    @series begin
        seriestype := :scatter
        label --> "Local maxima"
        markercolor --> :red
        h.dv, h.av
    end
end

"""
    est, H = duet( x1, x2, n_sources, n = 1024;
        p            = 1, # amplitude power used to weight histogram
        q            = 0, # delay power used to weight histogram
        amax         = 0.7,
        dmax         = 3.6,
        abins        = 35, # number of histogram bins
        dbins        = 50, # number of histogram bins
        kernel_size  = 1, # Controls the smoothing of the histogram.
        window       = hanning,
        bigdelay     = false,
        kernel_sizeδ = 0.1,
        kwargs..., # These are sent to the stft function
    )

DUET is an algorithm for blind source separation. It works on stereo mixtures and can separate any number of sources as long as they do not overlap in the time-frequency domain.

## `p` and `q`
The paper gives the following guidelines:
- `p = 0, q = 0`: the counting histogram proposed in the original DUET algorithm
- `p = 1, q = 0`: motivated by the ML symmetric attenuation estimator
- `p = 1, q = 2`: motivated by the ML delay estimator
- `p = 2, q = 0`: in order to reduce delay estimator bias
- `p = 2, q = 2`: for low signal-to-noise ratio or speech mixtures
From the paper referenced below:
"`p = 1, q = 0` is a good default choice. When the sources are not
equal power, we would suggest `p = 0.5, q = 0` as it prevents the dominant
source from hiding the smaller source peaks in the histogram."

For large delays, it appears to be more efficient to set `dmax` to something small (20), `nfft ≈ 8n, p=0.5, q=2`

- `bigdelay` indicates whether or not the two microphones are far apart. If `true`, the delay `δ` is estimated using the differential method (see the paper sec 8.4.1) and the delay map is smoothed using `kernel_sizeδ`.

Implementation based on *Rickard, Scott. (2007). The DUET blind source separation algorithm. 10.1007/978-1-4020-6479-1_8.*
"""
function duet(
    x1,
    x2,
    n_sources,
    n     = 1024;
    nfft  = n,
    p     = 1,
    q     = 0,
    amax  = 0.7,
    dmax  = 3.6,
    abins = 35,
    dbins = 50,
    kernel_size = 1,
    window = hanning,
    onesided = true,
    bigdelay = false,
    kernel_sizeδ = 0.1,
    kwargs...,
)

    S1 = stft(x1, n, n÷2; nfft = nfft, window = window, onesided=onesided, kwargs...)
    S2 = stft(x2, n, n÷2; nfft = nfft, window = window, onesided=onesided, kwargs...)
    # S1, S2 = S1[2:end, :], S2[2:end, :] # remove dc
    S1[1, :] .= 0; S2[1, :] .= 0 # remove dc
    @views S1w, S2w = S1[2:end,:], S2[2:end,:]
    if onesided
        freqs = FFTW.rfftfreq(nfft)[2:end] .* (2pi)
    else
        freqs = FFTW.fftfreq(nfft)[2:end] .* (2pi)
        # freqs = freqs[1:size(S1, 1)]
    end

    R21 = (S2w .+ eps()) ./ (S1w .+ eps()) # ratio of spectrograms
    α = abs.(R21) # relative attenuation
    α = @. α - 1 / α # symmetric attenuation

    if bigdelay
        Δω = 2pi / nfft
        Δω*dmax > π && @warn("frequency resolution not sufficient for the chosen maximum delay.")
        δ0 = @. R21 * conj(S2w)/S1w / abs2(R21)
        if kernel_sizeδ > 0
            δ0 = imfilt(δ0, kernel_sizeδ) # It appears to be slightly better to filter before angle, but more tests needed
        end
        δ = @. -angle(δ0) / freqs # NOTE: I removed the Δω divisor here and it seems to have made it more stable
    else
        δ = @. -angle(R21) / freqs # δ relative delay
    end

    Sweight = @. (abs(S1w) * abs(S2w))^p * abs(freqs)^q # weights

    αmask = @. (abs(α) < amax) & (abs(δ) < dmax)
    α_vec = α[αmask]
    δ_vec = δ[αmask]
    Sweight = Sweight[αmask]
    # determine histogram indices
    αind = @. round(Int, 1 + (abins - 1) * (α_vec + amax) / (2amax))
    δind = @. round(Int, 1 + (dbins - 1) * (δ_vec + dmax) / (2dmax))
    # FULL-SPARSE TRICK TO CREATE 2D WEIGHTED HISTOGRAM
    A = Matrix(sparse(αind, δind, Sweight, abins, dbins))
    if kernel_size > 0
        A = imfilt(A, kernel_size)
    end
    # abins, dbins = size(A)
    ar,dr = LinRange(-amax , amax , abins ), LinRange(-dmax , dmax , dbins )
    inds = find_extrema(A, n_sources)
    av = [ar[i[1]] for i in inds]
    dv = [dr[i[2]] for i in inds]

    est, masks, spectrograms = separate(av, dv, S1, S2, freqs, n; (nfft=nfft, window=window, onesided=onesided, kwargs...)...)
    H = DUET(A,ar,dr,av,dv,α,δ,masks,spectrograms)

    est, H
end

function imfilt(x,kernel_size)
    K = KernelFactors.gaussian(kernel_size)
    imfilter(x, kernelfactors((K,K)))
end

function imfilt1(x,kernel_size)
    K = KernelFactors.gaussian(kernel_size)
    imfilter(x, K[:,1:1])
end

"""
    mix = mixture(signals, amps, [delays::Vector{Int}])
    mixes = mixture(signals, amps::Vector{Vector}, [delays::Vector{Vector{Int}}])

Mix together signals using amplitudes `amps` and `delays` (`delays` is specified in samples). If `amps` and the optional `delays` are vectors of vectors, then a vector of mixtures is returned. A vector of vectors is converted to a matrix using `M = reduce(hcat, mixes)`.

The returned signals will be shorter than the input signals corresponding to the maximum difference between delays, e.g., if `delays = [-2,0,5]`, the returned mixture will be 7 samples shorter.
"""
function mixture(signals, amps, delays::Vector{Int}=zeros(Int, length(amps)))
    length(signals) == length(amps) == length(delays) || throw(ArgumentError("The length of the inputs must be the same"))
    N = length(signals[1])
    mind, maxd = extrema(delays)
    maxdiff = maxd-mind
    Nn = N-maxdiff-mind
    ind = -mind+1:Nn
    inds = [ind .+ d for d in delays]
    sum(amps[i]*signals[i][inds[i]] for i in eachindex(signals))
end

function mixture(signals, amps::Vector{<:Vector{<:Number}}, delays::Vector{Vector{Int}}=[zeros(Int, length(a)) for a in amps])
    Nmix = length(amps)
    N = length(signals[1])
    mind, maxd = extrema(reduce(vcat,delays))
    maxdiff = maxd-mind
    Nn = N-maxdiff-mind
    ind = -mind+1:Nn
    map(amps, delays) do amps, delays
        length(signals) == length(amps) == length(delays) || throw(ArgumentError("The length of the inputs must be the same"))
        inds = [ind .+ d for d in delays]
        sum(amps[i]*signals[i][inds[i]] for i in eachindex(signals))
    end
end


text:
######################################################
# Created by Alex Dowling (alexdowling.net)
# while at the University of Wisconsin-Madison
# ####################################################

# Define problem
using JuMP
using Ipopt
using DegeneracyHunter

# Define problem
function lp1()

	m = Model()

	@variable(m, 0 <= x[1:3] <= 5)
	@variable(m, 0 <= y <= 0)
	@constraint(m, x[1] + x[2] >= 1)
	@constraint(m, x[1] + x[2] + x[3] == 1)
	@constraint(m, x[2] - 2*x[3] <= 1)
	@constraint(m, x[1] + x[3] >= 1)
	@constraint(m, x[1] + x[2] + x[3] == 1) # Redundant constraint - makes problem degenerate
	@objective(m, Min, sum{x[i],i=1:3})

	return m
end

function initialize!(m::Model)
	x = getindex(m,:x)
	setvalue(x[1], 1.0)
	setvalue(x[2], 5.0)
	setvalue(x[3], -1.0)

	# Intentionally do not initialize y

	return nothing
end

# Create and initialize model
m = lp1()
JuMP.setsolver(m,IpoptSolver())
initialize!(m)

# Check initialization
DegeneracyHunter.printVariableDiagnostics(m)

# Print equations with absolute residuals greater than 0.001
DegeneracyHunter.printInfeasibleEquations(m, 0.001)

# Check for degeneracy constraints
DegeneracyHunter.degeneracyHunter(m, IpoptSolver(), includeBounds=true)

# Solve model
t0 = time()
status = solve(m)
tm = time() - t0

# Print inactive constraints at solution
DegeneracyHunter.printInactiveEquations(m)

# Package and print problem size information
ps = DegeneracyHunter.assembleProblemStats(m,status,tm)
DegeneracyHunter.printProblemStats(ps)


text:
# Autogenerated wrapper script for IOAPI_jll for x86_64-apple-darwin-libgfortran5
export airs2m3, bcwndw, camxtom3, datshift, dayagg, factor, findwndw, greg2jul, gregdate, gridprobe, insertgrid, jul2greg, juldate, juldiff, julshift, kfxtract, latlon, libioapi, m3agmask, m3agmax, m3combo, m3cple, m3diff, m3edhdr, m3fake, m3hdr, m3interp, m3mask, m3merge, m3pair, m3probe, m3stat, m3totxt, m3tproc, m3tshift, m3wndw, m3xtract, mpasdiff, mpasstat, mpastom3, mtxblend, mtxbuild, mtxcalc, mtxcple, presterp, presz, projtool, randomstat, selmrg2d, timediff, timeshift, vertimeproc, vertintegral, vertot, wndwdesc, wrfgriddesc, wrftom3

using NetCDF_jll
using NetCDFF_jll
using CompilerSupportLibraries_jll
JLLWrappers.@generate_wrapper_header("IOAPI")
JLLWrappers.@declare_executable_product(airs2m3)
JLLWrappers.@declare_executable_product(bcwndw)
JLLWrappers.@declare_executable_product(camxtom3)
JLLWrappers.@declare_executable_product(datshift)
JLLWrappers.@declare_executable_product(dayagg)
JLLWrappers.@declare_executable_product(factor)
JLLWrappers.@declare_executable_product(findwndw)
JLLWrappers.@declare_executable_product(greg2jul)
JLLWrappers.@declare_executable_product(gregdate)
JLLWrappers.@declare_executable_product(gridprobe)
JLLWrappers.@declare_executable_product(insertgrid)
JLLWrappers.@declare_executable_product(jul2greg)
JLLWrappers.@declare_executable_product(juldate)
JLLWrappers.@declare_executable_product(juldiff)
JLLWrappers.@declare_executable_product(julshift)
JLLWrappers.@declare_executable_product(kfxtract)
JLLWrappers.@declare_executable_product(latlon)
JLLWrappers.@declare_library_product(libioapi, "@rpath/libioapi.dylib")
JLLWrappers.@declare_executable_product(m3agmask)
JLLWrappers.@declare_executable_product(m3agmax)
JLLWrappers.@declare_executable_product(m3combo)
JLLWrappers.@declare_executable_product(m3cple)
JLLWrappers.@declare_executable_product(m3diff)
JLLWrappers.@declare_executable_product(m3edhdr)
JLLWrappers.@declare_executable_product(m3fake)
JLLWrappers.@declare_executable_product(m3hdr)
JLLWrappers.@declare_executable_product(m3interp)
JLLWrappers.@declare_executable_product(m3mask)
JLLWrappers.@declare_executable_product(m3merge)
JLLWrappers.@declare_executable_product(m3pair)
JLLWrappers.@declare_executable_product(m3probe)
JLLWrappers.@declare_executable_product(m3stat)
JLLWrappers.@declare_executable_product(m3totxt)
JLLWrappers.@declare_executable_product(m3tproc)
JLLWrappers.@declare_executable_product(m3tshift)
JLLWrappers.@declare_executable_product(m3wndw)
JLLWrappers.@declare_executable_product(m3xtract)
JLLWrappers.@declare_executable_product(mpasdiff)
JLLWrappers.@declare_executable_product(mpasstat)
JLLWrappers.@declare_executable_product(mpastom3)
JLLWrappers.@declare_executable_product(mtxblend)
JLLWrappers.@declare_executable_product(mtxbuild)
JLLWrappers.@declare_executable_product(mtxcalc)
JLLWrappers.@declare_executable_product(mtxcple)
JLLWrappers.@declare_executable_product(presterp)
JLLWrappers.@declare_executable_product(presz)
JLLWrappers.@declare_executable_product(projtool)
JLLWrappers.@declare_executable_product(randomstat)
JLLWrappers.@declare_executable_product(selmrg2d)
JLLWrappers.@declare_executable_product(timediff)
JLLWrappers.@declare_executable_product(timeshift)
JLLWrappers.@declare_executable_product(vertimeproc)
JLLWrappers.@declare_executable_product(vertintegral)
JLLWrappers.@declare_executable_product(vertot)
JLLWrappers.@declare_executable_product(wndwdesc)
JLLWrappers.@declare_executable_product(wrfgriddesc)
JLLWrappers.@declare_executable_product(wrftom3)
function __init__()
    JLLWrappers.@generate_init_header(NetCDF_jll, NetCDFF_jll, CompilerSupportLibraries_jll)
    JLLWrappers.@init_executable_product(
        airs2m3,
        "bin/airs2m3",
    )

    JLLWrappers.@init_executable_product(
        bcwndw,
        "bin/bcwndw",
    )

    JLLWrappers.@init_executable_product(
        camxtom3,
        "bin/camxtom3",
    )

    JLLWrappers.@init_executable_product(
        datshift,
        "bin/datshift",
    )

    JLLWrappers.@init_executable_product(
        dayagg,
        "bin/dayagg",
    )

    JLLWrappers.@init_executable_product(
        factor,
        "bin/factor",
    )

    JLLWrappers.@init_executable_product(
        findwndw,
        "bin/findwndw",
    )

    JLLWrappers.@init_executable_product(
        greg2jul,
        "bin/greg2jul",
    )

    JLLWrappers.@init_executable_product(
        gregdate,
        "bin/gregdate",
    )

    JLLWrappers.@init_executable_product(
        gridprobe,
        "bin/gridprobe",
    )

    JLLWrappers.@init_executable_product(
        insertgrid,
        "bin/insertgrid",
    )

    JLLWrappers.@init_executable_product(
        jul2greg,
        "bin/jul2greg",
    )

    JLLWrappers.@init_executable_product(
        juldate,
        "bin/juldate",
    )

    JLLWrappers.@init_executable_product(
        juldiff,
        "bin/juldiff",
    )

    JLLWrappers.@init_executable_product(
        julshift,
        "bin/julshift",
    )

    JLLWrappers.@init_executable_product(
        kfxtract,
        "bin/kfxtract",
    )

    JLLWrappers.@init_executable_product(
        latlon,
        "bin/latlon",
    )

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

    JLLWrappers.@init_executable_product(
        m3agmask,
        "bin/m3agmask",
    )

    JLLWrappers.@init_executable_product(
        m3agmax,
        "bin/m3agmax",
    )

    JLLWrappers.@init_executable_product(
        m3combo,
        "bin/m3combo",
    )

    JLLWrappers.@init_executable_product(
        m3cple,
        "bin/m3cple",
    )

    JLLWrappers.@init_executable_product(
        m3diff,
        "bin/m3diff",
    )

    JLLWrappers.@init_executable_product(
        m3edhdr,
        "bin/m3edhdr",
    )

    JLLWrappers.@init_executable_product(
        m3fake,
        "bin/m3fake",
    )

    JLLWrappers.@init_executable_product(
        m3hdr,
        "bin/m3hdr",
    )

    JLLWrappers.@init_executable_product(
        m3interp,
        "bin/m3interp",
    )

    JLLWrappers.@init_executable_product(
        m3mask,
        "bin/m3mask",
    )

    JLLWrappers.@init_executable_product(
        m3merge,
        "bin/m3merge",
    )

    JLLWrappers.@init_executable_product(
        m3pair,
        "bin/m3pair",
    )

    JLLWrappers.@init_executable_product(
        m3probe,
        "bin/m3probe",
    )

    JLLWrappers.@init_executable_product(
        m3stat,
        "bin/m3stat",
    )

    JLLWrappers.@init_executable_product(
        m3totxt,
        "bin/m3totxt",
    )

    JLLWrappers.@init_executable_product(
        m3tproc,
        "bin/m3tproc",
    )

    JLLWrappers.@init_executable_product(
        m3tshift,
        "bin/m3tshift",
    )

    JLLWrappers.@init_executable_product(
        m3wndw,
        "bin/m3wndw",
    )

    JLLWrappers.@init_executable_product(
        m3xtract,
        "bin/m3xtract",
    )

    JLLWrappers.@init_executable_product(
        mpasdiff,
        "bin/mpasdiff",
    )

    JLLWrappers.@init_executable_product(
        mpasstat,
        "bin/mpasstat",
    )

    JLLWrappers.@init_executable_product(
        mpastom3,
        "bin/mpastom3",
    )

    JLLWrappers.@init_executable_product(
        mtxblend,
        "bin/mtxblend",
    )

    JLLWrappers.@init_executable_product(
        mtxbuild,
        "bin/mtxbuild",
    )

    JLLWrappers.@init_executable_product(
        mtxcalc,
        "bin/mtxcalc",
    )

    JLLWrappers.@init_executable_product(
        mtxcple,
        "bin/mtxcple",
    )

    JLLWrappers.@init_executable_product(
        presterp,
        "bin/presterp",
    )

    JLLWrappers.@init_executable_product(
        presz,
        "bin/presz",
    )

    JLLWrappers.@init_executable_product(
        projtool,
        "bin/projtool",
    )

    JLLWrappers.@init_executable_product(
        randomstat,
        "bin/randomstat",
    )

    JLLWrappers.@init_executable_product(
        selmrg2d,
        "bin/selmrg2d",
    )

    JLLWrappers.@init_executable_product(
        timediff,
        "bin/timediff",
    )

    JLLWrappers.@init_executable_product(
        timeshift,
        "bin/timeshift",
    )

    JLLWrappers.@init_executable_product(
        vertimeproc,
        "bin/vertimeproc",
    )

    JLLWrappers.@init_executable_product(
        vertintegral,
        "bin/vertintegral",
    )

    JLLWrappers.@init_executable_product(
        vertot,
        "bin/vertot",
    )

    JLLWrappers.@init_executable_product(
        wndwdesc,
        "bin/wndwdesc",
    )

    JLLWrappers.@init_executable_product(
        wrfgriddesc,
        "bin/wrfgriddesc",
    )

    JLLWrappers.@init_executable_product(
        wrftom3,
        "bin/wrftom3",
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
using TreeShaker
using Test

@testset "TreeShaker.jl" begin
    # Write your own tests here.
end


text:
using Revise, Pkg

#NOTE: Consider getting both indstl and fs records in ccm table

if  pwd() ∉ LOAD_PATH
  push!(LOAD_PATH,pwd())
end

Pkg.activate(pwd())
using QFactors

#QFactors.recompressqfactorfiles()
#QFactors.recompressqfactorfile("CRSP-D")
#QFactors.recompressqfactorfile("CRSP-M")

#refreshcomp (PreprocessCOMP): refreshes the compustat data from the csv.lz4 file
#refreshcrsp (PreprocessCRSP): refreshes the crsp data from the csv.lz4 file
#refreshmerge[1] (MergeCRSPCOMP): refreshes the compustat-crsp merge.
#refreshfactors (FormFactors): refreshes the compustat-crsp merge.
#[1]True if refreshcomp, refreshcrsp, or validatemerged are true
#[2]True if refreshcomp, refreshcrsp, refreshmerge, or validatemerged are true

#=full flow (updated 7/27)
Key: function (FileName)
... means substantial code and/or subordinate (minor) function calls

formfactors (FormFactors)
  makeuniv (MergeCRSPCOMP)
    loadcrspdata (PreprocessCRSP)
      preprocesscrsp
        ...
    loadcompdata (PreprocessCOMP)
      preprocesscomp(2 arg)
        preprocesscomp(1 arg)
        ...
        dedupcomp!
        ...
        formannouncementintervals!
      mergecompccm (MergeCOMPCCM)
        loadccm
        preprocessccm!
        ...
        reconcilecompccmintervals!
    mergecrspcomp!
      testmerged
  formportfolios!

=#

function qfactorscript()
  QFactors.formfactors(
    refreshcomp = true,
    refreshcrsp = false,
    refreshmerge = true,
    validatemerged = true,
    refreshfactors = true,
    parallel = true)

  QFactors.analyzeqfactors()
end

@time qfactorscript()

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
frac_orig_size = 0.2

@testset "Recursive splitting" begin
    @testset "E = $E" for E in 2:4
        @testset "k = $k" for k in 2:3
            canonical_simplex_vertices = zeros(E + 1, E)
            canonical_simplex_vertices[2:(E + 1), :] = eye(E)
            canonical_simplex_vol = abs(det(hcat(canonical_simplex_vertices, ones(E + 1))))

            # Refine the canonical simplex
            simplex_indices = zeros(Int, 1, E + 1)
            simplex_indices[1, :] = collect(1:E+1)

            centroids, radii = centroids_radii2(canonical_simplex_vertices, simplex_indices)
            volumes_before = simplex_volumes(canonical_simplex_vertices, simplex_indices)

            # Generic splitting rules for this combination of dimension and splitting
            # factor.
            radiusmax = maximum(radii)

            maxradius_allowed = radiusmax * 0.4
            refined = refine_recursive(canonical_simplex_vertices, simplex_indices, maxradius_allowed, 2)

            radiusmax_after_refinement = maximum(refined[4])
            volumes_after_refinement = refined[5]

            @testset "Desired refinement level reached" begin
                @test radiusmax_after_refinement < maxradius_allowed
            end

            @testset "Triangulation volumes are preserved" begin
                @test sum(volumes_before) ≈ sum(volumes_after_refinement)
            end
        end
    end
end

max_stars_repo_name: kahaaga/SimplexSplitting

max_issues_repo_name: kahaaga/SimplexSplitting

max_forks_repo_name: kahaaga/SimplexSplitting


text:
"""
    optimize(model::DEModel, de::DE, n_iter::Int; progress=false, kwargs...)

Finds optimal set of parameters.

# Arguments

- `model`: a model containing likelihood function with data and priors
- `de`: differential evolution object
- `n_iter`: number of iterations or samples

# Keywords

- `progress=false`: show progress of algorithm
- `kwargs...`: optional keyword arguments
"""
optimize(model::DEModel, de::DE, n_iter::Int; progress=false, kwargs...) = _optimize(model::DEModel, de::DE, n_iter::Int; progress, stepfun=step!, kwargs...)

function _optimize(model::DEModel, de::DE, n_iter::Int; progress=false, stepfun=step!, kwargs...)
    meter = Progress(n_iter)
    # initialize particles based on prior distribution
    groups = sample_init(model, de, n_iter)
    for iter in 1:n_iter
        de.iter = iter
        # explicitly pass groups so parallel works
        groups = stepfun(model, de, groups)
        progress ? next!(meter) : nothing
    end
    return vcat(groups...)
end

"""
    optimize(model::DEModel, de::DE, ::MCMCThreads, n_iter::Int; progress=false, kwargs...)

Finds optimal set of parameters.

# Arguments

- `model`: a model containing likelihood function with data and priors
- `de`: differential evolution object
- `MCMCThreads`: pass MCMCThreads() object to run on multiple threads
- `n_iter`: number of iterations or samples

# Keywords

- `progress=false`: show progress of algorithm
- `kwargs...`: optional keyword arguments
"""
function optimize(model::DEModel, de::DE, ::MCMCThreads, n_iter::Int; progress=false, kwargs...)
    _optimize(model::DEModel, de::DE, n_iter::Int; progress, stepfun=pstep!, kwargs...)
end


text:
import Base: +, -
import Base.broadcast
import Base.diff

(+)(ta::TimeArray) = .+ta
(-)(ta::TimeArray) = .-ta

# ND TimeArray <--> MD TimeArray
function broadcast(f, ta1::TimeArray, ta2::TimeArray)
    # first test metadata matches
    meta = ta1.meta == ta2.meta ? ta1.meta : Void

    # determine array widths and name cols accordingly
    w1, w2  = length(ta1.colnames), length(ta2.colnames)
    if w1 == w2
        cnames = [ta1.colnames[i] * "_" * ta2.colnames[i] for i = 1:w1]
    elseif w1 == 1
        cnames = [ta1.colnames[1] * "_" * ta2.colnames[i] for i = 1:w2]
    elseif w2 == 1
        cnames = [ta1.colnames[i] * "_" * ta2.colnames[1] for i = 1:w1]
    else
        error("arrays must have the same number of columns, or one must be a single column")
    end

    # obtain shared timestamp
    idx1, idx2 = overlaps(ta1.timestamp, ta2.timestamp)
    tstamp = ta1[idx1].timestamp

    # retrieve values that match the Int array matching dates
    vals1, vals2 = ta1[idx1].values, ta2[idx2].values

    # compute output values
    vals = broadcast(f, vals1, vals2)
    TimeArray(tstamp, vals, cnames, meta)
end

function broadcast(f, ta::TimeArray, args...)
    vals = broadcast(f, ta.values, args...)
    TimeArray(ta.timestamp, vals, ta.colnames, ta.meta)
end

# FIXME: How to deal with f(Number, Number, ..., TimeArray)?
function broadcast(f, n::Number, ta::TimeArray, args...)
    vals = broadcast(f, n, ta.values, args...)
    TimeArray(ta.timestamp, vals, ta.colnames, ta.meta)
end

###### lag, lead ################

function lag(ta::TimeArray{T, N}, n::Int=1;
             padding::Bool=false, period::Int=0) where {T, N}

    if period != 0
        warn("the period kwarg is deprecated, use lag(ta::TimeArray, period::Int) instead")
        n = period
    end

    if padding
        paddedvals = [NaN*ones(n, length(ta.colnames)); ta.values[1:end-n, :]]
        ta = TimeArray(ta.timestamp, paddedvals, ta.colnames, ta.meta)
    else
        ta = TimeArray(ta.timestamp[1+n:end], ta.values[1:end-n, :], ta.colnames, ta.meta)
    end

    N == 1 && (ta = ta[ta.colnames[1]])
    return ta

end  # lag

function lead(ta::TimeArray{T, N}, n::Int=1;
              padding::Bool=false, period::Int=0) where {T, N}

    if period != 0
      warn("the period kwarg is deprecated, use lead(ta::TimeArray, period::Int) instead")
      n = period
    end

    if padding
        paddedvals = [ta.values[1+n:end, :]; NaN*ones(n, length(ta.colnames))]
        ta = TimeArray(ta.timestamp, paddedvals, ta.colnames, ta.meta)
    else
        ta = TimeArray(ta.timestamp[1:end-n], ta.values[1+n:end, :], ta.colnames, ta.meta)
    end

    N == 1 && (ta = ta[ta.colnames[1]])
    return ta

end  # lead

###### diff #####################

# TODO: Support higher-order differencing?
function diff(ta::TimeArray; padding::Bool=false)
    cols = ta.colnames
    ta = ta .- lag(ta, padding=padding)
    ta.colnames[:] = cols
    return ta
end  # diff

###### percentchange ############

function percentchange(ta::TimeArray, returns::Symbol=:simple;
                       padding::Bool=false, method::AbstractString="")

    if method != ""
        warn("the method kwarg is deprecated, use percentchange(ta, :methodname) instead")
        returns = Symbol(method)
    end

    cols = ta.colnames
    ta = returns == :log ? diff(log.(ta), padding=padding) :
         returns == :simple ? expm1.(percentchange(ta, :log, padding=padding)) :
         error("returns must be either :simple or :log")
    ta.colnames[:] = cols

   return ta

end  # percentchange

###### moving ###################

function moving(ta::TimeArray{T, 1}, f::Function, window::Int;
                padding::Bool=false) where {T}
    tstamps = padding ? ta.timestamp : ta.timestamp[window:end]
    vals    = zeros(ta.values[window:end])
    for i=1:length(vals)
        vals[i] = f(ta.values[i:i+(window-1)])
    end
    padding && (vals = [NaN*ones(window-1); vals])
    TimeArray(tstamps, vals, ta.colnames, ta.meta)
end

function moving(ta::TimeArray{T, 2}, f::Function, window::Int;
                padding::Bool=false) where {T}
    tstamps = padding ? ta.timestamp : ta.timestamp[window:end]
    vals    = zeros(ta.values[window:end, :])
    for i=1:size(vals, 1), j=1:size(vals, 2)
        vals[i, j] = f(ta.values[i:i+(window-1), j])
    end
    padding && (vals = [NaN*ones(ta.values[1:(window-1), :]); vals])
    TimeArray(tstamps, vals, ta.colnames, ta.meta)
end

###### upto #####################

function upto(ta::TimeArray{T, 1}, f::Function) where {T}
    vals = zeros(ta.values)
    for i=1:length(vals)
        vals[i] = f(ta.values[1:i])
    end
    TimeArray(ta.timestamp, vals, ta.colnames, ta.meta)
end

function upto(ta::TimeArray{T, 2}, f::Function) where {T}
    vals = zeros(ta.values)
    for i=1:size(vals, 1), j=1:size(vals, 2)
        vals[i, j] = f(ta.values[1:i, j])
    end
    TimeArray(ta.timestamp, vals, ta.colnames, ta.meta)
end

###### basecall #################

basecall(ta::TimeArray, f::Function; cnames=ta.colnames) =
    TimeArray(ta.timestamp, f(ta.values), cnames, ta.meta)

###### uniform observations #####

function uniformspaced(ta::TimeArray)
    gap1 = ta.timestamp[2] - ta.timestamp[1]
    i, n, is_uniform = 2, length(ta), true
    while is_uniform & (i < n)
        is_uniform = gap1 == (ta.timestamp[i+1] - ta.timestamp[i])
        i += 1
    end
    return is_uniform
end  # uniformspaced

function uniformspace(ta::TimeArray{T, N}) where {T, N}
    min_gap = minimum(ta.timestamp[2:end] - ta.timestamp[1:end-1])
    newtimestamp = ta.timestamp[1]:min_gap:ta.timestamp[end]
    emptyta = TimeArray(collect(newtimestamp), zeros(length(newtimestamp), 0), String[], ta.meta)
    ta = merge(emptyta, ta, :left)
    N == 1 && (ta = ta[ta.colnames[1]])
    return ta
end  # uniformspace

###### dropnan ####################

dropnan(ta::TimeArray, method::Symbol=:all) =
    method == :all ? ta[find(any(.!isnan.(ta.values), 2))] :
    method == :any ? ta[find(all(.!isnan.(ta.values), 2))] :
    error("dropnan method must be :all or :any")


text:
using BlockArrays, FillArrays, Test

@testset "broadcast" begin
    @testset "BlockArray" begin
        A = BlockArray(randn(6), 1:3)

        @test BlockArrays.BroadcastStyle(typeof(A)) == BlockArrays.BlockStyle{1}()
        bc = Base.broadcasted(exp,A)
        @test axes(A) === axes(bc) === axes(similar(bc, Float64))
        @test exp.(A) == exp.(Vector(A))
        @test axes(A) === axes(exp.(A))

        @test A+A isa BlockArray
        @test axes(A + A,1).lasts == axes(A .+ A,1).lasts == axes(A,1).lasts
        @test axes(A .+ 1,1).lasts == axes(A,1).lasts

        A = BlockArray(randn(6,6), 1:3,1:3)

        @test BlockArrays.BroadcastStyle(typeof(A)) == BlockArrays.BlockStyle{2}()

        @test exp.(A) == exp.(Matrix(A))
        @test axes(A) == axes(exp.(A))

        @test axes(A + A) == axes(A .+ A) == axes(A)
        @test axes(A .+ 1) == axes(A)
    end

    @testset "PseudoBlockArray" begin
        A = PseudoBlockArray(randn(6), 1:3)

        @test BlockArrays.BroadcastStyle(typeof(A)) == BlockArrays.PseudoBlockStyle{1}()

        @test exp.(A) == exp.(Vector(A))
        @test axes(A,1).lasts == axes(exp.(A),1).lasts

        @test A+A isa PseudoBlockArray
        @test axes(A + A,1).lasts == axes(A .+ A,1).lasts == axes(A,1).lasts
        @test axes(A .+ 1,1).lasts == axes(A,1).lasts

        B = PseudoBlockArray(randn(6,6), 1:3,1:3)

        @test BlockArrays.BroadcastStyle(typeof(B)) == BlockArrays.PseudoBlockStyle{2}()

        @test exp.(B) == exp.(Matrix(B))
        @test axes(B) == axes(exp.(B))

        @test axes(B + B) == axes(B .+ B) == axes(B)
        @test axes(B .+ 1) == axes(B)
        @test axes(A .+ 1 .+ B) == axes(B)
        @test A .+ 1 .+ B == Vector(A) .+ 1 .+ B == Vector(A) .+ 1 .+ Matrix(B)
    end

    @testset "Mixed" begin
        A = BlockArray(randn(6), 1:3)
        B = PseudoBlockArray(randn(6), 1:3)

        @test A + B isa BlockArray
        @test B + A isa BlockArray

        @test axes(A + B,1).lasts == axes(A,1).lasts

        C = randn(6)

        @test A + C isa BlockVector{Float64}
        @test C + A isa BlockVector{Float64}
        @test B + C isa PseudoBlockVector{Float64}
        @test C + B isa PseudoBlockVector{Float64}

        @test blocksize(A+C) == blocksize(C+A) == blocksize(A)
        @test blocksize(B+C) == blocksize(C+B) == blocksize(B)

        A = BlockArray(randn(6,6), 1:3, 1:3)
        D = Diagonal(ones(6))
        @test blocksize(A + D) == blocksize(A)
        @test blocksize(B .+ D) == (3,1)
    end

    @testset "Mixed block sizes" begin
        A = BlockArray(randn(6), 1:3)
        B = BlockArray(randn(6), fill(2,3))
        @test blocksize(A+B) == (5,)

        A = BlockArray(randn(6,6), 1:3, 1:3)
        B = BlockArray(randn(6,6), fill(2,3), fill(3,2))

        @test blocksize(A+B) == (5,3)
    end

    @testset "UnitRange" begin
        n = 3
        x = mortar([1:4n, 1:n])
        @test eltype(x.blocks) <: UnitRange
        y = 1:length(x)
        z = randn(size(x))
        x2 = vcat(x.blocks...)
        y2 = copy(y)
        z2 = copy(z)
        @test (@. z = x + y + z; z) == (@. z2 = x2 + y2 + z2; z2)
    end

    @testset "Special broadcast" begin
        v = mortar([1:3,4:7])
        @test broadcast(+, v) isa BlockVector{Int,Vector{UnitRange{Int}}}
        @test broadcast(+, v) == v
        @test broadcast(-, v) isa BlockVector{Int,Vector{StepRange{Int,Int}}}
        @test broadcast(-, v) == -v == -Vector(v)
        @test broadcast(+, v, 1) isa BlockVector{Int,Vector{UnitRange{Int}}}
        @test broadcast(+, v, 1) == Vector(v).+1
        @test broadcast(*, 2, v) isa BlockVector{Int,Vector{StepRange{Int,Int}}}
        @test broadcast(*, 2, v) == 2Vector(v)
    end

    @testset "special axes" begin
        A = BlockArray(randn(6), Ones{Int}(6))
        B = BlockArray(randn(6), Ones{Int}(6))
        @test axes(A+B,1) === axes(A,1)

        C = BlockArray(randn(6), (BlockArrays._BlockedUnitRange(1,2:6),))
        @test axes(A+C,1) === BlockArrays._BlockedUnitRange(1,1:6)
    end
end


text:
#!/usr/bin/env julia
import Mortchartgen
using PGFPlotsX
frames = Mortchartgen.load_frames()
cts4564 = Mortchartgen.propplot_sexesyrs("circ", "tum", [2;1], 4290, 15, 18,
	1951:2016, false, frames, "sv", "", false)
cts6574 = Mortchartgen.propplot_sexesyrs("circ", "tum", [2;1], 4290, 19, 20,
	1951:2016, false, frames, "sv", "", false)

pf2pgftable(pf, sex) =
	Table([pf[sex][:propframe][:Year], pf[sex][:propframe][:value]])

p = @pgf Axis({"yticklabel style"={"/pgf/number format/use comma"},
		"xticklabel style"={"/pgf/number format/set thousands separator={}"},
		"legend style"={"font=\\tiny"},
		legend_pos="north east", ylabel="cirkulation/tumör",
		"xtick distance"=10, "ytick distance"=.5,
		title="Antal dödsfall cirkulation/tumör Sverige",
		xmajorgrids, ymajorgrids},
	PlotInc(pf2pgftable(cts4564, 2)), LegendEntry("kvinnor 45–64"),
	PlotInc(pf2pgftable(cts4564, 1)), LegendEntry("män 45–64"),
	PlotInc(pf2pgftable(cts6574, 2)), LegendEntry("kvinnor 65–74"),
	PlotInc(pf2pgftable(cts6574, 1)), LegendEntry("män 65–74"))

pgfsave("../../images/secirctum5116_4574.svg", p)

max_stars_repo_name: klpn/static-dust

max_issues_repo_name: klpn/static-dust

max_forks_repo_name: klpn/static-dust


text:
@testset "plotting" begin
    x = rand(10)
    f = GP(SqExponentialKernel())
    gp = f(x, 0.1)

    z = rand(10)
    plt1 = sampleplot(z, gp)
    @test plt1.n == 1
    @test isequal(plt1.series_list[1].plotattributes[:x], vcat(z, NaN))

    plt2 = sampleplot(gp; samples=3)
    @test plt2.n == 1
    plt2_x = plt2.series_list[1].plotattributes[:x]
    plt2_y = plt2.series_list[1].plotattributes[:y]
    @test isequal(plt2_x, vcat(x, NaN, x, NaN, x, NaN))
    @test length(plt2_y) == length(plt2_x)
    @test isnan(plt2_y[length(z) + 1]) && isnan(plt2_y[2length(z) + 2])

    z3 = rand(7)
    plt3 = sampleplot(z3, f; samples=2)
    @test plt3.n == 1
    @test isequal(plt3.series_list[1].plotattributes[:x], vcat(z3, NaN, z3, NaN))

    # Check recipe dispatches for `FiniteGP`s
    rec = RecipesBase.apply_recipe(Dict{Symbol,Any}(), gp)
    @test length(rec) == 1 && length(rec[1].args) == 2 # one series with two arguments
    @test rec[1].args[1] == x
    @test rec[1].args[2] == gp
    @test isempty(rec[1].plotattributes) # no default attributes

    z = 1 .+ x
    for kwargs in (
        Dict{Symbol,Any}(),
        Dict{Symbol,Any}(:ribbon_scale => 3),
        Dict{Symbol,Any}(:ribbon_scale => rand()),
    )
        scale = get(kwargs, :ribbon_scale, 1.0)
        rec = RecipesBase.apply_recipe(kwargs, z, gp)
        @test length(rec) == 1 && length(rec[1].args) == 2 # one series with two arguments
        @test rec[1].args[1] == z
        @test rec[1].args[2] == zero(x)
        # 3 default attributes
        attributes = rec[1].plotattributes
        @test sort!(collect(keys(attributes))) == [:fillalpha, :linewidth, :ribbon]
        @test attributes[:fillalpha] == 0.3
        @test attributes[:linewidth] == 2
        @test attributes[:ribbon] == scale .* sqrt.(var(gp))
    end

    # Check recipe dispatches for `AbstractGP`s
    # with `AbstractVector` and `AbstractRange`:
    for x in (rand(10), 0:0.01:1)
        rec = RecipesBase.apply_recipe(Dict{Symbol,Any}(), x, f)
        @test length(rec) == 1 && length(rec[1].args) == 1 # one series with one argument
        @test rec[1].args[1] isa AbstractGPs.FiniteGP
        @test rec[1].args[1].x == x
        @test rec[1].args[1].f == f
        @test isempty(rec[1].plotattributes) # no default attributes
    end

    # Checks
    @test_throws DimensionMismatch plot(rand(5), gp)
    @test_throws ErrorException plot(rand(10), gp; ribbon_scale=-0.5)
end

max_stars_repo_name: JuliaGaussianProcesses/AbstractGP

max_issues_repo_name: JuliaGaussianProcesses/AbstractGP

max_forks_repo_name: JuliaGaussianProcesses/AbstractGP


text:
function _precompile_()
    ccall(:jl_generating_output, Cint, ()) == 1 || return nothing
    Base.precompile(Tuple{Core.kwftype(typeof(AcuteML.render2file)),NamedTuple{(:id, :age, :field, :GPA, :courses), Tuple{Int64, Int64, String, Float64, Vector{String}}},typeof(render2file),String,Bool})
    Base.precompile(Tuple{typeof(AcuteML.aml_create),Expr,Vector{Union{Expr, Symbol}},Vector{Any},Vector{Union{Expr, Symbol, Type}},Vector{Union{Expr, Symbol}},Vector{Union{Missing, String}},Vector{Union{Missing, Function, Symbol}},Vector{Union{Missing, Type}},String,Type,Array{Union{Missing, Function, Symbol}, 0},Bool,Vector{Union{Nothing, Expr}},Vector{Union{Nothing, Expr}},Vector{Union{Nothing, Expr}},Expr})
    Base.precompile(Tuple{typeof(AcuteML.aml_create),Expr,Vector{Union{Expr, Symbol}},Vector{Any},Vector{Union{Expr, Symbol, Type}},Vector{Union{Expr, Symbol}},Vector{Union{Missing, String}},Vector{Union{Missing, Function, Symbol}},Vector{Union{Missing, Type}},String,Type,Array{Union{Missing, Function, Symbol}, 0},Bool,Vector{Union{Nothing, Expr}},Vector{Union{Nothing, Expr}},Vector{Union{Nothing, Expr}},Symbol})
    Base.precompile(Tuple{typeof(AcuteML.aml_parse),Expr})
    Base.precompile(Tuple{typeof(AcuteML.multiString),Vector{String}})
    Base.precompile(Tuple{typeof(addnode!),Node,String,Node,Type{AbsNormal}})
    Base.precompile(Tuple{typeof(addnode!),Node,String,Vector{String},Type{AbsText}})
    Base.precompile(Tuple{typeof(findcontent),Type{Vector{Float64}},String,Node,Type{AbsNormal}})
    Base.precompile(Tuple{typeof(findcontent),Type{Vector{Int64}},String,Node,Type{AbsNormal}})
    Base.precompile(Tuple{typeof(findcontent),Type{Vector{String}},String,Node,Type{AbsText}})
    Base.precompile(Tuple{typeof(updatecontent!),Dict{String, String},String,Node,Type{AbsNormal}})
    Base.precompile(Tuple{typeof(updatecontent!),Vector{Any},String,Node,Type{AbsNormal}})
    Base.precompile(Tuple{typeof(updatecontent!),Vector{String},String,Node,Type{AbsText}})
end

max_stars_repo_name: aminya/AML

max_forks_repo_name: aminya/AML


text:
StaticArrays.Size(::Type{LinSRange{T,B,E,L,D}}) where {T,B,E,L,D} = Size{(L,)}()

StaticArrays.Size(::Type{StepSRangeLen{T,Tr,Ts,R,S,L,F}}) where {T,Tr,Ts,R,S,L,F} = Size{(L,)}()

function StaticArrays.Size(::Type{StepSRange{T,Ts,F,S,L}})  where {T,Ts,F,S,L}
    return Size{(start_step_stop_to_length(T, F, S, L),)}()
end

function StaticArrays.Size(::Type{UnitSRange{T,F,L}})  where {T<:Union{UInt,UInt64,UInt128},F,L}
    return Size{(L < F ? zero(T) : Base.Checked.checked_add(L - F, one(T)),)}()
end

function StaticArrays.Size(::Type{UnitSRange{T,F,L}}) where {T<:Union{Int,Int64,Int128},F,L}
    return Size{(Base.Checked.checked_add(Base.Checked.checked_sub(L, F), one(T)),)}()
end

StaticArrays.Size(::Type{OneToSRange{T,L}}) where {T<:Union{Int,Int64},L} = Size{(T(L),)}()

StaticArrays.Size(::Type{OneToSRange{T,L}}) where {T,L} = Size{(Integer(L - zero(T)),)}()

StaticArrays.Size(::Type{T}) where {T<:MRange} = Size{(Dynamic(),)}()

# TODO would be better to have this implemented in StaticArrays
_Size(::Type{T}) where {T<:Tuple} = Size{StaticArrays.get.(Tuple(Length.(T.parameters)))}()

Base.checkindex(::Type{Bool}, r::SRange, i) = _checkindex(Length(r), i)
_checkindex(::Length{L}, i::Integer) where {L} = (i < 1 || i > L) ? false : true
function _checkindex(::Length{L}, i::AbstractVector{<:Integer}) where {L}
    return minimum(i) < 1 || maximum(i) > L ? false : true
end


text:
using PrettyNumber
using Test

# Check test validity
import PrettyNumber.max_int_len

@assert (max_int_len) == (5) "
This test is prepared for other default constants.
"

import PrettyNumber.pretty

@testset "Integers" begin
    @test pretty(1)      == "1"
    @test pretty(-11111) == "-11111"
    @test pretty(-1e4)   == "-10000"
    @test pretty(+1e6)   == "1⋅10⁶"
    @test pretty(-1e6)   == "-1⋅10⁶"
    @test pretty(-5e+60) == "-5⋅10⁶⁰"
end

@testset "Floating point" begin
    @test pretty(3e-4)                == "3⋅10⁻⁴"
    @test pretty(1.289e2)             == "128.9"
    @test pretty(1.22e1)              == "12.2"
    @test pretty(0.22e1)              == "2.2"
    @test pretty(1.12345678)          == "1.123"
    @test pretty(198.12345678)        == "198.123"
    @test pretty(1988.12345678)       == "1.988⋅10³"
    @test pretty(0.00000198812345678) == "1.988⋅10⁻⁶"
end

@testset "General macro" begin
    @test "1"                       |> prettify == "1"
    @test "01"                      |> prettify == "1"
    @test "010"                     |> prettify == "10"
    @test "100000000"               |> prettify == "1⋅10⁸"
    @test "1e20"                    |> prettify == "1⋅10²⁰"
    @test "1e2"                     |> prettify == "100"
    @test "+1"                      |> prettify == "1"
    @test "+01"                     |> prettify == "1"
    @test "+010"                    |> prettify == "10"
    @test "+100000000"              |> prettify == "1⋅10⁸"
    @test "+1e20"                   |> prettify == "1⋅10²⁰"
    @test "+1e+2"                   |> prettify == "100"
    @test "-1"                      |> prettify == "-1"
    @test "-01"                     |> prettify == "-1"
    @test "-010"                    |> prettify == "-10"
    @test "-100000000"              |> prettify == "-1⋅10⁸"
    @test "-1e20"                   |> prettify == "-1⋅10²⁰"
    @test "-1e-2"                   |> prettify == "-0.01"
    @test "+1.42"                   |> prettify == "1.42"
    @test "-01.42"                  |> prettify == "-1.42"
    @test "+010.42"                 |> prettify == "10.42"
    @test "-100000.42"              |> prettify == "-1.000⋅10⁵"
    @test "2.42e-20"                |> prettify == "2.42⋅10⁻²⁰"
    @test "-0.000000242"            |> prettify == "-2.42⋅10⁻⁷"
    @test "+0.000000242"            |> prettify == "2.42⋅10⁻⁷"
end

@testset "Unicode expresions" begin
    @test "ρ ≃ 0.4 ± 2e-5"         |> prettify == "ρ ≃ 0.4 ± 2⋅10⁻⁵"
    @test "ρ ≃ 00.04 ± 2e-3"       |> prettify == "ρ ≃ 0.04 ± 0.002"
    @test "γ=4.05e-20, β=5e-3"     |> prettify == "γ=4.05⋅10⁻²⁰, β=0.005"
    @test "∫σ dx = 0.00004"        |> prettify == "∫σ dx = 4⋅10⁻⁵"
    @test "0.42924042 ± 0.420e-20" |> prettify == "0.429 ± 4.2⋅10⁻²¹"
end

@testset "Escaping sequences" begin
    num = Float64(π)
    @test prettify("$num") == "3.142"
end

printstyled("All test passed ✓\n", color=:green)


text:
# Purpose: Here we provide functions to calculate various transforms of the
# window function, as needed for shot noise, for the SFB transform itself, and
# for mode decoupling.
#
# We define the abstract type 'WindowFunction' so that we can specialize on the
# following subtypes. All of them follow the same philosophy that the window
# function gives the probability of a type of galaxy being in the survey in a
# voxel, and the 'WindowFunction' object acts like a 2D array with the first
# dimension being indexing the radius, the second dimension a healpix on the
# sky.
#
#   * 'WindowFunction3D': This is a window function that specifies the window
#   in every voxel of the survey. The voxels are given at each radius by a
#   healpix map.
#
#   * 'WindowFunction3DSeparable': This window function type assumes that the
#   radial selection and angular mask are separable, i.e., W(r⃗) = M(r̂)⋅ϕ(r).
#
#
# Todo:
#
# - power_win_mix() is rather slow on large problems. Generally, there are
# several ways to improve the speed.
#
#   * First, improve single-core performance.
#
#   * Second, specialize on window functions that have separable radial and
#   angular selection functions, and then take a perturbative or similar
#   approach to calculate corrections.
#
#   * Third, parallelize.
#
#   * Fourth, exploit symmetry of (2ℓ+1)*M.

module Windows

export ConfigurationSpaceModes
export window_r, apply_window, apodize_window
export win_rhat_ln, integrate_window, calc_wmix, power_win_mix, win_lnn
export check_nsamp
export calc_fvol

using ..Modes
#using ..HealPy
using ..SeparableArrays
using ..HealpixHelpers
using Healpix
using LinearAlgebra
using SpecialFunctions
using WignerSymbols
using WignerFamilies
using SparseArrays
using Random

using Distributed
using SharedArrays
using ProgressMeter
#using Base.Threads  # Threads.@threads macro, export JULIA_NUM_THREADS=8

#using QuadGK  # for testing

#using Profile

#SparseArrays.rowvals(mat) = 1:size(mat,1)  # we also want to use it for full vectors

struct ConfigurationSpaceModes{Tarr}
    rmin::Real
    rmax::Real
    Δr::Real
    r::Tarr
    nr::Integer
    npix::Integer
    nside::Integer
end

# for the @__dot syntax:
Base.length(w::ConfigurationSpaceModes) = 1
Base.iterate(w::ConfigurationSpaceModes) = s, nothing
Base.iterate(w::ConfigurationSpaceModes, x) = nothing

@doc raw"""
    ConfigurationSpaceModes(rmin, rmax, nr, nside)

A struct to describe and define the voxelization scheme.
"""
function ConfigurationSpaceModes(rmin, rmax, nr, nside)
    Δr = (rmax - rmin) / nr
    r = range(rmin+Δr/2, rmax-Δr/2, length=nr)  # midpoints
    npix = nside2npix(nside)
    return ConfigurationSpaceModes(rmin, rmax, Δr, r, nr, npix, nside)
end

@doc raw"""
    window_r(wmodes::ConfigurationSpaceModes)

Get the $r$-values of the radial bins and corresponding widths $\Delta r$, e.g.,
```julia
r, Δr = SFB.window_r(wmodes)
```
"""
window_r(wmodes::ConfigurationSpaceModes) = wmodes.r, wmodes.Δr

# only a basic implementation, with lots of edge cases poorly handled
function apodize_window(win, wmodes::ConfigurationSpaceModes, smooth=50.0)
    winapod = deepcopy(win)
    r, Δr = window_r(wmodes)
    nr = length(r)
    npix = size(win,2)
    nside = npix2nside(npix)
    Δi = ceil(Int, smooth / Δr)
    if iseven(Δi)
        Δi += 1
    end
    weights = @. exp(-((1:Δi) - Δi/2)^2 / (2 * 10^2))
    weights .*= Δi / sum(weights)
    @show weights
    for j=1:npix, i=Δi:nr-Δi
        mmin = i - Δi ÷ 2
        mmax = i + Δi ÷ 2
        w = win[mmin:mmax,j] .* weights
        winapod[i,j] = mean(w)
    end
    return winapod
end

function apply_window(rθϕ, win, rmin, rmax, win_r, win_Δr; rng=Random.GLOBAL_RNG)
    nside = npix2nside(size(win,2))
    r_out = Float32[]
    θ_out = Float32[]
    ϕ_out = Float32[]
    r = rθϕ[1,:]
    θ = rθϕ[2,:]
    ϕ = rθϕ[3,:]
    reso = Resolution(nside)
    idx_ang = ang2pixRing.(Ref(reso), θ, ϕ) # .+ 1
    Wmax = maximum(win)
    @show extrema(r)
    for i=1:length(r)
        !(rmin <= r[i] <= rmax) && continue
        idx_r = ceil(Int, (r[i] - rmin) / win_Δr)
        (idx_r == 0) && (idx_r += 1)
        (idx_r == size(win,1)+1) && (idx_r -= 1)
        if rand(rng) <= win[idx_r,idx_ang[i]] / Wmax
            # include in survey
            push!(r_out, r[i])
            push!(θ_out, θ[i])
            push!(ϕ_out, ϕ[i])
        end
    end
    return [r_out θ_out ϕ_out]'
end

apply_window(rθϕ, win, wmodes::ConfigurationSpaceModes; rng=Random.GLOBAL_RNG) = begin
    apply_window(rθϕ, win, wmodes.rmin, wmodes.rmax, wmodes.r, wmodes.Δr; rng)
end

function integrate_window(win, wmodes::ConfigurationSpaceModes)
    nr = size(win, 1)
    npix = size(win,2)
    r = wmodes.r
    Δr = wmodes.Δr
    ΔΩpix = 4*π / npix
    radial = [sum(win[i,:]) for i=1:size(win,1)]
    Veff = Δr * ΔΩpix * sum(@. radial * r^2)
    return Veff
end

function calc_fvol(win, wmodes::ConfigurationSpaceModes; Wthreshold=0.1)
    nr = size(win, 1)
    npix = size(win,2)
    r = wmodes.r
    Δr = wmodes.Δr
    ΔΩpix = 4*π / npix
    radial = [sum(w->w>Wthreshold, win[i,:]) for i=1:size(win,1)]
    V = Δr * ΔΩpix * sum(@. radial * r^2)
    Vsfb = (4*π/3) * (wmodes.rmax^3 - wmodes.rmin^3)
    return V / Vsfb
end

function win_rhat_ln(win, wmodes::ConfigurationSpaceModes, amodes::AnlmModes)
    gnl = amodes.basisfunctions
    nr = size(win, 1)
    r, Δr = window_r(wmodes)
    W_rhat_ln = fill(NaN, size(win,2), amodes.lmax+1, amodes.nmax)
    check_nsamp_1gnl(amodes, wmodes)
    for n=1:amodes.nmax, l=0:amodes.lmax_n[n]
        l==0 && @show n,l
        int_nowin = @. r^2 * gnl(n, l, r)
        W_rhat_ln[:,l+1,n] .= win' * int_nowin
    end
    @. W_rhat_ln *= Δr
    #@assert all(isfinite.(W_rhat_ln))  # not all needs to be finite if we limit by kmax
    return W_rhat_ln
end

# specialize on separable windows
function win_rhat_ln(win::SeparableArray, wmodes::ConfigurationSpaceModes, amodes::AnlmModes)
    gnl = amodes.basisfunctions
    nr = size(win, 1)
    r, Δr = window_r(wmodes)
    W_ln = fill(NaN, amodes.lmax+1, amodes.nmax)
    check_nsamp_1gnl(amodes, wmodes)
    for n=1:amodes.nmax, l=0:amodes.lmax_n[n]
        l==0 && @show n,l,amodes.nmax
        W_ln[l+1,n] = Δr * sum(@. r^2 * gnl(n, l, r) * win.phi)
    end
    return SeparableArray(win.mask, W_ln, name1=:mask, name2=:w_ln)
end

# This should be very performant
function calc_wmix(win, wmodes::ConfigurationSpaceModes, amodes::AnlmModes; neg_m=false)
    nlmsize = getnlmsize(amodes)
    wmix = fill(NaN*im, nlmsize, nlmsize)
    @debug "wmix" length(wmix), size(wmix)

    r, Δr = window_r(wmodes)
    nr = length(r)

    LMAX = 2 * amodes.lmax
    Wr_lm = calc_Wr_lm(win, LMAX, amodes.nside)
    @debug "Wr_lm" LMAX amodes.nside size(Wr_lm) Wr_lm[:,1]

    alm = Alm(LMAX, LMAX)
    LMLM = fill(0, LMAX+1, LMAX+1)
    for L=0:LMAX, M=0:L
        LMLM[L+1,M+1] = almIndex(alm, L, M) # + 1
    end
    @debug "LMLM" size(LMLM)

    check_nsamp(amodes, wmodes)

    gnl = amodes.basisfunctions
    @time for n′=1:amodes.nmax, n=1:amodes.nmax, l′=0:amodes.lmax_n[n′], l=0:amodes.lmax_n[n]
        ibase = getidx(amodes, n, l, 0)
        i′base = getidx(amodes, n′, l′, 0)
        ibase==1 && @show ibase,i′base, n,n′, l,l′, nlmsize

        gg1 = @. r^2 * gnl(n,l,r) * gnl(n′,l′,r)
        ## debug
        #gg1_quadgk,E = quadgk(r->r^2 * gnl(n,l,r) * gnl(n′,l′,r), amodes.rmin, amodes.rmax)
        #@debug "gg1" sum(gg1)*Δr gg1_quadgk,E
        gg1sum = sum(gg1)

        for m=0:l, m′=0:l′
            if neg_m
                m = -m
            end
            i = ibase + abs(m)
            i′ = i′base + m′
            M = m - m′

            #L1 = max(abs(l-l′),abs(M)):(l+l′)
            #w3j1 = wigner3j.(Float64, l, l′, L1, -m, m′)

            w3j_f = wigner3j_f(Float64, l, l′, -m, m′)
            w3j2 = w3j_f.symbols
            L2 = eachindex(w3j_f)

            #@show l,l′ m,m′ M L1 L2
            #@show w3j1 w3j2
            #@assert all(@. abs(w3j1 - w3j2) < eps(1.0))
            #@assert all(L1 .== L2)

            L = L2
            w3j = w3j2
            #@show w3j

            w3j000 = @. wigner3j000(L, l, l′)
            #@show w3j000
            #@show size(w3j) size(L) size(w3j000)
            #@show eachindex(w3j) eachindex(L) eachindex(w3j000)
            gaunt = @. √((2*L+1) * (2*l+1) * (2*l′+1) / (4*π)) * w3j000 * w3j
            #@debug "Gaunt" l,m l′,m′ L M gaunt[1] 1/√(4*π) √(4*π)*gaunt[1]

            w_ang = 0.0im
            for j=1:length(L)
                #@assert -L[j] <= M <= L[j]
                #LM = almIndex(Alm(LMAX, LMAX), L[j], abs(M)) # + 1
                LM = LMLM[L[j]+1,abs(M)+1]
                w_ang += gaunt[j] * (gg1' * Wr_lm[:,LM])
                #@debug "Wr_lm" L[j],M Wr_lm[1,LM] Wr_lm[1,LM]/√(4*π)
            end
            if M < 0
                w_ang = (-1)^M * conj(w_ang)
            end
            #@debug "wmix" i,i′ l,m l′,m′ w_ang gg1sum
            wmix[i,i′] = (-1)^m * w_ang * Δr
            #@show wmix[i,i′]
            @assert isfinite(wmix[i,i′])
        end
    end
    @assert all(isfinite.(wmix))
    return wmix
end

# This should be very performant
function win_lnn(win, wmodes::ConfigurationSpaceModes, cmodes::ClnnModes)
    lnnsize = getlnnsize(cmodes)
    Wlnn = fill(NaN, lnnsize)
    @show length(Wlnn), size(Wlnn)

    nr = size(win, 1)
    r, Δr = window_r(wmodes)

    # Note: the maximum ℓ we need is 0. However, healpy changes precision, and
    # for comparison we use the same lmax as elsewhere.
    Wr_00 = Array{Float64,1}(calc_Wr_lm(win, 2*cmodes.amodes.lmax, cmodes.amodes.nside)[:,1])
    @show size(Wr_00) typeof(Wr_00)
    @assert all(isfinite.(Wr_00))

    check_nsamp(cmodes.amodes, wmodes)

    gnl = cmodes.amodes.basisfunctions
    @time for i=1:lnnsize
        l, n, n′ = getlnn(cmodes, i)
        #@show i, l,n,n′, lnnsize

        gg = @. r^2 * gnl(n,l,r) * gnl(n′,l,r) * Wr_00

        Wlnn[i] = Δr * sum(gg) / √(4π)
        if !isfinite(Wlnn[i])
            @error "not finite" i l,n,n′ extrema(gg) Δr sum(gg) Wlnn[i]
            break
        end
    end
    @assert all(isfinite.(Wlnn))
    return Wlnn
end

function wigner3j000(l, l′, L)
    (abs(l-l′) <= L <= l+l′) || return 0.0
    J = l + l′ + L
    iseven(J) || return 0.0
    wig3j = (-1)^(J÷2) * exp(0.5*loggamma(1+J-2l) + 0.5*loggamma(1+J-2l′)
                             + 0.5*loggamma(1+J-2L) - 0.5*loggamma(1+J+1)
                             + loggamma(1+J/2)
                             - loggamma(1+J/2-l) - loggamma(1+J/2-l′)
                             - loggamma(1+J/2-L))
    return wig3j
end

function power_win_mix(wmix, wmix_negm, cmodes)
    amodes = cmodes.amodes
    nlmsize = getnlmsize(amodes)
    lnnsize = getlnnsize(cmodes)
    mmix = fill(NaN, lnnsize, lnnsize)
    for i′=1:lnnsize, i=1:lnnsize
        l, n, n′ = getlnn(cmodes, i)
        L, N, N′ = getlnn(cmodes, i′)
        mix = 0.0
        for m=0:l, M=0:L
            j = getidx(amodes, n, l, m)
            j′ = getidx(amodes, n′, l, m)
            J = getidx(amodes, N, L, M)
            J′ = getidx(amodes, N′, L, M)
            mix += real(wmix[j,J] * conj(wmix[j′,J′]))
            if m>0
                mix += real(conj(wmix_negm[j,J]) * wmix_negm[j′,J′])
            end
            if M>0
                mix += real(wmix_negm[j,J] * conj(wmix_negm[j′,J′]))
            end
            if m>0 && M>0
                mix += real(conj(wmix[j,J]) * wmix[j′,J′])
            end
        end
        mmix[i,i′] = mix / (2*l + 1)
    end
    return mmix
end

function calc_Wr_lm(win, LMAX, Wnside)
    nr = size(win,1)
    Wr_lm = fill(NaN*im, nr, getlmsize(LMAX))
    @time for i=1:nr
        W = udgrade(win[i,:], Wnside)
        Wr_lm[i,:] .= mymap2alm(W, lmax=LMAX)
    end
    return Wr_lm
end

# specialize
function calc_Wr_lm(win::SeparableArray, LMAX, Wnside)
    mask = udgrade(win.mask, Wnside)
    wlm = mymap2alm(mask, lmax=LMAX)
    return SeparableArray(win.phi, wlm, name1=:phi, name2=:wlm)
end

function optimize_Wr_lm_layout(Wr_lm)
    return Wr_lm  # Don't change the semantics! If speed is a problem, we will have to revisit this.
    #LMmax = size(Wr_lm,2)
    #w = Array{eltype(Wr_lm),1}[]
    #for lm=1:LMmax
    #    wr = Wr_lm[:,lm]
    #    push!(w, wr)
    #end
    #return w
end

optimize_Wr_lm_layout(Wr_lm::SeparableArray) = Wr_lm  # noop for SeparableArray

function cmix_kernel(gg1, gg2, wr)
    real(dot(gg1, wr) * conj(dot(gg2, wr)))
end

function test_cmix_kernel()
    gg1 = rand(100)
    gg2 = rand(100)
    wr = rand(Complex{Float64}, 100)
    l = 5
    L = 10
    LMAX = l + L
    alm = Alm(LMAX, LMAX)
    L1M1cache = [almIndex(alm, L, 0:L) for L=0:LMAX]
    Wnside = estimate_nside(LMAX)
    win = rand(100, nside2npix(Wnside))
    @show typeof(win)
    Wr_lm = optimize_Wr_lm_layout(calc_Wr_lm(win, LMAX, Wnside))
    @show typeof(Wr_lm)

    # compile
    cmix_kernel(gg1, gg2, wr)
    calc_cmix_ang(l, L, L1M1cache, gg1, gg2, Wr_lm)

    s = 0.0
    @time cmix_kernel(gg1, gg2, wr)
    @time cmix_kernel(gg1, gg2, Wr_lm[1])

    @show typeof(l) typeof(L) typeof(L1M1cache) typeof(gg1) typeof(gg2) typeof(Wr_lm)

    @time calc_cmix_ang(l, L, L1M1cache, gg1, gg2, Wr_lm)
    @time m_ang = calc_cmix_ang(l, L, L1M1cache, gg1, gg2, Wr_lm)
    @show typeof(m_ang)
end

function calc_cmix_ang(l, L, L1M1cache, gg1, gg2, Wr_lm)
    #@debug "calc_cmix_ang" l L size(gg1) size(gg2) size(Wr_lm) size(L1M1cache) size(Wr_lm[1])
    m_ang = 0.0
    for L1=abs(l-L):2:(l+L)
        L1M1 = L1M1cache[L1+1][1]
        s = cmix_kernel(gg1, gg2, Wr_lm[:,L1M1])
        for M1=1:L1
            #@debug "" L1 M1
            L1M1 = L1M1cache[L1+1][M1+1]
            s += 2 * cmix_kernel(gg1, gg2, Wr_lm[:,L1M1])
        end
        m_ang += s * wigner3j000(l, L, L1)^2
    end
    return m_ang
end

function calc_cmixii(i, i′, cmodes, r, Δr, gnlr, Wr_lm, L1M1cache,
                     div2Lp1, interchange_NN′)
    l, n, n′ = getlnn(cmodes, i)
    L, N, N′ = getlnn(cmodes, i′)
    if interchange_NN′
        tmp = N′
        N′ = N
        N = tmp
    end

    showthis = ((l==L==0 && n==N′==1 && n′==N==2) || (l==L==0 && n==N′==2 && n′==N==1))
    showthis && @show "huzzah",i,i′, n,n′, N,N′, l,L
    #@show i,i′, (l,n,n′), (L,N,N′)

    gg1 = @. r^2 * gnlr[:,n,l+1] * gnlr[:,N,L+1]
    gg2 = @. r^2 * gnlr[:,n′,l+1] * gnlr[:,N′,L+1]

    mix = calc_cmix_ang(l, L, L1M1cache, gg1, gg2, Wr_lm)

    mix = 1 / (4*π) * mix * Δr^2
    if !div2Lp1
        mix *= (2*L+1)
    end

    return mix
end

# specialize separable
function calc_cmixii(i, i′, cmodes, r, Δr, gnlgNLϕ, ang_mix::AbstractMatrix,
                     div2Lp1, interchange_NN′)
    l, n, n′ = getlnn(cmodes, i)
    L, N, N′ = getlnn(cmodes, i′)
    if interchange_NN′
        tmp = N′
        N′ = N
        N = tmp
    end
    #@show i,i′, n,n′, N,N′, l,l′

    #@show typeof(r) typeof(phi)
    gg1 = gnlgNLϕ[n,l+1,N,L+1]
    gg2 = gnlgNLϕ[n′,l+1,N′,L+1]

    m_ang = ang_mix[l+1,L+1]

    mix = m_ang * gg1 * gg2
    if !div2Lp1
        mix *= (2*L+1)
    end

    return mix
end

@doc raw"""
    power_win_mix(win, wmodes, cmodes; div2Lp1=false, interchange_NN′=false)
    power_win_mix(win, w̃mat, vmat, wmodes, bcmodes; div2Lp1=false, interchange_NN′=false)
    power_win_mix(wmix, wmix_negm, cmodes)

This function is used to calculate the coupling matrix $\mathcal{M}_{\ell nn'}^{LNN'}$
-- the first version without any binning, the second version also takes the
binning matrices `w̃mat` and `vmat` to calculate the coupling matrix of the
binned modes $\mathcal{N}_{LNN'}^{lnn'}$. These assume the symmetry between $N$
and $N'$.

The last version is probably not useful. It takes a fully calculated window
mixing matrix to calculate the coupling matrix brute-force.

If `div2Lp1=true` then the whole matrix is divided by $2L+1$.

If `interchange_NN′=true` then calculate the same, but with $N$ and $N'$
interchanged, which might be useful for the covariance matrix.

Either version of `power_win_mix()` will specialize to a separable window
function if `win` is a `SeparableArray`.
"""
function power_win_mix(win, wmodes::ConfigurationSpaceModes, cmodes::ClnnModes;
                       div2Lp1=false, interchange_NN′=false)
    amodes = cmodes.amodes
    lnnsize = getlnnsize(cmodes)
    #mix = fill(NaN, lnnsize, lnnsize)
    mix = SharedArray{Float64}(lnnsize, lnnsize)
    @show length(mix), size(mix)
    @show lnnsize^2, lnnsize

    r, Δr = window_r(wmodes)

    LMAX = 2 * amodes.lmax
    Wr_lm = optimize_Wr_lm_layout(calc_Wr_lm(win, LMAX, amodes.nside))

    # the gnl precomputation only saves about 10% time
    gnl = amodes.basisfunctions
    gnlr = fill(NaN, length(r), size(gnl.knl)...)
    @time for l=0:amodes.lmax, n=1:amodes.nmax_l[l+1]
        @. gnlr[:,n,l+1] = gnl(n,l,r)
    end
    check_nsamp(amodes, wmodes)

    alm = Alm(LMAX, LMAX)
    L1M1cache = [almIndex(alm, L, 0:L) for L=0:LMAX]

    ## crashes at end:
    #@time @threads for i′=1:lnnsize
    #    @show i′, lnnsize
    #    @time for i=1:lnnsize
    #        mix[i,i′] = calc_cmixii(i, i′, cmodes, r, Δr, gnlr, Wr_lm, L1M1cache)
    #        #@show i,i′, mix[i,i′]
    #    end
    #end

    ## too slow:
    #@time mix = pmap(idx -> calc_cmixii(idx[1], idx[2], cmodes, r, Δr, gnlr, Wr_lm, L1M1cache),
    #                 CartesianIndices((1:lnnsize, 1:lnnsize)))

    #@show "full"
    #@time @sync @distributed for i′=1:lnnsize
    #    @time for i=1:lnnsize
    #        mix[i,i′] = calc_cmixii(i, i′, cmodes, r, Δr, gnlr, Wr_lm, L1M1cache)
    #        #@show i,i′, mix[i,i′]
    #    end
    #end
    #mix1 = deepcopy(mix)

    #@show "symmetric"
    #@time @sync @distributed for i′=1:lnnsize
    #    L, = getlnn(cmodes, i′)
    #    @time for i=i′:lnnsize
    #        l, = getlnn(cmodes, i)
    #        mix[i,i′] = calc_cmixii(i, i′, cmodes, r, Δr, gnlr, Wr_lm, L1M1cache)
    #        mix[i′,i] = (2*l+1) / (2*L+1) * mix[i,i′]
    #        #@show i,i′, mix[i,i′]
    #    end
    #end
    #@assert mix == mix1

    #@show "symmetric pmap"
    @showprogress pmap(i′ -> begin
                   L, = getlnn(cmodes, i′)
                   #@show i′,L,lnnsize
                   for i=i′:lnnsize
                       l, = getlnn(cmodes, i)
                       mix[i,i′] = calc_cmixii(i, i′, cmodes, r, Δr, gnlr,
                                               Wr_lm, L1M1cache, div2Lp1,
                                               interchange_NN′)
                       mix[i′,i] = (2*l+1) / (2*L+1) * mix[i,i′]
                       #@show i,i′, mix[i,i′]
                   end
                   return i′  # return something that doesn't take much memory
               end,
               1:lnnsize)
    #@assert mix == mix1
    @assert all(isfinite.(mix))
    return mix
end

# specialize to Separable window
function power_win_mix(win::SeparableArray, wmodes::ConfigurationSpaceModes,
        cmodes::ClnnModes; div2Lp1=false, interchange_NN′=false)
    # Rather than writing a new specialized method, re-use what we already have:
    bcmodes = ClnnBinnedModes(I, I, cmodes)
    return power_win_mix(win, I, I, wmodes, bcmodes;
        div2Lp1=div2Lp1, interchange_NN′=interchange_NN′)
end

nzind(vec::AbstractVector) = 1:length(vec)
nzind(vec::SparseVector) = vec.nzind

Base.getindex(::UniformScaling{T}, ::Colon, m::Integer) where {T} = sparsevec([m], [T(1)])
Base.getindex(::UniformScaling{T}, m::Integer, ::Colon) where {T} = sparsevec([m], [T(1)])

# binned cmix
function _power_win_mix(w̃mat, vmat, r, Δr, gnlr, Wr_lm, L1M1cache, bcmodes;
                       div2Lp1=false, interchange_NN′=false)
    cmodes = bcmodes.cmodes
    lnnsize = getlnnsize(cmodes)
    LNNsize1 = (typeof(w̃mat) <: UniformScaling) ? lnnsize : size(w̃mat,1)
    LNNsize2 = (typeof(vmat) <: UniformScaling) ? lnnsize : size(vmat,2)

    # Use pmap() to allow distributed parallel computing:
    n_idxs = SeparableArray(LNNsize1:-1:1, ones(Int, LNNsize2))
    m_idxs = SeparableArray(ones(Int, LNNsize1), LNNsize2:-1:1)
    batchsize = (LNNsize1 * LNNsize2) ÷ (nworkers()^2) + 1
    @show LNNsize1, LNNsize2, batchsize
    mix = @showprogress pmap((n,m) -> begin
            w̃mat_n = w̃mat[n,:]
            vmat_m = vmat[:,m]
            w̃nzrange = nzind(w̃mat_n)
            vnzrange = nzind(vmat_m)
            c = 0.0
            for i in w̃nzrange, i′ in vnzrange
                v = vmat_m[i′]
                v==0 && continue
                w̃ = w̃mat_n[i]
                w̃==0 && continue
                c += w̃ * v * calc_cmixii(i, i′, cmodes, r, Δr, gnlr, Wr_lm,
                                         L1M1cache, div2Lp1, interchange_NN′)
            end
            return c
        end,
        n_idxs,
        m_idxs,
        batch_size=batchsize,
    )

    return collect(mix[end:-1:1,end:-1:1])
end

function calc_angular_mixing_matrix(lmax, wlm)
    Wℓ = alm2cl(wlm)
    ang_mix = fill(NaN, lmax+1, lmax+1)
    for L=0:lmax, l=0:lmax
        s = 0.0
        for L1=abs(L-l):2:(L+l)
            wig = wigner3j000(l, L, L1)
            s += wig^2 * (2*L1+1) * Wℓ[L1+1]
        end
        ang_mix[l+1,L+1] = 1 / (4π) * s  # 2*L+1 will be included later, if not symmetric
    end
    return ang_mix
end

check_nsamp_1gnl(amodes, wmodes::ConfigurationSpaceModes) = check_nsamp_1gnl(amodes, wmodes.nr)
function check_nsamp_1gnl(amodes, nr)
    num_imprecise = 0
    max_Nsamp = 0
    lmax = amodes.lmax
    nmax_l = amodes.nmax_l
    for L=0:lmax, N=1:nmax_l[L+1]
        Nsamp = 8 * N  # + L
        max_Nsamp = max(max_Nsamp, Nsamp)
        if Nsamp > nr
            num_imprecise += 1
        end
    end
    if num_imprecise > 0
        @warn "Radial integral over one gnl(r) unlikely to converge" num_imprecise max_Nsamp nr amodes.rmin amodes.rmax amodes.lmax amodes.nmax amodes.nmax_l
        #throw(ErrorException("Nsamp > nr"))
    end
end

check_nsamp(amodes, wmodes::ConfigurationSpaceModes) = check_nsamp(amodes, wmodes.nr)
function check_nsamp(amodes, nr)
    num_imprecise = 0
    max_Nsamp = 0
    lmax = amodes.lmax
    nmax_l = amodes.nmax_l
    for L=0:lmax, N=1:nmax_l[L+1]
        for l=0:lmax, n=1:nmax_l[l+1]
            Nsamp = 8 * (n + N)  # + l + L
            max_Nsamp = max(max_Nsamp, Nsamp)
            if Nsamp > nr
                num_imprecise += 1
            end
        end
    end
    if num_imprecise > 0
        @warn "Radial integrals unlikely to converge" num_imprecise max_Nsamp nr amodes.rmin amodes.rmax amodes.lmax amodes.nmax amodes.nmax_l
        #throw(ErrorException("Nsamp > nr"))
    end
end

# calculate radial mixers
function calc_radial_mixing(lmax, nmax_l, gnlr, phi, r, Δr)
    nmax = maximum(nmax_l)
    gnlgNLϕ = fill(NaN, nmax, lmax+1, nmax, lmax+1)
    ggϕint = fill(NaN, length(phi))
    for L=0:lmax, N=1:nmax_l[L+1]
        for l=0:lmax, n=1:nmax_l[l+1]
            !isnan(gnlgNLϕ[n,l+1,N,L+1]) && continue
            @. ggϕint = r^2 * gnlr[:,n,l+1] * gnlr[:,N,L+1] * phi
            gg = Δr * sum(ggϕint)
            gnlgNLϕ[n,l+1,N,L+1] = gg
            gnlgNLϕ[N,L+1,n,l+1] = gg
        end
    end
    return gnlgNLϕ
end

# specialized for separable window
function _power_win_mix(w̃mat, vmat, r, Δr, gnlr, Wr_lm::SeparableArray, L1M1cache, bcmodes;
                       div2Lp1=false, interchange_NN′=false)
    cmodes = bcmodes.cmodes
    lnnsize = getlnnsize(cmodes)
    LNNsize1 = (typeof(w̃mat) <: UniformScaling) ? lnnsize : size(w̃mat,1)
    LNNsize2 = (typeof(vmat) <: UniformScaling) ? lnnsize : size(vmat,2)
    mix = fill(NaN, LNNsize1, LNNsize2)
    @show length(mix), size(mix)

    println("Calculate angular and radial mixing:")
    check_nsamp(cmodes.amodes, length(gnlr[:,1,1]))
    lmax = bcmodes.cmodes.amodes.lmax
    nmax_l = bcmodes.cmodes.amodes.nmax_l
    @time ang_mix = calc_angular_mixing_matrix(lmax, Wr_lm.wlm)
    @time gnlgNLϕ = calc_radial_mixing(lmax, nmax_l, gnlr, Wr_lm.phi, r, Δr)

    println("Calculate binned mixing matrix:")
    #@time @sync @distributed for m=1:LNNsize
    @time for m=1:LNNsize2
        #@show m, LNNsize
        vmat_m = vmat[:,m]
        vnzrange = nzind(vmat_m)
        #@show typeof(vmat_m) typeof(vnzrange) size(vmat) vnzrange length(vnzrange) vmat_m[vnzrange]
        for n=1:LNNsize1
            #@show m,n,LNNsize
            w̃mat_n = w̃mat[n,:]
            w̃nzrange = nzind(w̃mat_n)
            c = 0.0
            for i in w̃nzrange, i′ in vnzrange
                v = vmat_m[i′]
                v==0 && continue
                w̃ = w̃mat_n[i]
                w̃==0 && continue
                c += w̃ * v * calc_cmixii(i, i′, cmodes, r, Δr, gnlgNLϕ,
                                         ang_mix, div2Lp1, interchange_NN′)
            end
            mix[n,m] = c
        end
    end
    return mix
end

# calculate binned power spectrum mode-coupling matrix
function power_win_mix(win, w̃mat, vmat, wmodes::ConfigurationSpaceModes, bcmodes::ClnnBinnedModes; div2Lp1=false, interchange_NN′=false)
    cmodes = bcmodes.cmodes
    amodes = cmodes.amodes
    lnnsize = getlnnsize(cmodes)
    LNNsize = getlnnsize(bcmodes)
    @show LNNsize^2, LNNsize, lnnsize

    r, Δr = window_r(wmodes)

    println("Calculate Wr_lm:")
    LMAX = 2 * amodes.lmax
    @time Wr_lm = optimize_Wr_lm_layout(calc_Wr_lm(win, LMAX, amodes.nside))

    println("Calculate gnlr:")
    gnl = amodes.basisfunctions
    gnlr = fill(NaN, length(r), size(gnl.knl)...)
    @time for l=0:amodes.lmax, n=1:amodes.nmax_l[l+1]
        @. gnlr[:,n,l+1] = gnl(n,l,r)
    end
    check_nsamp(amodes, wmodes)

    alm = Alm(LMAX, LMAX)
    L1M1cache = [almIndex(alm, L, 0:L) for L=0:LMAX]

    mix = _power_win_mix(w̃mat, vmat, r, Δr, gnlr, Wr_lm, L1M1cache, bcmodes; div2Lp1=div2Lp1, interchange_NN′=interchange_NN′)

    @assert all(isfinite.(mix))
    return mix
end

end

# vim: set sw=4 et sts=4 :


text:
using Pkg; Pkg.activate(@__DIR__)
using MuJoCo
mj_activate("/home/taylor/.mujoco/bin/mjkey.txt") # set location to MuJoCo key path

using LyceumMuJoCo, LyceumMuJoCoViz
using FiniteDiff

using IterativeLQR
using LinearAlgebra
using Random

# ## load MuJoCo model
path = joinpath(@__DIR__, "../../../env/cartpole/deps/cartpole.xml")

include("mujoco_model.jl")
cartpole = MuJoCoModel(path)
sim = LyceumMuJoCo.MJSim(cartpole.m, cartpole.d)

@show cartpole.m.body_mass
@show cartpole.m.body_inertia[:, :]

LyceumMuJoCoViz.visualize(sim) #, trajectories=[states]) # ctrl + LEFT (access trajectory mode)

# ## horizon
T = 26

# ## cartpole
nx = cartpole.nx
nu = cartpole.nu

# ## model
dyn = Dynamics(
    (y, x, u, w) -> f!(y, cartpole, x, u),
    (dx, x, u, w) -> fx!(dx, cartpole, x, u),
    (du, x, u, w) -> fu!(du, cartpole, x, u),
    nx, nx, nu)

model = [dyn for t = 1:T-1]

# ## initial conditions
x1 = [0.0; π; 0.0; 0.0]
xM = zeros(cartpole.nx)
xT = zeros(cartpole.nx)
x_ref = zeros(cartpole.nx)

# # rollout
ū = [t < 5 ? 1.0 * rand(nu) : (t < 10 ? -1.0 * rand(nu) : zeros(nu)) for t = 1:T-1]
w = [zeros(0) for t = 1:T-1]
x̄ = rollout(model, x1, ū)

# ## objective

function objt(x, u, w)
	J = 0.0
	J += transpose(x - x_ref) * Diagonal(1.0e-1 * ones(nx)) * (x - x_ref)
	J += transpose(u) * Diagonal(1.0e-3 * ones(nu)) * u
	return J
end

function objT(x, u, w)
	J = 0.0
	J += transpose(x - x_ref) * Diagonal(1.0e-1 * ones(nx)) * (x - x_ref)
	return J
end

ct = IterativeLQR.Cost(objt, nx, nu)
cT = IterativeLQR.Cost(objT, nx, 0)
obj = [[ct for t = 1:T-1]..., cT]

# ## constraints
goal(x, u, w) = x - xT

cont = Constraint()
conT = Constraint(goal, nx, 0)
cons = [[cont for t = 1:T-1]..., conT]

# ## problem
prob = IterativeLQR.solver(model, obj, cons,
    opts=Options(
        verbose=true,
        max_al_iter=10))
initialize_controls!(prob, ū)
initialize_states!(prob, x̄)

# ## solve
@time solve!(prob)

# ## solution
x_sol, u_sol = get_trajectory(prob)
@show IterativeLQR.eval_obj(prob.m_data.obj.costs, prob.m_data.x, prob.m_data.u, prob.m_data.w)
@show prob.s_data.iter[1]
@show norm(goal(prob.m_data.x[T], zeros(0), zeros(0)), Inf)

# ##
# using Plots
# plot(hcat(x_sol...)')
# plot(hcat(u_sol..., u_sol[end])', linetype=:steppost)

# ## visualize
x_vis = [[x_sol[1] for t = 1:10]..., x_sol..., [x_sol[end] for t = 1:10]...]
states = Array(undef, statespace(sim), length(x_vis))
for t = 1:length(x_vis)
    sim.d.qpos .= x_vis[t][1:2]
    sim.d.qvel .= x_vis[t][3:4]
    sim.d.ctrl .= [0.0]
    # LyceumMuJoCo.step!(sim)
    states[:, t] .= getstate(sim)
end
visualize(sim, trajectories=[states])


text:
# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.10.3
#   kernelspec:
#     display_name: Julia 1.6.4
#     language: julia
#     name: julia-1.6
# ---

# %%
using Distributions
using Plots
using Optim

traceon = Optim.Options(
    store_trace = true,
    extended_trace = true
)

# model
Distributions.TDist(μ, ρ, ν) = LocationScale(μ, ρ, TDist(ν))

# test data
X = [-0.01, 0.01, 1.0]

# Warning: Very Slow!
lower = [-0.1, -15.0, -2.5]
upper = [ 0.4,  10.0,  8.0]
@show lower upper
r = @time optimize(x -> -loglikelihood(TDist(x[1], 10^x[2], 10^x[3]), X), lower, upper, [0.33, 0.0, 0.27], Fminbox(LBFGS()), traceon)

# %%
function plot_trace(r)
    x = r.initial_x
    c = hcat(x, (t.metadata["x"] for t in r.trace)...)
    l = -minimum(r)
    m = r.minimizer
    t = "initial = $x\nfinal = $(round.(m; digits=3))\nfinal loglikelihood = $(round(l; digits=3))"

    plot(; legend=:topleft)
    scatter3d!([x[1]], [x[2]], [x[3]]; label="initial", c=:blue, ms=3)
    plot!(c[1,:], c[2,:], c[3,:]; label="path", c=:darkcyan)
    scatter3d!([m[1]], [m[2]], [m[3]]; label="final", c=:yellow, ms=3)
    plot!(; xlabel="μ", ylabel="log₁₀ρ", zlabel="log₁₀ν")
    title!(t; titlefontsize=10)
end

plot_trace(r)
plot!(; fmt=:png)

# %%
function f(μ, log10ν, X)
    log10rhos = range(-15, 10; length=200)
    -maximum(loglikelihood(TDist(μ, 10^log10ρ, 10^log10ν), X) for log10ρ in log10rhos)
end

function plot_trace2d(r)
    x = r.initial_x
    c = hcat(x, (t.metadata["x"] for t in r.trace)...)
    l = -minimum(r)
    m = r.minimizer
    t = "initial = $x\nfinal = $(round.(m; digits=3))\nfinal loglikelihood = $(round(l; digits=3))"

    mus = range(-0.1, 0.4; length=200)
    log10nus = range(-2.5, 8.0; length=200)
    z = f.(mus', log10nus, Ref(X))

    plot(; legend=:topleft, colorbar=false)
    plot!(; xlim=extrema(mus), ylim=extrema(log10nus))
    heatmap!(mus, log10nus, z; clim=(-1, 5), c=reverse(cgrad(:CMRmap)))
    scatter!([x[1]], [x[3]]; label="initial", ms=4, c=:blue)
    plot!(c[1,:], c[3,:]; label="path", c=:cyan, lw=1.5)
    scatter!([m[1]], [m[3]]; label="final", ms=4, c=:yellow)
    plot!(; xlabel="μ", ylabel="log₁₀ν")
    title!(t; titlefontsize=10)
end

plot_trace2d(r)
plot!(; fmt=:png)

# %%

max_stars_repo_path: 0024/TDist model is not log-convex - Fminbox with LBFGS.jl

max_stars_repo_name: genkuroki/public

max_issues_repo_name: genkuroki/public

max_forks_repo_name: genkuroki/public


text:
"""
    binary()
'labels' contains labels of boundaries that are not generated exclusively by contained vertices.
For each of these label, at least one of the generating vertices are not contained in the other simplex.
This function translates each label into a binary array. Each number is transformed
into its binar yexpression as if you had N+1 bits.

The array of 0's and 1's

Input arguments
---------------
m::                  Column array of numbers whose binary expressions we want to extract
n::                 Maximum power in the binary decomposition
labels::Vec{Int}     Contains a label ordering for all possible boundaries. Labels take values in the range 1:2^(dim+1)-2 (but not all of them), except for the first  2^Nin - 1, where Nin is the minimum number of contained vertices. Exclude the maximum because it corresponds to the whole simplex.
dimension::Int      Previuously 'n' The dimension of the space.
Returns
-------
boundary_binary_labels::Array{Int, 2} Matrix of 0's and 1's. Each row represents a number contained in 'labels'. The columns for this row is the binary expansion of this number in D+1 bits. Each column of this matrix contains a vertex generating a generic simplex. If there is a zero in a particular row, it means that the corresponding boundary is not generated by that vertex.
"""
function Binary(m, n)
    # m: column array of numbers whose binary expressions we want to extract
    # n: maximum power in the binary decomposition

    # OUTCOME:
    # array of dimension size(m,2) x (n+1)
    # each row contains the binary decomposition of the corresponding entry
    # in m, that is
    # (thea(1),...,theta(n+1)) \in {0,1}^{(n+1)}
    # such that m = \sum_{k=0}^{n} theta(k+1)*2^k

    exponent = repeat(transpose(float(2) .^ (-(collect(0:n)))), size(m, 1), 1)
    m = repeat(m, 1, n + 1)

    binary = m .* exponent
    binary = floor.(Int64, binary)
    binary[:, 1:n] = binary[:, 1:n] - 2 * binary[:, 2:(n + 1)]

    return binary
end


text:
@testset "Initial Convergence Handling" begin
    f(x) = x[1]^2
    function g!(out, x)
        out[1] = 2.0 * x[1]
    end
    function h!(out, x)
        out[1,1] = 2.0
    end
    for Optimizer in (AcceleratedGradientDescent,
                      GradientDescent, ConjugateGradient, LBFGS, BFGS,
                      MomentumGradientDescent)

        res = optimize(f, g!, [0.], Optimizer())
        @test Optim.minimizer(res)[1] ≈ 0.
    end

    for Optimizer in (Newton, NewtonTrustRegion)
        res = optimize(f, g!, h!, [0.], Optimizer())
        @test Optim.minimizer(res)[1] ≈ 0.
    end
end


text:
using Pkg
using LinearAlgebra
using DataFrames
import Statistics.mean
import Statistics.std
using Printf
using CSV

println("start run script for cpmmh")

# load functions
include(pwd()*"/src/SDEMEM OU neuron data/ou_sdemem.jl")
include(pwd()*"/src/SDEMEM OU neuron data/mcmc.jl")

seed = 100

M_subjects = 100

y, prior_parameters_η, prior_parameters_σ_ϵ = set_up(M_subjects, seed)

startval_ϕ = ones(M_subjects,3)

for j = 1:3

    μ_0_j, M_0_j, α_j, β_j = prior_parameters_η[j,:]

    startval_ϕ[:,j,1] = μ_0_j*ones(M_subjects)

end

startval_log_σ_ϵ = log(0.5)

# set time step
dt = diff(y[1].mSec)[1]

posterior_mean_kalman = CSV.read("data/SDEMEM OU neuron data/posterior_mean_sigma_epsilon_and_phi.csv")

log_σ_ϵ  = posterior_mean_kalman[1,1]
ϕ  = Matrix(posterior_mean_kalman[2:end,:])

function run_cpf(particels)

    nbr_particles = particels*ones(Int64, M_subjects)
    ρ = 0.0

    # set number of obs for each data set
    T_vec = zeros(Int64, M_subjects)
    for i = 1:M_subjects; T_vec[i] = length(y[i].mV); end

    u_prop_old = [randn(nbr_particles[1], T_vec[1]+1)]
    u_resample_old = [randn(T_vec[1], 2)]

    for i in 2:M_subjects;
        append!(u_prop_old, [randn(nbr_particles[i],T_vec[i]+1)])
        append!(u_resample_old, [randn(T_vec[i],2)])
    end

    run_time = @elapsed ll1 =  cpf(y, exp(log_σ_ϵ), ϕ, dt, u_prop_old, u_resample_old, nbr_particles, false)

    return ll1, run_time

end

function loglik_est(nbr_particles_pf, nbr_pf_eval)

    ll = zeros(nbr_pf_eval)

    for i in 1:nbr_pf_eval
        ll_part, ~ = run_cpf(nbr_particles_pf)
        ll[i] = sum(ll_part)
    end

    return ll

end

nbr_pf_eval = 50
nbr_particles_pf = [20, 50, 100, 200, 250, 500, 1000, 2000, 3000, 4000]

loglik_matrix = zeros(nbr_pf_eval, length(nbr_particles_pf)+1)
run_times = zeros(length(nbr_particles_pf))

#i = 1
#run_times[i]  = @elapsed loglik_matrix[:,i] = loglik_est(nbr_particles_pf[i], nbr_pf_eval)

Random.seed!(1)

for i in 1:length(nbr_particles_pf)
    @printf "Current nbr of particels: %.2f\n" nbr_particles_pf[i]
    run_times[i]  = @elapsed loglik_matrix[:,i+1]  = loglik_est(nbr_particles_pf[i], nbr_pf_eval)
end

loglik_matrix[1:length(nbr_particles_pf),1] = run_times

CSV.write("data/SDEMEM OU neuron data/tune_nbr_particles_bootstrap_pf.csv", DataFrame(loglik_matrix))

max_stars_repo_path: src/SDEMEM OU neuron data/tune_nbr_particles_bootstrap_pf.jl