diff --git a/Project.toml b/Project.toml index 44f89f80..8700a16a 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" @@ -20,6 +23,7 @@ LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" LogarithmicNumbers = "aa2f6b4e-9042-5d33-9679-40d3a6b85899" MappedArrays = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900" NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3" +OneTwoMany = "762dc654-8631-413a-a342-372a7419ad9d" PrettyPrinting = "54e16d92-306c-5ea0-a30b-337be88ac337" PropertyFunctions = "09e99361-2bb8-48a2-a80f-de58f0739eb4" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" @@ -27,25 +31,49 @@ 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" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" +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" +MeasureBaseReactantExt = "Reactant" [compat] +ArgCheck = "1, 2" +ArraysOfArrays = "0.6, 0.7" 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 +82,22 @@ LinearAlgebra = "1" LogExpFunctions = "0.3, 1" LogarithmicNumbers = "1" MappedArrays = "0.4" +Mooncake = "0.5.34" NaNMath = "0.3, 1" +OneTwoMany = "0.1.2" +PDMats = "0.11" PrettyPrinting = "0.3, 0.4" PropertyFunctions = "0.2.2" Random = "1" +Reactant = "0.2" 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..a230386a --- /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, _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..933a3f6f --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/measure_interface.jl @@ -0,0 +1,27 @@ +# 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.mintegrate(f, base::Distribution) = + MeasureBase.mintegrate(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..f607eac3 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/univariate.jl @@ -0,0 +1,137 @@ +# 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::Number) = + _trafo_cdf_impl(_dist_params_numtype(d), d, x) + +@inline _trafo_cdf_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, x::Number) = + Distributions.cdf(d, x) + + +@inline _trafo_quantile(d::Distribution{Univariate,Continuous}, u::Number) = + _trafo_quantile_impl(_dist_params_numtype(d), d, u) + +@inline _trafo_quantile_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, u::Number) = + _trafo_quantile_impl_generic(d, u) + + +@inline _trafo_quantile_impl_generic(d::Distribution{Univariate,Continuous}, u::Number) = + 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<:Number} + float(promote_type(T, _dist_params_numtype(d))) +end + + +@inline function MeasureBase.transport_def(::StdUniform, μ::Distribution{Univariate,Continuous}, x) + R = _result_numtype(μ, x) + y = _trafo_cdf(μ, x) + ifelse(Distributions.insupport(μ, x), convert(R, y), convert(R, NaN)) +end + + +@inline function MeasureBase.transport_def(ν::Distribution{Univariate,Continuous}, ::StdUniform, x::T) where {T} + R = _result_numtype(ν, x) + TF = float(T) + # Avoid x ≈ 0 and x ≈ 1 to avoid infinite variate values for target + # distributions with infinite support, keep the quantile argument valid + # for out-of-range x (the result is masked to NaN then): + clamped_x = clamp(convert(TF, x), zero(TF), one(TF)) + mod_x = ifelse(x ≈ 0, zero(TF) + eps(TF), ifelse(x ≈ 1, one(TF) - eps(TF), clamped_x)) + y = _trafo_quantile(ν, mod_x) + ifelse((zero(x) <= x) & (x <= one(x)), convert(R, y), convert(R, NaN)) +end + + +# Use standard measures as transformation origin for scaled/translated equivalents: + +function _origin_to_affine(ν::Distribution{Univariate}, y::T) where {T<:Number} + 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<:Number} + 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..0d91f8b8 --- /dev/null +++ b/ext/MeasureBaseForwardDiffExt.jl @@ -0,0 +1,16 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +module MeasureBaseForwardDiffExt + +using MeasureBase +using MeasureBase: containsnan +import ForwardDiff + +function MeasureBase.containsnan(x::ForwardDiff.Dual) + a = containsnan(x.value) + b = containsnan(x.partials) + return a || b +end + + +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/ext/MeasureBaseReactantExt.jl b/ext/MeasureBaseReactantExt.jl new file mode 100644 index 00000000..53bc6b4b --- /dev/null +++ b/ext/MeasureBaseReactantExt.jl @@ -0,0 +1,12 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +module MeasureBaseReactantExt + +using Reactant: TracedRNumber +import MeasureBase +using MeasureBase: RealValues, IntegerValues + +Base.in(::TracedRNumber{<:Real}, ::RealValues) = true +Base.in(::TracedRNumber{<:Integer}, ::IntegerValues) = true + +end # module MeasureBaseReactantExt diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index 2bad7d92..4b3adc4e 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,8 +44,29 @@ 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 + +using ArraysOfArrays: + ArrayOfSimilarArrays, VectorOfArrays, VectorOfVectors, VectorOfSimilarArrays, + VectorOfSimilarVectors, flatview + +using OneTwoMany: firstarg, secondarg + export gentype -export rebase export AbstractMeasure @@ -64,6 +86,42 @@ 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 + +Base.:(==)(a::AsMeasure, b::AsMeasure) = a.obj == b.obj +Base.isapprox(a::AsMeasure, b::AsMeasure; kwargs...) = isapprox(a.obj, b.obj; kwargs...) + function Pretty.quoteof(d::M) where {M<:AbstractMeasure} the_names = fieldnames(typeof(d)) :($M($([getfield(d, n) for n in the_names]...))) @@ -105,6 +163,14 @@ Compute the log-density of the measure m at the point `x`, relative to Compute the log-density of `m1` relative to `m2` at the point `x`, assuming `insupport(m1, x)` and `insupport(m2, x)`. + +The generic implementation descends the base measure chains of both +measures in lockstep, so it terminates at the first pair of base measures +for which a specialized relative density is available (in particular at +pairs of identical primitive measures) and any shared chain suffix cancels +symbolically. To provide specialized relative densities for pairs of +measure types, add methods to [`MeasureBase.logdensity_rel_def`](@ref), +not to `logdensity_def` itself. """ function logdensity_def end @@ -113,37 +179,37 @@ using Compat using IrrationalConstants using IrrationalConstants: loghalf -include("static.jl") +include("collection_utils.jl") include("smf.jl") +include("mspace.jl") include("getdof.jl") include("transport.jl") -include("schema.jl") -include("splat.jl") include("proxies.jl") -include("kernel.jl") include("parameterized.jl") include("domains.jl") include("primitive.jl") include("utils.jl") include("mass-interface.jl") +include("density.jl") +include("density-core.jl") + include("primitives/counting.jl") include("primitives/lebesgue.jl") include("primitives/dirac.jl") 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") include("combinators/power.jl") +include("density-batched.jl") include("combinators/spikemixture.jl") include("combinators/likelihood.jl") -include("combinators/pointwise.jl") include("combinators/restricted.jl") include("combinators/smart-constructors.jl") -include("combinators/powerweighted.jl") include("combinators/conditional.jl") include("combinators/implicitlymapped.jl") @@ -152,14 +218,16 @@ include("standard/stduniform.jl") include("standard/stdexponential.jl") include("standard/stdlogistic.jl") include("standard/stdnormal.jl") +include("combinators/product_transport.jl") +include("combinators/combined.jl") +include("combinators/bind.jl") include("combinators/half.jl") #include("implicitmaps.jl") include("rand.jl") -include("density.jl") -include("density-core.jl") +include("measure_operators.jl") include("interface.jl") diff --git a/src/collection_utils.jl b/src/collection_utils.jl new file mode 100644 index 00000000..5b4d915a --- /dev/null +++ b/src/collection_utils.jl @@ -0,0 +1,124 @@ +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))) + +Base.@propagate_inbounds _as_tuple(v::AbstractVector, ::Val{N}) where {N} = Tuple(SVector{N}(v)) + + +Base.@propagate_inbounds function _get_or_view(A::AbstractVector, from::IntegerLike, until::IntegerLike) + view(A, dynamic(from):dynamic(until)) +end + +Base.@propagate_inbounds function _get_or_view( + A::AbstractVector, + ::StaticInteger{from}, + ::StaticInteger{until}, +) where {from,until} + SVector{until - from + 1}(view(A, from:until)) +end + +# ToDo: Specialize for StaticVector instead of SVector? +Base.@propagate_inbounds function _get_or_view( + A::SVector, + from::StaticInteger, + until::StaticInteger, +) + # ToDo: Improve implementation: + SVector(_get_or_view(Tuple(A), from, until)) +end + +Base.@propagate_inbounds function _get_or_view(tpl::Tuple, from::IntegerLike, until::IntegerLike) + ntuple(i -> tpl[from + i - 1], Val(until - from + 1)) +end + + +@inline function _split_after(x::AbstractVector, n::IntegerLike) + idxs = maybestatic_eachindex(x) + i_first = maybestatic_first(idxs) + i_last = maybestatic_last(idxs) + _get_or_view(x, i_first, i_first + n - one(n)), _get_or_view(x, i_first + n, i_last) +end + +@inline _split_after(x::Tuple, n) = _split_after(x::Tuple, Val{n}()) +@inline _split_after(x::Tuple, ::Val{N}) where {N} = x[begin:(begin+N-1)], x[(begin+N):end] + +@generated function _split_after(x::NamedTuple{names}, ::Val{names_a}) where {names,names_a} + n = length(names_a) + if names[begin:(begin+n-1)] == names_a + names_b = names[(begin+n):end] + quote + a, b = _split_after(values(x), Val($n)) + NamedTuple{$names_a}(a), NamedTuple{$names_b}(b) + end + else + quote + throw(ArgumentError("Can't split NamedTuple{$names} after {$names_a}")) + end + end +end + + +# Field access functions for Fill: +_fill_value(x::FillArrays.Fill) = x.value +_fill_axes(x::FillArrays.Fill) = x.axes + + +_flatten_to_rv(VV::AbstractVector{<:AbstractVector{<:Number}}) = flatview(VectorOfArrays(VV)) +_flatten_to_rv(VV::AbstractVector{<:StaticVector{N,<:Number}}) where {N} = + flatview(VectorOfSimilarArrays(VV)) + +_flatten_to_rv(VV::VectorOfSimilarVectors{<:Number}) = flatview(VV) +_flatten_to_rv(VV::VectorOfVectors{<:Number}) = flatview(VV) + +_flatten_to_rv(::Tuple{}) = [] +_flatten_to_rv(tpl::Tuple{Vararg{AbstractVector}}) = vcat(tpl...) +_flatten_to_rv(tpl::Tuple{Vararg{StaticVector}}) = vcat(tpl...) + + +# Non-mutating concatenation of measure collections: +_cat_measures(a::Tuple, b::Tuple) = (a..., b...) +_cat_measures(a::AbstractVector, b::Tuple) = vcat(a, [b...]) +_cat_measures(a::Tuple, b::AbstractVector) = vcat([a...], b) +_cat_measures(a::AbstractVector, b::AbstractVector) = vcat(a, b) + + +# Take the beginning of a flat vector stream as a variate of size `sz`, +# scalar variates have size `()` and multi-rank variates are reshaped: +Base.@propagate_inbounds _consume_from_stream(x::AbstractVector, sz::Tuple{IntegerLike}) = + _split_after(x, sz[1]) + +Base.@propagate_inbounds function _consume_from_stream(x::AbstractVector, ::Tuple{}) + idxs = maybestatic_eachindex(x) + i_first = maybestatic_first(idxs) + x[i_first], _get_or_view(x, i_first + one(i_first), maybestatic_last(idxs)) +end + +function _consume_from_stream(x::AbstractVector, sz::Tuple{Vararg{IntegerLike}}) + a_flat, x_rest = _split_after(x, size2length(sz)) + return maybestatic_reshape(a_flat, sz), x_rest +end + +function _consume_from_stream(x::AbstractVector, @nospecialize(sz)) + throw(ArgumentError("Can't consume a variate of size $sz from a flat vector stream")) +end diff --git a/src/combinators/bind.jl b/src/combinators/bind.jl index cc2022f2..ab7ab708 100644 --- a/src/combinators/bind.jl +++ b/src/combinators/bind.jl @@ -1,36 +1,362 @@ -struct Bind{M,K} <: AbstractMeasure - μ::M - k::K -end +@doc raw""" + mkernel(f_β, f_c = OneTwoMany.secondarg)::Function + +Constructs a generalized monadic transition kernel from a primary transition +kernel function `f_β` and a value combination function `f_c`. + +`f_β` must behave like `β = f_β(a)`, taking a value `a` from a primary +measurable space and returning a measure-like object `β`. + +`f_c` must behave like `c = f_c(a, b)`, taking a value `a` (like `f_β`) and +a value `b` from the measurable space of `β` and returning a value `c`. + +`f_k = mkernel(f_β, f_c)` then acts like + +```julia +f_k(a) ≡ pushfwd(c -> f_c(c[1], c[2]), productmeasure((Dirac(a), f_β(a)))) +``` + +(`≡` denoting pseudocode-equivalency here). So with the default +`f_c == OneTwoMany.secondarg`, we just have `f_k(a) ≡ f_β(a)`. + +Also, + +```julia +mbind(mkernel(f_β, f_c), α) == mbind(f_β, α, f_c) +``` + +See also [`mbind`](@ref). +""" +function mkernel end +export mkernel + + +""" + struct MeasureBase.MKernel <: Function -export ↣ +Represents a generalized monadic transition kernel. +User code should not create instances of `MKernel` directly, but should +call [`mkernel`](@ref) instead. """ -If -- μ is an `AbstractMeasure` or satisfies the Measure interface, and -- k is a function taking values from the support of μ and returning a measure +struct MKernel{FT,FC} <: Function + f_β::FT + f_c::FC +end + +(f_k::MKernel)(a) = mbind(f_k, Dirac(a)) + +@inline mkernel(f_β::MKernel) = f_β +@inline mkernel(f_β, f_c = secondarg) = _generic_mkernel_impl(f_β, f_c) + +@inline _generic_mkernel_impl(f_β, f_c) = MKernel(f_β, f_c) +@inline _generic_mkernel_impl(f_β::MKernel, ::typeof(secondarg)) = f_β + + +@doc raw""" + mbind(f_β, α::AbstractMeasure, f_c = OneTwoMany.secondarg) + mbind(f_β::MeasureBase.MKernel, α::AbstractMeasure) + +Constructs a monadic bind, resp. a hierarchical measure, from a transition +kernel function `f_β`, a primary measure `α` and a value combination +function `f_c`. + +`f_β` must be a function that maps a point `a` from the space of the primary +measure `α` to a dependent secondary measure `β_a = f_β(a)`. +`ab = f_c(a, b)` must map such a point `a` and a point `b` from the +space of measure `β_a` to a combined value `ab = f_c(a, b)`. + +The resulting measure + +```julia +μ = mbind(f_β, α, f_c) +``` + +has the mathematical interpretation (on sets $$A$$ and $$B$$) + +```math +\mu(f_c(A, B)) = \int_A \beta_a(B)\, \mathrm{d}\, \alpha(a) +``` + +When using the default `f_c = OneTwoMany.secondarg` (so `ab == b`) this +simplifies to -Then `μ ↣ k` is a measure, called a *monadic bind*. In a -probabilistic programming language like Soss.jl, this could be expressed as +```math +\mu(B) = \int_A \beta_a(B)\, \mathrm{d}\, \alpha(a) +``` + +which is equivalent to a monadic bind, viewing measures as monads. -Note that bind is usually written `>>=`, but this symbol is unavailable in Julia. +Computationally, `ab = rand(μ)` is equivalent to +```julia +a = rand(α) +β_a = f_β(a) +b = rand(β_a) +ab = f_c(a, b) ``` -bind = @model μ,k begin - x ~ μ - y ~ k(x) - return y + +The measure `α` that went into the bind can be retrieved via +`boundmeasure(mbind(f_β, α, f_c)) == α` and the kernel via +`bindkernel(mbind(f_β, α, f_c)) == mkernel(f_β, f_c)`. + +Densities on hierarchical measures can only be evaluated if `ab = f_c(a, b)` +can be unambiguously split into `a` and `b` again, knowing `α`. This is +currently implemented for `f_c` that is either `tuple` or `=>`/`Pair` (these +work for any combination of variate types), `vcat` (for tuple- or +vector-like variates) and `merge` (`NamedTuple` variates). +[`MeasureBase.tpmeasure_split_combined`](@ref) can be specialized to +support other choices for `f_c`. + +# Extended help + +Bayesian example with a correlated prior: Mathematically, let + + position = a1 ~ StdNormal() + noise = a2 ~ pushforward(h(a1, ·), StdExponential()) + +where `h(a1, a2) = √(abs(a1) * a2)`. Because this prior on the space of +`A = A1 × A2 = (position, noise)` is a hierarchical measure (`a2` depends +on `a1`), we can construct it using `mbind` with `merge` as `f_c`: + +```julia +using MeasureBase, AffineMaps + +prior = mbind( + productmeasure(( + position = StdNormal(), + )), merge +) do a + productmeasure(( + noise = pushfwd(setinverse(sqrt, setladj(x -> x^2, x -> log(2))) ∘ Mul(abs(a.position)), StdExponential()), + )) end + +model = θ -> pushfwd(MulAdd(θ.noise, θ.position), StdNormal())^10 + +joint_θ_obs = mbind(model, prior, tuple) +prior_predictive = mbind(model, prior) + +observation = rand(prior_predictive) +likelihood = likelihoodof(model, observation) + +posterior = mintegrate(likelihood, prior) + +θ = rand(prior) +logdensityof(posterior, θ) ``` +""" +function mbind end +export mbind -See also `bind` and `Bind` +@inline mbind(f_β) = Base.Fix1(mbind, f_β) + +@inline function mbind(f_β, α::AbstractMeasure, f_c = secondarg) + _generic_mbind_impl(f_β, asmeasure(α), f_c) +end + +@inline function _generic_mbind_impl(f_β, α::AbstractMeasure, f_c) + F, M, G = Core.Typeof(f_β), Core.Typeof(α), Core.Typeof(f_c) + Bind{F,M,G}(f_β, α, f_c) +end + +@inline _generic_mbind_impl(f_β, α::Dirac, f_c) = mcombine(f_c, α, asmeasure(f_β(α.x))) + +@inline _generic_mbind_impl(@nospecialize(f_β), α::AbstractMeasure, ::typeof(firstarg)) = α +@inline _generic_mbind_impl(@nospecialize(f_β), α::Dirac, ::typeof(firstarg)) = α + +@inline _generic_mbind_impl(f_k::MKernel, α::AbstractMeasure, ::typeof(secondarg)) = + mbind(f_k.f_β, α, f_k.f_c) +@inline _generic_mbind_impl(f_k::MKernel, α::Dirac, ::typeof(secondarg)) = + mbind(f_k.f_β, α, f_k.f_c) + + +""" + struct MeasureBase.Bind <: AbstractMeasure + +Represents a monadic bind resp. a hierarchical measure in general. + +User code should not create instances of `Bind` directly, but should call +[`mbind`](@ref) instead. """ -↣(μ, k) = bind(μ, k) +struct Bind{FT,M<:AbstractMeasure,FC} <: AbstractMeasure + f_β::FT + α::M + f_c::FC +end + +# ToDo: Store MKernel in Bind instead of separate fields f_β and f_c? + + +""" + bindkernel(μ::Bind)::MKernel + +Returns the monadic transition kernel of a monadic bind, so that +`bindkernel(mbind(f_k::MKernel, α)) == f_k`. + +See [`mbind`](@ref) and [`mkernel`](@ref) for details. +""" +function bindkernel end +export bindkernel + +bindkernel(μ::Bind) = mkernel(μ.f_β, μ.f_c) + + +""" + boundmeasure(μ::Bind)::AbstractMeasure + +Returns the measure that went into a monadic bind, so that +`boundmeasure(mbind(f_k, α)) == α`. + +See [`mbind`](@ref) and [`mkernel`](@ref) for details. +""" +function boundmeasure end +export boundmeasure + +boundmeasure(μ::Bind) = μ.α + + +_get_β_a(μ::Bind, a) = asmeasure(μ.f_β(a)) + +function transportmeasure(μ::Bind, x) + tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, x) + tpm_β_a = transportmeasure(_get_β_a(μ, a), b) + mcombine(μ.f_c, tpm_α, tpm_β_a) +end + +localmeasure(μ::Bind, x) = transportmeasure(μ, x) + +tpmeasure_split_combined(f_c, μ::Bind, xy) = _bind_tpm_sc(f_c, μ, xy) + +function _bind_tpm_sc(::typeof(tuple), μ::Bind, xy::Tuple{Vararg{Any,2}}) + x, y = xy[1], xy[2] + tpm_μ = transportmeasure(μ, x) + return tpm_μ, x, y +end + +function _bind_tpm_sc(::Type{Pair}, μ::Bind, xy::Pair) + x, y = xy.first, xy.second + tpm_μ = transportmeasure(μ, x) + return tpm_μ, x, y +end + +const _BindBy{FC} = Bind{<:Any,<:AbstractMeasure,FC} +_bind_tpm_sc(f_c::typeof(vcat), μ::_BindBy{typeof(vcat)}, xy::AbstractVector) = + _bind_tpm_sc_cat(f_c, μ, xy) +_bind_tpm_sc(f_c::typeof(merge), μ::_BindBy{typeof(merge)}, xy::NamedTuple) = + _bind_tpm_sc_cat(f_c, μ, xy) + +function _bind_tpm_sc_cat_lμabyxy(f_c, μ, xy) + tpm_α, a, by = tpmeasure_split_combined(μ.f_c, μ.α, xy) + β_a = _get_β_a(μ, a) + tpm_β_a, b, y = tpmeasure_split_combined(f_c, β_a, by) + tpm_μ = mcombine(μ.f_c, tpm_α, tpm_β_a) + return tpm_μ, a, b, y, xy +end + +function _bind_tpm_sc_cat(f_c::typeof(vcat), μ::_BindBy{typeof(vcat)}, xy::AbstractVector) + tpm_μ, a, b, y, xy = _bind_tpm_sc_cat_lμabyxy(f_c, μ, xy) + # Don't use `x = f_c(a, b)` here, would allocate, splitting xy can use views: + x, y = _split_after(xy, length(a) + length(b)) + return tpm_μ, x, y +end + +function _bind_tpm_sc_cat(f_c::typeof(merge), μ::_BindBy{typeof(merge)}, xy::NamedTuple) + tpm_μ, a, b, y, xy = _bind_tpm_sc_cat_lμabyxy(f_c, μ, xy) + return tpm_μ, f_c(a, b), y +end + + +@inline insupport(μ::Bind, ::Any) = NoFastInsupport{typeof(μ)}() + +@inline getdof(μ::Bind) = NoDOF{typeof(μ)}() + +# Bypass `checked_arg`, would require potentially costly evaluation of f_β: +@inline checked_arg(::Bind, x) = x + +rootmeasure(::Bind) = + throw(ArgumentError("root measure is implicit, but can't be instantiated, for Bind")) + +basemeasure(::Bind) = throw(ArgumentError("basemeasure is not available for Bind")) + +testvalue(::Bind) = throw(ArgumentError("testvalue is not available for Bind")) + +logdensity_def(::Bind, x) = + throw(ArgumentError("logdensity_def is not available for Bind")) + +# Density evaluation consumes the variate parts of the primary and secondary +# measure in a single pass, using the with-rest protocol for value-dependent +# variate sizes: + +logdensityof_impl(μ::Bind, x) = _bind_ld_impl(μ.f_c, μ, x) + +unsafe_logdensityof(μ::Bind, x) = logdensityof_impl(μ, x) + +function _bind_ld_impl(::typeof(tuple), μ::Bind, xy::Tuple{Vararg{Any,2}}) + a, b = xy[1], xy[2] + logdensityof(μ.α, a) + logdensityof(_get_β_a(μ, a), b) +end + +function _bind_ld_impl(::Type{Pair}, μ::Bind, xy::Pair) + a, b = xy.first, xy.second + logdensityof(μ.α, a) + logdensityof(_get_β_a(μ, a), b) +end + +function _bind_ld_impl(::Union{typeof(vcat),typeof(merge)}, μ::Bind, xy) + ℓ, x_μ, x_rest = logdensityof_with_rest(μ, xy) + if !isempty(x_rest) + throw(ArgumentError("Variate too long during density evaluation of a bind")) + end + return ℓ +end + +function _bind_ld_impl(@nospecialize(f_c), @nospecialize(μ::Bind), @nospecialize(xy)) + throw( + ArgumentError( + "Can't compute density of a bind with value combination function of type $(nameof(typeof(f_c)))", + ), + ) +end + +function logdensityof_with_rest(μ::_BindBy{typeof(vcat)}, x::AbstractVector) + ℓ_a, a, x2 = logdensityof_with_rest(μ.α, x) + ℓ_b, b, x_rest = logdensityof_with_rest(_get_β_a(μ, a), x2) + x_μ, _ = _split_after(x, maybestatic_length(x) - maybestatic_length(x_rest)) + return ℓ_a + ℓ_b, x_μ, x_rest +end + +function logdensityof_with_rest(μ::_BindBy{typeof(merge)}, x::NamedTuple) + ℓ_a, a, x2 = logdensityof_with_rest(μ.α, x) + ℓ_b, b, x_rest = logdensityof_with_rest(_get_β_a(μ, a), x2) + return ℓ_a + ℓ_b, merge(a, b), x_rest +end + + +function Base.rand(rng::Random.AbstractRNG, ::Type{T}, μ::Bind) where {T<:Real} + a = rand(rng, T, μ.α) + b = rand(rng, T, _get_β_a(μ, a)) + return μ.f_c(a, b) +end + +function Base.rand(rng::Random.AbstractRNG, μ::Bind) + a = rand(rng, μ.α) + b = rand(rng, _get_β_a(μ, a)) + return μ.f_c(a, b) +end + + +function transport_to_mvstd(ν_inner::StdMeasure, μ::Bind, ab) + tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, ab) + β_a = _get_β_a(μ, a) + y1 = transport_to_mvstd(ν_inner, tpm_α, a) + y2 = transport_to_mvstd(ν_inner, β_a, b) + return vcat(y1, y2) +end -bind(μ, k) = Bind(μ, k) -function Base.rand(rng::AbstractRNG, ::Type{T}, d::Bind) where {T} - x = rand(rng, T, d.μ) - y = rand(rng, T, d.k(x)) - return y +function transport_from_mvstd_with_rest(ν::Bind, μ_inner::StdMeasure, x) + a, x2 = transport_from_mvstd_with_rest(ν.α, μ_inner, x) + β_a = _get_β_a(ν, a) + b, x_rest = transport_from_mvstd_with_rest(β_a, μ_inner, x2) + return ν.f_c(a, b), x_rest end diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl new file mode 100644 index 00000000..d425ae0a --- /dev/null +++ b/src/combinators/combined.jl @@ -0,0 +1,214 @@ +""" + MeasureBase.tpmeasure_split_combined(f_c, α::AbstractMeasure, ab) + +Splits a combined value `ab` that originated from combining a point `a` +from the space of a measure `α` with a point `b` from the space of +another measure `β` via `ab = f_c(a, b)`. + +Returns a semantic equivalent of +`(MeasureBase.transportmeasure(α, a), a, b)`. + +With `a_orig = rand(α)`, `b_orig = rand(β)` and +`ab = f_c(a_orig, b_orig)`, the following must hold true: + +```julia +tpm_α, a, b = tpmeasure_split_combined(f_c, α, ab) +a ≈ a_orig && b ≈ b_orig +``` +""" +function tpmeasure_split_combined end + +function tpmeasure_split_combined(f_c, α::AbstractMeasure, ab) + a, b = _generic_split_combined(f_c, α, ab) + return transportmeasure(α, a), a, b +end + +@inline _generic_split_combined(::typeof(tuple), ::AbstractMeasure, x::Tuple{Vararg{Any,2}}) = x +@inline _generic_split_combined(::Type{Pair}, ::AbstractMeasure, ab::Pair) = (ab...,) + +function _generic_split_combined(f_c::FC, α::AbstractMeasure, ab) where {FC} + _split_variate_byvalue(f_c, testvalue(α), ab) +end + +_split_variate_byvalue(::typeof(vcat), test_a::AbstractVector, ab::AbstractVector) = + _split_after(ab, length(test_a)) + +_split_variate_byvalue(::typeof(vcat), ::Number, ab::AbstractVector) = + _consume_from_stream(ab, ()) + +_split_variate_byvalue(::typeof(vcat), ::NTuple{N,Any}, ab::Tuple) where {N} = + _split_after(ab, Val{N}()) + +function _split_variate_byvalue(::typeof(merge), ::NamedTuple{names_a}, ab::NamedTuple) where {names_a} + _split_after(ab, Val(names_a)) +end + + +@doc raw""" + mcombine(f_c, α::AbstractMeasure, β::AbstractMeasure) + +Combines two measures `α` and `β` to a combined measure via a point +combination function `f_c`. + +`f_c` must combine a given point `a` from the space of measure `α` with a +given point `b` from the space of measure `β` to a single value +`ab = f_c(a, b)` in the space of the combined measure +`μ = mcombine(f_c, α, β)`. + +The combined measure has the mathematical interpretation (on sets +$$A$$ and $$B$$) + +```math +\mu(f_c(A, B)) = \alpha(A)\, \beta(B) +``` +""" +function mcombine end +export mcombine + +@inline function mcombine(f_c, α::AbstractMeasure, β::AbstractMeasure) + _generic_mcombine_impl_stage1(f_c, α, β) +end + +@inline _generic_mcombine_impl_stage1(::typeof(firstarg), α::AbstractMeasure, β::AbstractMeasure) = α +@inline _generic_mcombine_impl_stage1(::typeof(secondarg), α::AbstractMeasure, β::AbstractMeasure) = β + +@inline function _generic_mcombine_impl_stage1(::typeof(tuple), α::AbstractMeasure, β::AbstractMeasure) + productmeasure((α, β)) +end + +@inline function _generic_mcombine_impl_stage1( + f_c::Union{typeof(vcat),typeof(merge)}, + α::AbstractProductMeasure, + β::AbstractProductMeasure, +) + _mcombine_product_shortcut(f_c, marginals(α), marginals(β), α, β) +end + +_mcombine_product_shortcut(::typeof(vcat), ma::AbstractVector, mb::AbstractVector, α, β) = + productmeasure(vcat(ma, mb)) +_mcombine_product_shortcut(::typeof(merge), ma::NamedTuple, mb::NamedTuple, α, β) = + productmeasure(merge(ma, mb)) +_mcombine_product_shortcut(f_c, ma, mb, α, β) = _generic_mcombine_impl_stage2(f_c, α, β) + +@inline function _generic_mcombine_impl_stage1(f_c, α::AbstractMeasure, β::AbstractMeasure) + _generic_mcombine_impl_stage2(f_c, α, β) +end + +@inline function _generic_mcombine_impl_stage2(f_c, α::AbstractMeasure, β::AbstractMeasure) + FC, MA, MB = Core.Typeof(f_c), Core.Typeof(α), Core.Typeof(β) + CombinedMeasure{FC,MA,MB}(f_c, α, β) +end + +@inline function _generic_mcombine_impl_stage2(f_c, α::Dirac, β::Dirac) + Dirac(f_c(α.x, β.x)) +end + + +""" + struct CombinedMeasure <: AbstractMeasure + +Represents a combination of two measures. + +User code should not create instances of `CombinedMeasure` directly, but +should call [`mcombine(f_c, α, β)`](@ref) instead. +""" +struct CombinedMeasure{FC,MA<:AbstractMeasure,MB<:AbstractMeasure} <: AbstractMeasure + f_c::FC + α::MA + β::MB +end + + +@inline insupport(μ::CombinedMeasure, ab) = NoFastInsupport{typeof(μ)}() + +@inline getdof(μ::CombinedMeasure) = getdof(μ.α) + getdof(μ.β) +@inline fast_dof(μ::CombinedMeasure) = fast_dof(μ.α) + fast_dof(μ.β) + +# Bypass `checked_arg`, would require splitting ab: +@inline checked_arg(::CombinedMeasure, ab) = ab + +mdomain(μ::CombinedMeasure) = combinesets(μ.f_c, mdomain(μ.α), mdomain(μ.β)) + +rootmeasure(μ::CombinedMeasure) = mcombine(μ.f_c, rootmeasure(μ.α), rootmeasure(μ.β)) + +basemeasure(μ::CombinedMeasure) = mcombine(μ.f_c, basemeasure(μ.α), basemeasure(μ.β)) + +function logdensity_def(μ::CombinedMeasure, ab) + # Use tpmeasure_split_combined to avoid duplicate calculation of transportmeasure(α): + tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, ab) + return logdensity_def(tpm_α, a) + logdensity_def(μ.β, b) +end + +# Density evaluation consumes the variate parts of both component measures +# in a single pass, using the with-rest protocol for value-dependent +# variate sizes: + +logdensityof_impl(μ::CombinedMeasure, ab) = _combined_ld_impl(μ.f_c, μ, ab) + +unsafe_logdensityof(μ::CombinedMeasure, ab) = logdensityof_impl(μ, ab) + +function _combined_ld_impl(::typeof(tuple), μ::CombinedMeasure, ab::Tuple{Vararg{Any,2}}) + logdensityof(μ.α, ab[1]) + logdensityof(μ.β, ab[2]) +end + +function _combined_ld_impl(::Type{Pair}, μ::CombinedMeasure, ab::Pair) + logdensityof(μ.α, ab.first) + logdensityof(μ.β, ab.second) +end + +function _combined_ld_impl(::Union{typeof(vcat),typeof(merge)}, μ::CombinedMeasure, ab) + ℓ, x_μ, x_rest = logdensityof_with_rest(μ, ab) + if !isempty(x_rest) + throw( + ArgumentError( + "Variate too long during density evaluation of a combined measure", + ), + ) + end + return ℓ +end + +function _combined_ld_impl(f_c, μ::CombinedMeasure, ab) + tpm_α, a, b = tpmeasure_split_combined(f_c, μ.α, ab) + return logdensityof(tpm_α, a) + logdensityof(μ.β, b) +end + +function logdensityof_with_rest(μ::CombinedMeasure{typeof(vcat)}, x::AbstractVector) + ℓ_a, a, x2 = logdensityof_with_rest(μ.α, x) + ℓ_b, b, x_rest = logdensityof_with_rest(μ.β, x2) + x_μ, _ = _split_after(x, maybestatic_length(x) - maybestatic_length(x_rest)) + return ℓ_a + ℓ_b, x_μ, x_rest +end + +function logdensityof_with_rest(μ::CombinedMeasure{typeof(merge)}, x::NamedTuple) + ℓ_a, a, x2 = logdensityof_with_rest(μ.α, x) + ℓ_b, b, x_rest = logdensityof_with_rest(μ.β, x2) + return ℓ_a + ℓ_b, merge(a, b), x_rest +end + + +function Base.rand(rng::Random.AbstractRNG, ::Type{T}, μ::CombinedMeasure) where {T<:Real} + a = rand(rng, T, μ.α) + b = rand(rng, T, μ.β) + return μ.f_c(a, b) +end + +function Base.rand(rng::Random.AbstractRNG, μ::CombinedMeasure) + a = rand(rng, μ.α) + b = rand(rng, μ.β) + return μ.f_c(a, b) +end + + +function transport_to_mvstd(ν_inner::StdMeasure, μ::CombinedMeasure, ab) + tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, ab) + y1 = transport_to_mvstd(ν_inner, tpm_α, a) + y2 = transport_to_mvstd(ν_inner, μ.β, b) + return vcat(y1, y2) +end + + +function transport_from_mvstd_with_rest(ν::CombinedMeasure, μ_inner::StdMeasure, x) + a, x2 = transport_from_mvstd_with_rest(ν.α, μ_inner, x) + b, x_rest = transport_from_mvstd_with_rest(ν.β, μ_inner, x2) + return ν.f_c(a, b), x_rest +end diff --git a/src/combinators/half.jl b/src/combinators/half.jl index 24063b76..c93327b0 100644 --- a/src/combinators/half.jl +++ b/src/combinators/half.jl @@ -19,7 +19,7 @@ function Base.rand(rng::AbstractRNG, ::Type{T}, μ::Half) where {T} return abs(rand(rng, T, unhalf(μ))) end -function logdensityof(μ::Half, x) +function logdensityof_impl(μ::Half, x) ld = logdensityof(unhalf(μ), x) - loghalf return x ≥ 0 ? ld : oftype(ld, -Inf) end 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/likelihood.jl b/src/combinators/likelihood.jl index 6dfd164f..40001007 100644 --- a/src/combinators/likelihood.jl +++ b/src/combinators/likelihood.jl @@ -1,207 +1,214 @@ -export AbstractLikelihood, Likelihood +""" + abstract type AbstractLikelihood <: Function -abstract type AbstractLikelihood end +Abstract supertype for likelihood objects. -# @inline function logdensityof(ℓ::AbstractLikelihood, p) -# t() = dynamic(unsafe_logdensityof(ℓ, p)) -# f() = -Inf -# ifelse(insupport(ℓ, p), t, f)() -# end +Likelihoods are *not* measures, but density functions. They are callable +and also support the DensityInterface API. If `ℒ isa AbstractLikelihood`, +then -# insupport(ℓ::AbstractLikelihood, p) = insupport(ℓ.k(p), ℓ.x) +```julia +DensityInterface.DensityKind(ℒ) == IsDensity() +log(ℒ(θ)) ≈ logdensityof(ℒ, θ) +``` -@doc raw""" - Likelihood(k::AbstractTransitionKernel, x) +Given a transition kernel `k(θ)` (a function that takes a parameter object +and returns a measure) and an observation `x`, the recommended way to create +a likelihood object is -"Observe" a value `x`, yielding a function from the parameters to ℝ. +```julia +ℒ = likelihoodof(k, x) +ℒ isa AbstractLikelihood +``` -Likelihoods are most commonly used in conjunction with an existing _prior_ -measure to yield a new measure, the _posterior_. In Bayes's Law, we have +Then -``P(θ|x) ∝ P(θ) P(x|θ)`` +```julia +log(ℒ(θ)) ≈ logdensityof(ℒ, θ) ≈ logdensityof(k(θ), x) +``` -Here ``P(θ)`` is the prior. If we consider ``P(x|θ)`` as a function on ``θ``, -then it is called a likelihood. +See [`likelihoodof`](@ref) for details on the mathematical semantics of +`k` and `x`. -Since measures are most commonly manipulated using `density` and `logdensity`, -it's awkward to commit a (log-)likelihood to using one or the other. To evaluate -a `Likelihood`, we therefore use `density` or `logdensity`, depending on the -circumstances. In the latter case, it is of course acting as a log-density. +Likelihood-like types that are not subtypes of `AbstractLikelihood` can be +made compatible with the `MeasureBase` likelihoods and Lebesgue integrals by +specializing [`MeasureBase.as_likelihood`](@ref) and +[`MeasureBase.as_integrand`](@ref). +""" +abstract type AbstractLikelihood <: Function end +export AbstractLikelihood -For example, +@inline AbstractLikelihood(l) = as_likelihood(l)::AbstractLikelihood - julia> ℓ = Likelihood(Normal{(:μ,)}, 2.0) - Likelihood(Normal{(:μ,), T} where T, 2.0) +Base.convert(::Type{AbstractLikelihood}, l::AbstractLikelihood) = l +Base.convert(::Type{AbstractLikelihood}, l) = AbstractLikelihood(l) - julia> density_def(ℓ, (μ=2.0,)) - 1.0 - julia> logdensity_def(ℓ, (μ=2.0,)) - -0.0 +""" + likelihood_kernel(ℒ::AbstractLikelihood) -If, as above, the measure includes the parameter information, we can optionally -leave it out of the second argument in the call to `density` or `logdensity`. +Return the transition kernel that is part of likelihood `ℒ`. - julia> density_def(ℓ, 2.0) - 1.0 +If `ℒ = likelihoodof(k, x)` then `likelihood_kernel(ℒ)` must return an +equivalent of `k` (typically but not necessarily `k` itself). +""" +function likelihood_kernel end +export likelihood_kernel - julia> logdensity_def(ℓ, 2.0) - -0.0 -With several parameters, things work as expected: - - julia> ℓ = Likelihood(Normal{(:μ,:σ)}, 2.0) - Likelihood(Normal{(:μ, :σ), T} where T, 2.0) - - julia> logdensity_def(ℓ, (μ=2, σ=3)) - -1.0986122886681098 - - julia> logdensity_def(ℓ, (2,3)) - -1.0986122886681098 - - julia> logdensity_def(ℓ, [2, 3]) - -1.0986122886681098 +""" + likelihood_obs(ℒ::AbstractLikelihood) ---------- +Return the observation that is part of likelihood `ℒ`. - Likelihood(M<:ParameterizedMeasure, constraint::NamedTuple, x) +If `ℒ = likelihoodof(k, x)` then `likelihood_obs(ℒ)` must return an +equivalent of `x` (typically but not necessarily `x` itself). +""" +function likelihood_obs end +export likelihood_obs -In some cases the measure might have several parameters, and we may want the -(log-)likelihood with respect to some subset of them. In this case, we can use -the three-argument form, where the second argument is a constraint. For example, - julia> ℓ = Likelihood(Normal{(:μ,:σ)}, (σ=3.0,), 2.0) - Likelihood(Normal{(:μ, :σ), T} where T, (σ = 3.0,), 2.0) +""" + MeasureBase.as_likelihood(l)::AbstractLikelihood -Similarly to the above, we have +Turn a likelihood-like object `l` into an `AbstractLikelihood`. - julia> density_def(ℓ, (μ=2.0,)) - 0.3333333333333333 +Likelihood-like types that are not subtypes of `AbstractLikelihood` can be +made compatible by specializing - julia> logdensity_def(ℓ, (μ=2.0,)) - -1.0986122886681098 +```julia +MeasureBase.as_likelihood(l::MyLikelihoodType) = likelihoodof(..., ...) +MeasureBase.as_integrand(l::MyLikelihoodType) = MeasureBase.as_likelihood(l) +``` - julia> density_def(ℓ, 2.0) - 0.3333333333333333 +By default, this is implemented for objects like - julia> logdensity_def(ℓ, 2.0) - -1.0986122886681098 +```julia +l = Base.Fix2(densityof, x) ∘ f +l = FuncDensity(Base.Fix2(densityof, x) ∘ f) +l = LogFuncDensity(Base.Fix2(logdensityof, x) ∘ f) +``` +""" +function as_likelihood end +export as_likelihood ------------------------ +@inline as_likelihood(l::AbstractLikelihood) = l -Finally, let's return to the expression for Bayes's Law, +@inline as_integrand(l::AbstractLikelihood) = l -``P(θ|x) ∝ P(θ) P(x|θ)`` -The product on the right side is computed pointwise. To work with this in -MeasureBase, we have a "pointwise product" `⊙`, which takes a measure and a -likelihood, and returns a new measure, that is, the unnormalized posterior that -has density ``P(θ) P(x|θ)`` with respect to the base measure of the prior. +(ℒ::AbstractLikelihood)(p) = densityof(ℒ, p) -For example, say we have - μ ~ Normal() - x ~ Normal(μ,σ) - σ = 1 +DensityInterface.DensityKind(::AbstractLikelihood) = IsDensity() -and we observe `x=3`. We can compute the posterior measure on `μ` as - julia> post = Normal() ⊙ Likelihood(Normal{(:μ, :σ)}, (σ=1,), 3) - Normal() ⊙ Likelihood(Normal{(:μ, :σ), T} where T, (σ = 1,), 3) +_eval_k(ℒ::AbstractLikelihood, p) = asmeasure(likelihood_kernel(ℒ)(p)) - julia> logdensity_def(post, 2) - -2.5 -""" -struct Likelihood{K,X} <: AbstractLikelihood - k::K - x::X +function DensityInterface.logdensityof(ℒ::AbstractLikelihood, p) + logdensityof(_eval_k(ℒ, p), likelihood_obs(ℒ)) +end - Likelihood(k::K, x::X) where {K<:AbstractTransitionKernel,X} = new{K,X}(k, x) - Likelihood(k::K, x::X) where {K<:Function,X} = new{K,X}(k, x) - Likelihood(μ, x) = Likelihood(kernel(μ), x) +function DensityInterface.densityof(ℒ::AbstractLikelihood, p) + exp(ULogarithmic, logdensityof(_eval_k(ℒ, p), likelihood_obs(ℒ))) end -(lik::AbstractLikelihood)(p) = exp(ULogarithmic, logdensityof(lik.k(p), lik.x)) -DensityInterface.DensityKind(::AbstractLikelihood) = IsDensity() +const _SimpleLikelihood1 = ComposedFunction{<:Base.Fix2{typeof(densityof),<:Any},<:Any} +as_likelihood(l::_SimpleLikelihood1) = likelihoodof(l.inner, l.outer.x) +as_integrand(l::_SimpleLikelihood1) = as_likelihood(l) -function Pretty.quoteof(ℓ::Likelihood) - k = Pretty.quoteof(ℓ.k) - x = Pretty.quoteof(ℓ.x) - :(Likelihood($k, $x)) -end +const _SimpleLikelihood2 = DensityInterface.FuncDensity{ + <:ComposedFunction{<:Base.Fix2{typeof(densityof),<:Any},<:Any}, +} +as_likelihood(l::_SimpleLikelihood2) = likelihoodof(l._f.inner, l._f.outer.x) +as_integrand(l::_SimpleLikelihood2) = as_likelihood(l) -function Base.show(io::IO, ℓ::Likelihood) - io = IOContext(io, :compact => true) - Pretty.pprint(io, ℓ) -end +const _SimpleLikelihood3 = DensityInterface.LogFuncDensity{ + <:ComposedFunction{<:Base.Fix2{typeof(logdensityof),<:Any},<:Any}, +} +as_likelihood(l::_SimpleLikelihood3) = likelihoodof(l._log_f.inner, l._log_f.outer.x) +as_integrand(l::_SimpleLikelihood3) = as_likelihood(l) -insupport(ℓ::AbstractLikelihood, p) = insupport(ℓ.k(p), ℓ.x) +const _SimpleLogLikelihood1 = ComposedFunction{<:Base.Fix2{typeof(logdensityof),<:Any},<:Any} +as_integrand_exp(l::_SimpleLogLikelihood1) = likelihoodof(l.inner, l.outer.x) -@inline function logdensityof(ℓ::AbstractLikelihood, p) - logdensityof(ℓ.k(p), ℓ.x) -end -@inline function unsafe_logdensityof(ℓ::AbstractLikelihood, p) - return unsafe_logdensityof(ℓ.k(p), ℓ.x) -end - -# basemeasure(ℓ::Likelihood) = @error "Likelihood requires local base measure" +@doc raw""" + struct Likelihood <: AbstractLikelihood -export likelihoodof +Default result of [`likelihoodof(k, x)`](@ref). +See [`AbstractLikelihood`](@ref) and [`likelihoodof`](@ref) for details. """ - likelihoodof(k::AbstractTransitionKernel, x; constraints...) - likelihoodof(k::AbstractTransitionKernel, x, constraints::NamedTuple) +struct Likelihood{K,X} <: AbstractLikelihood + k::K + x::X -A likelihood is *not* a measure. Rather, a likelihood acts on a measure, through -the "pointwise product" `⊙`, yielding another measure. -""" -function likelihoodof end + Likelihood{K,X}(k, x) where {K,X} = new{K,X}(k, x) +end +export Likelihood -likelihoodof(k, x, ::NamedTuple{()}) = Likelihood(k, x) +# For type stability, in case k is a type (resp. a constructor): +Likelihood(k, x::X) where {X} = Likelihood{Core.Typeof(k),X}(k, x) -likelihoodof(k, x; kwargs...) = likelihoodof(k, x, NamedTuple(kwargs)) +likelihood_kernel(ℒ::Likelihood) = ℒ.k +likelihood_obs(ℒ::Likelihood) = ℒ.x -likelihoodof(k, x, pars::NamedTuple) = likelihoodof(kernel(k, pars), x) +function Pretty.quoteof(ℒ::Likelihood) + k = Pretty.quoteof(ℒ.k) + x = Pretty.quoteof(ℒ.x) + :(Likelihood($k, $x)) +end -likelihoodof(k::AbstractTransitionKernel, x) = Likelihood(k, x) +function Base.show(io::IO, ℒ::Likelihood) + io = IOContext(io, :compact => true) + Pretty.pprint(io, ℒ) +end -export log_likelihood_ratio -""" - log_likelihood_ratio(ℓ::Likelihood, p, q) +@doc raw""" + likelihoodof(k, x)::AbstractLikelihood -Compute the log of the likelihood ratio, in order to compare two choices for -parameters. This is computed as +Returns the likelihood of observing `x` under a family of probability +measures that is generated by a transition kernel `k(θ)`. - logdensity_rel(ℓ.k(p), ℓ.k(q), ℓ.x) +`k(θ)` maps points in the parameter space to measures (resp. objects that can +be converted to measures) on an implicit set `Χ` that contains values like +`x`. -Since `logdensity_rel` can leave common base measure unevaluated, this can be -more efficient than +`likelihoodof(k, x)` returns a likelihood object. A likelihood is **not** a +measure, it is a function from the parameter space to `ℝ₊`. Likelihood +objects can also be interpreted as "generic densities" (but **not** as +probability densities). - logdensityof(ℓ.k(p), ℓ.x) - logdensityof(ℓ.k(q), ℓ.x) -""" -log_likelihood_ratio(ℓ::Likelihood, p, q) = logdensity_rel(ℓ.k(p), ℓ.k(q), ℓ.x) +`likelihoodof(k, x)` implicitly chooses `ξ = rootmeasure(k(θ))` as the +reference measure on the observation set `Χ`. Note that this implicit +`ξ` **must** be independent of `θ`. -# likelihoodof(k, x; kwargs...) = likelihoodof(k, x, NamedTuple(kwargs)) +`ℒ = likelihoodof(k, x)` has the mathematical interpretation -export likelihood_ratio +```math +\mathcal{L}_x(\theta) = \frac{\rm{d}\, k(\theta)}{\rm{d}\, \chi}(x) +``` -""" - likelihood_ratio(ℓ::Likelihood, p, q) +`likelihoodof` must return an object that implements the +[`DensityInterface`](https://github.com/JuliaMath/DensityInterface.jl) API +and `ℒ = likelihoodof(k, x)` must satisfy -Compute the log of the likelihood ratio, in order to compare two choices for -parameters. This is equal to +```julia +log(ℒ(θ)) == logdensityof(ℒ, θ) ≈ logdensityof(k(θ), x) - density_rel(ℓ.k(p), ℓ.k(q), ℓ.x) +DensityKind(ℒ) isa IsDensity +``` -but is computed using LogarithmicNumbers.jl to avoid underflow and overflow. -Since `density_rel` can leave common base measure unevaluated, this can be -more efficient than +[`likelihood_kernel(ℒ)`](@ref) must return an equivalent of `k` and +[`likelihood_obs(ℒ)`](@ref) must return an equivalent of `x` (typically, but +not necessarily, `k` and `x` themselves). - logdensityof(ℓ.k(p), ℓ.x) - logdensityof(ℓ.k(q), ℓ.x) +By default, an instance of [`MeasureBase.Likelihood`](@ref) is returned. """ -function likelihood_ratio(ℓ::Likelihood, p, q) - exp(ULogarithmic, logdensity_rel(ℓ.k(p), ℓ.k(q), ℓ.x)) -end +function likelihoodof end +export likelihoodof + +likelihoodof(k, x) = Likelihood(k, x) diff --git a/src/combinators/pointwise.jl b/src/combinators/pointwise.jl deleted file mode 100644 index 778e7f4e..00000000 --- a/src/combinators/pointwise.jl +++ /dev/null @@ -1,30 +0,0 @@ -export ⊙ - -struct PointwiseProductMeasure{P,L} <: AbstractMeasure - prior::P - likelihood::L -end - -iterate(p::PointwiseProductMeasure, i = 1) = iterate((p.prior, p.likelihood), i) - -function Pretty.tile(d::PointwiseProductMeasure) - Pretty.pair_layout(Pretty.tile(d.prior), Pretty.tile(d.likelihood), sep = " ⊙ ") -end - -⊙(prior, ℓ) = pointwiseproduct(prior, ℓ) - -@inbounds function insupport(d::PointwiseProductMeasure, p) - prior, ℓ = d - istrue(insupport(prior, p)) && istrue(insupport(ℓ, p)) -end - -@inline function logdensity_def(d::PointwiseProductMeasure, p) - prior, ℓ = d - unsafe_logdensityof(ℓ, p) -end - -basemeasure(d::PointwiseProductMeasure) = d.prior - -function gentype(d::PointwiseProductMeasure) - gentype(d.prior) -end diff --git a/src/combinators/power.jl b/src/combinators/power.jl index e6397c3f..62e6ba4b 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,31 @@ 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)) -end +marginals(d::PowerMeasure) = maybestatic_fill(d.parent, d.axes) -marginals(d::PowerMeasure) = fill_with(d.parent, d.axes) +# Powers of scalar-variate measures have array-valued variates of known size: +@inline mspace_elsize(μ::PowerMeasure) = _pwr_mspace_elsize(μ, mspace_elsize(pwr_base(μ))) +@inline _pwr_mspace_elsize(μ::PowerMeasure, ::Tuple{}) = pwr_size(μ) +@inline _pwr_mspace_elsize(μ::PowerMeasure, ::Any) = NoMSpaceElementSize{typeof(μ)}() 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)) @@ -78,25 +102,19 @@ params(d::PowerMeasure) = params(first(marginals(d))) basemeasure(d.parent)^d.axes 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) - end - end +# Power structure is unwrapped into the power axes arguments of the batched +# density machinery (see density-batched.jl), which fuses evaluation over +# flat variate storage: - @eval @inline function $func(d::PowerMeasure{M,Tuple{Static.SOneTo{N}}}, x) where {M,N} - parent = d.parent - sum(1:N) do j - @inbounds $func(parent, x[j]) - end +for head in [:logdensityof_impl, :logdensity_def] + @eval @inline function $head(d::PowerMeasure, x) + _powered_ld($head, pwr_base(d), x, pwr_axes(d)) end - @eval @inline function $func( - d::PowerMeasure{M,NTuple{N,Static.SOneTo{0}}}, + @eval @inline function $head( + ::PowerMeasure{<:Any,<:Tuple{Vararg{StaticOneToLike{0}}}}, x, - ) where {M,N} + ) static(0.0) end end @@ -109,23 +127,32 @@ end end end +_all(A) = all(A) +_all(::AbstractArray{NoFastInsupport{T}}) where {T} = NoFastInsupport{T}() + @inline function insupport(μ::PowerMeasure, x::AbstractArray) p = μ.parent - all(x) do xj + insupp = broadcast(x) do xj # https://github.com/SciML/Static.jl/issues/36 dynamic(insupport(p, xj)) end + _all(insupp) end -@inline getdof(μ::PowerMeasure) = getdof(μ.parent) * prod(map(length, μ.axes)) +@inline getdof(μ::PowerMeasure) = getdof(μ.parent) * size2length(axes2size(μ.axes)) +@inline fast_dof(μ::PowerMeasure) = fast_dof(μ.parent) * size2length(axes2size(μ.axes)) -@inline function getdof(::PowerMeasure{<:Any,NTuple{N,Static.SOneTo{0}}}) where {N} +# Static.SOneTo(0) is not static (yet): +@inline function getdof(::PowerMeasure{<:Any,<:NTuple{N,StaticOneToLike{0}}}) where {N} + static(0) +end +@inline function fast_dof(::PowerMeasure{<:Any,<:NTuple{N,StaticOneToLike{0}}}) where {N} static(0) end @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")) @@ -142,10 +169,11 @@ massof(m::PowerMeasure) = massof(m.parent)^prod(m.axes) logdensity_def(::PowerMeasure{P}, x) where {P<:PrimitiveMeasure} = static(0.0) -# To avoid ambiguities +# Disambiguation with the static-zero-size power density method: function logdensity_def( - ::PowerMeasure{P,Tuple{Vararg{Static.SOneTo{0},N}}}, - x, -) where {P<:PrimitiveMeasure,N} + ::PowerMeasure{P,<:Tuple{Vararg{StaticOneToLike{0}}}}, + ::Any, +) where {P<:PrimitiveMeasure} static(0.0) end + diff --git a/src/combinators/powerweighted.jl b/src/combinators/powerweighted.jl deleted file mode 100644 index 47f50da4..00000000 --- a/src/combinators/powerweighted.jl +++ /dev/null @@ -1,37 +0,0 @@ -export ↑ - -struct PowerWeightedMeasure{M,A} <: AbstractMeasure - parent::M - exponent::A -end - -logdensity_def(d::PowerWeightedMeasure, x) = d.exponent * logdensity_def(d.parent, x) - -basemeasure(d::PowerWeightedMeasure, x) = basemeasure(d.parent, x)↑d.exponent - -basemeasure(d::PowerWeightedMeasure) = basemeasure(d.parent)↑d.exponent - -function powerweightedmeasure(d, α) - isone(α) && return d - PowerWeightedMeasure(d, α) -end - -(d::AbstractMeasure)↑α = powerweightedmeasure(d, α) - -insupport(d::PowerWeightedMeasure, x) = insupport(d.parent, x) - -function Base.show(io::IO, d::PowerWeightedMeasure) - print(io, d.parent, " ↑ ", d.exponent) -end - -function powerweightedmeasure(d::PowerWeightedMeasure, α) - powerweightedmeasure(d.parent, α * d.exponent) -end - -function powerweightedmeasure(d::WeightedMeasure, α) - weightedmeasure(α * d.logweight, powerweightedmeasure(d.base, α)) -end - -function Pretty.tile(d::PowerWeightedMeasure) - Pretty.pair_layout(Pretty.tile(d.parent), Pretty.tile(d.exponent), sep = " ↑ ") -end diff --git a/src/combinators/product.jl b/src/combinators/product.jl index 0290419d..22dae865 100644 --- a/src/combinators/product.jl +++ b/src/combinators/product.jl @@ -72,8 +72,8 @@ function _rand_product( end |> collect end -for func in [:logdensityof, :logdensity_def] - @eval @inline function $func(d::AbstractProductMeasure, x) +for (head, func) in [(:logdensityof_impl, :logdensityof), (:logdensity_def, :logdensity_def)] + @eval @inline function $head(d::AbstractProductMeasure, x) mapreduce($func, +, marginals(d), x) end end @@ -82,10 +82,27 @@ struct ProductMeasure{M} <: AbstractProductMeasure marginals::M end -@inline function logdensity_rel(μ::ProductMeasure, ν::ProductMeasure, x) - mapreduce(logdensity_rel, +, marginals(μ), marginals(ν), x) +proxy(μ::ProductMeasure{<:FillArrays.Fill}) = + powermeasure(_fill_value(marginals(μ)), _fill_axes(marginals(μ))) + +# Relative densities between products evaluate marginal-wise. Support +# checks happen at the logdensity_rel level for the whole products, so the +# unsafe marginal evaluation suffices here: +@inline function logdensity_rel_def(μ::ProductMeasure, ν::ProductMeasure, x) + mapreduce(unsafe_logdensity_rel, +, marginals(μ), marginals(ν), x) +end + +# For tuples, `mapreduce` has trouble with type inference: +@inline function logdensity_rel_def( + μ::ProductMeasure{<:Tuple}, + ν::ProductMeasure{<:Tuple}, + x, +) + sum(map(unsafe_logdensity_rel, marginals(μ), marginals(ν), x)) end +_mspace_names(μ::ProductMeasure{<:NamedTuple{names}}) where {names} = names + function Pretty.tile(d::ProductMeasure{T}) where {T<:Tuple} Pretty.list_layout(Pretty.tile.([marginals(d)...]), sep = " ⊗ ") end @@ -109,14 +126,14 @@ end return q end -for func in [:logdensityof, :logdensity_def] +for (head, func) in [(:logdensityof_impl, :logdensityof), (:logdensity_def, :logdensity_def)] # For tuples, `mapreduce` has trouble with type inference - @eval @inline function $func(d::ProductMeasure{T}, x) where {T<:Tuple} + @eval @inline function $head(d::ProductMeasure{T}, x) where {T<:Tuple} ℓs = map($func, marginals(d), x) sum(ℓs) end - @eval function $func(d::ProductMeasure{NamedTuple{N,T}}, x) where {N,T} + @eval function $head(d::ProductMeasure{NamedTuple{N,T}}, x) where {N,T} _product_gen_impl(Val($func), d, x) end end @@ -160,6 +177,17 @@ function _basemeasure( productmeasure(mappedarray(basemeasure, mar)) end +""" + MeasureBase.basekernel(f) + +For a function `f` that returns a measure, return the function that returns +the base measure instead, satisfying `basekernel(f)(p) == basemeasure(f(p))`. +""" +function basekernel end + +basekernel(f) = basemeasure ∘ f +basekernel(f::Returns) = Returns(basemeasure(f.value)) + function _basemeasure( μ::ProductMeasure{Base.Generator{I,F}}, ::Type{B}, @@ -179,18 +207,6 @@ function testvalue(::Type{T}, d::AbstractProductMeasure) where {T} _map(m -> testvalue(T, m), marginals(d)) end -export ⊗ - -""" - ⊗(μs::AbstractMeasure...) - -`⊗` is a binary operator for building product measures. This satisfies the law - -``` - basemeasure(μ ⊗ ν) == basemeasure(μ) ⊗ basemeasure(ν) -``` -""" -⊗(μs::AbstractMeasure...) = productmeasure(μs) ############################################################################### # I <: Base.Generator @@ -233,17 +249,25 @@ end @inline function insupport(d::AbstractProductMeasure, x) for (mj, xj) in zip(marginals(d), x) - dynamic(insupport(mj, xj)) || return false + insup = dynamic(insupport(mj, xj)) + if insup isa NoFastInsupport || insup == false + return insup + end end return true end -getdof(d::AbstractProductMeasure) = mapreduce(getdof, +, marginals(d)) +getdof(d::AbstractProductMeasure) = sum(getdof, marginals(d)) +fast_dof(d::AbstractProductMeasure) = sum(fast_dof, marginals(d)) function checked_arg(μ::ProductMeasure{<:NTuple{N,Any}}, x::NTuple{N,Any}) where {N} 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/product_transport.jl b/src/combinators/product_transport.jl new file mode 100644 index 00000000..a534300b --- /dev/null +++ b/src/combinators/product_transport.jl @@ -0,0 +1,447 @@ +""" + transport_to(ν, ::Type{MU}) where {MU<:StdMeasure} + transport_to(::Type{NU}, μ) where {NU<:StdMeasure} + +As a user convenience, a standard measure type like [`StdUniform`](@ref), +[`StdExponential`](@ref), [`StdNormal`](@ref) or [`StdLogistic`](@ref) +may be used directly as the source or target of a measure transport. + +Depending on [`MeasureBase.some_dof(μ)`](@ref) (resp. `ν`), an instance of +the standard measure itself or a power of it will be automatically chosen as +the transport partner. + +Example: + +```julia +transport_to(StdNormal, μ) +transport_to(ν, StdNormal) +``` +""" +function transport_to(ν, ::Type{MU}) where {MU<:StdMeasure} + transport_to(ν, _std_tp_partner(MU, ν)) +end + +function transport_to(::Type{NU}, μ) where {NU<:StdMeasure} + transport_to(_std_tp_partner(NU, μ), μ) +end + +function transport_to(::Type{NU}, ::Type{MU}) where {NU<:StdMeasure,MU<:StdMeasure} + throw( + ArgumentError( + "Can't construct a transport function between the types of two standard measures, need a measure instance on one side", + ), + ) +end + +_std_tp_partner(::Type{M}, μ) where {M<:StdMeasure} = _std_tp_partner_bydof(M, some_dof(μ)) +_std_tp_partner_bydof(::Type{M}, ::StaticInteger{1}) where {M<:StdMeasure} = M() +_std_tp_partner_bydof(::Type{M}, dof::IntegerLike) where {M<:StdMeasure} = M()^dof +function _std_tp_partner_bydof(::Type{M}, ::AbstractNoDOF{MU}) where {M<:StdMeasure,MU} + throw( + ArgumentError( + "Can't determine a standard transport partner for measures of type $(nameof(MU))", + ), + ) +end + + +# For transport, always pull a multi-dimensional PowerMeasure back to a +# one-dimensional PowerMeasure first: + +const _PowerMeasureRank1{M} = PowerMeasure{M,<:NTuple{1,OneToLike}} + +function transport_origin(μ::PowerMeasure) + pwr_base(μ)^prod(pwr_size(μ)) +end + +function to_origin(μ::PowerMeasure, x) + maybestatic_reshape(x, (prod(pwr_size(μ)),)) +end + +function from_origin(μ::PowerMeasure, x_origin) + # Sanity check, should never fail: + @assert x_origin isa AbstractVector + maybestatic_reshape(x_origin, pwr_size(μ)) +end + + +# A one-dimensional PowerMeasure has an origin if its parent has an origin: + +function transport_origin(μ::_PowerMeasureRank1) + _pwr_origin(typeof(μ), transport_origin(pwr_base(μ)), pwr_axes(μ)) +end +_pwr_origin(::Type{MU}, parent_origin, axes) where {MU} = parent_origin^axes +_pwr_origin(::Type{MU}, ::NoTransportOrigin, axes) where {MU} = NoTransportOrigin{MU}() + +function to_origin(μ::_PowerMeasureRank1, x) + to_origin.(Ref(pwr_base(μ)), x) +end + +function from_origin(μ::_PowerMeasureRank1, x_origin) + # Sanity check, should never fail: + @assert x_origin isa AbstractVector + from_origin.(Ref(pwr_base(μ)), x_origin) +end + + +# Transport between powers of standard measures, of any rank: + +function _stdpow_transport_def(ν::StdPowerMeasure{NU}, μ::StdPowerMeasure{MU}, x) where {NU,MU} + y = transport_to(pwr_base(ν), pwr_base(μ)).(x) + maybestatic_reshape(y, pwr_size(ν)) +end + +function _stdpow_transport_def(ν::StdPowerMeasure{MU}, μ::StdPowerMeasure{MU}, x) where {MU} + maybestatic_reshape(x, pwr_size(ν)) +end + +transport_def(ν::StdPowerMeasure{NU}, μ::StdPowerMeasure{MU}, x) where {NU,MU} = + _stdpow_transport_def(ν, μ, x) + +# Disambiguation with the mvstd gateway methods below: +transport_def(ν::StdPowerMeasure{NU,1}, μ::StdPowerMeasure{MU}, x) where {NU,MU} = + _stdpow_transport_def(ν, μ, x) +transport_def(ν::StdPowerMeasure{NU}, μ::StdPowerMeasure{MU,1}, x) where {NU,MU} = + _stdpow_transport_def(ν, μ, x) +transport_def(ν::StdPowerMeasure{NU,1}, μ::StdPowerMeasure{MU,1}, x) where {NU,MU} = + _stdpow_transport_def(ν, μ, x) + + +# Transport between univariate standard measures and one-dimensional power +# measures of size one: + +function transport_def(ν::StdMeasure, μ::StdPowerMeasure{MU,1}, x) where {MU} + return transport_def(ν, pwr_base(μ), only(x)) +end + +function transport_def(ν::StdPowerMeasure{NU,1}, μ::StdMeasure, x) where {NU} + sz_ν = pwr_size(ν) + @assert prod(sz_ν) == 1 + return maybestatic_fill(transport_def(pwr_base(ν), μ, x), sz_ν) +end + + +# Transport to a multivariate standard measure from any measure: + +function transport_def(ν::StdPowerMeasure{NU,1}, μ::AbstractMeasure, x) where {NU} + y = transport_to_mvstd(pwr_base(ν), μ, x) + if maybestatic_length(y) != maybestatic_length(ν) + throw(ArgumentError("Length of transport target doesn't match variate DOF during transport")) + end + return y +end + +function transport_to_mvstd(ν_inner::StdMeasure, μ::AbstractMeasure, x) + return _to_mvstd_withdof(ν_inner, μ, fast_dof(μ), x) +end + +# For standard measures and their powers specialized `transport_def` methods +# exist, for other measures the origin-based machinery must be used directly +# instead of `transport_def`, to prevent infinite dispatch recursion via the +# gateway methods above: +const _StdOrStdPowerMeasure = Union{StdMeasure,StdPowerMeasure} + +_transport_def_nongateway(ν::_StdOrStdPowerMeasure, μ::_StdOrStdPowerMeasure, x) = + transport_def(ν, μ, x) +function _transport_def_nongateway(ν, μ, x) + _transport_between_origins(ν, _origin_depth(ν), _origin_depth(μ), μ, x) +end + +function _to_mvstd_withdof(ν_inner::StdMeasure, μ::AbstractMeasure, dof_μ::IntegerLike, x) + _transport_def_nongateway(ν_inner^dof_μ, μ, x) +end + +function _to_mvstd_withdof(ν_inner::StdMeasure, μ::AbstractMeasure, ::AbstractNoDOF, x) + _to_mvstd_withorigin(ν_inner, μ, transport_origin(μ), x) +end + +function _to_mvstd_withorigin(ν_inner::StdMeasure, μ::AbstractMeasure, μ_origin, x) + x_origin = to_origin(μ, x) + transport_to_mvstd(ν_inner, μ_origin, x_origin) +end + +function _to_mvstd_withorigin(ν_inner::StdMeasure, μ::AbstractMeasure, ::NoTransportOrigin, x) + throw( + ArgumentError( + "Don't know how to transport values of type $(nameof(typeof(x))) from $(nameof(typeof(μ))) to a power of $(nameof(typeof(ν_inner)))", + ), + ) +end + + +# Transport from a multivariate standard measure to any measure: + +function transport_def(ν::AbstractMeasure, μ::StdPowerMeasure{MU,1}, x) where {MU} + _transport_from_mvstd(ν, pwr_base(μ), x) +end + +function _transport_from_mvstd(ν::AbstractMeasure, μ_inner::StdMeasure, x) + y, x_rest = transport_from_mvstd_with_rest(ν, μ_inner, x) + if !isempty(x_rest) + throw(ArgumentError("Input value too long during transport")) + end + return y +end + +function transport_from_mvstd_with_rest(ν::AbstractMeasure, μ_inner::StdMeasure, x) + dof_ν = fast_dof(ν) + return _from_mvstd_with_rest_withdof(ν, dof_ν, μ_inner, x) +end + +function _from_mvstd_with_rest_withdof( + ν::AbstractMeasure, + dof_ν::IntegerLike, + μ_inner::StdMeasure, + x, +) + len_x = maybestatic_length(x) + + # Since we can't check the DOF of measures like Bind, we could "run out + # of x" if the original x was too short. `transport_to` below will detect + # this, but better throw a more informative exception here: + if len_x < dof_ν + throw(ArgumentError("Variate too short during transport")) + end + + x_inner_dof, x_rest = _split_after(x, dof_ν) + y = _transport_def_nongateway(ν, μ_inner^dof_ν, x_inner_dof) + return y, x_rest +end + +function _from_mvstd_with_rest_withdof( + ν::AbstractMeasure, + ::AbstractNoDOF, + μ_inner::StdMeasure, + x, +) + _from_mvstd_with_rest_withorigin(ν, transport_origin(ν), μ_inner, x) +end + +function _from_mvstd_with_rest_withorigin( + ν::AbstractMeasure, + ν_origin, + μ_inner::StdMeasure, + x, +) + x_origin, x_rest = transport_from_mvstd_with_rest(ν_origin, μ_inner, x) + from_origin(ν, x_origin), x_rest +end + +function _from_mvstd_with_rest_withorigin( + ν::AbstractMeasure, + ::NoTransportOrigin, + μ_inner::StdMeasure, + x, +) + throw( + ArgumentError( + "Don't know how to transport a value of type $(nameof(typeof(x))) from a power of $(nameof(typeof(μ_inner))) to $(nameof(typeof(ν)))", + ), + ) +end + + +# Transport between a standard measure and Dirac: + +@inline transport_from_mvstd_with_rest(ν::Dirac, ::StdMeasure, x::Any) = ν.x, x + +@inline transport_to_mvstd(::StdMeasure, ::Dirac, ::Any) = FillArrays.Zeros{Bool}(0) + + +# Pull back from a product over a Fill to a power measure: + +@inline transport_origin(μ::ProductMeasure) = _marginals_tp_origin(marginals(μ)) +@inline to_origin(μ::ProductMeasure, x) = _marginals_to_origin(marginals(μ), x) +@inline from_origin(μ::ProductMeasure, x_origin) = + _marginals_from_origin(marginals(μ), x_origin) + +_marginals_tp_origin(::Ms) where {Ms} = NoTransportOrigin{ProductMeasure{Ms}}() + +_marginals_tp_origin(marginals_μ::FillArrays.Fill) = + _fill_value(marginals_μ)^_fill_axes(marginals_μ) +_marginals_to_origin(::FillArrays.Fill, x) = x +_marginals_from_origin(::FillArrays.Fill, x_origin) = x_origin + + +# Pull back from a NamedTuple product measure to a Tuple product measure: +# +# Maybe ToDo (breaking): For transport between NamedTuple-marginals we could +# match names where possible, even if given in different order, and transport +# between the remaining non-matching names in the order given. This may not +# be worth the additional complexity, though, since transport is typically +# used with a (power of a) standard measure on one side. + +_marginals_tp_origin(marginals_μ::NamedTuple{names}) where {names} = + productmeasure(values(marginals_μ)) +_marginals_to_origin(::NamedTuple{names}, x::NamedTuple{names}) where {names} = values(x) +_marginals_from_origin(::NamedTuple{names}, x_origin::Tuple) where {names} = + NamedTuple{names}(x_origin) + + +# Transport between two instances of ProductMeasure: + +transport_def(ν::ProductMeasure, μ::ProductMeasure, x) = + _marginal_transport_def(marginals(ν), marginals(μ), x) + +function _marginal_transport_def(marginals_ν, marginals_μ, x) + @assert size(marginals_ν) == size(marginals_μ) == size(x) # Sanity check, should not fail + transport_def.(marginals_ν, marginals_μ, x) +end + +function _marginal_transport_def( + marginals_ν::Tuple{Vararg{AbstractMeasure,N}}, + marginals_μ::Tuple{Vararg{AbstractMeasure,N}}, + x, +) where {N} + map(transport_def, marginals_ν, marginals_μ, x) +end + +function _marginal_transport_def( + marginals_ν::AbstractVector{<:AbstractMeasure}, + marginals_μ::Tuple{Vararg{AbstractMeasure,N}}, + x, +) where {N} + _marginal_transport_def(_as_tuple(marginals_ν, Val(N)), marginals_μ, x) +end + +function _marginal_transport_def( + marginals_ν::Tuple{Vararg{AbstractMeasure,N}}, + marginals_μ::AbstractVector{<:AbstractMeasure}, + x, +) where {N} + _marginal_transport_def(marginals_ν, _as_tuple(marginals_μ, Val(N)), _as_tuple(x, Val(N))) +end + + +# Transport from a ProductMeasure to a standard measure: + +function transport_to_mvstd(ν_inner::StdMeasure, μ::ProductMeasure, x) + _marginals_to_mvstd(ν_inner, marginals(μ), x) +end + +struct _TransportToMvStd{NU<:StdMeasure} <: Function end +(::_TransportToMvStd{NU})(μ, x) where {NU} = transport_to_mvstd(NU(), μ, x) + +function _marginals_to_mvstd(::NU, marginals_μ::Tuple, x::Tuple) where {NU<:StdMeasure} + _flatten_to_rv(map(_TransportToMvStd{NU}(), marginals_μ, x)) +end + +function _marginals_to_mvstd( + ν::NU, + marginals_μ::NamedTuple{names}, + x::NamedTuple{names}, +) where {NU<:StdMeasure,names} + _marginals_to_mvstd(ν, values(marginals_μ), values(x)) +end + +function _marginals_to_mvstd(::NU, marginals_μ, x) where {NU<:StdMeasure} + _flatten_to_rv(broadcast(_TransportToMvStd{NU}(), marginals_μ, x)) +end + + +# Transport from a standard measure to a ProductMeasure, with rest: + +const _MaybeUnknownDOF = Union{IntegerLike,AbstractNoDOF} + +const _KnownDOFs = Union{Tuple{Vararg{IntegerLike,N}} where N,StaticVector{<:IntegerLike}} + +function transport_from_mvstd_with_rest(ν::ProductMeasure, μ_inner::StdMeasure, x) + νs = marginals(ν) + dofs = map(fast_dof, νs) + return _marginals_from_mvstd_with_rest(νs, dofs, μ_inner, x) +end + +function transport_from_mvstd_with_rest( + ν::ProductMeasure{<:NamedTuple{names}}, + μ_inner::StdMeasure, + x, +) where {names} + ys, x_rest = + transport_from_mvstd_with_rest(productmeasure(values(marginals(ν))), μ_inner, x) + return NamedTuple{names}(ys), x_rest +end + +function _dof_access_firstidxs(dofs::Tuple{Vararg{IntegerLike,N}}, first_idx) where {N} + cumsum((first_idx, dofs[begin:(end-1)]...)) +end + +function _dof_access_firstidxs(dofs::AbstractVector{<:IntegerLike}, first_idx) + # ToDo: Improve implementation (reduce memory allocations): + cumsum(vcat([eltype(dofs)(first_idx)], dofs[begin:(end-1)])) +end + +function _split_x_by_marginals_with_rest( + dofs::Union{Tuple,AbstractVector}, + x::AbstractVector{<:Number}, +) + x_idxs = maybestatic_eachindex(x) + first_idxs = _dof_access_firstidxs(dofs, maybestatic_first(x_idxs)) + xs = map((from, n) -> _get_or_view(x, from, from + n - one(n)), first_idxs, dofs) + x_rest = _get_or_view(x, first_idxs[end] + dofs[end], maybestatic_last(x_idxs)) + return xs, x_rest +end + +function _marginals_from_mvstd_with_rest( + νs, + dofs::_KnownDOFs, + μ_inner::StdMeasure, + x::AbstractVector{<:Number}, +) + xs, x_rest = _split_x_by_marginals_with_rest(dofs, x) + μs = map(n -> μ_inner^n, dofs) + ys = map(transport_def, νs, μs, xs) + return ys, x_rest +end + +function _marginals_from_mvstd_with_rest( + νs, + dofs, + μ_inner::StdMeasure, + x::AbstractVector{<:Number}, +) + _marginals_from_mvstd_with_rest_nodof(νs, μ_inner, x) +end + +function _marginals_from_mvstd_with_rest_nodof( + νs::Tuple{Vararg{AbstractMeasure}}, + μ_inner::StdMeasure, + x::AbstractVector{<:Number}, +) + # ToDo: Check for type stability, may need a generated function: + y1, x_rest = transport_from_mvstd_with_rest(νs[1], μ_inner, x) + y2_end, x_final_rest = _marginals_from_mvstd_with_rest_nodof(Base.tail(νs), μ_inner, x_rest) + return (y1, y2_end...), x_final_rest +end + +_marginals_from_mvstd_with_rest_nodof(::Tuple{}, ::StdMeasure, x::AbstractVector{<:Number}) = + (), x + +function _marginals_from_mvstd_with_rest_nodof( + νs::AbstractVector{M}, + μ_inner::StdMeasure, + x::AbstractVector{<:Number}, +) where {M<:AbstractMeasure} + if isconcretetype(M) + # Marginals of concrete type produce variates of uniform type, so + # the loop below is type stable (the type of the remaining stream + # stays invariant under repeated view-taking): + idxs = eachindex(νs) + y1, x_rest = transport_from_mvstd_with_rest(νs[first(idxs)], μ_inner, x) + ys = Vector{typeof(y1)}(undef, length(idxs)) + ys[begin] = y1 + j = firstindex(ys) + 1 + for i in Iterators.drop(idxs, 1) + ys[j], x_rest = transport_from_mvstd_with_rest(νs[i], μ_inner, x_rest) + j += 1 + end + return ys, x_rest + else + # Fallback for marginals of mixed type: + ys_any = Vector{Any}(undef, length(eachindex(νs))) + x_rest = x + for (i, ν) in zip(eachindex(ys_any), νs) + ys_any[i], x_rest = transport_from_mvstd_with_rest(ν, μ_inner, x_rest) + end + return [y for y in ys_any], x_rest + end +end diff --git a/src/combinators/reshape.jl b/src/combinators/reshape.jl new file mode 100644 index 00000000..f77dbb68 --- /dev/null +++ b/src/combinators/reshape.jl @@ -0,0 +1,56 @@ +# 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, some_mspace_elsize(m)), m) diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl index 26ba3948..be69d142 100644 --- a/src/combinators/smart-constructors.jl +++ b/src/combinators/smart-constructors.jl @@ -1,187 +1,267 @@ +# Canonical measure type nesting, outer to inner: +# +# WeightedMeasure, Dirac, PowerMeasure, ProductMeasure + + ############################################################################### # Half +""" + half(μ::AbstractMeasure) + +Constructs the half-measure of a measure `μ` that is symmetric around zero: +`μ` folded onto the non-negative half-line. +""" half(μ::AbstractMeasure) = Half(μ) +export half ############################################################################### -# PointwiseProductMeasure +# PowerMeaure -function pointwiseproduct(μ::AbstractMeasure, ℓ::Likelihood) - T = Core.Compiler.return_type(ℓ.k, Tuple{gentype(μ)}) - return pointwiseproduct(T, μ, ℓ) -end +""" + powermeasure(μ, dims) + powermeasure(μ, axes) -function pointwiseproduct(::Type{T}, μ::AbstractMeasure, ℓ::Likelihood) where {T} - return PointwiseProductMeasure(μ, ℓ) -end +Constructs a power of a measure `μ`. -############################################################################### -# PowerMeaure +`powermeasure(μ, exponent)` is semantically equivalent to +`productmeasure(Fill(μ, exponent))`, but more efficient. +""" +function powermeasure end +export powermeasure + +@inline powermeasure(μ, exponent) = _generic_powermeasure_stage1(asmeasure(μ), asaxes(exponent)) -powermeasure(m::AbstractMeasure, ::Tuple{}) = m +@inline _generic_powermeasure_stage1(μ::AbstractMeasure, ::Tuple{}) = μ -function powermeasure( - μ::WeightedMeasure, - dims::Tuple{<:AbstractArray,Vararg{AbstractArray}}, -) - k = mapreduce(length, *, dims) * μ.logweight - return weightedmeasure(k, μ.base^dims) +@inline function _generic_powermeasure_stage1(μ::AbstractMeasure, exponent::Tuple) + _generic_powermeasure_stage2(μ, exponent) end -function powermeasure(μ::WeightedMeasure, dims::NonEmptyTuple) - k = prod(dims) * μ.logweight - return weightedmeasure(k, μ.base^dims) +@inline _generic_powermeasure_stage2(μ::AbstractMeasure, exponent::Tuple) = + PowerMeasure(μ, exponent) + +@inline function _generic_powermeasure_stage2(μ::Dirac, exponent::Tuple) + Dirac(maybestatic_fill(μ.x, exponent)) +end + +@inline function _generic_powermeasure_stage2(μ::WeightedMeasure, exponent::Tuple) + ν = μ.base^exponent + k = size2length(axes2size(exponent)) * μ.logweight + return weightedmeasure(k, ν) end ############################################################################### # ProductMeasure -productmeasure(mar::FillArrays.Fill) = powermeasure(mar.value, mar.axes) - -function productmeasure(mar::ReadonlyMappedArray{T,N,A,Returns{M}}) where {T,N,A,M} - return powermeasure(mar.f.value, axes(mar.data)) +""" + productmeasure(μs) + +Constructs a product over a collection `μs` of measures. + +Examples: + +```julia +productmeasure((StdNormal(), StdExponential())) +productmeasure((a = StdNormal(), b = StdExponential())) +productmeasure([pushfwd(Base.Fix1(*, scale), StdExponential()) for scale in 0.1:0.2:2]) +``` +""" +function productmeasure end +export productmeasure + +@inline productmeasure(mar) = _generic_productmeasure_impl(mar) + +@inline _generic_productmeasure_impl(mar::FillArrays.Fill) = + powermeasure(_fill_value(mar), _fill_axes(mar)) + +# Empty products are unit measures: +@inline _generic_productmeasure_impl(::Tuple{}) = Dirac(()) +@inline _generic_productmeasure_impl(::NamedTuple{()}) = Dirac(NamedTuple()) + +@inline _generic_productmeasure_impl(mar::Tuple{Vararg{AbstractMeasure}}) = + ProductMeasure(mar) +_generic_productmeasure_impl(mar::Tuple{Dirac,Vararg{Dirac}}) = Dirac(map(m -> m.x, mar)) +_generic_productmeasure_impl(mar::Tuple{WeightedMeasure,Vararg{WeightedMeasure}}) = + weightedmeasure(sum(map(_logweight, mar)), productmeasure(map(m -> m.base, mar))) +_generic_productmeasure_impl(mar::Tuple) = productmeasure(map(asmeasure, mar)) + +@inline _generic_productmeasure_impl( + mar::NamedTuple{names,<:Tuple{Vararg{AbstractMeasure}}}, +) where {names} = ProductMeasure(mar) +_generic_productmeasure_impl( + mar::NamedTuple{names,<:Tuple{Dirac,Vararg{Dirac}}}, +) where {names} = Dirac(map(m -> m.x, mar)) +_generic_productmeasure_impl( + mar::NamedTuple{names,<:Tuple{WeightedMeasure,Vararg{WeightedMeasure}}}, +) where {names} = + weightedmeasure(sum(map(_logweight, values(mar))), productmeasure(map(m -> m.base, mar))) +_generic_productmeasure_impl(mar::NamedTuple) = productmeasure(map(asmeasure, mar)) + +_generic_productmeasure_impl(mar::AbstractArray{<:Dirac}) = Dirac((m -> m.x).(mar)) + +_generic_productmeasure_impl(mar::AbstractArray{<:WeightedMeasure}) = + weightedmeasure(sum(_logweight, mar), productmeasure((m -> m.base).(mar))) + +@inline function _generic_productmeasure_impl( + mar::AbstractArray{<:WeightedMeasure{StaticFloat64{W},M}}, +) where {W,M} + return weightedmeasure( + static(W) * maybestatic_length(mar), + productmeasure((m -> m.base).(mar)), + ) end -productmeasure(mar::Base.Generator) = ProductMeasure(mar) -productmeasure(mar::AbstractArray) = ProductMeasure(mar) +function _generic_productmeasure_impl(mar::AbstractArray{T}) where {T} + if Base.issingletontype(T) + powermeasure(instance(T), axes(mar)) + elseif T <: AbstractMeasure + ProductMeasure(mar) + else + ProductMeasure(map(asmeasure, mar)) + end +end -# TODO: Make this static when its length is static -@inline function productmeasure( - mar::AbstractArray{WeightedMeasure{StaticFloat64{W},M}}, -) where {W,M} - return weightedmeasure(W * length(mar), productmeasure(map(basemeasure, mar))) +@inline function _generic_productmeasure_impl( + mar::ReadonlyMappedArray{T,N,A,Returns{M}}, +) where {T,N,A,M} + return powermeasure(mar.f.value, axes(mar.data)) end -productmeasure(nt::NamedTuple) = ProductMeasure(nt) -productmeasure(tup::Tuple) = ProductMeasure(tup) +@inline _generic_productmeasure_impl(mar::Base.Generator) = ProductMeasure(mar) -productmeasure(f, param_maps, pars) = ProductMeasure(kernel(f, param_maps), pars) +############################################################################### +# PushforwardMeasure -function productmeasure(k::ParameterizedTransitionKernel, pars) - productmeasure(k.suff, k.param_maps, pars) -end +# The pushforward of a point mass is a point mass. Note that no density +# volume correction applies, Dirac measures are density-defined relative +# to counting measure: +_pushfwd_impl(f, μ::Dirac, ::PushFwdStyle) = Dirac(f(μ.x)) -function productmeasure(f::Returns{W}, ::typeof(identity), pars) where {W<:WeightedMeasure} - ℓ = _logweight(f.value) - base = basemeasure(f.value) - newbase = productmeasure(Returns(base), identity, pars) - weightedmeasure(length(pars) * ℓ, newbase) +# Pushforward and weighting commute: +function _pushfwd_impl(f, μ::WeightedMeasure, style::PushFwdStyle) + weightedmeasure(μ.logweight, _pushfwd_impl(f, μ.base, style)) end ############################################################################### # RestrictedMeasure export restrict -restrict(f, b) = RestrictedMeasure(f, b) +@inline restrict(f) = Base.Fix1(restrict, f) + +restrict(f, μ) = RestrictedMeasure(f, asmeasure(μ)) + +# Nested restrictions fuse into a single predicate: +restrict(f, μ::RestrictedMeasure) = + RestrictedMeasure(x -> μ.predicate(x) && f(x), μ.base) ############################################################################### # SuperpositionMeasure -superpose(a::AbstractArray) = SuperpositionMeasure(a) +""" + superpose(μs...) + superpose(μs) -superpose(t::Tuple) = SuperpositionMeasure(t) -superpose(nt::NamedTuple) = SuperpositionMeasure(nt) +Constructs a superposition of measures, given either as separate arguments or +as a collection (array, tuple or named tuple) of measures. -function superpose(μ::T, ν::T) where {T<:AbstractMeasure} - if μ == ν - return weightedmeasure(static(float(logtwo)), μ) - else - return superpose((μ, ν)) - end -end - -function superpose(μ::AbstractMeasure, μs...) - if all(==(μ), μs) - return weightedmeasure(log(length(μs) + 1), μ) - else - return superpose((μ, μs...)) - end -end +The vararg form simplifies algebraically: equal measures combine into weighted +measures (`superpose(μ, μ) == weightedmeasure(log(2), μ)`), weighted measures +with equal bases add their weights, and superpositions merge their components. +To keep `superpose` type stable, such simplifications only happen when +equality of the measures involved can be decided from their types alone. +Collections are wrapped as-is, apart from cost-free structural simplifications. +""" +function superpose end +export superpose -add_measures(μs::AbstractVector, νs) = push!(μs, νs...) -add_measures(μs::Tuple, νs) = (μs..., νs...) +superpose(μ::AbstractMeasure) = μ -function superpose(μ::SuperpositionMeasure, μs...) - SuperpositionMeasure(add_measures(μ.components, μs)) +function superpose(μ::AbstractMeasure, ν::AbstractMeasure, more::AbstractMeasure...) + superpose(_superpose_two(μ, ν), more...) end -superpose(μ::SuperpositionMeasure) = μ - -############################################################################### -# WeightedMeasure - -function weightedmeasure(ℓ::R, b::M) where {R,M} - WeightedMeasure{R,M}(ℓ, b) +function superpose(a::AbstractArray{T}) where {T} + if Base.issingletontype(T) + weightedmeasure(log(length(a)), asmeasure(instance(T))) + else + SuperpositionMeasure(a) + end end -function weightedmeasure(ℓ, b::WeightedMeasure) - weightedmeasure(ℓ + _logweight(b), b.base) -end +superpose(a::FillArrays.Fill) = weightedmeasure(log(length(a)), asmeasure(_fill_value(a))) -############################################################################### -# TransitionKernel +superpose(t::Tuple) = SuperpositionMeasure(t) +superpose(nt::NamedTuple) = SuperpositionMeasure(nt) -# kernel(Normal(μ=2)) -function kernel(μ::M) where {M<:ParameterizedMeasure} - kernel(M) -end +# Measure equality can typically only be established at runtime, but measure +# construction must be type stable, so simplifications may only depend on +# measure equality that is decidable from the measure types alone: +@inline _static_isequal(::T, ::T) where {T} = static(Base.issingletontype(T)) +@inline _static_isequal(::Any, ::Any) = static(false) -function kernel(d::PowerMeasure) - Base.Fix2(powermeasure, d.axes) ∘ kernel(d.parent) +function _superpose_two(μ::AbstractMeasure, ν::AbstractMeasure) + if _static_isequal(μ, ν) isa True + weightedmeasure(static(float(logtwo)), μ) + else + SuperpositionMeasure((μ, ν)) + end end -function kernel(f) - T = Core.Compiler.return_type(f, Tuple{Any}) - _kernel(f, T) +function _superpose_two(μ::WeightedMeasure, ν::WeightedMeasure) + if _static_isequal(μ.base, ν.base) isa True + weightedmeasure(logaddexp(asnonstatic(μ.logweight), asnonstatic(ν.logweight)), μ.base) + else + SuperpositionMeasure((μ, ν)) + end end -function _kernel(f, ::Type{T}) where {T} - GenericTransitionKernel(f) +function _superpose_two(μ::WeightedMeasure, ν::AbstractMeasure) + if _static_isequal(μ.base, ν) isa True + weightedmeasure(log1pexp(asnonstatic(μ.logweight)), μ.base) + else + SuperpositionMeasure((μ, ν)) + end end -function _kernel(f, ::Type{P}) where {N,P<:ParameterizedMeasure{N}} - k = length(N) - C = constructorof(P) - maps = ntuple(Val(k)) do i - x -> @inbounds x[i] +function _superpose_two(μ::AbstractMeasure, ν::WeightedMeasure) + if _static_isequal(μ, ν.base) isa True + weightedmeasure(log1pexp(asnonstatic(ν.logweight)), ν.base) + else + SuperpositionMeasure((μ, ν)) end - - kernel(params ∘ f, C, NamedTuple{N}(maps)) end -kernel(f::F, ::Type{M}; kwargs...) where {F<:Function,M} = kernel(f, M, NamedTuple(kwargs)) +_superpose_two(μ::SuperpositionMeasure, ν::SuperpositionMeasure) = + SuperpositionMeasure(_cat_measures(μ.components, ν.components)) +_superpose_two(μ::SuperpositionMeasure, ν::AbstractMeasure) = + SuperpositionMeasure(_cat_measures(μ.components, (ν,))) +_superpose_two(μ::AbstractMeasure, ν::SuperpositionMeasure) = + SuperpositionMeasure(_cat_measures((μ,), ν.components)) +_superpose_two(μ::SuperpositionMeasure, ν::WeightedMeasure) = + SuperpositionMeasure(_cat_measures(μ.components, (ν,))) +_superpose_two(μ::WeightedMeasure, ν::SuperpositionMeasure) = + SuperpositionMeasure(_cat_measures((μ,), ν.components)) -function kernel(f::F, ::Type{M}, nt::NamedTuple) where {F<:Function,M} - ParameterizedTransitionKernel(M, f, nt) -end - -function kernel(f::F, ::Type{M}, ::NamedTuple{()}) where {F<:Function,M} - T = Core.Compiler.return_type(f, Tuple{Any}) - _kernel(f, M, T) -end +############################################################################### +# WeightedMeasure -kernel(::Type{P}, nt::NamedTuple) where {P<:ParameterizedMeasure} = kernel(identity, P, nt) +""" + weightedmeasure(logweight::Real, μ) -kernel(::Type{T}; kwargs...) where {T} = kernel(T, NamedTuple(kwargs)) +Constructs a measure that behaves like the measure `μ`, but with its density +scaled by `exp(logweight)`. Weights of nested weighted measures combine +additively. +""" +function weightedmeasure end +export weightedmeasure -function kernel(::Type{M}, ::NamedTuple{()}) where {M} - C = constructorof(M) - TypedTransitionKernel(C, identity) +function weightedmeasure(ℓ::R, b::M) where {R,M} + WeightedMeasure{R,M}(ℓ, b) end -function _kernel(f::F, ::Type{M}, ::Type{NT}) where {M,F,N,NT<:NamedTuple{N}} - k = length(N) - maps = ntuple(Val(k)) do i - x -> @inbounds x[i] - end - - ParameterizedTransitionKernel(M, values ∘ f, NamedTuple{N}(maps)) +function weightedmeasure(ℓ, b::WeightedMeasure) + weightedmeasure(ℓ + _logweight(b), b.base) end -kernel(f::F; kwargs...) where {F<:Function} = kernel(f, NamedTuple(kwargs)) - -function kernel(f::F, nt::NamedTuple{()}) where {F<:Function} - T = Core.Compiler.return_type(f, Tuple{Any}) - _kernel(f, T) -end diff --git a/src/combinators/superpose.jl b/src/combinators/superpose.jl index 099ee806..d9b53bd5 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_rel_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_rel_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_rel_def(s::SuperpositionMeasure, β::SuperpositionMeasure, x) = + _superpos_logdensity_rel(s, β, x) + +@inline logdensity_rel_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/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl index c9db7a6b..706784e6 100644 --- a/src/combinators/transformedmeasure.jl +++ b/src/combinators/transformedmeasure.jl @@ -8,6 +8,7 @@ pushforward. Either [`AdaptRootMeasure()`](@ref) or abstract type PushFwdStyle end export PushFwdStyle +# Backward compatibility with user code, do not use in MeasureBase itself: const TransformVolCorr = PushFwdStyle """ @@ -22,9 +23,10 @@ Density calculations for pushforward measures constructed with transform (typically via the log-abs-det-Jacobian of the transform) into account. """ -struct AdaptRootMeasure <: TransformVolCorr end +struct AdaptRootMeasure <: PushFwdStyle end export AdaptRootMeasure +# Backward compatibility with user code, do not use in MeasureBase itself: const WithVolCorr = AdaptRootMeasure """ @@ -37,9 +39,10 @@ Density calculations for pushforward measures constructed with `PushfwdRootMeasure()` will ignore the volume element of the variate transform. """ -struct PushfwdRootMeasure <: TransformVolCorr end +struct PushfwdRootMeasure <: PushFwdStyle end export PushfwdRootMeasure +# Backward compatibility with user code, do not use in MeasureBase itself: const NoVolCorr = PushfwdRootMeasure abstract type AbstractTransformedMeasure <: AbstractMeasure end @@ -115,24 +118,21 @@ end # end # TODO: Would profit from custom pullback: -function _combine_logd_with_ladj(logd_orig::Real, ladj::Real) +function _combine_logd_with_ladj(logd_orig::Number, ladj::Number) logd_result = logd_orig + ladj R = typeof(logd_result) - if isnan(logd_result) && isneginf(logd_orig) && isposinf(ladj) - # Zero μ wins against infinite volume: - R(-Inf)::R - elseif isfinite(logd_orig) && isneginf(ladj) - # Maybe also for isneginf(logd_orig) && isfinite(ladj) ? - # Return constant -Inf to prevent problems with ForwardDiff: - #R(-Inf) - near_neg_inf(R)::R # Avoids AdvancedHMC warnings - else - logd_result::R - end + # Zero μ wins against infinite volume: + zero_wins = isnan(logd_result) & isneginf(logd_orig) & isposinf(ladj) + # Maybe also for isneginf(logd_orig) && isfinite(ladj) ? + # Return near_neg_inf instead of constant -Inf to prevent problems + # with ForwardDiff and to avoid AdvancedHMC warnings: + fades_out = isfinite(logd_orig) & isneginf(ladj) + + ifelse(zero_wins, R(-Inf), ifelse(fades_out, near_neg_inf(R), logd_result))::R end -function logdensityof( +function logdensityof_impl( @nospecialize(μ::_NonBijectivePusfwdMeasure{M,<:PushfwdRootMeasure}), @nospecialize(v::Any) ) where {M} @@ -143,7 +143,7 @@ function logdensityof( ) end -function logdensityof( +function logdensityof_impl( @nospecialize(μ::_NonBijectivePusfwdMeasure{M,<:AdaptRootMeasure}), @nospecialize(v::Any) ) where {M} @@ -154,15 +154,15 @@ function logdensityof( ) end -for func in [:logdensityof, :logdensity_def] - @eval function $func(ν::PushforwardMeasure{F,I,M,<:AdaptRootMeasure}, y) where {F,I,M} +for (head, func) in [(:logdensityof_impl, :logdensityof), (:logdensity_def, :logdensity_def)] + @eval function $head(ν::PushforwardMeasure{F,I,M,<:AdaptRootMeasure}, y) where {F,I,M} f_inv = unwrap(ν.finv) x, inv_ladj = with_logabsdet_jacobian(f_inv, y) logd_orig = $func(ν.origin, x) return _combine_logd_with_ladj(logd_orig, inv_ladj) end - @eval function $func(ν::PushforwardMeasure{F,I,M,<:PushfwdRootMeasure}, y) where {F,I,M} + @eval function $head(ν::PushforwardMeasure{F,I,M,<:PushfwdRootMeasure}, y) where {F,I,M} f_inv = unwrap(ν.finv) x = f_inv(y) logd_orig = $func(ν.origin, x) @@ -170,7 +170,9 @@ for func in [:logdensityof, :logdensity_def] end end -insupport(m::PushforwardMeasure, x) = insupport(transport_origin(m), to_origin(m, x)) +# Checking insupport via the origin would require a potentially costly +# transformation of x: +insupport(m::PushforwardMeasure, x) = NoFastInsupport{typeof(m)}() function testvalue(::Type{T}, ν::PushforwardMeasure) where {T} ν.f(testvalue(T, parent(ν))) @@ -193,6 +195,9 @@ _pushfwd_dof(::Type{MU}, ::Type{<:Tuple{Any,Real}}, dof) where {MU} = dof @inline getdof(ν::MU) where {MU<:PushforwardMeasure} = getdof(ν.origin) @inline getdof(m::_NonBijectivePusfwdMeasure) = MeasureBase.NoDOF{typeof(m)}() +@inline fast_dof(ν::PushforwardMeasure) = fast_dof(ν.origin) +@inline fast_dof(m::_NonBijectivePusfwdMeasure) = MeasureBase.NoDOF{typeof(m)}() + # Bypass `checked_arg`, would require potentially costly transformation: @inline checked_arg(::PushforwardMeasure, x) = x @@ -222,9 +227,12 @@ To manually specify an inverse, call function pushfwd end export pushfwd +@inline pushfwd(f) = Base.Fix1(pushfwd, f) @inline pushfwd(f, μ) = _pushfwd_impl(f, μ, AdaptRootMeasure()) -@inline pushfwd(f, μ, style::AdaptRootMeasure) = _pushfwd_impl(f, μ, style) -@inline pushfwd(f, μ, style::PushfwdRootMeasure) = _pushfwd_impl(f, μ, style) +@inline pushfwd(f, μ, style::PushFwdStyle) = _pushfwd_impl(f, μ, style) + +@inline pushfwd(::typeof(identity), μ) = μ +@inline pushfwd(::typeof(identity), μ, ::PushFwdStyle) = μ _pushfwd_impl(f, μ, style) = PushforwardMeasure(f, inverse(f), μ, style) @@ -239,8 +247,8 @@ function _pushfwd_impl( PushforwardMeasure(new_f, new_f_inv, orig_μ, style) end -_pushfwd_impl(::typeof(identity), μ, ::AdaptRootMeasure) = μ -_pushfwd_impl(::typeof(identity), μ, ::PushfwdRootMeasure) = μ +# Simplifications for Dirac and WeightedMeasure origins are defined in +# smart-constructors.jl. ############################################################################### # pullback @@ -263,9 +271,9 @@ To manually specify an inverse, call function pullbck end export pullbck +@inline pullbck(f) = Base.Fix1(pullbck, f) @inline pullbck(f, μ) = _pullback_impl(f, μ, AdaptRootMeasure()) -@inline pullbck(f, μ, style::AdaptRootMeasure) = _pullback_impl(f, μ, style) -@inline pullbck(f, μ, style::PushfwdRootMeasure) = _pullback_impl(f, μ, style) +@inline pullbck(f, μ, style::PushFwdStyle) = _pullback_impl(f, μ, style) function _pullback_impl(f, μ, style = AdaptRootMeasure()) pushfwd(inverse(f), μ, style) diff --git a/src/combinators/weighted.jl b/src/combinators/weighted.jl index 124662b6..033e334f 100644 --- a/src/combinators/weighted.jl +++ b/src/combinators/weighted.jl @@ -16,6 +16,12 @@ _logweight(::AbstractMeasure) = 0 d.logweight end +# The weight-shifted density of a support-safe base density is support-safe, +# no explicit support check required: +@inline function logdensityof_impl(d::AbstractWeightedMeasure, x) + d.logweight + logdensityof_impl(basemeasure(d), x) +end + function Base.rand(rng::AbstractRNG, ::Type{T}, μ::AbstractWeightedMeasure) where {T} rand(rng, T, basemeasure(μ)) end @@ -44,7 +50,7 @@ function Base.:*(k::T, m::AbstractMeasure) where {T<:Number} return weightedmeasure(logk, m) end -Base.:*(m::AbstractMeasure, k::Real) = k * m +Base.:*(m::AbstractMeasure, k::Number) = k * m gentype(μ::WeightedMeasure) = gentype(μ.base) diff --git a/src/density-batched.jl b/src/density-batched.jl new file mode 100644 index 00000000..a29effd4 --- /dev/null +++ b/src/density-batched.jl @@ -0,0 +1,199 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +export logdensities + +""" + logdensities(μ::AbstractMeasure, X::AbstractArray) + +Compute the log-density of `μ` at each point in `X`. + +Returns an array of the shape of `X`, semantically equivalent to +`logdensityof.(Ref(μ), X)`. The computation may be fused across points, +though: power measures with flat variate storage (e.g. based on +`ArraysOfArrays.ArrayOfSimilarArrays`) evaluate as a single flat broadcast +plus a segmented reduction over the underlying flat data, compatible with +GPU-backed storage. + +Measure types should specialize [`MeasureBase.batched_logdensityof_impl`](@ref) +instead of `logdensities` itself. +""" +function logdensities end + +@inline function logdensities(μ::AbstractMeasure, X::AbstractArray) + _logdensities(logdensityof_impl, μ, X) +end + +""" + MeasureBase.batched_logdensityof_impl(μ::AbstractMeasure, X::AbstractArray) + +Implements [`logdensities(μ, X)`](@ref logdensities) for arrays `X` of +plain `μ`-variates. Power measures never reach `batched_logdensityof_impl`, their +power structure is processed generically beforehand. + +Measure types that support fused multi-point evaluation should specialize +`batched_logdensityof_impl`. Implementations must preserve the shape of `X` and +must handle points outside the support of `μ` (the result must be `-Inf` +at such points). + +The default implementation broadcasts the log-density over `X` for +measures with scalar variates and falls back to a `map` over `X` +otherwise. +""" +function batched_logdensityof_impl end + +function batched_logdensityof_impl(μ::AbstractMeasure, X::AbstractArray) + _logdensities_generic(logdensityof_impl, μ, X) +end + +# Batched density machinery, parameterized over the point-level density +# function `f` (`logdensityof_impl` or `logdensity_def`). +# +# `_logdensities(f, μ, X, powers...)` treats each element of `X` as a +# variate of `(μ^pN)^…^p1` for `powers = (p1, …, pN)` (power axes ordered +# outermost first, i.e. in the order in which they are encountered when +# descending into a variate) and returns the density sum for each element, +# preserving the shape of `X`. Power measures are unwrapped into the power +# axes arguments before any other dispatch happens, so implementation +# methods only ever dispatch on plain measure types. + +@inline function _logdensities(f::F, μ, X::AbstractArray, powers::Vararg{Any,N}) where {F,N} + _logdensities_stripped(f, μ, X, powers...) +end + +@inline function _logdensities( + f::F, + μ::PowerMeasure, + X::AbstractArray, + powers::Vararg{Any,N}, +) where {F,N} + _logdensities(f, pwr_base(μ), X, powers..., pwr_axes(μ)) +end + +@inline function _logdensities_stripped(f::F, μ, X::AbstractArray) where {F} + _batched_logdensityof_impl(f, μ, X) +end + +function _logdensities_stripped( + f::F, + μ, + X::AbstractArray, + p1, + powers::Vararg{Any,N}, +) where {F,N} + map(x -> _powered_ld(f, μ, x, p1, powers...), X) +end + +function _logdensities_stripped( + f::F, + μ, + X::ArrayOfSimilarArrays{<:Number}, + p1, + powers::Vararg{Any,N}, +) where {F,N} + _logdensities_fused(f, μ, X, mspace_elsize(μ), p1, powers...) +end + +# Absolute densities go through the `batched_logdensityof_impl` extension point: +@inline function _batched_logdensityof_impl(::typeof(logdensityof_impl), μ, X::AbstractArray) + batched_logdensityof_impl(μ, X) +end + +@inline function _batched_logdensityof_impl(f::F, μ, X::AbstractArray) where {F} + _logdensities_generic(f, μ, X) +end + +@inline function _logdensities_generic(f::F, μ, X::AbstractArray) where {F} + _logdensities_byelsize(f, μ, X, mspace_elsize(μ)) +end + +# Scalar variates evaluate as a single flat broadcast: +@inline function _logdensities_byelsize( + f::F, + μ, + X::AbstractArray{<:Number}, + ::Tuple{}, +) where {F} + broadcast(Base.Fix1(f, μ), X) +end + +@inline function _logdensities_byelsize(f::F, μ, X::AbstractArray, ::Any) where {F} + map(Base.Fix1(f, μ), X) +end + +# Scalar-variate measure with flat variate storage: evaluate as a single +# flat broadcast followed by a segmented reduction over the per-point +# power structure: +function _logdensities_fused( + f::F, + μ, + X::ArrayOfSimilarArrays{<:Number,M}, + ::Tuple{}, + powers::Vararg{Any,N}, +) where {F,M,N} + sz_inner = _flat_powers_size(powers...) + if length(sz_inner) == M + X_flat = flatview(X) + if ntuple(i -> size(X_flat, i), Val(M)) != sz_inner + throw(ArgumentError("Size of variates doesn't match size of power measure")) + end + ld_flat = broadcast(Base.Fix1(f, μ), X_flat) + reshape(sum(ld_flat, dims = ntuple(identity, Val(M))), size(X)) + else + map(x -> _powered_ld(f, μ, x, powers...), X) + end +end + +function _logdensities_fused( + f::F, + μ, + X::AbstractArray, + ::Any, + p1, + powers::Vararg{Any,N}, +) where {F,N} + map(x -> _powered_ld(f, μ, x, p1, powers...), X) +end + +# The flat size of a variate of `(μ^pN)^…^p1` for a scalar-variate `μ`, +# innermost power axes vary fastest: +@inline _flat_powers_size() = () +@inline function _flat_powers_size(p1, powers::Vararg{Any,N}) where {N} + (_flat_powers_size(powers...)..., axes2size(p1)...) +end + +# Scalar counterpart of `_logdensities`: log-density of `(μ^pN)^…^p1` at a +# single variate `x`. +@inline _powered_ld(f::F, μ, x) where {F} = f(μ, x) + +@inline function _powered_ld( + f::F, + μ::PowerMeasure, + x, + p1, + powers::Vararg{Any,N}, +) where {F,N} + _powered_ld(f, pwr_base(μ), x, p1, powers..., pwr_axes(μ)) +end + +@inline function _powered_ld(f::F, μ, x, p1, powers::Vararg{Any,N}) where {F,N} + if axes2size(p1) != maybestatic_size(x) + throw(ArgumentError("Size of variate doesn't match size of power measure")) + end + R = _powered_ld_type(f, μ, x, powers...) + if isempty(x) + zero(R)::R + else + # Conversion needed since summation can turn static into dynamic values: + convert(R, _powered_ld_sum(f, μ, x, powers...))::R + end +end + +@inline _powered_ld_sum(f::F, μ, x) where {F} = sum(Base.Fix1(f, μ), x) + +@inline function _powered_ld_sum(f::F, μ, x, p1, powers::Vararg{Any,N}) where {F,N} + sum(_logdensities(f, μ, x, p1, powers...)) +end + +@inline function _powered_ld_type(f::F, ::MU, x, powers::Vararg{Any,N}) where {F,MU,N} + Core.Compiler.return_type(_powered_ld, Tuple{F,MU,eltype(x),map(typeof, powers)...}) +end diff --git a/src/density-core.jl b/src/density-core.jl index 6ac3d01e..9a6b0649 100644 --- a/src/density-core.jl +++ b/src/density-core.jl @@ -10,25 +10,41 @@ export density_rel export density_def """ - logdensityof(m::AbstractMeasure, x) + logdensityof(m::AbstractMeasure, x) Compute the log-density of the measure `m` at `x`. Density is always relative, but `DensityInterface.jl` does not account for this. For compatibility with this, `logdensityof` for a measure is always implicitly relative to -[`rootmeasure(x)`](@ref rootmeasure). +[`rootmeasure(x)`](@ref rootmeasure). -`logdensityof` works by first computing `insupport(m, x)`. If this is true, then -`unsafe_logdensityof` is called. If `insupport(m, x)` is known to be `true`, it -can be a little faster to directly call `unsafe_logdensityof(m, x)`. +`logdensityof(m, x)` is implemented via +[`MeasureBase.logdensityof_impl`](@ref), measure types should specialize +`logdensityof_impl` instead of `logdensityof` itself. To compute log-density relative to `basemeasure(m)` or *define* a log-density (relative to `basemeasure(m)` or another measure given explicitly), see -`logdensity_def`. +`logdensity_def`. To compute a log-density relative to a specific base-measure, see -`logdensity_rel`. +`logdensity_rel`. """ -@inline function logdensityof(μ::AbstractMeasure, x) +@inline logdensityof(μ::AbstractMeasure, x) = logdensityof_impl(μ, x) + +""" + MeasureBase.logdensityof_impl(μ::AbstractMeasure, x) + +Implements [`logdensityof(μ, x)`](@ref logdensityof). + +Measure types should specialize `logdensityof_impl` instead of +`logdensityof` itself. Implementations must return the log-density of `μ` +at `x` relative to [`rootmeasure(μ)`](@ref) and must handle `x` outside of +the support of `μ` (the result must be `-Inf` then). + +The default implementation checks `insupport(μ, x)` (unless the result is +a [`MeasureBase.NoFastInsupport`](@ref)) and computes the density via +[`unsafe_logdensityof`](@ref). +""" +@inline function logdensityof_impl(μ::AbstractMeasure, x) result = dynamic(unsafe_logdensityof(μ, x)) _checksupport(insupport(μ, x), result) end @@ -38,6 +54,77 @@ end end _checksupport(cond, result) = ifelse(cond == true, result, oftype(result, -Inf)) +@inline _checksupport(::NoFastInsupport, result) = result + +""" + MeasureBase.logdensityof_with_rest(μ::AbstractMeasure, x) + +Compute the log-density of `μ` at the beginning of `x`, a flat stream of +variate content that may extend beyond the variate of `μ`. + +`x` must either be a vector that starts with the (flattened) variate of `μ` +(for measures combined via `vcat`) or a `NamedTuple` whose first properties +constitute the variate of `μ` (for measures combined via `merge`). + +Returns a tuple `(ℓ, x_μ, x_rest)` of the log-density `ℓ`, the variate +`x_μ` of `μ` consumed from the stream, and the unconsumed rest of the +stream. + +Measure types whose variate size depends on measure values, like +[`mbind`](@ref) results, implement density calculation via +`logdensityof_with_rest` instead of +[`logdensityof_impl`](@ref MeasureBase.logdensityof_impl). + +The default implementation determines the size resp. the property names of +the variate via [`some_mspace_elsize`](@ref MeasureBase.some_mspace_elsize) +resp. `testvalue` and delegates to `logdensityof_impl`. +""" +function logdensityof_with_rest end + +function logdensityof_with_rest(μ::AbstractMeasure, x::AbstractVector) + a, x_rest = _consume_from_stream(x, some_mspace_elsize(μ)) + return logdensityof_impl(μ, a), a, x_rest +end + +function logdensityof_with_rest(μ::AbstractMeasure, x::NamedTuple) + a, x_rest = _split_after(x, Val(_mspace_names(μ))) + return logdensityof_impl(μ, a), a, x_rest +end + +_mspace_names(μ::AbstractMeasure) = keys(testvalue(μ)) + + +""" + localmeasure(m::AbstractMeasure, x)::AbstractMeasure + +Return a measure that behaves like `m` in the infinitesimal neighborhood +of `x` in respect to density calculation. + +Note that the resulting measure may not be well defined outside of the +infinitesimal neighborhood of `x`. + +For most measure types simply returns `m` itself. [`mbind`](@ref), +for example, generates measures for which `localmeasure(m, x)` depends +on `x`. +""" +localmeasure(m::AbstractMeasure, x) = m +export localmeasure + + +""" + MeasureBase.transportmeasure(m::AbstractMeasure, x)::AbstractMeasure + +Return a measure that behaves like `m` in the infinitesimal neighborhood +of `x` in respect to both transport and density calculation. + +Note that the resulting measure may not be well defined outside of the +infinitesimal neighborhood of `x`. + +For most measure types simply returns `m` itself. [`mbind`](@ref), +for example, generates measures for which `transportmeasure(m, x)` depends +on `x`. +""" +transportmeasure(m::AbstractMeasure, x) = m export unsafe_logdensityof @@ -50,11 +137,17 @@ This is "unsafe" because it does not check `insupport(m, x)`. See also `logdensityof`. """ -@inline function unsafe_logdensityof(μ::M, x) where {M} +@inline function unsafe_logdensityof(μ::AbstractMeasure, x) + μ_local = localmeasure(μ, x) + # Extra dispatch boundary to reduce number of required specializations of implementation: + return _unsafe_logdensityof_local(μ_local, x) +end + +@inline function _unsafe_logdensityof_local(μ::M, x) where {M} ℓ_0 = logdensity_def(μ, x) b_0 = μ Base.Cartesian.@nexprs 10 i -> begin # 10 is just some "big enough" number - b_{i} = basemeasure(b_{i - 1}, x) + b_{i} = basemeasure(b_{i - 1}) # The below makes the evaluated code shorter, but screws up Zygote # if b_{i} isa typeof(b_{i - 1}) @@ -76,20 +169,56 @@ known to be in the support of both, it can be more efficient to call `unsafe_logdensity_rel`. """ @inline function logdensity_rel(μ::M, ν::N, x::X) where {M,N,X} + inμ = insupport(μ, x) + inν = insupport(ν, x) + return _logdensity_rel_impl(μ, ν, x, inμ, inν) +end + +@inline function _logdensity_rel_impl(μ::M, ν::N, x::X, inμ::Bool, inν::Bool) where {M,N,X} T = unstatic( promote_type( return_type(logdensity_def, (μ, x)), return_type(logdensity_def, (ν, x)), ), ) - inμ = insupport(μ, x) - inν = insupport(ν, x) istrue(inμ) || return convert(T, ifelse(inν, -Inf, NaN)) istrue(inν) || return convert(T, Inf) return unsafe_logdensity_rel(μ, ν, x) end +@inline function _logdensity_rel_impl( + μ::M, + ν::N, + x::X, + @nospecialize(::NoFastInsupport), + @nospecialize(::NoFastInsupport) +) where {M,N,X} + unsafe_logdensity_rel(μ, ν, x) +end + +@inline function _logdensity_rel_impl( + μ::M, + ν::N, + x::X, + inμ::Bool, + @nospecialize(::NoFastInsupport) +) where {M,N,X} + logd = unsafe_logdensity_rel(μ, ν, x) + return istrue(inμ) ? logd : oftype(logd, -Inf) +end + +@inline function _logdensity_rel_impl( + μ::M, + ν::N, + x::X, + @nospecialize(::NoFastInsupport), + inν::Bool +) where {M,N,X} + logd = unsafe_logdensity_rel(μ, ν, x) + return istrue(inν) ? logd : oftype(logd, +Inf) +end + """ unsafe_logdensity_rel(m1, m2, x) @@ -98,71 +227,121 @@ known to be in the support of both `m1` and `m2`. See also `logdensity_rel`. """ -@inline function unsafe_logdensity_rel(μ::M, ν::N, x::X) where {M,N,X} - if static_hasmethod(logdensity_def, Tuple{M,N,X}) - return logdensity_def(μ, ν, x) - end - μs = basemeasure_sequence(μ) - νs = basemeasure_sequence(ν) - cb = commonbase(μs, νs, X) - # _logdensity_rel(μ, ν) - isnothing(cb) && begin - μ = μs[end] - ν = νs[end] - @warn """ - No common base measure for - $μ - and - $ν - - Returning a relative log-density of NaN. If this is incorrect, add a - three-argument method - logdensity_def($μ, $ν, x) - """ - return NaN - end - return _logdensity_rel(μs, νs, cb, x) +@inline function unsafe_logdensity_rel(μ::AbstractMeasure, ν::AbstractMeasure, x) + μ_local = localmeasure(μ, x) + ν_local = localmeasure(ν, x) + return logdensity_def(μ_local, ν_local, x) end -# Note that this method assumes `μ` and `ν` to have the same type -function logdensity_def(μ::T, ν::T, x) where {T} - if μ === ν - return zero(logdensity_def(μ, x)) - else - α = basemeasure(μ) - β = basemeasure(ν) - return logdensity_def(μ, x) - logdensity_def(ν, x) + logdensity_rel(α, β, x) - end -end +# Indicates that no specialized method is available to compute the +# log-density between a given pair of measures: +struct _NoLogdensityRel end -@generated function _logdensity_rel( - μs::Tμ, - νs::Tν, - ::Tuple{<:StaticInteger{M},<:StaticInteger{N}}, - x::X, -) where {Tμ,Tν,M,N,X} - sμ = schema(Tμ) - sν = schema(Tν) +""" + MeasureBase.logdensity_rel_def(μ, ν, x) - q = quote - $(Expr(:meta, :inline)) - ℓ = logdensity_def(μs[$M], νs[$N], x) - end +Specialization point for the log-density of `μ` relative to `ν` at `x`. - for i in 1:M-1 - push!(q.args, :(Δℓ = logdensity_def(μs[$i], x))) - # push!(q.args, :(println("Adding", Δℓ))) - push!(q.args, :(ℓ += Δℓ)) - end +Measure types may add methods for pairs of measure types whose relative +density can be computed directly. The generic implementation of +[`logdensity_def(μ, ν, x)`](@ref logdensity_def) descends the base measure +chains of both measures in lockstep and uses the first specialized +`logdensity_rel_def` method it encounters along the way. + +Do not call `logdensity_rel_def` directly, call +[`logdensity_rel`](@ref) (or `logdensity_def`) instead. +""" +@inline logdensity_rel_def(μ, ν, x) = _NoLogdensityRel() + +# Generic relative density: descend the base measure chains of both measures +# in lockstep, after equalizing their depths. Since the members of a shared +# chain suffix have the same depth-from-root on both sides, the descent +# terminates at a specialized `logdensity_rel_def` method as soon as one +# becomes applicable (in particular for pairs of identical primitive +# measures), so any shared chain suffix cancels symbolically instead of +# numerically. The descent is fully unrolled at compile time based on the +# static base measure depths, only the base measures actually visited are +# constructed, and whether a specialized method applies at a given level is +# decided purely by dispatch (on the sentinel type `_NoLogdensityRel`). +@inline function logdensity_def(μ, ν, x) + _logdensity_rel_descent(μ, basemeasure_depth(μ), ν, basemeasure_depth(ν), x) +end - for j in 1:N-1 - push!(q.args, :(Δℓ = logdensity_def(νs[$j], x))) - # push!(q.args, :(println("Subtracting", Δℓ))) - push!(q.args, :(ℓ -= Δℓ)) +@generated function _logdensity_rel_descent( + μ, + ::StaticInteger{M}, + ν, + ::StaticInteger{N}, + x, +) where {M,N} + μsym(i) = Symbol(:μ_, i) + νsym(j) = Symbol(:ν_, j) + prog = Expr(:block, Expr(:meta, :inline), :(μ_0 = μ), :(ν_0 = ν)) + terms = Any[] + n_checks = 0 + # Return via a specialized `logdensity_rel_def` method for the current + # measure pair, if available. Whether one is available is decided purely + # by type, so unsuccessful checks are free at run time: + function emit_check!(i, j) + r = Symbol(:r_, n_checks) + n_checks += 1 + push!(prog.args, :($r = logdensity_rel_def($(μsym(i)), $(νsym(j)), x))) + ret = isempty(terms) ? r : :(+($(terms...), $r)) + push!(prog.args, :(if !($r isa _NoLogdensityRel) + return $ret + end)) end + i = j = 0 + emit_check!(i, j) + # Equalize depths, accumulating one-sided density terms: + while M - i > N - j + ℓ = Symbol(:ℓμ_, i) + push!(prog.args, :($ℓ = logdensity_def($(μsym(i)), x))) + push!(prog.args, :($(μsym(i + 1)) = basemeasure($(μsym(i))))) + push!(terms, ℓ) + i += 1 + emit_check!(i, j) + end + while N - j > M - i + ℓ = Symbol(:ℓν_, j) + push!(prog.args, :($ℓ = -logdensity_def($(νsym(j)), x))) + push!(prog.args, :($(νsym(j + 1)) = basemeasure($(νsym(j))))) + push!(terms, ℓ) + j += 1 + emit_check!(i, j) + end + # Lockstep descent at equal depth: + for _ in 1:(M-i) + ℓμ, ℓν = Symbol(:ℓμ_, i), Symbol(:ℓν_, j) + push!(prog.args, :($ℓμ = logdensity_def($(μsym(i)), x))) + push!(prog.args, :($ℓν = -logdensity_def($(νsym(j)), x))) + push!(terms, ℓμ, ℓν) + push!(prog.args, :($(μsym(i + 1)) = basemeasure($(μsym(i))))) + push!(prog.args, :($(νsym(j + 1)) = basemeasure($(νsym(j))))) + i += 1 + j += 1 + emit_check!(i, j) + end + # Both measures are at root level now: + push!( + prog.args, + :(r_root = _root_logdensity_rel($(μsym(i)), $(νsym(j)), x)), + ) + ret = isempty(terms) ? :r_root : :(+($(terms...), r_root)) + push!(prog.args, :(return $ret)) + return prog +end + +# Root measures of the same type are equal almost everywhere for the +# purpose of pointwise relative densities: +_root_logdensity_rel(μ::M, ν::M, x) where {M} = zero(logdensity_def(μ, x)) - push!(q.args, :(return ℓ)) - return q +function _root_logdensity_rel(@nospecialize(μ), @nospecialize(ν), @nospecialize(x)) + throw( + ArgumentError( + "No method available to compute the log-density between measures with root measures of type $(nameof(typeof(μ))) and $(nameof(typeof(ν)))", + ), + ) end @inline density_rel(μ, ν, x) = exp(logdensity_rel(μ, ν, x)) diff --git a/src/density.jl b/src/density.jl index a79021de..67d34464 100644 --- a/src/density.jl +++ b/src/density.jl @@ -20,8 +20,7 @@ For measures `μ` and `ν`, `Density(μ,ν)` represents the _density function_ `dμ/dν`, also called the _Radon-Nikodym derivative_: https://en.wikipedia.org/wiki/Radon%E2%80%93Nikodym_theorem#Radon%E2%80%93Nikodym_derivative -Instead of calling this directly, users should call `density_rel(μ, ν)` or -its abbreviated form, `𝒹(μ,ν)`. +Instead of calling this directly, users should call `density_rel(μ, ν)`. """ struct Density{M,B} <: AbstractDensity μ::M @@ -32,16 +31,6 @@ Base.:∘(::typeof(log), d::Density) = logdensity_rel(d.μ, d.base) Base.log(d::Density) = log ∘ d -export 𝒹 - -""" - 𝒹(μ, base) - -Compute the density (Radon-Nikodym derivative) of μ with respect to `base`. This -is a shorthand form for `density_rel(μ, base)`. -""" -𝒹(μ, base) = density_rel(μ, base) - density_rel(μ, base) = Density(μ, base) (f::Density)(x) = density_rel(f.μ, f.base, x) @@ -61,8 +50,7 @@ For measures `μ` and `ν`, `LogDensity(μ,ν)` represents the _log-density func `log(dμ/dν)`, also called the _Radon-Nikodym derivative_: https://en.wikipedia.org/wiki/Radon%E2%80%93Nikodym_theorem#Radon%E2%80%93Nikodym_derivative -Instead of calling this directly, users should call `logdensity_rel(μ, ν)` or -its abbreviated form, `log𝒹(μ,ν)`. +Instead of calling this directly, users should call `logdensity_rel(μ, ν)`. """ struct LogDensity{M,B} <: AbstractDensity μ::M @@ -73,16 +61,6 @@ Base.:∘(::typeof(exp), d::LogDensity) = density_rel(d.μ, d.base) Base.exp(d::LogDensity) = exp ∘ d -export log𝒹 - -""" - log𝒹(μ, base) - -Compute the log-density (Radon-Nikodym derivative) of μ with respect to `base`. -This is a shorthand form for `logdensity_rel(μ, base)` -""" -log𝒹(μ, base) = logdensity_rel(μ, base) - logdensity_rel(μ, base) = LogDensity(μ, base) (f::LogDensity)(x) = logdensity_rel(f.μ, f.base, x) @@ -98,12 +76,13 @@ DensityInterface.funcdensity(d::LogDensity) = throw(MethodError(funcdensity, (d, base :: B end -A `DensityMeasure` is a measure defined by a density or log-density with respect -to some other "base" measure. +A `DensityMeasure` is a measure defined by a density or log-density with +respect to some other "base" measure. -Users should not call `DensityMeasure` directly, but should instead call `∫(f, -base)` (if `f` is a density function or `DensityInterface.IsDensity` object) or -`∫exp(f, base)` (if `f` is a log-density function). +Users should not instantiate `DensityMeasure` directly, but should instead +call `mintegrate(f, base)` (if `f` is a density function or +`DensityInterface.IsDensity` object) or `mintegrate_exp(f, base)` (if `f` +is a log-density function). """ struct DensityMeasure{F,B} <: AbstractMeasure f::F @@ -120,42 +99,124 @@ end end function Pretty.tile(μ::DensityMeasure{F,B}) where {F,B} - result = Pretty.literal("DensityMeasure ∫(") + result = Pretty.literal("mintegrate(") result *= Pretty.pair_layout(Pretty.tile(μ.f), Pretty.tile(μ.base); sep = ", ") result *= Pretty.literal(")") end -export ∫ """ - ∫(f, base::AbstractMeasure) + MeasureBase.as_integrand(f) + MeasureBase.as_integrand(density) + +Make `f` or `density` (more) suitable as an integrand for +[`mintegrate`](@ref). + +`mintegrate(obj, μ::AbstractMeasure)` automatically calls +`as_integrand(obj)` internally. -Define a new measure in terms of a density `f` over some measure `base`. +If a density is passed, it must implement the DensityInterface API. + +By default just returns `f` resp. `density`, but may be specialized for +functions and densities that can profit from conversion to a form optimized +for use in `mintegrate`. + +See also [`MeasureBase.as_likelihood`](@ref). """ -∫(f, base) = _densitymeasure(f, base, DensityKind(f)) +function as_integrand end + +@inline as_integrand(obj) = _as_integrand_default_impl(obj, DensityKind(obj)) -_densitymeasure(f, base, ::IsDensity) = DensityMeasure(f, base) -function _densitymeasure(f, base, ::HasDensity) - @error "`∫(f, base)` requires `DensityKind(f)` to be `IsDensity()` or `NoDensity()`." +@inline _as_integrand_default_impl(f, ::NoDensity) = funcdensity(f) + +@inline _as_integrand_default_impl(density, ::IsDensity) = density + +function _as_integrand_default_impl(obj, ::HasDensity) + throw( + ArgumentError( + "`MeasureBase.as_integrand(obj)` requires `DensityKind(obj)` to be `IsDensity()` or `NoDensity()`.", + ), + ) end -_densitymeasure(f, base, ::NoDensity) = DensityMeasure(funcdensity(f), base) -export ∫exp +@doc raw""" + mintegrate(f, μ::AbstractMeasure)::AbstractMeasure + mintegrate(density, μ::AbstractMeasure)::AbstractMeasure + +Returns a new measure that represents the indefinite +[integral](https://en.wikipedia.org/wiki/Radon%E2%80%93Nikodym_theorem) +of `f` with respect to `μ`. + +If a density is passed, it must implement the DensityInterface API. + +`ν = mintegrate(f, μ)` generates a measure `ν` that has the mathematical +interpretation + +```math +\nu(A) = \int_A f(a) \, \rm{d}\mu(a) +``` """ - ∫exp(f, base::AbstractMeasure) +function mintegrate end +export mintegrate + +@inline mintegrate(obj, μ::AbstractMeasure) = DensityMeasure(as_integrand(obj), μ) -Define a new measure in terms of a log-density `f` over some measure `base`. + +""" + MeasureBase.as_integrand_exp(log_f) + +Convert the logarithm of an integrand to an integrand. + +See also [`MeasureBase.as_integrand`](@ref). """ -∫exp(f, base) = _logdensitymeasure(f, base, DensityKind(f)) +function as_integrand_exp end + +@inline as_integrand_exp(log_f) = _as_integrand_exp_default_impl(log_f, DensityKind(log_f)) + +@inline _as_integrand_exp_default_impl(log_f, ::NoDensity) = logfuncdensity(log_f) -function _logdensitymeasure(f, base, ::IsDensity) - @error "`∫exp(f, base)` is not valid when `DensityKind(f) == IsDensity()`. Use `∫(f, base)` instead." +function _as_integrand_exp_default_impl(log_f, ::IsDensity) + throw( + ArgumentError( + "`as_integrand_exp(log_f)` is not valid when `DensityKind(log_f) == IsDensity()`. Use `as_integrand(log_f)` instead.", + ), + ) end -function _logdensitymeasure(f, base, ::HasDensity) - @error "`∫exp(f, base)` is not valid when `DensityKind(f) == HasDensity()`." + +function _as_integrand_exp_default_impl(log_f, ::HasDensity) + throw( + ArgumentError( + "`as_integrand_exp(log_f)` is not valid when `DensityKind(log_f) == HasDensity()`.", + ), + ) end -_logdensitymeasure(f, base, ::NoDensity) = DensityMeasure(logfuncdensity(f), base) + + +@doc raw""" + mintegrate_exp(log_f, μ::AbstractMeasure) + +Given a function `log_f` that semantically represents the log of a function +`f`, `mintegrate_exp` returns a new measure that represents the indefinite +[integral](https://en.wikipedia.org/wiki/Radon%E2%80%93Nikodym_theorem) +of `f` with respect to `μ`. + +`ν = mintegrate_exp(log_f, μ)` generates a measure `ν` that has the +mathematical interpretation + +```math +\nu(A) = \int_A e^{log(f(a))} \, \rm{d}\mu(a) = \int_A f(a) \, \rm{d}\mu(a) +``` + +Note that `exp(log_f(...))` is usually not run explicitly, calculations that +involve the resulting measure are typically performed in log-space, +internally. +""" +function mintegrate_exp end +export mintegrate_exp + +mintegrate_exp(log_f, μ::AbstractMeasure) = DensityMeasure(as_integrand_exp(log_f), μ) + basemeasure(μ::DensityMeasure) = μ.base @@ -163,7 +224,7 @@ logdensity_def(μ::DensityMeasure, x) = logdensityof(μ.f, x) density_def(μ::DensityMeasure, x) = densityof(μ.f, x) -function logdensityof(μ::DensityMeasure, x::Any) +function logdensityof_impl(μ::DensityMeasure, x::Any) integrand, μ_base = μ.f, μ.base base_logval = logdensityof(μ_base, x) @@ -180,14 +241,3 @@ function logdensityof(μ::DensityMeasure, x::Any) convert(R, integrand_logval + base_logval)::R end end - -""" - rebase(μ, ν) - -Express `μ` in terms of a density over `ν`. Satisfies -``` -basemeasure(rebase(μ, ν)) == ν -density(rebase(μ, ν)) == 𝒹(μ,ν) -``` -""" -rebase(μ, ν) = ∫(𝒹(μ, ν), ν) diff --git a/src/domains.jl b/src/domains.jl index e03f753c..1f3535ef 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 + +@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 + +maybe_in(@nospecialize(x), ::ImplicitDomain) = true + +function Base.isempty(::ImplicitDomain) + throw(ArgumentError("Can't test if an ImplicitDomain is empty")) +end + +""" + struct MeasureBase.UnknownDomain{T} <: MeasureBase.ValueSet + +Represents the unknown domain of a value of type `T`. + +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 -Base.in(x, b::BoundedReals) = b.lower ≤ x ≤ b.upper +maybe_in(@nospecialize(x), ::UnknownDomain) = true + +Base.isempty(::UnknownDomain) = false + +""" + RealValues() isa MeasureBase.ValueSet + +The real numbers. +""" +struct RealValues <: ValueSet end + +@inline Base.in(x::Real, ::RealValues) = true +@inline Base.in(x, ::RealValues) = isreal(x) + +@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.ℝ") -export ℝ, ℝ₊, 𝕀, ℤ +""" + MeasureBase.IntegerValues() isa MeasureBase.ValueSet +""" +struct IntegerValues <: ValueSet end -const ℝ₊ = BoundedReals(static(0.0), static(Inf)) -const 𝕀 = BoundedReals(static(0.0), static(1.0)) +@inline Base.in(x::Integer, ::IntegerValues) = true +@inline Base.in(x, ::IntegerValues) = isinteger(x) -Base.minimum(b::BoundedReals) = b.lower -Base.maximum(b::BoundedReals) = b.upper +@inline Base.isempty(::IntegerValues) = false -Base.show(io::IO, ::typeof(ℝ₊)) = print(io, "ℝ₊") -Base.show(io::IO, ::typeof(𝕀)) = print(io, "𝕀") +@inline Base.union(s::IntegerValues, ::IntegerValues...) = s -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 +testvalue(::Type{T}, ::IntegerValues) where {T} = zero(T) -abstract type IntegerDomain <: AbstractDomain end +# # 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)) -struct IntegerNumbers <: IntegerDomain end +""" + const ℤ = IntegerValues() -Base.in(x, ::IntegerNumbers) = isinteger(x) +The set of all integers, see [`MeasureBase.IntegerValues`](@ref). +""" +const ℤ = IntegerValues() +export ℤ -const ℤ = IntegerNumbers() +Base.show(io::IO, ::IntegerValues) = print(io, "ℤ") +Base.show(io::IO, ::MIME"text/plain", ::IntegerValues) = print(io, "MeasureBase.ℤ") -Base.show(io::IO, ::typeof(ℤ)) = print(io, "ℤ") +""" + struct MeasureBase.BoundedInts{L,U} <: MeasureBase.ValueSet -Base.minimum(::IntegerNumbers) = static(-Inf) -Base.maximum(::IntegerNumbers) = static(Inf) -struct BoundedInts{L,U} <: IntegerDomain +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,195 @@ 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)...) + +""" + struct MeasureBase.AbstractCartSetProd <: ValueSet -########################################################### -# ZeroSet +Supertype for cartesian products of sets. +""" +abstract type AbstractCartSetProd <: ValueSet end -export ZeroSet +""" + struct CartesianProduct <: AbstractCartSetProd -struct ZeroSet{F,G} <: AbstractDomain - f::F - ∇f::G +A cartesian product over a collection of sets. + +Constructor: + +```julia +prodset = CartesianProduct(sets) +``` + +`sets` may be a `Tuple`, `NamedTuple` or `AbstractArray` of sets/domains. +""" +struct CartesianProduct{S<:Union{Tuple,NamedTuple,AbstractArray}} <: AbstractCartSetProd + _sets::S end -# 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) +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}}}) where {names} = + 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 -########################################################### -# CodimOne +@inline Base.isempty(s::CartesianProduct) = any(isempty, componentsets(s)) -export CodimOne +@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 -abstract type CodimOne <: AbstractDomain end +@inline function Base.union( + s::CartesianProduct{<:NamedTuple{names}}, + others::CartesianProduct{<:NamedTuple{names}}..., +) where {names} + CartesianProduct(map(union, componentsets(s), map(componentsets, others)...)) +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 +function Base.union( + s::CartesianProduct{<:AbstractArray{<:Any,N}}, + others::CartesianProduct{<:AbstractArray{<:Any,N}}..., +) where {N} + CartesianProduct(union.(componentsets(s), map(componentsets, others)...)) +end - # Get the normal vectors - g1 = a.∇f(x) - g2 = b.∇f(x) +""" + struct CartesianPower <: AbstractCartSetProd - # See if one is a multiple of the other - one(T) - Statistics.corm(g1, zero(T), g2, zero(T)) < tol +Represents the n-fold Cartesian product of a set. +""" +struct CartesianPower{S,A} <: AbstractCartSetProd + _base::S + _axes::A end -function zeroset(::CodimOne)::ZeroSet end +@inline setcartpower(s::SetLike, dims) = CartesianPower(s, asaxes(dims)) -########################################################### -# Simplex -export Simplex +@inline pwr_base(s::CartesianPower) = s._base +@inline pwr_axes(s::CartesianPower) = s._axes +@inline pwr_size(s::CartesianPower) = axes2size(pwr_axes(s)) -struct Simplex <: CodimOne end +componentsets(s::CartesianPower) = maybestatic_fill(pwr_base(s), pwr_axes(s)) -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) +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.Fix2(in, pwr_base(s)), x)::Bool end -function Base.in(x::AbstractArray{T}, ::Simplex) where {T} - all(≥(zero(eltype(x))), x) || return false - return x ∈ zeroset(Simplex()) +Base.isempty(s::CartesianPower) = isempty(pwr_base(s)) || size2length(pwr_size(s)) == 0 + +function Base.union(s::CartesianPower, others::CartesianPower...) + axs = pwr_axes(s) + + all(isequal(axs), map(pwr_axes, others)) || throw( + ArgumentError("Cannot create union of CartesianPower sets with different axes."), + ) + + setcartpower(union(pwr_base(s), map(pwr_base, others)...), axs) end -projectto!(x, ::Simplex) = normalize!(x, 1) -########################################################### -# Sphere +""" + struct CombinedSet <: ValueSet -struct Sphere <: CodimOne end +Represents a combination of two sets. + +User code should not create instances of `CombinedSet` directly, but should +call [`combinesets(f_c, α, β)`](@ref) instead. +""" +struct CombinedSet{FC,MA<:SetLike,MB<:SetLike} <: ValueSet + f_c::FC + α::MA + β::MB +end -function zeroset(::Sphere) - f(x::AbstractArray{T}) where {T} = dot(x, x) - one(T) - ∇f(x::AbstractArray{T}) where {T} = x - ZeroSet(f, ∇f) +function Base.in(@nospecialize(x), ::CombinedSet) + throw(ArgumentError("Cannot test if a value lies within a combined set.")) end -function Base.in(x::AbstractArray{T}, ::Sphere) where {T} - return x ∈ zeroset(Sphere()) +maybe_in(@nospecialize(x), ::CombinedSet) = true + +Base.isempty(s::CombinedSet) = isempty(s.α) || isempty(s.β) + +""" + combinesets(f_c, α, β) + +Combine two sets `α` and `β` into the set of all values `f_c(a, b)` with +`a ∈ α` and `b ∈ β`. + +`f_c` must combine values as described in [`mcombine`](@ref). Uses set +representations more specific than [`MeasureBase.CombinedSet`](@ref) where +possible. +""" +function combinesets end +export combinesets + +@inline combinesets(f_c, α::SetLike, β::SetLike) = _generic_combinesets(f_c, α, β) + +# Combining the implicit domains of two measures yields the implicit domain +# of the combined measure: +@inline combinesets(f_c, α::ImplicitDomain, β::ImplicitDomain) = + ImplicitDomain(mcombine(f_c, α.m, β.m)) + +@inline _generic_combinesets(::typeof(firstarg), α::SetLike, β::SetLike) = α +@inline _generic_combinesets(::typeof(secondarg), α::SetLike, β::SetLike) = β +@inline _generic_combinesets(::typeof(tuple), α::SetLike, β::SetLike) = + setcartprod((α, β)) +@inline _generic_combinesets(f_c::typeof(vcat), α::SetLike, β::SetLike) = + _combinesets_cat(f_c, α, β) +@inline _generic_combinesets(f_c::typeof(merge), α::SetLike, β::SetLike) = + _combinesets_cat(f_c, α, β) +@inline _generic_combinesets(f_c, α::SetLike, β::SetLike) = CombinedSet(f_c, α, β) + +_combinesets_cat( + ::typeof(vcat), + α::CartesianProduct{<:AbstractVector}, + β::CartesianProduct{<:AbstractVector}, +) = setcartprod(vcat(componentsets(α), componentsets(β))) + +_combinesets_cat( + ::typeof(merge), + α::CartesianProduct{<:NamedTuple}, + β::CartesianProduct{<:NamedTuple}, +) = setcartprod(merge(componentsets(α), componentsets(β))) + +# Concatenating one-dimensional powers of equal base sets yields a longer +# power. Set equality can typically only be established at runtime, so this +# simplification only happens when it is decidable from the set types alone: +function _combinesets_cat( + ::typeof(vcat), + α::CartesianPower{<:Any,<:Tuple{Any}}, + β::CartesianPower{<:Any,<:Tuple{Any}}, +) + if _static_isequal(pwr_base(α), pwr_base(β)) isa True + n = size2length(pwr_size(α)) + size2length(pwr_size(β)) + setcartpower(pwr_base(α), (n,)) + else + CombinedSet(vcat, α, β) + end end -projectto!(x, ::Sphere) = normalize!(x, 2) +_combinesets_cat(f_c, α::SetLike, β::SetLike) = CombinedSet(f_c, α, β) diff --git a/src/getdof.jl b/src/getdof.jl index dbce2202..c2edd15c 100644 --- a/src/getdof.jl +++ b/src/getdof.jl @@ -1,11 +1,30 @@ """ - MeasureBase.NoDOF{MU} + abstract type MeasureBase.AbstractNoDOF{MU} + +Abstract supertype for [`NoDOF`](@ref) and [`NoFastDOF`](@ref). +""" +abstract type AbstractNoDOF{MU} end + +Base.:+(nodof::AbstractNoDOF) = nodof +Base.:+(::IntegerLike, nodof::AbstractNoDOF) = nodof +Base.:+(nodof::AbstractNoDOF, ::IntegerLike) = nodof +Base.:+(nodof::AbstractNoDOF, ::AbstractNoDOF) = nodof + +Base.:*(nodof::AbstractNoDOF) = nodof +Base.:*(::IntegerLike, nodof::AbstractNoDOF) = nodof +Base.:*(nodof::AbstractNoDOF, ::IntegerLike) = nodof +Base.:*(nodof::AbstractNoDOF, ::AbstractNoDOF) = nodof + + +""" + MeasureBase.NoDOF{MU} <: AbstractNoDOF{MU} Indicates that there is no way to compute degrees of freedom of a measure of type `MU` with the given information, e.g. because the DOF are not a global property of the measure. """ -struct NoDOF{MU} end +struct NoDOF{MU} <: AbstractNoDOF{MU} end + """ getdof(μ) @@ -22,22 +41,86 @@ Also see [`check_dof`](@ref). function getdof end # Prevent infinite recursion: -@inline _default_getdof(::Type{MU}, ::MU) where {MU} = NoDOF{MU} +@inline _default_getdof(::Type{MU}, ::MU) where {MU} = NoDOF{MU}() @inline _default_getdof(::Type{MU}, mu_base) where {MU} = getdof(mu_base) @inline getdof(μ::MU) where {MU} = _default_getdof(MU, basemeasure(μ)) + +""" + MeasureBase.NoFastDOF{MU} <: AbstractNoDOF{MU} + +Indicates that there is no way to compute the degrees of freedom of a +measure of type `MU` efficiently. +""" +struct NoFastDOF{MU} <: AbstractNoDOF{MU} end + + +""" + fast_dof(μ::MU) + +Returns the effective number of degrees of freedom of variates of +measure `μ`, if it can be computed efficiently, otherwise +returns [`NoFastDOF{MU}()`](@ref). + +Defaults to `getdof(μ)` and should be specialized for measures for +which DOF can't be computed instantly. + +Also see [`getdof`](@ref) and [`check_dof`](@ref). +""" +function fast_dof end +export fast_dof + +fast_dof(μ) = getdof(μ) + + +""" + MeasureBase.some_dof(μ::AbstractMeasure) + +Get the DOF at some unspecified point of measure `μ`. + +Use with caution! + +In general, use [`getdof(μ)`](@ref) instead. `some_dof` is useful for +measures that are expected to have a constant DOF over their whole +space, but for which there is no way to compute it (or prove that +the DOF is constant over the measurable space). +""" +function some_dof end + +function some_dof(μ) + m = asmeasure(μ) + _try_direct_dof(m, getdof(m)) +end + +_try_direct_dof(::AbstractMeasure, dof::IntegerLike) = dof +_try_direct_dof(μ::AbstractMeasure, ::AbstractNoDOF) = + _try_local_dof(μ, some_dof(_some_localmeasure(μ))) + +_try_local_dof(::AbstractMeasure, dof::IntegerLike) = dof +_try_local_dof(μ::AbstractMeasure, ::AbstractNoDOF) = + throw(ArgumentError("Can't determine DOF for measure of type $(nameof(typeof(μ)))")) + +_some_localmeasure(μ::AbstractMeasure) = localmeasure(μ, testvalue(μ)) + + """ MeasureBase.check_dof(ν, μ)::Nothing Check if `ν` and `μ` have the same effective number of degrees of freedom -according to [`MeasureBase.getdof`](@ref). +according to [`MeasureBase.fast_dof`](@ref). + +Does not throw an exception if the DOF of `ν` or `μ` can't be computed +efficiently. """ function check_dof end function check_dof(ν, μ) - n_ν = getdof(ν) - n_μ = getdof(μ) + n_ν = fast_dof(ν) + n_μ = fast_dof(μ) + if n_ν isa AbstractNoDOF || n_μ isa AbstractNoDOF + return nothing + end if n_ν != n_μ throw( ArgumentError( @@ -66,7 +149,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/insupport.jl b/src/insupport.jl index a9a96363..bb09b1e6 100644 --- a/src/insupport.jl +++ b/src/insupport.jl @@ -1,12 +1,21 @@ """ - inssupport(m, x) + MeasureBase.NoFastInsupport{MU} + +Indicates that there is no fast way to compute if a point lies within the +support of measures of type `MU`. +""" +struct NoFastInsupport{MU} end + + +""" + insupport(m, x) insupport(m) -`insupport(m,x)` computes whether `x` is in the support of `m`. +`insupport(m, x)` computes whether `x` is in the support of `m` and +returns either a `Bool` or an instance of [`NoFastInsupport`](@ref). `insupport(m)` returns a function, and satisfies - -insupport(m)(x) == insupport(m, x) +`insupport(m)(x) == insupport(m, x)`. """ function insupport end @@ -15,12 +24,18 @@ function insupport end Checks if `x` is in the support of distribution/measure `μ`, throws an `ArgumentError` if not. + +Will not throw an exception if `insupport` returns an instance of +[`NoFastInsupport`](@ref). """ function require_insupport end function require_insupport(μ, x) - if !insupport(μ, x) - throw(ArgumentError("x is not within the support of μ")) + ins = insupport(μ, x) + if !(ins isa NoFastInsupport) + if !ins + throw(ArgumentError("x is not within the support of μ")) + end end return nothing end diff --git a/src/interface.jl b/src/interface.jl index 18080ac7..31e68ca0 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -5,7 +5,7 @@ using Reexport @reexport using MeasureBase using MeasureBase: basemeasure_depth, proxy, istrue -using MeasureBase: insupport, basemeasure_sequence, commonbase +using MeasureBase: insupport, basemeasure_sequence using MeasureBase: transport_to, NoTransport using DensityInterface: logdensityof @@ -21,7 +21,6 @@ export basemeasure_depth export proxy export insupport export basemeasure_sequence -export commonbase using Test @@ -64,7 +63,7 @@ function test_interface(μ::M) where {M} # testvalue, logdensityof x = @inferred testvalue(Float64, μ) - β = @inferred basemeasure(μ, x) + β = @inferred basemeasure(μ) ℓμ = @inferred logdensityof(μ, x) ℓβ = @inferred logdensityof(β, x) @@ -110,7 +109,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) @@ -119,7 +118,8 @@ function test_smf(μ, n = 100) @assert issorted(p) x = invsmf.(μ, p) @test issorted(x) - @test all(istrue ∘ insupport(μ), x) + # insupport may return a non-Bool "don't know" (NoFastInsupport): + @test all(x_i -> insupport(μ, x_i) != false, x) @test all((Finv ∘ F).(x) .≈ x) diff --git a/src/kernel.jl b/src/kernel.jl deleted file mode 100644 index d6667c7b..00000000 --- a/src/kernel.jl +++ /dev/null @@ -1,113 +0,0 @@ -export AbstractTransitionKernel, - GenericTransitionKernel, TypedTransitionKernel, ParameterizedTransitionKernel - -abstract type AbstractTransitionKernel <: AbstractMeasure end - -struct GenericTransitionKernel{F} <: AbstractTransitionKernel - f::F -end - -(k::GenericTransitionKernel)(x) = k.f(x) - -struct TypedTransitionKernel{M,F} <: AbstractTransitionKernel - m::M - f::F -end - -(k::TypedTransitionKernel)(x) = (k.m ∘ k.f)(x) -struct ParameterizedTransitionKernel{M,S,N,T} <: AbstractTransitionKernel - m::M - suff::S - param_maps::NamedTuple{N,T} - - function ParameterizedTransitionKernel( - ::Type{M}, - suff::S, - param_maps::NamedTuple{N,T}, - ) where {M,S,N,T} - new{Type{M},S,N,T}(M, suff, param_maps) - end - function ParameterizedTransitionKernel( - m::M, - suff::S, - param_maps::NamedTuple{N,T}, - ) where {M,S,N,T} - new{M,S,N,T}(m, suff, param_maps) - end -end - -""" -A *kernel* is a function that returns a measure. - - k1 = kernel() do x - Normal(x, x^2) - end - - k2 = kernel(Normal) do x - (μ = x, σ = x^2) - end - - k3 = kernel(Normal; μ = identity, σ = abs2) - - k4 = kernel(Normal; μ = first, σ = last) do x - (x, x^2) - end - - x = randn(); k1(x) == k2(x) == k3(x) == k4(x) - -This function is not exported, because "kernel" can have so many other meanings. -See for example https://github.com/JuliaGaussianProcesses/KernelFunctions.jl for -another common use of this term. - -# Reference - -* https://en.wikipedia.org/wiki/Markov_kernel -""" -function kernel end - -mapcall(t, x) = map(func -> func(x), t) - -function (k::ParameterizedTransitionKernel)(x) - s = k.suff(x) - k.m(; mapcall(k.param_maps, s)...) -end - -(k::AbstractTransitionKernel)(x1, x2, xs...) = k((x1, x2, xs...)) - -(k::AbstractTransitionKernel)(; kwargs...) = k(NamedTuple(kwargs)) - -""" -For any `k::TransitionKernel`, `basekernel` is expected to satisfy -``` -basekernel(k)(p) == (basemeasure ∘ k)(p) -``` - -The main purpose of `basekernel` is to make it efficient to compute -``` -basemeasure(d::ProductMeasure) == productmeasure(basekernel(d.f), d.xs) -``` -""" -function basekernel end - -# TODO: Find a way to do better than this -basekernel(f) = basemeasure ∘ f - -basekernel(f::Returns) = Returns(basemeasure(f.value)) - -function Base.show(io::IO, μ::AbstractTransitionKernel) - io = IOContext(io, :compact => true) - Pretty.pprint(io, μ) -end - -function Pretty.tile(k::K) where {K<:AbstractTransitionKernel} - Pretty.list_layout( - Pretty.tile.([getproperty(k, p) for p in propertynames(k)]), - prefix = nameof(constructorof(K)), - ) -end - -const kleisli = kernel - -export kleisli - -kernel(k::AbstractTransitionKernel) = k diff --git a/src/mass-interface.jl b/src/mass-interface.jl index 7b0518f9..33937498 100644 --- a/src/mass-interface.jl +++ b/src/mass-interface.jl @@ -22,7 +22,10 @@ for T in (:UnknownFiniteMass, :UnknownMass) @eval begin Base.:+(::$T, ::$T) = $T() Base.:*(::$T, ::$T) = $T() - Base.:^(::$T, k::Number) = isfinite(k) ? $T() : UnknownMass() + Base.:^(::$T, k::Real) = isfinite(k) ? $T() : UnknownMass() + # Disambiguation: + Base.:^(::$T, k::Integer) = isfinite(k) ? $T() : UnknownMass() + Base.:^(::$T, k::Rational) = isfinite(k) ? $T() : UnknownMass() end end @@ -65,7 +68,7 @@ finite, or we may know nothing at all about it. For these cases, it will return `UnknownFiniteMass` or `UnknownMass`, respectively. When no `massof` method exists, it defaults to `UnknownMass`. """ -massof(m::AbstractMeasure) = UnknownMass(m) +massof(::AbstractMeasure) = UnknownMass() struct NormalizedMeasure{P,M} <: AbstractMeasure parent::P @@ -105,7 +108,7 @@ isnormalized(x, p::Real = 2) = isone(norm(x, p)) isone(::AbstractUnknownMass) = false function massof(m, s) - _massof(m, s, rootmeasure(m)) + _default_massof_impl(m, s, rootmeasure(m)) end """ @@ -116,4 +119,16 @@ in this way, users should add the corresponding `massof` method. """ (m::AbstractMeasure)(s) = massof(m, s) -massof(μ, a_b::AbstractInterval) = smf(μ, rightendpoint(a_b)) - smf(μ, leftendpoint(a_b)) +function massof(μ, a_b::AbstractInterval) + _smf_interval_massof(μ, smf(μ, rightendpoint(a_b)), smf(μ, leftendpoint(a_b))) +end + +_smf_interval_massof(μ, smf_r, smf_l) = smf_r - smf_l + +function _smf_interval_massof(μ, ::NoSMF, ::NoSMF) + throw( + ArgumentError( + "Can't compute the mass over an interval for a measure of type $(nameof(typeof(μ))), no statistical measure function available", + ), + ) +end diff --git a/src/measure_operators.jl b/src/measure_operators.jl new file mode 100644 index 00000000..41606367 --- /dev/null +++ b/src/measure_operators.jl @@ -0,0 +1,135 @@ +""" + module MeasureOperators + +Defines the following operators for measures: + +* `f ⋄ μ == pushfwd(f, μ)` +* `μ ⊙ f == pullbck(f, μ)` +* `μ ▷ k == mbind(k, μ)` +* `⊗(μs...) == productmeasure(μs)` +* `∫(f, μ) == mintegrate(f, μ)` +* `∫exp(f, μ) == mintegrate_exp(f, μ)` +* `𝒹(ν, μ) == density_rel(ν, μ)` +* `log𝒹(ν, μ) == logdensity_rel(ν, μ)` +""" +module MeasureOperators + +using MeasureBase: AbstractMeasure +using MeasureBase: pushfwd, pullbck, mbind, productmeasure +using MeasureBase: mintegrate, mintegrate_exp, density_rel, logdensity_rel +using InverseFunctions: inverse + +@doc raw""" + ⋄(f, μ::AbstractMeasure) = pushfwd(f, μ) + +The `\\diamond` operator denotes a pushforward operation: `ν = f ⋄ μ` +generates a +[pushforward measure](https://en.wikipedia.org/wiki/Pushforward_measure). + +A common mathematical notation for a pushforward is ``f_*μ``, but as +there is no "subscript-star" operator in Julia, we use `⋄`. + +See [`pushfwd(f, μ)`](@ref) for details. + +Also see [`ν ⊙ f`](@ref), the pullback operator. +""" +⋄(f, μ::AbstractMeasure) = pushfwd(f, μ) +export ⋄ + +@doc raw""" + ⊙(ν::AbstractMeasure, f) = pullbck(f, ν) + +The `\\odot` operator denotes a pullback operation. + +See also [`pullbck(ν, f)`](@ref) for details. Note that `pullbck` takes it's +arguments in different order, in keeping with the Julia convention of +passing functions as the first argument. A pullback is mathematically the +precomposition of a measure `μ`` with the function `f` applied to sets. so +`⊙` takes the measure as the first and the function as the second argument, +as common in mathematical notation for precomposition. + +A common mathematical notation for pullback in measure theory is +``f \circ μ``, but as `∘` is used for function composition in Julia and as +`f` semantically acts point-wise on sets, we use `⊙`. + +Also see [f ⋄ μ](@ref), the pushforward operator. +""" +⊙(ν::AbstractMeasure, f) = pullbck(f, ν) +export ⊙ + +""" + μ ▷ k = mbind(k, μ) + +The `\\triangleright` operator denotes a measure monadic bind operation. + +A common operator choice for a monadic bind operator is `>>=` (e.g. in +the Haskell programming language), but this has a different meaning in +Julia and there is no close equivalent, so we use `▷`. + +See [`mbind(k, μ)`](@ref) for details. Note that `mbind` takes its +arguments in different order, in keeping with the Julia convention of +passing functions as the first argument. `▷`, on the other hand, takes +its arguments in the order common for monadic binds in functional +programming (like the Haskell `>>=` operator) and mathematics. +""" +▷(μ::AbstractMeasure, k) = mbind(k, μ) +export ▷ + +# ToDo: Use `⨂` instead of `⊗` for better readability? +""" + ⊗(μs::AbstractMeasure...) = productmeasure(μs) + +`⊗` is an operator for building product measures. + +See [`productmeasure(μs)`](@ref) for details. +""" +⊗(μs::AbstractMeasure...) = productmeasure(μs) +export ⊗ + +""" + ∫(f, μ::AbstractMeasure) = mintegrate(f, μ) + +Denotes an indefinite integral of the function `f` with respect to the +measure `μ`. + +See [`mintegrate(f, μ)`](@ref) for details. +""" +∫(f, μ::AbstractMeasure) = mintegrate(f, μ) +export ∫ + +""" + ∫exp(f, μ::AbstractMeasure) = mintegrate_exp(f, μ) + +Generates a new measure that is the indefinite integral of `exp` of `f` +with respect to the measure `μ`. + +See [`mintegrate_exp(f, μ)`](@ref) for details. +""" +∫exp(f, μ::AbstractMeasure) = mintegrate_exp(f, μ) +export ∫exp + +""" + 𝒹(ν, μ) = density_rel(ν, μ) + +Compute the density, i.e. the +[Radom-Nikodym derivative](https://en.wikipedia.org/wiki/Radon%E2%80%93Nikodym_theorem) +of `ν`` with respect to `μ`. + +For details, see [`density_rel(ν, μ)`}(@ref). +""" +𝒹(ν, μ::AbstractMeasure) = density_rel(ν, μ) +export 𝒹 + +""" + log𝒹(ν, μ) = logdensity_rel(ν, μ) + +Compute the log-density, i.e. the logarithm of the +[Radom-Nikodym derivative](https://en.wikipedia.org/wiki/Radon%E2%80%93Nikodym_theorem) +of `ν`` with respect to `μ`. + +For details, see [`logdensity_rel(ν, μ)`}(@ref). +""" +log𝒹(ν, μ::AbstractMeasure) = logdensity_rel(ν, μ) +export log𝒹 + +end # module MeasureOperators diff --git a/src/mspace.jl b/src/mspace.jl new file mode 100644 index 00000000..6227537d --- /dev/null +++ b/src/mspace.jl @@ -0,0 +1,46 @@ +""" + MeasureBase.NoMSpaceElementSize{MU} + +Indicates that either the measurable space of measures of type `MU` is not +a space over arrays, or that the size of the arrays is not fixed or can not +be easily/efficiently determined. +""" +struct NoMSpaceElementSize{MU} end + + +""" + mspace_elsize(μ) + +For a measure `μ` over an array-valued measurable space, return the size of +the arrays that are the elements of the space. + +May return [`NoMSpaceElementSize{typeof(μ)}()`](@ref). +""" +function mspace_elsize end +export mspace_elsize + +@inline mspace_elsize(μ::AbstractMeasure) = NoMSpaceElementSize{typeof(μ)}() + + +""" + MeasureBase.some_mspace_elsize(μ::AbstractMeasure) + +For a measure `μ` over an array-valued measurable space, return the size of +an arbitrary element of the space. + +Use with caution, the space of some measures is made up of arrays of +different sizes! + +In general, use [`mspace_elsize(μ)`](@ref) instead. `some_mspace_elsize` is +useful if the measurable space is expected to contain only arrays of the +same size but there is no way to prove this automatically. Algorithms that +use the returned size should always check that it matches the size of each +point of the space that is processed. +""" +function some_mspace_elsize end + +@inline some_mspace_elsize(μ) = _mspace_some_elsize_impl(μ, mspace_elsize(μ)) + +@inline _mspace_some_elsize_impl(::AbstractMeasure, sz::SizeLike) = sz +_mspace_some_elsize_impl(μ::AbstractMeasure, ::NoMSpaceElementSize) = + maybestatic_size(testvalue(μ)) diff --git a/src/parameterized.jl b/src/parameterized.jl index 78e43995..4412ebdd 100644 --- a/src/parameterized.jl +++ b/src/parameterized.jl @@ -24,27 +24,6 @@ function Pretty.tile(d::ParameterizedMeasure{()}) result end -# Allow things like -# -# julia> Normal{(:μ,)}(2) -# Normal(μ = 2,) -function kernel(::Type{P}) where {N,P<:ParameterizedMeasure{N}} - C = constructorof(P) - _kernel(C, Val(N)) -end - -@inline function _kernel(::Type{C}, ::Val{N}) where {C,N} - @inline function f(args::T) where {T<:Tuple} - C(NamedTuple{N,T}(args))::C{N,T} - end - - @inline function f(arg::T) where {T} - C(NamedTuple{N,Tuple{T}}((arg,)))::C{N,Tuple{T}} - end - - kernel(f) -end - function (::Type{P})(nt::NamedTuple{K,T}) where {K,T,N,P<:ParameterizedMeasure{N}} C = constructorof(P) arg = NamedTuple{N}(nt) @@ -127,14 +106,3 @@ params(::Type{PM}) where {N,PM<:ParameterizedMeasure{N}} = N function paramnames(μ, constraints::NamedTuple{N}) where {N} tuple((k for k in paramnames(μ) if k ∉ N)...) end - -############################################################################### -# kernelfactor - -function kernelfactor(::Type{P}) where {N,P<:ParameterizedMeasure{N}} - (constructorof(P), N) -end - -function kernelfactor(::P) where {N,P<:ParameterizedMeasure{N}} - (constructorof(P), N) -end diff --git a/src/primitive.jl b/src/primitive.jl index 85cf2beb..0a2aba37 100644 --- a/src/primitive.jl +++ b/src/primitive.jl @@ -11,7 +11,7 @@ measures satisfy the following laws: logdensity_def(μ::PrimitiveMeasure, x) = 0.0 - logdensity_def(μ::M, ν::M, x) where {M<:PrimitiveMeasure} = 0.0 + logdensity_rel_def(μ::M, ν::M, x) where {M<:PrimitiveMeasure} = 0.0 """ abstract type PrimitiveMeasure <: AbstractMeasure end @@ -19,12 +19,12 @@ basemeasure(μ::PrimitiveMeasure) = μ @inline basemeasure_depth(::PrimitiveMeasure) = static(0) -@inline logdensityof(::PrimitiveMeasure, x::Real) = zero(float(typeof(x))) -@inline logdensityof(::PrimitiveMeasure, x) = false +@inline logdensityof_impl(::PrimitiveMeasure, x::Number) = zero(float(typeof(x))) +@inline logdensityof_impl(::PrimitiveMeasure, x) = static(0.0) logdensity_def(::PrimitiveMeasure, x) = static(0.0) -logdensity_def(μ::M, ν::M, x) where {M<:PrimitiveMeasure} = 0.0 +logdensity_rel_def(μ::M, ν::M, x) where {M<:PrimitiveMeasure} = 0.0 function Pretty.quoteof(μ::M) where {M<:PrimitiveMeasure} :($M()) diff --git a/src/primitives/counting.jl b/src/primitives/counting.jl index c61d0624..74e7f8c2 100644 --- a/src/primitives/counting.jl +++ b/src/primitives/counting.jl @@ -12,12 +12,12 @@ struct Counting{T} <: AbstractMeasure Counting(supp) = new{Core.Typeof(supp)}(supp) end -@inline function logdensityof(μ::Counting, x::Real) +@inline function logdensityof_impl(μ::Counting, x::Number) R = float(typeof(x)) - insupport(μ, x) ? zero(R) : R(-Inf) + _checksupport(insupport(μ, x), zero(R)) end -@inline logdensityof(μ::Counting, x) = insupport(μ, x) ? 0.0 : -Inf +@inline logdensityof_impl(μ::Counting, x) = _checksupport(insupport(μ, x), 0.0) @inline logdensity_def(μ::Counting, x) = logdensityof(μ, x) @@ -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/dirac.jl b/src/primitives/dirac.jl index 01297486..96077696 100644 --- a/src/primitives/dirac.jl +++ b/src/primitives/dirac.jl @@ -9,6 +9,9 @@ function Pretty.tile(d::Dirac) Pretty.literal("Dirac(") * Pretty.tile(d.x) * Pretty.literal(")") end +Base.:(==)(a::Dirac, b::Dirac) = a.x == b.x +Base.isapprox(a::Dirac, b::Dirac; kwargs...) = isapprox(a.x, b.x; kwargs...) + gentype(μ::Dirac{X}) where {X} = X function (μ::Dirac{X})(s) where {X} @@ -20,14 +23,14 @@ basemeasure(d::Dirac) = CountingBase() massof(::Dirac) = static(1.0) -function logdensityof(μ::Dirac, x::Real) +function logdensityof_impl(μ::Dirac, x::Number) R = float(typeof(x)) - insupport(μ, x) ? zero(R) : R(-Inf) + _checksupport(insupport(μ, x), zero(R)) end -logdensityof(μ::Dirac, x) = insupport(μ, x) ? 0.0 : -Inf +logdensityof_impl(μ::Dirac, x) = _checksupport(insupport(μ, x), 0.0) -logdensity_def(::Dirac, x::Real) = zero(float(typeof(x))) +logdensity_def(::Dirac, x::Number) = zero(float(typeof(x))) logdensity_def(::Dirac, x) = 0.0 Base.rand(::Random.AbstractRNG, T::Type, μ::Dirac) = μ.x @@ -40,6 +43,8 @@ insupport(d::Dirac, x) = x == d.x @inline getdof(::Dirac) = static(0) +@inline mspace_elsize(μ::Dirac) = maybestatic_size(μ.x) + @propagate_inbounds function checked_arg(μ::Dirac, x) @boundscheck insupport(μ, x) || throw(ArgumentError("Invalid variate for measure")) x diff --git a/src/primitives/lebesgue.jl b/src/primitives/lebesgue.jl index 3846eaf5..770ed03e 100644 --- a/src/primitives/lebesgue.jl +++ b/src/primitives/lebesgue.jl @@ -12,13 +12,13 @@ insupport(::LebesgueBase, x) = true insupport(::LebesgueBase) = Returns(true) -logdensity_def(::LebesgueBase, ::CountingBase, x) = -Inf +logdensity_rel_def(::LebesgueBase, ::CountingBase, x) = -Inf -logdensity_def(::CountingBase, ::LebesgueBase, x) = Inf +logdensity_rel_def(::CountingBase, ::LebesgueBase, x) = Inf @inline getdof(::LebesgueBase) = static(1) -@inline checked_arg(::LebesgueBase, x::Real) = x +@inline checked_arg(::LebesgueBase, x::Number) = x @propagate_inbounds function checked_arg(::LebesgueBase, x::Any) @boundscheck throw(ArgumentError("Invalid variate type for measure")) @@ -26,12 +26,12 @@ end massof(::LebesgueBase) = static(Inf) -function _massof(m, s::Interval, ::LebesgueBase) +function _default_massof_impl(m, s::AbstractInterval, ::LebesgueBase) mass = massof(m) nu = mass * StdUniform() f = transport_to(nu, m) - a = f(minimum(s)) - b = f(maximum(s)) + a = f(leftendpoint(s)) + b = f(rightendpoint(s)) return mass * abs(b - a) end @@ -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,37 +61,37 @@ 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) +@inline function logdensityof_impl(μ::Lebesgue, x::Number) R = float(typeof(x)) - insupport(μ, x) ? zero(R) : R(-Inf) + _checksupport(insupport(μ, x), zero(R)) end -@inline logdensityof(μ::Lebesgue, x) = insupport(μ, x) ? 0.0 : -Inf +@inline logdensityof_impl(μ::Lebesgue, x) = _checksupport(insupport(μ, x), 0.0) -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/proxies.jl b/src/proxies.jl index 95aed270..f2805176 100644 --- a/src/proxies.jl +++ b/src/proxies.jl @@ -15,17 +15,29 @@ macro useproxy(M) M = esc(M) quote @inline $MeasureBase.logdensity_def(μ::$M, x) = logdensity_def(proxy(μ), x) + @inline $MeasureBase.unsafe_logdensityof(μ::$M, x) = unsafe_logdensityof(proxy(μ), x) @inline $MeasureBase.basemeasure(μ::$M) = basemeasure(proxy(μ)) - @inline $MeasureBase.basemeasure_depth(μ::$M) = basemeasure_depth(proxy(μ)) + @inline $MeasureBase.rootmeasure(μ::$M) = rootmeasure(proxy(μ)) + + @inline $MeasureBase.insupport(μ::$M) = insupport(proxy(μ)) + + @inline $MeasureBase.getdof(μ::$M) = getdof(proxy(μ)) + @inline $MeasureBase.fast_dof(μ::$M) = fast_dof(proxy(μ)) @inline $MeasureBase.transport_origin(μ::$M) = transport_origin(proxy(μ)) @inline $MeasureBase.to_origin(μ::$M, y) = to_origin(proxy(μ), y) @inline $MeasureBase.from_origin(μ::$M, x) = from_origin(proxy(μ), x) + @inline $MeasureBase.localmeasure(μ::$M, x) = localmeasure(proxy(μ), x) + @inline $MeasureBase.transportmeasure(μ::$M, x) = transportmeasure(proxy(μ), x) + @inline $MeasureBase.massof(μ::$M) = massof(proxy(μ)) @inline $MeasureBase.massof(μ::$M, s) = massof(proxy(μ), s) + # Disambiguation with massof(μ, ::AbstractInterval): + @inline $MeasureBase.massof(μ::$M, s::$(IntervalSets.AbstractInterval)) = + massof(proxy(μ), s) @inline $MeasureBase.smf(μ::$M, x) = smf(proxy(μ), x) @inline $MeasureBase.invsmf(μ::$M, x) = invsmf(proxy(μ), x) diff --git a/src/schema.jl b/src/schema.jl deleted file mode 100644 index 70c85577..00000000 --- a/src/schema.jl +++ /dev/null @@ -1,34 +0,0 @@ -# Taken from https://github.com/cscherrer/NestedTuples.jl/blob/cd298fd1e5f7e571701a6fee916d2d47c19f32f5/src/typelevel.jl - -ntkeys(::Type{NamedTuple{K,V}}) where {K,V} = K -ntvaltype(::Type{NamedTuple{K,V}}) where {K,V} = V - -""" - schema(::Type) - -`schema` turns a type into a value that's easier to work with. -Example: - julia> nt = (a=(b=[1,2],c=(d=[3,4],e=[5,6])),f=[7,8]); - julia> NT = typeof(nt) - NamedTuple{(:a, :f),Tuple{NamedTuple{(:b, :c),Tuple{Array{Int64,1},NamedTuple{(:d, :e),Tuple{Array{Int64,1},Array{Int64,1}}}}},Array{Int64,1}}} - julia> schema(NT) - (a = (b = Array{Int64,1}, c = (d = Array{Int64,1}, e = Array{Int64,1})), f = Array{Int64,1}) -""" -function schema end - -schema(::NamedTuple{(),Tuple{}}) = NamedTuple() -schema(::Type{NamedTuple{(),Tuple{}}}) = NamedTuple() - -function schema(NT::Type{NamedTuple{names,T}}) where {names,T} - return NamedTuple{ntkeys(NT)}(schema(ntvaltype(NT))) -end - -function schema(TT::Type{T}) where {T<:Tuple} - return schema.(Tuple(TT.types)) -end - -schema(t::T) where {T<:Tuple} = schema(T) - -schema(t::T) where {T<:NamedTuple} = schema(T) - -schema(T) = T diff --git a/src/splat.jl b/src/splat.jl deleted file mode 100644 index d1df4f17..00000000 --- a/src/splat.jl +++ /dev/null @@ -1,11 +0,0 @@ -struct Splat{F} - f::F -end - -function (s::Splat{F})(x) where {F} - s.f(x...) -end - -unsplat(s::Splat) = s.f - -splat(f) = Splat(f) diff --git a/src/standard/stdexponential.jl b/src/standard/stdexponential.jl index c985c224..1e10bb88 100644 --- a/src/standard/stdexponential.jl +++ b/src/standard/stdexponential.jl @@ -4,9 +4,9 @@ export StdExponential insupport(::StdExponential, x) = x ≥ zero(x) -@inline function logdensityof(::StdExponential, x) +@inline function logdensityof_impl(d::StdExponential, x) R = float(typeof(x)) - x ≥ zero(R) ? convert(R, -x) : R(-Inf) + _checksupport(insupport(d, x), convert(R, -x)) end @inline logdensity_def(::StdExponential, x) = -x diff --git a/src/standard/stdlogistic.jl b/src/standard/stdlogistic.jl index 58a1ba67..b28dd618 100644 --- a/src/standard/stdlogistic.jl +++ b/src/standard/stdlogistic.jl @@ -4,7 +4,7 @@ export StdLogistic @inline insupport(d::StdLogistic, x) = true -@inline logdensityof(::StdLogistic, x) = (u = -abs(x); u - 2 * log1pexp(u)) +@inline logdensityof_impl(::StdLogistic, x) = (u = -abs(x); u - 2 * log1pexp(u)) @inline logdensity_def(::StdLogistic, x) = logdensityof(StdLogistic(), x) @inline basemeasure(::StdLogistic) = LebesgueBase() diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl index 833f280e..db0e42fe 100644 --- a/src/standard/stdmeasure.jl +++ b/src/standard/stdmeasure.jl @@ -4,120 +4,15 @@ StdMeasure(::typeof(rand)) = StdUniform() StdMeasure(::typeof(randexp)) = StdExponential() StdMeasure(::typeof(randn)) = StdNormal() -@inline check_dof(::StdMeasure, ::StdMeasure) = nothing - -@inline transport_def(::MU, μ::MU, x) where {MU<:StdMeasure} = x - -function transport_def(ν::StdMeasure, μ::PowerMeasure{<:StdMeasure}, x) - return transport_def(ν, μ.parent, only(x)) -end - -function transport_def(ν::PowerMeasure{<:StdMeasure}, μ::StdMeasure, x) - return fill_with(transport_def(ν.parent, μ, only(x)), map(length, ν.axes)) -end - -function transport_def( - ν::PowerMeasure{<:StdMeasure,<:NTuple{1,Base.OneTo}}, - μ::PowerMeasure{<:StdMeasure,<:NTuple{1,Base.OneTo}}, - x, -) - return transport_to(ν.parent, μ.parent).(x) -end - -function transport_def( - ν::PowerMeasure{<:StdMeasure,<:NTuple{N,Base.OneTo}}, - μ::PowerMeasure{<:StdMeasure,<:NTuple{M,Base.OneTo}}, - x, -) where {N,M} - return reshape(transport_to(ν.parent, μ.parent).(x), map(length, ν.axes)...) -end - -# Implement transport_to(NU::Type{<:StdMeasure}, μ) and transport_to(ν, MU::Type{<:StdMeasure}): - -_std_measure(::Type{M}, ::StaticInteger{1}) where {M<:StdMeasure} = M() -_std_measure(::Type{M}, dof::IntegerLike) where {M<:StdMeasure} = M()^dof -_std_measure_for(::Type{M}, μ::Any) where {M<:StdMeasure} = _std_measure(M, getdof(μ)) - -function transport_to(::Type{NU}, μ) where {NU<:StdMeasure} - transport_to(_std_measure_for(NU, μ), μ) -end - -function transport_to(ν, ::Type{MU}) where {MU<:StdMeasure} - transport_to(ν, _std_measure_for(MU, ν)) -end - -# Transform between standard measures and Dirac: +""" + MeasureBase.StdPowerMeasure{MU<:StdMeasure,N} -@inline transport_def(ν::Dirac, ::PowerMeasure{<:StdMeasure}, ::Any) = ν.x +The type of an `N`-dimensional power of a standard measure of type `MU`. +""" +const StdPowerMeasure{MU<:StdMeasure,N} = PowerMeasure{MU,<:NTuple{N,OneToLike}} -@inline function transport_def(ν::PowerMeasure{<:StdMeasure}, ::Dirac, ::Any) - Zeros{Bool}(map(_ -> 0, ν.axes)) -end +@inline mspace_elsize(::StdMeasure) = () -# Helpers for product transforms and similar: - -struct _TransportToStd{NU<:StdMeasure} <: Function end -_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) - -function _tuple_transport_def( - ν::PowerMeasure{NU}, - μs::Tuple, - xs::Tuple, -) where {NU<:StdMeasure} - reshape(vcat(map(_TransportToStd{NU}, μs, xs)...), ν.axes) -end - -function transport_def( - ν::PowerMeasure{NU}, - μ::ProductMeasure{<:Tuple}, - x, -) where {NU<:StdMeasure} - _tuple_transport_def(ν, marginals(μ), x) -end - -function transport_def( - ν::PowerMeasure{NU}, - μ::ProductMeasure{<:NamedTuple{names}}, - x, -) where {NU<:StdMeasure,names} - _tuple_transport_def(ν, values(marginals(μ)), values(x)) -end - -@inline _offset_cumsum(s, x, y, rest...) = (s, _offset_cumsum(s + x, y, rest...)...) -@inline _offset_cumsum(s, x) = (s,) -@inline _offset_cumsum(s) = () - -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) -end - -function _tuple_transport_def( - νs::Tuple, - μ::PowerMeasure{MU}, - x::AbstractArray{<:Real}, -) where {MU<:StdMeasure} - vrs = _stdvar_viewranges(νs, firstindex(x)) - xs = map(r -> view(x, r), vrs) - map(_TransportFromStd{MU}, νs, xs) -end - -function transport_def( - ν::ProductMeasure{<:Tuple}, - μ::PowerMeasure{MU}, - x, -) where {MU<:StdMeasure} - _tuple_transport_def(marginals(ν), μ, x) -end +@inline check_dof(::StdMeasure, ::StdMeasure) = nothing -function transport_def( - ν::ProductMeasure{<:NamedTuple{names}}, - μ::PowerMeasure{MU}, - x, -) where {MU<:StdMeasure,names} - NamedTuple{names}(_tuple_transport_def(values(marginals(ν)), μ, x)) -end +@inline transport_def(::MU, μ::MU, x) where {MU<:StdMeasure} = x diff --git a/src/standard/stdnormal.jl b/src/standard/stdnormal.jl index 057a8629..f636d311 100644 --- a/src/standard/stdnormal.jl +++ b/src/standard/stdnormal.jl @@ -7,7 +7,7 @@ export StdNormal @inline insupport(::StdNormal, x) = true -@inline logdensityof(::StdNormal, x) = (-x^2 - log2π) / 2 +@inline logdensityof_impl(::StdNormal, x) = (-x^2 - log2π) / 2 @inline logdensity_def(::StdNormal, x) = -x^2 / 2 @inline basemeasure(::StdNormal) = WeightedMeasure(static(-0.5 * log2π), LebesgueBase()) diff --git a/src/standard/stduniform.jl b/src/standard/stduniform.jl index 7bbe15ed..f9b07bf1 100644 --- a/src/standard/stduniform.jl +++ b/src/standard/stduniform.jl @@ -2,11 +2,11 @@ struct StdUniform <: StdMeasure end export StdUniform -insupport(::StdUniform, x) = zero(x) ≤ x ≤ one(x) +insupport(::StdUniform, x) = (zero(x) ≤ x) & (x ≤ one(x)) -@inline function logdensityof(::StdUniform, x) +@inline function logdensityof_impl(d::StdUniform, x) R = float(typeof(x)) - zero(x) ≤ x ≤ one(x) ? zero(R) : R(-Inf) + _checksupport(insupport(d, x), zero(R)) end @inline logdensity_def(::StdUniform, x) = zero(x) @@ -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/transport.jl b/src/transport.jl index b0c8ed41..65c5f386 100644 --- a/src/transport.jl +++ b/src/transport.jl @@ -8,6 +8,8 @@ See [`MeasureBase.transport_origin`](@ref). """ struct NoTransportOrigin{NU} end +Base.:^(origin::NoTransportOrigin, ::IntegerLike) = origin + """ MeasureBase.transport_origin(ν) @@ -135,11 +137,29 @@ end return static(10) end -# If both both measures have no origin: +# If both measures have no origin: function _transport_between_origins(ν, ::StaticInteger{0}, ::StaticInteger{0}, μ, x) + _transport_between_noorigins(ν, fast_dof(ν), fast_dof(μ), μ, x) +end + +function _transport_between_noorigins(ν, ::IntegerLike, ::IntegerLike, μ, x) _transport_with_intermediate(ν, _transport_intermediate(ν, μ), μ, x) end +# If the DOF of either side is not known, pivot through a flat vector of +# standard-measure variates, using the with-rest transport protocol: +_transport_between_noorigins(ν, ::AbstractNoDOF, ::IntegerLike, μ, x) = + _transport_mvstd_pivot(ν, μ, x) +_transport_between_noorigins(ν, ::IntegerLike, ::AbstractNoDOF, μ, x) = + _transport_mvstd_pivot(ν, μ, x) +_transport_between_noorigins(ν, ::AbstractNoDOF, ::AbstractNoDOF, μ, x) = + _transport_mvstd_pivot(ν, μ, x) + +function _transport_mvstd_pivot(ν, μ, x) + z = transport_to_mvstd(StdUniform(), μ, x) + return _transport_from_mvstd(ν, StdUniform(), z) +end + @generated function _transport_between_origins( ν, ::StaticInteger{n_ν}, @@ -147,6 +167,21 @@ end μ, x, ) where {n_ν,n_μ} + if n_ν == 10 + return :(throw( + ArgumentError( + "Transport to measure of type $(nameof(typeof(ν))) not supported, origin stack too deep.", + ), + )) + end + if n_μ == 10 + return :(throw( + ArgumentError( + "Transport from measure of type $(nameof(typeof(μ))) not supported, origin stack too deep.", + ), + )) + end + prog = quote μ0 = μ x0 = x @@ -183,8 +218,8 @@ end return prog end -@inline _transport_intermediate(ν, μ) = _transport_intermediate(getdof(ν), getdof(μ)) -@inline _transport_intermediate(::Integer, n_μ::Integer) = StdUniform()^n_μ +@inline _transport_intermediate(ν, μ) = _transport_intermediate(fast_dof(ν), fast_dof(μ)) +@inline _transport_intermediate(::IntegerLike, n_μ::IntegerLike) = StdUniform()^n_μ @inline _transport_intermediate(::StaticInteger{1}, ::StaticInteger{1}) = StdUniform() _call_transport_def(ν, μ, x) = transport_def(ν, μ, x) @@ -227,7 +262,7 @@ struct TransportFunction{NU,MU} <: Function end end -@inline transport_to(ν, μ) = TransportFunction(ν, μ) +@inline transport_to(ν, μ) = TransportFunction(asmeasure(ν), asmeasure(μ)) function Base.:(==)(a::TransportFunction, b::TransportFunction) return a.ν == b.ν && a.μ == b.μ diff --git a/src/utils.jl b/src/utils.jl index 0ec81a50..b6035e96 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -19,8 +19,6 @@ testvalue(::Type{T}) where {T} = zero(T) export rootmeasure -basemeasure(μ, x) = basemeasure(μ) - """ rootmeasure(μ::AbstractMeasure) @@ -95,36 +93,6 @@ measure of the previous term, and with no repeated entries. return filter(!isnothing, Base.Cartesian.@ntuple 10 b) end -commonbase(μ, ν) = commonbase(μ, ν, Any) - -""" - commonbase(μ, ν, T) -> Tuple{StaticInt{i}, StaticInt{j}} - -Find minimal (with respect to their sum) `i` and `j` such that there is a method - - logdensity_def(basemeasure_sequence(μ)[i], basemeasure_sequence(ν)[j], ::T) - -This is used in `logdensity_rel` to help make that function efficient. -""" -@inline function commonbase(μ, ν, ::Type{T}) where {T} - return commonbase(basemeasure_sequence(μ), basemeasure_sequence(ν), T) -end - -@generated function commonbase(μ::M, ν::N, ::Type{T}) where {M<:Tuple,N<:Tuple,T} - m = schema(M) - n = schema(N) - - sols = Iterators.filter( - ((i, j),) -> static_hasmethod(logdensity_def, Tuple{m[i],n[j],T}), - Iterators.product(1:length(m), 1:length(n)), - ) - isempty(sols) && return :(nothing) - minsol = static.(argmin(((i, j),) -> i + j, sols)) - quote - $minsol - end -end - mymap(f, gen::Base.Generator) = mymap(f ∘ gen.f, gen.iter) mymap(f, inds...) = Iterators.map(f, inds...) @@ -170,7 +138,7 @@ fcomp(::typeof(identity), g) = g fcomp(f, ::typeof(identity)) = f fcomp(::typeof(identity), ::typeof(identity)) = identity -near_neg_inf(::Type{T}) where {T<:Real} = T(-1E38) # Still fits into Float32 +near_neg_inf(::Type{T}) where {T<:Number} = T(-1E38) # Still fits into Float32 isneginf(x) = isinf(x) && x < zero(x) isposinf(x) = isinf(x) && x > zero(x) @@ -180,3 +148,43 @@ isapproxzero(A::AbstractArray) = all(isapproxzero, A) isapproxone(x::T) where {T<:Real} = x ≈ one(T) isapproxone(A::AbstractArray) = all(isapproxone, A) + +containsnan(x::Number) = 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<:Number} = T(x) +convert_realtype(::Type{T}, x::AbstractArray{U}) where {T<:Real,U<:Number} = 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) + +# 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..19645d48 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,21 +1,31 @@ [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" +JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" LogarithmicNumbers = "aa2f6b4e-9042-5d33-9679-40d3a6b85899" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" +OneTwoMany = "762dc654-8631-413a-a342-372a7419ad9d" +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/bind.jl b/test/combinators/bind.jl new file mode 100644 index 00000000..efd2cabd --- /dev/null +++ b/test/combinators/bind.jl @@ -0,0 +1,158 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using Random +using StableRNGs: StableRNG +using AffineMaps: Mul + +using MeasureBase +using MeasureBase: StdExponential, StdNormal, StdUniform +using MeasureBase: mbind, mkernel, bindkernel, boundmeasure +using MeasureBase: pushfwd, productmeasure, transport_to, transportmeasure, localmeasure + +@testset "bind" begin + stblrng() = StableRNG(789990641) + + f_β(σ) = pushfwd(Mul(σ + 0.5), StdNormal()) + α = StdExponential() + + @testset "monadic bind" begin + μ = mbind(f_β, α) + @test μ isa MeasureBase.Bind + @test boundmeasure(μ) === α + @test bindkernel(μ) isa MeasureBase.MKernel + @test mbind(bindkernel(μ), α) == μ + @test mbind(f_β)(α) == μ + + a = rand(stblrng(), Float64, α) + b = rand(copy(stblrng()), Float64, μ) # not comparable directly, just smoke: + @test rand(stblrng(), Float64, μ) isa Real + + @test MeasureBase.insupport(μ, 0.4) isa MeasureBase.NoFastInsupport + @test MeasureBase.getdof(μ) isa MeasureBase.NoDOF + @test_throws ArgumentError basemeasure(μ) + @test_throws ArgumentError MeasureBase.rootmeasure(μ) + end + + @testset "mbind with tuple and Pair" begin + for f_c in (tuple, Pair) + μ = mbind(f_β, α, f_c) + ab = rand(stblrng(), Float64, μ) + a, b = f_c === tuple ? ab : (ab.first, ab.second) + @test logdensityof(μ, ab) ≈ logdensityof(α, a) + logdensityof(f_β(a), b) + + tpm = transportmeasure(μ, ab) + @test logdensityof(tpm, ab) ≈ logdensityof(μ, ab) + @test localmeasure(μ, ab) == tpm + end + end + + @testset "mbind with vcat" begin + αv = StdExponential()^1 + f_βv(a) = pushfwd(Mul(a[1] + 0.5), StdNormal())^2 + μ = mbind(f_βv, αv, vcat) + + xy = rand(stblrng(), Float64, μ) + @test xy isa AbstractVector{<:Real} && length(xy) == 3 + a, b = xy[1:1], xy[2:3] + @test logdensityof(μ, xy) ≈ logdensityof(αv, a) + logdensityof(f_βv(a), b) + + # Transport to and from a standard measure, dof of μ is not fast-computable: + y = transport_to(StdUniform()^3, μ)(xy) + @test y isa AbstractVector{<:Real} && length(y) == 3 + @test all(u -> 0 <= u <= 1, y) + xy_reco = transport_to(μ, StdUniform()^3)(y) + @test xy_reco ≈ xy + + # Transport between two measures of unknown DOF (mvstd pivot): + μ2 = mbind(f_βv, StdUniform()^1, vcat) + xy2 = transport_to(μ2, μ)(xy) + @test xy2 isa AbstractVector{<:Real} && length(xy2) == 3 + @test transport_to(μ, μ2)(xy2) ≈ xy + @test logdensityof(μ2, xy2) isa Real + + # Transport between known-DOF and unknown-DOF measures (mvstd pivot): + ν_known = productmeasure((StdNormal(), StdNormal(), StdNormal())) + z = transport_to(ν_known, μ)(xy) + @test z isa Tuple{Vararg{Real,3}} + @test collect(transport_to(μ, ν_known)(z)) ≈ xy + + # DOF mismatches must not go unnoticed: + @test_throws ArgumentError transport_to(StdUniform()^5, μ)(xy) + @test_throws ArgumentError transport_to(StdUniform()^2, μ)(xy) + + # Scalar-variate primary measures work in vcat streams: + f_βs(a) = pushfwd(Mul(abs(a) + 0.5), StdNormal())^2 + μsc = mbind(f_βs, StdExponential(), vcat) + xys = rand(stblrng(), Float64, μsc) + @test xys isa AbstractVector{<:Real} && length(xys) == 3 + @test logdensityof(μsc, xys) ≈ + logdensityof(StdExponential(), xys[1]) + logdensityof(f_βs(xys[1]), xys[2:3]) + ysc = transport_to(StdUniform()^3, μsc)(xys) + @test transport_to(μsc, StdUniform()^3)(ysc) ≈ xys + + # Nested binds evaluate in a single with-rest pass: + μnest = mbind(f_βv, μ, vcat) + xyz = rand(stblrng(), Float64, μnest) + @test length(xyz) == 5 + an, bn = xyz[1:3], xyz[4:5] + @test logdensityof(μnest, xyz) ≈ logdensityof(μ, an) + logdensityof(f_βv(an), bn) + + # Variates that are too long must not go unnoticed: + @test_throws ArgumentError logdensityof(μ, vcat(xy, [0.5])) + + # Products of same-typed unknown-DOF marginals transport type-stably: + P = productmeasure([μ, μ]) + yP = rand(stblrng(), Float64, P) + z = transport_to(StdUniform()^6, P)(yP) + @test z isa AbstractVector{<:Real} && length(z) == 6 + yP_reco, rest = MeasureBase.transport_from_mvstd_with_rest(P, StdUniform(), z) + @test yP_reco isa Vector{<:AbstractVector{Float64}} + @test yP_reco ≈ yP && isempty(rest) + @test logdensityof(P, yP) ≈ logdensityof(μ, yP[1]) + logdensityof(μ, yP[2]) + + # Products with value-dependent marginal sizes, in all marginal + # container flavors: + Pt = productmeasure((μ, μ)) + yt = rand(stblrng(), Float64, Pt) + @test logdensityof(Pt, yt) ≈ logdensityof(μ, yt[1]) + logdensityof(μ, yt[2]) + zt = transport_to(StdUniform()^6, Pt)(yt) + yt_reco = transport_to(Pt, StdUniform()^6)(zt) + @test all(map(≈, yt_reco, yt)) + + Pnt = productmeasure((a = StdNormal(), b = μ)) + ynt = rand(stblrng(), Float64, Pnt) + @test logdensityof(Pnt, ynt) ≈ + logdensityof(StdNormal(), ynt.a) + logdensityof(μ, ynt.b) + znt = transport_to(StdUniform()^4, Pnt)(ynt) + ynt_reco = transport_to(Pnt, StdUniform()^4)(znt) + @test ynt_reco.a ≈ ynt.a && ynt_reco.b ≈ ynt.b + end + + @testset "mbind with merge" begin + αnt = productmeasure((position = StdNormal(),)) + f_βnt(a) = productmeasure(( + noise = pushfwd(Mul(abs(a.position) + 0.5), StdExponential()), + )) + μ = mbind(f_βnt, αnt, merge) + + x = rand(stblrng(), Float64, μ) + @test x isa NamedTuple{(:position, :noise)} + @test logdensityof(μ, x) ≈ + logdensityof(αnt, (position = x.position,)) + + logdensityof(f_βnt(x), (noise = x.noise,)) + + y = transport_to(StdUniform()^2, μ)(x) + @test y isa AbstractVector{<:Real} && length(y) == 2 + x_reco = transport_to(μ, StdUniform()^2)(y) + @test x_reco.position ≈ x.position && x_reco.noise ≈ x.noise + end + + @testset "mbind with Dirac" begin + @test mbind(f_β, MeasureBase.Dirac(1.5)) == asmeasure(f_β(1.5)) + μ = mbind(f_β, MeasureBase.Dirac(1.5), tuple) + ab = rand(stblrng(), Float64, μ) + @test ab[1] == 1.5 + end +end diff --git a/test/combinators/combined.jl b/test/combinators/combined.jl new file mode 100644 index 00000000..9a79f225 --- /dev/null +++ b/test/combinators/combined.jl @@ -0,0 +1,52 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using Random +using StableRNGs: StableRNG +using OneTwoMany: firstarg, secondarg + +using MeasureBase +using MeasureBase: StdExponential, StdLogistic, StdNormal, StdUniform +using MeasureBase: mcombine, productmeasure, transport_to + +@testset "mcombine" begin + stblrng() = StableRNG(789990641) + + α = StdExponential() + β = StdLogistic() + + @testset "combination shortcuts" begin + @test mcombine(firstarg, α, β) === α + @test mcombine(secondarg, α, β) === β + @test mcombine(tuple, α, β) == productmeasure((α, β)) + @test mcombine(vcat, StdNormal()^2, StdNormal()^1) == StdNormal()^3 + @test mcombine(vcat, productmeasure([α, α]), productmeasure([α])) == + productmeasure([α, α, α]) + @test mcombine( + merge, + productmeasure((a = α,)), + productmeasure((b = β,)), + ) == productmeasure((a = α, b = β)) + @test mcombine(tuple, MeasureBase.Dirac(1), MeasureBase.Dirac(2)) == + MeasureBase.Dirac((1, 2)) + end + + @testset "CombinedMeasure" begin + μ = mcombine(Pair, α, β) + @test μ isa MeasureBase.CombinedMeasure + + ab = rand(stblrng(), Float64, μ) + @test ab isa Pair + @test logdensityof(μ, ab) ≈ logdensityof(α, ab.first) + logdensityof(β, ab.second) + + @test MeasureBase.getdof(μ) == 2 + @test MeasureBase.fast_dof(μ) == 2 + @test MeasureBase.insupport(μ, ab) isa MeasureBase.NoFastInsupport + + y = transport_to(StdUniform()^2, μ)(ab) + @test y isa AbstractVector{<:Real} && length(y) == 2 + ab_reco = transport_to(μ, StdUniform()^2)(y) + @test ab_reco.first ≈ ab.first && ab_reco.second ≈ ab.second + end +end 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/combinators/smart_constructors.jl b/test/combinators/smart_constructors.jl new file mode 100644 index 00000000..06d7906b --- /dev/null +++ b/test/combinators/smart_constructors.jl @@ -0,0 +1,151 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using MeasureBase +using MeasureBase: + weightedmeasure, superpose, productmeasure, powermeasure, pushfwd, pullbck, restrict +using MeasureBase: + WeightedMeasure, + SuperpositionMeasure, + ProductMeasure, + PowerMeasure, + PushforwardMeasure, + RestrictedMeasure, + Dirac, + StdNormal, + StdUniform, + StdExponential, + PushfwdRootMeasure, + AdaptRootMeasure +using FillArrays: Fill +using Static: static + +@testset "smart constructors" begin + @testset "powermeasure" begin + @test powermeasure(StdNormal(), ()) === StdNormal() + + @test powermeasure(Dirac(4.2), (3,)) == Dirac(Fill(4.2, 3)) + + wpw = weightedmeasure(0.3, StdNormal())^(2, 3) + @test wpw isa WeightedMeasure + @test wpw.logweight ≈ 6 * 0.3 + @test wpw.base == StdNormal()^(2, 3) + + # Weight pull-out must not depend on the collapsed base type: + wd = weightedmeasure(0.3, Dirac(1.5))^3 + @test wd isa WeightedMeasure + @test wd.logweight ≈ 3 * 0.3 + @test wd.base == Dirac(Fill(1.5, 3)) + + ws = weightedmeasure(static(0.5), StdNormal())^static(4) + @test ws.logweight ≈ 2.0 + end + + @testset "productmeasure" begin + @test productmeasure(Fill(StdUniform(), 3)) == StdUniform()^3 + + @test productmeasure(()) === Dirac(()) + @test productmeasure(NamedTuple()) === Dirac(NamedTuple()) + + @test productmeasure((Dirac(1), Dirac(2))) === Dirac((1, 2)) + @test productmeasure((a = Dirac(1), b = Dirac(2))) === Dirac((a = 1, b = 2)) + @test productmeasure([Dirac(1), Dirac(2)]) == Dirac([1, 2]) + + pt = productmeasure((2.0 * StdNormal(), 3.0 * StdUniform())) + @test pt isa WeightedMeasure + @test exp(pt.logweight) ≈ 6 + @test pt.base == ProductMeasure((StdNormal(), StdUniform())) + + pnt = productmeasure((a = 2.0 * StdNormal(), b = 3.0 * StdUniform())) + @test pnt isa WeightedMeasure + @test exp(pnt.logweight) ≈ 6 + @test pnt.base == ProductMeasure((a = StdNormal(), b = StdUniform())) + + pa = productmeasure([2.0 * StdNormal(), 3.0 * StdNormal()]) + @test pa isa WeightedMeasure + @test exp(pa.logweight) ≈ 6 + @test pa.base == StdNormal()^2 + + @test logdensityof(pt, (0.3, 0.5)) ≈ log(6) + logdensityof(StdNormal(), 0.3) + + @test productmeasure([StdNormal(), StdNormal()]) == StdNormal()^2 + @test productmeasure([StdNormal()^2, StdUniform()^3]) isa ProductMeasure + end + + @testset "superpose" begin + μ, ν = StdNormal(), StdUniform() + + @test superpose(μ) === μ + @test superpose(μ, ν) == SuperpositionMeasure((μ, ν)) + + s2 = superpose(μ, μ) + @test s2 isa WeightedMeasure && exp(s2.logweight) ≈ 2 && s2.base === μ + + s4 = superpose(μ, μ, μ, μ) + @test s4 isa WeightedMeasure && exp(s4.logweight) ≈ 4 + + c = superpose(2.0 * μ, 3.0 * μ) + @test c isa WeightedMeasure && exp(c.logweight) ≈ 5 && c.base === μ + @test exp(superpose(2.0 * μ, μ).logweight) ≈ 3 + @test exp(superpose(μ, 2.0 * μ).logweight) ≈ 3 + @test superpose(2.0 * μ, 3.0 * ν) == SuperpositionMeasure((2.0 * μ, 3.0 * ν)) + + ss = superpose(superpose((μ, ν)), superpose((ν, StdExponential()))) + @test ss.components === (μ, ν, ν, StdExponential()) + @test superpose(superpose((μ, ν)), StdExponential()).components === + (μ, ν, StdExponential()) + @test superpose(StdExponential(), superpose((μ, ν))).components === + (StdExponential(), μ, ν) + + # Merging must not mutate existing superpositions: + sv = superpose(AbstractMeasure[μ, ν]) + sv2 = superpose(sv, StdExponential()) + @test length(sv.components) == 2 && length(sv2.components) == 3 + + # Simplifications must be type stable, so they only happen when + # measure equality is decidable from the measure types: + @inferred superpose(μ, ν) + @inferred superpose(μ, μ) + @inferred superpose(2.0 * μ, 3.0 * μ) + @inferred superpose(2.0 * μ, μ) + @inferred superpose(μ, μ, μ, μ) + @inferred superpose(Dirac(1), Dirac(1)) + @test superpose(Dirac(1), Dirac(1)) isa SuperpositionMeasure + @inferred superpose(2.0 * Dirac(1), 3.0 * Dirac(1)) + @test superpose(2.0 * Dirac(1), 3.0 * Dirac(1)) isa SuperpositionMeasure + + @test superpose(Fill(μ, 4)) == weightedmeasure(log(4), μ) + @test superpose([μ, μ, μ]) == weightedmeasure(log(3), μ) + + @test logdensityof(c, 0.3) ≈ log(5) + logdensityof(μ, 0.3) + end + + @testset "pushfwd" begin + @test pushfwd(identity, StdNormal()) === StdNormal() + @test pushfwd(identity, StdNormal(), PushfwdRootMeasure()) === StdNormal() + + @test pushfwd(sqrt, Dirac(4.0)) === Dirac(2.0) + @test pushfwd(sqrt, Dirac(4.0), PushfwdRootMeasure()) === Dirac(2.0) + + pw = pushfwd(sqrt, 3.0 * StdExponential()) + @test pw isa WeightedMeasure && exp(pw.logweight) ≈ 3 + @test pw.base isa PushforwardMeasure + + pp = pushfwd(exp, pushfwd(sqrt, StdExponential())) + @test pp isa PushforwardMeasure && pp.origin === StdExponential() + + @test pullbck(log, Dirac(4.0)) === Dirac(exp(4.0)) + end + + @testset "restrict" begin + r = restrict(x -> x > 0, StdNormal()) + @test r isa RestrictedMeasure && r.base === StdNormal() + + r2 = restrict(x -> x < 1, r) + @test r2 isa RestrictedMeasure && r2.base === StdNormal() + @test r2.predicate(0.5) && !r2.predicate(-1.0) && !r2.predicate(2.0) + + @test restrict(x -> x > 0)(StdNormal()) isa RestrictedMeasure + end +end diff --git a/test/combinators/superpose.jl b/test/combinators/superpose.jl index ed4c6996..753d5b8c 100644 --- a/test/combinators/superpose.jl +++ b/test/combinators/superpose.jl @@ -1,7 +1,7 @@ using Test using MeasureBase -using MeasureBase: superpose +using MeasureBase: superpose, weightedmeasure, StdNormal @testset "superpose.jl" begin μ = Dirac(0) @@ -14,12 +14,18 @@ using MeasureBase: superpose μs = SuperpositionMeasure([μ, ν]) @test μs isa SuperpositionMeasure{<:AbstractVector{<:AbstractMeasure}} - @test_throws ErrorException density_def(μs, 0) - @test basemeasure(μs).components == - SuperpositionMeasure([CountingBase(), CountingBase()]).components + @test density_def(μs, 0) == 1.0 + @test basemeasure(μs) == weightedmeasure(log(2), CountingBase()) + # Dirac equality is not decidable from types, so no weighted collapse: μ2 = μ + μ - @test μ2 isa WeightedMeasure + @test μ2 isa SuperpositionMeasure @test μ2 == superpose(μ, μ) - @test basemeasure(μ2) == μ + @test density_def(μ2, 0) == 1.0 + + # For singleton measure types equal measures combine into weighted measures: + s2 = StdNormal() + StdNormal() + @test s2 isa WeightedMeasure + @test exp(s2.logweight) ≈ 2 + @test basemeasure(s2) == StdNormal() end diff --git a/test/combinators/transformedmeasure.jl b/test/combinators/transformedmeasure.jl index 28ddbb50..6261c762 100644 --- a/test/combinators/transformedmeasure.jl +++ b/test/combinators/transformedmeasure.jl @@ -152,9 +152,10 @@ using ChangesOfVariables @test rootmeasure(ν) === rootmeasure(μ) # AdaptRootMeasure @test rootmeasure(ν_no_corr) isa PushforwardMeasure # PushfwdRootMeasure - # Test basemeasure - @test basemeasure(ν) isa PushforwardMeasure - @test basemeasure(ν).style isa PushfwdRootMeasure + # Test basemeasure. The base measure of μ is a weighted Lebesgue measure, + # so the weight gets pulled out of the pushforward: + @test basemeasure(ν) isa WeightedMeasure{<:Any,<:PushforwardMeasure} + @test basemeasure(ν).base.style isa PushfwdRootMeasure # Test massof # TODO: mass interface is very incomplete @@ -162,7 +163,7 @@ using ChangesOfVariables # Test rand @test rand(ν) isa Real - @test insupport(ν, rand(ν)) + @test insupport(ν, rand(ν)) != false # Test pullback pb = pullbck(f, ν) 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..a8244b4d --- /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) != false + 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) != false + 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..f3d6c237 --- /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.mintegrate(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/domains.jl b/test/domains.jl new file mode 100644 index 00000000..8d465eb9 --- /dev/null +++ b/test/domains.jl @@ -0,0 +1,65 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using MeasureBase +using MeasureBase: combinesets, setcartprod, setcartpower +using MeasureBase: CombinedSet, CartesianProduct, CartesianPower, ImplicitDomain +using MeasureBase: maybe_in, mdomain, mcombine, mbind, pushfwd +using MeasureBase: StdNormal, StdUniform, StdExponential +using MeasureBase: ℝ, ℤ +using OneTwoMany: firstarg, secondarg +using AffineMaps: Mul + +@testset "domains" begin + @testset "combinesets" begin + @test combinesets(firstarg, ℝ, ℤ) === ℝ + @test combinesets(secondarg, ℝ, ℤ) === ℤ + + s_tuple = combinesets(tuple, ℝ, ℤ) + @test s_tuple isa CartesianProduct + @test (1.5, 2) ∈ s_tuple + @test !((1.5, 2.5) ∈ s_tuple) + + pv = setcartprod([ℝ, ℝ]) + sv = combinesets(vcat, pv, pv) + @test sv isa CartesianProduct + @test [1.0, 2.0, 3.0, 4.0] ∈ sv + + snt = combinesets(merge, setcartprod((a = ℝ,)), setcartprod((b = ℤ,))) + @test snt isa CartesianProduct + @test (a = 1.5, b = 2) ∈ snt + + # One-dimensional powers of equal singleton base sets concatenate: + spw = combinesets(vcat, setcartpower(ℝ, (2,)), setcartpower(ℝ, (3,))) + @test spw isa CartesianPower + @test [1.0, 2.0, 3.0, 4.0, 5.0] ∈ spw + @test combinesets(vcat, setcartpower(ℝ, (2,)), setcartpower(ℤ, (3,))) isa + CombinedSet + + # No specific representation available: + sc = combinesets(vcat, ℝ, setcartpower(ℝ, (2,))) + @test sc isa CombinedSet + @test maybe_in([1.0, 2.0, 3.0], sc) + @test !isempty(sc) + @test_throws ArgumentError [1.0, 2.0, 3.0] ∈ sc + + # Implicit domains combine into the implicit domain of the + # combined measure: + sid = combinesets( + vcat, + ImplicitDomain(StdNormal()^2), + ImplicitDomain(StdUniform()^1), + ) + @test sid isa ImplicitDomain + @test maybe_in([1.0, 2.0, 3.0], sid) + end + + @testset "mdomain of combined measures" begin + f_β(a) = pushfwd(Mul(a[1] + 0.5), StdNormal())^2 + μc = mcombine(vcat, StdNormal()^2, mbind(f_β, StdExponential()^1, vcat)) + @test μc isa MeasureBase.CombinedMeasure + @test mdomain(μc) isa ImplicitDomain + @test maybe_in(rand(Float64, μc), mdomain(μc)) + end +end diff --git a/test/logdensities.jl b/test/logdensities.jl new file mode 100644 index 00000000..89ff253f --- /dev/null +++ b/test/logdensities.jl @@ -0,0 +1,79 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using MeasureBase +using MeasureBase: logdensities, StdNormal, StdUniform +using ArraysOfArrays: VectorOfSimilarVectors, nestedview, flatview +using IrrationalConstants: log2π +import JLArrays +using JLArrays: JLArray + +stdnormal_ld(x) = -(x^2 + log2π) / 2 + +@testset "logdensities" begin + @testset "scalar variates" begin + X = randn(10) + @test @inferred(logdensities(StdNormal(), X)) ≈ stdnormal_ld.(X) + Xm = randn(2, 3) + @test logdensities(StdNormal(), Xm) ≈ stdnormal_ld.(Xm) + end + + @testset "powers with nested variates" begin + m3 = StdNormal()^3 + X = [randn(3) for _ in 1:10] + @test @inferred(logdensities(m3, X)) ≈ [sum(stdnormal_ld, x) for x in X] + @test only(logdensities(m3, [X[1]])) ≈ logdensityof(m3, X[1]) + + m23 = StdNormal()^(2, 3) + X23 = [randn(2, 3) for _ in 1:5] + @test logdensities(m23, X23) ≈ [sum(stdnormal_ld, x) for x in X23] + + mpp = (StdNormal()^(2, 3))^4 + Xpp = [[randn(2, 3) for _ in 1:4] for _ in 1:6] + @test logdensities(mpp, Xpp) ≈ [sum(x -> sum(stdnormal_ld, x), xs) for xs in Xpp] + end + + @testset "powers with flat variate storage" begin + m3 = StdNormal()^3 + X = VectorOfSimilarVectors(randn(3, 10)) + @test @inferred(logdensities(m3, X)) ≈ + vec(sum(stdnormal_ld.(flatview(X)), dims = 1)) + + # Power structure may be stored flattened out within each point: + mpp = (StdNormal()^(2, 3))^4 + Xpp = nestedview(randn(2, 3, 4, 7), 3) + @test logdensities(mpp, Xpp) ≈ [sum(stdnormal_ld, x) for x in Xpp] + end + + @testset "non-scalar-variate fallback" begin + mprod = productmeasure((StdUniform(), StdNormal())) + X = [(rand(), randn()) for _ in 1:5] + @test logdensities(mprod, X) ≈ logdensityof.(Ref(mprod), X) + end + + @testset "size mismatch" begin + @test_throws ArgumentError logdensities(StdNormal()^3, [randn(3), randn(2)]) + @test_throws ArgumentError logdensities( + StdNormal()^3, + VectorOfSimilarVectors(randn(2, 5)), + ) + end + + @testset "GPU array semantics" begin + JLArrays.allowscalar(false) + + X = JLArray(randn(100)) + ld = logdensities(StdNormal(), X) + @test ld isa JLArray + @test Array(ld) ≈ stdnormal_ld.(Array(X)) + + Xb = VectorOfSimilarVectors(JLArray(randn(3, 50))) + ldb = logdensities(StdNormal()^3, Xb) + @test ldb isa JLArray + @test Array(ldb) ≈ vec(sum(stdnormal_ld.(Array(flatview(Xb))), dims = 1)) + + xj = JLArray(randn(10)) + @test logdensityof(StdNormal()^10, xj) ≈ logdensityof(StdNormal()^10, Array(xj)) + end +end diff --git a/test/measure_operators.jl b/test/measure_operators.jl new file mode 100644 index 00000000..1530f191 --- /dev/null +++ b/test/measure_operators.jl @@ -0,0 +1,24 @@ +using Test + +using MeasureBase: AbstractMeasure +using MeasureBase: StdExponential, StdLogistic, StdNormal, StdUniform +using MeasureBase: pushfwd, pullbck, mbind, productmeasure +using MeasureBase: mintegrate, mintegrate_exp, density_rel, logdensity_rel +using MeasureBase.MeasureOperators: ⋄, ⊙, ▷, ⊗, ∫, ∫exp, 𝒹, log𝒹 + +@testset "MeasureOperators" begin + μ = StdExponential() + ν = StdUniform() + k(σ) = pushfwd(x -> σ * x, StdNormal()) + μs = (StdExponential(), StdLogistic(), StdUniform()) + f = sqrt + + @test @inferred(f ⋄ μ) == pushfwd(f, μ) + @test @inferred(ν ⊙ f) == pullbck(f, ν) + @test @inferred(μ ▷ k) == mbind(k, μ) + @test @inferred(⊗(μs...)) == productmeasure(μs) + @test @inferred(∫(f, μ)) == mintegrate(f, μ) + @test @inferred(∫exp(f, μ)) == mintegrate_exp(f, μ) + @test @inferred(𝒹(ν, μ)) == density_rel(ν, μ) + @test @inferred(log𝒹(ν, μ)) == logdensity_rel(ν, μ) +end diff --git a/test/runtests.jl b/test/runtests.jl index c2f63c4e..b6073d76 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -9,18 +9,30 @@ using MeasureBase: test_interface, test_smf include("test_aqua.jl") -include("static.jl") include("test_primitive.jl") include("test_standard.jl") include("test_basics.jl") include("getdof.jl") +include("logdensities.jl") include("transport.jl") include("smf.jl") +include("domains.jl") +include("test_mooncake.jl") + +include("measure_operators.jl") + +include("combinators/smart_constructors.jl") include("combinators/weighted.jl") +include("combinators/superpose.jl") include("combinators/transformedmeasure.jl") +include("combinators/reshape.jl") include("combinators/implicitlymapped.jl") +include("combinators/combined.jl") +include("combinators/bind.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..f683e2c1 100644 --- a/test/test_aqua.jl +++ b/test/test_aqua.jl @@ -4,10 +4,14 @@ import Test import Aqua import MeasureBase -#Test.@testset "Package ambiguities" begin -# Test.@test isempty(Test.detect_ambiguities(MeasureBase)) -#end # testset +Test.@testset "Package ambiguities" begin + Test.@test isempty(Test.detect_ambiguities(MeasureBase, recursive = true)) +end # testset Test.@testset "Aqua tests" begin - Aqua.test_all(MeasureBase, ambiguities = false) + Aqua.test_all( + MeasureBase, + # 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..8dbdbe4d 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -1,4 +1,4 @@ -d = ∫exp(x -> -x^2, Lebesgue(ℝ)) +d = mintegrate_exp(x -> -x^2, Lebesgue(ℝ)) # function draw2(μ) # x = rand(μ) @@ -75,11 +75,6 @@ testbroken_measures = [ end end -# @testset "TransitionKernel" begin -# κ = MeasureBase.kernel(MeasureBase.Dirac, identity) -# @test rand(κ(1.1)) == 1.1 -# end - @testset "SpikeMixture" begin @test rand(SpikeMixture(Dirac(0), 0.5)) == 0 @test rand(SpikeMixture(Dirac(1), 1.0)) == 1 @@ -120,11 +115,12 @@ 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(ℝ)) +NormalMeasure() = mintegrate_exp(x -> -0.5x^2, Lebesgue(ℝ)) @testset "Half" begin HalfNormal() = Half(NormalMeasure()) @@ -135,12 +131,10 @@ end @testset "Likelihood" begin ℓ = Likelihood(3) do (μ,) - ∫exp(Lebesgue(ℝ)) do x + mintegrate_exp(Lebesgue(ℝ)) do x -(x - μ)^2 end end - - @inferred logdensityof(Lebesgue() ⊙ ℓ, 2.0) end # @testset "Likelihood" begin @@ -196,7 +190,7 @@ end f2 = x -> sqrt(abs(sum(x))) f3 = x -> 2 * sum(x) f4 = x -> sum(sqrt.(abs.(x))) - m = @inferred ∫exp(f1, ∫exp(f2, ∫exp(f3, ∫exp(f4, StdUniform()^3)))) + m = @inferred mintegrate_exp(f1, mintegrate_exp(f2, mintegrate_exp(f3, mintegrate_exp(f4, StdUniform()^3)))) for x in [Float32[0.7, 0.2, 0.5], Float32[-0.7, 0.2, 0.5]] @test @inferred(logdensityof(m, x)) isa Float32 @@ -205,6 +199,31 @@ end end end +@testset "logdensityof_with_rest" begin + StdNormal = MeasureBase.StdNormal + x = [0.3, 0.7, 0.2, 0.9, 0.5] + + # Scalar variates consume one stream element: + @test MeasureBase.mspace_elsize(StdNormal()) == () + ℓ, a, x_rest = MeasureBase.logdensityof_with_rest(StdNormal(), x) + @test a == 0.3 && length(x_rest) == 4 + @test ℓ ≈ logdensityof(StdNormal(), 0.3) + + # Vector variates: + @test MeasureBase.mspace_elsize(StdNormal()^2) == (2,) + ℓ, a, x_rest = MeasureBase.logdensityof_with_rest(StdNormal()^2, x) + @test a == [0.3, 0.7] && length(x_rest) == 3 + @test ℓ ≈ logdensityof(StdNormal()^2, [0.3, 0.7]) + + # Multi-rank variates are consumed in flattened form and reshaped: + @test MeasureBase.mspace_elsize(StdNormal()^(2, 2)) == (2, 2) + ℓ, a, x_rest = MeasureBase.logdensityof_with_rest(StdNormal()^(2, 2), x) + @test a == [0.3 0.2; 0.7 0.9] && length(x_rest) == 1 + @test ℓ ≈ logdensityof(StdNormal()^(2, 2), a) + + @test MeasureBase.mspace_elsize(Dirac([1, 2])) == (2,) +end + @testset "logdensity_rel" begin @test logdensity_rel(Dirac(0.0) + Lebesgue(), Dirac(1.0), 0.0) == Inf @test logdensity_rel(Dirac(0.0) + Lebesgue(), Dirac(1.0), 1.0) == -Inf @@ -223,18 +242,41 @@ end @test logdensity_rel(Lebesgue(), Dirac(0.0) + Lebesgue(), 1.0) == 0.0 @test isnan(logdensity_rel(Dirac(0), Dirac(1), 2)) + + # The generic implementation descends the base measure chains of both + # measures in lockstep, type-stably and with symbolic cancellation of + # shared chain suffixes: + let μW = MeasureBase.weightedmeasure(0.7, MeasureBase.StdNormal()) + StdNormal, StdUniform, StdExponential = + MeasureBase.StdNormal, MeasureBase.StdUniform, MeasureBase.StdExponential + @test @inferred(logdensity_rel(μW, StdNormal(), 0.5)) ≈ 0.7 + @test @inferred(logdensity_rel(StdNormal(), μW, 0.5)) ≈ -0.7 + @test @inferred(logdensity_rel(StdNormal(), StdUniform(), 0.5)) ≈ + logdensityof(StdNormal(), 0.5) + p1 = productmeasure((StdNormal(), StdExponential())) + p2 = productmeasure((StdUniform(), StdExponential())) + @test @inferred(logdensity_rel(p1, p2, (0.5, 0.5))) ≈ + logdensityof(StdNormal(), 0.5) + + # Incompatible root measures result in an informative exception: + @test_throws ArgumentError logdensity_rel( + productmeasure((StdNormal(),)), + StdNormal()^1, + (0.5,), + ) + end end @testset "Density measures and Radon-Nikodym" begin x = randn() f(x) = x^2 - @test log(𝒹(∫exp(f, Lebesgue()), Lebesgue())(x)) ≈ f(x) + @test log(density_rel(mintegrate_exp(f, Lebesgue()), Lebesgue())(x)) ≈ f(x) - let f = 𝒹(∫exp(x -> x^2, Lebesgue()), Lebesgue()) + let f = density_rel(mintegrate_exp(x -> x^2, Lebesgue()), Lebesgue()) @test log(f(x)) ≈ x^2 end - let f = log𝒹(∫exp(x -> x^2, NormalMeasure()), NormalMeasure()) + let f = logdensity_rel(mintegrate_exp(x -> x^2, NormalMeasure()), NormalMeasure()) @test f(x) ≈ x^2 end end 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