hdavid16 / DisjunctiveProgramming.jl

A JuMP extension for Generalized Disjunctive Programming
MIT License
27 stars 3 forks source link

Checking variables in a constraint #53

Closed odow closed 1 year ago

odow commented 1 year ago

I hope this helps. You shouldn't need to do the string parsing stuff

using JuMP

function _variable_list(
    list::Set{Int}, 
    nlp::MOI.Nonlinear.Model, 
    expr::MOI.Nonlinear.Expression,
)
    for node in expr.nodes
        if node.type == MOI.Nonlinear.NODE_MOI_VARIABLE
            push!(list, node.index)
        elseif node.type == MOI.Nonlinear.NODE_SUBEXPRESSION
            _variable_list(list, nlp, nlp.expressions[node.index].nodes)
        end
    end
    return
end

function variable_list(c::NonlinearConstraintRef)
    list = Set{Int}()
    nlp = nonlinear_model(c.model)
    constraint = nlp[index(c)]
    _variable_list(list, nlp, constraint.expression)
    return map(sort!(collect(list))) do i
        return VariableRef(c.model, MOI.VariableIndex(i))
    end
end

variable_list(c::ConstraintRef) = variable_list(constraint_object(c).func)

variable_list(c::AffExpr) = collect(keys(c.terms))

model = Model()
@variable(model, x)
@variable(model, y)
@variable(model, z)
c1 = @NLconstraint(model, sin(x) <= z)
c2 = @constraint(model, x + y <= 0)
variable_list(c1)
variable_list(c2)