diff --git a/Project.toml b/Project.toml index 44f89f80..7ea88557 100644 --- a/Project.toml +++ b/Project.toml @@ -1,9 +1,12 @@ name = "MeasureBase" uuid = "fa1605e6-acd5-459c-a1e6-7e635759db14" -version = "0.14.12" +version = "0.15.0" authors = ["Chad Scherrer ", "Oliver Schulz ", "contributors"] [deps] +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" +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 +14,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 +24,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 +32,54 @@ 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" +StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775" [weakdeps] +AffineMaps = "2c83c9a8-abf5-4329-a0d7-deffaf474661" 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] +MeasureBaseAffineMapsExt = "AffineMaps" 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] +Adapt = "3.7, 4" +AffineMaps = "0.3" +ArgCheck = "1, 2" +ArraysOfArrays = "1.3" 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 +88,23 @@ 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" +PropertyFunctions = "0.3" Random = "1" +Reactant = "0.2" Reexport = "1" -SpecialFunctions = "2" +SpecialFunctions = "2.1.4" Static = "0.8, 1" StaticArrays = "1.5" +StaticThings = "0.3" Statistics = "1" +StructArrays = "0.6.18, 0.7" +StatsBase = "0.33, 0.34" +StatsFuns = "0.9, 1, 2" Test = "1" Tricks = "0.1" julia = "1.10" diff --git a/docs/make.jl b/docs/make.jl index 3407cf93..96f147eb 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -12,7 +12,7 @@ DocMeta.setdocmeta!(MeasureBase, :DocTestSetup, :(using MeasureBase); recursive makedocs( sitename = "MeasureBase", - modules = [MeasureBase], + modules = [MeasureBase, MeasureBase.MeasureOperators], format = Documenter.HTML( prettyurls = !("local" in ARGS), canonical = "https://juliamath.github.io/MeasureBase.jl/stable/", diff --git a/docs/src/api.md b/docs/src/api.md index 83c27f70..fc4a8a33 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -21,6 +21,6 @@ Order = [:macro, :function] # Documentation ```@autodocs -Modules = [MeasureBase] +Modules = [MeasureBase, MeasureBase.MeasureOperators] Order = [:module, :type, :constant, :macro, :function] ``` diff --git a/ext/MeasureBaseAffineMapsExt.jl b/ext/MeasureBaseAffineMapsExt.jl new file mode 100644 index 00000000..1861b273 --- /dev/null +++ b/ext/MeasureBaseAffineMapsExt.jl @@ -0,0 +1,45 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +module MeasureBaseAffineMapsExt + +using MeasureBase +using MeasureBase: PushforwardMeasure, AdaptRootMeasure, PushfwdRootMeasure +using MeasureBase: StaticInteger +using AffineMaps: AbstractAffineMap +using ChangesOfVariables: with_logabsdet_jacobian + +# Affine maps treat matrices as batches of column vectors, so flat batches +# of vector variates are applied as `(n, :)` matrices: +function MeasureBase._apply_generic(f::AbstractAffineMap, X::AbstractArray, ::StaticInteger{1}) + _columns_back(f(_as_columns(X)), X) +end + +@inline _as_columns(x::AbstractVector) = x +@inline _as_columns(X::AbstractArray) = reshape(X, (size(X, 1), :)) +@inline _columns_back(y::AbstractVector, ::AbstractVector) = y +@inline _columns_back(Y::AbstractMatrix, X::AbstractArray) = reshape(Y, (size(Y, 1), Base.tail(size(X))...)) + +const _AffinePushfwd{M,S} = PushforwardMeasure{<:AbstractAffineMap,<:AbstractAffineMap,M,S} + +# Densities of affine pushforwards of vector variates use the per-column +# log-abs-det-Jacobians of the inverse map: +for (bhead, head) in [(:batched_logdensityof_impl, :logdensityof_impl), (:batched_logdensity_def, :logdensity_def)] + @eval function MeasureBase.$bhead(ν::_AffinePushfwd{M,<:AdaptRootMeasure}, Y::AbstractArray) where {M} + _affine_pushfwd_ld(MeasureBase.$head, ν, Y, MeasureBase._static_ndims(ν)) + end + @eval function MeasureBase.$bhead(ν::_AffinePushfwd{M,<:PushfwdRootMeasure}, Y::AbstractArray) where {M} + MeasureBase._batched_kernel(MeasureBase.$head, ν.origin, MeasureBase._apply_batched(ν.finv, Y, MeasureBase._static_ndims(ν))) + end +end + +function _affine_pushfwd_ld(f::F, ν::PushforwardMeasure, Y::AbstractArray, ::StaticInteger{1}) where {F} + X2, ladj2 = with_logabsdet_jacobian(ν.finv, _as_columns(Y)) + ℓ = MeasureBase._batched_kernel(f, ν.origin, _columns_back(X2, Y)) + return MeasureBase._lazy_combine_ladj(ℓ, _ladj_back(ladj2, Y)) +end +_affine_pushfwd_ld(f::F, ν::PushforwardMeasure, Y::AbstractArray, k) where {F} = MeasureBase._default_batched_kernel(f, ν, Y, k) + +@inline _ladj_back(ladj::Number, ::AbstractVector) = ladj +@inline _ladj_back(ladj::AbstractMatrix, Y::AbstractArray) = reshape(ladj, Base.tail(size(Y))) + +end # module MeasureBaseAffineMapsExt diff --git a/ext/MeasureBaseChainRulesCoreExt.jl b/ext/MeasureBaseChainRulesCoreExt.jl index 57ed25fa..ddac715d 100644 --- a/ext/MeasureBaseChainRulesCoreExt.jl +++ b/ext/MeasureBaseChainRulesCoreExt.jl @@ -3,7 +3,7 @@ module MeasureBaseChainRulesCoreExt using MeasureBase -using ChainRulesCore: NoTangent, ZeroTangent +using ChainRulesCore: NoTangent, ZeroTangent, @non_differentiable import ChainRulesCore # = utils ==================================================================== @@ -11,14 +11,76 @@ 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 +# Variate layouts are type-level information: +using MeasureBase: fixed_stream_size, _pushfwd_varsize, _unit_bounds, _prob_floor +@non_differentiable fixed_stream_size(::Type) +@non_differentiable _pushfwd_varsize(f, μ) +@non_differentiable _unit_bounds(p) +@non_differentiable _prob_floor(p) + +using MeasureBase: check_dof, require_insupport, checked_arg, _checksupport @inline function ChainRulesCore.rrule(::typeof(_checksupport), cond, result) y = _checksupport(cond, result) @@ -33,22 +95,10 @@ function ChainRulesCore.rrule(::typeof(require_insupport), μ, x) return require_insupport(μ, x), _require_insupport_pullback end -_origin_depth_pullback(ΔΩ) = NoTangent(), NoTangent() -ChainRulesCore.rrule(::typeof(_origin_depth), ν) = _origin_depth(ν), _origin_depth_pullback - _check_dof_pullback(ΔΩ) = NoTangent(), NoTangent(), NoTangent() ChainRulesCore.rrule(::typeof(check_dof), ν, μ) = check_dof(ν, μ), _check_dof_pullback _checked_arg_pullback(ΔΩ) = NoTangent(), NoTangent(), ΔΩ ChainRulesCore.rrule(::typeof(checked_arg), ν, x) = checked_arg(ν, x), _checked_arg_pullback -# = return type inference ==================================================== - -using MeasureBase: logdensityof_rt - -_logdensityof_rt_pullback(::Any) = (NoTangent(), NoTangent(), ZeroTangent()) -function ChainRulesCore.rrule(::typeof(logdensityof_rt), target, v) - logdensityof_rt(target, v), _logdensityof_rt_pullback -end - end # module MeasureBaseChainRulesCoreExt diff --git a/ext/MeasureBaseDistributionsChainRulesCoreExt.jl b/ext/MeasureBaseDistributionsChainRulesCoreExt.jl new file mode 100644 index 00000000..02fa145d --- /dev/null +++ b/ext/MeasureBaseDistributionsChainRulesCoreExt.jl @@ -0,0 +1,43 @@ +# 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()) +using MeasureBase: _gamma_cdf, _gamma_quantile, _beta_cdf, _beta_quantile, _gamma_logpdf, _beta_logpdf + +# Derivatives with respect to the variate resp. probability argument of the +# regularized incomplete gamma and beta functions and their inverses: +function ChainRulesCore.rrule(::typeof(_gamma_cdf), α::Real, x::Real) + y = _gamma_cdf(α, x) + dy_dx = exp(_gamma_logpdf(α, x)) + return y, ȳ -> (NoTangent(), NoTangent(), dy_dx * ȳ) +end +function ChainRulesCore.rrule(::typeof(_gamma_quantile), α::Real, p::Real) + x = _gamma_quantile(α, p) + dx_dp = exp(-_gamma_logpdf(α, x)) + return x, x̄ -> (NoTangent(), NoTangent(), dx_dp * x̄) +end +function ChainRulesCore.rrule(::typeof(_beta_cdf), α::Real, β::Real, x::Real) + y = _beta_cdf(α, β, x) + dy_dx = exp(_beta_logpdf(α, β, x)) + return y, ȳ -> (NoTangent(), NoTangent(), NoTangent(), dy_dx * ȳ) +end +function ChainRulesCore.rrule(::typeof(_beta_quantile), α::Real, β::Real, p::Real) + x = _beta_quantile(α, β, p) + dx_dp = exp(-_beta_logpdf(α, β, x)) + return x, x̄ -> (NoTangent(), NoTangent(), NoTangent(), dx_dp * x̄) +end + +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..3733c025 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl @@ -0,0 +1,75 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +module MeasureBaseDistributionsExt + +using LinearAlgebra: Diagonal, Cholesky, LowerTriangular, UpperTriangular, diag, dot, cholesky +import Adapt + +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: _unit_interior +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_to_std, transport_from_std +using MeasureBase: batched_transport_to_std, batched_transport_from_std +using MeasureBase: Reshape +using MeasureBase: convert_realtype, _fwddiff, @_adignore +import MeasureBase: + _dist_params_numtype, _trafo_logcdf_impl, _trafo_logccdf_impl, + _trafo_quantile_impl, _trafo_cquantile_impl, _dist_quantile, _dist_cquantile +using MeasureBase: _pushfront, _pushback, _dropfront, _dropback, _rev_cumsum, _exp_cumsum_log +using MeasureBase: _gamma_cdf, _gamma_quantile, _beta_cdf, _beta_quantile, _gamma_logpdf, _beta_logpdf, _dualtag + +import Distributions +using Distributions: Distribution, VariateForm, ValueSupport, ContinuousDistribution +using Distributions: Univariate, Multivariate, ArrayLikeVariate, Continuous, Discrete +using Distributions: Uniform, Exponential, Logistic, Normal +using Distributions: Cauchy, Laplace, LogNormal, Weibull, Gamma, Poisson, Bernoulli +using Distributions: MvNormal, AbstractMvNormal, Beta, Dirichlet +using Distributions: ReshapedDistribution, AbstractMixtureModel + +import Statistics +import StatsBase +import StatsFuns +import PDMats + +using IrrationalConstants: log2π, invsqrt2π +using LogExpFunctions: logistic, log1pexp +using SpecialFunctions: loggamma, logbeta, gamma_inc, gamma_inc_inv, beta_inc, beta_inc_inv + +using HeterogeneousComputing: real_numtype, GenContext, get_rng, get_precision, get_compute_unit, CPUnit, AbstractComputeUnit + +using Static: True, False, StaticInt, static, dynamic +using StaticThings: SizeLike, 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("families.jl") +include("standardmv.jl") +include("product.jl") +include("reshaped.jl") +include("mixture.jl") +include("dirichlet.jl") +include("multivariate.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..b156d4a7 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/dirichlet.jl @@ -0,0 +1,8 @@ +# 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) + +@inline MeasureBase.preferred_stdmeasure(::Type{<:Dirichlet}) = StdUniform diff --git a/ext/MeasureBaseDistributionsExt/dist_vartransform.jl b/ext/MeasureBaseDistributionsExt/dist_vartransform.jl new file mode 100644 index 00000000..b3e45d2b --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/dist_vartransform.jl @@ -0,0 +1,34 @@ +# 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, ν)) + +# Disambiguation between the type forms of standard measures and distributions: +function MeasureBase.transport_to(::Type{NU}, ::Type{MU}) where {NU<:_AnyStdDistribution,MU<:_AnyStdDistribution} + _throw_two_std_types() +end +function MeasureBase.transport_to(::Type{NU}, ::Type{MU}) where {NU<:StdMeasure,MU<:_AnyStdDistribution} + _throw_two_std_types() +end +function MeasureBase.transport_to(::Type{NU}, ::Type{MU}) where {NU<:_AnyStdDistribution,MU<:StdMeasure} + _throw_two_std_types() +end +function _throw_two_std_types() + throw( + ArgumentError( + "Can't construct a transport function between the types of two standard measures, need a measure instance on one side", + ), + ) +end diff --git a/ext/MeasureBaseDistributionsExt/distribution_measure.jl b/ext/MeasureBaseDistributionsExt/distribution_measure.jl new file mode 100644 index 00000000..acf9fa53 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/distribution_measure.jl @@ -0,0 +1,116 @@ +# 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) + + +# Distributions' samplers run on the CPU, variates on other compute units +# are generated from standard variates via the transports: +MeasureBase.rand_impl(ctx::GenContext, m::DistributionMeasure) = _dist_rand(ctx, m, get_compute_unit(ctx)) +MeasureBase.batched_rand_impl(ctx::GenContext, m::DistributionMeasure, sz::SizeLike) = _dist_batched_rand(ctx, m, sz, get_compute_unit(ctx)) + +_dist_rand(ctx::GenContext, m::DistributionMeasure, ::CPUnit) = + convert_realtype(get_precision(ctx), rand(get_rng(ctx), m.obj)) +_dist_rand(ctx::GenContext, m::DistributionMeasure, ::AbstractComputeUnit) = + MeasureBase._rand_default(ctx, m, (), MeasureBase._NoRandImpl()) +_dist_batched_rand(ctx::GenContext, m::DistributionMeasure, sz::SizeLike, ::CPUnit) = + _flat_powrand(get_rng(ctx), get_precision(ctx), m.obj, asnonstatic(sz)) +_dist_batched_rand(ctx::GenContext, m::DistributionMeasure, sz::SizeLike, ::AbstractComputeUnit) = + MeasureBase._rand_default(ctx, m, sz, MeasureBase._NoRandImpl()) + +# A single variate for zero batch dimensions, flat batches otherwise: +_flat_powrand(rng::AbstractRNG, ::Type{T}, d::Distribution, ::Tuple{}) where {T<:Real} = convert_realtype(T, rand(rng, d)) +_flat_powrand(rng::AbstractRNG, ::Type{T}, d::Distribution, sz::Dims) where {T<:Real} = _flat_powrand_batch(rng, T, d, sz) + +function _flat_powrand_batch(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_batch(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_batch(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_batch(rng::AbstractRNG, ::Type{T}, d::Distribution, sz::Dims) where {T<:Real} + flatview(ArrayOfSimilarArrays(convert_realtype(T, rand(rng, d, sz)))) +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) + +# Distributions evaluate flat batches of array variates (the trailing +# dimensions are batch dimensions) directly, univariate wrappers broadcast +# their point kernels: +for (bhead, phead) in ((:batched_logdensityof_impl, :logdensityof_impl), (:batched_logdensity_def, :logdensity_def)) + @eval function MeasureBase.$bhead(m::DistributionMeasure{<:ArrayLikeVariate{N}}, X::AbstractArray) where {N} + Distributions.logpdf(m.obj, X) + end + @eval function MeasureBase.$bhead(m::DistributionMeasure{<:ArrayLikeVariate{0}}, X::AbstractArray) + MeasureBase._scalar_kernel_broadcast(MeasureBase.$phead, m, X) + end +end +@inline MeasureBase.unsafe_logdensityof(m::DistributionMeasure, x) = DensityInterface.logdensityof(m.obj, x) +@inline MeasureBase.insupport(m::DistributionMeasure, x) = Distributions.insupport(m.obj, x) & _finite_variate(m.obj, x) +# Infinite values lie outside the support of univariate distributions, +# where Distributions may evaluate to NaN: +@inline _finite_variate(::Distribution{Univariate}, x) = isfinite(x) +@inline _finite_variate(::Distribution, x) = true + +@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(d::Distribution) = MeasureBase.NoMSpaceElementSize{typeof(d)}() +@inline MeasureBase.mspace_flatsize(d::Distribution) = MeasureBase.NoMSpaceElementSize{typeof(d)}() +@inline MeasureBase.mspace_elsize(d::Distribution{Univariate}) = () +@inline MeasureBase.mspace_elsize(d::Distribution{<:ArrayLikeVariate}) = size(d) +@inline MeasureBase.mspace_flatsize(d::Distribution{Univariate}) = () +@inline MeasureBase.mspace_flatsize(d::Distribution{<:ArrayLikeVariate}) = size(d) +@inline MeasureBase.mspace_elsize(m::DistributionMeasure) = MeasureBase.mspace_elsize(m.obj) +@inline MeasureBase.mspace_flatsize(m::DistributionMeasure) = MeasureBase.mspace_flatsize(m.obj) +@inline MeasureBase.mspace_flatsize(::Type{<:Distribution{Univariate}}) = () +@inline MeasureBase.mspace_ndims(::Type{<:Distribution{<:ArrayLikeVariate{N}}}) where {N} = N +@inline MeasureBase.mspace_ndims(::Type{AsMeasure{D}}) where {D<:Distribution} = MeasureBase.mspace_ndims(D) +@inline MeasureBase.mspace_flatsize(::Type{AsMeasure{D}}) where {D<:Distribution} = MeasureBase.mspace_flatsize(D) + +@inline MeasureBase.preferred_stdmeasure(::Type{AsMeasure{D}}) where {D<:Distribution} = MeasureBase.preferred_stdmeasure(D) + +@inline MeasureBase.getdof(m::DistributionMeasure{<:ArrayLikeVariate{0}}) = 1 + +# Delegate transport to the wrapped distribution: +@inline MeasureBase.transport_to_std(::Type{S}, m::DistributionMeasure, x) where {S<:StdMeasure} = + MeasureBase.transport_to_std(S, m.obj, x) +@inline MeasureBase.transport_from_std(::Type{S}, m::DistributionMeasure, z) where {S<:StdMeasure} = + MeasureBase.transport_from_std(S, m.obj, z) +@inline MeasureBase.batched_transport_to_std(::Type{S}, m::DistributionMeasure, X::AbstractArray) where {S<:StdMeasure} = + MeasureBase.batched_transport_to_std(S, m.obj, X) +@inline MeasureBase.batched_transport_from_std(::Type{S}, m::DistributionMeasure, Z::AbstractArray) where {S<:StdMeasure} = + MeasureBase.batched_transport_from_std(S, m.obj, Z) + +@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/families.jl b/ext/MeasureBaseDistributionsExt/families.jl new file mode 100644 index 00000000..54a29633 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/families.jl @@ -0,0 +1,136 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +# Densities and standard transports of the main distribution families, +# implemented as plain arithmetic on the parameters without branches or +# foreign calls, so that the batched kernels of the wrapped measures run on +# devices and in traced code. Distributions' own implementations remain in +# use for other families. + +const _Families = Union{Normal,Uniform,Exponential,Logistic,Cauchy,Laplace,LogNormal,Weibull,Gamma,Beta,Poisson,Bernoulli} + +# Densities relative to the base measures (Lebesgue resp. counting +# measure), support checks are separate masks. The density formulas must +# not throw outside of the support, where their results are masked: +@inline MeasureBase.logdensity_def(m::AsMeasure{<:_Families}, x) = _family_logd(m.obj, x) +@inline MeasureBase.unsafe_logdensityof(m::AsMeasure{<:_Families}, x) = _family_logd(m.obj, x) +@inline MeasureBase.insupport(m::AsMeasure{<:_Families}, x) = _family_insupport(m.obj, x) & _finite_variate(m.obj, x) + +# `c * log(y)`, zero for `c == 0` also where `y == 0`: +@inline _clog(c, y) = ifelse(iszero(c), zero(c * log(one(y))), c * log(y)) + +@inline function _family_logd(d::Normal, x) + z = (x - d.μ) / d.σ + -z * z / 2 - log(d.σ) - log2π / 2 +end +@inline _family_insupport(::Normal, x) = true + +@inline _family_logd(d::Uniform, x) = -log(d.b - d.a) + zero(x) +@inline _family_insupport(d::Uniform, x) = (d.a <= x) & (x <= d.b) + +@inline _family_logd(d::Exponential, x) = -x / d.θ - log(d.θ) +@inline _family_insupport(::Exponential, x) = x >= 0 + +@inline function _family_logd(d::Logistic, x) + z = (x - d.μ) / d.θ + -z - 2 * log1pexp(-z) - log(d.θ) +end +@inline _family_insupport(::Logistic, x) = true + +@inline function _family_logd(d::Cauchy, x) + z = (x - d.μ) / d.σ + -log1p(z * z) - log(π * d.σ) +end +@inline _family_insupport(::Cauchy, x) = true + +@inline _family_logd(d::Laplace, x) = -abs((x - d.μ) / d.θ) - log(2 * d.θ) +@inline _family_insupport(::Laplace, x) = true + +@inline function _family_logd(d::LogNormal, x) + lx = log(abs(x)) + z = (lx - d.μ) / d.σ + ℓ = -z * z / 2 - log(d.σ) - log2π / 2 - lx + ifelse(x > 0, ℓ, oftype(ℓ, -Inf)) +end +@inline _family_insupport(::LogNormal, x) = x >= 0 + +@inline function _family_logd(d::Weibull, x) + xθ = abs(x / d.θ) + ℓ = log(d.α / d.θ) + _clog(d.α - 1, xθ) - xθ^d.α + ifelse(isinf(xθ), oftype(ℓ, -Inf), ℓ) +end +@inline _family_insupport(::Weibull, x) = x >= 0 + +@inline function _family_logd(d::Gamma, x) + ℓ = _clog(d.α - 1, abs(x)) - x / d.θ - loggamma(d.α) - d.α * log(d.θ) + ifelse(isinf(x), oftype(ℓ, -Inf), ℓ) +end +@inline _family_insupport(::Gamma, x) = x >= 0 + +@inline function _family_logd(d::Beta, x) + _clog(d.α - 1, abs(x)) + _clog(d.β - 1, abs(1 - x)) - logbeta(d.α, d.β) +end +@inline _family_insupport(::Beta, x) = (0 <= x) & (x <= 1) + +@inline _family_logd(d::Poisson, x) = _clog(x, d.λ) - d.λ - loggamma(abs(x) + 1) +@inline _family_insupport(::Poisson, x) = (x >= 0) & (x == floor(x)) + +@inline _family_logd(d::Bernoulli, x) = ifelse(x == 1, log(d.p), log1p(-d.p)) +@inline _family_insupport(::Bernoulli, x) = (x == 0) | (x == 1) + + +# Standard transports of the non-affine families: Cauchy, Laplace, Gamma +# and Beta pivot on the uniform measure, the log-normal and Weibull +# families on the normal resp. exponential measure. + +# The regularized incomplete gamma and beta functions and their inverses +# (from SpecialFunctions), with derivatives with respect to the variate +# resp. probability argument provided by the autodiff extensions: +@inline MeasureBase._gamma_cdf(α, x) = MeasureBase._gamma_cdf_impl(MeasureBase._dualtag(α, x), α, x) +@inline MeasureBase._gamma_quantile(α, p) = MeasureBase._gamma_quantile_impl(MeasureBase._dualtag(α, p), α, p) +@inline MeasureBase._beta_cdf(α, β, x) = MeasureBase._beta_cdf_impl(MeasureBase._dualtag(α, β, x), α, β, x) +@inline MeasureBase._beta_quantile(α, β, p) = MeasureBase._beta_quantile_impl(MeasureBase._dualtag(α, β, p), α, β, p) +@inline MeasureBase._gamma_cdf_impl(::Type{Nothing}, α, x) = first(gamma_inc(α, x)) +# The complementary probability is formed in the common float type, as +# `gamma_inc_inv` requires `p + q == 1` exactly: +@inline function MeasureBase._gamma_quantile_impl(::Type{Nothing}, α, p) + T = float(promote_type(typeof(α), typeof(p))) + pp = convert(T, p) + gamma_inc_inv(convert(T, α), pp, one(T) - pp) +end +@inline MeasureBase._beta_cdf_impl(::Type{Nothing}, α, β, x) = first(beta_inc(α, β, x)) +@inline MeasureBase._beta_quantile_impl(::Type{Nothing}, α, β, p) = first(beta_inc_inv(α, β, p)) +@inline MeasureBase._gamma_logpdf(α, x) = _clog(α - 1, x) - x - loggamma(α) +@inline MeasureBase._beta_logpdf(α, β, x) = _clog(α - 1, x) + _clog(β - 1, 1 - x) - logbeta(α, β) + +@inline MeasureBase.preferred_stdmeasure(::Type{<:Cauchy}) = StdUniform +@inline MeasureBase.transport_to_std(::Type{StdUniform}, d::Cauchy, x) = 1 // 2 + atan((x - d.μ) / d.σ) / π +@inline MeasureBase.transport_from_std(::Type{StdUniform}, d::Cauchy, p) = + _nan_outside(StdUniform, p, muladd(d.σ, tan(π * (_unit_interior(p) - 1 // 2)), d.μ)) + +@inline MeasureBase.preferred_stdmeasure(::Type{<:Laplace}) = StdUniform +@inline function MeasureBase.transport_to_std(::Type{StdUniform}, d::Laplace, x) + z = (x - d.μ) / d.θ + ifelse(z < 0, exp(z) / 2, 1 - exp(-z) / 2) +end +@inline function MeasureBase.transport_from_std(::Type{StdUniform}, d::Laplace, p) + u = p - 1 // 2 + _nan_outside(StdUniform, p, muladd(-d.θ * sign(u), log1p(-abs(2 * _unit_interior(p) - 1)), d.μ)) +end + +@inline MeasureBase.preferred_stdmeasure(::Type{<:LogNormal}) = StdNormal +@inline MeasureBase.transport_to_std(::Type{StdNormal}, d::LogNormal, x) = _nan_outside(d, x, (log(abs(x)) - d.μ) / d.σ) +@inline MeasureBase.transport_from_std(::Type{StdNormal}, d::LogNormal, z) = exp(muladd(d.σ, z, d.μ)) + +@inline MeasureBase.preferred_stdmeasure(::Type{<:Weibull}) = StdExponential +@inline MeasureBase.transport_to_std(::Type{StdExponential}, d::Weibull, x) = _nan_outside(d, x, abs(x / d.θ)^d.α) +@inline MeasureBase.transport_from_std(::Type{StdExponential}, d::Weibull, z) = _nan_outside(StdExponential, z, d.θ * abs(z)^(1 / d.α)) + +@inline MeasureBase.preferred_stdmeasure(::Type{<:Gamma}) = StdUniform +@inline MeasureBase.transport_to_std(::Type{StdUniform}, d::Gamma, x) = _nan_outside(d, x, _gamma_cdf(d.α, abs(x / d.θ))) +@inline MeasureBase.transport_from_std(::Type{StdUniform}, d::Gamma, p) = _nan_outside(StdUniform, p, d.θ * _gamma_quantile(d.α, _unit_interior(p))) + +@inline MeasureBase.preferred_stdmeasure(::Type{<:Beta}) = StdUniform +@inline MeasureBase.transport_to_std(::Type{StdUniform}, d::Beta, x) = _nan_outside(d, x, _beta_cdf(d.α, d.β, _unit_clamp(x))) +@inline MeasureBase.transport_from_std(::Type{StdUniform}, d::Beta, p) = _nan_outside(StdUniform, p, _beta_quantile(d.α, d.β, _unit_interior(p))) + +@inline _unit_clamp(x) = clamp(x, zero(x), one(x)) 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/multivariate.jl b/ext/MeasureBaseDistributionsExt/multivariate.jl new file mode 100644 index 00000000..1ddfe237 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/multivariate.jl @@ -0,0 +1,120 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +# Multivariate normal and Dirichlet measures with density kernels and +# transports over flat batches `(n, batch dims...)`, in terms of array +# operations on the parameters, so that they run on devices and in traced +# code. Single variates are vectors. + +# Column batches: flat batches as `(n, :)` matrices, single variates stay +# vectors. Reductions over the columns give arrays over the batch +# dimensions, numbers for single variates. +@inline _as_columns(x::AbstractVector) = x +@inline _as_columns(X::AbstractArray) = reshape(X, (size(X, 1), :)) +@inline _from_columns(y::AbstractVector, ::AbstractVector) = y +@inline _from_columns(Y::AbstractMatrix, X::AbstractArray) = reshape(Y, (size(Y, 1), Base.tail(size(X))...)) +@inline _column_sums(f, z::AbstractVector) = sum(f, z) +@inline _column_sums(f, Z::AbstractMatrix) = vec(sum(f, Z; dims = 1)) +@inline _column_all(z::AbstractVector) = all(z) +@inline _column_all(Z::AbstractMatrix) = vec(all(Z; dims = 1)) +@inline _batch_results(r::Number, ::AbstractVector) = r +@inline _batch_results(r::AbstractVector, X::AbstractArray) = reshape(r, Base.tail(size(X))) +@inline _rows(z::AbstractVector, r) = view(z, r) +@inline _rows(Z::AbstractMatrix, r) = view(Z, r, :) +@inline _masked(ℓ::Number, ins) = ifelse(ins, ℓ, oftype(ℓ, -Inf)) +@inline _masked(ℓ::AbstractArray, ins) = ifelse.(ins, ℓ, eltype(ℓ)(-Inf)) +@inline _nan_columns(Z::AbstractArray, ins) = ifelse.(_as_row(ins), Z, eltype(Z)(NaN)) +@inline _as_row(ins::Bool) = ins +@inline _as_row(ins::AbstractVector) = reshape(ins, 1, :) + + +# Multivariate normal: densities via the Cholesky factor of the covariance. + +const MvNormalMeasure = AsMeasure{<:MvNormal} + +_logdet_cov(Σ::PDMats.PDMat) = 2 * sum(log, diag(Σ.chol.factors)) +_logdet_cov(Σ::PDMats.PDiagMat) = sum(log, Σ.diag) +_logdet_cov(Σ::PDMats.ScalMat) = Σ.dim * log(Σ.value) + +for bhead in (:batched_logdensityof_impl, :batched_logdensity_def) + @eval function MeasureBase.$bhead(m::MvNormalMeasure, X::AbstractArray) + d = m.obj + Z = _cholesky_L(d.Σ) \ (_as_columns(X) .- d.μ) + sq = _column_sums(abs2, Z) + _batch_results(-sq ./ 2 .- (_logdet_cov(d.Σ) + length(d) * log2π) / 2, X) + end +end +MeasureBase.logdensity_def(m::MvNormalMeasure, x::AbstractVector) = MeasureBase.batched_logdensity_def(m, x) +MeasureBase.unsafe_logdensityof(m::MvNormalMeasure, x::AbstractVector) = MeasureBase.batched_logdensityof_impl(m, x) + +function MeasureBase.batched_transport_to_std(::Type{StdNormal}, d::MvNormal, X::AbstractArray) + _from_columns(_cholesky_L(d.Σ) \ (_as_columns(X) .- d.μ), X) +end +function MeasureBase.batched_transport_from_std(::Type{StdNormal}, d::MvNormal, Z::AbstractArray) + _from_columns(_cholesky_L(d.Σ) * _as_columns(Z) .+ d.μ, Z) +end +MeasureBase.transport_to_std(::Type{StdNormal}, d::MvNormal, x) = MeasureBase.batched_transport_to_std(StdNormal, d, x) +MeasureBase.transport_from_std(::Type{StdNormal}, d::MvNormal, z) = MeasureBase.batched_transport_from_std(StdNormal, d, z) + +# Parameters follow the batches to the device: +function Adapt.adapt_structure(to, m::MvNormalMeasure) + d = m.obj + asmeasure(MvNormal(Adapt.adapt(to, d.μ), _adapt_cov(to, d.Σ))) +end +function _adapt_cov(to, Σ::PDMats.PDMat) + chol = Σ.chol + PDMats.PDMat(Adapt.adapt(to, Σ.mat), Cholesky(Adapt.adapt(to, chol.factors), chol.uplo, chol.info)) +end +_adapt_cov(to, Σ::PDMats.PDiagMat) = PDMats.PDiagMat(Adapt.adapt(to, Σ.diag)) +_adapt_cov(to, Σ::PDMats.ScalMat) = Σ + + +# Dirichlet: densities over the simplex, transports via the stick-breaking +# Beta transports (M. J. Betancourt, "Cruising The Simplex: Hamiltonian +# Monte Carlo and the Dirichlet Distribution", arXiv:1010.3436), with the +# cumulative sums and products running along the variate dimension. + +for bhead in (:batched_logdensityof_impl, :batched_logdensity_def) + @eval function MeasureBase.$bhead(m::DirichletMeasure, X::AbstractArray) + d = m.obj + Xc = _as_columns(X) + ℓ = _column_sums(identity, _clog.(d.alpha .- 1, abs.(Xc))) .- d.lmnB + _batch_results(_masked(ℓ, _simplex_mask(Xc)), X) + end +end +MeasureBase.logdensity_def(m::DirichletMeasure, x::AbstractVector) = MeasureBase.batched_logdensity_def(m, x) +MeasureBase.unsafe_logdensityof(m::DirichletMeasure, x::AbstractVector) = MeasureBase.batched_logdensityof_impl(m, x) + +@inline function _simplex_mask(Xc::AbstractArray) + tol = sqrt(eps(float(eltype(Xc)))) + _column_all(Xc .>= 0) .& (abs.(_column_sums(identity, Xc) .- 1) .<= tol) +end + +# The stick-breaking Beta parameters, for the first `K - 1` components: +@inline _stick_breaking_params(d::Dirichlet) = (_dropfront(_rev_cumsum(d.alpha)), _dropback(d.alpha)) + +function MeasureBase.batched_transport_to_std(::Type{StdUniform}, d::Dirichlet, X::AbstractArray) + K = length(d) + αs, βs = _stick_breaking_params(d) + Xc = _as_columns(X) + rem = 1 .- cumsum(Xc; dims = 1) + # The remaining mass before each component is the mass after it plus + # the component itself: + beta_v = _rows(rem, 1:(K - 1)) ./ (_rows(rem, 1:(K - 1)) .+ _rows(Xc, 1:(K - 1))) + Z = _beta_cdf.(αs, βs, _unit_clamp.(beta_v)) + _from_columns(_nan_columns(Z, _simplex_mask(Xc)), X) +end + +function MeasureBase.batched_transport_from_std(::Type{StdUniform}, d::Dirichlet, Z::AbstractArray) + K = length(d) + αs, βs = _stick_breaking_params(d) + Zc = _as_columns(Z) + beta_v = _beta_quantile.(αs, βs, _unit_interior.(Zc)) + cp = cumprod(beta_v; dims = 1) + # Each component takes what its Beta variate leaves of the remaining mass: + X = vcat(1 .- _rows(cp, 1:1), _rows(cp, 1:(K - 2)) .- _rows(cp, 2:(K - 1)), _rows(cp, (K - 1):(K - 1))) + _from_columns(_nan_columns(X, _column_all((Zc .>= 0) .& (Zc .<= 1))), Z) +end +MeasureBase.transport_to_std(::Type{StdUniform}, d::Dirichlet, x) = MeasureBase.batched_transport_to_std(StdUniform, d, x) +MeasureBase.transport_from_std(::Type{StdUniform}, d::Dirichlet, z) = MeasureBase.batched_transport_from_std(StdUniform, d, z) + +Adapt.adapt_structure(to, m::DirichletMeasure) = asmeasure(Dirichlet(Adapt.adapt(to, m.obj.alpha))) diff --git a/ext/MeasureBaseDistributionsExt/product.jl b/ext/MeasureBaseDistributionsExt/product.jl new file mode 100644 index 00000000..35817c4b --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/product.jl @@ -0,0 +1,64 @@ +# 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)) + + @inline function MeasureBase.preferred_stdmeasure(::Type{<:Distributions.Product{<:Any,T}}) where {T} + MeasureBase.preferred_stdmeasure(T) + end + + 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)) + + @inline function MeasureBase.preferred_stdmeasure(::Type{<:Distributions.ProductDistribution{N,M,D}}) where {N,M,D<:AbstractArray} + MeasureBase.preferred_stdmeasure(eltype(D)) + end + + 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..8f55c34c --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/reshaped.jl @@ -0,0 +1,21 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +MeasureBase.getdof(μ::ReshapedDistribution) = MeasureBase.getdof(μ.dist) + + +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..a51280b3 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/standard_dist.jl @@ -0,0 +1,219 @@ +# 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 + + +for (A, B) in [ + (Uniform, StdUniform), + (Exponential, StdExponential), + (Logistic, StdLogistic), + (Normal, StdNormal) +] + @eval begin + @inline MeasureBase.preferred_stdmeasure(::Type{<:StandardDist{$A}}) = $B + @inline MeasureBase.transport_to_std(::Type{$B}, ::StandardDist{$A,0}, x::Number) = x + @inline MeasureBase.transport_from_std(::Type{$B}, ::StandardDist{$A,0}, z::Number) = z + @inline MeasureBase.transport_to_std(::Type{$B}, ::StandardDist{$A}, x::AbstractArray) = vec(x) + @inline MeasureBase.transport_from_std(::Type{$B}, d::StandardDist{$A}, z::AbstractVector) = reshape(z, 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 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..1c9293aa --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/standardmv.jl @@ -0,0 +1,25 @@ +# 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) + +@inline MeasureBase.preferred_stdmeasure(::Type{<:AbstractMvNormal}) = StdNormal + +_cholesky_L(A) = _lower_factor(cholesky(A)) +# `Cholesky.L` copies the transposed factor and scalar-indexes on GPUs. +_lower_factor(C::Cholesky) = C.uplo === 'L' ? LowerTriangular(C.factors) : UpperTriangular(C.factors)' +_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)) + +#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..27d5f523 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/univariate.jl @@ -0,0 +1,143 @@ +# 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 + +_dist_params_numtype(d::Distribution) = real_numtype(typeof(Distributions.params(d))) + +@inline function _result_numtype(d::Distribution{Univariate}, x::T) where {T<:Number} + float(promote_type(T, _dist_params_numtype(d))) +end + + +# Generic transports between univariate continuous distributions and +# StdLogistic: the log-cdf and log-ccdf keep both tails accurate on the way +# to the standard measure, quantile and complementary quantile on the way +# back. The implementation hooks are specialized for dual numbers in the +# ForwardDiff extension. + +@inline MeasureBase.preferred_stdmeasure(::Type{<:Distribution{Univariate,Continuous}}) = StdLogistic + +@inline _trafo_logcdf(d::Distribution{Univariate,Continuous}, x::Number) = + _trafo_logcdf_impl(_dist_params_numtype(d), d, x) +@inline _trafo_logccdf(d::Distribution{Univariate,Continuous}, x::Number) = + _trafo_logccdf_impl(_dist_params_numtype(d), d, x) +@inline _trafo_quantile(d::Distribution{Univariate,Continuous}, p::Number) = + _trafo_quantile_impl(_dist_params_numtype(d), d, p) +@inline _trafo_cquantile(d::Distribution{Univariate,Continuous}, p::Number) = + _trafo_cquantile_impl(_dist_params_numtype(d), d, p) + +@inline _trafo_logcdf_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, x::Number) = + Distributions.logcdf(d, x) +@inline _trafo_logccdf_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, x::Number) = + Distributions.logccdf(d, x) +@inline _trafo_quantile_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, p::Number) = + _dist_quantile(d, p) +@inline _trafo_cquantile_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, p::Number) = + _dist_cquantile(d, p) + +@inline _dist_quantile(d::Distribution{Univariate,Continuous}, p::Number) = Distributions.quantile(d, p) +@inline _dist_cquantile(d::Distribution{Univariate,Continuous}, p::Number) = Distributions.cquantile(d, p) + +# The quantile implementation of Beta only supports Float64: +const _Float64Compatible = Union{Integer,AbstractFloat} +@inline function _dist_quantile(d::Beta{<:_Float64Compatible}, p::_Float64Compatible) + Distributions.quantile(d, convert(promote_type(Float64, typeof(p)), p)) +end +@inline function _dist_cquantile(d::Beta{<:_Float64Compatible}, p::_Float64Compatible) + Distributions.cquantile(d, convert(promote_type(Float64, typeof(p)), p)) +end + +# Rounding errors can push quantiles of truncated distributions slightly +# outside of their support: +const _Truncated = Distributions.Truncated{<:Distribution{Univariate,Continuous}} +@inline _dist_quantile(d::_Truncated, p::Real) = _clamp_to_support(d, Distributions.quantile(d, p)) +@inline _dist_cquantile(d::_Truncated, p::Real) = _clamp_to_support(d, Distributions.cquantile(d, p)) + +function _clamp_to_support(d::_Truncated, x::T) where {T<:Real} + 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 MeasureBase.transport_to_std(::Type{StdLogistic}, d::Distribution{Univariate,Continuous}, x) + R = _result_numtype(d, x) + l = _trafo_logcdf(d, x) - _trafo_logccdf(d, x) + ifelse(Distributions.insupport(d, x), convert(R, l), convert(R, NaN)) +end + +@inline function MeasureBase.transport_from_std(::Type{StdLogistic}, d::Distribution{Univariate,Continuous}, l) + R = _result_numtype(d, l) + # From the side that keeps the tail: + x = l < zero(l) ? _trafo_quantile(d, _unit_interior(logistic(l))) : _trafo_cquantile(d, _unit_interior(logistic(-l))) + convert(R, x) +end + + +# Location-scale families of standard measures transport by their affine map: + +@inline function _affine_to_std(d::Distribution{Univariate}, x::Number) + z = (x - Distributions.location(d)) / Distributions.scale(d) + convert(_result_numtype(d, x), z) +end + +@inline function _std_to_affine(d::Distribution{Univariate}, z::Number) + x = muladd(z, Distributions.scale(d), Distributions.location(d)) + convert(_result_numtype(d, z), x) +end + +# Transports outside the support of the source give NaN, the formulas +# must not throw outside the support (see `MeasureBase._nan_outside`): +@inline _nan_outside(d::Distribution, x, y) = MeasureBase._nan_outside(asmeasure(d), x, y) +@inline _nan_outside(::Type{S}, z, y) where {S<:StdMeasure} = MeasureBase._nan_outside(S(), z, y) + +for (D, S) in [ + (Logistic, StdLogistic), + (Normal, StdNormal) +] + @eval begin + @inline MeasureBase.preferred_stdmeasure(::Type{<:$D}) = $S + @inline MeasureBase.transport_to_std(::Type{$S}, d::$D, x) = _affine_to_std(d, x) + @inline MeasureBase.transport_from_std(::Type{$S}, d::$D, z) = _std_to_affine(d, z) + end +end + +@inline MeasureBase.preferred_stdmeasure(::Type{<:Uniform}) = StdUniform +@inline MeasureBase.transport_to_std(::Type{StdUniform}, d::Uniform, x) = _nan_outside(d, x, _affine_to_std(d, x)) +@inline MeasureBase.transport_from_std(::Type{StdUniform}, d::Uniform, z) = _nan_outside(StdUniform, z, _std_to_affine(d, z)) + +@inline MeasureBase.preferred_stdmeasure(::Type{<:Exponential}) = StdExponential +@inline MeasureBase.transport_to_std(::Type{StdExponential}, d::Exponential, x) = + _nan_outside(d, x, convert(_result_numtype(d, x), Distributions.scale(d) \ x)) +@inline MeasureBase.transport_from_std(::Type{StdExponential}, d::Exponential, z) = + _nan_outside(StdExponential, z, convert(_result_numtype(d, z), Distributions.scale(d) * z)) + + +# Affine transformed distributions transport via the underlying distribution: + +const _AffineDist = Distributions.AffineDistribution + +@inline MeasureBase.preferred_stdmeasure(::Type{<:_AffineDist{<:Any,<:Any,D}}) where {D} = + MeasureBase.preferred_stdmeasure(D) + +@inline function MeasureBase.transport_to_std(::Type{S}, d::_AffineDist, x) where {S<:StdMeasure} + transport_to_std(S, d.ρ, d.σ \ (x - d.μ)) +end +@inline function MeasureBase.transport_from_std(::Type{S}, d::_AffineDist, z) where {S<:StdMeasure} + muladd(d.σ, transport_from_std(S, d.ρ, z), d.μ) +end +# Disambiguation with the generic univariate transports: +@inline function MeasureBase.transport_to_std(::Type{StdLogistic}, d::_AffineDist, x) + transport_to_std(StdLogistic, d.ρ, d.σ \ (x - d.μ)) +end +@inline function MeasureBase.transport_from_std(::Type{StdLogistic}, d::_AffineDist, z) + muladd(d.σ, transport_from_std(StdLogistic, d.ρ, z), d.μ) +end diff --git a/ext/MeasureBaseDistributionsForwardDiffExt.jl b/ext/MeasureBaseDistributionsForwardDiffExt.jl new file mode 100644 index 00000000..181fc410 --- /dev/null +++ b/ext/MeasureBaseDistributionsForwardDiffExt.jl @@ -0,0 +1,108 @@ +# 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 + +# Dual-number transports for distributions with plain parameters, via the +# derivatives of cdf and quantile: + +const _PlainParams = Type{<:Union{Integer,AbstractFloat}} + +@inline function MeasureBase._trafo_logcdf_impl( + ::_PlainParams, + d::Distribution{Univariate,Continuous}, + x::ForwardDiff.Dual{TAG}, +) where {TAG} + x_v = ForwardDiff.value(x) + lp = Distributions.logcdf(d, x_v) + dlp_dx = exp(Distributions.logpdf(d, x_v) - lp) + ForwardDiff.Dual{TAG}(lp, dlp_dx * ForwardDiff.partials(x)) +end + +@inline function MeasureBase._trafo_logccdf_impl( + ::_PlainParams, + d::Distribution{Univariate,Continuous}, + x::ForwardDiff.Dual{TAG}, +) where {TAG} + x_v = ForwardDiff.value(x) + lp = Distributions.logccdf(d, x_v) + dlp_dx = -exp(Distributions.logpdf(d, x_v) - lp) + ForwardDiff.Dual{TAG}(lp, dlp_dx * ForwardDiff.partials(x)) +end + +@inline function MeasureBase._trafo_quantile_impl( + ::_PlainParams, + d::Distribution{Univariate,Continuous}, + p::ForwardDiff.Dual{TAG}, +) where {TAG} + p_v = ForwardDiff.value(p) + x = MeasureBase._dist_quantile(d, p_v) + dx_dp = inv(Distributions.pdf(d, x)) + ForwardDiff.Dual{TAG}(x, dx_dp * ForwardDiff.partials(p)) +end + +@inline function MeasureBase._trafo_cquantile_impl( + ::_PlainParams, + d::Distribution{Univariate,Continuous}, + p::ForwardDiff.Dual{TAG}, +) where {TAG} + p_v = ForwardDiff.value(p) + x = MeasureBase._dist_cquantile(d, p_v) + dx_dp = -inv(Distributions.pdf(d, x)) + ForwardDiff.Dual{TAG}(x, dx_dp * ForwardDiff.partials(p)) +end + +# Dual numbers through the regularized incomplete gamma and beta functions +# and their inverses: the derivative with respect to the variate resp. +# probability argument follows from the density, derivatives with respect +# to the parameters are not available and yield NaN partials. +using MeasureBase: _gamma_cdf, _gamma_quantile, _beta_cdf, _beta_quantile, _gamma_logpdf, _beta_logpdf + +const _Dual = ForwardDiff.Dual + +@inline MeasureBase._dualtag(::_Dual{TAG}, ::Number...) where {TAG} = _Dual{TAG} + +# The derivative through the last argument; the partials of parameters +# are marked NaN where they are nonzero (their derivatives are unknown): +@inline _arg_partials(::Type{_Dual{TAG}}, x::_Dual{TAG}, params...) where {TAG} = ForwardDiff.partials(x) +@inline _arg_partials(::Type{_Dual{TAG}}, ::Real, params...) where {TAG} = zero(ForwardDiff.partials(_first_dual(params...))) +@inline _first_dual(x::_Dual, rest...) = x +@inline _first_dual(::Real, rest...) = _first_dual(rest...) +@inline _nan_partials(∂, ::Real) = ∂ +@inline function _nan_partials(∂, x::_Dual) + ∂ + ForwardDiff.Partials(map(v -> ifelse(iszero(v), zero(v), oftype(v, NaN)), ForwardDiff.partials(x).values)) +end +@inline function _through_last(::Type{_Dual{TAG}}, value, dvalue, last, params...) where {TAG} + ∂ = dvalue * _arg_partials(_Dual{TAG}, last, params...) + ForwardDiff.Dual{TAG}(value, foldl(_nan_partials, params; init = ∂)) +end + +@inline function MeasureBase._gamma_cdf_impl(::Type{_Dual{TAG}}, α, x) where {TAG} + αv, xv = ForwardDiff.value(α), ForwardDiff.value(x) + _through_last(_Dual{TAG}, _gamma_cdf(αv, xv), exp(_gamma_logpdf(αv, xv)), x, α) +end +@inline function MeasureBase._gamma_quantile_impl(::Type{_Dual{TAG}}, α, p) where {TAG} + αv, pv = ForwardDiff.value(α), ForwardDiff.value(p) + xv = _gamma_quantile(αv, pv) + _through_last(_Dual{TAG}, xv, exp(-_gamma_logpdf(αv, xv)), p, α) +end +@inline function MeasureBase._beta_cdf_impl(::Type{_Dual{TAG}}, α, β, x) where {TAG} + αv, βv, xv = ForwardDiff.value(α), ForwardDiff.value(β), ForwardDiff.value(x) + _through_last(_Dual{TAG}, _beta_cdf(αv, βv, xv), exp(_beta_logpdf(αv, βv, xv)), x, α, β) +end +@inline function MeasureBase._beta_quantile_impl(::Type{_Dual{TAG}}, α, β, p) where {TAG} + αv, βv, pv = ForwardDiff.value(α), ForwardDiff.value(β), ForwardDiff.value(p) + xv = _beta_quantile(αv, βv, pv) + _through_last(_Dual{TAG}, xv, exp(-_beta_logpdf(αv, βv, xv)), p, α, β) +end + +# The quantile of Beta doesn't support dual parameters: +@inline MeasureBase._dist_quantile(d::Beta{<:ForwardDiff.Dual}, p::Real) = convert(float(typeof(p)), NaN) +@inline MeasureBase._dist_cquantile(d::Beta{<:ForwardDiff.Dual}, p::Real) = convert(float(typeof(p)), NaN) + +end # module MeasureBaseDistributionsForwardDiffExt diff --git a/ext/MeasureBaseDistributionsForwardDiffPullbacksExt.jl b/ext/MeasureBaseDistributionsForwardDiffPullbacksExt.jl new file mode 100644 index 00000000..d7ee12b2 --- /dev/null +++ b/ext/MeasureBaseDistributionsForwardDiffPullbacksExt.jl @@ -0,0 +1,22 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +module MeasureBaseDistributionsForwardDiffPullbacksExt + +import MeasureBase +using MeasureBase: StdMeasure, transport_to_std, transport_from_std + +import Distributions +using Distributions: Distribution, Univariate + +import ChainRulesCore +using ForwardDiffPullbacks: fwddiff + +# Use ForwardDiff for univariate transports: +@inline function ChainRulesCore.rrule(::typeof(transport_to_std), ::Type{S}, d::Distribution{Univariate}, x::Any) where {S<:StdMeasure} + ChainRulesCore.rrule(fwddiff(transport_to_std), S, d, x) +end +@inline function ChainRulesCore.rrule(::typeof(transport_from_std), ::Type{S}, d::Distribution{Univariate}, z::Any) where {S<:StdMeasure} + ChainRulesCore.rrule(fwddiff(transport_from_std), S, d, z) +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..899e2036 --- /dev/null +++ b/ext/MeasureBaseMooncakeExt.jl @@ -0,0 +1,25 @@ +# 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 + +# 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(check_dof),Any,Any} + +end # module MeasureBaseMooncakeExt diff --git a/ext/MeasureBaseReactantExt.jl b/ext/MeasureBaseReactantExt.jl new file mode 100644 index 00000000..c39be494 --- /dev/null +++ b/ext/MeasureBaseReactantExt.jl @@ -0,0 +1,17 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +module MeasureBaseReactantExt + +using Reactant: Reactant, TracedRNumber +using IrrationalConstants: sqrt2 +import MeasureBase +using MeasureBase: RealValues, IntegerValues + +Base.in(::TracedRNumber{<:Real}, ::RealValues) = true +Base.in(::TracedRNumber{<:Integer}, ::IntegerValues) = true + +# CHLO provides erf_inv but no erfc_inv, so the standard normal quantile +# loses precision for arguments close to 0 and 1 in traced code: +MeasureBase.Φinv(p::TracedRNumber) = Reactant.Ops.erf_inv(2 * p - 1) * sqrt2 + +end # module MeasureBaseReactantExt diff --git a/redesign.md b/redesign.md new file mode 100644 index 00000000..71b8cbf8 --- /dev/null +++ b/redesign.md @@ -0,0 +1,230 @@ +# MeasureBase redesign notes (branch `major-upgrade`) + +Working notes on the approach behind this branch, for reviews and for +guiding the next steps. Kept up to date while the branch evolves, to be +removed before the merge. + +## Goals + +- A breaking release that runs on GPUs (CUDA, JLArrays) and under + Reactant, with batching built into the foundation. +- One implementation per measure type for densities, transports and + random variates, so scalar and batched paths can't drift apart. +- Composable, structural solutions: powers, products, combinations, + binds and pushforwards implement their behavior once in terms of their + components. No shape inference, no per-call DOF sums, no function + traits in the core. + +## Design philosophy + +- **Batched first.** A single variate is a batch with zero batch + dimensions. Every kernel handles both, the point API is the batched API + at zero batch dimensions. Static arrays keep the scalar path + allocation-free. +- **Ranks, not sizes.** A flat batch is an array + `(variate dims..., batch dims...)`. Kernels only need the variate rank + of their measure; sizes are optional declarations for validation, + stream consumption and static fast paths, never for routing. Unknown + sizes are safe. +- **Standard measures as pivots.** Transport goes through a standard + measure type chosen by promoting the measures' preferences. Measure + types implement transport to and from their preferred standard measure + only. +- **Branch-free, device-friendly kernels.** Broadcasts, reductions and + masks instead of branches; host loops only where documented. +- **Entry points normalize layouts.** Users pass nested arrays, tuples of + batches, struct arrays or flat arrays; the kernels only see flat + batches. + +## Concepts + +**Variate rank.** `mspace_ndims(::Type{M})`, 0 for scalar variates. +Derived from `mspace_flatsize(::Type{M})` where known, declared by +array-variate leaves, derived by structural measures. Without a rank the +batched defaults throw with a message naming the declaration; point +kernels keep working. + +**Flat batches.** Kernels return arrays over the batch dimensions, a +number for a single variate, possibly lazily. ArraysOfArrays containers +are fused into their flat storage at the entry points, ragged +containers are evaluated variate by variate. Tuple and named tuple +variates batch as tuples of batches; struct arrays and arrays of tuples +are accepted, their flat storage is the tuple of component storages. +`rand(Pt^n)` of a tuple product is a struct array. + +**Streams.** Variates of `mcombine(vcat, ...)`, binds and tuple products +inside such streams are flat vectors consumed with the with-rest +protocol. Point forms return `(result, x_μ, x_rest)` (binds need the +consumed variate), batched forms take streams `(rows, batch dims...)` +and a multiplicity `sz::SizeLike` of variates per stream and return +`(result, rest)`. Powers pass their size as multiplicity to their base; +combined measures and tuple products split rows by their fixed stream +lengths. `fixed_stream_size(::Type{M})` decides whether a batch of +streams is consumed in fused operations or stream by stream by the +outermost combinator (binds never fuse). Scalar leaves consume +`(1, sz..., batch dims...)` and drop the leading dimension. Nested +element variates in vcat streams are flattened. + +**Transport.** `transport_to(ν, μ)` gives a `TransportFunction`; `f(x)` +transports a variate, `f.(X)` a whole batch. Leaves implement +`transport_to_std`/`transport_from_std` for their preferred standard +measure (`preferred_stdmeasure`, `promote_stdmeasure`, `AnyStdMeasure`, +`NoStdTransport`), array-variate leaves also the `batched_` forms. +`stdconvert` converts between standard measures in log form, +`transport_def` may be specialized for direct pairs. Standard streams are +`(dof, batch dims...)`; the from-side default consumes `fast_dof(μ)` +entries per variate, composed measures implement the with-rest forms. +`getdof`/`fast_dof` are declaration-derived, used at construction time +and for chunking, never inside kernels. + +**Static variates.** Sizes are `StaticThings.SizeLike` throughout, the +with-rest multiplicity and the batch size of `rand` included, so a +statically sized power stays static through streams. On the CPU, bulk +draws of a fully static size are static arrays and the stream splits and +reshapes keep them static: `rand`, `transport_to` and `logdensityof` of a +measure whose variate sizes are all static are type stable and allocation +free, for powers, tuple and named tuple products and binds over them. +Sizes that depend on variate values stay dynamic, and other compute units +allocate their own arrays. + +**Random variates.** `rand(ctx::GenContext, μ)` with RNG, precision and +compute unit (`rand(μ)`, `rand(rng, μ)`, `rand(T, μ)` are wrappers). +`batched_rand_impl(ctx, μ, sz::SizeLike)` returns a flat batch, a single +variate for `sz == ()`; `rand_impl` defaults to it. Defaults draw +standard variates in bulk on the compute unit and transport them, or +generate variate by variate without a standard transport. + +**Array products.** `productmeasure(::AbstractArray)` stores isbits +parameterized marginals as `StructArrays` (nested parameter structs +unwrapped; numbers, arrays, tuples, strings, symbols and function objects +stay opaque columns), in one place: `_marginal_storage`. Fused kernels +broadcast over the leaf columns and rebuild marginals via +`ConstructionBase.constructorof`, which works on CUDA and under Reactant. +Fusion needs concrete scalar-variate marginals (one DOF for transport); +other array products loop over host-resident marginals. Measures holding +arrays have `Adapt` rules. + +**Pushforwards.** `pushfwd(f, μ)` learns its output size once at +construction from a test value when the origin has a size. Batched +application needs no traits: elementwise for `Base.BroadcastFunction` +(fused density kernels), column batches for AffineMaps types (weak +dependency, `MeasureBaseAffineMapsExt`), a host loop otherwise. + +## Extension points + +| Aspect | Scalar-variate leaf | Array-variate leaf | Composed measure | +|---|---|---|---| +| Density | `logdensity_def` | `mspace_ndims`, `batched_logdensityof_impl` (+`_def`) | kernels via components, with-rest forms | +| Transport | `preferred_stdmeasure`, `transport_to_std`, `transport_from_std` | + `batched_transport_to_std`, `batched_transport_from_std` | with-rest forms, `fixed_stream_size` | +| Random | `batched_rand_impl` (default via transport) | `batched_rand_impl` | derived | +| Declarations | none | `mspace_ndims`, optionally `mspace_flatsize`, `getdof` | derived | + +## Layouts per combinator + +- Powers: `(base dims..., power dims..., batch dims...)`, innermost base + first; standard variates are the flat vector of the base's. Results + are nested views over the flat storage. Powers of tuple products treat + numeric flat variates as streams and accept tuples of batches. +- Array products: `(marginal dims..., product dims..., batch dims...)`. +- Tuple products: tuples of batches; marginal by marginal in streams. +- Combined `vcat`: streams; `merge`: merged named tuples. Binds: single + streams, value-dependent sizes. +- Weighted, restricted, density measures, Half: forward plus weights or + masks. Superpositions and spike mixtures: one batch per component, + masks aligned with the variate dimensions. Dirac: constant batches. +- Distributions extension: the main univariate families (Normal, + Uniform, Exponential, Logistic, Cauchy, Laplace, LogNormal, Weibull, + Gamma, Beta, Poisson, Bernoulli) get branch-free density formulas and + transports to the standard measure matching their shape (`families.jl`), + other univariate distributions go via `StdLogistic` (log-cdf and + quantile). `MvNormal` works on column batches through its Cholesky + factor, `Dirichlet` through stick-breaking Beta transports with `cumsum` + and `cumprod` along the variate dimension (`multivariate.jl`). Incomplete + gamma and beta functions go through the `_gamma_cdf`/`_beta_cdf` hooks + (and quantiles), whose ForwardDiff and ChainRules derivatives come from + the densities. Draws use the Distributions samplers on the CPU and the + standard transports on other compute units. Remaining array-variate + distributions batch via `logpdf(d, X)` on the host. + +## Changes relative to `master` + +Commits since `c773afe`: variate size contract and +`preferred_stdmeasure`, densities over flat storage, branch-free +evaluation, Reactant smoke tests, structural batched kernels, transport +rebuilt on standard measures, batched transport and the broadcast hook, +rand via `GenContext`, then the batched-first redesign (density core, +struct array products, transport, random variates, structured batches), +the review fixes, and device-friendly kernels for wrapped Distributions. + +Removed: `transport_origin`/`to_origin`/`from_origin` and the origin +machinery, `NoTransport`, `transport_to_mvstd`, per-measure +`Base.rand(rng, T, μ)` methods, the `_trafo_cdf`/`_trafo_quantile` hooks. + +Behavior changes for NEWS: univariate Distributions pivot on +`StdLogistic`; nested powers and array products return ArraysOfArrays +views, powers of tuple products struct arrays; `Half` transports via +`StdUniform` (tails limited); `mcombine(vcat, ::Product, ::Product)` +merges only concrete homogeneous marginals; `rand(rng, Int, μ)` +unsupported; isbits marginal arrays become struct arrays (`rand` of such +products gives plain arrays); vcat-combined and bind variates are flat; +`transport_to(StdUniform, m)` with binds inside errors (use +`StdUniform()^n`); batched kernels of rank-less measures throw. + +## Verification + +Full suite (Aqua, extensions, doctests) on CPU with JLArrays cases and +a FixedSizeArrays transparency check (fixed-size inputs give fixed-size +outputs; FixedSizeArrays stays a test dependency, allocating fixed-size +variates by default is a HeterogeneousComputing decision for later). +`test/static_variates.jl` checks that statically sized measures stay type +stable and allocation free. +`test/test_reactant.jl` runs as part of the suite on 64-bit Linux and +macOS with stable Julia, adding Reactant on demand as MGVI does (backend +via `MEASUREBASE_REACTANT_BACKEND`). `test/cuda` is opt-in. Both run +locally on the GB10, green at HEAD except one expected-broken CUDA case +(AffineMaps Jacobian on device). + +## Known gaps and open decisions + +- Upstream: AffineMaps lacks `Adapt` rules and device/traced Jacobians; + Distributions isn't device-aware and its parameter structs (PDMats, + `Dirichlet`) can't hold traced arrays, so under Reactant distribution + parameters stay constants; SpecialFunctions' incomplete beta and gamma + functions don't compile for GPUs or Reactant, so Beta and Dirichlet + transports and draws (Gamma under Reactant too) stay on the CPU; + ChangesOfVariables has no rules for `Base.Fix1`/`Fix2` arithmetic; + HeterogeneousComputing has no Reactant compute unit; JLArrays has no + RNG and no triangular solves; Reactant rejects traced `VectorOfArrays` + and empty batches. +- Out-of-support convention (decided 2026-09-19): structural errors (wrong + rank, size, container or element kind) throw an `ArgumentError` from + `checked_arg` at the entry points; variates of the right shape never + throw: densities are `-Inf` outside the support (including non-integers + for counting-based measures and infinite values), transports are `NaN` + outside the support of the source, `NaN` inputs give `NaN` or `-Inf`, + relative densities keep `+Inf`/`-Inf`/`NaN`. Finite inputs give finite + transports: uniform inputs are clamped into the open unit interval + before quantiles (the endpoints stand for their nearest interior grid + points, a null-set convention like BAT's `[eps, 1 - eps]` clamping in + practice), tail probabilities of log-space conversions floor at the + smallest positive float. + Kernels must not throw outside the support, since the masks evaluate + both branches (`abs`, `clamp`, `min` guards instead of `NaNMath`, which + isn't device-compatible). Downstream checks such as BAT's + `checked_logdensityof` stay downstream. +- BAT's boundary tweaks: adopted are infinite variates outside the + support of continuous wrapped distributions, the clamping of uniform + inputs into the open unit interval (for inputs inside `[0, 1]`, outside + stays `NaN`), `-Inf + Inf` giving `-Inf` in pushforward densities and a zero + Jacobian term where both densities vanish. Not adopted: replacing finite + densities with infinite Jacobian terms by `-1e38`, and re-evaluating + densities an `eps` inside the support where Distributions returns `NaN` + (the family kernels are exact there). + Quantile results within `4 eps` of the support edges snap to the edges + only on the generic logistic path of wrapped distributions. +- Decisions pending: `Half` tails via log-ccdf, device random variate + infrastructure and `rand!`, the `smart-constructors.jl` review + (location-scale arrays as affine pushforwards of powers), + `_static_ndims` type-first vs. instance-first. +- Polish before merge: docs pass, NEWS, history curation, version bump, + remove this file. diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index 2bad7d92..f4a72f7e 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,33 @@ using Static: StaticInteger using FunctionChains using PropertyFunctions: PropSelFunction +using StaticThings: + SizeLike, StaticSizeLike, OneToLike, StaticOneToLike, IntegerLike, + asaxes, asnonstatic, canonical_size, size_dims, + maybestatic_eachindex, maybestatic_length, maybestatic_size, + maybestatic_first, maybestatic_last, maybestatic_view, + maybestatic_oneto, maybestatic_fill, maybestatic_reshape, + axes2size, size2length, split_at, staticarray_type, + static_all, static_any, static_mapreduce, + sum_leading_dims, drop_leading_dims, merge_leading_dims, all_leading_dims + +import HeterogeneousComputing +using HeterogeneousComputing: real_numtype +using HeterogeneousComputing: + GenContext, AbstractComputeUnit, CPUnit, get_rng, get_precision, get_compute_unit, allocate_array + +using ArraysOfArrays: + ArrayOfSimilarArrays, VectorOfArrays, VectorOfVectors, VectorOfSimilarArrays, + VectorOfSimilarVectors, flatview, fused, stacked, sliced, getsplitmode, + is_memordered_splitmode, AbstractSplitMode, AbstractPartMode, UnknownSplitMode, NonSplitMode + +using OneTwoMany: firstarg, secondarg + +using StructArrays: StructArray +import StructArrays +import Adapt + export gentype -export rebase export AbstractMeasure @@ -64,6 +90,47 @@ 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 + +# Struct arrays of wrapped objects rebuild elements via ConstructionBase: +ConstructionBase.constructorof(::Type{<:AsMeasure}) = _asmeasure +_asmeasure(obj) = AsMeasure{typeof(obj)}(obj) + +Base.:(==)(a::AsMeasure, b::AsMeasure) = a.obj == b.obj +Base.hash(a::AsMeasure, h::UInt) = hash(a.obj, hash(:AsMeasure, h)) +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 +172,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,53 +188,56 @@ using Compat using IrrationalConstants using IrrationalConstants: loghalf -include("static.jl") +include("collection_utils.jl") include("smf.jl") +include("mspace.jl") include("getdof.jl") +include("standard/stdmeasure.jl") include("transport.jl") -include("schema.jl") -include("splat.jl") +include("rand.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("density-batched.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("transport-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") -include("standard/stdmeasure.jl") include("standard/stduniform.jl") include("standard/stdexponential.jl") include("standard/stdlogistic.jl") include("standard/stdnormal.jl") +include("standard/stdconvert.jl") +include("standard/stdtraits.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..9b4f2155 --- /dev/null +++ b/src/collection_utils.jl @@ -0,0 +1,92 @@ +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)) + + +@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_at(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], maybestatic_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_at(x, size2length(sz)) + return maybestatic_reshape(a_flat, sz), x_rest +end + +Base.@propagate_inbounds _consume_from_stream(x::AbstractVector, sz::StaticArrays.Size) = + _consume_from_stream(x, size_dims(sz)) + +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..2900d999 100644 --- a/src/combinators/bind.jl +++ b/src/combinators/bind.jl @@ -1,36 +1,424 @@ -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, -export ↣ +```julia +mbind(mkernel(f_β, f_c), α) == mbind(f_β, α, f_c) +``` +See also [`mbind`](@ref). """ -If -- μ is an `AbstractMeasure` or satisfies the Measure interface, and -- k is a function taking values from the support of μ and returning a measure +function mkernel end +export mkernel + + +""" + struct MeasureBase.MKernel <: Function + +Represents a generalized monadic transition kernel. + +User code should not create instances of `MKernel` directly, but should +call [`mkernel`](@ref) instead. +""" +struct MKernel{FT,FC} <: Function + f_β::FT + f_c::FC +end + +(f_k::MKernel)(a) = mbind(f_k, Dirac(a)) -Then `μ ↣ k` is a measure, called a *monadic bind*. In a -probabilistic programming language like Soss.jl, this could be expressed as +@inline mkernel(f_β::MKernel) = f_β +@inline mkernel(f_β, f_c = secondarg) = _generic_mkernel_impl(f_β, f_c) -Note that bind is usually written `>>=`, but this symbol is unavailable in Julia. +@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 + +```math +\mu(B) = \int_A \beta_a(B)\, \mathrm{d}\, \alpha(a) ``` -bind = @model μ,k begin - x ~ μ - y ~ k(x) - return y + +which is equivalent to a monadic bind, viewing measures as monads. + +Computationally, `ab = rand(μ)` is equivalent to + +```julia +a = rand(α) +β_a = f_β(a) +b = rand(β_a) +ab = f_c(a, b) +``` + +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 + +@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. +""" +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? + -See also `bind` and `Bind` """ -↣(μ, k) = bind(μ, k) + 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) -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 +""" + 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} + +@inline preferred_stdmeasure(::Type{<:Bind{<:Any,M}}) where {M} = preferred_stdmeasure(M) +_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_at(xy, maybestatic_length(a) + maybestatic_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")) + +# Test values follow the primary test value through the secondary measure: +function testvalue(::Type{T}, μ::Bind) where {T} + a = testvalue(T, μ.α) + _combine_variates(μ.f_c, a, testvalue(T, _get_β_a(μ, a))) +end + +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_rest = logdensityof_with_rest(μ, xy) + isempty(x_rest) || _throw_stream_too_long() + 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 + +# The secondary measure depends on the primary variate, so streams are +# consumed one by one: +@inline fixed_stream_size(::Type{<:Bind}) = static(false) +@inline mspace_ndims(::Type{<:_BindBy{typeof(vcat)}}) = 1 + +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_at(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 batched_logdensityof_with_rest(μ::Bind, x::AbstractVector, ::Tuple{}) + ℓ, _, x_rest = logdensityof_with_rest(μ, x) + return ℓ, x_rest +end + +batched_logdensityof_impl(μ::_BindBy{typeof(vcat)}, X::AbstractArray) = _streamwise_ld(logdensityof_impl, μ, X) + +# Batches of streams containing binds are consumed stream by stream (by +# the outermost stream combinator, see `fixed_stream_size`): +@noinline function batched_logdensityof_with_rest(::Bind, ::AbstractArray, ::SizeLike) + throw(ArgumentError("Batches of variate streams containing binds must be consumed stream by stream")) +end +batched_logdensityof_impl(μ::_BindBy{typeof(vcat)}, x::AbstractVector) = _bind_ld_impl(vcat, μ, x) + + +function rand_impl(ctx::GenContext, μ::Bind) + a = rand_impl(ctx, μ.α) + b = rand_impl(ctx, _get_β_a(μ, a)) + return _combine_variates(μ.f_c, a, b) +end + +# The secondary measure depends on the primary variate, so batches are +# generated variate by variate: +batched_rand_impl(ctx::GenContext, μ::Bind, sz::SizeLike) = _batched_rand_pointwise(ctx, μ, sz) + + +# Transport consumes the variate parts of the primary and secondary +# measure in a single pass, analogous to density evaluation: + +transport_to_std(::Type{S}, μ::Bind, ab) where {S<:StdMeasure} = _bind_to_std(S, μ.f_c, μ, ab) + +function _bind_to_std(::Type{S}, f_c, μ::Bind, ab) where {S} + tpm_α, a, b = tpmeasure_split_combined(f_c, μ.α, ab) + vcat(_as_stdstream(transport_to_std(S, tpm_α, a)), _as_stdstream(transport_to_std(S, _get_β_a(μ, a), b))) +end + +function _bind_to_std(::Type{S}, ::Union{typeof(vcat),typeof(merge)}, μ::Bind, ab) where {S} + z, _, x_rest = transport_to_std_with_rest(S, μ, ab) + if !isempty(x_rest) + throw(ArgumentError("Variate too long during transport of a bind")) + end + return z +end + +function transport_to_std_with_rest(::Type{S}, μ::_BindBy{typeof(vcat)}, x::AbstractVector) where {S<:StdMeasure} + z_a, a, x2 = transport_to_std_with_rest(S, μ.α, x) + z_b, _, x_rest = transport_to_std_with_rest(S, _get_β_a(μ, a), x2) + x_μ, _ = split_at(x, maybestatic_length(x) - maybestatic_length(x_rest)) + return vcat(z_a, z_b), x_μ, x_rest +end + +function transport_to_std_with_rest(::Type{S}, μ::_BindBy{typeof(merge)}, x::NamedTuple) where {S<:StdMeasure} + z_a, a, x2 = transport_to_std_with_rest(S, μ.α, x) + z_b, b, x_rest = transport_to_std_with_rest(S, _get_β_a(μ, a), x2) + return vcat(z_a, z_b), merge(a, b), x_rest +end + +function transport_from_std_with_rest(::Type{S}, μ::Bind, z::AbstractVector) where {S<:StdMeasure} + a, z2 = transport_from_std_with_rest(S, μ.α, z) + b, z_rest = transport_from_std_with_rest(S, _get_β_a(μ, a), z2) + return _combine_variates(μ.f_c, a, b), z_rest +end + +function transport_from_std(::Type{S}, μ::Bind, z::AbstractVector) where {S<:StdMeasure} + x, z_rest = transport_from_std_with_rest(S, μ, z) + isempty(z_rest) || _throw_std_length_mismatch() + return x +end + +# The secondary measure depends on the primary variate, so batches of +# streams are consumed stream by stream (by the outermost stream +# combinator, see `fixed_stream_size`): +function batched_transport_to_std_with_rest(::Type{S}, μ::Bind, X::AbstractArray, sz::SizeLike) where {S<:StdMeasure} + _bind_to_std_with_rest(S, μ, X, sz) +end +function _bind_to_std_with_rest(::Type{S}, μ::Bind, x::AbstractVector, ::Tuple{}) where {S} + z, _, x_rest = transport_to_std_with_rest(S, μ, x) + return z, x_rest +end +@noinline function _bind_to_std_with_rest(::Type{S}, ::Bind, ::AbstractArray, ::SizeLike) where {S} + throw(ArgumentError("Batches of variate streams containing binds must be consumed stream by stream")) end diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl new file mode 100644 index 00000000..3e85a3ae --- /dev/null +++ b/src/combinators/combined.jl @@ -0,0 +1,400 @@ +""" + 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_at(ab, maybestatic_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 + +function _mcombine_product_shortcut(::typeof(vcat), ma::AbstractVector{T}, mb::AbstractVector{T}, α, β) where {T} + isconcretetype(T) ? productmeasure(vcat(ma, mb)) : _generic_mcombine_impl_stage2(vcat, α, β) +end +_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 function preferred_stdmeasure(::Type{<:CombinedMeasure{<:Any,MA,MB}}) where {MA,MB} + promote_stdmeasure(preferred_stdmeasure(MA), preferred_stdmeasure(MB)) +end + + +@inline insupport(μ::CombinedMeasure, ab) = NoFastInsupport{typeof(μ)}() + +@inline function mspace_flatsize(μ::CombinedMeasure{typeof(vcat)}) + _vcat_flatsize(mspace_flatsize(μ.α), mspace_flatsize(μ.β)) +end + +@inline _vcat_flatsize(a::SizeLike, b::SizeLike) = canonical_size((size2length(a) + size2length(b),)) +@inline _vcat_flatsize(a::NoMSpaceElementSize, ::SizeLike) = a +@inline _vcat_flatsize(::SizeLike, b::NoMSpaceElementSize) = b +@inline _vcat_flatsize(a::NoMSpaceElementSize, ::NoMSpaceElementSize) = a + +@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(::typeof(vcat), μ::CombinedMeasure, ab::AbstractVector) + _point_result(_materialize(_combined_batched_ld(μ, ab, static(true))), μ) +end + +function _combined_ld_impl(::typeof(merge), μ::CombinedMeasure, ab::NamedTuple) + ℓ, _, x_rest = logdensityof_with_rest(μ, ab) + isempty(x_rest) || _throw_stream_too_long() + 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 + +@inline mspace_ndims(::Type{<:CombinedMeasure{typeof(vcat)}}) = 1 +@inline function fixed_stream_size(::Type{<:CombinedMeasure{<:Any,MA,MB}}) where {MA,MB} + fixed_stream_size(MA) & fixed_stream_size(MB) +end + +# Batches of vcat-combined variates are batches of streams: with fixed +# component sizes the whole batch is consumed in fused operations, +# otherwise stream by stream. +@inline function batched_logdensityof_impl(μ::CombinedMeasure{typeof(vcat)}, X::AbstractArray) + _combined_batched_ld(μ, X, fixed_stream_size(μ)) +end +function _combined_batched_ld(μ::CombinedMeasure, X::AbstractArray, ::True) + ℓ, X_rest = batched_logdensityof_with_rest(μ, X, ()) + size(X_rest, 1) == 0 || _throw_stream_too_long() + return ℓ +end +_combined_batched_ld(μ::CombinedMeasure, X::AbstractVector, ::False) = _combined_batched_ld(μ, X, static(true)) +_combined_batched_ld(μ::CombinedMeasure, X::AbstractArray, ::False) = _streamwise_ld(logdensityof_impl, μ, X) + +function batched_logdensityof_with_rest(μ::CombinedMeasure{typeof(vcat)}, X::AbstractArray, ::Tuple{}) + _combined_ld_with_rest(μ, X) +end +batched_logdensityof_with_rest(μ::CombinedMeasure{typeof(vcat)}, x::AbstractVector, ::Tuple{}) = _combined_ld_with_rest(μ, x) +function _combined_ld_with_rest(μ::CombinedMeasure, X::AbstractArray) + ℓ_a, X2 = batched_logdensityof_with_rest(μ.α, X, ()) + ℓ_b, X_rest = batched_logdensityof_with_rest(μ.β, X2, ()) + return _lazy_add(ℓ_a, ℓ_b), X_rest +end + +# Several variates per stream interleave the component parts, so the rows +# of each variate are split by the fixed component sizes: +function batched_logdensityof_with_rest(μ::CombinedMeasure{typeof(vcat)}, X::AbstractArray, sz::SizeLike) + n_a, n_b = _fixed_stream_length(μ.α), _fixed_stream_length(μ.β) + X_μ, X_rest = _batched_split(X, _chunk_rows(n_a + n_b, sz)) + X_v = maybestatic_reshape(X_μ, (n_a + n_b, size_dims(sz)..., Base.tail(_batch_dims(X_μ))...)) + X_a, X_b = _batched_split(X_v, n_a) + ℓ_a, _ = batched_logdensityof_with_rest(μ.α, X_a, ()) + ℓ_b, _ = batched_logdensityof_with_rest(μ.β, X_b, ()) + return _lazy_add(ℓ_a, ℓ_b), X_rest +end + +@inline _fixed_stream_length(μ) = _fixed_stream_length(μ, mspace_flatsize(μ)) +@inline _fixed_stream_length(μ, sz::SizeLike) = size2length(sz) +@noinline function _fixed_stream_length(μ, ::NoMSpaceElementSize) + throw(ArgumentError("Consuming several variates per stream requires measures of type $(nameof(typeof(μ))) to have a known variate size")) +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_at(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 + + +rand_impl(ctx::GenContext, μ::CombinedMeasure) = _combine_variates(μ.f_c, rand_impl(ctx, μ.α), rand_impl(ctx, μ.β)) + +# Variates of vcat-combined measures are flat streams, nested variates of +# the components (e.g. of powers of measures with value-dependent sizes) +# are flattened: +@inline _combine_variates(f_c, a, b) = f_c(a, b) +@inline _combine_variates(::typeof(vcat), a, b) = vcat(_flat_stream(a), _flat_stream(b)) +@inline _flat_stream(x::Number) = x +@inline _flat_stream(x::AbstractArray{<:Number}) = vec(x) +@inline _flat_stream(x::AbstractArray) = reduce(vcat, map(_flat_stream, x)) +@inline _flat_stream(x::Union{Tuple,NamedTuple}) = reduce(vcat, map(_flat_stream, values(x))) + +batched_rand_impl(ctx::GenContext, μ::CombinedMeasure, sz::SizeLike) = _batched_rand_pointwise(ctx, μ, sz) + +# Batches of merge-combined measures merge the named tuples of batches: +function batched_rand_impl(ctx::GenContext, μ::CombinedMeasure{typeof(merge)}, sz::SizeLike) + merge(batched_rand_impl(ctx, μ.α, sz), batched_rand_impl(ctx, μ.β, sz)) +end + +# Batches of vcat-combined measures are concatenated along the streams: +function batched_rand_impl(ctx::GenContext, μ::CombinedMeasure{typeof(vcat)}, sz::SizeLike) + _combined_batched_rand(ctx, μ, sz, fixed_stream_size(μ)) +end +function _combined_batched_rand(ctx::GenContext, μ::CombinedMeasure, sz::SizeLike, ::True) + A = _as_stream_batch(batched_rand_impl(ctx, μ.α, sz), μ.α) + B = _as_stream_batch(batched_rand_impl(ctx, μ.β, sz), μ.β) + return vcat(A, B) +end +_combined_batched_rand(ctx::GenContext, μ::CombinedMeasure, sz::SizeLike, ::False) = _batched_rand_pointwise(ctx, μ, sz) + + +# Transport consumes the variate parts of both component measures in a +# single pass, analogous to density evaluation: + +transport_to_std(::Type{S}, μ::CombinedMeasure, ab) where {S<:StdMeasure} = _combined_to_std(S, μ.f_c, μ, ab) + +function _combined_to_std(::Type{S}, f_c, μ::CombinedMeasure, ab) where {S} + tpm_α, a, b = tpmeasure_split_combined(f_c, μ.α, ab) + vcat(_as_stdstream(transport_to_std(S, tpm_α, a)), _as_stdstream(transport_to_std(S, μ.β, b))) +end + +function _combined_to_std(::Type{S}, ::Union{typeof(vcat),typeof(merge)}, μ::CombinedMeasure, ab) where {S} + z, _, x_rest = transport_to_std_with_rest(S, μ, ab) + if !isempty(x_rest) + throw(ArgumentError("Variate too long during transport of a combined measure")) + end + return z +end + +function transport_to_std_with_rest(::Type{S}, μ::CombinedMeasure{typeof(vcat)}, x::AbstractVector) where {S<:StdMeasure} + z_a, _, x2 = transport_to_std_with_rest(S, μ.α, x) + z_b, _, x_rest = transport_to_std_with_rest(S, μ.β, x2) + x_μ, _ = split_at(x, maybestatic_length(x) - maybestatic_length(x_rest)) + return vcat(z_a, z_b), x_μ, x_rest +end + +function transport_to_std_with_rest(::Type{S}, μ::CombinedMeasure{typeof(merge)}, x::NamedTuple) where {S<:StdMeasure} + z_a, a, x2 = transport_to_std_with_rest(S, μ.α, x) + z_b, b, x_rest = transport_to_std_with_rest(S, μ.β, x2) + return vcat(z_a, z_b), merge(a, b), x_rest +end + +function transport_from_std_with_rest(::Type{S}, μ::CombinedMeasure, z::AbstractVector) where {S<:StdMeasure} + a, z2 = transport_from_std_with_rest(S, μ.α, z) + b, z_rest = transport_from_std_with_rest(S, μ.β, z2) + return _combine_variates(μ.f_c, a, b), z_rest +end + +function transport_from_std(::Type{S}, μ::CombinedMeasure, z::AbstractVector) where {S<:StdMeasure} + x, z_rest = transport_from_std_with_rest(S, μ, z) + isempty(z_rest) || _throw_std_length_mismatch() + return x +end + + +# Batches of vcat-combined variates are batches of streams: with fixed +# component sizes the whole batch is consumed in fused operations, +# otherwise stream by stream. + +function batched_transport_to_std(::Type{S}, μ::CombinedMeasure{typeof(vcat)}, X::AbstractArray) where {S<:StdMeasure} + _combined_batched_to_std(S, μ, X, fixed_stream_size(μ)) +end +function _combined_batched_to_std(::Type{S}, μ::CombinedMeasure, X::AbstractArray, ::True) where {S} + Z, X_rest = batched_transport_to_std_with_rest(S, μ, X, ()) + size(X_rest, 1) == 0 || _throw_stream_too_long() + return Z +end +function _combined_batched_to_std(::Type{S}, μ::CombinedMeasure, x::AbstractVector, ::False) where {S} + _combined_batched_to_std(S, μ, x, static(true)) +end +function _combined_batched_to_std(::Type{S}, μ::CombinedMeasure, X::AbstractArray, ::False) where {S} + stacked(map(x -> _combined_batched_to_std(S, μ, x, static(true)), sliced(X, Val(1)))) +end + +function batched_transport_to_std_with_rest(::Type{S}, μ::CombinedMeasure{typeof(vcat)}, X::AbstractArray, sz::SizeLike) where {S<:StdMeasure} + _combined_to_std_with_rest(S, μ, X, sz) +end +function _combined_to_std_with_rest(::Type{S}, μ::CombinedMeasure, X::AbstractArray, ::Tuple{}) where {S} + Z_a, X2 = batched_transport_to_std_with_rest(S, μ.α, X, ()) + Z_b, X_rest = batched_transport_to_std_with_rest(S, μ.β, X2, ()) + return vcat(Z_a, Z_b), X_rest +end + +# Several variates per stream interleave the component parts, so the rows +# of each variate are split by the fixed component sizes: +function _combined_to_std_with_rest(::Type{S}, μ::CombinedMeasure, X::AbstractArray, sz::SizeLike) where {S} + n_rows = _fixed_stream_length(μ.α) + _fixed_stream_length(μ.β) + X_μ, X_rest = _batched_split(X, _chunk_rows(n_rows, sz)) + X_v = maybestatic_reshape(X_μ, (n_rows, size_dims(sz)..., Base.tail(_batch_dims(X_μ))...)) + Z, _ = _combined_to_std_with_rest(S, μ, X_v, ()) + return _merge_multiplicity(Z, sz), X_rest +end + +function batched_transport_from_std(::Type{S}, μ::CombinedMeasure{typeof(vcat)}, Z::AbstractArray) where {S<:StdMeasure} + _combined_batched_from_std(S, μ, Z, fixed_stream_size(μ)) +end +function _combined_batched_from_std(::Type{S}, μ::CombinedMeasure, Z::AbstractArray, ::True) where {S} + X, Z_rest = batched_transport_from_std_with_rest(S, μ, Z, ()) + size(Z_rest, 1) == 0 || _throw_std_length_mismatch() + return X +end +function _combined_batched_from_std(::Type{S}, μ::CombinedMeasure, z::AbstractVector, ::False) where {S} + transport_from_std(S, μ, z) +end +function _combined_batched_from_std(::Type{S}, μ::CombinedMeasure, Z::AbstractArray, ::False) where {S} + stacked(map(z -> transport_from_std(S, μ, z), sliced(Z, Val(1)))) +end + +function batched_transport_from_std_with_rest(::Type{S}, μ::CombinedMeasure{typeof(vcat)}, Z::AbstractArray, sz::SizeLike) where {S<:StdMeasure} + _combined_from_std_with_rest(S, μ, Z, sz) +end +# Single streams yield a variate via the point protocol: +function _combined_from_std_with_rest(::Type{S}, μ::CombinedMeasure, z::AbstractVector, ::Tuple{}) where {S} + transport_from_std_with_rest(S, μ, z) +end +function _combined_from_std_with_rest(::Type{S}, μ::CombinedMeasure, Z::AbstractArray, ::Tuple{}) where {S} + _combined_batch_from_std_with_rest(S, μ, Z, fixed_stream_size(μ)) +end +function _combined_batch_from_std_with_rest(::Type{S}, μ::CombinedMeasure, Z::AbstractArray, ::True) where {S} + A, Z2 = batched_transport_from_std_with_rest(S, μ.α, Z, ()) + B, Z_rest = batched_transport_from_std_with_rest(S, μ.β, Z2, ()) + return vcat(_as_stream_batch(A, μ.α), _as_stream_batch(B, μ.β)), Z_rest +end +# Components without fixed variate sizes are consumed stream by stream: +function _combined_batch_from_std_with_rest(::Type{S}, μ::CombinedMeasure, Z::AbstractArray, ::False) where {S} + results = map(z -> transport_from_std_with_rest(S, μ, z), sliced(Z, Val(1))) + return stacked(map(first, results)), stacked(map(last, results)) +end +function _combined_from_std_with_rest(::Type{S}, μ::CombinedMeasure, Z::AbstractArray, sz::SizeLike) where {S} + _batched_from_std_bydof(S, μ, Z, sz, fast_dof(μ)) +end + +Adapt.adapt_structure(to, μ::CombinedMeasure) = mcombine(μ.f_c, Adapt.adapt(to, μ.α), Adapt.adapt(to, μ.β)) diff --git a/src/combinators/half.jl b/src/combinators/half.jl index 24063b76..9578366d 100644 --- a/src/combinators/half.jl +++ b/src/combinators/half.jl @@ -4,6 +4,13 @@ struct Half{M} <: AbstractMeasure parent::M end +@inline mspace_elsize(μ::Half) = mspace_elsize(μ.parent) +@inline mspace_flatsize(μ::Half) = mspace_flatsize(μ.parent) +@inline mspace_flatsize(::Type{<:Half{M}}) where {M} = mspace_flatsize(M) +@inline mspace_ndims(::Type{<:Half{M}}) where {M} = mspace_ndims(M) +@inline fixed_stream_size(::Type{<:Half{M}}) where {M} = fixed_stream_size(M) +@inline preferred_stdmeasure(::Type{<:Half}) = StdUniform + function Base.show(io::IO, μ::Half) print(io, "Half") show(io, μ.parent) @@ -15,11 +22,10 @@ unhalf(μ::Half) = μ.parent weightedmeasure(logtwo, basemeasure(unhalf(μ))) end -function Base.rand(rng::AbstractRNG, ::Type{T}, μ::Half) where {T} - return abs(rand(rng, T, unhalf(μ))) -end +@inline rand_impl(ctx::GenContext, μ::Half) = abs(rand_impl(ctx, unhalf(μ))) +@inline batched_rand_impl(ctx::GenContext, μ::Half, sz::SizeLike) = abs.(batched_rand_impl(ctx, unhalf(μ), sz)) -function logdensityof(μ::Half, x) +function logdensityof_impl(μ::Half, x) ld = logdensityof(unhalf(μ), x) - loghalf return x ≥ 0 ? ld : oftype(ld, -Inf) end @@ -40,9 +46,8 @@ function smf(μ::Half, x) end function invsmf(μ::Half, p) - @assert zero(p) ≤ p ≤ one(p) - invsmf(μ.parent, (p + 1) / 2) + _nan_outside(StdUniform(), p, invsmf(μ.parent, _unit_interior((p + 1) / 2))) end -transport_def(μ::Half, ::StdUniform, p) = invsmf(μ, p) -transport_def(::StdUniform, μ::Half, x) = smf(μ, x) +@inline transport_to_std(::Type{StdUniform}, μ::Half, x) = _nan_outside(μ, x, smf(μ, x)) +@inline transport_from_std(::Type{StdUniform}, μ::Half, p) = invsmf(μ, p) diff --git a/src/combinators/implicitlymapped.jl b/src/combinators/implicitlymapped.jl index 964ea466..721d7268 100644 --- a/src/combinators/implicitlymapped.jl +++ b/src/combinators/implicitlymapped.jl @@ -84,7 +84,7 @@ export ImplicitlyMapped Get the original object (a measure or transition/Markov kernel) that was implicitly mapped. -See [ImplicitlyMapped](@ref) for detailed semantics. +See [`ImplicitlyMapped`](@ref) for detailed semantics. # Implementation @@ -100,7 +100,7 @@ export implicit_origin Get an explicit map/function based on an implicitly mapped object and an observation. -See [ImplicitlyMapped](@ref) for detailed semantics. +See [`ImplicitlyMapped`](@ref) for detailed semantics. # Implementation @@ -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) @@ -206,7 +206,7 @@ Constructors: * `Marginalized(mu)` * `Marginalized(f_kernel)` -See [ImplicitlyMapped](@ref) for detailed semantics. +See [`ImplicitlyMapped`](@ref) for detailed semantics. Example: @@ -227,7 +227,7 @@ export Marginalized implicit_origin(mapped::Marginalized) = mapped.obj function explicit_mapfunc(::Marginalized, obs::NamedTuple{names}) where {names} - PropSelFunction{names,names}() + PropSelFunction(names...) end function pushfwd(f::PropSelFunction, mu::ProductMeasure{<:NamedTuple}, ::PushfwdRootMeasure) productmeasure(f(marginals(mu))) 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..a86a06ea 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) @@ -27,42 +51,42 @@ function Pretty.tile(μ::PowerMeasure) return Pretty.pair_layout(arg1, arg2; sep = " ^ ") 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( - rng::AbstractRNG, - ::Type{T}, - d::PowerMeasure{M}, -) where {T,M<:AbstractMeasure} - map(_cartidxs(d.axes)) do _ - rand(rng, T, d.parent) - end +# Variates of powers are generated as one flat batch of variates of the +# base measure, with the power's size as additional batch dimensions. Base +# measures without fixed variate sizes generate their variates one by one. + +rand_impl(ctx::GenContext, μ::PowerMeasure) = _pwr_rand(ctx, μ, fixed_stream_size(pwr_base(μ))) +_pwr_rand(ctx::GenContext, μ::PowerMeasure, ::True) = _pwr_variate(μ, batched_rand_impl(ctx, μ, ())) +function _pwr_rand(ctx::GenContext, μ::PowerMeasure, ::False) + ν = pwr_base(μ) + map(_ -> rand_impl(ctx, ν), _cartidxs(pwr_axes(μ))) end -function Base.rand(rng::AbstractRNG, ::Type{T}, d::PowerMeasure) where {T} - map(_cartidxs(d.axes)) do _ - rand(rng, d.parent) - end +function batched_rand_impl(ctx::GenContext, μ::PowerMeasure, sz::SizeLike) + _pwr_batched_rand(ctx, μ, sz, fixed_stream_size(pwr_base(μ))) end +function _pwr_batched_rand(ctx::GenContext, μ::PowerMeasure, sz::SizeLike, ::True) + batched_rand_impl(ctx, pwr_base(μ), (size_dims(pwr_size(μ))..., size_dims(sz)...)) +end +_pwr_batched_rand(ctx::GenContext, μ::PowerMeasure, sz::SizeLike, ::False) = _batched_rand_pointwise(ctx, μ, sz) -@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 +marginals(d::PowerMeasure) = maybestatic_fill(d.parent, d.axes) -@inline function powermeasure(x::T, sz::Tuple{Vararg{Any,N}}) where {T,N} - PowerMeasure(x, _pm_axes(sz)) +@inline mspace_elsize(μ::PowerMeasure) = pwr_size(μ) +@inline mspace_flatsize(μ::PowerMeasure) = _cat_sizes(mspace_flatsize(pwr_base(μ)), pwr_size(μ)) +@inline function mspace_flatsize(::Type{<:PowerMeasure{M,A}}) where {M,A<:Tuple{Vararg{StaticOneToLike}}} + _cat_sizes(mspace_flatsize(M), axes2size(A)) end -marginals(d::PowerMeasure) = fill_with(d.parent, d.axes) - function Base.:^(μ::AbstractMeasure, dims::Tuple{Vararg{AbstractArray,N}}) where {N} powermeasure(μ, dims) end -Base.:^(μ::AbstractMeasure, dims::Tuple) = powermeasure(μ, one_to.(dims)) +Base.:^(μ::AbstractMeasure, dims::Tuple) = powermeasure(μ, maybestatic_oneto.(dims)) Base.:^(μ::AbstractMeasure, n) = powermeasure(μ, (n,)) # Base.show(io::IO, d::PowerMeasure) = print(io, d.parent, " ^ ", size(d.xs)) @@ -78,74 +102,326 @@ 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 +# Numeric flat variates of powers of measures with fixed stream sizes but +# no variate rank (e.g. tuple products) are streams: +@inline function mspace_ndims(::Type{<:PowerMeasure{M,A}}) where {M,A<:Tuple} + _pwr_ndims(mspace_ndims(M), fieldcount(A), fixed_stream_size(M)) +end +@inline _pwr_ndims(n::Integer, k::Integer, ::Any) = n + k +@inline _pwr_ndims(::NoMSpaceElementSize, ::Integer, ::True) = 1 +@inline _pwr_ndims(n::NoMSpaceElementSize, ::Integer, ::False) = n + +# Local measures of powers at nested variates are products of the local +# measures of the elements: +@inline localmeasure(μ::PowerMeasure, ::AbstractArray{<:Number}) = μ +function localmeasure(μ::PowerMeasure, x::AbstractArray) + size(x) == asnonstatic(pwr_size(μ)) || return μ + productmeasure(map(Base.Fix1(localmeasure, pwr_base(μ)), x)) +end +@inline fixed_stream_size(::Type{<:PowerMeasure{M}}) where {M} = fixed_stream_size(M) + +# The innermost base measure of nested powers and the total number of power +# dimensions: +@inline _pwr_unwrap(μ) = (μ, static(0)) +@inline function _pwr_unwrap(μ::PowerMeasure) + ν, n = _pwr_unwrap(pwr_base(μ)) + ν, n + static(length(pwr_axes(μ))) +end + +# Batched kernels: the base kernel runs over the flat batch, the power then +# sums the leading dimensions of the result that belong to its axes. +@inline function _powered_kernel(f::F, μ::PowerMeasure, X) where {F} + _check_pwr_batch(X, μ) + _powered_kernel_impl(f, μ, X, _static_ndims(pwr_base(μ))) +end +@inline function _powered_kernel_impl(f::F, μ::PowerMeasure, X, ::Any) where {F} + _pwr_sum_dims(_batched_kernel(f, pwr_base(μ), X), static(length(pwr_axes(μ)))) +end - @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 +# `sum_leading_dims`, but naming the measure when the variate isn't an array: +@inline _pwr_sum_dims(ℓ, n::StaticInteger) = sum_leading_dims(ℓ, n) +@inline _pwr_sum_dims(ℓ::Number, ::StaticInteger{0}) = ℓ +@noinline _pwr_sum_dims(::Number, ::StaticInteger) = _throw_pwr_variate_not_array() +# Numeric batches of powers of bases without a variate rank are batches of +# streams: +@inline function _powered_kernel_impl(::typeof(logdensityof_impl), μ::PowerMeasure, X::AbstractArray{<:Number}, ::NoMSpaceElementSize) + _powered_stream_kernel(μ, X, fixed_stream_size(pwr_base(μ))) +end +function _powered_stream_kernel(μ::PowerMeasure, X::AbstractArray, ::True) + ℓ, X_rest = batched_logdensityof_with_rest(μ, X, ()) + size(X_rest, 1) == 0 || _throw_stream_too_long() + return ℓ +end +_powered_stream_kernel(μ::PowerMeasure, X::AbstractArray, ::False) = _streamwise_ld(logdensityof_impl, μ, X) + +# Flat batches of powers have the power dimensions after the variate +# dimensions of the base measure (where the rank of the base is known): +@inline _check_pwr_batch(X::AbstractArray, μ::PowerMeasure) = _check_pwr_dims(X, _static_ndims(pwr_base(μ)), asnonstatic(pwr_size(μ)), false) +@inline _check_pwr_batch(::Any, ::PowerMeasure) = nothing +@inline function _check_pwr_dims(X::AbstractArray, ::StaticInteger{K}, dims::Dims, exact::Bool) where {K} + n = length(dims) + if (exact ? ndims(X) != K + n : ndims(X) < K + n) || ntuple(i -> size(X, K + i), Val(length(dims))) != dims + _throw_size_mismatch() end + return nothing +end +@inline _check_pwr_dims(::AbstractArray, ::NoMSpaceElementSize, ::Dims, ::Bool) = nothing +@inline batched_logdensityof_impl(μ::PowerMeasure, X) = _powered_kernel(logdensityof_impl, μ, X) +@inline batched_logdensity_def(μ::PowerMeasure, X) = _powered_kernel(logdensity_def, μ, X) + +# Point evaluation: flat variates are batches with zero batch dimensions, +# nested variates without flat storage sum the point densities of the base. +@inline _point_ld(f::F, μ::PowerMeasure, x::AbstractArray{<:Number}) where {F} = f(μ, x) +@inline logdensityof_impl(μ::PowerMeasure, x) = _powered_point(logdensityof_impl, μ, x) +@inline logdensity_def(μ::PowerMeasure, x) = _powered_point(logdensity_def, μ, x) + +@inline function _powered_point(f::F, μ::PowerMeasure, x::AbstractArray{<:Number}) where {F} + _point_result(_materialize(_batched_kernel(f, μ, x)), μ) +end +@inline function _powered_point(f::F, μ::PowerMeasure, x::AbstractArray) where {F} + _check_pwr_shape(μ, x) + _powered_point_nested(f, μ, x, _flat_storage(x)) +end +@inline function _powered_point_nested(f::F, μ::PowerMeasure, x, x_flat::Union{AbstractArray,Tuple,NamedTuple}) where {F} + _point_result(_materialize(_batched_kernel(f, μ, x_flat)), μ) +end +function _powered_point_nested(f::F, μ::PowerMeasure, x::AbstractArray, ::NoFlatStorage) where {F} + ν = pwr_base(μ) + sum(_PointLogd(f, ν), x; init = zero(_logd_numtype(x))) +end +@noinline _powered_point(::F, ::PowerMeasure, x) where {F} = _throw_pwr_variate_not_array() - @eval @inline function $func( - d::PowerMeasure{M,NTuple{N,Static.SOneTo{0}}}, - x, - ) where {M,N} - static(0.0) +@noinline _throw_pwr_variate_not_array() = + throw(ArgumentError("Variates of powers of measures must be arrays")) + +# Nested variates have the power's shape: +@inline function _check_pwr_shape(μ::PowerMeasure, x::AbstractArray) + if maybestatic_size(x) != pwr_size(μ) + _throw_size_mismatch() end + return nothing end -@inline function insupport(μ::PowerMeasure, x) - p = μ.parent - all(x) do xj - # https://github.com/SciML/Static.jl/issues/36 - dynamic(insupport(p, xj)) +# Streams: a power consumes its size times the variates of the base measure +# and sums the base results over its axes. Bases without fixed variate +# sizes are consumed element by element, for single streams. +function batched_logdensityof_with_rest(μ::PowerMeasure, X::AbstractArray, sz::SizeLike) + _powered_ld_with_rest(μ, X, sz, fixed_stream_size(pwr_base(μ))) +end +function batched_logdensityof_with_rest(μ::PowerMeasure, x::AbstractVector, sz::Tuple{}) + _powered_ld_with_rest(μ, x, sz, fixed_stream_size(pwr_base(μ))) +end +function _powered_ld_with_rest(μ::PowerMeasure, X::AbstractArray, sz::SizeLike, ::True) + ℓ, X_rest = batched_logdensityof_with_rest(pwr_base(μ), X, (size_dims(pwr_size(μ))..., size_dims(sz)...)) + return _pwr_sum_dims(ℓ, static(length(pwr_axes(μ)))), X_rest +end +function _powered_ld_with_rest(μ::PowerMeasure, x::AbstractVector, ::Tuple{}, ::False) + ν = pwr_base(μ) + ℓ = zero(_logd_numtype(x)) + x_rest = x + for _ in 1:length(marginals(μ)) + ℓ_i, _, x_rest = logdensityof_with_rest(ν, x_rest) + ℓ += ℓ_i end + return ℓ, x_rest +end +@noinline function _powered_ld_with_rest(μ::PowerMeasure, ::AbstractArray, ::SizeLike, ::False) + throw(ArgumentError("Batches of variate streams containing powers of measures of type $(nameof(typeof(pwr_base(μ)))) must be consumed stream by stream")) end +# Support checks of powers run over the flat variate storage where the +# innermost base measure has scalar variates, elementwise otherwise: @inline function insupport(μ::PowerMeasure, x::AbstractArray) - p = μ.parent - all(x) do xj - # https://github.com/SciML/Static.jl/issues/36 - dynamic(insupport(p, xj)) - end + ν, _ = _pwr_unwrap(μ) + _powered_insupport(μ, x, _flat_storage(x), _static_ndims(ν)) +end + +@inline function _powered_insupport(μ::PowerMeasure, x, x_flat::AbstractArray, ::StaticInteger{0}) + ν, _ = _pwr_unwrap(μ) + _all_insupport(broadcast(_insupport_bool ∘ Base.Fix1(insupport, ν), x_flat)) +end +@inline function _powered_insupport(μ::PowerMeasure, x, x_flat::AbstractArray, ::StaticInteger{K}) where {K} + ν, _ = _pwr_unwrap(μ) + _all_insupport(map(_insupport_bool ∘ Base.Fix1(insupport, ν), sliced(x_flat, Val(K)))) +end +@inline function _powered_insupport(μ::PowerMeasure, x, ::AbstractArray{<:Number}, ::NoMSpaceElementSize) + NoFastInsupport{typeof(μ)}() end +@inline _powered_insupport(μ::PowerMeasure, x, ::Any, ::Any) = _powered_insupport_elementwise(pwr_base(μ), x) -@inline getdof(μ::PowerMeasure) = getdof(μ.parent) * prod(map(length, μ.axes)) +@inline function _powered_insupport_elementwise(ν, x::AbstractArray) + _all_insupport(broadcast(_insupport_bool ∘ Base.Fix1(insupport, ν), x)) +end + +function insupport(μ::PowerMeasure, x) + mapreduce(_insupport_bool ∘ Base.Fix1(insupport, pwr_base(μ)), _insupport_and, x) +end -@inline function getdof(::PowerMeasure{<:Any,NTuple{N,Static.SOneTo{0}}}) where {N} +@inline getdof(μ::PowerMeasure) = getdof(μ.parent) * size2length(axes2size(μ.axes)) +@inline fast_dof(μ::PowerMeasure) = fast_dof(μ.parent) * size2length(axes2size(μ.axes)) + +# 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 +# Variates may be nested arrays of the power's shape or their flat storage: @propagate_inbounds function checked_arg(μ::PowerMeasure, x::AbstractArray{<:Any}) - @boundscheck begin - sz_μ = map(length, μ.axes) - sz_x = size(x) - if sz_μ != sz_x - throw(ArgumentError("Size of variate doesn't match size of power measure")) - end - end + @boundscheck _check_pwr_variate(μ, x) return x end -function checked_arg(μ::PowerMeasure, x::Any) - throw(ArgumentError("Size of variate doesn't match size of power measure")) +@inline function _check_pwr_variate(μ::PowerMeasure, x::AbstractArray) + if maybestatic_size(x) != pwr_size(μ) + _check_pwr_flat(x, _static_ndims(pwr_base(μ)), asnonstatic(pwr_size(μ))) + end + return nothing +end +@inline _check_pwr_flat(x::AbstractArray, k::StaticInteger, dims::Dims) = _check_pwr_dims(x, k, dims, true) +@inline _check_pwr_flat(::AbstractArray, ::NoMSpaceElementSize, ::Dims) = _throw_size_mismatch() + +checked_arg(μ::PowerMeasure, x::Any) = _throw_size_mismatch() + +massof(m::PowerMeasure) = massof(m.parent)^dynamic(size2length(pwr_size(m))) + + +# Transport: the standard variate of a power is the flat vector of the +# standard variates of its innermost base measure, in the order of the flat +# variate storage. Batches transport over the flat storage `(base variate +# dims..., power dims..., batch dims...)`. + +function batched_transport_to_std(::Type{S}, μ::PowerMeasure, X::AbstractArray) where {S<:StdMeasure} + _check_pwr_batch(X, μ) + _pwr_batched_to_std(S, μ, X, _static_ndims(pwr_base(μ))) +end +function batched_transport_to_std(::Type{S}, μ::PowerMeasure, X::Union{Tuple,NamedTuple}) where {S<:StdMeasure} + _pwr_batched_to_std(S, μ, X, nothing) +end +@inline function _pwr_batched_to_std(::Type{S}, μ::PowerMeasure, X, ::Any) where {S} + ν, n = _pwr_unwrap(μ) + merge_leading_dims(batched_transport_to_std(S, ν, X), static(1) + n) +end +# Numeric batches of powers of bases without a variate rank are batches of +# streams: +function _pwr_batched_to_std(::Type{S}, μ::PowerMeasure, X::AbstractArray{<:Number}, ::NoMSpaceElementSize) where {S} + Z, X_rest = batched_transport_to_std_with_rest(S, μ, X, ()) + size(X_rest, 1) == 0 || _throw_stream_too_long() + return Z +end + +function batched_transport_from_std(::Type{S}, μ::PowerMeasure, Z::AbstractArray) where {S<:StdMeasure} + ν, _ = _pwr_unwrap(μ) + dims = _pwr_dims(μ) + n_rows = _batch_dims(Z)[1] + dof_ν = _base_dof(n_rows, prod(dims)) + dof_ν * prod(dims) == n_rows || _throw_std_length_mismatch() + batched_transport_from_std(S, ν, maybestatic_reshape(Z, (dof_ν, dims..., Base.tail(_batch_dims(Z))...))) end -massof(m::PowerMeasure) = massof(m.parent)^prod(m.axes) +# Empty powers leave the degrees of freedom of the base undetermined: +@inline _base_dof(n_rows::IntegerLike, n_pwr::IntegerLike) = n_rows ÷ max(n_pwr, one(n_pwr)) -logdensity_def(::PowerMeasure{P}, x) where {P<:PrimitiveMeasure} = static(0.0) +# All power dimensions of nested powers, innermost first: +@inline _pwr_dims(μ::PowerMeasure) = (_pwr_dims(pwr_base(μ))..., size_dims(pwr_size(μ))...) +@inline _pwr_dims(ν) = () -# To avoid ambiguities -function logdensity_def( - ::PowerMeasure{P,Tuple{Vararg{Static.SOneTo{0},N}}}, - x, -) where {P<:PrimitiveMeasure,N} - static(0.0) +# Point transport: flat variates are batches with zero batch dimensions, +# nested variates without flat storage transport element by element. +function transport_to_std(::Type{S}, μ::PowerMeasure, x::AbstractArray) where {S<:StdMeasure} + _pwr_to_std(S, μ, x, _flat_storage(x)) +end +@inline function _pwr_to_std(::Type{S}, μ::PowerMeasure, x, x_flat::AbstractArray) where {S} + _single_std(batched_transport_to_std(S, μ, x_flat)) +end +function _pwr_to_std(::Type{S}, μ::PowerMeasure, x::AbstractArray, ::NoFlatStorage) where {S} + _check_pwr_shape(μ, x) + _flat_std_of(map(Base.Fix1(_ToStd{S}(), pwr_base(μ)), x)) +end + +function transport_from_std(::Type{S}, μ::PowerMeasure, z::AbstractVector) where {S<:StdMeasure} + _pwr_variate(μ, batched_transport_from_std(S, μ, z)) +end + +# Streams: a power consumes the variates of its base measure with its size +# as additional multiplicity. Bases without fixed variate sizes are +# consumed element by element, for single streams. +function batched_transport_to_std_with_rest(::Type{S}, μ::PowerMeasure, X::AbstractArray, sz::SizeLike) where {S<:StdMeasure} + _pwr_to_std_with_rest(S, μ, X, sz, fixed_stream_size(pwr_base(μ))) +end +@inline function _pwr_to_std_with_rest(::Type{S}, μ::PowerMeasure, X::AbstractArray, sz::SizeLike, ::True) where {S} + batched_transport_to_std_with_rest(S, pwr_base(μ), X, (size_dims(pwr_size(μ))..., size_dims(sz)...)) +end +function _pwr_to_std_with_rest(::Type{S}, μ::PowerMeasure, x::AbstractVector, ::Tuple{}, ::False) where {S} + z, _, x_rest = transport_to_std_with_rest(S, μ, x) + return z, x_rest +end +@noinline function _pwr_to_std_with_rest(::Type{S}, μ::PowerMeasure, ::AbstractArray, ::SizeLike, ::False) where {S} + throw(ArgumentError("Batches of variate streams containing powers of measures of type $(nameof(typeof(pwr_base(μ)))) must be consumed stream by stream")) +end + +function transport_to_std_with_rest(::Type{S}, μ::PowerMeasure, x::AbstractVector) where {S<:StdMeasure} + _pwr_point_to_std_with_rest(S, μ, x, fixed_stream_size(pwr_base(μ))) +end +function _pwr_point_to_std_with_rest(::Type{S}, μ::PowerMeasure, x::AbstractVector, ::True) where {S} + _pwr_point_to_std_with_rest(S, μ, x, _static_ndims(pwr_base(μ))) +end +function _pwr_point_to_std_with_rest(::Type{S}, μ::PowerMeasure, x::AbstractVector, ::StaticInteger) where {S} + x_μ, x_rest = _consume_from_stream(x, _stream_consume_size(μ)) + return _as_stdstream(transport_to_std(S, μ, x_μ)), x_μ, x_rest +end +# Bases without a variate rank (tuple products) consume streams via the +# batched protocol: +function _pwr_point_to_std_with_rest(::Type{S}, μ::PowerMeasure, x::AbstractVector, ::NoMSpaceElementSize) where {S} + z, x_rest = _pwr_to_std_with_rest(S, μ, x, (), static(true)) + x_μ, _ = split_at(x, maybestatic_length(x) - maybestatic_length(x_rest)) + return z, x_μ, x_rest +end +function _pwr_point_to_std_with_rest(::Type{S}, μ::PowerMeasure, x::AbstractVector, ::False) where {S} + ν = pwr_base(μ) + zs = Vector{Any}(undef, length(marginals(μ))) + x_rest = x + for i in eachindex(zs) + zs[i], _, x_rest = transport_to_std_with_rest(S, ν, x_rest) + end + x_μ, _ = split_at(x, maybestatic_length(x) - maybestatic_length(x_rest)) + return reduce(vcat, [z for z in zs]), x_μ, x_rest +end + +# Powers of measures without fast degrees of freedom transport their +# elements sequentially: +function transport_from_std_with_rest(::Type{S}, μ::PowerMeasure, z::AbstractVector) where {S<:StdMeasure} + _pwr_from_std_with_rest(S, μ, z, fast_dof(μ)) +end +@inline _pwr_from_std_with_rest(::Type{S}, μ, z, n::IntegerLike) where {S} = _from_std_with_rest_bydof(S, μ, z, n) +function _pwr_from_std_with_rest(::Type{S}, μ, z, ::AbstractNoDOF) where {S} + _marginals_from_std_with_rest(S, marginals(μ), z) +end + +# The stream length of a power with a base of fixed stream length: +@inline function _fixed_stream_length(μ::PowerMeasure) + _fixed_stream_length(pwr_base(μ)) * size2length(pwr_size(μ)) +end + +# The nested variate layout of a power over its flat storage, batches of +# tuple variates become struct arrays: +@inline _pwr_variate(μ::PowerMeasure, A::AbstractArray) = _pwr_variate_impl(μ, A) +@inline _pwr_variate(μ::PowerMeasure, A::Union{Tuple,NamedTuple}) = _pwr_variate_impl(μ, A) +@inline _pwr_variate_impl(μ::PowerMeasure, A) = _pwr_nest(pwr_base(μ), _pwr_variate(pwr_base(μ), A)) +@inline _pwr_variate(ν, A::AbstractArray) = _nest_leaf(A, _static_ndims(ν)) +@inline function _pwr_variate(ν::ProductMeasure{<:Tuple}, X::Tuple) + StructArray(map((m, Xi) -> _nest_leaf(Xi, _static_ndims(m)), marginals(ν), X)) +end +@inline function _pwr_variate(ν::ProductMeasure{<:NamedTuple{names}}, X::NamedTuple{names}) where {names} + StructArray(NamedTuple{names}(map((m, Xi) -> _nest_leaf(Xi, _static_ndims(m)), values(marginals(ν)), values(X)))) end +@inline _nest_leaf(A::AbstractArray, ::StaticInteger{0}) = A +@inline _nest_leaf(A::AbstractArray, ::NoMSpaceElementSize) = A +@inline _nest_leaf(A::AbstractArray{<:Any,N}, ::StaticInteger{N}) where {N} = A +@inline _nest_leaf(A::AbstractArray, ::StaticInteger{K}) where {K} = sliced(A, Val(K)) +@inline _pwr_nest(ν::PowerMeasure, B::AbstractArray) = sliced(B, Val(length(pwr_axes(ν)))) +@inline _pwr_nest(ν, B::AbstractArray) = B + +Adapt.adapt_structure(to, μ::PowerMeasure) = PowerMeasure(Adapt.adapt(to, pwr_base(μ)), pwr_axes(μ)) 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..e6a61be6 100644 --- a/src/combinators/product.jl +++ b/src/combinators/product.jl @@ -23,69 +23,95 @@ export marginals function Base.:(==)(a::AbstractProductMeasure, b::AbstractProductMeasure) marginals(a) == marginals(b) end +Base.hash(a::AbstractProductMeasure, h::UInt) = hash(marginals(a), hash(:AbstractProductMeasure, h)) Base.length(μ::AbstractProductMeasure) = length(marginals(μ)) Base.size(μ::AbstractProductMeasure) = size(marginals(μ)) basemeasure(d::AbstractProductMeasure) = productmeasure(map(basemeasure, marginals(d))) -function Base.rand(rng::AbstractRNG, ::Type{T}, d::AbstractProductMeasure) where {T} - mar = marginals(d) - _rand_product(rng, T, mar, eltype(mar)) -end +rand_impl(ctx::GenContext, d::AbstractProductMeasure) = _map(Base.Fix1(_marginal_rand, ctx), marginals(d)) -function _rand_product( - rng::AbstractRNG, - ::Type{T}, - mar, - ::Type{M}, -) where {T,M<:AbstractMeasure} - map(mar) do dⱼ - rand(rng, T, dⱼ) +@inline _marginal_rand(ctx::GenContext, m::AbstractMeasure) = rand_impl(ctx, m) +@inline _marginal_rand(ctx::GenContext, d) = convert_realtype(get_precision(ctx), rand(get_rng(ctx), d)) + +for (head, func) in [(:logdensityof_impl, :logdensityof), (:logdensity_def, :logdensity_def)] + @eval @inline function $head(d::AbstractProductMeasure, x) + _check_marginal_count(marginals(d), x) + mapreduce($func, +, marginals(d), x) end end -function _rand_product( - rng::AbstractRNG, - ::Type{T}, - mar::ReadonlyMappedArray, - ::Type{M}, -) where {T,M<:AbstractMeasure} - mappedarray(mar.data) do dⱼ - rand(rng, T, mar.f(dⱼ)) - end |> collect -end - -function _rand_product(rng::AbstractRNG, ::Type{T}, mar, ::Type{M}) where {T,M} - map(mar) do dⱼ - rand(rng, dⱼ) - end + +# Variates of products are collections of marginal variates, with the same +# structure as the marginals: +@inline function _check_marginal_count(mar::AbstractArray, x::AbstractArray) + size(mar) == size(x) || _throw_marginal_mismatch() + return nothing +end +@inline _check_marginal_count(::AbstractArray, x) = _throw_marginal_mismatch() +# Tuple products also take vector variates (e.g. from converted product +# distributions): +@inline function _check_marginal_count(mar::Tuple, x::Union{Tuple,AbstractVector}) + length(mar) == length(x) || _throw_marginal_mismatch() + return nothing end +@inline _check_marginal_count(::Tuple, x) = _throw_marginal_mismatch() +@inline _check_marginal_count(mar, x) = nothing -function _rand_product( - rng::AbstractRNG, - ::Type{T}, - mar::ReadonlyMappedArray, - ::Type{M}, -) where {T,M} - mappedarray(mar.data) do dⱼ - rand(rng, mar.f(dⱼ)) - end |> collect +@noinline _throw_marginal_mismatch() = + throw(ArgumentError("Variate doesn't match the structure of the marginals of a product measure")) + +struct ProductMeasure{M} <: AbstractProductMeasure + marginals::M end -for func in [:logdensityof, :logdensity_def] - @eval @inline function $func(d::AbstractProductMeasure, x) - mapreduce($func, +, marginals(d), x) +proxy(μ::ProductMeasure{<:FillArrays.Fill}) = + powermeasure(_fill_value(marginals(μ)), _fill_axes(marginals(μ))) + +# Array products with fused kernels draw their variates in one batch (also +# on devices): +function rand_impl(ctx::GenContext, d::ProductMeasure{<:AbstractArray{M}}) where {M} + _array_product_rand(ctx, d, _fused_marginals(M)) +end +_array_product_rand(ctx::GenContext, d::ProductMeasure, ::Val{true}) = batched_rand_impl(ctx, d, ()) +_array_product_rand(ctx::GenContext, d::ProductMeasure, ::Val{false}) = _map(Base.Fix1(_marginal_rand, ctx), marginals(d)) + +# Batches of tuple and named tuple products are tuples resp. named tuples +# of marginal batches: +function batched_rand_impl(ctx::GenContext, μ::ProductMeasure{<:Union{Tuple,NamedTuple}}, sz::SizeLike) + map(m -> batched_rand_impl(ctx, m, sz), marginals(μ)) +end + +# Batches of tuple and named tuple variates are tuples resp. named tuples +# of batches, the marginal densities add up lazily: +for (bhead, head) in [(:batched_logdensityof_impl, :logdensityof_impl), (:batched_logdensity_def, :logdensity_def)] + @eval @inline function $bhead(μ::ProductMeasure{<:Tuple}, X::Tuple) + _lazy_sum(map((m, Xi) -> _batched_kernel($head, m, Xi), marginals(μ), X)) + end + @eval @inline function $bhead(μ::ProductMeasure{<:NamedTuple{names}}, X::NamedTuple{names}) where {names} + _lazy_sum(map((m, Xi) -> _batched_kernel($head, m, Xi), values(marginals(μ)), values(X))) end end +@inline _lazy_sum(ℓs::Tuple) = reduce(_lazy_add, ℓs) -struct ProductMeasure{M} <: AbstractProductMeasure - marginals::M +# 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 -@inline function logdensity_rel(μ::ProductMeasure, ν::ProductMeasure, x) - mapreduce(logdensity_rel, +, marginals(μ), marginals(ν), x) +# 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 +135,15 @@ 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} + _check_marginal_count(marginals(d), x) ℓ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 +187,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}, @@ -171,26 +209,97 @@ end marginals(μ::ProductMeasure) = μ.marginals -# TODO: Better `map` support in MappedArrays -_map(f, args...) = map(f, args...) -_map(f, x::MappedArrays.ReadonlyMappedArray) = mappedarray(fchain((x.f, f)), x.data) +@inline mspace_elsize(μ::ProductMeasure{<:AbstractArray}) = maybestatic_size(marginals(μ)) -function testvalue(::Type{T}, d::AbstractProductMeasure) where {T} - _map(m -> testvalue(T, m), marginals(d)) +@inline function mspace_flatsize(μ::ProductMeasure{<:AbstractArray{M}}) where {M} + _cat_sizes(mspace_flatsize(M), maybestatic_size(marginals(μ))) end -export ⊗ +@inline function mspace_ndims(::Type{<:ProductMeasure{<:AbstractArray{M,N}}}) where {M,N} + _add_ndims(mspace_ndims(M), N) +end +@inline fixed_stream_size(::Type{<:ProductMeasure{<:AbstractArray{M}}}) where {M} = fixed_stream_size(M) + +# Batched kernels over flat storage `(marginal variate dims..., product +# dims..., batch dims...)`. Marginals with scalar variates align with the +# leading dimensions of the batch, so one broadcast evaluates all marginal +# densities. Marginals with array variates are evaluated one by one over +# their slices of the batch. +for (bhead, head) in [(:batched_logdensityof_impl, :logdensityof_impl), (:batched_logdensity_def, :logdensity_def)] + @eval @inline function $bhead(μ::ProductMeasure{<:AbstractArray{M}}, X::AbstractArray) where {M} + _array_product_kernel($head, μ, X, _static_ndims_of(mspace_ndims(M))) + end +end -""" - ⊗(μs::AbstractMeasure...) +@inline function _array_product_kernel(f::F, μ::ProductMeasure, X::AbstractArray, ::StaticInteger{0}) where {F} + mar = marginals(μ) + _check_flatsize(X, maybestatic_size(mar)) + sum_leading_dims(_marginal_broadcast(_DynamicPointLogd(f), mar, X), static(ndims(mar))) +end +@inline function _array_product_kernel(f::F, μ::ProductMeasure, X::AbstractArray, ::StaticInteger{K}) where {F,K} + _marginal_slices_ld(f, marginals(μ), X, Val(K), Val(ndims(X) - K - ndims(marginals(μ)))) +end +@noinline function _array_product_kernel(::F, μ::ProductMeasure{<:AbstractArray{M}}, ::AbstractArray, ::NoMSpaceElementSize) where {F,M} + throw(ArgumentError("Batched density evaluation of products over arrays of marginals of type $(nameof(M)) requires MeasureBase.mspace_ndims to be declared for that type")) +end -`⊗` is a binary operator for building product measures. This satisfies the law +struct _DynamicPointLogd{F} <: Function + f::F +end +@inline (k::_DynamicPointLogd)(m, x) = _dynamic_logd(k.f(m, x), x) + +function _marginal_slices_ld(f::F, mar::AbstractArray{<:Any,N}, A::AbstractArray, ::Val{K}, ::Val{B}) where {F,N,K,B} + if ndims(A) != K + N + B || ntuple(i -> size(A, K + i), Val(N)) != size(mar) + _throw_size_mismatch() + end + lead = ntuple(_ -> Colon(), Val(K)) + trail = ntuple(_ -> Colon(), Val(B)) + ld(i) = _materialize(_batched_kernel(f, mar[i], view(A, lead..., Tuple(i)..., trail...))) + init = _zero_logd(A, ntuple(i -> size(A, K + N + i), Val(B))) + return mapreduce(ld, +, CartesianIndices(mar); init = init) +end + +@inline _zero_logd(A::AbstractArray, ::Tuple{}) = zero(_logd_numtype(A)) +@inline _zero_logd(A::AbstractArray, dims::Tuple) = fill!(similar(A, _logd_numtype(A), dims), 0) + +# Point densities of array products at numeric variates go through the +# batched kernel where the variate rank of the marginals is known, +# marginal by marginal otherwise: +@inline _point_ld(f::F, μ::AbstractProductMeasure, x::AbstractArray{<:Number}) where {F} = f(μ, x) +@inline function logdensityof_impl(μ::ProductMeasure{<:AbstractArray{M}}, x::AbstractArray{<:Number}) where {M} + _array_product_ld(logdensityof_impl, μ, x, mspace_ndims(M)) +end +@inline function logdensity_def(μ::ProductMeasure{<:AbstractArray{M}}, x::AbstractArray{<:Number}) where {M} + _array_product_ld(logdensity_def, μ, x, mspace_ndims(M)) +end +@inline function _array_product_ld(f::F, μ::ProductMeasure, x::AbstractArray, ::Integer) where {F} + _point_result(_materialize(_batched_kernel(f, μ, x)), μ) +end +@inline function _array_product_ld(f::F, μ::ProductMeasure, x::AbstractArray, ::NoMSpaceElementSize) where {F} + _array_product_ld_nested(f, μ, x) +end +@inline function _array_product_ld_nested(f::F, μ::ProductMeasure, x::AbstractArray) where {F} + _check_marginal_count(marginals(μ), x) + mapreduce(_PointLogd(f, nothing), +, marginals(μ), x) +end +@inline (k::_PointLogd{F,Nothing})(m, x) where {F} = _point_ld(k.f, m, x) + +# TODO: Better `map` support in MappedArrays +# `F` keeps the mapped function specialized, it would be passed on +# unspecialized otherwise and its results boxed: +_map(f::F, args...) where {F} = map(f, args...) +# `map` over a named tuple splats its values and boxes the results, +# mapping over the values themselves doesn't: +@inline _map(f::F, nt::NamedTuple{names}) where {F,names} = NamedTuple{names}(map(f, values(nt))) +_map(f::F, x::MappedArrays.ReadonlyMappedArray) where {F} = mappedarray(fchain((x.f, f)), x.data) +# `map` over a struct array builds struct arrays of the results, variates +# of the marginals are wanted as plain arrays: +_map(f::F, x::StructArray) where {F} = [f(m) for m in x] + +function testvalue(::Type{T}, d::AbstractProductMeasure) where {T} + _map(m -> testvalue(T, m), marginals(d)) +end -``` - basemeasure(μ ⊗ ν) == basemeasure(μ) ⊗ basemeasure(ν) -``` -""" -⊗(μs::AbstractMeasure...) = productmeasure(μs) ############################################################################### # I <: Base.Generator @@ -214,39 +323,422 @@ end export rand! using Random: rand!, GLOBAL_RNG -function _rand(rng::AbstractRNG, ::Type{T}, d::ProductMeasure, mar::AbstractArray) where {T} - elT = typeof(rand(rng, T, first(mar))) - - sz = size(mar) - x = Array{elT,length(sz)}(undef, sz) - rand!(rng, d, x) -end - @inline function insupport(d::AbstractProductMeasure, x::AbstractArray) - mar = marginals(d) - # We might get lucky and know statically that everything is inbounds - T = Core.Compiler.return_type(insupport, Tuple{eltype(mar),eltype(x)}) - T <: True || all(zip(x, mar)) do (xj, mj) - insupport(mj, xj) == true - end + _all_insupport(broadcast(_insupport_bool ∘ insupport, marginals(d), x)) end @inline function insupport(d::AbstractProductMeasure, x) - for (mj, xj) in zip(marginals(d), x) - dynamic(insupport(mj, xj)) || return false - end - return true + mapreduce(insupport, _insupport_and, marginals(d), x) end -getdof(d::AbstractProductMeasure) = mapreduce(getdof, +, marginals(d)) +@inline _all_insupport(A::AbstractArray{<:NoFastInsupport{T}}) where {T} = NoFastInsupport{T}() +@inline _all_insupport(A::AbstractArray) = all(A) + +getdof(d::AbstractProductMeasure) = _sum_dofs(getdof, marginals(d)) +fast_dof(d::AbstractProductMeasure) = _sum_dofs(fast_dof, marginals(d)) + +# Sums over static DOFs of tuples fold at compile time, arrays of marginals +# are summed dynamically (also on GPU arrays): +@inline _sum_dofs(f, mar) = sum(f, mar) +@inline _sum_dofs(f, mar::AbstractArray{M}) where {M} = _sum_dofs(f, mar, _unit_dof(M)) +@inline _sum_dofs(f, mar::AbstractArray, ::True) = length(mar) +@inline _sum_dofs(f, mar::AbstractArray, ::False) = mapreduce(_dynamic_dof ∘ f, +, mar; init = 0) + +# Marginals with scalar variates and a standard transport have one degree +# of freedom each, so their total needs no reduction over the marginals +# (which may live on a device): +@inline function _unit_dof(::Type{M}) where {M} + static(mspace_ndims(M) === 0 && preferred_stdmeasure(M) isa Type{<:StdMeasure}) +end +@inline _sum_dofs(f, mar::StaticArray) = mapreduce(f, +, mar; init = static(0)) +@inline _dynamic_dof(n::IntegerLike) = dynamic(n) +@inline _dynamic_dof(nodof::AbstractNoDOF) = nodof function checked_arg(μ::ProductMeasure{<:NTuple{N,Any}}, x::NTuple{N,Any}) where {N} map(checked_arg, marginals(μ), x) end +# Variates of array products are arrays of marginal variates or, for +# marginals with array variates of declared rank, their flat storage: +@propagate_inbounds function checked_arg(μ::ProductMeasure{<:AbstractArray{M}}, x::AbstractArray) where {M} + @boundscheck _check_product_arg(marginals(μ), x, _static_ndims_of(mspace_ndims(M))) + return x +end + +@inline function _check_product_arg(mar, x::AbstractArray, ::Any) + _check_marginal_count(mar, x) + foreach(checked_arg, mar, x) + return nothing +end +@inline function _check_product_arg(mar::AbstractArray{<:Any,N}, x::AbstractArray{<:Number}, ::StaticInteger{0}) where {N} + _check_marginal_count(mar, x) +end +@inline function _check_product_arg(mar::AbstractArray{<:Any,N}, x::AbstractArray{<:Number}, ::StaticInteger{K}) where {N,K} + if ndims(x) != K + N || ntuple(i -> size(x, K + i), Val(N)) != size(mar) + _throw_marginal_mismatch() + end + return nothing +end +@inline _check_product_arg(mar, x::AbstractArray{<:Number}, ::NoMSpaceElementSize) = _check_marginal_count(mar, x) + + function checked_arg( μ::ProductMeasure{<:NamedTuple{names}}, x::NamedTuple{names}, ) where {names} NamedTuple{names}(map(checked_arg, values(marginals(μ)), values(x))) end + + +# Transport marginal by marginal, the standard variates of the marginals +# are concatenated in order. Batches of tuple and named tuple variates are +# tuples resp. named tuples of batches. + +function transport_to_std(::Type{S}, μ::ProductMeasure{<:Tuple}, x::Tuple) where {S<:StdMeasure} + _flatten_to_rv(map((m, xi) -> _as_stdstream(transport_to_std(S, m, xi)), marginals(μ), x)) +end + +function transport_to_std(::Type{S}, μ::ProductMeasure{<:NamedTuple{names}}, x::NamedTuple{names}) where {S<:StdMeasure,names} + transport_to_std(S, productmeasure(values(marginals(μ))), values(x)) +end + +function batched_transport_to_std(::Type{S}, μ::ProductMeasure{<:Tuple}, X::Tuple) where {S<:StdMeasure} + _vcat_std(map((m, Xi) -> batched_transport_to_std(S, m, Xi), marginals(μ), X)) +end + +function batched_transport_to_std(::Type{S}, μ::ProductMeasure{<:NamedTuple{names}}, X::NamedTuple{names}) where {S<:StdMeasure,names} + batched_transport_to_std(S, productmeasure(values(marginals(μ))), values(X)) +end + +@inline _vcat_std(Zs::Tuple) = vcat(Zs...) +@inline _vcat_std(::Tuple{}) = SVector{0,Bool}() + +function transport_from_std(::Type{S}, μ::ProductMeasure{<:Union{Tuple,NamedTuple}}, z::AbstractVector) where {S<:StdMeasure} + x, z_rest = transport_from_std_with_rest(S, μ, z) + isempty(z_rest) || _throw_std_length_mismatch() + return x +end + +# Streams of tuple product variates are consumed marginal by marginal: +function transport_to_std_with_rest(::Type{S}, μ::ProductMeasure{<:Union{Tuple,NamedTuple}}, x::AbstractVector) where {S<:StdMeasure} + z, x_rest = _marginals_to_std_with_rest(S, values(marginals(μ)), x) + x_μ, _ = split_at(x, maybestatic_length(x) - maybestatic_length(x_rest)) + return z, x_μ, x_rest +end + +function batched_transport_to_std_with_rest(::Type{S}, μ::ProductMeasure{<:Union{Tuple,NamedTuple}}, X::AbstractArray, sz::SizeLike) where {S<:StdMeasure} + _tuple_product_to_std_with_rest(S, μ, X, sz) +end +function _tuple_product_to_std_with_rest(::Type{S}, μ, X::AbstractArray, ::Tuple{}) where {S} + _marginals_to_std_with_rest(S, values(marginals(μ)), X) +end +function _tuple_product_to_std_with_rest(::Type{S}, μ, X::AbstractArray, sz::SizeLike) where {S} + X_v, X_rest = _split_stream_variates(μ, X, sz) + Z, _ = _marginals_to_std_with_rest(S, values(marginals(μ)), X_v) + return _merge_multiplicity(Z, sz), X_rest +end + +function _marginals_to_std_with_rest(::Type{S}, νs::Tuple{Vararg{Any}}, X::AbstractArray) where {S} + Z1, X_rest = batched_transport_to_std_with_rest(S, νs[1], X, ()) + Z2_end, X_final_rest = _marginals_to_std_with_rest(S, Base.tail(νs), X_rest) + return vcat(Z1, Z2_end), X_final_rest +end +function _marginals_to_std_with_rest(::Type{S}, νs::Tuple{Any}, X::AbstractArray) where {S} + batched_transport_to_std_with_rest(S, νs[1], X, ()) +end + +function transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:Tuple}, z::AbstractVector) where {S<:StdMeasure} + _marginals_from_std_with_rest(S, marginals(μ), z) +end + +function transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:NamedTuple{names}}, z::AbstractVector) where {S<:StdMeasure,names} + ys, z_rest = _marginals_from_std_with_rest(S, values(marginals(μ)), z) + return NamedTuple{names}(ys), z_rest +end + +function batched_transport_from_std(::Type{S}, μ::ProductMeasure{<:Union{Tuple,NamedTuple}}, Z::AbstractArray) where {S<:StdMeasure} + X, Z_rest = batched_transport_from_std_with_rest(S, μ, Z, ()) + size(Z_rest, 1) == 0 || _throw_std_length_mismatch() + return X +end + +function batched_transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:Tuple}, Z::AbstractArray, sz::SizeLike) where {S<:StdMeasure} + _tuple_product_from_std_with_rest(S, μ, Z, sz) +end + +function batched_transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:NamedTuple{names}}, Z::AbstractArray, sz::SizeLike) where {S<:StdMeasure,names} + Xs, Z_rest = _tuple_product_from_std_with_rest(S, productmeasure(values(marginals(μ))), Z, sz) + return NamedTuple{names}(Xs), Z_rest +end + +# One variate per stream is consumed marginal by marginal, several per +# stream via the degrees of freedom of the whole product: +function _tuple_product_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, ::Tuple{}) where {S} + _marginals_batched_from_std_with_rest(S, marginals(μ), Z) +end +function _tuple_product_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, sz::SizeLike) where {S} + _batched_from_std_bydof(S, μ, Z, sz, fast_dof(μ)) +end + +function _marginals_batched_from_std_with_rest(::Type{S}, νs::Tuple{Vararg{Any}}, Z::AbstractArray) where {S} + X1, Z_rest = batched_transport_from_std_with_rest(S, νs[1], Z, ()) + X2_end, Z_final_rest = _marginals_batched_from_std_with_rest(S, Base.tail(νs), Z_rest) + return (X1, X2_end...), Z_final_rest +end + +_marginals_batched_from_std_with_rest(::Type{S}, ::Tuple{}, Z::AbstractArray) where {S} = (), Z + +function _marginals_from_std_with_rest(::Type{S}, νs::Tuple{Vararg{Any}}, z::AbstractVector) where {S} + y1, z_rest = transport_from_std_with_rest(S, νs[1], z) + y2_end, z_final_rest = _marginals_from_std_with_rest(S, Base.tail(νs), z_rest) + return (y1, y2_end...), z_final_rest +end + +_marginals_from_std_with_rest(::Type{S}, ::Tuple{}, z::AbstractVector) where {S} = (), z + +function _marginals_from_std_with_rest(::Type{S}, νs::AbstractArray{M}, z::AbstractVector) where {S,M} + idxs = eachindex(νs) + if isconcretetype(M) + # The variate type is uniform, so the loop is type stable (the type + # of the remaining stream stays invariant under repeated view-taking): + y1, z_rest = transport_from_std_with_rest(S, νs[first(idxs)], z) + ys = similar(Array{typeof(y1)}, axes(νs)) + ys[first(idxs)] = y1 + for i in Iterators.drop(idxs, 1) + ys[i], z_rest = transport_from_std_with_rest(S, νs[i], z_rest) + end + return ys, z_rest + else + ys_any = Vector{Any}(undef, length(idxs)) + z_rest = z + for (j, i) in enumerate(idxs) + ys_any[j], z_rest = transport_from_std_with_rest(S, νs[i], z_rest) + end + return [y for y in ys_any], z_rest + end +end + + +# Array products: marginals of scalar variates with one degree of freedom +# each transport in a single broadcast (the marginals align with the leading +# dimensions of the batch), other marginals one by one over their slices of +# the batch. + +@inline _fused_marginals(::Type{M}) where {M} = Val(isconcretetype(M) && _unit_dof(M) === static(true)) + +function batched_transport_to_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, X::AbstractArray) where {S<:StdMeasure,M} + _array_product_batched_to_std(S, μ, X, _fused_marginals(M), _static_ndims_of(mspace_ndims(M))) +end +function _array_product_batched_to_std(::Type{S}, μ, X::AbstractArray, ::Val{true}, ::Any) where {S} + mar = marginals(μ) + _check_flatsize(X, maybestatic_size(mar)) + _as_stream_batch(_materialize(_marginal_broadcast(_ToStd{S}(), mar, X)), static(ndims(mar))) +end +function _array_product_batched_to_std(::Type{S}, μ, X::AbstractArray, ::Val{false}, ::StaticInteger{K}) where {S,K} + mar = marginals(μ) + n_batch = ndims(X) - K - ndims(mar) + n_batch >= 0 || _throw_size_mismatch() + _marginals_to_std_loop(S, mar, X, Val(K), Val(n_batch)) +end +@noinline function _array_product_batched_to_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, ::AbstractArray, ::Val{false}, ::NoMSpaceElementSize) where {S,M} + throw(ArgumentError("Batched transport of products over arrays of marginals of type $(nameof(M)) requires MeasureBase.mspace_ndims to be declared for that type")) +end + +function _marginals_to_std_loop(::Type{S}, mar::AbstractArray{<:Any,N}, X::AbstractArray, ::Val{K}, ::Val{B}) where {S,N,K,B} + ntuple(i -> size(X, K + i), Val(N)) == size(mar) || _throw_size_mismatch() + lead = ntuple(_ -> Colon(), Val(K)) + trail = ntuple(_ -> Colon(), Val(B)) + zs = map(i -> batched_transport_to_std(S, mar[i], view(X, lead..., Tuple(i)..., trail...)), vec(CartesianIndices(mar))) + isempty(zs) ? similar(X, (0, ntuple(i -> size(X, K + N + i), Val(B))...)) : reduce(vcat, zs) +end + +function batched_transport_from_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, Z::AbstractArray) where {S<:StdMeasure,M} + _array_product_batched_from_std(S, μ, Z, _fused_marginals(M), _static_ndims_of(mspace_ndims(M))) +end +function _array_product_batched_from_std(::Type{S}, μ, Z::AbstractArray, ::Val{true}, ::Any) where {S} + mar = marginals(μ) + size(Z, 1) == length(mar) || _throw_std_length_mismatch() + _materialize(_marginal_broadcast(_FromStd{S}(), mar, maybestatic_reshape(Z, (_batch_dims(mar)..., Base.tail(_batch_dims(Z))...)))) +end +function _array_product_batched_from_std(::Type{S}, μ, Z::AbstractArray, ::Val{false}, ::StaticInteger{K}) where {S,K} + X, Z_rest = _marginals_from_std_loop(S, marginals(μ), Z, Val(K)) + size(Z_rest, 1) == 0 || _throw_std_length_mismatch() + return X +end +@noinline function _array_product_batched_from_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, ::AbstractArray, ::Val{false}, ::NoMSpaceElementSize) where {S,M} + throw(ArgumentError("Batched transport to products over arrays of marginals of type $(nameof(M)) requires MeasureBase.mspace_ndims to be declared for that type")) +end + +function batched_transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, Z::AbstractArray, sz::SizeLike) where {S<:StdMeasure,M} + _array_product_batched_from_std_with_rest(S, μ, Z, sz, _fused_marginals(M), _static_ndims_of(mspace_ndims(M))) +end +function _array_product_batched_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, sz::SizeLike, ::Val{true}, ::Any) where {S} + _batched_from_std_bydof(S, μ, Z, sz, length(marginals(μ))) +end +function _array_product_batched_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, ::Tuple{}, ::Val{false}, ::StaticInteger{K}) where {S,K} + _marginals_from_std_loop(S, marginals(μ), Z, Val(K)) +end +function _array_product_batched_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, sz::SizeLike, ::Val{false}, ::StaticInteger{K}) where {S,K} + _batched_from_std_bydof(S, μ, Z, sz, fast_dof(μ)) +end +@noinline function _array_product_batched_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, ::AbstractArray, ::SizeLike, ::Val{false}, ::NoMSpaceElementSize) where {S,M} + throw(ArgumentError("Batched transport to products over arrays of marginals of type $(nameof(M)) requires MeasureBase.mspace_ndims to be declared for that type")) +end + +# The marginals consume the streams one after the other, their variates +# fill the batch `(marginal variate dims..., product dims..., batch dims...)`: +function _marginals_from_std_loop(::Type{S}, mar::AbstractArray{<:Any,N}, Z::AbstractArray, ::Val{K}) where {S,N,K} + idxs = vec(CartesianIndices(mar)) + batch_dims = Base.tail(size(Z)) + lead = ntuple(_ -> Colon(), Val(K)) + trail = ntuple(_ -> Colon(), Val(length(batch_dims))) + if isempty(idxs) + return similar(Z, (ntuple(_ -> 0, Val(K))..., size(mar)..., batch_dims...)), Z + end + X1, Z_rest = batched_transport_from_std_with_rest(S, mar[idxs[1]], Z, ()) + X = similar(Z, eltype(X1), (ntuple(i -> size(X1, i), Val(K))..., size(mar)..., batch_dims...)) + X[lead..., Tuple(idxs[1])..., trail...] = X1 + for i in idxs[2:end] + Xi, Z_rest = batched_transport_from_std_with_rest(S, mar[i], Z_rest, ()) + X[lead..., Tuple(i)..., trail...] = Xi + end + return X, Z_rest +end + +# Point transport of array products: arrays of marginal variates with flat +# storage and flat variates go through the batched kernels, marginals +# without a declared variate rank transport one by one. +function transport_to_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, x::AbstractArray) where {S<:StdMeasure,M} + _array_product_to_std(S, μ, x, _flat_storage(x), _static_ndims_of(mspace_ndims(M))) +end +@inline function _array_product_to_std(::Type{S}, μ, x::AbstractArray, x_flat::AbstractArray, ::StaticInteger) where {S} + _single_std(batched_transport_to_std(S, μ, x_flat)) +end +function _array_product_to_std(::Type{S}, μ, x::AbstractArray, ::Any, ::Any) where {S} + _check_marginal_count(marginals(μ), x) + zs = [_as_stdstream(transport_to_std(S, m, xi)) for (m, xi) in zip(marginals(μ), x)] + isempty(zs) ? SVector{0,Bool}() : reduce(vcat, zs) +end + +# Marginals with variates of fixed size and declared rank yield a nested +# view of the flat variate batch, others transport marginal by marginal: +function transport_from_std(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, z::AbstractVector) where {S<:StdMeasure,M} + _array_product_from_std(S, μ, z, fixed_stream_size(M), _static_ndims_of(mspace_ndims(M))) +end +@inline function _array_product_from_std(::Type{S}, μ, z::AbstractVector, ::True, k::StaticInteger) where {S} + _nest_leaf(batched_transport_from_std(S, μ, z), k) +end +function _array_product_from_std(::Type{S}, μ, z::AbstractVector, ::Any, ::Any) where {S} + ys, z_rest = _marginals_from_std_with_rest(S, marginals(μ), z) + isempty(z_rest) || _throw_std_length_mismatch() + return ys +end + +function transport_from_std_with_rest(::Type{S}, μ::ProductMeasure{<:AbstractArray{M}}, z::AbstractVector) where {S<:StdMeasure,M} + _array_product_from_std_with_rest(S, μ, z, fixed_stream_size(M), _static_ndims_of(mspace_ndims(M))) +end +function _array_product_from_std_with_rest(::Type{S}, μ, z::AbstractVector, ::True, k::StaticInteger) where {S} + X, z_rest = batched_transport_from_std_with_rest(S, μ, z, ()) + return _nest_leaf(X, k), z_rest +end +function _array_product_from_std_with_rest(::Type{S}, μ, z::AbstractVector, ::Any, ::Any) where {S} + _marginals_from_std_with_rest(S, marginals(μ), z) +end + + +# Streams: tuple products consume marginal by marginal, so marginals of +# value-dependent size are supported for a single variate per stream. +# Several variates per stream are split by the fixed stream length of the +# product. +function batched_logdensityof_with_rest(μ::ProductMeasure{<:Tuple}, X::AbstractArray, ::Tuple{}) + _marginals_ld_with_rest(marginals(μ), X) +end +function batched_logdensityof_with_rest(μ::ProductMeasure{<:Tuple}, x::AbstractVector, ::Tuple{}) + _marginals_ld_with_rest(marginals(μ), x) +end +function batched_logdensityof_with_rest(μ::ProductMeasure{<:Tuple}, X::AbstractArray, sz::SizeLike) + X_v, X_rest = _split_stream_variates(μ, X, sz) + ℓ, _ = _marginals_ld_with_rest(marginals(μ), X_v) + return ℓ, X_rest +end +function batched_logdensityof_with_rest(μ::ProductMeasure{<:NamedTuple{names}}, X::AbstractArray, sz::SizeLike) where {names} + batched_logdensityof_with_rest(productmeasure(values(marginals(μ))), X, sz) +end + +# The rows of `prod(sz)` variates of fixed stream length, as a batch of +# streams `(stream length, sz..., batch dims...)`: +function _split_stream_variates(μ, X::AbstractArray, sz::SizeLike) + n_rows = _fixed_stream_length(μ) + X_μ, X_rest = _batched_split(X, _chunk_rows(n_rows, sz)) + return maybestatic_reshape(X_μ, (n_rows, size_dims(sz)..., Base.tail(_batch_dims(X_μ))...)), X_rest +end + +@inline _fixed_stream_length(μ::ProductMeasure{<:Tuple}) = sum(_fixed_stream_length, marginals(μ)) +@inline _fixed_stream_length(μ::ProductMeasure{<:NamedTuple}) = sum(_fixed_stream_length, values(marginals(μ))) +function _marginals_ld_with_rest(ms::Tuple, X::AbstractArray) + ℓ1, X2 = batched_logdensityof_with_rest(ms[1], X, ()) + ℓ_rest, X_rest = _marginals_ld_with_rest(Base.tail(ms), X2) + return _lazy_add(ℓ1, ℓ_rest), X_rest +end +function _marginals_ld_with_rest(ms::Tuple{Any}, X::AbstractArray) + batched_logdensityof_with_rest(ms[1], X, ()) +end +function batched_logdensityof_with_rest(μ::ProductMeasure{<:NamedTuple{names}}, X::AbstractArray, sz::Tuple{}) where {names} + batched_logdensityof_with_rest(productmeasure(values(marginals(μ))), X, sz) +end +function batched_logdensityof_with_rest(μ::ProductMeasure{<:NamedTuple{names}}, x::AbstractVector, sz::Tuple{}) where {names} + batched_logdensityof_with_rest(productmeasure(values(marginals(μ))), x, sz) +end + +# Folded pairwise over the marginal types, so that the result is a constant: +@inline fixed_stream_size(::Type{<:ProductMeasure{M}}) where {M<:Tuple} = static_all(fixed_stream_size, M) +@inline function fixed_stream_size(::Type{<:ProductMeasure{NamedTuple{names,M}}}) where {names,M<:Tuple} + fixed_stream_size(ProductMeasure{M}) +end + + +# Broadcasts over struct arrays of marginals run over their leaf columns, +# the marginals are rebuilt from the column values inside the kernel: + +@inline _leaf_columns(sa::StructArray) = _leaf_columns_of(values(StructArrays.components(sa))) +@inline _leaf_columns_of(cs::Tuple) = (_leaf_columns_of(first(cs))..., _leaf_columns_of(Base.tail(cs))...) +@inline _leaf_columns_of(::Tuple{}) = () +@inline _leaf_columns_of(c::StructArray) = _leaf_columns(c) +@inline _leaf_columns_of(c::AbstractArray{T}) where {T} = Base.issingletontype(T) ? () : (c,) + +@generated function _rebuild_element(::Type{SA}, vals::Tuple) where {SA<:StructArray} + expr, _ = _rebuild_expr(SA, 1) + return expr +end +function _rebuild_expr(::Type{SA}, i::Int) where {T,N,C,SA<:StructArray{T,N,C}} + args = Any[] + for CT in C.parameters[2].parameters + if CT <: StructArray + e, i = _rebuild_expr(CT, i) + push!(args, e) + elseif Base.issingletontype(eltype(CT)) + push!(args, :($(eltype(CT).instance))) + else + push!(args, :(vals[$i])) + i += 1 + end + end + return :(constructorof($T)($(args...))), i +end + +struct _WithElement{SA,G} <: Function + g::G +end +_WithElement{SA}(g::G) where {SA,G} = _WithElement{SA,G}(g) +@inline function (k::_WithElement{SA})(args::Vararg{Any,N}) where {SA,N} + k.g(_rebuild_element(SA, Base.front(args)), args[end]) +end + +@inline function _marginal_broadcast(g::G, mar::AbstractArray, X) where {G} + Broadcast.instantiate(Broadcast.broadcasted(g, mar, X)) +end +@inline function _marginal_broadcast(g::G, mar::StructArray, X) where {G} + Broadcast.instantiate(Broadcast.broadcasted(_WithElement{typeof(mar)}(g), _leaf_columns(mar)..., X)) +end + +Adapt.adapt_structure(to, μ::ProductMeasure) = ProductMeasure(Adapt.adapt(to, marginals(μ))) diff --git a/src/combinators/reshape.jl b/src/combinators/reshape.jl new file mode 100644 index 00000000..e20a0c5f --- /dev/null +++ b/src/combinators/reshape.jl @@ -0,0 +1,70 @@ +# 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) + +@inline mspace_elsize(μ::PushforwardMeasure{<:Reshape}) = μ.f.output_size +# Reshaped variates have flat storage only if the reshaped elements are +# numbers, the reshape of a nested variate has no flat form its density +# kernel could consume: +@inline function mspace_flatsize(μ::PushforwardMeasure{<:Reshape}) + _reshaped_flatsize(mspace_flatsize(μ.origin), mspace_elsize(μ.origin), μ.f.output_size) +end +@inline function _reshaped_flatsize(sz_flat::SizeLike, sz_outer::SizeLike, sz_out) + _reshaped_flatsize(Val(length(size_dims(sz_flat)) == length(size_dims(sz_outer))), sz_out) +end +@inline _reshaped_flatsize(::Val{true}, sz_out) = sz_out +@inline _reshaped_flatsize(::Val{false}, sz_out) = NoMSpaceElementSize{typeof(sz_out)}() +@inline _reshaped_flatsize(::Any, ::Any, sz_out) = NoMSpaceElementSize{typeof(sz_out)}() diff --git a/src/combinators/restricted.jl b/src/combinators/restricted.jl index e3b66212..b4339a59 100644 --- a/src/combinators/restricted.jl +++ b/src/combinators/restricted.jl @@ -3,11 +3,17 @@ struct RestrictedMeasure{P,M} <: AbstractMeasure base::M end +@inline mspace_elsize(μ::RestrictedMeasure) = mspace_elsize(μ.base) +@inline mspace_flatsize(μ::RestrictedMeasure) = mspace_flatsize(μ.base) +@inline mspace_flatsize(::Type{<:RestrictedMeasure{<:Any,M}}) where {M} = mspace_flatsize(M) +@inline mspace_ndims(::Type{<:RestrictedMeasure{<:Any,M}}) where {M} = mspace_ndims(M) +@inline fixed_stream_size(::Type{<:RestrictedMeasure{<:Any,M}}) where {M} = fixed_stream_size(M) + @inline logdensity_def(d::RestrictedMeasure, x) = logdensity_def(d.base, x) basemeasure(μ::RestrictedMeasure) = μ.base -insupport(μ::RestrictedMeasure, x) = μ.predicate(x) && insupport(μ.base, x) +insupport(μ::RestrictedMeasure, x) = _insupport_and(μ.predicate(x), insupport(μ.base, x)) function Pretty.quoteof(d::RestrictedMeasure) qf = Pretty.quoteof(d.predicate) diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl index 26ba3948..8a402036 100644 --- a/src/combinators/smart-constructors.jl +++ b/src/combinators/smart-constructors.jl @@ -1,187 +1,287 @@ +# 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{<:Number}, 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) +""" + productmeasure(μs) + +Constructs a product over a collection `μs` of measures. + +Examples: + +```julia +productmeasure((StdNormal(), StdExponential())) +productmeasure((a = StdNormal(), b = StdExponential())) +productmeasure([pushfwd(AffineMaps.Mul(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 -function productmeasure(mar::ReadonlyMappedArray{T,N,A,Returns{M}}) where {T,N,A,M} - return powermeasure(mar.f.value, axes(mar.data)) +function _generic_productmeasure_impl(mar::AbstractArray{T}) where {T} + if Base.issingletontype(T) + powermeasure(instance(T), axes(mar)) + elseif T <: AbstractMeasure + ProductMeasure(_marginal_storage(mar)) + else + ProductMeasure(_marginal_storage(map(asmeasure, mar))) + end end -productmeasure(mar::Base.Generator) = ProductMeasure(mar) -productmeasure(mar::AbstractArray) = ProductMeasure(mar) +# Arrays of parameterized isbits marginals are stored as struct arrays, so +# that batched kernels broadcast over the numeric parameter columns (also +# on GPUs and in traced code): +function _marginal_storage(mar::AbstractArray{T}) where {T} + if isconcretetype(T) && isbitstype(T) && !Base.issingletontype(T) + StructArray(mar; unwrap = _unwrap_field) + else + mar + end +end +_marginal_storage(mar::StructArray) = mar + +# Function objects stay opaque columns, their wrappers (e.g. `Base.Fix1`) +# can't be rebuilt from their fields via ConstructionBase: +@inline function _unwrap_field(::Type{T}) where {T} + isstructtype(T) && !Base.issingletontype(T) && fieldcount(T) > 0 && + !(T <: Number) && !(T <: AbstractArray) && !(T <: Tuple) && !(T <: AbstractString) && + !(T <: Symbol) && !(T <: Function) +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/spikemixture.jl b/src/combinators/spikemixture.jl index e39d4230..02e8e57a 100644 --- a/src/combinators/spikemixture.jl +++ b/src/combinators/spikemixture.jl @@ -23,23 +23,9 @@ end for func in [:logdensityof, :logdensity_def] @eval @inline function $func(μ::SpikeMixture, x) - # NOTE: We could instead write this as - # R1 = typeof(log(one(μ.s))) - # R2 = typeof(log(one(μ.w))) - - # which would rely on constant propagation insteadof type inference. - # We'll try this for now and come back to the question if we see - # problems. - - R1 = Core.Compiler.return_type(log, Tuple{typeof(μ.s)}) - R2 = Core.Compiler.return_type(log, Tuple{typeof(μ.w)}) - R3 = Core.Compiler.return_type($func, Tuple{typeof(μ.m),typeof(x)}) - R = promote_type(R1, R2, R3) - if iszero(x) - return convert(R, log(μ.s))::R - else - return convert(R, log(μ.w) + $func(μ.m, x))::R - end + ℓ_spike = dynamic(log(μ.s)) + ℓ_parent = dynamic(log(μ.w)) + dynamic($func(μ.m, x)) + ifelse(iszero(x), oftype(ℓ_parent, ℓ_spike), ℓ_parent) end end @@ -47,10 +33,24 @@ function gentype(μ::SpikeMixture) gentype(μ.m) end -function Base.rand(rng::AbstractRNG, T::Type, μ::SpikeMixture) - return (rand(rng, T) < μ.w) * rand(rng, T, μ.m) +function rand_impl(ctx::GenContext, μ::SpikeMixture) + return (rand(get_rng(ctx), get_precision(ctx)) < μ.w) * rand_impl(ctx, μ.m) end +function batched_rand_impl(ctx::GenContext, μ::SpikeMixture, sz::SizeLike) + _spike_batched_rand(ctx, μ, sz, _static_ndims(μ.m)) +end +function _spike_batched_rand(ctx::GenContext, μ::SpikeMixture, sz::SizeLike, k::StaticInteger) + X = batched_rand_impl(ctx, μ.m, sz) + return ifelse.(_batch_mask(_rand_bulk(ctx, sz) .< μ.w, k), X, zero(eltype(X))) +end +_spike_batched_rand(ctx::GenContext, μ::SpikeMixture, sz::SizeLike, ::NoMSpaceElementSize) = + _batched_rand_pointwise(ctx, μ, sz) + testvalue(::Type{T}, μ::SpikeMixture) where {T} = zero(T) -insupport(μ::SpikeMixture, x) = dynamic(insupport(μ.m, x)) || iszero(x) +insupport(μ::SpikeMixture, x) = _insupport_mask(insupport(μ.m, x)) | iszero(x) + + +@inline mspace_flatsize(μ::SpikeMixture) = _scalar_or_unknown(mspace_flatsize(μ.m)) +@inline mspace_flatsize(::Type{<:SpikeMixture{M}}) where {M} = _scalar_or_unknown(mspace_flatsize(M)) diff --git a/src/combinators/superpose.jl b/src/combinators/superpose.jl index 099ee806..5952a760 100644 --- a/src/combinators/superpose.jl +++ b/src/combinators/superpose.jl @@ -65,85 +65,126 @@ 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 - -function density_def(s::SuperpositionMeasure, x) - T = typeof(s) - msg = """ - Not implemented: There is no method - density_def(::$T, x) - """ - error(msg) -end - -@inline function logdensity_def( - μ::T, - ν::T, - x, -) where {T<:(SuperpositionMeasure{Tuple{A,B}} where {A,B})} - if μ === ν - return zero(return_type(logdensity_def, (μ, x))) - else - return logdensity_def(μ, x) - logdensity_def(ν, x) - end +# Masks components outside of their support with -Inf: +@inline _masked_logd(ℓ, ins) = ifelse(_insupport_mask(ins), ℓ, oftype(ℓ, -Inf)) + +# Branch-free logsumexp over the components, valid for infinite entries: +@inline function _logsumexp_components(ℓs) + m = reduce(max, ℓs) + m_finite = ifelse(isfinite(m), m, zero(m)) + m_finite + log(sum(map(ℓ -> exp(ℓ - m_finite), ℓs))) 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)) +# The density of a superposition relative to the superposition of the +# component base measures, in log space: each component contributes its +# own density, divided by the density of the superposed base measures +# relative to its own base measure. +function logdensity_def(s::SuperpositionMeasure, x) + cs = values(s.components) + αs = map(basemeasure, cs) + terms = map(cs, αs) do cᵢ, αᵢ + ℓᵢ = _dynamic_logd(logdensity_def(cᵢ, x), x) + log_dΣα_dαᵢ = _logsumexp_components(map(αs) do αⱼ + _masked_logd(logdensity_rel(αⱼ, αᵢ, x), insupport(αⱼ, x)) + end) + _masked_logd(ℓᵢ - log_dΣα_dαᵢ, insupport(cᵢ, x)) + end + _logsumexp_components(terms) 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)) +@inline function logdensity_rel_def(μ::T, ν::T, x) where {T<:SuperpositionMeasure} + ℓ = logdensity_def(μ, x) - logdensity_def(ν, x) + ifelse(μ === ν, zero(ℓ), ℓ) end -@inline function logdensity_def(s, β::(SuperpositionMeasure{Tuple{A,B}} where {A,B}), x) - -logdensity_def(β, s, x) +function _superpos_logdensity_rel(s::SuperpositionMeasure, β, x) + cs = values(s.components) + ds = map(cs) do μ + _masked_logd(logdensity_rel(μ, β, x), insupport(μ, x)) + end + _logsumexp_components(ds) end -@inline logdensity_def(s::SuperpositionMeasure, x) = log(density_def(s, x)) +@inline logdensity_rel_def(s::SuperpositionMeasure, β, x) = _superpos_logdensity_rel(s, β, x) + +@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 density_def(s::SuperpositionMeasure, x) = exp(logdensity_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 _component_masses(μ::SuperpositionMeasure) + masses = map(massof, values(μ.components)) + total = sum(masses) + total isa AbstractUnknownMass && throw( + ArgumentError("Cannot sample from a superposition of measures of unknown mass"), + ) + return map(dynamic, masses), dynamic(total) +end -@inline function insupport(d::SuperpositionMeasure, x) - any(d.components) do c - dynamic(insupport(c, x)) +function rand_impl(ctx::GenContext, μ::SuperpositionMeasure) + components = values(μ.components) + masses, total = _component_masses(μ) + threshold = rand(get_rng(ctx), get_precision(ctx)) * total + csum = zero(threshold) + for (mass, c) in zip(masses, components) + csum += mass + csum >= threshold && return rand_impl(ctx, c) end + return rand_impl(ctx, last(components)) +end + +# Batches of superpositions draw a batch from each component and select +# by mass, branch-free: +function batched_rand_impl(ctx::GenContext, μ::SuperpositionMeasure, sz::SizeLike) + _superpose_batched_rand(ctx, μ, sz, _static_ndims(μ)) +end +function _superpose_batched_rand(ctx::GenContext, μ::SuperpositionMeasure, sz::SizeLike, k::StaticInteger) + components = values(μ.components) + masses, total = _component_masses(μ) + thresholds = _batch_mask(_rand_bulk(ctx, sz) .* total, k) + X = batched_rand_impl(ctx, first(components), sz) + csum = first(masses) + for (mass, c) in Iterators.drop(zip(masses, components), 1) + X = ifelse.(thresholds .<= csum, X, batched_rand_impl(ctx, c, sz)) + csum += mass + end + return X +end +_superpose_batched_rand(ctx::GenContext, μ::SuperpositionMeasure, sz::SizeLike, ::NoMSpaceElementSize) = + _batched_rand_pointwise(ctx, μ, sz) + +@inline function insupport(d::SuperpositionMeasure, x) + mapreduce(c -> _insupport_mask(insupport(c, x)), |, values(d.components)) +end + + +@inline mspace_flatsize(μ::SuperpositionMeasure) = mspace_flatsize(typeof(μ)) + +# The variate rank of a superposition is the common rank of its components, +# folded pairwise over the component types so that it stays a constant: +@inline mspace_ndims(::Type{<:SuperpositionMeasure{C}}) where {C<:AbstractArray} = mspace_ndims(eltype(C)) +@inline function mspace_ndims(::Type{MU}) where {C<:Tuple,MU<:SuperpositionMeasure{C}} + static_mapreduce(mspace_ndims, _CommonNDims{MU}(), C) +end +struct _CommonNDims{MU} <: Function end +@inline (::_CommonNDims{MU})(a::Integer, b::Integer) where {MU} = a == b ? a : NoMSpaceElementSize{MU}() +@inline (::_CommonNDims{MU})(::Any, ::Any) where {MU} = NoMSpaceElementSize{MU}() + +@inline mspace_flatsize(::Type{<:SuperpositionMeasure{C}}) where {C<:AbstractArray} = _scalar_or_unknown(mspace_flatsize(eltype(C))) +@inline function mspace_flatsize(::Type{<:SuperpositionMeasure{C}}) where {C<:Tuple} + _scalar_or_unknown(static_mapreduce(mspace_flatsize, _common_flatsize, C)) end +@inline _common_flatsize(a, b) = a === b ? a : NoMSpaceElementSize{typeof((a, b))}() diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl index c9db7a6b..84aec842 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 @@ -69,25 +72,46 @@ export PushforwardMeasure Users should not call `PushforwardMeasure` directly. Instead call or add methods to `pushfwd`. """ -struct PushforwardMeasure{F,I,M,S<:PushFwdStyle} <: AbstractPushforward +struct PushforwardMeasure{F,I,M,S<:PushFwdStyle,VS} <: AbstractPushforward f::F finv::I origin::M style::S + varsize::VS - function PushforwardMeasure{F,I,M,S}( - f::F, - finv::I, - origin::M, - style::S, - ) where {F,I,M,S<:PushFwdStyle} - new{F,I,M,S}(f, finv, origin, style) + function PushforwardMeasure(f, finv, origin::M, style::S, varsize::VS) where {M,S<:PushFwdStyle,VS} + new{Core.Typeof(f),Core.Typeof(finv),M,S,VS}(f, finv, origin, style, varsize) end +end - function PushforwardMeasure(f, finv, origin::M, style::S) where {M,S<:PushFwdStyle} - new{Core.Typeof(f),Core.Typeof(finv),M,S}(f, finv, origin, style) - end +# The size of the variates of a pushforward follows from a test value of +# the origin, where the origin has variates of known size: +# The output size is learned from a test value whenever the origin's +# variates have a fixed layout, which includes tuple products: +@inline function _pushfwd_varsize(f, μ::MU) where {MU} + _pushfwd_varsize(f, μ, fixed_stream_size(MU)) end +@inline _pushfwd_varsize(f, μ, ::True) = _value_flatsize(f(testvalue(μ))) +@inline _pushfwd_varsize(f, μ, ::False) = NoMSpaceElementSize{typeof(μ)}() + +@inline mspace_elsize(ν::PushforwardMeasure) = _value_or_unknown(ν.varsize, ν) +@inline mspace_flatsize(ν::PushforwardMeasure) = _value_or_unknown(ν.varsize, ν) +@inline _value_or_unknown(sz::SizeLike, ν) = sz +@inline _value_or_unknown(::NoMSpaceElementSize, ν) = NoMSpaceElementSize{typeof(ν)}() +@inline fixed_stream_size(::Type{<:PushforwardMeasure{<:Any,<:Any,<:Any,<:Any,VS}}) where {VS} = static(VS <: SizeLike) +@inline function mspace_ndims(::Type{MU}) where {VS,MU<:PushforwardMeasure{<:Any,<:Any,<:Any,<:Any,VS}} + _ndims_of_size_type(VS, MU) +end +@inline _ndims_of_size_type(::Type{<:Tuple{Vararg{Any,N}}}, ::Type) where {N} = N +@inline mspace_flatsize(::Type{<:PushforwardMeasure{<:Any,<:Any,<:Any,<:Any,Tuple{}}}) = () +@inline mspace_flatsize(::Type{<:PushforwardMeasure{<:Any,<:Any,<:Any,<:Any,StaticArrays.Size{S}}}) where {S} = StaticArrays.Size(S) +@inline _ndims_of_size_type(::Type{StaticArrays.Size{S}}, ::Type) where {S} = length(S) +@inline _ndims_of_size_type(::Type, ::Type{MU}) where {MU} = NoMSpaceElementSize{MU}() + +# Pushforwards by elementwise functions keep the variate rank of their +# origin: +const _ElementwisePushfwd{M,S} = PushforwardMeasure{<:Base.BroadcastFunction,<:Base.BroadcastFunction,M,S} +@inline mspace_ndims(::Type{MU}) where {M,MU<:_ElementwisePushfwd{M}} = mspace_ndims(M) const _NonBijectivePusfwdMeasure{M<:PushforwardMeasure,S<:PushFwdStyle} = Union{ PushforwardMeasure{<:Any,<:NoInverse,M,S}, @@ -115,24 +139,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 +164,7 @@ function logdensityof( ) end -function logdensityof( +function logdensityof_impl( @nospecialize(μ::_NonBijectivePusfwdMeasure{M,<:AdaptRootMeasure}), @nospecialize(v::Any) ) where {M} @@ -154,15 +175,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 +191,39 @@ for func in [:logdensityof, :logdensity_def] end end -insupport(m::PushforwardMeasure, x) = insupport(transport_origin(m), to_origin(m, x)) +# Pushforwards by elementwise functions evaluate densities over flat +# batches, the log-abs-det-Jacobian terms sum over the variate dimensions +# of the origin: +for (bhead, head) in [(:batched_logdensityof_impl, :logdensityof_impl), (:batched_logdensity_def, :logdensity_def)] + @eval function $bhead(ν::_ElementwisePushfwd{M,<:AdaptRootMeasure}, Y) where {M} + _elementwise_pushfwd_ld($head, ν, Y, _static_ndims(ν.origin)) + end + @eval function $bhead(ν::_ElementwisePushfwd{M,<:PushfwdRootMeasure}, Y) where {M} + _batched_kernel($head, ν.origin, broadcast(ν.finv.f, Y)) + end +end + +function _elementwise_pushfwd_ld(f::F, ν::PushforwardMeasure, Y, k::StaticInteger) where {F} + f_inv = ν.finv.f + ℓ = _batched_kernel(f, ν.origin, broadcast(f_inv, Y)) + ladj = sum_leading_dims(broadcast(_LadjOf(f_inv), Y), k) + return _lazy_combine_ladj(ℓ, ladj) +end +function _elementwise_pushfwd_ld(f::F, ν::PushforwardMeasure, Y, ::NoMSpaceElementSize) where {F} + _default_batched_kernel(f, ν, Y, _static_ndims(ν)) +end + +struct _LadjOf{F} <: Function + f::F +end +@inline (k::_LadjOf)(y) = last(with_logabsdet_jacobian(k.f, y)) + +@inline _lazy_combine_ladj(ℓ::Number, ladj::Number) = _combine_logd_with_ladj(ℓ, ladj) +@inline _lazy_combine_ladj(ℓ, ladj) = Broadcast.instantiate(Broadcast.broadcasted(_combine_logd_with_ladj, ℓ, ladj)) + +# 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,17 +246,59 @@ _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 -@inline transport_origin(ν::PushforwardMeasure) = ν.origin -@inline from_origin(ν::PushforwardMeasure, x) = ν.f(x) -@inline to_origin(ν::PushforwardMeasure, y) = ν.finv(y) +# Pushforwards transport via their origin: +@inline transport_to_std(::Type{S}, ν::PushforwardMeasure, y) where {S<:StdMeasure} = + transport_to_std(S, ν.origin, ν.finv(y)) +@inline transport_from_std(::Type{S}, ν::PushforwardMeasure, z) where {S<:StdMeasure} = + ν.f(transport_from_std(S, ν.origin, z)) +@inline function transport_from_std_with_rest(::Type{S}, ν::PushforwardMeasure, z::AbstractVector) where {S<:StdMeasure} + x, z_rest = transport_from_std_with_rest(S, ν.origin, z) + return ν.f(x), z_rest +end -massof(m::PushforwardMeasure) = massof(transport_origin(m)) +# Batches of pushforwards apply the functions to flat batches of the +# origin, elementwise for `Base.BroadcastFunction`s and variate by variate +# (in a host loop) otherwise. The AffineMaps extension adds affine maps. +function batched_transport_to_std(::Type{S}, ν::PushforwardMeasure, Y) where {S<:StdMeasure} + batched_transport_to_std(S, ν.origin, _apply_batched(ν.finv, Y, _static_ndims(ν))) +end +function batched_transport_from_std(::Type{S}, ν::PushforwardMeasure, Z::AbstractArray) where {S<:StdMeasure} + _apply_batched(ν.f, batched_transport_from_std(S, ν.origin, Z), _static_ndims(ν.origin)) +end -function Base.rand(rng::AbstractRNG, ::Type{T}, ν::PushforwardMeasure) where {T} - return ν.f(rand(rng, T, ν.origin)) +# Apply `f` to a flat batch of variates of rank `k`: +@inline _apply_batched(f, X, k) = _apply_generic(unwrap(f), X, k) +@inline _apply_generic(f, X, k) = _apply_by_rank(f, X, k) +@inline _apply_generic(f::Base.BroadcastFunction, X, k) = broadcast(f.f, X) +@inline _apply_by_rank(f, X, ::StaticInteger{0}) = broadcast(f, X) +@inline _apply_by_rank(f, X::AbstractArray, ::StaticInteger{0}) = broadcast(f, X) +@inline _apply_by_rank(f, X::AbstractArray, ::StaticInteger{K}) where {K} = _apply_to_slices(f, X, Val(K)) +@inline _apply_to_slices(f, X::AbstractArray{<:Any,K}, ::Val{K}) where {K} = f(X) +@inline _apply_to_slices(f, X::AbstractArray, ::Val{K}) where {K} = stacked(map(f, sliced(X, Val(K)))) +@noinline function _apply_by_rank(f, X, ::NoMSpaceElementSize) + throw(ArgumentError("Applying functions of type $(nameof(typeof(f))) to batches of variates requires MeasureBase.mspace_ndims to be declared for the measure")) +end + +massof(m::PushforwardMeasure) = massof(m.origin) + +rand_impl(ctx::GenContext, ν::PushforwardMeasure) = ν.f(rand_impl(ctx, ν.origin)) + +# Batches of pushforwards apply the function to the variates of a batch of +# the origin: +function batched_rand_impl(ctx::GenContext, ν::PushforwardMeasure, sz::SizeLike) + _pushfwd_batched_rand(ctx, ν, sz, _static_ndims(ν.origin)) +end +@inline function _pushfwd_batched_rand(ctx::GenContext, ν::PushforwardMeasure, sz::SizeLike, k::StaticInteger) + _apply_batched(ν.f, batched_rand_impl(ctx, ν.origin, sz), k) +end +@inline function _pushfwd_batched_rand(ctx::GenContext, ν::PushforwardMeasure, sz::SizeLike, ::NoMSpaceElementSize) + _batched_rand_pointwise(ctx, ν, sz) end ############################################################################### @@ -222,11 +317,14 @@ 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) -_pushfwd_impl(f, μ, style) = PushforwardMeasure(f, inverse(f), μ, style) +@inline pushfwd(::typeof(identity), μ) = μ +@inline pushfwd(::typeof(identity), μ, ::PushFwdStyle) = μ + +_pushfwd_impl(f, μ, style) = PushforwardMeasure(f, inverse(f), μ, style, _pushfwd_varsize(f, μ)) function _pushfwd_impl( f, @@ -236,11 +334,11 @@ function _pushfwd_impl( orig_μ = μ.origin new_f = fcomp(f, μ.f) new_f_inv = fcomp(μ.finv, inverse(f)) - PushforwardMeasure(new_f, new_f_inv, orig_μ, style) + PushforwardMeasure(new_f, new_f_inv, orig_μ, style, _pushfwd_varsize(new_f, orig_μ)) 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,12 +361,16 @@ 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) end @deprecate pullback(f, μ, style::PushFwdStyle = AdaptRootMeasure()) pullbck(f, μ, style) + +function Adapt.adapt_structure(to, ν::PushforwardMeasure) + PushforwardMeasure(Adapt.adapt(to, ν.f), Adapt.adapt(to, ν.finv), Adapt.adapt(to, ν.origin), ν.style, ν.varsize) +end diff --git a/src/combinators/weighted.jl b/src/combinators/weighted.jl index 124662b6..f7d3b979 100644 --- a/src/combinators/weighted.jl +++ b/src/combinators/weighted.jl @@ -12,13 +12,30 @@ abstract type AbstractWeightedMeasure <: AbstractMeasure end # By default the weight for all measure is 1 _logweight(::AbstractMeasure) = 0 -@inline function logdensity_def(d::AbstractWeightedMeasure, _) - d.logweight +@inline logdensity_def(d::AbstractWeightedMeasure, x) = _logweight_for(d.logweight, x) + +# Plain floating-point log-weights adopt the number type of the variate, +# log-weights that carry more information (dual numbers, traced values) +# promote as usual: +@inline _logweight_for(w, x) = w +@inline _logweight_for(w::Union{AbstractFloat,StaticFloat64}, x) = _logd_numtype(x)(dynamic(w)) + +# 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) + _logweight_for(d.logweight, x) + logdensityof_impl(basemeasure(d), x) end -function Base.rand(rng::AbstractRNG, ::Type{T}, μ::AbstractWeightedMeasure) where {T} - rand(rng, T, basemeasure(μ)) +@inline function batched_logdensityof_impl(d::AbstractWeightedMeasure, X) + _lazy_add(_logweight_for(d.logweight, X), batched_logdensityof_impl(basemeasure(d), X)) end +@inline function batched_logdensity_def(d::AbstractWeightedMeasure, X) + _lazy_add(_logweight_for(d.logweight, X), _zero_logd_batch(X, _static_ndims(basemeasure(d)))) +end + +@inline rand_impl(ctx::GenContext, μ::AbstractWeightedMeasure) = rand_impl(ctx, basemeasure(μ)) +@inline batched_rand_impl(ctx::GenContext, μ::AbstractWeightedMeasure, sz::SizeLike) = + batched_rand_impl(ctx, basemeasure(μ), sz) testvalue(::Type{T}, μ::AbstractWeightedMeasure) where {T} = testvalue(T, basemeasure(μ)) @@ -29,6 +46,12 @@ struct WeightedMeasure{R,M} <: AbstractWeightedMeasure base::M end +@inline mspace_elsize(μ::WeightedMeasure) = mspace_elsize(μ.base) +@inline mspace_flatsize(μ::WeightedMeasure) = mspace_flatsize(μ.base) +@inline mspace_flatsize(::Type{<:WeightedMeasure{<:Any,M}}) where {M} = mspace_flatsize(M) +@inline mspace_ndims(::Type{<:WeightedMeasure{<:Any,M}}) where {M} = mspace_ndims(M) +@inline fixed_stream_size(::Type{<:WeightedMeasure{<:Any,M}}) where {M} = fixed_stream_size(M) + massof(w::WeightedMeasure) = exp(w.logweight) * massof(w.base) _logweight(μ::WeightedMeasure) = μ.logweight @@ -44,14 +67,31 @@ 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) insupport(μ::WeightedMeasure, x) = insupport(μ.base, x) -# TODO: Transports must preserve mass -transport_origin(ν::WeightedMeasure) = ν.base - -to_origin(w::WeightedMeasure, y) = y -from_origin(w::WeightedMeasure, x) = x +# Weighted measures transport like their base: +@inline transport_to_std(::Type{S}, μ::AbstractWeightedMeasure, x) where {S<:StdMeasure} = + transport_to_std(S, basemeasure(μ), x) +@inline transport_from_std(::Type{S}, μ::AbstractWeightedMeasure, z) where {S<:StdMeasure} = + transport_from_std(S, basemeasure(μ), z) +@inline transport_to_std_with_rest(::Type{S}, μ::AbstractWeightedMeasure, x::AbstractVector) where {S<:StdMeasure} = + transport_to_std_with_rest(S, basemeasure(μ), x) +@inline transport_to_std_with_rest(::Type{S}, μ::AbstractWeightedMeasure, x::NamedTuple) where {S<:StdMeasure} = + transport_to_std_with_rest(S, basemeasure(μ), x) +@inline transport_from_std_with_rest(::Type{S}, μ::AbstractWeightedMeasure, z::AbstractVector) where {S<:StdMeasure} = + transport_from_std_with_rest(S, basemeasure(μ), z) + +@inline batched_transport_to_std(::Type{S}, μ::AbstractWeightedMeasure, X) where {S<:StdMeasure} = + batched_transport_to_std(S, basemeasure(μ), X) +@inline batched_transport_from_std(::Type{S}, μ::AbstractWeightedMeasure, Z::AbstractArray) where {S<:StdMeasure} = + batched_transport_from_std(S, basemeasure(μ), Z) +@inline batched_transport_to_std_with_rest(::Type{S}, μ::AbstractWeightedMeasure, X::AbstractArray, sz::SizeLike) where {S<:StdMeasure} = + batched_transport_to_std_with_rest(S, basemeasure(μ), X, sz) +@inline batched_transport_from_std_with_rest(::Type{S}, μ::AbstractWeightedMeasure, Z::AbstractArray, sz::SizeLike) where {S<:StdMeasure} = + batched_transport_from_std_with_rest(S, basemeasure(μ), Z, sz) + +Adapt.adapt_structure(to, μ::WeightedMeasure) = WeightedMeasure(μ.logweight, Adapt.adapt(to, μ.base)) diff --git a/src/density-batched.jl b/src/density-batched.jl new file mode 100644 index 00000000..a7f0d660 --- /dev/null +++ b/src/density-batched.jl @@ -0,0 +1,333 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +# Batched-first density evaluation. +# +# A flat batch of variates is an array `(variate dims..., batch dims...)`, +# zero batch dims meaning a single variate. Nested batches (arrays of +# variates with flat storage, see ArraysOfArrays) are fused at the entry +# points, the container then determines the batch dimensions. Batched +# kernels know the variate rank of their measure (see `mspace_ndims`) and +# return arrays over the batch dimensions, a number for a single variate. +# Structural measures implement their kernels once, in terms of the kernels +# of their components. + +export logdensities + +""" + logdensities(μ::AbstractMeasure, X) + +Compute the log-density of `μ` at each variate in the batch `X`. + +`X` is an array of variates (e.g. an `ArraysOfArrays.ArrayOfSimilarArrays`, +or an array of numbers for measures with scalar variates), the flat storage +of a batch (an array of numbers with the variate dimensions leading, see +[`MeasureBase.mspace_ndims`](@ref)), or a tuple resp. `NamedTuple` of +batches for measures with tuple resp. `NamedTuple` variates. Returns an +array over the batch dimensions (a number for a single variate), +semantically equivalent to `logdensityof.(Ref(μ), X)` for arrays of +variates. Batches with flat storage are evaluated in fused operations, +compatible with GPU and traced arrays. + +Measure types implement [`MeasureBase.batched_logdensityof_impl`](@ref). +""" +function logdensities end + +@inline logdensities(μ::AbstractMeasure, X) = _materialize(_batched_ld(logdensityof_impl, μ, X)) + +""" + MeasureBase.batched_logdensityof_impl(μ::AbstractMeasure, X) + +Log-densities of `μ` (relative to its root measure) at the variates of the +flat batch `X`, an array `(variate dims..., batch dims...)`, returned as an +array over the batch dimensions. `X` may be a single variate (zero batch +dimensions), the result is a number then. Results may be lazy broadcasts, +callers materialize them where necessary. The results must be `-Inf` for +variates outside the support of `μ`. + +This is the primary density extension point. The default implementation +broadcasts the point kernel [`MeasureBase.logdensityof_impl`](@ref) over +`X` for measures with scalar variates and maps it over the variate slices +of `X` (in a host loop) for measures with array variates of a declared +number of dimensions (see [`MeasureBase.mspace_ndims`](@ref)). Measure +types with array variates should implement `batched_logdensityof_impl` +directly. Structural measures evaluate their components via +`batched_logdensityof_impl` as well. +""" +function batched_logdensityof_impl end + +@inline function batched_logdensityof_impl(μ::AbstractMeasure, X) + _default_batched_kernel(logdensityof_impl, μ, X, _static_ndims(μ)) +end + +# The variate rank as a static integer, from the type where known: +@inline _static_ndims(μ::MU) where {MU} = _static_ndims(mspace_ndims(MU), μ) +@inline _static_ndims(n::Integer, μ) = static(n) +@inline _static_ndims(::NoMSpaceElementSize, μ) = _static_ndims_of(mspace_ndims(μ)) +@inline _static_ndims_of(n::Integer) = static(n) +@inline _static_ndims_of(n::NoMSpaceElementSize) = n + +""" + MeasureBase.batched_logdensity_def(μ::AbstractMeasure, X) + +Batched form of [`logdensity_def`](@ref): log-densities of `μ` relative to +`basemeasure(μ)` at the variates of the flat batch `X`, with the same +conventions and defaults as [`MeasureBase.batched_logdensityof_impl`](@ref). +""" +function batched_logdensity_def end + +@inline function batched_logdensity_def(μ::AbstractMeasure, X) + _default_batched_kernel(logdensity_def, μ, X, _static_ndims(μ)) +end + +@inline _default_batched_kernel(f::F, μ, X, n::Integer) where {F} = _default_batched_kernel(f, μ, X, static(n)) +@inline _default_batched_kernel(f::F, μ, X, ::StaticInteger{0}) where {F} = _scalar_kernel_broadcast(f, μ, X) +@inline _default_batched_kernel(f::F, μ, X::AbstractArray, ::StaticInteger{0}) where {F} = _scalar_kernel_broadcast(f, μ, X) +@inline function _default_batched_kernel(f::F, μ, X::AbstractArray, ::StaticInteger{K}) where {F,K} + _map_variate_slices(f, μ, X, Val(K)) +end +@noinline function _default_batched_kernel(f::F, μ, X, ::NoMSpaceElementSize) where {F} + throw(ArgumentError("Batched density evaluation for measures of type $(nameof(typeof(μ))) requires MeasureBase.mspace_ndims to be declared for the type or MeasureBase.batched_logdensityof_impl to be implemented")) +end + +# Point kernels of scalar-variate measures broadcast over the batch. Static +# results are made dynamic to keep reductions type stable. +@inline function _scalar_kernel_broadcast(f::F, μ, X::AbstractArray) where {F} + Broadcast.instantiate(Broadcast.broadcasted(_DynamicLogd(f, μ), X)) +end +@inline _scalar_kernel_broadcast(f::F, μ, x) where {F} = _dynamic_logd(f(μ, x), x) + +struct _DynamicLogd{F,M} <: Function + f::F + μ::M +end +@inline (k::_DynamicLogd)(x) = _dynamic_logd(k.f(k.μ, x), x) + +# Point kernels of array-variate measures map over the variate slices of +# the batch, a batch of a single variate is evaluated directly: +@inline _map_variate_slices(f::F, μ, X::AbstractArray{<:Any,K}, ::Val{K}) where {F,K} = _dynamic_logd(f(μ, X), X) +@inline function _map_variate_slices(f::F, μ, X::AbstractArray, ::Val{K}) where {F,K} + map(_DynamicLogd(f, μ), sliced(X, Val(K))) +end + +# The batched kernel for a point-level density function: +@inline _batched_kernel(::typeof(logdensityof_impl), μ, X) = batched_logdensityof_impl(μ, X) +@inline _batched_kernel(::typeof(logdensity_def), μ, X) = batched_logdensity_def(μ, X) + +# Flat storage of a (nested) batch: the underlying array of memory-ordered +# split arrays, a stacked copy for other known split modes. Struct arrays +# of tuple and named tuple variates have the flat storage of their +# components (copied where necessary, the batch dimensions are shared). +struct NoFlatStorage end +@inline _flat_storage(X::AbstractArray{<:Number}) = X +@inline _flat_storage(X::AbstractArray) = _flat_storage_bymode(X, getsplitmode(X)) +@inline _flat_storage(x) = NoFlatStorage() +@inline function _flat_storage_bymode(X::AbstractArray, smode::AbstractSplitMode) + _flat_storage(is_memordered_splitmode(smode) ? fused(X) : stacked(X)) +end +@inline _flat_storage_bymode(::AbstractArray, ::UnknownSplitMode) = NoFlatStorage() +@inline _flat_storage_bymode(::AbstractArray, ::NonSplitMode) = NoFlatStorage() +# Ragged batches (e.g. `VectorOfArrays`) are evaluated variate by variate: +@inline _flat_storage_bymode(::AbstractArray, ::AbstractPartMode) = NoFlatStorage() + +@inline _flat_storage(X::StructArray{<:Union{Tuple,NamedTuple}}) = _components_storage(StructArrays.components(X)) +@inline function _flat_storage(X::AbstractArray{<:Union{Tuple,NamedTuple}}) + _flat_storage(StructArray(X; unwrap = T -> T <: Union{Tuple,NamedTuple})) +end +@inline _components_storage(cs::Tuple) = map(_component_storage, cs) +@inline _components_storage(cs::NamedTuple{names}) where {names} = NamedTuple{names}(map(_component_storage, values(cs))) +@inline _component_storage(c::StructArray{<:Union{Tuple,NamedTuple}}) = _flat_storage(c) +@inline _component_storage(c::AbstractArray{<:Number}) = c +@inline _component_storage(c::AbstractArray{<:AbstractArray}) = _component_flat(c, _flat_storage(c)) +@inline _component_flat(c, c_flat::AbstractArray) = c_flat +@inline _component_flat(c, ::NoFlatStorage) = stacked(c) + +# Entry: arrays of numbers are flat storage, arrays of variates are fused +# into their flat storage (else evaluated variate by variate), tuples and +# named tuples of batches go to the kernels directly. +@inline _batched_ld(f::F, μ, X::AbstractArray{<:Number}) where {F} = _batched_kernel(f, μ, X) +@inline _batched_ld(f::F, μ, X::Union{Tuple,NamedTuple}) where {F} = _batched_kernel(f, μ, X) +@inline _batched_ld(f::F, μ, X::AbstractArray) where {F} = _batched_ld_nested(f, μ, X, _flat_storage(X), _static_ndims(μ)) +@inline function _batched_ld_nested(f::F, μ, X::AbstractArray, X_flat::AbstractArray, ::StaticInteger) where {F} + _check_batch_shape(_batched_kernel(f, μ, X_flat), X) +end +@inline function _batched_ld_nested(f::F, μ, X::AbstractArray, X_flat::Union{Tuple,NamedTuple}, ::Any) where {F} + _check_batch_shape(_batched_kernel(f, μ, X_flat), X) +end +@inline function _batched_ld_nested(f::F, μ, X::AbstractArray, ::Any, ::Any) where {F} + Broadcast.instantiate(Broadcast.broadcasted(_PointLogd(f, μ), X)) +end + +struct _PointLogd{F,M} <: Function + f::F + μ::M +end +@inline (k::_PointLogd)(x) = _point_ld(k.f, k.μ, x) + +# Results over an array of variates must have the shape of the array, a +# mismatch means the variate rank of the measure doesn't match the batch: +@inline function _check_batch_shape(result, X::AbstractArray) + if size(result) != size(X) + _throw_batch_shape(size(result), size(X)) + end + return result +end +@noinline function _throw_batch_shape(sz_result, sz_batch) + throw(ArgumentError("Batched density kernel returned a result of size $sz_result for a batch of size $sz_batch, the variate dimensions of the measure don't match the batch")) +end + +# Point evaluation: array variates of measures with a declared variate rank +# go through the batched kernel with zero batch dimensions, other variates +# through the point kernel. A batched kernel that returns an array for a +# single variate has taken the variate for a batch: the variate doesn't +# fit the measure, or the measure lacks a batched kernel for array +# variates. +@inline _point_ld(f::F, μ, x::AbstractArray{<:Number}) where {F} = _point_ld_byrank(f, μ, x, _static_ndims(μ)) +@inline _point_ld(f::F, μ, x) where {F} = f(μ, x) +@inline _point_ld_byrank(f::F, μ, x, ::StaticInteger) where {F} = _point_result(_materialize(_batched_kernel(f, μ, x)), μ) +@inline _point_ld_byrank(f::F, μ, x, ::NoMSpaceElementSize) where {F} = f(μ, x) +@inline _point_ld(f::F, μ::PrimitiveMeasure, x::AbstractArray{<:Number}) where {F} = f(μ, x) + +@inline _point_result(ℓ::Number, μ) = ℓ +@inline _point_result(ℓ::AbstractArray{<:Number,0}, μ) = ℓ[] +@noinline function _point_result(ℓ, μ) + throw(ArgumentError("Density evaluation of measures of type $(nameof(typeof(μ))) at an array variate resulted in a batch of densities: the variate doesn't fit the measure, or the measure lacks a batched kernel for array variates")) +end + +const _LazyBroadcast = Broadcast.Broadcasted + +@noinline _throw_size_mismatch() = throw(ArgumentError("Size of variate doesn't match size of measure")) + +# The leading dimensions of a flat batch must match a flat variate size: +@inline function _check_flatsize(A::AbstractArray, sz_flat::SizeLike) + n = length(size_dims(sz_flat)) + if ndims(A) < n || ntuple(i -> size(A, i), Val(n)) != asnonstatic(sz_flat) + _throw_size_mismatch() + end + return nothing +end + +@inline _materialize(bc::_LazyBroadcast) = copy(bc) +@inline _materialize(x) = x + + +@inline _lazy_add(a::Number, b::Number) = a + b +@inline _lazy_add(a, b) = Broadcast.instantiate(Broadcast.broadcasted(+, a, b)) + +# Zero log-densities over the batch dimensions of a flat batch of variates +# with `n` variate dimensions: +@inline function _zero_logd_batch(X::AbstractArray, n::Integer) + FillArrays.Zeros{_logd_numtype(X)}(ntuple(i -> size(X, n + i), ndims(X) - n)) +end +@inline _zero_logd_batch(X::AbstractArray{<:Any,N}, ::StaticInteger{N}) where {N} = zero(_logd_numtype(X)) +@inline _zero_logd_batch(x::Number, ::StaticInteger{0}) = zero(_logd_numtype(x)) +@noinline function _zero_logd_batch(X, ::NoMSpaceElementSize{MU}) where {MU} + throw(ArgumentError("Batched density evaluation for measures of type $(nameof(MU)) requires MeasureBase.mspace_ndims to be declared for the type or MeasureBase.batched_logdensity_def to be implemented")) +end + + +# Streams: variates of composed measures are consumed from flat vector +# streams, batched as `(rows, batch dims...)`. + +""" + MeasureBase.batched_logdensityof_with_rest(μ::AbstractMeasure, X, sz::SizeLike) + +Consume variates of `μ` from the batch `X` of flat vector streams (first +dimension along the streams, further dimensions are batch dimensions), a +batch of variates of size `sz` per stream, and compute their +log-densities. + +Returns a tuple `(ℓ, X_rest)` of the log-densities, an array over +`(sz..., batch dims...)` (possibly lazy, a number for a single stream and +`sz == ()`), and the unconsumed rest of the streams. Measure types whose +variates have a fixed size consume `prod(sz)` variates in one batched +kernel evaluation, the default implementation does so for the variate size +given by [`MeasureBase.mspace_flatsize`](@ref) or +[`MeasureBase.some_mspace_elsize`](@ref). Measure types with variates of +value-dependent size implement `batched_logdensityof_with_rest` for +single streams and `sz == ()` themselves and report +`MeasureBase.fixed_stream_size` as false, so that the enclosing stream +combinators consume batches stream by stream. +""" +function batched_logdensityof_with_rest end + +function batched_logdensityof_with_rest(μ::AbstractMeasure, X::AbstractArray, sz::SizeLike) + _stream_ld_with_rest(logdensityof_impl, μ, X, sz) +end + +# A single stream consumes one variate via the point path: +function batched_logdensityof_with_rest(μ::AbstractMeasure, x::AbstractVector, ::Tuple{}) + ℓ, _, x_rest = logdensityof_with_rest(μ, x) + return ℓ, x_rest +end + +function _stream_ld_with_rest(f::F, μ, X::AbstractArray, sz::SizeLike) where {F} + vsz = _stream_consume_size(μ) + X_μ, X_rest = _batched_consume(X, vsz, sz) + return _consumed_ld(f, μ, X_μ, vsz), X_rest +end + +# Scalar variates are consumed as `(1, sz..., batch dims...)`, their +# leading dimension is dropped before the kernel runs: +@inline _consumed_ld(f::F, μ, X_μ, ::Tuple{}) where {F} = _batched_kernel(f, μ, _drop_stdstream_dim(X_μ)) +@inline _consumed_ld(f::F, μ, X_μ, ::SizeLike) where {F} = _batched_kernel(f, μ, X_μ) + +# Consume `prod(sz)` variates of flat size `vsz` from the leading rows of a +# batch of streams as a flat batch `(vsz..., sz..., batch dims...)`; scalar +# variates as `(1, sz..., batch dims...)`. Static sizes keep static +# streams static. +@inline function _batched_consume(X::AbstractArray, vsz::SizeLike, sz::SizeLike) + dims = _consumed_dims(vsz) + X_flat, X_rest = _batched_split(X, _chunk_rows(prod(dims), sz)) + return _reshape_consumed(X_flat, (dims..., size_dims(sz)...)), X_rest +end +@inline _consumed_dims(::Tuple{}) = (static(1),) +@inline _consumed_dims(vsz::SizeLike) = size_dims(vsz) +@inline _chunk_rows(n::IntegerLike, sz::SizeLike) = n * size2length(sz) + +@inline _reshape_consumed(X_flat::AbstractArray, ::Tuple{IntegerLike}) = X_flat +@inline function _reshape_consumed(X_flat::AbstractArray, dims::Tuple{Vararg{IntegerLike}}) + maybestatic_reshape(X_flat, (dims..., Base.tail(_batch_dims(X_flat))...)) +end + +# Sizes as tuples of (maybe static) integers and the leading dimension of +# a batch of streams: +@inline _batch_dims(A::AbstractArray) = size_dims(maybestatic_size(A)) +@inline _as_stdstream_batch(Z::AbstractArray) = merge_leading_dims(Z, static(0)) +@inline _as_stdstream_batch(z::Number) = SVector(z) +@inline _drop_stdstream_dim(Z::AbstractArray) = drop_leading_dims(Z, static(1)) + +@inline function _batched_split(A::AbstractArray, n::IntegerLike) + n_rows = dynamic(n) + stream_idxs = axes(A, 1) + if length(stream_idxs) < n_rows + throw(ArgumentError("Variate streams too short during batched evaluation")) + end + batch_axes = Base.tail(axes(A)) + i0 = first(stream_idxs) + A_flat = view(A, i0:(i0 + n_rows - 1), batch_axes...) + A_rest = view(A, (i0 + n_rows):last(stream_idxs), batch_axes...) + return A_flat, A_rest +end + +# Static streams split into static chunks for static row counts: +@inline _batched_split(A::StaticVector, n_rows::StaticInteger) = split_at(A, n_rows) + +@noinline function _throw_stream_too_long() + throw(ArgumentError("Variate streams too long during density evaluation")) +end + +# Whether all variates consumed by a measure from streams have sizes that +# are fixed at the type level, so that batches of streams can be consumed +# in fused operations; otherwise a batch of streams is consumed stream by +# stream by the outermost stream combinator. +@inline fixed_stream_size(μ::MU) where {MU} = fixed_stream_size(MU) +@inline fixed_stream_size(::Type{MU}) where {MU} = static(mspace_ndims(MU) isa Integer) + +# Batches of streams consumed stream by stream (host loop): +function _streamwise_ld(f::F, μ, X::AbstractArray) where {F} + map(sliced(X, Val(1))) do x + ℓ, x_rest = batched_logdensityof_with_rest(μ, x, ()) + isempty(x_rest) || _throw_stream_too_long() + _dynamic_logd(_materialize(ℓ), x) + end +end diff --git a/src/density-core.jl b/src/density-core.jl index 6ac3d01e..3859f28a 100644 --- a/src/density-core.jl +++ b/src/density-core.jl @@ -10,34 +10,156 @@ 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`. + +# Extended help + +Variates of the right shape and element type never throw: outside the +support of `m` the result is `-Inf`, also for non-integer values of +measures over counting measures and for infinite values. `NaN` inputs give +`NaN` or `-Inf`. Variates of the wrong shape throw an `ArgumentError`. +Implementations of `logdensityof_impl` and `unsafe_logdensityof` must not +throw outside the support, since support masks evaluate both branches. +""" +@inline logdensityof(μ::AbstractMeasure, x) = _point_ld(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(μ::AbstractMeasure, x) - result = dynamic(unsafe_logdensityof(μ, x)) +@inline function logdensityof_impl(μ::AbstractMeasure, x) + result = _dynamic_logd(unsafe_logdensityof(μ, x), x) _checksupport(insupport(μ, x), result) end -@inline function logdensityof_rt(::T, ::U) where {T,U} - Core.Compiler.return_type(logdensityof, Tuple{T,U}) +# Log-density kernels return numbers of the number type of the variate, +# never static numbers, so that automatic differentiation and tracing see +# ordinary floating point values throughout: +@inline _logd_numtype(x) = float(real_numtype(typeof(x))) +@inline _dynamic_logd(ℓ, x) = dynamic(ℓ) +@inline _neg_inf_logd(x) = _logd_numtype(x)(-Inf) + +# Support checks as masks: `NoFastInsupport` means the density is evaluated +# unconditionally. +@inline _insupport_mask(ins) = ins == true +@inline _insupport_mask(::NoFastInsupport) = true + +# Support checks as booleans, keeping `NoFastInsupport`: +@inline _insupport_bool(ins) = ins == true +@inline _insupport_bool(ins::NoFastInsupport) = ins + +# Combining support checks of components, `NoFastInsupport` is absorbing: +@inline _insupport_and(a, b) = _insupport_bool(a) & _insupport_bool(b) +@inline _insupport_and(a::NoFastInsupport, ::Any) = a +@inline _insupport_and(::Any, b::NoFastInsupport) = b +@inline _insupport_and(a::NoFastInsupport, ::NoFastInsupport) = a + +@inline _checksupport(cond, result) = ifelse(_insupport_mask(cond), result, oftype(result, -Inf)) + +# Transports of variates outside the support of the source measure give +# NaN. Both branches are evaluated, formulas must not throw outside the +# support: +@inline _nan_outside(μ, x, y) = ifelse(_insupport_mask(insupport(μ, x)), y, oftype(y, NaN)) + +# On the floating-point grid the endpoints of the unit interval stand for +# their nearest interior points (the smallest normal float above zero, +# since devices may flush subnormals, and the grid point below one), so +# that quantiles stay finite, and tail probabilities in log-space +# conversions never underflow to zero: +@inline _unit_interior(p) = clamp(p, _unit_bounds(p)...) +@inline _unit_bounds(p) = (_prob_floor(p), prevfloat(one(p))) +@inline _positive_prob(p) = max(p, _prob_floor(p)) +@inline _prob_floor(p) = floatmin(typeof(one(p))) + +""" + MeasureBase.logdensityof_with_rest(μ::AbstractMeasure, x) + +Consume the variate of `μ` at the beginning of the flat vector stream `x` +(or the named entries of the `NamedTuple` `x`) and compute its log-density. + +Returns a tuple `(ℓ, x_μ, x_rest)` of the log-density, the consumed variate +`x_μ` and the unconsumed rest of `x`. Measures whose variates have a fixed +size consume that size (see [`MeasureBase.mspace_flatsize`](@ref) and +[`MeasureBase.some_mspace_elsize`](@ref)), measures with variates of +value-dependent size implement the consumption themselves. Batches of +streams are consumed by [`MeasureBase.batched_logdensityof_with_rest`](@ref). +""" +function logdensityof_with_rest end + +function logdensityof_with_rest(μ::AbstractMeasure, x::AbstractVector) + a, x_rest = _consume_from_stream(x, _stream_consume_size(μ)) + return _point_ld(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 -_checksupport(cond, result) = ifelse(cond == true, result, oftype(result, -Inf)) +@inline _stream_consume_size(μ) = _stream_consume_size(μ, mspace_flatsize(μ)) +@inline _stream_consume_size(μ, sz::SizeLike) = sz +@inline _stream_consume_size(μ, ::NoMSpaceElementSize) = some_mspace_elsize(μ) + +_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,19 +172,23 @@ 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}) # return ℓ_{i - 1} # end - ℓ_{i} = let Δℓ_{i} = logdensity_def(b_{i}, x) - ℓ_{i - 1} + Δℓ_{i} - end + ℓ_{i} = ℓ_{i - 1} + logdensity_def(b_{i}, x) end return ℓ_10 end @@ -75,19 +201,12 @@ whether `x` is in the support of `m1` or `m2` (or both, or neither). If `x` is 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} - 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) +@inline function logdensity_rel(μ, ν, x) + inμ = _insupport_mask(insupport(μ, x)) + inν = _insupport_mask(insupport(ν, x)) + logd = _dynamic_logd(unsafe_logdensity_rel(μ, ν, x), x) + outside = ifelse(inμ, oftype(logd, +Inf), ifelse(inν, oftype(logd, -Inf), oftype(logd, NaN))) + return ifelse(inμ & inν, logd, outside) end """ @@ -98,71 +217,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 +# Indicates that no specialized method is available to compute the +# log-density between a given pair of measures: +struct _NoLogdensityRel end + +""" + MeasureBase.logdensity_rel_def(μ, ν, x) + +Specialization point for the log-density of `μ` relative to `ν` at `x`. + +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 -@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ν) - - q = quote - $(Expr(:meta, :inline)) - ℓ = logdensity_def(μs[$M], νs[$N], x) +@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 - - for i in 1:M-1 - push!(q.args, :(Δℓ = logdensity_def(μs[$i], x))) - # push!(q.args, :(println("Adding", Δℓ))) - push!(q.args, :(ℓ += Δℓ)) + 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 - - for j in 1:N-1 - push!(q.args, :(Δℓ = logdensity_def(νs[$j], x))) - # push!(q.args, :(println("Subtracting", Δℓ))) - push!(q.args, :(ℓ -= Δℓ)) + 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 - push!(q.args, :(return ℓ)) - return q +# 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)) + +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..f6dabe46 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 @@ -116,78 +95,151 @@ struct DensityMeasure{F,B} <: AbstractMeasure end @inline function insupport(d::DensityMeasure, x) - insupport(d.base, x) == true && isfinite(logdensityof(getfield(d, :f), x)) + _insupport_mask(insupport(d.base, x)) & isfinite(logdensityof(getfield(d, :f), x)) 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. + +If a density is passed, it must implement the DensityInterface API. -Define a new measure in terms of a density `f` over some measure `base`. +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)) + +@inline _as_integrand_default_impl(f, ::NoDensity) = funcdensity(f) + +@inline _as_integrand_default_impl(density, ::IsDensity) = density -_densitymeasure(f, base, ::IsDensity) = DensityMeasure(f, base) -function _densitymeasure(f, base, ::HasDensity) - @error "`∫(f, base)` requires `DensityKind(f)` to be `IsDensity()` or `NoDensity()`." +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 -""" - ∫exp(f, base::AbstractMeasure) +@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. -Define a new measure in terms of a log-density `f` over some measure `base`. +`ν = mintegrate(f, μ)` generates a measure `ν` that has the mathematical +interpretation + +```math +\nu(A) = \int_A f(a) \, \rm{d}\mu(a) +``` """ -∫exp(f, base) = _logdensitymeasure(f, base, DensityKind(f)) +function mintegrate end +export mintegrate -function _logdensitymeasure(f, base, ::IsDensity) - @error "`∫exp(f, base)` is not valid when `DensityKind(f) == IsDensity()`. Use `∫(f, base)` instead." -end -function _logdensitymeasure(f, base, ::HasDensity) - @error "`∫exp(f, base)` is not valid when `DensityKind(f) == HasDensity()`." -end -_logdensitymeasure(f, base, ::NoDensity) = DensityMeasure(logfuncdensity(f), base) +@inline mintegrate(obj, μ::AbstractMeasure) = DensityMeasure(as_integrand(obj), μ) -basemeasure(μ::DensityMeasure) = μ.base -logdensity_def(μ::DensityMeasure, x) = logdensityof(μ.f, x) +""" + MeasureBase.as_integrand_exp(log_f) -density_def(μ::DensityMeasure, x) = densityof(μ.f, x) +Convert the logarithm of an integrand to an integrand. -function logdensityof(μ::DensityMeasure, x::Any) - integrand, μ_base = μ.f, μ.base +See also [`MeasureBase.as_integrand`](@ref). +""" +function as_integrand_exp end - base_logval = logdensityof(μ_base, x) +@inline as_integrand_exp(log_f) = _as_integrand_exp_default_impl(log_f, DensityKind(log_f)) - T = typeof(base_logval) - U = logdensityof_rt(integrand, x) - R = promote_type(T, U) +@inline _as_integrand_exp_default_impl(log_f, ::NoDensity) = logfuncdensity(log_f) - # Don't evaluate base measure if integrand is zero or NaN - if isneginf(base_logval) - R(-Inf) - else - integrand_logval = logdensityof(integrand, x) - convert(R, integrand_logval + base_logval)::R - end +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 -""" - rebase(μ, ν) +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 -Express `μ` in terms of a density over `ν`. Satisfies + +@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) ``` -basemeasure(rebase(μ, ν)) == ν -density(rebase(μ, ν)) == 𝒹(μ,ν) -``` + +Note that `exp(log_f(...))` is usually not run explicitly, calculations that +involve the resulting measure are typically performed in log-space, +internally. """ -rebase(μ, ν) = ∫(𝒹(μ, ν), ν) +function mintegrate_exp end +export mintegrate_exp + +mintegrate_exp(log_f, μ::AbstractMeasure) = DensityMeasure(as_integrand_exp(log_f), μ) + + +basemeasure(μ::DensityMeasure) = μ.base + +@inline mspace_elsize(μ::DensityMeasure) = mspace_elsize(μ.base) +@inline mspace_flatsize(μ::DensityMeasure) = mspace_flatsize(μ.base) +@inline mspace_flatsize(::Type{<:DensityMeasure{<:Any,B}}) where {B} = mspace_flatsize(B) +@inline mspace_ndims(::Type{<:DensityMeasure{<:Any,B}}) where {B} = mspace_ndims(B) +@inline fixed_stream_size(::Type{<:DensityMeasure{<:Any,B}}) where {B} = fixed_stream_size(B) + +logdensity_def(μ::DensityMeasure, x) = logdensityof(μ.f, x) + +density_def(μ::DensityMeasure, x) = densityof(μ.f, x) + +# Density measures evaluate the base measure and the integrand directly, +# the base measure validates array variates: +@inline _point_ld(f::F, μ::DensityMeasure, x::AbstractArray{<:Number}) where {F} = f(μ, x) + +function logdensityof_impl(μ::DensityMeasure, x::Any) + integrand, μ_base = μ.f, μ.base + base_logval = dynamic(logdensityof(μ_base, x)) + integrand_logval = dynamic(logdensityof(integrand, x)) + logval = integrand_logval + base_logval + # Outside of the support of the base measure the integrand may be + # anything, including NaN: + ifelse(isneginf(base_logval), oftype(logval, -Inf), logval) +end diff --git a/src/domains.jl b/src/domains.jl index e03f753c..643f1ec0 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 -Base.in(x, b::BoundedReals) = b.lower ≤ x ≤ b.upper +""" + struct MeasureBase.UnknownDomain{T} <: MeasureBase.ValueSet + +Represents the unknown domain of a value of type `T`. + +Constructors: + +``` +MeasureBase.UnknownDomain(x::T) +``` -export ℝ, ℝ₊, 𝕀, ℤ +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). -const ℝ₊ = BoundedReals(static(0.0), static(Inf)) -const 𝕀 = BoundedReals(static(0.0), static(1.0)) +`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 -Base.minimum(b::BoundedReals) = b.lower -Base.maximum(b::BoundedReals) = b.upper +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.show(io::IO, ::typeof(ℝ₊)) = print(io, "ℝ₊") -Base.show(io::IO, ::typeof(𝕀)) = print(io, "𝕀") +maybe_in(@nospecialize(x), ::UnknownDomain) = true -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 +Base.isempty(::UnknownDomain) = false -abstract type IntegerDomain <: AbstractDomain end +""" + RealValues() isa MeasureBase.ValueSet -struct IntegerNumbers <: IntegerDomain end +The real numbers. +""" +struct RealValues <: ValueSet end -Base.in(x, ::IntegerNumbers) = isinteger(x) +@inline Base.in(x::Real, ::RealValues) = true +@inline Base.in(x, ::RealValues) = isreal(x) -const ℤ = IntegerNumbers() +@inline Base.isempty(::RealValues) = false -Base.show(io::IO, ::typeof(ℤ)) = print(io, "ℤ") +@inline Base.union(s::RealValues, ::RealValues...) = s -Base.minimum(::IntegerNumbers) = static(-Inf) -Base.maximum(::IntegerNumbers) = static(Inf) -struct BoundedInts{L,U} <: IntegerDomain +@inline Base.minimum(::RealValues) = static(-Inf) +@inline Base.maximum(::RealValues) = static(Inf) + +testvalue(::Type{T}, ::RealValues) where {T} = zero(T) + +""" + const MeasureBase.ℝ = RealValues() + +The set of all real numbers, see [`MeasureBase.RealValues`](@ref). +""" +const ℝ = RealValues() +export ℝ + +Base.show(io::IO, ::RealValues) = print(io, "ℝ") +Base.show(io::IO, ::MIME"text/plain", ::RealValues) = print(io, "MeasureBase.ℝ") + +""" + MeasureBase.IntegerValues() isa MeasureBase.ValueSet +""" +struct IntegerValues <: ValueSet end + +@inline Base.in(x::Integer, ::IntegerValues) = true +@inline Base.in(x, ::IntegerValues) = isinteger(x) + +@inline Base.isempty(::IntegerValues) = false + +@inline Base.union(s::IntegerValues, ::IntegerValues...) = s + +testvalue(::Type{T}, ::IntegerValues) where {T} = zero(T) + +# # This could get tricky with mixed-precision code. Probably needs some +# # special AbstractInteger infinity type (but custom AbstractInteger types +# # may cause a lot of method invalidations, which is why Static.StaticInteger +# # is not an AbstractInteger). +# @inline Base.minimum(::RealValues) = static(typemax(Int64)) +# @inline Base.maximum(::RealValues) = static(typemin(Int64)) + +""" + const ℤ = IntegerValues() + +The set of all integers, see [`MeasureBase.IntegerValues`](@ref). +""" +const ℤ = IntegerValues() +export ℤ + +Base.show(io::IO, ::IntegerValues) = print(io, "ℤ") +Base.show(io::IO, ::MIME"text/plain", ::IntegerValues) = print(io, "MeasureBase.ℤ") + +""" + struct MeasureBase.BoundedInts{L,U} <: MeasureBase.ValueSet + +The integers from `lower` to `upper` (bounds may be infinite). + +Constructors: + +```julia +BoundedInts(lower, upper) +ℤ[lower:upper] +``` +""" +struct BoundedInts{L,U} <: ValueSet lower::L upper::U end -Base.in(x, b::BoundedInts) = x ∈ ℤ && b.lower ≤ x ≤ b.upper +@inline Base.in(x, b::BoundedInts) = x ∈ ℤ && b.lower <= x <= b.upper + +Base.isempty(b::BoundedInts) = b.lower > b.upper Base.minimum(b::BoundedInts) = b.lower Base.maximum(b::BoundedInts) = b.upper @@ -63,83 +232,207 @@ 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 + +Supertype for cartesian products of sets. +""" +abstract type AbstractCartSetProd <: ValueSet end + +""" + struct CartesianProduct <: AbstractCartSetProd + +A cartesian product over a collection of sets. -########################################################### -# ZeroSet +Constructor: -export ZeroSet +```julia +prodset = CartesianProduct(sets) +``` -struct ZeroSet{F,G} <: AbstractDomain - f::F - ∇f::G +`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 + +Represents a combination of two sets. -struct Sphere <: CodimOne end +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, α, β) + + +# Element size of the arrays in array-valued sets, for mspace_elsize: +@inline _valueset_elsize(::RealValues) = () +@inline _valueset_elsize(::IntegerValues) = () +@inline _valueset_elsize(::BoundedInts) = () +@inline _valueset_elsize(::IntervalSets.AbstractInterval) = () +@inline _valueset_elsize(s::CartesianPower) = pwr_size(s) +@inline _valueset_elsize(s) = NoMSpaceElementSize{typeof(s)}() + +@inline _valueset_flatsize(s::CartesianPower) = _cat_sizes(_valueset_flatsize(pwr_base(s)), pwr_size(s)) +@inline _valueset_flatsize(s) = _valueset_elsize(s) diff --git a/src/getdof.jl b/src/getdof.jl index dbce2202..fd9b4a20 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,90 @@ 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 +function _try_direct_dof(μ::AbstractMeasure, ::AbstractNoDOF) + μ_local = _some_localmeasure(μ) + # A local measure of the same type would recurse forever: + typeof(μ_local) === typeof(μ) && _try_local_dof(μ, NoDOF{typeof(μ)}()) + _try_local_dof(μ, some_dof(μ_local)) +end + +_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( @@ -61,12 +148,15 @@ struct NoArgCheck{MU,T} end Return `x` if `x` is a valid variate of `μ`, throw an `ArgumentError` if not, return `NoArgCheck{MU,T}()` if not check can be performed. + +Only the shape and type of `x` are checked, never its value: values +outside the support are valid arguments of densities and transports. """ 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..d938d033 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -5,8 +5,8 @@ using Reexport @reexport using MeasureBase using MeasureBase: basemeasure_depth, proxy, istrue -using MeasureBase: insupport, basemeasure_sequence, commonbase -using MeasureBase: transport_to, NoTransport +using MeasureBase: insupport, basemeasure_sequence +using MeasureBase: transport_to using DensityInterface: logdensityof using InverseFunctions: inverse @@ -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) @@ -91,9 +90,8 @@ function test_transport(ν, μ) @testset "transport_to $μ to $ν" begin x = rand(μ) - @test !(@inferred(transport_to(ν, μ)(x)) isa NoTransport) f = transport_to(ν, μ) - y = f(x) + y = @inferred f(x) @test structisapprox(@inferred(inverse(f)(y)), x) @test @inferred(with_logabsdet_jacobian(f, x)) isa Tuple{supertype(y),Real} @test @inferred(with_logabsdet_jacobian(inverse(f), y)) isa Tuple{supertype(x),Real} @@ -110,7 +108,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 +117,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..0807b6e6 100644 --- a/src/mass-interface.jl +++ b/src/mass-interface.jl @@ -22,7 +22,12 @@ for T in (:UnknownFiniteMass, :UnknownMass) @eval begin Base.:+(::$T, ::$T) = $T() Base.:*(::$T, ::$T) = $T() - Base.:^(::$T, k::Number) = isfinite(k) ? $T() : UnknownMass() + Base.:*(x::Real, ::$T) = isfinite(x) ? $T() : UnknownMass() + Base.:*(::$T, x::Real) = isfinite(x) ? $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 +70,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 @@ -102,10 +107,10 @@ Check whether `norm(x, p) == 1`. """ isnormalized(x, p::Real = 2) = isone(norm(x, p)) -isone(::AbstractUnknownMass) = false +Base.isone(::AbstractUnknownMass) = false function massof(m, s) - _massof(m, s, rootmeasure(m)) + _default_massof_impl(m, s, rootmeasure(m)) end """ @@ -116,4 +121,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..42fbb837 --- /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 ⋄ μ`, 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..0ed034e5 --- /dev/null +++ b/src/mspace.jl @@ -0,0 +1,127 @@ +""" + 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, `()` for scalar variates. + +The size is static where it is known statically. Returns +[`NoMSpaceElementSize{typeof(μ)}()`](@ref NoMSpaceElementSize) if the elements of the space +are not arrays of one common size, e.g. for structured variates or variates +whose size depends on the value, or if the size can not be determined +efficiently. + +See also [`MeasureBase.mspace_flatsize`](@ref). +""" +function mspace_elsize end +export mspace_elsize + +@inline mspace_elsize(μ::AbstractMeasure) = NoMSpaceElementSize{typeof(μ)}() + + +""" + MeasureBase.mspace_flatsize(μ) + +Return the size of the flat storage of a variate of `μ`, `()` for scalar +variates. + +Variates of powers of measures with array-valued variates are nested +arrays, their flat storage has the size of the inner arrays followed by +the size of the power. Returns [`NoMSpaceElementSize{typeof(μ)}()`](@ref NoMSpaceElementSize) +if the variates of `μ` have no flat storage of a common size. + +See also [`mspace_elsize`](@ref). +""" +function mspace_flatsize end + +@inline mspace_flatsize(μ::AbstractMeasure) = NoMSpaceElementSize{typeof(μ)}() + +@inline _cat_sizes(a::SizeLike, b::SizeLike) = canonical_size((size_dims(a)..., size_dims(b)...)) +@inline _cat_sizes(a::NoMSpaceElementSize, ::SizeLike) = a +@inline _cat_sizes(::SizeLike, b::NoMSpaceElementSize) = b +@inline _cat_sizes(a::NoMSpaceElementSize, ::NoMSpaceElementSize) = a + + +""" + 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(μ)) + +@inline _value_elsize(::Number) = () +@inline _value_elsize(x::AbstractArray) = maybestatic_size(x) +@inline _value_elsize(x) = NoMSpaceElementSize{typeof(x)}() + +@inline _value_flatsize(::Number) = () +@inline _value_flatsize(x::AbstractArray{<:Number}) = maybestatic_size(x) +@inline _value_flatsize(x) = NoMSpaceElementSize{typeof(x)}() + +@inline _scalar_or_unknown(::Tuple{}) = () +@inline _scalar_or_unknown(sz::NoMSpaceElementSize) = sz +@inline _scalar_or_unknown(sz) = NoMSpaceElementSize{typeof(sz)}() + + +""" + MeasureBase.mspace_flatsize(::Type{MU}) + +The flat variate size of measures of type `MU`, if it is determined by the +type alone, e.g. `()` for measures with scalar variates. Returns +`NoMSpaceElementSize{MU}()` otherwise. + +Composite measures use it to determine the flat size of their variates +without inspecting each component. +""" +@inline mspace_flatsize(::Type{MU}) where {MU} = NoMSpaceElementSize{MU}() + + +""" + MeasureBase.mspace_ndims(::Type{MU}) + MeasureBase.mspace_ndims(μ) + +The number of dimensions of the flat variates of measures of type `MU`, +`0` for scalar variates, or a [`MeasureBase.NoMSpaceElementSize`](@ref) +if unknown. + +Batched kernels rely on it to tell the variate dimensions of a flat batch +from its batch dimensions. It follows from +[`MeasureBase.mspace_flatsize`](@ref) where that is known, measure types +with array variates of dynamic size declare it directly. +""" +function mspace_ndims end + +@inline mspace_ndims(::Type{MU}) where {MU} = _ndims_of_size(mspace_flatsize(MU), MU) +@inline mspace_ndims(μ::MU) where {MU} = _ndims_of_size(mspace_flatsize(μ), MU, mspace_ndims(MU)) + +@inline _ndims_of_size(sz::SizeLike, ::Type) = length(size_dims(sz)) +@inline _ndims_of_size(::NoMSpaceElementSize, ::Type{MU}) where {MU} = NoMSpaceElementSize{MU}() +@inline _ndims_of_size(sz::SizeLike, ::Type, ::Any) = length(size_dims(sz)) +@inline _ndims_of_size(::NoMSpaceElementSize, ::Type, n) = n + +@inline _add_ndims(n::Integer, k::Integer) = n + k +@inline _add_ndims(n::NoMSpaceElementSize, ::Integer) = n 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..9589c8ca 100644 --- a/src/primitive.jl +++ b/src/primitive.jl @@ -9,9 +9,9 @@ measures satisfy the following laws: basemeasure(μ::PrimitiveMeasure) = μ - logdensity_def(μ::PrimitiveMeasure, x) = 0.0 + logdensity_def(μ::PrimitiveMeasure, x) == 0 - logdensity_def(μ::M, ν::M, x) where {M<:PrimitiveMeasure} = 0.0 + logdensity_rel_def(μ::M, ν::M, x) where {M<:PrimitiveMeasure} == 0 """ abstract type PrimitiveMeasure <: AbstractMeasure end @@ -19,12 +19,11 @@ 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) = zero(_logd_numtype(x)) -logdensity_def(::PrimitiveMeasure, x) = static(0.0) +logdensity_def(::PrimitiveMeasure, x) = zero(_logd_numtype(x)) -logdensity_def(μ::M, ν::M, x) where {M<:PrimitiveMeasure} = 0.0 +logdensity_rel_def(μ::M, ν::M, x) where {M<:PrimitiveMeasure} = zero(_logd_numtype(x)) function Pretty.quoteof(μ::M) where {M<:PrimitiveMeasure} :($M()) diff --git a/src/primitives/counting.jl b/src/primitives/counting.jl index c61d0624..f4335526 100644 --- a/src/primitives/counting.jl +++ b/src/primitives/counting.jl @@ -4,6 +4,10 @@ export Counting, CountingBase struct CountingBase <: PrimitiveMeasure end +@inline mspace_elsize(::CountingBase) = () +@inline mspace_flatsize(::CountingBase) = () +@inline mspace_flatsize(::Type{CountingBase}) = () + insupport(::CountingBase, x) = true struct Counting{T} <: AbstractMeasure @@ -12,12 +16,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) @@ -25,6 +29,11 @@ basemeasure(::Counting) = CountingBase() Counting() = Counting(ℤ) +@inline mspace_elsize(μ::Counting) = _valueset_elsize(μ.support) +@inline mspace_flatsize(μ::Counting) = _valueset_flatsize(μ.support) +@inline mspace_flatsize(::Type{<:Counting{IntegerValues}}) = () +@inline mspace_flatsize(::Type{<:Counting{<:BoundedInts}}) = () + testvalue(::Type{T}, d::Counting) where {T} = testvalue(T, d.support) proxy(d::Counting) = restrict(in(d.support), CountingBase()) @@ -40,3 +49,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..dea1c087 100644 --- a/src/primitives/dirac.jl +++ b/src/primitives/dirac.jl @@ -9,6 +9,10 @@ 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.hash(a::Dirac, h::UInt) = hash(a.x, hash(:Dirac, h)) +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,17 +24,18 @@ 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), zero(_logd_numtype(x))) -logdensity_def(::Dirac, x::Real) = zero(float(typeof(x))) -logdensity_def(::Dirac, x) = 0.0 +logdensity_def(::Dirac, x::Number) = zero(float(typeof(x))) +logdensity_def(::Dirac, x) = zero(_logd_numtype(x)) -Base.rand(::Random.AbstractRNG, T::Type, μ::Dirac) = μ.x +@inline rand_impl(::GenContext, μ::Dirac) = μ.x +@inline batched_rand_impl(ctx::GenContext, μ::Dirac, sz::SizeLike) = _const_batch(ctx, μ.x, sz) export dirac @@ -40,7 +45,49 @@ insupport(d::Dirac, x) = x == d.x @inline getdof(::Dirac) = static(0) +@inline mspace_elsize(μ::Dirac) = _value_elsize(μ.x) +@inline mspace_flatsize(μ::Dirac) = _value_flatsize(μ.x) +@inline mspace_flatsize(::Type{<:Dirac{<:Number}}) = () + @propagate_inbounds function checked_arg(μ::Dirac, x) @boundscheck insupport(μ, x) || throw(ArgumentError("Invalid variate for measure")) x end + +# Dirac measures have no degrees of freedom: +@inline transport_to_std(::Type{S}, ::Dirac, x) where {S<:StdMeasure} = SVector{0,Bool}() +@inline transport_from_std(::Type{S}, μ::Dirac, z::AbstractVector) where {S<:StdMeasure} = μ.x +@inline transport_from_std_with_rest(::Type{S}, μ::Dirac, z::AbstractVector) where {S<:StdMeasure} = μ.x, z + +# Batched kernels cover Dirac measures with numbers and numeric arrays as +# flat variates, others have no declared variate rank: +const _FlatDirac = Dirac{<:Union{Number,AbstractArray{<:Number}}} + +@inline batched_transport_to_std(::Type{S}, ::Dirac, ::Number) where {S<:StdMeasure} = SVector{0,Bool}() +function batched_transport_to_std(::Type{S}, μ::_FlatDirac, X::AbstractArray) where {S<:StdMeasure} + n = length(_value_flatsize(μ.x)) + similar(X, Bool, (0, ntuple(i -> size(X, n + i), Val(ndims(X) - n))...)) +end + +function batched_transport_from_std(::Type{S}, μ::_FlatDirac, Z::AbstractArray) where {S<:StdMeasure} + _const_variates(μ.x, Z) +end +@inline _const_variates(x::Number, ::AbstractVector) = x +function _const_variates(x, Z::AbstractArray) + X = similar(Z, eltype(x), (size(x)..., Base.tail(size(Z))...)) + X .= x + return X +end + +@inline mspace_ndims(::Type{<:Dirac{<:AbstractArray{<:Number,N}}}) where {N} = N + +# Batches of array variates: all elements of a variate must match. +function batched_logdensityof_impl(μ::Dirac{<:AbstractArray{<:Number,N}}, X::AbstractArray) where {N} + matches = all_leading_dims(X .== μ.x, static(N)) + ifelse.(matches, zero(_logd_numtype(X)), _neg_inf_logd(X)) +end +function batched_logdensity_def(μ::Dirac{<:AbstractArray{<:Number}}, X::AbstractArray) + _zero_logd_batch(X, static(ndims(μ.x))) +end + +Adapt.adapt_structure(to, μ::Dirac) = Dirac(Adapt.adapt(to, μ.x)) diff --git a/src/primitives/lebesgue.jl b/src/primitives/lebesgue.jl index 3846eaf5..1604b112 100644 --- a/src/primitives/lebesgue.jl +++ b/src/primitives/lebesgue.jl @@ -4,6 +4,10 @@ export Lebesgue struct LebesgueBase <: PrimitiveMeasure end +@inline mspace_elsize(::LebesgueBase) = () +@inline mspace_flatsize(::LebesgueBase) = () +@inline mspace_flatsize(::Type{LebesgueBase}) = () + massof(::LebesgueBase, s::Interval) = width(s) testvalue(::LebesgueBase) = 0.0 @@ -12,13 +16,13 @@ insupport(::LebesgueBase, x) = true insupport(::LebesgueBase) = Returns(true) -logdensity_def(::LebesgueBase, ::CountingBase, x) = -Inf +logdensity_rel_def(::LebesgueBase, ::CountingBase, x) = _neg_inf_logd(x) -logdensity_def(::CountingBase, ::LebesgueBase, x) = Inf +logdensity_rel_def(::CountingBase, ::LebesgueBase, x) = -_neg_inf_logd(x) @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 +30,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 @@ -48,10 +52,15 @@ gentype(::Lebesgue) = Float64 Lebesgue() = Lebesgue(ℝ) +@inline mspace_elsize(μ::Lebesgue) = _valueset_elsize(μ.support) +@inline mspace_flatsize(μ::Lebesgue) = _valueset_flatsize(μ.support) +@inline mspace_flatsize(::Type{<:Lebesgue{RealValues}}) = () +@inline mspace_flatsize(::Type{<:Lebesgue{<:IntervalSets.AbstractInterval}}) = () + 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 +70,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..bfdd13c4 100644 --- a/src/proxies.jl +++ b/src/proxies.jl @@ -15,17 +15,25 @@ 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/rand.jl b/src/rand.jl index f92cb16a..ec0e4499 100644 --- a/src/rand.jl +++ b/src/rand.jl @@ -1,24 +1,154 @@ -import Base +# Random variate generation is parameterized by a `GenContext` carrying the +# random number generator, the numerical precision and the compute unit. +# Variates are generated as flat batches `(variate dims..., batch dims...)` +# on the compute unit, a single variate is a batch with zero batch +# dimensions. -Base.rand(d::AbstractMeasure) = rand(Random.GLOBAL_RNG, Float64, d) +""" + rand([rng::AbstractRNG], [T::Type{<:AbstractFloat}], μ::AbstractMeasure) + rand(ctx::GenContext, μ::AbstractMeasure) -Base.rand(T::Type, μ::AbstractMeasure) = rand(Random.GLOBAL_RNG, T, μ) +Generate a random variate of `μ`. -Base.rand(rng::AbstractRNG, d::AbstractMeasure) = rand(rng, Float64, d) +The generative context `ctx` (see `HeterogeneousComputing.GenContext`) +determines the random number generator, the numerical precision (`Float64` +by default) and the compute unit that array-valued variates are generated +on. Variates of powers of measures are generated as one flat batch of +variates of the base measure. -@inline Random.rand!(d::AbstractMeasure, args...) = rand!(GLOBAL_RNG, d, args...) +Measure types should specialize [`MeasureBase.batched_rand_impl`](@ref) +instead of `rand`. +""" +Base.rand(ctx::GenContext, μ::AbstractMeasure) = rand_impl(ctx, μ) -# TODO: Make this work -# function Base.rand(rng::AbstractRNG, ::Type{T}, d::AbstractMeasure) where {T} -# x = testvalue(d) -# rand!(d, x) -# end +Base.rand(μ::AbstractMeasure) = rand(GenContext{Float64}(), μ) +Base.rand(rng::AbstractRNG, μ::AbstractMeasure) = rand(GenContext{Float64}(rng), μ) +Base.rand(::Type{T}, μ::AbstractMeasure) where {T<:AbstractFloat} = rand(GenContext{T}(), μ) +Base.rand(rng::AbstractRNG, ::Type{T}, μ::AbstractMeasure) where {T<:AbstractFloat} = rand(GenContext{T}(rng), μ) -# struct ArraySlot{A,I} -# arr::A -# i::I -# end +@inline Random.rand!(d::AbstractMeasure, args...) = rand!(Random.default_rng(), d, args...) -# function rand!(rng::AbstractRNG, d::AbstractMeasure, x::ArraySlot) -# x.arr[x.i...] = rand(rng, d) -# end + +""" + MeasureBase.batched_rand_impl(ctx::GenContext, μ, sz::SizeLike) + +Generate a batch of random variates of `μ` of batch size `sz` in flat +form, an array `(variate dims..., sz...)`, or a single variate for +`sz == ()`. Batches of tuple and named tuple variates are tuples resp. +named tuples of batches. Fully static batch sizes give static arrays on +the CPU. + +This is the primary extension point for random variate generation. The +default implementation draws a batch of variates of the preferred +standard measure of `μ` and transports it to `μ` (see +[`MeasureBase.batched_transport_from_std`](@ref)), or generates the +variates one by one via [`MeasureBase.rand_impl`](@ref) if `μ` has no +standard transport. +""" +function batched_rand_impl end + +""" + MeasureBase.rand_impl(ctx::GenContext, μ) + +Generate one random variate of `μ` in the generative context `ctx`. + +The default implementation generates a batch with zero batch dimensions +via [`MeasureBase.batched_rand_impl`](@ref). Measure types with a more +direct way of generating single variates may specialize `rand_impl`. +""" +function rand_impl end + +# The marker tells the defaults whether `rand_impl` may be specialized +# for the measure (coming from the default `rand_impl` itself, it is not): +struct _NoRandImpl end +struct _MaybeRandImpl end + +@inline rand_impl(ctx::GenContext, μ) = _rand_default(ctx, μ, (), _NoRandImpl()) +@inline batched_rand_impl(ctx::GenContext, μ, sz::SizeLike) = _rand_default(ctx, μ, sz, _MaybeRandImpl()) + +@inline _rand_default(ctx::GenContext, μ, sz::SizeLike, m) = _rand_via_std(ctx, μ, sz, preferred_stdmeasure(μ), m) + +@inline function _rand_via_std(ctx::GenContext, μ, sz::SizeLike, ::Type{S}, m) where {S<:StdMeasure} + _rand_via_std_dof(ctx, μ, sz, S, fast_dof(μ), m) +end +@inline _rand_via_std(ctx::GenContext, μ, sz::SizeLike, ::Type{AnyStdMeasure}, m) = _rand_via_std(ctx, μ, sz, StdUniform, m) +@inline _rand_via_std(ctx::GenContext, μ, sz::SizeLike, ::Any, m) = _rand_pointwise(ctx, μ, sz, m) + +function _rand_via_std_dof(ctx::GenContext, μ, sz::SizeLike, ::Type{S}, n::IntegerLike, ::Any) where {S<:StdMeasure} + convert_realtype(get_precision(ctx), batched_transport_from_std(S, μ, _rand_std(ctx, S, (n, size_dims(sz)...)))) +end +@inline _rand_via_std_dof(ctx::GenContext, μ, sz::SizeLike, ::Type, ::Any, m) = _rand_pointwise(ctx, μ, sz, m) + +# Variates generated one by one, stacked into a flat batch: +@inline _rand_pointwise(ctx::GenContext, μ, sz::SizeLike, ::Any) = _batched_rand_pointwise(ctx, μ, sz) +@inline _rand_pointwise(ctx::GenContext, μ, ::Tuple{}, ::_MaybeRandImpl) = rand_impl(ctx, μ) +@noinline function _rand_pointwise(::GenContext, μ, ::Tuple{}, ::_NoRandImpl) + throw(ArgumentError("Random variate generation is not implemented for measures of type $(nameof(typeof(μ))), define MeasureBase.batched_rand_impl or MeasureBase.rand_impl")) +end + +function _batched_rand_pointwise(ctx::GenContext, μ, sz::SizeLike) + _stack_variates(map(_ -> rand_impl(ctx, μ), CartesianIndices(asnonstatic(sz)))) +end +@inline _batched_rand_pointwise(ctx::GenContext, μ, ::Tuple{}) = rand_impl(ctx, μ) + +@inline _stack_variates(xs::AbstractArray{<:Number}) = xs +@inline _stack_variates(xs::AbstractArray{<:AbstractArray}) = stacked(map(_stack_variates, xs)) +@inline _stack_variates(xs::AbstractArray{<:Union{Tuple,NamedTuple}}) = StructArrays.components(StructArray(xs)) + + +# Bulk draws of standard variates on the compute unit, single draws for +# zero batch dimensions: + +@inline _rand_std(ctx::GenContext, ::Type{S}, dims::SizeLike) where {S<:StdMeasure} = batched_rand_impl(ctx, S(), dims) + +@inline _rand_bulk(ctx::GenContext, sz::SizeLike) = _bulk_draw(rand, ctx, sz) +@inline _randn_bulk(ctx::GenContext, sz::SizeLike) = _bulk_draw(randn, ctx, sz) +@inline _randexp_bulk(ctx::GenContext, sz::SizeLike) = _randexp_bulk(ctx, sz, get_compute_unit(ctx)) +@inline _randexp_bulk(ctx::GenContext, sz::SizeLike, ::CPUnit) = _bulk_draw(randexp, ctx, sz) +# Not all compute units provide exponential draws, derive them from uniform draws then: +@inline _randexp_bulk(ctx::GenContext, sz::SizeLike, ::AbstractComputeUnit) = -log1p.(-_rand_bulk(ctx, sz)) + +# Fully static batch sizes draw static arrays on the CPU, so that variates +# of statically sized measures are allocation-free. Other compute units +# allocate their own arrays. +@inline _bulk_draw(f::F, ctx::GenContext, sz::SizeLike) where {F} = f(ctx, asnonstatic(sz)) +@inline _bulk_draw(f::F, ctx::GenContext, sz::StaticSizeLike) where {F} = + _bulk_draw(f, ctx, sz, get_compute_unit(ctx)) +@inline _bulk_draw(f::F, ctx::GenContext, sz::StaticSizeLike, ::AbstractComputeUnit) where {F} = + f(ctx, asnonstatic(sz)) +@inline _bulk_draw(f::F, ctx::GenContext, sz::StaticSizeLike, ::CPUnit) where {F} = + f(get_rng(ctx), staticarray_type(get_precision(ctx), canonical_size(sz))) + +@inline _rand_bulk(ctx::GenContext, ::Tuple{}) = rand(get_rng(ctx), get_precision(ctx)) +@inline _randn_bulk(ctx::GenContext, ::Tuple{}) = randn(get_rng(ctx), get_precision(ctx)) +@inline _randexp_bulk(ctx::GenContext, ::Tuple{}) = randexp(get_rng(ctx), get_precision(ctx)) + +# Test values use a constant RNG, which only draws single values: +const _ConstantContext = GenContext{<:AbstractFloat,<:AbstractComputeUnit,ConstantRNG} +@inline _rand_bulk(ctx::_ConstantContext, sz::SizeLike) = _const_bulk(ctx, rand(ConstantRNG(), get_precision(ctx)), sz) +@inline _randn_bulk(ctx::_ConstantContext, sz::SizeLike) = _const_bulk(ctx, randn(ConstantRNG(), get_precision(ctx)), sz) +@inline _randexp_bulk(ctx::_ConstantContext, sz::SizeLike) = _const_bulk(ctx, randexp(ConstantRNG(), get_precision(ctx)), sz) +@inline _rand_bulk(ctx::_ConstantContext, ::Tuple{}) = rand(ConstantRNG(), get_precision(ctx)) +@inline _randn_bulk(ctx::_ConstantContext, ::Tuple{}) = randn(ConstantRNG(), get_precision(ctx)) +@inline _randexp_bulk(ctx::_ConstantContext, ::Tuple{}) = randexp(ConstantRNG(), get_precision(ctx)) +@inline _const_bulk(ctx::GenContext, x, sz::SizeLike) = _const_bulk(ctx, x, sz, get_compute_unit(ctx)) +@inline _const_bulk(ctx::GenContext, x, sz::SizeLike, ::AbstractComputeUnit) = + fill!(allocate_array(ctx, typeof(x), asnonstatic(sz)), x) +@inline _const_bulk(ctx::GenContext, x, sz::StaticSizeLike, ::CPUnit) = maybestatic_fill(x, sz) + +# A mask over the batch dimensions, aligned with a flat batch of variates +# of rank `k`: +@inline _batch_mask(mask::Number, ::Any) = mask +@inline _batch_mask(mask::AbstractArray, ::StaticInteger{0}) = mask +@inline function _batch_mask(mask::AbstractArray, ::StaticInteger{K}) where {K} + reshape(mask, (ntuple(_ -> 1, Val(K))..., size(mask)...)) +end + +# A batch of copies of a constant variate: +function _const_batch(ctx::GenContext, x, sz::SizeLike) + X = allocate_array(ctx, eltype(x), (size(x)..., asnonstatic(sz)...)) + X .= x + return X +end +@inline _const_batch(ctx::GenContext, x::Number, sz::SizeLike) = _const_bulk(ctx, x, sz) +@inline _const_batch(::GenContext, x::Number, ::Tuple{}) = 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/stdconvert.jl b/src/standard/stdconvert.jl new file mode 100644 index 00000000..8ddf841d --- /dev/null +++ b/src/standard/stdconvert.jl @@ -0,0 +1,63 @@ +# Direct transports between standard measures, tail-accurate in both +# directions. Transports via StdUniform lose the upper tail of unbounded +# measures, since the uniform variate saturates at one. + +# Standard normal log-cdf and log-ccdf: +@inline _normlogcdf(z) = logerfc(-z * invsqrt2) - logtwo +@inline _normlogccdf(z) = logerfc(z * invsqrt2) - logtwo + +# Complementary standard normal cdf, accurate for large positive arguments: +@inline _normccdf(z) = erfc(z * invsqrt2) / 2 + +@inline function transport_def(::StdExponential, ::StdNormal, z) + ifelse(z < zero(z), -log1p(-Φ(z)), -log(_positive_prob(_normccdf(z)))) +end + +@inline function transport_def(::StdNormal, μ::StdExponential, x) + _nan_outside(μ, x, ifelse(x < oftype(x, logtwo), Φinv(_positive_prob(-expm1(-x))), -Φinv(min(_positive_prob(exp(-x)), one(x))))) +end + +@inline transport_def(::StdLogistic, ::StdNormal, z) = _normlogcdf(z) - _normlogccdf(z) + +@inline function transport_def(::StdNormal, ::StdLogistic, l) + ifelse(l < zero(l), Φinv(_positive_prob(logistic(l))), -Φinv(_positive_prob(logistic(-l)))) +end + +@inline transport_def(::StdLogistic, μ::StdExponential, x) = _nan_outside(μ, x, log(abs(expm1(-x))) + x) + +@inline transport_def(::StdExponential, ::StdLogistic, l) = log1pexp(l) + + +""" + MeasureBase.stdconvert(::Type{S}, ::Type{T}, x) + +Convert a variate `x` of the standard measure type `T` into a variate of +the standard measure type `S`, elementwise for arrays. +""" +function stdconvert end + +@inline stdconvert(::Type{S}, ::Type{S}, x) where {S<:StdMeasure} = x +@inline stdconvert(::Type{S}, ::Type{T}, x) where {S<:StdMeasure,T<:StdMeasure} = _StdConvert{S,T}()(x) + +struct _StdConvert{S,T} <: Function end +@inline (::_StdConvert{S,T})(x::Number) where {S,T} = transport_def(S(), T(), x) +@inline (k::_StdConvert)(x::AbstractArray) = broadcast(k, x) + + +""" + MeasureBase.StdPowerMeasure{MU<:StdMeasure,N} + +The type of an `N`-dimensional power of a standard measure of type `MU`. +""" +const StdPowerMeasure{MU<:StdMeasure,N} = PowerMeasure{MU,<:NTuple{N,OneToLike}} + +# Powers of standard measures transport directly, by elementwise conversion: +function transport_def(ν::StdPowerMeasure{NU}, μ::StdPowerMeasure{MU}, x) where {NU<:StdMeasure,MU<:StdMeasure} + _pwr_variate(ν, maybestatic_reshape(stdconvert(NU, MU, x), mspace_flatsize(ν))) +end + +function batched_transport_def(ν::StdPowerMeasure{NU}, μ::StdPowerMeasure{MU}, X::AbstractArray) where {NU<:StdMeasure,MU<:StdMeasure} + n_μ = length(mspace_flatsize(μ)) + batch_dims = ntuple(i -> size(X, n_μ + i), Val(ndims(X) - n_μ)) + reshape(stdconvert(NU, MU, X), (asnonstatic(mspace_flatsize(ν))..., batch_dims...)) +end diff --git a/src/standard/stdexponential.jl b/src/standard/stdexponential.jl index c985c224..124ce62a 100644 --- a/src/standard/stdexponential.jl +++ b/src/standard/stdexponential.jl @@ -1,18 +1,25 @@ +""" + StdExponential <: StdMeasure + +The standard exponential measure, the exponential distribution with unit +scale as a measure. +""" struct StdExponential <: StdMeasure end 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 @inline basemeasure(::StdExponential) = LebesgueBase() -@inline transport_def(::StdUniform, μ::StdExponential, x) = -expm1(-x) -@inline transport_def(::StdExponential, μ::StdUniform, x) = -log1p(-x) +@inline transport_def(::StdUniform, μ::StdExponential, x) = _nan_outside(μ, x, -expm1(-x)) +@inline transport_def(::StdExponential, μ::StdUniform, x) = _nan_outside(μ, x, -log1p(-_unit_interior(x))) -Base.rand(rng::Random.AbstractRNG, ::Type{T}, ::StdExponential) where {T} = randexp(rng, T) +@inline rand_impl(ctx::GenContext, ::StdExponential) = randexp(get_rng(ctx), get_precision(ctx)) +@inline batched_rand_impl(ctx::GenContext, ::StdExponential, sz::SizeLike) = _randexp_bulk(ctx, sz) diff --git a/src/standard/stdlogistic.jl b/src/standard/stdlogistic.jl index 58a1ba67..018beed3 100644 --- a/src/standard/stdlogistic.jl +++ b/src/standard/stdlogistic.jl @@ -1,20 +1,25 @@ +""" + StdLogistic <: StdMeasure + +The standard logistic measure, the logistic distribution with zero +location and unit scale as a measure. +""" struct StdLogistic <: StdMeasure end 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() @inline transport_def(::StdUniform, μ::StdLogistic, x) = logistic(x) -@inline transport_def(::StdLogistic, μ::StdUniform, p) = logit(p) +@inline transport_def(::StdLogistic, μ::StdUniform, p) = _nan_outside(μ, p, logit(_unit_interior(p))) -@inline function Base.rand(rng::Random.AbstractRNG, ::Type{T}, ::StdLogistic) where {T} - logit(rand(rng, T)) -end +@inline rand_impl(ctx::GenContext, ::StdLogistic) = logit(rand(get_rng(ctx), get_precision(ctx))) +@inline batched_rand_impl(ctx::GenContext, ::StdLogistic, sz::SizeLike) = logit.(_rand_bulk(ctx, sz)) smf(::StdLogistic, x) = logistic(x) smf(::StdLogistic) = logistic diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl index 833f280e..19110167 100644 --- a/src/standard/stdmeasure.jl +++ b/src/standard/stdmeasure.jl @@ -1,123 +1,119 @@ +""" + abstract type MeasureBase.StdMeasure <: AbstractMeasure + +Supertype of the standard measures that transports pivot on. + +Variates of standard measures are `Real` numbers. +""" abstract type StdMeasure <: AbstractMeasure end StdMeasure(::typeof(rand)) = StdUniform() StdMeasure(::typeof(randexp)) = StdExponential() StdMeasure(::typeof(randn)) = StdNormal() +@inline mspace_elsize(::StdMeasure) = () +@inline mspace_flatsize(::StdMeasure) = () +@inline mspace_flatsize(::Type{<:StdMeasure}) = () + @inline check_dof(::StdMeasure, ::StdMeasure) = nothing -@inline transport_def(::MU, μ::MU, x) where {MU<:StdMeasure} = x +# Standard measures have real scalar variates, checking them directly keeps +# the recursion over base measures (and its boxed arguments) out of the +# kernels: +@inline checked_arg(::StdMeasure, x::Real) = x -function transport_def(ν::StdMeasure, μ::PowerMeasure{<:StdMeasure}, x) - return transport_def(ν, μ.parent, only(x)) -end +@inline massof(::StdMeasure) = static(1.0) -function transport_def(ν::PowerMeasure{<:StdMeasure}, μ::StdMeasure, x) - return fill_with(transport_def(ν.parent, μ, only(x)), map(length, ν.axes)) -end +@inline transport_def(::MU, μ::MU, x) where {MU<:StdMeasure} = x -function transport_def( - ν::PowerMeasure{<:StdMeasure,<:NTuple{1,Base.OneTo}}, - μ::PowerMeasure{<:StdMeasure,<:NTuple{1,Base.OneTo}}, - x, -) - return transport_to(ν.parent, μ.parent).(x) -end +@inline transport_to_std(::Type{S}, ::S, x) where {S<:StdMeasure} = _std_identity(S, x) +@inline transport_from_std(::Type{S}, ::S, z) where {S<:StdMeasure} = _std_identity(S, z) -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)...) +# Only concrete standard measure types identify a transport partner: +@inline function _std_identity(::Type{S}, x) where {S<:StdMeasure} + isconcretetype(S) || _throw_abstract_std(S) + return x end +@noinline _throw_abstract_std(::Type{S}) where {S} = + throw(ArgumentError("$(S) is not a concrete standard measure type")) -# 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(μ)) +""" + struct MeasureBase.NoStdTransport{MU} -function transport_to(::Type{NU}, μ) where {NU<:StdMeasure} - transport_to(_std_measure_for(NU, μ), μ) -end +Indicates that measures of type `MU` can't be transported to or from a +standard measure. +""" +struct NoStdTransport{MU} end -function transport_to(ν, ::Type{MU}) where {MU<:StdMeasure} - transport_to(ν, _std_measure_for(MU, ν)) -end +""" + struct MeasureBase.AnyStdMeasure -# Transform between standard measures and Dirac: +Indicates that any standard measure serves as transport partner, e.g. +for measures with zero degrees of freedom. +""" +struct AnyStdMeasure end -@inline transport_def(ν::Dirac, ::PowerMeasure{<:StdMeasure}, ::Any) = ν.x +const _StdTransportPartner = Union{Type{<:StdMeasure},Type{AnyStdMeasure},Type{<:NoStdTransport}} -@inline function transport_def(ν::PowerMeasure{<:StdMeasure}, ::Dirac, ::Any) - Zeros{Bool}(map(_ -> 0, ν.axes)) -end +""" + MeasureBase.preferred_stdmeasure(μ)::Type + MeasureBase.preferred_stdmeasure(::Type{MU})::Type -# Helpers for product transforms and similar: +The type of standard measure that variates of `μ` are transported to and +from by default. -struct _TransportToStd{NU<:StdMeasure} <: Function end -_TransportToStd{NU}(μ, x) where {NU} = transport_to(NU()^getdof(μ), μ)(x) +Returns `MeasureBase.AnyStdMeasure` if any standard measure serves and +`MeasureBase.NoStdTransport{MU}` if measures of type `MU` have no +standard-measure transport. Composite measures combine the preferences of +their components via [`MeasureBase.promote_stdmeasure`](@ref). -struct _TransportFromStd{MU<:StdMeasure} <: Function end -_TransportFromStd{MU}(ν, x) where {MU} = transport_to(ν, MU()^getdof(ν))(x) +Measure types that support transport to and from standard measures should +specialize the type-based method. +""" +function preferred_stdmeasure end -function _tuple_transport_def( - ν::PowerMeasure{NU}, - μs::Tuple, - xs::Tuple, -) where {NU<:StdMeasure} - reshape(vcat(map(_TransportToStd{NU}, μs, xs)...), ν.axes) -end +@inline preferred_stdmeasure(μ) = preferred_stdmeasure(typeof(μ)) +@inline preferred_stdmeasure(::Type{MU}) where {MU} = NoStdTransport{MU} -function transport_def( - ν::PowerMeasure{NU}, - μ::ProductMeasure{<:Tuple}, - x, -) where {NU<:StdMeasure} - _tuple_transport_def(ν, marginals(μ), x) +@inline function preferred_stdmeasure(::Type{MU}) where {MU<:StdMeasure} + isconcretetype(MU) ? MU : NoStdTransport{MU} end -function transport_def( - ν::PowerMeasure{NU}, - μ::ProductMeasure{<:NamedTuple{names}}, - x, -) where {NU<:StdMeasure,names} - _tuple_transport_def(ν, values(marginals(μ)), values(x)) -end +""" + MeasureBase.promote_stdmeasure(A::Type, B::Type)::Type -@inline _offset_cumsum(s, x, y, rest...) = (s, _offset_cumsum(s + x, y, rest...)...) -@inline _offset_cumsum(s, x) = (s,) -@inline _offset_cumsum(s) = () +Combine two results of [`MeasureBase.preferred_stdmeasure`](@ref) into +the preferred standard measure type of a measure composed of both. -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 +Standard measure types promote to the one with the wider range of +values that remain distinguishable in floating point arithmetic: +`StdUniform` promotes to any other standard measure type, +`StdExponential` to `StdLogistic` and `StdNormal`, and `StdLogistic` +to `StdNormal`. +""" +function promote_stdmeasure 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) +@inline function promote_stdmeasure(::Type{A}, ::Type{B}) where {A<:StdMeasure,B<:StdMeasure} + _stdmeasure_rank(A) >= _stdmeasure_rank(B) ? A : B end -function transport_def( - ν::ProductMeasure{<:Tuple}, - μ::PowerMeasure{MU}, - x, -) where {MU<:StdMeasure} - _tuple_transport_def(marginals(ν), μ, x) +@inline promote_stdmeasure(::Type{AnyStdMeasure}, ::Type{B}) where {B} = B +@inline promote_stdmeasure(::Type{A}, ::Type{AnyStdMeasure}) where {A} = A +@inline promote_stdmeasure(::Type{AnyStdMeasure}, ::Type{AnyStdMeasure}) = AnyStdMeasure +@inline promote_stdmeasure(::Type{A}, ::Type{B}) where {A<:NoStdTransport,B} = A +@inline promote_stdmeasure(::Type{A}, ::Type{B}) where {A,B<:NoStdTransport} = B +@inline promote_stdmeasure(::Type{A}, ::Type{B}) where {A<:NoStdTransport,B<:NoStdTransport} = A +@inline promote_stdmeasure(::Type{A}, ::Type{AnyStdMeasure}) where {A<:NoStdTransport} = A +@inline promote_stdmeasure(::Type{AnyStdMeasure}, ::Type{B}) where {B<:NoStdTransport} = B + +@inline promote_stdmeasure(::Type{A}) where {A} = A +@inline function promote_stdmeasure(::Type{A}, ::Type{B}, Cs::Vararg{Type,N}) where {A,B,N} + promote_stdmeasure(promote_stdmeasure(A, B), Cs...) end -function transport_def( - ν::ProductMeasure{<:NamedTuple{names}}, - μ::PowerMeasure{MU}, - x, -) where {MU<:StdMeasure,names} - NamedTuple{names}(_tuple_transport_def(values(marginals(ν)), μ, x)) -end +@inline _stdmeasure_rank(::Type{<:StdMeasure}) = 0 + +@inline batched_transport_to_std(::Type{S}, ::S, X::AbstractArray) where {S<:StdMeasure} = _as_stdstream_batch(X) +@inline batched_transport_from_std(::Type{S}, ::S, Z::AbstractArray) where {S<:StdMeasure} = _drop_stdstream_dim(Z) +@inline batched_transport_from_std(::Type{S}, ::S, z::AbstractVector) where {S<:StdMeasure} = z[begin] diff --git a/src/standard/stdnormal.jl b/src/standard/stdnormal.jl index 057a8629..a94d92e4 100644 --- a/src/standard/stdnormal.jl +++ b/src/standard/stdnormal.jl @@ -1,20 +1,27 @@ -using SpecialFunctions: erfc, erfcinv -using IrrationalConstants: invsqrt2, log2π +using SpecialFunctions: erfc, erfcinv, logerfc +using IrrationalConstants: invsqrt2, log2π, logtwo +""" + StdNormal <: StdMeasure + +The standard normal measure, the normal distribution with zero mean and +unit variance as a measure. +""" struct StdNormal <: StdMeasure end 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()) @inline getdof(::StdNormal) = static(1) -@inline Base.rand(rng::Random.AbstractRNG, ::Type{T}, ::StdNormal) where {T} = randn(rng, T) +@inline rand_impl(ctx::GenContext, ::StdNormal) = randn(get_rng(ctx), get_precision(ctx)) +@inline batched_rand_impl(ctx::GenContext, ::StdNormal, sz::SizeLike) = _randn_bulk(ctx, sz) Φ(z) = erfc(-z * invsqrt2) / 2 Φinv(p) = -erfcinv(2 * p) * sqrt2 @@ -23,10 +30,10 @@ InverseFunctions.inverse(::typeof(Φ)) = Φinv InverseFunctions.inverse(::typeof(Φinv)) = Φ smf(::StdNormal, x) = Φ(x) -invsmf(::StdNormal, p) = Φinv(p) +invsmf(::StdNormal, p) = _nan_outside(StdUniform(), p, Φinv(_unit_interior(p))) smf(::StdNormal) = Φ invsmf(::StdNormal) = Φinv -transport_def(::StdNormal, ::StdUniform, p) = Φinv(p) +transport_def(::StdNormal, μ::StdUniform, p) = _nan_outside(μ, p, Φinv(_unit_interior(p))) transport_def(::StdUniform, ::StdNormal, x) = Φ(x) diff --git a/src/standard/stdtraits.jl b/src/standard/stdtraits.jl new file mode 100644 index 00000000..d9b7a235 --- /dev/null +++ b/src/standard/stdtraits.jl @@ -0,0 +1,40 @@ +# Standard measure preferences of the combinators and standard measure ranks: + +@inline _stdmeasure_rank(::Type{StdUniform}) = 1 +@inline _stdmeasure_rank(::Type{StdExponential}) = 2 +@inline _stdmeasure_rank(::Type{StdLogistic}) = 3 +@inline _stdmeasure_rank(::Type{StdNormal}) = 4 + + +@inline preferred_stdmeasure(::Type{<:PowerMeasure{M}}) where {M} = preferred_stdmeasure(M) +@inline preferred_stdmeasure(::Type{<:WeightedMeasure{<:Any,M}}) where {M} = preferred_stdmeasure(M) +# Transports of the base don't transport the restricted measure: +@inline preferred_stdmeasure(::Type{MU}) where {MU<:RestrictedMeasure} = NoStdTransport{MU} +@inline preferred_stdmeasure(::Type{<:PushforwardMeasure{<:Any,<:Any,M}}) where {M} = preferred_stdmeasure(M) +@inline preferred_stdmeasure(::Type{<:Dirac}) = AnyStdMeasure + +@inline preferred_stdmeasure(::Type{<:ProductMeasure{M}}) where {M<:AbstractArray} = preferred_stdmeasure(eltype(M)) + +# Arrays of marginals of mixed types combine their preferences at run time: +@inline function preferred_stdmeasure(μ::ProductMeasure{<:AbstractArray}) + _array_product_stdmeasure(μ, preferred_stdmeasure(typeof(μ))) +end +@inline _array_product_stdmeasure(μ, S::Type) = S +function _array_product_stdmeasure(μ::ProductMeasure{<:AbstractArray{M}}, ::Type{NoStdTransport{MU}}) where {M,MU} + isconcretetype(M) && return NoStdTransport{MU} + mapreduce(preferred_stdmeasure, promote_stdmeasure, marginals(μ); init = AnyStdMeasure) +end + +@inline function preferred_stdmeasure(::Type{<:ProductMeasure{M}}) where {M<:Tuple} + _promote_stdmeasure_oftypes(M) +end + +@inline function preferred_stdmeasure(::Type{<:ProductMeasure{NamedTuple{names,M}}}) where {names,M<:Tuple} + _promote_stdmeasure_oftypes(M) +end + +@inline _promote_stdmeasure_oftypes(::Type{Tuple{}}) = AnyStdMeasure +@generated function _promote_stdmeasure_oftypes(::Type{M}) where {M<:Tuple} + args = [:(preferred_stdmeasure($T)) for T in M.parameters] + :(promote_stdmeasure($(args...))) +end diff --git a/src/standard/stduniform.jl b/src/standard/stduniform.jl index 7bbe15ed..5971cd0f 100644 --- a/src/standard/stduniform.jl +++ b/src/standard/stduniform.jl @@ -1,24 +1,28 @@ +""" + StdUniform <: StdMeasure + +The standard uniform measure on the unit interval, the uniform +distribution on `[0, 1]` as a measure. +""" 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) @inline basemeasure(::StdUniform) = LebesgueBase() -Base.rand(rng::Random.AbstractRNG, ::Type{T}, ::StdUniform) where {T} = rand(rng, T) +@inline rand_impl(ctx::GenContext, ::StdUniform) = rand(get_rng(ctx), get_precision(ctx)) +@inline batched_rand_impl(ctx::GenContext, ::StdUniform, sz::SizeLike) = _rand_bulk(ctx, sz) -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)) -function invsmf(::StdUniform, p) - @assert zero(p) ≤ p ≤ one(p) - p -end +invsmf(d::StdUniform, p) = _nan_outside(d, p, p) 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-batched.jl b/src/transport-batched.jl new file mode 100644 index 00000000..9d0fa2e3 --- /dev/null +++ b/src/transport-batched.jl @@ -0,0 +1,275 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +# Batched-first transport over flat batches of variates `(variate dims..., +# batch dims...)`, zero batch dims meaning a single variate. Streams of +# standard variates are batches `(dof, batch dims...)`, consumed along their +# first dimension. Kernels know the variate rank of their measure (see +# `mspace_ndims`), structural measures implement them once in terms of the +# kernels of their components. + +""" + MeasureBase.batched_transport_to_std(::Type{S}, μ, X) + +Transport the flat batch `X` of variates of `μ` to a batch `(getdof(μ), +batch dims...)` of variates of the standard measure type `S`. `X` may be a +single variate, the result is a vector then. + +The default implementation broadcasts the point transport +[`MeasureBase.transport_to_std`](@ref) for measures with scalar variates +and maps it over the variate slices of `X` (in a host loop) for measures +with array variates of a declared number of dimensions (see +[`MeasureBase.mspace_ndims`](@ref)). Measure types with array variates +should implement `batched_transport_to_std` directly. +""" +function batched_transport_to_std end + +@inline function batched_transport_to_std(::Type{S}, μ, X) where {S<:StdMeasure} + _batched_to_std(S, μ, X, _static_ndims(μ)) +end + +@inline function _batched_to_std(::Type{S}, μ, X, ::StaticInteger{0}) where {S} + _as_stdstream_batch(broadcast(Base.Fix1(_ToStd{S}(), μ), X)) +end +@inline function _batched_to_std(::Type{S}, μ, X::AbstractArray, ::StaticInteger{0}) where {S} + _as_stdstream_batch(broadcast(Base.Fix1(_ToStd{S}(), μ), X)) +end +@inline function _batched_to_std(::Type{S}, μ, X::AbstractArray, ::StaticInteger{K}) where {S,K} + _to_std_slices(S, μ, X, Val(K)) +end +@noinline function _batched_to_std(::Type{S}, μ, X, ::NoMSpaceElementSize) where {S} + throw(ArgumentError("Batched transport requires MeasureBase.mspace_ndims to be declared for measures of type $(nameof(typeof(μ))) or MeasureBase.batched_transport_to_std to be implemented")) +end + +@inline _to_std_slices(::Type{S}, μ, X::AbstractArray{<:Any,K}, ::Val{K}) where {S,K} = _as_stdstream(transport_to_std(S, μ, X)) +@inline function _to_std_slices(::Type{S}, μ, X::AbstractArray, ::Val{K}) where {S,K} + stacked(map(Base.Fix1(_ToStd{S}(), μ), sliced(X, Val(K)))) +end + +# A flat batch of variates of `μ` as a batch of streams, the variate +# dimensions merged into the first dimension. Tuples of batches (tuple +# products and their powers) interleave the rows of their components +# variate by variate. +@inline _as_stream_batch(X, μ) = _as_stream_batch(X, _static_ndims(μ)) +@inline function _as_stream_batch(X::AbstractArray, ::StaticInteger{K}) where {K} + ndims(X) >= K || _throw_size_mismatch() + merge_leading_dims(X, static(K)) +end +@inline _as_stream_batch(x::Number, ::StaticInteger{0}) = SVector(x) +@noinline function _as_stream_batch(X, ::NoMSpaceElementSize) + throw(ArgumentError("Concatenating batches of variates requires MeasureBase.mspace_ndims to be declared for the measures involved")) +end +@inline function _as_stream_batch(X::Union{Tuple,NamedTuple}, μ::ProductMeasure) + vcat(map(_as_stream_batch, values(X), values(marginals(μ)))...) +end +function _as_stream_batch(X::Union{Tuple,NamedTuple}, μ::PowerMeasure) + ν, _ = _pwr_unwrap(μ) + n_pwr = length(_pwr_dims(μ)) + n = prod(asnonstatic(_pwr_dims(μ))) + parts = map(values(X), values(marginals(ν))) do Xi, m + A = _as_stream_batch(Xi, m) + reshape(A, (size(A, 1), n, ntuple(i -> size(A, 1 + n_pwr + i), Val(ndims(A) - 1 - n_pwr))...)) + end + merge_leading_dims(vcat(parts...), static(2)) +end + +# The standard variates of a single variate must form a vector: +@inline _single_std(z::AbstractVector) = z +@noinline function _single_std(z) + throw(ArgumentError("Transport of a single variate resulted in a batch of standard variates, the variate doesn't fit the measure")) +end + +@noinline function _throw_std_length_mismatch() + throw(ArgumentError("Length of standard variates doesn't match the degrees of freedom of the measure")) +end + + +""" + MeasureBase.batched_transport_from_std(::Type{S}, μ, Z::AbstractArray) + +Transport the batch `Z` of variates of the standard measure type `S`, of +size `(getdof(μ), batch dims...)`, to a flat batch of variates of `μ`. A +single stream `Z` yields a single variate. + +The default implementation broadcasts the point transport +[`MeasureBase.transport_from_std`](@ref) for measures with scalar +variates and maps it over the columns of `Z` (in a host loop) for measures +with array variates of a declared number of dimensions. +""" +function batched_transport_from_std end + +@inline function batched_transport_from_std(::Type{S}, μ, Z::AbstractArray) where {S<:StdMeasure} + _batched_from_std(S, μ, Z, _static_ndims(μ)) +end + +@inline _batched_from_std(::Type{S}, μ, Z::AbstractArray, ::StaticInteger{0}) where {S} = _from_std_scalar(S, μ, Z) +@inline _batched_from_std(::Type{S}, μ, Z::AbstractArray, ::StaticInteger{K}) where {S,K} = _from_std_columns(S, μ, Z) +@noinline function _batched_from_std(::Type{S}, μ, ::AbstractArray, ::NoMSpaceElementSize) where {S} + throw(ArgumentError("Batched transport requires MeasureBase.mspace_ndims to be declared for measures of type $(nameof(typeof(μ))) or MeasureBase.batched_transport_from_std to be implemented")) +end + +@inline _from_std_scalar(::Type{S}, μ, z::AbstractVector) where {S} = transport_from_std(S, μ, z[begin]) +@inline function _from_std_scalar(::Type{S}, μ, Z::AbstractArray) where {S} + broadcast(Base.Fix1(_FromStd{S}(), μ), _drop_stdstream_dim(Z)) +end +@inline _from_std_columns(::Type{S}, μ, z::AbstractVector) where {S} = transport_from_std(S, μ, z) +@inline function _from_std_columns(::Type{S}, μ, Z::AbstractArray) where {S} + stacked(map(Base.Fix1(_FromStd{S}(), μ), sliced(Z, Val(1)))) +end + + +""" + MeasureBase.batched_transport_to_std_with_rest(::Type{S}, μ, X::AbstractArray, sz::SizeLike) + +Consume variates of `μ` from the batch `X` of flat vector streams (first +dimension along the streams, further dimensions are batch dimensions), a +batch of variates of size `sz` per stream, and transport them to the +standard measure type `S`. + +Returns a tuple `(Z, X_rest)` of the standard variates as a batch +`(getdof(μ) * prod(sz), batch dims...)` and the unconsumed rest of the +streams. The default implementation consumes variates of the size given by +[`MeasureBase.mspace_flatsize`](@ref) or +[`MeasureBase.some_mspace_elsize`](@ref), a single stream with `sz == ()` +goes through [`MeasureBase.transport_to_std_with_rest`](@ref). Measures +whose variates are composed of the variates of other measures implement +`batched_transport_to_std_with_rest` in terms of their components. +""" +function batched_transport_to_std_with_rest end + +function batched_transport_to_std_with_rest(::Type{S}, μ, X::AbstractArray, sz::SizeLike) where {S<:StdMeasure} + _to_std_with_rest_default(S, μ, X, sz) +end + +function _to_std_with_rest_default(::Type{S}, μ, x::AbstractVector, ::Tuple{}) where {S} + z, _, x_rest = transport_to_std_with_rest(S, μ, x) + return z, x_rest +end +function _to_std_with_rest_default(::Type{S}, μ, X::AbstractArray, sz::SizeLike) where {S} + vsz = _stream_consume_size(μ) + X_μ, X_rest = _batched_consume(X, vsz, sz) + Z = batched_transport_to_std(S, μ, _consumed_variates(X_μ, vsz)) + return _merge_multiplicity(Z, sz), X_rest +end +@inline _consumed_variates(X_μ::AbstractArray, ::Tuple{}) = _drop_stdstream_dim(X_μ) +@inline _consumed_variates(X_μ::AbstractArray, ::SizeLike) = X_μ + +# Standard variates of `prod(sz)` variates per stream, `(dof, sz..., batch +# dims...)`, as one stream chunk `(dof * prod(sz), batch dims...)`, and +# back: +@inline _merge_multiplicity(Z::AbstractArray, sz::SizeLike) = merge_leading_dims(Z, static(1) + maybestatic_length(size_dims(sz))) +@inline _split_multiplicity(Z::AbstractArray, ::Tuple{}, n) = Z +@inline function _split_multiplicity(Z::AbstractArray, sz::SizeLike, n) + maybestatic_reshape(Z, (n, size_dims(sz)..., Base.tail(_batch_dims(Z))...)) +end + + +""" + MeasureBase.batched_transport_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, sz::SizeLike) + +Consume standard variates of type `S` for a batch of variates of size +`sz` per stream from the batch `Z` of streams of standard variates (first +dimension along the streams) and transport them to `μ`. + +Returns a tuple `(X, Z_rest)` of the flat batch `(flat variate dims..., +sz..., batch dims...)` of variates of `μ` and the unconsumed rest of the +streams. The default implementation consumes [`MeasureBase.fast_dof(μ)`](@ref) +entries per variate, a single stream with `sz == ()` goes through +[`MeasureBase.transport_from_std_with_rest`](@ref). Measures whose +variates are composed of the variates of other measures implement +`batched_transport_from_std_with_rest` in terms of their components. +""" +function batched_transport_from_std_with_rest end + +function batched_transport_from_std_with_rest(::Type{S}, μ, Z::AbstractArray, sz::SizeLike) where {S<:StdMeasure} + _from_std_with_rest_default(S, μ, Z, sz) +end + +_from_std_with_rest_default(::Type{S}, μ, z::AbstractVector, ::Tuple{}) where {S} = transport_from_std_with_rest(S, μ, z) +function _from_std_with_rest_default(::Type{S}, μ, Z::AbstractArray, sz::SizeLike) where {S} + _batched_from_std_bydof(S, μ, Z, sz, fast_dof(μ)) +end + +function _batched_from_std_bydof(::Type{S}, μ, Z::AbstractArray, sz::SizeLike, n::IntegerLike) where {S} + Z_μ, Z_rest = _batched_split(Z, _chunk_rows(n, sz)) + return batched_transport_from_std(S, μ, _split_multiplicity(Z_μ, sz, n)), Z_rest +end +@noinline function _batched_from_std_bydof(::Type{S}, μ, ::AbstractArray, ::SizeLike, ::AbstractNoDOF) where {S} + throw(ArgumentError("Batched transport from standard measures requires measures of type $(nameof(typeof(μ))) to have fast degrees of freedom or to implement MeasureBase.batched_transport_from_std_with_rest")) +end + + +""" + MeasureBase.batched_transport_def(ν, μ, X) + +Transport the flat batch `X` of variates of `μ` to a flat batch of +variates of `ν`, via the standard measure type the preferences of `ν` and +`μ` promote to. Specialize for pairs of measure types with a direct +batched transport. +""" +function batched_transport_def end + +function batched_transport_def(ν, μ, X) + S = _transport_pivot(ν, μ) + Z = batched_transport_to_std(S, μ, X) + Y, Z_rest = batched_transport_from_std_with_rest(S, ν, Z, ()) + if size(Z_rest, 1) != 0 + throw(ArgumentError("Degrees of freedom of source and target measure of a transport don't match")) + end + return Y +end + + +# Broadcasting a transport function over an array of variates with flat +# storage, or over the flat storage of a batch, transports the batch as a +# whole. Fused broadcast arguments are materialized first, static arrays +# are transported point by point. +function Broadcast.broadcasted(f::TransportFunction, X::AbstractArray) + _broadcast_transport(f, X, _flat_storage(X), _static_ndims(f.μ), _static_ndims(f.ν)) +end + +function Broadcast.broadcasted(f::TransportFunction, bc::Broadcast.Broadcasted) + Broadcast.broadcasted(f, Broadcast.materialize(bc)) +end + +# Static arrays of scalar variates are transported point by point: +Broadcast.broadcasted(f::TransportFunction, X::StaticArray) = _broadcast_static(f, X, _static_ndims(f.μ)) +_broadcast_static(f::TransportFunction, X::StaticArray, ::StaticInteger{0}) = map(_Pointwise(f), X) +function _broadcast_static(f::TransportFunction, X::StaticArray, k) + _broadcast_transport(f, X, X, k, _static_ndims(f.ν)) +end + +function _broadcast_transport(f::TransportFunction, X, X_flat::AbstractArray, ::StaticInteger, ::StaticInteger{K}) where {K} + Y_flat = batched_transport_def(f.ν, f.μ, X_flat) + return _batch_variates(Y_flat, f.ν, Val(K)) +end + +# Batches of tuple and named tuple variates are tuples of batches, the +# target layout follows from the target measure: +function _broadcast_transport(f::TransportFunction, X, X_flat::Union{Tuple,NamedTuple}, ::Any, ::Any) + _structured_variates(batched_transport_def(f.ν, f.μ, X_flat), f.ν) +end +function _broadcast_transport(f::TransportFunction{<:ProductMeasure{<:Union{Tuple,NamedTuple}}}, X, X_flat::AbstractArray, ::StaticInteger, ::NoMSpaceElementSize) + _structured_variates(batched_transport_def(f.ν, f.μ, X_flat), f.ν) +end +@inline _structured_variates(Y::Union{Tuple,NamedTuple}, ν) = _pwr_variate(ν, Y) +@inline _structured_variates(Y::AbstractArray, ν) = _batch_variates(Y, ν, Val(dynamic(_static_ndims(ν)))) + +_broadcast_transport(f::TransportFunction, X, ::Any, ::Any, ::Any) = map(_Pointwise(f), X) + +# Prevents re-entering the broadcast hook from `map` implementations that +# broadcast (e.g. GPU arrays): +struct _Pointwise{F} <: Function + f::F +end +@inline (p::_Pointwise)(x) = p.f(x) + +# The batch of variates in the layout of the target measure over the flat +# result, nested powers included, batches of tuple variates as struct +# arrays: +@inline _batch_variates(Y::AbstractArray, ν, ::Val{K}) where {K} = _nest_batch(Y, Val(K)) +@inline function _batch_variates(Y::AbstractArray, ν::PowerMeasure, ::Val) + sliced(_pwr_variate(ν, Y), Val(length(pwr_axes(ν)))) +end +@inline _nest_batch(Y::AbstractArray, ::Val{0}) = Y +@inline _nest_batch(Y::AbstractArray{<:Any,K}, ::Val{K}) where {K} = Y +@inline _nest_batch(Y::AbstractArray, ::Val{K}) where {K} = sliced(Y, Val(K)) diff --git a/src/transport.jl b/src/transport.jl index b0c8ed41..6061a21c 100644 --- a/src/transport.jl +++ b/src/transport.jl @@ -1,49 +1,3 @@ -""" - struct MeasureBase.NoTransportOrigin{NU} - -Indicates that no (default) pullback measure is available for measures of -type `NU`. - -See [`MeasureBase.transport_origin`](@ref). -""" -struct NoTransportOrigin{NU} end - -""" - MeasureBase.transport_origin(ν) - -Default measure to pullback to resp. pushforward from when transforming -between `ν` and another measure. -""" -function transport_origin end - -transport_origin(ν::NU) where {NU} = NoTransportOrigin{NU}() - -""" - MeasureBase.from_origin(ν, x) - -Push `x` from `MeasureBase.transport_origin(μ)` forward to `ν`. -""" -function from_origin end - -from_origin(ν::NU, ::Any) where {NU} = NoTransportOrigin{NU}() - -""" - MeasureBase.to_origin(ν, y) - -Pull `y` from `ν` back to `MeasureBase.transport_origin(ν)`. -""" -function to_origin end - -to_origin(ν::NU, ::Any) where {NU} = NoTransportOrigin{NU}() - -""" - struct MeasureBase.NoTransport{NU,MU} end - -Indicates that no transformation from a measure of type `MU` to a measure of -type `NU` could be found. -""" -struct NoTransport{NU,MU} end - """ f = transport_to(ν, μ) @@ -52,161 +6,70 @@ Generates a [measurable function](https://en.wikipedia.org/wiki/Measurable_funct a value `y = f(x)` distributed according to a measure `ν`. The [pushforward measure](https://en.wikipedia.org/wiki/Pushforward_measure) -from `μ` under `f` is is equivalent to `ν`. - -If terms of random values this implies that `f(rand(μ))` is equivalent to -`rand(ν)` (if `rand(μ)` and `rand(ν)` are supported). - -The resulting function `f` should support -`ChangesOfVariables.with_logabsdet_jacobian(f, x)` if mathematically well-defined, -so that densities of `ν` can be derived from densities of `μ` via `f` (using -appropriate base measures). - -Returns NoTransportOrigin{typeof(ν),typeof(μ)} if no transformation from -`μ` to `ν` can be found. - -To add transformation rules for a measure type `MyMeasure`, specialize - -* `MeasureBase.transport_def(ν::SomeStdMeasure, μ::CustomMeasure, x) = ...` -* `MeasureBase.transport_def(ν::MyMeasure, μ::SomeStdMeasure, x) = ...` - -and/or - -* `MeasureBase.transport_origin(ν::MyMeasure) = SomeMeasure(...)` -* `MeasureBase.from_origin(μ::MyMeasure, x) = y` -* `MeasureBase.to_origin(μ::MyMeasure, y) = x` - -and ensure `MeasureBase.getdof(μ::MyMeasure)` is defined correctly. - -A standard measure type like `StdUniform`, `StdExponential` or -`StdLogistic` may also be used as the source or target of the transform: +from `μ` under `f` is equivalent to `ν`, so `f(rand(μ))` is equivalent +to `rand(ν)`. `f` supports `InverseFunctions.inverse` and +`ChangesOfVariables.with_logabsdet_jacobian`. + +Measures are transported via standard measures: `x` is transported to the +standard measure type that the preferences of `ν` and `μ` promote to (see +[`MeasureBase.preferred_stdmeasure`](@ref)) and from there to `ν`. +Broadcasting `f` over an array of variates with flat storage, or over +the flat storage of a batch of variates (see +[`MeasureBase.mspace_ndims`](@ref)), transports the whole batch at once. A standard measure +type like `StdUniform` or `StdNormal` may also be used directly as the +source or target: ```julia -f_to_uniform(StdUniform, μ) -f_to_uniform(ν, StdUniform) +transport_to(StdNormal, μ) +transport_to(ν, StdNormal) ``` -Depending on [`getdof(μ)`](@ref) (resp. `ν`), an instance of the standard -distribution itself or a power of it (e.g. `StdUniform()` or -`StdUniform()^dof`) will be chosen as the transformation partner. +The transport partner is then an instance of the standard measure for +measures with scalar variates, and a power of it with as many elements as +the measure has degrees of freedom otherwise. + +# Extended help + +To support transport for a measure type, specialize +[`MeasureBase.transport_to_std`](@ref) and +[`MeasureBase.transport_from_std`](@ref) for its preferred standard measure +type, and declare [`MeasureBase.mspace_ndims`](@ref) for array variates. +Measure types with array variates should also implement the batched forms +[`MeasureBase.batched_transport_to_std`](@ref) and +[`MeasureBase.batched_transport_from_std`](@ref), which transport whole +batches of variates. Measures whose variates are composed of the variates +of other measures specialize the stream forms +[`MeasureBase.transport_to_std_with_rest`](@ref) and +[`MeasureBase.transport_from_std_with_rest`](@ref) instead (and their +batched forms). [`MeasureBase.transport_def`](@ref) may be specialized +for pairs of measure types with a direct transport. """ function transport_to end +export transport_to """ transport_to(ν, μ, x) -Transport `x` from the measure `μ` to the measure `ν` -""" -transport_to(ν, μ, x) = transport_to(ν, μ)(x) - -""" - transport_def(ν, μ, x) - -Transforms a value `x` distributed according to `μ` to a value `y` distributed -according to `ν`. - -If no specialized `transport_def(::MU, ::NU, ...)` is available then -the default implementation of`transport_def(ν, μ, x)` uses the following -strategy: +Transport `x` from the measure `μ` to the measure `ν`, equivalent to +`transport_to(ν, μ)(x)`. -* Evaluate [`transport_origin`](@ref) for μ and ν. Transform between - each and it's origin, if available, and use the origin(s) as intermediate - measures for another transformation. +# Extended help -* If all else fails, try to transform from μ to a standard multivariate - uniform measure and then to ν. +Variates of the right shape never throw: outside the support of `μ` the +result is `NaN` (elementwise for powers and products). Variates of the +wrong shape throw an `ArgumentError`. Transport implementations must not +throw outside the support, since the `NaN` masks evaluate both branches. -See [`transport_to`](@ref). +Finite inputs give finite results: on the floating-point grid the +endpoints of the unit interval stand for their nearest interior grid +points (uniform inputs are clamped into the open interval before +quantiles), and tail probabilities in log-space conversions never +underflow to zero. """ -function transport_def end - -function transport_def(ν, μ, x) - _transport_between_origins(ν, _origin_depth(ν), _origin_depth(μ), μ, x) -end - -@inline function _origin_depth(ν::NU) where {NU} - ν_0 = ν - Base.Cartesian.@nexprs 10 i -> begin # 10 is just some "big enough" number - ν_{i} = transport_origin(ν_{i - 1}) - if ν_{i} isa NoTransportOrigin - return static(i - 1) - end - end - return static(10) -end - -# If both both measures have no origin: -function _transport_between_origins(ν, ::StaticInteger{0}, ::StaticInteger{0}, μ, x) - _transport_with_intermediate(ν, _transport_intermediate(ν, μ), μ, x) -end - -@generated function _transport_between_origins( - ν, - ::StaticInteger{n_ν}, - ::StaticInteger{n_μ}, - μ, - x, -) where {n_ν,n_μ} - prog = quote - μ0 = μ - x0 = x - ν0 = ν - end - for i in 1:n_μ - μ_i = Symbol(:μ, i) - μ_last = Symbol(:μ, i - 1) - push!(prog.args, :($μ_i = transport_origin($μ_last))) - end - for i in 1:n_μ - x_i = Symbol(:x, i) - x_last = Symbol(:x, i - 1) - μ_last = Symbol(:μ, i - 1) - push!(prog.args, :($x_i = to_origin($μ_last, $x_last))) - end - for i in 1:(n_ν) - ν_i = Symbol(:ν, i) - ν_last = Symbol(:ν, i - 1) - push!(prog.args, :($ν_i = transport_origin($ν_last))) - end - μ_im = Symbol(:μ, n_μ) - x_im = Symbol(:x, n_μ) - ν_im = Symbol(:ν, n_ν) - y_im = Symbol(:y, n_ν) - push!(prog.args, :($y_im = transport_def($ν_im, $μ_im, $x_im))) - for i in (n_ν-1):-1:0 - y_i = Symbol(:y, i) - y_last = Symbol(:y, i + 1) - ν_last = Symbol(:ν, i) - push!(prog.args, :($y_i = from_origin($ν_last, $y_last))) - end - push!(prog.args, :(return y0)) - return prog -end - -@inline _transport_intermediate(ν, μ) = _transport_intermediate(getdof(ν), getdof(μ)) -@inline _transport_intermediate(::Integer, n_μ::Integer) = StdUniform()^n_μ -@inline _transport_intermediate(::StaticInteger{1}, ::StaticInteger{1}) = StdUniform() - -_call_transport_def(ν, μ, x) = transport_def(ν, μ, x) -_call_transport_def(::Any, ::Any, x::NoTransportOrigin) = x -_call_transport_def(::Any, ::Any, x::NoTransport) = x - -function _transport_with_intermediate(ν, m, μ, x) - z = _call_transport_def(m, μ, x) - y = _call_transport_def(ν, m, z) - return y -end - -# Prevent infinite recursion in case vartransform_intermediate doesn't change type: -@inline function _transport_with_intermediate(::NU, ::NU, ::MU, ::Any) where {NU,MU} - NoTransport{NU,MU}() -end -@inline function _transport_with_intermediate(::NU, ::MU, ::MU, ::Any) where {NU,MU} - NoTransport{NU,MU}() -end +transport_to(ν, μ, x) = transport_to(ν, μ)(x) """ - struct TransportFunction <: Function + struct MeasureBase.TransportFunction <: Function Transforms a variate from one measure to a variate of another. @@ -227,14 +90,15 @@ 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.μ end +Base.hash(f::TransportFunction, h::UInt) = hash(f.ν, hash(f.μ, hash(:TransportFunction, h))) Base.@propagate_inbounds function (f::TransportFunction)(x) - return _call_transport_def(f.ν, f.μ, checked_arg(f.μ, x)) + return transport_def(f.ν, f.μ, checked_arg(f.μ, x)) end @inline function InverseFunctions.inverse(f::TransportFunction{NU,MU}) where {NU,MU} @@ -243,11 +107,11 @@ end function ChangesOfVariables.with_logabsdet_jacobian(f::TransportFunction, x) y = f(x) - logpdf_src = logdensityof(f.μ, x) - logpdf_trg = logdensityof(f.ν, y) - ladj = logpdf_src - logpdf_trg - # If logpdf_src and logpdf_trg are -Inf setting lafj to zero is safe: - fixed_ladj = logpdf_src == logpdf_trg == -Inf ? zero(ladj) : ladj + logd_src = logdensityof(f.μ, x) + logd_trg = logdensityof(f.ν, y) + ladj = logd_src - logd_trg + # Both densities being -Inf leaves the Jacobian undefined, zero is a safe choice then: + fixed_ladj = ifelse(isneginf(logd_src) & isneginf(logd_trg), zero(ladj), ladj) return y, fixed_ladj end @@ -274,3 +138,214 @@ function Base.show(io::IO, f::TransportFunction) end Base.show(io::IO, M::MIME"text/plain", f::TransportFunction) = show(io, f) + + +""" + MeasureBase.transport_def(ν, μ, x) + +Transport a variate `x` of `μ` to a variate of `ν`. + +The default implementation transports `x` via the standard measure type the +preferences of `ν` and `μ` promote to. Specialize `transport_def` for pairs +of measure types with a direct transport. +""" +function transport_def end + +@inline transport_def(ν, μ, x) = _transport_via_std(_transport_pivot(ν, μ), ν, μ, x) + +function _transport_via_std(::Type{S}, ν, μ, x) where {S<:StdMeasure} + z = transport_to_std(S, μ, x) + y, z_rest = transport_from_std_with_rest(S, ν, _as_stdstream(z)) + if !isempty(z_rest) + throw(ArgumentError("Degrees of freedom of source and target measure of a transport don't match")) + end + return y +end + +@inline function _transport_pivot(ν, μ) + _concrete_pivot(promote_stdmeasure(preferred_stdmeasure(ν), preferred_stdmeasure(μ)), ν, μ) +end +@inline function _concrete_pivot(::Type{S}, ν, μ) where {S<:StdMeasure} + isconcretetype(S) || _throw_abstract_std(S) + return S +end +@inline _concrete_pivot(::Type{AnyStdMeasure}, ν, μ) = StdUniform +function _concrete_pivot(::Type{<:NoStdTransport{MU}}, ν, μ) where {MU} + throw(ArgumentError("No transport between measures of type $(nameof(typeof(ν))) and $(nameof(typeof(μ))), measures of type $(nameof(MU)) have no transport via standard measures")) +end + +# Standard variates of scalar-variate measures are scalars, streams of +# standard variates are vectors: +@inline _as_stdstream(z::AbstractVector) = z +@inline _as_stdstream(z::Number) = SVector(z) + + +""" + MeasureBase.transport_to_std(::Type{S}, μ, x) + +Transport a variate `x` of `μ` to a variate of the standard measure type +`S`: a number if the variates of `μ` are scalars, a flat vector of length +[`getdof(μ)`](@ref) otherwise. + +Measure types specialize `transport_to_std` for their preferred standard +measure type (see [`MeasureBase.preferred_stdmeasure`](@ref)), the generic +implementation converts between standard measure types. +""" +function transport_to_std end + +@inline function transport_to_std(::Type{S}, μ, x) where {S<:StdMeasure} + _to_std_via(S, preferred_stdmeasure(μ), μ, x) +end + +@inline function _to_std_via(::Type{S}, ::Type{T}, μ, x) where {S<:StdMeasure,T<:StdMeasure} + stdconvert(S, T, transport_to_std(T, μ, x)) +end +function _to_std_via(::Type{S}, ::Type{S}, μ, x) where {S<:StdMeasure} + throw(ArgumentError("Transport to $(nameof(S)) is not implemented for measures of type $(nameof(typeof(μ)))")) +end +function _to_std_via(::Type{S}, ::Type{AnyStdMeasure}, μ, x) where {S<:StdMeasure} + throw(ArgumentError("Transport to standard measures is not implemented for measures of type $(nameof(typeof(μ)))")) +end +function _to_std_via(::Type{S}, ::Type, μ, x) where {S<:StdMeasure} + throw(ArgumentError("Measures of type $(nameof(typeof(μ))) have no transport via standard measures")) +end + + +""" + MeasureBase.transport_from_std(::Type{S}, μ, z) + +Transport a variate `z` of the standard measure type `S` to a variate of +`μ`, the inverse of [`MeasureBase.transport_to_std`](@ref). +""" +function transport_from_std end + +@inline function transport_from_std(::Type{S}, μ, z) where {S<:StdMeasure} + _from_std_via(S, preferred_stdmeasure(μ), μ, z) +end + +@inline function _from_std_via(::Type{S}, ::Type{T}, μ, z) where {S<:StdMeasure,T<:StdMeasure} + transport_from_std(T, μ, stdconvert(T, S, z)) +end +function _from_std_via(::Type{S}, ::Type{S}, μ, z) where {S<:StdMeasure} + throw(ArgumentError("Transport from $(nameof(S)) is not implemented for measures of type $(nameof(typeof(μ)))")) +end +function _from_std_via(::Type{S}, ::Type{AnyStdMeasure}, μ, z) where {S<:StdMeasure} + throw(ArgumentError("Transport from standard measures is not implemented for measures of type $(nameof(typeof(μ)))")) +end +function _from_std_via(::Type{S}, ::Type, μ, z) where {S<:StdMeasure} + throw(ArgumentError("Measures of type $(nameof(typeof(μ))) have no transport via standard measures")) +end + + +""" + MeasureBase.transport_to_std_with_rest(::Type{S}, μ, x) + +Transport the variate of `μ` at the beginning of the stream `x` of +variate content to the standard measure type `S`. + +Returns a tuple `(z, x_μ, x_rest)` of the flat vector `z` of standard +variates, the variate `x_μ` of `μ` consumed from the stream and the +unconsumed rest of the stream. See +[`MeasureBase.logdensityof_with_rest`](@ref) for the stream conventions. + +The default implementation consumes a variate of the size given by +[`MeasureBase.mspace_flatsize`](@ref) or +[`MeasureBase.some_mspace_elsize`](@ref). Measure types whose variates are +composed of the variates of other measures implement +`transport_to_std_with_rest` instead of +[`MeasureBase.transport_to_std`](@ref). +""" +function transport_to_std_with_rest end + +function transport_to_std_with_rest(::Type{S}, μ, x::AbstractVector) where {S<:StdMeasure} + x_μ, x_rest = _consume_from_stream(x, _stream_consume_size(μ)) + return _as_stdstream(transport_to_std(S, μ, x_μ)), x_μ, x_rest +end + +function transport_to_std_with_rest(::Type{S}, μ, x::NamedTuple) where {S<:StdMeasure} + x_μ, x_rest = _split_after(x, Val(_mspace_names(μ))) + return _as_stdstream(transport_to_std(S, μ, x_μ)), x_μ, x_rest +end + + +""" + MeasureBase.transport_from_std_with_rest(::Type{S}, μ, z) + +Transport the beginning of the flat stream `z` of standard variates of type +`S` to a variate of `μ`, consuming as many entries as `μ` requires. + +Returns a tuple `(x, z_rest)` of the variate `x` and the unconsumed rest of +the stream. Measure types whose degrees of freedom depend on variate values +implement `transport_from_std_with_rest` instead of +[`MeasureBase.transport_from_std`](@ref). +""" +function transport_from_std_with_rest end + +function transport_from_std_with_rest(::Type{S}, μ, z::AbstractVector) where {S<:StdMeasure} + _from_std_with_rest_bydof(S, μ, z, fast_dof(μ)) +end + +function _from_std_with_rest_bydof(::Type{S}, μ, z::AbstractVector, n::IntegerLike) where {S} + if maybestatic_length(z) < n + throw(ArgumentError("Stream of standard variates too short during transport")) + end + z_μ, z_rest = split_at(z, n) + return transport_from_std(S, μ, _chunk_as_variate(μ, z_μ)), z_rest +end + +function _from_std_with_rest_bydof(::Type{S}, μ, z::AbstractVector, ::AbstractNoDOF) where {S} + throw(ArgumentError("Transport from standard measures requires measures of type $(nameof(typeof(μ))) to implement MeasureBase.transport_from_std_with_rest")) +end + +# Scalar-variate measures take their standard variate as a number: +@inline _chunk_as_variate(μ, z) = _chunk_as_variate(z, _static_ndims(μ)) +@inline _chunk_as_variate(z::AbstractVector, ::StaticInteger{0}) = z[begin] +@inline _chunk_as_variate(z::AbstractVector, ::Any) = z + + +""" + 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. + +The transport partner is an instance of the standard measure for measures +with scalar variates, and a power of it with +[`MeasureBase.some_dof(μ)`](@ref) (resp. `ν`) elements otherwise. +""" +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 + +function _std_tp_partner(::Type{M}, μ) where {M<:StdMeasure} + m = asmeasure(μ) + _std_tp_partner_byrank(M, _static_ndims(m), m) +end +_std_tp_partner_byrank(::Type{M}, ::StaticInteger{0}, μ) where {M<:StdMeasure} = M() +_std_tp_partner_byrank(::Type{M}, ::Any, μ) where {M<:StdMeasure} = M()^some_dof(μ) + + +# Element-wise transport kernels for broadcasts and maps: +struct _ToStd{S} <: Function end +@inline (::_ToStd{S})(μ, x) where {S} = transport_to_std(S, μ, x) +struct _FromStd{S} <: Function end +@inline (::_FromStd{S})(μ, z) where {S} = transport_from_std(S, μ, z) + +# Flat vector of standard variates from an array of standard variates: +@inline _flat_std_of(A::AbstractArray{<:Number}) = vec(A) +@inline _flat_std_of(A::AbstractArray{<:AbstractVector}) = _flatten_to_rv(vec(A)) +@inline _flat_std_of(A::AbstractVector{<:AbstractVector}) = _flatten_to_rv(A) diff --git a/src/utils.jl b/src/utils.jl index 0ec81a50..397c379a 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,13 +138,78 @@ 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) +isneginf(x) = isinf(x) & (x < zero(x)) +isposinf(x) = isinf(x) & (x > zero(x)) isapproxzero(x::T) where {T<:Real} = x ≈ zero(T) 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) + +# Regularized incomplete gamma and beta functions and their inverses, with +# the log-densities of the standard gamma and beta distributions for their +# derivatives. Implemented in the Distributions extension, differentiated +# with respect to the variate resp. probability argument in the autodiff +# extensions: +function _gamma_cdf end +function _gamma_quantile end +function _beta_cdf end +function _beta_quantile end +function _gamma_logpdf end +function _beta_logpdf end +function _gamma_cdf_impl end +function _gamma_quantile_impl end +function _beta_cdf_impl end +function _beta_quantile_impl end + +# The dual number type among the arguments of such a function, `Nothing` +# for plain numbers (the ForwardDiff extension adds dual numbers): +@inline _dualtag() = Nothing +@inline _dualtag(::Number, rest::Number...) = _dualtag(rest...) + +# Distributions implementation hooks, specialized for dual numbers in the +# ForwardDiff extension: +function _trafo_logcdf_impl end +function _trafo_logccdf_impl end +function _trafo_quantile_impl end +function _trafo_cquantile_impl end +function _dist_quantile end +function _dist_cquantile end +function _dist_params_numtype end diff --git a/test/Project.toml b/test/Project.toml index 376c1b05..c518a35a 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,21 +1,39 @@ [deps] +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" 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" +FixedSizeArrays = "3821ddf9-e5b5-40d5-8e25-6813ab96b5e2" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +ForwardDiffPullbacks = "450a3b6d-2448-4ee1-8e34-e4eb8713b605" +IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953" 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" +Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" 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" +StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[compat] +Aqua = "0.8 - 0.8.16" diff --git a/test/batched_regressions.jl b/test/batched_regressions.jl new file mode 100644 index 00000000..8471dfe3 --- /dev/null +++ b/test/batched_regressions.jl @@ -0,0 +1,136 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +# Regression tests for the batched-first review findings. + +using Test + +using MeasureBase +using MeasureBase: StdNormal, StdUniform, StdExponential, Dirac, Lebesgue, GenContext +using MeasureBase: productmeasure, pushfwd, mcombine, mbind, weightedmeasure, insupport, testvalue +using MeasureBase: batched_rand_impl, batched_transport_to_std_with_rest, batched_transport_from_std_with_rest +using MeasureBase.InverseFunctions: inverse +using ArraysOfArrays: VectorOfVectors, nestedview, flatview, sliced +using StaticArrays: SVector +using Static: static +using Distributions: MvNormal, logpdf +using LinearAlgebra: I +using JLArrays + +struct UnknownRankMeasure <: AbstractMeasure end +MeasureBase.basemeasure(::UnknownRankMeasure) = Lebesgue() +MeasureBase.logdensity_def(::UnknownRankMeasure, x) = -sum(abs, x) +MeasureBase.insupport(::UnknownRankMeasure, x) = true + +# A parameterized function object (not unwrapped into struct array columns): +struct Scale <: Function + s::Float64 +end +(f::Scale)(x) = f.s * x +MeasureBase.InverseFunctions.inverse(f::Scale) = Scale(inv(f.s)) +MeasureBase.ChangesOfVariables.with_logabsdet_jacobian(f::Scale, x) = (f(x), log(abs(f.s))) + +include("testutils.jl") + +@testset "batched regressions" begin + @testset "products of function-wrapper marginals" begin + pm = productmeasure([pushfwd(Scale(s), StdExponential()) for s in 0.1:0.2:0.9]) + x = rand(pm) + @test logdensityof(pm, x) ≈ sum(logdensityof.(MeasureBase.marginals(pm), x)) + X = rand(pm^4) + @test logdensities(pm, X) ≈ [logdensityof(pm, X[j]) for j in 1:4] + f = transport_to(StdUniform()^5, pm) + @test flatview(inverse(f).(f.(sliced(flatview(X), Val(1))))) ≈ flatview(X) + JLArrays.allowscalar(false) + @test Array(logdensities(MeasureBase.Adapt.adapt(JLArray, pm), JLArray(flatview(X)))) ≈ logdensities(pm, X) + end + + @testset "support of powers with array-variate bases" begin + mv = MeasureBase.AsMeasure{typeof(MvNormal(zeros(2), I(2)))}(MvNormal(zeros(2), I(2))) + Xf = randn(2, 3) + @test insupport(mv^3, Xf) === true + @test insupport(mv^3, nestedview(Xf)) === true + @test logdensities(mv, Xf) ≈ [logpdf(mv.obj, Xf[:, j]) for j in 1:3] + @test logdensityof(mv^3, Xf) ≈ sum(logpdf(mv.obj, Xf)) + end + + @testset "powers of array Diracs" begin + D = Dirac([1.0, 2.0]) + μ = D^3 + x = [1.0 1.0 1.0; 2.0 2.0 2.0] + @test MeasureBase.mspace_ndims(typeof(μ)) == 2 + @test logdensityof(μ, x) == 0 && insupport(μ, x) + @test logdensityof(μ, [[1.0, 2.0] for _ in 1:3]) == 0 + @test logdensityof(μ, 2 .* x) == -Inf + @test logdensities(μ, cat(x, 2 .* x; dims = 3)) == [0.0, -Inf] + @test rand(μ) == nestedview(x) + end + + @testset "streams with value-dependent sizes inside powers" begin + bnd = mbind(a -> StdNormal()^(a > 0.5 ? 2 : 1), StdUniform(), vcat) + @test testvalue(bnd) isa AbstractVector + m = mcombine(vcat, StdNormal(), bnd^2) + x = [0.5, 0.7, 0.1, 0.2, 0.3, 0.1] + ℓ = logdensityof(m, x) + @test ℓ ≈ logdensityof(StdNormal(), 0.5) + logdensityof(bnd, [0.7, 0.1, 0.2]) + logdensityof(bnd, [0.3, 0.1]) + @test logdensities(m, hcat(x, x)) ≈ [ℓ, ℓ] + f = transport_to(StdUniform()^6, m) + @test inverse(f)(f(x)) ≈ x + @test flatview(inverse(f).(f.(sliced(hcat(x, x), Val(1))))) ≈ hcat(x, x) + end + + @testset "powers of tuple products in streams" begin + Pt = productmeasure((StdNormal(), StdExponential()^2)) + m = mcombine(vcat, StdNormal(), Pt^2) + x = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7] + ℓ = logdensityof(StdNormal(), 0.1) + logdensityof(Pt, (0.2, [0.3, 0.4])) + logdensityof(Pt, (0.5, [0.6, 0.7])) + @test logdensityof(m, x) ≈ ℓ + X = hcat(x, 2 .* x) + @test logdensities(m, X) ≈ [logdensityof(m, X[:, j]) for j in 1:2] + f = transport_to(StdUniform()^7, m) + @test inverse(f)(f(x)) ≈ x + @test flatview(inverse(f).(f.(sliced(X, Val(1))))) ≈ X + Z, R = batched_transport_to_std_with_rest(StdUniform, Pt, hcat(x, x), (2,)) + @test size(Z) == (6, 2) && size(R) == (1, 2) + Xr, _ = batched_transport_from_std_with_rest(StdUniform, Pt, Z, (2,)) + @test Xr[1] ≈ [0.1 0.1; 0.4 0.4] && size(Xr[2]) == (2, 2, 2) + end + + @testset "batched kernels of measures without a declared rank" begin + w = weightedmeasure(0.3, UnknownRankMeasure()) + @test logdensityof(w, [1.0, 2.0]) ≈ 0.3 - 3 + @test_throws ArgumentError MeasureBase.batched_logdensity_def(w, randn(2, 2)) + @test MeasureBase.batched_logdensity_def(weightedmeasure(0.3, StdNormal()^2), [1.0, 2.0]) === 0.3 + end + + @testset "ragged batches are evaluated variate by variate" begin + V = VectorOfVectors([randn(2) for _ in 1:4]) + @test logdensities(StdNormal()^2, V) ≈ logdensityof.(Ref(StdNormal()^2), V) + @test_throws ArgumentError logdensities(StdNormal(), VectorOfVectors([[1.0], [2.0], [3.0], [4.0]])) + @test logdensities(StdNormal()^2, nestedview(randn(2, 4))) isa AbstractVector + end + + @testset "static streams stay allocation-free" begin + m1 = mcombine(vcat, StdNormal(), StdExponential()^static(2)) + x1 = SVector(0.1, 0.2, 0.3) + @test logdensityof(m1, x1) ≈ logdensityof(StdNormal(), 0.1) + logdensityof(StdExponential()^2, [0.2, 0.3]) + @test allocations_of(logdensityof, m1, x1) == 0 + m2 = mcombine(vcat, StdNormal()^static(2), StdExponential()^static(3)) + x2 = SVector(0.1, 0.2, 0.3, 0.4, 0.5) + @test allocations_of(logdensityof, m2, x2) == 0 + g = transport_to(StdUniform()^3, StdNormal()^3) + v = randn(3) + @test g.(SVector{3}(v))[] ≈ g(v) + gs = transport_to(StdUniform(), StdNormal()) + @test gs.(SVector{3}(v)) isa SVector{3,Float64} + end + + @testset "random variates of fused array products" begin + P = productmeasure([pushfwd(Base.Fix1(*, s), StdExponential()) for s in 0.1:0.2:0.9]) + x = rand(GenContext{Float64}(), P) + @test x isa Vector{Float64} && length(x) == 5 + X = batched_rand_impl(GenContext{Float64}(), P, (7,)) + @test size(X) == (5, 7) + Pj = productmeasure(JLArray([weightedmeasure(log(i), StdNormal()) for i in 1:3])) + @test length(rand(Pj)) == 3 + end +end diff --git a/test/combinators/bind.jl b/test/combinators/bind.jl new file mode 100644 index 00000000..413f2890 --- /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 (standard 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 (standard 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_std_with_rest(StdUniform, P, 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..3b6e3e1c --- /dev/null +++ b/test/combinators/combined.jl @@ -0,0 +1,74 @@ +# 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, pushfwd +using AffineMaps: Mul + +@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 "mcombine of products" begin + p1 = productmeasure([pushfwd(Mul(1.0), StdNormal()), pushfwd(Mul(2.0), StdNormal())]) + p2 = productmeasure([pushfwd(Mul(3.0), StdNormal())]) + p12 = mcombine(vcat, p1, p2) + @test p12 isa MeasureBase.ProductMeasure && length(MeasureBase.marginals(p12)) == 3 + x3 = randn(3) + @test logdensityof(p12, x3) ≈ logdensityof(p1, x3[1:2]) + logdensityof(p2, x3[3:3]) + + # Products of different or abstract marginal types stay combined measures: + mab = mcombine(vcat, StdNormal()^2, StdUniform()^2) + @test mab isa MeasureBase.CombinedMeasure + pa = productmeasure(AbstractMeasure[StdNormal(), StdNormal()]) + pb = productmeasure(AbstractMeasure[StdUniform(), StdUniform()]) + mpab = mcombine(vcat, pa, pb) + @test mpab isa MeasureBase.CombinedMeasure + x = vcat(randn(2), rand(2)) + @test logdensityof(mpab, x) ≈ logdensityof(mab, x) + Xab = vcat(randn(2, 3), rand(2, 3)) + @test logdensities(mpab, Xab) ≈ [logdensityof(mpab, x) for x in eachcol(Xab)] + 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/product.jl b/test/combinators/product.jl new file mode 100644 index 00000000..5a4ee6e1 --- /dev/null +++ b/test/combinators/product.jl @@ -0,0 +1,65 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using Random, Statistics +using StableRNGs: StableRNG +using StructArrays: StructArray +using Adapt: adapt +using JLArrays +using AffineMaps: Mul +using ArraysOfArrays: sliced, flatview +using InverseFunctions: inverse + +using MeasureBase +using MeasureBase: StdNormal, StdUniform, productmeasure, pushfwd, marginals, transport_to, logdensities + +@testset "products over arrays of marginals" begin + stblrng() = StableRNG(789990641) + + @testset "struct array storage" begin + P = productmeasure([pushfwd(Mul(s), StdNormal()) for s in (1.0, 2.0, 3.0)]) + mar = marginals(P) + @test mar isa StructArray + @test length(mar) == 3 && mar[2] == pushfwd(Mul(2.0), StdNormal()) + @test productmeasure(mar) == P + + x = randn(stblrng(), 3) + ℓ_ref = sum(logdensityof(m, xi) for (m, xi) in zip(mar, x)) + @test @inferred(logdensityof(P, x)) ≈ ℓ_ref + @test @inferred(MeasureBase.logdensity_def(P, x)) isa Real + X = randn(stblrng(), 3, 5) + @test @inferred(logdensities(P, X)) ≈ [logdensityof(P, x) for x in eachcol(X)] + @test logdensities(P, sliced(X, Val(1))) ≈ logdensities(P, X) + + f = transport_to(StdUniform()^3, P) + @test inverse(f)(f(x)) ≈ x + @test flatview(f.(sliced(X, Val(1)))) ≈ stack(map(f, eachcol(X))) + @test flatview(inverse(f).(f.(sliced(X, Val(1))))) ≈ X + + @test rand(stblrng(), P) isa Vector{Float64} + Xr = rand(stblrng(), P^100) + @test size(flatview(Xr)) == (3, 100) + @test isapprox(vec(mean(flatview(Xr), dims = 2)), zeros(3), atol = 1.0) + @test isapprox(vec(std(flatview(Xr), dims = 2)), [1.0, 2.0, 3.0], rtol = 0.4) + + # Non-isbits marginals keep their container: + Pv = productmeasure([pushfwd(Mul(randn(stblrng(), 2, 2)), StdNormal()^2) for _ in 1:2]) + @test !(marginals(Pv) isa StructArray) + end + + @testset "device arrays" begin + JLArrays.allowscalar(false) + P = productmeasure([pushfwd(Mul(s), StdNormal()) for s in (1.0, 2.0, 3.0)]) + Pj = adapt(JLArray, P) + @test marginals(Pj) isa StructArray + X = randn(stblrng(), 3, 5) + Xj = JLArray(X) + ℓj = logdensities(Pj, Xj) + @test ℓj isa JLArray && Array(ℓj) ≈ logdensities(P, X) + f = transport_to(StdUniform()^3, Pj) + Yj = f.(sliced(Xj, Val(1))) + @test flatview(Yj) isa JLArray + @test Array(flatview(Yj)) ≈ flatview(transport_to(StdUniform()^3, P).(sliced(X, Val(1)))) + 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..2369ae5e 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) @@ -9,17 +9,29 @@ using MeasureBase: superpose μs = μ + ν @test μs isa SuperpositionMeasure{<:Tuple{Dirac,Dirac}} @test μs == SuperpositionMeasure((μ, ν)) == superpose(μ, ν) - @test density_def(μs, 0) == 1.0 + @test density_def(μs, 0) == 0.5 @test basemeasure(μs) == CountingBase() + CountingBase() + @test densityof(μs, 0) == 1.0 μ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) == 0.5 + @test basemeasure(μs) == weightedmeasure(log(2), CountingBase()) + @test densityof(μs, 0) == 1.0 + # Base measures of components count wherever they have mass, not only + # where the component itself does: + @test logdensityof(superpose(StdNormal(), StdUniform()), -1.0) ≈ logdensityof(StdNormal(), -1.0) + @test logdensityof(superpose(StdNormal(), StdUniform()), 0.5) ≈ log(exp(logdensityof(StdNormal(), 0.5)) + 1) + # 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..a7f80ac5 100644 --- a/test/combinators/transformedmeasure.jl +++ b/test/combinators/transformedmeasure.jl @@ -4,6 +4,10 @@ using MeasureBase using MeasureBase: pushfwd, StdUniform, StdExponential, StdLogistic using MeasureBase: pushfwd, PushforwardMeasure using MeasureBase: transport_to, unsafe_logdensityof +using MeasureBase: productmeasure, mbind +import Zygote +using InverseFunctions: inverse +using ChangesOfVariables: with_logabsdet_jacobian import Statistics: var using DensityInterface: logdensityof using LogExpFunctions @@ -152,9 +156,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 +167,7 @@ using ChangesOfVariables # Test rand @test rand(ν) isa Real - @test insupport(ν, rand(ν)) + @test insupport(ν, rand(ν)) != false # Test pullback pb = pullbck(f, ν) @@ -178,4 +183,25 @@ end @test PushfwdRootMeasure() isa PushFwdStyle @test MeasureBase.WithVolCorr === AdaptRootMeasure @test MeasureBase.NoVolCorr === PushfwdRootMeasure + + @testset "output size of pushforwards of tuple products" begin + Pt = productmeasure((StdNormal(), StdUniform()^2)) + ν = pushfwd(x -> vcat(x[1], x[2]), Pt) + @test MeasureBase.mspace_flatsize(ν) == (3,) + @test MeasureBase.mspace_ndims(typeof(ν)) == 1 + @test MeasureBase.mspace_flatsize(pushfwd(x -> x, mbind(x -> StdNormal()^(x > 0 ? 1 : 2), StdNormal()))) isa MeasureBase.NoMSpaceElementSize + end + + @testset "construction inside differentiated functions" begin + flat(x::Tuple) = vcat(x[1], x[2]) + unflat(v::AbstractVector) = (v[1], v[2:end]) + InverseFunctions.inverse(::typeof(flat)) = unflat + InverseFunctions.inverse(::typeof(unflat)) = flat + ChangesOfVariables.with_logabsdet_jacobian(::typeof(flat), x) = (flat(x), zero(eltype(x[2]))) + ChangesOfVariables.with_logabsdet_jacobian(::typeof(unflat), v) = (unflat(v), zero(eltype(v))) + Pt = productmeasure((StdNormal(), StdUniform()^2)) + g(v) = logdensityof(pushfwd(flat, Pt), v) + v = [0.3, 0.2, 0.7] + @test Zygote.gradient(g, v)[1] ≈ [-0.3, 0.0, 0.0] + end end diff --git a/test/cuda/Project.toml b/test/cuda/Project.toml new file mode 100644 index 00000000..b9159db5 --- /dev/null +++ b/test/cuda/Project.toml @@ -0,0 +1,13 @@ +[deps] +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" +AffineMaps = "2c83c9a8-abf5-4329-a0d7-deffaf474661" +ArraysOfArrays = "65a8f2f4-9b39-5baf-92e2-a9cc46fdf018" +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +HeterogeneousComputing = "2182be2a-124f-4a91-8389-f06db5907a21" +MeasureBase = "fa1605e6-acd5-459c-a1e6-7e635759db14" +StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[sources] +MeasureBase = {path = "../.."} diff --git a/test/cuda/runtests.jl b/test/cuda/runtests.jl new file mode 100644 index 00000000..aaaba760 --- /dev/null +++ b/test/cuda/runtests.jl @@ -0,0 +1,147 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +# CUDA tests, not part of the default test suite (they need a CUDA GPU). +# Run with `julia --project=test/cuda test/cuda/runtests.jl` after +# instantiating that project. + +using Test +using CUDA +using Adapt: adapt +using HeterogeneousComputing: GenContext, AbstractComputeUnit +using MeasureBase +using MeasureBase: StdNormal, StdUniform, StdExponential, Dirac, asmeasure +using MeasureBase: productmeasure, pushfwd, mcombine, weightedmeasure, superpose, SpikeMixture +using MeasureBase: batched_rand_impl +using MeasureBase.InverseFunctions: inverse +using ArraysOfArrays: sliced, flatview +using AffineMaps: Mul, MulAdd +using Distributions: Normal, Uniform, Exponential, Logistic, Cauchy, Laplace, LogNormal, Weibull, Gamma, Beta +using Distributions: Poisson, Bernoulli, MvNormal, Dirichlet + +CUDA.allowscalar(false) + +# Evaluates `f` on device copies of `args` and compares with the plain +# result, the result must live on the device: +function test_cuda(f, args...) + expected = f(args...) + result = f(map(cu_copy, args)...) + @test _device_array(result) + @test _plain(result) ≈ _plain(expected) nans = true + return result +end + +cu_copy(x::AbstractArray) = CuArray(x) +cu_copy(μ::AbstractMeasure) = adapt(CuArray, μ) +cu_copy(x) = x +_device_array(x::AbstractArray) = parent_array(x) isa CuArray +_device_array(x::Tuple) = all(_device_array, x) +parent_array(x::CuArray) = x +parent_array(x::AbstractArray) = parent_array(parent(x)) +parent_array(x::Base.ReshapedArray) = parent_array(parent(x)) +_plain(x::AbstractArray) = Array(flatview(x)) +_plain(x::Tuple) = map(_plain, x) + +@testset "CUDA" begin + X = randn(3, 20) + Xc = vcat(randn(2, 20), rand(1, 20)) + + @testset "densities" begin + test_cuda(X -> logdensities(StdNormal()^3, X), X) + test_cuda(X -> logdensities(StdNormal()^3, sliced(X, Val(1))), X) + test_cuda(X -> logdensities((StdNormal()^3)^4, reshape(X[:, 1:16], 3, 4, 4)), X) + test_cuda(X -> logdensities(weightedmeasure(log(0.3), StdNormal()^3), X), X) + mc = mcombine(vcat, StdNormal()^2, StdUniform()^1) + test_cuda(X -> logdensities(mc, X), Xc) + mix = superpose(weightedmeasure(log(0.3), StdNormal()), weightedmeasure(log(0.7), StdUniform())) + test_cuda(x -> logdensities(mix, x), rand(20)) + test_cuda(x -> logdensities(SpikeMixture(StdNormal(), 0.2), x), vcat(randn(19), 0.0)) + test_cuda(x -> logdensities(Dirac(0.5), x), vcat(rand(19), 0.5)) + νe = pushfwd(Base.BroadcastFunction(exp), StdNormal()^3) + test_cuda(Y -> logdensities(νe, Y), exp.(X)) + P = productmeasure([pushfwd(Mul(s), StdNormal()) for s in (1.0, 2.0, 3.0)]) + test_cuda((P, X) -> logdensities(P, X), P, X) + test_cuda((P, X) -> logdensities(P, sliced(X, Val(1))), P, X) + Pa = productmeasure([MeasureBase.AsMeasure{Normal{Float64}}(Normal(μ, 1.0)) for μ in (0.0, 1.0, 2.0)]) + test_cuda((P, X) -> logdensities(P, X), Pa, X) + # Array products of array-variate marginals loop over the marginals + # on the host, the marginals stay host arrays: + Pv = productmeasure([weightedmeasure(log(i), StdNormal()^2) for i in 1:3]) + test_cuda(X -> logdensities(Pv, X), randn(2, 3, 5)) + end + + @testset "transports" begin + g = transport_to(StdUniform()^3, StdNormal()^3) + test_cuda(X -> flatview(g.(sliced(X, Val(1)))), X) + test_cuda(X -> flatview(inverse(g).(g.(sliced(X, Val(1))))), X) + h = transport_to(StdUniform()^(2, 3), (StdNormal()^2)^3) + test_cuda(X -> flatview(h.(sliced(X, Val(2)))), randn(2, 3, 4)) + mc = mcombine(vcat, StdNormal()^2, StdUniform()^1) + c = transport_to(StdExponential()^3, mc) + test_cuda(X -> flatview(c.(sliced(X, Val(1)))), Xc) + test_cuda(X -> flatview(inverse(c).(c.(sliced(X, Val(1))))), Xc) + P = productmeasure([pushfwd(Mul(s), StdNormal()) for s in (1.0, 2.0, 3.0)]) + f = transport_to(StdUniform()^3, P) + test_cuda((P, X) -> flatview(transport_to(StdUniform()^3, P).(sliced(X, Val(1)))), P, X) + test_cuda((P, X) -> flatview(transport_to(P, StdUniform()^3).(sliced(X, Val(1)))), P, rand(3, 20)) + Pv = productmeasure([weightedmeasure(log(i), StdNormal()^2) for i in 1:3]) + fv = transport_to(StdUniform()^6, Pv) + test_cuda(X -> flatview(fv.(sliced(X, Val(2)))), randn(2, 3, 5)) + test_cuda(X -> flatview(inverse(fv).(fv.(sliced(X, Val(2))))), randn(2, 3, 5)) + νe = pushfwd(Base.BroadcastFunction(exp), StdNormal()^3) + fe = transport_to(StdUniform()^3, νe) + test_cuda(Y -> flatview(fe.(sliced(Y, Val(1)))), exp.(X)) + A = [2.0 0.5; 0.0 1.5] + b = [1.0, -1.0] + νa = pushfwd(MulAdd(A, b), StdNormal()^2) + νac = pushfwd(MulAdd(CuArray(A), CuArray(b)), StdNormal()^2) + Ya = randn(2, 20) + ra = flatview(transport_to(StdUniform()^2, νac).(sliced(CuArray(Ya), Val(1)))) + @test ra isa CuArray && Array(ra) ≈ flatview(transport_to(StdUniform()^2, νa).(sliced(Ya, Val(1)))) + # AffineMaps has no device support for log-abs-det-Jacobians yet: + @test_broken Array(logdensities(νac, CuArray(Ya))) ≈ logdensities(νa, Ya) + end + + @testset "wrapped distributions" begin + for d in (Normal(0.3, 1.7), Uniform(-1.0, 2.5), Exponential(0.7), Logistic(0.2, 1.3), Cauchy(0.1, 0.8), Laplace(-0.4, 1.1), LogNormal(0.2, 0.6), Weibull(1.4, 0.9), Gamma(2.3, 1.2), Beta(2.5, 3.5)) + m = asmeasure(d) + xd = rand(d, 20) + test_cuda(x -> logdensities(m, x), xd) + f = transport_to(StdNormal(), m) + if d isa Beta + # SpecialFunctions' incomplete beta function doesn't compile for GPUs: + @test_broken Array(f.(CuArray(xd))) ≈ f.(xd) + else + test_cuda(x -> f.(x), xd) + test_cuda(z -> inverse(f).(z), randn(20)) + end + end + for d in (Poisson(2.7), Bernoulli(0.3)) + test_cuda(x -> logdensities(asmeasure(d), x), Float64.(rand(d, 20))) + end + mvn = MvNormal([0.3, -2.9], [1.7 0.5; 0.5 2.3]) + mm = asmeasure(mvn) + Xm = rand(mvn, 20) + test_cuda((m, X) -> logdensities(m, X), mm, Xm) + test_cuda((m, X) -> flatview(transport_to(StdNormal()^2, m).(sliced(X, Val(1)))), mm, Xm) + test_cuda((m, Z) -> flatview(transport_to(m, StdNormal()^2).(sliced(Z, Val(1)))), mm, randn(2, 20)) + dir = Dirichlet([2.0, 3.0, 4.0, 1.5]) + md = asmeasure(dir) + Xd = rand(dir, 20) + test_cuda((m, X) -> logdensities(m, X), md, Xd) + @test_broken flatview(transport_to(StdUniform()^3, cu_copy(md)).(sliced(CuArray(Xd), Val(1)))) isa CuArray + end + + @testset "random variates" begin + ctx = GenContext{Float32}(AbstractComputeUnit(CUDA.device()), CUDA.default_rng()) + for μ in (StdNormal(), StdUniform(), StdExponential(), StdNormal()^3, (StdNormal()^2)^3, weightedmeasure(0.3, StdNormal()^2), mcombine(vcat, StdNormal()^2, StdUniform()^1), pushfwd(Base.BroadcastFunction(exp), StdNormal()^2), superpose(StdNormal(), StdUniform()), SpikeMixture(StdNormal(), 0.5), asmeasure(Normal(0.3, 1.7)), asmeasure(Weibull(1.4, 0.9)), cu_copy(asmeasure(MvNormal([0.3, -2.9], [1.7 0.5; 0.5 2.3])))) + X = batched_rand_impl(ctx, μ, (100,)) + @test X isa CuArray{Float32} + ℓ = logdensities(μ, X) + @test ℓ isa CuArray && all(isfinite, Array(ℓ)) + end + Pt = productmeasure((StdNormal(), StdExponential()^2)) + Xt = batched_rand_impl(ctx, Pt, (50,)) + @test Xt isa Tuple && all(x -> x isa CuArray{Float32}, Xt) + @test size(Xt[2]) == (2, 50) + end +end diff --git a/test/distributions/getjacobian.jl b/test/distributions/getjacobian.jl new file mode 100644 index 00000000..87de7b86 --- /dev/null +++ b/test/distributions/getjacobian.jl @@ -0,0 +1,34 @@ +# This file is a part of ChangesOfVariables.jl, licensed under the MIT License (MIT). + +import ForwardDiff + +torv_and_back(V::AbstractVector{<:Real}) = V, identity +torv_and_back(x::Real) = [x], V -> V[1] +torv_and_back(x::Complex) = [real(x), imag(x)], V -> Complex(V[1], V[2]) +torv_and_back(x::NTuple{N}) where N = [x...], V -> ntuple(i -> V[i], Val(N)) + +function torv_and_back(x::Ref) + xval = x[] + V, to_xval = torv_and_back(xval) + back_to_ref(V) = Ref(to_xval(V)) + return (V, back_to_ref) +end + +torv_and_back(A::AbstractArray{<:Real}) = vec(A), V -> reshape(V, size(A)) + +function torv_and_back(A::AbstractArray{Complex{T}, N}) where {T<:Real, N} + RA = cat(real.(A), imag.(A), dims = N+1) + V, to_array = torv_and_back(RA) + function back_to_complex(V) + RA = to_array(V) + Complex.(view(RA, map(_ -> :, size(A))..., 1), view(RA, map(_ -> :, size(A))..., 2)) + end + return (V, back_to_complex) +end + + +function getjacobian(f, x) + V, to_x = torv_and_back(x) + vf(V) = torv_and_back(f(to_x(V)))[1] + ForwardDiff.jacobian(vf, V) +end diff --git a/test/distributions/test_autodiff_utils.jl b/test/distributions/test_autodiff_utils.jl new file mode 100644 index 00000000..5197725b --- /dev/null +++ b/test/distributions/test_autodiff_utils.jl @@ -0,0 +1,18 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using LinearAlgebra +using Distributions, ArraysOfArrays +import ForwardDiff, Zygote + + +@testset "trafo_utils" begin + xs = rand(5) + @test Zygote.jacobian(MeasureBase._pushfront, xs, 42)[1] ≈ ForwardDiff.jacobian(xs -> MeasureBase._pushfront(xs, 1), xs) + @test Zygote.jacobian(MeasureBase._pushfront, xs, 42)[2] ≈ vec(ForwardDiff.jacobian(x -> MeasureBase._pushfront(xs, x[1]), [42])) + @test Zygote.jacobian(MeasureBase._pushback, xs, 42)[1] ≈ ForwardDiff.jacobian(xs -> MeasureBase._pushback(xs, 1), xs) + @test Zygote.jacobian(MeasureBase._pushback, xs, 42)[2] ≈ vec(ForwardDiff.jacobian(x -> MeasureBase._pushback(xs, x[1]), [42])) + @test Zygote.jacobian(MeasureBase._rev_cumsum, xs)[1] ≈ ForwardDiff.jacobian(MeasureBase._rev_cumsum, xs) + @test Zygote.jacobian(MeasureBase._exp_cumsum_log, xs)[1] ≈ ForwardDiff.jacobian(MeasureBase._exp_cumsum_log, xs) ≈ ForwardDiff.jacobian(cumprod, xs) +end diff --git a/test/distributions/test_conversions.jl b/test/distributions/test_conversions.jl new file mode 100644 index 00000000..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_device_kernels.jl b/test/distributions/test_device_kernels.jl new file mode 100644 index 00000000..c75dffcf --- /dev/null +++ b/test/distributions/test_device_kernels.jl @@ -0,0 +1,145 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +# The device-friendly kernels of the wrapped distribution families: their +# densities and transports agree with Distributions, and they run on flat +# batches, also of device arrays. + +using Test +using Distributions, LinearAlgebra, StableRNGs, Statistics +using MeasureBase +using MeasureBase: asmeasure, GenContext, StdUniform, StdNormal, StdExponential, StdLogistic +using MeasureBase: batched_rand_impl, logdensities, insupport +using MeasureBase.InverseFunctions: inverse +using ArraysOfArrays: sliced, flatview +import Adapt +using JLArrays + +@testset "device kernels of distribution families" begin + JLArrays.allowscalar(false) + stblrng() = StableRNG(28734) + + families = [ + Normal(0.3, 1.7), Uniform(-1.0, 2.5), Exponential(0.7), Logistic(0.2, 1.3), Cauchy(0.1, 0.8), + Laplace(-0.4, 1.1), LogNormal(0.2, 0.6), Weibull(1.4, 0.9), Gamma(2.3, 1.2), Beta(2.5, 3.5), + Normal(0.3f0, 1.7f0), Gamma(0.7, 2.0), Beta(0.6, 0.8), + ] + + @testset "$(nameof(typeof(d)))" for d in families + m = asmeasure(d) + xs = vcat(rand(stblrng(), d, 20), [-1.0, 0.0, 1.0, 5.0, Inf]) + xs = eltype(d) == Float32 ? Float32.(xs) : xs + ℓ_ref = logpdf.(d, xs) + @test all(map((a, b) -> a == b || a ≈ b || (isnan(a) && isnan(b)), logdensityof.(Ref(m), xs), ℓ_ref)) + @test logdensities(m, xs) ≈ ℓ_ref nans = true + @test Array(logdensities(m, JLArray(xs))) ≈ ℓ_ref nans = true + @test insupport.(Ref(m), xs) == (Distributions.insupport.(d, xs) .& isfinite.(xs)) + + if d isa ContinuousUnivariateDistribution + x = rand(stblrng(), d, 12) + f = transport_to(StdUniform(), m) + p = f.(x) + @test p ≈ cdf.(d, x) + @test inverse(f).(p) ≈ x + @test Array(inverse(f).(f.(JLArray(x)))) ≈ x + g = transport_to(StdNormal(), m) + @test inverse(g).(g.(x)) ≈ x + X = batched_rand_impl(GenContext{Float64}(stblrng()), m, (5,)) + @test X isa Vector{Float64} && length(X) == 5 + # The transport-based generation used on devices, in single precision: + Xs = MeasureBase._rand_default(GenContext{Float32}(stblrng()), m, (2000,), MeasureBase._NoRandImpl()) + @test Xs isa Vector{Float32} && all(insupport.(Ref(m), Xs)) + if isfinite(mean(d)) && isfinite(var(d)) + @test isapprox(mean(Xs), mean(d), atol = 5 * sqrt(var(d) / 2000) + 1e-3) + end + end + end + + @testset "discrete families" begin + for d in (Poisson(2.7), Bernoulli(0.3)) + m = asmeasure(d) + xs = [0, 1, 2, 3, 7] + @test logdensityof.(Ref(m), xs) ≈ logpdf.(d, xs) + @test logdensities(m, xs) ≈ logpdf.(d, xs) + @test Array(logdensities(m, JLArray(xs))) ≈ logpdf.(d, xs) + @test logdensityof(m, -1) == -Inf && logdensityof(m, 1.5) == -Inf + end + end + + @testset "MvNormal" begin + for Σ in [[1.7 0.5; 0.5 2.3], PDMats.PDiagMat([0.5, 2.0]), PDMats.ScalMat(2, 1.5)] + d = MvNormal([0.3, -2.9], Σ) + m = asmeasure(d) + X = rand(stblrng(), d, 6) + ℓ_ref = logpdf(d, X) + @test logdensityof(m, X[:, 1]) ≈ ℓ_ref[1] + @test logdensities(m, X) ≈ ℓ_ref + @test logdensities(m, sliced(X, Val(1))) ≈ ℓ_ref + @test logdensityof(m^6, X) ≈ sum(ℓ_ref) + f = transport_to(StdNormal()^2, m) + Y = f.(sliced(X, Val(1))) + @test flatview(Y) ≈ stack(map(f, eachcol(X))) + @test flatview(inverse(f).(Y)) ≈ X + # JLArrays have no triangular solves, so only diagonal + # covariances run on them (CUDA covers the general case): + if !(Σ isa AbstractMatrix) + mj = Adapt.adapt(JLArray, m) + @test Array(logdensities(mj, JLArray(X))) ≈ ℓ_ref + fj = transport_to(StdNormal()^2, mj) + @test Array(flatview(fj.(sliced(JLArray(X), Val(1))))) ≈ flatview(Y) + @test Array(flatview(inverse(fj).(sliced(JLArray(flatview(Y)), Val(1))))) ≈ X + end + @test size(batched_rand_impl(GenContext{Float64}(stblrng()), m, (7,))) == (2, 7) + end + end + + @testset "Dirichlet" begin + d = Dirichlet([2.0, 3.0, 4.0, 1.5]) + m = asmeasure(d) + X = rand(stblrng(), d, 6) + ℓ_ref = logpdf(d, X) + @test logdensityof(m, X[:, 1]) ≈ ℓ_ref[1] + @test logdensities(m, X) ≈ ℓ_ref + @test logdensityof(m, [0.5, 0.5, 0.2, 0.1]) == -Inf + @test logdensityof(m, [0.5, 0.6, -0.1, 0.0]) == -Inf + mj = Adapt.adapt(JLArray, m) + @test Array(logdensities(mj, JLArray(X))) ≈ ℓ_ref + f = transport_to(StdUniform()^3, m) + Y = f.(sliced(X, Val(1))) + @test flatview(Y) ≈ stack(map(f, eachcol(X))) + @test flatview(inverse(f).(Y)) ≈ X + @test all(0 .<= flatview(Y) .<= 1) + Xr = batched_rand_impl(GenContext{Float64}(stblrng()), m, (200,)) + @test size(Xr) == (4, 200) && all(sum(Xr; dims = 1) .≈ 1) + end +end + +@testset "wrapped distributions outside the support" begin + for d in (Uniform(-1.0, 2.5), Exponential(0.7), LogNormal(0.2, 0.6), Weibull(1.4, 0.9), Gamma(2.3, 1.2), Beta(2.5, 3.5)) + m = asmeasure(d) + x = minimum(d) - 1 + @test logdensityof(m, x) == -Inf + @test isnan(transport_to(StdNormal(), m)(x)) + @test isnan(transport_to(StdNormal(), m)(Inf)) || d isa Union{Exponential,LogNormal,Weibull,Gamma} + Y = Array(transport_to(StdNormal(), m).(JLArray([x, mean(d)]))) + @test isnan(Y[1]) && !isnan(Y[2]) + end + for d in (Cauchy(0.1, 0.8), Laplace(-0.4, 1.1), Gamma(2.3, 1.2), Beta(2.5, 3.5), Uniform(-1.0, 2.5)) + f = transport_to(asmeasure(d), StdUniform()) + @test isnan(f(1.5)) && isnan(f(-0.5)) && !isnan(f(0.3)) + @test isequal(Array(f.(JLArray([1.5, 0.3]))), f.([1.5, 0.3])) + end + for d in (Exponential(0.7), Weibull(1.4, 0.9)) + @test isnan(transport_to(asmeasure(d), StdExponential())(-1.0)) + end + @test logdensityof(asmeasure(Rayleigh(2.0)), Inf) == -Inf + @test logdensityof(asmeasure(Poisson(2.7)), Inf) == -Inf + @test logdensityof(asmeasure(Poisson(2.7)), 1.5) == -Inf + @test logdensityof(asmeasure(Poisson(2.7)), -1.0) == -Inf + @test logdensityof(asmeasure(Bernoulli(0.3)), 0.5) == -Inf + md = asmeasure(Dirichlet([2.0, 3.0, 4.0, 1.5])) + @test all(isnan, transport_to(StdUniform()^3, md)([0.5, 0.5, 0.5, 0.5])) + @test all(isnan, transport_to(md, StdUniform()^3)([1.5, 0.5, 0.5])) + @test logdensityof(md, [0.5, 0.5, 0.5, 0.5]) == -Inf + Zd = flatview(transport_to(StdUniform()^3, md).(sliced([0.5 0.1; 0.5 0.2; 0.5 0.3; 0.5 0.4], Val(1)))) + @test all(isnan, Zd[:, 1]) && !any(isnan, Zd[:, 2]) +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..9e5777b0 --- /dev/null +++ b/test/distributions/test_distributions.jl @@ -0,0 +1,26 @@ +# 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_shape_contract.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_device_kernels.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..17c08427 --- /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.transport_from_std(StdUniform, dirich, u)[1], [0.3, 0.7]) + _test_gradient(x -> sum(MeasureBase.transport_to_std(StdUniform, 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_shape_contract.jl b/test/distributions/test_shape_contract.jl new file mode 100644 index 00000000..d4865fc1 --- /dev/null +++ b/test/distributions/test_shape_contract.jl @@ -0,0 +1,41 @@ +using Test + +using MeasureBase +using MeasureBase: mspace_elsize, mspace_flatsize, preferred_stdmeasure, NoStdTransport +using MeasureBase: StdNormal, StdUniform, StdExponential, StdLogistic, Dirac, productmeasure +using Distributions +using LinearAlgebra: I + +@testset "shape contract for Distributions" begin + @test @inferred(mspace_elsize(Normal(1, 2))) === () + @test @inferred(mspace_elsize(asmeasure(Normal(1, 2)))) === () + @test @inferred(mspace_flatsize(asmeasure(MvNormal(zeros(3), I(3))))) == (3,) + @test @inferred(mspace_elsize(asmeasure(MvNormal(zeros(3), I(3)))^4)) == (4,) + @test @inferred(mspace_flatsize(asmeasure(MvNormal(zeros(3), I(3)))^4)) == (3, 4) + + @test @inferred(preferred_stdmeasure(Normal(1, 2))) === StdNormal + @test @inferred(preferred_stdmeasure(asmeasure(Normal(1, 2)))) === StdNormal + @test @inferred(preferred_stdmeasure(3 + 2 * Normal())) === StdNormal + @test @inferred(preferred_stdmeasure(Uniform(1, 2))) === StdUniform + @test @inferred(preferred_stdmeasure(Exponential(2.0))) === StdExponential + @test @inferred(preferred_stdmeasure(Logistic(1, 2))) === StdLogistic + @test @inferred(preferred_stdmeasure(Beta(2, 3))) === StdUniform + @test @inferred(preferred_stdmeasure(truncated(Normal(), 0, 1))) === StdLogistic + @test @inferred(preferred_stdmeasure(MvNormal(zeros(2), I(2)))) === StdNormal + @test @inferred(preferred_stdmeasure(Dirichlet([1.0, 2.0]))) === StdUniform + @test @inferred(preferred_stdmeasure(Poisson(3))) <: NoStdTransport + @test @inferred(preferred_stdmeasure(StandardDist{Normal}(3))) === StdNormal + @test @inferred(preferred_stdmeasure(StandardDist{Uniform}())) === StdUniform + + @test @inferred(preferred_stdmeasure(productmeasure((asmeasure(Beta(2, 3)), asmeasure(Normal()))))) === StdNormal + @test @inferred(preferred_stdmeasure(productmeasure((a = Dirac(1.0), b = asmeasure(Beta(2, 3)))))) === StdUniform + @test @inferred(preferred_stdmeasure(productmeasure((a = asmeasure(Poisson(2)), b = asmeasure(Beta(2, 3)))))) <: NoStdTransport + @test @inferred(preferred_stdmeasure(productmeasure([asmeasure(Normal(i, 1)) for i in 1:3]))) === StdNormal + @test @inferred(preferred_stdmeasure(product_distribution([Beta(2, 3), Beta(1, 1)]))) === StdUniform + + lkj = asmeasure(LKJCholesky(3, 1.0)) + @test @inferred(mspace_elsize(lkj)) isa MeasureBase.NoMSpaceElementSize + @test @inferred(mspace_flatsize(lkj)) isa MeasureBase.NoMSpaceElementSize + X_lkj = [rand(LKJCholesky(3, 1.0)) for _ in 1:3] + @test logdensities(lkj, X_lkj) ≈ logdensityof.(Ref(lkj), X_lkj) +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..da14b275 --- /dev/null +++ b/test/distributions/test_transport.jl @@ -0,0 +1,257 @@ +# 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 ArraysOfArrays: sliced, flatview +using StableRNGs +import ForwardDiff, Zygote +import PDMats + +using MeasureBase: transport_to +using MeasureBase: StdUniform, StdNormal, StdExponential, StdLogistic +using .MeasureBaseDistributionsExt: _trafo_logcdf, _trafo_logccdf, _trafo_quantile, _trafo_cquantile + +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_to(trg, src)(x) + src_v_reco = transport_to(src, trg)(y) + + @test x ≈ src_v_reco + + f = x -> transport_to(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) = sliced(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{Normal}() + + 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 + dual_normal = Normal(Dual(0, 1, 0, 0), Dual(1, 0, 1, 0)) + dual_x = Dual(0.5, 0, 0, 1) + dual_p = Dual(0.3, 0, 0, 1) + + @test isapprox(_trafo_logcdf(dual_normal, dual_x), logcdf(dual_normal, dual_x), rtol = 10^-6) + @test isapprox(_trafo_logcdf(Normal(0, 1), Dual(0.5, 1)), logcdf(Normal(0, 1), Dual(0.5, 1)), rtol = 10^-6) + @test isapprox(_trafo_logccdf(dual_normal, dual_x), logccdf(dual_normal, dual_x), rtol = 10^-6) + @test isapprox(_trafo_logccdf(Normal(0, 1), Dual(0.5, 1)), logccdf(Normal(0, 1), Dual(0.5, 1)), rtol = 10^-6) + + @test isapprox(_trafo_quantile(Normal(0, 1), Dual(0.3, 1)), quantile(Normal(0, 1), Dual(0.3, 1)), rtol = 10^-6) + @test isapprox(_trafo_quantile(dual_normal, dual_p), quantile(dual_normal, dual_p), rtol = 10^-6) + @test isapprox(_trafo_cquantile(Normal(0, 1), Dual(0.3, 1)), cquantile(Normal(0, 1), Dual(0.3, 1)), rtol = 10^-6) + @test isapprox(_trafo_cquantile(dual_normal, dual_p), cquantile(dual_normal, dual_p), rtol = 10^-6) + + # Distributions whose cdf doesn't support dual numbers natively: + beta = Beta(2.0, 3.0) + dlogitcdf(d, x) = pdf(d, x) / (cdf(d, x) * ccdf(d, x)) + @test ForwardDiff.derivative(x -> transport_to(StdLogistic(), beta)(x), 0.3) ≈ dlogitcdf(beta, 0.3) + x_b = transport_to(beta, StdLogistic())(-0.4) + @test ForwardDiff.derivative(l -> transport_to(beta, StdLogistic())(l), -0.4) ≈ inv(dlogitcdf(beta, x_b)) + end + + @testset "tails of univariate transports" begin + # Bounded and heavy-lower-tailed distributions lose the lower tail in + # their quantile functions, so the ranges differ: + for (d, ls) in [ + (Normal(0.3, 1.7), [-700.0, -40.0, -8.0, 0.0, 8.0, 40.0, 700.0]), + (Weibull(0.7, 1.3), [-40.0, -8.0, 0.0, 8.0, 40.0, 700.0]), + (truncated(Normal(0.2, 1.1), -3.0, 2.5), [-8.0, 0.0, 8.0]), + ] + for l in ls + x = transport_to(d, StdLogistic())(l) + @test insupport(d, x) + @test isapprox(transport_to(StdLogistic(), d)(x), l, rtol = 1e-6, atol = 1e-12) + end + end + for z in [-8.0, 8.0, 37.0] + x = transport_to(Weibull(0.7, 1.3), StdNormal())(z) + @test isfinite(x) && x > 0 + @test transport_to(StdNormal(), Weibull(0.7, 1.3))(x) ≈ z rtol = 1e-6 + end + 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 "batched transport" begin + mvn = MvNormal([0.3, -2.9], [1.7 0.5; 0.5 2.3]) + f = transport_to(StdNormal()^2, mvn) + X = rand(StableRNG(789990641), mvn, 6) + Y = f.(sliced(X, Val(1))) + @test flatview(Y) ≈ stack(map(f, eachcol(X))) + @test flatview(inverse(f).(Y)) ≈ X + g = transport_to(StdNormal(), Weibull(0.7, 1.3)) + x = rand(StableRNG(789990641), Weibull(0.7, 1.3), 10) + @test g.(x) ≈ map(g, x) + pd = product_distribution([Weibull(0.7), Exponential(1.3), Normal(0.5, 2.0)]) + h = transport_to(StdNormal()^3, asmeasure(pd)) + Xp = rand(StableRNG(789990641), pd, 5) + @test stack(h.(sliced(Xp, Val(1)))) ≈ stack(map(h, eachcol(Xp))) + pn = product_distribution([Normal(1.0, 2.0), Normal(0.0, 3.0), Normal(2.0, 1.0)]) + mn = asmeasure(pn) + @test MeasureBase.mspace_flatsize(mn) == (3,) + hn = transport_to(StdUniform()^3, mn) + Xn = rand(StableRNG(789990641), pn, 4) + Yn = hn.(sliced(Xn, Val(1))) + @test flatview(Yn) ≈ stack(map(hn, eachcol(Xn))) + @test flatview(inverse(hn).(Yn)) ≈ Xn + @test eltype(rand(StableRNG(1), Float32, mn)) == Float32 + @test eltype(flatview(rand(StableRNG(1), Float32, mn^3))) == Float32 + @test eltype(rand(StableRNG(1), Float32, asmeasure(pd))) == Float32 + 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/fixed_size_arrays.jl b/test/fixed_size_arrays.jl new file mode 100644 index 00000000..6006f24e --- /dev/null +++ b/test/fixed_size_arrays.jl @@ -0,0 +1,48 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using MeasureBase +using MeasureBase: StdNormal, StdUniform, StdExponential, logdensities, weightedmeasure, productmeasure, pushfwd +using MeasureBase: batched_transport_to_std, batched_transport_from_std +using InverseFunctions: inverse +using FixedSizeArrays: FixedSizeArrayDefault +using ArraysOfArrays: flatview, sliced +using AffineMaps: Mul +using Distributions: Normal + +# Fixed-size inputs give fixed-size outputs, the kernels allocate via +# `similar` and never fall back to plain arrays: +@testset "fixed-size arrays" begin + fixed(A) = FixedSizeArrayDefault(A) + isfixed(A) = A isa FixedSizeArrayDefault + isfixed(A::Union{SubArray,Base.ReshapedArray}) = isfixed(parent(A)) + X = fixed(randn(3, 20)) + x = fixed(randn(3)) + m3 = StdNormal()^3 + + @testset "densities" begin + ℓ = logdensities(m3, X) + @test isfixed(ℓ) && ℓ ≈ logdensities(m3, Array(X)) + @test logdensityof(m3, x) ≈ logdensityof(m3, Array(x)) + @test isfixed(logdensities(m3, sliced(X, Val(1)))) + @test isfixed(logdensities((StdNormal()^2)^3, fixed(randn(2, 3, 4)))) + @test isfixed(logdensities(weightedmeasure(0.3, m3), X)) + P = productmeasure(fixed([pushfwd(Mul(s), StdNormal()) for s in (1.0, 2.0, 3.0)])) + @test isfixed(logdensities(P, X)) + Pn = productmeasure(fixed([Normal(μ, 1.0) for μ in (0.0, 1.0, 2.0)])) + @test isfixed(logdensities(Pn, X)) && logdensities(Pn, X) ≈ logdensities(Pn, Array(X)) + end + + @testset "transports" begin + f = transport_to(StdUniform()^3, m3) + y = f(x) + @test isfixed(y) && y ≈ f(Array(x)) + Y = f.(X) + @test isfixed(flatview(Y)) && flatview(Y) ≈ flatview(f.(Array(X))) + @test flatview(inverse(f).(Y)) ≈ X + @test isfixed(batched_transport_to_std(StdNormal, StdExponential()^3, fixed(rand(3, 5)))) + @test isfixed(batched_transport_from_std(StdNormal, StdExponential()^3, fixed(randn(3, 5)))) + @test isfixed(MeasureBase.convert_realtype(Float32, x)) + end +end diff --git a/test/logdensities.jl b/test/logdensities.jl new file mode 100644 index 00000000..d1c8a01d --- /dev/null +++ b/test/logdensities.jl @@ -0,0 +1,295 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using MeasureBase +using MeasureBase: logdensities, logdensity_def, StdNormal, StdUniform, StdExponential, StdLogistic, Dirac, Lebesgue, LebesgueBase, superpose, weightedmeasure, mcombine, productmeasure +using ArraysOfArrays: VectorOfSimilarVectors, sliced, flatview +using StaticArrays: SVector, @SVector, @SMatrix +using Static: static +using IrrationalConstants: log2π +import JLArrays +using JLArrays: JLArray + +stdnormal_ld(x) = -(x^2 + log2π) / 2 + +# A measure with array variates of a known flat size at the type level: +struct VecTestMeasure{T} <: AbstractMeasure + s::T +end +MeasureBase.mspace_elsize(::VecTestMeasure) = (2,) +MeasureBase.mspace_flatsize(::VecTestMeasure) = (2,) +MeasureBase.mspace_flatsize(::Type{<:VecTestMeasure}) = (2,) +MeasureBase.basemeasure(::VecTestMeasure) = LebesgueBase()^2 +MeasureBase.insupport(::VecTestMeasure, x) = true +MeasureBase.logdensityof_impl(m::VecTestMeasure, x) = -sum(abs2, x) / (2 * m.s) + +include("testutils.jl") + +@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 = sliced(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 "unknown variate size" begin + mix = superpose(weightedmeasure(log(0.3), StdNormal()), weightedmeasure(log(0.7), StdUniform())) + X = randn(4, 5) + @test logdensities(mix, X) ≈ logdensityof.(Ref(mix), X) + @test logdensityof(mix^4, X[:, 1]) ≈ sum(logdensityof.(Ref(mix), X[:, 1])) + @test logdensities(mix^4, sliced(X, 1)) ≈ vec(sum(logdensityof.(Ref(mix), X), dims = 1)) + @test @inferred(logdensityof(mix^0, Float64[])) == 0 + @test logdensities(mix^0, [Float64[], Float64[]]) == [0.0, 0.0] + end + + @testset "flat and nested variate forms agree" begin + for (μ, x_flat) in ( + ((StdNormal()^3)^2, randn(3, 2)), + ((StdNormal()^(2, 3))^4, randn(2, 3, 4)), + ) + x_nested = sliced(x_flat, length(MeasureBase.mspace_flatsize(MeasureBase.pwr_base(μ)))) + @test logdensityof(μ, x_flat) ≈ logdensityof(μ, x_nested) + @test MeasureBase.checked_arg(μ, x_flat) === x_flat + @test MeasureBase.checked_arg(μ, x_nested) === x_nested + @test_throws ArgumentError MeasureBase.checked_arg(μ, randn(7)) + end + end + + @testset "size mismatch" begin + @test_throws ArgumentError logdensityof(StdNormal()^3, randn(3, 4)) + @test_throws ArgumentError logdensityof(StdNormal()^(2, 3), randn(2, 3, 1)) + @test_throws ArgumentError logdensityof(StdNormal()^3, randn(4)) + @test_throws ArgumentError logdensityof(StdNormal()^3, 1.0) + @test_throws ArgumentError logdensities(StdNormal(), VectorOfSimilarVectors(randn(3, 5))) + @test_throws ArgumentError logdensities(StdNormal()^3, [randn(3), randn(2)]) + @test_throws ArgumentError logdensities( + StdNormal()^3, + VectorOfSimilarVectors(randn(2, 5)), + ) + end + + @testset "flat batch storage and array-variate bases" begin + m3 = StdNormal()^3 + Xf = randn(3, 10) + @test @inferred(logdensities(m3, Xf)) ≈ vec(sum(stdnormal_ld.(Xf), dims = 1)) + x = randn(3) + @test @inferred(logdensities(m3, x)) ≈ sum(stdnormal_ld, x) + + mpp = (StdNormal()^(2, 3))^4 + Xpp_flat = randn(2, 3, 4, 7) + @test @inferred(logdensities(mpp, Xpp_flat)) ≈ vec(sum(stdnormal_ld.(Xpp_flat), dims = (1, 2, 3))) + Xpp_nested = sliced(sliced(Xpp_flat, 2), 1) + @test logdensities(mpp, Xpp_nested) ≈ logdensities(mpp, Xpp_flat) + xpp = randn(2, 3, 4) + @test @inferred(logdensityof(mpp, xpp)) ≈ logdensityof(mpp, [xpp[:, :, i] for i in 1:4]) + @test logdensityof(mpp, sliced(xpp, 2)) ≈ logdensityof(mpp, xpp) + + mvec = Dirac([1.0, 2.0])^3 + @test @inferred(logdensities(mvec, [fill([1.0, 2.0], 3) for _ in 1:2])) == [0.0, 0.0] + end + + @testset "static variates" begin + m3 = StdNormal()^static(3) + xs = @SVector randn(3) + @test @inferred(logdensityof(m3, xs)) ≈ sum(stdnormal_ld, xs) + @test allocations_of(logdensityof, m3, xs) == 0 + xd = randn(3) + @test @inferred(logdensityof(StdNormal()^3, xd)) ≈ sum(stdnormal_ld, xd) + @test allocations_of(logdensityof, StdNormal()^3, xd) == 0 + Xs = @SMatrix randn(3, 4) + @test @inferred(logdensities(m3, Xs)) ≈ vec(sum(stdnormal_ld.(Xs), dims = 1)) + @test logdensities(m3, Xs) isa SVector{4} + @test @inferred(logdensityof(StdNormal()^static(0), SVector{0,Float64}())) == 0 + @test @inferred(logdensityof(StdNormal()^0, Float64[])) == 0 + end + + @testset "powers of primitive measures" begin + @test @inferred(logdensity_def(Lebesgue()^3, randn(3))) == 0 + @test @inferred(logdensity_def(LebesgueBase()^(2, 2), randn(2, 2))) == 0 + @test @inferred(logdensityof(Lebesgue()^3, randn(3))) == 0 + end + + @testset "structural batched kernels" begin + w = weightedmeasure(log(0.3), StdNormal()^3) + X = randn(3, 10) + @test @inferred(logdensities(w, X)) ≈ [logdensityof(w, x) for x in eachcol(X)] + xw = randn(3) + @test @inferred(logdensityof(w, xw)) ≈ log(0.3) + sum(stdnormal_ld, xw) + @test allocations_of(logdensityof, w, xw) == 0 + + ms = [weightedmeasure(log(i), StdNormal()) for i in 1:4] + prod4 = productmeasure(ms) + @test @inferred(MeasureBase.mspace_flatsize(prod4)) == (4,) + @test @inferred(MeasureBase.mspace_elsize(prod4)) == (4,) + xp = randn(4) + @test @inferred(logdensityof(prod4, xp)) ≈ sum(log(i) + stdnormal_ld(xp[i]) for i in 1:4) + @test allocations_of(logdensityof, prod4, xp) == 0 + Xp = randn(4, 7) + @test @inferred(logdensities(prod4, Xp)) ≈ [logdensityof(prod4, x) for x in eachcol(Xp)] + @test logdensities(prod4, sliced(Xp, 1)) ≈ logdensities(prod4, Xp) + @test @inferred(logdensityof(prod4^2, randn(4, 2))) isa Float64 + Xpp = randn(4, 2, 5) + @test logdensities(prod4^2, Xpp) ≈ [logdensityof(prod4^2, Xpp[:, :, i]) for i in 1:5] + + ms2 = reshape([weightedmeasure(log(i), StdUniform()) for i in 1:6], 2, 3) + prod23 = productmeasure(ms2) + @test @inferred(MeasureBase.mspace_flatsize(prod23)) == (2, 3) + x23 = rand(2, 3) + @test @inferred(logdensityof(prod23, x23)) ≈ sum(log(i) for i in 1:6) + @test logdensities(prod23, rand(2, 3, 4)) ≈ fill(sum(log(i) for i in 1:6), 4) + + mvec = productmeasure([StdNormal()^2, StdNormal()^2]) + @test @inferred(MeasureBase.mspace_flatsize(mvec)) isa MeasureBase.NoMSpaceElementSize + end + + @testset "batched with-rest for combined measures" begin + m = mcombine(vcat, StdNormal()^2, StdUniform()^3) + @test @inferred(MeasureBase.mspace_flatsize(m)) == (5,) + x = vcat(randn(2), rand(3)) + @test @inferred(logdensityof(m, x)) ≈ sum(stdnormal_ld, x[1:2]) + X = vcat(randn(2, 6), rand(3, 6)) + @test @inferred(logdensities(m, X)) ≈ [logdensityof(m, x) for x in eachcol(X)] + @test logdensities(m, sliced(X, 1)) ≈ logdensities(m, X) + ℓ, A_rest = MeasureBase.batched_logdensityof_with_rest(StdNormal()^2, X, ()) + @test ℓ ≈ vec(sum(stdnormal_ld.(X[1:2, :]), dims = 1)) + @test size(A_rest) == (3, 6) + ℓ2, A_rest2 = MeasureBase.batched_logdensityof_with_rest(StdNormal(), X, (2,)) + @test MeasureBase._materialize(ℓ2) ≈ stdnormal_ld.(X[1:2, :]) && size(A_rest2) == (3, 6) + @test_throws ArgumentError logdensities(m, vcat(X, rand(1, 6))) + + m3 = mcombine(vcat, StdNormal(), mcombine(vcat, StdExponential()^2, StdLogistic())) + @test @inferred(MeasureBase.mspace_flatsize(m3)) == (4,) + X3 = vcat(randn(1, 5), rand(2, 5), randn(1, 5)) + @test @inferred(logdensities(m3, X3)) ≈ [logdensityof(m3, x) for x in eachcol(X3)] + end + + @testset "static sizes in combined measures" begin + m = mcombine(vcat, StdNormal()^static(2), StdUniform()^3) + @test @inferred(MeasureBase.mspace_flatsize(m)) == (5,) + x = vcat(randn(2), rand(3)) + @test @inferred(logdensityof(m, x)) ≈ sum(stdnormal_ld, x[1:2]) + X = vcat(randn(2, 4), rand(3, 4)) + @test @inferred(logdensities(m, X)) ≈ [logdensityof(m, x) for x in eachcol(X)] + ms = mcombine(vcat, StdNormal()^static(2), StdUniform()^static(3)) + @test @inferred(MeasureBase.mspace_flatsize(ms)) == MeasureBase.mspace_flatsize(StdNormal()^static(5)) + @test @inferred(logdensityof(ms, SVector{5}(x))) ≈ logdensityof(m, x) + @test @inferred(logdensityof(ms, x)) ≈ logdensityof(m, x) + @test logdensities(ms, X) ≈ logdensities(m, X) + @test_throws ArgumentError MeasureBase.batched_logdensityof_with_rest(StdNormal(), zeros(0, 4), ()) + end + + @testset "array products of array-variate marginals" begin + p = productmeasure([VecTestMeasure(1.0), VecTestMeasure(2.0), VecTestMeasure(0.5)]) + @test @inferred(MeasureBase.mspace_flatsize(p)) == (2, 3) + xs = [randn(2) for _ in 1:3] + X = stack(xs) + ℓ = sum(map(logdensityof, MeasureBase.marginals(p), xs)) + @test @inferred(logdensityof(p, xs)) ≈ ℓ + @test @inferred(logdensityof(p, X)) ≈ ℓ + A = randn(2, 3, 5) + @test @inferred(logdensities(p, A)) ≈ [logdensityof(p, A[:, :, i]) for i in 1:5] + @test logdensities(p, sliced(A, Val(2))) ≈ logdensities(p, A) + @test logdensities(p, zeros(2, 3, 0)) == Float64[] + @test_throws ArgumentError logdensityof(p, randn(2, 2)) + @test_throws ArgumentError logdensities(p, randn(2, 2, 5)) + + # Powers with static axes have a flat size at the type level: + pp = MeasureBase.ProductMeasure([StdNormal()^static(2), StdNormal()^static(2)]) + @test @inferred(MeasureBase.mspace_flatsize(pp)) == (2, 2) + Xp = randn(2, 2) + @test @inferred(logdensityof(pp, Xp)) ≈ sum(stdnormal_ld, Xp) + A3 = randn(2, 2, 3) + @test @inferred(logdensities(pp, A3)) ≈ [sum(stdnormal_ld, A3[:, :, i]) for i in 1:3] + end + + @testset "products with mixed marginal types" begin + pa = productmeasure(AbstractMeasure[StdNormal(), StdUniform()]) + X = vcat(randn(1, 4), rand(1, 4)) + xs = [X[:, i] for i in 1:4] + @test logdensities(pa, xs) ≈ [logdensityof(pa, x) for x in xs] + @test logdensities(pa, sliced(X, Val(1))) ≈ [logdensityof(pa, x) for x in xs] + @test_throws ArgumentError logdensityof(pa, 0.5) + @test_throws ArgumentError logdensityof(productmeasure((StdNormal(), StdUniform())), 0.5) + @test_throws ArgumentError logdensityof(pa, X[:, 1:1]) + end + + @testset "out-of-support and empty batches of structural kernels" begin + w = weightedmeasure(log(0.3), StdUniform()^2) + X = [0.5 -0.5 0.5; 0.5 0.5 1.5] + @test logdensities(w, X) == [log(0.3), -Inf, -Inf] + @test logdensities(w, zeros(2, 0)) == Float64[] + pu = productmeasure([weightedmeasure(log(i), StdUniform()) for i in 1:2]) + @test logdensities(pu, X) == [log(2), -Inf, -Inf] + mc = mcombine(vcat, StdUniform()^1, StdExponential()^1) + @test logdensities(mc, [0.5 -0.5 0.5; 0.5 0.5 -1.0]) == [-0.5, -Inf, -Inf] + @test logdensities(mc, zeros(2, 0)) == Float64[] + @test MeasureBase.logdensity_def(pu, [0.5, 0.5]) ≈ log(2) + end + + @testset "GPU array semantics for structural kernels" begin + JLArrays.allowscalar(false) + ms = JLArray([weightedmeasure(log(i), StdNormal()) for i in 1:4]) + prodj = productmeasure(ms) + Xj = JLArray(randn(4, 7)) + ldj = logdensities(prodj, Xj) + @test ldj isa JLArray + @test Array(ldj) ≈ logdensities(productmeasure(Array(ms)), Array(Xj)) + mc = mcombine(vcat, StdNormal()^2, StdUniform()^3) + Xc = JLArray(vcat(randn(2, 6), rand(3, 6))) + ldc = logdensities(mc, Xc) + @test ldc isa JLArray + @test Array(ldc) ≈ logdensities(mc, Array(Xc)) + 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/numtype.jl b/test/numtype.jl new file mode 100644 index 00000000..5a962ae9 --- /dev/null +++ b/test/numtype.jl @@ -0,0 +1,39 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using MeasureBase +using MeasureBase: StdNormal, StdUniform, LebesgueBase, Lebesgue, Dirac +using MeasureBase: logdensity_def, logdensity_rel, weightedmeasure, superpose, logdensities, mbind +using Static: static +import ForwardDiff + +@testset "number type of log-densities" begin + x = 0.5f0 + xf = randn(Float32, 3) + uf = rand(Float32, 3) + mix = superpose(weightedmeasure(log(0.3f0), StdNormal()), weightedmeasure(log(0.7f0), StdUniform())) + + @test @inferred(logdensityof(StdNormal(), x)) isa Float32 + @test @inferred(logdensityof(StdNormal()^3, xf)) isa Float32 + @test @inferred(MeasureBase.unsafe_logdensityof(StdNormal()^3, xf)) isa Float32 + @test @inferred(logdensity_def(LebesgueBase()^3, xf)) isa Float32 + @test @inferred(logdensity_def(LebesgueBase(), x)) isa Float32 + @test @inferred(logdensity_rel(StdNormal(), StdUniform(), x)) isa Float32 + @test @inferred(logdensity_rel(StdNormal()^3, StdUniform()^3, uf)) isa Float32 + @test @inferred(logdensity_rel(StdNormal()^3, StdUniform()^3, xf)) isa Float32 + @test @inferred(logdensity_rel(Lebesgue(), Dirac(1f0), 2f0)) isa Float32 + @test @inferred(logdensity_rel(Dirac(1f0), Lebesgue(), 1f0)) isa Float32 + @test @inferred(logdensityof(Dirac(1f0), 1f0)) isa Float32 + @test @inferred(logdensityof(weightedmeasure(0.5f0, StdNormal())^3, xf)) isa Float32 + @test @inferred(logdensityof(weightedmeasure(static(0.5), StdNormal())^3, xf)) isa Float32 + @test @inferred(logdensityof(weightedmeasure(0.5, StdNormal()), x)) isa Float32 + @test @inferred(logdensityof(mix, x)) isa Float32 + @test @inferred(logdensityof(mix^3, uf)) isa Float32 + @test @inferred(logdensities(StdNormal(), randn(Float32, 4))) isa Vector{Float32} + @test @inferred(logdensities(StdNormal()^3, randn(Float32, 3, 4))) isa Vector{Float32} + + # Weights that carry derivatives keep them: + @test ForwardDiff.derivative(w -> logdensityof(weightedmeasure(w, StdNormal()), 0.3), 0.1) ≈ 1 + @test ForwardDiff.derivative(w -> logdensityof(weightedmeasure(w, StdNormal())^2, [0.3, 0.1]), 0.1) ≈ 2 +end diff --git a/test/rand.jl b/test/rand.jl new file mode 100644 index 00000000..867c6276 --- /dev/null +++ b/test/rand.jl @@ -0,0 +1,122 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using Random, Statistics +using StableRNGs: StableRNG +using Static: static +using ArraysOfArrays: flatview +using AffineMaps: Add + +using MeasureBase +using MeasureBase: GenContext +using MeasureBase: StdUniform, StdExponential, StdLogistic, StdNormal, Dirac, Lebesgue +using MeasureBase: weightedmeasure, superpose, mcombine, mbind, productmeasure, pushfwd, testvalue +using MeasureBase: rand_impl, batched_rand_impl, massof, isnormalized + +@testset "rand" begin + stblrng() = StableRNG(789990641) + + @testset "generative contexts" begin + @test rand(stblrng(), StdNormal()) == rand(stblrng(), StdNormal()) + @test rand(stblrng(), StdNormal()) == rand(GenContext{Float64}(stblrng()), StdNormal()) + @test rand(stblrng(), StdNormal()^3) == rand(stblrng(), StdNormal()^3) + @test @inferred(rand(stblrng(), Float32, StdNormal())) isa Float32 + @test @inferred(rand(Float32, StdNormal()^3)) isa Vector{Float32} + @test @inferred(rand(GenContext{Float32}(stblrng()), StdUniform()^(2, 3))) isa Matrix{Float32} + @test @inferred(rand(StdExponential())) isa Float64 + @test_throws ArgumentError rand(Lebesgue()) + end + + @testset "layout of power variates" begin + x = rand(StdNormal()^(2, 3)) + @test x isa Matrix{Float64} && size(x) == (2, 3) + xs = rand(StdNormal()^static(3)) + @test xs isa AbstractVector{Float64} && length(xs) == 3 + xn = rand((StdNormal()^2)^3) + @test xn isa AbstractVector && length(xn) == 3 && all(x -> length(x) == 2, xn) + @test size(flatview(xn)) == (2, 3) + @test logdensityof((StdNormal()^2)^3, xn) ≈ logdensityof(StdNormal()^6, vec(flatview(xn))) + @test batched_rand_impl(GenContext{Float64}(stblrng()), StdNormal()^2, (4, 5)) isa Array{Float64,3} + @test size(batched_rand_impl(GenContext{Float64}(stblrng()), (StdNormal()^2)^3, (4,))) == (2, 3, 4) + end + + @testset "test values" begin + @test testvalue(mcombine(vcat, StdNormal()^2, StdUniform()^3)) == [0.0, 0.0, 0.5, 0.5, 0.5] + @test rand(MeasureBase.ConstantRNG(), Float64, StdUniform()^3) == fill(0.5, 3) + @test rand(MeasureBase.ConstantRNG(), Float32, StdLogistic()^2) == zeros(Float32, 2) + @test testvalue(StdNormal()) == 0 + @test testvalue(Float32, StdUniform()) === 0.5f0 + @test testvalue(StdExponential()^3) == ones(3) + @test testvalue((StdLogistic()^2)^2) == [zeros(2), zeros(2)] + @test testvalue(productmeasure((a = StdNormal(), b = StdUniform()^2))) == (a = 0.0, b = [0.5, 0.5]) + end + + @testset "distribution of variates" begin + n = 10^5 + for (μ, m, v) in [ + (StdUniform(), 0.5, 1 / 12), + (StdExponential(), 1.0, 1.0), + (StdLogistic(), 0.0, π^2 / 3), + (StdNormal(), 0.0, 1.0), + (weightedmeasure(0.3, StdNormal()), 0.0, 1.0), + (pushfwd(exp, StdNormal()), exp(0.5), (exp(1) - 1) * exp(1)), + (MeasureBase.Half(StdNormal()), sqrt(2 / π), 1 - 2 / π), + (SpikeMixture(Dirac(1.0), 0.25), 0.25, 0.25 * 0.75), + ] + X = rand(stblrng(), μ^n) + @test isapprox(mean(X), m, atol = 5 * sqrt(v / n) + 1e-3) + @test isapprox(var(X), v, rtol = 0.05) + xs = [rand_impl(GenContext{Float64}(stblrng()), μ) for _ in 1:20] + @test all(x -> insupport(μ, x) != false, xs) + end + + mix = superpose(weightedmeasure(log(0.3), Dirac(0.0)), weightedmeasure(log(0.7), Dirac(1.0))) + @test isapprox(mean(rand(stblrng(), mix^n)), 0.7, atol = 0.01) + @test isapprox(mean([rand(mix) for _ in 1:n]), 0.7, atol = 0.01) + mixn = superpose(weightedmeasure(log(0.5), StdNormal()), weightedmeasure(log(0.5), pushfwd(Add(4.0), StdNormal()))) + @test isapprox(mean(rand(stblrng(), mixn^n)), 2.0, atol = 0.02) + end + + @testset "structural measures" begin + P = MeasureBase.ProductMeasure([weightedmeasure(log(i), StdNormal()) for i in 1:3]) + @test @inferred(rand(stblrng(), P)) isa Vector{Float64} + XP = rand(stblrng(), P^100) + @test size(flatview(XP)) == (3, 100) + @test logdensities(P, XP) ≈ [logdensityof(P, x) for x in XP] + + Pt = productmeasure((StdNormal(), StdUniform()^2)) + xt = rand(stblrng(), Pt) + @test xt isa Tuple{Float64,Vector{Float64}} + Xt = rand(stblrng(), Pt^5) + @test Xt isa AbstractVector && length(Xt) == 5 + + mc = mcombine(vcat, StdNormal()^2, StdUniform()^3) + xc = rand(stblrng(), mc) + @test xc isa Vector{Float64} && length(xc) == 5 && all(0 .<= xc[3:5] .<= 1) + Xc = rand(stblrng(), mc^50) + @test size(flatview(Xc)) == (5, 50) + @test logdensities(mc, Xc) ≈ [logdensityof(mc, x) for x in Xc] + + f_β(a) = StdNormal()^length(a) + μb = mbind(f_β, StdUniform()^2, vcat) + xb = rand(stblrng(), μb) + @test xb isa AbstractVector && length(xb) == 4 + Xb = rand(stblrng(), μb^3) + @test Xb isa AbstractVector && length(Xb) == 3 && all(x -> length(x) == 4, Xb) + + ctx = GenContext{Float64}(stblrng()) + spd = superpose(weightedmeasure(log(0.5), Dirac([1.0, 2.0])), weightedmeasure(log(0.5), Dirac([3.0, 4.0]))) + Xspd = batched_rand_impl(ctx, spd, (6,)) + @test size(Xspd) == (2, 6) && all(c -> c == [1.0, 2.0] || c == [3.0, 4.0], eachcol(Xspd)) + Xsm = batched_rand_impl(ctx, SpikeMixture(StdNormal()^3, 0.5), (4,)) + @test size(Xsm) == (3, 4) && all(c -> all(iszero, c) || !any(iszero, c), eachcol(Xsm)) + @test massof(StdNormal()) == 1 && massof(StdUniform()^3) == 1 && massof(weightedmeasure(log(2.0), StdNormal()^2)) ≈ 2 + @test isnormalized(StdNormal()) && isnormalized(StdExponential()^(2, 2)) && !isnormalized(2.0 * StdNormal()) + @test !isnormalized(Lebesgue()) + + d = Dirac([1.0, 2.0]) + @test rand(d^2) == [[1.0, 2.0], [1.0, 2.0]] + @test batched_rand_impl(GenContext{Float64}(stblrng()), d, (2,)) == [1.0 1.0; 2.0 2.0] + end +end diff --git a/test/rand_batched.jl b/test/rand_batched.jl new file mode 100644 index 00000000..ee63ef75 --- /dev/null +++ b/test/rand_batched.jl @@ -0,0 +1,102 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using Random, Statistics +using StableRNGs: StableRNG +using StructArrays: StructArray +using ArraysOfArrays: flatview + +using MeasureBase +using MeasureBase: GenContext +using MeasureBase: StdUniform, StdExponential, StdLogistic, StdNormal, Dirac +using MeasureBase: weightedmeasure, superpose, mcombine, mbind, productmeasure, pushfwd, SpikeMixture +using MeasureBase: rand_impl, batched_rand_impl +using Distributions: Normal, MvNormal, logpdf + +struct NoBatchRandMeasure <: AbstractMeasure end +MeasureBase.rand_impl(ctx::GenContext, ::NoBatchRandMeasure) = 3 * rand(MeasureBase.get_rng(ctx), MeasureBase.get_precision(ctx)) + +struct NoRandMeasure <: AbstractMeasure end + +@testset "batched rand" begin + stblrng() = StableRNG(789) + ctx() = GenContext{Float64}(stblrng()) + + @testset "single variates as batches with zero batch dimensions" begin + @test @inferred(batched_rand_impl(ctx(), StdNormal(), ())) isa Float64 + @test batched_rand_impl(ctx(), StdNormal(), ()) == rand(stblrng(), StdNormal()) + @test @inferred(batched_rand_impl(ctx(), StdExponential(), ())) isa Float64 + @test @inferred(batched_rand_impl(ctx(), StdLogistic(), ())) isa Float64 + @test @inferred(batched_rand_impl(ctx(), StdUniform(), ())) isa Float64 + @test @inferred(batched_rand_impl(ctx(), StdNormal()^3, ())) isa Vector{Float64} + @test batched_rand_impl(ctx(), StdNormal()^3, ()) == rand(stblrng(), StdNormal()^3) + @test batched_rand_impl(ctx(), (StdNormal()^2)^3, ()) == flatview(rand(stblrng(), (StdNormal()^2)^3)) + @test batched_rand_impl(ctx(), Dirac(2.5), ()) === 2.5 + @test batched_rand_impl(ctx(), Dirac([1.0, 2.0]), ()) == [1.0, 2.0] + @test batched_rand_impl(ctx(), weightedmeasure(0.3, StdNormal()), ()) isa Float64 + @test batched_rand_impl(ctx(), pushfwd(exp, StdNormal()), ()) isa Float64 + @test batched_rand_impl(ctx(), pushfwd(Base.BroadcastFunction(exp), StdNormal()^2), ()) isa Vector{Float64} + @test batched_rand_impl(ctx(), SpikeMixture(StdNormal(), 0.5), ()) isa Float64 + @test batched_rand_impl(ctx(), superpose(StdNormal(), StdUniform()), ()) isa Float64 + @test batched_rand_impl(ctx(), NoBatchRandMeasure(), ()) isa Float64 + @test batched_rand_impl(ctx(), NoBatchRandMeasure(), (4,)) isa Vector{Float64} + @test_throws ArgumentError rand(NoRandMeasure()) + @test_throws ArgumentError batched_rand_impl(ctx(), NoRandMeasure(), ()) + @test_throws ArgumentError batched_rand_impl(ctx(), NoRandMeasure(), (3,)) + end + + @testset "structured batches" begin + Pt = productmeasure((StdNormal(), StdExponential()^2)) + Xt = batched_rand_impl(ctx(), Pt, (5,)) + @test Xt isa Tuple && size(Xt[1]) == (5,) && size(Xt[2]) == (2, 5) + xt = batched_rand_impl(ctx(), Pt, ()) + @test xt isa Tuple{Float64,Vector{Float64}} && xt == rand(stblrng(), Pt) + XPt = rand(stblrng(), Pt^5) + @test XPt isa StructArray && length(XPt) == 5 + @test XPt[2] isa Tuple{Float64,<:AbstractVector{Float64}} && length(XPt[2][2]) == 2 + @test logdensityof(Pt^5, collect(XPt)) ≈ sum(logdensityof.(Ref(Pt), XPt)) + Pn = productmeasure((a = StdNormal(), b = StdExponential()^2)) + Xn = batched_rand_impl(ctx(), Pn, (2, 3)) + @test Xn isa NamedTuple{(:a, :b)} && size(Xn.a) == (2, 3) && size(Xn.b) == (2, 2, 3) + XPn = rand(stblrng(), Pn^4) + @test XPn isa StructArray && XPn[1] isa NamedTuple{(:a, :b)} + + mm = mcombine(merge, productmeasure((a = StdNormal(),)), productmeasure((b = StdUniform()^2,))) + Xm = batched_rand_impl(ctx(), mm, (3,)) + @test Xm isa NamedTuple{(:a, :b)} && size(Xm.a) == (3,) && size(Xm.b) == (2, 3) + @test batched_rand_impl(ctx(), mm, ()) == rand(stblrng(), mm) + mt = mcombine(tuple, StdNormal(), StdUniform()^2) + Xtt = batched_rand_impl(ctx(), mt, (3,)) + @test Xtt isa Tuple && size(Xtt[1]) == (3,) && size(Xtt[2]) == (2, 3) + @test batched_rand_impl(ctx(), mt, ()) == rand(stblrng(), mt) + end + + @testset "powers and value-dependent sizes" begin + f_β(a) = StdNormal()^length(a) + μb = mbind(f_β, StdUniform()^1, vcat) + x = rand(stblrng(), μb^3) + @test x isa Vector{Vector{Float64}} && length(x) == 3 && all(length.(x) .== 2) + X = batched_rand_impl(ctx(), μb^3, (4,)) + @test size(X) == (2, 3, 4) + @test batched_rand_impl(ctx(), μb, ()) == rand(stblrng(), μb) + @test size(batched_rand_impl(ctx(), (StdNormal()^2)^3, (4, 5))) == (2, 3, 4, 5) + @test size(batched_rand_impl(ctx(), StdNormal()^(2, 3), (4,))) == (2, 3, 4) + end + + @testset "moments of batches" begin + n = 20_000 + X = batched_rand_impl(ctx(), superpose(StdNormal(), Dirac(3.0)), (n,)) + @test isapprox(mean(X), 1.5, atol = 0.05) + Xs = batched_rand_impl(ctx(), SpikeMixture(StdNormal()^2, 0.5), (n,)) + @test size(Xs) == (2, n) && isapprox(mean(Xs .== 0), 0.5, atol = 0.02) + Xp = batched_rand_impl(ctx(), pushfwd(Base.BroadcastFunction(exp), StdNormal()^2), (n,)) + @test isapprox(mean(log.(Xp)), 0.0, atol = 0.03) + Xd = batched_rand_impl(ctx(), MeasureBase.AsMeasure{Normal{Float64}}(Normal(1.0, 2.0)), (n,)) + @test Xd isa Vector{Float64} && isapprox(mean(Xd), 1.0, atol = 0.05) + xd = batched_rand_impl(ctx(), MeasureBase.AsMeasure{Normal{Float64}}(Normal(1.0, 2.0)), ()) + @test xd isa Float64 + Xmv = batched_rand_impl(ctx(), MeasureBase.AsMeasure{typeof(MvNormal([1.0, 2.0], [1.0, 0.5]))}(MvNormal([1.0, 2.0], [1.0, 0.5])), (n,)) + @test size(Xmv) == (2, n) && isapprox(vec(mean(Xmv, dims = 2)), [1.0, 2.0], atol = 0.05) + end +end diff --git a/test/runtests.jl b/test/runtests.jl index c2f63c4e..e53acc85 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -9,18 +9,51 @@ 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("shape_contract.jl") +include("logdensities.jl") +include("structured_batches.jl") +include("batched_regressions.jl") +include("static_variates.jl") +include("fixed_size_arrays.jl") +include("support_conventions.jl") +include("numtype.jl") include("transport.jl") +include("transport_batched.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("combinators/product.jl") + +include("rand.jl") +include("rand_batched.jl") + +include("distributions/test_distributions.jl") + +# Reactant only supports 64-bit Linux and macOS, and some of its +# dependencies break already during precompilation on other platforms, +# so it can't be a static test dependency: +if Sys.WORD_SIZE == 64 && (Sys.islinux() || Sys.isapple()) && isempty(VERSION.prerelease) + import Pkg + Base.identify_package("Reactant") === nothing && Pkg.add("Reactant") + include("test_reactant.jl") +end include("test_docs.jl") diff --git a/test/shape_contract.jl b/test/shape_contract.jl new file mode 100644 index 00000000..6e34973b --- /dev/null +++ b/test/shape_contract.jl @@ -0,0 +1,106 @@ +using Test + +using MeasureBase +using MeasureBase: mspace_elsize, mspace_flatsize, NoMSpaceElementSize +using MeasureBase: preferred_stdmeasure, promote_stdmeasure, AnyStdMeasure, NoStdTransport +using MeasureBase: StdNormal, StdUniform, StdExponential, StdLogistic +using MeasureBase: Dirac, Lebesgue, Counting, LebesgueBase, CountingBase +using MeasureBase: mreshape, productmeasure, weightedmeasure, pushfwd, mbind, restrict +using IntervalSets: (..) +using StaticArrays: SVector, Size +using Static: static +using MeasureBase: size2length +using MeasureBase: setcartpower, ℝ, testvalue + +_flat_iter(x::Number) = (x,) +_flat_iter(x::AbstractArray) = Iterators.flatten(map(_flat_iter, x)) + +struct _CustomStd <: MeasureBase.StdMeasure end + +@testset "shape contract" begin + @testset "mspace_elsize and mspace_flatsize" begin + for μ in (StdNormal(), StdUniform(), Lebesgue(), Lebesgue(0..1), Counting(), LebesgueBase(), CountingBase(), Dirac(1.5)) + @test @inferred(mspace_elsize(μ)) === () + @test @inferred(mspace_flatsize(μ)) === () + end + + @test @inferred(mspace_elsize(StdNormal()^3)) == (3,) + @test @inferred(mspace_flatsize(StdNormal()^3)) == (3,) + @test @inferred(mspace_elsize(StdNormal()^(2, 3))) == (2, 3) + @test @inferred(mspace_flatsize(StdNormal()^(2, 3))) == (2, 3) + + @test @inferred(mspace_elsize((StdNormal()^3)^(2, 4))) == (2, 4) + @test @inferred(mspace_flatsize((StdNormal()^3)^(2, 4))) == (3, 2, 4) + @test @inferred(mspace_flatsize((StdNormal()^static(3))^static(2))) === Size(3, 2) + @test @inferred(mspace_flatsize((StdNormal()^static(3))^2)) === (static(3), 2) + + @test @inferred(mspace_elsize(Dirac([1, 2]))) == (2,) + @test @inferred(mspace_flatsize(Dirac([1, 2]))) == (2,) + @test @inferred(mspace_flatsize(Dirac(SVector(1, 2)))) === Size(2) + @test @inferred(mspace_elsize(Dirac([[1], [2]]))) == (2,) + @test @inferred(mspace_flatsize(Dirac([[1], [2]]))) isa NoMSpaceElementSize + @test @inferred(mspace_elsize(Dirac((a = 1, b = 2)))) isa NoMSpaceElementSize + + @test @inferred(mspace_elsize(weightedmeasure(0.3, StdNormal()^2))) == (2,) + @test @inferred(mspace_flatsize(weightedmeasure(0.3, (StdNormal()^2)^3))) == (2, 3) + @test @inferred(mspace_elsize(restrict(x -> x > 0, StdNormal()))) === () + @test @inferred(mspace_elsize(mreshape(StdNormal()^6, (2, 3)))) == (2, 3) + @test @inferred(mspace_flatsize(mreshape(StdNormal()^6, (2, 3)))) == (2, 3) + @test @inferred(mspace_elsize(mreshape((StdNormal()^2)^6, (2, 3)))) == (2, 3) + @test @inferred(mspace_flatsize(mreshape((StdNormal()^2)^6, (2, 3)))) isa NoMSpaceElementSize + + s = setcartpower(setcartpower(ℝ, 2), 3) + @test @inferred(mspace_elsize(Lebesgue(s))) == (3,) + @test @inferred(mspace_flatsize(Lebesgue(s))) == (2, 3) + for μ in (StdNormal(), StdNormal()^3, (StdNormal()^2)^3, Dirac([1.0, 2.0]), Dirac(3.0)) + @test size2length(mspace_flatsize(μ)) == length(vec(collect(Iterators.flatten(_flat_iter(testvalue(μ)))))) + end + + @test @inferred(mspace_elsize(productmeasure((a = StdNormal(), b = StdUniform())))) isa NoMSpaceElementSize + @test @inferred(mspace_flatsize(mbind(x -> StdNormal()^2, StdUniform()))) isa NoMSpaceElementSize + end + + @testset "preferred_stdmeasure" begin + for S in (StdNormal, StdUniform, StdExponential, StdLogistic) + @test @inferred(preferred_stdmeasure(S())) === S + @test @inferred(preferred_stdmeasure(S()^3)) === S + @test @inferred(preferred_stdmeasure(weightedmeasure(0.1, S()))) === S + @test @inferred(preferred_stdmeasure(pushfwd(exp, S()))) === S + end + # Transports of the base measure don't transport restricted measures: + @test @inferred(preferred_stdmeasure(restrict(x -> x > 0, StdNormal()))) <: NoStdTransport + @test @inferred(preferred_stdmeasure(MeasureBase.Half(StdNormal()))) === StdUniform + + @test @inferred(preferred_stdmeasure(Dirac(2.0))) === AnyStdMeasure + @test @inferred(preferred_stdmeasure(Lebesgue())) <: NoStdTransport + @test @inferred(preferred_stdmeasure(Counting())) <: NoStdTransport + + @test @inferred(preferred_stdmeasure(productmeasure((StdUniform(), StdNormal())))) === StdNormal + @test @inferred(preferred_stdmeasure(productmeasure((a = StdUniform(), b = StdExponential())))) === StdExponential + @test @inferred(preferred_stdmeasure(productmeasure((a = Dirac(1.0), b = StdUniform())))) === StdUniform + @test @inferred(preferred_stdmeasure(productmeasure((a = Lebesgue(), b = StdUniform())))) <: NoStdTransport + @test @inferred(preferred_stdmeasure(productmeasure(fill(StdLogistic(), 3)))) === StdLogistic + @test @inferred(preferred_stdmeasure(productmeasure(()))) === AnyStdMeasure + @test @inferred(preferred_stdmeasure(mbind(x -> StdNormal()^2, StdUniform()))) === StdUniform + end + + @testset "promote_stdmeasure" begin + @test @inferred(promote_stdmeasure(StdUniform, StdNormal)) === StdNormal + @test @inferred(promote_stdmeasure(StdNormal, StdUniform)) === StdNormal + @test @inferred(promote_stdmeasure(StdUniform, StdExponential)) === StdExponential + @test @inferred(promote_stdmeasure(StdExponential, StdLogistic)) === StdLogistic + @test @inferred(promote_stdmeasure(StdLogistic, StdNormal)) === StdNormal + @test @inferred(promote_stdmeasure(StdLogistic, StdLogistic)) === StdLogistic + @test @inferred(promote_stdmeasure(AnyStdMeasure, StdUniform)) === StdUniform + @test @inferred(promote_stdmeasure(StdUniform, AnyStdMeasure)) === StdUniform + @test @inferred(promote_stdmeasure(AnyStdMeasure, AnyStdMeasure)) === AnyStdMeasure + @test @inferred(promote_stdmeasure(NoStdTransport{Int}, StdNormal)) === NoStdTransport{Int} + @test @inferred(promote_stdmeasure(StdNormal, NoStdTransport{Int})) === NoStdTransport{Int} + @test @inferred(promote_stdmeasure(NoStdTransport{Int}, AnyStdMeasure)) === NoStdTransport{Int} + @test @inferred(promote_stdmeasure(AnyStdMeasure, NoStdTransport{Int})) === NoStdTransport{Int} + @test @inferred(promote_stdmeasure(StdUniform, StdExponential, AnyStdMeasure, StdLogistic)) === StdLogistic + @test @inferred(promote_stdmeasure(_CustomStd, StdUniform)) === StdUniform + @test @inferred(promote_stdmeasure(_CustomStd, AnyStdMeasure)) === _CustomStd + @test @inferred(preferred_stdmeasure(_CustomStd())) === _CustomStd + end +end 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/static_variates.jl b/test/static_variates.jl new file mode 100644 index 00000000..453504cd --- /dev/null +++ b/test/static_variates.jl @@ -0,0 +1,174 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +# Static variates end to end: measures whose variate sizes are statically +# known generate, transport and evaluate static arrays, type stable and +# allocation free. + +using Test + +using MeasureBase +using MeasureBase: StdNormal, StdUniform, StdExponential +using MeasureBase: productmeasure, mbind, mcombine, weightedmeasure, transport_to, logdensityof +using MeasureBase: batched_logdensityof_with_rest, batched_transport_to_std, + batched_transport_from_std, batched_transport_to_std_with_rest, + batched_transport_from_std_with_rest, _materialize +using MeasureBase.InverseFunctions: inverse +using ArraysOfArrays: flatview, sliced +using StaticArrays: SVector, SMatrix, Size +using Static: static + +include("testutils.jl") + +# A hierarchical model over a named tuple of scalar and static-array +# marginals: the secondary marginals depend on the primary variate, their +# sizes don't. +const static_primary = productmeasure(( + a = StdNormal(), + b = weightedmeasure(-0.5, StdExponential()), +)) + +static_kernel(x) = productmeasure(( + c = StdUniform()^static(2), + d = weightedmeasure(-abs(x.a), StdNormal()^static(3)), +)) + +const static_model = mbind(static_kernel, static_primary, merge) + +@testset "static variates" begin + @testset "static powers" begin + @test @inferred(rand(StdNormal()^static(3))) isa SVector{3,Float64} + @test @inferred(rand(StdUniform()^static(2))) isa SVector{2,Float64} + @test @inferred(rand(StdExponential()^static(4))) isa SVector{4,Float64} + + @test allocations_of(rand, StdNormal()^static(3)) == 0 + + # Nested powers keep the flat `(base dims..., power dims...)` rule: + x = @inferred(rand((StdNormal()^static(2))^static(3))) + @test flatview(x) isa SMatrix{2,3,Float64} + @test length(x) == 3 && all(xi -> xi isa SVector{2,Float64}, x) + @test reduce(hcat, x) == flatview(x) + + μ = StdNormal()^static(3) + z = SVector(0.1, 0.2, 0.3) + @test @inferred(transport_to(StdUniform()^static(3), μ)(z)) isa SVector{3,Float64} + @test @inferred(transport_to(μ, StdUniform()^static(3))(SVector(0.1, 0.5, 0.9))) isa + SVector{3,Float64} + @test @inferred(logdensityof(μ, z)) ≈ sum(logdensityof.(Ref(StdNormal()), z)) + @test allocations_of(logdensityof, μ, z) == 0 + @test allocations_of(transport_to(StdUniform()^static(3), μ), z) == 0 + end + + @testset "hierarchical model with static marginals" begin + μ = static_model + x = @inferred rand(μ) + @test x isa NamedTuple{(:a, :b, :c, :d)} + @test x.a isa Float64 + @test x.b isa Float64 + @test x.c isa SVector{2,Float64} + @test x.d isa SVector{3,Float64} + + ν = StdNormal()^static(7) + f = transport_to(ν, μ) + f_inv = inverse(f) + + z = @inferred f(x) + @test z isa SVector{7,Float64} + y = @inferred f_inv(z) + @test y isa NamedTuple{(:a, :b, :c, :d)} + @test y.c isa SVector{2,Float64} + @test y.d isa SVector{3,Float64} + @test all(map((u, v) -> u ≈ v, values(y), values(x))) + + @test allocations_of(f, x) == 0 + @test allocations_of(f_inv, z) == 0 + @test allocations_of(logdensityof, μ, x) == 0 + @test allocations_of(rand, μ) == 0 + + # The model density is the sum of the component densities: + x_a = (a = x.a, b = x.b) + x_b = (c = x.c, d = x.d) + @test @inferred(logdensityof(μ, x)) ≈ + logdensityof(static_primary, x_a) + logdensityof(static_kernel(x_a), x_b) + end + + @testset "products of scalar and static-array marginals" begin + μ = productmeasure(( + a = StdNormal(), + b = StdUniform()^static(2), + c = weightedmeasure(-0.25, StdExponential()^static(3)), + )) + x = @inferred rand(μ) + @test x isa NamedTuple{(:a, :b, :c)} + @test x.a isa Float64 + @test x.b isa SVector{2,Float64} + @test x.c isa SVector{3,Float64} + + f = transport_to(StdNormal()^static(6), μ) + f_inv = inverse(f) + z = @inferred f(x) + @test z isa SVector{6,Float64} + y = @inferred f_inv(z) + @test y isa NamedTuple{(:a, :b, :c)} + @test all(map((u, v) -> u ≈ v, values(y), values(x))) + @test allocations_of(f, x) == 0 + @test allocations_of(f_inv, z) == 0 + @test allocations_of(logdensityof, μ, x) == 0 + @test allocations_of(rand, μ) == 0 + @test allocations_of(rand, productmeasure((a = StdNormal(), b = StdExponential()))) == 0 + + # Tuple products behave the same way: + μ_t = productmeasure((StdNormal(), StdUniform()^static(2))) + x_t = @inferred rand(μ_t) + @test x_t isa Tuple{Float64,SVector{2,Float64}} + f_t = transport_to(StdNormal()^static(3), μ_t) + @test @inferred(f_t(x_t)) isa SVector{3,Float64} + @test @inferred(inverse(f_t)(f_t(x_t))) isa Tuple{Float64,SVector{2,Float64}} + end + + @testset "batches of static variates" begin + μ = StdNormal()^static(3) + ν = StdUniform()^static(3) + X = SMatrix{3,4}(reshape(collect(1:12) ./ 10, 3, 4)) + + @test @inferred(logdensities(μ, X)) ≈ + [logdensityof(μ, SVector{3}(X[:, i])) for i in 1:4] + @test allocations_of(logdensities, μ, X) == 0 + + Z = @inferred batched_transport_to_std(StdUniform, μ, X) + @test Z ≈ reduce(hcat, [transport_to(ν, μ)(SVector{3}(X[:, i])) for i in 1:4]) + @test @inferred(batched_transport_from_std(StdUniform, μ, Z)) ≈ X + + # The broadcast hook transports the whole batch at once: + Y = transport_to(ν, μ).(sliced(X, Val(1))) + @test flatview(Y) ≈ Z + @test Y[2] ≈ transport_to(ν, μ)(SVector{3}(X[:, 2])) + end + + # Several variates per stream, with the multiplicity as a tuple of + # static integers and as a `StaticArrays.Size`: + @testset "static stream multiplicity" begin + μ = StdNormal()^static(2) + mc = mcombine(vcat, StdNormal()^static(2), StdUniform()^static(3)) + x = SVector{4}(randn(4)) + xc = SVector{10}(vcat(randn(2), rand(3), randn(2), rand(3))) + to_u = transport_to(StdUniform(), StdNormal()) + + for sz in ((static(2),), Size(2)) + z, x_rest = batched_transport_to_std_with_rest(StdUniform, μ, x, sz) + @test z isa SVector{4,Float64} && isempty(x_rest) + @test z ≈ to_u.(x) + x_back, z_rest = batched_transport_from_std_with_rest(StdUniform, μ, z, sz) + @test x_back ≈ reshape(x, (2, 2)) && size(z_rest, 1) == 0 + + ℓ, x_ld_rest = batched_logdensityof_with_rest(μ, x, sz) + @test _materialize(ℓ) isa SVector{2,Float64} && isempty(x_ld_rest) + @test _materialize(ℓ) ≈ [logdensityof(μ, x[(2i - 1):(2i)]) for i in 1:2] + + zc, xc_rest = batched_transport_to_std_with_rest(StdUniform, mc, xc, sz) + @test length(zc) == 10 && isempty(xc_rest) + ℓc, xc_ld_rest = batched_logdensityof_with_rest(mc, xc, sz) + @test isempty(xc_ld_rest) + @test _materialize(ℓc) ≈ [logdensityof(mc, xc[(5i - 4):(5i)]) for i in 1:2] + end + end +end diff --git a/test/structured_batches.jl b/test/structured_batches.jl new file mode 100644 index 00000000..184f01da --- /dev/null +++ b/test/structured_batches.jl @@ -0,0 +1,45 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using MeasureBase +using MeasureBase: StdNormal, StdUniform, StdExponential, productmeasure, insupport +using MeasureBase.InverseFunctions: inverse +using StructArrays: StructArray +using ArraysOfArrays: flatview + +@testset "structured batches" begin + Pt = productmeasure((StdNormal(), StdExponential()^2)) + xs = [rand(Pt) for _ in 1:6] + ℓ_ref = logdensityof.(Ref(Pt), xs) + @test logdensities(Pt, xs) ≈ ℓ_ref + @test logdensities(Pt, StructArray(xs)) ≈ ℓ_ref + X = rand(Pt^6) + @test X isa StructArray + @test logdensities(Pt, X) ≈ logdensityof.(Ref(Pt), X) + @test logdensityof(Pt^6, X) ≈ sum(logdensityof.(Ref(Pt), X)) + @test logdensityof(Pt^6, collect(X)) ≈ logdensityof(Pt^6, X) + Xm = rand(Pt^(2, 3)) + @test size(Xm) == (2, 3) && logdensityof(Pt^(2, 3), Xm) ≈ sum(logdensityof.(Ref(Pt), Xm)) + @test_throws ArgumentError logdensityof(Pt^5, X) + + Pn = productmeasure((a = StdNormal(), b = StdExponential()^2)) + Xn = rand(Pn^5) + @test logdensities(Pn, Xn) ≈ logdensityof.(Ref(Pn), Xn) + @test logdensities(Pn, collect(Xn)) ≈ logdensityof.(Ref(Pn), Xn) + + f = transport_to(StdUniform()^3, Pt) + Y = f.(X) + @test Y isa AbstractVector && length(Y) == 6 && flatview(Y) ≈ stack(map(f, X)) + Xr = inverse(f).(Y) + @test Xr isa StructArray && all(map((a, b) -> all(map(≈, a, b)), Xr, X)) + h = transport_to(Pn, Pt) + Yn = h.(X) + @test Yn isa StructArray && Yn[1] isa NamedTuple{(:a, :b)} + @test all(Yn[i].a ≈ X[i][1] && Yn[i].b ≈ X[i][2] for i in 1:6) + @test inverse(h).(Yn) isa StructArray + + @test insupport(StdUniform()^3, [0.1, 0.5, 0.9]) && !insupport(StdUniform()^3, [0.1, 1.5, 0.9]) + @test insupport((StdUniform()^2)^3, rand(2, 3)) && !insupport((StdUniform()^2)^3, fill(2.0, 2, 3)) + @test insupport(StdUniform()^3, [0.1, 0.5, 0.9]) isa Bool +end diff --git a/test/support_conventions.jl b/test/support_conventions.jl new file mode 100644 index 00000000..72b36837 --- /dev/null +++ b/test/support_conventions.jl @@ -0,0 +1,68 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using MeasureBase +using MeasureBase: StdNormal, StdUniform, StdExponential, StdLogistic, Half, logdensities +using MeasureBase: weightedmeasure, productmeasure, mcombine +using JLArrays: JLArray + +# Variates of the right shape never throw: densities are -Inf and +# transports are NaN outside the support, wrong shapes throw. +@testset "support conventions" begin + outside = Dict(StdUniform() => -0.5, StdExponential() => -1.0, Half(StdNormal()) => -1.0) + + @testset "densities" begin + for (μ, x) in outside + @test logdensityof(μ, x) == -Inf + @test logdensityof(weightedmeasure(0.3, μ), x) == -Inf + X = [x 0.5 0.5; 0.5 x 0.5; 0.5 0.5 0.5] + ℓ = logdensities(μ^3, X) + @test ℓ[1:2] == [-Inf, -Inf] && isfinite(ℓ[3]) + @test Array(logdensities(μ^3, JLArray(X))) == ℓ + @test Array(logdensities(μ, JLArray(vec(X)))) == logdensities(μ, vec(X)) + end + @test logdensityof(StdUniform(), Inf) == -Inf + @test logdensityof(StdExponential(), Inf) == -Inf + @test logdensityof(StdNormal(), -Inf) == -Inf + @test logdensityof(productmeasure((StdUniform(), StdExponential())), (0.5, -1.0)) == -Inf + @test logdensityof(mcombine(vcat, StdUniform()^2, StdExponential()^1), [0.5, 1.5, 0.5]) == -Inf + end + + @testset "transports" begin + stds = (StdUniform(), StdExponential(), StdLogistic(), StdNormal()) + for (μ, x) in outside, ν in stds + μ === ν && continue + f = transport_to(ν, μ) + @test isnan(f(x)) + @test isnan(transport_to(ν^2, μ^2)([x, 0.5])[1]) + Y = f.([x, 0.5, 0.5]) + @test isnan(Y[1]) && !isnan(Y[2]) + @test isequal(Array(f.(JLArray([x, 0.5, 0.5]))), Y) + @test isnan(transport_to(μ, ν)(NaN)) + end + for ν in (StdExponential(), StdLogistic(), StdNormal(), Half(StdNormal())) + @test isnan(transport_to(ν, StdUniform())(1.5)) + @test isnan(transport_to(ν, StdUniform())(-0.5)) + end + # Endpoints of the unit interval stand for their nearest interior + # points, tails never underflow to infinite variates: + for ν in (StdExponential(), StdLogistic(), StdNormal(), Half(StdNormal())) + f = transport_to(ν, StdUniform()) + @test isfinite(f(0.0)) && isfinite(f(1.0)) && f(0.0) <= f(0.5) <= f(1.0) + @test f(1.0) == f(prevfloat(1.0)) && f(0.0) == f(floatmin(Float64)) + end + for (ν, μ) in ((StdExponential(), StdNormal()), (StdNormal(), StdExponential()), (StdNormal(), StdLogistic())) + @test all(isfinite, transport_to(ν, μ).([-1e6, -40.0, 40.0, 1e6][MeasureBase.insupport.(Ref(μ), [-1e6, -40.0, 40.0, 1e6])])) + end + @test !isnan(transport_to(StdUniform(), StdNormal())(-37.0)) + @test !isnan(transport_to(StdUniform(), StdLogistic())(-800.0)) + end + + @testset "wrong shapes throw" begin + @test_throws ArgumentError logdensityof(StdNormal()^3, randn(2)) + @test_throws ArgumentError logdensities(StdNormal()^3, randn(2, 5)) + @test_throws ArgumentError transport_to(StdUniform()^3, StdNormal()^3)(randn(2)) + @test_throws ArgumentError logdensityof(StdNormal(), randn(2)) + end +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..71540e22 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,39 @@ 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 "hash follows equality" begin + a = productmeasure([Dirac([1.0, 2.0]), StdNormal()]) + b = productmeasure([Dirac([1.0, 2.0]), StdNormal()]) + @test a == b && hash(a) == hash(b) + f, g = transport_to(StdUniform()^1, a), transport_to(StdUniform()^1, b) + @test f == g && hash(f) == hash(g) +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 +250,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..7d10fb2c --- /dev/null +++ b/test/test_mooncake.jl @@ -0,0 +1,53 @@ +# 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 + +_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) + 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 diff --git a/test/test_reactant.jl b/test/test_reactant.jl new file mode 100644 index 00000000..4df1f7db --- /dev/null +++ b/test/test_reactant.jl @@ -0,0 +1,140 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +# Reactant tests. Reactant isn't a static test dependency (it only +# supports 64-bit Linux and macOS), runtests.jl adds it on the fly where +# supported. The backend defaults to the CPU, set the environment variable +# `MEASUREBASE_REACTANT_BACKEND` (e.g. to "gpu") to change it; the file can +# also be run standalone in an environment that provides Reactant. + +using Test +using Reactant +using MeasureBase +using MeasureBase: StdNormal, StdUniform, StdExponential, StdLogistic, Lebesgue, Dirac, asmeasure +using MeasureBase: logdensities, logdensity_rel, weightedmeasure, superpose, restrict, mintegrate_exp +using MeasureBase: mcombine +using ArraysOfArrays: VectorOfSimilarVectors, sliced, flatview +using Distributions: Normal, Uniform, Exponential, Logistic, Cauchy, Laplace, LogNormal, Weibull, Gamma, Beta +using Distributions: Poisson, MvNormal, Dirichlet +using MeasureBase.InverseFunctions: inverse + +Reactant.set_default_backend(get(ENV, "MEASUREBASE_REACTANT_BACKEND", "cpu")) + +# Compiles `f` for traced copies of `args` and compares with the plain +# result. Array results are copied inside the compiled function, so that +# views and reshapes of device arrays come back as plain device arrays: +function test_traced(f, args...; kwargs...) + expected = f(args...) + traced_args = map(Reactant.to_rarray, args) + g = (xs...) -> _contiguous(f(xs...)) + result = @jit g(traced_args...) + @test _plain(result) ≈ _plain(expected) nans = true + return result +end + +_contiguous(x::AbstractArray) = copy(x) +_contiguous(x) = x +_plain(x::AbstractArray) = Array(x) +_plain(x::Number) = Float64(x) + +@testset "Reactant" begin + x = randn(10) + X = randn(3, 20) + + @testset "powers and batches" begin + test_traced(x -> logdensityof(StdNormal()^10, x), x) + test_traced(X -> logdensities(StdNormal(), X), X) + test_traced(X -> logdensities(StdNormal()^3, X), X) + test_traced(X -> logdensities(StdNormal()^3, sliced(X, 1)), X) + test_traced(X -> logdensities((StdNormal()^3)^4, reshape(X[:, 1:16], 3, 4, 4)), X) + test_traced(x -> logdensityof(StdUniform()^10, x), rand(10)) + test_traced(x -> logdensityof(StdExponential()^10, x), rand(10)) + test_traced(x -> logdensityof(weightedmeasure(0.3, StdNormal())^10, x), x) + test_traced(x -> logdensityof(Lebesgue()^10, x), x) + end + + @testset "support masks" begin + xu = 2 .* rand(10) .- 0.5 + test_traced(x -> logdensities(StdUniform(), x), xu) + test_traced(x -> logdensities(StdExponential(), x), xu) + test_traced(x -> logdensities(restrict(x -> x > 0, StdNormal()), x), xu) + end + + @testset "relative densities" begin + xu = 2 .* rand(10) .- 0.5 + test_traced(x -> logdensity_rel.(Ref(StdUniform()), Ref(StdExponential()), x), xu) + test_traced(x -> logdensity_rel.(Ref(StdNormal()), Ref(StdLogistic()), x), x) + test_traced(x -> logdensity_rel.(Ref(StdNormal()^10), Ref(StdLogistic()^10), Ref(x)), x) + end + + @testset "superposition, density measures and spike mixtures" begin + mix = superpose(weightedmeasure(log(0.3), StdNormal()), weightedmeasure(log(0.7), StdLogistic())) + test_traced(x -> logdensities(mix, x), x) + test_traced(x -> logdensityof(mix^10, x), x) + dm = mintegrate_exp(x -> -abs(x), StdNormal()) + test_traced(x -> logdensities(dm, x), x) + sm = SpikeMixture(StdNormal(), 0.2) + test_traced(x -> logdensities(sm, x), vcat(x, 0.0)) + end + + # Products over arrays of marginals are not covered: Reactant can't + # broadcast over arrays of measures together with traced arrays. + @testset "structural batched kernels" begin + w = weightedmeasure(log(0.3), StdNormal()^3) + test_traced(X -> logdensities(w, X), X) + mc = mcombine(vcat, StdNormal()^2, StdUniform()^1) + test_traced(X -> logdensities(mc, X), vcat(X[1:2, :], rand(1, 20))) + test_traced(X -> logdensities((StdNormal()^2)^3, X), reshape(X[1:2, 1:6], 2, 3, 2)) + end + + @testset "transport of powers and products" begin + test_traced(x -> transport_to(StdNormal()^10, StdUniform()^10)(x), rand(10)) + test_traced(x -> transport_to(StdExponential()^10, StdNormal()^10)(x), x) + test_traced(x -> transport_to(StdLogistic()^10, StdExponential()^10)(x), rand(10)) + test_traced(x -> transport_to(StdNormal()^6, (StdUniform()^2)^3)(x), rand(2, 3)) + mc = mcombine(vcat, StdNormal()^2, StdUniform()^1) + test_traced(x -> transport_to(StdNormal()^3, mc)(x), vcat(randn(2), rand(1))) + test_traced(z -> transport_to(mc, StdNormal()^3)(z), randn(3)) + end + + @testset "batched transport" begin + test_traced(X -> transport_to(StdNormal(), StdUniform()).(X), rand(10)) + test_traced(X -> flatview(transport_to(StdExponential()^3, StdNormal()^3).(sliced(X, Val(1)))), X) + mc = mcombine(vcat, StdNormal()^2, StdUniform()^1) + test_traced(X -> flatview(transport_to(StdLogistic()^3, mc).(sliced(X, Val(1)))), vcat(X[1:2, :], rand(1, 20))) + end + + @testset "transport" begin + test_traced(x -> transport_to(StdUniform(), StdNormal()).(x), x) + test_traced(x -> transport_to(StdNormal(), StdUniform()).(x), rand(10)) + test_traced(x -> transport_to(Normal(2, 3), StdNormal()).(x), x) + test_traced(x -> transport_to(StdNormal(), Exponential(2.0)).(x), rand(10)) + end + + # Distribution parameters stay constants, Distributions' parameter + # structs can't hold traced arrays: + @testset "wrapped distributions" begin + for d in (Normal(0.3, 1.7), Uniform(-1.0, 2.5), Exponential(0.7), Logistic(0.2, 1.3), Cauchy(0.1, 0.8), Laplace(-0.4, 1.1), LogNormal(0.2, 0.6), Weibull(1.4, 0.9), Gamma(2.3, 1.2), Beta(2.5, 3.5)) + m = asmeasure(d) + xd = rand(d, 10) + test_traced(x -> logdensities(m, x), xd) + f = transport_to(StdNormal(), m) + if d isa Union{Gamma,Beta} + # SpecialFunctions' incomplete gamma and beta functions have no Reactant methods: + @test_broken @jit((x -> copy(f.(x)))(Reactant.to_rarray(xd))) isa AbstractArray + else + test_traced(x -> f.(x), xd) + test_traced(z -> inverse(f).(z), randn(10)) + end + end + test_traced(x -> logdensities(asmeasure(Poisson(2.7)), x), Float64.(rand(Poisson(2.7), 10))) + mvn = MvNormal([0.3, -2.9], [1.7 0.5; 0.5 2.3]) + mm = asmeasure(mvn) + Xm = rand(mvn, 10) + test_traced(X -> logdensities(mm, X), Xm) + test_traced(X -> flatview(transport_to(StdNormal()^2, mm).(sliced(X, Val(1)))), Xm) + test_traced(Z -> flatview(transport_to(mm, StdNormal()^2).(sliced(Z, Val(1)))), randn(2, 10)) + dir = Dirichlet([2.0, 3.0, 4.0, 1.5]) + md = asmeasure(dir) + test_traced(X -> logdensities(md, X), rand(dir, 10)) + end +end diff --git a/test/testutils.jl b/test/testutils.jl new file mode 100644 index 00000000..5fa9c86f --- /dev/null +++ b/test/testutils.jl @@ -0,0 +1,9 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +# Allocations of `f(args...)` after a warm-up call, measured inside a +# function since `@allocated` at top level reports a boxed result on +# Julia 1.10: +function allocations_of(f::F, args::Vararg{Any,N}) where {F,N} + f(args...) + @allocated f(args...) +end diff --git a/test/transport.jl b/test/transport.jl index 34359784..24734fa3 100644 --- a/test/transport.jl +++ b/test/transport.jl @@ -1,19 +1,22 @@ using Test +import MeasureBase using MeasureBase.Interface: transport_to, test_transport +using MeasureBase: AbstractMeasure using MeasureBase: StdUniform, StdExponential, StdLogistic, StdNormal -using MeasureBase: Dirac +using MeasureBase: Dirac, Half, restrict, mbind, productmeasure, pushfwd +using MeasureBase: transport_to_std, transport_from_std, transport_from_std_with_rest +using InverseFunctions: inverse +using MeasureBase: weightedmeasure, mcombine +using StaticArrays: SVector +using Static: static using LogExpFunctions: logit +using ArraysOfArrays: sliced, flatview, fused +using JLArrays -using ChainRulesTestUtils +include("testutils.jl") @testset "transport_to" begin - test_rrule( - MeasureBase._origin_depth, - pushfwd(exp, StdUniform()), - output_tangent = static(0), - ) - for (f, μ) in [ (logit, StdUniform()) (log, StdExponential()) @@ -60,6 +63,173 @@ using ChainRulesTestUtils transport_to(StdUniform()^(2, 3), StdExponential()^6) end + # Tail accuracy of transports between standard measures: + + @testset "transports between standard measures" begin + stds = (StdUniform(), StdExponential(), StdLogistic(), StdNormal()) + for ν in stds, μ in stds + f = transport_to(ν, μ) + for x in [rand(μ) for _ in 1:5] + @test inverse(f)(f(x)) ≈ x + end + end + + # Round trips between the unbounded standard measures keep the tails: + for z in (-37.0, -20.0, -8.0, -6.0, 6.0, 8.0, 20.0, 37.0) + for ν in (StdExponential(), StdLogistic()) + y = transport_to(ν, StdNormal())(z) + @test isfinite(y) + @test transport_to(StdNormal(), ν)(y) ≈ z rtol = 1e-8 + end + end + for l in (-700.0, -40.0, -8.0, 8.0, 40.0, 700.0) + y = transport_to(StdExponential(), StdLogistic())(l) + @test isfinite(y) && y >= 0 + @test transport_to(StdLogistic(), StdExponential())(y) ≈ l rtol = 1e-8 + end + # The lower tail survives a uniform pivot, the upper tail saturates: + @test transport_to(StdNormal(), StdUniform())(transport_to(StdUniform(), StdNormal())(-37.0)) ≈ -37.0 rtol = 1e-8 + end + + @testset "scalar and static transports" begin + f = transport_to(StdNormal(), StdUniform()) + @test @inferred(f(0.3)) isa Float64 + @test allocations_of(f, 0.3) == 0 + g = transport_to(StdExponential()^static(3), StdNormal()^static(3)) + xs = SVector(0.1, -0.4, 2.0) + @test @inferred(g(xs)) isa SVector{3,Float64} + @test allocations_of(g, xs) == 0 + @test inverse(g)(g(xs)) ≈ xs + h = transport_to(StdNormal()^3, StdUniform()^3) + @test h(Float32[0.1, 0.5, 0.9]) isa Vector{Float32} + end + + @testset "nested powers" begin + μ = (StdNormal()^2)^3 + x = rand(μ) + f = transport_to(StdUniform()^6, μ) + y = f(x) + @test y isa AbstractVector{<:Real} && length(y) == 6 + x_reco = inverse(f)(y) + @test all(map(≈, x_reco, x)) + test_transport(StdExponential()^(3, 2), μ) + end + + @testset "powers of measures without fast DOF" begin + f_β(a) = StdNormal()^length(a) + μ = mbind(f_β, StdUniform()^1, vcat) + P = μ^2 + x = [rand(μ), rand(μ)] + z = transport_to(StdUniform()^4, P)(x) + @test z isa AbstractVector{<:Real} && length(z) == 4 + x_reco = transport_to(P, StdUniform()^4)(z) + @test x_reco isa AbstractVector && all(map(≈, x_reco, x)) + end + + @testset "Half" begin + μ = Half(StdNormal()) + test_transport(StdUniform(), μ) + test_transport(StdLogistic(), μ) + test_transport(μ, StdNormal()) + @test transport_to(StdUniform(), μ)(0.0) ≈ 0 + end + + @testset "array products of mixed standard measures" begin + src = productmeasure([StdNormal(), StdExponential()]) + trg = productmeasure([StdUniform(), StdLogistic()]) + @test MeasureBase.preferred_stdmeasure(src) === StdNormal + f = transport_to(trg, src) + x = [0.5, 1.0] + y = f(x) + @test y ≈ [transport_to(StdUniform(), StdNormal())(0.5), transport_to(StdLogistic(), StdExponential())(1.0)] + @test inverse(f)(y) ≈ x + @test_throws ArgumentError transport_to_std(MeasureBase.StdMeasure, StdNormal(), 0.5) + + pm = productmeasure(AbstractMeasure[StdNormal(), StdNormal()^2, Dirac(1.0)]) + xm = [0.5, [0.1, 0.2], 1.0] + z = transport_to(StdUniform()^3, pm)(xm) + @test z isa AbstractVector{<:Real} && length(z) == 3 + xm_reco = transport_to(pm, StdUniform()^3)(z) + @test xm_reco[1] ≈ xm[1] && xm_reco[2] ≈ xm[2] && xm_reco[3] == 1.0 + end + + @testset "measures without standard transport" begin + μ = restrict(x -> x > 0, StdNormal()) + @test_throws ArgumentError transport_to(StdUniform(), μ)(0.5) + @test_throws ArgumentError transport_to(μ, StdUniform())(0.5) + end + + @testset "batched transport" begin + f = transport_to(StdNormal(), StdUniform()) + X = rand(7) + @test f.(X) ≈ map(f, X) + @test inverse(f).(f.(X)) ≈ X + @test eltype(f.(rand(Float32, 5))) == Float32 + + g = transport_to(StdExponential()^3, StdNormal()^3) + Xn = randn(3, 5) + Yn = g.(sliced(Xn, Val(1))) + @test Yn isa AbstractVector && length(Yn) == 5 + @test flatview(Yn) ≈ stack(map(g, eachcol(Xn))) + @test flatview(g.(Xn)) ≈ flatview(Yn) + @test flatview(inverse(g).(Yn)) ≈ Xn + Xv = [randn(3) for _ in 1:4] + @test g.(Xv) == map(g, Xv) + + h = transport_to(StdUniform()^(2, 3), (StdNormal()^2)^3) + Xh = randn(2, 3, 4) + @test flatview(h.(Xh)) ≈ stack([h(Xh[:, :, i]) for i in 1:4]) + @test flatview(fused(inverse(h).(h.(Xh)))) ≈ Xh + Yh = h.(Xh) + Xh_reco = inverse(h).(Yh) + @test Xh_reco[2] == inverse(h)(Yh[2]) + @test Xh_reco[2] isa AbstractVector && length(Xh_reco[2]) == 3 && Xh_reco[2][1] isa AbstractVector + + X3 = randn(3, 4, 5) + Y3 = g.(X3) + @test size(Y3) == (4, 5) && size(flatview(Y3)) == (3, 4, 5) + @test Y3[2, 3] ≈ g(X3[:, 2, 3]) + @test f.(SVector(0.3, 0.6, 0.9)) isa SVector{3,Float64} + @test g.(Xn .+ 0.0) == g.(Xn) + + P = MeasureBase.ProductMeasure([weightedmeasure(log(i), StdNormal()) for i in 1:3]) + p = transport_to(StdUniform()^3, P) + Xp = randn(3, 6) + Yp = p.(sliced(Xp, Val(1))) + @test flatview(Yp) ≈ stack(map(p, eachcol(Xp))) + @test flatview(inverse(p).(Yp)) ≈ Xp + + mc = mcombine(vcat, StdNormal()^2, StdUniform()^3) + c = transport_to(StdExponential()^5, mc) + Xc = vcat(randn(2, 4), rand(3, 4)) + Yc = c.(sliced(Xc, Val(1))) + @test flatview(Yc) ≈ stack(map(c, eachcol(Xc))) + @test flatview(inverse(c).(Yc)) ≈ Xc + cd = transport_to(mcombine(vcat, Dirac(0.5), StdUniform()^2), StdNormal()^2) + @test flatview(cd.(randn(2, 3)))[1, :] == fill(0.5, 3) + + pf = transport_to(StdUniform(), pushfwd(exp, StdNormal())) + Xe = exp.(randn(8)) + @test pf.(Xe) ≈ map(pf, Xe) + @test inverse(pf).(pf.(Xe)) ≈ Xe + + w = transport_to(StdLogistic()^2, weightedmeasure(0.3, StdNormal()^2)) + Xw = randn(2, 5) + @test flatview(w.(sliced(Xw, Val(1)))) ≈ stack(map(w, eachcol(Xw))) + + JLArrays.allowscalar(false) + Xj = JLArray(Xn) + Yj = g.(sliced(Xj, Val(1))) + @test flatview(Yj) isa JLArray + @test Array(flatview(Yj)) ≈ flatview(Yn) + @test Array(flatview(c.(sliced(JLArray(Xc), Val(1))))) ≈ flatview(Yc) + Pj = MeasureBase.ProductMeasure(JLArray([weightedmeasure(log(i), StdNormal()) for i in 1:3])) + pj = transport_to(StdUniform()^3, Pj) + @test Array(flatview(pj.(sliced(JLArray(Xp), Val(1))))) ≈ flatview(Yp) + Yej = pf.(JLArray(Xe)) + @test Yej isa JLArray && Array(Yej) ≈ pf.(Xe) + end + @testset "transport for products" begin test_transport( StdUniform()^(2, 2), diff --git a/test/transport_batched.jl b/test/transport_batched.jl new file mode 100644 index 00000000..db91ef5a --- /dev/null +++ b/test/transport_batched.jl @@ -0,0 +1,171 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using MeasureBase +using MeasureBase: StdUniform, StdExponential, StdLogistic, StdNormal, Dirac, GenContext +using MeasureBase: productmeasure, pushfwd, mcombine, weightedmeasure, mbind, PushfwdRootMeasure +using MeasureBase: transport_to_std, transport_from_std, transport_to_std_with_rest +using MeasureBase: batched_transport_to_std, batched_transport_from_std +using MeasureBase: batched_transport_to_std_with_rest, batched_transport_from_std_with_rest +using MeasureBase: batched_rand_impl +using InverseFunctions: inverse +using ArraysOfArrays: sliced, flatview +using StaticArrays: SVector +using Distributions: MvNormal, LogNormal, logpdf +using AffineMaps: Mul, MulAdd +using JLArrays + +@testset "batched transport" begin + stdn_to_u = transport_to(StdUniform(), StdNormal()) + + @testset "several variates per stream" begin + X = vcat(randn(2, 6), rand(3, 6)) + Z, R = batched_transport_to_std_with_rest(StdUniform, StdNormal(), X, (2,)) + @test size(Z) == (2, 6) && size(R) == (3, 6) + @test Z ≈ stdn_to_u.(X[1:2, :]) + Xb, Rb = batched_transport_from_std_with_rest(StdUniform, StdNormal(), Z, (2,)) + @test Xb ≈ X[1:2, :] && size(Rb) == (0, 6) + + # Powers consume their base with their size as multiplicity: + Zp, Rp = batched_transport_to_std_with_rest(StdUniform, StdNormal()^2, X, ()) + @test Zp ≈ Z && size(Rp) == (3, 6) + X2 = vcat(X, X) + Zp2, Rp2 = batched_transport_to_std_with_rest(StdUniform, StdNormal()^2, X2, (2,)) + @test size(Zp2) == (4, 6) && size(Rp2) == (6, 6) + @test Zp2 ≈ stdn_to_u.(X2[1:4, :]) + Xp2, _ = batched_transport_from_std_with_rest(StdUniform, StdNormal()^2, Zp2, (2,)) + @test Xp2 ≈ reshape(X2[1:4, :], (2, 2, 6)) + + # Combined measures split the rows of each variate by component: + m = mcombine(vcat, StdNormal()^2, StdUniform()^3) + Z1 = batched_transport_to_std(StdUniform, m, X) + Zm, Rm = batched_transport_to_std_with_rest(StdUniform, m, X2, (2,)) + @test size(Rm) == (0, 6) && Zm ≈ vcat(Z1, Z1) + Xm, _ = batched_transport_from_std_with_rest(StdUniform, m, Zm, (2,)) + @test Xm ≈ reshape(X2, (5, 2, 6)) + + # Tuple products consume several variates via their degrees of freedom: + Pt = productmeasure((StdNormal(), StdExponential()^2)) + Zt = rand(6, 4) + Xt, Rt = batched_transport_from_std_with_rest(StdUniform, Pt, Zt, (2,)) + @test size(Xt[1]) == (2, 4) && size(Xt[2]) == (2, 2, 4) && size(Rt) == (0, 4) + for j in 1:4, i in 1:2 + a, b = transport_from_std(StdUniform, Pt, Zt[(3i - 2):(3i), j]) + @test_throws ArgumentError transport_from_std(StdUniform, Pt, Zt[:, j]) + @test Xt[1][i, j] ≈ a && Xt[2][:, i, j] ≈ b + end + + @test_throws ArgumentError batched_transport_to_std_with_rest(StdUniform, StdNormal()^2, X, (3,)) + end + + @testset "tuple and named tuple products" begin + Pt = productmeasure((StdNormal(), StdExponential()^2)) + Xt = (randn(4), rand(2, 4)) + Zt = batched_transport_to_std(StdUniform, Pt, Xt) + @test size(Zt) == (3, 4) + @test Zt ≈ stack([transport_to_std(StdUniform, Pt, (Xt[1][j], Xt[2][:, j])) for j in 1:4]) + Xr = batched_transport_from_std(StdUniform, Pt, Zt) + @test Xr[1] ≈ Xt[1] && Xr[2] ≈ Xt[2] + Pn = productmeasure((a = StdNormal(), b = StdExponential()^2)) + Zn = batched_transport_to_std(StdUniform, Pn, (a = Xt[1], b = Xt[2])) + @test Zn ≈ Zt + Xn = batched_transport_from_std(StdUniform, Pn, Zn) + @test Xn.a ≈ Xt[1] && Xn.b ≈ Xt[2] + @test batched_transport_to_std(StdUniform, Pt, (Xt[1][1], Xt[2][:, 1])) ≈ Zt[:, 1] + x1 = batched_transport_from_std(StdUniform, Pt, Zt[:, 1]) + @test x1[1] ≈ Xt[1][1] && x1[2] ≈ Xt[2][:, 1] + @test_throws ArgumentError batched_transport_from_std(StdUniform, Pt, rand(4, 4)) + end + + @testset "array products of array-variate marginals" begin + P = productmeasure([weightedmeasure(log(i), StdNormal()^2) for i in 1:3]) + X = randn(2, 3, 5) + Z = batched_transport_to_std(StdUniform, P, X) + f = transport_to(StdUniform()^6, P) + @test size(Z) == (6, 5) && Z ≈ stack([f(X[:, :, j]) for j in 1:5]) + @test batched_transport_from_std(StdUniform, P, Z) ≈ X + Y = f.(sliced(X, Val(2))) + @test flatview(Y) ≈ Z + @test flatview(inverse(f).(Y)) ≈ X + xs = [randn(2) for _ in 1:3] + z = transport_to_std(StdUniform, P, xs) + @test z ≈ f(stack(xs)) + xn = transport_from_std(StdUniform, P, z) + @test length(xn) == 3 && all(xn[i] ≈ xs[i] for i in 1:3) + @test_throws ArgumentError batched_transport_to_std(StdUniform, P, randn(2, 2, 5)) + @test_throws ArgumentError batched_transport_from_std(StdUniform, P, rand(5, 5)) + end + + @testset "streams with value-dependent sizes" begin + f_β(a) = StdNormal()^length(a) + μb = mbind(f_β, StdUniform()^1, vcat) + m = mcombine(vcat, μb, StdExponential()) + X = vcat(rand(1, 4), randn(1, 4), rand(1, 4)) + Z = batched_transport_to_std(StdUniform, m, X) + @test size(Z) == (3, 4) + @test Z ≈ stack([transport_to_std(StdUniform, m, X[:, j]) for j in 1:4]) + @test batched_transport_from_std(StdUniform, m, Z) ≈ X + P = μb^2 + x = vcat(rand(1), randn(1), rand(1), randn(1), rand(2)) + z, x_μ, x_rest = transport_to_std_with_rest(StdUniform, P, x) + @test length(z) == 4 && length(x_μ) == 4 && length(x_rest) == 2 + @test z ≈ transport_to_std(StdUniform, P, [x[1:2], x[3:4]]) + end + + @testset "elementwise pushforwards" begin + νe = pushfwd(Base.BroadcastFunction(exp), StdNormal()^3) + @test MeasureBase.mspace_ndims(typeof(νe)) == 1 + Ye = exp.(randn(3, 4)) + @test logdensities(νe, Ye) ≈ [logdensityof(νe, Ye[:, j]) for j in 1:4] + @test logdensityof(νe, Ye[:, 1]) ≈ sum(logpdf.(LogNormal(), Ye[:, 1])) + fe = transport_to(StdUniform()^3, νe) + @test flatview(fe.(sliced(Ye, Val(1)))) ≈ stack(map(fe, eachcol(Ye))) + @test flatview(inverse(fe).(fe.(sliced(Ye, Val(1))))) ≈ Ye + @test size(batched_rand_impl(GenContext{Float64}(), νe, (5,))) == (3, 5) + νr = pushfwd(Base.BroadcastFunction(exp), StdNormal()^3, PushfwdRootMeasure()) + @test logdensities(νr, Ye) ≈ [logdensityof(νr, Ye[:, j]) for j in 1:4] + JLArrays.allowscalar(false) + @test Array(logdensities(νe, JLArray(Ye))) ≈ logdensities(νe, Ye) + @test Array(flatview(fe.(sliced(JLArray(Ye), Val(1))))) ≈ flatview(fe.(sliced(Ye, Val(1)))) + end + + @testset "affine pushforwards" begin + A = [2.0 0.5; 0.0 1.5] + b = [1.0, -1.0] + ν = pushfwd(MulAdd(A, b), StdNormal()^2) + @test MeasureBase.mspace_ndims(typeof(ν)) == 1 + Y = randn(2, 5) + @test logdensities(ν, Y) ≈ [logpdf(MvNormal(b, A * A'), Y[:, j]) for j in 1:5] + @test logdensityof(ν, Y[:, 1]) ≈ logpdf(MvNormal(b, A * A'), Y[:, 1]) + f = transport_to(StdUniform()^2, ν) + @test flatview(f.(sliced(Y, Val(1)))) ≈ stack(map(f, eachcol(Y))) + @test flatview(inverse(f).(f.(sliced(Y, Val(1))))) ≈ Y + @test size(batched_rand_impl(GenContext{Float64}(), ν, (7,))) == (2, 7) + νs = pushfwd(Mul(2.0), StdNormal()) + @test logdensities(νs, Y[1, :]) ≈ logdensityof.(Ref(νs), Y[1, :]) + end + + @testset "single variates through batched forms" begin + approx(a::Tuple, b::Tuple) = all(map(approx, a, b)) + approx(a, b) = a ≈ b + for μ in ( + StdNormal(), + StdNormal()^3, + productmeasure((StdNormal(), StdExponential()^2)), + mcombine(vcat, StdNormal()^2, StdUniform()^3), + weightedmeasure(0.3, StdNormal()^2), + pushfwd(Base.BroadcastFunction(exp), StdNormal()^2), + ) + x = rand(μ) + z = MeasureBase._as_stdstream(transport_to_std(StdUniform, μ, x)) + zb = batched_transport_to_std(StdUniform, μ, x) + @test zb isa AbstractVector && zb ≈ z + @test approx(batched_transport_from_std(StdUniform, μ, z), x) + @test approx(transport_from_std(StdUniform, μ, MeasureBase._chunk_as_variate(μ, z)), x) + end + @test batched_transport_to_std(StdUniform, Dirac(1.0), 1.0) == SVector{0,Bool}() + @test batched_transport_from_std(StdUniform, Dirac(2.0), SVector{0,Bool}()) == 2.0 + @test_throws ArgumentError transport_to_std(StdUniform, StdNormal()^3, randn(3, 2)) + end +end