diff --git a/src/plans/augmented_lagrangian_plan.jl b/src/plans/augmented_lagrangian_plan.jl index e1ccd235c7..66deba6d17 100644 --- a/src/plans/augmented_lagrangian_plan.jl +++ b/src/plans/augmented_lagrangian_plan.jl @@ -44,8 +44,8 @@ function set_manopt_parameter!(alc::AugmentedLagrangianCost, ::Val{:λ}, λ) return alc end function (L::AugmentedLagrangianCost)(M::AbstractManifold, p) - gp = get_inequality_constraints(M, L.co, p) - hp = get_equality_constraints(M, L.co, p) + gp = get_inequality_constraint(M, L.co, p, :) + hp = get_equality_constraint(M, L.co, p, :) m = length(gp) n = length(hp) c = get_cost(M, L.co, p) @@ -124,7 +124,9 @@ function (LG::AugmentedLagrangianGrad)( end end if n > 0 - X .+= sum((hp .* LG.ρ .+ LG.λ) .* get_grad_equality_constraints(M, LG.co, p, range)) + X .+= sum( + (hp .* LG.ρ .+ LG.λ) .* get_grad_equality_constraint(M, LG.co, p, :, range) + ) end return X end diff --git a/src/plans/count.jl b/src/plans/count.jl index 95f69d0a54..27bb48941e 100644 --- a/src/plans/count.jl +++ b/src/plans/count.jl @@ -307,7 +307,7 @@ end # Constraint function get_constraints(M::AbstractManifold, co::ManifoldCountObjective, p) _count_if_exists(co, :Constraints) - return get_constraint(M, co.objective, p) + return get_constraints(M, co.objective, p) end function get_equality_constraint( M::AbstractManifold, co::ManifoldCountObjective, p, c::Colon @@ -340,7 +340,7 @@ function get_inequality_constraint( M::AbstractManifold, co::ManifoldCountObjective, p, i::Colon ) _count_if_exists(co, :InequalityConstraints) - return get_inequality_constraints(M, co.objective, p) + return get_inequality_constraint(M, co.objective, p, i) end function get_inequality_constraint( M::AbstractManifold, co::ManifoldCountObjective, p, i::Integer @@ -390,7 +390,7 @@ function get_grad_inequality_constraint( M::AbstractManifold, co::ManifoldCountObjective, p, i::Colon ) _count_if_exists(co, :GradInequalityConstraints) - return get_grad_inequality_constraints(M, co.objective, p, i) + return get_grad_inequality_constraint(M, co.objective, p, i) end function get_grad_inequality_constraint!( M::AbstractManifold, X, co::ManifoldCountObjective, p, i::Colon diff --git a/src/plans/exact_penalty_method_plan.jl b/src/plans/exact_penalty_method_plan.jl index c0c1852413..fad94e1a74 100644 --- a/src/plans/exact_penalty_method_plan.jl +++ b/src/plans/exact_penalty_method_plan.jl @@ -71,8 +71,8 @@ function set_manopt_parameter!(epc::ExactPenaltyCost, ::Val{:u}, u) return epc end function (L::ExactPenaltyCost{<:LogarithmicSumOfExponentials})(M::AbstractManifold, p) - gp = get_inequality_constraints(M, L.co, p) - hp = get_equality_constraints(M, L.co, p) + gp = get_inequality_constraint(M, L.co, p, :) + hp = get_equality_constraint(M, L.co, p, :) m = length(gp) n = length(hp) cost_ineq = (m > 0) ? sum(L.u .* log.(1 .+ exp.(gp ./ L.u))) : 0.0 @@ -80,8 +80,8 @@ function (L::ExactPenaltyCost{<:LogarithmicSumOfExponentials})(M::AbstractManifo return get_cost(M, L.co, p) + (L.ρ) * (cost_ineq + cost_eq) end function (L::ExactPenaltyCost{<:LinearQuadraticHuber})(M::AbstractManifold, p) - gp = get_inequality_constraints(M, L.co, p) - hp = get_equality_constraints(M, L.co, p) + gp = get_inequality_constraint(M, L.co, p, :) + hp = get_equality_constraint(M, L.co, p, :) m = length(gp) n = length(hp) cost_eq_greater_u = (m > 0) ? sum((gp .- L.u / 2) .* (gp .> L.u)) : 0.0 @@ -135,8 +135,8 @@ function (EG::ExactPenaltyGrad)(M::AbstractManifold, p) return EG(M, X, p) end function (EG::ExactPenaltyGrad{<:LogarithmicSumOfExponentials})(M::AbstractManifold, X, p) - gp = get_inequality_constraints(M, EG.co, p) - hp = get_equality_constraints(M, EG.co, p) + gp = get_inequality_constraint(M, EG.co, p, :) + hp = get_equality_constraint(M, EG.co, p, :) m = length(gp) n = length(hp) # start with `gradf` @@ -144,14 +144,15 @@ function (EG::ExactPenaltyGrad{<:LogarithmicSumOfExponentials})(M::AbstractManif c = 0 # add gradient of the components of g (m > 0) && (c = EG.ρ .* exp.(gp ./ EG.u) ./ (1 .+ exp.(gp ./ EG.u))) - (m > 0) && (X .+= sum(get_grad_inequality_constraints(M, EG.co, p) .* c)) + (m > 0) && (X .+= sum(get_grad_inequality_constraint(M, EG.co, p, :) .* c)) # add gradient of the components of h + # TODO: This could be reduced and only active gradients be evaluated (n > 0) && ( c = EG.ρ .* (exp.(hp ./ EG.u) .- exp.(-hp ./ EG.u)) ./ (exp.(hp ./ EG.u) .+ exp.(-hp ./ EG.u)) ) - (n > 0) && (X .+= sum(get_grad_equality_constraints(M, EG.co, p) .* c)) + (n > 0) && (X .+= sum(get_grad_equality_constraint(M, EG.co, p, :) .* c)) return X end @@ -159,83 +160,20 @@ end function (EG::ExactPenaltyGrad{<:LinearQuadraticHuber})( M::AbstractManifold, X, p::P ) where {P} - gp = get_inequality_constraints(M, EG.co, p) - hp = get_equality_constraints(M, EG.co, p) + gp = get_inequality_constraint(M, EG.co, p, :) + hp = get_equality_constraint(M, EG.co, p, :) m = length(gp) n = length(hp) get_gradient!(M, X, EG.co, p) if m > 0 - gradgp = get_grad_inequality_constraints(M, EG.co, p) + gradgp = get_grad_inequality_constraint(M, EG.co, p, :) X .+= sum(gradgp .* (gp .>= EG.u) .* EG.ρ) # add the ones >= u X .+= sum(gradgp .* (gp ./ EG.u .* (0 .<= gp .< EG.u)) .* EG.ρ) # add < u end if n > 0 + # TODO: This could be reduced and only active gradients be evaluated c = (hp ./ sqrt.(hp .^ 2 .+ EG.u^2)) .* EG.ρ - X .+= sum(get_grad_equality_constraints(M, EG.co, p) .* c) + X .+= sum(get_grad_equality_constraint(M, EG.co, p, :) .* c) end return X end -# Variant 2: vectors of allocating gradients -#= TODO: rework to easier access with substructures -function ( - EG::ExactPenaltyGrad{ - <:LinearQuadraticHuber, - <:ConstrainedManifoldObjective{AllocatingEvaluation,<:VectorConstraint}, - } -)( - M::AbstractManifold, X, p::P -) where {P} - m = length(EG.co.g) - n = length(EG.co.h) - get_gradient!(M, X, EG.co, p) - for i in 1:m - gpi = get_inequality_constraint(M, EG.co, p, i) - if (gpi >= 0) # these are the only necessary allocations. - (gpi .>= EG.u) && (X .+= gpi .* EG.ρ) - (0 < gpi < EG.u) && (X .+= gpi .* (gpi / EG.u) * EG.ρ) - end - end - for j in 1:n - hpj = get_equality_constraint(M, EG.co, p, j) - if hpj > 0 - c = hpj / sqrt(hpj^2 + EG.u^2) - X .+= get_grad_equality_constraint(M, EG.co, p, j) .* (c * EG.ρ) - end - end - return X -end - -# Variant 3: vectors of mutating gradients -function ( - EG::ExactPenaltyGrad{ - <:LinearQuadraticHuber, - <:ConstrainedManifoldObjective{InplaceEvaluation,<:VectorConstraint}, - } -)( - M::AbstractManifold, X, p::P -) where {P} - m = length(EG.co.g) - n = length(EG.co.h) - get_gradient!(M, X, EG.co, p) - Y = zero_vector(M, p) - for i in 1:m - gpi = get_inequality_constraint(M, EG.co, p, i) - if (gpi >= 0) # the cases where to evaluate the gradient - # only evaluate the gradient if `gpi > 0` - get_grad_inequality_constraint!(M, Y, EG.co, p, i) - # just add the gradient scaled by ρ - (gpi >= EG.u) && (X .+= EG.ρ .* Y) - # use a different factor, but exclude the case `g = 0` as well - (0 < gpi < EG.u) && (X .+= ((gpi / EG.u) * EG.ρ) .* Y) - end - end - for j in 1:n - hpj = get_equality_constraint(M, EG.co, p, j) - if hpj > 0 - get_grad_equality_constraint!(M, Y, EG.co, p, j) - X .+= ((hpj / sqrt(hpj^2 + EG.u^2)) * EG.ρ) .* Y - end - end - return X -end -=# diff --git a/src/plans/vectorial_plan.jl b/src/plans/vectorial_plan.jl index 40466ca4ac..136b2954f1 100644 --- a/src/plans/vectorial_plan.jl +++ b/src/plans/vectorial_plan.jl @@ -153,11 +153,11 @@ Since `i` is assumed to be a linear index, you can provide get_cost(M::AbstractManifold, vgf::VectorGradientFunction, p, i, range=nothing) function get_cost( M::AbstractManifold, - vgf::VectorGradientFunction{<:AllocatingEvaluation,<:FunctionVectorialType}, + vgf::VectorGradientFunction{E,<:FunctionVectorialType}, p, i, range=nothing, -) +) where {E} c = vgf.costs!!(M, p) if isa(c, Number) return c @@ -167,20 +167,20 @@ function get_cost( end function get_cost( M::AbstractManifold, - vgf::VectorGradientFunction{<:AllocatingEvaluation,<:ComponentVectorialType}, + vgf::VectorGradientFunction{E,<:ComponentVectorialType}, p, i::Integer, range=nothing, -) +) where {E} return vgf.costs!![i](M, p) end function get_cost( M::AbstractManifold, - vgf::VectorGradientFunction{<:AllocatingEvaluation,<:ComponentVectorialType}, + vgf::VectorGradientFunction{E,<:ComponentVectorialType}, p, i, range=nothing, -) +) where {E} return [f(M, p) for f in vgf.costs!![i]] end @@ -351,6 +351,18 @@ function get_gradient!( copyto!(mP, X, vgf.jacobian!!(M, p)[mP, i]) return X end +function get_gradient!( + M::AbstractManifold, + X, + vgf::VectorGradientFunction{<:AllocatingEvaluation,FT,<:FunctionVectorialType}, + p, + i::Integer, + range::Union{AbstractPowerRepresentation,Nothing}=NestedPowerRepresentation(), +) where {FT} + mP = PowerManifold(M, range, vgf.range_dimension) + copyto!(mP, X, vgf.jacobian!!(M, p)[mP, i]) + return X +end # # # Part II: In-place evaluations @@ -419,7 +431,7 @@ function get_gradient!( ) where {FT} n = _vgf_index_to_length(i, vgf.range_dimension) pM = PowerManifold(M, range, n) - for (j, f) in zip(j, vgf.jacobian!![i]) + for (j, f) in zip(i, vgf.jacobian!![i]) f(M, X[pM, j], p) end return X @@ -437,7 +449,7 @@ function get_gradient!( P = fill(p, pM) x = zero_vector(pM, P) vgf.jacobian!!(M, x, p) - copyto!(M, X, p, x[mP, i]) + copyto!(M, X, p, x[pM, i]) return X end function get_gradient!( diff --git a/src/solvers/augmented_Lagrangian_method.jl b/src/solvers/augmented_Lagrangian_method.jl index c5b811bc5b..694dcb0473 100644 --- a/src/solvers/augmented_Lagrangian_method.jl +++ b/src/solvers/augmented_Lagrangian_method.jl @@ -76,8 +76,8 @@ mutable struct AugmentedLagrangianMethodState{ λ_max::R=20.0, λ_min::R=-λ_max, μ_max::R=20.0, - μ::V=ones(length(get_inequality_constraints(M, co, p))), - λ::V=ones(length(get_equality_constraints(M, co, p))), + μ::V=ones(length(get_inequality_constraint(M, co, p, :))), + λ::V=ones(length(get_equality_constraint(M, co, p, :))), ρ::R=1.0, τ::R=0.8, θ_ρ::R=0.3, @@ -391,9 +391,9 @@ function augmented_Lagrangian_method!( ϵ_min::Real=1e-6, ϵ_exponent::Real=1 / 100, θ_ϵ::Real=(ϵ_min / ϵ)^(ϵ_exponent), - μ::Vector=ones(length(get_inequality_constraints(M, cmo, p))), + μ::Vector=ones(length(get_inequality_constraint(M, cmo, p, :))), μ_max::Real=20.0, - λ::Vector=ones(length(get_equality_constraints(M, cmo, p))), + λ::Vector=ones(length(get_equality_constraint(M, cmo, p, :))), λ_max::Real=20.0, λ_min::Real=-λ_max, τ::Real=0.8, diff --git a/test/plans/test_constrained_plan.jl b/test/plans/test_constrained_plan.jl index a3381f0ea2..1120dc8b9c 100644 --- a/test/plans/test_constrained_plan.jl +++ b/test/plans/test_constrained_plan.jl @@ -91,49 +91,49 @@ include("../utils/dummy_types.jl") ) co1f = ConstrainedManifoldObjective(f, grad_f!; g=g, grad_g=grad_g) @test get_constraints(M, co1f, p) == [c[1], []] - @test get_grad_equality_constraints(M, co1f, p) == [] - @test get_grad_inequality_constraints(M, co1f, p) == gg + @test get_grad_equality_constraint(M, co1f, p, :) == [] + @test get_grad_inequality_constraint(M, co1f, p, :) == gg co1v = ConstrainedManifoldObjective( f, grad_f!; g=[g1, g2], grad_g=[grad_g1, grad_g2] ) @test get_constraints(M, co1v, p) == [c[1], []] - @test get_grad_equality_constraints(M, co1v, p) == [] - @test get_grad_inequality_constraints(M, co1v, p) == gg + @test get_grad_equality_constraint(M, co1v, p, :) == [] + @test get_grad_inequality_constraint(M, co1v, p, :) == gg co2f = ConstrainedManifoldObjective(f, grad_f!; h=h, grad_h=grad_h) @test get_constraints(M, co2f, p) == [[], c[2]] - @test get_grad_equality_constraints(M, co2f, p) == gh - @test get_grad_inequality_constraints(M, co2f, p) == [] + @test get_grad_equality_constraint(M, co2f, p, :) == gh + @test get_grad_inequality_constraint(M, co2f, p, :) == [] co2v = ConstrainedManifoldObjective(f, grad_f!; h=[h1], grad_h=[grad_h1]) @test get_constraints(M, co2v, p) == [[], c[2]] - @test get_grad_equality_constraints(M, co2v, p) == gh - @test get_grad_inequality_constraints(M, co2v, p) == [] + @test get_grad_equality_constraint(M, co2v, p, :) == gh + @test get_grad_inequality_constraint(M, co2v, p, :) == [] end for co in [cofa, cofm, cova, covm] @testset "$co" begin dmp = DefaultManoptProblem(M, co) @test get_constraints(dmp, p) == c - @test get_equality_constraints(dmp, p) == c[2] + @test get_equality_constraint(dmp, p, :) == c[2] @test get_equality_constraint(dmp, p, 1) == c[2][1] - @test get_inequality_constraints(dmp, p) == c[1] + @test get_inequality_constraint(dmp, p, :) == c[1] @test get_inequality_constraint(dmp, p, 1) == c[1][1] @test get_inequality_constraint(dmp, p, 2) == c[1][2] - @test get_grad_equality_constraints(dmp, p) == gh + @test get_grad_equality_constraint(dmp, p, :) == gh Xh = [zeros(3)] - @test get_grad_equality_constraints!(dmp, Xh, p) == gh + @test get_grad_equality_constraint!(dmp, Xh, p, :) == gh @test Xh == gh X = zeros(3) @test get_grad_equality_constraint(dmp, p, 1) == gh[1] @test get_grad_equality_constraint!(dmp, X, p, 1) == gh[1] @test X == gh[1] - @test get_grad_inequality_constraints(dmp, p) == gg + @test get_grad_inequality_constraint(dmp, p, :) == gg Xg = [zeros(3), zeros(3)] - @test get_grad_inequality_constraints!(dmp, Xg, p) == gg + @test get_grad_inequality_constraint!(dmp, Xg, p, :) == gg @test Xg == gg @test get_grad_inequality_constraint(dmp, p, 1) == gg[1] @test get_grad_inequality_constraint!(dmp, X, p, 1) == gg[1] @@ -205,11 +205,12 @@ include("../utils/dummy_types.jl") for obj in [cofa, cofm, cova, covm] ddo = DummyDecoratedObjective(obj) @test get_constraints(M, ddo, p) == get_constraints(M, obj, p) - @test get_equality_constraints(M, ddo, p) == get_equality_constraints(M, obj, p) - @test get_inequality_constraints(M, ddo, p) == - get_inequality_constraints(M, obj, p) - Xe = get_grad_equality_constraints(M, ddo, p) - Ye = get_grad_equality_constraints(M, obj, p) + @test get_equality_constraint(M, ddo, p, :) == + get_equality_constraint(M, obj, p, :) + @test get_inequality_constraint(M, ddo, p, :) == + get_inequality_constraint(M, obj, p, :) + Xe = get_grad_equality_constraint(M, ddo, p, :) + Ye = get_grad_equality_constraint(M, obj, p, :) @test Ye == Xe get_grad_equality_constraints!(M, Xe, ddo, p) get_grad_equality_constraints!(M, Ye, obj, p) @@ -234,11 +235,11 @@ include("../utils/dummy_types.jl") Y = get_grad_inequality_constraint!(M, Y, obj, p, j) @test X == Y end - Xe = get_grad_inequality_constraints(M, ddo, p) - Ye = get_grad_inequality_constraints(M, obj, p) + Xe = get_grad_inequality_constraint(M, ddo, p, :) + Ye = get_grad_inequality_constraint(M, obj, p, :) @test Ye == Xe - get_grad_inequality_constraints!(M, Xe, ddo, p) - get_grad_inequality_constraints!(M, Ye, obj, p) + get_grad_inequality_constraint!(M, Xe, ddo, p, :) + get_grad_inequality_constraint!(M, Ye, obj, p, :) @test Ye == Xe end end @@ -260,14 +261,15 @@ include("../utils/dummy_types.jl") ) @test get_constraints(M, ccofa, p) == get_constraints(M, cofa, p) @test get_count(ccofa, :Constraints) == 1 - @test get_equality_constraints(M, ccofa, p) == get_equality_constraints(M, cofa, p) + @test get_equality_constraint(M, ccofa, p, :) == + get_equality_constraint(M, cofa, p, :) @test get_count(ccofa, :EqualityConstraints) == 1 @test get_equality_constraint(M, ccofa, p, 1) == get_equality_constraint(M, cofa, p, 1) @test get_count(ccofa, :EqualityConstraint) == 1 @test get_count(ccofa, :EqualityConstraint, 1) == 1 - @test get_inequality_constraints(M, ccofa, p) == - get_inequality_constraints(M, cofa, p) + @test get_inequality_constraint(M, ccofa, p, :) == + get_inequality_constraint(M, cofa, p, :) @test get_count(ccofa, :InequalityConstraints) == 1 @test get_inequality_constraint(M, ccofa, p, 1) == get_inequality_constraint(M, cofa, p, 1) @@ -278,10 +280,10 @@ include("../utils/dummy_types.jl") @test get_count(ccofa, :InequalityConstraint, 2) == 1 @test get_count(ccofa, :InequalityConstraint, [1, 2, 3]) == -1 - Xe = get_grad_equality_constraints(M, cofa, p) - @test get_grad_equality_constraints(M, ccofa, p) == Xe + Xe = get_grad_equality_constraint(M, cofa, p, :) + @test get_grad_equality_constraint(M, ccofa, p, :) == Xe Ye = copy.(Ref(M), Ref(p), Xe) - get_grad_equality_constraints!(M, Ye, ccofa, p) + get_grad_equality_constraint!(M, Ye, ccofa, p, :) @test Ye == Xe @test get_count(ccofa, :GradEqualityConstraints) == 2 X = get_grad_equality_constraint(M, cofa, p, 1) @@ -291,10 +293,10 @@ include("../utils/dummy_types.jl") @test Y == X @test get_count(ccofa, :GradEqualityConstraint) == 2 @test get_count(ccofa, :GradEqualityConstraint, 1) == 2 - Xi = get_grad_inequality_constraints(M, cofa, p) - @test get_grad_inequality_constraints(M, ccofa, p) == Xi + Xi = get_grad_inequality_constraint(M, cofa, p, :) + @test get_grad_inequality_constraint(M, ccofa, p, :) == Xi Yi = copy.(Ref(M), Ref(p), Xi) - @test get_grad_inequality_constraints!(M, Yi, ccofa, p) == Xi + @test get_grad_inequality_constraint!(M, Yi, ccofa, p, :) == Xi @test get_count(ccofa, :GradInequalityConstraints) == 2 X1 = get_grad_inequality_constraint(M, cofa, p, 1) @test get_grad_inequality_constraint(M, ccofa, p, 1) == X1 @@ -324,8 +326,8 @@ include("../utils/dummy_types.jl") ] ccofa = Manopt.objective_count_factory(M, cofa, cache_and_count) cccofa = Manopt.objective_cache_factory(M, ccofa, (:LRU, cache_and_count)) - @test get_constraint(M, cofa, p, :) == get_constraint(M, cccofa, p, :) # counts - @test get_constraint(M, cofa, p, :) == get_constraint(M, cccofa, p, :) # cached + @test get_constraints(M, cofa, p) == get_constraints(M, cccofa, p) # counts + @test get_constraints(M, cofa, p) == get_constraints(M, cccofa, p) # cached @test get_count(cccofa, :Constraints) == 1 ce = get_equality_constraint(M, cofa, p, :)