Add Distributions extension

Project.toml

@@ -4,13 +4,16 @@ version = "0.14.12"
 authors = ["Chad Scherrer <chad.scherrer@gmail.com>", "Oliver Schulz <oschulz@mpp.mpg.de>", "contributors"]
 
 [deps]
+ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
+ArraysOfArrays = "65a8f2f4-9b39-5baf-92e2-a9cc46fdf018"
 ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 ConstantRNGs = "aa9b60e7-6b1c-4c29-a6e5-e43521412437"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 FunctionChains = "8e6b2b91-af83-483e-ba35-d00930e4cf9b"
+HeterogeneousComputing = "2182be2a-124f-4a91-8389-f06db5907a21"
 IfElse = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
 IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
 InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112"
@@ -27,25 +30,47 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+StaticThings = "7e4b4f32-fbf9-4b74-9510-4d15222ac973"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775"
 
 [weakdeps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+ForwardDiffPullbacks = "450a3b6d-2448-4ee1-8e34-e4eb8713b605"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 
 [extensions]
 MeasureBaseChainRulesCoreExt = "ChainRulesCore"
+MeasureBaseDistributionsExt = ["Distributions", "StatsBase", "StatsFuns", "PDMats"]
+MeasureBaseDistributionsChainRulesCoreExt = ["Distributions", "ChainRulesCore"]
+MeasureBaseDistributionsForwardDiffExt = ["Distributions", "ForwardDiff"]
+MeasureBaseDistributionsForwardDiffPullbacksExt = ["Distributions", "ForwardDiffPullbacks", "ChainRulesCore"]
+MeasureBaseDistributionsMooncakeExt = ["Distributions", "Mooncake"]
+MeasureBaseForwardDiffExt = "ForwardDiff"
+MeasureBaseForwardDiffPullbacksExt = "ForwardDiffPullbacks"
+MeasureBaseMooncakeExt = "Mooncake"
 
 [compat]
+ArgCheck = "1, 2"
+ArraysOfArrays = "0.6"
 ChainRulesCore = "1"
 ChangesOfVariables = "0.1.3"
 Compat = "3.35, 4"
 ConstantRNGs = "0.1.1"
 ConstructionBase = "1.3"
 DensityInterface = "0.4"
+Distributions = "0.25.1"
 FillArrays = "0.12, 0.13, 1"
-FunctionChains = "0.2"
+ForwardDiff = "0.10, 1"
+ForwardDiffPullbacks = "0.2"
+FunctionChains = "0.2.3"
+HeterogeneousComputing = "0.2.3"
 IfElse = "0.1"
 IntervalSets = "0.7"
 InverseFunctions = "0.1.8"
@@ -54,15 +79,20 @@ LinearAlgebra = "1"
 LogExpFunctions = "0.3, 1"
 LogarithmicNumbers = "1"
 MappedArrays = "0.4"
+Mooncake = "0.5.34"
 NaNMath = "0.3, 1"
+PDMats = "0.11"
 PrettyPrinting = "0.3, 0.4"
 PropertyFunctions = "0.2.2"
 Random = "1"
 Reexport = "1"
-SpecialFunctions = "2"
+SpecialFunctions = "2.1.4"
 Static = "0.8, 1"
 StaticArrays = "1.5"
+StaticThings = "0.2"
 Statistics = "1"
+StatsBase = "0.33, 0.34"
+StatsFuns = "0.9, 1, 2"
 Test = "1"
 Tricks = "0.1"
 julia = "1.10"

ext/MeasureBaseChainRulesCoreExt.jl

@@ -11,11 +11,66 @@ import ChainRulesCore
 using MeasureBase: isneginf, isposinf
 
 _isneginf_pullback(::Any) = (NoTangent(), ZeroTangent())
-ChainRulesCore.rrule(::typeof(isneginf), x) = isneginf(x), _logdensityof_rt_pullback
+ChainRulesCore.rrule(::typeof(isneginf), x) = isneginf(x), _isneginf_pullback
 
 _isposinf_pullback(::Any) = (NoTangent(), ZeroTangent())
 ChainRulesCore.rrule(::typeof(isposinf), x) = isposinf(x), _isposinf_pullback
 
+using MeasureBase: _adignore_call
+
+@inline _adignore_call_pullback(@nospecialize ΔΩ) = (NoTangent(), NoTangent())
+ChainRulesCore.rrule(::typeof(_adignore_call), f) = _adignore_call(f), _adignore_call_pullback
+
+using MeasureBase: convert_realtype
+
+_convert_realtype_pullback(ΔΩ) = NoTangent(), NoTangent(), ΔΩ
+ChainRulesCore.rrule(::typeof(convert_realtype), ::Type{T}, x) where {T} =
+    convert_realtype(T, x), _convert_realtype_pullback
+
+# = collection utils =========================================================
+
+using MeasureBase: _pushfront, _pushback, _rev_cumsum, _exp_cumsum_log
+
+function ChainRulesCore.rrule(::typeof(_pushfront), v::AbstractVector, x)
+    result = _pushfront(v, x)
+    function _pushfront_pullback(thunked_ΔΩ)
+        ΔΩ = ChainRulesCore.unthunk(thunked_ΔΩ)
+        (NoTangent(), ΔΩ[firstindex(ΔΩ)+1:lastindex(ΔΩ)], ΔΩ[firstindex(ΔΩ)])
+    end
+    return result, _pushfront_pullback
+end
+
+
+function ChainRulesCore.rrule(::typeof(_pushback), v::AbstractVector, x)
+    result = _pushback(v, x)
+    function _pushback_pullback(thunked_ΔΩ)
+        ΔΩ = ChainRulesCore.unthunk(thunked_ΔΩ)
+        (NoTangent(), ΔΩ[firstindex(ΔΩ):lastindex(ΔΩ)-1], ΔΩ[lastindex(ΔΩ)])
+    end
+    return result, _pushback_pullback
+end
+
+
+function ChainRulesCore.rrule(::typeof(_rev_cumsum), xs::AbstractVector)
+    result = _rev_cumsum(xs)
+    function _rev_cumsum_pullback(ΔΩ)
+        ∂xs = ChainRulesCore.@thunk cumsum(ChainRulesCore.unthunk(ΔΩ))
+        (NoTangent(), ∂xs)
+    end
+    return result, _rev_cumsum_pullback
+end
+
+
+function ChainRulesCore.rrule(::typeof(_exp_cumsum_log), xs::AbstractVector)
+    result = _exp_cumsum_log(xs)
+    function _exp_cumsum_log_pullback(ΔΩ)
+        ∂xs = inv.(xs) .* _rev_cumsum(exp.(cumsum(log.(xs))) .* ChainRulesCore.unthunk(ΔΩ))
+        (NoTangent(), ∂xs)
+    end
+    return result, _exp_cumsum_log_pullback
+end
+
+
 # = insupport & friends ======================================================
 
 using MeasureBase: check_dof, require_insupport, checked_arg, _checksupport, _origin_depth

ext/MeasureBaseDistributionsChainRulesCoreExt.jl

@@ -0,0 +1,18 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+module MeasureBaseDistributionsChainRulesCoreExt
+
+using MeasureBase
+import Distributions
+import ChainRulesCore
+using ChainRulesCore: NoTangent
+
+using MeasureBase: _dist_params_numtype
+using Distributions: Distribution
+
+_dist_params_numtype_pullback(ΔΩ) = (NoTangent(), NoTangent())
+function ChainRulesCore.rrule(::typeof(_dist_params_numtype), d::Distribution)
+    _dist_params_numtype(d), _dist_params_numtype_pullback
+end
+
+end # module MeasureBaseDistributionsChainRulesCoreExt

ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl

@@ -0,0 +1,66 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+module MeasureBaseDistributionsExt
+
+using LinearAlgebra: Diagonal, diag, dot, cholesky
+
+import Random
+using Random: AbstractRNG, rand!
+
+import DensityInterface
+using DensityInterface: logdensityof, densityof
+
+import MeasureBase
+using MeasureBase: AbstractMeasure, AsMeasure, asmeasure
+using MeasureBase: Lebesgue, Counting, ℝ
+using MeasureBase: StdMeasure, StdUniform, StdExponential, StdLogistic, StdNormal
+using MeasureBase: PowerMeasure, WeightedMeasure, SuperpositionMeasure, PushforwardMeasure
+using MeasureBase: basemeasure, rootmeasure, testvalue, productmeasure, pushfwd, superpose
+using MeasureBase: getdof, checked_arg, massof
+using MeasureBase: transport_to, transport_def, transport_origin, from_origin, to_origin
+using MeasureBase: NoTransportOrigin, NoTransport
+using MeasureBase: Reshape
+using MeasureBase: convert_realtype, firsttype, _fwddiff, @_adignore
+import MeasureBase:
+    _dist_params_numtype, _trafo_cdf_impl, _trafo_quantile_impl, _trafo_quantile_impl_generic
+using MeasureBase: _pushfront, _pushback, _dropfront, _dropback, _rev_cumsum, _exp_cumsum_log
+
+import Distributions
+using Distributions: Distribution, VariateForm, ValueSupport, ContinuousDistribution
+using Distributions: Univariate, Multivariate, ArrayLikeVariate, Continuous, Discrete
+using Distributions: Uniform, Exponential, Logistic, Normal
+using Distributions: MvNormal, AbstractMvNormal, Beta, Dirichlet
+using Distributions: ReshapedDistribution, AbstractMixtureModel
+
+import Statistics
+import StatsBase
+import StatsFuns
+import PDMats
+
+using IrrationalConstants: log2π, invsqrt2π
+
+using HeterogeneousComputing: real_numtype
+
+using Static: True, False, StaticInt, static, dynamic
+using StaticThings: asnonstatic
+using FillArrays: Fill, Ones, Zeros
+
+using ArgCheck: @argcheck
+
+using ArraysOfArrays: ArrayOfSimilarArrays, flatview
+
+include("measure_interface.jl")
+include("standard_dist.jl")
+include("standard_uniform.jl")
+include("standard_normal.jl")
+include("distribution_measure.jl")
+include("dist_vartransform.jl")
+include("univariate.jl")
+include("standardmv.jl")
+include("product.jl")
+include("reshaped.jl")
+include("mixture.jl")
+include("dirichlet.jl")
+include("dirac.jl")
+
+end # module MeasureBaseDistributionsExt

ext/MeasureBaseDistributionsExt/dirac.jl

@@ -0,0 +1,14 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+MeasureBase.AbstractMeasure(obj::Distributions.Dirac) = MeasureBase.Dirac(obj.value)
+
+function AsMeasure{D}(::D) where {D<:Distributions.Dirac}
+    throw(ArgumentError("Don't wrap Distributions.Dirac into MeasureBase.AsMeasure, use asmeasure to convert instead."))
+end
+
+
+Distributions.Distribution(m::MeasureBase.Dirac{<:Real}) = Distributions.Dirac(m.x)
+
+function Distributions.Distribution(@nospecialize(m::MeasureBase.Dirac{T})) where T
+    throw(ArgumentError("Can only convert MeasureBase.Dirac{<:Real} to Distributions.Dirac, but not MeasureBase.Dirac{<:$(nameof(T))}"))
+end

ext/MeasureBaseDistributionsExt/dirichlet.jl

@@ -0,0 +1,59 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+const DirichletMeasure = AsMeasure{<:Dirichlet}
+
+MeasureBase.getdof(d::Dirichlet) = length(d) - 1
+MeasureBase.getdof(m::DirichletMeasure) = getdof(m.obj)
+
+MeasureBase.transport_origin(d::Dirichlet) = StdUniform()^getdof(d)
+
+
+
+function _dirichlet_beta_trafo(α::Real, β::Real, x::Real)
+    R = float(promote_type(typeof(α), typeof(β), typeof(x)))
+    convert(R, transport_def(Beta(α, β), StdUniform(), x))::R
+end
+
+_a_times_one_minus_b(a::Real, b::Real) = a * (1 - b)
+
+function MeasureBase.from_origin(ν::Dirichlet, x)
+    # See M. J. Betancourt, "Cruising The Simplex: Hamiltonian Monte Carlo and the Dirichlet Distribution",
+    # https://arxiv.org/abs/1010.3436
+
+    @_adignore @argcheck length(ν) == length(x) + 1
+
+    αs = _dropfront(_rev_cumsum(ν.alpha))
+    βs = _dropback(ν.alpha)
+    beta_v = _fwddiff(_dirichlet_beta_trafo).(αs, βs, x)
+    beta_v_cp = _exp_cumsum_log(_pushfront(beta_v, 1))
+    beta_v_ext = _pushback(beta_v, 0)
+    _fwddiff(_a_times_one_minus_b).(beta_v_cp, beta_v_ext)
+end
+
+
+function _inv_dirichlet_beta_trafo(α::Real, β::Real, beta_v::Real)
+    R = float(promote_type(typeof(α), typeof(β), typeof(beta_v)))
+    convert(R, transport_def(StdUniform(), Beta(α, β), beta_v))::R
+end
+
+# ToDo: Find efficient pullback for this:
+function _dirichlet_variate_to_beta_v(y::AbstractVector{<:Real})
+    beta_v = similar(y, length(eachindex(y)) - 1)
+    @assert firstindex(beta_v) == firstindex(y)
+    @assert lastindex(beta_v) == lastindex(y) - 1
+    T = eltype(y)
+    sum_log_beta_v::T = 0
+    @inbounds for i in eachindex(beta_v)
+        beta_v[i] = 1 - y[i] / exp(sum_log_beta_v)
+        sum_log_beta_v += log(beta_v[i])
+    end
+    return beta_v
+end
+
+function MeasureBase.to_origin(ν::Dirichlet, y)
+    @_adignore @argcheck length(ν) == length(y)
+    αs = _dropfront(_rev_cumsum(ν.alpha))
+    βs = _dropback(ν.alpha)
+    beta_v = _dirichlet_variate_to_beta_v(y)
+    _fwddiff(_inv_dirichlet_beta_trafo).(αs, βs, beta_v)
+end

ext/MeasureBaseDistributionsExt/dist_vartransform.jl

@@ -0,0 +1,16 @@
+# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT).
+
+const _AnyStdUniform = Union{StandardUniform,Uniform}
+const _AnyStdNormal = Union{StandardNormal,Normal}
+
+const _AnyStdDistribution = Union{_AnyStdUniform,_AnyStdNormal}
+
+_std_dist(::Type{<:_AnyStdUniform}) = StandardUniform
+_std_dist(::Type{<:_AnyStdNormal}) = StandardNormal
+
+_std_dist(::Type{D}, ::StaticInt{1}) where {D<:_AnyStdDistribution} = D()
+_std_dist(::Type{D}, dof) where {D<:_AnyStdDistribution} = D(dynamic(dof))
+_std_dist_for(::Type{D}, μ::Any) where {D<:_AnyStdDistribution} = _std_dist(_std_dist(D), getdof(μ))
+
+MeasureBase.transport_to(::Type{NU}, μ) where {NU<:_AnyStdDistribution} = transport_to(_std_dist_for(NU, μ), μ)
+MeasureBase.transport_to(ν, ::Type{MU}) where {MU<:_AnyStdDistribution} = transport_to(ν, _std_dist_for(MU, ν))