Stochastic Tunneling

[cpfiffer]

@mohamed82008 wrote a nifty thing to do stochastic tunneling and it might be nice to do a little write up on it.

using Distributions, LinearAlgebra

mutable struct ObjDist{F, Tobj, Tsol, Tstep, Tbound} <: Distribution{Multivariate, Continuous}
    f::F
    best_obj::Tobj
    best_sol::Tsol
    step::Tstep
    lb::Tbound
    ub::Tbound
end

function ObjDist(f, N=1; step = 1.0, lb=-Inf, ub=Inf)
    x0 = rand.(TruncatedNormal.(zeros(N), step, lb, ub))
    obj = f(x0)
    return ObjDist(f, obj, x0, step, lb, ub)
end

function Base.rand(dist::ObjDist)
    N = length(dist.best_sol)
    r = rand.(TruncatedNormal.(dist.best_sol, dist.step, dist.lb, dist.ub))
    return r
end

function Distributions.logpdf(dist::ObjDist, x::AbstractVector)
    obj = dist.f(x)
    if obj > dist.best_obj || isnan(dist.best_obj)
        dist.best_obj = obj
        dist.best_sol .= x
    end
    return obj
end

using Turing

function STUN(f, N, alg = MH(10000))
    dist = ObjDist(f, N)
    @model obj_model() = begin
        obj ~ dist
    end
    sample(obj_model(), alg)
    return dist.best_sol, dist.best_obj
end

STUN(x->-norm(x .- 20), 3)