text:
abstract type AbstractSubOptimizer end

struct QpData{T,Tv<:AbstractArray{T},Tm<:AbstractMatrix{T}}
    sense::MOI.OptimizationSense
    Q::Union{Nothing,Tm}
    c::Tv
    c0::T # objective functiton constant term
    A::Tm
    b::Tv
    c_lb::Tv
    c_ub::Tv
    v_lb::Tv
    v_ub::Tv
end

mutable struct QpModel{T,Tv<:AbstractArray{T},Tm<:AbstractMatrix{T}} <: AbstractSubOptimizer
    model::MOI.AbstractOptimizer
    data::QpData{T,Tv,Tm}
    j_row::Vector{Int}
    j_col::Vector{Int}
    adj::Vector{Int}
    x::Vector{MOI.VariableIndex}
    constr_v_ub::Vector{MOI.ConstraintIndex}
    constr_v_lb::Vector{MOI.ConstraintIndex}
    constr::Vector{MOI.ConstraintIndex}
    slack_vars::Dict{Int,Vector{MOI.VariableIndex}}
    constr_slack::Vector{MOI.ConstraintIndex}

    function QpModel(
        model::MOI.AbstractOptimizer,
        data::QpData{T,Tv,Tm},
        j_row::Vector{Int},
        j_col::Vector{Int},
    ) where {T,Tv,Tm}
        qp = new{T,Tv,Tm}()
        qp.model = model
        qp.data = data
        qp.j_row = j_row
        qp.j_col = j_col
        qp.adj = []
        qp.x = []
        qp.constr_v_ub = []
        qp.constr_v_lb = []
        qp.constr = []
        qp.constr_slack = []
        qp.slack_vars = Dict()
        return qp
    end
end

function create_model!(qp::QpModel{T,Tv,Tm}, x_k::Tv, Δ::T, tol_error = 0.0) where {T,Tv,Tm}

    # empty optimizer just in case
    MOI.empty!(qp.model)
    qp.adj = []
    qp.constr_v_ub = []
    qp.constr_v_lb = []
    qp.constr = []
    qp.constr_slack = []
    empty!(qp.slack_vars)

    n = length(qp.data.c)
    m = length(qp.data.c_lb)

    @assert n > 0
    @assert m >= 0
    @assert length(qp.data.c) == n
    @assert length(qp.data.c_lb) == m
    @assert length(qp.data.c_ub) == m
    @assert length(qp.data.v_lb) == n
    @assert length(qp.data.v_ub) == n
    @assert length(x_k) == n

    # variables
    qp.x = MOI.add_variables(qp.model, n)

    # objective function
    obj_terms = Array{MOI.ScalarAffineTerm{T},1}()
    for i = 1:n
        push!(obj_terms, MOI.ScalarAffineTerm{T}(qp.data.c[i], MOI.VariableIndex(i)))
    end

    for i = 1:m
        # add slack variables
        qp.slack_vars[i] = []
        push!(qp.slack_vars[i], MOI.add_variable(qp.model))
        if qp.data.c_lb[i] > -Inf && qp.data.c_ub[i] < Inf
            push!(qp.slack_vars[i], MOI.add_variable(qp.model))
        end

        # Set slack bounds and objective coefficient
        push!(
            qp.constr_slack,
            MOI.add_constraint(
                qp.model,
                MOI.SingleVariable(qp.slack_vars[i][1]),
                MOI.GreaterThan(0.0),
            ),
        )
        push!(obj_terms, MOI.ScalarAffineTerm{T}(1.0, qp.slack_vars[i][1]))
        if length(qp.slack_vars[i]) == 2
            push!(
                qp.constr_slack,
                MOI.add_constraint(
                    qp.model,
                    MOI.SingleVariable(qp.slack_vars[i][2]),
                    MOI.GreaterThan(0.0),
                ),
            )
            push!(obj_terms, MOI.ScalarAffineTerm{T}(1.0, qp.slack_vars[i][2]))
        end
    end

    # set constant term to the objective function
    MOI.set(
        qp.model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}}(),
        MOI.ScalarAffineFunction(obj_terms, qp.data.c0),
    )
    MOI.set(qp.model, MOI.ObjectiveSense(), qp.data.sense)

    for i = 1:n
        ub = min(Δ, qp.data.v_ub[i] - x_k[i])
        lb = max(-Δ, qp.data.v_lb[i] - x_k[i])
        ub = (abs(ub) <= tol_error) ? 0.0 : ub
        lb = (abs(lb) <= tol_error) ? 0.0 : lb
        push!(
            qp.constr_v_ub,
            MOI.add_constraint(qp.model, MOI.SingleVariable(qp.x[i]), MOI.LessThan(ub)),
        )
        push!(
            qp.constr_v_lb,
            MOI.add_constraint(qp.model, MOI.SingleVariable(qp.x[i]), MOI.GreaterThan(lb)),
        )
    end

    for i = 1:m
        c_ub = qp.data.c_ub[i] - qp.data.b[i]
        c_lb = qp.data.c_lb[i] - qp.data.b[i]
        c_ub = (abs(c_ub) <= tol_error) ? 0.0 : c_ub
        c_lb = (abs(c_lb) <= tol_error) ? 0.0 : c_lb

        if qp.data.c_lb[i] == qp.data.c_ub[i] #This means the constraint is equality
            push!(
                qp.constr,
                MOI.add_constraint(
                    qp.model,
                    MOI.ScalarAffineFunction(
                        MOI.ScalarAffineTerm.(
                            [1.0; -1.0],
                            [qp.slack_vars[i][1]; qp.slack_vars[i][2]],
                        ),
                        0.0,
                    ),
                    MOI.EqualTo(c_lb),
                ),
            )
        elseif qp.data.c_lb[i] != -Inf &&
               qp.data.c_ub[i] != Inf &&
               qp.data.c_lb[i] < qp.data.c_ub[i]
            push!(
                qp.constr,
                MOI.add_constraint(
                    qp.model,
                    MOI.ScalarAffineFunction(
                        MOI.ScalarAffineTerm.([1.0], [qp.slack_vars[i][1]]),
                        0.0,
                    ),
                    MOI.GreaterThan(c_lb),
                ),
            )
            push!(qp.adj, i)
        elseif qp.data.c_lb[i] != -Inf
            push!(
                qp.constr,
                MOI.add_constraint(
                    qp.model,
                    MOI.ScalarAffineFunction(
                        MOI.ScalarAffineTerm.([1.0], [qp.slack_vars[i][1]]),
                        0.0,
                    ),
                    MOI.GreaterThan(c_lb),
                ),
            )
        elseif qp.data.c_ub[i] != Inf
            push!(
                qp.constr,
                MOI.add_constraint(
                    qp.model,
                    MOI.ScalarAffineFunction(
                        MOI.ScalarAffineTerm.([-1.0], [qp.slack_vars[i][1]]),
                        0.0,
                    ),
                    MOI.LessThan(c_ub),
                ),
            )
        end
    end

    for i in qp.adj
        c_ub = qp.data.c_ub[i] - qp.data.b[i]
        c_ub = (abs(c_ub) <= tol_error) ? 0.0 : c_ub
        push!(
            qp.constr,
            MOI.add_constraint(
                qp.model,
                MOI.ScalarAffineFunction(
                    MOI.ScalarAffineTerm.([-1.0], [qp.slack_vars[i][2]]),
                    0.0,
                ),
                MOI.LessThan(c_ub),
            ),
        )
    end
end

"""
	sub_optimize!

Solve subproblem

# Arguments
- `qp`: QP model
- `x_k`: trust region center
- `Δ`: trust region size
- `feasibility`: indicator for feasibility restoration phase
- `tol_error`: threshold to drop small numbers to zeros
"""
function sub_optimize!(
    qp::QpModel{T,Tv,Tm},
    x_k::Tv,
    Δ::T,
    feasibility = false,
    tol_error = 0.0,
) where {T,Tv,Tm}

    # dimension of LP
    m, n = size(qp.data.A)
    @assert n > 0
    @assert m >= 0
    @assert length(qp.data.c) == n
    @assert length(qp.data.c_lb) == m
    @assert length(qp.data.c_ub) == m
    @assert length(qp.data.v_lb) == n
    @assert length(qp.data.v_ub) == n
    @assert length(x_k) == n

    b = deepcopy(qp.data.b)

    if feasibility
        # set objective coefficient
        MOI.modify(
            qp.model,
            MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}}(),
            MOI.ScalarConstantChange(0.0),
        )
        for i = 1:n
            MOI.modify(
                qp.model,
                MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}}(),
                MOI.ScalarCoefficientChange(MOI.VariableIndex(i), 0.0),
            )
        end

        # set slack objective coefficient
        for (_, slacks) in qp.slack_vars, s in slacks
            MOI.modify(
                qp.model,
                MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}}(),
                MOI.ScalarCoefficientChange(s, 1.0),
            )
        end

        # set optimization sense
        MOI.set(qp.model, MOI.ObjectiveSense(), MOI.MIN_SENSE)

        do_transform = false
        for cons in qp.constr_slack
            if typeof(cons) == MOI.ConstraintIndex{MOI.SingleVariable,MOI.EqualTo{T}}
                do_transform = true
                break
            end
        end

        # set slack variable bounds
        constr_index = 1
        for i = 1:m
            # Adjust parameters for feasibility problem
            viol = 0.0
            if qp.data.b[i] > qp.data.c_ub[i]
                viol = qp.data.c_ub[i] - qp.data.b[i]
            elseif qp.data.b[i] < qp.data.c_lb[i]
                viol = qp.data.c_lb[i] - qp.data.b[i]
            end
            b[i] -= abs(viol)

            # Add bound constraints
            if length(qp.slack_vars[i]) == 2
                if viol < 0
                    if do_transform
                        qp.constr_slack[constr_index] = MOI.transform(
                            qp.model,
                            qp.constr_slack[constr_index],
                            MOI.GreaterThan(0.0),
                        )
                    else
                        MOI.set(
                            qp.model,
                            MOI.ConstraintSet(),
                            qp.constr_slack[constr_index],
                            MOI.GreaterThan(0.0),
                        )
                    end
                    constr_index += 1
                    if do_transform
                        qp.constr_slack[constr_index] = MOI.transform(
                            qp.model,
                            qp.constr_slack[constr_index],
                            MOI.GreaterThan(viol),
                        )
                    else
                        MOI.set(
                            qp.model,
                            MOI.ConstraintSet(),
                            qp.constr_slack[constr_index],
                            MOI.GreaterThan(viol),
                        )
                    end
                    constr_index += 1
                else
                    if do_transform
                        qp.constr_slack[constr_index] = MOI.transform(
                            qp.model,
                            qp.constr_slack[constr_index],
                            MOI.GreaterThan(-viol),
                        )
                    else
                        MOI.set(
                            qp.model,
                            MOI.ConstraintSet(),
                            qp.constr_slack[constr_index],
                            MOI.GreaterThan(-viol),
                        )
                    end
                    constr_index += 1
                    if do_transform
                        qp.constr_slack[constr_index] = MOI.transform(
                            qp.model,
                            qp.constr_slack[constr_index],
                            MOI.GreaterThan(0.0),
                        )
                    else
                        MOI.set(
                            qp.model,
                            MOI.ConstraintSet(),
                            qp.constr_slack[constr_index],
                            MOI.GreaterThan(0.0),
                        )
                    end
                    constr_index += 1
                end
            elseif length(qp.slack_vars[i]) == 1
                if do_transform
                    qp.constr_slack[constr_index] = MOI.transform(
                        qp.model,
                        qp.constr_slack[constr_index],
                        MOI.GreaterThan(-abs(viol)),
                    )
                else
                    MOI.set(
                        qp.model,
                        MOI.ConstraintSet(),
                        qp.constr_slack[constr_index],
                        MOI.GreaterThan(-abs(viol)),
                    )
                end
                # @show i, viol, length(qp.slack_vars[i]), qp.constr_slack[constr_index]
                constr_index += 1
            else
                @error "unexpected slack_vars"
            end
        end
    else
        # set objective coefficient
        MOI.modify(
            qp.model,
            MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}}(),
            MOI.ScalarConstantChange(qp.data.c0),
        )
        for i = 1:n
            MOI.modify(
                qp.model,
                MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}}(),
                MOI.ScalarCoefficientChange(MOI.VariableIndex(i), qp.data.c[i]),
            )
        end

        # set slack objective coefficient
        for (_, slacks) in qp.slack_vars, s in slacks
            MOI.modify(
                qp.model,
                MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}}(),
                MOI.ScalarCoefficientChange(s, 0.0),
            )
        end

        # set optimization sense
        MOI.set(qp.model, MOI.ObjectiveSense(), qp.data.sense)

        # set slack variable bounds
        do_transform = false
        for cons in qp.constr_slack
            if typeof(cons) != MOI.ConstraintIndex{MOI.SingleVariable,MOI.EqualTo{T}}
                do_transform = true
                break
            end
        end
        if do_transform
            for i in eachindex(qp.constr_slack)
                qp.constr_slack[i] =
                    MOI.transform(qp.model, qp.constr_slack[i], MOI.EqualTo(0.0))
            end
        end
    end

    # set variable bounds
    for i = 1:n
        ub = min(Δ, qp.data.v_ub[i] - x_k[i])
        lb = max(-Δ, qp.data.v_lb[i] - x_k[i])
        ub = (abs(ub) <= tol_error) ? 0.0 : ub
        lb = (abs(lb) <= tol_error) ? 0.0 : lb
        MOI.set(qp.model, MOI.ConstraintSet(), qp.constr_v_ub[i], MOI.LessThan(ub))
        MOI.set(qp.model, MOI.ConstraintSet(), qp.constr_v_lb[i], MOI.GreaterThan(lb))
    end
    # @show Δ, qp.data.v_lb, qp.data.v_ub, x_k

    # modify the constraint coefficients
    for i = 1:length(qp.j_row)
        coeff =
            (abs(qp.data.A[qp.j_row[i], qp.j_col[i]]) <= tol_error) ? 0.0 :
            qp.data.A[qp.j_row[i], qp.j_col[i]]
        MOI.modify(
            qp.model,
            qp.constr[qp.j_row[i]],
            MOI.ScalarCoefficientChange(MOI.VariableIndex(qp.j_col[i]), coeff),
        )
    end
    for (ind, val) in enumerate(qp.adj), i = 1:length(qp.data.A[val, :].nzind)
        i_col = qp.data.A[val, :].nzind[i]
        value = qp.data.A[val, :].nzval[i]
        coeff = (abs(value) <= tol_error) ? 0.0 : value
        MOI.modify(
            qp.model,
            qp.constr[m+ind],
            MOI.ScalarCoefficientChange(MOI.VariableIndex(i_col), coeff),
        )
    end
    # @show qp.data.A, qp.j_row, qp.j_col

    # modify RHS
    for i = 1:m
        c_ub = qp.data.c_ub[i] - b[i]
        c_lb = qp.data.c_lb[i] - b[i]
        c_ub = (abs(c_ub) <= tol_error) ? 0.0 : c_ub
        c_lb = (abs(c_lb) <= tol_error) ? 0.0 : c_lb

        if qp.data.c_lb[i] == qp.data.c_ub[i]
            MOI.set(qp.model, MOI.ConstraintSet(), qp.constr[i], MOI.EqualTo(c_lb))
        elseif qp.data.c_lb[i] != -Inf &&
               qp.data.c_ub[i] != Inf &&
               qp.data.c_lb[i] < qp.data.c_ub[i]
            MOI.set(qp.model, MOI.ConstraintSet(), qp.constr[i], MOI.GreaterThan(c_lb))
        elseif qp.data.c_lb[i] != -Inf
            MOI.set(qp.model, MOI.ConstraintSet(), qp.constr[i], MOI.GreaterThan(c_lb))
        elseif qp.data.c_ub[i] != Inf
            MOI.set(qp.model, MOI.ConstraintSet(), qp.constr[i], MOI.LessThan(c_ub))
        end
    end
    # @show qp.data.c_lb-b, qp.data.c_ub-b, b
    for (i, val) in enumerate(qp.adj)
        c_ub = qp.data.c_ub[val] - b[val]
        c_ub = (abs(c_ub) <= tol_error) ? 0.0 : c_ub
        MOI.set(qp.model, MOI.ConstraintSet(), qp.constr[i+m], MOI.LessThan(c_ub))
    end

    # dest = MOI.FileFormats.Model(format = MOI.FileFormats.FORMAT_LP)
    # MOI.copy_to(dest, qp.model)
    # MOI.write_to_file(dest, "debug.lp")

    MOI.optimize!(qp.model)
    status = MOI.get(qp.model, MOI.TerminationStatus())

    # TODO: These can be part of data.
    Xsol = Tv(undef, n)
    p_slack = Dict{Int,Vector{Float64}}()
    lambda = Tv(undef, m)
    mult_x_U = Tv(undef, n)
    mult_x_L = Tv(undef, n)

    if status == MOI.OPTIMAL
        # @show MOI.get(qp.model, MOI.ObjectiveValue())
        Xsol .= MOI.get(qp.model, MOI.VariablePrimal(), qp.x)
        for (i, slacks) in qp.slack_vars
            p_slack[i] = MOI.get(qp.model, MOI.VariablePrimal(), slacks)
        end
        # @show Xsol
        # @show p_slack

        # extract the multipliers to constraints
        for i = 1:m
            lambda[i] = MOI.get(qp.model, MOI.ConstraintDual(1), qp.constr[i])
        end
        for (i, val) in enumerate(qp.adj)
            lambda[val] += MOI.get(qp.model, MOI.ConstraintDual(1), qp.constr[i+m])
        end
        # @show MOI.get(qp.model, MOI.ConstraintDual(1), qp.constr)

        # extract the multipliers to column bounds
        mult_x_U .= MOI.get(qp.model, MOI.ConstraintDual(1), qp.constr_v_ub)
        mult_x_L .= MOI.get(qp.model, MOI.ConstraintDual(1), qp.constr_v_lb)
        # careful because of the trust region
        for j = 1:n
            if Xsol[j] < qp.data.v_ub[j] - x_k[j]
                mult_x_U[j] = 0.0
            end
            if Xsol[j] > qp.data.v_lb[j] - x_k[j]
                mult_x_L[j] = 0.0
            end
        end
    elseif status == MOI.DUAL_INFEASIBLE
        @error "Trust region must be employed."
    elseif status == MOI.INFEASIBLE
        fill!(Xsol, 0.0)
        fill!(lambda, 0.0)
        fill!(mult_x_U, 0.0)
        fill!(mult_x_L, 0.0)
    else
        @error "Unexpected status: $(status)"
    end

    return Xsol, lambda, mult_x_U, mult_x_L, p_slack, status
end

"""
	get_scalar_quadratic_terms

Collect the quadratic terms of the objective function
"""
function get_scalar_quadratic_terms(Q::Tm) where {T,Tm<:AbstractSparseMatrix{T,Int}}
    terms = Array{MOI.ScalarQuadraticTerm{T},1}()
    rows = rowvals(Q)
    vals = nonzeros(Q)
    m, n = size(Q)
    for j = 1:n
        for i in nzrange(Q, j)
            if i == j
                push!(terms, MOI.ScalarQuadraticTerm{T}(0.5 * vals[i], rows[i], j))
            elseif i > j
                push!(terms, MOI.ScalarQuadraticTerm{T}(vals[i], rows[i], j))
            end
        end
    end
    return terms
end

max_stars_repo_name: exanauts/ActiveSetMethods

max_issues_repo_name: exanauts/ActiveSetMethods

max_forks_repo_name: exanauts/ActiveSetMethods


text:
"""
    abstract type AbstractPolynomialSpace{D}

A polynomial space over `D`. This is a vector space under polynomial addition
and scalar multiplication.
"""
abstract type AbstractPolynomialSpace{D} end

"""
    struct Pk{D,K} <: AbstractPolynomialSpace{D}

The space of all polynomials over `D` of degree `≤K`.

When `D` is a hypercube in `d` dimensions, the precise definition is `Pk{D,K} =
span{𝐱ᶿ : max(θ)≤ K}`; when `D` is a `d`-dimensional simplex, the space is
`Pk{D,K} = span{𝐱ᶿ : sum(θ)≤ K}`, where `θ ∈ 𝐍ᵈ` is a multi-index.
"""
struct Pk{D,K} <: AbstractPolynomialSpace{D} end
Pk(d::AbstractReferenceShape,k::Int) = Pk{typeof(d),k}()

function Base.show(io::IO, pk::Pk{D,K}) where {D,K}
    print(io, "Pk{$D,$K} : space of all polynomials over $D of degree ≤ $K")
end

function dimension(::SType{Pk{D,K}}) where {D,K}
    if D == ReferenceLine
        return K + 1
    elseif D == ReferenceTriangle
        return (K + 1) * (K + 2) ÷ 2
    elseif D == ReferenceTetrahedron
        return (K + 1) * (K + 2) * (K + 3) ÷ 6
    elseif D == ReferenceSquare
        return (K + 1)^2
    elseif D == ReferenceCube
        return (K + 1)^3
    else
        notimplemented()
    end
end

"""
    monomial_basis(sp::Pk)

Return an `NTuple` containing a basis of monomials `𝐱ᶿ` spanning the polynomial
space [`Pk`](@ref).
"""
function monomial_basis end

@generated function monomial_basis(::Pk{ReferenceLine,K}) where {K}
    # NOTE: it would be more efficient to interpolate the variable i directly
    # here.
    # the K+1 monomials x^0, x^1, ..., x^K
    b = ntuple(K + 1) do i
        x -> x^(i - 1)
    end
    return :($b)
end

@generated function monomial_basis(::Pk{ReferenceSquare,K}) where {K}
    # the K+1 monomials x^(0,0), x^(0,1),x^(1,0), ..., x^(K,K)
    I = CartesianIndices((K + 1, K + 1)) .- CartesianIndex(1, 1)
    N = length(I)
    b = ntuple(N) do i
        θ = Tuple(I[i]) # map linear to cartesian index
        x -> prod(x.^θ)
    end
    return :($b)
end

@generated function monomial_basis(::Pk{ReferenceTriangle,K}) where {K}
    # the (K+1)*(K+2)/2 monomials x^(a,b) with a+b ≤ K
    # construct first the indices for the square, then filter only those for
    # which the sum is less than K.
    I =  CartesianIndices((K + 1, K + 1)) .- CartesianIndex(1, 1)
    I = filter(I) do idx
        sum(Tuple(idx)) ≤ K
    end
    N = (K + 1) * (K + 2) ÷ 2
    b = ntuple(N) do i
        θ = Tuple(I[i])
        x -> prod(x.^θ)
    end
    return :($b)
end

"""
    lagrange_basis(nodes,[sp::AbstractPolynomialSpace])

Return the set of `n` polynomials in `sp` taking the value of `1` on node `i`
and `0` on nodes `j ≂̸ i` for `1 ≤ i ≤ n`. For `N`-dimensional
tensor-product nodes represented in the form of an `SVector{N,Vector{T}}`, the
argument `sp` may be ommited.

!!! danger
    It is assumed that the value of a function on `nodes` uniquely determine a
    polynomial in `sp`.
"""
function lagrange_basis(nodes, sp::AbstractPolynomialSpace)
    N = dimension(sp)
    @assert length(nodes) == N
    basis = monomial_basis(sp)
    # compute the matrix of coeffcients of the lagrange polynomials over the
    # monomomial basis
    V     = [p(x) for x in nodes, p in basis]
    C     = V\I
    lag_basis = ntuple(N) do j
        x-> sum(1:N) do i
            C[i,j]*basis[i](x)
        end
    end
    return lag_basis
end

function lagrange_basis(nodes::Vector{<:Number})
    l = x -> prod(nodes) do xi
        x-xi
    end
    w = barycentric_lagrange_weights(nodes)
    map(nodes,w) do xj,wj
        x -> x == xj ? 1.0 : l(x) * wj/(x-xj)
    end
end
# other possible one-dimenional syntaxes
lagrange_basis(nodes::SVector{1,<:Vector}) = lagrange_basis(nodes[1])
function lagrange_basis(nodes::SVector{<:Any,SVector{1,T}}) where {T}
    x = reinterpret(T,nodes) |> Vector
    lagrange_basis(x)
end

ext: jl

max_stars_repo_name: WaveProp/WavePropBase

max_issues_repo_name: WaveProp/WavePropBase

max_forks_repo_name: WaveProp/WavePropBase


text:
# Types defined in types.h
# https://github.com/oxfordcontrol/osqp/blob/master/include/types.h

# Integer type from C
if Sys.WORD_SIZE == 64   # 64bit system
    const Cc_int = Clonglong
else  # 32bit system
    const Cc_int = Cint
end

struct Ccsc
    nzmax::Cc_int
    m::Cc_int
    n::Cc_int
    p::Ptr{Cc_int}
    i::Ptr{Cc_int}
    x::Ptr{Cdouble}
    nz::Cc_int
end

struct ManagedCcsc
    nzmax::Cc_int
    m::Cc_int
    n::Cc_int
    p::Vector{Cc_int}
    i::Vector{Cc_int}
    x::Vector{Cdouble}
    nz::Cc_int

end

# Construct ManagedCcsc matrix from SparseMatrixCSC
function ManagedCcsc(M::SparseMatrixCSC)

    # Get dimensions
    m = M.m
    n = M.n

    # Get vectors of data, rows indices and column pointers
    x = convert(Array{Float64,1}, M.nzval)
    # C is 0 indexed
    i = convert(Array{Cc_int,1}, M.rowval .- 1)
    # C is 0 indexed
    p = convert(Array{Cc_int,1}, M.colptr .- 1)

    # Create new ManagedCcsc matrix
    ManagedCcsc(length(M.nzval), m, n, p, i, x, -1)
end

# Returns an Ccsc matrix. The vectors are *not* GC tracked in the struct.
# Use this only when you know that the managed matrix will outlive the Ccsc
# matrix.
Ccsc(m::ManagedCcsc) =
    Ccsc(m.nzmax, m.m, m.n, pointer(m.p), pointer(m.i), pointer(m.x), m.nz)

struct Solution
    x::Ptr{Cdouble}
    y::Ptr{Cdouble}
end

# Internal C type for info
# N.B. This is not the one returned to the user!
struct CInfo
    iter::Cc_int
    # We need to allocate 32 bytes for a character string, so we allocate 256 bits
    # of integer instead
    # TODO: Find a better way to do this
    status::NTuple{32,Cchar}
    status_val::Cc_int
    status_polish::Cc_int
    obj_val::Cdouble
    pri_res::Cdouble
    dua_res::Cdouble
    setup_time::Cdouble
    solve_time::Cdouble
    polish_time::Cdouble
    run_time::Cdouble
    rho_updates::Cc_int
    rho_estimate::Cdouble
end

struct Data
    n::Cc_int
    m::Cc_int
    P::Ptr{Ccsc}
    A::Ptr{Ccsc}
    q::Ptr{Cdouble}
    l::Ptr{Cdouble}
    u::Ptr{Cdouble}
end

struct Settings
    rho::Cdouble
    sigma::Cdouble
    scaling::Cc_int
    adaptive_rho::Cc_int
    adaptive_rho_interval::Cc_int
    adaptive_rho_tolerance::Cdouble
    adaptive_rho_fraction::Cdouble
    max_iter::Cc_int
    eps_abs::Cdouble
    eps_rel::Cdouble
    eps_prim_inf::Cdouble
    eps_dual_inf::Cdouble
    alpha::Cdouble
    linsys_solver::Cint  # Enum type
    delta::Cdouble
    polish::Cc_int
    polish_refine_iter::Cc_int
    verbose::Cc_int
    scaled_termination::Cc_int
    check_termination::Cc_int
    warm_start::Cc_int
    time_limit::Cdouble
end

function Settings()
    s = Ref{OSQP.Settings}()
    ccall((:osqp_set_default_settings, OSQP.osqp), Nothing,
          (Ref{OSQP.Settings},), s)
    return s[]
end

function Settings(settings_dict::Dict{Symbol,Any})
#  function Settings(settings::Base.Iterators.IndexValue)
#  function Settings(settings::Array{Any, 1})
    default_settings = OSQP.Settings()

       # Convert linsys_solver string to number
    linsys_solver_str_to_int!(settings_dict)

    # Get list with elements of default and user settings
    # If setting is in the passed settings (settings_dict),
    # then convert type to the right type. Otherwise just take
    # the default one
    settings_list = [setting in keys(settings_dict) ?
             convert(fieldtype(typeof(default_settings), setting), settings_dict[setting]) :
             getfield(default_settings, setting)
             for setting in fieldnames(typeof(default_settings))]

    # Create new settings with new dictionary
    s = OSQP.Settings(settings_list...)
    return s

end

struct Workspace
    data::Ptr{OSQP.Data}
    linsys_solver::Ptr{Nothing}
    pol::Ptr{Nothing}

    rho_vec::Ptr{Cdouble}
    rho_inv_vec::Ptr{Cdouble}
    constr_type::Ptr{Cc_int}

    # Iterates
    x::Ptr{Cdouble}
    y::Ptr{Cdouble}
    z::Ptr{Cdouble}
    xz_tilde::Ptr{Cdouble}
    x_prev::Ptr{Cdouble}
    z_prev::Ptr{Cdouble}

    # Primal and dual residuals
    Ax::Ptr{Cdouble}
    Px::Ptr{Cdouble}
    Aty::Ptr{Cdouble}

    # Primal infeasibility
    delta_y::Ptr{Cdouble}
    Atdelta_y::Ptr{Cdouble}

    # Dual infeasibility
    delta_x::Ptr{Cdouble}
    Pdelta_x::Ptr{Cdouble}
    Adelta_x::Ptr{Cdouble}

    # Scaling
    D_temp::Ptr{Cdouble}
    D_temp_A::Ptr{Cdouble}
    E_temp::Ptr{Cdouble}

    settings::Ptr{OSQP.Settings}
    scaling::Ptr{Nothing}
    solution::Ptr{OSQP.Solution}
    info::Ptr{OSQP.CInfo}

    timer::Ptr{Nothing}
    first_run::Cc_int
    summary_printed::Cc_int

end

mutable struct Info
    iter::Int64
    status::Symbol
    status_val::Int64
    status_polish::Int64
    obj_val::Float64
    pri_res::Float64
    dua_res::Float64
    setup_time::Float64
    solve_time::Float64
    polish_time::Float64
    run_time::Float64
    rho_updates::Int64
    rho_estimate::Float64

    Info() = new()
end

function Compat.copyto!(info::Info, cinfo::CInfo)
    info.iter = cinfo.iter
    info.status = OSQP.status_map[cinfo.status_val]
    info.status_val = cinfo.status_val
    info.status_polish = cinfo.status_polish
    info.obj_val = cinfo.obj_val
    info.pri_res = cinfo.pri_res
    info.dua_res = cinfo.dua_res
    info.setup_time = cinfo.setup_time
    info.solve_time = cinfo.solve_time
    info.polish_time = cinfo.polish_time
    info.run_time = cinfo.run_time
    info.rho_updates = cinfo.rho_updates
    info.rho_estimate
    info
end

mutable struct Results
    x::Vector{Float64}
    y::Vector{Float64}
    info::OSQP.Info
    prim_inf_cert::Vector{Float64}
    dual_inf_cert::Vector{Float64}
end

Results() = Results(Float64[], Float64[], Info(), Float64[], Float64[])

function Base.resize!(results::Results, n::Int, m::Int)
    resize!(results.x, n)
    resize!(results.y, m)
    resize!(results.prim_inf_cert, m)
    resize!(results.dual_inf_cert, n)
    results
end


text:
#####################################################################
# Creates the SPN architecture for images, suggested in
# Poon and Domingos - Sum-Product Networks: A new deep architecture
#####################################################################

mutable struct PoonParameters
    width::Int
    height::Int
    "the size of coarse region is a multiple of this"
    baseres::Int
    "number of nodes that represent each region"
    nsum::Int
    "number of leaf nodes that represent each unit region."
    nleaf::Int
end

"""
A decomposition describes how a region is split into two parts.
The two parts are represented by their regionIDs `r1_id` and `r2_id`.
"""
mutable struct Decomposition
    "The id of the first region."
    r1_id::Int
    "The id of the second region."
    r2_id::Int

    "Each decomposition has a set of product nodes."
    nodes::Vector{ProdNode}

    """
    Creates a new decompositon object.
    """
    function Decomposition(r1_id::Int, r2_id::Int)
        new(r1_id,r2_id,ProdNode[])
    end
end

"""
A region is a rectangular image patch.
"""
mutable struct Region
    "Top left position."
    a1::Int
    "Top right position"
    a2::Int
    "Bottom left position"
    b1::Int
    "Bottom right position"
    b2::Int

    "Decompositions for this region."
    decompositions::Vector{Decomposition}
    decomposed::Bool

    """
    Each region is represented by a set of sum nodes.
    Single pixel regions are represented by leaf nodes instead.
    """
    nodes::Vector{Node}

    """
    Creates a new Region object.
    """
    function Region(a1::Int, a2::Int, b1::Int, b2::Int)
        decompositions = Vector{Decomposition}()
        new(a1,a2,b1,b2,decompositions,false,Node[])
    end
end

"""
    regionID(a1::Int, a2::Int, b1::Int, b2::Int, ps::PoonParameters) -> Int

Creates a unique id for every region. The region is specified by its top left,
top right, bottom left, and bottom right position.
"""
function regionID(a1::Int, a2::Int, b1::Int, b2::Int, ps::PoonParameters)
    id = ((a1*ps.width + a2 - 1)*ps.height + b1)*ps.height+b2-1
    return id
end

"""
    regionID(a1::Int, a2::Int, b1::Int, b2::Int, ps::PoonParameters) -> Int

Creates a unique id for an existing region object.
"""
regionID(r::Region, ps::PoonParameters) = regionID(r.a1, r.a2, r.b1, r.b2, ps::PoonParameters)

"""
    generateRegionGraph!(ps::PoonParameters) -> Dict{Int,Region}

Creates a graph of regions and decompositions by decomposing an image into rectangular
subregions.
"""
function generateRegionGraph!(ps::PoonParameters)

    # width and height must be multiples of the base resolution for
    # coarse decomposition:
    @assert ps.width % ps.baseres == 0
    @assert ps.height % ps.baseres == 0

    # shorter notation for better readability:
    w = ps.width
    h = ps.height
    br = ps.baseres

    # here we collect all "visited" regions:
    regions = Dict{Int,Region}()

    # coarse regions:
    for width in br:br:w
        for height in br:br:h
            for a1 in 0:br:(w-width)
                a2 = a1 + width
                for b1 in 0:br:(h-height)
                    b2 = b1 + height
                    r = Region(a1,a2,b1,b2)
                    rid = regionID(r, ps)
                    regions[rid] = r
                end
            end
        end
    end

    # fine regions
    # ca is the left boundary of the coarse region
    # cb is the upper boundary of the coarse region
    for width in 1:br
        for height in 1:br
            if !(width==height==br)
                for ca in 0:br:w-br
                    for cb in 0:br:h-br
                        for a1 in ca:ca+br-width
                            a2 = a1 + width
                            for b1 in cb:cb+br-height
                                b2 = b1 + height
                                r = Region(a1,a2,b1,b2)
                                rid = regionID(r, ps)
                                regions[rid] = r
                            end
                        end
                    end
                end
            end
        end
    end

    # decompositions:
    for r in values(regions)
        if r.a2 - r.a1 <= br && r.b2 - r.b1 <= br
            splitres = 1  # fine decomposition
        else
            splitres = br  # coarse decomposition
        end

        # vertical splits:
        for splitpos in r.a1+splitres:splitres:r.a2-splitres
            # left region:
            r1_id = regionID(r.a1, splitpos, r.b1, r.b2, ps)
            # right region:
            r2_id = regionID(splitpos, r.a2, r.b1, r.b2, ps)
            # add decomposition to r's list of decompositions:
            d = Decomposition(r1_id, r2_id)
            push!(r.decompositions, d)
        end

        # horizontal splits:
        for splitpos in r.b1+splitres:splitres:r.b2-splitres
            # top region:
            r1_id = regionID(r.a1, r.a2, r.b1, splitpos, ps)
            # bottom region:
            r2_id = regionID(r.a1, r.a2, splitpos, r.b2, ps)
            # add decomposition to r's list of decompositions:
            d = Decomposition(r1_id, r2_id)
            push!(r.decompositions, d)
        end
    end

    return regions
end

"""
    generateSPN!(regions::Dict{Int,Region}, nsum::Int)

Generates an SPN from a region graph by adding sum nodes
to regions and product nodes to decompositions.

## Arguments
* `regions::Dict{Int,Region}`: region graph
* `nsum::Int`: number of sum nodes per region
"""
function generateSPN!(x::AbstractMatrix, regions::Dict{Int,Region}, ps::PoonParameters)
    # first phase. Generate sum nodes:
    for r in values(regions)
        if isRootRegion(r, ps)
            # node is root, create only one sum node:
            push!(r.nodes, SumNode())
        elseif isUnitRegion(r)
            # get variable index from region position (assumption: variables are arranged row-wise):
            varidx = r.b1 * ps.width + r.a1 + 1
            # obtain quantile means for current variable:
            means = quantileMeans(x[:,varidx], ps.nleaf)
            for i in 1:ps.nleaf
                g = GaussianNode(varidx, means[i], 1.0)
                push!(r.nodes, g)
            end
        else
            for _ in 1:ps.nsum
                push!(r.nodes, SumNode())
            end
        end
    end

    # Second phase. Generate product nodes and connect them to sum nodes:
    for r in values(regions)
        for d in r.decompositions
            # a product node for every possible pairing of sum nodes, one from each subregion:

            # get nodes from decomposition:
            r1 = regions[d.r1_id]
            r2 = regions[d.r2_id]
            for i in 1:length(r1.nodes)#ps.nsum
                for j in 1:length(r2.nodes)
                    p = ProdNode()
                    # add as a child to all of r's sum nodes:
                    for node in r.nodes
                        connect!(node,p)
                    end

                    connect!(p, r1.nodes[i])
                    connect!(p, r2.nodes[j])

                    push!(d.nodes, p)
                end
            end
        end
    end
end

"""
    isRootRegion(r::Region, ps::PoonParameters) -> Bool

Checks, whether a region is the root region, i.e. the whole image.
"""
function isRootRegion(r::Region, ps::PoonParameters)
    return r.a1==0 && r.a2==ps.width && r.b1==0 && r.b2==ps.height
end

"""
    isUnitRegion(r::Region) -> Bool

Checks, whether a region is a unit region, i.e. a single pixel.
"""
function isUnitRegion(r::Region)
    return r.a2-r.a1==1 && r.b2-r.b1==1
end

"""
    generatePoonStructure() -> SumProductNetwork

Creates a SPN with the architecture suggested in
Poon and Domingos - Sum-Product Networks: A new deep architecture.
"""
function structureLearnPoon(x::AbstractMatrix, height::Int, width::Int; baseres::Int=1, nsum::Int=5, nleaf::Int=4)
    # Test validity of arguments:
    width > 0 || throw(DomainError())
    height > 0 || throw(DomainError())
    nsum > 0 || throw(DomainError())
    nleaf > 0 || throw(DomainError())
    baseres > 0 || throw(DomainError())
    width * height == size(x,2) || throw(DomainError())
    width % baseres == 0 || throw(DomainError())
    height % baseres == 0 || throw(DomainError())

    ps = PoonParameters(width, height, baseres, nsum, nleaf)

    print("Generating decompositions ... ")
    rs = generateRegionGraph!(ps)
    println("done")

    print("Generating SPN ...")
    generateSPN!(x, rs, ps)
    println("done")

    rootid = regionID(0, ps.width, 0, ps.height, ps)
    r = rs[rootid]

    return SumProductNetwork(r.nodes[1])  # turn the root node into an SPN
end


text:
module testdata

# types
include("types.jl")
import .types.FL, .types.SquareSparseMatrix, .types.CustomSquareSparseMatrixRawPairs

include("test-data/external.jl")
using .matrixgen

include("utilities.jl")
using .utilities

include("blocksys.jl")
using .blocksys

"""
### Generates a set of test data
— two files containing the `A` matrix and the `b` vector. Generated files are contained in directory of name `test-data/n‹n›/` with filenames `A.txt` and `b.txt`.

### Arguments:
- `n` size of the matrix
- `l` size of each chunk (submatrix)

### Returns a tuple of paths to the testfiles:
- `[1]` — path to the `A` matrix file
- `[2]` — path to the `b` vector file

"""
function generateTestData(n::Int, l::Int)::Tuple{String, String}

    # first check if there are no files
    outputdirectory = string("test-data/n", string(n), "/")
    outputfile_A = string(outputdirectory, "A.txt")
    outputfile_b = string(outputdirectory, "b.txt")

    if isfile(outputfile_A) && isfile(outputfile_b)
        # there are already some testfiles generated
        return outputfile_A, outputfile_b
    end

    # generate a new directory for these testfiles
    mkpath(outputdirectory)

    # first generate the main matrix using an external script by D.Phil. hab. Paweł Zieliński
    matrixgen.blockmat(n, l, 16.0, outputfile_A)

    # read the `A` matrix
    A = utilities.Amatrix(outputfile_A)

    # generate the `b` vector
    # b = blocksys.generateB(A, n, l)

    # write the `b` vector to another file
    open(outputfile_b, "w") do f

        # print the header
        println(f, n)

        for k in 1:Int(n/l)
            for p in (k-1)*l+1 : k*l
                b = 0
                for j in k : k*l + l
                    if j > n
                        break
                    end
                    b += A[p, j]
                end

                println(f, b)

            end
        end

        # write the contents of the `b` vector in `n` following lines
        # for i in 1:n
        #     println(f, b[i])
        # end

    end

    return outputfile_A, outputfile_b

end

end

max_stars_repo_name: jerry-sky/academic-notebook

max_issues_repo_name: jerry-sky/academic-notebook

max_forks_repo_name: jerry-sky/academic-notebook


text:
function calc_a03d(surf::ThreeDSurfSimple)

    lhs = zeros(surf.nspan, surf.nspan)
    rhs = zeros(surf.nspan)

    for i = 1:surf.nspan
        for n = 1:surf.nspan
            nn = 2*n - 1
            lhs[i,n] = sin(nn*surf.psi[i])*(sin(surf.psi[i]) + (nn*pi/(2*surf.AR)))
        end
        rhs[i] = pi*sin(surf.psi[i])*surf.bc[i]/(2*surf.AR)
    end

    surf.bcoeff[:] = lhs \ rhs

    for i = 1:surf.nspan
        surf.a03d[i] = 0

        for n = 1:surf.nspan
            nn = 2*n - 1
            surf.a03d[i] -= real(nn)*surf.bcoeff[n]*sin(nn*surf.psi[i])/sin(surf.psi[i])
        end
    end
    return surf
end

function calc_a0a13d(surf::ThreeDSurfSimple)

    lhs = zeros(surf.nspan, surf.nspan)
    rhs = zeros(surf.nspan)

    for i = 1:surf.nspan
        integ0 = simpleTrapz(surf.s2d[i].cam_slope, surf.s2d[i].theta)
        integ1 = simpleTrapz(surf.s2d[i].cam_slope.*cos.(surf.s2d[i].theta), surf.s2d[i].theta)
        for n = 1:surf.nspan
            nn = 2*n - 1
            lhs[i,n] = sin(nn*surf.psi[i])*(sin(surf.psi[i]) + nn*pi/(2*surf.AR)*
                                            (cos(surf.s2d[i].kinem.alpha) + sin(surf.s2d[i].kinem.alpha)*(integ0 - integ1)/pi))
        end
        rhs[i] = pi*sin(surf.psi[i])*surf.bc[i]/(2*surf.AR)
    end

    surf.bcoeff[:] = lhs \ rhs

    for i = 1:surf.nspan
        integ0 = simpleTrapz(surf.s2d[i].cam_slope, surf.s2d[i].theta)
        integ1 = simpleTrapz(surf.s2d[i].cam_slope.*cos.(surf.s2d[i].theta), surf.s2d[i].theta)
        surf.a03d[i] = 0
        surf.aterm3d[1,i] = 0

        for n = 1:surf.nspan
            nn = 2*n - 1
            surf.a03d[i] -= real(nn)*surf.bcoeff[n]*sin(nn*surf.psi[i])/sin(surf.psi[i])*(cos(surf.s2d[i].kinem.alpha) + sin(surf.s2d[i].kinem.alpha)*integ0/pi)
            surf.aterm3d[1,i] += 2*real(nn)*surf.bcoeff[n]*sin(nn*surf.psi[i])/(sin(surf.psi[i])*pi)*sin(surf.s2d[i].kinem.alpha)*integ1
        end
    end
    return surf
end

function calc_a0a13d_wlev(surf::ThreeDSurfSimple)

    lhs = zeros(surf.nspan, surf.nspan)
    rhs = zeros(surf.nspan)

    levcount = 0
    for i = 1:surf.nspan
        integ0 = simpleTrapz(surf.s2d[i].cam_slope, surf.s2d[i].theta)
        integ1 = simpleTrapz(surf.s2d[i].cam_slope.*cos.(surf.s2d[i].theta), surf.s2d[i].theta)
        for n = 1:surf.nspan
            nn = 2*n - 1
            lhs[i,n] = sin(nn*surf.psi[i])*(sin(surf.psi[i]) + nn*pi/(2*surf.AR)*
                                            (cos(surf.s2d[i].kinem.alpha) + sin(surf.s2d[i].kinem.alpha)*(integ0 - integ1)/pi))
        end
        rhs[i] = (surf.levstr[i] + pi*surf.bc[i])*sin(surf.psi[i])/(2*surf.AR)
    end

    surf.bcoeff[:] = lhs \ rhs

    for i = 1:surf.nspan
        integ0 = simpleTrapz(surf.s2d[i].cam_slope, surf.s2d[i].theta)
        integ1 = simpleTrapz(surf.s2d[i].cam_slope.*cos.(surf.s2d[i].theta), surf.s2d[i].theta)
        surf.a03d[i] = 0
        surf.aterm3d[1,i] = 0

        for n = 1:surf.nspan
            nn = 2*n - 1
            surf.a03d[i] -= real(nn)*surf.bcoeff[n]*sin(nn*surf.psi[i])/sin(surf.psi[i])*(cos(surf.s2d[i].kinem.alpha) + sin(surf.s2d[i].kinem.alpha)*integ0/pi)
            surf.aterm3d[1,i] += 2*real(nn)*surf.bcoeff[n]*sin(nn*surf.psi[i])/(sin(surf.psi[i])*pi)*sin(surf.s2d[i].kinem.alpha)*integ1
        end
    end
    return surf
end

function calc_a0a13d_wlev2(surf::ThreeDSurfSimple, levstr)

    lhs = zeros(surf.nspan, surf.nspan)
    rhs = zeros(surf.nspan)

    for i = 1:surf.nspan
        integ0 = simpleTrapz(surf.s2d[i].cam_slope, surf.s2d[i].theta)
        integ1 = simpleTrapz(surf.s2d[i].cam_slope.*cos.(surf.s2d[i].theta), surf.s2d[i].theta)
        for n = 1:surf.nspan
            nn = 2*n - 1
            lhs[i,n] = sin(nn*surf.psi[i])*(sin(surf.psi[i]) + nn*pi/(2*surf.AR)*
                                            (cos(surf.s2d[i].kinem.alpha) + sin(surf.s2d[i].kinem.alpha)*(integ0 - integ1)/pi))
        end

        rhs[i] = (levstr[i] + pi*surf.bc[i])*sin(surf.psi[i])/(2*surf. AR)
    end

    surf.bcoeff[:] = lhs \ rhs

    for i = 1:surf.nspan
        integ0 = simpleTrapz(surf.s2d[i].cam_slope, surf.s2d[i].theta)
        integ1 = simpleTrapz(surf.s2d[i].cam_slope.*cos.(surf.s2d[i].theta), surf.s2d[i].theta)
        surf.a03d[i] = 0
        surf.aterm3d[1,i] = 0

        for n = 1:surf.nspan
            nn = 2*n - 1
            surf.a03d[i] -= real(nn)*surf.bcoeff[n]*sin(nn*surf.psi[i])/sin(surf.psi[i])*(cos(surf.s2d[i].kinem.alpha) + sin(surf.s2d[i].kinem.alpha)*integ0/pi)
            surf.aterm3d[1,i] += 2*real(nn)*surf.bcoeff[n]*sin(nn*surf.psi[i])/(sin(surf.psi[i])*pi)*sin(surf.s2d[i].kinem.alpha)*integ1
        end
    end
    return surf
end

function calc_a2toan3d(surf::ThreeDSurfSimple)
    for ia = 2:surf.naterm
        for i = 1:surf.nspan
            surf.aterm3d[ia,i] = 0
            integ = simpleTrapz(surf.s2d[i].cam_slope.*cos.(ia*surf.s2d[i].theta), surf.s2d[i].theta)

            for n = 1:surf.nspan
                nn = 2*n - 1
                surf.aterm3d[ia,i] += 2*real(nn)*surf.bcoeff[n]*sin(nn*surf.psi[i])/(sin(surf.psi[i])*pi)*sin(surf.s2d[i].kinem.alpha)*integ
            end
        end
    end
end

# Function for calculating the wake rollup
function wakeroll(surf3d::ThreeDSurfSimple, field3d::ThreeDFieldSimple, dt)

    for ispan = 1:surf3d.nspan
	surf = surf3d.s2d[ispan]
	curfield = field3d.f2d[ispan]

	wi = 0
        for n = 1:surf3d.nspan
            nn = 2*n -2
            wi += real(nn)*surf3d.bcoeff[n]*sin(nn*surf3d.psi[ispan])/sin(surf3d.psi[ispan])
        end

        nlev = length(curfield.lev)
        ntev = length(curfield.tev)
        nextv = length(curfield.extv)

        #Clean induced velocities
        for i = 1:ntev
            curfield.tev[i].vx = 0
            curfield.tev[i].vz = -wi
        end

        for i = 1:nlev
            curfield.lev[i].vx = 0
            curfield.lev[i].vz = -wi
        end

        for i = 1:nextv
            curfield.extv[i].vx = 0
            curfield.extv[i].vz = -wi
        end

        #Velocities induced by free vortices on each other
        mutual_ind([curfield.tev; curfield.lev; curfield.extv])

        #Add the influence of velocities induced by bound vortices
        utemp = zeros(ntev + nlev + nextv)
        wtemp = zeros(ntev + nlev + nextv)
        utemp, wtemp = ind_vel(surf.bv, [map(q -> q.x, curfield.tev); map(q -> q.x, curfield.lev); map(q -> q.x, curfield.extv)], [map(q -> q.z, curfield.tev); map(q -> q.z, curfield.lev); map(q -> q.z, curfield.extv) ])

        for i = 1:ntev
            curfield.tev[i].vx += utemp[i]
            curfield.tev[i].vz += wtemp[i]
        end
        for i = ntev+1:ntev+nlev
            curfield.lev[i-ntev].vx += utemp[i]
            curfield.lev[i-ntev].vz += wtemp[i]
        end
        for i = ntev+nlev+1:ntev+nlev+nextv
            curfield.extv[i-ntev-nlev].vx += utemp[i]
            curfield.extv[i-ntev-nlev].vz += wtemp[i]
        end

            #Add the influence of freestream velocities
        for i = 1:ntev
            curfield.tev[i].vx += curfield.u[1]
            curfield.tev[i].vz += curfield.w[1]
        end
        for i = 1:nlev
        curfield.lev[i].vx += curfield.u[1]
            curfield.lev[i].vz += curfield.w[1]
        end
        for i = 1:nextv
            curfield.extv[i].vx += curfield.u[1]
            curfield.extv[i].vz += curfield.w[1]
        end

        #Convect free vortices with their induced velocities
        for i = 1:ntev
            curfield.tev[i].x += dt*curfield.tev[i].vx
            curfield.tev[i].z += dt*curfield.tev[i].vz
        end
        for i = 1:nlev
            curfield.lev[i].x += dt*curfield.lev[i].vx
            curfield.lev[i].z += dt*curfield.lev[i].vz
        end
        for i = 1:nextv
            curfield.extv[i].x += dt*curfield.extv[i].vx
            curfield.extv[i].z += dt*curfield.extv[i].vz
        end
    end

    return field3d
end

function update_kinemCantilever(surf::ThreeDSurfSimple, dt::Float64, kinem::KinemCantilever, Phi::Array{Float64}, m_eta_s_etadd::Array{Float64}, m_eta_s_etad::Array{Float64}, m_eta_s_eta::Array{Float64}, cl::Array{Float64}, cm::Array{Float64})

    Qvecloc = zeros(4, surf.nspan)

    rho = 1.1
    q = 1. *pi/180

    pdyn = 0.5*rho*surf.uref^2

    #Update previous terms
    kinem.eta_pr = kinem.eta
    kinem.etad_pr3 = kinem.etad_pr2
    kinem.etad_pr2 = kinem.etad_pr
    kinem.etad_pr = kinem.etad
    kinem.etadd_pr3 = kinem.etadd_pr2
    kinem.etadd_pr2 = kinem.etadd_pr

    for i = 1:surf.nspan
        if i == 1
            yl = -0.5*surf.AR*surf.cref
            yu = 0.5*(surf.yle[i] + surf.yle[i+1])
        elseif i == surf.nspan
            yl = 0.5*(surf.yle[i] + surf.yle[i-1])
            yu = 0.
        else
            yl = 0.5*(surf.yle[i] + surf.yle[i-1])
            yu = 0.5*(surf.yle[i] + surf.yle[i+1])
        end
        dy = yu - yl
        Qvecloc[:,surf.nspan-i+1] = Phi[:,:,surf.nspan-i+1]'*pdyn*[-cl[i]*dy*surf.s2d[i].c; cm[i]*dy*surf.s2d[i].c^2; 0];
    end

    Qvec = sum(Qvecloc,dims=2)

    kinem.etadd_pr = m_eta_s_etadd\(-m_eta_s_etad*kinem.etad - m_eta_s_eta*kinem.eta + Qvec)

    kinem.etad = kinem.etad_pr + (dt/12.)*(23*kinem.etadd_pr - 16*kinem.etadd_pr2 + 5*kinem.etadd_pr3)

    kinem.eta = kinem.eta_pr + (dt/12)*(23*kinem.etad_pr - 16*kinem.etad_pr2 + 5*kinem.etad_pr3)

    for i = 1:surf.nspan
        vecphys = Phi[:,:,i]*kinem.eta
        vecphys_dot = Phi[:,:,i]*kinem.etad
        j = surf.nspan - i + 1
        surf.s2d[j].kinem.h = -vecphys[1]
        surf.s2d[j].kinem.alpha = vecphys[2] + q
        surf.s2d[j].kinem.hdot = -vecphys_dot[1]
        surf.s2d[j].kinem.alphadot = vecphys_dot[2]
    end

end


text:
# -*- encoding: utf-8 -*-
#
# The MIT License (MIT)
#
# Copyright © 2021 Matteo Foglieni and Riccardo Gervasoni
#

@testset "test_Rays" begin
     @testset "test_is_close" begin
          ray1 = Ray(Point(1.0, 2.0, 3.0), Vec(5.0, 4.0, -1.0))
          ray2 = Ray(Point(1.0, 2.0, 3.0), Vec(5.0, 4.0, -1.0))
          ray3 = Ray(Point(5.0, 1.0, 4.0), Vec(3.0, 9.0, 4.0))

          @test ray1 ≈ ray2
          @test !( ray1 ≈ ray3 )
     end

     @testset "test_at" begin
          ray = Ray(Point(1.0, 2.0, 4.0),Vec(4.0, 2.0, 1.0))
          @test at(ray, 0.0) ≈ ray.origin
          @test at(ray, 1.0) ≈ Point(5.0, 4.0, 5.0)
          @test at(ray, 2.0) ≈ Point(9.0, 6.0, 6.0)
     end

     @testset "test_transform" begin
          ray = Ray(Point(1.0, 2.0, 3.0), Vec(6.0, 5.0, 4.0))
          T = translation(Vec(10.0, 11.0, 12.0)) * rotation_x(π/2)
          transformed = T*ray

          @test transformed.origin ≈ Point(11.0, 8.0, 14.0)
          @test transformed.dir ≈ Vec(6.0, -4.0, 5.0)
     end
end

@testset "test_Camera" begin

     @testset "test_OrthogonalCamera" begin
          cam = OrthogonalCamera(2.0)
          ray1 = fire_ray(cam, 0.0, 0.0)
          ray2 = fire_ray(cam, 1.0, 0.0)
          ray3 = fire_ray(cam, 0.0, 1.0)
          ray4 = fire_ray(cam, 1.0, 1.0)

          # Verify that the rays are parallel by verifying that cross-products vanish
          @test 0.0 ≈ squared_norm(ray1.dir × ray2.dir)
          @test 0.0 ≈ squared_norm(ray1.dir × ray3.dir)
          @test 0.0 ≈ squared_norm(ray1.dir × ray4.dir)

          # Verify that the ray hitting the corners have the right coordinates
          @test at(ray1, 1.0) ≈ Point(0.0, 2.0, -1.0)
          @test at(ray2, 1.0) ≈ Point(0.0, -2.0, -1.0)
          @test at(ray3, 1.0) ≈ Point(0.0, 2.0, 1.0)
     end

     @testset "test_PerspectiveCamera" begin
          cam = PerspectiveCamera(1.0, 2.0)
          ray1 = fire_ray(cam, 0.0, 0.0)
          ray2 = fire_ray(cam, 1.0, 0.0)
          ray3 = fire_ray(cam, 0.0, 1.0)
          ray4 = fire_ray(cam, 1.0, 1.0)

          # Verify that all the rays depart from the same point
          @test ray1.origin ≈ ray2.origin
          @test ray1.origin ≈ ray3.origin
          @test ray1.origin ≈ ray4.origin

          # Verify that the ray hitting the corners have the right coordinates
          @test at(ray1, 1.0) ≈ Point(0.0, 2.0, -1.0)
          @test at(ray2, 1.0) ≈ Point(0.0, -2.0, -1.0)
          @test at(ray3, 1.0) ≈ Point(0.0, 2.0, 1.0)
          @test at(ray4, 1.0) ≈ Point(0.0, -2.0, 1.0)
     end

end

@testset "test_ImageTracer" begin

     @testset "test_uv_sub_mapping" begin
          img = HDRimage(4, 2)
          Pcam = PerspectiveCamera(1., 2.) # PAY ATTENSTION TO POSITIONAL ARGUMENTS!!!!!!!!!!!!!!!!!
          tracer = ImageTracer(img, Pcam)
          r1 = fire_ray(tracer, 0, 0, 2.5, 1.5)
          r2 = fire_ray(tracer, 2, 1)
          @test r1 ≈ r2
     end

     @testset "test_orientation" begin
          img = HDRimage(4, 2)
          Pcam = PerspectiveCamera(1., 2.) # PAY ATTENSTION TO POSITIONAL ARGUMENTS!!!!!!!!!!!!!!!!!
          tracer = ImageTracer(img, Pcam)

          top_left_ray = fire_ray(tracer, 0, 0, 0., 0.)
          bottom_right_ray = fire_ray(tracer, 3, 1, 1.0, 1.0)

          @test Point(0., 2., 1.) ≈ at(top_left_ray, 1.)
          @test Point(0., -2., -1.) ≈ at(bottom_right_ray, 1.)
     end

     @testset "test_image_coverage" begin
          img = HDRimage(4, 2)
          Pcam = PerspectiveCamera(1., 2.) # PAY ATTENSTION TO POSITIONAL ARGUMENTS!!!!!!!!!!!!!!!!!
          tracer = ImageTracer(img, Pcam)
          fire_all_rays!(tracer, x->RGB{Float32}(1.0, 2.0, 3.0))
          for row in tracer.img.height-1:-1:0, col in 0:tracer.img.width-1
               @test Raytracing.get_pixel(img, col, row) == RGB{Float32}(1.0, 2.0, 3.0)
          end
     end

end


text:
function PTVer(A,b; inf = 1e9)
    # Ax <= b
    (m,n) = size(A)

    # Initial big space
    Ver = [VerticeType(-inf*ones(n),GenNeighbor(zeros(n)),GenCuts(zeros(n)))]
    for i = 1:2^n-1
        P = zeros(n);
        point = Int2Bin(i,n)
        P[point .== 1] = inf
        P[point .== 0] = -inf
        push!(Ver,VerticeType(P,GenNeighbor(point),GenCuts(point)))
    end
    # The polytope
    CutNum = 2*n + 1
    for i = 1:m
        Cut = A[i,:]
        push!(Cut, -b[i])
        Ver = Vert(deepcopy(Ver),Cut,CutNum)
        CutNum += 1
    end

    # Removing infinity corners
    i = 1
    while i<=length(Ver)
        if any(Ver[i].Point .== inf) || any(Ver[i].Point .== -inf)
            deleteat!(Ver,i)
        else
            i += 1
        end
    end
    ver = zeros(length(Ver),n)
    for i = 1:length(Ver)
        ver[i,:] = Ver[i].Point
    end
    ver
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, Test

# zero eigenmode model
X = Dict(
         1 => [2.0, 3.0, 4.0],
         2 => [6.0, 3.0, 2.0],
         3 => [2.0, 5.0, 1.0],
         4 => [4.0, 3.0, 6.0])
X[5] = 1/2*(X[1] + X[2])
X[6] = 1/2*(X[2] + X[3])
X[7] = 1/2*(X[3] + X[1])
X[8] = 1/2*(X[1] + X[4])
X[9] = 1/2*(X[2] + X[4])
X[10] = 1/2*(X[3] + X[4])

element = Element(Tet10, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
update!(element, "youngs modulus", 480.0)
update!(element, "poissons ratio", 1/3)
update!(element, "geometry", X)
update!(element, "density", 105.0)

body = Problem(Elasticity, "TET", 3)
add_elements!(body, element)
analysis = Analysis(Modal)
add_problems!(analysis, body)
analysis.properties.nev = 30
run!(analysis)
w2 = analysis.properties.eigvals
w2_expected = [5.66054e-16, 1.08925e-15, 1.95035e-15, -2.18372e-15, -8.01069e-15, 9.26902e-15, 0.401407, 0.88248, 1.3185, 2.55833, 3.3538, 5.02371, 7.71933, 8.43733, 11.0521, 19.3262, 26.0944, 28.7033, 51.5814, 59.1422, 83.3597, 91.1507, 121.689, 126.753, 157.938, 160.344, 205.634, 324.608, 468.331]
@test isapprox(w2, w2; atol=1.0e-4)


text:
wrapper_config(; kwargs...) = WrapperConfig(; destfile="", kwargs...)

@test isempty(extensions(wrapper_config(wrap_core=false)))
@test length(extensions(wrapper_config())) > 200
@test length(extensions(wrapper_config(wrap_core=false, include_platforms=[PLATFORM_VI]))) < 5
@test length(extensions(wrapper_config(include_provisional_exts=true))) > length(extensions(wrapper_config()))
@test extensions(wrapper_config(include_provisional_exts=true)) == extensions(wrapper_config(include_platforms=[PLATFORM_PROVISIONAL]))


text:
#=
Julia has a few packages for visualization, depending on your tastes.

\toc

## Plots

Plots is probably the most well-known, and also one of the oldest.
It provides an abstraction layer and you can select the backend that actually does the plotting.
See the backend list [here](https://docs.juliaplots.org/latest/backends/).

When you install Plots it also install GR, which is quite fast and has some decent graphics.
GR is the default choice, but if you alternate from it, you can activate it again with `gr()`.

The PyPlot package provides another common backend.
It uses Python's matplotlib, which provides very nice graphics, although it is a bit slower than GR.
Change to PyPlot using `pyplot()`.

You can also use UnicodePlots for plotting directly into the REPL.
Change using `unicodeplots()`.

PlotlyJS is also quite well-liked, because it usually allows some interactivity in the plots.
Use `plotlyjs()` to change.

Some packages are more fickle than others to install.
GR, for instance, decided to not work with GitHub Actions for this site.
For that reason, I use mostly PyPlot.

Here are some examples with PyPlot.
Try them with other backends.
=#
using LinearAlgebra, Plots, Random
pyplot()

Random.seed!(0)
n = 100
x = range(-1, 1, length=n)
y = 5 * exp.(-10x.^2) + x + 0.2 * randn(n)
Etr, Ete = Float64[], Float64[]
Ite = shuffle(1:n)[1:2 * div(n, 5)]
Itr = setdiff(1:n, Ite)
pmax = 6
plt = plot(title="ML stuff", layout=grid(1, 2), size=(800,400))
scatter!(plt[1], x, y, leg=false)
for p = 1:pmax
  X = [x[i]^j for i = 1:n, j = 0:p]
  β = X[Itr,:] \ y[Itr]
  plot!(plt[1], t -> β[1] + sum(β[j+1] * t^j for j = 1:p), -1, 1)
  r = y - X * β
  push!(Etr, norm(r[Itr])^2 / length(Itr))
  push!(Ete, norm(r[Ite])^2 / length(Ite))
end
plot!(title="MSE")
plot!(plt[2], Etr, m=(3,:blue,:square), lab="Train")
plot!(plt[2], Ete, m=(3,:red,:circle), lab="Test")
png(joinpath(@OUTPUT, "vis-plots-1"))

# \fig{vis-plots-1.png}

X = [
  0.4 * randn(div(n, 2), 2) .+ 3;
  1.2 * randn(div(n, 2), 2) .- 1;
]
h = 1.0
K(x) = exp(-norm(x)^2 / h) / h
f(x) = sum(K(x - X[i,:]) for i = 1:n) / n

plt = plot(layout=grid(2,2), leg=false)
scatter!(plt[1], X[:,1], X[:,2], title="scatter")
contour!(plt[2], -5:0.1:5, -5:0.1:5, (x,y) -> f([x;y]), title="countour", levels=50)
surface!(plt[3], -5:0.1:5, -5:0.1:5, (x,y) -> f([x;y]), title="surface")
surface!(plt[4], -5:0.1:5, -5:0.1:5, (x,y) -> f([x;y]), title="surface", camera=(70,10))
plt
png(joinpath(@OUTPUT, "vis-plots-2"))

# \fig{vis-plots-2.png}

θg = range(0, 2π, length=120)
anim = Animation()
n = 3
plot(leg=false, axis=false, grid=false, size=(600,600))
for i = 1:6
  θr = range(0, 2π, length=n + 1)
  plot!(cos.(θr), sin.(θr), c=:blue, opacity=0.5, lw=3)
  radius = 1 / cos(π / n)
  plot!(radius * cos.(θr .+ π / n), radius * sin.(θr .+ π / n), c=:red, opacity=0.5, lw=3)
  xlims!(-2, 2)
  ylims!(-2, 2)
  frame(anim)
  global n = 2n
end

gif(anim, joinpath(@OUTPUT, "vis-plots-anim-1.gif"), fps=1)

# \fig{vis-plots-anim-1.gif}

#=
## PyPlot

PyPlot uses Python's matplotlib, as said above.
We can use it as a backend to Plots, but it can also be used directly.
The main advantage are that the syntax is closer to matplotlib's syntax, and Plots doesn't support everything that PyPlot does.
=#

using PyPlot
x = range(0, 2π, length=1000)
y = sin.(3 * x + 4 * cos.(2 * x))
PyPlot.figure()
PyPlot.plot(x, y, color="red", linewidth=2.0, linestyle="--")
PyPlot.title("A sinusoidally modulated sinusoid")
PyPlot.savefig(joinpath(@OUTPUT, "vis-pyplot-1"))

# \fig{vis-pyplot-1.png}

#=
## Makie

Makie is a complete alternative to Plots.
It tries to achieve high performance using heavier backends, such as OpenGL.

## StatsPlots

This package provides "recipes for Plots.jl" - i.e., new plotting functions - for statistical plots.
=#
using RDatasets, StatsPlots

df = dataset("datasets", "iris")
@df df cornerplot(cols(1:4), compact=true)
png(joinpath(@OUTPUT, "vis-statsplots-1"))

# \fig{vis-statsplots-1.png}

plt = plot(layout=grid(2, 2))
for i = 1:4
  StatsPlots.violin!(plt[i], df.Species, df[!,i], lab="", ylabel=names(df)[i])
  StatsPlots.boxplot!(plt[i], df.Species, df[!,i], lab="", fillalpha=0.6, m=(0))
  StatsPlots.dotplot!(plt[i], df.Species, df[!,i], lab="", m=(:white, 2), opacity=0.6)
end
png(plt, joinpath(@OUTPUT, "vis-statsplots-2"))

# \fig{vis-statsplots-2.png}

#=
## GraphRecipes

Plot recipes for Graphs.
=#
using GraphRecipes
g = rand(0:1, 5, 5)
graphplot(g)
png(joinpath(@OUTPUT, "vis-graphrecipes-1"))

# \fig{vis-graphrecipes-1.png}

max_stars_repo_name: chunggnuhc/julia-tutorial-at-nlesc

max_issues_repo_name: chunggnuhc/julia-tutorial-at-nlesc

max_forks_repo_name: chunggnuhc/julia-tutorial-at-nlesc


text:
dsts = Dict(
    0 => cu([3, 4, 5, 6, 7]),
    1 => cu([3 3 4 4 5;
             5 5 6 6 7]),
)
srcs = Dict(
    (0, true) => cu(ones(Int, 3, 4)),
    (0, false) => cu(ones(Int, 3) * collect(1:4)'),
    (1, true) => cu(ones(Int, 2, 3, 4)),
    (1, false) => cu([1, 2] .* reshape(ones(Int, 3) * collect(1:4)', 1,3,4)),
)
idxs = [
    cu([1 2 3 4;
        4 2 1 3;
        3 5 5 3]),  # integer index
    cu([(1,) (2,) (3,) (4,);
        (4,) (2,) (1,) (3,);
        (3,) (5,) (5,) (3,)]),  # tuple index
]

types = [CuArray{Int32}, CuArray{Int64}, CuArray{Float32}, CuArray{Float64}]

@testset "scatter" begin
    for T = types
        @testset "$(T)" begin
            @testset "+" begin
                for idx = idxs, dims = [0, 1]
                    mutated = true
                    gputest((dst, src) -> NNlib.scatter!(+, dst, src, idx), T(copy(dsts[dims])), T(srcs[(dims, mutated)]), checkgrad=true)

                    mutated = false
                    gputest(src -> NNlib.scatter(+, src, idx), T(srcs[(dims, mutated)]), checkgrad=true)
                end
            end

            @testset "-" begin
                for idx = idxs, dims = [0, 1]
                    mutated = true
                    gputest((dst, src) -> NNlib.scatter!(-, dst, src, idx), T(copy(dsts[dims])), T(srcs[(dims, mutated)]), checkgrad=true)

                    mutated = false
                    gputest(src -> NNlib.scatter(-, src, idx), T(srcs[(dims, mutated)]), checkgrad=true)
                end
            end

            @testset "max" begin
                for idx = idxs, dims = [0, 1]
                    mutated = true
                    gputest((dst, src) -> NNlib.scatter!(max, dst, src, idx), T(copy(dsts[dims])), T(srcs[(dims, mutated)]), checkgrad=true)

                    mutated = false
                    gputest(src -> NNlib.scatter(max, src, idx), T(srcs[(dims, mutated)]), checkgrad=true)
                end
            end

            @testset "min" begin
                for idx = idxs, dims = [0, 1]
                    mutated = true
                    gputest((dst, src) -> NNlib.scatter!(min, dst, src, idx), T(copy(dsts[dims])), T(srcs[(dims, mutated)]), checkgrad=true)

                    mutated = false
                    gputest(src -> NNlib.scatter(min, src, idx), T(srcs[(dims, mutated)]), checkgrad=true)
                end
            end
        end
    end

    for T = [CuArray{Float32}, CuArray{Float64}]
        @testset "$(T)" begin
            @testset "*" begin
                for idx = idxs, dims = [0, 1]
                    mutated = true
                    gputest((dst, src) -> NNlib.scatter!(*, dst, src, idx), T(copy(dsts[dims])), T(srcs[(dims, mutated)]), checkgrad=true)

                    mutated = false
                    gputest(src -> NNlib.scatter(*, src, idx), T(srcs[(dims, mutated)]), checkgrad=true)
                end
            end

            @testset "/" begin
                for idx = idxs, dims = [0, 1]
                    mutated = true
                    gputest((dst, src) -> NNlib.scatter!(/, dst, src, idx), T(copy(dsts[dims])), T(srcs[(dims, mutated)]), checkgrad=true)

                    mutated = false
                    gputest(src -> NNlib.scatter(/, src, idx), T(srcs[(dims, mutated)]), checkgrad=true)
                end
            end

            @testset "mean" begin
                for idx = idxs, dims = [0, 1]
                    mutated = true
                    gputest((dst, src) -> NNlib.scatter!(mean, dst, src, idx), T(copy(dsts[dims])), T(srcs[(dims, mutated)]), checkgrad=true)

                    mutated = false
                    gputest(src -> NNlib.scatter(mean, src, idx), T(srcs[(dims, mutated)]), checkgrad=true)
                end
            end
        end
    end
end


text:
atreplinit() do repl
    try
        @eval using Revise
        @async Revise.wait_steal_repl_backend()
    catch e
        @warn(e.msg)
    end

push!(Revise.dont_watch_pkgs, :Gadfly)
push!(Revise.dont_watch_pkgs, :Compose)
end

max_stars_repo_name: Greg-R/vimrctmuxconfig

max_issues_repo_name: Greg-R/vimrctmuxconfig

max_forks_repo_name: Greg-R/vimrctmuxconfig


text:
# ====================================================================
# Component structure
#
mutable struct Gaussian_1D <: AbstractComponent
    norm::Parameter
    center::Parameter
    sigma::Parameter

    function Gaussian_1D(norm::Number, center::Number, sigma::Number)
        @assert norm  > 0
        @assert sigma > 0

        out = new(Parameter(norm), Parameter(center), Parameter(sigma))
        out.norm.low = 0
        out.sigma.low = 0
        return out
    end
end

mutable struct Gaussian_2D <: AbstractComponent
    norm::Parameter
    centerX::Parameter
    centerY::Parameter
    sigmaX::Parameter
    sigmaY::Parameter
    angle::Parameter

    function Gaussian_2D(norm::Number, centerX::Number, centerY::Number, sigmaX::Number, sigmaY::Number, angle::Number)
        @assert norm   > 0
        @assert sigmaX > 0
        @assert sigmaY > 0

        out = new(Parameter(norm), Parameter(centerX), Parameter(centerY), Parameter(sigmaX), Parameter(sigmaY), Parameter(angle))
        out.norm.low = 0
        out.sigmaX.low = 0
        out.sigmaY.low = 0
        return out
    end
end

Gaussian(norm, center, sigma) = Gaussian_1D(norm, center, sigma)
Gaussian(norm, centerX, centerY, sigmaX, sigmaY, angle) = Gaussian_2D(norm, centerX, centerY, sigmaX, sigmaY, angle)
function Gaussian(norm, centerX, centerY, sigma)
    out = Gaussian_2D(norm, centerX, centerY, sigma, sigma, 0.)
    out.sigmaX.free = false
    out.sigmaX.expr = "this_sigmaY"
    out.angle.free = false
    return out
end

# ====================================================================
# Prepare component `cdata`
mutable struct Gaussian_cdata
    ix::Vector{Int}
    iy::Vector{Int}
    Gaussian_cdata() = new(Vector{Int}(), Vector{Int}())
end

ceval_data(domain::Domain_1D, comp::Gaussian_1D) = (Gaussian_cdata(), length(domain))
ceval_data(domain::Domain_2D, comp::Gaussian_2D) = (Gaussian_cdata(), length(domain))

# ====================================================================
# Evaluate component
function evaluate(c::CompEval{Domain_1D, Gaussian_1D},
                  norm, center, sigma)
    # TODO: optimize using cdata
    x = c.domain[1]
    @. (c.eval = exp( ((x - center) / sigma)^2. / (-2.)) /
        2.5066282746310002 / sigma * norm) # sqrt(2pi) = 2.5066282746310002
end

function evaluate(c::CompEval{Domain_2D, Gaussian_2D},
                   norm, centerX, centerY, sigmaX, sigmaY, angle)
    angle *= -pi / 180.
    a =  (cos(angle) / sigmaX)^2 / 2  +  (sin(angle) / sigmaY)^2 / 2
    b = -sin(2angle) / sigmaX^2  / 2  +  sin(2angle) / sigmaY^2  / 2
    c =  (sin(angle) / sigmaX)^2 / 2  +  (cos(angle) / sigmaY)^2 / 2

    # TODO: optimize using cdata
    x = c.domain[1]
    y = c.domain[2]

    @. (c.eval = norm *
        exp(
            -(
                a * (x - centerX)^2. +
                b * (x - centerX) * (y - centerY) +
                c *                 (y - centerY)^2.
            )
        ) / 6.283185307179586 / sigmaX / sigmaY) # 2pi = 6.283185307179586
end


text:
## estimae the stellar mass using a known age of the ]OC
##
rootdir = ENV["GAIA_ROOT"]
push!(LOAD_PATH,"$rootdir/run/src")
using GaiaClustering

using PyCall
## PopStar modules...
synthetic= pyimport("popstar.synthetic")
evolution= pyimport("popstar.evolution")
atmospheres= pyimport("popstar.atmospheres")
reddening= pyimport("popstar.reddening")

using Distributions , Statistics , DataFrames , CSV
using Distances, Random , Printf, Glob

import PyPlot

## directory
wdir=    "$rootdir/e2e_products"
plotdir= "$wdir/plotsIndividuals"
ocdir=   "$wdir/oc"
isodir=  "$wdir/isochrones"

let
## ISOCHRONES parameters
VOTAGE= "votlist-age.csv"    ## csv file with votname,log_age

## Main Loop ####################################################
cd(ocdir)
files= glob("*-oc.csv")
cd(wdir)
votage= CSV.read(VOTAGE, delim=";")

for f in files
    println("## $f")
    fvot= f[1:end-7] ; vname= "$fvot.vot" ; fname= f[1:end-4]
    ocmassfile= "$ocdir/$fname-mass.csv"
    logAge= votage[(votage.votname .== vname), :log_age][1]
    println("$ocdir/$f")
    df= CSV.read("$ocdir/$f" , delim= ";")

    dist= median(df.distance)
    AKs= 0.0                        # extinction in mags
    metallicity= 0.0                # Metallicity in [M/H]
    nstar= length(df.distance)

    ## Define evolution/atmosphere models and extinction law
    if logAge>6.5
        evo_model = evolution.MergedBaraffePisaEkstromParsec()
    else
        evo_model= evolution.MISTv1()
    end

    atm_func = atmospheres.get_merged_atmosphere
    red_law = reddening.RedLawHosek18b()

    ## Filter
    filt_list = [ "gaia,dr2_rev,Gbp", "gaia,dr2_rev,G"  ,"gaia,dr2_rev,Grp"]

    #####################################
    println("### Computing isochrones..")
    iso = synthetic.IsochronePhot(logAge, AKs, dist, metallicity=metallicity,
                            evo_model=evo_model, atm_func=atm_func,
                            red_law=red_law, filters=filt_list, iso_dir= isodir)

    data= iso.points
    ndata= size(data)[1]

    mass= zeros(ndata)
    Gmag= zeros(ndata)
    BPmag= zeros(ndata)
    RPmag= zeros(ndata)
    color= zeros(ndata)

    for i in 1:ndata
        mass[i]= data[i][4]
        Gmag[i]= data[i][10]
        BPmag[i]= data[i][9]
        RPmag[i]= data[i][11]
        color[i]= BPmag[i]-RPmag[i]
    end

    ## Cluster
    BmR= df.bp .- df.rp
    GMAG= df.gbar + 5 .* log10.(df.distance) .- 17.

    ### Stellar mass
    ################
    nxy= length(mass)
    A= Array{Float64}(undef,2,nxy)
    for i in 1:nxy
        A[1,i]= color[i]
        A[2,i]= Gmag[i]
    end

    d= Distances.Euclidean()
    mstar= zeros(length(BmR))

    for i in 1:length(BmR)
        starm= [BmR[i]; GMAG[i]]
        r= Distances.colwise(d, A, starm)

        idx= argmin(r)
        mstar[i]= mass[idx]
    end

    df.mass= mstar

    ### save the OC df
    CSV.write(ocmassfile, df, delim=";")
    println("\n\n")
end
###
println("## Done")
end

max_stars_repo_name: bosscha/gaia-shock

max_issues_repo_name: bosscha/gaia-shock

max_forks_repo_name: bosscha/gaia-shock


text:
using Brillouin: wignerseitz

@testset "Test `WignerSeitzCell` against `Brillouin`" begin
    # Data from https://thchr.github.io/Brillouin.jl/dev/wignerseitz/
    lattice = Lattice([0.0, 0.5, 0.5], [0.5, 0.0, 0.5], [0.5, 0.5, 0.0])
    ws = WignerSeitzCell(lattice)
    cᴿ = wignerseitz([[0.0, 0.5, 0.5], [0.5, 0.0, 0.5], [0.5, 0.5, 0.0]])
    @test vertices(ws) == [
        [0.25, -0.75, 0.25],
        [-0.5, -0.5, 0.5],
        [-0.75, 0.25, 0.25],
        [0.75, -0.25, -0.25],
        [-0.25, -0.25, -0.25],
        [0.25, 0.25, -0.75],
        [-0.5, 0.5, -0.5],
        [0.5, -0.5, -0.5],
        [-0.25, -0.25, 0.75],
        [0.5, -0.5, 0.5],
        [-0.25, 0.75, -0.25],
        [0.25, 0.25, 0.25],
        [0.5, 0.5, -0.5],
        [-0.5, 0.5, 0.5],
    ]
    @test ws.faces == [
        [3, 7, 5, 2],
        [6, 8, 5, 7],
        [7, 3, 14, 11],
        [11, 13, 6, 7],
        [14, 3, 2, 9],
        [1, 2, 5, 8],
        [1, 10, 9, 2],
        [13, 11, 14, 12],
        [14, 9, 10, 12],
        [4, 8, 6, 13],
        [4, 10, 1, 8],
        [13, 12, 10, 4],
    ]
    @test faces(ws) == [
        [[-0.75, 0.25, 0.25], [-0.5, 0.5, -0.5], [-0.25, -0.25, -0.25], [-0.5, -0.5, 0.5]],
        [[0.25, 0.25, -0.75], [0.5, -0.5, -0.5], [-0.25, -0.25, -0.25], [-0.5, 0.5, -0.5]],
        [[-0.5, 0.5, -0.5], [-0.75, 0.25, 0.25], [-0.5, 0.5, 0.5], [-0.25, 0.75, -0.25]],
        [[-0.25, 0.75, -0.25], [0.5, 0.5, -0.5], [0.25, 0.25, -0.75], [-0.5, 0.5, -0.5]],
        [[-0.5, 0.5, 0.5], [-0.75, 0.25, 0.25], [-0.5, -0.5, 0.5], [-0.25, -0.25, 0.75]],
        [[0.25, -0.75, 0.25], [-0.5, -0.5, 0.5], [-0.25, -0.25, -0.25], [0.5, -0.5, -0.5]],
        [[0.25, -0.75, 0.25], [0.5, -0.5, 0.5], [-0.25, -0.25, 0.75], [-0.5, -0.5, 0.5]],
        [[0.5, 0.5, -0.5], [-0.25, 0.75, -0.25], [-0.5, 0.5, 0.5], [0.25, 0.25, 0.25]],
        [[-0.5, 0.5, 0.5], [-0.25, -0.25, 0.75], [0.5, -0.5, 0.5], [0.25, 0.25, 0.25]],
        [[0.75, -0.25, -0.25], [0.5, -0.5, -0.5], [0.25, 0.25, -0.75], [0.5, 0.5, -0.5]],
        [[0.75, -0.25, -0.25], [0.5, -0.5, 0.5], [0.25, -0.75, 0.25], [0.5, -0.5, -0.5]],
        [[0.5, 0.5, -0.5], [0.25, 0.25, 0.25], [0.5, -0.5, 0.5], [0.75, -0.25, -0.25]],
    ]
    @test cartesianize(cᴿ).verts == vertices(WignerSeitzCell(lattice), true)
    @testset "Construct Brillouin zones from `WignerSeitzCell`" begin
        # Data from https://thchr.github.io/Brillouin.jl/dev/wignerseitz/
        reci_lattice = inv(lattice)
        ws = WignerSeitzCell(reci_lattice)
        cᴳ = wignerseitz([[-1.0, 1.0, 1.0], [1.0, -1.0, 1.0], [1.0, 1.0, -1.0]])
        @test vertices(ws) == [
            [0.25, -0.5, -0.25],
            [0.5, -0.25, 0.25],
            [0.25, -0.25, 0.5],
            [-0.25, 0.25, -0.5],
            [0.25, -0.25, -0.5],
            [-0.25, -0.5, 0.25],
            [0.25, 0.75, 0.5],
            [0.75, 0.25, 0.5],
            [0.25, 0.5, 0.75],
            [0.5, 0.25, 0.75],
            [-0.25, -0.75, -0.5],
            [-0.5, -0.75, -0.25],
            [-0.5, -0.25, -0.75],
            [-0.25, -0.5, -0.75],
            [-0.75, -0.25, -0.5],
            [-0.75, -0.5, -0.25],
            [-0.25, 0.25, 0.5],
            [-0.5, -0.25, 0.25],
            [-0.5, 0.25, -0.25],
            [-0.25, 0.5, 0.25],
            [0.5, 0.25, -0.25],
            [0.25, 0.5, -0.25],
            [0.5, 0.75, 0.25],
            [0.75, 0.5, 0.25],
        ]
        @test ws.faces == [
            [9, 7, 20, 17],
            [24, 21, 22, 23],
            [8, 24, 23, 7, 9, 10],
            [10, 3, 2, 8],
            [5, 21, 24, 8, 2, 1],
            [1, 11, 14, 5],
            [2, 3, 6, 12, 11, 1],
            [5, 14, 13, 4, 22, 21],
            [19, 20, 7, 23, 22, 4],
            [4, 13, 15, 19],
            [14, 11, 12, 16, 15, 13],
            [9, 17, 18, 6, 3, 10],
            [17, 20, 19, 15, 16, 18],
            [18, 16, 12, 6],
        ]
        @test faces(ws) == [
            [[0.25, 0.5, 0.75], [0.25, 0.75, 0.5], [-0.25, 0.5, 0.25], [-0.25, 0.25, 0.5]],
            [[0.75, 0.5, 0.25], [0.5, 0.25, -0.25], [0.25, 0.5, -0.25], [0.5, 0.75, 0.25]],
            [
                [0.75, 0.25, 0.5],
                [0.75, 0.5, 0.25],
                [0.5, 0.75, 0.25],
                [0.25, 0.75, 0.5],
                [0.25, 0.5, 0.75],
                [0.5, 0.25, 0.75],
            ],
            [[0.5, 0.25, 0.75], [0.25, -0.25, 0.5], [0.5, -0.25, 0.25], [0.75, 0.25, 0.5]],
            [
                [0.25, -0.25, -0.5],
                [0.5, 0.25, -0.25],
                [0.75, 0.5, 0.25],
                [0.75, 0.25, 0.5],
                [0.5, -0.25, 0.25],
                [0.25, -0.5, -0.25],
            ],
            [
                [0.25, -0.5, -0.25],
                [-0.25, -0.75, -0.5],
                [-0.25, -0.5, -0.75],
                [0.25, -0.25, -0.5],
            ],
            [
                [0.5, -0.25, 0.25],
                [0.25, -0.25, 0.5],
                [-0.25, -0.5, 0.25],
                [-0.5, -0.75, -0.25],
                [-0.25, -0.75, -0.5],
                [0.25, -0.5, -0.25],
            ],
            [
                [0.25, -0.25, -0.5],
                [-0.25, -0.5, -0.75],
                [-0.5, -0.25, -0.75],
                [-0.25, 0.25, -0.5],
                [0.25, 0.5, -0.25],
                [0.5, 0.25, -0.25],
            ],
            [
                [-0.5, 0.25, -0.25],
                [-0.25, 0.5, 0.25],
                [0.25, 0.75, 0.5],
                [0.5, 0.75, 0.25],
                [0.25, 0.5, -0.25],
                [-0.25, 0.25, -0.5],
            ],
            [
                [-0.25, 0.25, -0.5],
                [-0.5, -0.25, -0.75],
                [-0.75, -0.25, -0.5],
                [-0.5, 0.25, -0.25],
            ],
            [
                [-0.25, -0.5, -0.75],
                [-0.25, -0.75, -0.5],
                [-0.5, -0.75, -0.25],
                [-0.75, -0.5, -0.25],
                [-0.75, -0.25, -0.5],
                [-0.5, -0.25, -0.75],
            ],
            [
                [0.25, 0.5, 0.75],
                [-0.25, 0.25, 0.5],
                [-0.5, -0.25, 0.25],
                [-0.25, -0.5, 0.25],
                [0.25, -0.25, 0.5],
                [0.5, 0.25, 0.75],
            ],
            [
                [-0.25, 0.25, 0.5],
                [-0.25, 0.5, 0.25],
                [-0.5, 0.25, -0.25],
                [-0.75, -0.25, -0.5],
                [-0.75, -0.5, -0.25],
                [-0.5, -0.25, 0.25],
            ],
            [
                [-0.5, -0.25, 0.25],
                [-0.75, -0.5, -0.25],
                [-0.5, -0.75, -0.25],
                [-0.25, -0.5, 0.25],
            ],
        ]
        @test cartesianize(cᴳ).verts == vertices(WignerSeitzCell(reci_lattice), true)
    end
end


text:
using Revise # reduce need for recompile
using Plots
using LinearAlgebra

push!(LOAD_PATH, "./src") # user defined modules
using CommonUtils
using Basis2DTri
using UniformTriMesh

using SetupDG
using UnPack

"Define physical and discretization parameters"
N   = 3 # The order of approximation
K1D = 16 # number of elements along each edge of a rectangle
CFL = 1 # relative size of a time-step
T   = 4 # final time

"=========== Setup reference element and mesh ============="

# construct mesh
VX,VY,EToV = uniform_tri_mesh(K1D,K1D)

# intialize reference operators
rd = init_reference_tri(N)

# initialize physical mesh data
md = init_mesh((VX,VY),EToV,rd)

#Make boundary maps periodic
@unpack Nfaces,Vf = rd
@unpack x,y,K,mapM,mapP,mapB = md
xf,yf = (x->Vf*x).((x,y))
LX,LY = (x->maximum(x)-minimum(x)).((VX,VY)) # find lengths of domain
mapPB = build_periodic_boundary_maps(xf,yf,LX,LY,Nfaces*K,mapM,mapP,mapB)
mapP[mapB] = mapPB
@pack! md = mapP

"======== Define initial coefficients and time-stepping =========="

"Define the initial conditions by interpolation"
p = @. exp(-10*(x^2+y^2))
u = @. 0*x
v = @. 0*x

"Define function to evaluate the RHS: Q = (p,u,v)"
function rhs(Q, rd::RefElemData, md::MeshData)

    @unpack Dr,Ds,LIFT,Vf = rd
    @unpack rxJ,sxJ,ryJ,syJ,J = md
    @unpack nxJ,nyJ,sJ = md
    @unpack mapP,mapB = md

    Qf = (x->Vf*x).(Q)
    QP = (x->x[mapP]).(Qf)
    dp,du,dv = QP.-Qf

    tau = .5
    dun = @. (du*nxJ + dv*nyJ)/sJ
    pflux = @. .5*(du*nxJ + dv*nyJ) - tau*dp*sJ
    uflux = @. .5*dp*nxJ - tau*dun*nxJ
    vflux = @. .5*dp*nyJ - tau*dun*nyJ

    pr,ur,vr = (x->Dr*x).(Q)
    ps,us,vs = (x->Ds*x).(Q)
    dudx = @. rxJ*ur + sxJ*us
    dvdy = @. ryJ*vr + syJ*vs
    dpdx = @. rxJ*pr + sxJ*ps
    dpdy = @. ryJ*pr + syJ*ps
    rhsp = dudx + dvdy + LIFT*pflux
    rhsu = dpdx        + LIFT*uflux
    rhsv = dpdy        + LIFT*vflux

    return (x->-x./J).((rhsp,rhsu,rhsv))
end

"Time integration coefficients"
rk4a,rk4b,rk4c = rk45_coeffs()
CN = (N+1)*(N+2)/2  # estimated trace constant
dt = CFL * 2 / (CN*K1D)
Nsteps = convert(Int,ceil(T/dt))
dt = T/Nsteps

"Perform time-stepping"
Q = [p,u,v] # use mutable arrays for Q, resQ for scoping
resQ = [zeros(size(x)) for i in eachindex(Q)]
for i = 1:Nsteps
    for INTRK = 1:5
        time = i*dt + rk4c[INTRK]*dt
        rhsQ = rhs(Q,rd,md)
        @. resQ = rk4a[INTRK]*resQ + dt*rhsQ
        @. Q    = Q + rk4b[INTRK]*resQ
    end

    if i%10==0 || i==Nsteps
        println("Number of time steps $i out of $Nsteps")
    end
end

# plotting nodes
gr(aspect_ratio=1, legend=false,
   markerstrokewidth=0,markersize=2,
   camera=(0,90))

@unpack Vp = rd
vv = Vp*Q[1]
scatter(Vp*x,Vp*y,vv,zcolor=vv)

max_stars_repo_name: jlchan/ESDG-CNS

max_issues_repo_name: jlchan/ESDG-CNS

max_forks_repo_name: jlchan/ESDG-CNS


text:
using NeuralQuantum
using Test
using NeuralQuantum: flipat!

function test_state_properties(s, _ldim, _n_sites, _type, _ctype)
    @test local_dimension(s) == _ldim
    @test nsites(s) == _n_sites
    @test spacedimension(s) == local_dimension(s)^nsites(s)
    @test eltype(s) == _type
    @test typeof(config(s)) <: _ctype
end

function test_state_dyn_single(s)
    l = local_dimension(s)
    n = nsites(s)
    dim = spacedimension(s)

    configurations = [];
    indices = [];
    for i=1:dim
        set_index!(s, i)
        push!(indices, index(s))
        push!(configurations, config(i))
    end

    @test indices == 1:dim
    @test length(unique(configurations)) == dim

    for rep=1:30
        rand!(s)
        for i=1:length(config(s))
            v = config(s)[i]
            old, new = flipat!(s, i)
            @test old == v
            @test new != v
            setat!(s, i, old)
            @test config(s)[i] == v
        end
    end
end

function test_state_dyn_double(s)
    l = local_dimension(s)
    n = nsites(s)
    dim = spacedimension(s)

    configurations = [];
    indices = [];
    for i=1:dim
        set_index!(s, i)
        push!(indices, index(s))
        push!(configurations, config(i))
    end

    @test indices == 1:dim
    @test length(unique(configurations)) == dim

    for rep=1:30
        rand!(s)
        for i=1:length(config(s))
            fl = length(first(config(s)))
            if i > fl
                v = first(config(s))[i-fl]
                old, new = flipat!(s, i)
                @test old == v
                @test new != v
                setat!(s, i, old)
                @test first(config(s))[i-fl] == v
            else
                v = last(config(s))[i]
                old, new = flipat!(s, i)
                @test old == v
                @test new != v
                setat!(s, i, old)
                @test last(config(s))[i] == v
            end
        end
    end
end

@testset "NAryState" begin
    s = NAryState(3, 4)
    test_state_properties(s, 3, 4, Float32, AbstractArray)
    test_state_dyn_single(s)
end

@testset "Double-NAryState" begin
    s = DoubleState(NAryState(3, 4))
    test_state_properties(s, 3, 4*2, Float32, Tuple{AbstractArray, AbstractArray})
    test_state_dyn_double(s)
end

@testset "WrappedChanges : NAryState" begin
    sb = NAryState(3, 4)
    s = ModifiedState(sb)
    test_state_properties(s, 3, 4, Float32, AbstractArray)
    test_state_dyn_single(s)
end

@testset "WrappedChanges : Double-NAryState" begin
    sb = ModifiedState(NAryState(3, 4))
    s = DoubleState(sb)
    test_state_properties(s, 3, 4*2, Float32, Tuple{AbstractArray, AbstractArray})
    test_state_dyn_double(s)
end

@testset "LU: WrappedChanges : NAryState" begin
    sb = NAryState(3, 4)
    s = LUState(ModifiedState(sb), nothing)
    test_state_properties(s, 3, 4, Float32, AbstractArray)
    test_state_dyn_single(s)
end

@testset "LU: WrappedChanges : Double-NAryState" begin
    sb = ModifiedState(NAryState(3, 4))
    s = LUState(DoubleState(sb), nothing)
    test_state_properties(s, 3, 4*2, Float32, Tuple{AbstractArray, AbstractArray})
    test_state_dyn_double(s)
end


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

using Markdown
using InteractiveUtils

# ╔═╡ 6f58c330-cce4-11eb-3e45-6787cd75141f
begin
	using Futbolista: DataLoaders, Plotter
end

# ╔═╡ e850ce3c-f57e-41ab-acb3-958a1b5951a9
begin
	const sb = DataLoaders.Statsbomb
	const pm = Plotter.Passmap

	ldr = sb.SBLoader("local", raw"C:\repository\open-data\data")
	events_df = sb.load_events(ldr, 7298)
	p = pm.plot_passmap(events_df, "home", "mean")
end

# ╔═╡ Cell order:
# ╠═6f58c330-cce4-11eb-3e45-6787cd75141f
# ╠═e850ce3c-f57e-41ab-acb3-958a1b5951a9