text:
IVERBOSE && @info "ArrayAllez loaded in-place code for Tracker"

using .Tracker
using .Tracker: track, TrackedArray, TrackedReal, @grad, data, nobacksies # also for prod....jl

"""
It is safe to mutate the forward output of `exp!` and `exp_`,
as they keep a copy for backwards use.

    exp!!(A::TrackedArray)
The gradient function of `exp!!` mutates its backward `Δ`, no copies.
Whether or not this is safe is for you to decide.
It tends to lead to Inf problems when used inside `@btime`.
"""
exp!(A::TrackedArray) = track(exp!, A)
exp!!(A::TrackedArray) = track(exp!!, A)
exp_(A::TrackedArray) = track(exp_, A)

exp0(A::TrackedArray) = exp.(A)

@grad function exp_(A::TrackedArray)
    expA = exp_(A.data)
    expA_copy = copy(expA) # ensures that mutating output won't damage the gradient
    expA, Δ -> ( scale!(expA_copy, Δ) ,) # and we can then mutate the copy here... for 1st deriv?
end
@grad function exp!(A::TrackedArray)
    expA = exp!(A.data)
    expA_copy = copy(expA)
    expA, Δ -> ( scale!(expA_copy, Δ) ,)
end
@grad function exp!!(A::TrackedArray)
    expA = exp!(A.data)
    expA, Δ -> ( scale!(Δ, expA) ,)
end

exp_(name::Symbol, A::TrackedArray) = track(exp_, name, A)
@grad function exp_(name::Symbol, A::TrackedArray)
    expA = exp_(name, A.data)
    expA_copy = copy_(:exp_copy, expA) # opts in but has a distinct name? NOT safe
    expA, Δ -> (nothing, scale!(expA_copy, Δ) ,)
end

"""
For `log!` it is safe to mutate both the input and the forward output,
as the `inv_(A)` needed for the gradient is computed ahead of time.
For `log_` it is safe to mutate the output but not its input.

    log!!(A::TrackedArray)
Note that the gradient function of `log!!` mutates its backward `Δ`.
Even less of a good idea than `exp!!` as we must copy `inv_(A)` anyway.
"""
log!(A::TrackedArray) = track(log!, A)
log!!(A::TrackedArray) = track(log!!, A)
log_(A::TrackedArray) = track(log_, A)

log0(A::TrackedArray) = log.(A)

@grad log_(A::TrackedArray) =
    log_(A.data), Δ -> ( iscale_(Δ, A.data) ,)
@grad function log!(A::TrackedArray)
    invA = inv_(A.data)
    log!(A.data), Δ -> ( scale!(invA, Δ) ,)
end
@grad function log!!(A::TrackedArray)
    invA = inv_(A.data)
    logA = log!(A.data)
    logA, Δ -> ( scale!(Δ, invA) ,)
end

log_(name::Symbol, A::TrackedArray) = track(log_, name, A)
@grad function log_(name::Symbol, A::TrackedArray)
    invA = inv_(:log_copy, A.data)
    logA = log_(name, A.data)
    logA, Δ -> (nothing, scale!(invA, Δ) ,)
end

"""
    scale!!(A::TrackedArray, b)
This may mutate its backward `Δ`, watch out.
"""
scale!(A::TrackedArray, b) = track(scale!, A, b)
scale!(A::TrackedArray, b::Number) = track(scale!, A, b) # just to avoid ambiguty
scale!(A::Array, b::TrackedArray) = track(scale!, A, b)

scale_(A::TrackedArray, b) = track(scale_, A, b)
scale_(A::TrackedArray, b::AbstractArray) = track(scale_, A, b) # avoid an ambiguty?
scale0(A::TrackedArray, b) = A .* b

@grad scale_(A::TrackedArray, b::Number) = scale_(A.data, b), Δ -> (scale_(Δ, b), nothing)
@grad scale!(A::TrackedArray, b::Number) = scale!(A.data, b), Δ -> (scale_(Δ, b), nothing)
@grad scale!!(A::TrackedArray, b::Number) = scale!(A.data, b), Δ -> (scale!(Δ, b), nothing)

# @grad scale_(A::TrackedArray, b::TrackedReal) = scale_(A.data, b), Δ -> (scale_(Δ, b), dot(A,Δ))
# @grad scale!(A::TrackedArray, b::TrackedReal) = @error "hmm"

@grad scale_(A::TrackedArray, B::Union{Array, RVector}) = scale_(A.data, B), Δ -> (scale_(Δ, B), nothing)
@grad scale!(A::TrackedArray, B::Union{Array, RVector}) = scale!(A.data, B), Δ -> (scale_(Δ, B), nothing)
@grad scale!!(A::TrackedArray, B::Union{Array, RVector}) = scale!(A.data, B), Δ -> (scale!(Δ, B), nothing)

@grad  scale_(A::TrackedArray, B::TrackedArray) =
    scale_(A.data, B.data), Δ ->  (scale_(Δ, B), sum!(similar(B), scale_(Δ,A)) )
@grad function scale!(A::TrackedArray, B::TrackedArray)
    Ac = copy(A.data)
    scale!(A.data, B.data), Δ -> (scale_(Δ, B), sum!(similar(B), scale!(Ac,Δ)) )
end
@grad function scale!!(A::TrackedArray, B::TrackedArray)
    Ac = copy(A.data)
    scale!(A.data, B.data), function(Δ)
        ∇B = sum!(similar(B), scale!(Ac,Δ))
        (scale!(Δ, B),  ∇B)
    end
end

@grad  scale_(A::Array, B::TrackedArray) =
    scale_(A, B.data), Δ ->  (nothing, sum!(similar(B), scale_(Δ,A)) )
@grad function scale!(A::Array, B::TrackedArray)
    Ac = copy(A)
    scale!(A, B.data), Δ -> (nothing, sum!(similar(B), scale!(Ac,Δ)) )
end

iscale!(A::TrackedArray, b) = track(iscale!, A, b)
iscale_(A::TrackedArray, b) = track(iscale_, A, b)
iscale0(A::TrackedArray, b) = A ./ b

@grad iscale_(A::TrackedArray, b::Number) = iscale_(A.data, b), Δ -> (iscale_(Δ, b), nothing)
@grad iscale!(A::TrackedArray, b::Number) = iscale!(A.data, b), Δ -> (iscale!(Δ, b), nothing)

# @grad iscale_(A::TrackedArray, b::TrackedReal) = iscale_(A.data, b.data), Δ -> (iscale_(Δ, b), nothing)
# @grad iscale!(A::TrackedArray, b::TrackedReal) = iscale!(A.data, b.data), Δ -> (iscale!(Δ, b), nothing)

@grad iscale_(A::TrackedArray, B::Union{Array, RVector}) = iscale_(A.data, B), Δ -> (iscale_(Δ, B), nothing)
@grad iscale!(A::TrackedArray, B::Union{Array, RVector}) = iscale!(A.data, B), Δ -> (iscale!(Δ, B), nothing)
@grad iscale!!(A::TrackedArray, B::Union{Array, RVector}) = scale!(A.data, inv!(B)), Δ -> (scale!(Δ, B), nothing)

@grad iscale_(A::TrackedArray, B::TrackedArray) = iscale_(A.data, B.data), Δ -> (iscale_(Δ, B), sum!(similar(B), -1 .* Δ .* A ./ B .^2))
# @grad function iscale!(A::TrackedArray, B::TrackedArray)
#     Ac = copy(A.data)
#     iscale!(A.data, B.data), Δ -> (scale(Δ, B.data), sum!(similar(B), scale!(Ac,Δ)) )
# end

@grad iscale_(A::Array, B::TrackedArray) = iscale_(A, B.data), Δ -> (nothing, sum!(similar(B), -1 .* Δ .* A ./ B .^2))

"""
    inv!!(A::TrackedArray)
This may mutate its backward `Δ`, watch out.
"""
inv!(A::TrackedArray, b::Number=1) = track(inv!, A, b)
inv_(A::TrackedArray, b::Number=1) = track(inv_, A, b)
inv!!(A::TrackedArray, b::Number=1) = track(inv!!, A, b)

inv0(A::TrackedArray, b) = 1 ./ A

# @grad inv!(A::TrackedArray, b::Number=1) = inv_(A.data, b), Δ -> (-1 .* Δ .* b .* A.data .^ (-2) , nothing) # one copy
@grad function inv_(A::TrackedArray, b=1)
    invA = inv_(A.data, b)
    invA_copy = copy(invA) # don't invert twice?
    invA, Δ -> (invA_copy .= -1 .* Δ .* b .* invA_copy .^ 2 , nothing) # total one copy
end
# @grad inv!(A::TrackedArray, b::Number=1) = inv!(A.data, b), Δ -> (-1 .* Δ .* b .* A.data .^ 2 , nothing) # one copy
@grad function inv!(A::TrackedArray, b=1)
    invA = inv!(A.data, b)
    invA_copy = copy(invA) # keep gradient calc safe from downstream mutations
    invA, Δ -> (invA_copy .= -1 .* Δ .* b .* invA_copy .^ 2 , nothing) # total one copy
end
@grad inv!!(A::TrackedArray, b::Number=1) = inv!(A.data, b), Δ -> (scale!(Δ,A,A,-b), nothing)


text:
using Test
using SymDiff

@test differentiate(1, :x) == 0
@test differentiate(1.0, :x) == 0
@test differentiate(1//2, :x) == 0

ext: jl


text:
@testset "bicgstb" begin
	@testset "real matrix" begin
		A  = sprandn(100,100,.1) + 10*I
		D  = Array(diag(A))
		Af = x -> A*x
		M  = x -> D.\x
		rhs = randn(100)

		# test flag for early stopping
		xt = bicgstb(A,rhs,tol=1e-6,maxIter=3,out=2,storeInterm=true)
		@test xt[2]==-1

		# test handling of zero rhs
		xt = bicgstb(A,zeros(size(A,2)),tol=1e-6,maxIter=3,out=2)
		@test  all(xt[1].==0)
		@test xt[2]==-9

		x0 = bicgstb(Matrix(A),rhs,tol=1e-6)
		x1 = bicgstb(A,rhs,tol=1e-6,out=1)
		x2 = bicgstb(Af,rhs,tol=1e-6)
		x3 = bicgstb(Af,rhs,tol=1e-6,maxIter=100,x=randn(size(rhs)))
		x5 = bicgstb(Af,rhs,tol=1e-6,maxIter=100,M1=M)

		@test norm(A*x0[1]-rhs)/norm(rhs) < 1e-6
		@test norm(A*x1[1]-rhs)/norm(rhs) < 1e-6
		@test norm(A*x2[1]-rhs)/norm(rhs) < 1e-6
		@test norm(A*x3[1]-rhs)/norm(rhs) < 1e-6
		@test norm(A*x5[1]-rhs)/norm(rhs) < 1e-6
		@test norm(x2[1]-x1[1])/norm(x1[1]) < 1e-12
		@test norm(x3[1]-x1[1])/norm(x1[1]) < 1e-5
		@test norm(x0[1]-x1[1])/norm(x1[1]) < 1e-5
	end
	@testset "complex matrix" begin
		A  = sprandn(100,100,.1) + 10*I + im*(sprandn(100,100,.1) + 10*I )
		D  = Array(diag(A))
		Af = x -> A*x
		M  = x -> D.\x
		rhs = randn(100) + 1im * randn(100)

		x1 = bicgstb(A,rhs,tol=1e-6)
		x2 = bicgstb(Af,rhs,tol=1e-6)
		x3 = bicgstb(Af,rhs,tol=1e-6,maxIter=200,x=randn(size(rhs))+im*randn(size(rhs)))
		x4 = bicgstb(Af,rhs,tol=1e-6,maxIter=200,M1=M)

		@test norm(A*x1[1]-rhs)/norm(rhs) < 1e-6
		@test norm(A*x2[1]-rhs)/norm(rhs) < 1e-6
		@test norm(A*x3[1]-rhs)/norm(rhs) < 1e-6
		@test norm(A*x4[1]-rhs)/norm(rhs) < 1e-6
		@test norm(x2[1]-x1[1])/norm(x1[1]) < 1e-12
		@test norm(x3[1]-x1[1])/norm(x1[1]) < 1e-5

	end
end


text:
# package up, and have a coordiante transform from shift,a,b to u,1,length(s).

#include("../src/FractionalDelayFilters.jl")
#import .FractionalDelayFilters
import FractionalDelayFilters

using LinearAlgebra
using FFTW
import PyPlot

import Utilities # https://gitlab.com/RoyCCWang/utilities

PyPlot.close("all")
fig_num = 1

PyPlot.matplotlib["rcParams"][:update](["font.size" => 22, "font.family" => "serif"])

# sine for debug.
f = sin
P = LinRange(-pi, pi, 3000)
s = f.(P)

## set up.
a = P[1]
b = P[end]

S = FractionalDelayFilters.FiniteSignalType(s, a, b)
N_thirian = 15
b_thirian = Vector{Float64}(undef, N_thirian)

######## get shift coordinates from Δp, which is in units of P.
Δp = 2.0

FractionalDelayFilters.resetsignal!(S)
shift = Δp*length(S.s)/(b-a)
FractionalDelayFilters.applyshift!(S, shift, b_thirian)

# visualize.
evals_S = collect( FractionalDelayFilters.evalsignal(S, j) for j = 1:length(S.s) )

PyPlot.figure(fig_num)
fig_num += 1

PyPlot.plot(P, s, label = "s")
PyPlot.plot(P, evals_S, label = "evals_S", "--")

PyPlot.legend()
PyPlot.xlabel("ppm")
PyPlot.ylabel("real")
PyPlot.title("fractional shift with Δp = $(Δp)")

### shift back to show it is reversible.
FractionalDelayFilters.applyshift!(S, -shift, b_thirian)

# visualize.
evals_S = collect( FractionalDelayFilters.evalsignal(S, j) for j = 1:length(S.s) )

PyPlot.figure(fig_num)
fig_num += 1

PyPlot.plot(P, s, label = "s")
PyPlot.plot(P, evals_S, label = "evals_S", "--")

PyPlot.legend()
PyPlot.xlabel("ppm")
PyPlot.ylabel("real")
PyPlot.title("fractional shift with Δp = -$(Δp)")

println("Discrepancy of the shifted back sequence with the original: ", norm(s-evals_S))

max_stars_repo_name: RoyCCWang/FractionalDelayFilters

max_issues_repo_name: RoyCCWang/FractionalDelayFilters

max_forks_repo_name: RoyCCWang/FractionalDelayFilters


text:
function main()
  println("hello world")
end
main()

max_stars_repo_name: justletterh/hithere

max_issues_repo_name: justletterh/hithere

max_forks_repo_name: justletterh/hithere


text:
mutable struct Budget{Kmax_T} <: AbstractBudget
    Kmax::Kmax_T
end
Budget(; Kmax=10) = Budget(Kmax)


text:
"""
1. Velocity, pressure and strain rate analytically
f₀(x) = f_eq(ρ₀(x), u₀(x), p₀(x)) + f¹(ρ₀(x), u₀(x), p₀(x), σ₀(x))
ρ₀ = ρ₀(x)
u₀ = u₀(x)
p₀ = p₀(x)
σ₀ = σ₀(x)
"""
struct AnalyticalEquilibriumAndOffEquilibrium <: InitializationStrategy end
function initial_condition(
    ::AnalyticalEquilibriumAndOffEquilibrium,
    q::Quadrature,
    problem::FluidFlowProblem,
    x::T,
    y::T,
) where {T <: Real}
    f = equilibrium(q, problem, x, y)

    σ = problem.u_max^2 * deviatoric_tensor(q, problem, x, y, 0.0)
    ∇u = problem.u_max^2 * velocity_gradient(problem, x, y, 0.0)
    τ = q.speed_of_sound_squared * LatticeBoltzmann.lattice_viscosity(problem)

    cs = q.speed_of_sound_squared

    # HMM..
    factor = (problem.domain_size[1] * problem.domain_size[2])

    # For Taylor Green Vortex...
    factor = 2 * (problem.domain_size[1] * problem.domain_size[2])
    factor = (problem.domain_size[1] * problem.domain_size[2])

    for f_idx in 1:length(f)
        Q = hermite(Val{2}, q.abscissae[:, f_idx], q)

        # TODO: Check if we are using the correct viscosity here (we aren't...),
        # since we are setting the off equilibrium components of the shifted distribution funcitons
        f[f_idx] += -factor * q.weights[f_idx] * 0.5 * ((τ + 0.5) * cs) * dot(Q, ∇u + ∇u') #/ 0.785398163397454
        # f[f_idx] += q.weights[f_idx] * (cs^2 /factorial(2)) * dot(Q, σ)
    end

    velocity(q, f) = begin
        v = zeros(2)
        velocity!(q, f, density(q, f), v)
        v
    end
    σ = problem.u_max^2 * deviatoric_tensor(q, problem, x, y, 0.0)

    return f
end
function initial_condition(
    ::AnalyticalEquilibriumAndOffEquilibrium,
    q::Quadrature,
    problem::TGV,
    x::T,
    y::T,
) where {T <: Real}
    ρ = lattice_density(q, problem, x, y)
    f = equilibrium(
        q,
        ρ,
        lattice_velocity(q, problem, x, y),
        lattice_temperature(q, problem, x, y),
    )

    u = lattice_velocity(q, problem, x, y)
    σ = problem.u_max^2 * deviatoric_tensor(q, problem, x, y, 0.0)
    ∇u = problem.u_max^2 * velocity_gradient(problem, x, y, 0.0)

    σ = deviatoric_tensor(q, problem, x, y, 0.0)
    ∇u = velocity_gradient(problem, x, y, 0.0)
    p = pressure(q, problem, x, y, 0.0)

    # NEW
    cs = q.speed_of_sound_squared
    τ = cs * lattice_viscosity(problem)
    τ_eff = τ + 0.5
    for f_idx in 1:length(f)
        w_i = q.weights[f_idx]
        H_2 = hermite(Val{2}, q.abscissae[:, f_idx], q)

        temperature = 1.0
        f_neq = w_i * (cs * τ_eff * ρ * temperature) / (2) * dot(H_2, ∇u + ∇u')

        f[f_idx] -= f_neq
    end
    return f
end


text:
@testset "Index" begin
    @testset "NoIndex" begin
        @test project(NoIndex(), ()) === NoIndex()
        @test project(NoIndex(), (:a,)) === NoIndex()
        @test project(NoIndex(), (:a, :b)) === NoIndex()
    end

    @testset "UniqueIndex" begin
        @test project(UniqueIndex{(:a, :b)}(), (:a, :b)) === UniqueIndex{(:a, :b)}()
        @test project(UniqueIndex{(:a, :b)}(), (:b, :a)) === UniqueIndex{(:a, :b)}()
        @test project(UniqueIndex{(:a, :b)}(), (:a,)) === NoIndex()
        @test project(UniqueIndex{(:a, :b)}(), (:b,)) === NoIndex()
    end

    @testset "SortIndex" begin
        v = [3,2,1]
        @test project(SortIndex{(:a, :b)}(v), (:a, :b)) === SortIndex{(:a, :b)}(v)
        @test project(SortIndex{(:a, :b)}(v), (:b, :a)) === SortIndex{(:a, :b)}(v)
        @test project(SortIndex{(:a, :b)}(v), (:a,)) === SortIndex{(:a,)}(v)
        @test project(SortIndex{(:a, :b)}(v), (:b,)) === NoIndex()
    end

    @testset "UniqueSortIndex" begin
        v = [3,2,1]
        @test project(UniqueSortIndex{(:a, :b)}(v), (:a, :b)) === UniqueSortIndex{(:a, :b)}(v)
        @test project(UniqueSortIndex{(:a, :b)}(v), (:b, :a)) === UniqueSortIndex{(:a, :b)}(v)
        @test project(UniqueSortIndex{(:a, :b)}(v), (:a,)) === SortIndex{(:a,)}(v)
        @test project(UniqueSortIndex{(:a, :b)}(v), (:b,)) === NoIndex()
    end

    @testset "HashIndex" begin
        d = Dict((3, 2.0) => [1, 3], (4, 1.0) => [2])
        @test project(HashIndex{(:a, :b)}(d), (:a, :b)) === HashIndex{(:a, :b)}(d)
        @test project(HashIndex{(:a, :b)}(d), (:b, :a)) === HashIndex{(:a, :b)}(d)
        @test project(HashIndex{(:a, :b)}(d), (:a,)) === NoIndex()
        @test project(HashIndex{(:a, :b)}(d), (:b,)) === NoIndex()
    end

    @testset "UniqueHashIndex" begin
        d = Dict((3, 2.0) => 1, (4, 1.0) => 2)
        @test project(UniqueHashIndex{(:a, :b)}(d), (:a, :b)) === UniqueHashIndex{(:a, :b)}(d)
        @test project(UniqueHashIndex{(:a, :b)}(d), (:b, :a)) === UniqueHashIndex{(:a, :b)}(d)
        @test project(UniqueHashIndex{(:a, :b)}(d), (:a,)) === NoIndex()
        @test project(UniqueHashIndex{(:a, :b)}(d), (:b,)) === NoIndex()
    end
end

@testset "Multiple indexes" begin
    @testset "Removing NoIndex" begin
        @test clean(()) === ()
        @test clean((NoIndex(),)) === ()
        @test clean((UniqueIndex{:a}(),)) === (UniqueIndex{:a}(),)
        @test clean((NoIndex(), UniqueIndex{:a}(), NoIndex(), UniqueIndex{:b}())) === (UniqueIndex{:a}(), UniqueIndex{:b}())
    end

    @testset "Projecting multiple indexes" begin
        p = Project{(:a,)}()

        @test p(()) === ()
        @test p((NoIndex(),)) === ()
        @test p((UniqueIndex{(:a,)}(),)) === (UniqueIndex{(:a,)}(),)
        @test p((UniqueIndex{(:b,)}(),)) === ()
        @test p((UniqueIndex{(:a,:b)}(),)) === ()
        @test p((UniqueIndex{(:a,)}(), UniqueIndex{(:b,)}())) === (UniqueIndex{(:a,)}(),)
    end
end


text:
#### Turbulence closures
using DocStringExtensions
using CLIMA.PlanetParameters
using CLIMA.SubgridScaleParameters
export ConstantViscosityWithDivergence, SmagorinskyLilly, Vreman, AnisoMinDiss
export turbulence_tensors

abstract type TurbulenceClosure end

vars_state(::TurbulenceClosure, FT) = @vars()
vars_aux(::TurbulenceClosure, FT) = @vars()

function atmos_init_aux!(::TurbulenceClosure, ::AtmosModel, aux::Vars, geom::LocalGeometry)
end
function atmos_nodal_update_aux!(::TurbulenceClosure, ::AtmosModel, state::Vars, aux::Vars, t::Real)
end
function gradvariables!(::TurbulenceClosure, transform::Vars, state::Vars, aux::Vars, t::Real)
end
function diffusive!(::TurbulenceClosure, ::Orientation, diffusive, ∇transform, state, aux, t)
end

"""
    ν, τ = turbulence_tensors(::TurbulenceClosure, state::Vars, diffusive::Vars, aux::Vars, t::Real)

Compute the kinematic viscosity tensor (`ν`) and SGS momentum flux tensor (`τ`).
"""
function turbulence_tensors end

"""
    principal_invariants(X)

Calculates principal invariants of a tensor `X`. Returns 3 element tuple containing the invariants.
"""
function principal_invariants(X)
  first = tr(X)
  second = (first^2 - tr(X .^ 2))/2
  third = det(X)
  return (first, second, third)
end

"""
    symmetrize(X)

Compute `(X + X')/2`, returning a `SHermitianCompact` object.
"""
function symmetrize(X::StaticArray{Tuple{3,3}})
  SHermitianCompact(SVector(X[1,1], (X[2,1] + X[1,2])/2, (X[3,1] + X[1,3])/2, X[2,2], (X[3,2] + X[2,3])/2, X[3,3]))
end

"""
    norm2(X)

Compute
```math
\\sum_{i,j} S_{ij}^2
```
"""
function norm2(X::SMatrix{3,3,FT}) where {FT}
  abs2(X[1,1]) + abs2(X[2,1]) + abs2(X[3,1]) +
  abs2(X[1,2]) + abs2(X[2,2]) + abs2(X[3,2]) +
  abs2(X[1,3]) + abs2(X[2,3]) + abs2(X[3,3])
end
function norm2(X::SHermitianCompact{3,FT,6}) where {FT}
  abs2(X[1,1]) + 2*abs2(X[2,1]) + 2*abs2(X[3,1]) +
                   abs2(X[2,2]) + 2*abs2(X[3,2]) +
                                    abs2(X[3,3])
end

"""
    strain_rate_magnitude(S)

Compute
```math
|S| = \\sqrt{2\\sum_{i,j} S_{ij}^2}
```
"""
function strain_rate_magnitude(S::SHermitianCompact{3,FT,6}) where {FT}
  return sqrt(2*norm2(S))
end

"""
    ConstantViscosityWithDivergence <: TurbulenceClosure

Turbulence with constant dynamic viscosity (`ρν`).
Divergence terms are included in the momentum flux tensor.

# Fields

$(DocStringExtensions.FIELDS)
"""
struct ConstantViscosityWithDivergence{FT} <: TurbulenceClosure
  "Dynamic Viscosity [kg/m/s]"
  ρν::FT
end

vars_gradient(::ConstantViscosityWithDivergence,FT) = @vars()
vars_diffusive(::ConstantViscosityWithDivergence, FT) =
  @vars(S::SHermitianCompact{3,FT,6})

function diffusive!(::ConstantViscosityWithDivergence, ::Orientation,
    diffusive::Vars, ∇transform::Grad, state::Vars, aux::Vars, t::Real)

  diffusive.turbulence.S = symmetrize(∇transform.u)
end

function turbulence_tensors(m::ConstantViscosityWithDivergence,
    state::Vars, diffusive::Vars, aux::Vars, t::Real)

  S = diffusive.turbulence.S
  ν = m.ρν / state.ρ
  τ = (-2*ν) * S + (2*ν/3)*tr(S) * I
  return ν, τ
end

"""
    SmagorinskyLilly <: TurbulenceClosure

See § 1.3.2 in CliMA documentation

article{doi:10.1175/1520-0493(1963)091<0099:GCEWTP>2.3.CO;2,
  author = {Smagorinksy, J.},
  title = {General circulation experiments with the primitive equations},
  journal = {Monthly Weather Review},
  volume = {91},
  number = {3},
  pages = {99-164},
  year = {1963},
  doi = {10.1175/1520-0493(1963)091<0099:GCEWTP>2.3.CO;2},
  URL = {https://doi.org/10.1175/1520-0493(1963)091<0099:GCEWTP>2.3.CO;2},
  eprint = {https://doi.org/10.1175/1520-0493(1963)091<0099:GCEWTP>2.3.CO;2}
  }

article{doi:10.1111/j.2153-3490.1962.tb00128.x,
  author = {LILLY, D. K.},
  title = {On the numerical simulation of buoyant convection},
  journal = {Tellus},
  volume = {14},
  number = {2},
  pages = {148-172},
  doi = {10.1111/j.2153-3490.1962.tb00128.x},
  url = {https://onlinelibrary.wiley.com/doi/abs/10.1111/j.2153-3490.1962.tb00128.x},
  eprint = {https://onlinelibrary.wiley.com/doi/pdf/10.1111/j.2153-3490.1962.tb00128.x},
  year = {1962}
  }

Brunt-Vaisala frequency N² defined as in equation (1b) in
  Durran, D.R. and J.B. Klemp, 1982:
  On the Effects of Moisture on the Brunt-Väisälä Frequency.
  J. Atmos. Sci., 39, 2152–2158,
  https://doi.org/10.1175/1520-0469(1982)039<2152:OTEOMO>2.0.CO;2

# Fields

$(DocStringExtensions.FIELDS)
"""
struct SmagorinskyLilly{FT} <: TurbulenceClosure
  "Smagorinsky Coefficient [dimensionless]"
  C_smag::FT
end

vars_aux(::SmagorinskyLilly,FT) = @vars(Δ::FT)
vars_gradient(::SmagorinskyLilly,FT) = @vars(θ_v::FT)
vars_diffusive(::SmagorinskyLilly,FT) = @vars(S::SHermitianCompact{3,FT,6}, N²::FT)

function atmos_init_aux!(::SmagorinskyLilly, ::AtmosModel, aux::Vars, geom::LocalGeometry)
  aux.turbulence.Δ = lengthscale(geom)
end

function gradvariables!(m::SmagorinskyLilly, transform::Vars, state::Vars, aux::Vars, t::Real)
  transform.turbulence.θ_v = aux.moisture.θ_v
end

function diffusive!(::SmagorinskyLilly, orientation::Orientation,
    diffusive::Vars, ∇transform::Grad, state::Vars, aux::Vars, t::Real)

  diffusive.turbulence.S = symmetrize(∇transform.u)
  ∇Φ = ∇gravitational_potential(orientation, aux)
  diffusive.turbulence.N² = dot(∇transform.turbulence.θ_v, ∇Φ) / aux.moisture.θ_v
end

function turbulence_tensors(m::SmagorinskyLilly, state::Vars, diffusive::Vars, aux::Vars, t::Real)

  FT = eltype(state)
  S = diffusive.turbulence.S
  normS = strain_rate_magnitude(S)

  # squared buoyancy correction
  Richardson = diffusive.turbulence.N² / (normS^2 + eps(normS))
  f_b² = sqrt(clamp(1 - Richardson*inv_Pr_turb, 0, 1))
  ν = normS * f_b² * FT(m.C_smag * aux.turbulence.Δ)^2
  τ = (-2*ν) * S
  return ν, τ
end

"""
    Vreman{FT} <: TurbulenceClosure

  §1.3.2 in CLIMA documentation
Filter width Δ is the local grid resolution calculated from the mesh metric tensor. A Smagorinsky coefficient
is specified and used to compute the equivalent Vreman coefficient.

1) ν_e = √(Bᵦ/(αᵢⱼαᵢⱼ)) where αᵢⱼ = ∂uⱼ∂uᵢ with uᵢ the resolved scale velocity component.
2) βij = Δ²αₘᵢαₘⱼ
3) Bᵦ = β₁₁β₂₂ + β₂₂β₃₃ + β₁₁β₃₃ - β₁₂² - β₁₃² - β₂₃²
βᵢⱼ is symmetric, positive-definite.
If Δᵢ = Δ, then β = Δ²αᵀα

@article{Vreman2004,
  title={An eddy-viscosity subgrid-scale model for turbulent shear flow: Algebraic theory and applications},
  author={Vreman, AW},
  journal={Physics of fluids},
  volume={16},
  number={10},
  pages={3670--3681},
  year={2004},
  publisher={AIP}
}

# Fields

$(DocStringExtensions.FIELDS)
"""
struct Vreman{FT} <: TurbulenceClosure
  "Smagorinsky Coefficient [dimensionless]"
  C_smag::FT
end
vars_aux(::Vreman,FT) = @vars(Δ::FT)
vars_gradient(::Vreman,FT) = @vars(θ_v::FT)
vars_diffusive(::Vreman,FT) = @vars(∇u::SMatrix{3,3,FT,9}, N²::FT)

function atmos_init_aux!(::Vreman, ::AtmosModel, aux::Vars, geom::LocalGeometry)
  aux.turbulence.Δ = lengthscale(geom)
end
function gradvariables!(m::Vreman, transform::Vars, state::Vars, aux::Vars, t::Real)
  transform.turbulence.θ_v = aux.moisture.θ_v
end
function diffusive!(::Vreman, orientation::Orientation,
                    diffusive::Vars, ∇transform::Grad, state::Vars, aux::Vars, t::Real)
  diffusive.turbulence.∇u = ∇transform.u
  ∇Φ = ∇gravitational_potential(orientation, aux)
  diffusive.turbulence.N² = dot(∇transform.turbulence.θ_v, ∇Φ) / aux.moisture.θ_v
end

function turbulence_tensors(m::Vreman, state::Vars, diffusive::Vars, aux::Vars, t::Real)
  FT = eltype(state)
  α = diffusive.turbulence.∇u
  S = symmetrize(α)

  normS = strain_rate_magnitude(S)
  Richardson = diffusive.turbulence.N² / (normS^2 + eps(normS))
  f_b² = sqrt(clamp(1 - Richardson*inv_Pr_turb, 0, 1))

  β = f_b² * (aux.turbulence.Δ)^2 * (α' * α)
  Bβ = principal_invariants(β)[2]

  ν = max(0, m.C_smag^2 * FT(2.5) * sqrt(abs(Bβ/(norm2(α)+eps(FT)))))
  τ = (-2*ν) * S

  return ν, τ
end

"""
    AnisoMinDiss{FT} <: TurbulenceClosure

  §1.3.2 in CLIMA documentation
Filter width Δ is the local grid resolution calculated from the mesh metric tensor. A Poincare coefficient
is specified and used to compute the equivalent AnisoMinDiss coefficient (computed as the solution to the
eigenvalue problem for the Laplacian operator).

@article{doi:10.1063/1.4928700,
author = {Rozema,Wybe  and Bae,Hyun J.  and Moin,Parviz  and Verstappen,Roel },
title = {Minimum-dissipation models for large-eddy simulation},
journal = {Physics of Fluids},
volume = {27},
number = {8},
pages = {085107},
year = {2015},
doi = {10.1063/1.4928700},
URL = {https://aip.scitation.org/doi/abs/10.1063/1.4928700},
eprint = {https://aip.scitation.org/doi/pdf/10.1063/1.4928700}
}
-------------------------------------------------------------------------------------
# TODO: Future versions will include modifications of Abkar(2016), Verstappen(2018)
@article{PhysRevFluids.1.041701,
title = {Minimum-dissipation scalar transport model for large-eddy simulation of turbulent flows},
author = {Abkar, Mahdi and Bae, Hyun J. and Moin, Parviz},
journal = {Phys. Rev. Fluids},
volume = {1},
issue = {4},
pages = {041701},
numpages = {10},
year = {2016},
month = {Aug},
publisher = {American Physical Society},
doi = {10.1103/PhysRevFluids.1.041701},
url = {https://link.aps.org/doi/10.1103/PhysRevFluids.1.041701}
}

"""
struct AnisoMinDiss{FT} <: TurbulenceClosure
  C_poincare::FT
end
vars_aux(::AnisoMinDiss,FT) = @vars(Δ::FT)
vars_gradient(::AnisoMinDiss,FT) = @vars()
vars_diffusive(::AnisoMinDiss,FT) = @vars(∇u::SMatrix{3,3,FT,9})

function atmos_init_aux!(::AnisoMinDiss, ::AtmosModel, aux::Vars, geom::LocalGeometry)
  aux.turbulence.Δ = lengthscale(geom)
end
function gradvariables!(m::AnisoMinDiss, transform::Vars, state::Vars, aux::Vars, t::Real)
end
function diffusive!(::AnisoMinDiss, ::Orientation,
                    diffusive::Vars, ∇transform::Grad, state::Vars, aux::Vars, t::Real)
  diffusive.turbulence.∇u = ∇transform.u
end

function turbulence_tensors(m::AnisoMinDiss, state::Vars, diffusive::Vars, aux::Vars, t::Real)
  FT = eltype(state)
  α = diffusive.turbulence.∇u
  S = symmetrize(α)

  coeff = (aux.turbulence.Δ * m.C_poincare)^2
  βij = -(α' * α)
  ν = max(0, coeff * (dot(βij, S) / (norm2(α) + eps(FT))))
  τ = (-2*ν) * S

  return ν, τ
end

max_stars_repo_name: favba/CLIMA

max_issues_repo_name: favba/CLIMA

max_forks_repo_name: favba/CLIMA


text:
import Base: *, \
import LinearAlgebra: mul!, ldiv!

_get_data(A::LowerTriangular) = tril(A.data)
_get_data(A::UpperTriangular) = triu(A.data)
_get_data(A::UnitLowerTriangular) = tril(A.data)
_get_data(A::UnitUpperTriangular) = triu(A.data)
_get_data(A::Symmetric) = A.data

_unwrap_adj(x::Union{Adjoint,Transpose}) = parent(x)
_unwrap_adj(x) = x

const SparseMatrices{T} = Union{SparseMatrixCSC{T,BlasInt},
                        Symmetric{T,SparseMatrixCSC{T,BlasInt}},
                        LowerTriangular{T, SparseMatrixCSC{T,BlasInt}},
                        UnitLowerTriangular{T, SparseMatrixCSC{T,BlasInt}},
                        UpperTriangular{T, SparseMatrixCSC{T,BlasInt}},
                        UnitUpperTriangular{T, SparseMatrixCSC{T,BlasInt}}}

for T in [Complex{Float32}, Complex{Float64}, Float32, Float64]
for mat in (:StridedVector, :StridedMatrix)
for (tchar, ttype) in (('N', :()),
                       ('C', :Adjoint),
                       ('T', :Transpose))
    AT = tchar == 'N' ? :(SparseMatrixCSC{$T,BlasInt}) : :($ttype{$T,SparseMatrixCSC{$T,BlasInt}})
    @eval begin
        function mul!(C::$mat{$T}, adjA::$AT, B::$mat{$T}, α::Number, β::Number)
            A = _unwrap_adj(adjA)
            if isa(B, AbstractVector)
                return cscmv!($tchar, $T(α), matdescra(A), A, B, $T(β), C)
            else
                return cscmm!($tchar, $T(α), matdescra(A), A, B, $T(β), C)
            end
        end

        mul!(C::$mat{$T}, adjA::$AT, B::$mat{$T}) = mul!(C, adjA, B, one($T), zero($T))

        function (*)(adjA::$AT, B::$mat{$T})
            A = _unwrap_adj(adjA)
            if isa(B,AbstractVector)
                return mul!(zeros($T, mkl_size($tchar, A)[1]),            adjA, B)
            else
                return mul!(zeros($T, mkl_size($tchar, A)[1], size(B,2)), adjA, B)
            end
        end
    end

    for w in (:Symmetric, :LowerTriangular, :UnitLowerTriangular, :UpperTriangular, :UnitUpperTriangular)
        AT = tchar == 'N' ?
            :($w{$T,SparseMatrixCSC{$T,BlasInt}}) :
            :($ttype{$T,$w{$T,SparseMatrixCSC{$T,BlasInt}}})
        @eval begin
            function mul!(C::$mat{$T}, adjA::$AT, B::$mat{$T}, α::Number, β::Number)
                A = _unwrap_adj(adjA)
                if isa(B,AbstractVector)
                    return cscmv!($tchar, $T(α), matdescra(A), _get_data(A), B, $T(β), C)
                else
                    return cscmm!($tchar, $T(α), matdescra(A), _get_data(A), B, $T(β), C)
                end
            end

            mul!(C::$mat{$T}, adjA::$AT, B::$mat{$T}) = mul!(C, adjA, B, one($T), zero($T))

            function (*)(adjA::$AT, B::$mat{$T})
                A = _unwrap_adj(adjA)
                if isa(B,AbstractVector)
                    return mul!(zeros($T, mkl_size($tchar, A)[1]),            adjA, B)
                else
                    return mul!(zeros($T, mkl_size($tchar, A)[1], size(B,2)), adjA, B)
                end
            end
        end

        if w != :Symmetric
            @eval begin
                function ldiv!(α::Number, adjA::$AT,
                               B::$mat{$T}, C::$mat{$T})
                    A = _unwrap_adj(adjA)
                    if isa(B,AbstractVector)
                        return cscsv!($tchar, $T(α), matdescra(A), _get_data(A), B, C)
                    else
                        return cscsm!($tchar, $T(α), matdescra(A), _get_data(A), B, C)
                    end
                end

                ldiv!(C::$mat{$T}, A::$AT, B::$mat{$T}) =
                    ldiv!(one($T), A, B, C)

                function (\)(A::$AT, B::$mat{$T})
                    if isa(B,AbstractVector)
                        return ldiv!(zeros($T, size(A,1)),            A, B)
                    else
                        return ldiv!(zeros($T, size(A,1), size(B,2)), A, B)
                    end
                end
            end
        end
    end
end
end # mat
end # T


text:
using BenchmarkTools
using SASLib
using ReadStat

if length(ARGS) != 1
	println("Usage: julia ", PROGRAM_FILE, " <output-dir>")
	exit(1)
end

dir = ARGS[1]
if !isdir(dir)
	println("Error: ", dir, " does not exist")
	exit(2)
end

function performtest(io, f, samples, seconds)
    println(io, "\n\n================ $f =================")
	mime = MIME("text/plain")
    try
	    info("testing $f with ReadStat")
        println(io, "ReadStat:")
        b1 = @benchmark read_sas7bdat($f) samples=samples seconds=seconds
        show(io, mime, b1)
        println(io)

	    info("testing $f with SASLib")
        println(io, "SASLib:")
        b2 = @benchmark readsas($f, verbose_level=0) samples=samples seconds=seconds
        show(io, mime, b2)
        println(io)

	    info("testing $f with SASLib regular string")
        println(io, "SASLib (regular string):")
        b3 = @benchmark readsas($f, string_array_fn=Dict(:_all_ => REGULAR_STR_ARRAY), verbose_level=0)  samples=samples seconds=seconds
        show(io, mime, b3)
        println(io)

        md = metadata(f)
        nd = count(x -> last(x) == Float64, md.columnsinfo)
        ns = count(x -> last(x) == String, md.columnsinfo)
        nt = md.ncols - ns - nd
		println("debug: meta nd=$nd ns=$ns nt=$nt")

        t1 = minimum(b1).time / 1000000
        t2 = minimum(b2).time / 1000000
        t3 = minimum(b3).time / 1000000
        p2 = round(Int, t2/t1*100)
        p3 = round(Int, t3/t1*100)
        comp = md.compression
		info("Results: ", join(string.([f,t1,t2,p2,t3,p3,nd,ns,nt,comp]), ","))
        @printf io "%-40s: %8.3f ms %8.3f ms (%3d%%) %8.3f ms (%3d%%) %4d %4d %4d %4s\n" f t1 t2 p2 t3 p3 nd ns nt comp
    catch ex
        println(ex)
        b = " "
        @printf io "%-40s: %8s ms %8s ms (%3s%%) %8s ms (%3s%%) %4s %4s %4s %4s\n" f b b b b b b b b b
    finally
        flush(io)
    end
end

open("$dir/saslib_vs_readstat.log", "w") do io
	@printf io "%-40s: %8s    %8s    %3s    %8s    %3s   %4s %4s %4s %4s\n" "Filename" "ReadStat" "SASLib" "S/R" "SASLibA" "SA/R" "F64" "STR" "DT" "COMP"
	files = [
		("data_pandas/test3.sas7bdat", 10000, 5),
		("data_misc/numeric_1000000_2.sas7bdat", 100, 5),
		("data_misc/types.sas7bdat", 10000, 5),
		("data_AHS2013/homimp.sas7bdat", 10000, 5),
		("data_AHS2013/omov.sas7bdat", 10000, 5),
		("data_AHS2013/owner.sas7bdat", 10000, 5),
        ("data_AHS2013/ratiov.sas7bdat", 10000,5),
        ("data_AHS2013/rmov.sas7bdat", 10000, 5),
        ("data_AHS2013/topical.sas7bdat", 10000, 30),
        ("data_pandas/airline.sas7bdat", 10000, 5),
        ("data_pandas/datetime.sas7bdat", 10000, 5),
        ("data_pandas/productsales.sas7bdat", 10000, 10),
        ("data_pandas/test1.sas7bdat", 10000, 5),
        ("data_pandas/test2.sas7bdat", 10000, 5),
        ("data_pandas/test4.sas7bdat", 10000, 5),
        ("data_pandas/test5.sas7bdat", 10000, 5),
        ("data_pandas/test6.sas7bdat", 10000, 5),
        ("data_pandas/test7.sas7bdat", 10000, 5),
        ("data_pandas/test8.sas7bdat", 10000, 5),
        ("data_pandas/test9.sas7bdat", 10000, 5),
        ("data_pandas/test11.sas7bdat", 10000, 5),
        ("data_pandas/test10.sas7bdat", 10000, 5),
        ("data_pandas/test12.sas7bdat", 10000, 5),
        ("data_pandas/test13.sas7bdat", 10000, 5),
        ("data_pandas/test14.sas7bdat", 10000, 5),
        ("data_pandas/test15.sas7bdat", 10000, 5),
        ("data_pandas/test16.sas7bdat", 10000, 5),
        ("data_pandas/zero_variables.sas7bdat", 10000, 5),
        ("data_reikoch/barrows.sas7bdat", 10000, 5),
        ("data_reikoch/binary.sas7bdat", 10000, 5),
        ("data_reikoch/extr.sas7bdat", 10000, 5),
        ("data_reikoch/ietest2.sas7bdat", 10000, 5),
    ]
    for (f, x, y) in files
        performtest(io, f, x, y)
    end

end


text:
module WoodburyBenchmarks
using LinearAlgebra
using BenchmarkTools
using WoodburyFactorizations
using WoodburyFactorizations: Woodbury

# TODO: benchmark factorize against factorize without logdet computation
# TODO: figure out parameter c in WoodburyFactorizations, which
# TODO: benchmark pre-allocated solve, mul!
function judge_allocated_multiply(n = 64, m = 3)
    suite = BenchmarkGroup()
    U = randn(n, m)
    W = Woodbury(I(n), U, I(m), U')
    # factorization
    x = randn(n)
    F = factorize(W)
    # suite["solve"] = @benchmarkable $F \ $x
    suite["multiply"] = @benchmarkable $W * $x

    # suite["ternary dot"] = @benchmarkable dot($x, $W, $x)
    return suite
end

# controls when the Woodbury representation is more efficient than dense
function woodbury(n = 64, m = 3)
    suite = BenchmarkGroup()
    U = randn(n, m)
    W = Woodbury(I(n), U, I(m), U')
    # factorization
    suite["factorize"] = @benchmarkable factorize($W)
    x = randn(n)
    F = factorize(W)
    suite["solve"] = @benchmarkable $F \ $x
    suite["multiply"] = @benchmarkable $W * $x
    suite["ternary dot"] = @benchmarkable dot($x, $W, $x)
    suite["logdet"] = @benchmarkable logdet($W)
    return suite
end

########  to compare
function dense(n = 64, m = 3)
    suite = BenchmarkGroup()
    U = randn(n, m)
    W = Woodbury(I(n), U, I(m), U')
    MW = Matrix(W)
    suite["factorize"] = @benchmarkable factorize($MW)
    MWF = factorize(MW)
    x = randn(n)
    suite["solve"] = @benchmarkable  $MWF \ $x
    suite["multiply"] = @benchmarkable $MW * $x
    suite["ternary dot"] = @benchmarkable dot($x, $MW, $x)
    suite["logdet"] = @benchmarkable logdet($MW)
    return suite
end

# comparing Woodbury identity to dense matrix algebra
function wood_vs_dense(n = 64, m = 3)
    wsuite = woodbury(n, m)
    wr = run(wsuite)
    dsuite = dense(n, m)
    dr = run(dsuite)
    judge(minimum(wr), minimum(dr))
end

end # WoodburyBenchmarks


text:
using Test, Flux , StochasticDiffEq , LinearAlgebra
println("Optimal Stopping Time Test")
using NeuralPDE

using Random
Random.seed!(100)

d = 1
r = 0.04f0
beta = 0.2f0
T = 1
u0 = fill(80.00 , d , 1)
sdealg = EM()
ensemblealg = EnsembleThreads()
f(du,u,p,t) = (du .= r*u)
sigma(du,u,p,t)  = (du .= Diagonal(beta*u))
tspan = (0.0 , 1.0)
N = 50
dt = tspan[2]/49
K = 100.00
function g(t , x)
  return exp(-r*t)*(max(K -  maximum(x)  , 0))
end

prob  = SDEProblem(f , sigma , u0 , tspan ; g = g)
opt = Flux.ADAM(0.1)
m = Chain(Dense(d , 5, tanh), Dense(5, 16 , tanh)  , Dense(16 , N ), softmax)
sol = solve(prob, NeuralPDE.NNStopping( m, opt , sdealg , ensemblealg), verbose = true, dt = dt,
            abstol=1e-6, maxiters = 20 , trajectories = 200)

##Analytical Binomial Tree approach for American Options
function BinomialTreeAM1D(S0 , N , r , beta)
    V = zeros(N+1)
    dT = T/N
    u = exp(beta*sqrt(dT))
    d = 1/u
    S_T = [S0*(u^j)* (d^(N-j)) for j in 0:N]
    a = exp(r*dT)
    p = (a - d)/(u - d)
    q = 1.0 - p
    V = [max(K - x , 0) for x in S_T]
    for i in N-1:-1:0
      V[1:end-1] = exp(-r*dT).*(p*V[2:end] + q*V[1:end-1])
      S_T = S_T*u
      V = [max(K - S_T[i] , V[i]) for i in 1:size(S_T)[1]]
    end
    return V[1]
end
real_sol = BinomialTreeAM1D(u0[1] , N , r , beta)
error = abs(sol - real_sol)
@test error < 0.5


text:
# # Opinion spread

# ![](opinion.gif)

# This is a simple model of how an opinion spreads through a community. Each individual has a number of opinions as a list of integers. They can change their opinion by changing the numbers in the list.

# Agents can change their opinion at each step. They choose one of their neighbors randomly, and adopt one of the neighbor's opinion. They are more likely to adopt their neighbors opinion if the share more opinions with each other.

using Agents
using AgentsPlots
using Random

# ## Building the model
# ### 1. Model creation

mutable struct Citizen <: AbstractAgent
	id::Int
	pos::Tuple{Int, Int}
	stabilized::Bool
	opinion::Array{Int, 1}
	prev_opinion::Array{Int, 1}
end

function create_model(;dims=(10, 10), nopinions=3, levels_per_opinion=4)
	space = GridSpace(dims, periodic=true, moore=true)
	properties = Dict(:nopinions=>nopinions)
	model = AgentBasedModel(Citizen, space, scheduler=random_activation, properties=properties)
	for cell in 1:nv(model)
		add_agent!(cell, model, false, rand(1:levels_per_opinion, nopinions), rand(1:levels_per_opinion, nopinions))
	end
	return model
end

# ### 2. Stepping functions

function adopt!(agent, model)
	neighbor = rand(space_neighbors(agent, model))
	matches = model[neighbor].opinion .== agent.opinion
	nmatches = count(matches)
	if nmatches < model.nopinions && rand() < nmatches/model.nopinions
		switchId = rand(findall(x-> x==false, matches))
		agent.opinion[switchId] = model[neighbor].opinion[switchId]
	end
end

function update_prev_opinion!(agent, model)
	for i in 1:model.nopinions
		agent.prev_opinion[i] = agent.opinion[i]
	end
end

function is_stabilized!(agent, model)
	if agent.prev_opinion == agent.opinion
		agent.stabilized = true
	else
		agent.stabilized = false
	end
end

function agent_step!(agent, model)
	update_prev_opinion!(agent, model)
	adopt!(agent, model)
	is_stabilized!(agent, model)
end

# ## Running the model

levels_per_opinion=4
model = create_model(nopinions=3,levels_per_opinion=levels_per_opinion)

"run the model until all agents stabilize"
rununtil(model, s) = count(getproperty.(values(model.agents), :stabilized)) == prod(model.space.dimensions)

agentdata, _ = run!(model, agent_step!, dummystep, rununtil, adata=[(:stabilized, count)])

# ## Plotting

# The plot shows the number of stable agents, that is, number of agents whose opinions don't change from one step to the next. Note that the number of stable agents can fluctuate before the final convergence.

plot(1:size(agentdata, 1), agentdata.count_stabilized, legend=false, xlabel="generation", ylabel="# of stabilized agents")

# ### Animation

# Here is an animation that shows change of agent opinions over time. The first three opinions of an agent determines its color in RGB.
levels_per_opinion = 3
ac(agent) = RGB((agent.opinion[1:3] ./ levels_per_opinion)...)
model = create_model(nopinions=3, levels_per_opinion=levels_per_opinion)
anim = @animate for sp in 1:500
  step!(model, agent_step!)
  p = plotabm(model, ac=ac, as=12, am=:square)
  title!(p, "Step $(sp)")
  if rununtil(model, 1)
    break
  end
end
gif(anim, "opinion.gif")


text:
using Flux
using Zygote
using CUDA

include("utils/utils.jl")
include("representation/representation.jl")
include("nn_structures/nn_structures.jl")


text:
@testset "Differentials" begin
    @testset "Wirtinger" begin
        w = Wirtinger(1+1im, 2+2im)
        @test wirtinger_primal(w) == 1+1im
        @test wirtinger_conjugate(w) == 2+2im
        @test w + w == Wirtinger(2+2im, 4+4im)

        @test w + One() == w + 1 == w + Thunk(()->1) == Wirtinger(2+1im, 2+2im)
        @test w * One() == One() * w == w
        @test w * 2 == 2 * w == Wirtinger(2 + 2im, 4 + 4im)

        # TODO: other + methods stack overflow
        @test_throws ErrorException w*w
        @test_throws ArgumentError extern(w)
        for x in w
            @test x === w
        end
        @test broadcastable(w) == w
        @test_throws MethodError conj(w)
    end
    @testset "Zero" begin
        z = Zero()
        @test extern(z) === false
        @test z + z == z
        @test z + 1 == 1
        @test 1 + z == 1
        @test z * z == z
        @test z * 1 == z
        @test 1 * z == z
        for x in z
            @test x === z
        end
        @test broadcastable(z) isa Ref{Zero}
        @test conj(z) == z
    end
    @testset "One" begin
        o = One()
        @test extern(o) === true
        @test o + o == 2
        @test o + 1 == 2
        @test 1 + o == 2
        @test o * o == o
        @test o * 1 == 1
        @test 1 * o == 1
        for x in o
            @test x === o
        end
        @test broadcastable(o) isa Ref{One}
        @test conj(o) == o
    end

    @testset "Thunk" begin
        @test @thunk(3) isa Thunk

        @testset "show" begin
            rep = repr(Thunk(rand))
            @test occursin(r"Thunk\(.*rand.*\)", rep)
        end

        @testset "Externing" begin
            @test extern(@thunk(3)) == 3
            @test extern(@thunk(@thunk(3))) == 3
        end

        @testset "calling thunks should call inner function" begin
            @test (@thunk(3))() == 3
            @test (@thunk(@thunk(3)))() isa Thunk
        end

        @testset "erroring thunks should include the source in the backtrack" begin
            expected_line = (@__LINE__) + 2  # for testing it is at right palce
            try
                x = @thunk(error())
                extern(x)
            catch err
                err isa ErrorException || rethrow()
                st = stacktrace(catch_backtrace())
                # Should be 2nd last line, as last line will be the `error` function
                stackframe = st[2]
                @test stackframe.line == expected_line
                @test stackframe.file == Symbol(@__FILE__)
            end
        end
    end

    @testset "No ambiguities in $f" for f in (+, *)
        # We don't use `Test.detect_ambiguities` as we are only interested in
        # the +, and * operations. We also would catch any that are unrelated
        # to this package. but that is not a problem. Since no such failings
        # occur in our dependencies.

        ambig_methods = [
            (m1, m2) for m1 in methods(f), m2 in methods(f) if Base.isambiguous(m1, m2)
        ]
        @test isempty(ambig_methods)
    end

    @testset "Refine Differential" begin
        @test refine_differential(typeof(1.0 + 1im), Wirtinger(2,2)) == Wirtinger(2,2)
        @test refine_differential(typeof([1.0 + 1im]), Wirtinger(2,2)) == Wirtinger(2,2)

        @test refine_differential(typeof(1.2), Wirtinger(2,2)) == 4
        @test refine_differential(typeof([1.2]), Wirtinger(2,2)) == 4

        # For most differentials, in most domains, this does nothing
        for der in (DoesNotExist(), @thunk(23), @thunk(Wirtinger(2,2)), [1 2], One(), Zero(), 0.0)
            for 𝒟 in typeof.((1.0 + 1im, [1.0 + 1im], 1.2, [1.2]))
                @test refine_differential(𝒟, der) === der
            end
        end
    end
end


text:
@testset "KnotVector" begin
    @testset "Construction" begin
        @test KnotVector([1.0, 1.5, 2.5], 3) == [1.0, 1.0, 1.0, 1.0, 1.5, 2.5, 2.5, 2.5, 2.5]
        @test KnotVector(1:5, 2, 3) == [1, 1, 1, 2, 3, 4, 5, 5, 5, 5]
        @test KnotVector(big(-1):big(1), 2) == [-1, -1, -1, 0, 1, 1, 1]
        @test KnotVector(1//2:5//2, 2, 0) == [1//2, 1//2, 1//2, 3//2, 5//2]
        @test KnotVector([true, false, true], 0) == [true, false, true]
        @test KnotVector(["test", "string"], 0, 1) == ["test", "string", "string"]
        @test KnotVector(KnotVector(UInt8(0):UInt8(3), 2, 1), 2) == [0, 0, 0, 0, 0, 1, 2, 3, 3, 3, 3]
        @test KnotVector([1.0, 1.5, 2.5], 3) isa AbstractVector{Float64}
        @test KnotVector(1:5, 2, 3) isa AbstractVector{Int}
        @test KnotVector(big(-1):big(1), 2) isa AbstractVector{BigInt}
        @test KnotVector(1//2:5//2, 2, 0) isa AbstractVector{Rational{Int}}
        @test KnotVector([true, false, true], 0) isa AbstractVector{Bool}
        @test KnotVector(["test", "string"], 0, 1) isa AbstractVector{String}
        @test KnotVector(KnotVector(UInt8(0):UInt8(3), 2, 1), 2) isa AbstractVector{UInt8}
        @test_throws ArgumentError KnotVector(1:0, 2)
        @test_throws ArgumentError KnotVector(Float64[], 2)
        @test_throws DomainError KnotVector(1:5, -1)
        @test_throws DomainError KnotVector(1:5, -1, 2)
        @test_throws DomainError KnotVector(1:5, 2, -1)
        @test_throws ArgumentError KnotVector(OffsetArray(1:5, 1), 2)
    end

    @testset "copy" begin
        @test copy(KnotVector(1:5, 2)) === KnotVector(1:5, 2)
        a = KnotVector([1,2,3], 1, 2)
        @test a == copy(a)
        @test a !== copy(a)
        @test typeof(a) === typeof(copy(a))
    end

    @testset "Dimensions" begin
        @test axes(KnotVector(1:5, 3)) == (1:11,)
        @test axes(KnotVector(1:5, 3), 1) == 1:11
        @test axes(KnotVector(1:5, 3), 2) == 1:1
        @test axes(KnotVector([1.0, 4.0], 2, 1)) == (1:5,)
        @test axes(KnotVector([1.0, 4.0], 2, 1), 1) == 1:5
        @test axes(KnotVector([1.0, 4.0], 2, 1), 2) == 1:1
        @test length(KnotVector(1:5, 3)) == 11
        @test length(KnotVector([1.0, 4.0], 2, 1)) == 5
        @test size(KnotVector(1:5, 3)) == (11,)
        @test size(KnotVector(1:5, 3), 1) == 11
        @test size(KnotVector(1:5, 3), 2) == 1
        @test size(KnotVector([1.0, 4.0], 2, 1)) == (5,)
        @test size(KnotVector([1.0, 4.0], 2, 1), 1) == 5
        @test size(KnotVector([1.0, 4.0], 2, 1), 2) == 1
    end

    @testset "Indices" begin
        @test firstindex(KnotVector(1:5, 3)) == 1
        @test firstindex(KnotVector(1:5, 3), 1) == 1
        @test firstindex(KnotVector(1:5, 3), 2) == 1
        @test firstindex(KnotVector([1.0, 4.0], 2, 1)) == 1
        @test firstindex(KnotVector([1.0, 4.0], 2, 1), 1) == 1
        @test firstindex(KnotVector([1.0, 4.0], 2, 1), 2) == 1
        @test lastindex(KnotVector(1:5, 3)) == 11
        @test lastindex(KnotVector(1:5, 3), 1) == 11
        @test lastindex(KnotVector(1:5, 3), 2) == 1
        @test lastindex(KnotVector([1.0, 4.0], 2, 1)) == 5
        @test lastindex(KnotVector([1.0, 4.0], 2, 1), 1) == 5
        @test lastindex(KnotVector([1.0, 4.0], 2, 1), 2) == 1
        @test eachindex(KnotVector(1:5, 3)) == 1:11
        @test eachindex(KnotVector([1.0, 4.0], 2, 1)) == 1:5
        @test keys(KnotVector(1:5, 3)) == 1:11
        @test keys(KnotVector([1.0, 4.0], 2, 1)) == 1:5
    end

    @testset "Indexing and iteration" begin
        @test collect(KnotVector(1:5, 3)) == [1, 1, 1, 1, 2, 3, 4, 5, 5, 5, 5]
        @test collect(KnotVector([1.0, 4.0], 2, 1)) == [1.0, 1.0, 1.0, 4.0, 4.0]
        @test first(KnotVector(1:5, 3)) == 1
        @test first(KnotVector([1.0, 4.0], 2, 1)) == 1.0
        @test last(KnotVector(1:5, 3)) == 5
        @test last(KnotVector([1.0, 4.0], 2, 1)) == 4.0
        @test KnotVector(1:5, 3)[:] === KnotVector(1:5, 3)
        @test KnotVector(1:5, 3)[1] == 1
        @test KnotVector(1:5, 3)[5] == 2
        @test KnotVector(1:5, 3)[end] == 5
        @test KnotVector(1:5, 3)[1:5] == [1, 1, 1, 1, 2]
        @test KnotVector(1:5, 3)[2,1] == 1
        @test KnotVector(1:5, 3)[1:5,1] == [1, 1, 1, 1, 2]
        @test_throws BoundsError KnotVector(1:5, 3)[0]
        @test_throws BoundsError KnotVector(1:5, 3)[12]
        @test_throws BoundsError KnotVector(1:5, 3)[1,2]
        @test_throws BoundsError KnotVector(1:5, 3)[1:5,2]
        @test_throws BoundsError KnotVector(1:5, 3)[1:20,1]
        @test KnotVector([1.0, 4.0], 2, 1)[:] == KnotVector([1.0, 4.0], 2, 1)
        @test KnotVector([1.0, 4.0], 2, 1)[:] !== KnotVector([1.0, 4.0], 2, 1)
        @test parent(KnotVector([1.0, 4.0], 2, 1)[:]) == parent(KnotVector([1.0, 4.0], 2, 1))
        @test typeof(KnotVector([1.0, 4.0], 2, 1)[:]) === typeof(KnotVector([1.0, 4.0], 2, 1))
        @test KnotVector([1.0, 4.0], 2, 1)[1] == 1.0
        @test KnotVector([1.0, 4.0], 2, 1)[4] == 4.0
        @test KnotVector([1.0, 4.0], 2, 1)[end] == 4.0
        @test KnotVector([1.0, 4.0], 2, 1)[1:4] == [1.0, 1.0, 1.0, 4.0]
        @test KnotVector([1.0, 4.0], 2, 1)[4,1] == 4.0
        @test KnotVector([1.0, 4.0], 2, 1)[1:4,1] == [1.0, 1.0, 1.0, 4.0]
        @test_throws BoundsError KnotVector([1.0, 4.0], 2, 1)[0]
        @test_throws BoundsError KnotVector([1.0, 4.0], 2, 1)[6]
        @test_throws BoundsError KnotVector([1.0, 4.0], 2, 1)[1,2]
        @test_throws BoundsError KnotVector([1.0, 4.0], 2, 1)[1:4,2]
        @test_throws BoundsError KnotVector([1.0, 4.0], 2, 1)[1:10,1]
    end

    @testset "issorted" begin
        @test issorted(KnotVector(1:10, 5))
        @test issorted(KnotVector([1.0, 2.5, 5.0], 3, 5))
        @test !issorted(KnotVector(5:-1:1, 5))
        @test !issorted(KnotVector([1.0, 3.0, 2.0], 2, 1))
    end

    @testset "parent" begin
        @test (vec = [1.0, 1.5, 2.5]; parent(KnotVector(vec, 3)) === vec)
        @test (vec = 1:5; parent(KnotVector(vec, 2, 3)) === vec)
        @test (vec = big(-1):big(1); parent(KnotVector(vec, 2)) === vec)
        @test (vec = UInt8(0):UInt8(3); parent(KnotVector(KnotVector(vec, 2, 1), 2)) === vec)
    end

    @testset "unique/allunique" begin
        @test unique(KnotVector(1:3, 5)) == [1,2,3]
        @test unique(KnotVector([0], 2, 3)) == [0]
        @test unique(KnotVector([1,1,2,1], 0)) == [1,2]
        @test unique(KnotVector([1,1,2,1], 3)) == [1,2]
        @test allunique(KnotVector(1:5, 0))
        @test !allunique(KnotVector(1:5, 1))
        @test !allunique(KnotVector(1:5, 0, 1))
        @test !allunique(KnotVector(1:5, 2, 0))
        @test !allunique(KnotVector([1,1], 0))
    end
end


text:
abstract type AbstractCurve <: Function end

struct FunctionCurve <: AbstractCurve
    γ::Function
    tspan::Tuple{Float64, Float64}
end

function (γ::FunctionCurve)(t::Real)
    γ.γ(t)
end

max_stars_repo_name: zoemcc/geometry

max_issues_repo_name: zoemcc/geometry

max_forks_repo_name: zoemcc/geometry


text:
#Finds the steady states of a biochemical reaction network.
function steady_states(model::DiffEqBase.AbstractReactionNetwork,p::Vector,replacements::Expr...)
    expr = meta_program(deepcopy(model.f_func), model.syms, Dict(zip(model.params, p)),replacements...)
    return eval(expr)
end

#Finds the steady states of a biochemical reaction network and the corresponding stability, returns the result in an array.
function steady_states_stabiity(model::DiffEqBase.AbstractReactionNetwork,p::Vector,replacements::Expr...)
    expr = meta_program(deepcopy(model.f_func), model.syms, Dict(zip(model.params, p)),replacements...)
    return [(fp,stability(fp,model,p)) for fp in eval(expr)]
end

#Takes the reaction rates in a biochemical reaction networks and pieces together an expression which can be evaluated to generate the steady states.
function meta_program(eqs::Vector{Expr},vars::Vector{Symbol},pars::Dict,replacements::Expr...;allowed_exp_params::Set{Symbol}=Set{Symbol}())
    poly_vars = :(@polyvar)
    foreach(var -> push!(poly_vars.args,var), vars)
    polynomial = :(internal___var___poly = [])
    foreach(eq -> push!(polynomial.args[2].args,eq), eqs)
    my_replace!(polynomial,vars,replacements...)
    recursive_check_powers!(recursive_replace!(polynomial,pars),allowed_exp_params)
    return Expr(:block, poly_vars, polynomial, :(solve_poly(internal___var___poly)))
end

#Takes the reaction rates in a biochemical reaction networks and pieces together an expression which can be evaluated to generate the steady states.
function my_replace!(polynomial::Expr,syms::Vector{Symbol},replacements::Expr...)
    replaced = Set{Symbol}()
    for rep in replacements
        next_replace = recursive_find_next_replacement(rep,setdiff(syms,replaced))
        (next_replace!=nothing) ? push!(replaced,next_replace) : continue
        polynomial.args[2].args[findfirst(syms.==next_replace)] = balance_poly(deepcopy(rep))
    end
end

#Take a polynomial enterd with a '=' and makes it to a single side.
function balance_poly(poly::Expr)
    return :($(poly.args[1])-$(poly.args[2].args[2]))
end

#Solves a polynomial system using HomotopyContinuation, but discards solutions not of interest for BRNs.
function solve_poly(poly)
    in(:num,fieldnames(typeof(poly[1]))) ? (polynomial = [poly[1].num]) : (polynomial = [poly[1]])
    for i = 2:length(poly)
        in(:num,fieldnames(typeof(poly[i]))) ? (push!(polynomial, poly[i].num)) : (push!(polynomial, poly[i]))
    end
    results = HomotopyContinuation.solve(polynomial,report_progress=false)
    results = realsolutions(results; tol=1e-6, onlynonsingular=false, singulartol=1e10, onlyfinite=true)
    final_results = Vector{Vector{Float64}}()
    foreach(res -> (minimum(res)>-0.001) && push!(final_results,res), results)
    return final_results
end

#Parses an expression and inserts number instead of symbols.
function recursive_replace!(expr::Any,pars::Dict)
    haskey(pars,expr) && return pars[expr]
    (typeof(expr) != Expr) && return expr
    for i = 1:length(expr.args)
        expr.args[i] = recursive_replace!(expr.args[i], pars)
    end
    return expr
end

#Find the reactant in a fixed concentration equation for which to replace its reaction rate equation.
function recursive_find_next_replacement(expr::Any,syms::Array{Symbol,1})
    (typeof(expr)==Symbol)&&in(expr,syms)&&(return expr)
    (typeof(expr) != Expr)&&(return nothing)
    for arg in expr.args
        output = recursive_find_next_replacement(arg,syms)
        (output!=nothing) && (return output)
    end #Returns nothing if for loop passed.
end

#Checks all of the powers in an expressions, ensuring they have the correct form (i.e. not a float).
function recursive_check_powers!(expr::Any,allowed_params::Set{Symbol}=Set{Symbol}())
    !(typeof(expr) <: Expr) && return
    foreach(arg -> recursive_check_powers!(arg,allowed_params), expr.args)
    if (expr.args[1] == :^)
        exp = expr.args[3]
        if (typeof(exp)==Expr)&&(exp.head==:call)&&(exp.args[1]==:-)&&(length(exp.args)==2)&&(typeof(exp.args[2])<:Number)
            expr.args[3] = -exp.args[2]
            exp = expr.args[3]
        end
        in(exp,allowed_params) && return
        !(typeof(exp)<:Number) && error("Error: Bad exponent: ", exp)
        (exp%1.0 != 0) && error("Error: Non integer exponent.")
        expr.args[3] = Int64(exp)
    end
end

#Calculates the stability of a given fixed point for a given reaction network and parameter values.
function stability(fp::Vector{Float64},model::DiffEqBase.AbstractReactionNetwork,pars::Vector)
    jacobian = map(val->recursive_replace!(recursive_replace!(val,Dict(zip(model.syms,fp))),Dict(zip(model.params,pars))),deepcopy(model.symjac))
    eigenvalues = eigvals(eval.(jacobian))
    return (maximum(real(eigenvalues))<1e-6,any(imag(eigenvalues).>1e-6))
end

max_stars_repo_name: TorkelE/BRNEquilibrate

max_issues_repo_name: TorkelE/BRNEquilibrate

max_forks_repo_name: TorkelE/BRNEquilibrate


text:
MOI.supports(::Optimizer, ::MOI.VariableName, ::Type{MOI.VariableIndex}) = false

MOI.supports_constraint(::Optimizer, ::Type{<:ALLV}, ::Type{<:SS}) = true
MOI.supports_constraint(::Optimizer, ::Type{<:ALLV}, ::Type{<:NCS}) = true

MOI.supports_constraint(::Optimizer, ::Type{<:SF}, ::Type{<:LS}) = true
MOI.supports_constraint(::Optimizer, ::Type{VAF}, ::Type{<:VLS}) = true
MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{<:VLS}) = true

# TODO:
MOI.supports(::Optimizer, ::MOI.ConstraintName, ::Type{MOI.ConstraintIndex}) = false
# TODO we can support those
MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{NPC}) = false

MOI.supports(::Optimizer, ::MOI.VariablePrimalStart, ::Type{MOI.VariableIndex}) = true
MOI.supports(::Optimizer, ::MOI.ConstraintDualStart, ::Type{MOI.ConstraintIndex}) = true

MOI.supports(::Optimizer, ::MOI.NLPBlock) = true

MOI.supports(::Optimizer, ::MOI.ObjectiveFunction{MOI.SingleVariable}) = true
MOI.supports(::Optimizer, ::MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}) = true
MOI.supports(::Optimizer, ::MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{Float64}}) = true
MOI.supports(::Optimizer, ::MOI.ObjectiveSense) = true


text:
# spike generation function for integrator experiment
using Distributions, DifferentialEquations

function sample_inhomogeneous_poisson(rate, trange; kwargs...)
    spiketimes = Float64[]
    rate_wrapped(u,p,t) = rate(t)
    effect!(integrator) = push!(integrator.p, integrator.t)

    prob = ODEProblem((du,u,p,t)->du.=0, Float64[],trange, spiketimes)
    jump = VariableRateJump(rate_wrapped, effect!)
    jump_prob=JumpProblem(prob, Direct(),jump)
    solve(jump_prob,Tsit5(); kwargs...)

    return spiketimes
end

function generateContiguousVolleys(trange, population_size, λ, blob_size, background_rate=0; dt_spike=0.0, dt_volley=0)
    T = trange[2]-trange[1]
    spikes = [rand(rand(Poisson(background_rate*T))).*T.+trange[1] for i ∈ 1:population_size]
    volleys = NamedTuple{(:time,:idxs),Tuple{Float64,Vector{Int}}}[]

    smp=if isa(λ,Real)
        rand(rand(Poisson(λ*T))).*T.+trange[1]
    else
        sample_inhomogeneous_poisson(λ, trange)
    end

    sort!(smp)
    dels = Int[]
    for (k,t) ∈ enumerate(smp[2:end])
        if t ≤ smp[k] + dt_volley
            push!(dels, k+1)
        end
    end
    deleteat!(smp, dels)

    for st ∈ smp
        idxs = mod.(rand(1:population_size).+ (0:blob_size-1), Ref(1:population_size))
        for idx ∈ idxs
            push!(spikes[idx], st)
        end

        push!(volleys, (time=st, idxs=idxs))
    end

    foreach(sort!, spikes)
    for j ∈ 1:population_size
        sort!(spikes[j])
        dels = Int[]
        for (k,t) ∈ enumerate(spikes[j][2:end])
            if t ≤ spikes[j][k] + dt_spike
                push!(dels, k+1)
            end
        end
        deleteat!(spikes[j], dels)
    end
    return spikes, volleys
end

"""
    plot_spike_raster(trange, spikes, spike_width; p=plot(), colors=1:sum(length.(spikes)), kwargs...)
Plots spiketrains belonging to several different groups given by `spikes::Vector{Vector{Vector{Float64}}}` with the given `spike_width` over a time interval `trange`.
Each group is colorcoded by the corresponding element of `colors`.
"""
function plot_spike_raster(trange, spikes, spike_width; p=plot(), colors=1:length(spikes), highlight=[NamedTuple{(:time,:idxs),Tuple{Float64,Vector{Int}}}[] for i ∈ eachindex(spikes)], highlight_color=:red, ε=1e-3, kwargs...)
    plot!(p)
    k=0
    for (i,group) ∈ enumerate(spikes)

        # Set spike colors for each spike in each train of this group
        spike_colors = [convert(Vector{Any},fill(colors[i], length(spiketrain))) for spiketrain ∈ group]
        for (k,volley) ∈ enumerate(highlight[i])
            c = if isa(highlight_color, AbstractArray)
                highlight_color[i]
            else
                highlight_color
            end

            for j ∈ volley.idxs
                cc = if isa(c, AbstractArray)
                    c[j]
                else
                    c
                end

                # overwrite spike color if highlighted
                for s ∈ searchsortedfirst(group[j], volley.time-ε/2):searchsortedlast(group[j], volley.time+ε/2)
                    spike_colors[j][s] = c
                end
            end
        end

        # plot the actual spikes
        for (j,synapse) ∈ enumerate(group)
            plot!([Shape([t,t+spike_width,t+spike_width,t],[k-0.45, k-0.45, k+0.45, k+0.45]) for t ∈ synapse], color=spike_colors[j], linecolor=nothing)
            k-=1
        end
    end
    plot!(; legend=false, grid=false, yticks=false, xlim=trange, ylim=(k+1-0.5,0.5), kwargs...)
    return p
end


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct ApplicationGatewaySslPolicy <: SwaggerModel
    disabledSslProtocols::Any # spec type: Union{ Nothing, Vector{String} } # spec name: disabledSslProtocols

    function ApplicationGatewaySslPolicy(;disabledSslProtocols=nothing)
        o = new()
        validate_property(ApplicationGatewaySslPolicy, Symbol("disabledSslProtocols"), disabledSslProtocols)
        setfield!(o, Symbol("disabledSslProtocols"), disabledSslProtocols)
        o
    end
end # type ApplicationGatewaySslPolicy

const _property_map_ApplicationGatewaySslPolicy = Dict{Symbol,Symbol}(Symbol("disabledSslProtocols")=>Symbol("disabledSslProtocols"))
const _property_types_ApplicationGatewaySslPolicy = Dict{Symbol,String}(Symbol("disabledSslProtocols")=>"Vector{String}")
Base.propertynames(::Type{ ApplicationGatewaySslPolicy }) = collect(keys(_property_map_ApplicationGatewaySslPolicy))
Swagger.property_type(::Type{ ApplicationGatewaySslPolicy }, name::Symbol) = Union{Nothing,eval(Meta.parse(_property_types_ApplicationGatewaySslPolicy[name]))}
Swagger.field_name(::Type{ ApplicationGatewaySslPolicy }, property_name::Symbol) =  _property_map_ApplicationGatewaySslPolicy[property_name]

const _allowed_ApplicationGatewaySslPolicy_disabledSslProtocols = ["TLSv1_0", "TLSv1_1", "TLSv1_2"]

function check_required(o::ApplicationGatewaySslPolicy)
    true
end

function validate_property(::Type{ ApplicationGatewaySslPolicy }, name::Symbol, val)
    if name === Symbol("disabledSslProtocols")
        Swagger.validate_param(name, "ApplicationGatewaySslPolicy", :enum, val, _allowed_ApplicationGatewaySslPolicy_disabledSslProtocols)
    end
end


text:
using BackwardsLinalg
using Random
using Test

@testset "symeigen real" begin
    A = randn(4,4)
    A = A+A'
    op = randn(4, 4)
    op += op'
    function f(A)
        E, U = symeigen(A)
        E |> sum
    end
    function g(A)
        E, U = symeigen(A)
        v = U[:,1]
        (v'*op*v)[]|>real
    end
    @test gradient_check(f, A)
    @test gradient_check(g, A)
end

@testset "symeigen complex" begin
    Random.seed!(6)
    A = randn(ComplexF64, 4,4)
    A = A+A'
    op = randn(ComplexF64, 4, 4)
    op += op'
    function f(A)
        E, U = symeigen(A)
        E |> sum
    end
    function g(A)
        E, U = symeigen(A)
        v = U[:,1]
        (v'*op*v)[]|>real
    end
    @test gradient_check(f, A)
    @test gradient_check(g, A)
end


text:
struct SingleIssue
    x
    path::String
    reason::String
    val
end

function Base.show(io::IO, issue::SingleIssue)
    println(io, """Validation failed:
    path:         $(isempty(issue.path) ? "top-level" : issue.path)
    instance:     $(issue.x)
    schema key:   $(issue.reason)
    schema value: $(issue.val)""")
end

"""
    validate(x, s::Schema)

Validate document `x` is valid against the Schema `s`. If valid, return `nothing`, else
return a `SingleIssue`. When printed, the returned `SingleIssue` describes the reason why
the validation failed.

## Examples

    julia> schema = Schema(
               Dict(
                   "properties" => Dict(
                       "foo" => Dict(),
                       "bar" => Dict()
                   ),
                   "required" => ["foo"]
               )
           )
    Schema

    julia> data_pass = Dict("foo" => true)
    Dict{String,Bool} with 1 entry:
      "foo" => true

    julia> data_fail = Dict("bar" => 12.5)
    Dict{String,Float64} with 1 entry:
      "bar" => 12.5

    julia> validate(data_pass, schema)

    julia> validate(data_fail, schema)
    Validation failed:
    path:         top-level
    instance:     Dict("bar"=>12.5)
    schema key:   required
    schema value: ["foo"]
"""
function validate(x, schema::Schema)
    return _validate(x, schema.data, "")
end

Base.isvalid(x, schema::Schema) = validate(x, schema) === nothing

function _validate(x, schema, path::String)
    schema = _resolve_refs(schema)
    return _validate_entry(x, schema, path)
end

function _validate_entry(x, schema::AbstractDict, path)
    for (k, v) in schema
        ret = _validate(x, schema, Val{Symbol(k)}(), v, path)
        if ret !== nothing
            return ret
        end
    end
    return
end

function _validate_entry(x, schema::Bool, path::String)
    return schema ? nothing : SingleIssue(x, path, "schema", schema)
end

function _resolve_refs(schema::AbstractDict, explored_refs = Any[schema])
    if !haskey(schema, "\$ref")
        return schema
    end
    schema = schema["\$ref"]
    if any(x -> x === schema, explored_refs)
        error("cannot support circular references in schema.")
    end
    push!(explored_refs, schema)
    return _resolve_refs(schema, explored_refs)
end
_resolve_refs(schema, explored_refs = Any[]) = schema

# Default fallback
_validate(::Any, ::Any, ::Val, ::Any, ::String) = nothing

###
### Core JSON Schema
###

# 9.2.1.1
function _validate(x, schema, ::Val{:allOf}, val::Vector, path::String)
    for v in val
        ret = _validate(x, v, path)
        if ret !== nothing
            return ret
        end
    end
    return
end

# 9.2.1.2
function _validate(x, schema, ::Val{:anyOf}, val::Vector, path::String)
    for v in val
        if _validate(x, v, path) === nothing
            return
        end
    end
    return SingleIssue(x, path, "anyOf", val)
end

# 9.2.1.3
function _validate(x, schema, ::Val{:oneOf}, val::Vector, path::String)
    found_match = false
    for v in val
        if _validate(x, v, path) === nothing
            if found_match # Found more than one match!
                return SingleIssue(x, path, "oneOf", val)
            end
            found_match = true
        end
    end
    return found_match ? nothing : SingleIssue(x, path, "oneOf", val)
end

# 9.2.1.4
function _validate(x, schema, ::Val{:not}, val, path::String)
    ret = _validate(x, val, path)
    return ret === nothing ? SingleIssue(x, path, "not", val) : nothing
end

# 9.2.2.1: if

# 9.2.2.2: then

# 9.2.2.3: else

###
### Checks for Arrays.
###

# 9.3.1.1
function _validate(x::Vector, schema, ::Val{:items}, val::AbstractDict, path::String)
    items = fill(false, length(x))
    for (i, xi) in enumerate(x)
        ret = _validate(xi, val, path * "[$(i)]")
        if ret !== nothing
            return ret
        end
        items[i] = true
    end
    additionalItems = get(schema, "additionalItems", nothing)
    return _additional_items(x, schema, items, additionalItems, path)
end

function _validate(x::Vector, schema, ::Val{:items}, val::Vector, path::String)
    items = fill(false, length(x))
    for (i, xi) in enumerate(x)
        if i > length(val)
            break
        end
        ret = _validate(xi, val[i], path * "[$(i)]")
        if ret !== nothing
            return ret
        end
        items[i] = true
    end
    additionalItems = get(schema, "additionalItems", nothing)
    return _additional_items(x, schema, items, additionalItems, path)
end

function _validate(x::Vector, schema, ::Val{:items}, val::Bool, path::String)
    return val || (!val && length(x) == 0) ? nothing : SingleIssue(x, path, "items", val)
end

function _additional_items(x, schema, items, val, path)
    for i = 1:length(x)
        if items[i]
            continue  # Validated against 'items'.
        end
        ret = _validate(x[i], val, path * "[$(i)]")
        if ret !== nothing
            return ret
        end
    end
    return
end

function _additional_items(x, schema, items, val::Bool, path)
    return !val && !all(items) ? SingleIssue(x, path, "additionalItems", val) : nothing
end

_additional_items(x, schema, items, val::Nothing, path) = nothing

# 9.3.1.2
function _validate(x::Vector, schema, ::Val{:additionalItems}, val, path::String)
    return  # Supported in `items`.
end

# 9.3.1.3: unevaluatedProperties

# 9.3.1.4
function _validate(x::Vector, schema, ::Val{:contains}, val, path::String)
    for (i, xi) in enumerate(x)
        ret = _validate(xi, val, path * "[$(i)]")
        if ret === nothing
            return
        end
    end
    return SingleIssue(x, path, "contains", val)
end

###
### Checks for Objects
###

# 9.3.2.1
function _validate(x::AbstractDict, schema, ::Val{:properties}, val::AbstractDict, path::String)
    for (k, v) in x
        if haskey(val, k)
            ret = _validate(v, val[k], path * "[$(k)]")
            if ret !== nothing
                return ret
            end
        end
    end
    return
end

# 9.3.2.2
function _validate(x::AbstractDict, schema, ::Val{:patternProperties}, val::AbstractDict, path::String)
    for (k_val, v_val) in val
        r = Regex(k_val)
        for (k_x, v_x) in x
            if match(r, k_x) === nothing
                continue
            end
            ret = _validate(v_x, v_val, path * "[$(k_x)")
            if ret !== nothing
                return ret
            end
        end
    end
    return
end

# 9.3.2.3
function _validate(x::AbstractDict, schema, ::Val{:additionalProperties}, val::AbstractDict, path::String)
    properties = get(schema, "properties", Dict{String,Any}())
    patternProperties = get(schema, "patternProperties", Dict{String,Any}())
    for (k, v) in x
        if k in keys(properties) || any(r -> match(Regex(r), k) !== nothing, keys(patternProperties))
            continue
        end
        ret = _validate(v, val, path * "[$(k)]")
        if ret !== nothing
            return ret
        end
    end
    return
end

function _validate(x::AbstractDict, schema, ::Val{:additionalProperties}, val::Bool, path::String)
    if val
        return
    end
    properties = get(schema, "properties", Dict{String,Any}())
    patternProperties = get(schema, "patternProperties", Dict{String,Any}())
    for (k, v) in x
        if k in keys(properties) || any(r -> match(Regex(r), k) !== nothing, keys(patternProperties))
            continue
        end
        return SingleIssue(x, path, "additionalProperties", val)
    end
    return
end

# 9.3.2.4: unevaluatedProperties

# 9.3.2.5
function _validate(x::AbstractDict, schema, ::Val{:propertyNames}, val, path::String)
    for k in keys(x)
        ret = _validate(k, val, path)
        if ret !== nothing
            return ret
        end
    end
    return
end

###
### Checks for generic types.
###

# 6.1.1
function _validate(x, schema, ::Val{:type}, val::String, path::String)
    return !_is_type(x, Val{Symbol(val)}()) ? SingleIssue(x, path, "type", val) : nothing
end

function _validate(x, schema, ::Val{:type}, val::Vector, path::String)
    if !any(v -> _is_type(x, Val{Symbol(v)}()), val)
        return SingleIssue(x, path, "type", val)
    end
    return
end

_is_type(::Any, ::Val) = false
_is_type(::Array, ::Val{:array}) = true
_is_type(::Bool, ::Val{:boolean}) = true
_is_type(::Integer, ::Val{:integer}) = true
_is_type(::Real, ::Val{:number}) = true
_is_type(::Nothing, ::Val{:null}) = true
_is_type(::Missing, ::Val{:null}) = true
_is_type(::AbstractDict, ::Val{:object}) = true
_is_type(::String, ::Val{:string}) = true
# Note that Julia treat's Bool <: Number, but JSON-Schema distinguishes them.
_is_type(::Bool, ::Val{:number}) = false
_is_type(::Bool, ::Val{:integer}) = false

# 6.1.2
function _validate(x, schema, ::Val{:enum}, val, path::String)
    return !any(x == v for v in val) ? SingleIssue(x, path, "enum", val) : nothing
end

# 6.1.3
function _validate(x, schema, ::Val{:const}, val, path::String)
    return x != val ? SingleIssue(x, path, "const", val) : nothing
end

###
### Checks for numbers.
###

# 6.2.1
function _validate(x::Number, schema, ::Val{:multipleOf}, val::Number, path::String)
    y = isapprox(x / val, round(x / val))
    return !y ? SingleIssue(x, path, "multipleOf", val) : nothing
end

# 6.2.2
function _validate(x::Number, schema, ::Val{:maximum}, val::Number, path::String)
    return x > val ? SingleIssue(x, path, "maximum", val) : nothing
end

# 6.2.3
function _validate(x::Number, schema, ::Val{:exclusiveMaximum}, val::Number, path::String)
    return x >= val ? SingleIssue(x, path, "exclusiveMaximum", val) : nothing
end

function _validate(x::Number, schema, ::Val{:exclusiveMaximum}, val::Bool, path::String)
    if !val
        return
    end
    max = get(schema, "maximum", Inf)
    return x >= max ? SingleIssue(x, path, "exclusiveMaximum", val) : nothing
end

# 6.2.4
function _validate(x::Number, schema, ::Val{:minimum}, val::Number, path::String)
    return x < val ? SingleIssue(x, path, "minimum", val) : nothing
end

# 6.2.5
function _validate(x::Number, schema, ::Val{:exclusiveMinimum}, val::Number, path::String)
    return x <= val ? SingleIssue(x, path, "exclusiveMinimum", val) : nothing
end

function _validate(x::Number, schema, ::Val{:exclusiveMinimum}, val::Bool, path::String)
    if !val
        return
    end
    max = get(schema, "minimum", -Inf)
    return x <= max ? SingleIssue(x, path, "exclusiveMinimum", val) : nothing
end

###
### Checks for strings.
###

# 6.3.1
function _validate(x::String, schema, ::Val{:maxLength}, val::Integer, path::String)
    return length(x) > val ? SingleIssue(x, path, "maxLength", val) : nothing
end

# 6.3.2
function _validate(x::String, schema, ::Val{:minLength}, val::Integer, path::String)
    return length(x) < val ? SingleIssue(x, path, "minLength", val) : nothing
end

# 6.3.3
function _validate(x::String, schema, ::Val{:pattern}, val::String, path::String)
    y = occursin(Regex(val), x)
    return !y ? SingleIssue(x, path, "pattern", val) : nothing
end

###
### Checks for arrays.
###

# 6.4.1
function _validate(x::Vector, schema, ::Val{:maxItems}, val::Integer, path::String)
    return length(x) > val ? SingleIssue(x, path, "maxItems", val) : nothing
end

# 6.4.2
function _validate(x::Vector, schema, ::Val{:minItems}, val::Integer, path::String)
    return length(x) < val ? SingleIssue(x, path, "minItems", val) : nothing
end

# 6.4.3
function _validate(x::Vector, schema, ::Val{:uniqueItems}, val::Bool, path::String)
    # It isn't sufficient to just compare allunique on x, because Julia treats 0 == false,
    # but JSON distinguishes them.
    y = [(xx, typeof(xx)) for xx in x]
    return val && !allunique(y) ? SingleIssue(x, path, "uniqueItems", val) : nothing
end

# 6.4.4: maxContains

# 6.4.5: minContains

###
### Checks for objects.
###

# 6.5.1
function _validate(x::AbstractDict, schema, ::Val{:maxProperties}, val::Integer, path::String)
    return length(x) > val ? SingleIssue(x, path, "maxProperties", val) : nothing
end

# 6.5.2
function _validate(x::AbstractDict, schema, ::Val{:minProperties}, val::Integer, path::String)
    return length(x) < val ? SingleIssue(x, path, "minProperties", val) : nothing
end

# 6.5.3
function _validate(x::AbstractDict, schema, ::Val{:required}, val::Vector, path::String)
    return any(v -> !haskey(x, v), val) ? SingleIssue(x, path, "required", val) : nothing
end

# 6.5.4
function _validate(x::AbstractDict, schema, ::Val{:dependencies}, val::AbstractDict, path::String)
    for (k, v) in val
        if !haskey(x, k)
            continue
        elseif !_dependencies(x, path, v)
            return SingleIssue(x, path, "dependencies", val)
        end
    end
    return
end

function _dependencies(x::AbstractDict, path::String, val::Union{Bool,Dict})
    return _validate(x, val, path) === nothing
end
_dependencies(x::AbstractDict, path::String, val::Array) = all(v -> haskey(x, v), val)


text:
mutable struct Atoms
 # atoms struct to store the data from the atoms
     Natoms::Int64
     R::Array{Float64,2}
     sigma::Array{Float64,2}
     omega::Array{Float64,2}
     Eqdist::Array{Float64,2}
     mass::Array{Float64,2}
     Z::Array{Float64,2}
     nocc::Array{Int64,2}
     force::Array{Float64,2}
     function Atoms(Natoms, R, sigma, omega, Eqdist, mass, Z, nocc, force)
         return new(Natoms, R, sigma, omega, Eqdist, mass, Z, nocc, force)
     end
     #TODO:   this is a hack to make this thing work we need to come back and find
     #         how to properly define this function
     function Atoms(Natoms, R, sigma, omega, Eqdist, mass, Z, nocc)
         return new(Natoms, R, sigma, omega, Eqdist, mass, Z, nocc, 0*nocc)
     end
end

max_stars_repo_name: Forgotten/modelscf

max_issues_repo_name: Forgotten/modelscf

max_forks_repo_name: Forgotten/modelscf


text:
import Base: +, -, /, *, maximum, minimum, isequal, ==
using Statistics

# +

function +(vi1::VolumeImage, vi2::VolumeImage)

    dimensions_equal(vi1, vi2)

    @debug("Adding two volume images with $(size(vi1.data, 4)) time instances")

    vout = deepcopy(vi1)

    vout.data = vi1.data .+ vi2.data

    return vout
end

# -

function -(vi1::VolumeImage, vi2::VolumeImage)

    dimensions_equal(vi1, vi2)

    @debug("Subtracting two volume images with $(size(vi1.data, 4)) time instances")

    vout = deepcopy(vi1)

    vout.data = vi1.data .- vi2.data

    return vout
end

# /

function /(vi1::VolumeImage, vi2::VolumeImage)

    dimensions_equal(vi1, vi2)

    @debug("Dividing two volume images with $(size(vi1.data, 4)) time instances")

    vout = deepcopy(vi1)

    vout.data = vi1.data ./ vi2.data

    return vout
end

function /(vi::VolumeImage, c::Number)

    vout = deepcopy(vi)

    vout.data = vi.data ./ c

    return vout
end

# *

function *(vi::VolumeImage, c::Number)

    vout = deepcopy(vi)

    vout.data = vi.data .* c

    return vout
end

# mean

function mean(vi::VolumeImage)

    @debug("Taking mean of one volume images with $(size(vi.data, 4)) time instances")

    vout = deepcopy(vi)

    vout.data = Statistics.mean(vout.data, dims = 4)

    # Store time as 0 to indicate its been averaged
    vout.t = [NaN] * u"s"

    return vout
end

function mean(va::Array{VolumeImage,1})

    @debug(
        "Taking mean of $(length(va)) volume images with $(size(va[1].data, 4)) time instances"
    )

    mean_va = deepcopy(va[1])

    for i = 2:length(va)

        mean_va = mean_va + va[i]

    end

    return mean_va / length(va)
end

# maximum

function maximum(vi::VolumeImage)

    maximum(vi.data)
end

function maximum(vis::Array{VolumeImage})

    maximum([maximum(vi) for vi in vis])
end

# minimum

function minimum(vi::VolumeImage)

    minimum(vi.data)
end

function minimum(vis::Array{VolumeImage})

    minimum([minimum(vi) for vi in vis])
end

# normalise

function normalise(vi::VolumeImage)

    @debug("Normalising one volume images with $(size(vi.data, 4)) time instances")

    normalisation_constant = maximum(vi)

    vi = deepcopy(vi) / normalisation_constant

    vi.info["NormalisationConstant"] = normalisation_constant

    return vi
end

function normalise(va::Array{VolumeImage,1})

    @debug(
        "Normalising $(length(va)) volume images with $(size(va[1].data, 4)) time instances"
    )

    vo = deepcopy(va)
    for i = 1:length(vo)
        vo[i] = normalise(vo[i])
    end
    return vo
end

# isequal

function isequal(a::VolumeImage, b::VolumeImage)

    a.data == b.data ||
        a.units == b.units ||
        a.x == b.x ||
        a.y == b.y ||
        a.z == b.z ||
        a.t == b.t ||
        a.method == b.method ||
        a.coord_system == b.coord_system

end

function ==(a::VolumeImage, b::VolumeImage)

    isequal(a, b)
end

#
# Helper functions
# ----------------
#

function dimensions_equal(
    vi1::VolumeImage,
    vi2::VolumeImage;
    x::Bool = true,
    y::Bool = true,
    z::Bool = true,
    t::Bool = true,
    units::Bool = true,
    kwargs...,
)

    matching = true
    if x .& .!(vi1.x == vi2.x)
        throw(KeyError("X dimensions do not match"))
    end
    if y .& .!(vi1.y == vi2.y)
        throw(KeyError("Y dimensions do not match"))
    end
    if z .& .!(vi1.z == vi2.z)
        throw(KeyError("Z dimensions do not match"))
    end
    if t .& .!(vi1.t == vi2.t)
        throw(KeyError("T dimensions do not match"))
    end
    if units .& .!(vi1.units == vi2.units)
        throw(KeyError("Units do not match"))
    end

    if matching
        return true
    end
end

"""
Find indicies of location in VolumeImage
"""
function find_location(vi::VolumeImage, x::Real, y::Real, z::Real)

    x_loc = findall(minimum(abs.(ustrip(vi.x) .- x)) .== abs.(ustrip(vi.x) .- x))[1]
    y_loc = findall(minimum(abs.(ustrip(vi.y) .- y)) .== abs.(ustrip(vi.y) .- y))[1]
    z_loc = findall(minimum(abs.(ustrip(vi.z) .- z)) .== abs.(ustrip(vi.z) .- z))[1]

    if length(size(vi.data)) == 3
        return [x_loc, y_loc, z_loc]
    elseif length(size(vi.data)) == 4
        return [x_loc, y_loc, z_loc, 1]
    else
        return [NaN]
    end
end


text:
function mdl_tg_ode!(device_states,
                    output_ode,
                    device::DynGenerator{M, S, A, TGFixed, P})  where {M <: Machine,
                                                                       S <: Shaft,
                                                                       A <: AVR,
                                                                       P <: PSS}

    #Update inner vars
    device.inner_vars[τm_var] = get_P_ref(device) * device.tg.efficiency
    #@show device
    #@show device.inner_vars[τm_var]

    return

end

function mdl_tg_ode!(device_states,
                    output_ode,
                    device::DynGenerator{M, S, A, TGTypeI, P})  where {M <: Machine,
                                                                       S <: Shaft,
                                                                       A <: AVR,
                                                                       P <: PSS}

    #Obtain references
    ω_ref = get_ω_ref(device)
    P_ref = get_P_ref(device)

    #Obtain indices for component w/r to device
    local_ix = device.local_state_ix[device.tg]

    #Define internal states for component
    internal_states = @view device_states[local_ix]
    x_g1 = internal_states[1]
    x_g2 = internal_states[2]
    x_g3 = internal_states[3]

    #Obtain external states inputs for component
    external_ix = device.input_port_mapping[device.tg]
    ω = @view device_states[external_ix]

    #Get Parameters
    R = device.tg.R
    Ts = device.tg.Ts
    Tc = device.tg.Tc
    T3 = device.tg.T3
    T4 = device.tg.T4
    T5 = device.tg.T5
    P_min = device.tg.P_min
    P_max = device.tg.P_max

    #Compute auxiliary parameters
    P_in = P_ref + (1.0/R)*(ω_ref - ω[1])

    #Set anti-windup for P_in. NOT WORKING
    #if P_in > P_max
    #    P_in = P_max
    #elseif P_in < P_min
    #    P_in = P_min
    #end

    #Compute 3 States TG ODE:
    output_ode[local_ix[1]] = (1.0/Ts)*(P_in - x_g1)
    output_ode[local_ix[2]] = (1.0/Tc)*((1.0 - T3/Tc)*x_g1 - x_g2)
    output_ode[local_ix[3]] = (1.0/T5)*((1.0 - T4/T5)*(x_g2 + (T3/Tc)*x_g1) - x_g3)

    #Update mechanical torque
    device.inner_vars[τm_var] = x_g3 + (T4/T5)*(x_g2 + (T3/Tc)*x_g1)

    return
end

function mdl_tg_ode!(device_states,
                    output_ode,
                    device::DynGenerator{M, S, A, TGTypeII, P})  where {M <: Machine,
                                                                       S <: Shaft,
                                                                       A <: AVR,
                                                                       P <: PSS}

    #Obtain references
    ω_ref = get_ω_ref(device)
    P_ref = get_P_ref(device)

    #Obtain indices for component w/r to device
    local_ix = device.local_state_ix[device.tg]

    #Define internal states for component
    internal_states = @view device_states[local_ix]
    xg = internal_states[1]

    #Obtain external states inputs for component
    external_ix = device.input_port_mapping[device.tg]
    ω = @view device_states[external_ix]

    #Get Parameters
    R = device.tg.R
    T1 = device.tg.T1
    T2 = device.tg.T2
    τ_min = device.tg.τ_min
    τ_max = device.tg.τ_max

    #Compute auxiliary parameters
    τ_m = xg + (1.0/R)*(ω_ref - ω[1]) + P_ref/1.0

    #Set anti-windup for τ_m. NOT WORKING
    #if τ_m > τ_max
    #    τ_m = τ_max
    #elseif τ_m < τ_min
    #    τ_m = τ_min
    #end

    #Compute 1 State TG ODE:
    output_ode[local_ix[1]] = (1.0/T2)*( (1.0/R)*(1- T2/T1)*(ω_ref - ω[1]) - xg )

    #Update mechanical torque
    device.inner_vars[τm_var] = τ_m

    return
end


text:
include("../MNIST/MNIST.jl")

module ANNDemo

using MNIST, AF

export runDemo

include("ANN.jl")

function accuracy(predicted, target)
    pMaxArray, pMaxIndex = imax(predicted, 1)
    tMaxArray, tMaxIndex = imax(target, 1)
    (100.0f0 * count(Float32, pMaxIndex .== tMaxIndex)) / elements(tMaxIndex)
end

function runDemo(af)
    println("Setting up training data.")

    data = loadSubset(af)

    featureSize = DimT(elements(data.trainImages) / data.numTrain)

    println("Reshaping images into feature vectors.")
    trainFeats = transpose(moddims(data.trainImages, featureSize, data.numTrain))
    testFeats = transpose(moddims(data.testImages, featureSize, data.numTest))

    println("Creating targets.")
    trainTarget = transpose(data.trainLabels)
    testTarget = transpose(data.testLabels)

    println("Creating Network.")
    network = ANN(af, [dims(trainFeats, 1), 100, 50, data.numClasses])

    println("Warming up.")
    train(
        network,
        trainFeats,
        trainTarget,
        ANNTrainOptions(2.0f0, 1, 100, 0.0001f0))
    sync(af)

    println("Starting.")
    sec = @elapsed begin
        train(
            network,
            trainFeats,
            trainTarget,
            ANNTrainOptions(2.0f0, 250, 100, 0.0001f0))
        sync(af)
    end

    # Run the trained network and test accuracy.
    trainOutput = predict(network, trainFeats)
    testOutput = predict(network, testFeats)

    println("Training set:")
    println("Accuracy on training data: $(accuracy(trainOutput, trainTarget))")

    println("Test set:")
    println("Accuracy on testing  data: $(accuracy(testOutput, testTarget))")

    println("Training time: $sec seconds")
end

end

max_stars_repo_name: unbornchikken/julia-ml-proto

max_issues_repo_name: unbornchikken/julia-ml-proto

max_forks_repo_name: unbornchikken/julia-ml-proto


text:
export Canvas, canvas, png, svg, pdf

mutable struct Canvas
    size::Tuple{Int, Int}
    toplayer::Layer
    rect::Shape{Rect}
    color::Colors.Colorant
end

function canvas(
    width, height, toplayer::Union{Layer,Nothing}=nothing; color::Union{String, Colors.Colorant}="white")

    l = if isnothing(toplayer)
        layer() +
            Visible(true) +
            Markersize(3) +
            Marker(:o) +
            Fill("transparent") +
            Stroke("black") +
            Linewidth(1) +
            Linestyle(:solid) +
            Textfill("black")
    else
        toplayer
    end
    r = rect!(l, (0, 0), width, height, deg(0)) + Visible(false)
    color = typeof(color) <: String ? parse(Colors.Colorant, color) : color
    Canvas((width, height), l, r, color), l
end

function png(c::Canvas, filename::String; px_per_pt=1/0.75)
    csurface = draw_rgba(c, px_per_pt=px_per_pt)
    bufdata = UInt8[]
    iobuf = IOBuffer(bufdata, read=true, write=true)
    Cairo.write_to_png(csurface, iobuf)
    Cairo.finish(csurface)
    Cairo.destroy(csurface)
    write(filename, bufdata)
    nothing
end

function svg(c::Canvas, filename::String)
    draw_svg(c, filename)
end

function pdf(c::Canvas, filename::String)
    csurface = draw_pdf(c, filename)
    Cairo.finish(csurface)
    Cairo.destroy(csurface)
    nothing
end

function Base.show(io::IO, ::MIME"image/svg+xml", c::Canvas)
    svgbuffer = draw_svg(c)
    print(io, String(take!(svgbuffer)))
end

# function Base.show(io::IO, ::MIME"image/png", c::Canvas)
#     mktempdir() do path
#         p = joinpath(path, "layered.png")
#         png(c, p)
#         write(io, read(p))
#     end
# end


text:
# # Spiking-based plasticity rules under stimulation protocols (2 neurons)

#=

[Same as the original, but using NtwExpKernel struct and Recorder for spikes]

In this example, I show the effects of plasticity rules on a single
pre-post connection between the two neurons. The neural activity is
entirely regulated from the outside (e.g. the neurons do not interact
through their weights). This is to better illustrate plasticity rules
in their simplest form.
=#

# ## Initialization
using LinearAlgebra,Statistics,StatsBase,Distributions
using Plots,NamedColors ; theme(:default) ; gr()

using ProgressMeter
using Random
Random.seed!(0)

using HawkesSimulator; const global H = HawkesSimulator

## #src
# Pairwise STPD

# ## Stimulation protocol

# The neurons are forced to spike as in typical STDP measures. That is,
# they both fire at the same rate, but with a time difference set
# to a certain $\Delta t$.
function post_pre_spiketrains(rate::R,Δt_ro::R,Ttot::R;
    tstart::R = 0.05) where R
  post = collect(range(tstart,Ttot; step=inv(rate)))
  pre = post .- Δt_ro
  return [pre,post]
end

# let's take a look at one example: rate of 5.0 Hz, 20 ms difference in spike time
_ = let trains = post_pre_spiketrains(5.0,20E-3,1.0)
  plt = plot()
  scatter!(plt,trains[1], 0.1 .+ zero(trains[1]),markersize=30,
      markercolor=:black,markershape=:vline,leg=false)
  scatter!(plt,trains[2], 0.2 .+ zero(trains[2]),markersize=30,
      markercolor=:black,markershape=:vline,leg=false)
  plot!(plt,ylims=(0.0,0.3),xlims=(0,1),xlabel="time (s)")
end

# Now I define a population of input neurons using this input protocol.
# The Population object includes the weight matrix and the
# plasticity rules associated to it, both wrapped in the `Connection`
# object. I consider the connection as an input parameter of this function,
# to be set externally with the desired plasticity rule.
#
# Note that I define the population as `PopulationInputTestWeights` to
# indicate non-interacting weights
function post_pre_population(rate::Real,nreps::Integer,Δt_ro::Real,connection::H.Connection)
  Ttot = nreps/rate
  prepostspikes = post_pre_spiketrains(rate,Δt_ro,Ttot)
  gen = H.SGTrains(prepostspikes)
  state = H.PopulationState(H.InputUnit(gen),2)
  return H.PopulationInputTestWeights(state,connection)
end

# ## Constructors
# Here I define a function that tests the plasticity rule.
# That is, given a $Δt$, it creates the input neurons, then
# it defines a network, and iterates it dynamically for
# a certain number of pre-post repetitions.
# Finally it outputs the weight total change divided by time
# (so it's a weight change per second)
#
# Since both neurons have the same plasticity rule, a
# positive pre-post $\Delta t$ for neuron A, impacting the
#  $w_{\text{AB}}$ weight, is the
# equivalent of a negative pre-post $\Delta t$ for neuron B,
# impacting the $w_{\text{BA}}$ weight.

function test_stpd_rule(Δt::Real,connection::H.Connection;
    nreps=510,wstart=100.0,rate=0.5)
  num_spikes = nreps*2 - 10 # a bit less than total, for safety
  recorder = H.RecFullTrain(num_spikes)
  population = post_pre_population(rate,nreps,Δt,connection)
  network = H.RecurrentNetworkExpKernel(population,recorder)
  wmat = connection.weights
  fill!(wmat,wstart)
  wmat[diagind(wmat)] .= 0.0
  t_now = 0.0
  H.reset!.((network,recorder,connection)) # clear spike trains etc
  for _ in 1:num_spikes
    t_now = H.dynamics_step_singlepopulation!(t_now,network)
  end
  w12,w21 = wmat[1,2],wmat[2,1]
  H.reset!(network)
  dw12 = (w12-wstart)/t_now
  dw21 = (w21-wstart)/t_now
  return dw12,dw21
end

# ## STDP rule
# Here I define the plasticty type and the parameters that I want to test.
# I choose the pairwise STDP rule.
# I initialize the weight matrix to 100.0
#
# Note once again that neurons are not interacting. The sole purpose
# of the "dummy" weights it to be changed by plasticity.

connection_test = let wmat =  fill(100.0,2,2)
  wmat[diagind(wmat)] .= 0.0
  τplus = 10E-3
  τminus = 10E-3
  Aplus = 1E-1
  Aminus = -1E-1
  npost,npre = size(wmat)
  stdp_plasticity = H.PairSTDP(τplus,τminus,Aplus,Aminus,npost,npre)
  H.ConnectionWeights(wmat,stdp_plasticity)
end

# let's do one run
dw12,dw21 = test_stpd_rule(2E-3,connection_test)
println("change in w12 : $(dw12), change in w21 $(dw21)")

# Neuron 1 spikes before neuron 2, therefore connection from 1 to 2, dw21 is
# potentiated, while connection from 2 to 1, dw12 is depressed.

# ## STDP curve

# Compute and plot the weight changes due to STDP for varying $\Delta t$s

nsteps = 200
deltats = range(0.1E-3,100E-3;length=nsteps)
deltats_all = vcat(-reverse(deltats),deltats)
out = Vector{Float64}(undef,2*nsteps)

for (i,Δt) in enumerate(deltats)
  dw12,dw21 = test_stpd_rule(Δt,connection_test)
  out[nsteps+i]=dw21
  out[nsteps-i+1]=dw12
end

plot(deltats_all,out;leg=false,xlabel="Delta t",ylabel="dw/dt",linewidth=3)

# **THE END**

# Literate.markdown("examples/plasticty_STDP.jl","docs/src";documenter=true,repo_root_url="https://github.com/dylanfesta/HawkesSimulator.jl/blob/master") #src


text:
# This file is part of IntegerSequences.
# Copyright Peter Luschny. License is MIT.

(@__DIR__) ∉ LOAD_PATH && push!(LOAD_PATH, (@__DIR__))

module NodesAndEdges

using Nemo, NumberTheory, IntPartitions
export V000055, V000081, V000106, V209397, V217420, V095350, V002494
export V055542, V055543, V055544, V000088, L087803, L036039

"""
Rooted trees and similar topics.

* V000055, V000081, V000106, V209397, V217420, V095350, V002494, V055542, V055543, V055544, V000088, L087803, L036039
"""
const ModuleNodesAndEdges = ""

"""

Return the number of unlabeled rooted trees with ``n`` nodes.
"""
function V000081(n::Int)
    global CacheV000081
    n <= 1 && return BigInt(n)
    haskey(CacheV000081, n) && return CacheV000081[n]
    S = sum(V209397(j)*V000081(n-j) for j in 1:n-1)
    return CacheV000081[n] = div(S, n-1)
end

const CacheV000081 = Dict{Int, BigInt}()

"""

Return ``∑_{d|n}`` d V000081(d).
"""
function V209397(n::Int)
    global CacheV209397
    haskey(CacheV209397, n) && return CacheV209397[n]
    S = sum(d*V000081(d) for d in divisors(n))
    return CacheV209397[n] = BigInt(S)
end

const CacheV209397 = Dict{Int, BigInt}()

"""

Return the number of linear forests of 2-rooted trees.
"""
V000106(n) = sum(V000081(k) * V000081(n-k) for k in 0:n)

"""

Return the number of rooted unlabeled trees where the root node has degree 2 and both branches are distinct.
"""
function V217420(n)
    n <= 0 && return BigInt(0)
    c = isodd(n) ? V000081(div(n-1, 2)) : 0
    div(V000106(n-1) - c, 2)
end

"""

Return the number of trees with ``n`` unlabeled nodes.
"""
function V000055(n)
    n == 0 && return BigInt(1)
    V000081(n) - V217420(n+1)
end

"""

Return total number of edges in all rooted trees with ``n`` nodes.
"""
V095350(n) = (n-1)*V000081(n)

"""

Total number of nodes in all rooted trees with ``n`` nodes.
"""
V055544(n) = n*V000081(n)

"""

Total number of nodes in all trees with n nodes.
"""
V055543(n) = n*V000055(n)

"""

Return the number of permutations whose cycle structure is the given list.
"""
function PermutationCount(v)
    m, k = 1, 0
    s = BigInt.(0)
    for i ∈ 1:length(v)
        t = v[i]
        k = i > 1 && t == v[i-1] ? k + 1 : 1
        m *= t*k
        s += t
    end
    div(factorial(s), m)
end

"""

Return the multinomial coefficients of the integer partitions of ``n``. Gives the number of permutations whose cycle structure is the given partition.
# Examples
```julia-repl
julia> L036039(5)
[24, 30, 20, 20, 15, 10, 1]
```
"""
L036039(n) = [PermutationCount(p) for p in IntegerPartitions(n)]

#=
[1]
[1, 1]
[2, 3, 1]
[6, 8, 3, 6, 1]
[24, 30, 20, 20, 15, 10, 1]
=#

# Missing in the OEIS
function Edges(v)
    (len = length(v)) == 1 && return div(v[1], 2)
    (sum(gcd(v[i], v[j]) for i in 2:len for j in 1:i-1 ) +
        sum(div(v[i], 2) for i in 1:len))
end

#for n in 1:6
#    [Edges(p) for p in IntegerPartitions(n)] |> println
#end

"""

Return the number of graphs on ``n`` unlabeled nodes.
"""
function V000088(n)
    n == 0 && return BigInt(1)
    s = sum(PermutationCount(p)*2^Edges(p) for p in IntegerPartitions(n))
    div(s, factorial(n))
end

"""

Return the total number of nodes in all simple graphs of ``n`` nodes.
"""
V055542(n) = n*V000088(n)

"""

Return the number of ``n``-node graphs without isolated nodes.
"""
V002494(n) = n == 0 ? BigInt(1) : V000088(n) - V000088(n-1)

"""

Number of unlabeled rooted trees with at most n nodes.
"""
L087803(len) = cumsum([V000081(k) for k in 0:len])

# See also:: V000595, V001349, V275331,
# V275330, V008406, V000273, V003024, V057500, V240955

#START-TEST-########################################################

using Test

function test()
    @testset "Nodes" begin
        @test V000055(10) == 106
        @test V000081(10) == 719
        @test V000106(10) == 1235
        @test V095350(10) == 6471
        @test V209397(10) == 7238
        @test V217420(10) == 239
        @test V055544(10) == 7190
        @test V055543(10) == 1060
        @test V055542(10) == 120051680
        @test V000088(10) == 12005168
        @test V002494(10) == 11730500
        @test L087803(5)  == BigInt[0, 1, 2, 4, 8, 17]
        @test L036039(5)  == BigInt[24, 30, 20, 20, 15, 10, 1]
    end
end

function demo()
    println("V000106")
    [V000106(n) for n ∈ 0:12] |> println

    println("V209397")
    [V209397(n) for n ∈ 1:12] |> println

    println("V000081")
    [V000081(n) for n ∈ 0:12] |> println

    println("V000055")
    [V000055(n) for n ∈ 0:12] |> println

    println("V000088")
    [V000088(n) for n ∈ 0:12] |> println

    println("V095350")
    [V095350(n) for n ∈ 0:12] |> println

    println("V217420")
    [V217420(n) for n ∈ 0:12] |> println

    println("V055542")
    [V055542(n) for n ∈ 0:12] |> println

    println("V055543")
    [V055543(n) for n ∈ 0:12] |> println

    println("V055544")
    [V055544(n) for n ∈ 0:12] |> println

    println("V002494")
    [V002494(n) for n ∈ 0:12] |> println

    println("L087803")
    L087803(5) |> println

    println("L036039")
    L036039(5) |> println

end

#  0.101582 seconds (285.83 k allocations: 13.003 MiB)
function perf()
    @time [V000055(n) for n in 0:99]
end

function main()
    test()
    demo()
    perf()
end

main()

end # module

max_stars_repo_name: OpenLibMathSeq/Sequences

max_issues_repo_name: OpenLibMathSeq/Sequences


text:
using MergeIterator
using Base.Test

# write your own tests here
@test collect(merge_iterator([1,3,4,6],[2,5,7])) == [1,2,3,4,5,6,7]