TuringLang / AdvancedHMC.jl

Robust, modular and efficient implementation of advanced Hamiltonian Monte Carlo algorithms
https://turinglang.org/AdvancedHMC.jl/
MIT License
228 stars 39 forks source link

Define a model #363

Open Vilin97 opened 8 months ago

Vilin97 commented 8 months ago

It would be great if there was an easy interface to define a model without having to create a custom struct and wrap it in LogDensityModel. Its signature would be LogDensityModel(logpdf, D) and it would be something like

function LogDensityModel(logpdf, D)
    struct LogTargetDensity
        dim::Int
    end
    LogDensityProblems.logdensity(p::LogTargetDensity, θ) = logpdf(θ)
    LogDensityProblems.dimension(p::LogTargetDensity) = p.dim
    LogDensityProblems.capabilities(::Type{LogTargetDensity}) = LogDensityProblems.LogDensityOrder{0}()

    model = AdvancedHMC.LogDensityModel(LogDensityProblemsAD.ADgradient(Val(:ForwardDiff), LogTargetDensity(D))) 
end

With such interface the lengthy example from the README would become

D = 3; initial_θ = rand(D)
model = LogDensityModel(x -> -sum(abs2, x)/2, D)
n_samples, n_adapts, δ = 1_000, 2_000, 0.8
sampler = NUTS(δ)
samples = AbstractMCMC.sample(
      model,
      sampler,
      n_adapts + n_samples;
      nadapts = n_adapts,
      initial_params = initial_θ,
  )
torfjelde commented 8 months ago

We should probably do something like:

struct LogDensityFunction{F}
    f::F
    dim::Int
end

LogDensityProblems.logdensity(ldf::LogDensityFunction, x) = ldf.f(x)
LogDensityProblems.dimension(ldf::LogDensityFunction) = ldf.dim
LogDensityProblems.capabilities(::Type{<:LogDensityFunction}) = LogDensityProblems.LogDensityOrder{0}()

which is going to be type-stable.

But IMO this is not something we should put in AdvancedHMC.jl :confused: Probably would be better suited to go in LogDensityProblems.jl itself or maaaybe AbstractMCMC.jl or something like this. @devmotion thoughts?