fix issue with wrong objective function

jalving · Jan 30, 2024 · 49a4883 · 49a4883
1 parent b8d8205
commit 49a4883
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 9 deletions.
diff --git a/src/optigraph.jl b/src/optigraph.jl
@@ -277,12 +277,12 @@ function _moi_set_objective_function(
     _add_backend_variables(graph_backend(graph), expr)
 
     # update the moi function variable indices
-    _create_graph_moi_func(graph_backend(graph), moi_func, expr)
+    graph_moi_func = _create_graph_moi_func(graph_backend(graph), moi_func, expr)
 
     MOI.set(
         graph_backend(graph),
         MOI.ObjectiveFunction{MOI.ScalarAffineFunction{C}}(),
-        moi_func,
+        graph_moi_func,
     )
     return
 end
@@ -297,12 +297,12 @@ function _moi_set_objective_function(
     _add_backend_variables(graph_backend(graph), expr)
 
     # update the moi function variable indices
-    _create_graph_moi_func(graph_backend(graph), moi_func, expr)
+    graph_moi_func = _create_graph_moi_func(graph_backend(graph), moi_func, expr)
 
     MOI.set(
         graph_backend(graph),
         MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{C}}(),
-        moi_func,
+        graph_moi_func,
     )
     return
 end
@@ -311,7 +311,7 @@ function _moi_set_objective_function(
     graph::OptiGraph, 
     expr::JuMP.GenericNonlinearExpr{NodeVariableRef}
 )
-    moi_func = JuMP.moi_function(expr)
+    graph_moi_func = JuMP.moi_function(expr)
 
     # add variables to backend if using subgraphs
     _add_backend_variables(graph_backend(graph), expr)
@@ -322,7 +322,7 @@ function _moi_set_objective_function(
     MOI.set(
         graph_backend(graph),
         MOI.ObjectiveFunction{MOI.ScalarNonlinearFunction}(),
-        moi_func,
+        graph_moi_func,
     )
     return
 end

diff --git a/src/optinode.jl b/src/optinode.jl
@@ -91,6 +91,28 @@ function JuMP.backend(node::OptiNode)
     return JuMP.backend(graph_backend(node))
 end
 
+# TODO: nonlinear operators on nodes
+function JuMP.add_nonlinear_operator(
+    node::OptiNode,
+    dim::Int,
+    f::Function,
+    args::Vararg{Function,N};
+    name::Symbol = Symbol(f),
+) where {N}
+    nargs = 1 + N
+    if !(1 <= nargs <= 3)
+        error(
+            "Unable to add operator $name: invalid number of functions " *
+            "provided. Got $nargs, but expected 1 (if function only), 2 (if " *
+            "function and gradient), or 3 (if function, gradient, and " *
+            "hesssian provided)",
+        )
+    end
+    name = Symbol(node.label, ".", name)
+    MOI.set(graph_backend(node), MOI.UserDefinedFunction(name, dim), tuple(f, args...))
+    return JuMP.NonlinearOperator(f, name)
+end
+
 function _set_dirty(node::OptiNode)
     for graph in containing_optigraphs(node)
         graph.is_model_dirty = true
@@ -121,7 +143,7 @@ function MOI.get(
     return con_inds
 end
 
-# TODO: store objective functions on nodes and query node attributes
+# TODO: store objective functions on nodes and query as node attributes
 # function MOI.get(node::OptiNode, attr::MOI.AnyAttribute)
 #     return MOI.get(graph_backend(node), attr)
 # end
@@ -193,8 +215,9 @@ function MOI.set(
     args...
 )
     for graph in containing_optigraphs(JuMP.owner_model(cref))
-        graph_index = graph_backend(graph).element_to_graph_map[cref]
-        MOI.set(graph_backend(graph), attr, graph_index, args...)
+        gb = graph_backend(graph)
+        graph_index = gb.element_to_graph_map[cref]
+        MOI.set(gb, attr, graph_index, args...)
     end
     return
 end