diff --git a/src/Nonlinear/SymbolicAD/SymbolicAD.jl b/src/Nonlinear/SymbolicAD/SymbolicAD.jl index 8a94707400..169fed6218 100644 --- a/src/Nonlinear/SymbolicAD/SymbolicAD.jl +++ b/src/Nonlinear/SymbolicAD/SymbolicAD.jl @@ -77,16 +77,19 @@ function simplify!(f::MOI.ScalarAffineFunction{T}) where {T} if isempty(f.terms) return f.constant end + if iszero(f.constant) && length(f.terms) == 1 + term = only(f.terms) + if isone(term.coefficient) + return term.variable + end + end return f end function simplify!(f::MOI.ScalarQuadraticFunction{T}) where {T} f = MOI.Utilities.canonicalize!(f) if isempty(f.quadratic_terms) - if isempty(f.affine_terms) - return f.constant - end - return MOI.ScalarAffineFunction(f.affine_terms, f.constant) + return simplify!(MOI.ScalarAffineFunction(f.affine_terms, f.constant)) end return f end @@ -117,7 +120,7 @@ function simplify!(f::MOI.ScalarNonlinearFunction) push!(result_stack, arg) end end - return _simplify_if_affine!(only(result_stack)) + return _simplify_if_quadratic!(only(result_stack)) end function simplify!(f::MOI.VectorAffineFunction{T}) where {T} @@ -140,10 +143,12 @@ function simplify!(f::MOI.VectorQuadraticFunction{T}) where {T} end function simplify!(f::MOI.VectorNonlinearFunction) - for (i, row) in enumerate(f.rows) - f.rows[i] = simplify!(row) + rows = simplify!.(f.rows) + Y = reduce(promote_type, typeof.(rows)) + if isconcretetype(Y) + return MOI.Utilities.vectorize(convert(Vector{Y}, rows)) end - return f + return MOI.VectorNonlinearFunction(rows) end # If a ScalarNonlinearFunction has only constant arguments, we should return @@ -1507,22 +1512,28 @@ function MOI.eval_hessian_lagrangian(model::Evaluator, H, x, σ, μ) end # A default fallback for all types -_add_to_affine!(::Any, ::Any, ::T) where {T} = nothing +_add_to_quadratic!(::Any, ::Real, ::Any) = nothing +_add_to_quadratic!(::Any, ::Real, ::Any, ::Any) = nothing -# The creation of `ret::MOI.ScalarAffineFunction` has been delayed until now. -function _add_to_affine!( - ::Nothing, - f::Union{Real,MOI.VariableIndex,MOI.ScalarAffineFunction}, +# The creation of `ret::MOI.ScalarQuadraticFunction` has been delayed until now. +function _add_to_quadratic!( + ::Missing, scale::T, -) where {T} - return _add_to_affine!(zero(MOI.ScalarAffineFunction{T}), f, scale) -end - -function _add_to_affine!( - ret::MOI.ScalarAffineFunction{T}, - x::S, + f::Union{ + Real, + MOI.VariableIndex, + MOI.ScalarAffineFunction, + MOI.ScalarQuadraticFunction, + }..., +) where {T<:Real} + return _add_to_quadratic!(zero(MOI.ScalarQuadraticFunction{T}), scale, f...) +end + +function _add_to_quadratic!( + ret::MOI.ScalarQuadraticFunction{T}, scale::T, -) where {T,S<:Real} + x::S, +) where {T<:Real,S<:Real} if promote_type(T, S) != T return # We can't store `S` in `T`. end @@ -1530,38 +1541,122 @@ function _add_to_affine!( return ret end -function _add_to_affine!( - ret::MOI.ScalarAffineFunction{T}, - x::MOI.VariableIndex, +function _add_to_quadratic!( + ret::MOI.ScalarQuadraticFunction{T}, scale::T, -) where {T} - push!(ret.terms, MOI.ScalarAffineTerm(scale, x)) + f::MOI.ScalarAffineTerm{S}, +) where {T<:Real,S} + @assert promote_type(T, S) == T + push!( + ret.affine_terms, + MOI.ScalarAffineTerm{T}(scale * f.coefficient, f.variable), + ) return ret end -function _add_to_affine!( - ret::MOI.ScalarAffineFunction{T}, - f::MOI.ScalarAffineFunction{S}, +function _add_to_quadratic!( + ret::MOI.ScalarQuadraticFunction{T}, scale::T, -) where {T,S} + f::MOI.ScalarQuadraticTerm{S}, +) where {T<:Real,S} + @assert promote_type(T, S) == T + push!( + ret.quadratic_terms, + MOI.ScalarQuadraticTerm{T}( + scale * f.coefficient, + f.variable_1, + f.variable_2, + ), + ) + return ret +end + +function _add_to_quadratic!( + ret::MOI.ScalarQuadraticFunction{T}, + scale::T, + x::MOI.VariableIndex, +) where {T<:Real} + return _add_to_quadratic!(ret, scale, MOI.ScalarAffineTerm(one(T), x)) +end + +function _add_to_quadratic!( + ret::MOI.ScalarQuadraticFunction{T}, + scale::T, + f::MOI.ScalarAffineFunction{S}, +) where {T<:Real,S} if promote_type(T, S) != T return # We can't store `S` in `T`. end - ret = _add_to_affine!(ret, f.constant, scale) + ret = _add_to_quadratic!(ret, scale, f.constant) for term in f.terms - ret = _add_to_affine!(ret, term.variable, scale * term.coefficient) + ret = _add_to_quadratic!(ret, scale, term) end return ret end -function _add_to_affine!( - ret::Union{Nothing,MOI.ScalarAffineFunction{T}}, - f::MOI.ScalarNonlinearFunction, +function _add_to_quadratic!( + ret::MOI.ScalarQuadraticFunction{T}, scale::T, -) where {T} + f::MOI.ScalarQuadraticFunction{S}, +) where {T<:Real,S} + if promote_type(T, S) != T + return # We can't store `S` in `T`. + end + ret = _add_to_quadratic!(ret, scale, f.constant) + for term in f.affine_terms + ret = _add_to_quadratic!(ret, scale, term) + end + for q_term in f.quadratic_terms + ret = _add_to_quadratic!(ret, scale, q_term) + end + return ret +end + +function _add_to_quadratic!( + ret::MOI.ScalarQuadraticFunction{T}, + scale::T, + f::MOI.VariableIndex, + g::MOI.VariableIndex, +) where {T<:Real} + return _add_to_quadratic!(ret, scale, one(T) * f * g) +end + +function _add_to_quadratic!( + ret::MOI.ScalarQuadraticFunction{T}, + scale::T, + f::MOI.ScalarAffineFunction{F}, + g::MOI.ScalarAffineFunction{G}, +) where {T<:Real,F,G} + H = MOI.ScalarAffineFunction{promote_type(F, G)} + return _add_to_quadratic!(ret, scale, convert(H, f) * convert(H, g)) +end + +function _add_to_quadratic!( + ret::MOI.ScalarQuadraticFunction{T}, + scale::T, + f::MOI.VariableIndex, + g::MOI.ScalarAffineFunction, +) where {T<:Real} + return _add_to_quadratic!(ret, scale, f * g) +end + +function _add_to_quadratic!( + ret::MOI.ScalarQuadraticFunction{T}, + scale::T, + f::MOI.ScalarAffineFunction, + g::MOI.VariableIndex, +) where {T<:Real} + return _add_to_quadratic!(ret, scale, g, f) +end + +function _add_to_quadratic!( + ret::Union{Missing,MOI.ScalarQuadraticFunction{T}}, + scale::T, + f::MOI.ScalarNonlinearFunction, +) where {T<:Real} if f.head == :+ for arg in f.args - ret = _add_to_affine!(ret, arg, scale) + ret = _add_to_quadratic!(ret, scale, arg) if ret === nothing return end @@ -1569,38 +1664,51 @@ function _add_to_affine!( return ret elseif f.head == :- if length(f.args) == 1 - return _add_to_affine!(ret, only(f.args), -scale) + return _add_to_quadratic!(ret, -scale, only(f.args)) end @assert length(f.args) == 2 - ret = _add_to_affine!(ret, f.args[1], scale) + ret = _add_to_quadratic!(ret, scale, f.args[1]) if ret === nothing return end - return _add_to_affine!(ret, f.args[2], -scale) + return _add_to_quadratic!(ret, -scale, f.args[2]) elseif f.head == :* - y = nothing + y1, y2 = nothing, nothing for arg in f.args if arg isa Real scale *= arg - elseif y === nothing - y = arg + elseif y1 === nothing + y1 = arg + elseif y2 === nothing + y2 = arg else return # We already have a `y`. Can't multiple factors. end end - return _add_to_affine!(ret, something(y, one(T)), convert(T, scale)) + if y1 === nothing + @assert y2 === nothing + return _add_to_quadratic!(ret, one(T), scale) + elseif y2 === nothing + return _add_to_quadratic!(ret, scale, y1) + else + return _add_to_quadratic!(ret, scale, y1, y2) + end + elseif f.head == :^ && f.args[2] isa Real && f.args[2] == 2 + return _add_to_quadratic!(ret, scale, f.args[1], f.args[1]) + elseif f.head == :/ && f.args[2] isa Real + return _add_to_quadratic!(ret, convert(T, scale / f.args[2]), f.args[1]) end return # An unsupported f.head end -function _simplify_if_affine!(f::MOI.ScalarNonlinearFunction) - ret = _add_to_affine!(nothing, f, 1.0) +function _simplify_if_quadratic!(f::MOI.ScalarNonlinearFunction) + ret = _add_to_quadratic!(missing, 1.0, f) if ret === nothing return f end - return simplify!(ret::MOI.ScalarAffineFunction{Float64}) + return simplify!(ret::MOI.ScalarQuadraticFunction{Float64}) end -_simplify_if_affine!(f::Any) = f +_simplify_if_quadratic!(f::Any) = f end # module diff --git a/test/Nonlinear/SymbolicAD.jl b/test/Nonlinear/SymbolicAD.jl index d94b941623..9b53473982 100644 --- a/test/Nonlinear/SymbolicAD.jl +++ b/test/Nonlinear/SymbolicAD.jl @@ -46,8 +46,8 @@ function test_derivative() # derivative(f::MOI.ScalarQuadraticFunction{T}, x::MOI.VariableIndex) convert(MOI.ScalarQuadraticFunction{Float64}, 1.0 * x)=>1.0, convert(MOI.ScalarQuadraticFunction{Float64}, 1.0 * x + 0.0 * y)=>1.0, - 1.0*x*y=>1.0*y, - 1.0*y*x=>1.0*y, + 1.0*x*y=>y, + 1.0*y*x=>y, 1.0*x*x=>2.0*x, 1.0*x*x+3.0*x+4.0=>2.0*x+3.0, 2.0*x*x-2.0*x+1.0=>4.0*x-2.0, @@ -76,9 +76,9 @@ function test_derivative() # f.head == :- op(:-, sin_x, cos_x)=>op(:+, cos_x, sin_x), # f.head == :* - op(:*, x, y, z)=>op(:*, y, z), - op(:*, y, x, z)=>op(:*, y, z), - op(:*, y, z, x)=>op(:*, y, z), + op(:*, x, y, z)=>1.0*y*z, + op(:*, y, x, z)=>1.0*y*z, + op(:*, y, z, x)=>1.0*y*z, # :^ op(:^, sin_x, 2)=>op(:*, 2.0, sin_x, cos_x), op(:^, sin_x, 1)=>cos_x, @@ -91,6 +91,7 @@ function test_derivative() op(:*, x, op(:^, x, x - 1)), op(:*, op(:^, x, x), op(:log, x)), ), + op(:^, x, 3)=>3.0*x*x, # :/ op(:/, x, 2)=>0.5, op( @@ -199,22 +200,21 @@ function test_simplify_ScalarAffineFunction() x = MOI.VariableIndex(1) @test SymbolicAD.simplify(1.0 * x + 1.0) ≈ 1.0 * x + 1.0 @test SymbolicAD.simplify(1.0 * x + 2.0 * x + 1.0) ≈ 3.0 * x + 1.0 + @test SymbolicAD.simplify(1.0 * x) == x + @test SymbolicAD.simplify(2.0 * x) ≈ 2.0 * x return end function test_simplify_ScalarQuadraticFunction() x = MOI.VariableIndex(1) - f = MOI.ScalarQuadraticFunction( - MOI.ScalarQuadraticTerm{Float64}[], - [MOI.ScalarAffineTerm{Float64}(1.0, x)], - 1.0, - ) + terms = MOI.ScalarQuadraticTerm{Float64}[] + f = MOI.ScalarQuadraticFunction(terms, [MOI.ScalarAffineTerm(1.0, x)], 1.0) @test SymbolicAD.simplify(f) ≈ 1.0 * x + 1.0 - f = MOI.ScalarQuadraticFunction( - MOI.ScalarQuadraticTerm{Float64}[], - MOI.ScalarAffineTerm{Float64}[], - 2.0, - ) + f = MOI.ScalarQuadraticFunction(terms, [MOI.ScalarAffineTerm(2.0, x)], 0.0) + @test SymbolicAD.simplify(f) ≈ 2.0 * x + 0.0 + f = MOI.ScalarQuadraticFunction(terms, [MOI.ScalarAffineTerm(1.0, x)], 0.0) + @test SymbolicAD.simplify(f) == x + f = MOI.ScalarQuadraticFunction(terms, MOI.ScalarAffineTerm{Float64}[], 2.0) @test SymbolicAD.simplify(f) === 2.0 @test SymbolicAD.simplify(1.0 * x * x + 1.0) ≈ 1.0 * x * x + 1.0 g = 1.0 * x * x + 2.0 * x * x + 1.0 @@ -401,7 +401,15 @@ function test_simplify_VectorNonlinearFunction() return op(:+, op(:-, f, 0.0), 0.0) end f = MOI.VectorNonlinearFunction(wrap.([y; x_plus])) - g = MOI.VectorNonlinearFunction([y; x_plus]) + g = MOI.Utilities.vectorize([sum(1.0 .* x .* x); x]) + @test SymbolicAD.simplify(f) ≈ g + return +end + +function test_simplify_VectorNonlinearFunction_abstract() + x = MOI.VariableIndex(1) + f = MOI.VectorNonlinearFunction([x, 1.0 * x, op(:sin, x)]) + g = MOI.VectorNonlinearFunction([x, x, op(:sin, x)]) @test SymbolicAD.simplify(f) ≈ g return end @@ -703,16 +711,18 @@ function test_SymbolicAD_univariate_registered() return end -function test_simplify_if_affine() +function test_simplify_if_quadratic() x = MOI.VariableIndex(1) for (f, ret) in Any[ + # Constants op(:*, 2)=>2, op(:*, 2 // 3)=>2/3, op(:*, 2, 3)=>6, - op(:*, 2, x, 3)=>6*x, op(:+, 2, 3)=>5, op(:-, 2)=>-2, op(:-, 2, 3)=>-1, + # Affine + op(:*, 2, x, 3)=>6*x, op(:-, x)=>-1*x, op(:-, x, 2)=>1.0*x-2.0, op(:-, 2, x)=>2.0+-1.0*x, @@ -723,19 +733,37 @@ function test_simplify_if_affine() op(:+, 1.0 * x, 1.0 * x + 2.0)=>2.0*x+2.0, op(:-, 1.0 * x + 2.0)=>-1.0*x-2.0, op(:*, 2, 1.0 * x + 2.0, 3)=>6.0*x+12.0, + # Quadratic + op(:+, 1.0 * x * x)=>1.0*x*x, + op(:*, 2, x, 3, x)=>6.0*x*x, + op(:*, 2, x, 3, 1.0 * x)=>6.0*x*x, + op(:*, 2, x, 3, 1.0 * x + 2.0)=>6.0*x*x+12.0*x, + op(:*, 2, 1.0 * x + 2.0, 3, x)=>6.0*x*x+12.0*x, + op(:*, 1.0 * x + 2.0, 1.0 * x)=>1.0*x*x+2.0*x, + op(:*, 1.0 * x + 2.0, 1.0 * x + 1.0)=>1.0*x*x+3.0*x+2.0, + # Power + op(:^, x, 2)=>1.0*x*x, + op(:^, 1.0 * x, 2)=>1.0*x*x, + op(:^, 1.0 * x + 1.0, 2)=>1.0*x*x+2.0*x+1.0, + op(:^, -2.0 * x - 2.0, 2)=>4.0*x*x+8.0*x+4.0, + # Divide + op(:/, x, 2)=>0.5*x, + op(:/, 1.0 * x, 2)=>0.5*x, + op(:/, 2.0 * x, 2)=>x, + op(:/, 2.0 * x * x, 2)=>1.0*x*x, # Early termination because not affine op(:+, op(:sin, x))=>nothing, op(:-, op(:sin, x))=>nothing, op(:-, op(:sin, x), 1)=>nothing, op(:-, x, op(:sin, x))=>nothing, - op(:*, 2, x, 3, x)=>nothing, op(:*, 2, 3, op(:sin, x))=>nothing, op(:log, x)=>nothing, + op(:+, big(1) * x * x, big(2))=>nothing, op(:+, big(1) * x, big(2))=>nothing, op(:+, x, big(2))=>nothing, - op(:+, 1.0 * x * x)=>nothing, + op(:+, x, 1 + 2im)=>nothing, ] - @test SymbolicAD._simplify_if_affine!(f) ≈ something(ret, f) + @test SymbolicAD._simplify_if_quadratic!(f) ≈ something(ret, f) end return end