text:
## R code 7.3
m7.1 <- quap(
    alist(
        brain_std ~ dnorm( mu , exp(log_sigma) ),
        mu <- a + b*mass_std,
        a ~ dnorm( 0.5 , 1 ),
        b ~ dnorm( 0 , 10 ),
        log_sigma ~ dnorm( 0 , 1 )
    ), data=d )

## R code 7.7
m7.2 <- quap(
    alist(
        brain_std ~ dnorm( mu , exp(log_sigma) ),
        mu <- a + b[1]*mass_std + b[2]*mass_std^2,
        a ~ dnorm( 0.5 , 1 ),
        b ~ dnorm( 0 , 10 ),
        log_sigma ~ dnorm( 0 , 1 )
    ), data=d , start=list(b=rep(0,2)) )

## R code 7.8
m7.3 <- quap(
    alist(
        brain_std ~ dnorm( mu , exp(log_sigma) ),
        mu <- a + b[1]*mass_std + b[2]*mass_std^2 +
                  b[3]*mass_std^3,
        a ~ dnorm( 0.5 , 1 ),
        b ~ dnorm( 0 , 10 ),
        log_sigma ~ dnorm( 0 , 1 )
    ), data=d , start=list(b=rep(0,3)) )

m7.4 <- quap(
    alist(
        brain_std ~ dnorm( mu , exp(log_sigma) ),
        mu <- a + b[1]*mass_std + b[2]*mass_std^2 +
                  b[3]*mass_std^3 + b[4]*mass_std^4,
        a ~ dnorm( 0.5 , 1 ),
        b ~ dnorm( 0 , 10 ),
        log_sigma ~ dnorm( 0 , 1 )
    ), data=d , start=list(b=rep(0,4)) )

m7.5 <- quap(
    alist(
        brain_std ~ dnorm( mu , exp(log_sigma) ),
        mu <- a + b[1]*mass_std + b[2]*mass_std^2 +
                  b[3]*mass_std^3 + b[4]*mass_std^4 +
                  b[5]*mass_std^5,
        a ~ dnorm( 0.5 , 1 ),
        b ~ dnorm( 0 , 10 ),
        log_sigma ~ dnorm( 0 , 1 )
    ), data=d , start=list(b=rep(0,5)) )

## R code 7.9
m7.6 <- quap(
    alist(
        brain_std ~ dnorm( mu , 0.001 ),
        mu <- a + b[1]*mass_std + b[2]*mass_std^2 +
                  b[3]*mass_std^3 + b[4]*mass_std^4 +
                  b[5]*mass_std^5 + b[6]*mass_std^6,
        a ~ dnorm( 0.5 , 1 ),
        b ~ dnorm( 0 , 10 )
    ), data=d , start=list(b=rep(0,6)) )

## R code 7.19
data(cars)
m <- quap(
    alist(
        dist ~ dnorm(mu,sigma),
        mu <- a + b*speed,
        a ~ dnorm(0,100),
        b ~ dnorm(0,10),
        sigma ~ dexp(1)
    ) , data=cars )

m5.1 <- quap(
    alist(
        D ~ dnorm( mu , sigma ) ,
        mu <- a + bA * A ,
        a ~ dnorm( 0 , 0.2 ) ,
        bA ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) , data = d )

m5.2 <- quap(
    alist(
        D ~ dnorm( mu , sigma ) ,
        mu <- a + bM * M ,
        a ~ dnorm( 0 , 0.2 ) ,
        bM ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) , data = d )

m5.3 <- quap(
    alist(
        D ~ dnorm( mu , sigma ) ,
        mu <- a + bM*M + bA*A ,
        a ~ dnorm( 0 , 0.2 ) ,
        bM ~ dnorm( 0 , 0.5 ) ,
        bA ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) , data = d )

m5.3t <- quap(
    alist(
        D ~ dstudent( 2 , mu , sigma ) ,
        mu <- a + bM*M + bA*A ,
        a ~ dnorm( 0 , 0.2 ) ,
        bM ~ dnorm( 0 , 0.5 ) ,
        bA ~ dnorm( 0 , 0.5 ) ,
        sigma ~ dexp( 1 )
    ) , data = d )


text:
export iauUtcut1
"""
Time scale transformation:  Coordinated Universal Time, UTC, to
Universal Time, UT1.

This function is part of the International Astronomical Union's
SOFA (Standards of Fundamental Astronomy) software collection.

Status:  canonical.

Given:
   utc1,utc2  double   UTC as a 2-part quasi Julian Date (Notes 1-4)
   dut1       double   Delta UT1 = UT1-UTC in seconds (Note 5)

Returned:
   ut11,ut12  double   UT1 as a 2-part Julian Date (Note 6)

Returned (function value):
               int      status: +1 = dubious year (Note 3)
                                 0 = OK
                              -1 = unacceptable date

Notes:

   1. utc1+utc2 is quasi Julian Date (see Note 2), apportioned in any
   convenient way between the two arguments, for example where utc1
   is the Julian Day Number and utc2 is the fraction of a day.

   2. JD cannot unambiguously represent UTC during a leap second unless
   special measures are taken.  The convention in the present
   function is that the JD day represents UTC days whether the
   length is 86399, 86400 or 86401 SI seconds.

   3. The warning status "dubious year" flags UTCs that predate the
   introduction of the time scale or that are too far in the future
   to be trusted.  See iauDat for further details.

   4. The function iauDtf2d converts from calendar date and time of
   day into 2-part Julian Date, and in the case of UTC implements
   the leap-second-ambiguity convention described above.

   5. Delta UT1 can be obtained from tabulations provided by the
   International Earth Rotation and Reference Systems Service.
   It is the caller's responsibility to supply a dut1 argument
   containing the UT1-UTC value that matches the given UTC.

   6. The returned ut11,ut12 are such that their sum is the UT1 Julian
   Date.

References:

   McCarthy, D. D., Petit, G. (eds.), IERS Conventions (2003),
   IERS Technical Note No. 32, BKG (2004)

   Explanatory Supplement to the Astronomical Almanac,
   P. Kenneth Seidelmann (ed), University Science Books (1992)

Called:
   iauJd2cal    JD to Gregorian calendar
   iauDat       delta(AT) = TAI-UTC
   iauUtctai    UTC to TAI
   iauTaiut1    TAI to UT1

This revision:  2013 August 12

SOFA release 2018-01-30

Copyright (C) 2018 IAU SOFA Board.  See notes at end.
"""

# int iauUtcut1(double utc1, double utc2, double dut1,
#               double *ut11, double *ut12)

function iauUtcut1(utc1::Real, utc2::Real, dut1::Real)
   ref_ut11 = Ref{Float64}(0.0)
   ref_ut12 = Ref{Float64}(0.0)

   status = ccall((:iauUtcut1, libsofa_c), Cint,
                   (Cdouble, Cdouble, Cdouble, Ref{Cdouble}, Ref{Cdouble}),
                   convert(Float64, utc1), convert(Float64, utc2), convert(Float64, dut1),
                   ref_ut11, ref_ut12)

   return status, ref_ut11[], ref_ut12[]
end

ext: jl


text:
using CPIMC, CPIMC.PlaneWaves, CPIMC.UniformElectronGas, CPIMC.DefaultUpdates, DataStructures

import CPIMC: update!, orbs, adjacent_kinks_affecting_orbs, kinks_affecting_orbs, prev, next, next_affecting, prev_affecting, τ_prev_affecting, τ_next_affecting, τ_borders, isunaffected, isunaffected_in_interval, time_ordered_orbs, occupations_at, longest_type_1_chain_length, right_type_1_count, kinks_from_periodic_interval, times_from_periodic_interval, Δ, Woffdiag_element, ΔWoffdiag_element, ΔWdiag_element, ΔW_diag, add_orbs, add_orbs!, drop_orbs, drop_orbs!, drop_kinks, drop_kinks!, add_kinks, add_kinks!

S = sphere_with_same_spin(PlaneWave((0,0,0)),dk=1)
a = PlaneWave((-2,0,0))
b = PlaneWave((3,0,0))
c = PlaneWave((0,0,0))
d = PlaneWave((1,0,0))
e = PlaneWave((5,9,9))
f = PlaneWave((1,1,1))

g = PlaneWave(a.vec + e.vec - f.vec, Up)
h = PlaneWave(c.vec + d.vec - g.vec, Up)

sd = Kinks( ImgTime(0.2) => T4(a,b,c,d),
       ImgTime(0.5) => T4(c,d,a,b),
       ImgTime(0.6) => T4(b,a,d,c),
       ImgTime(0.8) => T4(d,c,b,a) )

conf = Configuration(S,sd)

@testset "apply_step" begin
    @test apply_step(conf, Step()) == conf
    @test apply_step(conf, Step((c,d),nothing,(e,f),nothing)).occupations == union(setdiff(S,(c,d)),(e,f))
end

@testset "apply_step!" begin
    c2 = Configuration(S,sd)
    apply_step!(c2, Step())
    @test c2 == conf
    apply_step!(c2, Step((c,d),nothing,(e,f),nothing))
    @test c2.occupations == union(setdiff(S,(c,d)),(e,f))
end

@testset "apply_step" for _ in (1:200)
    c1 = Configuration(sphere(PlaneWave((0,0,0),Up),dk=1))
    c2 = Configuration(sphere(PlaneWave((0,0,0),Up),dk=1))
    m = UEG()
    ens = CEnsemble(2,2,7)
    step_list = Array{Step,1}()
    for _ in (1:5)
        dv, Δ = add_type_B!(m, ens, c1)

        if dv == 0
            apply_back_step!(c1, Δ)
        else
            push!(step_list, Δ)
        end
        dv, Δ = add_type_C!(m, ens, c1)

        if dv == 0
            apply_back_step!(c1, Δ)
        else
            push!(step_list, Δ)
        end
        dv, Δ = add_type_D!(m, ens, c1)

        if dv == 0
            apply_back_step!(c1, Δ)
        else
            push!(step_list, Δ)
        end
        dv, Δ = add_type_E!(m, ens, c1)

        if dv == 0
            apply_back_step!(c1, Δ)
        else
            push!(step_list, Δ)
        end
    end
    for i in (1:3)
        dv, Δ = remove_type_B!(m, ens, c1)

        if dv == 0
            apply_back_step!(c1, Δ)
        else
            push!(step_list, Δ)
        end
        dv, Δ = remove_type_C!(m, ens, c1)

        if dv == 0
            apply_back_step!(c1, Δ)
        else
            push!(step_list, Δ)
        end
        dv, Δ = remove_type_D!(m, ens, c1)

        if dv == 0
            apply_back_step!(c1, Δ)
        else
            push!(step_list, Δ)
        end
        dv, Δ = remove_type_E!(m, ens, c1)

        if dv == 0
            apply_back_step!(c1, Δ)
        else
            push!(step_list, Δ)
        end
    end
    @test apply_step(c2,step_list).kinks == c1.kinks
    @test apply_step(c2,step_list).occupations == c1.occupations
end

@testset "update!" begin
    ens = CEnsemble(2.0, 5.680898543560106, 7)# θ: 0.125, λ: 0.09945178864947428
    c1 = Configuration(sphere_with_same_spin(PlaneWave((0,0,0)),dk=1))
    for _ in 1:100
        old_occs = copy(c1.occupations)
        old_kinks = copy(c1.kinks)
        coin = last(update!(UEG(), ens, c1, [move_particle!, add_type_B!, remove_type_B!, add_type_C!, remove_type_C!, add_type_D!, remove_type_D!, add_type_E!, remove_type_E!, shuffle_indices!]))
        if coin == :accept
            @test ((c1.occupations != old_occs) || (c1.kinks != old_kinks))
        elseif coin == :reject
            @test c1.occupations == old_occs
            @test c1.kinks == old_kinks
        elseif coin == :trivial
            @test c1.occupations == old_occs
            @test c1.kinks == old_kinks
        else
            @test(false)
        end
    end
end


text:
using RCall, JLD, DataFramesMeta

R"
library(gridExtra)
library(cowplot)
library(tidyverse)
library(latex2exp)
library(viridis)
"

DF = load("/Users/dariabystrova/Documents/GitHub/approximatingBNPpriors/saves_for_figures/figure_6_4_2.jld")
DF_NGG_FK = DF["DF_NGG_FK_slow"]

R"
range01 <- function(x){(x-min(x))/(max(x)-min(x))}"
R"
df_r= $DF_NGG_FK
df_r$beta_scaled= range01(df_r$beta_log)
df_r$color = rgb(df_r$beta_scaled, df_r$sigma,0.5,1)
df_r$color_sigma = rgb(0, df_r$sigma,0.5,1)
df_r$color_beta = rgb(df_r$beta_scaled, 0,0.5,1)
"
R"
p <- df_r %>% ggplot(aes(x=Exp, y = std, group=sigma))  + geom_line( alpha = 0.8,linetype = 'longdash') + geom_point(aes(colour=color), size=4)
p_ngg_fk <- p  + geom_line(data =df_r, aes(x=Exp, y = std, group=beta), alpha = 0.9,linetype = 'dotted') +
         xlim(0, 100)+ ylim(0,15)+labs(y='Std', x='Expectation')+scale_colour_identity()+theme_classic()+ theme(plot.title = element_text(hjust = 0.5,size = 15), axis.text.x = element_text(size=15), axis.text.y = element_text(size=15),legend.position='none')
"

R"p_ngg_fk
pdf('/Users/dariabystrova/Documents/GitHub/approximatingBNPpriors/figures/figure_6_4_2.pdf', width= 4, height = 4)
plot(p_ngg_fk)
dev.off()"

max_stars_repo_name: konkam/approximatingBNPpriors

max_issues_repo_name: konkam/approximatingBNPpriors

max_forks_repo_name: konkam/approximatingBNPpriors


text:
# ------------ Choice sites ------------ #

@inline function (ctx::ParameterBackpropagateContext)(call::typeof(rand),
                                                      addr::T,
                                                      d::Distribution{K}) where {T <: Address, K}
    s = get_top(ctx.tr, addr).val
    ctx.weight += logpdf(d, s)
    return s
end

@inline function (ctx::ChoiceBackpropagateContext)(call::typeof(rand),
                                                   addr::T,
                                                   d::Distribution{K}) where {T <: Address, K}
    s = get_top(ctx.tr, addr).val
    ctx.weight += logpdf(d, s)
    return s
end

# ------------ Learnable ------------ #

@inline function (ctx::ParameterBackpropagateContext)(fn::typeof(learnable), addr::Address)
    return read_parameter(ctx, addr)
end

@inline function (ctx::ChoiceBackpropagateContext)(fn::typeof(learnable), addr::Address)
    return read_parameter(ctx, addr)
end

# ------------ Call sites ------------ #

@inline function (ctx::ParameterBackpropagateContext)(c::typeof(rand),
                                                      addr::T,
                                                      call::Function,
                                                      args...) where T <: Address
    cl = get_sub(ctx.tr, addr)
    param_grads = Gradients()
    ps = get_sub(ctx.initial_params, addr)
    ret = simulate_call_pullback(ps, param_grads, cl, args)
    ctx.param_grads.tree[addr] = param_grads
    return ret
end

@inline function (ctx::ChoiceBackpropagateContext)(c::typeof(rand),
                                                   addr::T,
                                                   call::Function,
                                                   args...) where T <: Address
    cl = get_sub(ctx.tr, addr)
    choice_grads = Gradients()
    ps = get_sub(ctx.initial_params, addr)
    ret = simulate_choice_pullback(ps, choice_grads, get_sub(ctx.select, addr), cl, args)
    ctx.choice_grads.tree[addr] = choice_grads
    return ret
end


text:
@testset "963.minimum-area-rectangle-ii.jl" begin
    @test min_area_free_rect([[3,1],[1,1],[0,1],[2,1],[3,3],[3,2],[0,2],[2,3]]) == 2.0
    @test min_area_free_rect([[0,1],[2,1],[1,1],[1,0],[2,0]]) == 1.0
    @test min_area_free_rect([[1,2],[2,1],[1,0],[0,1]]) == 2.0
    @test min_area_free_rect([[0,3],[1,2],[3,1],[1,3],[2,1]]) == 0.0
end


text:
#
# This function computes distance between two selections
# (that is, the minimum distance between their atoms)
#

function distance(sel1,sel2,sides,x,y,z)

  n1 = length(sel1)
  n2 = length(sel2)

  distance = Inf
  for i in 1:n1
    center = [ x[sel1[i]], y[sel1[i]], z[sel1[i]] ]
    wrap!(sides,x,y,z,center=center,sel=sel2)
    for j in 1:n2
      distance = min(distance, (x[sel1[i]] - x[sel2[j]])^2 +
                               (y[sel1[i]] - y[sel2[j]])^2 +
                               (z[sel1[i]] - z[sel2[j]])^2 )
    end
  end
  distance = sqrt(distance)

  return distance
end


text:
import Pluto: ImmutableVector, ImmutableSet, ImmutableDefaultDict, setdiffkeys

@testset "ImmutableCollections" begin

# ╔═╡ bd27d82e-62d6-422c-8fbe-61993dc4c268
@test isempty(ImmutableVector{Int}())

# ╔═╡ d4f2016a-b093-4619-9ccb-3e99bf6fdc9b
@test ImmutableVector{Int32}([1,2,3]).c |> eltype == Int32

# ╔═╡ 055f21c0-3741-4762-ac4e-4c89633afbc4
let
	x = [1,2,3]
	y = ImmutableVector(x)
	push!(x,4)
	@test y == [1,2,3]
end

# ╔═╡ 52310ade-6e06-4ab8-8589-444c161cd93b
let
	x = [1,2,3]
	y = ImmutableVector{Int32}(x)
	push!(x,4)
	@test y == [1,2,3]
end

# ╔═╡ d61600f0-2202-4228-8d35-380f732214e7
ImmutableSet()

# ╔═╡ d3871580-cd22-48c1-a1fd-d13a7f2f2135
ImmutableSet{Int}()

# ╔═╡ 2af00467-8bbf-49d8-bfe3-6f8d6307e900
ImmutableSet{Int64}(Set([1,2]); skip_copy=true)

# ╔═╡ 46836112-7c5c-4ffd-8d1e-93a2c8990b20
@test ImmutableSet{Int64}(Set([1,2]); skip_copy=true) == Set([1,2])

# ╔═╡ fd687b2e-8bec-48b2-810e-38ef00bf567b
let
	x = [1.1,2,3]
	y = ImmutableVector(x; skip_copy=true)
	push!(x,4)
	@test y == [1.1,2,3,4]
end

# ╔═╡ f4dddf0b-cf0a-41d0-880e-6a8fac7c60cb
let
	x = [1.1,2,3]
	y = ImmutableVector{Float64}(x; skip_copy=true)
	push!(x,4)
	@test y == [1.1,2,3,4]
end

# ╔═╡ 25c78371-f12d-44ae-b180-32b88d3aa4f5
@test eltype(ImmutableSet([2,3,4])) == Int

# ╔═╡ 45115ac6-6586-458c-83e6-d661c2ce8db2
let
	x = Set([1,2,3])
	y = ImmutableSet(x)
	push!(x,4)
	@test y == Set([1,2,3])
end

# ╔═╡ 4f26640d-31d2-44c4-bbba-82c18d7497ae
let
	x = Set([1.1,2,3])
	y = ImmutableSet(x; skip_copy=true)
	push!(x,4)
	@test y == Set([1.1,2,3,4])
end

# ╔═╡ eac9c95b-a2b6-4f1f-8cce-a2ad4c0972c5
@test union(ImmutableSet([1,2]),[2,3]) == ImmutableSet([1,2,3])

# ╔═╡ bff65a2c-8654-4403-8e34-58aac8616729
@test filter(x -> true, ImmutableVector([1,2,3])) == [1,2,3]

# ╔═╡ ce3cdb24-e851-4cc3-9955-b34fe358b41a
@test ImmutableVector([1,2,3])[2:end] isa ImmutableVector

# ╔═╡ c61196d6-f529-4883-b334-ed1b0f653acf

# ╔═╡ 5c2b3440-7231-42df-b4e5-619001d225a8
ImmutableSet([123,234])

@test setdiffkeys(Dict(1=>2,3=>4),[3]) == Dict(1=>2)

let
	d = setdiffkeys(ImmutableDefaultDict(() -> 7, Dict(1=>2,3=>4)),[3])
	@test d[1] == 2 && d[3] == 7
end

@test setdiff(ImmutableSet([1,2]), [2]) isa ImmutableSet

end

max_stars_repo_path: test/data structures.jl

max_issues_repo_path: test/data structures.jl

max_forks_repo_path: test/data structures.jl


text:
@testset "679.24-game.jl" begin
    @test least_interval(['A','A','A','B','B','B'], 2) == 8
    @test least_interval(['A','A','A','B','B','B'], 0) == 6
    @test least_interval(['A','A','A','B','B','B'], 1) == 6
    @test least_interval(['A','A','A','A','A','A','B','C','D','E','F','G'], 2) == 16
end


text:
function pcd!(
    white_rbm::AffineRBM,
    data::AbstractArray;
    batchsize::Int = 1,
    epochs::Int = 1,
    optim = Flux.ADAM(),
    history::MVHistory = MVHistory(),
    wts = nothing,
    steps::Int = 1,
    vm::AbstractArray = initial_fantasy_v(white_rbm, data, batchsize),
    damping::Real = 1//100, ϵv::Real = 0, ϵh::Real = 0,
    transform_v::AbstractTransform = Whiten(),
    transform_h::AbstractTransform = Stdize()
)
    @assert size(data) == (size(visible(white_rbm))..., size(data)[end])
    @assert isnothing(wts) || _nobs(data) == _nobs(wts)

    @assert 0 ≤ damping ≤ 1

    whiten_visible_from_data!(white_rbm, data, transform_v; wts, ϵ = ϵv)
    stats = RBMs.suffstats(visible(white_rbm), data; wts)

    for epoch in 1:epochs
        batches = RBMs.minibatches(data, wts; batchsize)
        Δt = @elapsed for (vd, wd) in batches
            # update fantasy chains
            vm .= RBMs.sample_v_from_v(white_rbm, vm; steps)
            # update hidden affine transform
            if !(transform_h isa Identity)
                inputs = RBMs.inputs_v_to_h(white_rbm, vd)
                whiten_hidden_from_inputs!(white_rbm, inputs, transform_h; damping, wts=wd, ϵ=ϵh)
            end
            # compute contrastive divergence gradient
            ∂ = RBMs.∂contrastive_divergence(white_rbm, vd, vm; wd, stats)
            # compute parameter step according to optimization algorithm
            Δ = RBMs.update!(∂, white_rbm, optim)
            # update parameters using gradient
            RBMs.update!(white_rbm, Δ)
            # store gradient and update step norms
            push!(history, :∂, RBMs.gradnorms(∂))
            push!(history, :Δ, RBMs.gradnorms(Δ))
        end

        push!(history, :epoch, epoch)
        push!(history, :Δt, Δt)
        @debug "epoch $epoch/$epochs ($(round(Δt, digits=2))s)"
    end

    return white_rbm, history
end

function RBMs.∂contrastive_divergence(
    rbm::WhiteRBM, vd::AbstractArray, vm::AbstractArray;
    wd = nothing, wm = nothing,
    stats = RBMs.suffstats(visible(rbm), vd; wts = wd),
)
    ∂d = RBMs.∂free_energy(rbm, vd; wts = wd, stats)
    ∂m = RBMs.∂free_energy(rbm, vm; wts = wm)
    ∂ = RBMs.subtract_gradients(∂d, ∂m)
    return ∂
end

function RBMs.update!(white_rbm::WhiteRBM, ∂::NamedTuple)
    RBMs.update!(RBMs.RBM(white_rbm), ∂)
    return white_rbm
end

RBMs.update!(∂::NamedTuple, rbm::WhiteRBM, optim) = RBMs.update!(∂, RBM(rbm), optim)

function initial_fantasy_v(rbm, data::AbstractArray, batchsize::Int)
    inputs = falses(size(data)[1:(end - 1)]..., batchsize)
    vm = RBMs.sample_from_inputs(visible(rbm), inputs)
    return oftype(data, vm)
end


text:
using Plots

function kwt_alfa(degalfa_p)
    eff = 0.15
    area = 30

    teta_p = deg2rad(170)
    alfa_p = deg2rad(degalfa_p)

    lat = deg2rad(59.3)

    delta_s(n) = deg2rad(-23.44) * cos((2pi / 365) * n)

    omega(t) = deg2rad((15 * t) - 180)

    alfa_s(n, t) = asin(sin(lat) * sin(delta_s(n)) + cos(lat) * cos(delta_s(n)) * cos(omega(t)))

    teta_s(n, t) = pi + (omega(t) / abs(omega(t))) * acos((sin(lat) * sin(alfa_s(n, t)) - sin(delta_s(n))) / (cos(lat) * cos(alfa_s(n, t))))

    I_0 = 1360
    I(n, t) = 1.1 * I_0 * 0.7^((1 / (abs(sin(alfa_s(n, t)))+sin(alfa_s(n, t)))/2)^0.678)  # underlig abs medelvärde för att få 0 om negativt värde

    I_p(n, t) = I(n, t) * (cos(teta_p - teta_s(n, t)) * cos(alfa_p - alfa_s(n, t)) + sin(alfa_s(n, t)) * sin(alfa_p) * (1 - cos(teta_p - teta_s(n, t))))

    weather_loss = 1-(1821/(365*12))

    P(x) = (*(((abs(I_p(floor(x/24), x))+I_p(floor(x/24), x))/2), Weather_loss, eff, area))

    s_date = 152
    e_date = 181
    step = 0.1

    x = 24*s_date:step:24*(e_date)+24
    x = filter(x->x%12!=0,x)

    kwh = 0

    for i in x
        kwh += i*P(i)
    end

    return kwh
end

plot(kwt_alfa, 0, 180)
# plot kwh as a function of alfa_p. We can later change this to output the biggest alfa_p with relative ease.


text:
vis = Visualizer()
open(vis)

# Create Roadway
roadway_opts = MergingRoadwayOptions(merging_point=1.4, lane_width=0.28)
roadway = build_roadway(roadway_opts)

# Create players
T = Float64
p = 3
model = BicycleGame(p=p)
n = model.n
m = model.m

# Define the horizon of the problem
N = 10 # N time steps
dt = 0.10 # each step lasts 0.1 second
probsize = ProblemSize(N,model) # Structure holding the relevant sizes of the problem

# Define the objective of each player
# We use a LQR cost
Q = [Diagonal(3*SVector{model.ni[i],T}([1., 2.0, 1., 1.])) for i=1:p] # Quadratic state cost
R = [Diagonal(0.3*ones(SVector{model.mi[i],T})) for i=1:p] # Quadratic control cost
# Desrired state
xf = [SVector{model.ni[1],T}([2.0, 0.05, 0.0, 0]),
      SVector{model.ni[2],T}([2.0, 0.05, 0.0, 0]),
      SVector{model.ni[2],T}([2.0, 0.05, 0.0, 0]),
      ]
# Desired control
uf = [zeros(SVector{model.mi[i],T}) for i=1:p]
# Objectives of the game
game_obj = Algames.GameObjective(Q,R,xf,uf,N,model)
# radius = 0.22*ones(p)
# μ = 40.0*ones(p)
# add_collision_cost!(game_obj, radius, μ)

# Define the constraints that each player must respect
game_con = Algames.GameConstraintValues(probsize)
radius = 0.08
add_collision_avoidance!(game_con, radius)

# Define the initial state of the system
x0 = SVector{model.n,T}([
    0.2,     0.2,    0.2,
   -0.45,   -0.13,   0.13,
    0.45,    0.60,   0.60,
    0.15,    0.0,    0.0,
    ])

# Define the Options of the solver
opts = Options()
opts.ls_iter = 15
opts.outer_iter = 20
opts.inner_iter = 20
opts.ρ_0 = 1e0
opts.reg_0 = 1e-5
opts.α_dual = 1.0
opts.λ_max = 1.0*1e7
opts.ϵ_dyn = 1e-6
opts.ϵ_sta = 1e-6
opts.ϵ_con = 1e-6
opts.ϵ_opt = 1e-6
opts.regularize = true
# Define the game problem
prob = Algames.GameProblem(N,dt,x0,model,opts,game_obj,game_con)

# Solve the problem
newton_solve!(prob)

plot!(prob.model, prob.pdtraj.pr)
plot!(prob.stats)

players = Vector{Player{T}}(undef, p)
player_opts = HighwayPlayerOptions{model.ni[1],model.mi[1],T}()
players[1] = Player(model, roadway.lane[3])
players[2] = Player(model, roadway.lane[3])
players[3] = Player(model, roadway.lane[3])

# Create Scenario
sce = Scenario(model, roadway, players)

# Initialize visualizers
# vis = Visualizer()
# open(vis)

# Visualize trajectories
set_scenario!(vis, sce)
set_env!(vis, VehicleState(1.3, 0.0, 0.0, 0.0))
set_camera_birdseye!(vis, height=5.0)

build_waypoint!(vis, sce.player, N, key=0)
set_waypoint_traj!(vis, model, sce, prob.pdtraj.pr, key=0)
set_traj!(vis, model, sce, prob.pdtraj.pr)

get_num_active_constraint(prob)

include("explorer.jl")
prob_copy = deepcopy(prob)
ascore = ActiveSetCore(probsize)
active_vertical_mask!(ascore, prob_copy.game_con)
vmask = deepcopy(ascore.vmask)

# Tests
Sh = probsize.S+p*(p-1)*(N-1)
Sv = probsize.S+Int(p*(p-1)*(N-1)/2)

xp = get_primal_dual(ascore, prob_copy)
α = 1e-3
x = get_primal_dual(ascore, prob_copy)
update_nullspace!(ascore, prob_copy, prob_copy.pdtraj)
ns = length(ascore.null.vec)
v = deepcopy(ascore.null.vec[1])
β = 1.0
# active_vertical_mask!(ascore, prob_copy.game_con)
# vmask = deepcopy(ascore.vmask)
λ = 1e-1*rand(Sv)
f(xp, α, x, v, β)
mean(abs.(c(x,ascore,prob_copy)))
mean(abs.(c(x,ascore,prob_copy)[vmask]))
mean(abs.(c(x,ascore,prob_copy)[setdiff(1:Sv,vmask)]))
mean(max.(0, c(x,ascore,prob_copy)[setdiff(1:Sv,vmask)]))

λ[vmask]'*c(x, ascore, prob_copy)[vmask]
L(xp, α, λ, x, v, β, ascore, prob_copy, vmask)

# g = grad(xp, α, λ, x, v, β, ascore, prob_copy, vmask)
# fdg = finite_difference_grad(xp, α, λ, x, v, β, ascore, prob_copy, vmask)
# nn = Int(length(g)/2)
# mean(abs.(fdg[1:nn] - g[1:nn]))
# mean(abs.(fdg[nn+1:end] - g[nn+1:end]))

jac(xp, α, λ, x, v, β, ascore, prob_copy, vmask)

α = 1e-4
maximum(abs.(α*v))

plot!(model, prob_copy.pdtraj.pr)
prob_copy2 = deepcopy(prob)
λ = 1e-10*rand(Sv)
xp, λ = newton_solver(α, λ, x, v, β, ascore, prob_copy2, vmask)
c(xp, ascore, prob_copy2)
f(xp, α, x, v, β)
λ[vmask]'*c(xp, ascore, prob_copy2)[vmask]
mean(abs.(λ[vmask]))

plot!(model, prob_copy.pdtraj.pr)
plot!(model, prob_copy2.pdtraj.pr)
set_primal_dual!(ascore, prob_copy2, x+α*v)
plot!(model, prob_copy2.pdtraj.pr)

λ[vmask][end-28+1:end]'*c(xp, ascore, prob_copy2)[vmask][end-28+1:end]
λ[vmask][end-30:end]
maximum(abs.(λ[vmask][end-30:end]))
maximum(abs.(c(xp, ascore, prob_copy2)[vmask][end-30:end]))

mean(abs.(c(xp, ascore, prob_copy2)[vmask]))

set_primal_dual!(ascore, prob_copy2, x+1e-1*v)
plot!(model, prob.pdtraj.pr)
plot!(model, prob_copy2.pdtraj.pr)

mean(abs.(xp))
mean(abs.(x))
mean(abs.(xp[vmask]))
mean(abs.(x[vmask]))

mean(abs.(xp[vmask]-x[vmask]))

prob_copy2 = deepcopy(prob)
set_primal_dual!(ascore, prob_copy2, xp)
residual!(prob_copy2)
prob_copy2.core.res

a = 10
a = 10
a = 10
a = 10
a = 10
a = 10

#
# include("explorer.jl")
# prob_copy = deepcopy(prob)
# # plot!(model, prob_copy.pdtraj.pr)
# ascore = ActiveSetCore(probsize)
# active_vertical_mask!(ascore, prob_copy.game_con)
# vmask = deepcopy(ascore.vmask)
# # set_primal_dual!(ascore, prob_copy, x)
#
# # Tests
# Sh = probsize.S+p*(p-1)*(N-1)
# Sv = probsize.S+Int(p*(p-1)*(N-1)/2)
# # xr = rand(Sh)
# # set_primal_dual!(ascore, prob_copy, xr)
# # xr_ = get_primal_dual(ascore, prob_copy)
# # @test xr == xr_
#
# xp = get_primal_dual(ascore, prob_copy)
# α = 0.2
# x = get_primal_dual(ascore, prob_copy)
# update_nullspace!(ascore, prob_copy, prob_copy.pdtraj)
# ns = length(ascore.null.vec)
# v = deepcopy(ascore.null.vec[1])
# β = 10.0
# # active_vertical_mask!(ascore, prob_copy.game_con)
# # vmask = deepcopy(ascore.vmask)
# λ = 1e-1*rand(Sv)
# f(xp, α, x, v, β)
# λ'*c(x, ascore, prob_copy)
# L(xp, α, λ, x, v, β, ascore, prob_copy, vmask)
# grad(xp, α, λ, x, v, β, ascore, prob_copy, vmask)
# jac(xp, α, λ, x, v, β, ascore, prob_copy, vmask)
# newton_solver(λ, x, v, β, ascore, prob_copy, vmask)

prob_copy = deepcopy(prob)
ascore = ActiveSetCore(probsize)
get_primal_dual(ascore, probsize)

prob_copy = deepcopy(prob)
ascore, subspace = subspace_dimension(prob_copy, α=1e-3) # Bicycle
subspace
vals, ref_pdtraj, eig_pdtraj_1, eig_pdtraj_2, = pca(prob, subspace, β=1.2*10^5)
display_eigvals(vals)

β = 0.8*1e5
display_arrow2(prob_copy, sce, ref_pdtraj, vals[1], eig_pdtraj_1, color=:yellow, key=1, β=β)
display_arrow2(prob_copy, sce, ref_pdtraj, vals[2], eig_pdtraj_2, color=:cornflowerblue, key=2, β=β)

a = 10

function display_arrow2(prob::Algames.GameProblem, sce::Scenario, ref_pdtraj::PrimalDualTraj, val::T,
	eig_pdtraj::PrimalDualTraj; vis_opts::PlayerVisualizationOptions=PlayerVisualizationOptions{T}(),
	color=:orange, key::Int=0, β::T=1.2*10^5) where {T}

	probsize = prob.probsize
	N = probsize.N
	α = vis_opts.α
	for k = 2:N
		vs_ref = standardize(prob.model, Algames.state(ref_pdtraj.pr[k]))
		vs_eig = standardize(prob.model, Algames.state(eig_pdtraj.pr[k]))
		for i = 1:p
			player = sce.player[i]
			arr_vis = ArrowVisualizer(vis["env/arrow/arrow_traj$key/player$(player.id)/$k"])
			col = vis_opts.colors[color]
			mat = MeshPhongMaterial(color=RGBA(col, α))

			setobject!(arr_vis, mat)
			settransform!(arr_vis,
				Point(0.0, 0.0, 0.0),
				Vec(β*val*norm([vs_eig[i].x, vs_eig[i].y]), 0.0, 0.0),
				shaft_radius=0.002,
				max_head_radius=0.01)

			t = Translation(vs_ref[i].x, vs_ref[i].y, 0.01)
			r1 = LinearMap(AngleAxis(pi/2, 0.0, 1.0, 0.0))
			r2 = LinearMap(AngleAxis(-atan(vs_eig[i].y,vs_eig[i].x), 1.0, 0.0, 0.0))
			tr = compose(t, compose(r1, r2))
			settransform!(vis["env/arrow/arrow_traj$key/player$(player.id)/$k"], tr)
		end
	end
	return nothing
end

function display_ghost(vis::Visualizer, prob::Algames.GameProblem, sce::Scenario)
    N = prob.probsize.N
    for k ∈ vcat(Vector((1:10:N)),N)
        @show k
        vis_opts = PlayerVisualizationOptions{T}(α=0.3+0.2*k/10)
        set_player!(vis, sce.player, vis_opts=vis_opts, key=k)
        set_state!(vis, prob.model, sce, Algames.state(prob.pdtraj.pr[k]), key=k)
    end
    return nothing
end

function get_num_active_constraint(prob::Algames.GameProblem)
    probsize = prob.probsize
    p = probsize.p
	update_active_set!(prob.game_con, prob.pdtraj.pr)
    na = 0
    for i = 1:p
        for j = i+1:p
            conval, v = get_collision_conval(prob.game_con, i, j)
            if v
				for (l,k) in enumerate(conval.inds)
					na += conval.active[l][1]
				end
			end
        end
    end
    return na
end

function update_λ!(ascore::ActiveSetCore, prob::Algames.GameProblem, Δλcol::AbstractVector)
    probsize = prob.probsize
    N = probsize.N
    p = probsize.p
    px = probsize.px

    stamp = CStamp()
    for i = 1:p
        for j ∈ setdiff(1:p,i)
            conval, v = get_collision_conval(prob.game_con, i, j)
            if v
                for (l,k) in enumerate(conval.inds)
                    stampify!(stamp, :h, :col, i, j, k)
                    if conval.active[l][1] == 1
                        conval.λ[l] += Δλcol[horizontal_idx(ascore, stamp) .- probsize.S]
                        conval.λ[l] = max.(0, conval.λ[l])
                    end
                end
            end
        end
    end
    return nothing
end

function subspace_dimension(prob::Algames.GameProblem; α::T=2e-3) where {T}
    prob_ = deepcopy(prob)
    probsize = prob.probsize
    ascore = ActiveSetCore(probsize)
    update_nullspace!(ascore, prob, prob.pdtraj)
    ns = length(ascore.null.vec)
    val = 0
    subspace = Vector{PrimalDualTraj}([deepcopy(prob.pdtraj)])
    plt = plot(legend=false)
    for l = 1:ns
        @show l
        prob_ = deepcopy(prob)
        prob_.opts.shift = 0
        prob_.opts.α_dual = 0.0
        prob_.opts.αx_dual = 1e-3*ones(p)
        prob_.opts.dual_reset = false
        prob_.opts.ρ_0 = prob_.pen.ρ[1]
        prob_.opts.ρ_max = prob_.pen.ρ[1]
        prob_.opts.inner_print = false
        prob_.opts.outer_print = false
        prob_.opts.inner_iter = 20
        prob_.opts.outer_iter = 40
        prob_.opts.reg_0 = 1e-6
        ϵ = 1e-6
        prob_.opts.ϵ_opt = ϵ
        prob_.opts.ϵ_dyn = ϵ
        prob_.opts.ϵ_con = ϵ
        prob_.opts.ϵ_sta = ϵ
        set_traj!(prob_.core, prob_.Δpdtraj, α*ascore.null.Δtraj[l])
        update_traj!(prob.pdtraj, prob_.pdtraj, 1.0, prob_.Δpdtraj)
        update_λ!(ascore, prob_, α*ascore.null.Δλ[l])
        newton_solve!(prob_)
        # plot!(prob_.stats)

        opt = mean(abs.(prob_.stats.opt_vio[end].vio))
        sta = mean(abs.(prob_.stats.sta_vio[end].vio))
        dyn = mean(abs.(prob_.stats.dyn_vio[end].vio))
        con = mean(abs.(prob_.stats.con_vio[end].vio))
        opts = prob_.opts
        if opt < opts.ϵ_opt && sta < opts.ϵ_sta && dyn < opts.ϵ_dyn && con < opts.ϵ_con
            val += 1
			@show val
            # plot!(prob_.model, prob_.pdtraj.pr, plt=plt)
            push!(subspace, deepcopy(prob_.pdtraj))
        end
    end
    @show ns
    @show val
    return ascore, subspace
end

function pca(prob::Algames.GameProblem, subspace::Vector{PrimalDualTraj}; β::T=2.0) where {T}
    probsize = prob.probsize
    N = probsize.N
    n = probsize.n
    m = probsize.m
    p = probsize.p
    px = probsize.px
    S = probsize.S
    M = length(subspace)-1

    V = []
    v = zeros(S)
    v_ref = zeros(S)
    get_traj!(prob.core, subspace[1], v_ref)
    for k = 1:M
        get_traj!(prob.core, subspace[k+1], v)
        v = v - v_ref
        push!(V, deepcopy(v))
    end

    pr_mask = primals_mask(prob.core, probsize)
	R = length(pr_mask)
	W = [v[pr_mask] for v in V]
    m = sum(W) ./ M
    Wc = [w .- m for w in W]
    cov = sum([wc*wc' for wc in Wc])/(M-1)
    vecs = eigvecs(cov)
    vals = eigvals(cov)
    eigs = [(vecs[:,i], vals[i]) for i=1:R]
    sort!(eigs, by = x -> x[2], rev=true)
	vecs = [e[1] for e in eigs[1:M]]
    vals = [e[2] for e in eigs[1:M]]

    @show round.(log.(10, abs.(vals)), digits=3)

    plt = plot(legend=false)
    pdtraj_ref = subspace[1]
    plot!(prob.model, pdtraj_ref.pr, plt=plt)

    pdtraj_1 = deepcopy(pdtraj_ref)
    v1 = zeros(S)
    v1[pr_mask] = vecs[1]
    set_traj!(prob.core, pdtraj_1, v1)

    pdtraj_2 = deepcopy(pdtraj_ref)
    v2 = zeros(S)
    v2[pr_mask] = vecs[2]
    set_traj!(prob.core, pdtraj_2, v2)

    pdtraj_3 = deepcopy(pdtraj_ref)
    v3 = zeros(S)
    v3[pr_mask] = vecs[3]
    set_traj!(prob.core, pdtraj_3, v3)

    x = [[Algames.state(pdtraj_ref.pr[k])[px[i][1]] for k=2:N] for i=1:p]
    y = [[Algames.state(pdtraj_ref.pr[k])[px[i][2]] for k=2:N] for i=1:p]
    u1 = [[vals[1]*Algames.state(pdtraj_1.pr[k])[px[i][1]]  for k=2:N] for i=1:p]
    v1 = [[vals[1]*Algames.state(pdtraj_1.pr[k])[px[i][2]]  for k=2:N] for i=1:p]
    u2 = [[vals[2]*Algames.state(pdtraj_2.pr[k])[px[i][1]]  for k=2:N] for i=1:p]
    v2 = [[vals[2]*Algames.state(pdtraj_2.pr[k])[px[i][2]]  for k=2:N] for i=1:p]
    u3 = [[vals[3]*Algames.state(pdtraj_3.pr[k])[px[i][1]]  for k=2:N] for i=1:p]
    v3 = [[vals[3]*Algames.state(pdtraj_3.pr[k])[px[i][2]]  for k=2:N] for i=1:p]

    for i = 1:p
        quiver!(x[i],y[i],quiver=(β*u1[i],β*v1[i]), linewidth=1.5, color=:cyan)
        quiver!(x[i],y[i],quiver=(β*u2[i],β*v2[i]), linewidth=1.5, color=:red)
        # quiver!(x[i],y[i],quiver=(β*u3[i],β*v3[i]), linewidth=1.5, color=:black)
    end

    display(plt)
    return vals, pdtraj_ref, pdtraj_1, pdtraj_2
end

function primals_mask(core::NewtonCore, probsize::ProblemSize)
    N = probsize.N
    inds = Vector{Int}()
    stamp = HStamp()
    for k = 1:N-1
        # States
        stampify!(stamp, :x, 1, k+1)
        ind = horizontal_idx(core, stamp)
        inds = vcat(inds, Vector(ind))
        # Controls
        for i = 1:p
            stampify!(stamp, :u, i, k)
            ind = horizontal_idx(core, stamp)
            inds = vcat(inds, Vector(ind))
        end
    end
    return inds
end

function display_eigvals(vals)
    plt = plot(legend=false)
    # plot!(sqrt.(vals),
    plot!(vals,
        linetype=:steppost,
        linewidth=3.0,
        title="PCA on the Nash Equilibrium subspace \n around a nominal NE trajectory.",
        xlabel="eigval_index",
        ylabel="eigval_magnitude")
    display(plt)
	for (i,val) in enumerate(vals)
		println("($i,$val)")
	end
    return nothing
end

#
# ################################################################################
# # Toy Example
# ################################################################################
# using Plots
# T = Float64
#
# # Define the dynamics of the system
# p = 3 # Number of players
# # model = DoubleIntegratorGame(p=p) # game with 3 players with unicycle dynamics
# model = UnicycleGame(p=p) # game with 3 players with unicycle dynamics
# model = BicycleGame(p=p) # game with 3 players with unicycle dynamics
# n = model.n
# m = model.m
#
# # Define the horizon of the problem
# N = 20 # N time steps
# dt = 0.1 # each step lasts 0.1 second
# probsize = ProblemSize(N,model) # Structure holding the relevant sizes of the problem
#
# # Define the objective of each player
# # We use a LQR cost
# Q = [Diagonal(1*ones(SVector{model.ni[i],T})) for i=1:p] # Quadratic state cost
# R = [Diagonal(0.1*ones(SVector{model.mi[i],T})) for i=1:p] # Quadratic control cost
# # Desrired state
# xf = [SVector{model.ni[1],T}([2, 0.0,0,0]),
#       SVector{model.ni[2],T}([2, 0.0,0,0]),
#       SVector{model.ni[3],T}([2, 0.0,0,0]),
#       # SVector{model.ni[4],T}([3,+0.8,0,0]),
#       ]
# # Desired control
# uf = [zeros(SVector{model.mi[i],T}) for i=1:p]
#
# # # Objectives of the game
# game_obj = GameObjective(Q,R,xf,uf,N,model)
# # radius = 0.5*ones(p)
# # μ = 4.0*ones(p)
# # μ = 4.0*[1.0, 0.1]
# # μ = 4.0*[0.1, 1.0]
# # add_collision_cost!(game_obj, radius, μ)
#
# # Define the constraints that each player must respect
# game_con = GameConstraintValues(probsize)
# # Add collision avoidance
# radius = 0.19
# add_collision_avoidance!(game_con, radius)
#
# # # Add control bounds
# # u_max =  5*ones(SVector{m,T})
# # u_min = -5*ones(SVector{m,T})
# # add_control_bound!(game_con, u_max, u_min)
#
# # # Add wall constraint
# # walls = [Wall([-1.0,-1.75], [1.0,0.25], [1.,-1.]/sqrt(2))]
# # add_wall_constraint!(game_con, walls)
#
# # # Add circle constraint
# # xc = [1., 2., 3.]
# # yc = [1., 2., 3.]
# # radius = [0.1, 0.2, 0.3]
# # add_circle_constraint!(game_con, xc, yc, radius)
#
# # Define the initial state of the system
# x0 = SVector{model.n,T}([
#     0.0, 0.0, 0.0, #0.0,
#    -0.4, 0.0, 0.4, #0.6,
#     0.0, 0.0, 0.0, #0.0,
#     0.0, 0.0, 0.0, #0.0,
#     ])
#
# # Define the Options of the solver
# opts = Options()
# opts.ls_iter = 15
# opts.outer_iter = 20
# opts.inner_iter = 20
# opts.ρ_0 = 1e0
# opts.reg_0 = 1e-5
# opts.α_dual = 1.0
# opts.λ_max = 1.0*1e7
# opts.ϵ_dyn = 1e-5
# opts.ϵ_sta = 1e-5
# opts.ϵ_con = 1e-5
# opts.ϵ_opt = 1e-5
# opts.regularize = true
# opts.αx_dual = ones(p)

# s = 0.5
# opts.αx_dual = [sqrt(s), 1/sqrt(s)]
# Define the game problem
# prob = GameProblem(N,dt,x0,model,opts,game_obj,game_con)
#
# # Solve the problem
# @time newton_solve!(prob)
# # @profiler newton_solve!(prob)
#
# plot!(prob.model, prob.pdtraj.pr)
# plot!(prob.stats)


text:
module Introducao end

export soma

soma(a,b) = a + b

max_stars_repo_name: raphaelpanta/julia-lang-exemplos

max_issues_repo_name: raphaelpanta/julia-lang-exemplos

max_forks_repo_name: raphaelpanta/julia-lang-exemplos


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

using Markdown
using InteractiveUtils

# ╔═╡ 6c340276-d41e-11ea-02b3-358efeaf544b
println("Julia 👩🦰 sagt: Hallo, Jörg!")

# ╔═╡ 02869d36-d41f-11ea-374c-e7070aed0782
a = 10

# ╔═╡ 0cfea038-d41f-11ea-0f40-db0fa69142eb
b = 5

# ╔═╡ 169080b0-d41f-11ea-18ee-6bee7fba1a92
a + b

# ╔═╡ Cell order:
# ╠═6c340276-d41e-11ea-02b3-358efeaf544b
# ╠═02869d36-d41f-11ea-374c-e7070aed0782
# ╠═0cfea038-d41f-11ea-0f40-db0fa69142eb
# ╠═169080b0-d41f-11ea-18ee-6bee7fba1a92

max_stars_repo_name: kantel/julia

max_issues_repo_name: kantel/julia

max_forks_repo_name: kantel/julia


text:
# Rules for automatic differentiation

## Constructor
### These rules designate QDHT as non-differentiable
function ChainRulesCore.frule(Δargs, ::Type{T}, args...; kwargs...) where {T<:QDHT}
    return T(args...; kwargs...), NoTangent()
end

function ChainRulesCore.rrule(::Type{T}, args...; kwargs...) where {T<:QDHT}
    function QDHT_pullback(ΔQ)
        return (NoTangent(), map(_ -> NoTangent(), args)...)
    end
    return T(args...; kwargs...), QDHT_pullback
end

## rules for fwd/rev transform
ChainRulesCore.frule((_, _, ΔA), ::typeof(*), Q::QDHT, A) = (Q * A, Q * unthunk(ΔA))
ChainRulesCore.frule((_, _, ΔA), ::typeof(\), Q::QDHT, A) = (Q \ A, Q \ unthunk(ΔA))
function ChainRulesCore.frule((_, ΔY, _, ΔA), ::typeof(mul!), Y, Q::QDHT, A)
    return mul!(Y, Q, A), mul!(unthunk(ΔY), Q, unthunk(ΔA))
end
function ChainRulesCore.frule((_, ΔY, _, ΔA), ::typeof(ldiv!), Y, Q::QDHT, A)
    return ldiv!(Y, Q, A), ldiv!(unthunk(ΔY), Q, unthunk(ΔA))
end

function ChainRulesCore.rrule(::typeof(*), Q::QDHT, A)
    Y = Q * A
    function mul_pullback(ΔY)
        ∂Q = NoTangent()
        ∂A = _mul_back(unthunk(ΔY), Q, A, Q.scaleRK)
        return NoTangent(), ∂Q, ∂A
    end
    return Y, mul_pullback
end

function ChainRulesCore.rrule(::typeof(\), Q::QDHT, A)
    Y = Q \ A
    function ldiv_pullback(ΔY)
        ∂Q = NoTangent()
        ∂A = _mul_back(unthunk(ΔY), Q, A, inv(Q.scaleRK))
        return NoTangent(), ∂Q, ∂A
    end
    return Y, ldiv_pullback
end

function _mul_back(ΔY, Q, A, s)
    T = typeof(one(eltype(ΔY)) * one(eltype(Q.T)) * one(eltype(s)))
    ∂A = Array{T}(undef, size(A))
    dot!(∂A, Q.T', ΔY, dim = Q.dim)
    ∂A .*= s
    return ∂A
end

## rules for integrateR/integrateK
function ChainRulesCore.frule((_, ΔA, _), ::typeof(integrateR), A, Q::QDHT; kwargs...)
    return integrateR(A, Q; kwargs...), integrateR(unthunk(ΔA), Q; kwargs...)
end

function ChainRulesCore.frule((_, ΔA, _), ::typeof(integrateK), A, Q::QDHT; kwargs...)
    return integrateK(A, Q; kwargs...), integrateK(unthunk(ΔA), Q; kwargs...)
end

function ChainRulesCore.rrule(::typeof(integrateR), A, Q::QDHT; dim = 1)
    function integrateR_pullback(ΔΩ)
        ∂A = _integrateRK_back(unthunk(ΔΩ), A, Q.scaleR; dim = dim)
        return NoTangent(), ∂A, NoTangent()
    end
    return integrateR(A, Q; dim = dim), integrateR_pullback
end

function ChainRulesCore.rrule(::typeof(integrateK), A, Q::QDHT; dim = 1)
    function integrateK_pullback(ΔΩ)
        ∂A = _integrateRK_back(unthunk(ΔΩ), A, Q.scaleK; dim = dim)
        return NoTangent(), ∂A, NoTangent()
    end
    return integrateK(A, Q; dim = dim), integrateK_pullback
end

_integrateRK_back(ΔΩ, A::AbstractVector, scale; dim = 1) = ΔΩ .* scale
function _integrateRK_back(ΔΩ, A::AbstractMatrix, scale; dim = 1)
    return dim == 1 ? scale .* ΔΩ : ΔΩ * scale'
end
function _integrateRK_back(ΔΩ, A, scale; dim = 1)
    N = size(A, dim)
    sz = ntuple(_ -> 1, ndims(A))
    sz = Base.setindex(sz, N, dim)
    scalearray = reshape(scale, sz)
    ∂A = ΔΩ .* scalearray
    return ∂A
end

max_stars_repo_name: chrisbrahms/Hankel

max_issues_repo_name: chrisbrahms/Hankel

max_forks_repo_name: chrisbrahms/Hankel


text:
#md# This example shows how to generate a dynamic word cloud on series data.
#md# ### Prepare some data
#md# We use the GDP of countries from 2000 to 2020 as an example.
using WorldBankData
using DataFrames
df0 = wdi("NY.GDP.MKTP.CD", "all", 2000, 2020) #GDP from 2000 to 2020
df0.year = round.(Int, df0[!, :year])
df = unstack(df0, :year, :NY_GDP_MKTP_CD)
country_groups = ["XC", "EU", "XE", "XD", "XF", "ZB", "ZT", "XH", "XI", "XG", "ZJ", "XJ", "XL", "XO",
    "XM", "XN", "ZQ", "XQ", "XP", "XU", "XY", "OE", "ZG", "ZF", "XT"]|>Set # https://rpubs.com/federicoganz/749793
countrymask = .!occursin.(r"\d" , df.iso2c) .& .!in.(df.iso2c, Ref(country_groups))
df = df[countrymask, 2:end]
df[!, 2:end] .= sqrt.(df[!, 2:end])
df.country = replace.(df.country, " "=>"\n") #some names are too long
for i in 1:size(df, 1) #interpolation, fill missing
    for j in 3:size(df, 2)
        if ismissing(df[i, j])
            df[i, j] = df[i, j-1]
        end
    end
end
#md# ### Generate the wordcloud
using WordCloud
function gather!(wc, i=1:length(wc), r=1)
    O = reverse(size(getmask(wc))) .÷ 2
    order1 = sortperm(sortperm(getpositions(wc, i, type=getcenter), by=p->sum((p .- O).^2)))
    order2 = sortperm(getweights(wc, i))
    radial = Int.(order2 .> order1)
    gatheritem!.(wc, WordCloud.index(wc, i), r.*radial);
end
function gatheritem!(wc, i, r=1)
    r == 0 && return
    C = getpositions(wc, i, type=getcenter)
    Δ = sign.(reverse(size(getmask(wc))) .÷ 2 .- C)
    setpositions!(wc, i, C .+ r .* Δ, type=setcenter!)
end
#md# These two functions are ussed to forcing large words to the center, but are not essential.
wc = wordcloud(df.country, 1, angles=0)
gif = WordCloud.GIF("series")
@assert length(unique(df[!, 1])) == length(df[!, 1])
initialized = false
for name in names(df)[2:end] #the first column is word list
    words = df[!, 1]
    weights = df[!, name]
    missingmask = ismissing.(weights)
    setweights!.(wc, words[.!missingmask], weights[.!missingmask])
    println("#"^9, name, "#"^9)
    ignore(wc, words[missingmask]) do
        global initialized
        words = getwords(wc)
        pos = getpositions(wc, type=getcenter)
        initwords!(wc)
        if !initialized
            placewords!(wc, style=:gathering)
            initialized = true
        else
            setpositions!(wc, words, pos, type=setcenter!)
            setstate!(wc, :placewords!)
        end
        generate!(wc, callback=ep->(ep%5==0 && gather!(wc, 1:length(wc)÷5, 5)))
        WordCloud.frame(wc, name) |> gif
        # display(wc)
    end
end
WordCloud.Render.generate(gif, framerate=2)
wc
#eval# runexample(:series)
#md# ![](series/animation.gif)


text:
# ------------------------------------------------------------
# Regularization: Algorithm uses some regularization method
# ------------------------------------------------------------
@define_trait Regularization = begin
    RD   # Algorithm uses the regularized decomposition method of Ruszczyński
    TR   # Algorithm uses the trust-region method of Linderoth/Wright
    LV      # Algorithm uses the level set method of Fábián/Szőke
end

@define_traitfn Regularization init_solver!(lshaped::AbstractLShapedSolver) = begin
    function init_solver!(lshaped::AbstractLShapedSolver, !Regularization)
        nothing
    end
end

@define_traitfn Regularization log_regularization!(lshaped::AbstractLShapedSolver) = begin
    function log_regularization!(lshaped::AbstractLShapedSolver, !Regularization)
        nothing
    end
end

@define_traitfn Regularization log_regularization!(lshaped::AbstractLShapedSolver, t::Integer) = begin
    function log_regularization!(lshaped::AbstractLShapedSolver, t::Integer, !Regularization)
        nothing
    end
end

@define_traitfn Regularization take_step!(lshaped::AbstractLShapedSolver) = begin
    function take_step!(lshaped::AbstractLShapedSolver, !Regularization)
        nothing
    end
end

@define_traitfn Regularization decision(lshaped::AbstractLShapedSolver) = begin
    function decision(lshaped::AbstractLShapedSolver, !Regularization)
        return lshaped.x
    end

    function decision(lshaped::AbstractLShapedSolver, Regularization)
        return lshaped.ξ
    end
end

@define_traitfn Regularization solve_problem!(lshaped::AbstractLShapedSolver, solver::LQSolver) = begin
    function solve_problem!(lshaped::AbstractLShapedSolver, solver::LQSolver, !Regularization)
        solver(lshaped.mastervector)
    end

    function solve_problem!(lshaped::AbstractLShapedSolver, solver::LQSolver, Regularization)
        solver(lshaped.mastervector)
    end
end

@define_traitfn Regularization gap(lshaped::AbstractLShapedSolver) = begin
    function gap(lshaped::AbstractLShapedSolver, !Regularization)
        @unpack τ = lshaped.parameters
        @unpack Q,θ = lshaped.solverdata
        return abs(θ-Q)/(abs(Q)+1e-10)
    end

    function gap(lshaped::AbstractLShapedSolver, Regularization)
        @unpack τ = lshaped.parameters
        @unpack Q̃,θ = lshaped.solverdata
        return abs(θ-Q̃)/(abs(Q̃)+1e-10)
    end
end

@define_traitfn Regularization process_cut!(lshaped::AbstractLShapedSolver, cut::HyperPlane) = begin
    function process_cut!(lshaped::AbstractLShapedSolver, cut::HyperPlane, !Regularization)
        nothing
    end

    function process_cut!(lshaped::AbstractLShapedSolver, cut::HyperPlane, Regularization)
        nothing
    end
end

@define_traitfn Regularization project!(lshaped::AbstractLShapedSolver) = begin
    function project!(lshaped::AbstractLShapedSolver, !Regularization)
        nothing
    end

    function project!(lshaped::AbstractLShapedSolver, Regularization)
        nothing
    end
end

# RD
# ------------------------------------------------------------
@implement_traitfn function init_solver!(lshaped::AbstractLShapedSolver, RD)
    if lshaped.parameters.autotune
        if lshaped.parameters.linearize
            σ̅ = max(4.0,0.01*norm(lshaped.x,Inf))
            σ̲ = min(2.0,0.001*norm(lshaped.x,Inf))
            σ = min(3.0,0.005*norm(lshaped.x,Inf))
            @pack! lshaped.parameters = σ,σ̅,σ̲
        else
            σ̅ = max(4.0,norm(lshaped.x))
            σ̲ = min(2.0,0.005*norm(lshaped.x))
            σ = σ̲
            @pack! lshaped.parameters = σ,σ̅,σ̲
        end
    end
    lshaped.solverdata.σ = lshaped.parameters.σ
    push!(lshaped.σ_history,lshaped.solverdata.σ)
    # Add ∞-norm auxilliary variable
    if lshaped.parameters.linearize
        # t
        MPB.addvar!(lshaped.mastersolver.lqmodel, -Inf, Inf, 1.0)
        if typeof(lshaped.mastersolver.optimsolver) == GurobiSolver
            updatemodel!(lshaped.mastersolver.lqmodel)
        end
    end
    # Add quadratic penalty
    c = copy(lshaped.c)
    append!(c, MPB.getobj(lshaped.mastersolver.lqmodel)[end-nbundles(lshaped)+1:end])
    add_penalty!(lshaped, lshaped.mastersolver.lqmodel, c, 1/lshaped.solverdata.σ, lshaped.ξ)
end

@implement_traitfn function log_regularization!(lshaped::AbstractLShapedSolver, RD)
    @unpack Q̃,σ = lshaped.solverdata
    push!(lshaped.Q̃_history, Q̃)
    push!(lshaped.σ_history, σ)
end

@implement_traitfn function log_regularization!(lshaped::AbstractLShapedSolver, t::Integer, RD)
    @unpack Q̃,σ = lshaped.solverdata
    lshaped.Q̃_history[t] = Q̃
    lshaped.σ_history[t] = σ
end

@implement_traitfn function take_step!(lshaped::AbstractLShapedSolver, RD)
    @unpack Q,Q̃,θ,σ = lshaped.solverdata
    @unpack τ,γ,σ̅,σ̲ = lshaped.parameters
    need_update = false
    if abs(θ-Q) <= τ*(1+abs(θ)) || lshaped.solverdata.major_iterations == 0
        lshaped.ξ[:] = lshaped.x[:]
        lshaped.solverdata.Q̃ = Q
        need_update = true
        lshaped.solverdata.major_iterations += 1
    else
        lshaped.solverdata.minor_iterations += 1
    end
    new_σ = if Q + τ <= (1-γ)*Q̃ + γ*θ
        min(σ̅,2*σ)
    elseif Q - τ >= γ*Q̃ + (1-γ)*θ
        max(σ̲,0.5*σ)
    else
        σ
    end
    if abs(new_σ-σ) > τ
        need_update = true
    end
    lshaped.solverdata.σ = new_σ
    if need_update
        c = copy(lshaped.c)
	append!(c, MPB.getobj(lshaped.mastersolver.lqmodel)[end-nbundles(lshaped)+1:end])
        add_penalty!(lshaped, lshaped.mastersolver.lqmodel, c, 1/lshaped.solverdata.σ, lshaped.ξ)
    end
    nothing
end

@implement_traitfn function solve_problem!(lshaped::AbstractLShapedSolver, solver::LQSolver, RD)
    if lshaped.parameters.linearize
        solve_linearized_problem!(lshaped, solver)
    else
        solver(lshaped.mastervector)
    end
end

# TR
# ------------------------------------------------------------
@define_traitfn TR set_trustregion!(lshaped::AbstractLShapedSolver)
@define_traitfn TR enlarge_trustregion!(lshaped::AbstractLShapedSolver)
@define_traitfn TR reduce_trustregion!(lshaped::AbstractLShapedSolver)

@implement_traitfn function init_solver!(lshaped::AbstractLShapedSolver, TR)
    if lshaped.parameters.autotune
        Δ = max(1.0,0.01*norm(lshaped.x,Inf))
        Δ̅ = max(1000.0,norm(lshaped.x,Inf))
        @pack! lshaped.parameters = Δ,Δ̅
    end
    lshaped.solverdata.Δ = lshaped.parameters.Δ
    set_trustregion!(lshaped)
end

@implement_traitfn function log_regularization!(lshaped::AbstractLShapedSolver, TR)
    @unpack Q̃,Δ = lshaped.solverdata
    push!(lshaped.Q̃_history,Q̃)
    push!(lshaped.Δ_history,Δ)
end

@implement_traitfn function log_regularization!(lshaped::AbstractLShapedSolver, t::Integer, TR)
    @unpack Q̃,Δ = lshaped.solverdata
    lshaped.Q̃_history[t] = Q̃
    lshaped.Δ_history[t] = Δ
end

@implement_traitfn function take_step!(lshaped::AbstractLShapedSolver, TR)
    @unpack Q,Q̃,θ = lshaped.solverdata
    @unpack γ = lshaped.parameters
    need_update = false
    if Q <= Q̃ - γ*abs(Q̃-θ)
        need_update = true
        enlarge_trustregion!(lshaped)
        lshaped.solverdata.cΔ = 0
        lshaped.ξ[:] = lshaped.x[:]
        lshaped.solverdata.Q̃ = Q
        lshaped.solverdata.major_iterations += 1
    else
        need_update = reduce_trustregion!(lshaped)
        lshaped.solverdata.minor_iterations += 1
    end
    if need_update
        set_trustregion!(lshaped)
    end
    nothing
end

@implement_traitfn function set_trustregion!(lshaped::AbstractLShapedSolver, TR)
    nb = nbundles(lshaped)
    l = max.(StochasticPrograms.get_stage_one(lshaped.stochasticprogram).colLower, lshaped.ξ .- lshaped.solverdata.Δ)
    append!(l, fill(-Inf,nb))
    u = min.(StochasticPrograms.get_stage_one(lshaped.stochasticprogram).colUpper, lshaped.ξ .+ lshaped.solverdata.Δ)
    append!(u, fill(Inf,nb))
    MPB.setvarLB!(lshaped.mastersolver.lqmodel, l)
    MPB.setvarUB!(lshaped.mastersolver.lqmodel, u)
end

@implement_traitfn function enlarge_trustregion!(lshaped::AbstractLShapedSolver, TR)
    @unpack Q,Q̃,θ = lshaped.solverdata
    @unpack τ,Δ̅ = lshaped.parameters
    if Q̃ - Q >= 0.5*(Q̃-θ) && abs(norm(lshaped.ξ-lshaped.x,Inf) - lshaped.solverdata.Δ) <= τ
        # Enlarge the trust-region radius
        lshaped.solverdata.Δ = min(Δ̅, 2*lshaped.solverdata.Δ)
        return true
    else
        return false
    end
end

@implement_traitfn function reduce_trustregion!(lshaped::AbstractLShapedSolver, TR)
    @unpack Q,Q̃,θ = lshaped.solverdata
    ρ = min(1, lshaped.solverdata.Δ)*(Q-Q̃)/(Q̃-θ)
    if ρ > 0
        lshaped.solverdata.cΔ += 1
    end
    if ρ > 3 || (lshaped.solverdata.cΔ >= 3 && 1 < ρ <= 3)
        # Reduce the trust-region radius
        lshaped.solverdata.cΔ = 0
        lshaped.solverdata.Δ = (1/min(ρ,4))*lshaped.solverdata.Δ
        return true
    else
        return false
    end
end

@implement_traitfn function process_cut!(lshaped::AbstractLShapedSolver, cut::HyperPlane{FeasibilityCut}, TR)
    @unpack τ = lshaped.parameters
    if !satisfied(cut,lshaped.ξ,τ)
        A = [I cut.δQ; cut.δQ' 0*I]
        b = [zeros(length(lshaped.ξ)); -gap(cut, lshaped.ξ)]
        t = A\b
        lshaped.ξ[:] = lshaped.ξ + t[1:length(lshaped.ξ)]
    end
end

# Level-set
# ------------------------------------------------------------

@implement_traitfn function init_solver!(lshaped::AbstractLShapedSolver, LV)
    # θs
    for i = 1:nbundles(lshaped)
        MPB.addvar!(lshaped.projectionsolver.lqmodel, -Inf, Inf, 0.0)
    end
    if lshaped.parameters.linearize
        # t
        MPB.addvar!(lshaped.projectionsolver.lqmodel, -Inf, Inf, 1.0)
    end
    if typeof(lshaped.projectionsolver.optimsolver) == GurobiSolver
        updatemodel!(lshaped.projectionsolver.lqmodel)
    end
end

@implement_traitfn function log_regularization!(lshaped::AbstractLShapedSolver, LV)
    @unpack Q̃ = lshaped.solverdata
    push!(lshaped.Q̃_history, Q̃)
end

@implement_traitfn function log_regularization!(lshaped::AbstractLShapedSolver, t::Integer, LV)
    @unpack Q̃ = lshaped.solverdata
    lshaped.Q̃_history[t] = Q̃
end

@implement_traitfn function take_step!(lshaped::AbstractLShapedSolver, LV)
    @unpack Q,Q̃ = lshaped.solverdata
    @unpack τ = lshaped.parameters
    if Q + τ <= Q̃
        lshaped.solverdata.Q̃ = Q
        lshaped.ξ[:] = lshaped.x[:]
    end
    nothing
end

@implement_traitfn function solve_problem!(lshaped::AbstractLShapedSolver, solver::LQSolver, LV)
    if lshaped.parameters.linearize
        solve_linearized_problem!(lshaped, solver)
    else
        solver(lshaped.mastervector)
    end
end

@implement_traitfn function process_cut!(lshaped::AbstractLShapedSolver, cut::HyperPlane, LV)
    MPB.addconstr!(lshaped.projectionsolver.lqmodel, lowlevel(cut)...)
    # TODO: Rewrite with MathOptInterface
end

@implement_traitfn function project!(lshaped::AbstractLShapedSolver{true}, LV)
    @unpack Q = lshaped.solverdata
    if Q < Inf
        _project!(lshaped)
    end
end

@implement_traitfn function project!(lshaped::AbstractLShapedSolver{false}, LV)
    _project!(lshaped)
end

function _project!(lshaped::AbstractLShapedSolver)
    @unpack θ,Q̃ = lshaped.solverdata
    @unpack λ = lshaped.parameters
    # Update level (TODO: Rewrite with MathOptInterface)
    nb = nbundles(lshaped)
    c = sparse(MPB.getobj(lshaped.mastersolver.lqmodel))
    L = (1-λ)*θ + λ*Q̃
    push!(lshaped.levels,L)
    if lshaped.solverdata.levelindex == -1
        MPB.addconstr!(lshaped.projectionsolver.lqmodel, c.nzind, c.nzval, -Inf, L)
        lshaped.solverdata.levelindex = first_stage_nconstraints(lshaped.stochasticprogram)+length(lshaped.cuts)+1
    else
        MPB.delconstrs!(lshaped.projectionsolver.lqmodel, [lshaped.solverdata.levelindex])
        MPB.addconstr!(lshaped.projectionsolver.lqmodel, c.nzind, c.nzval, -Inf, L)
        lshaped.solverdata.levelindex = first_stage_nconstraints(lshaped.stochasticprogram)+length(lshaped.cuts)+1
        if lshaped.parameters.linearize
            if typeof(lshaped.projectionsolver.optimsolver) == GurobiSolver
                updatemodel!(lshaped.projectionsolver.lqmodel)
            end
            lshaped.solverdata.regularizerindex -= 1
        end
    end
    # Update regularizer
    add_penalty!(lshaped, lshaped.projectionsolver.lqmodel, zeros(length(lshaped.ξ)+nb), 1.0, lshaped.ξ)
    # Solve projection problem
    solve_problem!(lshaped, lshaped.projectionsolver)
    if status(lshaped.projectionsolver) == :Infeasible
        @warn "Projection problem is infeasible, unprojected solution will be used"
        if Q̃ <= θ
            # If the upper objective bound is lower than the model lower bound for some reason, reset it.
            lshaped.solverdata.Q̃ = Inf
        end
    else
        # Update master solution
        ncols = decision_length(lshaped.stochasticprogram)
        x = getsolution(lshaped.projectionsolver)
        lshaped.mastervector[:] = x[1:ncols+nb]
        lshaped.x[1:ncols] = x[1:ncols]
        lshaped.θs[:] = x[ncols+1:ncols+nb]
    end
    nothing
end

function add_penalty!(lshaped::AbstractLShapedSolver, model::MPB.AbstractLinearQuadraticModel, c::AbstractVector, α::Real, ξ::AbstractVector)
    nb = nbundles(lshaped)
    if lshaped.parameters.linearize
        ncols = decision_length(lshaped.stochasticprogram)
        tidx = ncols+nb+1
        j = lshaped.solverdata.regularizerindex
        if j == -1
            for i in 1:ncols
                MPB.addconstr!(model, [i,tidx], [-α,1], -α*ξ[i], Inf)
                MPB.addconstr!(model, [i,tidx], [-α,-1], -Inf, -ξ[i])
            end
        else
            MPB.delconstrs!(model, collect(j:j+2*ncols-1))
            for i in 1:ncols
                MPB.addconstr!(model, [i,tidx], [-α,1], -ξ[i], Inf)
                MPB.addconstr!(model, [i,tidx], [-α,-1], -Inf, -ξ[i])
            end
        end
        lshaped.solverdata.regularizerindex = first_stage_nconstraints(lshaped.stochasticprogram)+length(lshaped.cuts)+1
        if hastrait(lshaped,LV)
            lshaped.solverdata.regularizerindex += 1
        end
    else
        # Linear part
        c[1:length(ξ)] -= α*ξ
        MPB.setobj!(model,c)
        # Quadratic part
        qidx = collect(1:length(ξ)+nb)
        qval = fill(α, length(lshaped.ξ))
        append!(qval, zeros(nb))
        if applicable(MPB.setquadobj!, model, qidx, qidx, qval)
            MPB.setquadobj!(model, qidx, qidx, qval)
        else
            error("Setting a quadratic penalty requires a solver that handles quadratic objectives")
        end
    end
end

function solve_linearized_problem!(lshaped::AbstractLShapedSolver, solver::LQSolver)
    push!(lshaped.mastervector, norm(lshaped.x-lshaped.ξ,Inf))
    solver(lshaped.mastervector)
    pop!(lshaped.mastervector)
    return nothing
end


text:
"""
	minimal_to_maximal(mechanism, y)

	convert minimal to maximal representation

	mechanism: Mechanism
	y: minimal state
"""
function minimal_to_maximal(mechanism::Mechanism{T,Nn,Ne,Nb,Ni}, y::AbstractVector) where {T,Nn,Ne,Nb,Ni}
	off = 0
	for id in mechanism.root_to_leaves
		(id > Ne) && continue # only treat joints
		joint = mechanism.joints[id]
		nu = input_dimension(joint)
		set_minimal_coordinates_velocities!(mechanism, joint, xmin=y[off .+ SUnitRange(1, 2nu)])
		off += 2nu
	end
	return get_maximal_state(mechanism)
end

"""
	maximal_to_minimal(mechanism, z)

	convert maximal to minimal representation

	mechanism: Mechanism
	z: maximal state
"""
function maximal_to_minimal(mechanism::Mechanism{T,Nn,Ne,Nb,Ni}, z::AbstractVector{Tz}) where {T,Nn,Ne,Nb,Ni,Tz}
	x = []
	for id in mechanism.root_to_leaves
		(id > Ne) && continue # only treat joints
		joint = mechanism.joints[id]
		c = zeros(Tz,0)
		v = zeros(Tz,0)
		ichild = joint.child_id - Ne
		for element in (joint.translational, joint.rotational)
			xb, vb, qb, ωb = unpack_maximal_state(z, ichild)
			if joint.parent_id != 0
				iparent = joint.parent_id - Ne
				xa, va, qa, ωa = unpack_maximal_state(z, iparent)
			else
				xa, va, qa, ωa = current_configuration_velocity(mechanism.origin.state)
			end
			push!(c, minimal_coordinates(element, xa, qa, xb, qb)...)
			push!(v, minimal_velocities(element, xa, va, qa, ωa, xb, vb, qb, ωb, mechanism.timestep)...)
		end
		push!(x, [c; v]...)
	end
	x = [x...]
	return x
end

function unpack_maximal_state(z::AbstractVector, i::Int)
	zi = z[(i-1) * 13 .+ (1:13)]
	x2 = zi[SUnitRange(1,3)]
	v15 = zi[SUnitRange(4,6)]
	q2 = Quaternion(zi[7:10]...)
	ω15 = zi[SUnitRange(11,13)]
	return x2, v15, q2, ω15
end

function pack_maximal_state!(z::AbstractVector,
	x2::AbstractVector, v15::AbstractVector, q2::Quaternion, ω15::AbstractVector,
	i::Int)

	z[(i-1) * 13 .+ (1:3)] = x2
	z[(i-1) * 13 .+ (4:6)] = v15
	z[(i-1) * 13 .+ (7:10)] = vector(q2)
	z[(i-1) * 13 .+ (11:13)] = ω15

	return nothing
end


text:
"The perfect sidekick to your scientific inquiries"
module DrWatson
import Pkg, LibGit2

const PATH_SEPARATOR = joinpath("_", "_")[2]

# Pure Julia implementation
include("project_setup.jl")
include("naming.jl")
include("saving_tools.jl")

# Functionality that saves/loads
using FileIO: save, load
export save, load
wsave = save
wload = load

include("saving_files.jl")

# Functionality that requires Dataframes and other heavy dependencies:
using Requires
function __init__()
    @require DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" begin
        include("result_collection.jl")
    end
end

# Update messages
const display_update = true
const update_version = "1.0.0"
const update_name = "update_v$update_version"
if display_update
if !isfile(joinpath(@__DIR__, update_name))
printstyled(stdout,
"""
\nUpdate message: DrWatson v$update_version

Welcome to the first major release!

Checkout the CHANGELOG for breaking changes.
\n
"""; color = :light_magenta)
touch(joinpath(@__DIR__, update_name))
end
end

end


text:
const Cftok = UInt32

const IPC_RMID = 0 # remove identifier
const IPC_SET  = 1 # set options
const IPC_STAT = 2 # get options

const IPC_CREAT  = 0o1000 # create entry if key does not exist
const IPC_EXCL   = 0o2000 # fail if key exists
const IPC_NOWAIT = 0o4000 # error if request must wait

ftok(path::String, id::Int) = ccall(:ftok, Cftok, (Cstring,Cint), path, id)

function shmget(tok::Cftok, size::Int; create::Bool=true, create_exclusive::Bool=false, permission::Integer=0o660)
    flags = Cint(permission)
    if create
        flags |= IPC_CREAT
        if create_exclusive
            flags |= IPC_EXCL
        end
    end
    id = ccall(:shmget, Cint, (Cftok, Csize_t, Cint), tok, size, flags)
    systemerror("error creating shared memory segment", id < 0)
    id
end

function shmat(id::Cint)
    ptr = ccall(:shmat, Ptr{Void}, (Cint, Ptr{Void}, Cint), id, C_NULL, 0)
    systemerror("error attaching shared memory segment", convert(Int, ptr) == -1)
    ptr
end

function shmdt(ptr::Ptr{Void})
    ret = ccall(:shmdt, Cint, (Ptr{Void},), ptr)
    systemerror("error detaching shared memory segment", ret == -1)
    nothing
end

function shmrm(id::Cint)
    ret = ccall(:shmctl, Cint, (Cint, Cint, Ptr{Void}), id, IPC_RMID, C_NULL)
    systemerror("error removing shared memory segment", ret == -1)
    nothing
end

struct SysVSharedArray{T,N}
    path::String
    id::Cint
    size::Int
    shmid::Cint
    shmptr::Ptr{Void}
    A::Array{T,N}
end

function SysVSharedVector{T}(path::Union{String,Tuple{String,Integer}}, size::Integer, ::Type{T}; create::Bool=true, create_exclusive::Bool=false, permission::Integer=0o660)
    id = 0
    if isa(path, Tuple)
        path,id = path
    end
    tok = ftok(path, id)
    shmid = shmget(tok, size*sizeof(T); create=create, create_exclusive=create_exclusive, permission=permission)
    shmptr = shmat(shmid)
    A = unsafe_wrap(Array, convert(Ptr{T}, shmptr), (size,))
    SysVSharedArray{T,1}(path, id, size, shmid, shmptr, A)
end

close(vec::SysVSharedArray) = shmdt(vec.shmptr)
delete!(vec::SysVSharedArray) = shmrm(vec.shmid)


text:
function open_file(filename)
    if Sys.isapple()
        run(`open $(filename)`)
    elseif Sys.islinux() || Sys.isbsd()
        run(`xdg-open $(filename)`)
    elseif Sys.iswindows()
        run(`$(ENV["COMSPEC"]) /c start $(filename)`)
    else
        @warn "Showing plots is not supported on OS $(string(Sys.KERNEL))"
    end
end


text:
"""
This function finds the most frequent labels of neighbors
and randomly returns one of them
"""
function maxvote{T}(label_count::Dict{Int,T})
  max_count = maximum(values(label_count))
  candidate_labels = Int[]
  for (k,v) in label_count
    if v == max_count
      push!(candidate_labels, k)
    end
  end
  sample(candidate_labels)
end

"""
Applies the value of label for each neighbor
and calls maxvote function
"""
function applyvote!{V}(v::V, g::AbstractGraph{V}, Label::Vector{Int})
  v_neighbors = out_neighbors(v, g)
  if length(v_neighbors)>0
    nLabels = Array(Int, length(v_neighbors))
    for (i, u) in enumerate(v_neighbors)
      nLabels[i] = Label[vertex_index(u, g)]
    end
    Label[vertex_index(v, g)] = maxvote(countmap(nLabels))
  end
end

function applyvote!{V, T<:Real}(v::V, g::AbstractGraph{V}, Label::Vector{Int}, ns::Vector{T}, λ::Real=1.0)
  v_edges = out_edges(v, g)
  if length(v_edges)>0
    label_count = Dict{Int,Float64}()
    for e in v_edges
      e_idx = edge_index(e, g)
      u = target(e, g)
      u_idx = vertex_index(u, g)
      if haskey(label_count, Label[u_idx])
        label_count[Label[u_idx]] += 1.0 + λ*ns[e_idx]
      else
        label_count[Label[u_idx]] = 1.0 + λ*ns[e_idx]
      end
    end
    Label[vertex_index(v, g)] = maxvote(label_count)
  end
end

"""
Perform label propagation alghrithm for each node in a random order
up to `iterations` times
"""
function label_propagation{V}(g::AbstractGraph{V}; iterations::Int=20)
  Label = [1:num_vertices(g)]
  for i=1:iterations
    order = shuffle(collect(vertices(g)))
    for v in order
      applyvote!(v, g, Label)
    end
  end
  permute_labels!(Label)
  Label
end

function label_propagation{V, T<:Real}(g::AbstractGraph{V}, ns::Vector{T}; λ::Real=1.0, iterations::Int=20)
  Label = [1:num_vertices(g)]
  for i=1:iterations
    order = shuffle(collect(vertices(g)))
    for v in order
      applyvote!(v, g, Label, ns, λ)
    end
  end
  permute_labels!(Label)
  Label
end


text: my_f(x, y) = 2x+3y


text:
abstract type EnergyCarrier end

mutable struct Electricity <: EnergyCarrier
    power::AbstractArray{Float64,3}
    Electricity() = new()
end

mutable struct Heat <: EnergyCarrier
    power::AbstractArray{Float64,3}
    Heat() = new()
end

mutable struct Hydrogen <: EnergyCarrier
    power::AbstractArray{Float64,3}
    Hydrogen() = new()
end


text:
################################################################################

using MindReader
using HiddenMarkovModelReaders

################################################################################

import Flux: cpu, gpu, flatten, leakyrelu

################################################################################

import Parameters: @with_kw

################################################################################

# read parameters
include( "Parameters.jl");

################################################################################

# argument parser
include( "Utilities/argParser.jl" );

################################################################################

# read data
begin
  # read edf file
  edfDf, _, _ = getSignals(shArgs)

  # calculate fft
  freqDc = extractFFT(edfDf, shArgs)
end;

################################################################################

# build autoencoder & train hidden Markov model
begin

  errDc = Dict{String, Tuple{Array{Int64, 1}, Array{Array{Float64, 1}, 1}}}()

  for (κ, υ) ∈ freqDc
    println()
    @info κ

    #  build & train autoencoder
    freqAr = shifter(υ)
    model = buildAutoencoder(length(freqAr[1]), nnParams = NNParams)
    model = modelTrain!(model, freqAr, nnParams = NNParams)

    ################################################################################

    # calculate post autoencoder
    postAr = cpu(model).(freqAr)

    ################################################################################

    begin
      @info "Creating Hidden Markov Model..."
      # error
      aErr = reshifter(postAr - freqAr) |> π -> flatten(π) |> permutedims

      # setup
      hmm = setup(aErr)

      # process
      for _ ∈ 1:4
        errDc[κ] = process!(hmm, aErr, true, params = hmmParams)
      end

      # final
      for _ ∈ 1:2
        errDc[κ] = process!(hmm, aErr, false, params = hmmParams)
      end
    end;

    ################################################################################

  end

  println()

end;

################################################################################

# write traceback & states
writeHMM(errDc, shArgs)

################################################################################

# graphic rendering
mindGraphics(errDc, shArgs)

################################################################################

max_stars_repo_name: DanielRivasMD/MindReader

max_issues_repo_name: DanielRivasMD/MindReader

max_forks_repo_name: DanielRivasMD/MindReader


text:
using ExamplePackage
ExamplePackage.greet()

b = 1.0
function g(a)
    global b
    for i ∈ 1:1_000_000
        tmp = a + b
    end
end

using BenchmarkTools
@show @btime g(1.0) #globals are bad for memory allocation
#@show @code_native g(1.0)

#w/o globals

function f(a, b)
    for i ∈ 1:1_000_000
        tmp = a + b
    end
end

@btime f(1.0,1.0)

const b_const = 1.0
function g(a)
    global b_const
    for i ∈ 1:1_000_000
        tmp = a + b_const
    end
end

struct Foo_generic
    a
end

struct Foo_abstract
    a::Real
end

struct Foo_concrete{T<:Real}
    a::T
end

fg = Foo_generic(1.0)
fa = Foo_abstract(1.0)
fc = Foo_concrete(1.0)

function f(foo)
    for i ∈ 1:1_000_000
        tmp = i + foo.a
    end
end

@show @btime f($fc)
@code_native f(fc)

max_stars_repo_name: shanemcmiken/quantecon-notebooks-julia

max_issues_repo_name: shanemcmiken/quantecon-notebooks-julia

max_forks_repo_name: shanemcmiken/quantecon-notebooks-julia


text:
struct Rejection{Tode} <: AbstractCHVIterator
	ode::Tode	# ODE solver to use for the flow in between jumps
end

function (rej::Rejection{Tode})(xdot, x, prob::Tpb, t) where {Tode, Tpb <: PDMPCaracteristics}
	prob.F(xdot, x, prob.xd, prob.parms, t)
	return nothing
end
###################################################################################################
### implementation of the rejection algo using DiffEq

# The following function is a callback to discrete jump. Its role is to perform the jump on the solution given by the ODE solver
# callable struct
function rejectionjump(integrator, prob::PDMPProblem, save_pre_jump, save_rate, verbose)
	# we declare the characteristics for convenience
	caract = prob.caract
	ratecache = caract.ratecache
	simjptimes = prob.simjptimes

	# final simulation time
	tf = prob.tspan[2]

	# find the next jump time
	t = integrator.t
	simjptimes.lastjumptime = t

	verbose && printstyled(color=:red,"--> REJECTION JUMP, t = $t\n")
	verbose && printstyled(color=:red,"----> xc = $(integrator.u)\n")

	verbose && printstyled(color=:green,"--> Fictitous jump at t = $t, # = ",simjptimes.fictitous_jumps," !!\n")

	simjptimes.ppf .= caract.R(ratecache.rate, integrator.u, caract.xd, caract.parms, t, true)
	@assert simjptimes.ppf[1] < simjptimes.ppf[2] "Error, your bound on the rates is not high enough!, $(simjptimes.ppf)"

	simjptimes.reject = rand() < 1 - simjptimes.ppf[1] / simjptimes.ppf[2]
	δt = -log(rand()) / simjptimes.ppf[2]

	verbose && printstyled(color=:green,"----> xc = ",caract.xc ,", xd = ",caract.xd,", reject = ",simjptimes.reject,", rates = ",simjptimes.ppf,"\n")

	# execute the jump
	if t < tf && simjptimes.reject == false
		verbose && println("----> Jump!, ratio = ", simjptimes.ppf[1] / simjptimes.ppf[2], ", xd = ", caract.xd)
		simjptimes.ppf .= caract.R(ratecache.rate, integrator.u, caract.xd, caract.parms, t, false)

		if (save_pre_jump) && (t <= tf)
			verbose && printstyled(color=:green,"----> save pre-jump\n")
			pushXc!(prob, integrator.u)
			pushXd!(prob, copy(caract.xd))
			pushTime!(prob, t)
			#save rates for debugging
			save_rate && push!(prob.rate_hist, sum(ratecache.rate))
		end

		# Update event
		ev = pfsample(ratecache.rate)

		# we perform the jump
		affect!(caract.pdmpjump, ev, integrator.u, caract.xd, caract.parms, t)
		u_modified!(integrator, true)

		@inbounds for ii in eachindex(caract.xc)
			caract.xc[ii] = integrator.u[ii]
		end
		simjptimes.njumps += 1
	else
		simjptimes.fictitous_jumps += 1
	end
	verbose && printstyled(color=:green,"----> jump effectued, xd = ",caract.xd,"\n")

	# we register the next time interval to solve the extended ode
	simjptimes.tstop_extended += δt
	add_tstop!(integrator, simjptimes.tstop_extended)
	verbose && printstyled(color=:green,"--> End jump\n\n")
end

function rejection_diffeq!(problem::PDMPProblem,
				ti::Tc, tf::Tc, verbose = false; ode = Tsit5(),
				save_positions = (false, true), n_jumps::Td = Inf64, reltol=1e-7, abstol=1e-9, save_rate = false, finalizer = finalize_dummy) where {Tc, Td}
	verbose && println("#"^30)
	verbose && printstyled(color=:red,"Entry in rejection_diffeq\n")
	ti, tf = problem.tspan
	algopdmp = Rejection(ode)

#ISSUE HERE, IF USING A PROBLEM p MAKE SURE THE TIMES in p.sim ARE WELL SET
	# set up the current time as the initial time
	t = ti
	# previous jump time, needed because problem.simjptimes.lastjumptime contains next jump time even if above tf
	tprev = t

	# initialise the problem. If I call twice this function, it should give the same result...
	init!(problem)

	# we declare the characteristics for convenience
	caract = problem.caract
	ratecache = caract.ratecache
	simjptimes = problem.simjptimes

	# vector to hold the state space
	X0 = copy(caract.xc)

	# current jump number
	# njumps = 0
	simjptimes.njumps = 1
	simjptimes.lambda_star = 0.0	# this is the bound for the rejection method
	simjptimes.ppf .= caract.R(ratecache.rate, X0, caract.xd, caract.parms, t, true)

	simjptimes.tstop_extended = simjptimes.tstop_extended / simjptimes.ppf[2] + ti

	simjptimes.reject = true

	# definition of the callback structure passed to DiffEq
	cb = DiscreteCallback(problem, integrator -> rejectionjump(integrator, problem, save_positions[1], save_rate, verbose), save_positions = (false, false))

	# define the ODE flow, this leads to big memory saving
	prob_REJ = ODEProblem((xdot, x, data, tt) -> algopdmp(xdot, x, caract, tt), X0, (ti, 1e9))
	integrator = init(prob_REJ, ode, tstops = simjptimes.tstop_extended, callback = cb, save_everystep = false, reltol = reltol, abstol = abstol, advance_to_tstop = true)

	while (t < tf) && simjptimes.njumps < n_jumps #&& simjptimes.fictitous_jumps < 10
		verbose && println("--> n = $(simjptimes.njumps), t = $t -> ", simjptimes.tstop_extended)
		step!(integrator)
		if t >= simjptimes.lastjumptime
			@warn "Could not compute next jump time $(simjptimes.njumps).\nReturn code = $(integrator.sol.retcode)\n $t < $(simjptimes.lastjumptime),\n solver = $ode"
			return PDMPResult(problem, save_positions)
		end
		t, tprev = simjptimes.lastjumptime, t

		# the previous step was a jump!
		if save_positions[2] && (t <= tf) && simjptimes.reject == false
			verbose && println("----> save post-jump, xd = ", problem.Xd)
			pushTime!(problem ,t)
			pushXc!(problem, copy(caract.xc))
			pushXd!(problem, copy(caract.xd))
			#save rates for debugging
			save_rate && push!(problem.rate_hist, sum(ratecache.rate))
			verbose && println("----> end save post-jump, ")
			#put the flag for rejection
			simjptimes.reject = true
		end
		finalizer(ratecache.rate, caract.xc, caract.xd, caract.parms, t)
	end
	# we check whether the last bit [t_last_jump, tf] is missing
	if t>tf
		verbose && println("----> LAST BIT!!, xc = ", caract.xc[end], ", xd = ",caract.xd)
		prob_last_bit = ODEProblem((xdot,x,data,tt) -> caract.F(xdot, x, caract.xd, caract.parms, tt), copy(caract.xc), (tprev, tf))
		sol = DiffEqBase.solve(prob_last_bit, ode)
		verbose && println("-------> xc[end] = ", sol.u[end])
		pushXc!(problem, sol.u[end])
		pushXd!(problem, copy(caract.xd))
		pushTime!(problem, sol.t[end])
	end
	return PDMPResult(problem, save_positions)
end

function solve(problem::PDMPProblem{Tc, Td, vectype_xc, vectype_xd, Tcar}, algo::Rejection{Tode}; verbose = false, n_jumps = Inf64, save_positions = (false, true), reltol = 1e-7, abstol = 1e-9, save_rate = true, finalizer = finalize_dummy) where {Tc, Td, vectype_xc, vectype_xd, vectype_rate, Tnu, Tp, TF, TR, Tcar, Tode <: DiffEqBase.DEAlgorithm}

	return rejection_diffeq!(problem, problem.tspan[1], problem.tspan[2], verbose; ode = algo.ode, save_positions = save_positions, n_jumps = n_jumps, reltol = reltol, abstol = abstol, save_rate = save_rate, finalizer = finalizer )
end


text:
function write!(fn::AbstractString, ga::GeoArray, nodata = nothing, shortname = find_shortname(fn))
    options = String[]
    w, h, b = size(ga)
    dtype = eltype(ga)
    data = copy(ga.A)
    use_nodata = false

    # Set compression options for GeoTIFFs
    shortname == "GTiff" && (options = ["COMPRESS=DEFLATE","TILED=YES", "NUM_THREADS=4"]) # NUM_THREADS=4

    # Slice data and replace missing by nodata
    if isa(dtype, Union) && dtype.a == Missing
        dtype = dtype.b
        try convert(ArchGDAL.GDALDataType, dtype)
            nothing
        catch
            dtype, data = cast_to_gdal(data)
        end
        nodata === nothing && (nodata = typemax(dtype))
        m = ismissing.(data)
        data[m] .= nodata
        data = Array{dtype}(data)
        use_nodata = true
    end

    try convert(ArchGDAL.GDALDataType, dtype)
        nothing
    catch
        dtype, data = cast_to_gdal(data)
    end

    ArchGDAL.create(fn, driver=ArchGDAL.getdriver(shortname), width=w, height=h, nbands=b, dtype=dtype, options=options) do dataset
        for i = 1:b
            band = ArchGDAL.getband(dataset, i)
            ArchGDAL.write!(band, data[:,:,i])
            use_nodata && ArchGDAL.GDAL.gdalsetrasternodatavalue(band.ptr, nodata)
        end

        # Set geotransform and crs
        gt = affine_to_geotransform(ga.f)
        ArchGDAL.GDAL.gdalsetgeotransform(dataset.ptr, gt)
        ArchGDAL.GDAL.gdalsetprojection(dataset.ptr, GFT.val(ga.crs))

    end
    fn
end

"""
    st_write(ga::GeoArray, fn::AbstractString; nodata = nothing)
    st_write(A::AbstractArray{T}, b::bbox, fn::AbstractString;nodata = nothing)

# Arguments
- `A`: 2d or 3d array
- `fn`: outfile path
- `nodata`: should be in the same data type as ga.A (or A)
"""
st_write(ga::GeoArray, fn::AbstractString; nodata = nothing) = write!(fn, ga, nodata)

function st_write(A::AbstractArray{T}, b::bbox, fn::AbstractString; nodata = nothing) where T <: Real
    ga = GeoArray(A, b)
    st_write(ga, fn; nodata = nodata)
end

export st_write


text:
"""
Update prior in `PFWI` either using mediumm or mediumi

* `priori` : prior model on inversion grid
"""
function update_prior!(pa; priori=nothing, priorm=nothing)

	if(priori===nothing) # use current mediumi as prior!
		if(priorm===nothing)
			copyto!(pa.priori, pa.mediumi)
		else
			interp_spray!(priorm, pa.priori, :interp)
			copyto!(pa.priori, pa.mediumi)
		end
	else
		!(isequal(priori.mgrid,pa.mediumi.mgrid)) && error("priori model has to be on igrid")
		copyto!(pa.priori, priori)
	end
	# update pa.mx.prior
	copyto!(pa.mx.prior,pa.priori,pa.parameterization)
end


text:
#! /usr/bin/env julia

using Pkg

using libfive_jll

using Clang.Generators

include_dir = joinpath(libfive_jll.artifact_dir, "include") |> normpath
clang_dir = joinpath(include_dir, "clang-c")

options = load_options(joinpath(@__DIR__, "generator.toml"))

# add compiler flags, e.g. "-DXXXXXXXXX"
args = get_default_args()
push!(args, "-I$include_dir")

headers = [joinpath(include_dir, "libfive.h")]
# there is also an experimental `detect_headers` function for auto-detecting top-level headers in the directory
# headers = detect_headers(clang_dir, args)

# create context
ctx = create_context(headers, args, options)

# run generator
build!(ctx)


text:
using LogDensityProblems, Test, Distributions, TransformVariables
import LogDensityProblems: capabilities, dimension, logdensity
using LogDensityProblems: logdensity_and_gradient, LogDensityOrder

import ForwardDiff, Flux, Zygote, TransformVariables, Random
using Parameters: @unpack

using Plots
gr()

Random.seed!(3)

####
#### transformed Bayesian problem
####

ProjDir = @__DIR__
cd(ProjDir)
include("../src/HMC.jl")
include("../src/generate_n_samples.jl")

# ### snippet 9.3

# Construct the logdensity problem

struct clip_9_3_model{TY <: AbstractVector, TX <: AbstractVector}
    "Observations."
    y::TY
    "Covariate"
    x::TX
end

# Make the type callable with the parameters *as a single argument*.

function (problem:: clip_9_3_model)(θ)
    @unpack y, x, = problem    # extract the data
    @unpack muy, mux = θ     # works on the named tuple too
    ll = 0.0
    ll += loglikelihood(Normal(mux, 1), x)
    ll += loglikelihood(Normal(muy, 1), y)
    ll += logpdf(Normal(0, 1), mux)
    ll += logpdf(Normal(0, 1), muy)
    ll
end

# Instantiate the model with data and inits.

Random.seed!(3)

N = 100
x = rand(Normal(0, 1), N)
#=
x = [0.83975036,  0.70534183,  1.30596472, -1.38799622,  1.27291686,  0.18419277,
  0.75227990,  0.59174505, -0.98305260, -0.27606396, -0.87085102,  0.71871055,
  0.11065288, -0.07846677, -0.42049046, -0.56212588,  0.99751344, -1.10513006,
  -0.14228783,  0.31499490]
=#
y = rand(Normal(0, 1), N)
#=
y = [2.287247161, -1.196771682, -0.694292510, -0.412292951, -0.970673341, -0.947279945,
  0.748139340, -0.116955226,  0.152657626,  2.189978107,  0.356986230,  2.716751783,
  2.281451926,  0.324020540,  1.896067067,  0.467680511, -0.893800723, -0.307328300,
  -0.004822422,  0.988164149]
=#

p = clip_9_3_model(y, x)
θ = (muy = 0.0, mux=0.0)

# Write a function to return properly dimensioned transformation.

t = as((muy = asℝ, mux = asℝ))
#t1 = as((muy = asℝ₊, mux = asℝ₊))

problem_transformation(p::clip_9_3_model) = t

# Wrap the problem with a transformation, then use Flux for the gradient.

P = TransformedLogDensity(problem_transformation(p), p)
∇P = ADgradient(:ForwardDiff, P);

println()
p(θ) |> display
v, g = logdensity_and_gradient(∇P, [0.0, 0.0])
println()
v  |> display
g  |> display
value_U = -v
grad_U = -g
v, g = logdensity_and_gradient(∇P, [-0.1, 0.2])
println()
v  |> display
g  |> display
value_U = -v
grad_U = -g

@testset "transformed Bayesian problem 2D" begin

    for _ in 1:100
        x = random_arg(t)
        θ, lj = transform_and_logjac(t, x)
        px = logdensity(P, x)
        px2, ∇px = logdensity_and_gradient(∇P, x)
        @test px == px2
    end

end

println()
new_q, ptraj, qtraj, accept, dH = HMC(p, ∇P, 0.03, 11, [-0.1, 0.2])

new_q |> display
println()
ptraj |> display
println()
qtraj |> display
println()
dH |> display
println()

fig = generate_n_samples(p, ∇P, q =  [-0.1, 0.2]);
savefig(fig, "fig-9-3.pdf")


text:
@testset "FastNAU" begin
    nau = NAU(2,3)
    x = rand(Float32,2,10)
    z = nau(x)
    ps = params(nau)

    loss() = sum(abs2, nau(x) .- reshape(x[1,:],1,:))
    gs = Flux.gradient(loss, ps)
    gs = vcat([vec(gs[p]) for p in ps]...)

    fastnau = FastNAU(2,3)
    loss1(p) = sum(abs2, fastnau(x,p) .- reshape(x[1,:],1,:))
    @test DiffEqFlux.paramlength(fastnau) == 6
    @test length(DiffEqFlux.initial_params(fastnau)) == 6

    p = vec(nau.W)
    gf = Flux.gradient(loss1, p)[1]
    @test all(gs .== gf)

    loss2(x) = sum(abs2, nau(x) .- reshape(x[1,:],1,:))
    gs = Flux.gradient(loss2, x)

    loss3(x) = sum(abs2, fastnau(x,p) .- reshape(x[1,:],1,:))
    gf = Flux.gradient(loss3, x)

    @test all(gs .== gf)
end


text:
include("util.jl")
include("graph.jl")
include("opinion.jl")

# function convex_approach(Q::Array{WEIGHT, 2}, y::OPINIONS, budget::WEIGHT)
#     z = Q * y
#     N = size(y)
#     N = N[1]
#     δ = Convex.Variable(N)
#     problem = minimize(sumsquares(z + (Q * δ)))
#     problem.constraints += (z + δ) <= 1
#     problem.constraints += (z + δ) >= -1
#     problem.constraints += sum((dot(^)(δ,2))) <= budget
#     solve!(problem, SCS.Optimizer)
#     @show problem.status
#     @show (z + (Q * δ.value))
#     return δ.value
# end

cost_of_removing(x) = x * x

function convex_approach(Q::Matrix{WEIGHT}, y::OPINIONS,
                        α::PSYCHOLOGICAL_FACTOR, budget::WEIGHT)
    z = Q * (α .* y)
    N = size(y)
    N = N[1]
    δ = Variable(N)
    problem = minimize(sumsquares(z + (Q * (α .* δ))))
    problem.constraints += (z + δ) <= 1
    problem.constraints += (z + δ) >= -1
    problem.constraints += sum((dot(^)(δ,2))) <= budget
    solve!(problem, SCS.Optimizer)
    @show problem.status
    @show Q * (α .* (y + δ))
    return δ.value
end

function mult_but_exclude(M::Matrix{WEIGHT}, y::OPINIONS,
    exclude::Array{Int64, 1})

    temp_M = M[:, setdiff(1:end, exclude)]
    temp_y = y[setdiff(1:end, exclude)]

    return (temp_M * temp_y)

end

function greedy_approach(Q::Matrix{WEIGHT}, y::OPINIONS, α::PSYCHOLOGICAL_FACTOR, budget::WEIGHT)

    (n, m) = size(Q)
    T = Q * diagm(α)
    total_cost = 0
    prev_len = 0
    to_remove = fill(0, 0)

    # while we can still potentially remove a node's influence
    while total_cost < budget
        current_cost = total_cost
        best_so_far = WEIGHT(Inf)
        current_to_remove = -1

        # consider each node
        for i = 1:n
            # get cost of adding node
            cost_to_add_i = (cost_of_removing(y[i]) + current_cost)

            # if we have not considered adding that node and
            # our budget would allow us to add them
            if i ∉ to_remove && (cost_to_add_i <= budget)
                # consider them
                push!(to_remove, i)
                # get change to the norm
                marginal_norm = norm(mult_but_exclude(T, y, to_remove))

                # if the norm obtained is lower than our current best in this
                # iteration
                if marginal_norm < best_so_far
                    # flag it as the best to remove
                    current_to_remove = i
                    best_so_far = marginal_norm
                end
                # pop off to remove temporarily
                pop!(to_remove)
            end
        end

        # if we found a viable node to remove
        if current_to_remove != -1
            # add it to remove list
            push!(to_remove, current_to_remove)
             #@show current_to_remove
        end

        # if (current_cost == total_cost) || (length(to_remove) == n)
        #     break
        if (length(to_remove) == n)
            break
        else
            total_cost = current_cost + cost_of_removing(y[to_remove[end]])
        end
    end

    δ = WEIGHT.(zeros(n))
    δ[to_remove] = -y[to_remove]

    return δ, total_cost, (T * (y + δ))
end

function select_using_centrality(Q::Matrix{WEIGHT}, y::OPINIONS,
    α::PSYCHOLOGICAL_FACTOR, budget::WEIGHT,
    G::SimpleGraph{Int64}, centrality::String)

    centralities = CENTRALITIES[centrality](G)
    (n, m) = size(Q)
    T = Q * diagm(α)

    value = fill(WEIGHT(0.0), n)
    cost = fill(WEIGHT(0.0), n)
    z_inf = T * y
    original = norm(z_inf)
    mean_y = mean(y)
    for i = 1:n
        value[i] = 1 * centralities[i]
        cost[i] = cost_of_removing(y[i])
    end
    perm = sortperm(-value)
    δ = fill(WEIGHT(0.0), n)
    total_cost = 0
    for i ∈ perm
        if (cost[i] + total_cost) <= budget
            total_cost += cost[i]
            δ[i] = -y[i]
        else
            remaining = budget - total_cost
            fraction = sqrt(remaining)
            δ[i] = -(sign(y[i])) * fraction
            break
        end
    end
    z = T * (y + δ)
    return (δ, total_cost, z)

end

function select_using_pagerank(Q::Matrix{WEIGHT}, y::OPINIONS,
    α::PSYCHOLOGICAL_FACTOR, budget::WEIGHT, A)

    g = construct_graph(A)
    centralities = closeness_centrality(g)

    (n, m) = size(Q)
    T = Q * diagm(α)
    value = fill(WEIGHT(0.0), n)
    cost = fill(WEIGHT(0.0), n)
    z_inf = T * y
    original = norm(z_inf)
    for i = 1:n
        value[i] = centralities[i]
        cost[i] = cost_of_removing(y[i])
    end

    perm = sortperm(-value)
    δ = fill(WEIGHT(0.0), n)
    total_cost = 0
    for i ∈ perm
        if (cost[i] + total_cost) <= budget
            total_cost += cost[i]
            δ[i] = -y[i]
        else
            remaining = budget - total_cost
            fraction = sqrt(remaining)
            δ[i] = -(sign(y[i])) * fraction
            break
        end
    end
    z = T * (y + δ)
    return (δ, total_cost, z)

end

function fractional_knapsack(Q::Matrix{WEIGHT}, y::OPINIONS,
    α::PSYCHOLOGICAL_FACTOR, budget::WEIGHT)
    (n, m) = size(Q)
    T = Q * diagm(α)
    value = fill(WEIGHT(0.0), n)
    cost = fill(WEIGHT(0.0), n)
    z_inf = T * y
    original = norm(z_inf)
    for i = 1:n
        v = mult_but_exclude(T, y, [i])
        v = norm(z_inf) - norm(v)
        value[i] = 0
        if value[i] > 0
            value[i] = v
        end
        cost[i] = cost_of_removing(y[i])
    end

    perm = sortperm(-value)
    δ = fill(WEIGHT(0.0), n)
    total_cost = 0
    for i ∈ perm
        if (cost[i] + total_cost) <= budget
            total_cost += cost[i]
            δ[i] = -y[i]
        else
            remaining = budget - total_cost
            fraction = sqrt(remaining)
            δ[i] = -(sign(y[i])) * fraction
            break
        end
    end
    z = T * (y + δ)
    return (δ, total_cost, z)
end

function plot_multiple_series(series)
    p = plot()
    n, m = size(series)
    x = 1:m
    for i = 1:n
        plot!(p, x, series[i, :])
    end
end

function approx_equlibrium_vector(M::SparseMatrixCSC, y::OPINIONS,
     α::PSYCHOLOGICAL_FACTOR, k::Int64)
    #k = Int64(ceil(compute_num_iters(p.y, p.M, epsilon, p.α)))
    z = y
    for i = 1:k
        z = (α .* y) + (M * z)
    end
    return z
 end

 function greedy_approach_approx(M::SparseMatrixCSC{WEIGHT}, y::OPINIONS,
     α::PSYCHOLOGICAL_FACTOR, budget::WEIGHT, epsilon::WEIGHT)

     k = Int64(ceil(compute_num_iters(y, M, epsilon, α)))

     (n, m) = size(M)
     #T = Q * diagm(α)
     total_cost = 0
     to_remove = fill(0, 0)

     # while we can still potentially remove a node's influence
     while (total_cost < budget) && (length(to_remove) < n)
         current_cost = total_cost
         best_so_far = WEIGHT(Inf)
         current_to_remove = -1

         # consider each node
         for i = 1:n
             # get cost of adding node
             cost_to_add_i = (cost_of_removing(y[i]) + current_cost)

             # if we have not considered adding that node and
             # our budget would allow us to add them
             if i ∉ to_remove && (cost_to_add_i <= budget)
                 # consider them
                 push!(to_remove, i)
                 # get change to the norm
                 temp_y = copy(y)
                 temp_y[to_remove] .= 0
                 temp_eqilibrium = approx_equlibrium_vector(M, temp_y, α, k)
                 marginal_norm = norm(temp_eqilibrium)

                 # if the norm obtained is lower than our current best in this
                 # iteration
                 if marginal_norm < best_so_far
                     # flag it as the best to remove
                     current_to_remove = i
                     best_so_far = marginal_norm
                 end
                 # pop off to remove temporarily
                 pop!(to_remove)
             end
         end

         # if we found a viable node to remove
         if current_to_remove != -1
             # add it to remove list
             push!(to_remove, current_to_remove)
             #@show current_to_remove
         end

         # if (current_cost == total_cost) || (length(to_remove) == n)
         #     break
         total_cost = current_cost + cost_of_removing(y[to_remove[end]])
     end

     δ = WEIGHT.(zeros(n))
     δ[to_remove] = -y[to_remove]
     # @show M
     # @show y + δ
     z = approx_equlibrium_vector(M, y + δ, α, k)
     return δ, total_cost, z
 end

function get_approx_mat(M::SparseMatrixCSC{WEIGHT}, α::PSYCHOLOGICAL_FACTOR, k::Int64)

    p1 = M ^ 0
    powers = Array{SparseMatrixCSC{WEIGHT}, 1}()
    push!(powers, p1)
    for i = 1:k
        temp = powers[i] * M
        push!(powers, temp)
    end

    p2 = sum(powers[1:end-1]) * diagm(α)
    p2 += powers[end]
    return p2

end

 function fractional_knapsack_approx(M::SparseMatrixCSC{WEIGHT}, y::OPINIONS,
     α::PSYCHOLOGICAL_FACTOR, budget::WEIGHT, epsilon::WEIGHT)

     k = Int64(ceil(compute_num_iters(y, M, epsilon, α)))

     (n, m) = size(M)
     value = fill(WEIGHT(0.0), n)
     cost = fill(WEIGHT(0.0), n)
     # T = get_approx_mat(M, α, k)
     # @show T
     # z_inf = norm(T * y)
     z_inf = approx_equlibrium_vector(M, y, α, k)
     original = norm(z_inf)
     for i = 1:n
         temp_y = copy(y)
         temp_y[i] = 0
         v = approx_equlibrium_vector(M, temp_y, α, k)
         #v = T * temp_y
         v = norm(z_inf) - norm(v)
         value[i] = 0
         if value[i] > 0
             value[i] = v
         end
         cost[i] = cost_of_removing(y[i])
     end

     perm = sortperm(-value)
     δ = fill(WEIGHT(0.0), n)
     total_cost = 0
     for i ∈ perm
         if (cost[i] + total_cost) <= budget
             total_cost += cost[i]
             δ[i] = -y[i]
         else
             remaining = budget - total_cost
             fraction = sqrt(remaining)
             δ[i] = -(sign(y[i])) * fraction
             break
         end
     end
     z = approx_equlibrium_vector(M, y + δ, α, k)
     return (δ, total_cost, z)
 end

# function main()
#
#     DATASET = "congress"
#     BUDGET = WEIGHT(3.0)
#     filepath = "./data/raw/$(DATASET).edgelist"
#
#
#     @show filepath
#     (W, U, n) = read_edgelist(filepath)
#     normalize_edgelist!(W)
#     normalize_edgelist!(U)
#
#     y = generate_opinions(n)
#     α, β, γ = generate_factors(n)
#     M = (β .* W) - (γ .* U)
#     G = construct_graph(M)
#
#     M = Matrix(M)
#     Q = inv(I - M)
#
#     z = Q * (α .* y)
#
#
#
#     (δ, total_cost, z_prime) =  greedy_approach(Q, y, α,  BUDGET)
#
#
#
#     δ1 = convex_approach(Q, y, α,  BUDGET)
#     res = Q * (α .* (y + δ1))
#
#     (δ, total_cost, z_knap) = fractional_knapsack(Q, y, α,  BUDGET)
#
#     (δ, total_cost, z_page_rank) = select_using_pagerank(Q, y, α,  BUDGET, M)
#
#     @show norm(y)
#     @show norm(z)
#     @show norm(z_prime)
#     @show norm(z_knap)
#     @show norm(z_page_rank)
#     @show norm(res)
#
#     times = n * 2
#
#     trace = fill(WEIGHT(0.0), n, times)
#     trace[:,1] = y
#     for i = 2:times
#         next_opinion = get_next_opinion(α, β, γ,W, U, y, trace[:,i - 1])
#         trace[:,i] = next_opinion
#     end
#
#
#     # plot_multiple_series(trace)
#     # plot(1:times)
#     return (y, Q, α, δ, total_cost, z_prime, G, trace)
#     # return y, z, res
# end
#
# #y, z, res = main()
# (y, Q, α, δ, total_cost, z_prime, G, trace) = main()

max_stars_repo_name: InzamamRahaman/SignedDegroot

max_issues_repo_name: InzamamRahaman/SignedDegroot

max_forks_repo_name: InzamamRahaman/SignedDegroot