diff --git a/Project.toml b/Project.toml index 44f89f80..e2bfee4a 100644 --- a/Project.toml +++ b/Project.toml @@ -4,6 +4,8 @@ version = "0.14.12" authors = ["Chad Scherrer ", "Oliver Schulz ", "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" @@ -11,6 +13,7 @@ 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" diff --git a/ext/MeasureBaseChainRulesCoreExt.jl b/ext/MeasureBaseChainRulesCoreExt.jl index 57ed25fa..25019da1 100644 --- a/ext/MeasureBaseChainRulesCoreExt.jl +++ b/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 diff --git a/ext/MeasureBaseDistributionsChainRulesCoreExt.jl b/ext/MeasureBaseDistributionsChainRulesCoreExt.jl new file mode 100644 index 00000000..63cc3e93 --- /dev/null +++ b/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 diff --git a/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl b/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl new file mode 100644 index 00000000..177ec306 --- /dev/null +++ b/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 diff --git a/ext/MeasureBaseDistributionsExt/dirac.jl b/ext/MeasureBaseDistributionsExt/dirac.jl new file mode 100644 index 00000000..8580df8c --- /dev/null +++ b/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 diff --git a/ext/MeasureBaseDistributionsExt/dirichlet.jl b/ext/MeasureBaseDistributionsExt/dirichlet.jl new file mode 100644 index 00000000..c60eeecd --- /dev/null +++ b/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 diff --git a/ext/MeasureBaseDistributionsExt/dist_vartransform.jl b/ext/MeasureBaseDistributionsExt/dist_vartransform.jl new file mode 100644 index 00000000..ceedabe9 --- /dev/null +++ b/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, ν)) diff --git a/ext/MeasureBaseDistributionsExt/distribution_measure.jl b/ext/MeasureBaseDistributionsExt/distribution_measure.jl new file mode 100644 index 00000000..bcbfd558 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/distribution_measure.jl @@ -0,0 +1,74 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + + +const DistributionMeasure{F<:VariateForm,S<:ValueSupport,D<:Distribution{F,S}} = AsMeasure{D} + +@inline MeasureBase.AbstractMeasure(obj::Distribution) = AsMeasure{typeof(obj)}(obj) +@inline Base.convert(::Type{AbstractMeasure}, obj::Distribution) = AbstractMeasure(obj) + +@inline Distributions.Distribution(m::DistributionMeasure) = m.obj +@inline Distributions.Distribution{F}(m::DistributionMeasure{F}) where {F<:VariateForm} = Distribution(m) +@inline Distributions.Distribution{F,S}(m::DistributionMeasure{F,S}) where {F<:VariateForm,S<:ValueSupport} = Distribution(m) + +@inline Base.convert(::Type{Distribution}, m::DistributionMeasure) = Distribution(m) +@inline Base.convert(::Type{Distribution{F}}, m::DistributionMeasure{F}) where {F<:VariateForm} = Distribution(m) +@inline Base.convert(::Type{Distribution{F,S}}, m::DistributionMeasure{F,S}) where {F<:VariateForm,S<:ValueSupport} = Distribution(m) + + +Base.rand(rng::AbstractRNG, ::Type{T}, m::DistributionMeasure) where {T<:Real} = convert_realtype(T, rand(m.obj)) + +function _flat_powrand(rng::AbstractRNG, ::Type{T}, d::Distribution{<:ArrayLikeVariate{0}}, sz::Dims) where {T<:Real} + convert_realtype(T, reshape(rand(rng, d, prod(sz)), sz...)) +end + +function _flat_powrand(rng::AbstractRNG, ::Type{T}, d::Distribution{<:ArrayLikeVariate{1}}, sz::Dims) where {T<:Real} + convert_realtype(T, reshape(rand(rng, d, prod(sz)), size(d)..., sz...)) +end + +function _flat_powrand(rng::AbstractRNG, ::Type{T}, d::ReshapedDistribution{N,<:Any,<:Distribution{<:ArrayLikeVariate{1}}}, sz::Dims) where {T<:Real,N} + convert_realtype(T, reshape(rand(rng, d.dist, prod(sz)), d.dims..., sz...)) +end + +function _flat_powrand(rng::AbstractRNG, ::Type{T}, d::Distribution, sz::Dims) where {T<:Real} + flatview(ArrayOfSimilarArrays(convert_realtype(T, rand(rng, d, sz)))) +end + +function Base.rand(rng::AbstractRNG, ::Type{T}, m::PowerMeasure{<:DistributionMeasure{<:ArrayLikeVariate{0}}, NTuple{N,Base.OneTo{Int}}}) where {T<:Real,N} + _flat_powrand(rng, T, m.parent.obj, map(length, m.axes)) +end + +function Base.rand(rng::AbstractRNG, ::Type{T}, m::PowerMeasure{<:DistributionMeasure{<:ArrayLikeVariate{M}}, NTuple{N,Base.OneTo{Int}}}) where {T<:Real,M,N} + flat_data = _flat_powrand(rng, T, m.parent.obj, map(length, m.axes)) + ArrayOfSimilarArrays{T,M,N}(flat_data) +end + + +@inline DensityInterface.densityof(m::DistributionMeasure) = densityof(m.obj) +@inline DensityInterface.logdensityof(m::DistributionMeasure) = logdensityof(m.obj) + +@inline MeasureBase.logdensity_def(m::DistributionMeasure, x) = DensityInterface.logdensityof(m.obj, x) +@inline MeasureBase.unsafe_logdensityof(m::DistributionMeasure, x) = DensityInterface.logdensityof(m.obj, x) +@inline MeasureBase.insupport(m::DistributionMeasure, x) = Distributions.insupport(m.obj, x) + +@inline MeasureBase.rootmeasure(m::DistributionMeasure{<:ArrayLikeVariate{0},<:Continuous}) = Lebesgue() +@inline MeasureBase.rootmeasure(m::DistributionMeasure{<:ArrayLikeVariate,<:Continuous}) = Lebesgue()^size(m.obj) +@inline MeasureBase.rootmeasure(m::DistributionMeasure{<:ArrayLikeVariate{0},<:Discrete}) = Counting() +@inline MeasureBase.rootmeasure(m::DistributionMeasure{<:ArrayLikeVariate,<:Discrete}) = Counting()^size(m.obj) + +@inline MeasureBase.basemeasure(m::DistributionMeasure) = rootmeasure(m) + +@inline MeasureBase.massof(::DistributionMeasure) = static(1.0) + +@inline MeasureBase.mspace_elsize(m::DistributionMeasure{<:ArrayLikeVariate}) = size(m.obj) + +@inline MeasureBase.getdof(m::DistributionMeasure{<:ArrayLikeVariate{0}}) = 1 + +# Delegate transport to the wrapped distribution: +@inline MeasureBase.transport_origin(m::DistributionMeasure) = m.obj +@inline MeasureBase.to_origin(::DistributionMeasure, y) = y +@inline MeasureBase.from_origin(::DistributionMeasure, x) = x + +@inline MeasureBase.paramnames(m::DistributionMeasure) = propertynames(m.obj) +@inline MeasureBase.params(m::DistributionMeasure) = NamedTuple{propertynames(m.obj)}(Distributions.params(m.obj)) + +# @inline MeasureBase.testvalue(m::DistributionMeasure) = testvalue(basemeasure(d)) diff --git a/ext/MeasureBaseDistributionsExt/measure_interface.jl b/ext/MeasureBaseDistributionsExt/measure_interface.jl new file mode 100644 index 00000000..6fed5d4a --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/measure_interface.jl @@ -0,0 +1,26 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +@inline MeasureBase.logdensity_def(d::Distribution, x) = DensityInterface.logdensityof(d, x) +@inline MeasureBase.unsafe_logdensityof(d::Distribution, x) = DensityInterface.logdensityof(d, x) + +@inline MeasureBase.insupport(d::Distribution, x) = Distributions.insupport(d, x) + +@inline MeasureBase.basemeasure(d::Distribution{<:ArrayLikeVariate{0},<:Continuous}) = Lebesgue() +@inline MeasureBase.basemeasure(d::Distribution{<:ArrayLikeVariate,<:Continuous}) = Lebesgue()^size(d) +@inline MeasureBase.basemeasure(d::Distribution{<:ArrayLikeVariate{0},<:Discrete}) = Counting() +@inline MeasureBase.basemeasure(d::Distribution{<:ArrayLikeVariate,<:Discrete}) = Counting()^size(d) + +@inline MeasureBase.paramnames(d::Distribution) = propertynames(d) +@inline MeasureBase.params(d::Distribution) = NamedTuple{propertynames(d)}(Distributions.params(d)) + +@inline MeasureBase.testvalue(d::Distribution) = testvalue(basemeasure(d)) +@inline MeasureBase.testvalue(::Type{T}, d::Distribution) where {T} = testvalue(T, basemeasure(d)) + + +@inline MeasureBase.basemeasure(d::Distributions.Poisson) = + Counting(MeasureBase.BoundedInts(static(0), static(Inf))) +@inline MeasureBase.basemeasure(d::Distributions.Product{<:Any,<:Distributions.Poisson}) = + Counting(MeasureBase.BoundedInts(static(0), static(Inf)))^size(d) + + +MeasureBase.∫(f, base::Distribution) = MeasureBase.∫(f, convert(AbstractMeasure, base)) diff --git a/ext/MeasureBaseDistributionsExt/mixture.jl b/ext/MeasureBaseDistributionsExt/mixture.jl new file mode 100644 index 00000000..89d41acf --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/mixture.jl @@ -0,0 +1,24 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +function MeasureBase.AbstractMeasure(d::Distributions.AbstractMixtureModel) + superpose(map((w, c) -> w * asmeasure(c), Distributions.probs(d), Distributions.components(d))) +end + +function AsMeasure{D}(::D) where {D<:Distributions.AbstractMixtureModel} + throw(ArgumentError("Don't wrap Distributions.AbstractMixtureModel into MeasureBase.AsMeasure, use asmeasure to convert instead.")) +end + + +const _MixtureMeasure = SuperpositionMeasure{ + <:Union{Tuple{Vararg{WeightedMeasure}},AbstractVector{<:WeightedMeasure}}, +} + +_mixture_component(m::AsMeasure{<:Distribution}) = m.obj + +function Distributions.Distribution(m::_MixtureMeasure) + components = map(c -> _mixture_component(c.base), collect(values(m.components))) + prior = map(c -> exp(c.logweight), collect(values(m.components))) + Distributions.MixtureModel(components, prior) +end + +Base.convert(::Type{Distribution}, m::_MixtureMeasure) = Distributions.Distribution(m) diff --git a/ext/MeasureBaseDistributionsExt/product.jl b/ext/MeasureBaseDistributionsExt/product.jl new file mode 100644 index 00000000..a050dc97 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/product.jl @@ -0,0 +1,56 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +@static if isdefined(Distributions, :Product) + MeasureBase.AbstractMeasure(obj::Distributions.Product) = productmeasure(map(asmeasure, obj.v)) + + function AsMeasure{D}(::D) where {D<:Distributions.Product} + throw(ArgumentError("Don't wrap Distributions.Product into MeasureBase.AsMeasure, use asmeasure to convert instead.")) + end +end + +function Distributions.Distribution( + m::MeasureBase.ProductMeasure{<:AbstractArray{<:AsMeasure{<:Distribution{Univariate}}}}, +) + Distributions.product_distribution(map(x -> x.obj, MeasureBase.marginals(m))) +end + +function Base.convert( + ::Type{Distribution}, + m::MeasureBase.ProductMeasure{<:AbstractArray{<:AsMeasure{<:Distribution{Univariate}}}}, +) + Distributions.Distribution(m) +end + +@static if isdefined(Distributions, :ProductDistribution) + MeasureBase.AbstractMeasure(obj::Distributions.ProductDistribution) = productmeasure(map(asmeasure, obj.dists)) + + function AsMeasure{D}(::D) where {D<:Distributions.ProductDistribution} + throw(ArgumentError("Don't wrap Distributions.ProductDistribution into MeasureBase.AsMeasure, use asmeasure to convert instead.")) + end + + function Distributions.Distribution( + m::MeasureBase.ProductMeasure{<:AbstractArray{<:AsMeasure{<:Distribution}}}, + ) + Distributions.product_distribution(map(x -> x.obj, MeasureBase.marginals(m))) + end + + function Distributions.Distribution( + m::MeasureBase.ProductMeasure{<:Tuple{Vararg{AsMeasure{<:Distribution}}}}, + ) + Distributions.product_distribution(map(x -> x.obj, MeasureBase.marginals(m))...) + end + + function Base.convert( + ::Type{Distribution}, + m::MeasureBase.ProductMeasure{<:AbstractArray{<:AsMeasure{<:Distribution}}}, + ) + Distributions.Distribution(m) + end + + function Base.convert( + ::Type{Distribution}, + m::MeasureBase.ProductMeasure{<:Tuple{Vararg{AsMeasure{<:Distribution}}}}, + ) + Distributions.Distribution(m) + end +end diff --git a/ext/MeasureBaseDistributionsExt/reshaped.jl b/ext/MeasureBaseDistributionsExt/reshaped.jl new file mode 100644 index 00000000..6efde609 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/reshaped.jl @@ -0,0 +1,27 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +MeasureBase.getdof(μ::ReshapedDistribution) = MeasureBase.getdof(μ.dist) + +MeasureBase.transport_origin(μ::ReshapedDistribution) = μ.dist + +MeasureBase.to_origin(ν::ReshapedDistribution, y) = reshape(y, size(ν.dist)) + +MeasureBase.from_origin(ν::ReshapedDistribution, x) = reshape(x, ν.dims) + + +function MeasureBase.AbstractMeasure(d::Distributions.ReshapedDistribution) + orig_dist = d.dist + pushfwd(Reshape(size(d), size(orig_dist)), AbstractMeasure(orig_dist)) +end + +function AsMeasure{D}(::D) where {D<:Distributions.ReshapedDistribution} + throw(ArgumentError("Don't wrap Distributions.ReshapedDistribution into MeasureBase.AsMeasure, use asmeasure to convert instead.")) +end + + +function Distributions.Distribution(m::PushforwardMeasure{<:Reshape}) + reshape(Distributions.Distribution(m.origin), asnonstatic(m.f.output_size)...) +end + +Base.convert(::Type{Distribution}, m::PushforwardMeasure{<:Reshape}) = + Distributions.Distribution(m) diff --git a/ext/MeasureBaseDistributionsExt/standard_dist.jl b/ext/MeasureBaseDistributionsExt/standard_dist.jl new file mode 100644 index 00000000..010b20a9 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/standard_dist.jl @@ -0,0 +1,222 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +""" + struct StandardDist{D<:Distribution{Univariate,Continuous},N} <: Distributions.Distribution{ArrayLikeVariate{N},Continuous} + +Represents `D()` or a product distribution of `D()` in a dispatchable fashion. + +Constructor: +``` + StandardDist{Uniform}(size...) + StandardDist{Normal}(size...) +``` +""" +struct StandardDist{D<:Distribution{Univariate,Continuous},N,U<:Integer} <: + Distributions.Distribution{ArrayLikeVariate{N},Continuous} + _size::NTuple{N,U} +end +export StandardDist + +StandardDist{D}() where {D<:Distribution{Univariate,Continuous}} = + StandardDist{D,0,Int}(()) +StandardDist{D}(dims::Vararg{U,N}) where {D<:Distribution{Univariate,Continuous},N,U<:Integer} = + StandardDist{D,N,U}((dims...,)) + + +const StandardUnivariateDist{D<:Distribution{Univariate,Continuous},U<:Integer} = StandardDist{D,0,U} +const StandardMultivariteDist{D<:Distribution{Univariate,Continuous},U<:Integer} = StandardDist{D,1,U} + + +function Base.show(io::IO, d::StandardDist{D}) where {D} + print(io, nameof(typeof(d)), "{", D, "}") + show(io, d._size) +end + + +@inline MeasureBase.transport_def(::MU, μ::MU, x) where {MU<:StandardDist{<:Any,0}} = x + +for (A, B) in [ + (Uniform, StdUniform), + (Exponential, StdExponential), + (Logistic, StdLogistic), + (Normal, StdNormal) +] + @eval begin + @inline MeasureBase.transport_origin(d::StandardDist{$A,0}) = $B() + @inline MeasureBase.transport_origin(d::StandardDist{$A,N}) where {N} = $B()^size(d) + + # StandardDist{$A} and $B are equivalent as measures, so convert + # instead of wrapping: + MeasureBase.AbstractMeasure(::StandardDist{$A,0}) = $B() + MeasureBase.AbstractMeasure(d::StandardDist{$A,N}) where {N} = $B()^size(d) + + Distributions.Distribution(::$B) = StandardDist{$A}() + Base.convert(::Type{Distribution}, ::$B) = StandardDist{$A}() + + function Distributions.Distribution(m::PowerMeasure{$B}) + StandardDist{$A}(map(dynamic ∘ length, m.axes)...) + end + Base.convert(::Type{Distribution}, m::PowerMeasure{$B}) = Distributions.Distribution(m) + end +end + +@inline MeasureBase.to_origin(ν::StandardDist, y) = y +@inline MeasureBase.from_origin(ν::StandardDist, x) = x + + +@inline nonstddist(::StandardDist{D,0}) where {D} = D(Distributions.params(D())...) +@inline function nonstddist(d::StandardDist{D,N}) where {D,N} + nonstd0 = nonstddist(StandardDist{D}()) + reshape(Distributions.product_distribution(fill(nonstd0, length(d))), size(d)) +end + + +(::Type{D})(d::StandardDist{D,0}) where {D<:Distribution{Univariate,Continuous}} = nonstddist(d) + +# TODO: Replace `fill` by `FillArrays.Fill` once Distributions fully supports this: +(::Type{Distributions.Product})(d::StandardDist{D,1}) where {D} = + Distributions.Product(fill(StandardDist{D}(), length(d))) + +Base.convert(::Type{D}, d::StandardDist{D,0}) where {D<:Distribution{Univariate,Continuous}} = D(d) +Base.convert(::Type{Distributions.Product}, d::StandardDist{D,1}) where {D} = + Distributions.Product(d) + + + +@inline Base.size(d::StandardDist) = d._size +@inline Base.length(d::StandardDist) = prod(size(d)) + +Base.eltype(::Type{<:StandardDist}) = Float64 + +@inline Distributions.partype(d::StandardDist{D}) where {D} = Float64 + +@inline StatsBase.params(d::StandardDist) = () + +for f in ( + :(Base.minimum), + :(Base.maximum), + :(Statistics.mean), + :(Statistics.median), + :(StatsBase.mode), + :(Statistics.var), + :(Statistics.std), + :(StatsBase.skewness), + :(StatsBase.kurtosis), + :(Distributions.location), + :(Distributions.scale), +) + @eval begin + ($f)(d::StandardDist{D,0}) where {D} = ($f)(nonstddist(d)) + ($f)(d::StandardDist{D,N}) where {D,N} = Fill(($f)(StandardDist{D}()), size(d)...) + end +end + +StatsBase.modes(d::StandardDist) = [StatsBase.mode(d)] + +# ToDo: Define cov for N!=1? +Statistics.cov(d::StandardDist{D,1}) where {D} = Diagonal(Statistics.var(d)) +Distributions.invcov(d::StandardDist{D,1}) where {D} = + Diagonal(Fill(inv(Statistics.var(StandardDist{D}())), length(d))) +Distributions.logdetcov(d::StandardDist{D,1}) where {D} = + length(d) * log(Statistics.var(StandardDist{D}())) + +StatsBase.entropy(d::StandardDist{D,0}) where {D} = StatsBase.entropy(nonstddist(d)) +StatsBase.entropy(d::StandardDist{D,N}) where {D,N} = + length(d) * StatsBase.entropy(StandardDist{D}()) + + +Distributions.insupport(d::StandardDist{D,0}, x::Real) where {D} = + Distributions.insupport(nonstddist(d), x) + +function Distributions.insupport(d::StandardDist{D,N}, x::AbstractArray{<:Real,N}) where {D,N} + all(Base.Fix1(Distributions.insupport, StandardDist{D}()), checked_arg(d, x)) +end + + +@inline Distributions.logpdf(d::StandardDist{D,0}, x::U) where {D,U} = + Distributions.logpdf(nonstddist(d), x) + +function Distributions.logpdf(d::StandardDist{D,N}, x::AbstractArray{<:Real,N}) where {D,N} + Distributions._logpdf(d, checked_arg(d, x)) +end + +# Explicit N=1/N=2 methods to avoid dispatch ambiguities with Distributions: +function Distributions._logpdf(::StandardDist{D,1}, x::AbstractArray{<:Real,1}) where {D} + sum(Base.Fix1(Distributions.logpdf, StandardDist{D}()), x) +end + +function Distributions._logpdf(::StandardDist{D,2}, x::AbstractArray{<:Real,2}) where {D} + sum(Base.Fix1(Distributions.logpdf, StandardDist{D}()), x) +end + +function Distributions._logpdf(::StandardDist{D,N}, x::AbstractArray{<:Real,N}) where {D,N} + sum(Base.Fix1(Distributions.logpdf, StandardDist{D}()), x) +end + + +Distributions.gradlogpdf(d::StandardDist{D,0}, x::Real) where {D} = + Distributions.gradlogpdf(nonstddist(d), x) + +function Distributions.gradlogpdf(d::StandardDist{D,N}, x::AbstractArray{<:Real,N}) where {D,N} + Distributions.gradlogpdf.(StandardDist{D}(), checked_arg(d, x)) +end + + +# Explicit N=1/N=2 methods to avoid dispatch ambiguities with Distributions: +function Distributions.pdf(d::StandardDist{D,1}, x::AbstractVector{U}) where {D,U<:Real} + Distributions._pdf(d, checked_arg(d, x)) +end + +function Distributions._pdf(d::StandardDist{D,1}, x::AbstractVector{U}) where {D,U<:Real} + exp(Distributions._logpdf(d, x)) +end + +function Distributions.pdf(d::StandardDist{D,2}, x::AbstractMatrix{U}) where {D,U<:Real} + Distributions._pdf(d, checked_arg(d, x)) +end + +function Distributions._pdf(d::StandardDist{D,2}, x::AbstractMatrix{U}) where {D,U<:Real} + exp(Distributions._logpdf(d, x)) +end + +function Distributions.pdf(d::StandardDist{D,N}, x::AbstractArray{U,N}) where {D,N,U<:Real} + Distributions._pdf(d, checked_arg(d, x)) +end + +function Distributions._pdf(d::StandardDist{D,N}, x::AbstractArray{U,N}) where {D,N,U<:Real} + exp(Distributions._logpdf(d, x)) +end + + +for f in ( + :(Distributions.logcdf), + :(Distributions.cdf), + :(Distributions.logccdf), + :(Distributions.ccdf), + :(Distributions.quantile), + :(Distributions.cquantile), + :(Distributions.invlogcdf), + :(Distributions.invlogccdf), + :(Distributions.mgf), + :(Distributions.cf), +) + @eval begin + @inline ($f)(d::StandardDist, x::Real) = ($f)(nonstddist(d), x) + end +end + + +Base.rand(rng::AbstractRNG, d::StandardDist{D,0}) where {D} = rand(rng, nonstddist(d)) +Random.rand!(rng::AbstractRNG, d::StandardDist{D,0}, x::AbstractArray{<:Real,0}) where {D} = + (x[] = rand(rng, d); return x) +Random.rand!(rng::AbstractRNG, d::StandardDist{D,N}, x::AbstractArray{<:Real,N}) where {D,N} = + rand!(rng, StandardDist{D}(), x) + + +Distributions.truncated(d::StandardDist{D,0}, l::Real, u::Real) where {D} = + Distributions.truncated(nonstddist(d), l, u) + +Distributions.product_distribution(dists::AbstractVector{<:StandardDist{D,0}}) where {D} = + StandardDist{D}(size(dists)...) +Distributions.product_distribution(dists::AbstractArray{<:StandardDist{D,0}}) where {D} = + StandardDist{D}(size(dists)...) diff --git a/ext/MeasureBaseDistributionsExt/standard_normal.jl b/ext/MeasureBaseDistributionsExt/standard_normal.jl new file mode 100644 index 00000000..6bc27d04 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/standard_normal.jl @@ -0,0 +1,77 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +""" + const StandardNormal{N} = StandardDist{Normal,N} + +The standard normal distribution, scalar (`N == 0`) or as a product over an +array of rank `N`. +""" +const StandardNormal{N} = StandardDist{Normal,N} +export StandardNormal + +Distributions.Normal(d::StandardDist{Normal,0}) = Distributions.Normal() + +Distributions.MvNormal(d::StandardDist{Normal,1}) = MvNormal(PDMats.ScalMat(length(d), 1)) +Base.convert(::Type{Distributions.MvNormal}, d::StandardDist{Normal,1}) = + Distributions.MvNormal(d) + +Base.minimum(d::StandardDist{Normal,0}) = -Inf +Base.maximum(d::StandardDist{Normal,0}) = +Inf + +Distributions.insupport(d::StandardDist{Normal,0}, x::Real) = !isnan(x) + +Distributions.location(d::StandardDist{Normal,0}) = Statistics.mean(d) +Distributions.scale(d::StandardDist{Normal,0}) = Statistics.var(d) + +Statistics.mean(d::StandardDist{Normal,0}) = 0 +Statistics.mean(d::StandardDist{Normal,N}) where {N} = Zeros{Int}(size(d)...) + +Statistics.median(d::StandardDist{Normal}) = Statistics.mean(d) +StatsBase.mode(d::StandardDist{Normal}) = Statistics.mean(d) + +StatsBase.modes(d::StandardDist{Normal,0}) = Zeros{Int}(1) + +Statistics.var(d::StandardDist{Normal,0}) = 1 +Statistics.var(d::StandardDist{Normal,N}) where {N} = Ones{Int}(size(d)...) + +Statistics.std(d::StandardDist{Normal,0}) = 1 +Statistics.std(d::StandardDist{Normal,N}) where {N} = Ones{Int}(size(d)...) + +StatsBase.skewness(d::StandardDist{Normal,0}) = 0 +StatsBase.kurtosis(d::StandardDist{Normal,0}) = 0 + +StatsBase.entropy(d::StandardDist{Normal,0}) = muladd(log2π, 1 / 2, 1 / 2) + +Distributions.logpdf(d::StandardDist{Normal,0}, x::U) where {U<:Real} = + muladd(abs2(x), -U(1) / U(2), -log2π / U(2)) +Distributions.pdf(d::StandardDist{Normal,0}, x::U) where {U<:Real} = + invsqrt2π * exp(-abs2(x) / U(2)) + +@inline Distributions.gradlogpdf(d::StandardDist{Normal,0}, x::Real) = -x + +@inline Distributions.logcdf(d::StandardDist{Normal,0}, x::Real) = StatsFuns.normlogcdf(x) +@inline Distributions.cdf(d::StandardDist{Normal,0}, x::Real) = StatsFuns.normcdf(x) +@inline Distributions.logccdf(d::StandardDist{Normal,0}, x::Real) = StatsFuns.normlogccdf(x) +@inline Distributions.ccdf(d::StandardDist{Normal,0}, x::Real) = StatsFuns.normccdf(x) +@inline Distributions.quantile(d::StandardDist{Normal,0}, p::Real) = StatsFuns.norminvcdf(p) +@inline Distributions.cquantile(d::StandardDist{Normal,0}, p::Real) = StatsFuns.norminvccdf(p) +@inline Distributions.invlogcdf(d::StandardDist{Normal,0}, p::Real) = StatsFuns.norminvlogcdf(p) +@inline Distributions.invlogccdf(d::StandardDist{Normal,0}, p::Real) = StatsFuns.norminvlogccdf(p) + +Base.rand(rng::AbstractRNG, d::StandardDist{Normal,0}) = randn(rng) +Base.rand(rng::AbstractRNG, d::StandardDist{Normal,N}) where {N} = randn(rng, size(d)...) +Random.rand!(rng::AbstractRNG, d::StandardDist{Normal,N}, x::AbstractArray{<:Real,N}) where {N} = + Random.randn!(rng, x) + +Distributions.invcov(d::StandardDist{Normal,1}) = Distributions.cov(d) +Distributions.logdetcov(d::StandardDist{Normal,1}) = 0 + + +function Distributions.sqmahal(d::StandardDist{Normal,N}, x::AbstractArray{<:Real,N}) where {N} + dot(x, checked_arg(d, x)) +end + +function Distributions.sqmahal!(r::AbstractVector, d::StandardDist{Normal,N}, x::AbstractMatrix) where {N} + x_cols = eachcol(checked_arg(d, first(eachcol(x)))) + r .= dot.(x_cols, x_cols) +end diff --git a/ext/MeasureBaseDistributionsExt/standard_uniform.jl b/ext/MeasureBaseDistributionsExt/standard_uniform.jl new file mode 100644 index 00000000..51398fcb --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/standard_uniform.jl @@ -0,0 +1,79 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +""" + const StandardUniform{N} = StandardDist{Uniform,N} + +The standard uniform distribution, scalar (`N == 0`) or as a product over an +array of rank `N`. +""" +const StandardUniform{N} = StandardDist{Uniform,N} +export StandardUniform + +Distributions.Uniform(d::StandardDist{Uniform,0}) = Distributions.Uniform() + +Base.minimum(::StandardDist{Uniform,0}) = 0 +Base.maximum(::StandardDist{Uniform,0}) = 1 + +Distributions.location(::StandardDist{Uniform,0}) = 0 +Distributions.scale(::StandardDist{Uniform,0}) = 1 + +Statistics.mean(d::StandardDist{Uniform,0}) = 1 // 2 +Statistics.median(d::StandardDist{Uniform,0}) = Statistics.mean(d) +StatsBase.mode(d::StandardDist{Uniform,0}) = Statistics.mean(d) +StatsBase.modes(d::StandardDist{Uniform,0}) = Zeros{Int}(0) +StatsBase.modes(d::StandardDist{Uniform,N}) where {N} = Fill(Zeros{Int}(size(d))) + +Statistics.var(d::StandardDist{Uniform,0}) = 1 // 12 +Statistics.std(d::StandardDist{Uniform,0}) = sqrt(Statistics.var(d)) +StatsBase.skewness(d::StandardDist{Uniform,0}) = 0 +StatsBase.kurtosis(d::StandardDist{Uniform,0}) = -6 // 5 + +StatsBase.entropy(d::StandardDist{Uniform,0}) = 0 + + +function Distributions.logpdf(d::StandardDist{Uniform,0}, x::U) where {U<:Real} + ifelse(Distributions.insupport(d, x), U(0), U(-Inf)) +end + +function Distributions.pdf(d::StandardDist{Uniform,0}, x::U) where {U<:Real} + ifelse(Distributions.insupport(d, x), one(U), zero(U)) +end + + +Distributions.logcdf(d::StandardDist{Uniform,0}, x::U) where {U<:Real} = + log(Distributions.cdf(d, x)) + +function Distributions.cdf(d::StandardDist{Uniform,0}, x::U) where {U<:Real} + ifelse(x < zero(U), zero(U), ifelse(x < one(U), x, one(U))) +end + +Distributions.logccdf(d::StandardDist{Uniform,0}, x::U) where {U<:Real} = + log(Distributions.ccdf(d, x)) + +Distributions.ccdf(d::StandardDist{Uniform,0}, x::U) where {U<:Real} = + one(x) - Distributions.cdf(d, x) + + +function Distributions.quantile(d::StandardDist{Uniform,0}, p::U) where {U<:Real} + convert(float(U), p) +end + +function Distributions.cquantile(d::StandardDist{Uniform,0}, p::U) where {U<:Real} + y = Distributions.quantile(d, p) + one(y) - y +end + + +Distributions.mgf(d::StandardDist{Uniform,0}, t::Real) = Distributions.mgf(nonstddist(d), t) +Distributions.cf(d::StandardDist{Uniform,0}, t::Real) = Distributions.cf(nonstddist(d), t) + +Distributions.gradlogpdf(d::StandardDist{Uniform,0}, x::Real) = zero(x) + +function Distributions.gradlogpdf(d::StandardDist{Uniform,N}, x::AbstractArray{<:Real,N}) where {N} + zero(checked_arg(d, x)) +end + +Base.rand(rng::AbstractRNG, d::StandardDist{Uniform,0}) = rand(rng) +Base.rand(rng::AbstractRNG, d::StandardDist{Uniform,N}) where {N} = rand(rng, size(d)...) +Random.rand!(rng::AbstractRNG, d::StandardDist{Uniform,N}, x::AbstractArray{<:Real,N}) where {N} = + rand!(rng, x) diff --git a/ext/MeasureBaseDistributionsExt/standardmv.jl b/ext/MeasureBaseDistributionsExt/standardmv.jl new file mode 100644 index 00000000..c8e99039 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/standardmv.jl @@ -0,0 +1,36 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + + +MeasureBase.getdof(d::AbstractMvNormal) = length(d) +MeasureBase.getdof(m::AsMeasure{<:AbstractMvNormal}) = getdof(m.obj) + +MeasureBase.transport_origin(ν::MvNormal) = StandardDist{Normal}(length(ν)) + +_cholesky_L(A) = cholesky(A).L +_cholesky_L(A::Diagonal{<:Real}) = Diagonal(sqrt.(diag(A))) +_cholesky_L(A::PDMats.PDiagMat{<:Real}) = Diagonal(sqrt.(A.diag)) +_cholesky_L(A::PDMats.ScalMat{<:Real}) = Diagonal(Fill(sqrt(A.value), A.dim)) + +function MeasureBase.from_origin(ν::MvNormal, x) + A = _cholesky_L(ν.Σ) + b = ν.μ + muladd(A, x, b) +end + +function MeasureBase.to_origin(ν::MvNormal, y) + A = _cholesky_L(ν.Σ) + b = ν.μ + A \ (y - b) +end + + +#DirichletMultinomial +#Distributions.AbstractMvLogNormal +#Distributions.AbstractMvTDist +#Distributions.ProductDistribution{1} +#Distributions.ReshapedDistribution{1, S, D} where {S<:ValueSupport, D<:(Distribution{<:ArrayLikeVariate, S})} +#JointOrderStatistics +#Multinomial +#MultivariateMixture (alias for AbstractMixtureModel{ArrayLikeVariate{1}}) +#MvLogitNormal +#VonMisesFisher diff --git a/ext/MeasureBaseDistributionsExt/univariate.jl b/ext/MeasureBaseDistributionsExt/univariate.jl new file mode 100644 index 00000000..126fe36d --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/univariate.jl @@ -0,0 +1,142 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + + +@inline MeasureBase.getdof(::Distribution{Univariate}) = static(1) + +@inline MeasureBase.check_dof(a::Distribution{Univariate}, b::Distribution{Univariate}) = nothing + + +# Generic transformations to/from StdUniform via cdf/quantile: + + +_dist_params_numtype(d::Distribution) = real_numtype(typeof(Distributions.params(d))) + + +@inline _trafo_cdf(d::Distribution{Univariate,Continuous}, x::Real) = + _trafo_cdf_impl(_dist_params_numtype(d), d, x) + +@inline _trafo_cdf_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, x::Real) = + Distributions.cdf(d, x) + + +@inline _trafo_quantile(d::Distribution{Univariate,Continuous}, u::Real) = + _trafo_quantile_impl(_dist_params_numtype(d), d, u) + +@inline _trafo_quantile_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, u::Real) = + _trafo_quantile_impl_generic(d, u) + + +@inline _trafo_quantile_impl_generic(d::Distribution{Univariate,Continuous}, u::Real) = + Distributions.quantile(d, u) + +# Workaround for Beta dist, current quantile implementation only supports Float64: +@inline function _trafo_quantile_impl_generic(d::Beta{T}, u::Union{Integer,AbstractFloat}) where {T<:Union{Integer,AbstractFloat}} + Distributions.quantile(d, convert(promote_type(Float64, typeof(u)), u)) +end + +# Workaround for rounding errors that can result in quantile values outside of support of Truncated: +@inline function _trafo_quantile_impl_generic(d::Distributions.Truncated{<:Distribution{Univariate,Continuous}}, u::Real) + x = Distributions.quantile(d, u) + T = typeof(x) + min_x = T(minimum(d)) + max_x = T(maximum(d)) + if x < min_x && isapprox(x, min_x, atol = 4 * eps(T)) + min_x + elseif x > max_x && isapprox(x, max_x, atol = 4 * eps(T)) + max_x + else + x + end +end + + +@inline function _result_numtype(d::Distribution{Univariate}, x::T) where {T<:Real} + float(promote_type(T, _dist_params_numtype(d))) +end + + +@inline function MeasureBase.transport_def(::StdUniform, μ::Distribution{Univariate,Continuous}, x) + R = _result_numtype(μ, x) + if Distributions.insupport(μ, x) + y = _trafo_cdf(μ, x) + convert(R, y) + else + convert(R, NaN) + end +end + + +@inline function MeasureBase.transport_def(ν::Distribution{Univariate,Continuous}, ::StdUniform, x::T) where {T} + R = _result_numtype(ν, x) + TF = float(T) + if 0 <= x <= 1 + # Avoid x ≈ 0 and x ≈ 1 to avoid infinite variate values for target distributions with infinite support: + mod_x = ifelse(x ≈ 0, zero(TF) + eps(TF), ifelse(x ≈ 1, one(TF) - eps(TF), convert(TF, x))) + y = _trafo_quantile(ν, mod_x) + convert(R, y) + else + convert(R, NaN) + end +end + + +# Use standard measures as transformation origin for scaled/translated equivalents: + +function _origin_to_affine(ν::Distribution{Univariate}, y::T) where {T<:Real} + trg_offs, trg_scale = Distributions.location(ν), Distributions.scale(ν) + x = muladd(y, trg_scale, trg_offs) + convert(_result_numtype(ν, y), x) +end + +function _affine_to_origin(μ::Distribution{Univariate}, x::T) where {T<:Real} + src_offs, src_scale = Distributions.location(μ), Distributions.scale(μ) + y = (x - src_offs) / src_scale + convert(_result_numtype(μ, x), y) +end + +for (A, B) in [ + (Uniform, StdUniform), + (Logistic, StdLogistic), + (Normal, StdNormal) +] + @eval begin + @inline MeasureBase.transport_origin(::$A) = $B() + @inline MeasureBase.to_origin(ν::$A, y) = _affine_to_origin(ν, y) + @inline MeasureBase.from_origin(ν::$A, x) = _origin_to_affine(ν, x) + end +end + +@inline MeasureBase.transport_origin(::Exponential) = StdExponential() +@inline MeasureBase.to_origin(ν::Exponential, y) = Distributions.scale(ν) \ y +@inline MeasureBase.from_origin(ν::Exponential, x) = Distributions.scale(ν) * x + + +# Use the underlying distribution as transformation origin for affine +# transformed distributions: + +@inline MeasureBase.transport_origin(d::Distributions.AffineDistribution) = d.ρ +@inline MeasureBase.from_origin(d::Distributions.AffineDistribution, x) = muladd(d.σ, x, d.μ) +@inline MeasureBase.to_origin(d::Distributions.AffineDistribution, y) = d.σ \ (y - d.μ) + + + +# Transform between univariate and single-element power measure + +function MeasureBase.transport_def(ν::Distribution{Univariate}, μ::PowerMeasure{<:StdMeasure}, x) + return transport_def(ν, μ.parent, only(x)) +end + +function MeasureBase.transport_def(ν::PowerMeasure{<:StdMeasure}, μ::Distribution{Univariate}, x) + return Fill(transport_def(ν.parent, μ, only(x)), map(length, ν.axes)...) +end + + +# Transform between univariate and single-element standard multivariate + +function MeasureBase.transport_def(ν::Distribution{Univariate}, μ::StandardDist{D,1}, x) where {D} + return transport_def(ν, StandardDist{D}(), only(x)) +end + +function MeasureBase.transport_def(ν::StandardDist{D,1}, μ::Distribution{Univariate}, x) where {D} + return Fill(transport_def(StandardDist{D}(), μ, only(x)), size(ν)...) +end diff --git a/ext/MeasureBaseDistributionsForwardDiffExt.jl b/ext/MeasureBaseDistributionsForwardDiffExt.jl new file mode 100644 index 00000000..245c8e60 --- /dev/null +++ b/ext/MeasureBaseDistributionsForwardDiffExt.jl @@ -0,0 +1,36 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +module MeasureBaseDistributionsForwardDiffExt + +import MeasureBase +import Distributions +import ForwardDiff + +using Distributions: Distribution, Univariate, Continuous, Beta + +@inline function MeasureBase._trafo_cdf_impl( + ::Type{<:Union{Integer,AbstractFloat}}, + d::Distribution{Univariate,Continuous}, + x::ForwardDiff.Dual{TAG}, +) where {TAG} + x_v = ForwardDiff.value(x) + u = Distributions.cdf(d, x_v) + dudx = Distributions.pdf(d, x_v) + ForwardDiff.Dual{TAG}(u, dudx * ForwardDiff.partials(x)) +end + +@inline function MeasureBase._trafo_quantile_impl( + ::Type{<:Union{Integer,AbstractFloat}}, + d::Distribution{Univariate,Continuous}, + u::ForwardDiff.Dual{TAG}, +) where {TAG} + x = MeasureBase._trafo_quantile_impl_generic(d, ForwardDiff.value(u)) + dxdu = inv(Distributions.pdf(d, x)) + ForwardDiff.Dual{TAG}(x, dxdu * ForwardDiff.partials(u)) +end + +# Workaround for Beta dist, ForwardDiff doesn't work for parameters: +@inline MeasureBase._trafo_quantile_impl_generic(d::Beta{T}, u::Real) where {T<:ForwardDiff.Dual} = + convert(float(typeof(u)), NaN) + +end # module MeasureBaseDistributionsForwardDiffExt diff --git a/ext/MeasureBaseDistributionsForwardDiffPullbacksExt.jl b/ext/MeasureBaseDistributionsForwardDiffPullbacksExt.jl new file mode 100644 index 00000000..667af023 --- /dev/null +++ b/ext/MeasureBaseDistributionsForwardDiffPullbacksExt.jl @@ -0,0 +1,25 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +module MeasureBaseDistributionsForwardDiffPullbacksExt + +import MeasureBase +using MeasureBase: StdMeasure, transport_def + +import Distributions +using Distributions: Distribution, Univariate + +import ChainRulesCore +using ForwardDiffPullbacks: fwddiff + +# Use ForwardDiff for univariate transformations: +@inline function ChainRulesCore.rrule(::typeof(transport_def), ν::Distribution{Univariate}, μ::Distribution{Univariate}, x::Any) + ChainRulesCore.rrule(fwddiff(transport_def), ν, μ, x) +end +@inline function ChainRulesCore.rrule(::typeof(transport_def), ν::StdMeasure, μ::Distribution{Univariate}, x::Any) + ChainRulesCore.rrule(fwddiff(transport_def), ν, μ, x) +end +@inline function ChainRulesCore.rrule(::typeof(transport_def), ν::Distribution{Univariate}, μ::StdMeasure, x::Any) + ChainRulesCore.rrule(fwddiff(transport_def), ν, μ, x) +end + +end # module MeasureBaseDistributionsForwardDiffPullbacksExt diff --git a/ext/MeasureBaseDistributionsMooncakeExt.jl b/ext/MeasureBaseDistributionsMooncakeExt.jl new file mode 100644 index 00000000..75bde91b --- /dev/null +++ b/ext/MeasureBaseDistributionsMooncakeExt.jl @@ -0,0 +1,21 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +module MeasureBaseDistributionsMooncakeExt + +using MeasureBase +import Distributions +import Mooncake +using Mooncake: @zero_derivative, MinimalCtx + +using Distributions: Distribution +using MeasureBase: _dist_params_numtype + +# The distribution transports themselves need no rules here: Mooncake +# provides rules for Distributions and StatsFuns/SpecialFunctions, so it +# differentiates the cdf/quantile-based transports natively. The +# ForwardDiffPullbacks-based rules for `transport_def` are a +# Zygote/ChainRules pathway. + +@zero_derivative MinimalCtx Tuple{typeof(_dist_params_numtype),Distribution} + +end # module MeasureBaseDistributionsMooncakeExt diff --git a/ext/MeasureBaseForwardDiffExt.jl b/ext/MeasureBaseForwardDiffExt.jl new file mode 100644 index 00000000..20113c7f --- /dev/null +++ b/ext/MeasureBaseForwardDiffExt.jl @@ -0,0 +1,18 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +module MeasureBaseForwardDiffExt + +using MeasureBase +using MeasureBase: containsnan, firsttype +import ForwardDiff + +function MeasureBase.containsnan(x::ForwardDiff.Dual) + a = containsnan(x.value) + b = containsnan(x.partials) + return a || b +end + +MeasureBase.firsttype(::Type{T}, ::Type{<:ForwardDiff.Dual{tag,<:Real,N}}) where {T<:Real,tag,N} = + ForwardDiff.Dual{tag,T,N} + +end # module MeasureBaseForwardDiffExt diff --git a/ext/MeasureBaseForwardDiffPullbacksExt.jl b/ext/MeasureBaseForwardDiffPullbacksExt.jl new file mode 100644 index 00000000..72ffc751 --- /dev/null +++ b/ext/MeasureBaseForwardDiffPullbacksExt.jl @@ -0,0 +1,10 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +module MeasureBaseForwardDiffPullbacksExt + +import MeasureBase +using ForwardDiffPullbacks: fwddiff + +MeasureBase._fwddiff(f::Function) = fwddiff(f) + +end # module MeasureBaseForwardDiffPullbacksExt diff --git a/ext/MeasureBaseMooncakeExt.jl b/ext/MeasureBaseMooncakeExt.jl new file mode 100644 index 00000000..e401cb7d --- /dev/null +++ b/ext/MeasureBaseMooncakeExt.jl @@ -0,0 +1,29 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +module MeasureBaseMooncakeExt + +using MeasureBase +import Mooncake +using Mooncake: @zero_derivative, MinimalCtx + +using MeasureBase: isneginf, isposinf, _adignore_call +using MeasureBase: check_dof, require_insupport, _origin_depth +using MeasureBase: logdensityof_rt + +# Unlike Zygote, Mooncake differentiates the collection utilities +# (`_pushfront`, etc., mutating code in general), `checked_arg` and +# `_checksupport` natively, so only the non-differentiable functions +# need rules: + +@zero_derivative MinimalCtx Tuple{typeof(isneginf),Any} +@zero_derivative MinimalCtx Tuple{typeof(isposinf),Any} + +@zero_derivative MinimalCtx Tuple{typeof(_adignore_call),Any} + +@zero_derivative MinimalCtx Tuple{typeof(require_insupport),Any,Any} +@zero_derivative MinimalCtx Tuple{typeof(_origin_depth),Any} +@zero_derivative MinimalCtx Tuple{typeof(check_dof),Any,Any} + +@zero_derivative MinimalCtx Tuple{typeof(logdensityof_rt),Any,Any} + +end # module MeasureBaseMooncakeExt diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index 2bad7d92..4ab59601 100644 --- a/src/MeasureBase.jl +++ b/src/MeasureBase.jl @@ -31,6 +31,7 @@ import ConstructionBase using ConstructionBase: constructorof using IntervalSets +import StaticArrays using StaticArrays: StaticArray, StaticVector, StaticMatrix, SArray, SVector, SMatrix, SOneTo @@ -43,6 +44,22 @@ using Static: StaticInteger using FunctionChains using PropertyFunctions: PropSelFunction +using StaticThings: + AxesLike, StaticAxesLike, SizeLike, StaticSizeLike, + OneToLike, StaticOneTo, StaticOneToLike, RealLike, + IntegerLike, StaticUnitRange, StaticUnitRangeLike, + NoTypeSize, + asaxes, asnonstatic, + canonical_axes, canonical_indices, canonical_size, + maybestatic_axes, maybestatic_eachindex, + maybestatic_length, maybestatic_size, maybestatic_first, maybestatic_last, + maybestatic_oneto, maybestatic_fill, maybestatic_reshape, + size_from_type, axes2size, size2axes, size2length, + staticarray_type + +import HeterogeneousComputing +using HeterogeneousComputing: real_numtype + export gentype export rebase @@ -64,6 +81,39 @@ abstract type AbstractMeasure end AbstractMeasure(m::AbstractMeasure) = m +""" + asmeasure(m) + +Turns a measure-like object `m` into an `AbstractMeasure`. + +Calls `convert(AbstractMeasure, m)` by default +""" +function asmeasure end + +@inline asmeasure(m::AbstractMeasure) = m +asmeasure(m) = convert(AbstractMeasure, m) +export asmeasure + +""" + struct AsMeasure{T} + +Wrapes a measure-like object into an `AbstractMeasure`. + +Constructor: + +``` +AsMeasure{T}(obj::T) +``` + +User code should not create instances of `AsMeasure` directly, but should +call `asmeasure(obj)` instead. +""" +struct AsMeasure{T} <: AbstractMeasure + obj::T + + AsMeasure{T}(obj::T) where {T} = new(obj) +end + function Pretty.quoteof(d::M) where {M<:AbstractMeasure} the_names = fieldnames(typeof(d)) :($M($([getfield(d, n) for n in the_names]...))) @@ -113,7 +163,7 @@ using Compat using IrrationalConstants using IrrationalConstants: loghalf -include("static.jl") +include("collection_utils.jl") include("smf.jl") include("getdof.jl") include("transport.jl") @@ -134,6 +184,7 @@ include("primitives/trivial.jl") include("combinators/bind.jl") include("combinators/transformedmeasure.jl") +include("combinators/reshape.jl") include("combinators/weighted.jl") include("combinators/superpose.jl") include("combinators/product.jl") diff --git a/src/collection_utils.jl b/src/collection_utils.jl new file mode 100644 index 00000000..1de51f7e --- /dev/null +++ b/src/collection_utils.jl @@ -0,0 +1,24 @@ +function _pushfront(v::AbstractVector, x) + T = promote_type(eltype(v), typeof(x)) + r = similar(v, T, length(eachindex(v)) + 1) + r[firstindex(r)] = x + r[firstindex(r)+1:lastindex(r)] = v + r +end + +function _pushback(v::AbstractVector, x) + T = promote_type(eltype(v), typeof(x)) + r = similar(v, T, length(eachindex(v)) + 1) + r[lastindex(r)] = x + r[firstindex(r):lastindex(r)-1] = v + r +end + +_dropfront(v::AbstractVector) = v[firstindex(v)+1:lastindex(v)] + +_dropback(v::AbstractVector) = v[firstindex(v):lastindex(v)-1] + +_rev_cumsum(xs::AbstractVector) = reverse(cumsum(reverse(xs))) + +# Equivalent to `cumprod(xs)``: +_exp_cumsum_log(xs::AbstractVector) = exp.(cumsum(log.(xs))) diff --git a/src/combinators/implicitlymapped.jl b/src/combinators/implicitlymapped.jl index 964ea466..3966b10a 100644 --- a/src/combinators/implicitlymapped.jl +++ b/src/combinators/implicitlymapped.jl @@ -179,13 +179,13 @@ struct TakeAny{T<:IntegerLike} n::T end -_takeany_range(f::TakeAny, idxs) = first(idxs):first(idxs)+dynamic(f.n)-1 +_takeany_range(f::TakeAny, idxs) = first(idxs):(first(idxs)+dynamic(f.n)-1) @inline _takeany_range(f::TakeAny, ::OneTo) = OneTo(dynamic(f.n)) @inline _takeany_range(::TakeAny{<:Static.StaticInteger{N}}, ::OneTo) where {N} = SOneTo(N) @inline _takeany_range(::TakeAny{<:Static.StaticInteger{N}}, ::SOneTo) where {N} = SOneTo(N) -@inline (f::TakeAny)(xs::Tuple) = xs[begin:begin+f.n-1] +@inline (f::TakeAny)(xs::Tuple) = xs[begin:(begin+f.n-1)] @inline (f::TakeAny)(xs::AbstractVector) = xs[_takeany_range(f, eachindex(xs))] function (f::TakeAny)(xs) diff --git a/src/combinators/power.jl b/src/combinators/power.jl index e6397c3f..6f065c3f 100644 --- a/src/combinators/power.jl +++ b/src/combinators/power.jl @@ -1,4 +1,5 @@ import Base +import StaticThings export PowerMeasure @@ -11,14 +12,37 @@ the product determines the dimensionality of the resulting support. Note that power measures are only well-defined for integer powers. The nth power of a measure μ can be written μ^n. + +See also [`pwr_base`](@ref), [`pwr_axes`](@ref) and [`pwr_size`](@ref). """ struct PowerMeasure{M,A} <: AbstractProductMeasure parent::M axes::A end -maybestatic_length(μ::PowerMeasure) = prod(maybestatic_size(μ)) -maybestatic_size(μ::PowerMeasure) = map(maybestatic_length, μ.axes) +StaticThings.maybestatic_length(μ::PowerMeasure) = size2length(maybestatic_size(μ)) +StaticThings.maybestatic_size(μ::PowerMeasure) = axes2size(μ.axes) + +""" + MeasureBase.pwr_base(μ::PowerMeasure) + +Returns `ν` for `μ = ν^axs` +""" +@inline pwr_base(μ::PowerMeasure) = μ.parent + +""" + MeasureBase.pwr_axes(μ::PowerMeasure) + +Returns `axs` for `μ = ν^axs`, `axs` being a tuple of integer ranges. +""" +@inline pwr_axes(μ::PowerMeasure) = μ.axes + +""" + MeasureBase.pwr_size(μ::PowerMeasure) + +Returns `sz` for `μ = ν^sz`, `sz` being a tuple of integers. +""" +@inline pwr_size(μ::PowerMeasure) = axes2size(μ.axes) function Pretty.tile(μ::PowerMeasure) sz = length.(μ.axes) @@ -30,7 +54,7 @@ end # ToDo: Make rand return static arrays for statically-sized power measures. function _cartidxs(axs::Tuple{Vararg{AbstractUnitRange,N}}) where {N} - CartesianIndices(map(_dynamic, axs)) + CartesianIndices(map(asnonstatic, axs)) end function Base.rand( @@ -38,31 +62,30 @@ function Base.rand( ::Type{T}, d::PowerMeasure{M}, ) where {T,M<:AbstractMeasure} - map(_cartidxs(d.axes)) do _ - rand(rng, T, d.parent) + axs, base_d = pwr_axes(d), pwr_base(d) + map(_cartidxs(axs)) do _ + rand(rng, T, base_d) end end function Base.rand(rng::AbstractRNG, ::Type{T}, d::PowerMeasure) where {T} - map(_cartidxs(d.axes)) do _ - rand(rng, d.parent) + axs, base_d = pwr_axes(d), pwr_base(d) + map(_cartidxs(axs)) do _ + rand(rng, base_d) end end -@inline _pm_axes(sz::Tuple{Vararg{IntegerLike,N}}) where {N} = map(one_to, sz) -@inline _pm_axes(axs::Tuple{Vararg{AbstractUnitRange,N}}) where {N} = axs - @inline function powermeasure(x::T, sz::Tuple{Vararg{Any,N}}) where {T,N} - PowerMeasure(x, _pm_axes(sz)) + PowerMeasure(x, asaxes(sz)) end -marginals(d::PowerMeasure) = fill_with(d.parent, d.axes) +marginals(d::PowerMeasure) = maybestatic_fill(d.parent, d.axes) function Base.:^(μ::AbstractMeasure, dims::Tuple{Vararg{AbstractArray,N}}) where {N} powermeasure(μ, dims) end -Base.:^(μ::AbstractMeasure, dims::Tuple) = powermeasure(μ, one_to.(dims)) +Base.:^(μ::AbstractMeasure, dims::Tuple) = powermeasure(μ, maybestatic_oneto.(dims)) Base.:^(μ::AbstractMeasure, n) = powermeasure(μ, (n,)) # Base.show(io::IO, d::PowerMeasure) = print(io, d.parent, " ^ ", size(d.xs)) @@ -80,13 +103,25 @@ end for func in [:logdensityof, :logdensity_def] @eval @inline function $func(d::PowerMeasure{M}, x) where {M} - parent = d.parent - sum(x) do xj - $func(parent, xj) + parent_m = d.parent + sz_parent = axes2size(d.axes) + sz_x = maybestatic_size(x) + if sz_parent != sz_x + throw(ArgumentError("Size of variate doesn't match size of power measure")) + end + R = infer_logdensity_type($func, parent_m, eltype(x)) + if isempty(x) + return zero(R)::R + else + # Need to convert since sum can turn static into dynamic values: + return convert(R, sum(Base.Fix1($func, parent_m), x))::R end end - @eval @inline function $func(d::PowerMeasure{M,Tuple{Static.SOneTo{N}}}, x) where {M,N} + @eval @inline function $func( + d::PowerMeasure{<:Any,Tuple{<:StaticOneToLike{N}}}, + x, + ) where {N} parent = d.parent sum(1:N) do j @inbounds $func(parent, x[j]) @@ -94,9 +129,9 @@ for func in [:logdensityof, :logdensity_def] end @eval @inline function $func( - d::PowerMeasure{M,NTuple{N,Static.SOneTo{0}}}, + ::PowerMeasure{<:Any,<:Tuple{Vararg{StaticOneToLike{0}}}}, x, - ) where {M,N} + ) static(0.0) end end @@ -117,15 +152,11 @@ end end end -@inline getdof(μ::PowerMeasure) = getdof(μ.parent) * prod(map(length, μ.axes)) - -@inline function getdof(::PowerMeasure{<:Any,NTuple{N,Static.SOneTo{0}}}) where {N} - static(0) -end +@inline getdof(μ::PowerMeasure) = getdof(μ.parent) * size2length(axes2size(μ.axes)) @propagate_inbounds function checked_arg(μ::PowerMeasure, x::AbstractArray{<:Any}) @boundscheck begin - sz_μ = map(length, μ.axes) + sz_μ = pwr_size(μ) sz_x = size(x) if sz_μ != sz_x throw(ArgumentError("Size of variate doesn't match size of power measure")) @@ -144,8 +175,10 @@ logdensity_def(::PowerMeasure{P}, x) where {P<:PrimitiveMeasure} = static(0.0) # To avoid ambiguities function logdensity_def( - ::PowerMeasure{P,Tuple{Vararg{Static.SOneTo{0},N}}}, - x, + ::PowerMeasure{P,<:Tuple{Vararg{StaticOneToLike{0},N}}}, ::Any, ) where {P<:PrimitiveMeasure,N} static(0.0) end + + +@inline mspace_elsize(m::PowerMeasure) = axes2size(m.axes) diff --git a/src/combinators/product.jl b/src/combinators/product.jl index 0290419d..9135dc2b 100644 --- a/src/combinators/product.jl +++ b/src/combinators/product.jl @@ -244,6 +244,10 @@ function checked_arg(μ::ProductMeasure{<:NTuple{N,Any}}, x::NTuple{N,Any}) wher map(checked_arg, marginals(μ), x) end +function checked_arg(μ::ProductMeasure{<:AbstractArray}, x::AbstractArray) + map(checked_arg, marginals(μ), x) +end + function checked_arg( μ::ProductMeasure{<:NamedTuple{names}}, x::NamedTuple{names}, diff --git a/src/combinators/reshape.jl b/src/combinators/reshape.jl new file mode 100644 index 00000000..dddae55b --- /dev/null +++ b/src/combinators/reshape.jl @@ -0,0 +1,69 @@ +# ToDo: Support static resizes for static arrays + +""" + struct MeasureBase.Reshape <: Function + +Represents a function that reshapes an array. + +Supports `InverseFunctions.inverse` and +`ChangesOfVariables.with_logabsdet_jacobian`. + +Constructor: + +```julia +Reshape(output_size::Dims, input_size::Dims) +``` +""" +struct Reshape{M<:SizeLike,N<:SizeLike} <: Function + output_size::M + input_size::N + + Reshape{M,N}(out_sz::M, in_sz::N) where {M<:SizeLike,N<:SizeLike} = + new{M,N}(out_sz, in_sz) +end + +function Reshape(output_size::SizeLike, input_size::SizeLike) + out_sz = canonical_size(output_size) + in_sz = canonical_size(input_size) + return Reshape{typeof(out_sz), typeof(in_sz)}(out_sz, in_sz) +end + +_throw_reshape_mismatch(sz, sz_x) = throw(DimensionMismatch("Reshape input size is $sz but got input of size $sz_x")) + +function (f::Reshape)(x::AbstractArray) + sz_x = maybestatic_size(x) + f.input_size == sz_x || _throw_reshape_mismatch(f.input_size, sz_x) + return reshape(x, f.output_size) +end + +InverseFunctions.inverse(f::Reshape{M,N}) where {M,N} = Reshape{N,M}(f.input_size, f.output_size) + +function ChangesOfVariables.with_logabsdet_jacobian(f::Reshape, x::AbstractArray) + return f(x), zero(real_numtype(typeof(x))) +end + + +""" + mreshape(m::AbstractMeasure, sz::Vararg{N,IntegerLike}) where N + mreshape(m::AbstractMeasure, sz::NTuple{N,IntegerLike}) where N + +Reshape a measure `m` over an array-valued space, returning a measure over +a space of arrays with shape `sz`. +""" +function mreshape end + +mreshape(m::AbstractMeasure, sz::IntegerLike...) = mreshape(m, sz) +mreshape(m::AbstractMeasure, sz::SizeLike) = pushfwd(Reshape(sz, mspace_elsize(m)), m) + + +""" + MeasureBase.mspace_elsize(m::AbstractMeasure)::MeasureBase.SizeLike + +Return the size of the elements of the measurable space of `m`. + +Defaults to the size of a test value of `m`, may be specialized for +measure types where this is inefficient. +""" +function mspace_elsize end + +mspace_elsize(m::AbstractMeasure) = maybestatic_size(testvalue(m)) diff --git a/src/combinators/superpose.jl b/src/combinators/superpose.jl index 099ee806..aa7b6e20 100644 --- a/src/combinators/superpose.jl +++ b/src/combinators/superpose.jl @@ -65,37 +65,24 @@ function Base.:+(μ::AbstractMeasure, ν::AbstractMeasure) superpose(μ, ν) end -oneplus(x::ULogarithmic) = exp(ULogarithmic, log1pexp(x.log)) - -@inline function density_def(s::SuperpositionMeasure{Tuple{A,B}}, x) where {A,B} - (μ, ν) = s.components - - istrue(insupport(μ, x)) || return exp(ULogarithmic, logdensity_def(ν, x)) - istrue(insupport(ν, x)) || return exp(ULogarithmic, logdensity_def(μ, x)) - - α = basemeasure(μ) - β = basemeasure(ν) - dμ_dα = exp(ULogarithmic, logdensity_def(μ, x)) - dν_dβ = exp(ULogarithmic, logdensity_def(ν, x)) - dα_dβ = exp(ULogarithmic, logdensity_rel(α, β, x)) - dβ_dα = inv(dα_dβ) - return dμ_dα / oneplus(dβ_dα) + dν_dβ / oneplus(dα_dβ) -end +@inline _ulogexp(x) = exp(ULogarithmic, dynamic(x)) function density_def(s::SuperpositionMeasure, x) - T = typeof(s) - msg = """ - Not implemented: There is no method - density_def(::$T, x) - """ - error(msg) + cs = values(s.components) + αs = map(basemeasure, cs) + idxs = eachindex(cs) + sum(idxs) do i + dμᵢ_dαᵢ = _ulogexp(logdensity_def(cs[i], x)) + istrue(insupport(cs[i], x)) || return zero(dμᵢ_dαᵢ) + dΣα_dαᵢ = sum(idxs) do j + dαⱼ_dαᵢ = _ulogexp(logdensity_rel(αs[j], αs[i], x)) + istrue(insupport(cs[j], x)) ? dαⱼ_dαᵢ : zero(dαⱼ_dαᵢ) + end + dμᵢ_dαᵢ / dΣα_dαᵢ + end end -@inline function logdensity_def( - μ::T, - ν::T, - x, -) where {T<:(SuperpositionMeasure{Tuple{A,B}} where {A,B})} +@inline function logdensity_def(μ::T, ν::T, x) where {T<:SuperpositionMeasure} if μ === ν return zero(return_type(logdensity_def, (μ, x))) else @@ -103,44 +90,49 @@ end end end -@inline function logdensity_def( - s::T, - β, - x, -) where {T<:(SuperpositionMeasure{Tuple{A,B}} where {A,B})} - (μ, ν) = s.components - - istrue(insupport(μ, x)) || return logdensity_rel(ν, β, x) - istrue(insupport(ν, x)) || return logdensity_rel(μ, β, x) - return logaddexp(logdensity_rel(μ, β, x), logdensity_rel(ν, β, x)) +function _superpos_logdensity_rel(s::SuperpositionMeasure, β, x) + cs = values(s.components) + ds = map(cs) do μ + istrue(insupport(μ, x)) ? dynamic(logdensity_rel(μ, β, x)) : -Inf + end + logsumexp(ds) end -@inline function logdensity_def( - s::SuperpositionMeasure{Tuple{A,B}}, - β::SuperpositionMeasure, - x, -) where {A,B} - (μ, ν) = s.components - istrue(insupport(μ, x)) || return logdensity_rel(ν, β, x) - istrue(insupport(ν, x)) || return logdensity_rel(μ, β, x) - return logaddexp(logdensity_rel(μ, β, x), logdensity_rel(ν, β, x)) -end +@inline logdensity_def(s::SuperpositionMeasure, β, x) = _superpos_logdensity_rel(s, β, x) -@inline function logdensity_def(s, β::(SuperpositionMeasure{Tuple{A,B}} where {A,B}), x) - -logdensity_def(β, s, x) -end +@inline logdensity_def(s::SuperpositionMeasure, β::SuperpositionMeasure, x) = + _superpos_logdensity_rel(s, β, x) + +@inline logdensity_def(s, β::SuperpositionMeasure, x) = -_superpos_logdensity_rel(β, s, x) @inline logdensity_def(s::SuperpositionMeasure, x) = log(density_def(s, x)) -function basemeasure(μ::SuperpositionMeasure{Tuple{A,B}}) where {A,B} +function basemeasure(μ::SuperpositionMeasure{<:Tuple}) superpose(map(basemeasure, μ.components)...) end + +function basemeasure(μ::SuperpositionMeasure{<:AbstractArray}) + bases = map(basemeasure, μ.components) + allequal(bases) ? weightedmeasure(log(length(bases)), first(bases)) : superpose(bases) +end + basemeasure(μ::SuperpositionMeasure) = superpose(map(basemeasure, μ.components)) -# TODO: Fix `rand` method (this one is wrong) -# function Base.rand(μ::SuperpositionMeasure{X,N}) where {X,N} -# return rand(rand(μ.components)) -# end +function Base.rand(rng::AbstractRNG, ::Type{T}, μ::SuperpositionMeasure) where {T} + components = values(μ.components) + masses = map(massof, components) + total = sum(masses) + total isa AbstractUnknownMass && throw( + ArgumentError("Cannot sample from a superposition of measures of unknown mass"), + ) + threshold = rand(rng) * dynamic(total) + csum = zero(threshold) + for (mass, c) in zip(masses, components) + csum += dynamic(mass) + csum >= threshold && return rand(rng, T, c) + end + return rand(rng, T, last(components)) +end @inline function insupport(d::SuperpositionMeasure, x) any(d.components) do c diff --git a/src/density-core.jl b/src/density-core.jl index 6ac3d01e..f3b2db2b 100644 --- a/src/density-core.jl +++ b/src/density-core.jl @@ -149,13 +149,13 @@ end ℓ = logdensity_def(μs[$M], νs[$N], x) end - for i in 1:M-1 + for i in 1:(M-1) push!(q.args, :(Δℓ = logdensity_def(μs[$i], x))) # push!(q.args, :(println("Adding", Δℓ))) push!(q.args, :(ℓ += Δℓ)) end - for j in 1:N-1 + for j in 1:(N-1) push!(q.args, :(Δℓ = logdensity_def(νs[$j], x))) # push!(q.args, :(println("Subtracting", Δℓ))) push!(q.args, :(ℓ -= Δℓ)) diff --git a/src/domains.jl b/src/domains.jl index e03f753c..7458b450 100644 --- a/src/domains.jl +++ b/src/domains.jl @@ -1,59 +1,228 @@ -abstract type AbstractDomain end +""" + mdomain(m)::MeasureBase.SetLike -abstract type RealDomain <: AbstractDomain end +Return the domain, i.e. the measurable set, of the measure `m`. -# TODO: Use IntervalSets -struct RealNumbers <: RealDomain end +The measure must allow for evaluating densities and the like over the whole +domain, even if the support of the measure is only a subset of the domain. -const ℝ = RealNumbers() +May return [`MeasureBase.ImplicitDomain`](@ref) if the domain cannot be +computed (efficiently). +""" +function mdomain end +export mdomain -Base.minimum(::RealNumbers) = static(-Inf) -Base.maximum(::RealNumbers) = static(Inf) +@inline mdomain(m) = ImplicitDomain(m) -Base.in(x, ::RealNumbers) = isreal(x) -Base.show(io::IO, ::typeof(ℝ)) = print(io, "ℝ") +# Custom abstract set type. Design reserve to be able to switch to +#`Base.AbstractSet` or another set type hierarchy in the future: +""" + MeasureBase.ValueSet -struct BoundedReals{L,U} <: RealDomain - lower::L - upper::U +Abstract type for some measurable sets. + +Not every measurable set needs to be a a subtype of +`MeasureBase.ValueSet`. + +See also [`MeasureBase.SetLike`](@ref). +""" +abstract type ValueSet end + +""" + const MeasureBase.SetLike = Union{MeasureBase.ValueSet, Base.AbstractSet, IntervalSets.Domain} + +Any kind of (measurable) set. + +There needs to be an implicit sigma-algebra for subtypes of +`MeasureBase.SetLike` to make them useable for measures. This can't easily be +imposed via type constraints, though, so it is by-contract. +""" +const SetLike = Union{MeasureBase.ValueSet,Base.AbstractSet,IntervalSets.Domain} + +""" + valdomain(x)::MeasureBase.SetLike + +Return the domain of a given value. + +May return [`MeasureBase.UnknownDomain`](@ref) if no domain type is available +that can represents values like `x`. +""" +function valdomain end +export valdomain + +@inline valdomain(x) = UnknownDomain(x) + +""" + MeasureBase.maybe_in(x, s) + +Test if `x` may be a member of `s`. + +Defaults to `in(x, s)`, but may be specialized for certain types of `s`, +e.g. for `s::MeasureBase.ImplicitDomain`. +""" +function maybe_in end + +maybe_in(x, s) = in(x, s) + +""" + struct MeasureBase.ImplicitDomain{M} <: MeasureBase.ValueSet + +Represents the domain (i.e. the measurable set) of a measure `m::M`. + +Constructors: + +``` +MeasureBase.ImplicitDomain(m) +``` + +Fields: + +* `m::M`: The measure. + +For many pushforward measures and similar, the measureable space can not be +computed efficiently or at all. In such cases, [`mdomain(m)`](@ref) should +return `ImplicitDomain(m)`. + +Does not support `Base.in(x, s::MeasureBase.ImplicitDomain)`, and +`MeasureBase.maybe_in(x, s::MeasureBase.ImplicitDomain)` always return `true` +(unless specialized for the measure type). +""" +struct ImplicitDomain{M} <: ValueSet + m::M end -Base.in(x, b::BoundedReals) = b.lower ≤ x ≤ b.upper +@inline Base.union(s::ImplicitDomain, others::ImplicitDomain...) = + ImplicitDomain(+(s.m, map(x -> x.m, others)...)) + +function Base.in(@nospecialize(x), ::ImplicitDomain) + throw(ArgumentError("Cannot test if a value lies withing an implicit domain.")) +end -export ℝ, ℝ₊, 𝕀, ℤ +maybe_in(@nospecialize(x), ::ImplicitDomain) = true -const ℝ₊ = BoundedReals(static(0.0), static(Inf)) -const 𝕀 = BoundedReals(static(0.0), static(1.0)) +function Base.isempty(::ImplicitDomain) + throw(ArgumentError("Can't test if an ImplicitDomain is empty")) +end -Base.minimum(b::BoundedReals) = b.lower -Base.maximum(b::BoundedReals) = b.upper +""" + struct MeasureBase.UnknownDomain{T} <: MeasureBase.ValueSet -Base.show(io::IO, ::typeof(ℝ₊)) = print(io, "ℝ₊") -Base.show(io::IO, ::typeof(𝕀)) = print(io, "𝕀") +Represents the unknown domain of a value of type `T`. -testvalue(::Type{T}, ::typeof(ℝ)) where {T} = zero(T) -testvalue(::Type{T}, ::typeof(ℝ₊)) where {T} = one(T) -testvalue(::Type{T}, ::typeof(𝕀)) where {T} = one(T) / 2 +Constructors: + +``` +MeasureBase.UnknownDomain(x::T) +``` + +Does not support `Base.in(x, s::MeasureBase.UnknownDomain)`, and +`MeasureBase.maybe_in(x, s::MeasureBase.UnknownDomain)` always return `true` +(unless specialized for the measure type). + +`isempty` will always return false, `UnknownDomain` should only be created +if a value of type `T` existed in the first place, which implies that the +domain can not be empty. +""" +struct UnknownDomain{T} <: ValueSet end + +UnknownDomain(::T) where {T} = UnknownDomain{T}() + +Base.eltype(::UnknownDomain{T}) where {T} = T + +@inline Base.union(s::UnknownDomain, others::UnknownDomain...) = + UnknownDomain{promote_type(eltype(s), map(eltype, others)...)}() + +function Base.in(@nospecialize(x), ::UnknownDomain) + throw(ArgumentError("Cannot test if a value lies withing an unknown domain.")) +end -abstract type IntegerDomain <: AbstractDomain end +maybe_in(@nospecialize(x), ::UnknownDomain) = true -struct IntegerNumbers <: IntegerDomain end +Base.isempty(::UnknownDomain) = false -Base.in(x, ::IntegerNumbers) = isinteger(x) +""" + RealValues() isa MeasureBase.ValueSet -const ℤ = IntegerNumbers() +The real numbers. +""" +struct RealValues <: ValueSet end -Base.show(io::IO, ::typeof(ℤ)) = print(io, "ℤ") +@inline Base.in(x::Real, ::RealValues) = true +@inline Base.in(x, ::RealValues) = isreal(x) -Base.minimum(::IntegerNumbers) = static(-Inf) -Base.maximum(::IntegerNumbers) = static(Inf) -struct BoundedInts{L,U} <: IntegerDomain +@inline Base.isempty(::RealValues) = false + +@inline Base.union(s::RealValues, ::RealValues...) = s + +@inline Base.minimum(::RealValues) = static(-Inf) +@inline Base.maximum(::RealValues) = static(Inf) + +testvalue(::Type{T}, ::RealValues) where {T} = zero(T) + +""" + const MeasureBase.ℝ = RealValues() + +The set of all real numbers, see [`MeasureBase.RealValues`](@ref). +""" +const ℝ = RealValues() +export ℝ + +Base.show(io::IO, ::RealValues) = print(io, "ℝ") +Base.show(io::IO, ::MIME"text/plain", ::RealValues) = print(io, "MeasureBase.ℝ") + +""" + MeasureBase.IntegerValues() isa MeasureBase.ValueSet +""" +struct IntegerValues <: ValueSet end + +@inline Base.in(x::Integer, ::IntegerValues) = true +@inline Base.in(x, ::IntegerValues) = isinteger(x) + +@inline Base.isempty(::IntegerValues) = false + +@inline Base.union(s::IntegerValues, ::IntegerValues...) = s + +testvalue(::Type{T}, ::IntegerValues) where {T} = zero(T) + +# # This could get tricky with mixed-precision code. Probably needs some +# # special AbstractInteger infinity type (but custom AbstractInteger types +# # may cause a lot of method invalidations, which is why Static.StaticInteger +# # is not an AbstractInteger). +# @inline Base.minimum(::RealValues) = static(typemax(Int64)) +# @inline Base.maximum(::RealValues) = static(typemin(Int64)) + +""" + const ℤ = IntegerValues() + +The set of all integers, see [`MeasureBase.IntegerValues`](@ref). +""" +const ℤ = IntegerValues() +export ℤ + +Base.show(io::IO, ::IntegerValues) = print(io, "ℤ") +Base.show(io::IO, ::MIME"text/plain", ::IntegerValues) = print(io, "MeasureBase.ℤ") + +""" + struct MeasureBase.BoundedInts{L,U} <: MeasureBase.ValueSet + +The integers from `lower` to `upper` (bounds may be infinite). + +Constructors: + +```julia +BoundedInts(lower, upper) +ℤ[lower:upper] +``` +""" +struct BoundedInts{L,U} <: ValueSet lower::L upper::U end -Base.in(x, b::BoundedInts) = x ∈ ℤ && b.lower ≤ x ≤ b.upper +@inline Base.in(x, b::BoundedInts) = x ∈ ℤ && b.lower <= x <= b.upper + +Base.isempty(b::BoundedInts) = b.lower > b.upper Base.minimum(b::BoundedInts) = b.lower Base.maximum(b::BoundedInts) = b.upper @@ -63,83 +232,126 @@ function Base.show(io::IO, b::BoundedInts) print(io, "ℤ[", b.lower, ":", b.upper, "]") end -testvalue(b::BoundedInts) = min(b.lower, 0) +testvalue(b::BoundedInts) = convert(Int, clamp(0, dynamic(b.lower), dynamic(b.upper))) +testvalue(::Type{T}, b::BoundedInts) where {T} = convert(T, testvalue(b)) -function Base.getindex(::typeof(ℤ), r::AbstractUnitRange) - BoundedInts(extrema(r)...) -end +Base.getindex(::typeof(ℤ), r::AbstractUnitRange) = BoundedInts(extrema(r)...) -########################################################### -# ZeroSet +""" + struct MeasureBase.AbstractCartSetProd <: ValueSet -export ZeroSet +Supertype for cartesian products of sets. +""" +abstract type AbstractCartSetProd <: ValueSet end -struct ZeroSet{F,G} <: AbstractDomain - f::F - ∇f::G -end +""" + struct CartesianProduct <: AbstractCartSetProd -# Based on some quick tests, but may need some adjustment -Base.in(x::AbstractArray{T}, z::ZeroSet) where {T} = abs(z.f(x)) < ldexp(eps(float(T)), 6) +A cartesian product over a collection of sets. -########################################################### -# CodimOne +Constructor: -export CodimOne +```julia +prodset = CartesianProduct(sets) +``` -abstract type CodimOne <: AbstractDomain end +`sets` may be a `Tuple`, `NamedTuple` or `AbstractArray` of sets/domains. +""" +struct CartesianProduct{S<:Union{Tuple,NamedTuple,AbstractArray}} <: AbstractCartSetProd + _sets::S +end -function tangentat( - a::CodimOne, - b::CodimOne, - x::AbstractArray{T}; - tol = ldexp(eps(float(T)), 6), -) where {T} - # Sometimes you get lucky - a == b && return true +componentsets(s::CartesianProduct) = s._sets + +setcartprod(sets::AbstractArray{<:SetLike}) = CartesianProduct(sets) +setcartprod(sets::Tuple{Vararg{SetLike}}) = CartesianProduct(sets) +setcartprod(sets::NamedTuple{names,<:Tuple{Vararg{SetLike}}}) = CartesianProduct(sets) + +@inline Base.in(x::Tuple{}, s::CartesianProduct{Tuple{}}) = true +@inline Base.in(x::Tuple{Vararg{Any,N}}, s::CartesianProduct{<:Tuple{Vararg{Any,N}}}) where {N} = + prod(map(in, x, componentsets(s)))::Bool +@inline Base.in(x::NamedTuple{names}, s::CartesianProduct{<:NamedTuple{names}}) where {names} = + prod(map(in, values(x), values(componentsets(s))))::Bool +# ToDo: Allow this? +# Base.in(x::AbstractVector, s::CartesianProduct{<:Tuple}) = all(in.(x,componentsets(s)))::Bool +function Base.in( + x::AbstractArray{<:Any,N}, + s::CartesianProduct{<:AbstractArray{<:Any,N}}, +) where {N} + sets = componentsets(s) + isempty(x) && isempty(sets) ? true : all(in.(x, sets))::Bool +end - # Get the normal vectors - g1 = a.∇f(x) - g2 = b.∇f(x) +@inline Base.isempty(s::CartesianProduct) = any(isempty, componentsets(s)) - # See if one is a multiple of the other - one(T) - Statistics.corm(g1, zero(T), g2, zero(T)) < tol +@inline function Base.union( + s::CartesianProduct{<:Tuple{Vararg{Any,N}}}, + others::CartesianProduct{<:Tuple{Vararg{Any,N}}}..., +) where {N} + CartesianProduct(map(union, componentsets(s), map(componentsets, others)...)) end -function zeroset(::CodimOne)::ZeroSet end +@inline function Base.union( + s::CartesianProduct{<:NamedTuple{names}}, + others::CartesianProduct{<:NamedTuple{names}}..., +) where {names} + CartesianProduct(map(union, componentsets(s), map(componentsets, others)...)) +end -########################################################### -# Simplex -export Simplex +function Base.union( + s::CartesianProduct{<:AbstractArray{<:Any,N}}, + others::CartesianProduct{<:AbstractArray{<:Any,N}}..., +) where {N} + CartesianProduct(union.(componentsets(s), map(componentsets, others)...)) +end -struct Simplex <: CodimOne end +""" + struct CartesianPower <: AbstractCartSetProd -function zeroset(::Simplex) - f(x::AbstractArray{T}) where {T} = sum(x) - one(T) - ∇f(x::AbstractArray{T}) where {T} = fill_with(one(T), size(x)) - ZeroSet(f, ∇f) +Represents the n-fold Cartesian product of a set. +""" +struct CartesianPower{S,A} <: AbstractCartSetProd + _base::S + _axes::A end -function Base.in(x::AbstractArray{T}, ::Simplex) where {T} - all(≥(zero(eltype(x))), x) || return false - return x ∈ zeroset(Simplex()) +@inline setcartpower(s::SetLike, dims) = CartesianPower(s, asaxes(dims)) + +@inline pwr_base(s::CartesianPower) = s._base +@inline pwr_axes(s::CartesianPower) = s._axes +@inline pwr_size(s::CartesianPower) = axes2size(pwr_axes(s)) + +componentsets(s::CartesianPower) = maybestatic_fill(pwr_base(s), pwr_axes(s)) + +function Base.in(x::AbstractArray, s::CartesianPower) + pwr_size(s) == size(x) || + throw(ArgumentError("Size of CartesianPower and given point are incompatible.")) + isempty(x) ? true : all(Base.Fix1(in, pwr_base(s)), x)::Bool end -projectto!(x, ::Simplex) = normalize!(x, 1) +Base.isempty(s::CartesianPower) = isempty(pwr_base(s)) || size2length(pwr_size(s)) == 0 -########################################################### -# Sphere +function Base.union(s::CartesianPower, others::CartesianPower...) + axs = pwr_axes(s) -struct Sphere <: CodimOne end + all(isequal(axs), map(pwr_axes, others)) || throw( + ArgumentError("Cannot create union of CartesianPower sets with different axes."), + ) -function zeroset(::Sphere) - f(x::AbstractArray{T}) where {T} = dot(x, x) - one(T) - ∇f(x::AbstractArray{T}) where {T} = x - ZeroSet(f, ∇f) + setcartpower(union(pwr_base(s), map(pwr_base, others)...), axs) end -function Base.in(x::AbstractArray{T}, ::Sphere) where {T} - return x ∈ zeroset(Sphere()) -end -projectto!(x, ::Sphere) = normalize!(x, 2) +""" + struct CombinedSet <: ValueSet + +Represents a combination of two sets. + +User code should not create instances of `CombinedMeasure` directly, but should call +[`combinesets(f_c, α, β)`](@ref) instead. +""" +struct CombinedSet{FC,MA<:SetLike,MB<:SetLike} <: ValueSet + f_c::FC + α::MA + β::MB +end diff --git a/src/getdof.jl b/src/getdof.jl index dbce2202..16ae7cc6 100644 --- a/src/getdof.jl +++ b/src/getdof.jl @@ -66,7 +66,7 @@ function checked_arg end # Prevent infinite recursion: @propagate_inbounds function _default_checked_arg(::Type{MU}, ::MU, ::T) where {MU,T} - NoArgCheck{MU,T} + NoArgCheck{MU,T}() end @propagate_inbounds function _default_checked_arg(::Type{MU}, mu_base, x) where {MU} checked_arg(mu_base, x) diff --git a/src/interface.jl b/src/interface.jl index 18080ac7..4890ddd6 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -110,7 +110,7 @@ function test_smf(μ, n = 100) @testset "smf($μ)" begin # Get `n` sorted uniforms in O(n) time p = rand(n) - p .+= 0:n-1 + p .+= 0:(n-1) p .*= inv(n) F(x) = smf(μ, x) diff --git a/src/primitives/counting.jl b/src/primitives/counting.jl index c61d0624..31398f32 100644 --- a/src/primitives/counting.jl +++ b/src/primitives/counting.jl @@ -40,3 +40,8 @@ insupport(μ::Counting{T}, x) where {T<:Type} = x isa μ.support massof(c::Counting, s::Set) = massof(CountingBase(), filter(insupport(c), s)) massof(::CountingBase, s::Set) = length(s) + +# ToDo: Would this be correct? +# @inline mdomain(::CountingBase) = IntegerValues() + +@inline mdomain(::Counting{DomainType}) where {DomainType} = DomainType() diff --git a/src/primitives/lebesgue.jl b/src/primitives/lebesgue.jl index 3846eaf5..3d92a2ed 100644 --- a/src/primitives/lebesgue.jl +++ b/src/primitives/lebesgue.jl @@ -51,7 +51,7 @@ Lebesgue() = Lebesgue(ℝ) testvalue(::Type{T}, d::Lebesgue) where {T} = testvalue(T, d.support)::T proxy(d::Lebesgue) = restrict(in(d.support), LebesgueBase()) -proxy(::Lebesgue{MeasureBase.RealNumbers}) = LebesgueBase() +proxy(::Lebesgue{MeasureBase.RealValues}) = LebesgueBase() @useproxy Lebesgue @@ -61,7 +61,7 @@ Base.show(io::IO, d::Lebesgue) = print(io, "Lebesgue(", d.support, ")") insupport(μ::Lebesgue, x) = x ∈ μ.support -insupport(::Lebesgue{RealNumbers}, ::Real) = true +insupport(::Lebesgue{RealValues}, ::Real) = true @inline function logdensityof(μ::Lebesgue, x::Real) R = float(typeof(x)) @@ -70,28 +70,28 @@ end @inline logdensityof(μ::Lebesgue, x) = insupport(μ, x) ? 0.0 : -Inf -massof(::Lebesgue{RealNumbers}, s::Interval) = width(s) +massof(::Lebesgue{RealValues}, s::Interval) = width(s) # Example: # julia> Lebesgue(𝕀)(0.2..5) # 0.8 -function massof(μ::Lebesgue{<:BoundedReals}, s::Interval) - a = μ.support.lower - b = μ.support.upper +function massof(μ::Lebesgue{<:AbstractInterval}, s::Interval) + a, b = endpoints(μ.support) left = max(s.left, a) right = min(s.right, b) w = right - left max(w, zero(w)) end -function smf(μ::Lebesgue{<:BoundedReals}, x) - clamp(x, μ.support.lower, μ.support.upper) +function smf(μ::Lebesgue{<:AbstractInterval}, x) + a, b = endpoints(μ.support) + clamp(x, a, b) end -smf(::Lebesgue{RealNumbers}, x) = x -smf(::Lebesgue{RealNumbers}) = identity -invsmf(::Lebesgue{RealNumbers}, x) = x -invsmf(::Lebesgue{RealNumbers}) = identity +smf(::Lebesgue{<:RealValues}, x) = x +smf(::Lebesgue{<:RealValues}) = identity +invsmf(::Lebesgue{<:RealValues}, x) = x +invsmf(::Lebesgue{<:RealValues}) = identity smf(::LebesgueBase, x) = x smf(::LebesgueBase) = identity diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl index 833f280e..a9c09b5c 100644 --- a/src/standard/stdmeasure.jl +++ b/src/standard/stdmeasure.jl @@ -13,7 +13,7 @@ function transport_def(ν::StdMeasure, μ::PowerMeasure{<:StdMeasure}, x) end function transport_def(ν::PowerMeasure{<:StdMeasure}, μ::StdMeasure, x) - return fill_with(transport_def(ν.parent, μ, only(x)), map(length, ν.axes)) + return maybestatic_fill(transport_def(ν.parent, μ, only(x)), map(length, ν.axes)) end function transport_def( @@ -57,7 +57,7 @@ end # Helpers for product transforms and similar: struct _TransportToStd{NU<:StdMeasure} <: Function end -_TransportToStd{NU}(μ, x) where {NU} = transport_to(NU()^getdof(μ), μ)(x) +(::_TransportToStd{NU})(μ, x) where {NU} = transport_to(NU()^getdof(μ), μ)(x) struct _TransportFromStd{MU<:StdMeasure} <: Function end _TransportFromStd{MU}(ν, x) where {MU} = transport_to(ν, MU()^getdof(ν))(x) @@ -67,7 +67,7 @@ function _tuple_transport_def( μs::Tuple, xs::Tuple, ) where {NU<:StdMeasure} - reshape(vcat(map(_TransportToStd{NU}, μs, xs)...), ν.axes) + reshape(vcat(map(_TransportToStd{NU}(), μs, xs)...), ν.axes) end function transport_def( @@ -93,7 +93,7 @@ end function _stdvar_viewranges(μs::Tuple, startidx::IntegerLike) N = map(getdof, μs) offs = _offset_cumsum(startidx, N...) - map((o, n) -> o:o+n-1, offs, N) + map((o, n) -> o:(o+n-1), offs, N) end function _tuple_transport_def( @@ -121,3 +121,28 @@ function transport_def( ) where {MU<:StdMeasure,names} NamedTuple{names}(_tuple_transport_def(values(marginals(ν)), μ, x)) end + +function transport_def( + ν::PowerMeasure{NU}, + μ::ProductMeasure{<:AbstractArray}, + x, +) where {NU<:StdMeasure} + reshape(vcat(map(_TransportToStd{NU}(), marginals(μ), x)...), ν.axes) +end + +function _marginal_viewranges(μs::AbstractArray, startidx::IntegerLike) + ns = map(m -> dynamic(getdof(m)), μs) + offs = cumsum(vcat(dynamic(startidx), ns[begin:(end-1)])) + map((o, n) -> o:(o+n-1), offs, ns) +end + +function transport_def( + ν::ProductMeasure{<:AbstractArray}, + μ::PowerMeasure{MU}, + x::AbstractArray{<:Real}, +) where {MU<:StdMeasure} + νs = marginals(ν) + vrs = _marginal_viewranges(νs, firstindex(x)) + xs = map(r -> view(x, r), vrs) + map(_TransportFromStd{MU}, νs, xs) +end diff --git a/src/standard/stduniform.jl b/src/standard/stduniform.jl index 7bbe15ed..e3702656 100644 --- a/src/standard/stduniform.jl +++ b/src/standard/stduniform.jl @@ -14,7 +14,7 @@ end Base.rand(rng::Random.AbstractRNG, ::Type{T}, ::StdUniform) where {T} = rand(rng, T) -massof(::StdUniform, s::Interval) = massof(Lebesgue(𝕀), s::Interval) +massof(::StdUniform, s::Interval) = massof(Lebesgue(0.0 .. 1.0), s) smf(::StdUniform, x) = clamp(x, zero(x), one(x)) diff --git a/src/static.jl b/src/static.jl deleted file mode 100644 index da471b62..00000000 --- a/src/static.jl +++ /dev/null @@ -1,63 +0,0 @@ -""" - MeasureBase.IntegerLike - -Equivalent to `Union{Integer,Static.StaticInteger}`. -""" -const IntegerLike = Union{Integer,Static.StaticInteger} - -""" - MeasureBase.one_to(n::IntegerLike) - -Creates a range from one to n. - -Returns an instance of `Base.OneTo` or `Static.SOneTo`, depending -on the type of `n`. -""" -@inline one_to(n::Integer) = Base.OneTo(n) -@inline one_to(::Static.StaticInteger{N}) where {N} = Static.SOneTo{N}() - -_dynamic(x::Number) = dynamic(x) -_dynamic(::Static.SOneTo{N}) where {N} = Base.OneTo(N) -_dynamic(r::AbstractUnitRange) = minimum(r):maximum(r) - -""" - MeasureBase.fill_with(x, sz::NTuple{N,<:IntegerLike}) where N - -Creates an array of size `sz` filled with `x`. - -Returns an instance of `FillArrays.Fill`. -""" -function fill_with end - -@inline function fill_with(x::T, sz::Tuple{Vararg{IntegerLike,N}}) where {T,N} - fill_with(x, map(one_to, sz)) -end - -@inline function fill_with(x::T, axs::Tuple{Vararg{AbstractUnitRange,N}}) where {T,N} - # While `FillArrays.Fill` (mostly?) works with axes that are static unit - # ranges, some operations that automatic differentiation requires do fail - # on such instances of `Fill` (e.g. `reshape` from dynamic to static size). - # So need to use standard ranges for the axes for now: - dyn_axs = map(_dynamic, axs) - FillArrays.Fill(x, dyn_axs) -end - -""" - MeasureBase.maybestatic_length(x)::IntegerLike - -Returns the length of `x` as a dynamic or static integer. -""" -maybestatic_length(x) = length(x) -maybestatic_length(x::AbstractUnitRange) = length(x) -function maybestatic_length( - ::Static.OptionallyStaticUnitRange{<:StaticInteger{A},<:StaticInteger{B}}, -) where {A,B} - StaticInt{B - A + 1}() -end - -""" - MeasureBase.maybestatic_size(x)::Tuple{Vararg{IntegerLike}} - -Returns the size of `x` as a tuple of dynamic or static integers. -""" -maybestatic_size(x) = size(x) diff --git a/src/utils.jl b/src/utils.jl index 0ec81a50..c1e97034 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -133,6 +133,11 @@ function infer_zero(f, args...) zero(typeintersect(AbstractFloat, inferred_type)) end +function infer_logdensity_type(f::F, ::M, ::Type{T}) where {F,M,T} + inferred_type = Core.Compiler.return_type(f, Tuple{M,T}) + return inferred_type +end + @inline function allequal(f, x::AbstractArray) val = f(first(x)) @simd for xj in x @@ -180,3 +185,54 @@ isapproxzero(A::AbstractArray) = all(isapproxzero, A) isapproxone(x::T) where {T<:Real} = x ≈ one(T) isapproxone(A::AbstractArray) = all(isapproxone, A) + +containsnan(x::Real) = isnan(x) +containsnan(x) = any(containsnan, x) + + +# ForwardDiffPullbacks dummy `fwddiff`, overloaded by +# ForwardDiffPullbacks extension when loaded: +@inline _fwddiff(f) = f + + +# Autodiff ignore: + +@inline _adignore_call(f) = f() + +macro _adignore(expr) + :(_adignore_call(() -> $(esc(expr)))) +end + + +""" + MeasureBase.convert_realtype(::Type{T}, x) where {T<:Real} + +Convert `x` to use `T` as its underlying type for real numbers. +""" +function convert_realtype end + +@inline convert_realtype(::Type{T}, x::T) where {T<:Real} = x +@inline convert_realtype(::Type{T}, x::AbstractArray{T}) where {T<:Real} = x +@inline convert_realtype(::Type{T}, x::U) where {T<:Real,U<:Real} = T(x) +convert_realtype(::Type{T}, x::AbstractArray{U}) where {T<:Real,U<:Real} = T.(x) +convert_realtype(::Type{T}, x::Union{Tuple,NamedTuple}) where {T<:Real} = + map(Base.Fix1(convert_realtype, T), x) +convert_realtype(::Type{T}, x::AbstractArray) where {T<:Real} = + map(Base.Fix1(convert_realtype, T), x) + +""" + MeasureBase.firsttype(::Type{T}, ::Type{U}) where {T<:Real,U<:Real} + +Return the first type, but as a dual number type if the second one is dual. +""" +function firsttype end + +firsttype(::Type{T}, ::Type{U}) where {T<:Real,U<:Real} = T + + + +# Distributions implementation hooks: +function _trafo_cdf_impl end +function _trafo_quantile_impl end +function _trafo_quantile_impl_generic end +function _dist_params_numtype end diff --git a/test/Project.toml b/test/Project.toml index 376c1b05..cb68fc93 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,21 +1,29 @@ [deps] AffineMaps = "2c83c9a8-abf5-4329-a0d7-deffaf474661" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +ArraysOfArrays = "65a8f2f4-9b39-5baf-92e2-a9cc46fdf018" +ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0" DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +ForwardDiffPullbacks = "450a3b6d-2448-4ee1-8e34-e4eb8713b605" InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112" IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" LogarithmicNumbers = "aa2f6b4e-9042-5d33-9679-40d3a6b85899" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" +PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" PropertyFunctions = "09e99361-2bb8-48a2-a80f-de58f0739eb4" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" +StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" diff --git a/test/combinators/reshape.jl b/test/combinators/reshape.jl new file mode 100644 index 00000000..c6624582 --- /dev/null +++ b/test/combinators/reshape.jl @@ -0,0 +1,7 @@ +using Test + +using MeasureBase + +@testset "reshape" begin + +end diff --git a/test/distributions/getjacobian.jl b/test/distributions/getjacobian.jl new file mode 100644 index 00000000..87de7b86 --- /dev/null +++ b/test/distributions/getjacobian.jl @@ -0,0 +1,34 @@ +# This file is a part of ChangesOfVariables.jl, licensed under the MIT License (MIT). + +import ForwardDiff + +torv_and_back(V::AbstractVector{<:Real}) = V, identity +torv_and_back(x::Real) = [x], V -> V[1] +torv_and_back(x::Complex) = [real(x), imag(x)], V -> Complex(V[1], V[2]) +torv_and_back(x::NTuple{N}) where N = [x...], V -> ntuple(i -> V[i], Val(N)) + +function torv_and_back(x::Ref) + xval = x[] + V, to_xval = torv_and_back(xval) + back_to_ref(V) = Ref(to_xval(V)) + return (V, back_to_ref) +end + +torv_and_back(A::AbstractArray{<:Real}) = vec(A), V -> reshape(V, size(A)) + +function torv_and_back(A::AbstractArray{Complex{T}, N}) where {T<:Real, N} + RA = cat(real.(A), imag.(A), dims = N+1) + V, to_array = torv_and_back(RA) + function back_to_complex(V) + RA = to_array(V) + Complex.(view(RA, map(_ -> :, size(A))..., 1), view(RA, map(_ -> :, size(A))..., 2)) + end + return (V, back_to_complex) +end + + +function getjacobian(f, x) + V, to_x = torv_and_back(x) + vf(V) = torv_and_back(f(to_x(V)))[1] + ForwardDiff.jacobian(vf, V) +end diff --git a/test/distributions/test_autodiff_utils.jl b/test/distributions/test_autodiff_utils.jl new file mode 100644 index 00000000..5197725b --- /dev/null +++ b/test/distributions/test_autodiff_utils.jl @@ -0,0 +1,18 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using LinearAlgebra +using Distributions, ArraysOfArrays +import ForwardDiff, Zygote + + +@testset "trafo_utils" begin + xs = rand(5) + @test Zygote.jacobian(MeasureBase._pushfront, xs, 42)[1] ≈ ForwardDiff.jacobian(xs -> MeasureBase._pushfront(xs, 1), xs) + @test Zygote.jacobian(MeasureBase._pushfront, xs, 42)[2] ≈ vec(ForwardDiff.jacobian(x -> MeasureBase._pushfront(xs, x[1]), [42])) + @test Zygote.jacobian(MeasureBase._pushback, xs, 42)[1] ≈ ForwardDiff.jacobian(xs -> MeasureBase._pushback(xs, 1), xs) + @test Zygote.jacobian(MeasureBase._pushback, xs, 42)[2] ≈ vec(ForwardDiff.jacobian(x -> MeasureBase._pushback(xs, x[1]), [42])) + @test Zygote.jacobian(MeasureBase._rev_cumsum, xs)[1] ≈ ForwardDiff.jacobian(MeasureBase._rev_cumsum, xs) + @test Zygote.jacobian(MeasureBase._exp_cumsum_log, xs)[1] ≈ ForwardDiff.jacobian(MeasureBase._exp_cumsum_log, xs) ≈ ForwardDiff.jacobian(cumprod, xs) +end diff --git a/test/distributions/test_conversions.jl b/test/distributions/test_conversions.jl new file mode 100644 index 00000000..0f9c4d80 --- /dev/null +++ b/test/distributions/test_conversions.jl @@ -0,0 +1,113 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using Random, Statistics, LinearAlgebra +using Distributions +using StableRNGs + +import MeasureBase +using MeasureBase: AbstractMeasure, AsMeasure, asmeasure +using MeasureBase: StdUniform, StdNormal, StdExponential, StdLogistic +using MeasureBase: SuperpositionMeasure, PushforwardMeasure, ProductMeasure +using MeasureBase: logdensityof, massof, insupport + + +@testset "conversions" begin + stblrng() = StableRNG(789990641) + + function test_conversion(d::Distribution, ::Type{M}) where {M} + @testset "conversion $(typeof(d).name) <-> $M" begin + m = asmeasure(d) + @test m isa M + @test typeof(convert(AbstractMeasure, d)) === typeof(m) + @test_throws ArgumentError AsMeasure{typeof(d)}(d) + + d2 = convert(Distribution, m) + @test d2 isa Distribution + @test typeof(Distributions.Distribution(m)) === typeof(d2) + + for x in (rand(stblrng(), d) for _ in 1:10) + @test logdensityof(m, x) ≈ logpdf(d, x) + @test logpdf(d2, x) ≈ logpdf(d, x) + @test insupport(m, x) + end + + x = rand(stblrng(), Float64, m) + # Tuple-marginal product measures have tuple variates: + x isa Tuple ? (@test length(x) == length(d)) : (@test size(x) == size(d)) + @test insupport(m, x) + end + end + + @testset "Dirac" begin + d = Distributions.Dirac(4.2) + m = @inferred asmeasure(d) + @test m === MeasureBase.Dirac(4.2) + @test_throws ArgumentError AsMeasure{typeof(d)}(d) + @test @inferred(Distributions.Distribution(m)) === d + end + + @testset "products" begin + test_conversion(product_distribution(Weibull.([0.7, 1.1, 1.3])), ProductMeasure) + test_conversion(product_distribution(Poisson.([0.7, 1.4])), ProductMeasure) + + if isdefined(Distributions, :ProductDistribution) + test_conversion(product_distribution(Weibull(0.7), Exponential(1.3)), ProductMeasure) + end + end + + @testset "reshaped" begin + test_conversion(reshape(MvNormal([0.7, 0.9], [1.4 0.5; 0.5 1.1]), 1, 2), PushforwardMeasure) + test_conversion(reshape(product_distribution(Weibull.([0.7, 1.1, 1.3, 0.9, 1.2, 0.8])), 2, 3), PushforwardMeasure) + end + + @testset "mixtures" begin + test_conversion(MixtureModel([Normal(-1.0, 1.0), Normal(2.0, 3.0)], [0.3, 0.7]), SuperpositionMeasure) + test_conversion(MixtureModel([Normal(-2.0, 1.0), Normal(0.0, 2.0), Normal(3.0, 1.0)], [0.2, 0.5, 0.3]), SuperpositionMeasure) + test_conversion(MixtureModel([Exponential(0.3), Weibull(2.0, 1.0)], [0.4, 0.6]), SuperpositionMeasure) + test_conversion(MixtureModel([MvNormal([0.0, 0.0], I(2)), MvNormal([2.0, 2.0], 2 * I(2))], [0.3, 0.7]), SuperpositionMeasure) + test_conversion(UnivariateGMM([-1.0, 2.0], [1.0, 0.5], Categorical([0.4, 0.6])), SuperpositionMeasure) + + d = MixtureModel([Normal(-1.0, 1.0), Normal(2.0, 3.0)], [0.3, 0.7]) + m = asmeasure(d) + @test probs(convert(Distribution, m)) ≈ probs(d) + @test massof(m) ≈ 1 + @test massof(asmeasure(Normal())) == 1 + @test mean(rand(stblrng(), Float64, m^1000)) ≈ mean(d) atol = 0.3 + + # Hand-built superpositions of weighted probability measures behave + # like mixtures: + m2 = 0.3 * asmeasure(Normal(-1.0, 1.0)) + 0.7 * asmeasure(Normal(2.0, 3.0)) + for x in (rand(stblrng(), d) for _ in 1:10) + @test logdensityof(m2, x) ≈ logpdf(d, x) + end + end + + @testset "standard distributions" begin + @test StandardUniform === StandardDist{Uniform} + @test StandardNormal === StandardDist{Normal} + @test StandardUniform{0} === StandardDist{Uniform,0} + @test StandardNormal{1} === StandardDist{Normal,1} + + for (D, B) in [ + (Uniform, StdUniform()), + (Exponential, StdExponential()), + (Logistic, StdLogistic()), + (Normal, StdNormal()), + ] + @test @inferred(asmeasure(StandardDist{D}())) === B + @test @inferred(asmeasure(StandardDist{D}(3))) == B^3 + @test @inferred(asmeasure(StandardDist{D}(2, 3))) == B^(2, 3) + + @test @inferred(Distributions.Distribution(B)) === StandardDist{D}() + @test @inferred(convert(Distribution, B)) === StandardDist{D}() + @test @inferred(Distributions.Distribution(B^3)) == StandardDist{D}(3) + @test @inferred(convert(Distribution, B^(2, 3))) == StandardDist{D}(2, 3) + + d = StandardDist{D}(3) + x = rand(stblrng(), d) + @test logdensityof(asmeasure(d), x) ≈ logpdf(d, x) + end + end +end diff --git a/test/distributions/test_distribution_measure.jl b/test/distributions/test_distribution_measure.jl new file mode 100644 index 00000000..7715e58a --- /dev/null +++ b/test/distributions/test_distribution_measure.jl @@ -0,0 +1,53 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +import Distributions +using Distributions: Distribution +import MeasureBase +using MeasureBase: AbstractMeasure + +@testset "Measure interface" begin + d = Distributions.Weibull() + @test @inferred(AbstractMeasure(d)) isa AbstractMeasure + @test @inferred(AbstractMeasure(d)) isa DistributionMeasure + @test @inferred(convert(AbstractMeasure, d)) isa AbstractMeasure + @test @inferred(convert(AbstractMeasure, d)) isa DistributionMeasure + @test @inferred(Distribution(AbstractMeasure(d))) === d + @test @inferred(convert(Distribution, convert(AbstractMeasure, d))) === d + + + c0 = AbstractMeasure(Distributions.Weibull(0.7, 1.3)) + c1 = AbstractMeasure(Distributions.MvNormal([0.7, 0.9], [1.4 0.5; 0.5 1.1])) + + d0 = AbstractMeasure(Distributions.Poisson(0.7)) + d1 = AbstractMeasure(Distributions.product_distribution(Distributions.Poisson.([0.7, 1.4]))) + + for μ in [c0, c1, d0, d1] + d = Distribution(μ) + x = rand(μ) + @test @inferred(MeasureBase.logdensity_def(μ, x)) == Distributions.logpdf(d, x) + @test @inferred(MeasureBase.unsafe_logdensityof(μ, x)) == Distributions.logpdf(d, x) + + MeasureBase.Interface.test_interface(d) + end + + @test @inferred(MeasureBase.basemeasure(c0)) == MeasureBase.Lebesgue(MeasureBase.ℝ) + @test @inferred(MeasureBase.basemeasure(c1)) == MeasureBase.Lebesgue(MeasureBase.ℝ) ^ 2 + + @test @inferred(MeasureBase.insupport(c0, 3)) == true + @test @inferred(MeasureBase.insupport(c0, -3)) == false + @test @inferred(MeasureBase.insupport(c1, [0.1, 0.2])) == true + @test @inferred(MeasureBase.insupport(d0, 3)) == true + @test @inferred(MeasureBase.insupport(d0, 3.2)) == false + @test @inferred(MeasureBase.insupport(d1, [1, 2])) == true + @test @inferred(MeasureBase.insupport(d1, [1.1, 2.2])) == false + + @test MeasureBase.paramnames(c0) == (:α, :θ) + if VERSION >= v"1.8" + @test @inferred(MeasureBase.params(c0)) == (α = 0.7, θ = 1.3) + else + # v1.6 can't type-infer this: + @test (MeasureBase.params(c0)) == (α = 0.7, θ = 1.3) + end +end diff --git a/test/distributions/test_distributions.jl b/test/distributions/test_distributions.jl new file mode 100644 index 00000000..65fab330 --- /dev/null +++ b/test/distributions/test_distributions.jl @@ -0,0 +1,24 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test +using MeasureBase +using Distributions +import ForwardDiff, ForwardDiffPullbacks, ChainRulesCore + +const MeasureBaseDistributionsExt = Base.get_extension(MeasureBase, :MeasureBaseDistributionsExt) +@test MeasureBaseDistributionsExt isa Module + +using .MeasureBaseDistributionsExt: + StandardDist, StandardUniform, StandardNormal, DistributionMeasure, nonstddist + +@testset "Distributions extension" begin + include("test_autodiff_utils.jl") + include("test_measure_interface.jl") + include("test_distribution_measure.jl") + include("test_standard_dist.jl") + include("test_standard_uniform.jl") + include("test_standard_normal.jl") + include("test_conversions.jl") + include("test_transport.jl") + include("test_mooncake.jl") +end diff --git a/test/distributions/test_measure_interface.jl b/test/distributions/test_measure_interface.jl new file mode 100644 index 00000000..d11d2889 --- /dev/null +++ b/test/distributions/test_measure_interface.jl @@ -0,0 +1,43 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +import Distributions +import MeasureBase + +@testset "Measure interface" begin + c0 = Distributions.Weibull(0.7, 1.3) + c1 = Distributions.MvNormal([0.7, 0.9], [1.4 0.5; 0.5 1.1]) + + d0 = Distributions.Poisson(0.7) + d1 = Distributions.product_distribution(Distributions.Poisson.([0.7, 1.4])) + + for d in [c0, c1, d0, d1] + x = rand(d) + @test @inferred(MeasureBase.logdensity_def(d, x)) == Distributions.logpdf(d, x) + @test @inferred(MeasureBase.unsafe_logdensityof(d, x)) == Distributions.logpdf(d, x) + + MeasureBase.Interface.test_interface(d) + end + + @test @inferred(MeasureBase.basemeasure(c0)) == MeasureBase.Lebesgue(MeasureBase.ℝ) + @test @inferred(MeasureBase.basemeasure(c1)) == MeasureBase.Lebesgue(MeasureBase.ℝ) ^ 2 + + @test @inferred(MeasureBase.insupport(c0, 3)) == true + @test @inferred(MeasureBase.insupport(c0, -3)) == false + @test @inferred(MeasureBase.insupport(c1, [0.1, 0.2])) == true + @test @inferred(MeasureBase.insupport(d0, 3)) == true + @test @inferred(MeasureBase.insupport(d0, 3.2)) == false + @test @inferred(MeasureBase.insupport(d1, [1, 2])) == true + @test @inferred(MeasureBase.insupport(d1, [1.1, 2.2])) == false + + @test MeasureBase.paramnames(c0) == (:α, :θ) + if VERSION >= v"1.8" + @test @inferred(MeasureBase.params(c0)) == (α = 0.7, θ = 1.3) + else + # v1.6 can't type-infer this: + @test (MeasureBase.params(c0)) == (α = 0.7, θ = 1.3) + end + + @test MeasureBase.∫(x -> Distributions.Normal(x, 0), Distributions.Normal()) isa MeasureBase.DensityMeasure +end diff --git a/test/distributions/test_mooncake.jl b/test/distributions/test_mooncake.jl new file mode 100644 index 00000000..4a20e46c --- /dev/null +++ b/test/distributions/test_mooncake.jl @@ -0,0 +1,68 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using Random, LinearAlgebra +using Distributions +import Mooncake +import ForwardDiff + +using MeasureBase +using MeasureBase: transport_to, transport_def, asmeasure +using MeasureBase: StdUniform, StdNormal + +_mooncake_gradient(f, x) = Mooncake.value_and_gradient!!( + Mooncake.prepare_gradient_cache(f, x), f, x +)[2][2] + +_test_gradient(f, x::Real) = @test _mooncake_gradient(f, x) ≈ ForwardDiff.derivative(f, x) +_test_gradient(f, x::AbstractVector) = @test _mooncake_gradient(f, x) ≈ ForwardDiff.gradient(f, x) + +@testset "Mooncake AD with Distributions" begin + @test Base.get_extension(MeasureBase, :MeasureBaseDistributionsMooncakeExt) isa Module + + @testset "zero-derivative primitives" begin + rng = Random.Xoshiro(789990641) + Mooncake.TestUtils.test_rule( + rng, MeasureBase._dist_params_numtype, Normal(0.2, 1.3); + is_primitive = true, + ) + end + + @testset "univariate transport gradients" begin + _test_gradient(x -> transport_def(StdUniform(), Normal(1.0, 2.0), x), 0.5) + _test_gradient(u -> transport_def(Normal(1.0, 2.0), StdUniform(), u), 0.3) + _test_gradient(u -> transport_def(Beta(2.0, 3.0), StdUniform(), u), 0.3) + _test_gradient(x -> transport_def(StdUniform(), Gamma(2.0, 1.0), x), 0.7) + _test_gradient(x -> transport_def(StdUniform(), truncated(Normal(0.3, 1.2), -0.5, 1.5), x), 0.4) + _test_gradient(x -> transport_def(StdUniform(), 2.0 * Weibull(0.7) + 1.0, x), 3.0) + _test_gradient(x -> transport_def(StdNormal(), StandardDist{Uniform}(), x), 0.4) + end + + @testset "multivariate transport gradients" begin + mvn = MvNormal([0.3, -2.9], [1.7 0.5; 0.5 2.3]) + _test_gradient(x -> sum(transport_to(StdNormal()^2, mvn)(x)), [0.1, -2.0]) + _test_gradient(y -> sum(transport_to(mvn, StdNormal()^2)(y)), [0.2, 0.7]) + + pd = product_distribution([Weibull(0.7), Exponential(1.3), Normal(0.5, 2.0)]) + _test_gradient(x -> sum(transport_to(StdNormal()^3, asmeasure(pd))(x)), [0.4, 0.8, 1.5]) + + dirich = Dirichlet([2.0, 3.0, 4.0]) + _test_gradient(u -> MeasureBase.from_origin(dirich, u)[1], [0.3, 0.7]) + _test_gradient(x -> sum(MeasureBase.to_origin(dirich, vcat(x, 1 - sum(x)))), [0.28, 0.23]) + end + + @testset "logdensityof gradients" begin + for d in [ + Weibull(0.7, 1.3), + MixtureModel([Normal(-1.0, 1.0), Normal(2.0, 3.0)], [0.3, 0.7]), + MixtureModel([Normal(-2.0, 1.0), Normal(0.0, 2.0), Normal(3.0, 1.0)], [0.2, 0.5, 0.3]), + ] + m = asmeasure(d) + _test_gradient(x -> logdensityof(m, x[1]), [0.5]) + end + + mvn = MvNormal([0.3, -2.9], [1.7 0.5; 0.5 2.3]) + _test_gradient(x -> logdensityof(asmeasure(mvn), x), [0.1, -2.0]) + end +end diff --git a/test/distributions/test_standard_dist.jl b/test/distributions/test_standard_dist.jl new file mode 100644 index 00000000..64b9f655 --- /dev/null +++ b/test/distributions/test_standard_dist.jl @@ -0,0 +1,128 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using Random, Statistics, LinearAlgebra +using Distributions, PDMats +using StableRNGs +import ForwardDiff, ChainRulesTestUtils + + +@testset "standard_dist" begin + stblrng() = StableRNG(789990641) + + for (D, sz, dref) in [ + (Uniform, (), Uniform()), + (Uniform, (5,), product_distribution(fill(Uniform(0.0, 1.0), 5))), + (Uniform, (2, 3), reshape(product_distribution(fill(Uniform(0.0, 1.0), 6)), 2, 3)), + (Normal, (), Normal()), + (Normal, (), Normal(0., 1.0)), + (Normal, (5,), MvNormal(Diagonal(fill(1.0, 5)))), + (Normal, (2, 3), reshape(MvNormal(Diagonal(fill(1.0, 6))), 2, 3)), + (Exponential, (), Exponential()), + (Exponential, (5,), product_distribution(fill(Exponential(1.0), 5))), + (Exponential, (2, 3), reshape(product_distribution(fill(Exponential(1.0), 6)), 2, 3)), + ] + @testset "StandardDist{$D}($(join(sz,",")))" begin + N = length(sz) + + @test @inferred(StandardDist{D}(sz...)) isa StandardDist{D} + @test @inferred(StandardDist{D}(sz...)) isa StandardDist{D} + @test @inferred(size(StandardDist{D}(sz...))) == size(dref) + @test @inferred(size(StandardDist{D}(sz...))) == size(dref) + + d = StandardDist{D}(sz...) + + if size(d) == () + @test @inferred(MeasureBaseDistributionsExt.nonstddist(d)) == dref + end + + @test @inferred(length(d)) == length(dref) + @test @inferred(size(d)) == size(dref) + + @test @inferred(eltype(typeof(d))) == eltype(typeof(dref)) + @test @inferred(eltype(d)) == eltype(dref) + + @test @inferred(Distributions.params(d)) == () + @test @inferred(partype(d)) == partype(dref) + + for f in [minimum, maximum, mean, median, mode, modes, var, std, skewness, kurtosis, location, scale, entropy] + supported_by_dref = try f(dref); true catch MethodError; false; end + if supported_by_dref + @test @inferred(f(d)) ≈ f(dref) + end + end + + for x in [rand(dref) for i in 1:10] + ref_gradlogpdf = try + gradlogpdf(dref, x) + catch MethodError + ForwardDiff.gradient(x -> logpdf(dref, x), x) + end + @test @inferred(gradlogpdf(d, x)) ≈ ref_gradlogpdf + @test @inferred(logpdf(d, x)) ≈ logpdf(dref, x) + @test @inferred(pdf(d, x)) ≈ pdf(dref, x) + end + + if size(d) == () + for x in [minimum(dref), quantile(dref, 1//3), quantile(dref, 1//2), quantile(dref, 2//3), maximum(dref)] + for f in [logpdf, pdf, gradlogpdf, logcdf, cdf, logccdf, ccdf] + @test @inferred(f(d, x)) ≈ f(dref, x) + end + end + + for x in [0, 1//3, 1//2, 2//3, 1] + for f in [quantile, cquantile] + @test @inferred(f(d, x)) ≈ f(dref, x) + end + end + + for x in log.([0, 1//3, 1//2, 2//3, 1]) + for f in [invlogcdf, invlogccdf] + @test @inferred(f(d, x)) ≈ f(dref, x) + end + end + + for p in [0.0, 0.25, 0.75, 1.0] + @test @inferred(quantile(d, p)) == quantile(dref, p) + @test @inferred(cquantile(d, p)) == cquantile(dref, p) + end + + for t in [-3, 0, 3] + @test isapprox(@inferred(mgf(d, t)), mgf(dref, t), rtol = 1e-5) + @test isapprox(@inferred(cf(d, t)), cf(dref, t), rtol = 1e-5) + end + + @test @inferred(truncated(d, quantile(dref, 1//3), quantile(dref, 2//3))) == truncated(dref, quantile(dref, 1//3), quantile(dref, 2//3)) + + @test @inferred(product_distribution(fill(d, 3))) == StandardDist{typeof(d)}(3) + @test @inferred(product_distribution(fill(d, 3, 4))) == StandardDist{typeof(d)}(3, 4) + end + + if length(size(d)) == 1 + @test @inferred(convert(Distributions.Product, d)) isa Distributions.Product + d_as_prod = convert(Distributions.Product, d) + @test d_as_prod.v == fill(StandardDist{D}(), size(d)...) + end + + @test @inferred(rand(stblrng(), d)) == rand(stblrng(), d) + @test @inferred(rand(stblrng(), d, 5)) == rand(stblrng(), d, 5) + + @test @inferred(rand(stblrng(), d)) == rand(stblrng(), dref) + @test @inferred(rand(stblrng(), d, 5)) == rand(stblrng(), dref, 5) + @test @inferred(rand!(stblrng(), d, zeros(size(d)...))) == rand!(stblrng(), dref, zeros(size(dref)...)) + if length(size(d)) == 1 + @test @inferred(rand!(stblrng(), d, zeros(size(d)..., 5))) == rand!(stblrng(), dref, zeros(size(dref)..., 5)) + end + end + end + + @testset "StandardDist{Normal}()" begin + # TODO: Add @inferred + d = StandardDist{Normal}(4) + d_uv = StandardDist{Normal}() + dref = MvNormal(Diagonal(fill(1.0, 4))) + @test (MvNormal(d)) == dref + @test (Base.convert(MvNormal, d)) == dref + end +end diff --git a/test/distributions/test_standard_normal.jl b/test/distributions/test_standard_normal.jl new file mode 100644 index 00000000..3d77f583 --- /dev/null +++ b/test/distributions/test_standard_normal.jl @@ -0,0 +1,129 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using Random, Statistics, LinearAlgebra +using Distributions, PDMats +using StableRNGs + + +@testset "StandardDist{Normal}" begin + stblrng() = StableRNG(789990641) + + @testset "StandardDist{Normal,0}" begin + @test @inferred(Normal(StandardDist{Normal}())) isa Normal{Float64} + @test @inferred(Normal(StandardDist{Normal}())) == Normal() + @test @inferred(convert(Normal, StandardDist{Normal}())) == Normal() + + d = StandardDist{Normal}() + dref = Normal() + + @test @inferred(minimum(d)) == minimum(dref) + @test @inferred(maximum(d)) == maximum(dref) + + @test @inferred(Distributions.params(d)) == () + @test @inferred(partype(d)) == partype(dref) + + @test @inferred(location(d)) == location(dref) + @test @inferred(scale(d)) == scale(dref) + + @test @inferred(eltype(typeof(d))) == eltype(typeof(dref)) + @test @inferred(eltype(d)) == eltype(dref) + + @test @inferred(length(d)) == length(dref) + @test @inferred(size(d)) == size(dref) + + @test @inferred(mean(d)) == mean(dref) + @test @inferred(median(d)) == median(dref) + @test @inferred(mode(d)) == mode(dref) + @test @inferred(modes(d)) ≈ modes(dref) + + @test @inferred(var(d)) == var(dref) + @test @inferred(std(d)) == std(dref) + @test @inferred(skewness(d)) == skewness(dref) + @test @inferred(kurtosis(d)) == kurtosis(dref) + + @test @inferred(entropy(d)) == entropy(dref) + + for x in [-Inf, -1.3, 0.0, 1.3, +Inf] + @test @inferred(gradlogpdf(d, x)) == gradlogpdf(dref, x) + + @test @inferred(logpdf(d, x)) == logpdf(dref, x) + @test @inferred(pdf(d, x)) == pdf(dref, x) + @test @inferred(logcdf(d, x)) == logcdf(dref, x) + @test @inferred(cdf(d, x)) == cdf(dref, x) + @test @inferred(logccdf(d, x)) == logccdf(dref, x) + @test @inferred(ccdf(d, x)) == ccdf(dref, x) + end + + for p in [0.0, 0.25, 0.75, 1.0] + @test @inferred(quantile(d, p)) == quantile(dref, p) + @test @inferred(cquantile(d, p)) == cquantile(dref, p) + end + + for t in [-3, 0, 3] + @test @inferred(mgf(d, t)) == mgf(dref, t) + @test @inferred(cf(d, t)) == cf(dref, t) + end + + @test @inferred(rand(stblrng(), d)) == rand(stblrng(), dref) + @test @inferred(rand!(stblrng(), d, fill(0.0))) == rand!(stblrng(), dref, fill(0.0)) + @test @inferred(rand(stblrng(), d, 5)) == rand(stblrng(), dref, 5) + + @test @inferred(truncated(StandardDist{Normal}(), -2.2f0, 3.1f0)) isa Truncated{Normal{Float64}} + @test truncated(StandardDist{Normal}(), -2.2f0, 3.1f0) == truncated(Normal(0.0, 1.0), -2.2f0, 3.1f0) + + @test @inferred(product_distribution(fill(StandardDist{Normal}(), 3))) isa StandardDist{Normal,1} + @test product_distribution(fill(StandardDist{Normal}(), 3)) == StandardDist{Normal}(3) + end + + + @testset "StandardDist{Normal,1}" begin + @test @inferred(StandardDist{Normal}(3)) isa StandardDist{Normal,1} + @test @inferred(StandardDist{Normal}(3)) isa StandardDist{Normal,1} + @test @inferred(StandardDist{Normal}(3)) isa StandardDist{Normal,1} + + @test @inferred(MvNormal(StandardDist{Normal}(3))) isa MvNormal{Int} + @test @inferred(MvNormal(StandardDist{Normal}(3))) == MvNormal(ScalMat(3, 1.0)) + @test @inferred(convert(MvNormal, StandardDist{Normal}(3))) == MvNormal(ScalMat(3, 1.0)) + + d = StandardDist{Normal}(3) + dref = MvNormal(ScalMat(3, 1.0)) + + @test @inferred(eltype(typeof(d))) == eltype(typeof(dref)) + @test @inferred(eltype(d)) == eltype(dref) + + @test @inferred(length(d)) == length(dref) + @test @inferred(size(d)) == size(dref) + + @test @inferred(Distributions.params(d)) == () + @test @inferred(partype(d)) == partype(dref) + + @test @inferred(mean(d)) == mean(dref) + @test @inferred(var(d)) == var(dref) + @test @inferred(cov(d)) == cov(dref) + + @test @inferred(mode(d)) == mode(dref) + @test @inferred(modes(d)) == modes(dref) + + @test @inferred(invcov(d)) == invcov(dref) + @test @inferred(logdetcov(d)) == logdetcov(dref) + + @test @inferred(entropy(d)) == entropy(dref) + + for x in fill.([-Inf, -1.3, 0.0, 1.3, +Inf], 3) + # Distributions.insupport is inconsistent at +- Inf between Normal and MvNormal + if !any(isinf, x) + @test @inferred(Distributions.insupport(d, x)) == Distributions.insupport(dref, x) + end + @test @inferred(logpdf(d, x)) == logpdf(dref, x) + @test @inferred(pdf(d, x)) == pdf(dref, x) + @test @inferred(sqmahal(d, x)) == sqmahal(dref, x) + @test @inferred(gradlogpdf(d, x)) == gradlogpdf(dref, x) + end + + @test @inferred(rand(stblrng(), d)) == rand(stblrng(), d) + @test @inferred(rand!(stblrng(), d, zeros(3))) == rand!(stblrng(), d, zeros(3)) + @test @inferred(rand!(stblrng(), d, zeros(3, 10))) == rand!(stblrng(), d, zeros(3, 10)) + end +end diff --git a/test/distributions/test_standard_uniform.jl b/test/distributions/test_standard_uniform.jl new file mode 100644 index 00000000..bcb0fb3e --- /dev/null +++ b/test/distributions/test_standard_uniform.jl @@ -0,0 +1,118 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using Random, Statistics, LinearAlgebra +using Distributions, PDMats +using StableRNGs +using FillArrays +using ForwardDiff + + +@testset "StandardDist{Uniform}" begin + stblrng() = StableRNG(789990641) + + @testset "StandardDist{Uniform,0}" begin + @test @inferred(Uniform(StandardDist{Uniform}())) isa Uniform{Float64} + @test @inferred(Uniform(StandardDist{Uniform}())) == Uniform() + @test @inferred(convert(Uniform, StandardDist{Uniform}())) == Uniform() + + d = StandardDist{Uniform}() + dref = Uniform() + + @test @inferred(minimum(d)) == minimum(dref) + @test @inferred(maximum(d)) == maximum(dref) + + @test @inferred(Distributions.params(d)) == () + @test @inferred(partype(d)) == partype(dref) + + @test @inferred(location(d)) == location(dref) + @test @inferred(scale(d)) == scale(dref) + + @test @inferred(eltype(typeof(d))) == eltype(typeof(dref)) + @test @inferred(eltype(d)) == eltype(dref) + + @test @inferred(length(d)) == length(dref) + @test @inferred(size(d)) == size(dref) + + @test @inferred(mean(d)) == mean(dref) + @test @inferred(median(d)) == median(dref) + @test @inferred(mode(d)) == mode(dref) + @test @inferred(modes(d)) ≈ modes(dref) + + @test @inferred(var(d)) ≈ var(dref) + @test @inferred(std(d)) ≈ std(dref) + @test @inferred(skewness(d)) == skewness(dref) + @test @inferred(kurtosis(d)) ≈ kurtosis(dref) + + @test @inferred(entropy(d)) == entropy(dref) + + for x in [-0.5, 0.0, 0.25, 0.75, 1.0, 1.5] + @test @inferred(logpdf(d, x)) == logpdf(dref, x) + @test @inferred(pdf(d, x)) == pdf(dref, x) + @test @inferred(logcdf(d, x)) == logcdf(dref, x) + @test @inferred(cdf(d, x)) == cdf(dref, x) + @test @inferred(logccdf(d, x)) == logccdf(dref, x) + @test @inferred(ccdf(d, x)) == ccdf(dref, x) + end + + for p in [0.0, 0.25, 0.75, 1.0] + @test @inferred(quantile(d, p)) == quantile(dref, p) + @test @inferred(cquantile(d, p)) == cquantile(dref, p) + end + + for t in [-3, 0, 3] + @test @inferred(mgf(d, t)) == mgf(dref, t) + @test @inferred(cf(d, t)) == cf(dref, t) + end + + @test @inferred(rand(stblrng(), d)) == rand(stblrng(), dref) + @test @inferred(rand!(stblrng(), d, fill(0.0))) == rand!(stblrng(), dref, fill(0.0)) + @test @inferred(rand(stblrng(), d, 5)) == rand(stblrng(), dref, 5) + + @test @inferred(truncated(StandardDist{Uniform}(), -0.5f0, 0.7f0)) isa Uniform{Float64} + @test truncated(StandardDist{Uniform}(), -0.5f0, 0.7f0) == Uniform(0.0f0, 0.7f0) + @test truncated(StandardDist{Uniform}(), 0.2f0, 0.7f0) == Uniform(0.2f0, 0.7f0) + + @test @inferred(product_distribution(fill(StandardDist{Uniform}(), 3))) isa MeasureBaseDistributionsExt.StandardDist{Uniform,1} + @test product_distribution(fill(StandardDist{Uniform}(), 3)) == MeasureBaseDistributionsExt.StandardDist{Uniform}(3) + end + + + @testset "StandardDist{Uniform,1}" begin + d = MeasureBaseDistributionsExt.StandardDist{Uniform}(3) + dref = product_distribution(fill(Uniform(), 3)) + + @test @inferred(eltype(typeof(d))) == eltype(typeof(dref)) + @test @inferred(eltype(d)) == eltype(dref) + + @test @inferred(length(d)) == length(dref) + @test @inferred(size(d)) == size(dref) + + @test @inferred(Distributions.params(d)) == () + @test @inferred(partype(d)) == partype(dref) + + @test @inferred(mean(d)) == mean(dref) + @test @inferred(var(d)) ≈ var(dref) + @test @inferred(cov(d)) ≈ cov(dref) + + @test @inferred(mode(d)) == [0.5, 0.5, 0.5] + @test @inferred(modes(d)) == fill([0, 0,0 ]) + + @test @inferred(invcov(d)) == inv(cov(dref)) + @test @inferred(logdetcov(d)) == logdet(cov(dref)) + + @test @inferred(entropy(d)) == entropy(dref) + + for x in fill.([-Inf, -1.3, 0.0, 1.3, +Inf], 3) + @test @inferred(Distributions.insupport(d, x)) == Distributions.insupport(dref, x) + @test @inferred(logpdf(d, x)) == logpdf(dref, x) + @test @inferred(pdf(d, x)) == pdf(dref, x) + @test @inferred(gradlogpdf(d, x)) == ForwardDiff.gradient(x -> logpdf(d, x), x) + end + + @test @inferred(rand(stblrng(), d)) == rand(stblrng(), d) + @test @inferred(rand!(stblrng(), d, zeros(3))) == rand!(stblrng(), d, zeros(3)) + @test @inferred(rand!(stblrng(), d, zeros(3, 10))) == rand!(stblrng(), d, zeros(3, 10)) + end +end diff --git a/test/distributions/test_transport.jl b/test/distributions/test_transport.jl new file mode 100644 index 00000000..4ae50831 --- /dev/null +++ b/test/distributions/test_transport.jl @@ -0,0 +1,194 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using LinearAlgebra +using InverseFunctions, ChangesOfVariables +using Distributions, ArraysOfArrays +using StableRNGs +import ForwardDiff, Zygote +import PDMats + +using MeasureBase: transport_to, transport_def, transport_origin +using MeasureBase: StdUniform, StdNormal, StdExponential +using .MeasureBaseDistributionsExt: _trafo_cdf, _trafo_quantile + +include("getjacobian.jl") + + +@testset "test_distribution_transform" begin + function test_back_and_forth(trg, src) + @testset "transform $(typeof(trg).name) <-> $(typeof(src).name)" begin + x = rand(src) + y = transport_def(trg, src, x) + src_v_reco = transport_def(src, trg, y) + + @test x ≈ src_v_reco + + f = x -> transport_def(trg, src, x) + ref_ladj = logpdf(src, x) - logpdf(trg, y) + @test ref_ladj ≈ logabsdet(getjacobian(f, x))[1] + end + end + + reshaped_rand(d::Distribution{Univariate}, n) = rand(d, n) + reshaped_rand(d::Distribution{Multivariate}, n) = nestedview(rand(d, n)) + + function test_dist_trafo_moments(trg, src) + unshaped(x) = first(torv_and_back(x)) + @testset "check moments of trafo $(typeof(trg).name) <- $(typeof(src).name)" begin + X = reshaped_rand(src, 10^5) + Y = transport_to(trg, src).(X) + Y_ref = reshaped_rand(trg, 10^6) + @test isapprox(mean(unshaped.(Y)), mean(unshaped.(Y_ref)), rtol = 0.5) + @test isapprox(cov(unshaped.(Y)), cov(unshaped.(Y_ref)), rtol = 0.5) + end + end + + @testset "transforms-tests" begin + stduvuni = StandardDist{Uniform}() + stduvnorm = StandardDist{Uniform}() + + uniform1 = Uniform(-5.0, -0.01) + uniform2 = Uniform(0.01, 5.0) + + normal1 = Normal(-10, 1) + normal2 = Normal(10, 5) + + stdmvnorm1 = StandardDist{Normal}(1) + stdmvnorm2 = StandardDist{Normal}(2) + + stdmvuni2 = StandardDist{Uniform}(2) + + standnorm2_reshaped = reshape(stdmvnorm2, 1, 2) + + mvnorm = MvNormal([0.3, -2.9], [1.7 0.5; 0.5 2.3]) + beta = Beta(3,1) + gamma = Gamma(0.1,0.7) + dirich = Dirichlet([0.1,4]) + + test_back_and_forth(stduvuni, stduvuni) + test_back_and_forth(stduvnorm, stduvnorm) + test_back_and_forth(stduvuni, stduvnorm) + test_back_and_forth(stduvnorm, stduvuni) + + test_back_and_forth(stdmvuni2, stdmvuni2) + test_back_and_forth(stdmvnorm2, stdmvnorm2) + test_back_and_forth(stdmvuni2, stdmvnorm2) + test_back_and_forth(stdmvnorm2, stdmvuni2) + + test_back_and_forth(beta, stduvnorm) + test_back_and_forth(gamma, stduvnorm) + test_back_and_forth(gamma, beta) + + test_back_and_forth(mvnorm, stdmvuni2) + test_back_and_forth(stdmvuni2, mvnorm) + + test_back_and_forth(mvnorm, standnorm2_reshaped) + test_back_and_forth(standnorm2_reshaped, mvnorm) + test_back_and_forth(stdmvnorm2, standnorm2_reshaped) + test_back_and_forth(standnorm2_reshaped, standnorm2_reshaped) + + test_dist_trafo_moments(normal2, normal1) + test_dist_trafo_moments(uniform2, uniform1) + + test_dist_trafo_moments(beta, stduvnorm) + test_dist_trafo_moments(gamma, stduvnorm) + + test_dist_trafo_moments(mvnorm, stdmvnorm2) + test_dist_trafo_moments(dirich, stdmvnorm1) + + let + mvuni = product_distribution([Uniform(), Uniform()]) + + x = rand() + @test_throws ArgumentError transport_to(stduvnorm, mvnorm)(x) + @test_throws ArgumentError transport_to(stduvnorm, stdmvnorm1)(x) + @test_throws ArgumentError transport_to(stduvnorm, stdmvnorm2)(x) + + x = rand(2) + @test_throws ArgumentError transport_to(stduvnorm, mvnorm)(x) + @test_throws ArgumentError transport_to(stduvnorm, stdmvnorm1)(x) + @test_throws ArgumentError transport_to(stduvnorm, stdmvnorm2)(x) + end + end + + @testset "Custom cdf and quantile for dual numbers" begin + Dual = ForwardDiff.Dual + + @test isapprox(_trafo_cdf(Normal(Dual(0, 1, 0, 0), Dual(1, 0, 1, 0)), Dual(0.5, 0, 0, 1)), cdf(Normal(Dual(0, 1, 0, 0), Dual(1, 0, 1, 0)), Dual(0.5, 0, 0, 1)), rtol = 10^-6) + @test isapprox(_trafo_cdf(Normal(0, 1), Dual(0.5, 1)), cdf(Normal(0, 1), Dual(0.5, 1)), rtol = 10^-6) + + @test isapprox(_trafo_quantile(Normal(0, 1), Dual(0.5, 1)), quantile(Normal(0, 1), Dual(0.5, 1)), rtol = 10^-6) + @test isapprox(_trafo_quantile(Normal(Dual(0, 1, 0, 0), Dual(1, 0, 1, 0)), Dual(0.5, 0, 0, 1)), quantile(Normal(Dual(0, 1, 0, 0), Dual(1, 0, 1, 0)), Dual(0.5, 0, 0, 1)), rtol = 10^-6) + end + + @testset "trafo autodiff pullbacks" begin + x = [0.6, 0.7, 0.8, 0.9] + f = transport_to(Dirichlet([3.0, 4.0, 5.0, 6.0, 7.0]), Uniform) + @test isapprox(ForwardDiff.jacobian(f, x), Zygote.jacobian(f, x)[1], rtol = 10^-4) + f = inverse(transport_to(Normal, Dirichlet([3.0, 4.0, 5.0, 6.0, 7.0]))) + @test isapprox(ForwardDiff.jacobian(f, x), Zygote.jacobian(f, x)[1], rtol = 10^-4) + end + + + @testset "transport_to autosel" begin + for (M,R) in [ + (StandardDist{Normal}, StandardDist{Normal}) + (Normal, StandardDist{Normal}) + (StandardDist{Uniform}, StandardDist{Uniform}) + (Uniform, StandardDist{Uniform}) + ] + @test @inferred(transport_to(M, Weibull())) == transport_to(R(), Weibull()) + @test @inferred(transport_to(Weibull(), M)) == transport_to(Weibull(), R()) + @test @inferred(transport_to(M, MvNormal(float(I(5))))) == transport_to(R(5), MvNormal(float(I(5)))) + @test @inferred(transport_to(MvNormal(float(I(5))), M)) == transport_to(MvNormal(float(I(5))), R(5)) + @test @inferred(transport_to(M, StdExponential()^(2,3))) == transport_to(R(6), StdExponential()^(2,3)) + @test @inferred(transport_to(StdExponential()^(2,3), M)) == transport_to(StdExponential()^(2,3), R(6)) + end + end + + @testset "affine transformed distributions" begin + d = 2.0 * Weibull(0.7) + 1.0 + x = rand(StableRNG(789990641), d) + u = transport_to(StdUniform(), d)(x) + @test u ≈ cdf(d, x) + @test transport_to(d, StdUniform())(u) ≈ x + test_back_and_forth(StandardDist{Normal}(), d) + end + + @testset "truncated distributions" begin + d = truncated(Normal(0.3, 1.2), -0.5, 1.5) + for u in [0.0, 0.25, 0.75, 1.0, prevfloat(1.0)] + x = transport_to(d, StdUniform())(u) + @test minimum(d) <= x <= maximum(d) + end + test_back_and_forth(StandardDist{Uniform}(), d) + end + + @testset "products of distributions" begin + pd = product_distribution([Weibull(0.7), Exponential(1.3), Normal(0.5, 2.0)]) + m = MeasureBase.asmeasure(pd) + x = rand(StableRNG(789990641), pd) + for trg in [StdUniform()^3, StdNormal()^3] + y = transport_to(trg, m)(x) + y_ref = map((d_i, x_i) -> transport_to(trg.parent, d_i)(x_i), pd.v, x) + @test y ≈ y_ref + @test transport_to(m, trg)(y) ≈ x + end + + pd2 = product_distribution([Normal(2.0, 0.5), Weibull(1.2), Uniform(-1.0, 3.0)]) + m2 = MeasureBase.asmeasure(pd2) + y = transport_to(m2, m)(x) + @test transport_to(m, m2)(y) ≈ x + end + + @testset "MvNormal covariance representations" begin + for Σ in [PDMats.ScalMat(3, 2.5), PDMats.PDiagMat([0.5, 1.0, 2.5]), Diagonal([0.5, 1.0, 2.5])] + mvn = MvNormal([0.2, -0.4, 0.6], Σ) + x = rand(StableRNG(789990641), mvn) + y = transport_to(StandardDist{Normal}(3), mvn)(x) + @test transport_to(mvn, StandardDist{Normal}(3))(y) ≈ x + end + end +end diff --git a/test/runtests.jl b/test/runtests.jl index c2f63c4e..f0d4488a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -9,7 +9,6 @@ using MeasureBase: test_interface, test_smf include("test_aqua.jl") -include("static.jl") include("test_primitive.jl") include("test_standard.jl") @@ -19,8 +18,13 @@ include("getdof.jl") include("transport.jl") include("smf.jl") +include("test_mooncake.jl") + include("combinators/weighted.jl") include("combinators/transformedmeasure.jl") +include("combinators/reshape.jl") include("combinators/implicitlymapped.jl") +include("distributions/test_distributions.jl") + include("test_docs.jl") diff --git a/test/static.jl b/test/static.jl deleted file mode 100644 index f618124b..00000000 --- a/test/static.jl +++ /dev/null @@ -1,34 +0,0 @@ -using Test - -import MeasureBase - -import Static -using Static: static -import FillArrays - -@testset "static" begin - @test 2 isa MeasureBase.IntegerLike - @test static(2) isa MeasureBase.IntegerLike - @test true isa MeasureBase.IntegerLike - @test static(true) isa MeasureBase.IntegerLike - - @test @inferred(MeasureBase.one_to(7)) isa Base.OneTo - @test @inferred(MeasureBase.one_to(7)) == 1:7 - @test @inferred(MeasureBase.one_to(static(7))) isa Static.SOneTo - @test @inferred(MeasureBase.one_to(static(7))) == static(1):static(7) - - @test @inferred(MeasureBase.fill_with(4.2, (7,))) == FillArrays.Fill(4.2, 7) - @test @inferred(MeasureBase.fill_with(4.2, (static(7),))) == FillArrays.Fill(4.2, 7) - @test @inferred(MeasureBase.fill_with(4.2, (3, static(7)))) == - FillArrays.Fill(4.2, 3, 7) - @test @inferred(MeasureBase.fill_with(4.2, (3:7,))) == FillArrays.Fill(4.2, (3:7,)) - @test @inferred(MeasureBase.fill_with(4.2, (static(3):static(7),))) == - FillArrays.Fill(4.2, (3:7,)) - @test @inferred(MeasureBase.fill_with(4.2, (3:7, static(2):static(5)))) == - FillArrays.Fill(4.2, (3:7, 2:5)) - - @test MeasureBase.maybestatic_length(MeasureBase.one_to(7)) isa Int - @test MeasureBase.maybestatic_length(MeasureBase.one_to(7)) == 7 - @test MeasureBase.maybestatic_length(MeasureBase.one_to(static(7))) isa Static.StaticInt - @test MeasureBase.maybestatic_length(MeasureBase.one_to(static(7))) == static(7) -end diff --git a/test/test_aqua.jl b/test/test_aqua.jl index b6290e31..d4546ac2 100644 --- a/test/test_aqua.jl +++ b/test/test_aqua.jl @@ -9,5 +9,10 @@ import MeasureBase #end # testset Test.@testset "Aqua tests" begin - Aqua.test_all(MeasureBase, ambiguities = false) + Aqua.test_all( + MeasureBase, + ambiguities = false, + # Only used by package extensions: + stale_deps = (ignore = [:ArgCheck, :ArraysOfArrays],), + ) end # testset diff --git a/test/test_basics.jl b/test/test_basics.jl index 7ac29dc1..bd5a409c 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -120,8 +120,9 @@ end end @testset "powers" begin - @test logdensityof(Lebesgue()^3, 2) == logdensityof(Lebesgue()^(3,), 2) - @test logdensityof(Lebesgue()^3, 2) == logdensityof(Lebesgue()^(3, 1), (2, 0)) + @test logdensityof(Lebesgue()^3, [2, 2, 2]) == logdensityof(Lebesgue()^(3,), fill(2, 3)) + @test logdensityof(Lebesgue()^3, fill(2, 3)) == + logdensityof(Lebesgue()^(3, 1), fill(2, 3, 1)) end NormalMeasure() = ∫exp(x -> -0.5x^2, Lebesgue(ℝ)) diff --git a/test/test_mooncake.jl b/test/test_mooncake.jl new file mode 100644 index 00000000..95ce84b5 --- /dev/null +++ b/test/test_mooncake.jl @@ -0,0 +1,56 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using Random +import Mooncake +import ForwardDiff + +using MeasureBase +using MeasureBase: transport_to +using MeasureBase: isneginf, isposinf, _adignore_call +using MeasureBase: check_dof, require_insupport, _origin_depth +using MeasureBase: logdensityof_rt + +_mooncake_gradient(f, x) = Mooncake.value_and_gradient!!( + Mooncake.prepare_gradient_cache(f, x), f, x +)[2][2] + +@testset "Mooncake AD rules" begin + @test Base.get_extension(MeasureBase, :MeasureBaseMooncakeExt) isa Module + + @testset "zero-derivative primitives" begin + rng = Random.Xoshiro(789990641) + Mooncake.TestUtils.test_rule(rng, isneginf, 0.5; is_primitive = true) + Mooncake.TestUtils.test_rule(rng, isposinf, 0.5; is_primitive = true) + Mooncake.TestUtils.test_rule(rng, _adignore_call, () -> 42.0; is_primitive = true) + Mooncake.TestUtils.test_rule(rng, check_dof, StdNormal(), StdUniform(); is_primitive = true) + Mooncake.TestUtils.test_rule(rng, require_insupport, StdNormal(), 0.5; is_primitive = true) + Mooncake.TestUtils.test_rule(rng, _origin_depth, StdNormal(); is_primitive = true) + Mooncake.TestUtils.test_rule(rng, logdensityof_rt, StdNormal(), 0.5; is_primitive = true) + end + + @testset "@_adignore is ignored" begin + f_adignore(x) = (MeasureBase.@_adignore x^3; x^2) + @test _mooncake_gradient(f_adignore, 3.0) ≈ 6.0 + end + + @testset "logdensityof gradients" begin + x = [0.1, -0.2, 0.3] + f_ld = x -> logdensityof(StdNormal()^3, x) + @test _mooncake_gradient(f_ld, x) ≈ ForwardDiff.gradient(f_ld, x) + + f_ldu = x -> logdensityof(StdExponential()^3, x) + @test _mooncake_gradient(f_ldu, abs.(x)) ≈ ForwardDiff.gradient(f_ldu, abs.(x)) + end + + @testset "transport gradients" begin + x = [0.1, -0.2, 0.3] + f_t = x -> sum(transport_to(StdUniform()^3, StdNormal()^3)(x)) + @test _mooncake_gradient(f_t, x) ≈ ForwardDiff.gradient(f_t, x) + + u = [0.3, 0.5, 0.7] + f_ti = u -> sum(transport_to(StdNormal()^3, StdUniform()^3)(u)) + @test _mooncake_gradient(f_ti, u) ≈ ForwardDiff.gradient(f_ti, u) + end +end