text:
module CellularBase
import Base: getindex, ndims, size, CartesianIndices, show, eltype
import Markdown:parse

export BoundaryCondition, Periodic, FixedMax, FixedMin, Clamp, boundary_conditions,
		AbstractNeighborhood, VonNeumannNeighborhood, MooreNeighborhood,
		radius, cartesians, neighborhood,
		AbstractGrid, state, state!, possible_states, newstate,
		evolve!, tabular_evolution!, simulate!

abstract type BoundaryCondition end
struct Periodic <: BoundaryCondition end
struct FixedMin <: BoundaryCondition end
struct FixedMax <: BoundaryCondition end
struct Clamp <: BoundaryCondition end

# Neighborhoods
abstract type AbstractNeighborhood{N} end
radius(neighborhood::AbstractNeighborhood) = neighborhood.radius
function cartesians(neighborhood::AbstractNeighborhood{N})::Array{CartesianIndex{N},1} where N
	neighborhood.cartesians
end
(dimensions(::AbstractNeighborhood{N}) where N) = N

struct VonNeumannNeighborhood{N} <: AbstractNeighborhood{N}
	cartesians::Array{CartesianIndex{N},1}
	radius::Int
	function VonNeumannNeighborhood{T}(radius::Int) where T
		arr = CartesianIndex[]
		for i in 1:radius
			for j in 1:T
				zs = zeros(Int, T)
				zs[j] = i
				push!(arr, CartesianIndex(Tuple(zs)))
				zs[j] = -i
				push!(arr, CartesianIndex(Tuple(zs)))
			end
		end
		new(arr, radius)
	end
end

struct MooreNeighborhood{N} <: AbstractNeighborhood{N}
	cartesians::Array{CartesianIndex{N},1}
	radius::Int
	function MooreNeighborhood{T}(radius::Int) where T
		new(
			CartesianIndices(Tuple(-radius:radius for _ in 1:dimensions)) |>
			x -> reshape(x, :),
			radius
		)
	end
end

function Base.show(io::IO, ::MIME"text/html", neighborhood::AbstractNeighborhood)
	show(
		io,
		MIME("text/html") ,
		parse("**Type:** " * repr(typeof(neighborhood), context=io) *
		"\n\n **Radius:** " * repr(radius(neighborhood), context=io) *
		"\n\n **List of Neighbors:** " * repr(Tuple.(cartesians(neighborhood)), context=io))
	)
end

# Grids
abstract type AbstractGrid{T,N} <: AbstractArray{T,N} end

function Base.eltype(::Type{AbstractGrid{T,N}}) where {T,N}
	return T
end
(Base.ndims(grid::AbstractGrid{T,N}) where {T,N}) = N
Base.size(grid::AbstractGrid) = size(state(grid))
Base.IndexStyle(::Type{AbstractGrid}) = Base.IndexCartesian

function wrap(::Periodic, grid::AbstractGrid{T,N}, index::CartesianIndex{N})::T where {T,N}
	@inbounds grid[CartesianIndex(mod1.(Tuple(index), size(grid)))]
end
function wrap(::FixedMax, grid::AbstractGrid{T,N}, index::CartesianIndex{N})::T where {T,N}
	maximum(possible_states(grid))
end
function wrap(::FixedMin, grid::AbstractGrid{T,N}, index::CartesianIndex{N})::T where {T,N}
	minimum(possible_states(grid))
end
function wrap(::Clamp, grid::AbstractGrid{T,N}, index::CartesianIndex{N})::T where {T,N}
	@inbounds grid[
		CartesianIndex(
			Tuple(
				clamp(Tuple(index)[i], 1, size(grid)[i])
				for i in 1:ndims(grid)
			)
		)
	]
end
@inline function Base.getindex(grid::AbstractGrid{T,N}, index::CartesianIndex{N})::T where {T,N}
    if checkbounds(Bool, grid, index)
       	@inbounds state(grid)[index]
	else
        wrap(boundary_conditions(grid), grid, index)
    end
end

function Base.show(io::IO, ::MIME"text/html", grid::AbstractGrid)
	println(
		io,
		"Type: ", typeof(grid),
		"Possible States: ", possible_states(grid),
		"
Boundary Condition: ", boundary_conditions(grid),
		"
Neighborhood: ", neighborhood(grid),
		"
Size: ", size(grid)
	)
end

function state(grid::AbstractGrid{T,N})::Array{T,N} where {T,N}
	grid.state
end

function state!(grid::AbstractGrid, newstate)::Nothing
	grid.state = newstate
	return nothing
end
boundary_conditions(grid::AbstractGrid)::BoundaryCondition = grid.bc
(neighborhood(grid::AbstractGrid{T,N})::Array{CartesianIndex{N},1}) where {T,N} = cartesians(grid.neighborhood)
function possible_states(grid::AbstractGrid{T,N})::Array{T,1} where {T,N}
	grid.possible_states
end
function newstate(grid::AbstractGrid{T,N})::Array{T,N} where {T,N}
	similar(state(grid))
end
function neighbors(grid::AbstractGrid{T,N}, i::CartesianIndex{N})::Array{T,1} where {T,N}
    map(neighborhood(grid)) do ind
		grid[ind + i]
	end
end

function evolve!(grid::AbstractGrid{T,N}; kwargs...)::Nothing where {T,N}
    _newstate = newstate(grid)
    for i in eachindex(grid)
        _newstate[i] = grid(grid[i], neighbors(grid, i); kwargs...)
	end
	state!(grid, _newstate)
	return nothing
end

function tabular_evolution!(grid::AbstractGrid, table::Dict)
    newstate = newstate(grid)
    for i in CartesianIndices(grid)
        newstate[i] = table[grid(neighbors(grid, i))]
	end
	state!(grid, newstate)
	return nothing
end

function simulate!(grid::AbstractGrid{T,N}, max_steps::Int; store_results=true, postrunhook=nothing, kwargs...) where {T,N}
	if store_results
		results = Array{T,N}[]
		push!(results, copy(state(grid)))
	end
	if isnothing(postrunhook)
		postrunhook = (grid, step; kwargs...) -> true
	end
	postrunhook(grid, 0; kwargs...)
	ss = false
	step = 0
	while step < max_steps
		if (!ss) evolve!(grid; kwargs...) end
		if store_results push!(results, copy(state(grid))) end
		message = postrunhook(grid, step; kwargs...)
		ss = message == :shortcurcuit
		if message == :interrupt
			break
		end
		step += 1
	end
	if store_results results else Array{T,N}[] end
end

end


text:
"""
    DiagnosticsMachine

This module provides the infrastructure to extract diagnostics from a
ClimateMachine simulation. Two key abstractions are defined: diagnostic
variables and diagnostic groups. The `StdDiagnostics` module makes use of
these to define many standard variables and groups which may be used
directly by experiments. `DiagnosticsMachine` may be used by experiments
to define specialized variables and groups.
"""
module DiagnosticsMachine

export DiagnosticVar,
    PointwiseDiagnostic,
    @pointwise_diagnostic,
    dv_PointwiseDiagnostic,
    HorizontalAverage,
    @horizontal_average,
    dv_HorizontalAverage,
    States,
    #DiagnosticsGroup,
    @diagnostics_group
#DiagnosticsGroupParams

using CUDA
using Dates
using InteractiveUtils
using KernelAbstractions
using MacroTools
using MacroTools: prewalk
using MPI
using OrderedCollections
using Printf
using StaticArrays

using ..Diagnostics # until old diagnostics groups are removed
using ..Atmos
using ..BalanceLaws
using ..ConfigTypes
using ..DGMethods
using ..GenericCallbacks
using ..Mesh.Grids
using ..Mesh.Interpolation
using ..Mesh.Topologies
using ..MPIStateArrays
using ..Spectra
using ..TicToc
using ..VariableTemplates
using ..Writers

using CLIMAParameters
using CLIMAParameters.Planet: planet_radius

# Container to store simulation information necessary for all
# diagnostics groups.
Base.@kwdef mutable struct DiagnosticSettings
    mpicomm::MPI.Comm = MPI.COMM_WORLD
    param_set::Union{Nothing, AbstractParameterSet} = nothing
    dg::Union{Nothing, SpaceDiscretization} = nothing
    Q::Union{Nothing, MPIStateArray} = nothing
    starttime::Union{Nothing, String} = nothing
    output_dir::Union{Nothing, String} = nothing
    no_overwrite::Bool = false
end
const Settings = DiagnosticSettings()

include("onetime.jl")
include("variables.jl")
include("groups.jl")

const AllDiagnosticVars = OrderedDict{
    Type{<:ClimateMachineConfigType},
    OrderedDict{String, DiagnosticVar},
}()
function add_all_dvar_dicts(T::DataType)
    AllDiagnosticVars[T] = OrderedDict{String, DiagnosticVar}()
    for t in subtypes(T)
        add_all_dvar_dicts(t)
    end
end
add_all_dvar_dicts(ClimateMachineConfigType)

"""
    init(mpicomm, param_set, dg, Q, starttime, output_dir)

Save the parameters into `Settings`, a container for simulation
information necessary for all diagnostics groups.
"""
function init(
    mpicomm::MPI.Comm,
    param_set::AbstractParameterSet,
    dg::SpaceDiscretization,
    Q::MPIStateArray,
    starttime::String,
    output_dir::String,
    no_overwrite::Bool,
)
    Settings.mpicomm = mpicomm
    Settings.param_set = param_set
    Settings.dg = dg
    Settings.Q = Q
    Settings.starttime = starttime
    Settings.output_dir = output_dir
    Settings.no_overwrite = no_overwrite
end

include("atmos_diagnostic_funs.jl")

end # module DiagnosticsMachine

ext: jl


text:
# レイヤ定義ファイルの読み込み
include("../common/layers.jl")

using Random

# TwoLayerNetの定義
mutable struct TwoLayerNet
    layers
    loss_layer
    params
    grads

    function TwoLayerNet(;input_size, hidden_size, output_size, weight_init=0.01) # Keyword Argumentのみの場合はセミコロンが必要 cf. https://docs.julialang.org/en/v1/manual/functions/#Keyword-Arguments
        self = new()
        I_s, H_s, O_s = input_size, hidden_size, output_size
        # 重みとバイアスの初期化
        W_1 = weight_init .* Random.randn(I_s, H_s)
        b_1 = zeros(H_s)' # JuliaのブロードキャストはNumPyと挙動が異なり，次元を追加しないため，予め転置して次元数を揃える
        W_2 = weight_init .* Random.randn(H_s, O_s)
        b_2 = zeros(O_s)'

        # レイヤの生成
        self.layers = [
            Affine(W_1, b_1),
            Sigmoid(),
            Affine(W_2, b_2)
        ]

        # すべての重みと勾配をまとめる
        self.params, self.grads = [], []
        for layer in self.layers
            push!(self.params, layer.params)
            push!(self.grads, layer.grads)
        end
        self.loss_layer = SoftmaxWithLoss()
        return self
    end
end

# Python版のTwoLayerNetのpredict()
function predict!(model::TwoLayerNet, x)
    for layer in model.layers
        x = forward!(layer, x)
    end
    return x
end

# Python版のTwoLayersNetのforward()
function forward!(model::TwoLayerNet, x, t)
    score = predict!(model, x)
    loss = forward!(model.loss_layer, score, t) # SoftmaxWithLoss()のforward()
    return loss
end

# Python版のTwoLayersNetのbackward()
function backward!(model::TwoLayerNet; dout=1)
    dout = backward!(model.loss_layer, dout=dout) # SoftmaxWithLoss()のbackward()
    for layer in reverse(model.layers)
        dout = backward!(layer, dout=dout)
    end
    return dout
end


text:
# we are making a Python model if this gets called, so lets load the Python strategy -
include("./strategy/PythonStrategy.jl")

function make_python_model(path_to_model_file::String, path_to_output_dir::String; model_type::Symbol = :static)
    error("Ooops! Python model generation has not yet been implemented.")
end

max_stars_repo_name: varnerlab/CFMG

max_issues_repo_name: varnerlab/CFMG

max_forks_repo_name: varnerlab/CFMG


text:
function tsvd(A::NamedArray, k::Integer)
    m, n = size(A.array)
    F = svds(A.array, nsv=k)[1]
    U::Array{Float32, 2}, Z::Array{Float32, 1}, V::Array{Float32, 2} = F.U, F.S, F.V
    @assert size(U, 2) == length(Z) == size(V, 2) == k
    return U, Z, V
end

scale(A) = mapslices(normalize!, A .- mean(A,dims=1), dims=1)
X= flunch.assays["RNA"].norm_count

function RunPCA(X::NamedArray)
    mean(X.array, dims =)
    Y = broadcast(-, X.array, mean(X.array,1))
    tsvd
end

function RunPCA(X::scJuliaObject)
    Y = broadcast(-, X, mean(X,1))
    tsvd
end

# H. Li, G. C. Linderman, et al. "Algorithm 971: An Implementation of a
# Randomized Algorithm for Principal Component Analysis", ACM Transactions on
# Mathematical Software (TOMS), 2017
# DOI: https://doi.org/10.1145/3004053
function rsvd(A::SparseArrays.SparseMatrixCSC{Int32, Int32}, k::Integer; its=3, l=k+5)
    m, n = size(A)
    l = min(l, m, n)
    @assert 0 < k ≤ l ≤ min(m, n)
    Q::SparseArrays.SparseMatrixCSC{Int32, Int32} = rand(eltype(Int32), n, l) .- Int32(0.5)
    Y = A*Q
    F = lu!(Y)
    for i in 1:its
        Y = A'F.L
        F = lu!(Y)
        Y = A*F.L
        if i < its
            F = lu!(Y)
        else
            F = qr!(Y)
        end
    end
    Q = Matrix(F.Q)
    B = Q'A
    W, Σ, V = svd(B)
    U = Q*W
    return U[:,1:k], Σ[1:k], V[:,1:k]
end

Z = count(flunch)
Dispersion = var(Z.array,dims=2) ./ mean(Z.array,dims=2)

feature_data(flunch)

df = DataFrame(
  colour = ["green","blue"],
  shape = ["circle", "triangle"],
  border = ["dotted", "line"]
)

X = flunch
flunch.assays["RNA"].feature_data
insertcols!(feature_data(flunch), 2, :dispersion => Dispersion[:,1], makeunique=false)
feature_data(flunch)
FindVariableFeatures(flunch)
feature_data(flunch)[:dispersion] = Dispersion[:,1]

Dispersion = var(norm_count(flunch).array, dims = 2) ./ mean(norm_count(flunch).array,dims = 2)
GetFeatureData(flunch)[!, :dispersion] = Dispersion[:,1]
varfeatures = @pipe (flunch |> feature_data |> sort(_, [:dispersion]) |>  last(_, nfeatures).features |> reverse)
GetFeatureData(flunch)[!, :varfeatures] = in.(GetFeatureData(flunch)[!, :features], Ref(varfeatures))
flunch
nfeatures = 2000

Dispersion[:,1]

max_stars_repo_name: Cortalak/Pissaro

max_issues_repo_name: Cortalak/Pissaro

max_forks_repo_name: Cortalak/Pissaro


text:
module TestShare

using Test

module PkgA
export @f
using ContextualMacros
ContextualMacros.@def @f
macro context(ex)
    ContextualMacros.expandwith(__module__, ex; f = ctx -> Expr(:call, :+, ctx.args...))
end
end

module PkgB
export @f
using ContextualMacros
ContextualMacros.@def @f
macro context(ex)
    ContextualMacros.expandwith(__module__, ex; f = ctx -> Expr(:call, :*, ctx.args...))
end
end

using .PkgA
using .PkgB

name = Symbol("@f")
@test (@eval $name) === (@eval PkgA $name) === (@eval PkgB $name)
@test (PkgA.@context @f 1 2 3 4) == 10
@test (PkgB.@context @f 1 2 3 4) == 24

end  # module


text:
arr = map(PCM16Sample, rand(100, 2) .- 0.5)
testbuf = SampleBuf(arr, srate)
for ext in extensions
  fname = string(tempname(), ext)
  FileIO.save(fname, testbuf)
  buf = FileIO.load(fname)
  @test buf isa SampleBuf
end

if !Sys.iswindows()
  # testing with unicode
  file = joinpath(tempdir(),"α.flac")
  FileIO.save(file, testbuf)
  FileIO.load(file)
end


text:
function row_echelon_form!(A::AbstractMatrix{T}) where T <: Number
    c, d = size(A)
    l = Int[]
    pos = 0
    i = 0
    for i = 1:d
        j = findfirst(x -> !iszero(x), @view A[pos+1:end, i])
        j === nothing && continue
        j += pos
        pos += 1
        push!(l, i)

        if (pos != j)
            A[[pos, j], :] = A[[j, pos], :]
        end
        A[pos, :] ./= A[pos, i]

        for j = 1:c
            if j!=pos
                A[j, :] -= A[j, i].*A[pos, :]
            end
        end
    end
    return A, l
end

function row_echelon_form(A::AbstractMatrix)
    return row_echelon_form!(deepcopy(A))
end

function right_nullspace(M::AbstractMatrix{T}) where T
    A, l = row_echelon_form(M)
    c, d = size(A)
    (length(l) == d) && return zeros(T, d)
    W = zeros(T, d, d-length(l))
    i = 0
    for el in setdiff(1:d, l)
        i+= 1
        W[el, i] += 1
        for (j, k) in enumerate(l)
            if j < el
                W[k, i] -= A[j, el]
            end
        end
    end
    return W
end

function left_nullspace(M::AbstractMatrix)
    return Matrix(transpose(right_nullspace(Matrix(transpose(M)))))
end

function left_eigen(M::AbstractMatrix{T}) where T <: FiniteFields.GF
    @assert ==(size(M)...)
    Id = Matrix{T}(I, size(M)...)
    eigen = Dict{T, typeof(M)}()
    cumdim = 0
    for i in T # brute force; TODO: use factorization of characteristic/minimal polynomials
        cumdim >= size(M, 1) && break
        #do left eigenspaces!
        basis = first(row_echelon_form!(left_nullspace(M - i*Id)))
        nullity = size(basis, 1)
        if (nullity == 1) && all(iszero, basis)
            nullity -= 1 # left_nullspace returns trivial kernel if kernel is empty
        end
        if nullity > 0
            cumdim += nullity
            eigen[i] = basis
        end
    end
    return eigen
end

function _find_l(M::AbstractMatrix)
    # this function should be redundant when defining a better structure for echelonized subspaces
    l = Int[]
    for i = 1:size(M, 2)
        j = findfirst(isone, @view M[length(l)+1:end,i])
        if j !== nothing
            push!(l, i)
        end
    end
    return l
end

# EigenSpaceDecomposition

function eigen_decomposition!(M::Matrix{T}) where T <: FiniteFields.GF
    eigspace_ptrs = Vector{Int}()
    eigen = left_eigen(M)
    sizehint!(eigspace_ptrs, length(eigen) + 1)
    push!(eigspace_ptrs, 1)
    for val in sort!(collect(keys(eigen))) #to get deterministic behaviour
        basis = eigen[val]
        dim = size(basis, 1)
        cd = eigspace_ptrs[end]
        ran = cd:cd+dim-1
        M[ran, :] = basis
        push!(eigspace_ptrs, cd+dim)
    end
    @assert eigspace_ptrs[end] == size(M, 1) + 1 "Matrix does not split over $T"
    return M, eigspace_ptrs
end

mutable struct EigenSpaceDecomposition{T <: FiniteFields.GF}
    basis::Matrix{T}
    eigspace_ptrs::Vector{Int}

    function EigenSpaceDecomposition(
                                     basis::Matrix{T},
                                     eigspace_ptrs::AbstractVector{<:Integer}
                                    ) where T <: FiniteFields.GF

        @assert eigspace_ptrs[1] == 1
        @assert eigspace_ptrs[end] == size(basis, 1) + 1
        return new{T}(basis, eigspace_ptrs)
    end
end

EigenSpaceDecomposition(M::Matrix{T}) where T <: FiniteFields.GF =
    EigenSpaceDecomposition(eigen_decomposition!(deepcopy(M))...)

function Base.show(io::IO, ::MIME"text/plain", esd::EigenSpaceDecomposition{T}) where T
    println(io, tuple(diff(esd.eigspace_ptrs)...), "-splitting over ", T)
    print(io, esd.basis)
end

function Base.show(io::IO, esd::EigenSpaceDecomposition{T}) where T
    print(io, tuple(diff(esd.eigspace_ptrs)...), "-splitting over ", T)
end

Base.length(esd::EigenSpaceDecomposition) = length(esd.eigspace_ptrs)-1

function Base.getindex(esd::EigenSpaceDecomposition, i::Int)
    @boundscheck 1 <= i <= length(esd)
    return esd.basis[esd.eigspace_ptrs[i]:esd.eigspace_ptrs[i+1]-1, :]
end

function Base.iterate(esd::EigenSpaceDecomposition, s=1)
    s > length(esd) && return nothing
    first_last = esd.eigspace_ptrs[s]:esd.eigspace_ptrs[s+1]-1
    return (esd.basis[first_last, :], s+1)
end

Base.eltype(esd::EigenSpaceDecomposition{T}) where T = Matrix{T}

LinearAlgebra.isdiag(esd::EigenSpaceDecomposition) = esd.eigspace_ptrs == 1:length(esd)+1

function refine(esd::EigenSpaceDecomposition{T}, M::Matrix{T}) where T
    nbasis = Array{T}(undef, 0, size(first(esd), 2))
    nptrs = [1]
    for (i, e) in enumerate(esd)
        if size(e, 1) > 1
            esd2, ptrs = eigen_decomposition!(e*M[:, _find_l(e)])
            nbasis = vcat(nbasis, esd2*e)
            append!(nptrs, ptrs.+(pop!(nptrs)-1))
        else
            nbasis = vcat(nbasis, e)
            push!(nptrs, nptrs[end]+1)
        end
    end
    return EigenSpaceDecomposition(nbasis,nptrs)
end


text:
export AbstractModel, Heisenberg
export heisenberg_ij, hamiltonian, heisenberg_term, ground_state, energy, energy_exact, get_bonds, energy, heisenberg_2d, nspin

abstract type AbstractModel{D} end
abstract type AbstractHeisenberg{D} <: AbstractModel{D} end

nspin(model::AbstractModel) = prod(size(model))

"""
Exact ground state energy.
"""
function energy_exact(tc::QuantumMPS, model::AbstractModel)
    nbit = nbit_simulated(tc)
    expect(hamiltonian(model), state_exact(tc)) |> real
end

"""
Ground state energy by sampling Quantum MPS.
The hamiltonian is limited to Heisenberg and J1J2 Type.
"""
function energy(chem::QuantumMPS, model::AbstractHeisenberg)
    energy(chem, X, model) + energy(chem, Y, model) + energy(chem, Z, model)
end

function energy(chem::QuantumMPS, pauli::PauliGate, model::AbstractHeisenberg)
    res = gensample(chem, pauli)
    local eng = 0.0
    for bond in get_bonds(model)
        eng += bond[3]*mean(res[:,bond[1]].*res[:,bond[2]])
    end
    eng/4
end

function ground_state(model::AbstractModel)
    # get the ground state
    hami = hamiltonian(model)
    E, v = eigsolve(mat(hami), 1, :SR)
    register(v[1])
end

include("Heisenberg.jl")
include("J1J2.jl")

max_stars_repo_name: Roger-luo/QuantumMPS

max_issues_repo_name: Roger-luo/QuantumMPS

max_forks_repo_name: Roger-luo/QuantumMPS


text:
using Pkg
Pkg.activate(@__DIR__)
Pkg.instantiate()

using Logging
using Statistics
using Gadfly
using Cairo
using DataFrames
using DifferentialEquations
import Distributions: Normal, Uniform

struct IVP
    f::Function
    ini::Vector
    t::AbstractArray
    phi::Vector
    data::Vector
    IVP(f, ini, t, phi, data) = new(f, ini, t, phi, data)
    IVP(f, ini, t, phi) = new(f, ini, t, phi,
                              (solve_ivp(new(f, ini, t, phi, []))
                               |> transpose_vector))
end

function get_ivps()::Dict
    var_ini = 0.01
    ivps = Dict()

    function fic(dz_dt, z, phi, t)
        r_1 = phi[1]*z[1]
        r_2 = phi[2]*z[2]

        dz_dt[1] = -r_1
        dz_dt[2] = r_1 - r_2
    end

    ivps[:fic] = Dict()
    ini = [1., 0.]
    t = range(0.0, 1.0, length=10)
    phi = [5.0035, 1.0]

    ivps[:fic][:nom] = IVP(fic, ini, t, phi)
    ivps_test = IVP(fic, add_noise(ini, var_ini), t, phi)
    ivps[:fic][:test] = IVP(fic, ivps_test.ini, t, phi,
                            add_noise(ivps_test.data, 0.1))

    ivps[:fic][:ss] = IVP(fic, ini, t,
                          [5.011535232847903,
                           0.9490955193480339], [])
    ivps[:fic][:ssr] = IVP(fic, ini, t,
                           [5.662374165061859,
                            4.600853424625171], [])
    ivps[:fic][:ds] = IVP(fic, ini, t,
                          [4.767811504376019,
                           0.9620285702765027], [])
    ivps[:fic][:dss] = IVP(fic, ini, t,
                           [5.057700471184158,
                            1.1660869766794006], [])

    function ccgo(dz_dt, z, phi, t)
        r_1 = phi[1]*z[1]^2
        r_2 = phi[2]*z[2]
        r_3 = phi[3]*z[1]^2

        dz_dt[1] = - r_1 - r_3
        dz_dt[2] = r_1 - r_2
    end

    ivps[:ccgo] = Dict()
    ini = [1., 0.]
    t = range(0.0, 0.95, length=10)
    phi = [12.214, 7.9798, 2.2216]

    ivps[:ccgo][:nom] = IVP(ccgo, ini, t, phi)
    ivps_test = IVP(ccgo, add_noise(ini, var_ini), t, phi)
    ivps[:ccgo][:test] = IVP(ccgo, ivps_test.ini, t, phi,
                            add_noise(ivps_test.data, 0.1))

    ivps[:ccgo][:ss] = IVP(ccgo, ini, t,
                           [10.601173576959617,
                            7.024603700551421,
                            5.296705518161371], [])
    ivps[:ccgo][:ssr] = IVP(ccgo, ini, t,
                            [12.694091102822359,
                             7.7461038473135275,
                             2.3847262326540015], [])
    ivps[:ccgo][:ds] = IVP(ccgo, ini, t,
                           [3.261725472047141,
                            9.713554040149868,
                            15.845961062140784], [])
    ivps[:ccgo][:dss] = IVP(ccgo, ini, t,
                            [11.701892961901253,
                             7.718104344633692,
                             2.748291378149207], [])

    function fhn(dz_dt, z, phi, t)
        a, b, c = phi
        V, R = z
        u = 0
        dz_dt[1] = c*(V - (V^3)/3 + R) + u
        dz_dt[2] = -(1/c)*(V - a + b*R)
    end

    ivps[:fhn] = Dict()
    ini = [-1.0, 1.0]
    t = range(0.0, 20.0, length=10)
    phi = [0.2, 0.2, 3.0]

    ivps[:fhn][:nom] = IVP(fhn, ini, t, phi)
    ivps_test = IVP(fhn, add_noise(ini, var_ini, false), t, phi)
    ivps[:fhn][:test] = IVP(fhn, ivps_test.ini, t, phi,
                            add_noise(ivps_test.data, 0.01, false))

    ivps[:fhn][:ss] = IVP(fhn, ini, t,
                          [1.0006533637635783e-5,
                           36762.161906015695,
                           29720.269687392283], [])
    ivps[:fhn][:ssr] = IVP(fhn, ini, t,
                           [0.1973538954650349,
                            0.11390293281470151,
                            3.0111993244825737], [])
    ivps[:fhn][:ds] = IVP(fhn, ini, t,
                          [0.06942005151231272,
                           0.4361123400524667,
                           1.696854260779007], [])
    ivps[:fhn][:dss] = IVP(fhn, ini, t,
                           [0.17695871191283072,
                            0.18851523532377604,
                            3.0158958746888795], [])
    return ivps
end

function transpose_vector(v::Vector{Vector{T}})::Vector{Vector{T}} where T
    [[v[i][j] for i in 1:length(v)] for j in 1:length(v[1])]
end

function solve_ivp(ivp::IVP)::Vector
    prob = ODEProblem(ivp.f, ivp.ini,
                      (ivp.t[1], ivp.t[end]), ivp.phi)
    osol  = solve(prob, AutoVern7(Rodas5()), saveat=ivp.t)
    return osol.u
end

function add_noise(data::Vector{T}, percent::Float64,
                   non_negative::Bool=true)::Vector{T} where T
    if percent == 0.0
        return data
    end

    epsilon_map = function(x)
        (isa(x, Vector{Vector{T}} where T) ?
         map(epsilon_map, x) :
         repeat([0.01*mean(x)], length(x)))
    end

    epsilon_iter = (map(epsilon_map, data)
                    |> Iterators.Flatten
                    |> Iterators.Stateful)

    noise_map = function(x)
        if isa(x, Vector)
            map(noise_map, x)
        else
            sigma = abs(percent*x + iterate(epsilon_iter)[1])
            d = Normal(0,sigma)
            noisy_x = x + rand(d)
            non_negative ? abs(noisy_x) : noisy_x
        end
    end

    noisy_data = map(noise_map, data)

    return noisy_data
end

"""
Normalized Mean Square Error
"""
function nmse(data::Vector{T}, data_est::Vector{T})::T where T
    normalizer = abs2(maximum(data) - minimum(data))
    res = mean(abs2.(data_est-data))/normalizer
    return res
end

"""
Normalized Root Mean Square Error
"""
function nrmse(data_nominal::Vector, data_estimated::Vector)::Float64
    state_pairs = zip(data_nominal, data_estimated)
    total_nmse = map(pair -> nmse(pair...),
                     state_pairs)
    # (mean) Normalized Root Mean Squared Error
    nrmse = mean(total_nmse) |> mean |> sqrt
    return nrmse
end

function gen_plt_err(ivp_comp::IVP,
                     ivp_est::IVP)::NamedTuple

    ini_cond = getindex.(ivp_comp.data, 1)
    ivp_est = IVP(ivp_est.f, ini_cond, ivp_est.t, ivp_est.phi)

    df_est = map(data -> DataFrame(x=ivp_est.t,
                                   y=data,
                                   Legend=repeat(["Estimated"],
                                                 length(ivp_est.t))),
                 ivp_est.data)

    df_comp = map(data -> DataFrame(x=ivp_comp.t,
                                   y=data,
                                   Legend=repeat(["Data"],
                                                 length(ivp_comp.t))),
                 ivp_comp.data)

    df_pairs = zip(df_comp, df_est)

    gen_layer(df) = layer(vcat(df...), x=:x, y=:y,
                          color=:Legend,
                          linestyle=:Legend,
                          Geom.line, Geom.point)

    layers = map(gen_layer, df_pairs)

    gen_plot(layers) = plot((layers...),
                            Guide.XLabel("t"),
                            Guide.YLabel("x"))
    plt = gen_plot(layers)

    nrmse_result = nrmse(ivp_comp.data, ivp_est.data)

    return (plt=plt, err=nrmse_result)
end

function get_res_from_dict(ivps::Dict)::Vector{NamedTuple}
    # Unpack first Dictionary layer into Vector of NamedTuple
    isest(x) = x != :nom && x != :test

    est_ivps = map(k -> (ivp=ivps[k], method=string(k)),
                   (keys(ivps) |> collect |> x -> filter(isest, x)))

    nom_test_ivps = map(k -> (ivp=ivps[k], method=string(k)),
                        (keys(ivps) |> collect |> x -> filter(!isest, x)))

    compare_pairs = [(i, j) for i in nom_test_ivps for j in est_ivps]

    results_list = map(x -> (gen_plt_err(x[1].ivp, x[2].ivp)
                             |> res -> (plt=res.plt,
                                        err=res.err,
                                        method=x[2].method,
                                        plt_type=x[1].method)),
                             compare_pairs)
    return results_list
end

function save_results(results::NamedTuple, path::AbstractString)::String
    base_path = joinpath(path, results.model)

    # Save plots
    map(x ->
        draw(PDF(joinpath(base_path, "$(x.plt_type)_$(x.method).pdf")), x.plt),
        results.res)

    # Save info
    io = joinpath(path, "info.log") |> f -> open(f, "a+")
    with_logger(SimpleLogger(io)) do
        map(x ->
            @info("NRMSE", x.plt_type, results.model, x.method, x.err),
            results.res)
    end
    close(io)
    return base_path
end

function main(args::Array{<:String})::Nothing
    path = string(args[1])

    str_mkdir(x) = (x
                    |> string
                    |> y -> joinpath(path, y)
                    |> z -> (isdir(z) ?
                             z : mkdir(z)))

    ivps = get_ivps()

    folders = map(str_mkdir,
                  (keys(ivps) |> collect))

    res_list = map(k -> (res=get_res_from_dict(ivps[k]),
                         model=string(k)),
                   (keys(ivps) |> collect))

    saved_files = map(r -> save_results(r, path), res_list)

    with_logger(ConsoleLogger()) do
        @debug res_list
        @debug saved_files
    end
    nothing
end

main(ARGS)

max_stars_repo_name: atgmello/DEParamEst

max_issues_repo_name: atgmello/DEParamEst

max_forks_repo_name: atgmello/DEParamEst


text:
using GPUArrays
using Base.Test, GPUArrays.TestSuite

function run_random(Typ)
    @testset "Random" begin
        @testset "rand" begin  # uniform
            for T in (Float32, Float64)
                @test length(rand(Typ{T,1}, (4,))) == 4
                @test length(rand(Typ, T, 4)) == 4
                @test length(rand(Typ{T,2}, (4,5))) == 20
                @test length(rand(Typ, T, 4, 5)) == 20
                A = rand(Typ{T,2}, (2,2))
                B = copy(A)
                @test all(A .== B)
                rand!(B)
                @test !any(A .== B)
            end
        end
    end
end


text:
# Day 15: Chiton
load(::Day{15}, path) = permutedims(parse.(Int, hcat(split.(readlines(path), "")...)))

# Part 1 – Path finding
# This old bookmark ended up being a useful reference:
# https://www.redblobgames.com/pathfinding/a-star/introduction.html

const VN_OFFSETS = [ # von Neumann neighborhood # TODO: refactor Day 9 to use this!
    CartesianIndex(1, 0),
    CartesianIndex(-1, 0),
    CartesianIndex(0, 1),
    CartesianIndex(0, -1),
]

function pathfind(risk)
    # Allocate matrix tracking total risk
    h, w = size(risk)
    totalrisk = fill(typemax(Int), h, w)

    # Add starting point to queue
    risk_start = 0
    totalrisk[1, 1] = risk_start
    queue = [(CartesianIndex(1, 1), risk_start)] # index, total risk

    while !isempty(queue)
        current, risk_so_far = pop!(queue)

        # Return the total risk when a path is found
        if Tuple(current) == (h, w)
            return risk_so_far
        end

        for offset in VN_OFFSETS
            # Check all neighbors iteratively
            neighbor = current + offset
            if checkbounds(Bool, risk, neighbor)
                new_total = risk_so_far + risk[neighbor]
                if new_total < totalrisk[neighbor]
                    totalrisk[neighbor] = new_total
                    push!(queue, (neighbor, new_total))
                end
            end
        end
        # Ideally I'd implement a priority queue, but we'll just sort
        sort!(queue; by=c -> -c[2])
    end
    return totalrisk[h, w]
end

solve1(::Day{15}, data) = pathfind(data)

# Part 2 – tile shifted risk maps and repeat pathfinding
function solve2(::Day{15}, data; n=5)
    mats = [mod1.(data .+ shift, 9) for shift in 0:(2 * n - 1)] # tiles
    risk = hvcat((n), [mats[x + y - 1] for x in 1:n, y in 1:n]...)
    return pathfind(risk)
end

testresult1(::Day{15}) = 40
testresult2(::Day{15}) = 315


text:
using OpenQuantumTools, Test, OrdinaryDiffEq, QuadGK, Random

H = DenseHamiltonian([(s) -> 1.0], [σi], unit=:ħ)
u0 = PauliVec[1][1]
coupling = ConstantCouplings(["Z"], unit=:ħ)

# test for OhmicBath
bath = Ohmic(1e-4, 4, 16)
annealing = Annealing(H, u0; coupling=coupling, bath=bath)
tf = 10
U = solve_unitary(annealing, tf, alg=Tsit5(), abstol=1e-8, reltol=1e-8)
sol = solve_redfield(annealing, tf, InplaceUnitary(U), alg=Tsit5(),
    reltol=1e-6)
cfun(t) = correlation(t, bath)
f(t) = quadgk(cfun, 0, t)[1]
γ = real(quadgk(f, 0, tf)[1])
@test sol(10)[1, 2] ≈ exp(-4 * γ) * 0.5 atol = 1e-5 rtol = 1e-5

sol = solve_redfield(annealing, tf, InplaceUnitary(U), vectorize=true,
    alg=TRBDF2(), reltol=1e-6)
@test sol(10)[2] ≈ exp(-4 * γ) * 0.5 atol = 1e-5 rtol = 1e-5

# test for hybrid spin-fluctuator and Redfield equation
fbath = EnsembleFluctuator([0.1], [1.0])
interactions = InteractionSet(Interaction(coupling, bath), Interaction(coupling, fbath))
annealing = Annealing(H, u0; interactions=interactions)
tf = 10
U = solve_unitary(annealing, tf, alg=Tsit5(), abstol=1e-8, reltol=1e-8)
prob1 = build_ensembles(annealing, tf, U, :redfield)
prob2 = build_ensembles(annealing, tf, U, :redfield, save_positions=(true, true))
Random.seed!(1234)
sol1 = solve(prob1, Tsit5(), EnsembleSerial(), trajectories=1, save_everystep=false)[1]
Random.seed!(1234)
sol2 = solve(prob2, Tsit5(), EnsembleSerial(), trajectories=1, save_everystep=false)[1]
@test length(sol1) == 2
@test !(length(sol2) == 2)

# test for CustomBath
H = DenseHamiltonian([(s) -> 0.0], [σi], unit=:ħ)
u0 = PauliVec[1][1]
coupling = ConstantCouplings(["Z"], unit=:ħ)
tf = 20
cfun(x) = x <= 20 ? 1e-4 : 0
bath = CustomBath(correlation=cfun)
annealing = Annealing(H, u0; coupling=coupling, bath=bath)
U = solve_unitary(annealing, tf, alg=Tsit5(), abstol=1e-8, reltol=1e-8);
sol = solve_redfield(
    annealing,
    tf,
    InplaceUnitary(U);
    alg=Tsit5(),
    abstol=1e-8,
    reltol=1e-8,
);
@test sol(tf)[1, 2] ≈ exp(-4 * 0.02) * 0.5 atol = 1e-4 rtol = 1e-4

# test for InstPulseCallback
cbu = InstPulseCallback([0.5 * tf], (c, x) -> c .= σx * c)
U = solve_unitary(
    annealing,
    tf,
    alg=Tsit5(),
    abstol=1e-8,
    reltol=1e-8,
    callback=cbu,
);
cb = InstPulseCallback([0.5 * tf], (c, x) -> c .= σx * c * σx)
sol = solve_redfield(
    annealing,
    tf,
    InplaceUnitary(U),
    alg=Tsit5(),
    abstol=1e-8,
    reltol=1e-8,
    callback=cb,
)
@test sol(tf)[1, 2] ≈ 0.5 atol = 1e-5 rtol = 1e-5

# test for CustomCouplings
H = DenseHamiltonian([(s) -> 1.0], [σz], unit=:ħ)
Uz(t) = exp(-1.0im * t * σz)
u0 = PauliVec[3][1]
coupling = ConstantCouplings(["X"], unit=:ħ)
bath = Ohmic(1e-4, 4, 16)
annealing = Annealing(H, u0; coupling=coupling, bath=bath)
tf = 10
sol = solve_redfield(annealing, tf, Uz, alg=Tsit5(), reltol=1e-6)

iH = DenseHamiltonian([(s) -> 0.0], [σz], unit=:ħ)
icoupling = CustomCouplings([(s) -> Uz(10 * s)' * σx * Uz(10 * s)], unit=:ħ)
iannealing = Annealing(iH, u0; coupling=icoupling, bath=bath)
isol = solve_redfield(iannealing, tf, (t) -> σi, alg=Tsit5(), reltol=1e-6)

@test sol[end][1,1] ≈ isol[end][1,1] atol = 1e-6 rtol = 1e-6

# Test of non-positivity check
β = 4
T = β_2_temperature(β)
η = 0.1
fc = 10 / (2π)
bath = Ohmic(η, fc, T)

Hp = 0.5 * σz ⊗ σi - 0.7 * σi ⊗ σz + 0.3 * σz ⊗ σz
Hd = standard_driver(2)
H = DenseHamiltonian([(s) -> 1 - s, (s) -> s], [-Hd, Hp], unit=:ħ)

tf = 20
ρ0 = (σi + σx) ⊗ (σi + σx) / 4
coupling = ConstantCouplings([σz ⊗ σi, σi ⊗ σz], unit=:ħ)
annealing = Annealing(H, ρ0, bath=bath, coupling=coupling)

U = solve_unitary(annealing, tf, alg=Tsit5(), abstol=1e-7, reltol=1e-7)

redfield_sol = solve_redfield(annealing, tf, U, alg=Tsit5(), abstol=1e-7, reltol=1e-7, callback=PositivityCheckCallback())

@test redfield_sol.t[end] ≈ 3.45278167 atol = 1e-5 rtol = 1e-5


text:
function CH(p)
 	#A = []
	D = Dict{Plane3,Bool}()
	F = Tuple{Int64,Int64,Int64}[]
	scatter3d(getindex.(p,2), getindex.(p,3), getindex.(p,4), markershape=:circle,legend=false)
	for i in 1:length(p), j in (i+1):length(p), k in (i+1):length(p)
		(j==k) && continue
		pln = plane(p[i],p[j],p[k])
		if (~(get(D,pln,false)) && ~(all(pln.==0)) && ~(ptsAbove(p,pln)))
			#push!(A,ptsA,ptsB,ptsC)
			#@warn string("""Points: """,ptsA,ptsB,ptsC)
			pts = [p[i];p[j];p[k];p[i]]
			D[pln] = true;
			push!(F,((i-1),(j-1),(k-1)))
			#push!(A,pts)
			#@info string("""Face Values: """,F)
			display(plot3d!(getindex.(pts,2), getindex.(pts,3), getindex.(pts,4), markershape=:circle))
		end
	end
	@warn string("""The length of the vertices array is: """,length(p))
	@info """Please input a filename without the file type (i.e. loremipsum will produce loremipsum.py):"""
	Export("Blender",[string(readline(stdin),""".py"""),p,F])
	#return A
	return p, F
end

function ptsAbove(pts,pl::Plane3)
	for p in pts
		abovePl(p,pl) && return true
	end
		return false
end

max_stars_repo_name: snoeyink/CGinJulia

max_issues_repo_name: snoeyink/CGinJulia

max_forks_repo_name: snoeyink/CGinJulia


text:
using ZigZagBoomerang
using Random
Random.seed!(1)
# negative log-density with respect to Lebesgue
ϕ(x) = cos(π*x) + x^2/2 # not needed

# gradient of ϕ(x)
∇ϕ(x) = -π*sin(π*x) + x # (REPLACE IT WITH AUTOMATIC DIFFERENTIATION)

# Example: ZigZag
x0, θ0 = randn(), 1.0
T = 300.0
out1, acc = ZigZagBoomerang.pdmp(∇ϕ, x0, θ0, T, 1.2π, ZigZag1d())
@show acc

# Example: Boomerang
B = Boomerang1d(0.1)
out2, acc = ZigZagBoomerang.pdmp(∇ϕ, x0, θ0, T, 0.00001, B)
@show acc

using Makie, CairoMakie
p1 = Makie.lines(eventtime.(out1), eventposition.(out1))
save("figures/zigzag.png", title(p1, "ZigZag 1d"))

dt = 0.01
xx = ZigZagBoomerang.discretize(out2, B, dt)
p2 = Makie.lines(xx.t, xx.x, linewidth=0.4)
save("figures/boomerang.png", title(p2, "Boomerang 1d"))

dt = 0.1
xx = ZigZagBoomerang.discretize(out2, B, dt)
p2 = Makie.lines(xx.t[1:end÷3], xx.x[1:end÷3], linewidth=0.2)
Makie.scatter!(p2, xx.t[1:end÷3], xx.x[1:end÷3], markersize=0.3)
save("figures/boomerangshort.png", title(p2, "Boomerang 1d"))


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder, Pkg

name = "libftd2xx"
version = v"1.0.0"

# Collection of sources required to complete build
sources = [
    ArchiveSource("https://www.ftdichip.com/Drivers/CDM/CDM v2.12.28 WHQL Certified.zip",
    "82db36f089d391f194c8ad6494b0bf44c508b176f9d3302777c041dad1ef7fe6", unpack_target="windows"),

    FileSource("https://www.ftdichip.com/Drivers/D2XX/Linux/libftd2xx-i386-1.4.8.gz",
    "84c9aaf7288a154faf0e2814ba590ec965141c698b2a00bffc94d8e0c7ebeb4c"),

    FileSource("https://www.ftdichip.com/Drivers/D2XX/Linux/libftd2xx-x86_64-1.4.8.gz",
    "815d880c5ec40904f062373e52de07b2acaa428e54fece98b31e6573f5d261a0"),

    FileSource("https://www.ftdichip.com/Drivers/D2XX/Linux/libftd2xx-arm-v7-hf-1.4.8.gz",
    "81c8556184e9532a3a19ee6915c3a43110dc208116967a4d3e159f00db5d16e1"),

    FileSource("https://www.ftdichip.com/Drivers/D2XX/Linux/libftd2xx-arm-v8-1.4.8.gz",
    "e353cfa94069dee6d5bba1c4d8a19b0fd2bf3db1e8bbe0c3b9534fdfaf7a55ed"),

    FileSource("https://www.ftdichip.com/Drivers/D2XX/MacOSX/D2XX1.4.16.dmg",
    "757ef22c3e48c2022974c2110d25ee45dd09bff8f397c8432018c50fb4b2d785"),

    DirectorySource("./bundled"),
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir

# Add license file
install_license LICENSE

mkdir ${libdir}

if [[ ${target} == x86_64-linux-* ]]; then
    tar zxvf libftd2xx-x86_64-1.4.8.gz;
    cp -v release/build/*.${dlext}* ${libdir};
fi

if [[ ${target} == i686-linux-* ]]; then
    tar zxvf libftd2xx-i386-1.4.8.gz;
    cp -v release/build/*.${dlext}* ${libdir};
fi

if [[ ${target} == aarch64-linux-* ]]; then
    tar zxvf libftd2xx-arm-v8-1.4.8.gz;
    cp -v release/build/*.${dlext}* ${libdir};
fi

if [[ ${target} == arm-linux-* ]]; then
    tar zxvf libftd2xx-arm-v7-hf-1.4.8.gz;
    cp -v release/build/*.${dlext}* ${libdir};
fi

if [[ ${target} == x86_64-apple-darwin* ]]; then
    apk add p7zip;
    7z x D2XX1.4.16.dmg;
    cp -v release/D2XX/*.${dlext}* ${libdir};
fi

if [[ ${target} == x86_64-w64-mingw32 ]]; then
    cp -v windows/amd64/*.${dlext} ${libdir};
fi

if [[ ${target} == i686-w64-mingw32 ]]; then
    cp -v windows/i386/*.${dlext} ${libdir};
fi
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = [
    Linux(:armv7l, libc=:glibc, call_abi=:eabihf),
    Linux(:x86_64, libc=:glibc),
    Linux(:aarch64, libc=:glibc),
    Linux(:i686, libc=:glibc),
    MacOS(:x86_64),
    Windows(:x86_64),
]

# The products that we will ensure are always built
products = [
    LibraryProduct(["libftd2xx", "ftd2xx", "ftd2xx64"], :libftd2xx)
]

# Dependencies that must be installed before this package can be built
dependencies = Dependency[
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies)


text:
@doc raw"""
    ImplicitMF(dim::Int64, use_bias::Bool, reg_coeff::Float64)

Matrix factorization model for implicit feedback. The predicted rating for item ``i`` by user ``u`` is expreseed as
```math
\hat r_{ui} = \mu + b_i + b_u + \bm u_u \cdot \bm v_i\,,
```
Unlike the model for explicit feedback, the model treats all the (user, item) pairs in the train dataset as positive interaction with label 1, and sample negative (user, item) pairs from the corpus. Currently only the uniform item sampling is implemented. The fitting criteria is the ordinary logloss function
```math
    L = -r_{ui}\log(\hat r_{ui}) - (1 - r_{ui})\log(1 - \hat r_{ui}).
```

# Constructor arguments
- `dim`: dimension of user/item vectors.
- `use_bias`: if set to false, the bias terms (``\mu``, ``b_i``, ``b_u``) are set to zero.
- `reg_coeff`: ``L_2`` regularization coefficients for model parameters.

# References
For instance, Rendle et. al. (2020), [Neural Collaborative Filtering vs. Matrix Factorization Revisited
](http://arxiv.org/abs/2005.09683).
"""
mutable struct ImplicitMF <: AbstractRecommender
    dim::Int64
    loss::LossFunction
    use_bias::Bool
    reg_coeff::Float64

    μ::Union{Nothing,Float64}
    user_bias::Union{Nothing,Vector{Float64}}
    item_bias::Union{Nothing,Vector{Float64}}
    user_embedding::Union{Nothing,Matrix{Float64}}
    item_embedding::Union{Nothing,Matrix{Float64}}

    user2uidx::Union{Dict,Nothing}
    item2iidx::Union{Dict,Nothing}
    iidx2item::Union{Dict,Nothing}
    user_history::Union{Dict,Nothing}

    ImplicitMF(dim::Int64, use_bias::Bool, reg_coeff::Float64) = new(
        dim,
        Logloss(), # default
        use_bias,
        reg_coeff,
        nothing,
        nothing,
        nothing,
        nothing,
        nothing,
        nothing,
        nothing,
        nothing,
        nothing,
    )
end

function predict(model::ImplicitMF, uidx, iidx)::Float64
    pred = model.user_embedding[:, uidx]' * model.item_embedding[:, iidx]
    if model.use_bias
        pred += model.μ + model.user_bias[uidx] + model.item_bias[iidx]
    end
    return pred
end

function sgd!(model::ImplicitMF, uidx, iidx, grad_value, lr)
    reg = model.reg_coeff

    if model.use_bias
        model.user_bias[uidx] -= lr * (grad_value + reg * model.user_bias[uidx])
        model.item_bias[iidx] -= lr * (grad_value + reg * model.item_bias[iidx])
        model.μ -= lr * (grad_value + reg * model.μ)
    end
    uvec = @view model.user_embedding[:, uidx]
    ivec = @view model.item_embedding[:, iidx]
    model.user_embedding[:, uidx] -= lr * (grad_value * ivec + reg * uvec)
    model.item_embedding[:, iidx] -= lr * (grad_value * uvec + reg * ivec)
end

"""
    fit!(model::ImplicitMF, table; callbacks = Any[], col_user = :userid, col_item = :item_id, n_epochs = 2, learning_rate = 0.01, n_negatives = 1, verbose = -1)

Fit the `ImplicitMF` model by stochastic grandient descent (with no batching).

# Model-specific arguments
- `n_epochs`: number of epochs. During one epoch, all the row in `table` is read once.
- `learning_rate`: Learing rate of SGD.
- `n_negatives`: Number of negative item samples per positive (user, item) pair.
- `verbose`: If set to positive integer, the training info is printed once per `verbose`.
- `callbacks`: Additional callback functions during SGD. One can implement, for instance, monitoring the validation metrics and the early stopping. See [Callbacks](@ref).
"""
function fit!(
    model::ImplicitMF,
    table;
    callbacks = Any[],
    col_user = :userid,
    col_item = :item_id,
    n_epochs = 2,
    learning_rate = 0.01,
    n_negatives = 1,
    verbose = -1,
    kwargs...,
)
    model.user_history = Dict()
    for (userid, history) in
        zip(make_u2i_dataset(table, col_user = col_user, col_item = col_item)...)
        model.user_history[userid] = history
    end

    table, model.user2uidx, model.item2iidx, model.iidx2item =
        make_idmap(table, col_user = col_user, col_item = col_item)

    n_user = length(keys(model.user2uidx))

    unique_items = collect(keys(model.iidx2item))
    n_item = length(unique_items)

    model.user_embedding = rand(model.dim, n_user)
    model.item_embedding = rand(model.dim, n_item)

    if model.use_bias
        model.user_bias = rand(n_user)
        model.item_bias = rand(n_item)
        model.μ = rand()
    end

    # callback is any callbale with same interface
    callbacks = append!(Any[LogTrainLoss()], callbacks)
    for cb in callbacks
        if typeof(cb) <: AbstractCallback
            initialize!(cb, model, col_user = col_user, col_item = col_item)
        end
    end

    for epoch = 1:n_epochs
        train_loss = 0
        n_sample = 0
        for row in Tables.rows(table)
            uidx = row[col_user]

            # update positive
            iidx = row[col_item]

            pred = predict(model, uidx, iidx)
            train_loss = (model.loss(pred, 1) + train_loss * n_sample) / (n_sample + 1)
            n_sample += 1

            grad_value = grad(model.loss, pred, 1)
            sgd!(model, uidx, iidx, grad_value, learning_rate)

            # negative samples
            for _ = 1:n_negatives
                iidx = rand(unique_items)

                pred = predict(model, uidx, iidx)
                train_loss = (model.loss(pred, 0) + train_loss * n_sample) / (n_sample + 1)
                n_sample += 1

                grad_value = grad(model.loss, pred, 0)
                sgd!(model, uidx, iidx, grad_value, learning_rate)
            end
        end

        try
            for cb in callbacks
                cb(model, train_loss, epoch, verbose)
            end
        catch StopTrain
            break
        end
    end

end

@doc raw"""
    predict_u2i(model::ImplicitMF, userid::Union{AbstractString,Int}, n::Int64; drop_history = false)

Make predictions by using ``\hat r_{ui}``.
"""
function predict_u2i(
    model::ImplicitMF,
    userid::Union{AbstractString,Int},
    n::Int64;
    drop_history = false,
    kwargs...,
)
    if !(userid in keys(model.user2uidx))
        return []
    end
    uidx = model.user2uidx[userid]
    unique_iidx = collect(keys(model.iidx2item))
    preds = [predict(model, uidx, iidx) for iidx in unique_iidx]
    pred_iidx = unique_iidx[sortperm(preds, rev = true)]
    pred_items = [model.iidx2item[iidx] for iidx in pred_iidx]
    if drop_history
        filter!(e -> !(e in model.user_history[userid]), pred_items)
    end
    n = min(n, length(pred_items))
    return pred_items[1:n]
end


text:
const var_type_map = Dict(
KTR_VARTYPE_CONTINUOUS => :Cont,
KTR_VARTYPE_INTEGER => :Int,
KTR_VARTYPE_BINARY => :Bin
)

const rev_var_type_map = Dict(
:Cont => KTR_VARTYPE_CONTINUOUS,
:Int => KTR_VARTYPE_INTEGER,
:Bin => KTR_VARTYPE_BINARY
)

# grep "#define" knitro.h | grep "KTR_PARAM_" | awk '{ printf("\"%s\" => Int32(%s),\n",$2,$3) }'
const paramName2Indx = Dict(
"KTR_PARAM_NEWPOINT" => Int32(1001),
"KTR_PARAM_HONORBNDS" => Int32(1002),
"KTR_PARAM_ALGORITHM" => Int32(1003),
"KTR_PARAM_ALG" => Int32(1003),
"KTR_PARAM_BAR_MURULE" => Int32(1004),
"KTR_PARAM_BARRULE" => Int32(1004),
"KTR_PARAM_BAR_FEASIBLE" => Int32(1006),
"KTR_PARAM_FEASIBLE" => Int32(1006),
"KTR_PARAM_GRADOPT" => Int32(1007),
"KTR_PARAM_HESSOPT" => Int32(1008),
"KTR_PARAM_BAR_INITPT" => Int32(1009),
"KTR_PARAM_ACT_LPSOLVER" => Int32(1012),
"KTR_PARAM_LPSOLVER" => Int32(1012),
"KTR_PARAM_MAXCGIT" => Int32(1013),
"KTR_PARAM_MAXIT" => Int32(1014),
"KTR_PARAM_OUTLEV" => Int32(1015),
"KTR_PARAM_OUTMODE" => Int32(1016),
"KTR_PARAM_SCALE" => Int32(1017),
"KTR_PARAM_SHIFTINIT" => Int32(1018),
"KTR_PARAM_SOC" => Int32(1019),
"KTR_PARAM_DELTA" => Int32(1020),
"KTR_PARAM_BAR_FEASMODETOL" => Int32(1021),
"KTR_PARAM_FEASMODETOL" => Int32(1021),
"KTR_PARAM_FEASTOL" => Int32(1022),
"KTR_PARAM_FEASTOLABS" => Int32(1023),
"KTR_PARAM_MAXTIMECPU" => Int32(1024),
"KTR_PARAM_BAR_INITMU" => Int32(1025),
"KTR_PARAM_MU" => Int32(1025),
"KTR_PARAM_OBJRANGE" => Int32(1026),
"KTR_PARAM_OPTTOL" => Int32(1027),
"KTR_PARAM_OPTTOLABS" => Int32(1028),
"KTR_PARAM_PIVOT" => Int32(1029),
"KTR_PARAM_XTOL" => Int32(1030),
"KTR_PARAM_DEBUG" => Int32(1031),
"KTR_PARAM_MULTISTART" => Int32(1033),
"KTR_PARAM_MSENABLE" => Int32(1033),
"KTR_PARAM_MSMAXSOLVES" => Int32(1034),
"KTR_PARAM_MSMAXBNDRANGE" => Int32(1035),
"KTR_PARAM_MSMAXTIMECPU" => Int32(1036),
"KTR_PARAM_MSMAXTIMEREAL" => Int32(1037),
"KTR_PARAM_LMSIZE" => Int32(1038),
"KTR_PARAM_BAR_MAXCROSSIT" => Int32(1039),
"KTR_PARAM_MAXCROSSIT" => Int32(1039),
"KTR_PARAM_MAXTIMEREAL" => Int32(1040),
"KTR_PARAM_PRECOND" => Int32(1041),
"KTR_PARAM_BLASOPTION" => Int32(1042),
"KTR_PARAM_BAR_MAXREFACTOR" => Int32(1043),
"KTR_PARAM_BAR_MAXBACKTRACK" => Int32(1044),
"KTR_PARAM_BLASOPTIONLIB" => Int32(1045),
"KTR_PARAM_OUTAPPEND" => Int32(1046),
"KTR_PARAM_OUTDIR" => Int32(1047),
"KTR_PARAM_CPLEXLIB" => Int32(1048),
"KTR_PARAM_BAR_PENRULE" => Int32(1049),
"KTR_PARAM_BAR_PENCONS" => Int32(1050),
"KTR_PARAM_MSNUMTOSAVE" => Int32(1051),
"KTR_PARAM_MSSAVETOL" => Int32(1052),
"KTR_PARAM_PRESOLVEDEBUG" => Int32(1053),
"KTR_PARAM_MSTERMINATE" => Int32(1054),
"KTR_PARAM_MSSTARTPTRANGE" => Int32(1055),
"KTR_PARAM_INFEASTOL" => Int32(1056),
"KTR_PARAM_LINSOLVER" => Int32(1057),
"KTR_PARAM_BAR_DIRECTINTERVAL" => Int32(1058),
"KTR_PARAM_PRESOLVE" => Int32(1059),
"KTR_PARAM_PRESOLVE_TOL" => Int32(1060),
"KTR_PARAM_BAR_SWITCHRULE" => Int32(1061),
"KTR_PARAM_HESSIAN_NO_F" => Int32(1062),
"KTR_PARAM_MA_TERMINATE" => Int32(1063),
"KTR_PARAM_MA_MAXTIMECPU" => Int32(1064),
"KTR_PARAM_MA_MAXTIMEREAL" => Int32(1065),
"KTR_PARAM_MSSEED" => Int32(1066),
"KTR_PARAM_MA_OUTSUB" => Int32(1067),
"KTR_PARAM_MS_OUTSUB" => Int32(1068),
"KTR_PARAM_XPRESSLIB" => Int32(1069),
"KTR_PARAM_TUNER" => Int32(1070),
"KTR_PARAM_TUNER_OPTIONSFILE" => Int32(1071),
"KTR_PARAM_TUNER_MAXTIMECPU" => Int32(1072),
"KTR_PARAM_TUNER_MAXTIMEREAL" => Int32(1073),
"KTR_PARAM_TUNER_OUTSUB" => Int32(1074),
"KTR_PARAM_TUNER_TERMINATE" => Int32(1075),
"KTR_PARAM_LINSOLVER_OOC" => Int32(1076),
"KTR_PARAM_BAR_RELAXCONS" => Int32(1077),
"KTR_PARAM_MSDETERMINISTIC" => Int32(1078),
"KTR_PARAM_BAR_REFINEMENT" => Int32(1079),
"KTR_PARAM_DERIVCHECK" => Int32(1080),
"KTR_PARAM_DERIVCHECK_TYPE" => Int32(1081),
"KTR_PARAM_DERIVCHECK_TOL" => Int32(1082),
"KTR_PARAM_LINSOLVER_INEXACT" => Int32(1083),
"KTR_PARAM_LINSOLVER_INEXACTTOL" => Int32(1084),
"KTR_PARAM_MAXFEVALS" => Int32(1085),
"KTR_PARAM_FSTOPVAL" => Int32(1086),
"KTR_PARAM_DATACHECK" => Int32(1087),
"KTR_PARAM_DERIVCHECK_TERMINATE" => Int32(1088),
"KTR_PARAM_BAR_WATCHDOG" => Int32(1089),
"KTR_PARAM_FTOL" => Int32(1090),
"KTR_PARAM_FTOL_ITERS" => Int32(1091),
"KTR_PARAM_ACT_QPALG" => Int32(1092),
"KTR_PARAM_MIP_METHOD" => Int32(2001),
"KTR_PARAM_MIP_BRANCHRULE" => Int32(2002),
"KTR_PARAM_MIP_SELECTRULE" => Int32(2003),
"KTR_PARAM_MIP_INTGAPABS" => Int32(2004),
"KTR_PARAM_MIP_INTGAPREL" => Int32(2005),
"KTR_PARAM_MIP_MAXTIMECPU" => Int32(2006),
"KTR_PARAM_MIP_MAXTIMEREAL" => Int32(2007),
"KTR_PARAM_MIP_MAXSOLVES" => Int32(2008),
"KTR_PARAM_MIP_INTEGERTOL" => Int32(2009),
"KTR_PARAM_MIP_OUTLEVEL" => Int32(2010),
"KTR_PARAM_MIP_OUTINTERVAL" => Int32(2011),
"KTR_PARAM_MIP_OUTSUB" => Int32(2012),
"KTR_PARAM_MIP_DEBUG" => Int32(2013),
"KTR_PARAM_MIP_IMPLICATNS" => Int32(2014),
"KTR_PARAM_MIP_GUB_BRANCH" => Int32(2015),
"KTR_PARAM_MIP_KNAPSACK" => Int32(2016),
"KTR_PARAM_MIP_ROUNDING" => Int32(2017),
"KTR_PARAM_MIP_ROOTALG" => Int32(2018),
"KTR_PARAM_MIP_LPALG" => Int32(2019),
"KTR_PARAM_MIP_TERMINATE" => Int32(2020),
"KTR_PARAM_MIP_MAXNODES" => Int32(2021),
"KTR_PARAM_MIP_HEURISTIC" => Int32(2022),
"KTR_PARAM_MIP_HEUR_MAXIT" => Int32(2023),
"KTR_PARAM_MIP_HEUR_MAXTIMECPU" => Int32(2024),
"KTR_PARAM_MIP_HEUR_MAXTIMEREAL" => Int32(2025),
"KTR_PARAM_MIP_PSEUDOINIT" => Int32(2026),
"KTR_PARAM_MIP_STRONG_MAXIT" => Int32(2027),
"KTR_PARAM_MIP_STRONG_CANDLIM" => Int32(2028),
"KTR_PARAM_MIP_STRONG_LEVEL" => Int32(2029),
"KTR_PARAM_PAR_NUMTHREADS" => Int32(3001),
"KTR_PARAM_PAR_CONCURRENT_EVALS" => Int32(3002),
"KTR_PARAM_PAR_BLASNUMTHREADS" => Int32(3003),
"KTR_PARAM_PAR_LSNUMTHREADS" => Int32(3004)
)


text:
module SFS

export sfs, sfs_TVS, sfs_CV

include("../data.jl")
using .DataUtil
using Statistics, LinearAlgebra
using InvertedIndices

# forward stepwise feature selection with step size k
function sfs(𝐗, 𝐲, k)
    @assert k ≤ size(𝐗, 2)

    𝐴 = Array{Int}(undef, k)
    𝐬 = Array{Float64}(undef, k)
    for t = 1:k
        𝐴_t = 𝐴[1:(t - 1)]
        𝐴ᶜ_t = collect(1:size(𝐗, 2))[Not(𝐴_t)]
        𝐗_𝐴 = 𝐗[:, 𝐴_t]
        𝐫 = 𝐲 - 𝐗_𝐴 / (𝐗_𝐴' * 𝐗_𝐴) * (𝐗_𝐴' * 𝐲)
        j_t = 𝐴ᶜ_t[argmax(abs.(𝐗[:, 𝐴ᶜ_t]' * 𝐫))] # new feature
        𝐴[t] = j_t
        𝐬[t] = float(sign(𝐗[:,j_t]' * 𝐫))
    end
    return (𝐴, 𝐬)
end

# train-val-split for determining the hyperparameter "k" in candidate list K
function sfs_TVS(𝐗, 𝐲, K; ratio=0.5)
    𝐗_train, 𝐗_val = train_val_split(𝐗; ratio=ratio, do_centering=true)
    𝐲_train, 𝐲_val = train_val_split(𝐲; ratio=ratio, do_centering=true)
    k = K[argmin((k -> validation_error(k, 𝐗_train, 𝐗_val, 𝐲_train, 𝐲_val)).(K))]
    return (sfs(𝐗, 𝐲, k)..., k)
end

# cross-validation for determining the hyperparameter "k" in candidate list K
function sfs_CV(𝐗, 𝐲, K; 𝑘=length(𝐲))
    k = K[argmin((k -> cv_error(k, 𝐗, 𝐲, 𝑘)).(K))]
    return (sfs(𝐗, 𝐲, k)..., k)
end

function validation_error(k, 𝐗_train, 𝐗_val, 𝐲_train, 𝐲_val)
    𝐴_k, _ = sfs(𝐗_train, 𝐲_train, k)
    𝐗_train_𝐴 = 𝐗_train[:, 𝐴_k] # active train input
    𝐗_val_𝐴 = 𝐗_val[:, 𝐴_k] # active validation input
    𝐗₊ = (𝐗_train_𝐴' * 𝐗_train_𝐴) \ 𝐗_train_𝐴' # pseudo inverse matrix of 𝐗_train_𝐴
    𝐞 = 𝐲_val - 𝐗_val_𝐴 * 𝐗₊ * 𝐲_train
    return 𝐞' * 𝐞
end

function cv_error(k, 𝐗, 𝐲, 𝑘)
    𝑘_fold = KFold(length(𝐲), 𝑘) # 𝑘-fold Cross Validation
    return mean(
        begin
        𝐗_train, 𝐗_val = train_val_split(𝐗, 𝑘_fold.parts[i], 𝑘_fold.parts[i + 1] - 1)
        𝐲_train, 𝐲_val = train_val_split(𝐲, 𝑘_fold.parts[i], 𝑘_fold.parts[i + 1] - 1)
        validation_error(k, 𝐗_train, 𝐗_val, 𝐲_train, 𝐲_val)
    end for i = 1:𝑘)
end

end # module SFS


text:
include("6.1 Hyperbolic sine/.jl")
include("6.3 Hyperbolic tangent/.jl")
include("6.5 Hyperbolic secant/.jl")
include("6.7 Miscellaneous/.jl")

max_stars_repo_path: src/rules/6 Hyperbolic functions/.jl

max_issues_repo_path: src/rules/6 Hyperbolic functions/.jl

max_forks_repo_path: src/rules/6 Hyperbolic functions/.jl


text:
"""
$(TYPEDEF)

A collection of DataDrivenProblems used to concatenate different trajectories or experiments.

Can be called with either a `NTuple` of problems or a `NamedTuple` of `NamedTuples`.
Similar to the `DataDrivenProblem`, it has three constructors available:

+ `DirectDataset` for direct problems
+ `DiscreteDataset` for discrete problems
+ `ContinuousDataset` for continuous problems

# Fields
$(FIELDS)

# Signatures
$(SIGNATURES)

"""
struct DataDrivenDataset{N, U, C} <: AbstractDataDrivenProblem{N,U,C}
    """Name of the dataset"""
    name::Symbol
    """The problems"""
    probs::NTuple{M, AbstractDataDrivenProblem{N,U,C}} where M
    """The length of each problem - for internal use"""
    sizes::NTuple{M, Int} where M
end

# Constructor

function DataDrivenDataset(probs::Vararg{T, N}; name = gensym(:DDSet), kwargs...) where {T <: AbstractDataDrivenProblem, N}
    return DataDrivenDataset(name, probs, map(length, probs))
end

function DataDrivenDataset(solutions::Vararg{T, N}; name = gensym(:DDSet), kwargs...) where {T <: DiffEqBase.DESolution, N}
    probs = map(solutions) do s
        DataDrivenProblem(s; kwargs...)
    end
    return DataDrivenDataset(probs..., name = name)
end

"""
A direct `DataDrivenDataset` useable for problems of the form `f(x,p,t,u) ↦ y`.

$(SIGNATURES)
"""
function DirectDataset(s::NamedTuple; name = gensym(:DDSet), kwargs...)
    probs = map(keys(s)) do k
        si = s[k]
        _kwargs = collect_problem_kwargs(si; kwargs...)
        DataDrivenProblem(si[:X]; probtype = DDProbType(1), _kwargs...)
    end

    DataDrivenDataset(probs...; name = name)
end

"""
A time discrete `DataDrivenDataset` useable for problems of the form `f(x,p,t,u) ↦ x(t+1)`.

$(SIGNATURES)
"""
function DiscreteDataset(s::NamedTuple; name = gensym(:DDSet), kwargs...)
    probs = map(keys(s)) do k
        si = s[k]
        _kwargs = collect_problem_kwargs(si; kwargs...)
        DataDrivenProblem(si[:X]; probtype = DDProbType(2), _kwargs...)
    end
    DataDrivenDataset(probs...; name = name)
end

"""
A time continuous `DataDrivenDataset` useable for problems of the form `f(x,p,t,u) ↦ dx/dt`.

$(SIGNATURES)

Automatically constructs derivatives via an additional collocation method, which can be either a collocation
or an interpolation from `DataInterpolations.jl` wrapped by an `InterpolationMethod` provided by the `collocation` keyworded argument.
"""
function ContinuousDataset(s::NamedTuple; name = gensym(:DDSet), collocation = InterpolationMethod(), kwargs...)
    probs = map(keys(s)) do k
        si = s[k]
        # Check for differential states
        _kwargs = collect_problem_kwargs(si; kwargs...)
        if :DX ∈ keys(_kwargs)
            return DataDrivenProblem(si[:X]; probtype = DDProbType(3), _kwargs...)
        elseif :t ∈ keys(_kwargs)
            dx, x = collocate_data(si[:X], si[:t], collocation; kwargs...)
            return DataDrivenProblem(x; DX = dx, probtype = DDProbType(3), _kwargs...)
        else
            throw(ArgumentError("A continuous problem $(k) needs to have either derivative or time information specified!"))
        end
    end
    DataDrivenDataset(probs...; name = name)
end

collect_problem_kwargs(s ; kwargs...) = begin
    _kwargs = Dict()
    for k in keys(s)
        if k ∈ [:DX, :t, :Y, :U, :p] # Very specific subset
            _kwargs[k] = s[k]
        end
    end
    merge(_kwargs, kwargs)
end

Base.length(s::DataDrivenDataset) = sum(s.sizes)
Base.size(s::DataDrivenDataset) = (first(size(first(s.probs))), length(s))

function Base.summary(io::IO, x::DataDrivenDataset{N,C,P}) where {N,C,P}
    print(io, "$P Dataset{$N} $(x.name) with $(length(x.probs)) problems")
    n,m = size(x)
    print(io, " in $n dimensions and $m samples")
    C ? nothing : print(io, " with controls")
    return
end

Base.show(io::IO, x::DataDrivenDataset{N,C,P}) where {N,C,P} = summary(io, x)

function Base.print(io::IO, x::DataDrivenDataset{N,C,P}) where {N,C,P}
    println(io, "$P DataDrivenDataset{$N} $(x.name)")
    println(io, "Summary")
    map(x.probs) do p
        print(io, p)
        println(io)
    end
    return
end

function is_valid(x::DataDrivenDataset)
    all(map(is_valid, x.probs))
end

function get_target(x::DataDrivenDataset)
    reduce(hcat, map(get_target, x.probs))
end

function init_implicits(x::DataDrivenDataset{N, W, C}) where {N, W, C}
    first_prob = first(x.probs)
    n_x, m = size(x)
    n_u = size(first_prob.U, 1)
    n_y = size(get_target(first_prob), 1)
    return (
        zeros(N, n_y+n_x, m),
        parameters(first_prob),
        zeros(N, m),
        n_u > 0 ? zeros(N, n_u, m) : N[]
    )
end

function init_explicit(x::DataDrivenDataset{N, W, C}) where {N, W, C}
    first_prob = first(x.probs)
    n_x, m = size(x)
    n_u = size(first_prob.U, 1)
    n_y = size(get_target(first_prob), 1)
    return (
        zeros(N, n_x, m),
        parameters(first_prob),
        zeros(N, m),
        n_u > 0 ? zeros(N, n_u, m) : N[]
    )
end

function get_oop_args(x::DataDrivenDataset{N,W,C}) where {N,W,C}
    X, p, t, U = init_explicit(x)
    last = 1
    @views for (i, s) in enumerate(cumsum(x.sizes))
        # Only copy if U is present
        if !W
            map(copyto!, (X[:, last:s], p, t[last:s],U[:, last:s]), get_oop_args(x.probs[i]))
        else
            map(copyto!, (X[:, last:s], p, t[last:s]), get_oop_args(x.probs[i])[1:3])
        end
        last += s
    end
    return (X, p, t, U)
end

function get_implicit_oop_args(x::DataDrivenDataset{N,W,C}) where {N,W,C}
    X, p, t, U = init_implicits(x)
    last = 1
    @views for (i, s) in enumerate(cumsum(x.sizes))
        # Only copy if U is present
        if !W
            map(copyto!, (X[:, last:s], p, t[last:s],U[:, last:s]), get_implicit_oop_args(x.probs[i]))
        else
            map(copyto!, (X[:, last:s], p, t[last:s]), get_implicit_oop_args(x.probs[i])[1:3])
        end
        last += s
    end
    return (X, p, t, U)
end

function (b::AbstractBasis)(d::DataDrivenDataset)
    reduce(hcat, map(b, d.probs))
end

function (b::AbstractBasis)(dx::AbstractMatrix, d::DataDrivenDataset)
    last = 1
    @views for (i,s) in enumerate(cumsum(d.sizes))
        b(dx[:, last:s], d.probs[i])
        last += s
    end
    return
end


text:
using SparseIR
using Test

import PyCall: pyimport, PyNULL, PyVector
import LegendrePolynomials: Pl

sparse_ir = pyimport("sparse_ir")

@testset "augment.LegendreBasis" begin
    for stat in [fermion, boson]
        beta = 10.0
        Nl = 100
        cl = sqrt.(2 .* collect(0:(Nl - 1)) .+ 1)
        basis = SparseIR.LegendreBasis(stat, beta, Nl; cl=cl)

        tau = Float64[0, 0.1 * beta, 0.4 * beta, beta]
        uval = basis.u(tau)

        ref = Matrix{Float64}(undef, Nl, length(tau))
        for l in 0:(Nl - 1)
            x = @. 2 * tau / beta - 1.0
            ref[l + 1, :] = cl[l + 1] .* (sqrt(2 * l + 1) / beta) * Pl.(x, l)
        end
        @test isapprox(uval, ref)

        sign = statistics_sign(stat)

        # G(iv) = 1/(iv-pole)
        # G(tau) = -e^{-tau*pole}/(1 + e^{-beta*pole}) [F]
        #        = -e^{-tau*pole}/(1 - e^{-beta*pole}) [B]
        pole = 1.0
        tau_smpl = TauSampling(basis)
        gtau = -exp.(-tau_smpl.sampling_points .* pole) ./ (1 - sign * exp(-beta * pole))
        gl_from_tau = fit(tau_smpl, gtau)

        matsu_smpl = MatsubaraSampling(basis)
        giv = 1 ./ ((im * π / beta) .* matsu_smpl.sampling_points .- pole)
        gl_from_matsu = fit(matsu_smpl, giv)

        @test isapprox(gl_from_tau, gl_from_matsu;
                       atol=1e-10 * maximum(abs.(gl_from_matsu)), rtol=0)
    end
end