Merge pull request #120 from JuliaDecisionFocusedLearning/perturbed-overhaul

General overhaul

Author: BatyLeo

.gitignore

@@ -1,3 +1,4 @@
+.DS_Store
 *.jl.*.cov
 *.jl.cov
 *.jl.mem

CITATION.bib

@@ -2,7 +2,7 @@ @misc{InferOpt.jl
 	author  = {Guillaume Dalle, Léo Baty, Louis Bouvier and Axel Parmentier},
 	title   = {InferOpt.jl},
 	url     = {https://github.com/axelparmentier/InferOpt.jl},
-	version = {v0.6.1},
-	year    = {2023},
-	month   = {9}
+	version = {v0.7.0},
+	year    = {2025},
+	month   = {4}
 }

Project.toml

@@ -1,37 +1,47 @@
 name = "InferOpt"
 uuid = "4846b161-c94e-4150-8dac-c7ae193c601f"
 authors = ["Guillaume Dalle", "Léo Baty", "Louis Bouvier", "Axel Parmentier"]
-version = "0.6.1"
+version = "0.7.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d"
+DifferentiableExpectations = "fc55d66b-b2a8-4ccc-9d64-c0c2166ceb36"
 DifferentiableFrankWolfe = "b383313e-5450-4164-a800-befbd27b574d"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+FrankWolfe = "f55ce6ea-fdc5-4628-88c5-0087fe54bd30"
+ImplicitDifferentiation = "57b37032-215b-411a-8a7c-41a003a55207"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RequiredInterfaces = "97f35ef4-7bc5-4ec1-a41a-dcc69c7308c6"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
-ThreadsX = "ac1d9e8a-700a-412c-b207-f0111f4b6c0d"
 
 [weakdeps]
 DifferentiableFrankWolfe = "b383313e-5450-4164-a800-befbd27b574d"
+FrankWolfe = "f55ce6ea-fdc5-4628-88c5-0087fe54bd30"
+ImplicitDifferentiation = "57b37032-215b-411a-8a7c-41a003a55207"
 
 [extensions]
-InferOptFrankWolfeExt = "DifferentiableFrankWolfe"
+InferOptFrankWolfeExt = ["DifferentiableFrankWolfe", "FrankWolfe", "ImplicitDifferentiation"]
 
 [compat]
 ChainRulesCore = "1"
 DensityInterface = "0.4.0"
-DifferentiableFrankWolfe = "0.2"
+DifferentiableExpectations = "0.2"
+DifferentiableFrankWolfe = "0.3"
+Distributions = "0.25"
+DocStringExtensions = "0.9"
+FrankWolfe = "0.3"
+ImplicitDifferentiation = "0.6"
 LinearAlgebra = "1"
 Random = "1"
 RequiredInterfaces = "0.1.3"
 Statistics = "1"
 StatsBase = "0.33, 0.34"
 StatsFuns = "1.3"
-ThreadsX = "0.1.11"
 julia = "1.10"
 
 [extras]
@@ -45,6 +55,7 @@ FrankWolfe = "f55ce6ea-fdc5-4628-88c5-0087fe54bd30"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 GridGraphs = "dd2b58c7-5af7-4f17-9e46-57c68ac813fb"
 HiGHS = "87dc4568-4c63-4d18-b0c0-bb2238e4078b"
+ImplicitDifferentiation = "57b37032-215b-411a-8a7c-41a003a55207"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
@@ -61,4 +72,4 @@ UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "DifferentiableFrankWolfe", "Distributions", "Documenter", "FiniteDifferences", "Flux", "FrankWolfe", "Graphs", "GridGraphs", "HiGHS", "JET", "JuliaFormatter", "JuMP", "LinearAlgebra", "Literate", "Pkg", "ProgressMeter", "Random", "Revise", "Statistics", "Test", "TestItemRunner", "UnicodePlots", "Zygote"]
+test = ["Aqua", "DifferentiableFrankWolfe", "Distributions", "Documenter", "FiniteDifferences", "Flux", "FrankWolfe", "Graphs", "GridGraphs", "HiGHS", "ImplicitDifferentiation", "JET", "JuliaFormatter", "JuMP", "LinearAlgebra", "Literate", "Pkg", "ProgressMeter", "Random", "Revise", "Statistics", "Test", "TestItemRunner", "UnicodePlots", "Zygote"]

docs/make.jl

@@ -6,11 +6,20 @@ DocMeta.setdocmeta!(InferOpt, :DocTestSetup, :(using InferOpt); recursive=true)
 
 # Copy README.md into docs/src/index.md (overwriting)
 
-cp(
-    joinpath(dirname(@__DIR__), "README.md"),
-    joinpath(@__DIR__, "src", "index.md");
-    force=true,
-)
+open(joinpath(@__DIR__, "src", "index.md"), "w") do io
+    println(
+        io,
+        """
+        ```@meta
+        EditURL = "https://github.com/JuliaDecisionFocusedLearning/InferOpt.jl/blob/main/README.md"
+        ```
+        """,
+    )
+    # Write the contents out below the meta bloc
+    for line in eachline(joinpath(dirname(@__DIR__), "README.md"))
+        println(io, line)
+    end
+end
 
 # Parse test/tutorial.jl into docs/src/tutorial.md (overwriting)
 
@@ -21,8 +30,14 @@ Literate.markdown(tuto_jl_file, tuto_md_dir; documenter=true, execute=false)
 makedocs(;
     modules=[InferOpt],
     authors="Guillaume Dalle, Léo Baty, Louis Bouvier, Axel Parmentier",
+    repo="https://github.com/JuliaDecisionFocusedLearning/InferOpt.jl/blob/{commit}{path}#{line}",
     sitename="InferOpt.jl",
-    format=Documenter.HTML(),
+    format=Documenter.HTML(;
+        prettyurls=get(ENV, "CI", "false") == "true",
+        canonical="https://juliadecisionfocusedlearning.github.io/InferOpt.jl",
+        assets=String[],
+        repolink="https://github.com/JuliaDecisionFocusedLearning/InferOpt.jl",
+    ),
     pages=[
         "Home" => "index.md",
         "Background" => "background.md",

ext/InferOptFrankWolfeExt.jl

@@ -1,12 +1,31 @@
 module InferOptFrankWolfeExt
 
+using DifferentiableExpectations:
+    DifferentiableExpectations, FixedAtomsProbabilityDistribution
 using DifferentiableFrankWolfe: DifferentiableFrankWolfe, DiffFW
-using DifferentiableFrankWolfe: LinearMinimizationOracle  # from FrankWolfe
-using DifferentiableFrankWolfe: IterativeLinearSolver  # from ImplicitDifferentiation
-using InferOpt: InferOpt, RegularizedFrankWolfe, FixedAtomsProbabilityDistribution
-using InferOpt: compute_expectation, compute_probability_distribution
+using FrankWolfe: LinearMinimizationOracle
+using ImplicitDifferentiation: KrylovLinearSolver
+using InferOpt: InferOpt, RegularizedFrankWolfe
 using LinearAlgebra: dot
 
+"""
+    RegularizedFrankWolfe(linear_maximizer; Ω, Ω_grad, frank_wolfe_kwargs=(;), implicit_kwargs=(; linear_solver=KrylovLinearSolver(; verbose=false)))
+
+Construct a `RegularizedFrankWolfe` struct with a linear maximizer and the necessary components for the Frank-Wolfe algorithm.
+Set `implicit_kwargs` to `(; linear_solver=KrylovLinearSolver(; verbose=true))` if you want to see the solver potential warnings.
+"""
+function RegularizedFrankWolfe(
+    linear_maximizer;
+    Ω,
+    Ω_grad,
+    frank_wolfe_kwargs=NamedTuple(),
+    implicit_kwargs=(; linear_solver=KrylovLinearSolver(; verbose=false)),
+)
+    return RegularizedFrankWolfe(
+        linear_maximizer, Ω, Ω_grad, frank_wolfe_kwargs, implicit_kwargs
+    )
+end
+
 """
     LinearMaximizationOracleWithKwargs{F,K}
 Wraps a linear maximizer as a `FrankWolfe.LinearMinimizationOracle` with a sign switch and predefined keyword arguments.
@@ -40,14 +59,17 @@ Keyword arguments are passed to the underlying linear maximizer.
 function InferOpt.compute_probability_distribution(
     regularized::RegularizedFrankWolfe, θ::AbstractArray; kwargs...
 )
-    (; linear_maximizer, Ω, Ω_grad, frank_wolfe_kwargs) = regularized
+    shape = size(θ)
+    (; linear_maximizer, Ω, Ω_grad, frank_wolfe_kwargs, implicit_kwargs) = regularized
     f(y, θ) = Ω(y) - dot(θ, y)
     f_grad1(y, θ) = Ω_grad(y) - θ
-    lmo = LinearMaximizationOracleWithKwargs(linear_maximizer, kwargs)
-    implicit_kwargs = (; linear_solver=IterativeLinearSolver(; accept_inconsistent=true))
+    maximizer(θ; shape, kwargs...) = vec(linear_maximizer(reshape(θ, shape); kwargs...))
+    lmo = LinearMaximizationOracleWithKwargs(maximizer, (; shape, kwargs...))
     dfw = DiffFW(f, f_grad1, lmo; implicit_kwargs)
-    weights, atoms = dfw.implicit(θ; frank_wolfe_kwargs=frank_wolfe_kwargs)
-    probadist = FixedAtomsProbabilityDistribution(atoms, weights)
+    weights, atoms = dfw.implicit(vec(θ); frank_wolfe_kwargs=frank_wolfe_kwargs)
+    probadist = FixedAtomsProbabilityDistribution(
+        map(atom -> reshape(atom, shape), atoms), weights
+    )
     return probadist
 end
 

src/InferOpt.jl

@@ -10,71 +10,80 @@ module InferOpt
 using ChainRulesCore: ChainRulesCore, NoTangent, RuleConfig, Tangent, ZeroTangent
 using ChainRulesCore: rrule, rrule_via_ad, unthunk
 using DensityInterface: logdensityof
+using DifferentiableExpectations:
+    DifferentiableExpectations, Reinforce, empirical_predistribution, empirical_distribution
+using Distributions:
+    Distributions,
+    ContinuousUnivariateDistribution,
+    LogNormal,
+    Normal,
+    product_distribution,
+    logpdf
+using DocStringExtensions: TYPEDEF, TYPEDFIELDS, TYPEDSIGNATURES
 using LinearAlgebra: dot
-using Random: AbstractRNG, GLOBAL_RNG, MersenneTwister, rand, seed!
+using Random: Random, AbstractRNG, GLOBAL_RNG, MersenneTwister, rand, seed!
 using Statistics: mean
 using StatsBase: StatsBase, sample
 using StatsFuns: logaddexp, softmax
-using ThreadsX: ThreadsX
 using RequiredInterfaces
 
 include("interface.jl")
 
+include("utils/utils.jl")
 include("utils/some_functions.jl")
-include("utils/probability_distribution.jl")
 include("utils/pushforward.jl")
-include("utils/generalized_maximizer.jl")
+include("utils/linear_maximizer.jl")
 include("utils/isotonic_regression/isotonic_l2.jl")
 include("utils/isotonic_regression/isotonic_kl.jl")
 include("utils/isotonic_regression/projection.jl")
 
-include("simple/interpolation.jl")
-include("simple/identity.jl")
+# Layers
+include("layers/simple/interpolation.jl")
+include("layers/simple/identity.jl")
 
-include("regularized/abstract_regularized.jl")
-include("regularized/soft_argmax.jl")
-include("regularized/sparse_argmax.jl")
-include("regularized/soft_rank.jl")
-include("regularized/regularized_frank_wolfe.jl")
+include("layers/perturbed/utils.jl")
+include("layers/perturbed/perturbation.jl")
+include("layers/perturbed/perturbed.jl")
 
-include("perturbed/abstract_perturbed.jl")
-include("perturbed/additive.jl")
-include("perturbed/multiplicative.jl")
-include("perturbed/perturbed_oracle.jl")
-
-include("imitation/spoplus_loss.jl")
-include("imitation/ssvm_loss.jl")
-include("imitation/fenchel_young_loss.jl")
-include("imitation/imitation_loss.jl")
-include("imitation/zero_one_loss.jl")
+include("layers/regularized/abstract_regularized.jl")
+include("layers/regularized/soft_argmax.jl")
+include("layers/regularized/sparse_argmax.jl")
+include("layers/regularized/soft_rank.jl")
+include("layers/regularized/regularized_frank_wolfe.jl")
 
 if !isdefined(Base, :get_extension)
     include("../ext/InferOptFrankWolfeExt.jl")
 end
 
+# Losses
+include("losses/fenchel_young_loss.jl")
+include("losses/spoplus_loss.jl")
+include("losses/ssvm_loss.jl")
+include("losses/zero_one_loss.jl")
+include("losses/imitation_loss.jl")
+
+export compute_probability_distribution
+
 export half_square_norm
 export shannon_entropy, negative_shannon_entropy
 export one_hot_argmax, ranking
-export GeneralizedMaximizer, objective_value
+export LinearMaximizer, apply_g, objective_value
 
-export FixedAtomsProbabilityDistribution
-export compute_expectation
-export compute_probability_distribution
 export Pushforward
 
 export IdentityRelaxation
 export Interpolation
 
-export AbstractRegularized, AbstractRegularizedGeneralizedMaximizer
+export AbstractRegularized
 export SoftArgmax, soft_argmax
 export SparseArgmax, sparse_argmax
 export SoftRank, soft_rank, soft_rank_l2, soft_rank_kl
 export SoftSort, soft_sort, soft_sort_l2, soft_sort_kl
 export RegularizedFrankWolfe
 
+export PerturbedOracle
 export PerturbedAdditive
 export PerturbedMultiplicative
-export PerturbedOracle
 
 export FenchelYoungLoss
 export StructuredSVMLoss