Turing custom distribution implementation

Specific Domains Probabilistic programming
,
1

Hi, I am fairly new to Julia and have been having trouble implementing custom distributions for use within Turing.

The model is defined as follows. Let Y denote the data which we assume are generated by Y = \theta + X,

where \theta > 0 is a model parameter and X \sim \text{Exponential}(\lambda) where \lambda > 0 is the rate. The density of the data Y then is f_Y(y \mid \theta, \lambda) = f_X(y - \theta \mid \lambda), i.e. a shifted exponential distribution which is defined for all y \geq \theta and zero otherwise.

I have tried to follow the documentation in the Advanced Usage section “How to Define a Customized Distribution” where I first defined the distribution type, implemented methods for sampling and evaluating the log pdf and defined the domain. However, I am receiving type errors, which seem to be a result of ForwardDiff.Dual objects being passed to the distribution constructor. I have tried messing around with various type conversions and comparing with other standard implementations in the Distributions.jl library however have had no luck. This model is deliberately simple to be minimally reproducible. Please see an example code snippet below.

Any guidance or advice would be greatly appreciated.

Many thanks,
Joseph

using Distributions
using Random
using Turing

## 1. Define distribution type
struct ShiftedExponential{T<:Real} <: ContinuousUnivariateDistribution
    rate::Float64
    shift::Float64
    ShiftedExponential{T}(rate::T, shift::T) where {T} = new{T}(rate, shift)
end

function ShiftedExponential(rate::T, shift::T; check_args = true) where {T<:Real}
    Distributions.@check_args ShiftedExponential (rate, rate > zero(rate)) (shift, shift > zero(shift))
    return ShiftedExponential{T}(rate, shift)
end

ShiftedExponential(rate::Real, shift::Real; check_args = true) = ShiftedExponential(promote(rate, shift)...; check_args = check_args)
ShiftedExponential(rate::Integer, shift::Integer; check_args = true) = ShiftedExponential(float(rate), float(shift); check_args = check_args)

Distributions.@distr_support ShiftedExponential d.shift Inf

Base.eltype(::Type{ShiftedExponential{T}}) where {T} = T

## Conversions
Base.convert(::Type{ShiftedExponential{T}}, rate::S, shift::S) where {T <: Real, S <: Real} = ShiftedExponential(T(rate), T(shift))
Base.convert(::Type{ShiftedExponential{T}}, d::ShiftedExponential) where {T} = ShiftedExponential{T}(T(d.rate), T(d.shift))
Base.convert(::Type{ShiftedExponential{T}}, d::ShiftedExponential{T}) where {T} = d

## Parameters
Distributions.rate(d::ShiftedExponential) = d.rate
Distributions.params(d::ShiftedExponential) = (d.rate, d.shift)
Distributions.partype(::ShiftedExponential{T}) where {T} = T

# 2. Implement sampling and evaluation of the log pdf
function Distributions.logpdf(d::ShiftedExponential, y::Real)
    if y < d.shift
        return oftype(y, -Inf)
    else
        return log(d.rate) - d.rate * (y - d.shift)
    end
end

function Distributions.rand(rng::AbstractRNG, d::ShiftedExponential)
    return rand(rng, Exponential(1/d.rate)) .+ d.shift
end

function Distributions.cdf(d::ShiftedExponential, y::Real)
    if x < d.shift
        return zero(x)
    else
        return 1 - exp(-d.rate * (y - d.shift))
    end
end

function Distributions.logcdf(d::ShiftedExponential, y::Real)
    if x < d.shift
        return log(zero(x))
    else
        return log1p(-exp(-d.rate * (y - d.shift)))
    end
end

function Distributions.ccdf(d::ShiftedExponential, y::Real)
    if x < d.shift
        return one(x)
    else
        return exp(-d.rate * (y - d.shift))
    end
end

function Distributions.logccdf(d::ShiftedExponential, y::Real)
    if x < d.shift
        return log(one(x))
    else
        return -d.rate * (y - d.shift)
    end
end

function Distributions.pdf(d::ShiftedExponential, y::Real)
    if x < d.shift
        return zero(x)
    else
        return d.rate * exp(-d.rate * (y - d.shift))
    end
end

function Distributions.logpdf(d::ShiftedExponential, y::Real)
    if y < d.shift
        return oftype(y, -Inf)
    else
        return log(d.rate) - d.rate * (y - d.shift)
    end
end

function Distributions.gradlogpdf(d::ShiftedExponential, y::Real)
    if y < d.shift
        return zero(y)
    else
        return -d.rate
    end
end



## Define helper Functions for min and maximum
Distributions.minimum(d::ShiftedExponential) = d.shift
Distributions.maximum(d::ShiftedExponential) = Inf


## Simulate some data from the distribution
rate = 0.7
shift = 1.45
nobs = 100
dist = ShiftedExponential(rate, shift)
y = rand(dist, nobs)

@model function shifted_exponential(y)
    rate ~ Exponential(1/0.01)
    shift ~ Exponential(1/0.01)
    y ~ ShiftedExponential(rate, shift)
end

chain = sample(shifted_exponential(y), NUTS(), 5000, progress=true)

with stack trace

ERROR: MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{DynamicPPL.DynamicPPLTag, Float64}, Float64, 2})

Closest candidates are:
  (::Type{T})(::Real, ::RoundingMode) where T<:AbstractFloat
   @ Base rounding.jl:207
  (::Type{T})(::T) where T<:Number
   @ Core boot.jl:792
  Float64(::IrrationalConstants.Inv2π)
   @ IrrationalConstants C:\Users\joesm\.julia\packages\IrrationalConstants\vp5v4\src\macro.jl:112
  ...

Stacktrace:
  [1] convert(::Type{Float64}, x::ForwardDiff.Dual{ForwardDiff.Tag{DynamicPPL.DynamicPPLTag, Float64}, Float64, 2})
    @ Base .\number.jl:7
  [2] ShiftedExponential{…}(rate::ForwardDiff.Dual{…}, shift::ForwardDiff.Dual{…})
    @ Main c:\Users\joesm\OneDrive - The University of Nottingham (1)\Research\Infinity\bayesian-inference-lexical-decision\Code\simple_shift\shifted_exponential.jl:9
  [3] ShiftedExponential(rate::ForwardDiff.Dual{…}, shift::ForwardDiff.Dual{…}; check_args::Bool)
    @ Main c:\Users\joesm\OneDrive - The University of Nottingham (1)\Research\Infinity\bayesian-inference-lexical-decision\Code\simple_shift\shifted_exponential.jl:14
  [4] ShiftedExponential(rate::ForwardDiff.Dual{…}, shift::ForwardDiff.Dual{…})
    @ Main c:\Users\joesm\OneDrive - The University of Nottingham (1)\Research\Infinity\bayesian-inference-lexical-decision\Code\simple_shift\shifted_exponential.jl:12
  [5] macro expansion
    @ C:\Users\joesm\.julia\packages\DynamicPPL\93t4P\src\compiler.jl:579 [inlined]
  [6] shifted_exponential(__model__::DynamicPPL.Model{…}, __varinfo__::DynamicPPL.TypedVarInfo{…}, __context__::DynamicPPL.SamplingContext{…}, y::Vector{…})
    @ Main c:\Users\joesm\OneDrive - The University of Nottingham (1)\Research\Infinity\bayesian-inference-lexical-decision\Code\simple_shift\shifted_exponential.jl:120
  [7] _evaluate!!(model::DynamicPPL.Model{…}, varinfo::DynamicPPL.TypedVarInfo{…}, context::DynamicPPL.SamplingContext{…})
    @ DynamicPPL C:\Users\joesm\.julia\packages\DynamicPPL\93t4P\src\model.jl:973
  [8] evaluate_threadunsafe!!(model::DynamicPPL.Model{…}, varinfo::DynamicPPL.TypedVarInfo{…}, context::DynamicPPL.SamplingContext{…})
    @ DynamicPPL C:\Users\joesm\.julia\packages\DynamicPPL\93t4P\src\model.jl:946
  [9] evaluate!!(model::DynamicPPL.Model{…}, varinfo::DynamicPPL.TypedVarInfo{…}, context::DynamicPPL.SamplingContext{…})
    @ DynamicPPL C:\Users\joesm\.julia\packages\DynamicPPL\93t4P\src\model.jl:894
 [10] logdensity(f::LogDensityFunction{…}, θ::Vector{…})
    @ DynamicPPL C:\Users\joesm\.julia\packages\DynamicPPL\93t4P\src\logdensityfunction.jl:140
 [11] (::Base.Fix1{typeof(LogDensityProblems.logdensity), LogDensityFunction{…}})(y::Vector{ForwardDiff.Dual{…}})
    @ Base .\operators.jl:1118
 [12] vector_mode_dual_eval!(f::Base.Fix1{…}, cfg::ForwardDiff.GradientConfig{…}, x::Vector{…})
    @ ForwardDiff C:\Users\joesm\.julia\packages\ForwardDiff\PcZ48\src\apiutils.jl:24
 [13] vector_mode_gradient!(result::DiffResults.MutableDiffResult{…}, f::Base.Fix1{…}, x::Vector{…}, cfg::ForwardDiff.GradientConfig{…})
    @ ForwardDiff C:\Users\joesm\.julia\packages\ForwardDiff\PcZ48\src\gradient.jl:96
 [14] gradient!(result::DiffResults.MutableDiffResult{…}, f::Base.Fix1{…}, x::Vector{…}, cfg::ForwardDiff.GradientConfig{…}, ::Val{…})
    @ ForwardDiff C:\Users\joesm\.julia\packages\ForwardDiff\PcZ48\src\gradient.jl:37
 [15] gradient!(result::DiffResults.MutableDiffResult{…}, f::Base.Fix1{…}, x::Vector{…}, cfg::ForwardDiff.GradientConfig{…})
    @ ForwardDiff C:\Users\joesm\.julia\packages\ForwardDiff\PcZ48\src\gradient.jl:35
 [16] logdensity_and_gradient(fℓ::LogDensityProblemsADForwardDiffExt.ForwardDiffLogDensity{…}, x::Vector{…})
    @ LogDensityProblemsADForwardDiffExt C:\Users\joesm\.julia\packages\LogDensityProblemsAD\rBlLq\ext\LogDensityProblemsADForwardDiffExt.jl:118
 [17] (::Turing.Inference.var"#∂logπ∂θ#32"{LogDensityProblemsADForwardDiffExt.ForwardDiffLogDensity{…}})(x::Vector{Float64})
    @ Turing.Inference C:\Users\joesm\.julia\packages\Turing\r3Hmj\src\mcmc\hmc.jl:180
 [18] ∂H∂θ(h::AdvancedHMC.Hamiltonian{…}, θ::Vector{…})
    @ AdvancedHMC C:\Users\joesm\.julia\packages\AdvancedHMC\AlvV4\src\hamiltonian.jl:38
 [19] phasepoint(h::AdvancedHMC.Hamiltonian{…}, θ::Vector{…}, r::Vector{…})
    @ AdvancedHMC C:\Users\joesm\.julia\packages\AdvancedHMC\AlvV4\src\hamiltonian.jl:74
 [20] phasepoint(rng::TaskLocalRNG, θ::Vector{…}, h::AdvancedHMC.Hamiltonian{…})
    @ AdvancedHMC C:\Users\joesm\.julia\packages\AdvancedHMC\AlvV4\src\hamiltonian.jl:155
 [21] initialstep(rng::TaskLocalRNG, model::DynamicPPL.Model{…}, spl::DynamicPPL.Sampler{…}, vi_original::DynamicPPL.TypedVarInfo{…}; initial_params::Nothing, nadapts::Int64, kwargs::@Kwargs{})
    @ Turing.Inference C:\Users\joesm\.julia\packages\Turing\r3Hmj\src\mcmc\hmc.jl:184
 [22] step(rng::TaskLocalRNG, model::DynamicPPL.Model{…}, spl::DynamicPPL.Sampler{…}; initial_params::Nothing, kwargs::@Kwargs{…})
    @ DynamicPPL C:\Users\joesm\.julia\packages\DynamicPPL\93t4P\src\sampler.jl:116
 [23] step
    @ C:\Users\joesm\.julia\packages\DynamicPPL\93t4P\src\sampler.jl:99 [inlined]
 [24] macro expansion
    @ C:\Users\joesm\.julia\packages\AbstractMCMC\YrmkI\src\sample.jl:130 [inlined]
 [25] macro expansion
    @ C:\Users\joesm\.julia\packages\ProgressLogging\6KXlp\src\ProgressLogging.jl:328 [inlined]
 [26] macro expansion
    @ C:\Users\joesm\.julia\packages\AbstractMCMC\YrmkI\src\logging.jl:9 [inlined]
 [27] mcmcsample(rng::TaskLocalRNG, model::DynamicPPL.Model{…}, sampler::DynamicPPL.Sampler{…}, N::Int64; progress::Bool, progressname::String, callback::Nothing, discard_initial::Int64, thinning::Int64, chain_type::Type, initial_state::Nothing, kwargs::@Kwargs{…})
    @ AbstractMCMC C:\Users\joesm\.julia\packages\AbstractMCMC\YrmkI\src\sample.jl:120
 [28] sample(rng::TaskLocalRNG, model::DynamicPPL.Model{…}, sampler::DynamicPPL.Sampler{…}, N::Int64; chain_type::Type, resume_from::Nothing, initial_state::Nothing, progress::Bool,
 nadapts::Int64, discard_adapt::Bool, discard_initial::Int64, kwargs::@Kwargs{})
    @ Turing.Inference C:\Users\joesm\.julia\packages\Turing\r3Hmj\src\mcmc\hmc.jl:123
 [29] sample
    @ C:\Users\joesm\.julia\packages\Turing\r3Hmj\src\mcmc\hmc.jl:92 [inlined]
 [30] #sample#4
    @ C:\Users\joesm\.julia\packages\Turing\r3Hmj\src\mcmc\Inference.jl:274 [inlined]
 [31] sample
    @ C:\Users\joesm\.julia\packages\Turing\r3Hmj\src\mcmc\Inference.jl:265 [inlined]
 [32] #sample#3
    @ C:\Users\joesm\.julia\packages\Turing\r3Hmj\src\mcmc\Inference.jl:262 [inlined]
 [33] top-level scope
    @ c:\Users\joesm\simple_shift\shifted_exponential.jl:123
2

For your needs you can use LocationScale

LocationScale(θ, 1, Exponential(λ))

or even shorter

θ + Exponential(λ)
3

Thank you for your response, this does indeed solve the issue in the original example.

However, what about for a more general custom univariate probability distribution, that is not necessarily an affine transformation of another?

The recipe above should be similar, that is define the distribution, define the sampler, and how to evaluate the log prob, and derivatives etc, but the autodiff still has issues with type conversions.

4

Had also not seen it immediately, but your struct ShiftedExponential{T<:Real} has hardcoded its fields as Float64. Seems your code works fine, when you change it to rate::T and shift::T.

5

Great spot, thanks for all your help.