Add keyword n_tries to maximum_likelihood and maximum_a_posteriori

src/optimisation/Optimisation.jl

@@ -506,34 +506,46 @@ end
     maximum_a_posteriori(
         model::DynamicPPL.Model,
         [solver];
+        n_tries = 1,
         kwargs...
     )
 
 Find the maximum a posteriori estimate of a model.
 
-This is a convenience function that calls `estimate_mode` with `MAP()` as the estimator.
-Please see the documentation of [`Turing.Optimisation.estimate_mode`](@ref) for more
+This is a convenience function that calls `estimate_mode` with `MAP()` as the estimator. The keyword `n_tries = 1`
+determines the number of times the mode is estimated from different starting points, which can mitigate the problems with local maxima. The result with the maximum lp is selected from multiple attempts. Please see the documentation of [`Turing.Optimisation.estimate_mode`](@ref) for more
 details.
 """
-function maximum_a_posteriori(model::DynamicPPL.Model, args...; kwargs...)
-    return estimate_mode(model, MAP(), args...; kwargs...)
+function maximum_a_posteriori(model::DynamicPPL.Model, args...; n_tries = 1, kwargs...)
+    result = estimate_mode(model, MAP(), args...; kwargs...)
+    for _ in 2:n_tries 
+        _result = estimate_mode(model, MAP(), args...; kwargs...)
+        result = _result.lp > result.lp ? _result : result 
+    end
+    return result
 end
 
 """
     maximum_likelihood(
         model::DynamicPPL.Model,
         [solver];
+        n_tries = 1,
         kwargs...
     )
 
 Find the maximum likelihood estimate of a model.
 
-This is a convenience function that calls `estimate_mode` with `MLE()` as the estimator.
-Please see the documentation of [`Turing.Optimisation.estimate_mode`](@ref) for more
+This is a convenience function that calls `estimate_mode` with `MLE()` as the estimator. The keyword `n_tries = 1`
+determines the number of times the mode is estimated from different starting points, which can mitigate problems with local maxima. The result with the maximum lp is selected from multiple attempts. Please see the documentation of [`Turing.Optimisation.estimate_mode`](@ref) for more
 details.
 """
-function maximum_likelihood(model::DynamicPPL.Model, args...; kwargs...)
-    return estimate_mode(model, MLE(), args...; kwargs...)
+function maximum_likelihood(model::DynamicPPL.Model, args...; n_tries = 1, kwargs...)
+    mle = estimate_mode(model, MLE(), args...; kwargs...)
+    for _ in 2:n_tries 
+        _mle = estimate_mode(model, MLE(), args...; kwargs...)
+        mle = _mle.lp > mle.lp ? _mle : mle 
+    end
+    return mle
 end
 
 end

test/optimisation/Optimisation.jl

@@ -1,6 +1,6 @@
 module OptimisationTests
 
-using ..Models: gdemo, gdemo_default
+using ..Models: gdemo, gdemo_default, lognormal
 using Distributions
 using Distributions.FillArrays: Zeros
 using DynamicPPL: DynamicPPL
@@ -436,6 +436,45 @@ using Turing
         end
     end
 
+    @testset "MLE with multiple trys" begin    
+        Random.seed!(8454)
+        y = rand(LogNormal(-1, 1), 50) .+ .3
+
+        lb = [-10,0, 0]
+        ub = [10, 10, minimum(y)]
+
+        Random.seed!(80)
+        mle = maximum_likelihood(lognormal(y); lb, ub, n_tries = 20)
+
+        # # Generate a MLE estimate.
+        # #initial_params = round.(rand(Uniform(0, 900), 2), digits = 2)
+        # mle_estimate = maximum_likelihood(inverse_guassian(y); lb, ub, initial_params = [28.97,489.41])
+        Random.seed!(80)
+        lps = map(_ -> maximum_likelihood(lognormal(y); lb, ub).lp, 1:20)
+
+        # test whether there is significant variation in lp 
+        @test var(lps) ≥ 1  
+        @test mle.lp ≈ maximum(lps)
+    end
+
+    @testset "MAP with multiple trys" begin
+        Random.seed!(8454)
+        y = rand(LogNormal(-1, 1), 50) .+ .3
+
+        lb = [-10,0, 0]
+        ub = [10, 10, minimum(y)]
+
+        Random.seed!(80)
+        result = maximum_a_posteriori(lognormal(y); lb, ub, n_tries = 20)
+
+        Random.seed!(80)
+        lps = map(_ -> maximum_a_posteriori(lognormal(y); lb, ub).lp, 1:20)
+
+        # test whether there is significant variation in lp 
+        @test var(lps) ≥ 1  
+        @test result.lp ≈ maximum(lps)
+    end
+
     @testset "StatsBase integration" begin
         Random.seed!(54321)
         mle_est = maximum_likelihood(gdemo_default)

test/runtests.jl

@@ -82,4 +82,4 @@ end
     end
 end
 
-show(TIMEROUTPUT; compact=true, sortby=:firstexec)
+show(TIMEROUTPUT; compact=true, sortby=:firstexec)

test/test_utils/models.jl

@@ -96,4 +96,10 @@ end
 
 MoGtest_default_z_vector = MoGtest_z_vector([1.0 1.0 4.0 4.0])
 
+@model function lognormal(y, min_obs = minimum(y))
+    μ ~ Normal(-1, 2)
+    σ ~ truncated(Normal(.8, 2), 0, Inf)
+    τ ~ Uniform(0, min_obs)
+    y ~ LogNormal(μ, σ) + τ
+end
 end