From 05b9a68f8ed5e278f086dd4fa9d1b71a84a10850 Mon Sep 17 00:00:00 2001 From: odow Date: Wed, 23 Aug 2023 09:44:13 +1200 Subject: [PATCH 1/3] [Utilities] improve performance of ScalarNonlinearFunction utilities --- src/Utilities/functions.jl | 34 ++++++++++++++++++++++++++++------ 1 file changed, 28 insertions(+), 6 deletions(-) diff --git a/src/Utilities/functions.jl b/src/Utilities/functions.jl index 017b0a7421..8f6b8f0ccf 100644 --- a/src/Utilities/functions.jl +++ b/src/Utilities/functions.jl @@ -252,7 +252,17 @@ function map_indices(index_map::F, ci::MOI.ConstraintIndex) where {F<:Function} end function map_indices(index_map::F, x::AbstractArray) where {F<:Function} - return [map_indices(index_map, xi) for xi in x] + # @odow tried other alternatives here, like + # [map_indices(index_map, xi) for xi in x] + # and + # map(Base.Fix1(map_indices, index_map), x) + # but this was the fastest. The insight is that we know that map_indices + # does not modify the element type of `x`. + y = similar(x) + for i in eachindex(x) + y[i] = map_indices(index_map, x[i]) + end + return y end function map_indices(index_map::F, t::MOI.ScalarAffineTerm) where {F<:Function} @@ -337,11 +347,23 @@ function map_indices( index_map::F, f::MOI.ScalarNonlinearFunction, ) where {F<:Function} - # TODO(odow): this uses recursion. We should remove at some point. - return MOI.ScalarNonlinearFunction( - f.head, - convert(Vector{Any}, map_indices(index_map, f.args)), - ) + root = MOI.ScalarNonlinearFunction(f.head, similar(f.args)) + stack = Tuple{MOI.ScalarNonlinearFunction,Int,Any}[ + (root, i, f.args[i]) for i in length(f.args):-1:1 + ] + while !isempty(stack) + parent, i, arg = pop!(stack) + if arg isa MOI.ScalarNonlinearFunction + child = MOI.ScalarNonlinearFunction(arg.head, similar(arg.args)) + for i in length(arg.args):-1:1 + push!(stack, (child, i, arg.args[i])) + end + parent.args[i] = child + else + parent.args[i] = MOI.Utilities.map_indices(index_map, arg) + end + end + return root end function map_indices( From 558ca44cc72425a753061758764c92f9090fbd1b Mon Sep 17 00:00:00 2001 From: odow Date: Wed, 23 Aug 2023 21:44:07 +1200 Subject: [PATCH 2/3] Tweak and add tests --- src/Utilities/functions.jl | 19 +++++++++---- test/Utilities/functions.jl | 54 ++++++++++++++++++++++++++++++++++++- 2 files changed, 67 insertions(+), 6 deletions(-) diff --git a/src/Utilities/functions.jl b/src/Utilities/functions.jl index 8f6b8f0ccf..03baca4d20 100644 --- a/src/Utilities/functions.jl +++ b/src/Utilities/functions.jl @@ -348,15 +348,24 @@ function map_indices( f::MOI.ScalarNonlinearFunction, ) where {F<:Function} root = MOI.ScalarNonlinearFunction(f.head, similar(f.args)) - stack = Tuple{MOI.ScalarNonlinearFunction,Int,Any}[ - (root, i, f.args[i]) for i in length(f.args):-1:1 - ] + stack = Tuple{MOI.ScalarNonlinearFunction,Int,MOI.ScalarNonlinearFunction}[] + for (i, fi) in enumerate(f.args) + if fi isa MOI.ScalarNonlinearFunction + push!(stack, (root, i, fi)) + else + root.args[i] = MOI.Utilities.map_indices(index_map, fi) + end + end while !isempty(stack) parent, i, arg = pop!(stack) if arg isa MOI.ScalarNonlinearFunction child = MOI.ScalarNonlinearFunction(arg.head, similar(arg.args)) - for i in length(arg.args):-1:1 - push!(stack, (child, i, arg.args[i])) + for (j, argj) in enumerate(arg.args) + if argj isa MOI.ScalarNonlinearFunction + push!(stack, (child, j, argj)) + else + child.args[j] = MOI.Utilities.map_indices(index_map, argj) + end end parent.args[i] = child else diff --git a/test/Utilities/functions.jl b/test/Utilities/functions.jl index b98c090783..75072f28db 100644 --- a/test/Utilities/functions.jl +++ b/test/Utilities/functions.jl @@ -1954,13 +1954,65 @@ function test_ScalarNonlinearFunction_map_indices() x = MOI.add_variable(src) f = MOI.ScalarNonlinearFunction(:log, Any[x]) c = MOI.add_constraint(src, f, MOI.LessThan(1.0)) - dest = MOI.Utilities.Model{Float64}() + dest = MOI.Utilities.MockOptimizer(MOI.Utilities.Model{Float64}()) index_map = MOI.copy_to(dest, src) new_f = MOI.Utilities.map_indices(index_map, f) @test new_f ≈ MOI.ScalarNonlinearFunction(:log, Any[index_map[x]]) return end +function test_ScalarNonlinearFunction_map_indices_nested() + src = MOI.Utilities.Model{Float64}() + x = MOI.add_variables(src, 5) + f = MOI.ScalarNonlinearFunction( + :+, + Any[ + x[1], + MOI.ScalarNonlinearFunction( + :+, + Any[MOI.ScalarNonlinearFunction(:+, Any[x[2], x[3]]), x[4]], + ), + x[5], + ], + ) + c = MOI.add_constraint(src, f, MOI.LessThan(1.0)) + dest = MOI.Utilities.MockOptimizer(MOI.Utilities.Model{Float64}()) + index_map = MOI.copy_to(dest, src) + new_f = MOI.Utilities.map_indices(index_map, f) + y = [index_map[xi] for xi in x] + @test new_f ≈ MOI.ScalarNonlinearFunction( + :+, + Any[ + y[1], + MOI.ScalarNonlinearFunction( + :+, + Any[MOI.ScalarNonlinearFunction(:+, Any[y[2], y[3]]), y[4]], + ), + y[5], + ], + ) + return +end + +function test_ScalarNonlinearFunction_map_indices_deep_recursion() + src = MOI.Utilities.Model{Float64}() + x = MOI.add_variable(src) + f = MOI.ScalarNonlinearFunction(:log, Any[x]) + for _ in 1:50_000 + f = MOI.ScalarNonlinearFunction(:log, Any[f]) + end + c = MOI.add_constraint(src, f, MOI.LessThan(1.0)) + dest = MOI.Utilities.MockOptimizer(MOI.Utilities.Model{Float64}()) + index_map = MOI.copy_to(dest, src) + new_f = MOI.Utilities.map_indices(index_map, f) + g = new_f + for _ in 1:50_000 + g = g.args[1] + end + @test g ≈ MOI.ScalarNonlinearFunction(:log, Any[index_map[x]]) + return +end + function test_ScalarNonlinearFunction_substitute_variables() x = MOI.VariableIndex(1) f = MOI.ScalarNonlinearFunction(:log, Any[1.0*x]) From 3cf419151fc5c36f3ce8340a336c0244515d0be8 Mon Sep 17 00:00:00 2001 From: Oscar Dowson Date: Thu, 24 Aug 2023 07:57:46 +1200 Subject: [PATCH 3/3] Update test/Utilities/functions.jl --- test/Utilities/functions.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/Utilities/functions.jl b/test/Utilities/functions.jl index 75072f28db..3cd785dd6e 100644 --- a/test/Utilities/functions.jl +++ b/test/Utilities/functions.jl @@ -1980,7 +1980,7 @@ function test_ScalarNonlinearFunction_map_indices_nested() index_map = MOI.copy_to(dest, src) new_f = MOI.Utilities.map_indices(index_map, f) y = [index_map[xi] for xi in x] - @test new_f ≈ MOI.ScalarNonlinearFunction( + @test new_f ≈ MOI.ScalarNonlinearFunction( :+, Any[ y[1],