Register NLparameters and pretty print them (#2516)

jump-dev · Mar 7, 2021 · 3ba8d43 · 3ba8d43
1 parent b688373
commit 3ba8d43
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 5 deletions.
diff --git a/src/nlp.jl b/src/nlp.jl
@@ -1377,14 +1377,21 @@ function Base.show(io::IO, v::_VariablePrintWrapper)
     print(io, function_string(v.mode, v.v))
 end
 mutable struct _ParameterPrintWrapper
+    model::Model
     idx::Int
     mode::Any
 end
 function Base.show(io::IO, p::_ParameterPrintWrapper)
-    if p.mode == IJuliaMode
-        print(io, "parameter_{$(p.idx)}")
+    relevant_parameters = filter(i->i[2] isa NonlinearParameter && i[2].index==p.idx, p.model.obj_dict)
+    if length(relevant_parameters) == 1
+        par_name = first(relevant_parameters)[1]
+        print(io, par_name)
     else
-        print(io, "parameter[$(p.idx)]")
+        if p.mode == IJuliaMode
+            print(io, "parameter_{$(p.idx)}")
+        else
+            print(io, "parameter[$(p.idx)]")
+        end
     end
 end
 mutable struct _SubexpressionPrintWrapper
@@ -1438,7 +1445,7 @@ function _tape_to_expr(
         if splat_subexpressions
             return parameter_values[nod.index]
         else
-            return _ParameterPrintWrapper(nod.index, print_mode)
+            return _ParameterPrintWrapper(m, nod.index, print_mode)
         end
     elseif nod.nodetype == CALL
         op = nod.index

diff --git a/src/print.jl b/src/print.jl
@@ -754,7 +754,13 @@ function function_string(::Type{<:PrintMode}, p::NonlinearExpression)
 end
 
 function function_string(::Type{<:PrintMode}, p::NonlinearParameter)
-    return "Reference to nonlinear parameter #$(p.index)"
+    relevant_parameters = filter(i->i[2] isa NonlinearParameter && i[2].index==p.index, p.m.obj_dict)
+    if length(relevant_parameters) == 1
+        par_name = first(relevant_parameters)[1]
+        return "Reference to nonlinear parameter $(par_name)"
+    else
+        return "Reference to nonlinear parameter #$(p.index)"
+    end
 end
 
 function Base.show(io::IO, ex::Union{NonlinearExpression,NonlinearParameter})

diff --git a/test/print.jl b/test/print.jl
@@ -180,6 +180,12 @@ end
         io_test(REPLMode, param, "\"Reference to nonlinear parameter #1\"")
     end
 
+    @testset "Registered nonlinear parameters" begin
+        model = Model()
+        model[:param] = @NLparameter(model, param == 1.0)
+        io_test(REPLMode, param, "\"Reference to nonlinear parameter param\"")
+    end
+
     @testset "NLPEvaluator" begin
         model = Model()
         evaluator = JuMP.NLPEvaluator(model)
@@ -256,6 +262,27 @@ end
         )
     end
 
+    @testset "Nonlinear constraints with embedded registered parameters/expressions" begin
+        le = JuMP._math_symbol(REPLMode, :leq)
+
+        model = Model()
+        @variable(model, x)
+        expr = @NLexpression(model, x + 1)
+        model[:param] = @NLparameter(model, param == 1.0)
+
+        constr = @NLconstraint(model, expr - param <= 0)
+        io_test(
+            REPLMode,
+            constr,
+            "(subexpression[1] - param) - 0.0 $le 0",
+        )
+        io_test(
+            IJuliaMode,
+            constr,
+            "(subexpression_{1} - param) - 0.0 \\leq 0",
+        )
+    end
+
     @testset "Custom constraint" begin
         model = Model()
         function test_constraint(function_str, in_set_str, name)