Rewrite everything using probability distributions

Author: gdalle

CITATION.bib

@@ -1,8 +1,8 @@
 @misc{InferOpt.jl,
-	author  = {Axel Parmentier, Guillaume Dalle, Léo Baty, Louis Bouvier},
+	author  = {Guillaume Dalle, Léo Baty, Louis Bouvier and Axel Parmentier},
 	title   = {InferOpt.jl},
 	url     = {https://github.com/axelparmentier/InferOpt.jl},
-	version = {v0.1.0},
+	version = {v0.2.0},
 	year    = {2022},
-	month   = {4}
+	month   = {6}
 }

Project.toml

@@ -1,6 +1,6 @@
 name = "InferOpt"
 uuid = "4846b161-c94e-4150-8dac-c7ae193c601f"
-authors = ["Axel Parmentier", "Guillaume Dalle", "Léo Baty", "Louis Bouvier"]
+authors = ["Guillaume Dalle", "Léo Baty", "Louis Bouvier", "Axel Parmentier"]
 version = "0.2.0"
 
 [deps]
@@ -13,14 +13,16 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SimpleTraits = "699a6c99-e7fa-54fc-8d76-47d257e15c1d"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
 ChainRulesCore = "1"
-FrankWolfe = "0.2"
-Krylov = "0.8.2"
-LinearOperators = "2.3.2"
+FrankWolfe = "0.2.1"
+Krylov = "0.8"
+LinearOperators = "2.3"
 SimpleTraits = "0.9"
+StatsBase = "0.33"
 julia = "1.7"
 
 [extras]

docs/Manifest.toml

@@ -16,9 +16,9 @@ version = "0.0.1"
 
 [[deps.AbstractFFTs]]
 deps = ["ChainRulesCore", "LinearAlgebra"]
-git-tree-sha1 = "6f1d9bc1c08f9f4a8fa92e3ea3cb50153a1b40d4"
+git-tree-sha1 = "69f7020bd72f069c219b5e8c236c1fa90d2cb409"
 uuid = "621f4979-c628-5d54-868e-fcf4e3e8185c"
-version = "1.1.0"
+version = "1.2.1"
 
 [[deps.Adapt]]
 deps = ["LinearAlgebra"]
@@ -49,15 +49,15 @@ version = "3.5.1+1"
 
 [[deps.ArrayInterface]]
 deps = ["ArrayInterfaceCore", "Compat", "IfElse", "LinearAlgebra", "Static"]
-git-tree-sha1 = "1d062b8ab719670c16024105ace35e6d32988d4f"
+git-tree-sha1 = "6ccb71b40b04ad69152f1f83d5925de13911417e"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "6.0.18"
+version = "6.0.19"
 
 [[deps.ArrayInterfaceCore]]
 deps = ["LinearAlgebra", "SparseArrays", "SuiteSparse"]
-git-tree-sha1 = "5e732808bcf7bbf730e810a9eaafc52705b38bb5"
+git-tree-sha1 = "7d255eb1d2e409335835dc8624c35d97453011eb"
 uuid = "30b0a656-2188-435a-8636-2ec0e6a096e2"
-version = "0.1.13"
+version = "0.1.14"
 
 [[deps.Artifacts]]
 uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
@@ -96,15 +96,15 @@ version = "3.11.0"
 
 [[deps.ChainRules]]
 deps = ["ChainRulesCore", "Compat", "IrrationalConstants", "LinearAlgebra", "Random", "RealDot", "SparseArrays", "Statistics"]
-git-tree-sha1 = "97fd0a3b7703948a847265156a41079730805c77"
+git-tree-sha1 = "b06ed86d99c982cbe9047a45a93ac62d9605a361"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "1.36.0"
+version = "1.36.2"
 
 [[deps.ChainRulesCore]]
 deps = ["Compat", "LinearAlgebra", "SparseArrays"]
-git-tree-sha1 = "9489214b993cd42d17f44c36e359bf6a7c919abf"
+git-tree-sha1 = "2dd813e5f2f7eec2d1268c57cf2373d3ee91fcea"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "1.15.0"
+version = "1.15.1"
 
 [[deps.ChangesOfVariables]]
 deps = ["ChainRulesCore", "LinearAlgebra", "Test"]
@@ -172,15 +172,14 @@ uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"
 
 [[deps.ConstructionBase]]
 deps = ["LinearAlgebra"]
-git-tree-sha1 = "f74e9d5388b8620b4cee35d4c5a618dd4dc547f4"
+git-tree-sha1 = "59d00b3139a9de4eb961057eabb65ac6522be954"
 uuid = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
-version = "1.3.0"
+version = "1.4.0"
 
 [[deps.Contour]]
-deps = ["StaticArrays"]
-git-tree-sha1 = "9f02045d934dc030edad45944ea80dbd1f0ebea7"
+git-tree-sha1 = "a599cfb8b1909b0f97c5e1b923ab92e1c0406076"
 uuid = "d38c429a-6771-53c6-b99e-75d170b6e991"
-version = "0.5.7"
+version = "0.6.1"
 
 [[deps.Crayons]]
 git-tree-sha1 = "249fe38abf76d48563e2f4556bebd215aa317e15"
@@ -294,9 +293,9 @@ version = "0.10.30"
 
 [[deps.FrankWolfe]]
 deps = ["Arpack", "GenericSchur", "Hungarian", "LinearAlgebra", "MathOptInterface", "Printf", "ProgressMeter", "Random", "Setfield", "SparseArrays", "TimerOutputs"]
-git-tree-sha1 = "75f6e18896767729d011bd473958043dc9948641"
+git-tree-sha1 = "87622f91e2256920debd418da506a9490f5c3efc"
 uuid = "f55ce6ea-fdc5-4628-88c5-0087fe54bd30"
-version = "0.2.0"
+version = "0.2.1"
 
 [[deps.FreeType]]
 deps = ["CEnum", "FreeType2_jll"]
@@ -391,7 +390,7 @@ uuid = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
 version = "0.1.1"
 
 [[deps.InferOpt]]
-deps = ["ChainRulesCore", "FrankWolfe", "Krylov", "LinearAlgebra", "LinearOperators", "Random", "SimpleTraits", "SparseArrays", "Statistics", "Test"]
+deps = ["ChainRulesCore", "FrankWolfe", "Krylov", "LinearAlgebra", "LinearOperators", "Random", "SimpleTraits", "SparseArrays", "Statistics", "StatsBase", "Test"]
 path = ".."
 uuid = "4846b161-c94e-4150-8dac-c7ae193c601f"
 version = "0.2.0"
@@ -743,15 +742,15 @@ uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [[deps.SpecialFunctions]]
 deps = ["ChainRulesCore", "IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_jll"]
-git-tree-sha1 = "a9e798cae4867e3a41cae2dd9eb60c047f1212db"
+git-tree-sha1 = "d75bda01f8c31ebb72df80a46c88b25d1c79c56d"
 uuid = "276daf66-3868-5448-9aa4-cd146d93841b"
-version = "2.1.6"
+version = "2.1.7"
 
 [[deps.Static]]
 deps = ["IfElse"]
-git-tree-sha1 = "11f1b69a28b6e4ca1cc18342bfab7adb7ff3a090"
+git-tree-sha1 = "46638763d3a25ad7818a15d441e0c3446a10742d"
 uuid = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
-version = "0.7.3"
+version = "0.7.5"
 
 [[deps.StaticArrays]]
 deps = ["LinearAlgebra", "Random", "StaticArraysCore", "Statistics"]
@@ -760,9 +759,9 @@ uuid = "90137ffa-7385-5640-81b9-e52037218182"
 version = "1.5.0"
 
 [[deps.StaticArraysCore]]
-git-tree-sha1 = "6edcea211d224fa551ec8a85debdc6d732f155dc"
+git-tree-sha1 = "66fe9eb253f910fe8cf161953880cfdaef01cdf0"
 uuid = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
-version = "1.0.0"
+version = "1.0.1"
 
 [[deps.Statistics]]
 deps = ["LinearAlgebra", "SparseArrays"]
@@ -776,15 +775,15 @@ version = "1.4.0"
 
 [[deps.StatsBase]]
 deps = ["DataAPI", "DataStructures", "LinearAlgebra", "LogExpFunctions", "Missings", "Printf", "Random", "SortingAlgorithms", "SparseArrays", "Statistics", "StatsAPI"]
-git-tree-sha1 = "642f08bf9ff9e39ccc7b710b2eb9a24971b52b1a"
+git-tree-sha1 = "48598584bacbebf7d30e20880438ed1d24b7c7d6"
 uuid = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
-version = "0.33.17"
+version = "0.33.18"
 
 [[deps.StructArrays]]
 deps = ["Adapt", "DataAPI", "StaticArrays", "Tables"]
-git-tree-sha1 = "55ef24d228f9396fa9e2317eb60c953b8cec1ae7"
+git-tree-sha1 = "ec47fb6069c57f1cee2f67541bf8f23415146de7"
 uuid = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
-version = "0.6.10"
+version = "0.6.11"
 
 [[deps.SuiteSparse]]
 deps = ["Libdl", "LinearAlgebra", "Serialization", "SparseArrays"]
@@ -841,9 +840,9 @@ uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
 
 [[deps.UnicodePlots]]
 deps = ["ColorSchemes", "ColorTypes", "Contour", "Crayons", "Dates", "FileIO", "FreeTypeAbstraction", "LazyModules", "LinearAlgebra", "MarchingCubes", "NaNMath", "Printf", "SparseArrays", "StaticArrays", "StatsBase", "Unitful"]
-git-tree-sha1 = "3cfa7a287c202c4c574ab535971c4c5b6c104371"
+git-tree-sha1 = "e8192bf70f28cf0e79ae9c215008e2f1464edbd6"
 uuid = "b8865327-cd53-5732-bb35-84acbb429228"
-version = "3.0.0"
+version = "3.0.3"
 
 [[deps.Unitful]]
 deps = ["ConstructionBase", "Dates", "LinearAlgebra", "Random"]
@@ -856,10 +855,10 @@ deps = ["Libdl"]
 uuid = "83775a58-1f1d-513f-b197-d71354ab007a"
 
 [[deps.Zygote]]
-deps = ["AbstractFFTs", "ChainRules", "ChainRulesCore", "DiffRules", "Distributed", "FillArrays", "ForwardDiff", "IRTools", "InteractiveUtils", "LinearAlgebra", "MacroTools", "NaNMath", "Random", "Requires", "SparseArrays", "SpecialFunctions", "Statistics", "ZygoteRules"]
-git-tree-sha1 = "a49267a2e5f113c7afe93843deea7461c0f6b206"
+deps = ["AbstractFFTs", "ChainRules", "ChainRulesCore", "DiffRules", "Distributed", "FillArrays", "ForwardDiff", "IRTools", "InteractiveUtils", "LinearAlgebra", "LogExpFunctions", "MacroTools", "NaNMath", "Random", "Requires", "SparseArrays", "SpecialFunctions", "Statistics", "ZygoteRules"]
+git-tree-sha1 = "3cfdb31b517eec4173584fba2b1aa65daad46e09"
 uuid = "e88e6eb3-aa80-5325-afca-941959d7151f"
-version = "0.6.40"
+version = "0.6.41"
 
 [[deps.ZygoteRules]]
 deps = ["MacroTools"]

docs/src/algorithms.md

@@ -6,6 +6,13 @@
 Modules = [InferOpt]
 ```
 
+## Probability distributions
+
+```@autodocs
+Modules = [InferOpt]
+Pages = ["utils/probability_distribution.jl", "utils/composition.jl"]
+```
+
 ## Interpolation
 
 !!! note "Reference"

src/InferOpt.jl

@@ -7,13 +7,17 @@ using FrankWolfe: away_frank_wolfe, compute_extreme_point
 using Krylov: gmres
 using LinearAlgebra
 using LinearOperators: LinearOperator
-using Random
+using Random: AbstractRNG, GLOBAL_RNG, MersenneTwister, rand, seed!
 using SimpleTraits: SimpleTraits
 using SimpleTraits: @traitdef, @traitfn, @traitimpl
 using SparseArrays
 using Statistics
+using StatsBase: StatsBase, sample
 using Test
 
+include("utils/probability_distribution.jl")
+include("utils/composition.jl")
+
 include("interpolation/interpolation.jl")
 
 include("frank_wolfe/frank_wolfe_utils.jl")
@@ -26,10 +30,10 @@ include("regularized/sparse_argmax.jl")
 include("regularized/regularized_generic.jl")
 
 include("perturbed/abstract_perturbed.jl")
-include("perturbed/composition.jl")
 include("perturbed/additive.jl")
 include("perturbed/multiplicative.jl")
 
+include("fenchel_young/perturbed.jl")
 include("fenchel_young/fenchel_young.jl")
 
 include("spo/spoplus_loss.jl")
@@ -38,7 +42,9 @@ include("ssvm/isbaseloss.jl")
 include("ssvm/zeroone_baseloss.jl")
 include("ssvm/ssvm_loss.jl")
 
-export get_probability_distribution
+export FixedAtomsProbabilityDistribution, sample, compute_expectation
+export ProbabilisticComposition
+export compute_probability_distribution
 
 export Interpolation
 
@@ -53,7 +59,6 @@ export RegularizedGeneric
 
 export PerturbedAdditive
 export PerturbedMultiplicative
-export PerturbedComposition
 
 export FenchelYoungLoss
 

src/fenchel_young/fenchel_young.jl

@@ -20,13 +20,13 @@ end
 ## Forward pass
 
 function (fyl::FenchelYoungLoss)(
-    θ::AbstractArray{<:Real}, y::AbstractArray{<:Real}; kwargs...
+    θ::AbstractArray{<:Real}, y_true::AbstractArray{<:Real}; kwargs...
 )
-    _, l = prediction_and_loss(fyl, θ, y; kwargs...)
+    l, _ = fenchel_young_loss_and_grad(fyl, θ, y_true; kwargs...)
     return l
 end
 
-@traitfn function prediction_and_loss(
+@traitfn function fenchel_young_loss_and_grad(
     fyl::FenchelYoungLoss{P},
     θ::AbstractArray{<:Real},
     y_true::AbstractArray{<:Real};
@@ -37,24 +37,32 @@ end
     Ωy_true = compute_regularization(predictor, y_true)
     Ωŷ = compute_regularization(predictor, ŷ)
     l = (Ωy_true - dot(θ, y_true)) - (Ωŷ - dot(θ, ŷ))
-    return ŷ, l
+    g = ŷ - y_true
+    return l, g
 end
 
-function prediction_and_loss(
-    fyl::FenchelYoungLoss{P}, θ::AbstractArray{<:Real}, y::AbstractArray{<:Real}; kwargs...
+function fenchel_young_loss_and_grad(
+    fyl::FenchelYoungLoss{P},
+    θ::AbstractArray{<:Real},
+    y_true::AbstractArray{<:Real};
+    kwargs...,
 ) where {P<:AbstractPerturbed}
     (; predictor) = fyl
-    ŷ, F = compute_y_and_F(predictor, θ; kwargs...)
-    l = F - dot(θ, y)
-    return ŷ, l
+    F, almost_ŷ = fenchel_young_F_and_first_part_of_grad(predictor, θ; kwargs...)
+    l = F - dot(θ, y_true)
+    g = almost_ŷ - y_true
+    return l, g
 end
 
 ## Backward pass
 
 function ChainRulesCore.rrule(
-    fyl::FenchelYoungLoss, θ::AbstractArray{<:Real}, y::AbstractArray{<:Real}; kwargs...
+    fyl::FenchelYoungLoss,
+    θ::AbstractArray{<:Real},
+    y_true::AbstractArray{<:Real};
+    kwargs...,
 )
-    ŷ, l = prediction_and_loss(fyl, θ, y; kwargs...)
-    fyl_pullback(dl) = NoTangent(), dl * (ŷ - y), NoTangent()
+    l, g = fenchel_young_loss_and_grad(fyl, θ, y_true; kwargs...)
+    fyl_pullback(dl) = NoTangent(), dl * g, NoTangent()
     return l, fyl_pullback
 end

src/fenchel_young/perturbed.jl

@@ -0,0 +1,38 @@
+function fenchel_young_F_and_first_part_of_grad(
+    perturbed::AbstractPerturbed, θ::AbstractArray{<:Real}; kwargs...
+)
+    Z_samples = sample_perturbations(perturbed, θ)
+    F_and_y_samples = [
+        fenchel_young_F_and_first_part_of_grad(perturbed, θ, Z; kwargs...) for
+        Z in Z_samples
+    ]
+    return mean(first, F_and_y_samples), mean(last, F_and_y_samples)
+end
+
+function fenchel_young_F_and_first_part_of_grad(
+    perturbed::PerturbedAdditive,
+    θ::AbstractArray{<:Real},
+    Z::AbstractArray{<:Real};
+    kwargs...,
+)
+    (; maximizer, ε) = perturbed
+    θ_perturbed = θ .+ ε .* Z
+    y = maximizer(θ_perturbed; kwargs...)
+    F = dot(θ_perturbed, y)
+    return F, y
+end
+
+function fenchel_young_F_and_first_part_of_grad(
+    perturbed::PerturbedMultiplicative,
+    θ::AbstractArray{<:Real},
+    Z::AbstractArray{<:Real};
+    kwargs...,
+)
+    (; maximizer, ε) = perturbed
+    eZ = exp.(ε .* Z .- ε^2)
+    θ_perturbed = θ .* eZ
+    y = maximizer(θ_perturbed; kwargs...)
+    F = dot(θ_perturbed, y)
+    y_scaled = y .* eZ
+    return F, y_scaled
+end