Generically supporting wrapped distributions

[sethaxen]

#1715 proposes adding the Wrapped Normal distribution. Other common named wrapped distributions include:

• Wrapped Poisson
• Wrapped Cauchy
• k-wrapped distributions (i.e. wrapping a circular distribution onto itself to represent k-fold symmetry).

One could support these generically, as well as arbitrary wrappings, with a "wrapper" distribution Wrapped, that could take a distribution with support on the real line and an interval around which to wrap it. e.g.

using Distributions

struct Wrapped{D<:UnivariateDistribution, S<:ValueSupport, T <: Real} <: UnivariateDistribution{S}
    unwrapped::D
    lower::T      # lower bound
    upper::T      # upper bound
    k::Int
    function Wrapped(d::UnivariateDistribution, l::T, u::T, k::Int=1) where {T <: Real}
        new{typeof(d), Distributions.value_support(typeof(d)), T}(d, l, u, k)
    end
end

wrapped(d::UnivariateDistribution, l::T, u::T, k::Int=1) where {T<:Real} = Wrapped(d, l, u, k)
wrapped(d::UnivariateDistribution, l::Real, u::Real, k::Int=1) = Wrapped(d, Base.promote(l, u)..., k)

Base.minimum(d::Wrapped) = d.lower
Base.maximum(d::Wrapped) = d.upper

function Distributions.pdf(d::Wrapped, x::Real; tol=sqrt(eps(float(typeof(x)))), maxiter=1_000)
    u = d.upper
    l = d.lower
    period = u - l
    z = mod(x, period)
    p = pdf(d.unwrapped, z)
    for i in 1:maxiter
        r = i // d.k
        zl = muladd(-r, period, z)
        zu = muladd(r, period, z)
        pl = insupport(d.unwrapped, zl) ? pdf(d.unwrapped, zl) : zero(p)
        pu = insupport(d.unwrapped, zu) ? pdf(d.unwrapped, zu) : zero(p)
        Δp = pl + pu
        p += Δp
        abs(Δp) < p * tol && break
    end
    return p / d.k
end

Advantages

• resembling what we currently do with truncated and censored
• supporting wrapping to arbitrary periods, including integer periods
• still supporting dispatch for cases like WrappedNormal where we have more efficient ways to evaluate the series

Disadvantages:

• In some cases, like WrappedNormal in #1715, depending on the parameters of unwrapped, we might at construction time be able to compute some heuristic to decide which method to use e.g. for density evaluation, which the above approach would require us do during every density evaluation. One solution would be to add some field of arbitrary type to Wrapped that could in an overloaded constructor be filled with any relevant precomputations that would be used in pdf, etc.

An example usage

using StatsPlots

fig = plot(
    plot(wrapped(Normal(π/3, 1), -π, π); title="Wrapped Normal", proj=:polar),
    plot(wrapped(Cauchy(π/3, 1), 0, 2π); title="Wrapped Cauchy", proj=:polar),
    plot(wrapped(Uniform(-1, 1), -π, π, 8); title="8-times Wrapped Uniform", proj=:polar),
    plot(wrapped(Poisson(10), 0, 15); title="Wrapped Poisson");
    label="", dpi=180, titlefontsize=10
)

adapted from original post by @sethaxen in https://github.com/JuliaStats/Distributions.jl/issues/1715#issuecomment-1536129396

[sethaxen]

This also makes it easy to create axial distributions from directional distributions:

plot(wrapped(VonMises(π/4, 10), -π, π, 2); proj=:polar, label="")

[sethaxen]

In some cases, like WrappedNormal in https://github.com/JuliaStats/Distributions.jl/issues/1715, depending on the parameters of unwrapped, we might at construction time be able to compute some heuristic to decide which method to use e.g. for density evaluation, which the above approach would require us do during every density evaluation. One solution would be to add some field of arbitrary type to Wrapped that could in an overloaded constructor be filled with any relevant precomputations that would be used in pdf, etc.

After more thought, this can also be handled generically. For some arbitrary wrapped distribution, we can compute the density by:

• for narrow unwrapped distributions: explicitly sum the density along intervals. Few terms are needed for convergence.
• for broad unwrapped distributions: use the Fourier series representation, where the terms of the series are the characteristic function of the unwrapped distribution. Again, few terms are needed for convergence.

A field use_characteristic::Bool, defaulting to False, would specify which approach should be used. One can then overload wrapped(::MyDist, l, u) to specify that the characteristic function should instead be used, which can in special cases like WrappedNormal be chosen by the parameters.