text:
function AlgQuat(K::Field, a::T, b::T) where { T <: FieldElem }

  M = zeros(K, 4, 4, 4)

  M[1, 1, 1] = one(K) # 1*1=1
  M[1, 2, 2] = one(K) # 1*i=i
  M[1, 3, 3] = one(K) # 1*j=j
  M[1, 4, 4] = one(K) # 1*ij=1

  M[2, 1, 2] = one(K)
  M[2, 2, 1] = a
  M[2, 3, 4] = one(K)
  M[2, 4, 3] = a

  M[3, 1, 3] = one(K)
  M[3, 2, 4] = -one(K)
  M[3, 3, 1] = b
  M[3, 4, 2] = -b

  M[4, 1, 4] = one(K)
  M[4, 2, 3] = -a
  M[4, 3, 2] = b
  M[4, 4, 1] = -a*b

  z = AlgQuat{T}()

  z.base_ring = K
  z.mult_table = M
  z.one = T[ one(K), zero(K), zero(K), zero(K) ]
  z.std = a, b
  return z
end

function show(io::IO, A::AlgQuat)
  print(io, "Quaternion algebra over\n")
  println(IOContext(io, :compact => true), base_ring(A))
  a, b = standard_form(A)
  print(io, "defined by i^2 = ", a, ", j^2 = ", b)
end

function show(io::IO, a::AlgAssElem{T, AlgQuat{T}}) where {T}
  v = a.coeffs
  print(io, "(", v[1], ") + (", v[2], ") * i + (", v[3], ") * j + (", v[4], ") * k")
end

dim(A::AlgQuat) = 4

base_ring(A::AlgQuat{T}) where {T} = A.base_ring::parent_type(T)

multiplication_table(A::AlgQuat; copy = false) = A.mult_table

standard_form(A::AlgQuat) = A.std

has_one(A::AlgQuat) = true

elem_type(A::AlgQuat{T}) where {T} = AlgAssElem{T, AlgQuat{T}}

elem_type(::Type{AlgQuat{T}}) where {T} = AlgAssElem{T, AlgQuat{T}}

iscommutative(A::AlgQuat) = false

issimple(A::AlgQuat) = true

issimple_known(A::AlgQuat) = true

dimension_of_center(A::AlgQuat) = 1

(A::AlgQuat{T})(a::Union{Integer, fmpz}) where {T} = A(map(base_ring(A), [a, 0, 0, 0]))

(A::AlgQuat{nf_elem})(a::nf_elem) = A(map(base_ring(A), [a, 0, 0, 0]))

(A::AlgQuat{T})(a::fmpq) where {T} = A(map(base_ring(A), [a, 0, 0, 0]))

order_type(::AlgQuat{T}) where { T <: NumFieldElem} = order_type(AlgQuat{T})

order_type(::Type{AlgQuat{T}}) where {T <: NumFieldElem} = AlgAssRelOrd{T, fractional_ideal_type(order_type(parent_type(T))), AlgQuat{T}}

################################################################################
#
#  Conjugation
#
################################################################################

function standard_involution(A::AlgQuat{T}) where {T}
  if isdefined(A, :std_inv)
    return A.std_inv::morphism_type(AlgQuat{T}, AlgQuat{T})
  else
    f = __standard_involution(A)
    A.std_inv = f
    return f
  end
end

function conjugate(a::AlgAssElem{T, AlgQuat{T}}) where {T}
  return standard_involution(parent(a))(a)
end

function trred(a::AlgAssElem{T, AlgQuat{T}}) where {T}
  return base_ring(parent(a))(a + conjugate(a))
end

function normred(a::AlgAssElem{T, AlgQuat{T}}) where {T}
  return base_ring(parent(a))(a * conjugate(a))
end

function reduced_charpoly(a::AlgAssElem{T, AlgQuat{T}}) where {T}
  A = parent(a)
  R = PolynomialRing(base_ring(A), "x", cached = false)[1]
  return x^2 - trred(a) * x + normred(a)
end

function standard_involution(A::AlgAss{T}) where {T}
  return __standard_involution(A)
end

function __standard_involution(A)
  BB = basis(A)

  if isone(BB[1])
    return ___standard_involution(A)
  end

  B = copy(basis(A))
  n = dim(A)
  o = one(A)
  for i in 1:n
    if dot(o.coeffs, B[i].coeffs) != 0
      B[i] = o
      B[i], B[1] = B[1], B[i]
      break
    end
  end

  _A, _AtoA = _change_basis(A, B)

  _f = ___standard_involution(_A)

  M = zero_matrix(base_ring(A), n, n)

  g = inv(_AtoA) * _f * _AtoA

  for i in 1:n
    elem_to_mat_row!(M, i, g(BB[i]))
  end

  return hom(A, A, M, inv(M))
end

# John Voight, "Quaternion algebra companion", Algorithm 4.6.1
# https://math.dartmouth.edu/~jvoight/hints-solns.pdf
function isquaternion_algebra(A::AlgAss)
  @assert dim(A) == 4
  @assert dimension_of_center(A) == 1

  f = standard_involution(A)
  K = base_ring(A)
  G = zero_matrix(K, 4, 4)
  B = copy(basis(A))
  @assert dot(B[1].coeffs, one(A).coeffs) != 0
  B[1] = one(A)
  for i in 1:4
    for j in 1:4
      G[i, j] = trred(B[i] * f(B[j]))//2
    end
  end

  F, S = _gram_schmidt(G, identity)
  stdbasis = elem_type(A)[]
  for i in 1:4
    push!(stdbasis, dot(B, [S[i, j] for j in 1:4]))
  end

  @assert stdbasis[1] == 1

  if iszero(det(F))
    return false
  end

  a = stdbasis[2]^2
  b = stdbasis[3]^2

  (scalea, scaleb, newa, newb) = _reduce_standard_form(K(a), K(b))

  stdbasis[2] = scalea * stdbasis[2]
  stdbasis[3] = scaleb * stdbasis[3]

  stdbasis[4] = stdbasis[2] * stdbasis[3]

  @assert stdbasis[2]^2 == A(newa)
  @assert stdbasis[3]^2 == A(newb)

  @assert stdbasis[2] * stdbasis[3] == -stdbasis[3] * stdbasis[2]
  @assert stdbasis[2] * stdbasis[3] == stdbasis[4]

  QA = AlgQuat(K, newa, newb)

  #@show stdbasis

  #@show basis(A)

  SB = basis_matrix(stdbasis)# * basis_matrix(B)

  #@show SB

  for i in 1:4
    @assert sum(SB[i, j] * basis(A)[j] for j in 1:4) == stdbasis[i]
  end

  SBinv = inv(SB)

  return QA, hom(QA, A, SB, SBinv)
end

################################################################################
#
#  Reduce standard generators
#
################################################################################

function _reduce_standard_form(a::nf_elem, b::nf_elem)
  K = parent(a)
  if isrational(a) && isrational(b)
    n, m, ap, bp = _reduce_standard_form(FlintQQ(a), FlintQQ(b))
    return K(n), K(m), K(ap), K(bp)
  else
    return one(K), one(K), a, b
  end
end

function _reduce_standard_form(a::fmpq, b::fmpq)
  n = denominator(a)
  ap = a * denominator(a)^2
  m = denominator(b)
  bp = b * denominator(b)^2

  apabs = abs(ap)

  while apabs > 1 && issquare(numerator(apabs))
    sq = sqrt(numerator(apabs))
    n = n//sq
    apabs = apabs//sq^2
  end

  ap = sign(ap) * apabs

  bpabs = abs(bp)

  while bpabs > 1 && issquare(numerator(bpabs))
    #@show numerator(bpabs)
    sq = sqrt(numerator(bpabs))
    m = m//sq
    bpabs = bpabs//sq^2
  end

  bp = sign(bp) * bpabs

  return (n, m, ap, bp)
end

# TODO: Think about this.
_reduce_standard_form(a::T, b::T) where {T} = (one(parent(a)), one(parent(a)), a, b)

################################################################################
#
#  Enumerate
#
################################################################################

# return elements x with absolute_tr(rednorm(x)) <= B (including zero) up to sign!
function Base.enumerate(O::Union{AlgAssRelOrd, AlgAssAbsOrd}, b::Int, equal::Bool = false)
  A = algebra(O)
  f = standard_involution(A)
  B = elem_in_algebra.(absolute_basis(O))
  d = length(B)
  G = zero_matrix(FlintQQ, d, d)
  for i in 1:d
    for j in 1:d
      G[i, j] = FlintZZ(absolute_tr(trred(B[i] * f(B[j]))))//2
    end
  end

  # TODO: Replace this by short_vectors_gram(M, nrr) once it works
  @assert !iszero(det(G))
  V = _short_vectors_gram(G, fmpz(b), hard = true)
  res = elem_type(O)[]
  for i in 1:length(V)
    y = sum(V[i][1][j] * B[j] for j in 1:d)
    @assert absolute_tr(normred(y)) <= b
    if equal
      if absolute_tr(normred(y)) == b
        push!(res, O(y))
      end
    else
      push!(res, O(y))
    end
  end

  return res
end

# Thanks Aurel!
# https://mathoverflow.net/questions/250753/finite-group-of-units-in-quaternion-orders
# TODO: There is a faster version in Magma.
function unit_group_modulo_scalars(O::AlgAssRelOrd)
  A = algebra(O)
  @assert A isa AlgQuat
  OF = base_ring(O)
  u, mu = unit_group(lll(OF))
  q, mq = quo(u, 2)
  norms = fmpz[]
  gens = elem_type(O)[]
  for e in q
    _x = mu(mq\e)
    # Reduce modulo squares, so that the trace is hopefully small
    x = evaluate(reduce_mod_powers(elem_in_nf(_x), 2))
    n = abs(FlintZZ(absolute_tr(x)))
    if !(n in norms)
      newel = enumerate(O, Int(n), true)
      for un in newel
        if isunit(un) && !(un in gens)
          isnew = true
          for oldunits in gens
            if (all(k -> iszero((elem_in_algebra(un) * inv(elem_in_algebra(oldunits))).coeffs[k]), 2:4))
              isnew = false
              break
            end
          end

          if isnew
            push!(gens, un)
          end
        end
      end
      push!(norms, n)
    end
  end

  @assert all(isunit(u) for u in gens)

  return gens
end

function unit_group_modulo_scalars(O::AlgAssAbsOrd)
  A = algebra(O)
  @assert A isa AlgQuat
  return enumerate(O, 1)
end

function _unit_group_generators_quaternion(O::Union{AlgAssRelOrd, AlgAssAbsOrd})
  gens1 = unit_group_modulo_scalars(O)
  u, mu = unit_group(base_ring(O))
  A = algebra(O)
  gens2 = [ O(A(elem_in_nf(mu(u[i])))) for i in 1:ngens(u) ]
  return append!(gens1, gens2)
end

### change basis

function _change_basis(A::AlgAss, bas)
  n = dim(A)
  M = zero_matrix(base_ring(A), n, n)
  N = zero_matrix(base_ring(A), n, n)

  for i in 1:n
    elem_to_mat_row!(M, i, bas[i])
  end

  # This is the "adjusted" basis matrix
  invM = inv(M)

  K = base_ring(A)

  mt = Array{elem_type(K), 3}(undef, n, n, n)

  for i in 1:n
    for j in 1:n
      c = bas[i] * bas[j]
      t = matrix(base_ring(A), 1, dim(A), c.coeffs) * invM
      @assert sum(t[1, i] * bas[i] for i in 1:n) == c
      for k in 1:n
        mt[i, j, k] = t[1, k]
      end
    end
  end

  B = AlgAss(K, mt)
  h = hom(B, A, M, invM)
  return B, h
end

function ___standard_involution(A)
  n = dim(A)
  o = one(A)

  @assert isone(basis(A)[1])

  N = zero_matrix(base_ring(A), n, n)

  K = base_ring(A)

  B = basis(A)

  @assert isone(B[1])
  t = elem_type(base_ring(A))[]
  push!(t, zero(K))
  for i in 2:n
    v = matrix(K, 1, n, (B[i]^2).coeffs)
    ti = v[1, i]
    push!(t, ti)
    _ni = B[i]^2 - ti * B[i]
    @assert all(i -> iszero(_ni.coeffs[i]), 2:n)
  end

  for i in 2:n
    for j in (i+1):n
      nij = (B[i] + B[j])^2 - (t[i] + t[j])*(B[i] + B[j])
      @assert all(i -> iszero(nij.coeffs[i]), 2:n)
    end
  end

  for i in 1:n
    b = i == 1 ? B[i] : t[i] - B[i]
    v = matrix(base_ring(A), 1, n, b.coeffs)
    for j in 1:n
      N[i, j] = v[1, j]
    end
  end
  invol = N
  return hom(A, A, invol, inv(invol))
end

global _debug = []

function _is_principal_maximal_quaternion_generic_proper(a, M, side = :right)
  A = algebra(M)
  f = standard_involution(A)
  K = base_ring(A)
  #@assert right_order(a) == M
  @assert _test_ideal_sidedness(a, M, :right)
  nr = normred(a)
  nr = simplify(nr)
  #@show norm(nr)
  #@show nr
  fl, c = isprincipal(nr)
  if !fl
    return false, zero(A)
  end
  #@show c
  fl, u, reps = _reps_for_totally_positive(c, K)
  if !fl
    return false, zero(A)
  end

  #@show u
  #@show istotally_positive(u * c)
  #@show u * c

  Babs = absolute_basis(a)
  d = length(Babs)
  G = zero_matrix(FlintZZ, d, d)
  #@show reps
  for z in reps
    Nnu = z * u * c

    alpha = inv(Nnu)

    _d = denominator(alpha, maximal_order(K))
    alpha = _d * alpha

    #@show isintegral(alpha)

    for i in 1:d
      for j in 1:d
        G[i, j] = FlintZZ(absolute_tr(alpha * trred(Babs[i] * f(Babs[j]))))
      end
    end

    B = 2 * trace(alpha * Nnu)

    @assert isintegral(B)

    ##@show Hecke._eltseq(G)
    #
    #@show denominator(G)

    #_d = degree(base_ring(A))

    #@show B

    v = _short_vectors_gram_integral(G, FlintZZ(B), hard = true)

    #if min == degree(base_ring(A))
    for w in v
      xi = sum(w[1][i] * Babs[i] for i in 1:length(Babs))
      if abs(norm(normred(xi))) == norm(normred(a))
        @assert xi in a
        @assert xi * M == a
        return true, xi
      end
    end
  end
  # TODO: Replace this by short_vectors_gram(M, nrr) once it works
  return false, zero(A)
end


text:
"""
    gewekediag(x::AbstractVector{<:Real}; first::Real=0.1, last::Real=0.5, kwargs...)

Compute the Geweke diagnostic [^Geweke1991] from the `first` and `last` proportion of
samples `x`.

The diagnostic is designed to asses convergence of posterior means estimated with
autocorrelated samples.  It computes a normal-based test statistic comparing the sample
means in two windows containing proportions of the first and last iterations.  Users should
ensure that there is sufficient separation between the two windows to assume that their
samples are independent.  A non-significant test p-value indicates convergence.  Significant
p-values indicate non-convergence and the possible need to discard initial samples as a
burn-in sequence or to simulate additional samples.

[^Geweke1991]: Geweke, J. F. (1991). Evaluating the accuracy of sampling-based approaches to the calculation of posterior moments (No. 148). Federal Reserve Bank of Minneapolis.
"""
function gewekediag(x::AbstractVector{<:Real}; first::Real=0.1, last::Real=0.5, kwargs...)
    0 < first < 1 || throw(ArgumentError("`first` is not in (0, 1)"))
    0 < last < 1 || throw(ArgumentError("`last` is not in (0, 1)"))
    first + last <= 1 || throw(ArgumentError("`first` and `last` proportions overlap"))

    n = length(x)
    x1 = x[1:round(Int, first * n)]
    x2 = x[round(Int, n - last * n + 1):n]
    z =
        (Statistics.mean(x1) - Statistics.mean(x2)) /
        hypot(mcse(x1; kwargs...), mcse(x2; kwargs...))
    p = SpecialFunctions.erfc(abs(z) / sqrt(2))

    return (zscore=z, pvalue=p)
end

ext: jl


text:
# Level-set
# ------------------------------------------------------------
@with_kw mutable struct LVData{T <: AbstractFloat}
    Q̃::T = 1e10
    incumbent::Int = 1
    major_iterations::Int = 0
    minor_iterations::Int = 0
    levelindex::Int = -1
end

@with_kw mutable struct LVParameters{T <: AbstractFloat}
    λ::T = 0.5
end

"""
    LevelSet

Functor object for using level-set regularization in an L-shaped algorithm. Create by supplying an [`LV`](@ref) object through `regularize ` in the `LShapedSolver` factory function and then pass to a `StochasticPrograms.jl` model.

...
# Parameters
- `λ::AbstractFloat = 0.5`: Controls the level position L = (1-λ)*θ + λ*Q̃, a convex combination of the current lower and upper bound.
- `penaltyterm::PenaltyTerm = Quadratic`: Specify penaltyterm variant ([`Quadratic`](@ref), [`Linearized`](@ref), [`InfNorm`](@ref), [`ManhattanNorm`][@ref])
...
"""
struct LevelSet{T <: AbstractFloat, A <: AbstractVector, P <: LQSolver, PT <: PenaltyTerm} <: AbstractRegularization
    data::LVData{T}
    parameters::LVParameters{T}

    ξ::A
    Q̃_history::A
    levels::A
    incumbents::Vector{Int}

    projectionsolver::P
    penaltyterm::PT

    function LevelSet(ξ₀::AbstractVector, projectionsolver::MPB.AbstractMathProgSolver, penaltyterm::PenaltyTerm; kw...)
        T = promote_type(eltype(ξ₀), Float32)
        ξ₀_ = convert(AbstractVector{T}, copy(ξ₀))
        A = typeof(ξ₀_)
        psolver = LQSolver(Model(), projectionsolver)
        P = typeof(psolver)
        PT = typeof(penaltyterm)
        return new{T, A, P, PT}(LVData{T}(), LVParameters{T}(;kw...), ξ₀_, A(), A(), Vector{Int}(), psolver, penaltyterm)
    end
end

function initialize_regularization!(lshaped::AbstractLShapedSolver, lv::LevelSet)
    MPB.loadproblem!(lv.projectionsolver.lqmodel, loadLP(StochasticPrograms.get_stage_one(lshaped.stochasticprogram))...)
    MPB.setobj!(lv.projectionsolver.lqmodel, zeros(decision_length(lshaped.stochasticprogram)))
    # θs
    for i = 1:nthetas(lshaped)
        MPB.addvar!(lv.projectionsolver.lqmodel, -Inf, Inf, 0.0)
    end
    initialize_penaltyterm!(lv.penaltyterm, lv.projectionsolver, lshaped.x)
    return nothing
end

function log_regularization!(lshaped::AbstractLShapedSolver, lv::LevelSet)
    @unpack Q̃,incumbent = lv.data
    push!(lv.Q̃_history, Q̃)
    push!(lv.incumbents, incumbent)
    return nothing
end

function log_regularization!(lshaped::AbstractLShapedSolver, t::Integer, lv::LevelSet)
    @unpack Q̃,incumbent = lv.data
    lv.Q̃_history[t] = Q̃
    lv.incumbents[t] = incumbent
    return nothing
end

function take_step!(lshaped::AbstractLShapedSolver, lv::LevelSet)
    @unpack Q = lshaped.data
    @unpack τ = lshaped.parameters
    @unpack Q̃ = lv.data
    if Q <= Q̃ - τ
        lv.data.Q̃ = Q
        lv.ξ .= current_decision(lshaped)
        lv.data.incumbent = timestamp(lshaped)
        lv.data.major_iterations += 1
    else
        lv.data.minor_iterations += 1
    end
    return nothing
end

function process_cut!(lshaped::AbstractLShapedSolver, cut::AbstractHyperPlane, lv::LevelSet)
    MPB.addconstr!(lv.projectionsolver.lqmodel, lowlevel(cut)...)
    return nothing
end

function project!(lshaped::AbstractLShapedSolver, lv::LevelSet)
    @unpack Q = lshaped.data
    if count(active_model_objectives(lshaped)) == nthetas(lshaped)
        _project!(lshaped, lv)
    end
    return nothing
end

function _project!(lshaped::AbstractLShapedSolver, lv::LevelSet)
    @unpack θ = lshaped.data
    @unpack Q̃ = lv.data
    @unpack λ = lv.parameters
    nt = nthetas(lshaped)
    c = sparse(MPB.getobj(lshaped.mastersolver.lqmodel))
    L = (1-λ)*θ + λ*Q̃
    push!(lv.levels, L)
    if lv.data.levelindex == -1
        MPB.addconstr!(lv.projectionsolver.lqmodel, c.nzind, c.nzval, -Inf, L)
        lv.data.levelindex = first_stage_nconstraints(lshaped.stochasticprogram)+ncuts(lshaped)+1
    else
        ub = MPB.getconstrUB(lv.projectionsolver.lqmodel)
        ub[lv.data.levelindex] = L
        MPB.setconstrUB!(lv.projectionsolver.lqmodel, ub)
    end
    # Update regularizer
    c = vcat(get_obj(lshaped), active_model_objectives(lshaped))
    update_penaltyterm!(lv.penaltyterm, lv.projectionsolver, c, 1.0, lshaped.x)
    # Solve projection problem
    solve_penalized!(lv.penaltyterm, lv.projectionsolver, lshaped.mastervector, lshaped.x, lshaped.x)
    if status(lv.projectionsolver) == :Infeasible
        @warn "Projection problem is infeasible, unprojected solution will be used"
        if Q̃ <= θ
            # If the upper objective bound is lower than the model lower bound for some reason, reset it.
            lv.data.Q̃ = Inf
        end
    elseif status(lv.projectionsolver) == :Optimal
        # Update master solution
        update_solution!(lshaped, lv.projectionsolver)
    else
        @warn "Projection problem could not be solved, unprojected solution will be used"
    end
    return nothing
end

function str(::LevelSet)
    return "L-shaped using level sets"
end

# API
# ------------------------------------------------------------
"""
    LV

Factory object for [`LevelSet`](@ref). Pass to `regularize ` in the `LShapedSolver` factory function. Equivalent factory calls: `LV`, `WithLV`, `LevelSet`, `WithLevelSets`. See ?LevelSet for parameter descriptions.

"""
struct LV <: AbstractRegularizer
    projectionsolver::MPB.AbstractMathProgSolver
    penaltyterm::PenaltyTerm
    parameters::Dict{Symbol,Any}
end
LV(; projectionsolver::MPB.AbstractMathProgSolver, penaltyterm = Quadratic(), kw...) = LV(projectionsolver, penaltyterm, Dict{Symbol,Any}(kw))
WithLV(; projectionsolver::MPB.AbstractMathProgSolver, penaltyterm = Quadratic(), kw...) = LV(projectionsolver, penaltyterm, Dict{Symbol,Any}(kw))
LevelSet(; projectionsolver::MPB.AbstractMathProgSolver, penaltyterm = Quadratic(), kw...) = LV(projectionsolver, penaltyterm, Dict{Symbol,Any}(kw))
WithLevelSets(; projectionsolver::MPB.AbstractMathProgSolver, penaltyterm = Quadratic(), kw...) = LV(projectionsolver, penaltyterm, Dict{Symbol,Any}(kw))

function add_regularization_params!(regularizer::LV; kwargs...)
    push!(regularizer.parameters, kwargs...)
    for (k,v) in kwargs
        if k ∈ [:projectionsolver, :penaltyterm]
            setfield!(regularizer, k, v)
            delete!(regularizer.parameters, k)
        end
    end
    return nothing
end

function (lv::LV)(x::AbstractVector)
    return LevelSet(x, lv.projectionsolver, lv.penaltyterm; lv.parameters...)
end

function str(::LV)
    return "L-shaped using level sets"
end


text:
using BlackBoxOptim
using FactCheck

NumTestRepetitions = 100


text:
# utility functions for the CUSPARSE wrappers

"""
check that the dimensions of matrix `X` and vector `Y` make sense for a multiplication
"""
function chkmvdims(X, n, Y, m)
    if length(X) != n
        throw(DimensionMismatch("X must have length $n, but has length $(length(X))"))
    elseif length(Y) != m
        throw(DimensionMismatch("Y must have length $m, but has length $(length(Y))"))
    end
end

"""
check that the dimensions of matrices `X` and `Y` make sense for a multiplication
"""
function chkmmdims( B, C, k, l, m, n )
    if size(B) != (k,l)
        throw(DimensionMismatch("B has dimensions $(size(B)) but needs ($k,$l)"))
    elseif size(C) != (m,n)
        throw(DimensionMismatch("C has dimensions $(size(C)) but needs ($m,$n)"))
    end
end


text:
using DelayEmbeddings, DynamicalSystemsBase
using Test, StaticArrays

println("\nTesting neighborhoods...")

@testset "Neighborhoods" begin

    ds = Systems.towel()
    data = trajectory(ds, 10000)

    tree = KDTree(data)

    @testset "neighborhood" begin

        point1 = data[100]
        point2 = point1 + 0.000001one(SVector{3})

        kn1 = [1724, 100, 6765]

        ntype = FixedMassNeighborhood(3)
        n1 = neighborhood(point1, tree, ntype)
        @test n1 == kn1
        n11 = neighborhood(point2, tree, ntype)
        @test n11 == kn1

        kn1 = [100, 6765]
        ntype = FixedSizeNeighborhood(0.01)
        n1 = neighborhood(point1, tree, ntype)
        @test n1 == kn1
        n11 = neighborhood(point2, tree, ntype)
        @test n11 == kn1
    end

    @testset "Theiler" begin

        for ntype in [FixedMassNeighborhood(3), FixedSizeNeighborhood(0.1)]
            n1 = neighborhood(point1, tree, ntype, 100, 1)
            @test 100 ∉ n1
            n11 = neighborhood(point2, tree, ntype, 100, 0)
            @test 100 ∈ n11
        end

        ds = Systems.lorenz()
        data = trajectory(ds, 200, dt = 0.001, alg = Vern9())
        tree = KDTree(data)

        point1 = data[10000]

        for ntype in [FixedMassNeighborhood(3), FixedSizeNeighborhood(0.1)]

            n1 = neighborhood(point1, tree, ntype)

            typeof(ntype) <: FixedMassNeighborhood &&
                @test sort(n1) == [9999, 10000, 10001]

            n2 = neighborhood(point1, tree, ntype, 10000, 10)
            for i in 1:3
                @test abs(10000 - i) ≥ 10
            end
        end
    end
end


text:
new_SIG = "Hello SIG"
println(new_SIG)


text:
##################################################################################
# This file is part of ModelBaseEcon.jl
# BSD 3-Clause License
# Copyright (c) 2020, Bank of Canada
# All rights reserved.
##################################################################################

"""
    linearize!(model::Model; <keyword arguments>)

Transform model into its linear approximation about its steady state.

### Keyword arguments
  * `sstate` - linearize about the provided steady state solution
  * `deviation`::Bool - whether or not the linearized model will tread data passed
to is as deviation from the steady state

See also: [`linearized`](@ref) and [`with_linearized`](@ref)
"""
function linearize! end
export linearize!

"""
    LinearizedModelEvaluationData <: AbstractModelEvaluationData

Model evaluation data for the linearized model case.
"""
struct LinearizedModelEvaluationData <: AbstractModelEvaluationData
    deviation::Bool
    sspt::Array{Float64,2}
    med::ModelEvaluationData
end

"""
    islinearized(m::Model)

Return `true` if the given model is linearized and `false` otherwise.
"""
islinearized(m::Model) = isa(m.evaldata, LinearizedModelEvaluationData)
export islinearized

# Specialize eval_R! for the new model evaluation type
function eval_R!(RES::AbstractVector{Float64}, point::AbstractMatrix{Float64}, lmed::LinearizedModelEvaluationData)
    med = lmed.med
    if lmed.deviation
        RES .= med.R .+ med.J * vec(point)
    else
        RES .= med.R .+ med.J * vec(point .- lmed.sspt)
    end
    return nothing
end

# Specialize eval_RJ for the new model evaluation type
function eval_RJ(point::AbstractMatrix{Float64}, lmed::LinearizedModelEvaluationData)
    med = lmed.med
    RES = similar(med.R)
    if lmed.deviation
        RES .= med.R .+ med.J * vec(point)
    else
        RES .= med.R .+ med.J * vec(point - lmed.sspt)
    end
    return RES, med.J
end

"""
    LinearizationError <: ModelErrorBase

A concrete error type used when a model cannot be linearized for some reason.
"""
struct LinearizationError <: ModelErrorBase
    reason
    hint::String
    LinearizationError(r) = new(r, "")
    LinearizationError(r, h) = new(r, string(h))
end
msg(le::LinearizationError) = "Cannot linearize model because $(le.reason)"
hint(le::LinearizationError) = le.hint
export LinearizationError

export linearize!
function linearize!(model::Model;
    # Idea:
    #
    # We compute the residual and the Jacobian matrix at the steady state and
    # store them in the model evaluation data. We store that into our new
    # linearized model evaluation data, which we place in the Model instance.
    # When we evaluate the linearized model residual, all we need to do is multiply
    # the stored steady state Jacobian by the deviation of the given point from
    # the steady state and add the stored steady state residual.

    sstate::SteadyStateData = model.sstate,
    deviation::Bool = false)

    if !isempty(model.auxvars) || !isempty(model.auxeqns)
        throw(LinearizationError("there are auxiliary variables.",
                    "Try setting `model.options.substitutions=false` in your model file."))
    end
    if !all(sstate.mask)
        throw(LinearizationError("the steady state is unknown.", "Solve for the steady state first."))
    end
    if maximum(abs, sstate.values[2:2:end]) > 1e-10
        throw(LinearizationError("the steady state has a non-zero linear growth."))
    end

    # We need a ModelEvaluationData in order to proceed
    if isa(model.evaldata, ModelEvaluationData)
        med = model.evaldata
    elseif isa(model.evaldata, LinearizedModelEvaluationData)
        # Model has already been linearized
        # We're not done though, because it may have been liearized about another point
        # So we simply retrieve the old ModelEvaluationData and keep going with it.
        med = model.evaldata.med
    else
        # What?! Didn't expect to get here. Oh well, we make a new ModelEvaluationData
        med = @timer ModelEvaluationData(model)
    end

    ntimes = 1 + model.maxlag + model.maxlead
    nvars = length(model.variables)
    nshks = length(model.shocks)
    @timer sspt = [repeat(sstate.values[1:2:(2 * nvars)], inner = ntimes); zeros(ntimes * nshks)]
    @timer sspt = reshape(sspt, ntimes, nvars + nshks)

    @timer eval_RJ(sspt, med)  # updates med.R and med.J in place
    model.evaldata = LinearizedModelEvaluationData(deviation, sspt, med)
    return model
end
@assert precompile(linearize!, (Model,))
@assert precompile(deepcopy, (Model,))

export linearized
"""
    linearized(model::Model; <arguments>)

Create a new model that is the linear approximation of the given model about its steady state.

### Keyword arguments
  * `sstate` - linearize about the provided steady state solution
  * `deviation`::Bool - whether or not the linearized model will tread data passed
to is as deviation from the steady state

See also: [`linearize!`](@ref) and [`with_linearized`](@ref)
"""
function linearized(model::Model; kwargs...)
    m = deepcopy(model)
    linearize!(m; kwargs...)
end

export with_linearized
"""
    with_linearized(F::Function, model::Model; <arguments>)

Apply the given function on a new model that is the linear approximation
of the given model about its steady state.  This is meant to be used
with the `do` syntax, as in the example below.

### Keyword arguments
  * `sstate` - linearize about the provided steady state solution
  * `deviation`::Bool - whether or not the linearized model will tread data passed
to is as deviation from the steady state

See also: [`linearize!`](@ref) and [`with_linearized`](@ref)

### Example

```julia
with_linearized(m) do lm
    # do something awesome with linearized model `lm`
end
# model `m` is still non-linear.
```
"""
function with_linearized(F::Function, model::Model; kwargs...)
    # store the evaluation data
    ed = model.evaldata
    # linearize
    linearize!(model; kwargs...)
    # do what we have to do
    ret = F(model)
    # restore the original model evaluation data
    model.evaldata = ed
    return ret
end


text:
# Combinations iterator with tabu
using Combinatorics
import Base: iterate, eltype, length

struct CombinationsWithout{T}
     a::T
     t::Int
     w::Int
     l::Int
     L::Int
 end

 function Base.iterate(c::CombinationsWithout)
     s = zeros(Int, c.t)
     if 0 < c.t <= c.l
         s[1] = 1
         if c.w == 1
             s[1] = 2
         end
         i = 1
         while i < c.t
             i += 1
             s[i] = s[i-1] + 1
             if s[i] == c.w
                 s[i] += 1
             end
         end
     end
     Base.iterate(c, (s, Vector{eltype(c.a)}(undef, c.t), 1))
 end

 function Base.iterate(c::CombinationsWithout, state)
     s, comb, k =  state
     k > c.L && return nothing
     k += 1
     for i in 1:c.t
         comb[i] = c.a[s[i]]
     end
     for i in length(s):-1:1
         s[i] += 1
         if s[i] > c.l - (length(s)-i) - (s[i] <= c.w)
             continue
         end
         if s[i] == c.w
             s[i] += 1
         end
         for j = i+1:length(s)
             s[j] = s[j-1]+1
             if s[j] == c.w
                 s[j] += 1
             end
         end
         break
     end
     comb, (s, comb, k)
 end

 length(c::CombinationsWithout) = c.L
 eltype(::Type{CombinationsWithout{T}}) where {T} = Vector{eltype(T)}

 """
     combinations_without(a, n::Integer, w)

 Generate all combinations of `n` elements from an indexable object except those with index `w`. Note that the combinations
 are modified inplace.
 """
 function combinations_without(a, t::Integer, w)
     l = length(a)
     ins = 1 <= w <= l
     L = binomial(l - ins, t)
     if ins && length(a) == w # w is last index
         CombinationsWithout(a, t, 0, l - 1, L)
     elseif ins && l == t # t too big
         CombinationsWithout(a, t, w, l, 0)
     elseif ins
         CombinationsWithout(a, t, w, l, L)
     else
         CombinationsWithout(a, t, 0, l, L)
     end
 end


text:
"""
 Holds the tableau of a stochastic implicit partitioned Runge-Kutta method.
 qdrift, pdrift hold the RK coefficients for the drift part,
 and qdiff, pdiff hold the RK coefficients for the diffusion part of the SDE.
"""
struct TableauSIPRK{T} <: AbstractTableauIRK{T}
    name::Symbol
    s::Int
    qdrift::CoefficientsRK{T}
    qdiff::CoefficientsRK{T}
    pdrift::CoefficientsRK{T}
    pdiff::CoefficientsRK{T}

    # Order of the tableau is not included, because unlike in the deterministic
    # setting, it depends on the properties of the noise (e.g., the dimension of
    # the Wiener process and the commutativity properties of the diffusion matrix)
    #
    # Orders stored in qdrift and qdiff are understood as the classical orders of these methods.

    function TableauSIPRK{T}(name, s, qdrift, qdiff, pdrift, pdiff) where {T}
        # THE COMMENTED OUT PART WAS FOR TableauFIRK. MAY IMPLEMENT SOMETHING
        # SIMILAR FOR TableauSFIRK LATER.

        # if (q.s > 1 && istrilstrict(q.a)) || (q.s==1 && q.a[1,1] == 0)
        #     warn("Initializing TableauFIRK with explicit tableau ", q.name, ".\n",
        #          "You might want to use TableauERK instead.")
        # elseif q.s > 1 && istril(q.a)
        #     warn("Initializing TableauFIRK with diagonally implicit tableau ", q.name, ".\n",
        #          "You might want to use TableauDIRK instead.")
        # end

        @assert s == qdrift.s == qdiff.s == pdrift.s == pdiff.s

        new(name, s, qdrift, qdiff, pdrift, pdiff)
    end
end

function TableauSIPRK(name::Symbol, qdrift::CoefficientsRK{T}, qdiff::CoefficientsRK{T}, pdrift::CoefficientsRK{T}, pdiff::CoefficientsRK{T}) where {T}
    TableauSIPRK{T}(name, qdrift.s, qdrift, qdiff, pdrift, pdiff)
end

function TableauSIPRK(name::Symbol, qorder_drift::Int, qa_drift::Matrix{T}, qb_drift::Vector{T}, qc_drift::Vector{T}, qorder_diff::Int, qa_diff::Matrix{T}, qb_diff::Vector{T}, qc_diff::Vector{T},
                                    porder_drift::Int, pa_drift::Matrix{T}, pb_drift::Vector{T}, pc_drift::Vector{T}, porder_diff::Int, pa_diff::Matrix{T}, pb_diff::Vector{T}, pc_diff::Vector{T}) where {T}
    TableauSIPRK{T}(name, length(qc_drift), CoefficientsRK(name, qorder_drift, qa_drift, qb_drift, qc_drift), CoefficientsRK(name, qorder_diff, qa_diff, qb_diff, qc_diff), CoefficientsRK(name, porder_drift, pa_drift, pb_drift, pc_drift), CoefficientsRK(name, porder_diff, pa_diff, pb_diff, pc_diff))
end

# TODO function readTableauSFIRKFromFile(dir::AbstractString, name::AbstractString)

"""
Parameters for right-hand side function of implicit Runge-Kutta methods.
  A - if positive, the upper bound of the Wiener process increments; if A=0.0, no truncation
"""
mutable struct ParametersSIPRK{DT, TT, ET <: PSDE{DT,TT}, D, M, S} <: Parameters{DT,TT}
    equ::ET
    tab::TableauSIPRK{TT}
    Δt::TT
    ΔW::Vector{DT}
    ΔZ::Vector{DT}
    A::DT

    t::TT
    q::Vector{DT}
    p::Vector{DT}
end

function ParametersSIPRK(equ::ET, tab::TableauSIPRK{TT}, Δt::TT, ΔW::Vector{DT}, ΔZ::Vector{DT}, A::DT) where {DT, TT, ET <: PSDE{DT,TT}}
    @assert equ.m == length(ΔW) == length(ΔZ)
    ParametersSIPRK{DT, TT, ET, equ.d, equ.m, tab.s}(equ, tab, Δt, ΔW, ΔZ, A, 0, zeros(DT, equ.d), zeros(DT, equ.d))
end

"""
Structure for holding the internal stages Q, the values of the drift vector
and the diffusion matrix evaluated at the internal stages VQ=v(Q), BQ=B(Q),
and the increments Y = Δt*a_drift*v(Q) + a_diff*B(Q)*ΔW
"""
struct IntegratorCacheSIPRK{DT}
    Q::Vector{Vector{DT}}
    P::Vector{Vector{DT}}
    V::Vector{Vector{DT}}
    F::Vector{Vector{DT}}
    B::Vector{Matrix{DT}}
    G::Vector{Matrix{DT}}
    Y::Vector{Vector{DT}}
    Z::Vector{Vector{DT}}

    v::Vector{DT}
    f::Vector{DT}
    b::Matrix{DT}
    g::Matrix{DT}
    y::Vector{DT}
    z::Vector{DT}

    V1::Vector{DT}
    V2::Vector{DT}
    F1::Vector{DT}
    F2::Vector{DT}
    B1::Matrix{DT}
    B2::Matrix{DT}
    G1::Matrix{DT}
    G2::Matrix{DT}
    Δw::Vector{DT}
    ΔQ::Vector{DT}
    Y1::Vector{DT}
    Y2::Vector{DT}
    ΔP::Vector{DT}
    Z1::Vector{DT}
    Z2::Vector{DT}
    Δy::Vector{DT}
    Δz::Vector{DT}

    function IntegratorCacheSIPRK{DT}(D, M, S) where {DT}

        # create internal stage vectors
        Q = create_internal_stage_vector(DT, D, S)
        P = create_internal_stage_vector(DT, D, S)
        V = create_internal_stage_vector(DT, D, S)
        F = create_internal_stage_vector(DT, D, S)
        B = create_internal_stage_vector(DT, D, M, S)
        G = create_internal_stage_vector(DT, D, M, S)
        Y = create_internal_stage_vector(DT, D, S)
        Z = create_internal_stage_vector(DT, D, S)

        # create velocity and update vector
        v = zeros(DT,D)
        f = zeros(DT,D)
        b = zeros(DT,D,M)
        g = zeros(DT,D,M)
        y = zeros(DT,D)
        z = zeros(DT,D)

        # create temporary arrays
        V1 = zeros(DT,D)
        V2 = zeros(DT,D)
        F1 = zeros(DT,D)
        F2 = zeros(DT,D)
        B1 = zeros(DT,D,M)
        B2 = zeros(DT,D,M)
        G1 = zeros(DT,D,M)
        G2 = zeros(DT,D,M)
        ΔQ = zeros(DT,D)
        Y1 = zeros(DT,D)
        Y2 = zeros(DT,D)
        ΔP = zeros(DT,D)
        Z1 = zeros(DT,D)
        Z2 = zeros(DT,D)
        Δw = zeros(DT,M)
        Δy = zeros(DT,M)
        Δz = zeros(DT,M)

        new(Q, P, V, F, B, G, Y, Z, v, f, b, g, y, z, V1, V2, F1, F2, B1, B2, G1, G2, Δw, ΔQ, Y1, Y2, ΔP, Z1, Z2, Δy, Δz)
    end
end

"Stochastic implicit partitioned Runge-Kutta integrator."
struct IntegratorSIPRK{DT, TT, PT <: ParametersSIPRK{DT,TT},
                               ST <: NonlinearSolver{DT}, N} <: StochasticIntegrator{DT,TT}
    params::PT
    solver::ST
    cache::IntegratorCacheSIPRK{DT}
end

# K - the integer in the bound A = √(2 K Δt |log Δt|) due to Milstein & Tretyakov; K=0 no truncation
function IntegratorSIPRK(equation::PSDE{DT,TT,VT,BT,N}, tableau::TableauSIPRK{TT}, Δt::TT; K::Int=0) where {DT,TT,VT,BT,N}
    D = equation.d
    M = equation.m
    NS= max(equation.ns,equation.ni)
    S = tableau.s

    # create params
    K==0 ? A = 0.0 : A = sqrt( 2*K*Δt*abs(log(Δt)) )
    params = ParametersSIPRK(equation, tableau, Δt, zeros(DT,M), zeros(DT,M), A)

    # create solver
    solver = create_nonlinear_solver(DT, 2*D*S, params)

    # create cache for internal stage vectors and update vectors
    cache = IntegratorCacheSIPRK{DT}(D, M, S)

    # create integrator
    IntegratorSIPRK{DT, TT, typeof(params), typeof(solver), N}(params, solver, cache)
end

@inline equation(integrator::IntegratorSIPRK) = integrator.params.equ
@inline timestep(integrator::IntegratorSIPRK) = integrator.params.Δt
@inline tableau(integrator::IntegratorSIPRK) = integrator.params.tab
@inline nstages(integrator::IntegratorSIPRK)  = nstages(tableau(integrator))
@inline eachstage(integrator::IntegratorSIPRK) = 1:nstages(integrator)
Base.eltype(integrator::IntegratorSIPRK{DT, TT, PT, ST, N}) where {DT, TT, PT, ST, N} = DT

function update_params!(int::IntegratorSIPRK, sol::AtomicSolutionPSDE)
    # set time for nonlinear solver and copy previous solution
    int.params.t  = sol.t
    int.params.q .= sol.q
    int.params.p .= sol.p
    int.params.ΔW .= sol.ΔW
    int.params.ΔZ .= sol.ΔZ

    # truncate the increments ΔW with A
    if int.params.A > 0
        for i in eachindex(int.params.ΔW)
            if int.params.ΔW[i] <- int.params.A
                int.params.ΔW[i] = -int.params.A
            elseif int.params.ΔW[i] > int.params.A
                int.params.ΔW[i] = int.params.A
            end
        end
    end
end

"""
Unpacks the data stored in x = (Y[1][1], Y[1][2], ... Y[1][D], Y[2][1], ..., Z[1][1], Z[1][2], ... Z[1][D], Z[2][1], ...)
into Y, Z::Vector{Vector}, calculates the internal stages Q, P, the values of the RHS
of the SDE ( v(Q,P), f(Q,P), B(Q,P) and G(Q,P) ), and assigns them to V, F, B and G.
Unlike for FIRK, here
Y = Δt a_drift v(Q,P) + a_diff B(Q,P) ΔW,
Z = Δt ̃a_drift v(Q,P) + ̃a_diff B(Q,P) ΔW.
"""
function compute_stages!(x::Vector{ST}, Q::Vector{Vector{ST}}, P::Vector{Vector{ST}},
                                        V::Vector{Vector{ST}}, F::Vector{Vector{ST}},
                                        B::Vector{Matrix{ST}}, G::Vector{Matrix{ST}},
                                        Y::Vector{Vector{ST}}, Z::Vector{Vector{ST}},
                                        params::ParametersSIPRK{DT,TT,ET,D,M,S}) where {ST,DT,TT,ET,D,M,S}

    local tqᵢ::TT       #times for the q internal stages
    local tpᵢ::TT       #times for the p internal stages

    @assert S == length(Q) == length(V) == length(B) == length(P) == length(F) == length(G)

    # copy x to Y, Z and calculate Q, P
    for i in 1:S
        @assert D == length(Q[i]) == length(V[i]) == size(B[i],1) == length(P[i]) == length(F[i]) == size(G[i],1)
        @assert M == size(B[i],2) == size(G[i],2)

        for k in 1:D
            Y[i][k] = x[D*(  i-1)+k]
            Z[i][k] = x[D*(S+i-1)+k]
        end

        Q[i] .= params.q .+ Y[i]
        P[i] .= params.p .+ Z[i]
    end

    # compute V=v(t,Q,P), F=f(t,Q,P), B=B(t,Q,P) and G=G(t,Q,P)
    for i in 1:S
        tqᵢ = params.t + params.Δt * params.tab.qdrift.c[i]
        tpᵢ = params.t + params.Δt * params.tab.pdrift.c[i]
        # calculates v(t,Q,P) and f(t,Q,P) and assigns to the i-th column of V and F
        params.equ.v(tqᵢ, Q[i], P[i], V[i])
        params.equ.f(tpᵢ, Q[i], P[i], F[i])
        # calculates B(t,Q,P) and assigns to the matrices B[i] and G[i]
        params.equ.B(tqᵢ, Q[i], P[i], B[i])
        params.equ.G(tpᵢ, Q[i], P[i], G[i])
    end
end

"Compute stages of implicit Runge-Kutta methods."
@generated function function_stages!(x::Vector{ST}, b::Vector{ST}, params::ParametersSIPRK{DT,TT,ET,D,M,S}) where {ST,DT,TT,ET,D,M,S}

    cache = IntegratorCacheSIPRK{ST}(D, M, S)

    quote
        compute_stages!(x, $cache.Q, $cache.P, $cache.V, $cache.F, $cache.B, $cache.G, $cache.Y, $cache.Z, params)

        local y1::ST
        local y2::ST
        local z1::ST
        local z2::ST

        # compute b = - (Y-AV)
        for i in 1:S
            for k in 1:D
                y1 = 0
                y2 = 0
                z1 = 0
                z2 = 0
                for j in 1:S
                    y1 += params.tab.qdrift.a[i,j] * $cache.V[j][k] * params.Δt
                    y2 += params.tab.qdrift.â[i,j] * $cache.V[j][k] * params.Δt
                    z1 += params.tab.pdrift.a[i,j] * $cache.F[j][k] * params.Δt
                    z2 += params.tab.pdrift.â[i,j] * $cache.F[j][k] * params.Δt
                    for l in 1:M
                        y1 += params.tab.qdiff.a[i,j] * $cache.B[j][k,l] * params.ΔW[l]
                        y2 += params.tab.qdiff.â[i,j] * $cache.B[j][k,l] * params.ΔW[l]
                        z1 += params.tab.pdiff.a[i,j] * $cache.G[j][k,l] * params.ΔW[l]
                        z2 += params.tab.pdiff.â[i,j] * $cache.G[j][k,l] * params.ΔW[l]
                    end
                end
                b[D*(  i-1)+k] = - $cache.Y[i][k] + (y1 + y2)
                b[D*(S+i-1)+k] = - $cache.Z[i][k] + (z1 + z2)
            end
        end
    end
end

"""
This function computes initial guesses for Y, Z and assigns them to int.solver.x
The prediction is calculated using an explicit integrator.
"""
function initial_guess!(int::IntegratorSIPRK{DT,TT}, sol::AtomicSolutionPSDE{DT,TT}) where {DT,TT}

    # SIMPLE SOLUTION
    # The simplest initial guess for Y, Z is 0
    # int.solver.x .= zeros(eltype(int), 2*int.params.tab.s*ndims(int))

    # USING AN EXPLICIT INTEGRATOR TO COMPUTE AN INITIAL GUESS
    # Below we use the R2 method of Burrage & Burrage to calculate
    # the internal stages at the times c[1]...c[s].
    # This approach seems to give very good approximations if the time step
    # and magnitude of noise are not too large. If the noise intensity is too big,
    # one may have to perform a few iterations of the explicit method with a smaller
    # time step, use a higher-order explicit method (e.g. CL or G5), or use
    # the simple solution above.

    local t2::TT
    local Δt_local::TT

    # When calling this function, int.params should contain the data:
    # int.params.q - the q solution at the previous time step
    # int.params.p - the p solution at the previous time step
    # int.params.t - the time of the previous step
    # int.params.ΔW- the increment of the Brownian motion for the current step

    # Evaluating the functions v and B at t,q - same for all stages
    int.params.equ.v(int.params.t, int.params.q, int.params.p, int.cache.V1)
    int.params.equ.B(int.params.t, int.params.q, int.params.p, int.cache.B1)
    int.params.equ.f(int.params.t, int.params.q, int.params.p, int.cache.F1)
    int.params.equ.G(int.params.t, int.params.q, int.params.p, int.cache.G1)

    # Calculating the positions q at the points qdrift.c[i]
    # if qdrift.c==pdrift.c, then also calculating the momenta
    for i in 1:int.params.tab.s

        # Taking the c[i] from the qdrift tableau.
        Δt_local = int.params.tab.qdrift.c[i]  * int.params.Δt
        int.cache.Δw .= int.params.tab.qdrift.c[i] .* int.params.ΔW

        simd_mult!(int.cache.ΔQ, int.cache.B1, int.cache.Δw)
        simd_mult!(int.cache.ΔP, int.cache.G1, int.cache.Δw)
        @. int.cache.Q[i] = int.params.q + 2. / 3. * Δt_local * int.cache.V1 + 2. / 3. * int.cache.ΔQ
        @. int.cache.P[i] = int.params.p + 2. / 3. * Δt_local * int.cache.F1 + 2. / 3. * int.cache.ΔP

        t2 = int.params.t + 2. / 3. * Δt_local

        int.params.equ.v(t2, int.cache.Q[i], int.cache.P[i], int.cache.V2)
        int.params.equ.B(t2, int.cache.Q[i], int.cache.P[i], int.cache.B2)
        int.params.equ.f(t2, int.cache.Q[i], int.cache.P[i], int.cache.F2)
        int.params.equ.G(t2, int.cache.Q[i], int.cache.P[i], int.cache.G2)

        #Calculating the Y's and assigning them to the array int.solver.x as initial guesses
        simd_mult!(int.cache.Y1, int.cache.B1, int.cache.Δw)
        simd_mult!(int.cache.Y2, int.cache.B2, int.cache.Δw)
        for j in 1:int.params.equ.d
            int.solver.x[(i-1)*int.params.equ.d+j] =  Δt_local*(1. / 4. * int.cache.V1[j] + 3. / 4. * int.cache.V2[j]) + 1. / 4. * int.cache.Y1[j] + 3. / 4. * int.cache.Y2[j]
        end

        # if the collocation points are the same for both q and p parts
        simd_mult!(int.cache.Z1, int.cache.G1, int.cache.Δw)
        simd_mult!(int.cache.Z2, int.cache.G2, int.cache.Δw)
        if int.params.tab.qdrift.c==int.params.tab.pdrift.c
            for j in 1:int.params.equ.d
                int.solver.x[(int.params.tab.s+i-1)*int.params.equ.d+j] =  Δt_local*(1. / 4. * int.cache.F1[j] + 3. / 4. * int.cache.F2[j]) + 1. / 4. * int.cache.Z1[j] + 3. / 4. * int.cache.Z2[j]
            end
        end
    end

    # If qdrift.c != pdrift.c, then calculating the momenta p at the points pdrift.c[i]
    if int.params.tab.qdrift.c != int.params.tab.pdrift.c

        for i in 1:int.params.tab.s
            # Taking the c[i] from the pdrift tableau.
            Δt_local = int.params.tab.pdrift.c[i]  * int.params.Δt
            int.cache.Δw .= int.params.tab.pdrift.c[i] .* int.params.ΔW

            simd_mult!(int.cache.ΔQ, int.cache.B1, int.cache.Δw)
            simd_mult!(int.cache.ΔP, int.cache.G1, int.cache.Δw)
            @. int.cache.Q[i] = int.params.q + 2. / 3. * Δt_local * int.cache.V1 + 2. / 3. * int.cache.ΔQ
            @. int.cache.P[i] = int.params.p + 2. / 3. * Δt_local * int.cache.F1 + 2. / 3. * int.cache.ΔP

            t2 = int.params.t + 2. / 3. * Δt_local

            int.params.equ.v(t2, int.cache.Q[i], int.cache.P[i], int.cache.V2)
            int.params.equ.B(t2, int.cache.Q[i], int.cache.P[i], int.cache.B2)
            int.params.equ.f(t2, int.cache.Q[i], int.cache.P[i], int.cache.F2)
            int.params.equ.G(t2, int.cache.Q[i], int.cache.P[i], int.cache.G2)

            # Calculating the Z's and assigning them to the array int.solver.x as initial guesses
            # The guesses for the Y's have already been written to x above
            simd_mult!(int.cache.Z1, int.cache.G1, int.cache.Δw)
            simd_mult!(int.cache.Z2, int.cache.G2, int.cache.Δw)
            for j in 1:int.params.equ.d
                int.solver.x[(int.params.tab.s+i-1)*int.params.equ.d+j] =  Δt_local*(1. / 4. * int.cache.F1[j] + 3. / 4. * int.cache.F2[j]) + 1. / 4. * int.cache.Z1[j] + 3. / 4. * int.cache.Z2[j]
            end
        end
    end

end

"""
Integrate PSDE with a stochastic implicit partitioned Runge-Kutta integrator.
Integrating the m-th sample path
"""
function integrate_step!(int::IntegratorSIPRK{DT,TT}, sol::AtomicSolutionPSDE{DT,TT}) where {DT,TT}
    # update nonlinear solver parameters from cache
    update_params!(int, sol)

    # compute initial guess and assign to int.solver.x
    initial_guess!(int, sol)

    # reset cache
    reset!(sol, timestep(int))

    # call nonlinear solver
    solve!(int.solver)

    # print solver status
    print_solver_status(int.solver.status, int.solver.params)

    # check if solution contains NaNs or error bounds are violated
    check_solver_status(int.solver.status, int.solver.params)

    # compute the drift vector field and the diffusion matrix at internal stages
    compute_stages!(int.solver.x, int.cache.Q, int.cache.P, int.cache.V, int.cache.F, int.cache.B, int.cache.G, int.cache.Y, int.cache.Z, int.params)

    # compute final update
    update_solution!(sol.q, sol.p, int.cache.V, int.cache.F, int.cache.B, int.cache.G,
                    int.params.tab.qdrift.b, int.params.tab.qdrift.b̂, int.params.tab.qdiff.b, int.params.tab.qdiff.b̂,
                    int.params.tab.pdrift.b, int.params.tab.pdrift.b̂, int.params.tab.pdiff.b, int.params.tab.pdiff.b̂,
                    int.params.Δt, int.params.ΔW, int.cache.Δy, int.cache.Δz)
end


text:
using Test
import .TensorTrains:one_rdm,two_rdm

@testset "Orbital RDM" begin
    C = randn(2,2,2,2)
    C = 1/norm(C)*C
    C_tt =ttv_decomp(C)
    #one_rdm
    γ1 = one_rdm(C_tt)
    i = rand(1:4)
    A = zeros(2,2)
    A[1,1] = dot(selectdim(C,i,1),selectdim(C,i,1))
    A[2,2] = 1-A[1,1]
    @test(isapprox(γ1[i,:,:],A,atol=1e-12))
    #two_rdm
    γ2 = two_rdm(C_tt;fermion=false)
    i = rand(1:4)
    j = rand(setdiff(1:4,i))
    i,j = min(i,j),max(i,j)
    A = zeros(2,2,2,2)
    A[1,1,1,1] = norm(selectdim(selectdim(C,j,1),i,1))^2
    A[1,2,1,2] = norm(selectdim(selectdim(C,j,2),i,1))^2
    A[2,1,2,1] = norm(selectdim(selectdim(C,j,1),i,2))^2
    A[2,2,2,2] = norm(selectdim(selectdim(C,j,2),i,2))^2
    A[1,2,2,1] = dot(selectdim(selectdim(C,j,2),i,1),selectdim(selectdim(C,j,1),i,2))
    A[2,1,1,2] = A[1,2,2,1]
    @test(isapprox(γ2[i,j,:,:,:,:],A,atol=1e-12))
end

max_stars_repo_name: msdupuy/Tensor-Train-Julia

max_issues_repo_name: msdupuy/Tensor-Train-Julia

max_forks_repo_name: msdupuy/Tensor-Train-Julia


text:
# Memory copy 3D to return T_peak
using AMDGPU

function copy3D!(T2, T, Ci)
    ix = (workgroupIdx().x - 1) * workgroupDim().x + workitemIdx().x  # only workgroupIdx starts at 1
    iy = (workgroupIdx().y - 1) * workgroupDim().y + workitemIdx().y  # only workgroupIdx starts at 1
    iz = (workgroupIdx().z - 1) * workgroupDim().z + workitemIdx().z  # only workgroupIdx starts at 1
    T2[ix,iy,iz] = T[ix,iy,iz] + Ci[ix,iy,iz]
    return
end

function memcopy3D()
    # Numerics
    nx, ny, nz = 1024, 512, 512                            # Number of grid points in dimensions x, y and z
    TBSize = (32, 8, 1)
    GSize  = (nx, ny, nz)

    nt  = 100                                              # Number of time steps

    # Array initializations
    T   = ROCArray(zeros(Float64, nx, ny, nz))
    T2  = ROCArray(zeros(Float64, nx, ny, nz))
    Ci  = ROCArray(zeros(Float64, nx, ny, nz))

    # Initial conditions
    Ci .= 0.5
    T  .= 1.7
    T2 .= T

    t_tic = 0.0
    # Time loop
    for it = 1:nt
        if (it == 11) t_tic=time() end  # Start measuring time.
        wait( @roc groupsize = TBSize gridsize = GSize copy3D!(T2, T, Ci) )
        T, T2 = T2, T
    end
    t_toc=time()-t_tic

    # Performance
    @show A_eff = (2*1+1)*1/1e9*nx*ny*nz*sizeof(Float64)      # Effective main memory access per iteration [GB] (Lower bound of required memory access: T has to be read and written: 2 whole-array memaccess; Ci has to be read: : 1 whole-array memaccess)
    t_it  = t_toc/(nt-10)                               # Execution time per iteration [s]
    T_eff = A_eff/t_it                                  # Effective memory throughput [GB/s]
    println("time_s=$t_toc T_eff=$T_eff")
end

memcopy3D()

max_stars_repo_name: luraess/amdgpu-test

max_issues_repo_name: luraess/amdgpu-test

max_forks_repo_name: luraess/amdgpu-test