types.jl

# Copyright (c) 2017: Miles Lubin and contributors
# Copyright (c) 2017: Google Inc.
#
# Use of this source code is governed by an MIT-style license that can be found
# in the LICENSE.md file or at https://opensource.org/licenses/MIT.

struct _SubexpressionStorage
    nodes::Vector{Nonlinear.Node}
    adj::SparseArrays.SparseMatrixCSC{Bool,Int}
    sizes::Sizes
    const_values::Vector{Float64}
    forward_storage::Vector{Float64}
    partials_storage::Vector{Float64}
    reverse_storage::Vector{Float64}
    partials_storage_ϵ::Vector{Float64}
    linearity::Linearity

    function _SubexpressionStorage(
        nodes::Vector{Nonlinear.Node},
        adj::SparseArrays.SparseMatrixCSC{Bool,Int},
        const_values::Vector{Float64},
        partials_storage_ϵ::Vector{Float64},
        linearity::Linearity,
    )
        sizes = _infer_sizes(nodes, adj)
        N = _length(sizes)
        return new(
            nodes,
            adj,
            _infer_sizes(nodes, adj),
            const_values,
            zeros(N),  # forward_storage,
            zeros(N),  # partials_storage,
            zeros(N),  # reverse_storage,
            partials_storage_ϵ,
            linearity,
        )
    end
end

# We don't need to store the full vector of `linearity` but we return
# it because it is needed in `compute_hessian_sparsity`.
function _subexpression_and_linearity(
    expr::Nonlinear.Expression,
    moi_index_to_consecutive_index,
    partials_storage_ϵ::Vector{Float64},
    d,
)
    nodes = _replace_moi_variables(expr.nodes, moi_index_to_consecutive_index)
    adj = Nonlinear.adjacency_matrix(nodes)
    linearity = if d.want_hess
        _classify_linearity(nodes, adj, d.subexpression_linearity)
    else
        [NONLINEAR]
    end
    return _SubexpressionStorage(
        nodes,
        adj,
        expr.values,
        partials_storage_ϵ,
        linearity[1],
    ),
    linearity
end

struct _FunctionStorage{R<:SMC.AbstractColoringResult}
    expr::_SubexpressionStorage
    grad_sparsity::Vector{Int}
    # Nonzero pattern of Hessian matrix
    hess_colptr::Vector{Int}
    hess_I::Vector{Int}
    hess_J::Vector{Int}
    rinfo::Union{Nothing,ColoringResult{R}}
    seed_matrix::Matrix{Float64}
    # subexpressions which this function depends on, ordered for forward pass.
    dependent_subexpressions::Vector{Int}

    function _FunctionStorage{R}(
        expr::_SubexpressionStorage,
        num_variables,
        coloring_storage::MOI.Nonlinear.ReverseAD.Coloring.IndexedSet,
        coloring_algorithm::Union{
            Nothing,
            SMC.GreedyColoringAlgorithm,
        },
        subexpressions::Vector{_SubexpressionStorage},
        dependent_subexpressions,
        subexpression_edgelist,
        subexpression_variables,
        linearity::Vector{Linearity},
    ) where {R}
        empty!(coloring_storage)
        _compute_gradient_sparsity!(coloring_storage, expr.nodes)
        for k in dependent_subexpressions
            _compute_gradient_sparsity!(
                coloring_storage,
                subexpressions[k].nodes,
            )
        end
        grad_sparsity = sort!(collect(coloring_storage))
        empty!(coloring_storage)
        if !isnothing(coloring_algorithm)
            edgelist = _compute_hessian_sparsity(
                expr.nodes,
                expr.adj,
                linearity,
                coloring_storage,
                subexpression_edgelist,
                subexpression_variables,
            )
            hess_colptr, hess_I, hess_J, rinfo = _hessian_color_preprocess(
                edgelist,
                num_variables,
                coloring_algorithm,
                coloring_storage,
            )
            seed_matrix = _seed_matrix(rinfo)
            return new{R}(
                expr,
                grad_sparsity,
                hess_colptr,
                hess_I,
                hess_J,
                rinfo,
                seed_matrix,
                dependent_subexpressions,
            )
        else
            return new{R}(
                expr,
                grad_sparsity,
                Int[],
                Int[],
                nothing,
                Array{Float64}(undef, 0, 0),
                dependent_subexpressions,
            )
        end
    end
end

"""
    NLPEvaluator(
        model::Nonlinear.Model,
        ordered_variables::Vector{MOI.VariableIndex},
        coloring_algorithm::SMC.AbstractColoringAlgorithm = SMC.GreedyColoringAlgorithm(; decompression=:substitution),
    )

Return an `NLPEvaluator` object that implements the `MOI.AbstractNLPEvaluator`
interface.

!!! warning
    Before using, you must initialize the evaluator using `MOI.initialize`.
"""
mutable struct NLPEvaluator{
    R,
    C<:SMC.GreedyColoringAlgorithm,
} <: MOI.AbstractNLPEvaluator
    data::Nonlinear.Model
    ordered_variables::Vector{MOI.VariableIndex}
    coloring_algorithm::C

    objective::Union{Nothing,_FunctionStorage{R}}
    constraints::Vector{_FunctionStorage{R}}
    subexpressions::Vector{_SubexpressionStorage}
    subexpression_order::Vector{Int}
    # Storage for the subexpressions in reverse-mode automatic differentiation.
    subexpression_forward_values::Vector{Float64}
    subexpression_reverse_values::Vector{Float64}
    subexpression_linearity::Vector{Linearity}

    # A cache of the last x. This is used to guide whether we need to re-run
    # reverse-mode automatic differentiation.
    last_x::Vector{Float64}

    # Temporary storage for computing Jacobians. This is also used as temporary
    # storage for the input of multivariate functions.
    jac_storage::Vector{Float64}
    # Temporary storage for the gradient of multivariate functions
    user_output_buffer::Vector{Float64}

    # storage for computing hessians
    # these Float64 vectors are reinterpreted to hold multiple epsilon components
    # so the length should be multiplied by the maximum number of epsilon components
    disable_2ndorder::Bool # don't offer Hess or HessVec
    want_hess::Bool
    storage_ϵ::Vector{Float64} # (longest expression including subexpressions)
    input_ϵ::Vector{Float64} # (number of variables)
    output_ϵ::Vector{Float64} # (number of variables)
    subexpression_forward_values_ϵ::Vector{Float64} # (number of subexpressions)
    subexpression_reverse_values_ϵ::Vector{Float64} # (number of subexpressions)
    hessian_sparsity::Vector{Tuple{Int64,Int64}}
    max_chunk::Int # chunk size for which we've allocated storage

    function NLPEvaluator(
        data::Nonlinear.Model,
        ordered_variables::Vector{MOI.VariableIndex},
        coloring_algorithm::SMC.GreedyColoringAlgorithm = SMC.GreedyColoringAlgorithm(;
            decompression = :substitution,
        ),
    )
        problem = SMC.ColoringProblem(;
            structure = :symmetric,
            partition = :column,
        )
        C = typeof(coloring_algorithm)
        R = Base.promote_op(
            SMC.coloring,
            SMC.SparsityPatternCSC{Int},
            typeof(problem),
            C,
        )
        return new{R,C}(data, ordered_variables, coloring_algorithm)
    end
end