text:
struct NonconvexADWrapper{F,P}
    f::F
    prob::P
end

(f::NonconvexADWrapper)(θ) = begin
    x = f.f.f(θ, f.prob.p)
    return first(x)
end

## Create ChainRule
function ChainRulesCore.rrule(f::NonconvexADWrapper, θ::AbstractVector)
    val =  f.f.f(θ, f.prob.p)
    G = similar(θ)
    f.f.grad(G, θ)
    val, Δ -> (ChainRulesCore.NoTangent(), Δ * G)
end

function separate_kwargs(orig_kwargs)
    sub_options = :sub_options .∈ Ref(keys(orig_kwargs)) ? orig_kwargs[:sub_options] : nothing
    convergence_criteria = :convcriteria .∈ Ref(keys(orig_kwargs)) ? orig_kwargs[:convcriteria] : nothing
    opt_kwargs_names = keys(orig_kwargs)[[(keys(orig_kwargs) .∉ Ref([:sub_options, :convcriteria]))...]]
    opt_kwargs = (; zip(opt_kwargs_names, [orig_kwargs[j] for j in opt_kwargs_names])...)

    return opt_kwargs, sub_options, convergence_criteria
end

function __map_optimizer_args(prob::OptimizationProblem, opt::Nonconvex.NonconvexCore.AbstractOptimizer;
    cb=nothing,
    maxiters::Union{Number, Nothing}=nothing,
    maxtime::Union{Number, Nothing}=nothing,
    abstol::Union{Number, Nothing}=nothing,
    reltol::Union{Number, Nothing}=nothing,
    integer=nothing,
    kwargs...)

    opt_kwargs, sub_options, convergence_criteria = separate_kwargs(kwargs)
    opt_kwargs = convert_common_kwargs(opt, opt_kwargs, cb=cb, maxiters=maxiters, maxtime=maxtime, abstol=abstol, reltol=reltol)
    mapped_args = _create_options(opt, opt_kwargs=opt_kwargs, sub_options=sub_options, convergence_criteria=convergence_criteria)

    integer = isnothing(integer) ? fill(false, length(prob.u0)) : integer

    return mapped_args, integer
end

function __solve(prob::OptimizationProblem, opt::Nonconvex.NonconvexCore.AbstractOptimizer;
                 cb = nothing,
                 maxiters::Union{Number, Nothing} = nothing,
                 maxtime::Union{Number, Nothing} = nothing,
                 abstol::Union{Number, Nothing}=nothing,
                 reltol::Union{Number, Nothing}=nothing,
                 progress = false,
                 surrogate_objective::Union{Symbol, Nothing}=nothing,
                 integer=nothing,
                 kwargs...)

    # local x

    maxiters = _check_and_convert_maxiters(maxiters)
    maxtime = _check_and_convert_maxtime(maxtime)

    f = instantiate_function(prob.f,prob.u0,prob.f.adtype,prob.p)

    _loss = NonconvexADWrapper(f, prob)

    opt_f(θ) =_loss(θ)

    opt_args, integer = _map_optimizer_args(prob, opt, cb=cb, maxiters=maxiters, maxtime=maxtime,abstol=abstol, reltol=reltol, integer=integer; kwargs...)

    opt_set = Nonconvex.Model()
    if isnothing(surrogate_objective)
        Nonconvex.set_objective!(opt_set, opt_f)
    else
        Nonconvex.set_objective!(opt_set, opt_f, flags=[surrogate_objective])
    end

    Nonconvex.addvar!(opt_set, prob.lb, prob.ub, integer = integer)

    if !isnothing(prob.f.cons)
        @warn "Equality constraints are current not passed on by GalacticOptim"
        #add_ineq_constraint!(opt_set, f)
    end

    if !isnothing(prob.lcons)
        @warn "Inequality constraints are current not passed on by GalacticOptim"
        #add_ineq_constraint!(opt_set, f)
    end

    if !isnothing(prob.ucons)
        @warn "Inequality constraints are current not passed on by GalacticOptim"
        #add_ineq_constraint!(opt_set, f)
    end

    t0 = time()
    if check_optimizer_backend(opt)
        opt_res = Nonconvex.optimize(opt_set, opt; opt_args...)
    else
        opt_res = Nonconvex.optimize(opt_set, opt, prob.u0; opt_args...)
    end
    t1 = time()

    opt_ret = hasproperty(opt_res,:status) ? Symbol(string(opt_res.status)) : nothing

    SciMLBase.build_solution(prob, opt, opt_res.minimizer, opt_res.minimum; (isnothing(opt_ret) ? (; original=opt_res) : (; original=opt_res, retcode=opt_ret))...)
end


text:
"""
    mm(X,v,K) = v*X / (X + K)

A Michaelis-Menten rate function.
"""
mm(X,v,K) = v*X / (X + K)
@register mm(X,v,K);
Symbolics.derivative(::typeof(mm), args::NTuple{3,Any}, ::Val{1}) = (args[2]*args[3]) / (args[1]+args[3])^2
Symbolics.derivative(::typeof(mm), args::NTuple{3,Any}, ::Val{2}) = args[1]/(args[1]+args[3])
Symbolics.derivative(::typeof(mm), args::NTuple{3,Any}, ::Val{3}) = - args[2]*args[1]/(args[1]+args[3])^2

# Registers the repressing Michaelis-Menten function.
"""
    mmr(X,v,K) = v*K / (X + K)

A repressive Michaelis-Menten rate function.
"""
mmr(X,v,K) = v*K / (X + K)
@register mmr(X,v,K);
Symbolics.derivative(::typeof(mmr), args::NTuple{3,Any}, ::Val{1}) = - (args[2]*args[3]) / (args[1]+args[3])^2
Symbolics.derivative(::typeof(mmr), args::NTuple{3,Any}, ::Val{2}) = args[3]/(args[1]+args[3])
Symbolics.derivative(::typeof(mmr), args::NTuple{3,Any}, ::Val{3}) = args[2]*args[1]/(args[1]+args[3])^2

"""
    hill(X,v,K,n) = v*(X^n) / (X^n + K^n)

A Hill rate function.
"""
hill(X,v,K,n) = v*(X^n) / (X^n + K^n)
@register hill(X,v,K,n);
Symbolics.derivative(::typeof(hill), args::NTuple{4,Any}, ::Val{1}) = args[2] * args[4] * (args[3]^args[4]) * (args[1]^(args[4]-1))  /  (args[1]^args[4] + args[3]^args[4])^2
Symbolics.derivative(::typeof(hill), args::NTuple{4,Any}, ::Val{2}) = (args[1]^args[4])  /  (args[1]^args[4] + args[3]^args[4])
Symbolics.derivative(::typeof(hill), args::NTuple{4,Any}, ::Val{3}) = - args[2] * args[4] * (args[1]^args[4]) * (args[3]^(args[4]-1))  /  (args[1]^args[4] + args[3]^args[4])^2
Symbolics.derivative(::typeof(hill), args::NTuple{4,Any}, ::Val{4}) = args[2] * (args[1]^args[4]) * (args[3]^args[4]) * (log(args[1])-log(args[3]))  /  (args[1]^args[4] + args[3]^args[4])^2

"""
    hillr(X,v,K,n) = v*(K^n) / (X^n + K^n)

A repressive Hill rate function.
"""
hillr(X,v,K,n) = v*(K^n) / (X^n + K^n)
@register hillr(X,v,K,n);
Symbolics.derivative(::typeof(hillr), args::NTuple{4,Any}, ::Val{1}) = - args[2] * args[4] * (args[3]^args[4]) * (args[1]^(args[4]-1))  /  (args[1]^args[4] + args[3]^args[4])^2
Symbolics.derivative(::typeof(hillr), args::NTuple{4,Any}, ::Val{2}) = (args[3]^args[4])  /  (args[1]^args[4] + args[3]^args[4])
Symbolics.derivative(::typeof(hillr), args::NTuple{4,Any}, ::Val{3}) = args[2] * args[4] * (args[1]^args[4]) * (args[3]^(args[4]-1))  /  (args[1]^args[4] + args[3]^args[4])^2
Symbolics.derivative(::typeof(hillr), args::NTuple{4,Any}, ::Val{4}) = args[2] * (args[1]^args[4]) * (args[3]^args[4]) * (log(args[3])-log(args[1]))  /  (args[1]^args[4] + args[3]^args[4])^2

"""
    hillar(X,Y,v,K,n) = v*(X^n) / (X^n + Y^n + K^n)

An activation/repressing Hill rate function.
"""
hillar(X,Y,v,K,n) = v*(X^n) / (X^n + Y^n + K^n)
@register hillar(X,Y,v,K,n);
Symbolics.derivative(::typeof(hillar), args::NTuple{5,Any}, ::Val{1}) = args[3] * args[5] * (args[1]^(args[5]-1)) * (args[2]^args[5]+args[4]^args[5])  /  (args[1]^args[5] + args[2]^args[5] + args[4]^args[5])^2
Symbolics.derivative(::typeof(hillar), args::NTuple{5,Any}, ::Val{2}) = - args[3] * args[5] * (args[2]^(args[5]-1)) * (args[1]^args[5])  /  (args[1]^args[5] + args[2]^args[5] + args[4]^args[5])^2
Symbolics.derivative(::typeof(hillar), args::NTuple{5,Any}, ::Val{3}) = (args[1]^args[5])   /  (args[1]^args[5] + args[2]^args[5] + args[4]^args[5])
Symbolics.derivative(::typeof(hillar), args::NTuple{5,Any}, ::Val{4}) = - args[3] * args[5] * (args[3]^(args[5]-1)) * (args[1]^args[5])  /  (args[1]^args[5] + args[2]^args[5] + args[4]^args[5])^2
Symbolics.derivative(::typeof(hillar), args::NTuple{5,Any}, ::Val{5}) = args[3] * (args[1]^args[5])  *  (log(args[1])*(args[2]^args[5] + args[4]^args[5]) - (args[2]^args[5])*log(args[2]) - (args[4]^args[5])*log(args[4]))   /   (args[1]^args[5] + args[2]^args[5] + args[4]^args[5])^2

ext: jl


text:
"""
Uses `hunting-circular-orbits.jl` and the Page & Thorne (1974) eq. (12)
to simulate a flux model.

Works for a < 0.97, then the orbit finder starts to panic close to the ISCO.
"""

# using Revise

using GeodesicTracer
using ComputedGeodesicEquations
using GeodesicBase

import CarterBoyerLindquist # for rms function
import LinearAlgebra

using StaticArrays
using Plots
gr()

# hoist other functions
include("hunting-circular-orbits.jl")

# for Kerr, this is just `u[2]`, but to illustrate
# how you could do this for other metrics
function sqrt_g_tilde(m, u)
    metric = GeodesicBase.metric(m, u)
    metric_no_θ = metric[setdiff(1:4, 3), setdiff(1:4, 3)]

    sqrt(-LinearAlgebra.det(metric_no_θ))
end

# gradient function
#
# we just compute discrete derivatives, since this is
# sufficient when we have small `r`.
# using interpolations can make the resolution better
# and increase the accuracy, but this is a good quick-and-dirty way
# to see if our results are going in the right direction
function ∇r(array, r)
    deriv = diff(array) ./ diff(r)
    # ensure same dimension using a lazy extension
    push!(deriv, deriv[end])
    deriv
end

"""
    f_approx(Lz, ∇rwt, ∇rLz_array, wt_array)

A discretized version of `f` from Page & Thorne (1974) eq. (12):

```math
f(r_n) \\approx - \\left. \\left[
   \\frac{1}{L_z} \\left(\\frac{\\partial \\dot{u}^t}{\\partial r} \\right)
\\right] \\right\\rvert_{r = r_n}
\\sum_{i=0}^{n} \\left\\{ \\left. \\left[
   \\frac{1}{\\dot{u}^t} \\left( \\frac{\\partial L_z}{\\partial r}\\right)
\\right] \\right\\rvert_{r = r_i}
\\right\\}
```

"""
function f_approx(Lz, ∇rwt, ∇rLz_array, wt_array)
    -(∇rwt / Lz) * sum(∇rLz_array ./ wt_array)
end

# use this to see where things break down / if they break down
function plot_derivative_test(r_range, res)
    lzs = res[:, 2]
    # naive derivatives
    derivs = ∇r(lzs, r_range)

    plot(r_range, lzs, label = "Lz", c = :black, legend = false)
    plot!(r_range, derivs, label = "derivative", c = :black, style = :dot, lw = 1.2)
end

# test plot
function f_approx_test(r_range, res)
    wts = res[:, 4]
    Lzs = res[:, 2]

    ∇rwt_array = ∇r(wts, r_range)
    ∇rLz_array = ∇r(Lzs, r_range)

    fs = map(2:length(r_range)) do i
        @views f_approx(Lzs[i], ∇rwt_array[i], ∇rLz_array[1:i], wts[1:i])
    end

    # we rescale here for a better comparison
    fs = fs ./ maximum(fs)

    # we lose a point, so select all but last `r_range`
    # since our derivative function needs at least 2 points to work
    plot(r_range[2:end], fs, c = :black, xlabel = "r", ylabel = "f", label = "simulated f")
end

# this works up until about 0.97 then, the orbit search
# starts to break down
m = BoyerLindquist(M = 1.0, a = 0.97)
u = @SVector [0.0, 13.0, deg2rad(90), 0.0]

println("Starting approximation...")

# making the steps too big over-estimates the gradient
r_range = CarterBoyerLindquist.rms(m.M, m.a):0.01:30.0
res = find_orbit_range(; r_range = r_range, a = m.a)

# use this to check if the derivatives are okay
# plot_derivative_test(r_range, res)

f_approx_test(r_range, res)
println("Approximation done!")

# compare to "true" value
import AccretionFormulae
f_true = AccretionFormulae.flux.(r_range, m.a, m.M)
# rescale as well
f_true = f_true ./ maximum(f_true)
plot!(r_range, f_true, label = "true f")

# savefig("simulated-f-not-log.svg")


text:
function hankel_operators(m,a::T,b::T) where T
    if iseven(m)
        # m = 2n
        n = Int(m//2)
        H_low = zeros(T,(n+1,n+1,m+1))
        H_up = zeros(T,(n,n,m+1))
        for i in 1:n+1
            for j in 1:n+1
                H_low[i,j,i+j-1] = 1
                if (i <= n) & (j <= n)
                    H_up[i,j,i+j] = a+b
                    H_up[i,j,i+j+1] = -1
                    H_up[i,j,i+j-1] = -a*b
                end
            end
        end
    else
        # m = 2n+1
        n = Int((m-1)//2)
        H_low = zeros(T,(n+1,n+1,m+1))
        H_up = zeros(T,(n+1,n+1,m+1))
        for i in 1:n+1
            for j in 1:n+1
                H_low[i,j,i+j] = 1
                H_low[i,j,i+j-1] = -a
                H_up[i,j,i+j-1] = b
                H_up[i,j,i+j] = -1
            end
        end
    end
    return H_low, H_up
end
function get_hankels(s,a,b)
    m = length(s)-1
    H_low, H_up = hankel_operators(m,a,b)
    s_reshaped = reshape(s,(1,1,length(s)))
    return dropdims(sum(H_low .*s_reshaped, dims=3),dims=3),dropdims(sum(H_up .*s_reshaped, dims=3),dims=3)
end
vec_H(H) = reshape(H,(size(H,1)*size(H,2),size(H,3)))
check_moment_sequence(x,a,b) = LinearAlgebra.isposdef.(get_hankels(x,a,b))


text:
# This file was generated, do not modify it. # hide
const dim = 8

db = DynamicMatrixDatabase(Float32, dim) # or VectorDatabase(Vector{Float32})
dist = L1Distance()

max_stars_repo_name: sadit/SimilaritySearchDemos

max_issues_repo_name: sadit/SimilaritySearchDemos

max_forks_repo_name: sadit/SimilaritySearchDemos


text:
using Gadfly

plot((x,y) -> x*exp(-(x-int(x))^2-y^2), -8., 8, -2., 2)


text:
include("intcode.jl")

# Part 1: count number of blocks drawn
prog = open(f -> parse.(Int64, split(readline(f), ",")), "day13_input.txt")
command_triples = reshape(run_intcode!(copy(prog), []), 3, :)'
drawn_tiles = Dict((row[1], row[2]) => row[3] for row in eachrow(command_triples))
println("Number of blocks: $(length(filter(isequal(2), collect(values(drawn_tiles)))))")

# Part 2: fake the payment and play the game!
function render(tiles)
    cols = maximum(first.(keys(tiles))) + 1
    rows = maximum(last.(keys(tiles))) + 1
    screen = zeros(Int32, rows, cols)
    for (coord, val) in tiles
        screen[last(coord) + 1, first(coord) + 1] = val
    end
    return [join(getindex.(Ref(Dict(0 => ' ', 1 => 'X', 2 => '█', 3 => '-', 4 => 'o')), row)) for row in eachrow(screen)]
end

prog[1] = 2
input = Channel(1)
output = Channel(30000)
drawn_tiles = Dict((0, 0) => 0)
@async begin
    run_intcode!(copy(prog), input, output)
    close(input)
    close(output)
end

# Screen renderer / input bot
try
    ball_x = 0
    paddle_x = 0
    score = 0
    while true
        put!(input, cmp(ball_x, paddle_x))
        yield()
        while isready(output)
            cmd = copy([take!(output) for i in 1:3])
            if cmd[1] == -1 && cmd[2] == 0
                score = cmd[3]
            else
                drawn_tiles[(cmd[1], cmd[2])] = cmd[3]
                if cmd[3] == 3 paddle_x = cmd[1] end
                if cmd[3] == 4 ball_x = cmd[1] end
            end
        end
        println.(["Score: $score, BX: $ball_x, PX: $paddle_x"; render(drawn_tiles); ""])
    end
catch InvalidStateException
end

max_stars_repo_name: ericgreveson/adventofcode

max_issues_repo_name: ericgreveson/adventofcode

max_forks_repo_name: ericgreveson/adventofcode


text:
include("CIFAR10.jl")

@testset "CIFAR Train/val test" begin
    train, val = get_processed_data(0.5f0, 8)

end

@testset "CIFAR Test test" begin
    unbatched = get_test_data()

    batched = get_test_data(1f0, 8)
end

max_stars_repo_name: huangyxi/DifferentiableNAS

max_issues_repo_name: huangyxi/DifferentiableNAS

max_forks_repo_name: huangyxi/DifferentiableNAS


text: # Data Loader


text:
# Autogenerated wrapper script for URCU_jll for aarch64-linux-gnu
export liburcu, liburcu_bp, liburcu_cds, liburcu_common, liburcu_mb, liburcu_memb, liburcu_qsbr, liburcu_signal

JLLWrappers.@generate_wrapper_header("URCU")
JLLWrappers.@declare_library_product(liburcu, "liburcu.so.8")
JLLWrappers.@declare_library_product(liburcu_bp, "liburcu-bp.so.8")
JLLWrappers.@declare_library_product(liburcu_cds, "liburcu-cds.so.8")
JLLWrappers.@declare_library_product(liburcu_common, "liburcu-common.so.8")
JLLWrappers.@declare_library_product(liburcu_mb, "liburcu-mb.so.8")
JLLWrappers.@declare_library_product(liburcu_memb, "liburcu-memb.so.8")
JLLWrappers.@declare_library_product(liburcu_qsbr, "liburcu-qsbr.so.8")
JLLWrappers.@declare_library_product(liburcu_signal, "liburcu-signal.so.8")
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_library_product(
        liburcu,
        "lib/liburcu.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        liburcu_bp,
        "lib/liburcu-bp.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        liburcu_cds,
        "lib/liburcu-cds.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        liburcu_common,
        "lib/liburcu-common.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        liburcu_mb,
        "lib/liburcu-mb.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        liburcu_memb,
        "lib/liburcu-memb.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        liburcu_qsbr,
        "lib/liburcu-qsbr.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        liburcu_signal,
        "lib/liburcu-signal.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
# simplified model from Chakrawal21

function chak21_phys_system()
    #using ModelingToolkit, DifferentialEquations, Plots
    @parameters t ks km kd Y HS HB kmr ksm
    # t in hours
    # masses per soilmass in mol/g
    @variables s(t) b(t) cr_tot(t) r_tot(t) q(t) dec_s(t) tvr_b(t)
    @variables mm_s(t) u1(t) u2(t) r(t) r_gr(t) r_m(t) db(t)
    D = Differential(t)
    HG = HS - Y * HB # can take out of system because does not involve t
    eqs = [
        D(s) ~ -dec_s + tvr_b,
        db ~ Y * u1 - tvr_b, # store output
        D(b) ~ db,
        D(cr_tot) ~ r_tot,
        D(r) ~ Y * u1/b * (s/(s+kmr) - r),
        u1 ~ b * ks * r * s/(s + km),
        u2 ~ b * ksm * (1-r) * s/(s + km),
        dec_s ~ u1 + u2,
        tvr_b ~ kd*b,
        r_gr ~ (1-Y) * u1,
        r_m ~ u2,
        r_tot ~ r_gr + r_m,
        q ~ -HG * u1 - HS * u2,
        ]
    @named chak21_phys = ODESystem(eqs)
    chak21_phys_s = structural_simplify(chak21_phys)
    p_straw = Dict(
        :parms => [ks => 0.15, km => 2.62e-6, kd => 1.35e-2, Y => 0.72,
            HB => -492, HS => -469,
            kmr => 1e-9, # half-saturation low so that optimum r -> 1
            ksm => 0.15, # uptake for P (maint comp) equal that for U (growth)
            ],
        :x0 => [s => 4.17e-5, b => 1.14e-5, cr_tot => 0, r => 1]
    ) # see S5
    chak21_phys_s, chak21_phys, p_straw
end


text:
function compute()::Int
    months = [31, 0, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
    day = 2
    result = 0
    for year ∈ 1901:2000
        months[2] = 28 + (year % 4 == 0 && (year % 100 ≠ 0 || year % 400 == 0))
        for month ∈ 1:12
            day += months[month] % 7
            if day % 7 == 0 result += 1 end
        end
    end
    return result
end

max_stars_repo_name: Dynortice/Project-Euler

max_issues_repo_name: Dynortice/Project-Euler

max_forks_repo_name: Dynortice/Project-Euler


text:
using Test
using LogicCircuits.Logical
using LogicCircuits.Utils

@testset "Trimmed SDD Test" begin

    num_vars = 7
    mgr = balanced_vtree(TrimMgrNode, num_vars)

    @test num_variables(mgr) == num_vars
    @test num_nodes(mgr) == 2*num_vars-1
    @test num_edges(mgr) == 2*num_vars-2
    @test mgr isa TrimSddMgr

    @test descends_from(mgr[1], mgr[end])
    @test descends_from(mgr[end].left, mgr[end])
    @test descends_from(mgr[end].right, mgr[end])
    @test descends_left_from(mgr[end].left, mgr[end])
    @test descends_left_from(mgr[end].left.left, mgr[end])
    @test descends_left_from(mgr[end].left.right, mgr[end])
    @test descends_right_from(mgr[end].right, mgr[end])
    @test descends_right_from(mgr[end].right.right, mgr[end])
    @test descends_right_from(mgr[end].right.left, mgr[end])

    @test !descends_from(mgr[end], mgr[1])
    @test !descends_left_from(mgr[end].right, mgr[end])
    @test !descends_left_from(mgr[end].right.left, mgr[end])
    @test !descends_left_from(mgr[end].right.right, mgr[end])
    @test !descends_left_from(mgr[end], mgr[end])
    @test !descends_left_from(mgr[end], mgr[end].left)
    @test !descends_left_from(mgr[end], mgr[end].right)
    @test !descends_right_from(mgr[end].left, mgr[end])
    @test !descends_right_from(mgr[end].left.right, mgr[end])
    @test !descends_right_from(mgr[end].left.left, mgr[end])
    @test !descends_right_from(mgr[end], mgr[end])
    @test !descends_right_from(mgr[end], mgr[end].left)
    @test !descends_right_from(mgr[end], mgr[end].right)

    x = Var(1)
    y = Var(2)

    x_c = compile(mgr, x)
    y_c = compile(mgr, y)

    @test x_c != y_c

    @test variable(x_c) == x
    @test literal(x_c) == var2lit(x)
    @test vtree(x_c) ∈ mgr
    @test positive(x_c)
    @test x_c == compile(mgr, x)

    @test variable(y_c) == y
    @test literal(y_c) == var2lit(y)
    @test vtree(y_c) ∈ mgr
    @test positive(y_c)
    @test y_c == compile(mgr, y)

    notx = -var2lit(x)

    notx_c = compile(mgr,notx)

    @test sat_prob(node2dag(x_c)) == 1//2
    @test sat_prob(node2dag(notx_c)) == 1//2
    @test model_count(node2dag(x_c),num_vars) == BigInt(2)^(num_vars-1)
    @test model_count(node2dag(notx_c),num_vars) == BigInt(2)^(num_vars-1)

    @test variable(notx_c) == x
    @test literal(notx_c) == notx
    @test vtree(notx_c) ∈ mgr
    @test negative(notx_c)
    @test notx_c == compile(mgr, notx)

    true_c = compile(true)

    @test is_true(true_c)
    @test constant(true_c) == true

    false_c = compile(false)

    @test is_false(false_c)
    @test constant(false_c) == false

    @test !true_c == false_c
    @test !false_c == true_c
    @test !x_c == notx_c
    @test !notx_c == x_c

    @test false_c & true_c == false_c
    @test false_c & notx_c == false_c
    @test false_c & x_c == false_c
    @test true_c & notx_c == notx_c
    @test x_c & notx_c == false_c
    @test true_c & true_c == true_c
    @test false_c & false_c == false_c
    @test x_c & x_c == x_c
    @test !x_c & !x_c == !x_c

    @test false_c | true_c == true_c
    @test false_c | notx_c == notx_c
    @test false_c | x_c == x_c
    @test true_c | notx_c == true_c
    @test x_c | notx_c == true_c
    @test true_c | true_c == true_c
    @test false_c | false_c == false_c
    @test x_c | x_c == x_c
    @test !x_c | !x_c == !x_c

    @test model_count(node2dag(true_c),num_vars) == BigInt(2)^(num_vars)
    @test model_count(node2dag(false_c),num_vars) == BigInt(0)

    v1 = compile(mgr, Var(1))
    v2 = compile(mgr, Var(2))
    v3 = compile(mgr, Var(3))
    v4 = compile(mgr, Var(4))
    v5 = compile(mgr, Var(5))
    v6 = compile(mgr, Var(6))
    v7 = compile(mgr, Var(7))

    p1 = XYPartition([Element(true_c,v3)])
    @test canonicalize(p1) === v3
    p2 = XYPartition([Element(v1,true_c), Element(!v1,false_c)])
    @test canonicalize(p2) === v1

    p3 = XYPartition([Element(v1,v3), Element(!v1,v7)])
    n1 = canonicalize(p3)
    p4 = XYPartition([Element(!v1,v7), Element(v1,v3)])
    n2 = canonicalize(p4)
    @test n1 === n2

    t1 = v1 & v3
    t2 = v3 & v1

    @test t1 === t2
    @test model_count(node2dag(t1),num_vars) == BigInt(2)^(num_vars-2)

    @test model_count(node2dag(v1 & v2 & v6),num_vars) == BigInt(2)^(num_vars-3)

    c1 = v1 | v3
    c2 = v3 | v1

    @test c1 === c2
    @test model_count(node2dag(c1),num_vars) == BigInt(2)^(num_vars-2) * 3

    f1 = (c1 & c2)

    @test f1 === (c2 & c1)
    @test f1 === (c1 & c2 & c2)
    @test f1 === (c1 & c2 & c2 & true_c)

    @test (v3 | v7) !== false_c
    @test (v3 | v7) & (!v3 | v7) !== false_c
    @test (v3 | v7) & (!v3 | v7) & (v3 | !v7) !== false_c
    @test (v3 | v7) & (!v3 | v7) & (v3 | !v7) & (!v3 | !v7) === false_c

    f2 = (c1 | c2)

    @test f2 === (c2 | c1)
    @test f2 === (c1 | c2 | c2)
    @test f2 === (c1 | c2 | c2 | false_c)

    @test (v3 & v7) !== true_c
    @test (v3 & v7) | (!v3 & v7) !== true_c
    @test (v3 & v7) | (!v3 & v7) | (v3 & !v7) !== true_c
    @test (v3 & v7) | (!v3 & v7) | (v3 & !v7) | (!v3 & !v7) === true_c

    @test (v1 & v7 & v3) === !(!v1 | !v7 | !v3)

    @test f2 & !f2 === false_c
    @test f2 | !f2 === true_c

    f3 = (v1 | !v3 | v5) & (!v2 | !v3 | v7) & (v2 | !v6 | !v5) & (v1 | v2 | v3) & (!v5 | v6 | v7) & (v2 | !v3 | !v6) & (!v1 | !v6 | !v7)

    @test f3 & !f3 === false_c
    @test f3 | !f3 === true_c
    @test f3 & !(v2 | !v6 | !v5) === false_c

    f4 = ((v1 | !v3 | v5) & (!v5 | v6 | v7| v1) & (!v5 | v6 | v7)  & (!v1 | !v6 | !v7)  & (!v2 | !v3 | v7)) & ((v2 | !v3 | !v6)  & (v2 | !v6 | !v5) & (v1 | v2 | v3) & (v1 | v2 | v3 | !v6) )

    @test f3 === f4

    fΔ = node2dag(f4, TrimNode)

    @test TrimSdd <: Sdd
    @test fΔ isa TrimSdd
    @test fΔ isa Sdd

end


text:
"""
constrained: return the violation of the KKT conditions
length(lambda) > 0
"""
function _grad_lagrangian(pb    :: AbstractMPCCModel,
                          state :: MPCCAtX)

 if (pb.meta.ncon + pb.meta.ncc) == 0 & !has_bounds(pb)
  return state.gx
 elseif (pb.meta.ncon + pb.meta.ncc) == 0
  return state.gx + state.mu
 elseif pb.meta.ncc == 0
  return state.gx + state.mu + state.Jx' * state.lambda
 else #pb.meta.ncon + pb.meta.ncc > 0
  return state.gx + state.mu + state.Jx' * state.lambda - state.JGx' * state.lambdaG - state.JHx' * state.lambdaH
 end
end

function _sign_multipliers_bounds(pb    :: AbstractMPCCModel,
                                  state :: MPCCAtX)
 if has_bounds(pb)
  return vcat(min.(max.( state.mu,0.0), - state.x + pb.meta.uvar),
              min.(max.(-state.mu,0.0),   state.x - pb.meta.lvar))
 else
  return zeros(0)
 end
end

function _sign_multipliers_nonlin(pb    :: AbstractMPCCModel,
                                  state :: MPCCAtX)
 if pb.meta.ncon == 0
  return zeros(0)
 else
  return vcat(min.(max.(  state.lambda, 0.0), - state.cx + pb.meta.ucon),
              min.(max.(- state.lambda, 0.0),   state.cx - pb.meta.lcon))
 end
end

function _sign_multipliers_comp(pb    :: AbstractMPCCModel,
                                state :: MPCCAtX)
 if pb.meta.ncc == 0
  return zeros(0)
 else
  return vcat(min.(abs.(state.lambdaG), abs.(state.cGx - pb.meta.lccG)),
              min.(abs.(state.lambdaH), abs.(state.cHx - pb.meta.lccH)))
 end
end

function _sign_multipliers_comp_Sstat(pb    :: AbstractMPCCModel,
                                      state :: MPCCAtX,
                                      actif :: Float64)

 I = findall(x -> x<=actif, 0.5*(state.cGx + state.cHx))
 pos = max.(max.(- state.lambdaG[I], 0.0), max.(- state.lambdaH[I], 0.0))

 return pos
end

function _sign_multipliers_comp_Mstat(pb    :: AbstractMPCCModel,
                                      state :: MPCCAtX)

 eq  = max.(abs.(state.lambdaG .* state.cGx), abs.(state.lambdaH .* state.cHx))
 pos = max.(- state.lambdaG .* state.lambdaH, 0.0)
 mpo = max.(- max.(state.lambdaG, state.lambdaH), 0.0)

 return max.(eq, pos, mpo)
end

function _sign_multipliers_comp_Cstat(pb    :: AbstractMPCCModel,
                                      state :: MPCCAtX)

 eq  = max.(abs.(state.lambdaG .* state.cGx), abs.(state.lambdaH .* state.cHx))
 pos = max.(- state.lambdaG .* state.lambdaH, 0.0)

 return max.(eq, pos)
end

function _feasibility(pb    :: AbstractMPCCModel,
                      state :: MPCCAtX)
 if (pb.meta.ncon + pb.meta.ncc) == 0
  return vcat(max.(  state.x  - pb.meta.uvar, 0.0),
              max.(- state.x  + pb.meta.lvar, 0.0))
 elseif pb.meta.ncc == 0
  return vcat(max.(  state.cx  - pb.meta.ucon, 0.0),
              max.(- state.cx  + pb.meta.lcon, 0.0),
              max.(  state.x   - pb.meta.uvar, 0.0),
              max.(- state.x   + pb.meta.lvar, 0.0))
 else #pb.meta.ncon + pb.meta.ncc > 0
  return vcat(max.(  state.cx  - pb.meta.ucon, 0.0),
              max.(- state.cx  + pb.meta.lcon, 0.0),
              max.(  state.x   - pb.meta.uvar, 0.0),
              max.(- state.x   + pb.meta.lvar, 0.0),
              max.(- state.cGx + pb.meta.lccG, 0.0),
              max.(- state.cHx + pb.meta.lccH, 0.0),
              min.(  state.cGx, state.cHx))
 end
end

"""
W-stationary check:
WStat: verifies the Weak-stationary conditions
required: state.gx
+ if bounds: state.mu
+ if constraints: state.cx, state.Jx, state.lambda
+ if complementarity: state.cGx, state.JGx, state.lambdaG
                      state.cHx, state.JHx, state.lambdaH
"""
function WStat(pb    :: AbstractMPCCModel,
               state :: MPCCAtX;
               pnorm :: Float64 = Inf,
               kwargs...)

    #Check the gradient of the Lagrangian
    gLagx      = _grad_lagrangian(pb, state)
    #Check the complementarity condition for the bounds
    res_bounds = _sign_multipliers_bounds(pb, state)
    #Check the complementarity condition for the constraints
    res_nonlin = _sign_multipliers_nonlin(pb, state)
    #Check the complementarity condition for the complementarity constraints
    res_comp   = _sign_multipliers_comp(pb, state)
    #Check the feasibility
    feas       = _feasibility(pb, state)

    res = vcat(gLagx, feas, res_bounds, res_nonlin, res_comp)

    return norm(res, pnorm)
end

"""
S-stationary check:
Compute the score corresponding to:
a) W-stationarity
b) the correct sign of multipliers for indices
I_00 = { i : 0.5(G_i(x^*)+H_i(x^*)) <= actif }

required: state.gx
+ if bounds: state.mu
+ if constraints: state.cx, state.Jx, state.lambda
+ if complementarity: state.cGx, state.JGx, state.lambdaG
                      state.cHx, state.JHx, state.lambdaH
"""
function SStat(pb    :: AbstractMPCCModel,
               state :: MPCCAtX;
               pnorm :: Float64 = Inf,
               actif :: Float64 = 1e-6,
               kwargs...)

    res = WStat(pb, state, pnorm = pnorm; kwargs...)
    sign = pb.meta.ncc > 0 ? norm(_sign_multipliers_comp_Sstat(pb, state, actif), pnorm) : 0.0

    return max(res, sign)
end

"""
M-stationary check:
Compute the score corresponding to:
a) W-stationarity
b) the correct sign of multipliers for indices I_00 by checking:
         lambdaG .* lambdaH => 0
(i.e. max.(- state.lambdaG .* state.lambdaH, 0.0))
and
         max(lambdaG, lambdaH) => 0
(i.e. max.(- max.(state.lambdaG, state.lambdaH), 0.0))

required: state.gx
+ if bounds: state.mu
+ if constraints: state.cx, state.Jx, state.lambda
+ if complementarity: state.cGx, state.JGx, state.lambdaG
                      state.cHx, state.JHx, state.lambdaH
"""
function MStat(pb    :: AbstractMPCCModel,
               state :: MPCCAtX;
               pnorm :: Float64 = Inf,
               kwargs...)

    res = WStat(pb, state, pnorm = pnorm; kwargs...)
    sign = pb.meta.ncc > 0 ? norm(_sign_multipliers_comp_Mstat(pb, state), pnorm) : 0.0

    return max(res, sign)
end

"""
C-stationary check:
Compute the score corresponding to:
a) W-stationarity
b) the correct sign of multipliers for indices I_00 by checking:
         lambdaG .* lambdaH => 0
(i.e. max.(- state.lambdaG .* state.lambdaH, 0.0))

required: state.gx
+ if bounds: state.mu
+ if constraints: state.cx, state.Jx, state.lambda
+ if complementarity: state.cGx, state.JGx, state.lambdaG
                      state.cHx, state.JHx, state.lambdaH
"""
function CStat(pb    :: AbstractMPCCModel,
               state :: MPCCAtX;
               pnorm :: Float64 = Inf,
               kwargs...)

    res = WStat(pb, state, pnorm = pnorm; kwargs...)
    sign = pb.meta.ncc > 0 ? norm(_sign_multipliers_comp_Cstat(pb, state), pnorm) : 0.0

    return max(res, sign)
end


text:
"""
AbstractChannelData.jl defines the AbstractChannelData type and methods

------------------------------------------------------------------------------
COPYRIGHT/LICENSE. This file is part of the XYZ package. It is subject to
the license terms in the LICENSE file found in the top-level directory of
this distribution. No part of the XYZ package, including this file, may be
copied, modified, propagated, or distributed except according to the terms
contained in the LICENSE file.
------------------------------------------------------------------------------
"""
# --- Exports

# ------ Types

export AbstractChannelData

# ------ Functions

import ZMQ.get_type
export get_data, set_data!
export encode_value, decode_value

# --- Type definitions

"""
    AbstractChannelData

Supertype for types that store and encode/decode data transmitted by input and
output channels. Concrete subtypes should

* possess a default constructor with no arguments and

* be mutable or contain mutable fields so that they can be modified to
  represent the most recent data sent or received by a data channel.

Interface
=========

    get_type(channel_data::AbstractChannelData)

    get_data(channel_data::AbstractChannelData)

    set_data!(channel_data::AbstractChannelData, value)

    encode_value(::Type{<:AbstractChannelData}, data)::Vector{UInt8}

    decode_value(::Type{<:AbstractChannelData}, bytes::Vector{UInt8})
"""
abstract type AbstractChannelData end

# --- Method definitions
#
# Note: the following method definitions are no-op place holders to provide
#       a central location for docstrings.

"""
    get_type(channel_data::AbstractChannelData)

Return the type of the data handled by `channel_data`.
"""
get_type(channel_data::AbstractChannelData) = nothing

"""
    get_data(channel_data::AbstractChannelData; encode::Bool=false)

Return the data value stored by `channel_data`. When `encode` is true, the data
value is encoded as a byte vector using the `encode_value()` function before
being returned. When `encode` is false, the data value is returned directly.
"""
get_data(channel_data::AbstractChannelData; encode::Bool=false) = nothing

"""
    set_data!(channel_data::AbstractChannelData, value; decode::Bool=false)

Set the data stored by `channel_data` to `value`. When `decode` is true,
`value` is intepreted as a byte vector and decoded using the `decode_value()`
function before being stored. When `decode` is false, `value` is stored as-is.
"""
set_data!(channel_data::AbstractChannelData, value; decode::Bool=false) =
    nothing

"""
    encode_value(value)::Vector{UInt8}

Encode `value` as a byte vector.
"""
encode_value(::Type{<:AbstractChannelData}, value)::Vector{UInt8} = nothing

"""
    decode_value(::Type{<:AbstractChannelData}, bytes::Vector{UInt8})

Convert `bytes` to a value.
"""
decode_value(::Type{<:AbstractChannelData}, bytes::Vector{UInt8}) = nothing


text:
using MXNet

#--------------------------------------------------------------------------------
# Helper functions to construct larger networks

# basic Conv + BN + ReLU factory
function conv_factory(data, num_filter, kernel; stride=(1,1), pad=(0,0), act_type=:relu)
  conv = mx.Convolution(data=data, num_filter=num_filter, kernel=kernel, stride=stride, pad=pad)
  bn   = mx.BatchNorm(data=conv)
  act  = mx.Activation(data=bn, act_type=act_type)
  return act
end

# simple downsampling factory
function downsample_factory(data, ch_3x3)
  # conv 3x3
  conv = conv_factory(data, ch_3x3, (3,3), stride=(2,2), pad=(1,1))
  # pool
  pool = mx.Pooling(data=data, kernel=(3,3), stride=(2,2), pool_type=:max)
  # concat
  concat = mx.Concat(conv, pool)
  return concat
end

# a simple module
function simple_factory(data, ch_1x1, ch_3x3)
  # 1x1
  conv1x1 = conv_factory(data, ch_1x1, (1,1); pad=(0,0))
  # 3x3
  conv3x3 = conv_factory(data, ch_3x3, (3,3); pad=(1,1))
  # concat
  concat = mx.Concat(conv1x1, conv3x3)
  return concat
end

#--------------------------------------------------------------------------------
# Actual architecture
data    = mx.Variable(:data)
conv1   = conv_factory(data, 96, (3,3); pad=(1,1), act_type=:relu)
in3a    = simple_factory(conv1, 32, 32)
in3b    = simple_factory(in3a, 32, 48)
in3c    = downsample_factory(in3b, 80)
in4a    = simple_factory(in3c, 112, 48)
in4b    = simple_factory(in4a, 96, 64)
in4c    = simple_factory(in4b, 80, 80)
in4d    = simple_factory(in4b, 48, 96)
in4e    = downsample_factory(in4d, 96)
in5a    = simple_factory(in4e, 176, 160)
in5b    = simple_factory(in5a, 176, 160)
pool    = mx.Pooling(data=in5b, pool_type=:avg, kernel=(7,7), name=:global_pool)
flatten = mx.Flatten(data=pool, name=:flatten1)
fc      = mx.FullyConnected(data=flatten, num_hidden=10, name=:fc1)
softmax = mx.SoftmaxOutput(data=fc, name=:loss)

#--------------------------------------------------------------------------------
# Prepare data
filenames = mx.get_cifar10()
batch_size = 128
num_epoch  = 10
num_gpus   = 8

train_provider = mx.ImageRecordProvider(label_name=:loss_label,
        path_imgrec=filenames[:train], mean_img=filenames[:mean],
        rand_crop=true, rand_mirror=true, data_shape=(28,28,3),
        batch_size=batch_size, preprocess_threads=1)
test_provider = mx.ImageRecordProvider(label_name=:loss_label,
        path_imgrec=filenames[:test], mean_img=filenames[:mean],
        rand_crop=false, rand_mirror=false, data_shape=(28,28,3),
        batch_size=batch_size, preprocess_threads=1)

#--------------------------------------------------------------------------------
# Training model
gpus = [mx.Context(mx.GPU, i) for i = 0:num_gpus-1]
model = mx.FeedForward(softmax, context=gpus)

# optimizer
optimizer = mx.SGD(lr=0.05, momentum=0.9, weight_decay=0.0001)

# fit parameters
mx.fit(model, optimizer, train_provider, n_epoch=num_epoch, eval_data=test_provider,
       initializer=mx.UniformInitializer(0.07), callbacks=[mx.speedometer()])


text:
using AuditoryFilters, DSP, Base.Test

include("test_erb_fb.jl")
include("test_modulation_fb.jl")
include("test_gammatonegram.jl")


text:
@testset "1 Algebraic functions" begin
    include("1.1 Binomial products/.jl")
    include("1.2 Trinomial products/.jl")
    include("1.3 Miscellaneous/.jl")
end

max_stars_repo_path: test/1 Algebraic functions/.jl

max_issues_repo_path: test/1 Algebraic functions/.jl

max_forks_repo_path: test/1 Algebraic functions/.jl


text:
using DrWatson
@quickactivate :FractalDimension # uses DynamicalSystems, PyPlot
include(srcdir("style.jl"))

using DelimitedFiles
n = 2
name = "Shinriki_data_$n"
file = datadir("experimental", name*".txt")
X = readdlm(file, ','; skipstart = 21) # for Shinriki data
L = size(X)[1]
x = X[1:2:L÷10, 2]

newname = "shinriki$(n)"
writedlm(datadir("experimental", newname*".txt"), X[1:2:L÷10, 2:end])

# Run this after `x` is obtained
fig = figure()
ax = subplot(2,1,1)
ax.plot(x)

mi = selfmutualinfo(x, 0:100)
w = estimate_delay(x, "mi_min")
ax = subplot(2,2,3)
ax.plot(0:100, mi; label = "mut. inf.")
ax.axvline(w; color = "C2")
ax.legend()

Es = delay_afnn(x, w, 2:10)
ax = subplot(2,2,4)
ax.plot(2:10, Es; label = "Cao")
ax.legend()

# Y, τs = pecuzal_embedding(x .+ 1e9randn(length(x)); w, τs = 0:1:4w)

max_stars_repo_name: Datseris/FractalDimension

max_issues_repo_name: Datseris/FractalDimension

max_forks_repo_name: Datseris/FractalDimension


text:
export MachOSegmentCmd, MachOSegment32Cmd, MachOSegment64Cmd

export segment_name, segment_offset, segment_file_size, segment_memory_size,
       segment_num_sections

"""
    MachOSegmentCmd

Mach-O Segment load command type, containing information about the virtual
memory layout of a chunk of the program.  This type is very MachO-specific, so
it is not comparable to other object file formats directly.  However, note that
this data structure is the gateway through which the file sections are
accessed, and while there is a convenience `Sections(::MachOHandle)` method
that will abstract all that away for you, you can also directly call
`Sections(::MachOLoadCmdRef{MachOSegmentCmd})` to get the sections belonging to
the given segment.  Note that the `Sections` call works only upon a
`MachOLoadCmdRef{MachOSegmentCmd}`, it will not work on the `MachOSegmentCmd`
itself directly.

### Creation:
  - MachOSegmentCmd()

### Format-specific properties:
  - segment_name()
  - segment_offset()
  - segment_file_size()
  - segment_memory_size()
  - segment_num_sections()
"""
abstract type MachOSegmentCmd{H <: MachOHandle} <: MachOLoadCmd{H} end
@io struct MachOSegment32Cmd{H <: MachOHandle} <: MachOSegmentCmd{H}
    segname::fixed_string{UInt128}
    vmaddr::UInt32
    vmsize::UInt32
    fileoff::UInt32
    filesize::UInt32
    maxprot::UInt32
    initprot::UInt32
    nsects::UInt32
    flags::UInt32
end

@io struct MachOSegment64Cmd{H <: MachOHandle} <: MachOSegmentCmd{H}
    segname::fixed_string{UInt128}
    vmaddr::UInt64
    vmsize::UInt64
    fileoff::UInt64
    filesize::UInt64
    maxprot::UInt32
    initprot::UInt32
    nsects::UInt32
    flags::UInt32
end

segment_name(cmd::MachOSegmentCmd) = cmd.segname
segment_offset(cmd::MachOSegmentCmd) = cmd.fileoff
segment_file_size(cmd::MachOSegmentCmd) = cmd.filesize
segment_memory_size(cmd::MachOSegmentCmd) = cmd.vmsize
segment_num_sections(cmd::MachOSegmentCmd) = cmd.nsects

@derefmethod segment_name(cmd::MachOLoadCmdRef{H,S}) where {H <: MachOHandle} where {S <: MachOSegmentCmd}
@derefmethod segment_offset(cmd::MachOLoadCmdRef{H,S}) where {H <: MachOHandle} where {S <: MachOSegmentCmd}
@derefmethod segment_file_size(cmd::MachOLoadCmdRef{H,S}) where {H <: MachOHandle} where {S <: MachOSegmentCmd}
@derefmethod segment_memory_size(cmd::MachOLoadCmdRef{H,S}) where {H <: MachOHandle} where {S <: MachOSegmentCmd}
@derefmethod segment_num_sections(cmd::MachOLoadCmdRef{H,S}) where {H <: MachOHandle} where {S <: MachOSegmentCmd}


text:
"""
Free space in Gb at the given path.
"""
free_space(path) = shutil.disk_usage(path)[3] / 1000_000_000

settings = ArgParseSettings()

@add_arg_table settings begin
    "--file"
    help = """
    path to the file.
    ...
    """
end

@add_arg_table settings begin
    "--output", "-o"
    help = "output..."
    default = "."
    """
    --file
    """
    help = """
    path to the file.
    ...
    """
    default = ""
end


text:
module MultidimensionalSpectroscopy

export create_colormap, zeropad, interpt, make2Dspectra, correlations,
        view_dm_evo, save_2d, load_2d, plot2d, crop2d, round2d, tri, absorptionSpectrum, plot2d_comps,
         pretty_show_mat, vib_analysis, create_subspace, rand_normal, plot_levels, logger, out2d, corr2spec

using QuantumOptics, FFTW, LinearAlgebra, PyPlot, Colors, DelimitedFiles, Printf, Random, JLD2, Interact, Blink

"""
    plot2d(ω, data; repr = "absorptive", norm_spec=false, scaling="lin")

Plots the 2D data.

# Arguments
ω         : frequency / energy axis (must be equal in both dimensions)
data      : data structure of 2D calculation output
repr      : pick a representation ("absorptive", "dispersive", "absolute")
norm_spec : normalize to abs.(maximum(data)), where data is the 2D spectrum to
            be plotted
scaling   : pick z-scaling ("lin", "tan", "cube", "asinh")
norm      : if norm == 0 normalize each 2D spectrum individually, else normalize
            all spectra to norm = <val>, where <val> can be the overall maximum of
            the dataset (for all evolution time steps, T)

"""
function plot2d(ω, data; repr = "absorptive", norm_spec=false, scaling="lin", norm=0)

    ## select the representation
    if repr in ["absorptive", "abst", "real"]
        data = real(data)
        cmp = create_colormap()
    elseif repr in ["dispersive", "disp", "imaginary", "imag", "i"]
        data = imag(data)
        cmp = create_colormap()
    elseif repr in ["absolute", "absl"]
        #data = sqrt.(real(data).^2 + imag(data).^2) #BUG doing this calculation results in Inf values when data{ComplexF16}
        data = abs.(data)
        cmp = create_colormap()
    elseif repr in ["phase", "ph"]
        data_cntl = real(data) ./ maximum(abs.(data))
        data = atan.(imag(data) ./ real(data))
        cmp = "rainbow"
    end

    ## normalize 2D spectrum to maximum
    if norm_spec
        data = data/maximum(abs.(data));
    end

    #TODO: let user set the sensitivity of contourlines
    sensitivity = "++"

    ## make levels for contourplot with different sensitivity...
    lvls = [-1.025:0.05:1.025;]
    if sensitivity == "+++"
        lvls = [-1,-.9,-.8,-.7,-.6,-.5,-.4,-.3,-.2,-.1,-.08,-.06,-.04,-.02,-.01,-0.001, 0,
                .001,.01,.02,.04,.06,.08,.1,.2,.3,.4,.5,.6,.7,.8,.9,1];
    elseif sensitivity == "++"
         lvls = [-1,-.9,-.8,-.7,-.6,-.5,-.4,-.3,-.2,-.1,-.08,-.06,-.04,-.02, 0,
                .02,.04,.06,.08,.1,.2,.3,.4,.5,.6,.7,.8,.9,1];
    elseif sensitivity == "+"
        lvls = [-1,-.9,-.8,-.7,-.6,-.5,-.4,-.3,-.2,-.1,-.075,-.05, 0,
                .05,.075,.1,.2,.3,.4,.5,.6,.7,.8,.9,1];
    else
    end

    # ... and scaling
    if scaling == "lin"
        # do nothing
    elseif scaling == "tan"
        lvls = tan.(lvls); lvls = lvls ./ maximum(lvls);
    elseif scaling == "cube"
        lvls = lvls.^3;
    elseif scaling == "asinh"
        lvls = asinh.(0.1*lvls); lvls = lvls ./ maximum(lvls);
    end

    # rescale lvls to data (don't rescale data, as it takes more time)
    # - redundant if data is normalized to 1.
    m = -1 * (maximum(abs.(data))); M = (maximum(abs.(data)))
    if m == 0 && M == 0             # set minimum level in case m == 0 and M == 0
        lvls = 0.001 * lvls;
    elseif m == Inf || M == Inf     # throw error in case of inf
        println("\nERROR: z-value is Inf!\n")
    else                            # scale lvls
        lvls = lvls * maximum(abs.([m, M]));
        lvls_ticks = [-1:0.2:1;] * maximum(abs.([m, M]))
    end

    # normalize to global maximum with evolution time T scan
    if norm == 0
        # do nothing
    else
        lvls = lvls ./ maximum(lvls) * norm;
    end

    ## plot data as filled contour

    # choose convention
    conv = "ω3ω1"

    if conv == "ω1ω3"
        cs  = contourf(ω,ω,data,lvls,cmap=create_colormap()); colorbar();
        if repr == "phase"
            cs2 = contour(ω,ω,data_cntl,lvls,colors="white",linewidths=(.5,))
        else
            cs2 = contour(cs,levels=cs.levels[:],colors="gray",linewidths=(.5,))
        end
        xlabel("detection (ω₃)"); ylabel("excitation (ω₁)"); #clim([m, M]);
    elseif conv == "ω3ω1"
        cs  = contourf(ω,ω,transpose(data),lvls,cmap=cmp);
        if repr == "phase"
            cs2 = contour(ω,ω,transpose(data_cntl),lvls,colors="white",linewidths=(.5,))
        else
            cs2 = contour(cs,levels=cs.levels[:],colors="gray",linewidths=(.5,))
        end
        xlabel("excitation (ω₁)"); ylabel("detection (ω₃)");
    end

    # add colorbar
    cbar = colorbar(cs)
    cbar.add_lines(cs2)
    # add diagonal line
    plot([ω[1], ω[end]], [ω[1], ω[end]],"k--");

end

"""
    plot_timeTrace(dat2d,T,ω,w1,w3)

Plots a time trace in T from a set of 2D spectra.

# Arguments
- dat2d:    2D dataset
- T:        evolution time
- ω:        frequency / energy axis
- w1, w3:   combination of excitation and detection frequency to plot.
            Can be a list of values.

# Example

tbd

"""
function plot_timeTrace(dat2d,T,ω,w1,w3)

    ## create the figure
    figure(figsize=(8.5,3.5))

    ## plot 2D map
    subplot(121)
    plot2d(ω,dat2d[1])
    for i in 1:length(w1)
        plot(w1[i],w3[i],"o",fillstyle="none",mew=3,ms=10,marker="s")
    end

    ## plot time trace
    subplot(122)
    # get indices from (lists) w1 and w3
    idx1 = [argmin(abs.(ω .- i)) for i in w1]
    idx3 = [argmin(abs.(ω .- i)) for i in w3]
    for j in 1:length(idx1)
        I = [dat2d[i][idx1[j],idx3[j]] for i in 1:length(T)]
        plot(T,I,"o-",label=string(w1[j])*";"*string(w3[j]))
    end
    xlabel("Time"); ylabel("Intensity at w1;w3")
    legend()
    tight_layout()

    ##TODO: slider to flip through 2D spectra
    #sel = slider(1:length(dat2d))
    #w = Window()
    #Interact.@on
    #body!(w, sel);
end

"""
    save_2d(spec2d,T,fn_base)

Saves 2D spectra to specified directory. Can be used in a loop to store all
data.

# Arguments
- spec2d  : 2D array of data to be saved
- T       : array of T times
- fn_base : directory to store data in

# Example
```julia
    save_2d([out2d[i].full2d for i = 1:length(T)],T,fn)
```
"""
function save_2d(spec2d,T,fn_base)

    ## check if directory exists
    if isdir(fn_base)
        # do nothing
    else
        mkdir(fn_base)
    end

    ## save data
    # write T and ω to files
    writedlm(fn_base * "T_steps.dat"    , T, ',')
    writedlm(fn_base * "energy_axis.dat", spec2d[1].ω, ',')
    # save different representations of 2D spectra
    for i = 1:length(T)
        fn = fn_base * @sprintf("2Dspec_%03i.dat",i)
        writedlm(fn, round.(real.(spec2d[i].full2d), digits=2), ',')
        fn = fn_base * @sprintf("2Dspec_%03i_gsb.dat",i)
        writedlm(fn, round.(real.(spec2d[i].gsb), digits=2), ',')
        fn = fn_base * @sprintf("2Dspec_%03i_se.dat",i)
        writedlm(fn, round.(real.(spec2d[i].se), digits=2), ',')
        fn = fn_base * @sprintf("2Dspec_%03i_esa.dat",i)
        writedlm(fn, round.(real.(spec2d[i].esa), digits=2), ',')
        println("\nFile saved as: " * fn)
    end
    println("\nFiles saved!")
end

"""
save as .h5 instead

"""
function save_2d(spec2d,T,fn_base,type)

    ## check if directory exists
    if isdir(fn_base)
        # do nothing
    else
        mkdir(fn_base)
    end

    #TODO: correctly save struct to .h5
    fn = fn_base * "2Dspec.h5"
    h5open(fn, "w") do file
        write(file, "timesteps", T)
        for i = 1:length(T)
            write(file, @sprintf("2Dspec_%03i",i), spec2d[i])  # alternatively, say "@write file A"
        end
    end

    println("\nFiles saved!")
end

"""
Loads a 2D spectrum.
"""
function load_2d(fn_base;type="full")
    #
    T = readdlm(fn_base * "T_steps.dat")
    ω = readdlm(fn_base * "energy_axis.dat")
    #
    dat2d = []
    for i in 1:length(T)
        if type == "full"
            push!(dat2d,readdlm(fn_base * @sprintf("2Dspec_%03i.dat",i),','));
        end
    end
    return T, ω, dat2d;
end

"""
    view_dm_evo(rhot,dt)

Quick function to view evolution of density matrix elements as 2D plot.

# Arguments
rhot : time evolution of density matrix
dt   : time steps between subplots

"""
function view_dm_evo(rhot,dt)

    ## create a figure
    figure()
    for i = 1:12 # limit to 12 subplots
        subplot(3,4,i,aspect="equal")
        j = (i-1) * dt + 1
        pcolormesh(real(rhot[j].data),cmap="seismic"); clim([-1,1]);
        xticks([1:1:9;]); yticks([1:1:9;])
        grid("on",which="major")
        #title("tlist # $j")
    end
    tight_layout()
end

"""
    rand_normal(mean, stdev)

Return a random sample from a normal (Gaussian) distribution. See
https://www.johndcook.com/blog/services-2/

Can be useful for simulating the effect of random processes on the
2D spectra.

"""
function rand_normal(mean, stdev)
    if stdev <= 0.0
        error("standard deviation must be positive")
    end
    u1 = rand()
    u2 = rand()
    r = sqrt( -2.0*log(u1) )
    theta = 2.0*pi*u2
    mean + stdev*r*sin(theta)
end

"""
    cmp = create_colormap(scheme)

Create a "bright" or "dark" custom colormap to display 2D spectra.
"""
function create_colormap(scheme="bright")

    ## define colors from most negative (MMM) ... to ... most positive (PPP)
    MMM = RGB(0.718 , 0     , 0.718 );
    MM  = RGB(0.516 , 0.516 , 0.991 );
    M   = RGB(0     , 0.559 , 0.559 );
    P   = RGB(0     , 0.592 , 0     );
    PP  = RGB(0.827 , 0.827 , 0     );
    PPP = RGB(0.947 , 0.057 , 0.057 );

    ## zero value is either "bright" or "dark"
    if scheme == "bright"
        Z = RGB(.93,.93,.93);
    elseif scheme == "dark"
        Z = RGB(.23,.23,.23);
    else
        error("MultidimensionalSpectroscopy.jl -- No colormap scheme selected!")
    end

    return cmp  = ColorMap("cmps",[MMM,MM,M,Z,P,PP,PPP]);
end

"""
    dat = zeropad(dat, N)

Add zeros to end end of correlation data (after it has decayed to zero) in order
to increase frequency/energy resolution in the spectral domain (i.e. after Fourier
transform).

# Arguments
* 'dat' : correlation data (1D or 2D)
* 'N'   : data will be zeropadded up to 2^N, if N = 0 return input data

"""
function zeropad(dat, N)

    if 0 < N < log2(size(dat,1))
        # if N smaller than log2(size of data), increase N by 1
        N = ceil(log2(size(dat,1)))
    elseif N == 0
        # if N == zero, return original data
        println("\nNo zero padding.\n")
        return dat;
    end

    # Sanity check that N is not too large
    if N > 2^12
        error("Large N detected. Check if N isn't too large.")
    end

    if N > 0
        # switch between 1D and 2D array
        if size(dat,2) == 1
            out = complex(zeros(2^N));
            println("\nZeropadding 1D-data to $(2^N) values\n")
        elseif size(dat,2) > 1
            out = complex(zeros(2^N,2^N));
            println("\nZeropadding 2D-data to $(2^N) x $(2^N) values\n")
        end

        out[1:size(dat,1),1:size(dat,2)] = dat;

        return out;
    end
end

"""
    tlist = interpt(tlist, N)

Interpolate frequency/energy vector after zeropadding of data using "zeropad".

# Arguments
* 'tlist'   :
* 'N'       :

# Output
* 'tlist'   : extrapolated time vector
* 'ω'       : interpolated frequency/energy vector
"""
function interpt(tlist, N)

    # if N smaller than log2(size of data), increase N by 1
    if 0 < N < log2(length(tlist))
        N = ceil(log2(size(dat,1)))
    elseif N == 0
        tlist = tlist;
    end

    # Sanity check that N is not too large
    if N > 2^12
        error("Large N detected. Check if N isn't too large.")
    end

    if N > 0
        # extrapolate tlist / get new tlist[end]
        dt = abs(tlist[end]-tlist[1]) / (length(tlist)-1)
        tlist = [0:(2^N)-1;] * dt .+ tlist[1];
    end

    # create ω axis corresponding to zeropadded data
    dω = 1/tlist[end]
    ω = ([1:length(tlist);] .- length(tlist)/2 ) * dω * 2 * π;
    # return zeropadded data, corresponding ω and tlist

    return tlist, ω;
end

"""
    out2d

Structure that stores the output of "make2Dspectra"

# Elements:
* 'ω'           : energy or frequency axis
* 'T'           : evolution time step
* 'full2d'      : complex 2D spectrum, take real() or imag() to get absorptive
                  or refractive 2D spectrum
* 'full2d_r'    : rephasing complex 2D spectrum
* 'full2d_nr'   : non-rephasing complex 2D spectrum
* 'gsb'         : complex 2D spectrum with only GSB pathways
* 'se'          : complex 2D spectrum with only SE pathways
* 'esa'         : complex 2D spectrum with only ESA pathways
* '..._r'       : rephasing part
* '..._nr'      : non-rephasing part
* 'corr'        : correlation map (in t1 and t3)
"""
mutable struct out2d{T}
    ω
    full2d::T
    full2d_r::T
    full2d_nr::T
    gsb::T
    gsb_r::T
    gsb_nr::T
    se::T
    se_r::T
    se_nr::T
    esa::T
    esa_r::T
    esa_nr::T
    corr
end

"""
    create_subspace(H, manifold, F, dat...)

Function accepts variable number of arguments (dat) in addition to H and returns tuple
of said arguments in Eigen basis (exciton basis/polariton basis/...).

# Arguments
* 'H'        : Hamiltonian
* 'manifold' : up to which excitation manifold/sector should the subspace extent, "si", "bi", "bi_lowest", "bi_polariton"
* 'F'        : operator list, e.g. list of Lindblad dissipation operators
* 'dat'      : either a single operator, or a tuple of operators, e.g.:

H, ..., dat_a                = create_subspace(H, ..., dat_a )

H, ..., dat_a, dat_b         = create_subspace(H, ..., dat_a, dat_b)

H, ..., dat_a, dat_b, dat_c  = create_subspace(H, ..., dat_a, dat_b, dat_c)

...

# Output
* 'H_sub'       : subspace Hamiltonian
* 'transf_op'   : operator to convert between Eigen and site basis
* 'P'           : projection operator
* 'F_sub'       : operator list, e.g. Lindblad dissipation operators, in subspace
* 'dat_sub'     : single or tuple of operators in subspace

`transf_op` can be used to convert between Eigen basis and site basis, e.g.:

`H_site    = transf_op     * H_eigen     * transf_op'`
"""
function create_subspace(H, manifold, dat...)
    Eivecs = eigvecs(dense(H[1]).data)
    L = length(H[1].basis_l.bases)
    # size of subspace
    if manifold == "bi"
        n_sub = 1 + L + binomial(L,2)
    elseif manifold == "si"
        n_sub = 1 + L + 1
    end

    # eigenvectors of subspace
    Eivecs_sub = [Ket(H[1].basis_l,Eivecs[:,i]) for i in 1:n_sub]
    # subspace basis
    b_sub = SubspaceBasis(H[1].basis_l,Eivecs_sub)
    # projection operators
    P = sparse(projector(b_sub, (H[1].basis_l)))
    Pt = dagger(P)
    println("dim(superspace): ", *(H[1].basis_l.shape...))
    # output new dimension
    println("dim(subspace): ", length(b_sub))
    H_sub   = [P * H[i] * Pt for i in 1:length(H)]
    dat_sub = [P * dat[i] * Pt for i in 1:length(dat)]
    println(dat)
#    F_sub    = [P * F[i] * Pt for i in 1:length(F)]
    return H_sub, Eivecs, P, dat_sub...;
end

function create_subspace(H, manifold, F, dat...)
    Eivecs = eigvecs(dense(H[1]).data)
    L = length(H[1].basis_l.bases)
    # size of subspace
    if manifold == "bi"
        #n_sub = 1 + L + binomial(L,2)
        n_sub = collect(1:1+L+binomial(L,2)+4)   #TODO: +1 for 2nd fock state... w/o +1 only true for all TLSs (cavity and matter)
    elseif manifold == "bi_lowest"
        #n_sub = 1 + L + 1                       # just use a single bi-excitonic state as an approximation
        n_sub = collect(1:1+L+1)
    elseif manifold == "bi_polariton"
        n_sub = append!(collect(1:1+L+1),1+L+binomial(L,2)+1)   # just use bi-excitonic states that are coupled to field ...
    elseif manifold == "si"
        #n_sub = 1 + L
        n_sub = collect(1:1+L)
    end

    # eigenvectors of subspace
    Eivecs_sub = [Ket(H[1].basis_l,Eivecs[:,i]) for i in n_sub]
    # subspace basis
    b_sub = SubspaceBasis(H[1].basis_l,Eivecs_sub)
    # projection operators
    P = sparse(projector(b_sub, (H[1].basis_l)))
    Pt = dagger(P)
    println("dim(superspace): ", *(H[1].basis_l.shape...))
    # output new dimension
    println("dim(subspace): ", length(b_sub))
    H_sub   = [P * H[i] * Pt for i in 1:length(H)]
    dat_sub = [P * dat[i] * Pt for i in 1:length(dat)]
    F_sub   = [P * F[i] * Pt for i in 1:length(F)]

    transf_op = Operator(H[1].basis_l, H[1].basis_r, Eivecs)
    transf_op = P * transf_op * Pt
    #transf_op = Eivecs
    return H_sub, transf_op, P, F_sub, dat_sub...;
end

"""
    out2d = make2dspectra(tlist, rho0, H, F, μ12, μ23, method; debug=false, zp=0)

make2Dspectra.jl calculates 2D spectra (GSB, SE, and ESA) using 2D correlation
functions by evaluating the Lindblad master equation.

# Arguments
* 'tlist' : time steps at which to evaluate correlation functions τ and t
* 'rho0'  : ground state density matrix
* 'H'     : Hamiltonian of the system
* 'F'     : collapse operators (i.e. relaxation pathways during τ and t) OR
            relaxation tensor R for Redfield
* 'μ12'   : transition dipole operator between ground- and excited states,
            responsible for GSB and SE
* 'μ23'   : transition dipole operator between excited- and higher excited states,
            responsible for ESA
* 'method': select either Lindblad or Redfield master equation solver
* 'debug' : if true, plot response functions and spectra of GSB, SE and ESA and
            show intermediate density matrices
* 'zp'    : zero pad data up to 2^zp in both dimensions. Improves data quality
            after fft, whenever timeevolution has already decayed to 0. Remember
            to also scale tlist and ω via function "interpt.jl". Default zp=0,
            data is not zeropadded.

Hidden is a lineshape function (work in progress).
"""
function make2Dspectra(tlist, rho0, H, F, μ12, μ23, T, method; debug=false, use_sub=true, zp=0, t2coh="kin")

    if use_sub
    H_si    = [[pop!(H)]]
    H_bi    = [[pop!(H)]]
    rho0_si = pop!(rho0)
    rho0_bi = pop!(rho0)
    F_si    = pop!(F)
    F_bi    = pop!(F)
    μ12_si  = pop!(μ12)
    μ12_bi  = pop!(μ12)
    μ23_si  = pop!(μ23)
    μ23_bi  = pop!(μ23)
    else
        H_si = [[H]]
        H_bi = [[H]]
        rho0_si = rho0
        rho0_bi = rho0
        F_si = F
        F_bi = F
        μ12_si = μ12
        μ12_bi = μ12
        μ23_si = μ23
        μ23_bi = μ23
    end

    println(rho0)

    println("############# ########################### #############");
    println("############# calculate R and NR pathways #############");
    println("############# ########################### #############\n");

    #TODO: implement T-dependent H
    N = 1
    # if !isempty(methods(H))
    #     H = [H(T) for i in 1:N]
    #     println.(H)
    # elseif length(H) == 2
    #     println("L2")
    #     #println(H)
    #     N = length(H[1])
    #     H = [[H[1][i], H[2][i]] for i in 1:N]
    # else
    #     N = 1
    #     H = [[H]]
    # end

    # keep around for now
    if use_sub
        println("Using system subspace!")
        #H_bi, μ12_bi, μ23_bi, rho0_bi, F_bi =  create_subspace(H, μ12, μ23, rho0, F, "bi")
    else
    #    H_bi, μ12_bi, μ23_bi, rho0_bi, F_bi = H, μ12, μ23, rho0, F
    end

    # TODO: can use a 2D lineshape function to cause inhomogeneous broadening, but need to decrease rate
    # of Lindblad operators beforehand ... very narrow transition, then broadened by lineshape function
    # inhomogeneously
    D = .0001
    τc = 20
    tc = 20 #+ T
    #τc = D / 0.000125
    #tc = D / 0.000125
    #gτ = D^2 * τc^2 * (exp.(-tlist./τc) .+ tlist./τc .- 1)
    gτ_(τ) = D^2 * τc^2 * (exp.(-τ./τc) .+ τ./τc .- 1)
    #gt = D^2 * tc^2 * (exp.(-tlist./tc) .+ tlist./tc .- 1)
    gt_(t) = D^2 * tc^2 * (exp.(-t./tc) .+ t./tc .- 1)
    test = true

    # excited state absorption
    corr_func_NR_esa = sum(correlations(tlist,rho0_bi,H_bi[i],F_bi,μ12_bi,μ23_bi,T,
                            "NR_esa",method,false; t2coh=t2coh) for i in 1:N) ./ N
    corr_func_R_esa  = sum(correlations(tlist,rho0_bi,H_bi[i],F_bi,μ12_bi,μ23_bi,T,
                            "R_esa", method,false; t2coh=t2coh) for i in 1:N) ./ N

    # use lineshape function if test == true
    if test
        corr_func_NR_esa = corr_func_NR_esa .* (exp.(-gτ_(-tlist)) * exp.(-gt_(-tlist))')
        corr_func_R_esa  = corr_func_R_esa  .* (exp.(-gτ_( tlist)) * exp.(-gt_( tlist))')
    end

    # excited state absorption ... from GS #TODO!: Can't use new Ham(T) here, needs to be same as for GSB
    corr_func_NR_esax = sum(correlations(tlist,rho0_bi,H_bi[i],F_bi,μ12_bi,μ23_bi,T,
                                "NR_esax",method,false; t2coh=t2coh) for i in 1:N) ./ N
    corr_func_R_esax  = sum(correlations(tlist,rho0_bi,H_bi[i],F_bi,μ12_bi,μ23_bi,T,
                                "R_esax", method,false; t2coh=t2coh) for i in 1:N) ./ N

    # use lineshape function if test == true
    if test
        corr_func_NR_esax = corr_func_NR_esax .* (exp.(-gτ_(-tlist)) * exp.(-gt_(-tlist))')
        corr_func_R_esax  = corr_func_R_esax  .* (exp.(-gτ_( tlist)) * exp.(-gt_( tlist))')
    end

    # keep this around for now
    if use_sub
       # H_si, μ12_si, μ23_si, rho0_si, F_si =  create_subspace(H, μ12, μ23, rho0, F, "si")
    else
    #    H_si, μ12_si, μ23_si, rho0_si, F_si = H, μ12, μ23, rho0, F
    end

    # ground state bleach
    corr_func_NR_gsb = sum(correlations(tlist,rho0_si,H_si[i],F_si,μ12_si,μ23_si,T,
                            "NR_gsb",method,false; t2coh=t2coh) for i in 1:N) ./ N
    corr_func_R_gsb  = sum(correlations(tlist,rho0_si,H_si[i],F_si,μ12_si,μ23_si,T,
                            "R_gsb", method,false; t2coh=t2coh) for i in 1:N) ./ N

    # use lineshape function if test == true
    if test
        corr_func_NR_gsb = corr_func_NR_gsb .* (exp.(-gτ_(-tlist)) * exp.(-gt_(-tlist))')
        corr_func_R_gsb  = corr_func_R_gsb  .* (exp.(-gτ_( tlist)) * exp.(-gt_( tlist))')
    end

    # stimulated emission
    corr_func_NR_se = sum(correlations(tlist,rho0_si,H_si[i],F_si,μ12_si,μ23_si,T,
                            "NR_se",method,false; t2coh=t2coh) for i in 1:N) ./ N
    corr_func_R_se  = sum(correlations(tlist,rho0_si,H_si[i],F_si,μ12_si,μ23_si,T,
                            "R_se", method,false; t2coh=t2coh) for i in 1:N) ./ N

    # use lineshape function if test == true
    if test
        corr_func_NR_se = corr_func_NR_se .* (exp.(-gτ_(-tlist)) * exp.(-gt_(-tlist))')
        corr_func_R_se  = corr_func_R_se  .* (exp.(-gτ_( tlist)) * exp.(-gt_( tlist))')
    end

    # use alternative way of calculating signal
    # was used for testing correct calculation of full 2D signals from components
    use_alt = false

    # calculate 2D spectrum by adding time-domain (t1,t3) coherence maps
    if use_alt
        # combine coherences into going from neg to pos t1
        corr_gsb = hcat(reverse(corr_func_R_gsb,dims=2), corr_func_NR_gsb[:,2:end])
        corr_se  = hcat(reverse(corr_func_R_se,dims=2),  corr_func_NR_se[:,2:end])
        corr_esa = hcat(reverse(corr_func_R_esa,dims=2), corr_func_NR_esa[:,2:end])
        corr_esax = hcat(reverse(corr_func_R_esax,dims=2),corr_func_NR_esax[:,2:end])
        # add components
        corr = corr_gsb + corr_se - corr_esa - corr_esax
        # convert to 2D spectra
        spec2dd        = corr2spec(corr, zp)
        spec2dd_gsb    = corr2spec(corr_gsb, zp)
        #spec2dd_R_gsb  = fft(corr_func_R_gsb)
        #spec2dd_NR_gsb = corr2spec(corr_func_NR_gsb, zp)
        spec2dd_se     = corr2spec(corr_se, zp)
        #spec2dd_R_se   = corr2spec(corr_func_R_se, zp)
        #spec2dd_NR_se  = corr2spec(corr_func_NR_se, zp)
        spec2dd_esa    = -corr2spec(corr_esa, zp)
        #spec2dd_R_esa  = corr2spec(corr_func_R_esa, zp)
        #spec2dd_NR_esa = corr2spec(corr_func_NR_esa, zp)
        spec2dd_esax = -corr2spec(corr_esax, zp)

        #spec2d_NR_gsb = []
        #spec2d_R_gsb  = []
        #spec2d_NR_se  = []
        #spec2d_R_se   = []
        #spec2d_NR_esa = []
        #spec2d_R_esa  = []
        #spec2d_NR_esax = []
        #spec2d_R_esax  = []

    else
        # define corr=[] in case
        corr = []

    end

        # divide first value by .5 (see Hamm and Zanni, section 9.5.3
        # Fourier transform, "http://2d-ir-spectroscopy.com/".)
        corr_func_NR_gsb[1,:] = corr_func_NR_gsb[1,:] / 2
        corr_func_NR_gsb[:,1] = corr_func_NR_gsb[:,1] / 2
        corr_func_R_gsb[1,:]  = corr_func_NR_gsb[1,:] / 2
        corr_func_R_gsb[:,1]  = corr_func_NR_gsb[:,1] / 2

        corr_func_NR_se[1,:] = corr_func_NR_se[1,:] / 2
        corr_func_NR_se[:,1] = corr_func_NR_se[:,1] / 2
        corr_func_R_se[1,:]  = corr_func_NR_se[1,:] / 2
        corr_func_R_se[:,1]  = corr_func_NR_se[:,1] / 2

        corr_func_NR_esa[1,:] = corr_func_NR_esa[1,:] / 2
        corr_func_NR_esa[:,1] = corr_func_NR_esa[:,1] / 2
        corr_func_R_esa[1,:]  = corr_func_NR_esa[1,:] / 2
        corr_func_R_esa[:,1]  = corr_func_NR_esa[:,1] / 2

        corr_func_NR_esax[1,:] = corr_func_NR_esax[1,:] / 2
        corr_func_NR_esax[:,1] = corr_func_NR_esax[:,1] / 2
        corr_func_R_esax[1,:]  = corr_func_NR_esax[1,:] / 2
        corr_func_R_esax[:,1]  = corr_func_NR_esax[:,1] / 2

        # zeropad data prior to fft to increase resolution
        corr_func_NR_gsb = zeropad(corr_func_NR_gsb,zp)
        corr_func_R_gsb  = zeropad(corr_func_R_gsb ,zp)

        corr_func_NR_se  = zeropad(corr_func_NR_se,zp)
        corr_func_R_se   = zeropad(corr_func_R_se ,zp)

        corr_func_NR_esa = zeropad(corr_func_NR_esa,zp)
        corr_func_R_esa  = zeropad(corr_func_R_esa ,zp)

        corr_func_NR_esax = zeropad(corr_func_NR_esax,zp)
        corr_func_R_esax  = zeropad(corr_func_R_esax ,zp)

        #######################################################
        ######### plot 2nd order correlation function #########
        #######################################################

        println("done\n")

        println("############# ########################### #############");
        println("############ Fourier transform of pathways ############");
        println("############# ########################### #############\n");

        spec2d_NR_gsb = fftshift(fft((corr_func_NR_gsb)))
        spec2d_R_gsb  = fftshift(fft((corr_func_R_gsb)))

        spec2d_NR_se  = fftshift(fft((corr_func_NR_se)))
        spec2d_R_se   = fftshift(fft((corr_func_R_se)))

        spec2d_NR_esa = -fftshift(fft((corr_func_NR_esa)))
        spec2d_R_esa  = -fftshift(fft((corr_func_R_esa)))

        spec2d_NR_esax = -fftshift(fft((corr_func_NR_esax)))
        spec2d_R_esax  = -fftshift(fft((corr_func_R_esax)))

        println("done\n")

        println("############# ########################### #############");
        println("########## Construct absorptive 2D spectrum ###########");
        println("############# ########################### #############\n");

        spec2d_R_gsb  = circshift(reverse(spec2d_R_gsb ,dims=1),(1,0))
        spec2d_R_se   = circshift(reverse(spec2d_R_se  ,dims=1),(1,0))
        spec2d_R_esa  = circshift(reverse(spec2d_R_esa ,dims=1),(1,0))
        spec2d_R_esax = circshift(reverse(spec2d_R_esax ,dims=1),(1,0))

    # calculate the absorptive spectrum of GSB, ESA, and SE
    spec2d_gsb  = spec2d_NR_gsb  + spec2d_R_gsb
    spec2d_se   = spec2d_NR_se   + spec2d_R_se
    spec2d_esa  = spec2d_NR_esa  + spec2d_R_esa
    spec2d_esax = spec2d_NR_esax + spec2d_R_esax

    # sum rephasing and non-rephasing spectra individually
    spec2d_r =  spec2d_R_gsb +
                spec2d_R_se  +
                spec2d_R_esa +
                spec2d_R_esax

    spec2d_nr = spec2d_NR_gsb +
                spec2d_NR_se  +
                spec2d_NR_esa +
                spec2d_NR_esax;

    # Subtract esax from gsb. Additional channel to simulate ground state recovery
    spec2d_gsb = spec2d_gsb + spec2d_esax

    # calculate the total 2D spectrum (complex)
    spec2d = spec2d_gsb +
             spec2d_se  +
             spec2d_esa

    println("done\n")

    #tlist = [tlist; tlist[2:end] .+ tlist[end]] # BUG: has to do with alternative way of calculating 2D signal
    tlist, ω = interpt(tlist,zp)

    # could return corr for looking at it ... set [] to save space
    #corr = []

    #FACT: spec2dd behaves correctly (for coupled dimer) when plotting absorptive, absolute, dispersive and phase // spec2d now as well ! ...
    #TODO: absl is stretched too much in ω3!
    if use_alt
        #BUG #TODO: So far I need to change tlist and then crop ω when using alternative way of calculating 2D signals
        ω = ω[end÷2+1:end]

        out = out2d{Array{ComplexF32,2}}(ω, spec2dd, spec2d_r, spec2d_nr, spec2dd_gsb, spec2d_R_gsb,
                                     spec2d_NR_gsb, spec2d_se, spec2d_R_se, spec2d_NR_se, spec2d_esa,
                                     spec2d_R_esa, spec2d_NR_esa, corr)
    else
        out = out2d{Array{ComplexF32,2}}(ω, spec2d, spec2d_r, spec2d_nr, spec2d_gsb, spec2d_R_gsb,
                                     spec2d_NR_gsb, spec2d_se, spec2d_R_se, spec2d_NR_se, spec2d_esa,
                                     spec2d_R_esa, spec2d_NR_esa, corr)
    end

    out = round2d(out,2)

    return out
end

"""
    function crop2d(data_struct,w_min;w_max=100,step=1)

# Arguments
* 'data_struct'     : full data structure out2d
* 'w_min'           : lower cut in frequency / energy
* 'w_max'           : upper cut in frequency / energy
* 'step'            : step in frequency dimension (in data points)

Use to crop 2D spectra to smaller size.
"""
function crop2d(data_struct,w_min;w_max=100,step=1)

    min_i = argmin(abs.(data_struct.ω .- w_min))
    max_i = argmin(abs.(data_struct.ω .- w_max))

    data_struct.ω         = data_struct.ω[min_i:step:max_i]
    data_struct.full2d    = data_struct.full2d[min_i:step:max_i,min_i:step:max_i]
    data_struct.full2d_r  = data_struct.full2d_r[min_i:step:max_i,min_i:step:max_i]
    data_struct.full2d_nr = data_struct.full2d_nr[min_i:step:max_i,min_i:step:max_i]
    data_struct.gsb       = data_struct.gsb[min_i:step:max_i,min_i:step:max_i]
    data_struct.gsb_r     = data_struct.gsb_r[min_i:step:max_i,min_i:step:max_i]
    data_struct.gsb_nr    = data_struct.gsb_nr[min_i:step:max_i,min_i:step:max_i]
    data_struct.se        = data_struct.se[min_i:step:max_i,min_i:step:max_i]
    data_struct.se_r      = data_struct.se_r[min_i:step:max_i,min_i:step:max_i]
    data_struct.se_nr     = data_struct.se_nr[min_i:step:max_i,min_i:step:max_i]
    data_struct.esa       = data_struct.esa[min_i:step:max_i,min_i:step:max_i]
    data_struct.esa_r     = data_struct.esa_r[min_i:step:max_i,min_i:step:max_i]
    data_struct.esa_nr    = data_struct.esa_nr[min_i:step:max_i,min_i:step:max_i]

    return data_struct

end

"""
    function round2d(data_struct,digits=2)

# Arguments
* 'data_struct' : from out2d
* 'digits'      : the precision

Use to round 2D spectra to smaller size.
"""
function round2d(data_struct,digits=2)

    data_struct.ω         = round.(data_struct.ω,digits=6)
    data_struct.full2d    = round.(data_struct.full2d,digits=digits)
    data_struct.full2d_r  = round.(data_struct.full2d_r,digits=digits)
    data_struct.full2d_nr = round.(data_struct.full2d_nr,digits=digits)
    data_struct.gsb       = round.(data_struct.gsb,digits=digits)
    data_struct.gsb_r     = round.(data_struct.gsb_r,digits=digits)
    data_struct.gsb_nr    = round.(data_struct.gsb_nr,digits=digits)
    data_struct.se        = round.(data_struct.se,digits=digits)
    data_struct.se_r      = round.(data_struct.se_r,digits=digits)
    data_struct.se_nr     = round.(data_struct.se_nr,digits=digits)
    data_struct.esa       = round.(data_struct.esa,digits=digits)
    data_struct.esa_r     = round.(data_struct.esa_r,digits=digits)
    data_struct.esa_nr    = round.(data_struct.esa_nr,digits=digits)

    return data_struct

end

#CHECK I am considering R + conj(R) ... seems that 2D spectra turn out distorted
"""
    corr = correlations(tlist, rho0, H, F, μa, μb, pathway, method, debug; t2)

Calculate the response functions necessary for 2D spectroscopy.

# Arguments
* 'tlist'   : time steps at which to evaluate correlation functions τ and t
* 'rho0'    : ground state density matrix
* 'H'       : Hamiltonian of the system
* 'F'       : collapse operators (i.e. relaxation pathways during τ and t) OR
              relaxation tensor R for Redfield
* 'μa'      : transition dipole operator between ground- and excited states,
              responsible for GSB and SE
* 'μb'      : transition dipole operator between excited- and higher excited
              states, responsible for ESA
* 'pathway' : select response pathway 'R', or 'NR' _ 'gsb', 'se', or 'esa'
              'NR_gsb', 'R_esa', etc.
* 'method'  : select either Lindblad or Redfield master equation solver
* 'debug'   : if true, plot response functions and spectra of GSB, SE and ESA and
              show intermediate density matrices
* 't2'      : behaviour during t2 time "kin", "vib" or "full"

# Output
* 'corr'    : 2D correlation function

# Additional considerations:

* GS        : ground state
* ESs       : excited state or states
* FS        : final state
* μa and μb : TDMs, coupling states, such that GS <-μa-> ESs <-μb-> FS
* F         : is either list of collapse operators (c_ops; Lindblad) or
              a relaxation tensor (R; Redfield)

# Pathways calculated (w/o cc)
rephasing:
```
GSB               { μ * [(rho0 * μ)  * μ]}  + cc
            copy: ( μ * ((rho0 * μ)  * μ))
SE                {[μ *  (rho0 * μ)] * μ}   + cc
            copy: ((μ *  (rho0 * μ)) * μ)
ESA               { μ *  [μ * (rho0  * μ)]} + cc
            copy: ( μ *  (μ * (rho0  * μ)))
```

non-rephasing:
```
GSB               {μ * [μ * (μ * rho0)]} + cc
            copy: (μ * (μ * (μ * rho0)))
SE                {[(μ * rho0) * μ] * μ} + cc
            copy: (((μ * rho0) * μ) * μ)
ESA               {μ * [(μ * rho0) * μ]} + cc
            copy: (μ * ((μ * rho0) * μ))
```
"""
function correlations(tlist, rho0, H, F, μ_ge, μ_ef, T, pathway, method, debug; t2coh=false)

    #=
    #IDEA either create subspace in function that generates Hamiltonian, or create in t_ev
    try
        if pathway[end] == 'a' || pathway[end] == 'x'
            H, μa, μb, rho0, F =  create_subspace(H, μa, μb, rho0, F, "bi")
        elseif pathway[end] == 'b' || pathway[end] == 'e'
            H, μa, μb, rho0, F =  create_subspace(H, μa, μb, rho0, F, "si")
        else
            println("error creating subspace")
        end
    catch
        println("error creating subspace")
    end
    =#

    ## use full transition dipole operator matrix
    #μ = μa + μb #DELETE
    ## or not ...
    #μ_ge = μa
    #μ_ef = μb

    #TODO!: T-dependent H not working with Redfield
    if length(H) == 1
        H   = H[1]
        H_T = H
    elseif length(H) == 2
        H_T = H[2]
        H   = H[1]
    else
        println("error T-dep H")
    end

    ## switch between Lindblad and Redfield master equation
    t_ev(A,B,C,D) =
        if method == "lindblad"
            #if isempty(methods(H))
                timeevolution.master(A,B,C,D)
            #elseif !isempty(methods(H))
            #    timeevolution.master_dynamic(A,B,C) #TODO: implement T-dep Hamiltonian, need to figure out how to create subspace when H is passed as a function
            #end
        elseif method == "redfield"
            timeevolution.master_bloch_redfield(A,B,D,C)
        else
            error("MultidimensionalSpectroscopy.jl(correlations) -- No valid method selected!")
        end

    ## add functionality to vary T
    T = [0.:1.:T;]

    # initialize output matrix
    corr    = complex(zeros(length(tlist),length(tlist)))
    corr_cc = complex(zeros(length(tlist),length(tlist)))

    # use right-hand excitation for rephasing and left-hand excitation for
    # non-rephasing to ensure that emission is from the left of the double-
    # sided vertical Feynman diagram
    if pathway[1] == 'R'        # first interaction from right if rephasing
        #rho1    = rho0 * μ_ge
        rho1    = rho0 * tri(μ_ge,"U")
    elseif pathway[1] == 'N'    # first interaction from left if non-rephasing
        #rho1    = μ_ge * rho0
        rho1    = tri(μ_ge,"L") * rho0
    else
        error("MultidimensionalSpectroscopy.jl(correlations) -- No pathway selected!")
    end

    # define transition dipole operator to evaluate expectation value
    if pathway[end] == 'a' # if esa pathway use μb
        #μexp = -μb         # minus sign follows from Feynman diagrams
        μexp    = tri(μ_ef,"U") # ... removed -1
    elseif pathway[end] == 'e'
        μexp    = tri(μ_ge,"U")
    elseif pathway[end] == 'x'
        μ_ef    =  μ_ge
        μexp    =  tri(μ_ge,"U")
        pathway = chop(pathway)
    else
        μexp    =  tri(μ_ge,"U")
    end

    # calculate evolution during t1(τ)
    τ = tlist; tout, rho1_τ    = t_ev(τ,rho1,H,F)

    # add progress bar #TODO: make compatible with multiple Threads
    R = length(collect(1:20:length(rho1_τ)))
    Tpr = T[end]
    print(Threads.threadid())
    print(string(rpad(pathway,6), " at T = $Tpr fs. Progress: [" * repeat(" ", R-1)) * "]")
    print(repeat("\b", R))

    # TODO: damping function τ, implement somehow in redfield ??
    """
    τc = 8
    tc = 20
    gτ = (exp.(-τ./τc) .+ τ./τc .- 1)
    gt = (exp.(-τ./tc) .+ τ./tc .- 1)
    rho1_τ = rho1_τ .* exp.(-gτ)
    rho1_τ_cc = rho1_τ_cc .* exp.(-gt)
    """

    # now, use every rho(τ) as starting point to calculate second coherence
    # time (t) dependence. Note that first interaction with μ puts density
    # matrix back into a population without any time delay to the next
    # interaction (T=0)!
    for i=1:length(rho1_τ)

        if pathway == "R_gsb"
            #------------------------#
            # {μ * [(rho0 * μ) * μ]}
            # copy: (μ * ((rho0 * μ) * μ))
            #------------------------#
            #rho2    =   rho1_τ[i]     * μ_ge
            rho2    =   rho1_τ[i]    * tri(μ_ge,"L")
            rho2a   =   rho2
            # eliminate on-diagonal or off-diagonal elements
            if t2coh == "kin"
                rho2.data = tril(triu(rho2.data))
            elseif t2coh == "vib"
                rho2.data = tril(rho2.data,-1) + triu(rho2.data,1)
            elseif t2coh == "full"
                # do nothing
            end
            rho2b = rho2
            # time evolution during t2(T)-time
            if T[end] != 0
                tout, rho2_T = t_ev(T,rho2,H,F)
                rho2 = rho2_T[end]
            end
            # third field interaction
            #rho3   = μ_ge * rho2
            rho3   = tri(μ_ge,"L")    * rho2

        elseif pathway == "NR_gsb"
            #-------------------------#
            # {μ * [μ * (μ * rho0)]}
            # copy: (μ * (μ * (μ * rho0)))
            #-------------------------#
            #rho2 = μ_ge * rho1_τ[i]
            rho2 = tri(μ_ge,"U") * rho1_τ[i]

            if t2coh == "kin"
                rho2.data = tril(triu(rho2.data))
            elseif t2coh == "vib"
                rho2.data = tril(rho2.data,-1) + triu(rho2.data,1)
            elseif t2coh == "full"
                # do nothing
            end

            if T[end] != 0
                tout, rho2_T = t_ev(T,rho2,H,F)
                rho2 = rho2_T[end]
            end

            #rho3   = μ_ge * rho2
            rho3   = tri(μ_ge,"L") * rho2

        elseif pathway == "R_se"
            #------------------------------#
            # {[μ * (rho0 * μ)] * μ} + cc
            # copy: ((μ * (rho0 * μ)) * μ)
            #------------------------------#
            #rho2 = μ_ge * rho1_τ[i]
            rho2 = tri(μ_ge,"L") * rho1_τ[i]

            if t2coh == "kin"
                rho2.data = tril(triu(rho2.data))
            elseif t2coh == "vib"
                rho2.data = tril(rho2.data,-1) + triu(rho2.data,1)
            elseif t2coh == "full"
                # do nothing
            end

            if T[end] != 0
                tout, rho2_T = t_ev(T,rho2,H,F)
                rho2 = rho2_T[end]
            end

            #rho3   =       rho2 * μ_ge
            rho3   =       rho2 * tri(μ_ge,"L")

        elseif pathway == "NR_se"
            #------------------------------#
            # {[(μ * rho0) * μ] * μ} + cc
            # copy: (((μ * rho0) * μ) * μ)
            #------------------------------#
            #rho2 =      rho1_τ[i]   * μ_ge
            rho2 =      rho1_τ[i]   * tri(μ_ge,"U")

            if t2coh == "kin"
                rho2.data = tril(triu(rho2.data))
            elseif t2coh == "vib"
                rho2.data = tril(rho2.data,-1) + triu(rho2.data,1)
            elseif t2coh == "full"
                # do nothing
            end

            if T[end] != 0
                tout, rho2_T = t_ev(T,rho2,H,F)
                rho2 = rho2_T[end]
            end

            #rho2.data[1,1] = 0

            #rho3   =       rho2 * μ_ge
            rho3   =       rho2 * tri(μ_ge,"L")

        elseif pathway == "R_esa"
            #μ_ef = μ_ge# + μ_ef # after relaxation, reabsorption is part of esa (!?)
            #------------------------------#
            # {μ * [μ * (rho0 * μ)]} + cc
            # copy: (μ * (μ * (rho0 * μ)))
            #------------------------------#
            #rho2 = μ_ge * rho1_τ[i]
            rho2 = tri(μ_ge,"L")    * rho1_τ[i]

            if t2coh == "kin"
                rho2.data = tril(triu(rho2.data))
            elseif t2coh == "vib"
                rho2.data = tril(rho2.data,-1) + triu(rho2.data,1)
            elseif t2coh == "full"
                # do nothing
            end

            if T[end] != 0
                tout, rho2_T = t_ev(T,rho2,H,F)
                rho2 = rho2_T[end]
            end

            #rho3   = μ_ef * rho2
            rho3    = tri(μ_ef,"L") * rho2

        elseif pathway == "NR_esa"
            #μ_ef = μ_ge #+ μ_ef
            #-----------------------------#
            # {μ * [(μ * rho0) * μ]} + cc
            # copy: (μ * ((μ * rho0) * μ))
            #-----------------------------#
            #rho2 =      rho1_τ[i]   * μ_ge
            rho2 =      rho1_τ[i]   * tri(μ_ge,"U")
            #rho2 =      μ_ge * rho1_τ[i]

            if t2coh == "kin"
                rho2.data = tril(triu(rho2.data))
            elseif t2coh == "vib"
                rho2.data = tril(rho2.data,-1) + triu(rho2.data,1)
            elseif t2coh == "full"
                # do nothing
            end

            if T[end] != 0
                tout, rho2_T = t_ev(T,rho2,H,F)
                rho2 = rho2_T[end]
            end

            #rho3    = μ_ef * rho2
            rho3    = tri(μ_ef,"L") * rho2
        end

        # calculate time evolution of ρ during detection time t
        t = tlist; tout, rho3_t     = t_ev(t,rho3,H_T,F)

        # calc. exp. v. for GSB and SE (μexp = μ12) and for ESA (μexp = μ23)
        corr[i,:]    = (expect(μexp,conj(rho3_t))) # .* exp.(-gt) #TODO: lineshape f
        # ALTERNATIVELY take the trace: Tr(Aρ)
        # for j=1:length(t)
        #     corr[i, j] = tr(μexp.data * real(rhot_temp[j].data))
        # end

        #TODO: damping function t
        #D = .4
        #τc = 10
        #gt = (exp.(-t./τc) .+ t./τc .- 1)
        #corr[i,:] = corr[i,:] .* exp.(-gt)

        # progress bar update
        if mod(i,20) == 0
            print("=")
        end

        # print a few combinations to check in debug mode
        if debug
            if i in [1]
                println("\n")
                println("###################################################")
                println("###################### DEBUG ######################")
                println("###################################################\n")
                println("showing density matrices from during the evolution.\n")
                println("Which path ?")
                println(pathway)
                println(" ")
                println("μ_ge")
                println(dense(μ_ge))
                println(" ")
                println("μ_ef")
                println(dense(μ_ef))
                println(" ")
                println("rho0")
                println(rho0)
                println(" ")
                println("rho1")
                println(rho1)
                println(" ")
                if T[end] != 0
                    println(" ")
                    println("rho2_T[1] (before T evolution)")
                    println(rho2_T[1])
                end
                println("rho2a")
                println(rho2a)
                println("rho2b")
                println(rho2b)
                println("rho2")
                println(rho2)
                println(" ")
                println("rho3")
                println(rho3)
                println(" ")
                println("μexp")
                println(μexp)
                println(" ")
                print("μexp * rho3")
                println(μexp * rho3)
                println(" ")
                println("\n")
            end
        end
    end
    print("\n")
    return corr
end

"""
    ω, corr, spec = absorptionSpectrum(t, rho, H, F, μ; zp=11)

Quick function to calculate absorption spectrum.

# Arguments
* 't'   :
* 'rho' :
* 'H'   :
* 'F'   :
* 'μ'   :
* 'zp'  :

"""
function absorptionSpectrum(t, rho, H, F, μ; zp=11)

    corr = timecorrelations.correlation(t, rho, H, F, μ, μ)
    corr_intp = zeropad(corr,zp)
    t_intp, ω_intp = interpt(t,zp)
    ω, spec = timecorrelations.correlation2spectrum(t, corr; normalize_spec=true)
    return ω, corr, spec
end

#FACT: tri is required for properly making ... σ⁺ ρ ... (tf does this mean?)
#FACT: deactivate tri for displaced harmonic oscillator
"""
    out = tri(dat, uplo)

function used internally by cmds.jl to select upper triagular (uplo="U") or lower
triangular (uplo="L") matrix.
"""
function tri(dat,uplo)
    out = copy(dat)
    if uplo == "L"
        out.data = tril(out.data)
    elseif uplo == "U"
        out.data = triu(out.data)
    end
    return out;
end

"""
    plot2d_comps(data)

Plots GSB, SE and ESA component of 2D spectrum.

# Arguments
* 'data'  : data struct for single time step T (e.g. out2d[1])

"""
function plot2d_comps(data)

    fig, ax = subplots(1,4,sharex=true,sharey=true,figsize=(15,3))

    sca(ax[1])
    plot2d(data.ω, data.full2d)
    PyPlot.title("absorptive")

    sca(ax[2])
    plot2d(data.ω, data.gsb)
    PyPlot.title("GSB")

    sca(ax[3])
    plot2d(data.ω, data.se)
    PyPlot.title("SE")

    sca(ax[4])
    plot2d(data.ω, data.esa)
    PyPlot.title("ESA")

    tight_layout()
end

#TODO: is this useful ?
function pretty_show_mat(data)
    fig, ax = plt.subplots()
    im = ax.imshow(data)
    for i in 1:size(data,1)
        for j in 1:size(data,2)
            ax.text(i-1,j-1,round(data[i,j],digits=2),ha="center",va="center")
        end
    end
end

"""
    vib_analysis(data)

Used to quickly plot GSB, SE and ESA rephasing (R) and non-rephasing (NR) spectra. Useful for checking
for beating patterns, but could also be used to look at individual components of "normal" 2D dataset.

# Arguments
* 'data'   : just the 2D data structure as (e.g. out2d[1])
"""
function vib_analysis(data;norm=0)
     fig, ax = subplots(3,2,sharex=true,sharey=true,figsize=(2*3.2,2.5*3))
     sca(ax[1,1])
     plot2d(data.ω,abs.(data.gsb_r); norm=norm)
     PyPlot.title("GSB R")
     sca(ax[1,2])
     plot2d(data.ω,abs.(data.gsb_nr); norm=norm)
     PyPlot.title("GSB NR")
     sca(ax[2,1])
     plot2d(data.ω,abs.(data.se_r); norm=norm)
     PyPlot.title("SE R")
     sca(ax[2,2])
     plot2d(data.ω,abs.(data.se_nr); norm=norm)
     PyPlot.title("SE NR")
     sca(ax[3,1])
     plot2d(data.ω,abs.(data.esa_r); norm=norm)
     PyPlot.title("ESA R")
     sca(ax[3,2])
     plot2d(data.ω,abs.(data.esa_nr); norm=norm)
     PyPlot.title("ESA NR")
     tight_layout()
end

"""
    plot_levels(H,center_x;col,ls)

Plots the energy levels of H at position center_x within an axis. Create a figure() first,
to plot the energy levels of different systems, subspaces, monomers, etc. together. E.g.

```julia
figure(); title("Energy level diagram")
cmds.plot_levels(H,0)
cmds.plot_levels(Hcav,-1)
cmds.plot_levels(Hexc,1)
cmds.plot_levels(Hmon,2)
xticks([-1, 0, 1, 2], ["cavity", "full", "matter", "monomer"])
```

"""
function plot_levels(H, center_x; col="k", ls="solid")
    # get energies and eigenvectors of full system
    E      = eigvals(dense(H).data)
    vecs   = eigvecs(dense(H).data)
    eivecs = [Ket(H.basis_l,vecs[:,i]) for i in 1:length(E)]
    # plot eigenenergies
    xn = 0
    lx = center_x - .45
    rx = center_x + .45
    for i in 1:length(E)
        if i > 1 #&& E[i] - E[i-1] < 0.1
            hlines(E[i],lx,rx,col,ls)
            xn += .02
        else
            hlines(E[i],lx,rx,col,ls)
            xn = 0
        end
    end

end

function logger(what, fn)
    open(fn,"a") do io
        # do stuff with the open file
        println(io,what)
    end
end

function logger(what,dat, fn)
    open(fn,"a") do io
        # do stuff with the open file
        println(io,what,dat)
    end
end

"""
    corr2spec(corr, zp)

Auxiliary function to convert correlation map to 2D spectrum.
"""
function corr2spec(corr, zp)
        corr = vcat(corr[end:-1:1,end:-1:1],corr[2:end,:])
        corr = vcat(zeros(size(corr,2))',corr)
        corr = hcat(corr,zeros(size(corr,1)))

        N = 2^zp
        temp = complex(zeros(N,N));

        temp[(N-size(corr,1))÷2+1:end-(N-size(corr,1))÷2,(N-size(corr,1))÷2+1:end-(N-size(corr,1))÷2] = corr
        corr = temp

        corr = circshift(corr,(0,1))

        spec = fft(fftshift(corr))
        spec = fftshift(spec)
        #spec = circshift(spec,(-100,-100))
        spec = spec[1:end÷2,1:end÷2]
        spec = reverse(spec,dims=1)
        spec = reverse(spec,dims=2)
        return spec
    end

end # module end

max_stars_repo_name: lbmnky/CMDS

max_issues_repo_name: lbmnky/CMDS

max_forks_repo_name: lbmnky/CMDS


text:
# Simple ridge regression based upon the sindy-mpc
# repository, see https://arxiv.org/abs/1711.05501
# and https://github.com/eurika-kaiser/SINDY-MPC/blob/master/LICENSE

mutable struct STRRidge{T} <: AbstractOptimiser
    λ::T
end

STRRidge() = STRRidge(0.1)

function set_threshold!(opt::STRRidge, threshold)
    opt.λ = threshold
end

get_threshold(opt::STRRidge) = opt.λ

init(o::STRRidge, A::AbstractArray, Y::AbstractArray) = A \ Y
init!(X::AbstractArray, o::STRRidge, A::AbstractArray, Y::AbstractArray) =  ldiv!(X, qr(A, Val(true)), Y)

function fit!(X::AbstractArray, A::AbstractArray, Y::AbstractArray, opt::STRRidge; maxiter::Int64 = 1)
    smallinds = abs.(X) .<= opt.λ
    biginds = @. ! smallinds[:, 1]
    for i in 1:maxiter
        smallinds .= abs.(X) .<= opt.λ
        X[smallinds] .= zero(eltype(X))
        for j in 1:size(Y, 2)
            biginds .= @. ! smallinds[:, j]
            X[biginds, j] .= A[:, biginds] \ Y[:,j]
        end
    end

    X[abs.(X) .< get_threshold(opt)] .= zero(eltype(X))
end


text:
"""Abstract index supertype"""
abstract type Index end

"""An index whose value is not known. Identified by its name. There should be
no need to set the id by the user.
"""
struct FreeIndex{T<:AbstractVectorSpace} <: Index
    V::T
    name::String
    id::Int
end

FreeIndex(V::T, name::String) where {T<:AbstractVectorSpace} = FreeIndex(V, name, 0)

"""Creates an index with the same name in the dual space."""
Base.adjoint(i::FreeIndex) = FreeIndex(dual(i.V), i.name, i.id)

"""An index whose value is known."""
struct FixedIndex{T<:AbstractVectorSpace} <: Index
    V::T
    value::Int
    function FixedIndex(V::T, value::Int) where {T<:AbstractVectorSpace}
        if dim(V) === nothing || value < 1 || value > dim(V)
            throw(DomainError(value, string("Value not in range for the vectorspace ", V)))
        end
        new{T}(V, value)
    end
end

"""Creates an index with the same value in the dual space."""
Base.adjoint(i::FixedIndex) = FixedIndex(dual(i.V), i.value)

"""A node for holding indices."""
struct Indices <: Node
    indices
    children
end

Indices(indices::Vararg{Index}) = Indices(indices, ())

"""Creates an index from an integer and VectorSpace."""
function toindex(V::AbstractVectorSpace, i::Int)
    return FixedIndex(V, i)
end

"""Creates an index from a string and VectorSpace."""
function toindex(V::AbstractVectorSpace, s::String)
    return FreeIndex(V, s)
end

"""Added for convenience."""
function toindex(V::AbstractVectorSpace, i::Index)
    return i
end

Base.show(io::IO, indices::Indices) = print(io, indices.indices)

function Base.show(io::IO, i::FreeIndex)
    if i.name != ""
        print(io, i.name)
    else
        print(io, "FreeIndex")
    end
    if i.V isa DualVectorSpace
        print(io, "*")
    end
end

Base.show(io::IO, i::FixedIndex) = print(io, i.value)

max_stars_repo_name: Sohva/TensorDSL

max_issues_repo_name: Sohva/TensorDSL

max_forks_repo_name: Sohva/TensorDSL


text:
reverse_dict(dict::AbstractDict) = Dict(v => k for (k, v) ∈ dict)

function_pointer(instance::Instance, name) = get_instance_proc_addr(name; instance)
function_pointer(device::Device, name) = get_device_proc_addr(device, name)


text:
using .Groebner: _normal_form_reference, _reducegb_reference, _reducegb_reference!,
                    change_ordering, rootn

@testset "Normal Form" begin

    R, (x, y, z) = PolynomialRing(QQ, ["x", "y", "z"])

    Gs = [
        [x], [x, y], [x, y, z]
    ]

    @test _normal_form_reference(x, Gs[1]) == 0
    @test _normal_form_reference(y, Gs[1]) == y
    @test _normal_form_reference(x + y^2, Gs[3]) == 0
    @test _normal_form_reference(5x^2 - y^3 - 1, Gs[2]) == -1
    @test _normal_form_reference(x + y + 3z, Gs[2]) == 3z

    G = [
        x^2 + y*x + 1,
        y^2 - y
    ]

    @test _normal_form_reference(x + y + z, G) == x + y + z
    @test _normal_form_reference(x^2 + y^2, G) == y - y*x - 1
    @test _normal_form_reference(y^3, G) == _normal_form_reference(y^4, G) == y

    G = [
        x*y + y,
        z + 5,
    ]
    @test _normal_form_reference(x*y, G) == -y
    @test _normal_form_reference(3z - 1, G) == -16

    @test _normal_form_reference(8x*y + 88z, [R(1)]) == 0

    R, (x, y, z) = PolynomialRing(QQ, ["x","y","z"], ordering=:degrevlex)
    G = [
        x^2 + y,
        y^2 + x
    ]
    @test _normal_form_reference(x^2 + y^2, G) == -x - y

end

@testset "Change ordering" begin
    R, (x, y, z) = PolynomialRing(QQ, ["x", "y", "z"])

    fs = [1x + 2y*z + 3z^3, y*z, 5x]
    fs_deg = change_ordering(fs, :degrevlex)
    R_deg = parent(first(fs_deg))

    @test ordering(R_deg) == :degrevlex
    @test collect(exponent_vectors(fs_deg[1])) == [[0, 0, 3], [0, 1, 1], [1, 0, 0]]
    @test collect(exponent_vectors(fs_deg[2])) == [[0, 1, 1]]
    @test collect(exponent_vectors(fs_deg[3])) == [[1, 0, 0]]

    fs_vv = change_ordering(fs, :lex)
    @test fs_vv == fs
end