text:
module FoodCachingExperiments

using BSON, StatsModels, DataFrames, Random, LinearAlgebra, Distributions,
      HypothesisTests, DataFramesMeta, Unitful, StatsBase, CodecZstd

export run!, summarize, statistical_tests, results, nbirds, target

include("language.jl")
include("experiments.jl")
include("anova.jl")
include("statisticaltests.jl")
include("fileio.jl")

const EXPERIMENTS = Dict{Symbol, Experiment}()

for filename in readdir(joinpath(@__DIR__, "protocols"))
    if splitext(filename)[end] == ".jl"
        include(joinpath(@__DIR__, "protocols", filename))
    end
end

run!(name::Symbol, models) = run!(EXPERIMENTS[name], models)
summarize(name::Symbol, data) = summarize(EXPERIMENTS[name], data)
function summarize(e::Experiment, data)
    ag = _summarize(e, data)
    sort!(ag, default_keycols(ag))
end
statistical_tests(name::Symbol, data) = statistical_tests(EXPERIMENTS[name], data)
results(name::Symbol, data) = results(EXPERIMENTS[name], data)
nbirds(name::Symbol) = nbirds(EXPERIMENTS[name])
target(name::Symbol) = target(EXPERIMENTS[name])
foodtypes(name::Symbol) = foodtypes(EXPERIMENTS[name])
asdataframe(name::Symbol, x, k = Ref(0)) = asdataframe(EXPERIMENTS[name], x, k)

function __init__()
    filename = joinpath(@__DIR__, "..", "data", "processed", "experiments.bson.zstd")
    if !isfile(filename)
        @warn "$filename not found. Please run `julia script/process_experiments.jl`."
    else
        for (name, experiment) in bload(filename)
            EXPERIMENTS[name] = experiment
        end
    end
end

end


text:
using CoupledDipoles
using Test

@testset "CoupledDipoles.jl" begin
    @testset "3D scalar: 2 atoms" begin
        r =[1 2 0;
            1 0 1.0]
        atoms = Atom(Cube(), Array(transpose(r)), 10)
        laser = Laser(PlaneWave3D([0,0,1]), 1e-6, 1.0)
        problem = LinearOptics(Scalar(), atoms, laser)
        ωₙ, Γₙ = get_spectrum(problem)

        Γₙ_ans = [0.675922453937850, 0.324077546062151]
        ωₙ_ans = [0.861973588757495, 1.13802641124251]
        @test all(sort(ωₙ_ans) .≈ sort(ωₙ))
        @test all(sort(Γₙ_ans) .≈ sort(Γₙ))

        βₛₛ = get_steady_state(problem)
        βₛₛ_ans = [0.000695478188875652-0.000416450295287958*im,
                   0.000554066244758147+0.000208376732488666*im ]
        @test all(βₛₛ_ans .≈ βₛₛ)
    end
    @testset "3D scalar: 3 atoms" begin
        r =[ 0 0 0;
             1 1 1;
            -1 1 0.5]
        atoms = Atom(Cube(), Array(transpose(r)), 10)
        laser = Laser(PlaneWave3D([0,0,1]), 1e-6, 0.0)
        problem = LinearOptics(Scalar(), atoms, laser)
        ωₙ, Γₙ = get_spectrum(problem)

        Γₙ_ans = [0.293248907595325, 0.15447403278592, 1.05227705967608]
        ωₙ_ans = [-0.0853652197846037, -0.0371827945713394,0.12254801435943]
        @test all(sort(ωₙ_ans) .≈ sort(ωₙ))
        @test all(sort(Γₙ_ans) .≈ sort(Γₙ))

        R_nm = get_pairwise_matrix(atoms.r)
        R_nm_ans = [0 1.7320508075688 1.5;
                    1.73205080756888 0 2.06155281280883;
                    1.5 2.06155281280883 0]
        @test all(R_nm_ans .≈ R_nm)

        G = get_interaction_matrix(problem)
        G_ans = [-0.5 -0.284930049591257-0.04634868038121*im -0.332498328868018+0.0235790672225676*im;
                 -0.284930049591257-0.046348680381261*im -0.5 -0.213910742134572-0.11430539712927*im;
                 -0.332498328868018+0.0235790672225676*im -0.213910742134572-0.11430539712197*im -0.5]
        @test all(G .≈ G_ans)
    end
end

ext: jl


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

using Markdown
using InteractiveUtils

# ╔═╡ a0736b60-3c7d-11ec-1b71-27fd38cbf6e4
begin

# Import the necessary packages

using Distributions
using DelimitedFiles
using DataFrames
using Optim
using CSV
using Random
using Statistics
using FiniteDiff
using LinearAlgebra
using ForwardDiff
using JuMP
import Ipopt

end

# ╔═╡ 7b7526d0-3c9d-11ec-3575-37a66c520386
# Replication of Bruhin, Fehr, and Schunk, 2019, "Many Faces of Human Sociality: Uncovering the Distribution and Stability of Social Preferences"

# ╔═╡ 7b36490e-3c9d-11ec-1ea6-dbf0082362dc
# Authors:
# Massimiliano Pozzi (Bocconi University, pozzi.massimiliano@studbocconi.it)
# Salvatore Nunnari (Bocconi University, https://snunnari.github.io, salvatore.nunnari@unibocconi.it)

# ╔═╡ c7fb5b79-7d87-49c5-9459-a360e1057715
# This notebook works with both Neptune and Pluto; the "begin" and "end" tags in each cell are necessary for Pluto notebooks but do not play any role in Neptune

# ╔═╡ 7d517760-4913-11ec-24fb-03122e4dda9a
# This notebook was tested with the following packages versions:

# Pozzi: julia 1.6.1, Pluto 0.14.7, Neptune 0.14.0, Distributions 0.25.17, DataFrames 1.2.2, Optim 1.4.1, CSV 0.9.5, FiniteDiff 2.8.1, ForwardDiff 0.10.20, JuMP 0.21.10

# ╔═╡ e0965590-3c7d-11ec-13b6-a7d15db3fd9f
begin

# 1. Data Cleaning and Preparation

# We import the relevant datasets containing the data on the 39 dictator games and 78 reciprocity games in Session 1 and Session 2. As the authors, we remove from these datasets those subjects that behaved very inconsistenly throughout the games. These subjects are identified from the individual estimates that we do not run in this notebook. The file "dropped_subjects_section4paragraph2.csv" contains the ID of these subjects.

# Import data for session 1 and create a DataFrame
data, header = readdlm(raw"../input/choices_exp1.csv",',',header=true)
dt1 = DataFrame(data, vec(header))

# Import data with ID of subjects to drop
data_drop = readdlm(raw"../input/dropped_subjects_section4paragraph2.csv",',',header=false)

# Drop the session 1 subjects whose IDs are listed in the data_drop dataframe (14 individuals)

dt1 = dt1[[i for i=1:length(dt1.sid) if !(dt1.sid[i] in(data_drop))],:]

end

# ╔═╡ d8928282-3c85-11ec-00ca-c5215c731b33
begin

# Create indicators_x and indicators_y. These are the columns s, r, q, v

indicators_x = hcat(dt1.s_x,dt1.r_x,dt1.q,dt1.v)
indicators_y = hcat(dt1.s_y,dt1.r_y,dt1.q,dt1.v)

end;

# ╔═╡ 56a26400-422c-11ec-3599-77a4d3d0e180
# Define the function to minimize. This is the negative of the log likelihood of observing our data given the parameters of the model.

# For more information we refer the reader to equation (4) and (5) in the paper or to the python notebook in this repository

# v is the vector of parameters (θ,σ)
# y is the choice of the player
# self_x is the payoff when choosing x (left)
# self_y is the payoff when choosing y (right)
# indicators are the ones explained above

function loglikeM(v,y,self_x, other_x, self_y,other_y,indicators_x,indicators_y)

	beta  = v[1:(length(v)-1)]  # vector of parameters alpha,beta,gamma,delta
    sigma = exp(v[length(v)])   # choice sensitivity σ. We take the exp since it can only take positive values
    lli = (indicators_x * beta) # we obtain a 1x18720 vector. Each element is (αs+βr+γq+δv) for the single game
    rli = (indicators_y * beta)
	uleft  = sigma.*( (1 .-lli) .*self_x .+lli .*other_x) # utility when choosing allocation X (left)
	uright = sigma.*( (1 .-rli) .*self_y .+rli .*other_y) # utility when choosing 	allocation Y (right)

	# probs is a 1x18720 vector containing the likelihood of observing the data of a single game
	probs  = (exp.(uleft)./(exp.(uleft).+exp.(uright))).^y .* (exp.(uright)./
		(exp.(uleft).+exp.(uright))).^(1 .-y)
	nll = -sum(log.(probs)) # negative log-likelihood
	return nll
	end

# ╔═╡ dfcceb10-3c7d-11ec-2664-9d36301add75
begin

# initialize random starting guesses. Same random initial guesses as in the authors' R code

beta_init = rand(Uniform(0.01,0.02),4) # close to zero for α, β, γ, δ
sigma_init = log.(rand(Uniform(0.05,0.8),1)/mean([mean(dt1.self_x),mean(dt1.other_x),mean(dt1.self_y),mean(dt1.other_y)])) # log since in function we take exp
v0 = [beta_init;sigma_init]

end

# ╔═╡ a3f76eb0-3d51-11ec-20e7-0d245610572b
# Wrap our function. Optim.optimize can only take one argument

function ll(v)
	fll = loglikeM(v,dt1.choice_x,dt1.self_x,dt1.other_x,dt1.self_y,dt1.other_y,indicators_x,indicators_y)
	return fll
end

# ╔═╡ 1baa58c0-4494-11ec-1794-6b4d74f6812b
begin

# This cell computes the aggregate estimate for session 1 using the BFGS algortihm. We do not provide the analytical gradient since the one computed by the algorithm is precise enough for this problem

sol = optimize(ll,v0,method=BFGS(),iterations = 10000,store_trace=true, extended_trace=true, autodiff=:forward)
res_s1 = Optim.minimizer(sol) # estimates
results_s1 = [res_s1[1:4];exp(res_s1[5])] # Need to take exp for sigma
vloglike_s1 = -Optim.minimum(sol) # Value of loglikelihood evaluated in the minimum

end

# ╔═╡ f576544e-3d55-11ec-009e-5ddcc502108b
# Compute the individual cluster robust standard errors.

# For a more detailed explanation on how to compute cluster robust standard errors we refer the reader to the python notebook or to David A. Freedman, 2006, "On The So-Called 'Huber Sandwich Estimator' and 'Robust Standard Errors'", (The American Statistician 60:4, 299-302)

# The formula is the following:

# adj * (inv_hessian * grad_con * inv_hessian).

# adj is a correction for the degree of freedoms and the number of clusters. inv_hessian is the inverse of the hessian of the negative log-likelihood evaluated at the minimum. grad_con is is a 5x5 matrix of gradient contributions. It is computed as the sum over all clusters of the outer product of the sum of the single observation gradient evaluated in the minimum for all observations of an individual cluster.

# Define the function that computes the matrix of individual gradient contribution

function congradloglike(v,y,self_x,other_x,self_y,other_y,indicators_x,indicators_y,clusters)

	# parameters. Same as before
	beta  = v[1:(length(v)-1)]
    sigma = exp(v[length(v)])

	# This part is the same as before
    lli = (indicators_x * beta)
    rli = (indicators_y * beta)
	utl  = sigma.*( (1 .-lli) .*self_x .+lli .*other_x)
	utr  = sigma.*( (1 .-rli) .*self_y .+rli .*other_y)
	probs  = (exp.(utl)./(exp.(utl).+exp.(utr))).^y .* (exp.(utr)./(exp.(utl).+exp.(utr))).^(1 .-y)

	# Compute the analytical gradient
	probsm = ones(size(indicators_x,1), size(indicators_x,2)) .* probs
	u = ones(size(indicators_x,1), size(indicators_x,2)) .*exp.(utl)
	w = ones(size(indicators_x,1), size(indicators_x,2)) .*(exp.(utl) .+ exp.(utr))
	up = sigma * indicators_x .* (ones(size(indicators_x,1), size(indicators_x,2)) .* (other_x .-self_x)) .* u
	wp = up .+ sigma .* indicators_y .* (ones(size(indicators_x,1), size(indicators_x,2)) .*((other_y.-self_y).*exp.(utr)))

	bgradi = (-1).^(1 .-y).*((up.*w .- u.*wp)./w.^2)./probsm
	up2 = utl .* u[:,1]
	wp2 = up2 .+ utr .* exp.(utr)
	ggradi = ((-1).^(1 .-y) .* (((up2.*w[:,1] - u[:,1].*wp2)./w[:,1].^2)./probs))

	gradi = [bgradi ggradi] # 18720x5 matrix of partial derivatives for each obs.

	cl = unique(clusters)
	j = length(cl) # nr of clusters is length of unique individual = 160
	k = size(gradi,2) # nr of columns in gradi = 5
	sandwich = zeros(k,k)
	for i=1:j # sum columns by columns for individual i
		sel = clusters .== cl[i] # vector. 1 if same cluster else 0
		sel = findall(>(0), sel) # indices of 1
		gradsel = sum(gradi[sel,:],dims=1)
		sandwich .+= (gradsel' * gradsel) # sum of outer product
	end
	return sandwich
end

# ╔═╡ f558e140-3d55-11ec-075d-9d31f531f525
begin

# Compute the cluster robust standard errors.

# This is the code to use the inverse of the hessian stored in the BFGS algorithm. We use the one computed using the Pkg ForwardDiff but they should be the same since we used the option autodiff=:forward

inverse_hessian_s1 = sol.trace[end].metadata["~inv(H)"]

# Compute the hessian using the ForwardDiff package

inv_hess_num = inv(ForwardDiff.hessian(ll,res_s1))

# degree of freedom and cluster adjustment

adj = (length(dt1.self_x)-1) / (length(dt1.self_x)-length(res_s1)) * (160/159)

grad_contribution = congradloglike(res_s1,dt1.choice_x,dt1.self_x,dt1.other_x,dt1.self_y,dt1.other_y,indicators_x,indicators_y,dt1.sid)

# Sandwich formula for the clustered standard errors

varcov_s1 = adj .* (inv_hess_num * grad_contribution * inv_hess_num)
se_s1 = sqrt.(diag(varcov_s1))
se_s1 = [se_s1[1:4]; sqrt(exp(res_s1[5])^2*se_s1[5]^2)]

# sqrt(exp(res_s1[5])^2*se_s1[5]^2) is the formula to go from standard error log(sigma) to standard error of sigma. We use the Delta method

end

# ╔═╡ f970b080-3cb2-11ec-0e09-65acfacc9837
# Compute aggregate estimates for session 2

# ╔═╡ 9e5d90a0-3cb2-11ec-2efc-8189c078aee7
begin

# Import session 2 data and remove subjects to drop as we did for session 1 above

data2, header2 = readdlm(raw"../input/choices_exp2.csv",',',header=true)
dt2 = DataFrame(data2, vec(header2))
dt2 = dt2[[i for i=1:length(dt2.sid) if !(dt2.sid[i] in(data_drop))],:]

end

# ╔═╡ e6ab8b50-3cb2-11ec-0b39-e5424b4fdbd2
begin

# Create indicators_x and indicators_y. These are the columns s, r, q, v

indicators_x2 = hcat(dt2.s_x,dt2.r_x,dt2.q,dt2.v)
indicators_y2 = hcat(dt2.s_y,dt2.r_y,dt2.q,dt2.v)

end;

# ╔═╡ d554ea40-3cb2-11ec-1642-99c3aadfd3e4
begin

# This cell computes the aggregate estimate for session 2

function ll2(v)
	fll = loglikeM(v,dt2.choice_x,dt2.self_x,dt2.other_x,dt2.self_y,dt2.other_y,indicators_x2,indicators_y2)
	return fll
end

sol2 = optimize(ll2,v0,method=BFGS(),iterations = 10000,store_trace=true, extended_trace=true, autodiff=:forward)
res_s2 = Optim.minimizer(sol2)
results_s2 = [res_s2[1:4];exp(res_s2[5])]
vloglike_s2 = -Optim.minimum(sol2)

end

# ╔═╡ 586df880-3cd7-11ec-348c-6133b9ddba73
begin

# Compute the cluster robust standard errors.

inv_hessian_num2 = inv(ForwardDiff.hessian(ll2,res_s2))

adj2 = (length(dt2.self_x)-1) / (length(dt2.self_x)-length(res_s2)) * (160/159) # degree of freedom adjustment

grad_contribution2 = congradloglike(res_s2,dt2.choice_x,dt2.self_x,dt2.other_x,dt2.self_y,dt2.other_y,indicators_x2,indicators_y2,dt2.sid)

# Sandwich formula for the clustered standard errors

varcov_s2 = adj2 .* (inv_hessian_num2 * grad_contribution2 * inv_hessian_num2)
se_s2 = sqrt.(diag(varcov_s2))
se_s2 = [se_s2[1:4]; sqrt(exp(res_s2[5])^2*se_s2[5]^2)]

end

# ╔═╡ c32cedfe-daf7-4b0b-81bc-e8028a050294
begin

# Hypothesis Testing

# We now do some hypothesis testing on the parameters we obtained in session 1 and session 2. We first compute the z-test statistics and the corresponding p-values to check if each parameter we obtained is statistically different from zero. We then compute the p-value of a z-test to check if the parameters we obtained in session 1 are statistically different from the parameters we obtained in session 2.

# Compute the z-test statistics and the corresponding p-values to check if the parameters are statistically different from zero

# Session 1

zvalues_s1 = results_s1 ./ se_s1
pvalues_s1 = 2*(1 .-cdf.(Normal(0,1),zvalues_s1))

# Session 2

zvalues_s2 = results_s2 ./ se_s2
pvalues_s2 = 2*(1 .-cdf.(Normal(0,1),zvalues_s2))

# Check if parameters obtained in session 1 are statistically different from parameters in session 2

# First we need the variance for the parameters in session 1 and session2

var_s1 = se_s1.^2
var_s2 = se_s2.^2

# Now we compute the p-values of the z-test statistics

zvalues_s1s2 = abs.(results_s1.-results_s2) ./ sqrt.(var_s1.+var_s2)
pvalues_s1s2 = 2*(1 .-cdf.(Normal(0,1),zvalues_s1s2))

end

# ╔═╡ 802c89c0-3d60-11ec-1d69-0dcf94880bf3
begin

# Create a new table and save the results in a csv file. We round results up to the third decimal

parameters_name = ["α: Weight on other's payoff when behind",
                   "β: Weight on other's payoff when ahead",
                   "γ: Measure of positive reciprocity",
                   "δ: Measure of negative reciprocity",
                   "σ: Choice sensitivity"]

table1 = DataFrame()
table1.parameters = parameters_name
table1.estimates_s1 = round.(results_s1, digits=3)
table1.standarderr_s1 = round.(se_s1, digits=3)
table1.zstat_s1 = round.(zvalues_s1, digits=3)
table1.pval_s1 = round.(pvalues_s1, digits=3)
table1.estimates_s2 = round.(results_s2, digits=3)
table1.standarderr_s2 = round.(se_s2, digits=3)
table1.zstat_s2 = round.(zvalues_s2, digits=3)
table1.pval_s2 = round.(pvalues_s2, digits=3)
table1.pval_s1s2 = round.(pvalues_s1s2, digits=3)

CSV.write(raw"../output/table1_julia.csv", table1)

table1

end

# ╔═╡ 4ee3a2da-1a91-450b-97d8-4108fa08944e
begin

# We now compute the point estimates for session 1 using JuMP.

# Define a function that takes the variables separately and returns the negative log likelihood. This simplifies the way we can define the variables in the JuMP model later on

function lljump(alpha,beta,gamma,delta,sigma)
		v = [alpha,beta,gamma,delta,sigma]
		flljump = loglikeM(v,dt1.choice_x,dt1.self_x,dt1.other_x,dt1.self_y,dt1.other_y,indicators_x,indicators_y)
	return flljump
end

# Define a new model using JuMP. We use the Ipopt optimizer

model = Model(Ipopt.Optimizer; bridge_constraints = false)

# You need to clear the julia variables, else you will not be able to run this cell a second time. This is a known problem when using Pluto

alpha,beta,gamma,delta,sigma= nothing, nothing, nothing, nothing, nothing

# Define the variables and the starting values. Our problem is unconstrained.

@variable(model, alpha, start = v0[1])
@variable(model, beta,  start = v0[2])
@variable(model, gamma, start = v0[3])
@variable(model, delta, start = v0[4])
@variable(model, sigma, start = v0[5])
@NLobjective(
        model,
        Min,
        lljump(alpha,beta,gamma,delta,sigma) # the function to minimize
    )

optimize!(model)
solution_summary(model)

end

# ╔═╡ bbe59950-4915-11ec-289a-c97076f93ca1
begin

# In this cell we show that the point estimates for Session 1 are the same when using Optim or JuMP

JuMP_s1 = [round(value(alpha), digits=3); round(value(beta), digits=3); round(value(gamma), digits=3); round(value(delta), digits=3); round(exp(value(sigma)), digits=3)]

table_diff = DataFrame()
table_diff.parameters = parameters_name
table_diff.Optim_s1 = round.(results_s1, digits=3)
table_diff.JuMP_s1  = JuMP_s1
table_diff

end

# ╔═╡ e071b690-3c7d-11ec-0911-bd86ef5668aa
# We write here the same negative log-likelihood function as before written using a loop over all observations. This is slower but could be interesting to see

function loglike_dummy_ind(v,y,self_x,other_x,self_y,other_y,indicators_x,indicators_y)
	beta  = v[1:(length(v)-1)] # vector of parameters alpha,beta,gamma,delta
    sigma = exp(v[length(v)])  # sigma. exp to force it positive during search
	len = length(dt1.self_x)
	probs = Vector{Float64}(undef, len)
	for j=1:len
    lli = (indicators_x[j,1:4] .* beta)  # We get the vector [αs,βr,γq,δv]
    rli = (indicators_y[j,1:4] .* beta)
    uleft  = (sigma*((1-lli[1]-lli[2]-lli[3]-lli[4])*self_x[j]+ (lli[1]+lli[2]+lli[3]+lli[4]) *other_x[j])) # Utility when choosing left
    uright = (sigma*((1-rli[1]-rli[2]-rli[3]-rli[4])*self_y[j]+ (rli[1]+rli[2]+rli[3]+rli[4]) *other_y[j])) # Utility when choosing right
    probability = log((exp(uleft)/(exp(uleft)+exp(uright)))^y[j] * (exp(uright)/(exp(uleft)+exp(uright)))^(1-y[j])) # probability. See equation (4) in the paper
	probs[j]=probability
	end
    nsumloglike = -sum((probs)) # negative log likelihood
    return nsumloglike
end

# ╔═╡ Cell order:
# ╠═7b7526d0-3c9d-11ec-3575-37a66c520386
# ╠═7b36490e-3c9d-11ec-1ea6-dbf0082362dc
# ╠═c7fb5b79-7d87-49c5-9459-a360e1057715
# ╠═7d517760-4913-11ec-24fb-03122e4dda9a
# ╠═a0736b60-3c7d-11ec-1b71-27fd38cbf6e4
# ╠═e0965590-3c7d-11ec-13b6-a7d15db3fd9f
# ╠═d8928282-3c85-11ec-00ca-c5215c731b33
# ╠═56a26400-422c-11ec-3599-77a4d3d0e180
# ╠═dfcceb10-3c7d-11ec-2664-9d36301add75
# ╠═a3f76eb0-3d51-11ec-20e7-0d245610572b
# ╠═1baa58c0-4494-11ec-1794-6b4d74f6812b
# ╠═f576544e-3d55-11ec-009e-5ddcc502108b
# ╠═f558e140-3d55-11ec-075d-9d31f531f525
# ╠═f970b080-3cb2-11ec-0e09-65acfacc9837
# ╠═9e5d90a0-3cb2-11ec-2efc-8189c078aee7
# ╠═e6ab8b50-3cb2-11ec-0b39-e5424b4fdbd2
# ╠═d554ea40-3cb2-11ec-1642-99c3aadfd3e4
# ╠═586df880-3cd7-11ec-348c-6133b9ddba73
# ╠═c32cedfe-daf7-4b0b-81bc-e8028a050294
# ╠═802c89c0-3d60-11ec-1d69-0dcf94880bf3
# ╠═4ee3a2da-1a91-450b-97d8-4108fa08944e
# ╠═bbe59950-4915-11ec-289a-c97076f93ca1
# ╠═e071b690-3c7d-11ec-0911-bd86ef5668aa


text:
# This file was generated, do not modify it. # hide
dcr = @load DecisionTreeRegressor pkg=DecisionTree

dcrm = machine(dcr, X, y)

fit!(dcrm, rows=train)
pred_dcrm = MLJ.predict(dcrm, rows=test);


text:
using Merlin

type Token
    form::Int
    cat::Int
    head::Int
end

const nulltoken = Token(0, 0, 0)

function main()
    worddict = IntDict{String}()
    catdict = IntDict{String}()
    #unkform = worddict["UNKNOWN"]
    function encode(items::Vector{String})
        form0 = replace(items[2], r"[0-9]", '0')
        #formid = get(worddict, lowercase(form0), unkform)
        formid = push!(worddict, lowercase(form0))
        catid = push!(catdict, items[5])
        head = parse(Int, items[7])
        Token(formid, catid, head)
    end
    trainx = CoNLL.read(encode, ".data/wsj_02-21.conll")
    trainy = map(s -> map(x -> x.head, s), trainx)
    testx = CoNLL.read(encode, ".data/wsj_23.conll")
    testy = map(s -> map(x -> x.head, s), testx)

    trainstates = map(State, trainx)
    teststates = map(State, testx)

    function f()
        for i = 1:1
            for s in trainstates
                beamsearch(s, 16, nextpred, a->0.0)
            end
        end
    end
    @time f()
    #n = beamsearch(trainstates[1], 1, nextgold, a->0.0)
end

include("intdict.jl")
include("io.jl")
include("beamsearch.jl")
include("state.jl")

main()

max_stars_repo_name: hshindo/Merlin-Examples

max_issues_repo_name: hshindo/Merlin-Examples

max_forks_repo_name: hshindo/Merlin-Examples


text:
function Base.show(io::IO, tbl::AbstractTable)
    if column_indexable(tbl)
        show_ci_tbl(io, tbl)
    else
        @printf(io, "A %s", typof(tbl))
    end
end


text:
#= MULTIVECTOR TYPES

Type hierarchy:

           AbstractMultivector
            /               \
HomogeneousMultivector     MixedMultivector
   /        \
Blade   Multivector                 │
                                    │
           ╰─ CompositeMultivector ─╯

- `Blade`: a scalar multiple of a wedge product of orthogonal basis vectors.
   Note that the mathematical definition of a ``k``-blade is the wedge product
   of ``k`` different _vectors_, not necessarily basis vectors. Thus, not all
   ``k``-blades are representable as a `Blade` (but always as a `Multivector`).
- `Multivector`: a homogeneous multivector; a sum of same-grade blades.
- `MixedMultivector`: an inhomogeneous multivector. All elements in a geometric
   algebra can be represented as this type.

=#

"""
	AbstractMultivector{sig}

Supertype of all elements in the geometric algebra over the vector space with
metric signature `sig`, retrieved with the `signature` method.
"""
abstract type AbstractMultivector{sig} <: Number end

signature(::Type{<:AbstractMultivector{sig}}) where sig = sig
signature(::T) where {T<:AbstractMultivector} = signature(T)

dimension(::Union{T,Type{T}}) where {T<:AbstractMultivector} = dimension(signature(T))

Base.broadcastable(a::AbstractMultivector) = Ref(a)

"""
	HomogeneousMultivector{sig,k} <: AbstractMultivector{sig}

Supertype of grade `k` elements in the geometric algebra with metric signature `sig`.
"""
abstract type HomogeneousMultivector{sig,k} <: AbstractMultivector{sig} end

grade(::Type{<:HomogeneousMultivector{sig,k} where sig}) where k = k
grade(::T) where {T<:HomogeneousMultivector} = grade(T)

"""
	Blade{sig,k,bits,T} <: HomogeneousMultivector{sig,k}

A blade of grade `k` with basis blade `bits` and scalar coefficient of type `T`.

Parameters
===
- `sig`: metric signature defining the parent geometric algebra
- `k`: grade of the blade, equal to `grade(bits) === count_ones(bits)`
- `bits`: unsigned integer representing the indices of the basis vectors whose wedge product is the unit blade
- `T`: type of the scalar coefficient
"""
struct Blade{sig,k,bits,T} <: HomogeneousMultivector{sig,k}
	coeff::T
end
Blade{sig}(coeff::T, bits) where {sig,T} = Blade{sig,grade(bits),bits,T}(coeff)

# for internal use only (assumes `k == grade(bits)` without checking)
Blade{sig,k}(coeff::T, bits) where {sig,k,T} = Blade{sig,k,bits,T}(coeff)
Blade{sig,k,bits}(coeff::T) where {sig,k,bits,T} = Blade{sig,k,bits,T}(coeff)

"""
	bitsof(a::Blade{sig,k,bits,T}) = bits

The bits parameter of a blade, representing its basis `k`-blade.
"""
bitsof(::Type{<:Blade{sig,k,bits} where {sig,k}}) where bits = bits
bitsof(::T) where {T<:Blade} = bitsof(T)

"""
	Multivector{sig,k,S} <: HomogeneousMultivector{sig,k}

A homogeneous multivector of grade `k` with storage type `S`.

Parameters
===
- `sig`: metric signature defining the parent geometric algebra
- `k`: grade of the multivector
- `S`: type in which the multivector components are stored; usually a vector-like or dictionary-like type
"""
struct Multivector{sig,k,S} <: HomogeneousMultivector{sig,k}
	components::S
end
Multivector{sig,k}(comps::S) where {sig,k,S} = Multivector{sig,k,S}(comps)

"""
	MixedMultivector{sig,S} <: AbstractMultivector{sig}

An generally inhomogeneous multivector.

All elements of a geometric algebra are representable as a `MixedMultivector`.

Parameters
===
- `sig`: metric signature defining the parent geometric algebra
- `S`: type in which the multivector components are stored; usually a vector-like or dictionary-like type
"""
struct MixedMultivector{sig,S} <: AbstractMultivector{sig}
	components::S
end
MixedMultivector{sig}(comps::S) where {sig,S} = MixedMultivector{sig,S}(comps)

# const UnionMultivectorTypes = Union{Blade,Multivector,AbstractMultivector}
const CompositeMultivector{S} = Union{Multivector{sig,k,S},MixedMultivector{sig,S}} where {sig,k}
const Scalar = Union{filter(T -> !(T <: AbstractMultivector), subtypes(Number))...}

"""
	eltype(a)

The numerical type of the components of a multivector or multivector type.
"""
Base.eltype(::Type{<:Blade{sig,k,bits,T} where {sig,k,bits}}) where T = T
Base.eltype(::Type{<:CompositeMultivector{S}}) where S = valtype(S)

ncomponents(::Type{<:Multivector{sig,k}}) where {sig,k} = binomial(dimension(sig), k)
ncomponents(::Type{<:MixedMultivector{sig}}) where sig = 2^dimension(sig)
ncomponents(::T) where {T<:AbstractMultivector} = ncomponents(T)

#= ZERO & ONE CONSTRUCTORS =#

zeroslike(::Type{Vector{T}}, I...) where T = zeros(T, I...)
zeroslike(::Type{SparseVector{Tv,Ti}}, I...) where {Tv,Ti} = spzeros(Tv, Ti, I...)
zeroslike(::Type{S}, I...) where {S<:StaticVector} = zeros(S)

Base.zero(::Type{<:Blade{sig,k,bits,T}}) where {sig,k,bits,T} = Blade{sig,k,bits,T}(zero(T))
Base.zero(::Type{<:Multivector{sig,k,S}}) where {sig,k,S} = Multivector{sig,k}(zeroslike(S, binomial(dimension(sig), k)))
Base.zero(::Type{<:MixedMultivector{sig,S}}) where {sig,S} = MixedMultivector{sig}(zeroslike(S, 2^dimension(sig)))

# needed to override default implementation for `<:Number`
Base.zero(::T) where {T<:AbstractMultivector} = zero(T)

Base.iszero(a::Blade) = iszero(a.coeff)
Base.iszero(a::CompositeMultivector{<:AbstractVector}) = iszero(a.components)

Base.one(a::Type{<:Blade}) = Blade{signature(a)}(one(eltype(a)), bits_scalar())

Base.oneunit(a::Type{<:Blade}) = Blade{signature(a)}(one(eltype(a)), bitsof(a))
Base.oneunit(::T) where {T<:AbstractMultivector} = oneunit(T)

#= MULTIVECTOR TYPE INFERENCE UTILITIES

The `<:AbstractMultivector` type system is complex -- any multivector type consists of four orthogonal aspects:

 - `multivectortype`: one of `Blade, Multivector, MixedMultivector`
 - `signature`: an all-bits type parameter defining the geometric algebra
 - `storagetype`: the container type `S` which stores components of a `<:CompositeMultivector{S}` (not applicable for `Blade`)
 - `eltype`: the numerical type of the components

The direction of promotion for these aspects are, respectively:

 - `Blade` -> `Multivector` -> `MixedMultivector`
 -  no promotion occurs -- an error is thrown if types of different signature are compared
 - `StaticVector` -> `Vector` -> `SparseVector` - the storage type shouldn't be affected by `Blade` types
 -  usual promotion behaviour of `<:Number` types

The master function `best_type` performs `<:AbstractMultivector` type promotion, and exposes ways of
controlling various type parameters (e.g., setting the grade or requiring a minimum `multivectortype` or `eltype`).
=#

multivectortype(::Type{<:Blade}) = Blade
multivectortype(::Type{<:Multivector}) = Multivector
multivectortype(::Type{<:MixedMultivector}) = MixedMultivector
multivectortype(::T) where {T<:AbstractMultivector} = multivectortype(T)

shared_sig(::Type{<:AbstractMultivector{sig}}...) where sig = sig
shared_sig(::Type{<:AbstractMultivector}...) = error("multivectors must share the same metric signature")
shared_sig(As::AbstractMultivector...) = shared_sig(typeof.(As)...)

# harness built in type promotion system to promote storage types with more semantic code
struct StorageType{S} end
unwrap(::Type{StorageType{S}}) where S = S

storagetype(::Type{<:Blade}) = Union{} # indented to be type promoted with other storage types
storagetype(::Type{<:CompositeMultivector{S}}) where S = StorageType{S}
storagetype(::T) where {T<:AbstractMultivector} = storagetype(T)

Base.promote_rule(::Type{StorageType{Vector{T}}}, ::Type{StorageType{Vector{S}}}) where {T,S} = StorageType{Vector{promote_type(T, S)}}
Base.promote_rule(::Type{StorageType{SparseVector{Tv,Ti}}}, ::Type{StorageType{Vector{T}}}) where {Tv,Ti,T} = StorageType{SparseVector{promote_type(Tv, T),Ti}}
Base.promote_rule(::Type{StorageType{SparseVector{Tv,Ti}}}, ::Type{StorageType{SparseVector{Sv,Si}}}) where {Tv,Ti,Sv,Si} = StorageType{SparseVector{promote_type(Tv, Sv),promote_type(Ti, Si)}}
# fallback
Base.promote_rule(::Type{StorageType{T}}, ::Type{StorageType{S}}) where {T,S} = StorageType{promote_type(T, S)}

set_eltype_parameter(::Type{<:Vector}, ::Type{T}) where {T} = Vector{T}
set_eltype_parameter(::Type{<:SparseVector{Tv,Ti} where Tv}, ::Type{T}) where {T,Ti} = SparseVector{T,Ti}
set_eltype_parameter(::Type{<:SVector}, ::Type{T}) where {T} = SVector{N,T} where N
set_eltype_parameter(::Type{<:MVector}, ::Type{T}) where {T} = MVector{N,T} where N

set_size_parameter(::Type{T}, ::Val{N}) where {T,N} = T
set_size_parameter(::Type{<:SVector{N′,T} where N′}, ::Val{N}) where {N,T} = SVector{N,T}
set_size_parameter(::Type{<:MVector{N′,T} where N′}, ::Val{N}) where {N,T} = MVector{N,T}

# TODO: eventually `default_storagetype` should chose types appropriately by taking into account
# the algebra's dimension for optimal memory use/speed
default_storagetype(sig, Tv, Ti) = Vector{Tv}

# NOTE: best_type must work with non-concrete types in order for promotion to work with more than two objects.
# E.g., promote_type(x, x*y, x) == promote_type(promote_type(x, x*y), x)
#                                               \__ non-concrete __/
# Parameters which are allowed to be unspecified are {k,bits}. If any of {sig,S} are unspecified, do not use best_type to promote.

unwrap(::Type{Type{T}}) where T = T
unwrap(::Type{T}) where T = T
# unwrap(::T) where T = T

# must work where each `a...` is a multivector instance *or* a multivector type
@generated function best_type_parameters(a...; set_eltype::Type{SetT}=Nothing, promote_eltype_with::Type{PromT}=Union{}) where {SetT,PromT}
	a = unwrap.(a) # normalize Type{Type{A}} -> Type{A}
	sig = shared_sig(a...)
	T = SetT == Nothing ? promote_type(eltype.(a)..., PromT) : SetT
	S = promote_type(storagetype.(a)...)
	if S === Union{}
		S = default_storagetype(sig, T, UInt)
	end
	S = unwrap(S)
	if S === StorageType error("storagetype promotion failed") end
	S = set_eltype_parameter(S, T)
	sig, T, S
end

# best type if fastest when returning concrete types (as opposed to `UnionAll`s)
function best_type(::Type{Blade}, a...; bits::Val{b}=Val(missing), kwargs...) where b
	sig, T = best_type_parameters(a...; kwargs...)
	ismissing(b) ? Blade{sig,k,bits,T} where {k,bits} : Blade{sig,grade(b),b,T}
end
function best_type(::Type{Multivector}, a...; grade::Val{k}=Val(missing), kwargs...) where k
	sig, T, S = best_type_parameters(a...; kwargs...)
	if ismissing(k)
		Multivector{sig,k,S} where k
	else
		N = binomial(dimension(sig), k)
		Multivector{sig,k,set_size_parameter(S, Val(N))}
	end
end
function best_type(::Type{MixedMultivector}, a...; kwargs...)
	sig, T, S = best_type_parameters(a...; kwargs...)
	S = set_size_parameter(S, Val(2^dimension(sig)))
	MixedMultivector{sig,S}
end

# single argument form is used to modify parameters (e.g., to change the eltype)
best_type(a::Type{<:Blade}; kwargs...) = best_type(Blade, a; bits=Val(bitsof(a)), kwargs...)
best_type(a::Type{<:Multivector{sig,k} where sig}; kwargs...) where k = best_type(Multivector, a; grade=Val(k), kwargs...)
best_type(a::Type{<:Multivector}; kwargs...) = best_type(Multivector, a; kwargs...)
best_type(a::Type{<:MixedMultivector}; kwargs...) = best_type(MixedMultivector, a; kwargs...)
best_type(a::AbstractMultivector; kwargs...) = best_type(typeof(a); kwargs...)

#= CONVERSION

Conversion to more general types is possible:
 Blade ⊂ Multivector ⊂ MixedMultivector

The grade parameter of the target type is *not* respected. For subtypes of `HomogeneousMultivector{sig,k}`,
conversion to a type of a differing grade will still produce an object of grade `k`.
=#

# conversion of element/storage type
Base.convert(::Type{<:Blade{sig,k,bits,T} where {k,bits}}, a::Blade) where {sig,T} = Blade{sig,grade(a),bitsof(a),T}(a.coeff)
Base.convert(T::Type{<:Multivector}, a::Multivector) = best_type(T; grade=Val(grade(a)))(a.components)
Base.convert(::Type{<:MixedMultivector{sig,C}}, a::MixedMultivector) where {sig,C} = MixedMultivector{sig,C}(a.components)

# get (2^dim)-element vector of components for every basis blade in the algebra
full_components_vector(a::MixedMultivector{sig,<:AbstractVector}) where sig = a.components
function full_components_vector(a::HomogeneousMultivector)
	fcv = zeros(eltype(a), 2^dimension(a)) # TODO: use sparse vector? fcv is large and will remain mostly empty
	for b ∈ blades(a)
		fcv[begin + bitsof(b)] = b.coeff
	end
	fcv
end

# conversion from lower multivector type
function Base.convert(T::Type{<:Multivector}, a::Blade)
	T = best_type(T; grade=Val(grade(a)))
	i = bits_to_key(T, bitsof(a))
	comps = [i == j ? convert(eltype(T), a.coeff) : zero(eltype(T)) for j ∈ 1:ncomponents(T)]
	T(comps)
end
Base.convert(T::Type{<:MixedMultivector}, a::AbstractMultivector) = T(full_components_vector(a))

# conversion from scalar
Base.convert(::Type{<:Blade{sig,k,bits,T} where {k,bits}}, a::Scalar) where {sig,T} = Blade{sig,0,bits_scalar(),T}(a)
Base.convert(T::Type{<:Multivector}, a::Scalar) = best_type(T; grade=Val(0))([a])
Base.convert(T::Type{<:MixedMultivector}, a::Scalar) = zero(T) + a

#= PROMOTION

Promotion should result in objects of identical types *except* for the parameters {k,bits}, which should not be changed.
Promotion should only occur between objects of the same signature -- otherwise, fall back on default promotion behaviour.

TODO: is it actually appropriate to define these promote rules? maybe it's more Julian to not do implicit conversions like this...
=#

# same multivector type
Base.promote_rule(T::Type{<:Blade{sig}}, S::Type{<:Blade{sig}}) where sig = best_type(Blade, T, S)
Base.promote_rule(T::Type{<:Multivector{sig}}, S::Type{<:Multivector{sig}}) where sig = best_type(Multivector, T, S)
Base.promote_rule(T::Type{<:MixedMultivector{sig}}, S::Type{<:MixedMultivector{sig}}) where sig = best_type(MixedMultivector, T, S)

# promotion to higher multivector type
Base.promote_rule(T::Type{<:Multivector{sig}}, S::Type{<:Blade{sig}}) where sig = best_type(Multivector, T, S)
Base.promote_rule(T::Type{<:MixedMultivector{sig}}, S::Type{<:AbstractMultivector{sig}}) where sig = best_type(MixedMultivector, T, S)

# promotion from scalars
# note that the grade/bits parameters should *not* be preserved
Base.promote_rule(T::Type{<:AbstractMultivector{sig}}, S::Type{<:Scalar}) where sig = best_type(multivectortype(T), T; promote_eltype_with=S)

#= CONVENIENCE CONSTRUCTORS AND CONVERTERS =#

blade_like(a::AbstractMultivector, coeff=1, bits::Unsigned=bits_scalar()) = Blade{signature(a),grade(bits),bits,eltype(a)}(coeff)

Multivector(a::Blade) = convert(best_type(Multivector, a; grade=Val(grade(a))), a)
MixedMultivector(a::HomogeneousMultivector) = convert(best_type(MixedMultivector, a), a)

"""
	mapcomponents(f, a::AbstractMultivector; kwargs...)

Apply a mapping `f :: Blade -> Scalar` component-wise to `a`.
The function `f` is given a `Blade` representing each component,
and should return a scalar (not a blade).

The same keyword arguments as [`best_type`](@ref) are accepted,
including `set_eltype` and `promote_eltype_with` which can be
used to specify the eltype of the resultant multivector.
"""
mapcomponents(f, a::Blade; kwargs...) = best_type(a; kwargs...)(f(a)) # best_type preserves {k,bits}
function mapcomponents(f, a::CompositeMultivector; kwargs...)
	a′ = zero(best_type(a; kwargs...))
	for b ∈ blades(a)
		a′.components[bits_to_key(a, bitsof(b))] = f(b)
	end
	a′
end

Base.float(T::Type{<:AbstractMultivector}) = best_type(T, set_eltype=float(eltype(T)))
Base.big(T::Type{<:AbstractMultivector}) = best_type(T, set_eltype=big(eltype(T)))
Base.complex(T::Type{<:AbstractMultivector}) = best_type(T, set_eltype=complex(eltype(T)))
Base.real(T::Type{<:AbstractMultivector}) = best_type(T, set_eltype=real(eltype(T)))

Base.float(a::AbstractMultivector) = convert(float(typeof(a)), a)
Base.big(a::AbstractMultivector) = convert(big(typeof(a)), a)
Base.complex(a::AbstractMultivector) = convert(complex(typeof(a)), a)

Base.real(a::AbstractMultivector) = mapcomponents(b -> real(b.coeff), a; set_eltype=real(eltype(a)))
Base.imag(a::AbstractMultivector) = mapcomponents(b -> imag(b.coeff), a; set_eltype=real(eltype(a)))

#= EQUALITY

Only equality between multivectors of the *almost the same type* is defined,
leveraging the promotion interface to enable comparison between differing types.

Multivector promotion (intentionally) does *not* unify the grade and bits parameters,
so the following methods must cover types which are identical except for those parameters.

Note that two blades of differing basis or two multivectors of differing grade are
treated as equal if they are both zero.
=#

==(a::Blade{sig}, b::Blade{sig}) where sig = bitsof(a) == bitsof(b) ? a.coeff == b.coeff : iszero(a) && iszero(b)
==(a::(Multivector{sig,k,S} where k), b::(Multivector{sig,k,S} where k)) where {sig,S} = grade(a) == grade(b) ? a.components == b.components : iszero(a) && iszero(b)
==(a::MixedMultivector{sig,S}, b::MixedMultivector{sig,S}) where {sig,S} = a.components == b.components

==(a::AbstractMultivector, b::Scalar) = iszero(b) && iszero(a) || isscalar(a) && scalar(a) == b
==(a::Scalar, b::AbstractMultivector) = iszero(a) && iszero(b) || isscalar(b) && a == scalar(b)

==(a::AbstractMultivector{sig}...) where sig = ==(promote(a...)...)
==(a::AbstractMultivector...) = false # multivectors with non-identical signatures are teated as non-equal

function Base.isapprox(a::Blade{sig}, b::Blade{sig}; kwargs...) where sig
	if bitsof(a) == bitsof(b)
		isapprox(a.coeff, b.coeff; kwargs...)
	else
		isapprox(a.coeff, zero(a.coeff); kwargs...) && isapprox(b.coeff, zero(a.coeff); kwargs...)
	end
end
Base.isapprox(a::T, b::T; kwargs...) where {T<:MixedMultivector} = isapprox(a.components, b.components; kwargs...)
function Base.isapprox(a::Multivector{sig}, b::Multivector{sig}; kwargs...) where sig
	if grade(a) == grade(b)
		isapprox(a.components, b.components; kwargs...)
	else
		# multivectors of different grade are approximately equal is they are both approximately zero
		isapprox(a, zero(a); kwargs...) && isapprox(b, zero(b); kwargs...)
	end
end
Base.isapprox(a::AbstractMultivector, b::AbstractMultivector; kwargs...) = isapprox(promote(a, b)...; kwargs...)

Base.isapprox(a::AbstractMultivector, b::Scalar; kwargs...) = isapprox(promote(a, b)...; kwargs...)
Base.isapprox(a::Scalar, b::AbstractMultivector; kwargs...) = isapprox(promote(a, b)...; kwargs...)

#= COMPONENT ACCESS =#

blades(a::Blade) = (a,)
blades(a::Multivector{sig,k,<:AbstractVector}) where {sig,k} = (Blade{sig,k}(coeff, bits) for (coeff, bits) ∈ zip(a.components, bits_of_grade(k)))
blades(a::MixedMultivector{sig,<:AbstractVector}) where sig = (Blade{sig}(coeff, unsigned(i - 1)) for (i, coeff) ∈ enumerate(a.components))
# blades(a::CompositeMultivector{<:AbstractDict}) where sig = (Blade{sig}(coeff, bits) for (bits, coeff) ∈ a.components)

"""
	getcomponent(a::AbstractMultivector, I...)
	a[I...]

Return the component of the multivector `a` corresponding to the basis blade
defined by `I`.

Examples
===
```jldoctest
julia> v = basis((1, 1, 1));

julia> a = 1 + 2v[1] + 3v[2]v[3];

julia> a[], a[1], a[2], a[2,3]
(1, 2, 0, 3)

julia> GeometricAlgebra.getcomponent(a, 3, 2)
-3
```
"""
getcomponent

bits_to_key(::Type{<:Multivector{sig,k,<:AbstractVector}}, bits::Unsigned) where {sig,k} = bits_to_linear_index(bits)
bits_to_key(::Type{<:MixedMultivector{sig,<:AbstractVector}}, bits::Unsigned) where sig = 1 + bits
# bits_to_key(::Type{<:CompositeMultivector{<:AbstractDict}}, bits::Unsigned) = bits
bits_to_key(::T, bits) where {T<:AbstractMultivector} = bits_to_key(T, bits)

parity_sign(I) = iseven(parity(sortperm(collect(I)))) ? +1 : -1

function getcomponent(a::Blade, I...)
	indices_to_bits(I) == bitsof(a) || return zero(eltype(a))
	s = parity_sign(I)
	s*a.coeff
end

getcomponent(a::Multivector, I...) = zero(eltype(a))
function getcomponent(a::Multivector{sig,k}, I::Vararg{Any,k}) where {sig,k}
	s = parity_sign(I)
	s*a.components[bits_to_key(a, indices_to_bits(I))]
end

function getcomponent(a::MixedMultivector, I...)
	s = parity_sign(I)
	s*a.components[bits_to_key(a, indices_to_bits(I))]
end

getcomponent(a::MixedMultivector) = first(a.components)

for T ∈ (Integer, Symbol)
	@eval Base.getindex(a::AbstractMultivector{sig}, I::$T...) where sig = getcomponent(a, normalize_bv_index.(Ref(sig), I)...)
end
Base.getindex(a::AbstractMultivector) = getcomponent(a) # for ambiguity

# experimental notations for grade selection
# Base.getindex(a::AbstractMultivector, I::Vector{<:Integer}) = grade(a, Iterators.only(I))
# Base.getindex(a::AbstractMultivector; grade) = GeometricAlgebra2.grade(a, grade)


text:
# Factor Analysis

"""Factor Analysis type"""
immutable FactorAnalysis{T<:AbstractFloat}
    mean::Vector{T}       # sample mean: of length d (mean can be empty, which indicates zero mean)
    W::Matrix{T}          # factor loadings matrix: of size d x p
    Ψ::Vector{T}          # noise covariance: diagonal of size d x d
end

indim(M::FactorAnalysis) = size(M.W, 1)
outdim(M::FactorAnalysis) = size(M.W, 2)

Base.mean(M::FactorAnalysis) = fullmean(indim(M), M.mean)
projection(M::FactorAnalysis) = svdfact(M.W)[:U] # recover principle components from the weight matrix
Base.cov(M::FactorAnalysis) = M.W*M.W'+ Diagonal(M.Ψ)
Base.var(M::FactorAnalysis) = M.Ψ
loadings(M::FactorAnalysis) = M.W

## use

function transform{T<:AbstractFloat}(m::FactorAnalysis{T}, x::AbstractVecOrMat{T})
    xn = centralize(x, mean(m))
    W = m.W
    WᵀΨ⁻¹ = W'*diagm(1./m.Ψ)  # (q x d) * (d x d) = (q x d)
    return inv(I+WᵀΨ⁻¹*W)*(WᵀΨ⁻¹*xn)  # (q x q) * (q x d) * (d x 1) = (q x 1)
end

function reconstruct{T<:AbstractFloat}(m::FactorAnalysis{T}, z::AbstractVecOrMat{T})
    W  = m.W
    # ΣW(W'W)⁻¹z+μ = ΣW(W'W)⁻¹W'Σ⁻¹(x-μ)+μ = Σ(WW⁻¹)((W')⁻¹W')Σ⁻¹(x-μ)+μ = ΣΣ⁻¹(x-μ)+μ = (x-μ)+μ = x
    return cov(m)*W*inv(W'W)*z .+ mean(m)
end

## show

function Base.show(io::IO, M::FactorAnalysis)
    print(io, "Factor Analysis(indim = $(indim(M)), outdim = $(outdim(M)))")
end

## core algorithms

""" Expectation-maximization (EM) algorithms for maximum likelihood factor analysis.

    Rubin, Donald B., and Dorothy T. Thayer. "EM algorithms for ML factor analysis." Psychometrika 47.1 (1982): 69-76.
"""
function faem{T<:AbstractFloat}(S::DenseMatrix{T}, mv::Vector{T}, n::Int;
             maxoutdim::Int=size(X,1)-1,
             tol::Real=1.0e-6,   # convergence tolerance
             tot::Integer=1000)  # maximum number of iterations

    d = size(S,1)
    q = maxoutdim
    W = eye(T,d,q)
    Ψ = fill(T(0.01),d)

    L_old = 0.
    for c in 1:tot
        # EM-steps
        Ψ⁻¹W = diagm(1./Ψ)*W
        SΨ⁻¹W = S*Ψ⁻¹W
        H = SΨ⁻¹W*inv(I + W'*Ψ⁻¹W)
        W⁺ = SΨ⁻¹W*inv(I + H'*Ψ⁻¹W)
        Ψ⁺ = diag(S - H*(W⁺)')
        # new parameters
        W = W⁺
        Ψ = Ψ⁺
        # log likelihood
        Ψ⁻¹ = diagm(1./Ψ)
        WᵀΨ⁻¹ = W'*Ψ⁻¹
        detΣ = prod(Ψ)*det(I + WᵀΨ⁻¹*W)
        Σ⁻¹ = Ψ⁻¹ - Ψ⁻¹*W*inv(I + WᵀΨ⁻¹*W)*WᵀΨ⁻¹
        L = (-n/2)*(d*log(2π) + log(detΣ) + trace(Σ⁻¹*S))
        # println("$c] ΔL: $(abs(L_old - L)), L: $L")
        if abs(L_old - L) < tol
            break
        end
        L_old = L
    end

    return FactorAnalysis(mv, W, Ψ)
end

""" Fast conditional maximization (CM) algorithm for factor analysis.

    Zhao, J-H., Philip LH Yu, and Qibao Jiang. "ML estimation for factor analysis: EM or non-EM?." Statistics and computing 18.2 (2008): 109-123.
"""
function facm{T<:AbstractFloat}(S::DenseMatrix{T}, mv::Vector{T}, n::Int;
             maxoutdim::Int=size(X,1)-1,
             tol::Real=1.0e-6,   # convergence tolerance
             tot::Integer=1000,  # maximum number of iterations
             η = tol)            # variance low bound
    d = size(S,1)

    q = maxoutdim
    W = eye(T,d,q)
    Ψ = fill(T(0.01),d)
    V = zeros(T,q)
    eᵢeᵢ = zeros(T,d,d)
    Bᵢ⁻¹ = zeros(T,d,d)
    addconst = d*log(2π)

    L_old = 0.
    for c in 1:tot
        # CM-step 1
        Ψ⁻ʰ = Diagonal(1./sqrt.(Ψ))
        S⁺ = Ψ⁻ʰ*S*Ψ⁻ʰ

        F = eigfact(S⁺)
        λ = real(F[:values])
        ord = sortperm(λ, rev=true)
        λ = λ[ord]

        q′ = λ[q] > 0 ? q : findlast(λ .> 0)

        λq = 0.0
        @inbounds for i in 1:q′
            λq += log(λ[i]) - λ[i] + 1.
            V[i] = sqrt(λ[i] - 1.)
        end

        L = (addconst + log(prod(Ψ)) + trace(S⁺) + λq)*(-n/2)

        U = convert(Matrix{T}, F[:vectors][:,ord[1:q′]])
        W = U*diagm(V[1:q′]) # set new parameter

        # CM-step 2
        @inbounds for i in 1:q′
            V[i] = 1./λ[i] - 1.
        end
        eBe⁻¹ = ωᵢᵗ⁺¹ = 0.0
        for i in 1:d
            if i == 1
                Bᵢ⁻¹ = U*diagm(V[1:q′])*U' + I
            else
                eᵢeᵢ[i-1,i-1] = 1.
                Bᵢ⁻¹ = Bᵢ⁻¹ -  ωᵢᵗ⁺¹*Bᵢ⁻¹*eᵢeᵢ*Bᵢ⁻¹ / (1. + ωᵢᵗ⁺¹ * Bᵢ⁻¹[i-1,i-1])
                eᵢeᵢ[i-1,i-1] = 0.
            end
            eBe⁻¹ = 1./Bᵢ⁻¹[i,i]
            ωᵢᵗ⁺¹ = (Bᵢ⁻¹*S⁺*Bᵢ⁻¹)[i,i] * eBe⁻¹ * eBe⁻¹ - eBe⁻¹
            Ψ[i] = max(η, (ωᵢᵗ⁺¹ + 1.)*Ψ[i]) # set new parameter
        end

        # println("$c] ΔL: $(abs(L_old - L)), L: $L")
        if abs(L_old - L) < tol
            break
        end
        L_old = L
    end

    return FactorAnalysis(mv, Diagonal(sqrt.(Ψ))*W, Ψ)
end

## interface functions
function fit{T<:AbstractFloat}(::Type{FactorAnalysis}, X::DenseMatrix{T};
             method::Symbol=:cm,
             maxoutdim::Int=size(X,1)-1,
             mean=nothing,
             tol::Real=1.0e-6,   # convergence tolerance
             tot::Integer=1000,  # maximum number of iterations
             η = tol)            # variance low bound

    d, n = size(X)

    # process mean
    mv = preprocess_mean(X, mean)
    S = Base.covm(X, isempty(mv) ? 0 : mv, 2)
    if method == :em
        M = faem(S, mv, n, maxoutdim=maxoutdim, tol=tol, tot=tot)
    elseif method == :cm
        M = facm(S, mv, n, maxoutdim=maxoutdim, tol=tol, tot=tot, η = η)
    else
        throw(ArgumentError("Invalid method name $(method)"))
    end

    return M::FactorAnalysis
end


text:
"""
    link!(p => c, l = brlength(c))

Connect a `Pair` of nodes: From `p`arent to `c`hild with a branch of length `ℓ`.
"""
function link!(pc::Pair{RNode{T},RNode{T}}, l=brlength(pc[2])) where T
    p, c = pc

    push!(p.children, c)
    c.parent = p
    brlength!(c, l)

    return nothing
end

"""
    cut!(c)

Disconnect a node `c` from its parent.
"""
function cut!(c::RNode)
    p = parent(c)
    ic = findfirst(==(c), p.children)

    deleteat!(p.children, ic)

    c.parent = nothing

    return nothing
end

"""
    swap!(c1, c2)

Exchange the parents of nodes `c1` and `c2`
"""
function swap!(c1::RNode, c2::RNode)
    p1 = parent(c1)
    p2 = parent(c2)
    ic1 = findfirst(==(c1), p1.children)
    ic2 = findfirst(==(c2), p2.children)

    p1.children[ic1] = c2
    p2.children[ic2] = c1

    c1.parent = p2
    c2.parent = p1

    return nothing
end

"""
    graft!(x, c, at=0.5)

Insert node `x` between node `c` and its parent.

            graft!(x, c, 0.25)

p --------> c   =========>  p ------> x --> c
    1.0                        0.75    0.25

"""

max_stars_repo_name: eascarrunz/RootedTrees

max_issues_repo_name: eascarrunz/RootedTrees

max_forks_repo_name: eascarrunz/RootedTrees


text:
export bandpass, bandpass!, bandstop, bandstop!, lowpass, lowpass!,
        highpass, highpass!, taper, taper!, envelope

"""
    bandpass!(C,freqmin,freqmax,fs,corners=4,zerophase=false)

Butterworth-Bandpass Filter.

Filter data in `C` from `freqmin` to `freqmax` using `corners` corners.

# Arguments
- `C::SeisChannel`: SeisChannel to filter.
- `freqmin::Float64`: Pass band low corner frequency.
- `freqmax::Float64`: Pass band high corner frequency.
- `fs::Float64`: Sampling rate in Hz.
- `fs::Int`: Filter corners / order.
- `zerophase::Bool`: If True, apply filter once forwards and once backwards.
This results in twice the filter order but zero phase shift in
the resulting filtered trace.
"""
function bandpass!(C::SeisChannel, freqmin::Float64, freqmax::Float64;
                   corners::Int=4, zerophase::Bool=false)
    bandpass!(C.x,freqmin,freqmax,C.fs,corners=corners,zerophase=zerophase)
    return nothing
end
bandpass(C::SeisChannel,freqmin::Float64, freqmax::Float64; corners::Int=4,
         zerophase::Bool=false) = (U = deepcopy(C);bandpass!(U,freqmin,freqmax,
                                  corners=corners,zerophase=zerophase);return U)

"""
  bandpass!(S,freqmin,freqmax,fs,corners=4,zerophase=false)

Butterworth-Bandpass Filter.

Filter channels in `S` from `freqmin` to `freqmax` using `corners` corners.

# Arguments
- `S::SeisData`: SeisData to filter.
- `freqmin::Float64`: Pass band low corner frequency.
- `freqmax::Float64`: Pass band high corner frequency.
- `fs::Float64`: Sampling rate in Hz.
- `fs::Int`: Filter corners / order.
- `zerophase::Bool`: If True, apply filter once forwards and once backwards.
This results in twice the filter order but zero phase shift in
the resulting filtered trace.
"""
function bandpass!(S::SeisData, freqmin::Float64, freqmax::Float64;
                   corners::Int=4, zerophase::Bool=false)
    @inbounds for i = 1:S.n
        bandpass!(S[i].x,freqmin,freqmax,S[i].fs,corners=corners,
                             zerophase=zerophase)
    end
    return nothing
end
bandpass(S::SeisData,freqmin::Float64, freqmax::Float64; corners::Int=4,
         zerophase::Bool=false) = (U = deepcopy(S);bandpass!(U,freqmin,freqmax,
                                  corners=corners,zerophase=zerophase);return U)

"""
    bandstop!(C,freqmin,freqmax,fs,corners=4,zerophase=false)

Butterworth-Bandstop Filter.

Filter data in `C` removing data between frequencies `freqmin` to `freqmax` using
`corners` corners.

# Arguments
- `C::SeisChannel`: SeisChannel to filter.
- `freqmin::Float64`: Stop band low corner frequency.
- `freqmax::Float64`: Stop band high corner frequency.
- `fs::Float64`: Sampling rate in Hz.
- `fs::Int`: Filter corners / order.
- `zerophase::Bool`: If True, apply filter once forwards and once backwards.
This results in twice the filter order but zero phase shift in
the resulting filtered trace.
"""
function bandstop!(C::SeisChannel, freqmin::Float64, freqmax::Float64;
                   corners::Int=4, zerophase::Bool=false)
    bandstop!(C.x,freqmin,freqmax,C.fs,corners=corners,zerophase=zerophase)
    return nothing
end
bandstop(C::SeisChannel,freqmin::Float64, freqmax::Float64; corners::Int=4,
         zerophase::Bool=false) = (U = deepcopy(C);bandstop!(U,freqmin,freqmax,
                                  corners=corners,zerophase=zerophase);return U)

function bandstop!(S::SeisData, freqmin::Float64, freqmax::Float64;
                   corners::Int=4, zerophase::Bool=false)
    @inbounds for i = 1:S.n
        bandstop!(S[i].x,freqmin,freqmax,s[i].fs,corners=corners,zerophase=zerophase)
    end
    return nothing
end

"""
    bandstop!(S,freqmin,freqmax,fs,corners=4,zerophase=false)

Butterworth-Bandstop Filter.

Filter channels in `S` removing data between frequencies `freqmin` to `freqmax` using
`corners` corners.

# Arguments
- `S::SeisData`: SeisData to filter.
- `freqmin::Float64`: Stop band low corner frequency.
- `freqmax::Float64`: Stop band high corner frequency.
- `fs::Float64`: Sampling rate in Hz.
- `fs::Int`: Filter corners / order.
- `zerophase::Bool`: If True, apply filter once forwards and once backwards.
This results in twice the filter order but zero phase shift in
the resulting filtered trace.
"""
bandstop(S::SeisData,freqmin::Float64, freqmax::Float64; corners::Int=4,
         zerophase::Bool=false) = (U = deepcopy(S);bandstop!(U,freqmin,freqmax,
                                  corners=corners,zerophase=zerophase);return U)

"""
    lowpass(C,freq,fs,corners=4,zerophase=false)

Butterworth-Lowpass Filter.

Filter data in `C` over certain frequency `freq` using `corners` corners.

# Arguments
- `C::SeisChannel`: SeisChannel to filter.
- `freq::Float64`: Filter corner frequency.
- `fs::Float64`: Sampling rate in Hz.
- `fs::Int`: Filter corners / order.
- `zerophase::Bool`: If True, apply filter once forwards and once backwards.
This results in twice the filter order but zero phase shift in
the resulting filtered trace.
"""
function lowpass!(C::SeisChannel, freq::Float64;corners::Int=4, zerophase::Bool=false)
    lowpass!(C.x,freq,C.fs,corners=corners,zerophase=zerophase)
    return nothing
end
lowpass(C::SeisChannel,freq::Float64; corners::Int=4,zerophase::Bool=false) =
       (U = deepcopy(C);lowpass!(U,freq,corners=corners,zerophase=zerophase);
        return U)

"""
    lowpass(S,freq,fs,corners=4,zerophase=false)

Butterworth-Lowpass Filter.

Filter channels in `S` over certain frequency `freq` using `corners` corners.

# Arguments
- `S::SeisData`: SeisData to filter.
- `freq::Float64`: Filter corner frequency.
- `fs::Float64`: Sampling rate in Hz.
- `fs::Int`: Filter corners / order.
- `zerophase::Bool`: If True, apply filter once forwards and once backwards.
This results in twice the filter order but zero phase shift in
the resulting filtered trace.
"""
function lowpass!(S::SeisData, freq::Float64;corners::Int=4, zerophase::Bool=false)
    @inbounds for i = 1:S.n
        lowpass!(S[i].x,freq,S[i].fs,corners=corners,
                            zerophase=zerophase)
    end
    return nothing
end

lowpass(S::SeisData,freq::Float64; corners::Int=4,zerophase::Bool=false) =
       (U = deepcopy(S);lowpass!(U,freq,corners=corners,zerophase=zerophase);
       return U)

"""
    highpass(C,freq,fs,corners=4,zerophase=false)

Butterworth-Highpass Filter.

Filter data in `C` removing data below certain frequency `freq` using `corners` corners.

# Arguments
- `C::SeisChannel`: SeisChannel to filter.
- `freq::Float64`: Filter corner frequency.
- `fs::Float64`: Sampling rate in Hz.
- `fs::Int`: Filter corners / order.
- `zerophase::Bool`: If True, apply filter once forwards and once backwards.
This results in twice the filter order but zero phase shift in
the resulting filtered trace.
"""
function highpass!(C::SeisChannel, freq::Float64;corners::Int=4, zerophase::Bool=false)
    highpass!(C.x,freq,C.fs,corners=corners,zerophase=zerophase)
    return nothing
end
highpass(C::SeisChannel,freq::Float64; corners::Int=4,zerophase::Bool=false) =
       (U = deepcopy(C);highpass!(U,freq,corners=corners,zerophase=zerophase);
        return U)

"""
    highpass(S,freq,fs,corners=4,zerophase=false)

Butterworth-Highpass Filter.

Filter channels in `S` removing data below certain frequency `freq` using `corners` corners.

# Arguments
- `S::SeisData`: SeisData to filter.
- `freq::Float64`: Filter corner frequency.
- `fs::Float64`: Sampling rate in Hz.
- `fs::Int`: Filter corners / order.
- `zerophase::Bool`: If True, apply filter once forwards and once backwards.
This results in twice the filter order but zero phase shift in
the resulting filtered trace.
"""
function highpass!(S::SeisData, freq::Float64;corners::Int=4, zerophase::Bool=false)
    @inbounds for i = 1:S.n
        highpass!(S[i].x,freq,S[i].fs,corners=corners,
                             zerophase=zerophase)
    end
    return nothing
end
highpass(S::SeisData,freq::Float64; corners::Int=4,zerophase::Bool=false) =
       (U = deepcopy(S);highpass!(U,freq,corners=corners,zerophase=zerophase);
       return U)

"""
    envelope(A)

Envelope of a function.

Computes the upper and lower envelopes of the given function.

# Arguments
- `A::AbstractArray`: Data to make envelope of.
"""
function envelope(A::AbstractArray)
    Amean = mean(A)
    Acentered = A .- Amean
    env = abs.(hilbert(Acentered))
    upper = env .+ Amean
    lower = -env .+ Amean
    return upper, lower
end


text:
function solve_steady_state(m::DTCSCCModel, mc::ModelCalibration=m.calibration)
    p, s0, x0 = mc[:parameters, :states, :controls]
    e_ = zeros(mc[:shocks])
    ns = length(s0)
    nx = length(x0)

    function obj!(sx, out)
        s = sub(sx, 1:ns)
        x = sub(sx, ns+1:ns+nx)
        s_out = sub(out, 1:ns)
        x_out = sub(out, ns+1:ns+nx)

        # update state part of residual
        transition!(m, s_out, s, x, e_, p)
        broadcast!(-, s_out, s_out, s)

        # now update control part
        arbitrage!(m, x_out, s, x, e_, s, x, p)
        out
    end

    sol = nlsolve(obj!, vcat(s0, x0))
    !converged(sol) &&  error("Nonlinear solver failed to find steady state")

    # otherwise set controls
    out = deepcopy(mc)
    out[:states] = sol.zero[1:ns]
    out[:controls] = sol.zero[ns+1:end]
    out
end

# define custom exception types that hold some data for us to catch/explore
struct GeneralizedEigenvaluesError <: Exception
    msg::AbstractString
    diag_S::Vector{Float64}
    diag_T::Vector{Float64}
end

struct BlanchardKahnError <: Exception
    msg::AbstractString
    n_found::Int
    n_expected::Int
end

function _check_gev(diag_S, diag_T, tol)
    ok = sum((abs(diag_S) .< tol) .* (abs(diag_T) .< tol)) == 0
    if !ok
        msg = "Eigenvalues are not uniquely defined."
        throw(GeneralizedEigenvaluesError(msg, diag_S, diag_T))
    end
    true
end

function _check_bk_conditions(eigval, tol, n_expected)
    n_big = sum(eigval .> tol)
    if n_expected != n_big
        msg = """
        There are $(n_big) eigenvalues greater than one.
        There should be exactly $(n_expected) to meet Blanchard-Kahn conditions.
        """
        throw(BlanchardKahnError(msg, n_big, n_expected))
    end
    true
end

function linear_solve(m::AbstractDTCSCC, calib::ModelCalibration=m.calibration;
                      verbose::Bool=false, eigtol::Float64=1.0+1e-6)
    p, s, x, e = calib[:parameters, :states, :controls, :shocks]
    ns = length(s)
    nx = length(x)
    ne = length(e)
    nv = ns + nx

    # evaluate Jacobians and stack into A,B matrix for qz
    g_s, g_x, g_e = eval_jacobians(m, transition!, ns, (s, x, e), p)
    f_s, f_x, f_e, f_S, f_X = eval_jacobians(m, arbitrage!, nx,
                                             (s, x, e, s, x), p)

    A = [eye(ns) zeros(ns, nx)
         -f_S     -f_X]
    B = [g_s g_x
         f_s f_x]

    # do orderd QZ decomposition
    gs = schurfact(A, B)
    select = (abs(gs[:alpha]) .> eigtol*abs(gs[:beta]))
    ordschur!(gs, select)
    S, T, Q, Z = gs[:S], gs[:T], gs[:Q], gs[:Z]
    diag_S = diag(S)
    diag_T = diag(T)
    eigval = abs(diag_S./diag_T)

    # check eigen values and BK conditions
    _check_gev(diag_S, diag_T, 1e-10)
    _check_bk_conditions(eigval, eigtol, nx)

    # now obtain first order solution
    Z11 = Z[1:ns, 1:ns]
    Z12 = Z[1:ns, ns+1:end]
    Z21 = Z[ns+1:end, 1:ns]
    Z22 = Z[ns+1:end, ns+1:end]
    S11 = S[1:ns, 1:ns]
    T11 = T[1:ns, 1:ns]

    # first order solution
    C = (Z11'\Z21')'
    # P = (S11'\Z11')'*(Z11'\T11')'
    # Q = g_e
    # A = g_s + g_x*C
    # B = g_e

    dr = TaylorExpansion(s, x, C)
    # dr.A = A
    # dr.B = B
    # dr.sigma = get(m.distribution, :Normal, zeros(ne, ne))
end


text:
#=
This file is auto-generated. Do not edit.
=#
"""
    mutable struct AVRSimple <: AVR
        Kv::Float64
        V_ref::Float64
        ext::Dict{String, Any}
        states::Vector{Symbol}
        n_states::Int64
        internal::InfrastructureSystemsInternal
    end

Parameters of a simple proportional AVR in the derivative of EMF
i.e. an integrator controller on EMF

# Arguments
- `Kv::Float64`: Proportional Gain, validation range: (0, nothing)
- `V_ref::Float64`: Reference Voltage Set-point, validation range: (0, nothing)
- `ext::Dict{String, Any}`
- `states::Vector{Symbol}`
- `n_states::Int64`: Fixed AVR has 1 states
- `internal::InfrastructureSystemsInternal`: power system internal reference, do not modify
"""
mutable struct AVRSimple <: AVR
    "Proportional Gain"
    Kv::Float64
    "Reference Voltage Set-point"
    V_ref::Float64
    ext::Dict{String, Any}
    states::Vector{Symbol}
    "Fixed AVR has 1 states"
    n_states::Int64
    "power system internal reference, do not modify"
    internal::InfrastructureSystemsInternal
end

function AVRSimple(Kv, V_ref=1.0, ext=Dict{String, Any}(), )
    AVRSimple(Kv, V_ref, ext, [:Vf], 1, InfrastructureSystemsInternal(), )
end

function AVRSimple(; Kv, V_ref=1.0, ext=Dict{String, Any}(), )
    AVRSimple(Kv, V_ref, ext, )
end

# Constructor for demo purposes; non-functional.
function AVRSimple(::Nothing)
    AVRSimple(;
        Kv=0,
        V_ref=0,
        ext=Dict{String, Any}(),
    )
end

"""Get AVRSimple Kv."""
get_Kv(value::AVRSimple) = value.Kv
"""Get AVRSimple V_ref."""
get_V_ref(value::AVRSimple) = value.V_ref
"""Get AVRSimple ext."""
get_ext(value::AVRSimple) = value.ext
"""Get AVRSimple states."""
get_states(value::AVRSimple) = value.states
"""Get AVRSimple n_states."""
get_n_states(value::AVRSimple) = value.n_states
"""Get AVRSimple internal."""
get_internal(value::AVRSimple) = value.internal

"""Set AVRSimple Kv."""
set_Kv!(value::AVRSimple, val::Float64) = value.Kv = val
"""Set AVRSimple V_ref."""
set_V_ref!(value::AVRSimple, val::Float64) = value.V_ref = val
"""Set AVRSimple ext."""
set_ext!(value::AVRSimple, val::Dict{String, Any}) = value.ext = val
"""Set AVRSimple states."""
set_states!(value::AVRSimple, val::Vector{Symbol}) = value.states = val
"""Set AVRSimple n_states."""
set_n_states!(value::AVRSimple, val::Int64) = value.n_states = val
"""Set AVRSimple internal."""
set_internal!(value::AVRSimple, val::InfrastructureSystemsInternal) = value.internal = val


text:
include("../../common/sparse/src/sparse_module.jl")
include("../../../std/util/src/util_module.jl")

module linalg
using ..sparse
using ..util
using LinearAlgebra

include("norm.jl")
include("triangular.jl")
include("ldl.jl")
include("lu.jl")
include("tridiagonal.jl")

end # module

max_stars_repo_name: HomoModelicus/julia

max_issues_repo_name: HomoModelicus/julia

max_forks_repo_name: HomoModelicus/julia


text:
using OrdinaryDiffEq, CUDA, LinearAlgebra, Test
function f(u,p,t)
    A*u
end
function f(du,u,p,t)
    mul!(du,A,u)
end
function jac(J,u,p,t)
    J .= A
end
function jac(u,p,t)
    A
end
function tgrad(du,u,p,t)
    du .= 0
end
function tgrad(u,p,t)
    zero(u)
end
ff = ODEFunction(f,jac=jac,tgrad=tgrad)
CUDA.allowscalar(false)
A = cu(-rand(3,3))
u0 = cu([1.0;0.0;0.0])
tspan = (0.,100.)

prob = ODEProblem(ff,u0,tspan)
sol = solve(prob,Tsit5())
@test solve(prob,Rosenbrock23()).retcode == :Success
solve(prob,Rosenbrock23(autodiff=false));

prob_oop = ODEProblem{false}(ff,u0,tspan)
CUDA.allowscalar(false)
sol = solve(prob_oop,Tsit5())
@test solve(prob_oop,Rosenbrock23()).retcode == :Success
@test solve(prob_oop,Rosenbrock23(autodiff=false)).retcode == :Success

prob_nojac = ODEProblem(f,u0,tspan)
@test solve(prob_nojac,Rosenbrock23()).retcode == :Success
@test solve(prob_nojac,Rosenbrock23(autodiff=false)).retcode == :Success
@test solve(prob_nojac,Rosenbrock23(autodiff=false,diff_type = Val{:central})).retcode == :Success
@test solve(prob_nojac,Rosenbrock23(autodiff=false,diff_type = Val{:complex})).retcode == :Success

#=
prob_nojac_oop = ODEProblem{false}(f,u0,tspan)
DiffEqBase.prob2dtmin(prob_nojac_oop)
@test_broken solve(prob_nojac_oop,Rosenbrock23()).retcode == :Success
@test_broken solve(prob_nojac_oop,Rosenbrock23(autodiff=false)).retcode == :Success
@test_broken solve(prob_nojac_oop,Rosenbrock23(autodiff=false,diff_type = Val{:central})).retcode == :Success
@test_broken solve(prob_nojac_oop,Rosenbrock23(autodiff=false,diff_type = Val{:complex})).retcode == :Success
=#

# Complex Numbers Adaptivity DifferentialEquations.jl#460
f_complex(u,nothing,t) = 1/2 .*u
u0 = cu(rand(32,32).+ 1im*rand(32,32));
prob = ODEProblem(f_complex,u0,(0.0,1.0))
@test_nowarn sol = solve(prob,Tsit5())


text:
function solve_lp(mapf::MAPF; T::Integer, integer=false, capacity=true)
    @assert all(<=(T), mapf.starting_times)

    g = mapf.graph
    w_vec = [mapf.edge_weights[src(ed), dst(ed)] for ed in edges(g)]

    A = nb_agents(mapf)
    V, E = nv(g), ne(g)

    model = Model(SCIP.Optimizer)

    @variable(model, x[1:A, 1:(T - 1), 1:E])
    @variable(model, y[1:A, 1:T, 1:V])

    @objective(model, Min, dot(w_vec, sum(x; dims=(1, 2))))

    if integer
        set_binary.(x)
        set_binary.(y)
    else
        @constraint(model, 0 .<= x .<= 1)
        @constraint(model, 0 .<= y .<= 1)
    end

    for a in 1:A
        s = mapf.sources[a]
        d = mapf.destinations[a]
        t0 = mapf.starting_times[a]
        for t = 1:t0-1
            @constraint(model, y[a, t, :] .== 0)
            @constraint(model, x[a, t, :] .== 0)
        end
        for t in t0:T
            @constraint(model, sum(@view y[a, t, :]) == 1)
        end
        @constraint(model, y[a, t0, s] == 1)
        @constraint(model, y[a, T, d] == 1)
    end

    for v in 1:V
        e_in = [e for (e, ed) in enumerate(edges(g)) if dst(ed) == v]
        e_out = [e for (e, ed) in enumerate(edges(g)) if src(ed) == v]
        for a in 1:A
            t0 = mapf.starting_times[a]
            for t in t0:(T - 1)
                @constraint(model, y[a, t, v] == sum(@view x[a, t, e_out]))
            end
            for t in (t0 + 1):T
                @constraint(model, y[a, t, v] == sum(@view x[a, t - 1, e_in]))
            end
        end
    end

    if capacity
        for group in mapf.conflict_groups
            for t in 1:T
                @constraint(model, sum(@view y[:, t, group]) <= 1)
            end
        end
    end

    optimize!(model)

    stat = termination_status(model)
    solution = Solution()

    if stat == MOI.OPTIMAL
        val = objective_value(model)
        yopt = value.(y)
        for a = 1:A
            path = Path()
            t0 = mapf.starting_times[a]
            for t = t0:T
                v = argmax(@view yopt[a, t, :])
                push!(path, (t, v))
                v == mapf.destinations[a] && break
            end
            push!(solution, path)
        end
    else
        val = Inf
    end

    return stat, val, solution
end


text:
module LogEnergies

using Spec
using ..Traits, ..RNG, ..Var, ..Tagging, ..Space
using ..Util: Box
using ..Basis: AbstractΩ
using ..SeenVars

export logenergy, ℓ, logenergyexo

## Simple Case
# In the simple case, the values of all the exogenous variables
# are within FIXME

# zt: fix me precision of these two are different,
logenergy(::StdNormal{T}, v) where T = - v^2/2 - log(sqrt(2π))
logenergy(::StdUniform{T}, v) where T = zero(T)
logenergy(::UniformInt{Int32}, v) = -22.18070977791824990135142788666165017841600429952816813186176030378859590303026
logenergy(::UniformInt{Int64}, v) = -44.36141955583649980270285577332330035683200859905633626372352060757719180606051
logenergy(::UniformInt{UInt64}, v) = -44.36141955583649980270285577332330035683200859905633626372352060757719180606051

logintsize(::Type{Int64}) = log(1) - log(BigInt(2)^(64))
logenergy(dist::PrimRandVar, v) = logenergy(dist.class, v)

"`logenergyexo(ω)` Log energy of `ω` only on exogenous variables"
function logenergyexo(ω)
  reduce(pairs(ω); init = 0.0) do logpdf_, (dist, val)
    logpdf_ + logenergy(dist, val)
  end
end

"""
`logenergy(rng::AbstractRNG, x, ω)`

# Returns
- joint log probability of
"""
function logenergy(x, ω::AbstractΩ)
  ω_ = tagseen(taglogenergy(ω))
  ret = x(ω_)
  ω_.tags.logenergy.val
  # ctxapply(x, ω)
end

# logenergy(::LogThenComplete, x, ω) = logenergy(complete(x, ω))

## Complex Case

@inline taglogenergy(ω, logenergy_ = 0.0) =
  tagseen(tag(ω, (logenergy = Box(logenergy_),)))

# """
# `logenergy(rng::AbstractRNG, x, ω)`

# # Returns
# - joint log probability of
# """
# function logenergy(x, ω::AbstractΩ)
#   ω_ = tagseen(taglogenergy(ω))
#   ret = x(ω_)
#   ω_.tags.logenergy.val
#   # ctxapply(x, ω)
# end

function Var.posthook(::trait(LogEnergy, Seen), ret, f::PrimRandVar, ω)
  if f ∉ ω.tags.seen
    # @show "SEEN"
    ω.tags.logenergy.val += logenergy(f, ret)
    push!(ω.tags.logenergy.seen, id)
  end
  nothing
end

const ℓ = logenergy

## Complete
function complete(x, ω)
  ω_ = tagcomplete(ω)
  ret = x(ω_)
  ω_.tags.complete.val
end
@post complete(x, ω) = (ω_ = __ret__; isvalid(x, ω_) & (ω_ ⊆ ω))

# function Var.posthook(::trait(Complete), ret, f::PrimRandVar, ω)
#   if f ∉ ω.tags.seen
#     # @show "SEEN"
#     ω.tags.logenergy.val += logenergy(f, ret)
#     push!(ω.tags.logenergy.seen, id)
#   end
#   nothing
# end

end

max_stars_repo_name: zenna/expect

max_issues_repo_name: zenna/expect

max_forks_repo_name: zenna/expect


text:
using DIVAnd, PyPlot

mask,(pm,pn),(xi,yi) = DIVAnd_squaredom(2,-1:0.01:1)

x        = [-0.5,0.5]; # x-positions
y        = [ 0. ,0. ]; # y-positions
f        = [ 1. ,1. ]; # values
epsilon2 = [ 0.5,1. ]; # expected error variances (normalized)

len = 0.2 # correlation length-scale
fi,s = DIVAndrun(mask,(pm,pn),(xi,yi),(x,y),f,len,epsilon2);

# Plotting
clf();
pcolor(xi,yi,fi; cmap="jet");
axis("equal")
colorbar();

max_stars_repo_name: gher-ulg/DIVAnd-presentation

max_issues_repo_name: gher-ulg/DIVAnd-presentation

max_forks_repo_name: gher-ulg/DIVAnd-presentation


text:
@testset "compute_bc.jl" begin

end


text:
"""
    TypeP(var_1::TypeO)

Lorem ipsum dolor sit amet.
"""
struct TypeP

  var_1::TypeO

end


text:
import OmegaCore.Proposal
using OmegaCore: PrimRandVar, Member, like, propose, StdNormal, StdUniform, UniformInt

using OmegaCore.Util: update

export SSProposal

"Single-site proposal: change a single variable in `ω`"
struct SSProposal end

# Normally distribuetd proposal around `val`
function subpropose(qω, ::Member{<:StdNormal}, val; σ = 0.01)
  (randn(qω) * σ) + val
end

# Indepdent re-sample
function subpropose(qω, ::Member{<:UniformInt{T}}, val; σ = 0.01) where {T}
  # @assert false
  rand(qω, T)
  # (randn(qω) + val) * σ
end

function subpropose(qω, ::Member{<:StdUniform}, val; σ = 0.1)
  @assert false "not implemented"
end

function Proposal.propose(qω, f, ω, ::SSProposal)
  # Uniformly choose over sites
  k = rand(qω, keys(ω))
  qv = subpropose(qω, k, ω[k])

  # What about keeping most things constant?
  # ω_ = like(ω, k => qv)
  ω_ = update(ω, k, qv)

  # zt: is this ignore really necessary
  f(OmegaCore.tagignorecondition(OmegaCore.tagrng(ω_, qω)))
  ω_
end

function Proposal.propose_and_logratio(qω, f, ω, ss::SSProposal)
  ω_ = propose(qω, f, ω, ss)
  (ω_, 0.0) # Symmetric proposal
end

defpropose_and_logratio(f) = (qω, ω) -> Proposal.propose_and_logratio(qω, f, ω,  SSProposal())

# filter(x, ωO)


text:
# Types to help with differentiation of Abstract Objectives
# Needs `AbstractMOP` and `AbstractObjective` to be defined.
# `TransformerFn` is also needed.
# Note: Differatiates the objective handle for scaled input.
struct GradWrapper{
        TT<:TransformerFn,
        FT<:Union{AbstractVector{<:Function},Nothing},
        JT<:Union{Function, Nothing}
    } <: DiffFn
    tfn :: TT
    list_of_grads :: FT
    jacobian_handle :: JT   # JT either Nothing or <:Function
end

GradWrapper(tfn :: TransformerFn, g :: Function, jh :: Union{Nothing, Function} ) = GradWrapper(tfn, [g,], jh );

# true =̂ gw.list_of_grads not empty
function _get_gradient( :: Val{true}, gw :: GradWrapper, x̂ :: Vec, ℓ :: Int = 1)
    Jtfn = _jacobian_unscaling( gw.tfn, x̂ )
    g = gw.list_of_grads[ℓ](gw.tfn(x̂))
    @assert length(g) == length(x̂) "Gradient $(ℓ) must have length $(length(x̂))."
    return ensure_vec(Jtfn'g);
end

# false =̂ list_of_grads empty but jacobian handle provided
function _get_gradient( :: Val{false}, gw :: GradWrapper, x̂ :: Vec, ℓ :: Int = 1)
    Jtfn = _jacobian_unscaling( gw.tfn, x̂ )
    return ensure_vec( Jtfn[:,ℓ]'gw.jacobian_handle( gw.tfn(x̂) ) )
end

function get_gradient( gw :: GradWrapper, x̂ :: Vec, ℓ :: Int = 1)
    return _get_gradient( Val(!isempty(gw.list_of_grads)), gw, x̂, ℓ );
end

# prefer user defined handle over gradient iteration
function get_jacobian( gw :: GradWrapper{TT,FT,JT}, x̂ :: Vec ) where{TT,FT,JT<:Function}
    Jtfn = _jacobian_unscaling( gw.tfn, x̂ )
    return Jtfn'gw.jacobian_handle( gw.tfn(x̂) )
end

function get_jacobian( gw :: GradWrapper{TT,FT,JT}, x̂ :: Vec ) where{TT,FT,JT<:Nothing}
    return transpose( hcat(
        ( _get_gradient( Val(true), gw, x̂, ℓ ) for ℓ = eachindex(gw.list_of_grads) )...
    ))
end

struct AutoDiffWrapper{
        O <: AbstractObjective,
        TT <: TransformerFn,
        JT <: Union{Nothing,Function}
    } <: DiffFn
    objf :: O
    tfn :: TT
    jacobian :: JT
end

# no jacobian handle set
function get_jacobian( adw :: AutoDiffWrapper{O,TT,JT}, x̂ :: Vec ) where{O,TT,JT<:Nothing}
    #return AD.jacobian( eval_handle(adw.objf) ∘ adw.tfn, x̂ );
    return AD.jacobian( _eval_handle(adw), x̂ )
end

# jacobian handle set
function get_jacobian( adw :: AutoDiffWrapper{O,TT,JT}, x̂ :: Vec ) where{O,TT,JT<:Function}
    Jtfn = _get_transformer_jacobian( adw, x̂ )
    return Jtfn'gw.jacobian_handle( adw.tfn(x̂) )
end

# jacobian handle set
function get_gradient( adw :: AutoDiffWrapper{O,TT,JT}, x̂ :: Vec, ℓ :: Int = 1 ) where{O,TT,JT<:Function}
    Jtfn = _get_transformer_jacobian( adw, x̂ )
    return ensure_vec(Jtfn[:,ℓ]'gw.jacobian_handle( adw.tfn(x̂) ))
end

# fallbacks, if no jacobian is set; not optimized for multiple outputs
function get_gradient( adw :: AutoDiffWrapper, x̂ :: Vec, ℓ :: Int = 1)
    #return AD.gradient( x -> eval_handle(adw.objf)( adw.tfn(x) )[ℓ], x̂ )
    return AD.gradient( _eval_handle(adw, ℓ), x̂ )
end

#=
# not optimized for multiple outputs
function get_hessian( adw :: AutoDiffWrapper, x̂ :: Vec, ℓ :: Int = 1)
    #return AD.hessian( x -> eval_handle(adw.objf)( adw.tfn(x) )[ℓ], x̂ )
    return AD.hessian( x -> eval_objf(adw.objf, adw.tfn, x )[ℓ], x̂ )
end
=#

struct FiniteDiffWrapper{
        O<:AbstractObjective,
        TT<:TransformerFn,
        JT<:Union{Function,Nothing}
    } <: DiffFn
    objf :: O
    tfn :: TT
    jacobian :: JT
end

function _eval_handle( fdw :: Union{AutoDiffWrapper,FiniteDiffWrapper}, ℓ = nothing )
    if isnothing(ℓ)
        return x -> eval_objf( fdw.objf, fdw.tfn, x )
    else
        return x -> eval_objf( fdw.objf, fdw.tfn, x )[ℓ]
    end
end

# no jacobian handle set: apply finite differencing to the composed function(s)
# F_ℓ ∘ unscale
function get_jacobian( fdw :: FiniteDiffWrapper{O,TT,JT}, x̂ :: Vec ) where{O,TT,JT<:Nothing}
    return FD.finite_difference_jacobian(_eval_handle(fdw), x̂ )
end

# jacobian handle set: use the jacobian handle and apply chain rule for unscaling transformation
function get_jacobian( fdw :: FiniteDiffWrapper{O,TT,JT}, x̂ :: Vec ) where{O,TT,JT<:Function}
    Jtfn = _jacobian_unscaling( fdw.tfn, x̂ )
    return Jtfn'gw.jacobian_handle( fdw.tfn(x̂) )
end

# gradient from jacobian if jacobian handle is set:
function get_gradient(fdw :: FiniteDiffWrapper{O,TT,JT}, x̂ :: Vec, ℓ :: Int = 1 ) where{O,TT,JT<:Function}
    Jtfn = _jacobian_unscaling( fdw.tfn, x̂ )
    return Jtfn[:,ℓ]'gw.jacobian_handle( fdw.tfn(x̂) )
end

function get_gradient( fdw :: FiniteDiffWrapper, x̂ :: Vec, ℓ :: Int = 1)
    return FD.finite_difference_gradient( _eval_handle(fdw, ℓ), x̂ );
end

#=
# NOTE this is clearly is not optimized for multiple outputs
# (and atm is not meant to be optimized)
# one could differentiate the gradients here and use the same evaluation sites
# for all hessians
function get_hessian( fdw :: FiniteDiffWrapper, x̂ :: Vec, ℓ :: Int = 1 )
    return FD.finite_difference_jacobian(x -> get_gradient( fdw, x, ℓ ), x̂ );
end
=#

###### Hessians ########################

struct HessWrapper{TT<:TransformerFn, FT<:AbstractVector{<:Function}} <: DiffFn
    tfn :: TT
    list_of_hess_fns :: FT
end

HessWrapper(tfn :: TransformerFn, hf :: Function ) = HessWrapper(tfn, [hf,] );

function get_hessian( hw :: HessWrapper, x̂ :: Vec, ℓ :: Int = 1)
    Jtfn = _jacobian_unscaling( hw.tfn, x̂ )
    return Jtfn'hw.list_of_hess_fns[ℓ](hw.tfn(x̂))*Jtfn;
end

struct HessFromGrads{GW} <: DiffFn
    gw :: GW
end

#=
# if gw is a vector of gradient callbacks
function _get_grad_func( gw :: AbstractVector{<:Function}, ℓ :: Int )
    return x -> gw[ℓ](x)
end

# if gw is jacobian callback
function _get_grad_func( gw :: Function, ℓ :: Int )
    return x -> ensure_vec( gw(x)[ℓ, :] )
end
=# # Use a GradWrapper instead

function _get_grad_func( gw, ℓ :: Int)
    return x -> get_gradient( gw, x, ℓ)
end

function get_hessian( hw :: HessFromGrads{<:AutoDiffWrapper}, x̂ :: Vec, ℓ :: Int = 1 )
    AD.jacobian( _get_grad_func( hw.gw, ℓ), x̂ )
end

function get_hessian( hw :: HessFromGrads{<:FiniteDiffWrapper}, x̂ :: Vec, ℓ :: Int = 1 )
    FD.finite_difference_jacobian( _get_grad_func( hw.gw, ℓ), x̂ )
end

#=
function get_hessian( hw :: HessFromGrads, x̂ :: Vec, ℓ :: Int = 1)
    if hw.method == :autodiff
        return AD.jacobian( _get_grad_func( hw.gw, ℓ), x̂ )
    elseif hw.method == :fdm
        return FD.finite_difference_jacobian( _get_grad_func( hw.gw, ℓ), x̂ )
    end
end
=#


text:
using LinearAlgebra
using Test

using SetProg, SetProg.Sets
using Polyhedra
using MultivariatePolynomials

import DynamicPolynomials
import MultivariateBases
const MB = MultivariateBases
const MonoBasis = MB.MonomialBasis{DynamicPolynomials.Monomial{true}, DynamicPolynomials.MonomialVector{true}}

using JuMP
const MOIT = MOI.Test

const □ = polyhedron(HalfSpace([1, 0], 1.0) ∩ HalfSpace([-1, 0], 1) ∩ HalfSpace([0, 1], 1) ∩ HalfSpace([0, -1], 1))
const ◇ = polyhedron(convexhull([1.0, 0], [0, 1], [-1, 0], [0, -1]), lib)

function square_test(optimizer, config::MOIT.Config,
                     inner::Bool, variable::SetProg.AbstractVariable,
                     metric::Function, objective_value, set_test)
    model = _model(optimizer)

    @variable(model, ◯, variable)
    if inner
        cref = @constraint(model, ◯ ⊆ □)
    else
        cref = @constraint(model, □ ⊆ ◯)
    end
    @objective(model, inner ? MOI.MAX_SENSE : MOI.MIN_SENSE,
               metric(volume(◯)))

    SetProg.optimize!(model)
    @test JuMP.termination_status(model) == MOI.OPTIMAL
    @test JuMP.objective_sense(model) == MOI.MAX_SENSE
    @test JuMP.objective_value(model) ≈ objective_value atol=config.atol rtol=config.rtol
    set_test(JuMP.value(◯))
end

function john_homogeneous_square_test(optimizer, config)
    square_test(optimizer, config, true,
                Ellipsoid(symmetric=true, dimension=2),
                nth_root, 1.0,
                ◯ -> begin
                    @test ◯ isa Sets.Polar{Float64, Sets.Ellipsoid{Float64}}
                    @test Sets.polar(◯).Q ≈ Symmetric([1.0 0.0; 0.0 1.0]) atol=config.atol rtol=config.rtol
                end)
end

function john_nonhomogeneous_ell_square_test(optimizer, config)
    square_test(optimizer, config, true,
                Ellipsoid(point=SetProg.InteriorPoint([0.0, 0.0])),
                nth_root, 1.0,
                ◯ -> begin
                    @test ◯ isa Sets.PerspectiveDual{Float64, Sets.Householder{Float64, Sets.ShiftedEllipsoid{Float64}, Float64}}
                    z = Sets.perspective_variable(◯)
                    x, y = Sets.space_variables(◯)
                    ◯_dual = Sets.perspective_dual(◯)
                    @test ◯_dual.p ≈ -z^2 + x^2 + y^2 atol=config.atol rtol=config.rtol
                    @test Sets._householder(◯_dual.h) ≈ [-1.0 0.0 0.0
                                                          0.0 1.0 0.0
                                                          0.0 0.0 1.0] atol=config.atol rtol=config.rtol
                    @test ◯_dual.set.Q ≈ Symmetric([1.0 0.0; 0.0 1.0]) atol=config.atol rtol=config.rtol
                    @test ◯_dual.set.b ≈ [0.0, 0.0] atol=config.atol rtol=config.rtol
                    @test ◯_dual.set.β ≈ -1.0 atol=config.atol rtol=config.rtol
                end)
end

function john_nonhomogeneous_quad_square_test(optimizer, config)
    square_test(optimizer, config, true,
                PolySet(degree=2, convex=true, point=SetProg.InteriorPoint([0.0, 0.0])),
                set -> L1_heuristic(set, [1.0, 1.0]),
                8/3,
                ◯ -> begin
                    @test ◯ isa Sets.PerspectiveDual{Float64, Sets.Householder{Float64, Sets.ConvexPolynomialSet{Float64, MonoBasis, Float64}, Float64}}
                    z = Sets.perspective_variable(◯)
                    x, y = Sets.space_variables(◯)
                    ◯_dual = Sets.perspective_dual(◯)
                    @test ◯_dual.p ≈ -z^2 + x^2 + y^2 atol=config.atol rtol=config.rtol
                end)
end

function löwner_homogeneous_square_test(optimizer, config)
    square_test(optimizer, config, false,
                Ellipsoid(symmetric=true, dimension=2),
                nth_root, 0.5,
                ◯ -> begin
                    @test ◯ isa Sets.Ellipsoid
                    @test ◯.Q ≈ Symmetric([0.5 0.0
                                           0.0 0.5]) atol=config.atol rtol=config.rtol
                end)
end

function piecewise_semiell_inner_homogeneous_◇_square_test(optimizer, config)
    square_test(optimizer, config, true,
                Ellipsoid(symmetric=true, piecewise=◇),
                set -> L1_heuristic(set), 4,
                p◯ -> begin
                    @test p◯ isa Sets.Polar
                    ◯ = p◯.set
                    @test ◯ isa Sets.Piecewise{Float64, Sets.Ellipsoid{Float64}}
                    @test length(◯.sets) == 4
                    Q1 = Symmetric([ 1.0  1.0
                                     1.0  1.0])
                    Q2 = Symmetric([ 1.0 -1.0
                                    -1.0  1.0])
                    @test ◯.sets[1].Q ≈ Q1 atol=config.atol rtol=config.rtol
                    @test ◯.sets[2].Q ≈ Q2 atol=config.atol rtol=config.rtol
                    @test ◯.sets[3].Q ≈ Q2 atol=config.atol rtol=config.rtol
                    @test ◯.sets[4].Q ≈ Q1 atol=config.atol rtol=config.rtol
                end)
end

function piecewise_semiell_inner_homogeneous_□_square_test(optimizer, config)
    square_test(optimizer, config, true,
                Ellipsoid(symmetric=true, piecewise=□),
                set -> L1_heuristic(set), 8/3,
                p◯ -> begin
                    @test p◯ isa Sets.Polar
                    ◯ = p◯.set
                    @test ◯ isa Sets.Piecewise{Float64, Sets.Ellipsoid{Float64}}
                    @test length(◯.sets) == 4
                    Q = Symmetric([1.0 0.0
                                   0.0 1.0])
                    @test ◯.sets[1].Q ≈ Q atol=config.atol rtol=config.rtol
                    @test ◯.sets[2].Q ≈ Q atol=config.atol rtol=config.rtol
                    @test ◯.sets[3].Q ≈ Q atol=config.atol rtol=config.rtol
                    @test ◯.sets[4].Q ≈ Q atol=config.atol rtol=config.rtol
                end)
end

const quartic_inner_poly = [3.1518541833100864, -0.1617384194869734]
const quartic_inner_obj = 6.447419478140056
const quartic_inner_α = 5.6567546886722795
const quartic_inner_convexity = [12.0, 0.0, quartic_inner_α, 0.0, quartic_inner_poly[1]+2quartic_inner_poly[2],
                                 quartic_inner_α, 8.48516455194103, 0.0, 0.0, 12.0]

function quartic_inner_homogeneous_square_test(optimizer, config)
    square_test(optimizer, config, true,
                PolySet(symmetric=true, degree=4, dimension=2, convex=true),
                nth_root, quartic_inner_obj,
                ◯ -> begin
                    @test ◯ isa Sets.Polar{Float64, Sets.ConvexPolySet{Float64, MonoBasis, Float64}}
                    ◯_polar = Sets.polar(◯)
                    @test ◯_polar.degree == 4
                    x, y = variables(◯_polar.p)
                    @test polynomial(◯_polar.p) ≈ x^4 + quartic_inner_convexity[5]*x^2*y^2 + y^4 atol=config.atol rtol=config.rtol
                    convexity_proof = Sets.convexity_proof(◯)
                    @test convexity_proof.n == 4
                    @test convexity_proof.Q ≈ quartic_inner_convexity atol=config.atol rtol=config.rtol
                end)
end

const quartic_outer_β = 0.30177574048813055
const quartic_outer_γ = 0.5936049698923986
const quartic_outer_λ = -0.09857757888257276
const quartic_outer_obj = 1.611854896946893
const quartic_outer_α = 0.7928996242545062
const quartic_outer_convexity = [3.621308885857567, 0.0, quartic_outer_α, 0.0, 0.08578956499151169,
                                 quartic_outer_α, 1.5, 0.0, 0.0, 3.6212933687704307]

function quartic_outer_homogeneous_square_test(optimizer, config)
    square_test(optimizer, config, false,
                PolySet(symmetric=true, degree=4, dimension=2, convex=true),
                nth_root,
                quartic_outer_obj,
                ◯ -> begin
                    @test ◯ isa Sets.ConvexPolySet{Float64, MonoBasis, Float64}
                    @test ◯.degree == 4
                    x, y = variables(◯.p)
                    @test polynomial(◯.p) ≈ quartic_outer_β*x^4 + (quartic_outer_γ+2quartic_outer_λ)*x^2*y^2 + quartic_outer_β*y^4 atol=config.atol rtol=config.rtol
                    convexity_proof = Sets.convexity_proof(◯)
                    @test convexity_proof.n == 4
                    @test convexity_proof.Q ≈ quartic_outer_convexity atol=config.atol rtol=config.rtol
                end)
end

const square_tests = Dict(
    "john_homogeneous_square" =>
     john_homogeneous_square_test,
    "john_nonhomogeneous_ell_square" =>
     john_nonhomogeneous_ell_square_test,
    "john_nonhomogeneous_quad_square" =>
     john_nonhomogeneous_quad_square_test,
    "löwner_homogeneous_square" =>
     löwner_homogeneous_square_test,
    "piecewise_semiell_inner_homogeneous_◇_square" =>
     piecewise_semiell_inner_homogeneous_◇_square_test,
    "piecewise_semiell_inner_homogeneous_□_square" =>
     piecewise_semiell_inner_homogeneous_□_square_test,
    "quartic_inner_homogeneous_square" =>
     quartic_inner_homogeneous_square_test,
    "quartic_outer_homogeneous_square" =>
     quartic_outer_homogeneous_square_test
)

@test_suite square


text:
module UserSitgnature
import MLStyle.AbstractPatterns: PComp
export and, or, literal, and, wildcard, decons, guard, effect, self

PatternUse = NamedTuple{(:and, :or, :literal, :wildcard, :decons, :guard, :effect)}

@nospecialize
_empty_ntuple = NamedTuple()
and(ps::Vector, config::NamedTuple = _empty_ntuple) = function apply(impl::PatternUse)
    xs = [p(impl) for p in ps]
    impl.and(xs, config)
end

or(ps::Vector, config::NamedTuple = _empty_ntuple) = function apply(impl::PatternUse)
    xs = [p(impl) for p in ps]
    impl.or(xs, config)
end
literal(val, config::NamedTuple = _empty_ntuple) = function apply(impl::PatternUse)
    impl.literal(val, config)
end
wildcard(config::NamedTuple = _empty_ntuple) = function apply(impl::PatternUse)
    impl.wildcard(config)
end

decons(comp::PComp, extract::Function, ps, config::NamedTuple = _empty_ntuple) =
    function apply(impl::PatternUse)
        xs = [p(impl) for p in ps]
        impl.decons(comp, extract, xs, config)
    end

guard(pred, config::NamedTuple = _empty_ntuple) = function apply(impl::PatternUse)
    impl.guard(pred, config)
end

effect(ctx_perf, config::NamedTuple = _empty_ntuple) = function apply(impl::PatternUse)
    impl.effect(ctx_perf, config)
end

const self = (
    and = and,
    or = or,
    literal = literal,
    wildcard = wildcard,
    decons = decons,
    guard = guard,
    effect = effect,
)

end
@specialize


text:
export on_call!
export combineLatest

# TODO: WIP
# TODO: Untested and undocumented

# Speficic method for map operator over SingleObservable
function on_call!(::Type{L}, ::Type{R}, operator::MapOperator{R}, source::SingleObservable{L}) where L where R
    return of(convert(R, operator.mappingFn(source.value)))
end

# Speficic method for filter operator over SingleObservable
function on_call!(::Type{L}, ::Type{L}, operator::FilterOperator, source::SingleObservable{L}) where L
    if operator.filterFn(source.value)
        return of(source.value)
    else
        return completed(L)
    end
end


text:
export asConst
struct AsConst{T} <: TV.VectorTransform
    value::T
end

asConst(x) = AsConst(x)

TV.as(c::Returns) = AsConst(c.value)

TV.dimension(t::AsConst) = 0

function TV.transform_with(flag::TV.NoLogJac, t::AsConst, x, index)
    return (t.value, TV.NoLogJac(), index)
end

function TV.transform_with(flag::TV.LogJac, t::AsConst, x, index)
    return (t.value, 0.0, index)
end

TV.inverse_eltype(t::AsConst, x::AbstractVector) = Any

function TV.inverse_at!(x::AbstractVector, index, t::AsConst, y::AbstractVector)
    return index
end