added product LMO

src/norm_oracles.jl

@@ -13,7 +13,7 @@ end
 
 LpNormLMO{p}(right_hand_side::T) where {T,p} = LpNormLMO{T,p}(right_hand_side)
 
-function compute_extreme_point(lmo::LpNormLMO{T,2}, direction) where {T}
+function compute_extreme_point(lmo::LpNormLMO{T,2}, direction; kwargs...) where {T}
     dir_norm = norm(direction, 2)
     res = similar(direction)
     n = length(direction)
@@ -26,11 +26,11 @@ function compute_extreme_point(lmo::LpNormLMO{T,2}, direction) where {T}
     return res
 end
 
-function compute_extreme_point(lmo::LpNormLMO{T,Inf}, direction) where {T}
+function compute_extreme_point(lmo::LpNormLMO{T,Inf}, direction; kwargs...) where {T}
     return -[lmo.right_hand_side * sign(d) for d in direction]
 end
 
-function compute_extreme_point(lmo::LpNormLMO{T,1}, direction) where {T}
+function compute_extreme_point(lmo::LpNormLMO{T,1}, direction; kwargs...) where {T}
     idx = 0
     v = -one(eltype(direction))
     for i in eachindex(direction)
@@ -42,7 +42,7 @@ function compute_extreme_point(lmo::LpNormLMO{T,1}, direction) where {T}
     return MaybeHotVector(-lmo.right_hand_side * sign(direction[idx]), idx, length(direction))
 end
 
-function compute_extreme_point(lmo::LpNormLMO{T,p}, direction) where {T,p}
+function compute_extreme_point(lmo::LpNormLMO{T,p}, direction; kwargs...) where {T,p}
     # covers the case where the Inf or 1 is of another type
     if p == Inf
         return compute_extreme_point(LpNormLMO{T,Inf}(lmo.right_hand_side), direction)
@@ -69,7 +69,7 @@ struct L1ballDense{T} <: LinearMinimizationOracle
 end
 
 
-function compute_extreme_point(lmo::L1ballDense{T}, direction) where {T}
+function compute_extreme_point(lmo::L1ballDense{T}, direction; kwargs...) where {T}
     idx = 0
     v = -1.0
     for i in eachindex(direction)
@@ -98,7 +98,7 @@ struct KNormBallLMO{T} <: LinearMinimizationOracle
     right_hand_side::T
 end
 
-function compute_extreme_point(lmo::KNormBallLMO{T}, direction) where {T}
+function compute_extreme_point(lmo::KNormBallLMO{T}, direction; kwargs...) where {T}
     K = max(min(lmo.K, length(direction)), 1)
 
     oinf = zero(eltype(direction))

src/oracles.jl

@@ -12,6 +12,7 @@ abstract type LinearMinimizationOracle end
 
 Computes the point `argmin_{v ∈ C} v ⋅ direction`
 with `C` the set represented by the LMO.
+All LMOs should accept keyword arguments that they can ignore.
 """
 function compute_extreme_point end
 
@@ -246,3 +247,46 @@ function compute_extreme_point(
     end
     return v
 end
+
+"""
+    ProductLMO(lmos...)
+
+Linear minimization oracle over the Cartesian product of multiple LMOs.
+"""
+struct ProductLMO{N, TL <: NTuple{N, LinearMinimizationOracle}} <: LinearMinimizationOracle
+    lmos::TL
+end
+
+function ProductLMO{N}(lmos::TL) where {N, TL <: NTuple{N, LinearMinimizationOracle}}
+    return ProductLMO{N, TL}(lmos)
+end
+
+function ProductLMO(lmos::Vararg{LinearMinimizationOracle, N}) where {N}
+    return ProductLMO{N}(lmos)
+end
+
+"""
+    compute_extreme_point(lmo::ProductLMO, direction::Tuple; kwargs...)
+
+Extreme point computation on Cartesian product, with a direction `(d1, d2, ...)` given as a tuple of directions.
+All keyword arguments are passed to all LMOs.
+"""
+function compute_extreme_point(lmo::ProductLMO, direction::Tuple; kwargs...)
+    return compute_extreme_point.(lmo.lmos, direction; kwargs...)
+end
+
+"""
+    compute_extreme_point(lmo::ProductLMO, direction::AbstractArray; direction_indices, storage=similar(direction))
+
+Extreme point computation, with a direction array and `direction_indices` provided such that:
+`direction[direction_indices[i]]` is passed to the i-th LMO.
+The result is stored in the optional `storage` container.
+
+All keyword arguments are passed to all LMOs.
+"""
+function compute_extreme_point(lmo::ProductLMO{N}, direction::AbstractArray; storage=similar(direction), direction_indices, kwargs...) where {N}
+    for idx in 1:N
+        storage[direction_indices[idx]] .= compute_extreme_point(lmo.lmos[idx], direction[direction_indices[idx]]; kwargs...)
+    end
+    return storage
+end

src/polytope_oracles.jl

@@ -15,7 +15,7 @@ struct KSparseLMO{T} <: LinearMinimizationOracle
     right_hand_side::T
 end
 
-function compute_extreme_point(lmo::KSparseLMO{T}, direction) where {T}
+function compute_extreme_point(lmo::KSparseLMO{T}, direction; kwargs...) where {T}
     K = min(lmo.K, length(direction))
     K_indices = sortperm(direction[1:K], by=abs, rev=true)
     K_values = direction[K_indices]
@@ -87,7 +87,7 @@ Its extreme vertices are all permutation matrices of side-dimension `dimension`.
 struct BirkhoffPolytopeLMO <: LinearMinimizationOracle
 end
 
-function compute_extreme_point(::BirkhoffPolytopeLMO, direction::AbstractMatrix{T}) where {T}
+function compute_extreme_point(::BirkhoffPolytopeLMO, direction::AbstractMatrix{T}; kwargs...) where {T}
     n = size(direction, 1)
     n == size(direction, 2) ||
         DimensionMismatch("direction should be square and matching BirkhoffPolytopeLMO dimension")

src/simplex_oracles.jl

@@ -87,7 +87,7 @@ LMO for scaled probability simplex.
 Returns a vector with one active value equal to RHS in the
 most improving (or least degrading) direction.
 """
-function compute_extreme_point(lmo::ProbabilitySimplexOracle{T}, direction) where {T}
+function compute_extreme_point(lmo::ProbabilitySimplexOracle{T}, direction; kwargs...) where {T}
     idx = argmin(direction)
     return MaybeHotVector(lmo.right_side, idx, length(direction))
 end
@@ -120,7 +120,7 @@ for scaled probability simplex.
 Returns two vectors. The first one is the dual costs associated with the constraints 
 and the second is the reduced costs for the variables.
 """
-function compute_dual_solution(::ProbabilitySimplexOracle{T}, direction, primal_solution) where {T}
+function compute_dual_solution(::ProbabilitySimplexOracle{T}, direction, primal_solution; kwargs...) where {T}
     idx = argmax(primal_solution)
     lambda = [direction[idx]]
     mu = direction .- lambda

test/lmo.jl

@@ -573,3 +573,21 @@ end
         end
     end
 end
+
+@testset "Product LMO" begin
+    # 
+    lmo = FrankWolfe.ProductLMO(
+        FrankWolfe.LpNormLMO{Inf}(3.0),
+        FrankWolfe.LpNormLMO{1}(2.0),
+    )
+    dinf = randn(10)
+    d1 = randn(5)
+    vtup = FrankWolfe.compute_extreme_point(lmo, (dinf, d1))
+    @test length(vtup) == 2
+    (vinf, v1) = vtup
+    @test sum(abs, vinf) ≈ 10 * 3.0
+    @test sum(!iszero, v1) == 1
+
+    vvec = FrankWolfe.compute_extreme_point(lmo, [dinf;d1]; direction_indices=(1:10, 11:15))
+    @test vvec ≈ [vinf;v1]
+end