text:
### mcvar and mcse stand for Monte Carlo variance and Monte Carlo error respectively

## Monte Carlo variance assuming IID samples

mcvar(v::AbstractArray, ::Type{Val{:iid}}) = var(v)/length(v)

mcvar(v::AbstractArray, ::Type{Val{:iid}}, region) = mapslices(x -> mcvar(x, Val{:iid}), v, region)

mcvar(s::VariableNState{Univariate}, ::Type{Val{:iid}}) = mcvar(s.value, Val{:iid})

mcvar(s::VariableNState{Multivariate}, ::Type{Val{:iid}}, i::Integer) = mcvar(s.value[i, :], Val{:iid})

mcvar(s::VariableNState{Multivariate}, ::Type{Val{:iid}}, r::AbstractVector=1:s.size) =
  eltype(s)[mcvar(s, Val{:iid}, i) for i in r]

## Monte Carlo standard error assuming IID samples

mcse(v::AbstractArray, ::Type{Val{:iid}}) = sqrt(mcvar(v, Val{:iid}))

mcse(v::AbstractArray, ::Type{Val{:iid}}, region) = mapslices(x -> mcse(x, Val{:iid}), v, region)

mcse(s::VariableNState{Univariate}, ::Type{Val{:iid}}) = mcse(s.value, Val{:iid})

mcse(s::VariableNState{Multivariate}, ::Type{Val{:iid}}, i::Integer) = mcse(s.value[i, :], Val{:iid})

mcse(s::VariableNState{Multivariate}, ::Type{Val{:iid}}, r::AbstractVector=1:s.size) =
  eltype(s)[mcse(s, Val{:iid}, i) for i in r]

## Monte Carlo variance using batch means
## Reference:
## Flegal J.M and Jones G.L
## Batch Means and Spectral Variance Estimators in Markov chain Monte Carlo
## Annals of Statistics, 2010, 38 (2), pp 1034-1070

function mcvar(v::AbstractVector{T}, ::Type{Val{:bm}}, batchlen::Integer=100) where T
  nbatches = div(length(v), batchlen)
  @assert nbatches > 1 "Choose batch size such that the number of batches is greather than one"
  nbsamples = nbatches*batchlen
  batchmeans = T[mean(v[((j-1)*batchlen+1):(j*batchlen)]) for j = 1:nbatches]
  return batchlen*var(batchmeans)/nbsamples
end

mcvar(v::AbstractArray, ::Type{Val{:bm}}, region, batchlen::Integer=100) =
  mapslices(x -> mcvar(vec(x), Val{:bm}, batchlen), v, region)

mcvar(s::VariableNState{Univariate}, ::Type{Val{:bm}}, batchlen::Integer=100) = mcvar(s.value, Val{:bm}, batchlen)

mcvar(s::VariableNState{Multivariate}, ::Type{Val{:bm}}, i::Integer, batchlen::Integer=100) =
  mcvar(s.value[i, :], Val{:bm}, batchlen)

mcvar(s::VariableNState{Multivariate}, ::Type{Val{:bm}}, r::AbstractVector=1:s.size, batchlen::Integer=100) =
  eltype(s)[mcvar(s, Val{:bm}, i, batchlen) for i in r]

## Monte Carlo standard error using batch means

mcse(v::AbstractVector, ::Type{Val{:bm}}, batchlen::Integer=100) = sqrt(mcvar(v, Val{:bm}, batchlen))

mcse(v::AbstractArray, ::Type{Val{:bm}}, region, batchlen::Integer=100) =
  mapslices(x -> mcse(vec(x), Val{:bm}, batchlen), v, region)

mcse(s::VariableNState{Univariate}, ::Type{Val{:bm}}, batchlen::Integer=100) = mcse(s.value, Val{:bm}, batchlen)

mcse(s::VariableNState{Multivariate}, ::Type{Val{:bm}}, i::Integer, batchlen::Integer=100) =
  mcse(s.value[i, :], Val{:bm}, batchlen)

mcse(s::VariableNState{Multivariate}, ::Type{Val{:bm}}, r::AbstractVector=1:s.size, batchlen::Integer=100) =
  eltype(s)[mcse(s, Val{:bm}, i, batchlen) for i in r]

## Initial monotone sequence estimator (IMSE) of Monte Carlo variance
## Reference:
## Geyer C.J
## Practical Markov Chain Monte Carlo
## Statistical Science, 1992, 7 (4), pp 473-483

function mcvar(v::AbstractVector{T}, ::Type{Val{:imse}}, maxlag::Integer=length(v)-1) where T
  k = convert(Integer, floor((maxlag-1)/2))
  m = k+1

  # Preallocate memory
  g = Array(T, k+1)

  # Calculate empirical autocovariance
  acv = autocov(v, 0:maxlag)

  # Calculate \hat{G}_{n, m} from the autocovariances, see pp 477 in Geyer
  for j = 0:k
    g[j+1] = acv[2*j+1]+acv[2*j+2]
    if g[j+1] <= 0
      m = j
      break
    end
  end

  # Create the monotone sequence of g
  if m > 1
    for j = 2:m
      if g[j] > g[j-1]
        g[j] = g[j-1]
      end
    end
  end

  # Calculate the initial monotone sequence estimator
  return (-acv[1]+2*sum(g[1:m]))/length(v)
end

mcvar(v::AbstractArray, ::Type{Val{:imse}}, region) = mapslices(x -> mcvar(vec(x), Val{:imse}, length(x)-1), v, region)

mcvar(s::VariableNState{Univariate}, ::Type{Val{:imse}}, maxlag::Integer=s.n-1) = mcvar(s.value, Val{:imse}, maxlag)

mcvar(s::VariableNState{Multivariate}, ::Type{Val{:imse}}, i::Integer, maxlag::Integer=s.n-1) =
  mcvar(s.value[i, :], Val{:imse}, maxlag)

mcvar(s::VariableNState{Multivariate}, ::Type{Val{:imse}}, r::AbstractVector=1:s.size, maxlag::Integer=s.n-1) =
  eltype(s)[mcvar(s, Val{:imse}, i, maxlag) for i in r]

## Initial monotone sequence estimator (IMSE) of Monte Carlo standard error

mcse(v::AbstractVector, ::Type{Val{:imse}}, maxlag::Integer=length(v)-1) = sqrt(mcvar(v, Val{:imse}, maxlag))

mcse(v::AbstractArray, ::Type{Val{:imse}}, region) = mapslices(x -> mcse(vec(x), Val{:imse}, length(x)-1), v, region)

mcse(s::VariableNState{Univariate}, ::Type{Val{:imse}}, maxlag::Integer=s.n-1) = mcse(s.value, Val{:imse}, maxlag)

mcse(s::VariableNState{Multivariate}, ::Type{Val{:imse}}, i::Integer, maxlag::Integer=s.n-1) =
  mcse(s.value[i, :], Val{:imse}, maxlag)

mcse(s::VariableNState{Multivariate}, ::Type{Val{:imse}}, r::AbstractVector=1:s.size, maxlag::Integer=s.n-1) =
  eltype(s)[mcse(s, Val{:imse}, i, maxlag) for i in r]

## Initial positive sequence estimator (IPSE) of Monte Carlo variance
## Reference:
## Geyer C.J
## Practical Markov Chain Monte Carlo
## Statistical Science, 1992, 7 (4), pp 473-483

function mcvar(v::AbstractVector{T}, ::Type{Val{:ipse}}, maxlag::Integer=length(v)-1) where T
  k = convert(Integer, floor((maxlag-1)/2))

  # Preallocate memory
  g = Array(T, k+1)
  m = k+1

  # Calculate empirical autocovariance
  acv = autocov(v, 0:maxlag)

  # Calculate \hat{G}_{n, m} from the autocovariances, see pp. 477 in Geyer
  for j = 0:k
    g[j+1] = acv[2*j+1]+acv[2*j+2]
    if g[j+1] <= 0
      m = j
      break
    end
  end

  # Calculate the initial positive sequence estimator
  return (-acv[1]+2*sum(g[1:m]))/length(v)
end

mcvar(v::AbstractArray, ::Type{Val{:ipse}}, region) = mapslices(x -> mcvar(vec(x), Val{:ipse}, length(x)-1), v, region)

mcvar(s::VariableNState{Univariate}, ::Type{Val{:ipse}}, maxlag::Integer=s.n-1) = mcvar(s.value, Val{:ipse}, maxlag)

mcvar(s::VariableNState{Multivariate}, ::Type{Val{:ipse}}, i::Integer, maxlag::Integer=s.n-1) =
  mcvar(s.value[i, :], Val{:ipse}, maxlag)

mcvar(s::VariableNState{Multivariate}, ::Type{Val{:ipse}}, r::AbstractVector=1:s.size, maxlag::Integer=s.n-1) =
  eltype(s)[mcvar(s, Val{:ipse}, i, maxlag) for i in r]

## Initial positive sequence estimator (IPSE) of Monte Carlo standard error

mcse(v::AbstractVector, ::Type{Val{:ipse}}, maxlag::Integer=length(v)-1) = sqrt(mcvar(v, Val{:ipse}, maxlag))

mcse(v::AbstractArray, ::Type{Val{:ipse}}, region) = mapslices(x -> mcse(vec(x), Val{:ipse}, length(x)-1), v, region)

mcse(s::VariableNState{Univariate}, ::Type{Val{:ipse}}, maxlag::Integer=s.n-1) = mcse(s.value, Val{:ipse}, maxlag)

mcse(s::VariableNState{Multivariate}, ::Type{Val{:ipse}}, i::Integer, maxlag::Integer=s.n-1) =
  mcse(s.value[i, :], Val{:ipse}, maxlag)

mcse(s::VariableNState{Multivariate}, ::Type{Val{:ipse}}, r::AbstractVector=1:s.size, maxlag::Integer=s.n-1) =
  eltype(s)[mcse(s, Val{:ipse}, i, maxlag) for i in r]

## Wrapper function for estimating Monte Carlo variance using various approaches

mcvar(v::AbstractArray, vtype::Symbol, args...) = mcvar(v, Val{vtype}, args...)

mcvar(s::VariableNState{Univariate}, vtype::Symbol, args...) = mcvar(s, Val{vtype}, args...)

mcvar(s::VariableNState{Multivariate}, vtype::Symbol, i::Integer, args...) = mcvar(s, Val{vtype}, i, args...)

mcvar(s::VariableNState{Multivariate}, vtype::Symbol, r::AbstractVector=1:s.size, args...) = mcvar(s, Val{vtype}, r, args...)

mcvar(v::AbstractArray, args...) = mcvar(v, Val{:imse}, args...)

mcvar(s::VariableNState{Univariate}, args...) = mcvar(s, Val{:imse}, args...)

mcvar(s::VariableNState{Multivariate}, i::Integer, args...) = mcvar(s, Val{:imse}, i, args...)

mcvar(s::VariableNState{Multivariate}, r::AbstractVector=1:s.size, args...) = mcvar(s, Val{:imse}, r, args...)

## Wrapper function for estimating Monte Carlo standard error using various approaches

mcse(v::AbstractArray, vtype::Symbol, args...) = mcse(v, Val{vtype}, args...)

mcse(s::VariableNState{Univariate}, vtype::Symbol, args...) = mcse(s, Val{vtype}, args...)

mcse(s::VariableNState{Multivariate}, vtype::Symbol, i::Integer, args...) = mcse(s, Val{vtype}, i, args...)

mcse(s::VariableNState{Multivariate}, vtype::Symbol, r::AbstractVector=1:s.size, args...) = mcse(s, Val{vtype}, r, args...)

mcse(v::AbstractArray, args...) = mcse(v, Val{:imse}, args...)

mcse(s::VariableNState{Univariate}, args...) = mcse(s, Val{:imse}, args...)

mcse(s::VariableNState{Multivariate}, i::Integer, args...) = mcse(s, Val{:imse}, i, args...)

mcse(s::VariableNState{Multivariate}, r::AbstractVector=1:s.size, args...) = mcse(s, Val{:imse}, r, args...)


text:
# This file is a part of Julia. License is MIT: https://julialang.org/license

module Test_LibGit2_https

using Test, LibGit2, NetworkOptions

# we currently use system SSL/TLS on macOS and Windows platforms
# and libgit2 cannot set the CA roots path on those systems
# if that changes, this may need to be adjusted
const CAN_SET_CA_ROOTS_PATH = !Sys.isapple() && !Sys.iswindows()

@testset "empty CA roots file" begin
    # these fail for different reasons on different platforms:
    # - on Apple & Windows you cannot set the CA roots path location
    # - on Linux & FreeBSD you you can but these are invalid files
    ENV["JULIA_SSL_CA_ROOTS_PATH"] = "/dev/null"
    @test_throws LibGit2.GitError LibGit2.ensure_initialized()
    ENV["JULIA_SSL_CA_ROOTS_PATH"] = tempname()
    @test_throws LibGit2.GitError LibGit2.ensure_initialized()
    # test that it still fails if called a second time
    @test_throws LibGit2.GitError LibGit2.ensure_initialized()
    if !CAN_SET_CA_ROOTS_PATH
        # test that this doesn't work on macOS & Windows
        ENV["JULIA_SSL_CA_ROOTS_PATH"] = NetworkOptions.bundled_ca_roots()
        @test_throws LibGit2.GitError LibGit2.ensure_initialized()
        delete!(ENV, "JULIA_SSL_CA_ROOTS_PATH")
        @test LibGit2.ensure_initialized() === nothing
    end
end

if CAN_SET_CA_ROOTS_PATH
    @testset "non-empty but bad CA roots file" begin
        # should still be possible to initialize
        ENV["JULIA_SSL_CA_ROOTS_PATH"] = joinpath(@__DIR__, "bad_ca_roots.pem")
        @test LibGit2.ensure_initialized() === nothing
    end
    mktempdir() do dir
        repo_url = "https://github.com/JuliaLang/Example.jl"
        @testset "HTTPS clone with bad CA roots fails" begin
            repo_path = joinpath(dir, "Example.HTTPS")
            c = LibGit2.CredentialPayload(allow_prompt=false, allow_git_helpers=false)
            redirect_stderr(devnull)
            err = try LibGit2.clone(repo_url, repo_path, credentials=c)
            catch err
                err
            end
            @test err isa LibGit2.GitError
            @test err.msg == "user rejected certificate for github.com"
        end
    end
end

end # module

ext: jl

max_stars_repo_name: jonas-schulze/julia

max_issues_repo_name: jonas-schulze/julia

max_forks_repo_name: jonas-schulze/julia


text:
"""
Program to calculate gradient and QFIM of quantum circuits.

Based on Julia and Yao.

Tobias Haug @ Imperial College London

t.haug@imperial.ac.uk

"""

using Yao
using LinearAlgebra, Plots

using Statistics
using Random

##warn: fidelity function of yao is lacking ^2 for some reason. Make sure to add yourself

@const_gate SQISWAP = ComplexF64[1 0 0 0;0 1/sqrt(2) -1/sqrt(2)*1.0im 0;0 -1/sqrt(2)*1.0im 1/sqrt(2) 0;0 0 0 1]

@const_gate InvSQISWAP = ComplexF64[1 0 0 0;0 1/sqrt(2) 1/sqrt(2)*1.0im 0;0 1/sqrt(2)*1.0im 1/sqrt(2) 0;0 0 0 1]

@const_gate RyHalf = mat(Ry(pi/2))
@const_gate InvRyHalf = mat(Ry(-pi/2))

@const_gate SqrtH = mat(Ry(pi/4))

@const_gate InvSqrtH = mat(Ry(-pi/4))

#NOTE: Yao internal function for daggered is very slow. Thus, we construct the conjugated circuit ourselves

#Yao: zero_state(n_qubits) |> circuit applies circuit onto state immdeiatly
# zero_state(n_qubits) => circuit delays application, is evaluated only when e.g. expecation values are needed
#use => when calculating gradients via expect'()

##set parameters for calculation
n_qubits=4 #number qubits
depth=10##2^(div(n_qubits,2)) #depth of circuit

rng_seed=1 #seed for random generator

circuit_type=0 #0: CPHASE with random single qubit rotations, 1: CNOT with random single qubit rotations 2: Sqrt[iswap] with random single qubit rotations

choose_param_type=0#0: random parameters, 1: fix to 0

cache_entangler=false #whether to cache entangling gates. Causes issues for constant gates, also seems to be slower except for cnot gates. Suggest to keep false.

epsilon_inv=10^-12#threshold of inversion of Gtensor, e.g. limits the size of the larget inverted eigenvalue. Use this to filter out eigenvalues that are actually zero, but rounded up due to numerical rounding errors
epsilon_inv_replace=0#1/epsilon_inv#if threshold is hit, replace inverse with this

use_yao_adjoint=false #set this false, else it will be very slow.. But take care to implement daggered properly first!

starttime=time()

println("")

function get_pauli_circuit(curr_depth,curr_qubit,circuit_spec,n_qubits)
    """
    get pauli operator from circuit definition. Required for the analytic gradient for the QFIM
    """
    element=circuit_spec[(curr_depth-1)*n_qubits+curr_qubit]

    if element==0
        return I2
    elseif element==1
        return X
    elseif element==2
        return Y
    elseif element==3
        return Z
    elseif element==5
        return 1/sqrt(2)*(X+Y)
    elseif element==6
        return 1/sqrt(2)*(Z+X)
    else
        error("not defined ",element)
    end

end

function get_pauli_op(pauli_string)
    """
    turns pauli string into operator
    """
    n_qubits=length(pauli_string)
    pauli_circuit = chain(n_qubits)
    for i in 1:length(pauli_string)
        if pauli_string[i]!=0
            if pauli_string[i]==1
                push!(pauli_circuit, put(i=>X))
            elseif pauli_string[i]==2
                push!(pauli_circuit, put(i=>Y))
            elseif pauli_string[i]==3
                push!(pauli_circuit, put(i=>Z))
            end
        end
    end

    return pauli_circuit
end

#creates entangling gate
#daggered whether to create adjoint of it or not
function entangler(entanglingGate,n_qubits,daggered)
    if daggered==false
        ent=chain(n_qubits,if type==0; control(ctrl, target=>Z) elseif type==1;control(ctrl, target=>X) elseif type==2;put((ctrl,target)=>SQISWAP)   end for (ctrl,target,type) in entanglingGate)

    else#make sure its properly daggered, i.e. also reverse order and conjugate
        ent=chain(n_qubits,if type==0; control(ctrl, target=>Z) elseif type==1;control(ctrl, target=>X) elseif type==2;put((ctrl,target)=>InvSQISWAP)  end for (ctrl,target,type) in entanglingGate[length(entanglingGate):-1:1])
    end
    return ent
end

function build_circuit(n_qubits, nlayers, entanglingGates,circuit_spec,daggered=false)
    """
    construct circuit according to specifications
    """
    circuit = chain(n_qubits)
    param_type=[]

    column=1
    if(daggered==false)
        counter=0
        #H squareroot, use matblock(mat()) to fix parameters. There may be better way to do this
        push!(circuit, chain(n_qubits, put(i=>SqrtH) for i = 1:n_qubits))
        #push!(circuit, cache(chain(n_qubits, put(i=>H) for i = 1:n_qubits)))
        countEntanglingLayers=0
        for i in 1:nlayers
            countEntanglingLayers+=1
            push!(circuit, chain(n_qubits, put(p=>chain(RotationGate(get_pauli_circuit(i,p,circuit_spec,n_qubits),0))) for p = 1:n_qubits))

            counter+=n_qubits

            entgl=entangler(entanglingGates[countEntanglingLayers],n_qubits,daggered)

            #whether to cache entanglers. Seems to cause issues for constant gates
            if cache_entangler==true
                push!(circuit, cache(entgl))
            else
                push!(circuit, entgl)
            end

        end
    else #daggered version of circuit
        counter=length(circuit_spec)

        countEntanglingLayers=nlayers+1
        for i in nlayers:-1:1
            countEntanglingLayers-=1
            #do the daggered version of the entangler
            entgl=entangler(entanglingGates[countEntanglingLayers],n_qubits,daggered)

            if cache_entangler==true
                push!(circuit, cache(entgl))
            else
                push!(circuit, entgl)
            end

            #daggered rotations
            push!(circuit, chain(n_qubits, put(p=>chain(RotationGate(-1*get_pauli_circuit(i,p,circuit_spec,n_qubits),0))) for p = n_qubits:-1:1))

            counter-=n_qubits

        end
        #H squareroot, use matblock(mat()) to fix parameters. There may be better way to do this
        push!(circuit, chain(n_qubits, put(i=>InvSqrtH) for i = n_qubits:-1:1))

    end

    return circuit
end

function get_HC_from_pauli(weight,pauli_strings)
    ab=[]
    for i in 1:length(pauli_strings)
        push!(ab,weight[i]*get_pauli_op(pauli_strings[i]))
    end

    return sum(ab)

end

function calculate_circuit(qcbm,cost_op,circuit_spec,qcbm_daggered)
    n_qubits=nqubits(qcbm)

    function get_loss(qcbm,hc)
        reg=zero_state(n_qubits) |>qcbm
        loss= expect(hc, reg)|> real
        return loss
    end

    function get_gradient(qcbm,hc)
        ##do_gtensor_here is just additional switch to calculate QFIM in case qng==0
        reg=zero_state(n_qubits) => qcbm #whenver get gradients use =>
        g_reg,pure_gradient=expect'(hc, reg)
        return pure_gradient
    end

    function get_QFIM(qcbm,qcbm_daggered)

        n_parameters=nparameters(qcbm)
        Gtensor=zeros(n_parameters,n_parameters)

        expect_single=zeros(depth,n_qubits)
        #do diagonal entries

        for j in 1:depth
            sub_qcbm=qcbm[1:2*j-1]
            reg=zero_state(n_qubits) |> sub_qcbm #|> immediatley evaluate register, instead of => (is evaluated only when expect is calculated)

            for k in 1:n_qubits
                op_p=get_pauli_circuit(j,k,circuit_spec,n_qubits)
                op_y=1/2*put(n_qubits, k=>op_p)
                expct= expect(op_y, reg)|> real
                expect_single[j,k]=expct
            end

            expect_double=zeros(n_qubits,n_qubits)
            for k in 1:n_qubits
                for m in k:n_qubits
                    if k!=m #if same, resolves to identity for paulis
                        op_pk=get_pauli_circuit(j,k,circuit_spec,n_qubits)
                        op_pm=get_pauli_circuit(j,m,circuit_spec,n_qubits)
                        op_yy=1/4*chain(n_qubits, [put(k=>op_pk),put(m=>op_pm)] )
                    else
                        op_yy=1/4*igate(n_qubits)
                    end
                    expct= expect(op_yy, reg)|> real
                    expect_double[k,m]=expct
                end

            end

            #geometric tensor components needed are real hermitian
            for k in 1:n_qubits
                for m in k+1:n_qubits
                    expect_double[m,k]=expect_double[k,m]
                end
            end

            for k in 1:n_qubits
                for m in 1:n_qubits
                    Gtensor[(j-1)*n_qubits+k,(j-1)*n_qubits+m]=expect_double[k,m]-expect_single[j,k]*expect_single[j,m]
                end
            end

        end

        #gradients for depth n>j
        #build gradients for U_n=R_y(theta_n)*Cphase, U_0=R_y(pi/4)=qcbm[1]
        #<0|U_0^\dag U_1^\dag U_j^\dag U_2 (\partial_n U_n^\dag) (U_3^\dag U_3) U_n U_2 U_j (\partial_j U_j)U_1 U_0|0>
        for j in 1:depth
            sub_qcbm=qcbm[1:2*j-1]
            reg=zero_state(n_qubits) |> sub_qcbm #|> immediatley evaluate register, instead of => (is evaluated only when expect is calculated)
            for n in j+1:depth#min(j,depth)#depth#min(j+1,depth)

                right_qcbm=qcbm[2*j:2*n-1]#goes from depth j to beginning of depth n

                #get sub circuit from daggered circuit. DO not use daggered function since it is horrible slow..
                if use_yao_adjoint==true
                    left_qcbm=Daggered(qcbm[2*j:2*n-1]) #goes from depth n to depth j as adjoint, Daggered is very slow!
                else
                    circuit_length=length(qcbm_daggered)
                    left_qcbm=qcbm_daggered[circuit_length+1-(2*n-1):circuit_length+1-2*j] #goes from depth n to depth j as adjoint
                end

                for k in 1:n_qubits
                    for m in 1:n_qubits
                        op_pjk=get_pauli_circuit(j,k,circuit_spec,n_qubits)
                        op_pnm=get_pauli_circuit(n,m,circuit_spec,n_qubits)
                        op_y_k=1/2*put(n_qubits, k=>op_pjk)#sits at qubit k at depth j
                        op_y_m=1/2*put(n_qubits, m=>op_pnm) # sits at qubit m at  depth n
                        op=left_qcbm*op_y_m*right_qcbm*op_y_k

                        expct= expect(op, reg)|> real #only real part needed for gtensor
                        Gtensor[(j-1)*n_qubits+k,(n-1)*n_qubits+m]=expct-expect_single[j,k]*expect_single[n,m]

                    end
                end
            end
        end
        for j in 1:depth
            for n in j+1:depth
                for k in 1:n_qubits
                    for m in 1:n_qubits
                        Gtensor[(n-1)*n_qubits+m,(j-1)*n_qubits+k]=Gtensor[(j-1)*n_qubits+k,(n-1)*n_qubits+m]
                    end
                end
            end
        end

        return Gtensor
    end

    loss=get_loss(qcbm,cost_op)
    gradient= get_gradient(qcbm,cost_op)
    QFIM= get_QFIM(qcbm,qcbm_daggered)

    return loss,gradient,QFIM
end

rng_param = MersenneTwister(rng_seed);
rng_circ= MersenneTwister(rng_seed+1);

if circuit_type==0
    entangler_type=0 #cphase
elseif circuit_type==1
    entangler_type=1 #CNOT
elseif circuit_type==2
    entangler_type=2#sqrt(iswap)
end

entanglingGates=[]
tempEnt=vcat([[2*i-1, ((2*i-1+1-1) %(n_qubits))+1,entangler_type] for i = 1:fld(n_qubits,2)],
[[2*i, ((2*i+1-1) %(n_qubits))+1,entangler_type] for i = 1:fld(n_qubits-1,2)])
entanglingGates=[tempEnt for p =1:depth]

#x,y,z rotations randomly
circuit_spec=[rand(rng_circ,1:3) for i in 1:depth*n_qubits]

n_var_parameters=depth*n_qubits

if(choose_param_type==0)
    var_param=rand(rng_param,n_var_parameters)*2*pi #list of random parameters
elseif(choose_param_type==1)
    var_param=zeros(rng_param,n_var_parameters) #list of 0 parameters
end

hilbertspace=2^n_qubits

##get Hamiltonian. Is the ZZ pauli operator acting on first 2 qubits
hamilton_paulis=[]
n_terms=1
hc_pauli_weight=ones(n_terms)
hc_pauli_strings=[zeros(Int,n_qubits) for i in 1:n_terms]
hc_pauli_strings[1][1:2]=[3,3]

hamilton_paulis=hc_pauli_strings
cost_op=get_HC_from_pauli(hc_pauli_weight,hamilton_paulis)

println("Nparameters ",n_var_parameters)

qcbm = build_circuit(n_qubits, depth,entanglingGates,circuit_spec)
dispatch!(qcbm,var_param) # initialize the parameters for circuit
##construct inverse of circuit. This is needed for QFIM
qcbm_daggered=build_circuit(n_qubits, depth,entanglingGates,circuit_spec,true)#daggered circuit
dispatch!(qcbm_daggered, var_param[length(var_param):-1:1]) # initialize the parameters for daggered circuit

loss,gradient, QFIM = calculate_circuit(qcbm, cost_op,circuit_spec,qcbm_daggered)

println("Loss ",loss)

println("Gradient mean ",mean(gradient))
println("Gradient variance ",var(gradient))

##2D plot of the QFIM
gr()
data = QFIM
fig2D1=heatmap(1:size(data,1),
    1:size(data,2), data,clim=(-0.25,0.25),
    c=cgrad([:blue, :white,:red]),
    xlabel="params", ylabel="params")
display(fig2D1)

#get eigenvalues of the QFIM
eigenvalues=eigvals(QFIM)

count_threshold=0
e_val_threshold=[] #eigenvalues greater than threshold
for i in 1:length(eigenvalues)
    if eigenvalues[i]>=epsilon_inv
        push!(e_val_threshold,eigenvalues[i])
    end
end

println("Non-zero eigenvalues of QFIM ",length(e_val_threshold))

println("Total time ",time()-starttime)

max_stars_repo_name: txhaug/quantum-geometry

max_issues_repo_name: txhaug/quantum-geometry

max_forks_repo_name: txhaug/quantum-geometry


text:
#------------------------------------------------------------------------
# Structure
#------------------------------------------------------------------------
# ModulusModel contains all JuMP related information
# AMARI contains optimization information
# SteinMinimaxEstimator contains the fit result

#------------------------------------------------------------------------
# Call structure for fitting AMARI
#------------------------------------------------------------------------
#  StatsBase.fit(::AMARI, target, Zs; initialize=true)
#  -- initialize_method(method, target, Zs) -> returns AMARI
#     -- initialize_modulus_model(method, modulus_model, target, δ1) -> returns AbstractModulusModel
#  -- fit_initialized!(::AMARI, target, Zs; kwargs...) -> returns SteinMinimaxEstimator
#     -- set_target!(::AbstractModulusModel, target)
#     -- set_δ!(::AbstractModulusModel, δ)

abstract type AbstractModulusModel end

Base.@kwdef struct ModulusModelWithF <: AbstractModulusModel
    method
    model
    g1
    g2
    f1
    f2
    f_sqrt
    estimated_marginal_density
    Δf
    δ_max
    δ_up
    bound_delta
    target
end

Base.@kwdef struct ModulusModelWithoutF <: AbstractModulusModel
    method
    model
    g1
    g2
    δ_max
    δ_up
    bound_delta
    target
    discretizer
    representative_eb_samples
end

function Base.show(io::IO, model::AbstractModulusModel)
    println(io, "Modulus model with target:")
    show(io, model.target)
end

function get_δ(Δf)
    norm(JuMP.value.(Δf))
end

# TODO: recalculate δ for general case?
function get_δ(model::AbstractModulusModel)
      JuMP.value(model.δ_up)
end

function get_δ(model::ModulusModelWithF; recalculate_δ = false)
    if recalculate_δ
        get_δ(model.Δf)
    else
      JuMP.value(model.δ_up)
    end
end

"""
    DeltaTuner

Abstract type used to represent ways of picking
``\\delta`` at which to solve the modulus problem, cf.
Manuscript. Different choices of ``\\delta`` correspond
to different choices of the Bias-Variance tradeoff with
every choice leading to Pareto-optimal tradeoff.
"""
abstract type DeltaTuner end

abstract type BiasVarAggregate <: DeltaTuner end

function get_bias_var(modulus_model::AbstractModulusModel)
    @unpack model = modulus_model
    δ = get_δ(modulus_model)
    ω_δ = objective_value(model)
    ω_δ_prime = -JuMP.dual(modulus_model.bound_delta)
    max_bias = (ω_δ - δ*ω_δ_prime)/2
    unit_var_proxy = ω_δ_prime^2
    max_bias, unit_var_proxy
end

function (bv::BiasVarAggregate)(modulus_model::AbstractModulusModel)
    bv(get_bias_var(modulus_model)...)
end

"""
    RMSE(n::Integer) <: DeltaTuner

A `DeltaTuner` to optimizes
the worst-case (root) mean squared error.  Here `n` is the sample
size used for estimation.
"""
struct RMSE{N} <: BiasVarAggregate
    n::N
end

RMSE() = RMSE(DataBasedDefault())
(rmse::RMSE)(bias, unit_var_proxy) =  sqrt(bias^2 + unit_var_proxy/rmse.n)

function Empirikos._set_defaults(rmse::RMSE, Zs; hints)
    RMSE(nobs(Zs))
end

"""
    HalfCIWidth(n::Integer, α::Float64) <: DeltaTuner

A `DeltaTuner` that chooses the `δ ≧ δ_min` the optimizes
the worst-case confidence interval width.  Here `n` is the sample
size used for estimation.
"""
Base.@kwdef struct HalfCIWidth{N} <: BiasVarAggregate
    n::N = DataBasedDefault()
    α::Float64 = 0.05
end

function Empirikos._set_defaults(halfci::HalfCIWidth, Zs; hints)
    @set halfci.n = nobs(Zs)
end

function (half_ci::HalfCIWidth)(bias, unit_var_proxy)
    se = sqrt(unit_var_proxy/half_ci.n)
    gaussian_ci(se; maxbias=bias, α=half_ci.α)
end

"""
    AMARI(convexclass::Empirikos.ConvexPriorClass,
          flocalization::Empirikos.FLocalization,
          solver,
          plugin_G = KolmogorovSmirnovMinimumDistance(convexclass, solver))

Affine Minimax Anderson-Rubin intervals for empirical Bayes estimands.
Here `flocalization` is a  pilot [`Empirikos.FLocalization`](@ref), `convexclass` is
a [`Empirikos.ConvexPriorClass`](@ref), `solver` is a JuMP.jl compatible solver.
`plugin_G` is a [`Empirikos.EBayesMethod`](@ref) used as an initial estimate of the marginal
distribution of the i.i.d. samples ``Z``.
"""
Base.@kwdef struct AMARI{N, G, M, EB}
    convexclass::G
    flocalization::N
    solver
    discretizer = nothing#DataBasedDefault()
    plugin_G = KolmogorovSmirnovMinimumDistance(convexclass, solver)
    data_split = :none
    delta_grid = 0.2:0.5:6.7
    delta_objective = RMSE()
    modulus_model::M = ModulusModelWithoutF
    representative_eb_samples::EB = nothing
    n = nothing
end

function Base.show(io::IO, amari::AMARI)
    print(io, "AMARI with")
    print(io, "  F-Localization: ")
    show(io, amari.flocalization)
    print(io, "\n")
    print(io, "            𝒢: ")
    show(io, amari.convexclass)
end

function initialize_modulus_model(method::AMARI, ::Type{ModulusModelWithF}, target::Empirikos.LinearEBayesTarget, δ)

    #TODO: perhaps this can also be moved..
    @unpack discretizer, solver, convexclass, flocalization, representative_eb_samples = method

    length(representative_eb_samples.vec) == 1 ||
            throw(ArgumentError("ModulusModelWithF only works for Homoskedastic samples."))
    Z = first(representative_eb_samples.vec)
    if isa(method.plugin_G, Distribution) #TODO SPECIAL CASE this elsewhere?
        estimated_marginal_density = Empirikos.dictfun(discretizer, Z, z-> pdf(method.plugin_G, z))
    else
        estimated_marginal_density = Empirikos.dictfun(discretizer, Z, z-> pdf(method.plugin_G.prior, z))
    end

    model = Model(solver)

    g1 = Empirikos.prior_variable!(model, convexclass)
    g2 = Empirikos.prior_variable!(model, convexclass)

    Empirikos.flocalization_constraint!(model, flocalization, g1)
    Empirikos.flocalization_constraint!(model, flocalization, g2)

    Zs = collect(keys(estimated_marginal_density))
    f_sqrt = sqrt.(collect(values(estimated_marginal_density)))

    f1 = pdf.(g1, Zs)
    f2 = pdf.(g2, Zs)

    @variable(model, δ_up)

    Δf = @expression(model, (f1 - f2)./f_sqrt)

    @constraint(model, pseudo_chisq_constraint,
           [δ_up; Δf] in SecondOrderCone())

    @objective(model, Max, target(g2) - target(g1))

    @constraint(model, bound_delta, δ_up <= δ)

    ModulusModelWithF(method=method, model=model, g1=g1, g2=g2, f1=f1, f2=f2,
        f_sqrt=f_sqrt, estimated_marginal_density=estimated_marginal_density,
        Δf=Δf, δ_max=Inf, δ_up=δ_up,
        bound_delta=bound_delta, target=target)
end

function modulus_cholesky_factor(convexclass::AbstractSimplexPriorClass, plugin_G, discr,
            eb_samples::HeteroskedasticSamples)
    K = nparams(convexclass)
    chr = cholesky(zeros(K, K) + I)
    cache_vec = zeros(K)
    fill!(chr.factors, 0)
    cache_vec  = zeros(K)
    for _interval in discr.sorted_intervals
        for (z, pr) in zip(eb_samples.vec, eb_samples.probs)
            z = set_response(z, _interval)
            cache_vec .= sqrt(pr) .* exp.(logpdf.(components(convexclass), z) .- logpdf(plugin_G, z)/2)
            lowrankupdate!(chr, cache_vec)
        end
    end
    chr
end

function initialize_modulus_model(method::AMARI, ::Type{ModulusModelWithoutF}, target::Empirikos.LinearEBayesTarget, δ)

    @unpack flocalization, convexclass, discretizer, plugin_G, solver, representative_eb_samples = method

    model = Model(solver)

    g1 = Empirikos.prior_variable!(model, convexclass)
    g2 = Empirikos.prior_variable!(model, convexclass)

    Empirikos.flocalization_constraint!(model, flocalization, g1)
    Empirikos.flocalization_constraint!(model, flocalization, g2)

    @variable(model, δ_up)

    K_chol = modulus_cholesky_factor(convexclass, plugin_G, discretizer, representative_eb_samples)
    KΔg = @expression(model, K_chol.U * (g1.finite_param - g2.finite_param))

    @constraint(model, pseudo_chisq_constraint,
           [δ_up; KΔg] in SecondOrderCone())

    @objective(model, Max, target(g2) - target(g1))
    @constraint(model, bound_delta, δ_up <= δ)

    ModulusModelWithoutF(method=method, model=model, g1=g1, g2=g2,
        δ_max=Inf, δ_up=δ_up,
        bound_delta=bound_delta, target=target,
        discretizer = discretizer, representative_eb_samples = representative_eb_samples
        )
end

function set_δ!(modulus_model::AbstractModulusModel, δ)
    set_normalized_rhs(modulus_model.bound_delta, δ)
    optimize!(modulus_model.model)
    modulus_model
end

function set_target!(modulus_model::AbstractModulusModel, target::Empirikos.LinearEBayesTarget)
    @unpack model, g1, g2 = modulus_model
    @objective(model, Max, target(g2) - target(g1))
    modulus_model = @set modulus_model.target = target
    optimize!(model)
    modulus_model
end

function default_support_discretizer(Zs::AbstractVector{<:AbstractNormalSample})
    _low,_up = quantile(response.(Zs), (0.02, 0.98))
    _step = mean( std.(Zs))/100
    interval_discretizer(range(_low; stop=_up, step=_step))
end

function initialize_method(method::AMARI, target::Empirikos.LinearEBayesTarget, Zs; kwargs...)

    method = Empirikos.set_defaults(method, Zs; kwargs...)
    fitted_floc = StatsBase.fit(method.flocalization, Zs; kwargs...)
    if isa(method.plugin_G, Distribution) #TODO SPECIAL CASE this elsewhere?
        fitted_plugin_G = method.plugin_G
    else
        fitted_plugin_G = StatsBase.fit(method.plugin_G, Zs; kwargs...)
    end
    discr = method.discretizer #TODO SPECIAL CASE for ::Distribution
    if isnothing(discr)
        discr = default_support_discretizer(Zs)
        method = @set method.discretizer = discr
    end
    modulus_model = method.modulus_model

    method = @set method.representative_eb_samples = heteroskedastic(Zs)

    method = @set method.flocalization = fitted_floc
    method = @set method.plugin_G = fitted_plugin_G

    n = nobs(Zs)
    method = @set method.n = n

    @unpack delta_grid = method

    δ1 = delta_grid[1]
    modulus_model = initialize_modulus_model(method, modulus_model, target, δ1)
    method = @set method.modulus_model = modulus_model
    method
end

Base.@kwdef mutable struct SteinMinimaxEstimator{M, T, D}
    target::T
    δ::Float64
    ω_δ::Float64
    ω_δ_prime::Float64
    g1
    g2
    Q::D
    max_bias::Float64
    unit_var_proxy::Float64
    modulus_model::M
    method
    δs = zeros(Float64,5)
    δs_objective = zeros(Float64, length(δs))
end

Base.@kwdef struct QDonoho{G,H,T,D}
    g1::G
    g2::G
    plugin_G::H
    mult::T
    offset::T
    discretizer::D
end

function (Q::QDonoho)(Z::EBayesSample)
    @unpack g1, g2, plugin_G, mult, offset, discretizer  = Q
    Z = discretizer(Z)
    _g1_val = exp(logpdf(g1, Z) - logpdf(plugin_G, Z))
    _g2_val = exp(logpdf(g2, Z) - logpdf(plugin_G, Z))
    offset + mult*(_g2_val - _g1_val)
end

function SteinMinimaxEstimator(modulus_model::ModulusModelWithoutF)
    @unpack model, method, target, discretizer, representative_eb_samples  = modulus_model
    @unpack convexclass, plugin_G = method

    discretizer == method.discretizer || throw("Internal discretizer modified.")

    δ = get_δ(modulus_model)
    ω_δ = objective_value(model)
    ω_δ_prime = -JuMP.dual(modulus_model.bound_delta)

    g1 = modulus_model.g1()
    g2 = modulus_model.g2()

    L1 = target(g1)
    L2 = target(g2)

    offset_sum = sum(discretizer.sorted_intervals) do _int
        sum(zip(representative_eb_samples.vec, representative_eb_samples.probs)) do (z,pr)
            z = set_response(z, _int)
            f2_z = pdf(g2, z)
            f1_z = pdf(g1, z)
            barf_z = pdf(plugin_G, z)
            pr*(f2_z - f1_z) * (f2_z + f1_z) / barf_z
        end
    end

    offset =  (L1+L2)/2 - ω_δ_prime/(2*δ)*offset_sum
    mult = ω_δ_prime/δ

    Q = QDonoho(;g1 = g1, g2 = g2, plugin_G = plugin_G,
                offset = offset, mult = mult,
                discretizer = discretizer)

    max_bias = (ω_δ - δ*ω_δ_prime)/2
    unit_var_proxy = ω_δ_prime^2

    SteinMinimaxEstimator(
              target=target,
              δ=δ,
              ω_δ=ω_δ,
              ω_δ_prime=ω_δ_prime,
              g1=g1,
              g2=g2,
              Q=Q,
              max_bias=max_bias,
              unit_var_proxy=unit_var_proxy,
              method=method,
              modulus_model=modulus_model
        )
end

function SteinMinimaxEstimator(modulus_model::ModulusModelWithF)
    @unpack model, method, target, estimated_marginal_density = modulus_model
    @unpack convexclass = method

    δ = get_δ(modulus_model)
    ω_δ = objective_value(model)
    ω_δ_prime = -JuMP.dual(modulus_model.bound_delta)

    g1 = modulus_model.g1()
    g2 = modulus_model.g2()

    L1 = target(g1)
    L2 = target(g2)

    Zs = collect(keys(estimated_marginal_density))

    f1 = pdf.(g1, Zs)
    f2 = pdf.(g2, Zs)

    f̄s = collect(values(estimated_marginal_density))

    Q = ω_δ_prime/δ*(f2 .- f1)./f̄s
    Q_0  = (L1+L2)/2 -
        ω_δ_prime/(2*δ)*sum( (f2 .- f1).* (f2 .+ f1) ./ f̄s)

    Q = Empirikos.dictfun(method.discretizer, Zs, Q .+ Q_0)

    max_bias = (ω_δ - δ*ω_δ_prime)/2
    unit_var_proxy = ω_δ_prime^2

    SteinMinimaxEstimator(
              target=target,
              δ=δ,
              ω_δ=ω_δ,
              ω_δ_prime=ω_δ_prime,
              g1=g1,
              g2=g2,
              Q=Q,
              max_bias=max_bias,
              unit_var_proxy=unit_var_proxy,
              method=method,
              modulus_model=modulus_model
        )
end

function fit_initialized!(method::AMARI, target, Zs; kwargs...)
    @unpack modulus_model, delta_grid, delta_objective, n = method

    modulus_model = set_target!(modulus_model, target)

    δs_objective = zeros(Float64, length(delta_grid))

    for (index, δ) in enumerate(delta_grid)
        set_δ!(modulus_model, δ/sqrt(n))
        δs_objective[index] = delta_objective(modulus_model)
    end

    if length(delta_grid) > 1 # resolve to get best objective
        idx_best = argmin(δs_objective)
        set_δ!(modulus_model, delta_grid[idx_best]/sqrt(n))
    end

    sm = SteinMinimaxEstimator(modulus_model)
    sm = @set sm.δs = collect(delta_grid)
    sm = @set sm.δs_objective = δs_objective
    sm = @set sm.method = method
    sm
end

function confint(Q::SteinMinimaxEstimator, target, Zs; level=0.95)
    target == Q.modulus_model.target ||
           error("Target has changed")
    α = 1- level
    _bias = Q.max_bias
    _Qs = Q.Q.(Zs)
    _wts = StatsBase.weights(Zs)
    _se = std(Q.Q.(Zs), _wts; corrected=true)/sqrt(nobs(Zs))
    point_estimate = mean(Q.Q.(Zs), _wts)
    halfwidth = gaussian_ci(_se; maxbias=_bias, α=α)
    BiasVarianceConfidenceInterval(estimate = point_estimate,
                                   maxbias = _bias,
                                   se = _se,
                                   α = α, method = nothing, target = target)
end

function confint(method::AMARI, target::Empirikos.LinearEBayesTarget, Zs; initialize=true, kwargs...)
    _fit = StatsBase.fit(method, target, Zs; initialize=initialize)
    confint(_fit, target, Zs; kwargs...)
end

function StatsBase.fit(method::AMARI, target, Zs; initialize=true, kwargs...)
    if initialize
        method = initialize_method(method, target, Zs; kwargs...)
    end
    fit_initialized!(method::AMARI, target, Zs; kwargs...)
end

function Base.broadcasted(::typeof(confint), amari::AMARI,
    targets::AbstractArray{<:Empirikos.EBayesTarget}, Zs)

    init_target = isa(targets[1], LinearEBayesTarget) ? targets[1] : denominator(targets[1])
    method = initialize_method(amari, init_target, Zs)

    _ci =  confint(method, targets[1], Zs; initialize=false)
    confint_vec = fill(_ci, axes(targets))
    for (index, target) in enumerate(targets[2:end])
        confint_vec[index+1] = confint(method, target, Zs; initialize=false)
    end
    confint_vec
end

function Base.broadcasted_kwsyntax(::typeof(confint), amari::AMARI,
    targets::AbstractArray{<:Empirikos.EBayesTarget}, Zs; level=0.95)

    init_target = isa(targets[1], LinearEBayesTarget) ? targets[1] : denominator(targets[1])
    method = initialize_method(amari, init_target, Zs)

    _ci =  confint(method, targets[1], Zs; initialize=false, level=level)
    confint_vec = fill(_ci, axes(targets))
    for (index, target) in enumerate(targets[2:end])
        confint_vec[index+1] = confint(method, target, Zs; initialize=false, level=level)
    end
    confint_vec
end

"""
    StatsBase.confint(method::AMARI,
                      target::Empirikos.EBayesTarget,
                      Zs;
                      level=0.95)

Form a confidence interval for the [`Empirikos.EBayesTarget`](@ref) `target` with coverage
    `level` based on the samples `Zs` using the [`AMARI`](@ref) `method`.
"""
function confint(method::AMARI, target::Empirikos.AbstractPosteriorTarget, Zs;
                          initialize=true, level=0.95, kwargs...)
    if initialize
        init_target = Empirikos.PosteriorTargetNullHypothesis(target, 0.0)
        method = initialize_method(method, init_target, Zs; kwargs...)
    end

    floc_worst_case = FLocalizationInterval(flocalization = method.flocalization,
            convexclass = method.convexclass,
            solver= method.solver)

    α = 1 - level
    outer_ci = confint(floc_worst_case, target)
    outer_ci =  @set outer_ci.α = α

    c_lower = outer_ci.lower
    c_upper = outer_ci.upper

    target_lower = Empirikos.PosteriorTargetNullHypothesis(target, c_lower)
    target_upper = Empirikos.PosteriorTargetNullHypothesis(target, c_upper)

    fit_lower = fit_initialized!(method, target_lower, Zs) #SteinMinimax
    Q_lower = fit_lower.Q.(Zs)
    confint_lower = confint(fit_lower, target_lower, Zs; level=level)
    max_bias_lower = confint_lower.maxbias
    var_Q_lower = abs2(confint_lower.se)
    estimate_lower = confint_lower.estimate

    fit_upper = fit_initialized!(method, target_upper, Zs)
    Q_upper = fit_upper.Q.(Zs)
    confint_upper = confint(fit_upper, target_upper, Zs; level=level)
    max_bias_upper = confint_upper.maxbias
    var_Q_upper = abs2(confint_upper.se)
    estimate_upper = confint_upper.estimate

    _wts = StatsBase.weights(Zs)
    cov_lower_upper = cov([Q_lower Q_upper], _wts; corrected=true)[1,2] / nobs(Zs)

    bisection_pair = BisectionPair(var1 = var_Q_lower, max_bias1= max_bias_lower, estimate1= estimate_lower,
                         var2 = var_Q_upper, max_bias2 = max_bias_upper, estimate2= estimate_upper,
                         cov = cov_lower_upper)

    λs = range(0, stop=1, length=10_000)
    all_cis = confint.(Ref(bisection_pair), λs; α = α)
    zero_in_ci = first.(all_cis) .<= 0.0 .<= last.(all_cis)

    idx_lhs = findfirst(zero_in_ci)
    idx_rhs = findlast(zero_in_ci)

    if isnothing(idx_lhs)
        return(outer_ci)
    end

    λ_lhs = λs[idx_lhs]
    λ_rhs = λs[idx_rhs]

    c_lower_updated = (1-λ_lhs)*c_lower + λ_lhs*c_upper
    c_upper_updated = (1-λ_rhs)*c_lower + λ_rhs*c_upper

    LowerUpperConfidenceInterval(α=α, target=target, method=nothing,
                                 lower=c_lower_updated,
                                 upper=c_upper_updated)
end

# used right now only for sanity check in tests
function worst_case_bias_lp(fitted_amari::AMARI, Q::QDonoho, target; max=true)
    @unpack convexclass, solver, discretizer, flocalization, representative_eb_samples = fitted_amari

    transposed_intervals = reshape(discretizer.sorted_intervals, 1, length(discretizer.sorted_intervals))
    Zs = set_response.(representative_eb_samples.vec, transposed_intervals)
    model = Model(solver)

    g1 = prior_variable!(model, convexclass)
    flocalization_constraint!(model, flocalization, g1)

    Qs = Q.(Zs)
    f_G = @expression(model, pdf.(g1, Zs) .* representative_eb_samples.probs)

    @objective(model, Max, dot(f_G, Qs) - target(g1))
    optimize!(model)
    maxbias = JuMP.objective_value(model)

    @objective(model, Min, dot(f_G, Qs) - target(g1))
    optimize!(model)
    minbias = JuMP.objective_value(model)

    expected_var = dot(Qs.^2, pdf.(fitted_amari.plugin_G, Zs) .* representative_eb_samples.probs)

    (maxbias = maxbias, minbias = minbias, expected_var = expected_var)
end


text:
"""

```
writexyz(x::Vector{T}, file::String) where T <: AbstractVector
```

Print test xyz file.

"""
function writexyz(x::Vector{T}, file::String) where T <: AbstractVector
  f = open(file,"w")
  nx = length(x)
  println(f,nx)
  println(f,"title")
  for i in 1:nx
    println(f,"H $(x[i][1]) $(x[i][2]) $(x[i][3])")
  end
  close(f)
  nothing
end


text:
@testset "Thinning" begin
    # for □ like figure
    img = falses(20,20)
    for i=7:7+4, j=7:7+4
        img[i,j] = img[j,i] = true
    end
    thin = thinning(img)
    @test count(thin .== 1) == 1
    @test thin[9,9] == 1

    # for + like figure
    img = falses(20,20)
    for i=8:8+5, j=4:13+4
        img[i,j] = true
        img[j,i] = true
    end
    ans = falses(size(img))
    ans[[111,131,151,171,191,192,193,194,195,211,231,251,271]] .= 1
    @test thinning(img) == ans

    img = Bool.([1 1 1 1 1 1
                 0 0 0 0 0 0
                 1 1 1 1 1 1
                 1 1 1 1 1 1
                 0 0 1 1 0 0
                 0 0 0 0 0 0])
    ans = Bool.([1 1 1 1 1 1
                 0 0 0 0 0 0
                 0 0 0 0 0 0
                 1 1 1 1 1 0
                 0 0 0 0 0 0
                 0 0 0 0 0 0])
    @test thinning(img) == ans

    # already thinned
    img = Bool.([0 0 0 0 0 0
                 0 0 0 1 1 0
                 0 0 1 0 0 0
                 0 0 1 0 0 0
                 0 0 0 1 0 0
                 0 0 0 0 0 0
                 0 0 0 0 0 0])
    @test thinning(img) == img
end


text:
struct CUGResult
    obsval::Real
    randval::AbstractVector
    iter::Int
end

function cugtest_edges(g::AbstractGraph, f::Function, iter::Integer)

    e = ne(g); n = nv(g)
    tv = f(g)

    res = zeros(iter)
    for i in 1:iter
        gi = erdos_renyi(n, e)
        res[i] = f(gi)
    end

    return CUGResult(tv, res, iter)
end

function plot(t::CUGResult; xlabel="Test statistic")
    histogram(t.randval, label="Replications", xlabel=xlabel)
    vline!([t.obsval], label="Observed", linewidth=2, color=:red, linestyle=:dash)
end

max_stars_repo_name: jfb-h/EIIndex

max_issues_repo_name: jfb-h/EIIndex

max_forks_repo_name: jfb-h/EIIndex


text:
module  ChernovA1P2

#this is useful if we have mdoules
if pwd() ∉ LOAD_PATH
  push!(LOAD_PATH,pwd())
end

using DataFrames, Distributions, StatsBase, GZip, JLD,
  Gadfly, CTModCopy, JuMP, NLopt, ForwardDiff

const DATA_PATH = pwd() * "\\data" #path where we store the data
const OUTPUT_PATH = pwd() * "\\output"
const DATE_FORMAT_STR = "m/d/y"
const SP500_NAME = "GSPC"

const DEF_NUM_MOMENTS = 6

function preProcessSP500()::Void
  dateFormat::String = DATE_FORMAT_STR
  sp500DF::DataFrame = readtable("$DATA_PATH\\$SP500_NAME.csv")

  #fix the dates
  rename!(sp500DF, :Date, :DATE)
  sp500DF[:,:DATE] = ((s::String)->Date(s, dateFormat)::Date).(sp500DF[:,:DATE])::DataVector{Date}
  sort!(sp500DF, cols = [:DATE])
  sp500DF[:,:returns] = similar(sp500DF[:,:Adj_Close])
  sp500DF[2:end,:returns] = log.(sp500DF[2:end,:Adj_Close] ./  sp500DF[1:(end-1),:Adj_Close])
  #showcols(sp500DF)

  stream::IOStream = open("$DATA_PATH\\$SP500_NAME.jls", "w")
  serialize(stream, sp500DF)
  close(stream)

  return nothing
end

#this function gets the sample moments of a given vector of data
function getSampleMoments!(x::Vector{Float64}, N::Int, moments::Vector{Float64} = Vector{Float64}(N))

  #for each moment we are getting
  for i::Int ∈ 1:N
    moments[i] = sum(x.^Float64(i))/Float64(length(x))
  end

  println(moments)
  println("Var: ")

  return moments
end

function getCumulantsOfJumpNorm(α::Real, σ²::Real, θ::Real, δ²::Real, ω::Real)::Vector{Real}
  κ::Vector{Real} = Vector{Real}(DEF_NUM_MOMENTS)

  #δ²::Real = δ*δ #helpers
  θ²::Real = θ*θ
  θ³::Real = θ*θ²

  κ[1] = α + ω*θ
  κ[2] = σ²+ω*(δ²+θ²) #σ^2 +ω*(δ^2+θ^2)
  κ[3] = 3.*ω*δ²*θ+ω*θ³ #3*ω*δ^2*θ+ω*θ^3
  κ[4] = ω*(6.*δ²*θ²+θ²*θ²+3.*δ²*δ²) #6*ω*δ^2*θ^2+ω*θ^4 +3*ω*δ^4
  κ[5] = ω*(15.*θ*δ²*δ²+10.*δ²*θ³+θ³*θ²) #θ*ω*(15*δ^4+10*δ^2*θ^2+θ^4)
  κ[6] = ω*(15.*δ²*δ²*δ²+45.*δ²*δ²*θ²+15.*δ²*θ²*θ²+θ³*θ³) #ω*(15*δ^6+45*δ^4*θ^2+15*δ^2*θ^4+θ^6)

  return κ
end

function get6CumulantMoments(κ...)::Vector{Real}
  μ::Vector{Real} = Vector{Real}(DEF_NUM_MOMENTS)

  κ1²::Real = κ[1] * κ[1]
  κ1³::Real = κ[1] * κ1²
  κ2²::Real = κ[2] * κ[2]

  μ[1] = κ[1]
  μ[2] = κ[2] + κ1² #κ[2]+κ[1]^2
  μ[3] = κ[3] + 3.*κ[2]*κ[1] + κ1³
  μ[4] = κ[4] + 4.*κ[3]*κ[1] + 3.*κ2² + 6.*κ[2]*κ1²+κ1²*κ1²
  μ[5] = κ[5] + 5.*κ[4]*κ[1] + 10.*κ[3]*κ[2] + 10.*κ[3]*κ1² + 15.*κ2² *κ[1] + 10.*κ[2]*κ1³+κ1³*κ1²
  μ[6] = κ[6] + 6.*κ[5]*κ[1] + 15.*κ[4]*κ[2] + 15.*κ[4]*κ1² + 10.*κ[3]*κ[3]+60*κ[3]*κ[2]*κ[1] +
    20.*κ[3]*κ1³+15.*κ2²*κ[2] + 45.*κ2²*κ1² + 15.*κ[2]*κ1²*κ1²+κ1³*κ1³

  return μ
end

function getParametersOfJumpNorm(x::Vector{Float64})::Void

  numMoments::Int = DEF_NUM_MOMENTS

  #get the sample moments
  M::Vector{Float64} = getSampleMoments!(x, numMoments)

  #this is our objective function
  function obj(α::Real, σ²::Real, θ::Real, δ²::Real, ω::Real)::Real

    κ::Vector{Real} = getCumulantsOfJumpNorm(α, σ², θ, δ², ω)
    μ::Vector{Real} = get6CumulantMoments(κ...)
    d::Real = sum((M .- μ) .* (M .- μ))

    return d
  end

  αBest::Float64=0.0; σ²Best::Float64=0.0; θBest::Float64=0.0
  δ²Best::Float64=0.0; ωBest::Float64=0.0; dBest::Float64=100.0

  solveFails::Int = 0
  solveAttempts::Int = 0
  failed::Bool = false

  #encapsulate the routing for searching a grid
  #Algorithms: :LD_LBFGS, :LN_COBYLA, :LN_BOBYQA
  #  :LD_SLSQP, :LD_TNEWTON, :LD_VAR2/:LD_VAR1,
  #  :LD_MMA, :GN_ESCH, :GN_DIRECT, GN_DIRECT_L, :GD_STOGO
  # GN_ISRES, :LN_NELDERMEAD, :LN_SBPLX,
  #GN_CRS2_LM
  function solveProblem(αMin::Float64, σ²Min::Float64, θMin::Float64, δ²Min::Float64, ωMin::Float64, Δ::Float64)
    solveAttempts += 1
    mod = Model(solver=NLoptSolver(algorithm=:LD_SLSQP, maxtime=1))

    #JuMP.build(mod)
    #maxtime!(internalmodel(mod), 1.0)

    @variable(mod, αMin <= α <= αMin + Δ)
    @variable(mod, σ²Min <= σ² <= σ²Min + Δ)
    @variable(mod, θMin <= θ <= θMin + Δ)
    @variable(mod, δ²Min <= δ² <= δ²Min + Δ)
    @variable(mod, ωMin <= ω <= ωMin + Δ)

    #register the outside functions
    JuMP.register(mod, :getCumulantsOfJumpNorm, 5, getCumulantsOfJumpNorm, autodiff=true)
    JuMP.register(mod, :get6CumulantMoments, 6, get6CumulantMoments, autodiff=true)
    JuMP.register(mod, :obj, 5, obj, autodiff=true)

    @NLobjective(mod, Min, obj(α, σ², θ, δ², ω))

    #set the starting point at the midpoint of the grid
    setvalue(α , (αMin + Δ/2.0))
    setvalue(σ², (σ²Min + Δ/2.0))
    setvalue(θ , (θMin + Δ/2.0))
    setvalue(δ², (δ²Min + Δ/2.0))
    setvalue(ω , (ωMin + Δ/2.0))

    #indicator
    failed = false
    try #solve the model
      output = solve(mod)
    catch
      solveFails+=1
      failed = true
    end

    objValue::Float64 = obj(getvalue(α), getvalue(σ²), getvalue(θ), getvalue(δ²), getvalue(ω))
    if !failed && objValue<dBest #if both statements are true, update the best parameters
      dBest = objValue
      αBest, σ²Best, θBest, δ²Best, ωBest = getvalue(α), getvalue(σ²), getvalue(θ), getvalue(δ²), getvalue(ω)
    end
  end

  #iterate the parameter space
  Δ::Float64=0.1
  @fastmath for αTry::Float64 ∈ 0.0:Δ:(0.1-Δ),
     σ²Try::Float64 ∈ 0.0:Δ:(0.2-Δ),
      θTry::Float64 ∈ -1.0:Δ:(0.0-Δ),
      δ²Try::Float64 ∈ 0.0:Δ:(0.2-Δ),
      ωTry::Float64 ∈ 0.0:Δ:(1.0-Δ)
        solveProblem(αTry, σ²Try, θTry, δ²Try, ωTry, Δ)
  end

  #= Output code here
  println("Trys: $solveAttempts")
  println("Fails: $solveFails")
  println("α=$(αBest), σ²=$(σ²Best), θ=$(θBest), δ²=$(δ²Best), ω=$(ωBest)")
  println("Also: σ=$(σ²Best^0.5) annualized: $((σ²Best)^0.5*255.0^0.5)")
  println("Objective: $dBest")
  println("For Reference: x̄=$(mean(x)), std(x)= $(std(x))")
  =#

  return nothing
end

#the main method for running the optimziation routines
function problemScript(;refreshData::Bool = true)::Void
  if !isfile("$DATA_PATH\\$SP500_NAME.jls") || refreshData
    preProcessSP500()
  end

  #read the binary file
  stream::IOStream = open("$DATA_PATH\\$SP500_NAME.jls")
  sp500DF::DataFrame = deserialize(stream)
  close(stream)

  x::Vector{Float64} = sp500DF[2:end, :returns]
  getParametersOfJumpNorm(x)

  return nothing
end

@time problemScript()
end

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
using ProjAssistant
@quickactivate

# ------------------------------------------------------------------
@time begin
    import MAT

    import SparseArrays

    import Chemostat_Folsom2014
    const ChF = Chemostat_Folsom2014

    const iJR = ChF.iJR904
    const Fd = ChF.FolsomData # experimental data
    const Bd = ChF.BegData    # cost data

    import Chemostat
    const Ch = Chemostat
    const ChU = Ch.Utils
    const ChSS = Ch.SteadyState
    const ChLP = Ch.LP

end

## ------------------------------------------------------------------
# LOAD RAW MODEL
src_file = rawdir(iJR, "iJR904.mat")
mat_model = MAT.matread(src_file)["model"]
model = ChU.MetNet(mat_model; reshape=true)
ChU.tagprintln_inmw("MAT MODEL LOADED",
    "\nfile:             ", relpath(src_file),
    "\nfile size:        ", filesize(src_file), " bytes",
    "\nmodel size:       ", size(model),
    "\nChU.nzabs_range:      ", ChU.nzabs_range(model.S),
)
ChF.test_fba(model, iJR.BIOMASS_IDER; summary = false)

## -------------------------------------------------------------------
# Set bounds
# The abs maximum bounds will be set to 100
ChU.tagprintln_inmw("CLAMP BOUNDS",
    "\nabs max bound: ", iJR.ABS_MAX_BOUND
)
foreach(model.rxns) do ider
        ChU.isfixxed(model, ider) && return # fixxed reaction are untouched

        old_ub = ChU.ub(model, ider)
        new_ub = old_ub == 0.0 ? 0.0 : iJR.ABS_MAX_BOUND
        ChU.ub!(model, ider, new_ub)

        old_lb = ChU.lb(model, ider)
        new_lb = old_lb == 0.0 ? 0.0 : -iJR.ABS_MAX_BOUND
        ChU.lb!(model, ider, new_lb)
end

## -------------------------------------------------------------------
# CLOSING EXCHANGES
exchs = ChU.exchanges(model)
ChU.tagprintln_inmw("CLOSE EXCANGES",
    "\nChU.exchanges: ", exchs |> length
)
# Close, for now, all ChU.exchanges for avoiding it to be in revs
# The reversible reactions will be splited for modeling cost
# Exchanges have not associated cost, so, we do not split them
foreach(exchs) do idx
    ChU.ub!(model, idx, 0.0) # Closing all outtakes
    ChU.lb!(model, idx, 0.0) # Closing all intakes
end

## -------------------------------------------------------------------
# ENZYMATIC COST INFO
# The cost will be introduced as a reaction, we follow the same cost models as
# Beg et al. (2007): https://doi.org/10.1073/pnas.0609845104.
# A new balance equations is then added:
#        Σ(rᵢ*costᵢ) + tot_cost = 0
#    Because the cost coefficients (costᵢ) < 0 (it resamble a reactant), the system must allocate
#    the fluxes (rᵢ) so that Σ(rᵢ*costᵢ) = tot_cost, and tot_cost
#    are usually bounded [0.0, 1.0]
cost_info = Dict()
fwd_ider(rxn) = string(rxn, ChU.FWD_SUFFIX);
bkwd_ider(rxn) = string(rxn, ChU.BKWD_SUFFIX);
for rxn in model.rxns
    # The ChU.exchanges, the atpm and the biomass are synthetic reactions, so,
    # they have should not have an associated enzimatic cost
    any(startswith.(rxn, ["EX_", "DM_"])) && continue
    rxn == iJR.BIOMASS_IDER && continue
    rxn == iJR.BIOMASS_IDER && continue
    rxn == iJR.ATPM_IDER && continue

    # Only the internal, non reversible reactions have an associated cost
    # We will split the rev reactions, so we save the cost for both versions (fwd, bkwd)
    if ChU.isrev(model, rxn)
        cost_info[fwd_ider(rxn)] = -iJR.beg_enz_cost(rxn)
        cost_info[bkwd_ider(rxn)] = -iJR.beg_enz_cost(rxn)
    else
        cost_info[rxn] = -iJR.beg_enz_cost(rxn)
    end
end

## -------------------------------------------------------------------
# SPLITING REVS
ChU.tagprintln_inmw("SPLITING REVS",
    "\nfwd_suffix:      ", ChU.FWD_SUFFIX,
    "\nbkwd_suffix:     ", ChU.BKWD_SUFFIX,
)
model = ChU.split_revs(model;
    get_fwd_ider = fwd_ider,
    get_bkwd_ider = bkwd_ider,
);

## -------------------------------------------------------------------
# ADDING COST REACCION
cost_met_id = "cost"
cost_exch_id = iJR.COST_IDER
ChU.tagprintln_inmw("ADDING COST",
    "\ncosts to add: ", cost_info |> length,
    "\nmin abs coe:  ", cost_info |> values .|> abs |> minimum,
    "\nmax abs coe:  ", cost_info |> values .|> abs |> maximum,
    "\ncost met id:  ", cost_met_id,
    "\ncost exch id: ", cost_exch_id
)

M, N = size(model)
cost_met = ChU.Met(cost_met_id, S = collect(values(cost_info)), rxns = collect(keys(cost_info)), b = 0.0)
model = ChU.expanded_model(model, M + 1, N + 1)
ChU.set_met!(model, ChU.findempty(model, :mets), cost_met)
cost_exch = ChU.Rxn(cost_exch_id, S = [1.0], mets = [cost_met_id], lb = -iJR.ABS_MAX_BOUND, ub = 0.0, c = 0.0)
ChU.set_rxn!(model, ChU.findempty(model, :rxns), cost_exch);

## -------------------------------------------------------------------
# SET BASE EXCHANGE
ChU.tagprintln_inmw("SETTING EXCHANGES")
# To control the intakes just the metabolites defined in the
# base_intake_info (The minimum medium) will be opened.
# The base model will be constraint as in a cultivation with
# experimental minimum xi
# see Cossios paper (see README)

foreach(exchs) do idx
    ChU.ub!(model, idx, iJR.ABS_MAX_BOUND) # Opening all outakes
    ChU.lb!(model, idx, 0.0) # Closing all intakes
end

# see Cossios paper (see README) for details in the Chemostat bound constraint
xi = minimum(Fd.val(:xi))
intake_info = iJR.load_base_intake_info()
ChSS.apply_bound!(model, xi, intake_info; emptyfirst = true)

# tot_cost is the exchange that controls the bounds of the
# enzimatic cost contraint, we bound it to [0, 1.0]
ChU.lb!(model, cost_exch_id, 0.0);
ChU.ub!(model, cost_exch_id, 1.0);

## -------------------------------------------------------------------
model = ChU.fix_dims(model)
ChF.test_fba(model, iJR.BIOMASS_IDER, iJR.COST_IDER)

## -------------------------------------------------------------------
# FVA PREPROCESSING
MODELS_FILE = procdir(iJR, "base_models.bson")
const BASE_MODELS = ldat(MODELS_FILE) do
    Dict{Any, Any}("base_model" => ChU.compressed_model(model))
end

## -------------------------------------------------------------------
let
    for (exp, D) in Fd.val(:D) |> enumerate

        DAT = get!(BASE_MODELS, "fva_models", Dict())
        ChU.tagprintln_inmw("DOING FVA",
            "\nexp:             ", exp,
            "\nD:               ", D,
            "\ncProgress:       ", length(DAT),
            "\n"
        )
        haskey(DAT, exp) && continue # cached

        ## -------------------------------------------------------------------
        # prepare model
        model0 = deepcopy(model)
        M, N = size(model0)
        exp_xi = Fd.val(:xi, exp)
        intake_info = iJR.intake_info(exp)
        ChSS.apply_bound!(model0, exp_xi, intake_info;
            emptyfirst = true)

        ChF.test_fba(exp, model0, iJR.BIOMASS_IDER, iJR.COST_IDER)
        fva_model = ChLP.fva_preprocess(model0,
            check_obj = iJR.BIOMASS_IDER,
            verbose = true
        );
        ChF.test_fba(exp, fva_model, iJR.BIOMASS_IDER, iJR.COST_IDER)

        # storing
        DAT[exp] = ChU.compressed_model(fva_model)

        ## -------------------------------------------------------------------
        # caching
        sdat(BASE_MODELS, MODELS_FILE);
        GC.gc()
    end
end

## -------------------------------------------------------------------
# MAX MODEL
let
    # This model is bounded by the maximum rates found for EColi.
    # Data From:
    # Varma, (1993): 2465–73. https://doi.org/10.1128/AEM.59.8.2465-2473.1993.
    # Extract max exchages from FIG 3 to form the maximum polytope

    ChU.tagprintln_inmw("DOING MAX MODEL", "\n")

    max_model = deepcopy(model)

    # Biomass
    # 2.2 1/ h
    ChU.bounds!(max_model, iJR.BIOMASS_IDER, 0.0, 2.2)

    Fd_rxns_map = iJR.load_rxns_map()
    # 40 mmol / gDW h
    ChU.bounds!(max_model, Fd_rxns_map["GLC"], -40.0, 0.0)
    # 45 mmol/ gDW
    ChU.bounds!(max_model, Fd_rxns_map["AC"], 0.0, 40.0)
    # 55 mmol/ gDW h
    ChU.bounds!(max_model, Fd_rxns_map["FORM"], 0.0, 55.0)
    # 20 mmol/ gDW h
    ChU.bounds!(max_model, Fd_rxns_map["O2"], -20.0, 0.0)

    # fva
    max_model = ChLP.fva_preprocess(max_model,
        check_obj = iJR.BIOMASS_IDER,
        verbose = true
    );

    ## -------------------------------------------------------------------
    test_model = deepcopy(max_model)
    for exp in 1:4
        D = Fd.val(:D, exp)
        cgD_X = Fd.cval(:GLC, exp) * Fd.val(:D, exp) / Fd.val(:X, exp)
        ChU.lb!(test_model, iJR.EX_GLC_IDER, -cgD_X)
        fbaout = ChLP.fba(test_model, iJR.BIOMASS_IDER, iJR.COST_IDER)
        biom = ChU.av(test_model, fbaout, iJR.BIOMASS_IDER)
        cost = ChU.av(test_model, fbaout, iJR.COST_IDER)
        @info("Test", exp, cgD_X, D, biom, cost); println()
    end

    ## -------------------------------------------------------------------
    # saving
    BASE_MODELS["max_model"] = ChU.compressed_model(max_model)
    sdat(BASE_MODELS, MODELS_FILE)
end


text:
using ArgParse

include("functions_align.jl")
include("functions_deletions.jl")

function parse_commandline()
    s = ArgParseSettings()

    @add_arg_table! s begin
        "--infile", "-i"
            help = "the alignment to type, in fasta format"
            required = true
        "--csv-out"
            help = "the name of csv file to write with deletion genotypes"
            default = "out.csv"
        "--dels-file"
            help = "file of deletions to type and include as variants when collapsing"
            required = true
        "--fasta-out"
            help = "the name of fasta file to write if --append-as-snp is invoked"
            default = "out.fasta"
        "--append-as-snp"
            action = "store_true"
            help = "if invoked, append the deletion genotypes as SNPs in the output alignment"
    end

    return parse_args(s)
end

function main()
    parsed_args = parse_commandline()

    dels = read_del_file(parsed_args["dels-file"])

    nucleotide_array, fasta_IDs = populate_byte_array_get_names(parsed_args["infile"])

    if parsed_args["append-as-snp"]
        nucleotide_array_dels, genotypes_table = type_deletions_and_append_as_SNP(nucleotide_array, dels)

        open(parsed_args["csv-out"], "w") do io
            println(io, "sequence," * join([join(x, "_") for x in dels], ","))
            for i in 1:size(genotypes_table, 2)
                println(io, fasta_IDs[i] * "," * join(genotypes_table[:,i], ","))
            end
        close(io)
        end

        open(parsed_args["fasta-out"], "w") do io
            for i in 1:size(nucleotide_array_dels, 2)
                println(io, ">" * fasta_IDs[i])
                println(io, get_seq_from_1D_byte_array(nucleotide_array_dels[:,i]))
            end
        close(io)
        end

    else
        genotypes_table = type_deletions(nucleotide_array, dels)

        open(parsed_args["csv-out"], "w") do io
            println(io, "sequence," * join([join(x, "_") for x in dels], ","))
            for i in 1:size(genotypes_table, 2)
                println(io, fasta_IDs[i] * "," * join(genotypes_table[:,i], ","))
            end
        close(io)
        end
    end
end

main()

max_stars_repo_name: cov-ert/julialign

max_issues_repo_name: cov-ert/julialign

max_forks_repo_name: cov-ert/julialign


text:
"""
pde_lin_test(n)

Test the linear operators.
"""
function pde_lin_test(n)
    h = 1.0 / (n + 1)
    x = collect(h:h:1.0-h)
    fdata = fishinit(n)
    z = rand(n, n)
    L2d = Lap2d(n)
    DX = Dx2d(n)
    DY = Dy2d(n)
    lapok = lap_test(n, x, z, L2d, fdata)
    discok = disc_test(n, x, L2d, DX, DY, fdata)
    lapok && discok
end

"""
disc_test()
Have I broken the discretizations?
"""
function disc_test(n, x, L2d, DX, DY, fdata)
    n2 = n * n
    ue = solexact(x)
    ux = dxexact(x)
    uy = dyexact(x)
    D2u = l2dexact(x)
    ue1 = reshape(ue, (n2,))
    uex1 = reshape(ux, (n2,))
    uey1 = reshape(uy, (n2,))
    ued21 = reshape(D2u, (n2,))
    # test DX
    dx21 = DX * ue1
    dxerr = norm(dx21 - uex1, Inf)
    # test DY
    dy21 = DY * ue1
    dyerr = norm(dy21 - uey1, Inf)
    # test Laplacian
    du21 = L2d * ue1
    d2err = norm(du21 - ued21, Inf)
    pass = (d2err < 0.75) && (dxerr < 0.1) && (dyerr < 1.e-12)
    pass || println("Discretization test fails")
    return pass
end

"""
lap_test()

Does the FFT invert the discrete Laplacian?
Does the discrete Laplacian pass a simple eigenvalue test.
"""
function lap_test(n, x, z, L2d, fdata)
    randok = rand_test(z, n, L2d, fdata)
    eigok = eig_test(n, x, fdata)
    pass = randok && eigok
    return pass
end

function rand_test(z, n, L2d, fdata)
    n2 = n * n
    z1 = reshape(z, (n2,))
    y1 = L2d * z1
    y = reshape(y1, (n, n))
    mz = fish2d(y, fdata)
    q = reshape(mz, (n2,))
    pass = (norm(q - z1, Inf) < 1.e-12)
    pass || println("rand_test fails, norm =", norm(q - z1))
    return pass
end

function eig_test(n, x, fdata)
    lambda = pi * pi * 5
    efunx = sin.(pi * x)
    efuny = sin.(2 * pi * x)
    efunu = efunx * efuny'
    vfun = fish2d(efunu, fdata)
    pass = (norm(lambda * vfun - efunu, Inf) < 1.e-2)
    pass || println("eig test fails")
    return pass
end


text:
function build_square_circle(
    sbp_order,
    Nqr,
    Nqs,
    gDfun,
    gDdotfun,
    gNfun,
    body_force,
    bc_map,
    do_output = false;
    rho = 1,
    cxx = nothing,
    cxy = nothing,
    cyy = nothing,
)
    @assert Nqr == Nqs

    # EToV defines the element by its vertices
    # EToF defines element by its four faces, in global face number
    # FToB defines whether face is Dirichlet (1), Neumann (2), frictional interface (7),
    # interior jump (8),  or just an interior interface (0)
    # EToDomain is 1 if element is inside circle; 2 otherwise
    (verts, EToV, EToF, FToB, EToDomain) =
        read_inp_2d("square_circle.inp"; bc_map = bc_map)

    # Plot the original connectivity before mesh warping
    # plot_connectivity(verts, EToV)

    # This is needed to fix up points that should be on the boundary of the
    # circle, but the mesh didn't quite put them there
    for v in 1:size(verts, 2)
        x, y = verts[1, v], verts[2, v]
        if abs(hypot(x, y) - 1) < 1e-5
            Q = atan(y, x)
            verts[1, v], verts[2, v] = cos(Q), sin(Q)
        end
    end

    # Determine secondary arrays
    # FToE : Unique Global Face to Element Number
    #        (the i'th column of this stores the element numbers that share the
    #        global face number i)
    # FToLF: Unique Global Face to Element local face number
    #        (the i'th column of this stores the element local face numbers that
    #        shares the global face number i)
    # EToO : Element to Unique Global Faces Orientation
    #        (the i'th column of this stores the whether the element and global
    #        face are oriented in the same way in physical memory or need to be
    #        rotated)
    # EToS : Element to Unique Global Face Side
    #        (the i'th column of this stores whether an element face is on the
    #        plus side or minus side of the global face)
    (FToE, FToLF, EToO, EToS) = connectivityarrays(EToV, EToF)

    # number of elements and faces
    (nelems, nfaces) = (size(EToV, 2), size(FToB, 1))
    if do_output
        @show (nelems, nfaces)
    end

    # Dictionary to store the operators (independent of element/block)
    mets = create_metrics(
        sbp_order,
        11,
        11,
        (r, s) -> (r, ones(size(r)), zeros(size(s))),
        (r, s) -> (s, zeros(size(r)), ones(size(s))),
    )

    OPTYPE = typeof(
        rhsoperators(
            sbp_order,
            12,
            12,
            mets,
            gDfun,
            gDdotfun,
            gNfun,
            body_force;
            rho = rho,
        ),
    )

    METTYPE = typeof(mets)
    rhsops = Dict{Int64, OPTYPE}()
    metrics = Dict{Int64, METTYPE}()

    # Build the local operators for each block
    for e in 1:nelems

        # Get domain corners
        (x1, x2, x3, x4) = verts[1, EToV[:, e]]
        (y1, y2, y3, y4) = verts[2, EToV[:, e]]

        # Initialize the domain transformations as transfinite between the corners
        ex = [
            (α) -> x1 * (1 .- α) / 2 + x3 * (1 .+ α) / 2,
            (α) -> x2 * (1 .- α) / 2 + x4 * (1 .+ α) / 2,
            (α) -> x1 * (1 .- α) / 2 + x2 * (1 .+ α) / 2,
            (α) -> x3 * (1 .- α) / 2 + x4 * (1 .+ α) / 2,
        ]
        exα = [
            (α) -> -x1 / 2 + x3 / 2,
            (α) -> -x2 / 2 + x4 / 2,
            (α) -> -x1 / 2 + x2 / 2,
            (α) -> -x3 / 2 + x4 / 2,
        ]
        ey = [
            (α) -> y1 * (1 .- α) / 2 + y3 * (1 .+ α) / 2,
            (α) -> y2 * (1 .- α) / 2 + y4 * (1 .+ α) / 2,
            (α) -> y1 * (1 .- α) / 2 + y2 * (1 .+ α) / 2,
            (α) -> y3 * (1 .- α) / 2 + y4 * (1 .+ α) / 2,
        ]
        eyα = [
            (α) -> -y1 / 2 + y3 / 2,
            (α) -> -y2 / 2 + y4 / 2,
            (α) -> -y1 / 2 + y2 / 2,
            (α) -> -y3 / 2 + y4 / 2,
        ]

        # For blocks on the circle, put in the curved edge transform
        if abs(FToB[EToF[1, e]]) == BC_JUMP_INTERFACE
            error("curved face 1 not implemented yet")
        end
        if abs(FToB[EToF[2, e]]) == BC_JUMP_INTERFACE
            error("curved face 2 not implemented yet")
        end
        if abs(FToB[EToF[3, e]]) == BC_JUMP_INTERFACE
            Q1 = atan(y1, x1)
            Q2 = atan(y2, x2)
            if !(-π / 2 < Q1 - Q2 < π / 2)
                Q2 -= sign(Q2) * 2 * π
            end
            ex[3] = (α) -> cos.(Q1 * (1 .- α) / 2 + Q2 * (1 .+ α) / 2)
            ey[3] = (α) -> sin.(Q1 * (1 .- α) / 2 + Q2 * (1 .+ α) / 2)
            β3 = (Q2 - Q1) / 2
            exα[3] = (α) -> -β3 .* sin.(Q1 * (1 .- α) / 2 + Q2 * (1 .+ α) / 2)
            eyα[3] = (α) -> +β3 .* cos.(Q1 * (1 .- α) / 2 + Q2 * (1 .+ α) / 2)
        end
        if abs(FToB[EToF[4, e]]) == BC_JUMP_INTERFACE
            Q3 = atan(y3, x3)
            Q4 = atan(y4, x4)
            if !(-π / 2 < Q3 - Q4 < π / 2)
                error("curved face 4 angle correction not implemented yet")
            end
            ex[4] = (α) -> cos.(Q3 * (1 .- α) / 2 + Q4 * (1 .+ α) / 2)
            ey[4] = (α) -> sin.(Q3 * (1 .- α) / 2 + Q4 * (1 .+ α) / 2)
            β4 = (Q4 - Q3) / 2
            exα[4] = (α) -> -β4 .* sin.(Q3 * (1 .- α) / 2 + Q4 * (1 .+ α) / 2)
            eyα[4] = (α) -> +β4 .* cos.(Q3 * (1 .- α) / 2 + Q4 * (1 .+ α) / 2)
        end

        # Create the volume transform as the transfinite blending of the edge
        # transformations
        xt(r, s) = transfinite_blend(
            ex[1],
            ex[2],
            ex[3],
            ex[4],
            exα[1],
            exα[2],
            exα[3],
            exα[4],
            r,
            s,
        )
        yt(r, s) = transfinite_blend(
            ey[1],
            ey[2],
            ey[3],
            ey[4],
            eyα[1],
            eyα[2],
            eyα[3],
            eyα[4],
            r,
            s,
        )

        metrics[e] = create_metrics(sbp_order, Nqr - 1, Nqs - 1, xt, yt)

        if isnothing(cxx) && isnothing(cyy) && isnothing(cyy)
            crr = metrics[e].crr
            crs = metrics[e].crs
            css = metrics[e].css
        else
            J = metrics[e].J
            (x, y) = metrics[e].coord
            rx = metrics[e].rx
            ry = metrics[e].ry
            sx = metrics[e].sx
            sy = metrics[e].sy
            cxx_ = cxx.(x, y)
            cyy_ = cyy.(x, y)
            cyx_ = cxy_ = cxy.(x, y)
            crr =
                J .* (
                    rx .* cxx_ .* rx +
                    rx .* cxy_ .* ry +
                    ry .* cyx_ .* rx +
                    ry .* cyy_ .* ry
                )
            crs =
                J .* (
                    rx .* cxx_ .* sx +
                    rx .* cxy_ .* sy +
                    ry .* cyx_ .* sx +
                    ry .* cyy_ .* sy
                )
            css =
                J .* (
                    sx .* cxx_ .* sx +
                    sx .* cxy_ .* sy +
                    sy .* cyx_ .* sx +
                    sy .* cyy_ .* sy
                )
        end

        # Linear operators:
        rhsops[e] = rhsoperators(
            sbp_order,
            Nqr,
            Nqs,
            metrics[e],
            gDfun,
            gDdotfun,
            gNfun,
            body_force;
            rho = rho,
            crr = crr,
            crs = crs,
            css = css,
        )
    end

    do_output && plot_blocks(metrics, abs.(FToB), EToF)

    return (
        metrics = metrics,
        rhsops = rhsops,
        EToDomain = EToDomain,
        EToF = EToF,
        EToO = EToO,
        EToS = EToS,
        FToB = FToB,
        FToE = FToE,
        FToLF = FToLF,
    )
end


text:
# notebook will not be generated
y = 1 + 2


text:
#===============================================================================

    AmpModelResults.jl

    output from fit_model!()

    Author: Tom Price
    Date:   July 2019

===============================================================================#

import DataStructures.OrderedDict
import Ipopt: IpoptSolver #, NLoptSolver

#===============================================================================
    Field definitions >>
===============================================================================#

## issue: rename `rbbs_ary3` as `calibrated` once juliaapi_new has been updated
const AMPLONGMODELRESULTS_FIELD_DEFS = [
    Field(:rbbs_ary3,        Vector{Float_T}),
    Field(:bl_fit,           Union{AmpModelFit,Symbol}),
    Field(:bl_notes,         Vector{String}),
    Field(:blsub_fluos,      Vector{Float_T}),
    Field(:blsub_fluos_flb,  Vector{Float_T}),
    Field(:quant_fit,        AmpModelFit),
    Field(:quant_status,     Symbol),
    Field(:coefs,            Vector{Float_T}),
    Field(:d0,               Float_T),
    Field(:quant_fluos,      Vector{Float_T}),
    Field(:dr1_pred,         Vector{Float_T}),
    Field(:dr2_pred,         Vector{Float_T}),
    Field(:dr1_pred1,        Vector{Float_T}),
    Field(:dr2_pred1,        Vector{Float_T}),
    Field(:max_dr1,          Float_T),
    Field(:max_dr2,          Float_T),
    Field(:cyc_vals_4cq,     OrderedDict{Symbol,Float_T}),
    Field(:eff_vals_4cq,     OrderedDict{Symbol,Float_T}),
    Field(:cq_raw,           Float_T),
    Field(:cq,               Float_T),
    Field(:eff,              Float_T),
    Field(:cq_fluo,          Float_T)]

const AMPSHORTMODELRESULTS_FIELDNAMES = [
    # rbbs_ary3,
    :blsub_fluos,
    :blsub_fluos_flb,
    :dr1_pred,
    :dr2_pred,
    :dr1_pred1,
    :dr2_pred1,
    :quant_status,
    :coefs,
    :cq,
    :d0]

const AMPCQFLUOMODELRESULTS_FIELDNAMES = [
    :quant_status,
    :cq_fluo]

#===============================================================================
    struct and constructor generation >>
===============================================================================#

abstract type AmpModelResults end

SCHEMA = AMPLONGMODELRESULTS_FIELD_DEFS
@make_struct_from_SCHEMA AmpLongModelResults AmpModelResults
@make_constructor_from_SCHEMA AmpLongModelResults

SCHEMA = subset_schema(
    AMPLONGMODELRESULTS_FIELD_DEFS,
    AMPSHORTMODELRESULTS_FIELDNAMES)
@make_struct_from_SCHEMA AmpShortModelResults AmpModelResults
@make_constructor_from_SCHEMA AmpShortModelResults

SCHEMA = subset_schema(
    AMPLONGMODELRESULTS_FIELD_DEFS,
    AMPCQFLUOMODELRESULTS_FIELDNAMES)
@make_struct_from_SCHEMA AmpCqFluoModelResults AmpModelResults
@make_constructor_from_SCHEMA AmpCqFluoModelResults

max_stars_repo_name: MakerButt/chaipcr

max_issues_repo_name: MakerButt/chaipcr

max_forks_repo_name: MakerButt/chaipcr


text:
"""
	tree_approximation!(newtree::Tree, path::Function, nIterations::Int64, p::Int64=2, r::Int64=2)

Returns a valuated probability scenario tree approximating the input stochastic process.

Args:
- *newtree* - Tree with a certain branching structure,

- *path* - function generating samples from the stochastic process to be approximated,

- *nIterations* - number of iterations for stochastic approximation procedure,

- *p* - choice of norm (default p = 2 (Euclidean distance)), and,

- *r* - transportation distance parameter
"""
function tree_approximation!(newtree::Tree{A,B,C,D}, path::Function, nIterations::Int64, p::Int64=2, r::Int64=2) where {A,B,C,D}
    leaf, ~, probaLeaf = leaves(newtree)      # leaves, indexes and probabilities of the leaves of the tree
    T = height(newtree)                            # height of the tree = number of stages - 1
    n = length(leaf)                               # number of leaves = no of omegas
    d = zeros(Float64, length(leaf))
    samplepath = zeros(Float64, T+1)           # T + 1 = the number of stages in the tree.
    probaLeaf = zero(probaLeaf)
    probaNode = nodes(newtree)                                # all nodes of the tree
    path_to_leaves = [root(newtree, i) for i in leaf]         # all the paths from root to the leaves
    path_to_all_nodes = [root(newtree, j) for j in probaNode] # all paths to other nodes
    @showprogress 100 "Computing..." for k = 1 : nIterations

        # Critical == 0 if n >4!!! This is almost always the case
        critical = max(0.0, 0.2 * sqrt(k) - 0.1 * n)
        #tmp = findall(xi -> xi <= critical, probaLeaf)

        tmp = Int64[inx for (inx, ppf) in enumerate(probaLeaf) if ppf <= critical]
        samplepath .= vec(path())  # a new trajectory to update the values on the nodes

        #The following part addresses the critical probabilities of the tree so that we don't loose the branches
        if !isempty(tmp) && !iszero(tmp)
            probaNode = zero(probaNode)
            probaNode[leaf] = probaLeaf
            for i = leaf
                while newtree.parent[i] > 0
                    probaNode[newtree.parent[i]] = probaNode[newtree.parent[i]] + probaNode[i]
                    i = newtree.parent[i]
                end
            end
            for tmpi = tmp
                rt = path_to_leaves[tmpi]
                #tmpi = findall(pnt -> pnt <= critical, probaNode[rt])
                tmpi = Int64[ind for (ind, pnt) in enumerate(probaNode[rt]) if pnt <= critical]
                newtree.state[rt[tmpi]] .= samplepath[tmpi]
            end
        end
        #To the step  of STOCHASTIC COMPUTATIONS
        endleaf = 0 #start from the root
        for t = 1 : T+1
            tmpleaves = newtree.children[endleaf + 1]
            disttemp = Inf #or fill(Inf,dm)
            for i = tmpleaves
                dist = norm(view(samplepath, 1 : t) - view(newtree.state, path_to_all_nodes[i]), p)
                if dist < disttemp
                    disttemp = dist
                    endleaf = i
                end
            end
        end
        #istar = findall(lf -> lf == endleaf, leaf)
        istar = Int64[idx for (idx, lf) in enumerate(leaf) if lf == endleaf]
        probaLeaf[istar] .+= 1.0                                                            #counter  of probabilities
        StPath = path_to_leaves[endleaf - (leaf[1] - 1)]
        delta = newtree.state[StPath] - samplepath
        tmpNorm = sum(delta.^p).^(1/p)
        d[istar] .+= tmpNorm.^(r)
        delta .=  r .* tmpNorm.^(r - p) .* abs.(delta).^(p - 1) .* sign.(delta)
        ak = 1.0 ./ (30.0 .+ probaLeaf[istar]) #.^ 0.75        # step size function - sequence for convergence
        newtree.state[StPath] = newtree.state[StPath] - delta .* ak
    end

    probabilities = probaLeaf./sum(probaLeaf)   # this is faster than map()
	# probaLeaf = probabilities * nIterations?
	# this can be used for continuing approximation

    t_dist = (sum(d .* hcat(probabilities)) / nIterations) .^ (1 / r)
    newtree.name = "$(newtree.name) with d=$(t_dist) at $(nIterations) iterations"
    newtree.probability .= build_probabilities!(newtree, hcat(probabilities)) #build the probabilities of this tree

    return newtree
end

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

"""
	lattice_approximation(bstructure::Array{Int64,1}, path::Function, nIterations::Int64, r::Int64 = 2, dim::Int64 = 1)
Returns a valuated approximated lattice for the stochastic process provided in any dimension. The default dimension is dim = 1 for a lattice in 1 dimension.

Args:
- bstructure - Branching structure of the scenario lattice e.g., bstructure = [1,2,3,4,5] represents a 5-staged lattice
- path - Function generating samples from a known distribution with length equal to the length of bstructure of the lattice.
- nIterations - Number of iterations to be performed.
- r - Parameter for the transportation distance.
- dim - dimension of the lattice to be generated. This depends entirely on the dimension of the samples generated.
"""
function lattice_approximation(bstructure::Array{Int64,1}, path::Function, nIterations::Int64, r::Int64 = 2, dim::Int64 = 1)
    tdist = zeros(dim)         # multistage distance
    T = length(bstructure)     # number of stages in the lattice
    states = [zeros(bstructure[j], 1, dim) for j = 1 : T] # States of the lattice at each time t
    probabilities = vcat([zeros(bstructure[1], 1, dim)],[zeros(bstructure[j-1], bstructure[j], dim) for j = 2 : T]) # Probabilities of the lattice at each time t
    init_path = path()

    #Check the dimension of the sample path if it is equal to the dimension of the lattice.
    #if dim != size(init_path)[2]
    #    @error("Dimension of lattice ($dim) is not equal to the dimension of the input array ($(size(init_path)[2]))")
    #end
    # Initialize the states of the nodes of the lattice
    for t = 1 : T
        states[t] .= init_path[t]
    end
    Z = Array{Float64}(undef, T, dim) # Array to hold the new samples generated
    #Stochastic approximation step starts here
    @showprogress 100 "Computing..." for n = 1 : nIterations
        Z .= path() # Draw a new sample Gaussian path
        last_index = Int64.(ones(dim))
        dist = zeros(dim)
        for t = 1 : T
            for i = 1 : dim
                # corrective action to include lost nodes
                states[t][:, :, i][findall(sum(probabilities[t][:, :, i], dims = 2) .< 1.3 * sqrt(n) / bstructure[t])] .= Z[t,i]
                min_dist, new_index = findmin(abs.(vec(states[t][:, :, i] .- Z[t,i]))) # find the closest lattice entry
                dist[i] += min_dist^2  # Euclidean distance for the paths
                probabilities[t][last_index[i], new_index, i] += 1.0  # increase the counter for the nodes
                # use stochastic approximation algorithm to calculate for the new states of the nodes in the lattice.
                states[t][new_index, :, i] = states[t][new_index, :, i] - r * (min_dist)^(r-1) * (states[t][new_index, :, i] .- Z[t,i]) / ((Z[1] + 30 + n))
                last_index[i] = new_index
            end
        end
        # calculate the multistage distance
        dist = dist .^ (1/2); tdist = (tdist .* (n - 1) + dist .^r )/ n
        #tdist = (tdist*(n-1) + dist)/n
    end
    # calculate the probabilities
    probabilities = probabilities ./ nIterations

    # functions to print the results of the lattices
    # these functions depends on the dimension of the lattice
    if dim == 1
        # Round the results to make sure that they are neat.
        # Remember it is a 3D array, so a loop will do.
        for i=length(probabilities):-1:2
            for j=1:size(probabilities[i])[1]
                probabilities[i][j,:,1] = round.(probabilities[i][j,:,1] ./ sum(probabilities[i-1][:,j,1]) , digits=6)
            end
        end

        return Lattice("Approximated Lattice with $bstructure branching structure and distance=$(tdist .^ (1 / r)) at $(nIterations) iterations",
        [round.(states[i], digits = 6) for i = 1 : length(states)], [round.(probabilities[j], digits = 6) for j = 1 : length(probabilities)])
    end
end

function path_ident(trr, path::Function, nIterations::Int64, r::Int64=2, p::Int64=2)
    leaf, ~, probaLeaf = leaves(trr)      # leaves, indexes and probabilities of the leaves of the tree
    T = height(trr)                            # height of the tree = number of stages - 1
    n = length(leaf)                               # number of leaves = no of omegas
    d = zeros(Float64, length(leaf))
    samplepath = zeros(Float64, T+1)           # T + 1 = the number of stages in the tree.
    chosenPath = zeros(Int64, T+1)

    probaLeaf = zero(probaLeaf)
    probaNode = nodes(trr)                                # all nodes of the tree
    path_to_leaves = [root(trr, i) for i in leaf]         # all the paths from root to the leaves
    path_to_all_nodes = [root(trr, j) for j in probaNode] # all paths to other nodes
    for k = 1 : nIterations

       samplepath .= vec(path())  # a new trajectory to update the values on the nodes

        #To the step  of STOCHASTIC COMPUTATIONS
        endleaf = 0 #start from the root
        for t = 1 : T+1
            tmpleaves = trr.children[endleaf + 1]
            disttemp = Inf #or fill(Inf,dm)
            for i = tmpleaves
                dist = norm(view(samplepath, 1 : t) - view(trr.state, path_to_all_nodes[i]), p)
                if dist < disttemp
                    disttemp = dist
                    endleaf = i
                    chosenPath[t] = i
                end
            end
        end

    end
    return(chosenPath, trr.state[chosenPath])

end

function path_ident(lat::Lattice, path::Function, nIterations::Int64, r::Int64=2)
    tdist = 0.0
    T = length(lat.state)

    states = lat.state
    probabilities = lat.probability

    Z = Array{Float64}(undef, T) # Array to hold the new samples generated7
    chosenPath = Array{Float64}(undef,T)

    #Stochastic approximation step starts here
    for n = 1 : nIterations
        Z .= vec(path()) # Draw a new sample Gaussian path
        dist = 0
        for t = 1 : T
            # corrective action to include lost nodes
            min_dist, new_index = findmin(abs.(vec(states[t][:, :, 1] .- Z[t]))) # find the closest lattice entry
            dist += min_dist^2  # Euclidean distance for the paths
            chosenPath[t] = new_index
        end
        # calculate the multistage distance
	tdist = (tdist .* (n - 1) + dist .^(r/2) )/ n
    end

    return(chosenPath, [ states[t][Int(chosenPath[t]),1,1 ] for t=1:T]  )

end