text:
using Gen
using FunctionalCollections: PersistentVector
import Distributions

####################
# generative model #
####################

function x_mean(x_prev::Real, t::Int)
    (x_prev / 2.) + 25 * (x_prev / (1 + x_prev * x_prev)) + 8 * cos(1.2 * t)
end

y_mean(x::Real) = (x * x / 20.)

@gen function hmm(var_x, var_y, T::Int)
    xs = Vector{Float64}(undef, T)
    ys = Vector{Float64}(undef, T)
    xs[1] = @trace(normal(0, 5), :x => 1)
    ys[1] = @trace(normal(y_mean(xs[1]), sqrt(var_y)), :y => 1)
    for t=2:T
        xs[t] = @trace(normal(x_mean(xs[t-1], t), sqrt(var_x)), :x => t)
        ys[t] = @trace(normal(y_mean(xs[t]), sqrt(var_y)), :y => t)
    end
    return (xs, ys)
 end

struct State
    x::Float64
    y::Float64
end

struct Params
    var_x::Float64
    var_y::Float64
end

# the kernel accepts two special arguments (the time step t and the return
# value of the previous tieration, followed by the rest of the arguments,
# specified as an argument to unfold)
@gen (static) function kernel(t::Int, state::State, params::Params)
    x::Float64 = @trace(normal(t > 1 ? x_mean(state.x, t) : 0.,
                              t > 1 ? sqrt(params.var_x) : 5.), :x)
    y::Float64 = @trace(normal(y_mean(x), sqrt(params.var_y)), :y)
    ret::State = State(x, y)
    return ret
end

# arguments: (T::Int, init_state, common)
# i.e. the kernel needs to accept three arguments..
hmm2 = Unfold(kernel)

#################################
# sequential Monte Carlo in hmm #
#################################

#const num_particles = 200
#const ess = 100

# implement of SMC using forward simulation as the proposal
# it's okay to be specialized to the model 'hmm'

@gen function single_step_observer(t::Int, y::Float64)
    @trace(dirac(y), :y => t)
end

@gen (static) function obs_sub(y::Float64)
    @trace(dirac(y), :y)
end

single_step_observer2 = at_dynamic(obs_sub, Int)

#@gen (static) function single_step_observer2(t::Int, y::Float64)
    #@trace(at_dynamic(obs_sub,Int)(t, (y,)))
    #@trace(dirac(y), t => :y)
#end

function logsumexp(arr)
    min_arr = maximum(arr)
    min_arr + log(sum(exp.(arr .- min_arr)))
end

function effective_sample_size(log_weights::Vector{Float64})
    # assumes weights are normalized
    log_ess = -logsumexp(2. * log_weights)
    exp(log_ess)
end

# TODO need to implement generate and extend for the unfold generator type
# it should have a custom argchange argument that is indicates whether only the
# time length argument has changed, which will allow extend to only visit new
# time steps

function smc(var_x, var_y, T::Int, N, ess_threshold, ys::AbstractArray{Float64,1})
    log_unnormalized_weights = Vector{Float64}(undef, N)
    log_ml_estimate = 0.
    obs = get_choices(simulate(single_step_observer2, (1, (ys[1],))))
    traces = Vector{Gen.get_trace_type(hmm2)}(undef, N)
    next_traces = Vector{Gen.get_trace_type(hmm2)}(undef, N)
    args = (1, State(NaN, NaN), Params(var_x, var_y))
    for i=1:N
        (traces[i], log_unnormalized_weights[i]) = generate(hmm2, args, obs)
    end
    num_resamples = 0
    args_change = UnfoldCustomArgDiff(true, false, false)
    for t=2:T
        #println(t)
        log_total_weight = logsumexp(log_unnormalized_weights)
        log_normalized_weights = log_unnormalized_weights .- log_total_weight
        if effective_sample_size(log_normalized_weights) < ess_threshold
            weights = exp.(log_normalized_weights)
            parents = rand(Distributions.Categorical(weights / sum(weights)), N)
            #parents = map((i) -> categorical(weights / sum(weights)), 1:N)
            log_ml_estimate += log_total_weight - log(N)
            log_unnormalized_weights = zeros(N)
            num_resamples += 1
        else
            parents = 1:N
        end
        obs = get_choices(simulate(single_step_observer2, (t, (ys[t],))))
        args = (t, State(NaN, NaN), Params(var_x, var_y))
        for i=1:N
            parent = parents[i]
            (next_traces[i], weight) = extend(trace[i], hmm2, args, args_change, obs)
            log_unnormalized_weights[i] += weight
        end
        tmp = traces
        traces = next_traces
        next_traces = tmp
    end
    log_total_weight = logsumexp(log_unnormalized_weights)
    log_normalized_weights = log_unnormalized_weights .- log_total_weight
    log_ml_estimate += log_total_weight - log(N)
    log_ml_estimate # just return the log ML estimate
end

##########################################
# collapsed generative model (handcoded) #
##########################################

using Gen: VectorDistTrace, VectorDistTraceChoiceMap
import Gen: get_call_record, has_choices, get_choices, simulate, assess

struct CollapsedHMMTrace
    vector::VectorDistTrace{Float64}
end

get_call_record(trace::CollapsedHMMTrace) = trace.vector.call
has_choices(trace::CollapsedHMMTrace) = length(trace.vector.call.retval) > 0
get_choices(trace::CollapsedHMMTrace) = CollapsedHMMChoiceMap(get_choices(trace.vector))

struct CollapsedHMMChoiceMap <: ChoiceMap
    y_assignment::VectorDistTraceChoiceMap
end

# addrs are: :y => i

Base.isempty(assignment::CollapsedHMMChoiceMap) = (length(assignment.trace.call.retval) > 0)
Gen.get_address_schema(::Type{CollapsedHMMChoiceMap}) = DynamicAddressSchema()
Gen.has_internal_node(assignment::CollapsedHMMChoiceMap, addr) = (addr == :y)
function Gen.get_internal_node(assignment::CollapsedHMMChoiceMap, addr)
    addr == :y ? assignment.y_assignment : throw(KeyError(addr))
end
Gen.get_internal_nodes(assignment::CollapsedHMMChoiceMap) = ((:y, assignment.y_assignment),)

struct CollapsedHMM <: GenerativeFunction{PersistentVector{Float64},CollapsedHMMTrace} end
collapsed_hmm = CollapsedHMM()

function unbiased_logpdf_est(args, ys::PersistentVector{Float64})
    (var_x, var_y, T, num_particles, ess) = args
	local lml_estimate::Float64
	lml_estimate = smc(var_x, var_y, T, num_particles, ess, ys)
    retval = PersistentVector{Float64}(ys)
    vector = VectorDistTrace(retval, args, lml_estimate, length(ys))
    CollapsedHMMTrace(vector)
end

function Gen.generate(generator::CollapsedHMM, args, constraints)
    (var_x, var_y, T, num_particles, ess) = args
    if isempty(constraints)
        error("Unsupported constraints")
    end
    ys = PersistentVector{Float64}(get_leaf_node(constraints, :y => t) for t=1:T)
    if var_x < 0 || var_y < 0
        retval = PersistentVector{Float64}(ys)
        vector = VectorDistTrace(retval, args, -Inf, length(ys))
        trace = CollapsedHMMTrace(vector)
        weight = -Inf
    else
        trace = unbiased_logpdf_est(args, ys)
        weight = get_call_record(trace).score
    end
    (trace, weight)
end

function Gen.update(generator::CollapsedHMM, new_args, args_change, trace, constraints)
    (var_x, var_y, T, num_particles, ess) = new_args
    if !isempty(constraints)
        error("Unsupported constraints")
    end
    ys = trace.vector.values
    if var_x < 0 || var_y < 0
        retval = get_call_record(trace).retval
        T != length(retval) && error("cannot change length")
        vector = VectorDistTrace(retval, new_args, -Inf, length(retval))
        new_trace = CollapsedHMMTrace(vector)
        weight = -Inf
    else
        new_trace = unbiased_logpdf_est(new_args, ys)
        weight = get_call_record(new_trace).score - get_call_record(trace).score
    end
    (new_trace, weight, EmptyChoiceMap(), nothing)
end

@gen (static) function model_collapsed(T::Int)
    var_x::Float64 = @trace(gamma(1, 1), :var_x)
    var_y::Float64 = @trace(gamma(1, 1), :var_y)
    @trace(collapsed_hmm(var_x, var_y, T, 4096, 2048), :hmm)
end

max_stars_repo_name: georgematheos/Gen

max_issues_repo_name: georgematheos/Gen

max_forks_repo_name: georgematheos/Gen


text:
module TintiNet

    #####
    # 0. Package Dependencies
    #####

    using CUDA, Flux, Transformers, FASTX, DelimitedFiles, JSON, BSON, Random, Dates, Printf;
    using IterTools: ncycle;
    using Transformers.Basic, Transformers.Pretrain, Transformers.Stacks, Transformers.Datasets, Transformers.BidirectionalEncoder;
    using BSON: @load, @save;
    using Flux: @functor, onehot, onecold, mse, shuffle, sigmoid, gradient, unsqueeze;
    using Flux.Optimise: update!;
    using Statistics: mean, std;

    #####
    # 1. Source Files
    #####

    include("./tinti_io.jl")        # Functions to deal with reading/writing from/to files
    include("./definitions.jl")     # Stactic dictionaries
    include("./preprocessing.jl")   # Functions to prepare data for training/evaluation
    include("./architectures.jl")   # Layer architecture definitions
    include("./training.jl")        # Functions for training the neural network
    include("./evaluation.jl")      # Functions for evaluating samples using a model

    #####
    # 2. Namespace export
    #####

    # From tinti_io.jl
    export read_sequences_from_file,
        write_csv_predictions

    # From definitions.jl
    export fastaCategories,
        fastaVocab,
        ss3Categories,
        ss3Vocab,
        ss8Categories,
        ss8Vocab,
        rsa_denominator_dictionary

    # From preprocessing.jl
    export pad_aa_seq,
        preprocess,
        prepare_mask,
        split_train_test

    # From architectures.jl
    export Inception4,
        Inception8

    # From training.jl
    export generic_smooth,
        train_loss,
        eval_loss,
        accuracy_hits

    # From evaluation.jl
    export network_pass_single_GPU,
        network_pass_single_CPU,
        network_pass_batched_GPU,
        network_pass_batched_CPU,
        compute_classifier_predictions,
        compute_regressor_predictions

end

ext: jl


text:
function computus(year; servois=false)

    # Year's position on the 19 year metonic cycle
    a = mod(year, 19)

    # Century index
    k = fld(year, 100)

    # Shift of metonic cycle, add a day offset every 300 years
    p = fld(13 + 8 * k, 25)

    # Correction for non-observed leap days
    q = fld(k, 4)

    # Correction to starting point of calculation each century
    M = mod(15 - p + k - q, 30)

    # Number of days from March 21st until the full moon
    d = mod(19 * a + M, 30)

    # Returning if user wants value for Servois' table
    if servois
        return string(mod(21 + d,31))
    end

    # Finding the next Sunday
    # Century-based offset in weekly calculation
    N = mod(4 + k - q, 7)

    # Correction for leap days
    b = mod(year, 4)
    c = mod(year, 7)

    # Days from d to next Sunday
    e = mod(2 * b + 4 * c + 6 * d + N, 7)

    # Historical corrections for April 26 and 25
    if (d == 29 && e == 6) || (d == 28 && e == 6 && a > 10)
        e = -1
    end

    # Determination of the correct month for Easter
    if(22 + d + e > 31)
        return "April " * string(d + e - 9)
    else
        return "March " * string(22 + d + e)
    end
end

# Here, we will output the date of the Paschal full moon
# (using Servois notation), and Easter for 2020-2030

a = collect(2020:2030)
servois_numbers = computus.(a; servois=true)
easter_dates = computus.(a)

println("The following are the dates of the Paschal full moon (using Servois " *
        "notation) and the date of Easter for 2020-2030 AD:")
println("Year\tServois number\tEaster")
for i = 1:length(a)
    println("$(a[i])\t$(servois_numbers[i])\t\t$(easter_dates[i])")
end


text:
using SparseArrays
export choi2so, so2choi

"""
    permutedims_matrix(dims, p)

Returns a matrix `M` so that for any vector `v` of length `prod(dims)`,

    M*v == vec(permutedims(reshape(v, dims), p))

"""
function permutedims_matrix(dims, p)
    d = prod(dims)
    n = length(dims)
    sparse(reshape(PermutedDimsArray(reshape(I(d), (dims..., dims...)), (p..., (n + 1:2n)...)), (d, d)))
end

struct Choi2SOAtom <: Convex.AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Tuple{Convex.AbstractExpr}
    size::Tuple{Int,Int}
    dA::Int
    dB::Int
    function Choi2SOAtom(x::Convex.AbstractExpr, dA::Int, dB::Int)
        if dA^2 * dB^2 != length(x)
            error("Cannot Choi2SOAtom expression of size $(x.size) with dimensions ($(dA), $(dB))")
        end
        return new(:choi2so, objectid(x), (x,), (dB^2, dA^2), dA, dB)
    end
end

function Convex.sign(x::Choi2SOAtom)
    return Convex.sign(x.children[1])
end

function Convex.monotonicity(x::Choi2SOAtom)
    return (Convex.Nondecreasing(),)
end

function Convex.curvature(x::Choi2SOAtom)
    return Convex.ConstVexity()
end

function Convex.evaluate(x::Choi2SOAtom)
    J = Convex.evaluate(x.children[1])
    return choi2so(J, x.dA, x.dB)
end

function Convex.conic_form!(x::Choi2SOAtom, unique_conic_forms::Convex.UniqueConicForms = Convex.UniqueConicForms())
    if !Convex.has_conic_form(unique_conic_forms, x)
        objective = Convex.conic_form!(x.children[1], unique_conic_forms)
        dA = x.dA; dB = x.dB;
        M = permutedims_matrix((dA, dB, dA, dB), (2, 4, 1, 3))
        objective = M * objective
        Convex.cache_conic_form!(unique_conic_forms, x, objective)
    end

    return Convex.get_conic_form(unique_conic_forms, x)
end

"""
    choi2so(x::Convex.AbstractExpr, dA::Int, dB::Int)

Convert from an (unnormalized) Choi matrix of a quantum channel `T` to a matrix representation.
It is assumed `T` is a quantum channel from `A` to `B`, where `A` is `dA`-dimensional space and `B` a `dB`-dimensional space.
"""
choi2so(x::Convex.AbstractExpr, dA::Int, dB::Int) = Choi2SOAtom(x, dA, dB)
choi2so(J, dA, dB) = reshape(PermutedDimsArray(reshape(J, (dA, dB, dA, dB)), (2, 4, 1, 3)), (dB^2, dA^2))

struct SO2ChoiAtom <: Convex.AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Tuple{Convex.AbstractExpr}
    size::Tuple{Int,Int}
    dA::Int
    dB::Int
    function SO2ChoiAtom(x::Convex.AbstractExpr, dA::Int, dB::Int)
        if !(dA^2 * dB^2 == length(x) && size(x) == (dB^2, dA^2))
            error("Cannot Choi2SOAtom expression of size $(x.size) with dimensions ($(dA), $(dB))")
        end
        return new(:so2choi, objectid(x), (x,), (dA*dB, dA*dB), dA, dB)
    end
end

function Convex.sign(x::SO2ChoiAtom)
    return Convex.sign(x.children[1])
end

function Convex.monotonicity(x::SO2ChoiAtom)
    return (Convex.Nondecreasing(),)
end

function Convex.curvature(x::SO2ChoiAtom)
    return Convex.ConstVexity()
end

function Convex.evaluate(x::SO2ChoiAtom)
    T = Convex.evaluate(x.children[1])
    return so2choi(T, x.dA, x.dB)
end

function Convex.conic_form!(x::SO2ChoiAtom, unique_conic_forms::Convex.UniqueConicForms = Convex.UniqueConicForms())
    if !Convex.has_conic_form(unique_conic_forms, x)
        objective = Convex.conic_form!(x.children[1], unique_conic_forms)
        dA = x.dA; dB = x.dB;
        M = permutedims_matrix((dB, dB, dA, dA), (3, 1, 4, 2))
        objective = M * objective
        Convex.cache_conic_form!(unique_conic_forms, x, objective)
    end

    return Convex.get_conic_form(unique_conic_forms, x)
end

"""
    so2choi(x::Convex.AbstractExpr, dA::Int, dB::Int)

Convert from a matrix representation of a quantum channel `T` to the Choi representation.
It is assumed `T` is a quantum channel from `A` to `B`, where `A` is `dA`-dimensional space and `B` a `dB`-dimensional space.
"""
so2choi(x::Convex.AbstractExpr, dA::Int, dB::Int) = SO2Choi(x, dA, dB)
so2choi(T, dA, dB) = reshape(PermutedDimsArray(reshape(T, (dB, dB, dA, dA)), (3, 1, 4, 2)), (dA*dB, dA*dB))


text:
######################################################################
# Non-static enums which can be expanded at runtime. Once a value is
# defined it cannot be undefined.
# -----
# Licensed under MIT License
module RuntimeEnums
import Base.Enums

export RuntimeEnum
abstract type RuntimeEnum{I<:Integer} <: Enum{I} end

export @runtime_enum
macro runtime_enum(T::Union{Symbol, Expr}, exprs...)
    typename, basetype = extract_runtime_enum_type(__module__, T)
    esctypename = esc(typename)

    namemap = Dict{basetype, Symbol}()

    blk = quote
        Base.@__doc__(primitive type $esctypename <: RuntimeEnum{$basetype} $(sizeof(basetype) * 8) end)
        function ($esctypename)(x::Integer)
            if x ∉ keys(Enums.namemap($esctypename))
                throw(ArgumentError("invalid value for runtime enum $(string($typename)): $x"))
            end
            return Base.bitcast($esctypename, convert($basetype, x))
        end

        Base.Enums.namemap(::Type{$esctypename}) = $namemap
    end

    if length(exprs) > 0
        append!(blk.args, make_constants(typename, namemap, exprs))
    end

    push!(blk.args, nothing)
    blk
end

macro register(E::Symbol, exprs::Union{Symbol, Expr}...)
    if length(exprs) == 0
        throw(ArgumentError("no values provided"))
    end

    namemap = Core.eval(__module__, :(Base.Enums.namemap($E)))
    E = esc(E)

    Expr(:block, make_constants(E, namemap, exprs)...)
end

function Base.setindex!(E::Type{<:RuntimeEnum}, value::Integer, name::Symbol)
    namemap = Enums.namemap(E)
    if name ∈ values(namemap)
        throw(ArgumentError("runtime enum name $E::$name already in use"))
    end
    if value ∈ keys(namemap)
        throw(ArgumentError("runtime enum value $E::$(namemap[value]) = $value already in use"))
    end

    Enums.namemap(E)[value] = name
    return E(value)
end

Base.typemin(E::Type{<:RuntimeEnum}) = (assert_nonempty(E); E(minkey(Enums.namemap(E))))
Base.typemax(E::Type{<:RuntimeEnum}) = (assert_nonempty(E); E(maxkey(Enums.namemap(E))))
Base.instances(E::Type{<:RuntimeEnum}) = E.(keys(Enums.namemap(E)))
Base.isempty(E::Type{<:RuntimeEnum}) = isempty(Enums.namemap(E))

maxkey(keyset::Base.KeySet) = reduce((a, b) -> max(a, b), keyset)
minkey(keyset::Base.KeySet) = reduce((a, b) -> min(a, b), keyset)
maxkey(dict::Dict) = maxkey(keys(dict))
minkey(dict::Dict) = minkey(keys(dict))

function assert_nonempty(E::Type{<:RuntimeEnum})
    @assert(!isempty(keys(Enums.namemap(E))), "empty runtime enum")
end

extract_runtime_enum_type(_, name::Symbol) = name, Int32
function extract_runtime_enum_type(mod::Module, T::Expr)
    if !(T.head === :(::) && length(T.args) == 2 && T.args[1] isa Symbol)
        throw(ArgumentError("invalid type expression for runtime enum $T"))
    end

    typename = T.args[1]
    basetype = Core.eval(mod, T.args[2])

    if !(basetype isa DataType && basetype <: Integer && isbitstype(basetype))
        throw(ArgumentError("invalid base type for runtime enum $typename, $T=::$basetype; base type must be an integer primitive type"))
    end

    return typename, basetype
end

normalize_value_expr(namemap::Dict{I, Symbol} where {I<:Integer}, name::Symbol) = (name, isempty(keys(namemap)) ? 0 : maxkey(namemap)+1)
function normalize_value_expr(_, expr::Expr)
    if !(expr.head === :(=) && expr.args[1] isa Symbol && expr.args[2] isa Integer)
        throw(ArgumentError("invalid runtime enum value expression: $expr"))
    end
    return expr.args[1], expr.args[2]
end

function collect_value_exprs(exprs)
    # Collects all valid value expressions & preserves line numbers.
    result = []
    for expr in exprs
        if expr isa Symbol
            push!(result, expr)
            continue
        elseif expr isa Expr
            if expr.head === :block
                append!(result, collect_value_exprs(expr.args))
            elseif expr.head === :(=)
                if expr.args[1] isa Symbol && expr.args[2] isa Integer
                    push!(result, expr)
                    continue
                end
            end
            continue
        elseif expr isa LineNumberNode
            push!(result, expr)
        end

        throw(ArgumentError("invalid runtime enum value expression $expr"))
    end

    return result
end

function make_constants(E::Symbol, namemap, exprs)
    result = []
    for expr ∈ exprs
        if expr isa LineNumberNode
            push!(result, expr)
        elseif expr isa Expr && expr.head === :block
            append!(result, make_constants(E, namemap, expr.args))
        else
            varname, value = normalize_value_expr(namemap, expr)
            if varname ∈ values(namemap)
                throw(ArgumentError("duplicate runtime enum name $varname"))
            elseif value ∈ keys(namemap)
                throw(ArgumentError("duplicate runtime enum value $value"))
            end

            namemap[value] = varname
            push!(result, :(const $(esc(varname)) = $(esc(E))($value)))
        end
    end
    return result
end

end # module RuntimeEnums


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

######################################################################################
# Upper-Hessenberg matrices H+μI, analogous to the UpperTriangular type

"""
    UpperHessenberg(A::AbstractMatrix)

Construct an `UpperHessenberg` view of the matrix `A`.
Entries of `A` below the first subdiagonal are ignored.

Efficient algorithms are implemented for `H \\ b`, `det(H)`, and similar.

See also the [`hessenberg`](@ref) function to factor any matrix into a similar
upper-Hessenberg matrix.

If `F::Hessenberg` is the factorization object, the unitary matrix can be accessed
with `F.Q` and the Hessenberg matrix with `F.H`. When `Q` is extracted, the resulting
type is the `HessenbergQ` object, and may be converted to a regular matrix with
[`convert(Array, _)`](@ref) (or `Array(_)` for short).

Iterating the decomposition produces the factors `F.Q` and `F.H`.

# Examples
```jldoctest
julia> A = [1 2 3 4; 5 6 7 8; 9 10 11 12; 13 14 15 16]
4×4 Matrix{Int64}:
  1   2   3   4
  5   6   7   8
  9  10  11  12
 13  14  15  16

julia> UpperHessenberg(A)
4×4 UpperHessenberg{Int64, Matrix{Int64}}:
 1   2   3   4
 5   6   7   8
 ⋅  10  11  12
 ⋅   ⋅  15  16
```
"""
struct UpperHessenberg{T,S<:AbstractMatrix{T}} <: AbstractMatrix{T}
    data::S

    function UpperHessenberg{T,S}(data) where {T,S<:AbstractMatrix{T}}
        require_one_based_indexing(data)
        new{T,S}(data)
    end
end
UpperHessenberg(H::UpperHessenberg) = H
UpperHessenberg{T}(A::AbstractMatrix) where {T} = UpperHessenberg(convert(AbstractMatrix{T}, A))
UpperHessenberg{T}(H::UpperHessenberg) where {T} = UpperHessenberg{T}(H.data)
UpperHessenberg(A::AbstractMatrix) = UpperHessenberg{eltype(A),typeof(A)}(A)
Matrix(H::UpperHessenberg{T}) where {T} = Matrix{T}(H)
Array(H::UpperHessenberg) = Matrix(H)
size(H::UpperHessenberg, d) = size(H.data, d)
size(H::UpperHessenberg) = size(H.data)
parent(H::UpperHessenberg) = H.data

# similar behaves like UpperTriangular
similar(H::UpperHessenberg, ::Type{T}) where {T} = UpperHessenberg(similar(H.data, T))
similar(H::UpperHessenberg, ::Type{T}, dims::Dims{N}) where {T,N} = similar(H.data, T, dims)

AbstractMatrix{T}(H::UpperHessenberg) where {T} = UpperHessenberg(AbstractMatrix{T}(H.data))

copy(H::UpperHessenberg) = UpperHessenberg(copy(H.data))
real(H::UpperHessenberg{<:Real}) = H
real(H::UpperHessenberg{<:Complex}) = UpperHessenberg(triu!(real(H.data),-1))
imag(H::UpperHessenberg) = UpperHessenberg(triu!(imag(H.data),-1))

function Matrix{T}(H::UpperHessenberg) where T
    m,n = size(H)
    return triu!(copyto!(Matrix{T}(undef, m, n), H.data), -1)
end

getindex(H::UpperHessenberg{T}, i::Integer, j::Integer) where {T} =
    i <= j+1 ? convert(T, H.data[i,j]) : zero(T)

function setindex!(A::UpperHessenberg, x, i::Integer, j::Integer)
    if i > j+1
        x == 0 || throw(ArgumentError("cannot set index in the lower triangular part " *
            "($i, $j) of an UpperHessenberg matrix to a nonzero value ($x)"))
    else
        A.data[i,j] = x
    end
    return A
end

function Base.replace_in_print_matrix(A::UpperHessenberg, i::Integer, j::Integer, s::AbstractString)
    return i <= j+1 ? s : Base.replace_with_centered_mark(s)
end

Base.copy(A::Adjoint{<:Any,<:UpperHessenberg}) = tril!(adjoint!(similar(A.parent.data), A.parent.data), 1)
Base.copy(A::Transpose{<:Any,<:UpperHessenberg}) = tril!(transpose!(similar(A.parent.data), A.parent.data), 1)

-(A::UpperHessenberg) = UpperHessenberg(-A.data)
rmul!(H::UpperHessenberg, x::Number) = (rmul!(H.data, x); H)
lmul!(x::Number, H::UpperHessenberg) = (lmul!(x, H.data); H)

fillstored!(H::UpperHessenberg, x) = (fillband!(H.data, x, -1, size(H,2)-1); H)

+(A::UpperHessenberg, B::UpperHessenberg) = UpperHessenberg(A.data+B.data)
-(A::UpperHessenberg, B::UpperHessenberg) = UpperHessenberg(A.data-B.data)

for T = (:UniformScaling, :Diagonal, :Bidiagonal, :Tridiagonal, :SymTridiagonal,
         :UpperTriangular, :UnitUpperTriangular)
    for op = (:+, :-)
        @eval begin
            $op(H::UpperHessenberg, x::$T) = UpperHessenberg($op(H.data, x))
            $op(x::$T, H::UpperHessenberg) = UpperHessenberg($op(x, H.data))
        end
    end
end

for T = (:Number, :UniformScaling, :Diagonal)
    @eval begin
        *(H::UpperHessenberg, x::$T) = UpperHessenberg(H.data * x)
        *(x::$T, H::UpperHessenberg) = UpperHessenberg(x * H.data)
        /(H::UpperHessenberg, x::$T) = UpperHessenberg(H.data / x)
        \(x::$T, H::UpperHessenberg) = UpperHessenberg(x \ H.data)
    end
end

function *(H::UpperHessenberg, U::UpperOrUnitUpperTriangular)
    T = typeof(oneunit(eltype(H))*oneunit(eltype(U)))
    HH = similar(H.data, T, size(H))
    copyto!(HH, H)
    rmul!(HH, U)
    UpperHessenberg(HH)
end
function *(U::UpperOrUnitUpperTriangular, H::UpperHessenberg)
    T = typeof(oneunit(eltype(H))*oneunit(eltype(U)))
    HH = similar(H.data, T, size(H))
    copyto!(HH, H)
    lmul!(U, HH)
    UpperHessenberg(HH)
end

function /(H::UpperHessenberg, U::UpperTriangular)
    T = typeof(oneunit(eltype(H))/oneunit(eltype(U)))
    HH = similar(H.data, T, size(H))
    copyto!(HH, H)
    rdiv!(HH, U)
    UpperHessenberg(HH)
end
function /(H::UpperHessenberg, U::UnitUpperTriangular)
    T = typeof(oneunit(eltype(H))/oneunit(eltype(U)))
    HH = similar(H.data, T, size(H))
    copyto!(HH, H)
    rdiv!(HH, U)
    UpperHessenberg(HH)
end

function \(U::UpperTriangular, H::UpperHessenberg)
    T = typeof(oneunit(eltype(U))\oneunit(eltype(H)))
    HH = similar(H.data, T, size(H))
    copyto!(HH, H)
    ldiv!(U, HH)
    UpperHessenberg(HH)
end
function \(U::UnitUpperTriangular, H::UpperHessenberg)
    T = typeof(oneunit(eltype(U))\oneunit(eltype(H)))
    HH = similar(H.data, T, size(H))
    copyto!(HH, H)
    ldiv!(U, HH)
    UpperHessenberg(HH)
end

function *(H::UpperHessenberg, B::Bidiagonal)
    TS = promote_op(matprod, eltype(H), eltype(B))
    if B.uplo == 'U'
        A_mul_B_td!(UpperHessenberg(zeros(TS, size(H)...)), H, B)
    else
        A_mul_B_td!(zeros(TS, size(H)...), H, B)
    end
end
function *(B::Bidiagonal, H::UpperHessenberg)
    TS = promote_op(matprod, eltype(B), eltype(H))
    if B.uplo == 'U'
        A_mul_B_td!(UpperHessenberg(zeros(TS, size(B)...)), B, H)
    else
        A_mul_B_td!(zeros(TS, size(B)...), B, H)
    end
end

function /(H::UpperHessenberg, B::Bidiagonal)
    A = Base.@invoke /(H::AbstractMatrix, B::Bidiagonal)
    B.uplo == 'U' ? UpperHessenberg(A) : A
end

# Solving (H+µI)x = b: we can do this in O(m²) time and O(m) memory
# (in-place in x) by the RQ algorithm from:
#
#    G. Henry, "The shifted Hessenberg system solve computation," Tech. Rep. 94–163,
#    Center for Appl. Math., Cornell University (1994).
#
# as reviewed in
#
#    C. Beattie et al., "A note on shifted Hessenberg systems and frequency
#    response computation," ACM Trans. Math. Soft. 38, pp. 12:6–12:16 (2011)
#
# (Note, however, that there is apparently a typo in Algorithm 1 of the
#  Beattie paper: the Givens rotation uses u(k), not H(k,k) - σ.)
#
# Essentially, it works by doing a Givens RQ factorization of H+µI from
# right to left, and doing backsubstitution *simultaneously*.

# solve (H+μI)X = B, storing result in B
function ldiv!(F::UpperHessenberg, B::AbstractVecOrMat; shift::Number=false)
    checksquare(F)
    m = size(F,1)
    m != size(B,1) && throw(DimensionMismatch("wrong right-hand-side # rows != $m"))
    require_one_based_indexing(B)
    n = size(B,2)
    H = F.data
    μ = shift
    u = Vector{typeof(zero(eltype(H))+μ)}(undef, m) # for last rotated col of H-μI
    copyto!(u, 1, H, m*(m-1)+1, m) # u .= H[:,m]
    u[m] += μ
    X = B # not a copy, just rename to match paper
    cs = Vector{Tuple{real(eltype(u)),eltype(u)}}(undef, length(u)) # store Givens rotations
    @inbounds for k = m:-1:2
        c, s, ρ = givensAlgorithm(u[k], H[k,k-1])
        cs[k] = (c, s)
        for i = 1:n
            X[k,i] /= ρ
            t₁ = s * X[k,i]; t₂ = c * X[k,i]
            @simd for j = 1:k-2
                X[j,i] -= u[j]*t₂ + H[j,k-1]*t₁
            end
            X[k-1,i] -= u[k-1]*t₂ + (H[k-1,k-1] + μ) * t₁
        end
        @simd for j = 1:k-2
            u[j] = H[j,k-1]*c - u[j]*s'
        end
        u[k-1] = (H[k-1,k-1] + μ) * c - u[k-1]*s'
    end
    for i = 1:n
        τ₁ = X[1,i] / u[1]
        @inbounds for j = 2:m
            τ₂ = X[j,i]
            c, s = cs[j]
            X[j-1,i] = c*τ₁ + s*τ₂
            τ₁ = c*τ₂ - s'τ₁
        end
        X[m,i] = τ₁
    end
    return X
end

# solve X(H+μI) = B, storing result in B
#
# Note: this can be derived from the Henry (1994) algorithm
# by transformation to F(Hᵀ+µI)F FXᵀ = FBᵀ, where
# F is the permutation matrix that reverses the order
# of rows/cols.  Essentially, we take the ldiv! algorithm,
# swap indices of H and X to transpose, and reverse the
# order of the H indices (or the order of the loops).
function rdiv!(B::AbstractMatrix, F::UpperHessenberg; shift::Number=false)
    checksquare(F)
    m = size(F,1)
    m != size(B,2) && throw(DimensionMismatch("wrong right-hand-side # cols != $m"))
    require_one_based_indexing(B)
    n = size(B,1)
    H = F.data
    μ = shift
    u = Vector{typeof(zero(eltype(H))+μ)}(undef, m) # for last rotated row of H-μI
    u .= @view H[1,:]
    u[1] += μ
    X = B # not a copy, just rename to match paper
    cs = Vector{Tuple{real(eltype(u)),eltype(u)}}(undef, length(u)) # store Givens rotations
    @inbounds for k = 1:m-1
        c, s, ρ = givensAlgorithm(u[k], H[k+1,k])
        cs[k] = (c, s)
        for i = 1:n
            X[i,k] /= ρ
            t₁ = s * X[i,k]; t₂ = c * X[i,k]
            @simd for j = k+2:m
                X[i,j] -= u[j]*t₂ + H[k+1,j]*t₁
            end
            X[i,k+1] -= u[k+1]*t₂ + (H[k+1,k+1] + μ) * t₁
        end
        @simd for j = k+2:m
            u[j] = H[k+1,j]*c - u[j]*s'
        end
        u[k+1] = (H[k+1,k+1] + μ) * c - u[k+1]*s'
    end
    for i = 1:n
        τ₁ = X[i,m] / u[m]
        @inbounds for j = m-1:-1:1
            τ₂ = X[i,j]
            c, s = cs[j]
            X[i,j+1] = c*τ₁ + s*τ₂
            τ₁ = c*τ₂ - s'τ₁
        end
        X[i,1] = τ₁
    end
    return X
end

# Hessenberg-matrix determinant formula for H+μI based on:
#
#    N. D. Cahill, J. R. D’Errico, D. A. Narayan, and J. Y. Narayan, "Fibonacci determinants,"
#    College Math. J. 33, pp. 221-225 (2003).
#
# as reviewed in Theorem 2.1 of:
#
#    K. Kaygisiz and A. Sahin, "Determinant and permanent of Hessenberg matrix and generalized Lucas polynomials,"
#    arXiv:1111.4067 (2011).
#
# Cost is O(m²) with O(m) storage.
function det(F::UpperHessenberg; shift::Number=false)
    checksquare(F)
    H = F.data
    m = size(H,1)
    μ = shift
    m == 0 && return one(zero(eltype(H)) + μ)
    determinant = H[1,1] + μ
    prevdeterminant = one(determinant)
    m == 1 && return determinant
    prods = Vector{typeof(determinant)}(undef, m-1) # temporary storage for partial products
    @inbounds for n = 2:m
        prods[n-1] = prevdeterminant
        prevdeterminant = determinant
        determinant *= H[n,n] + μ
        h = H[n,n-1]
        @simd for r = n-1:-2:2
            determinant -= H[r,n] * (prods[r] *= h) - H[r-1,n] * (prods[r-1] *= h)
        end
        if iseven(n)
            determinant -= H[1,n] * (prods[1] *= h)
        end
    end
    return determinant
end

# O(m²) log-determinant based on first doing Givens RQ to put H+μI into upper-triangular form and then
# taking the product of the diagonal entries.   The trick is that we only need O(m) temporary storage,
# because we don't need to store the whole Givens-rotated matrix, only the most recent column.
# We do RQ (column rotations) rather than QR (row rotations) for more consecutive memory access.
# (We could also use it for det instead of the Cahill algorithm above.  Cahill is slightly faster
#  for very small matrices where you are likely to use det, and also uses only ± and * so it can
#  be applied to Hessenberg matrices over other number fields.)
function logabsdet(F::UpperHessenberg; shift::Number=false)
    checksquare(F)
    H = F.data
    m = size(H,1)
    μ = shift
    P = one(zero(eltype(H)) + μ)
    logdeterminant = zero(real(P))
    m == 0 && return (logdeterminant, P)
    g = Vector{typeof(P)}(undef, m) # below, g is the k-th col of Givens-rotated H+μI matrix
    copyto!(g, 1, H, m*(m-1)+1, m) # g .= H[:,m]
    g[m] += μ
    @inbounds for k = m:-1:2
        c, s, ρ = givensAlgorithm(g[k], H[k,k-1])
        logdeterminant += log(abs(ρ))
        P *= sign(ρ)
        g[k-1] = c*(H[k-1,k-1] + μ) - s'*g[k-1]
        @simd for j = 1:k-2
            g[j] = c*H[j,k-1] - s'*g[j]
        end
    end
    logdeterminant += log(abs(g[1]))
    P *= sign(g[1])
    return (logdeterminant, P)
end

function dot(x::AbstractVector, H::UpperHessenberg, y::AbstractVector)
    require_one_based_indexing(x, y)
    m = size(H, 1)
    (length(x) == m == length(y)) || throw(DimensionMismatch())
    if iszero(m)
        return dot(zero(eltype(x)), zero(eltype(H)), zero(eltype(y)))
    end
    x₁ = x[1]
    r = dot(x₁, H[1,1], y[1])
    r += dot(x[2], H[2,1], y[1])
    @inbounds for j in 2:m-1
        yj = y[j]
        if !iszero(yj)
            temp = adjoint(H[1,j]) * x₁
            @simd for i in 2:j+1
                temp += adjoint(H[i,j]) * x[i]
            end
            r += dot(temp, yj)
        end
    end
    ym = y[m]
    if !iszero(ym)
        temp = adjoint(H[1,m]) * x₁
        @simd for i in 2:m
            temp += adjoint(H[i,m]) * x[i]
        end
        r += dot(temp, ym)
    end
    return r
end

######################################################################################
# Hessenberg factorizations Q(H+μI)Q' of A+μI:

"""
    Hessenberg <: Factorization

A `Hessenberg` object represents the Hessenberg factorization `QHQ'` of a square
matrix, or a shift `Q(H+μI)Q'` thereof, which is produced by the [`hessenberg`](@ref) function.
"""
struct Hessenberg{T,SH<:AbstractMatrix,S<:AbstractMatrix,W<:AbstractVector,V<:Number} <: Factorization{T}
    H::SH # UpperHessenberg or SymTridiagonal
    uplo::Char
    factors::S # reflector data in uplo triangle, may share data with H
    τ::W # more Q (reflector) data
    μ::V # diagonal shift for copy-free (F+μI) \ b solves and similar
end
Hessenberg(factors::AbstractMatrix, τ::AbstractVector, H::AbstractMatrix=UpperHessenberg(factors), uplo::AbstractChar='L'; μ::Number=false) =
    Hessenberg{typeof(zero(eltype(factors))+μ),typeof(H),typeof(factors),typeof(τ),typeof(μ)}(H, uplo, factors, τ, μ)
Hessenberg(F::Hessenberg) = F
Hessenberg(F::Hessenberg, μ::Number) = Hessenberg(F.factors, F.τ, F.H, F.uplo; μ=μ)

copy(F::Hessenberg{<:Any,<:UpperHessenberg}) = Hessenberg(copy(F.factors), copy(F.τ); μ=F.μ)
copy(F::Hessenberg{<:Any,<:SymTridiagonal}) = Hessenberg(copy(F.factors), copy(F.τ), copy(F.H), F.uplo; μ=F.μ)
size(F::Hessenberg, d) = size(F.H, d)
size(F::Hessenberg) = size(F.H)

adjoint(F::Hessenberg) = Adjoint(F)

# iteration for destructuring into components
Base.iterate(S::Hessenberg) = (S.Q, Val(:H))
Base.iterate(S::Hessenberg, ::Val{:H}) = (S.H, Val(:μ))
Base.iterate(S::Hessenberg, ::Val{:μ}) = (S.μ, Val(:done))
Base.iterate(S::Hessenberg, ::Val{:done}) = nothing

hessenberg!(A::StridedMatrix{<:BlasFloat}) = Hessenberg(LAPACK.gehrd!(A)...)

function hessenberg!(A::Union{Symmetric{<:BlasReal},Hermitian{<:BlasFloat}})
    factors, τ, d, e = LAPACK.hetrd!(A.uplo, A.data)
    return Hessenberg(factors, τ, SymTridiagonal(d, e), A.uplo)
end

"""
    hessenberg!(A) -> Hessenberg

`hessenberg!` is the same as [`hessenberg`](@ref), but saves space by overwriting
the input `A`, instead of creating a copy.
"""
hessenberg!(A::AbstractMatrix)

"""
    hessenberg(A) -> Hessenberg

Compute the Hessenberg decomposition of `A` and return a `Hessenberg` object. If `F` is the
factorization object, the unitary matrix can be accessed with `F.Q` (of type `LinearAlgebra.HessenbergQ`)
and the Hessenberg matrix with `F.H` (of type [`UpperHessenberg`](@ref)), either of
which may be converted to a regular matrix with `Matrix(F.H)` or `Matrix(F.Q)`.

If `A` is [`Hermitian`](@ref) or real-[`Symmetric`](@ref), then the Hessenberg
decomposition produces a real-symmetric tridiagonal matrix and `F.H` is of type
[`SymTridiagonal`](@ref).

Note that the shifted factorization `A+μI = Q (H+μI) Q'` can be
constructed efficiently by `F + μ*I` using the [`UniformScaling`](@ref)
object [`I`](@ref), which creates a new `Hessenberg` object with shared storage
and a modified shift.   The shift of a given `F` is obtained by `F.μ`.
This is useful because multiple shifted solves `(F + μ*I) \\ b`
(for different `μ` and/or `b`) can be performed efficiently once `F` is created.

Iterating the decomposition produces the factors `F.Q, F.H, F.μ`.

# Examples
```jldoctest
julia> A = [4. 9. 7.; 4. 4. 1.; 4. 3. 2.]
3×3 Matrix{Float64}:
 4.0  9.0  7.0
 4.0  4.0  1.0
 4.0  3.0  2.0

julia> F = hessenberg(A)
Hessenberg{Float64, UpperHessenberg{Float64, Matrix{Float64}}, Matrix{Float64}, Vector{Float64}, Bool}
Q factor:
3×3 LinearAlgebra.HessenbergQ{Float64, Matrix{Float64}, Vector{Float64}, false}:
 1.0   0.0        0.0
 0.0  -0.707107  -0.707107
 0.0  -0.707107   0.707107
H factor:
3×3 UpperHessenberg{Float64, Matrix{Float64}}:
  4.0      -11.3137       -1.41421
 -5.65685    5.0           2.0
   ⋅        -8.88178e-16   1.0

julia> F.Q * F.H * F.Q'
3×3 Matrix{Float64}:
 4.0  9.0  7.0
 4.0  4.0  1.0
 4.0  3.0  2.0

julia> q, h = F; # destructuring via iteration

julia> q == F.Q && h == F.H
true
```
"""
hessenberg(A::AbstractMatrix{T}) where T =
    hessenberg!(copy_oftype(A, eigtype(T)))

function show(io::IO, mime::MIME"text/plain", F::Hessenberg)
    summary(io, F)
    if !iszero(F.μ)
        print("\nwith shift μI for μ = ", F.μ)
    end
    println(io, "\nQ factor:")
    show(io, mime, F.Q)
    println(io, "\nH factor:")
    show(io, mime, F.H)
end

"""
    HessenbergQ <: AbstractQ

Given a [`Hessenberg`](@ref) factorization object `F`, `F.Q` returns
a `HessenbergQ` object, which is an implicit representation of the unitary
matrix `Q` in the Hessenberg factorization `QHQ'` represented by `F`.
This `F.Q` object can be efficiently multiplied by matrices or vectors,
and can be converted to an ordinary matrix type with `Matrix(F.Q)`.
"""
struct HessenbergQ{T,S<:AbstractMatrix,W<:AbstractVector,sym} <: AbstractQ{T}
    uplo::Char
    factors::S
    τ::W
    function HessenbergQ{T,S,W,sym}(uplo::AbstractChar, factors, τ) where {T,S<:AbstractMatrix,W<:AbstractVector,sym}
        new(uplo, factors, τ)
    end
end
HessenbergQ(F::Hessenberg{<:Any,<:UpperHessenberg,S,W}) where {S,W} = HessenbergQ{eltype(F.factors),S,W,false}(F.uplo, F.factors, F.τ)
HessenbergQ(F::Hessenberg{<:Any,<:SymTridiagonal,S,W}) where {S,W} = HessenbergQ{eltype(F.factors),S,W,true}(F.uplo, F.factors, F.τ)

function getproperty(F::Hessenberg, d::Symbol)
    d === :Q && return HessenbergQ(F)
    return getfield(F, d)
end

Base.propertynames(F::Hessenberg, private::Bool=false) =
    (:Q, :H, :μ, (private ? (:τ, :factors, :uplo) : ())...)

# HessenbergQ from LAPACK/BLAS (as opposed to Julia libraries like GenericLinearAlgebra)
const BlasHessenbergQ{T,sym} = HessenbergQ{T,<:StridedMatrix{T},<:StridedVector{T},sym} where {T<:BlasFloat,sym}

## reconstruct the original matrix
Matrix{T}(Q::BlasHessenbergQ{<:Any,false}) where {T} = convert(Matrix{T}, LAPACK.orghr!(1, size(Q.factors, 1), copy(Q.factors), Q.τ))
Matrix{T}(Q::BlasHessenbergQ{<:Any,true}) where {T} = convert(Matrix{T}, LAPACK.orgtr!(Q.uplo, copy(Q.factors), Q.τ))
AbstractArray(F::Hessenberg) = AbstractMatrix(F)
Matrix(F::Hessenberg) = Array(AbstractArray(F))
Array(F::Hessenberg) = Matrix(F)
function AbstractMatrix(F::Hessenberg)
    Q = F.Q
    A = rmul!(lmul!(Q, Matrix{eltype(Q)}(F.H)), Q')
    μ = F.μ
    if iszero(μ)
        return A
    elseif typeof(zero(eltype(A))+μ) <: eltype(A) # can shift A in-place
        for i = 1:size(A,1)
            @inbounds A[i,i] += μ
        end
        return A
    else
        return A + μ*I # allocate another matrix, e.g. if A is real and μ is complex
    end
end

# adjoint(Q::HessenbergQ{<:Real})

lmul!(Q::BlasHessenbergQ{T,false}, X::StridedVecOrMat{T}) where {T<:BlasFloat} =
    LAPACK.ormhr!('L', 'N', 1, size(Q.factors, 1), Q.factors, Q.τ, X)
rmul!(X::StridedVecOrMat{T}, Q::BlasHessenbergQ{T,false}) where {T<:BlasFloat} =
    LAPACK.ormhr!('R', 'N', 1, size(Q.factors, 1), Q.factors, Q.τ, X)
lmul!(adjQ::Adjoint{<:Any,<:BlasHessenbergQ{T,false}}, X::StridedVecOrMat{T}) where {T<:BlasFloat} =
    (Q = adjQ.parent; LAPACK.ormhr!('L', ifelse(T<:Real, 'T', 'C'), 1, size(Q.factors, 1), Q.factors, Q.τ, X))
rmul!(X::StridedVecOrMat{T}, adjQ::Adjoint{<:Any,<:BlasHessenbergQ{T,false}}) where {T<:BlasFloat} =
    (Q = adjQ.parent; LAPACK.ormhr!('R', ifelse(T<:Real, 'T', 'C'), 1, size(Q.factors, 1), Q.factors, Q.τ, X))

lmul!(Q::BlasHessenbergQ{T,true}, X::StridedVecOrMat{T}) where {T<:BlasFloat} =
    LAPACK.ormtr!('L', Q.uplo, 'N', Q.factors, Q.τ, X)
rmul!(X::StridedVecOrMat{T}, Q::BlasHessenbergQ{T,true}) where {T<:BlasFloat} =
    LAPACK.ormtr!('R', Q.uplo, 'N', Q.factors, Q.τ, X)
lmul!(adjQ::Adjoint{<:Any,<:BlasHessenbergQ{T,true}}, X::StridedVecOrMat{T}) where {T<:BlasFloat} =
    (Q = adjQ.parent; LAPACK.ormtr!('L', Q.uplo, ifelse(T<:Real, 'T', 'C'), Q.factors, Q.τ, X))
rmul!(X::StridedVecOrMat{T}, adjQ::Adjoint{<:Any,<:BlasHessenbergQ{T,true}}) where {T<:BlasFloat} =
    (Q = adjQ.parent; LAPACK.ormtr!('R', Q.uplo, ifelse(T<:Real, 'T', 'C'), Q.factors, Q.τ, X))

lmul!(Q::HessenbergQ{T}, X::Adjoint{T,<:StridedVecOrMat{T}}) where {T} = rmul!(X', Q')'
rmul!(X::Adjoint{T,<:StridedVecOrMat{T}}, Q::HessenbergQ{T}) where {T} = lmul!(Q', X')'
lmul!(adjQ::Adjoint{<:Any,<:HessenbergQ{T}}, X::Adjoint{T,<:StridedVecOrMat{T}}) where {T}  = rmul!(X', adjQ')'
rmul!(X::Adjoint{T,<:StridedVecOrMat{T}}, adjQ::Adjoint{<:Any,<:HessenbergQ{T}}) where {T} = lmul!(adjQ', X')'

# multiply x by the entries of M in the upper-k triangle, which contains
# the entries of the upper-Hessenberg matrix H for k=-1
function rmul_triu!(M::AbstractMatrix, x, k::Integer=0)
    require_one_based_indexing(M)
    m, n = size(M)
    for j = 1:n, i = 1:min(j-k,m)
        @inbounds M[i,j] *= x
    end
    return M
end
function lmul_triu!(x, M::AbstractMatrix, k::Integer=0)
    require_one_based_indexing(M)
    m, n = size(M)
    for j = 1:n, i = 1:min(j-k,m)
        @inbounds M[i,j] = x * M[i,j]
    end
    return M
end

# when H is UpperHessenberg, it shares data with F.factors
# multiply Hessenberg by scalar (but don't modify lower triangle of F.H.data)
rmul!(F::Hessenberg{<:Any,<:UpperHessenberg{T}}, x::T) where {T<:Number} = Hessenberg(rmul_triu!(F.factors, x, -1), F.τ; μ=F.μ*x)
lmul!(x::T, F::Hessenberg{<:Any,<:UpperHessenberg{T}}) where {T<:Number} = Hessenberg(lmul_triu!(x, F.factors, -1), F.τ; μ=x*F.μ)

rmul!(F::Hessenberg{<:Any,<:SymTridiagonal{T}}, x::T) where {T<:Number} = Hessenberg(F.factors, F.τ, SymTridiagonal(F.H.dv*x, F.H.ev*x), F.uplo; μ=F.μ*x)
lmul!(x::T, F::Hessenberg{<:Any,<:SymTridiagonal{T}}) where {T<:Number} = Hessenberg(F.factors, F.τ, SymTridiagonal(x*F.H.dv, x*F.H.ev), F.uplo; μ=x*F.μ)

# Promote F * x or x * F.  In general, we don't know how to do promotions
# that would change the element type of F.H, however.
function (*)(F::Hessenberg{<:Any,<:AbstractMatrix{T}}, x::S) where {T,S<:Number}
    TS = typeof(zero(T) * x)
    if TS === T
        return rmul!(copy(F), convert(T, x))
    else
        throw(MethodError(*, (F, x)))
    end
end
function (*)(x::S, F::Hessenberg{<:Any,<:AbstractMatrix{T}}) where {T,S<:Number}
    TS = typeof(zero(T) * x)
    if TS === T
        return lmul!(convert(T, x), copy(F))
    else
        throw(MethodError(*, (x, F)))
    end
end
-(F::Hessenberg) = F * -one(eltype(F.H))

# shift Hessenberg by λI
+(F::Hessenberg, J::UniformScaling) = Hessenberg(F, F.μ + J.λ)
+(J::UniformScaling, F::Hessenberg) = Hessenberg(F, J.λ + F.μ)
-(F::Hessenberg, J::UniformScaling) = Hessenberg(F, F.μ - J.λ)
-(J::UniformScaling, F::Hessenberg) = Hessenberg(-F, J.λ - F.μ)

function ldiv!(F::Hessenberg, B::AbstractVecOrMat)
    Q = F.Q
    if iszero(F.μ)
        return lmul!(Q, ldiv!(F.H, lmul!(Q', B)))
    else
        return lmul!(Q, ldiv!(F.H, lmul!(Q', B); shift=F.μ))
    end
end

function rdiv!(B::AbstractMatrix, F::Hessenberg)
    Q = F.Q
    return rmul!(rdiv!(rmul!(B, Q), F.H; shift=F.μ), Q')
end

# handle case of real H and complex μ — we need to work around the
# fact that we can't multiple a real F.Q by a complex matrix directly in LAPACK
function ldiv!(F::Hessenberg{<:Complex,<:Any,<:AbstractMatrix{<:Real}}, B::AbstractVecOrMat{<:Complex})
    Q = F.Q
    Br = lmul!(Q', real(B))
    Bi = lmul!(Q', imag(B))
    ldiv!(F.H, B .= Complex.(Br,Bi); shift=F.μ)
    Br .= real.(B); Bi .= imag.(B)
    Br = lmul!(Q, Br)
    Bi = lmul!(Q, Bi)
    return B .= Complex.(Br,Bi)
end
function rdiv!(B::AbstractVecOrMat{<:Complex}, F::Hessenberg{<:Complex,<:Any,<:AbstractMatrix{<:Real}})
    Q = F.Q
    Br = rmul!(real(B), Q)
    Bi = rmul!(imag(B), Q)
    rdiv!(B .= Complex.(Br,Bi), F.H; shift=F.μ)
    Br .= real.(B); Bi .= imag.(B)
    Br = rmul!(Br, Q')
    Bi = rmul!(Bi, Q')
    return B .= Complex.(Br,Bi)
end

ldiv!(F::Adjoint{<:Any,<:Hessenberg}, B::AbstractVecOrMat) = rdiv!(B', F')'
rdiv!(B::AbstractMatrix, F::Adjoint{<:Any,<:Hessenberg}) = ldiv!(F', B')'

det(F::Hessenberg) = det(F.H; shift=F.μ)
logabsdet(F::Hessenberg) = logabsdet(F.H; shift=F.μ)
function logdet(F::Hessenberg)
    d,s = logabsdet(F)
    return d + log(s)
end

max_stars_repo_name: jonas-schulze/julia

max_issues_repo_name: jonas-schulze/julia

max_forks_repo_name: jonas-schulze/julia


text:
type Field{Format}
  name          :: String
  tableID       :: Int32
  columnID      :: Int16
  datatypeID    :: Int32
  datatypeSize  :: Int16
  typeModifier  :: Int32
end


text:
### A Pluto.jl notebook ###
# v0.16.1

using Markdown
using InteractiveUtils

# ╔═╡ 4df289c2-b460-49a6-b1b3-dac6c557397d
begin
	# TODO: remove this lines to use the notebook package manager.
	using Pkg
	Pkg.activate("../../")

	md"""
	# Lattice-Free MMI part II: training the model
	*[Lucas Ondel](https://lucasondel.github.io/index), September 2021*

	This notebook is the second step of the LF-MMI recipe. It implements the creation and the training of a *Time-Delay Neural Network* (TDNN) with the *Lattice-Free Maximum Mutual Information* (LF-MMI) objective function.

	Once the model is trained, the output of the neural network on the test set is dumped to an archive.
	"""
end

# ╔═╡ ee11cbaf-73c0-45db-84f7-db9eab7b6005
begin
	using AutoGrad
	using CUDA
	using Dates
	using JLD2
	using Knet
	using Logging
	using MarkovModels
	using HDF5
	using PlutoUI
	using Random
	using ShortCodes
	using TOML
end

# ╔═╡ fc878252-e6cd-4663-ad45-e42cae318ffb
TableOfContents()

# ╔═╡ 2ac818e7-7c52-4dfb-9f79-519fbce9d651
md"""
## Setup

Depending on your setup, select the CUDA version you want to use. You can also deactivate this cell and let the julia environment decide for itself.
"""

# ╔═╡ 876cffca-c46b-4805-a00f-5ac30d0631fe
ENV["JULIA_CUDA_VERSION"] = "11.3"

# ╔═╡ c10dd5eb-7bbe-4427-8a04-a3219d69942f
md"""
!!! warning
	If you change the CUDA version in the line above, you will need to restart the
	notebook to make the change effective.
"""

# ╔═╡ 850ac54c-b2e8-4e8e-a67a-b83f8d4d5905
md"""
Import the dependencies. Importantly, we use:
  * [MarkovModels](https://github.com/lucasondel/MarkovModels.jl) to implement the LF-MMI function and it's gradient
  * [KNet](https://github.com/denizyuret/Knet.jl) for the neural-network functions
  * [AutoGrad](https://github.com/denizyuret/AutoGrad.jl) for automatic differentiation.

!!! note
	It is possible to use [Flux](https://github.com/FluxML/Flux.jl)/[Zygote](https://github.com/FluxML/Zygote.jl)
    for the neural-network/automatic differentiation backend with little modifications
    to this notebook. The most notable changes will be to specify the gradient with
	Zygote API and to adapt the creation of the network.
"""

# ╔═╡ ad6986d7-71b3-49cc-92e4-dadf42953b19
md"""
For information, we print the CUDA configuration:
"""

# ╔═╡ bbde1167-ac11-4c15-8374-daa3f807cf3f
with_terminal() do
	CUDA.versioninfo()
end

# ╔═╡ 016eabc9-127c-4e12-98ab-8cd5004edfa0
md"""
We load the configuration file of the experiment. By default, we look for the file  `config.toml` in the root directory (i.e. the directory containing this notebook). Alternatively, when calling this notebook as a julia script, you can specify another configuration file by setting the environment variable `SPEECHLAB_LFMMI_CONFIG=/path/to/file`.
"""

# ╔═╡ 20d81979-4fd8-4b51-9b27-70f54fb2f886
rootdir = @__DIR__

# ╔═╡ 8ef896f4-c786-49f2-967b-72db140c933d
configfile = get(ENV, "SPEECHLAB_LFMMI_CONFIG", joinpath(rootdir, "config.toml"))

# ╔═╡ 104fc433-def8-44b1-b309-13d7683e0b33
begin
	@info "Reading configuration from $configfile."
	config = TOML.parsefile(configfile)
end

# ╔═╡ 1e09402a-a8b9-4e40-8bc2-2c0c153333c8
md"""
We use single precision.
"""

# ╔═╡ 65590fe8-f0e0-497f-bd8b-e128e23fd60b
const T = Float32

# ╔═╡ 8bf454e4-7707-4184-a9b3-13a9d072576a
md"""
## Input

Training with LF-MMI requires the numerator and denominator graphs (prepared by the previous step of the recipe) and the features for the training and development set.

The graphs should be organized as:
```
<rootdir>/
+-- <graphs.dir>/
|   +-- <dataset.name>/
|   |   +-- denominator_fsm.jld2
|   |   +-- dev_alignments_fsms.jld2
|   |   +-- pdfid_mapping.jld2
|   |   +-- train_alignments_fsms.jld2
```

As for the features, they should be in [HDF5](https://en.wikipedia.org/wiki/Hierarchical_Data_Format) archive organized as follows:
```
<features.dir>/
+-- <dataset.name>/
|   +-- train/
|   |   +-- <features.name>.h5
|   +-- dev/
|   |   +-- <features.name>.h5
|   +-- test/
|   |   +-- <features.name>.h5
```
where `<features.dir>`, `<features.name>` and `<dataset.name>` are taken from the configuration file.
!!! note
	The features directory path should be absolute and not relative to the root directory.
"""

# ╔═╡ 8ab11e2e-af68-43d1-a766-7095719133ed
graphsdir = joinpath(rootdir, config["graphs"]["dir"], config["dataset"]["name"])

# ╔═╡ 8936b2c7-68b6-4845-94b5-6512a1cb16a2
feadir = joinpath(config["features"]["dir"], config["dataset"]["name"])

# ╔═╡ 0b77b0f3-451f-41a9-a39a-da0525268ec5
feaname = config["features"]["name"]

# ╔═╡ 252656b1-0182-4616-9909-a7bc535244cf
featrain = joinpath(feadir, "train", feaname * ".h5")

# ╔═╡ 208c0a95-43ca-420c-b349-ac12da9631e4
feadev = joinpath(feadir, "dev", feaname * ".h5")

# ╔═╡ 907ad799-2f95-4b68-9033-7cc1d8853f00
featest = joinpath(feadir, "test", feaname * ".h5")

# ╔═╡ 270cb75f-62ad-4435-9932-3ff37d6db89f
use_gpu = config["training"]["use_gpu"]

# ╔═╡ aa271670-9e13-4838-968c-e15b34e47123
md"""
## Output

The notebook output the trained model and the output of the last layer on the train data.

```
<rootdir>/
+-- <experiment.dir>/
|   +-- <dataset.name>/
|   |   +-- <training.dir>/
|   |   |   +-- checkpoint.jld2
|   |   |   +-- best.jld2
|   |   +-- <output.dir>/
|   |   |   +-- test.h5
```
The keys `<experiment.dir>`, `<dataset.name>, `<training.dir>` and `<output.dir>` are taken from configuration file.

!!! note
	`best.jld2` is the model with the highest loss function on the development set whereas `checkpoint.jld2` is the model at the last epoch..

"""

# ╔═╡ eaa248f5-b8f3-4e3a-93a3-d4a275929878
expdir = joinpath(config["experiment"]["dir"], config["dataset"]["name"])

# ╔═╡ 6e8c6723-d6d9-4dbe-acd7-c0575b1f832e
traindir = joinpath(expdir, config["training"]["dir"])

# ╔═╡ d8de3103-bbe8-4858-8f97-ecebc3dd9103
outdir = joinpath(expdir, config["output"]["dir"])

# ╔═╡ 71f17a4a-88d5-449a-b0d0-fdee2956c96c
mkpath.([traindir, outdir])

# ╔═╡ d137c3f8-ceb5-4fbb-8e46-72bdb397d072
md"""
## Model creation

We use a TDNN model with ReLU activations. The input dimension corresponds to the input features dimension.
"""

# ╔═╡ f3e1352d-9b13-4378-9b5d-6ff7f65a4005
indim = h5open(featrain, "r") do f
	uttid, _ = iterate(keys(f))
	dim = size(read(f[uttid]), 1)
	@info "Neural network input dimension: $dim"
	dim
end

# ╔═╡ 69f0e6a3-e620-4b55-951f-d32488a48fc5
md"""
The output dimension corresponds to the number of pdf-ids.
"""

# ╔═╡ 110a9455-f3f4-49bf-9d26-f6ed94440076
outdim = jldopen(joinpath(graphsdir, "pdfid_mapping.jld2"), "r") do f
	dim = maximum( t -> t[2], f["pdfid_mapping"])
	@info "Neural network output dimension: $dim"
	dim
end

# ╔═╡ a2d14acd-fed2-42b9-85dd-77c5e04b982d
md"""
Now we build the model.
"""

# ╔═╡ 0b7d14dc-d580-4e5d-96c7-fd185deed2ed
ckpt_path = joinpath(traindir, "checkpoint.jld2")

# ╔═╡ 67fd7845-47ca-4b3a-9b78-68db1586ea9e
md"""
## Training
"""

# ╔═╡ e15b3f0f-6ec9-4bdf-9998-7b6359f6724f
md"""
First we load the phonotactic language model, i.e. the *denominator graph* of the MMI objective.
"""

# ╔═╡ 357e7c9d-9537-46f3-9153-b0bbccfd93b8
# Don't remove the end ';' as this object is too big to be displayed
den_fsm = load(joinpath(graphsdir, "denominator_fsm.jld2"), "fsm");

# ╔═╡ 19f0f3e2-dbce-4397-a9a2-12d6d0e9ea42
md"""
Now, we initialize the "training state". In our case, this comprises:
  * the parameters' optimizer
  * the best loss (on the validation set) achieved so far
  * the current epoch
  * the scheduler of the learning rate.
"""

# ╔═╡ c8dc633e-81a6-41ee-b61f-90b26dc04494
md"""
We define here the main function of the training. The input arguments are:

|Argument       | Description                                           |
|:--------------|:------------------------------------------------------|
| `model`       | Model to train                                        |
| `trainconfig` | Training part of the configuration of the experiment  |
| `trainstate`  | Current state of the training 						|
| `trainfea`    | Features archive for the train set                    |
| `trainnums`   | Numerator graphs for the train set                    |
| `devfea`      | Features archive for the train set                    |
| `devnums`     | Numerator graphs for the train set                    |
| `den_fsm`     | Phonotactic language model FSM for the denominator    |
| `use_gpu`     | Whether or not to use a GPU                           |
| `traindir`    | Output directory for the training 	                |
"""

# ╔═╡ 35aeec87-3267-4d07-8552-70d2c5da8d88
md"""
Let's train our model !
"""

# ╔═╡ 04c6afc4-c4cc-4df3-9eae-6b8e89d794f5
md"""
## Testing

Now the neural-network is trained and can be connected with a speech decoder.
As an example, we feed the test data to our model and store it in an HDF5 archive.
Because the HDF5 format can be read from most programming languages, you can probably use it with your favorite decoder 😉!

!!! note
	If you want to use a model already trained and stored on disk, you can simply load
    it by using:
	```julia
	load("/path/to/model.jld2", "model")
	```
"""

# ╔═╡ b66d91c2-9178-4074-bd20-e70d6475a0d5
md"""
Main function to dump the output of the network.
"""

# ╔═╡ 19166346-942f-4956-a83e-82ede6dfd2fa
md"""
Write the output to `<expdir>/output/test.h5`.
"""

# ╔═╡ d3ef660a-981a-4329-8565-9b1eeb393760
md"""
!!! tip
	With [python](https://www.python.org/), You can easily convert the HDF5 archive
	to a [Kaldi](https://github.com/kaldi-asr/kaldi) archive using
	[kaldi_io](https://github.com/vesis84/kaldi-io-for-python),
	[h5py](https://github.com/h5py/h5py) and
	[numpy](https://github.com/numpy/numpy). Here is an example:
	```python
	import h5py
	import kaldi_io
	import numpy as np

	h5archive = '<expdir>/output/test.h5'
	karchive = '<expdir>/output/test.ark'
	with h5py.File(h5archive) as fin:
		with open(karchive, 'wb') as fout:
			kaldi_io.write_mat(fout, np.array(fin[key]), key=key)
	```

"""

# ╔═╡ e881dd00-a629-4b72-8849-2615554053da
md"""
## References

Some references about TDNN and LF-MMI:
"""

# ╔═╡ 176d281d-d901-4fc5-ac44-2a1861793b50
[
	DOI("10.1109/29.21701"),
	DOI("10.21437/INTERSPEECH.2016-595"),
	DOI("10.21437/Interspeech.2018-1423"),
	DOI("10.21437/Interspeech.2020-3053")

]

# ╔═╡ f755a17b-6873-477d-9551-204b1ec858d9
md"""
## Utilities

Here is the implementations of all the utility functions we use in this notebook.
"""

# ╔═╡ 7df9cc64-79bb-45b5-b79b-98affd7d9d24
function update_scheduler!(s, opts, loss)
	s[:nsteps] += 1
	if loss < s[:best_loss] * (1 - s[:threshold])
		s[:nsteps] = 0
		s[:best_loss] = loss
	elseif s[:nsteps] > s[:patience]
		for opt in opts
			opt.lr = max(opt.lr * s[:factor], s[:min_lr])
		end
		@info "Decreasing the learning rate to $(opts[1].lr)."
		s[:nsteps] = 0
	end
	s
end

# ╔═╡ 154f2532-4e84-4a61-8860-5525be685243
md"""
### Neural Network API

Since KNet provides a rather low-level API to build neural network,
we define a few utilities that will make things easier for people
accustomed to pytorch/tensorflow. Unfold this cell if you want to
"""

# ╔═╡ f94e8e95-3ba7-437f-b3d1-ec4f2317a1fa
begin
	"""
		reallocate(obj, atype)

	Returns a copy of `obj'` with its buffers of type `atype`.
	This function is used to move the model to/from CPU/GPU.
	"""
	reallocate

	"""
		gpu(obj)

	Return a copy `obj` allocated on GPU.
	"""
	gpu

	"""
		cpu(obj)

	Return a copy `obj` allocated on CPU.
	"""
	cpu

	"""
		trainmode!(model)

	Set the model on "train mode" (i.e. activate dropout, batch-norm, ...).
	"""
	trainmode!

	"""
		testmode!(model)

	Set the model on "test mode".
	"""
	testmode!

	trainmode!(x) = x
	testmode!(x) = x
	# By default, the function does nothing but to copy the object.
	reallocate(obj, atype) = deepcopy(obj)
	gpu(obj) = reallocate(obj, CuArray{T})
	cpu(obj) = reallocate(obj, Array{T})

	function reallocate(opt::Adam, atype)
        fstm = isnothing(opt.fstm) ? nothing : atype(opt.fstm)
        scndm = isnothing(opt.scndm) ? nothing : atype(opt.scndm)
		Adam(opt.lr, opt.beta1, opt.beta2, opt.eps, opt.t, opt.gclip,
			 fstm, scndm)
	end

	"""
		struct PermuteDims
			perm
		end

	Neural network layer that permutes the dimension of the input layer.
	"""
	struct PermuteDims perm end
	(p::PermuteDims)(X) = permutedims(X, p.perm)
	getstate(p::PermuteDims) = Dict()
	setstate!(p::PermuteDims, s) = nothing
	Base.show(io::IO, p::PermuteDims) = print(io, "PermuteDims($(p.perm))")

	"""
		struct AddExtraDim end

	Neural network layer that adds an extra dimension to the 3D input-tensor.
	This is used to simuate 1D convolution with 2D convolution functions.
	"""
	struct AddExtraDim end
	(::AddExtraDim)(X) = reshape(X, size(X, 1), 1, size(X, 2), size(X, 3))
	getstate(::AddExtraDim) = Dict()
	setstate!(::AddExtraDim, s) = nothing
	Base.show(io::IO, ::AddExtraDim) = print(io, "AddExtraDim()")

	"""
		struct RemoveExtraDim end

	Neural network layer that removes the extra dimension added by
	[`AddExtraDim`](@ref).
	"""
	struct RemoveExtraDim end
	(::RemoveExtraDim)(X) = reshape(X, size(X, 1), size(X, 3), size(X, 4))
	getstate(::RemoveExtraDim) = Dict()
	setstate!(::RemoveExtraDim, s) = nothing
	Base.show(io::IO, ::RemoveExtraDim) = print(io, "RemoveExtraDim()")

	"""
		mutable struct Dropout
			pdrop
			training
		end

	Drop-out layers. If `training = true`, input dimensions are dropped
	with probability `pdrop`.
	"""
	mutable struct Dropout pdrop; training end
	Dropout(pdrop) = Dropout(pdrop, true)
	(d::Dropout)(X) = dropout(X, d.pdrop, drop = d.training)
	getstate(d::Dropout) = Dict(:pdrop => d.pdrop, :training => d.training)
	setstate!(d::Dropout, s) = (d.pdrop = s[:pdrop]; d.training = s[:training])
	Base.show(io::IO, d::Dropout) = print(io, "Dropout($(d.pdrop))")
	trainmode!(d::Dropout) = d.training = true
	testmode!(d::Dropout) = d.training = false

	"""
		struct Dense
			W
			b
			σ
		end

	Standard feed-forward layer with parameters `W` and `b` and activation
	function `σ`.
	"""
	mutable struct Dense W; b; σ end
	function (d::Dense)(X)
		rX = reshape(X, size(X, 1), :)
		Y = d.σ.(d.W * rX .+ d.b)
		reshape(Y, :, size(X)[2:end]...)
	end
	Dense(indim::Int, outdim::Int, σ = identity) = Dense(
		param(outdim, indim, atype = Array{T}),
		param0(outdim, atype = Array{T}), # `param0` initialize the buffer to 0.
		σ
	)
	reallocate(d::Dense, atype) = Dense(
		Param(atype(value(d.W))),
		Param(atype(value(d.b))),
		d.σ
	)

	getstate(d::Dense) = Dict(:W => d.W, :b => d.b, :σ => d.σ)
	setstate!(d::Dense, s) = (d.W = s[:W]; d.b = s[:b]; d.σ = s[:σ])

	function Base.show(io::IO, d::Dense)
		outdim, indim = size(d.W)
		print(io, "Dense($indim,  $outdim, $(d.σ))")
	end

	"""
		Conv(kernelsize, inchannels => outchannels; stride = (1, 1),
			 padding = (0, 0), dilation = (1, 1))

	2D convolution layer.
	"""
	mutable struct Conv W; b; dims; stride; padding; dilation end
	function Conv(ksize, in_out; stride = (1, 1), padding = (0, 0),
				  dilation = (1, 1))
		W = param(ksize..., in_out.first, in_out.second, atype = Array{T})
		b = param0(1, 1, in_out.second, 1, atype = Array{T})
		Conv(W, b, in_out, stride, padding, dilation)
	end
	(c::Conv)(X) = conv4(c.W, X, padding = c.padding, stride = c.stride,
						 dilation = c.dilation) .+ c.b
	function reallocate(c::Conv, atype)
		Conv(
			Param(atype(value(c.W))),
			Param(atype(value(c.b))),
			c.dims, c.stride, c.padding, c.dilation
		)
	end
	getstate(d::Conv) = Dict(:W => d.W, :b => d.b, :dims => d.dims,
							 :stride => d.stride, :padding => d.padding,
							 :dilation => d.dilation)
	setstate!(d::Conv, s) = (d.W = s[:W]; d.b = s[:b]; d.dims = s[:dims];
							 d.stride = s[:stride]; d.padding = s[:padding];
							 d.dilation = s[:dilation])
	function Base.show(io::IO, c::Conv)
		s = (size(c.W, 1), size(c.W, 2))
		print(io, "Conv($s, $(c.dims), stride = $(c.stride), padding = $(c.padding), dilation = $(c.dilation))")
	end

	"""
		mutable struct BatchNorm
			σ
			training
			...
		end

	Batch normalization layer.
	"""
	mutable struct BatchNorm
		moments
		params
		σ
		training
	end
	BatchNorm(dim, σ = identity) = BatchNorm(bnmoments(), Param(bnparams(T, dim)),
											 σ, true)
	(bn::BatchNorm)(X) = bn.σ.(batchnorm(X, bn.moments, bn.params;
										 training = bn.training))
	trainmode!(bn::BatchNorm) = bn.training = true
	testmode!(bn::BatchNorm) = bn.training = false

	function reallocate(bn::BatchNorm, atype)
		mean, var = bn.moments.mean, bn.moments.var
		BatchNorm(
			bnmoments(;
				momentum = bn.moments.momentum,
				mean = isnothing(mean) ? nothing : atype(mean),
				var = isnothing(var) ? nothing : atype(var),
				meaninit = bn.moments.meaninit,
				varinit = bn.moments.varinit
			),
			Param(atype(value(bn.params))),
			bn.σ, bn.training
		)
	end

	getstate(d::BatchNorm) = Dict(:moments => d.moments, :params => d.params,
								  :σ => d.σ, :training => d.training)
	setstate!(d::BatchNorm, s) = (d.moments = s[:moments]; d.params = s[:params];
								  d.σ = s[:σ]; d.training = s[:training])

	Base.show(io::IO, bn::BatchNorm) =
		print(io, "BatchNorm($(length(bn.params) ÷ 2), $(bn.σ))")

	"""
		struct Chain
			layers
		end

	Build a sequence of transformations.
	"""
	struct Chain
		layers
		Chain(layers...) = new(layers)
	end
	function (c::Chain)(X)
		for layer in c.layers
			X = layer(X)
		end
		X
	end
	trainmode!(c::Chain) = trainmode!.(c.layers)
	testmode!(c::Chain) = testmode!.(c.layers)
	reallocate(c::Chain, atype) = Chain(reallocate.(c.layers, atype)...)

	getstate(c::Chain) = getstate.(c.layers)
	setstate!(c::Chain, s) = setstate!.(c.layers, s)

	Base.show(io::IO, c::Chain) = print(io, "Chain($(c.layers))")
	function Base.show(io::IO, ::MIME"text/plain", c::Chain)
		println(io, "Chain(")
		for layer in c.layers
			println(io, "    $(layer),")
		end
		print(io, ")")
	end
end

# ╔═╡ a82a8460-0e7e-4a0a-b22b-b1760796204e
md"""
### TDNN construction
"""

# ╔═╡ 3d95e008-866b-4670-a9a1-c948f6733464
function buildmodel(modelconfig, indim, outdim)
	dropout = modelconfig["dropout"]
	dilations = modelconfig["dilations"]
	hdims = modelconfig["hidden-dims"]
	ksizes = modelconfig["kernel-sizes"]
	strides = modelconfig["strides"]

	layers = []
	push!(layers, PermuteDims([2, 1, 3])) # Features dimensions as channels.
	push!(layers, AddExtraDim()) # To simulate 1D convolution.
	for i in 1:length(hdims)
		layer = Conv(
			(ksizes[i], 1),
			indim => hdims[i],
			stride = (strides[i], 1),
			dilation = (dilations[i], 1),
			padding = (dilations[i]*(ksizes[i] - 1) ÷ 2, 0)
		)
		push!(layers, layer)
		push!(layers, BatchNorm(hdims[i], relu))
		push!(layers, Dropout(dropout))
		indim = hdims[i]
	end
	push!(layers, RemoveExtraDim())
	push!(layers, PermuteDims([2, 1, 3]))
	push!(layers, Dense(hdims[end], outdim))

	model = Chain(layers...)
end

# ╔═╡ ed5d482b-0504-4716-bfc0-3a635db50704
initmodel = buildmodel(config["model"], indim, outdim)

# ╔═╡ 21585c4e-463a-4489-9c73-ac6d585f9c46
if ispath(ckpt_path)
	modelstate = load(ckpt_path, "model")
	setstate!(initmodel, modelstate)
end;

# ╔═╡ 1273528f-6d95-4b80-924c-e6a8eb7c7d34
trainstate = if ispath(ckpt_path)
	ckpt = load(ckpt_path)

	@info "Found existing checkpoint at $(ckpt_path)"
	@info "Training will start at epoch $(load(ckpt_path, "epoch"))"

	Dict(
		"optimizers" => load(ckpt_path, "optimizers"),
		"epoch" => load(ckpt_path, "epoch"),
		"bestloss" => load(ckpt_path, "bestloss"),
		"scheduler" => load(ckpt_path, "scheduler"),
	)
else
	local trainconfig = config["training"]
	Dict(
		"optimizers" => [Adam(lr = trainconfig["optimizer"]["learning_rate"],
	                          beta1 = trainconfig["optimizer"]["beta1"],
			                  beta2 = trainconfig["optimizer"]["beta2"])
		                 for p in params(initmodel)] ,
		"epoch" => 1,
		"bestloss" => Inf,
		"scheduler" => Dict(
			:factor => config["training"]["scheduler"]["factor"],
			:patience => config["training"]["scheduler"]["patience"],
			:threshold => config["training"]["scheduler"]["threshold"],
			:best_loss => Inf, # best loss so far
			:nsteps => 0, # current number of steps without improvment
			:min_lr => config["training"]["scheduler"]["min_learning_rate"]
		)
	)
end;

# ╔═╡ 1ad2d2e4-b23f-464c-b53a-47b7a8ea9649
md"""
The next function is needed to calculate the length of the sequences at the end of the neural network (after the downsampling).
"""

# ╔═╡ 0c081ba2-7b97-4aca-8608-36fc74114066
function getlengths(modelconfig, inlengths)
	dilations = modelconfig["dilations"]
	hdims = modelconfig["hidden-dims"]
	ksizes = modelconfig["kernel-sizes"]
	strides = modelconfig["strides"]

	seqlengths = inlengths
	for i in 1:length(dilations)
		ksize = ksizes[i]
		pad = dilations[i] * (ksize - 1) ÷ 2
		newlengths = seqlengths .+ 2*pad .- dilations[i]*(ksize - 1) .+ strides[i] .- 1
		seqlengths = newlengths .÷ strides[i]
	end
	seqlengths
end

# ╔═╡ 8af48863-6b5f-496b-9c8f-2543fb412b59
md"""
### TDNN training
"""

# ╔═╡ 12f7ef0c-3b5a-4d24-8077-876d3718afa8
function train_epoch!(model, batchloader, denfsm, opts, updatefreq, use_gpu)
	trainmode!(model)
	θ = params(model)
	acc_loss = 0
	acc_etime = 0
	N = 0
	gs = nothing
	for (i, (batch_data, batch_nums, inlengths)) in enumerate(batchloader)

		seqlengths = getlengths(config["model"], inlengths)
		batch_dens = union(repeat([denfsm], size(batch_data, 3))...)

		if use_gpu
			batch_data = CuArray(batch_data)
			batch_nums = MarkovModels.gpu(batch_nums)
		end

		GC.gc()
		CUDA.reclaim()

		t₀ = now()
		L = @diff begin
            Y = model(batch_data)
            lfmmi_loss(Y, batch_nums, batch_dens, seqlengths)
        end
		t₁ = now()

		@debug "[batch=$i] loss+backprop time = $((t₁ - t₀).value / 1000)"
		gs = isnothing(gs) ? grad.([L], θ) : gs .+ grad.([L], θ)

		if i % updatefreq == 0
			update!(value.(θ), gs, opts)
            gs = nothing
		end
		acc_loss += value(L)
		N += sum(seqlengths)

        if use_gpu CUDA.unsafe_free!(batch_data) end
        batch_nums = nothing
		L = nothing
	end

    # If the gradients of the last batches were not used, make a
    # final update.
    if ! isnothing(gs)
        update!(value.(θ), gs, opts)
    end

	acc_loss / N
end

# ╔═╡ f5e0931c-66a5-4bde-b7e0-d0391a234b23
md"""
### Loading data
"""

# ╔═╡ 5964d2d5-a61b-43a2-b37a-97e2e06331eb
function padutterance(utt, Nmax)
	D, N = size(utt)
	pad = similar(utt, D, Nmax - N)
	pad .= utt[:, end]
	hcat(utt, pad)
end

# ╔═╡ dcf4e7ed-15f1-46b9-9516-c082d7f482b2
function BatchLoader(h5data, alifsms, batchsize; shuffledata = false)
	uttids = sort(keys(h5data), by = k -> size(h5data[k], 2))
	bins = [uttids[i:min(i+batchsize-1, length(uttids))]
		    for i in 1:batchsize:length(uttids)]

	# We only shuffle at the batch level.
	if shuffledata shuffle!(bins) end

	Channel() do ch
		for bin in bins
			# Load the features' utterance and build the batch.
			feas = [read(h5data[uttid]) for uttid in bin]
			seqlengths = [size(fea, 2) for fea in feas]
			Nmax = maximum(seqlengths)
			batch_feas = cat(padutterance.(feas, Nmax)..., dims = 3)

			# Load the numerator graphs.
			batch_alis = union([alifsms[uttid] for uttid in bin]...)

			put!(ch, (batch_feas, batch_alis, seqlengths))
		end
	end
end

# ╔═╡ 1de913ef-d987-42ce-858a-05b3213da1a2
function DumpBatchLoader(h5data, batchsize)
	uttids = sort(keys(h5data), by = k -> size(h5data[k], 2))
	bins = [uttids[i:min(i+batchsize-1, length(uttids))]
		    for i in 1:batchsize:length(uttids)]

	Channel() do ch
		for bin in bins
			# Load the features' utterance and build the batch.
			feas = [read(h5data[uttid]) for uttid in bin]
			seqlengths = [size(fea, 2) for fea in feas]
			Nmax = maximum(seqlengths)
			batch_feas = cat(padutterance.(feas, Nmax)..., dims = 3)

			put!(ch, (batch_feas, seqlengths, bin))
		end
	end
end

# ╔═╡ 8ec6d305-7a49-47f1-a051-b2a1195e7555
function dump_output(model, feafile, outfile, mbsize, use_gpu)
	testmode!(model)

	h5open(feafile, "r") do fin
		h5open(outfile, "w") do fout
			bl = DumpBatchLoader(fin, mbsize)
			for (batch_data, inlengths, uttids) in bl
				if use_gpu
					batch_data = CuArray(batch_data)
				end

				seqlengths = getlengths(config["model"], inlengths)
				Y = Array(model(batch_data))
				for i in 1:size(batch_data, 3)
					Yᵢ = Y[:, 1:seqlengths[i], i]
					fout[uttids[i]] = Yᵢ
				end
			end
		end
	end
end

# ╔═╡ bb446ac1-e770-41fb-8177-4d5b59d0bb8b
md"""
### Loss function

The LF-MMI loss and its gradients.
"""

# ╔═╡ e6c62410-4439-4fc5-bf7c-a6f5abcdc082
function _lfmmi_loss(ϕ, numfsms, denfsms, seqlengths)
	t₀ = now()
	γ_num, ttl_num = pdfposteriors(numfsms, ϕ, seqlengths)
	γ_den, ttl_den = pdfposteriors(denfsms, ϕ, seqlengths)
	t₁ = now()
	etime = (t₁ - t₀).value / 1000 # Forward-Backward time for debugging.
	@debug "forward-backward time = $etime"

	loss = -sum(ttl_num .- ttl_den)
	loss, γ_num, γ_den
end

# ╔═╡ 41f82983-0c37-4e88-a5e6-3a08ae51cf7b
_∇lfmmi_loss(γ_num, γ_den) = -(γ_num .- γ_den)

# ╔═╡ 24710344-74fe-4d6e-b003-c1bd49b529d8
@primitive1 _lfmmi_loss(ϕ, numf, denf, slen),dy,y (dy[1]*_∇lfmmi_loss(y[2:end]...))

# ╔═╡ 9bd4a7e7-4699-4306-9c61-1ad5ebad8e0a
lfmmi_loss(ϕ, numfsms, denfsms, seqlengths) =
	_lfmmi_loss(ϕ, numfsms, denfsms, seqlengths)[1]

# ╔═╡ 7222e9b4-8786-454c-9e2b-86c80bc27b29
function validate!(model, batchloader, denfsm, use_gpu)
	testmode!(model)
	acc_loss = 0
	N = 0
	for (i, (batch_data, batch_nums, inlengths)) in enumerate(batchloader)
		#GC.gc()
		#CUDA.reclaim()

		seqlengths = getlengths(config["model"], inlengths)
		batch_dens = union(repeat([denfsm], size(batch_data, 3))...)

		if use_gpu
			batch_data = CuArray(batch_data)
			batch_nums = MarkovModels.gpu(batch_nums)
		end

		L = lfmmi_loss(model(batch_data), batch_nums, batch_dens, seqlengths)
		acc_loss += value(L) / sum(seqlengths)
		N += 1
	end
	acc_loss / N
end

# ╔═╡ 3041378b-a4f7-48d6-a4f2-3f19020104e0
function train_model!(model, trainconfig, trainstate, trainfea, trainnums,
					  devfea, devnums, den_fsm, use_gpu, traindir)
    total_t₀ = now()

	epochs = trainconfig["epochs"]
	curriculum = trainconfig["curriculum"]
	mbsize = trainconfig["minibatch_size"]
	ckpt_path = joinpath(traindir, "checkpoint.jld2")
	best_path = joinpath(traindir, "best.jld2")

	startepoch = trainstate["epoch"]
	opts = trainstate["optimizers"]
	scheduler = trainstate["scheduler"]
	bestloss = trainstate["bestloss"]
	bestmodel = nothing

	opts = use_gpu ? opts .|> gpu : opts .|> cpu

	for epoch in startepoch:epochs
		t₀ = now()

		shuffledata = epoch > curriculum ? true : false
		train_loss = dev_loss = 0

		jldopen(trainnums, "r") do train_numfsms
			h5open(trainfea, "r") do f
				trainbl = BatchLoader(f, train_numfsms, mbsize; shuffledata)
				train_loss = train_epoch!(model, trainbl, den_fsm, opts,
                                          trainconfig["update_freq"], use_gpu)
			end
		end

		jldopen(devnums, "r") do dev_numfsms
			h5open(devfea, "r") do f
				devbl = BatchLoader(f, dev_numfsms, mbsize)
				dev_loss = validate!(model, devbl, den_fsm, use_gpu)
			end
		end

		scheduler = update_scheduler!(scheduler, opts, dev_loss)

		checkpoint = Dict(
			"model" => getstate(model |> cpu),
			"optimizers" => opts .|> cpu,
			"epoch" => epoch+1,
			"bestloss" => bestloss,
			"scheduler" => scheduler
		)
		save(ckpt_path, checkpoint)
		if dev_loss < bestloss
			bestmodel = model |> cpu
			save(best_path, checkpoint)
		end

		t₁ = now()
		@info "epoch=$epoch/$epochs train_loss=$train_loss dev_loss=$dev_loss epoch_time=$((t₁ - t₀).value / 1000)"
	end
    total_t₁ = now()
    @debug "Total training time = $((total_t₁ - total_t₀).value / 1000) seconds."

	bestmodel
end

# ╔═╡ f6ecb7d9-9bcc-4076-94b8-7755a17f2a49
model = train_model!(
	use_gpu ? initmodel |> gpu : initmodel |> cpu,
	config["training"],
	trainstate,
	featrain,
	joinpath(graphsdir, "train_alignments_fsms.jld2"),
	feadev,
	joinpath(graphsdir, "dev_alignments_fsms.jld2"),
	use_gpu ? den_fsm |> MarkovModels.gpu : den_fsm,
	use_gpu,
	traindir
)

# ╔═╡ 988480f8-9666-4626-b404-ea0f69f96d4e
dump_output(
	use_gpu ? model |> gpu : model |> cpu,
	featest,
	joinpath(outdir, "test.h5"),
	config["training"]["minibatch_size"],
	use_gpu
)

# ╔═╡ Cell order:
# ╠═4df289c2-b460-49a6-b1b3-dac6c557397d
# ╠═fc878252-e6cd-4663-ad45-e42cae318ffb
# ╟─2ac818e7-7c52-4dfb-9f79-519fbce9d651
# ╠═876cffca-c46b-4805-a00f-5ac30d0631fe
# ╟─c10dd5eb-7bbe-4427-8a04-a3219d69942f
# ╟─850ac54c-b2e8-4e8e-a67a-b83f8d4d5905
# ╠═ee11cbaf-73c0-45db-84f7-db9eab7b6005
# ╟─ad6986d7-71b3-49cc-92e4-dadf42953b19
# ╠═bbde1167-ac11-4c15-8374-daa3f807cf3f
# ╟─016eabc9-127c-4e12-98ab-8cd5004edfa0
# ╠═20d81979-4fd8-4b51-9b27-70f54fb2f886
# ╠═8ef896f4-c786-49f2-967b-72db140c933d
# ╠═104fc433-def8-44b1-b309-13d7683e0b33
# ╟─1e09402a-a8b9-4e40-8bc2-2c0c153333c8
# ╠═65590fe8-f0e0-497f-bd8b-e128e23fd60b
# ╠═8bf454e4-7707-4184-a9b3-13a9d072576a
# ╠═8ab11e2e-af68-43d1-a766-7095719133ed
# ╠═8936b2c7-68b6-4845-94b5-6512a1cb16a2
# ╠═0b77b0f3-451f-41a9-a39a-da0525268ec5
# ╠═252656b1-0182-4616-9909-a7bc535244cf
# ╠═208c0a95-43ca-420c-b349-ac12da9631e4
# ╠═907ad799-2f95-4b68-9033-7cc1d8853f00
# ╠═270cb75f-62ad-4435-9932-3ff37d6db89f
# ╟─aa271670-9e13-4838-968c-e15b34e47123
# ╠═eaa248f5-b8f3-4e3a-93a3-d4a275929878
# ╠═6e8c6723-d6d9-4dbe-acd7-c0575b1f832e
# ╠═d8de3103-bbe8-4858-8f97-ecebc3dd9103
# ╠═71f17a4a-88d5-449a-b0d0-fdee2956c96c
# ╟─d137c3f8-ceb5-4fbb-8e46-72bdb397d072
# ╠═f3e1352d-9b13-4378-9b5d-6ff7f65a4005
# ╟─69f0e6a3-e620-4b55-951f-d32488a48fc5
# ╠═110a9455-f3f4-49bf-9d26-f6ed94440076
# ╟─a2d14acd-fed2-42b9-85dd-77c5e04b982d
# ╠═ed5d482b-0504-4716-bfc0-3a635db50704
# ╠═0b7d14dc-d580-4e5d-96c7-fd185deed2ed
# ╠═21585c4e-463a-4489-9c73-ac6d585f9c46
# ╟─67fd7845-47ca-4b3a-9b78-68db1586ea9e
# ╟─e15b3f0f-6ec9-4bdf-9998-7b6359f6724f
# ╠═357e7c9d-9537-46f3-9153-b0bbccfd93b8
# ╟─19f0f3e2-dbce-4397-a9a2-12d6d0e9ea42
# ╠═1273528f-6d95-4b80-924c-e6a8eb7c7d34
# ╟─c8dc633e-81a6-41ee-b61f-90b26dc04494
# ╠═3041378b-a4f7-48d6-a4f2-3f19020104e0
# ╟─35aeec87-3267-4d07-8552-70d2c5da8d88
# ╠═f6ecb7d9-9bcc-4076-94b8-7755a17f2a49
# ╟─04c6afc4-c4cc-4df3-9eae-6b8e89d794f5
# ╟─b66d91c2-9178-4074-bd20-e70d6475a0d5
# ╠═8ec6d305-7a49-47f1-a051-b2a1195e7555
# ╟─19166346-942f-4956-a83e-82ede6dfd2fa
# ╠═988480f8-9666-4626-b404-ea0f69f96d4e
# ╟─d3ef660a-981a-4329-8565-9b1eeb393760
# ╟─e881dd00-a629-4b72-8849-2615554053da
# ╟─176d281d-d901-4fc5-ac44-2a1861793b50
# ╟─f755a17b-6873-477d-9551-204b1ec858d9
# ╠═7df9cc64-79bb-45b5-b79b-98affd7d9d24
# ╟─154f2532-4e84-4a61-8860-5525be685243
# ╠═f94e8e95-3ba7-437f-b3d1-ec4f2317a1fa
# ╟─a82a8460-0e7e-4a0a-b22b-b1760796204e
# ╠═3d95e008-866b-4670-a9a1-c948f6733464
# ╟─1ad2d2e4-b23f-464c-b53a-47b7a8ea9649
# ╠═0c081ba2-7b97-4aca-8608-36fc74114066
# ╟─8af48863-6b5f-496b-9c8f-2543fb412b59
# ╠═12f7ef0c-3b5a-4d24-8077-876d3718afa8
# ╠═7222e9b4-8786-454c-9e2b-86c80bc27b29
# ╟─f5e0931c-66a5-4bde-b7e0-d0391a234b23
# ╠═5964d2d5-a61b-43a2-b37a-97e2e06331eb
# ╠═dcf4e7ed-15f1-46b9-9516-c082d7f482b2
# ╠═1de913ef-d987-42ce-858a-05b3213da1a2
# ╟─bb446ac1-e770-41fb-8177-4d5b59d0bb8b
# ╠═e6c62410-4439-4fc5-bf7c-a6f5abcdc082
# ╠═41f82983-0c37-4e88-a5e6-3a08ae51cf7b
# ╠═24710344-74fe-4d6e-b003-c1bd49b529d8
# ╠═9bd4a7e7-4699-4306-9c61-1ad5ebad8e0a

max_stars_repo_name: lucasondel/SpeechRecipes

max_issues_repo_name: lucasondel/SpeechRecipes

max_forks_repo_name: lucasondel/SpeechRecipes


text:
### Modeling operators

export seismod_fun

## Wrapper for FunctionMod2Dat type

function seismod_fun(dat_pars::DataParams, src_pars::SrcParams; abc_geom::ABC_Geom = ABC_Geom(), acc::Int64 = 4)

	return FunctionMod2Dat(
		(m; store_extracomp = false) -> forward_FD(m, dat_pars, src_pars; abc_geom = abc_geom, acc = acc, store_wav = store_extracomp),
		(m, Δd) -> jacobian_adjoint_FD(m, Δd, dat_pars, src_pars; abc_geom = abc_geom, acc = acc),
		(m, Δd) -> backward_FD(m, Δd, dat_pars, src_pars; abc_geom = abc_geom, acc = acc))

end

## Forward & backward (finite-differences)

function forward_FD(
	m::Model, dat_pars::DataParams, src_pars::SrcParams;
	abc_geom::ABC_Geom = ABC_Geom(), acc::Int64 = 4, store_wav::Bool = false)

	# Extending domain to include ABC
	nz, nb = size(m)
	m_ext = reshape(extend(m.m, abc_geom), 1, :, nb)
	tap = reshape(taper(abc_geom), 1, :)
	nz_ext = length(tap)

	# Initialize output
	nt, np = size(dat_pars)
	dat = Array{Float32, 3}(undef, nt, np, nb)
	if store_wav
		dttu = Array{Float32, 4}(undef, nt, np, nz_ext, nb)
	end

	# Initial conditions
	u0 = zeros(Float32, np, nz_ext, nb)
	u1 = zeros(Float32, np, nz_ext, nb)

	# Loop over time
	dz = m.geom.dz
	dt = dat_pars.dt
	wav = src_pars.wav
	abc_size = abc_geom.size
	p = dat_pars.p
	@inbounds for idx_t = 1:nt

		# Time step
		u2 = ((2f0*u1-u0).*(m_ext.-p.^2f0)/dt^2f0.+Dzz(u1, dz; acc = acc).+tap/dt.*u1)./((m_ext.-p.^2f0)/dt^2f0.+tap/dt)

		# Source injection
		isrc1 = Int64(floor(src_pars.zsrc/dz))+1
		u2[:, abc_size[1]+isrc1, :] += wav[idx_t]./((m_ext[:, abc_size[1]+1, :].-p.^2f0)/dt^2f0.+tap[:, abc_size[1]+1]/dt)

		# Float32ecording wavefield at surface position
		dat[idx_t, :, :] = reshape(u1[:, abc_size[1]+1, :], np, nb)

		# Float32ecording full wavefield
		if store_wav
			dttu[idx_t, :, :, :] = ((u2-2f0*u1+u0)/dt^2f0)
		end

		# Update wavefields
		u0 = u1
		u1 = u2

	end

	if !store_wav
		return Data(dat, dat_pars, nothing)
	else
		return Data(dat, dat_pars, dttu)
	end

end

function backward_FD(
	m::Model, Δd::Data,
	dat_pars::DataParams, src_pars::SrcParams;
	abc_geom::ABC_Geom = ABC_Geom(), acc::Int64 = 4)

	# Extending domain to include ABC
	nz, nb = size(m)
	m_ext = reshape(extend(m.m, abc_geom), 1, :, nb)
	tap = reshape(taper(abc_geom), 1, :)
	nz_ext = length(tap)

	# Initialize output
	nt, np = size(dat_pars)
	q = zeros(Float32, np, nz, nb)

	# Final conditions
	v0 = zeros(Float32, np, nz_ext, nb)
	v1 = zeros(Float32, np, nz_ext, nb)

	# Loop over time
	dz = m.geom.dz
	dt = dat_pars.dt
	wav = src_pars.wav
	abc_size = abc_geom.size
	p = dat_pars.p
	@inbounds for idx_t = nt:-1:1

		# Time step
		v2 = ((2f0*v1-v0).*(m_ext.-p.^2f0)/dt^2f0.+Dzz(v1, dz; acc = acc).+tap/dt.*v1)./((m_ext.-p.^2f0)/dt^2f0.+tap/dt)

		# Source injection
		v2[:, abc_size[1]+1, :] += Δd.d[idx_t, :, :]./((m_ext[:, abc_size[1]+1, :].-p.^2f0)/dt^2f0.+tap[abc_size[1]+1]/dt)

		# Update source
		q += v2[:, abc_size[1]+1:end-abc_size[2], :]*wav[idx_t]

		# Update wavefields
		v0 = v1
		v1 = v2

	end

	return q

end

## Jacobian adjoint

function jacobian_adjoint_FD(
	m::Model, Δd::Data,
	dat_pars::DataParams, src_pars::SrcParams;
	abc_geom::ABC_Geom = ABC_Geom(), acc::Int64 = 4)

	# Initialize output
	nz, nb = size(m)
	Δm = zeros(Float32, nz, nb)

	# Float32epeat forward computation, if not previously computed
	if Δd.dttu == nothing
		Δd.dttu = forward_FD(m, dat_pars, src_pars; abc_geom = abc_geom, acc = acc, store_wav = true).dttu
	end

	# Extending domain to include ABC
	nz, nb = size(m)
	m_ext = reshape(extend(m.m, abc_geom), 1, :, nb)
	tap = reshape(taper(abc_geom), 1, :)
	nz_ext = length(tap)

	# Final conditions
	nt, np = size(dat_pars)
	v0 = zeros(Float32, np, nz_ext, nb)
	v1 = zeros(Float32, np, nz_ext, nb)

	# Loop over time
	dz = m.geom.dz
	dt = dat_pars.dt
	wav = src_pars.wav
	abc_size = abc_geom.size
	p = dat_pars.p
	@inbounds for idx_t = nt:-1:1

		# Time step
		v2 = ((2f0*v1-v0).*(m_ext.-p.^2f0)/dt^2f0.+Dzz(v1, dz; acc = acc).+tap/dt.*v1)./((m_ext.-p.^2f0)/dt^2f0.+tap/dt)

		# Source injection
		v2[:, abc_size[1]+1, :] += Δd.d[idx_t, :, :]./((m_ext[:, abc_size[1]+1, :].-p.^2f0)/dt^2f0.+tap[abc_size[1]+1]/dt)

		# Update gradient
		Δm_ = -reshape(sum(v1.*Δd.dttu[idx_t, :, :, :], dims = 1), :, nb)
		Δm[1,       :] += vec(sum(Δm_[1:abc_size[1]+1,                 :], dims = 1))
		Δm[2:end-1, :] +=         Δm_[abc_size[1]+2:end-abc_size[2]-1, :]
		Δm[end,     :] += vec(sum(Δm_[end-abc_size[2]:end,             :], dims = 1))

		# Update wavefields
		v0 = v1
		v1 = v2

	end

	return Model(m.geom, Δm)

end


text:
using DataDeps

register(DataDep(
  "climate-model-simulation-crashes",
  "http://archive.ics.uci.edu/ml/datasets/Climate+Model+Simulation+Crashes",
  "http://archive.ics.uci.edu/ml/machine-learning-databases/00252/pop_failures.dat",
  "8a892cbb056072e84749aed7409d6f17c7b4a1c899fc98ce52dc2af8b727ff9c",
  post_fetch_method=(path -> begin
    # The whitespace in this file is weird. Read/write to normalize whitespace
    writedlm(path, readdlm(path))

    UCIData.process_dataset(path,
        target_index=21,
        feature_indices=3:20,
        id_indices=1:2,
        delim='\t',
        skipto=2,
    )
  end),
))

max_stars_repo_name: JackDunnNZ/uci-data

max_issues_repo_name: JackDunnNZ/uci-data

max_forks_repo_name: JackDunnNZ/uci-data


text:
module TestIR

using IRTest
IRTest.@include("__test_ir.jl")

end  # module


text:
mutable struct RationalFunctionField
    generating_set       # present always
    groebner_coeffs      # lazy field, not to be addressed directly
    groebner_ideal       # lazy field, not to be addressed directly
end

function RationalFunctionField(genset)
    RationalFunctionField(genset, [], [])
end

function gens(FF::RationalFunctionField)
    return FF.groebner_coeffs
end

function contains_randomized(
            FF::RationalFunctionField, elem;
            p=0.99)

    # Theorem 3.3 from
    # Structural identifiability via input-output projections, Pogudin et alia

    n = length(gens(parent(numerator(elem))))
    d = maximum([
            maximum([degrees(numerator(f))..., degrees(denominator(f))...])
            for f in FF.generating_set
    ])
    M = round(Int, 6d^(n+3) / (1 - p))

    @info "Random values bound M = $M"

    # TODO: merge into a funuction
    It, yoverx, basepolyring, nvariables, ground, ystrings, Q, t = generators_to_saturated_ideal(FF.generating_set)

    @info "" typeof(It[1])

    grounds = tosingular(ground)
    eval_ring, evalvars = Singular.PolynomialRing(
                              grounds,
                              [ystrings..., "t"],
                              # Gleb: why and what is this? To discuss
                              ordering=ordering_lp(nvariables)*ordering_c())

    eval_ring_2, = Nemo.PolynomialRing(ground, [ystrings..., "t"])

    G = GroebnerEvaluator(It, eval_ring, basepolyring, ground)

    p = generate_point(G, M=M)

    gb = evaluate(G, p)
    elem = idealize_and_eval_element(elem, eval_ring_2, [p..., one(ground)])
    elem = change_base_ring(grounds, elem, parent=eval_ring)

    @debug "Evaluated eLement is $elem"
    @debug "Evaluated groebner basis is $gb"

    I = Singular.Ideal(eval_ring, gb)

    return iszero(Singular.reduce(elem, GroebnerBasis.f4(I)))
end

function contains_using_groebner(FF::RationalFunctionField, elem)

    Is, yoverx, basepolyring, yoverxs, basepolyrings = groebner_ideal_to_singular(FF.groebner_ideal)

    f = idealize_element(elem, basepolyring, basepolyrings, yoverx, yoverxs)

    fs = Ideal(yoverxs, f)
    Is = Ideal(yoverxs, Is)

    return Singular.contains(Is, fs)
end

function compute_groebner!(
                               FF::RationalFunctionField;
                               backend_algorithm=new_generating_set)

    FF.groebner_ideal = backend_algorithm(FF.generating_set)
    FF.groebner_coeffs = field_generators(FF.groebner_ideal)
end

function simple_preprocess(FF::RationalFunctionField)
    generators = FF.groebner_coeffs
    tmp = copy(generators)
    tmp = filter((c) -> !isconstant(c), tmp)
    sort!(tmp, by=(c) -> sum(degrees(c)), rev=true)
    ans = [ first(tmp) ]
    for elem in tmp
        field = RationalFunctionField(ans)
        compute_groebner!(field)
        if !contains(field, elem)
            push!(ans, elem)
        end
    end
    ans = [ x * inv(Nemo.lc(numerator(x))) for x in ans ]
    ans
end

function simplify_generators!(
                              FF::RationalFunctionField;
                              backend_algorithm=simple_preprocess)
    FF.groebner_coeffs = backend_algorithm(FF)
end

function contains(FF::RationalFunctionField, elem; proved=true)
    if proved
        contains_using_groebner(FF, elem)
    else
        contains_randomized(FF, elem)
    end
end

function field_generators(FF::RationalFunctionField)
    FF.groebner_coeffs
end


text:
include(joinpath(@__DIR__, "types.jl"));

abstract type DipoleResponse end

(dr::DipoleResponse)(E0::Real) = begin; error("Not yet implemented"); E0; end

function μ_dielectric(E0::Real, K1::Real, K2::Real, cϕ::Real, sϕ::Real)
  n̂i = n̂(cϕ, sϕ);
  return ((K1 - K2) * E0 * cϕ * n̂i) + (K2 * [E0; 0.0]);
end

struct DielectricResponse <: DipoleResponse
  K1::Real;
  K2::Real;
end

function (dr::DielectricResponse)(E0::Real, cϕ::Real, sϕ::Real)
  return μ_dielectric(E0, dr.K1, dr.K2, cϕ, sϕ);
end

struct PolarResponse <: DipoleResponse
  M::FdMatrix;
end

function (pr::PolarResponse)(::Real, cϕ::Real, sϕ::Real)
  return pr.M * n̂(cϕ, sϕ);
end

max_stars_repo_name: grasingerm/polymer-stats

max_issues_repo_name: grasingerm/polymer-stats

max_forks_repo_name: grasingerm/polymer-stats


text:
const IMGUI_DEMO_COMBO_STATES = Dict(
    "basic_combo" => Combos.State(
        "combo",
        items=[
            "AAAA",
            "BBBB",
            "CCCC",
            "DDDD",
            "EEEE",
            "FFFF",
            "GGGG",
            "HHHH",
            "IIII",
            "JJJJ",
            "KKKK",
            "LLLLLLL",
            "MMMM",
            "OOOOOOO",
        ],
    ),
)


text:
import Distributions: Uniform, Weibull

uniform_dist(l, u) = return () -> rand(Uniform(l, u))
weibull_dist(a, mu) = return () -> rand(Weibull(a, mu))

function parse_dist(arg)
    # Parses input string for given distribution.
    # Returns a distribution, and the average
    d, params = split(arg, ':')
    params = [parse(Float64, p) for p in split(params, ',')]
    if d == "U"
        return uniform_dist(params...), sum(params)/2
    elseif d == "W"
        a, mu = params
        return weibull_dist(a, mu), mu
    elseif d == "C"
        return () -> params[1], params[1]
    else
        throw(DomainError(d, "Unrecognized distribution function"))
    end
end

mass_fraction(vol_frac, ρ_i, ρ_m) = vol_frac*ρ_i / (vol_frac*ρ_i + (1-vol_frac)*ρ_m)


text:
l_min = 1.e3;
l_loss = [];
l_loss_val = [];
l_grad = [];
l_pnorm = [];
iter = 1;
cb = function (p, loss, loss_val, g_norm; doplot=true)
    global l_loss, l_grad, iter
    push!(l_loss, loss)
    push!(l_loss_val, loss_val)
    push!(l_grad, g_norm)
    push!(l_pnorm, norm(p))

    if doplot & (iter % n_plot == 0)
        regression_plot(;max=10)

        plot_loss()

        @save string(ckpt_path, "/mymodel.bson") p opt l_loss l_loss_val l_grad l_pnorm iter
        if loss_val < l_min
            @save string(ckpt_path, "/model_$(length(l_loss)).bson") p opt l_loss l_loss_val l_grad l_pnorm iter
            global l_min = loss_val
        end
    end
    iter += 1
    return false
end

if is_restart
    @load string(ckpt_path, "/mymodel.bson") p opt l_loss l_loss_val l_grad l_pnorm iter
    iter += 1
end

function load_model(epoch)
    @load string(ckpt_path, "/model_$(epoch).bson") p opt l_loss l_loss_val l_grad l_pnorm iter
end

max_stars_repo_name: DENG-MIT/DeepReduction

max_issues_repo_name: DENG-MIT/DeepReduction

max_forks_repo_name: DENG-MIT/DeepReduction


text:
export
ruleSPSampleListOutNPP,
ruleVBSampleListOut

function ruleSPSampleListOutNPP(msg_out::Nothing,
	                            msg_s::Message{PointMass, Multivariate}, # Multivariate, because vectors of samples and weight are passed as parameters
	                            msg_w::Message{PointMass, Multivariate})

    s = deepcopy(msg_s.dist.params[:m])
    w = deepcopy(msg_w.dist.params[:m])

    return Message(variateType(s[1]), SampleList, s=s, w=w)
end

function ruleVBSampleListOut(dist_out::Any,
                             dist_s::ProbabilityDistribution{Multivariate, PointMass}, # Multivariate, because vectors of samples and weight are passed as parameters
                             dist_w::ProbabilityDistribution{Multivariate, PointMass})

    s = deepcopy(dist_s.params[:m])
    w = deepcopy(dist_w.params[:m])

    return Message(variateType(s[1]), SampleList, s=s, w=w)
end


text:
export Marginal, getdata, is_clamped, is_initial, as_marginal
export SkipClamped, SkipInitial, SkipClampedAndInitial, IncludeAll

using Distributions
using Rocket

import Rocket: getrecent
import Base: ndims, precision, length, size

struct Marginal{D}
    data       :: D
    is_clamped :: Bool
    is_initial :: Bool
end

Base.show(io::IO, marginal::Marginal) = print(io, string("Marginal(", getdata(marginal), ")"))

getdata(marginal::Marginal)    = marginal.data
is_clamped(marginal::Marginal) = marginal.is_clamped
is_initial(marginal::Marginal) = marginal.is_initial

# TupleTools.prod is a more efficient version of Base.all for NTuple here
is_clamped(marginals::Tuple)  = TupleTools.prod(map(is_clamped, marginals))
is_initial(marginals::Tuple)  = TupleTools.prod(map(is_initial, marginals))

## Statistics

Distributions.pdf(marginal::Marginal, x)    = Distributions.pdf(getdata(marginal), x)
Distributions.logpdf(marginal::Marginal, x) = Distributions.logpdf(getdata(marginal), x)

MacroHelpers.@proxy_methods Marginal getdata [
    Distributions.mean,
    Distributions.median,
    Distributions.mode,
    Distributions.shape,
    Distributions.scale,
    Distributions.rate,
    Distributions.var,
    Distributions.std,
    Distributions.cov,
    Distributions.invcov,
    Distributions.logdetcov,
    Distributions.entropy,
    Distributions.params,
    Base.precision,
    Base.length,
    Base.ndims,
    Base.size,
    mean_cov,
    mean_var,
    mean_invcov,
    mean_precision,
    weightedmean_cov,
    weightedmean_var,
    weightedmean_invcov,
    weightedmean_precision,
    probvec,
    weightedmean,
    inversemean,
    logmean,
    meanlogmean,
    mirroredlogmean,
    loggammamean
]

## Utility functions

getdata(marginals::Tuple) = map(getdata, marginals)

as_marginal(marginal::Marginal) = marginal

## Marginal observable

abstract type MarginalSkipStrategy end

struct SkipClamped           <: MarginalSkipStrategy end
struct SkipInitial           <: MarginalSkipStrategy end
struct SkipClampedAndInitial <: MarginalSkipStrategy end
struct IncludeAll            <: MarginalSkipStrategy end

apply_skip_filter(observable, ::SkipClamped)           = observable |> filter(v -> !is_clamped(v))
apply_skip_filter(observable, ::SkipInitial)           = observable |> filter(v -> !is_initial(v))
apply_skip_filter(observable, ::SkipClampedAndInitial) = observable |> filter(v -> !is_initial(v) && !is_clamped(v))
apply_skip_filter(observable, ::IncludeAll)            = observable

struct MarginalObservable <: Subscribable{ Marginal }
    subject :: Rocket.RecentSubjectInstance{ Marginal, Subject{ Marginal, AsapScheduler, AsapScheduler } }
    stream  :: LazyObservable{ Marginal }
end

MarginalObservable() = MarginalObservable(RecentSubject(Marginal), lazy(Marginal))

as_marginal_observable(observable::MarginalObservable, skip_strategy::MarginalSkipStrategy) = apply_skip_filter(observable, skip_strategy)
as_marginal_observable(observable)                                                          = as_marginal_observable(observable, IncludeAll())

function as_marginal_observable(observable, skip_strategy::MarginalSkipStrategy)
    output = MarginalObservable()
    connect!(output, observable)
    return as_marginal_observable(output, skip_strategy)
end

Rocket.getrecent(observable::MarginalObservable) = Rocket.getrecent(observable.subject)
Rocket.getrecent(observables::Tuple)             = Rocket.getrecent.(observables)
Rocket.getrecent(::Nothing)                      = nothing

@inline Rocket.on_subscribe!(observable::MarginalObservable, actor) = subscribe!(observable.stream, actor)

@inline Rocket.subscribe!(observable::MarginalObservable, actor::Rocket.Actor{ <: Marginal })           = Rocket.on_subscribe!(observable.stream, actor)
@inline Rocket.subscribe!(observable::MarginalObservable, actor::Rocket.NextActor{ <: Marginal })       = Rocket.on_subscribe!(observable.stream, actor)
@inline Rocket.subscribe!(observable::MarginalObservable, actor::Rocket.ErrorActor{ <: Marginal })      = Rocket.on_subscribe!(observable.stream, actor)
@inline Rocket.subscribe!(observable::MarginalObservable, actor::Rocket.CompletionActor{ <: Marginal }) = Rocket.on_subscribe!(observable.stream, actor)

@inline Rocket.subscribe!(observable::MarginalObservable, actor::Rocket.Subject{ <: Marginal })                 = Rocket.on_subscribe!(observable.stream, actor)
@inline Rocket.subscribe!(observable::MarginalObservable, actor::Rocket.BehaviorSubjectInstance{ <: Marginal }) = Rocket.on_subscribe!(observable.stream, actor)
@inline Rocket.subscribe!(observable::MarginalObservable, actor::Rocket.PendingSubjectInstance{ <: Marginal })  = Rocket.on_subscribe!(observable.stream, actor)
@inline Rocket.subscribe!(observable::MarginalObservable, actor::Rocket.RecentSubjectInstance{ <: Marginal })   = Rocket.on_subscribe!(observable.stream, actor)
@inline Rocket.subscribe!(observable::MarginalObservable, actor::Rocket.ReplaySubjectInstance{ <: Marginal })   = Rocket.on_subscribe!(observable.stream, actor)

function connect!(marginal::MarginalObservable, source)
    set!(marginal.stream, source |> multicast(marginal.subject) |> ref_count())
    return nothing
end

function setmarginal!(marginal::MarginalObservable, value)
    next!(marginal.subject, Marginal(value, false, true))
    return nothing
end


text:
module JAnki

using HTTP
using JSON

include("adders.jl")
include("getters.jl")
include("utils.jl")

export anki_get_card, anki_get_tags
export anki_find_notes, anki_find_note, anki_find_notess
export anki_add_tag, anki_remove_tag
export anki_suspend_card

end


text:
module TestTextAnalysis
    using Base.Test
    using Languages
    using TextAnalysis
    using Compat

    my_tests = [
        "tokenizer.jl",
        "ngramizer.jl",
        "document.jl",
        "metadata.jl",
        "corpus.jl",
        "preprocessing.jl",
        "dtm.jl",
        "stemmer.jl"
    ]

    println("Running tests:")
    println(typeof(Compat.String))

    for my_test in my_tests
        println(" * $(my_test)")
        include(my_test)
    end
end