text:
# module EquivariantNeuralOperators

using LinearAlgebra
# using Rotations
using Zygote
using Random
using Flux

include("./utils.jl")
Random.seed!(1)

"""
equivariant convolution params
"""
# struct EquivConv{F,S<:AbstractArray,T<:Union{Zeros, AbstractVector}}
struct EquivConv{T<:AbstractFloat}
    paths::AbstractVector{LinearOperator}
    pathsmap::AbstractVector
    dx::T
    rmax::T
    sz
end

Flux.trainable(m::EquivConv) = [x.params for x in m.paths]

"""
initialize params
"""

function EquivConv(
    inranks::AbstractVector,
    outranks::AbstractVector,
    dx::Float64;
    rmax = 1/SMALL,
    sz=nothing,
    lmax = 1,
    paths = Vector{LinearOperator}(),
)
    Zygote.ignore() do
        nin = length(inranks)
        nout = length(outranks)
        inparities = (-1) .^ inranks
        outparities = (-1) .^ outranks

        # if sz!==nothing
        #     sz=2 .*sz .-1
        # end

        # iterate all tensor product paths
        if isempty(paths)
            for lr = 0:lmax
                for (i, li, si) in zip(1:nin, inranks, inparities)
                    loutrange = abs(li - lr):min(li + lr, lmax)
                    parity = si * (-1)^lr
                    for (o, lo) in enumerate(outranks)
                        if lo in loutrange && outparities[o] == parity
                            path = LinearOperator(
                                dx;
                                li = li,
                                lo = lo,
                                lr = lr,
                                i = i,
                                o = o,
                                rmax = rmax,
                                sz=sz
                            )
                            push!(paths, path)
                        end
                    end
                end
            end
        end

        pathsmap = [[] for i = 1:nout]
        for (i, path) in enumerate(paths)
            push!(pathsmap[path.o], i)
        end

        return EquivConv(paths, pathsmap, dx, rmax,sz)
    end
end

"""
layer function
"""
function (f::EquivConv)(
    X::AbstractVector;
    remake = true,
) where {T<:AbstractFloat}
    paths, pathsmap, dx, rmax,sz = f.paths, f.pathsmap, f.dx, f.rmax,f.sz
    sx = size(X[1][1])

if remake
    for x in paths
        remake!(x;dx=dx,rmax=rmax)
    end
    end

    Y = [
        dropdims(
            sum(
                hcat([paths[i](X[paths[i].i]) for i in pathsmap[o]]...),
                dims = 2,
            ),
            dims = 2,
        ) for o in eachindex(pathsmap)
    ]
    return Y
end

##
mutable struct Prod
    i1::Int
    i2::Int
    o::Int
    li1::Int
    li2::Int
    lo::Int
    prod::Function
    params::Any
end

function Prod(l1, l2, l)
    if l1 == 0
        return (x, y) -> [x[1] .* y for y in y]
    elseif l2 == 0
        return (x, y) -> [x .* y[1] for x in x]
    elseif l1 == l2 && l == 0
        return (x, y) -> [sum([x .* y for (x, y) in zip(x, y)])]
    end
end

"""
equivariant params
"""
struct EquivAttn
    paths::Vector{Prod}
    pathsmap::AbstractVector
end
Flux.trainable(m::EquivAttn) = [x.params for x in m.paths]

"""
initialize params
"""

function EquivAttn(
    inranks::AbstractVector{Int},
    outranks::AbstractVector{Int},
    ;
    lmax = 1,
    paths = Vector{Prod}(),
)
    Zygote.ignore() do
        nin = length(inranks)
        nout = length(outranks)
        println(inranks)
        inparities = (-1) .^ inranks
        outparities = (-1) .^ outranks

        # iterate all tensor product paths
        if isempty(paths)
                for (i1, li1, si1) in zip(1:nin, inranks, inparities)
                    for (i2, li2, si2) in
                        zip(1:i1, inranks[1:i1], inparities[1:i1])
                        # for (i2, li2, si2) in zip(
                        #     0:nin,
                        #     vcat([0], inranks[1:i1]),
                        #     vcat([1], inparities[1:i1]),
                        # )
                        loutrange = abs(li1 - li2):min(li1 + li2, lmax)
                        parity = si1 * si2
                        for (o, lo) in enumerate(outranks)
                            if lo in loutrange && outparities[o] == parity
                                path = Prod(
                                    i1,
                                    i2,
                                    o,
                                    li1,
                                    li2,
                                    lo,
                                    Prod(li1, li2, lo),
                                    ones(4, 1),
                                )
                                push!(paths, path)
                        end
                    end
                end
            end
        end

        pathsmap = [[] for i = 1:nout]
        for (i, path) in enumerate(paths)
            push!(pathsmap[path.o], i)
        end

        return EquivAttn(paths, pathsmap)
    end
end

"""
layer function
"""
function (f::EquivAttn)(X::AbstractVector;) where {T<:AbstractFloat}
    paths, pathsmap = f.paths, f.pathsmap
    sz = size(X[1][1])

    Y = [
        dropdims(
            sum(
                hcat(
                    [
                        paths[i].params[1] * paths[i].prod(
                            X[paths[i].i1],
                            paths[i].i2 === 0 ? [ones(sz...)] : X[paths[i].i2],
                        ) for i in pathsmap[o]
                    ]...,
                ),
                dims = 2,
            ),
            dims = 2,
        ) for o in eachindex(pathsmap)
    ]
    return Y
end

# end # module

max_stars_repo_name: paulxshen/eno

max_issues_repo_name: paulxshen/eno

max_forks_repo_name: paulxshen/eno


text:
using Test
using DataFrames, TableOperations, Tables, Random
using Recommenders:
    Movielens100k,
    load_dataset,
    ratio_split,
    SLIM,
    evaluate_u2i,
    MeanPrecision,
    MeanRecall,
    MeanNDCG,
    EvaluateValidData

ml100k = Movielens100k()
download(ml100k)
rating, _, _ = load_dataset(ml100k)
rating = rating |> TableOperations.transform(Dict(:rating => x -> 1.0))

Random.seed!(1234);
train_table, test_table = ratio_split(rating, 0.8)

prec10 = MeanPrecision(10)
recall10 = MeanRecall(10)
ndcg10 = MeanNDCG(10)
metrics = [prec10, recall10, ndcg10]

model = SLIM(1e-4, 1e-4, 10)

result = evaluate_u2i(
    model,
    train_table,
    test_table,
    metrics,
    10,
    drop_history = true,
    col_item = :movieid,
    cd_tol = 1e-1,
    nλ = 2,
    verbose = -1,
)

@test !(model.W === nothing)

@test result[:ndcg10] >= 0
@test result[:precision10] >= 0
@test result[:recall10] >= 0

pred_u2i = predict_u2i(model, 1, 10)
@test length(pred_u2i) == 10

pred_u2i = predict_u2i(model, [1, 10], 10)
for p in pred_u2i
    @test length(p) == 10
end

pred_i2i = predict_i2i(model, 222, 10)
@test length(pred_i2i) == 10
@show pred_i2i

pred_i2i = predict_i2i(model, [1, 10], 10)
for p in pred_i2i
    @test length(p) == 10
end

tempfile = tempname()
save_model(model, tempfile)
@test isfile(tempfile)
model = load_model(tempfile)
@test model.W !== nothing

ext: jl


text:
abstract type AbstractNetwork end

show(network::AbstractNetwork) = println("Network\n- N = $(network.N)")
display(network::AbstractNetwork) = show(network)


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

using Markdown
using InteractiveUtils

# ╔═╡ f394ac3a-f787-11ea-314f-bf1bac91a4e9
using Pkg, DrWatson

# ╔═╡ f394e7e2-f787-11ea-1c4e-71f58d4abbea
begin
	@quickactivate "SR2TuringPluto"
	using StatisticalRethinking
end

# ╔═╡ 7c07959e-f787-11ea-14c6-23630e43ea67
md"## Fig3.4t.jl"

# ╔═╡ f3956618-f787-11ea-2032-fbf3b79ef810
p_grid = range(0, stop=1, length=10000)

# ╔═╡ 8171efda-f788-11ea-2c6b-834d66a19f38
prior = ones(length(p_grid));

# ╔═╡ 8172298c-f788-11ea-1cc5-4779c3e53066
likelihood = pdf.(Binomial.(3, p_grid), 3)

# ╔═╡ 8172c05c-f788-11ea-0fe9-d513d34fd0cd
begin
	posterior = likelihood .* prior
	posterior = posterior / sum(posterior)
end

# ╔═╡ 8182be1c-f788-11ea-3fba-8966e136c7d4
samples = sample(p_grid, Weights(posterior), length(p_grid));

# ╔═╡ f3a1f566-f787-11ea-1901-21b24e86c6ba
begin
	p1 = density(samples;
	  xlab="Proportion water (p)", ylab="Density", lab="density",
	  title="Mean, median & mode", leg=:topleft);
	vline!(p1, [mode(samples)], lab="mode")
	vline!(p1, [median(samples)], lab="median")
	vline!(p1, [mean(samples)], lab="mean")
end;

# ╔═╡ f3a27664-f787-11ea-0a07-0f5212c9a2b7
begin
	loss = map(p -> sum(posterior .* abs.(p .- p_grid)), p_grid)
	m = findmin(loss)
	p2 = plot(loss;
	  xlab="Decision (p)", ylab="Expected proportional loss",
	  title="Loss value", lab="Loss function");
	scatter!([m[2]], [m[1]], lab="p_grid[$(m[2])]=$(round(m[1], digits=3))")
end;

# ╔═╡ f3b12006-f787-11ea-3f0c-27f71c4ec2da
plot(p1, p2, layout=(1, 2))

# ╔═╡ f3b1aed6-f787-11ea-30cf-515030a997c3
md"## End of Fig3.4t.jl"

# ╔═╡ Cell order:
# ╟─7c07959e-f787-11ea-14c6-23630e43ea67
# ╠═f394ac3a-f787-11ea-314f-bf1bac91a4e9
# ╠═f394e7e2-f787-11ea-1c4e-71f58d4abbea
# ╠═f3956618-f787-11ea-2032-fbf3b79ef810
# ╠═8171efda-f788-11ea-2c6b-834d66a19f38
# ╠═8172298c-f788-11ea-1cc5-4779c3e53066
# ╠═8172c05c-f788-11ea-0fe9-d513d34fd0cd
# ╠═8182be1c-f788-11ea-3fba-8966e136c7d4
# ╠═f3a1f566-f787-11ea-1901-21b24e86c6ba
# ╠═f3a27664-f787-11ea-0a07-0f5212c9a2b7
# ╠═f3b12006-f787-11ea-3f0c-27f71c4ec2da
# ╟─f3b1aed6-f787-11ea-30cf-515030a997c3


text:
using JuMP
using Clp
using Test
using JuMPLPUtils
import Random
Random.seed!(1234)
A = rand(40, 10)
c = rand(10)
b = rand(40)
m = Model()
J = 1:10
I = 1:40
@variable(m, x[J] >= 0)

#minimize the cost
z = @objective(m, Min, sum(c[j]*x[j] for j in J))
#Satisfy the constraints
@constraint(m, constraint[i in I], sum(A[i, j]*x[j] for j in J) >= b[i])
setsolver(m, ClpSolver(LogLevel = 0))
solve(m)

N = [1,3,6,7,8,9]
B = [2,4,5,10]
x_val =getvalue(x).innerArray

@test sum(x_val[N]) <= 1e-6
@test abs(sum(x_val[B]) -3.179) <= 1e-3
@test abs(getslack(constraint[10])-.775) <= 1e-3

s = A*x_val - b
@test sum(abs(s[i] -getslack(constraint[i])) for i in I) <= 1e-10

c_int, b_int = getoptranges(m)
z_val = getvalue(z)
pi = getdual(constraint)
for j in J
    dc = .8*max(c_int[j,1],-1.0)
    c[j] += dc
    @objective(m, Min, sum(c[j]*x[j] for j in J))
    solve(m)
    @test sum(abs(x_val[j] - getvalue(x)[j]) for j in J) <= 1e-10
    c[j] -= dc
    dc = .8*min(c_int[j,2],1.0)
    c[j] += dc
    @objective(m, Min, sum(c[j]*x[j] for j in J))
    solve(m)
    @test sum(abs(x_val[j] - getvalue(x)[j]) for j in J) <= 1e-10
    c[j] -= dc
end
@objective(m, Min, sum(c[j]*x[j] for j in J))

for i in I
    db = .8*max(b_int[i,1],-1.0)
    JuMP.setRHS(constraint[i], b[i] + db)
    solve(m)
    @test sum(getvalue(x).innerArray[N]) <= 1e-10
    @test abs(z_val - getvalue(z) + db*pi[i]) <= 1e-10
    db= .8*min(b_int[i,2], 1.0)
    JuMP.setRHS(constraint[i], b[i] + db)
    solve(m)
    @test sum(getvalue(x).innerArray[N]) <= 1e-10
    @test abs(z_val - getvalue(z) + db*pi[i]) <= 1e-10
    JuMP.setRHS(constraint[i], b[i])
end

m = Model()
@variable(m, x[J] >= 0)
#minimize the cost
z = @objective(m, Max, sum(c[j]*x[j] for j in J))
#Satisfy the constraints
@constraint(m, constraint[i in I], sum(A[i, j]*x[j] for j in J) <= b[i])
setsolver(m, ClpSolver(LogLevel = 0))
solve(m)

N = [1,2,4,5,6,7,8,9,10]
B = [3]
x_val =getvalue(x).innerArray

@test sum(x_val[N]) <= 1e-6
@test abs(sum(x_val[B]) - .063688) <= 1e-3
@test abs(getslack(constraint[10])-.5704) <= 1e-3

s = b- A*x_val
@test sum(abs(s[i] -getslack(constraint[i])) for i in I) <= 1e-10

c_int, b_int = getoptranges(m)
z_val = getvalue(z)
pi = getdual(constraint)
for j in J
    dc = .99*max(c_int[j,1],-1.0)
    c[j] += dc
    @objective(m, Max, sum(c[j]*x[j] for j in J))
    solve(m)
    @test sum(abs(x_val[j] - getvalue(x)[j]) for j in J) <= 1e-10
    c[j] -= dc
    dc = .99*min(c_int[j,2],1.0)
    c[j] += dc
    @objective(m, Max, sum(c[j]*x[j] for j in J))
    solve(m)
    @test sum(abs(x_val[j] - getvalue(x)[j]) for j in J) <= 1e-10
    c[j] -= dc
end
@objective(m, Max, sum(c[j]*x[j] for j in J))

for i in I
    db = .99*max(b_int[i,1],-1.0)
    JuMP.setRHS(constraint[i], b[i] + db)
    solve(m)
    @test sum(getvalue(x).innerArray[N]) <= 1e-10
    @test abs(z_val - getvalue(z) + db*pi[i]) <= 1e-10
    db= .99*min(b_int[i,2], 1.0)
    JuMP.setRHS(constraint[i], b[i] + db)
    solve(m)
    @test sum(getvalue(x).innerArray[N]) <= 1e-10
    @test abs(z_val - getvalue(z) + db*pi[i]) <= 1e-10
    JuMP.setRHS(constraint[i], b[i])
end

max_stars_repo_name: EdvinAblad/JuMPLPUtils

max_issues_repo_name: EdvinAblad/JuMPLPUtils

max_forks_repo_name: EdvinAblad/JuMPLPUtils


text:
using SparseArrays
using LinearAlgebra

struct svparams
	K::Function
	S::Function
	In::Array{Float64,2}
end

# Julia versions prior to v"1.3.1" can't use LinearAlgebra's 5 argument mul!, routines
# included here for backwards compatability
if(VERSION < v"1.3.1")
    include("backwards_compat.jl")
end

# Method to generate the coefficients for the adjoint
# method for a given RK scheme (not PRK)
function adjoint_tableau(A,b,c)
	#If any weights vanish, exit immediately
	if(any(b.==0.0))
		println("Please feed me non-zero weights.")
		return
	end

	m,n = size(A)
	B = zeros(m,n)
	for j = 1:n
		for i = 1:m
			B[i,j] = b[j] - b[j]*A[j,i]/b[i]
		end
	end
	return B
end

# Force function
@inline function step(stepper::svparams, t::Float64, u::Array{Float64,N},
	v::Array{Float64,N}, h::Float64, uforce::Function, vforce::Function) where N
	uforce0 = uforce(t)
	vforce0 = vforce(t)

	vforce05 = vforce(t+0.5*h)
	vforce1 = vforce(t + h)
	uforce1 = uforce(t + h)

    t,u,v,v05 = step(stepper, t, u, v, h, uforce0, vforce05, uforce1)

	return t, u, v, v05
end

# wWthout forcing
@inline function step(stepper::svparams, t::Float64, u::Array{Float64,N}, v::Array{Float64,N}, h::Float64) where N
	Ntot = size(u)
	vforce0 = zeros(Ntot)
	uforce05 = zeros(Ntot)
	vforce1 = zeros(Ntot)

    t,u,v, v05 = step(stepper, t, u, v, h, vforce0, uforce05, vforce1)

	return t, u, v, v05
end

# Forcing as array
@inline function step(stepper::svparams, t::Float64, u::Array{Float64,N}, v::Array{Float64,N},
	h::Float64, uforce0::Array{Float64,N}, vforce05::Array{Float64,N}, uforce1::Array{Float64,N}) where N
	S = stepper.S
	K = stepper.K
	In = stepper.In

	K0  = K(t)
	S0  = S(t)
	K05 = K(t +0.5*h)
	S05 = S(t + 0.5*h)
	K1  = K(t+h)
	S1  = S(t+h)

 	rhs = K05*u .+  S05*v + vforce05
	l1 = (In .-  0.5*h.*S05)\rhs
	v05 = (v .+ 0.5*h.*l1)
	kappa1 = S0*u .- K0*v05 + uforce0
	rhs = S1*(u .+ 0.5*h*kappa1) .- K1*v05 + uforce1
	kappa2 = (In .- (0.5*h).*S1)\rhs

	u = u .+ (0.5*h).*(kappa1 .+ kappa2)
	l2 = K05*u .+ S05*v05 + vforce05

	v = v .+ 0.5*h.*(l1 .+ l2)
	t = t + h
	return t, u, v, v05
end

###########################################################
# FMG
# Here we attempt an implementation of an explicit RK time
# integrator. Remember: return intermediate stages
@inline function explicit_step(stepper::svparams,t::Float64, u::Array{Float64,N},
	    v::Array{Float64,N}, h::Float64, uforce::Function, vforce::Function) where N
	S  = stepper.S
	K  = stepper.K
	In = stepper.In
	A  = stepper.A
	b  = stepper.b
	c  = stepper.c
	n_stages = stepper.n_stages

	# Arrays for storing intermediate stages
	u_stage = zeros(length(u),n_stages)
	k_mat_u = zeros(length(u),n_stages)
	v_stage = zeros(length(v),n_stages)
	k_mat_v = zeros(length(v),n_stages)

	St = S(t)
	Kt = K(t)

	# First stage is always the current solution
	k_mat_u[:,1] = St*u - Kt*v
	k_mat_v[:,1] = Kt*u + St*v
	u_stage[:,1] = u
	v_stage[:,1] = v
	for s = 2:n_stages
		# Build intermediate stage
		u_stage[:,s] = u
		v_stage[:,s] = v
		for ss = 1:s-1
			u_stage[:,s] += h*A[s,ss]*k_mat_u[:,ss]
			v_stage[:,s] += h*A[s,ss]*k_mat_v[:,ss]
		end

		St = S(t+c[s]*h)
		Kt = K(t+c[s]*h)
		k_mat_u[:,s] = St*u_stage[:,s] - Kt*v_stage[:,s]
		k_mat_v[:,s] = Kt*u_stage[:,s] + St*v_stage[:,s]
	end

	# Step forward in time
	for s = 1:n_stages
		u += h*b[s]*k_mat_u[:,s]
		v += h*b[s]*k_mat_v[:,s]
	end

	# FMG: Must return intermediate stages.
	t = t + h
	return t, u, v
end
###########################################################

# Forward gradient calculation.
@inline function step_fwdGrad!(t::Float64, nNeumann::Int64, u::Array{Float64,N}, v::Array{Float64,N},v05::Array{Float64,N},
	h::Float64, uforce0::Array{Float64,N}, vforce0::Array{Float64,N},uforce1::Array{Float64,N},
	vforce1::Array{Float64,N},K0::Array{Float64,N},
	S0::Array{Float64,N},K05::Array{Float64,N},S05::Array{Float64,N},
 	K1::Array{Float64,N},S1::Array{Float64,N},In::Array{Float64,N},
 	κ₁::Array{Float64,N}, κ₂::Array{Float64,N},ℓ₁::Array{Float64,N},
 	ℓ₂::Array{Float64,N},rhs::Array{Float64,N}) where N

 	# rhs .= K05*u .+  S05*v .+ vforce0
 	LinearAlgebra.mul!(rhs,K05,u)
 	mul!(rhs,S05,v,1.0,1.0)
	LinearAlgebra.axpy!(1.0,vforce0,rhs)

	# ℓ₁     .= (In .-  0.5*h.*S05)\rhs
	neumann!(nNeumann,h,S05,rhs,v05,ℓ₁)

	# v05    .= v .+ 0.5*h*ℓ₁
	copy!(v05,v)
	LinearAlgebra.axpy!(0.5*h,ℓ₁,v05)

	# κ₁  .= S0*u .- K0*v05 .+ uforce0
	LinearAlgebra.mul!(κ₁,S0,u)
	mul!(κ₁,K0,v05,-1.0,1.0)
	LinearAlgebra.axpy!(1.0,uforce0,κ₁)

	# rhs    .= S1*(u .+ 0.5*h*κ₁) .- K1*v05
	LinearAlgebra.mul!(rhs, S1,u)
	mul!(rhs,S1,κ₁,0.5*h,1.0)
	mul!(rhs,K1,v05,-1.0,1.0)
	LinearAlgebra.axpy!(1.0,uforce1,rhs)

	# u      .= u .+ (0.5*h).*κ₁
	LinearAlgebra.axpy!(0.5*h,κ₁,u)

	# κ₂     .= (In .- (0.5*h)*S1)\rhs
	neumann!(nNeumann,h,S1,rhs,κ₁,κ₂)

	# u      .= u .+ (0.5*h).*κ₂
	LinearAlgebra.axpy!(0.5*h,κ₂,u)

	# ℓ₂  .= K05*u .+ S05*v05 .+ vforce1
	LinearAlgebra.mul!(ℓ₂,K05,u)
	mul!(ℓ₂,S05,v05,1.0,1.0)
	LinearAlgebra.axpy!(1.0,vforce1,ℓ₂)

	v      .= v .+ 0.5*h.*(ℓ₁ .+ ℓ₂)

	t      = t + h
end

# sparse version of step_fwdGrad above
@inline function step_fwdGrad!(t::Float64, nNeumann::Int64, u::Array{Float64,N}, v::Array{Float64,N}, v05::Array{Float64,N},
	h::Float64, uforce0::Array{Float64,N}, vforce0::Array{Float64,N},uforce1::Array{Float64,N},
	vforce1::Array{Float64,N},K0::SparseMatrixCSC{Float64,Int64},
	S0::SparseMatrixCSC{Float64,Int64},K05::SparseMatrixCSC{Float64,Int64},S05::SparseMatrixCSC{Float64,Int64},
 	K1::SparseMatrixCSC{Float64,Int64},S1::SparseMatrixCSC{Float64,Int64},In::SparseMatrixCSC{Float64,Int64},
 	κ₁::Array{Float64,N}, κ₂::Array{Float64,N},ℓ₁::Array{Float64,N},
 	ℓ₂::Array{Float64,N},rhs::Array{Float64,N}) where N

 	# rhs .= K05*u .+  S05*v .+ vforce0
 	LinearAlgebra.mul!(rhs,K05,u)
 	mul!(rhs,S05,v,1.0,1.0)
	LinearAlgebra.axpy!(1.0,vforce0,rhs)

	# ℓ₁     .= (In .-  0.5*h.*S05)\rhs
	neumann!(nNeumann,h,S05,rhs,v05,ℓ₁)

	# v05    .= v .+ 0.5*h*ℓ₁
	copy!(v05,v)
	LinearAlgebra.axpy!(0.5*h,ℓ₁,v05)

	# κ₁  .= S0*u .- K0*v05 .+ uforce0
	LinearAlgebra.mul!(κ₁,S0,u)
	mul!(κ₁,K0,v05,-1.0,1.0)
	LinearAlgebra.axpy!(1.0,uforce0,κ₁)

	# rhs    .= S1*(u .+ 0.5*h*κ₁) .- K1*v05
	LinearAlgebra.mul!(rhs, S1,u)
	mul!(rhs,S1,κ₁,0.5*h,1.0)
	mul!(rhs,K1,v05,-1.0,1.0)
	LinearAlgebra.axpy!(1.0,uforce1,rhs)

	# u      .= u .+ (0.5*h).*κ₁
	LinearAlgebra.axpy!(0.5*h,κ₁,u)

	# κ₂     .= (In .- (0.5*h)*S1)\rhs
	neumann!(nNeumann,h,S1,rhs,κ₁,κ₂)

	# u      .= u .+ (0.5*h).*κ₂
	LinearAlgebra.axpy!(0.5*h,κ₂,u)

	# ℓ₂  .= K05*u .+ S05*v05 .+ vforce1
	LinearAlgebra.mul!(ℓ₂,K05,u)
	mul!(ℓ₂,S05,v05,1.0,1.0)
	LinearAlgebra.axpy!(1.0,vforce1,ℓ₂)

	v      .= v .+ 0.5*h.*(ℓ₁ .+ ℓ₂)

	t      = t + h
end

# FMG: This routine has been modified. This is for the adjoint solve which contains
# forcing. We note that h is negative in this case.
@inline function step!(t::Float64, nNeumann::Int64, μ::Array{Float64,N}, ν::Array{Float64,N}, X::Array{Float64,N}, h::Float64,
  					uforce0::Array{Float64,N}, vforce05::Array{Float64,N}, uforce1::Array{Float64,N},
  					K0::Array{Float64,N},S0::Array{Float64,N},K05::Array{Float64,N},S05::Array{Float64,N},
 					K1::Array{Float64,N},S1::Array{Float64,N},In::Array{Float64,N},
 					κ₁::Array{Float64,N}, κ₂::Array{Float64,N},ℓ₁::Array{Float64,N}, ℓ₂::Array{Float64,N},
  					rhs::Array{Float64,N}) where N

    # rhs .= S0*μ .- K05*ν .+ uforce0
	LinearAlgebra.mul!(rhs,S0,μ)
	mul!(rhs,K05,ν,-1.0,1.0)
	LinearAlgebra.axpy!(1.0, uforce0, rhs)

	# Neumann Series to invert linear system
	neumann!(nNeumann,h,S0,rhs,κ₁,κ₂)

	# μ  .= μ .+ 0.5*h.*κ₂
	LinearAlgebra.axpy!(0.5*h, κ₂, μ)

	# X   .= (μ .+ 0.5*h.*κ₂)
	copy!(X,μ)

	# ℓ₂  .= K0*X .+ S05*ν .+ vforce05
	LinearAlgebra.mul!(ℓ₂,K0,X)
	mul!(ℓ₂,S05,ν,1.0,1.0)
	LinearAlgebra.axpy!(1.0, vforce05, ℓ₂)

	# rhs .= S05*(ν .+ 0.5*h*ℓ₂) .+ K1*X .+ vforce05
	LinearAlgebra.mul!(rhs,S05,ν)
	mul!(rhs,S05,ℓ₂,0.5*h,1.0)
	mul!(rhs,K1,X,1.0,1.0)
	LinearAlgebra.axpy!(1.0, vforce05, rhs)

	# Neumann Series to invert linear system
	neumann!(nNeumann,h,S05,rhs,κ₂,ℓ₁)

	ν  .= ν .+ (0.5*h).*(ℓ₂ .+ ℓ₁)

	# κ₁ .= -K05*ν .+ S1*X .+ uforce1
	LinearAlgebra.mul!(κ₁,S1,X)
	mul!(κ₁,K05,ν,-1.0,1.0)
	LinearAlgebra.axpy!(1.0, uforce1, κ₁)

	# μ  .= μ .+ 0.5*h.*κ₁
	LinearAlgebra.axpy!(0.5*h, κ₁, μ)

	t  = t + h
end

# sparse version of step function above
@inline function step!(t::Float64, nNeumann::Int64, μ::Array{Float64,N}, ν::Array{Float64,N}, X::Array{Float64,N}, h::Float64,
  					uforce0::Array{Float64,N}, vforce05::Array{Float64,N}, uforce1::Array{Float64,N},
  					K0::SparseMatrixCSC{Float64,Int64},S0::SparseMatrixCSC{Float64,Int64},K05::SparseMatrixCSC{Float64,Int64},S05::SparseMatrixCSC{Float64,Int64},
 					K1::SparseMatrixCSC{Float64,Int64},S1::SparseMatrixCSC{Float64,Int64},In::SparseMatrixCSC{Float64,Int64},
					κ₁::Array{Float64,N}, κ₂::Array{Float64,N},ℓ₁::Array{Float64,N}, ℓ₂::Array{Float64,N},
  					rhs::Array{Float64,N}) where N

    # rhs .= S0*μ .- K05*ν .+ uforce0
	LinearAlgebra.mul!(rhs,S0,μ)
	mul!(rhs,K05,ν,-1.0,1.0)
	LinearAlgebra.axpy!(1.0, uforce0, rhs)

	# Neumann Series to invert linear system
	neumann!(nNeumann,h,S0,rhs,κ₁,κ₂)

	# μ  .= μ .+ 0.5*h.*κ₂
	LinearAlgebra.axpy!(0.5*h, κ₂, μ)

	# X   .= (μ .+ 0.5*h.*κ₂)
	copy!(X,μ)

	# ℓ₂  .= K0*X .+ S05*ν .+ vforce05
	LinearAlgebra.mul!(ℓ₂,K0,X)
	mul!(ℓ₂,S05,ν,1.0,1.0)
	LinearAlgebra.axpy!(1.0, vforce05, ℓ₂)

	# rhs .= S05*(ν .+ 0.5*h*ℓ₂) .+ K1*X .+ vforce05
	LinearAlgebra.mul!(rhs,S05,ν)
	mul!(rhs,S05,ℓ₂,0.5*h,1.0)
	mul!(rhs,K1,X,1.0,1.0)
	LinearAlgebra.axpy!(1.0, vforce05, rhs)

	# Neumann Series to invert linear system
	neumann!(nNeumann,h,S05,rhs,κ₂,ℓ₁)

	ν  .= ν .+ (0.5*h).*(ℓ₂ .+ ℓ₁)

	# κ₁ .= -K05*ν .+ S1*X .+ uforce1
	LinearAlgebra.mul!(κ₁,S1,X)
	mul!(κ₁,K05,ν,-1.0,1.0)
	LinearAlgebra.axpy!(1.0, uforce1, κ₁)

	# μ  .= μ .+ 0.5*h.*κ₁
	LinearAlgebra.axpy!(0.5*h, κ₁, μ)

	t  = t + h

end

# FMG: This routine has been modified. This is for the forward evolution with no forcing.
@inline function step!(t::Float64, nNeumann::Int64, u::Array{Float64,N}, v::Array{Float64,N}, v05::Array{Float64,N}, h::Float64,
  					  K0::Array{Float64,N},S0::Array{Float64,N},K05::Array{Float64,N},S05::Array{Float64,N},
  					  K1::Array{Float64,N},S1::Array{Float64,N},In::Array{Float64,N},
  					  κ₁::Array{Float64,N}, κ₂::Array{Float64,N},ℓ₁::Array{Float64,N}, ℓ₂::Array{Float64,N},
  					  rhs::Array{Float64,N}) where N

 	# rhs    .= (K05*u .+  S05*v)
 	LinearAlgebra.mul!(rhs,K05,u)
 	mul!(rhs,S05,v,1.0,1.0)

	# ℓ₁     .= (In .-  0.5*h.*S05)\rhs
	neumann!(nNeumann,h,S05,rhs,v05,ℓ₁)

	# v05    .= v .+ 0.5*h*ℓ₁
	copy!(v05,v)
	LinearAlgebra.axpy!(0.5*h,ℓ₁,v05)

	# κ₁     .= S0*u .- K0*v05
	LinearAlgebra.mul!(κ₁,S0,u)
	mul!(κ₁,K0,v05,-1.0,1.0)

	# rhs    .= S1*(u .+ 0.5*h*κ₁) .- K1*v05
	LinearAlgebra.mul!(rhs, S1,u)
	mul!(rhs,S1,κ₁,0.5*h,1.0)
	mul!(rhs,K1,v05,-1.0,1.0)

	# u      .= u .+ (0.5*h).*κ₁
	LinearAlgebra.axpy!(0.5*h,κ₁,u)

	# κ₂     .= (In .- (0.5*h)*S1)\rhs
	neumann!(nNeumann,h,S1,rhs,κ₁,κ₂)

	# u      .= u .+ (0.5*h).*κ₂
	LinearAlgebra.axpy!(0.5*h,κ₂,u)

	# ℓ₂     .= K05*u .+ S05*v05
	LinearAlgebra.mul!(ℓ₂,K05,u)
	mul!(ℓ₂,S05,v05,1.0,1.0)

	v      .= v .+ 0.5*h.*(ℓ₁ .+ ℓ₂)

	t      = t + h
	return t
end

# sparse version of step function above
@inline function step!(t::Float64, nNeumann::Int64, u::Array{Float64,N}, v::Array{Float64,N}, v05::Array{Float64,N}, h::Float64,
  					  K0::SparseMatrixCSC{Float64,Int64},S0::SparseMatrixCSC{Float64,Int64},K05::SparseMatrixCSC{Float64,Int64},S05::SparseMatrixCSC{Float64,Int64},
  					  K1::SparseMatrixCSC{Float64,Int64},S1::SparseMatrixCSC{Float64,Int64},In::SparseMatrixCSC{Float64,Int64},
  					  κ₁::Array{Float64,N}, κ₂::Array{Float64,N},ℓ₁::Array{Float64,N}, ℓ₂::Array{Float64,N},
  					  rhs::Array{Float64,N}) where N

 	# rhs    .= (K05*u .+  S05*v)
 	LinearAlgebra.mul!(rhs,K05,u)
 	mul!(rhs,S05,v,1.0,1.0)

	# ℓ₁     .= (In .-  0.5*h.*S05)\rhs
	neumann!(nNeumann,h,S05,rhs,v05,ℓ₁)

	# v05    .= v .+ 0.5*h*ℓ₁
	copy!(v05,v)
	LinearAlgebra.axpy!(0.5*h,ℓ₁,v05)

	# κ₁     .= S0*u .- K0*v05
	LinearAlgebra.mul!(κ₁,S0,u)
	mul!(κ₁,K0,v05,-1.0,1.0)

	# rhs    .= S1*(u .+ 0.5*h*κ₁) .- K1*v05
	LinearAlgebra.mul!(rhs, S1,u)
	mul!(rhs,S1,κ₁,0.5*h,1.0)
	mul!(rhs,K1,v05,-1.0,1.0)

	# u      .= u .+ (0.5*h).*κ₁
	LinearAlgebra.axpy!(0.5*h,κ₁,u)

	# κ₂     .= (In .- (0.5*h)*S1)\rhs
	neumann!(nNeumann,h,S1,rhs,κ₁,κ₂)

	# u      .= u .+ (0.5*h).*κ₂
	LinearAlgebra.axpy!(0.5*h,κ₂,u)

	# ℓ₂     .= K05*u .+ S05*v05
	LinearAlgebra.mul!(ℓ₂,K05,u)
	mul!(ℓ₂,S05,v05,1.0,1.0)

	v      .= v .+ 0.5*h.*(ℓ₁ .+ ℓ₂)

	t      = t + h
	return t
end

function stepseparable(stepper::svparams,u,v,t,h)
  if S(stepper.t) != 0
  	error("S has to be zero for separble time-stepping to ework")
  end

	S = stepper.S
	K = stepper.K
	In = stepper.In
	h = stepper.h
	uforce = stepper.uforce
	vforce = stepper.vforce

	l1 = (K(t)*u .+ vforce(t))
	kappa1 = - K(t + 0.5*h)*(v .+ 0.5*h.*l1) .+ uforce(t + 0.5*h)
	kappa2 = kappa1

	u = u .+ (h/2).*(kappa1 .+ kappa2)
	l2 = K(t +h)*u  .+ vforce(t + h)

	v = v .+ 0.5*h.*(l1 .+ l2)
	t = t + h

	return t, u, v
end

function getgamma(order,stages = [])
 if order == 2	# 2nd order basic verlet
	stages = 1
    gamma = [1]
  elseif order == 4 # 4th order Composition of Stromer-Verlet methods
    stages=3
    gamma = zeros(stages)
    gamma[1] = gamma[3] = 1/(2 - 2^(1/3))
    gamma[2] = -2^(1/3)*gamma[1]
  elseif order == 6 # Yoshida (1990) 6th order, 7 stage method
    if stages==7
      gamma = zeros(stages)
      gamma[2] = gamma[6] = 0.23557321335935813368479318
      gamma[1] = gamma[7] = 0.78451361047755726381949763
      gamma[3] = gamma[5] = -1.17767998417887100694641568
      gamma[4] = 1.31518632068391121888424973
    else # Kahan + Li 6th order, 9 stage method
      stages=9;
      gamma = zeros(stages)
      gamma[1]= gamma[9]= 0.39216144400731413927925056
      gamma[2]= gamma[8]= 0.33259913678935943859974864
      gamma[3]= gamma[7]= -0.70624617255763935980996482
      gamma[4]= gamma[6]= 0.08221359629355080023149045
      gamma[5]= 0.79854399093482996339895035
    end
   end

   return gamma, stages
end

# Routine to advance the solution by one time step with a second order Magnus integrator. Note that
# we use this in a very special case where we have an evolution matrix that is
# real, orthogonal, and skew-symmetric which gives a single complex conjugate
# pair of eigenvalues of ±i. This allows an incredibly simplified approach
# to compute the matrix exponential.
@inline function magnus(stepper::svparams, t::Float64, u::Array{Float64,N},
	v::Array{Float64,N}, h::Float64, uforce::Function, vforce::Function) where N
	uforce = uforce(t+0.5*h)
	vforce = vforce(t+0.5*h)

	S = stepper.S
	K = stepper.K
	In = stepper.In

	# Only evaluate the matrix at half-steps
	K = K(t +0.5*h)
	S = S(t + 0.5*h)

	# Single application of matrix
	# u1 = S*u - K*v
	# v1 = K*u + S*v
	# u += sin(h)*u1 + (1.0-cos(h))*(S*u1-K*v1)
	# v += sin(h)*v1 + (1.0-cos(h))*(K*u1+S*v1)

	# # Add in forcing
	# u += sin(h)*uforce + (1.0-cos(h))*(S*uforce - K*vforce)
	# v += sin(h)*vforce + (1.0-cos(h))*(K*uforce + S*vforce)
	# t = t + h

	# Brute force
	A = [S -K;K S]
	E = exp(h*A)
	Ndeg = length(u)
	tmp = E*[u;v] + inv(A)*(E - [In zeros(Ndeg,Ndeg);zeros(Ndeg,Ndeg) In])*[uforce;vforce]

	u = tmp[1:length(u)]
	v = tmp[length(u)+1:end]
	t = t + h
	return t, u, v
end

@inline function neumann!(nterms::Int64, h::Float64, S::SparseMatrixCSC{Float64,Int64},
						 B::Array{Float64,N}, T::Array{Float64,N}, X::Array{Float64,N}) where N
	copy!(X,B)
	copy!(T,B)
	coeff = 1.0
	for j = 1:nterms
		LinearAlgebra.mul!(T,S,B)
		coeff *= (0.5*h)
		LinearAlgebra.axpy!(coeff, T, X)
		copy!(B,T)
	end
end

@inline function neumann!(nterms::Int64, h::Float64, S::Array{Float64,N}, B::Array{Float64,N},
						T::Array{Float64,N}, X::Array{Float64,N}) where N
	copy!(X,B)
	copy!(T,B)
	coeff = 1.0
	for j = 1:nterms
		LinearAlgebra.mul!(T,S,B)
		coeff *= (0.5*h)
		LinearAlgebra.axpy!(coeff, T, X)
		copy!(B,T)
	end
end


text:
# Logistic regression model for the Prima Indian data

# load packages
using Pkg
using PyPlot
using LinearAlgebra
using Random
using Printf
using GLM
using DataFrames

# include files
include(pwd()*"/utilities.jl")

Random.seed!(1234) # set random numbers

# load data
x_train,y_train,x_val,y_val,x_test,y_test = load_data([0,1],0.2)

# create data frame
training_data = zeros(size(x_train,2), size(x_train,1)+1)
training_data[:,1:end-1] = x_train'
training_data[:,end] = y_train
train_df = DataFrame(training_data)
test_df = DataFrame(x_test')

model = @formula(x9 ~ x1+x2+x3+x4+x5+x6+x7+x8)

@time logisticregresson = glm(model, train_df, Binomial(), LogitLink())

print(logisticregresson)

# compute predictions
y_pred_train = class_logistic(predict(logisticregresson))
y_pred_test = class_logistic(predict(logisticregresson,test_df))

# check classification results
class_res_training = classification_results(y_pred_train, y_train, [0,1])
class_res_test = classification_results(y_pred_test, y_test, [0,1])

max_stars_repo_name: SamuelWiqvist/dl-compute-project

max_issues_repo_name: SamuelWiqvist/dl-compute-project

max_forks_repo_name: SamuelWiqvist/dl-compute-project


text:
struct PyConvertRule
    type :: Type
    func :: Function
    priority :: Int
end

const PYCONVERT_RULES = Dict{String, Vector{PyConvertRule}}()
const PYCONVERT_EXTRATYPES = Py[]

function pyconvert_add_rule(pytypename::String, type::Type, func::Function, priority::Int=0)
    @nospecialize type func
    push!(get!(Vector{PyConvertRule}, PYCONVERT_RULES, pytypename), PyConvertRule(type, func, priority))
    return
end

if false
    # this scheme returns either the result or Unconverted()
    struct Unconverted end
    pyconvert_return(x) = x
    pyconvert_unconverted() = Unconverted()
    pyconvert_returntype(::Type{T}) where {T} = Union{T,Unconverted}
    pyconvert_isunconverted(r) = r === Unconverted()
    pyconvert_result(::Type{T}, r) where {T} = r::T
elseif true
    # this scheme stores the result in PYCONVERT_RESULT
    const PYCONVERT_RESULT = Ref{Any}(nothing)
    pyconvert_return(x) = (PYCONVERT_RESULT[] = x; true)
    pyconvert_unconverted() = false
    pyconvert_returntype(::Type{T}) where {T} = Bool
    pyconvert_isunconverted(r::Bool) = !r
    pyconvert_result(::Type{T}, r::Bool) where {T} = (ans = PYCONVERT_RESULT[]::T; PYCONVERT_RESULT[] = nothing; ans)
else
    # same as the previous scheme, but with special handling for bits types
    const PYCONVERT_RESULT = Ref{Any}(nothing)
    const PYCONVERT_RESULT_ISBITS = Ref{Bool}(false)
    const PYCONVERT_RESULT_TYPE = Ref{Type}(Union{})
    const PYCONVERT_RESULT_BITSLEN = 1024
    const PYCONVERT_RESULT_BITS = fill(0x00, PYCONVERT_RESULT_BITSLEN)
    function pyconvert_return(x::T) where {T}
        if isbitstype(T) && sizeof(T) ≤ PYCONVERT_RESULT_BITSLEN
            unsafe_store!(Ptr{T}(pointer(PYCONVERT_RESULT_BITS)), x)
            PYCONVERT_RESULT_ISBITS[] = true
            PYCONVERT_RESULT_TYPE[] = T
        else
            PYCONVERT_RESULT[] = x
            PYCONVERT_RESULT_ISBITS[] = false
        end
        return true
    end
    pyconvert_unconverted() = false
    pyconvert_returntype(::Type{T}) where {T} = Bool
    pyconvert_isunconverted(r::Bool) = !r
    function pyconvert_result(::Type{T}, r::Bool) where {T}
        if isbitstype(T)
            if sizeof(T) ≤ PYCONVERT_RESULT_BITSLEN
                @assert PYCONVERT_RESULT_ISBITS[]
                @assert PYCONVERT_RESULT_TYPE[] == T
                return unsafe_load(Ptr{T}(pointer(PYCONVERT_RESULT_BITS)))::T
            end
        elseif PYCONVERT_RESULT_ISBITS[]
            t = PYCONVERT_RESULT_TYPE[]
            @assert isbitstype(t)
            @assert sizeof(t) ≤ PYCONVERT_RESULT_BITSLEN
            @assert t <: T
            @assert isconcretetype(t)
            return unsafe_load(Ptr{t}(pointer(PYCONVERT_RESULT_BITS)))::T
        end
        # general case
        ans = PYCONVERT_RESULT[]::T
        PYCONVERT_RESULT[] = nothing
        return ans::T
    end
end

pyconvert_result(r) = pyconvert_result(Any, r)

pyconvert_tryconvert(::Type{T}, x::T) where {T} = pyconvert_return(x)
pyconvert_tryconvert(::Type{T}, x) where {T} =
    try
        pyconvert_return(convert(T, x)::T)
    catch
        pyconvert_unconverted()
    end

function pyconvert_typename(t::Py)
    m = pygetattr(t, "__module__", "<unknown>")
    n = pygetattr(t, "__name__", "<name>")
    return "$m/$n"
end

function pyconvert_get_rules(type::Type, pytype::Py)
    @nospecialize type

    pyisin(x, ys) = any(pyis(x, y) for y in ys)

    # get the MROs of all base types we are considering
    omro = collect(pytype.__mro__)
    basetypes = Py[pytype]
    basemros = Vector{Py}[omro]
    for xtype in PYCONVERT_EXTRATYPES
        # find the topmost supertype of
        xbase = PyNULL
        for base in omro
            if pyissubclass(base, xtype)
                xbase = base
            end
        end
        if !ispynull(xbase)
            push!(basetypes, xtype)
            xmro = collect(xtype.__mro__)
            pyisin(xbase, xmro) || pushfirst!(xmro, xbase)
            push!(basemros, xmro)
        end
    end
    for xbase in basetypes[2:end]
        push!(basemros, [xbase])
    end

    # merge the MROs
    # this is a port of the merge() function at the bottom of:
    # https://www.python.org/download/releases/2.3/mro/
    mro = Py[]
    while !isempty(basemros)
        # find the first head not contained in any tail
        ok = false
        b = PyNULL
        for bmro in basemros
            b = bmro[1]
            if all(bmro -> !pyisin(b, bmro[2:end]), basemros)
                ok = true
                break
            end
        end
        ok || error("Fatal inheritence error: could not merge MROs (mro=$mro, basemros=$basemros)")
        # add it to the list
        push!(mro, b)
        # remove it from consideration
        for bmro in basemros
            filter!(t -> !pyis(t, b), bmro)
        end
        # remove empty lists
        filter!(x -> !isempty(x), basemros)
    end
    # check the original MRO is preserved
    omro_ = filter(t -> pyisin(t, omro), mro)
    @assert length(omro) == length(omro_)
    @assert all(pyis(x,y) for (x,y) in zip(omro, omro_))

    # get the names of the types in the MRO of pytype
    xmro = [String[pyconvert_typename(t)] for t in mro]

    # add special names corresponding to certain interfaces
    # these get inserted just above the topmost type satisfying the interface
    for (t, x) in reverse(collect(zip(mro, xmro)))
        if pyhasattr(t, "__array_struct__")
            push!(x, "<arraystruct>")
            break
        end
    end
    for (t, x) in reverse(collect(zip(mro, xmro)))
        if pyhasattr(t, "__array_interface__")
            push!(x, "<arrayinterface>")
            break
        end
    end
    for (t, x) in reverse(collect(zip(mro, xmro)))
        if pyhasattr(t, "__array__")
            push!(x, "<array>")
            break
        end
    end
    for (t, x) in reverse(collect(zip(mro, xmro)))
        if C.PyType_CheckBuffer(getptr(t))
            push!(x, "<buffer>")
            break
        end
    end

    # flatten to get the MRO as a list of strings
    mro = String[x for xs in xmro for x in xs]

    # get corresponding rules
    rules = PyConvertRule[rule for tname in mro for rule in get!(Vector{PyConvertRule}, PYCONVERT_RULES, tname)]

    # order the rules by priority, then by original order
    order = sort(axes(rules, 1), by = i -> (-rules[i].priority, i))
    rules = rules[order]

    # TODO: everything up to here does not depend on the julia type and could be cached

    # intersect rules with type
    rules = PyConvertRule[PyConvertRule(typeintersect(rule.type, type), rule.func, rule.priority) for rule in rules]

    # explode out unions
    rules = [PyConvertRule(type, rule.func, rule.priority) for rule in rules for type in Utils.explode_union(rule.type)]

    # filter out empty rules
    rules = [rule for rule in rules if rule.type != Union{}]

    # filter out repeated rules
    rules = [rule for (i, rule) in enumerate(rules) if !any((rule.func === rules[j].func) && ((rule.type) <: (rules[j].type)) for j in 1:(i-1))]

    @debug "pyconvert" type pytype mro=join(mro, " ") rules
    return Function[pyconvert_fix(rule.type, rule.func) for rule in rules]
end

pyconvert_fix(::Type{T}, func) where {T} = x -> func(T, x)

const PYCONVERT_RULES_CACHE = Dict{C.PyPtr, Dict{Type, Vector{Function}}}()

function pyconvert_rule_fast(::Type{T}, x::Py) where {T}
    if T isa Union
        a = pyconvert_rule_fast(T.a, x) :: pyconvert_returntype(T.a)
        pyconvert_isunconverted(a) || return a
        b = pyconvert_rule_fast(T.b, x) :: pyconvert_returntype(T.b)
        pyconvert_isunconverted(b) || return b
    end
    pyconvert_unconverted()
end

pytryconvert(::Type{T}, x) where {T} = @autopy x begin
    # We can optimize the conversion for some types by overloading pytryconvert_fast.
    # It MUST give the same results as via the slower route using rules.
    ans = pyconvert_rule_fast(T, getpy(x_)) :: pyconvert_returntype(T)
    pyconvert_isunconverted(ans) || return ans
    # get rules from the cache
    # TODO: we should hold weak references and clear the cache if types get deleted
    tptr = C.Py_Type(getptr(x_))
    trules = get!(Dict{Type, Vector{PyConvertRule}}, PYCONVERT_RULES_CACHE, tptr)
    if !haskey(trules, T)
        t = pynew(incref(tptr))
        trules[T] = pyconvert_get_rules(T, t)
        pydel!(t)
    end
    rules = trules[T]
    # apply the rules
    for rule in rules
        ans = rule(getpy(x_)) :: pyconvert_returntype(T)
        pyconvert_isunconverted(ans) || return ans
    end
    return pyconvert_unconverted()
end

"""
    @pyconvert(T, x, [onfail])

Convert the Python object `x` to a `T`.

On failure, evaluates to `onfail`, which defaults to `return pyconvert_unconverted()` (mainly useful for writing conversion rules).
"""
macro pyconvert(T, x, onfail=:(return $pyconvert_unconverted()))
    quote
        ans = pytryconvert($(esc(T)), $(esc(x)))
        if pyconvert_isunconverted(ans)
            $(esc(onfail))
        else
            pyconvert_result($(esc(T)), ans)
        end
    end
end
export @pyconvert

"""
    pyconvert(T, x, [d])

Convert the Python object `x` to a `T`.

If `d` is specified, it is returned on failure instead of throwing an error.
"""
pyconvert(::Type{T}, x) where {T} = @autopy x @pyconvert T x_ error("cannot convert this Python '$(pytype(x_).__name__)' to a Julia '$T'")
pyconvert(::Type{T}, x, d) where {T} = @autopy x @pyconvert T x_ d
export pyconvert

"""
    pyconvertarg(T, x, name)

Convert the Python object `x` to a `T`.

On failure, throws a Python `TypeError` saying that the argument `name` could not be converted.
"""
pyconvertarg(::Type{T}, x, name) where {T} = @autopy x @pyconvert T x_ begin
    errset(pybuiltins.TypeError, "Cannot convert argument '$name' to a Julia '$T', got a '$(pytype(x_).__name__)'")
    pythrow()
end

macro pyconvert_and_del(T, x, onfail=:(return $pyconvert_unconverted()))
    quote
        ans = pytryconvert($(esc(T)), $(esc(x)))
        pydel!($(esc(x)))
        if pyconvert_isunconverted(ans)
            $(esc(onfail))
        else
            pyconvert_result($(esc(T)), ans)
        end
    end
end

pyconvert_and_del(::Type{T}, x) where {T} = begin
    ans = pyconvert(T, x)
    pydel!(x)
    ans
end

function init_pyconvert()
    push!(PYCONVERT_EXTRATYPES, pyimport("io"=>"IOBase"))
    push!(PYCONVERT_EXTRATYPES, pyimport("numbers"=>("Number", "Complex", "Real", "Rational", "Integral"))...)
    push!(PYCONVERT_EXTRATYPES, pyimport("collections.abc" => ("Iterable", "Sequence", "Set", "Mapping"))...)

    # priority 300: wrapped julia values
    pyconvert_add_rule("juliacall/As", Any, pyconvert_rule_jlas, 300)
    pyconvert_add_rule("juliacall/ValueBase", Any, pyconvert_rule_jlvalue, 300)
    # priority 200: arrays
    pyconvert_add_rule("<arraystruct>", PyArray, pyconvert_rule_array_nocopy, 200)
    pyconvert_add_rule("<arrayinterface>", PyArray, pyconvert_rule_array_nocopy, 200)
    pyconvert_add_rule("<array>", PyArray, pyconvert_rule_array_nocopy, 200)
    pyconvert_add_rule("<buffer>", PyArray, pyconvert_rule_array_nocopy, 200)
    # priority 100: canonical
    pyconvert_add_rule("builtins/NoneType", Nothing, pyconvert_rule_none, 100)
    pyconvert_add_rule("builtins/bool", Bool, pyconvert_rule_bool, 100)
    pyconvert_add_rule("builtins/float", Float64, pyconvert_rule_float, 100)
    pyconvert_add_rule("builtins/complex", Complex{Float64}, pyconvert_rule_complex, 100)
    pyconvert_add_rule("numbers/Integral", Integer, pyconvert_rule_int, 100)
    pyconvert_add_rule("builtins/str", String, pyconvert_rule_str, 100)
    pyconvert_add_rule("builtins/bytes", Base.CodeUnits{UInt8,String}, pyconvert_rule_bytes, 100)
    pyconvert_add_rule("builtins/range", StepRange{<:Integer,<:Integer}, pyconvert_rule_range, 100)
    pyconvert_add_rule("numbers/Rational", Rational{<:Integer}, pyconvert_rule_fraction, 100)
    pyconvert_add_rule("collections.abc/Iterable", PyIterable, pyconvert_rule_iterable, 100)
    pyconvert_add_rule("collections.abc/Sequence", PyList, pyconvert_rule_sequence, 100)
    pyconvert_add_rule("collections.abc/Set", PySet, pyconvert_rule_set, 100)
    pyconvert_add_rule("collections.abc/Mapping", PyDict, pyconvert_rule_mapping, 100)
    pyconvert_add_rule("io/IOBase", PyIO, pyconvert_rule_io, 100)
    pyconvert_add_rule("pandas.core.frame/DataFrame", PyPandasDataFrame, pyconvert_rule_pandasdataframe, 100)
    pyconvert_add_rule("builtins/tuple", Tuple, pyconvert_rule_iterable, 100)
    pyconvert_add_rule("datetime/datetime", DateTime, pyconvert_rule_datetime, 100)
    pyconvert_add_rule("datetime/date", Date, pyconvert_rule_date, 100)
    pyconvert_add_rule("datetime/time", Time, pyconvert_rule_time, 100)
    # priority 0: reasonable
    pyconvert_add_rule("builtins/NoneType", Missing, pyconvert_rule_none)
    pyconvert_add_rule("builtins/bool", Number, pyconvert_rule_bool)
    pyconvert_add_rule("numbers/Real", Number, pyconvert_rule_float)
    pyconvert_add_rule("builtins/float", Nothing, pyconvert_rule_float)
    pyconvert_add_rule("builtins/float", Missing, pyconvert_rule_float)
    pyconvert_add_rule("numbers/Complex", Number, pyconvert_rule_complex)
    pyconvert_add_rule("numbers/Integral", Number, pyconvert_rule_int)
    pyconvert_add_rule("builtins/str", Symbol, pyconvert_rule_str)
    pyconvert_add_rule("builtins/str", Char, pyconvert_rule_str)
    pyconvert_add_rule("builtins/bytes", Vector{UInt8}, pyconvert_rule_bytes)
    pyconvert_add_rule("builtins/range", UnitRange{<:Integer}, pyconvert_rule_range)
    pyconvert_add_rule("numbers/Rational", Number, pyconvert_rule_fraction)
    pyconvert_add_rule("collections.abc/Iterable", Vector, pyconvert_rule_iterable)
    pyconvert_add_rule("collections.abc/Iterable", Tuple, pyconvert_rule_iterable)
    pyconvert_add_rule("collections.abc/Iterable", Pair, pyconvert_rule_iterable)
    pyconvert_add_rule("collections.abc/Iterable", Set, pyconvert_rule_iterable)
    pyconvert_add_rule("collections.abc/Sequence", Vector, pyconvert_rule_iterable)
    pyconvert_add_rule("collections.abc/Sequence", Tuple, pyconvert_rule_iterable)
    pyconvert_add_rule("collections.abc/Set", Set, pyconvert_rule_iterable)
    pyconvert_add_rule("collections.abc/Mapping", Dict, pyconvert_rule_mapping)
    pyconvert_add_rule("<arraystruct>", Array, pyconvert_rule_array)
    pyconvert_add_rule("<arrayinterface>", Array, pyconvert_rule_array)
    pyconvert_add_rule("<array>", Array, pyconvert_rule_array)
    pyconvert_add_rule("<buffer>", Array, pyconvert_rule_array)
    pyconvert_add_rule("<arraystruct>", AbstractArray, pyconvert_rule_array)
    pyconvert_add_rule("<arrayinterface>", AbstractArray, pyconvert_rule_array)
    pyconvert_add_rule("<array>", AbstractArray, pyconvert_rule_array)
    pyconvert_add_rule("<buffer>", AbstractArray, pyconvert_rule_array)
    # priority -100: fallbacks
    pyconvert_add_rule("builtins/object", Py, pyconvert_rule_object, -100)
    # priority -200: explicit
end


text:
using Documenter, StencilInterpolation

push!(LOAD_PATH, "../src/")
makedocs(; sitename="Stencil Interpolations")


text:
using BenchmarkTools
using Parameters
using Roots
using Interpolations
using LinearAlgebra

"""
    compute_steady_state!(gl::ModelGL)

Computes steady-state policy functions with given set of parameters.
"""
function compute_steady_state!(gl::ModelGL)
    initilize!(gl)
    EGM!(gl)
    compute_distribution!(gl)
    aggregate!(gl)
end

"""
    initilize!(gl::ModelGL)

Adds borrowing constraint as gridpoint
"""
function initilize!(gl::ModelGL)

    # Ensure that constraint is on the grid
    gl.b_grid = gl.bmin .+ ((1:gl.nb)/gl.nb).^2 .* (gl.bmax - gl.bmin); # denser for low values
    i = maximum(  [  searchsortedlast(gl.b_grid, -gl.ϕ) ,  1] )
    gl.b_grid[i] =  -gl.ϕ
end

"""
    EGM!(gl::ModelGL)

Calculates optimal policy functions with the EGM algorithm.
"""
function EGM!(gl::ModelGL)

    @unpack r,ϕ  = gl
    @unpack β, ν,θ, η,ψ,γ, B,S,nb,gameta,fin,sep = gl               # model parameters
    @unpack cmin,tol_pol = gl                                       # numerical paras
    @unpack b_grid,Ic,db = gl                                       # grids
    @unpack pr,Pr = gl                                              # processes
    @unpack cl,n_pol,y_pol,b_pol,c_pol,tol_dist,mpcs = gl           # solutions

    # Update budget constraint
    τ   = (pr[1]*ν + r/(1+r)*B) / (1 - pr[1]);      # labor tax
    z   = [ν, -τ.*ones(S-1)...];                    # full transfer scheme (tau tilde in paper)

    # Find consumption at the lower bound of the state space
    fac = (ψ ./ θ) .^ (1/η);
    for s = 1:S

        cl[s] =   find_zero( x -> find_cl(x,s, -ϕ, -ϕ, r, θ, z, fac, gameta)
                                , (cmin, 100), Bisection()  )
    end

    # A) Solve for consumption policy
    #----------------------------------
    # Initial guess for policy function
    c_pol[:,:]  = r .* ones(S,nb) .* b_grid'
    c_pol[c_pol .<= cmin] .= cmin
    c_poli = copy(c_pol)

    # policy Convergence
    dif = 1;
    while dif > tol_pol
      # expected marginal utility tomorrow
      ui = Pr *  c_pol.^(-γ) ;
      ui = ui[:, b_grid .>= -ϕ];

      for s=1:S
        # unconstrained
        c = ((1+r) * β * ui[s,:]) .^ (-1/γ);                   # Euler
        n =   (1 .- fac[s]*c.^gameta )                         # labor supply
        n = [ n[i] < 0. ? 0. : n[i] for i = 1:length(n)]
        b = b_grid[b_grid .>= -ϕ] / (1+r) .+ c .- θ[s] .*n .- z[s]; # budget

        # constrained
        if b[1] > -ϕ
            c_c = range(cl[s], stop=c[1], length=Ic);
            n_c =   (1 .- fac[s]*c_c.^gameta )                         # labor supply
            n_c = [ n_c[i] < 0. ? 0. : n_c[i] for i = 1:length(n_c)]
            b_c = -ϕ/(1+r) .+ c_c .- θ[s] .*n_c .- z[s]; # budget
            b   = [b_c[1:Ic-1]..., b...];
            c   = [c_c[1:Ic-1]..., c...];
        end
        itp = extrapolate(    interpolate((b,),c, Gridded(Linear())), Interpolations.Flat() )
        c_poli[s,:] = itp(  b_grid   )
        # interp1(b, c, b_grid, 'linear', 'extrap');
      end

      # check convergence
      c_poli[c_poli .<= cmin] .= cmin
      dif    = norm(c_poli - c_pol)/(1+norm(c_pol) );

      # update
      c_pol[:,:] = copy(c_poli);
    end

    # Save other policy functions and MPCs
    for s = 1:S
      npp   = 1 .- fac[s] .*c_pol[s,:] .^gameta
      npp[npp .<= 0.] .= 0.
      n_pol[s,:] .= npp
      y_pol[s,:] = θ[s] .* n_pol[s,:];
      bpp        = (1+r) * (b_grid .+ y_pol[s,:] .- c_pol[s,:] .+ z[s])
      bpp[bpp .<= -ϕ] .= -ϕ
      b_pol[s,:] .= bpp
      itp = extrapolate(    interpolate((b_grid,),c_pol[s,:], Gridded(Linear())), Interpolations.Flat() )
      mpcs[s,:]  = (    itp( b_grid .+ db )  .- c_pol[s,:]  ) ./ db;
      # (interp1(b_grid, c_pol(s,:), b_grid + db, 'linear', 'extrap') - c_pol(s,:)) / db;
    end
end

"""
    compute_distribution!(gl::ModelGL)

Computes the joint distribution of productivity and bond holdings.
Requires that optimal policies have already been solved (EGM!)
"""
function compute_distribution!(gl::ModelGL)

    @unpack b_pol, b_grid,tol_dist,S,nb, JD,Pr,fin,pr,sep,ϕ = gl

    # B) Find invariant distribution
    #----------------------------------
    # Assign weights to adjacent grid points proportionally to distance
    idxes       = [ searchsortedlast(b_grid,b_pol[s, b]) for s = 1:size(b_pol)[1], b = 1:size(b_pol)[2] ]
    weights     = zeros(size(b_pol))
    for s = 1:size(b_pol)[1]
    for b = 1:size(b_pol)[2]
        idx             = idxes[s,b]
        idx_prime       = minimum(  [ idx+1,size(b_pol)[2] ]    )
        #db              = maximum(  [ b_grid[idx_prime] - b_grid[idx]  , eps() ] )
        db              =  b_grid[idx_prime] - b_grid[idx]
        weights[s,b]    = ( b_pol[s,b]  - b_grid[idx] ) / db;
    end
    end

    # Iterate asset transition matrix starting from uniform distribution
    dif = 1;
    iter =1
    JD = ones(S,nb) ./ (S*nb);
    while (dif > tol_dist) && iter < 1000

        iter += 1
        JDp = zeros( size(JD) )
        for s = 1:S
        for b = 1:nb
            for si = 1:S
                JDp[si, idxes[s, b]]     = (1 - weights[s, b]) * Pr[s, si] * JD[s, b]  + JDp[si, idxes[s, b]]
                JDp[si, idxes[s, b] + 1] =      weights[s, b]  * Pr[s, si] * JD[s, b]  + JDp[si, idxes[s, b] + 1]
            end
        end
        end

        # check convergence
        dif = norm( JDp - JD ) /( 1+norm(JD) )

        # make sure that distribution integrates to 1
        JD[:,:]  = JDp / sum( JDp );

    end

    gl.JD = copy(JD)
    gl.steadystate_solved = true
    return JD
end

"""
    aggregate!(gl::ModelGL;terminal::Bool = false)

Computes aggregate values and their distance from target.
Requires that the joint distribution has already been solved (compute_distribution!)
If ´terminal = true´ only the terminal debt target will be considered.
"""
function aggregate!(gl::ModelGL;terminal::Bool = false)

    # Compute Steady-state aggregates
    @unpack ν,β,ϕ,ψ,η  = gl
    @unpack NE,νY,B_4Y,D1_4Y,D2_4Y  = gl
    @unpack JD,b_grid,B,y_pol,n_pol  = gl
    @unpack tol_cali,tol_mkt  = gl

    # C) Check market clearing and calibration
    #-------------------------------------------------
    # Bond market clearing, i refers to current iteration
    Bi      = (sum(JD,dims=1) * b_grid)[1]
    res_mkt = abs(B - Bi)

    # Calibration statistics, i refers to current iteration
    Yi    = sum(JD  .* y_pol  )                                     # GDP
    NEi   = sum(  JD .* (n_pol.*(n_pol.>0) ) )  / sum(  JD .*(n_pol.>0)  )       # avg hours of employed
    B_4Yi =  Bi / Yi / 4;                                                # debt ratio
    Di    = - (   sum(JD,dims=1) *  ( b_grid .* (b_grid .< 0.) )  )[1]
    D_4Yi =  Di / Yi / 4;                                                # debt ratio
    νYi    =  ν / Yi;                                                    # UI benefit ratio
    if terminal == false
        res_cali = maximum(abs.([B_4Yi, D_4Yi, νYi, NEi] .- [B_4Y, D1_4Y, νY, NE]));
    else
        res_cali = maximum(abs.([  D_4Yi ] .- [  D2_4Y ]));
    end

    #store actual values
    gl.B_4Y_actual  = B_4Yi
    gl.D_4Y_actual  = D_4Yi
    gl.νY_actual    = νYi
    gl.NE_actual    = NEi
    gl.Y_actual     = Yi
    return Bi ,res_mkt ,Yi ,NEi,B_4Yi ,Di,D_4Yi,νYi,res_cali
end

"""
    calibrate!(gl::ModelGL)

Calibrate the model to the target values.
Does not require any prior function calls. If load_params = true, the
calibrated parameters are preloaded to speed up the calibration.
"""
function calibrate!(gl::ModelGL;load_params::Bool=false)

    if load_params == true
        load_parameters_iss!(gl)
    end

    @unpack NE,νY,η,B_4Y,D1_4Y  = gl
    @unpack tol_cali,tol_mkt  = gl
    ite = 0
    while true
        ite +=1
        initilize!(gl)
        EGM!(gl)
        compute_distribution!(gl)
        Bi ,res_mkt ,Yi ,NEi,B_4Yi ,Di,D_4Yi,νYi,res_cali = aggregate!(gl)
        println( string(ite)*"  "*string(res_mkt)*"   "*string(res_cali) )
        # Check convergence of both, update if necessary
        if (res_cali < tol_cali) && (res_mkt < tol_mkt)
            break
        else
            # discount factor
            btemp   = -log(1- gl.β);
            btemp   = btemp - .1*(Bi - B_4Y*4*Yi);
            gl.β       = 1 - exp(-btemp);

            # rest will be updated based on these
            ϕ_d    = gl.ϕ * D1_4Y / D_4Yi;
            ν_d    = νY * Yi;
            ψ_d    = gl.ψ * ((1-NE) / (1-NEi))^η;

            gl.ϕ   = gl.ϕ  + 0.1 * (ϕ_d   - gl.ϕ);
            gl.ν   = gl.ν  + 1   * (ν_d   - gl.ν);
            gl.ψ   = gl.ψ  + 1   * (ψ_d   - gl.ψ);

            # Update aggregates
            gl.B = gl.B + 0.1 * (Bi - gl.B);
        end
    end
    gl.calibrated = true
end

"""
    gl_tss = calibrate_terminal(gl_initial::ModelGL)

Calibrate the model to the terminal target value (only borrowing constraint)
Requires a calibrated initital steady-state object.
Returns a new terminal steady-state object. If load_params = true, the
calibrated parameters are preloaded to speed up the calibration.
"""
function calibrate_terminal(gl_initial::ModelGL;load_params::Bool=false)

    gl = deepcopy(gl_initial)   # terminal steadystate
    gl.ϕ = gl.ϕ₂                # update initial guess

    if load_params == true
        load_parameters_tss!(gl)
    end
    @unpack D2_4Y  = gl
    @unpack tol_cali,tol_mkt  = gl
    ite = 0
    while true
        ite +=1
        initilize!(gl)
        EGM!(gl)
        compute_distribution!(gl)
        Bi ,res_mkt ,Yi ,NEi,B_4Yi ,Di,D_4Yi,νYi,res_cali = aggregate!(gl,terminal = true)
        println( string(ite)*"  "*string(res_mkt)*"   "*string(res_cali) )
        # Check convergence of both, update if necessary
        if (res_cali < tol_cali) && (res_mkt < tol_mkt)
            break
        else
            gl.r = gl.r - 1/400*(Bi-gl.B)

            ϕ_d    = gl.ϕ * D2_4Y / D_4Yi;
            gl.ϕ   = gl.ϕ  + 0.1 * (ϕ_d   - gl.ϕ);

        end
    end
    gl.calibrated = true
    gl.steadystate_type = "terminal"
    return gl
end

"""
    transition!(gl::ModelGL,gl_tss::ModelGL,Tgl::TransGL)

Computes transition path.
"""
function transition!(gl::ModelGL,gl_tss::ModelGL,Tgl::TransGL)
    @unpack T, maxit_trans, tol_mkt_trans, weight, JD_t = Tgl
    @unpack Pr,B,nb = gl

    gl_t = deepcopy(gl)  # define new object for each time period

    # Define credit crunch parameters
    t_shock = 6                           # quarters to get to new constraint
    dϕ      = (gl.ϕ - gl_tss.ϕ)/t_shock   # step size
    Tgl.ϕ_t  = [max(gl_tss.ϕ,gl.ϕ .- dϕ*t) for t in 0:T] # full path

    # Initial guess of r_t
    Tgl.r_t = gl_tss.r .* ones(T)  # interest rate in terminal ss
    Tgl.JD_t[:,:,end] = gl_tss.JD # distribution in terminal ss
    for it = 1:maxit_trans

        # 1) Iterate HH problem backwards
        gl_t.c_pol = gl_tss.c_pol   # start from terminal ss
        gl_t.JD = gl_tss.JD # distribution in terminal ss

        for t = T:-1:1
            # current value of r and ϕ
            gl_t.r = Tgl.r_t[t]
            gl_t.ϕ = Tgl.ϕ_t[t+1]
            # Ensure that constraint is on the grid
            initilize!(gl_t)
            # Compute policy functions and save distribution weights
            EGM_trans!(gl_t,Tgl,t)
        end

        # 2) Iterate distribution forward

        JD = gl.JD;  # start from initial ss

        for t = 1:T
            # store distribution over assets and productivity states
            Tgl.JD_t[:,:,t] = JD
            gl_t.JD = JD
#            if t==1
#                println("is the distribution correct?", Tgl.JD_t[:,:,t] == gl.JD)
#                println("is the second version of the distribution correct?", gl_t.JD == JD)
#            end

            # Ensure that constraint is on the grid
            gl_t.ϕ = Tgl.ϕ_t[t+1]
            initilize!(gl_t)

            # calculate bond market clearing
            Tgl.Bdem_t[t] = (sum(Tgl.JD_t[:,:,t],dims=1) * gl_t.b_grid)[1]

            # calculate other aggregates
            Tgl.Y_t[t]    = sum(Tgl.JD_t[:,:,t]  .* Tgl.y_pol_t[:,:,t])                         # GDP
            Tgl.D_t[t]    = -(sum(Tgl.JD_t[:,:,t],dims=1) * (gl_t.b_grid .* (gl_t.b_grid .< 0.)))[1]  # Debt
            Tgl.D_4Y_t[t] =  Tgl.D_t[t] / Tgl.Y_t[t] / 4                               # Debt2GDP annual

            # iterate on distribution
            JDp = zeros(size(JD))
            for s = 1:S
            for b = 1:nb
                for si = 1:S
                    JDp[si, Tgl.ib_pol_t[s, b, t]]     = (1 - Tgl.wei_t[s, b, t]) * Pr[s, si] * Tgl.JD_t[s,b,t]  + JDp[si, Tgl.ib_pol_t[s, b, t]]
                    JDp[si, Tgl.ib_pol_t[s, b, t] + 1] =      Tgl.wei_t[s, b, t]  * Pr[s, si] * Tgl.JD_t[s,b,t]  + JDp[si, Tgl.ib_pol_t[s, b, t] + 1]
                end
            end
            end

            # make sure that distribution integrates to 1
            #JD[:,:]  = JDp / sum(JDp)
            JD = JDp / sum(JDp)

        end

        # 3) Check convergence

        # Market clearing
        BM_t = Tgl.Bdem_t .- B

        # Metric of deviation
        res_BM = sqrt(BM_t' * BM_t / T)

        # Report progress
        println( string(it)*"  "*string(res_BM))

        if res_BM < tol_mkt_trans
            break
        end

        # Update
        Tgl.r_t = Tgl.r_t - weight.*BM_t
    end

end

"""
    EGM_trans!(gl::ModelGL,Tgl::TransGL,t::Int)

Calculates optimal policy functions with one step of the EGM algorithm.
"""
function EGM_trans!(gl::ModelGL,Tgl::TransGL,t::Int)

        @unpack r,ϕ  = gl
        @unpack β, ν,θ, η,ψ,γ, B,S,gameta = gl            # model parameters
        @unpack cmin = gl                                 # numerical paras
        @unpack b_grid,Ic,bmax = gl                       # grids
        @unpack pr,Pr = gl                                # processes
        @unpack cl, c_pol = gl                            # solutions

        # Update budget constraint
        τ   = (pr[1]*ν + r/(1+r)*B) / (1 - pr[1]);      # labor tax
        z   = [ν, -τ.*ones(S-1)...];                    # full transfer scheme (tau tilde in paper)

        # Find consumption at the lower bound of the state space
        fac = (ψ ./ θ) .^ (1/η);
        for s = 1:S
                cl[s] =   find_zero( x -> find_cl(x,s, -Tgl.ϕ_t[t], -Tgl.ϕ_t[t+1], r, θ, z, fac, gameta)
                , (cmin, 100), Bisection()  )
        end

        # expected marginal utility tomorrow
        ui = Pr *  c_pol.^(-γ) ;
        ui = ui[:, b_grid .>= -ϕ];

        # previous consumption and savings policy
        for s=1:S
            # unconstrained
            c = ((1+r) * β * ui[s,:]) .^ (-1/γ);                   # Euler
            n =   (1 .- fac[s]*c.^gameta )                         # labor supply
            n = [ n[i] < 0. ? 0. : n[i] for i = 1:length(n)]
            b = b_grid[b_grid .>= -ϕ] / (1+r) .+ c .- θ[s] .*n .- z[s]; # budget

            # constrained
            if b[1] > -ϕ
                    c_c = range(cl[s], stop=c[1], length=Ic);
                    n_c =   (1 .- fac[s]*c_c.^gameta )                         # labor supply
                    n_c = [ n_c[i] < 0. ? 0. : n_c[i] for i = 1:length(n_c)]
                    b_c = -ϕ/(1+r) .+ c_c .- θ[s] .*n_c .- z[s]; # budget
                    b   = [b_c[1:Ic-1]..., b...];
                    c   = [c_c[1:Ic-1]..., c...];
            end

            # update policies
            itp = extrapolate(    interpolate((b,),c, Gridded(Linear())), Interpolations.Flat() )
            gl.c_pol[s,:] = itp(b_grid)
            gl.c_pol[c_pol .<= cmin] .= cmin

            npp   = 1 .- fac[s] .*gl.c_pol[s,:] .^gameta
            npp[npp .<= 0.] .= 0.
            gl.n_pol[s,:] .= npp
            gl.y_pol[s,:] = θ[s] .* gl.n_pol[s,:];
            bpp        = (1+r) * (b_grid .+ gl.y_pol[s,:] .- gl.c_pol[s,:] .+ z[s])
            bpp[bpp .<= -ϕ] .= -ϕ
            bpp[bpp .>= bmax] .= bmax - 0.000001
            gl.b_pol[s,:] .= bpp

            # save policies
            Tgl.c_pol_t[s, :, t] = gl.c_pol[s, :];
            Tgl.n_pol_t[s, :, t] = gl.n_pol[s, :];
            Tgl.y_pol_t[s, :, t] = gl.y_pol[s, :];
        end

        # save weights needed to iterate distribution
        idxes       = [searchsortedlast(b_grid,gl.b_pol[s, b]) for s = 1:size(gl.b_pol)[1], b = 1:size(gl.b_pol)[2] ]
        weights     = zeros(size(gl.b_pol))
        for s = 1:size(gl.b_pol)[1]
        for b = 1:size(gl.b_pol)[2]
            idx             = idxes[s,b]
            idx_prime       = minimum(  [ idx+1,size(gl.b_pol)[2] ]    )
            db              =  b_grid[idx_prime] - b_grid[idx]
            weights[s,b]    = ( gl.b_pol[s,b]  - b_grid[idx] ) / db;
        end
        end
        Tgl.ib_pol_t[:,:,t] = idxes
        Tgl.wei_t[:,:,t]    = weights
end

"""
    F = find_cl(c, j, b1, b2, r, θ, z, fac, gameta)

Helper function to find consumption at the constraint.
Returns the residual the budget constraint.

* `c` - consumption
* `j` - state index
* `b1` - bond holdings today
* `b2` - bond holdings tomorrow
* `r` - interest rate
* `θ` - state value array
* `fac` - parameter
* `gameta` - parameter
Return
* `F` - residual budget constraint
"""
function find_cl(c, j, b1, b2, r, θ, z, fac, gameta)
    # Find consumption for current assets b1 and next assets b2

    n = max(0, 1 - fac[j]*c.^gameta);
    F = b1 - b2/(1+r) - c + θ[j]*n + z[j];

    return F
end

"""
    hello(who::String)

Return "Hello, `who`".
"""
hello(who::String) = "Hello, $who"

"""
    domath(x::Number)

Return `x + 5`.
"""
domath(x::Number) = x + 5

"""
    load_parameters_iss!(gl))

Pre-load calibrated parameters to speed up calibration.
"""
function load_parameters_iss!(gl)
    gl.β = 0.9774431455410588
    gl.ϕ = 1.6043143416048709
    gl.ν = 0.1669508673012496
    gl.ψ = 15.882767548488852
    gl.B = 2.6712129167713377
end

"""
    load_parameters_tss!(gl))

Pre-load calibrated parameters to speed up calibration (terminal).
"""
function load_parameters_tss!(gl)
    gl.r = 0.003733964921583098
    gl.ϕ = 0.8983295288221897
end


text:
"""
# Description

Julia version of Spectral Methods in Matlab

# Argument
- 'N': polynomial order

# Return
- 'D': Chebyshev differentiation matrix
- 'x': Guass-Lobatto points

"""
function cheb(N)
    if N==0
        return [0], [0]
    else
        x = @. cos(pi*(0:N)/N);
        c = [2; ones(N-1); 2] .* (-1).^(0:N);
        dX = x .- x';
        D =  (c ./ c') ./ (dX + I);                # off-diagonal entries
        D = D - Diagonal(sum(D', dims = 1)[1:N+1]) # diagonal entries
        return D, x
    end
end

max_stars_repo_name: upiterbarg/QuasiGeostrophy

max_issues_repo_name: sandreza/QuasiGeostrophy

max_forks_repo_name: upiterbarg/QuasiGeostrophy


text:
idxplot = 1:50:length(soundings)
burninfrac = 0.5
nbins = 100
computeforwards, nforwards = true, 100
interpolatedr = 10
vmax, vmin = .- extrema(fbounds)
## get the summary images of resistivity and nuisances over the profile
transD_GP.TEMPEST1DInversion.summaryimages(soundings,
                            qp1=0.05,
                            qp2=0.95,
                            burninfrac = burninfrac,
                            zstart = zstart,
                            extendfrac = extendfrac,
                            dz = dz,
                            nlayers = nlayers,
                            useML = useML,
                            nmin = nmin,
                            nmax = nmax,
                            K = K,
                            demean = demean,
                            sampledc = sampledc,
                            sddc = sddc,
                            sdpos = sdpos,
                            sdprop = sdprop,
                            fbounds = fbounds,
                            λ = λ,
                            δ = δ,
                            save_freq = save_freq,
                            nuisance_sdev   = nuisance_sdev,
                            nuisance_bounds = nuisance_bounds,
                            updatenuisances = updatenuisances,
                            dr = interpolatedr,
                            fontsize = 10,
                            vmin = vmin,
                            vmax = vmax,
                            cmap="viridis",
                            figsize=(8,10),
                            topowidth=1,
                            idx = idxplot,
                            showderivs = false,
                            omitconvergence = false,
                            preferEright = false,
                            preferNright = false,
                            saveplot = true,
                            labelnu = ["zRx m", "xRx m", "pitch "*L"^\circ"]
                        )

## now plot individual soundings in in idxplot
transD_GP.TEMPEST1DInversion.plotindividualsoundings(soundings,
                        burninfrac = burninfrac,
                        zstart = zstart,
                        extendfrac = extendfrac,
                        dz = dz,
                        nlayers = nlayers,
                        nmin = nmin,
                        nmax = nmax,
                        K = K,
                        fbounds = fbounds,
                        λ = λ,
                        δ = δ,
                        nuisance_sdev   = nuisance_sdev,
                        nuisance_bounds = nuisance_bounds,
                        updatenuisances = updatenuisances,
                        nbins = nbins,
                        figsize  = (12,6),
                        zfixed   = zfixed,
                        ρfixed   = ρfixed,
                        ntimesperdecade = ntimesperdecade,
                        nfreqsperdecade = nfreqsperdecade,
                        computeforwards = computeforwards,
                        nforwards = nforwards,
                        idxcompute = idxplot)

max_stars_repo_name: GeoscienceAustralia/HiQGA

max_issues_repo_name: GeoscienceAustralia/HiQGA


text:
import ..KERNEL.System

"Save information needed to identify a SSBond"
struct SSBond    #Tuple of chain_id and residue_number
    chain_id::Char
    res_num::Int64
end

"Packages the parsed fields from a conect line in a struct to avoid allocations"
mutable struct ConectLineParams
    bonding_atoms::Vector{Int}
    tokens::Vector{String}

end

"Packages the parsed fields from an atom line in a struct to avoid allocations"
mutable struct AtomLineParams
    name::String
    x::Float64
    y::Float64
    z::Float64
    elem::String
    occupancy::Union{Float64,Nothing}
    temp_factor::Union{Float64,Nothing}
    charge::Union{Float64,Nothing}
    serial::Int64

    AtomLineParams() = new()
end

"""
    parseConectLine(line::String)
Parses a line starting with "CONNECT" in a PDB file.
"""
parseConectLine(params::ConectLineParams, line::String) = begin

    if length(line) == 16
        push!(params.tokens, line[12:16])
    elseif length(line) == 21
       push!(params.tokens, line[12:16], line[17:21])
    elseif length(line) == 26
        push!(params.tokens, line[12:16], line[17:21],
                line[22:26])
    else
        push!(params.tokens, line[12:16], line[17:21],
                line[22:26], line[27:31])
    end

    for x in params.tokens
        if strip(x) != ""
            push!(params.bonding_atoms, parse(Int,strip(x)))
        end
    end
    return params
end

"""
    parseAtomLine(line::String)
Parses a line starting with "ATOM" or "HET" in a PDB file.
"""
parseAtomLine(params::AtomLineParams, line::String) = begin

    params.name = strip(line[13:16])
    params.x = parse(Float64,line[31:38])
    params.y = parse(Float64,line[39:46])
    params.z = parse(Float64,line[47:54])
    params.elem = strip(line[77:78])
    params.occupancy = strip(line[55:60]) == "" ? nothing : parse(Float64,line[55:60])
    params.temp_factor = strip(line[61:66]) == "" ? nothing : parse(Float64,line[61:66])

    if length(line) >= 80
        params.charge = strip(line[79:80]) == "" ? nothing : parse(Float64,line[79:80])
    else
        params.charge = nothing
    end

    result = Atom(params.name,params.x,params.y,params.z,params.elem,params.charge,
                    params.occupancy,params.serial,params.temp_factor)::Atom
    if startswith(line,"HET")
        setProperty(result,("hetero",true))
    end

    return result
end

"""
    compare_ssbonds(b1::SSBond, b2::SSBond)
Comparison function for sorting.
"""
compare_ssbonds(b1::SSBond, b2::SSBond) = begin
    b1.chain_id < b2.chain_id && return true
    b1.chain_id == b2.chain_id && b1.res_num < b2.res_num && return true
    return false
end

"""
    parseSSBondLine(line::String)
Parses a line starting with "SSOBND" ina PDB file.
"""
parseSSBondLine(line::String) = begin
    ssbond1::SSBond = SSBond(line[16], parse(Int,strip(line[18:21])))
    ssbond2::SSBond = SSBond(line[30] ,parse(Int,strip(line[32:35])))
    return (ssbond1, ssbond2)
end #returns 2 ssbonds
#only use following Function to build a System initially

"""
    adjust_ter_indices(indices::Vector{Int64}, ter_pos::Vector{Int64})
Adjusts the indices in a list of stom serials read from a PDB file. The "TER" entries in PDB files
have a serial number like the Atoms, which makes the serial numbers not continuous.
"""
adjust_ter_indices(indices::Vector{Int64}, ter_pos::Vector{Int64}) = begin
    for i in 1:length(indices)
         c = 0
         for ter in ter_pos
            if indices[i] > ter
                c += 1
            end
         end
         indices[i] -= c
    end
    return indices
end
adjust_ter_indices(x::Int64, ter_pos::Vector{Int64}) = adjust_ter_indices([x],ter_pos)

"""
    parsePDB(path::String)
Parses a PDB file from `path`. Creates a representation using [`KERNEL`](@ref KERNEL_header) types.\n
If the file has multiple models the property \"fromNMR\" is added to `root`.
See [Hint](https://www.wwpdb.org/documentation/file-format-content/format33/sect9.html#MODEL).
"""
parsePDB(path::String) = begin
     # `ssbonds` holds all the ssbonds in order of appearance in the file
    ssbonds::Vector{SSBond} = Vector{SSBond}()

    # When reading the atoms from the file, relate the atom to an SSBond endpoint
    ssbonds_to_atoms::Dict{SSBond, Int64} = Dict{SSBond, Int64}() #mapping ssbonds to atom serial number

    #List of pairs of atoms bonded by a SSBond
    ssbond_pairs::Vector{Tuple{SSBond,SSBond}} = Vector{Tuple{SSBond,SSBond}}()

    current_ssbond_index::Int = 1

    root::System = System()
    target_system::System = System()
    atoms::Vector{Atom} = Atom[]
    ter_positions::Vector{Int} = Int64[]

    latest_chain::Chain = Chain()  #assume chains are named in alphabetical order with capitalized letters
    latest_residue::Residue = Residue()
    latest_residue.res_number_ = -127   #init value, which hopefully is NOT the same as the first residue-number in the file
    atom_counter::Int = 0

    latest_chain.id_ = '-'  #init value, which hopefully is NOT the same as the first chain id in the file
    appendChild(target_system, latest_chain)

    seen_header::Bool = false::Bool
    ready_to_sort_ssbonds::Bool = false
    seen_ssbonds::Bool = false
    seen_atoms::Bool = false
    seen_model::Bool = false

    #init fields for parsing a "CONECT" line
    conect_line_params::ConectLineParams = ConectLineParams(Int[],String[])

    #init fields for parsing a "ATOM" line
    atom_line_params::AtomLineParams = AtomLineParams()

    record_residue_name::String = ""

    open(path) do file
        for line in eachline(file)
            #header
            if !seen_header && startswith(line,"HEADER")
                if length(line) >= 66
                    root.name_ = strip(line[63:66])
                end
                seen_header = true::Bool
            end

            if startswith(line,"TER")
                push!(ter_positions, parse(Int64,strip(line[7:11])))
            end

            #save infos about the atoms which participate in ssbonds
            if !seen_ssbonds && startswith(line, "SSBOND")
                pair = parseSSBondLine(line)
                push!(ssbond_pairs, pair)
                push!(ssbonds, pair...)
                ready_to_sort_ssbonds = true
            end

            #after having seen all of the ssbonds sort them to later find the correct atoms easier
            if ready_to_sort_ssbonds && !startswith(line,"SSBOND")
                seen_ssbonds = true
                ready_to_sort_ssbonds = false
                sort!(ssbonds, alg=InsertionSort, lt = compare_ssbonds)
            end

            #atoms of name "SG" belong to ssbonds. associate atom to ssbond entry
            if (current_ssbond_index <= length(ssbonds)) && startswith(line,"ATOM") &&
                                                           (strip(line[13:16]) == "SG")
                #find items from SSBOND
                current_ssbond = ssbonds[current_ssbond_index]
                #replace "SG" with a list of possible names of atoms which engage in ssbonds
                if line[22] == current_ssbond.chain_id && current_ssbond.res_num ==  parse(Int,strip(line[23:26]))
                    ssbonds_to_atoms[current_ssbond] = adjust_ter_indices(parse(Int64,strip(line[7:11])), ter_positions)[1]
                    current_ssbond_index += 1
                end
            end

            #model lines - means we need to construct a new model
            if startswith(line,"MODEL")
                if !seen_model
                    nothing
                else
                    #finalize target_system and append it to system
                    appendChild(root, target_system)
                    target_system = System()
                    latest_chain = Chain()
                    latest_chain.id_ = 'A'

                    prev_residue_name = latest_residue.name_
                    latest_residue = Residue()
                    latest_residue.res_number_ = 1
                    latest_residue.name_ = prev_residue_name
                    latest_residue.is_hetero_ = false

                    appendChild(target_system, latest_chain)
                    appendChild(latest_chain, latest_residue)
                    atom_counter = 0
                end
                seen_model = true
            end

            if (!seen_atoms && startswith(line,"ATOM") ||
                                !seen_atoms && startswith(line,"HETATM"))  &&
                                line[17] in (' ','A')
                atom_counter += 1
                record_chain_id = line[22]
                latest_chain.id_ == '-' && (latest_chain.id_ = record_chain_id)
                record_residue_name = strip(line[18:20])
                record_residue_number = parse(Int64,strip(line[23:26]))

                #create chain if new chain
                if record_chain_id != latest_chain.id_
                    #is the next residues belong to another chain, create the chain
                    latest_chain = Chain()
                    latest_chain.id_ = record_chain_id
                    appendChild(target_system, latest_chain)
                end

               if record_residue_number != latest_residue.res_number_
                    latest_residue = Residue()
                    latest_residue.res_number_ = record_residue_number
                    latest_residue.name_ = record_residue_name
                    latest_residue.is_hetero_ = false
                    if latest_residue.name_ in Amino_Acids
                        setProperty(latest_residue, ("amino_acid",true))
                    end
                    appendChild(latest_chain, latest_residue)

                end
               #create atoms
                atom_line_params.serial = atom_counter
                parsed_atom::Atom = parseAtomLine(atom_line_params, line)
                #the only property right now can be the hetero property
                if length(parsed_atom.properties_) != 0
                    latest_residue.is_hetero_ = true
                end

                appendChild(latest_residue, parsed_atom)
            end

            if startswith(line, "CONECT")
                if seen_atoms == false
                    seen_atoms = true
                    atoms = collectAtoms(target_system)
                end

                if seen_model
                    appendChild(root, target_system)
                end

                empty!(conect_line_params.bonding_atoms)
                empty!(conect_line_params.tokens)
                bonding_atom_serials = parseConectLine(conect_line_params, line).bonding_atoms
                bonding_atom_serials = adjust_ter_indices(bonding_atom_serials, ter_positions)
                for serial in bonding_atom_serials[2:end]
                    createBond(atoms[bonding_atom_serials[1]], atoms[serial],
                                order=ORDER__SINGLE, type=TYPE__COVALENT)
                end
            end
        end
    end

    #for each ssbond pair, make a bond betwen the two corresponding atoms using the dict
    for (ssbond1, ssbond2) in ssbond_pairs
        deleteBond(atoms[ssbonds_to_atoms[ssbond1]], atoms[ssbonds_to_atoms[ssbond2]])
        a = createBond(atoms[ssbonds_to_atoms[ssbond1]], atoms[ssbonds_to_atoms[ssbond2]],
                                order=ORDER__SINGLE, type=TYPE__DISULPHIDE_BRIDGE)
    end

    if root.name_ == ""
        temp = findlast("/",path)
        if isnothing(temp)
           root_name_ = path[1:end-4]
        else
            root.name_ = path[temp[1]+1:end-4]
        end
    end

    if !seen_model
        target_system.name_ = root.name_
        i = 1
        for chain in collectChains(target_system)
            if chain.id_ < 'A'+(i-1)
                idx = Int(chain.id_) - Int('A') +1
                if idx < 1      #catch chains with empty/too low char
                    chain.id_ = 'A'+(i-1)
                    i += 1
                    continue
                end
                ch_l = collectChains(target_system)[idx]
                appendSibling(ch_l.last_child_, chain.first_child_)
                removeChild(chain)
            else
                chain.id_ = 'A'+(i-1)
                i += 1
            end
        end

        return target_system
    else
        setProperty(root,("fromNMR",true))
        for (num,system) in enumerate(collectSystems(root)[2:end])
            system.name_ = root.name_*"-"*string(num)
            i = 1
            for chain in collectChains(target_system)
                if chain.id_ < 'A'+(i-1)
                    idx = Int(chain.id_) - Int('A') +1
                    ch_l = collectChains(target_system)[idx]
                    appendSibling(ch_l.last_child_, chain.first_child_)
                    removeChild(chain)
                else
                    chain.id_ = 'A'+(i-1)
                    i += 1
                end
            end
        end

        #copy bonds over
        last_sys = last(getChildren(root))
        last_sys_bonds = collectBonds(last_sys)
        atom_serials_pairs::Vector{Tuple{Int64, Int64}} = [(bond.source_.serial_, bond.target_.serial_)
                                                        for bond in last_sys_bonds]

        for sys in getChildren(root)
            atoms = collectAtoms(sys)
            for bond in last_sys_bonds
                createBond(atoms[bond.source_.serial_], atoms[bond.target_.serial_],
                            name = bond.name_, order = bond.bond_order_, type = bond.bond_type_,
                            properties = bond.properties_)
            end

        end
        return root
    end

    return root
end

# BioStructures reads structure and internal parser only the bonds
"""
    System(path::String)
Constructs a `System` from a PDB file at `path`.
"""
function System(path::String)
    if endswith(path, ".pdb")
        root = parsePDB(path)::System
        for sys in collectSystems(root)
            sys.number_of_children_ = countChildren(sys)

        end
        return root
    end
end


text:
##################################################################################################
# Integration with UncertainData.jl (sampling from full supports of the furnishing distributions)
#################################################################################################

#Vectors of uncertain values
uvals_x = [UncertainValue(Normal, rand(Normal(0, 5)), abs(rand(Normal(0, 3)))) for i = 1:100]
uvals_y = [UncertainValue(Normal, rand(Normal(0, 5)), abs(rand(Normal(0, 3)))) for i = 1:100];

# UncertainValueDataset
UVX = UncertainValueDataset(uvals_x)
UVY = UncertainValueDataset(uvals_y)

# UncertainIndexDataset
UVX_idx = UncertainIndexDataset(uvals_x)
UVY_idx = UncertainIndexDataset(uvals_y)

# Real-valued vectors
x = resample.(uvals_x);
y = resample.(uvals_y);

ctest = ConvergentCrossMappingTest(timeseries_lengths = [50, 60, 75, 80], n_reps = 10)

# One set of cross mappings per time series length
@test causality(x, y, ctest) isa Vector{Vector{T}} where T
@test causality(uvals_x, uvals_y, ctest) isa Vector{Vector{T}} where T
@test causality(x, uvals_y, ctest) isa Vector{Vector{T}} where T
@test causality(uvals_x, y, ctest) isa Vector{Vector{T}} where T
@test causality(UVX, UVY, ctest) isa Vector{Vector{T}} where T
@test causality(x, UVY, ctest) isa Vector{Vector{T}} where T
@test causality(UVX, y, ctest) isa Vector{Vector{T}} where T

# Four time series lengths, so there should be four vectors of cross map values
@test causality(x, y, ctest) |> length == 4
@test causality(uvals_x, uvals_y, ctest) |> length == 4
@test causality(x, uvals_y, ctest) |> length == 4
@test causality(uvals_x, y, ctest) |> length == 4
@test causality(UVX, UVY, ctest) |> length == 4
@test causality(x, UVY, ctest) |> length == 4
@test causality(UVX, y, ctest) |> length == 4

# There should be ten cross map values per time series length
@test causality(x, y, ctest)[1] |> length == 10
@test causality(uvals_x, uvals_y, ctest)[1] |> length == 10
@test causality(x, uvals_y, ctest)[1] |> length == 10
@test causality(uvals_x, y, ctest)[1] |> length == 10
@test causality(UVX, UVY, ctest)[1] |> length == 10
@test causality(x, UVY, ctest)[1] |> length == 10
@test causality(UVX, y, ctest)[1] |> length == 10

################################################################
# Integration with UncertainData.jl (with sampling constraints)
################################################################
onevar_constraints = ConstrainedResampling(TruncateStd(1))
twovar_constraints = ConstrainedResampling(TruncateStd(2), TruncateStd(1))

test_ccm = ConvergentCrossMappingTest(timeseries_lengths = [50, 60], n_reps = 5)
@test causality(x, y, test_ccm, twovar_constraints) isa Vector{Vector{T}} where T<:Real
@test causality(uvals_x, uvals_y, test_ccm, twovar_constraints) isa Vector{Vector{T}} where T<:Real
@test causality(x, uvals_y, test_ccm, onevar_constraints) isa Vector{Vector{T}} where T<:Real
@test causality(uvals_x, y, test_ccm, onevar_constraints) isa Vector{Vector{T}} where T<:Real
@test causality(UVX, UVY, test_ccm, twovar_constraints) isa Vector{Vector{T}} where T<:Real
@test causality(uvals_x, UVY, test_ccm, twovar_constraints) isa Vector{Vector{T}} where T<:Real
@test causality(UVX, uvals_y, test_ccm, twovar_constraints) isa Vector{Vector{T}} where T<:Real
@test causality(x, UVY, test_ccm, onevar_constraints) isa Vector{Vector{T}} where T<:Real
@test causality(UVX, y, test_ccm, onevar_constraints) isa Vector{Vector{T}} where T<:Real