text:
module SequentialFit
using LsqFit, RecipesBase, Calculus, Optim, LinearAlgebra

export SeqFit, update!, next_x

import Base.push!

mutable struct SeqFit
    f
    model
    X
    Y
    fit
    lims
end
function SeqFit(f, model, X, p0, lims=(-Inf, Inf))
    Y = f.(X)
    return SeqFit(f, model, copy(X), Y, curve_fit(model, X, Y, p0), lims)
end
function SeqFit(f, model, p0)
    return SeqFit(f, model, [-1.0, 0.0, 1.0], p0, lims)
end

function push!(s::SeqFit, x...)
    push!(s.X, x...)
    push!(s.Y, s.f.(x)...)
    return s.fit = curve_fit(s.model, s.X, s.Y, s.fit.param)
end

function update!(s::SeqFit)
    x_new = next_x(s)
    return push!(s, x_new)
end

function next_x(s::SeqFit)
    isempty(s.X) && return zero(eltype(s.X))
    o = optimize(
        x -> selectionFunction(x, s),
        max(2*minimum(s.X)-maximum(s.X), s.lims[1]),
        min(2*maximum(s.X)-minimum(s.X), s.lims[2]),
    )
    return o.minimizer
end

function selectionFunction(x, s)
    return -norm(Calculus.gradient(p -> s.model(x, p), s.fit.param))^2 *
           prod(abs.(s.X .- x) .^ 2)
end

@recipe function f(s::SeqFit; groundtruth=false, selection=false)
    @series begin # measurements
        subplot := 1
        seriestype := :scatter
        label --> "Samples"
        s.X, s.Y
    end

    @series begin # fit
        subplot := 1
        seriestype := :line
        label --> "Fit"
        x -> s.model(x, s.fit.param)
    end

    if groundtruth
        @series begin # Ground truth
            subplot := 1
            seriestype := :line
            linestyle --> :dash
            label --> "Ground truth"
            s.f
        end
    end

    if selection
        layout := (2, 1)
        link := :x
        @series begin # Selection function
            subplot := 2
            seriestype := :line
            linestyle --> :dash
            label --> "Selection function"
            ylims := :auto
            yticks := false
            legend := :bottomright
            x -> selectionFunction(x, s)
        end
    end
end

end


text:
module CPLEXCP

# TODO:
# - bridge for Circuit:
#     https://www.ibm.com/support/knowledgecenter/SSSA5P_12.5.1/ilog.odms.studio.help/OPL_Studio/oplmigration/topics/opl_mig_prev_3x4x_3xCP_constr_circuit.html
#     https://lost-contact.mit.edu/afs/pdc.kth.se/roots/ilse/v0.7/pdc/vol/cplex/12.5/amd64_co5/doc/html/en-US/OPL_Studio/oplmigration/topics/opl_mig_prev_3x4x_3xCP_constr_circuit.html
# - bridge for Cardinality/Sequence:
#     https://www.ibm.com/support/knowledgecenter/SSSA5P_12.5.1/ilog.odms.studio.help/OPL_Studio/oplmigration/topics/opl_mig_prev_3x4x_3xCP_constr_seq.html
#     https://lost-contact.mit.edu/afs/pdc.kth.se/roots/ilse/v0.7/pdc/vol/cplex/12.5/amd64_co5/doc/html/en-US/OPL_Studio/oplmigration/topics/opl_mig_prev_3x4x_3xCP_constr_seq.html
#     (https://sofdem.github.io/gccat/gccat/Cglobal_cardinality.html)
#     (https://sofdem.github.io/gccat/gccat/Csliding_distribution.html)
# Help for debugging, with `model` an MOI model:
#   cpo_java_exportmodel(model.inner, "somewhere/model.cpo")
#   -> cpoptimizer, r to read this file, opt to solve, etc.

using JavaCall
import MathOptInterface
import ConstraintProgrammingExtensions

const MOI = MathOptInterface
const MOIU = MOI.Utilities
const CleverDicts = MOIU.CleverDicts
const CP = ConstraintProgrammingExtensions

# Check if the package has been built correctly.
if isfile(joinpath(dirname(@__FILE__), "..", "deps", "deps.jl"))
    include("../deps/deps.jl")
else
    error(
        "CPLEXCP not properly installed. Please run `Pkg.build(\"CPLEXCP\")` or `]build CPLEXCP`",
    )
end

if !@isdefined(libcplexcpojava)
    error(
        "CPLEXCP not properly built. There probably was a problem when running `Pkg.build(\"CPLEXCP\")` or `]build CPLEXCP`",
    )
end

# Initialise the package by setting the right parameters for Java. This assumes
# no other code uses JavaCall...
function __init__()
    if get(ENV, "JULIA_REGISTRYCI_AUTOMERGE", "") != "true"
        cpo_java_init()
    end
end

# Export the Java API.
export cpo_java_init,
    JavaCPOModel,
    IloCP,
    IloIntVar,
    IloIntervalVar,
    IloIntervalSequenceVar,
    IloNumVar,
    IloIntExpr,
    IloNumExpr,
    IloIntTupleSet,
    IloNumToNumSegmentFunction,
    IloNumToNumStepFunction,
    IloCumulFunctionExpr,
    IloTransitionDistance,
    IloStateFunction,
    IloConstraint,
    IloAlternative,
    IloIsomorphism,
    IloNoOverlap,
    IloRange,
    IloSpan,
    IloSynchronize,
    IloObjective,
    IloMultiCriterionExpr,
    IloSolution,
    IloAddable,
    Callback,
    ConflictStatus,
    Constraint,
    IntExpr,
    NumVar,
    NumExpr,
    Addable,
    Variable,
    ConstraintArray,
    IntExprArray,
    NumVarArray,
    NumExprArray,
    AddableArray,
    IloInfinity,
    IloMaxInt,
    IloMinInt,
    cpo_java_release,
    cpo_java_boolvar,
    cpo_java_boolvararray,
    cpo_java_intvar,
    cpo_java_intvararray,
    cpo_java_numvar,
    cpo_java_numvararray,
    cpo_java_intervalvar,
    cpo_java_intervalsequencevar,
    cpo_java_numvar_getlb,
    cpo_java_numvar_getub,
    cpo_java_numvar_setlb,
    cpo_java_numvar_setub,
    cpo_java_intvar_getlb,
    cpo_java_intvar_getub,
    cpo_java_intvar_setlb,
    cpo_java_intvar_setub,
    cpo_java_abs,
    cpo_java_constant,
    cpo_java_count,
    cpo_java_countdifferent,
    cpo_java_div,
    cpo_java_element,
    cpo_java_endof,
    cpo_java_enfofnext,
    cpo_java_enfofprevious,
    cpo_java_exponent,
    cpo_java_heightatend,
    cpo_java_heightatstart,
    cpo_java_intexpr,
    cpo_java_intexprarray,
    cpo_java_lengthof,
    cpo_java_lengthofnext,
    cpo_java_lengthofprevious,
    cpo_java_log,
    cpo_java_max,
    cpo_java_min,
    cpo_java_modulo,
    cpo_java_negative,
    cpo_java_numexprarray,
    cpo_java_numvararray,
    cpo_java_overlaplength,
    cpo_java_piecewiselinear,
    cpo_java_power,
    cpo_java_prod,
    cpo_java_pulse,
    cpo_java_quot,
    cpo_java_scalprod,
    cpo_java_sizeeval,
    cpo_java_sizeof,
    cpo_java_sizeofnext,
    cpo_java_sizeofprevious,
    cpo_java_square,
    cpo_java_standarddeviation,
    cpo_java_starteval,
    cpo_java_startof,
    cpo_java_startofnext,
    cpo_java_startofprevious,
    cpo_java_sum,
    cpo_java_typeofnext,
    cpo_java_typeofprevious,
    cpo_java_linearintexpr,
    cpo_java_linearnumexpr,
    cpo_java_linearintexpr_add,
    cpo_java_linearintexpr_addterm,
    cpo_java_linearintexpr_addterms,
    cpo_java_linearintexpr_clear,
    cpo_java_linearintexpr_getconstant,
    cpo_java_linearintexpr_remove,
    cpo_java_linearintexpr_setconstant,
    cpo_java_linearnumexpr_add,
    cpo_java_linearnumexpr_addterm,
    cpo_java_linearnumexpr_addterms,
    cpo_java_linearnumexpr_clear,
    cpo_java_linearnumexpr_getconstant,
    cpo_java_linearnumexpr_remove,
    cpo_java_linearnumexpr_setconstant,
    cpo_java_inttable,
    cpo_java_inttupleset_addtuple,
    cpo_java_inttupleset_getarity,
    cpo_java_numtonumsegmentfunction,
    cpo_java_piecewiselinearfunction,
    cpo_java_numtonumsegmentfunction_add,
    cpo_java_numtonumsegmentfunction_addvalue,
    cpo_java_numtonumsegmentfunction_copy,
    cpo_java_numtonumsegmentfunction_dilate,
    cpo_java_numtonumsegmentfunction_getarea,
    cpo_java_numtonumsegmentfunction_getdefinitionintervalmax,
    cpo_java_numtonumsegmentfunction_getdefinitionintervalmin,
    cpo_java_numtonumsegmentfunction_getmax,
    cpo_java_numtonumsegmentfunction_getmin,
    cpo_java_numtonumsegmentfunction_getvalue,
    cpo_java_numtonumsegmentfunction_prod,
    cpo_java_numtonumsegmentfunction_setmax,
    cpo_java_numtonumsegmentfunction_setmax,
    cpo_java_numtonumsegmentfunction_setmin,
    cpo_java_numtonumsegmentfunction_setperiodic,
    cpo_java_numtonumsegmentfunction_setperiodicvalue,
    cpo_java_numtonumsegmentfunction_setslope,
    cpo_java_numtonumsegmentfunction_setvalue,
    cpo_java_numtonumsegmentfunction_shift,
    cpo_java_numtonumsegmentfunction_sub,
    cpo_java_numtonumstepfunction,
    cpo_java_numtonumstepfunction_add,
    cpo_java_numtonumstepfunction_addvalue,
    cpo_java_numtonumstepfunction_copy,
    cpo_java_numtonumstepfunction_dilate,
    cpo_java_numtonumstepfunction_getarea,
    cpo_java_numtonumstepfunction_getdefinitionintervalmax,
    cpo_java_numtonumstepfunction_getdefinitionintervalmin,
    cpo_java_numtonumstepfunction_getmax,
    cpo_java_numtonumstepfunction_getmin,
    cpo_java_numtonumstepfunction_getvalue,
    cpo_java_numtonumstepfunction_prod,
    cpo_java_numtonumstepfunction_setmax,
    cpo_java_numtonumstepfunction_setmin,
    cpo_java_numtonumstepfunction_setperiodic,
    cpo_java_numtonumstepfunction_setperiodicvalue,
    cpo_java_numtonumstepfunction_setvalue,
    cpo_java_numtonumstepfunction_shift,
    cpo_java_numtonumstepfunction_sub,
    cpo_java_statefunction,
    cpo_java_cumulfunctionexpr,
    cpo_java_diff,
    cpo_java_getnumberofsegments,
    cpo_java_getsegmentstart,
    cpo_java_getsegmentend,
    cpo_java_getsegmentvalue,
    cpo_java_getvalue_cumulfunctionexpr,
    cpo_java_step,
    cpo_java_stepatend,
    cpo_java_stepatend,
    cpo_java_stepatstart,
    cpo_java_stepatstart,
    cpo_java_sum_cumulfunctionexpr,
    cpo_java_transitiondistance,
    cpo_java_transitiondistance_getsize,
    cpo_java_transitiondistance_getvalue,
    cpo_java_transitiondistance_setvalue,
    cpo_java_add,
    cpo_java_alldiff,
    cpo_java_allmindistance,
    cpo_java_allowedassignments,
    cpo_java_alternative,
    cpo_java_alternative,
    cpo_java_alwaysconstant,
    cpo_java_alwaysequal,
    cpo_java_alwaysin,
    cpo_java_alwaysnostate,
    cpo_java_and,
    cpo_java_before,
    cpo_java_distribute,
    cpo_java_elementm,
    cpo_java_endatend,
    cpo_java_endatstart,
    cpo_java_endbeforeend,
    cpo_java_endbeforestart,
    cpo_java_eq,
    cpo_java_equiv,
    cpo_java_falseconstraint,
    cpo_java_first,
    cpo_java_forbiddenassignments,
    cpo_java_forbidend,
    cpo_java_forbidextent,
    cpo_java_forbidstart,
    cpo_java_ge,
    cpo_java_gt,
    cpo_java_ifthen,
    cpo_java_ifthenelse,
    cpo_java_imply,
    cpo_java_inverse,
    cpo_java_isomorphism,
    cpo_java_last,
    cpo_java_le,
    cpo_java_lexicographic,
    cpo_java_lt,
    cpo_java_neq,
    cpo_java_not,
    cpo_java_nooverlap,
    cpo_java_or,
    cpo_java_pack,
    cpo_java_presenceof,
    cpo_java_previous,
    cpo_java_range,
    cpo_java_samecommonsubsequence,
    cpo_java_samesubsequence,
    cpo_java_sequence,
    cpo_java_span,
    cpo_java_startatend,
    cpo_java_startatstart,
    cpo_java_startbeforeend,
    cpo_java_startbeforestart,
    cpo_java_strong,
    cpo_java_subcircuit,
    cpo_java_synchronize,
    cpo_java_trueconstraint,
    cpo_java_and_add,
    cpo_java_and_remove,
    cpo_java_or_add,
    cpo_java_or_remove,
    cpo_java_range_clearexpr,
    cpo_java_range_getexpr,
    cpo_java_range_getlb,
    cpo_java_range_getub,
    cpo_java_range_setbounds,
    cpo_java_range_setexpr,
    cpo_java_range_setlb,
    cpo_java_range_setub,
    cpo_java_addable_getname,
    cpo_java_addable_setname,
    cpo_java_maximize,
    cpo_java_maximize_multicriterion,
    cpo_java_minimize,
    cpo_java_minimize_multicriterion,
    cpo_java_staticlex,
    cpo_java_addeq,
    cpo_java_addge,
    cpo_java_addle,
    cpo_java_addmaximize,
    cpo_java_addminimize,
    cpo_java_addrange,
    cpo_java_getallconstrainedilocumulfunctionexprs,
    cpo_java_getallilointervalsequencevars,
    cpo_java_getallintervalvars,
    cpo_java_getallintvars,
    cpo_java_getallstatefunctions,
    cpo_java_getallkpinames,
    cpo_java_getdomain,
    cpo_java_getdomainsize,
    cpo_java_getend,
    cpo_java_getendmax,
    cpo_java_getendmin,
    cpo_java_getfirst,
    cpo_java_getincumbentvalue,
    cpo_java_getintvalue,
    cpo_java_getlast,
    cpo_java_getlength,
    cpo_java_getlengthmax,
    cpo_java_getlengthmin,
    cpo_java_getmax,
    cpo_java_getmin,
    cpo_java_getname,
    cpo_java_getnext,
    cpo_java_getobjbound,
    cpo_java_getobjbounds,
    cpo_java_getobjgap,
    cpo_java_getobjgaps,
    cpo_java_getobjvalue,
    cpo_java_getobjvalues,
    cpo_java_getprev,
    cpo_java_getsize,
    cpo_java_getsizemax,
    cpo_java_getsizemin,
    cpo_java_getstart,
    cpo_java_getstartmax,
    cpo_java_getstartmin,
    cpo_java_getvalue,
    cpo_java_getversion,
    cpo_java_hasobjective,
    cpo_java_isabsent,
    cpo_java_isfixed,
    cpo_java_isindomain,
    cpo_java_ispresent,
    cpo_java_next,
    cpo_java_propagate,
    cpo_java_refineconflict,
    cpo_java_restore,
    cpo_java_solve,
    cpo_java_startnewsearch,
    cpo_java_store,
    cpo_java_solution,
    cpo_java_solution_add,
    cpo_java_solution_contains,
    cpo_java_solution_end,
    cpo_java_solution_getend,
    cpo_java_solution_getendmin,
    cpo_java_solution_getlength,
    cpo_java_solution_getlengthmax,
    cpo_java_solution_getlengthmin,
    cpo_java_solution_getmax,
    cpo_java_solution_getmin,
    cpo_java_solution_getsize,
    cpo_java_solution_getsizemax,
    cpo_java_solution_getsizemin,
    cpo_java_solution_getstart,
    cpo_java_solution_getstartmax,
    cpo_java_solution_getstartmin,
    cpo_java_solution_getvalue,
    cpo_java_solution_isabsent,
    cpo_java_solution_isfixed,
    cpo_java_solution_isindomain,
    cpo_java_solution_ispresent,
    cpo_java_solution_remove_intervalvar,
    cpo_java_solution_remove_intervalvararray,
    cpo_java_solution_remove_intvar,
    cpo_java_solution_remove_intvararray,
    cpo_java_solution_setabsent,
    cpo_java_solution_setdomain,
    cpo_java_solution_setend,
    cpo_java_solution_setendmax,
    cpo_java_solution_setendmin,
    cpo_java_solution_setlength,
    cpo_java_solution_setlengthmax,
    cpo_java_solution_setlengthmin,
    cpo_java_solution_setmax,
    cpo_java_solution_setmin,
    cpo_java_solution_setoptional,
    cpo_java_solution_setpresent,
    cpo_java_solution_setsize,
    cpo_java_solution_setsizemax,
    cpo_java_solution_setsizemin,
    cpo_java_solution_setstart,
    cpo_java_solution_setstartmax,
    cpo_java_solution_setstartmin,
    cpo_java_solution_setvalue,
    cpo_java_solution_store,
    cpo_java_dumpmodel,
    cpo_java_exportmodel,
    cpo_java_getbuildid,
    cpo_java_getconflict_constraint,
    cpo_java_getconflict_intervalval,
    cpo_java_getconflict_numvar,
    cpo_java_importmodel,
    cpo_java_printinformation,
    cpo_java_remove,
    cpo_java_removeallcallbacks,
    cpo_java_removeallkpis,
    cpo_java_removecallback,
    cpo_java_removekpi,
    cpo_java_runseeds,
    cpo_java_getintinfo,
    cpo_java_getdoubleinfo,
    cpo_java_getdoubleparameter,
    cpo_java_getdoubledefaultparameter,
    cpo_java_getintparameter,
    cpo_java_getintdefaultparameter,
    cpo_java_setdoubleparameter,
    cpo_java_setintparameter

# Finally, the code.
include("api_java.jl")
include("MOI/wrapper.jl")
include("MOI/parse.jl")
include("MOI/helpers.jl")
include("MOI/wrapper_variables.jl")
include("MOI/wrapper_objective.jl")
include("MOI/wrapper_constraints.jl")
include("MOI/wrapper_constraints_singlevar.jl")
include("MOI/wrapper_constraints_names.jl")
include("MOI/wrapper_constraints_mo.jl")
include("MOI/wrapper_constraints_cp.jl")
include("MOI/wrapper_constraints_cp_reification.jl")

# Some deprecations.
@deprecate cpo_java_model JavaCPOModel

end # module

ext: jl


text:
using Random
using Plots
using Printf

function getAcc(pos, mass, G)
	# positions r = [x,y] for all particles
	x = pos[:,1]
	y = pos[:,2]

	# all pairwise particle separations
	Δx = x' .- x
	Δy = y' .- y

	# 1/r^3 for all particle pairwise particle separations
	inv_r3 = (Δx.^2 + Δy.^2)
	inv_r3[inv_r3.>0] = inv_r3[inv_r3.>0].^(-1.5)

	ax = G .* (Δx .* inv_r3) * mass
	ay = G .* (Δy .* inv_r3) * mass

	# return acceleration components
	a = [ax ay]
	return a;
end

function main(N, tEnd, dt, G, show=false, save=false)
	# Initialize time
	t = 0

	# Generate Initial Conditions
	seed = MersenneTwister(42)
	mass = randn(seed, N, 1)
	pos  = randn(seed, N, 2)
	vel  = randn(seed, N, 2)

	# calculate initial gravitational accelerations
	acc = getAcc(pos, mass, G);

	# number of timesteps
	Nt = convert(Int64,ceil(tEnd/dt));

	# Simulation Loop
	for i in 1:Nt
		vel += acc * dt/2
		pos += vel * dt
		acc = getAcc(pos, mass, G)
		vel += acc * dt/2
		t += dt;

		x = pos[:, 1]
		y = pos[:, 2]

		if show == true
			display(scatter(x, y, xlims=(-2.5, 2.5), ylims=(-2,2)))
		end
		if save == true
			savefig(scatter(x, y, xlims=(-2.5, 2.5), ylims=(-2,2)), "frames/$i.png")
		end
	end
end

main(25, 0.5, 0.01, 1, true, true)


text:
include("../diagonal_sbp.jl")
include("../deriv_ops_new.jl")
include("../split_matrix_free_dev.jl")
include("../split_matrix_free.jl")

if length(ARGS) != 0
    level = parse(Int,ARGS[1])
    Iterations = parse(Int,ARGS[2])
else
    level = 10
    Iterations = 10000
end

using LinearAlgebra
using SparseArrays
using Plots
using CUDA
using IterativeSolvers
using BenchmarkTools
using MAT

function e(i,n)
    # A = Matrix{Float64}(I,n,n)
    # return A[:,i]
    out = spzeros(n)
    out[i] = 1.0
    return out
end

function eyes(n)
    # return Matrix{Float64}(I,n,n)
    out = spzeros(n,n)
    for i in 1:n
        out[i,i] = 1.0
    end
    return out
end

function u(x,y)
    return sin.(π*x .+ π*y)
end

function Diag(A)
    # Self defined function that is similar to Matlab Diag
    return Diagonal(A[:])
end

function Operators_2d(i, j, p=2, h_list_x = ([1/2^1, 1/2^2, 1/2^3, 1/2^4, 1/2^5, 1/2^6, 1/2^7, 1/2^8, 1/2^9, 1/2^10, 1/2^11, 1/2^12, 1/2^13]),
			 h_list_y = ([1/2^1, 1/2^2, 1/2^3, 1/2^4, 1/2^5, 1/2^6, 1/2^7, 1/2^8, 1/2^9, 1/2^10, 1/2^11, 1/2^12, 1/2^13])
			 )
    hx = h_list_x[i];
    hy = h_list_y[j];

    x = range(0,step=hx,1);
    y = range(0,step=hy,1);
    m_list = 1 ./h_list_x;
    n_list = 1 ./h_list_y;

    # Matrix Size
    N_x = Integer(m_list[i]);
    N_y = Integer(n_list[j]);

    (D1x, HIx, H1x, r1x) = diagonal_sbp_D1(p,N_x,xc=(0,1));
    (D2x, S0x, SNx, HI2x, H2x, r2x) = diagonal_sbp_D2(p,N_x,xc=(0,1));

    (D1y, HIy, H1y, r1y) = diagonal_sbp_D1(p,N_y,xc=(0,1));
    (D2y, S0y, SNy, HI2y, H2y, r2y) = diagonal_sbp_D2(p,N_y,xc=(0,1));

    # BSx = sparse(SNx - S0x);
    # BSy = sparse(SNy - S0y);
    BSx = SNx - S0x
    BSy = SNy - S0y

    # Forming 2d Operators
    # e_1x = sparse(e(1,N_x+1));
    # e_Nx = sparse(e(N_x+1,N_x+1));
    # e_1y = sparse(e(1,N_y+1));
    # e_Ny = sparse(e(N_y+1,N_y+1));
    e_1x = e(1,N_x+1);
    e_Nx = e(N_x+1,N_x+1);
    e_1y = e(1,N_x+1);
    e_Ny = e(N_y+1,N_y+1);

    # I_Nx = sparse(eyes(N_x+1));
    # I_Ny = sparse(eyes(N_y+1));
    I_Nx = eyes(N_x+1);
    I_Ny = eyes(N_y+1);

    e_E = kron(e_Nx,I_Ny);
    e_W = kron(e_1x,I_Ny);
    e_S = kron(I_Nx,e_1y);
    e_N = kron(I_Nx,e_Ny);

    E_E = kron(sparse(Diag(e_Nx)),I_Ny);   # E_E = e_E * e_E'
    E_W = kron(sparse(Diag(e_1x)),I_Ny);
    E_S = kron(I_Nx,sparse(Diag(e_1y)));
    E_N = sparse(kron(I_Nx,sparse(Diag(e_Ny))));

    D1_x = kron(D1x,I_Ny);
    D1_y = kron(I_Nx,D1y);

    D2_x = kron(D2x,I_Ny);
    D2_y = kron(I_Nx,D2y);
    D2 = D2_x + D2_y

    HI_x = kron(HIx,I_Ny);
    HI_y = kron(I_Nx,HIy);

    H_x = kron(H1x,I_Ny);
    H_y = kron(I_Nx,H1y);

    BS_x = kron(BSx,I_Ny);
    BS_y = kron(I_Nx,BSy);

    HI_tilde = kron(HIx,HIx);
    H_tilde = kron(H1x,H1y);

    return (D1_x, D1_y, D2_x, D2_y, D2, HI_x, HI_y, BS_x, BS_y, HI_tilde, H_tilde, I_Nx, I_Ny, e_E, e_W, e_S, e_N, E_E, E_W, E_S, E_N)
end

h_list_x = [1/2^1, 1/2^2, 1/2^3, 1/2^4, 1/2^5, 1/2^6, 1/2^7, 1/2^8, 1/2^9, 1/2^10, 1/2^11, 1/2^12, 1/2^13]
h_list_y = [1/2^1, 1/2^2, 1/2^3, 1/2^4, 1/2^5, 1/2^6, 1/2^7, 1/2^8, 1/2^9, 1/2^10, 1/2^11, 1/2^12, 1/2^13]

let
    k = level
    i = j = k

    hx = h_list_x[i];
    hy = h_list_y[j];

    x = range(0,step=hx,1);
    y = range(0,step=hy,1);
    m_list = 1 ./h_list_x;
    n_list = 1 ./h_list_y;

    # Matrix Size
    N_x = Integer(m_list[i]);
    N_y = Integer(n_list[j]);

    Nx = N_x + 1;
    Ny = N_y + 1;

    println("Nx = $Nx, Ny=$Ny")
    (D1_x, D1_y, D2_x, D2_y, D2, HI_x, HI_y, BS_x, BS_y, HI_tilde, H_tilde, I_Nx, I_Ny, e_E, e_W, e_S, e_N, E_E, E_W, E_S, E_N) = Operators_2d(i,j);
    analy_sol = u(x,y');

    tau_E = -13/hx;
    tau_W = -13/hx;
    tau_N = -1;
    tau_S = -1;

    beta = 1;

     # Forming SAT terms

    ## Formulation 1
    SAT_W = tau_W*HI_x*E_W + beta*HI_x*BS_x'*E_W;
    SAT_E = tau_E*HI_x*E_E + beta*HI_x*BS_x'*E_E;

    # SAT_S = tau_S*HI_y*E_S*D1_y
    # SAT_N = tau_N*HI_y*E_N*D1_y

    SAT_S = tau_S*HI_y*E_S*BS_y;
    SAT_N = tau_N*HI_y*E_N*BS_y;

    SAT_W_r = tau_W*HI_x*E_W*e_W + beta*HI_x*BS_x'*E_W*e_W;
    SAT_E_r = tau_E*HI_x*E_E*e_E + beta*HI_x*BS_x'*E_E*e_E;
    SAT_S_r = tau_S*HI_y*E_S*e_S;
    SAT_N_r = tau_N*HI_y*E_N*e_N;

    (alpha1,alpha2,alpha3,alpha4,beta) = (tau_N,tau_S,tau_W,tau_E,beta);

    g_W = sin.(π*y);
    g_E = -sin.(π*y);
    g_S = -π*cos.(π*x);
    # g_N = -π*cos.(π*x)
    g_N = π*cos.(π*x .+ π);

    # Solving with CPU
    A = D2 + SAT_W + SAT_E + SAT_S + SAT_N;

    b = -2π^2*u(x,y')[:] + SAT_W_r*g_W + SAT_E_r*g_E + SAT_S_r*g_S + SAT_N_r*g_N;

    A = - H_tilde*A;
    b = - H_tilde*b;

    CUDA.reclaim()
    A_d = CUDA.CUSPARSE.CuSparseMatrixCSC(A);
    b_d = CuArray(b);

    for _ = 1:Iterations
        A_d * b_d
    end
    CUDA.memory_status()
    CUDA.unsafe_free!(A_d)
    CUDA.unsafe_free!(b_d)
    CUDA.reclaim()
    return 0
end


text:
""" 2-category of categories, functors, and natural transformations.

Categories in mathematics appear in the large, often as categories of sets with
extra structure, and in the small, as algebraic structures that generalize
groups, monoids, preorders, and graphs. This division manifests in Catlab as
well. Large categories (in spirit, if not in the [technical
sense](https://ncatlab.org/nlab/show/large+category)) occur throughout Catlab as
`@instance`s of the theory of categories. For computational reasons, small
categories are usually presented by generators and relations.

This module provides a minimal interface to accomodate both situations. Category
instances are supported through the wrapper type [`TypeCat`](@ref). Finitely
presented categories are provided by another module, [`FinCats`](@ref).
"""
module Categories
export Cat, TypeCat, Functor, Transformation, dom, codom, compose, id,
  ob, hom, is_hom_equal, ob_map, hom_map, dom_ob, codom_ob, component

using AutoHashEquals

using ...GAT
import ...Theories: Category2, ob, hom, dom, codom, compose, ⋅, ∘, id,
  composeH, *

# Categories
############

""" Abstract base type for a category.

The objects and morphisms in the category have Julia types `Ob` and `Hom`,
respectively. Note that these types do *not* necessarily form an `@instance` of
the theory of categories, as they may not meaningfully form a category outside
the context of this object. For example, a finite category regarded as a
reflexive graph with a composition operation might have type `Cat{Int,Int}`,
where the objects and morphisms are numerical identifiers for vertices and edges
in the graph.

The basic operations available in any category are: [`dom`](@ref),
[`codom`](@ref), [`id`](@ref), [`compose`](@ref).
"""
abstract type Cat{Ob,Hom} end

""" Coerce or look up object in category.
"""
ob(::Cat, x) = x

""" Coerce or look up morphism in category.
"""
hom(::Cat, f) = f

""" Domain of morphism in category.
"""
dom(C::Cat, f) = dom(f)

""" Codomain of morphism in category.
"""
codom(C::Cat, f) = codom(f)

""" Identity morphism on object in category.
"""
id(C::Cat, x) = id(x)

""" Compose morphisms in a category.
"""
compose(C::Cat, fs...) = compose(fs...)

""" Are two morphisms in a category equal?

By default, just checks for equality of Julia objects using ``==``. In some
categories, checking equality of morphisms may involve nontrivial reasoning.
"""
is_hom_equal(C::Cat, f, g) = is_hom_equal(f, g)
is_hom_equal(f, g) = f == g

# Instances
#----------

""" Pair of Julia types regarded as a category.

The Julia types should form an `@instance` of the theory of categories
(`Theories.Category`).
"""
struct TypeCat{Ob,Hom} <: Cat{Ob,Hom} end

TypeCat(Ob::Type, Hom::Type) = TypeCat{Ob,Hom}()

# FIXME: This isn't practical because types are often too tight.
#ob(::TypeCat{Ob,Hom}, x) where {Ob,Hom} = convert(Ob, x)
#hom(::TypeCat{Ob,Hom}, f) where {Ob,Hom} = convert(Hom, f)

Base.show(io::IO, ::TypeCat{Ob,Hom}) where {Ob,Hom} =
  print(io, "TypeCat($Ob, $Hom)")

# Functors
##########

""" Abstract base type for a functor between categories.

A functor has a domain and a codomain ([`dom`](@ref) and [`codom`](@ref)), which
are categories, and object and morphism maps, which can be evaluated using
[`ob_map`](@ref) and [`hom_map`](@ref). The functor object can also be called
directly when the objects and morphisms have distinct Julia types. This is
sometimes but not always the case (see [`Cat`](@ref)), so when writing generic
code one should prefer the `ob_map` and `hom_map` functions.
"""
abstract type Functor{Dom<:Cat,Codom<:Cat} end

""" Evaluate functor on object.
"""
@inline ob_map(F::Functor, x) = do_ob_map(F, x)

""" Evaluate functor on morphism.
"""
@inline hom_map(F::Functor, f) = do_hom_map(F, f)

""" Identity functor on a category.
"""
@auto_hash_equals struct IdentityFunctor{Dom<:Cat} <: Functor{Dom,Dom}
  dom::Dom
end

codom(F::IdentityFunctor) = F.dom

do_ob_map(F::IdentityFunctor, x) = ob(F.dom, x)
do_hom_map(F::IdentityFunctor, f) = hom(F.dom, f)

function Base.show(io::IO, F::IdentityFunctor)
  print(io, "id(")
  show_domains(io, F, codomain=false)
  print(io, ")")
end

""" Composite of functors.
"""
@auto_hash_equals struct CompositeFunctor{Dom,Codom,
    F<:Functor{Dom,<:Cat},G<:Functor{<:Cat,Codom}} <: Functor{Dom,Codom}
  fst::F
  snd::G
end
Base.first(F::CompositeFunctor) = F.fst
Base.last(F::CompositeFunctor) = F.snd

dom(F::CompositeFunctor) = dom(first(F))
codom(F::CompositeFunctor) = codom(last(F))

do_ob_map(F::CompositeFunctor, x) = ob_map(F.snd, ob_map(F.fst, x))
do_hom_map(F::CompositeFunctor, f) = hom_map(F.snd, hom_map(F.fst, f))

function Base.show(io::IO, F::CompositeFunctor)
  print(io, "compose(")
  show(io, first(F))
  print(io, ", ")
  show(io, last(F))
  print(io, ")")
end

show_type_constructor(io::IO, ::Type{<:Functor}) = print(io, "Functor")

function show_domains(io::IO, f; domain::Bool=true, codomain::Bool=true,
                      recurse::Bool=true)
  if get(io, :hide_domains, false)
    print(io, "…")
  else
    if domain
      show(IOContext(io, :compact=>true, :hide_domains=>!recurse), dom(f))
    end
    if codomain
      domain && print(io, ", ")
      show(IOContext(io, :compact=>true, :hide_domains=>!recurse), codom(f))
    end
  end
end

# Instances
#----------

(F::Functor{TypeCat{Ob,Hom}})(x::Ob) where {Ob,Hom} = ob_map(F, x)
(F::Functor{TypeCat{Ob,Hom}})(f::Hom) where {Ob,Hom} = hom_map(F, f)

# Natural transformations
#########################

""" Abstract base type for a natural transformation between functors.

A natural transformation ``α: F ⇒ G`` has a domain ``F`` and codomain ``G``
([`dom`](@ref) and [`codom`](@ref)), which are functors ``F,G: C → D`` having
the same domain ``C`` and codomain ``D``. The transformation consists of a
component ``αₓ: Fx → Gx`` in ``D`` for each object ``x ∈ C``, accessible using
[`component`](@ref) or indexing notation (`Base.getindex`).
"""
abstract type Transformation{C<:Cat,D<:Cat,Dom<:Functor,Codom<:Functor} end

""" Component of natural transformation.
"""
function component end

@inline Base.getindex(α::Transformation, c) = component(α, c)

""" Domain object of natural transformation.

Given ``α: F ⇒ G: C → D``, this function returns ``C``.
"""
dom_ob(α::Transformation) = dom(dom(α)) # == dom(codom(α))

""" Codomain object of natural transformation.

Given ``α: F ⇒ G: C → D``, this function returns ``D``.
"""
codom_ob(α::Transformation) = codom(dom(α)) # == codom(codom(α))

@auto_hash_equals struct IdentityTransformation{C<:Cat,D<:Cat,Dom<:Functor{C,D}} <:
     Transformation{C,D,Dom,Dom}
  dom::Dom
end

codom(α::IdentityTransformation) = α.dom

function component(α::IdentityTransformation, x)
  F = dom(α)
  id(codom(F), ob_map(F, x))
end

const IdIdTransformation{C<:Cat} = IdentityTransformation{C,C,IdentityFunctor{C}}

# 2-category of categories
##########################

@instance Category2{Cat,Functor,Transformation} begin
  dom(F::Functor) = F.dom
  codom(F::Functor) = F.codom
  id(C::Cat) = IdentityFunctor(C)

  function compose(F::Functor, G::Functor; strict::Bool=true)
    !strict || codom(F) == dom(G) ||
      error("Domain mismatch in composition $F ⋅ $G")
    compose_id(F, G)
  end

  dom(α::Transformation) = α.dom
  codom(α::Transformation) = α.codom
  id(F::Functor) = IdentityTransformation(F)

  function compose(α::Transformation, β::Transformation; strict::Bool=true)
    !strict || codom(α) == dom(β) ||
      error("Domain mismatch in vertical composition $α ⋅ $β")
    compose_id(α, β)
  end
  function composeH(α::Transformation, β::Transformation; strict::Bool=true)
    !strict || codom_ob(α) == dom_ob(β) ||
      error("Domain mismatch in horizontal composition $α * $β")
    composeH_id(α, β)
  end
  function composeH(α::Transformation, H::Functor; strict::Bool=true)
    !strict || codom_ob(α) == dom(H) ||
      error("Domain mismatch in whiskering $α * $H")
    composeH_id(α, H)
  end
  function composeH(F::Functor, β::Transformation; strict::Bool=true)
    !strict || codom(F) == dom_ob(β) ||
      error("Domain mismatch in whiskering $F * $β")
    composeH_id(F, β)
  end
end

# XXX: Is this normalization of identities using multiple dispatch a good idea?
# Unlike in `Sets`, it doesn't feel great since it requires so much boilerplate.

@inline compose_id(F::Functor, G::Functor) = do_compose(F, G)
@inline compose_id(F::Functor, ::IdentityFunctor) = F
@inline compose_id(::IdentityFunctor, G::Functor) = G
@inline compose_id(F::IdentityFunctor, ::IdentityFunctor) = F

@inline compose_id(α::Transformation, β::Transformation) = do_compose(α, β)
@inline compose_id(α::Transformation, ::IdentityTransformation) = α
@inline compose_id(::IdentityTransformation, β::Transformation) = β
@inline compose_id(α::IdentityTransformation, ::IdentityTransformation) = α

@inline composeH_id(α::Transformation, β::Transformation) = do_composeH(α, β)
@inline composeH_id(α::Transformation, ::IdIdTransformation) = α
@inline composeH_id(::IdIdTransformation, β::Transformation) = β
@inline composeH_id(α::IdIdTransformation, ::IdIdTransformation) = α

@inline composeH_id(α::Transformation, H::Functor) = do_composeH(α, H)
@inline composeH_id(α::Transformation, ::IdentityFunctor) = α
@inline composeH_id(α::IdentityTransformation, H::Functor) =
  id(compose_id(dom(α), H))
@inline composeH_id(α::IdentityTransformation, ::IdentityFunctor) = α

@inline composeH_id(F::Functor, β::Transformation) = do_composeH(F, β)
@inline composeH_id(::IdentityFunctor, β::Transformation) = β
@inline composeH_id(F::Functor, β::IdentityTransformation) =
  id(compose_id(F, dom(β)))
@inline composeH_id(::IdentityFunctor, β::IdentityTransformation) = β

do_compose(F::Functor, G::Functor) = CompositeFunctor(F, G)

@inline function do_composeH(α::Transformation, β::Transformation)
  do_composeH(α, β, Val{:covariant})
end
function do_composeH(α::Transformation, β::Transformation,
                     ::Type{Val{:covariant}})
  G, H = codom(α), dom(β)
  compose_id(composeH_id(α, H), composeH_id(G, β))
end
function do_composeH(α::Transformation, β::Transformation,
                     ::Type{Val{:contravariant}})
  F, K = dom(α), codom(β)
  compose_id(composeH_id(F, β), composeH_id(α, K))
end

end


text:
module NiceNumbers

import Primes: factor, prodfactors
import Base: showerror
import Base: +, -, *, inv, /, sqrt, ^
import Base: //
import Base: isless, <, <=, ==, hash
import Base: one, zero, isinteger, isfinite
import Base: promote_rule
import Base: isreal, real, imag, conj, abs, abs2
import LinearAlgebra: norm, norm2
import Base: copysign

export NiceNumber, nice, @nice
export isrational

struct NiceNumber <: Number
    a::Rational{Int}
    coeff::Rational{Int}
    radicand::Int
    function NiceNumber(a::S, coeff::T, radicand::Int) where {S<:Union{Integer,Rational}, T<:Union{Integer,Rational}}
        if radicand == 0
            coeff = zero(coeff)
        elseif coeff == 0
            radicand = zero(radicand)
        else
            coeff, radicand = reduce_root(coeff, radicand)
            if radicand == 1
                a += coeff
                coeff = zero(coeff)
                radicand = zero(radicand)
            end
        end
        new(a, coeff, radicand)
    end
end
NiceNumber(x::T) where {T<:Union{Integer,Rational}} = NiceNumber(x, 0, 0)
NiceNumber(a, coeff, radicand::Rational) =
    NiceNumber(a, coeff // denominator(radicand), numerator(radicand) * denominator(radicand))
NiceNumber(a, coeff, radicand) = NiceNumber(a)+NiceNumber(coeff)*sqrt(NiceNumber(radicand))
NiceNumber(x::AbstractFloat) = NiceNumber(rationalize(Int, x), 0, 0)
NiceNumber(z::Complex{T}) where T<:Real = NiceNumber(z.re) + NiceNumber(0,z.im,-1)
NiceNumber(n::NiceNumber) = n

"""
    reduce_root(coeff, radicand)

Simplifies the expression ``\\textrm{coeff}\\cdot\\sqrt{\\textrm{radicand}}`` by factoring all remaining squares
in `radicand` into `coeff`.
"""
function reduce_root(coeff::Union{Integer,Rational}, radicand::Integer)
    F = factor(radicand)
    for (k, v) in F
        if v ≥ 2
            coeff *= k^(v ÷ 2)
            F[k] = v % 2
        end
    end
    radicand = prodfactors(F)
    return coeff, radicand
end

struct NotNiceError <: Exception
    expr::String
end
showerror(io::IO, e::NotNiceError) = print(io, e.expr, " is not nice anymore!")

# always promote to NiceNumber; if that can't be done, it can't be nice
promote_rule(::Type{NiceNumber}, ::Type{T}) where {T<:Number} = NiceNumber
one(::NiceNumber) = NiceNumber(1, 0, 0)
zero(::NiceNumber) = NiceNumber(0, 0, 0)

AbstractFloat(n::NiceNumber) = float(n.a) + float(n.coeff) * √n.radicand
(::Type{T})(n::NiceNumber) where {T<:AbstractFloat} =
    convert(T, n.a) + convert(T, n.coeff) * convert(T, √n.radicand)

"""
    isrational(n::NiceNumber)

Whether `n` is purely rational, i.e. has no square root portion.
"""
isrational(n::NiceNumber) = iszero(n.coeff)
isinteger(n::NiceNumber) = isrational(n) && isinteger(n.a)
isfinite(n::NiceNumber) = isfinite(n.a) && isfinite(n.coeff)
isreal(n::NiceNumber) = n.radicand >= 0
real(n::NiceNumber) = isreal(n) ? n : NiceNumber(n.a)
imag(n::NiceNumber) = isreal(n) ? zero(n) : NiceNumber(0,n.coeff,-n.radicand)

function Base.show(io::IO, n::NiceNumber)
    pretty(r::Rational) = isinteger(r) ? numerator(r) : r

    if isrational(n)
        print(io, pretty(n.a))
    else
        print(
            io,
            iszero(n.a) ? "" : pretty(n.a),
            n.coeff>0 ? iszero(n.a) ? "" : "+" : "-",
            isone(abs(n.coeff)) ? "" : pretty(abs(n.coeff)),
            isone(abs(n.coeff)) ? "" : "⋅",
            "√",
            n.radicand,
        )
    end
end
Base.show(io::IO, ::MIME"text/plain", n::NiceNumber) = print(io, "Nice number:\n   ", n)

+(n::NiceNumber) = n
-(n::NiceNumber) = NiceNumber(-n.a, -n.coeff, n.radicand)
function +(n::NiceNumber, m::NiceNumber)
    if !(n.radicand == m.radicand || isrational(n) || isrational(m))
        throw(NotNiceError("$n + $m"))
    end
    if isrational(m)
        return NiceNumber(n.a + m.a, n.coeff, n.radicand)
    else
        # either n is rational (n.coeff==0) or the radicands match
        return NiceNumber(n.a + m.a, n.coeff + m.coeff, m.radicand)
    end
end
-(n::NiceNumber, m::NiceNumber) = n + (-m)

function *(n::NiceNumber, m::NiceNumber)
    # if both have a complex part we need a minus in the last term
    if n.radicand < 0 && m.radicand < 0
        return NiceNumber(m.a * n.a, m.a * n.coeff, n.radicand) +
               NiceNumber(0, m.coeff * n.a, m.radicand) +
               NiceNumber(0, -n.coeff * m.coeff, n.radicand * m.radicand)
    else
        return NiceNumber(m.a * n.a, m.a * n.coeff, n.radicand) +
               NiceNumber(0, m.coeff * n.a, m.radicand) +
               NiceNumber(0, n.coeff * m.coeff, n.radicand * m.radicand)
    end
end

inv(n::NiceNumber) = NiceNumber(n.a, -n.coeff, n.radicand) * inv(n.a^2 - n.coeff^2 * n.radicand)
/(n::NiceNumber, m::NiceNumber) = n * inv(m)

sqrt(n::NiceNumber) = isrational(n) ? NiceNumber(0, 1, n.a) : throw(NotNiceError("sqrt($n)"))
"""
    nthroot(m::NiceNumber, n)

Returns the nth root of `m`. Works by repeatedly determining the square root and thus only for
powers of two.
"""
function nthroot(m::NiceNumber, n)
    !ispow2(n) && throw(NotNiceError("nthroot($m,$n)"))
    while n > 1
        m = sqrt(m)
        n = n >> 1
    end
    return m
end
^(x::Number, n::NiceNumber) = isrational(n) ? NiceNumber(x)^n.a : throw(NotNiceError("$x^$n"))
^(n::NiceNumber, r::Rational) = nthroot(n^numerator(r), denominator(r))

# TODO: think about comparisons more (or less)
isless(n::NiceNumber, m::NiceNumber) = float(n) < float(m)
isless(n::Real, m::NiceNumber) = n < float(m)
isless(n::NiceNumber, m::Real) = !isless(m,n)
<(n::NiceNumber, m::NiceNumber) = float(n) < float(m)
<=(n::NiceNumber, m::NiceNumber) = n === m || n < m
==(n::NiceNumber, m::AbstractFloat) = float(n) == m
==(m::AbstractFloat, n::NiceNumber) = n == m
hash(n::NiceNumber, h::UInt) = hash(n.a, hash(n.coeff, hash(n.radicand, hash(:NiceNumber, h))))

//(n::S, m::T) where {S<:Union{NiceNumber,Integer,Rational},T<:Union{NiceNumber,Integer,Rational}} =
    n / m

conj(n::NiceNumber) = isreal(n) ? n : NiceNumber(n.a, -n.coeff, n.radicand)
abs(n::NiceNumber) = isreal(n) ? float(n)>0 ? n : -n : sqrt(n*conj(n))
abs2(n::NiceNumber) = n*conj(n)
norm(n::NiceNumber) = abs(n)
norm2(v::AbstractArray{NiceNumber,1}) = sqrt(v'v)
copysign(n::NiceNumber, m::NiceNumber) = sign(n) == sign(m) ? n : -n

## macro stuff
"""
    nice(x)

If `x` is an expression it replaces all occuring numbers by `NiceNumber`s.

If `x` is a number it turns it into a `NiceNumber`.

If `x` is the symbol `:im` it turns it into `NiceNumber(im)`.

Otherwise it does nothing.
"""
function nice end
nice(x) = x
nice(n::Number) = NiceNumber(n)
nice(s::Symbol) = s===:im ? NiceNumber(im) : s
nice(ex::Expr) = Expr(ex.head, map(nice, ex.args)...)

"""
    @nice

Return equivalent expression with all numbers converted to `NiceNumber`s.
"""
macro nice(code)
    return esc(nice(code))
end

include("NiceLinAlg.jl")

end # module


text:
export apply_back!, apply_back

const Rotor{N,T} = Union{RotationGate{N,T},PutBlock{N,<:Any,<:RotationGate{<:Any,T}}}

as_scalar(arr::AbstractArray{T,0}) where {T} = arr[]
as_scalar(arr) = arr

"""
    generator(rot::Rotor) -> AbstractBlock

Return the generator of rotation block.
"""
generator(rot::RotationGate) = rot.block
generator(rot::PutBlock{N,C,GT}) where {N,C,GT<:RotationGate} =
    PutBlock{N}(generator(rot |> content), rot |> occupied_locs)

"""
    apply_back!((ψ, ∂L/∂ψ*), circuit::AbstractBlock, collector) -> AbstractRegister

back propagate and calculate the gradient ∂L/∂θ = 2*Re(∂L/∂ψ*⋅∂ψ*/∂θ), given ∂L/∂ψ*.
`ψ` is the output register, ∂L/∂ψ* should also be register type.

Note: gradients are stored in `Diff` blocks, it can be access by either `diffblock.grad` or `gradient(circuit)`.
Note2: now `apply_back!` returns the inversed gradient!
"""
function apply_back!(st, block::AbstractBlock, collector) #,AbstractContainer{<:PrimitiveBlock}
    out, outδ = st
    if nparameters(block) == 0
        adjblock = block'
        in = apply!(out, adjblock)
        inδ = apply!(outδ, adjblock)
        return (in, inδ)
    else
        throw(MethodError(apply_back, (st, block, collector)))
    end
end

function apply_back!(st, block::Concentrator{N}, collector) where {N}
    out, outδ = st
    focus!(out, block.locs)
    focus!(outδ, block.locs)
    apply_back!((out, outδ), content(block), collector)
    relax!(out, block.locs; to_nactive = N)
    relax!(outδ, block.locs; to_nactive = N)
    return (out, outδ)
end

function apply_back!(st, block::Rotor{N}, collector) where {N}
    out, outδ = st
    adjblock = block'
    backward_params!((out, outδ), block, collector)
    in = apply!(out, adjblock)
    inδ = apply!(outδ, adjblock)
    return (in, inδ)
end

function apply_back!(st, block::TimeEvolution{N}, collector) where {N}
    out, outδ = st
    adjblock = block'

    out, outδ = st
    input = apply!(out, adjblock)
    for o in outδ
        !all(x -> x ≈ 0.0im, o.state) && apply!(o, adjblock)
    end
    pushfirst!(collector, -sum(imag(input' * apply!(copy(outδ), block.H))))
    return (input, outδ)
end

function apply_back!(st, block::PutBlock{N}, collector) where {N}
    out, outδ = st
    adjblock = block'
    in = apply!(out, adjblock)
    adjmat = outerprod(outδ, in)
    mat_back!(datatype(in), block, adjmat, collector)
    inδ = apply!(outδ, adjblock)
    return (in, inδ)
end

function apply_back!(st, block::KronBlock{N}, collector) where {N}
    apply_back!(st, chain(N, [put(loc => block[loc]) for loc in block.locs]), collector)
end

function apply_back!(st, block::ControlBlock{N}, collector) where {N}
    out, outδ = st
    adjblock = block'
    in = apply!(out, adjblock)
    #adjm = adjcunmat(outerprod(in, outδ), N, block.ctrl_locs, block.ctrl_config, mat(content(block)), block.locs)
    adjmat = outerprod(outδ, in)
    mat_back!(datatype(in), block, adjmat, collector)
    inδ = apply!(outδ, adjblock)
    return (in, inδ)
end

function apply_back!(st, block::Daggered, collector)
    out, outδ = st
    adjblock = block'
    in = apply!(out, adjblock)
    adjmat = outerprod(in, outδ)
    mat_back!(datatype(in), content(block), adjmat, collector)
    inδ = apply!(outδ, adjblock)
    return (in, inδ)
end

function apply_back!(st, block::Scale, collector)
    out, outδ = st
    apply_back!((out, outδ), content(block), collector)
    outδ.state .= outδ.state .* conj(factor(block))
    out.state .= out.state ./ factor(block)
    return (out, outδ)
end

function apply_back!(st, circuit::ChainBlock, collector)
    for blk in Base.Iterators.reverse(subblocks(circuit))
        st = apply_back!(st, blk, collector)
    end
    return st
end

function apply_back!(st, circuit::Add, collector; in)
    out, outδ = st
    adjmat = outerprod(outδ, in)
    for blk in Base.Iterators.reverse(subblocks(circuit))
        mat_back!(datatype(in), blk, adjmat, collector)
    end
    inδ = apply!(outδ, adjblock)
    (in, inδ)
end

function apply_back!(st, block::RepeatedBlock{N,C}, collector) where {N,C}
    if nparameters(content(block)) == 0
        return apply!.(st, Ref(block'))
    end
    res = Any[]
    st = apply_back!(st, chain(N, [put(loc => content(block)) for loc in block.locs]), res)
    res = dropdims(sum(reshape(res, :, C), dims = 2), dims = 2) |> as_scalar
    prepend!(collector, res)
    return st
end

# TODO: concentrator, repeat, kron
apply_back!(st, block::Measure, collector) = throw(MethodError(apply_back!, (st, block, collector)))

function backward_params!(st, block::Rotor, collector)
    in, outδ = st
    Σ = generator(block)
    g = dropdims(sum(conj.(state(in |> Σ)) .* state(outδ), dims = (1, 2)), dims = (1, 2)) |> as_scalar
    pushfirst!(collector, -imag(g) / 2)
    in |> Σ
    nothing
end

"""
    apply_back(st::Tuple{<:ArrayReg, <:ArrayReg}, block::AbstractBlock; kwargs...) -> (out, outδ), paramsδ

The backward function of `apply!`. Returns a tuple of ((input register, gradient of input register), parameter gradients)
"""
function apply_back(st::Tuple{<:ArrayReg,<:ArrayReg}, block::AbstractBlock; kwargs...)
    col = []
    in, inδ = apply_back!(st, block, col; kwargs...)
    (in, inδ), col
end

Base.adjoint(::typeof(expect)) = Adjoint(expect)
Base.show(io::IO, ::Adjoint{Any,typeof(expect)}) = print(io, "expect'")
Base.show(io::IO, ::MIME"text/plain", ::Adjoint{Any,typeof(expect)}) = print(io, "expect'")
"""
expect')(op::AbstractBlock, circuit::Pair{<:ArrayReg, <:AbstractBlock})

"""
function (::Adjoint{Any,typeof(expect)})(op::AbstractBlock, circuit::Pair{<:ArrayReg,<:AbstractBlock})
    reg, c = circuit
    out = copy(reg) |> c
    outδ = copy(out) |> op
    (in, inδ), paramsδ = apply_back((out, outδ), c)
    return outδ => paramsδ .* 2
end


text:
export substitude

function substitude(d::Dict{Variable}, t::Term)
    Term(t.head, map(x->substitude(d, x), t.args))
end

substitude(d::Dict{Variable}, t::Variable) = haskey(d, t) ? d[t] : t
substitude(d::Dict{Variable}, t) = t

Base.replace(t::Expression, ::Pair) = t

function Base.replace(t::Term, d::Pair{<:Term, <:Expression})
    p, s = d
    m = match(p, t)
    if m !== nothing
        for each in m
            s = substitude(each, s)
        end
        return s
    end

    return Term(replace(t.head, d), map(x->replace(x, d), t.args))
end


text:
using Juno
using Flux.Tracker: back!

runall(f) = f
runall(fs::AbstractVector) = () -> foreach(call, fs)

# The AD generates fairly large backtraces that are unhelpful if you interrupt
# while training; this just cleans that up.
macro interrupts(ex)
  :(try $(esc(ex))
    catch e
      e isa InterruptException || rethrow()
      throw(e)
    end)
end

"""
    train!(loss, data, opt)

For each datapoint `d` in `data` computes the gradient of `loss(d...)` through
backpropagation and calls the optimizer `opt`.

Takes a callback as keyword argument `cb`. For example, this will print "training"
every 10 seconds:

```julia
Flux.train!(loss, data, opt,
            cb = throttle(() -> println("training"), 10))
```

The callback can return `:stop` to interrupt the training loop.

Multiple optimisers and callbacks can be passed to `opt` and `cb` as arrays.
"""
function train!(loss, data, opt; cb = () -> ())
  cb = runall(cb)
  opt = runall(opt)
  @progress for d in data
    l = loss(d...)
    isinf(l) && error("Loss is Inf")
    isnan(l) && error("Loss is NaN")
    @interrupts back!(l)
    opt()
    cb() == :stop && break
  end
end

"""
    @epochs N body

Run `body` `N` times. Mainly useful for quickly doing multiple epochs of
training in a REPL.

```julia
julia> @epochs 2 println("hello")
INFO: Epoch 1
hello
INFO: Epoch 2
hello
```
"""
macro epochs(n, ex)
  :(@progress for i = 1:$(esc(n))
      info("Epoch $i")
      $(esc(ex))
    end)
end


text:
## Author:  Pearson
## Date: April 2020
## Julia version: 1.3.1
## Purpose: Solve EKNR(2016) Simple 2 Cities Case

# Precompile Packages
using Parameters
using MATLAB
using NLsolve
using Plots

# Introduce Functions
include("Simple_2Cty_Functions.jl")

# Program Parameters
T_data = 20
T_tail = 250
T = T_data+T_tail

# Initial Conditions and Time Invariant Parameters
K_init_rel_SS=[1/2 1/2]'
α = 0.55
βᴸ = 2/3
βᴷ = 1-βᴸ
δ = 0.1
ρ = 0.95
θ = 2
ω = [1/2, 1/2]
B = βᴸ^-βᴸ*βᴷ^-βᴷ

# Read data
fname = MatFile(joinpath(@__DIR__, "..", "input", "EKNR_Simple_2Cty_Shocks.mat"))
χ_data = get_variable(fname,"chi_data")
Aᴰ_data = get_variable(fname,"AD_data")
ϕ_data = get_variable(fname,"phi_data")
L_data = get_variable(fname,"L_data")
dₙᵢ_data = get_variable(fname,"dni_data")

# Hypothetical Data on Shocks(Assumed Constant During T_tail)
χ_tail = χ_data[:,T_data]*ones(1,T_tail)
Aᴰ_tail = Aᴰ_data[:,T_data]*ones(1,T_tail)
ϕ_tail = ϕ_data[:,T_data]*ones(1,T_tail)
L_tail = L_data[:,T_data]*ones(1,T_tail)
dₙᵢ = repeat(dₙᵢ_data[:,:,T_data],1,1,T)
χ = [χ_data χ_tail]
Aᴰ = [Aᴰ_data Aᴰ_tail]
ϕ = [ϕ_data ϕ_tail]
L = [L_data L_tail]
dₙᵢ[:,:,1:T_data] = dₙᵢ_data

# Hypothetical Data on Shocks in Change Form
χ̂ = [χ[:,2:end]./χ[:,1:end-1] ones(eltype(χ),size(χ[:,end]))]
Âᴰ = [Aᴰ[:,2:end]./Aᴰ[:,1:end-1] ones(eltype(Aᴰ),size(Aᴰ[:,end]))]
ϕ̂ = [ϕ[:,2:end]./ϕ[:,1:end-1] ones(eltype(ϕ),size(ϕ[:,end]))]
L̂ = [L[:,2:end]./L[:,1:end-1] ones(eltype(L),size(L[:,end]))]
d̂ₙᵢ = ones(eltype(dₙᵢ), size(dₙᵢ)); d̂ₙᵢ[:,:,1:end-1] = dₙᵢ[:,:,2:end]./dₙᵢ[:,:,1:end-1]

###########################################################################################
#                            Solve for the steady state
###########################################################################################
# Guess and Other Exogenous Variables
Yˢₛₛ = ω.*ϕ[:,T]
guess_ss = [1.5*Yˢₛₛ 200*Yˢₛₛ]

# Wrap All Exogenous Variables and Parameters
myexos_ss = @with_kw (χ = χ[:,T], Aᴰ = Aᴰ[:,T], L = L[:,T], dₙᵢ = dₙᵢ[:,:,T], Yˢₛₛ = Yˢₛₛ)
myparams_ss = @with_kw (α = α, βᴸ = βᴸ, βᴷ = βᴷ, δ = δ, ρ = ρ, θ = θ, B = B)
exos_ss = myexos_ss()
params_ss = myparams_ss()

# Solve for Steady State
println("\n\nStart to solve the steady state.")
println("Run time and memory cost:")
@time results_ss = nlsolve(
    (res_ss, guess_ss) -> Fun_2Cty_SS!(res_ss, guess_ss, exos_ss, params_ss),
    guess_ss,
    ftol=1e-6,
    method=:newton,
    autodiff=:forward,
    show_trace=false,
)

# Check Convergence
converged(results_ss) || error("Failed to converge in $(results_ss.iterations) iterations.")
println("Successfully solved steady state.\n")

# Catch Solutions for Steady State
res_ss = similar(results_ss.zero)
res_ss, Yₛₛ, Kₛₛ = Fun_2Cty_SS!(res_ss, results_ss.zero, exos_ss, params_ss)

###########################################################################################
#                         Solve for Solution to the Problem in Levels
###########################################################################################
# Initial Conditions, Guess and Other Exogenous Variables
K₁ = K_init_rel_SS.*Kₛₛ
Yˢ = ω.*ϕ
guess_lev = repeat(Yₛₛ,1,T)

# Wrap All Exogenous Variables and Parameters
myinit_lev = @with_kw (K₁ = K₁,)
myexos_lev = @with_kw (χ = χ, Aᴰ = Aᴰ, L = L, dₙᵢ = dₙᵢ, Kₛₛ = Kₛₛ, Yˢ = Yˢ)
myparams_lev = @with_kw (α = α, βᴷ = βᴷ, βᴸ = βᴸ, δ = δ, ρ = ρ, θ = θ, T = T)
init_lev = myinit_lev()
exos_lev = myexos_lev()
params_lev = myparams_lev()

# Solve the Problem in Levels
println("Start to solve the problem in levels.")
println("Run time and memory cost:")
@time results_lev =
    try
        results_lev = nlsolve(
            (res_lev, guess_lev) -> Fun_2Cty_Levels!(
                res_lev, guess_lev, init_lev, exos_lev, params_lev),
            guess_lev,
            ftol=1e-6,
            method=:newton,
            autodiff=:forward,
            show_trace=false,
        )
    catch err
        if isa(err, DomainError)
            error("Failed to solve the problem in levels, please try again.")
        end
    end

# Check Convergence
converged(results_lev) || error("Failed to converge in $(results_lev.iterations) iterations.")
println("Successfully solved the problem in levels.\n")

# Catch Solutions
res_lev = similar(results_lev.zero)
res_lev, K_lev, Y_lev, Yᴰ_lev, Xᴰ_lev, πₙᵢ_lev = Fun_2Cty_Levels!(
    res_lev, results_lev.zero, init_lev, exos_lev, params_lev)

###########################################################################################
#                         Solve for Solution to the Problem in Changes
###########################################################################################
# Guess
guess_hat = ones(2,T)

# Wrap All Exogenous Variables and Parameters
myinit_hat = @with_kw (Y₁ = Y_lev[:,1], Yᴰ₁ = Yᴰ_lev[:,1], πₙᵢ₁ = πₙᵢ_lev[:,:,1])
myexos_hat = @with_kw (χ̂ = χ̂, Âᴰ = Âᴰ, L̂ = L̂, d̂ₙᵢ = d̂ₙᵢ, ϕ̂ = ϕ̂)
myparams_hat = @with_kw (α = α, βᴷ = βᴷ, βᴸ = βᴸ, δ = δ, ρ = ρ, θ = θ, T = T)
init_hat = myinit_hat()
exos_hat = myexos_hat()
params_hat = myparams_hat()

# Solve the Problem in Changes
println("Start to solve the problem in changes.")
println("Run time and memory cost:")
@time results_hat =
    try
        results_hat = nlsolve(
            (res_hat, guess_hat) -> Fun_2Cty_Changes!(
                res_hat, guess_hat, init_hat, exos_hat, params_hat),
            guess_hat,
            ftol=1e-6,
            method=:newton,
            autodiff=:forward,
            show_trace=false,
        )
    catch err
        if isa(err, DomainError)
            error("Failed to solve the problem in changes, please try again.")
        end
    end

# Check Convergence
converged(results_hat) || error("Failed to converge in $(results_hat.iterations) iterations.")
println("Successfully solved the problem in changes.\n")

# Catch Solutions
res_hat = similar(results_hat.zero)
res_hat, K̂, Ŷ, Ŷᴰ, X̂ᴰ, π̂ₙᵢ = Fun_2Cty_Changes!(
    res_hat, results_hat.zero, init_hat, exos_hat, params_hat)
K_hat = similar(K_lev)
Y_hat = similar(Y_lev)
K_hat[:,1] = K₁
Y_hat[:,1] = Y_lev[:,1]
for t = 2:T
    K_hat[:,t] = K_hat[:,t-1].*K̂[:,t-1]
    Y_hat[:,t] = Y_hat[:,t-1].*Ŷ[:,t-1]
end

###########################################################################################
#                                      Plot Figures
###########################################################################################
println("Plotting the figure.")
l = @layout [a b; c d]
p1 = Plots.plot(K_lev[1,1:80], line=(:solid, 2), label="Level Solution");
    Plots.plot!(K_hat[1,1:80], line=(:dash, 2), label="Change Solution");
p2 = Plots.plot(K_lev[2,1:80], line=(:solid, 2), label="Level Solution");
    Plots.plot!(K_hat[2,1:80], line=(:dash, 2), label="Change Solution");
p3 = Plots.plot(Y_lev[1,1:80], line=(:solid, 2), label="Level Solution");
    Plots.plot!(Y_hat[1,1:80], line=(:dash, 2), label="Change Solution");
p4 = Plots.plot(Y_lev[2,1:80], line=(:solid, 2), label="Level Solution");
    Plots.plot!(Y_hat[2,1:80], line=(:dash, 2), label="Change Solution");
figures_all = plot(p1, p2, p3, p4, layout = l)
display(figures_all)
savefig(figures_all, joinpath(@__DIR__, "..", "output", "EKNR_2Cty.pdf"))
println("The figure is saved.")


text:
# functions for scaling features

"""
    standardization!(data::Matrix)

In-place standardization of the features in the matrix `data`.
"""
function standardization!(data::Matrix)

  # find dim. for data
  dim = minimum(size(data))

  # standardize data
  if size(data)[1] > size(data)[2]
    for i = 1:dim
        data[:,i] = (data[:,i] - mean(data[:,i]))/std(data[:,i])
    end
  else
    for i = 1:dim
        data[i,:] = (data[i,:] - mean(data[i,:]))/std(data[i,:])
    end
  end

end

"""
    standardization(data::Matrix, return_column_major_order::Bool=true)

Standardization of the features in the matrix `data`.
"""
function standardization(data::Matrix, return_column_major_order::Bool=true)

  # find dim. for data
  dim = minimum(size(data))

  # transform data to column-major order if necessary
  if size(data)[1] > size(data)[2]
      data = data'
  end
  data_standardized = copy(data)

  # standadize data
  for i = 1:dim
      data_standardized[i,:] = (data[i,:] - mean(data[i,:]))/std(data[i,:])
  end

  # return data in column major order
  if !return_column_major_order
      data_standardized = data_standardized'
  end

  return data_standardized

end

"""
    standardization!(data::Vector)

Standardization of the vector `data`.
"""
function standardization(data::Vector)
  data_standardized = (data - mean(data))/std(data)
  return data_standardized
end

"""
    standardization!(data::Vector)

In-place standardization of the vector `data`.
"""
function standardization!(data::Vector)
  data[:] = (data - mean(data))/std(data)
end

max_stars_repo_name: SamuelWiqvist/adamcmcpaper

max_issues_repo_name: SamuelWiqvist/adamcmcpaper

max_forks_repo_name: SamuelWiqvist/adamcmcpaper


text:
# This file is a part of BAT.jl, licensed under the MIT License (MIT).

abstract type SpatialVolume{T<:Real} end

Base.eltype(b::SpatialVolume{T}) where T = T

Random.rand(rng::AbstractRNG, vol::SpatialVolume) =
    rand!(rng, vol, Vector{float(eltype(vol))}(undef, ndims(vol)))

Random.rand(rng::AbstractRNG, vol::SpatialVolume, n::Integer) =
    rand!(rng, vol, VectorOfSimilarVectors(Matrix{float(eltype(vol))}(undef, ndims(vol), n)))

"""
    log_volume(vol::SpatialVolume)

*BAT-internal, not part of stable public API.*

Get the logarithm of the volume of the space in `vol`.
"""
function log_volume end

"""
    fromuhc!(Y::AbstractVector, X::AbstractVector, vol::SpatialVolume)
    fromuhc!(Y::VectorOfSimilarVectors, X::VectorOfSimilarVectors, vol::SpatialVolume)

*BAT-internal, not part of stable public API.*

Bijective transformation of coordinates `X` within the unit hypercube to
coordinates `Y` in `vol`. If `X` and `Y` are matrices, the transformation is
applied to the column vectors. Use `Y === X` to transform in-place.

Use `inv(fromuhc!)` to get the the inverse transformation.
"""
function fromuhc! end

function inv_fromuhc! end

Base.inv(::typeof(fromuhc!)) = inv_fromuhc!
Base.inv(::typeof(inv_fromuhc!)) = fromuhc!

"""
    fromuhc(X::AbstractVector, vol::SpatialVolume)
    fromuhc(X::VectorOfSimilarVectors, vol::SpatialVolume)

*BAT-internal, not part of stable public API.*

Bijective transformation from unit hypercube to `vol`. See `fromuhc!`.

Use `inv(fromuhc)` to get the the inverse transformation.
"""
function fromuhc(X::Union{AbstractVector,VectorOfSimilarVectors}, vol::SpatialVolume)
    fromuhc!(similar(X), X, vol)
end

function inv_fromuhc(X::Union{AbstractVector,VectorOfSimilarVectors}, vol::SpatialVolume)
    inv_fromuhc!(similar(X), X, vol)
end

Base.inv(::typeof(fromuhc)) = inv_fromuhc
Base.inv(::typeof(inv_fromuhc)) = fromuhc

struct HyperRectVolume{T<:Real} <: SpatialVolume{T}
    # ToDo: Use origin and widths instead of lo and hi, as in GeometryTypes.HyperRectangle?
    lo::Vector{T}
    hi::Vector{T}

    function HyperRectVolume{T}(lo::AbstractVector{T}, hi::AbstractVector{T}) where {T<:Real}
        (axes(lo) != axes(hi)) && throw(ArgumentError("lo and hi must have the same indices"))
        new{T}(lo, hi)
    end
end

HyperRectVolume(lo::AbstractVector{T}, hi::AbstractVector{T}) where {T<:Real} = HyperRectVolume{T}(lo, hi)

Base.:(==)(a::HyperRectVolume, b::HyperRectVolume) = a.lo == b.lo && a.hi == b.hi

Base.ndims(vol::HyperRectVolume) = size(vol.lo, 1)

function Base.isempty(vol::HyperRectVolume)
        @inbounds for i in eachindex(vol.lo, vol.hi)
             (vol.lo[i] > vol.hi[i]) && return true
        end
        isempty(vol.lo)
end

Base.similar(vol::HyperRectVolume) = HyperRectVolume(similar(vol.lo), similar(vol.hi))

Base.in(x::AbstractVector{<:Real}, vol::HyperRectVolume) =
    _all_lteq(vol.lo, x, vol.hi)

Base.in(x::Any, vol::HyperRectVolume) = unshaped(x) in vol

function Base.isinf(vol::HyperRectVolume)
    return (any(isinf.(vol.hi)) || any(isinf.(vol.lo)))
end

function Base.copy!(
    target::HyperRectVolume{T},
    src::HyperRectVolume{T}) where {T<:AbstractFloat}

    p = ndims(src)
    resize!(target.lo, p)
    copyto!(target.lo, src.lo)
    resize!(src.hi, p)
    copyto!(target.hi, src.hi)
    nothing
end

# TODO: Remove?
# function Base.in(X::AbstractMatrix, vol::HyperRectVolume, j::Integer)
#     lo = vol.lo
#     hi = vol.hi
#     for i in axes(X, 1)
#         (lo[i] <= X[i, j] <= hi[i]) || return false
#     end
#     return true
# end

# ToDo:
# Base.in(x::Matrix, vol::HyperRectVolume) =

function Base.intersect(a::HyperRectVolume, b::HyperRectVolume)
    c = similar(a)
    c.lo .= max.(a.lo, b.lo)
    c.hi .= min.(a.hi, b.hi)
    c
end

function Random.rand!(rng::AbstractRNG, vol::HyperRectVolume, x::Union{AbstractVector{<:Real},VectorOfSimilarVectors{<:Real}})
    x_flat = flatview(x)
    rand!(rng, x_flat)
    x_flat .= x_flat .* (vol.hi - vol.lo) .+ vol.lo # TODO: Avoid memory allocation
    x
end

function log_volume(vol::HyperRectVolume{T}) where {T}
    R = promote_type(Float64, T)
    S = promote_type(DoubleFloat{Float64}, T)
    s = zero(S)
    hi = vol.hi
    lo = vol.lo
    @assert axes(hi) == axes(lo)
    @inbounds @simd for i in eachindex(hi)
        d = max(zero(R), R(hi[i]) - R(lo[i]))
        s += log(d)
    end
    R(s)
end

# ToDo:
# log_intersect_volume(a::HyperRectVolume{T}, b::HyperRectVolume{T}) where {T} = ...

fromuhc!(Y::AbstractVector, X::AbstractVector, vol::HyperRectVolume) = _fromuhc_impl!(Y, X, vol)
fromuhc!(Y::VectorOfSimilarVectors, X::VectorOfSimilarVectors, vol::HyperRectVolume) = _fromuhc_impl!(Y, X, vol)

function _fromuhc_impl!(Y::Union{AbstractVector,VectorOfSimilarVectors}, X::Union{AbstractVector,VectorOfSimilarVectors}, vol::HyperRectVolume)
    _all_in_ui(flatview(X)) || throw(ArgumentError("X not in unit hypercube"))
    Y_flat = flatview(Y)
    Y_flat .= unsafe_fromui.(flatview(X), vol.lo, vol.hi)
    Y
end

inv_fromuhc!(Y::AbstractVector, X::AbstractVector, vol::HyperRectVolume) = _inv_fromuhc_impl!(Y, X, vol)
inv_fromuhc!(Y::VectorOfSimilarVectors, X::VectorOfSimilarVectors, vol::HyperRectVolume) = _inv_fromuhc_impl!(Y, X, vol)

function _inv_fromuhc_impl!(Y::Union{AbstractVector,VectorOfSimilarVectors}, X::Union{AbstractVector,VectorOfSimilarVectors}, vol::HyperRectVolume)
    X in vol || throw(ArgumentError("X not in vol"))
    Y_flat = flatview(Y)
    Y_flat .= unsafe_inv_fromui.(flatview(X), vol.lo, vol.hi)
    Y
end


text:
function treat_constants(MATRICES_with_CONSTANTS, CONSTRAINTS)
    #find constant rows
    index_none_column2 = findall(x -> x == "NONE", MATRICES_with_CONSTANTS[:,2])
    index_none_column3 = findall(x -> x == "NONE", MATRICES_with_CONSTANTS[:,3])
    rows_to_remove = intersect(index_none_column2,index_none_column3)
    CONSTANTS = MATRICES_with_CONSTANTS[rows_to_remove,:]
    #remove constants from MATRICES
    MATRICES_without_CONSTANTS = MATRICES_with_CONSTANTS[setdiff(1:end, rows_to_remove), :]
    #define dict constant per objctr
    constants_per_objconstraint = Dict(CONSTRAINTS[i,1] => 0.0 for i=1:size(CONSTRAINTS, 1))
    constants_per_objconstraint["OBJ"] = 0.0
    for l=1:size(CONSTANTS,1)
        name = CONSTANTS[l,1]
        constants_per_objconstraint[name] = CONSTANTS[l,4]
    end
    return MATRICES_without_CONSTANTS, constants_per_objconstraint
end #function treat_constants

# liste des variables de chaque bloc
function create_list_of_variables_per_block(BLOCKS)
    DIFFERENT_BLOCKS = Set{String}(BLOCKS[:,1])
    block_var = Dict{String, Set{String}}()
    for block in DIFFERENT_BLOCKS
        block_var[block] = Set{String}()
    end
    for line=1:size(BLOCKS,1)
        block = BLOCKS[line,1]
        var = BLOCKS[line,2]
        block_var[block] = union(block_var[block],Set{String}([var]))
    end
    return block_var
end

function creating_and_solving_model(DIFFERENT_VAR, NB_BLOCKS, DIFFERENT_BLOCKS, block_var,CLIQUE_TREE, CONSTRAINTS, MATRICES, constants_per_constraint,my_solver, objscale, time_limit)
    #initialize
    m = Model(with_optimizer(my_solver))
    NB_CONSTRAINTS = size(CONSTRAINTS,1)
    jumpX = Array{LinearAlgebra.Symmetric{VariableRef,Array{VariableRef,2}}}(undef,NB_BLOCKS)
    coeff_block = Dict{Tuple{String,String}, Set{String}}()
    mat_var = Dict{Tuple{String,String}, Dict{String,Any}}()
    println("Dictionary linking blocks with their set of variables created")

    for (block, var_list) in block_var
        for var1 in var_list
            for var2 in var_list
                mat_var[(var1,var2)] = Dict{String,Any}()
                coeff_block[(var1,var2)] = Set{String}()
            end
        end
    end

    println("Dictionary linking coeff with their couple of variables created")
    @printf("mat_var size is %8d\n", length(coeff_block))
    size_block = Dict{String, Int64}()

    i_block = 0
    for (block, var_list) in block_var
        i_block += 1
        i_var1 = 1
        size_block[block] = length(var_list)
        jumpX[i_block] =  @variable(m,[1:size_block[block],1:size_block[block]], base_name= "X_$block", PSD)
        for var1 in var_list
            i_var2 = 1
            for var2 in var_list
                mat_var[(var1,var2)][block] = jumpX[i_block][i_var1, i_var2]
                coeff_block[(var1,var2)] = union(coeff_block[(var1,var2)],Set{String}([block]))
                i_var2+=1
            end
            i_var1+=1
        end
    end
    println("Block variables defined in the model m : ", i_block, " - ", length(block_var))

    println("Dictionary linking variables, blocks and data coefficients done")
    xp = Dict{String,JuMP.GenericAffExpr}()
    xp["OBJ"] = 0*jumpX[1][1,1]
    for p = 1:NB_CONSTRAINTS
        xp[CONSTRAINTS[p,1]] = 0*jumpX[1][1,1]
    end
    NB_TERMS = size(MATRICES,1)

    #NO MEAN version
    for term = 1:NB_TERMS
        objctr = MATRICES[term,1]
        var1 = MATRICES[term,2]
        var2 = MATRICES[term,3]
        val = MATRICES[term,4]
        nb_block = length(coeff_block[(var1,var2)])
        min_size = Inf
        min_var = first(mat_var[(var1,var2)])[2]
        for (block, var) in mat_var[(var1,var2)]
            if length(block_var[block]) < min_size
                min_size = length(block_var[block])
                min_var = var
            end
        end
        add_to_expression!(xp[objctr], val*min_var)
    end
    println("Objective and constraints scalar products computed")

    my_timer = @elapsed @objective(m, Min , objscale*(xp["OBJ"]+constants_per_constraint["OBJ"]))
    @printf("%-35s%10.6f s\n", "Objective added to model m", my_timer)
    println("Scaling objective function : factor or 10^-6")

    constraints_ref=Dict{String,JuMP.ConstraintRef}()
    for p = 1:NB_CONSTRAINTS
        name = CONSTRAINTS[p,1]
        if (CONSTRAINTS[p,2]=="EQ")
            constraints_ref[name] = @constraint(m,xp[name]+constants_per_constraint[name] == CONSTRAINTS[p,3])
        elseif (CONSTRAINTS[p,2]=="LEQ")
            constraints_ref[name] = @constraint(m,xp[name]+constants_per_constraint[name]<=CONSTRAINTS[p,4])
        elseif (CONSTRAINTS[p,2]=="GEQ")
            constraints_ref[name] = @constraint(m,xp[name]+constants_per_constraint[name]>=CONSTRAINTS[p,3])
        else
            constraints_ref[*(name,"_LEQ")] = @constraint(m,xp[name]+constants_per_constraint[name]<=CONSTRAINTS[p,4])
            constraints_ref[*(name,"_GEQ")] = @constraint(m,CONSTRAINTS[p,3]<=xp[name]+constants_per_constraint[name])
        end
    end
    println("initial constraints added to model m")
    #constraints linking common terms in blocks
    nb_coupling_constraints=0
    if NB_BLOCKS > 1
        for i in 1:size(CLIQUE_TREE,1)
            B1 = CLIQUE_TREE[i,1]
            B2 = CLIQUE_TREE[i,2]
            vars_B1 = block_var[B1]
            vars_B2 = block_var[B2]
            common_vars = [ var for var in intersect(vars_B1,vars_B2)]
            for i in 1:length(common_vars)
                var1 = common_vars[i]
                for j in i:length(common_vars)
                    var2 = common_vars[j]
                    JuMPvar1 = mat_var[(var1,var2)][B1]
                    JuMPvar2 = mat_var[(var1,var2)][B2]
                    constraints_ref[*("cc$(B1)_$(B2)_$var1","_$var2")]=@constraint(m,JuMPvar1-JuMPvar2==0)
                    nb_coupling_constraints+=1
                end
            end
        end
    end
    println("nb_coupling_constraints primal : ", nb_coupling_constraints)
    println("nb variables primal : ", length(JuMP.all_variables(m)))
    for ctr_type in list_of_constraint_types(m)
        println("nb constraints type $ctr_type primal : ", num_constraints(m, ctr_type[1], ctr_type[2]))
    end
    set_time_limit_sec(m, time_limit)
    start_time = Dates.now()
    val, time, bytes, gctime, memallocs = @timed optimize!(m)
    end_time = Dates.now()
    elapsed = Dates.value(end_time - start_time) * 0.001
    mem = bytes*10^(-6)
    println("Status : ", JuMP.termination_status(m))
    println("Julia resolution time : ", time)
    println("Memory consumption $mem Mo")
    println("\nObjective value: ", 1/objscale*JuMP.objective_value(m))

    Xoptimal=Array{Any}(nothing,NB_BLOCKS)
    for i=1:NB_BLOCKS
        Xoptimal[i]=JuMP.value.(jumpX[i])
    end
    return Xoptimal,JuMP.objective_value(m), JuMP.termination_status(m), time, mem
end

function sdp_model_solver_to_specify(cliques,
                                     cliquetree,
                                     path_opf_ctr,
                                     path_opf_mat,
                                     objscale;
                                     extract_name = false,
                                     time_limit = 100000.0)
    my_solver = Mosek.Optimizer
    #read data files
    @debug "    Reading data"
    my_timer=@elapsed (BLOCKS, CLIQUE_TREE, CONSTRAINTS, MATRICES_with_CONSTANTS) = read_data(cliques, cliquetree, path_opf_ctr, path_opf_mat)
    @printf("%-35s%10.6f s\n", "read_data", my_timer)
    #treat constants
    @debug "    Treating constants"
    my_timer = @elapsed (MATRICES, constants_per_objconstraint) = treat_constants(MATRICES_with_CONSTANTS, CONSTRAINTS)
    @printf("%-35s%10.6f s\n", "constant coeffcients treated", my_timer)
    #CONSTRAINTS
    @debug "    Constraints"
    DIFFERENT_CONSTRAINTS = Set{String}(CONSTRAINTS[:,1])
    NB_CONSTRAINTS = length(DIFFERENT_CONSTRAINTS)
    #variables
    @debug "    Variables"
    DIFFERENT_VAR = Set{String}(BLOCKS[:,2])
    NB_VAR = length(DIFFERENT_VAR)
    #define block patterns
    @debug "    Blocks patterns"
    DIFFERENT_BLOCKS = Set{String}(BLOCKS[:,1])
    NB_BLOCKS = length(DIFFERENT_BLOCKS)
    println("\n NB BLOCKS = $NB_BLOCKS \n")
    my_timer = @elapsed block_var = create_list_of_variables_per_block(BLOCKS)
    @printf("%-35s%10.6f s\n", "create_list_of_variables_per_block", my_timer)
    @debug "    Create and solve model"
    my_timer= @elapsed (Xoptimal,obj,status, time, mem) = creating_and_solving_model(DIFFERENT_VAR, NB_BLOCKS, DIFFERENT_BLOCKS, block_var, CLIQUE_TREE, CONSTRAINTS, MATRICES, constants_per_objconstraint, my_solver, objscale, time_limit)
    @printf("%-35s%10.6f s\n", "creating_and_solving_model", my_timer)
    return Xoptimal, obj, status, time, mem
end

max_stars_repo_name: charlyalizadeh/OPFdecgen

max_issues_repo_name: charlyalizadeh/OPFdecgen

max_forks_repo_name: charlyalizadeh/OPFdecgen


text:
######################################################################
# gaussian.jl
#
# Creates gaussian distributions over domains
######################################################################
export gaussian

gaussian(domain::Domain, dm::VF, ds::MF) = gaussian(domain, [dm], [ds])

"""
`gaussian(domain, means, covs, weights)`

Generates a distribution over `domain` from a weighted collection of Gaussians. The `means` argument is a vector of mean vectors; `covs` is a vector of covariance matrices; `weights` is a vector of real-valued numbers weighting each Gaussian. The weights need not add to 1, and the argument is optional; if `weights` is not provided, all Gaussians are weighted equally.

The resulting distribution is normalized.

The simple constructor `gaussian(domain, mean, cov)` allows you to construct a distribution from a single Gaussian.
"""
function gaussian{T<:Real}(domain::Domain, means::VVF, covs::VMF, weights::Vector{T}=ones(length(means)))
	# TODO: check that provided dimensions are correct
	w = float(weights)
	n = domain.num_dims
	if n == 2
		return gaussian2D(domain, means, covs, w)
	elseif n == 3
		return gaussian3D(domain, means, covs, w)
	end
end

function gaussian2D(domain::Domain, means::VVF, covs::VMF, weights::VF)
	d = zeros(x_cells(domain), y_cells(domain))
	d_sum = 0.0
	num_gaussians = length(means)

	inv_covs, dens = get_fast_tools(means, covs)

	for xi = 1:x_cells(domain)
		x = x_min(domain) + (xi-0.5)*x_size(domain)
		for yi = 1:y_cells(domain)
			y = y_min(domain) + (yi-0.5)*y_size(domain)
			for gi = 1:num_gaussians
				d[xi,yi] += weights[gi] * fast_pdf([x,y], means[gi], inv_covs[gi], dens[gi])
			end
			d_sum += d[xi,yi]
		end
	end
	normalize!(d, domain.cell_size)
	return d
end

function gaussian3D(domain::Domain, means::VVF, covs::VMF, weights::VF)
	d = zeros(x_cells(domain), y_cells(domain), z_cells(domain))
	d_sum = 0.0
	num_gaussians = length(means)

	# create inv_covs matrix, matrix of inverse covariants
	inv_covs, dens = get_fast_tools(means, covs)

	for xi = 1:x_cells(domain)
		x = x_min(domain) + (xi-0.5)*x_size(domain)
		for yi = 1:y_cells(domain)
			y = y_min(domain) + (yi-0.5)*y_size(domain)
			for zi = 1:z_cells(domain)
				z = z_min(domain) + (zi-0.5)*z_size(domain)
				for gi = 1:num_gaussians
					d[xi,yi,zi] += weights[gi] * fast_pdf([x,y,z], means[gi], inv_covs[gi], dens[gi])
				end
				d_sum += d[xi,yi,zi]
			end
		end
	end
	normalize!(d, domain.cell_size)
	return d
end

# Must compute Gaussian everywhere in domain
# To speed this up, we precompute inverse of cov matrix and denominator
# Each of these take a long time
function get_fast_tools(means::VVF, covs::VMF)
	num_gaussians = length(means)
	k = length(means[1])

	inv_covs = deepcopy(covs)
	dens = zeros(num_gaussians)
	for gi = 1:num_gaussians
		inv_covs[gi] = inv(covs[gi])
		dens[gi] = sqrt( (2.0pi)^k * det(covs[gi]) )
	end
	return inv_covs, dens
end


text:
push!(LOAD_PATH, "../src/")
using Documenter, DigitalComm

makedocs(sitename="DigitalComm.jl",
		 format = Documenter.HTML(),
		 pages    = Any[
						"Introduction to DigitalComm"   => "index.md",
						"Function list"          => "base.md",
						"Examples"                => Any[
														 "Examples/example_AWGN.md",
														 "Examples/example_BER.md",
														 "Examples/example_PSD.md",
														 ],
						],
		 );

#makedocs(sitename="My Documentation", format = Documenter.HTML(prettyurls = false))

deploydocs(
    repo = "github.com/JuliaTelecom/DigitalComm.jl",
)


text:
using JuLIP, QMMM2, SHIPs
using JuLIP.MLIPs
using NeighbourLists, LinearAlgebra

# fname = @__DIR__() * "/C_test_data.json"
fname = "/Users/ortner/Dropbox/Work/Projects/SiteEFitting/data/C_Dim3.json"
D = load_json(fname)
sym = :C  # => can get this from atoms object

trans = PolyTransform(2, rnn(sym))
basis = IPSuperBasis(OneBody(1.0),
                     SHIPBasis(3, 15, 1.5, trans, 2, 0.7*rnn(:C), 5.0) )
length(basis)
# weights = Dict("Es" => 100.0, "dEs" => 10.0, "d2Es" => 1.0, "d3Es" => 0.1)
weights = Dict("Es" => 0.1, "dEs" => 0.1, "d2Esh" => 10.0, "d3Esh" => 0.0)

##

🚢, fitinfo = QMMM2.lsqfit(basis, D, weights; rtol=1e-6)

QMMM2.print_errors(fitinfo)

dat = D["data"][3]
d2Esh = hcat(dat["d2Esh"]...)
at = Atoms(D["at"])
fit = QMMM2._site_energy_d2h(🚢, at, 1, dat["l"], dat["i"], dat["h"])
Ilge = findall( sqrt.(sum(abs2, d2Esh - fit; dims=1))[:] .> 0.15 )

display([d2Esh; fit][:, Ilge])

🚢

positions(at)[Ilge] |> mat
dat

norm(d2Esh - fit, Inf)


text:
struct NewtonSolver
    cost::Function
    dynamics::Function
    cE::Function
    cI::Function

    ind1::NamedTuple
    ind2::NamedTuple
    sizes::NamedTuple
    pack::Function
    unpack::Function
    opts::Dict{Symbol,Any}
end

function NewtonSolver(solver::Solver)
    p,pI,pE = get_num_constraints(solver)
    p_N,pI_N,pE_N = get_num_terminal_constraints(solver)
    n,m,N = get_sizes(solver)

    c_function!, c_jacobian!, c_labels, cI!, cE! = generate_constraint_functions(solver.obj)

    Nx = N*n
    Nu = (N-1)*m
    Nz = Nx + Nu
    Nh = (N-1)*pE + pE_N
    Ng = (N-1)*pI + pI_N
    NN = 2Nx + Nu + Nh + Ng

    names = (:z,:ν,:λ,:μ)
    ind1 = TrajectoryOptimization.create_partition([Nz,Nx,Nh,Ng],names)
    ind2 = TrajectoryOptimization.create_partition2([Nz,Nx,Nh,Ng],names)
    indx = reshape(1:Nx,n,N)
    indu = reshape(1:Nu,m,N-1) .+ Nx

    sizes = (Nx=Nx,Nu=Nu,Nz=Nz,Nh=Nh,Ng=Ng,NN=NN,pI=pI,pE=pE,pI_N=pI_N,pE_N=pE_N,n=n,m=m,N=N,
        x=indx, u=indu)

    function mycost(Z)
        X,U = unpackZ(Z,sizes)
        cost(solver,X,U)
    end

    function cE(Z)
        X,U = unpackZ(Z,sizes)

        C = zeros(eltype(Z),pE,N-1)
        for k = 1:N-1
            cE!(view(C,1:pE,k),X[:,k],U[:,k])
        end
        CN = zeros(eltype(Z),pE_N)
        cE!(CN,X[:,N])
        CE = [vec(C); CN]
        return CE
    end

    function cI(Z)
        X,U = unpackZ(Z,sizes)

        C = zeros(eltype(Z),pI,N-1)
        for k = 1:N-1
            cI!(view(C,1:pI,k),X[:,k],U[:,k])
        end
        CN = zeros(eltype(Z),pI_N)
        cI!(CN,X[:,N])
        CI = [vec(C); CN]
        return CI
    end

    function dynamics(Z)
        X,U = unpackZ(Z,sizes)

        D = zeros(eltype(Z),n,N)
        D[:,1] = X[:,1] - solver.obj.x0
        for k = 2:N
            solver.fd(view(D,1:n,k),X[:,k-1],U[:,k-1])
            D[:,k] -= X[:,k]
        end
        return vec(D)
    end

    function unpack(V::AbstractVector)
        if length(V) == NN
            Z = view(V,ind1.z)
            ν = view(V,ind1.ν)
            λ = view(V,ind1.λ)
            μ = view(V,ind1.μ)
            return Z,ν,λ,μ
        elseif length(V) == Nz
            X = V[indx]
            U = V[indu]
        end
    end

    function pack(Z::Vector,ν::Vector,λ::Vector,μ::Vector)
        V = zeros(NN)
        V[ind1.z] = Z
        V[ind1.ν] = ν
        V[ind1.λ] = λ
        V[ind1.μ] = μ
        return V
    end

    ix = 1:Nx
    iu = Nx .+ (1:Nu)
    function pack(X::Matrix,U::Matrix)
        Z = zeros(Nz)
        Z[ix] = vec(X)
        Z[iu] = vec(U)
        return Z
    end
    pack(X::Matrix, U::Matrix, ν::T, λ::T, μ::T) where T <: Vector = pack(pack(X,U),ν,λ,μ)

    opts = Dict(:type=>:kkt,                # Type of hessian to assemble
                :ϵ_as=>1e-3,                # Active set tolerance
                :iters_linesearch=>10,      # Number of iterations to use in the linesearch
                :refinement=>:jacobian,     # Which refinement method to use
                :η=>0.1                     # Gradient descent criteria
                )

    NewtonSolver(mycost,dynamics,cE,cI,
        ind1,ind2,sizes,pack,unpack,opts)
end

function get_sizes(solver::NewtonSolver)
    solver.sizes.n, solver.sizes.m, solver.sizes.N
end

function get_num_constraints(solver::NewtonSolver)
    pI = solver.sizes.pI
    pE = solver.sizes.pE
    return pI+pE,pI,pE
end

function get_num_terminal_constraints(solver::NewtonSolver)
    pI = solver.sizes.pI_N
    pE = solver.sizes.pE_N
    return pI+pE,pI,pE
end

function max_violation(solver::NewtonSolver,V::AbstractVector,use_dynamics::Bool=true)
    c_max = max(solver.sizes.Ng > 0 ? maximum(solver.cI(V)) : 0, norm(solver.cE(V),Inf))
    if use_dynamics
        c_max = max(norm(solver.dynamics(V),Inf),c_max)
    end
    return c_max
end

function unpackZ(Z::AbstractVector, inds::NamedTuple)
    X = view(Z,inds.x)
    U = view(Z,inds.u)
    return X,U
end
unpackZ(Z::AbstractVector, solver::NewtonSolver) = unpackZ(Z,solver.sizes)

function create_V(solver::NewtonSolver,res::SolverIterResults)
    p,pI,pE = get_num_constraints(solver)
    p_N,pI_N,pE_N = get_num_terminal_constraints(solver)
    n,m,N = get_sizes(solver)

    x = to_array(res.X)
    u = to_array(res.U)
    ν = zeros(solver.sizes.Nx)

    if res isa ConstrainedVectorResults
        μ = zeros(pI,N-1)
        λ = zeros(pE,N-1)
        ineq = 1:pI
        eq = pI .+ (1:pE)
        for k = 1:N-1
            μ[:,k] = res.λ[k][ineq]
            λ[:,k] = res.λ[k][eq]
        end
        μ = vec(μ)
        λ = vec(λ)
        append!(μ, res.λ[N][1:pI_N])
        append!(λ, res.λ[N][pI_N .+ (1:pE_N)])
    else
        λ = Float64[]
        μ = Float64[]
    end
    return solver.pack(x,u,ν,λ,μ)
end

function gen_newton_functions(solver::NewtonSolver)
    ind1 = solver.ind1
    ind2 = solver.ind2

    cE,cI,dynamics = solver.cE, solver.cI, solver.dynamics
    mycost = solver.cost

    sizes = solver.sizes
    Nx = sizes.Nx
    Nz = sizes.Nz
    Nh = sizes.Nh
    Ng = sizes.Ng

    function buildKKT(V, ρ, type)
        Z, ν, λ, μ = solver.unpack(V)
        X,U = unpackZ(Z, solver)

        ∇²J = ForwardDiff.hessian(mycost,Z)
        ∇J = ForwardDiff.gradient(mycost,Z)
        D = ForwardDiff.jacobian(dynamics,Z)
        H = ForwardDiff.jacobian(cE,Z)
        G = ForwardDiff.jacobian(cI,Z)
        d = dynamics(Z)
        h = cE(Z)
        g = cI(Z)

        if type == :penalty
            A = [∇²J   D'         H'      G';
                 D   -1/ρ*I       zeros(Nx,Nh+Ng) ;
                 H   zeros(Nh,Nx) -1/ρ*I zeros(Nh,Ng);
                 G   zeros(Ng,Nx) zeros(Ng,Nh) -1/ρ*I]
            b = [∇J + D'nu + H'λ + G'μ; d; h; g]
        elseif type == :kkt
            A = [∇²J   D'    H'                G';
                 D   zeros(Nx,Nx) zeros(Nx,Nh) zeros(Nx,Ng);
                 H   zeros(Nh,Nx) zeros(Nh,Nh) zeros(Nh,Ng);
                 G   zeros(Ng,Nx) zeros(Ng,Nh) zeros(Ng,Ng)]
            b = [∇J + D'ν + H'λ + G'μ; d; h; g]
        end
        return A,b
    end

    a = ones(Bool,solver.sizes.NN)
    function active_set(Z,eps=0)
        X,U = unpackZ(Z, solver)
        ci = cI(Z)
        c_inds = ones(Bool,solver.sizes.Ng)
        c_inds = ci .>= -eps

        a[ind1.μ] = c_inds
        return a
    end

    function newton_step(V,ρ,meritfun::Function,type=:kkt; projection=:none,
        reg=Diagonal(zero(V)), η=solver.opts[:η])

        # Initial cost
        J0 = meritfun(V)
        V_ = copy(V)

        # Build and solve KKT
        A,b = buildKKT(V_,ρ,type)
        a = active_set(V_,1e-3)
        amu = a[ind1.μ]
        δV = zero(V)
        Ā = A[a,a] + reg[a,a]
        δV[a] = -Ā\b[a]

        if projection == :jacobian
            V1 = V + δV
            Z1 = V1[ind1.z]
            Z = V[ind1.z]
            D = ForwardDiff.jacobian(dynamics,Z)
            H = ForwardDiff.jacobian(cE,Z)
            G = ForwardDiff.jacobian(cI,Z)
            Y = [D;H;G[amu,:]]
            d1 = dynamics(Z1)
            h1 = cE(Z1)
            g1 = cI(Z1)
            y = [d1;h1;g1[amu]]
            # dv̂ = -D'*((D*D')\d1)
            δV̂ = -Y'*((Y*Y')\y)
            δV[ind1.z] += δV̂
        elseif projection == :constraints
            b̂ = copy(b)
            for j = 1:5
                V1 = V_ + δV
                d1 = dynamics(V1)
                h1 = cE(V1)
                g1 = cI(V1)
                b̂[ind1.ν] = d1
                b̂[ind1.λ] = h1
                b̂[ind1.μ] = g1
                δV[a] = -Ā\b̂[a]
                V_ = V_ + δV
            end
        end

        # Line Search
        α = armijo_line_search(meritfun,V,δV,b, max_iter=solver.opts[:iters_linesearch])
        V_ = V_ + α*δV

        stats = (α=α,condnum=cond(Ā),grad=norm(b))
        return V_, stats
    end

    return newton_step, buildKKT, active_set
end

function armijo_line_search(merit::Function,V,d,grad; max_iter=10, ϕ=0.01)
    α = 1
    J_prev = merit(V)
    J = merit(V+α*d)
    iter = 1
    while J > J_prev + α*ϕ*grad'd
        α *= 0.75
        J = merit(V + α*d)
        if iter > max_iter
            println("Max iterations")
            α = 0
            break
        end
        iter += 1
    end
    return α
end

function meritfunction(solver::NewtonSolver,V,ρ)
    eps = solver.opts[:ϵ_as]
    Z, ν, λ, μ = solver.unpack(V)
    d = solver.dynamics(Z)
    h = solver.cE(Z)
    g = solver.cI(Z)
    a = g .> eps
    solver.cost(Z) + ν'd + λ'h + μ'g + ρ/2*sqrt(d'd + h'h + g[a]'g[a])
end

function solve(solver::NewtonSolver,V::Vector, ρ=0; iters=1, verbose=true, meritfun=V->meritfunction(solver,V,ρ), projection=:none)

    newton_step, buildKKT, active_set = gen_newton_functions(solver)
    change(x,x0) = format((x0-x)/x0*100,precision=2) * "%"

    # Stats
    cost = zeros(iters)
    grad = zeros(iters)
    c_max = zeros(iters)
    c_max0 = zeros(iters)  # max violation without dynamics

    J_prev = meritfun(V)
    for i = 1:iters
        V, step_stats = newton_step(V,ρ,meritfun,projection=projection)
        J = meritfun(V)
        grad[i] = step_stats.grad
        c_max[i] = max_violation(solver,V)
        c_max0[i] = max_violation(solver,V,false)

        if verbose
            println("iter $i")
            println("  cost: $J $(change(J,J_prev))")
            println("  grad: $(grad[i])")
            println("  c_max: $(c_max[i])")
            println("  c_max0: $(c_max0[i]))")
            println("  α: $(step_stats.α)")
            println("  cond: $(step_stats.condnum)")
            println()
        end
        J_prev = J
    end
    return V
end
function solve(solver::NewtonSolver,res::ConstrainedIterResults; kwargs...)
    ρ = max_penalty(res)
    meritfun(V) = meritfunction(solver,V,ρ)
    solve(solver,create_V(solver,res),ρ)
end

# model, obj = Dynamics.pendulum
# solver = Solver(model, obj, N=21)
# res, stats = solve(solver)
#
# nsolver = NewtonSolver(solver)
# meritfun(V) = meritfunction(nsolver,V,0)
# solve(nsolver,res,iters=2,meritfun=meritfun)


text:
function bench_iteredges(g::AbstractGraph)
  i = 0
  for e in edges(g)
    i += 1
  end
  return i
end

function bench_has_edge(g::AbstractGraph)
    srand(1)
    nvg = nv(g)
    srcs = rand([1:nvg;], cld(nvg, 4))
    dsts = rand([1:nvg;], cld(nvg, 4))
    i = 0
    for (s, d) in zip(srcs, dsts)
        if has_edge(g, s, d)
            i += 1
        end
    end
    return i
end

EDGEFNS = [
    bench_iteredges,
    bench_has_edge
]

@benchgroup "edges" begin

  for fun in EDGEFNS
    @benchgroup "$fun" begin
      @benchgroup "graph" begin
        for (name, g) in GRAPHS
          @bench "$name" $fun($g)
        end
      end
      @benchgroup "digraph" begin
        for (name, g) in DIGRAPHS
          @bench "$name" $fun($g)
        end
      end # digraph
    end # fun
  end
end # edges

max_stars_repo_name: AsileBcd/JohnsonCycles

max_issues_repo_name: AsileBcd/JohnsonCycles

max_forks_repo_name: AsileBcd/JohnsonCycles


text:
# TODO: make bindings to KrylovKit eigenvalue solvers, probably via using LUfacts
module KryloveKitInterface

using KrylovKit
using ..Common

import ..AbstractEigenproblem
import ..AbstractLinearEigenproblem
import ..AbstractCFEigenproblem
import ..AbstractNonlinearEigenproblem
import ..HelmholtzProblem
import ..MaxwellProblem
import ..DEFAULT_LINEAR_EIGENSOLVER
import ..HelmholtzLEP
import ..MaxwellLEP

end

using .KryloveKitInterface


text:
module TokamakSmallCartesian

    import ElectromagneticFields.AxisymmetricTokamakCartesian

    export pauli_particle_3d_pode, hamiltonian, toroidal_momentum

    const equ = AxisymmetricTokamakCartesian.@code() # inject magnetic field code

    const qᵢ = [1.05,   0.0,    0.0]
    const vᵢ = [2.1E-3, 4.3E-4, 0.0]

    include("pauli_particle_3d.jl")

end


text:
# --------------------------------------------------------------------------
# ACE.jl and SHIPs.jl: Julia implementation of the Atomic Cluster Expansion
# Copyright (c) 2019 Christoph Ortner <christophortner0@gmail.com>
# All rights reserved.
# --------------------------------------------------------------------------

import ACE: standardevaluator, graphevaluator
using SparseArrays: SparseMatrixCSC, sparse
using LinearAlgebra: mul!

"""
`struct RPIBasis` formerly RPIBasis
"""

struct RPIBasis{T, BOP, NZ, TIN} <: IPBasis #struct RPIBasis{T, BOP, NZ, TIN,L} <: IPBasis
   pibasis::PIBasis{BOP, NZ, TIN}
   A2Bmaps::NTuple{NZ, SparseMatrixCSC{T, Int}}
   Bz0inds::NTuple{NZ, UnitRange{Int}}
   L::Int
   #valL::Val{L}
end
#getL(::RPIBasis{T, BOP, NZ, TIN,L}) where { R,BOP,NZ, TIN, L} = L
RPIBasis(pibasis,A2Bmaps, Bz0inds) = RPIBasis(pibasis,A2Bmaps, Bz0inds,0)

Base.length(basis::RPIBasis, iz0::Integer) = size(basis.A2Bmaps[iz0], 1)

Base.length(basis::RPIBasis) = sum(length(basis, iz0)
                                    for iz0 = 1:numz(basis.pibasis))

fltype(::RPIBasis{T}) where {T}  = T

zlist(basis::RPIBasis) = zlist(basis.pibasis)

cutoff(basis::RPIBasis) = cutoff(basis.pibasis)

standardevaluator(basis::RPIBasis) =
      RPIBasis( standardevaluator(basis.pibasis),
                basis.A2Bmaps, basis.Bz0inds )

graphevaluator(basis::RPIBasis) =
      RPIBasis( graphevaluator(basis.pibasis),
                basis.A2Bmaps, basis.Bz0inds )

# ------------------------------------------------------------------------
#    FIO code
# ------------------------------------------------------------------------

==(B1::RPIBasis, B2::RPIBasis) = (B1.pibasis == B2.pibasis)

write_dict(basis::RPIBasis) = Dict(
      "__id__" => "ACE_RPIBasis",
      "__v__"  => "v0_8_2",
      "pibasis" => write_dict(basis.pibasis),
      "A2Bmaps" => write_dict.(basis.A2Bmaps),
      "Bz0inds" => [ [ur.start, ur.stop] for ur in basis.Bz0inds ]
   )

# v0.8.2 onwards
function read_dict(::Val{:ACE_RPIBasis}, ::Val{:v0_8_2}, D::Dict)
   pibasis = read_dict(D["pibasis"])
   A2Bmaps = tuple( read_dict.(D["A2Bmaps"])... )
   Bz0inds = tuple( [ ur[1]:ur[2] for ur in D["Bz0inds"] ]... )
   return RPIBasis(pibasis, A2Bmaps, Bz0inds)
end

# old version
read_dict(::Val{:ACE_RPIBasis}, D::Dict) =
   RPIBasis(read_dict(D["pibasis"]))

read_dict(::Val{:SHIPs_RPIBasis}, D::Dict) =
   read_dict(Val{:ACE_RPIBasis}(), D)

# ------------------------------------------------------------------------
#    Basis construction code
# ------------------------------------------------------------------------

RPIBasis(basis1p::OneParticleBasis, N::Integer,
         D::AbstractDegree, maxdeg::Real; L=0, constants=false) =
   RPIBasis(PIBasis(basis1p, N, D, maxdeg; filter = RPIFilter{L}(constants)),L)

function RPIBasis(pibasis::PIBasis, L=0)
   basis1p = pibasis.basis1p

   # construct the cg matrices
   rotc = Rot3DCoeffs(L)
   A2Bmaps = ntuple(iz0 -> _rpi_A2B_matrix(rotc, pibasis, iz0), numz(pibasis)) ###

   # construct the indices within the B vector to which the A2Bmaps map.
   Bz0inds = UnitRange{Int}[]
   idx0 = 0
   for i = 1:length(A2Bmaps)
      len = size(A2Bmaps[i], 1)
      push!(Bz0inds, (idx0+1):(idx0+len))
      idx0 += len
   end

   return RPIBasis(pibasis, A2Bmaps, tuple(Bz0inds...),L) # Val(L) or Val{L}()
end

# TODO NOW: graphevaluator, standardevaluator

struct RPIFilter{L}
   constants::Bool
end

(f::RPIFilter)(pib::PIBasisFcn{0}) = f.constants
(f::RPIFilter{L})(pib::PIBasisFcn{1}) where {L} = (pib.oneps[1].l == L)
(f::RPIFilter{L})(pib::PIBasisFcn) where {L} = (
      iseven( L + sum(b.l for b in pib.oneps) ) && # for L odd we have covariance w.r.t. inversion
      (abs(sum(b.m for b in pib.oneps)) <= L) )

# _rpi_filter(pib::PIBasisFcn{0}) = false
# _rpi_filter(pib::PIBasisFcn{1}) = (pib.oneps[1].l == 0)
# _rpi_filter(pib::PIBasisFcn) = (
#       iseven( sum(b.l for b in pib.oneps) ) &&
#       (sum(b.m for b in pib.oneps) == 0) )

function _rpi_A2B_matrix(rotc::Rot3DCoeffs{T,L},
                         pibasis::PIBasis,
                         iz0) where {T,L}
   # allocate triplet format
   Irow, Jcol = Int[], Int[]

   if L == 0
      vals = real(fltype(pibasis.basis1p))[]
   elseif L==1
      vals = SVector{3,Complex{T}}[]
   end
   #real(fltype(pibasis.basis1p))[]
   # count the number of PI basis functions = number of rows
   idxB = 0
   # loop through all (zz, kk, ll) tuples; each specifies 1 to several B
   for i = 1:length(pibasis.inner[iz0])
      # get the specification of the ith basis function
      pib = get_basis_spec(pibasis, iz0, i)
      # skip it unless all m are zero, because we want to consider each
      # (nn, ll) block only once.
      if !all(b.m == 0 for b in pib.oneps)
         continue
      end
      # get the rotation-coefficients for this basis group
      # the bs are the basis functions corresponding to the columns
      U, bcols = _rpi_coupling_coeffs(pibasis, rotc, pib)
      # loop over the rows of Ull -> each specifies a basis function
      for irow = 1:size(U, 1)
         idxB += 1
         # loop over the columns of U / over brows
         for (icol, bcol) in enumerate(bcols)
            # this is a subtle step: bcol and bcol_ordered are equivalent
            # permutation-invariant basis functions. This means we will
            # add the same PI basis function several times, but in the call to
            # `sparse` the values will just be added.
            bcol_ordered = ACE._get_ordered(pibasis, bcol)
            idxAA = pibasis.inner[iz0].b2iAA[bcol_ordered]
            push!(Irow, idxB)
            push!(Jcol, idxAA)
            push!(vals, U[irow, icol])
         end
      end
   end
   # create CSC: [   triplet    ]  nrows   ncols
   return sparse(Irow, Jcol, vals, idxB, length(pibasis.inner[iz0]))
end

# U, bcols = rpi_coupling_coeffs(rotc, pib)

"""
this is essentially a wrapper function around Rotations3D.rpi_basis,
and is just meant to translate between different representations
"""
function _rpi_coupling_coeffs(pibasis, rotc::Rot3DCoeffs, pib::PIBasisFcn{N}
                              ) where {N}
   # convert to zz, ll, nn tuples
   zz, nn, ll, _ = _b2znlms(pib)
   # construct the RPI coupling coefficients
   U, Ms = Rotations3D.rpi_basis(rotc, zz, nn, ll)
   # convert the Ms into basis functions
   rpibs = [ _znlms2b(zz, nn, ll, mm, pib.z0) for mm in Ms ]
   return U, rpibs
end

_rpi_coupling_coeffs(pibasis, rotc::Rot3DCoeffs, pib::PIBasisFcn{0}) =
      [ 1.0 ], []

_b2znlms(pib::PIBasisFcn{N}) where {N} = (
   SVector(ntuple(n -> pib.oneps[n].z, N)...),
   SVector(ntuple(n -> pib.oneps[n].n, N)...),
   SVector(ntuple(n -> pib.oneps[n].l, N)...),
   SVector(ntuple(n -> pib.oneps[n].m, N)...) )

_znlms2b(zz, nn, ll, mm = zero(ll), z0 = AtomicNumber(0)) =
   PIBasisFcn( z0, ntuple(i -> PSH1pBasisFcn(nn[i], ll[i], mm[i], zz[i]),
                          length(zz)) )

function combine(basis::RPIBasis, coeffs)
   picoeffs = ntuple(iz0 -> (coeffs[basis.Bz0inds[iz0]]' * basis.A2Bmaps[iz0])[:],
                     numz(basis.pibasis))
   return PIPotential(basis.pibasis, picoeffs)
end

function scaling(basis::RPIBasis, p)
   wwpi = scaling(basis.pibasis, p)
   wwrpi = zeros(Float64, length(basis))
   for iz0 = 1:numz(basis)
      wwpi_iz0 = wwpi[basis.pibasis.inner[iz0].AAindices]
      wwrpi[basis.Bz0inds[iz0]] = basis.A2Bmaps[iz0] * wwpi_iz0
   end
   return wwrpi
end

# ------------------------------------------------------------------------
#    Evaluation code
# ------------------------------------------------------------------------

alloc_temp(basis::RPIBasis, args...) =
   ( AA = site_alloc_B(basis.pibasis, args...),
   #AAr = real(site_alloc_B(basis.pibasis, args...)),
     AAr = site_alloc_B(basis.pibasis, args...),
     tmp_pibasis = alloc_temp(basis.pibasis, args...)
   )

function evaluate!(B, tmp, basis::RPIBasis, Rs, Zs, z0)
   iz0 = z2i(basis, z0)
   AA = site_evaluate!(tmp.AA, tmp.tmp_pibasis, basis.pibasis, Rs, Zs, z0)
   # TODO: this could be done better maybe by adding the real function into
   #       site_evaluate!, or by writing a real version of it...
   AAr = @view tmp.AAr[1:length(AA)]
   AAr[:] .= AA
   #AAr[:] .= real.(AA)
   Bview = @view B[basis.Bz0inds[iz0]]

   mul!(Bview, basis.A2Bmaps[iz0], AAr)
   return real.(B)
end

# ------- gradient

alloc_temp_d(basis::RPIBasis, Rs::AbstractVector, args...) =
   alloc_temp_d(basis, length(Rs))

alloc_temp_d(basis::RPIBasis, nmax::Integer) =
    (
    AA = site_alloc_B(basis.pibasis),
    AAr = real(site_alloc_B(basis.pibasis)),
    tmp_pibasis = alloc_temp(basis.pibasis, nmax),
    dAA = site_alloc_dB(basis.pibasis, nmax),
    tmpd_pibasis = alloc_temp_d(basis.pibasis, nmax),
    )

# TODO: evaluate also B??? the interface seems to command it.
function evaluate_d!(B, dB, tmpd, basis::RPIBasis, Rs, Zs, z0)
   iz0 = z2i(basis, z0)
   AA, dAA = tmpd.AA, tmpd.dAA
   site_evaluate_d!(AA, dAA, tmpd.tmpd_pibasis, basis.pibasis, Rs, Zs, z0)
   len = length(basis.pibasis.inner[iz0])
   for i = 1:len
      AA[i] = real(AA[i])
   end
   mul!((@view B[basis.Bz0inds[iz0]]), basis.A2Bmaps[iz0], (@view AA[1:len]))
   for j = 1:length(Rs)
      # ∂∏A / ∂𝐫ⱼ
      dAAj = @view dAA[1:len, j]
      for i = 1:length(dAAj)
         # TODO: we should incorporate the `real` operation into the PIBasis?
         @inbounds dAAj[i] = real.(dAAj[i])
      end
      # copy into B
      dBview = @view dB[basis.Bz0inds[iz0], j]
      mul!(dBview, basis.A2Bmaps[iz0], dAAj)
   end
   return dB
end


text:
using LinearAlgebra, SparseArrays, Statistics
using Logging
using Test


text:
include("1.jl")
include("11.jl")
include("14.jl")
include("16.jl")
function updowntwicestep1111(r::Robot,side1::HorizonSide)
    angleofsquare1111(r)
    putmarker!(r)
    putmarkholeper(r,side1)
    while !isborder(r,Nord)
        disexit(r,Nord)
        putmarkholeper(r,inverse(side1))
        backingper(r, side1)
        disexit(r,inverse(side1))
    end
end

max_stars_repo_name: bubasooo/proggggggggg

max_issues_repo_name: bubasooo/proggggggggg

max_forks_repo_name: bubasooo/proggggggggg