text:
"""
    SubspaceBasis(basisstates)

A basis describing a subspace embedded a higher dimensional Hilbert space.
"""
struct SubspaceBasis{S,B<:Basis,T<:Ket,H,UT} <: Basis
    shape::S
    superbasis::B
    basisstates::Vector{T}
    basisstates_hash::UT

    function SubspaceBasis{S,B,T,H,UT}(shape::S,superbasis::B,basisstates::Vector{T},basisstates_hash::UT) where {S,B,T,H,UT}
        new{S,B,T,H,UT}(shape,superbasis,basisstates,basisstates_hash)
    end
end
function SubspaceBasis(superbasis::B, basisstates::Vector{T}) where {B<:Basis,T<:Ket}
    for state = basisstates
        if state.basis != superbasis
            throw(ArgumentError("The basis of the basisstates has to be the superbasis."))
        end
    end
    H = hash(hash.([hash.(x.data) for x=basisstates]))
    shape = Int[length(basisstates)]
    SubspaceBasis{typeof(shape),B,T,H,typeof(H)}(shape, superbasis, basisstates, H)
end
SubspaceBasis(basisstates::Vector{T}) where T<:Ket = SubspaceBasis(basisstates[1].basis, basisstates)

==(b1::SubspaceBasis, b2::SubspaceBasis) = b1.superbasis==b2.superbasis && b1.basisstates_hash==b2.basisstates_hash

proj(u::Ket, v::Ket) = dagger(v)*u/(dagger(u)*u)*u

"""
    orthonormalize(b::SubspaceBasis)

Orthonormalize the basis states of the given [`SubspaceBasis`](@ref)

A modified Gram-Schmidt process is used.
"""
function orthonormalize(b::SubspaceBasis)
    V = b.basisstates
    U = Ket[]
    for (k, v)=enumerate(V)
        u = copy(v)
        for i=1:k-1
            u -= proj(U[i], u)
        end
        normalize!(u)
        push!(U, u)
    end
    return SubspaceBasis(U)
end

"""
    projector(b1, b2; sparse=false)

Projection operator between subspaces and superspaces or between two subspaces.
"""
function projector(b1::SubspaceBasis, b2::SubspaceBasis; kwargs...)
    if b1.superbasis != b2.superbasis
        throw(ArgumentError("Both subspace bases have to have the same superbasis."))
    end
    T1 = projector(b1, b1.superbasis; kwargs...)
    T2 = projector(b2.superbasis, b2; kwargs...)
    return T1*T2
end

function projector(b1::SubspaceBasis, b2::Basis; sparse=false)
    if b1.superbasis != b2
        throw(ArgumentError("Second basis has to be the superbasis of the first one."))
    end
    T = promote_type(eltype.(b1.basisstates)...)
    if sparse
        data = spzeros(T, length(b1), length(b2))
    else
        data = zeros(T, length(b1), length(b2))
    end
    for (i, state) = enumerate(b1.basisstates)
        data[i,:] = state.data
    end
    return Operator(b1, b2, data)
end

function projector(b1::Basis, b2::SubspaceBasis; sparse=false)
    if b1 != b2.superbasis
        throw(ArgumentError("First basis has to be the superbasis of the second one."))
    end
    T = promote_type(eltype.(b2.basisstates)...)
    if sparse
        data = spzeros(T, length(b1), length(b2))
    else
        data = zeros(T, length(b1), length(b2))
    end
    for (i, state) = enumerate(b2.basisstates)
        data[:,i] = state.data
    end
    return Operator(b1, b2, data)
end


text:
import HCubature
export HCubatureAlg

struct HCubatureAlg{K}
    kw::K
end

function HCubatureAlg(;kw...)
    HCubatureAlg(kw)
end

function integral_kernel(f, dom::Domain, alg::HCubatureAlg)
    value, err = HCubature.hcubature(f, dom.lower, dom.upper; alg.kw...)
    (value=value, std=err) # TODO err is not really a standard deviation, what should be done?
end

ext: jl


text:
using PyPlot
using JLD
using Optim

# type Measurements{T}
#     i::Int64                # measurement counter
#     t::Array{Float64}       # time data (when it was received by this recorder)
#     t_msg::Array{Float64}   # time that the message was sent
#     z::Array{T}             # measurement values
# end

function main(code::AbstractString)
    #log_path_record = ("$(homedir())/open_loop/output-record-8bbb.jld","$(homedir())/open_loop/output-record-3593.jld")
    log_path_record = "$(homedir())/open_loop/output-record-$code.jld"
    d_rec = load(log_path_record)

    cmd_pwm_log = d_rec["cmd_pwm_log"]
    vel_est     = d_rec["vel_est"]
    pos_info    = d_rec["pos_info"]

    t0 = max(cmd_pwm_log.t[1],vel_est.t[1])
    t_end = min(cmd_pwm_log.t[end],vel_est.t[end])

    t = t0+0.1:.02:t_end-0.1
    v = zeros(length(t))
    cmd = zeros(length(t))

    for i=1:length(t)
        v[i] = vel_est.z[t[i].>vel_est.t,1][end]
        #v[i] = pos_info.z[t[i].>pos_info.t,15][end]
        cmd[i] = cmd_pwm_log.z[t[i].>cmd_pwm_log.t,1][end]
    end
    v_opt   = v[1:end]
    cmd_opt = cmd[1:end]
    v_opt   = v[cmd.>90]
    cmd_opt = cmd[cmd.>90]

    res = optimize(c->cost(cmd_opt,v_opt,c,0),[0.001,0,0.1,0.001,0])
    c = Optim.minimizer(res)
    c2 = [1/c[1], c[2]/c[1], 1/c[4], c[5]/c[4]]
    cost(cmd_opt,v_opt,c,1)
    println("c = $c")
    println("and c_2 = $c2")
end

function v_over_u(code::AbstractString)
    log_path_record = "$(homedir())/open_loop/output-record-$code.jld"
    d_rec = load(log_path_record)

    cmd_pwm_log = d_rec["cmd_pwm_log"]
    vel_est     = d_rec["vel_est"]
    pos_info    = d_rec["pos_info"]
    imu_meas    = d_rec["imu_meas"]

    t0 = max(cmd_pwm_log.t[1],vel_est.t[1])
    t_end = min(cmd_pwm_log.t[end],vel_est.t[end])

    L_b = 0.125
    t = t0+0.1:.02:t_end-0.1
    v = zeros(length(t))
    cmd = zeros(length(t),2)
    psiDot = zeros(length(t))

    for i=1:length(t)
        v[i] = vel_est.z[t[i].>vel_est.t,1][end]
        #v[i] = pos_info.z[t[i].>pos_info.t,15][end]
        cmd[i,:] = cmd_pwm_log.z[t[i].>cmd_pwm_log.t,:][end,:]
        psiDot[i] = imu_meas.z[t[i].>imu_meas.t,3][end]
    end
    v_x = real(sqrt(complex(v.^2-psiDot.^2*L_b^2)))
    v_y = L_b*psiDot
    delta = atan2(psiDot*0.25,v_x)
    c = ((cmd[:,1]-91)/6.5-0.5*v)./delta.^2
    c = c[abs(c).<100]
    plot(c)
    grid("on")
    figure()
    plot(v)
    grid("on")
    #plot(cmd[:,1],v)
end

function cost(cmd,v,c,p)
    cost = 0
    v_s = zeros(length(v))
    v_s[1] = v[1]
    v_s[2] = v[2]
    for i=2:length(cmd)-1
        if cmd[i]-cmd[i-1] != 0
            v_s[i] = v[i]
        end
        if cmd[i] > 90
            v_s[i+1] = v_s[i] + 0.02*(c[1]*cmd[i] + c[2] - c[3]*v_s[i])
        else
            v_s[i+1] = max(v_s[i] + 0.02*(c[4]*cmd[i] + c[5] - c[3]*v_s[i]),0)
        end
    end
    if p==1
        plot((1:length(v_s))*0.02,v_s,(1:length(v))*0.02,v)
        grid("on")
        legend(["v_sim","v_meas"])
        title("Curve fitting acceleration")
        xlabel("t [s]")
        ylabel("v [m/s]")
    end
    cost = norm(v_s-v)
    return cost
end

max_stars_repo_name: RoyaFiroozi/barc

max_issues_repo_name: RoyaFiroozi/barc

max_forks_repo_name: RoyaFiroozi/barc


text:
################################################################################
#
#	ABSTRACT TYPE
#
#   BondSpinHamiltonianKitaev <: AbstractBondHamiltonian{L,3}
#   --> L is the label type of bonds
#   --> N is the dimension the bond term matrix (NxN matrix)
#
#   FILE CONTAINS
#       - abstract type definition
#       - interface definition
#       - generator function
#
################################################################################

################################################################################
#
#   ABSTRACT TYPE DEFINITION
#
################################################################################
mutable struct BondSpinHamiltonianKitaev{L} <: AbstractBondHamiltonian{L,3}

    # couplings
    Jx :: Float64
    Jy :: Float64
    Jz :: Float64

    # bond labels of these couplings
    x_bonds :: Vector{L}
    y_bonds :: Vector{L}
    z_bonds :: Vector{L}

end

# export the type
export BondSpinHamiltonianKitaev

################################################################################
#
#	INTERFACING / ACCESSING BOND HAMILTONIANS
#	(functions have to be overwritten by concrete types)
#
################################################################################

# get bond term of Hamiltonian
function bondterm(
            h :: BondSpinHamiltonianKitaev{L},
            b :: AbstractBond{L,NB}
        ) :: Matrix{Complex{Float64}} where {L,NB}

    # get the bond label
    l = label(b)

    # define an empty coupling matrix
    coupling_matrix = zeros(3,3)

    # check for all neighbors if the label is in one of the categories
    if l in h.x_bonds
        coupling_matrix[1,1] += h.Jx
        return coupling_matrix
    elseif l in h.y_bonds
        coupling_matrix[2,2] += h.Jy
        return coupling_matrix
    elseif l in h.z_bonds
        coupling_matrix[3,3] += h.Jz
        return coupling_matrix
    end

    # if no bond found, return 0
    return coupling_matrix
end

# export the function
export bondterm

# FUNCTIONS TO FIND BOND LABELS

# FALLBACK ANY LABEL TYPE
# if specific function for label type L not found, denote all bonds as nearest neighbors
function getBondLabelsKitaevX(
        all_couplings :: Vector{L}
    ) :: Vector{L} where {L}

    # all couplings that somehow contain a x
    return filter(c->occursin('x',lowercase(string(c))) || c==getDefaultLabelX(L), all_couplings)
end
function getBondLabelsKitaevY(
        all_couplings :: Vector{L}
    ) :: Vector{L} where {L}

    # all couplings that somehow contain a y
    return filter(c->occursin('y',lowercase(string(c))) || c==getDefaultLabelY(L), all_couplings)
end
function getBondLabelsKitaevZ(
        all_couplings :: Vector{L}
    ) :: Vector{L} where {L}

    # all couplings that somehow contain a z
    return filter(c->occursin('z',lowercase(string(c))) || c==getDefaultLabelZ(L), all_couplings)
end

# CONSTRUCTION FUNCTION
function getSpinHamiltonianKitaev(
            unitcell :: U
        ) :: BondSpinHamiltonianKitaev{L} where {L,NB,S,B<:AbstractBond{L,NB},U<:AbstractUnitcell{S,B}}

    # obtain all couplings
    couplings = unique!(label.(bonds(unitcell)))

    # create default couplings
    Jx = 1.0
    Jy = 1.0
    Jz = 1.0
    # create list of labels
    x_bonds = getBondLabelsKitaevX(couplings)
    y_bonds = getBondLabelsKitaevY(couplings)
    z_bonds = getBondLabelsKitaevZ(couplings)

    # create and return a new object
    return BondSpinHamiltonianKitaev{L}(Jx,Jy,Jz,x_bonds,y_bonds,z_bonds)
end

# export the construction function
export getSpinHamiltonianKitaev

# CUSTOM SHOWING

# all labels
function show(io::IO, h::H) where {L,H<:BondSpinHamiltonianKitaev{L}}
    print(io, "Kitaev (Bond) Hamiltonian for nearest neighbors. Couplings strength and accepted labels are:")
    print(io, "\n--> Jx="*string(h.Jx)*", labels="*string(h.x_bonds))
    print(io, "\n--> Jy="*string(h.Jy)*", labels="*string(h.y_bonds))
    print(io, "\n--> Jz="*string(h.Jz)*", labels="*string(h.z_bonds))
end


text:
using PATH

const PTH_OPTIMIZER = PATH.Optimizer()#output = "no")
# MOI.set(PTH_OPTIMIZER, MOI.Silent(), true)
const PTH_CACHE = MOIU.UniversalFallback(MOIU.Model{Float64}())
const PTH_CACHED = MOIU.CachingOptimizer(PTH_CACHE, PTH_OPTIMIZER)
const PTH_BRIDGED = MOI.Bridges.full_bridge_optimizer(PTH_CACHED, Float64)

push!(solvers_nlp, (opt = PTH_BRIDGED, mode = BilevelJuMP.ComplementMode{Float64}(1e-9)))
# push!(solvers_nlp, (opt = PTH_BRIDGED, mode = BilevelJuMP.ProductWithSlackMode{Float64}(1e-9)))

push!(solvers, (opt = PTH_BRIDGED, mode = BilevelJuMP.ComplementMode{Float64}(1e-9)))
push!(solvers_quad, (opt = PTH_BRIDGED, mode = BilevelJuMP.ComplementMode{Float64}(1e-9)))


text:
##  Copyright 2022 Myrl Marmarelis
##
##  Licensed under the Apache License, Version 2.0 (the "License");
##  you may not use this file except in compliance with the License.
##  You may obtain a copy of the License at
##
##    http://www.apache.org/licenses/LICENSE-2.0
##
##  Unless required by applicable law or agreed to in writing, software
##  distributed under the License is distributed on an "AS IS" BASIS,
##  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
##  See the License for the specific language governing permissions and
##  limitations under the License.

# may or may not operate on a latent space from the covariates

using Flux
using Flux.Functors: @functor
# hitherto for differentiable models
abstract type Predictor{ In, Out, T<:Real } end

predict(::Predictor, ::AbstractVector) =
  error("unimplemented")
get_params(predictor::Predictor) = # ::Flux.Params
  Flux.params(predictor) # if @functor has been applied to the type

# idea: accommodate priors on the parametrization to do maximum a posteriori
#  ---with a hint of MaxEnt? favoring distributions of higher entropy

# predictor parametrizations
abstract type PredictionParam{ T<:Real } end

n_params(::PredictionParam) =
  error("unimplemented")
log_likelihood(::PredictionParam{T}, ::AbstractVector, ::T) where T =
  error("unimplemented")
sample_prediction(::PredictionParam{T}, ::AbstractVector) where T = # let the `AbstractVectors` containt `ForwardDiff.Dual`s
  error("optional unimplemented")
transform(::PredictionParam{T}, params::AbstractVector) where T =
  params

density(parametrization::PredictionParam{T}, params::AbstractVector, value::T) where T =
  log_likelihood(parametrization, params, value) |> exp
n_predictor_outputs(parametrization::PredictionParam, n_outputs::Int) =
  n_params(parametrization) * n_outputs
n_instance_outputs(::Predictor{In,POut}, parametrization::PredictionParam) where {In,POut} =
  div(POut, n_params(parametrization))

function log_likelihood(predictor::Predictor{In,POut,T},
    parametrization::PredictionParam{T}, instance::Instance{In,Out,T}
    )::T where {In,Out,POut,T}
  params = predict(predictor, instance.input) # do this in-place or as an SVector?
  param_size = n_params(parametrization)
  @assert Out * param_size == POut
  sum(instance.output|>enumerate) do (index, output)
    start = (index-1) * param_size + 1
    stop = start + param_size - 1
    slice = @view params[start:stop]
    trans_slice = transform(parametrization, slice)
    log_likelihood(parametrization, trans_slice, output)
  end
end

function map_prediction(func::Function, predictor::Predictor{In,POut,T},
    parametrization::PredictionParam{T}, input::SVector{In,T}) where {In,POut,T}
  params = predict(predictor, input)
  param_size = n_params(parametrization)
  Iterators.map(1:n_instance_outputs(predictor, parametrization)) do output_index
    start = (output_index-1) * param_size + 1
    stop = start + param_size - 1
    slice = @view params[start:stop]
    trans_slice = transform(parametrization, slice)
    func(trans_slice)
  end
end

function map_prediction(func::Function, predictor::Predictor{In,POut,T},
    parametrization::PredictionParam{T}, input::SVector{In}, index::Int) where {In,POut,T}
  params = predict(predictor, input)
  param_size = n_params(parametrization)
  start = (index-1) * param_size + 1
  stop = start + param_size - 1
  slice = @view params[start:stop]
  trans_slice = transform(parametrization, slice)
  func(trans_slice)
end

struct ConstPredictor{ In, Out, T } <: Predictor{In,Out,T}
  prediction :: MVector{Out, T}
end

ConstPredictor(In, Out, val) =
  ConstPredictor{In,Out,typeof(val)}(@MVector fill(val, Out))

@functor ConstPredictor

predict(model::ConstPredictor{In,Out,T}, ::AbstractVector) where {In,Out,T} =
  model.prediction

using Distributions
using Flux: softplus
using ChainRulesCore

# for SVector-type constructors that convert other `structure`s
function ChainRulesCore.rrule(constructor::Type{<:SArray}, structure)
  output = constructor(structure)
  project = ProjectTo(structure)
  pullback = tangent -> ( # `tangent` shares the structure (type?) of `output`
    NoTangent(), project(tangent) ) # basically convert tangent back to tuple or whatever
  output, pullback
end
#@Zygote.adjoint (V::Type{<:SArray})(x...) where T =
#  V(x...), y -> (y...,)

struct LomaxParam{ T } <: PredictionParam{T} end

n_params(::LomaxParam) = 2

function transform(::LomaxParam{T},
    (alpha, beta)::AbstractVector) where T
  softplus.((alpha, beta)) |> SVector{2} # represent logs directly? nah
end

function log_likelihood(::LomaxParam{T},
    (alpha, beta)::AbstractVector, value::T) where T
  log(alpha / beta) - (alpha+1) * log(1 + value/beta)
end

# exploit gamma-exponential mixture form
function sample_prediction(::LomaxParam{T}, (alpha, beta)::AbstractVector{T})::T where T
  gamma = Gamma(alpha, 1/beta)
  rate = rand(gamma)
  Exponential(1/rate) |> rand
end

# for a simple propensity model..
struct BetaParam{ T } <: PredictionParam{T}
  ceiling :: T
end

n_params(::BetaParam) = 2

using SpecialFunctions

function transform(param::BetaParam{T},
    (alpha, beta)::AbstractVector) where T
  curb = x -> param.ceiling * sigmoid(x/param.ceiling)
  curb.((alpha, beta)) |> SVector{2}
end

const beta_tolerance = 1e-2

function log_likelihood(::BetaParam{T},
    (alpha, beta)::AbstractVector, value::T) where T
  value = min(
    max(value,
      T(beta_tolerance) ),
    T(1-beta_tolerance) )
  #if value < beta_tolerance
  #  value += T(beta_tolerance)
  #elseif value > 1-beta_tolerance
  #  value -= T(beta_tolerance)
  #end
  (alpha-1)*log(value) + (beta-1)*log(1-value) - logbeta(alpha, beta)
end

function sample_prediction(::BetaParam{T}, (alpha, beta)::AbstractVector{T})::T where T
  Beta(alpha, beta) |> rand
end

struct GammaParam{ T } <: PredictionParam{T} end

n_params(::GammaParam) = 2

function transform(param::GammaParam{T},
    (alpha, beta)::AbstractVector) where T
  softplus.((alpha, beta)) |> SVector{2}
end

function log_likelihood(::GammaParam{T},
    (alpha, beta)::AbstractVector, value::T) where T
  alpha*log(beta) - loggamma(alpha) + (alpha-1)*log(value) - beta*value
end

function sample_prediction(::GammaParam{T}, (alpha, beta)::AbstractVector{T})::T where T
  Gamma(alpha, 1/beta) |> rand
end

struct ExpParam{ T } <: PredictionParam{T} end

n_params(::ExpParam) = 1

function transform(::ExpParam{T},
    (lambda,)::AbstractVector) where T
  softplus.((lambda,)) |> SVector{1}
end

function log_likelihood(::ExpParam{T},
    (lambda,)::AbstractVector, value::T) where T
  log(lambda) - lambda*value
end

function sample_prediction(::ExpParam{T}, (lambda,)::AbstractVector{T})::T where T
  Exponential(1/lambda) |> rand
end

struct BernoulliParam{ T } <: PredictionParam{T}
  smoothness :: T
end

n_params(::BernoulliParam) = 1

function transform(param::BernoulliParam{T},
    (probability,)::AbstractVector) where T
  sigmoid.((probability/param.smoothness,)) |> SVector{1}
end

# make it more nuanced later on
discretize(::BernoulliParam, value) = value > 0.5

function log_likelihood(param::BernoulliParam{T},
    (probability,)::AbstractVector, value::T) where T
  event = discretize(param, value)
  if event
    log(probability)
  else
    log(1 - probability)
  end
end

function sample_prediction(::BernoulliParam{T},
    (probability,)::AbstractVector{T})::T where T
  Bernoulli(probability) |> rand |> T
end

max_stars_repo_name: marmarelis/TreatmentCurves

max_issues_repo_name: marmarelis/TreatmentCurves

max_forks_repo_name: marmarelis/TreatmentCurves


text:
using Base.Test
using EGR

println("TestSSVRG")

createOracleOutputLevel = 1
numEquivalentPasses = 0.5
algOutputLevel = 0
maxOutputNum=20

sLikeGd(k,numTP) = k==0 ? 0 : numTP
uLikeGd(k,numTP) = k==0 ? numTP : 0

myREfunction(problem, opts, sd, expIndices,n) = false
myWriteFunction(problem, sd, opts, k, gnum, origWant, fromOutputsFunction) = false

for (gradientOracle, numVars, numTP, csDataType, LossFunctionString, myOutputter, L2reg, thisDataName, thisProblem) in Oracles

	println(" $thisDataName $LossFunctionString L2reg = $L2reg")

	maxG  = Int(round(numEquivalentPasses*numTP))

	expIndices=unique(round(Int, logspace(0,log10(maxG+1),maxOutputNum)))-1

	myOutputOpts =  OutputOpts(myOutputter,expIndices)
	myOptss(stepSizePower) = Opts(zeros(numVars,1); stepSizePower=stepSizePower, maxG=maxG, outputLevel=algOutputLevel, stepOutputLevel=0)
	k(k,mtilde)=k+1
	m=10
	alg(thisProblem(Task(() -> getSequentialFinite(numTP, gradientOracle))),  myOptss(0), SSVRG(k,m,numVars), myOutputOpts, myWriteFunction)
end

println("TestSSVRG successful!")
println()


text:
function PhaseBuilder(arg0::CorrelatedRandomVectorGenerator, arg1::GroundStation, arg2::jdouble, arg3::jdouble, arg4::jdouble, arg5::ObservableSatellite)
    return PhaseBuilder((CorrelatedRandomVectorGenerator, GroundStation, jdouble, jdouble, jdouble, ObservableSatellite), arg0, arg1, arg2, arg3, arg4, arg5)
end

function add_modifier(obj::AbstractMeasurementBuilder, arg0::EstimationModifier)
    return jcall(obj, "addModifier", void, (EstimationModifier,), arg0)
end

function build(obj::PhaseBuilder, arg0::Vector{SpacecraftState})
    return jcall(obj, "build", Phase, (Vector{SpacecraftState},), arg0)
end

function get_modifiers(obj::AbstractMeasurementBuilder)
    return jcall(obj, "getModifiers", List, ())
end

function init(obj::AbstractMeasurementBuilder, arg0::AbsoluteDate, arg1::AbsoluteDate)
    return jcall(obj, "init", void, (AbsoluteDate, AbsoluteDate), arg0, arg1)
end


text:
function test_mod_multiplicativeorder()
    @test multiplicativeorder(1, 11) == 1
    @test multiplicativeorder(2, 7) == 3
    @test multiplicativeorder(3, 7) == 6
    @test multiplicativeorder(7, 5) == 4
    @test multiplicativeorder(37, 1000) == 100

    @test_throws ErrorException multiplicativeorder(3, -6)
end

function test_mod_linearmod()
    @test linearmod(3, 5, 6) == []
    @test linearmod(3, 4, 5) == [3]
    @test linearmod(3, 6, 9) == [2, 5, 8]
    @test linearmod(-14, 30, 100) == [5, 55]

    @test_throws ErrorException linearmod(3, 4, 1)
end

function test_mod_primitiveroot()
    @test primitiveroot(2) == 1
    @test primitiveroot(17) == 3
    @test primitiveroot(71) == 7
    @test primitiveroot(101) == 2

    @test_throws ErrorException primitiveroot(4)
end

function test_mod_all()
    print(rpad("Math.NumberTheory.Mod...", 50, ' '))

    test_mod_multiplicativeorder()
    test_mod_linearmod()
    test_mod_primitiveroot()

    println("PASS")
end


text:
"""
    SteMDerivative(x::Vector{Float64}, y::Vector{Float64})

This function evaluates the numerical derivative according to the SteM
algorithm, [Camera et al. 2017](https://arxiv.org/abs/1606.03451):

- a linear regression over ``x`` and ``y`` is performed

- if the points obtained with the fit are close enough (less than 0.01 relative
difference) the linear ansatz is satisfied and the slope gives the derivative

- if the linear ansatz is not satisfied, the external couple of points and the
linear regression is performed again till the linear ansatz is satisfied
"""
function SteMDerivative(x::Vector{Float64}, y::Vector{Float64})
    x_copy = deepcopy(x)
    y_copy = deepcopy(y)
    if minimum(y) == maximum(y)
        der = 0
    else
        nonlinear = true
        while nonlinear
            coefficients = CustomRegression(x_copy, y_copy)
            y_fit = zeros(length(x_copy))
            percent_diff = zeros(length(x_copy))
            y_fit .= coefficients[1] .+ x_copy .* coefficients[2]
            percent_diff = abs.((y_copy .- y_fit) ./ y_copy)
            if all(percent_diff .<= 0.01) || length(x_copy) < 3
                nonlinear = false
                der = coefficients[2]
            else
                pop!(x_copy)
                pop!(y_copy)
                popfirst!(y_copy)
                popfirst!(x_copy)
            end
        end
    end
    return der
end


text:
using LinearAlgebra

#if VERSION < v"0.7-"
#  import Base: A_mul_B!
#  mul!(x,B,y) = A_mul_B!(x,B,y)
#end

@testset "Point-Field Routines" begin

  @testset "Point creation" begin
    @test_throws AssertionError VectorData([1,2,3],[1,2])

    @test_throws AssertionError TensorData([1,2,3],[1,2],[2,3],[4,5])

    f = ScalarData(10)
    f .= rand(10)

    ft = TensorData(f)
    ft[25] = 4
    @test ft.dvdx[5] == 4

  end

  @testset "Point operations" begin

    f = ScalarData(10)
    f .= rand(10)

    g = -f
    @test g[4] == -f[4]

    g = ScalarData(f)
    g .= rand(10)

    h = f + g

    @test h[8] == f[8]+g[8]

    h = 2.0*g

    @test h[4] == 2.0*g[4]

    Y = VectorData(4)
    Y .= rand(length(Y))
    X = 2*Y

    @test X.u[1] == 2*Y.u[1] && X.v[1] == 2*Y.v[1]

    X = Y/2

    @test X.u[3] == Y.u[3]/2 && X.v[3] == Y.v[3]/2

    fill!(Y,0.0)

    X = Y + (1,2)
    @test X.v[4] == 2.0

    X2 = (1,2) + Y
    @test X2 == X

    X = Y - (1,2)
    @test X.u[4] == -1.0

    X = (1,2) - Y
    @test X.u[4] == 1.0

    Z = 2.0 × X
    @test typeof(Z) <: VectorData
    @test Z.u[1] == -4.0

    X = TensorData(Y)
    fill!(X,1)
    Z = X + X
    @test Z.dvdx[4] == X.dvdx[4] + X.dvdx[4]

    Z = (1,2) ⋅ X
    @test typeof(Z) <: VectorData
    @test Z.u[3] == 3.0

  end

  @testset "Hadamard products" begin

    fcs = ScalarData(5,dtype=ComplexF64)
    fill!(fcs,2im)
    fcs .= rand(5)*im

    a = 2.0
    frs = TensorData(fcs,dtype=Float64)
    fill!(frs,a)

    out = similar(frs,element_type=ComplexF64)

    product!(out,frs,fcs)

    @test out[3] == frs[3]*fcs[3]

    out = frs ∘ fcs

    @test out[2] == frs[2]*fcs[2]

  end

  n = 10
  x = 0.5 .+ 0.2*rand(n)
  y = 0.5 .+ 0.2*rand(n)
  X = VectorData(x,y)

  nx = 12; ny = 12

  dx = 0.1
  H = Regularize(X,dx)
  H̃ = Regularize(X,dx,filter=true)

  @testset "Regularize vector to primal edges" begin

    f = VectorData(X)
    f.u .= rand(n)
    f.v .= rand(n)

    q = Edges(Primal,(nx,ny))

    H(q,f)
    @test sum(f.u) ≈ sum(q.u)*dx*dx
    @test sum(f.v) ≈ sum(q.v)*dx*dx

  end

  @testset "Regularize vector to dual edges" begin

  f = VectorData(X)
  f.u .= rand(n)
  f.v .= rand(n)

  p = Edges(Dual,(nx,ny))

  H(p,f)
  @test sum(f.u) ≈ sum(p.u)*dx*dx
  @test sum(f.v) ≈ sum(p.v)*dx*dx

  end

  @testset "Regularize vector to dual and primal nodes" begin

  f = VectorData(X)
  f.u .= rand(n)
  f.v .= rand(n)

  w = NodePair(Dual,(nx,ny))

  H(w,f)
  @test sum(f.u) ≈ sum(w.u)*dx*dx
  @test sum(f.v) ≈ sum(w.v)*dx*dx

  w = NodePair(Primal,(nx,ny))

  H(w,f)
  @test sum(f.u) ≈ sum(w.u)*dx*dx
  @test sum(f.v) ≈ sum(w.v)*dx*dx

  end

  @testset "Regularize tensor to edge gradient" begin

  f = TensorData(X)
  f.dudx .= rand(n)
  f.dudy .= rand(n)
  f.dvdx .= rand(n)
  f.dvdy .= rand(n)

  gradq = EdgeGradient(Dual,(nx,ny))

  H(gradq,f)

  @test sum(f.dudx) ≈ sum(gradq.dudx)*dx*dx
  @test sum(f.dudy) ≈ sum(gradq.dudy)*dx*dx
  @test sum(f.dvdx) ≈ sum(gradq.dvdx)*dx*dx
  @test sum(f.dvdy) ≈ sum(gradq.dvdy)*dx*dx

  gradq = EdgeGradient(Primal,(nx,ny))

  H(gradq,f)

  @test sum(f.dudx) ≈ sum(gradq.dudx)*dx*dx
  @test sum(f.dudy) ≈ sum(gradq.dudy)*dx*dx
  @test sum(f.dvdx) ≈ sum(gradq.dvdx)*dx*dx
  @test sum(f.dvdy) ≈ sum(gradq.dvdy)*dx*dx

  end

  @testset "Regularize scalar to primal nodes" begin

  f = ScalarData(X)
  f .= rand(n)

  w = Nodes(Primal,(nx,ny))

  H(w,f)
  @test sum(f) ≈ sum(w)*dx*dx

  end

  @testset "Regularize scalar to dual nodes" begin

  f = ScalarData(X)
  f .= rand(n)

  w = Nodes(Dual,(nx,ny))

  H(w,f)
  @test sum(f) ≈ sum(w)*dx*dx

  end

  @testset "Matrix representation" begin

  f = ScalarData(X)
  w = Nodes(Dual,(nx,ny))
  Hmat = RegularizationMatrix(H,f,w)
  Emat = InterpolationMatrix(H,w,f)

  f .= rand(n)

  w2 = Nodes(Dual,(nx,ny))
  mul!(w,Hmat,f)
  H(w2,f)
  @test w ≈ w2

  @test_throws MethodError mul!(f,Hmat,w)

  f2 = ScalarData(f)
  mul!(f,Emat,w)
  H(f2,w)
  @test f ≈ f2

  w .= rand(nx,ny)
  Ẽmat = InterpolationMatrix(H̃,w,f)

  mul!(f,Ẽmat,w)
  H̃(f2,w)
  @test f ≈ f2

  w = Nodes(Primal,(nx,ny))
  Hmat = RegularizationMatrix(H,f,w)
  Emat = InterpolationMatrix(H,w,f)
  Ẽmat = InterpolationMatrix(H̃,w,f)

  w2 = Nodes(Primal,(nx,ny))
  mul!(w,Hmat,f)
  H(w2,f)
  @test w ≈ w2

  w .= rand(Float64,size(w))
  f2 = ScalarData(f)
  mul!(f,Emat,w)
  H(f2,w)
  @test f ≈ f2

  f2 = ScalarData(f)
  mul!(f,Ẽmat,w)
  H̃(f2,w)
  @test f ≈ f2

  f = VectorData(X)
  f.u .= rand(n)
  f.v .= rand(n)

  p = Edges(Dual,(nx,ny))
  Hmat = RegularizationMatrix(H,f,p)
  Emat = InterpolationMatrix(H,p,f)
  Ẽmat = InterpolationMatrix(H̃,p,f)

  p2 = Edges(Dual,(nx,ny))
  mul!(p,Hmat,f)
  H(p2,f)
  @test p.u ≈ p2.u && p.v ≈ p2.v

  p.u .= rand(Float64,size(p.u))
  p.v .= rand(Float64,size(p.v))
  f2 = VectorData(f)
  mul!(f,Emat,p)
  H(f2,p)
  @test f.u ≈ f2.u && f.v ≈ f2.v

  mul!(f,Ẽmat,p)
  H̃(f2,p)
  @test f.u ≈ f2.u && f.v ≈ f2.v

  p = NodePair(Dual,(nx,ny))
  p2 = deepcopy(p)
  Hmat = RegularizationMatrix(H,f,p)
  Emat = InterpolationMatrix(H,p,f)
  Ẽmat = InterpolationMatrix(H̃,p,f)

  mul!(p,Hmat,f)
  H(p2,f)
  @test p.u ≈ p2.u && p.v ≈ p2.v

  p.u .= rand(Float64,size(p.u))
  p.v .= rand(Float64,size(p.v))
  mul!(f,Emat,p)
  H(f2,p)
  @test f.u ≈ f2.u && f.v ≈ f2.v

  mul!(f,Ẽmat,p)
  H̃(f2,p)
  @test f.u ≈ f2.u && f.v ≈ f2.v

  p = NodePair(Primal,(nx,ny))
  p2 = deepcopy(p)
  Hmat = RegularizationMatrix(H,f,p)
  Emat = InterpolationMatrix(H,p,f)
  Ẽmat = InterpolationMatrix(H̃,p,f)
  mul!(p,Hmat,f)
  H(p2,f)
  @test p.u ≈ p2.u && p.v ≈ p2.v

  p.u .= rand(Float64,size(p.u))
  p.v .= rand(Float64,size(p.v))
  mul!(f,Emat,p)
  H(f2,p)
  @test f.u ≈ f2.u && f.v ≈ f2.v

  mul!(f,Ẽmat,p)
  H̃(f2,p)
  @test f.u ≈ f2.u && f.v ≈ f2.v

  f = TensorData(X)
  f.dudx .= rand(n)
  f.dudy .= rand(n)
  f.dvdx .= rand(n)
  f.dvdy .= rand(n)
  f2 = TensorData(f)

  p = EdgeGradient(Dual,(nx,ny))
  p2 = deepcopy(p)
  Hmat = RegularizationMatrix(H,f,p)
  Emat = InterpolationMatrix(H,p,f)
  Ẽmat = InterpolationMatrix(H̃,p,f)

  mul!(p,Hmat,f)
  H(p2,f)
  @test p.dudx ≈ p2.dudx && p.dudy ≈ p2.dudy && p.dvdx ≈ p2.dvdx && p.dvdy ≈ p2.dvdy

  p.dudx .= rand(Float64,size(p.dudx))
  p.dudy .= rand(Float64,size(p.dudy))
  p.dvdx .= rand(Float64,size(p.dvdx))
  p.dvdy .= rand(Float64,size(p.dvdy))
  mul!(f,Emat,p)
  H(f2,p)
  @test f.dudx ≈ f2.dudx && f.dudy ≈ f2.dudy && f.dvdx ≈ f2.dvdx && f.dvdy ≈ f2.dvdy

  mul!(f,Ẽmat,p)
  H̃(f2,p)
  @test f.dudx ≈ f2.dudx && f.dudy ≈ f2.dudy && f.dvdx ≈ f2.dvdx && f.dvdy ≈ f2.dvdy

  p = EdgeGradient(Primal,(nx,ny))
  p2 = deepcopy(p)
  Hmat = RegularizationMatrix(H,f,p)
  Emat = InterpolationMatrix(H,p,f)
  Ẽmat = InterpolationMatrix(H̃,p,f)

  mul!(p,Hmat,f)
  H(p2,f)
  @test p.dudx ≈ p2.dudx && p.dudy ≈ p2.dudy && p.dvdx ≈ p2.dvdx && p.dvdy ≈ p2.dvdy

  p.dudx .= rand(Float64,size(p.dudx))
  p.dudy .= rand(Float64,size(p.dudy))
  p.dvdx .= rand(Float64,size(p.dvdx))
  p.dvdy .= rand(Float64,size(p.dvdy))
  mul!(f,Emat,p)
  H(f2,p)
  @test f.dudx ≈ f2.dudx && f.dudy ≈ f2.dudy && f.dvdx ≈ f2.dvdx && f.dvdy ≈ f2.dvdy

  mul!(f,Ẽmat,p)
  H̃(f2,p)
  @test f.dudx ≈ f2.dudx && f.dudy ≈ f2.dudy && f.dvdx ≈ f2.dvdx && f.dvdy ≈ f2.dvdy

  end

### COMPLEX ROUTINES

@testset "Complex point data" begin

  f = ScalarData(10)
  f .= rand(10)

  fc = ScalarData(f,dtype=ComplexF64)
  ftc = TensorData(f,dtype=ComplexF64)

  a = 1.0 + 2.0im
  fc[5] = a

  @test fc[5] == a

  fvc = VectorData(f,dtype=ComplexF64)
  b = (-2.0+0im,5.0im)
  gvc = fvc + b
  @test gvc.u[5] == b[1]
  @test gvc.v[5] == b[2]

  gvc *= 2im
  @test gvc.u[5] == 2im*b[1]
  @test gvc.v[5] == 2im*b[2]

end

end


text:
@testset "Envelope" begin
    @testset "Constructors" begin
        n = 15
        x1 = collect(range(0,stop = 10,length = n))
        L1 = MLine(x1,x1)

        en = Envelope(L1)
        @test isa(en,Envelope)
        @test length(en.env)==n
        @test length(gets(en))==0
        DCEGM.removed!(en,L1)
        @test length(getr(en))==0
        @test en.vbound == 0.0
    end
end

@testset "Testing upper_env" begin
    @testset "upper_env test 1" begin

        x1,x2,a = DCEGM.tplot1()
        e = DCEGM.upper_env(a)
        @test issorted(getx(e.env))

        # test actual values in the Envelope
        @test all( gety(filter(x -> x.x < 5.0 , e.env.v)) .== 5.0 )
        @test all( gety(filter(x -> x.x > 5.0 , e.env.v)) .== getx(filter(x -> x.x > 5.0 , e.env.v)) )

    end
    # @testset "upper_env test 2" begin
    #     n = 15
    #     x1 = collect(range(0,stop = 10,length = n))
    #     L1 = MLine(x1,x1)
    #     L2 = MLine(x1,ones(n)*5)
    #     e = Envelope([L1,L2])
    #     upper_env!(e)
    #     @test issorted(getx(e.env))
    #     @test all( gety(filter(x -> x.x < 5.0 , e.env.v)) .== 5.0 )
    #     @test all( gety(filter(x -> x.x > 5.0 , e.env.v)) .== getx(filter(x -> x.x > 5.0 , e.env.v)) )
    #
    #     upper_env!(e)
    #     @test issorted(getx(e.env))
    #     @test gets(e)[1].x ≈ 5.0
    #     @test gets(e)[1].y ≈ 5.0
    # end
    # @testset "upper_env test 3" begin
    #     n = 15
    #     x1 = collect(range(0,stop = 10,length = n))
    #     L1 = MLine(x1,x1)
    #     L2 = MLine(x1,5.0 .+ 0.3*x1)
    #     e = Envelope([L1,L2])
    #     upper_env!(e)
    #     @test issorted(getx(e.env))
    #     @test all( gety(filter(x -> x.x < 7.0 , e.env.v)) .== gety(filter(x -> x.x < 7.0 , L2.v)) )
    #
    #     upper_env!(e)
    #     @test issorted(getx(e.env))
    #     @test gets(e)[1].x ≈ 5.0/0.7
    #     @test gets(e)[1].y ≈ 5.0/0.7
    # end
    # @testset "upper_env test 4" begin
    #     n = 15
    #     x1 = collect(range(0,stop = 10,length = n))
    #     x2 = collect(range(-1,stop = 9,length = n))
    #     L1 = MLine(x1,x1)
    #     L2 = MLine(x2,5.0 .+ 0.3*x2)
    #     e = Envelope([L1,L2])
    #     upper_env!(e)
    #     @test issorted(getx(e.env))
    #     # @test isapprox(getx(e.env),sort(unique(vcat(x1,x2))), atol = 1e-15)
    #     @test gets(e)[1].x ≈ 5.0/0.7
    #     @test gets(e)[1].y ≈ 5.0/0.7
    # end
#     # poor test since we only do linear interpolation...
#     @testset "log(x) vs 2(x-1)" begin
#         x = collect(0.0:0.1:1.5)
#         L1 = MLine(x,log.(x))
#         L2 = MLine(x, 2 .* (x .- 1.0) )
#         e = Envelope([L1,L2])
#         upper_env!(e)
#         @test issorted(getx(e.env))
#         # @test gets(e)[1].x ≈ DCEGM.fzero(x-> 2*(x-1) - log(x),0.01,0.5)
#         # @test gets(e)[1].y ≈ log( gets(e)[1].x )
#         @test gets(e)[2].x == 1.0
#         @test gets(e)[2].y == 0.0
#     end
    @testset "y=x, y=2x, y=0.1x + 1.98" begin
        x1 = collect(-0.9:0.3:2.7)
        L1 = MLine(x1, x1)
        x2 = collect(0.0:0.1:1)
        L2 = MLine(x2, 2 .* x2)
        x3 = collect(1.0:0.45:2.9)
        L3 = MLine(x3, (0.1 .* x3) .+ 1.9)
        a = [L1,L2,L3]
        e = DCEGM.upper_env(a)
        @test issorted(getx(e.env))
        @test all( gety(filter(x -> x.x < 0.0 , e.env.v)) .== gety(filter(x -> x.x < 0.0 , L1.v)) )
        @test all( gety(filter(x -> (x.x > 0.0) & (x.x < 1.0) , e.env.v)) .== gety(filter(x -> (x.x > 0.0) & (x.x < 1.0)  , L2.v)) )

        @test gets(e)[1] ≈ Point(1.9/0.9,1.9/0.9)

    end
    @testset "upper_env test: decreasing " begin
        L,e,x1,x2,pl = DCEGM.test_upper_env_dec()
        # n = 15
        # x1 = collect(range(0,stop = 10,length = n))
        # x2 = collect(range(-1,stop = 9,length = n))
        # x = vcat(x1,x2)
        # y = vcat(x1[end:-1:1],ones(n)*5)
        # L = MLine(x,y)
        # e = splitLine(L)
        # upper_env!(e)
        @test issorted(getx(e.env))
    end
end
@testset "splitLine" begin
    @testset "test 1: simple" begin
        x = [1,2,3,1.5,2.1,2.9]
        y = [1,1.5,1.7,1.2,1.8,2.1]
        L1 = MLine(x,y)
        e = splitLine(L1)

        @test isa(e,Vector{MLine{Float64}})
        @test length(e) == 3
    end
    @testset "test 2: less simple" begin
        x = [1,2,3,2.9,2.5,1.9,1.8,1.5,2.1,2.9]
        y = [1,1.5,1.7,1.6,1.55,1.4,1.3,1.2,1.8,2.1]
        L1 = MLine(x,y)
        a = splitLine(L1)
        @test isa(a,Array{MLine{Float64}})
        e = DCEGM.upper_env(a)
        @test isa(e,Envelope)
        @test length(gets(e))==1
        @test length(a) == 3
        @test eltype(a)== MLine{Float64}

        @test extrema(getx(a[1].v)) == (1.0,3.0)
        @test extrema(getx(a[2].v)) == (1.5,3.0)
        @test extrema(getx(a[3].v)) == (1.5,2.9)
    end
end


text:
macro wandbinit(exs...)
    kwargs = Any[]
    for ex in exs
        if ex isa Expr && ex.head === :(=)
            k, v = ex.args
            if !(k isa Symbol)
                k = Symbol(k)
            end
            push!(kwargs, Expr(:kw, k, esc(v)))
        elseif ex isa Expr && ex.head === :... # Keyword splatting
            push!(kwargs, esc(ex))
        else
            k, v = ex, ex
            if !(k isa Symbol)
                k = Symbol(k)
            end
            push!(kwargs, Expr(:kw, k, esc(v)))
        end
    end
    return quote
        if !("JULIA_NO_WANDB" in keys(ENV)) || ENV["JULIA_NO_WANDB"] == "" || ENV["JULIA_NO_WANDB"] == "false" || ENV["JULIA_NO_WANDB"] == false
            kw = pairs((;$(kwargs...)))
            if Sys.iswindows()
                if :settings in keys(kw) && !isnothing(kw[:settings].start_method) && kw[:settings].start_method != "thread"
                    @warn "wandb initialization is known to crash on Windows OS if the setting `start_method` is not \"thread\". You have specified start_method=\"$(kw[:settings].start_method)\". But don't worry! This @wandbinit macro is going to override that setting and set it to \"thread\" for you. But next time, either specify settings=wandb.Settings(start_method=\"thread\") or don't set the start_method. In the latter case, this macro will set it to \"thread\" automatically for you and inform you gently."
                    kw[:settings].start_method = "thread"
                elseif :settings in keys(kw) && isnothing(kw[:settings].start_method)
                    @info "Adding kwarg start_method=\"thread\" to wandb.Settings to avoid a crash on Windows OS (a known issue)."
                    kw[:settings].start_method = "thread"
                elseif !(:settings in keys(kw))
                    @info "Adding kwarg settings=wandb.Settings(start_method=\"thread\") to avoid a crash on Windows OS (a known issue)."
                    kw = (kw..., settings=wandb.Settings(; start_method="thread"))
                end
            end
            wandb.init(; kw...)
        end
    end
end


text:
using Documenter
using EffectSizes

makedocs(;
    modules=[EffectSizes],
    format=Documenter.HTML(),
    pages=[
        "Home" => "index.md",
    ],
    repo="https://github.com/harryscholes/EffectSizes.jl/blob/{commit}{path}#L{line}",
    sitename="EffectSizes.jl",
)

deploydocs(;
    repo="github.com/harryscholes/EffectSizes.jl",
)


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

==(w::WeakRef, v::WeakRef) = isequal(w.value, v.value)
==(w::WeakRef, v) = isequal(w.value, v)
==(w, v::WeakRef) = isequal(w, v.value)

"""
    finalizer(x, f)

Register a function `f(x)` to be called when there are no program-accessible references to
`x`. The type of `x` must be a `mutable struct`, otherwise the behavior of this function is
unpredictable.
"""
function finalizer(@nospecialize(o), @nospecialize(f))
    if isimmutable(o)
        error("objects of type ", typeof(o), " cannot be finalized")
    end
    ccall(:jl_gc_add_finalizer_th, Void, (Ptr{Void}, Any, Any),
          Core.getptls(), o, f)
end

function finalizer(o::T, f::Ptr{Void}) where T
    @_inline_meta
    if isimmutable(T)
        error("objects of type ", T, " cannot be finalized")
    end
    ccall(:jl_gc_add_ptr_finalizer, Void, (Ptr{Void}, Any, Ptr{Void}),
          Core.getptls(), o, f)
end

"""
    finalize(x)

Immediately run finalizers registered for object `x`.
"""
finalize(@nospecialize(o)) = ccall(:jl_finalize_th, Void, (Ptr{Void}, Any,),
                                   Core.getptls(), o)

"""
    gc()

Perform garbage collection. This should not generally be used.
"""
gc(full::Bool=true) = ccall(:jl_gc_collect, Void, (Int32,), full)

"""
    gc_enable(on::Bool)

Control whether garbage collection is enabled using a boolean argument (`true` for enabled,
`false` for disabled). Returns previous GC state. Disabling garbage collection should be
used only with extreme caution, as it can cause memory use to grow without bound.
"""
gc_enable(on::Bool) = ccall(:jl_gc_enable, Int32, (Int32,), on) != 0

max_stars_repo_name: visr/julia

max_issues_repo_name: visr/julia

max_forks_repo_name: visr/julia


text:
using Merlin

include("bioes.jl")
include("ner.jl")
include("model.jl")

function train()
end

max_stars_repo_name: hshindo/Merlin-Examples

max_issues_repo_name: hshindo/Merlin-Examples

max_forks_repo_name: hshindo/Merlin-Examples


text:
### type definitions ###

"""
A space object considered to have negligible mass compared to the body around which it orbits
(e.g. a spacecraft or an asteroid).
"""
struct SpaceObject{P}
    name::AbstractString
    orbit::Orbit{P}
    primary::P
end

Base.hash(ob::SpaceObject, h::UInt) = hash_struct(ob, h, (:name, :orbit))

"""
A celestial body (e.g. a planet or moon) that is in orbit around a larger body.
"""
struct CelestialBody{P} <: CelestialObject
    name::AbstractString
    eqradius::Float64
    μ::Float64
    SoI::Float64
    rotperiod::Float64
    rotinitial::Float64
    orbit::Orbit{P}
    primary::CelestialObject
    color::Tuple{Int, Int, Int}
    satellite_bodies::Vector{<:CelestialBody}
    satellite_objects::Vector{<:SpaceObject}
end

Base.hash(bd::CelestialBody, h::UInt) = hash_struct(bd, h, (:name, :eqradius, :μ, :SoI, :rotperiod, :rotinitial, :orbit, :color))

"""
A star that is not in orbit around a larger body. Serves as the root of a tree defining a solar system.
"""
struct Star <: CelestialObject
    name::AbstractString
    eqradius::Float64
    μ::Float64
    color::Tuple{Int, Int, Int}
    satellite_bodies::Vector{CelestialBody{Star}}
    satellite_objects::Vector{SpaceObject{Star}}
end

Base.hash(bd::Star, h::UInt) = hash_struct(bd, h, (:name, :eqradius, :μ, :color))

### alternate constructors ###

# function SpaceObject(name::AbstractString, orbit::Orbit, primary::CelestialObject)
#     ob = SpaceObject(name::AbstractString, orbit::Orbit, primary::CelestialObject)
#     add_to_primary!(ob)
#     return bd
# end

"""
    soi = (μ, μprim, a)

Computes the sphere of influence of a body with semimajor axis `a` and gravity parameter `μ` in orbit around a larger body
with gravity parameter `μprim`.
"""
set_soi(μ, μprim, a) = a *(μ/μprim)^(2/5)

"""
Adds a `CelestialBody` or `SpaceObject` to the lists of satellites around its primary body.
"""
function add_to_primary!(bd::CelestialBody)
    # check if a (mostly) equivalent body hasn't already been added
    if hash(bd) ∉ [hash(sat) for sat in bd.primary.satellite_bodies]
        push!(bd.primary.satellite_bodies, bd)
        sort!(bd.primary.satellite_bodies)
    end
end

function add_to_primary!(ob::SpaceObject)
    # check if a (mostly) equivalent object hasn't already been added
    if hash(ob) ∉ [hash(sat) for sat in ob.primary.satellite_objects]
        push!(ob.primary.satellite_objects, ob)
        sort!(ob.primary.satellite_objects)
    end
end

"""
Creates a `CelestialBody` and adds it to the lists of satellites around its primary body.
"""
function CelestialBody(name::AbstractString, eqradius, μ, SoI, rotperiod, rotinitial, orbit::Orbit, color::Tuple{Int,Int,Int}=(255,255,255))
    bd = CelestialBody(name, eqradius, μ, SoI, rotperiod, rotinitial, orbit, orbit.primary, color, CelestialBody[], SpaceObject[])
    add_to_primary!(bd)
    return bd
end

function CelestialBody(name::AbstractString, eqradius, μ, rotperiod, rotinitial, orbit::Orbit, color::Tuple{Int,Int,Int}=(255,255,255))
    bd = CelestialBody(name, eqradius, μ, set_soi(μ, orbit.primary.μ, orbit.a), rotperiod, rotinitial, orbit, orbit.primary, color, CelestialBody{<:CelestialBody}[], SpaceObject{<:CelestialBody}[])
    add_to_primary!(bd)
    return bd
end

Star(name::AbstractString, eqradius, μ, color::Tuple{Int,Int,Int}=(255,255,255)) = Star(name, eqradius, μ, color, CelestialBody{Star}[], SpaceObject{Star}[])

### helper methods ###

"""
    sats = collect_satellite_bodies(bd)

Collects all planets or moons that are contained within the sphere of influence of `bd`.
"""
function collect_satellite_bodies(bd::CelestialObject)
    if isempty(bd.satellite_bodies)
        return CelestialBody[]
    else
        gchildren = CelestialBody[]
        for st in bd.satellite_bodies
            push!(gchildren, collect_satellite_bodies(st)...)
        end
        return CelestialBody[bd.satellite_bodies..., gchildren...]
    end
end

"""
    sats = collect_satellite_bodies(bd)

Collects all space objects that are contained within the sphere of influence of `bd`.
"""
function collect_satellite_objects(bd::CelestialObject)
    children = collect_satellite_bodies(bd)
    sat_objects = copy(bd.satellite_objects)
    for child in children
        push!(sat_objects, child.satellite_objects...)
    end
    return sat_objects
end

"""
    parent = closest_common_parent(bd1, bd2)

Returns the `CelestialObject` corresponding to the smallest system to which both `bd1` and `bd2` belong.
"""
function closest_common_parent(bd1::Union{CelestialObject,SpaceObject}, bd2::Union{CelestialObject,SpaceObject})
    if bd1 == bd2 || bd2 ∈ collect_satellite_bodies(bd1)
        return bd1
    end
    bd = bd1
    while hasfield(typeof(bd), :orbit)
        bd = bd.orbit.primary
        if bd2 ∈ collect_satellite_bodies(bd) || bd==bd2
            return bd
        end
    end
    return nothing
end

closest_common_parent(orb1::Orbit, orb2::Orbit) = closest_common_parent(orb1.primary, orb2.primary)

"""
    path = path_to_parent(bd, parent)

Returns a vector of `CelestialObject`s traversed in the solar system tree to arrive at `parent`
starting from `bd`. `parent` is assumed to be at a higher level than `bd`
"""
function path_to_parent(bd::Union{CelestialObject,SpaceObject}, parent::CelestialObject)
    if bd == parent return [bd] end
    bodies = Union{CelestialObject,SpaceObject}[bd]
    pbd = bd
    while hasfield(typeof(pbd), :orbit)
        pbd = pbd.orbit.primary
        push!(bodies, pbd)
        if pbd == parent
            return bodies
        end
    end
    return
end

"""
    path = path_to_parent(startbd, endbd)

Returns a vector of `CelestialObject`s traversed in the solar system tree to arrive at `endbd`
starting from `startbd`.
"""
function path_to_body(startbd::Union{CelestialObject,SpaceObject}, endbd::Union{CelestialObject,SpaceObject})
    commonparent = closest_common_parent(startbd, endbd)
    pathup = path_to_parent(startbd, commonparent)
    pathdown = reverse(path_to_parent(endbd, commonparent))
    return [pathup..., pathdown[2:end]...]
end

### comparison and sorting methods ###

function Base.isless(bd1::CelestialObject, bd2::CelestialObject)
    if bd1.orbit.a == bd2.orbit.a
        return isless(bd1.name, bd2.name)
    else
        return isless(bd1.orbit.a, bd2.orbit.a)
    end
end

### descriptive display ###

Base.show(io::IO, st::Star) = println(io,
    """Star "$(st.name)":\n""",
    "  $(length(st.satellite_bodies)) orbiting celestial bodies\n",
    "  $(length(st.satellite_objects)) orbiting space objects"
)

Base.show(io::IO, bd::CelestialBody) = println(io,
    """CelestialBody "$(bd.name)" orbiting $(bd.primary.name):\n""",
    "  $(length(bd.satellite_bodies)) orbiting celestial bodies\n",
    "  $(length(bd.satellite_objects)) orbiting space objects"
)

Base.show(io::IO, ob::SpaceObject) = println(io,
    """SpaceObject "$(ob.name)" orbiting $(ob.primary.name):\n""",
    "  semi-major axis:\t\t$(ob.orb.a)\n",
    "  eccentricity:\t\t\t$(ob.orb.e)\n",
    "  inclination:\t\t\t$(rad2deg(ob.orb.i))°\n",
    "  RA of ascending node:\t\t$(rad2deg(ob.orb.Ω))°\n",
    "  Argument of periapsis:\t$(rad2deg(ob.orb.ω))°\n",
    "  mean anomaly at epoch:\t$(rad2deg(orb.Mo))°\n",
    "  epoch:\t\t\t$(ob.orb.epoch) s"
)