text:
## Todo

import LinearAlgebra.checksquare
import Base.eltype
import Base.length
import Base.ndims
import Base.*
import LinearAlgebra.mul!
import Base.size
import KahanSummation

"""
A simple function that doesn't report any output
"""
function _noiterfunc(iter::Int, x::Vector{T}) where T
end

"""
`_applyv!`
--------

The fully generic function to compute
`x <- alpha*x + gamma*v`

Functions
---------
- `_applyv!(x, v, alpha, gamma)`

Example
-------
Internal function
"""
function _applyv! end

function _applyv!(x::Vector{T}, v::SparseVector{T},
                 alpha::T, gamma::T) where T
    @simd for i in 1:length(x)
        @inbounds x[i] *= alpha
    end
    vvals = v.nzval
    vrows = v.nzind
    @simd for j in 1:length(vrows)
        @inbounds x[vrows[j]] += gamma*vvals[j]
    end
end
 # version 0.5.0-dev

function _applyv!(x::Vector{T}, v, alpha::T, gamma::T) where T
    x *= alpha
    x += gamma*v
end

function _applyv!(x::Vector{T}, v::T, alpha::T, gamma::T) where T
    gv = gamma*v
    @simd for i in 1:length(x)
        @inbounds x[i] = alpha*x[i] + gv
    end
end

function _applyv!(x::Vector{T}, v::Vector{T}, alpha::T, gamma::T) where T
    @simd for i in 1:length(x)
        @inbounds x[i] = alpha*x[i] + gamma*v[i]
    end
end

function _applyv!(x::Vector{T}, v::SparseMatrixCSC{T,Int},
                 alpha::T, gamma::T) where T
    @simd for i in 1:length(x)
        @inbounds x[i] *= alpha
    end
    vvals = nonzeros(v)
    vrows = rowvals(v)
    @simd for j in nzrange(v,1)
        @inbounds x[vrows[j]] += gamma*vvals[j]
    end
end

"""
Convert a sparse representation into a dense vector.
"""
function _densevec(I::Vector{Int}, V::Vector{T}, n::Int) where T
    v = zeros(T,n)
    for ind in 1:length(I)
        v[ind] = V[ind]
    end
    return v
end
function _densevec(d::Dict{Int,T}, n::Int) where T
    v = zeros(T,n)
    for (ind,val) in d
        v[ind] = val
    end
    return v
end

"""
Return true if any element is negative.
"""
function _check_negative(v)
    for val in v
        if val < 0
            return true
        end
    end
    return false
end

"""
`pagerank_power!`
----------------

**Expert interface**

This function computes the strongly personalized PageRank
vector of a column sub-stochastic matrix P.

This is a powerful internal function and you most likely don't want
to use it unless you know what you are doing.

This call is duck-typed so that it can work with arbitrary input
matrix types as well as arbitrary vectors v. This enables it to
be efficient in the case of sparse or dense vectors v.

The solution returned will satisfy the strongly-personalized
PageRank equation to the given tolerance accuracy in the 1-norm.
Formally, it'll

This call allocates no extra memory and only uses the memory
included with the vectors x and y.

Functions
---------
- `pagerank_power!{T<Float}!(x::Vector{T}, y::Vector{T},
        P, alpha::T, v, tol::T, maxiter::Int, iterfunc::Function)`

Inputs
------
- `x`: a vector (of length n, where there are n nodes of the graph)
  that will store
- `y`: a vector that is used as part of the matrix-vector multiplication
  and the iteration procedure. It is just extra memory.
- `P`: a duck typed matrix to apply the stochastic operator
      the type P must support P*x and P*y
- `alpha`: the teleportation parameter, alpha must be between 0 and 1.
  This is the probability that the model follows the graph (so the
  problem is hard when alpha is close to 1).
-  `v`: a duck typed vector to apply the personalization
        the type v must support x += v where x is a Vector{T}
        examples of v include a scalar, a sparsevec, or a Vector{T}
- `tol`: the solution tolerance in the error 1-norm.
- `maxiter`: the maximum number of iterations
- `iterfunc`: a function that will be called on each iteration

Returns
-------
- `x`: The PageRank vector (just another reference to the input x)

Example
-------
~~~~

~~~~
"""
function pagerank_power!(x::Vector{T}, y::Vector{T},
    P, alpha::T, v, tol::T,
    maxiter::Int, iterfunc::Function) where T
    ialpha = 1.0/(1.0-alpha)
    xinit = x
    _applyv!(x,v,0.,1.0) # iteration number 0
    iterfunc(0,x)
    for iter=1:maxiter
        mul!(y,P,x)
        gamma = 1.0-alpha*KahanSummation.sum_kbn(y)

        delta = 0.
        _applyv!(y,v,alpha,gamma)
        @simd for i=1:length(x)
            @inbounds delta += abs(y[i] - x[i]) # TODO implement Kahan summation
        end
        x,y = y,x # swap
        iterfunc(iter,x)
        if delta*ialpha < tol
            break
        end
    end
    if !(x === xinit) # x is not xinit, so we need to copy it over
        xinit[:] .= x
    end
    xinit # always return xinit
end

"""
`stochastic_mult`
=================

**Expert interface**

This function computes `output` = \$A^T (b .* d)\$, i.e. the matrix \$A\$
tranposed times the vector \$b\$ element-wise multiplied by elements in \$d\$.
The idea is that \$d\$ is the inverse of the row-sums of \$A\$ so that this
operation computes the (column)-stochastic (or sub-stochastic) product
`output` = \$P b\$ where \$P_{i,j} = A_{j,i} / \$ `degree(j)`.

This function is used by the StochasticMult operators.

Functions
---------
- `stochastic_mult!!(output::StridedVecOrMat, A::SparseMatrixCSC,
        b::StridedVecOrMat, d::StridedVecOrMat)`
- `stochastic_mult!!(output::StridedVecOrMat, A::MatrixNetwork,
        b::StridedVecOrMat, d::StridedVecOrMat)`

Inputs
------
- `output`: See above
- `A`: See above
- `b`: See above
- `d`: the vector that will be multipled against b in-place

Returns
-------
- `output`: the output input is returned as well (as well as updated in place)
"""
function stochastic_mult end

function stochastic_mult!(output::StridedVecOrMat, A::SparseMatrixCSC, b::StridedVecOrMat,
    d::StridedVecOrMat)

    size(A,1) == size(b,1) || throw(DimensionMismatch())
    size(A,1) == size(output,1) || throw(DimensionMismatch())
    size(output,2) == size(b,2) || throw(DimenensionMismatch())
    nzv = A.nzval
    rv = A.rowval
    fill!(output, zero(eltype(output)))
    for col=1:A.n
        for k=1:size(output,2)
            tmp = zero(eltype(output))
            @inbounds for j=A.colptr[col]:(A.colptr[col+1]-1)
                tmp += (nzv[j])*b[rv[j],k]*d[rv[j]]
            end
            output[col,k] += tmp
        end
    end
    output
end

function stochastic_mult!(output::StridedVecOrMat, A::MatrixNetwork, b::StridedVecOrMat,
    d::StridedVecOrMat)

    # NOTE, for matrix networks, we have already transposed
    # the matrix, so we can do the matvec the other way.
    # that is, this is actually the tranpose one :)
    size(A,1) == size(b,1) || throw(DimensionMismatch())
    size(A,1) == size(output,1) || throw(DimensionMismatch())
    size(output,2) == size(b,2) || throw(DimenensionMismatch())
    fill!(output, zero(eltype(output)))
    # TODO fix the variable names here, they describe
    # the transpose matrix
    nzv = A.vals
    rv = A.ci
    for col=1:A.n
        for k=1:size(output,2)
            xj = b[col,k]*d[col]
            @inbounds for j=A.rp[col]:(A.rp[col+1]-1)
                output[rv[j],k] += nzv[j]*xj
            end
        end
    end
    output
end

"""
`SparseMatrixStochasticMult`
============================

**Expert interface**

This type is the result of creating
an implicit stochastic operator for
a sparse matrix.
"""
mutable struct SparseMatrixStochasticMult
    d::Vector{Float64}
    A::SparseMatrixCSC
end

eltype(op::SparseMatrixStochasticMult) = Float64
size(op::SparseMatrixStochasticMult) = size(op.A)
ndims(op::SparseMatrixStochasticMult) = 2

size(op::SparseMatrixStochasticMult, dim::Integer) = size(op.A,dim)
length(op::SparseMatrixStochasticMult) = length(op.A)

*(op::SparseMatrixStochasticMult, b::AbstractVector{S}) where {S} =
    mul!(Array{promote_type(Float64,S)}(undef, size(op.A,2)), op, b)
function mul!(output, op::SparseMatrixStochasticMult, b)
    stochastic_mult!(output, op.A, b, op.d)
end

"""
`MatrixNetworkStochasticMult`
============================

**Expert interface**

This type is the result of creating
an implicit stochastic operator for
a matrix network type.
"""
mutable struct MatrixNetworkStochasticMult
    d::Vector{Float64}
    A::MatrixNetwork
end

eltype(op::MatrixNetworkStochasticMult) = Float64
size(op::MatrixNetworkStochasticMult) = size(op.A)
ndims(op::MatrixNetworkStochasticMult) = 2

size(op::MatrixNetworkStochasticMult, dim::Integer) = size(op.A,dim)
length(op::MatrixNetworkStochasticMult) = prod(size(op.A))

*(op::MatrixNetworkStochasticMult, b::AbstractVector{S}) where {S} =
    mul!(Array{promote_type(Float64,S)}(undef, size(op.A,2)), op, b)
function mul!(output, op::MatrixNetworkStochasticMult, b)
    stochastic_mult!(output, op.A, b, op.d)
end

function _create_stochastic_mult(M::SparseMatrixCSC)
    n = checksquare(M)
    d = Array{Float64}(undef, n, 1)
    sum!(d,M) # compute out-degrees
    for i=1:length(d)
        if d[i]>0.
            d[i] = 1.0/d[i]
        elseif d[i] < 0.
            throw(DomainError(d[i]))
        end
    end
    Z = SparseMatrixStochasticMult(vec(d),M)
    return Z
end

function _create_stochastic_mult(M::MatrixNetwork)
    A = sparse_transpose(M) # this involves no work...
    n = checksquare(A)
    d = Array{Float64}(undef,1, n)
    sum!(d,A) # compute out-degrees
    for i=1:length(d)
        if d[i]>0.
            d[i] = 1.0/d[i]
        elseif d[i] < 0.
            throw(DomainError(d[i]))
        end
    end
    MatrixNetworkStochasticMult(vec(d),M)
end

"""
`pagerank`
---------

PageRank is the stationary distribution of a Markov chain defined
as follows. The behavior of the chain at state i is:
* with probability \$\alpha\$, randomly transition to an out-neighbor
  of the current node (based on a weighted probabilities if
  the graph has non-negative weights).
* with probability \$1-\alpha\$, jump to a random node chosen anywhere
  in the network.
* if there are no out-neighbors, then jump anywhere in the network
  with equal probability.

The solution satisfies a linear system of equations. This function
will solve that linear system to a 1-norm error of `tol` where
`tol` is chosen by default to be the machine precision.

The solution is always a probability distribution.

Functions
---------
* `x = pagerank(A::SparseMatrixCSC{V,Int},alpha::Float64)`
* `x = pagerank(A::MatrixNetwork{V},alpha::Float64)`
* `x = pagerank(A,alpha,eps::Float64)` specifies solution tolerance too.

Inputs
------
- `A`: The sparse matrix or matrix network that you wish to use
  to compute PageRank. In the case of a sparse matrix, it must
  have non-negative values and the values will impact the
  computation as we will compute a stochastic normalization
  as part of the algorithm.
- `alpha`: the teleportation parameter given above.
- `tol`: the tolerance, the default choice is machine precision
  for floating point, which is more than enough for almost all
  applications.

Examples
--------
~~~~
pagerank(A,alpha)
            # return the standard PageRank vector
            # with uniform teleportation and alpha=0.85
            # computed to machine precision
~~~~
"""
function pagerank end

function pagerank(A,alpha::Float64)
    pagerank(A,alpha,eps(Float64))
end

function pagerank(A,alpha::Float64,tol::Float64)
    _personalized_pagerank_validated(A,alpha,1.0/size(A,1),tol)
end

"""
`seeded_pagerank`
=================

PageRank is the stationary distribution of a Markov chain defined
as follows. The behavior of the chain at state i is:
* with probability \$\alpha\$, randomly transition to an out-neighbor
  of the current node (based on a weighted probabilities if
  the graph has non-negative weights).
* with probability \$1-\alpha\$, jump according to a distribution
  called the teleportation distribution and given by a vector \$v\$.
  (In the standard case, \$v\$ is the uniform distribution over nodes,
  but this is a parameter).
* if there are no out-neighbors, then transition according to
  the teleportation distribution (this is the strongly-personalized
  problem).

When \$v\$ is the uniform vector, then we compute the same thing
as the PageRank call itself.

The solution satisfies a linear system of equations. This function
will solve that linear system to a 1-norm error of `tol` where
`tol` is chosen by default to be the machine precision.

The solution is always a probability distribution.

Functions
---------
* For any input, A can be either `A::SparseMatrixCSC{V,Int}` for some value
  type V or `A::MatrixNetwork{V}`
* `x = seeded_pagerank(A,alpha::Float64,v::Float64)`
* `x = seeded_pagerank(A,alpha::Float64,v::Int)`
* `x = seeded_pagerank(A,alpha::Float64,v::Set{Int})`
* `x = seeded_pagerank(A,alpha::Float64,v::Dict{Int,Float64})`
* `x = seeded_pagerank(A,alpha::Float64,v::SparseMatrixCSC{Float64})`
* `x = seeded_pagerank(A,alpha::Float64,v::SparseVector{Float64})`
* `x = seeded_pagerank(A,alpha::Float64,v::Vector{Float64})`
* `x = seeded_pagerank(A,alpha::Float64,v::Vector{Float64})`
* `x = seeded_pagerank(A,alpha,v,eps::Float64)` specifies solution tolerance too.

Inputs
------
- `A`: The sparse matrix or matrix network that you wish to use
  to compute PageRank. In the case of a sparse matrix, it must
  have non-negative values and the values will impact the
  computation as we will compute a stochastic normalization
  as part of the algorithm.
- `alpha`: the teleportation parameter given above.
- `v`: the teleportation distribution vector. This can be
  an integer to teleport to a single node, a set to teleport
  (uniformly) to a set of nodes, a dictionary or sparse vector
  to weight the teleportation. Or a general dense vector.
- `tol`: the tolerance that we solve the linear system to
  (this is an error guarantee)

Examples
--------
~~~
seeded_pagerank(A,alpha,5)
            # return the seeded PageRank vector
            # with teleportation to node 5
            # computed to machine precision
~~~
"""
function seeded_pagerank end

seeded_pagerank(A,alpha,v) = personalized_pagerank(A,alpha,v)
seeded_pagerank(A,alpha,v,tol) = personalized_pagerank(A,alpha,v,tol)

function personalized_pagerank end

function personalized_pagerank(A,alpha::Float64,v)
    personalized_pagerank(A,alpha::Float64,v,eps(Float64))
end

function personalized_pagerank(A,alpha::Float64,v::Float64,tol::Float64)
    if abs(v - 1.0/size(A,1)) >= eps(Float64)
        throw(DomainError(-1))
    end
    _personalized_pagerank_validated(A,alpha,v,tol)
end

function personalized_pagerank(A,alpha::Float64,v::Int,tol::Float64)
    vecv = sparsevec([v], [1.], size(A,1))
    _personalized_pagerank_validated(A,alpha,vecv,tol)
end

function personalized_pagerank(A,alpha::Float64,v::Set{Int},tol::Float64)
    if isempty(v)
        throw(ArgumentError("the teleportation vector cannot be empty"))
    end
    n = size(A,1)
    if length(v) >= n/3 # TODO validate this size
        # use a dense call
        densev = _densevec(collect(v), ones(Float64, length(v)), size(A,1))
        return personalized_pagerank!(A,alpha,densev,tol) # normalized in the next call
    else
        sparsev = sparsevec(collect(v), ones(Float64, length(v)), size(A,1))
        return personalized_pagerank!(A,alpha,sparsev,tol) # normalized in the next call
    end
end

function personalized_pagerank(A,alpha::Float64,v::Dict{Int,Float64},tol::Float64)
    n = size(A,1)
    if isempty(v)
        throw(ArgumentError("the teleportation vector cannot be empty"))
    end
    if length(v) >= n/3 # TODO validate this size
        # use a dense call
        densev = _densevec(v, size(A,1))
        return personalized_pagerank!(A,alpha,densev,tol) # normalized in the next call
    else
        sparsev = sparsevec(v, size(A,1))
        return personalized_pagerank!(A,alpha,sparsev,tol) # normalized in the next call
    end
end

personalized_pagerank(A,alpha::Float64,v::SparseMatrixCSC{Float64},tol::Float64) =
    personalized_pagerank!(A,alpha,deepcopy(v),tol)

function personalized_pagerank!(A,alpha::Float64,v::SparseMatrixCSC{Float64},tol::Float64)
    n = size(A,1)
    if isempty(v)
        throw(ArgumentError("the teleportation vector cannot be empty"))
    end
    if size(v,1) != n
        throw(ArgumentError(@sprintf("as a sparsevector, v must be n-by-1 where n=%i", n)))
    end
    # This function automatically normalizes the values.
    vals = nonzeros(v)
    valisum = 1.0/KahanSummation.sum_kbn(vals)
    for ind in eachindex(vals)
        if vals[ind] < 0.
            throw(DomainError(vals[ind]))
        end
        vals[ind] *= valisum # normalize to sum to 1
    end
    _personalized_pagerank_validated(A,alpha,v,tol)
end

personalized_pagerank(A,alpha::Float64,v::SparseVector{Float64},tol::Float64) =
    personalized_pagerank!(A,alpha,deepcopy(v),tol)

function personalized_pagerank!(A,alpha::Float64,v::SparseVector{Float64}, tol::Float64)
    n = size(A,1)
    if isempty(v)
        throw(ArgumentError("the teleportation vector cannot be empty"))
    end
    if size(v,1) != n
        throw(ArgumentError(@sprintf("as a sparsevector, v must be n-by-1 where n=%i", n)))
    end
    # This function automatically normalizes the values.
    vals = nonzeros(v)
    valisum = 1.0/KahanSummation.sum_kbn(vals)
    for ind in eachindex(vals)
        if vals[ind] < 0.
            throw(DomainError(vals[ind]))
        end
        vals[ind] *= valisum # normalize to sum to 1
    end
    _personalized_pagerank_validated(A,alpha,v,tol)
end

personalized_pagerank(A,alpha::Float64,v::Vector{Float64},tol::Float64) =
    personalized_pagerank!(A,alpha,deepcopy(v),tol)

function personalized_pagerank!(A,alpha::Float64,v::Vector{Float64},tol::Float64)
    n = size(A,1)
    if isempty(v)
        throw(ArgumentError("the teleportation vector cannot be empty"))
    end
    if size(v,1) != n
        throw(ArgumentError(@sprintf("as a sparsevector, v must be n-by-1 where n=%i", n)))
    end
    # This function automatically normalizes the values.
    valisum = 1.0/KahanSummation.sum_kbn(v)
    @inbounds for ind in eachindex(v)
        if v[ind] < 0.
            throw(DomainError(v[ind]))
        end
        v[ind] *= valisum # normalize to sum to 1
    end
    _personalized_pagerank_validated(A,alpha,v,tol)
end

function _personalized_pagerank_validated(A,alpha::Float64,v,tol::Float64)
    #x = Vector{Float64}(size(A,1)) # these should suffice, do I still get bugs?
    #y = Vector{Float64}(size(A,1))
    x = zeros(size(A,1))
    y = zeros(size(A,1))
    if alpha < 0. || alpha >= 1.
         throw(ArgumentError(@sprintf("alpha must be in [0,1), but alpha=%f",alpha)))
    end
    maxiter = 2*ceil(Int,log(tol)/log(alpha))
    return pagerank_power!(x, y, _create_stochastic_mult(A),
                alpha, v, tol, maxiter, _noiterfunc)
end

#------------------------------------------------------------------

"""
`stochastic_heat_kernel_series!`
================================
Compute the stochastic heat kernel, the result
of \$x = \\exp(-t(I - P)) s \$ where \$s\$ is a stochastic
seed vector and \$t\$ is a time parameter and
\$P\$ is a column-stochastic or sub-stochastic
matrix.

This function will return a vector \$x\$ such that
\$ \\| x - x_{\\mbox{exact}} \\|_1 \\le \$ `eps`, and
this function further guarantees
\$x_{\\mbox{exact}} - x \\ge 0\$.

The algorithm is just a direct evaluation of the power
series represented by the heat kernel operation.

** Expert interface, see `seeded_stochastic_heat_kernel`
for more user-friendly interfaces.**

Inputs
------
-`x`: The solution vector
-`y`: An intermediate vector
-`z`: Another intermediate vector
-`P`: a variable that can be used with A_mul_B!
      in order to apply the stochastic or sub-stochastic matrix
      (column)
-`t`: the value of time in the heat kernel \$0 <= t\$
-`eps`: the solution tolerance \$ 0 < eps < 1 \$
-`maxiter`: the maximum number of series terms to use

Returns
-------
-`x`: the result of the stochastic heat kernel evaluation

Example
-------
~~~~
# TODO Show an example of how to use the expert interface
~~~~

"""
function stochastic_heat_kernel_series!(
    x::Vector{T}, y::Vector{T}, z::Vector{T},
    P, t::T, v, eps::T,
    maxiter::Int) where T

    iexpt = exp(-t)
    _applyv!(y,v,0.,1.) # iteration number 0
    # scale by iexpt
    @simd for i=1:length(x)
        @inbounds x[i] = iexpt*y[i]
    end

    eps_exp_t = eps*exp(t)
    err = exp(t)-1.
    coeff = 1.

    for k=1:maxiter
        mul!(z,P,y)       # compute z = P*y
        coeff = coeff*t/k
        BLAS.axpy!(iexpt*coeff, z, x)
        y,z = z,y # swap

        err = err - coeff
        if err < eps_exp_t
            break
        end
    end
    x
end

"""
`seeded_stochastic_heat_kernel`
===============================

Compute the stochastic heat kernel, the result
of \$x = \\exp(-t(I - P)) s \$ where \$s\$ is a stochastic
seed vector and \$t\$ is a time parameter and
\$P\$ is a column-stochastic or sub-stochastic
matrix.

This function will return a vector \$x\$ such that
\$ \\| x - x_{\\mbox{exact}} \\|_1 \\le \$ `eps`, and
this function further guarantees
\$x_{\\mbox{exact}} - x \\ge 0\$.

Functions
---------
* The input A can be either a `A::SparseMatrixCSC{V,Int}` or `A::MatrixNetwork{V}`
  for some value type `V`
* `seeded_stochastic_heat_kernel(A,t::Float64,s)` provides the stochastic seed `s` in any
  of the following forms:
* `seeded_stochastic_heat_kernel(A,t,s::Float64)`: s must be the number 1/size(A,1)
  and this corresponds to using the all-ones vector.
* `seeded_stochastic_heat_kernel(A,t,s::Int)` use the vector with just a single non-zero
* `seeded_stochastic_heat_kernel(A,t,s::Set{Int})`
* `seeded_stochastic_heat_kernel(A,t,s::Dict{Int,Float64})`
* `seeded_stochastic_heat_kernel(A,t,s::SparseMatrixCSC{Float64,Int})`
* `seeded_stochastic_heat_kernel(A,t,s::SparseVector{Float64,Int})`
* `seeded_stochastic_heat_kernel(A,t,s::Vector{Float64})`
* `seeded_stochastic_heat_kernel(A,t::Float64,s,eps::Float64)`: specifies
  the solution tolerance (default `eps=eps(Float64)`)

Inputs
------
-`A`: A sparse matrix with non-negative entries or a matrix network object
-`t`: the value of time in the heat kernel \$0 <= t\$
-`eps`: the solution tolerance \$ 0 < eps < 1 \$
-`s`: the stochastic seed vector, this will be normalized to be stochastic

Returns
-------
-`x`: the result of the stochastic heat kernel evaluation

Examples
--------
~~~

~~~
"""
function seeded_stochastic_heat_kernel end

function seeded_stochastic_heat_kernel(A,t::Float64,s)
    seeded_stochastic_heat_kernel(A,t::Float64,s,eps(Float64))
end

function seeded_stochastic_heat_kernel(A,t::Float64,s::Float64,tol::Float64)
    if abs(s - 1.0/size(A,1)) >= eps(Float64)
        throw(DomainError(-1))
    end
end

function seeded_stochastic_heat_kernel(A,t::Float64,s::Int,tol::Float64)
    vecv = sparsevec([s], [1.], size(A,1))
    _seeded_heat_kernel_validated(A,t,vecv,tol)
end

function seeded_stochastic_heat_kernel(A,t::Float64,s::Set{Int},tol::Float64)
    if isempty(s)
        throw(ArgumentError("the seed vector cannot be empty"))
    end
    n = size(A,1)
    if length(s) >= n/3 # TODO validate this size
        # use a dense call
        densev = _densevec(collect(s), ones(Float64, length(s)), size(A,1))
        return seeded_stochastic_heat_kernel!(A,t,densev,tol) # normalized in the next call
    else
        sparsev = sparsevec(collect(s), ones(Float64, length(s)), size(A,1))
        return seeded_stochastic_heat_kernel!(A,t,sparsev,tol) # normalized in the next call
    end
end

function seeded_stochastic_heat_kernel(A,t::Float64,s::Dict{Int,Float64},tol::Float64)
    n = size(A,1)
    if isempty(s)
        throw(ArgumentError("the seed vector cannot be empty"))
    end
    if length(s) >= n/3 # TODO validate this size
        # use a dense call
        densev = _densevec(s, size(A,1))
        return seeded_stochastic_heat_kernel!(A,t,densev,tol) # normalized in the next call
    else
        sparsev = sparsevec(s, size(A,1))
        return seeded_stochastic_heat_kernel!(A,t,sparsev,tol) # normalized in the next call
    end
end

seeded_stochastic_heat_kernel(A,t::Float64,s::SparseMatrixCSC{Float64},tol::Float64) =
    seeded_stochastic_heat_kernel!(A,t,deepcopy(s),tol)

function seeded_stochastic_heat_kernel!(A,t::Float64,s::SparseMatrixCSC{Float64},tol::Float64)
    n = size(A,1)
    if isempty(s)
        throw(ArgumentError("the seed vector cannot be empty"))
    end
    if size(s,1) != n
        throw(ArgumentError(@sprintf("as a sparsevector, s must be n-by-1 where n=%i", n)))
    end
    s = copy(s)
    # This function automatically normalizes the values.
    vals = nonzeros(s)
    valisum = 1.0/KahanSummation.sum_kbn(vals)
    for ind in eachindex(vals)
        if vals[ind] < 0.
            throw(DomainError(vals[ind]))
        end
        vals[ind] *= valisum # normalize to sum to 1
    end
    _seeded_heat_kernel_validated(A,t,s,tol)
end

seeded_stochastic_heat_kernel(A,t::Float64,s::SparseVector{Float64},tol::Float64) =
    seeded_stochastic_heat_kernel!(A,t,deepcopy(s),tol)

function seeded_stochastic_heat_kernel!(A,t::Float64,s::SparseVector{Float64}, tol::Float64)
    n = size(A,1)
    if isempty(s)
        throw(ArgumentError("the seed vector cannot be empty"))
    end
    if size(s,1) != n
        throw(ArgumentError(@sprintf("as a sparsevector, s must be n-by-1 where n=%i", n)))
    end
    # This function automatically normalizes the values.
    vals = nonzeros(s)
    valisum = 1.0/KahanSummation.sum_kbn(vals)
    for ind in eachindex(vals)
        if vals[ind] < 0.
            throw(DomainError(vals[ind]))
        end
        vals[ind] *= valisum # normalize to sum to 1
    end
    _seeded_heat_kernel_validated(A,t,s,tol)
end
 # sparsevector

seeded_stochastic_heat_kernel(A,t::Float64,s::Vector{Float64},tol::Float64) =
    seeded_stochastic_heat_kernel!(A,t,deepcopy(s),tol)

function seeded_stochastic_heat_kernel!(A,t::Float64,s::Vector{Float64},tol::Float64)
    n = size(A,1)
    if isempty(s)
        throw(ArgumentError("the seed vector cannot be empty"))
    end
    if size(s,1) != n
        throw(ArgumentError(@sprintf("as a vector, s must be n-by-1 where n=%i", n)))
    end
    # This function automatically normalizes the values.
    valisum = 1.0/KahanSummation.sum_kbn(s)
    @inbounds for ind in eachindex(s)
        if s[ind] < 0.
            throw(DomainError(s[ind]))
        end
        s[ind] *= valisum # normalize to sum to 1
    end
    _seeded_heat_kernel_validated(A,t,s,tol)
end

"""
Returns the number of terms of a Taylor series approximation
to exp(-t(I-P)) sufficient to guarantee eps accuracy in
the one norm when P is a stochastic matrix.

Since this is actually exp(-t) exp(tP) and P is a stochastic
matrix, if we want
\$ \\| e^{-t} \\exp(tP) - e^{-t} TaylorPoly(e^t) \\|_1 \\le \\eps \$
then it is equivalent to get
\$ \\| \\exp(tP) -  TaylorPoly(e^t) \\|_1 \\le \\eps e^t  \$
and that is what this code computes.

Returns -1 if maxdeg is insufficient
"""
function _hk_taylor_degree(t::Float64, eps::Float64, maxdeg::Int)
    eps_exp_t = eps*exp(t)
    err = exp(t)-1.
    coeff = 1.
    k::Int = 0
    while err > eps_exp_t && k < maxdeg
        k += 1
        coeff = coeff*t/k
        err = err - coeff
    end
    if k==maxdeg && err > eps_exp_t
        return -1
    end
    return max(k,1)
end

"""
This internal function actually allocates the memory and calls the
heat kernel algorithm. This could change in the future based on the
type of graph provided.
"""
function _seeded_heat_kernel_validated(A,t::Float64,s,tol::Float64)
    x = zeros(size(A,1))
    y = zeros(size(A,1))
    z = zeros(size(A,1))
    if t < 0. || !isfinite(t)
         throw(ArgumentError(@sprintf("t must be in [0,infty), but t=%f",t)))
    end
    #maxiter = 2*ceil(Int,log(tol)/log(alpha))
    bigmaxiter = 10*typemax(Int32) # 21 billion...
    maxiter = _hk_taylor_degree(t, tol, bigmaxiter)
    if maxiter == -1
        throw(ArgumentError("the value of t=$t is too large and requires over $bigmaxiter iterations"))
    end
    return stochastic_heat_kernel_series!(x, y, z,
                 _create_stochastic_mult(A),
                t, s, tol, maxiter)
end

#=

Just old code

import Base.eltype

type StochasticMult{T}
    n::Int
    mul::Function
end

eltype{T}(op::StochasticMult{T}) = T

ndims(op::StochasticMult) = 2

size(op::StochasticMult, dim::Integer) = (dim == 1) ? op.n :
                                            (dim == 2) ? op.n : 1
size(op::StochasticMult) = (op.m, op.n)

length(op::StochasticMult) = op.m*op.n

*(op::StochasticMult, b) = A_mul_B(op, b)

A_mul_B{R,S}(op::StochasticMult{R}, b::AbstractVector{S}) = A_mul_B!(Array{promote_type(R,S)}(undef, op.m), op, b)

A_mul_B!(output, op::StochasticMult, b) = op.mul(output, b)

function _create_stochastic_mult{V,F}(A::MatrixNetwork{V})
    M = sparse_transpose(A)
    d = Array{F,2}(1,size(M,1))
    sum!(d,M) # compute out-degrees
    for i=1:length(d)
        if d[i]>0.
            d[i] = 1./d[i]
        end
    end
    mul = (x) -> M*(d.*x)
    P = StochasticMult{F}(size(M,1),mul)
end

_create_stochastic_mult{V,F}(A::SparseMatrixCSC{V,Int}) =
    _create_stochastic_mult{V,F}(MatrixNetwork(A))

function pagerank{V,FA<:AbstractFloat}(A::MatrixNetwork{V}, alpha::FA)
    F = promote_type(V,FA)
    return pagerank_power(_create_stochastic_mult{V,F}(A),
                F(alpha), F(1./size(A,1)), eps(one(F)))
end

type StochasticMult{T}
    n::Int
    mul::Function
end

eltype{T}(op::StochasticMult{T}) = T
ndims(op::StochasticMult) = 2
size(op::StochasticMult, dim::Integer) = (dim == 1) ? op.n :
                                            (dim == 2) ? op.n : 1
size(op::StochasticMult) = (op.m, op.n)
length(op::StochasticMult) = op.m*op.n
*(op::StochasticMult, b) = A_mul_B(op, b)

A_mul_B{R,S}(op::StochasticMult{R}, b::AbstractVector{S}) =
    A_mul_B!(Array(promote_type(R,S), op.m), op, b)
A_mul_B!(output, op::StochasticMult, b) = op.mul(output, b)

function _create_stochastic_mult(A::MatrixNetwork{V})
    M = sparse_transpose(A)
    d = Array(Float64, 1, size(M,1))
    sum!(d,M) # compute out-degrees
    for i=1:length(d)
        if d[i]>0.
            d[i] = 1.0/d[i]
        end
    end
    mul = (x) -> M*(d.*x)
    P = StochasticMult{F}(size(M,1),mul)
end

function _handle_pagerank_network(A::SparseMatrixCSC)
    n = checksquare(A)
    d = Array(Float64, n, 1)
    sum!(d,A) # compute out-degrees
    for i=1:length(d)
        if d[i]>0.
            d[i] = 1.0/d[i]
        end
    end
    M = A'
    mul = (x) -> M*(d.*x)
    return StochasticMult{F}(size(M,1),mul)
end

function _create_stochastic_mult(A::SparseMatrixCSC)
    n = checksquare(A)
    vals = nonzeros(A)
    for i=1:length(vals)

    end
    d = Array(Float64, n, 1) # needs to be n,1 to get the sum right

    sum!(d,A) # compute out-degrees
    for i=1:length(d)
        if d[i]>0.
            d[i] = 1.0/d[i]
        end
    end
end

"""
`pagerank_power`
----------------

Example
"""
function pagerank_power{F}(P, alpha::F, v, tol::F, _iterfunc=_noiterfunc)
    x = Vector{F}(A.n)
    y = Vector{F}(A.n)
    maxiter = 2*ceil(Int,log(tol)/log(alpha))

    return pagerank_power!(x, y, P, v, tol, maxiter, _iterfunc)
end

function pagerank_power{V,FA<:AbstractFloat}(A::MatrixNetwork{V}, alpha::FA)
    F = promote_type(V,FA)
    P = _create_stochastic_mult{V,F}(A)
end

"""
`pagerank`
--------

Compute a PageRank vector of a graph with adjacency matrix A. This method
solves the linear system of equations

    (I - alpha A' D+) x = (1-alpha) v

such that the solution has a prescribed one-norm error. (The default
is machine precision)

Examples
--------
~~~
A = load_matrix_network("tapir")
pagerank(A,0.85)
pagerank(A,0.85;tol=1e-2) # compute a very low-accuracy PageRank vector
v = rand(size(A,1));
pagerank(A,0.85,v/sum(v)) # personalized PageRank
v = sparsevec([5],[1.],size(A,1)) #
pagerank(A,0.85,v) # personalized PageRank with a sparse right hand side
~~~
"""
function pagerank{V,FA<:AbstractFloat}(A::MatrixNetwork{V}, alpha::FA;
            kwargs...)
    pagerank()
end
function pagerank{V,FA<:AbstractFloat}(A::MatrixNetwork{V}, alpha::FA, v;
            kwargs...)
end
function pagerank{V,FA<:AbstractFloat}(A::SparseMatrixCSC{V}, alpha::FA)
    F = promote_type(V,FA)
    return pagerank_power(_create_stochastic_mult{V,F}(A),
                F(alpha), F(1.0/size(A,1)), eps(one(F)))
end

=#

#=

"""
This function computes the strongly personalized PageRank
vector of a column sub-stochastic matrix P using the
Neumann series expansion. It computes exactly

min(maxiter, ...)

pagerank_neumann
    x - the solution vector, initialized to v
    y - an extra vector of memory
    P - a duck typed matrix to apply the stochastic operator
        the type P must support size(P), eltype(P), P*x
    alpha - the value of alpha in PageRank
    tol - the solution tolerance
    maxiter - the maximum number of iterations
    iterfunc -
"""
function pagerank_neumann!{T}(x::Vector{T}, y::Vector{T},
    P, alpha::T, v, tol::T,
    maxiter::Int, iterfunc::Function)

    iterfunc(0, x)

    for iter=1:maxiter
        A_mul_B!(y,P,x)
        delta = 0.
        @simd for i=1:length(y)
            @inbounds y[i] *= alpha
            @inbounds delta += abs(y[i] - x[i])
        end
        x,y = y,x
        iterfunc(iter,x)
    end
    x
end

=#


text:
using Makie

function main()
    t = 0:0.01π:6π
    x = cos.(t)
    y = sin.(t)
    z = t/6pi
    s = Scene()
    lines!(s,x,y,z)
    scatter!(s,x,y,z)
    s
end

s=main()
s |> display

ext: jl

max_stars_repo_name: terasakisatoshi/juliaExer

max_issues_repo_name: terasakisatoshi/juliaExer

max_forks_repo_name: terasakisatoshi/juliaExer


text:
using DynamicalBilliards, InteractiveDynamics, GLMakie

N = 100
colors = :dense
colors = [GLMakie.RGBAf0(i/N, 0, 1 - i/N, 0.25) for i in 1:N]

# Uncomment any of the following to get the billiard you want:
bd = billiard_stadium()
bd = billiard_mushroom()
# bd = billiard_hexagonal_sinai(0.5, 1.0)
# bd = billiard_sinai(0.25f0, 1f0, 1f0)
# bd = Billiard(Antidot(Float32[0, 0], 0.5f0, false))
# bd, = billiard_logo(T = Float32)

ps = [MagneticParticle(1, 0.6 + 0.0005*i, 0, 1) for i in 1:N]
ps = [Particle(1, 0.6 + 0.00005*i, 0) for i in 1:N]
ps = particlebeam(0.8, 0.6, π/4, N, 0.01, nothing)

interactive_billiard(bd, ps; tail = 1000);

# %% make it a video
billiard_video("billiard.mp4", bd, ps; plot_particles=true, frames = 200, speed = 10);
billiard_video("billiard.mp4", bd, ps; plot_particles=false, frames = 500, speed = 30, framerate = 60);

# %% also interact with boundary map
interactive_billiard_bmap(bd);

# %% static plot of boundary map and billiard (several particles, same color)
ps = [randominside(bd) for i in 1:N]
scene, bmapax = billiard_bmap_plot(bd, ps; colors = colors, backgroundcolor = RGBf0(1,1,1))
save("static_billiard_plot.png", scene)

# %% 3b1b style video:
# Colors of 3b1b
BLUE = "#7BC3DC"
BROWN = "#8D6238"
colors = [BLUE, BROWN]
# Overwrite default color of obstacles to white (to fit with black)
InteractiveDynamics.obcolor(::Obstacle) = RGBf0(1,1,1)
bd = billiard_stadium(1.0f0, 1.0f0)
billiard_video(
    "3b1billiard.mp4", bd, 1.0, 0.6, 0;
    frames = 1200, speed = 8, backgroundcolor = :black, colors = colors
)


text:
module Day13

import ..data_dir # from parent module
input = read(joinpath(data_dir, "day13"), String)

export part1, part2

struct Paper
    coordinates::Vector{Tuple{Int64,Int64}}
    folds::Vector{Tuple{AbstractString,Int64}}
end

function Paper(input::AbstractString)
    raw_coordinates, raw_folds = split(input, "\n\n")
    coordinates = [Tuple(parse.(Int64, x)) for x in split.(split(raw_coordinates, "\n"), ",")]

    fold_equations = last.(split.(split(raw_folds, "\n"), " "))
    folds = [Tuple([string(axis), parse(Int64, amount)]) for (axis, amount) in split.(fold_equations, "=")]
    Paper(coordinates, folds)
end

function horizontal_fold(coordinates::Vector{Tuple{Int64,Int64}}, fold_line::Int64)
    # fold up
    below_fold = [coordinate for coordinate in coordinates if coordinate[2] > fold_line]
    above_fold = coordinates[(!in).(coordinates, Ref(below_fold))]

    below_fold_x = [coordinate[1] for coordinate in below_fold]
    below_fold_y = [coordinate[2] for coordinate in below_fold]
    mirrored_y_coordinates = [fold_line - (y - fold_line) for y in below_fold_y]
    mirrored = collect(zip(below_fold_x, mirrored_y_coordinates))
    unique([above_fold; mirrored])
end

function vertical_fold(coordinates::Vector{Tuple{Int64,Int64}}, fold_line::Int64)
    # fold left
    right_of_fold = [coordinate for coordinate in coordinates if coordinate[1] > fold_line]
    left_of_fold = coordinates[(!in).(coordinates, Ref(right_of_fold))]

    right_of_fold_x = [coordinate[1] for coordinate in right_of_fold]
    right_of_fold_y = [coordinate[2] for coordinate in right_of_fold]
    mirrored_x_coordinates = [fold_line - (x - fold_line) for x in right_of_fold_x]
    mirrored = collect(zip(mirrored_x_coordinates, right_of_fold_y))
    unique([left_of_fold; mirrored])
end

function fold(coordinates::Vector{Tuple{Int64,Int64}}, fold::Tuple{AbstractString,Int64})
    axis, amount = fold
    axis == "x" ? vertical_fold(coordinates, amount) : horizontal_fold(coordinates, amount)
end

function part1(input = input)
    paper = Paper(input)
    first_fold = first(paper.folds)
    new_coordinates = fold(paper.coordinates, first_fold)
    length(new_coordinates)
end

function print_coordinates(coordinates::Vector{Tuple{Int64,Int64}})
    max_x = maximum(x for (x, y) in coordinates)
    max_y = maximum(y for (x, y) in coordinates)
    display = ""
    for j = 0:max_y
        for i = 0:max_x
            if (i, j) in coordinates
                display = display * "█"
            else
                display = display * " "
            end
        end
        display = display * "\n"
    end
    print(display)
end

function part2(input = input)
    paper = Paper(input)
    coordinates = paper.coordinates
    folds = paper.folds
    while length(folds) > 0
        first_fold = popfirst!(folds)
        coordinates = fold(coordinates, first_fold)
    end
    print_coordinates(coordinates)
end

end #module


text:
@testset "variables macro" begin
    @rafvars x
    @test x == RevisedAffineForm(0, [1], 0)
    @test 1 + x ± 0.1 == RevisedAffineForm(1.0, @SVector([1.0]), 0.1)

    @rafvars x y z
    @test 1 + x + y + z ± 1 == RevisedAffineForm(1.0, @SVector([1.0, 1.0, 1.0]), 1.0)

    @rafvars x[1:4]
    p2 = 1 + x[1] + x[2] ± 2
    @test p2 == RevisedAffineForm(1.0, [1, 1, 0, 0], 2)
    @test string(p2) == "1.0 + 1.0x₁ + 1.0x₂ ± 2.0"
end

@testset "affinize" begin
    @affinize (a=1 .. 2, b=3 .. 4)
    @test a + b == RevisedAffineForm(5.0, @SVector([0.5, 0.5]), 0.0)
end


text:
## FIXME: broke from bitrot.
#@test typeof(primepi("10")) == Int128

@test typeof(primepi(10)) == Int128

@test nextprime(0) == 2
@test nextprime(2) == 3
@test prevprime(3) == 2
@test prevprime(2) == 0
@test nextprime(prevprime(nextprime(@bigint 10^100))) == nextprime(@bigint 10^100)
@test genprimes(10^6; alg = :sieve) == genprimes(10^6; alg = :next)

@test nextprime(2^20) == nextprime(BigInt(2^20))
@test prevprime(2^20) == prevprime(BigInt(2^20))


text:
module HeatCapacity

# This implements different methods to specify heat capacity
#
# If you want to add a new method here, feel free to do so.
# Remember to also export the function name in GeoParams.jl (in addition to here)

using Parameters, LaTeXStrings, Unitful
using ..Units
using GeoParams: AbstractMaterialParam, AbstractMaterialParamsStruct
import Base.show

abstract type AbstractHeatCapacity <: AbstractMaterialParam end

export  ComputeHeatCapacity,                # calculation routines
        ComputeHeatCapacity!,               # in-place routine
        ConstantHeatCapacity,               # constant
        T_HeatCapacity_Whittacker           # T-dependent heat capacity

# Constant Heat Capacity -------------------------------------------------------
"""
    ConstantHeatCapacity(cp=1050J/mol/kg)

Set a constant heat capacity:
```math
    cp  = cst
```
where ``cp`` is the thermal heat capacity [``J/kg/K``].
"""
@with_kw_noshow mutable struct ConstantHeatCapacity <: AbstractHeatCapacity
    equation::LaTeXString   =   L"cp = cst"
    cp::GeoUnit             =   1050J/kg/K                # heat capacity
end

# Calculation routine
function ComputeHeatCapacity(P, T, s::ConstantHeatCapacity)
    @unpack cp   = s
    if length(T)>1
        return Value(cp)*ones(size(T))
    else
        return Value(cp)
    end

end

# Calculation routine
function ComputeHeatCapacity!(P, T, s::ConstantHeatCapacity)
    @unpack cp   = s
    if length(T)>1
        return Value(cp)*ones(size(T))
    else
        return Value(cp)
    end

end

function ComputeHeatCapacity!(cp_array::AbstractArray{<:AbstractFloat,N},P::AbstractArray{<:AbstractFloat,N},T::AbstractArray{<:AbstractFloat,N}, s::ConstantHeatCapacity) where N
    @unpack cp   = s

    cp_array .= NumValue(cp)

    return nothing
end

# Print info
function show(io::IO, g::ConstantHeatCapacity)
    print(io, "Constant heat capacity: cp=$(g.cp.val)")
end
#-------------------------------------------------------------------------

# Temperature dependent heat capacity -------------------------------
"""
    T_HeatCapacity_Whittacker()

Sets a temperature-dependent heat capacity following the parameterization of Whittacker et al. (2009), Nature:
```math
    Cp = (a + b T - c/T^2)/m
```

where ``Cp`` is the heat capacity [``J/kg/K``], and ``a,b,c`` are parameters that dependent on the temperature `T`:
- a = 199.50 J/mol/K    if T<= 846 K
- a = 199.50 J/mol/K    if T> 846 K
- b = 0.0857J/mol/K^2   if T<= 846 K
- b = 0.0323J/mol/K^2   if T> 846 K
- c = 5e6J/mol*K        if T<= 846 K
- c = 47.9e-6J/mol*K    if T> 846 K
- molmass =   0.22178kg/mol

Note that this is slightly different than the equation in the manuscript, as Cp is in J/kg/K (rather than ``J/mol/K`` as in eq.3/4 of the paper)
"""
@with_kw_noshow mutable struct T_HeatCapacity_Whittacker <: AbstractHeatCapacity
    # Note: the resulting curve was visually compared with Fig. 2 of the paper
    equation::LaTeXString   =   L"cp = (a + b*T - c/T^2)/m"
    a0::GeoUnit             =   199.5J/mol/K                # prefactor for low T       (T<= 846 K)
    a1::GeoUnit             =   229.32J/mol/K               # prefactor for high T      (T>  846 K)
    b0::GeoUnit             =   0.0857J/mol/K^2             # linear term for low T     (T<= 846 K)
    b1::GeoUnit             =   0.0323J/mol/K^2             # linear term for high T    (T>  846 K)
    c0::GeoUnit             =   5e6J/mol*K                  # quadratic term for low T  (T<= 846 K)
    c1::GeoUnit             =   47.9e-6J/mol*K              # quadratic term for high T (T>  846 K)
    molmass::GeoUnit        =   0.22178kg/mol               # average molar mass
    Tcutoff::GeoUnit        =   846K                        # cutoff temperature
end

# Calculation routine
function ComputeHeatCapacity(P,T, s::T_HeatCapacity_Whittacker)
    @unpack a0,a1,b0,b1,c0,c1, molmass, Tcutoff   = s

    cp = zeros(size(T)).*Value(a0)./Value(molmass)

    for i in eachindex(T)
        if T[i] <= Value(Tcutoff)
            a,b,c = Value(a0),Value(b0),Value(c0)
        else
            a,b,c = Value(a1),Value(b1),Value(c1)
        end

        cp[i] = (a + b*T[i] - c/T[i]^2)/molmass
    end

    return cp
end

"""
    ComputeHeatCapacity!(cp_array::AbstractArray{<:AbstractFloat,N}, T::AbstractArray{<:AbstractFloat,N}, s::T_HeatCapacity_Whittacker) where N

Computes T-dependent heat capacity in-place
"""
function ComputeHeatCapacity!(cp_array::AbstractArray{<:AbstractFloat,N},T::AbstractArray{<:AbstractFloat,N}, s::T_HeatCapacity_Whittacker) where N
    @unpack a0,a1,b0,b1,c0,c1, molmass, Tcutoff   = s
    a0 = NumValue(a0)
    a1 = NumValue(a1)
    b0 = NumValue(b0)
    b1 = NumValue(b1)
    c0 = NumValue(c0)
    c1 = NumValue(c1)
    molmass = NumValue(molmass)
    Tcutoff = NumValue(Tcutoff)

    ind             =   (T .<= Tcutoff)
    T_local         =   view(T, ind )
    cp_array[ind]   =   (a0 .+ b0*T_local - c0./T_local.^2)/molmass

    ind             =   (T .> Tcutoff)
    T_local         =   view(T, ind )
    cp_array[ind]   =   (a1 .+ b1*T_local - c1./T_local.^2)/molmass

    return nothing
end

"""
    ComputeHeatCapacity!(cp_array::AbstractArray{<:AbstractFloat,N},P::AbstractArray{<:AbstractFloat,N},T::AbstractArray{<:AbstractFloat,N}, s::T_HeatCapacity_Whittacker) where N

Computes T-dependent heat capacity in-place
"""
function ComputeHeatCapacity!(cp_array::AbstractArray{<:AbstractFloat,N},P::AbstractArray{<:AbstractFloat,N},T::AbstractArray{<:AbstractFloat,N}, s::T_HeatCapacity_Whittacker) where N

    ComputeHeatCapacity!(cp_array,T, s)

    return nothing
end

# Print info
function show(io::IO, g::T_HeatCapacity_Whittacker)
    print(io, "T-dependent heat capacity: cp/$(g.molmass.val)=$(g.a0.val) + $(g.b0.val)*T - $(g.c0.val)/T^2 (for T<=$(g.Tcutoff.val)); ");
    print(io, " cp/$(g.molmass.val)=$(g.a1.val) + $(g.b1.val)*T - $(g.c1.val)/T^2 (for T>$(g.Tcutoff.val)) \n");
end
#-------------------------------------------------------------------------

#-------------------------------------------------------------------------
# Heat capacity from phase diagram

# to be implemented - see density implementation

#-------------------------------------------------------------------------

# Help info for the calculation routines
"""
    Cp = ComputeHeatCapacity(P, T, s:<AbstractHeatCapacity)

Returns the heat capacity `Cp` at any temperature `T` and pressure `P` using any of the heat capacity laws implemented.

Currently available:
- ConstantHeatCapacity
- T\\_HeatCapacity_Whittacker

# Example
Using dimensional units
```julia
julia> T  = (250:100:1250)*K;
julia> cp = T_HeatCapacity_Whittacker()
julia> Cp = ComputeHeatCapacity(0,T,cp)
11-element Vector{Unitful.Quantity{Float64, 𝐋² 𝚯⁻¹ 𝐓⁻², Unitful.FreeUnits{(kg⁻¹, J, K⁻¹), 𝐋² 𝚯⁻¹ 𝐓⁻², nothing}}}:
  635.4269997294616 J kg⁻¹ K⁻¹
  850.7470171764261 J kg⁻¹ K⁻¹
  962.0959598489883 J kg⁻¹ K⁻¹
 1037.5420433770641 J kg⁻¹ K⁻¹
 1097.3517921966488 J kg⁻¹ K⁻¹
 1149.2745563671706 J kg⁻¹ K⁻¹
 1157.7915050948404 J kg⁻¹ K⁻¹
 1172.3554874197264 J kg⁻¹ K⁻¹
 1186.9194697445964 J kg⁻¹ K⁻¹
  1201.483452069455 J kg⁻¹ K⁻¹
 1216.0474343943067 J kg⁻¹ K⁻¹
```

"""
ComputeHeatCapacity()

"""
    Cp = ComputeHeatCapacity(T::Any, s::AbstractHeatCapacity)

Computes heat capacity if only temperature (and not pressure) is specified
"""
ComputeHeatCapacity(T::Any, s::AbstractHeatCapacity) =  ComputeHeatCapacity(0,T, s)

"""
    Cp = ComputeHeatCapacity(s::ConstantHeatCapacity)

Returns heat capacity if we are sure that we will only employ constant heat capacity in the simulation
"""
ComputeHeatCapacity(s::ConstantHeatCapacity) =  ComputeHeatCapacity(0,0, s)

"""
    ComputeHeatCapacity!(Cp::AbstractArray{<:AbstractFloat}, Phases::AbstractArray{<:Integer}, P::AbstractArray{<:AbstractFloat},T::AbstractArray{<:AbstractFloat}, MatParam::AbstractArray{<:AbstractMaterialParamsStruct})

In-place computation of heat capacity `Cp` for the whole domain and all phases, in case a vector with phase properties `MatParam` is provided, along with `P` and `T` arrays.
This assumes that the `Phase` of every point is specified as an Integer in the `Phases` array.

"""
function ComputeHeatCapacity!(Cp::AbstractArray{<:AbstractFloat, N}, Phases::AbstractArray{<:Integer, N}, P::AbstractArray{<:AbstractFloat, N},T::AbstractArray{<:AbstractFloat, N}, MatParam::AbstractArray{<:AbstractMaterialParamsStruct, 1}) where N

    for i = 1:length(MatParam)

        if !isnothing(MatParam[i].HeatCapacity)
            # Create views into arrays (so we don't have to allocate)
            ind = Phases .== i;
            cp_local    =   view(Cp, ind )
            P_local     =   view(P , ind )
            T_local     =   view(T , ind )

            ComputeHeatCapacity!(cp_local, P_local, T_local, MatParam[i].HeatCapacity[1] )
        end

    end

end

"""
    ComputeHeatCapacity!(Cp::AbstractArray{<:AbstractFloat}, Phases::AbstractArray{<:Integer}, T::AbstractArray{<:AbstractFloat}, MatParam::AbstractArray{<:AbstractMaterialParamsStruct})

In-place computation of heat capacity `Cp` for the whole domain and all phases, in case a vector with phase properties `MatParam` is provided, along with `P` and `T` arrays.
This assumes that the `Phase` of every point is specified as an Integer in the `Phases` array.

"""
function ComputeHeatCapacity!(Cp::AbstractArray{<:AbstractFloat, N}, Phases::AbstractArray{<:Integer, N}, T::AbstractArray{<:AbstractFloat, N}, MatParam::AbstractArray{<:AbstractMaterialParamsStruct, 1}) where N

    for i = 1:length(MatParam)
        if !isnothing(MatParam[i].HeatCapacity)
            # Create views into arrays (so we don't have to allocate)
            ind = Phases .== i;
            cp_local    =   view(Cp, ind )
            T_local     =   view(T , ind )

            ComputeHeatCapacity!(cp_local, T_local, MatParam[i].HeatCapacity[1] )
        end

    end

end

"""
    ComputeDensity!(Cp::AbstractArray{<:AbstractFloat,N}, PhaseRatios::AbstractArray{<:AbstractFloat, M}, P::AbstractArray{<:AbstractFloat,N},T::AbstractArray{<:AbstractFloat,N}, MatParam::AbstractArray{<:AbstractMaterialParamsStruct})

In-place computation of heat capacity `Cp` for the whole domain and all phases, in case a vector with phase properties `MatParam` is provided, along with `P` and `T` arrays.
This assumes that the `PhaseRatio` of every point is specified as an Integer in the `PhaseRatios` array, which has one dimension more than the data arrays (and has a phase fraction between 0-1)

"""
function ComputeHeatCapacity!(Cp::AbstractArray{<:AbstractFloat, N}, PhaseRatios::AbstractArray{<:AbstractFloat, M}, P::AbstractArray{<:AbstractFloat, N},T::AbstractArray{<:AbstractFloat, N}, MatParam::AbstractArray{<:AbstractMaterialParamsStruct, 1}) where {N,M}

    if M!=(N+1)
        error("The PhaseRatios array should have one dimension more than the other arrays")
    end

    Cp .= 0.0;
    for i = 1:length(MatParam)

        Cp_local   = zeros(size(Cp))
        Fraction    = selectdim(PhaseRatios,M,i);
        if (maximum(Fraction)>0.0) & (!isnothing(MatParam[i].HeatCapacity))

            ComputeHeatCapacity!(Cp_local, P, T, MatParam[i].HeatCapacity[1] )

            Cp .= Cp .+ Cp_local.*Fraction
        end

    end

end

end


text:
import PyPlot; plt = PyPlot;
using CuArrays, CUDAnative
using CuDiff

function triangle_series(x)
    series_sum = 0f0

    for k = 1000:-1:1
        n = 2*k-1
        series_sum += (k & 1 == 1 ? +1f0 : -1f0) * CUDAnative.sin(n*x) / (n*n)
    end

    return series_sum
end

function kernel(x, y)
    i = (blockIdx().x-1) * blockDim().x + threadIdx().x
    y[i] = triangle_series(x[i])
    return nothing
end

function kernel_deriv(x, y, dy)
    i = (blockIdx().x-1) * blockDim().x + threadIdx().x
    y[i], dy[i] = CuDiff.derivative(triangle_series, x[i])
    return nothing
end

x = range(0, 4*pi, length=1024)
y = similar(x)
dy = similar(x)

d_x = CuArray(x)
d_y = similar(d_x)
d_dy = similar(d_x)

CUDAnative.@cuda threads=1024 kernel_deriv(d_x, d_y, d_dy)

copy!(y, d_y)
copy!(dy, d_dy)

plt.plot(x, y)
plt.plot(x, dy)
plt.xlabel("x")
plt.ylabel("y")


text:
export EmbeddingPatch

"""
    EmbeddingPatch()

Holds an `PVector` onto its affine patch. With this the effect
is basically the same as tracking in affine space.
"""
struct EmbeddingPatch <: AbstractAffinePatch end

struct EmbeddingPatchState{N} <: AbstractAffinePatchState{N} end

function state(::EmbeddingPatch, x::PVector{T,N}) where {T,N}
    EmbeddingPatchState{N}()
end

function onpatch!(x::PVector, state::EmbeddingPatchState)
    ProjectiveVectors.affine_chart!(x)
end

function evaluate!(u, state::EmbeddingPatchState{N}, x::PVector{T, N}) where {T, N}
    homvars = ProjectiveVectors.homvars(x)
    n = length(u) - N
    for i in 1:N
        u[n + i] = x[homvars[i]] - one(eltype(x))
    end
    nothing
end

function jacobian!(U, state::EmbeddingPatchState{N}, x::PVector{T, N}) where {T, N}
    homvars = ProjectiveVectors.homvars(x)
    n = size(U,1) - N
    for j in 1:size(U, 2), i in 1:N
        U[n + i, j] = j == homvars[i] ? one(eltype(x)) : zero(eltype(x))
    end
    nothing
end


text:
@block SimonDanisch ["short tests"] begin
    @cell begin
        scene = Scene(raw = true, camera = campixel!)
        text!(
            scene, "boundingbox",
            align = (:left, :center),
            position = (50, 50)
        )
        scale!(scene, Vec3f0(4, 1, 1))
        linesegments!(boundingbox(scene))
        offset = 0
        for a_lign in (:center, :left, :right), b_lign in (:center, :left, :right)
            global offset
            t = text!(
                "boundingbox",
                align = (a_lign, b_lign),
                position = (50, 100 + offset)
            )[end]
            linesegments!(boundingbox(t))
            offset += 50
        end
        scene
    end

    # @cell mesh(IRect(0, 0, 200, 200))

    @cell begin
        r = range(-3pi, stop = 3pi, length = 100)
        s = volume(r, r, r, (x, y, z)-> cos(x) + sin(y) + cos(z), algorithm = :iso, isorange = 0.1f0, show_axis = false)
        v2 = volume!(r, r, r, (x, y, z)-> cos(x) + sin(y) + cos(z), algorithm = :mip, show_axis = false)[end]
        translate!(v2, Vec3f0(6pi, 0, 0))
        s
    end
    @cell poly(IRect(0, 0, 200, 200), strokewidth = 20, strokecolor = :red, color = (:black, 0.4))

    @cell begin
        scene = poly([Rect(0, 0, 20, 20)])
        scatter!(Rect(0, 0, 20, 20), color = :red, markersize = 2)
    end

    @cell scatter(rand(10), color = rand(10), colormap = :Spectral)

    @cell begin
        lines(Rect(0, 0, 1, 1), linewidth = 4, scale_plot = false)
        scatter!([Point2f0(0.5, 0.5)], markersize = 1, marker = 'I', scale_plot = false)
    end

    @cell lines(rand(10), rand(10), color = rand(10), linewidth = 10)
    @cell lines(rand(10), rand(10), color = rand(RGBAf0, 10), linewidth = 10)
    @cell meshscatter(rand(10), rand(10), rand(10), color = rand(10))
    @cell meshscatter(rand(10), rand(10), rand(10), color = rand(RGBAf0, 10))
    @cell scatter(0..1, rand(10), markersize = rand(10) .* 0.1)
    @cell scatter(LinRange(0, 1, 10), rand(10))
    @cell scatter(rand(10), LinRange(0, 1, 10))

    @cell begin
        scene = Scene()
        cam2d!(scene)
        axis2d!(
            scene, IRect(Vec2f0(0), Vec2f0(1)),
            ticks = (
                ranges = ([0.1, 0.2, 0.9], [0.1, 0.2, 0.9]),
                labels = (["😸", "♡", "𝕴"], ["β ÷ δ", "22", "≙"])
            ), raw = true
        )
        center!(scene)
        scene
    end

    @cell begin
        angles = range(0, stop = 2pi, length = 20)
        pos = Point2f0.(sin.(angles), cos.(angles))
        scatter(pos, rotations = -angles , marker = '▲', scale_plot = false)
        scatter!(pos, markersize = 0.02, color = :red, scale_plot = false)
    end
    @cell begin
        using GeometryTypes
        s1 = GLNormalUVMesh(Sphere(Point3f0(0), 1f0))
        mesh(GLNormalUVMesh(Sphere(Point3f0(0), 1f0)), color = rand(50, 50))
        # ugh, bug In GeometryTypes for UVs of non unit spheres.
        s2 = GLNormalUVMesh(Sphere(Point3f0(0), 1f0))
        s2.vertices .= s2.vertices .+ (Point3f0(0, 2, 0),)
        mesh!(s2, color = rand(RGBAf0, 50, 50))
    end
    @cell heatmap(rand(50, 50), colormap = :RdBu, alpha = 0.2)

    @cell arc(Point2f0(0), 10f0, 0f0, pi, linewidth = 20)

    # themes
    @cell scatter(Theme(color = :green), rand(10), rand(10), markersize = 0.1)
    @cell scatter!(Scene(), Theme(color = :green), rand(10), rand(10), markersize = 0.01)
    @cell scatter!(Scene(), Theme(color = :green), rand(10), rand(10))
    @cell scatter(Theme(color = :green), rand(10), rand(10))
    @cell scatter(Theme(color = :green), rand(10), rand(10), markersize = 0.05)

    @cell contour(rand(10, 100))
    @cell contour(rand(100, 10))
    @cell contour(randn(100, 90), levels = 3)

    @cell contour(randn(100, 90), levels = [0.1, 0.5, 0.8])
    @cell contour(randn(100, 90), levels = [0.1, 0.5, 0.8], color = :black)
    @cell contour(randn(33, 30), levels = [0.1, 0.5, 0.9], color = [:black, :green, (:blue, 0.4)], linewidth = 2)
    @cell contour(randn(33, 30), levels = [0.1, 0.5, 0.9], colormap = :Spectral)
    @cell contour(
        rand(33, 30) .* 6 .- 3, levels = [-2.5, 0.4, 0.5, 0.6, 2.5],
        colormap = [(:black, 0.2), :red, :blue, :green, (:black, 0.2)],
        colorrange = (0.2, 0.8)
    )
    @cell lines(Circle(Point2f0(0), 1f0); scale_plot = false, resolution = (800, 1000)) # currently not working, fix is in PR

    @cell begin
        v(x::Point2{T}) where T = Point2{T}(x[2], 4*x[1])
        streamplot(v, -2..2, -2..2)
    end
end

#
# a = Point2f0.(200, 150:50:offset)
# b = Point2f0.(0, 150:50:offset)
# c = Point2f0.(500, 150:50:offset)
# yrange = 150:50:offset
# #axis2d!(range(0, stop = 500, length = length(yrange)), yrange)
#
#     import RDatasets
#     singers = RDatasets.dataset("lattice","singer")
#     x = singers[:VoicePart]
#     x2 = sort(collect(unique(x)))
#     xidx = map(x-> findfirst(x2, x), x)
#     boxplot(xidx, singers[:Height], strokewidth = 2, strokecolor = :green)
#     wireframe(Rect(0, 0, 10, 10), linewidth = 10, color = :gray)
#     using Makie
#     scene = poly([Rect(-2, -2, 9, 14)], strokewidth = 0.0, scale_plot = false, color = (:black, 0.4))
#     poly!([Rect(5, 0, -5, 10)], strokewidth = 2, strokecolor = (:gray, 0.4), color = :white, scale_plot = false)
#     histogram(rand(100000))
#     y = rand(10)
#     x = bar(1:10, y, color = y)


text:
using Base.Threads
import ClimaCore.Spaces
using CUDA

function spaceconfig(::Val{Nq}, ::Type{DA} = Array) where {Nq, DA}
    quad = Spaces.Quadratures.GLL{Nq}()
    ξ, W = Spaces.Quadratures.quadrature_points(Float64, quad)
    D = Spaces.Quadratures.differentiation_matrix(Float64, quad)
    return (DA(ξ), DA(W), DA(D))
end

# construct the coordinate array
function coordinates(::Val{Nq}, n1, n2) where {Nq}
    (ξ, W, D) = spaceconfig(Val(Nq))
    X = Array{Float64}(undef, (Nq, Nq, 2, n1, n2))
    for h2 in 1:n2, h1 in 1:n1
        x1_lo = 2pi * (2h1 - 2 - n1) / n1
        x1_hi = 2pi * (2h1 - n1) / n1
        x2_lo = 2pi * (2h2 - 2 - n2) / n2
        x2_hi = 2pi * (2h2 - n2) / n2
        for j in 1:Nq, i in 1:Nq
            X[i, j, 1, h1, h2] = (1 - ξ[i]) / 2 * x1_lo + (1 + ξ[i]) / 2 * x1_hi
            X[i, j, 2, h1, h2] = (1 - ξ[j]) / 2 * x2_lo + (1 + ξ[j]) / 2 * x2_hi
        end
    end
    return X
end

const Nstate = 4

function init_st_ref!(st_ref, X, ::Val{Nq}, parameters) where {Nq}
    @assert Nq == size(X, 1) == size(X, 2) == size(st_ref, 1) == size(st_ref, 2)
    n1 = size(X, 4)
    n2 = size(X, 5)
    @assert size(st_ref, 4) == n1
    @assert size(st_ref, 5) == n2

    @threads for h2 in 1:n2
        for h1 in 1:n1
            @inbounds for j in 1:Nq
                @simd for i in 1:Nq
                    coord = (x = X[i, j, 1, h1, h2], y = X[i, j, 2, h1, h2])
                    st = init_state(coord, parameters)
                    st_ref[i, j, 1, h1, h2] = st.ρ
                    st_ref[i, j, 2, h1, h2] = st.ρu.u1
                    st_ref[i, j, 3, h1, h2] = st.ρu.u2
                    st_ref[i, j, 4, h1, h2] = st.ρθ
                end
            end
        end
    end
    return st_ref
end

function init_st_ref(X, ::Val{Nq}, parameters) where {Nq}
    global Nstate
    n1 = size(X, 4)
    n2 = size(X, 5)
    st_ref = Array{Float64}(undef, (Nq, Nq, Nstate, n1, n2))
    init_st_ref!(st_ref, X, Val(Nq), parameters)
    return st_ref
end

struct TendencyState{DT, WJT, ST}
    ∂ξ∂x::DT
    WJv¹::WJT
    WJv²::WJT
    scratch::ST

    function TendencyState(n1, n2, ::Val{Nq}) where {Nq}
        global Nstate
        ∂ξ∂x = @SMatrix [n1/(2pi) 0; 0 n2/(2pi)]
        WJv¹ =
            MArray{Tuple{Nstate, Nq, Nq}, Float64, 3, Nstate * Nq * Nq}(undef)
        WJv² =
            MArray{Tuple{Nstate, Nq, Nq}, Float64, 3, Nstate * Nq * Nq}(undef)
        scratch =
            MArray{Tuple{Nq, Nq, Nstate}, Float64, 3, Nq * Nq * Nstate}(undef)
        return new{typeof(∂ξ∂x), typeof(WJv¹), typeof(scratch)}(
            ∂ξ∂x,
            WJv¹,
            WJv²,
            scratch,
        )
    end
end

function init_tendency_states(n1, n2, ::Val{Nq}) where {Nq}
    global Nstate
    return [TendencyState(n1, n2, Val(Nq)) for _ in 1:nthreads()]
end

getval(::Val{V}) where {V} = V

function volume_ref!(dydt_ref, y0_ref, (n1, n2, parameters, valNq, states), t)
    global Nstate
    # specialize on Nq
    Nq = getval(valNq)
    (_, W, D) = spaceconfig(Val(Nq))
    J = 2pi / n1 * 2pi / n2

    # "Volume" part
    @threads for h2 in 1:n2
        @inbounds begin
            g = parameters.g
            state = states[threadid()]
            WJv¹ = state.WJv¹
            WJv² = state.WJv²
            ∂ξ∂x = state.∂ξ∂x

            for h1 in 1:n1
                # compute volume flux
                for j in 1:Nq, i in 1:Nq
                    # 1. evaluate flux function at the point
                    y = (
                        ρ = y0_ref[i, j, 1, h1, h2],
                        ρu = Cartesian12Vector(
                            y0_ref[i, j, 2, h1, h2],
                            y0_ref[i, j, 3, h1, h2],
                        ),
                        ρθ = y0_ref[i, j, 4, h1, h2],
                    )
                    F = flux(y, parameters)
                    Fρ = F.ρ
                    Fρu1 = SVector(F.ρu[1, 1], F.ρu[1, 2])
                    Fρu2 = SVector(F.ρu[2, 1], F.ρu[2, 2])
                    Fρθ = F.ρθ

                    # 2. Convert to contravariant coordinates and store in work array
                    WJ = W[i] * W[j] * J
                    WJv¹[1, i, j], WJv²[1, i, j] = WJ * ∂ξ∂x * Fρ
                    WJv¹[2, i, j], WJv²[2, i, j] = WJ * ∂ξ∂x * Fρu1
                    WJv¹[3, i, j], WJv²[3, i, j] = WJ * ∂ξ∂x * Fρu2
                    WJv¹[4, i, j], WJv²[4, i, j] = WJ * ∂ξ∂x * Fρθ
                end

                # weak derivatives
                for j in 1:Nq, i in 1:Nq
                    WJ = W[i] * W[j] * J
                    for s in 1:Nstate
                        adj = 0.0
                        @simd for k in 1:Nq
                            # D'[i,:]*WJv¹[:,j]
                            adj += D[k, i] * WJv¹[s, k, j]
                            # D'[j,:]*WJv²[i,:]
                            adj += D[k, j] * WJv²[s, i, k]
                        end
                        dydt_ref[i, j, s, h1, h2] = adj / WJ
                    end
                end
            end
        end
    end

    return dydt_ref
end

function add_face_ref!(dydt_ref, y0_ref, (n1, n2, parameters, valNq), t)
    Nq = getval(valNq)
    (ξ, W, D) = spaceconfig(Val(Nq))
    J = 2pi / n1 * 2pi / n2

    # "Face" part
    sJ1 = 2pi / n1
    sJ2 = 2pi / n2

    @threads for h2 in 1:n2
        @inbounds begin
            for h1 in 1:n1
                # direction 1
                m1 = mod1(h1 - 1, n1)
                m2 = h2
                normal = SVector(-1.0, 0.0)
                for j in 1:Nq
                    sWJ = W[j] * sJ2
                    WJ⁻ = W[1] * W[j] * J
                    WJ⁺ = W[Nq] * W[j] * J

                    y⁻ = (
                        ρ = y0_ref[1, j, 1, h1, h2],
                        ρu = SVector(
                            y0_ref[1, j, 2, h1, h2],
                            y0_ref[1, j, 3, h1, h2],
                        ),
                        ρθ = y0_ref[1, j, 4, h1, h2],
                    )
                    y⁺ = (
                        ρ = y0_ref[Nq, j, 1, m1, m2],
                        ρu = SVector(
                            y0_ref[Nq, j, 2, m1, m2],
                            y0_ref[Nq, j, 3, m1, m2],
                        ),
                        ρθ = y0_ref[Nq, j, 4, m1, m2],
                    )
                    nf = roeflux(normal, (y⁻, parameters), (y⁺, parameters))

                    dydt_ref[1, j, 1, h1, h2] -= sWJ / WJ⁻ * nf.ρ
                    dydt_ref[1, j, 2, h1, h2] -= sWJ / WJ⁻ * nf.ρu[1]
                    dydt_ref[1, j, 3, h1, h2] -= sWJ / WJ⁻ * nf.ρu[2]
                    dydt_ref[1, j, 4, h1, h2] -= sWJ / WJ⁻ * nf.ρθ

                    dydt_ref[Nq, j, 1, m1, m2] += sWJ / WJ⁺ * nf.ρ
                    dydt_ref[Nq, j, 2, m1, m2] += sWJ / WJ⁺ * nf.ρu[1]
                    dydt_ref[Nq, j, 3, m1, m2] += sWJ / WJ⁺ * nf.ρu[2]
                    dydt_ref[Nq, j, 4, m1, m2] += sWJ / WJ⁺ * nf.ρθ
                end

                # direction 2
                m1 = h1
                m2 = mod1(h2 - 1, n2)
                normal = SVector(0.0, -1.0)
                for i in 1:Nq
                    sWJ = W[i] * sJ1
                    WJ⁻ = W[i] * W[1] * J
                    WJ⁺ = W[i] * W[Nq] * J

                    y⁻ = (
                        ρ = y0_ref[i, 1, 1, h1, h2],
                        ρu = SVector(
                            y0_ref[i, 1, 2, h1, h2],
                            y0_ref[i, 1, 3, h1, h2],
                        ),
                        ρθ = y0_ref[i, 1, 4, h1, h2],
                    )
                    y⁺ = (
                        ρ = y0_ref[i, Nq, 1, m1, m2],
                        ρu = SVector(
                            y0_ref[i, Nq, 2, m1, m2],
                            y0_ref[i, Nq, 3, m1, m2],
                        ),
                        ρθ = y0_ref[i, Nq, 4, m1, m2],
                    )
                    nf = roeflux(normal, (y⁻, parameters), (y⁺, parameters))

                    dydt_ref[i, 1, 1, h1, h2] -= sWJ / WJ⁻ * nf.ρ
                    dydt_ref[i, 1, 2, h1, h2] -= sWJ / WJ⁻ * nf.ρu[1]
                    dydt_ref[i, 1, 3, h1, h2] -= sWJ / WJ⁻ * nf.ρu[2]
                    dydt_ref[i, 1, 4, h1, h2] -= sWJ / WJ⁻ * nf.ρθ

                    dydt_ref[i, Nq, 1, m1, m2] += sWJ / WJ⁺ * nf.ρ
                    dydt_ref[i, Nq, 2, m1, m2] += sWJ / WJ⁺ * nf.ρu[1]
                    dydt_ref[i, Nq, 3, m1, m2] += sWJ / WJ⁺ * nf.ρu[2]
                    dydt_ref[i, Nq, 4, m1, m2] += sWJ / WJ⁺ * nf.ρθ
                end
            end
        end
    end

    return dydt_ref
end

function volume_ref_cuda!(dYdt, Y, parameters, Nq)
    Nstate = 4
    n1, n2 = size(Y, 4), size(Y, 5)

    (ξ, W, D) = spaceconfig(Val(Nq), CuArray)
    max_threads = 256
    @assert (Nq * Nq) ≤ max_threads

    @cuda threads = (Nq, Nq) blocks = (n1, n2) volume_ref_cuda_kernel!(
        dYdt,
        Y,
        parameters,
        W,
        D,
        Val(Nq),
        Val(Nstate),
    )
    return dYdt
end

function volume_ref_cuda_kernel!(
    dYdt,
    Y,
    parameters,
    W,
    D,
    ::Val{Nq},
    ::Val{Nstate},
) where {Nq, Nstate}
    h1, h2 = blockIdx().x, blockIdx().y
    i, j = threadIdx().x, threadIdx().y
    n1, n2 = size(Y, 4), size(Y, 5)
    FT = eltype(dYdt)
    J = 2pi / n1 * 2pi / n2
    ∂ξ∂x = @SMatrix [n1/(2pi) 0; 0 n2/(2pi)]

    WJv¹ = @cuStaticSharedMem FT (Nstate, Nq, Nq)
    WJv² = @cuStaticSharedMem FT (Nstate, Nq, Nq)

    # 1. evaluate flux function at the point
    y = (
        ρ = Y[i, j, 1, h1, h2],
        ρu = Cartesian12Vector(Y[i, j, 2, h1, h2], Y[i, j, 3, h1, h2]),
        ρθ = Y[i, j, 4, h1, h2],
    )

    F = flux(y, parameters)
    Fρ = F.ρ
    Fρu1 = SVector(F.ρu[1, 1], F.ρu[1, 2])
    Fρu2 = SVector(F.ρu[2, 1], F.ρu[2, 2])
    Fρθ = F.ρθ
    # 2. Convert to contravariant coordinates and store in work array
    WJ = W[i] * W[j] * J

    WJv¹[1, i, j], WJv²[1, i, j] = WJ * ∂ξ∂x * Fρ
    WJv¹[2, i, j], WJv²[2, i, j] = WJ * ∂ξ∂x * Fρu1
    WJv¹[3, i, j], WJv²[3, i, j] = WJ * ∂ξ∂x * Fρu2
    WJv¹[4, i, j], WJv²[4, i, j] = WJ * ∂ξ∂x * Fρθ
    sync_threads()
    # weak derivatives
    for s in 1:Nstate
        t = 0.0
        # D'[i,:]*WJv¹[:,j]
        for k in 1:Nq
            t += D[k, i] * WJv¹[s, k, j]
        end
        # D'[j,:]*WJv²[i,:]
        for k in 1:Nq
            t += D[k, j] * WJv²[s, i, k]
        end
        dYdt[i, j, s, h1, h2] = t / WJ
    end
    return nothing
end

function add_face_ref_cuda!(dYdt, Y, parameters, Nq)
    (ξ, W, D) = spaceconfig(Val(Nq), CuArray)
    n1, n2 = size(Y, 4), size(Y, 5)

    @cuda threads = (Nq) blocks = (n1, n2) add_face_ref_cuda_kernel!(
        dYdt,
        Y,
        parameters,
        Nq,
        W,
        D,
    )
    return dYdt
end

function add_face_ref_cuda_kernel!(dYdt, Y, parameters, Nq, W, D)
    h1, h2 = blockIdx().x, blockIdx().y
    i = j = threadIdx().x
    n1, n2 = size(Y, 4), size(Y, 5)
    FT = eltype(dYdt)

    # "Face" part
    sJ1 = 2pi / n1
    sJ2 = 2pi / n2

    J = 2pi / n1 * 2pi / n2

    # direction 1
    g1 = mod1(h1 - 1, n1)
    g2 = h2
    normal = SVector(-1.0, 0.0)

    sWJ = W[j] * sJ2
    WJ⁻ = W[1] * W[j] * J
    WJ⁺ = W[Nq] * W[j] * J

    y⁻ = (
        ρ = Y[1, j, 1, h1, h2],
        ρu = SVector(Y[1, j, 2, h1, h2], Y[1, j, 3, h1, h2]),
        ρθ = Y[1, j, 4, h1, h2],
    )
    y⁺ = (
        ρ = Y[Nq, j, 1, g1, g2],
        ρu = SVector(Y[Nq, j, 2, g1, g2], Y[Nq, j, 3, g1, g2]),
        ρθ = Y[Nq, j, 4, g1, g2],
    )
    nf = roeflux(normal, (y⁻, parameters), (y⁺, parameters))

    dYdt[1, j, 1, h1, h2] -= sWJ / WJ⁻ * nf.ρ
    dYdt[1, j, 2, h1, h2] -= sWJ / WJ⁻ * nf.ρu[1]
    dYdt[1, j, 3, h1, h2] -= sWJ / WJ⁻ * nf.ρu[2]
    dYdt[1, j, 4, h1, h2] -= sWJ / WJ⁻ * nf.ρθ

    dYdt[Nq, j, 1, g1, g2] += sWJ / WJ⁺ * nf.ρ
    dYdt[Nq, j, 2, g1, g2] += sWJ / WJ⁺ * nf.ρu[1]
    dYdt[Nq, j, 3, g1, g2] += sWJ / WJ⁺ * nf.ρu[2]
    dYdt[Nq, j, 4, g1, g2] += sWJ / WJ⁺ * nf.ρθ

    sync_threads()

    # direction 2
    g1 = h1
    g2 = mod1(h2 - 1, n2)
    normal = SVector(0.0, -1.0)

    sWJ = W[i] * sJ1
    WJ⁻ = W[i] * W[1] * J
    WJ⁺ = W[i] * W[Nq] * J

    y⁻ = (
        ρ = Y[i, 1, 1, h1, h2],
        ρu = SVector(Y[i, 1, 2, h1, h2], Y[i, 1, 3, h1, h2]),
        ρθ = Y[i, 1, 4, h1, h2],
    )
    y⁺ = (
        ρ = Y[i, Nq, 1, g1, g2],
        ρu = SVector(Y[i, Nq, 2, g1, g2], Y[i, Nq, 3, g1, g2]),
        ρθ = Y[i, Nq, 4, g1, g2],
    )
    nf = roeflux(normal, (y⁻, parameters), (y⁺, parameters))

    dYdt[i, 1, 1, h1, h2] -= sWJ / WJ⁻ * nf.ρ
    dYdt[i, 1, 2, h1, h2] -= sWJ / WJ⁻ * nf.ρu[1]
    dYdt[i, 1, 3, h1, h2] -= sWJ / WJ⁻ * nf.ρu[2]
    dYdt[i, 1, 4, h1, h2] -= sWJ / WJ⁻ * nf.ρθ

    dYdt[i, Nq, 1, g1, g2] += sWJ / WJ⁺ * nf.ρ
    dYdt[i, Nq, 2, g1, g2] += sWJ / WJ⁺ * nf.ρu[1]
    dYdt[i, Nq, 3, g1, g2] += sWJ / WJ⁺ * nf.ρu[2]
    dYdt[i, Nq, 4, g1, g2] += sWJ / WJ⁺ * nf.ρθ
    return nothing
end


text:
include("shared.jl")

@gen function covariance_prior(cur::Int)
    node_type = @trace(categorical(node_dist), (cur, :type))

    if node_type == CONSTANT
        param = @trace(uniform_continuous(0, 1), (cur, :param))
        node = Constant(param)

    # linear kernel
    elseif node_type == LINEAR
        param = @trace(uniform_continuous(0, 1), (cur, :param))
        node = Linear(param)

    # squared exponential kernel
    elseif node_type == SQUARED_EXP
        length_scale= @trace(uniform_continuous(0, 1), (cur, :length_scale))
        node = SquaredExponential(length_scale)

    # periodic kernel
    elseif node_type == PERIODIC
        scale = @trace(uniform_continuous(0, 1), (cur, :scale))
        period = @trace(uniform_continuous(0, 1), (cur, :period))
        node = Periodic(scale, period)

    # plus combinator
    elseif node_type == PLUS
        child1 = get_child(cur, 1, max_branch)
        child2 = get_child(cur, 2, max_branch)
        left = @trace(covariance_prior(child1))
        right = @trace(covariance_prior(child2))
        node = Plus(left, right)

    # times combinator
    elseif node_type == TIMES
        child1 = get_child(cur, 1, max_branch)
        child2 = get_child(cur, 2, max_branch)
        left = @trace(covariance_prior(child1))
        right = @trace(covariance_prior(child2))
        node = Times(left, right)

    # unknown node type
    else
        error("Unknown node type: $node_type")
    end

    return node
end

@gen function model(xs::Vector{Float64})
    n = length(xs)

    # sample covariance function
    covariance_fn::Node = @trace(covariance_prior(1), :tree)

    # sample diagonal noise
    noise = @trace(gamma(1, 1), :noise) + 0.01

    # compute covariance matrix
    cov_matrix = compute_cov_matrix_vectorized(covariance_fn, noise, xs)

    # sample from multivariate normal
    @trace(mvnormal(zeros(n), cov_matrix), :ys)

    return covariance_fn
end

@gen function noise_proposal(prev_trace)
    @trace(gamma(1, 1), :noise)
end

@gen function subtree_proposal(prev_trace)
    prev_subtree_node::Node = get_retval(prev_trace)
    (subtree_idx::Int, depth::Int) = @trace(pick_random_node(prev_subtree_node, 1, 0), :choose_subtree_root)
    new_subtree_node::Node = @trace(covariance_prior(subtree_idx), :subtree)
    (subtree_idx, depth, new_subtree_node)
end

function subtree_involution(trace, fwd_choices::ChoiceMap, fwd_ret::Tuple, proposal_args::Tuple)
    (subtree_idx, subtree_depth, new_subtree_node) = fwd_ret
    model_args = get_args(trace)

    # populate constraints with proposed subtree
    constraints = choicemap()
    set_submap!(constraints, :tree, get_submap(fwd_choices, :subtree))

    # populate backward assignment with choice of root
    bwd_choices = choicemap()
    set_submap!(bwd_choices, :choose_subtree_root => :recurse_left,
        get_submap(fwd_choices, :choose_subtree_root => :recurse_left))
    for depth=0:subtree_depth-1
        bwd_choices[:choose_subtree_root => :done => depth] = false
    end
    if !isa(new_subtree_node, LeafNode)
        bwd_choices[:choose_subtree_root => :done => subtree_depth] = true
    end

    # obtain new trace and discard, which contains the previous subtree
    (new_trace, weight, _, discard) = update(trace, model_args, (NoChange(),), constraints)

    # populate backward assignment with the previous subtree
    set_submap!(bwd_choices, :subtree, get_submap(discard, :tree))

    (new_trace, bwd_choices, weight)
end

function inference(xs::Vector{Float64}, ys::Vector{Float64}, num_iters::Int, callback)

    # observed data
    constraints = choicemap()
    constraints[:ys] = ys

    # generate initial trace consistent with observed data
    (trace, _) = generate(model, (xs,), constraints)

    # do MCMC
    local covariance_fn::Node
    local noise::Float64
    for iter=1:num_iters

        covariance_fn = get_retval(trace)
        noise = trace[:noise]
        callback(covariance_fn, noise)

        # do MH move on the subtree
        (trace, _) = metropolis_hastings(trace, subtree_proposal, (), subtree_involution)

        # do MH move on the top-level white noise
        (trace, _) = metropolis_hastings(trace, noise_proposal, ())
    end
    (covariance_fn, noise)
end

function experiment()

    # load and rescale the airline dataset
    (xs, ys) = get_airline_dataset()
    xs_train = xs[1:100]
    ys_train = ys[1:100]
    xs_test = xs[101:end]
    ys_test = ys[101:end]

    # set seed
    Random.seed!(0)

    # print MSE and predictive log likelihood on test data
    callback = (covariance_fn, noise) -> begin
        pred_ll = predictive_ll(covariance_fn, noise, xs_train, ys_train, xs_test, ys_test)
        mse = compute_mse(covariance_fn, noise, xs_train, ys_train, xs_test, ys_test)
        println("mse: $mse, predictive log likelihood: $pred_ll")
    end

    # do inference, time it
    @time (covariance_fn, noise) = inference(xs_train, ys_train, 1000, callback)

    # sample predictions
    pred_ys = predict_ys(covariance_fn, noise, xs_train, ys_train, xs_test)
end

experiment()

max_stars_repo_name: georgematheos/Gen

max_issues_repo_name: georgematheos/Gen

max_forks_repo_name: georgematheos/Gen


text:
using Random
using LinearAlgebra

@testset "LTMADS" begin
    all_in_range(l, arr, rng) = mapreduce(x -> x in rng, (x,y) -> x && y, arr)
    T = Float64
    #= Audet & Dennis 2006 pg. 203 generation of vector b(ℓ) =#
    @testset "b_l_generation" begin
        b = Dict(0.0 => [0.0, 1.0, 0.0])
        i = Dict(0.0 => 2)
        N = 3
        l = 0
        @test DS.b_l_generation(b,i,l,N) == (b[0.0],i[0.0])

        l = 1
        nb, ni = DS.b_l_generation(b,i,l,N)
        @test nb[ni] == 2.0 || nb[ni] == -2.0
        @test all_in_range(1.0, nb, -2^l:2^l)
    end

    #= Audet & Dennis 2006 pg. 203 generation of matrix L =#
    @testset "L_generation" begin
        l = 2
        N = 3
        mat = DS.L_generation(N, l)
        #Verify determinant
        @test abs(det(mat)) == 2^(l*(N-1))
        #Upper triangle is zero
        @test all([mat[i,j] == 0 for i=1:N-1, j=1:N-1 if i<j])
        #Lower triangle has all value in range -2ˡ+1:2ˡ-1
        @test all_in_range(l, [mat[i,j] for i=1:N-1, j=1:N-1 if i>j], -2^l+1:2^l-1)
        #Diagonal has all value equal to ±2ˡ
        @test all_in_range(l, [mat[i,j] for i=1:N-1, j=1:N-1 if i==j], [-2^l,2^l])
    end

    #= Audet & Dennis 2006 pg. 203 generation of matrix B =#
    @testset "B_generation" begin
        l = 3; N = 4
        L = DS.L_generation(N, l)
        b, I = DS.b_l_generation(Dict{T,Vector{T}}(), Dict{T,Int64}(), l, N)
        p = shuffle!(setdiff(1:N, I))

        B = DS.B_generation(N, I, b, L, perm=p)
        @test all([B[p[i],1:N-1] == L[i,:] for i=1:N-1])
        @test all(B[I, 1:N-1] .== 0 )
        @test B[:,N] == b
        @test abs(det(B)) == 2^(l * N)
    end

    @testset "Match_Example" begin
        b = [-3, 2, 4, -1, 0]
        I = 3
        L = [-4.0  0.0  0.0 0.0;
              3.0  4.0  0.0 0.0;
             -1.0  2.0 -4.0 0.0;
              1.0 -2.0  0.0 4.0]
        p = [4,1,2,5]
        q = [5,1,3,2,4]
        N = 5

        B = [ 3.0  4.0  0.0 0.0 -3.0;
             -1.0  2.0 -4.0 0.0  2.0;
              0.0  0.0  0.0 0.0  4.0;
             -4.0  0.0  0.0 0.0 -1.0;
              1.0 -2.0  0.0 4.0  0.0]
        _B = DS.B_generation(N, I, b, L, perm=p)
        @test _B == B

        B′ = [ 4.0 0.0  0.0 -3.0  3.0;
              2.0 0.0 -4.0  2.0 -1.0;
              0.0 0.0  0.0  4.0  0.0;
              0.0 0.0  0.0 -1.0 -4.0;
             -2.0 4.0  0.0  0.0  1.0]
        _B′ = DS.B′_generation(B, N, perm=q)
        @test _B′ == B′

        Dₖmin = [ 4.0 0.0  0.0 -3.0  3.0 -4.0;
                  2.0 0.0 -4.0  2.0 -1.0  1.0;
                  0.0 0.0  0.0  4.0  0.0 -4.0;
                  0.0 0.0  0.0 -1.0 -4.0  5.0;
                 -2.0 4.0  0.0  0.0  1.0 -3.0]
        @test Dₖmin == DS.form_basis_matrix(N, B′, false)

        Dₖmax = [ 4.0 0.0  0.0 -3.0  3.0 -4.0 -0.0 -0.0  3.0 -3.0;
                  2.0 0.0 -4.0  2.0 -1.0 -2.0 -0.0  4.0 -2.0  1.0;
                  0.0 0.0  0.0  4.0  0.0 -0.0 -0.0 -0.0 -4.0 -0.0;
                  0.0 0.0  0.0 -1.0 -4.0 -0.0 -0.0 -0.0  1.0  4.0;
                 -2.0 4.0  0.0  0.0  1.0  2.0 -4.0 -0.0 -0.0 -1.0]
        @test Dₖmax == DS.form_basis_matrix(N, B′, true)
    end

    @testset "LTMADS" begin
        p = DS.DSProblem{T}(3, poll=LTMADS{T}())
        @test !isdefined(p, :objective)
        @test isdefined(p, :constraints)
        @test p.sense == DS.Min
        @test p.N == 3
        @test p.status.optimization_status_string == "Unoptimized"
        @test p.status.optimization_status == DS.Unoptimized
    end

    @testset "MeshUpdate" begin
        p = DS.DSProblem{T}(3, poll=LTMADS{T}())
        @test p.config.mesh.Δᵐ == 1
        DS.MeshUpdate!(p, DS.Unsuccessful)
        @test p.config.mesh.Δᵐ == 1/4
        DS.MeshUpdate!(p, DS.Dominating)
        @test p.config.mesh.Δᵐ == 1
        DS.MeshUpdate!(p, DS.Improving)
        @test p.config.mesh.Δᵐ == 1
        DS.MeshUpdate!(p, DS.Dominating)
        @test p.config.mesh.Δᵐ == 1
        DS.MeshUpdate!(p, DS.Unsuccessful)
        DS.MeshUpdate!(p, DS.Unsuccessful)
        @test p.config.mesh.Δᵐ == 1/16
        DS.MeshUpdate!(p, DS.Improving)
        @test p.config.mesh.Δᵐ == 1/16
        DS.MeshUpdate!(p, DS.Dominating)
        @test p.config.mesh.Δᵐ == 1/4
    end
end


text:
using Manifolds, LinearAlgebra, PGFPlotsX, Colors

out_file = joinpath(
    @__DIR__,
    "..",
    "docs",
    "src",
    "assets",
    "images",
    "retraction_illustration.png",
)
n = 100

line_width = 1 / 2
light_color = colorant"#BBBBBB"
base_linewidth = line_width
point_color = colorant"#0077BB" #Blue
main_linewidth = 1.25
tangent_color = colorant"#33BBEE" #Cyan
tangent_linewidth = main_linewidth
geo_color = colorant"#009988" #teal
geo_linewidth = 1.25
geo_marker_size = 1.5
projection_color = colorant"#EE7733"
projection_color2 = colorant"#0077BB" #Blue
projection_color3 = colorant"#33BBEE" #Cyan
projection_marker_size = 1.5

C_to_2D(p) = [real(p), imag(p)]

M = Circle(ℂ)

p = (1 + 2im) / abs(1 + 2im)
X = imag(p) - real(p) * im
qE = exp(M, p, X)
qP = project(M, p + X)

circle = [[sin(φ), cos(φ)] for φ in range(0, 2π, length=n)]
tp = @pgf Axis({
    axis_lines = "none",
    axis_equal,
    xmin = -1.7,
    xmax = 1.7,
    ymin = -.25,
    ymax = 1.4,
    scale = 1.6,
})

plain_opts = @pgf {"no markers", color = light_color, line_width = base_linewidth}

push!(tp, Plot(plain_opts, Coordinates(Tuple.(circle))))
plaint_opts_a = @pgf {"only marks", mark_size = geo_marker_size / 2, color = light_color}
push!(tp, Plot(plaint_opts_a, Coordinates(Tuple.([[0.0, 0.0]]))))

tangent_line = [C_to_2D(p - 0.8 * X), C_to_2D(p + 1.3 * X)]
push!(tp, Plot(plain_opts, Coordinates(Tuple.(tangent_line))))

tangent_opts =
    @pgf {"no markers", "->", color = tangent_color, line_width = tangent_linewidth}
tangent_vec = [C_to_2D(p), C_to_2D(p + X)]
push!(tp, Plot(tangent_opts, Coordinates(Tuple.(tangent_vec))))

geo_opts = @pgf {"no markers", "-", color = geo_color, line_width = geo_linewidth}
geo_pts = C_to_2D.(shortest_geodesic(M, p, qE, range(0, 1, length=n)))
push!(tp, Plot(geo_opts, Coordinates(Tuple.(geo_pts))))
geo_opts2 = @pgf {"only marks", mark_size = geo_marker_size, color = geo_color}
push!(tp, Plot(geo_opts2, Coordinates(Tuple.([C_to_2D(p), C_to_2D(qE)]))))

projection_opts = @pgf {
    "no markers",
    "densely dotted",
    color = projection_color,
    line_width = base_linewidth,
}
projection_line = [C_to_2D(p + X), C_to_2D(qP)]
push!(tp, Plot(projection_opts, Coordinates(Tuple.(projection_line))))
projection_opts_a =
    @pgf {"only marks", mark_size = projection_marker_size, color = projection_color}
push!(tp, Plot(projection_opts_a, Coordinates(Tuple.([C_to_2D(qP)]))))

push!(
    tp,
    raw"\node[label ={[label distance=.05cm]above:{\color{gray}$p$}}] at (axis cs:" *
    "$(real(p)),$(imag(p))) {};",
)
push!(
    tp,
    raw"\node[label ={[label distance=.1cm]above:{\color{gray}$X\in T_p\mathcal C$}}] at (axis cs:" *
    "$(real(p+X/2)),$(imag(p+X/2))) {};",
)
push!(
    tp,
    raw"\node[label ={[label distance=.05cm]left:{\color{gray}$q'=\mathrm{retr}^{\mathrm{proj}}_pX=\mathrm{proj}_{\mathcal C}(p+X)$}}] at (axis cs:" *
    "$(real(qP)),$(imag(qP))) {};",
)
push!(
    tp,
    raw"\node[label ={[label distance=.05cm]below:{\color{gray}$q=\exp_pX$}}] at (axis cs:" *
    "$(real(qE)),$(imag(qE))) {};",
)

push!(tp, raw"\node at (axis cs: 0,-.6) {\color{gray}$\mathcal C$};")

pgfsave(out_file, tp; dpi=1200)


text:
import FactoredEikonalFastMarching.getAnalytic3DeikonalSolutionAll
import FactoredEikonalFastMarching.getAnalytic2DeikonalSolutionAll

function getConstModel(n)
m = ones(Float64,tuple(n...));
return m;
end

function getInvLinearModel(top,bottom,n)
	return getLinearModel(1./(top^2),1./(bottom^2),n);;
end

function getLinearModel(top,bottom,n)
	if length(n)==2
		rho2 = linspace(top,bottom,n[2]);
		rho1 = ones(Float64,n[1]);
		rho = rho1*rho2'
	else
		rho_t = linspace(top,bottom,n[3]);
		rho = ones(Float64,tuple(n...));
		for k=1:n[3]
			rho[:,:,k].*=rho_t[k];
		end
	end
	return rho;
end

function getWedgeModel(n::Array{Int64});
h = 1./(n-1);
x = linspace(0.0,1.0,n[1]);
y = linspace(0.0,1.0,n[2]);

X = x*ones(Float64,n[2])';
Y = ones(Float64,n[1])*y';
Z = 0.25*(tanh((4*Y-X - 0.75)*20)) + 0.75;
Z[:,end - div(n[2],2)+1:end] = Z[:,div(n[2],2) : -1 : 1];
return Z.^2,getRegularMesh([0.0,1.0,0.0,1.0],n-1);
end

function waveGuideModel(n::Array{Int64})
h = 1./(n-1);
x = linspace(0.0,1.0,n[1]);
y = linspace(0.0,1.0,n[2]);

X = x*ones(Float64,n[2])';
Y = ones(Float64,n[1])*y';
Z = 1.25*(1-0.4*exp(-32*abs(X - 0.5).^2));
return 1./(Z.^2),getRegularMesh([0.0,1.0,0.0,1.0],n-1);
end

function getSmoothGaussianMedium(n::Array{Int64,1},h::Array{Float64,1})

src_kappa = zeros(Int64,2);
xsrc = zeros(Float64,2);
src_kappa[1] = div(n[1],2);
src_kappa[2] = div(n[2],2);

xsrc[1] = (src_kappa[1]-1)*h[1];
xsrc[2] = (src_kappa[2]-1)*h[2];

X1,X2 = ndgrid((0:(n[1]-1))*h[1],(0:(n[2]-1))*h[2]);

kappaSquared = exp(-4*(X1 - xsrc[1]).^2 - 8*(X2 - xsrc[2]).^2);
return kappaSquared,1;
end

# function getSmoothGaussianMedium(n::Array{Int64,1},h::Array{Float64,1})
# src = div(n,2);
# kappaSquared = [];
# T_exact = [];
# if length(n)==2
	# (T1_exact,G11_exact,G12_exact) = getSmoothFactoredModel(n,h);
	# (T0,G01,G02,L0) = getAnalytic2DeikonalSolutionAll(n,h,src);
	# G1_exact = T0.*G11_exact + G01.*T1_exact;
	# G2_exact = T0.*G12_exact + G02.*T1_exact;
	# kappaSquared = G1_exact.*G1_exact + G2_exact.*G2_exact;
	# T_exact = T0.*T1_exact;
# else
	# (T1_exact,G11_exact,G12_exact,G13_exact) = getSmoothFactoredModel3D(n,h);
	# (T0,G01,G02,G03) = getAnalytic3DeikonalSolutionAll(n,h,src);
	# G1_exact = T0.*G11_exact + G01.*T1_exact;
	# G2_exact = T0.*G12_exact + G02.*T1_exact;
	# G3_exact = T0.*G13_exact + G03.*T1_exact;
	# kappaSquared = G1_exact.*G1_exact + G2_exact.*G2_exact + G3_exact.*G3_exact;
	# T_exact = T0.*T1_exact;
# end

# return (kappaSquared),src,T_exact;
# end

function getSmoothFactoredModel(n::Array{Int64,1},h::Array{Float64,1})
xsrc = zeros(2);
src_kappa = zeros(Int64,2);
src_kappa[1] = div(n[1],2);
src_kappa[2] = div(n[2],2);

xsrc[1] = (src_kappa[1]-1)*h[1];
xsrc[2] = (src_kappa[2]-1)*h[2];

X1,X2 = ndgrid((0:(n[1]-1))*h[1],(0:(n[2]-1))*h[2]);

sigma = 1.0;
T1_exact = (exp( - (sigma*((X1 - xsrc[1]).^2) + 4*sigma*((X2-xsrc[2]).^2))) + 1)/2;
G11_exact = -2*sigma*(X1 - xsrc[1]).*exp( - (sigma*((X1 - xsrc[1]).^2) + 4*sigma*((X2-xsrc[2]).^2)))/2;
G12_exact = -8*sigma*(X2 - xsrc[2]).*exp( - (sigma*((X1 - xsrc[1]).^2) + 4*sigma*((X2-xsrc[2]).^2)))/2;

return T1_exact,G11_exact,G12_exact;
end

# function getSmoothFactoredModel3D(n::Array{Int64,1},h::Array{Float64,1})
# xsrc = zeros(3);
# src_kappa = zeros(Int64,3);
# src_kappa[1] = div(n[1],3);
# src_kappa[2] = div(n[2],4);
# src_kappa[3] = div(n[3],2);

# xsrc[1] = (src_kappa[1]-1)*h[1];
# xsrc[2] = (src_kappa[2]-1)*h[2];
# xsrc[3] = (src_kappa[3]-1)*h[3];

# X1,X2,X3 = ndgrid((0:(n[1]-1))*h[1],(0:(n[2]-1))*h[2],(0:(n[3]-1))*h[3]);

# Sigma = [0.1,0.4,0.2];
# EXPRSkewedSquared = exp(-Sigma[1]*((X1 - xsrc[1]).^2) - Sigma[2]*((X2-xsrc[2]).^2) - Sigma[3]*((X3-xsrc[3]).^2));

# T1_exact = EXPRSkewedSquared/2 + 1/2;
# G11_exact = -Sigma[1]*(X1 - xsrc[1]).*EXPRSkewedSquared;
# G12_exact = -Sigma[2]*(X2 - xsrc[2]).*EXPRSkewedSquared;
# G13_exact = -Sigma[3]*(X3 - xsrc[3]).*EXPRSkewedSquared;

# return T1_exact,G11_exact,G12_exact,G13_exact;
# end


text:
using
    Test,
    Tracking,
    GNSSSignals,
    Random,
    StaticArrays,
    TrackingLoopFilters,
    StructArrays,
    Statistics

import Unitful: MHz, kHz, Hz, s, ms, dBHz

include("code_replica.jl")
include("carrier_replica.jl")
include("downconvert.jl")
include("discriminators.jl")
include("gps_l1.jl")
include("gps_l5.jl")
include("galileo_e1b.jl")
include("boc.jl")
include("secondary_code_or_bit_detector.jl")
include("bit_buffer.jl")
include("correlator.jl")
include("tracking_state.jl")
include("tracking_results.jl")
include("tracking_loop.jl")
include("cn0_estimation.jl")


text:
turnout = dataset("Zelig", "turnout")[1:100,:] # Take a subset of rows to reduce input size

links = [Logit(), Probit(), Cloglog(), Cauchit()]

priors = [
    Prior_Ridge(),
    Prior_Laplace(),
    Prior_Cauchy(),
    Prior_TDist(),
    Prior_Uniform(),
]

for link in links
    CRRao.set_rng(StableRNG(123))
    model = @fitmodel((Vote ~ Age + Race + Income + Educate), turnout, LogisticRegression(), link)
    @test sizeof(model) > 0
end

for prior in priors
    for link in links
        CRRao.set_rng(StableRNG(123))
        model = @fitmodel((Vote ~ Age + Race + Income + Educate), turnout, LogisticRegression(), link, prior)
        @test sizeof(model) > 0
    end
end


text:
using StatsBase
using Test

### Trimming outliers

@test collect(trim([8,2,3,4,5,6,7,1], prop=0.1)) == [8,2,3,4,5,6,7,1]
@test collect(trim([8,2,3,4,5,6,7,1], prop=0.2)) == [2,3,4,5,6,7]
@test collect(trim([1,2,3,4,5,6,7,8,9], prop=0.4)) == [4,5,6]
@test collect(trim([8,7,6,5,4,3,2,1], count=1)) == [7,6,5,4,3,2]
@test collect(trim([1,2,3,4,5,6,7,8,9], count=3)) == [4,5,6]

@test_throws ArgumentError trim([])
@test_throws ArgumentError trim([1,2,3,4,5], prop=0.5)

@test collect(trim!([8,2,3,4,5,6,7,1], prop=0.1)) == [8,2,3,4,5,6,7,1]
@test collect(trim!([8,2,3,4,5,6,7,1], prop=0.2)) == [2,3,4,5,6,7]
@test collect(trim!([1,2,3,4,5,6,7,8,9], prop=0.4)) == [4,5,6]
@test collect(trim!([8,7,6,5,4,3,2,1], count=1)) == [7,6,5,4,3,2]
@test collect(trim!([1,2,3,4,5,6,7,8,9], count=3)) == [4,5,6]

@test_throws ArgumentError trim!([])
@test_throws ArgumentError trim!([1,2,3,4,5], prop=0.5)

@test collect(winsor([8,2,3,4,5,6,7,1], prop=0.1)) == [8,2,3,4,5,6,7,1]
@test collect(winsor([8,2,3,4,5,6,7,1], prop=0.2)) == [7,2,3,4,5,6,7,2]
@test collect(winsor([1,2,3,4,5,6,7,8,9], prop=0.4)) == [4,4,4,4,5,6,6,6,6]
@test collect(winsor([1,2,3,4,5,6,7,8], count=1)) == [2,2,3,4,5,6,7,7]
@test collect(winsor([8,7,6,5,4,3,2,1], count=1)) == [7,7,6,5,4,3,2,2]
@test collect(winsor([1,2,3,4,5,6,7,8,9], count=3)) == [4,4,4,4,5,6,6,6,6]

@test_throws ArgumentError winsor([])
@test_throws ArgumentError winsor([1,2,3,4,5], prop=0.5)

@test collect(winsor!([8,2,3,4,5,6,7,1], prop=0.1)) == [8,2,3,4,5,6,7,1]
@test collect(winsor!([8,2,3,4,5,6,7,1], prop=0.2)) == [7,2,3,4,5,6,7,2]
@test collect(winsor!([1,2,3,4,5,6,7,8,9], prop=0.4)) == [4,4,4,4,5,6,6,6,6]
@test collect(winsor!([8,7,6,5,4,3,2,1], count=1)) == [7,7,6,5,4,3,2,2]
@test collect(winsor!([1,2,3,4,5,6,7,8,9], count=3)) == [4,4,4,4,5,6,6,6,6]

@test_throws ArgumentError winsor!([])
@test_throws ArgumentError winsor!([1,2,3,4,5], prop=0.5)

### Variance

@test trimvar([1,1,1,1,1]) ≈ 0.0
@test trimvar([2,3,4,5,6,7,8,9], prop=0.25) ≈ 1.0

@test_throws ArgumentError trimvar([])
@test_throws ArgumentError trimvar([1,2,3,4,5], prop=0.5)

### Other

@test mean(trim([-Inf,1,2,3,4], count=1)) == 2
@test mean(winsor([-Inf,1,2,3,4], count=1)) == 2


text:
# Example of MLE estimation. The data is really drawn from the
# Logit DGP, so the model is well specified, and the MLE has
# the properties discussed in class. E.g., if you make n very
# large you should see that the estimator is very close to the
# true value of theta used to generate data

using Econometrics
include("LogitDGP.jl")
n = 30 # sample size
θ  = [0, 0.5] # true theta for generating data
(y, x) = LogitDGP(n, θ) # generate the data

# now define things for estimation
model = θ -> logit(θ, y, x)
θstart = zeros(size(x,2)) # start values for estimation

# Perform the estimation - Make sure that you examine
# the MLE estimation programs so that you see how this works
θhat, objvalue, V, converged = mleresults(model, θstart, "estimate logit model");

max_stars_repo_name: mcreel/EconometricsNotes

max_issues_repo_name: mcreel/EconometricsNotes

max_forks_repo_name: mcreel/EconometricsNotes


text:
using Test, OrdinaryDiffEq, DiffEqCallbacks

tmin = 0.1
tmax = 5.2

for tmax_problem in [tmax; Inf]
    # Test with both a finite and an infinite tspan for the ODEProblem.
    #
    # Having support for infinite tspans is one of the main reasons for implementing PeriodicCallback
    # using add_tstop!, instead of just passing in a linspace as the tstops solve argument.
    # (the other being that the length of the internal tstops collection would otherwise become
    # linear in the length of the integration time interval.
    #
    # Testing a finite tspan is necessary because a naive implementation could add tstops after
    # tmax and thus integrate for too long (or even indefinitely).

    # Dynamics: two independent single integrators:
    du = [0; 0]
    u0 = [0.; 0.]
    dynamics = (u, p, t) -> eltype(u).(du)
    prob = ODEProblem(dynamics, u0, (tmin, tmax_problem))

    # Callbacks periodically increase the input to the integrators:
    Δt1 = 0.5
    increase_du_1 = integrator -> du[1] += 1
    periodic1_initialized = Ref(false)
    initialize1 = (c, u, t, integrator) -> periodic1_initialized[] = true
    periodic1 = PeriodicCallback(increase_du_1, Δt1; initialize = initialize1)

    Δt2 = 1.
    increase_du_2 = integrator -> du[2] += 1
    periodic2 = PeriodicCallback(increase_du_2, Δt2)

    # Terminate at tmax (regardless of whether the tspan of the ODE problem is infinite).
    terminator = DiscreteCallback((u, t, integrator) -> t == tmax, terminate!)

    # Solve.
    sol = solve(prob, Tsit5(); callback = CallbackSet(terminator, periodic1, periodic2), tstops = [tmax])

    # Ensure that initialize1 has been called
    @test periodic1_initialized[]

    # Make sure we haven't integrated past tmax:
    @test sol.t[end] == tmax

    # Make sure that the components of du have been incremented the appropriate number of times.
    Δts = [Δt1, Δt2]
    expected_num_calls = map(Δts) do Δt
        floor(Int, (tmax - tmin) / Δt)
    end
    @test du == expected_num_calls

    # Make sure that the final state matches manual integration of the piecewise linear function
    foreach(Δts, sol.u[end], du) do Δt, u_i, du_i
        @test u_i ≈ Δt * sum(1 : du_i - 1) + rem(tmax - tmin, Δt) * du_i atol = 1e-5
    end
end

function fff(du,u,p,t)
  du[1] = -u[1]
  du[2] = 0
end

u0 = [2.0, 0.0]
function periodic(integrator)
  integrator.u[2] = integrator.u[1]
end
cb = PeriodicCallback(periodic, 0.1, initial_affect = true, save_positions=(true,true))
tspan = (0.0, 10.0)
p = nothing

prob = ODEProblem(fff,u0,tspan,p)
sol = solve(prob, Tsit5(), callback = cb)
@test sol.u[2] == [2.0,2.0]


text: # BitBoard version of generate moves

max_stars_repo_name: kimrin/JapaneseChess

max_issues_repo_name: kimrin/MechajyoForWindows

max_forks_repo_name: kimrin/MechajyoForWindows