text:
module DickeBCE

export hamiltonian_operator, husimi, husimi_of_coherent, QuantumDickeSystem, diagonalization, coherent_state,
    coherent_state!, coherent_overlap, eigenstate_parities, Wigner, evolve, WignerProjQP, WignerProjqp,
    factor_R_of_coherent_state, participation_ratio, random_state!, random_state, random_cₖ_generator,
    survival_probability,dimension,density_of_states,energy_width_of_coherent_state
    import ..ClassicalSystems
    using LinearAlgebra
    using ProgressMeter
    using SparseArrays
    using Distributed
    import ..PhaseSpaces
    using GSL
    using WignerSymbols
    using Distributed
    using CSV
    using ..ClassicalDicke
    using Distributions
    import Random
    using Statistics
    using DataFrames,Tables

    """
    ```julia
    mutable struct QuantumDickeSystem <: ClassicalDicke.DickeSystem
    ```
    This object represents the quantum Dicke model. It stores the parameters of
    the system, and it may be passed to multiple functions
    in this module. To generate it, use

    ```julia
    function QuantumDickeSystem(classical_system::ClassicalDicke.ClassicalDickeSystem;
        j::Real,
        Nmax::Union{Integer,Nothing}=nothing)
    ```
    or [`DickeBCE.QuantumDickeSystem(;ω₀,ω,γ,j,Nmax)`](@ref DickeBCE.QuantumDickeSystem()).

    Because this object subtypes [`ClassicalDicke.DickeSystem`](@ref), it may be passed
    to any function that requires a `DickeSystem` or a [`ClassicalSystem`](@ref ClassicalSystems.ClassicalSystem),
    and the underlying `classical_system` will be used.
    # Arguments
    - `classical_system` should be generated with [`ClassicalDicke.ClassicalDickeSystem`](@ref).
    - `j` is the value of ``j``. It must be a positive half-integer.
    - `Nmax - 1` is the maximum excitation of the modified bosonic sector in the efficient coherent
      basis (see [Bastarrachea2014PSa](@cite), [Bastarrachea2014PSb](@cite)). `Nmax` must be a positive integer.
      It may be omitted. In this case, you may use functions that do not require `Nmax`. Moreover, if
      you call [`diagonalization`](@ref DickeBCE.diagonalization) with a `system` that has `Nmax = nothing`,
      and there is a  diagonalization saved in disk, `diagonalization(system)` will load
      the eigenstates with the largest `Nmax` and will set the value of  `Nmax` in `system`.
    """
    mutable struct QuantumDickeSystem <: ClassicalDicke.DickeSystem
       j::Real
       Nmax::Union{Integer,Nothing}
       classical_system::ClassicalDicke.ClassicalDickeSystem
       function QuantumDickeSystem(classical_system::ClassicalDicke.ClassicalDickeSystem;
           j::Real,
           Nmax::Union{Integer,Nothing}=nothing)
           if !isinteger(2*j) || j < 0
               error("j must be a positive half-integer. Got j = $j.")
           end
           if Nmax!== nothing && Nmax <= 0
               error("Nmax must be a positive integer. Got Nmax = $Nmax.")
           end
           if isinteger(j)
               j= Int(j)
           end
           return new(j,Nmax,classical_system)
       end
    end
    function ClassicalSystems.parameters(system::QuantumDickeSystem)
        return ClassicalSystems.parameters(system.classical_system)
    end
    """
    ```julia
    function QuantumDickeSystem(;ω₀,ω,γ,j,Nmax=nothing)
    ```
    Shorthand for
    [`QuantumDickeSystem`](@ref)`(ClassicalDicke.ClassicalDickeSystem(ω₀=ω₀, ω=ω, γ=γ), j=j, Nmax=Nmax)`.
    """
    function QuantumDickeSystem(;ω₀::Real,ω::Real,γ::Real, j::Real,Nmax::Union{Integer,Nothing}=nothing)
        return QuantumDickeSystem(ClassicalDicke.ClassicalDickeSystem(ω₀=ω₀,ω=ω,γ=γ), j= j, Nmax=Nmax)
    end
    function Base.show(io::IO, qds::QuantumDickeSystem)
        ω₀,ω,γ,j,Nmax=get_params(qds,need_Nmax=false)

        print(io,"QuantumDickeSystem(ω₀ = $ω₀, ω = $ω, γ = $γ, j = $j",if Nmax ==nothing "" else ", Nmax = $Nmax" end,")")
    end
    """
    ```julia
    function dimension(system::QuantumDickeSystem)
    ```
    Returns the integer ``(2j + 1)\\times N_\\text{max}``, which gives the dimension
    of the Hilbert space.
    """
    function dimension(system::QuantumDickeSystem)
        ω₀,ω,γ,j,Nmax=get_params(system)
        return Int(2j +1)*Nmax
    end
    ind(n,m,N,j)= Int(n)*Int(N+1)+Int(m+j) +1 #+1 because Julia counts from 1.
    function logfact(n::Real)::Float64
        if n==0
            return 0.0
        end
        return sum(log(i) for i in 1:n)
    end
    function _sqrt_factorial(n::Real)
        if n==0
            return 1
        end
        prod(sqrt(i) for i in 1:n)
    end
    function _sqrt_binomial(n::Real,k::Real)
        return exp(_sqrt_binomiallog(n,k))
    end
    function _sqrt_binomiallog(n::Real,k::Real)
        return (logfact(n)-logfact(k)-logfact(n-k))/2
    end
    function get_params(system::QuantumDickeSystem;need_Nmax=true)
        j=system.j
        Nmax=system.Nmax
        ω₀,ω,γ =ClassicalSystems.parameters(system)
        if Nmax!==nothing
            Nmax=Int(Nmax)
        else
            if need_Nmax
                error("You did not pass Nmax to QuantumDickeSystem and did not call diagonalization to load it from disk. This call requires to know Nmax")
            end
        end
        return ω₀,ω,γ,j,Nmax
    end
    """
    ```julia
    function parity_operator(system::QuantumDickeSystem)
    ```
    Returns a sparse matrix corresponding to the parity operator ``\\hat{\\Pi}=e^{i\\pi(\\hat{a}^\\dagger\\hat{a}+ \\hat{J}_z + j)}`` in the
    efficient coherent basis.
    # Arguments
    - `system` should be an instance of [`DickeBCE.QuantumDickeSystem`](@ref).
    """
    function parity_operator(system::QuantumDickeSystem)
        ω₀,ω,γ,j,Nmax=get_params(system)
        N=Int(2*j)
        I=Int64[]
        J=Int64[]
        V=Int8[]
        function pushv(i,j,v)
            push!(I,i)
            push!(J,j)
            push!(V,v)
        end

        for indexN in 0:(Nmax-1), m in -j:j
            pushv(ind(indexN,m,N,j),ind(indexN,-m,N,j),(-1)^indexN)
        end
        return sparse(I,J,V,Nmax*(N+1), Nmax*(N+1))
    end
    """
    ```julia
    function eigenstate_parities(system::QuantumDickeSystem,eigenstates::AbstractMatrix{<:Number})
    ```
    Returns a vector of `-1`s and `1`s contaning the parities of all of the eigenstates.
    # Arguments
    - `system` should be an instance of [`DickeBCE.QuantumDickeSystem`](@ref).
    - `eigenstates` is the matrix of eigenstates. (See  [`diagonalization`](@ref DickeBCE.diagonalization)).
    """
    function eigenstate_parities(system::QuantumDickeSystem,eigenstates::AbstractMatrix{<:Number})
        Π=parity_operator(system)
        ΠE=Π*eigenstates
        return Int8[Integer(round(dot(view(ΠE,:,k),view(eigenstates,:,k)))) for k in 1:size(eigenstates)[2]]
    end
    function population_of_level_n(system,state,n)
        ω₀,ω,γ,j,Nmax=get_params(system)
        N=Int(2*j)
        return sum(abs(state[ind(n,m,N,j)]) for m in -j:j)
    end
    """
    ```julia
    function eigenenergies(system::QuantumDickeSystem;kargs...)
    ```
    Returns a vector containing the eigenenergies of the system.

    This function is just shorthand for [`diagonalization`](@ref DickeBCE.diagonalization)`(system;kargs..., only_eigenenergies = true)`.
    """
    function eigenenergies(system::QuantumDickeSystem;kargs...)
        return diagonalization(system;kargs..., only_eigenenergies=true)
    end
    """
    ```julia
    function diagonalization(system::QuantumDickeSystem;
            load_cache::Bool = true,
            save_cache::Bool = true,
            cache_folder::AbstractString = joinpath(homedir(),"dicke_diagonalizations"),
            maxϵ::Real = 5.0,
            onlyload::Union{AbstractVector{<:Integer},Nothing} = nothing,
            only_eigenenergies::Bool = false,
            verbose::Bool = true,
            converged_tolerance::Real = 1e-3)
    ```
    Diagonalizes the Dicke Hamiltonian up to a maximum energy `maxϵ`. The resulting eigenstates
    are guaranteed to be converged, with a tolerance determined by `converged_tolerance`.
    Numerical degeneracies are also corrected, to ensure that the eigenstates have parity +1 or -1.
    If `only_eigenenergies` is `false` (default), a tuple `(eigenenergies,eigenstates)` is returned, where
    `eigenenergies` is real vector containing the eigenenergies and `eigenstates` is a real matrix that
    contains the eigenstates as columns.

    If `load_cache = true` it will try to load saved diagonalizations from `cache_folder`. In the case that
    `system` has `Nmax = nothing`, the diagonalization with the largest `Nmax` will be loaded, and `Nmax`
    will be set in `system`.

    # Arguments
    - `system` should be an instance of [`DickeBCE.QuantumDickeSystem`](@ref).
    # Keyword arguments
    - `load_cache` is a boolean indicating whether to try to load from cache folder. Defaults to `true`.
    - `save_cache` determines if the results of a computed diagonalization are saved to the cache folder. Defaults to `true`.
    - `cache_folder` is the cache folder where diagonalizations are saved. Default is `%HOME%/dicke_diagonalizations`
    - `maxϵ` is the maximum energy up to which we diagonalize. Keep this number
      higher than the maximum converged regime you want.
    - `onlyload` may be a vector of integers, indicating the indices of the eigenstates
      to load, or the default value, `nothing`, indicates that all eigenstates should be loaded.
    - `only_eigenenergies` should be `true` if you only to load the eigenenergies. Defaults
      to `false` (return eigenstates and eigenenergies).
    - `verbose` is `true` (default) if you want to see info messages.
    - `converged_tolerance` determines how strict we are in saying that an eigenstate is converged (see [Bastarrachea2014PSb](@cite)).
      The default value `1e-3` is usually the best.
    """
    function diagonalization(system::QuantumDickeSystem;
            load_cache::Bool = true,
            save_cache::Bool = true,
            cache_folder::AbstractString = joinpath(homedir(),"dicke_diagonalizations"),
            maxϵ::Real = 5.0,
            onlyload::Union{AbstractVector{<:Integer},Nothing} = nothing,
            only_eigenenergies::Bool = false,
            verbose::Bool = true,
            converged_tolerance::Real = 1e-3)

        ω₀,ω,γ,j,Nmax=get_params(system,need_Nmax=false)

        filename(ω₀,ω,γ,j,Nmax)="dicke_bce_ω₀=$(ω₀),ω=$(ω),γ=$(γ),j=$(j),Nmax=$(Nmax)"
        filelocation(ω₀,ω,γ,j,Nmax)=joinpath(cache_folder,filename(ω₀,ω,γ,j,Nmax))
        if Nmax===nothing
            if save_cache==false
                error("Nmax is nothing,  so we can't diagonalize, but save_cache is false. Give Neff or set save_cache to true")
            end
            reg=Regex(string(filename(ω₀,ω,γ,j,"(\\d+)"),"_E.csv"))
            diags=readdir(cache_folder)
            for d in diags
                m=match(reg,d)
                if m===nothing
                    continue
                end
                poss_nmax=parse(Int64,m.captures[1])
                flname=filelocation(ω₀,ω,γ,j,poss_nmax)
                if isfile("$(flname)_E.csv") && isfile("$(flname)_V.csv")
                    if Nmax===nothing
                        Nmax=poss_nmax
                    else
                        Nmax=max(poss_nmax,Nmax)
                    end
                end
            end
            if Nmax===nothing
                error("No diagonalizations found, and Nmax is nothing. Pass a value for Nmax to diagonalize.")
            else
                system.Nmax=Nmax
            end
        end

        flname=filelocation(ω₀,ω,γ,j,Nmax)
        if load_cache

             if isfile("$(flname)_E.csv") && isfile("$(flname)_V.csv")
                if verbose
                @info "Loading diagonalization: Nmax=$Nmax"
                end
                eigenenergies=DataFrame(CSV.File("$(flname)_E.csv")).E[:]
                if only_eigenenergies
                    return eigenenergies
                end
                eigenstates=Tables.matrix(DataFrame(CSV.File("$(flname)_V.csv",select=onlyload,type=Float64,threaded =true,tasks=4)))
                return eigenenergies,eigenstates
             end
        end

        Hq=hamiltonian_operator(system)

        if verbose
            @info "Diagonalizing..."
        end
        FullH=Symmetric(Matrix(Hq));
        Emin=-10*j
        Emax=maxϵ*j #maxenergy
        eigenenergies,eigenstates=eigen(FullH,Real(Emin),Real(Emax));

        convstates=(length(eigenenergies)-1)
        for k in 1:(length(eigenenergies)-1)
            if population_of_level_n(system,eigenstates[:,k],Nmax-1) > converged_tolerance
                convstates=k
                break
            end
        end
        eigenstates=eigenstates[:,1:convstates]
        eigenenergies=eigenenergies[1:convstates]
        if verbose
        @info "$convstates converged states were obtained up to ϵ=$(eigenenergies[end]/j)."

        @info "Fixing numerical degeneracies (correcting parity)."
        end
        #corregimos por degeneraciones numéricas
        Π=parity_operator(system)
        for k in 1:(length(eigenenergies)-1)
            if(eigenenergies[k+1]-eigenenergies[k]<1e-5)
                _correct_eigenstates(eigenstates,k,k+1,Π)
            end
        end
        if save_cache
            mkpath(cache_folder)
            CSV.write("$(flname)_V.csv",DataFrame(eigenstates,:auto))
            CSV.write("$(flname)_E.csv",DataFrame(E=eigenenergies))
        end
        return eigenenergies,eigenstates
    end

#    Assuming that Ek1 ≈ Ek2, checks whether |Ek1> and |Ek1> have different parities,
#    and, if so, it rotates them inside the degenerate energy subspace so that
#    the resulting two states have parity 1 and -1.
    function _correct_eigenstates(eigenstates,k1,k2,Π)
        v1=eigenstates[:,k1]
        v2=eigenstates[:,k2]
        eigenvals,mat =eigen(Transpose([v1 v2])*Π*[v1 v2])
        if !issetequal([1.0,-1.0],round.(eigenvals,digits=7))
            return
        end
        λ_1=mat[1,1]
        λ_2=mat[2,1]
        if λ_1<0
            λ_1,λ_2=-λ_1,-λ_2
        end
        μ_1=mat[1,2]
        μ_2=mat[2,2]
        if μ_2<0
            μ_1,μ_2=-μ_1,-μ_2
        end
        eigenstates[:,k1].= λ_1*v1+λ_2*v2;
        eigenstates[:,k2].=μ_1*v1+μ_2*v2;
        return true
    end
    """
    ```julia
    function hamiltonian_operator(system::QuantumDickeSystem)
    ```
    Returns a sparse matrix corresponding to the Dicke Hamiltonian (See Eq. (1) of [Pilatowsky2021NatCommun](@cite)) in the
    efficient coherent basis. (See Ref. [Bastarrachea2014PSa](@cite), [Bastarrachea2014PSb](@cite))
    # Arguments
    - `system` should be an instance of [`DickeBCE.QuantumDickeSystem`](@ref).
    """
    function hamiltonian_operator(system::QuantumDickeSystem;_ignorediagonalterms::Bool = false)
        ω₀,ω,γ,j,Nmax=get_params(system)
        N=Int(2*j)
        G=2*γ/(ω*sqrt(N))
        function DNNp(N,Np)

            f(k)=exp((logfact(N)+logfact(Np))/2 - (logfact(Np-k)+logfact(N-k)+logfact(k)))
            return exp(-G^2/2)*sum((f(k)*(-1)^k) * (G^(N+Np-2*k)) for k in 0:Np)
        end
        I=Int64[]
        J=Int64[]
        V=Float64[]
        function pushv(i,j,v)
            push!(I,i)
            push!(J,j)
            push!(V,v)
        end
        for indexN in 0:(Nmax-1)
            for indexNprime in 0:indexN
                D=DNNp(indexN,indexNprime)
                for m in -j:j
                    if !_ignorediagonalterms
                        if indexNprime==indexN
                            pushv(ind(indexN,m,N,j),ind(indexN,m,N,j),ω*(indexN-G^2*m^2))
                        end
                    end
                    if m<j
                        o=(-1)^indexN*D
                        v=-ω₀*sqrt(j*(j+1)-m*(m+1))*o/2
                        pushv(ind(indexNprime,m+1,N,j),ind(indexN,m,N,j),v)

                    end
                    if m>-j
                        o=(-1)^indexNprime*D
                        v=-ω₀*sqrt(j*(j+1)-m*(m-1))*o/2
                        pushv(ind(indexNprime,m-1,N,j),ind(indexN,m,N,j),v)

                    end
                end
            end
        end
        D=dimension(system)
        H = sparse(I,J,V, D,D)
        H=sparse(Symmetric(H));
        return H
    end
    """
    ```julia
    function Jz(system::QuantumDickeSystem)
    ```
    Returns a sparse matrix representing the operator ``\\hat{J}_z`` in the
    efficient coherent basis.
    # Arguments
    - `system` should be an instance of [`DickeBCE.QuantumDickeSystem`](@ref).
    """
    function Jz(system::QuantumDickeSystem)
        ω₀,ω,γ,j,Nmax=get_params(system)
        return hamiltonian_operator(QuantumDickeSystem(ClassicalDicke.ClassicalDickeSystem(ω₀=1,ω=1,γ=1),j=j,Nmax=Nmax),_ignorediagonalterms=true)
    end
    """
    ```julia
    function Jx(system::QuantumDickeSystem)
    ```
    Returns a diagonal matrix representing the operator ``\\hat{J}_x`` in the
    efficient coherent basis.
    # Arguments
    - `system` should be an instance of [`DickeBCE.QuantumDickeSystem`](@ref).
    """
    function Jx(system::QuantumDickeSystem)
        ω₀,ω,γ,j,Nmax=get_params(system)
        N=Int(2j)
        Jx = zeros(Float64,dimension(system))
        for indexN in 0:(Nmax-1)
            for m in -j:j
               Jx[ind(indexN,m,N,j)]=m
            end
        end
        return Diagonal(Jx)
    end
    """
    ```julia
    function participation_ratio(state::AbstractVector{<:Number};
        eigenstates::AbstractMatrix{<:Number},
        eigenenergies::Union{AbstractVector{<:Real},Nothing}=nothing,
        count_degeneracies::Bool=eigenenergies!=nothing,
        degentol::Real=1e-5)
    ```
    Returns the Participation Ratio (PR) of `state` in the eigenbasis. If `count_degeneracies` is `true`,
    then Eq. (7) of [Villasenor2020](@cite) is used. If `count_degeneracies` is `false`, then the inverse
    of Eq. (19) of [Villasenor2020](@cite) is used.
    # Arguments
    - `state` should be a complex vector representing the state in the efficient coherent basis.
    # Keyword arguments
    - `eigenstates` should be a matrix containing the eigenstates.
    - `eigenenergies` should be passed if `count_degeneracies` is  `true`. It is a list
      containing the eigenenergies.
    - `count_degeneracies` -- whether to modify the PR definition to account for degeneracies (see above).
      Default is `false` if `eigenenergies` is not passed, else it is `true`.
    - `degentol` minimum energy separation below which two eigenstates are considered degenerate. Default
      is `1e-5`.
    """
    function participation_ratio(state::AbstractVector{<:Number};
        eigenstates::AbstractMatrix{<:Number},
        eigenenergies::Union{AbstractVector{<:Real},Nothing}=nothing,
        count_degeneracies::Bool=eigenenergies!=nothing,
        degentol::Real=1e-5)
        eigenbasis_state=transpose(eigenstates)*state
        if count_degeneracies
            IPR=0.0
            degsbspace=0.0
            for i in 1:length(eigenenergies)

                if  i!=1 && eigenenergies[i]-eigenenergies[i-1]>degentol
                    IPR+=degsbspace^2
                    degsbspace=0.0
                end
                degsbspace+=abs2(eigenbasis_state[i])
                if i==length(eigenenergies)
                    IPR+=degsbspace^2
                end
            end
            return 1/IPR
        else
            return 1/sum(abs2(c)^2 for c in eigenbasis_state)
        end
    end
    """
    ```julia
    function factor_R_of_coherent_state(system::QuantumDickeSystem,
            x::AbstractVector{<:Real};
            eigenstates::AbstractMatrix{<:Number},
            eigenenergies::AbstractVector{<:Real},
            state::AbstractVector{<:Number} = coherent_state(system,x))
    ```
    Computes ``R`` for a coherent state, as defined in Eq. (30) of [Villasenor2020](@cite).
    # Arguments
    - `system` should be an instance of [`DickeBCE.QuantumDickeSystem`](@ref).
    - `x` is an array `[Q,q,P,p]` representing the center of the coherent state.
    # Keyword arguments
    - `eigenstates` should be a matrix containing the eigenstates as columns.
    - `eigenenergies`  is a list containing the eigenenergies.
    - `state` is a complex vector representing the coherent state ``\\left | \\mathbf{x} \\right \\rangle``.
      If it is not passed, it is computed using [`DickeBCE.coherent_state`](@ref)
    """
    function factor_R_of_coherent_state(system::QuantumDickeSystem,
            x::AbstractVector{<:Real};
            eigenstates::AbstractMatrix{<:Number},
            eigenenergies::AbstractVector{<:Real},
            state::AbstractVector{<:Number} = coherent_state(system,x))
        ω₀,ω,γ,j,Nmax=get_params(system)
        ϵ=ClassicalDicke.hamiltonian(system)(x)
        σ=energy_width_of_coherent_state(system, x)
        ν(ϵ)=density_of_states(system, ϵ)
        n=Distributions.Normal(ϵ,σ)
        f=1/(2*sqrt(pi)*σ*j*ν(ϵ))

        return 2*participation_ratio(state,eigenstates=eigenstates,eigenenergies=eigenenergies,count_degeneracies=true)*f
    end
    """
    ```julia
    function survival_probability(t::Union{Real,AbstractArray{<:Real}};
        state::AbstractVector{<:Number},
        eigenstates::AbstractMatrix{<:Number},
        eigenenergies::AbstractVector{<:Real},
        normwarning::Real=0.99)
    ```
    Computes the survival probability ``| \\langle`` `state` ``| e^{-i \\hat{H} t} |`` `state` ``\\rangle | ^2``
    at the times given.
    # Arguments
    - `t` may be a single time, or a vector of times. In the first case, a number will be
      returned, and in the second, a vector.
    # Keyword arguments
    - `state` is a vector representing the state in the BCE.
    - `eigenstates` and `eigenenergies` are the ones returned by [`diagonalization`](@ref).
    - `normwarning` is a tolerance (defaults to `0.99`). If the norm of the state
      in the eigenbasis is below this number, a warning will be thrown. This usually
      happens if `Nmax` is too small for the energy regime you are working on.
    """
    function survival_probability(t::Union{Real,AbstractArray{<:Real}};
        state::AbstractVector{<:Number},
        eigenstates::AbstractMatrix{<:Number},
        eigenenergies::AbstractVector{<:Real},
        normwarning::Real=0.99)
        if typeof(t)<:Real
            return survival_probability([t],state=state,eigenstates=eigenstates,eigenenergies=eigenenergies)[1]
        end
        eigenbasis_state=transpose(eigenstates)*state
        if norm(eigenbasis_state)<normwarning
            @warn "The state is not converged in the eigenbasis, this is probably because Nmax is too small"
        end
        return [abs2(sum(exp(-im*ti*Ek)*abs2(ck) for (ck,Ek) in zip(eigenbasis_state,eigenenergies))) for ti in t]
     end
     """
     ```julia
     function evolve(t::Union{Real,AbstractArray{<:Real}}
         state::AbstractVector{<:Number};
         eigenstates::AbstractMatrix{<:Number},
         eigenenergies::AbstractVector{<:Real},
         normwarning::Real=0.99)
     ```
     Computes the evolution ``e^{-i \\hat{H} t}|`` `state` ``\\rangle`` under the
     Dicke Hamiltonian ``\\hat{H}``. The result is returned in the BCE.

     # Arguments
     - `t` may be a single time, or a vector of times. In the first case, a vector will be
       returned, and in the second, a matrix with each column corresponding to each time.
     - `state` is a vector representing the state in the BCE.
     # Keyword arguments
     - `eigenstates` and `eigenenergies` are the ones returned by [`diagonalization`](@ref).
     - `normwarning` is a tolerance (defaults to 0.99). If the norm of the state
       in the eigenbasis is below this number, a warning will be thrown. This usually
       happens if `Nmax` is too small for the energy regime you are working on.
     """
     function evolve(t::Union{Real,AbstractArray{<:Real}},
         state::AbstractVector{<:Number};
         eigenstates::AbstractMatrix{<:Number},
         eigenenergies::AbstractVector{<:Real},
         normwarning::Real=0.99)
        if typeof(t)<:Real
            return evolve([t],state,eigenstates=eigenstates,eigenenergies=eigenenergies)
        end
        eigenbasis_state=transpose(eigenstates)*state
        if norm(eigenbasis_state)<normwarning
            @warn "The state is not converged in the eigenbasis, this is probably because Nmax is too small"
        end
        eigenstates*hcat((Diagonal(exp.(-im*ti*eigenenergies))*eigenbasis_state for ti in t)...)
    end
    """
    ```julia
    function coherent_state(system::QuantumDickeSystem,
        x::AbstractVector{<:Real};
        normwarning::Real=0.99,
        extra_phase::Complex=1.0+0im,
        chop::Real=0.0)
    ```
    Calls [`coherent_state!`](@ref) passing `data` as a new vector and returns the result.
    """
    function coherent_state(system::QuantumDickeSystem,
        x::AbstractVector{<:Real};
        normwarning::Real=0.99,
        extra_phase::Complex=1.0+0im,
        chop::Real=0.0)

        ω₀,ω,γ,j,Nmax=get_params(system)

        data = zeros(Complex{Float64},Nmax*(Int(2*j)+1))
        coherent_state!(system,x,data;normwarning=normwarning,extra_phase=extra_phase,chop=chop)
        return data
    end
    """
    ```julia
    function coherent_state!(system::QuantumDickeSystem,
                x::AbstractArray{<:Real,1},
                data::AbstractVector{<:Number};
                normwarning=0.99,
                add_state_to_data::Bool=false,
                extra_phase::Complex=1.0+0im,
                chop::Real=1e-6)
    ```
    This function computes the coefficients of a coherent state centered at
    `x` in the BCE, and stores the result in `data`. (See Ref. [Villasenor2020](@cite))
    # Arguments
    - `system` should be an instance of [`DickeBCE.QuantumDickeSystem`](@ref).
    - `x` is a vector with the coordinates `[Q, q, P, p]`.
    - `data` should be a complex vector of length [`dimension`](@ref)`(system)`.
    # Keyword arguments
    - `normwarning` is a tolerance (defaults to 0.99). If the norm of the resulting
      coherent state is below this number, a warning will be thrown. This usually
      happens if Nmax is too small for the energy regime you are working on.
    - `add_state_to_data` is a boolean. If it is `true`, the coefficients will be added
      to `data`. If it is `false` (default), this function will override any preexisting
      values in `data`. This is useful if you want to add many coherent states together
      without having to allocate that much memory.
    - `extra_phase` is a complex number that multiplies the resulting state overall.
      Defaults to `1`.
    - `chop` is a numerical tolerance between 0 and 1. If `chop=0`, then all the coefficients all computed.
      However, if `ε = 1 -  chop > 0.0`, then some coefficients at the tail of the distribution (whose total squared norm does not exceed `ε`) will be treated as `0.0`. This allows to
      significantly reduce computation time, but introduces a numerical error of order `ε`. The default is `1e-6`.
      See [this example](@ref ExampleEfficientHusimiFunctions) and Ref. [Pilatowsky2020Notes](@cite).
    """
    coherent_state!(system::QuantumDickeSystem,
                x::AbstractArray{<:Real,1},
                data::AbstractVector{<:Complex};
                normwarning=0.99,
                add_state_to_data::Bool=false,
                extra_phase::Complex=1.0+0im,
                chop::Real=1e-6)=
        _coherent_overlap_or_coherent_state(system,x;
                        datacache=data,chop=chop,normwarning=normwarning,add_state_to_data=add_state_to_data,extra_phase=extra_phase)

    function _Cohpart1(n,m,alpha,G)
        s=alpha+G*m
        if n==0 && s==0.0
            return 0.0im # to avoid 0*log(0)
        end
        return n*log(s) - logfact(n)/2
    end
    function _Cohpart2(m,N,j,w,alpha,G)
        a= _sqrt_binomiallog(N,Int(j + m)) + log(w)*m -abs2(G*m)/2 -G*m*(alpha)
        return a
    end

    function qtlmin(dist,chop,mn)
        vl=mn
        try
            dt=Distributions.quantile(dist,chop)
            vl=max(vl,dt)
        catch
        end
        return vl
    end
    function qtlmax(dist,chop,mx)
        vl=mx
        try
            dt=Distributions.quantile(dist,chop)
            vl=min(vl,dt)
        catch
        end
        return vl
    end
    """
    ```julia
    function density_of_states(system::QuantumDickeSystem, ϵ::Real)
    ```
    Returns the semiclassical density of states (DoS) ``ν(ϵ)``, in units of ``1/ϵ``, as given by
    Eq. (A1) of Ref. [Pilatowsky2021Identification](@cite).

    This is computed by multiplying  the volume of the classical energy shell at ϵ
     by ``(2\\pi \\hbar_\\text{eff})^2=(2 \\pi / j)^2`` (see
    [`ClassicalDicke.energy_shell_volume`](@ref)).

    # Arguments
    - `system` should be an instance of [`DickeBCE.QuantumDickeSystem`](@ref).
    - `ϵ` is the scaled energy ``ϵ=E/j``

    See [Plotting the semiclassical density of states](@ref) for an example.
    """
    function density_of_states(system::QuantumDickeSystem, ϵ::Real)
        ω₀,ω,γ,j,_=get_params(system, need_Nmax=false)
        ClassicalDicke.energy_shell_volume(system,ϵ)*(j/(2pi))^2
    end
    """
    ```julia
    function energy_width_of_coherent_state(system::QuantumDickeSystem,
                                            x::AbstractVector{<:Real})
    ```
    Returns the energy width ``\\sigma`` of the coherent state ``\\left | \\mathbf{x}\\right \\rangle``, in units of ``\\epsilon``. This
    quantity is given by ``\\sigma_D/j`` with ``\\sigma_D`` as in App. A of Ref. [Lerma2018](@cite).

    # Arguments
    - `system` should be an instance of [`DickeBCE.QuantumDickeSystem`](@ref).
    - `x` is the coordinate ``\\mathbf{x}`` of the coherent state in the format `[Q,q,P,p]`.
    """
    function energy_width_of_coherent_state(system::QuantumDickeSystem,x::AbstractVector{<:Real})
        ω₀,ω,γ,j,_=get_params(system, need_Nmax=false)
        Q,q,P,p=x
        θ=PhaseSpaces.θ_of_QP(Q,P)
        ϕ=PhaseSpaces.ϕ_of_QP(Q,P)

        Ω₁=j*(ω^2*(q^2+p^2)/2+ω₀^2*sin(θ)^2/2 +2*γ^2*((sin( θ)^2*sin(ϕ)^2+cos(θ)^2)*q^2 + sin(θ)^2*cos(ϕ)^2) +2*γ*q*(ω*cos(ϕ) + ω₀*cos(θ)*cos(ϕ))*sin(θ))
        #note that the sign of the third term is flipped with respect to Lerma2018,
        #because they use cosθ = jz and we use cosθ = - jz.

        Ω₂=γ^2*(sin(θ)^2*sin(ϕ)^2 + cos(θ)^2)
        return sqrt(Ω₁ + Ω₂)/j
    end

    """
    ```julia
    function husimi_of_coherent(system::QuantumDickeSystem,
        x::AbstractArray{<:Real,1},
        y::AbstractArray{<:Real,1})
    ```
    Returns  ``\\left | \\left \\langle \\mathbf{x}\\middle | \\mathbf{y}\\right \\rangle \\right | ^2``.
    Equation (3.14b) of Ref. [Arecchi1972](@cite) is used for the overlap of the Bloch coherent states.
    # Arguments
    - `system` should be an instance of [`DickeBCE.QuantumDickeSystem`](@ref).
    - `x` and `y` are vectors in the form `[Q, q, P, p]`.
    """
    function husimi_of_coherent(system::QuantumDickeSystem,x::AbstractArray{<:Real,1},y::AbstractArray{<:Real,1})
        ω₀,ω,γ,j,Nmax=get_params(system,need_Nmax=false)
        Q1,q1,P1,p1=x
        Q2,q2,P2,p2=y
        return exp(-(j/2)*((q1-q2)^2 + (p1-p2)^2))*cos(PhaseSpaces.arc_between_QP(Q1,P1,Q2,P2)/2)^(4j)
    end
    onelistcache=[0.0im]
    """
    ```julia
    function coherent_overlap(system::QuantumDickeSystem,
        x::AbstractVector{<:Real},
        state::AbstractVector{<:Number};
        chop::Real=1e-6,
        normwarning::Real=0.99,
        datacache::Union{AbstractArray{Complex{Float64},1},Nothing}=nothing)
    ```
    Returns  the overlap ``\\langle \\mathbf{x} | `` `state` ``\\rangle``,
    where ``\\left | \\mathbf{x} \\right \\rangle`` is a coherent state centered at `x`.
    # Arguments
    - `system` should be an instance of [`DickeBCE.QuantumDickeSystem`](@ref).
    - `x` is a vector in the form `[Q, q, P, p]`.
    - `state` is a vector in the BCE.
    # Keyword arguments
    - `chop` is a numerical tolerance between 0 and 1. If `chop = 0`, then all the products of coefficients are computed.
      However, if `ε = 1 -  chop > 0.0`, then some coefficients of the coherent state at the tail of the distribution  (whose total squared norm does not exceed `ε`)
      will be treated as `0.0`. This allows to
      significantly reduce computation time, but introduces a numerical error of order `ε`. The default is `1e-6`.
      See [this example](@ref ExampleEfficientHusimiFunctions)  and Ref. [Pilatowsky2020Notes](@cite).
    - `normwarning` is a tolerance (defaults to 0.99). If the norm of the coherent  state is below this number, a warning will be thrown.
      This usually happens if Nmax is too small for the energy regime you are working on.
    - `datacache` in an array where to store the result, or `nothing` for the result to be returned.
    """
    function coherent_overlap(system::QuantumDickeSystem,
        x::AbstractVector{<:Real},
        state::AbstractVector{<:Number};
        chop::Real=1e-6,
        normwarning::Real=0.99,
        datacache::Union{AbstractVector{<:Complex},Nothing}=nothing)
        if datacache===nothing
            datacache=onelistcache
        end
        return _coherent_overlap(system,x,state;datalength=1,datacache=datacache,chop=chop,normwarning=normwarning)[1]
    end
    """
    ```julia
    function coherent_overlap(system::QuantumDickeSystem,
        x::AbstractVector{<:Real},
        states::AbstractMatrix{<:Number};
        chop::Real=1e-6,
        normwarning::Real=0.99,
        datacache::Union{AbstractVector{<:Complex},Nothing}=nothing)
    ```
    Same as [`coherent_overlap(..., state,...)`](@ref coherent_overlap(::QuantumDickeSystem,::AbstractVector{<:Real},::AbstractVector{<:Number})), but
    instead of one `state`, it allows `states` to be a matrix with multiple states as columns. It is functionally equivalent to
    ```julia
    [coherent_overlap(..., states[:,k], ...) for k in 1:size(states)[2]]
    ```
    but it is (much) faster.
    The result is stored in `datacache` if provided.
    """
    function coherent_overlap(system::QuantumDickeSystem,
        x::AbstractVector{<:Real},
        states::AbstractMatrix{<:Number};
        chop::Real=1e-6,
        normwarning::Real=0.99,
        datacache::Union{AbstractVector{<:Complex},Nothing}=nothing)
        d=size(states)[2]
        if datacache===nothing
            datacache=zeros(Complex{Float64},d)
        end
        return _coherent_overlap(system,x,states;datacache=datacache,datalength=d,chop=chop,normwarning=normwarning)
    end
    #if you pass states it calculates overlap with coherent, else  coefficients of coherent.
    function _coherent_overlap_or_coherent_state(system::QuantumDickeSystem,point::AbstractArray{<:Real,1};
                        datacache,datalength=nothing,chop=1e-6,normwarning=0.99,
                        states::Union{AbstractArray{<:Number,1},AbstractArray{<:Number,2},Nothing}=nothing,add_state_to_data=false,extra_phase::Complex=1.0+0im)
            if !add_state_to_data
                datacache.=0.0im
            end

            ω₀,ω,γ,j,Nmax=get_params(system)
            N=Int(2*j)
            G=2*γ/(ω*sqrt(N))
            Q,q,P,p=point
            theta=PhaseSpaces.θ_of_QP(Q,P)
            phi=PhaseSpaces.φ_of_QP(Q,P)

            z=exp(-phi*im)*tan(theta/2)
            alpha=(q+im*p)*sqrt(j/2)
            m=-j

            nrm=0.0
            w=(1+z)/(1-z)
            part0=-abs2(alpha)/2 +j*log(w/(1+abs2(w)))

            if chop==0
               mmin=-j
               mmax=j
            else
                atomicDist=Distributions.Binomial(N,abs2(w)/(1+abs2(w)))
                mmin=qtlmin(atomicDist,chop/4,0)-j
                mmax=qtlmax(atomicDist,1-chop/4,N)-j
            end
            for m in mmin:mmax
                part2m=part0+_Cohpart2(m,N,j,w,alpha,G)
                if chop==0
                    nmin=0
                    nmax=Nmax-1
                else
                    efftol=chop/((mmax-mmin+1) *Distributions.pdf(atomicDist,m+j))
                    d=Distributions.Poisson(abs(G*m+alpha)^2)
                    nmin= qtlmin(d,efftol/4,0)
                    nmax=qtlmax(d,1-efftol/4,Nmax-1)
                end
                for n in nmin:nmax
                    v=conj(exp(_Cohpart1(n,m,alpha,G)+part2m))
                    if states!==nothing
                        for e in 1:datalength
                            datacache[e]+=states[ind(n,m,N,j),e]*v
                        end

                    else
                       datacache[ind(n,m,N,j)]+=conj(v)*extra_phase
                    end
                    nrm+=abs2(v)
                end
            end
            if nrm<normwarning
                @warn "The coherent state is not converged (norm=$nrm). You are probably working in an energy regime that is too high. You should make Nmax bigger."
            end
        return datacache
    end
    _coherent_overlap(system::QuantumDickeSystem,
    point::AbstractArray{<:Real,1},
    states::Union{AbstractArray{<:Number,1},AbstractArray{<:Number,2}};
    datacache,
    datalength,
    chop=1e-6,
    normwarning=0.99) =_coherent_overlap_or_coherent_state(system,point;
                                                        datacache=datacache,
                                                        datalength=datalength,
                                                        chop=chop,
                                                        states=states,
                                                        normwarning=normwarning)
    """
    ```julia
    function husimi(args...;kargs...)
    ```
    Computes `abs2.(`[`coherent_overlap`](@ref)`(args...;kargs...))`. The arguments
    and behavior are the same as [`coherent_overlap`](@ref).
    """
    function husimi(args...;kargs...)
        r=coherent_overlap(args...;kargs...)
        return abs2.(r)
    end

    function sph_harm(l,m,theta,phi;legendre_sphPlm_cache=:nothing)
        fact=1
        if m<0
            fact=(-1)^(-m)
        end
        d=0
        if legendre_sphPlm_cache==:nothing
            d=GSL.sf_legendre_sphPlm(l,abs(m),cos(theta))
        else
            d=legendre_sphPlm_cache[l,abs(m),theta]
        end
        return fact*d*exp(im*m*phi)
    end
    function generate_hermite_cache(n_max,pts)
        cache=Dict{NTuple{2,Real},Complex}()
        points=unique(pts)
        for x in points

            data=GSL.sf_hermite_phys_array(n_max,x)
            for (n,d) in enumerate(data)
                cache[n-1,x]=d
            end
        end
        return cache
    end
    function generate_legendre_sphPlm_cache(l_max,pts)
        cache=Dict{NTuple{3,Real},Complex}()
        thetas=unique(x[3] for x in pts)
        for theta in thetas

            data=GSL.sf_legendre_array_e(GSL.GSL_SF_LEGENDRE_SPHARM,l_max,cos(theta),-1)
            for l in 0:l_max,m in 0:l
                cache[l,m,theta]=data[GSL.sf_legendre_array_index(l, m)+1]
            end
        end
        return cache
    end
    function laguerre(n,a,x)

        if a<0
            return x^(-a)/GSL.sf_poch(-n,-a)*laguerre(n+a,-a,x)
        end
        return GSL.sf_laguerre_n(n,a,x)
    end
    function RotZtoX(θ,φ)
        a=-π/2
        z=[cos(θ/2) -exp(-im*φ)*sin(θ/2)]*[cos(a/2) -sin(a/2);
                                        sin(a/2) cos(a/2)]
        return 2*atan(abs(z[2]),abs(z[1])),angle(z[2])-angle(z[1])
    end
    function classical_QqPp_to_quantum_rot_qpθφ(j,QqPp)
        Q,q,P,p=QqPp
        q=sqrt(j)*q
        p=sqrt(j)*p
        θ,φ=RotZtoX(PhaseSpaces.θ_of_QP(Q,P),PhaseSpaces.φ_of_QP(Q,P))
        return [q,p,θ,φ]
    end

    function sph_harm_pt(k,i,pt;legendre_sphPlm_cache=:nothing)
        return sph_harm(k, i, pt[3],pt[4];legendre_sphPlm_cache=legendre_sphPlm_cache)
    end

    function Wigner_Fock(n,np,q,p)
        two_r_squared=(q^2+p^2)*2
        if two_r_squared==0
            two_r_squared=1e-30
        end
        return (exp((logfact(np)-logfact(n))/2 +
            im*(np-n)*atan(p,q)- two_r_squared/2)*(-1)^(np)/π * two_r_squared^((n-np)/2)*laguerre(np,n-np,two_r_squared))
    end
    function Wigner_Fock_displaced(n,np,m,mp,G,pt)
        q=pt[1]
        p=pt[2]

        qrec=q+G*(m+mp)/2 # alpha_m displacement
        return exp(-im*p*G*(mp-m))*Wigner_Fock(n,np,qrec,p)
    end
    """
    ```julia
    function Wigner(system::QuantumDickeSystem,
        state::AbstractArray{<:Number},
        points::Vector{<:AbstractVector{<:Real}})
    ```
    Evaluates the Wigner function of `state` (which is a vector in the BCE) in all `points = [[Q1,q1,P1,p1], [Q2,q2,P2,p2], ...]`,
    returning an array with the results.

    Note: This function has not been thoroughly tested. It is based on Ref. [Pilatowsky2019Notes](@cite).
    """
    function Wigner(system::QuantumDickeSystem,
        state::AbstractArray{<:Complex},
        points::Vector{<:AbstractVector{<:Real}})

        ω₀,ω,γ,j,Nmax=get_params(system)

        N=Int(2*j)
        G=2*γ/(ω*sqrt(N))

        points=classical_QqPp_to_quantum_rot_qpθφ.(j,points)
        legendre_sphPlm_cache=generate_legendre_sphPlm_cache(Int(2*j),points)
        #θ,φ=theta,phi
        function B(n,l,np,lp)
            altsm(n-l)*altsm(np-lp)*exp((logfact(n)+logfact(np))/2 -
                (logfact(l)+logfact(lp)))
        end

        W=fill!(similar(points,Complex{Float64}),0)
        _Wd=fill!(similar(points,Complex{Float64}),0)
        for m in -j:j, mp in -j:j

            mWd=Wigner_Dicke!(m,mp,j,points,_Wd,legendre_sphPlm_cache=legendre_sphPlm_cache)
            for n in 0:(Nmax-1),np in 0:n
                _factor=state[ind(n,m,N,j)]*conj(state[ind(np,mp,N,j)])
                if norm(_factor)<1e-6
                    continue
                end
                mW=_factor*mWd
                mW.*=Wigner_Fock_displaced.(n,np,m,mp,G,points)
                W+= real.(mW*(n==np ? 1 : 2))

            end

        end
        return real.(W)
    end
    function Wigner_Dicke!(m::Int64,mp::Int64,j::Int64,points::AbstractVector{<:AbstractVector{<:Real}},result::Array{<:Complex,1};legendre_sphPlm_cache=:nothing)

        result.=0.0+0.0im

        for k in abs(m-mp):Int(2 * j)
            i=m-mp
            _f=(-1)^(j - m - i) * clebschgordan(j,m, j,-mp, k,i)
            if _f !=0
                 result +=   _f*sph_harm_pt.(k, i, points,legendre_sphPlm_cache=legendre_sphPlm_cache)
            end
        end
        return result*sqrt((2*j+1)/(4*pi)) #normalize
    end
    function realmult(w2::Complex{Float64},factor)
        return real(factor*w2)
    end
    """
    ```julia
    function WignerProjqp(system::QuantumDickeSystem,
        states::AbstractArray{<:AbstractArray{<:Number}},
        ptsQP::Vector{<:AbstractVector{<:Real}};
        show_progress::Bool=true)
    ```
    Evaluates Wigner function, integrated over the bosonic variables ``q,p``, of each of the states in `states`
    (which is a vector of complex vectors)
    in all `ptsQP = [[Q1,P1], [Q2,P2], ...]`, returning an array with the results.
    See [this example](@ref wignerfuncexample). If `show_progress = true` it shows a progress bar.

    Note: This function has not been thoroughly tested. It is based on Ref. [Pilatowsky2019Notes](@cite).
    """
    function WignerProjqp(system::QuantumDickeSystem,
        states::AbstractArray{<:AbstractArray{<:Complex}},
        ptsQP::Vector{<:AbstractVector{<:Real}};
        show_progress::Bool=true)
        ω₀,ω,γ,j,Nmax=get_params(system)
        N=Int(2*j)
        G=2*γ/(ω*sqrt(N))
        points=[classical_QqPp_to_quantum_rot_qpθφ(j,[Q,0.0,P,0.0]) for (Q,P) in ptsQP]
        tmD=0.0
        tmF=0.0
        tmM=0.0
        _Wd=similar(points,Complex{Float64})
        AllWs=[similar(points,Float64) for s in states]
        fill!.(AllWs,0.0)

        legendre_sphPlm_cache=generate_legendre_sphPlm_cache(Int(2*j),points)
        if show_progress
            p = Progress(Int((2*j+1)^2*(Nmax-1)*Nmax/2),1)
        end
        for m in -j:j, mp in -j:j

            _Wd=Wigner_Dicke!(m,mp,j,points,_Wd,legendre_sphPlm_cache=legendre_sphPlm_cache)

            for n in 0:(Nmax-1),np in 0:n
                if np==n && mp==m
                    lag_fac=1.0
                else
                lag_fac=exp((logfact(np)-logfact(n))/2 - abs(G*(m-mp))^2/2 + (n-np)*log(G*abs(m-mp)))*sign(m-mp)^(n-np)*laguerre(np,n-np,abs(G*(m-mp))^2)

                end
                for i in 1:length(states)
                    _factor=states[i][ind(n,m,N,j)]*conj(states[i][ind(np,mp,N,j)])
                    finalfactor=_factor*lag_fac*(n==np ? 1 : 2)
                    if abs(finalfactor) ==0
                        continue
                    end
                    AllWs[i].+=realmult.(_Wd,finalfactor)
                end
                if show_progress
                    next!(p)
                end

            end

        end

        if show_progress
            finish!(p)
        end
        return AllWs
    end

    function WignerProjp(system,states,points)
        ω₀,ω,γ,j,Nmax=get_params(system)
        N=Int(2*j)
        G=2*γ/(ω*sqrt(N))
        points=classical_QqPp_to_quantum_rot_qpθφ.(j,points)
        tmD=0.0
        tmF=0.0
        tmM=0.0
        _Wd=similar(points,Complex{Float64})
        _FactorHermite=similar(points,Float64)
        AllWs=[similar(points,Float64) for s in states]
        fill!.(AllWs,0.0)

        legendre_sphPlm_cache=generate_legendre_sphPlm_cache(Int(2*j),points)
        hermite_cache=generate_hermite_cache(Nmax-1,Iterators.flatten((x[1]+G*m for m in -j:j) for x in points))
        function Fock_n_np_wavefunction(n::Int64,m::Int64,np::Int64,mp::Int64,x::Array{Float64,1})
            return exp(-(logfact(np)+logfact(n)+(np+n)*log(2))/2 - (x[1]+G*m)^2/2 - (x[1]+G*mp)^2/2)*hermite_cache[n,x[1]+G*m]*hermite_cache[np,x[1]+G*mp]/sqrt(pi)
        end
        p = Progress(Int((2*j+1)^2*(Nmax-1)*Nmax/2),1)
        for m in -j:j, mp in -j:j

            _Wd=Wigner_Dicke!(m,mp,j,points,_Wd,legendre_sphPlm_cache=legendre_sphPlm_cache)

            for n in 0:(Nmax-1),np in 0:n

                _FactorHermite.=Fock_n_np_wavefunction.(n,m,np,mp,points)

                for i in 1:length(states)
                    _factor=states[i][ind(n,m,N,j)]*conj(states[i][ind(np,mp,N,j)])
                    finalfactor=_factor*(n==np ? 1 : 2)
                    if abs(finalfactor) ==0
                        continue
                    end
                    AllWs[i].+=realmult.(_Wd,finalfactor).*_FactorHermite
                end
                next!(p)

            end

        end

        finish!(p)
        return AllWs
    end
    """
    ```julia
    function WignerProjQP(system::QuantumDickeSystem,
        state::AbstractArray{<:Number},
        ptsqp::Vector{<:AbstractVector{<:Real}})
    ```
    Evaluates Wigner function, integrated over the atomic variables ``Q,P``, of `state`
    in all `ptsqp = [[q1,p1], [q2,p2], ...]`, returning an array with the results.

    Note: This function has not been thoroughly tested. It is based on Ref. [Pilatowsky2019Notes](@cite).
    """
    function WignerProjQP(system::QuantumDickeSystem,
        state::AbstractArray{<:Number},
        ptsqp::Vector{<:AbstractVector{<:Real}})
        ω₀,ω,γ,j,Nmax=get_params(system)

        N=Int(2*j)
        G=2*γ/(ω*sqrt(N))
        points=[classical_QqPp_to_quantum_rot_qpθφ(j,[0.0,q,0.0,p]) for (q,p) in ptsqp]
        #θ,φ=theta,phi

        W=fill!(similar(points,Complex),0)
        _Wd=fill!(similar(points,Complex),0)
        for m in -j:j
            for n in 0:(Nmax-1),np in 0:n
                _factor=state[ind(n,m,N,j)]*conj(state[ind(np,m,N,j)])
                if norm(_factor)<1e-20
                    continue
                end
                mW=_factor*Wigner_Fock_displaced.(n,np,m,m,G,points)

                W+= real.(mW*(n==np ? 1 : 2))

            end

        end
        return real.(W)
    end
    """
    ```julia
    function random_state(system::QuantumDickeSystem,
                        n::Int=1;
                        kargs...)
    ```
    Generates `n` random states, returning them as columns in a matrix
    (or a vector if `n = 1`, which is the default). `kargs` are redirected to [`random_state!`](@ref): some
    additional keyword arguments are required: see [`random_state!`](@ref).
    """
    function random_state(system::QuantumDickeSystem,
        n::Int=1;kargs...)
        ω₀,ω,γ,j,Nmax=get_params(system)
        D=dimension(system)
        if n==1
            data = zeros(Complex{Float64},D)
        else
            data=zeros(Complex{Float64},D,n)
        end
        return random_state!(system,data;kargs...)
    end
    """
    ```julia
    function random_cₖ_generator(system::QuantumDickeSystem;
                                kargs...)
    ```
    Generates coefficients to build random states, which are given by
    ```math
        c_k(\\epsilon_k) \\sim e^{i \\theta_k} \\sqrt{\\frac{r_k\\rho(\\epsilon_k)}{\\nu(\\epsilon_k)}}.
    ```
    (See Refs. [Lerma2019](@cite), [Villasenor2020](@cite), and [Pilatowsky2021NatCommun](@cite)).
    It returns a function `cₖ(ϵ;number=1,cache=nothing)`, which produces n=`number` coefficients, storing them
    in `cache` (an n-vector), if given.

    The parameters of the equation above are determined as follows:

    # Required keyword arguments
    - To determine  ``\\rho(\\epsilon)`` above, pass either:
        - `σ` and `ϵ`, both real, in which case ``\\rho`` will be taken as a normal distribution
          centered at `ϵ` with standard deviation `σ`,
        - OR `envelope` as any of [`Distributions.UnivariateDistribution`](https://juliastats.org/Distributions.jl/v0.25/univariate/#Continuous-Distributions).
    - To determine ``r_k``, pass either:
        - `ensemble = :GUE`, which takes ``r_k`` from an exponential distribution `Exponential(0.91)` or `ensemble = :GOE`, which takes ``r_k`` from a ``\\chi^2``
          distribution with one degree of freedom,
        - OR `rₖ_distribution` as any of [`Distributions.UnivariateDistribution`](https://juliastats.org/Distributions.jl/v0.25/univariate/#Discrete-Distributions).
    - To determine ``\\theta_k``, pass either:
        - `phases = :real`, which takes ``\\theta_k`` to be random from ``\\{ -\\pi, \\pi \\}`` or `phases = :complex`, which takes ``\\theta_k`` to be uniform in ``[0, 2\\pi)``,
        - OR `ensemble = :GOE`, which automatically sets `phases = :real`, or `ensemble = :GUE`, which automatically sets `phases = :complex`.
    # Optional keyword arguments
    - `divide_by_DoS` is a boolean. If `true` (default), leaves the density of states, ``\\nu(\\epsilon_k)`` in the equation above. A value of `false` removes it.
    """
    function random_cₖ_generator(system::QuantumDickeSystem;
                                σ::Real,
                                ϵ::Real,
                                envelope::UnivariateDistribution=Distributions.Normal(ϵ,σ),
                                ensemble::Symbol=:GUE,
                                rₖ_distribution::UnivariateDistribution=
                                if ensemble==:GUE
                                        Exponential(0.91)
                                elseif ensemble==:GOE
                                        Chisq(1)
                                else
                                    error("ensemble must be :GOE or :GUE, or pass rₖ_distribution and phases directly")
                                end,
                                phases::Symbol=if ensemble==:GUE
                                        :complex
                                elseif ensemble==:GOE
                                        :real
                                else
                                    error("ensemble must be :GOE or :GUE, or pass rₖ_distribution and phases directly")
                                end,
                                divide_by_DoS::Bool=true)

        ρ(ϵ)=Distributions.pdf(envelope,ϵ)
        ν(ϵ)=density_of_states(system, ϵ)
        rₖ()=Distributions.rand(rₖ_distribution)
        rphase()=if phases==true || phases ==:complex
            exp(im*Random.rand()*2*pi)
        elseif phases==:real
            rand((1,-1))
        else
            1
        end
        function cₖ(ϵ;number=1,cache=nothing)
            if divide_by_DoS
                dos=ν(ϵ)
            else
                dos=1
            end
            if number>1 && cache==nothing
                cache=zeros(ComplexF64,number)
            end
            if dos==0
                if number==1
                    return 0
                else
                    cache*=0
                    return cache
                end
            else
                f()=sqrt(rₖ()*ρ(ϵ)/dos)*rphase()
                if number==1
                    return f()
                else
                    cache.=f.()
                    return cache
                end

            end
        end
        return cₖ
    end
    """
    ```julia
    function random_state!(system::QuantumDickeSystem,
        data::AbstractVecOrMat{<:Number};
        kargs...)
    ```
    Generates as many random states as columns in `data` saving the result to `data`.
    The general form of the random states in the eigenbasis is
    ``\\left | R \\right \\rangle = \\sum_k c_k \\left | E_k \\right \\rangle``, where
    the ``c_k`` are computed using [`random_cₖ_generator`](@ref). You may choose to select only
    positive or negative parity eigenstates (see below).
    # Arguments
    - `system` should be an instance of [`DickeBCE.QuantumDickeSystem`](@ref).
    - `data` should be a vector or matrix of numbers, where the result is stored. It
      must have as many columns as random numbers you want and [`dimension`](@ref)`(system)` rows.
    # Required keyword arguments
    - Please pass all the required keyword arguments of [`random_cₖ_generator`](@ref).
    - `eigenstates` and `eigenenergies` shoud be passed, being the ones returned by [`diagonalization`](@ref).
    # Optional keyword arguments
    - All the optional keyword arguments of [`random_cₖ_generator`](@ref) may be passed.
    - `tol` may be a real number which determines the tolerance for the state convergence.
      A value `tol=0` builds the state with all of the eigenstates, and a positive number
      cuts the tails of the energy envelope up to that number. Higher values are faster,
      but decrease precision (Defaults to `1e-6`)
    - `parity` may be `+`, `-`, or `nothing` (default). If it is not `nothing`, the random
      will be composed of only the eigenstates with this parity.
    - `parities` should be passed if `parity` is not `nothing`. It should be a vector of `-1`s
      and `1`s containing the parities of all of the eigenstates. If it is not passed, it will be
      computed with [`eigenstate_parities`](@ref), but this is slow. If you plan to use this
      function repeatedly, you should precompute `parities` by calling [`eigenstate_parities`](@ref)
      yourself and then pass it to `parities` to avoid repeated calls to [`eigenstate_parities`](@ref).
    """
    function random_state!(system::QuantumDickeSystem,
        data::AbstractVecOrMat{<:Number};
        eigenstates::AbstractMatrix{<:Number},
        eigenenergies::AbstractVector{<:Real},
        σ::Real,
        ϵ::Real,
        envelope::UnivariateDistribution=Distributions.Normal(ϵ,σ),
        tol::Real=1e-6,
        parity::Union{Nothing,typeof(-),typeof(+)}=nothing,
        parities::AbstractArray{<:Integer,1}=if parity==nothing
            Int8[]
        else #@warn "If you pass parity=+-, it is more efficient for you to compute parities=eigenstate_parities(system,eigenstates), and pass it to random_state!(... parities=parities, ...). This prevents repeated calls to eigenstate_parities wich can be slow."
            eigenstate_parities(system,eigenstates) end,kargs...)

        if length(size(data))==2
           for i in 1:size(data)[2]
                random_state!(system,view(data,:,i);
                                eigenenergies=eigenenergies,eigenstates=eigenstates,σ=σ,ϵ=ϵ,envelope=envelope,tol=tol,parity=parity,parities=parities,kargs...)

           end
           return data
        end
        ω₀,ω,γ,j,Nmax=get_params(system)
        cₖ=random_cₖ_generator(system;ϵ=ϵ,σ=σ,envelope=envelope,kargs...)
        data.=0.0
        nrm=0.0
        ϵ0,ϵf=Distributions.quantile(envelope,tol/2),Distributions.quantile(envelope,1-tol/2)
        if parity!=nothing
            par=parity(1)
            #this HAS to be done with an anonymous function, because
            # https://github.com/JuliaLang/julia/issues/15602
            is_correct_parity = k -> parities[k]==par

        else
            is_correct_parity = k -> true #idem
        end
        for (k,E) in enumerate(eigenenergies)

            ϵ=E/j
            if ϵ0<=ϵ<=ϵf && is_correct_parity(k)
                ck=cₖ(ϵ)
                nrm+= abs2(ck)
                data.+=ck.*view(eigenstates,:,k)
            end

        end
        data./=sqrt(nrm)
        return data
    end

    """
    ```julia
    function random_coherent_states_in_energy_shell(
        system::QuantumDickeSystem;
        ϵ::Real,
        N::Integer,
        dt::Real=3,
        chop::Real=1e-6,
        cachedata::Union{AbstractVector{<:Complex},Nothing}=nothing)
    ```
    Samples `N` points from the energy shell at `ϵ` using [`ClassicalDicke.classical_path_random_sampler`](@ref ClassicalDicke.classical_path_random_sampler),
    and then constructs the N-cat state of all the coherent states centered at those points.
    """
    function random_coherent_states_in_energy_shell(
        system::QuantumDickeSystem;
        ϵ::Real,
        N::Integer,
        dt::Real=3,
        chop::Real=1e-6,
        cachedata::Union{AbstractVector{<:Complex},Nothing}=nothing)

        if cachedata===nothing
            cachedata=zeros(Complex,dimension(system))
        end
        cachedata.=0.0im
        s=ClassicalDicke.classical_path_random_sampler(system,ϵ=ϵ,dt=dt)
        for i in 1:N
            θ=rand()*2*π
            p=s()
            coherent_state!(system,p,cachedata,extra_phase=exp(im*θ),chop=chop,add_state_to_data=true)
        end
        cachedata./=norm(cachedata)
        return cachedata
     end
end


text:
println(Sys.WORD_SIZE)

println("typeof(1): $(typeof(1))")
println("zero: $(zero(Float64))")
println("zero: $(one(Float64))")

for T in [Int8,Int16,Int32,Int64,Int128,UInt8,UInt16,UInt32,UInt64,UInt128]
    println("$(lpad(T,7)): [$(typemin(T)),$(typemax(T))]")
end

x = [1,2,3] .^ 3
println("vectorized operation using dot notation: $(x))")

# NaN is not equal to, not less than, and not greater than anything, including itself.
println(NaN == NaN)
println([1 NaN] == [1 NaN])
println([1 3] == [1 3])
println(isequal(NaN, NaN))
println(isfinite(8))
println(isinf(Inf))
println(isnan(7))

# can chain comparsions
v(x) = (println(x); x)
v(1) < v(2) <= v(3)

# rational numbers - note that number is reduced to smallest possible terms
x = 6//9
println("rational number: $(x), numerator: $(numerator(x)), denominator: $(denominator(x))")

# strings
println("""Can use triple "quotes" to escape embedded quotes""")
str = "Hello World!"
println("Single index: $(str[1]), $(typeof(str[1]))")  # also note indexing is 1-based
println("Same thing but using range: $(str[1:1]), $(typeof(str[1:1]))")
println("Using 'end' keyword: $(str[end])")

# function declaration - "traditional"
function f(x,y)
    x + y
end
println("1 + 2 = $(f(1, 2))")

# assignment form
g(x, y) = x + y
println("2 + 3 = $(g(2, 3))")

# anonymous function example - can also use more than 1 variable (or 0)
println(map(x -> x^2 + 2x - 1, [1,3,-1]))

# named tuples
x = (a=1, b=1+1)
println(x.a)

# argument destructuring
minmax2(x, y) = (y < x) ? (y, x) : (x, y)
range2((min, max)) = max - min  # note the extra set of parentheses
println("range of (10, 2): $(range2(minmax2(10, 2)))")

# varargs
bar(a,b,x...) = (a,b,x)
println("2 args (no varargs): $(bar(1,2))")
println("6 args: $(bar(1,2,3,4,5,6))")
x = (11, 12)
println("With splatting: $(bar(1, 2, x...))")

# Optional arguments
function Example(y::Int64, x::Int64=0, z::Int64=0)
    y + x + z
end

println(Example(2011, 2))

# Keyword arguments - note the semicolon
function circle(x, y; radius::Int=1)
    return (x=x, y=y, radius=radius)
end
println(circle(2, 3))
println(circle(4, 5; :radius => 6))

function keywords(; kwargs...)
    println(kwargs...)
end
keywords(x="Hello", y="World", z="!")

# do blocks - pass in function as the first argument
y = map([-2, -1, 0, 2]) do x
    if x < 0 && iseven(x)
        return 0
    elseif x == 0
        return 1
    else
        return x
    end
end
println(y)

# vectorizing functions (dot syntax)
f(x,y) = 3x + 4y
A = [1.0, 2.0, 3.0]
B = [4.0, 5.0, 6.0]
println(f.(pi, A))
println(f.(A, B))

Y = [1.0, 2.0, 3.0, 4.0]
X = similar(Y); # pre-allocate output array - more efficient than allocating for results
@. X = sin(cos(Y)) # equivalent to X .= sin.(cos.(Y))
println(X)

# compound expressions - begin blocks, note that they don't have to be multiline
z = begin
   x = 1
   y = 2
   x + y
end
println("compound expression: $(z)")

z = (x = 1; y = 2; x + y)  # using chain syntax - can be multiline
println("using chain syntax: $(z)")

# conditional statements - note that no local scope
function test(x,y)
   if x < y
       relation = "less than"
   elseif x == y
       relation = "equal to"
   else
       relation = "greater than"
   end
   println("x is ", relation, " y.")
end
test(1, 2)
test(2, 1)
test(2, 2)

# ternary operator
x = 1; y = 0
println(x < y ? "less than" : "not less than")

# short circuit evaluation
# you can use non-booleans as the last expression in a conditional chain
function fact(n::Int)
   n >= 0 || error("n must be non-negative")
   n == 0 && return 1
   n * fact(n-1)
end
println(fact(5))
println(fact(0))
#println(fact(-1))

# loops - while and for
i = 1
while i <= 5
   println(i)
   global i += 1
end

for i = 1:5  # 1:5 is a range object
   println(i)
end

for i in [1,4,0]
   println(i)
end

# same as using the keyword in
for s ∈ ["foo","bar","baz"]
   println(s)
end

for i ∈ [1,2,3,4,5,6]
    if iseven(i)
        continue
    elseif i == 5
        break
    end
    println(i)
end

# forms cartesian product of the iterables
for i = 1:2, j = 3:4
   println((i, j))
end

# exception handling
f(x) = x>=0 ? exp(-x) : throw(DomainError(x, "argument must be nonnegative"))
try
    f(-1)
catch e
    if isa(e, DomainError)
        println("you can only provide nonnegative numbers")
    else
        rethrow(e)
    end
finally
    println("we did it team")
end

ext: jl

max_stars_repo_name: suicidejack/julia-practice

max_issues_repo_name: suicidejack/julia-practice

max_forks_repo_name: suicidejack/julia-practice


text:
using CalciumScoring
using Test
using TestSetExtensions


text:
using DrWatson
@quickactivate
using ArgParse
using PyCall
using GenerativeAD
using StatsBase: fit!, predict, sample
using OrderedCollections
using BSON

s = ArgParseSettings()
@add_arg_table! s begin
	"max_seed"
		default = 1
		arg_type = Int
		help = "maximum number of seeds to run through"
	"dataset"
		default = "iris"
		arg_type = String
		help = "dataset"
	"sampling"
		default = "random"
		arg_type = String
		help = "sampling of hyperparameters"
	"contamination"
		arg_type = Float64
		help = "contamination rate of training data"
		default = 0.0
end
parsed_args = parse_args(ARGS, s)
@unpack dataset, max_seed, sampling, contamination = parsed_args

modelname = "pidforest"
function sample_params()
	parameters_rng = (;
		max_depth	= 6:2:10,
		n_trees		= 50:25:200,
		max_samples	= [50, 100, 250, 500, 1000, 5000],
		max_buckets	= 3:6,
		epsilon		= [0.05, 0.1, 0.2]
	)

	return (;zip(keys(parameters_rng), map(x->sample(x, 1)[1], parameters_rng))...)
end

function create_space()
	pyReal = pyimport("skopt.space")["Real"]
	pyInt = pyimport("skopt.space")["Integer"]
	pyCat = pyimport("skopt.space")["Categorical"]

	(;
		max_depth	= pyInt(6, 10, 							name="max_depth"),
		n_trees		= pyInt(50, 200,						name="n_trees"),
		max_samples	= pyInt(25, 5000, prior="log-uniform",	name="max_samples"),
		max_buckets	= pyInt(3, 6, 							name="max_buckets"),
		epsilon		= pyReal(0.05, 0.2, 					name="epsilon")
	)
end

function GenerativeAD.edit_params(data, parameters)
	D, N = size(data[1][1])
	if N < parameters.max_samples
		# if there are not enough samples, choose closest multiple of 50
		@info "Not enough samples in training, changing max_samples for each tree."
		return merge(parameters, (;max_samples = max(25, N - mod(N, 50))))
	else
		return parameters
	end
end

function fit(data, parameters)
	model = GenerativeAD.Models.PIDForest(Dict(pairs(parameters)))

	try
		global info, fit_t, _, _, _ = @timed fit!(model, data[1][1])
	catch e
		@info "Failed training due to \n$e"
		return (fit_t = NaN,), []
	end

	training_info = (
		fit_t = fit_t,
		model = nothing
		)

	training_info, [(x -> predict(model, x, pct=p), merge(parameters, Dict(:percentile => p))) for p in [10, 25, 50]]
end

function remove_constant_features(data)
	X = data[1][1]
	mask = (maximum(X, dims=2) .== minimum(X, dims=2))[:]
	if any(mask)
		@info "Removing $(sum(mask)) features with constant values."
		return Tuple((data[i][1][.~mask,:], data[i][2]) for i in 1:3)
	end
	data
end

try_counter = 0
max_tries = 10*max_seed
cont_string = (contamination == 0.0) ? "" : "_contamination-$contamination"
sampling_string = sampling == "bayes" ? "_bayes" : ""
prefix = "experiments$(sampling_string)/tabular$(cont_string)"
dataset_folder = datadir("$(prefix)/$(modelname)/$(dataset)")
while try_counter < max_tries
	if sampling == "bayes"
		parameters = GenerativeAD.bayes_params(
								create_space(),
								dataset_folder,
								sample_params)
	else
		parameters = sample_params()
	end

	for seed in 1:max_seed
		savepath = joinpath(dataset_folder, "seed=$(seed)")
		mkpath(savepath)

		# get data
		data = GenerativeAD.load_data(dataset, seed=seed, contamination=contamination)
		data = remove_constant_features(data)
		edited_parameters = sampling == "bayes" ? parameters : GenerativeAD.edit_params(data, parameters)

		if GenerativeAD.check_params(savepath, edited_parameters)
			@info "Started training PIDForest$(edited_parameters) on $(dataset):$(seed)"

			training_info, results = fit(data, edited_parameters)
			save_entries = merge(training_info, (modelname = modelname, seed = seed, dataset = dataset, contamination = contamination))

			all_scores = [GenerativeAD.experiment(result..., data, savepath; save_entries...) for result in results]
			if sampling == "bayes" && length(all_scores) > 0
				@info("Updating cache with $(length(all_scores)) results.")
				GenerativeAD.update_bayes_cache(dataset_folder,
						all_scores; ignore=Set([:percentile]))
			end
			global try_counter = max_tries + 1
		else
			@info "Model already present, trying new hyperparameters..."
			global try_counter += 1
		end
	end
end
(try_counter == max_tries) ? (@info "Reached $(max_tries) tries, giving up.") : nothing


text:
@assert begin
    acft = Aircraft()

    U∞ = 100.0
    ρ = 1.2

    b = 2.0
    c = 0.5

    xCG = 0.1
    e = 0.2

    θc = 1e-2
    hc = 1e-2
    θdc = 1.0
    hdc = 2.0

    EI = 1e4
    GJ = 1e4

    m = 2.0
    Ixx = m * c ^ 2 / 12

    s1 = WingSectInfo(
        [0.0, - b / 2, 0.0],
        c;
        e = e,
        xCG = xCG,
        EIy = EI,
        GJ = GJ,
        m = m,
        Ixx = Ixx
    )
    s2 = WingSectInfo(
        [0.0, 0.0, 0.0],
        c;
        e = e,
        xCG = xCG,
        EIy = EI,
        GJ = GJ,
        m = m,
        Ixx = Ixx
    )
    s3 = WingSectInfo(
        [0.0, b / 2, 0.0],
        c;
        e = e,
        xCG = xCG,
        EIy = EI,
        GJ = GJ,
        m = m,
        Ixx = Ixx
    )

    interpolate!(acft, [0.0, 1.0], s1, s2)
    interpolate!(acft, [0.0, 1.0], s2, s3)

    q = get_state(acft, U∞)

    set_θy(q, 2, θc)
    set_w(q, 2, hc)
    set_̇θy(q, 2, θdc)
    set_̇w(q, 2, hdc)

    qd, fcon = state_space(acft, q, U∞; ρ = ρ, structural_damping = false)

    α_AC = get_̇θy(qd, 2)
    a_AC = get_̇w(qd, 2)

    α_CG_num = α_AC
    a_CG_num = a_AC - α_AC * c * xCG

    # analytical

    m = (b / 2) * m
    Ixx = (b / 2) * Ixx

    L = b / 2

    Kh = 2 * 12 * EI / L ^ 3
    Kθ = 2 * GJ / L

    S = b * c / 2

    pdyn = ρ * (U∞ ^ 2 + hdc ^ 2) / 2
    Laed = ((fcon.sectional.CL[1] + fcon.sectional.CL[2]) / 2) * S * pdyn
    Maed = ((fcon.sectional.Cm[1] + fcon.sectional.Cm[2]) / 2) * S * pdyn

    h_e = - c * e * θc + hc

    M_e = - Kθ * θc
    L_e = - Kh * h_e

    M_CG = Maed + Laed * c * xCG + M_e + L_e * (xCG - e) * c
    L_CG = Laed + L_e

    a_CG = L_CG / m
    α_CG = M_CG / Ixx

    # Gotta make these values match!
    @assert isapprox(a_CG, a_CG_num; rtol = 1e-3)
    @assert isapprox(α_CG, α_CG_num; rtol = 1e-3)

    # plot_aircraft(acft, q)

    true
end

max_stars_repo_name: Equipe-AeroDesign-ITA/WingBiology

max_issues_repo_name: Equipe-AeroDesign-ITA/WingBiology

max_forks_repo_name: Equipe-AeroDesign-ITA/WingBiology


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

using ValueShapes
using Test

using Random
using ElasticArrays
using ArraysOfArrays
import TypedTables

@testset "abstract_value_shape" begin
    @testset "default_datatype" begin
        @test @inferred(ValueShapes.default_datatype(Integer)) == Int
        @test @inferred(ValueShapes.default_datatype(AbstractFloat)) == Float64
        @test @inferred(ValueShapes.default_datatype(Real)) == Float64
        @test @inferred(ValueShapes.default_datatype(Complex)) == Complex{Float64}

        @test @inferred(elshape(Complex(1.0, 2.0))) == ScalarShape{Complex{Float64}}()
        @test @inferred(elshape([[3, 5], [3, 2]])) == ArrayShape{Int,1}((2,))

        arrshape = ArrayShape{Real, 2}((2,3))
        vec = Vector{Real}(undef, arrshape)
        @test @inferred(length(vec)) == 6
        @test @inferred(size(vec)) == (6,)

        data1 = [1;2;3;4;7;8;9]
        scalarshape = ScalarShape{Real}()
        ntshape = NamedTupleShape(a=arrshape, b=scalarshape)
        shapedasnt = ntshape(data1)
        @test stripscalar(Ref(shapedasnt)) == Ref(shapedasnt)[]

        @test_throws ArgumentError Broadcast.broadcastable(ntshape)

        named_shapes = (
            a = ArrayShape{Real}(2, 3),
            b = ScalarShape{Real}(),
            c = ConstValueShape(4.2),
            x = ConstValueShape([11 21; 12 22]),
            y = ArrayShape{Real}(4)
        )
        shape = NamedTupleShape(;named_shapes...)
        data2 = VectorOfSimilarVectors(reshape(collect(1:22), 11, 2))
        @test_throws ArgumentError ValueShapes._checkcompat_inner(ntshape, data2)
        @test ValueShapes._checkcompat_inner(shape, data2) == nothing

    end
end


text:
module ITensorInfiniteMPS

# For `only`, which was introduced in Julia v1.4
using Compat
using ITensors
# For optional ITensorsVisualization dependency.
using Requires
# For using ∞ as lengths, ranges, etc.
using Infinities
# For functions like `isdiag`
using LinearAlgebra
# For indexing starting from something other than 0.
using OffsetArrays
using IterTools
# For HDF5 support
using HDF5
# For integration support when computing exact reference results
using QuadGK

using ITensors.NDTensors: eachdiagblock
using KrylovKit: eigsolve, linsolve, exponentiate

import Base: getindex, length, setindex!, +, -, *

import ITensors: AbstractMPS

include("ITensors.jl")
include("ITensorNetworks.jl")
include("itensormap.jl")
include("celledvectors.jl")
include("abstractinfinitemps.jl")
include("infinitemps.jl")
include("infinitempo.jl")
include("infinitecanonicalmps.jl")
include("transfermatrix.jl")
include("models/models.jl")
include("models/ising.jl")
include("models/ising_extended.jl")
include("models/heisenberg.jl")
include("models/hubbard.jl")
include("models/xx.jl")
include("orthogonalize.jl")
include("infinitemps_approx.jl")
include("nullspace.jl")
include("subspace_expansion.jl")
include("vumps_localham.jl")

export Cell,
  CelledVector,
  InfiniteMPS,
  InfiniteCanonicalMPS,
  InfMPS,
  InfiniteSum,
  InfiniteMPO,
  InfiniteSumLocalOps,
  ITensorMap,
  ITensorNetwork,
  TransferMatrix,
  @Model_str,
  Model,
  @Observable_str,
  Observable,
  input_inds,
  infinitemps_approx,
  infsiteinds,
  nsites,
  output_inds,
  reference,
  subspace_expansion,
  translatecell,
  tdvp,
  vumps,
  ⊕,
  ⊗,
  ×

function __init__()
  # This is used for debugging using visualizations
  @require ITensorsVisualization = "f2aed53d-2f32-47c3-a7b9-1ee253853786" @eval using ITensorsVisualization
end

end


text:
"""
Create `Evaluator` instance for a given `problem`.
"""
function make_evaluator(problem::OptimizationProblem, archive=nothing, params::Parameters=ParamsDict())
    workers = get(params, :Workers, Vector{Int}())
    nthreads = get(params, :NThreads, 0)
    if archive===nothing
        # make the default archive
        archiveCapacity = get(params, :ArchiveCapacity, 10)
        archive = TopListArchive(fitness_scheme(problem), numdims(problem), archiveCapacity)
    end
    if length(workers) > 0
        return ParallelEvaluator(problem, archive, pids=workers)
    elseif nthreads > 0
        return MultithreadEvaluator(problem, archive, nworkers=nthreads)
    else
        return ProblemEvaluator(problem, archive)
    end
end

"""
Optimization Run Controller.
Manages problem optimization using the specified method.

See `OptController`.
"""
mutable struct OptRunController{O<:Optimizer, E<:Evaluator}
    optimizer::O   # optimization algorithm
    evaluator::E   # problem evaluator

    trace_mode::Symbol # controller state trace mode (:verbose, :compact, :silent)
                       # :silent makes tr() generate no output)
    save_trace::Bool # FIXME if traces should be saved to a file
    trace_interval::Float64 # periodicity of calling trace_progress()

    callback_function::Function
    callback_interval::Float64

    # termination conditions
    max_steps::Int      # maximal number of steps
    max_fevals::Int # maximal number of fitness evaluations
    max_steps_without_fevals::Int # stop optimization if no func evals in this many steps (indicates a converged/degenerate search)
    max_steps_without_progress::Int # stop optimization if no improvement in this many steps (indicates a converged/degenerate search)
    max_time::Float64   # maximal time, 0 to ignore

    min_delta_fitness_tol::Float64 # minimal difference between current best fitness and second-best one
    fitness_tol::Float64  # minimal difference between current best fitness and the known optmimum

    num_steps::Int       # optimization steps done
    num_better::Int      # number of steps that improved best fitness

    num_better_since_last_report::Int # ditto, since last call to trace_progress()
    num_steps_since_last_report::Int # number of steps since last call to trace_progress()

    last_num_fevals::Int # the number of function evals on the previous step
    num_steps_without_fevals::Int # the number of steps without the function evals

    start_time::Float64 # time optimization started, 0 if not running yet
    stop_time::Float64  # time optimization stopped, 0 if still running
    last_report_time::Float64 # last time trace_progress() was called
    last_callback_time::Float64 # last time callback function was called

    stop_reason::String # the reason for algorithm termination, empty if it's not terminated
    shutdown::Bool # shutdown optimization run iff true
end

"""
Create `OptRunController` for a given problem using specified `optimizer`.

#Arguments
    * `optimizer` initialized optimization method
    * `evaluator` the evaluator of the problem fitness
    * `params` controller settings, see `DefaultParameters` for the default values:
        * `:MaxTime` max time in seconds (takes precedence over the other budget-related params if specified), 0.0 disables the check
        * `:MaxFuncEvals` max fitness evals (takes precedence over max iterations, but not max time), 0 disables the check
        * `:MaxSteps` max iterations gives the least control since different optimizers have different "size" of their "iterations"
        * `:MaxStepsWithoutProgress` max iterations without fitness improvement
        * `:MinDeltaFitnessTolerance` minimum delta fitness (difference between the two consecutive best fitness improvements) we can accept before terminating
        * `:FitnessTolerance` stop the optimization when the best fitness found is within this distance of the actual optimum (if known)
        * `:MaxNumStepsWithoutFuncEvals` stop optimization if no new fitness evals in this many steps (indicates a converged/degenerate search)
        * `:NumRepetitions` number of repetitions to run for each optimizer for each problem
        * `:TraceMode` how the optimizer state is traced to the STDOUT during the optimization (one of `:silent`, `:verbose`)
        * `:TraceInterval` the trace interval (in seconds)
        * `:SaveTrace` whether to save it to a file (defaults to `false`)
        * `:SaveFitnessTraceToCsv` whether the history of fitness changes during optimization should be save to a csv file
        * `:SaveParameters` save method/controller parameters to a JSON file
"""
function OptRunController(optimizer::O, evaluator::E, params) where {O<:Optimizer, E<:Evaluator}
    OptRunController{O,E}(optimizer, evaluator,
        [params[key] for key in Symbol[:TraceMode, :SaveTrace, :TraceInterval,
                      :CallbackFunction, :CallbackInterval,
                      :MaxSteps, :MaxFuncEvals, :MaxNumStepsWithoutFuncEvals, :MaxStepsWithoutProgress, :MaxTime,
                      :MinDeltaFitnessTolerance, :FitnessTolerance]]...,
        0, 0, 0, 0, 0, 0, 0.0, 0.0, 0.0, -1.0, "", false)
end

# stepping optimizer has it's own evaluator, get a reference
OptRunController(optimizer::SteppingOptimizer, problem::OptimizationProblem, params) =
    OptRunController(optimizer, evaluator(optimizer), params)
# all other optimizers are using make_evaluator() method to create evaluator by default
OptRunController(optimizer::Optimizer, problem::OptimizationProblem, params) =
    OptRunController(optimizer, make_evaluator(problem, nothing, params), params)
# We reuse evaluator and optimizer from a previous run if set up from such a run.
OptRunController(orc::OptRunController, params) =
    OptRunController(optimizer(orc), evaluator(orc), params)

# logging/tracing
function trace(ctrl::OptRunController, msg::AbstractString, obj = nothing)
    if ctrl.trace_mode != :silent
        print(msg)
        if obj !== nothing
            if isa(obj, AbstractString)
                print(obj)
            else
                show(IOContext(stdout, :compact => true), obj)
            end
        end
    end
    if ctrl.save_trace
        # No saving for now
    end
end

optimizer(ctrl::OptRunController) = ctrl.optimizer
evaluator(ctrl::OptRunController) = ctrl.evaluator
problem(ctrl::OptRunController) = problem(evaluator(ctrl))
isstarted(ctrl::OptRunController) = ctrl.start_time > 0
isrunning(ctrl::OptRunController) = isstarted(ctrl) && ctrl.stop_time == 0
isstopped(ctrl::OptRunController) = ctrl.stop_time > 0

num_steps(ctrl::OptRunController) = ctrl.num_steps
num_func_evals(ctrl::OptRunController) = num_evals(ctrl.evaluator)
stop_reason(ctrl::OptRunController) = ctrl.stop_reason
start_time(ctrl::OptRunController) = ctrl.start_time

best_candidate(ctrl::OptRunController) = best_candidate(ctrl.evaluator.archive)
best_fitness(ctrl::OptRunController) =  best_fitness(ctrl.evaluator.archive)

"""
    show_fitness(io, fit, [problem::OptimizationProblem])

Output fitness to the given IO stream.
`show_fitness()` method could be overridden for a specific problem, e.g.
to print the names of each objective.
"""
show_fitness(io::IO, fit::Number) = @printf(io, "%.9f", fit)

function show_fitness(io::IO, fit::NTuple{N}) where N
    print(io, "(")
    for i in 1:N
        if i > 1
            print(io, ", ")
        end
        @printf(io, "%.5f", fit[i])
    end
    print(io, ")")
end

function show_fitness(io::IO, fit::IndexedTupleFitness{N}) where N
    show_fitness(io, fit.orig)
    @printf(io, " agg=%.5f", fit.agg)
end

# problem-specific method defaults to problem-agnostic method
# note that the fitness type of the problem could be different from FA,
# because show_fitness() typically is called for the archived fitness
show_fitness(io::IO, fit::FA, problem::OptimizationProblem) where {FA} = show_fitness(io, fit)

"""
    format_fitness(fit, [problem::OptimizationProblem])

Format fitness into a string.
Calls `show_fitness()` under the hood.
"""
function format_fitness(fit::Any, problem::OptimizationProblem)
    buf = IOBuffer(false, true)
    show_fitness(buf, fit, problem)
    return bytestring(buf.data)
end

function format_fitness(fit::Any)
    buf = IOBuffer(false, true)
    show_fitness(buf, fit)
    return bytestring(buf.data)
end

function elapsed_time(ctrl::OptRunController)
    isrunning(ctrl) ? time() - ctrl.start_time : ctrl.stop_time - ctrl.start_time
end

function check_stop_condition(ctrl::OptRunController)
    if ctrl.shutdown
        return "Run explicitly stopped via shutdown method"
    end

    if ctrl.max_time > 0 && elapsed_time(ctrl) > ctrl.max_time
        return "Max time ($(ctrl.max_time) s) reached"
    end

    if ctrl.max_fevals > 0 && num_func_evals(ctrl) > ctrl.max_fevals
        return "Max number of function evaluations ($(ctrl.max_fevals)) reached"
    end

    if ctrl.max_steps_without_fevals > 0 && ctrl.num_steps_without_fevals > ctrl.max_steps_without_fevals
        return "Too many steps ($(ctrl.num_steps_without_fevals)) without any function evaluations (probably search has converged)"
    end

    if ctrl.max_steps > 0 && num_steps(ctrl) > ctrl.max_steps
        return "Max number of steps ($(ctrl.max_steps)) reached"
    end

    return check_stop_condition(ctrl.evaluator, ctrl)

    return "" # empty string, no termination
end

@inline function callback(ctrl::OptRunController)
    if ctrl.callback_interval >= 0.0
        ctrl.callback_function(ctrl)
        ctrl.last_callback_time = time()
    end
end

function trace_progress(ctrl::OptRunController)
    # update the counters
    ctrl.last_report_time = time()
    total_improvement_rate = ctrl.num_better/num_steps(ctrl)
    recent_improvement_rate = ctrl.num_better_since_last_report/ctrl.num_steps_since_last_report
    ctrl.num_better_since_last_report = 0
    ctrl.num_steps_since_last_report = 0

    if ctrl.trace_mode == :silent
        return
    end

    # Always print step number, num fevals and elapsed time
    trace(ctrl, @sprintf("%.2f secs, %d evals, %d steps",
            elapsed_time(ctrl), num_func_evals(ctrl), num_steps(ctrl)))

    # Only print if this optimizer reports on number of better. They return 0
    # if they do not.
    if total_improvement_rate > 0.0
        trace(ctrl, @sprintf(", improv/step: %.3f (last = %.4f)",
                total_improvement_rate, recent_improvement_rate))
    end

    # Always print fitness if num_evals > 0
    if num_func_evals(ctrl) > 0
        trace(ctrl, ", fitness=")
        show_fitness(stdout, best_fitness(ctrl), problem(ctrl))
    end

    trace(ctrl, "\n")

    trace_state(stdout, ctrl.optimizer, ctrl.trace_mode)
end

function step!(ctrl::OptRunController{<:AskTellOptimizer})
    # The ask()/tell() interface is more general since you can mix and max
    # elements from several optimizers using it. However, in this top-level
    # execution function we do not make use of this flexibility...
    candidates = ask(ctrl.optimizer)
    rank_by_fitness!(ctrl.evaluator, candidates)
    return tell!(ctrl.optimizer, candidates)
end

# step for SteppingOptimizers
function step!(ctrl::OptRunController{<:SteppingOptimizer})
    step!(ctrl.optimizer)
    return 0
end

setup_optimizer!(ctrl::OptRunController{<:SteppingOptimizer}) =
    setup!(ctrl.optimizer)
setup_optimizer!(ctrl::OptRunController{<:AskTellOptimizer}) =
    setup!(ctrl.optimizer, ctrl.evaluator)

shutdown!(ctrl::OptRunController) = ctrl.shutdown = true

function shutdown_optimizer!(ctrl::OptRunController{<:SteppingOptimizer})
    shutdown!(ctrl.optimizer)
    shutdown!(ctrl)
end

function shutdown_optimizer!(ctrl::OptRunController{<:AskTellOptimizer})
    shutdown!(ctrl.optimizer)
    shutdown!(ctrl.evaluator)
    shutdown!(ctrl)
end

"""
    run!(ctrl::OptRunController)

Run optimization until one of the stopping conditions are satisfied.
"""
function run!(ctrl::OptRunController)
    trace(ctrl, "Starting optimization with optimizer $(name(ctrl.optimizer))\n")
    try
        setup_optimizer!(ctrl)

        ctrl.start_time = time()
        ctrl.num_steps = 0
        while isempty(ctrl.stop_reason)
            # Report on progress every now and then...
            if (time() - ctrl.last_report_time) > ctrl.trace_interval
                trace_progress(ctrl)
            end

            # Take the step and then update the counters
            nstep_better = step!(ctrl)
            ctrl.num_better += nstep_better
            ctrl.num_better_since_last_report += nstep_better
            ctrl.num_steps += 1
            ctrl.num_steps_since_last_report += 1
            if num_evals(ctrl.evaluator) == ctrl.last_num_fevals
                ctrl.num_steps_without_fevals += 1
            else
                ctrl.num_steps_without_fevals = 0
            end
            ctrl.last_num_fevals = num_func_evals(ctrl)

            # Callback every now and then (if a callback interval has been set)...
            if ctrl.callback_interval >= 0.0
                if ctrl.callback_interval == 0.0 || (ctrl.last_callback_time <= 0.0 ||
                        (time() - ctrl.last_callback_time) > ctrl.callback_interval)
                    callback(ctrl)
                end
            end

            # Check if we have a reason to stop on next loop
            ctrl.stop_reason = check_stop_condition(ctrl)
        end
        ctrl.stop_time = time()
    finally
        shutdown_optimizer!(ctrl)
    end
    trace(ctrl, @sprintf "\nOptimization stopped after %d steps and %.2f seconds\n" ctrl.num_steps elapsed_time(ctrl))

    return ctrl.stop_reason
end

function show_report(ctrl::OptRunController, population_stats=false)
    final_elapsed_time = elapsed_time(ctrl)
    trace(ctrl, "Termination reason: $(ctrl.stop_reason)\n")
    trace(ctrl, @sprintf "Steps per second = %.2f\n" num_steps(ctrl)/final_elapsed_time)
    trace(ctrl, @sprintf "Function evals per second = %.2f\n" num_func_evals(ctrl)/final_elapsed_time)
    trace(ctrl, @sprintf "Improvements/step = %.5f\n" ctrl.num_better/ctrl.max_steps)
    trace(ctrl, "Total function evaluations = $(num_func_evals(ctrl))\n")

    if population_stats && isa(ctrl.optimizer, PopulationOptimizer)
        trace(ctrl, "\nMean value (in population) per position:",  mean(population(ctrl.optimizer),1))
        trace(ctrl, "\n\nStd dev (in population) per position:", std(population(ctrl.optimizer),1))
    end

    trace(ctrl, "\n\nBest candidate found: ", best_candidate(ctrl))
    trace(ctrl, "\n\nFitness: ")
    show_fitness(stdout, best_fitness(ctrl), problem(ctrl))
    trace(ctrl, "\n\n")
end

function write_result(ctrl::OptRunController, filename = "")
    if isempty(filename)
        timestamp = Libc.strftime("%y%m%d_%H%M%S", floor(Int, ctrl.start_time))
        filename = "$(timestamp)_$(problem_summary(ctrl.evaluator))_$(name(ctrl.optimizer)).csv"
        filename = replace(replace(filename, r"\s+", "_"), r"/", "_")
    end
    save_fitness_history_to_csv_file(ctrl.evaluator.archive, filename;
            header_prefix = "Problem,Dimension,Optimizer",
            line_prefix = "$(name(problem(ctrl.evaluator))),$(numdims(ctrl.evaluator)),$(name(ctrl.optimizer))",
            bestfitness = opt_value(problem(ctrl.evaluator)))
end

"""
Optimization Controller.

Applies specific optimization method to a given problem.
Supports restarts and modifying parameter of the method between runs.
`runcontrollers` field maintains the list of `OptRunController` instances applied so far.

See `OptRunController`.
"""
mutable struct OptController{O<:Optimizer, P<:OptimizationProblem}
    optimizer::O   # optimization algorithm
    problem::P     # opt problem
    parameters::ParamsDictChain
    runcontrollers::Vector{OptRunController{O}}
end

"""
Create `OptController` for a given `optimizer` and a `problem`.

`params` are any of `OptRunController` parameters plus
    * `:RngSeed` and `:RandomizeRngSeed` params for controlling the random seed
    * `:RecoverResults` if intermediate results are returned upon `InterruptException()` (on by default)
"""
OptController(
    optimizer::O, problem::P,
    params::ParamsDictChain) where {O<:Optimizer, P<:OptimizationProblem} =
    OptController{O, P}(optimizer, problem, params, OptRunController{O}[])

problem(oc::OptController) = oc.problem
numruns(oc::OptController) = length(oc.runcontrollers)
lastrun(oc::OptController) = oc.runcontrollers[end]

"""
    update_parameters!(oc::OptController, parameters::Associative)

Update the `OptController` parameters.
"""
function update_parameters!(oc::OptController, parameters::Parameters = EMPTY_DICT)
    parameters = convert(ParamsDict, parameters)

    # Most parameters cannot be changed since the problem and optimizer has already
    # been setup.
    valid_params = Set([:MaxTime, :MaxSteps, :MaxFuncEvals, :TraceMode])
    for k in keys(parameters)
        if k ∉ valid_params
            throw(ArgumentError("It is currently not supported to change parameters that can affect the original opt problem or optimizer (here: $(k))"))
        end
    end

    # Add new params in front if any are specified and they are valid.
    if length(parameters) > 0
        oc.parameters = chain(oc.parameters, parameters)
        check_valid!(oc.parameters) # We must recheck that new param settings are valid
    end
end

function init_rng!(parameters::Parameters)
    if parameters[:RandomizeRngSeed]
        parameters[:RngSeed] = rand(1:1_000_000)
        Random.seed!(parameters[:RngSeed])
    end
end

"""
    run!(oc::OptController)

Start a new optimization run, possibly with new parameters and report on results.
"""
function run!(oc::OptController)
    # If this is not the first run we create the controller based on the previous one
    # to reuse the archive etc, otherwise create it based on the problem.
    ctrl = if numruns(oc) > 0
        OptRunController(lastrun(oc), oc.parameters)
    else
        # No previous run controller so create a new one from optimizer and problem.
        OptRunController(oc.optimizer, oc.problem, oc.parameters)
    end
    push!(oc.runcontrollers, ctrl)

    # If this is the first run we might have to init the RNG.
    if numruns(oc) == 1
        init_rng!(oc.parameters)
    end

    try
        run!(ctrl)
        if ctrl.trace_mode != :silent
            show_report(ctrl)
        end

        if oc.parameters[:SaveFitnessTraceToCsv]
            write_result(ctrl)
        end

        return OptimizationResults(ctrl, oc)
    catch ex
        # If it was a ctrl-c interrupt we try to make a result and return it...
        if get(oc.parameters, :RecoverResults, true) && isa(ex, InterruptException)
            @warn("Optimization interrupted, recovering intermediate results...")
            ctrl.stop_reason = @sprintf("%s", ex)
            return OptimizationResults(ctrl, oc)
        else
            rethrow(ex)
        end
    end
end


text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

using JuliaFEM
using JuliaFEM.Testing

@testset "polygon clip case 1" begin
    S = Vector[
        [0.375, 0.0, 0.5],
        [0.6, 0.0, 0.5],
        [0.5, 0.25, 0.5]]
    M = Vector[
        [0.50, 0.0, 0.5],
        [0.25, 0.0, 0.5],
        [0.375, 0.25, 0.5]]
    n0 = [0.0, 0.0, 1.0]
    P = get_polygon_clip(S, M, n0)
    P_expected = Vector{Float64}[
        [0.500, 0.0, 0.5],
        [0.375, 0.0, 0.5],
        [0.4375, 0.125, 0.5]]
    @test length(P) == length(P_expected)
    for (Pi, Pj) in zip(P, P_expected)
        @test isapprox(Pi, Pj)
    end
end

@testset "polygon clip case 2" begin
    S = Vector[
        [0.25, 0.0, 0.5],
        [0.75, 0.0, 0.5],
        [0.50, 0.25, 0.5]]
    M = Vector[
        [0.50, 0.0, 0.5],
        [0.25, 0.0, 0.5],
        [0.375, 0.25, 0.5]]
    n0 = [0.0, 0.0, 1.0]
    P = get_polygon_clip(S, M, n0)
    @test length(P) == 3
end


text:
function binarysearch(arr, val; lt = <, rev = false, by = identity)
    n = length(arr)

    if n == 0
        return 1:0
    end

    val = by(val)
    arr = by.(arr)
    compare = rev ? (a, b) -> lt(b, a) : lt

    i = 1
    while n > 1
        m = div(n, 2)
        if compare(val, arr[m+1])
            arr = arr[1:m]
        else
            arr = arr[(m+1):end]
            i += m
        end
        n = length(arr)
    end

    if compare(val, arr[1])
        return i:(i-1)
    elseif compare(arr[1], val)
        return (i + 1):i
    else
        return i:i
    end
end

max_stars_repo_name: twolodzko/Learning

max_issues_repo_name: twolodzko/Learning

max_forks_repo_name: twolodzko/Learning


text:
using Test, OMEinsumContractionOrders, OMEinsum
using Graphs, Random
using SparseArrays

@testset "sa bipartition" begin
    Random.seed!(3)
    g = random_regular_graph(120, 3)
    adj, = OMEinsumContractionOrders.adjacency_matrix([(e.src,e.dst) for e in edges(g)])
    ws = fill(log2(2), ne(g))
    vertices = 1:110
    b = SABipartite(βs=0.1:0.2:20.0, niters=1000, ntrials=100, sc_target=40, initializer=:random)
    g1, g2 = OMEinsumContractionOrders.bipartite_sc(b, adj, vertices, ws)
    @test length(g1) + length(g2) == 110
end

@testset "sa" begin
    Random.seed!(2)
    function random_regular_eincode(n, k)
        g = Graphs.random_regular_graph(n, k)
        ixs = [minmax(e.src,e.dst) for e in Graphs.edges(g)]
        return EinCode((ixs..., [(i,) for i in Graphs.vertices(g)]...), ())
    end

    g = random_regular_graph(220, 3)
    rows = Int[]
    cols = Int[]
    for (i,edge) in enumerate(edges(g))
        push!(rows, edge.src, edge.dst)
        push!(cols, i, i)
    end
    graph = sparse(rows, cols, ones(Int, length(rows)))
    sc_target = 28.0
    log2_sizes = fill(1, size(graph, 2))
    βs = 0.01:0.05:10.0
    b = SABipartite(βs=βs, niters=1000, ntrials=100, sc_target=28)
    group1, group2 = OMEinsumContractionOrders.bipartite_sc(b, graph, collect(1:size(graph, 1)), log2_sizes)
    @test OMEinsumContractionOrders.group_sc(graph, group1, log2_sizes) <= sc_target+2
    @test OMEinsumContractionOrders.group_sc(graph, group2, log2_sizes) <= sc_target+2
    sc_target = 27.0
    group11, group12 = OMEinsumContractionOrders.bipartite_sc(b, graph, group1, log2_sizes)
    @test OMEinsumContractionOrders.group_sc(graph, group11, log2_sizes) <= sc_target+2
    @test OMEinsumContractionOrders.group_sc(graph, group12, log2_sizes) <= sc_target+2

    code = random_regular_eincode(220, 3)
    res = optimize_sa(code,uniformsize(code, 2); sc_target=30, βs=βs)
    tc, sc = OMEinsum.timespace_complexity(res, uniformsize(code, 2))
    @test sc <= 32

    # contraction test
    code = random_regular_eincode(50, 3)
    codeg = optimize_sa(code, uniformsize(code, 2); sc_target=12, βs=βs, ntrials=1, initializer=:greedy)
    codek = optimize_greedy(code, uniformsize(code, 2))
    tc, sc = OMEinsum.timespace_complexity(codek, uniformsize(code, 2))
    @test sc <= 12
    xs = [[2*randn(2, 2) for i=1:75]..., [randn(2) for i=1:50]...]
    resg = codeg(xs...)
    resk = codek(xs...)
    @test resg ≈ resk
end


text:
# ---
# title: 1018. Binary Prefix Divisible By 5
# id: problem1018
# author: Indigo
# date: 2022-03-19
# difficulty: Easy
# categories: Array
# link: <https://leetcode.com/problems/binary-prefix-divisible-by-5/description/>
# hidden: true
# ---
#
# Given an array `A` of `0`s and `1`s, consider `N_i`: the i-th subarray from
# `A[0]` to `A[i]` interpreted as a binary number (from most-significant-bit to
# least-significant-bit.)
#
# Return a list of booleans `answer`, where `answer[i]` is `true` if and only if
# `N_i` is divisible by 5.
#
# **Example 1:**
#
#
#
#     Input: [0,1,1]
#     Output: [true,false,false]
#     Explanation:
#     The input numbers in binary are 0, 01, 011; which are 0, 1, and 3 in base-10.  Only the first number is divisible by 5, so answer[0] is true.
#
#
# **Example 2:**
#
#
#
#     Input: [1,1,1]
#     Output: [false,false,false]
#
#
# **Example 3:**
#
#
#
#     Input: [0,1,1,1,1,1]
#     Output: [true,false,false,false,true,false]
#
#
# **Example 4:**
#
#
#
#     Input: [1,1,1,0,1]
#     Output: [false,false,false,false,false]
#
#
#
#
# **Note:**
#
#   1. `1 <= A.length <= 30000`
#   2. `A[i]` is `0` or `1`
#
#
## @lc code=start
using LeetCode

function prefixes_div_by5(nums::Vector{Int})
    len = length(nums)
    res = fill(false, len)
    cur = 0
    for i in 1:len
        cur = cur * 2 + nums[i]
        cur % 5 == 0 && (res[i] = true)
    end
    return res
end
## @lc code=end