text:
"""
The Gr4j type contains the state variables (st, st_uh1, st_uh2), the inputs
(epot, prec) for one time step, the parameters (param) and the time step
length (tstep) for the GR4J model.
"""
type Gr4j <: Hydro

  st::Array{Float64,1}
  st_uh1::Array{Float64,1}
  st_uh2::Array{Float64,1}
  ord_uh1::Array{Float64,1}
  ord_uh2::Array{Float64,1}
  epot::Float64
  prec::Float64
  q_sim::Float64
  param::Array{Float64,1}
  tstep::Float64
  time::DateTime

end

"""
    Gr4j(tstep)
Constructor for GR4J with predefined state variables, parameters and inputs.
The time step (tstep) is given as a fraction of one day. Thus, for hourly input
data tstep should be set to 1/24.
"""
function Gr4j(tstep::Float64, time::DateTime)

  n_ord = ceil(Int64, 20.0 * 24.0 / tstep)

  st      = zeros(Float64, 2)
  st_uh1  = zeros(Float64, n_ord)
  st_uh2  = zeros(Float64, 2 * n_ord)
  ord_uh1 = zeros(Float64, n_ord)
  ord_uh2 = zeros(Float64, 2 * n_ord)

  epot = 0.0
  prec = 0.0

  q_sim = 0.0

  param  = [257.238, 1.012, 88.235, 2.208]

  st[1] = 0.3 * param[1]
  st[2] = 0.5 * param[3]

  D = 1.30434782 * (tstep / 24.0) + 1.19565217

  UH1(ord_uh1, param[4] * 24.0 / tstep, D)
  UH2(ord_uh2, param[4] * 24.0 / tstep, D)

  Gr4j(st, st_uh1, st_uh2, ord_uh1, ord_uh2, epot, prec, q_sim, param, tstep, time)

end

"""
    Gr4j(tstep, time, param)
Constructor for GR4J with predefined state variables and inputs. The parameter
values are given as input. The time step (tstep) is given as a fraction of one
day. Thus, for hourly input data tstep should be set to 1/24.
"""
function Gr4j(tstep::Float64, time::DateTime, param::Array{Float64,1})

  n_ord = ceil(Int64, 20.0 * 24.0 / tstep)

  st      = zeros(Float64, 2)
  st_uh1  = zeros(Float64, n_ord)
  st_uh2  = zeros(Float64, 2 * n_ord)
  ord_uh1 = zeros(Float64, n_ord)
  ord_uh2 = zeros(Float64, 2 * n_ord)

  st[1] = 0.3 * param[1]
  st[2] = 0.5 * param[3]

  D = 1.30434782 * (tstep / 24.0) + 1.19565217

  UH1(ord_uh1, param[4] * 24.0 / tstep, D)
  UH2(ord_uh2, param[4] * 24.0 / tstep, D)

  epot = 0.0
  prec = 0.0

  q_sim = 0.0

  Gr4j(st, st_uh1, st_uh2, ord_uh1, ord_uh2, epot, prec, q_sim, param, tstep, time)

end

"""
    init_states(mdata::Gr4j)
Initilize the state variables of the model.
"""
function init_states(mdata::Gr4j)

  mdata.st[1] = 0.3 * mdata.param[1]
  mdata.st[2] = 0.5 * mdata.param[3]

  for i in eachindex(mdata.st_uh1)
    mdata.st_uh1[i] = 0.
  end

  for i in eachindex(mdata.st_uh2)
    mdata.st_uh2[i] = 0.
  end

end

"""
    get_param_range(mdata::Gr4j)
Get allowed parameter ranges for the calibration of the model.
"""
function get_param_range(mdata::Gr4j)

  param_range_hydro = [(1.0, 20000.0), (-100.0, 100.0), (1.0, 20000.0), (0.5, 10.0)]

end

"""
    get_states(mdata::Gr4j)

Get state variables for computing penelty during calibration.

"""
function get_states(mdata::Gr4j)

  return []

end

"""
    assign_param(mdata::Gr4j, param::Array{Float64,1})
Assign parameter values to the Gr4j type.
"""
function assign_param(mdata::Gr4j, param::Array{Float64,1})

  for i in eachindex(mdata.param)
    mdata.param[i] = param[i]
  end

  D = 1.30434782 * (mdata.tstep / 24.0) + 1.19565217

  UH1(mdata.ord_uh1, param[4] * 24.0 / mdata.tstep, D)
  UH2(mdata.ord_uh2, param[4] * 24.0 / mdata.tstep, D)

end

"""
    get_states(hydro_out::Array{Gr4j,1})
Get state variables
"""
function get_states(hydro_out::Array{Gr4j,1})

  states = zeros(length(hydro_out), 2)

  for i in 1:length(hydro_out)

    states[i, 1] = hydro_out[i].st[1]
    states[i, 2] = hydro_out[i].st[2]

  end

  return states

end

"""
    enkf_hydro(mdata::Array{Gr4j,1}, obs_ens, q_sim)
Implementation of the ensemble Kalman filter for the GR4J model.
"""
function enkf_hydro(mdata::Array{Gr4j,1}, obs_ens, q_sim)

  nens = length(mdata)

  # Allocate arrays

  st     = zeros(Float64, length(mdata[1].st), nens)
  st_uh1 = zeros(Float64, length(mdata[1].st_uh1), nens)
  st_uh2 = zeros(Float64, length(mdata[1].st_uh2), nens)

  # Add states to arrays

  for iens = 1:nens

    st[:, iens]     = mdata[iens].st
    st_uh1[:, iens] = mdata[iens].st_uh1
    st_uh2[:, iens] = mdata[iens].st_uh2

  end

  # Run ensemble kalman filter

  st     = enkf(st, obs_ens, q_sim)
  st_uh1 = enkf(st_uh1, obs_ens, q_sim)
  st_uh2 = enkf(st_uh2, obs_ens, q_sim)

  # Check limits of states

  st[st .< 0] = 0.
  st_uh1[st_uh1 .< 0] = 0.
  st_uh2[st_uh2 .< 0] = 0.

  st[1, st[1, :] .> mdata[1].param[1]] = mdata[1].param[1]
  st[2, st[2, :] .> mdata[1].param[3]] = mdata[1].param[3]

  # Add arrays to states

  for iens = 1:nens

    mdata[iens].st = st[:, iens]
    mdata[iens].st_uh1 = st_uh1[:, iens]
    mdata[iens].st_uh2 = st_uh2[:, iens]

  end

  nothing

end

"""
    run_timestep(mdata::Gr4j)
Propagate the model one time step and return simulated dischage.
"""
function run_timestep(mdata::Gr4j)

  St     = mdata.st
  StUH1  = mdata.st_uh1
  StUH2  = mdata.st_uh2
  OrdUH1 = mdata.ord_uh1
  OrdUH2 = mdata.ord_uh2
  Param  = mdata.param
  P1     = mdata.prec
  E      = mdata.epot
  tstep  = mdata.tstep

  A = Param[1]
  B = 0.9

  # Interception and production store

  if P1 <= E
    EN = E - P1
    PN = 0.0
    WS = EN / A
    if WS > 13.0
      WS = 13.0
    end
    TWS = tanh(WS)
    Sr = St[1] / A
    ER = St[1] * (2.0-Sr)*TWS/(1.0+(1.0-Sr)*TWS)
    AE = ER + P1
    St[1] = St[1] - ER
    PR = 0.0
  else
    EN = 0.0
    AE = E
    PN = P1 - E
    WS = PN / A
    if WS > 13.0
      WS = 13.0
    end
    TWS = tanh(WS)
    Sr = St[1] / A
    PS = A*(1.0-Sr*Sr)*TWS/(1.0+Sr*TWS)
    PR = PN - PS
    St[1] = St[1] + PS
  end

  # Percolation from production store

  if St[1] < 0.0
    St[1] = 0.0
  end

  scale_param = 25.62891 * (24.0 / tstep)

  Sr = St[1] / Param[1]
  Sr = Sr * Sr
  Sr = Sr * Sr
  PERC = St[1] * (1.0-1.0/sqrt(sqrt(1.0 + Sr/scale_param)))

  St[1] = St[1] - PERC

  PR = PR + PERC

  # Split of effective rainfall into the two routing components

  PRHU1 = PR * B
  PRHU2 = PR * (1.0-B)

  # Convolution of unit hydrograph UH1

  NH = length(OrdUH1)

  for K = 1:NH-1
    StUH1[K] = StUH1[K+1] + OrdUH1[K]*PRHU1
  end

  StUH1[NH] = OrdUH1[NH] * PRHU1

  # Convolution of unit hydrograph UH2

  NH = length(OrdUH2)

  for K = 1:NH-1
    StUH2[K] = StUH2[K+1] + OrdUH2[K]*PRHU2
  end

  StUH2[NH] = OrdUH2[NH] * PRHU2

  # Potential intercatchment semi-exchange

  scale_param = (tstep / 24.0)

  Rr = St[2]/Param[3]
  EXCH = scale_param * Param[2]*Rr*Rr*Rr*sqrt(Rr)

  # Routing store

  AEXCH1 = EXCH

  if (St[2]+StUH1[1]+EXCH) < 0.0
    AEXCH1 = -St[2] - StUH1[1]
  end

  St[2] = St[2] + StUH1[1] + EXCH

  if St[2] < 0.0
    St[2] = 0.0
  end

  scale_param = (24.0 / tstep)
  Rr = St[2]^4 / (Param[3]^4 * scale_param)

  QR = St[2] * (1.0-1.0/sqrt(sqrt(1.0+Rr)))

  St[2] = St[2] - QR

  # Runoff from direct branch QD

  AEXCH2 = EXCH

  if (StUH2[1]+EXCH) < 0.0
    AEXCH2 = -StUH2[1]
  end

  QD = max(0.0,StUH2[1]+EXCH)

  # Total runoff

  Q = QR + QD

  if Q < 0.0
    Q = 0.0
  end

  # Update struct

  mdata.time    += Dates.Hour(tstep)
  mdata.st      = St
  mdata.st_uh1  = StUH1
  mdata.st_uh2  = StUH2
  mdata.ord_uh1 = OrdUH1
  mdata.ord_uh2 = OrdUH2
  mdata.q_sim   = Q

  return nothing

end

function SS1(I,C,D)

  FI = I
  if FI <= 0.0
    SS1 = 0.0
  elseif FI < C
    SS1 = (FI/C)^D
  else
    SS1 = 1.0
  end

end

function SS2(I,C,D)

  FI = I
  if FI <= 0.0
    SS2 = 0.0
  elseif FI <= C
    SS2 = 0.5*(FI/C)^D
  elseif FI < 2.0*C
    SS2 = 1.0-0.5*(2.0-FI/C)^D
  else
    SS2=1.0
  end

end

function UH1(OrdUH1,C,D)

  NH = length(OrdUH1)

  for I in 1:NH

    OrdUH1[I] = SS1(I,C,D)-SS1(I-1,C,D)

  end

end

function UH2(OrdUH2,C,D)

  NH = length(OrdUH2)

  for I in 1:NH

    OrdUH2[I] = SS2(I,C,D)-SS2(I-1,C,D)

  end

end

#**********************************************************************
#     SUBROUTINE MOD_GR4J(St,StUH1,StUH2,OrdUH1,OrdUH2,Param,P1,E,Q)
# Run on a single time step with the GR4J model
# Inputs:
#       St     Vector of model states in stores at the beginning of the time step [mm]
#       StUH1  Vector of model states in Unit Hydrograph 1 at the beginning of the time step [mm]
#       StUH2  Vector of model states in Unit Hydrograph 2 at the beginning of the time step [mm]
#       OrdUH1 Vector of ordinates in UH1 [-]
#       OrdUH2 Vector of ordinates in UH2 [-]
#       Param  Vector of model parameters [various units]
#       P1     Value of rainfall during the time step [mm/day]
#       E      Value of potential evapotranspiration during the time step [mm/day]
# Outputs:
#       St     Vector of model states in stores at the end of the time step [mm]
#       StUH1  Vector of model states in Unit Hydrograph 1 at the end of the time step [mm]
#       StUH2  Vector of model states in Unit Hydrograph 2 at the end of the time step [mm]
#       Q      Value of simulated flow at the catchment outlet for the time step [mm]
#**********************************************************************

#**********************************************************************
#     FUNCTION SS1(I,C,D)
# Values of the S curve (cumulative HU curve) of GR unit hydrograph UH1
# Inputs:
#    C: time constant
#    D: exponent
#    I: time-step
# Outputs:
#    SS1: Values of the S curve for I
#**********************************************************************

#**********************************************************************
#     FUNCTION SS2(I,C,D)
# Values of the S curve (cumulative HU curve) of GR unit hydrograph UH2
# Inputs:
#    C: time constant
#    D: exponent
#    I: time-step
# Outputs:
#    SS2: Values of the S curve for I
#**********************************************************************

#**********************************************************************
#     SUBROUTINE UH1(OrdUH1,C,D)
# Computation of ordinates of GR unit hydrograph UH1 using successive differences on the S curve SS1
# Inputs:
#    C: time constant
#    D: exponent
# Outputs:
#    OrdUH1: NH ordinates of discrete hydrograph
#**********************************************************************

#**********************************************************************
#     SUBROUTINE UH2(OrdUH2,C,D)
# Computation of ordinates of GR unit hydrograph HU2 using successive differences on the S curve SS2
# Inputs:
#    C: time constant
#    D: exponent
# Outputs:
#    OrdUH2: 2*NH ordinates of discrete hydrograph
#**********************************************************************

ext: jl

max_stars_repo_name: jmgnve/Vann

max_forks_repo_name: jmgnve/Vann


text:
### Integrated autocorrelation time (IACT)

iact{N<:Real}(mcvariance::N, iidvariance::N) = mcvariance/iidvariance

iact(v::AbstractArray; vtype::Symbol=:imse, args...) = iact(mcvar(v; vtype=vtype, args...), mcvar_iid(v))

iact(v::AbstractArray, region; vtype::Symbol=:imse, args...) = mapslices(x -> iact(x; vtype=vtype, args...), v, region)

iact(s::VariableNState{Univariate}; vtype::Symbol=:imse, args...) = iact(s.value; vtype=vtype, args...)

iact(s::VariableNState{Multivariate}, i::Integer; vtype::Symbol=:imse, args...) = iact(s.value[i, :]; vtype=vtype, args...)

iact(s::VariableNState{Multivariate}, r::AbstractVector=1:s.size; vtype::Symbol=:imse, args...) =
  eltype(s)[iact(s, i; vtype=vtype, args...) for i in r]


text:
function initialize!(integrator,cache::SplitEulerConstantCache)
  integrator.kshortsize = 2
  integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)
  integrator.fsalfirst =
  integrator.f.f1(integrator.uprev,integrator.p,integrator.t) +
  integrator.f.f2(integrator.uprev,integrator.p,integrator.t) # Pre-start fsal

  # Avoid undefined entries if k is an array of arrays
  integrator.fsallast = zero(integrator.fsalfirst)
  integrator.k[1] = integrator.fsalfirst
  integrator.k[2] = integrator.fsallast
end

@muladd function perform_step!(integrator,cache::SplitEulerConstantCache,repeat_step=false)
  @unpack t,dt,uprev,u,f,p = integrator
  u = @. uprev + dt*integrator.fsalfirst
  integrator.fsallast = f.f1(u,p,t+dt) + f.f2(u,p,t+dt)  # For the interpolation, needs k at the updated point
  integrator.k[1] = integrator.fsalfirst
  integrator.k[2] = integrator.fsallast
  integrator.u = u
end

function initialize!(integrator,cache::SplitEulerCache)
  integrator.kshortsize = 2
  @unpack k,fsalfirst = cache
  integrator.fsalfirst = fsalfirst
  integrator.fsallast = k
  resize!(integrator.k, integrator.kshortsize)
  integrator.k[1] = integrator.fsalfirst
  integrator.k[2] = integrator.fsallast
  integrator.f.f1(integrator.fsalfirst,integrator.uprev,integrator.p,integrator.t) # For the interpolation, needs k at the updated point
  integrator.f.f2(cache.tmp,integrator.uprev,integrator.p,integrator.t) # For the interpolation, needs k at the updated point
  integrator.fsalfirst .+= cache.tmp
end

@muladd function perform_step!(integrator,cache::SplitEulerCache,repeat_step=false)
  @unpack t,dt,uprev,u,f,p = integrator
  @. u = uprev + dt*integrator.fsalfirst
  f.f1(integrator.fsallast,u,p,t+dt) # For the interpolation, needs k at the updated point
  f.f2(cache.tmp,u,p,t+dt) # For the interpolation, needs k at the updated point
  integrator.fsallast .+= cache.tmp
end


text:
using StatsBase
using Base.Test

X = randn(3, 8)

Z1 = X .- mean(X, 1)
Z2 = X .- mean(X, 2)

w1 = rand(3)
w2 = rand(8)

wv1 = weights(w1)
wv2 = weights(w2)

Z1w = X .- mean(X, wv1, 1)
Z2w = X .- mean(X, wv2, 2)

## reference results

S1 = Z1'Z1
S2 = Z2 * Z2'

Sz1 = X'X
Sz2 = X * X'

S1w = Z1w' * diagm(w1) * Z1w
S2w = Z2w * diagm(w2) * Z2w'

Sz1w = X' * diagm(w1) * X
Sz2w = X * diagm(w2) * X'

## scattermat

if VERSION < v"0.5.0-dev+679"
    @test scattermat(X)           ≈ S1
    @test scattermat(X; vardim=2) ≈ S2

    @test scattermat(X; mean=0)           ≈ Sz1
    @test scattermat(X; mean=0, vardim=2) ≈ Sz2

    @test scattermat(X; mean=mean(X,1))           ≈ S1
    @test scattermat(X; mean=mean(X,2), vardim=2) ≈ S2

    @test scattermat(X; mean=zeros(1,8))         ≈ Sz1
    @test scattermat(X; mean=zeros(3), vardim=2) ≈ Sz2

    ## weighted scatter mat

    @test scattermat(X, wv1)           ≈ S1w
    @test scattermat(X, wv2; vardim=2) ≈ S2w

    @test scattermat(X, wv1; mean=0)           ≈ Sz1w
    @test scattermat(X, wv2; mean=0, vardim=2) ≈ Sz2w

    @test scattermat(X, wv1; mean=mean(X, wv1, 1))           ≈ S1w
    @test scattermat(X, wv2; mean=mean(X, wv2, 2), vardim=2) ≈ S2w

    @test scattermat(X, wv1; mean=zeros(1,8))         ≈ Sz1w
    @test scattermat(X, wv2; mean=zeros(3), vardim=2) ≈ Sz2w
else
    @test scattermat(X)    ≈ S1
    @test scattermat(X, 2) ≈ S2

    @test StatsBase.scattermatm(X, 0)    ≈ Sz1
    @test StatsBase.scattermatm(X, 0, 2) ≈ Sz2

    @test StatsBase.scattermatm(X, mean(X,1))    ≈ S1
    @test StatsBase.scattermatm(X, mean(X,2), 2) ≈ S2

    @test StatsBase.scattermatm(X, zeros(1,8))  ≈ Sz1
    @test StatsBase.scattermatm(X, zeros(3), 2) ≈ Sz2

    ## weighted scatter mat

    @test scattermat(X, wv1)    ≈ S1w
    @test scattermat(X, wv2, 2) ≈ S2w

    @test StatsBase.scattermatm(X, 0, wv1)    ≈ Sz1w
    @test StatsBase.scattermatm(X, 0, wv2, 2) ≈ Sz2w

    @test StatsBase.scattermatm(X, mean(X, wv1, 1), wv1)    ≈ S1w
    @test StatsBase.scattermatm(X, mean(X, wv2, 2), wv2, 2) ≈ S2w

    @test StatsBase.scattermatm(X, zeros(1,8), wv1)  ≈ Sz1w
    @test StatsBase.scattermatm(X, zeros(3), wv2, 2) ≈ Sz2w
end

# weighted covariance

if VERSION < v"0.5.0-dev+679"
    @test cov(X, wv1)           ≈ S1w ./ sum(wv1)
    @test cov(X, wv2; vardim=2) ≈ S2w ./ sum(wv2)

    @test cov(X, wv1; mean=0)           ≈ Sz1w ./ sum(wv1)
    @test cov(X, wv2; mean=0, vardim=2) ≈ Sz2w ./ sum(wv2)

    @test cov(X, wv1; mean=mean(X, wv1, 1))           ≈ S1w ./ sum(wv1)
    @test cov(X, wv2; mean=mean(X, wv2, 2), vardim=2) ≈ S2w ./ sum(wv2)

    @test cov(X, wv1; mean=zeros(1,8))         ≈ Sz1w ./ sum(wv1)
    @test cov(X, wv2; mean=zeros(3), vardim=2) ≈ Sz2w ./ sum(wv2)
else
    @test cov(X, wv1)    ≈ S1w ./ sum(wv1)
    @test cov(X, wv2, 2) ≈ S2w ./ sum(wv2)

    @test Base.covm(X, 0, wv1)    ≈ Sz1w ./ sum(wv1)
    @test Base.covm(X, 0, wv2, 2) ≈ Sz2w ./ sum(wv2)

    @test Base.covm(X, mean(X, wv1, 1), wv1)    ≈ S1w ./ sum(wv1)
    @test Base.covm(X, mean(X, wv2, 2), wv2, 2) ≈ S2w ./ sum(wv2)

    @test Base.covm(X, zeros(1,8), wv1)  ≈ Sz1w ./ sum(wv1)
    @test Base.covm(X, zeros(3), wv2, 2) ≈ Sz2w ./ sum(wv2)
end

# mean_and_cov
if VERSION < v"0.5.0-dev+679"
    (m, C) = mean_and_cov(X; vardim=1)
    @test m == mean(X, 1)
    @test C == cov(X; vardim=1)

    (m, C) = mean_and_cov(X; vardim=2)
    @test m == mean(X, 2)
    @test C == cov(X; vardim=2)

    (m, C) = mean_and_cov(X, wv1; vardim=1)
    @test m == mean(X, wv1, 1)
    @test C == cov(X, wv1; vardim=1)

    (m, C) = mean_and_cov(X, wv2; vardim=2)
    @test m == mean(X, wv2, 2)
    @test C == cov(X, wv2; vardim=2)
else
    (m, C) = mean_and_cov(X, 1)
    @test m == mean(X, 1)
    @test C == cov(X, 1)

    (m, C) = mean_and_cov(X, 2)
    @test m == mean(X, 2)
    @test C == cov(X, 2)

    (m, C) = mean_and_cov(X, wv1, 1)
    @test m == mean(X, wv1, 1)
    @test C == cov(X, wv1, 1)

    (m, C) = mean_and_cov(X, wv2, 2)
    @test m == mean(X, wv2, 2)
    @test C == cov(X, wv2, 2)
end


text:
module DataPreprocessing

include("functions.jl")

export CounterfactualData, select_factual, apply_domain_constraints

end


text:
export iauTpxev
"""
In the tangent plane projection, given celestial direction cosines
for a star and the tangent point, solve for the star's rectangular
coordinates in the tangent plane.

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

Status:  support function.

Given:
   v         double[3]  direction cosines of star (Note 4)
   v0        double[3]  direction cosines of tangent point (Note 4)

Returned:
   *xi,*eta  double     tangent plane coordinates of star

Returned (function value):
            int        status: 0 = OK
                                 1 = star too far from axis
                                 2 = antistar on tangent plane
                                 3 = antistar too far from axis

Notes:

   1. The tangent plane projection is also called the "gnomonic
   projection" and the "central projection".

   2. The eta axis points due north in the adopted coordinate system.
   If the direction cosines represent observed (RA,Dec), the tangent
   plane coordinates (xi,eta) are conventionally called the
   "standard coordinates".  If the direction cosines are with
   respect to a right-handed triad, (xi,eta) are also right-handed.
   The units of (xi,eta) are, effectively, radians at the tangent
   point.

   3. The method used is to extend the star vector to the tangent
   plane and then rotate the triad so that (x,y) becomes (xi,eta).
   Writing (a,b) for the celestial spherical coordinates of the
   star, the sequence of rotations is (a+pi/2) around the z-axis
   followed by (pi/2-b) around the x-axis.

   4. If vector v0 is not of unit length, or if vector v is of zero
   length, the results will be wrong.

   5. If v0 points at a pole, the returned (xi,eta) will be based on
   the arbitrary assumption that the longitude coordinate of the
   tangent point is zero.

   6. This function is a member of the following set:

      spherical      vector         solve for

      iauTpxes    > iauTpxev <       xi,eta
      iauTpsts      iauTpstv          star
      iauTpors      iauTporv         origin

References:

   Calabretta M.R. & Greisen, E.W., 2002, "Representations of
   celestial coordinates in FITS", Astron.Astrophys. 395, 1077

   Green, R.M., "Spherical Astronomy", Cambridge University Press,
   1987, Chapter 13.

This revision:   2018 January 2

SOFA release 2018-01-30

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

# int iauTpxev(double v[3], double v0[3], double *xi, double *eta)

function iauTpxev(v::AbstractVector{<:Real}, v0::AbstractVector{<:Real})
   ref_xi = Ref{Float64}(0.0)
   ref_eta = Ref{Float64}(0.0)

   status = ccall((:iauTpxev, libsofa_c), Cint,
       (Ptr{Cdouble}, Ptr{Cdouble},
       Ref{Cdouble}, Ref{Cdouble}),
       convert(Array{Float64, 1}, v),
       convert(Array{Float64, 1}, v0),
       ref_xi, ref_eta)

   return status, ref_xi[], ref_eta[]
end


text:
using Test

# @testset "The LinearAlgebra Test Suite" begin
    include("cholesky.jl")
    include("qr.jl")
    include("eigenselfadjoint.jl")
    include("eigengeneral.jl")
    include("tridiag.jl")
    include("svd.jl")
    include("rectfullpacked.jl")
    include("lapack.jl")
# end


text:
module NestedSamplers

using LinearAlgebra
using Random
using Random: AbstractRNG, GLOBAL_RNG

using AbstractMCMC
using AbstractMCMC: AbstractSampler,
                    AbstractModel,
                    samples,
                    save!!
import AbstractMCMC: step,
                     bundle_samples,
                     mcmcsample
using Distributions: quantile, UnivariateDistribution
using LogExpFunctions: logaddexp, log1mexp
using MCMCChains: Chains
using ProgressLogging
import StatsBase

export Bounds,
       Proposals,
       Models,
       NestedModel,
       Nested

include("model.jl")         # The default model for nested sampling

# load submodules
include("bounds/Bounds.jl")
using .Bounds
include("proposals/Proposals.jl")
using .Proposals

include("staticsampler.jl") # The static nested sampler
include("step.jl")          # The stepping mechanics (extends AbstractMCMC)
include("sample.jl")        # Custom sampling (extends AbstractMCMC)

include("models/Models.jl")
using .Models

end


text:
using Gadfly
using CSV
using DataFrames

data = CSV.read("C:/Users/ztjsw/Desktop/data.txt", delim = ',', header = true)

data[:chromosome] = string.(data[:chromosome])
data[:neg_log_p] = -log10.(data[:pvalue])

#calculate adjusted position
max_pos_chrom = by(data, :chromosome, max_pos = :position => maximum)
max_pos_chrom[:cumsum] =  cumsum(max_pos_chrom[:max_pos])- max_pos_chrom[:max_pos]
data_new=join(data, max_pos_chrom, on = :chromosome, kind = :left)
data_new[:adj_pos] = data_new[:position] + data_new[:cumsum]

## Plot
#calculate postion for xticks
x_chrom_pos = by(data_new, :chromosome, mid_pos = :adj_pos => d -> (maximum(d) + minimum(d))/2)[2]
x_chrom_label = unique(data_new[:chromosome])

function pos_to_chrom(iPos)
    # return the name of its chromosome for a position
    return x_chrom_label[x_chrom_pos .== iPos][1]
end

p=plot(data_new, x=:adj_pos, y=:neg_log_p, color=:chromosome, Geom.point,
        Guide.xlabel("Chromosome"),
        Guide.ylabel("-log10(p)"),
        Guide.xticks(ticks=x_chrom_pos),
        Scale.x_continuous(labels = pos_to_chrom),  #labels is a function which maps a coordinate value to a string label
        Theme(key_position = :none, grid_line_width=0mm, point_size=0.1cm))

display(p)


text:
# Useful for testing indexing
struct ZeroBasedRange{T,A<:AbstractRange{T}} <: AbstractRange{T}
    a :: A
    function ZeroBasedRange(a::AbstractRange{T}) where {T}
        @assert !Base.has_offset_axes(a)
        new{T, typeof(a)}(a)
    end
end

struct ZeroBasedUnitRange{T,A<:AbstractUnitRange{T}} <: AbstractUnitRange{T}
    a :: A
    function ZeroBasedUnitRange(a::AbstractUnitRange{T}) where {T}
        @assert !Base.has_offset_axes(a)
        new{T, typeof(a)}(a)
    end
end

for Z in [:ZeroBasedRange, :ZeroBasedUnitRange]
    @eval Base.parent(A::$Z) = A.a
    @eval Base.first(A::$Z) = first(A.a)
    @eval Base.length(A::$Z) = length(A.a)
    @eval Base.last(A::$Z) = last(A.a)
    @eval Base.size(A::$Z) = size(A.a)
    @eval Base.axes(A::$Z) = map(x -> IdentityUnitRange(0:x-1), size(A.a))
    @eval Base.getindex(A::$Z, i::Int) = A.a[i + 1]
    @eval Base.getindex(A::$Z, i::Integer) = A.a[i + 1]
    @eval Base.firstindex(A::$Z) = 0
    @eval Base.step(A::$Z) = step(A.a)
    @eval OffsetArrays.no_offset_view(A::$Z) = A.a
    @eval function Base.show(io::IO, A::$Z)
        show(io, A.a)
        print(io, " with indices $(axes(A,1))")
    end
end

for Z in [:ZeroBasedRange, :ZeroBasedUnitRange]
    for R in [:AbstractRange, :AbstractUnitRange, :StepRange]
        @eval @inline function Base.getindex(A::$Z, r::$R{<:Integer})
            @boundscheck checkbounds(A, r)
            OffsetArrays._indexedby(A.a[r .+ 1], axes(r))
        end
    end

    for R in [:ZeroBasedUnitRange, :ZeroBasedRange]
        @eval @inline function Base.getindex(A::$Z, r::$R{<:Integer})
            @boundscheck checkbounds(A, r)
            OffsetArrays._indexedby(A.a[r.a .+ 1], axes(r))
        end
    end

    for R in [:IIUR, :IdOffsetRange]
        @eval @inline function Base.getindex(A::$Z, r::$R)
            @boundscheck checkbounds(A, r)
            OffsetArrays._indexedby(A.a[r .+ 1], axes(r))
        end
    end

    for R in [:AbstractUnitRange, :IdOffsetRange, :IdentityUnitRange, :SliceIntUR, :StepRange, :StepRangeLen, :LinRange]
        @eval @inline function Base.getindex(A::$R, r::$Z)
            @boundscheck checkbounds(A, r)
            OffsetArrays._indexedby(A[r.a], axes(r))
        end
    end
    @eval @inline function Base.getindex(A::StepRangeLen{<:Any,<:Base.TwicePrecision,<:Base.TwicePrecision}, r::$Z)
        @boundscheck checkbounds(A, r)
        OffsetArrays._indexedby(A[r.a], axes(r))
    end

    @eval Base.reshape(z::$Z, inds::Tuple{}) = reshape(parent(z), inds)
    @eval Base.reshape(z::$Z, inds::Tuple{Int, Vararg{Int}}) = reshape(parent(z), inds)
    @eval Base.reshape(z::$Z, inds::Tuple{Union{Int, AbstractUnitRange{<:Integer}}, Vararg{Union{Int, AbstractUnitRange{<:Integer}}}}) = reshape(parent(z), inds)
end

# A basic range that does not have specialized vector indexing methods defined
# In this case the best that we may do is to return an OffsetArray
# Despite this, an indexing operation involving this type should preserve the axes of the indices
struct CustomRange{T,A<:AbstractRange{T}} <: AbstractRange{T}
    a :: A
end
Base.parent(r::CustomRange) = r.a
Base.size(r::CustomRange) = size(parent(r))
Base.length(r::CustomRange) = length(parent(r))
Base.axes(r::CustomRange) = axes(parent(r))
Base.first(r::CustomRange) = first(parent(r))
Base.last(r::CustomRange) = last(parent(r))
Base.step(r::CustomRange) = step(parent(r))
Base.getindex(r::CustomRange, i::Int) = getindex(parent(r), i)


text:
module TestNamedNodes
using PatternDispatch

@pattern f((1,2))           = 3
@pattern f((x::Int,2))      = 2
@pattern f((1,y::Int))      = 1
@pattern f((x::Int,y::Int)) = x*y

show_dispatch(f)

end # module


text:
using BlackBoxOptim

function fit_quadrupolar_bb(
    experimental::Array{Float64, 2};
    sites::Int64 = 1,
    max_func_evals::Int64 = 10000,
    method::Symbol = :seperable_nes,
    I::Int64 = 3,
    samples::Int64 = 1_000_000,
    transitions::UnitRange{Int64} = 1:(2*I),
    range::Tuple{Float64,Float64} = (experimental[1, 1], experimental[end, 1]),
)
    experimental_ecdf = get_experimental_ecdf(experimental)
    ν0 =  get_ν0(experimental, experimental_ecdf)
    search_range = [(0.0, 9.0), (0.0, 1.0), (0.0, 1.0), (0.0, 1.0)]

    fitness_progress_history = Array{Tuple{Int, Float64, Array{Float64}},1}()
    callback = oc -> push!(fitness_progress_history,
        (BlackBoxOptim.num_func_evals(oc), best_fitness(oc), best_candidate(oc)))

    res =  bboptimize(x -> SpectraFit.ols_cdf(  # objective function
                Quadrupolar(x),
                experimental,
                experimental_ecdf,
                ν0,
                I = I,
                samples = samples,
                transitions = transitions,
                range = range,
            ), SearchRange = search_range, MaxFuncEvals = max_func_evals,
            CallbackFunction = callback, CallbackInterval = 0.0)

    return res, fitness_progress_history
end

function fit_chemicalshift_bb(
    experimental::Array{Float64, 2};
    sites::Int64 = 1,
    max_func_evals::Int64 = 10000,
    method::Symbol = :xnes,
    samples::Int64 = 1_000_000
)
    experimental_ecdf = get_experimental_ecdf(experimental)
    ν0 =  get_ν0(experimental, experimental_ecdf)
    search_range = repeat([(-4000.0, 4000.0), (0.000001, 800.0), (-4000.0, 4000.0), (0.000001, 400.0), (0.000001, 1.0), (0.000001, 1.0),  (0.000001, 1.0)], sites)

    fitness_progress_history = Array{Tuple{Int, Float64},1}()
    callback = oc -> push!(fitness_progress_history,
        (BlackBoxOptim.num_func_evals(oc), best_fitness(oc)))

    res =  bboptimize(x -> SpectraFit.ols_cdf(  # objective function
                ChemicalShift(x),
                experimental,
                experimental_ecdf,
                samples = samples
            ), SearchRange = search_range, MaxFuncEvals = max_func_evals,
            CallbackFunction = callback, CallbackInterval = 0.0)

    return res, fitness_progress_history
end


text:
using Pkg # hideall
Pkg.activate("_literate/EX-wine/Project.toml")
Pkg.update()
macro OUTPUT()
    return isdefined(Main, :Franklin) ? Franklin.OUT_PATH[] : "/tmp/"
end;

# ## Initial data processing
#
# In this example, we consider the [UCI "wine" dataset](http://archive.ics.uci.edu/ml/datasets/wine)
#
# > These data are the results of a chemical analysis of wines grown in the same region in Italy but derived from three different cultivars. The analysis determined the quantities of 13 constituents found in each of the three types of wines.
#
# ### Getting the data
#
# Let's download the data thanks to the [UrlDownload.jl](https://github.com/Arkoniak/UrlDownload.jl) package and load it into a DataFrame:

using HTTP
using MLJ
using PyPlot
import DataFrames: DataFrame, describe
using UrlDownload
MLJ.color_off() # hide

url = "http://archive.ics.uci.edu/ml/machine-learning-databases/wine/wine.data"
header = ["Class", "Alcool", "Malic acid", "Ash", "Alcalinity of ash",
          "Magnesium", "Total phenols", "Flavanoids",
          "Nonflavanoid phenols", "Proanthcyanins", "Color intensity",
          "Hue", "OD280/OD315 of diluted wines", "Proline"]
data = urldownload(url, true, format=:CSV, header=header);

# The second argument to `urldownload` adds a progress meter for the download,
# the `format` helps indicate the format of the file and the `header` helps
# pass the column names which are not in the file.

df = DataFrame(data)
describe(df)

# the target is the `Class` column, everything else is a feature; we can
# dissociate the two  using the `unpack` function:

y, X = unpack(df, ==(:Class));

# ### Setting the scientific type
#
# Let's explore the scientific type attributed by default to the target and the features

scitype(y)

# this should be changed as it should be considered as an ordered factor. The
# difference is as follows:
#
# * a `Count` corresponds to an integer between 0 and infinity
# * a `OrderedFactor` however is a categorical object (there are finitely many options) with ordering (`1 < 2 < 3`).
#

yc = coerce(y, OrderedFactor);

# Let's now consider the features

scitype(X)

# So there are `Continuous` values (encoded as floating point) and `Count` values (integer).
# Note also that there are no missing value (otherwise one of the scientific type would have been a `Union{Missing,*}`).
# Let's check which column is what:

schema(X)

# The two variable that are encoded as `Count` can  probably be re-interpreted; let's have a look at the `Proline` one to see what it looks like

X[1:5, :Proline]

# It can likely be interpreted as a Continuous as well (it would be better to know precisely what it is but for now let's just go with the hunch).
# We'll do the same with `:Magnesium`:

Xc = coerce(X, :Proline=>Continuous, :Magnesium=>Continuous);

# Finally, let's have a quick look at the mean and standard deviation of each feature to get a feel for their amplitude:

describe(Xc, :mean, :std)

# Right so it varies a fair bit which would invite to standardise the data.
#
# **Note**: to complete such a first step, one could explore histograms of the various features for instance, check that there is enough variation among the continuous features and that there does not seem to be problems in the encoding, we cut this out to shorten the tutorial. We could also have checked that the data is balanced.
#
# ## Getting a baseline
#
# It's a multiclass classification problem with continuous inputs so a sensible start is  to test two very simple classifiers to get a baseline.
# We'll train two simple pipelines:
# - a Standardizer + KNN classifier and
# - a Standardizer + Multinomial classifier (logistic regression).

KNNC = @load KNNClassifier
MNC = @load MultinomialClassifier pkg=MLJLinearModels;

KnnPipe = Standardizer |> KNNC
MnPipe = Standardizer |> MNC

# Note the `|>` syntax, which is syntactic sugar for creating a linear `Pipeline` from components models.

# We can now fit this on a train split of the data setting aside 20% of the data for eventual testing.

train, test = partition(collect(eachindex(yc)), 0.8, shuffle=true, rng=111)
Xtrain = selectrows(Xc, train)
Xtest = selectrows(Xc, test)
ytrain = selectrows(yc, train)
ytest = selectrows(yc, test);

# Let's now wrap an instance of these models with data (all hyperparameters are set to default here):

knn = machine(KnnPipe, Xtrain, ytrain)
multi = machine(MnPipe, Xtrain, ytrain)

# Let's train a KNNClassifier with default hyperparameters and get a baseline misclassification rate using 90% of the training data to train the model and the remaining 10% to evaluate it:

opts = (resampling=Holdout(fraction_train=0.9), measure=cross_entropy)
res = evaluate!(knn; opts...)
round(res.measurement[1], sigdigits=3)

# Now we do the same with a MultinomialClassifier

res = evaluate!(multi; opts...)
round(res.measurement[1], sigdigits=3)

# Both methods seem to offer comparable levels of performance.
# Let's check the misclassification over the full training set:

mcr_k = misclassification_rate(predict_mode(knn, Xtrain), ytrain)
mcr_m = misclassification_rate(predict_mode(multi, Xtrain), ytrain)
println(rpad("KNN mcr:", 10), round(mcr_k, sigdigits=3))
println(rpad("MNC mcr:", 10), round(mcr_m, sigdigits=3))

# So here we have done no hyperparameter training and already have a misclassification rate below 5%.
# Clearly the problem is not very difficult.
#
# ## Visualising the classes
#
# One way to get intuition for why the dataset is so easy to classify is to project it onto a 2D space using the PCA and display the two classes to see if they are well separated; we use the arrow-syntax here (if you're on Julia <= 1.2, use the commented-out lines as you won't be able to use the arrow-syntax)

PCA = @load PCA
pca_pipe = Standardizer() |> PCA(maxoutdim=2)
pca = machine(pca_pipe, Xtrain)
fit!(pca)
W = transform(pca, Xtrain)

# Let's now display this using different colours for the different classes:

x1 = W.x1
x2 = W.x2

mask_1 = ytrain .== 1
mask_2 = ytrain .== 2
mask_3 = ytrain .== 3

figure(figsize=(8, 6))
plot(x1[mask_1], x2[mask_1], linestyle="none", marker="o", color="red")
plot(x1[mask_2], x2[mask_2], linestyle="none", marker="o", color="blue")
plot(x1[mask_3], x2[mask_3], linestyle="none", marker="o", color="magenta")

xlabel("PCA dimension 1", fontsize=14)
ylabel("PCA dimension 2", fontsize=14)
legend(["Class 1", "Class 2", "Class 3"], fontsize=12)
xticks(fontsize=12)
yticks(fontsize=12)

savefig(joinpath(@OUTPUT, "EX-wine-pca.svg")) # hide

# \figalt{PCA}{EX-wine-pca.svg}
#
# On that figure it now becomes quite clear why we managed to achieve such high scores with very simple classifiers.
# At this point it's a bit pointless to dig much deaper into parameter tuning etc.
#
# As a last step, we can report performances of the models on the test set which we set aside earlier:

perf_k = misclassification_rate(predict_mode(knn, Xtest), ytest)
perf_m = misclassification_rate(predict_mode(multi, Xtest), ytest)
println(rpad("KNN mcr:", 10), round(perf_k, sigdigits=3))
println(rpad("MNC mcr:", 10), round(perf_m, sigdigits=3))

# Pretty good for so little work!
PyPlot.close_figs() # hide

max_issues_repo_name: alan-turing-institute/MLJTutorials

max_forks_repo_name: alan-turing-institute/MLJTutorials


text:
function ParameterDriversList()
    return ParameterDriversList(())
end

function add(obj::ParameterDriversList, arg0::ParameterDriver)
    return jcall(obj, "add", void, (ParameterDriver,), arg0)
end

function filter(obj::ParameterDriversList, arg0::jboolean)
    return jcall(obj, "filter", void, (jboolean,), arg0)
end

function find_by_name(obj::ParameterDriversList, arg0::JString)
    return jcall(obj, "findByName", ParameterDriversList_DelegatingDriver, (JString,), arg0)
end

function get_drivers(obj::ParameterDriversList)
    return jcall(obj, "getDrivers", List, ())
end

function get_nb_params(obj::ParameterDriversList)
    return jcall(obj, "getNbParams", jint, ())
end

function sort(obj::ParameterDriversList)
    return jcall(obj, "sort", void, ())
end


text:
module SpatialGraphs

using Graphs

import Graphs: vertex_index

export # Types
       Geometry,
       Point,
       SimpleSpatialVertex,
       SpatialVertex,
       # Graphs.jl interfaces
       vertex_index,
       # Spatial interfaces
       geometry,
       # Visualisation
       plot

abstract SpatialVertex
abstract Geometry{T<:Real}

vertex_index(v::SpatialVertex) = v.index
geometry(v::SpatialVertex) = v.geom

immutable Point{T} <: Geometry{T}
  x::T
  y::T
end

(+)(a::Point, b::Point) = Point(a.x + b.x, a.y + b.y)
(-)(a::Point, b::Point) = Point(a.x - b.x, a.y - b.y)
(*)(p::Point, s::Real) = Point(p.x*s, p.y*s)
(*)(s::Real, p::Point) = p*s
(==)(a::Point, b::Point) = ==(a.x, b.x) && ==(a.y, b.y)

type SimpleSpatialVertex{T} <: SpatialVertex
  index::Int
  geom::Geometry{T}
end
SimpleSpatialVertex{T}(id::Int, x::T, y::T) =
  SimpleSpatialVertex(id, Point(x, y))

end # module


text:
include("single_decision.jl")
include("vector_decision.jl")
include("affine.jl")
include("quadratic.jl")
include("vector_affine.jl")


text:
## List files to version control manually

# Raw files
# Normally these should not be changed
_raw_files = [
    "data/raw/ebd_sample.txt",
    "data/raw/ebd_warblers_cut.csv",
    "data/raw/ebd_warblers_head.csv",
    "data/raw/ebd_warblers.csv",
]

# Processed files
# The corresponding placeholder should also be created anywhere these files are written.
_jld_files = [
    # "data/jld2/bart-distributions_xtras.jld2",
    # "data/jld2/bart-distributions.jld2",
    # "data/jld2/bio-distributions.jld2",
    "data/jld2/comparison-results.jld2",
    # "data/jld2/raw-distributions.jld2",
    # "data/jld2/rf-distributions.jld2",
]
_proc_files = [
    # "data/proc/bart_predictions_lower.csv",
    # "data/proc/bart_predictions_pres.csv",
    # "data/proc/bart_predictions_prob.csv",
    # "data/proc/bart_predictions_upper.csv",
    # "data/proc/bart_summaries.csv",
    # "data/proc/bart_varimps.csv",
    "data/proc/ebd_warblers_prep.csv"
]
_raster_files = [
    "data/raster/bart_xtras_lower-distrib.tif",
    "data/raster/bart_xtras_prob-distrib.tif",
    "data/raster/bart_xtras_upper-distrib.tif",
]

## Create placeholders
_placeholder_files = vcat(_jld_files, _proc_files, _raster_files, _raw_files)
_placeholder_files = replace.(_placeholder_files, "." => "_placeholder.")

function _create_placeholder_files!(files)
    for f in files
        touch(f)
        open(f, "w") do io
            write(io, string(Dates.now()))
        end;
    end
end
# _create_placeholder_files!(_placeholder_files)

## Create function to manually version control files (separately for raw & proc)
"""
    verify_raw_files(; touch_placeholders=false)
    verify_proc_files(; touch_placeholders=false)

Verifies if the data files are more recent than their corresponding placeholder
files. Placeholder files were created for important data files which are too big
to be version controlled.

This function will throw a warning when the data files are more recent than
their placeholder files, indicating a possibly unwanted and dangerous change for
which verifications should be performed.

Using `touch_placeholders=true` will update the placeholders' timestampsso they
are more recent than the data files. The placeholders should then be committed
with a description of the changes made to the data files (or discarded if it's
confirmed not important change was made).
"""
function verify_raw_files(; touch_placeholders=false)
    files = _raw_files
    placeholder_files = filter(startswith("data/raw/"), _placeholder_files)
    more_recent = mtime.(files) .> mtime.(placeholder_files)
    if any(more_recent)
        @warn """\n
        The following raw data files are more recent than their corresponding placeholders.
        Raw data files should not normally change, so make sure these changes were made on
        purpose. These files are not version controlled.

            $(files[more_recent])

        Run the following command after verification to update the placeholder files
        timestamps and remove this warning.

            verify_raw_files(; touch_placeholders=true)
        """
        if touch_placeholders
            @info "Re-writing placeholder files to update timestaps and disable the warning"
            for p in placeholder_files[more_recent]
                open(p, "w") do io
                    write(io, string(Dates.now()))
                end;
            end
        end
    end
end
# verify_raw_files()

"""
See `verify_raw_files` for details.
"""
function verify_proc_files(; touch_placeholders=false)
    files = vcat(_proc_files, _jld_files, _raster_files)
    placeholder_files = filter(!startswith("data/raw/"), _placeholder_files)
    more_recent = mtime.(files) .> mtime.(placeholder_files)
    if any(more_recent)
        @warn """\n
        The following processed data files are more recent than their corresponding placeholders.
        Processed data files are expected to changed in the analyses, but make sure the
        changes were made on purpose. These files are not version controlled.

            $(files[more_recent])

        Run the following command after verification to update the placeholder files
        timestamps and remove this warning.

            verify_proc_files(; touch_placeholders=true)
        """
        if touch_placeholders
            @info "Re-writing placeholder files to update timestaps and disable the warning"
            for p in placeholder_files[more_recent]
                open(p, "w") do io
                    write(io, string(Dates.now()))
                end;
            end
        end
    end
end
# verify_proc_files()

max_stars_repo_name: gabrieldansereau/betadiversity-hotspots

max_issues_repo_name: gabrieldansereau/betadiversity-hotspots

max_forks_repo_name: gabrieldansereau/betadiversity-hotspots


text:
# 2D Helmoltz discretizations
# Curt Da Silva, 2016
#

export helm2d_7pt, helm2d_chen2013, param_to_wavenum

function helm2d_7pt(n,d,npml,freq,v,f0,unit::String)
# Standard 7pt discretization of
# d/dz( ex/ez du/dz ) + d/dx( ez/ex du/dx ) + ezexk^2u
#
# Usage:
#  [H,dH,ddH] = helm2d_7pt(n,d,npml,freq,v,f0)
#
# Input:
#   n     - [nz,nx] number of points in each direction (including PML)
#   d     - [dz,dx] grid spacing in each direciton (in meters)
#   npml  - [npz_bot,npz_top] - number of pml points on each edge of the domain
#           [npx_bot,npx_top]
#   freq  - frequency in Hz
#   v     - velocity (with pml included), of size [nz,nx]
#   f0    - peak frequency of source wavelet in Hz
#
# Output:
#   H     - Helmholtz matrix
#   dH    - derivative of Helmholtz matrix
#   ddH   - second derivative of Helmholtz matrix

    nz,nx = n
    Δz,Δx = d
    npz = npml[1,:]
    npx = npml[2,:]
    a0 = 1.79

    (k,dk,ddk) = param_to_wavenum(v,freq,unit)
    ex = pml_func1d(nx,npx,a0,f0,freq)
    ez = pml_func1d(nz,npz,a0,f0,freq)

    xg = 1:nx
    zg = 1:nz

    ez_up = ez(zg.+0.5)
    ez_down = ez(zg.-0.5)
    dz = 1/Δz^2*spdiagm(-1 => 1 ./ ez_down[2:end], 0 => -(1 ./ ez_down+1 ./ ez_up), 1 => 1 ./ ez_up[1:end-1])
    Dz = kron(spdiagm(ex(xg)),dz)

    ex_up = ex(xg.+0.5)
    ex_down = ex(xg.-0.5)
    dx = 1/Δx^2*spdiagm(-1 => 1 ./ ex_down[2:end], 0 => -(1 ./ ex_down+1 ./ ex_up), 1 => 1 ./ ex_up[1:end-1])
    Dx = kron(dx,spdiagm(ez(zg)))
    A = kron(spdiagm(ex(xg)),spdiagm(ez(zg)))

    H = Dz + Dx + A*spdiagm(vec(k))

    dH = A*spdiagm(vec(dk))
    ddH = A*spdiagm(vec(ddk))

  return H,dH,ddH
end

function helm2d_chen2013(n,d,npml,freq,v,f0,unit::String)
  # 9pt stencil from  Chen, et. al. "AN OPTIMAL 9-POINT FINITE DIFFERENCE
  # SCHEME FOR THE HELMHOLTZ EQUATION WITH PML", 2013
  #

    nz,nx = n
    N = nz*nx
    Δz,Δx = d
    Δz==Δx || error("z and x spacing must be identical")
    f0 > 0 || error("f0 must be positive")
    h = Δz
    npz = npml[1,:]
    npx = npml[2,:]
    a0 = 1.79

    (k,dk,ddk) = param_to_wavenum(v,freq,unit)

    vmin = minimum(vec(v))
    vmax = maximum(vec(v))
    Gmin = vmin/(h*freq)
    Gmax = vmax/(h*freq)
    IG = Gmin
    if IG >= 2.5 && IG <= 3
        b,d,e = (0.6803,0.444,0.0008)
    elseif IG >= 3 && IG <= 4
        b,d,e = (0.73427,0.4088,-0.0036)
    elseif IG >= 4 && IG <= 5
        b,d,e = (0.7840,0.3832,-0.0060)
    elseif IG >= 5 && IG <= 6
        b,d,e = (0.802,0.3712,-0.0072)
    elseif IG >= 6 && IG <= 8
        b,d,e = (0.8133,0.3637,-0.0075)
    elseif IG >= 8 && IG <= 10
        b,d,e = (0.8219,0.3578,-0.0078)
    elseif IG >= 10 && IG <= 400
        b,d,e = (0.8271,0.3540,-0.0080)
    end
    c = 1-d-e

    ex = pml_func1d(nx,npx,a0,f0,freq)
    ez = pml_func1d(nz,npz,a0,f0,freq)
    xg = 1:nx
    zg = 1:nz
    ez_up = ez(zg.+0.5)
    ez_down = ez(zg.-0.5)
    ex_up = ex(xg.+0.5)
    ex_down = ex(xg.-0.5)
    lx = 1/Δx^2*spdiagm(-1 => 1 ./ ex_down[2:end], 0 => -(1 ./ ex_down+1 ./ ex_up), 1 => 1 ./ ex_up[1:end-1])
    Lx = j->kron(lx,spdiagm(j => ez(zg.+j)[offset_range(j,nz)]))
    lz = 1/Δz^2*spdiagm(-1 => 1 ./ ez_down[2:end], 0 => -(1 ./ ez_down.+1 ./ ez_up), 1 => 1 ./ ez_up[1:end-1])
    Lz = j->kron(spdiagm(j => ex(xg.+j)[offset_range(j,nx)]),lz)

    L = b*Lx(0) + (1-b)/2*(Lx(-1) + Lx(1)) + b*Lz(0) + (1-b)/2*(Lz(-1) + Lz(1))

    # Neighbourhood notation around the point uNN
    #
    #    uMM  -- uMN -- uMP
    #             |
    #    uNM  -- uNN -- uNP
    #             |
    #    uPM  -- uPN -- uPP
    #
    #

    # Average along the points uMN, uPN,uNM,uNP
    # Eliminate points that don't have these offsets as neighbours
    eNM = ones(nz,nx)
    eNM[:,1] .= 0
    eNP = ones(nz,nx)
    eNP[:,end] .= 0
    eMN = ones(nz,nx)
    eMN[1,:] .= 0
    ePN = ones(nz,nx)
    ePN[end,:] .= 0
    I0 = 1/4 .* spdiagm(-nz => vec(eNM)[nz+1:N], nz => vec(eNP)[1:N-nz], -1 => vec(eMN)[2:N], 1 => vec(ePN)[1:N-1])

    # Average along the points uMM,uMP,uPM,uPP (corners of square above)
    # Eliminate points that don't have these offsets as neighbours
    eMM = ones(nz,nx)
    eMM[1,:] .= 0
    eMM[:,1] .= 0
    ePP = ones(nz,nx)
    ePP[end,:] .= 0
    ePP[:,end] .= 0
    ePM = ones(nz,nx)
    ePM[end,:] .= 0
    ePM[:,1] .= 0
    eMP = ones(nz,nx)
    eMP[1,:] .= 0
    eMP[:,end] .= 0

    I45 = 1/4 .* spdiagm(-nz-1 => vec(eMM)[(nz+2):N], -nz+1 => vec(ePM)[nz:N], nz-1 => vec(eMP)[1:N-nz+1], nz+1 => vec(ePP)[1:N-(nz+1)])

    A = (c*spdiagm(ones(N))+d*I0+e*I45)*kron(spdiagm(ex(xg)), spdiagm(ez(zg)))
    H = L + A*spdiagm(vec(k))

    dH = A*spdiagm(vec(dk))
    ddH = A*spdiagm(vec(ddk))

    return H,dH,ddH
end

function offset_range(j,nx)
# offset_range - Generates the range of indices to compensate for the differences
# between Matlab's spdiags and Julia's spdiagm
# j - offset in (-nx+1,nx-1)
# nx - number of points
  if j >= 0
    I = 1:nx-j
  else
    I = 1+abs(j):nx
  end
end

function pml_func1d(nx,nb,a0,f0,freq)
# Function that increases quadratically from the interior of the pml domain
# σ on p. 391 of Chen, et. al. "AN OPTIMAL 9-POINT FINITE DIFFERENCE
# SCHEME FOR THE HELMHOLTZ EQUATION WITH PML", 2013
#
  dist_from_int = x -> (x.<= nb[1]) .* (nb[1].-x) + (x .> nx.-nb[2]) .* (x.-(nx-nb[2]))
  σ = x-> 2*pi*a0*f0*(dist_from_int(x)/maximum(nb)).^2
  func = x -> 1 .- im*σ(x)/freq
end

function param_to_wavenum(v,freq,unit::String)
# Convert input parameter to rad^2*s2/m2
#
    ω2 = (2*pi*freq)^2
    if unit == "m/s"
        f = ω2*(v.^(-2))
        df = -2*ω2*(v.^(-3))
        ddf = 6*ω2*(v.^(-4))
    elseif unit == "s/m"
        f = ω2*(v.^2)
        df = 2*ω2*v
        ddf = 2*ω2*ones(size(v))
    elseif unit == "s2/m2"
        f = ω2*v
        df = zeros(size(v))
        ddf = zeros(size(v))
    elseif unit == "s2/km2"
        f = ω2*1e-6*v
        df = zeros(size(v))
        ddf = zeros(size(v))
    else
        error("Unknown unit type $unit")
    end
    return f,df,ddf
end


text:
using Test

@time begin
    @test include("jointmodels.jl")
    @test include("utils.jl")
    @test include("model.jl")
    @test include("correctness_ate.jl")
    @test include("correctness_iate.jl")
    @test include("report.jl")
end

max_stars_repo_name: olivierlabayle/TMLE

max_issues_repo_name: olivierlabayle/TMLE

max_forks_repo_name: olivierlabayle/TMLE


text:
## connect to Optim.jl as backend
function sem_wrap_optim(par, F, G, H, model::AbstractSem)
    if !isnothing(F)
        if !isnothing(G)
            if !isnothing(H)
                return objective_gradient_hessian!(G, H, model, par)
            else
                return objective_gradient!(G, model, par)
            end
        else
            if !isnothing(H)
                return objective_hessian!(H, model, par)
            else
                return objective!(model, par)
            end
        end
    else
        if !isnothing(G)
            if !isnothing(H)
                gradient_hessian!(G, H, model, par)
            else
                gradient!(G, model, par)
            end
        end
    end
    return nothing
end

function SemFit(
        optimization_result::Optim.MultivariateOptimizationResults,
        model::AbstractSem,
        start_val)
    return SemFit(
        optimization_result.minimum,
        optimization_result.minimizer,
        start_val,
        model,
        optimization_result
    )
end

optimizer(res::Optim.MultivariateOptimizationResults) = Optim.summary(res)
n_iterations(res::Optim.MultivariateOptimizationResults) = Optim.iterations(res)
convergence(res::Optim.MultivariateOptimizationResults) = Optim.converged(res)

function sem_fit(
        model::AbstractSemSingle{O, I, L, D};
        start_val = start_val,
        kwargs...) where {O, I, L, D <: SemOptimizerOptim}

    if !isa(start_val, Vector)
        start_val = start_val(model; kwargs...)
    end

    result = Optim.optimize(
                Optim.only_fgh!((F, G, H, par) -> sem_wrap_optim(par, F, G, H, model)),
                start_val,
                model.optimizer.algorithm,
                model.optimizer.options)
    return SemFit(result, model, start_val)

end

function sem_fit(
        model::SemEnsemble{N, T , V, D, S};
        start_val = start_val,
        kwargs...) where {N, T, V, D <: SemOptimizerOptim, S}

    if !isa(start_val, Vector)
        start_val = start_val(model; kwargs...)
    end

    result = Optim.optimize(
                Optim.only_fgh!((F, G, H, par) -> sem_wrap_optim(par, F, G, H, model)),
                start_val,
                model.optimizer.algorithm,
                model.optimizer.options)
    return SemFit(result, model, start_val)

end


text:
"""
    sample(chn::Chains, [wv::AbstractWeights,] n; replace=true, kwargs...)
    sample(rng::Random.AbstractRNG, chn::Chains, [wv::AbstractWeights,] n; kwargs...)

Sample `n` samples from `chn`; see also [`subset`](@ref).
Optionally, the samples can be weighted using `wv`.
Here, `kwargs` defaults to `replace=true` and `ordered=false`.
"""
sample(chn::Chains, n; kwargs...) = sample(Random.GLOBAL_RNG, chn, n; kwargs...)
function sample(chn::Chains, wv::AbstractWeights, n)
    return sample(Random.GLOBAL_RNG, chn, wv, n)
end

function sample(
    rng::Random.AbstractRNG,
    chn::Chains,
    n;
    replace = true,
    ordered = false
)
    return subset(chn, sample(rng, range(chn), n; replace = replace, ordered = ordered))
end

function sample(
    rng::Random.AbstractRNG,
    chn::Chains,
    wv::AbstractWeights,
    n
)
    return subset(chn, sample(rng, range(chn), wv, n))
end

"""
    subset(chn::Chains, samples)

Return a subset of samples.
"""
function subset(chn::Chains, samples)
    data = AxisArray(chn.value[samples, :, :].data;
                     iter = 1:length(samples), var = names(chn), chain = chains(chn))

    return Chains(data, missing, chn.name_map, chn.info)
end


text:
using GeometricBase.Config
using GeometricBase.Utils
using GeometricIntegrators.Integrators
using GeometricIntegrators.Tableaus
using GeometricProblems.HarmonicOscillator
using SimpleSolvers
using Test

SimpleSolvers.set_config(:nls_stol_break, 1E3)

using GeometricProblems.HarmonicOscillator: Δt, nt, refx

sode = harmonic_oscillator_sode()

sint = Integrator(sode, TableauLieA(), Δt)
ssol = integrate(sode, sint, nt)
@test relative_maximum_error(ssol.q, refx) < 5E-2

sintc = IntegratorComposition(sode, TableauLieA(), Δt)
ssolc = integrate(sode, sintc, nt)
@test ssol.q == ssolc.q

sint = Integrator(sode, TableauLieB(), Δt)
ssol = integrate(sode, sint, nt)
@test relative_maximum_error(ssol.q, refx) < 5E-2

sintc = IntegratorComposition(sode, TableauLieB(), Δt)
ssolc = integrate(sode, sintc, nt)
@test ssol.q == ssolc.q

sint = Integrator(sode, TableauStrang(), Δt)
ssol = integrate(sode, sint, nt)
@test relative_maximum_error(ssol.q, refx) < 1E-3

sintc = IntegratorComposition(sode, TableauStrang(), Δt)
ssolc = integrate(sode, sintc, nt)
@test ssol.q == ssolc.q

sint = Integrator(sode, TableauStrangA(), Δt)
ssol = integrate(sode, sint, nt)
@test relative_maximum_error(ssol.q, refx) < 1E-3

sintc = IntegratorComposition(sode, TableauStrangA(), Δt)
ssolc = integrate(sode, sintc, nt)
@test ssol.q == ssolc.q

sint = Integrator(sode, TableauStrangB(), Δt)
ssol = integrate(sode, sint, nt)
@test relative_maximum_error(ssol.q, refx) < 1E-3

sintc = IntegratorComposition(sode, TableauStrangB(), Δt)
ssolc = integrate(sode, sintc, nt)
@test ssol.q == ssolc.q

sint1 = Integrator(sode, TableauStrang(), Δt)
ssol1 = integrate(sode, sint, nt)
sint2 = Integrator(sode, TableauStrangA(), Δt)
ssol2 = integrate(sode, sint, nt)
sint3 = Integrator(sode, TableauStrangB(), Δt)
ssol3 = integrate(sode, sint, nt)
@test ssol1.q == ssol2.q
@test ssol1.q ≈  ssol3.q

sint = Integrator(sode, TableauMcLachlan2(), Δt)
ssol = integrate(sode, sint, nt)
@test relative_maximum_error(ssol.q, refx) < 1E-4

sintc = IntegratorComposition(sode, TableauMcLachlan2(), Δt)
ssolc = integrate(sode, sintc, nt)
@test ssol.q == ssolc.q

sint = Integrator(sode, TableauMcLachlan4(), Δt)
ssol = integrate(sode, sint, nt)
@test relative_maximum_error(ssol.q, refx) < 5E-4

sintc = IntegratorComposition(sode, TableauMcLachlan4(), Δt)
ssolc = integrate(sode, sintc, nt)
@test ssol.q == ssolc.q

sint = Integrator(sode, TableauTripleJump(), Δt)
ssol = integrate(sode, sint, nt)
@test relative_maximum_error(ssol.q, refx) < 5E-6

sintc = IntegratorComposition(sode, TableauTripleJump(), Δt)
ssolc = integrate(sode, sintc, nt)
@test ssol.q == ssolc.q

sint = Integrator(sode, TableauSuzukiFractal(), Δt)
ssol = integrate(sode, sint, nt)
@test relative_maximum_error(ssol.q, refx) < 5E-7

sintc = IntegratorComposition(sode, TableauSuzukiFractal(), Δt)
ssolc = integrate(sode, sintc, nt)
@test ssol.q == ssolc.q

DT = eltype(sode)
D  = ndims(sode)

ints_glrk1 = (IntegratorConstructor(DT, D, TableauGauss(1)), IntegratorConstructor(DT, D, TableauGauss(1)))
ints_erk4  = (IntegratorConstructor(DT, D, TableauRK4()), IntegratorConstructor(DT, D, TableauRK4()))

sint = IntegratorComposition(sode, ints_erk4, TableauLieA(), Δt)
ssol = integrate(sode, sint, nt)
# println(relative_maximum_error(ssol.q, refx))
@test relative_maximum_error(ssol.q, refx) < 5E-2

sint = IntegratorComposition(sode, ints_glrk1, TableauStrang(), Δt)
ssol = integrate(sode, sint, nt)
# println(relative_maximum_error(ssol.q, refx))
@test relative_maximum_error(ssol.q, refx) < 1E-3


text:
############################################################
# joAbstractLinearOperator #################################
############################################################

# includes
include("joAbstractLinearOperator/constructors.jl")
include("joAbstractLinearOperator/base_functions.jl")
include("joAbstractLinearOperator/extra_functions.jl")


text:
######################## common traits for the plots

max_width = 493.228346

set_theme!(
    font = "Helvetica", #
    fontsize = 10,
    resolution = (max_width, 500.0),
    linewidth = 0.3,
    strokewidth = 1px,
    markersize = 3px,
    rowgap = Fixed(10),
    colgap = Fixed(10),
    LLegend = (markersize = 10px, markerstrokewidth = 1, patchsize = (10, 10), rowgap = Fixed(2), titlegap = Fixed(5), groupgap = Fixed(10), titlehalign = :left, gridshalign = :left, framecolor = :transparent, padding = 0, linewidth = 0.3),
    LAxis = (xticklabelsize = 8, yticklabelsize = 8, xlabel = "X (cm)", ylabel = "Y (cm)", autolimitaspect = 1, xtickalign = 1, xticksize = 3, ytickalign = 1, yticksize = 3, xticklabelpad = 4)
)
markers = Dict("turning_point" => '•', "center_of_search" => '■')
brighten(c, p = 0.5) = weighted_color_mean(p, c, colorant"white")
mydecompose(origin, radii) = [origin + radii .* Iterators.reverse(sincos(t)) for t in range(0, stop = 2π, length = 51)]
mydecompose(x) = mydecompose(x.origin, x.radii)
legendmarkers = OrderedDict(
                            "track" => (linestyle = nothing, linewidth = 0.3, color = :black),
                            "burrow" => (color = :black, marker = '⋆', strokecolor = :black, markerstrokewidth = 0.5, strokewidth = 0.5, markersize = 15px),
                            "fictive burrow" => (color = :white, marker = '⋆', strokecolor = :black, markerstrokewidth = 0.5, strokewidth = 0.5, markersize = 15px),
                            "turning point" => (color = :black, marker = markers["turning_point"], strokecolor = :transparent, markersize = 15px),
                            "center of search" => (color = :black, marker = markers["center_of_search"], strokecolor = :transparent, markersize = 5px),
                            "mean ± FWHM" => [(color = brighten(colorant"black", 0.75), strokecolor = :transparent, polypoints = mydecompose(Point2f0(0.5, 0.5), Vec2f0(0.75, 0.5))),
                                           (color = :white, marker = '+', strokecolor = :transparent, markersize = 10px),
                                          ])
function getellipse(xy)
    n = length(xy)
    X = Array{Float64}(undef, 2, n)
    for i in 1:n
        X[:,i] = xy[i]
    end
    dis = fit(DiagNormal, X)
    radii = sqrt(2log(2))*sqrt.(var(dis)) # half the FWHM
    (origin = Point2f0(mean(dis)), radii = Vec2f0(radii))
end
function distance2nest(track)
    length(searching(track)) < 10 && return Inf
    t = homing(track)
    i = findfirst(>(0) ∘ last, t)
    isnothing(i) ? Inf : abs(first(t[i]))
end
apply_element(xs) = apply_element.(xs)
apply_element(x::NamedTuple) =  :marker ∈ keys(x) ? MarkerElement(; x...) :
                                :linestyle ∈ keys(x) ? LineElement(; x...) :
                                PolyElement(; x...)

label!(scene, ax, letter) = LText(scene, letter, fontsize = 12, padding = (10, 0, 0, 10), halign = :left, valign = :top, bbox = lift(FRect2D, ax.scene.px_area), font ="Noto Sans Bold")
function plottracks!(ax, g::GroupedDataFrame)
    for gg in g
        for r in eachrow(gg)
            lines!(ax, r.track.coords; legendmarkers["track"]..., color = r.color)
        end
    end
end
function plottracks!(ax, g::DataFrame)
    for r in eachrow(g)
        lines!(ax, r.track.coords; legendmarkers["track"]..., color = r.color)
    end
end
function plotpoints!(ax, g, point_type)
    if !ismissing(g[1].nest[1])
        scatter!(ax, [zero(Point2f0)]; legendmarkers["burrow"]...)
    end
    for (k, gg) in pairs(g)
        xy = gg[!, point_type]
        ellipse = getellipse(xy)
        c = gg.groupcolor[1]
        poly!(ax, mydecompose(ellipse), color = RGBA(brighten(c, 0.5), 0.5))
        scatter!(ax, [ellipse.origin]; legendmarkers["mean ± FWHM"][2]...)
        scatter!(ax, xy; legendmarkers[replace(point_type, "_" => " ")]..., color = RGBA(c, 0.75))
        if k.group ≠ "none"
            scatter!(ax, [Point2f0(intended(k.group))]; legendmarkers["fictive burrow"]..., strokecolor = gg.groupcolor[1])
        end
    end
end

######################## Figure 5 ################

function figure5(df)
    gdf = groupby(df, [:group, :nest_coverage])
    g = gdf[[(group = group, nest_coverage = "closed") for group in ("right","left","towards", "away")]]
    polys = OrderedDict(string(k.group) => (color = v.groupcolor[1], strokecolor = :transparent) for (k, v) in pairs(g))
    scene, layout = layoutscene(0, resolution = (max_width, 600.0))
    ax = layout[1,1] = LAxis(scene)
    plottracks!(ax, g)
    hidexdecorations!(ax, grid = false, ticklabels = false, ticks = false)
    hideydecorations!(ax, grid = false, ticklabels = false, ticks = false)
    label!(scene, ax, "a")
    axs = [LAxis(scene) for _ in 1:2]
    for (i, point_type) in enumerate(("turning_point", "center_of_search"))
        plotpoints!(axs[i], g, point_type)
    end
    layout[2,1:2] = axs
    hideydecorations!(axs[1], grid = false, ticklabels = false, ticks = false)
    hideydecorations!(axs[2], grid = false)
    hidexdecorations!.(axs, grid = false, ticklabels = false, ticks = false)
    layout[3, 1:2] = LText(scene, "X (cm)");
    linkaxes!(axs...)
    label!(scene, axs[1], "c")
    label!(scene, axs[2], "d")
    g = DataFrame(gdf[(group = "towards", nest_coverage = "open")], copycols = false)
    sort!(g, :track, by = distance2nest)
    d = g[1:3,:]
    c = d.groupcolor[1]
    ax = layout[1,2] = LAxis(scene)
    scatter!(ax, [zero(Point2f0)]; legendmarkers["burrow"]...)
    plottracks!(ax, d)
    scatter!(ax, turningpoint.(d.track); legendmarkers["turning point"]..., color = RGBA(c, 0.75))
    hidexdecorations!(ax, grid = false)
    hideydecorations!(ax, grid = false, ticklabels = false, ticks = false)
    linkxaxes!(ax, axs[2])
    label!(scene, ax, "b")
    layout[1:2, 0] = LText(scene, "Y (cm)", rotation = π/2);
    layout[4, 2:3] = LLegend(scene, apply_element.(values.([polys, legendmarkers])), collect.(keys.([polys, legendmarkers])), ["Direction of displacements", " "], orientation = :horizontal, nbanks = 2, tellheight = true, height = Auto(), groupgap = 30);
    scene
end
# FileIO.save("a.pdf", scene)

######################## Figure 6 ################
function figure6(df)
    gdf = groupby(df, :group)
    g = gdf[[(group = "zero",)]]
    polys = OrderedDict(string(k.group) => (color = v.groupcolor[1], strokecolor = :transparent) for (k, v) in pairs(g))
    scene, layout = layoutscene(0, resolution = (max_width, 400.0))
    ax = layout[1,1] = LAxis(scene)
    plottracks!(ax, g)
    hidexdecorations!(ax, grid = false, ticklabels = false, ticks = false)
    # hideydecorations!(ax, grid = false, ticklabels = false, ticks = false)
    label!(scene, ax, "a")
    axs = [LAxis(scene) for _ in 1:2]
    for (i, point_type) in enumerate(("turning_point", "center_of_search"))
        plotpoints!(axs[i], g, point_type)
    end
    layout[1,2:3] = axs
    hideydecorations!(axs[1], grid = false, ticklabels = false, ticks = false)
    hideydecorations!(axs[2], grid = false)
    hidexdecorations!.(axs, grid = false, ticklabels = false, ticks = false)
    layout[2, 1:3] = LText(scene, "X (cm)");
    linkaxes!(axs...)
    # linkyaxes!(ax, axs...)
    label!(scene, axs[1], "b")
    label!(scene, axs[2], "c")
    layout[3, 1:3] = LLegend(scene, apply_element(values(legendmarkers)), collect(keys(legendmarkers)), orientation = :horizontal, nbanks = 2, tellheight = true, height = Auto(), groupgap = 30);
    scene
end
# FileIO.save("a.pdf", scene)

######################## Figure 7 ################
function figure7(df)
    gdf = groupby(df, :group)
    g = gdf[[(group = "far",)]]
    polys = OrderedDict(string(k.group) => (color = v.groupcolor[1], strokecolor = :transparent) for (k, v) in pairs(g))
    scene, layout = layoutscene(0, resolution = (max_width, 300.0))
    ax = layout[1,1] = LAxis(scene)
    plottracks!(ax, g)
    hidexdecorations!(ax, grid = false, ticklabels = false, ticks = false)
    # hideydecorations!(ax, grid = false, ticklabels = false, ticks = false)
    label!(scene, ax, "a")
    axs = [LAxis(scene) for _ in 1:2]
    for (i, point_type) in enumerate(("turning_point", "center_of_search"))
        plotpoints!(axs[i], g, point_type)
    end
    layout[1,2:3] = axs
    hideydecorations!(axs[1], grid = false, ticklabels = false, ticks = false)
    hideydecorations!(axs[2], grid = false)
    hidexdecorations!.(axs, grid = false, ticklabels = false, ticks = false)
    layout[2, 1:3] = LText(scene, "X (cm)");
    linkaxes!(axs...)
    # linkaxes!(ax, axs...)
    label!(scene, axs[1], "b")
    label!(scene, axs[2], "c")
    x = copy(legendmarkers);
    delete!(x, "burrow")
    layout[3,1:3] = LLegend(scene, apply_element(values(x)), collect(keys(x)), orientation = :horizontal, nbanks = 2, tellheight = true, height = Auto(), groupgap = 30);
    scene
end
# FileIO.save("a.pdf", scene)

######################## Figure 4 ################

function binit(track, h, nbins, m, M)
    o = Union{Variance, Missing}[Variance() for _ in 1:nbins]
    d = track.rawcoords
    to = findfirst(x -> x.xy[2] > 0, d) - 1
    for (p1, p2) in Iterators.take(zip(d, lag(d, -1, default = d[to])), to)
        y = -(p2.xy[2] + p1.xy[2])/2
        if m < y < M
            i = StatsBase.binindex(h, y)
            v = norm(p2.xy - p1.xy)/(p2.t - p1.t)
            fit!(o[i], v)
        end
    end
    replace!(x -> nobs(x) < 2 ? missing : x, o)
end
function figure4(df)
    gdf = groupby(df, :group)
    g = gdf[[(group = "none",)]]
    polys = OrderedDict(string(k.group) => (color = v.groupcolor[1], strokecolor = :transparent) for (k, v) in pairs(g))
    scene = Scene(resolution = (max_width, 500.0), camera=campixel!);
    sz = round(Int, max_width/4)
    lo = GridLayout(
                    scene, 4, 3,
                    # colsizes = [Auto(), Fixed(sz), Auto(), Auto()],
                    # rowsizes = [Auto(), Auto(), Fixed(sz), Auto()],
                    alignmode = Outside(0, 0, 0, 0)
    )
    lo[2, 1:3] = LText(scene, "X (cm)", tellheight = true);
    ax1 = lo[1,1] = LAxis(scene, autolimitaspect = 1)
    # ax1 = lo[2,2] = LAxis(scene, xaxisposition = :top, xticklabelalign = (:center, :bottom), autolimitaspect = 1)
    # lo[1:2, 1] = LText(scene, "Y (cm)", rotation = π/2);
    plottracks!(ax1, g)
    rect = FRect2D(-10,-10,20,20)
    spinecolor = RGB(only(distinguishable_colors(1, [colorant"white"; g[1].color], dropseed = true))) #:yellow
    lines!(ax1, rect, color = spinecolor);
    hidexdecorations!(ax1, grid = false, ticklabels = false, ticks = false);
    # hideydecorations!(ax1, grid = false, ticklabels = false, ticks = false);
    # hideydecorations!(ax1, grid = false, ticklabels = false, ticks = false)
    label!(scene, ax1, "a")
    axs = [LAxis(scene, autolimitaspect = 1) for _ in 1:2]
    # axs = [LAxis(scene, xaxisposition = :top, xticklabelalign = (:center, :bottom), autolimitaspect = 1) for _ in 1:2]
    for (i, point_type) in enumerate(("turning_point", "center_of_search"))
        plotpoints!(axs[i], g, point_type)
    end
    lo[1,2:3] = axs
    hideydecorations!(axs[1], grid = false, ticklabels = false, ticks = false)
    hideydecorations!(axs[2], grid = false)
    hidexdecorations!.(axs, grid = false, ticklabels = false, ticks = false)
    linkaxes!(axs...)
    label!(scene, axs[1], "b")
    label!(scene, axs[2], "c")
    d = DataFrame(g[1], copycols = false)
    sort!(d, :turning_point, by = norm)
    d.color .= getcolor(d.groupcolor)
    d = d[1:3,:]
    # postTP = 25
    ax2 = lo[3,1] = LAxis(scene,
                          aspect = DataAspect(),
                          autolimitaspect = 1,
                          bottomspinecolor = spinecolor,
                          topspinecolor = spinecolor,
                          leftspinecolor = spinecolor,
                          rightspinecolor = spinecolor,
                          # xaxisposition = :top, xticklabelalign = (:center, :bottom)
                         )
    for r in eachrow(d)
        lines!(ax2, homing(r.track); legendmarkers["track"]..., color = r.color)
        postTP = findfirst(xy -> any(>(10)∘abs, xy), searching(r.track))
        lines!(ax2, searching(r.track)[1:postTP]; legendmarkers["track"]..., color = r.color)
        # lines!(ax2, searching(r.track)[1:postTP]; legendmarkers["track"]..., color = 1:postTP, colormap = [r.color, colorant"white"])
        scatter!(ax2, [r.turning_point]; legendmarkers["turning point"]..., color = RGBA(r.color, 0.75))
    end
    scatter!(ax2, [zero(Point2f0)]; legendmarkers["burrow"]...)
    ax2.targetlimits[] = rect
    hidexdecorations!(ax2, grid = false, ticklabels = false, ticks = false)
    # hideydecorations!(ax2, grid = false, ticklabels = false, ticks = false)
    label!(scene, ax2, "d")
    m, M = (0, 120)
    nbins = 6
    bins = range(m, stop = M, length = nbins + 1)
    mbins = StatsBase.midpoints(bins)
    h = Histogram(bins)
    g = DataFrame(g[1], copycols = false)
    g[!, :id] .= 1:nrow(g)
    DataFrames.transform!(g, :id, :track => ByRow(x -> binit(x, h, nbins, m, M)) => :yv)
    μ = [Variance() for _ in 1:nbins]
    for i in 1:nbins
        reduce(merge!, skipmissing(yv[i] for yv in g.yv), init = μ[i])
    end
    bandcolor = RGB(only(distinguishable_colors(1, [colorant"white"; g.color; spinecolor], dropseed = true))) #:yellow
    ax = lo[3,2:3] = LAxis(scene,
                           aspect = nothing,
                           autolimitaspect = nothing,
                           xlabel = "Distance to burrow (cm)",
    ylabel = "Speed (cm/s)",
    xticks = mbins,
    xreversed = true,
    yaxisposition = :right,
    yticklabelalign = (:left, :center)
   )
    bh = band!(ax, mbins, mean.(μ) .- std.(μ), mean.(μ) .+ std.(μ), color = RGBA(bandcolor, 0.25))
    lh = lines!(ax, mbins, mean.(μ), color = :white, linewidth = 5)
    for r in eachrow(g)
        xy = [Point2f0(x, mean(y)) for (x,y) in zip(mbins, r.yv) if !ismissing(y)]
        lines!(ax, xy; legendmarkers["track"]..., color = r.color)
        scatter!(ax, xy; legendmarkers["turning point"]..., color = r.color)
    end
    ylims!(ax, 0, ax.limits[].origin[2] + ax.limits[].widths[2])
    xlims!(ax, Iterators.reverse(extrema(mbins))...)
    label!(scene, ax, "e")
    x = copy(legendmarkers)
    delete!(x, "fictive burrow")
    x["mean speed ± std"] = [(color = RGBA(bandcolor, 0.25), strokecolor = :transparent), (linestyle = nothing, linewidth = 3, color = :white)]
    x["individual speed"] = [x["track"], x["turning point"]]
    lo[4, 1:3] = LLegend(scene, apply_element(values(x)), collect(keys(x)), orientation = :horizontal, nbanks = 2, tellheight = true, height = Auto(), groupgap = 30);
    tight_ticklabel_spacing!(ax)
    colsize!(lo, 2, Relative(1/3))
    rowsize!(lo, 3, Aspect(2, 1))
    overlay = Scene(scene, camera = campixel!, raw = true)
    for (x, corner) in zip((-10.0, 10.0), (:topleft, :topright))
        point = Node(Point2f0(x, -10.0))
        topleft_ax2 = lift(getfield(MakieLayout, corner), ax2.scene.px_area)
        point_screenspace = lift(ax1.scene.camera.projectionview, ax1.scene.camera.pixel_space, point) do pv, pspace, point
            projected = Point2f0(AbstractPlotting.project(inv(pspace) * pv, point)[1:2]...) .+ AbstractPlotting.origin(ax1.scene.px_area[])
        end
        lines!(overlay, @lift([$point_screenspace, $topleft_ax2]), color = spinecolor)
    end
    scene
end

function save_figures(df)
    # mkpath("figures")
    for figure in (:figure4, :figure5, :figure6, :figure7)
        @eval FileIO.save($("$figure.pdf"), $figure($df))
    end
end


text:
using Test, LinearAlgebra
using QBase

@testset "./src/constructors/states.jl" begin

@testset "pure_state()" begin
    @testset "valid inputs" begin
        ρ = pure_state(Ket([im,0]))
        @test ρ isa State{Complex{Int64}}
        @test rank(ρ) == 1
        @test ρ == [1 0;0 0]

        ρ = pure_state(Bra([0.,-1im]))
        @test ρ isa State{Complex{Float64}}
        @test rank(ρ) == 1
        @test ρ == [0 0;0 1]

        ρ = pure_state([1,1]/sqrt(2))
        @test ρ isa State{Float64}
        @test rank(ρ) == 1
        @test ρ ≈ [1 1;1 1]/2

        ρ = pure_state(([1,-im]/sqrt(2))')
        @test ρ isa State{Complex{Float64}}
        @test rank(ρ) == 1
        @test ρ ≈ [1 im;-im 1]/2
    end

    @test_throws DomainError pure_state([1;1])
end

@testset "mixed_state()" begin
    @testset "strongly typed inputs" begin
        ρ_mix = mixed_state(
            Probabilities([0.7,0.2,0.1]),
            State.([[1 0 0;0 0 0;0 0 0],[0 0 0;0 0 0;0 0 1],[0 0 0;0 1 0;0 0 0]])
        )
        @test ρ_mix == [0.7 0 0;0 0.1 0;0 0 0.2]
        @test ρ_mix isa State{Float64}
        @test rank(ρ_mix) == 3
    end

    @testset "loosely typed inputs" begin
        ρ_mix = mixed_state([0.7,0.2,0.1],[[1 0 0;0 0 0;0 0 0],[0 0 0;0 0 0;0 0 1],[0 0 0;0 1 0;0 0 0]])
        @test ρ_mix == [0.7 0 0;0 0.1 0;0 0 0.2]
        @test ρ_mix isa State{Float64}
        @test rank(ρ_mix) == 3

        ρ_mix = mixed_state(
            Probabilities([1/2, 1/2]),
            [0.5*[1 im;-im 1], [1 0;0 0]]
        )
        @test ρ_mix == [0.75 0.25im;-0.25im 0.25]
        @test ρ_mix isa State{ComplexF64}
        @test rank(ρ_mix) == 2

        @test_throws DomainError mixed_state([0.7,0.2], [[1 0;0 0],[0 0;0 1]])
        @test_throws DomainError mixed_state([0.7,0.3], [[1 0;0 .1],[0 0;0 1]])
    end
end

@testset "trine_qubit_states()" begin
    trine_states = trine_qubit_states()
    @test isapprox(trine_states, mirror_symmetric_qubit_states(π/3), atol=1e-7)
    @test trine_states isa Vector{State{Float64}}
end

@testset "mirror_symmetric_qubit_states()" begin
    @testset "range over θ" begin
        for θ in 0:0.1π:π/2
            @test mirror_symmetric_qubit_states(θ) isa Vector{State{Float64}}
        end
    end

    @test_throws DomainError mirror_symmetric_qubit_states(π)

    @test isapprox(mirror_symmetric_qubit_states(π/2),[[1 0; 0 0], [0 0; 0 1], [0 0; 0 1]], atol=1e-7)
end

@testset "bloch_qubit_state()" begin
    @test [1 0;0 0] == bloch_qubit_state(0,0,1)
    @test 0.5*[1 -im;im 1] == bloch_qubit_state(0,1,0)
    @test 0.5*[1 1;1 1] == bloch_qubit_state(1,0,0)
    @test 0.5*[1 0;0 1] == bloch_qubit_state(0,0,0)

    @test bloch_qubit_state(0,0,1) == bloch_qubit_state(0,0)
    @test isapprox(bloch_qubit_state(0,1,0),bloch_qubit_state(π/2,π/2),atol=1e-7)
    @test isapprox(bloch_qubit_state(1,0,0),bloch_qubit_state(π/2,0),atol=1e-7)

    @test bloch_qubit_state(0,0,1) isa State{Complex{Float64}}
    @test_throws DomainError bloch_qubit_state(1,1,1)
end

@testset "sic_qubit_states()" begin
    sic_states = sic_qubit_states()

    @test sic_states isa Vector{State{Complex{Float64}}}

    θ_sic = 2*acos(1/sqrt(3))

    @test sic_states[1] ≈ bloch_qubit_state(0, 0)
    @test sic_states[2] ≈ bloch_qubit_state(θ_sic, 0)
    @test sic_states[3] ≈ bloch_qubit_state(θ_sic, 2π/3)
    @test sic_states[4] ≈ bloch_qubit_state(θ_sic, 4π/3)
end

@testset "bb84_qubit_states()" begin
    bb84_states = bb84_qubit_states()
    @test bb84_states isa Vector{State{Float64}}

    @test isapprox(bb84_states, planar_symmetric_qubit_states(4), atol=1e-7)
end

@testset "computational_basis_states()" begin
    @test computational_basis_states(3) == [[1 0 0;0 0 0;0 0 0],[0 0 0;0 1 0;0 0 0],[0 0 0;0 0 0;0 0 1]]
    @test computational_basis_states(2) == [[1 0;0 0],[0 0;0 1]]
end

@testset "bell_states()" begin
    states = bell_states()

    @test states isa Vector{State{Float64}}
    @test states[1] ≈ 1/2*[1 0 0 1;0 0 0 0;0 0 0 0;1 0 0 1]
    @test states[2] ≈ 1/2*[1 0 0 -1;0 0 0 0;0 0 0 0;-1 0 0 1]
    @test states[3] ≈ 1/2*[0 0 0 0;0 1 1 0;0 1 1 0;0 0 0 0]
    @test states[4] ≈ 1/2*[0 0 0 0;0 1 -1 0;0 -1 1 0;0 0 0 0]
end

@testset "generalized_bell_states()" begin
    @testset "qubit bell states" begin
        states = generalized_bell_states(2)

        @test states isa Vector{State{Complex{Float64}}}
        @test isapprox(states, bell_states(),atol=1e-7)
    end
end

@testset "planar_symmetric_qubit_states()" begin
    @test isapprox(planar_symmetric_qubit_states(2), [[1 0;0 0],[0 0;0 1]], atol=1e-7)
    @test isapprox(planar_symmetric_qubit_states(3), trine_qubit_states(),atol=1e-7)

    @testset "scanning over cases verifying right number are computed" begin
        for n in 5:50
            qubits = planar_symmetric_qubit_states(n)

            @test length(qubits) == n
            @test qubits isa Vector{State{Float64}}
        end
    end
end

end


text:
using FactorGraph
using OMEinsum
using Zygote: @adjoint
using Zygote

include("sequential.jl")

"""
convert bonds to a incidence list - legmap
"""
function legmap_frombonds(bonds)
    legmap = zeros(Int, 60, length(bonds))
    occ = zeros(Int, 60)
    for (ib,bond) in enumerate(bonds)
        A, B = bond
        occ[A] += 1
        occ[B] += 1
        legmap[A, ib] = occ[A]
        legmap[B, ib] = occ[B]
    end
    legmap
end
@adjoint legmap_frombonds(bonds) = legmap_frombonds(bonds), adjy->nothing

"""the tensor defined on the site of a ferromagnetic ising model"""
function trg_T0(K)
    # sqrt(exp.(-K K; K -K))
    M = [sqrt(cosh(K)) sqrt(sinh(K));
         sqrt(cosh(K)) -sqrt(sinh(K))]
    T = ein"ij,ik,il->jkl"(M, M, M)
end

function fullerene(K)
    T0 = trg_T0(K)
    tensors = Tuple([T0 for i=1:60])

    # construct the factor graph
    bonds = ((1, 10), (1, 41), (1, 59), (2, 12), (2, 42), (2, 60), (3, 6), (3,
        43), (3, 57), (4, 8), (4, 44), (4, 58), (5, 13), (5, 56), (5,
        57), (6, 10), (6, 31), (7, 14), (7, 56), (7, 58), (8, 12), (8,
        32), (9, 23), (9, 53), (9, 59), (10, 15), (11, 24), (11, 53), (11,
        60), (12, 16), (13, 14), (13, 25), (14, 26), (15, 27), (15,
        49), (16, 28), (16, 50), (17, 18), (17, 19), (17, 54), (18,
        20), (18, 55), (19, 23), (19, 41), (20, 24), (20, 42), (21,
        31), (21, 33), (21, 57), (22, 32), (22, 34), (22, 58), (23,
        24), (25, 35), (25, 43), (26, 36), (26, 44), (27, 51), (27,
        59), (28, 52), (28, 60), (29, 33), (29, 34), (29, 56), (30,
        51), (30, 52), (30, 53), (31, 47), (32, 48), (33, 45), (34,
        46), (35, 36), (35, 37), (36, 38), (37, 39), (37, 49), (38,
        40), (38, 50), (39, 40), (39, 51), (40, 52), (41, 47), (42,
        48), (43, 49), (44, 50), (45, 46), (45, 54), (46, 55), (47,
        54), (48, 55))
    lm = legmap_frombonds(bonds)
    fg = FGraph(tensors, lm)

    # get the optimal contraction tree, optmaltree does not work, unfortunately
    order = (7, 14, 13, 5, 56,
        25, 26, 58, 29, 57,
        4, 44, 36, 35, 43, 3, 21, 33, 34, 22,
        50, 38, 37, 49, 6, 31, 45, 46, 32, 8,
        47, 54, 55, 48, 12, 16, 40, 39, 15, 10,
        1, 41, 19, 17, 18, 20, 42, 2, 60, 28, 52, 30, 51, 27, 59,
        9, 23, 24, 11, 53)

    y = sequential_contract(fg, order)
    (y[] |> log)/60
end

using Test
@testset "fullerene backward" begin
    ngradient(f, x) = (f(x+1e-5) - f(x-1e-5))/2e-5
    # ferromagnetic model
    @test fullerene(1.0) ≈ 1.5147392357988083
    @test ngradient(fullerene, 0.1) ≈ fullerene'(0.1)
end


text:
module VectorizedArgMinMax

using Polyester
using LoopVectorization
import LoopVectorization.length_one_axis

export vargmax, vargmin, vtargmax, vtargmin

include("vargmax.jl")

export vfindmax, vfindmin, vtfindmax, vtfindmin

include("vfindmax.jl")

end


text:
function findSS!(model::RSDSGEModel,guess::Array{Float64,1})
    ssparameters = Array{Float64}(model.meta.numparameters)
    for (p,affects) in enumerate(model.parameters.affectsSS)
        if affects
            ssparameters[p] = model.parameters.decomposition[p].bar
        end
    end
    ssparameters[.!model.parameters.affectsSS] = model.parameters.values[1, .!model.parameters.affectsSS]

    sswrapper = x->evaluateSS(model.steadystate.system,x,ssparameters)
    print("Solving for steady state...")
    nlout = NLsolve.nlsolve(NLsolve.not_in_place(sswrapper), guess)
    if !NLsolve.converged(nlout)
        println("FAILED.")
        println("Residuals at initial guess:")
        println(sswrapper(guess))
        error("Could not find steady state.")
    end
    model.steadystate.values = nlout.zero
    println("done.")
    return nlout.zero
end

function findSS!(model::RSDSGEModel)
    return findSS!(model,model.steadystate.values)
end

function updateparameters!(model::RSDSGEModel,name::String,value,reevaluateSS=true)
    index = model.parameters.dictionary[name]
    updateparameters!(model,index,value,reevaluateSS)
    return
end

function updateparameters!(model::RSDSGEModel,index::Int,value::Number,reevaluateSS=true)
    if model.parameters.isswitching[index]
        error("$(model.parameters.names[index]) is a state-dependent parameter, but you passed in only one value.")
    end
    model.parameters.values[:,index] = value
    # Re-evaluate SS. I currently don't track whether non-switching parameters affect the SS or not. I should, because I could be much more
    # efficient if I don't need to re-evaluate the SS on all parameters
    if reevaluateSS
        findSS!(model)
    end
    return
end

function updateparameters!(model::RSDSGEModel,index::Int,value::Tuple,reevaluateSS=true)
    updateparameters!(model,index,collect(value),reevaluateSS)
    return
end

function updateparameters!(model::RSDSGEModel,index::Int,value::Array{Float64,1},reevaluateSS=true)
    if !model.parameters.isswitching[index]
        error("$(model.parameters.names[index]) is not a regime switching parameter, but you passed in many values for it.")
    elseif length(value) != model.meta.numregimes
        error("Dimension mismatch. You passed in $(length(value)) values for $(model.parameters.name[index]), but there are $(model.meta.numregimes) regimes.")
    end
    model.parameters.values[:,index] = value
    if model.parameters.affectsSS[index]
	# I need to re do the parameter decomposition here
	newdecomp = RegimeParameterDecomposition(model.parameters.values[:,index],model.transmatrix.ergodic)
	model.parameters.decomposition[index] = newdecomp
	if reevaluateSS
	    findSS!(model)
	end
    end
    return
end

function updateparameters!(model::RSDSGEModel,names::Array{String,1},values::Array{Any},reevaluateSS=true)
    if length(names) != length(values)
	error("You supplied a different number of parameters to change ($(length(names)) than values ($(length(values)).")
    end
    for (i,name,value) in zip(1:length(names),names,values)
	if i < length(names)
	    updateparameters!(model,name,value,false)
	else
	    updateparameters!(model,name,value,reevaluateSS)
	end
    end
    return
end


text:
# features that are being used by parser, taken from :
# https://github.com/denizyuret/KUparser.jl/blob/conll17/src/features.jl

function getanchor(f::AbstractString, p::Parser, ldep, rdep)
    f1 = f[1]; f2 = f[2]; flen = length(f)
    if isdigit(f2)
        i = f2 - '0'            # index
        n = 3                   # next character
    else
        i = 0
        n = 2
    end
    a = 0                       # target word
    if f1 == 's'
        if (p.sptr - i >= 1) # 25
            a = Int(p.stack[p.sptr - i])             # 456
        end
    elseif f1 == 'n'
        if p.wptr + i <= p.nword
            a = p.wptr + i
        end
    else
        error("feature string should start with [sn]")
    end
    if a==0; return 0; end
    while n < flen
        f1 = f[n]; f2 = f[n+1]
        if isdigit(f2)
            i = f2 - '0'
            n = n+2
        else
            i = 1
            n = n+1
        end
        if i <= 0
            error("hlr indexing is one based") # 3 [lrh] is one based, [sn] was zero based
        end
        if f1 == 'l'                              # 2
            if a > p.wptr; error("buffer words other than n0 do not have ldeps"); end # 252
            if isassigned(ldep,a) && i <= length(ldep[a])
                a = Int(ldep[a][i])
            else
                return 0
            end
        elseif f1 == 'r'
            if a >= p.wptr; error("buffer words do not have rdeps"); end
            if isassigned(rdep,a) && i <= length(rdep[a])
                a = Int(rdep[a][i])
            else
                return 0
            end
        elseif f1 == 'h'
            for j=1:i       # 5
                a = Int(p.head[a]) # 147
                if a == 0
                    return 0
                end
            end
        else
            break
        end
    end
    if n != length(f); error("n!=length(f)"); end
    return a
end

function getdeps2{T<:Parser}(p::T)
    nw = p.nword
    nd = p.ndeps
    dr = p.deprel
    ldep = Array(Any,nw)
    rdep = Array(Any,nw)
    @inbounds for d=1:nw
        h=Int(p.head[d])
        if h==0
            continue
        elseif d<h
            if !isassigned(ldep,h)
                ldep[h]=[d]
            else
                push!(ldep[h],d)
            end
        elseif d>h
            if !isassigned(rdep,h)
                rdep[h]=[d]
            else
                unshift!(rdep[h],d)
            end
        else
            error("h==d")
        end
    end
    return (ldep,rdep)
end

function getrdist(f::AbstractString, p::Parser, a::Integer)
    if f[1]=='s' && a > 0
        if isdigit(f[2])
            i=f[2]-'0'
        else
            i=0
        end
        if i>0
            return p.stack[p.sptr - i + 1] - a
        elseif p.wptr <= p.nword
            return p.wptr - a
        else
            return 0
        end
    else
        return 0 # dist only defined for stack words
    end
end

function features(parsers, feats, model)
    pvecs,dvecs,lvecs,rvecs,xvecs = postagv(model),deprelv(model),lcountv(model),rcountv(model),distancev(model)
    pvec0,dvec0,lvec0,rvec0,xvec0 = zeros(pvecs[1]),zeros(dvecs[1]),zeros(lvecs[1]),zeros(rvecs[1]),zeros(xvecs[1])
    s = parsers[1].sentence
    wvec0,fvec0,bvec0 = zeros(s.wvec[1]),zeros(s.fvec[1]),zeros(s.bvec[1])
    fmatrix = []
    for p in parsers
        s = p.sentence
        ldep,rdep = getdeps2(p)
        for f in feats
            a = getanchor(f,p,ldep,rdep)
            fn = f[end]
            if fn == 'v'
                push!(fmatrix, a>0 ? s.wvec[a] : wvec0)
            elseif fn == 'c'
                push!(fmatrix, a>0 ? s.fvec[a] : fvec0)
                push!(fmatrix, a>0 ? s.bvec[a] : bvec0)
            elseif fn == 'p'
                push!(fmatrix, a>0 ? pvecs[s.postag[a]] : pvec0)
            elseif fn == 'L'
                push!(fmatrix, a>0 ? dvecs[p.deprel[a]] : dvec0)
            elseif fn == 'A'
                push!(fmatrix, a==0 ? dvec0 : !isassigned(ldep,a) ? dvec0 : length(ldep[a])==1 ? dvecs[p.deprel[ldep[a][1]]] : sum(dvecs[p.deprel[ldep[a]]]))
            elseif fn == 'B'
                push!(fmatrix, a==0 ? dvec0 : !isassigned(rdep,a) ? dvec0 : length(rdep[a])==1 ? dvecs[p.deprel[rdep[a][1]]] : sum(dvecs[p.deprel[rdep[a]]]))
            elseif fn == 'd'
                d = getrdist(f,p,a)
                push!(fmatrix, d>0 ? xvecs[min(d,10)] : xvec0)
            elseif fn == 'a'
                d = (a>0 && isassigned(ldep,a) ? length(ldep[a]) : 0)
                push!(fmatrix, a>0 ? lvecs[min(d+1,10)] : lvec0)
            elseif fn == 'b'
                d = (a>0 && isassigned(rdep,a) ? length(rdep[a]) : 0)
                push!(fmatrix, a>0 ? rvecs[min(d+1,10)] : rvec0)
            elseif fn == 'w'
                error("Dense features do not support 'w'")
            else
                error("Unknown feature $(fn)")
            end
        end
    end
    fmatrix = vcatn(fmatrix...)
    ncols = length(parsers)
    nrows = div(length(fmatrix), ncols)
    reshape(fmatrix, nrows, ncols)
end

max_stars_repo_name: kirnap/KU-dependencyParser

max_issues_repo_name: kirnap/KU-dependencyParser

max_forks_repo_name: kirnap/KU-dependencyParser


text:
using PeriodicTable
using DataFrames
# Implements Pouchou and Pichoir's XPP model from the description in the book
# Electron Probe Quantification edited by Kurt Heinrich and Dale E. Newbury

struct XPP <: MatrixCorrection
    # Configuration data
    subshell::AtomicSubShell
    material::Material
    E0 # Beam energy (eV)
    # Computed items
    ϕ0 # ϕ(ρz) at the surface
    A # Amplitude  factor
    a # Width factor
    B # Amplitude  factor
    b # Width factor
    F # Integral under the ϕ(ρz) curve

    """
        XPP(mat::Material, ashell::AtomicSubShell, E0::AbstractFloat)

    Construct an XPP object for the specified material, atomicsubshell,
    beam energy (in eV).
    """
    function XPP(mat::Material, ashell::AtomicSubShell, E0::AbstractFloat)
        # XPP calculations expect E0, eLv, J in keV
        e0, Mv, Jv, = 0.001 * E0, M(XPP, mat), 0.001 * J(XPP, mat)
        Zbarbv, eLv = Zbarb(XPP, mat), 0.001 * energy(ashell)
        U0v, V0v, ηbarv = e0 / eLv, e0 / Jv, ηbar(XPP, Zbarbv)
        @assert U0v > 1.0 "The beam energy must be larger than the subshell edge energy."
        @assert V0v > 1.0 "The beam energy must be larger than the mean energy loss. ($(mat), $(ashell), $(E0))"
        Dv, Pjv, mv, Wbarv = D(XPP, Jv), P(XPP, Jv), m(XPP, ashell), Wbar(XPP, ηbarv)
        invSv = invS(XPP, U0v, V0v, Mv, Dv, Pjv, T(XPP, Pjv, mv))
        qv, Rv, ϕ0v, QlAv = q(XPP, Wbarv), R(XPP, ηbarv, Wbarv, U0v), ϕ0(XPP, U0v, ηbarv), QlA(XPP, U0v, eLv, mv)
        Fv, FoverRbarv = F(XPP, Rv, invSv, QlAv), FoverRbar(XPP, X(XPP, Zbarbv), Y(XPP, Zbarbv), U0v)
        Rbarv = Rbar(XPP, Fv, FoverRbarv, ϕ0v)
        bv = b(XPP, Rbarv, ϕ0v, Fv)
        Pv = P(XPP, Zbarbv, U0v, Fv, Rbarv, ϕ0v, bv)
        ϵv = ϵ(XPP, Pv, bv, ϕ0v, Fv, Rbarv)
        av, Bv = a(XPP, bv, ϵv), B(XPP, bv, Fv, Pv, ϕ0v, ϵv)
        Av = A(XPP, Bv, bv, ϕ0v, Fv, ϵv)
        return new(ashell, mat, E0, ϕ0v, Av, av, Bv, bv, Fv)
    end
end

"""
    M(mat::Material)

P&P's M from top line in page 35
"""
M(::Type{XPP}, mat::Material) = #C1
    sum(NeXLCore.nonneg(mat, elm) * z(elm) / a(elm, mat) for elm in keys(mat))

"""
    J(::Type{XPP}, mat::Material)

Mean ionization potential for the specified material in eV. (PAP1991 Eqn 6)
"""
NeXLCore.J(::Type{XPP}, mat::Material) = #C1
    exp(sum(NeXLCore.nonneg(mat, elm) * (z(elm) / a(elm, mat)) * log(J(Zeller1973, elm)) for elm in keys(mat)) / M(XPP, mat))

# From PAP eqn 8
D(::Type{XPP}, J) = (6.6e-6, 1.12e-5 * (1.35 - 0.45 * J^2), 2.2e-6 / J) #C1

# From PAP eqn 8
P(::Type{XPP}, J) = (0.78, 0.1, -0.5 + 0.25 * J) #C1

# From PAP near Eqn 11
T(::Type{XPP}, P, m) = collect(1.0 + P[k] - m for k = 1:3) #C1

"""
    dEdρs(args::Dict{Label,AbstractFloat}, mat::MaterialLabel, Ekev::AbstractFloat, elms)

The function P&P uses to describe the deceleration of an electron in a material.
Output units are (keV/cm)/(g/cm^3) = keV cm^2/g. (PAP eqn 5)
"""
function dEdρs(::Type{XPP}, mat::Material, Ekev::AbstractFloat) #C1
    @assert Ekev < 50.0 "It appears that the beam energy is in keV not eV. ($Ekev)"
    Jkev = 0.001 * J(XPP, mat) # XPP expects in keV
    # PAP Eqn 8
    function f(mat, Ekev, Jkev)
        d, p, v = D(XPP, j), P(XPP, j), Ekev / Jkev
        return sum(i -> d[i] * v^p[i], 1:3)
    end
    return -(M(XPP, mat) / Jkev) / f(mat, Ekev, Jkev)
end

"""
    R0(J, D, P, M, Ekev)

Total trajectory (range) of an electron with initial energy Ekev. (in cm/(g/cm^3))
"""
R0(::Type{XPP}, J, D, P, M, Ekev) = #C1
    sum(J^(1.0 - P[k]) * D[k] * Ekev^(1.0 + P[k]) / (1.0 + P[k]) for k = 1:3) / M

"""
    ϕxpp(ρz, A, a, B, b, ϕ0)

Compute the shape of the ϕ(ρz) curve in the XPP model.
"""
ϕxpp(::Type{XPP}, ρz, A, a, B, b, ϕ0) = #C1
    A * exp(-a * ρz) + (B * ρz + ϕ0 - A) * exp(-b * ρz)

"""
    invS(U0, V0, M, D, P, T)

Computes 1/S where S is the stopping power.
"""
invS(::Type{XPP}, U0, V0, M, D, P, T) = #C1
    U0 / (V0 * M) * sum(D[k] * ((V0 / U0)^P[k]) * (T[k] * U0^T[k] * log(U0) - U0^T[k] + 1.0) / (T[k]^2) for k = 1:3)

"""
    S(mat, ashell, E)

Computes S, the stopping power at the specified energy (in keV)
"""
function S(::Type{XPP}, mat::Material, ashell::AtomicSubShell, Ekev)
    @assert Ekev < 50.0 "It appears that the beam energy is in keV not eV. ($Ekev)"
    jkev = 0.001 * J(XPP, mat) # XPP expects in keV
    return 1.0 / invS(XPP, Ekev / (0.001 * energy(ashell)), Ekev / jkev, M(XPP, mat), D(XPP, jkev), P(XPP, jkev), T(XPP, P(XPP, jkev), m(XPP, ashell)))
end

"""
    QlA(U,El,m)

Computes the relative ionization cross-section.
"""
QlA(::Type{XPP}, U, El, m) = log(U) / ((U^m) * (El^2)) #C1

"""
    m(ashell::AtomicSubShell)

Returns the ionization cross-section exponent for QlA(U,El,m(ashell))
"""
function m(::Type{XPP}, ashell::AtomicSubShell) #C1
    if isequal(n(ashell), 1)
        return 0.86 + 0.12 * exp(-(ashell.z / 5.0)^2)
    else
        return n(ashell) == 2 ? 0.82 : 0.78
    end
end

# PAP appendix 1
Zbarb(::Type{XPP}, mat::Material) = #C1
    sum(NeXLCore.nonneg(mat, elm) * sqrt(z(elm)) for elm in keys(mat))^2

# PAP appendix 1
ηbar(::Type{XPP}, Zbarb) = #C2
    1.75e-3 * Zbarb + 0.37 * (1.0 - exp(-0.015 * Zbarb^1.3))

# PAP appendix 1
Wbar(::Type{XPP}, ηbar) = #C2
    0.595 + ηbar / 3.7 + ηbar^4.55

q(::Type{XPP}, Wbar) = #C1
    (2.0 * Wbar - 1.0) / (1.0 - Wbar)

# PAP appendix 1
JU0(::Type{XPP}, U0) = #C2
    1.0 + U0 * (log(U0) - 1.0)

# PAP appendix 1
G(::Type{XPP}, U0, q) = #C2
    (U0 - 1.0 - (1.0 - 1.0 / U0^(1.0 + q)) / (1.0 + q)) / ((2.0 + q) * JU0(XPP, U0))

# PAP appendix 1
R(::Type{XPP}, ηbar, Wbar, U0) = #C2
    1.0 - ηbar * Wbar * (1.0 - G(XPP, U0, q(XPP, Wbar)))

"""
    R(mat::Material, u0)

Backscatter factor as a function of material and overvoltage.
Compares well to PAP Figure 23.
"""
function R(::Type{XPP}, mat::Material, u0)
    zbarbv = Zbarb(XPP, mat)
    ηbarv = ηbar(XPP, zbarbv)
    wbar = Wbar(XPP, ηbarv)
    return R(XPP, ηbarv, wbar, u0)
end

"""
   ϕ0(U0, ηbar)

The value of the ϕ(ρz) curve at ρz=0.0.
"""
ϕ0(::Type{XPP}, U0, ηbar) = #C1
    1.0 + 3.3 * (1.0 - 1.0 / (U0^(2.0 - 2.3 * ηbar))) * ηbar^1.2

X(::Type{XPP}, Zbarb) = #C3
    1.0 + 1.3 * log(Zbarb)

Y(::Type{XPP}, Zbarb) = #C3
    0.2 + Zbarb / 200.0

FoverRbar(::Type{XPP}, X, Y, U0) = #C3
    1.0 + X * log(1.0 + Y * (1.0 - 1.0 / (U0^0.42))) / log(1.0 + Y)

# PAP eqn 13
F(::Type{XPP}, R, invS, QlA) = #C1
    R * invS / QlA

# PAP eqn 28
Rbar(::Type{XPP}, F, FoverRbar, ϕ0) = # C1
    FoverRbar >= ϕ0 ? F / FoverRbar : F / ϕ0

# PAP eqn 28
g(::Type{XPP}, Zbarb, U0) = #C1
    0.22 * log(4.0 * Zbarb) * (1.0 - 2.0 * exp(-Zbarb * (U0 - 1.0) / 15.0))

# PAP eqn 28
h(::Type{XPP}, Zbarb, U0) = #C1
    1.0 - 10.0 * (1.0 - 1.0 / (1.0 + U0 / 10.0)) / (Zbarb^2)

"""
    b(Rbar, ϕ0, F)

XPP ϕ(ρz) model parameter b
""" # PAP Appendix 4
b(::Type{XPP}, Rbar, ϕ0, F) = #C1
    sqrt(2.0) * (1.0 + sqrt(1.0 - Rbar * ϕ0 / F)) / Rbar

# PAP eqn 29
P(::Type{XPP}, Zbarb, U0, F, Rbar, ϕ0, b) = #C1
    min(g(XPP, Zbarb, U0) * h(XPP, Zbarb, U0)^4, 0.9 * b * Rbar^2 * (b - 2.0 * ϕ0 / F)) * F / (Rbar^2)

function ϵ(::Type{XPP}, P, b, ϕ0, F, Rbar) #C1
    a = (P + b * (2.0 * ϕ0 - b * F)) / (b * F * (2.0 - b * Rbar) - ϕ0)
    tmp = (a - b) / b
    # print("a = ",a,"(a-b)/b = ",tmp)
    return abs(tmp) > 1.0e-6 ? tmp : sign(tmp) * 1.0e-6
end

"""
    a(b, ϵ)

XPP ϕ(ρz) model parameter a
"""
NeXLCore.a(::Type{XPP}, b, ϵ) = #C1
    b * (1.0 + ϵ)

"""
    B(b, F, P, ϕ0, ϵ)

XPP ϕ(ρz) model parameter B
"""
B(::Type{XPP}, b, F, P, ϕ0, ϵ) = #C1
    (b^2 * F * (1.0 + ϵ) - P - ϕ0 * b * (2.0 + ϵ)) / ϵ

"""
    A(B, b, ϕ0, F, ϵ)

XPP ϕ(ρz) model parameter A
"""
A(::Type{XPP}, B, b, ϕ0, F, ϵ) = #C1
    (B / b + ϕ0 - b * F) * (1 + ϵ) / ϵ

function Fχ(::Type{XPP}, χ, A, a, B, b, ϕ0)
    ϵ = (a - b) / b
    return (ϕ0 + B / (b + χ) - A * b * ϵ / (b * (1.0 + ϵ) + χ)) / (b + χ)
end

"""
    Fχp(::Type{XPP}, χ, A, a, B, b, ϕ0, τ)

The integral of the ϕ(ρz) exp(-χ ρz) from 0 to τ.
"""
Fχp(::Type{XPP}, χ, A, a, B, b, ϕ0, τ) =
    (A * (1.0 - exp(-(τ * (a + χ))))) / (a + χ) +
    (A * (-1.0 + exp(-(τ * (b + χ))))) / (b + χ) +
    ((1 - exp(-(τ * (b + χ)))) * ϕ0) / (b + χ) +
    (B * (-1 + exp(τ * (b + χ)) - τ * (b + χ))) / (exp(τ * (b + χ)) * ((b + χ)^2))

"""
Represents the essential intermediary values for an XPP matrix correction of
characteristic X-rays from a particular atomic sub-shell in a particular material.
"""

function NeXLUncertainties.asa(::Type{DataFrame}, xpps::XPP...)
    return DataFrame(
        SubShell = [xpp.subshell for xpp in xpps],
        Material = [name(xpp.material) for xpp in xpps],
        BeamEnergy = [xpp.E0 for xpp in xpps],
        Phi0 = [xpp.ϕ0 for xpp in xpps],
        A = [xpp.A for xpp in xpps],
        a = [xpp.a for xpp in xpps],
        B = [xpp.B for xpp in xpps],
        b = [xpp.b for xpp in xpps],
        F = [xpp.F for xpp in xpps],
    )
end

Base.show(io::IO, xpp::XPP) = print(io, "XPP[$(xpp.subshell) in $(name(xpp.material)) at $(0.001*xpp.E0) keV]")

Fχ(xpp::XPP, xray::CharXRay, θtoa::Real) = Fχ(XPP, χ(material(xpp), xray, θtoa), xpp.A, xpp.a, xpp.B, xpp.b, xpp.ϕ0)

Fχp(xpp::XPP, xray::CharXRay, θtoa::Real, τ::Real) =
    Fχp(XPP, χ(material(xpp), xray, θtoa), xpp.A, xpp.a, xpp.B, xpp.b, xpp.ϕ0, τ)

F(xpp::XPP) = xpp.F

NeXLCore.minproperties(::Type{XPP}) = (:BeamEnergy, :TakeOffAngle)

"""
    ϕ(ρz, xpp::XPP)

Computes the ϕ(ρz) curve according to the XPP algorithm.
"""
ϕ(xpp::XPP, ρz) = ϕxpp(XPP, ρz, xpp.A, xpp.a, xpp.B, xpp.b, xpp.ϕ0)

"""
    range(::Type{XPP}, mat::MaterialLabel, e0, inclDensity=true)
    range(ty::Type{<:BetheEnergyLoss}, mat::Material, e0::Float64, inclDensity=true; emin=50.0, mip::Type{<:NeXLMeanIonizationPotential}=Berger1982)

Total trajectory (range) of an electron with initial energy e0 (eV). (Units: inclDensity ? cm : cm/(g/cm³))
"""
function Base.range(::Type{XPP}, mat::Material, e0::Real, inclDensity=true)
    j = 0.001 * J(XPP, mat) # XPP expects in keV
    return R0(XPP, j, D(XPP, j), P(XPP, j), M(XPP, mat), 0.001 * e0) / (inclDensity ? density(mat) : 1.0)
end

"""
    matrixcorrection(::Type{XPP}, mat::Material, ashell::AtomicSubShell,e0)

Constructs an instance of the XPP algorithm.
"""
matrixcorrection(::Type{XPP}, mat::Material, ashell::AtomicSubShell, e0) = XPP(mat, ashell, e0)