From b699bdce5a0e75fc33a32fb4662a7710ddfc1bde Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:38:27 +0200 Subject: [PATCH 01/75] Extend static tools --- src/MeasureBase.jl | 1 + src/combinators/power.jl | 25 +-- src/static.jl | 344 ++++++++++++++++++++++++++++++++++++--- test/static.jl | 339 +++++++++++++++++++++++++++++++++++--- 4 files changed, 648 insertions(+), 61 deletions(-) diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index 2bad7d92..63149408 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 diff --git a/src/combinators/power.jl b/src/combinators/power.jl index e6397c3f..811e7d09 100644 --- a/src/combinators/power.jl +++ b/src/combinators/power.jl @@ -17,8 +17,8 @@ struct PowerMeasure{M,A} <: AbstractProductMeasure axes::A end -maybestatic_length(μ::PowerMeasure) = prod(maybestatic_size(μ)) -maybestatic_size(μ::PowerMeasure) = map(maybestatic_length, μ.axes) +maybestatic_length(μ::PowerMeasure) = size2length(maybestatic_size(μ)) +maybestatic_size(μ::PowerMeasure) = axes2size(μ.axes) function Pretty.tile(μ::PowerMeasure) sz = length.(μ.axes) @@ -30,7 +30,7 @@ end # ToDo: Make rand return static arrays for statically-sized power measures. function _cartidxs(axs::Tuple{Vararg{AbstractUnitRange,N}}) where {N} - CartesianIndices(map(_dynamic, axs)) + CartesianIndices(map(asnonstatic, axs)) end function Base.rand( @@ -49,11 +49,8 @@ function Base.rand(rng::AbstractRNG, ::Type{T}, d::PowerMeasure) where {T} end end -@inline _pm_axes(sz::Tuple{Vararg{IntegerLike,N}}) where {N} = map(one_to, sz) -@inline _pm_axes(axs::Tuple{Vararg{AbstractUnitRange,N}}) where {N} = axs - @inline function powermeasure(x::T, sz::Tuple{Vararg{Any,N}}) where {T,N} - PowerMeasure(x, _pm_axes(sz)) + PowerMeasure(x, asaxes(sz)) end marginals(d::PowerMeasure) = fill_with(d.parent, d.axes) @@ -86,7 +83,7 @@ for func in [:logdensityof, :logdensity_def] end end - @eval @inline function $func(d::PowerMeasure{M,Tuple{Static.SOneTo{N}}}, x) where {M,N} + @eval @inline function $func(d::PowerMeasure{<:Any,Tuple{<:StaticOneToLike}}, x) parent = d.parent sum(1:N) do j @inbounds $func(parent, x[j]) @@ -94,9 +91,9 @@ for func in [:logdensityof, :logdensity_def] end @eval @inline function $func( - d::PowerMeasure{M,NTuple{N,Static.SOneTo{0}}}, + ::PowerMeasure{<:Any,<:Tuple{Vararg{StaticOneToLike{0}}}}, x, - ) where {M,N} + ) static(0.0) end end @@ -117,11 +114,7 @@ end end end -@inline getdof(μ::PowerMeasure) = getdof(μ.parent) * prod(map(length, μ.axes)) - -@inline function getdof(::PowerMeasure{<:Any,NTuple{N,Static.SOneTo{0}}}) where {N} - static(0) -end +@inline getdof(μ::PowerMeasure) = getdof(μ.parent) * size2length(axes2size(μ.axes)) @propagate_inbounds function checked_arg(μ::PowerMeasure, x::AbstractArray{<:Any}) @boundscheck begin @@ -144,7 +137,7 @@ logdensity_def(::PowerMeasure{P}, x) where {P<:PrimitiveMeasure} = static(0.0) # To avoid ambiguities function logdensity_def( - ::PowerMeasure{P,Tuple{Vararg{Static.SOneTo{0},N}}}, + ::PowerMeasure{P,<:Tuple{Vararg{StaticOneToLike{0},N}}}, x, ) where {P<:PrimitiveMeasure,N} static(0.0) diff --git a/src/static.jl b/src/static.jl index da471b62..12db9585 100644 --- a/src/static.jl +++ b/src/static.jl @@ -1,3 +1,27 @@ +# A lots of this is about bridging Static and StaticArrays, both have their +# own SUnitRange and SOneTo. Also provides tools to control static vs dynamic +# array, size and axes handling. + +""" + MeasureBase.StaticUnitRange + +The MeasureBase default type for static unit ranges. +""" +const StaticUnitRange = @static if isdefined(StaticArrays, :SUnitRange) + # Unclear if StaticArrays.SUnitRange is part of StaticArrays stable API. + # Some packages use it, but let's be careful in case it disappears. + StaticArrays.SUnitRange +else + Static.SUnitRange +end + +""" + MeasureBase.StaticOneTo + +The MeasureBase default type for static one-based unit ranges. +""" +const StaticOneTo{T} = StaticArrays.SOneTo{T} + """ MeasureBase.IntegerLike @@ -5,6 +29,67 @@ Equivalent to `Union{Integer,Static.StaticInteger}`. """ const IntegerLike = Union{Integer,Static.StaticInteger} +""" + MeasureBase.SizeLike + +Something that can represent the size of a collection. +""" +const SizeLike = Union{Tuple{},Tuple{Vararg{IntegerLike}},StaticArrays.Size} + +""" + MeasureBase.StaticSizeLike + +Something that can represent the size of a statically sized collection. +""" +const StaticSizeLike = Union{Tuple{Vararg{StaticInteger}},StaticArrays.Size} + +""" + MeasureBase.AxesLike + +Something that can represent axes of a collection. +""" +const AxesLike = Union{Tuple{},Tuple{Vararg{AbstractVector{<:IntegerLike}}}} + +""" + MeasureBase.StaticAxesLike + +Something that can represent axes of a statically sized collection. +""" +@static if isdefined(StaticArrays, :SUnitRange) + const StaticAxesLike = Union{ + Tuple{Vararg{Union{StaticArrays.SOneTo,StaticArrays.SUnitRange,Static.SUnitRange}}}, + } +else + const StaticAxesLike = + Union{Tuple{Vararg{Union{StaticArrays.SOneTo,Static.SUnitRange}}}} +end + +""" + const OneToLike + +Alias for unit ranges that start at one. +""" +const OneToLike = Union{Base.OneTo,StaticArrays.SOneTo,Static.SOneTo} + +""" + const StaticOneToLike{N} + +A static unit range from one to N. +""" +const StaticOneToLike{N} = Union{StaticArrays.SOneTo{N},Static.SOneTo{N}} + +""" + const StaticUnitRangeLike + +A static unit range. +""" +@static if isdefined(StaticArrays, :SUnitRange) + const StaticUnitRangeLike = + Union{StaticArrays.SOneTo,StaticArrays.SUnitRange,Static.SUnitRange} +else + const StaticUnitRangeLike = Union{StaticArrays.SOneTo,Static.SUnitRange} +end + """ MeasureBase.one_to(n::IntegerLike) @@ -16,48 +101,265 @@ 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.asnonstatic(x) + +Return a non-static equivalent of `x`. + +Defaults to `Static.dynamic(x)`. +""" +@inline asnonstatic(x::Number) = dynamic(x) +@inline asnonstatic(::Tuple{}) = () +@static if isdefined(StaticArrays, :SUnitRange) + @inline asnonstatic(r::StaticArrays.SUnitRange) = r[begin]:r[end] +end +@inline asnonstatic(r::AbstractUnitRange) = asnonstatic(r[begin]):asnonstatic(r[end]) +@inline asnonstatic(r::Base.OneTo) = Base.OneTo(asnonstatic(r.stop)) +@inline asnonstatic(::StaticOneToLike{N}) where {N} = Base.OneTo(N) +@inline asnonstatic(x::SizeLike) = map(asnonstatic, x) +@inline asnonstatic(::StaticArrays.Size{TPL}) where {TPL} = TPL +@inline asnonstatic(x::AxesLike) = map(asnonstatic, x) """ 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`. +The result will typically be either a `FillArrays.Fill` or a static array, """ 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)) +@inline fill_with(x::T, n::IntegerLike) where {T} = fill_with(x, (n,)) + +@inline fill_with(x::T, ::Tuple{}) where {T} = FillArrays.Fill(x) + +@inline fill_with(x, sz::SizeLike) = fill_with(x, size2axes(sz)) + +@inline function fill_with(x::T, sz::StaticSizeLike) where {T} + fill(x, staticarray_type(T, canonical_size(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) +@inline function fill_with(x, axs::AxesLike) + dyn_axs = map(asnonstatic, axs) FillArrays.Fill(x, dyn_axs) end +# 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 build a filled static array: +@inline function fill_with(x::T, axs::Tuple{Vararg{StaticOneToLike}}) where {T} + sz = axes2size(axs) + fill(x, staticarray_type(T, sz)) +end + +""" + MeasureBase.staticarray_type(T, sz::StaticArrays.Size) + +Returns the type of a static array with element type `T` and size `sz`. +""" +function staticarray_type end + +@inline @generated function staticarray_type( + ::Type{T}, + ::StaticArrays.Size{sz}, +) where {T,sz} + N = length(sz) + len = prod(sz) + :(SArray{Tuple{$sz...},T,$N,$len}) +end + +""" + MeasureBase.maybestatic_reshape(A, sz) + +Reshapes array `A` to sizes `sz`. + +If `A` is a static array and `sz` is static, the result is a static array. +""" +function maybestatic_reshape end + +maybestatic_reshape(A, sz) = reshape(A, canonical_size(sz)) +function maybestatic_reshape(A, sz::StaticSizeLike) + StaticArrays.SArray(reshape(A, canonical_size(sz))) +end +function maybestatic_reshape(A::StaticArray, sz::Tuple{Vararg{StaticInteger}}) + staticarray_type(eltype(A), canonical_size(sz))(Tuple(A)) +end + """ - MeasureBase.maybestatic_length(x)::IntegerLike + MeasureBase.maybestatic_length(x) 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}() +@inline maybestatic_length(::Number) = static(1) +@inline maybestatic_length(::Tuple{}) = static(0) +@inline maybestatic_length(::Tuple{Vararg{Any,N}}) where {N} = static(N) +@inline maybestatic_length(nt::NamedTuple) = maybestatic_length(values(nt)) +@inline maybestatic_length(A::AbstractArray) = size2length(maybestatic_size(A)) +@static if isdefined(StaticArrays, :SUnitRange) + @inline maybestatic_length(r::StaticArrays.SUnitRange) = + maybestatic_last(r) - maybestatic_first(r) + static(1) end +@inline maybestatic_length(r::AbstractUnitRange) = + maybestatic_last(r) - maybestatic_first(r) + static(1) +@inline maybestatic_length(r::Base.OneTo) = length(r) +@inline maybestatic_length(::StaticArrays.SOneTo{N}) where {N} = static(N) +@inline maybestatic_length(::Static.SOneTo{N}) where {N} = static(N) """ - MeasureBase.maybestatic_size(x)::Tuple{Vararg{IntegerLike}} + + MeasureBase.maybestatic_size(x) Returns the size of `x` as a tuple of dynamic or static integers. """ -maybestatic_size(x) = size(x) +@inline maybestatic_size(::Number) = () +@inline maybestatic_size(::Tuple{}) = + throw(ArgumentError("Cannot determine (maybe-static) size of empty tuple")) +@inline maybestatic_size(::Tuple{Vararg{Any,N}}) where {N} = StaticArrays.Size{(N,)}() +@inline maybestatic_size(nt::NamedTuple) = maybestatic_size(values(nt)) +@inline maybestatic_size(A::AbstractArray) = axes2size(maybestatic_axes(A)) +@inline maybestatic_size(A::StaticArray) = StaticArrays.Size(A) + +""" + MeasureBase.maybestatic_axes(x)::Tuple{Vararg{IntegerLike}} + +Returns the size of `x` as a tuple of dynamic or static integers. +""" +@inline maybestatic_axes(::Number) = () + +@inline maybestatic_axes(::Tuple{}) = (StaticOneTo(0),) +@inline maybestatic_axes(::Tuple{Vararg{Any,N}}) where {N} = (StaticOneTo(N),) +@inline maybestatic_axes(nt::NamedTuple) = maybestatic_axes(values(nt)) +@inline maybestatic_axes(::StaticOneToLike{N}) where {N} = (StaticOneTo(N),) +@static if isdefined(StaticArrays, :SUnitRange) + @inline maybestatic_axes(r::StaticArrays.SUnitRange) = axes(r) +end +@inline maybestatic_axes(r::Static.OptionallyStaticUnitRange) = canonical_axes(axes(r)) +@inline maybestatic_axes(r::AbstractUnitRange) = axes(r) +@inline maybestatic_axes(A::AbstractArray) = axes(A) +@inline maybestatic_axes(A::StaticArray) = axes(A) + +""" + MeasureBase.axes2size(x::Tuple) + MeasureBase.axes2size(x::StaticArrays.Size) + +Get a length from a size (tuple). +""" +@inline axes2size(::Tuple{}) = () +@inline axes2size(axs::Tuple) = canonical_size(map(maybestatic_length, axs)) + +"""map(maybestatic_length, axs) + MeasureBase.size2axes(sz::Tuple) + MeasureBase.size2axes(sz::StaticArrays.Size) + +Get one-based indexing axes from a size. +""" +@inline size2axes(::Tuple{}) = () +@inline size2axes(sz::Tuple) = canonical_axes(map(one_to, sz)) +@inline size2axes(::StaticArrays.Size{TPL}) where {TPL} = map(StaticOneTo, TPL) + +""" + MeasureBase.size2length(sz::Tuple) + MeasureBase.size2length(sz::StaticArrays.Size) + +Get a length from a size (tuple). +""" +@inline size2length(::Tuple{}) = static(1) +@inline size2length(sz::Tuple) = prod(sz) +@inline size2length(::StaticArrays.Size{TPL}) where {TPL} = static(prod(TPL)) + +""" + MeasureBase.asaxes(axs::AxesLike) + MeasureBase.asaxes(sz::SizeLike) + MeasureBase.asaxes(len::IntegerLike) + +Converts axes or a size or a length of a collection to axes. + +One-based indexing will be used if the indexing offset can't be inferred from +the given dimensions. +""" +@inline asaxes(::Tuple{}) = () +@inline asaxes(axs::AxesLike) = axs +@inline asaxes(sz::SizeLike) = size2axes(sz) +@inline asaxes(len::IntegerLike) = size2axes((len,)) + +""" + MeasureBase.maybestatic_eachindex(x) + +Returns the the index range of `x` as a dynamic or static integer range +""" +maybestatic_eachindex(::Tuple{}) = StaticOneTo(0) +maybestatic_eachindex(::Tuple{Vararg{Any,N}}) where {N} = StaticOneTo(N) +maybestatic_eachindex(nt::NamedTuple) = maybestatic_eachindex(values(nt)) +maybestatic_eachindex(x::AbstractArray) = canonical_indices(eachindex(x)) + +""" + MeasureBase.maybestatic_first(A) + +Returns the first element of `A` as a dynamic or static value. +""" +maybestatic_first(tpl::Tuple) = tpl[begin] +maybestatic_first(nt::NamedTuple) = nt[begin] +maybestatic_first(A::AbstractArray) = A[begin] +maybestatic_first(::StaticArrays.Size{tpl}) where {tpl} = static(tpl[begin]) +maybestatic_first(::StaticArrays.SOneTo{N}) where {N} = static(1) +@static if isdefined(StaticArrays, :SUnitRange) + maybestatic_first(::StaticArrays.SUnitRange{B,L}) where {B,L} = static(B) +end +function maybestatic_first( + ::Static.OptionallyStaticUnitRange{<:Static.StaticInteger{from},<:Static.StaticInteger}, +) where {from} + static(from) +end + +""" + MeasureBase.maybestatic_last(A) + +Returns the last element of `A` as a dynamic or static value. +""" +maybestatic_last(tpl::Tuple) = tpl[end] +maybestatic_last(nt::NamedTuple) = nt[end] +maybestatic_last(A::AbstractArray) = A[end] +maybestatic_last(::StaticArrays.Size{tpl}) where {tpl} = static(tpl[end]) +maybestatic_last(::StaticArrays.SOneTo{N}) where {N} = static(N) +@static if isdefined(StaticArrays, :SUnitRange) + maybestatic_last(::StaticArrays.SUnitRange{B,L}) where {B,L} = static(B + L - 1) +end +function maybestatic_last( + ::Static.OptionallyStaticUnitRange{<:Any,<:Static.StaticInteger{until}}, +) where {until} + static(until) +end + +""" + MeasureBase.canonical_indices(idxs::AbstractVector{<:IntegerLike}) + +Return the canonical representation of a collection axis indices. +""" +@inline canonical_indices(idxs::AbstractVector{<:IntegerLike}) = idxs +@inline canonical_indices(idxs::AbstractArray{<:CartesianIndex}) = idxs +@inline canonical_indices( + ::Static.OptionallyStaticUnitRange{<:StaticInteger{1},<:StaticInteger{N}}, +) where {N} = StaticArrays.SOneTo{N}() +@inline canonical_indices( + ::Static.OptionallyStaticUnitRange{<:StaticInteger{A},<:StaticInteger{B}}, +) where {A,B} = StaticUnitRange(A, B) +@inline canonical_indices( + r::Static.OptionallyStaticUnitRange{<:StaticInteger{1},<:Integer}, +) = Base.OneTo(last(r)) + +""" + MeasureBase.canonical_size(sz::SizeLike) + +Return the canonical representation of a collection size. +""" +@inline canonical_size(sz::SizeLike) = sz +@inline canonical_size(sz::Tuple{Vararg{Static.StaticInteger}}) = + StaticArrays.Size{map(dynamic, sz)}() + +""" + MeasureBase.canonical_axes(sz::SizeLike) + +Return the canonical representation collection axes. +""" +@inline canonical_axes(axs::AxesLike) = map(canonical_indices, axs) diff --git a/test/static.jl b/test/static.jl index f618124b..83092ec2 100644 --- a/test/static.jl +++ b/test/static.jl @@ -1,34 +1,325 @@ using Test import MeasureBase +using MeasureBase: + StaticUnitRange, + StaticOneTo, + IntegerLike, + SizeLike, + StaticSizeLike, + AxesLike, + StaticAxesLike, + OneToLike, + StaticOneToLike, + StaticUnitRangeLike, + one_to, + asnonstatic, + fill_with, + staticarray_type, + maybestatic_reshape, + maybestatic_length, + maybestatic_size, + maybestatic_axes, + axes2size, + size2axes, + size2length, + asaxes, + maybestatic_eachindex, + maybestatic_first, + maybestatic_last, + canonical_indices, + canonical_size, + canonical_axes import Static using Static: static +import StaticArrays 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) + v = 4.2 + T = typeof(v) + + tpl = (7, 42, 5) + nt = (a = 7, b = 42, c = 5) + + i = 7 + si = static(7) + + @test i isa IntegerLike + @test si isa IntegerLike + + sz = (2, 4, 3) + sasz = StaticArrays.Size(2, 4, 3) + sisz = (static(2), static(4), static(3)) + + len = prod(sz) + slen = static(len) + + @test sz isa SizeLike + @test sasz isa SizeLike + @test sisz isa SizeLike + + @test !(sz isa StaticSizeLike) + @test sasz isa StaticSizeLike + @test sisz isa StaticSizeLike + + axs = (Base.OneTo(2), 2:5, Base.OneTo(3)) + axs1 = (Base.OneTo(2), Base.OneTo(4), Base.OneTo(3)) + saaxs = (StaticOneTo(2), StaticUnitRange(2, 5), StaticOneTo(3)) + saaxs1 = (StaticOneTo(2), StaticOneTo(4), StaticOneTo(3)) + siaxs = (Static.SOneTo(2), static(2):static(5), static(1):static(3)) + siaxs1 = (Static.SOneTo(2), static(1):static(4), static(1):static(3)) + + @test axs isa AxesLike + @test axs1 isa AxesLike + @test saaxs isa AxesLike + @test saaxs1 isa AxesLike + @test siaxs isa AxesLike + @test siaxs1 isa AxesLike + + @test !(axs isa StaticAxesLike) + @test saaxs isa StaticAxesLike + @test saaxs1 isa StaticAxesLike + @test siaxs isa StaticAxesLike + @test siaxs1 isa StaticAxesLike + + @test axs[1] isa OneToLike + @test !(axs[2] isa OneToLike) + @test axs[3] isa OneToLike + + @test saaxs[1] isa OneToLike + @test !(saaxs[2] isa OneToLike) + @test saaxs1[2] isa OneToLike + @test saaxs[3] isa OneToLike + + @test siaxs[1] isa OneToLike + @test !(siaxs[2] isa OneToLike) + @test siaxs1[2] isa OneToLike + @test siaxs[3] isa OneToLike + + @test !(axs[1] isa StaticOneToLike) + @test !(axs[2] isa StaticOneToLike) + @test !(axs[3] isa StaticOneToLike) + + @test saaxs[1] isa StaticOneToLike + @test !(saaxs[2] isa StaticOneToLike) + @test saaxs1[2] isa StaticOneToLike + @test saaxs[3] isa StaticOneToLike + + @test siaxs[1] isa StaticOneToLike + @test !(siaxs[2] isa StaticOneToLike) + @test siaxs1[2] isa StaticOneToLike + @test siaxs[3] isa StaticOneToLike + + @test !(axs[1] isa StaticUnitRangeLike) + @test !(axs[2] isa StaticUnitRangeLike) + @test !(axs[3] isa StaticUnitRangeLike) + + @test saaxs[1] isa StaticUnitRangeLike + @test saaxs[2] isa StaticUnitRangeLike + @test saaxs1[2] isa StaticUnitRangeLike + @test saaxs[3] isa StaticUnitRangeLike + + @test siaxs[1] isa StaticUnitRangeLike + @test siaxs[2] isa StaticUnitRangeLike + @test siaxs1[2] isa StaticUnitRangeLike + @test siaxs[3] isa StaticUnitRangeLike + + @test @inferred(one_to(i)) == Base.OneTo(i) + @test @inferred(one_to(si)) == StaticOneTo(i) + + @test @inferred(asnonstatic(i)) === i + @test @inferred(asnonstatic(si)) === i + @test @inferred(asnonstatic(sz)) === sz + @test @inferred(asnonstatic(sasz)) === sz + @test @inferred(asnonstatic(sisz)) === sz + @test @inferred(asnonstatic(axs)) === axs + @test @inferred(asnonstatic(saaxs)) === axs + @test @inferred(asnonstatic(saaxs1)) === (Base.OneTo(2), Base.OneTo(4), Base.OneTo(3)) + @test @inferred(asnonstatic(siaxs)) === axs + @test @inferred(asnonstatic(siaxs1)) === (Base.OneTo(2), Base.OneTo(4), Base.OneTo(3)) + + @test @inferred(fill_with(v, i)) === FillArrays.Fill(v, i) + @test @inferred(fill_with(v, si)) === StaticArrays.SVector(fill(v, i)...) + @test @inferred(fill_with(v, ())) === FillArrays.Fill(v) + + @test @inferred(fill_with(v, sz)) === FillArrays.Fill(v, sz) + @test @inferred(fill_with(v, sasz)) === StaticArrays.SArray{Tuple{sz...},T}(fill(v, sz)) + @test @inferred(fill_with(v, sisz)) === StaticArrays.SArray{Tuple{sz...},T}(fill(v, sz)) + + @test @inferred(fill_with(v, axs)) === FillArrays.Fill(v, axs) + @test @inferred(fill_with(v, saaxs)) === FillArrays.Fill(v, axs) + @test @inferred(fill_with(v, saaxs1)) === + StaticArrays.SArray{Tuple{sz...},T}(fill(v, sz)) + @test @inferred(fill_with(v, siaxs)) === FillArrays.Fill(v, axs) + @test @inferred(fill_with(v, siaxs1)) === + StaticArrays.SArray{Tuple{sz...},T}(fill(v, sz)) + + @test @inferred(staticarray_type(T, sasz)) <: StaticArrays.SArray{Tuple{2,4,3},T} + + A = rand(T, len) + FA = FillArrays.Fill(v, len) + SA = StaticArrays.SVector(A...) + + # Array with CartesianIndices + ciA = view(rand(5, 6, 6), 3:4, 2:5, 3:5) + ciidxs = eachindex(ciA) + + rshpA = reshape(A, sz) + rshpFA = FillArrays.Fill(v, sz) + rshpSA = StaticArrays.SArray{Tuple{sz...},T}(A) + + @test @inferred(maybestatic_reshape(A, sz)) == rshpA + @test typeof(maybestatic_reshape(A, sz)) == typeof(rshpA) + @test @inferred(maybestatic_reshape(A, sasz)) == rshpA + @test maybestatic_reshape(A, sasz) isa StaticArrays.SArray + @test @inferred(maybestatic_reshape(A, sisz)) == rshpA + @test maybestatic_reshape(A, sisz) isa StaticArrays.SArray + + @test @inferred(maybestatic_reshape(FA, sz)) == rshpFA + @test typeof(maybestatic_reshape(FA, sz)) == typeof(rshpFA) + @test @inferred(maybestatic_reshape(FA, sasz)) == rshpFA + @test maybestatic_reshape(FA, sasz) isa StaticArrays.SArray + @test @inferred(maybestatic_reshape(FA, sisz)) == rshpFA + @test maybestatic_reshape(FA, sisz) isa StaticArrays.SArray + + @test @inferred(maybestatic_reshape(SA, sz)) == rshpA + @test maybestatic_reshape(SA, sz) isa Base.ReshapedArray{T,3,<:StaticArrays.SVector} + @test @inferred(maybestatic_reshape(SA, sasz)) === rshpSA + @test @inferred(maybestatic_reshape(SA, sisz)) === rshpSA + + @test @inferred(maybestatic_length(5)) === static(1) + @test @inferred(maybestatic_length(())) === static(0) + @test @inferred(maybestatic_length((sz))) === static(3) + @test @inferred(maybestatic_length((a = 2, b = 4, c = 3))) === static(3) + @test @inferred(maybestatic_length(Base.OneTo(4))) === 4 + @test @inferred(maybestatic_length(StaticArrays.SOneTo(4))) === static(4) + @test @inferred(maybestatic_length(Static.SOneTo(4))) === static(4) + @test @inferred(maybestatic_length(static(2):static(5))) === static(4) + @test @inferred(maybestatic_length(rshpA)) === length(rshpA) + @test @inferred(maybestatic_length(rshpFA)) === length(rshpA) + @test @inferred(maybestatic_length(rshpSA)) === static(length(rshpA)) + + @test @inferred(maybestatic_size(5)) === () + @test_throws ArgumentError maybestatic_size(()) + @test @inferred(maybestatic_size((sz))) === StaticArrays.Size(3) + @test @inferred(maybestatic_size((a = 2, b = 4, c = 3))) === StaticArrays.Size(3) + @test @inferred(maybestatic_size(Base.OneTo(4))) === (4,) + @test @inferred(maybestatic_size(StaticArrays.SOneTo(4))) === StaticArrays.Size(4) + @test @inferred(maybestatic_size(StaticUnitRange(2, 5))) === StaticArrays.Size(4) + @test @inferred(maybestatic_size(Static.SOneTo(4))) === StaticArrays.Size(4) + @test @inferred(maybestatic_size(static(2):static(5))) === StaticArrays.Size(4) + @test @inferred(maybestatic_size(rshpA)) === size(rshpA) + @test @inferred(maybestatic_size(rshpFA)) === size(rshpA) + @test @inferred(maybestatic_size(rshpSA)) === StaticArrays.Size(size(rshpA)...) + + @test @inferred(maybestatic_axes(5)) === () + @test @inferred(maybestatic_axes(())) === (StaticOneTo(0),) + @test @inferred(maybestatic_axes((sz))) === (StaticOneTo(3),) + @test @inferred(maybestatic_axes((a = 2, b = 4, c = 3))) === (StaticOneTo(3),) + @test @inferred(maybestatic_axes(Base.OneTo(4))) === (Base.OneTo(4),) + @test @inferred(maybestatic_axes(StaticArrays.SOneTo(4))) === (StaticOneTo(4),) + @test @inferred(maybestatic_axes(Static.SOneTo(4))) === (StaticOneTo(4),) + @test @inferred(maybestatic_axes(static(2):static(5))) === (StaticOneTo(4),) + @test @inferred(maybestatic_axes(rshpA)) === axes(rshpA) + @test @inferred(maybestatic_axes(rshpFA)) === axes(rshpA) + @test @inferred(maybestatic_axes(rshpSA)) === saaxs1 + + @test @inferred(axes2size(())) === () + @test @inferred(axes2size(axs)) === sz + @test @inferred(axes2size(saaxs)) === sasz + @test @inferred(axes2size(saaxs1)) === sasz + @test @inferred(axes2size(siaxs)) === sasz + @test @inferred(axes2size(siaxs1)) === sasz + + @test @inferred(size2axes(())) === () + @test @inferred(size2axes(sz)) === axs1 + @test @inferred(size2axes(sasz)) === saaxs1 + @test @inferred(size2axes(sisz)) === saaxs1 + + @test @inferred(size2length(())) === static(1) + @test @inferred(size2length(sz)) === len + @test @inferred(size2length(sasz)) === slen + @test @inferred(size2length(sisz)) === slen + + @test @inferred(asaxes(())) === () + @test @inferred(asaxes(len)) === (Base.OneTo(len),) + @test @inferred(asaxes(slen)) === (StaticOneTo(len),) + @test @inferred(asaxes(sz)) === axs1 + @test @inferred(asaxes(sasz)) === saaxs1 + @test @inferred(asaxes(sisz)) === saaxs1 + @test @inferred(asaxes(axs)) === axs + @test @inferred(asaxes(axs1)) === axs1 + @test @inferred(asaxes(saaxs)) === saaxs + @test @inferred(asaxes(saaxs1)) === saaxs1 + @test @inferred(asaxes(siaxs)) === siaxs + @test @inferred(asaxes(siaxs1)) === siaxs1 + + @test @inferred(maybestatic_eachindex(())) === StaticOneTo(0) + @test @inferred(maybestatic_eachindex(tpl)) === StaticOneTo(3) + @test @inferred(maybestatic_eachindex(nt)) === StaticOneTo(3) + @test @inferred(maybestatic_eachindex(axs[1])) === Base.OneTo(length(axs[1])) + @test @inferred(maybestatic_eachindex(axs[2])) === Base.OneTo(length(axs[2])) + @test @inferred(maybestatic_eachindex(saaxs[1])) === StaticOneTo(length(axs[1])) + @test @inferred(maybestatic_eachindex(saaxs[2])) === StaticOneTo(length(axs[2])) + @test @inferred(maybestatic_eachindex(siaxs[1])) === StaticOneTo(length(axs[1])) + @test @inferred(maybestatic_eachindex(siaxs[2])) === StaticOneTo(length(axs[2])) + @test @inferred(maybestatic_eachindex(A)) === Base.OneTo(24) + @test @inferred(maybestatic_eachindex(ciA)) === eachindex(ciA) + @test @inferred(maybestatic_eachindex(FA)) === Base.OneTo(24) + @test @inferred(maybestatic_eachindex(SA)) === StaticOneTo(24) + + @test_throws BoundsError maybestatic_first(()) + @test @inferred(maybestatic_first(tpl)) === first(tpl) + @test @inferred(maybestatic_first(nt)) === first(nt) + @test @inferred(maybestatic_first(sz)) === first(sz) + @test @inferred(maybestatic_first(sasz)) === static(first(sz)) + @test @inferred(maybestatic_first(sisz)) === static(first(sz)) + @test @inferred(maybestatic_first(axs[1])) === first(axs[1]) + @test @inferred(maybestatic_first(axs[2])) === first(axs[2]) + @test @inferred(maybestatic_first(saaxs[1])) === static(first(axs[1])) + @test @inferred(maybestatic_first(saaxs[2])) === static(first(axs[2])) + @test @inferred(maybestatic_first(siaxs[1])) === static(first(axs[1])) + @test @inferred(maybestatic_first(siaxs[2])) === static(first(axs[2])) + @test @inferred(maybestatic_first(A)) === first(A) + @test @inferred(maybestatic_first(ciA)) === first(ciA) + @test @inferred(maybestatic_first(FA)) === first(FA) + @test @inferred(maybestatic_first(SA)) === first(SA) + + @test_throws BoundsError maybestatic_last(()) + @test @inferred(maybestatic_last(tpl)) === last(tpl) + @test @inferred(maybestatic_last(nt)) === last(nt) + @test @inferred(maybestatic_last(sz)) === last(sz) + @test @inferred(maybestatic_last(sasz)) === static(last(sz)) + @test @inferred(maybestatic_last(sisz)) === static(last(sz)) + @test @inferred(maybestatic_last(axs[1])) === last(axs[1]) + @test @inferred(maybestatic_last(axs[2])) === last(axs[2]) + @test @inferred(maybestatic_last(saaxs[1])) === static(last(axs[1])) + @test @inferred(maybestatic_last(saaxs[2])) === static(last(axs[2])) + @test @inferred(maybestatic_last(siaxs[1])) === static(last(axs[1])) + @test @inferred(maybestatic_last(siaxs[2])) === static(last(axs[2])) + @test @inferred(maybestatic_last(A)) === last(A) + @test @inferred(maybestatic_last(ciA)) === last(ciA) + @test @inferred(maybestatic_last(FA)) === last(FA) + @test @inferred(maybestatic_last(SA)) === last(SA) + + @test @inferred(canonical_indices(axs[1])) === axs[1] + @test @inferred(canonical_indices(axs[2])) === axs[2] + @test @inferred(canonical_indices(saaxs[1])) === saaxs[1] + @test @inferred(canonical_indices(saaxs[2])) === saaxs[2] + @test @inferred(canonical_indices(siaxs[1])) === saaxs[1] + @test @inferred(canonical_indices(siaxs[2])) === saaxs[2] + @test @inferred(canonical_indices(ciidxs)) === ciidxs + + @test @inferred(canonical_size(sz)) === sz + @test @inferred(canonical_size(sasz)) === sasz + @test @inferred(canonical_size(sisz)) === sasz + + @test @inferred(canonical_axes(axs)) === axs + @test @inferred(canonical_axes(axs1)) === axs1 + @test @inferred(canonical_axes(saaxs)) === saaxs + @test @inferred(canonical_axes(saaxs1)) === saaxs1 + @test @inferred(canonical_axes(siaxs)) === saaxs + @test @inferred(canonical_axes(siaxs1)) === saaxs1 end From 1eca09e9dcf43dd3cb3a1f630bc441cbd4c968e0 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:38:27 +0200 Subject: [PATCH 02/75] Add internal infer_logdensity_type Shouldn't call Core.Compiler.return_type directly in many places. --- src/utils.jl | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/utils.jl b/src/utils.jl index 0ec81a50..5d05d8b1 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -133,6 +133,11 @@ function infer_zero(f, args...) zero(typeintersect(AbstractFloat, inferred_type)) end +function infer_logdensity_type(f::F, ::M, ::Type{T}) where {F,M,T} + inferred_type = Core.Compiler.return_type(f, Tuple{M,T}) + return inferred_type +end + @inline function allequal(f, x::AbstractArray) val = f(first(x)) @simd for xj in x From 970ebdbb36635eab4bdb964835075fcb942abcb5 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:38:27 +0200 Subject: [PATCH 03/75] Make logdensityof for PowerMeasure handle empty powers/variates --- src/combinators/power.jl | 15 ++++++++++++--- test/test_basics.jl | 5 +++-- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/src/combinators/power.jl b/src/combinators/power.jl index 811e7d09..5c0b031d 100644 --- a/src/combinators/power.jl +++ b/src/combinators/power.jl @@ -77,9 +77,18 @@ end for func in [:logdensityof, :logdensity_def] @eval @inline function $func(d::PowerMeasure{M}, x) where {M} - parent = d.parent - sum(x) do xj - $func(parent, xj) + parent_m = d.parent + sz_parent = axes2size(d.axes) + sz_x = maybestatic_size(x) + if sz_parent != sz_x + throw(ArgumentError("Size of variate doesn't match size of power measure")) + end + R = infer_logdensity_type($func, parent_m, eltype(x)) + if isempty(x) + return zero(R)::R + else + # Need to convert since sum can turn static into dynamic values: + return convert(R, sum(Base.Fix1($func, parent_m), x))::R end end diff --git a/test/test_basics.jl b/test/test_basics.jl index 7ac29dc1..bd5a409c 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -120,8 +120,9 @@ end end @testset "powers" begin - @test logdensityof(Lebesgue()^3, 2) == logdensityof(Lebesgue()^(3,), 2) - @test logdensityof(Lebesgue()^3, 2) == logdensityof(Lebesgue()^(3, 1), (2, 0)) + @test logdensityof(Lebesgue()^3, [2, 2, 2]) == logdensityof(Lebesgue()^(3,), fill(2, 3)) + @test logdensityof(Lebesgue()^3, fill(2, 3)) == + logdensityof(Lebesgue()^(3, 1), fill(2, 3, 1)) end NormalMeasure() = ∫exp(x -> -0.5x^2, Lebesgue(ℝ)) From ce8d501121ead86f4bc426fcc8b0dea8375aa13f Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:38:27 +0200 Subject: [PATCH 04/75] Use internal _TransportToStd as a function --- src/standard/stdmeasure.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl index 833f280e..4b957651 100644 --- a/src/standard/stdmeasure.jl +++ b/src/standard/stdmeasure.jl @@ -57,7 +57,7 @@ end # Helpers for product transforms and similar: struct _TransportToStd{NU<:StdMeasure} <: Function end -_TransportToStd{NU}(μ, x) where {NU} = transport_to(NU()^getdof(μ), μ)(x) +(::_TransportToStd{NU})(μ, x) where {NU} = transport_to(NU()^getdof(μ), μ)(x) struct _TransportFromStd{MU<:StdMeasure} <: Function end _TransportFromStd{MU}(ν, x) where {MU} = transport_to(ν, MU()^getdof(ν))(x) @@ -67,7 +67,7 @@ function _tuple_transport_def( μs::Tuple, xs::Tuple, ) where {NU<:StdMeasure} - reshape(vcat(map(_TransportToStd{NU}, μs, xs)...), ν.axes) + reshape(vcat(map(_TransportToStd{NU}(), μs, xs)...), ν.axes) end function transport_def( @@ -93,7 +93,7 @@ end function _stdvar_viewranges(μs::Tuple, startidx::IntegerLike) N = map(getdof, μs) offs = _offset_cumsum(startidx, N...) - map((o, n) -> o:o+n-1, offs, N) + map((o, n) -> o:(o+n-1), offs, N) end function _tuple_transport_def( From 52a610ba1321fe5bab71e47f419180c075160c82 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:38:27 +0200 Subject: [PATCH 05/75] Add pwr_base, pwr_axes, pwr_size --- src/combinators/power.jl | 35 ++++++++++++++++++++++++++++++----- 1 file changed, 30 insertions(+), 5 deletions(-) diff --git a/src/combinators/power.jl b/src/combinators/power.jl index 5c0b031d..62b4336f 100644 --- a/src/combinators/power.jl +++ b/src/combinators/power.jl @@ -11,6 +11,8 @@ 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 @@ -20,6 +22,27 @@ end maybestatic_length(μ::PowerMeasure) = size2length(maybestatic_size(μ)) 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) arg1 = Pretty.tile(μ.parent) @@ -38,14 +61,16 @@ function Base.rand( ::Type{T}, d::PowerMeasure{M}, ) where {T,M<:AbstractMeasure} - map(_cartidxs(d.axes)) do _ - rand(rng, T, d.parent) + axs, base_d = pwr_axes(d), pwr_base(d) + map(_cartidxs(axs)) do _ + rand(rng, T, base_d) end end function Base.rand(rng::AbstractRNG, ::Type{T}, d::PowerMeasure) where {T} - map(_cartidxs(d.axes)) do _ - rand(rng, d.parent) + axs, base_d = pwr_axes(d), pwr_base(d) + map(_cartidxs(axs)) do _ + rand(rng, base_d) end end @@ -127,7 +152,7 @@ end @propagate_inbounds function checked_arg(μ::PowerMeasure, x::AbstractArray{<:Any}) @boundscheck begin - sz_μ = map(length, μ.axes) + sz_μ = pwr_size(μ) sz_x = size(x) if sz_μ != sz_x throw(ArgumentError("Size of variate doesn't match size of power measure")) From d4d3100620648ad9b0bc76a9bdd0bbe02c5ebd36 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:38:27 +0200 Subject: [PATCH 06/75] Code formatting --- src/combinators/implicitlymapped.jl | 4 ++-- src/density-core.jl | 4 ++-- src/interface.jl | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/combinators/implicitlymapped.jl b/src/combinators/implicitlymapped.jl index 964ea466..3966b10a 100644 --- a/src/combinators/implicitlymapped.jl +++ b/src/combinators/implicitlymapped.jl @@ -179,13 +179,13 @@ struct TakeAny{T<:IntegerLike} n::T end -_takeany_range(f::TakeAny, idxs) = first(idxs):first(idxs)+dynamic(f.n)-1 +_takeany_range(f::TakeAny, idxs) = first(idxs):(first(idxs)+dynamic(f.n)-1) @inline _takeany_range(f::TakeAny, ::OneTo) = OneTo(dynamic(f.n)) @inline _takeany_range(::TakeAny{<:Static.StaticInteger{N}}, ::OneTo) where {N} = SOneTo(N) @inline _takeany_range(::TakeAny{<:Static.StaticInteger{N}}, ::SOneTo) where {N} = SOneTo(N) -@inline (f::TakeAny)(xs::Tuple) = xs[begin:begin+f.n-1] +@inline (f::TakeAny)(xs::Tuple) = xs[begin:(begin+f.n-1)] @inline (f::TakeAny)(xs::AbstractVector) = xs[_takeany_range(f, eachindex(xs))] function (f::TakeAny)(xs) diff --git a/src/density-core.jl b/src/density-core.jl index 6ac3d01e..f3b2db2b 100644 --- a/src/density-core.jl +++ b/src/density-core.jl @@ -149,13 +149,13 @@ end ℓ = logdensity_def(μs[$M], νs[$N], x) end - for i in 1:M-1 + for i in 1:(M-1) push!(q.args, :(Δℓ = logdensity_def(μs[$i], x))) # push!(q.args, :(println("Adding", Δℓ))) push!(q.args, :(ℓ += Δℓ)) end - for j in 1:N-1 + for j in 1:(N-1) push!(q.args, :(Δℓ = logdensity_def(νs[$j], x))) # push!(q.args, :(println("Subtracting", Δℓ))) push!(q.args, :(ℓ -= Δℓ)) diff --git a/src/interface.jl b/src/interface.jl index 18080ac7..4890ddd6 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -110,7 +110,7 @@ function test_smf(μ, n = 100) @testset "smf($μ)" begin # Get `n` sorted uniforms in O(n) time p = rand(n) - p .+= 0:n-1 + p .+= 0:(n-1) p .*= inv(n) F(x) = smf(μ, x) From 6f6087d620a1a3df57700f71224162c4aca0b0c9 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:38:27 +0200 Subject: [PATCH 07/75] Add ForwardDiff extension --- Project.toml | 3 +++ ext/MeasureBaseForwardDiffExt.jl | 14 ++++++++++++++ 2 files changed, 17 insertions(+) create mode 100644 ext/MeasureBaseForwardDiffExt.jl diff --git a/Project.toml b/Project.toml index 44f89f80..501028be 100644 --- a/Project.toml +++ b/Project.toml @@ -33,9 +33,11 @@ Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775" [weakdeps] ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" [extensions] MeasureBaseChainRulesCoreExt = "ChainRulesCore" +MeasureBaseForwardDiffExt = "ForwardDiff" [compat] ChainRulesCore = "1" @@ -45,6 +47,7 @@ ConstantRNGs = "0.1.1" ConstructionBase = "1.3" DensityInterface = "0.4" FillArrays = "0.12, 0.13, 1" +ForwardDiff = "0.8, 0.9, 0.10" FunctionChains = "0.2" IfElse = "0.1" IntervalSets = "0.7" diff --git a/ext/MeasureBaseForwardDiffExt.jl b/ext/MeasureBaseForwardDiffExt.jl new file mode 100644 index 00000000..8a1cab44 --- /dev/null +++ b/ext/MeasureBaseForwardDiffExt.jl @@ -0,0 +1,14 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +module MeasureBaseForwardDiffExt + +using MeasureBase +import ForwardDiff + +function MeasureBase.containsnan(x::ForwardDiff.Dual) + a = containsnan(x.value) + b = containsnan(x.partials) + return a || b +end + +end # module MeasureBaseForwardDiffExt From 4986a9ae5266fd844f6df4d5ce4453d2e37113cb Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:40:04 +0200 Subject: [PATCH 08/75] Add Distributions and DistributionsForwardDiff extensions --- Project.toml | 5 +++++ ext/MeasureBaseDistributionsExt.jl | 8 ++++++++ ext/MeasureBaseDistributionsForwardDiffExt.jl | 9 +++++++++ 3 files changed, 22 insertions(+) create mode 100644 ext/MeasureBaseDistributionsExt.jl create mode 100644 ext/MeasureBaseDistributionsForwardDiffExt.jl diff --git a/Project.toml b/Project.toml index 501028be..17f621dc 100644 --- a/Project.toml +++ b/Project.toml @@ -33,10 +33,13 @@ Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775" [weakdeps] ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" +Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" [extensions] MeasureBaseChainRulesCoreExt = "ChainRulesCore" +MeasureBaseDistributionsExt = "Distributions" +MeasureBaseDistributionsForwardDiffExt = ["Distributions", "ForwardDiff"] MeasureBaseForwardDiffExt = "ForwardDiff" [compat] @@ -46,6 +49,8 @@ Compat = "3.35, 4" ConstantRNGs = "0.1.1" ConstructionBase = "1.3" DensityInterface = "0.4" +Distributions = "0.25.1" +Distributions = "0.25.111" FillArrays = "0.12, 0.13, 1" ForwardDiff = "0.8, 0.9, 0.10" FunctionChains = "0.2" diff --git a/ext/MeasureBaseDistributionsExt.jl b/ext/MeasureBaseDistributionsExt.jl new file mode 100644 index 00000000..beb47821 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt.jl @@ -0,0 +1,8 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +module MeasureBaseDistributionsExt + +using MeasureBase +import Distributions + +end # module MeasureBaseDistributionsExt diff --git a/ext/MeasureBaseDistributionsForwardDiffExt.jl b/ext/MeasureBaseDistributionsForwardDiffExt.jl new file mode 100644 index 00000000..36218eec --- /dev/null +++ b/ext/MeasureBaseDistributionsForwardDiffExt.jl @@ -0,0 +1,9 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +module MeasureBaseDistributionsForwardDiffExt + +using MeasureBase +import Distributions +import ForwardDiff + +end # module MeasureBaseDistributionsForwardDiffExt From a866672c0a4503b6adbd6281577f10c9fea473ac Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:40:10 +0200 Subject: [PATCH 09/75] Add DistributionsChainRulesCore extension --- Project.toml | 1 + ext/MeasureBaseDistributionsChainRulesCoreExt.jl | 9 +++++++++ 2 files changed, 10 insertions(+) create mode 100644 ext/MeasureBaseDistributionsChainRulesCoreExt.jl diff --git a/Project.toml b/Project.toml index 17f621dc..06a51ed9 100644 --- a/Project.toml +++ b/Project.toml @@ -39,6 +39,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" [extensions] MeasureBaseChainRulesCoreExt = "ChainRulesCore" MeasureBaseDistributionsExt = "Distributions" +MeasureBaseDistributionsChainRulesCoreExt = ["Distributions", "ChainRulesCore"] MeasureBaseDistributionsForwardDiffExt = ["Distributions", "ForwardDiff"] MeasureBaseForwardDiffExt = "ForwardDiff" diff --git a/ext/MeasureBaseDistributionsChainRulesCoreExt.jl b/ext/MeasureBaseDistributionsChainRulesCoreExt.jl new file mode 100644 index 00000000..4dd3f4ff --- /dev/null +++ b/ext/MeasureBaseDistributionsChainRulesCoreExt.jl @@ -0,0 +1,9 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +module MeasureBaseDistributionsChainRulesCoreExt + +using MeasureBase +import Distributions +import ChainRulesCore + +end # module MeasureBaseDistributionsChainRulesCoreExt From 88fe1680a62df947c7574d10f4b003048e56b001 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:40:10 +0200 Subject: [PATCH 10/75] Add function asmeasure Will be used a lot when bridging from Distributions to MeasureBase. --- src/MeasureBase.jl | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index 63149408..b871fd43 100644 --- a/src/MeasureBase.jl +++ b/src/MeasureBase.jl @@ -65,6 +65,21 @@ 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 + + function Pretty.quoteof(d::M) where {M<:AbstractMeasure} the_names = fieldnames(typeof(d)) :($M($([getfield(d, n) for n in the_names]...))) From e8fd530a7ee7fcb84189420775d4405b8a715b63 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:40:10 +0200 Subject: [PATCH 11/75] Add AsMeasure --- src/MeasureBase.jl | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index b871fd43..26eab64f 100644 --- a/src/MeasureBase.jl +++ b/src/MeasureBase.jl @@ -65,7 +65,6 @@ abstract type AbstractMeasure end AbstractMeasure(m::AbstractMeasure) = m - """ asmeasure(m) @@ -79,6 +78,25 @@ function asmeasure end asmeasure(m) = convert(AbstractMeasure, m) export asmeasure +""" + struct AsMeasure{T} + +Wrapes a measure-like object into an `AbstractMeasure`. + +Constructor: + +``` +AsMeasure{T}(obj::T) +``` + +User code should not create instances of `AsMeasure` directly, but should +call `asmeasure(obj)` instead. +""" +struct AsMeasure{T} <: AbstractMeasure + obj::T + + AsMeasure{T}(obj::T) where {T} = new(obj) +end function Pretty.quoteof(d::M) where {M<:AbstractMeasure} the_names = fieldnames(typeof(d)) From 73d29a1595848b3f76469dbdf30df933bf1f00f3 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:40:10 +0200 Subject: [PATCH 12/75] Add collection utils --- ext/MeasureBaseChainRulesCoreExt.jl | 44 +++++++++++++++++++++++++++++ src/MeasureBase.jl | 1 + src/collection_utils.jl | 24 ++++++++++++++++ 3 files changed, 69 insertions(+) create mode 100644 src/collection_utils.jl diff --git a/ext/MeasureBaseChainRulesCoreExt.jl b/ext/MeasureBaseChainRulesCoreExt.jl index 57ed25fa..0384a04b 100644 --- a/ext/MeasureBaseChainRulesCoreExt.jl +++ b/ext/MeasureBaseChainRulesCoreExt.jl @@ -16,6 +16,50 @@ ChainRulesCore.rrule(::typeof(isneginf), x) = isneginf(x), _logdensityof_rt_pull _isposinf_pullback(::Any) = (NoTangent(), ZeroTangent()) ChainRulesCore.rrule(::typeof(isposinf), x) = isposinf(x), _isposinf_pullback +# = collection utils ========================================================= + +using MeasureBase: _dropfront, _dropback, _rev_cumsum, _exp_cumsum_log + +function ChainRulesCore.rrule(::typeof(_pushfront), v::AbstractVector, x) + result = _pushfront(v, x) + function _pushfront_pullback(thunked_ΔΩ) + ΔΩ = ChainRulesCore.unthunk(thunked_ΔΩ) + (NoTangent(), ΔΩ[firstindex(ΔΩ)+1:lastindex(ΔΩ)], ΔΩ[firstindex(ΔΩ)]) + end + return result, _pushfront_pullback +end + + +function ChainRulesCore.rrule(::typeof(_pushback), v::AbstractVector, x) + result = _pushback(v, x) + function _pushback_pullback(thunked_ΔΩ) + ΔΩ = ChainRulesCore.unthunk(thunked_ΔΩ) + (NoTangent(), ΔΩ[firstindex(ΔΩ):lastindex(ΔΩ)-1], ΔΩ[lastindex(ΔΩ)]) + end + return result, _pushback_pullback +end + + +function ChainRulesCore.rrule(::typeof(_rev_cumsum), xs::AbstractVector) + result = _rev_cumsum(xs) + function _rev_cumsum_pullback(ΔΩ) + ∂xs = ChainRulesCore.@thunk cumsum(ChainRulesCore.unthunk(ΔΩ)) + (NoTangent(), ∂xs) + end + return result, _rev_cumsum_pullback +end + + +function ChainRulesCore.rrule(::typeof(_exp_cumsum_log), xs::AbstractVector) + result = _exp_cumsum_log(xs) + function _exp_cumsum_log_pullback(ΔΩ) + ∂xs = inv.(xs) .* _rev_cumsum(exp.(cumsum(log.(xs))) .* ChainRulesCore.unthunk(ΔΩ)) + (NoTangent(), ∂xs) + end + return result, _exp_cumsum_log_pullback +end + + # = insupport & friends ====================================================== using MeasureBase: check_dof, require_insupport, checked_arg, _checksupport, _origin_depth diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index 26eab64f..7957bb1a 100644 --- a/src/MeasureBase.jl +++ b/src/MeasureBase.jl @@ -148,6 +148,7 @@ using IrrationalConstants using IrrationalConstants: loghalf include("static.jl") +include("collection_utils.jl") include("smf.jl") include("getdof.jl") include("transport.jl") diff --git a/src/collection_utils.jl b/src/collection_utils.jl new file mode 100644 index 00000000..1de51f7e --- /dev/null +++ b/src/collection_utils.jl @@ -0,0 +1,24 @@ +function _pushfront(v::AbstractVector, x) + T = promote_type(eltype(v), typeof(x)) + r = similar(v, T, length(eachindex(v)) + 1) + r[firstindex(r)] = x + r[firstindex(r)+1:lastindex(r)] = v + r +end + +function _pushback(v::AbstractVector, x) + T = promote_type(eltype(v), typeof(x)) + r = similar(v, T, length(eachindex(v)) + 1) + r[lastindex(r)] = x + r[firstindex(r):lastindex(r)-1] = v + r +end + +_dropfront(v::AbstractVector) = v[firstindex(v)+1:lastindex(v)] + +_dropback(v::AbstractVector) = v[firstindex(v):lastindex(v)-1] + +_rev_cumsum(xs::AbstractVector) = reverse(cumsum(reverse(xs))) + +# Equivalent to `cumprod(xs)``: +_exp_cumsum_log(xs::AbstractVector) = exp.(cumsum(log.(xs))) From 4e0645d23e707b2a90aeb3f029f7c48b277182f9 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:40:10 +0200 Subject: [PATCH 13/75] Require FunctionChains v0.2.3 --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 06a51ed9..5a31c173 100644 --- a/Project.toml +++ b/Project.toml @@ -54,7 +54,7 @@ Distributions = "0.25.1" Distributions = "0.25.111" FillArrays = "0.12, 0.13, 1" ForwardDiff = "0.8, 0.9, 0.10" -FunctionChains = "0.2" +FunctionChains = "0.2.3" IfElse = "0.1" IntervalSets = "0.7" InverseFunctions = "0.1.8" From 19dd55d77d2c3628902a69025d03c8a4c18c63d5 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:40:10 +0200 Subject: [PATCH 14/75] Require HeterogeneousComputing --- Project.toml | 2 ++ src/MeasureBase.jl | 3 +++ 2 files changed, 5 insertions(+) diff --git a/Project.toml b/Project.toml index 5a31c173..fd3f862a 100644 --- a/Project.toml +++ b/Project.toml @@ -11,6 +11,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" @@ -55,6 +56,7 @@ Distributions = "0.25.111" FillArrays = "0.12, 0.13, 1" ForwardDiff = "0.8, 0.9, 0.10" FunctionChains = "0.2.3" +HeterogeneousComputing = "0.2.3" IfElse = "0.1" IntervalSets = "0.7" InverseFunctions = "0.1.8" diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index 7957bb1a..e1a3dd12 100644 --- a/src/MeasureBase.jl +++ b/src/MeasureBase.jl @@ -44,6 +44,9 @@ using Static: StaticInteger using FunctionChains using PropertyFunctions: PropSelFunction +import HeterogeneousComputing +using HeterogeneousComputing: real_numtype + export gentype export rebase From 1d52fdd587809d0e0ef144bbb235600abbb0133b Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:40:10 +0200 Subject: [PATCH 15/75] Use StaticThings.jl for static-size tooling --- Project.toml | 2 + src/MeasureBase.jl | 14 +- src/combinators/power.jl | 17 +- src/standard/stdmeasure.jl | 2 +- src/static.jl | 365 ------------------------------------- test/static.jl | 325 --------------------------------- 6 files changed, 26 insertions(+), 699 deletions(-) delete mode 100644 src/static.jl delete mode 100644 test/static.jl diff --git a/Project.toml b/Project.toml index fd3f862a..7dea2fc6 100644 --- a/Project.toml +++ b/Project.toml @@ -28,6 +28,7 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" +StaticThings = "7e4b4f32-fbf9-4b74-9510-4d15222ac973" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775" @@ -73,6 +74,7 @@ Reexport = "1" SpecialFunctions = "2" Static = "0.8, 1" StaticArrays = "1.5" +StaticThings = "0.2" Statistics = "1" Test = "1" Tricks = "0.1" diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index e1a3dd12..c4c0d955 100644 --- a/src/MeasureBase.jl +++ b/src/MeasureBase.jl @@ -44,6 +44,19 @@ using Static: StaticInteger using FunctionChains using PropertyFunctions: PropSelFunction +using StaticThings: + AxesLike, StaticAxesLike, SizeLike, StaticSizeLike, + OneToLike, StaticOneTo, StaticOneToLike, RealLike, + IntegerLike, StaticUnitRange, StaticUnitRangeLike, + NoTypeSize, + asaxes, asnonstatic, + canonical_axes, canonical_indices, canonical_size, + maybestatic_axes, maybestatic_eachindex, + maybestatic_length, maybestatic_size, maybestatic_first, maybestatic_last, + maybestatic_oneto, maybestatic_fill, maybestatic_reshape, + size_from_type, axes2size, size2axes, size2length, + staticarray_type + import HeterogeneousComputing using HeterogeneousComputing: real_numtype @@ -150,7 +163,6 @@ using Compat using IrrationalConstants using IrrationalConstants: loghalf -include("static.jl") include("collection_utils.jl") include("smf.jl") include("getdof.jl") diff --git a/src/combinators/power.jl b/src/combinators/power.jl index 62b4336f..4d4760ba 100644 --- a/src/combinators/power.jl +++ b/src/combinators/power.jl @@ -1,4 +1,5 @@ import Base +import StaticThings export PowerMeasure @@ -19,8 +20,8 @@ struct PowerMeasure{M,A} <: AbstractProductMeasure axes::A end -maybestatic_length(μ::PowerMeasure) = size2length(maybestatic_size(μ)) -maybestatic_size(μ::PowerMeasure) = axes2size(μ.axes) +StaticThings.maybestatic_length(μ::PowerMeasure) = size2length(maybestatic_size(μ)) +StaticThings.maybestatic_size(μ::PowerMeasure) = axes2size(μ.axes) """ MeasureBase.pwr_base(μ::PowerMeasure) @@ -78,13 +79,13 @@ end PowerMeasure(x, asaxes(sz)) end -marginals(d::PowerMeasure) = fill_with(d.parent, d.axes) +marginals(d::PowerMeasure) = maybestatic_fill(d.parent, d.axes) function Base.:^(μ::AbstractMeasure, dims::Tuple{Vararg{AbstractArray,N}}) where {N} powermeasure(μ, dims) end -Base.:^(μ::AbstractMeasure, dims::Tuple) = powermeasure(μ, one_to.(dims)) +Base.:^(μ::AbstractMeasure, dims::Tuple) = powermeasure(μ, maybestatic_oneto.(dims)) Base.:^(μ::AbstractMeasure, n) = powermeasure(μ, (n,)) # Base.show(io::IO, d::PowerMeasure) = print(io, d.parent, " ^ ", size(d.xs)) @@ -117,7 +118,10 @@ for func in [:logdensityof, :logdensity_def] end end - @eval @inline function $func(d::PowerMeasure{<:Any,Tuple{<:StaticOneToLike}}, x) + @eval @inline function $func( + d::PowerMeasure{<:Any,Tuple{<:StaticOneToLike{N}}}, + x, + ) where {N} parent = d.parent sum(1:N) do j @inbounds $func(parent, x[j]) @@ -171,8 +175,7 @@ logdensity_def(::PowerMeasure{P}, x) where {P<:PrimitiveMeasure} = static(0.0) # To avoid ambiguities function logdensity_def( - ::PowerMeasure{P,<:Tuple{Vararg{StaticOneToLike{0},N}}}, - x, + ::PowerMeasure{P,<:Tuple{Vararg{StaticOneToLike{0},N}}}, ::Any, ) where {P<:PrimitiveMeasure,N} static(0.0) end diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl index 4b957651..e7244fac 100644 --- a/src/standard/stdmeasure.jl +++ b/src/standard/stdmeasure.jl @@ -13,7 +13,7 @@ function transport_def(ν::StdMeasure, μ::PowerMeasure{<:StdMeasure}, x) end function transport_def(ν::PowerMeasure{<:StdMeasure}, μ::StdMeasure, x) - return fill_with(transport_def(ν.parent, μ, only(x)), map(length, ν.axes)) + return maybestatic_fill(transport_def(ν.parent, μ, only(x)), map(length, ν.axes)) end function transport_def( diff --git a/src/static.jl b/src/static.jl deleted file mode 100644 index 12db9585..00000000 --- a/src/static.jl +++ /dev/null @@ -1,365 +0,0 @@ -# A lots of this is about bridging Static and StaticArrays, both have their -# own SUnitRange and SOneTo. Also provides tools to control static vs dynamic -# array, size and axes handling. - -""" - MeasureBase.StaticUnitRange - -The MeasureBase default type for static unit ranges. -""" -const StaticUnitRange = @static if isdefined(StaticArrays, :SUnitRange) - # Unclear if StaticArrays.SUnitRange is part of StaticArrays stable API. - # Some packages use it, but let's be careful in case it disappears. - StaticArrays.SUnitRange -else - Static.SUnitRange -end - -""" - MeasureBase.StaticOneTo - -The MeasureBase default type for static one-based unit ranges. -""" -const StaticOneTo{T} = StaticArrays.SOneTo{T} - -""" - MeasureBase.IntegerLike - -Equivalent to `Union{Integer,Static.StaticInteger}`. -""" -const IntegerLike = Union{Integer,Static.StaticInteger} - -""" - MeasureBase.SizeLike - -Something that can represent the size of a collection. -""" -const SizeLike = Union{Tuple{},Tuple{Vararg{IntegerLike}},StaticArrays.Size} - -""" - MeasureBase.StaticSizeLike - -Something that can represent the size of a statically sized collection. -""" -const StaticSizeLike = Union{Tuple{Vararg{StaticInteger}},StaticArrays.Size} - -""" - MeasureBase.AxesLike - -Something that can represent axes of a collection. -""" -const AxesLike = Union{Tuple{},Tuple{Vararg{AbstractVector{<:IntegerLike}}}} - -""" - MeasureBase.StaticAxesLike - -Something that can represent axes of a statically sized collection. -""" -@static if isdefined(StaticArrays, :SUnitRange) - const StaticAxesLike = Union{ - Tuple{Vararg{Union{StaticArrays.SOneTo,StaticArrays.SUnitRange,Static.SUnitRange}}}, - } -else - const StaticAxesLike = - Union{Tuple{Vararg{Union{StaticArrays.SOneTo,Static.SUnitRange}}}} -end - -""" - const OneToLike - -Alias for unit ranges that start at one. -""" -const OneToLike = Union{Base.OneTo,StaticArrays.SOneTo,Static.SOneTo} - -""" - const StaticOneToLike{N} - -A static unit range from one to N. -""" -const StaticOneToLike{N} = Union{StaticArrays.SOneTo{N},Static.SOneTo{N}} - -""" - const StaticUnitRangeLike - -A static unit range. -""" -@static if isdefined(StaticArrays, :SUnitRange) - const StaticUnitRangeLike = - Union{StaticArrays.SOneTo,StaticArrays.SUnitRange,Static.SUnitRange} -else - const StaticUnitRangeLike = Union{StaticArrays.SOneTo,Static.SUnitRange} -end - -""" - 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}() - -""" - MeasureBase.asnonstatic(x) - -Return a non-static equivalent of `x`. - -Defaults to `Static.dynamic(x)`. -""" -@inline asnonstatic(x::Number) = dynamic(x) -@inline asnonstatic(::Tuple{}) = () -@static if isdefined(StaticArrays, :SUnitRange) - @inline asnonstatic(r::StaticArrays.SUnitRange) = r[begin]:r[end] -end -@inline asnonstatic(r::AbstractUnitRange) = asnonstatic(r[begin]):asnonstatic(r[end]) -@inline asnonstatic(r::Base.OneTo) = Base.OneTo(asnonstatic(r.stop)) -@inline asnonstatic(::StaticOneToLike{N}) where {N} = Base.OneTo(N) -@inline asnonstatic(x::SizeLike) = map(asnonstatic, x) -@inline asnonstatic(::StaticArrays.Size{TPL}) where {TPL} = TPL -@inline asnonstatic(x::AxesLike) = map(asnonstatic, x) - -""" - MeasureBase.fill_with(x, sz::NTuple{N,<:IntegerLike}) where N - -Creates an array of size `sz` filled with `x`. - -The result will typically be either a `FillArrays.Fill` or a static array, -""" -function fill_with end - -@inline fill_with(x::T, n::IntegerLike) where {T} = fill_with(x, (n,)) - -@inline fill_with(x::T, ::Tuple{}) where {T} = FillArrays.Fill(x) - -@inline fill_with(x, sz::SizeLike) = fill_with(x, size2axes(sz)) - -@inline function fill_with(x::T, sz::StaticSizeLike) where {T} - fill(x, staticarray_type(T, canonical_size(sz))) -end - -@inline function fill_with(x, axs::AxesLike) - dyn_axs = map(asnonstatic, axs) - FillArrays.Fill(x, dyn_axs) -end - -# 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 build a filled static array: -@inline function fill_with(x::T, axs::Tuple{Vararg{StaticOneToLike}}) where {T} - sz = axes2size(axs) - fill(x, staticarray_type(T, sz)) -end - -""" - MeasureBase.staticarray_type(T, sz::StaticArrays.Size) - -Returns the type of a static array with element type `T` and size `sz`. -""" -function staticarray_type end - -@inline @generated function staticarray_type( - ::Type{T}, - ::StaticArrays.Size{sz}, -) where {T,sz} - N = length(sz) - len = prod(sz) - :(SArray{Tuple{$sz...},T,$N,$len}) -end - -""" - MeasureBase.maybestatic_reshape(A, sz) - -Reshapes array `A` to sizes `sz`. - -If `A` is a static array and `sz` is static, the result is a static array. -""" -function maybestatic_reshape end - -maybestatic_reshape(A, sz) = reshape(A, canonical_size(sz)) -function maybestatic_reshape(A, sz::StaticSizeLike) - StaticArrays.SArray(reshape(A, canonical_size(sz))) -end -function maybestatic_reshape(A::StaticArray, sz::Tuple{Vararg{StaticInteger}}) - staticarray_type(eltype(A), canonical_size(sz))(Tuple(A)) -end - -""" - MeasureBase.maybestatic_length(x) - -Returns the length of `x` as a dynamic or static integer. -""" -@inline maybestatic_length(::Number) = static(1) -@inline maybestatic_length(::Tuple{}) = static(0) -@inline maybestatic_length(::Tuple{Vararg{Any,N}}) where {N} = static(N) -@inline maybestatic_length(nt::NamedTuple) = maybestatic_length(values(nt)) -@inline maybestatic_length(A::AbstractArray) = size2length(maybestatic_size(A)) -@static if isdefined(StaticArrays, :SUnitRange) - @inline maybestatic_length(r::StaticArrays.SUnitRange) = - maybestatic_last(r) - maybestatic_first(r) + static(1) -end -@inline maybestatic_length(r::AbstractUnitRange) = - maybestatic_last(r) - maybestatic_first(r) + static(1) -@inline maybestatic_length(r::Base.OneTo) = length(r) -@inline maybestatic_length(::StaticArrays.SOneTo{N}) where {N} = static(N) -@inline maybestatic_length(::Static.SOneTo{N}) where {N} = static(N) - -""" - - MeasureBase.maybestatic_size(x) - -Returns the size of `x` as a tuple of dynamic or static integers. -""" -@inline maybestatic_size(::Number) = () -@inline maybestatic_size(::Tuple{}) = - throw(ArgumentError("Cannot determine (maybe-static) size of empty tuple")) -@inline maybestatic_size(::Tuple{Vararg{Any,N}}) where {N} = StaticArrays.Size{(N,)}() -@inline maybestatic_size(nt::NamedTuple) = maybestatic_size(values(nt)) -@inline maybestatic_size(A::AbstractArray) = axes2size(maybestatic_axes(A)) -@inline maybestatic_size(A::StaticArray) = StaticArrays.Size(A) - -""" - MeasureBase.maybestatic_axes(x)::Tuple{Vararg{IntegerLike}} - -Returns the size of `x` as a tuple of dynamic or static integers. -""" -@inline maybestatic_axes(::Number) = () - -@inline maybestatic_axes(::Tuple{}) = (StaticOneTo(0),) -@inline maybestatic_axes(::Tuple{Vararg{Any,N}}) where {N} = (StaticOneTo(N),) -@inline maybestatic_axes(nt::NamedTuple) = maybestatic_axes(values(nt)) -@inline maybestatic_axes(::StaticOneToLike{N}) where {N} = (StaticOneTo(N),) -@static if isdefined(StaticArrays, :SUnitRange) - @inline maybestatic_axes(r::StaticArrays.SUnitRange) = axes(r) -end -@inline maybestatic_axes(r::Static.OptionallyStaticUnitRange) = canonical_axes(axes(r)) -@inline maybestatic_axes(r::AbstractUnitRange) = axes(r) -@inline maybestatic_axes(A::AbstractArray) = axes(A) -@inline maybestatic_axes(A::StaticArray) = axes(A) - -""" - MeasureBase.axes2size(x::Tuple) - MeasureBase.axes2size(x::StaticArrays.Size) - -Get a length from a size (tuple). -""" -@inline axes2size(::Tuple{}) = () -@inline axes2size(axs::Tuple) = canonical_size(map(maybestatic_length, axs)) - -"""map(maybestatic_length, axs) - MeasureBase.size2axes(sz::Tuple) - MeasureBase.size2axes(sz::StaticArrays.Size) - -Get one-based indexing axes from a size. -""" -@inline size2axes(::Tuple{}) = () -@inline size2axes(sz::Tuple) = canonical_axes(map(one_to, sz)) -@inline size2axes(::StaticArrays.Size{TPL}) where {TPL} = map(StaticOneTo, TPL) - -""" - MeasureBase.size2length(sz::Tuple) - MeasureBase.size2length(sz::StaticArrays.Size) - -Get a length from a size (tuple). -""" -@inline size2length(::Tuple{}) = static(1) -@inline size2length(sz::Tuple) = prod(sz) -@inline size2length(::StaticArrays.Size{TPL}) where {TPL} = static(prod(TPL)) - -""" - MeasureBase.asaxes(axs::AxesLike) - MeasureBase.asaxes(sz::SizeLike) - MeasureBase.asaxes(len::IntegerLike) - -Converts axes or a size or a length of a collection to axes. - -One-based indexing will be used if the indexing offset can't be inferred from -the given dimensions. -""" -@inline asaxes(::Tuple{}) = () -@inline asaxes(axs::AxesLike) = axs -@inline asaxes(sz::SizeLike) = size2axes(sz) -@inline asaxes(len::IntegerLike) = size2axes((len,)) - -""" - MeasureBase.maybestatic_eachindex(x) - -Returns the the index range of `x` as a dynamic or static integer range -""" -maybestatic_eachindex(::Tuple{}) = StaticOneTo(0) -maybestatic_eachindex(::Tuple{Vararg{Any,N}}) where {N} = StaticOneTo(N) -maybestatic_eachindex(nt::NamedTuple) = maybestatic_eachindex(values(nt)) -maybestatic_eachindex(x::AbstractArray) = canonical_indices(eachindex(x)) - -""" - MeasureBase.maybestatic_first(A) - -Returns the first element of `A` as a dynamic or static value. -""" -maybestatic_first(tpl::Tuple) = tpl[begin] -maybestatic_first(nt::NamedTuple) = nt[begin] -maybestatic_first(A::AbstractArray) = A[begin] -maybestatic_first(::StaticArrays.Size{tpl}) where {tpl} = static(tpl[begin]) -maybestatic_first(::StaticArrays.SOneTo{N}) where {N} = static(1) -@static if isdefined(StaticArrays, :SUnitRange) - maybestatic_first(::StaticArrays.SUnitRange{B,L}) where {B,L} = static(B) -end -function maybestatic_first( - ::Static.OptionallyStaticUnitRange{<:Static.StaticInteger{from},<:Static.StaticInteger}, -) where {from} - static(from) -end - -""" - MeasureBase.maybestatic_last(A) - -Returns the last element of `A` as a dynamic or static value. -""" -maybestatic_last(tpl::Tuple) = tpl[end] -maybestatic_last(nt::NamedTuple) = nt[end] -maybestatic_last(A::AbstractArray) = A[end] -maybestatic_last(::StaticArrays.Size{tpl}) where {tpl} = static(tpl[end]) -maybestatic_last(::StaticArrays.SOneTo{N}) where {N} = static(N) -@static if isdefined(StaticArrays, :SUnitRange) - maybestatic_last(::StaticArrays.SUnitRange{B,L}) where {B,L} = static(B + L - 1) -end -function maybestatic_last( - ::Static.OptionallyStaticUnitRange{<:Any,<:Static.StaticInteger{until}}, -) where {until} - static(until) -end - -""" - MeasureBase.canonical_indices(idxs::AbstractVector{<:IntegerLike}) - -Return the canonical representation of a collection axis indices. -""" -@inline canonical_indices(idxs::AbstractVector{<:IntegerLike}) = idxs -@inline canonical_indices(idxs::AbstractArray{<:CartesianIndex}) = idxs -@inline canonical_indices( - ::Static.OptionallyStaticUnitRange{<:StaticInteger{1},<:StaticInteger{N}}, -) where {N} = StaticArrays.SOneTo{N}() -@inline canonical_indices( - ::Static.OptionallyStaticUnitRange{<:StaticInteger{A},<:StaticInteger{B}}, -) where {A,B} = StaticUnitRange(A, B) -@inline canonical_indices( - r::Static.OptionallyStaticUnitRange{<:StaticInteger{1},<:Integer}, -) = Base.OneTo(last(r)) - -""" - MeasureBase.canonical_size(sz::SizeLike) - -Return the canonical representation of a collection size. -""" -@inline canonical_size(sz::SizeLike) = sz -@inline canonical_size(sz::Tuple{Vararg{Static.StaticInteger}}) = - StaticArrays.Size{map(dynamic, sz)}() - -""" - MeasureBase.canonical_axes(sz::SizeLike) - -Return the canonical representation collection axes. -""" -@inline canonical_axes(axs::AxesLike) = map(canonical_indices, axs) diff --git a/test/static.jl b/test/static.jl deleted file mode 100644 index 83092ec2..00000000 --- a/test/static.jl +++ /dev/null @@ -1,325 +0,0 @@ -using Test - -import MeasureBase -using MeasureBase: - StaticUnitRange, - StaticOneTo, - IntegerLike, - SizeLike, - StaticSizeLike, - AxesLike, - StaticAxesLike, - OneToLike, - StaticOneToLike, - StaticUnitRangeLike, - one_to, - asnonstatic, - fill_with, - staticarray_type, - maybestatic_reshape, - maybestatic_length, - maybestatic_size, - maybestatic_axes, - axes2size, - size2axes, - size2length, - asaxes, - maybestatic_eachindex, - maybestatic_first, - maybestatic_last, - canonical_indices, - canonical_size, - canonical_axes - -import Static -using Static: static -import StaticArrays -import FillArrays - -@testset "static" begin - v = 4.2 - T = typeof(v) - - tpl = (7, 42, 5) - nt = (a = 7, b = 42, c = 5) - - i = 7 - si = static(7) - - @test i isa IntegerLike - @test si isa IntegerLike - - sz = (2, 4, 3) - sasz = StaticArrays.Size(2, 4, 3) - sisz = (static(2), static(4), static(3)) - - len = prod(sz) - slen = static(len) - - @test sz isa SizeLike - @test sasz isa SizeLike - @test sisz isa SizeLike - - @test !(sz isa StaticSizeLike) - @test sasz isa StaticSizeLike - @test sisz isa StaticSizeLike - - axs = (Base.OneTo(2), 2:5, Base.OneTo(3)) - axs1 = (Base.OneTo(2), Base.OneTo(4), Base.OneTo(3)) - saaxs = (StaticOneTo(2), StaticUnitRange(2, 5), StaticOneTo(3)) - saaxs1 = (StaticOneTo(2), StaticOneTo(4), StaticOneTo(3)) - siaxs = (Static.SOneTo(2), static(2):static(5), static(1):static(3)) - siaxs1 = (Static.SOneTo(2), static(1):static(4), static(1):static(3)) - - @test axs isa AxesLike - @test axs1 isa AxesLike - @test saaxs isa AxesLike - @test saaxs1 isa AxesLike - @test siaxs isa AxesLike - @test siaxs1 isa AxesLike - - @test !(axs isa StaticAxesLike) - @test saaxs isa StaticAxesLike - @test saaxs1 isa StaticAxesLike - @test siaxs isa StaticAxesLike - @test siaxs1 isa StaticAxesLike - - @test axs[1] isa OneToLike - @test !(axs[2] isa OneToLike) - @test axs[3] isa OneToLike - - @test saaxs[1] isa OneToLike - @test !(saaxs[2] isa OneToLike) - @test saaxs1[2] isa OneToLike - @test saaxs[3] isa OneToLike - - @test siaxs[1] isa OneToLike - @test !(siaxs[2] isa OneToLike) - @test siaxs1[2] isa OneToLike - @test siaxs[3] isa OneToLike - - @test !(axs[1] isa StaticOneToLike) - @test !(axs[2] isa StaticOneToLike) - @test !(axs[3] isa StaticOneToLike) - - @test saaxs[1] isa StaticOneToLike - @test !(saaxs[2] isa StaticOneToLike) - @test saaxs1[2] isa StaticOneToLike - @test saaxs[3] isa StaticOneToLike - - @test siaxs[1] isa StaticOneToLike - @test !(siaxs[2] isa StaticOneToLike) - @test siaxs1[2] isa StaticOneToLike - @test siaxs[3] isa StaticOneToLike - - @test !(axs[1] isa StaticUnitRangeLike) - @test !(axs[2] isa StaticUnitRangeLike) - @test !(axs[3] isa StaticUnitRangeLike) - - @test saaxs[1] isa StaticUnitRangeLike - @test saaxs[2] isa StaticUnitRangeLike - @test saaxs1[2] isa StaticUnitRangeLike - @test saaxs[3] isa StaticUnitRangeLike - - @test siaxs[1] isa StaticUnitRangeLike - @test siaxs[2] isa StaticUnitRangeLike - @test siaxs1[2] isa StaticUnitRangeLike - @test siaxs[3] isa StaticUnitRangeLike - - @test @inferred(one_to(i)) == Base.OneTo(i) - @test @inferred(one_to(si)) == StaticOneTo(i) - - @test @inferred(asnonstatic(i)) === i - @test @inferred(asnonstatic(si)) === i - @test @inferred(asnonstatic(sz)) === sz - @test @inferred(asnonstatic(sasz)) === sz - @test @inferred(asnonstatic(sisz)) === sz - @test @inferred(asnonstatic(axs)) === axs - @test @inferred(asnonstatic(saaxs)) === axs - @test @inferred(asnonstatic(saaxs1)) === (Base.OneTo(2), Base.OneTo(4), Base.OneTo(3)) - @test @inferred(asnonstatic(siaxs)) === axs - @test @inferred(asnonstatic(siaxs1)) === (Base.OneTo(2), Base.OneTo(4), Base.OneTo(3)) - - @test @inferred(fill_with(v, i)) === FillArrays.Fill(v, i) - @test @inferred(fill_with(v, si)) === StaticArrays.SVector(fill(v, i)...) - @test @inferred(fill_with(v, ())) === FillArrays.Fill(v) - - @test @inferred(fill_with(v, sz)) === FillArrays.Fill(v, sz) - @test @inferred(fill_with(v, sasz)) === StaticArrays.SArray{Tuple{sz...},T}(fill(v, sz)) - @test @inferred(fill_with(v, sisz)) === StaticArrays.SArray{Tuple{sz...},T}(fill(v, sz)) - - @test @inferred(fill_with(v, axs)) === FillArrays.Fill(v, axs) - @test @inferred(fill_with(v, saaxs)) === FillArrays.Fill(v, axs) - @test @inferred(fill_with(v, saaxs1)) === - StaticArrays.SArray{Tuple{sz...},T}(fill(v, sz)) - @test @inferred(fill_with(v, siaxs)) === FillArrays.Fill(v, axs) - @test @inferred(fill_with(v, siaxs1)) === - StaticArrays.SArray{Tuple{sz...},T}(fill(v, sz)) - - @test @inferred(staticarray_type(T, sasz)) <: StaticArrays.SArray{Tuple{2,4,3},T} - - A = rand(T, len) - FA = FillArrays.Fill(v, len) - SA = StaticArrays.SVector(A...) - - # Array with CartesianIndices - ciA = view(rand(5, 6, 6), 3:4, 2:5, 3:5) - ciidxs = eachindex(ciA) - - rshpA = reshape(A, sz) - rshpFA = FillArrays.Fill(v, sz) - rshpSA = StaticArrays.SArray{Tuple{sz...},T}(A) - - @test @inferred(maybestatic_reshape(A, sz)) == rshpA - @test typeof(maybestatic_reshape(A, sz)) == typeof(rshpA) - @test @inferred(maybestatic_reshape(A, sasz)) == rshpA - @test maybestatic_reshape(A, sasz) isa StaticArrays.SArray - @test @inferred(maybestatic_reshape(A, sisz)) == rshpA - @test maybestatic_reshape(A, sisz) isa StaticArrays.SArray - - @test @inferred(maybestatic_reshape(FA, sz)) == rshpFA - @test typeof(maybestatic_reshape(FA, sz)) == typeof(rshpFA) - @test @inferred(maybestatic_reshape(FA, sasz)) == rshpFA - @test maybestatic_reshape(FA, sasz) isa StaticArrays.SArray - @test @inferred(maybestatic_reshape(FA, sisz)) == rshpFA - @test maybestatic_reshape(FA, sisz) isa StaticArrays.SArray - - @test @inferred(maybestatic_reshape(SA, sz)) == rshpA - @test maybestatic_reshape(SA, sz) isa Base.ReshapedArray{T,3,<:StaticArrays.SVector} - @test @inferred(maybestatic_reshape(SA, sasz)) === rshpSA - @test @inferred(maybestatic_reshape(SA, sisz)) === rshpSA - - @test @inferred(maybestatic_length(5)) === static(1) - @test @inferred(maybestatic_length(())) === static(0) - @test @inferred(maybestatic_length((sz))) === static(3) - @test @inferred(maybestatic_length((a = 2, b = 4, c = 3))) === static(3) - @test @inferred(maybestatic_length(Base.OneTo(4))) === 4 - @test @inferred(maybestatic_length(StaticArrays.SOneTo(4))) === static(4) - @test @inferred(maybestatic_length(Static.SOneTo(4))) === static(4) - @test @inferred(maybestatic_length(static(2):static(5))) === static(4) - @test @inferred(maybestatic_length(rshpA)) === length(rshpA) - @test @inferred(maybestatic_length(rshpFA)) === length(rshpA) - @test @inferred(maybestatic_length(rshpSA)) === static(length(rshpA)) - - @test @inferred(maybestatic_size(5)) === () - @test_throws ArgumentError maybestatic_size(()) - @test @inferred(maybestatic_size((sz))) === StaticArrays.Size(3) - @test @inferred(maybestatic_size((a = 2, b = 4, c = 3))) === StaticArrays.Size(3) - @test @inferred(maybestatic_size(Base.OneTo(4))) === (4,) - @test @inferred(maybestatic_size(StaticArrays.SOneTo(4))) === StaticArrays.Size(4) - @test @inferred(maybestatic_size(StaticUnitRange(2, 5))) === StaticArrays.Size(4) - @test @inferred(maybestatic_size(Static.SOneTo(4))) === StaticArrays.Size(4) - @test @inferred(maybestatic_size(static(2):static(5))) === StaticArrays.Size(4) - @test @inferred(maybestatic_size(rshpA)) === size(rshpA) - @test @inferred(maybestatic_size(rshpFA)) === size(rshpA) - @test @inferred(maybestatic_size(rshpSA)) === StaticArrays.Size(size(rshpA)...) - - @test @inferred(maybestatic_axes(5)) === () - @test @inferred(maybestatic_axes(())) === (StaticOneTo(0),) - @test @inferred(maybestatic_axes((sz))) === (StaticOneTo(3),) - @test @inferred(maybestatic_axes((a = 2, b = 4, c = 3))) === (StaticOneTo(3),) - @test @inferred(maybestatic_axes(Base.OneTo(4))) === (Base.OneTo(4),) - @test @inferred(maybestatic_axes(StaticArrays.SOneTo(4))) === (StaticOneTo(4),) - @test @inferred(maybestatic_axes(Static.SOneTo(4))) === (StaticOneTo(4),) - @test @inferred(maybestatic_axes(static(2):static(5))) === (StaticOneTo(4),) - @test @inferred(maybestatic_axes(rshpA)) === axes(rshpA) - @test @inferred(maybestatic_axes(rshpFA)) === axes(rshpA) - @test @inferred(maybestatic_axes(rshpSA)) === saaxs1 - - @test @inferred(axes2size(())) === () - @test @inferred(axes2size(axs)) === sz - @test @inferred(axes2size(saaxs)) === sasz - @test @inferred(axes2size(saaxs1)) === sasz - @test @inferred(axes2size(siaxs)) === sasz - @test @inferred(axes2size(siaxs1)) === sasz - - @test @inferred(size2axes(())) === () - @test @inferred(size2axes(sz)) === axs1 - @test @inferred(size2axes(sasz)) === saaxs1 - @test @inferred(size2axes(sisz)) === saaxs1 - - @test @inferred(size2length(())) === static(1) - @test @inferred(size2length(sz)) === len - @test @inferred(size2length(sasz)) === slen - @test @inferred(size2length(sisz)) === slen - - @test @inferred(asaxes(())) === () - @test @inferred(asaxes(len)) === (Base.OneTo(len),) - @test @inferred(asaxes(slen)) === (StaticOneTo(len),) - @test @inferred(asaxes(sz)) === axs1 - @test @inferred(asaxes(sasz)) === saaxs1 - @test @inferred(asaxes(sisz)) === saaxs1 - @test @inferred(asaxes(axs)) === axs - @test @inferred(asaxes(axs1)) === axs1 - @test @inferred(asaxes(saaxs)) === saaxs - @test @inferred(asaxes(saaxs1)) === saaxs1 - @test @inferred(asaxes(siaxs)) === siaxs - @test @inferred(asaxes(siaxs1)) === siaxs1 - - @test @inferred(maybestatic_eachindex(())) === StaticOneTo(0) - @test @inferred(maybestatic_eachindex(tpl)) === StaticOneTo(3) - @test @inferred(maybestatic_eachindex(nt)) === StaticOneTo(3) - @test @inferred(maybestatic_eachindex(axs[1])) === Base.OneTo(length(axs[1])) - @test @inferred(maybestatic_eachindex(axs[2])) === Base.OneTo(length(axs[2])) - @test @inferred(maybestatic_eachindex(saaxs[1])) === StaticOneTo(length(axs[1])) - @test @inferred(maybestatic_eachindex(saaxs[2])) === StaticOneTo(length(axs[2])) - @test @inferred(maybestatic_eachindex(siaxs[1])) === StaticOneTo(length(axs[1])) - @test @inferred(maybestatic_eachindex(siaxs[2])) === StaticOneTo(length(axs[2])) - @test @inferred(maybestatic_eachindex(A)) === Base.OneTo(24) - @test @inferred(maybestatic_eachindex(ciA)) === eachindex(ciA) - @test @inferred(maybestatic_eachindex(FA)) === Base.OneTo(24) - @test @inferred(maybestatic_eachindex(SA)) === StaticOneTo(24) - - @test_throws BoundsError maybestatic_first(()) - @test @inferred(maybestatic_first(tpl)) === first(tpl) - @test @inferred(maybestatic_first(nt)) === first(nt) - @test @inferred(maybestatic_first(sz)) === first(sz) - @test @inferred(maybestatic_first(sasz)) === static(first(sz)) - @test @inferred(maybestatic_first(sisz)) === static(first(sz)) - @test @inferred(maybestatic_first(axs[1])) === first(axs[1]) - @test @inferred(maybestatic_first(axs[2])) === first(axs[2]) - @test @inferred(maybestatic_first(saaxs[1])) === static(first(axs[1])) - @test @inferred(maybestatic_first(saaxs[2])) === static(first(axs[2])) - @test @inferred(maybestatic_first(siaxs[1])) === static(first(axs[1])) - @test @inferred(maybestatic_first(siaxs[2])) === static(first(axs[2])) - @test @inferred(maybestatic_first(A)) === first(A) - @test @inferred(maybestatic_first(ciA)) === first(ciA) - @test @inferred(maybestatic_first(FA)) === first(FA) - @test @inferred(maybestatic_first(SA)) === first(SA) - - @test_throws BoundsError maybestatic_last(()) - @test @inferred(maybestatic_last(tpl)) === last(tpl) - @test @inferred(maybestatic_last(nt)) === last(nt) - @test @inferred(maybestatic_last(sz)) === last(sz) - @test @inferred(maybestatic_last(sasz)) === static(last(sz)) - @test @inferred(maybestatic_last(sisz)) === static(last(sz)) - @test @inferred(maybestatic_last(axs[1])) === last(axs[1]) - @test @inferred(maybestatic_last(axs[2])) === last(axs[2]) - @test @inferred(maybestatic_last(saaxs[1])) === static(last(axs[1])) - @test @inferred(maybestatic_last(saaxs[2])) === static(last(axs[2])) - @test @inferred(maybestatic_last(siaxs[1])) === static(last(axs[1])) - @test @inferred(maybestatic_last(siaxs[2])) === static(last(axs[2])) - @test @inferred(maybestatic_last(A)) === last(A) - @test @inferred(maybestatic_last(ciA)) === last(ciA) - @test @inferred(maybestatic_last(FA)) === last(FA) - @test @inferred(maybestatic_last(SA)) === last(SA) - - @test @inferred(canonical_indices(axs[1])) === axs[1] - @test @inferred(canonical_indices(axs[2])) === axs[2] - @test @inferred(canonical_indices(saaxs[1])) === saaxs[1] - @test @inferred(canonical_indices(saaxs[2])) === saaxs[2] - @test @inferred(canonical_indices(siaxs[1])) === saaxs[1] - @test @inferred(canonical_indices(siaxs[2])) === saaxs[2] - @test @inferred(canonical_indices(ciidxs)) === ciidxs - - @test @inferred(canonical_size(sz)) === sz - @test @inferred(canonical_size(sasz)) === sasz - @test @inferred(canonical_size(sisz)) === sasz - - @test @inferred(canonical_axes(axs)) === axs - @test @inferred(canonical_axes(axs1)) === axs1 - @test @inferred(canonical_axes(saaxs)) === saaxs - @test @inferred(canonical_axes(saaxs1)) === saaxs1 - @test @inferred(canonical_axes(siaxs)) === saaxs - @test @inferred(canonical_axes(siaxs1)) === saaxs1 -end From c436e4b4dd599b068f652f4f7fdb7d496b7a17df Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:40:10 +0200 Subject: [PATCH 16/75] Remove ZeroSet and CodimOne Currently unused and undocumented, can add it back later when needed. --- src/domains.jl | 60 -------------------------------------------------- 1 file changed, 60 deletions(-) diff --git a/src/domains.jl b/src/domains.jl index e03f753c..c9912420 100644 --- a/src/domains.jl +++ b/src/domains.jl @@ -69,77 +69,17 @@ function Base.getindex(::typeof(ℤ), r::AbstractUnitRange) BoundedInts(extrema(r)...) end -########################################################### -# ZeroSet -export ZeroSet - -struct ZeroSet{F,G} <: AbstractDomain - f::F - ∇f::G -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) - -########################################################### -# CodimOne - -export CodimOne - -abstract type CodimOne <: AbstractDomain 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 - - # Get the normal vectors - g1 = a.∇f(x) - g2 = b.∇f(x) - - # See if one is a multiple of the other - one(T) - Statistics.corm(g1, zero(T), g2, zero(T)) < tol -end - -function zeroset(::CodimOne)::ZeroSet end - -########################################################### -# Simplex -export Simplex - -struct Simplex <: CodimOne end - -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) -end function Base.in(x::AbstractArray{T}, ::Simplex) where {T} all(≥(zero(eltype(x))), x) || return false return x ∈ zeroset(Simplex()) end -projectto!(x, ::Simplex) = normalize!(x, 1) -########################################################### -# Sphere struct Sphere <: CodimOne 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) -end - function Base.in(x::AbstractArray{T}, ::Sphere) where {T} return x ∈ zeroset(Sphere()) end - -projectto!(x, ::Sphere) = normalize!(x, 2) From 33b2407876436ae718107eedfde8ae49097ba175 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:40:10 +0200 Subject: [PATCH 17/75] Re-design of domains --- src/domains.jl | 327 +++++++++++++++++++++++++++++++------ src/primitives/counting.jl | 5 + src/primitives/lebesgue.jl | 18 +- 3 files changed, 293 insertions(+), 57 deletions(-) diff --git a/src/domains.jl b/src/domains.jl index c9912420..c69cc29e 100644 --- a/src/domains.jl +++ b/src/domains.jl @@ -1,85 +1,316 @@ -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{MeaureBase.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, to is is by-contract. +""" +const SetLike = Union{MeasureBase.ValueSet,Base.AbstractSet,IntervalSets.Domain} + +""" + valdomain(x)::MeasureBase.SetLike + +Return the domain of a given value. + +May return [`MeasureBase.UnknownDomain`](@ref) if no domain type is available +that can represents values like `x`. +""" +function valdomain end +export valdomain + +@inline valdomain(x) = UnknownDomain(x) + +""" + MeasureBase.maybe_in(x, s) + +Test if `x` may be a member of `s`. + +Defaults to `in(x, s)`, but may be specialized for certain types of `s`, +e.g. for `s::MeasureBase.ImplicitDomain`. +""" +function maybe_in end + +maybe_in(x, s) = in(x, s) + +""" + struct MeasureBase.ImplicitDomain{M} <: MeasureBase.ValueSet + +Represents the domain (i.e. the measurable set) of a measure `m::M`. + +Constructors: + +``` +MeasureBase.ImplicitDomain(m) +``` + +Fields: + +* `m::M`: The measure. + +For many pushforward measures and similar, the measureable space can not be +computed efficiently or at all. In such cases, [`mdomain(m)`](@ref) should +return `ImplicitDomain(m)`. + +Does not support `Base.in(x, s::MeasureBase.ImplicitDomain)`, and +`MeasureBase.maybe_in(x, s::MeasureBase.ImplicitDomain)` always return `true` +(unless specialized for the measure type). +""" +struct ImplicitDomain{M} <: ValueSet + m::M +end + +@inline Base.union(s::ImplicitDomain, others::ImplicitDomain...) = + ImplicitDomain(+(s.m, map(x -> x.m, others)...)) + +function Base.in(@nospecialize(x), ::ImplicitDomain) + throw(ArgumentError("Cannot test if a value lies withing an implicit domain.")) +end + +maybe_in(@nospecialize(x), ::ImplicitDomain) = true + +function Base.isempty(::ImplicitDomain) + throw(ArgumentError("Can't test if an ImplicitDomain is empty")) +end + +""" + struct MeasureBase.UnknownDomain{T} <: MeasureBase.ValueSet + +Represents the unknown domain of a value of type `T`. + +Constructors: + +``` +MeasureBase.UnknownDomain(x::T) +``` + +Does not support `Base.in(x, s::MeasureBase.UnknownDomain)`, and +`MeasureBase.maybe_in(x, s::MeasureBase.UnknownDomain)` always return `true` +(unless specialized for the measure type). + +`isempty` will always return false, `UnknownDomain` should only be created +if a value of type `T` existed in the first place, which implies that the +domain can not be empty. +""" +struct UnknownDomain{T} <: ValueSet end + +UnknownDomain(::T) where {T} = UnknownDomain{T}() + +Base.eltype(::UnknownDomain{T}) where {T} = T + +@inline Base.union(s::UnknownDomain, others::UnknownDomain...) = + UnknownDomain{promote_type(eltype(s), map(eltype, others)...)}() + +function Base.in(@nospecialize(x), ::UnknownDomain) + throw(ArgumentError("Cannot test if a value lies withing an unknown domain.")) end -Base.in(x, b::BoundedReals) = b.lower ≤ x ≤ b.upper +maybe_in(@nospecialize(x), ::UnknownDomain) = true + +Base.isempty(::UnknownDomain) = false + +""" + RealInterval() isa MeasureBase.ValueSet + +The real numbers. +""" +struct RealValues <: ValueSet end + +@inline Base.in(x::Real, ::RealValues) = true +@inline Base.in(x, ::RealValues) = isreal(x) + +@inline Base.isempty(::RealValues) = false + +@inline Base.union(s::RealValues, ::RealValues...) = s + +@inline Base.minimum(::RealValues) = static(-Inf) +@inline Base.maximum(::RealValues) = static(Inf) + +""" + const MeasureBase.ℝ = RealValues() + +The set of all real numbers, see [`MeasureBase.RealValues`](@ref). +""" +const ℝ = RealValues() + +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) -export ℝ, ℝ₊, 𝕀, ℤ +@inline Base.isempty(::IntegerValues) = false -const ℝ₊ = BoundedReals(static(0.0), static(Inf)) -const 𝕀 = BoundedReals(static(0.0), static(1.0)) +@inline Base.union(s::IntegerValues, ::IntegerValues...) = s -Base.minimum(b::BoundedReals) = b.lower -Base.maximum(b::BoundedReals) = b.upper +# # 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)) -Base.show(io::IO, ::typeof(ℝ₊)) = print(io, "ℝ₊") -Base.show(io::IO, ::typeof(𝕀)) = print(io, "𝕀") +""" + const ℤ = IntegerValues() -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 +The set of all integers, see [`MeasureBase.IntegerValues`](@ref). +""" +const ℤ = IntegerValues() -abstract type IntegerDomain <: AbstractDomain end +Base.show(io::IO, ::MIME"text/plain", ::IntegerValues) = print(io, "MeasureBase.ℤ") -struct IntegerNumbers <: IntegerDomain end +""" + struct MeasureBase.AbstractCartSetProd <: ValueSet -Base.in(x, ::IntegerNumbers) = isinteger(x) +Supertype for cartesian products of sets. +""" +abstract type AbstractCartSetProd <: ValueSet end -const ℤ = IntegerNumbers() +""" + struct CartesianProduct <: AbstractCartSetProd -Base.show(io::IO, ::typeof(ℤ)) = print(io, "ℤ") +A cartesian product over a collection of sets. -Base.minimum(::IntegerNumbers) = static(-Inf) -Base.maximum(::IntegerNumbers) = static(Inf) -struct BoundedInts{L,U} <: IntegerDomain - lower::L - upper::U +Constructor: + +```julia +prodset = CartesianProduct(sets) +``` + +`sets` may be a `Tuple`, `NamedTuple` or `AbstractArray` of sets/domains. +""" +struct CartesianProduct{S<:Union{Tuple,NamedTuple,AbstractArray}} <: AbstractCartSetProd + _sets::S +end + +componentsets(s::CartesianProduct) = s._sets + +setcartprod(sets::AbstractArray{<:SetLike}) = CartesianProduct(sets) +setcartprod(sets::Tuple{Vararg{SetLike}}) = CartesianProduct(sets) +setcartprod(sets::NamedTuple{names,<:Tuple{Vararg{SetLike}}}) = CartesianProduct(sets) + +@inline Base.in(x::Tuple{}, s::CartesianProduct{Tuple{}}) = true +@inline Base.in(x::Tuple{Vararg{Any,N}}, s::CartesianProduct{<:Tuple{Vararg{Any,N}}}) where {N} = + prod(map(in, x, componentsets(s)))::Bool +@inline Base.in(x::NamedTuple{names}, s::CartesianProduct{<:NamedTuple{names}}) where {names} = + prod(map(in, values(x), values(componentsets(s))))::Bool +# ToDo: Allow this? +# Base.in(x::AbstractVector, s::CartesianProduct{<:Tuple}) = all(in.(x,componentsets(s)))::Bool +function Base.in( + x::AbstractArray{<:Any,N}, + s::CartesianProduct{<:AbstractArray{<:Any,N}}, +) where {N} + sets = componentsets(s) + isempty(x) && isempty(sets) ? true : all(in.(x, sets))::Bool +end + +@inline Base.isempty(s::CartesianProduct) = all(!isempty, componentsets(s)) + +@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 + +@inline function Base.union( + s::CartesianProduct{<:NamedTuple{names}}, + others::CartesianProduct{<:NamedTuple{names}}..., +) where {names} + CartesianProduct(map(union, componentsets(s), map(componentsets, others)...)) end -Base.in(x, b::BoundedInts) = x ∈ ℤ && b.lower ≤ x ≤ b.upper +function Base.union( + s::CartesianProduct{<:AbstractArray{<:Any,N}}, + others::CartesianProduct{<:AbstractArray{<:Any,N}}..., +) where {N} + CartesianProduct(union.(componentsets(s), map(componentsets, others)...)) +end -Base.minimum(b::BoundedInts) = b.lower -Base.maximum(b::BoundedInts) = b.upper +""" + struct CartesianPower <: AbstractCartSetProd -function Base.show(io::IO, b::BoundedInts) - io = IOContext(io, :compact => true) - print(io, "ℤ[", b.lower, ":", b.upper, "]") +Represents the n-fold Cartesian product of a set. +""" +struct CartesianPower{S,A} <: AbstractCartSetProd + _base::S + _axes::A end -testvalue(b::BoundedInts) = min(b.lower, 0) +@inline setcartpower(s::SetLike, dims) = CartesianPower(s, asaxes(dims)) -function Base.getindex(::typeof(ℤ), r::AbstractUnitRange) - BoundedInts(extrema(r)...) +@inline pwr_base(s::CartesianPower) = s._base +@inline pwr_axes(s::CartesianPower) = s._axes +@inline pwr_size(s::CartesianPower) = axes2size(s.axes) + +componentsets(d::CartesianPower) = fill_with(d.parent, d.axes) + +function Base.in(x::AbstractArray, s::CartesianPower) + axes2size(s.axes) == size(x) || + throw(ArgumentError("Size of CartesianPower and given point are incompatible.")) + isempty(x) ? true : all(Base.Fix1(in, s.parent), x)::Bool end +Base.isempty(s::CartesianPower) = isempty(s.parent) || size2length(axes2size(s.axes)) == 0 +function Base.union(s::CartesianPower, others::CartesianPower...) + axs = s.axes -function Base.in(x::AbstractArray{T}, ::Simplex) where {T} - all(≥(zero(eltype(x))), x) || return false - return x ∈ zeroset(Simplex()) + all(isequal(axs), map(x -> x.axes, others)) || throw( + ArgumentError("Cannot create union of CartesianPower sets with different axes."), + ) + + setcartpower(union(s.parent, map(x -> x.parent, others)...)) end +""" + struct CombinedSet <: ValueSet + +Represents a combination of two sets. -struct Sphere <: CodimOne end +User code should not create instances of `CombinedMeasure` directly, but should call +[`combinesets(f_c, α, β)`](@ref) instead. +""" -function Base.in(x::AbstractArray{T}, ::Sphere) where {T} - return x ∈ zeroset(Sphere()) +struct CombinedSet{FC,MA<:SetLike,MB<:SetLike} <: ValueSet + f_c::FC + α::MA + β::MB end diff --git a/src/primitives/counting.jl b/src/primitives/counting.jl index c61d0624..31398f32 100644 --- a/src/primitives/counting.jl +++ b/src/primitives/counting.jl @@ -40,3 +40,8 @@ insupport(μ::Counting{T}, x) where {T<:Type} = x isa μ.support massof(c::Counting, s::Set) = massof(CountingBase(), filter(insupport(c), s)) massof(::CountingBase, s::Set) = length(s) + +# ToDo: Would this be correct? +# @inline mdomain(::CountingBase) = IntegerValues() + +@inline mdomain(::Counting{DomainType}) where {DomainType} = DomainType() diff --git a/src/primitives/lebesgue.jl b/src/primitives/lebesgue.jl index 3846eaf5..040d5bd2 100644 --- a/src/primitives/lebesgue.jl +++ b/src/primitives/lebesgue.jl @@ -75,23 +75,23 @@ massof(::Lebesgue{RealNumbers}, 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 = endpoinnts(μ.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 = endpoinnts(μ.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{<:RealNumbers}, x) = x +smf(::Lebesgue{<:RealNumbers}) = identity +invsmf(::Lebesgue{<:RealNumbers}, x) = x +invsmf(::Lebesgue{<:RealNumbers}) = identity smf(::LebesgueBase, x) = x smf(::LebesgueBase) = identity From bf2ad45bfd35a284f67a39ebd663fa7b12e04312 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 14:32:27 +0200 Subject: [PATCH 18/75] Complete domain redesign --- src/domains.jl | 67 ++++++++++++++++++++++++++++++-------- src/primitives/lebesgue.jl | 18 +++++----- src/standard/stduniform.jl | 2 +- 3 files changed, 64 insertions(+), 23 deletions(-) diff --git a/src/domains.jl b/src/domains.jl index c69cc29e..7458b450 100644 --- a/src/domains.jl +++ b/src/domains.jl @@ -30,13 +30,13 @@ See also [`MeasureBase.SetLike`](@ref). abstract type ValueSet end """ - const MeasureBase.SetLike = Union{MeaureBase.ValueSet, Base.AbstractSet, IntervalSets.Domain} + 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, to is is by-contract. +imposed via type constraints, though, so it is by-contract. """ const SetLike = Union{MeasureBase.ValueSet,Base.AbstractSet,IntervalSets.Domain} @@ -142,7 +142,7 @@ maybe_in(@nospecialize(x), ::UnknownDomain) = true Base.isempty(::UnknownDomain) = false """ - RealInterval() isa MeasureBase.ValueSet + RealValues() isa MeasureBase.ValueSet The real numbers. """ @@ -158,13 +158,17 @@ struct RealValues <: ValueSet end @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.ℝ") """ @@ -179,6 +183,8 @@ struct IntegerValues <: ValueSet end @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 @@ -192,9 +198,45 @@ struct IntegerValues <: ValueSet end 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 + +@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 + +function Base.show(io::IO, b::BoundedInts) + io = IOContext(io, :compact => true) + print(io, "ℤ[", b.lower, ":", b.upper, "]") +end + +testvalue(b::BoundedInts) = convert(Int, clamp(0, dynamic(b.lower), dynamic(b.upper))) +testvalue(::Type{T}, b::BoundedInts) where {T} = convert(T, testvalue(b)) + +Base.getindex(::typeof(ℤ), r::AbstractUnitRange) = BoundedInts(extrema(r)...) + """ struct MeasureBase.AbstractCartSetProd <: ValueSet @@ -240,7 +282,7 @@ function Base.in( isempty(x) && isempty(sets) ? true : all(in.(x, sets))::Bool end -@inline Base.isempty(s::CartesianProduct) = all(!isempty, componentsets(s)) +@inline Base.isempty(s::CartesianProduct) = any(isempty, componentsets(s)) @inline function Base.union( s::CartesianProduct{<:Tuple{Vararg{Any,N}}}, @@ -277,26 +319,26 @@ end @inline pwr_base(s::CartesianPower) = s._base @inline pwr_axes(s::CartesianPower) = s._axes -@inline pwr_size(s::CartesianPower) = axes2size(s.axes) +@inline pwr_size(s::CartesianPower) = axes2size(pwr_axes(s)) -componentsets(d::CartesianPower) = fill_with(d.parent, d.axes) +componentsets(s::CartesianPower) = maybestatic_fill(pwr_base(s), pwr_axes(s)) function Base.in(x::AbstractArray, s::CartesianPower) - axes2size(s.axes) == size(x) || + pwr_size(s) == size(x) || throw(ArgumentError("Size of CartesianPower and given point are incompatible.")) - isempty(x) ? true : all(Base.Fix1(in, s.parent), x)::Bool + isempty(x) ? true : all(Base.Fix1(in, pwr_base(s)), x)::Bool end -Base.isempty(s::CartesianPower) = isempty(s.parent) || size2length(axes2size(s.axes)) == 0 +Base.isempty(s::CartesianPower) = isempty(pwr_base(s)) || size2length(pwr_size(s)) == 0 function Base.union(s::CartesianPower, others::CartesianPower...) - axs = s.axes + axs = pwr_axes(s) - all(isequal(axs), map(x -> x.axes, others)) || throw( + all(isequal(axs), map(pwr_axes, others)) || throw( ArgumentError("Cannot create union of CartesianPower sets with different axes."), ) - setcartpower(union(s.parent, map(x -> x.parent, others)...)) + setcartpower(union(pwr_base(s), map(pwr_base, others)...), axs) end @@ -308,7 +350,6 @@ Represents a combination of two sets. User code should not create instances of `CombinedMeasure` directly, but should call [`combinesets(f_c, α, β)`](@ref) instead. """ - struct CombinedSet{FC,MA<:SetLike,MB<:SetLike} <: ValueSet f_c::FC α::MA diff --git a/src/primitives/lebesgue.jl b/src/primitives/lebesgue.jl index 040d5bd2..3d92a2ed 100644 --- a/src/primitives/lebesgue.jl +++ b/src/primitives/lebesgue.jl @@ -51,7 +51,7 @@ Lebesgue() = Lebesgue(ℝ) testvalue(::Type{T}, d::Lebesgue) where {T} = testvalue(T, d.support)::T proxy(d::Lebesgue) = restrict(in(d.support), LebesgueBase()) -proxy(::Lebesgue{MeasureBase.RealNumbers}) = LebesgueBase() +proxy(::Lebesgue{MeasureBase.RealValues}) = LebesgueBase() @useproxy Lebesgue @@ -61,7 +61,7 @@ Base.show(io::IO, d::Lebesgue) = print(io, "Lebesgue(", d.support, ")") insupport(μ::Lebesgue, x) = x ∈ μ.support -insupport(::Lebesgue{RealNumbers}, ::Real) = true +insupport(::Lebesgue{RealValues}, ::Real) = true @inline function logdensityof(μ::Lebesgue, x::Real) R = float(typeof(x)) @@ -70,13 +70,13 @@ end @inline logdensityof(μ::Lebesgue, x) = insupport(μ, x) ? 0.0 : -Inf -massof(::Lebesgue{RealNumbers}, s::Interval) = width(s) +massof(::Lebesgue{RealValues}, s::Interval) = width(s) # Example: # julia> Lebesgue(𝕀)(0.2..5) # 0.8 function massof(μ::Lebesgue{<:AbstractInterval}, s::Interval) - a, b = endpoinnts(μ.support) + a, b = endpoints(μ.support) left = max(s.left, a) right = min(s.right, b) w = right - left @@ -84,14 +84,14 @@ function massof(μ::Lebesgue{<:AbstractInterval}, s::Interval) end function smf(μ::Lebesgue{<:AbstractInterval}, x) - a, b = endpoinnts(μ.support) + a, b = endpoints(μ.support) clamp(x, a, b) end -smf(::Lebesgue{<:RealNumbers}, x) = x -smf(::Lebesgue{<:RealNumbers}) = identity -invsmf(::Lebesgue{<:RealNumbers}, x) = x -invsmf(::Lebesgue{<:RealNumbers}) = identity +smf(::Lebesgue{<:RealValues}, x) = x +smf(::Lebesgue{<:RealValues}) = identity +invsmf(::Lebesgue{<:RealValues}, x) = x +invsmf(::Lebesgue{<:RealValues}) = identity smf(::LebesgueBase, x) = x smf(::LebesgueBase) = identity diff --git a/src/standard/stduniform.jl b/src/standard/stduniform.jl index 7bbe15ed..e3702656 100644 --- a/src/standard/stduniform.jl +++ b/src/standard/stduniform.jl @@ -14,7 +14,7 @@ end Base.rand(rng::Random.AbstractRNG, ::Type{T}, ::StdUniform) where {T} = rand(rng, T) -massof(::StdUniform, s::Interval) = massof(Lebesgue(𝕀), s::Interval) +massof(::StdUniform, s::Interval) = massof(Lebesgue(0.0 .. 1.0), s) smf(::StdUniform, x) = clamp(x, zero(x), one(x)) From 7835f8492d2e2564eb8c27f4fc1864c06a66e0ba Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:40:10 +0200 Subject: [PATCH 19/75] Add mreshape --- src/MeasureBase.jl | 1 + src/combinators/power.jl | 3 ++ src/combinators/reshape.jl | 69 +++++++++++++++++++++++++++++++++++++ test/combinators/reshape.jl | 7 ++++ test/runtests.jl | 1 + 5 files changed, 81 insertions(+) create mode 100644 src/combinators/reshape.jl create mode 100644 test/combinators/reshape.jl diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index c4c0d955..4ab59601 100644 --- a/src/MeasureBase.jl +++ b/src/MeasureBase.jl @@ -184,6 +184,7 @@ include("primitives/trivial.jl") include("combinators/bind.jl") include("combinators/transformedmeasure.jl") +include("combinators/reshape.jl") include("combinators/weighted.jl") include("combinators/superpose.jl") include("combinators/product.jl") diff --git a/src/combinators/power.jl b/src/combinators/power.jl index 4d4760ba..6f065c3f 100644 --- a/src/combinators/power.jl +++ b/src/combinators/power.jl @@ -179,3 +179,6 @@ function logdensity_def( ) where {P<:PrimitiveMeasure,N} static(0.0) end + + +@inline mspace_elsize(m::PowerMeasure) = axes2size(m.axes) diff --git a/src/combinators/reshape.jl b/src/combinators/reshape.jl new file mode 100644 index 00000000..dddae55b --- /dev/null +++ b/src/combinators/reshape.jl @@ -0,0 +1,69 @@ +# ToDo: Support static resizes for static arrays + +""" + struct MeasureBase.Reshape <: Function + +Represents a function that reshapes an array. + +Supports `InverseFunctions.inverse` and +`ChangesOfVariables.with_logabsdet_jacobian`. + +Constructor: + +```julia +Reshape(output_size::Dims, input_size::Dims) +``` +""" +struct Reshape{M<:SizeLike,N<:SizeLike} <: Function + output_size::M + input_size::N + + Reshape{M,N}(out_sz::M, in_sz::N) where {M<:SizeLike,N<:SizeLike} = + new{M,N}(out_sz, in_sz) +end + +function Reshape(output_size::SizeLike, input_size::SizeLike) + out_sz = canonical_size(output_size) + in_sz = canonical_size(input_size) + return Reshape{typeof(out_sz), typeof(in_sz)}(out_sz, in_sz) +end + +_throw_reshape_mismatch(sz, sz_x) = throw(DimensionMismatch("Reshape input size is $sz but got input of size $sz_x")) + +function (f::Reshape)(x::AbstractArray) + sz_x = maybestatic_size(x) + f.input_size == sz_x || _throw_reshape_mismatch(f.input_size, sz_x) + return reshape(x, f.output_size) +end + +InverseFunctions.inverse(f::Reshape{M,N}) where {M,N} = Reshape{N,M}(f.input_size, f.output_size) + +function ChangesOfVariables.with_logabsdet_jacobian(f::Reshape, x::AbstractArray) + return f(x), zero(real_numtype(typeof(x))) +end + + +""" + mreshape(m::AbstractMeasure, sz::Vararg{N,IntegerLike}) where N + mreshape(m::AbstractMeasure, sz::NTuple{N,IntegerLike}) where N + +Reshape a measure `m` over an array-valued space, returning a measure over +a space of arrays with shape `sz`. +""" +function mreshape end + +mreshape(m::AbstractMeasure, sz::IntegerLike...) = mreshape(m, sz) +mreshape(m::AbstractMeasure, sz::SizeLike) = pushfwd(Reshape(sz, mspace_elsize(m)), m) + + +""" + MeasureBase.mspace_elsize(m::AbstractMeasure)::MeasureBase.SizeLike + +Return the size of the elements of the measurable space of `m`. + +Defaults to the size of a test value of `m`, may be specialized for +measure types where this is inefficient. +""" +function mspace_elsize end + +mspace_elsize(m::AbstractMeasure) = maybestatic_size(testvalue(m)) diff --git a/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/runtests.jl b/test/runtests.jl index c2f63c4e..dfd8b93a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -21,6 +21,7 @@ include("smf.jl") include("combinators/weighted.jl") include("combinators/transformedmeasure.jl") +include("combinators/reshape.jl") include("combinators/implicitlymapped.jl") include("test_docs.jl") From 78fc01ec4f95982ca72221368e50478027a71207 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 00:49:49 +0200 Subject: [PATCH 20/75] Require SpecialFunctions 2.1.4 --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 7dea2fc6..e729f511 100644 --- a/Project.toml +++ b/Project.toml @@ -71,7 +71,7 @@ PrettyPrinting = "0.3, 0.4" PropertyFunctions = "0.2.2" Random = "1" Reexport = "1" -SpecialFunctions = "2" +SpecialFunctions = "2.1.4" Static = "0.8, 1" StaticArrays = "1.5" StaticThings = "0.2" From 9853449b5118b0d05f83341f46a48344c4480010 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 00:50:31 +0200 Subject: [PATCH 21/75] Fix _default_checked_arg --- src/getdof.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/getdof.jl b/src/getdof.jl index dbce2202..16ae7cc6 100644 --- a/src/getdof.jl +++ b/src/getdof.jl @@ -66,7 +66,7 @@ function checked_arg end # Prevent infinite recursion: @propagate_inbounds function _default_checked_arg(::Type{MU}, ::MU, ::T) where {MU,T} - NoArgCheck{MU,T} + NoArgCheck{MU,T}() end @propagate_inbounds function _default_checked_arg(::Type{MU}, mu_base, x) where {MU} checked_arg(mu_base, x) From b9c53415b1385bfc892944f9cb4d90e15f6dbd34 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 00:52:05 +0200 Subject: [PATCH 22/75] Fix checked_arg for ProductMeasure --- src/combinators/product.jl | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/combinators/product.jl b/src/combinators/product.jl index 0290419d..9135dc2b 100644 --- a/src/combinators/product.jl +++ b/src/combinators/product.jl @@ -244,6 +244,10 @@ function checked_arg(μ::ProductMeasure{<:NTuple{N,Any}}, x::NTuple{N,Any}) wher map(checked_arg, marginals(μ), x) end +function checked_arg(μ::ProductMeasure{<:AbstractArray}, x::AbstractArray) + map(checked_arg, marginals(μ), x) +end + function checked_arg( μ::ProductMeasure{<:NamedTuple{names}}, x::NamedTuple{names}, From 87bde7e4bc81cfa1e804e1a45b093ecbbb23c4bb Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 00:52:48 +0200 Subject: [PATCH 23/75] More transport_def methods for PowerMeasure and ProductMeasure --- src/standard/stdmeasure.jl | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl index e7244fac..a9c09b5c 100644 --- a/src/standard/stdmeasure.jl +++ b/src/standard/stdmeasure.jl @@ -121,3 +121,28 @@ function transport_def( ) where {MU<:StdMeasure,names} NamedTuple{names}(_tuple_transport_def(values(marginals(ν)), μ, x)) end + +function transport_def( + ν::PowerMeasure{NU}, + μ::ProductMeasure{<:AbstractArray}, + x, +) where {NU<:StdMeasure} + reshape(vcat(map(_TransportToStd{NU}(), marginals(μ), x)...), ν.axes) +end + +function _marginal_viewranges(μs::AbstractArray, startidx::IntegerLike) + ns = map(m -> dynamic(getdof(m)), μs) + offs = cumsum(vcat(dynamic(startidx), ns[begin:(end-1)])) + map((o, n) -> o:(o+n-1), offs, ns) +end + +function transport_def( + ν::ProductMeasure{<:AbstractArray}, + μ::PowerMeasure{MU}, + x::AbstractArray{<:Real}, +) where {MU<:StdMeasure} + νs = marginals(ν) + vrs = _marginal_viewranges(νs, firstindex(x)) + xs = map(r -> view(x, r), vrs) + map(_TransportFromStd{MU}, νs, xs) +end From 43d355ffcbfb2edd92e0d50b5ad305bc3735d15a Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 08:37:30 +0200 Subject: [PATCH 24/75] Add Mooncake extension with AD rules Created by generative AI. --- Project.toml | 3 ++ ext/MeasureBaseMooncakeExt.jl | 29 ++++++++++++++++++ test/Project.toml | 1 + test/runtests.jl | 2 ++ test/test_mooncake.jl | 56 +++++++++++++++++++++++++++++++++++ 5 files changed, 91 insertions(+) create mode 100644 ext/MeasureBaseMooncakeExt.jl create mode 100644 test/test_mooncake.jl diff --git a/Project.toml b/Project.toml index e729f511..486afa73 100644 --- a/Project.toml +++ b/Project.toml @@ -37,6 +37,7 @@ Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" [extensions] MeasureBaseChainRulesCoreExt = "ChainRulesCore" @@ -44,6 +45,7 @@ MeasureBaseDistributionsExt = "Distributions" MeasureBaseDistributionsChainRulesCoreExt = ["Distributions", "ChainRulesCore"] MeasureBaseDistributionsForwardDiffExt = ["Distributions", "ForwardDiff"] MeasureBaseForwardDiffExt = "ForwardDiff" +MeasureBaseMooncakeExt = "Mooncake" [compat] ChainRulesCore = "1" @@ -66,6 +68,7 @@ LinearAlgebra = "1" LogExpFunctions = "0.3, 1" LogarithmicNumbers = "1" MappedArrays = "0.4" +Mooncake = "0.5.34" NaNMath = "0.3, 1" PrettyPrinting = "0.3, 0.4" PropertyFunctions = "0.2.2" diff --git a/ext/MeasureBaseMooncakeExt.jl b/ext/MeasureBaseMooncakeExt.jl new file mode 100644 index 00000000..e401cb7d --- /dev/null +++ b/ext/MeasureBaseMooncakeExt.jl @@ -0,0 +1,29 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +module MeasureBaseMooncakeExt + +using MeasureBase +import Mooncake +using Mooncake: @zero_derivative, MinimalCtx + +using MeasureBase: isneginf, isposinf, _adignore_call +using MeasureBase: check_dof, require_insupport, _origin_depth +using MeasureBase: logdensityof_rt + +# Unlike Zygote, Mooncake differentiates the collection utilities +# (`_pushfront`, etc., mutating code in general), `checked_arg` and +# `_checksupport` natively, so only the non-differentiable functions +# need rules: + +@zero_derivative MinimalCtx Tuple{typeof(isneginf),Any} +@zero_derivative MinimalCtx Tuple{typeof(isposinf),Any} + +@zero_derivative MinimalCtx Tuple{typeof(_adignore_call),Any} + +@zero_derivative MinimalCtx Tuple{typeof(require_insupport),Any,Any} +@zero_derivative MinimalCtx Tuple{typeof(_origin_depth),Any} +@zero_derivative MinimalCtx Tuple{typeof(check_dof),Any,Any} + +@zero_derivative MinimalCtx Tuple{typeof(logdensityof_rt),Any,Any} + +end # module MeasureBaseMooncakeExt diff --git a/test/Project.toml b/test/Project.toml index 376c1b05..fec229fd 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -12,6 +12,7 @@ IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" LogarithmicNumbers = "aa2f6b4e-9042-5d33-9679-40d3a6b85899" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" PropertyFunctions = "09e99361-2bb8-48a2-a80f-de58f0739eb4" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" diff --git a/test/runtests.jl b/test/runtests.jl index dfd8b93a..ff5c3140 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -19,6 +19,8 @@ include("getdof.jl") include("transport.jl") include("smf.jl") +include("test_mooncake.jl") + include("combinators/weighted.jl") include("combinators/transformedmeasure.jl") include("combinators/reshape.jl") diff --git a/test/test_mooncake.jl b/test/test_mooncake.jl new file mode 100644 index 00000000..95ce84b5 --- /dev/null +++ b/test/test_mooncake.jl @@ -0,0 +1,56 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using Random +import Mooncake +import ForwardDiff + +using MeasureBase +using MeasureBase: transport_to +using MeasureBase: isneginf, isposinf, _adignore_call +using MeasureBase: check_dof, require_insupport, _origin_depth +using MeasureBase: logdensityof_rt + +_mooncake_gradient(f, x) = Mooncake.value_and_gradient!!( + Mooncake.prepare_gradient_cache(f, x), f, x +)[2][2] + +@testset "Mooncake AD rules" begin + @test Base.get_extension(MeasureBase, :MeasureBaseMooncakeExt) isa Module + + @testset "zero-derivative primitives" begin + rng = Random.Xoshiro(789990641) + Mooncake.TestUtils.test_rule(rng, isneginf, 0.5; is_primitive = true) + Mooncake.TestUtils.test_rule(rng, isposinf, 0.5; is_primitive = true) + Mooncake.TestUtils.test_rule(rng, _adignore_call, () -> 42.0; is_primitive = true) + Mooncake.TestUtils.test_rule(rng, check_dof, StdNormal(), StdUniform(); is_primitive = true) + Mooncake.TestUtils.test_rule(rng, require_insupport, StdNormal(), 0.5; is_primitive = true) + Mooncake.TestUtils.test_rule(rng, _origin_depth, StdNormal(); is_primitive = true) + Mooncake.TestUtils.test_rule(rng, logdensityof_rt, StdNormal(), 0.5; is_primitive = true) + end + + @testset "@_adignore is ignored" begin + f_adignore(x) = (MeasureBase.@_adignore x^3; x^2) + @test _mooncake_gradient(f_adignore, 3.0) ≈ 6.0 + end + + @testset "logdensityof gradients" begin + x = [0.1, -0.2, 0.3] + f_ld = x -> logdensityof(StdNormal()^3, x) + @test _mooncake_gradient(f_ld, x) ≈ ForwardDiff.gradient(f_ld, x) + + f_ldu = x -> logdensityof(StdExponential()^3, x) + @test _mooncake_gradient(f_ldu, abs.(x)) ≈ ForwardDiff.gradient(f_ldu, abs.(x)) + end + + @testset "transport gradients" begin + x = [0.1, -0.2, 0.3] + f_t = x -> sum(transport_to(StdUniform()^3, StdNormal()^3)(x)) + @test _mooncake_gradient(f_t, x) ≈ ForwardDiff.gradient(f_t, x) + + u = [0.3, 0.5, 0.7] + f_ti = u -> sum(transport_to(StdNormal()^3, StdUniform()^3)(u)) + @test _mooncake_gradient(f_ti, u) ≈ ForwardDiff.gradient(f_ti, u) + end +end From ba5e85f92c6411ba68de7e3ca6123d8e2b7abed0 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Sun, 5 Jul 2026 12:40:10 +0200 Subject: [PATCH 25/75] Add Distributions extension Assisted by generative AI. --- Project.toml | 20 +- ext/MeasureBaseChainRulesCoreExt.jl | 15 +- ...asureBaseDistributionsChainRulesCoreExt.jl | 9 + ext/MeasureBaseDistributionsExt.jl | 8 - .../MeasureBaseDistributionsExt.jl | 66 ++++++ ext/MeasureBaseDistributionsExt/dirac.jl | 14 ++ ext/MeasureBaseDistributionsExt/dirichlet.jl | 59 +++++ .../dist_vartransform.jl | 16 ++ .../distribution_measure.jl | 74 ++++++ .../measure_interface.jl | 26 ++ ext/MeasureBaseDistributionsExt/mixture.jl | 24 ++ ext/MeasureBaseDistributionsExt/product.jl | 56 +++++ ext/MeasureBaseDistributionsExt/reshaped.jl | 27 +++ .../standard_dist.jl | 222 ++++++++++++++++++ .../standard_normal.jl | 77 ++++++ .../standard_uniform.jl | 79 +++++++ ext/MeasureBaseDistributionsExt/standardmv.jl | 36 +++ ext/MeasureBaseDistributionsExt/univariate.jl | 142 +++++++++++ ext/MeasureBaseDistributionsForwardDiffExt.jl | 29 ++- ...aseDistributionsForwardDiffPullbacksExt.jl | 25 ++ ext/MeasureBaseDistributionsMooncakeExt.jl | 21 ++ ext/MeasureBaseForwardDiffExt.jl | 4 + ext/MeasureBaseForwardDiffPullbacksExt.jl | 10 + src/combinators/superpose.jl | 102 ++++---- src/utils.jl | 51 ++++ test/Project.toml | 7 + test/distributions/getjacobian.jl | 34 +++ test/distributions/test_autodiff_utils.jl | 18 ++ test/distributions/test_conversions.jl | 113 +++++++++ .../test_distribution_measure.jl | 53 +++++ test/distributions/test_distributions.jl | 24 ++ test/distributions/test_measure_interface.jl | 43 ++++ test/distributions/test_mooncake.jl | 68 ++++++ test/distributions/test_standard_dist.jl | 128 ++++++++++ test/distributions/test_standard_normal.jl | 129 ++++++++++ test/distributions/test_standard_uniform.jl | 118 ++++++++++ test/distributions/test_transport.jl | 194 +++++++++++++++ test/runtests.jl | 3 +- test/test_aqua.jl | 7 +- 39 files changed, 2080 insertions(+), 71 deletions(-) delete mode 100644 ext/MeasureBaseDistributionsExt.jl create mode 100644 ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl create mode 100644 ext/MeasureBaseDistributionsExt/dirac.jl create mode 100644 ext/MeasureBaseDistributionsExt/dirichlet.jl create mode 100644 ext/MeasureBaseDistributionsExt/dist_vartransform.jl create mode 100644 ext/MeasureBaseDistributionsExt/distribution_measure.jl create mode 100644 ext/MeasureBaseDistributionsExt/measure_interface.jl create mode 100644 ext/MeasureBaseDistributionsExt/mixture.jl create mode 100644 ext/MeasureBaseDistributionsExt/product.jl create mode 100644 ext/MeasureBaseDistributionsExt/reshaped.jl create mode 100644 ext/MeasureBaseDistributionsExt/standard_dist.jl create mode 100644 ext/MeasureBaseDistributionsExt/standard_normal.jl create mode 100644 ext/MeasureBaseDistributionsExt/standard_uniform.jl create mode 100644 ext/MeasureBaseDistributionsExt/standardmv.jl create mode 100644 ext/MeasureBaseDistributionsExt/univariate.jl create mode 100644 ext/MeasureBaseDistributionsForwardDiffPullbacksExt.jl create mode 100644 ext/MeasureBaseDistributionsMooncakeExt.jl create mode 100644 ext/MeasureBaseForwardDiffPullbacksExt.jl create mode 100644 test/distributions/getjacobian.jl create mode 100644 test/distributions/test_autodiff_utils.jl create mode 100644 test/distributions/test_conversions.jl create mode 100644 test/distributions/test_distribution_measure.jl create mode 100644 test/distributions/test_distributions.jl create mode 100644 test/distributions/test_measure_interface.jl create mode 100644 test/distributions/test_mooncake.jl create mode 100644 test/distributions/test_standard_dist.jl create mode 100644 test/distributions/test_standard_normal.jl create mode 100644 test/distributions/test_standard_uniform.jl create mode 100644 test/distributions/test_transport.jl diff --git a/Project.toml b/Project.toml index 486afa73..e2bfee4a 100644 --- a/Project.toml +++ b/Project.toml @@ -4,6 +4,8 @@ version = "0.14.12" authors = ["Chad Scherrer ", "Oliver Schulz ", "contributors"] [deps] +ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197" +ArraysOfArrays = "65a8f2f4-9b39-5baf-92e2-a9cc46fdf018" ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0" Compat = "34da2185-b29b-5c13-b0c7-acf172513d20" ConstantRNGs = "aa9b60e7-6b1c-4c29-a6e5-e43521412437" @@ -37,17 +39,26 @@ Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +ForwardDiffPullbacks = "450a3b6d-2448-4ee1-8e34-e4eb8713b605" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" +PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" +StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" +StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c" [extensions] MeasureBaseChainRulesCoreExt = "ChainRulesCore" -MeasureBaseDistributionsExt = "Distributions" +MeasureBaseDistributionsExt = ["Distributions", "StatsBase", "StatsFuns", "PDMats"] MeasureBaseDistributionsChainRulesCoreExt = ["Distributions", "ChainRulesCore"] MeasureBaseDistributionsForwardDiffExt = ["Distributions", "ForwardDiff"] +MeasureBaseDistributionsForwardDiffPullbacksExt = ["Distributions", "ForwardDiffPullbacks", "ChainRulesCore"] +MeasureBaseDistributionsMooncakeExt = ["Distributions", "Mooncake"] MeasureBaseForwardDiffExt = "ForwardDiff" +MeasureBaseForwardDiffPullbacksExt = "ForwardDiffPullbacks" MeasureBaseMooncakeExt = "Mooncake" [compat] +ArgCheck = "1, 2" +ArraysOfArrays = "0.6" ChainRulesCore = "1" ChangesOfVariables = "0.1.3" Compat = "3.35, 4" @@ -55,9 +66,9 @@ ConstantRNGs = "0.1.1" ConstructionBase = "1.3" DensityInterface = "0.4" Distributions = "0.25.1" -Distributions = "0.25.111" FillArrays = "0.12, 0.13, 1" -ForwardDiff = "0.8, 0.9, 0.10" +ForwardDiff = "0.10, 1" +ForwardDiffPullbacks = "0.2" FunctionChains = "0.2.3" HeterogeneousComputing = "0.2.3" IfElse = "0.1" @@ -70,6 +81,7 @@ LogarithmicNumbers = "1" MappedArrays = "0.4" Mooncake = "0.5.34" NaNMath = "0.3, 1" +PDMats = "0.11" PrettyPrinting = "0.3, 0.4" PropertyFunctions = "0.2.2" Random = "1" @@ -79,6 +91,8 @@ Static = "0.8, 1" StaticArrays = "1.5" StaticThings = "0.2" Statistics = "1" +StatsBase = "0.33, 0.34" +StatsFuns = "0.9, 1, 2" Test = "1" Tricks = "0.1" julia = "1.10" diff --git a/ext/MeasureBaseChainRulesCoreExt.jl b/ext/MeasureBaseChainRulesCoreExt.jl index 0384a04b..25019da1 100644 --- a/ext/MeasureBaseChainRulesCoreExt.jl +++ b/ext/MeasureBaseChainRulesCoreExt.jl @@ -11,14 +11,25 @@ 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: _dropfront, _dropback, _rev_cumsum, _exp_cumsum_log +using MeasureBase: _pushfront, _pushback, _rev_cumsum, _exp_cumsum_log function ChainRulesCore.rrule(::typeof(_pushfront), v::AbstractVector, x) result = _pushfront(v, x) diff --git a/ext/MeasureBaseDistributionsChainRulesCoreExt.jl b/ext/MeasureBaseDistributionsChainRulesCoreExt.jl index 4dd3f4ff..63cc3e93 100644 --- a/ext/MeasureBaseDistributionsChainRulesCoreExt.jl +++ b/ext/MeasureBaseDistributionsChainRulesCoreExt.jl @@ -5,5 +5,14 @@ module MeasureBaseDistributionsChainRulesCoreExt using MeasureBase import Distributions import ChainRulesCore +using ChainRulesCore: NoTangent + +using MeasureBase: _dist_params_numtype +using Distributions: Distribution + +_dist_params_numtype_pullback(ΔΩ) = (NoTangent(), NoTangent()) +function ChainRulesCore.rrule(::typeof(_dist_params_numtype), d::Distribution) + _dist_params_numtype(d), _dist_params_numtype_pullback +end end # module MeasureBaseDistributionsChainRulesCoreExt diff --git a/ext/MeasureBaseDistributionsExt.jl b/ext/MeasureBaseDistributionsExt.jl deleted file mode 100644 index beb47821..00000000 --- a/ext/MeasureBaseDistributionsExt.jl +++ /dev/null @@ -1,8 +0,0 @@ -# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). - -module MeasureBaseDistributionsExt - -using MeasureBase -import Distributions - -end # module MeasureBaseDistributionsExt diff --git a/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl b/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl new file mode 100644 index 00000000..177ec306 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl @@ -0,0 +1,66 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +module MeasureBaseDistributionsExt + +using LinearAlgebra: Diagonal, diag, dot, cholesky + +import Random +using Random: AbstractRNG, rand! + +import DensityInterface +using DensityInterface: logdensityof, densityof + +import MeasureBase +using MeasureBase: AbstractMeasure, AsMeasure, asmeasure +using MeasureBase: Lebesgue, Counting, ℝ +using MeasureBase: StdMeasure, StdUniform, StdExponential, StdLogistic, StdNormal +using MeasureBase: PowerMeasure, WeightedMeasure, SuperpositionMeasure, PushforwardMeasure +using MeasureBase: basemeasure, rootmeasure, testvalue, productmeasure, pushfwd, superpose +using MeasureBase: getdof, checked_arg, massof +using MeasureBase: transport_to, transport_def, transport_origin, from_origin, to_origin +using MeasureBase: NoTransportOrigin, NoTransport +using MeasureBase: Reshape +using MeasureBase: convert_realtype, firsttype, _fwddiff, @_adignore +import MeasureBase: + _dist_params_numtype, _trafo_cdf_impl, _trafo_quantile_impl, _trafo_quantile_impl_generic +using MeasureBase: _pushfront, _pushback, _dropfront, _dropback, _rev_cumsum, _exp_cumsum_log + +import Distributions +using Distributions: Distribution, VariateForm, ValueSupport, ContinuousDistribution +using Distributions: Univariate, Multivariate, ArrayLikeVariate, Continuous, Discrete +using Distributions: Uniform, Exponential, Logistic, Normal +using Distributions: MvNormal, AbstractMvNormal, Beta, Dirichlet +using Distributions: ReshapedDistribution, AbstractMixtureModel + +import Statistics +import StatsBase +import StatsFuns +import PDMats + +using IrrationalConstants: log2π, invsqrt2π + +using HeterogeneousComputing: real_numtype + +using Static: True, False, StaticInt, static, dynamic +using StaticThings: asnonstatic +using FillArrays: Fill, Ones, Zeros + +using ArgCheck: @argcheck + +using ArraysOfArrays: ArrayOfSimilarArrays, flatview + +include("measure_interface.jl") +include("standard_dist.jl") +include("standard_uniform.jl") +include("standard_normal.jl") +include("distribution_measure.jl") +include("dist_vartransform.jl") +include("univariate.jl") +include("standardmv.jl") +include("product.jl") +include("reshaped.jl") +include("mixture.jl") +include("dirichlet.jl") +include("dirac.jl") + +end # module MeasureBaseDistributionsExt diff --git a/ext/MeasureBaseDistributionsExt/dirac.jl b/ext/MeasureBaseDistributionsExt/dirac.jl new file mode 100644 index 00000000..8580df8c --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/dirac.jl @@ -0,0 +1,14 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +MeasureBase.AbstractMeasure(obj::Distributions.Dirac) = MeasureBase.Dirac(obj.value) + +function AsMeasure{D}(::D) where {D<:Distributions.Dirac} + throw(ArgumentError("Don't wrap Distributions.Dirac into MeasureBase.AsMeasure, use asmeasure to convert instead.")) +end + + +Distributions.Distribution(m::MeasureBase.Dirac{<:Real}) = Distributions.Dirac(m.x) + +function Distributions.Distribution(@nospecialize(m::MeasureBase.Dirac{T})) where T + throw(ArgumentError("Can only convert MeasureBase.Dirac{<:Real} to Distributions.Dirac, but not MeasureBase.Dirac{<:$(nameof(T))}")) +end diff --git a/ext/MeasureBaseDistributionsExt/dirichlet.jl b/ext/MeasureBaseDistributionsExt/dirichlet.jl new file mode 100644 index 00000000..c60eeecd --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/dirichlet.jl @@ -0,0 +1,59 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +const DirichletMeasure = AsMeasure{<:Dirichlet} + +MeasureBase.getdof(d::Dirichlet) = length(d) - 1 +MeasureBase.getdof(m::DirichletMeasure) = getdof(m.obj) + +MeasureBase.transport_origin(d::Dirichlet) = StdUniform()^getdof(d) + + + +function _dirichlet_beta_trafo(α::Real, β::Real, x::Real) + R = float(promote_type(typeof(α), typeof(β), typeof(x))) + convert(R, transport_def(Beta(α, β), StdUniform(), x))::R +end + +_a_times_one_minus_b(a::Real, b::Real) = a * (1 - b) + +function MeasureBase.from_origin(ν::Dirichlet, x) + # See M. J. Betancourt, "Cruising The Simplex: Hamiltonian Monte Carlo and the Dirichlet Distribution", + # https://arxiv.org/abs/1010.3436 + + @_adignore @argcheck length(ν) == length(x) + 1 + + αs = _dropfront(_rev_cumsum(ν.alpha)) + βs = _dropback(ν.alpha) + beta_v = _fwddiff(_dirichlet_beta_trafo).(αs, βs, x) + beta_v_cp = _exp_cumsum_log(_pushfront(beta_v, 1)) + beta_v_ext = _pushback(beta_v, 0) + _fwddiff(_a_times_one_minus_b).(beta_v_cp, beta_v_ext) +end + + +function _inv_dirichlet_beta_trafo(α::Real, β::Real, beta_v::Real) + R = float(promote_type(typeof(α), typeof(β), typeof(beta_v))) + convert(R, transport_def(StdUniform(), Beta(α, β), beta_v))::R +end + +# ToDo: Find efficient pullback for this: +function _dirichlet_variate_to_beta_v(y::AbstractVector{<:Real}) + beta_v = similar(y, length(eachindex(y)) - 1) + @assert firstindex(beta_v) == firstindex(y) + @assert lastindex(beta_v) == lastindex(y) - 1 + T = eltype(y) + sum_log_beta_v::T = 0 + @inbounds for i in eachindex(beta_v) + beta_v[i] = 1 - y[i] / exp(sum_log_beta_v) + sum_log_beta_v += log(beta_v[i]) + end + return beta_v +end + +function MeasureBase.to_origin(ν::Dirichlet, y) + @_adignore @argcheck length(ν) == length(y) + αs = _dropfront(_rev_cumsum(ν.alpha)) + βs = _dropback(ν.alpha) + beta_v = _dirichlet_variate_to_beta_v(y) + _fwddiff(_inv_dirichlet_beta_trafo).(αs, βs, beta_v) +end diff --git a/ext/MeasureBaseDistributionsExt/dist_vartransform.jl b/ext/MeasureBaseDistributionsExt/dist_vartransform.jl new file mode 100644 index 00000000..ceedabe9 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/dist_vartransform.jl @@ -0,0 +1,16 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +const _AnyStdUniform = Union{StandardUniform,Uniform} +const _AnyStdNormal = Union{StandardNormal,Normal} + +const _AnyStdDistribution = Union{_AnyStdUniform,_AnyStdNormal} + +_std_dist(::Type{<:_AnyStdUniform}) = StandardUniform +_std_dist(::Type{<:_AnyStdNormal}) = StandardNormal + +_std_dist(::Type{D}, ::StaticInt{1}) where {D<:_AnyStdDistribution} = D() +_std_dist(::Type{D}, dof) where {D<:_AnyStdDistribution} = D(dynamic(dof)) +_std_dist_for(::Type{D}, μ::Any) where {D<:_AnyStdDistribution} = _std_dist(_std_dist(D), getdof(μ)) + +MeasureBase.transport_to(::Type{NU}, μ) where {NU<:_AnyStdDistribution} = transport_to(_std_dist_for(NU, μ), μ) +MeasureBase.transport_to(ν, ::Type{MU}) where {MU<:_AnyStdDistribution} = transport_to(ν, _std_dist_for(MU, ν)) diff --git a/ext/MeasureBaseDistributionsExt/distribution_measure.jl b/ext/MeasureBaseDistributionsExt/distribution_measure.jl new file mode 100644 index 00000000..bcbfd558 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/distribution_measure.jl @@ -0,0 +1,74 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + + +const DistributionMeasure{F<:VariateForm,S<:ValueSupport,D<:Distribution{F,S}} = AsMeasure{D} + +@inline MeasureBase.AbstractMeasure(obj::Distribution) = AsMeasure{typeof(obj)}(obj) +@inline Base.convert(::Type{AbstractMeasure}, obj::Distribution) = AbstractMeasure(obj) + +@inline Distributions.Distribution(m::DistributionMeasure) = m.obj +@inline Distributions.Distribution{F}(m::DistributionMeasure{F}) where {F<:VariateForm} = Distribution(m) +@inline Distributions.Distribution{F,S}(m::DistributionMeasure{F,S}) where {F<:VariateForm,S<:ValueSupport} = Distribution(m) + +@inline Base.convert(::Type{Distribution}, m::DistributionMeasure) = Distribution(m) +@inline Base.convert(::Type{Distribution{F}}, m::DistributionMeasure{F}) where {F<:VariateForm} = Distribution(m) +@inline Base.convert(::Type{Distribution{F,S}}, m::DistributionMeasure{F,S}) where {F<:VariateForm,S<:ValueSupport} = Distribution(m) + + +Base.rand(rng::AbstractRNG, ::Type{T}, m::DistributionMeasure) where {T<:Real} = convert_realtype(T, rand(m.obj)) + +function _flat_powrand(rng::AbstractRNG, ::Type{T}, d::Distribution{<:ArrayLikeVariate{0}}, sz::Dims) where {T<:Real} + convert_realtype(T, reshape(rand(rng, d, prod(sz)), sz...)) +end + +function _flat_powrand(rng::AbstractRNG, ::Type{T}, d::Distribution{<:ArrayLikeVariate{1}}, sz::Dims) where {T<:Real} + convert_realtype(T, reshape(rand(rng, d, prod(sz)), size(d)..., sz...)) +end + +function _flat_powrand(rng::AbstractRNG, ::Type{T}, d::ReshapedDistribution{N,<:Any,<:Distribution{<:ArrayLikeVariate{1}}}, sz::Dims) where {T<:Real,N} + convert_realtype(T, reshape(rand(rng, d.dist, prod(sz)), d.dims..., sz...)) +end + +function _flat_powrand(rng::AbstractRNG, ::Type{T}, d::Distribution, sz::Dims) where {T<:Real} + flatview(ArrayOfSimilarArrays(convert_realtype(T, rand(rng, d, sz)))) +end + +function Base.rand(rng::AbstractRNG, ::Type{T}, m::PowerMeasure{<:DistributionMeasure{<:ArrayLikeVariate{0}}, NTuple{N,Base.OneTo{Int}}}) where {T<:Real,N} + _flat_powrand(rng, T, m.parent.obj, map(length, m.axes)) +end + +function Base.rand(rng::AbstractRNG, ::Type{T}, m::PowerMeasure{<:DistributionMeasure{<:ArrayLikeVariate{M}}, NTuple{N,Base.OneTo{Int}}}) where {T<:Real,M,N} + flat_data = _flat_powrand(rng, T, m.parent.obj, map(length, m.axes)) + ArrayOfSimilarArrays{T,M,N}(flat_data) +end + + +@inline DensityInterface.densityof(m::DistributionMeasure) = densityof(m.obj) +@inline DensityInterface.logdensityof(m::DistributionMeasure) = logdensityof(m.obj) + +@inline MeasureBase.logdensity_def(m::DistributionMeasure, x) = DensityInterface.logdensityof(m.obj, x) +@inline MeasureBase.unsafe_logdensityof(m::DistributionMeasure, x) = DensityInterface.logdensityof(m.obj, x) +@inline MeasureBase.insupport(m::DistributionMeasure, x) = Distributions.insupport(m.obj, x) + +@inline MeasureBase.rootmeasure(m::DistributionMeasure{<:ArrayLikeVariate{0},<:Continuous}) = Lebesgue() +@inline MeasureBase.rootmeasure(m::DistributionMeasure{<:ArrayLikeVariate,<:Continuous}) = Lebesgue()^size(m.obj) +@inline MeasureBase.rootmeasure(m::DistributionMeasure{<:ArrayLikeVariate{0},<:Discrete}) = Counting() +@inline MeasureBase.rootmeasure(m::DistributionMeasure{<:ArrayLikeVariate,<:Discrete}) = Counting()^size(m.obj) + +@inline MeasureBase.basemeasure(m::DistributionMeasure) = rootmeasure(m) + +@inline MeasureBase.massof(::DistributionMeasure) = static(1.0) + +@inline MeasureBase.mspace_elsize(m::DistributionMeasure{<:ArrayLikeVariate}) = size(m.obj) + +@inline MeasureBase.getdof(m::DistributionMeasure{<:ArrayLikeVariate{0}}) = 1 + +# Delegate transport to the wrapped distribution: +@inline MeasureBase.transport_origin(m::DistributionMeasure) = m.obj +@inline MeasureBase.to_origin(::DistributionMeasure, y) = y +@inline MeasureBase.from_origin(::DistributionMeasure, x) = x + +@inline MeasureBase.paramnames(m::DistributionMeasure) = propertynames(m.obj) +@inline MeasureBase.params(m::DistributionMeasure) = NamedTuple{propertynames(m.obj)}(Distributions.params(m.obj)) + +# @inline MeasureBase.testvalue(m::DistributionMeasure) = testvalue(basemeasure(d)) diff --git a/ext/MeasureBaseDistributionsExt/measure_interface.jl b/ext/MeasureBaseDistributionsExt/measure_interface.jl new file mode 100644 index 00000000..6fed5d4a --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/measure_interface.jl @@ -0,0 +1,26 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +@inline MeasureBase.logdensity_def(d::Distribution, x) = DensityInterface.logdensityof(d, x) +@inline MeasureBase.unsafe_logdensityof(d::Distribution, x) = DensityInterface.logdensityof(d, x) + +@inline MeasureBase.insupport(d::Distribution, x) = Distributions.insupport(d, x) + +@inline MeasureBase.basemeasure(d::Distribution{<:ArrayLikeVariate{0},<:Continuous}) = Lebesgue() +@inline MeasureBase.basemeasure(d::Distribution{<:ArrayLikeVariate,<:Continuous}) = Lebesgue()^size(d) +@inline MeasureBase.basemeasure(d::Distribution{<:ArrayLikeVariate{0},<:Discrete}) = Counting() +@inline MeasureBase.basemeasure(d::Distribution{<:ArrayLikeVariate,<:Discrete}) = Counting()^size(d) + +@inline MeasureBase.paramnames(d::Distribution) = propertynames(d) +@inline MeasureBase.params(d::Distribution) = NamedTuple{propertynames(d)}(Distributions.params(d)) + +@inline MeasureBase.testvalue(d::Distribution) = testvalue(basemeasure(d)) +@inline MeasureBase.testvalue(::Type{T}, d::Distribution) where {T} = testvalue(T, basemeasure(d)) + + +@inline MeasureBase.basemeasure(d::Distributions.Poisson) = + Counting(MeasureBase.BoundedInts(static(0), static(Inf))) +@inline MeasureBase.basemeasure(d::Distributions.Product{<:Any,<:Distributions.Poisson}) = + Counting(MeasureBase.BoundedInts(static(0), static(Inf)))^size(d) + + +MeasureBase.∫(f, base::Distribution) = MeasureBase.∫(f, convert(AbstractMeasure, base)) diff --git a/ext/MeasureBaseDistributionsExt/mixture.jl b/ext/MeasureBaseDistributionsExt/mixture.jl new file mode 100644 index 00000000..89d41acf --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/mixture.jl @@ -0,0 +1,24 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +function MeasureBase.AbstractMeasure(d::Distributions.AbstractMixtureModel) + superpose(map((w, c) -> w * asmeasure(c), Distributions.probs(d), Distributions.components(d))) +end + +function AsMeasure{D}(::D) where {D<:Distributions.AbstractMixtureModel} + throw(ArgumentError("Don't wrap Distributions.AbstractMixtureModel into MeasureBase.AsMeasure, use asmeasure to convert instead.")) +end + + +const _MixtureMeasure = SuperpositionMeasure{ + <:Union{Tuple{Vararg{WeightedMeasure}},AbstractVector{<:WeightedMeasure}}, +} + +_mixture_component(m::AsMeasure{<:Distribution}) = m.obj + +function Distributions.Distribution(m::_MixtureMeasure) + components = map(c -> _mixture_component(c.base), collect(values(m.components))) + prior = map(c -> exp(c.logweight), collect(values(m.components))) + Distributions.MixtureModel(components, prior) +end + +Base.convert(::Type{Distribution}, m::_MixtureMeasure) = Distributions.Distribution(m) diff --git a/ext/MeasureBaseDistributionsExt/product.jl b/ext/MeasureBaseDistributionsExt/product.jl new file mode 100644 index 00000000..a050dc97 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/product.jl @@ -0,0 +1,56 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +@static if isdefined(Distributions, :Product) + MeasureBase.AbstractMeasure(obj::Distributions.Product) = productmeasure(map(asmeasure, obj.v)) + + function AsMeasure{D}(::D) where {D<:Distributions.Product} + throw(ArgumentError("Don't wrap Distributions.Product into MeasureBase.AsMeasure, use asmeasure to convert instead.")) + end +end + +function Distributions.Distribution( + m::MeasureBase.ProductMeasure{<:AbstractArray{<:AsMeasure{<:Distribution{Univariate}}}}, +) + Distributions.product_distribution(map(x -> x.obj, MeasureBase.marginals(m))) +end + +function Base.convert( + ::Type{Distribution}, + m::MeasureBase.ProductMeasure{<:AbstractArray{<:AsMeasure{<:Distribution{Univariate}}}}, +) + Distributions.Distribution(m) +end + +@static if isdefined(Distributions, :ProductDistribution) + MeasureBase.AbstractMeasure(obj::Distributions.ProductDistribution) = productmeasure(map(asmeasure, obj.dists)) + + function AsMeasure{D}(::D) where {D<:Distributions.ProductDistribution} + throw(ArgumentError("Don't wrap Distributions.ProductDistribution into MeasureBase.AsMeasure, use asmeasure to convert instead.")) + end + + function Distributions.Distribution( + m::MeasureBase.ProductMeasure{<:AbstractArray{<:AsMeasure{<:Distribution}}}, + ) + Distributions.product_distribution(map(x -> x.obj, MeasureBase.marginals(m))) + end + + function Distributions.Distribution( + m::MeasureBase.ProductMeasure{<:Tuple{Vararg{AsMeasure{<:Distribution}}}}, + ) + Distributions.product_distribution(map(x -> x.obj, MeasureBase.marginals(m))...) + end + + function Base.convert( + ::Type{Distribution}, + m::MeasureBase.ProductMeasure{<:AbstractArray{<:AsMeasure{<:Distribution}}}, + ) + Distributions.Distribution(m) + end + + function Base.convert( + ::Type{Distribution}, + m::MeasureBase.ProductMeasure{<:Tuple{Vararg{AsMeasure{<:Distribution}}}}, + ) + Distributions.Distribution(m) + end +end diff --git a/ext/MeasureBaseDistributionsExt/reshaped.jl b/ext/MeasureBaseDistributionsExt/reshaped.jl new file mode 100644 index 00000000..6efde609 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/reshaped.jl @@ -0,0 +1,27 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +MeasureBase.getdof(μ::ReshapedDistribution) = MeasureBase.getdof(μ.dist) + +MeasureBase.transport_origin(μ::ReshapedDistribution) = μ.dist + +MeasureBase.to_origin(ν::ReshapedDistribution, y) = reshape(y, size(ν.dist)) + +MeasureBase.from_origin(ν::ReshapedDistribution, x) = reshape(x, ν.dims) + + +function MeasureBase.AbstractMeasure(d::Distributions.ReshapedDistribution) + orig_dist = d.dist + pushfwd(Reshape(size(d), size(orig_dist)), AbstractMeasure(orig_dist)) +end + +function AsMeasure{D}(::D) where {D<:Distributions.ReshapedDistribution} + throw(ArgumentError("Don't wrap Distributions.ReshapedDistribution into MeasureBase.AsMeasure, use asmeasure to convert instead.")) +end + + +function Distributions.Distribution(m::PushforwardMeasure{<:Reshape}) + reshape(Distributions.Distribution(m.origin), asnonstatic(m.f.output_size)...) +end + +Base.convert(::Type{Distribution}, m::PushforwardMeasure{<:Reshape}) = + Distributions.Distribution(m) diff --git a/ext/MeasureBaseDistributionsExt/standard_dist.jl b/ext/MeasureBaseDistributionsExt/standard_dist.jl new file mode 100644 index 00000000..010b20a9 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/standard_dist.jl @@ -0,0 +1,222 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +""" + struct StandardDist{D<:Distribution{Univariate,Continuous},N} <: Distributions.Distribution{ArrayLikeVariate{N},Continuous} + +Represents `D()` or a product distribution of `D()` in a dispatchable fashion. + +Constructor: +``` + StandardDist{Uniform}(size...) + StandardDist{Normal}(size...) +``` +""" +struct StandardDist{D<:Distribution{Univariate,Continuous},N,U<:Integer} <: + Distributions.Distribution{ArrayLikeVariate{N},Continuous} + _size::NTuple{N,U} +end +export StandardDist + +StandardDist{D}() where {D<:Distribution{Univariate,Continuous}} = + StandardDist{D,0,Int}(()) +StandardDist{D}(dims::Vararg{U,N}) where {D<:Distribution{Univariate,Continuous},N,U<:Integer} = + StandardDist{D,N,U}((dims...,)) + + +const StandardUnivariateDist{D<:Distribution{Univariate,Continuous},U<:Integer} = StandardDist{D,0,U} +const StandardMultivariteDist{D<:Distribution{Univariate,Continuous},U<:Integer} = StandardDist{D,1,U} + + +function Base.show(io::IO, d::StandardDist{D}) where {D} + print(io, nameof(typeof(d)), "{", D, "}") + show(io, d._size) +end + + +@inline MeasureBase.transport_def(::MU, μ::MU, x) where {MU<:StandardDist{<:Any,0}} = x + +for (A, B) in [ + (Uniform, StdUniform), + (Exponential, StdExponential), + (Logistic, StdLogistic), + (Normal, StdNormal) +] + @eval begin + @inline MeasureBase.transport_origin(d::StandardDist{$A,0}) = $B() + @inline MeasureBase.transport_origin(d::StandardDist{$A,N}) where {N} = $B()^size(d) + + # StandardDist{$A} and $B are equivalent as measures, so convert + # instead of wrapping: + MeasureBase.AbstractMeasure(::StandardDist{$A,0}) = $B() + MeasureBase.AbstractMeasure(d::StandardDist{$A,N}) where {N} = $B()^size(d) + + Distributions.Distribution(::$B) = StandardDist{$A}() + Base.convert(::Type{Distribution}, ::$B) = StandardDist{$A}() + + function Distributions.Distribution(m::PowerMeasure{$B}) + StandardDist{$A}(map(dynamic ∘ length, m.axes)...) + end + Base.convert(::Type{Distribution}, m::PowerMeasure{$B}) = Distributions.Distribution(m) + end +end + +@inline MeasureBase.to_origin(ν::StandardDist, y) = y +@inline MeasureBase.from_origin(ν::StandardDist, x) = x + + +@inline nonstddist(::StandardDist{D,0}) where {D} = D(Distributions.params(D())...) +@inline function nonstddist(d::StandardDist{D,N}) where {D,N} + nonstd0 = nonstddist(StandardDist{D}()) + reshape(Distributions.product_distribution(fill(nonstd0, length(d))), size(d)) +end + + +(::Type{D})(d::StandardDist{D,0}) where {D<:Distribution{Univariate,Continuous}} = nonstddist(d) + +# TODO: Replace `fill` by `FillArrays.Fill` once Distributions fully supports this: +(::Type{Distributions.Product})(d::StandardDist{D,1}) where {D} = + Distributions.Product(fill(StandardDist{D}(), length(d))) + +Base.convert(::Type{D}, d::StandardDist{D,0}) where {D<:Distribution{Univariate,Continuous}} = D(d) +Base.convert(::Type{Distributions.Product}, d::StandardDist{D,1}) where {D} = + Distributions.Product(d) + + + +@inline Base.size(d::StandardDist) = d._size +@inline Base.length(d::StandardDist) = prod(size(d)) + +Base.eltype(::Type{<:StandardDist}) = Float64 + +@inline Distributions.partype(d::StandardDist{D}) where {D} = Float64 + +@inline StatsBase.params(d::StandardDist) = () + +for f in ( + :(Base.minimum), + :(Base.maximum), + :(Statistics.mean), + :(Statistics.median), + :(StatsBase.mode), + :(Statistics.var), + :(Statistics.std), + :(StatsBase.skewness), + :(StatsBase.kurtosis), + :(Distributions.location), + :(Distributions.scale), +) + @eval begin + ($f)(d::StandardDist{D,0}) where {D} = ($f)(nonstddist(d)) + ($f)(d::StandardDist{D,N}) where {D,N} = Fill(($f)(StandardDist{D}()), size(d)...) + end +end + +StatsBase.modes(d::StandardDist) = [StatsBase.mode(d)] + +# ToDo: Define cov for N!=1? +Statistics.cov(d::StandardDist{D,1}) where {D} = Diagonal(Statistics.var(d)) +Distributions.invcov(d::StandardDist{D,1}) where {D} = + Diagonal(Fill(inv(Statistics.var(StandardDist{D}())), length(d))) +Distributions.logdetcov(d::StandardDist{D,1}) where {D} = + length(d) * log(Statistics.var(StandardDist{D}())) + +StatsBase.entropy(d::StandardDist{D,0}) where {D} = StatsBase.entropy(nonstddist(d)) +StatsBase.entropy(d::StandardDist{D,N}) where {D,N} = + length(d) * StatsBase.entropy(StandardDist{D}()) + + +Distributions.insupport(d::StandardDist{D,0}, x::Real) where {D} = + Distributions.insupport(nonstddist(d), x) + +function Distributions.insupport(d::StandardDist{D,N}, x::AbstractArray{<:Real,N}) where {D,N} + all(Base.Fix1(Distributions.insupport, StandardDist{D}()), checked_arg(d, x)) +end + + +@inline Distributions.logpdf(d::StandardDist{D,0}, x::U) where {D,U} = + Distributions.logpdf(nonstddist(d), x) + +function Distributions.logpdf(d::StandardDist{D,N}, x::AbstractArray{<:Real,N}) where {D,N} + Distributions._logpdf(d, checked_arg(d, x)) +end + +# Explicit N=1/N=2 methods to avoid dispatch ambiguities with Distributions: +function Distributions._logpdf(::StandardDist{D,1}, x::AbstractArray{<:Real,1}) where {D} + sum(Base.Fix1(Distributions.logpdf, StandardDist{D}()), x) +end + +function Distributions._logpdf(::StandardDist{D,2}, x::AbstractArray{<:Real,2}) where {D} + sum(Base.Fix1(Distributions.logpdf, StandardDist{D}()), x) +end + +function Distributions._logpdf(::StandardDist{D,N}, x::AbstractArray{<:Real,N}) where {D,N} + sum(Base.Fix1(Distributions.logpdf, StandardDist{D}()), x) +end + + +Distributions.gradlogpdf(d::StandardDist{D,0}, x::Real) where {D} = + Distributions.gradlogpdf(nonstddist(d), x) + +function Distributions.gradlogpdf(d::StandardDist{D,N}, x::AbstractArray{<:Real,N}) where {D,N} + Distributions.gradlogpdf.(StandardDist{D}(), checked_arg(d, x)) +end + + +# Explicit N=1/N=2 methods to avoid dispatch ambiguities with Distributions: +function Distributions.pdf(d::StandardDist{D,1}, x::AbstractVector{U}) where {D,U<:Real} + Distributions._pdf(d, checked_arg(d, x)) +end + +function Distributions._pdf(d::StandardDist{D,1}, x::AbstractVector{U}) where {D,U<:Real} + exp(Distributions._logpdf(d, x)) +end + +function Distributions.pdf(d::StandardDist{D,2}, x::AbstractMatrix{U}) where {D,U<:Real} + Distributions._pdf(d, checked_arg(d, x)) +end + +function Distributions._pdf(d::StandardDist{D,2}, x::AbstractMatrix{U}) where {D,U<:Real} + exp(Distributions._logpdf(d, x)) +end + +function Distributions.pdf(d::StandardDist{D,N}, x::AbstractArray{U,N}) where {D,N,U<:Real} + Distributions._pdf(d, checked_arg(d, x)) +end + +function Distributions._pdf(d::StandardDist{D,N}, x::AbstractArray{U,N}) where {D,N,U<:Real} + exp(Distributions._logpdf(d, x)) +end + + +for f in ( + :(Distributions.logcdf), + :(Distributions.cdf), + :(Distributions.logccdf), + :(Distributions.ccdf), + :(Distributions.quantile), + :(Distributions.cquantile), + :(Distributions.invlogcdf), + :(Distributions.invlogccdf), + :(Distributions.mgf), + :(Distributions.cf), +) + @eval begin + @inline ($f)(d::StandardDist, x::Real) = ($f)(nonstddist(d), x) + end +end + + +Base.rand(rng::AbstractRNG, d::StandardDist{D,0}) where {D} = rand(rng, nonstddist(d)) +Random.rand!(rng::AbstractRNG, d::StandardDist{D,0}, x::AbstractArray{<:Real,0}) where {D} = + (x[] = rand(rng, d); return x) +Random.rand!(rng::AbstractRNG, d::StandardDist{D,N}, x::AbstractArray{<:Real,N}) where {D,N} = + rand!(rng, StandardDist{D}(), x) + + +Distributions.truncated(d::StandardDist{D,0}, l::Real, u::Real) where {D} = + Distributions.truncated(nonstddist(d), l, u) + +Distributions.product_distribution(dists::AbstractVector{<:StandardDist{D,0}}) where {D} = + StandardDist{D}(size(dists)...) +Distributions.product_distribution(dists::AbstractArray{<:StandardDist{D,0}}) where {D} = + StandardDist{D}(size(dists)...) diff --git a/ext/MeasureBaseDistributionsExt/standard_normal.jl b/ext/MeasureBaseDistributionsExt/standard_normal.jl new file mode 100644 index 00000000..6bc27d04 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/standard_normal.jl @@ -0,0 +1,77 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +""" + const StandardNormal{N} = StandardDist{Normal,N} + +The standard normal distribution, scalar (`N == 0`) or as a product over an +array of rank `N`. +""" +const StandardNormal{N} = StandardDist{Normal,N} +export StandardNormal + +Distributions.Normal(d::StandardDist{Normal,0}) = Distributions.Normal() + +Distributions.MvNormal(d::StandardDist{Normal,1}) = MvNormal(PDMats.ScalMat(length(d), 1)) +Base.convert(::Type{Distributions.MvNormal}, d::StandardDist{Normal,1}) = + Distributions.MvNormal(d) + +Base.minimum(d::StandardDist{Normal,0}) = -Inf +Base.maximum(d::StandardDist{Normal,0}) = +Inf + +Distributions.insupport(d::StandardDist{Normal,0}, x::Real) = !isnan(x) + +Distributions.location(d::StandardDist{Normal,0}) = Statistics.mean(d) +Distributions.scale(d::StandardDist{Normal,0}) = Statistics.var(d) + +Statistics.mean(d::StandardDist{Normal,0}) = 0 +Statistics.mean(d::StandardDist{Normal,N}) where {N} = Zeros{Int}(size(d)...) + +Statistics.median(d::StandardDist{Normal}) = Statistics.mean(d) +StatsBase.mode(d::StandardDist{Normal}) = Statistics.mean(d) + +StatsBase.modes(d::StandardDist{Normal,0}) = Zeros{Int}(1) + +Statistics.var(d::StandardDist{Normal,0}) = 1 +Statistics.var(d::StandardDist{Normal,N}) where {N} = Ones{Int}(size(d)...) + +Statistics.std(d::StandardDist{Normal,0}) = 1 +Statistics.std(d::StandardDist{Normal,N}) where {N} = Ones{Int}(size(d)...) + +StatsBase.skewness(d::StandardDist{Normal,0}) = 0 +StatsBase.kurtosis(d::StandardDist{Normal,0}) = 0 + +StatsBase.entropy(d::StandardDist{Normal,0}) = muladd(log2π, 1 / 2, 1 / 2) + +Distributions.logpdf(d::StandardDist{Normal,0}, x::U) where {U<:Real} = + muladd(abs2(x), -U(1) / U(2), -log2π / U(2)) +Distributions.pdf(d::StandardDist{Normal,0}, x::U) where {U<:Real} = + invsqrt2π * exp(-abs2(x) / U(2)) + +@inline Distributions.gradlogpdf(d::StandardDist{Normal,0}, x::Real) = -x + +@inline Distributions.logcdf(d::StandardDist{Normal,0}, x::Real) = StatsFuns.normlogcdf(x) +@inline Distributions.cdf(d::StandardDist{Normal,0}, x::Real) = StatsFuns.normcdf(x) +@inline Distributions.logccdf(d::StandardDist{Normal,0}, x::Real) = StatsFuns.normlogccdf(x) +@inline Distributions.ccdf(d::StandardDist{Normal,0}, x::Real) = StatsFuns.normccdf(x) +@inline Distributions.quantile(d::StandardDist{Normal,0}, p::Real) = StatsFuns.norminvcdf(p) +@inline Distributions.cquantile(d::StandardDist{Normal,0}, p::Real) = StatsFuns.norminvccdf(p) +@inline Distributions.invlogcdf(d::StandardDist{Normal,0}, p::Real) = StatsFuns.norminvlogcdf(p) +@inline Distributions.invlogccdf(d::StandardDist{Normal,0}, p::Real) = StatsFuns.norminvlogccdf(p) + +Base.rand(rng::AbstractRNG, d::StandardDist{Normal,0}) = randn(rng) +Base.rand(rng::AbstractRNG, d::StandardDist{Normal,N}) where {N} = randn(rng, size(d)...) +Random.rand!(rng::AbstractRNG, d::StandardDist{Normal,N}, x::AbstractArray{<:Real,N}) where {N} = + Random.randn!(rng, x) + +Distributions.invcov(d::StandardDist{Normal,1}) = Distributions.cov(d) +Distributions.logdetcov(d::StandardDist{Normal,1}) = 0 + + +function Distributions.sqmahal(d::StandardDist{Normal,N}, x::AbstractArray{<:Real,N}) where {N} + dot(x, checked_arg(d, x)) +end + +function Distributions.sqmahal!(r::AbstractVector, d::StandardDist{Normal,N}, x::AbstractMatrix) where {N} + x_cols = eachcol(checked_arg(d, first(eachcol(x)))) + r .= dot.(x_cols, x_cols) +end diff --git a/ext/MeasureBaseDistributionsExt/standard_uniform.jl b/ext/MeasureBaseDistributionsExt/standard_uniform.jl new file mode 100644 index 00000000..51398fcb --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/standard_uniform.jl @@ -0,0 +1,79 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +""" + const StandardUniform{N} = StandardDist{Uniform,N} + +The standard uniform distribution, scalar (`N == 0`) or as a product over an +array of rank `N`. +""" +const StandardUniform{N} = StandardDist{Uniform,N} +export StandardUniform + +Distributions.Uniform(d::StandardDist{Uniform,0}) = Distributions.Uniform() + +Base.minimum(::StandardDist{Uniform,0}) = 0 +Base.maximum(::StandardDist{Uniform,0}) = 1 + +Distributions.location(::StandardDist{Uniform,0}) = 0 +Distributions.scale(::StandardDist{Uniform,0}) = 1 + +Statistics.mean(d::StandardDist{Uniform,0}) = 1 // 2 +Statistics.median(d::StandardDist{Uniform,0}) = Statistics.mean(d) +StatsBase.mode(d::StandardDist{Uniform,0}) = Statistics.mean(d) +StatsBase.modes(d::StandardDist{Uniform,0}) = Zeros{Int}(0) +StatsBase.modes(d::StandardDist{Uniform,N}) where {N} = Fill(Zeros{Int}(size(d))) + +Statistics.var(d::StandardDist{Uniform,0}) = 1 // 12 +Statistics.std(d::StandardDist{Uniform,0}) = sqrt(Statistics.var(d)) +StatsBase.skewness(d::StandardDist{Uniform,0}) = 0 +StatsBase.kurtosis(d::StandardDist{Uniform,0}) = -6 // 5 + +StatsBase.entropy(d::StandardDist{Uniform,0}) = 0 + + +function Distributions.logpdf(d::StandardDist{Uniform,0}, x::U) where {U<:Real} + ifelse(Distributions.insupport(d, x), U(0), U(-Inf)) +end + +function Distributions.pdf(d::StandardDist{Uniform,0}, x::U) where {U<:Real} + ifelse(Distributions.insupport(d, x), one(U), zero(U)) +end + + +Distributions.logcdf(d::StandardDist{Uniform,0}, x::U) where {U<:Real} = + log(Distributions.cdf(d, x)) + +function Distributions.cdf(d::StandardDist{Uniform,0}, x::U) where {U<:Real} + ifelse(x < zero(U), zero(U), ifelse(x < one(U), x, one(U))) +end + +Distributions.logccdf(d::StandardDist{Uniform,0}, x::U) where {U<:Real} = + log(Distributions.ccdf(d, x)) + +Distributions.ccdf(d::StandardDist{Uniform,0}, x::U) where {U<:Real} = + one(x) - Distributions.cdf(d, x) + + +function Distributions.quantile(d::StandardDist{Uniform,0}, p::U) where {U<:Real} + convert(float(U), p) +end + +function Distributions.cquantile(d::StandardDist{Uniform,0}, p::U) where {U<:Real} + y = Distributions.quantile(d, p) + one(y) - y +end + + +Distributions.mgf(d::StandardDist{Uniform,0}, t::Real) = Distributions.mgf(nonstddist(d), t) +Distributions.cf(d::StandardDist{Uniform,0}, t::Real) = Distributions.cf(nonstddist(d), t) + +Distributions.gradlogpdf(d::StandardDist{Uniform,0}, x::Real) = zero(x) + +function Distributions.gradlogpdf(d::StandardDist{Uniform,N}, x::AbstractArray{<:Real,N}) where {N} + zero(checked_arg(d, x)) +end + +Base.rand(rng::AbstractRNG, d::StandardDist{Uniform,0}) = rand(rng) +Base.rand(rng::AbstractRNG, d::StandardDist{Uniform,N}) where {N} = rand(rng, size(d)...) +Random.rand!(rng::AbstractRNG, d::StandardDist{Uniform,N}, x::AbstractArray{<:Real,N}) where {N} = + rand!(rng, x) diff --git a/ext/MeasureBaseDistributionsExt/standardmv.jl b/ext/MeasureBaseDistributionsExt/standardmv.jl new file mode 100644 index 00000000..c8e99039 --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/standardmv.jl @@ -0,0 +1,36 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + + +MeasureBase.getdof(d::AbstractMvNormal) = length(d) +MeasureBase.getdof(m::AsMeasure{<:AbstractMvNormal}) = getdof(m.obj) + +MeasureBase.transport_origin(ν::MvNormal) = StandardDist{Normal}(length(ν)) + +_cholesky_L(A) = cholesky(A).L +_cholesky_L(A::Diagonal{<:Real}) = Diagonal(sqrt.(diag(A))) +_cholesky_L(A::PDMats.PDiagMat{<:Real}) = Diagonal(sqrt.(A.diag)) +_cholesky_L(A::PDMats.ScalMat{<:Real}) = Diagonal(Fill(sqrt(A.value), A.dim)) + +function MeasureBase.from_origin(ν::MvNormal, x) + A = _cholesky_L(ν.Σ) + b = ν.μ + muladd(A, x, b) +end + +function MeasureBase.to_origin(ν::MvNormal, y) + A = _cholesky_L(ν.Σ) + b = ν.μ + A \ (y - b) +end + + +#DirichletMultinomial +#Distributions.AbstractMvLogNormal +#Distributions.AbstractMvTDist +#Distributions.ProductDistribution{1} +#Distributions.ReshapedDistribution{1, S, D} where {S<:ValueSupport, D<:(Distribution{<:ArrayLikeVariate, S})} +#JointOrderStatistics +#Multinomial +#MultivariateMixture (alias for AbstractMixtureModel{ArrayLikeVariate{1}}) +#MvLogitNormal +#VonMisesFisher diff --git a/ext/MeasureBaseDistributionsExt/univariate.jl b/ext/MeasureBaseDistributionsExt/univariate.jl new file mode 100644 index 00000000..126fe36d --- /dev/null +++ b/ext/MeasureBaseDistributionsExt/univariate.jl @@ -0,0 +1,142 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + + +@inline MeasureBase.getdof(::Distribution{Univariate}) = static(1) + +@inline MeasureBase.check_dof(a::Distribution{Univariate}, b::Distribution{Univariate}) = nothing + + +# Generic transformations to/from StdUniform via cdf/quantile: + + +_dist_params_numtype(d::Distribution) = real_numtype(typeof(Distributions.params(d))) + + +@inline _trafo_cdf(d::Distribution{Univariate,Continuous}, x::Real) = + _trafo_cdf_impl(_dist_params_numtype(d), d, x) + +@inline _trafo_cdf_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, x::Real) = + Distributions.cdf(d, x) + + +@inline _trafo_quantile(d::Distribution{Univariate,Continuous}, u::Real) = + _trafo_quantile_impl(_dist_params_numtype(d), d, u) + +@inline _trafo_quantile_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, u::Real) = + _trafo_quantile_impl_generic(d, u) + + +@inline _trafo_quantile_impl_generic(d::Distribution{Univariate,Continuous}, u::Real) = + Distributions.quantile(d, u) + +# Workaround for Beta dist, current quantile implementation only supports Float64: +@inline function _trafo_quantile_impl_generic(d::Beta{T}, u::Union{Integer,AbstractFloat}) where {T<:Union{Integer,AbstractFloat}} + Distributions.quantile(d, convert(promote_type(Float64, typeof(u)), u)) +end + +# Workaround for rounding errors that can result in quantile values outside of support of Truncated: +@inline function _trafo_quantile_impl_generic(d::Distributions.Truncated{<:Distribution{Univariate,Continuous}}, u::Real) + x = Distributions.quantile(d, u) + T = typeof(x) + min_x = T(minimum(d)) + max_x = T(maximum(d)) + if x < min_x && isapprox(x, min_x, atol = 4 * eps(T)) + min_x + elseif x > max_x && isapprox(x, max_x, atol = 4 * eps(T)) + max_x + else + x + end +end + + +@inline function _result_numtype(d::Distribution{Univariate}, x::T) where {T<:Real} + float(promote_type(T, _dist_params_numtype(d))) +end + + +@inline function MeasureBase.transport_def(::StdUniform, μ::Distribution{Univariate,Continuous}, x) + R = _result_numtype(μ, x) + if Distributions.insupport(μ, x) + y = _trafo_cdf(μ, x) + convert(R, y) + else + convert(R, NaN) + end +end + + +@inline function MeasureBase.transport_def(ν::Distribution{Univariate,Continuous}, ::StdUniform, x::T) where {T} + R = _result_numtype(ν, x) + TF = float(T) + if 0 <= x <= 1 + # Avoid x ≈ 0 and x ≈ 1 to avoid infinite variate values for target distributions with infinite support: + mod_x = ifelse(x ≈ 0, zero(TF) + eps(TF), ifelse(x ≈ 1, one(TF) - eps(TF), convert(TF, x))) + y = _trafo_quantile(ν, mod_x) + convert(R, y) + else + convert(R, NaN) + end +end + + +# Use standard measures as transformation origin for scaled/translated equivalents: + +function _origin_to_affine(ν::Distribution{Univariate}, y::T) where {T<:Real} + trg_offs, trg_scale = Distributions.location(ν), Distributions.scale(ν) + x = muladd(y, trg_scale, trg_offs) + convert(_result_numtype(ν, y), x) +end + +function _affine_to_origin(μ::Distribution{Univariate}, x::T) where {T<:Real} + src_offs, src_scale = Distributions.location(μ), Distributions.scale(μ) + y = (x - src_offs) / src_scale + convert(_result_numtype(μ, x), y) +end + +for (A, B) in [ + (Uniform, StdUniform), + (Logistic, StdLogistic), + (Normal, StdNormal) +] + @eval begin + @inline MeasureBase.transport_origin(::$A) = $B() + @inline MeasureBase.to_origin(ν::$A, y) = _affine_to_origin(ν, y) + @inline MeasureBase.from_origin(ν::$A, x) = _origin_to_affine(ν, x) + end +end + +@inline MeasureBase.transport_origin(::Exponential) = StdExponential() +@inline MeasureBase.to_origin(ν::Exponential, y) = Distributions.scale(ν) \ y +@inline MeasureBase.from_origin(ν::Exponential, x) = Distributions.scale(ν) * x + + +# Use the underlying distribution as transformation origin for affine +# transformed distributions: + +@inline MeasureBase.transport_origin(d::Distributions.AffineDistribution) = d.ρ +@inline MeasureBase.from_origin(d::Distributions.AffineDistribution, x) = muladd(d.σ, x, d.μ) +@inline MeasureBase.to_origin(d::Distributions.AffineDistribution, y) = d.σ \ (y - d.μ) + + + +# Transform between univariate and single-element power measure + +function MeasureBase.transport_def(ν::Distribution{Univariate}, μ::PowerMeasure{<:StdMeasure}, x) + return transport_def(ν, μ.parent, only(x)) +end + +function MeasureBase.transport_def(ν::PowerMeasure{<:StdMeasure}, μ::Distribution{Univariate}, x) + return Fill(transport_def(ν.parent, μ, only(x)), map(length, ν.axes)...) +end + + +# Transform between univariate and single-element standard multivariate + +function MeasureBase.transport_def(ν::Distribution{Univariate}, μ::StandardDist{D,1}, x) where {D} + return transport_def(ν, StandardDist{D}(), only(x)) +end + +function MeasureBase.transport_def(ν::StandardDist{D,1}, μ::Distribution{Univariate}, x) where {D} + return Fill(transport_def(StandardDist{D}(), μ, only(x)), size(ν)...) +end diff --git a/ext/MeasureBaseDistributionsForwardDiffExt.jl b/ext/MeasureBaseDistributionsForwardDiffExt.jl index 36218eec..245c8e60 100644 --- a/ext/MeasureBaseDistributionsForwardDiffExt.jl +++ b/ext/MeasureBaseDistributionsForwardDiffExt.jl @@ -2,8 +2,35 @@ module MeasureBaseDistributionsForwardDiffExt -using MeasureBase +import MeasureBase import Distributions import ForwardDiff +using Distributions: Distribution, Univariate, Continuous, Beta + +@inline function MeasureBase._trafo_cdf_impl( + ::Type{<:Union{Integer,AbstractFloat}}, + d::Distribution{Univariate,Continuous}, + x::ForwardDiff.Dual{TAG}, +) where {TAG} + x_v = ForwardDiff.value(x) + u = Distributions.cdf(d, x_v) + dudx = Distributions.pdf(d, x_v) + ForwardDiff.Dual{TAG}(u, dudx * ForwardDiff.partials(x)) +end + +@inline function MeasureBase._trafo_quantile_impl( + ::Type{<:Union{Integer,AbstractFloat}}, + d::Distribution{Univariate,Continuous}, + u::ForwardDiff.Dual{TAG}, +) where {TAG} + x = MeasureBase._trafo_quantile_impl_generic(d, ForwardDiff.value(u)) + dxdu = inv(Distributions.pdf(d, x)) + ForwardDiff.Dual{TAG}(x, dxdu * ForwardDiff.partials(u)) +end + +# Workaround for Beta dist, ForwardDiff doesn't work for parameters: +@inline MeasureBase._trafo_quantile_impl_generic(d::Beta{T}, u::Real) where {T<:ForwardDiff.Dual} = + convert(float(typeof(u)), NaN) + end # module MeasureBaseDistributionsForwardDiffExt diff --git a/ext/MeasureBaseDistributionsForwardDiffPullbacksExt.jl b/ext/MeasureBaseDistributionsForwardDiffPullbacksExt.jl new file mode 100644 index 00000000..667af023 --- /dev/null +++ b/ext/MeasureBaseDistributionsForwardDiffPullbacksExt.jl @@ -0,0 +1,25 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +module MeasureBaseDistributionsForwardDiffPullbacksExt + +import MeasureBase +using MeasureBase: StdMeasure, transport_def + +import Distributions +using Distributions: Distribution, Univariate + +import ChainRulesCore +using ForwardDiffPullbacks: fwddiff + +# Use ForwardDiff for univariate transformations: +@inline function ChainRulesCore.rrule(::typeof(transport_def), ν::Distribution{Univariate}, μ::Distribution{Univariate}, x::Any) + ChainRulesCore.rrule(fwddiff(transport_def), ν, μ, x) +end +@inline function ChainRulesCore.rrule(::typeof(transport_def), ν::StdMeasure, μ::Distribution{Univariate}, x::Any) + ChainRulesCore.rrule(fwddiff(transport_def), ν, μ, x) +end +@inline function ChainRulesCore.rrule(::typeof(transport_def), ν::Distribution{Univariate}, μ::StdMeasure, x::Any) + ChainRulesCore.rrule(fwddiff(transport_def), ν, μ, x) +end + +end # module MeasureBaseDistributionsForwardDiffPullbacksExt diff --git a/ext/MeasureBaseDistributionsMooncakeExt.jl b/ext/MeasureBaseDistributionsMooncakeExt.jl new file mode 100644 index 00000000..75bde91b --- /dev/null +++ b/ext/MeasureBaseDistributionsMooncakeExt.jl @@ -0,0 +1,21 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +module MeasureBaseDistributionsMooncakeExt + +using MeasureBase +import Distributions +import Mooncake +using Mooncake: @zero_derivative, MinimalCtx + +using Distributions: Distribution +using MeasureBase: _dist_params_numtype + +# The distribution transports themselves need no rules here: Mooncake +# provides rules for Distributions and StatsFuns/SpecialFunctions, so it +# differentiates the cdf/quantile-based transports natively. The +# ForwardDiffPullbacks-based rules for `transport_def` are a +# Zygote/ChainRules pathway. + +@zero_derivative MinimalCtx Tuple{typeof(_dist_params_numtype),Distribution} + +end # module MeasureBaseDistributionsMooncakeExt diff --git a/ext/MeasureBaseForwardDiffExt.jl b/ext/MeasureBaseForwardDiffExt.jl index 8a1cab44..20113c7f 100644 --- a/ext/MeasureBaseForwardDiffExt.jl +++ b/ext/MeasureBaseForwardDiffExt.jl @@ -3,6 +3,7 @@ module MeasureBaseForwardDiffExt using MeasureBase +using MeasureBase: containsnan, firsttype import ForwardDiff function MeasureBase.containsnan(x::ForwardDiff.Dual) @@ -11,4 +12,7 @@ function MeasureBase.containsnan(x::ForwardDiff.Dual) return a || b end +MeasureBase.firsttype(::Type{T}, ::Type{<:ForwardDiff.Dual{tag,<:Real,N}}) where {T<:Real,tag,N} = + ForwardDiff.Dual{tag,T,N} + end # module MeasureBaseForwardDiffExt diff --git a/ext/MeasureBaseForwardDiffPullbacksExt.jl b/ext/MeasureBaseForwardDiffPullbacksExt.jl new file mode 100644 index 00000000..72ffc751 --- /dev/null +++ b/ext/MeasureBaseForwardDiffPullbacksExt.jl @@ -0,0 +1,10 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +module MeasureBaseForwardDiffPullbacksExt + +import MeasureBase +using ForwardDiffPullbacks: fwddiff + +MeasureBase._fwddiff(f::Function) = fwddiff(f) + +end # module MeasureBaseForwardDiffPullbacksExt diff --git a/src/combinators/superpose.jl b/src/combinators/superpose.jl index 099ee806..aa7b6e20 100644 --- a/src/combinators/superpose.jl +++ b/src/combinators/superpose.jl @@ -65,37 +65,24 @@ function Base.:+(μ::AbstractMeasure, ν::AbstractMeasure) superpose(μ, ν) end -oneplus(x::ULogarithmic) = exp(ULogarithmic, log1pexp(x.log)) - -@inline function density_def(s::SuperpositionMeasure{Tuple{A,B}}, x) where {A,B} - (μ, ν) = s.components - - istrue(insupport(μ, x)) || return exp(ULogarithmic, logdensity_def(ν, x)) - istrue(insupport(ν, x)) || return exp(ULogarithmic, logdensity_def(μ, x)) - - α = basemeasure(μ) - β = basemeasure(ν) - dμ_dα = exp(ULogarithmic, logdensity_def(μ, x)) - dν_dβ = exp(ULogarithmic, logdensity_def(ν, x)) - dα_dβ = exp(ULogarithmic, logdensity_rel(α, β, x)) - dβ_dα = inv(dα_dβ) - return dμ_dα / oneplus(dβ_dα) + dν_dβ / oneplus(dα_dβ) -end +@inline _ulogexp(x) = exp(ULogarithmic, dynamic(x)) function density_def(s::SuperpositionMeasure, x) - T = typeof(s) - msg = """ - Not implemented: There is no method - density_def(::$T, x) - """ - error(msg) + cs = values(s.components) + αs = map(basemeasure, cs) + idxs = eachindex(cs) + sum(idxs) do i + dμᵢ_dαᵢ = _ulogexp(logdensity_def(cs[i], x)) + istrue(insupport(cs[i], x)) || return zero(dμᵢ_dαᵢ) + dΣα_dαᵢ = sum(idxs) do j + dαⱼ_dαᵢ = _ulogexp(logdensity_rel(αs[j], αs[i], x)) + istrue(insupport(cs[j], x)) ? dαⱼ_dαᵢ : zero(dαⱼ_dαᵢ) + end + dμᵢ_dαᵢ / dΣα_dαᵢ + end end -@inline function logdensity_def( - μ::T, - ν::T, - x, -) where {T<:(SuperpositionMeasure{Tuple{A,B}} where {A,B})} +@inline function logdensity_def(μ::T, ν::T, x) where {T<:SuperpositionMeasure} if μ === ν return zero(return_type(logdensity_def, (μ, x))) else @@ -103,44 +90,49 @@ end end end -@inline function logdensity_def( - s::T, - β, - x, -) where {T<:(SuperpositionMeasure{Tuple{A,B}} where {A,B})} - (μ, ν) = s.components - - istrue(insupport(μ, x)) || return logdensity_rel(ν, β, x) - istrue(insupport(ν, x)) || return logdensity_rel(μ, β, x) - return logaddexp(logdensity_rel(μ, β, x), logdensity_rel(ν, β, x)) +function _superpos_logdensity_rel(s::SuperpositionMeasure, β, x) + cs = values(s.components) + ds = map(cs) do μ + istrue(insupport(μ, x)) ? dynamic(logdensity_rel(μ, β, x)) : -Inf + end + logsumexp(ds) end -@inline function logdensity_def( - s::SuperpositionMeasure{Tuple{A,B}}, - β::SuperpositionMeasure, - x, -) where {A,B} - (μ, ν) = s.components - istrue(insupport(μ, x)) || return logdensity_rel(ν, β, x) - istrue(insupport(ν, x)) || return logdensity_rel(μ, β, x) - return logaddexp(logdensity_rel(μ, β, x), logdensity_rel(ν, β, x)) -end +@inline logdensity_def(s::SuperpositionMeasure, β, x) = _superpos_logdensity_rel(s, β, x) -@inline function logdensity_def(s, β::(SuperpositionMeasure{Tuple{A,B}} where {A,B}), x) - -logdensity_def(β, s, x) -end +@inline logdensity_def(s::SuperpositionMeasure, β::SuperpositionMeasure, x) = + _superpos_logdensity_rel(s, β, x) + +@inline logdensity_def(s, β::SuperpositionMeasure, x) = -_superpos_logdensity_rel(β, s, x) @inline logdensity_def(s::SuperpositionMeasure, x) = log(density_def(s, x)) -function basemeasure(μ::SuperpositionMeasure{Tuple{A,B}}) where {A,B} +function basemeasure(μ::SuperpositionMeasure{<:Tuple}) superpose(map(basemeasure, μ.components)...) end + +function basemeasure(μ::SuperpositionMeasure{<:AbstractArray}) + bases = map(basemeasure, μ.components) + allequal(bases) ? weightedmeasure(log(length(bases)), first(bases)) : superpose(bases) +end + basemeasure(μ::SuperpositionMeasure) = superpose(map(basemeasure, μ.components)) -# TODO: Fix `rand` method (this one is wrong) -# function Base.rand(μ::SuperpositionMeasure{X,N}) where {X,N} -# return rand(rand(μ.components)) -# end +function Base.rand(rng::AbstractRNG, ::Type{T}, μ::SuperpositionMeasure) where {T} + components = values(μ.components) + masses = map(massof, components) + total = sum(masses) + total isa AbstractUnknownMass && throw( + ArgumentError("Cannot sample from a superposition of measures of unknown mass"), + ) + threshold = rand(rng) * dynamic(total) + csum = zero(threshold) + for (mass, c) in zip(masses, components) + csum += dynamic(mass) + csum >= threshold && return rand(rng, T, c) + end + return rand(rng, T, last(components)) +end @inline function insupport(d::SuperpositionMeasure, x) any(d.components) do c diff --git a/src/utils.jl b/src/utils.jl index 5d05d8b1..c1e97034 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -185,3 +185,54 @@ isapproxzero(A::AbstractArray) = all(isapproxzero, A) isapproxone(x::T) where {T<:Real} = x ≈ one(T) isapproxone(A::AbstractArray) = all(isapproxone, A) + +containsnan(x::Real) = isnan(x) +containsnan(x) = any(containsnan, x) + + +# ForwardDiffPullbacks dummy `fwddiff`, overloaded by +# ForwardDiffPullbacks extension when loaded: +@inline _fwddiff(f) = f + + +# Autodiff ignore: + +@inline _adignore_call(f) = f() + +macro _adignore(expr) + :(_adignore_call(() -> $(esc(expr)))) +end + + +""" + MeasureBase.convert_realtype(::Type{T}, x) where {T<:Real} + +Convert `x` to use `T` as its underlying type for real numbers. +""" +function convert_realtype end + +@inline convert_realtype(::Type{T}, x::T) where {T<:Real} = x +@inline convert_realtype(::Type{T}, x::AbstractArray{T}) where {T<:Real} = x +@inline convert_realtype(::Type{T}, x::U) where {T<:Real,U<:Real} = T(x) +convert_realtype(::Type{T}, x::AbstractArray{U}) where {T<:Real,U<:Real} = T.(x) +convert_realtype(::Type{T}, x::Union{Tuple,NamedTuple}) where {T<:Real} = + map(Base.Fix1(convert_realtype, T), x) +convert_realtype(::Type{T}, x::AbstractArray) where {T<:Real} = + map(Base.Fix1(convert_realtype, T), x) + +""" + MeasureBase.firsttype(::Type{T}, ::Type{U}) where {T<:Real,U<:Real} + +Return the first type, but as a dual number type if the second one is dual. +""" +function firsttype end + +firsttype(::Type{T}, ::Type{U}) where {T<:Real,U<:Real} = T + + + +# Distributions implementation hooks: +function _trafo_cdf_impl end +function _trafo_quantile_impl end +function _trafo_quantile_impl_generic end +function _dist_params_numtype end diff --git a/test/Project.toml b/test/Project.toml index fec229fd..cb68fc93 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,22 +1,29 @@ [deps] AffineMaps = "2c83c9a8-abf5-4329-a0d7-deffaf474661" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +ArraysOfArrays = "65a8f2f4-9b39-5baf-92e2-a9cc46fdf018" +ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0" DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +ForwardDiffPullbacks = "450a3b6d-2448-4ee1-8e34-e4eb8713b605" InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112" IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" LogarithmicNumbers = "aa2f6b4e-9042-5d33-9679-40d3a6b85899" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" +PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" PropertyFunctions = "09e99361-2bb8-48a2-a80f-de58f0739eb4" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" +StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" diff --git a/test/distributions/getjacobian.jl b/test/distributions/getjacobian.jl new file mode 100644 index 00000000..87de7b86 --- /dev/null +++ b/test/distributions/getjacobian.jl @@ -0,0 +1,34 @@ +# This file is a part of ChangesOfVariables.jl, licensed under the MIT License (MIT). + +import ForwardDiff + +torv_and_back(V::AbstractVector{<:Real}) = V, identity +torv_and_back(x::Real) = [x], V -> V[1] +torv_and_back(x::Complex) = [real(x), imag(x)], V -> Complex(V[1], V[2]) +torv_and_back(x::NTuple{N}) where N = [x...], V -> ntuple(i -> V[i], Val(N)) + +function torv_and_back(x::Ref) + xval = x[] + V, to_xval = torv_and_back(xval) + back_to_ref(V) = Ref(to_xval(V)) + return (V, back_to_ref) +end + +torv_and_back(A::AbstractArray{<:Real}) = vec(A), V -> reshape(V, size(A)) + +function torv_and_back(A::AbstractArray{Complex{T}, N}) where {T<:Real, N} + RA = cat(real.(A), imag.(A), dims = N+1) + V, to_array = torv_and_back(RA) + function back_to_complex(V) + RA = to_array(V) + Complex.(view(RA, map(_ -> :, size(A))..., 1), view(RA, map(_ -> :, size(A))..., 2)) + end + return (V, back_to_complex) +end + + +function getjacobian(f, x) + V, to_x = torv_and_back(x) + vf(V) = torv_and_back(f(to_x(V)))[1] + ForwardDiff.jacobian(vf, V) +end diff --git a/test/distributions/test_autodiff_utils.jl b/test/distributions/test_autodiff_utils.jl new file mode 100644 index 00000000..5197725b --- /dev/null +++ b/test/distributions/test_autodiff_utils.jl @@ -0,0 +1,18 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using LinearAlgebra +using Distributions, ArraysOfArrays +import ForwardDiff, Zygote + + +@testset "trafo_utils" begin + xs = rand(5) + @test Zygote.jacobian(MeasureBase._pushfront, xs, 42)[1] ≈ ForwardDiff.jacobian(xs -> MeasureBase._pushfront(xs, 1), xs) + @test Zygote.jacobian(MeasureBase._pushfront, xs, 42)[2] ≈ vec(ForwardDiff.jacobian(x -> MeasureBase._pushfront(xs, x[1]), [42])) + @test Zygote.jacobian(MeasureBase._pushback, xs, 42)[1] ≈ ForwardDiff.jacobian(xs -> MeasureBase._pushback(xs, 1), xs) + @test Zygote.jacobian(MeasureBase._pushback, xs, 42)[2] ≈ vec(ForwardDiff.jacobian(x -> MeasureBase._pushback(xs, x[1]), [42])) + @test Zygote.jacobian(MeasureBase._rev_cumsum, xs)[1] ≈ ForwardDiff.jacobian(MeasureBase._rev_cumsum, xs) + @test Zygote.jacobian(MeasureBase._exp_cumsum_log, xs)[1] ≈ ForwardDiff.jacobian(MeasureBase._exp_cumsum_log, xs) ≈ ForwardDiff.jacobian(cumprod, xs) +end diff --git a/test/distributions/test_conversions.jl b/test/distributions/test_conversions.jl new file mode 100644 index 00000000..0f9c4d80 --- /dev/null +++ b/test/distributions/test_conversions.jl @@ -0,0 +1,113 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using Random, Statistics, LinearAlgebra +using Distributions +using StableRNGs + +import MeasureBase +using MeasureBase: AbstractMeasure, AsMeasure, asmeasure +using MeasureBase: StdUniform, StdNormal, StdExponential, StdLogistic +using MeasureBase: SuperpositionMeasure, PushforwardMeasure, ProductMeasure +using MeasureBase: logdensityof, massof, insupport + + +@testset "conversions" begin + stblrng() = StableRNG(789990641) + + function test_conversion(d::Distribution, ::Type{M}) where {M} + @testset "conversion $(typeof(d).name) <-> $M" begin + m = asmeasure(d) + @test m isa M + @test typeof(convert(AbstractMeasure, d)) === typeof(m) + @test_throws ArgumentError AsMeasure{typeof(d)}(d) + + d2 = convert(Distribution, m) + @test d2 isa Distribution + @test typeof(Distributions.Distribution(m)) === typeof(d2) + + for x in (rand(stblrng(), d) for _ in 1:10) + @test logdensityof(m, x) ≈ logpdf(d, x) + @test logpdf(d2, x) ≈ logpdf(d, x) + @test insupport(m, x) + end + + x = rand(stblrng(), Float64, m) + # Tuple-marginal product measures have tuple variates: + x isa Tuple ? (@test length(x) == length(d)) : (@test size(x) == size(d)) + @test insupport(m, x) + end + end + + @testset "Dirac" begin + d = Distributions.Dirac(4.2) + m = @inferred asmeasure(d) + @test m === MeasureBase.Dirac(4.2) + @test_throws ArgumentError AsMeasure{typeof(d)}(d) + @test @inferred(Distributions.Distribution(m)) === d + end + + @testset "products" begin + test_conversion(product_distribution(Weibull.([0.7, 1.1, 1.3])), ProductMeasure) + test_conversion(product_distribution(Poisson.([0.7, 1.4])), ProductMeasure) + + if isdefined(Distributions, :ProductDistribution) + test_conversion(product_distribution(Weibull(0.7), Exponential(1.3)), ProductMeasure) + end + end + + @testset "reshaped" begin + test_conversion(reshape(MvNormal([0.7, 0.9], [1.4 0.5; 0.5 1.1]), 1, 2), PushforwardMeasure) + test_conversion(reshape(product_distribution(Weibull.([0.7, 1.1, 1.3, 0.9, 1.2, 0.8])), 2, 3), PushforwardMeasure) + end + + @testset "mixtures" begin + test_conversion(MixtureModel([Normal(-1.0, 1.0), Normal(2.0, 3.0)], [0.3, 0.7]), SuperpositionMeasure) + test_conversion(MixtureModel([Normal(-2.0, 1.0), Normal(0.0, 2.0), Normal(3.0, 1.0)], [0.2, 0.5, 0.3]), SuperpositionMeasure) + test_conversion(MixtureModel([Exponential(0.3), Weibull(2.0, 1.0)], [0.4, 0.6]), SuperpositionMeasure) + test_conversion(MixtureModel([MvNormal([0.0, 0.0], I(2)), MvNormal([2.0, 2.0], 2 * I(2))], [0.3, 0.7]), SuperpositionMeasure) + test_conversion(UnivariateGMM([-1.0, 2.0], [1.0, 0.5], Categorical([0.4, 0.6])), SuperpositionMeasure) + + d = MixtureModel([Normal(-1.0, 1.0), Normal(2.0, 3.0)], [0.3, 0.7]) + m = asmeasure(d) + @test probs(convert(Distribution, m)) ≈ probs(d) + @test massof(m) ≈ 1 + @test massof(asmeasure(Normal())) == 1 + @test mean(rand(stblrng(), Float64, m^1000)) ≈ mean(d) atol = 0.3 + + # Hand-built superpositions of weighted probability measures behave + # like mixtures: + m2 = 0.3 * asmeasure(Normal(-1.0, 1.0)) + 0.7 * asmeasure(Normal(2.0, 3.0)) + for x in (rand(stblrng(), d) for _ in 1:10) + @test logdensityof(m2, x) ≈ logpdf(d, x) + end + end + + @testset "standard distributions" begin + @test StandardUniform === StandardDist{Uniform} + @test StandardNormal === StandardDist{Normal} + @test StandardUniform{0} === StandardDist{Uniform,0} + @test StandardNormal{1} === StandardDist{Normal,1} + + for (D, B) in [ + (Uniform, StdUniform()), + (Exponential, StdExponential()), + (Logistic, StdLogistic()), + (Normal, StdNormal()), + ] + @test @inferred(asmeasure(StandardDist{D}())) === B + @test @inferred(asmeasure(StandardDist{D}(3))) == B^3 + @test @inferred(asmeasure(StandardDist{D}(2, 3))) == B^(2, 3) + + @test @inferred(Distributions.Distribution(B)) === StandardDist{D}() + @test @inferred(convert(Distribution, B)) === StandardDist{D}() + @test @inferred(Distributions.Distribution(B^3)) == StandardDist{D}(3) + @test @inferred(convert(Distribution, B^(2, 3))) == StandardDist{D}(2, 3) + + d = StandardDist{D}(3) + x = rand(stblrng(), d) + @test logdensityof(asmeasure(d), x) ≈ logpdf(d, x) + end + end +end diff --git a/test/distributions/test_distribution_measure.jl b/test/distributions/test_distribution_measure.jl new file mode 100644 index 00000000..7715e58a --- /dev/null +++ b/test/distributions/test_distribution_measure.jl @@ -0,0 +1,53 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +import Distributions +using Distributions: Distribution +import MeasureBase +using MeasureBase: AbstractMeasure + +@testset "Measure interface" begin + d = Distributions.Weibull() + @test @inferred(AbstractMeasure(d)) isa AbstractMeasure + @test @inferred(AbstractMeasure(d)) isa DistributionMeasure + @test @inferred(convert(AbstractMeasure, d)) isa AbstractMeasure + @test @inferred(convert(AbstractMeasure, d)) isa DistributionMeasure + @test @inferred(Distribution(AbstractMeasure(d))) === d + @test @inferred(convert(Distribution, convert(AbstractMeasure, d))) === d + + + c0 = AbstractMeasure(Distributions.Weibull(0.7, 1.3)) + c1 = AbstractMeasure(Distributions.MvNormal([0.7, 0.9], [1.4 0.5; 0.5 1.1])) + + d0 = AbstractMeasure(Distributions.Poisson(0.7)) + d1 = AbstractMeasure(Distributions.product_distribution(Distributions.Poisson.([0.7, 1.4]))) + + for μ in [c0, c1, d0, d1] + d = Distribution(μ) + x = rand(μ) + @test @inferred(MeasureBase.logdensity_def(μ, x)) == Distributions.logpdf(d, x) + @test @inferred(MeasureBase.unsafe_logdensityof(μ, x)) == Distributions.logpdf(d, x) + + MeasureBase.Interface.test_interface(d) + end + + @test @inferred(MeasureBase.basemeasure(c0)) == MeasureBase.Lebesgue(MeasureBase.ℝ) + @test @inferred(MeasureBase.basemeasure(c1)) == MeasureBase.Lebesgue(MeasureBase.ℝ) ^ 2 + + @test @inferred(MeasureBase.insupport(c0, 3)) == true + @test @inferred(MeasureBase.insupport(c0, -3)) == false + @test @inferred(MeasureBase.insupport(c1, [0.1, 0.2])) == true + @test @inferred(MeasureBase.insupport(d0, 3)) == true + @test @inferred(MeasureBase.insupport(d0, 3.2)) == false + @test @inferred(MeasureBase.insupport(d1, [1, 2])) == true + @test @inferred(MeasureBase.insupport(d1, [1.1, 2.2])) == false + + @test MeasureBase.paramnames(c0) == (:α, :θ) + if VERSION >= v"1.8" + @test @inferred(MeasureBase.params(c0)) == (α = 0.7, θ = 1.3) + else + # v1.6 can't type-infer this: + @test (MeasureBase.params(c0)) == (α = 0.7, θ = 1.3) + end +end diff --git a/test/distributions/test_distributions.jl b/test/distributions/test_distributions.jl new file mode 100644 index 00000000..65fab330 --- /dev/null +++ b/test/distributions/test_distributions.jl @@ -0,0 +1,24 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test +using MeasureBase +using Distributions +import ForwardDiff, ForwardDiffPullbacks, ChainRulesCore + +const MeasureBaseDistributionsExt = Base.get_extension(MeasureBase, :MeasureBaseDistributionsExt) +@test MeasureBaseDistributionsExt isa Module + +using .MeasureBaseDistributionsExt: + StandardDist, StandardUniform, StandardNormal, DistributionMeasure, nonstddist + +@testset "Distributions extension" begin + include("test_autodiff_utils.jl") + include("test_measure_interface.jl") + include("test_distribution_measure.jl") + include("test_standard_dist.jl") + include("test_standard_uniform.jl") + include("test_standard_normal.jl") + include("test_conversions.jl") + include("test_transport.jl") + include("test_mooncake.jl") +end diff --git a/test/distributions/test_measure_interface.jl b/test/distributions/test_measure_interface.jl new file mode 100644 index 00000000..d11d2889 --- /dev/null +++ b/test/distributions/test_measure_interface.jl @@ -0,0 +1,43 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +import Distributions +import MeasureBase + +@testset "Measure interface" begin + c0 = Distributions.Weibull(0.7, 1.3) + c1 = Distributions.MvNormal([0.7, 0.9], [1.4 0.5; 0.5 1.1]) + + d0 = Distributions.Poisson(0.7) + d1 = Distributions.product_distribution(Distributions.Poisson.([0.7, 1.4])) + + for d in [c0, c1, d0, d1] + x = rand(d) + @test @inferred(MeasureBase.logdensity_def(d, x)) == Distributions.logpdf(d, x) + @test @inferred(MeasureBase.unsafe_logdensityof(d, x)) == Distributions.logpdf(d, x) + + MeasureBase.Interface.test_interface(d) + end + + @test @inferred(MeasureBase.basemeasure(c0)) == MeasureBase.Lebesgue(MeasureBase.ℝ) + @test @inferred(MeasureBase.basemeasure(c1)) == MeasureBase.Lebesgue(MeasureBase.ℝ) ^ 2 + + @test @inferred(MeasureBase.insupport(c0, 3)) == true + @test @inferred(MeasureBase.insupport(c0, -3)) == false + @test @inferred(MeasureBase.insupport(c1, [0.1, 0.2])) == true + @test @inferred(MeasureBase.insupport(d0, 3)) == true + @test @inferred(MeasureBase.insupport(d0, 3.2)) == false + @test @inferred(MeasureBase.insupport(d1, [1, 2])) == true + @test @inferred(MeasureBase.insupport(d1, [1.1, 2.2])) == false + + @test MeasureBase.paramnames(c0) == (:α, :θ) + if VERSION >= v"1.8" + @test @inferred(MeasureBase.params(c0)) == (α = 0.7, θ = 1.3) + else + # v1.6 can't type-infer this: + @test (MeasureBase.params(c0)) == (α = 0.7, θ = 1.3) + end + + @test MeasureBase.∫(x -> Distributions.Normal(x, 0), Distributions.Normal()) isa MeasureBase.DensityMeasure +end diff --git a/test/distributions/test_mooncake.jl b/test/distributions/test_mooncake.jl new file mode 100644 index 00000000..4a20e46c --- /dev/null +++ b/test/distributions/test_mooncake.jl @@ -0,0 +1,68 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using Random, LinearAlgebra +using Distributions +import Mooncake +import ForwardDiff + +using MeasureBase +using MeasureBase: transport_to, transport_def, asmeasure +using MeasureBase: StdUniform, StdNormal + +_mooncake_gradient(f, x) = Mooncake.value_and_gradient!!( + Mooncake.prepare_gradient_cache(f, x), f, x +)[2][2] + +_test_gradient(f, x::Real) = @test _mooncake_gradient(f, x) ≈ ForwardDiff.derivative(f, x) +_test_gradient(f, x::AbstractVector) = @test _mooncake_gradient(f, x) ≈ ForwardDiff.gradient(f, x) + +@testset "Mooncake AD with Distributions" begin + @test Base.get_extension(MeasureBase, :MeasureBaseDistributionsMooncakeExt) isa Module + + @testset "zero-derivative primitives" begin + rng = Random.Xoshiro(789990641) + Mooncake.TestUtils.test_rule( + rng, MeasureBase._dist_params_numtype, Normal(0.2, 1.3); + is_primitive = true, + ) + end + + @testset "univariate transport gradients" begin + _test_gradient(x -> transport_def(StdUniform(), Normal(1.0, 2.0), x), 0.5) + _test_gradient(u -> transport_def(Normal(1.0, 2.0), StdUniform(), u), 0.3) + _test_gradient(u -> transport_def(Beta(2.0, 3.0), StdUniform(), u), 0.3) + _test_gradient(x -> transport_def(StdUniform(), Gamma(2.0, 1.0), x), 0.7) + _test_gradient(x -> transport_def(StdUniform(), truncated(Normal(0.3, 1.2), -0.5, 1.5), x), 0.4) + _test_gradient(x -> transport_def(StdUniform(), 2.0 * Weibull(0.7) + 1.0, x), 3.0) + _test_gradient(x -> transport_def(StdNormal(), StandardDist{Uniform}(), x), 0.4) + end + + @testset "multivariate transport gradients" begin + mvn = MvNormal([0.3, -2.9], [1.7 0.5; 0.5 2.3]) + _test_gradient(x -> sum(transport_to(StdNormal()^2, mvn)(x)), [0.1, -2.0]) + _test_gradient(y -> sum(transport_to(mvn, StdNormal()^2)(y)), [0.2, 0.7]) + + pd = product_distribution([Weibull(0.7), Exponential(1.3), Normal(0.5, 2.0)]) + _test_gradient(x -> sum(transport_to(StdNormal()^3, asmeasure(pd))(x)), [0.4, 0.8, 1.5]) + + dirich = Dirichlet([2.0, 3.0, 4.0]) + _test_gradient(u -> MeasureBase.from_origin(dirich, u)[1], [0.3, 0.7]) + _test_gradient(x -> sum(MeasureBase.to_origin(dirich, vcat(x, 1 - sum(x)))), [0.28, 0.23]) + end + + @testset "logdensityof gradients" begin + for d in [ + Weibull(0.7, 1.3), + MixtureModel([Normal(-1.0, 1.0), Normal(2.0, 3.0)], [0.3, 0.7]), + MixtureModel([Normal(-2.0, 1.0), Normal(0.0, 2.0), Normal(3.0, 1.0)], [0.2, 0.5, 0.3]), + ] + m = asmeasure(d) + _test_gradient(x -> logdensityof(m, x[1]), [0.5]) + end + + mvn = MvNormal([0.3, -2.9], [1.7 0.5; 0.5 2.3]) + _test_gradient(x -> logdensityof(asmeasure(mvn), x), [0.1, -2.0]) + end +end diff --git a/test/distributions/test_standard_dist.jl b/test/distributions/test_standard_dist.jl new file mode 100644 index 00000000..64b9f655 --- /dev/null +++ b/test/distributions/test_standard_dist.jl @@ -0,0 +1,128 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using Random, Statistics, LinearAlgebra +using Distributions, PDMats +using StableRNGs +import ForwardDiff, ChainRulesTestUtils + + +@testset "standard_dist" begin + stblrng() = StableRNG(789990641) + + for (D, sz, dref) in [ + (Uniform, (), Uniform()), + (Uniform, (5,), product_distribution(fill(Uniform(0.0, 1.0), 5))), + (Uniform, (2, 3), reshape(product_distribution(fill(Uniform(0.0, 1.0), 6)), 2, 3)), + (Normal, (), Normal()), + (Normal, (), Normal(0., 1.0)), + (Normal, (5,), MvNormal(Diagonal(fill(1.0, 5)))), + (Normal, (2, 3), reshape(MvNormal(Diagonal(fill(1.0, 6))), 2, 3)), + (Exponential, (), Exponential()), + (Exponential, (5,), product_distribution(fill(Exponential(1.0), 5))), + (Exponential, (2, 3), reshape(product_distribution(fill(Exponential(1.0), 6)), 2, 3)), + ] + @testset "StandardDist{$D}($(join(sz,",")))" begin + N = length(sz) + + @test @inferred(StandardDist{D}(sz...)) isa StandardDist{D} + @test @inferred(StandardDist{D}(sz...)) isa StandardDist{D} + @test @inferred(size(StandardDist{D}(sz...))) == size(dref) + @test @inferred(size(StandardDist{D}(sz...))) == size(dref) + + d = StandardDist{D}(sz...) + + if size(d) == () + @test @inferred(MeasureBaseDistributionsExt.nonstddist(d)) == dref + end + + @test @inferred(length(d)) == length(dref) + @test @inferred(size(d)) == size(dref) + + @test @inferred(eltype(typeof(d))) == eltype(typeof(dref)) + @test @inferred(eltype(d)) == eltype(dref) + + @test @inferred(Distributions.params(d)) == () + @test @inferred(partype(d)) == partype(dref) + + for f in [minimum, maximum, mean, median, mode, modes, var, std, skewness, kurtosis, location, scale, entropy] + supported_by_dref = try f(dref); true catch MethodError; false; end + if supported_by_dref + @test @inferred(f(d)) ≈ f(dref) + end + end + + for x in [rand(dref) for i in 1:10] + ref_gradlogpdf = try + gradlogpdf(dref, x) + catch MethodError + ForwardDiff.gradient(x -> logpdf(dref, x), x) + end + @test @inferred(gradlogpdf(d, x)) ≈ ref_gradlogpdf + @test @inferred(logpdf(d, x)) ≈ logpdf(dref, x) + @test @inferred(pdf(d, x)) ≈ pdf(dref, x) + end + + if size(d) == () + for x in [minimum(dref), quantile(dref, 1//3), quantile(dref, 1//2), quantile(dref, 2//3), maximum(dref)] + for f in [logpdf, pdf, gradlogpdf, logcdf, cdf, logccdf, ccdf] + @test @inferred(f(d, x)) ≈ f(dref, x) + end + end + + for x in [0, 1//3, 1//2, 2//3, 1] + for f in [quantile, cquantile] + @test @inferred(f(d, x)) ≈ f(dref, x) + end + end + + for x in log.([0, 1//3, 1//2, 2//3, 1]) + for f in [invlogcdf, invlogccdf] + @test @inferred(f(d, x)) ≈ f(dref, x) + end + end + + for p in [0.0, 0.25, 0.75, 1.0] + @test @inferred(quantile(d, p)) == quantile(dref, p) + @test @inferred(cquantile(d, p)) == cquantile(dref, p) + end + + for t in [-3, 0, 3] + @test isapprox(@inferred(mgf(d, t)), mgf(dref, t), rtol = 1e-5) + @test isapprox(@inferred(cf(d, t)), cf(dref, t), rtol = 1e-5) + end + + @test @inferred(truncated(d, quantile(dref, 1//3), quantile(dref, 2//3))) == truncated(dref, quantile(dref, 1//3), quantile(dref, 2//3)) + + @test @inferred(product_distribution(fill(d, 3))) == StandardDist{typeof(d)}(3) + @test @inferred(product_distribution(fill(d, 3, 4))) == StandardDist{typeof(d)}(3, 4) + end + + if length(size(d)) == 1 + @test @inferred(convert(Distributions.Product, d)) isa Distributions.Product + d_as_prod = convert(Distributions.Product, d) + @test d_as_prod.v == fill(StandardDist{D}(), size(d)...) + end + + @test @inferred(rand(stblrng(), d)) == rand(stblrng(), d) + @test @inferred(rand(stblrng(), d, 5)) == rand(stblrng(), d, 5) + + @test @inferred(rand(stblrng(), d)) == rand(stblrng(), dref) + @test @inferred(rand(stblrng(), d, 5)) == rand(stblrng(), dref, 5) + @test @inferred(rand!(stblrng(), d, zeros(size(d)...))) == rand!(stblrng(), dref, zeros(size(dref)...)) + if length(size(d)) == 1 + @test @inferred(rand!(stblrng(), d, zeros(size(d)..., 5))) == rand!(stblrng(), dref, zeros(size(dref)..., 5)) + end + end + end + + @testset "StandardDist{Normal}()" begin + # TODO: Add @inferred + d = StandardDist{Normal}(4) + d_uv = StandardDist{Normal}() + dref = MvNormal(Diagonal(fill(1.0, 4))) + @test (MvNormal(d)) == dref + @test (Base.convert(MvNormal, d)) == dref + end +end diff --git a/test/distributions/test_standard_normal.jl b/test/distributions/test_standard_normal.jl new file mode 100644 index 00000000..3d77f583 --- /dev/null +++ b/test/distributions/test_standard_normal.jl @@ -0,0 +1,129 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using Random, Statistics, LinearAlgebra +using Distributions, PDMats +using StableRNGs + + +@testset "StandardDist{Normal}" begin + stblrng() = StableRNG(789990641) + + @testset "StandardDist{Normal,0}" begin + @test @inferred(Normal(StandardDist{Normal}())) isa Normal{Float64} + @test @inferred(Normal(StandardDist{Normal}())) == Normal() + @test @inferred(convert(Normal, StandardDist{Normal}())) == Normal() + + d = StandardDist{Normal}() + dref = Normal() + + @test @inferred(minimum(d)) == minimum(dref) + @test @inferred(maximum(d)) == maximum(dref) + + @test @inferred(Distributions.params(d)) == () + @test @inferred(partype(d)) == partype(dref) + + @test @inferred(location(d)) == location(dref) + @test @inferred(scale(d)) == scale(dref) + + @test @inferred(eltype(typeof(d))) == eltype(typeof(dref)) + @test @inferred(eltype(d)) == eltype(dref) + + @test @inferred(length(d)) == length(dref) + @test @inferred(size(d)) == size(dref) + + @test @inferred(mean(d)) == mean(dref) + @test @inferred(median(d)) == median(dref) + @test @inferred(mode(d)) == mode(dref) + @test @inferred(modes(d)) ≈ modes(dref) + + @test @inferred(var(d)) == var(dref) + @test @inferred(std(d)) == std(dref) + @test @inferred(skewness(d)) == skewness(dref) + @test @inferred(kurtosis(d)) == kurtosis(dref) + + @test @inferred(entropy(d)) == entropy(dref) + + for x in [-Inf, -1.3, 0.0, 1.3, +Inf] + @test @inferred(gradlogpdf(d, x)) == gradlogpdf(dref, x) + + @test @inferred(logpdf(d, x)) == logpdf(dref, x) + @test @inferred(pdf(d, x)) == pdf(dref, x) + @test @inferred(logcdf(d, x)) == logcdf(dref, x) + @test @inferred(cdf(d, x)) == cdf(dref, x) + @test @inferred(logccdf(d, x)) == logccdf(dref, x) + @test @inferred(ccdf(d, x)) == ccdf(dref, x) + end + + for p in [0.0, 0.25, 0.75, 1.0] + @test @inferred(quantile(d, p)) == quantile(dref, p) + @test @inferred(cquantile(d, p)) == cquantile(dref, p) + end + + for t in [-3, 0, 3] + @test @inferred(mgf(d, t)) == mgf(dref, t) + @test @inferred(cf(d, t)) == cf(dref, t) + end + + @test @inferred(rand(stblrng(), d)) == rand(stblrng(), dref) + @test @inferred(rand!(stblrng(), d, fill(0.0))) == rand!(stblrng(), dref, fill(0.0)) + @test @inferred(rand(stblrng(), d, 5)) == rand(stblrng(), dref, 5) + + @test @inferred(truncated(StandardDist{Normal}(), -2.2f0, 3.1f0)) isa Truncated{Normal{Float64}} + @test truncated(StandardDist{Normal}(), -2.2f0, 3.1f0) == truncated(Normal(0.0, 1.0), -2.2f0, 3.1f0) + + @test @inferred(product_distribution(fill(StandardDist{Normal}(), 3))) isa StandardDist{Normal,1} + @test product_distribution(fill(StandardDist{Normal}(), 3)) == StandardDist{Normal}(3) + end + + + @testset "StandardDist{Normal,1}" begin + @test @inferred(StandardDist{Normal}(3)) isa StandardDist{Normal,1} + @test @inferred(StandardDist{Normal}(3)) isa StandardDist{Normal,1} + @test @inferred(StandardDist{Normal}(3)) isa StandardDist{Normal,1} + + @test @inferred(MvNormal(StandardDist{Normal}(3))) isa MvNormal{Int} + @test @inferred(MvNormal(StandardDist{Normal}(3))) == MvNormal(ScalMat(3, 1.0)) + @test @inferred(convert(MvNormal, StandardDist{Normal}(3))) == MvNormal(ScalMat(3, 1.0)) + + d = StandardDist{Normal}(3) + dref = MvNormal(ScalMat(3, 1.0)) + + @test @inferred(eltype(typeof(d))) == eltype(typeof(dref)) + @test @inferred(eltype(d)) == eltype(dref) + + @test @inferred(length(d)) == length(dref) + @test @inferred(size(d)) == size(dref) + + @test @inferred(Distributions.params(d)) == () + @test @inferred(partype(d)) == partype(dref) + + @test @inferred(mean(d)) == mean(dref) + @test @inferred(var(d)) == var(dref) + @test @inferred(cov(d)) == cov(dref) + + @test @inferred(mode(d)) == mode(dref) + @test @inferred(modes(d)) == modes(dref) + + @test @inferred(invcov(d)) == invcov(dref) + @test @inferred(logdetcov(d)) == logdetcov(dref) + + @test @inferred(entropy(d)) == entropy(dref) + + for x in fill.([-Inf, -1.3, 0.0, 1.3, +Inf], 3) + # Distributions.insupport is inconsistent at +- Inf between Normal and MvNormal + if !any(isinf, x) + @test @inferred(Distributions.insupport(d, x)) == Distributions.insupport(dref, x) + end + @test @inferred(logpdf(d, x)) == logpdf(dref, x) + @test @inferred(pdf(d, x)) == pdf(dref, x) + @test @inferred(sqmahal(d, x)) == sqmahal(dref, x) + @test @inferred(gradlogpdf(d, x)) == gradlogpdf(dref, x) + end + + @test @inferred(rand(stblrng(), d)) == rand(stblrng(), d) + @test @inferred(rand!(stblrng(), d, zeros(3))) == rand!(stblrng(), d, zeros(3)) + @test @inferred(rand!(stblrng(), d, zeros(3, 10))) == rand!(stblrng(), d, zeros(3, 10)) + end +end diff --git a/test/distributions/test_standard_uniform.jl b/test/distributions/test_standard_uniform.jl new file mode 100644 index 00000000..bcb0fb3e --- /dev/null +++ b/test/distributions/test_standard_uniform.jl @@ -0,0 +1,118 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using Random, Statistics, LinearAlgebra +using Distributions, PDMats +using StableRNGs +using FillArrays +using ForwardDiff + + +@testset "StandardDist{Uniform}" begin + stblrng() = StableRNG(789990641) + + @testset "StandardDist{Uniform,0}" begin + @test @inferred(Uniform(StandardDist{Uniform}())) isa Uniform{Float64} + @test @inferred(Uniform(StandardDist{Uniform}())) == Uniform() + @test @inferred(convert(Uniform, StandardDist{Uniform}())) == Uniform() + + d = StandardDist{Uniform}() + dref = Uniform() + + @test @inferred(minimum(d)) == minimum(dref) + @test @inferred(maximum(d)) == maximum(dref) + + @test @inferred(Distributions.params(d)) == () + @test @inferred(partype(d)) == partype(dref) + + @test @inferred(location(d)) == location(dref) + @test @inferred(scale(d)) == scale(dref) + + @test @inferred(eltype(typeof(d))) == eltype(typeof(dref)) + @test @inferred(eltype(d)) == eltype(dref) + + @test @inferred(length(d)) == length(dref) + @test @inferred(size(d)) == size(dref) + + @test @inferred(mean(d)) == mean(dref) + @test @inferred(median(d)) == median(dref) + @test @inferred(mode(d)) == mode(dref) + @test @inferred(modes(d)) ≈ modes(dref) + + @test @inferred(var(d)) ≈ var(dref) + @test @inferred(std(d)) ≈ std(dref) + @test @inferred(skewness(d)) == skewness(dref) + @test @inferred(kurtosis(d)) ≈ kurtosis(dref) + + @test @inferred(entropy(d)) == entropy(dref) + + for x in [-0.5, 0.0, 0.25, 0.75, 1.0, 1.5] + @test @inferred(logpdf(d, x)) == logpdf(dref, x) + @test @inferred(pdf(d, x)) == pdf(dref, x) + @test @inferred(logcdf(d, x)) == logcdf(dref, x) + @test @inferred(cdf(d, x)) == cdf(dref, x) + @test @inferred(logccdf(d, x)) == logccdf(dref, x) + @test @inferred(ccdf(d, x)) == ccdf(dref, x) + end + + for p in [0.0, 0.25, 0.75, 1.0] + @test @inferred(quantile(d, p)) == quantile(dref, p) + @test @inferred(cquantile(d, p)) == cquantile(dref, p) + end + + for t in [-3, 0, 3] + @test @inferred(mgf(d, t)) == mgf(dref, t) + @test @inferred(cf(d, t)) == cf(dref, t) + end + + @test @inferred(rand(stblrng(), d)) == rand(stblrng(), dref) + @test @inferred(rand!(stblrng(), d, fill(0.0))) == rand!(stblrng(), dref, fill(0.0)) + @test @inferred(rand(stblrng(), d, 5)) == rand(stblrng(), dref, 5) + + @test @inferred(truncated(StandardDist{Uniform}(), -0.5f0, 0.7f0)) isa Uniform{Float64} + @test truncated(StandardDist{Uniform}(), -0.5f0, 0.7f0) == Uniform(0.0f0, 0.7f0) + @test truncated(StandardDist{Uniform}(), 0.2f0, 0.7f0) == Uniform(0.2f0, 0.7f0) + + @test @inferred(product_distribution(fill(StandardDist{Uniform}(), 3))) isa MeasureBaseDistributionsExt.StandardDist{Uniform,1} + @test product_distribution(fill(StandardDist{Uniform}(), 3)) == MeasureBaseDistributionsExt.StandardDist{Uniform}(3) + end + + + @testset "StandardDist{Uniform,1}" begin + d = MeasureBaseDistributionsExt.StandardDist{Uniform}(3) + dref = product_distribution(fill(Uniform(), 3)) + + @test @inferred(eltype(typeof(d))) == eltype(typeof(dref)) + @test @inferred(eltype(d)) == eltype(dref) + + @test @inferred(length(d)) == length(dref) + @test @inferred(size(d)) == size(dref) + + @test @inferred(Distributions.params(d)) == () + @test @inferred(partype(d)) == partype(dref) + + @test @inferred(mean(d)) == mean(dref) + @test @inferred(var(d)) ≈ var(dref) + @test @inferred(cov(d)) ≈ cov(dref) + + @test @inferred(mode(d)) == [0.5, 0.5, 0.5] + @test @inferred(modes(d)) == fill([0, 0,0 ]) + + @test @inferred(invcov(d)) == inv(cov(dref)) + @test @inferred(logdetcov(d)) == logdet(cov(dref)) + + @test @inferred(entropy(d)) == entropy(dref) + + for x in fill.([-Inf, -1.3, 0.0, 1.3, +Inf], 3) + @test @inferred(Distributions.insupport(d, x)) == Distributions.insupport(dref, x) + @test @inferred(logpdf(d, x)) == logpdf(dref, x) + @test @inferred(pdf(d, x)) == pdf(dref, x) + @test @inferred(gradlogpdf(d, x)) == ForwardDiff.gradient(x -> logpdf(d, x), x) + end + + @test @inferred(rand(stblrng(), d)) == rand(stblrng(), d) + @test @inferred(rand!(stblrng(), d, zeros(3))) == rand!(stblrng(), d, zeros(3)) + @test @inferred(rand!(stblrng(), d, zeros(3, 10))) == rand!(stblrng(), d, zeros(3, 10)) + end +end diff --git a/test/distributions/test_transport.jl b/test/distributions/test_transport.jl new file mode 100644 index 00000000..4ae50831 --- /dev/null +++ b/test/distributions/test_transport.jl @@ -0,0 +1,194 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using LinearAlgebra +using InverseFunctions, ChangesOfVariables +using Distributions, ArraysOfArrays +using StableRNGs +import ForwardDiff, Zygote +import PDMats + +using MeasureBase: transport_to, transport_def, transport_origin +using MeasureBase: StdUniform, StdNormal, StdExponential +using .MeasureBaseDistributionsExt: _trafo_cdf, _trafo_quantile + +include("getjacobian.jl") + + +@testset "test_distribution_transform" begin + function test_back_and_forth(trg, src) + @testset "transform $(typeof(trg).name) <-> $(typeof(src).name)" begin + x = rand(src) + y = transport_def(trg, src, x) + src_v_reco = transport_def(src, trg, y) + + @test x ≈ src_v_reco + + f = x -> transport_def(trg, src, x) + ref_ladj = logpdf(src, x) - logpdf(trg, y) + @test ref_ladj ≈ logabsdet(getjacobian(f, x))[1] + end + end + + reshaped_rand(d::Distribution{Univariate}, n) = rand(d, n) + reshaped_rand(d::Distribution{Multivariate}, n) = nestedview(rand(d, n)) + + function test_dist_trafo_moments(trg, src) + unshaped(x) = first(torv_and_back(x)) + @testset "check moments of trafo $(typeof(trg).name) <- $(typeof(src).name)" begin + X = reshaped_rand(src, 10^5) + Y = transport_to(trg, src).(X) + Y_ref = reshaped_rand(trg, 10^6) + @test isapprox(mean(unshaped.(Y)), mean(unshaped.(Y_ref)), rtol = 0.5) + @test isapprox(cov(unshaped.(Y)), cov(unshaped.(Y_ref)), rtol = 0.5) + end + end + + @testset "transforms-tests" begin + stduvuni = StandardDist{Uniform}() + stduvnorm = StandardDist{Uniform}() + + uniform1 = Uniform(-5.0, -0.01) + uniform2 = Uniform(0.01, 5.0) + + normal1 = Normal(-10, 1) + normal2 = Normal(10, 5) + + stdmvnorm1 = StandardDist{Normal}(1) + stdmvnorm2 = StandardDist{Normal}(2) + + stdmvuni2 = StandardDist{Uniform}(2) + + standnorm2_reshaped = reshape(stdmvnorm2, 1, 2) + + mvnorm = MvNormal([0.3, -2.9], [1.7 0.5; 0.5 2.3]) + beta = Beta(3,1) + gamma = Gamma(0.1,0.7) + dirich = Dirichlet([0.1,4]) + + test_back_and_forth(stduvuni, stduvuni) + test_back_and_forth(stduvnorm, stduvnorm) + test_back_and_forth(stduvuni, stduvnorm) + test_back_and_forth(stduvnorm, stduvuni) + + test_back_and_forth(stdmvuni2, stdmvuni2) + test_back_and_forth(stdmvnorm2, stdmvnorm2) + test_back_and_forth(stdmvuni2, stdmvnorm2) + test_back_and_forth(stdmvnorm2, stdmvuni2) + + test_back_and_forth(beta, stduvnorm) + test_back_and_forth(gamma, stduvnorm) + test_back_and_forth(gamma, beta) + + test_back_and_forth(mvnorm, stdmvuni2) + test_back_and_forth(stdmvuni2, mvnorm) + + test_back_and_forth(mvnorm, standnorm2_reshaped) + test_back_and_forth(standnorm2_reshaped, mvnorm) + test_back_and_forth(stdmvnorm2, standnorm2_reshaped) + test_back_and_forth(standnorm2_reshaped, standnorm2_reshaped) + + test_dist_trafo_moments(normal2, normal1) + test_dist_trafo_moments(uniform2, uniform1) + + test_dist_trafo_moments(beta, stduvnorm) + test_dist_trafo_moments(gamma, stduvnorm) + + test_dist_trafo_moments(mvnorm, stdmvnorm2) + test_dist_trafo_moments(dirich, stdmvnorm1) + + let + mvuni = product_distribution([Uniform(), Uniform()]) + + x = rand() + @test_throws ArgumentError transport_to(stduvnorm, mvnorm)(x) + @test_throws ArgumentError transport_to(stduvnorm, stdmvnorm1)(x) + @test_throws ArgumentError transport_to(stduvnorm, stdmvnorm2)(x) + + x = rand(2) + @test_throws ArgumentError transport_to(stduvnorm, mvnorm)(x) + @test_throws ArgumentError transport_to(stduvnorm, stdmvnorm1)(x) + @test_throws ArgumentError transport_to(stduvnorm, stdmvnorm2)(x) + end + end + + @testset "Custom cdf and quantile for dual numbers" begin + Dual = ForwardDiff.Dual + + @test isapprox(_trafo_cdf(Normal(Dual(0, 1, 0, 0), Dual(1, 0, 1, 0)), Dual(0.5, 0, 0, 1)), cdf(Normal(Dual(0, 1, 0, 0), Dual(1, 0, 1, 0)), Dual(0.5, 0, 0, 1)), rtol = 10^-6) + @test isapprox(_trafo_cdf(Normal(0, 1), Dual(0.5, 1)), cdf(Normal(0, 1), Dual(0.5, 1)), rtol = 10^-6) + + @test isapprox(_trafo_quantile(Normal(0, 1), Dual(0.5, 1)), quantile(Normal(0, 1), Dual(0.5, 1)), rtol = 10^-6) + @test isapprox(_trafo_quantile(Normal(Dual(0, 1, 0, 0), Dual(1, 0, 1, 0)), Dual(0.5, 0, 0, 1)), quantile(Normal(Dual(0, 1, 0, 0), Dual(1, 0, 1, 0)), Dual(0.5, 0, 0, 1)), rtol = 10^-6) + end + + @testset "trafo autodiff pullbacks" begin + x = [0.6, 0.7, 0.8, 0.9] + f = transport_to(Dirichlet([3.0, 4.0, 5.0, 6.0, 7.0]), Uniform) + @test isapprox(ForwardDiff.jacobian(f, x), Zygote.jacobian(f, x)[1], rtol = 10^-4) + f = inverse(transport_to(Normal, Dirichlet([3.0, 4.0, 5.0, 6.0, 7.0]))) + @test isapprox(ForwardDiff.jacobian(f, x), Zygote.jacobian(f, x)[1], rtol = 10^-4) + end + + + @testset "transport_to autosel" begin + for (M,R) in [ + (StandardDist{Normal}, StandardDist{Normal}) + (Normal, StandardDist{Normal}) + (StandardDist{Uniform}, StandardDist{Uniform}) + (Uniform, StandardDist{Uniform}) + ] + @test @inferred(transport_to(M, Weibull())) == transport_to(R(), Weibull()) + @test @inferred(transport_to(Weibull(), M)) == transport_to(Weibull(), R()) + @test @inferred(transport_to(M, MvNormal(float(I(5))))) == transport_to(R(5), MvNormal(float(I(5)))) + @test @inferred(transport_to(MvNormal(float(I(5))), M)) == transport_to(MvNormal(float(I(5))), R(5)) + @test @inferred(transport_to(M, StdExponential()^(2,3))) == transport_to(R(6), StdExponential()^(2,3)) + @test @inferred(transport_to(StdExponential()^(2,3), M)) == transport_to(StdExponential()^(2,3), R(6)) + end + end + + @testset "affine transformed distributions" begin + d = 2.0 * Weibull(0.7) + 1.0 + x = rand(StableRNG(789990641), d) + u = transport_to(StdUniform(), d)(x) + @test u ≈ cdf(d, x) + @test transport_to(d, StdUniform())(u) ≈ x + test_back_and_forth(StandardDist{Normal}(), d) + end + + @testset "truncated distributions" begin + d = truncated(Normal(0.3, 1.2), -0.5, 1.5) + for u in [0.0, 0.25, 0.75, 1.0, prevfloat(1.0)] + x = transport_to(d, StdUniform())(u) + @test minimum(d) <= x <= maximum(d) + end + test_back_and_forth(StandardDist{Uniform}(), d) + end + + @testset "products of distributions" begin + pd = product_distribution([Weibull(0.7), Exponential(1.3), Normal(0.5, 2.0)]) + m = MeasureBase.asmeasure(pd) + x = rand(StableRNG(789990641), pd) + for trg in [StdUniform()^3, StdNormal()^3] + y = transport_to(trg, m)(x) + y_ref = map((d_i, x_i) -> transport_to(trg.parent, d_i)(x_i), pd.v, x) + @test y ≈ y_ref + @test transport_to(m, trg)(y) ≈ x + end + + pd2 = product_distribution([Normal(2.0, 0.5), Weibull(1.2), Uniform(-1.0, 3.0)]) + m2 = MeasureBase.asmeasure(pd2) + y = transport_to(m2, m)(x) + @test transport_to(m, m2)(y) ≈ x + end + + @testset "MvNormal covariance representations" begin + for Σ in [PDMats.ScalMat(3, 2.5), PDMats.PDiagMat([0.5, 1.0, 2.5]), Diagonal([0.5, 1.0, 2.5])] + mvn = MvNormal([0.2, -0.4, 0.6], Σ) + x = rand(StableRNG(789990641), mvn) + y = transport_to(StandardDist{Normal}(3), mvn)(x) + @test transport_to(mvn, StandardDist{Normal}(3))(y) ≈ x + end + end +end diff --git a/test/runtests.jl b/test/runtests.jl index ff5c3140..f0d4488a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -9,7 +9,6 @@ using MeasureBase: test_interface, test_smf include("test_aqua.jl") -include("static.jl") include("test_primitive.jl") include("test_standard.jl") @@ -26,4 +25,6 @@ include("combinators/transformedmeasure.jl") include("combinators/reshape.jl") include("combinators/implicitlymapped.jl") +include("distributions/test_distributions.jl") + include("test_docs.jl") diff --git a/test/test_aqua.jl b/test/test_aqua.jl index b6290e31..d4546ac2 100644 --- a/test/test_aqua.jl +++ b/test/test_aqua.jl @@ -9,5 +9,10 @@ import MeasureBase #end # testset Test.@testset "Aqua tests" begin - Aqua.test_all(MeasureBase, ambiguities = false) + Aqua.test_all( + MeasureBase, + ambiguities = false, + # Only used by package extensions: + stale_deps = (ignore = [:ArgCheck, :ArraysOfArrays],), + ) end # testset From 6f51ef626317e293d59b98eb52b6d38c4149bc18 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Fri, 1 Nov 2024 20:28:07 +0100 Subject: [PATCH 26/75] Remove PowerWeightedMeasure Unused and untested. (cherry picked from commit 1d271885dfeeac1149817af1b748588b67a4d696) --- src/MeasureBase.jl | 1 - src/combinators/powerweighted.jl | 37 -------------------------------- 2 files changed, 38 deletions(-) delete mode 100644 src/combinators/powerweighted.jl diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index 4ab59601..e3e4cf5e 100644 --- a/src/MeasureBase.jl +++ b/src/MeasureBase.jl @@ -194,7 +194,6 @@ 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") 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 From eac759b51282493f94b2f449c93b86c19fc8ab73 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Fri, 1 Nov 2024 20:28:07 +0100 Subject: [PATCH 27/75] Remove kernelfactor Not used currently. (cherry picked from commit 98445838a025ddfc310e59f0e4470baf6263f6b1) --- src/parameterized.jl | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/src/parameterized.jl b/src/parameterized.jl index 78e43995..8b1c8c88 100644 --- a/src/parameterized.jl +++ b/src/parameterized.jl @@ -127,14 +127,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 From dc94b4a474e2312257a40f8e6c5438c6b65bccf1 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Fri, 1 Nov 2024 20:28:08 +0100 Subject: [PATCH 28/75] Remove the rebase function A rebase can easily be written explicitly. (cherry picked from commit fb3c98cecb995e48b64367ccc7122081c9113b2b) --- src/MeasureBase.jl | 1 - src/density.jl | 11 ----------- 2 files changed, 12 deletions(-) diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index e3e4cf5e..ecfbab2b 100644 --- a/src/MeasureBase.jl +++ b/src/MeasureBase.jl @@ -61,7 +61,6 @@ import HeterogeneousComputing using HeterogeneousComputing: real_numtype export gentype -export rebase export AbstractMeasure diff --git a/src/density.jl b/src/density.jl index a79021de..35147645 100644 --- a/src/density.jl +++ b/src/density.jl @@ -180,14 +180,3 @@ function logdensityof(μ::DensityMeasure, x::Any) convert(R, integrand_logval + base_logval)::R end end - -""" - rebase(μ, ν) - -Express `μ` in terms of a density over `ν`. Satisfies -``` -basemeasure(rebase(μ, ν)) == ν -density(rebase(μ, ν)) == 𝒹(μ,ν) -``` -""" -rebase(μ, ν) = ∫(𝒹(μ, ν), ν) From bb9c0213a065f9b790d47d305d1595370940c34b Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Fri, 1 Nov 2024 20:28:08 +0100 Subject: [PATCH 29/75] Removes PointwiseProductMeasure `mintegral` should be used instead to express posteriors. (cherry picked from commit 3c611806cd54740d2458a5192c97f77085c186cb) --- src/MeasureBase.jl | 1 - src/combinators/pointwise.jl | 30 ------------------------------ 2 files changed, 31 deletions(-) delete mode 100644 src/combinators/pointwise.jl diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index ecfbab2b..405114aa 100644 --- a/src/MeasureBase.jl +++ b/src/MeasureBase.jl @@ -190,7 +190,6 @@ include("combinators/product.jl") include("combinators/power.jl") include("combinators/spikemixture.jl") include("combinators/likelihood.jl") -include("combinators/pointwise.jl") include("combinators/restricted.jl") include("combinators/smart-constructors.jl") include("combinators/conditional.jl") 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 From 48f5e934113816021a3c8a52e42d2263a1fed299 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Fri, 1 Nov 2024 20:28:08 +0100 Subject: [PATCH 30/75] Remove operator otimes To be re-introduced in sub-module MeasureOperators. (cherry picked from commit 0cdca3d443b14b28313264f17bfdc085a190c394) --- src/combinators/product.jl | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/src/combinators/product.jl b/src/combinators/product.jl index 9135dc2b..dbdfffa3 100644 --- a/src/combinators/product.jl +++ b/src/combinators/product.jl @@ -179,18 +179,6 @@ function testvalue(::Type{T}, d::AbstractProductMeasure) where {T} _map(m -> testvalue(T, m), marginals(d)) end -export ⊗ - -""" - ⊗(μs::AbstractMeasure...) - -`⊗` is a binary operator for building product measures. This satisfies the law - -``` - basemeasure(μ ⊗ ν) == basemeasure(μ) ⊗ basemeasure(ν) -``` -""" -⊗(μs::AbstractMeasure...) = productmeasure(μs) ############################################################################### # I <: Base.Generator From f8e1a509eb05a12e50e50c9c54c0c7d93fc86cdc Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Fri, 1 Nov 2024 20:28:08 +0100 Subject: [PATCH 31/75] Remove scrd operator To be reintroduced in submodule MeasureOperators (cherry picked from commit 867adbeeeb0fe2bad0a2d691ddb54b050d0aa8e5) --- src/density.jl | 23 +---------------------- 1 file changed, 1 insertion(+), 22 deletions(-) diff --git a/src/density.jl b/src/density.jl index 35147645..0db02c29 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) @@ -73,16 +62,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) From 7f32b7e66d2b9cc7d911346e12f94b1961040224 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 11:00:59 +0200 Subject: [PATCH 32/75] Rename bind to mbind and remove fish operator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Combines the original commits f500aea, 86db05e and 58f0746 from the measure-algebra branch: `mbind` takes the kernel as the first argument (do-block support), the `↣` operator is removed (it looks very similar to the `>=>` "fish" operator, which is not a monadic bind) and `Bind` stores the kernel first. Co-Authored-By: Claude Fable 5 --- src/combinators/bind.jl | 53 +++++++++++++++++++++++------------------ 1 file changed, 30 insertions(+), 23 deletions(-) diff --git a/src/combinators/bind.jl b/src/combinators/bind.jl index cc2022f2..60b7cfd8 100644 --- a/src/combinators/bind.jl +++ b/src/combinators/bind.jl @@ -1,36 +1,43 @@ +""" + struct MeasureBase.Bind{M,K} <: AbstractMeasure + +Represents a monatic bind. User code should not create instances of `Bind` +directly, but should call `mbind(k, μ)` instead. +""" struct Bind{M,K} <: AbstractMeasure - μ::M k::K + μ::M +end + +getdof(d::Bind) = NoDOF{typeof(d)}() + +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 end -export ↣ """ -If -- μ is an `AbstractMeasure` or satisfies the Measure interface, and -- k is a function taking values from the support of μ and returning a measure + mbind(k, μ)::AbstractMeasure -Then `μ ↣ k` is a measure, called a *monadic bind*. In a -probabilistic programming language like Soss.jl, this could be expressed as +Given -Note that bind is usually written `>>=`, but this symbol is unavailable in Julia. +- a measure μ +- a kernel function k that takes values from the support of μ and returns a + measure + +The *monadic bind* operation `mbind(k, μ)` returns is a new measure. + +A monadic bind is often written as `>>=` (e.g. in Haskell), but this symbol is +unavailable in Julia. ``` -bind = @model μ,k begin - x ~ μ - y ~ k(x) - return y +μ = StdExponential() +ν = mbind(μ) do scale + pushfwd(Base.Fix1(*, scale), StdNormal()) end ``` - -See also `bind` and `Bind` """ -↣(μ, k) = bind(μ, k) - -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 -end +mbind(k, μ) = Bind(k, μ) +export mbind From 81b0e7993b0fd9fc27e02b3cc95df0e155508655 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 11:04:12 +0200 Subject: [PATCH 33/75] Introduce mintegrate and mintegrate_exp MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the operators ∫ and ∫exp (which will return in the new MeasureOperators submodule) with the functions `mintegrate` and `mintegrate_exp`, and introduces `MeasureBase.as_integrand` and `MeasureBase.as_integrand_exp` as specializable integrand conversion hooks (folded forward from the later Likelihood/mintegrate refactor on the measure-algebra branch). Adapts the Distributions extension and the tests to use the function names. Combines the original measure-algebra commits fefe36d, 74df0a5 and the mintegrate-related parts of 1ed3f95. Co-Authored-By: Claude Fable 5 --- .../measure_interface.jl | 3 +- src/density.jl | 132 ++++++++++++++---- test/distributions/test_measure_interface.jl | 2 +- test/test_basics.jl | 16 +-- 4 files changed, 117 insertions(+), 36 deletions(-) diff --git a/ext/MeasureBaseDistributionsExt/measure_interface.jl b/ext/MeasureBaseDistributionsExt/measure_interface.jl index 6fed5d4a..933a3f6f 100644 --- a/ext/MeasureBaseDistributionsExt/measure_interface.jl +++ b/ext/MeasureBaseDistributionsExt/measure_interface.jl @@ -23,4 +23,5 @@ Counting(MeasureBase.BoundedInts(static(0), static(Inf)))^size(d) -MeasureBase.∫(f, base::Distribution) = MeasureBase.∫(f, convert(AbstractMeasure, base)) +MeasureBase.mintegrate(f, base::Distribution) = + MeasureBase.mintegrate(f, convert(AbstractMeasure, base)) diff --git a/src/density.jl b/src/density.jl index 0db02c29..06dc98e1 100644 --- a/src/density.jl +++ b/src/density.jl @@ -50,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 @@ -77,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 @@ -99,42 +99,124 @@ end end function Pretty.tile(μ::DensityMeasure{F,B}) where {F,B} - result = Pretty.literal("DensityMeasure ∫(") + result = Pretty.literal("mintegrate(") result *= Pretty.pair_layout(Pretty.tile(μ.f), Pretty.tile(μ.base); sep = ", ") result *= Pretty.literal(")") end -export ∫ """ - ∫(f, base::AbstractMeasure) + MeasureBase.as_integrand(f) + MeasureBase.as_integrand(density) -Define a new measure in terms of a density `f` over some measure `base`. +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. + +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 + +@doc raw""" + mintegrate(f, μ::AbstractMeasure)::AbstractMeasure + mintegrate(density, μ::AbstractMeasure)::AbstractMeasure + +Returns a new measure that represents the indefinite +[integral](https://en.wikipedia.org/wiki/Radon%E2%80%93Nikodym_theorem) +of `f` with respect to `μ`. + +If a density is passed, it must implement the DensityInterface API. + +`ν = mintegrate(f, μ)` generates a measure `ν` that has the mathematical +interpretation + +```math +\nu(A) = \int_A f(a) \, \rm{d}\mu(a) +``` +""" +function mintegrate end +export mintegrate + +@inline mintegrate(obj, μ::AbstractMeasure) = DensityMeasure(as_integrand(obj), μ) + """ - ∫exp(f, base::AbstractMeasure) + MeasureBase.as_integrand_exp(log_f) + +Convert the logarithm of an integrand to an integrand. -Define a new measure in terms of a log-density `f` over some measure `base`. +See also [`MeasureBase.as_integrand`](@ref). """ -∫exp(f, base) = _logdensitymeasure(f, base, DensityKind(f)) +function as_integrand_exp end + +@inline as_integrand_exp(log_f) = _as_integrand_exp_default_impl(log_f, DensityKind(log_f)) + +@inline _as_integrand_exp_default_impl(log_f, ::NoDensity) = logfuncdensity(log_f) -function _logdensitymeasure(f, base, ::IsDensity) - @error "`∫exp(f, base)` is not valid when `DensityKind(f) == IsDensity()`. Use `∫(f, base)` instead." +function _as_integrand_exp_default_impl(log_f, ::IsDensity) + throw( + ArgumentError( + "`as_integrand_exp(log_f)` is not valid when `DensityKind(log_f) == IsDensity()`. Use `as_integrand(log_f)` instead.", + ), + ) end -function _logdensitymeasure(f, base, ::HasDensity) - @error "`∫exp(f, base)` is not valid when `DensityKind(f) == HasDensity()`." + +function _as_integrand_exp_default_impl(log_f, ::HasDensity) + throw( + ArgumentError( + "`as_integrand_exp(log_f)` is not valid when `DensityKind(log_f) == HasDensity()`.", + ), + ) end -_logdensitymeasure(f, base, ::NoDensity) = DensityMeasure(logfuncdensity(f), base) + + +@doc raw""" + mintegrate_exp(log_f, μ::AbstractMeasure) + +Given a function `log_f` that semantically represents the log of a function +`f`, `mintegrate_exp` returns a new measure that represents the indefinite +[integral](https://en.wikipedia.org/wiki/Radon%E2%80%93Nikodym_theorem) +of `f` with respect to `μ`. + +`ν = mintegrate_exp(log_f, μ)` generates a measure `ν` that has the +mathematical interpretation + +```math +\nu(A) = \int_A e^{log(f(a))} \, \rm{d}\mu(a) = \int_A f(a) \, \rm{d}\mu(a) +``` + +Note that `exp(log_f(...))` is usually not run explicitly, calculations that +involve the resulting measure are typically performed in log-space, +internally. +""" +function mintegrate_exp end +export mintegrate_exp + +mintegrate_exp(log_f, μ::AbstractMeasure) = DensityMeasure(as_integrand_exp(log_f), μ) + basemeasure(μ::DensityMeasure) = μ.base diff --git a/test/distributions/test_measure_interface.jl b/test/distributions/test_measure_interface.jl index d11d2889..f3d6c237 100644 --- a/test/distributions/test_measure_interface.jl +++ b/test/distributions/test_measure_interface.jl @@ -39,5 +39,5 @@ import MeasureBase @test (MeasureBase.params(c0)) == (α = 0.7, θ = 1.3) end - @test MeasureBase.∫(x -> Distributions.Normal(x, 0), Distributions.Normal()) isa MeasureBase.DensityMeasure + @test MeasureBase.mintegrate(x -> Distributions.Normal(x, 0), Distributions.Normal()) isa MeasureBase.DensityMeasure end diff --git a/test/test_basics.jl b/test/test_basics.jl index bd5a409c..11f1a8fe 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(μ) @@ -125,7 +125,7 @@ end 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()) @@ -136,12 +136,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 @@ -197,7 +195,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 @@ -229,13 +227,13 @@ 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 From f41a3bbdbaff0eb28b9ced27ad885768b6cf7b27 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 11:05:49 +0200 Subject: [PATCH 34/75] Add measure operators in submodule MeasureOperators MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The operators ⋄ (pushfwd), ⊙ (pullbck), ▷ (mbind), ⊗ (productmeasure), ∫ (mintegrate), ∫exp (mintegrate_exp), 𝒹 (density_rel) and log𝒹 (logdensity_rel) now live in the submodule `MeasureBase.MeasureOperators`, so that users can opt into the operator syntax explicitly. Ports the original measure-algebra commit 5404ff1. Co-Authored-By: Claude Fable 5 --- src/MeasureBase.jl | 2 + src/measure_operators.jl | 135 ++++++++++++++++++++++++++++++++++++++ test/measure_operators.jl | 24 +++++++ test/runtests.jl | 2 + 4 files changed, 163 insertions(+) create mode 100644 src/measure_operators.jl create mode 100644 test/measure_operators.jl diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index 405114aa..8841488b 100644 --- a/src/MeasureBase.jl +++ b/src/MeasureBase.jl @@ -209,6 +209,8 @@ include("rand.jl") include("density.jl") include("density-core.jl") +include("measure_operators.jl") + include("interface.jl") using .Interface diff --git a/src/measure_operators.jl b/src/measure_operators.jl new file mode 100644 index 00000000..41606367 --- /dev/null +++ b/src/measure_operators.jl @@ -0,0 +1,135 @@ +""" + module MeasureOperators + +Defines the following operators for measures: + +* `f ⋄ μ == pushfwd(f, μ)` +* `μ ⊙ f == pullbck(f, μ)` +* `μ ▷ k == mbind(k, μ)` +* `⊗(μs...) == productmeasure(μs)` +* `∫(f, μ) == mintegrate(f, μ)` +* `∫exp(f, μ) == mintegrate_exp(f, μ)` +* `𝒹(ν, μ) == density_rel(ν, μ)` +* `log𝒹(ν, μ) == logdensity_rel(ν, μ)` +""" +module MeasureOperators + +using MeasureBase: AbstractMeasure +using MeasureBase: pushfwd, pullbck, mbind, productmeasure +using MeasureBase: mintegrate, mintegrate_exp, density_rel, logdensity_rel +using InverseFunctions: inverse + +@doc raw""" + ⋄(f, μ::AbstractMeasure) = pushfwd(f, μ) + +The `\\diamond` operator denotes a pushforward operation: `ν = f ⋄ μ` +generates a +[pushforward measure](https://en.wikipedia.org/wiki/Pushforward_measure). + +A common mathematical notation for a pushforward is ``f_*μ``, but as +there is no "subscript-star" operator in Julia, we use `⋄`. + +See [`pushfwd(f, μ)`](@ref) for details. + +Also see [`ν ⊙ f`](@ref), the pullback operator. +""" +⋄(f, μ::AbstractMeasure) = pushfwd(f, μ) +export ⋄ + +@doc raw""" + ⊙(ν::AbstractMeasure, f) = pullbck(f, ν) + +The `\\odot` operator denotes a pullback operation. + +See also [`pullbck(ν, f)`](@ref) for details. Note that `pullbck` takes it's +arguments in different order, in keeping with the Julia convention of +passing functions as the first argument. A pullback is mathematically the +precomposition of a measure `μ`` with the function `f` applied to sets. so +`⊙` takes the measure as the first and the function as the second argument, +as common in mathematical notation for precomposition. + +A common mathematical notation for pullback in measure theory is +``f \circ μ``, but as `∘` is used for function composition in Julia and as +`f` semantically acts point-wise on sets, we use `⊙`. + +Also see [f ⋄ μ](@ref), the pushforward operator. +""" +⊙(ν::AbstractMeasure, f) = pullbck(f, ν) +export ⊙ + +""" + μ ▷ k = mbind(k, μ) + +The `\\triangleright` operator denotes a measure monadic bind operation. + +A common operator choice for a monadic bind operator is `>>=` (e.g. in +the Haskell programming language), but this has a different meaning in +Julia and there is no close equivalent, so we use `▷`. + +See [`mbind(k, μ)`](@ref) for details. Note that `mbind` takes its +arguments in different order, in keeping with the Julia convention of +passing functions as the first argument. `▷`, on the other hand, takes +its arguments in the order common for monadic binds in functional +programming (like the Haskell `>>=` operator) and mathematics. +""" +▷(μ::AbstractMeasure, k) = mbind(k, μ) +export ▷ + +# ToDo: Use `⨂` instead of `⊗` for better readability? +""" + ⊗(μs::AbstractMeasure...) = productmeasure(μs) + +`⊗` is an operator for building product measures. + +See [`productmeasure(μs)`](@ref) for details. +""" +⊗(μs::AbstractMeasure...) = productmeasure(μs) +export ⊗ + +""" + ∫(f, μ::AbstractMeasure) = mintegrate(f, μ) + +Denotes an indefinite integral of the function `f` with respect to the +measure `μ`. + +See [`mintegrate(f, μ)`](@ref) for details. +""" +∫(f, μ::AbstractMeasure) = mintegrate(f, μ) +export ∫ + +""" + ∫exp(f, μ::AbstractMeasure) = mintegrate_exp(f, μ) + +Generates a new measure that is the indefinite integral of `exp` of `f` +with respect to the measure `μ`. + +See [`mintegrate_exp(f, μ)`](@ref) for details. +""" +∫exp(f, μ::AbstractMeasure) = mintegrate_exp(f, μ) +export ∫exp + +""" + 𝒹(ν, μ) = density_rel(ν, μ) + +Compute the density, i.e. the +[Radom-Nikodym derivative](https://en.wikipedia.org/wiki/Radon%E2%80%93Nikodym_theorem) +of `ν`` with respect to `μ`. + +For details, see [`density_rel(ν, μ)`}(@ref). +""" +𝒹(ν, μ::AbstractMeasure) = density_rel(ν, μ) +export 𝒹 + +""" + log𝒹(ν, μ) = logdensity_rel(ν, μ) + +Compute the log-density, i.e. the logarithm of the +[Radom-Nikodym derivative](https://en.wikipedia.org/wiki/Radon%E2%80%93Nikodym_theorem) +of `ν`` with respect to `μ`. + +For details, see [`logdensity_rel(ν, μ)`}(@ref). +""" +log𝒹(ν, μ::AbstractMeasure) = logdensity_rel(ν, μ) +export log𝒹 + +end # module MeasureOperators diff --git a/test/measure_operators.jl b/test/measure_operators.jl new file mode 100644 index 00000000..1530f191 --- /dev/null +++ b/test/measure_operators.jl @@ -0,0 +1,24 @@ +using Test + +using MeasureBase: AbstractMeasure +using MeasureBase: StdExponential, StdLogistic, StdNormal, StdUniform +using MeasureBase: pushfwd, pullbck, mbind, productmeasure +using MeasureBase: mintegrate, mintegrate_exp, density_rel, logdensity_rel +using MeasureBase.MeasureOperators: ⋄, ⊙, ▷, ⊗, ∫, ∫exp, 𝒹, log𝒹 + +@testset "MeasureOperators" begin + μ = StdExponential() + ν = StdUniform() + k(σ) = pushfwd(x -> σ * x, StdNormal()) + μs = (StdExponential(), StdLogistic(), StdUniform()) + f = sqrt + + @test @inferred(f ⋄ μ) == pushfwd(f, μ) + @test @inferred(ν ⊙ f) == pullbck(f, ν) + @test @inferred(μ ▷ k) == mbind(k, μ) + @test @inferred(⊗(μs...)) == productmeasure(μs) + @test @inferred(∫(f, μ)) == mintegrate(f, μ) + @test @inferred(∫exp(f, μ)) == mintegrate_exp(f, μ) + @test @inferred(𝒹(ν, μ)) == density_rel(ν, μ) + @test @inferred(log𝒹(ν, μ)) == logdensity_rel(ν, μ) +end diff --git a/test/runtests.jl b/test/runtests.jl index f0d4488a..7183c58f 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -20,6 +20,8 @@ include("smf.jl") include("test_mooncake.jl") +include("measure_operators.jl") + include("combinators/weighted.jl") include("combinators/transformedmeasure.jl") include("combinators/reshape.jl") From c5de45408c130a6ac8fd4fdbf8a792a6078026f1 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 11:08:55 +0200 Subject: [PATCH 35/75] Rework likelihoods around AbstractLikelihood MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce the `AbstractLikelihood <: Function` supertype with the `likelihood_kernel`/`likelihood_obs` accessor API and the `as_likelihood`/`as_integrand` conversion hooks (including default implementations for `Base.Fix2(densityof, x) ∘ f`-style objects, so likelihood-like functions and densities compose with `mintegrate`). `Likelihood` becomes a plain kernel/observation container; density evaluation goes through the generic `AbstractLikelihood` methods and converts kernel results via `asmeasure`. Combines the original measure-algebra commits bf35d40, aa7416b, 470db23 and the likelihood-related parts of 1ed3f95. Co-Authored-By: Claude Fable 5 --- src/MeasureBase.jl | 6 +- src/combinators/likelihood.jl | 289 +++++++++++++++++----------------- 2 files changed, 151 insertions(+), 144 deletions(-) diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index 8841488b..e9fcb312 100644 --- a/src/MeasureBase.jl +++ b/src/MeasureBase.jl @@ -176,6 +176,9 @@ include("primitive.jl") include("utils.jl") include("mass-interface.jl") +include("density.jl") +include("density-core.jl") + include("primitives/counting.jl") include("primitives/lebesgue.jl") include("primitives/dirac.jl") @@ -206,9 +209,6 @@ include("combinators/half.jl") include("rand.jl") -include("density.jl") -include("density-core.jl") - include("measure_operators.jl") include("interface.jl") 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) From 2fb782e9ce4d3ce0781e30776f3892377d6d07e8 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 11:09:29 +0200 Subject: [PATCH 36/75] Remove remnants of pointwiseproduct Ports the original measure-algebra commit 28c6d30. Co-Authored-By: Claude Fable 5 --- src/combinators/smart-constructors.jl | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl index 26ba3948..45b1c5fa 100644 --- a/src/combinators/smart-constructors.jl +++ b/src/combinators/smart-constructors.jl @@ -4,18 +4,6 @@ half(μ::AbstractMeasure) = Half(μ) -############################################################################### -# PointwiseProductMeasure - -function pointwiseproduct(μ::AbstractMeasure, ℓ::Likelihood) - T = Core.Compiler.return_type(ℓ.k, Tuple{gentype(μ)}) - return pointwiseproduct(T, μ, ℓ) -end - -function pointwiseproduct(::Type{T}, μ::AbstractMeasure, ℓ::Likelihood) where {T} - return PointwiseProductMeasure(μ, ℓ) -end - ############################################################################### # PowerMeaure From 15a0b9ec0f0b53fac14c32839cccb94fa14f2842 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 11:09:52 +0200 Subject: [PATCH 37/75] Remove splat (unused here, and Base has it now) Ports the original measure-algebra commit 81a7d93. Co-Authored-By: Claude Fable 5 --- src/MeasureBase.jl | 1 - src/splat.jl | 11 ----------- 2 files changed, 12 deletions(-) delete mode 100644 src/splat.jl diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index e9fcb312..e7c9f8a8 100644 --- a/src/MeasureBase.jl +++ b/src/MeasureBase.jl @@ -167,7 +167,6 @@ include("smf.jl") include("getdof.jl") include("transport.jl") include("schema.jl") -include("splat.jl") include("proxies.jl") include("kernel.jl") include("parameterized.jl") 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) From 5bc1389d1fcf4fdc4fd5a166051912dcaf4e2671 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 12:53:54 +0200 Subject: [PATCH 38/75] Add fast_dof, some_dof and NoFastInsupport Introduce the `AbstractNoDOF` type hierarchy with absorbing arithmetic, `NoFastDOF` and `fast_dof` (defaults to `getdof`, to be specialized by measures whose DOF can't be computed instantly) plus `some_dof` (DOF at an unspecified point, via `localmeasure`). `_default_getdof` now returns a `NoDOF` instance instead of the type, and `check_dof` uses `fast_dof` and tolerates unknown DOF. Add `NoFastInsupport` and make `require_insupport` tolerate it, so measures like monadic binds can declare that support checking is not cheap. Ports parts of the original measure-algebra commits 7e93050 and later STASH fixes. Co-Authored-By: Claude Fable 5 --- src/getdof.jl | 95 +++++++++++++++++++++++++++++++++++++++++++++--- src/insupport.jl | 27 +++++++++++--- 2 files changed, 110 insertions(+), 12 deletions(-) diff --git a/src/getdof.jl b/src/getdof.jl index 16ae7cc6..c2edd15c 100644 --- a/src/getdof.jl +++ b/src/getdof.jl @@ -1,11 +1,30 @@ """ - MeasureBase.NoDOF{MU} + abstract type MeasureBase.AbstractNoDOF{MU} + +Abstract supertype for [`NoDOF`](@ref) and [`NoFastDOF`](@ref). +""" +abstract type AbstractNoDOF{MU} end + +Base.:+(nodof::AbstractNoDOF) = nodof +Base.:+(::IntegerLike, nodof::AbstractNoDOF) = nodof +Base.:+(nodof::AbstractNoDOF, ::IntegerLike) = nodof +Base.:+(nodof::AbstractNoDOF, ::AbstractNoDOF) = nodof + +Base.:*(nodof::AbstractNoDOF) = nodof +Base.:*(::IntegerLike, nodof::AbstractNoDOF) = nodof +Base.:*(nodof::AbstractNoDOF, ::IntegerLike) = nodof +Base.:*(nodof::AbstractNoDOF, ::AbstractNoDOF) = nodof + + +""" + MeasureBase.NoDOF{MU} <: AbstractNoDOF{MU} Indicates that there is no way to compute degrees of freedom of a measure of type `MU` with the given information, e.g. because the DOF are not a global property of the measure. """ -struct NoDOF{MU} end +struct NoDOF{MU} <: AbstractNoDOF{MU} end + """ getdof(μ) @@ -22,22 +41,86 @@ Also see [`check_dof`](@ref). function getdof end # Prevent infinite recursion: -@inline _default_getdof(::Type{MU}, ::MU) where {MU} = NoDOF{MU} +@inline _default_getdof(::Type{MU}, ::MU) where {MU} = NoDOF{MU}() @inline _default_getdof(::Type{MU}, mu_base) where {MU} = getdof(mu_base) @inline getdof(μ::MU) where {MU} = _default_getdof(MU, basemeasure(μ)) + +""" + MeasureBase.NoFastDOF{MU} <: AbstractNoDOF{MU} + +Indicates that there is no way to compute the degrees of freedom of a +measure of type `MU` efficiently. +""" +struct NoFastDOF{MU} <: AbstractNoDOF{MU} end + + +""" + fast_dof(μ::MU) + +Returns the effective number of degrees of freedom of variates of +measure `μ`, if it can be computed efficiently, otherwise +returns [`NoFastDOF{MU}()`](@ref). + +Defaults to `getdof(μ)` and should be specialized for measures for +which DOF can't be computed instantly. + +Also see [`getdof`](@ref) and [`check_dof`](@ref). +""" +function fast_dof end +export fast_dof + +fast_dof(μ) = getdof(μ) + + +""" + MeasureBase.some_dof(μ::AbstractMeasure) + +Get the DOF at some unspecified point of measure `μ`. + +Use with caution! + +In general, use [`getdof(μ)`](@ref) instead. `some_dof` is useful for +measures that are expected to have a constant DOF over their whole +space, but for which there is no way to compute it (or prove that +the DOF is constant over the measurable space). +""" +function some_dof end + +function some_dof(μ) + m = asmeasure(μ) + _try_direct_dof(m, getdof(m)) +end + +_try_direct_dof(::AbstractMeasure, dof::IntegerLike) = dof +_try_direct_dof(μ::AbstractMeasure, ::AbstractNoDOF) = + _try_local_dof(μ, some_dof(_some_localmeasure(μ))) + +_try_local_dof(::AbstractMeasure, dof::IntegerLike) = dof +_try_local_dof(μ::AbstractMeasure, ::AbstractNoDOF) = + throw(ArgumentError("Can't determine DOF for measure of type $(nameof(typeof(μ)))")) + +_some_localmeasure(μ::AbstractMeasure) = localmeasure(μ, testvalue(μ)) + + """ MeasureBase.check_dof(ν, μ)::Nothing Check if `ν` and `μ` have the same effective number of degrees of freedom -according to [`MeasureBase.getdof`](@ref). +according to [`MeasureBase.fast_dof`](@ref). + +Does not throw an exception if the DOF of `ν` or `μ` can't be computed +efficiently. """ function check_dof end function check_dof(ν, μ) - n_ν = getdof(ν) - n_μ = getdof(μ) + n_ν = fast_dof(ν) + n_μ = fast_dof(μ) + if n_ν isa AbstractNoDOF || n_μ isa AbstractNoDOF + return nothing + end if n_ν != n_μ throw( ArgumentError( 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 From c123ef05da26200a92582cfee02b68c73d0ff01e Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 13:00:04 +0200 Subject: [PATCH 39/75] Add localmeasure and transportmeasure Introduce the `localmeasure`/`transportmeasure` interface: both default to the measure itself and return a measure that behaves like the original one in the infinitesimal neighborhood of a point, for density calculation resp. density calculation and transport. Measures like monadic binds specialize them. `unsafe_logdensityof` and `unsafe_logdensity_rel` now evaluate via the local measure, with an extra dispatch boundary to reduce the number of required specializations. `logdensity_rel` handles `NoFastInsupport`, and `_checksupport` passes results through unchanged in that case. Transport now converts its arguments via `asmeasure`, uses `fast_dof` for the standard intermediate measure, guards against overly deep transport-origin stacks, and powers of `NoTransportOrigin` stay `NoTransportOrigin`. The two-argument form of `basemeasure` is gone. Ports parts of the original measure-algebra commits fe28297, 95632f4, 67cf8ea and related STASH commits. Co-Authored-By: Claude Fable 5 --- src/density-core.jl | 93 ++++++++++++++++++++++++++++++++++++++++++--- src/interface.jl | 2 +- src/transport.jl | 23 +++++++++-- src/utils.jl | 2 - 4 files changed, 109 insertions(+), 11 deletions(-) diff --git a/src/density-core.jl b/src/density-core.jl index f3b2db2b..33a3c3cb 100644 --- a/src/density-core.jl +++ b/src/density-core.jl @@ -38,6 +38,40 @@ end end _checksupport(cond, result) = ifelse(cond == true, result, oftype(result, -Inf)) +@inline _checksupport(::NoFastInsupport, result) = result + + +""" + localmeasure(m::AbstractMeasure, x)::AbstractMeasure + +Return a measure that behaves like `m` in the infinitesimal neighborhood +of `x` in respect to density calculation. + +Note that the resulting measure may not be well defined outside of the +infinitesimal neighborhood of `x`. + +For most measure types simply returns `m` itself. [`mbind`](@ref), +for example, generates measures for which `localmeasure(m, x)` depends +on `x`. +""" +localmeasure(m::AbstractMeasure, x) = m +export localmeasure + + +""" + MeasureBase.transportmeasure(m::AbstractMeasure, x)::AbstractMeasure + +Return a measure that behaves like `m` in the infinitesimal neighborhood +of `x` in respect to both transport and density calculation. + +Note that the resulting measure may not be well defined outside of the +infinitesimal neighborhood of `x`. + +For most measure types simply returns `m` itself. [`mbind`](@ref), +for example, generates measures for which `transportmeasure(m, x)` depends +on `x`. +""" +transportmeasure(m::AbstractMeasure, x) = m export unsafe_logdensityof @@ -50,11 +84,17 @@ This is "unsafe" because it does not check `insupport(m, x)`. See also `logdensityof`. """ -@inline function unsafe_logdensityof(μ::M, x) where {M} +@inline function unsafe_logdensityof(μ::AbstractMeasure, x) + μ_local = localmeasure(μ, x) + # Extra dispatch boundary to reduce number of required specializations of implementation: + return _unsafe_logdensityof_local(μ_local, x) +end + +@inline function _unsafe_logdensityof_local(μ::M, x) where {M} ℓ_0 = logdensity_def(μ, x) b_0 = μ Base.Cartesian.@nexprs 10 i -> begin # 10 is just some "big enough" number - b_{i} = basemeasure(b_{i - 1}, x) + b_{i} = basemeasure(b_{i - 1}) # The below makes the evaluated code shorter, but screws up Zygote # if b_{i} isa typeof(b_{i - 1}) @@ -76,20 +116,56 @@ known to be in the support of both, it can be more efficient to call `unsafe_logdensity_rel`. """ @inline function logdensity_rel(μ::M, ν::N, x::X) where {M,N,X} + inμ = insupport(μ, x) + inν = insupport(ν, x) + return _logdensity_rel_impl(μ, ν, x, inμ, inν) +end + +@inline function _logdensity_rel_impl(μ::M, ν::N, x::X, inμ::Bool, inν::Bool) where {M,N,X} T = unstatic( promote_type( return_type(logdensity_def, (μ, x)), return_type(logdensity_def, (ν, x)), ), ) - inμ = insupport(μ, x) - inν = insupport(ν, x) istrue(inμ) || return convert(T, ifelse(inν, -Inf, NaN)) istrue(inν) || return convert(T, Inf) return unsafe_logdensity_rel(μ, ν, x) end +@inline function _logdensity_rel_impl( + μ::M, + ν::N, + x::X, + @nospecialize(::NoFastInsupport), + @nospecialize(::NoFastInsupport) +) where {M,N,X} + unsafe_logdensity_rel(μ, ν, x) +end + +@inline function _logdensity_rel_impl( + μ::M, + ν::N, + x::X, + inμ::Bool, + @nospecialize(::NoFastInsupport) +) where {M,N,X} + logd = unsafe_logdensity_rel(μ, ν, x) + return istrue(inμ) ? logd : oftype(logd, -Inf) +end + +@inline function _logdensity_rel_impl( + μ::M, + ν::N, + x::X, + @nospecialize(::NoFastInsupport), + inν::Bool +) where {M,N,X} + logd = unsafe_logdensity_rel(μ, ν, x) + return istrue(inν) ? logd : oftype(logd, +Inf) +end + """ unsafe_logdensity_rel(m1, m2, x) @@ -98,7 +174,14 @@ 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} +@inline function unsafe_logdensity_rel(μ::AbstractMeasure, ν::AbstractMeasure, x) + μ_local = localmeasure(μ, x) + ν_local = localmeasure(ν, x) + # Extra dispatch boundary to reduce number of required specializations of implementation: + return _unsafe_logdensity_rel_local(μ_local, ν_local, x) +end + +@inline function _unsafe_logdensity_rel_local(μ::M, ν::N, x::X) where {M,N,X} if static_hasmethod(logdensity_def, Tuple{M,N,X}) return logdensity_def(μ, ν, x) end diff --git a/src/interface.jl b/src/interface.jl index 4890ddd6..6003203d 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -64,7 +64,7 @@ function test_interface(μ::M) where {M} # testvalue, logdensityof x = @inferred testvalue(Float64, μ) - β = @inferred basemeasure(μ, x) + β = @inferred basemeasure(μ) ℓμ = @inferred logdensityof(μ, x) ℓβ = @inferred logdensityof(β, x) diff --git a/src/transport.jl b/src/transport.jl index b0c8ed41..9cea6fed 100644 --- a/src/transport.jl +++ b/src/transport.jl @@ -8,6 +8,8 @@ See [`MeasureBase.transport_origin`](@ref). """ struct NoTransportOrigin{NU} end +Base.:^(origin::NoTransportOrigin, ::IntegerLike) = origin + """ MeasureBase.transport_origin(ν) @@ -147,6 +149,21 @@ end μ, x, ) where {n_ν,n_μ} + if n_ν == 10 + return :(throw( + ArgumentError( + "Transport to measure of type $(nameof(typeof(ν))) not supported, origin stack too deep.", + ), + )) + end + if n_μ == 10 + return :(throw( + ArgumentError( + "Transport from measure of type $(nameof(typeof(μ))) not supported, origin stack too deep.", + ), + )) + end + prog = quote μ0 = μ x0 = x @@ -183,8 +200,8 @@ end return prog end -@inline _transport_intermediate(ν, μ) = _transport_intermediate(getdof(ν), getdof(μ)) -@inline _transport_intermediate(::Integer, n_μ::Integer) = StdUniform()^n_μ +@inline _transport_intermediate(ν, μ) = _transport_intermediate(fast_dof(ν), fast_dof(μ)) +@inline _transport_intermediate(::IntegerLike, n_μ::IntegerLike) = StdUniform()^n_μ @inline _transport_intermediate(::StaticInteger{1}, ::StaticInteger{1}) = StdUniform() _call_transport_def(ν, μ, x) = transport_def(ν, μ, x) @@ -227,7 +244,7 @@ struct TransportFunction{NU,MU} <: Function end end -@inline transport_to(ν, μ) = TransportFunction(ν, μ) +@inline transport_to(ν, μ) = TransportFunction(asmeasure(ν), asmeasure(μ)) function Base.:(==)(a::TransportFunction, b::TransportFunction) return a.ν == b.ν && a.μ == b.μ diff --git a/src/utils.jl b/src/utils.jl index c1e97034..e169c7c1 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) From 596593c8126e0ee7663f2e3ec0114a83dd6e298f Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 13:02:35 +0200 Subject: [PATCH 40/75] Add collection utils, StdPowerMeasure and power measure fast_dof Add the collection helpers `_as_tuple`, `_get_or_view`, `_split_after` (vector/tuple/NamedTuple), `_fill_value`/`_fill_axes` and `_flatten_to_rv`, in preparation for the generalized product transport and combined measures. Add the `StdPowerMeasure` type alias, `fast_dof` for power measures and make power measure `insupport` propagate `NoFastInsupport` from the parent measure. Ports parts of the original measure-algebra commits 763bd70, cd2db56 and 7e93050. Co-Authored-By: Claude Fable 5 --- src/MeasureBase.jl | 3 ++ src/collection_utils.jl | 72 ++++++++++++++++++++++++++++++++++++++ src/combinators/power.jl | 15 +++++++- src/standard/stdmeasure.jl | 7 ++++ 4 files changed, 96 insertions(+), 1 deletion(-) diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index e7c9f8a8..2aa92089 100644 --- a/src/MeasureBase.jl +++ b/src/MeasureBase.jl @@ -60,6 +60,9 @@ using StaticThings: import HeterogeneousComputing using HeterogeneousComputing: real_numtype +using ArraysOfArrays: + VectorOfArrays, VectorOfVectors, VectorOfSimilarArrays, VectorOfSimilarVectors, flatview + export gentype export AbstractMeasure diff --git a/src/collection_utils.jl b/src/collection_utils.jl index 1de51f7e..1ef9f9fd 100644 --- a/src/collection_utils.jl +++ b/src/collection_utils.jl @@ -22,3 +22,75 @@ _rev_cumsum(xs::AbstractVector) = reverse(cumsum(reverse(xs))) # Equivalent to `cumprod(xs)``: _exp_cumsum_log(xs::AbstractVector) = exp.(cumsum(log.(xs))) + +Base.@propagate_inbounds _as_tuple(v::AbstractVector, ::Val{N}) where {N} = Tuple(SVector{N}(v)) + + +Base.@propagate_inbounds function _get_or_view(A::AbstractVector, from::IntegerLike, until::IntegerLike) + view(A, from:until) +end + +Base.@propagate_inbounds function _get_or_view( + A::AbstractVector, + ::StaticInteger{from}, + ::StaticInteger{until}, +) where {from,until} + SVector{until - from + 1}(view(A, from:until)) +end + +# ToDo: Specialize for StaticVector instead of SVector? +Base.@propagate_inbounds function _get_or_view( + A::SVector, + from::StaticInteger, + until::StaticInteger, +) + # ToDo: Improve implementation: + SVector(_get_or_view(Tuple(A), from, until)) +end + +Base.@propagate_inbounds function _get_or_view(tpl::Tuple, from::IntegerLike, until::IntegerLike) + ntuple(i -> tpl[from + i - 1], Val(until - from + 1)) +end + + +@inline function _split_after(x::AbstractVector, n::IntegerLike) + idxs = maybestatic_eachindex(x) + i_first = maybestatic_first(idxs) + i_last = maybestatic_last(idxs) + _get_or_view(x, i_first, i_first + n - one(n)), _get_or_view(x, i_first + n, i_last) +end + +@inline _split_after(x::Tuple, n) = _split_after(x::Tuple, Val{n}()) +@inline _split_after(x::Tuple, ::Val{N}) where {N} = x[begin:(begin+N-1)], x[(begin+N):end] + +@generated function _split_after(x::NamedTuple{names}, ::Val{names_a}) where {names,names_a} + n = length(names_a) + if names[begin:(begin+n-1)] == names_a + names_b = names[(begin+n):end] + quote + a, b = _split_after(values(x), Val($n)) + NamedTuple{$names_a}(a), NamedTuple{$names_b}(b) + end + else + quote + throw(ArgumentError("Can't split NamedTuple{$names} after {$names_a}")) + end + end +end + + +# Field access functions for Fill: +_fill_value(x::FillArrays.Fill) = x.value +_fill_axes(x::FillArrays.Fill) = x.axes + + +_flatten_to_rv(VV::AbstractVector{<:AbstractVector{<:Real}}) = flatview(VectorOfArrays(VV)) +_flatten_to_rv(VV::AbstractVector{<:StaticVector{N,<:Real}}) where {N} = + flatview(VectorOfSimilarArrays(VV)) + +_flatten_to_rv(VV::VectorOfSimilarVectors{<:Real}) = flatview(VV) +_flatten_to_rv(VV::VectorOfVectors{<:Real}) = flatview(VV) + +_flatten_to_rv(::Tuple{}) = [] +_flatten_to_rv(tpl::Tuple{Vararg{AbstractVector}}) = vcat(tpl...) +_flatten_to_rv(tpl::Tuple{Vararg{StaticVector}}) = vcat(tpl...) diff --git a/src/combinators/power.jl b/src/combinators/power.jl index 6f065c3f..7e22fb36 100644 --- a/src/combinators/power.jl +++ b/src/combinators/power.jl @@ -144,15 +144,28 @@ end end end +_all(A) = all(A) +_all(::AbstractArray{NoFastInsupport{T}}) where {T} = NoFastInsupport{T}() + @inline function insupport(μ::PowerMeasure, x::AbstractArray) p = μ.parent - all(x) do xj + insupp = broadcast(x) do xj # https://github.com/SciML/Static.jl/issues/36 dynamic(insupport(p, xj)) end + _all(insupp) end @inline getdof(μ::PowerMeasure) = getdof(μ.parent) * 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 @propagate_inbounds function checked_arg(μ::PowerMeasure, x::AbstractArray{<:Any}) @boundscheck begin diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl index a9c09b5c..3dd30e24 100644 --- a/src/standard/stdmeasure.jl +++ b/src/standard/stdmeasure.jl @@ -4,6 +4,13 @@ StdMeasure(::typeof(rand)) = StdUniform() StdMeasure(::typeof(randexp)) = StdExponential() StdMeasure(::typeof(randn)) = StdNormal() +""" + 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}} + @inline check_dof(::StdMeasure, ::StdMeasure) = nothing @inline transport_def(::MU, μ::MU, x) where {MU<:StdMeasure} = x From c97bd3d872128a6dec690f04e55fe84f1f4871bb Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 13:07:57 +0200 Subject: [PATCH 41/75] Canonical measure nesting in smart constructors Rework the `powermeasure` and `productmeasure` smart constructors to maintain a canonical measure type nesting (WeightedMeasure over Dirac over PowerMeasure over ProductMeasure): powers and products of Dirac measures become Dirac measures of filled values, weights are pulled out of powers, marginal collections are converted via `asmeasure`, products over `Fill` and over singleton-typed arrays collapse to power measures. Product measures also get `fast_dof` and `NoFastInsupport`-aware `insupport`, plus a power-measure proxy for `Fill` marginals. `powermeasure` now accepts sizes, axes and plain integer exponents. Ports the original measure-algebra commits 00b4d61, b28a902 and parts of related STASH commits. Co-Authored-By: Claude Fable 5 --- src/combinators/power.jl | 4 - src/combinators/product.jl | 11 ++- src/combinators/smart-constructors.jl | 106 +++++++++++++++++++++----- 3 files changed, 96 insertions(+), 25 deletions(-) diff --git a/src/combinators/power.jl b/src/combinators/power.jl index 7e22fb36..1ce9ca98 100644 --- a/src/combinators/power.jl +++ b/src/combinators/power.jl @@ -75,10 +75,6 @@ function Base.rand(rng::AbstractRNG, ::Type{T}, d::PowerMeasure) where {T} end end -@inline function powermeasure(x::T, sz::Tuple{Vararg{Any,N}}) where {T,N} - PowerMeasure(x, asaxes(sz)) -end - marginals(d::PowerMeasure) = maybestatic_fill(d.parent, d.axes) function Base.:^(μ::AbstractMeasure, dims::Tuple{Vararg{AbstractArray,N}}) where {N} diff --git a/src/combinators/product.jl b/src/combinators/product.jl index dbdfffa3..468c38c0 100644 --- a/src/combinators/product.jl +++ b/src/combinators/product.jl @@ -28,6 +28,9 @@ Base.size(μ::AbstractProductMeasure) = size(marginals(μ)) basemeasure(d::AbstractProductMeasure) = productmeasure(map(basemeasure, marginals(d))) +proxy(μ::ProductMeasure{<:FillArrays.Fill}) = + powermeasure(_fill_value(marginals(μ)), _fill_axes(marginals(μ))) + function Base.rand(rng::AbstractRNG, ::Type{T}, d::AbstractProductMeasure) where {T} mar = marginals(d) _rand_product(rng, T, mar, eltype(mar)) @@ -221,12 +224,16 @@ end @inline function insupport(d::AbstractProductMeasure, x) for (mj, xj) in zip(marginals(d), x) - dynamic(insupport(mj, xj)) || return false + insup = dynamic(insupport(mj, xj)) + if insup isa NoFastInsupport || insup == false + return insup + end end return true end -getdof(d::AbstractProductMeasure) = mapreduce(getdof, +, marginals(d)) +getdof(d::AbstractProductMeasure) = sum(getdof, marginals(d)) +fast_dof(d::AbstractProductMeasure) = sum(fast_dof, marginals(d)) function checked_arg(μ::ProductMeasure{<:NTuple{N,Any}}, x::NTuple{N,Any}) where {N} map(checked_arg, marginals(μ), x) diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl index 45b1c5fa..cb15f6ad 100644 --- a/src/combinators/smart-constructors.jl +++ b/src/combinators/smart-constructors.jl @@ -1,4 +1,9 @@ +# Canonical measure type nesting, outer to inner: +# +# WeightedMeasure, Dirac, PowerMeasure, ProductMeasure + + ############################################################################### # Half @@ -7,44 +12,107 @@ half(μ::AbstractMeasure) = Half(μ) ############################################################################### # PowerMeaure -powermeasure(m::AbstractMeasure, ::Tuple{}) = m +""" + powermeasure(μ, dims) + powermeasure(μ, axes) + +Constructs a power of a measure `μ`. + +`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)) + +@inline _generic_powermeasure_stage1(μ::AbstractMeasure, ::Tuple{}) = μ -function powermeasure( - μ::WeightedMeasure, - dims::Tuple{<:AbstractArray,Vararg{AbstractArray}}, -) - k = mapreduce(length, *, dims) * μ.logweight - return weightedmeasure(k, μ.base^dims) +@inline function _generic_powermeasure_stage1(μ::AbstractMeasure, exponent::Tuple) + _generic_powermeasure_stage2(μ, exponent) end -function powermeasure(μ::WeightedMeasure, dims::NonEmptyTuple) - k = prod(dims) * μ.logweight - return weightedmeasure(k, μ.base^dims) +@inline _generic_powermeasure_stage2(μ::AbstractMeasure, exponent::Tuple) = + PowerMeasure(μ, exponent) + +@inline function _generic_powermeasure_stage2(μ::Dirac, exponent::Tuple) + Dirac(maybestatic_fill(μ.x, exponent)) +end + +@inline function _generic_powermeasure_stage2(μ::WeightedMeasure, exponent::Tuple) + ν = μ.base^exponent + k = maybestatic_length(ν) * μ.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(Base.Fix1(*, scale), StdExponential()) for scale in 0.1:0.2:2]) +``` +""" +function productmeasure end +export productmeasure + +@inline productmeasure(mar) = _generic_productmeasure_impl(mar) + +@inline _generic_productmeasure_impl(mar::FillArrays.Fill) = + powermeasure(_fill_value(mar), _fill_axes(mar)) + +@inline _generic_productmeasure_impl(mar::Tuple{Vararg{AbstractMeasure}}) = + ProductMeasure(mar) +_generic_productmeasure_impl(mar::Tuple{Vararg{Dirac}}) = Dirac(map(m -> m.x, 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{Vararg{Dirac}}}) where {names} = + Dirac(map(m -> m.x, mar)) +_generic_productmeasure_impl(mar::NamedTuple) = productmeasure(map(asmeasure, mar)) + +@inline _generic_productmeasure_impl(mar::AbstractArray{<:AbstractProductMeasure}) = + ProductMeasure(mar) + +_generic_productmeasure_impl(mar::AbstractArray{<:Dirac}) = Dirac((m -> m.x).(mar)) + +# TODO: We should be able to further optimize this +function _generic_productmeasure_impl(mar::AbstractArray{T}) where {T} + if Base.issingletontype(T) + first(mar)^size(mar) + else + ProductMeasure(asmeasure.(mar)) + end +end -function productmeasure(mar::ReadonlyMappedArray{T,N,A,Returns{M}}) where {T,N,A,M} +@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(mar::Base.Generator) = ProductMeasure(mar) -productmeasure(mar::AbstractArray) = ProductMeasure(mar) +@inline _generic_productmeasure_impl(mar::Base.Generator) = ProductMeasure(mar) # TODO: Make this static when its length is static -@inline function productmeasure( - mar::AbstractArray{WeightedMeasure{StaticFloat64{W},M}}, +@inline function _generic_productmeasure_impl( + mar::AbstractArray{<:WeightedMeasure{StaticFloat64{W},M}}, ) where {W,M} return weightedmeasure(W * length(mar), productmeasure(map(basemeasure, mar))) end -productmeasure(nt::NamedTuple) = ProductMeasure(nt) -productmeasure(tup::Tuple) = ProductMeasure(tup) +# ToDo: Remove or at least refactor this (ProductMeasure shouldn't take a kernel as its argument). -productmeasure(f, param_maps, pars) = ProductMeasure(kernel(f, param_maps), pars) +productmeasure(f, param_maps, pars) = productmeasure(kernel(f, param_maps), pars) function productmeasure(k::ParameterizedTransitionKernel, pars) productmeasure(k.suff, k.param_maps, pars) From fe1d136b245a3fbdd6d2d41dfcaa1aaf65209e9a Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 13:46:21 +0200 Subject: [PATCH 42/75] Generalized transport for products and unknown-DOF measures MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add the generalized product/power transport machinery: transport between arbitrary measures and powers of standard measures now goes through `transport_to_mvstd` and `transport_from_mvstd_with_rest`. The with-rest protocol consumes only as much of the flat standard variate as each component needs, so products (and later monadic binds) whose components have no efficiently computable DOF become transportable via their transport origins. Multi-dimensional power measures now flatten to one-dimensional powers as their transport origin, one-dimensional powers pull back to powers of their parent's origin, products over `Fill` pull back to power measures and NamedTuple-products to Tuple-products. Standard-measure type-based transport partner selection (`transport_to(StdNormal, μ)`) now uses `some_dof` and lives with the new machinery, replacing the DOF-slicing implementation in stdmeasure.jl. Also adds `NoMSpaceElementSize`/`mspace_elsize`/ `some_mspace_elsize`, equality for `AsMeasure` wrappers, and makes mixed static/dynamic vector views AD-friendly. Ports the original measure-algebra commits 09b12dc, ab140a1 and related STASH commits, with fixes: the power-measure and product-measure transport origins gained the missing `to_origin` implementations, and the mvstd gateway methods route through the origin machinery to avoid infinite dispatch recursion. Co-Authored-By: Claude Fable 5 --- src/MeasureBase.jl | 5 + src/collection_utils.jl | 2 +- src/combinators/product.jl | 6 +- src/combinators/product_transport.jl | 409 +++++++++++++++++++++++++++ src/combinators/reshape.jl | 15 +- src/mspace.jl | 46 +++ src/standard/stdmeasure.jl | 139 --------- 7 files changed, 465 insertions(+), 157 deletions(-) create mode 100644 src/combinators/product_transport.jl create mode 100644 src/mspace.jl diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index 2aa92089..67307979 100644 --- a/src/MeasureBase.jl +++ b/src/MeasureBase.jl @@ -116,6 +116,9 @@ struct AsMeasure{T} <: AbstractMeasure AsMeasure{T}(obj::T) where {T} = new(obj) end +Base.:(==)(a::AsMeasure, b::AsMeasure) = a.obj == b.obj +Base.isapprox(a::AsMeasure, b::AsMeasure; kwargs...) = isapprox(a.obj, b.obj; kwargs...) + function Pretty.quoteof(d::M) where {M<:AbstractMeasure} the_names = fieldnames(typeof(d)) :($M($([getfield(d, n) for n in the_names]...))) @@ -167,6 +170,7 @@ using IrrationalConstants: loghalf include("collection_utils.jl") include("smf.jl") +include("mspace.jl") include("getdof.jl") include("transport.jl") include("schema.jl") @@ -205,6 +209,7 @@ include("standard/stduniform.jl") include("standard/stdexponential.jl") include("standard/stdlogistic.jl") include("standard/stdnormal.jl") +include("combinators/product_transport.jl") include("combinators/half.jl") #include("implicitmaps.jl") diff --git a/src/collection_utils.jl b/src/collection_utils.jl index 1ef9f9fd..d15131d0 100644 --- a/src/collection_utils.jl +++ b/src/collection_utils.jl @@ -27,7 +27,7 @@ Base.@propagate_inbounds _as_tuple(v::AbstractVector, ::Val{N}) where {N} = Tupl Base.@propagate_inbounds function _get_or_view(A::AbstractVector, from::IntegerLike, until::IntegerLike) - view(A, from:until) + view(A, dynamic(from):dynamic(until)) end Base.@propagate_inbounds function _get_or_view( diff --git a/src/combinators/product.jl b/src/combinators/product.jl index 468c38c0..656ded0d 100644 --- a/src/combinators/product.jl +++ b/src/combinators/product.jl @@ -28,9 +28,6 @@ Base.size(μ::AbstractProductMeasure) = size(marginals(μ)) basemeasure(d::AbstractProductMeasure) = productmeasure(map(basemeasure, marginals(d))) -proxy(μ::ProductMeasure{<:FillArrays.Fill}) = - powermeasure(_fill_value(marginals(μ)), _fill_axes(marginals(μ))) - function Base.rand(rng::AbstractRNG, ::Type{T}, d::AbstractProductMeasure) where {T} mar = marginals(d) _rand_product(rng, T, mar, eltype(mar)) @@ -85,6 +82,9 @@ struct ProductMeasure{M} <: AbstractProductMeasure marginals::M end +proxy(μ::ProductMeasure{<:FillArrays.Fill}) = + powermeasure(_fill_value(marginals(μ)), _fill_axes(marginals(μ))) + @inline function logdensity_rel(μ::ProductMeasure, ν::ProductMeasure, x) mapreduce(logdensity_rel, +, marginals(μ), marginals(ν), x) end diff --git a/src/combinators/product_transport.jl b/src/combinators/product_transport.jl new file mode 100644 index 00000000..a2b44b1b --- /dev/null +++ b/src/combinators/product_transport.jl @@ -0,0 +1,409 @@ +""" + transport_to(ν, ::Type{MU}) where {MU<:StdMeasure} + transport_to(::Type{NU}, μ) where {NU<:StdMeasure} + +As a user convenience, a standard measure type like [`StdUniform`](@ref), +[`StdExponential`](@ref), [`StdNormal`](@ref) or [`StdLogistic`](@ref) +may be used directly as the source or target of a measure transport. + +Depending on [`MeasureBase.some_dof(μ)`](@ref) (resp. `ν`), an instance of +the standard measure itself or a power of it will be automatically chosen as +the transport partner. + +Example: + +```julia +transport_to(StdNormal, μ) +transport_to(ν, StdNormal) +``` +""" +function transport_to(ν, ::Type{MU}) where {MU<:StdMeasure} + transport_to(ν, _std_tp_partner(MU, ν)) +end + +function transport_to(::Type{NU}, μ) where {NU<:StdMeasure} + transport_to(_std_tp_partner(NU, μ), μ) +end + +function transport_to(::Type{NU}, ::Type{MU}) where {NU<:StdMeasure,MU<:StdMeasure} + throw( + ArgumentError( + "Can't construct a transport function between the types of two standard measures, need a measure instance on one side", + ), + ) +end + +_std_tp_partner(::Type{M}, μ) where {M<:StdMeasure} = _std_tp_partner_bydof(M, some_dof(μ)) +_std_tp_partner_bydof(::Type{M}, ::StaticInteger{1}) where {M<:StdMeasure} = M() +_std_tp_partner_bydof(::Type{M}, dof::IntegerLike) where {M<:StdMeasure} = M()^dof +function _std_tp_partner_bydof(::Type{M}, ::AbstractNoDOF{MU}) where {M<:StdMeasure,MU} + throw( + ArgumentError( + "Can't determine a standard transport partner for measures of type $(nameof(MU))", + ), + ) +end + + +# For transport, always pull a multi-dimensional PowerMeasure back to a +# one-dimensional PowerMeasure first: + +const _PowerMeasureRank1{M} = PowerMeasure{M,<:NTuple{1,OneToLike}} + +function transport_origin(μ::PowerMeasure) + pwr_base(μ)^prod(pwr_size(μ)) +end + +function to_origin(μ::PowerMeasure, x) + maybestatic_reshape(x, (prod(pwr_size(μ)),)) +end + +function from_origin(μ::PowerMeasure, x_origin) + # Sanity check, should never fail: + @assert x_origin isa AbstractVector + maybestatic_reshape(x_origin, pwr_size(μ)) +end + + +# A one-dimensional PowerMeasure has an origin if its parent has an origin: + +function transport_origin(μ::_PowerMeasureRank1) + _pwr_origin(typeof(μ), transport_origin(pwr_base(μ)), pwr_axes(μ)) +end +_pwr_origin(::Type{MU}, parent_origin, axes) where {MU} = parent_origin^axes +_pwr_origin(::Type{MU}, ::NoTransportOrigin, axes) where {MU} = NoTransportOrigin{MU}() + +function to_origin(μ::_PowerMeasureRank1, x) + to_origin.(Ref(pwr_base(μ)), x) +end + +function from_origin(μ::_PowerMeasureRank1, x_origin) + # Sanity check, should never fail: + @assert x_origin isa AbstractVector + from_origin.(Ref(pwr_base(μ)), x_origin) +end + + +# Transport between powers of standard measures, of any rank: + +function _stdpow_transport_def(ν::StdPowerMeasure{NU}, μ::StdPowerMeasure{MU}, x) where {NU,MU} + y = transport_to(pwr_base(ν), pwr_base(μ)).(x) + maybestatic_reshape(y, pwr_size(ν)) +end + +function _stdpow_transport_def(ν::StdPowerMeasure{MU}, μ::StdPowerMeasure{MU}, x) where {MU} + maybestatic_reshape(x, pwr_size(ν)) +end + +transport_def(ν::StdPowerMeasure{NU}, μ::StdPowerMeasure{MU}, x) where {NU,MU} = + _stdpow_transport_def(ν, μ, x) + +# Disambiguation with the mvstd gateway methods below: +transport_def(ν::StdPowerMeasure{NU,1}, μ::StdPowerMeasure{MU}, x) where {NU,MU} = + _stdpow_transport_def(ν, μ, x) +transport_def(ν::StdPowerMeasure{NU}, μ::StdPowerMeasure{MU,1}, x) where {NU,MU} = + _stdpow_transport_def(ν, μ, x) +transport_def(ν::StdPowerMeasure{NU,1}, μ::StdPowerMeasure{MU,1}, x) where {NU,MU} = + _stdpow_transport_def(ν, μ, x) + + +# Transport between univariate standard measures and one-dimensional power +# measures of size one: + +function transport_def(ν::StdMeasure, μ::StdPowerMeasure{MU,1}, x) where {MU} + return transport_def(ν, pwr_base(μ), only(x)) +end + +function transport_def(ν::StdPowerMeasure{NU,1}, μ::StdMeasure, x) where {NU} + sz_ν = pwr_size(ν) + @assert prod(sz_ν) == 1 + return maybestatic_fill(transport_def(pwr_base(ν), μ, x), sz_ν) +end + + +# Transport to a multivariate standard measure from any measure: + +function transport_def(ν::StdPowerMeasure{NU,1}, μ::AbstractMeasure, x) where {NU} + transport_to_mvstd(pwr_base(ν), μ, x) +end + +function transport_to_mvstd(ν_inner::StdMeasure, μ::AbstractMeasure, x) + return _to_mvstd_withdof(ν_inner, μ, fast_dof(μ), x) +end + +# For standard measures and their powers specialized `transport_def` methods +# exist, for other measures the origin-based machinery must be used directly +# instead of `transport_def`, to prevent infinite dispatch recursion via the +# gateway methods above: +const _StdOrStdPowerMeasure = Union{StdMeasure,StdPowerMeasure} + +_transport_def_nongateway(ν::_StdOrStdPowerMeasure, μ::_StdOrStdPowerMeasure, x) = + transport_def(ν, μ, x) +function _transport_def_nongateway(ν, μ, x) + _transport_between_origins(ν, _origin_depth(ν), _origin_depth(μ), μ, x) +end + +function _to_mvstd_withdof(ν_inner::StdMeasure, μ::AbstractMeasure, dof_μ::IntegerLike, x) + _transport_def_nongateway(ν_inner^dof_μ, μ, x) +end + +function _to_mvstd_withdof(ν_inner::StdMeasure, μ::AbstractMeasure, ::AbstractNoDOF, x) + _to_mvstd_withorigin(ν_inner, μ, transport_origin(μ), x) +end + +function _to_mvstd_withorigin(ν_inner::StdMeasure, μ::AbstractMeasure, μ_origin, x) + x_origin = to_origin(μ, x) + transport_to_mvstd(ν_inner, μ_origin, x_origin) +end + +function _to_mvstd_withorigin(ν_inner::StdMeasure, μ::AbstractMeasure, ::NoTransportOrigin, x) + throw( + ArgumentError( + "Don't know how to transport values of type $(nameof(typeof(x))) from $(nameof(typeof(μ))) to a power of $(nameof(typeof(ν_inner)))", + ), + ) +end + + +# Transport from a multivariate standard measure to any measure: + +function transport_def(ν::AbstractMeasure, μ::StdPowerMeasure{MU,1}, x) where {MU} + _transport_from_mvstd(ν, pwr_base(μ), x) +end + +function _transport_from_mvstd(ν::AbstractMeasure, μ_inner::StdMeasure, x) + y, x_rest = transport_from_mvstd_with_rest(ν, μ_inner, x) + if !isempty(x_rest) + throw(ArgumentError("Input value too long during transport")) + end + return y +end + +function transport_from_mvstd_with_rest(ν::AbstractMeasure, μ_inner::StdMeasure, x) + dof_ν = fast_dof(ν) + return _from_mvstd_with_rest_withdof(ν, dof_ν, μ_inner, x) +end + +function _from_mvstd_with_rest_withdof( + ν::AbstractMeasure, + dof_ν::IntegerLike, + μ_inner::StdMeasure, + x, +) + len_x = maybestatic_length(x) + + # Since we can't check the DOF of measures like Bind, we could "run out + # of x" if the original x was too short. `transport_to` below will detect + # this, but better throw a more informative exception here: + if len_x < dof_ν + throw(ArgumentError("Variate too short during transport")) + end + + x_inner_dof, x_rest = _split_after(x, dof_ν) + y = _transport_def_nongateway(ν, μ_inner^dof_ν, x_inner_dof) + return y, x_rest +end + +function _from_mvstd_with_rest_withdof( + ν::AbstractMeasure, + ::AbstractNoDOF, + μ_inner::StdMeasure, + x, +) + _from_mvstd_with_rest_withorigin(ν, transport_origin(ν), μ_inner, x) +end + +function _from_mvstd_with_rest_withorigin( + ν::AbstractMeasure, + ν_origin, + μ_inner::StdMeasure, + x, +) + x_origin, x_rest = transport_from_mvstd_with_rest(ν_origin, μ_inner, x) + from_origin(ν, x_origin), x_rest +end + +function _from_mvstd_with_rest_withorigin( + ν::AbstractMeasure, + ::NoTransportOrigin, + μ_inner::StdMeasure, + x, +) + throw( + ArgumentError( + "Don't know how to transport a value of type $(nameof(typeof(x))) from a power of $(nameof(typeof(μ_inner))) to $(nameof(typeof(ν)))", + ), + ) +end + + +# Transport between a standard measure and Dirac: + +@inline transport_from_mvstd_with_rest(ν::Dirac, ::StdMeasure, x::Any) = ν.x, x + +@inline transport_to_mvstd(::StdMeasure, ::Dirac, ::Any) = FillArrays.Zeros{Bool}(0) + + +# Pull back from a product over a Fill to a power measure: + +@inline transport_origin(μ::ProductMeasure) = _marginals_tp_origin(marginals(μ)) +@inline to_origin(μ::ProductMeasure, x) = _marginals_to_origin(marginals(μ), x) +@inline from_origin(μ::ProductMeasure, x_origin) = + _marginals_from_origin(marginals(μ), x_origin) + +_marginals_tp_origin(::Ms) where {Ms} = NoTransportOrigin{ProductMeasure{Ms}}() + +_marginals_tp_origin(marginals_μ::FillArrays.Fill) = + _fill_value(marginals_μ)^_fill_axes(marginals_μ) +_marginals_to_origin(::FillArrays.Fill, x) = x +_marginals_from_origin(::FillArrays.Fill, x_origin) = x_origin + + +# Pull back from a NamedTuple product measure to a Tuple product measure: +# +# Maybe ToDo (breaking): For transport between NamedTuple-marginals we could +# match names where possible, even if given in different order, and transport +# between the remaining non-matching names in the order given. This may not +# be worth the additional complexity, though, since transport is typically +# used with a (power of a) standard measure on one side. + +_marginals_tp_origin(marginals_μ::NamedTuple{names}) where {names} = + productmeasure(values(marginals_μ)) +_marginals_to_origin(::NamedTuple{names}, x::NamedTuple{names}) where {names} = values(x) +_marginals_from_origin(::NamedTuple{names}, x_origin::Tuple) where {names} = + NamedTuple{names}(x_origin) + + +# Transport between two instances of ProductMeasure: + +transport_def(ν::ProductMeasure, μ::ProductMeasure, x) = + _marginal_transport_def(marginals(ν), marginals(μ), x) + +function _marginal_transport_def(marginals_ν, marginals_μ, x) + @assert size(marginals_ν) == size(marginals_μ) == size(x) # Sanity check, should not fail + transport_def.(marginals_ν, marginals_μ, x) +end + +function _marginal_transport_def( + marginals_ν::Tuple{Vararg{AbstractMeasure,N}}, + marginals_μ::Tuple{Vararg{AbstractMeasure,N}}, + x, +) where {N} + map(transport_def, marginals_ν, marginals_μ, x) +end + +function _marginal_transport_def( + marginals_ν::AbstractVector{<:AbstractMeasure}, + marginals_μ::Tuple{Vararg{AbstractMeasure,N}}, + x, +) where {N} + _marginal_transport_def(_as_tuple(marginals_ν, Val(N)), marginals_μ, x) +end + +function _marginal_transport_def( + marginals_ν::Tuple{Vararg{AbstractMeasure,N}}, + marginals_μ::AbstractVector{<:AbstractMeasure}, + x, +) where {N} + _marginal_transport_def(marginals_ν, _as_tuple(marginals_μ, Val(N)), _as_tuple(x, Val(N))) +end + + +# Transport from a ProductMeasure to a standard measure: + +function transport_to_mvstd(ν_inner::StdMeasure, μ::ProductMeasure, x) + _marginals_to_mvstd(ν_inner, marginals(μ), x) +end + +struct _TransportToMvStd{NU<:StdMeasure} <: Function end +(::_TransportToMvStd{NU})(μ, x) where {NU} = transport_to_mvstd(NU(), μ, x) + +function _marginals_to_mvstd(::NU, marginals_μ::Tuple, x::Tuple) where {NU<:StdMeasure} + _flatten_to_rv(map(_TransportToMvStd{NU}(), marginals_μ, x)) +end + +function _marginals_to_mvstd(::NU, marginals_μ, x) where {NU<:StdMeasure} + _flatten_to_rv(broadcast(_TransportToMvStd{NU}(), marginals_μ, x)) +end + + +# Transport from a standard measure to a ProductMeasure, with rest: + +const _MaybeUnknownDOF = Union{IntegerLike,AbstractNoDOF} + +const _KnownDOFs = Union{Tuple{Vararg{IntegerLike,N}} where N,StaticVector{<:IntegerLike}} + +function transport_from_mvstd_with_rest(ν::ProductMeasure, μ_inner::StdMeasure, x) + νs = marginals(ν) + dofs = map(fast_dof, νs) + return _marginals_from_mvstd_with_rest(νs, dofs, μ_inner, x) +end + +function _dof_access_firstidxs(dofs::Tuple{Vararg{IntegerLike,N}}, first_idx) where {N} + cumsum((first_idx, dofs[begin:(end-1)]...)) +end + +function _dof_access_firstidxs(dofs::AbstractVector{<:IntegerLike}, first_idx) + # ToDo: Improve implementation (reduce memory allocations): + cumsum(vcat([eltype(dofs)(first_idx)], dofs[begin:(end-1)])) +end + +function _split_x_by_marginals_with_rest( + dofs::Union{Tuple,AbstractVector}, + x::AbstractVector{<:Real}, +) + x_idxs = maybestatic_eachindex(x) + first_idxs = _dof_access_firstidxs(dofs, maybestatic_first(x_idxs)) + xs = map((from, n) -> _get_or_view(x, from, from + n - one(n)), first_idxs, dofs) + x_rest = _get_or_view(x, first_idxs[end] + dofs[end], maybestatic_last(x_idxs)) + return xs, x_rest +end + +function _marginals_from_mvstd_with_rest( + νs, + dofs::_KnownDOFs, + μ_inner::StdMeasure, + x::AbstractVector{<:Real}, +) + xs, x_rest = _split_x_by_marginals_with_rest(dofs, x) + μs = map(n -> μ_inner^n, dofs) + ys = map(transport_def, νs, μs, xs) + return ys, x_rest +end + +function _marginals_from_mvstd_with_rest( + νs, + dofs, + μ_inner::StdMeasure, + x::AbstractVector{<:Real}, +) + _marginals_from_mvstd_with_rest_nodof(νs, μ_inner, x) +end + +function _marginals_from_mvstd_with_rest_nodof( + νs::Tuple{Vararg{AbstractMeasure}}, + μ_inner::StdMeasure, + x::AbstractVector{<:Real}, +) + # ToDo: Check for type stability, may need a generated function: + y1, x_rest = transport_from_mvstd_with_rest(νs[1], μ_inner, x) + y2_end, x_final_rest = _marginals_from_mvstd_with_rest_nodof(Base.tail(νs), μ_inner, x_rest) + return (y1, y2_end...), x_final_rest +end + +_marginals_from_mvstd_with_rest_nodof(::Tuple{}, ::StdMeasure, x::AbstractVector{<:Real}) = + (), x + +function _marginals_from_mvstd_with_rest_nodof( + νs::AbstractVector{<:AbstractMeasure}, + μ_inner::StdMeasure, + x::AbstractVector{<:Real}, +) + # ToDo: Check for type stability: + ys = Vector{Any}(undef, length(eachindex(νs))) + x_rest = x + for (i, ν) in zip(eachindex(ys), νs) + ys[i], x_rest = transport_from_mvstd_with_rest(ν, μ_inner, x_rest) + end + return [y for y in ys], x_rest +end diff --git a/src/combinators/reshape.jl b/src/combinators/reshape.jl index dddae55b..f77dbb68 100644 --- a/src/combinators/reshape.jl +++ b/src/combinators/reshape.jl @@ -53,17 +53,4 @@ a space of arrays with shape `sz`. function mreshape end mreshape(m::AbstractMeasure, sz::IntegerLike...) = mreshape(m, sz) -mreshape(m::AbstractMeasure, sz::SizeLike) = pushfwd(Reshape(sz, mspace_elsize(m)), m) - - -""" - MeasureBase.mspace_elsize(m::AbstractMeasure)::MeasureBase.SizeLike - -Return the size of the elements of the measurable space of `m`. - -Defaults to the size of a test value of `m`, may be specialized for -measure types where this is inefficient. -""" -function mspace_elsize end - -mspace_elsize(m::AbstractMeasure) = maybestatic_size(testvalue(m)) +mreshape(m::AbstractMeasure, sz::SizeLike) = pushfwd(Reshape(sz, some_mspace_elsize(m)), m) diff --git a/src/mspace.jl b/src/mspace.jl new file mode 100644 index 00000000..6227537d --- /dev/null +++ b/src/mspace.jl @@ -0,0 +1,46 @@ +""" + MeasureBase.NoMSpaceElementSize{MU} + +Indicates that either the measurable space of measures of type `MU` is not +a space over arrays, or that the size of the arrays is not fixed or can not +be easily/efficiently determined. +""" +struct NoMSpaceElementSize{MU} end + + +""" + mspace_elsize(μ) + +For a measure `μ` over an array-valued measurable space, return the size of +the arrays that are the elements of the space. + +May return [`NoMSpaceElementSize{typeof(μ)}()`](@ref). +""" +function mspace_elsize end +export mspace_elsize + +@inline mspace_elsize(μ::AbstractMeasure) = NoMSpaceElementSize{typeof(μ)}() + + +""" + MeasureBase.some_mspace_elsize(μ::AbstractMeasure) + +For a measure `μ` over an array-valued measurable space, return the size of +an arbitrary element of the space. + +Use with caution, the space of some measures is made up of arrays of +different sizes! + +In general, use [`mspace_elsize(μ)`](@ref) instead. `some_mspace_elsize` is +useful if the measurable space is expected to contain only arrays of the +same size but there is no way to prove this automatically. Algorithms that +use the returned size should always check that it matches the size of each +point of the space that is processed. +""" +function some_mspace_elsize end + +@inline some_mspace_elsize(μ) = _mspace_some_elsize_impl(μ, mspace_elsize(μ)) + +@inline _mspace_some_elsize_impl(::AbstractMeasure, sz::SizeLike) = sz +_mspace_some_elsize_impl(μ::AbstractMeasure, ::NoMSpaceElementSize) = + maybestatic_size(testvalue(μ)) diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl index 3dd30e24..81409796 100644 --- a/src/standard/stdmeasure.jl +++ b/src/standard/stdmeasure.jl @@ -14,142 +14,3 @@ const StdPowerMeasure{MU<:StdMeasure,N} = PowerMeasure{MU,<:NTuple{N,OneToLike}} @inline check_dof(::StdMeasure, ::StdMeasure) = nothing @inline transport_def(::MU, μ::MU, x) where {MU<:StdMeasure} = x - -function transport_def(ν::StdMeasure, μ::PowerMeasure{<:StdMeasure}, x) - return transport_def(ν, μ.parent, only(x)) -end - -function transport_def(ν::PowerMeasure{<:StdMeasure}, μ::StdMeasure, x) - return maybestatic_fill(transport_def(ν.parent, μ, only(x)), map(length, ν.axes)) -end - -function transport_def( - ν::PowerMeasure{<:StdMeasure,<:NTuple{1,Base.OneTo}}, - μ::PowerMeasure{<:StdMeasure,<:NTuple{1,Base.OneTo}}, - x, -) - return transport_to(ν.parent, μ.parent).(x) -end - -function transport_def( - ν::PowerMeasure{<:StdMeasure,<:NTuple{N,Base.OneTo}}, - μ::PowerMeasure{<:StdMeasure,<:NTuple{M,Base.OneTo}}, - x, -) where {N,M} - return reshape(transport_to(ν.parent, μ.parent).(x), map(length, ν.axes)...) -end - -# Implement transport_to(NU::Type{<:StdMeasure}, μ) and transport_to(ν, MU::Type{<:StdMeasure}): - -_std_measure(::Type{M}, ::StaticInteger{1}) where {M<:StdMeasure} = M() -_std_measure(::Type{M}, dof::IntegerLike) where {M<:StdMeasure} = M()^dof -_std_measure_for(::Type{M}, μ::Any) where {M<:StdMeasure} = _std_measure(M, getdof(μ)) - -function transport_to(::Type{NU}, μ) where {NU<:StdMeasure} - transport_to(_std_measure_for(NU, μ), μ) -end - -function transport_to(ν, ::Type{MU}) where {MU<:StdMeasure} - transport_to(ν, _std_measure_for(MU, ν)) -end - -# Transform between standard measures and Dirac: - -@inline transport_def(ν::Dirac, ::PowerMeasure{<:StdMeasure}, ::Any) = ν.x - -@inline function transport_def(ν::PowerMeasure{<:StdMeasure}, ::Dirac, ::Any) - Zeros{Bool}(map(_ -> 0, ν.axes)) -end - -# Helpers for product transforms and similar: - -struct _TransportToStd{NU<:StdMeasure} <: Function end -(::_TransportToStd{NU})(μ, x) where {NU} = transport_to(NU()^getdof(μ), μ)(x) - -struct _TransportFromStd{MU<:StdMeasure} <: Function end -_TransportFromStd{MU}(ν, x) where {MU} = transport_to(ν, MU()^getdof(ν))(x) - -function _tuple_transport_def( - ν::PowerMeasure{NU}, - μs::Tuple, - xs::Tuple, -) where {NU<:StdMeasure} - reshape(vcat(map(_TransportToStd{NU}(), μs, xs)...), ν.axes) -end - -function transport_def( - ν::PowerMeasure{NU}, - μ::ProductMeasure{<:Tuple}, - x, -) where {NU<:StdMeasure} - _tuple_transport_def(ν, marginals(μ), x) -end - -function transport_def( - ν::PowerMeasure{NU}, - μ::ProductMeasure{<:NamedTuple{names}}, - x, -) where {NU<:StdMeasure,names} - _tuple_transport_def(ν, values(marginals(μ)), values(x)) -end - -@inline _offset_cumsum(s, x, y, rest...) = (s, _offset_cumsum(s + x, y, rest...)...) -@inline _offset_cumsum(s, x) = (s,) -@inline _offset_cumsum(s) = () - -function _stdvar_viewranges(μs::Tuple, startidx::IntegerLike) - N = map(getdof, μs) - offs = _offset_cumsum(startidx, N...) - map((o, n) -> o:(o+n-1), offs, N) -end - -function _tuple_transport_def( - νs::Tuple, - μ::PowerMeasure{MU}, - x::AbstractArray{<:Real}, -) where {MU<:StdMeasure} - vrs = _stdvar_viewranges(νs, firstindex(x)) - xs = map(r -> view(x, r), vrs) - map(_TransportFromStd{MU}, νs, xs) -end - -function transport_def( - ν::ProductMeasure{<:Tuple}, - μ::PowerMeasure{MU}, - x, -) where {MU<:StdMeasure} - _tuple_transport_def(marginals(ν), μ, x) -end - -function transport_def( - ν::ProductMeasure{<:NamedTuple{names}}, - μ::PowerMeasure{MU}, - x, -) where {MU<:StdMeasure,names} - NamedTuple{names}(_tuple_transport_def(values(marginals(ν)), μ, x)) -end - -function transport_def( - ν::PowerMeasure{NU}, - μ::ProductMeasure{<:AbstractArray}, - x, -) where {NU<:StdMeasure} - reshape(vcat(map(_TransportToStd{NU}(), marginals(μ), x)...), ν.axes) -end - -function _marginal_viewranges(μs::AbstractArray, startidx::IntegerLike) - ns = map(m -> dynamic(getdof(m)), μs) - offs = cumsum(vcat(dynamic(startidx), ns[begin:(end-1)])) - map((o, n) -> o:(o+n-1), offs, ns) -end - -function transport_def( - ν::ProductMeasure{<:AbstractArray}, - μ::PowerMeasure{MU}, - x::AbstractArray{<:Real}, -) where {MU<:StdMeasure} - νs = marginals(ν) - vrs = _marginal_viewranges(νs, firstindex(x)) - xs = map(r -> view(x, r), vrs) - map(_TransportFromStd{MU}, νs, xs) -end From 8c3d680a2434835b875f7817b411af27581890a1 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 15:19:40 +0200 Subject: [PATCH 43/75] Rework mbind as combined-value bind, add mcombine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rework the monadic bind into a full hierarchical measure: `mbind(f_β, α, f_c)` combines points of the primary measure `α` and the dependent secondary measure `f_β(a)` via a value combination function `f_c` (`tuple`, `Pair`, `vcat` and `merge` are supported for density evaluation, `OneTwoMany.secondarg` gives the plain monadic bind). `mkernel(f_β, f_c)` represents the generalized transition kernel, `bindkernel` and `boundmeasure` recover the parts of a bind. Density evaluation goes through the new `transportmeasure`/ `tpmeasure_split_combined` mechanism, and transport to and from powers of standard measures uses the with-rest protocol, so hierarchical measures without a fast DOF are transportable. Add `mcombine` and `CombinedMeasure` for combining two independent measures via `f_c`, with product shortcuts for `tuple`, `vcat` and `merge`. MeasureBase now depends on OneTwoMany for `firstarg`/`secondarg`. Ports the original measure-algebra commits fe28297, 03724b7, 17e30ec, cce5b41, 8561e26, 4b3ed53 and related STASH commits, with fixes: `boundmeasure` returns the bound measure instead of the kernel, the tuple/Pair variate splitters no longer reference undefined variables, kernel results are converted via `asmeasure`, and NamedTuple products gained the missing to/from-mvstd transport paths. Co-Authored-By: Claude Fable 5 --- Project.toml | 2 + src/MeasureBase.jl | 5 +- src/combinators/bind.jl | 336 +++++++++++++++++++++++++-- src/combinators/combined.jl | 174 ++++++++++++++ src/combinators/product_transport.jl | 18 ++ test/Project.toml | 1 + test/combinators/bind.jl | 94 ++++++++ test/combinators/combined.jl | 52 +++++ test/runtests.jl | 2 + 9 files changed, 658 insertions(+), 26 deletions(-) create mode 100644 src/combinators/combined.jl create mode 100644 test/combinators/bind.jl create mode 100644 test/combinators/combined.jl diff --git a/Project.toml b/Project.toml index e2bfee4a..c0ba40c8 100644 --- a/Project.toml +++ b/Project.toml @@ -23,6 +23,7 @@ LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" LogarithmicNumbers = "aa2f6b4e-9042-5d33-9679-40d3a6b85899" MappedArrays = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900" NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3" +OneTwoMany = "762dc654-8631-413a-a342-372a7419ad9d" PrettyPrinting = "54e16d92-306c-5ea0-a30b-337be88ac337" PropertyFunctions = "09e99361-2bb8-48a2-a80f-de58f0739eb4" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" @@ -81,6 +82,7 @@ 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" diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index 67307979..9f8037eb 100644 --- a/src/MeasureBase.jl +++ b/src/MeasureBase.jl @@ -63,6 +63,8 @@ using HeterogeneousComputing: real_numtype using ArraysOfArrays: VectorOfArrays, VectorOfVectors, VectorOfSimilarArrays, VectorOfSimilarVectors, flatview +using OneTwoMany: firstarg, secondarg + export gentype export AbstractMeasure @@ -190,7 +192,6 @@ 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") @@ -210,6 +211,8 @@ include("standard/stdexponential.jl") include("standard/stdlogistic.jl") include("standard/stdnormal.jl") include("combinators/product_transport.jl") +include("combinators/combined.jl") +include("combinators/bind.jl") include("combinators/half.jl") #include("implicitmaps.jl") diff --git a/src/combinators/bind.jl b/src/combinators/bind.jl index 60b7cfd8..27cb02a4 100644 --- a/src/combinators/bind.jl +++ b/src/combinators/bind.jl @@ -1,43 +1,329 @@ +@doc raw""" + mkernel(f_β, f_c = OneTwoMany.secondarg)::Function + +Constructs a generalized monadic transition kernel from a primary transition +kernel function `f_β` and a value combination function `f_c`. + +`f_β` must behave like `β = f_β(a)`, taking a value `a` from a primary +measurable space and returning a measure-like object `β`. + +`f_c` must behave like `c = f_c(a, b)`, taking a value `a` (like `f_β`) and +a value `b` from the measurable space of `β` and returning a value `c`. + +`f_k = mkernel(f_β, f_c)` then acts like + +```julia +f_k(a) ≡ pushfwd(c -> f_c(c[1], c[2]), productmeasure((Dirac(a), f_β(a)))) +``` + +(`≡` denoting pseudocode-equivalency here). So with the default +`f_c == OneTwoMany.secondarg`, we just have `f_k(a) ≡ f_β(a)`. + +Also, + +```julia +mbind(mkernel(f_β, f_c), α) == mbind(f_β, α, f_c) +``` + +See also [`mbind`](@ref). """ - struct MeasureBase.Bind{M,K} <: AbstractMeasure +function mkernel end +export mkernel + -Represents a monatic bind. User code should not create instances of `Bind` -directly, but should call `mbind(k, μ)` instead. """ -struct Bind{M,K} <: AbstractMeasure - k::K - μ::M -end + struct MeasureBase.MKernel <: Function -getdof(d::Bind) = NoDOF{typeof(d)}() +Represents a generalized monadic transition kernel. -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 +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)) -""" - mbind(k, μ)::AbstractMeasure +@inline mkernel(f_β::MKernel) = f_β +@inline mkernel(f_β, f_c = secondarg) = _generic_mkernel_impl(f_β, f_c) + +@inline _generic_mkernel_impl(f_β, f_c) = MKernel(f_β, f_c) +@inline _generic_mkernel_impl(f_β::MKernel, ::typeof(secondarg)) = f_β + + +@doc raw""" + mbind(f_β, α::AbstractMeasure, f_c = OneTwoMany.secondarg) + mbind(f_β::MeasureBase.MKernel, α::AbstractMeasure) + +Constructs a monadic bind, resp. a hierarchical measure, from a transition +kernel function `f_β`, a primary measure `α` and a value combination +function `f_c`. + +`f_β` must be a function that maps a point `a` from the space of the primary +measure `α` to a dependent secondary measure `β_a = f_β(a)`. +`ab = f_c(a, b)` must map such a point `a` and a point `b` from the +space of measure `β_a` to a combined value `ab = f_c(a, b)`. + +The resulting measure -Given +```julia +μ = mbind(f_β, α, f_c) +``` -- a measure μ -- a kernel function k that takes values from the support of μ and returns a - measure +has the mathematical interpretation (on sets $$A$$ and $$B$$) -The *monadic bind* operation `mbind(k, μ)` returns is a new measure. +```math +\mu(f_c(A, B)) = \int_A \beta_a(B)\, \mathrm{d}\, \alpha(a) +``` -A monadic bind is often written as `>>=` (e.g. in Haskell), but this symbol is -unavailable in Julia. +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) ``` -μ = StdExponential() -ν = mbind(μ) do scale - pushfwd(Base.Fix1(*, scale), StdNormal()) + +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, θ) ``` """ -mbind(k, μ) = Bind(k, μ) +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? + + +""" + bindkernel(μ::Bind)::MKernel + +Returns the monadic transition kernel of a monadic bind, so that +`bindkernel(mbind(f_k::MKernel, α)) == f_k`. + +See [`mbind`](@ref) and [`mkernel`](@ref) for details. +""" +function bindkernel end +export bindkernel + +bindkernel(μ::Bind) = mkernel(μ.f_β, μ.f_c) + + +""" + boundmeasure(μ::Bind)::AbstractMeasure + +Returns the measure that went into a monadic bind, so that +`boundmeasure(mbind(f_k, α)) == α`. + +See [`mbind`](@ref) and [`mkernel`](@ref) for details. +""" +function boundmeasure end +export boundmeasure + +boundmeasure(μ::Bind) = μ.α + + +_get_β_a(μ::Bind, a) = asmeasure(μ.f_β(a)) + +function transportmeasure(μ::Bind, x) + tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, x) + tpm_β_a = transportmeasure(_get_β_a(μ, a), b) + mcombine(μ.f_c, tpm_α, tpm_β_a) +end + +localmeasure(μ::Bind, x) = transportmeasure(μ, x) + +tpmeasure_split_combined(f_c, μ::Bind, xy) = _bind_tpm_sc(f_c, μ, xy) + +function _bind_tpm_sc(::typeof(tuple), μ::Bind, xy::Tuple{Vararg{Any,2}}) + x, y = xy[1], xy[2] + tpm_μ = transportmeasure(μ, x) + return tpm_μ, x, y +end + +function _bind_tpm_sc(::Type{Pair}, μ::Bind, xy::Pair) + x, y = xy.first, xy.second + tpm_μ = transportmeasure(μ, x) + return tpm_μ, x, y +end + +const _BindBy{FC} = Bind{<:Any,<:AbstractMeasure,FC} +_bind_tpm_sc(f_c::typeof(vcat), μ::_BindBy{typeof(vcat)}, xy::AbstractVector) = + _bind_tpm_sc_cat(f_c, μ, xy) +_bind_tpm_sc(f_c::typeof(merge), μ::_BindBy{typeof(merge)}, xy::NamedTuple) = + _bind_tpm_sc_cat(f_c, μ, xy) + +function _bind_tpm_sc_cat_lμabyxy(f_c, μ, xy) + tpm_α, a, by = tpmeasure_split_combined(μ.f_c, μ.α, xy) + β_a = _get_β_a(μ, a) + tpm_β_a, b, y = tpmeasure_split_combined(f_c, β_a, by) + tpm_μ = mcombine(μ.f_c, tpm_α, tpm_β_a) + return tpm_μ, a, b, y, xy +end + +function _bind_tpm_sc_cat(f_c::typeof(vcat), μ::_BindBy{typeof(vcat)}, xy::AbstractVector) + tpm_μ, a, b, y, xy = _bind_tpm_sc_cat_lμabyxy(f_c, μ, xy) + # Don't use `x = f_c(a, b)` here, would allocate, splitting xy can use views: + x, y = _split_after(xy, length(a) + length(b)) + return tpm_μ, x, y +end + +function _bind_tpm_sc_cat(f_c::typeof(merge), μ::_BindBy{typeof(merge)}, xy::NamedTuple) + tpm_μ, a, b, y, xy = _bind_tpm_sc_cat_lμabyxy(f_c, μ, xy) + return tpm_μ, f_c(a, b), y +end + + +@inline insupport(μ::Bind, ::Any) = NoFastInsupport{typeof(μ)}() + +@inline getdof(μ::Bind) = NoDOF{typeof(μ)}() + +# Bypass `checked_arg`, would require potentially costly evaluation of f_β: +@inline checked_arg(::Bind, x) = x + +rootmeasure(::Bind) = + throw(ArgumentError("root measure is implicit, but can't be instantiated, for Bind")) + +basemeasure(::Bind) = throw(ArgumentError("basemeasure is not available for Bind")) + +testvalue(::Bind) = throw(ArgumentError("testvalue is not available for Bind")) + +logdensity_def(::Bind, x) = + throw(ArgumentError("logdensity_def is not available for Bind")) + +# Specialize logdensityof to avoid duplicate calculations: +function logdensityof(μ::Bind, x) + tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, x) + β_a = _get_β_a(μ, a) + logdensityof(tpm_α, a) + logdensityof(β_a, b) +end + +# Specialize unsafe_logdensityof to avoid duplicate calculations: +function unsafe_logdensityof(μ::Bind, x) + tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, x) + β_a = _get_β_a(μ, a) + unsafe_logdensityof(tpm_α, a) + unsafe_logdensityof(β_a, b) +end + + +function Base.rand(rng::Random.AbstractRNG, ::Type{T}, μ::Bind) where {T<:Real} + a = rand(rng, T, μ.α) + b = rand(rng, T, _get_β_a(μ, a)) + return μ.f_c(a, b) +end + +function Base.rand(rng::Random.AbstractRNG, μ::Bind) + a = rand(rng, μ.α) + b = rand(rng, _get_β_a(μ, a)) + return μ.f_c(a, b) +end + + +function transport_to_mvstd(ν_inner::StdMeasure, μ::Bind, ab) + tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, ab) + β_a = _get_β_a(μ, a) + y1 = transport_to_mvstd(ν_inner, tpm_α, a) + y2 = transport_to_mvstd(ν_inner, β_a, b) + return vcat(y1, y2) +end + + +function transport_from_mvstd_with_rest(ν::Bind, μ_inner::StdMeasure, x) + a, x2 = transport_from_mvstd_with_rest(ν.α, μ_inner, x) + β_a = _get_β_a(ν, a) + b, x_rest = transport_from_mvstd_with_rest(β_a, μ_inner, x2) + return ν.f_c(a, b), x_rest +end diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl new file mode 100644 index 00000000..c3ea079c --- /dev/null +++ b/src/combinators/combined.jl @@ -0,0 +1,174 @@ +""" + MeasureBase.tpmeasure_split_combined(f_c, α::AbstractMeasure, ab) + +Splits a combined value `ab` that originated from combining a point `a` +from the space of a measure `α` with a point `b` from the space of +another measure `β` via `ab = f_c(a, b)`. + +Returns a semantic equivalent of +`(MeasureBase.transportmeasure(α, a), a, b)`. + +With `a_orig = rand(α)`, `b_orig = rand(β)` and +`ab = f_c(a_orig, b_orig)`, the following must hold true: + +```julia +tpm_α, a, b = tpmeasure_split_combined(f_c, α, ab) +a ≈ a_orig && b ≈ b_orig +``` +""" +function tpmeasure_split_combined end + +function tpmeasure_split_combined(f_c, α::AbstractMeasure, ab) + a, b = _generic_split_combined(f_c, α, ab) + return transportmeasure(α, a), a, b +end + +@inline _generic_split_combined(::typeof(tuple), ::AbstractMeasure, x::Tuple{Vararg{Any,2}}) = x +@inline _generic_split_combined(::Type{Pair}, ::AbstractMeasure, ab::Pair) = (ab...,) + +function _generic_split_combined(f_c::FC, α::AbstractMeasure, ab) where {FC} + _split_variate_byvalue(f_c, testvalue(α), ab) +end + +_split_variate_byvalue(::typeof(vcat), test_a::AbstractVector, ab::AbstractVector) = + _split_after(ab, length(test_a)) + +_split_variate_byvalue(::typeof(vcat), ::NTuple{N,Any}, ab::Tuple) where {N} = + _split_after(ab, Val{N}()) + +function _split_variate_byvalue(::typeof(merge), ::NamedTuple{names_a}, ab::NamedTuple) where {names_a} + _split_after(ab, Val(names_a)) +end + + +@doc raw""" + mcombine(f_c, α::AbstractMeasure, β::AbstractMeasure) + +Combines two measures `α` and `β` to a combined measure via a point +combination function `f_c`. + +`f_c` must combine a given point `a` from the space of measure `α` with a +given point `b` from the space of measure `β` to a single value +`ab = f_c(a, b)` in the space of the combined measure +`μ = mcombine(f_c, α, β)`. + +The combined measure has the mathematical interpretation (on sets +$$A$$ and $$B$$) + +```math +\mu(f_c(A, B)) = \alpha(A)\, \beta(B) +``` +""" +function mcombine end +export mcombine + +@inline function mcombine(f_c, α::AbstractMeasure, β::AbstractMeasure) + _generic_mcombine_impl_stage1(f_c, α, β) +end + +@inline _generic_mcombine_impl_stage1(::typeof(firstarg), α::AbstractMeasure, β::AbstractMeasure) = α +@inline _generic_mcombine_impl_stage1(::typeof(secondarg), α::AbstractMeasure, β::AbstractMeasure) = β + +@inline function _generic_mcombine_impl_stage1(::typeof(tuple), α::AbstractMeasure, β::AbstractMeasure) + productmeasure((α, β)) +end + +@inline function _generic_mcombine_impl_stage1( + f_c::Union{typeof(vcat),typeof(merge)}, + α::AbstractProductMeasure, + β::AbstractProductMeasure, +) + _mcombine_product_shortcut(f_c, marginals(α), marginals(β), α, β) +end + +_mcombine_product_shortcut(::typeof(vcat), ma::AbstractVector, mb::AbstractVector, α, β) = + productmeasure(vcat(ma, mb)) +_mcombine_product_shortcut(::typeof(merge), ma::NamedTuple, mb::NamedTuple, α, β) = + productmeasure(merge(ma, mb)) +_mcombine_product_shortcut(f_c, ma, mb, α, β) = _generic_mcombine_impl_stage2(f_c, α, β) + +@inline function _generic_mcombine_impl_stage1(f_c, α::AbstractMeasure, β::AbstractMeasure) + _generic_mcombine_impl_stage2(f_c, α, β) +end + +@inline function _generic_mcombine_impl_stage2(f_c, α::AbstractMeasure, β::AbstractMeasure) + FC, MA, MB = Core.Typeof(f_c), Core.Typeof(α), Core.Typeof(β) + CombinedMeasure{FC,MA,MB}(f_c, α, β) +end + +@inline function _generic_mcombine_impl_stage2(f_c, α::Dirac, β::Dirac) + Dirac(f_c(α.x, β.x)) +end + + +""" + struct CombinedMeasure <: AbstractMeasure + +Represents a combination of two measures. + +User code should not create instances of `CombinedMeasure` directly, but +should call [`mcombine(f_c, α, β)`](@ref) instead. +""" +struct CombinedMeasure{FC,MA<:AbstractMeasure,MB<:AbstractMeasure} <: AbstractMeasure + f_c::FC + α::MA + β::MB +end + + +@inline insupport(μ::CombinedMeasure, ab) = NoFastInsupport{typeof(μ)}() + +@inline getdof(μ::CombinedMeasure) = getdof(μ.α) + getdof(μ.β) +@inline fast_dof(μ::CombinedMeasure) = fast_dof(μ.α) + fast_dof(μ.β) + +# Bypass `checked_arg`, would require splitting ab: +@inline checked_arg(::CombinedMeasure, ab) = ab + +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 + +# Specialize logdensityof directly to avoid creating temporary combined base measures: +function logdensityof(μ::CombinedMeasure, ab) + tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, ab) + return logdensityof(tpm_α, a) + logdensityof(μ.β, b) +end + +function unsafe_logdensityof(μ::CombinedMeasure, ab) + tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, ab) + return unsafe_logdensityof(tpm_α, a) + unsafe_logdensityof(μ.β, b) +end + + +function Base.rand(rng::Random.AbstractRNG, ::Type{T}, μ::CombinedMeasure) where {T<:Real} + a = rand(rng, T, μ.α) + b = rand(rng, T, μ.β) + return μ.f_c(a, b) +end + +function Base.rand(rng::Random.AbstractRNG, μ::CombinedMeasure) + a = rand(rng, μ.α) + b = rand(rng, μ.β) + return μ.f_c(a, b) +end + + +function transport_to_mvstd(ν_inner::StdMeasure, μ::CombinedMeasure, ab) + tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, ab) + y1 = transport_to_mvstd(ν_inner, tpm_α, a) + y2 = transport_to_mvstd(ν_inner, μ.β, b) + return vcat(y1, y2) +end + + +function transport_from_mvstd_with_rest(ν::CombinedMeasure, μ_inner::StdMeasure, x) + a, x2 = transport_from_mvstd_with_rest(ν.α, μ_inner, x) + b, x_rest = transport_from_mvstd_with_rest(ν.β, μ_inner, x2) + return ν.f_c(a, b), x_rest +end diff --git a/src/combinators/product_transport.jl b/src/combinators/product_transport.jl index a2b44b1b..01c1016c 100644 --- a/src/combinators/product_transport.jl +++ b/src/combinators/product_transport.jl @@ -322,6 +322,14 @@ function _marginals_to_mvstd(::NU, marginals_μ::Tuple, x::Tuple) where {NU<:Std _flatten_to_rv(map(_TransportToMvStd{NU}(), marginals_μ, x)) end +function _marginals_to_mvstd( + ν::NU, + marginals_μ::NamedTuple{names}, + x::NamedTuple{names}, +) where {NU<:StdMeasure,names} + _marginals_to_mvstd(ν, values(marginals_μ), values(x)) +end + function _marginals_to_mvstd(::NU, marginals_μ, x) where {NU<:StdMeasure} _flatten_to_rv(broadcast(_TransportToMvStd{NU}(), marginals_μ, x)) end @@ -339,6 +347,16 @@ function transport_from_mvstd_with_rest(ν::ProductMeasure, μ_inner::StdMeasure return _marginals_from_mvstd_with_rest(νs, dofs, μ_inner, x) end +function transport_from_mvstd_with_rest( + ν::ProductMeasure{<:NamedTuple{names}}, + μ_inner::StdMeasure, + x, +) where {names} + ys, x_rest = + transport_from_mvstd_with_rest(productmeasure(values(marginals(ν))), μ_inner, x) + return NamedTuple{names}(ys), x_rest +end + function _dof_access_firstidxs(dofs::Tuple{Vararg{IntegerLike,N}}, first_idx) where {N} cumsum((first_idx, dofs[begin:(end-1)]...)) end diff --git a/test/Project.toml b/test/Project.toml index cb68fc93..32833feb 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -17,6 +17,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" LogarithmicNumbers = "aa2f6b4e-9042-5d33-9679-40d3a6b85899" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" +OneTwoMany = "762dc654-8631-413a-a342-372a7419ad9d" PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150" PropertyFunctions = "09e99361-2bb8-48a2-a80f-de58f0739eb4" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" diff --git a/test/combinators/bind.jl b/test/combinators/bind.jl new file mode 100644 index 00000000..a1d8865a --- /dev/null +++ b/test/combinators/bind.jl @@ -0,0 +1,94 @@ +# 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 + end + + @testset "mbind with merge" begin + αnt = productmeasure((position = StdNormal(),)) + f_βnt(a) = productmeasure(( + noise = pushfwd(Mul(abs(a.position) + 0.5), StdExponential()), + )) + μ = mbind(f_βnt, αnt, merge) + + x = rand(stblrng(), Float64, μ) + @test x isa NamedTuple{(:position, :noise)} + @test logdensityof(μ, x) ≈ + logdensityof(αnt, (position = x.position,)) + + logdensityof(f_βnt(x), (noise = x.noise,)) + + y = transport_to(StdUniform()^2, μ)(x) + @test y isa AbstractVector{<:Real} && length(y) == 2 + x_reco = transport_to(μ, StdUniform()^2)(y) + @test x_reco.position ≈ x.position && x_reco.noise ≈ x.noise + end + + @testset "mbind with Dirac" begin + @test mbind(f_β, MeasureBase.Dirac(1.5)) == asmeasure(f_β(1.5)) + μ = mbind(f_β, MeasureBase.Dirac(1.5), tuple) + ab = rand(stblrng(), Float64, μ) + @test ab[1] == 1.5 + end +end diff --git a/test/combinators/combined.jl b/test/combinators/combined.jl new file mode 100644 index 00000000..9a79f225 --- /dev/null +++ b/test/combinators/combined.jl @@ -0,0 +1,52 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using Random +using StableRNGs: StableRNG +using OneTwoMany: firstarg, secondarg + +using MeasureBase +using MeasureBase: StdExponential, StdLogistic, StdNormal, StdUniform +using MeasureBase: mcombine, productmeasure, transport_to + +@testset "mcombine" begin + stblrng() = StableRNG(789990641) + + α = StdExponential() + β = StdLogistic() + + @testset "combination shortcuts" begin + @test mcombine(firstarg, α, β) === α + @test mcombine(secondarg, α, β) === β + @test mcombine(tuple, α, β) == productmeasure((α, β)) + @test mcombine(vcat, StdNormal()^2, StdNormal()^1) == StdNormal()^3 + @test mcombine(vcat, productmeasure([α, α]), productmeasure([α])) == + productmeasure([α, α, α]) + @test mcombine( + merge, + productmeasure((a = α,)), + productmeasure((b = β,)), + ) == productmeasure((a = α, b = β)) + @test mcombine(tuple, MeasureBase.Dirac(1), MeasureBase.Dirac(2)) == + MeasureBase.Dirac((1, 2)) + end + + @testset "CombinedMeasure" begin + μ = mcombine(Pair, α, β) + @test μ isa MeasureBase.CombinedMeasure + + ab = rand(stblrng(), Float64, μ) + @test ab isa Pair + @test logdensityof(μ, ab) ≈ logdensityof(α, ab.first) + logdensityof(β, ab.second) + + @test MeasureBase.getdof(μ) == 2 + @test MeasureBase.fast_dof(μ) == 2 + @test MeasureBase.insupport(μ, ab) isa MeasureBase.NoFastInsupport + + y = transport_to(StdUniform()^2, μ)(ab) + @test y isa AbstractVector{<:Real} && length(y) == 2 + ab_reco = transport_to(μ, StdUniform()^2)(y) + @test ab_reco.first ≈ ab.first && ab_reco.second ≈ ab.second + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 7183c58f..d24f58e2 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -26,6 +26,8 @@ include("combinators/weighted.jl") include("combinators/transformedmeasure.jl") include("combinators/reshape.jl") include("combinators/implicitlymapped.jl") +include("combinators/combined.jl") +include("combinators/bind.jl") include("distributions/test_distributions.jl") From 30973a7267a06d65b8d1914e13d34eb850430576 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 15:48:08 +0200 Subject: [PATCH 44/75] Extend proxy forwarding, pushforward and mass interface Forward `unsafe_logdensityof`, `rootmeasure`, `insupport`, `getdof`, `fast_dof`, `localmeasure` and `transportmeasure` through `@useproxy`. Pushforward measures now return `NoFastInsupport` from `insupport` (checking via the origin would require a potentially costly transformation), support `fast_dof` and gain curried `pushfwd(f)`/ `pullbck(f)` forms. `massof(::AbstractMeasure)` returns `UnknownMass()` (instead of an invalid constructor call), unknown-mass powers disambiguate integer and rational exponents, and the root-measure based `massof` implementation hook is now named `_default_massof_impl` and uses interval endpoints. Ports parts of the original measure-algebra commits 9517fb1, 8aec2f2 and related STASH commits. Co-Authored-By: Claude Fable 5 --- src/combinators/transformedmeasure.jl | 9 ++++++++- src/mass-interface.jl | 9 ++++++--- src/primitives/lebesgue.jl | 6 +++--- src/proxies.jl | 11 ++++++++++- test/combinators/transformedmeasure.jl | 2 +- test/distributions/test_conversions.jl | 4 ++-- 6 files changed, 30 insertions(+), 11 deletions(-) diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl index c9db7a6b..56112049 100644 --- a/src/combinators/transformedmeasure.jl +++ b/src/combinators/transformedmeasure.jl @@ -170,7 +170,9 @@ for func in [:logdensityof, :logdensity_def] end end -insupport(m::PushforwardMeasure, x) = insupport(transport_origin(m), to_origin(m, x)) +# Checking insupport via the origin would require a potentially costly +# transformation of x: +insupport(m::PushforwardMeasure, x) = NoFastInsupport{typeof(m)}() function testvalue(::Type{T}, ν::PushforwardMeasure) where {T} ν.f(testvalue(T, parent(ν))) @@ -193,6 +195,9 @@ _pushfwd_dof(::Type{MU}, ::Type{<:Tuple{Any,Real}}, dof) where {MU} = dof @inline getdof(ν::MU) where {MU<:PushforwardMeasure} = getdof(ν.origin) @inline getdof(m::_NonBijectivePusfwdMeasure) = MeasureBase.NoDOF{typeof(m)}() +@inline fast_dof(ν::PushforwardMeasure) = fast_dof(ν.origin) +@inline fast_dof(m::_NonBijectivePusfwdMeasure) = MeasureBase.NoDOF{typeof(m)}() + # Bypass `checked_arg`, would require potentially costly transformation: @inline checked_arg(::PushforwardMeasure, x) = x @@ -222,6 +227,7 @@ 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) @@ -263,6 +269,7 @@ 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) diff --git a/src/mass-interface.jl b/src/mass-interface.jl index 7b0518f9..59a5db88 100644 --- a/src/mass-interface.jl +++ b/src/mass-interface.jl @@ -22,7 +22,10 @@ for T in (:UnknownFiniteMass, :UnknownMass) @eval begin Base.:+(::$T, ::$T) = $T() Base.:*(::$T, ::$T) = $T() - Base.:^(::$T, k::Number) = isfinite(k) ? $T() : UnknownMass() + Base.:^(::$T, k::Real) = isfinite(k) ? $T() : UnknownMass() + # Disambiguation: + Base.:^(::$T, k::Integer) = isfinite(k) ? $T() : UnknownMass() + Base.:^(::$T, k::Rational) = isfinite(k) ? $T() : UnknownMass() end end @@ -65,7 +68,7 @@ finite, or we may know nothing at all about it. For these cases, it will return `UnknownFiniteMass` or `UnknownMass`, respectively. When no `massof` method exists, it defaults to `UnknownMass`. """ -massof(m::AbstractMeasure) = UnknownMass(m) +massof(::AbstractMeasure) = UnknownMass() struct NormalizedMeasure{P,M} <: AbstractMeasure parent::P @@ -105,7 +108,7 @@ isnormalized(x, p::Real = 2) = isone(norm(x, p)) isone(::AbstractUnknownMass) = false function massof(m, s) - _massof(m, s, rootmeasure(m)) + _default_massof_impl(m, s, rootmeasure(m)) end """ diff --git a/src/primitives/lebesgue.jl b/src/primitives/lebesgue.jl index 3d92a2ed..4b8bf7ab 100644 --- a/src/primitives/lebesgue.jl +++ b/src/primitives/lebesgue.jl @@ -26,12 +26,12 @@ end massof(::LebesgueBase) = static(Inf) -function _massof(m, s::Interval, ::LebesgueBase) +function _default_massof_impl(m, s::AbstractInterval, ::LebesgueBase) mass = massof(m) nu = mass * StdUniform() f = transport_to(nu, m) - a = f(minimum(s)) - b = f(maximum(s)) + a = f(leftendpoint(s)) + b = f(rightendpoint(s)) return mass * abs(b - a) end diff --git a/src/proxies.jl b/src/proxies.jl index 95aed270..109b9973 100644 --- a/src/proxies.jl +++ b/src/proxies.jl @@ -15,15 +15,24 @@ 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) diff --git a/test/combinators/transformedmeasure.jl b/test/combinators/transformedmeasure.jl index 28ddbb50..497f79c2 100644 --- a/test/combinators/transformedmeasure.jl +++ b/test/combinators/transformedmeasure.jl @@ -162,7 +162,7 @@ using ChangesOfVariables # Test rand @test rand(ν) isa Real - @test insupport(ν, rand(ν)) + @test insupport(ν, rand(ν)) != false # Test pullback pb = pullbck(f, ν) diff --git a/test/distributions/test_conversions.jl b/test/distributions/test_conversions.jl index 0f9c4d80..a8244b4d 100644 --- a/test/distributions/test_conversions.jl +++ b/test/distributions/test_conversions.jl @@ -30,13 +30,13 @@ using MeasureBase: logdensityof, massof, insupport 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) + @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) + @test insupport(m, x) != false end end From 27863de892ec77982ac5c8eea0cb8b4b18b13037 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 17:43:00 +0200 Subject: [PATCH 45/75] Use PushFwdStyle names directly, keep VolCorr names as compat aliases only Co-Authored-By: Claude Fable 5 --- src/combinators/transformedmeasure.jl | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl index 56112049..970b862d 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 From 451c95d019f38b0000557c840eeb93b866d899f6 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 20:02:13 +0200 Subject: [PATCH 46/75] Fix logweight scaling in powers of weighted measures The total weight was computed from the length of the new power measure, which fails when the base collapses to a non-power measure (e.g. for Dirac bases). Compute it from the exponent instead. Created by generative AI. --- src/combinators/smart-constructors.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl index cb15f6ad..a7fb9a80 100644 --- a/src/combinators/smart-constructors.jl +++ b/src/combinators/smart-constructors.jl @@ -41,7 +41,7 @@ end @inline function _generic_powermeasure_stage2(μ::WeightedMeasure, exponent::Tuple) ν = μ.base^exponent - k = maybestatic_length(ν) * μ.logweight + k = size2length(axes2size(exponent)) * μ.logweight return weightedmeasure(k, ν) end From da2ce09680c93d0903241bfb72e93f545f4653f4 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 20:02:44 +0200 Subject: [PATCH 47/75] Rework superpose into an optimizing pairwise algebra superpose now folds varargs pairwise: equal measures combine into weighted measures, weighted measures with equal bases add their weights, and superpositions merge their components. Merging no longer mutates existing superposition components (add_measures used push!). Collections of a single singleton measure type collapse to weighted measures. Also exports superpose, matching the other smart constructors. Created by generative AI. --- src/collection_utils.jl | 7 +++ src/combinators/smart-constructors.jl | 69 ++++++++++++++++++++------- 2 files changed, 60 insertions(+), 16 deletions(-) diff --git a/src/collection_utils.jl b/src/collection_utils.jl index d15131d0..be5cc934 100644 --- a/src/collection_utils.jl +++ b/src/collection_utils.jl @@ -94,3 +94,10 @@ _flatten_to_rv(VV::VectorOfVectors{<:Real}) = 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) diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl index a7fb9a80..2b52fede 100644 --- a/src/combinators/smart-constructors.jl +++ b/src/combinators/smart-constructors.jl @@ -134,35 +134,72 @@ restrict(f, b) = RestrictedMeasure(f, b) ############################################################################### # 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. + +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. +Collections are wrapped as-is, apart from cost-free structural simplifications. +""" +function superpose end +export superpose + +superpose(μ::AbstractMeasure) = μ + +function superpose(μ::AbstractMeasure, ν::AbstractMeasure, more::AbstractMeasure...) + superpose(_superpose_two(μ, ν), more...) +end -function superpose(μ::T, ν::T) where {T<:AbstractMeasure} - if μ == ν - return weightedmeasure(static(float(logtwo)), μ) +function superpose(a::AbstractArray{T}) where {T} + if Base.issingletontype(T) + weightedmeasure(log(length(a)), asmeasure(instance(T))) else - return superpose((μ, ν)) + SuperpositionMeasure(a) end end -function superpose(μ::AbstractMeasure, μs...) - if all(==(μ), μs) - return weightedmeasure(log(length(μs) + 1), μ) +superpose(a::FillArrays.Fill) = weightedmeasure(log(length(a)), asmeasure(_fill_value(a))) + +superpose(t::Tuple) = SuperpositionMeasure(t) +superpose(nt::NamedTuple) = SuperpositionMeasure(nt) + +function _superpose_two(μ::AbstractMeasure, ν::AbstractMeasure) + μ == ν ? weightedmeasure(static(float(logtwo)), μ) : SuperpositionMeasure((μ, ν)) +end + +function _superpose_two(μ::WeightedMeasure, ν::WeightedMeasure) + if μ.base == ν.base + weightedmeasure(logaddexp(asnonstatic(μ.logweight), asnonstatic(ν.logweight)), μ.base) else - return superpose((μ, μs...)) + SuperpositionMeasure((μ, ν)) end end -add_measures(μs::AbstractVector, νs) = push!(μs, νs...) -add_measures(μs::Tuple, νs) = (μs..., νs...) +function _superpose_two(μ::WeightedMeasure, ν::AbstractMeasure) + μ.base == ν ? weightedmeasure(log1pexp(asnonstatic(μ.logweight)), μ.base) : + SuperpositionMeasure((μ, ν)) +end -function superpose(μ::SuperpositionMeasure, μs...) - SuperpositionMeasure(add_measures(μ.components, μs)) +function _superpose_two(μ::AbstractMeasure, ν::WeightedMeasure) + μ == ν.base ? weightedmeasure(log1pexp(asnonstatic(ν.logweight)), ν.base) : + SuperpositionMeasure((μ, ν)) end -superpose(μ::SuperpositionMeasure) = μ +_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)) ############################################################################### # WeightedMeasure From 4263d57beedd75502fb779574f862f711fa655ba Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 20:02:53 +0200 Subject: [PATCH 48/75] Add value-based equality and isapprox for Dirac Dirac measures with equal points are equal measures; the default object identity failed for array-valued points. Created by generative AI. --- src/primitives/dirac.jl | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl index 01297486..f5e5931a 100644 --- a/src/primitives/dirac.jl +++ b/src/primitives/dirac.jl @@ -9,6 +9,9 @@ function Pretty.tile(d::Dirac) Pretty.literal("Dirac(") * Pretty.tile(d.x) * Pretty.literal(")") end +Base.:(==)(a::Dirac, b::Dirac) = a.x == b.x +Base.isapprox(a::Dirac, b::Dirac; kwargs...) = isapprox(a.x, b.x; kwargs...) + gentype(μ::Dirac{X}) where {X} = X function (μ::Dirac{X})(s) where {X} From d7b63b25e3217b4e5aafbe20b1e49cfac60b1f49 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 20:03:26 +0200 Subject: [PATCH 49/75] Extend product measure simplifications Products of weighted measures (tuples, named tuples and arrays) now pull the total weight out, following the canonical measure nesting. Empty products are explicit unit measures (Dirac of the empty variate). Arrays of measures are no longer copied via broadcast asmeasure, and arrays of a single singleton measure type collapse to power measures without requiring a non-empty array. The static-weight optimization now stays static when the array length is static. Created by generative AI. --- src/combinators/smart-constructors.jl | 46 +++++++++++++++++---------- 1 file changed, 30 insertions(+), 16 deletions(-) diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl index 2b52fede..b3a00496 100644 --- a/src/combinators/smart-constructors.jl +++ b/src/combinators/smart-constructors.jl @@ -69,29 +69,50 @@ export productmeasure @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{Vararg{Dirac}}) = Dirac(map(m -> m.x, 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{Vararg{Dirac}}}) where {names} = - Dirac(map(m -> m.x, 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)) -@inline _generic_productmeasure_impl(mar::AbstractArray{<:AbstractProductMeasure}) = - ProductMeasure(mar) - _generic_productmeasure_impl(mar::AbstractArray{<:Dirac}) = Dirac((m -> m.x).(mar)) -# TODO: We should be able to further optimize this +_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 _generic_productmeasure_impl(mar::AbstractArray{T}) where {T} if Base.issingletontype(T) - first(mar)^size(mar) + powermeasure(instance(T), axes(mar)) + elseif T <: AbstractMeasure + ProductMeasure(mar) else - ProductMeasure(asmeasure.(mar)) + ProductMeasure(map(asmeasure, mar)) end end @@ -103,13 +124,6 @@ end @inline _generic_productmeasure_impl(mar::Base.Generator) = ProductMeasure(mar) -# TODO: Make this static when its length is static -@inline function _generic_productmeasure_impl( - mar::AbstractArray{<:WeightedMeasure{StaticFloat64{W},M}}, -) where {W,M} - return weightedmeasure(W * length(mar), productmeasure(map(basemeasure, mar))) -end - # ToDo: Remove or at least refactor this (ProductMeasure shouldn't take a kernel as its argument). productmeasure(f, param_maps, pars) = productmeasure(kernel(f, param_maps), pars) From 32e5913ef96ec373be31c6c96abf64ad2dc8407b Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 20:03:50 +0200 Subject: [PATCH 50/75] Simplify pushforwards of Dirac and weighted measures Pushforwards of point masses collapse to point masses and pushforwards commute with weighting. The identity shortcut moves from _pushfwd_impl to pushfwd itself and the two style-specific pushfwd/pullbck methods merge into single PushFwdStyle methods, keeping dispatch unambiguous. Created by generative AI. --- src/combinators/smart-constructors.jl | 13 +++++++++++++ src/combinators/transformedmeasure.jl | 13 +++++++------ test/combinators/transformedmeasure.jl | 7 ++++--- 3 files changed, 24 insertions(+), 9 deletions(-) diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl index b3a00496..90b646c8 100644 --- a/src/combinators/smart-constructors.jl +++ b/src/combinators/smart-constructors.jl @@ -139,6 +139,19 @@ function productmeasure(f::Returns{W}, ::typeof(identity), pars) where {W<:Weigh weightedmeasure(length(pars) * ℓ, newbase) end +############################################################################### +# PushforwardMeasure + +# 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)) + +# Pushforward and weighting commute: +function _pushfwd_impl(f, μ::WeightedMeasure, style::PushFwdStyle) + weightedmeasure(μ.logweight, _pushfwd_impl(f, μ.base, style)) +end + ############################################################################### # RestrictedMeasure export restrict diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl index 970b862d..0017c86f 100644 --- a/src/combinators/transformedmeasure.jl +++ b/src/combinators/transformedmeasure.jl @@ -232,8 +232,10 @@ export pushfwd @inline pushfwd(f) = Base.Fix1(pushfwd, f) @inline pushfwd(f, μ) = _pushfwd_impl(f, μ, AdaptRootMeasure()) -@inline pushfwd(f, μ, style::AdaptRootMeasure) = _pushfwd_impl(f, μ, style) -@inline pushfwd(f, μ, style::PushfwdRootMeasure) = _pushfwd_impl(f, μ, style) +@inline pushfwd(f, μ, style::PushFwdStyle) = _pushfwd_impl(f, μ, style) + +@inline pushfwd(::typeof(identity), μ) = μ +@inline pushfwd(::typeof(identity), μ, ::PushFwdStyle) = μ _pushfwd_impl(f, μ, style) = PushforwardMeasure(f, inverse(f), μ, style) @@ -248,8 +250,8 @@ function _pushfwd_impl( PushforwardMeasure(new_f, new_f_inv, orig_μ, style) end -_pushfwd_impl(::typeof(identity), μ, ::AdaptRootMeasure) = μ -_pushfwd_impl(::typeof(identity), μ, ::PushfwdRootMeasure) = μ +# Simplifications for Dirac and WeightedMeasure origins are defined in +# smart-constructors.jl. ############################################################################### # pullback @@ -274,8 +276,7 @@ export pullbck @inline pullbck(f) = Base.Fix1(pullbck, f) @inline pullbck(f, μ) = _pullback_impl(f, μ, AdaptRootMeasure()) -@inline pullbck(f, μ, style::AdaptRootMeasure) = _pullback_impl(f, μ, style) -@inline pullbck(f, μ, style::PushfwdRootMeasure) = _pullback_impl(f, μ, style) +@inline pullbck(f, μ, style::PushFwdStyle) = _pullback_impl(f, μ, style) function _pullback_impl(f, μ, style = AdaptRootMeasure()) pushfwd(inverse(f), μ, style) diff --git a/test/combinators/transformedmeasure.jl b/test/combinators/transformedmeasure.jl index 497f79c2..6261c762 100644 --- a/test/combinators/transformedmeasure.jl +++ b/test/combinators/transformedmeasure.jl @@ -152,9 +152,10 @@ using ChangesOfVariables @test rootmeasure(ν) === rootmeasure(μ) # AdaptRootMeasure @test rootmeasure(ν_no_corr) isa PushforwardMeasure # PushfwdRootMeasure - # Test basemeasure - @test basemeasure(ν) isa PushforwardMeasure - @test basemeasure(ν).style isa PushfwdRootMeasure + # Test basemeasure. The base measure of μ is a weighted Lebesgue measure, + # so the weight gets pulled out of the pushforward: + @test basemeasure(ν) isa WeightedMeasure{<:Any,<:PushforwardMeasure} + @test basemeasure(ν).base.style isa PushfwdRootMeasure # Test massof # TODO: mass interface is very incomplete From 7ec90978e4b37eee01a6ce00093d5fcbcf0eb04e Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 20:04:07 +0200 Subject: [PATCH 51/75] Fuse nested measure restrictions and add curried restrict Created by generative AI. --- src/combinators/smart-constructors.jl | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl index 90b646c8..8292fa1c 100644 --- a/src/combinators/smart-constructors.jl +++ b/src/combinators/smart-constructors.jl @@ -156,7 +156,13 @@ 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 From af701872bf4e4c123d1ed9f4670766ef9e5df07d Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 20:04:29 +0200 Subject: [PATCH 52/75] Document and export half and weightedmeasure Matches the other smart constructors, which are documented and exported. Created by generative AI. --- src/combinators/smart-constructors.jl | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl index 8292fa1c..db355184 100644 --- a/src/combinators/smart-constructors.jl +++ b/src/combinators/smart-constructors.jl @@ -7,7 +7,14 @@ ############################################################################### # 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 ############################################################################### # PowerMeaure @@ -237,6 +244,16 @@ _superpose_two(μ::WeightedMeasure, ν::SuperpositionMeasure) = ############################################################################### # WeightedMeasure +""" + weightedmeasure(logweight::Real, μ) + +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 weightedmeasure(ℓ::R, b::M) where {R,M} WeightedMeasure{R,M}(ℓ, b) end From 0ce4beb8056209889655232510fb5a0f65e0cc1d Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 20:04:44 +0200 Subject: [PATCH 53/75] Add smart constructor tests Covers the power/product/superpose/pushfwd/restrict simplifications. Also wires the previously orphaned superpose tests into the test suite, updated to current basemeasure collapse behavior. Created by generative AI. --- test/combinators/smart_constructors.jl | 139 +++++++++++++++++++++++++ test/combinators/superpose.jl | 5 +- test/runtests.jl | 2 + 3 files changed, 143 insertions(+), 3 deletions(-) create mode 100644 test/combinators/smart_constructors.jl diff --git a/test/combinators/smart_constructors.jl b/test/combinators/smart_constructors.jl new file mode 100644 index 00000000..ac6da161 --- /dev/null +++ b/test/combinators/smart_constructors.jl @@ -0,0 +1,139 @@ +# 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 + + @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..17c677df 100644 --- a/test/combinators/superpose.jl +++ b/test/combinators/superpose.jl @@ -14,9 +14,8 @@ using MeasureBase: superpose μs = SuperpositionMeasure([μ, ν]) @test μs isa SuperpositionMeasure{<:AbstractVector{<:AbstractMeasure}} - @test_throws ErrorException density_def(μs, 0) - @test basemeasure(μs).components == - SuperpositionMeasure([CountingBase(), CountingBase()]).components + @test density_def(μs, 0) == 1.0 + @test basemeasure(μs) == weightedmeasure(log(2), CountingBase()) μ2 = μ + μ @test μ2 isa WeightedMeasure diff --git a/test/runtests.jl b/test/runtests.jl index d24f58e2..6168091a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -22,7 +22,9 @@ 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") From 5aa67b8363713bb273caae82fcdd735693fba509 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 20:18:01 +0200 Subject: [PATCH 54/75] Support transport between measures of unknown DOF Measure transport now falls back to a pivot through a flat vector of standard uniform variates, using the with-rest transport protocol, when the DOF of either side is not fast-computable. This enables transport between hierarchical measures (Bind) and between known-DOF and unknown-DOF measures, as already promised by the transport_def docstring. The mvstd gateway now also verifies that the produced variate length matches the target standard power measure, instead of silently returning a variate of the wrong length. Created by generative AI. --- src/combinators/product_transport.jl | 6 +++++- src/transport.jl | 20 +++++++++++++++++++- test/combinators/bind.jl | 17 +++++++++++++++++ 3 files changed, 41 insertions(+), 2 deletions(-) diff --git a/src/combinators/product_transport.jl b/src/combinators/product_transport.jl index 01c1016c..2d78ae2f 100644 --- a/src/combinators/product_transport.jl +++ b/src/combinators/product_transport.jl @@ -124,7 +124,11 @@ end # Transport to a multivariate standard measure from any measure: function transport_def(ν::StdPowerMeasure{NU,1}, μ::AbstractMeasure, x) where {NU} - transport_to_mvstd(pwr_base(ν), μ, x) + y = transport_to_mvstd(pwr_base(ν), μ, x) + if maybestatic_length(y) != maybestatic_length(ν) + throw(ArgumentError("Length of transport target doesn't match variate DOF during transport")) + end + return y end function transport_to_mvstd(ν_inner::StdMeasure, μ::AbstractMeasure, x) diff --git a/src/transport.jl b/src/transport.jl index 9cea6fed..65c5f386 100644 --- a/src/transport.jl +++ b/src/transport.jl @@ -137,11 +137,29 @@ end return static(10) end -# If both both measures have no origin: +# If both measures have no origin: function _transport_between_origins(ν, ::StaticInteger{0}, ::StaticInteger{0}, μ, x) + _transport_between_noorigins(ν, fast_dof(ν), fast_dof(μ), μ, x) +end + +function _transport_between_noorigins(ν, ::IntegerLike, ::IntegerLike, μ, x) _transport_with_intermediate(ν, _transport_intermediate(ν, μ), μ, x) end +# If the DOF of either side is not known, pivot through a flat vector of +# standard-measure variates, using the with-rest transport protocol: +_transport_between_noorigins(ν, ::AbstractNoDOF, ::IntegerLike, μ, x) = + _transport_mvstd_pivot(ν, μ, x) +_transport_between_noorigins(ν, ::IntegerLike, ::AbstractNoDOF, μ, x) = + _transport_mvstd_pivot(ν, μ, x) +_transport_between_noorigins(ν, ::AbstractNoDOF, ::AbstractNoDOF, μ, x) = + _transport_mvstd_pivot(ν, μ, x) + +function _transport_mvstd_pivot(ν, μ, x) + z = transport_to_mvstd(StdUniform(), μ, x) + return _transport_from_mvstd(ν, StdUniform(), z) +end + @generated function _transport_between_origins( ν, ::StaticInteger{n_ν}, diff --git a/test/combinators/bind.jl b/test/combinators/bind.jl index a1d8865a..1503f6e2 100644 --- a/test/combinators/bind.jl +++ b/test/combinators/bind.jl @@ -64,6 +64,23 @@ using MeasureBase: pushfwd, productmeasure, transport_to, transportmeasure, loca @test all(u -> 0 <= u <= 1, y) xy_reco = transport_to(μ, StdUniform()^3)(y) @test xy_reco ≈ xy + + # Transport between two measures of unknown DOF (mvstd pivot): + μ2 = mbind(f_βv, StdUniform()^1, vcat) + xy2 = transport_to(μ2, μ)(xy) + @test xy2 isa AbstractVector{<:Real} && length(xy2) == 3 + @test transport_to(μ, μ2)(xy2) ≈ xy + @test logdensityof(μ2, xy2) isa Real + + # Transport between known-DOF and unknown-DOF measures (mvstd pivot): + ν_known = productmeasure((StdNormal(), StdNormal(), StdNormal())) + z = transport_to(ν_known, μ)(xy) + @test z isa Tuple{Vararg{Real,3}} + @test collect(transport_to(μ, ν_known)(z)) ≈ xy + + # DOF mismatches must not go unnoticed: + @test_throws ArgumentError transport_to(StdUniform()^5, μ)(xy) + @test_throws ArgumentError transport_to(StdUniform()^2, μ)(xy) end @testset "mbind with merge" begin From a13275b8b9e338445a61c316b8cf17f3b23e74af Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 21:58:48 +0200 Subject: [PATCH 55/75] Make superpose simplifications type stable Superposition simplifications now only happen when measure equality is decidable from the measure types alone (via the new _static_isequal), instead of branching on runtime measure equality. As a consequence, value-equal measures of non-singleton type (e.g. equal Diracs) now stay superpositions instead of collapsing to weighted measures, and basemeasure of superpositions with equal-typed components is type stable. Created by generative AI. --- src/combinators/smart-constructors.jl | 30 ++++++++++++++++++++------ test/combinators/smart_constructors.jl | 12 +++++++++++ test/combinators/superpose.jl | 13 ++++++++--- 3 files changed, 46 insertions(+), 9 deletions(-) diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl index db355184..e2c4cfe3 100644 --- a/src/combinators/smart-constructors.jl +++ b/src/combinators/smart-constructors.jl @@ -184,6 +184,8 @@ as a collection (array, tuple or named tuple) of measures. 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 @@ -208,12 +210,22 @@ superpose(a::FillArrays.Fill) = weightedmeasure(log(length(a)), asmeasure(_fill_ superpose(t::Tuple) = SuperpositionMeasure(t) superpose(nt::NamedTuple) = SuperpositionMeasure(nt) +# 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 _superpose_two(μ::AbstractMeasure, ν::AbstractMeasure) - μ == ν ? weightedmeasure(static(float(logtwo)), μ) : SuperpositionMeasure((μ, ν)) + if _static_isequal(μ, ν) isa True + weightedmeasure(static(float(logtwo)), μ) + else + SuperpositionMeasure((μ, ν)) + end end function _superpose_two(μ::WeightedMeasure, ν::WeightedMeasure) - if μ.base == ν.base + if _static_isequal(μ.base, ν.base) isa True weightedmeasure(logaddexp(asnonstatic(μ.logweight), asnonstatic(ν.logweight)), μ.base) else SuperpositionMeasure((μ, ν)) @@ -221,13 +233,19 @@ function _superpose_two(μ::WeightedMeasure, ν::WeightedMeasure) end function _superpose_two(μ::WeightedMeasure, ν::AbstractMeasure) - μ.base == ν ? weightedmeasure(log1pexp(asnonstatic(μ.logweight)), μ.base) : - SuperpositionMeasure((μ, ν)) + if _static_isequal(μ.base, ν) isa True + weightedmeasure(log1pexp(asnonstatic(μ.logweight)), μ.base) + else + SuperpositionMeasure((μ, ν)) + end end function _superpose_two(μ::AbstractMeasure, ν::WeightedMeasure) - μ == ν.base ? weightedmeasure(log1pexp(asnonstatic(ν.logweight)), ν.base) : - SuperpositionMeasure((μ, ν)) + if _static_isequal(μ, ν.base) isa True + weightedmeasure(log1pexp(asnonstatic(ν.logweight)), ν.base) + else + SuperpositionMeasure((μ, ν)) + end end _superpose_two(μ::SuperpositionMeasure, ν::SuperpositionMeasure) = diff --git a/test/combinators/smart_constructors.jl b/test/combinators/smart_constructors.jl index ac6da161..06d7906b 100644 --- a/test/combinators/smart_constructors.jl +++ b/test/combinators/smart_constructors.jl @@ -103,6 +103,18 @@ using Static: static 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), μ) diff --git a/test/combinators/superpose.jl b/test/combinators/superpose.jl index 17c677df..753d5b8c 100644 --- a/test/combinators/superpose.jl +++ b/test/combinators/superpose.jl @@ -1,7 +1,7 @@ using Test using MeasureBase -using MeasureBase: superpose +using MeasureBase: superpose, weightedmeasure, StdNormal @testset "superpose.jl" begin μ = Dirac(0) @@ -17,8 +17,15 @@ using MeasureBase: superpose @test density_def(μs, 0) == 1.0 @test basemeasure(μs) == weightedmeasure(log(2), CountingBase()) + # Dirac equality is not decidable from types, so no weighted collapse: μ2 = μ + μ - @test μ2 isa WeightedMeasure + @test μ2 isa SuperpositionMeasure @test μ2 == superpose(μ, μ) - @test basemeasure(μ2) == μ + @test density_def(μ2, 0) == 1.0 + + # For singleton measure types equal measures combine into weighted measures: + s2 = StdNormal() + StdNormal() + @test s2 isa WeightedMeasure + @test exp(s2.logweight) ≈ 2 + @test basemeasure(s2) == StdNormal() end From 2068175a50a51d94bbe262323c73c21f97c341b7 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 22:06:55 +0200 Subject: [PATCH 56/75] Rework non-relative density evaluation around logdensityof_impl and with-rest logdensityof is now a single generic entry function; measure types specialize the new MeasureBase.logdensityof_impl instead of logdensityof itself. The new logdensityof_with_rest protocol evaluates densities of measures that live at the beginning of a flat variate stream (a vector for vcat-combined and a NamedTuple for merge-combined measures), returning the log-density, the consumed variate and the unconsumed rest of the stream. Its default implementation determines the variate size via mspace_elsize (falling back to testvalue). Bind and CombinedMeasure evaluate densities of vcat- and merge-combined variates through logdensityof_with_rest in a single pass now, without materializing intermediate transport measures and without splitting sub-variates twice. Variates that are too long now result in an informative exception, and so do binds whose value combination function does not support variate splitting. Created by generative AI. --- src/collection_utils.jl | 9 ++++ src/combinators/bind.jl | 53 ++++++++++++++++---- src/combinators/combined.jl | 47 +++++++++++++++--- src/combinators/half.jl | 2 +- src/combinators/power.jl | 8 ++-- src/combinators/product.jl | 10 ++-- src/combinators/transformedmeasure.jl | 10 ++-- src/density-core.jl | 69 +++++++++++++++++++++++---- src/density.jl | 2 +- src/primitive.jl | 4 +- src/primitives/counting.jl | 4 +- src/primitives/dirac.jl | 4 +- src/primitives/lebesgue.jl | 4 +- src/standard/stdexponential.jl | 2 +- src/standard/stdlogistic.jl | 2 +- src/standard/stdnormal.jl | 2 +- src/standard/stduniform.jl | 2 +- test/combinators/bind.jl | 10 ++++ 18 files changed, 192 insertions(+), 52 deletions(-) diff --git a/src/collection_utils.jl b/src/collection_utils.jl index be5cc934..6b9de0d5 100644 --- a/src/collection_utils.jl +++ b/src/collection_utils.jl @@ -101,3 +101,12 @@ _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`: +Base.@propagate_inbounds _consume_from_stream(x::AbstractVector, sz::Tuple{IntegerLike}) = + _split_after(x, sz[1]) + +function _consume_from_stream(x::AbstractVector, @nospecialize(sz::Tuple)) + 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 27cb02a4..ab7ab708 100644 --- a/src/combinators/bind.jl +++ b/src/combinators/bind.jl @@ -284,18 +284,51 @@ testvalue(::Bind) = throw(ArgumentError("testvalue is not available for Bind")) logdensity_def(::Bind, x) = throw(ArgumentError("logdensity_def is not available for Bind")) -# Specialize logdensityof to avoid duplicate calculations: -function logdensityof(μ::Bind, x) - tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, x) - β_a = _get_β_a(μ, a) - logdensityof(tpm_α, a) + logdensityof(β_a, b) +# 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 -# Specialize unsafe_logdensityof to avoid duplicate calculations: -function unsafe_logdensityof(μ::Bind, x) - tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, x) - β_a = _get_β_a(μ, a) - unsafe_logdensityof(tpm_α, a) + unsafe_logdensityof(β_a, b) +function _bind_ld_impl(::Type{Pair}, μ::Bind, xy::Pair) + a, b = xy.first, xy.second + logdensityof(μ.α, a) + logdensityof(_get_β_a(μ, a), b) +end + +function _bind_ld_impl(::Union{typeof(vcat),typeof(merge)}, μ::Bind, xy) + ℓ, x_μ, x_rest = logdensityof_with_rest(μ, xy) + if !isempty(x_rest) + throw(ArgumentError("Variate too long during density evaluation of a bind")) + end + return ℓ +end + +function _bind_ld_impl(@nospecialize(f_c), @nospecialize(μ::Bind), @nospecialize(xy)) + throw( + ArgumentError( + "Can't compute density of a bind with value combination function of type $(nameof(typeof(f_c)))", + ), + ) +end + +function logdensityof_with_rest(μ::_BindBy{typeof(vcat)}, x::AbstractVector) + ℓ_a, a, x2 = logdensityof_with_rest(μ.α, x) + ℓ_b, b, x_rest = logdensityof_with_rest(_get_β_a(μ, a), x2) + x_μ, _ = _split_after(x, maybestatic_length(x) - maybestatic_length(x_rest)) + return ℓ_a + ℓ_b, x_μ, x_rest +end + +function logdensityof_with_rest(μ::_BindBy{typeof(merge)}, x::NamedTuple) + ℓ_a, a, x2 = logdensityof_with_rest(μ.α, x) + ℓ_b, b, x_rest = logdensityof_with_rest(_get_β_a(μ, a), x2) + return ℓ_a + ℓ_b, merge(a, b), x_rest end diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl index c3ea079c..4325ced8 100644 --- a/src/combinators/combined.jl +++ b/src/combinators/combined.jl @@ -134,15 +134,50 @@ function logdensity_def(μ::CombinedMeasure, ab) return logdensity_def(tpm_α, a) + logdensity_def(μ.β, b) end -# Specialize logdensityof directly to avoid creating temporary combined base measures: -function logdensityof(μ::CombinedMeasure, ab) - tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, ab) +# Density evaluation consumes the variate parts of both component measures +# in a single pass, using the with-rest protocol for value-dependent +# variate sizes: + +logdensityof_impl(μ::CombinedMeasure, ab) = _combined_ld_impl(μ.f_c, μ, ab) + +unsafe_logdensityof(μ::CombinedMeasure, ab) = logdensityof_impl(μ, ab) + +function _combined_ld_impl(::typeof(tuple), μ::CombinedMeasure, ab::Tuple{Vararg{Any,2}}) + logdensityof(μ.α, ab[1]) + logdensityof(μ.β, ab[2]) +end + +function _combined_ld_impl(::Type{Pair}, μ::CombinedMeasure, ab::Pair) + logdensityof(μ.α, ab.first) + logdensityof(μ.β, ab.second) +end + +function _combined_ld_impl(::Union{typeof(vcat),typeof(merge)}, μ::CombinedMeasure, ab) + ℓ, x_μ, x_rest = logdensityof_with_rest(μ, ab) + if !isempty(x_rest) + throw( + ArgumentError( + "Variate too long during density evaluation of a combined measure", + ), + ) + end + return ℓ +end + +function _combined_ld_impl(f_c, μ::CombinedMeasure, ab) + tpm_α, a, b = tpmeasure_split_combined(f_c, μ.α, ab) return logdensityof(tpm_α, a) + logdensityof(μ.β, b) end -function unsafe_logdensityof(μ::CombinedMeasure, ab) - tpm_α, a, b = tpmeasure_split_combined(μ.f_c, μ.α, ab) - return unsafe_logdensityof(tpm_α, a) + unsafe_logdensityof(μ.β, b) +function logdensityof_with_rest(μ::CombinedMeasure{typeof(vcat)}, x::AbstractVector) + ℓ_a, a, x2 = logdensityof_with_rest(μ.α, x) + ℓ_b, b, x_rest = logdensityof_with_rest(μ.β, x2) + x_μ, _ = _split_after(x, maybestatic_length(x) - maybestatic_length(x_rest)) + return ℓ_a + ℓ_b, x_μ, x_rest +end + +function logdensityof_with_rest(μ::CombinedMeasure{typeof(merge)}, x::NamedTuple) + ℓ_a, a, x2 = logdensityof_with_rest(μ.α, x) + ℓ_b, b, x_rest = logdensityof_with_rest(μ.β, x2) + return ℓ_a + ℓ_b, merge(a, b), x_rest end diff --git a/src/combinators/half.jl b/src/combinators/half.jl index 24063b76..c93327b0 100644 --- a/src/combinators/half.jl +++ b/src/combinators/half.jl @@ -19,7 +19,7 @@ function Base.rand(rng::AbstractRNG, ::Type{T}, μ::Half) where {T} return abs(rand(rng, T, unhalf(μ))) end -function logdensityof(μ::Half, x) +function logdensityof_impl(μ::Half, x) ld = logdensityof(unhalf(μ), x) - loghalf return x ≥ 0 ? ld : oftype(ld, -Inf) end diff --git a/src/combinators/power.jl b/src/combinators/power.jl index 1ce9ca98..5145a80c 100644 --- a/src/combinators/power.jl +++ b/src/combinators/power.jl @@ -97,8 +97,8 @@ 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} +for (head, func) in [(:logdensityof_impl, :logdensityof), (:logdensity_def, :logdensity_def)] + @eval @inline function $head(d::PowerMeasure{M}, x) where {M} parent_m = d.parent sz_parent = axes2size(d.axes) sz_x = maybestatic_size(x) @@ -114,7 +114,7 @@ for func in [:logdensityof, :logdensity_def] end end - @eval @inline function $func( + @eval @inline function $head( d::PowerMeasure{<:Any,Tuple{<:StaticOneToLike{N}}}, x, ) where {N} @@ -124,7 +124,7 @@ for func in [:logdensityof, :logdensity_def] end end - @eval @inline function $func( + @eval @inline function $head( ::PowerMeasure{<:Any,<:Tuple{Vararg{StaticOneToLike{0}}}}, x, ) diff --git a/src/combinators/product.jl b/src/combinators/product.jl index 656ded0d..9899e0a9 100644 --- a/src/combinators/product.jl +++ b/src/combinators/product.jl @@ -72,8 +72,8 @@ function _rand_product( end |> collect end -for func in [:logdensityof, :logdensity_def] - @eval @inline function $func(d::AbstractProductMeasure, x) +for (head, func) in [(:logdensityof_impl, :logdensityof), (:logdensity_def, :logdensity_def)] + @eval @inline function $head(d::AbstractProductMeasure, x) mapreduce($func, +, marginals(d), x) end end @@ -112,14 +112,14 @@ end return q end -for func in [:logdensityof, :logdensity_def] +for (head, func) in [(:logdensityof_impl, :logdensityof), (:logdensity_def, :logdensity_def)] # For tuples, `mapreduce` has trouble with type inference - @eval @inline function $func(d::ProductMeasure{T}, x) where {T<:Tuple} + @eval @inline function $head(d::ProductMeasure{T}, x) where {T<:Tuple} ℓs = map($func, marginals(d), x) sum(ℓs) end - @eval function $func(d::ProductMeasure{NamedTuple{N,T}}, x) where {N,T} + @eval function $head(d::ProductMeasure{NamedTuple{N,T}}, x) where {N,T} _product_gen_impl(Val($func), d, x) end end diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl index 0017c86f..ee960f9a 100644 --- a/src/combinators/transformedmeasure.jl +++ b/src/combinators/transformedmeasure.jl @@ -135,7 +135,7 @@ function _combine_logd_with_ladj(logd_orig::Real, ladj::Real) end end -function logdensityof( +function logdensityof_impl( @nospecialize(μ::_NonBijectivePusfwdMeasure{M,<:PushfwdRootMeasure}), @nospecialize(v::Any) ) where {M} @@ -146,7 +146,7 @@ function logdensityof( ) end -function logdensityof( +function logdensityof_impl( @nospecialize(μ::_NonBijectivePusfwdMeasure{M,<:AdaptRootMeasure}), @nospecialize(v::Any) ) where {M} @@ -157,15 +157,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) diff --git a/src/density-core.jl b/src/density-core.jl index 33a3c3cb..806a33a6 100644 --- a/src/density-core.jl +++ b/src/density-core.jl @@ -10,25 +10,41 @@ export density_rel export density_def """ - logdensityof(m::AbstractMeasure, x) + logdensityof(m::AbstractMeasure, x) Compute the log-density of the measure `m` at `x`. Density is always relative, but `DensityInterface.jl` does not account for this. For compatibility with this, `logdensityof` for a measure is always implicitly relative to -[`rootmeasure(x)`](@ref rootmeasure). +[`rootmeasure(x)`](@ref rootmeasure). -`logdensityof` works by first computing `insupport(m, x)`. If this is true, then -`unsafe_logdensityof` is called. If `insupport(m, x)` is known to be `true`, it -can be a little faster to directly call `unsafe_logdensityof(m, x)`. +`logdensityof(m, x)` is implemented via +[`MeasureBase.logdensityof_impl`](@ref), measure types should specialize +`logdensityof_impl` instead of `logdensityof` itself. To compute log-density relative to `basemeasure(m)` or *define* a log-density (relative to `basemeasure(m)` or another measure given explicitly), see -`logdensity_def`. +`logdensity_def`. To compute a log-density relative to a specific base-measure, see -`logdensity_rel`. +`logdensity_rel`. """ -@inline function logdensityof(μ::AbstractMeasure, x) +@inline logdensityof(μ::AbstractMeasure, x) = logdensityof_impl(μ, x) + +""" + MeasureBase.logdensityof_impl(μ::AbstractMeasure, x) + +Implements [`logdensityof(μ, x)`](@ref logdensityof). + +Measure types should specialize `logdensityof_impl` instead of +`logdensityof` itself. Implementations must return the log-density of `μ` +at `x` relative to [`rootmeasure(μ)`](@ref) and must handle `x` outside of +the support of `μ` (the result must be `-Inf` then). + +The default implementation checks `insupport(μ, x)` (unless the result is +a [`MeasureBase.NoFastInsupport`](@ref)) and computes the density via +[`unsafe_logdensityof`](@ref). +""" +@inline function logdensityof_impl(μ::AbstractMeasure, x) result = dynamic(unsafe_logdensityof(μ, x)) _checksupport(insupport(μ, x), result) end @@ -40,6 +56,43 @@ end _checksupport(cond, result) = ifelse(cond == true, result, oftype(result, -Inf)) @inline _checksupport(::NoFastInsupport, result) = result +""" + MeasureBase.logdensityof_with_rest(μ::AbstractMeasure, x) + +Compute the log-density of `μ` at the beginning of `x`, a flat stream of +variate content that may extend beyond the variate of `μ`. + +`x` must either be a vector that starts with the (flattened) variate of `μ` +(for measures combined via `vcat`) or a `NamedTuple` whose first properties +constitute the variate of `μ` (for measures combined via `merge`). + +Returns a tuple `(ℓ, x_μ, x_rest)` of the log-density `ℓ`, the variate +`x_μ` of `μ` consumed from the stream, and the unconsumed rest of the +stream. + +Measure types whose variate size depends on measure values, like +[`mbind`](@ref) results, implement density calculation via +`logdensityof_with_rest` instead of +[`logdensityof_impl`](@ref MeasureBase.logdensityof_impl). + +The default implementation determines the size resp. the property names of +the variate via [`some_mspace_elsize`](@ref MeasureBase.some_mspace_elsize) +resp. `testvalue` and delegates to `logdensityof_impl`. +""" +function logdensityof_with_rest end + +function logdensityof_with_rest(μ::AbstractMeasure, x::AbstractVector) + a, x_rest = _consume_from_stream(x, some_mspace_elsize(μ)) + return logdensityof_impl(μ, a), a, x_rest +end + +function logdensityof_with_rest(μ::AbstractMeasure, x::NamedTuple) + a, x_rest = _split_after(x, Val(_mspace_names(μ))) + return logdensityof_impl(μ, a), a, x_rest +end + +_mspace_names(μ::AbstractMeasure) = keys(testvalue(μ)) + """ localmeasure(m::AbstractMeasure, x)::AbstractMeasure diff --git a/src/density.jl b/src/density.jl index 06dc98e1..67d34464 100644 --- a/src/density.jl +++ b/src/density.jl @@ -224,7 +224,7 @@ logdensity_def(μ::DensityMeasure, x) = logdensityof(μ.f, x) density_def(μ::DensityMeasure, x) = densityof(μ.f, x) -function logdensityof(μ::DensityMeasure, x::Any) +function logdensityof_impl(μ::DensityMeasure, x::Any) integrand, μ_base = μ.f, μ.base base_logval = logdensityof(μ_base, x) diff --git a/src/primitive.jl b/src/primitive.jl index 85cf2beb..f43485c2 100644 --- a/src/primitive.jl +++ b/src/primitive.jl @@ -19,8 +19,8 @@ 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::Real) = zero(float(typeof(x))) +@inline logdensityof_impl(::PrimitiveMeasure, x) = false logdensity_def(::PrimitiveMeasure, x) = static(0.0) diff --git a/src/primitives/counting.jl b/src/primitives/counting.jl index 31398f32..5a7a998e 100644 --- a/src/primitives/counting.jl +++ b/src/primitives/counting.jl @@ -12,12 +12,12 @@ struct Counting{T} <: AbstractMeasure Counting(supp) = new{Core.Typeof(supp)}(supp) end -@inline function logdensityof(μ::Counting, x::Real) +@inline function logdensityof_impl(μ::Counting, x::Real) R = float(typeof(x)) insupport(μ, x) ? zero(R) : R(-Inf) end -@inline logdensityof(μ::Counting, x) = insupport(μ, x) ? 0.0 : -Inf +@inline logdensityof_impl(μ::Counting, x) = insupport(μ, x) ? 0.0 : -Inf @inline logdensity_def(μ::Counting, x) = logdensityof(μ, x) diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl index f5e5931a..ba044c1d 100644 --- a/src/primitives/dirac.jl +++ b/src/primitives/dirac.jl @@ -23,12 +23,12 @@ basemeasure(d::Dirac) = CountingBase() massof(::Dirac) = static(1.0) -function logdensityof(μ::Dirac, x::Real) +function logdensityof_impl(μ::Dirac, x::Real) R = float(typeof(x)) insupport(μ, x) ? zero(R) : R(-Inf) end -logdensityof(μ::Dirac, x) = insupport(μ, x) ? 0.0 : -Inf +logdensityof_impl(μ::Dirac, x) = insupport(μ, x) ? 0.0 : -Inf logdensity_def(::Dirac, x::Real) = zero(float(typeof(x))) logdensity_def(::Dirac, x) = 0.0 diff --git a/src/primitives/lebesgue.jl b/src/primitives/lebesgue.jl index 4b8bf7ab..0e6b171a 100644 --- a/src/primitives/lebesgue.jl +++ b/src/primitives/lebesgue.jl @@ -63,12 +63,12 @@ insupport(μ::Lebesgue, x) = x ∈ μ.support insupport(::Lebesgue{RealValues}, ::Real) = true -@inline function logdensityof(μ::Lebesgue, x::Real) +@inline function logdensityof_impl(μ::Lebesgue, x::Real) R = float(typeof(x)) insupport(μ, x) ? zero(R) : R(-Inf) end -@inline logdensityof(μ::Lebesgue, x) = insupport(μ, x) ? 0.0 : -Inf +@inline logdensityof_impl(μ::Lebesgue, x) = insupport(μ, x) ? 0.0 : -Inf massof(::Lebesgue{RealValues}, s::Interval) = width(s) diff --git a/src/standard/stdexponential.jl b/src/standard/stdexponential.jl index c985c224..fda5b53a 100644 --- a/src/standard/stdexponential.jl +++ b/src/standard/stdexponential.jl @@ -4,7 +4,7 @@ export StdExponential insupport(::StdExponential, x) = x ≥ zero(x) -@inline function logdensityof(::StdExponential, x) +@inline function logdensityof_impl(::StdExponential, x) R = float(typeof(x)) x ≥ zero(R) ? convert(R, -x) : R(-Inf) end diff --git a/src/standard/stdlogistic.jl b/src/standard/stdlogistic.jl index 58a1ba67..b28dd618 100644 --- a/src/standard/stdlogistic.jl +++ b/src/standard/stdlogistic.jl @@ -4,7 +4,7 @@ export StdLogistic @inline insupport(d::StdLogistic, x) = true -@inline logdensityof(::StdLogistic, x) = (u = -abs(x); u - 2 * log1pexp(u)) +@inline logdensityof_impl(::StdLogistic, x) = (u = -abs(x); u - 2 * log1pexp(u)) @inline logdensity_def(::StdLogistic, x) = logdensityof(StdLogistic(), x) @inline basemeasure(::StdLogistic) = LebesgueBase() diff --git a/src/standard/stdnormal.jl b/src/standard/stdnormal.jl index 057a8629..f636d311 100644 --- a/src/standard/stdnormal.jl +++ b/src/standard/stdnormal.jl @@ -7,7 +7,7 @@ export StdNormal @inline insupport(::StdNormal, x) = true -@inline logdensityof(::StdNormal, x) = (-x^2 - log2π) / 2 +@inline logdensityof_impl(::StdNormal, x) = (-x^2 - log2π) / 2 @inline logdensity_def(::StdNormal, x) = -x^2 / 2 @inline basemeasure(::StdNormal) = WeightedMeasure(static(-0.5 * log2π), LebesgueBase()) diff --git a/src/standard/stduniform.jl b/src/standard/stduniform.jl index e3702656..d0fc236d 100644 --- a/src/standard/stduniform.jl +++ b/src/standard/stduniform.jl @@ -4,7 +4,7 @@ export StdUniform insupport(::StdUniform, x) = zero(x) ≤ x ≤ one(x) -@inline function logdensityof(::StdUniform, x) +@inline function logdensityof_impl(::StdUniform, x) R = float(typeof(x)) zero(x) ≤ x ≤ one(x) ? zero(R) : R(-Inf) end diff --git a/test/combinators/bind.jl b/test/combinators/bind.jl index 1503f6e2..f7621037 100644 --- a/test/combinators/bind.jl +++ b/test/combinators/bind.jl @@ -81,6 +81,16 @@ using MeasureBase: pushfwd, productmeasure, transport_to, transportmeasure, loca # DOF mismatches must not go unnoticed: @test_throws ArgumentError transport_to(StdUniform()^5, μ)(xy) @test_throws ArgumentError transport_to(StdUniform()^2, μ)(xy) + + # 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])) end @testset "mbind with merge" begin From e4bdb7b1fe47fcbb976d3470e46b60366119e7a5 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 22:23:58 +0200 Subject: [PATCH 57/75] Rework relative density evaluation as type-stable lockstep chain descent The generic three-argument logdensity_def now descends the base measure chains of both measures in lockstep after equalizing their static depths. Since the members of a shared chain suffix have the same depth-from-root on both sides, shared suffixes cancel symbolically. The descent is fully unrolled at compile time and constructs only the base measures it actually visits. Specialized relative densities move from three-argument logdensity_def methods to the new extension point MeasureBase.logdensity_rel_def. The descent checks for an applicable specialization at each visited measure pair; availability is decided purely by dispatch via a sentinel return type instead of method-table introspection, so the checks are free at run time and defining new specializations behaves like any ordinary method definition. This removes basemeasure_sequence-based chain materialization, the commonbase search, schema and the static_hasmethod gate from the relative density code path. Root measure pairs without a specialization now throw an informative exception instead of warning and returning NaN. Created by generative AI. --- src/MeasureBase.jl | 8 +- src/combinators/superpose.jl | 8 +- src/density-core.jl | 163 ++++++++++++++++++++++------------- src/interface.jl | 3 +- src/primitive.jl | 4 +- src/primitives/lebesgue.jl | 4 +- src/schema.jl | 34 -------- src/utils.jl | 30 ------- test/test_basics.jl | 23 +++++ 9 files changed, 142 insertions(+), 135 deletions(-) delete mode 100644 src/schema.jl diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index 9f8037eb..e4a3ce42 100644 --- a/src/MeasureBase.jl +++ b/src/MeasureBase.jl @@ -162,6 +162,13 @@ 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 specialized +`logdensity_def` method for a pair of base measures (in particular at pairs +of identical primitive measures) and any shared chain suffix cancels +symbolically. Measure types may add specialized three-argument methods for +measure pairs whose relative density can be computed more directly. """ function logdensity_def end @@ -175,7 +182,6 @@ include("smf.jl") include("mspace.jl") include("getdof.jl") include("transport.jl") -include("schema.jl") include("proxies.jl") include("kernel.jl") include("parameterized.jl") diff --git a/src/combinators/superpose.jl b/src/combinators/superpose.jl index aa7b6e20..d9b53bd5 100644 --- a/src/combinators/superpose.jl +++ b/src/combinators/superpose.jl @@ -82,7 +82,7 @@ function density_def(s::SuperpositionMeasure, x) end end -@inline function logdensity_def(μ::T, ν::T, x) where {T<:SuperpositionMeasure} +@inline function logdensity_rel_def(μ::T, ν::T, x) where {T<:SuperpositionMeasure} if μ === ν return zero(return_type(logdensity_def, (μ, x))) else @@ -98,12 +98,12 @@ function _superpos_logdensity_rel(s::SuperpositionMeasure, β, x) logsumexp(ds) end -@inline logdensity_def(s::SuperpositionMeasure, β, x) = _superpos_logdensity_rel(s, β, x) +@inline logdensity_rel_def(s::SuperpositionMeasure, β, x) = _superpos_logdensity_rel(s, β, x) -@inline logdensity_def(s::SuperpositionMeasure, β::SuperpositionMeasure, x) = +@inline logdensity_rel_def(s::SuperpositionMeasure, β::SuperpositionMeasure, x) = _superpos_logdensity_rel(s, β, x) -@inline logdensity_def(s, β::SuperpositionMeasure, x) = -_superpos_logdensity_rel(β, s, x) +@inline logdensity_rel_def(s, β::SuperpositionMeasure, x) = -_superpos_logdensity_rel(β, s, x) @inline logdensity_def(s::SuperpositionMeasure, x) = log(density_def(s, x)) diff --git a/src/density-core.jl b/src/density-core.jl index 806a33a6..9a6b0649 100644 --- a/src/density-core.jl +++ b/src/density-core.jl @@ -230,75 +230,118 @@ See also `logdensity_rel`. @inline function unsafe_logdensity_rel(μ::AbstractMeasure, ν::AbstractMeasure, x) μ_local = localmeasure(μ, x) ν_local = localmeasure(ν, x) - # Extra dispatch boundary to reduce number of required specializations of implementation: - return _unsafe_logdensity_rel_local(μ_local, ν_local, x) + return logdensity_def(μ_local, ν_local, x) end -@inline function _unsafe_logdensity_rel_local(μ::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) -end +# Indicates that no specialized method is available to compute the +# log-density between a given pair of measures: +struct _NoLogdensityRel end -# Note that this method assumes `μ` and `ν` to have the same type -function logdensity_def(μ::T, ν::T, x) where {T} - if μ === ν - return zero(logdensity_def(μ, x)) - else - α = basemeasure(μ) - β = basemeasure(ν) - return logdensity_def(μ, x) - logdensity_def(ν, x) + logdensity_rel(α, β, x) - end -end +""" + MeasureBase.logdensity_rel_def(μ, ν, x) -@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ν) +Specialization point for the log-density of `μ` relative to `ν` at `x`. - q = quote - $(Expr(:meta, :inline)) - ℓ = logdensity_def(μs[$M], νs[$N], x) - end +Measure types may add methods for pairs of measure types whose relative +density can be computed directly. The generic implementation of +[`logdensity_def(μ, ν, x)`](@ref logdensity_def) descends the base measure +chains of both measures in lockstep and uses the first specialized +`logdensity_rel_def` method it encounters along the way. - for i in 1:(M-1) - push!(q.args, :(Δℓ = logdensity_def(μs[$i], x))) - # push!(q.args, :(println("Adding", Δℓ))) - push!(q.args, :(ℓ += Δℓ)) - end +Do not call `logdensity_rel_def` directly, call +[`logdensity_rel`](@ref) (or `logdensity_def`) instead. +""" +@inline logdensity_rel_def(μ, ν, x) = _NoLogdensityRel() + +# Generic relative density: descend the base measure chains of both measures +# in lockstep, after equalizing their depths. Since the members of a shared +# chain suffix have the same depth-from-root on both sides, the descent +# terminates at a specialized `logdensity_rel_def` method as soon as one +# becomes applicable (in particular for pairs of identical primitive +# measures), so any shared chain suffix cancels symbolically instead of +# numerically. The descent is fully unrolled at compile time based on the +# static base measure depths, only the base measures actually visited are +# constructed, and whether a specialized method applies at a given level is +# decided purely by dispatch (on the sentinel type `_NoLogdensityRel`). +@inline function logdensity_def(μ, ν, x) + _logdensity_rel_descent(μ, basemeasure_depth(μ), ν, basemeasure_depth(ν), x) +end - for j in 1:(N-1) - push!(q.args, :(Δℓ = logdensity_def(νs[$j], x))) - # push!(q.args, :(println("Subtracting", Δℓ))) - push!(q.args, :(ℓ -= Δℓ)) +@generated function _logdensity_rel_descent( + μ, + ::StaticInteger{M}, + ν, + ::StaticInteger{N}, + x, +) where {M,N} + μsym(i) = Symbol(:μ_, i) + νsym(j) = Symbol(:ν_, j) + prog = Expr(:block, Expr(:meta, :inline), :(μ_0 = μ), :(ν_0 = ν)) + terms = Any[] + n_checks = 0 + # Return via a specialized `logdensity_rel_def` method for the current + # measure pair, if available. Whether one is available is decided purely + # by type, so unsuccessful checks are free at run time: + function emit_check!(i, j) + r = Symbol(:r_, n_checks) + n_checks += 1 + push!(prog.args, :($r = logdensity_rel_def($(μsym(i)), $(νsym(j)), x))) + ret = isempty(terms) ? r : :(+($(terms...), $r)) + push!(prog.args, :(if !($r isa _NoLogdensityRel) + return $ret + end)) + end + i = j = 0 + emit_check!(i, j) + # Equalize depths, accumulating one-sided density terms: + while M - i > N - j + ℓ = Symbol(:ℓμ_, i) + push!(prog.args, :($ℓ = logdensity_def($(μsym(i)), x))) + push!(prog.args, :($(μsym(i + 1)) = basemeasure($(μsym(i))))) + push!(terms, ℓ) + i += 1 + emit_check!(i, j) + end + while N - j > M - i + ℓ = Symbol(:ℓν_, j) + push!(prog.args, :($ℓ = -logdensity_def($(νsym(j)), x))) + push!(prog.args, :($(νsym(j + 1)) = basemeasure($(νsym(j))))) + push!(terms, ℓ) + j += 1 + emit_check!(i, j) + end + # Lockstep descent at equal depth: + for _ in 1:(M-i) + ℓμ, ℓν = Symbol(:ℓμ_, i), Symbol(:ℓν_, j) + push!(prog.args, :($ℓμ = logdensity_def($(μsym(i)), x))) + push!(prog.args, :($ℓν = -logdensity_def($(νsym(j)), x))) + push!(terms, ℓμ, ℓν) + push!(prog.args, :($(μsym(i + 1)) = basemeasure($(μsym(i))))) + push!(prog.args, :($(νsym(j + 1)) = basemeasure($(νsym(j))))) + i += 1 + j += 1 + emit_check!(i, j) end + # Both measures are at root level now: + push!( + prog.args, + :(r_root = _root_logdensity_rel($(μsym(i)), $(νsym(j)), x)), + ) + ret = isempty(terms) ? :r_root : :(+($(terms...), r_root)) + push!(prog.args, :(return $ret)) + return prog +end + +# Root measures of the same type are equal almost everywhere for the +# purpose of pointwise relative densities: +_root_logdensity_rel(μ::M, ν::M, x) where {M} = zero(logdensity_def(μ, x)) - push!(q.args, :(return ℓ)) - return q +function _root_logdensity_rel(@nospecialize(μ), @nospecialize(ν), @nospecialize(x)) + throw( + ArgumentError( + "No method available to compute the log-density between measures with root measures of type $(nameof(typeof(μ))) and $(nameof(typeof(ν)))", + ), + ) end @inline density_rel(μ, ν, x) = exp(logdensity_rel(μ, ν, x)) diff --git a/src/interface.jl b/src/interface.jl index 6003203d..f7f21004 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -5,7 +5,7 @@ using Reexport @reexport using MeasureBase using MeasureBase: basemeasure_depth, proxy, istrue -using MeasureBase: insupport, basemeasure_sequence, commonbase +using MeasureBase: insupport, basemeasure_sequence using MeasureBase: transport_to, NoTransport using DensityInterface: logdensityof @@ -21,7 +21,6 @@ export basemeasure_depth export proxy export insupport export basemeasure_sequence -export commonbase using Test diff --git a/src/primitive.jl b/src/primitive.jl index f43485c2..0e588334 100644 --- a/src/primitive.jl +++ b/src/primitive.jl @@ -11,7 +11,7 @@ measures satisfy the following laws: logdensity_def(μ::PrimitiveMeasure, x) = 0.0 - logdensity_def(μ::M, ν::M, x) where {M<:PrimitiveMeasure} = 0.0 + logdensity_rel_def(μ::M, ν::M, x) where {M<:PrimitiveMeasure} = 0.0 """ abstract type PrimitiveMeasure <: AbstractMeasure end @@ -24,7 +24,7 @@ basemeasure(μ::PrimitiveMeasure) = μ logdensity_def(::PrimitiveMeasure, x) = static(0.0) -logdensity_def(μ::M, ν::M, x) where {M<:PrimitiveMeasure} = 0.0 +logdensity_rel_def(μ::M, ν::M, x) where {M<:PrimitiveMeasure} = 0.0 function Pretty.quoteof(μ::M) where {M<:PrimitiveMeasure} :($M()) diff --git a/src/primitives/lebesgue.jl b/src/primitives/lebesgue.jl index 0e6b171a..7bfdf715 100644 --- a/src/primitives/lebesgue.jl +++ b/src/primitives/lebesgue.jl @@ -12,9 +12,9 @@ insupport(::LebesgueBase, x) = true insupport(::LebesgueBase) = Returns(true) -logdensity_def(::LebesgueBase, ::CountingBase, x) = -Inf +logdensity_rel_def(::LebesgueBase, ::CountingBase, x) = -Inf -logdensity_def(::CountingBase, ::LebesgueBase, x) = Inf +logdensity_rel_def(::CountingBase, ::LebesgueBase, x) = Inf @inline getdof(::LebesgueBase) = static(1) 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/utils.jl b/src/utils.jl index e169c7c1..06792c67 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -93,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...) diff --git a/test/test_basics.jl b/test/test_basics.jl index 11f1a8fe..168bac49 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -222,6 +222,29 @@ 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 From 130d6598ddf156d124fc099396ee95df87ee8af0 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 22:35:17 +0200 Subject: [PATCH 58/75] Make transport from mvstd to vector-marginal products type stable Products over vectors of same-typed unknown-DOF marginals (e.g. binds) now transport from multivariate standard measures via a typed loop instead of accumulating into a Vector{Any}; marginal vectors with abstract element type keep the untyped fallback. Created by generative AI. --- src/combinators/product_transport.jl | 32 +++++++++++++++++++++------- test/combinators/bind.jl | 9 ++++++++ 2 files changed, 33 insertions(+), 8 deletions(-) diff --git a/src/combinators/product_transport.jl b/src/combinators/product_transport.jl index 2d78ae2f..da2d3461 100644 --- a/src/combinators/product_transport.jl +++ b/src/combinators/product_transport.jl @@ -417,15 +417,31 @@ _marginals_from_mvstd_with_rest_nodof(::Tuple{}, ::StdMeasure, x::AbstractVector (), x function _marginals_from_mvstd_with_rest_nodof( - νs::AbstractVector{<:AbstractMeasure}, + νs::AbstractVector{M}, μ_inner::StdMeasure, x::AbstractVector{<:Real}, -) - # ToDo: Check for type stability: - ys = Vector{Any}(undef, length(eachindex(νs))) - x_rest = x - for (i, ν) in zip(eachindex(ys), νs) - ys[i], x_rest = transport_from_mvstd_with_rest(ν, μ_inner, x_rest) +) where {M<:AbstractMeasure} + if isconcretetype(M) + # Marginals of concrete type produce variates of uniform type, so + # the loop below is type stable (the type of the remaining stream + # stays invariant under repeated view-taking): + idxs = eachindex(νs) + y1, x_rest = transport_from_mvstd_with_rest(νs[first(idxs)], μ_inner, x) + ys = Vector{typeof(y1)}(undef, length(idxs)) + ys[begin] = y1 + j = firstindex(ys) + 1 + for i in Iterators.drop(idxs, 1) + ys[j], x_rest = transport_from_mvstd_with_rest(νs[i], μ_inner, x_rest) + j += 1 + end + return ys, x_rest + else + # Fallback for marginals of mixed type: + ys_any = Vector{Any}(undef, length(eachindex(νs))) + x_rest = x + for (i, ν) in zip(eachindex(ys_any), νs) + ys_any[i], x_rest = transport_from_mvstd_with_rest(ν, μ_inner, x_rest) + end + return [y for y in ys_any], x_rest end - return [y for y in ys], x_rest end diff --git a/test/combinators/bind.jl b/test/combinators/bind.jl index f7621037..147ddeb6 100644 --- a/test/combinators/bind.jl +++ b/test/combinators/bind.jl @@ -91,6 +91,15 @@ using MeasureBase: pushfwd, productmeasure, transport_to, transportmeasure, loca # Variates that are too long must not go unnoticed: @test_throws ArgumentError logdensityof(μ, vcat(xy, [0.5])) + + # Products of same-typed unknown-DOF marginals transport type-stably: + P = productmeasure([μ, μ]) + yP = rand(stblrng(), Float64, P) + z = transport_to(StdUniform()^6, P)(yP) + @test z isa AbstractVector{<:Real} && length(z) == 6 + yP_reco, rest = MeasureBase.transport_from_mvstd_with_rest(P, StdUniform(), z) + @test yP_reco isa Vector{<:AbstractVector{Float64}} + @test yP_reco ≈ yP && isempty(rest) end @testset "mbind with merge" begin From 60284898dc735608a89d27465e49d4073311099e Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 22:45:42 +0200 Subject: [PATCH 59/75] Fix method dispatch ambiguities and test for ambiguities Fixes ambiguities between the static-size power density methods and powers of primitive measures (the existing disambiguation method did not cover them), between the massof methods generated by @useproxy and massof over intervals, and in the legacy kernel constructors. Package ambiguity testing (including the Aqua ambiguity check) is now enabled in the test suite. Created by generative AI. --- src/combinators/power.jl | 11 +++++++++-- src/combinators/smart-constructors.jl | 4 ++++ src/proxies.jl | 3 +++ test/test_aqua.jl | 7 +++---- 4 files changed, 19 insertions(+), 6 deletions(-) diff --git a/src/combinators/power.jl b/src/combinators/power.jl index 5145a80c..128771cd 100644 --- a/src/combinators/power.jl +++ b/src/combinators/power.jl @@ -182,12 +182,19 @@ massof(m::PowerMeasure) = massof(m.parent)^prod(m.axes) logdensity_def(::PowerMeasure{P}, x) where {P<:PrimitiveMeasure} = static(0.0) -# To avoid ambiguities +# Disambiguation with the static-size power density methods: function logdensity_def( - ::PowerMeasure{P,<:Tuple{Vararg{StaticOneToLike{0},N}}}, ::Any, + ::PowerMeasure{P,Tuple{<:StaticOneToLike{N}}}, + ::Any, ) where {P<:PrimitiveMeasure,N} static(0.0) end +function logdensity_def( + ::PowerMeasure{P,<:Tuple{Vararg{StaticOneToLike{0}}}}, + ::Any, +) where {P<:PrimitiveMeasure} + static(0.0) +end @inline mspace_elsize(m::PowerMeasure) = axes2size(m.axes) diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl index e2c4cfe3..e035ca1e 100644 --- a/src/combinators/smart-constructors.jl +++ b/src/combinators/smart-constructors.jl @@ -324,6 +324,10 @@ end kernel(::Type{P}, nt::NamedTuple) where {P<:ParameterizedMeasure} = kernel(identity, P, nt) +# Disambiguation: +kernel(::Type{P}, ::NamedTuple{()}) where {P<:ParameterizedMeasure} = + TypedTransitionKernel(constructorof(P), identity) + kernel(::Type{T}; kwargs...) where {T} = kernel(T, NamedTuple(kwargs)) function kernel(::Type{M}, ::NamedTuple{()}) where {M} diff --git a/src/proxies.jl b/src/proxies.jl index 109b9973..f2805176 100644 --- a/src/proxies.jl +++ b/src/proxies.jl @@ -35,6 +35,9 @@ macro useproxy(M) @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/test/test_aqua.jl b/test/test_aqua.jl index d4546ac2..f683e2c1 100644 --- a/test/test_aqua.jl +++ b/test/test_aqua.jl @@ -4,14 +4,13 @@ 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, # Only used by package extensions: stale_deps = (ignore = [:ArgCheck, :ArraysOfArrays],), ) From b3e4e84513f7ab2b4191a76db371092c76f7cc23 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 22:57:41 +0200 Subject: [PATCH 60/75] Remove the legacy TransitionKernel machinery Removes kernel.jl (AbstractTransitionKernel and its subtypes, the kernel/kleisli constructors), the parameterized-measure kernel constructors and the kernel-based productmeasure methods. MKernel and mbind supersede this functionality. The basekernel helper stays, it is independent of the kernel types. Created by generative AI. --- src/MeasureBase.jl | 1 - src/combinators/product.jl | 11 +++ src/combinators/smart-constructors.jl | 85 ------------------- src/kernel.jl | 113 -------------------------- src/parameterized.jl | 21 ----- test/test_basics.jl | 5 -- 6 files changed, 11 insertions(+), 225 deletions(-) delete mode 100644 src/kernel.jl diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index e4a3ce42..2dc427b5 100644 --- a/src/MeasureBase.jl +++ b/src/MeasureBase.jl @@ -183,7 +183,6 @@ include("mspace.jl") include("getdof.jl") include("transport.jl") include("proxies.jl") -include("kernel.jl") include("parameterized.jl") include("domains.jl") include("primitive.jl") diff --git a/src/combinators/product.jl b/src/combinators/product.jl index 9899e0a9..46b4ad3f 100644 --- a/src/combinators/product.jl +++ b/src/combinators/product.jl @@ -163,6 +163,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}, diff --git a/src/combinators/smart-constructors.jl b/src/combinators/smart-constructors.jl index e035ca1e..be69d142 100644 --- a/src/combinators/smart-constructors.jl +++ b/src/combinators/smart-constructors.jl @@ -131,21 +131,6 @@ end @inline _generic_productmeasure_impl(mar::Base.Generator) = ProductMeasure(mar) -# ToDo: Remove or at least refactor this (ProductMeasure shouldn't take a kernel as its argument). - -productmeasure(f, param_maps, pars) = productmeasure(kernel(f, param_maps), pars) - -function productmeasure(k::ParameterizedTransitionKernel, pars) - productmeasure(k.suff, k.param_maps, pars) -end - -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) -end - ############################################################################### # PushforwardMeasure @@ -280,73 +265,3 @@ function weightedmeasure(ℓ, b::WeightedMeasure) weightedmeasure(ℓ + _logweight(b), b.base) end -############################################################################### -# TransitionKernel - -# kernel(Normal(μ=2)) -function kernel(μ::M) where {M<:ParameterizedMeasure} - kernel(M) -end - -function kernel(d::PowerMeasure) - Base.Fix2(powermeasure, d.axes) ∘ kernel(d.parent) -end - -function kernel(f) - T = Core.Compiler.return_type(f, Tuple{Any}) - _kernel(f, T) -end - -function _kernel(f, ::Type{T}) where {T} - GenericTransitionKernel(f) -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] - end - - kernel(params ∘ f, C, NamedTuple{N}(maps)) -end - -kernel(f::F, ::Type{M}; kwargs...) where {F<:Function,M} = kernel(f, M, NamedTuple(kwargs)) - -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 - -kernel(::Type{P}, nt::NamedTuple) where {P<:ParameterizedMeasure} = kernel(identity, P, nt) - -# Disambiguation: -kernel(::Type{P}, ::NamedTuple{()}) where {P<:ParameterizedMeasure} = - TypedTransitionKernel(constructorof(P), identity) - -kernel(::Type{T}; kwargs...) where {T} = kernel(T, NamedTuple(kwargs)) - -function kernel(::Type{M}, ::NamedTuple{()}) where {M} - C = constructorof(M) - TypedTransitionKernel(C, identity) -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)) -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/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/parameterized.jl b/src/parameterized.jl index 8b1c8c88..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) diff --git a/test/test_basics.jl b/test/test_basics.jl index 168bac49..5db95ad7 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -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 From 284c05071a399444055ed350e3893994936af86c Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 23:07:25 +0200 Subject: [PATCH 61/75] Flesh out mspace_elsize and flat-stream variate consumption Adds mspace_elsize methods for standard measures (scalar), Dirac and powers of scalar-variate measures (static where the axes are static), and NamedTuple-product variate names based on the marginal names alone. Flat vector streams now support scalar variates (consuming a single stream element, also in vcat variate splitting) and multi-rank variates (consumed in flattened form and reshaped). This enables densities and transport for binds with scalar-variate primary measures and density evaluation of multi-dimensional powers inside flat streams. Created by generative AI. --- src/collection_utils.jl | 16 ++++++++++++++-- src/combinators/combined.jl | 3 +++ src/combinators/power.jl | 7 +++++-- src/combinators/product.jl | 2 ++ src/primitives/dirac.jl | 2 ++ src/standard/stdmeasure.jl | 2 ++ test/combinators/bind.jl | 10 ++++++++++ test/test_basics.jl | 25 +++++++++++++++++++++++++ 8 files changed, 63 insertions(+), 4 deletions(-) diff --git a/src/collection_utils.jl b/src/collection_utils.jl index 6b9de0d5..a4c5e979 100644 --- a/src/collection_utils.jl +++ b/src/collection_utils.jl @@ -103,10 +103,22 @@ _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`: +# Take the beginning of a flat vector stream as a variate of size `sz`, +# scalar variates have size `()` and multi-rank variates are reshaped: Base.@propagate_inbounds _consume_from_stream(x::AbstractVector, sz::Tuple{IntegerLike}) = _split_after(x, sz[1]) -function _consume_from_stream(x::AbstractVector, @nospecialize(sz::Tuple)) +Base.@propagate_inbounds function _consume_from_stream(x::AbstractVector, ::Tuple{}) + idxs = maybestatic_eachindex(x) + i_first = maybestatic_first(idxs) + x[i_first], _get_or_view(x, i_first + one(i_first), maybestatic_last(idxs)) +end + +function _consume_from_stream(x::AbstractVector, sz::Tuple{Vararg{IntegerLike}}) + a_flat, x_rest = _split_after(x, size2length(sz)) + return maybestatic_reshape(a_flat, sz), x_rest +end + +function _consume_from_stream(x::AbstractVector, @nospecialize(sz)) throw(ArgumentError("Can't consume a variate of size $sz from a flat vector stream")) end diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl index 4325ced8..877e22ee 100644 --- a/src/combinators/combined.jl +++ b/src/combinators/combined.jl @@ -33,6 +33,9 @@ end _split_variate_byvalue(::typeof(vcat), test_a::AbstractVector, ab::AbstractVector) = _split_after(ab, length(test_a)) +_split_variate_byvalue(::typeof(vcat), ::Real, ab::AbstractVector) = + _consume_from_stream(ab, ()) + _split_variate_byvalue(::typeof(vcat), ::NTuple{N,Any}, ab::Tuple) where {N} = _split_after(ab, Val{N}()) diff --git a/src/combinators/power.jl b/src/combinators/power.jl index 128771cd..0ac4f3b6 100644 --- a/src/combinators/power.jl +++ b/src/combinators/power.jl @@ -77,6 +77,11 @@ end marginals(d::PowerMeasure) = maybestatic_fill(d.parent, d.axes) +# Powers of scalar-variate measures have array-valued variates of known size: +@inline mspace_elsize(μ::PowerMeasure) = _pwr_mspace_elsize(μ, mspace_elsize(pwr_base(μ))) +@inline _pwr_mspace_elsize(μ::PowerMeasure, ::Tuple{}) = pwr_size(μ) +@inline _pwr_mspace_elsize(μ::PowerMeasure, ::Any) = NoMSpaceElementSize{typeof(μ)}() + function Base.:^(μ::AbstractMeasure, dims::Tuple{Vararg{AbstractArray,N}}) where {N} powermeasure(μ, dims) end @@ -196,5 +201,3 @@ function logdensity_def( static(0.0) end - -@inline mspace_elsize(m::PowerMeasure) = axes2size(m.axes) diff --git a/src/combinators/product.jl b/src/combinators/product.jl index 46b4ad3f..0e02fdf6 100644 --- a/src/combinators/product.jl +++ b/src/combinators/product.jl @@ -89,6 +89,8 @@ proxy(μ::ProductMeasure{<:FillArrays.Fill}) = mapreduce(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 diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl index ba044c1d..e4e64a50 100644 --- a/src/primitives/dirac.jl +++ b/src/primitives/dirac.jl @@ -43,6 +43,8 @@ insupport(d::Dirac, x) = x == d.x @inline getdof(::Dirac) = static(0) +@inline mspace_elsize(μ::Dirac) = maybestatic_size(μ.x) + @propagate_inbounds function checked_arg(μ::Dirac, x) @boundscheck insupport(μ, x) || throw(ArgumentError("Invalid variate for measure")) x diff --git a/src/standard/stdmeasure.jl b/src/standard/stdmeasure.jl index 81409796..db0e42fe 100644 --- a/src/standard/stdmeasure.jl +++ b/src/standard/stdmeasure.jl @@ -11,6 +11,8 @@ The type of an `N`-dimensional power of a standard measure of type `MU`. """ const StdPowerMeasure{MU<:StdMeasure,N} = PowerMeasure{MU,<:NTuple{N,OneToLike}} +@inline mspace_elsize(::StdMeasure) = () + @inline check_dof(::StdMeasure, ::StdMeasure) = nothing @inline transport_def(::MU, μ::MU, x) where {MU<:StdMeasure} = x diff --git a/test/combinators/bind.jl b/test/combinators/bind.jl index 147ddeb6..1b62e3df 100644 --- a/test/combinators/bind.jl +++ b/test/combinators/bind.jl @@ -82,6 +82,16 @@ using MeasureBase: pushfwd, productmeasure, transport_to, transportmeasure, loca @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) diff --git a/test/test_basics.jl b/test/test_basics.jl index 5db95ad7..8dbdbe4d 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -199,6 +199,31 @@ end end end +@testset "logdensityof_with_rest" begin + StdNormal = MeasureBase.StdNormal + x = [0.3, 0.7, 0.2, 0.9, 0.5] + + # Scalar variates consume one stream element: + @test MeasureBase.mspace_elsize(StdNormal()) == () + ℓ, a, x_rest = MeasureBase.logdensityof_with_rest(StdNormal(), x) + @test a == 0.3 && length(x_rest) == 4 + @test ℓ ≈ logdensityof(StdNormal(), 0.3) + + # Vector variates: + @test MeasureBase.mspace_elsize(StdNormal()^2) == (2,) + ℓ, a, x_rest = MeasureBase.logdensityof_with_rest(StdNormal()^2, x) + @test a == [0.3, 0.7] && length(x_rest) == 3 + @test ℓ ≈ logdensityof(StdNormal()^2, [0.3, 0.7]) + + # Multi-rank variates are consumed in flattened form and reshaped: + @test MeasureBase.mspace_elsize(StdNormal()^(2, 2)) == (2, 2) + ℓ, a, x_rest = MeasureBase.logdensityof_with_rest(StdNormal()^(2, 2), x) + @test a == [0.3 0.2; 0.7 0.9] && length(x_rest) == 1 + @test ℓ ≈ logdensityof(StdNormal()^(2, 2), a) + + @test MeasureBase.mspace_elsize(Dirac([1, 2])) == (2,) +end + @testset "logdensity_rel" begin @test logdensity_rel(Dirac(0.0) + Lebesgue(), Dirac(1.0), 0.0) == Inf @test logdensity_rel(Dirac(0.0) + Lebesgue(), Dirac(1.0), 1.0) == -Inf From 753def4045619bd343b910e3e836996d09b4586a Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 23:07:25 +0200 Subject: [PATCH 62/75] Move product relative densities to logdensity_rel_def Relative densities between product measures were a specialization of logdensity_rel itself, bypassing the support-check layer at the product level. They are now logdensity_rel_def methods (with a type-stable tuple-marginals variant) evaluating marginals via unsafe_logdensity_rel, support checking happens for the products as a whole in logdensity_rel. Created by generative AI. --- src/combinators/product.jl | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/src/combinators/product.jl b/src/combinators/product.jl index 0e02fdf6..22dae865 100644 --- a/src/combinators/product.jl +++ b/src/combinators/product.jl @@ -85,8 +85,20 @@ end proxy(μ::ProductMeasure{<:FillArrays.Fill}) = powermeasure(_fill_value(marginals(μ)), _fill_axes(marginals(μ))) -@inline function logdensity_rel(μ::ProductMeasure, ν::ProductMeasure, x) - mapreduce(logdensity_rel, +, marginals(μ), marginals(ν), x) +# Relative densities between products evaluate marginal-wise. Support +# checks happen at the logdensity_rel level for the whole products, so the +# unsafe marginal evaluation suffices here: +@inline function logdensity_rel_def(μ::ProductMeasure, ν::ProductMeasure, x) + mapreduce(unsafe_logdensity_rel, +, marginals(μ), marginals(ν), x) +end + +# For tuples, `mapreduce` has trouble with type inference: +@inline function logdensity_rel_def( + μ::ProductMeasure{<:Tuple}, + ν::ProductMeasure{<:Tuple}, + x, +) + sum(map(unsafe_logdensity_rel, marginals(μ), marginals(ν), x)) end _mspace_names(μ::ProductMeasure{<:NamedTuple{names}}) where {names} = names From 81f0488031f85826f01f58e9cf207ee6647cf20f Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 23:07:25 +0200 Subject: [PATCH 63/75] Improve failure behavior of interval massof and generic primitive densities The smf-based interval mass now fails with an informative exception for measures without a statistical measure function instead of a MethodError on NoSMF values. The generic non-real-variate log-density of primitive measures returns a static zero instead of false. test_smf now tolerates insupport results that are not booleans (NoFastInsupport) and the logdensity_def docstring points to logdensity_rel_def as the extension point for specialized relative densities. Created by generative AI. --- src/MeasureBase.jl | 11 ++++++----- src/interface.jl | 3 ++- src/mass-interface.jl | 14 +++++++++++++- src/primitive.jl | 2 +- 4 files changed, 22 insertions(+), 8 deletions(-) diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index 2dc427b5..c5aee330 100644 --- a/src/MeasureBase.jl +++ b/src/MeasureBase.jl @@ -164,11 +164,12 @@ 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 specialized -`logdensity_def` method for a pair of base measures (in particular at pairs -of identical primitive measures) and any shared chain suffix cancels -symbolically. Measure types may add specialized three-argument methods for -measure pairs whose relative density can be computed more directly. +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 diff --git a/src/interface.jl b/src/interface.jl index f7f21004..31e68ca0 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -118,7 +118,8 @@ function test_smf(μ, n = 100) @assert issorted(p) x = invsmf.(μ, p) @test issorted(x) - @test all(istrue ∘ insupport(μ), x) + # insupport may return a non-Bool "don't know" (NoFastInsupport): + @test all(x_i -> insupport(μ, x_i) != false, x) @test all((Finv ∘ F).(x) .≈ x) diff --git a/src/mass-interface.jl b/src/mass-interface.jl index 59a5db88..33937498 100644 --- a/src/mass-interface.jl +++ b/src/mass-interface.jl @@ -119,4 +119,16 @@ in this way, users should add the corresponding `massof` method. """ (m::AbstractMeasure)(s) = massof(m, s) -massof(μ, a_b::AbstractInterval) = smf(μ, rightendpoint(a_b)) - smf(μ, leftendpoint(a_b)) +function massof(μ, a_b::AbstractInterval) + _smf_interval_massof(μ, smf(μ, rightendpoint(a_b)), smf(μ, leftendpoint(a_b))) +end + +_smf_interval_massof(μ, smf_r, smf_l) = smf_r - smf_l + +function _smf_interval_massof(μ, ::NoSMF, ::NoSMF) + throw( + ArgumentError( + "Can't compute the mass over an interval for a measure of type $(nameof(typeof(μ))), no statistical measure function available", + ), + ) +end diff --git a/src/primitive.jl b/src/primitive.jl index 0e588334..80f24847 100644 --- a/src/primitive.jl +++ b/src/primitive.jl @@ -20,7 +20,7 @@ basemeasure(μ::PrimitiveMeasure) = μ @inline basemeasure_depth(::PrimitiveMeasure) = static(0) @inline logdensityof_impl(::PrimitiveMeasure, x::Real) = zero(float(typeof(x))) -@inline logdensityof_impl(::PrimitiveMeasure, x) = false +@inline logdensityof_impl(::PrimitiveMeasure, x) = static(0.0) logdensity_def(::PrimitiveMeasure, x) = static(0.0) From 0dec09c36d737808da4d929a3f1e9830d724b226 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 23:37:29 +0200 Subject: [PATCH 64/75] Remove unused firsttype Created by generative AI. --- .../MeasureBaseDistributionsExt.jl | 2 +- ext/MeasureBaseForwardDiffExt.jl | 4 +--- src/utils.jl | 11 ----------- 3 files changed, 2 insertions(+), 15 deletions(-) diff --git a/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl b/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl index 177ec306..a230386a 100644 --- a/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl +++ b/ext/MeasureBaseDistributionsExt/MeasureBaseDistributionsExt.jl @@ -20,7 +20,7 @@ using MeasureBase: getdof, checked_arg, massof using MeasureBase: transport_to, transport_def, transport_origin, from_origin, to_origin using MeasureBase: NoTransportOrigin, NoTransport using MeasureBase: Reshape -using MeasureBase: convert_realtype, firsttype, _fwddiff, @_adignore +using MeasureBase: convert_realtype, _fwddiff, @_adignore import MeasureBase: _dist_params_numtype, _trafo_cdf_impl, _trafo_quantile_impl, _trafo_quantile_impl_generic using MeasureBase: _pushfront, _pushback, _dropfront, _dropback, _rev_cumsum, _exp_cumsum_log diff --git a/ext/MeasureBaseForwardDiffExt.jl b/ext/MeasureBaseForwardDiffExt.jl index 20113c7f..0d91f8b8 100644 --- a/ext/MeasureBaseForwardDiffExt.jl +++ b/ext/MeasureBaseForwardDiffExt.jl @@ -3,7 +3,7 @@ module MeasureBaseForwardDiffExt using MeasureBase -using MeasureBase: containsnan, firsttype +using MeasureBase: containsnan import ForwardDiff function MeasureBase.containsnan(x::ForwardDiff.Dual) @@ -12,7 +12,5 @@ function MeasureBase.containsnan(x::ForwardDiff.Dual) return a || b end -MeasureBase.firsttype(::Type{T}, ::Type{<:ForwardDiff.Dual{tag,<:Real,N}}) where {T<:Real,tag,N} = - ForwardDiff.Dual{tag,T,N} end # module MeasureBaseForwardDiffExt diff --git a/src/utils.jl b/src/utils.jl index 06792c67..6389343f 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -188,17 +188,6 @@ convert_realtype(::Type{T}, x::Union{Tuple,NamedTuple}) where {T<:Real} = convert_realtype(::Type{T}, x::AbstractArray) where {T<:Real} = map(Base.Fix1(convert_realtype, T), x) -""" - MeasureBase.firsttype(::Type{T}, ::Type{U}) where {T<:Real,U<:Real} - -Return the first type, but as a dual number type if the second one is dual. -""" -function firsttype end - -firsttype(::Type{T}, ::Type{U}) where {T<:Real,U<:Real} = T - - - # Distributions implementation hooks: function _trafo_cdf_impl end function _trafo_quantile_impl end From a9032fae327b37e90badf3def947fd75ad9f4f9d Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 23:37:29 +0200 Subject: [PATCH 65/75] Add regression tests for products of measures with value-dependent variate sizes Densities and standard-measure transport for products of hierarchical measures (in vector, tuple and named-tuple marginal form) work through the with-rest machinery; keep it that way. Created by generative AI. --- test/combinators/bind.jl | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/test/combinators/bind.jl b/test/combinators/bind.jl index 1b62e3df..efd2cabd 100644 --- a/test/combinators/bind.jl +++ b/test/combinators/bind.jl @@ -110,6 +110,24 @@ using MeasureBase: pushfwd, productmeasure, transport_to, transportmeasure, loca yP_reco, rest = MeasureBase.transport_from_mvstd_with_rest(P, StdUniform(), z) @test yP_reco isa Vector{<:AbstractVector{Float64}} @test yP_reco ≈ yP && isempty(rest) + @test logdensityof(P, yP) ≈ logdensityof(μ, yP[1]) + logdensityof(μ, yP[2]) + + # Products with value-dependent marginal sizes, in all marginal + # container flavors: + Pt = productmeasure((μ, μ)) + yt = rand(stblrng(), Float64, Pt) + @test logdensityof(Pt, yt) ≈ logdensityof(μ, yt[1]) + logdensityof(μ, yt[2]) + zt = transport_to(StdUniform()^6, Pt)(yt) + yt_reco = transport_to(Pt, StdUniform()^6)(zt) + @test all(map(≈, yt_reco, yt)) + + Pnt = productmeasure((a = StdNormal(), b = μ)) + ynt = rand(stblrng(), Float64, Pnt) + @test logdensityof(Pnt, ynt) ≈ + logdensityof(StdNormal(), ynt.a) + logdensityof(μ, ynt.b) + znt = transport_to(StdUniform()^4, Pnt)(ynt) + ynt_reco = transport_to(Pnt, StdUniform()^4)(znt) + @test ynt_reco.a ≈ ynt.a && ynt_reco.b ≈ ynt.b end @testset "mbind with merge" begin From 9f251d0008a1ae3dab635bb029719513f5490f61 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Mon, 6 Jul 2026 23:37:29 +0200 Subject: [PATCH 66/75] Add combinesets for combining measurable sets MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit combinesets(f_c, α, β) combines two sets along the value combination semantics of mcombine, with specific set representations where possible: cartesian products concatenate under vcat and merge, one-dimensional cartesian powers of equal (singleton-typed) base sets concatenate under vcat, and implicit domains of measures combine into the implicit domain of the combined measure. CombinedMeasure now provides mdomain. Also fixes two latent bugs uncovered by the new tests: setcartprod for NamedTuple sets had an unbound type parameter and never matched, and membership tests for CartesianPower used reversed argument order in Base.in. Created by generative AI. --- src/combinators/combined.jl | 2 + src/domains.jl | 77 +++++++++++++++++++++++++++++++++++-- test/domains.jl | 65 +++++++++++++++++++++++++++++++ test/runtests.jl | 1 + 4 files changed, 141 insertions(+), 4 deletions(-) create mode 100644 test/domains.jl diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl index 877e22ee..bf4e4291 100644 --- a/src/combinators/combined.jl +++ b/src/combinators/combined.jl @@ -127,6 +127,8 @@ end # 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(μ.β)) diff --git a/src/domains.jl b/src/domains.jl index 7458b450..1f3535ef 100644 --- a/src/domains.jl +++ b/src/domains.jl @@ -265,7 +265,8 @@ componentsets(s::CartesianProduct) = s._sets setcartprod(sets::AbstractArray{<:SetLike}) = CartesianProduct(sets) setcartprod(sets::Tuple{Vararg{SetLike}}) = CartesianProduct(sets) -setcartprod(sets::NamedTuple{names,<:Tuple{Vararg{SetLike}}}) = CartesianProduct(sets) +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} = @@ -326,7 +327,7 @@ componentsets(s::CartesianPower) = maybestatic_fill(pwr_base(s), pwr_axes(s)) function Base.in(x::AbstractArray, s::CartesianPower) pwr_size(s) == size(x) || throw(ArgumentError("Size of CartesianPower and given point are incompatible.")) - isempty(x) ? true : all(Base.Fix1(in, pwr_base(s)), x)::Bool + isempty(x) ? true : all(Base.Fix2(in, pwr_base(s)), x)::Bool end Base.isempty(s::CartesianPower) = isempty(pwr_base(s)) || size2length(pwr_size(s)) == 0 @@ -347,11 +348,79 @@ end Represents a combination of two sets. -User code should not create instances of `CombinedMeasure` directly, but should call -[`combinesets(f_c, α, β)`](@ref) instead. +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 Base.in(@nospecialize(x), ::CombinedSet) + throw(ArgumentError("Cannot test if a value lies within a combined set.")) +end + +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 + +_combinesets_cat(f_c, α::SetLike, β::SetLike) = CombinedSet(f_c, α, β) 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/runtests.jl b/test/runtests.jl index 6168091a..9bf40147 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -17,6 +17,7 @@ include("test_basics.jl") include("getdof.jl") include("transport.jl") include("smf.jl") +include("domains.jl") include("test_mooncake.jl") From deeef51d1cfbd3015b72b23700578d31b2fd1d5c Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Fri, 10 Jul 2026 01:23:08 +0200 Subject: [PATCH 67/75] Widen ArraysOfArrays compat to 0.6 and 0.7 Created by generative AI. --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index c0ba40c8..cdd9199c 100644 --- a/Project.toml +++ b/Project.toml @@ -59,7 +59,7 @@ MeasureBaseMooncakeExt = "Mooncake" [compat] ArgCheck = "1, 2" -ArraysOfArrays = "0.6" +ArraysOfArrays = "0.6, 0.7" ChainRulesCore = "1" ChangesOfVariables = "0.1.3" Compat = "3.35, 4" From b0ec47cf13ad4d6a57f3b78f44f9e308c8292f73 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Fri, 10 Jul 2026 01:28:29 +0200 Subject: [PATCH 68/75] Add logdensities for batched multi-point density evaluation logdensities(mu, X) computes the log-density of mu at each point in X, preserving the shape of X. Measure types specialize logdensities_impl. Power measures unwrap into power axes arguments of the internal machinery (ordered outermost first), so implementation methods never dispatch on nested PowerMeasure type signatures. Scalar-variate measures evaluate as a single flat broadcast; power batches with flat variate storage (ArrayOfSimilarArrays) additionally fuse the per-point reduction into a single segmented sum, keeping GPU-backed data on-device. Created by generative AI. --- src/MeasureBase.jl | 4 +- src/density-batched.jl | 199 +++++++++++++++++++++++++++++++++++++++++ test/Project.toml | 1 + test/logdensities.jl | 79 ++++++++++++++++ test/runtests.jl | 1 + 5 files changed, 283 insertions(+), 1 deletion(-) create mode 100644 src/density-batched.jl create mode 100644 test/logdensities.jl diff --git a/src/MeasureBase.jl b/src/MeasureBase.jl index c5aee330..4b3adc4e 100644 --- a/src/MeasureBase.jl +++ b/src/MeasureBase.jl @@ -61,7 +61,8 @@ import HeterogeneousComputing using HeterogeneousComputing: real_numtype using ArraysOfArrays: - VectorOfArrays, VectorOfVectors, VectorOfSimilarArrays, VectorOfSimilarVectors, flatview + ArrayOfSimilarArrays, VectorOfArrays, VectorOfVectors, VectorOfSimilarArrays, + VectorOfSimilarVectors, flatview using OneTwoMany: firstarg, secondarg @@ -204,6 +205,7 @@ include("combinators/weighted.jl") include("combinators/superpose.jl") include("combinators/product.jl") include("combinators/power.jl") +include("density-batched.jl") include("combinators/spikemixture.jl") include("combinators/likelihood.jl") include("combinators/restricted.jl") diff --git a/src/density-batched.jl b/src/density-batched.jl new file mode 100644 index 00000000..e827a179 --- /dev/null +++ b/src/density-batched.jl @@ -0,0 +1,199 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +export logdensities + +""" + logdensities(μ::AbstractMeasure, X::AbstractArray) + +Compute the log-density of `μ` at each point in `X`. + +Returns an array of the shape of `X`, semantically equivalent to +`logdensityof.(Ref(μ), X)`. The computation may be fused across points, +though: power measures with flat variate storage (e.g. based on +`ArraysOfArrays.ArrayOfSimilarArrays`) evaluate as a single flat broadcast +plus a segmented reduction over the underlying flat data, compatible with +GPU-backed storage. + +Measure types should specialize [`MeasureBase.logdensities_impl`](@ref) +instead of `logdensities` itself. +""" +function logdensities end + +@inline function logdensities(μ::AbstractMeasure, X::AbstractArray) + _logdensities(logdensityof_impl, μ, X) +end + +""" + MeasureBase.logdensities_impl(μ::AbstractMeasure, X::AbstractArray) + +Implements [`logdensities(μ, X)`](@ref logdensities) for arrays `X` of +plain `μ`-variates. Power measures never reach `logdensities_impl`, their +power structure is processed generically beforehand. + +Measure types that support fused multi-point evaluation should specialize +`logdensities_impl`. Implementations must preserve the shape of `X` and +must handle points outside the support of `μ` (the result must be `-Inf` +at such points). + +The default implementation broadcasts the log-density over `X` for +measures with scalar variates and falls back to a `map` over `X` +otherwise. +""" +function logdensities_impl end + +function logdensities_impl(μ::AbstractMeasure, X::AbstractArray) + _logdensities_generic(logdensityof_impl, μ, X) +end + +# Batched density machinery, parameterized over the point-level density +# function `f` (`logdensityof_impl` or `logdensity_def`). +# +# `_logdensities(f, μ, X, powers...)` treats each element of `X` as a +# variate of `(μ^pN)^…^p1` for `powers = (p1, …, pN)` (power axes ordered +# outermost first, i.e. in the order in which they are encountered when +# descending into a variate) and returns the density sum for each element, +# preserving the shape of `X`. Power measures are unwrapped into the power +# axes arguments before any other dispatch happens, so implementation +# methods only ever dispatch on plain measure types. + +@inline function _logdensities(f::F, μ, X::AbstractArray, powers::Vararg{Any,N}) where {F,N} + _logdensities_stripped(f, μ, X, powers...) +end + +@inline function _logdensities( + f::F, + μ::PowerMeasure, + X::AbstractArray, + powers::Vararg{Any,N}, +) where {F,N} + _logdensities(f, pwr_base(μ), X, powers..., pwr_axes(μ)) +end + +@inline function _logdensities_stripped(f::F, μ, X::AbstractArray) where {F} + _logdensities_impl(f, μ, X) +end + +function _logdensities_stripped( + f::F, + μ, + X::AbstractArray, + p1, + powers::Vararg{Any,N}, +) where {F,N} + map(x -> _powered_ld(f, μ, x, p1, powers...), X) +end + +function _logdensities_stripped( + f::F, + μ, + X::ArrayOfSimilarArrays{<:Real}, + p1, + powers::Vararg{Any,N}, +) where {F,N} + _logdensities_fused(f, μ, X, mspace_elsize(μ), p1, powers...) +end + +# Absolute densities go through the `logdensities_impl` extension point: +@inline function _logdensities_impl(::typeof(logdensityof_impl), μ, X::AbstractArray) + logdensities_impl(μ, X) +end + +@inline function _logdensities_impl(f::F, μ, X::AbstractArray) where {F} + _logdensities_generic(f, μ, X) +end + +@inline function _logdensities_generic(f::F, μ, X::AbstractArray) where {F} + _logdensities_byelsize(f, μ, X, mspace_elsize(μ)) +end + +# Scalar variates evaluate as a single flat broadcast: +@inline function _logdensities_byelsize( + f::F, + μ, + X::AbstractArray{<:Real}, + ::Tuple{}, +) where {F} + broadcast(Base.Fix1(f, μ), X) +end + +@inline function _logdensities_byelsize(f::F, μ, X::AbstractArray, ::Any) where {F} + map(Base.Fix1(f, μ), X) +end + +# Scalar-variate measure with flat variate storage: evaluate as a single +# flat broadcast followed by a segmented reduction over the per-point +# power structure: +function _logdensities_fused( + f::F, + μ, + X::ArrayOfSimilarArrays{<:Real,M}, + ::Tuple{}, + powers::Vararg{Any,N}, +) where {F,M,N} + sz_inner = _flat_powers_size(powers...) + if length(sz_inner) == M + X_flat = flatview(X) + if ntuple(i -> size(X_flat, i), Val(M)) != sz_inner + throw(ArgumentError("Size of variates doesn't match size of power measure")) + end + ld_flat = broadcast(Base.Fix1(f, μ), X_flat) + reshape(sum(ld_flat, dims = ntuple(identity, Val(M))), size(X)) + else + map(x -> _powered_ld(f, μ, x, powers...), X) + end +end + +function _logdensities_fused( + f::F, + μ, + X::AbstractArray, + ::Any, + p1, + powers::Vararg{Any,N}, +) where {F,N} + map(x -> _powered_ld(f, μ, x, p1, powers...), X) +end + +# The flat size of a variate of `(μ^pN)^…^p1` for a scalar-variate `μ`, +# innermost power axes vary fastest: +@inline _flat_powers_size() = () +@inline function _flat_powers_size(p1, powers::Vararg{Any,N}) where {N} + (_flat_powers_size(powers...)..., axes2size(p1)...) +end + +# Scalar counterpart of `_logdensities`: log-density of `(μ^pN)^…^p1` at a +# single variate `x`. +@inline _powered_ld(f::F, μ, x) where {F} = f(μ, x) + +@inline function _powered_ld( + f::F, + μ::PowerMeasure, + x, + p1, + powers::Vararg{Any,N}, +) where {F,N} + _powered_ld(f, pwr_base(μ), x, p1, powers..., pwr_axes(μ)) +end + +@inline function _powered_ld(f::F, μ, x, p1, powers::Vararg{Any,N}) where {F,N} + if axes2size(p1) != maybestatic_size(x) + throw(ArgumentError("Size of variate doesn't match size of power measure")) + end + R = _powered_ld_type(f, μ, x, powers...) + if isempty(x) + zero(R)::R + else + # Conversion needed since summation can turn static into dynamic values: + convert(R, _powered_ld_sum(f, μ, x, powers...))::R + end +end + +@inline _powered_ld_sum(f::F, μ, x) where {F} = sum(Base.Fix1(f, μ), x) + +@inline function _powered_ld_sum(f::F, μ, x, p1, powers::Vararg{Any,N}) where {F,N} + sum(_logdensities(f, μ, x, p1, powers...)) +end + +@inline function _powered_ld_type(f::F, ::MU, x, powers::Vararg{Any,N}) where {F,MU,N} + Core.Compiler.return_type(_powered_ld, Tuple{F,MU,eltype(x),map(typeof, powers)...}) +end diff --git a/test/Project.toml b/test/Project.toml index 32833feb..19645d48 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -13,6 +13,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" ForwardDiffPullbacks = "450a3b6d-2448-4ee1-8e34-e4eb8713b605" InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112" IrrationalConstants = "92d709cd-6900-40b7-9082-c6be49f344b6" +JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" LogarithmicNumbers = "aa2f6b4e-9042-5d33-9679-40d3a6b85899" diff --git a/test/logdensities.jl b/test/logdensities.jl new file mode 100644 index 00000000..89ff253f --- /dev/null +++ b/test/logdensities.jl @@ -0,0 +1,79 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +using Test + +using MeasureBase +using MeasureBase: logdensities, StdNormal, StdUniform +using ArraysOfArrays: VectorOfSimilarVectors, nestedview, flatview +using IrrationalConstants: log2π +import JLArrays +using JLArrays: JLArray + +stdnormal_ld(x) = -(x^2 + log2π) / 2 + +@testset "logdensities" begin + @testset "scalar variates" begin + X = randn(10) + @test @inferred(logdensities(StdNormal(), X)) ≈ stdnormal_ld.(X) + Xm = randn(2, 3) + @test logdensities(StdNormal(), Xm) ≈ stdnormal_ld.(Xm) + end + + @testset "powers with nested variates" begin + m3 = StdNormal()^3 + X = [randn(3) for _ in 1:10] + @test @inferred(logdensities(m3, X)) ≈ [sum(stdnormal_ld, x) for x in X] + @test only(logdensities(m3, [X[1]])) ≈ logdensityof(m3, X[1]) + + m23 = StdNormal()^(2, 3) + X23 = [randn(2, 3) for _ in 1:5] + @test logdensities(m23, X23) ≈ [sum(stdnormal_ld, x) for x in X23] + + mpp = (StdNormal()^(2, 3))^4 + Xpp = [[randn(2, 3) for _ in 1:4] for _ in 1:6] + @test logdensities(mpp, Xpp) ≈ [sum(x -> sum(stdnormal_ld, x), xs) for xs in Xpp] + end + + @testset "powers with flat variate storage" begin + m3 = StdNormal()^3 + X = VectorOfSimilarVectors(randn(3, 10)) + @test @inferred(logdensities(m3, X)) ≈ + vec(sum(stdnormal_ld.(flatview(X)), dims = 1)) + + # Power structure may be stored flattened out within each point: + mpp = (StdNormal()^(2, 3))^4 + Xpp = nestedview(randn(2, 3, 4, 7), 3) + @test logdensities(mpp, Xpp) ≈ [sum(stdnormal_ld, x) for x in Xpp] + end + + @testset "non-scalar-variate fallback" begin + mprod = productmeasure((StdUniform(), StdNormal())) + X = [(rand(), randn()) for _ in 1:5] + @test logdensities(mprod, X) ≈ logdensityof.(Ref(mprod), X) + end + + @testset "size mismatch" begin + @test_throws ArgumentError logdensities(StdNormal()^3, [randn(3), randn(2)]) + @test_throws ArgumentError logdensities( + StdNormal()^3, + VectorOfSimilarVectors(randn(2, 5)), + ) + end + + @testset "GPU array semantics" begin + JLArrays.allowscalar(false) + + X = JLArray(randn(100)) + ld = logdensities(StdNormal(), X) + @test ld isa JLArray + @test Array(ld) ≈ stdnormal_ld.(Array(X)) + + Xb = VectorOfSimilarVectors(JLArray(randn(3, 50))) + ldb = logdensities(StdNormal()^3, Xb) + @test ldb isa JLArray + @test Array(ldb) ≈ vec(sum(stdnormal_ld.(Array(flatview(Xb))), dims = 1)) + + xj = JLArray(randn(10)) + @test logdensityof(StdNormal()^10, xj) ≈ logdensityof(StdNormal()^10, Array(xj)) + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 9bf40147..b6073d76 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -15,6 +15,7 @@ include("test_standard.jl") include("test_basics.jl") include("getdof.jl") +include("logdensities.jl") include("transport.jl") include("smf.jl") include("domains.jl") From d2ccb547a64e01693c096dedf2abd5b762e665f2 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Fri, 10 Jul 2026 01:51:36 +0200 Subject: [PATCH 69/75] Base power measure density evaluation on the batched machinery logdensityof_impl and logdensity_def for PowerMeasure now unwrap the power structure into the power axes arguments of the batched density machinery. Nested powers with flat variate storage evaluate as a single fused broadcast and segmented reduction, GPU-compatible. Replaces the per-level power density methods, including the static-size specialization and its disambiguation methods, and removes the now unused infer_logdensity_type. Created by generative AI. --- src/combinators/power.jl | 38 +++++++------------------------------- src/utils.jl | 5 ----- 2 files changed, 7 insertions(+), 36 deletions(-) diff --git a/src/combinators/power.jl b/src/combinators/power.jl index 0ac4f3b6..62e6ba4b 100644 --- a/src/combinators/power.jl +++ b/src/combinators/power.jl @@ -102,31 +102,13 @@ params(d::PowerMeasure) = params(first(marginals(d))) basemeasure(d.parent)^d.axes end -for (head, func) in [(:logdensityof_impl, :logdensityof), (:logdensity_def, :logdensity_def)] - @eval @inline function $head(d::PowerMeasure{M}, x) where {M} - parent_m = d.parent - sz_parent = axes2size(d.axes) - sz_x = maybestatic_size(x) - if sz_parent != sz_x - throw(ArgumentError("Size of variate doesn't match size of power measure")) - end - R = infer_logdensity_type($func, parent_m, eltype(x)) - if isempty(x) - return zero(R)::R - else - # Need to convert since sum can turn static into dynamic values: - return convert(R, sum(Base.Fix1($func, parent_m), x))::R - end - end +# Power structure is unwrapped into the power axes arguments of the batched +# density machinery (see density-batched.jl), which fuses evaluation over +# flat variate storage: - @eval @inline function $head( - d::PowerMeasure{<:Any,Tuple{<:StaticOneToLike{N}}}, - x, - ) where {N} - parent = d.parent - sum(1:N) do j - @inbounds $func(parent, x[j]) - end +for head in [:logdensityof_impl, :logdensity_def] + @eval @inline function $head(d::PowerMeasure, x) + _powered_ld($head, pwr_base(d), x, pwr_axes(d)) end @eval @inline function $head( @@ -187,13 +169,7 @@ massof(m::PowerMeasure) = massof(m.parent)^prod(m.axes) logdensity_def(::PowerMeasure{P}, x) where {P<:PrimitiveMeasure} = static(0.0) -# Disambiguation with the static-size power density methods: -function logdensity_def( - ::PowerMeasure{P,Tuple{<:StaticOneToLike{N}}}, - ::Any, -) where {P<:PrimitiveMeasure,N} - static(0.0) -end +# Disambiguation with the static-zero-size power density method: function logdensity_def( ::PowerMeasure{P,<:Tuple{Vararg{StaticOneToLike{0}}}}, ::Any, diff --git a/src/utils.jl b/src/utils.jl index 6389343f..451d5045 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -101,11 +101,6 @@ function infer_zero(f, args...) zero(typeintersect(AbstractFloat, inferred_type)) end -function infer_logdensity_type(f::F, ::M, ::Type{T}) where {F,M,T} - inferred_type = Core.Compiler.return_type(f, Tuple{M,T}) - return inferred_type -end - @inline function allequal(f, x::AbstractArray) val = f(first(x)) @simd for xj in x From 6c0cc298a1a84c7d6b45a718712354aa6787d44f Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Fri, 10 Jul 2026 02:28:26 +0200 Subject: [PATCH 70/75] Rename logdensities_impl to batched_logdensityof_impl Batched protocol functions pair with their scalar counterparts by name: batched_logdensityof_impl implements logdensities, further batched_* functions (with-rest, transport, rand) will follow the same scheme. Created by generative AI. --- src/density-batched.jl | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/src/density-batched.jl b/src/density-batched.jl index e827a179..6a0aac44 100644 --- a/src/density-batched.jl +++ b/src/density-batched.jl @@ -14,7 +14,7 @@ though: power measures with flat variate storage (e.g. based on plus a segmented reduction over the underlying flat data, compatible with GPU-backed storage. -Measure types should specialize [`MeasureBase.logdensities_impl`](@ref) +Measure types should specialize [`MeasureBase.batched_logdensityof_impl`](@ref) instead of `logdensities` itself. """ function logdensities end @@ -24,14 +24,14 @@ function logdensities end end """ - MeasureBase.logdensities_impl(μ::AbstractMeasure, X::AbstractArray) + MeasureBase.batched_logdensityof_impl(μ::AbstractMeasure, X::AbstractArray) Implements [`logdensities(μ, X)`](@ref logdensities) for arrays `X` of -plain `μ`-variates. Power measures never reach `logdensities_impl`, their +plain `μ`-variates. Power measures never reach `batched_logdensityof_impl`, their power structure is processed generically beforehand. Measure types that support fused multi-point evaluation should specialize -`logdensities_impl`. Implementations must preserve the shape of `X` and +`batched_logdensityof_impl`. Implementations must preserve the shape of `X` and must handle points outside the support of `μ` (the result must be `-Inf` at such points). @@ -39,9 +39,9 @@ The default implementation broadcasts the log-density over `X` for measures with scalar variates and falls back to a `map` over `X` otherwise. """ -function logdensities_impl end +function batched_logdensityof_impl end -function logdensities_impl(μ::AbstractMeasure, X::AbstractArray) +function batched_logdensityof_impl(μ::AbstractMeasure, X::AbstractArray) _logdensities_generic(logdensityof_impl, μ, X) end @@ -70,7 +70,7 @@ end end @inline function _logdensities_stripped(f::F, μ, X::AbstractArray) where {F} - _logdensities_impl(f, μ, X) + _batched_logdensityof_impl(f, μ, X) end function _logdensities_stripped( @@ -93,12 +93,12 @@ function _logdensities_stripped( _logdensities_fused(f, μ, X, mspace_elsize(μ), p1, powers...) end -# Absolute densities go through the `logdensities_impl` extension point: -@inline function _logdensities_impl(::typeof(logdensityof_impl), μ, X::AbstractArray) - logdensities_impl(μ, X) +# Absolute densities go through the `batched_logdensityof_impl` extension point: +@inline function _batched_logdensityof_impl(::typeof(logdensityof_impl), μ, X::AbstractArray) + batched_logdensityof_impl(μ, X) end -@inline function _logdensities_impl(f::F, μ, X::AbstractArray) where {F} +@inline function _batched_logdensityof_impl(f::F, μ, X::AbstractArray) where {F} _logdensities_generic(f, μ, X) end From b549bf10ef5e391d37eba74105d2a46874ac2428 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Fri, 10 Jul 2026 03:11:13 +0200 Subject: [PATCH 71/75] Widen Real argument types to Number for traced-value compatibility Reactant traced scalars subtype Number, not Real. Widens primitive leaf density kernels, batched density eltype gates, with-rest stream signatures, log-weight arguments and numtype conversion sources. Real stays where realness is semantic (domain membership, numtype request parameters). Created by generative AI. --- src/collection_utils.jl | 8 ++++---- src/combinators/combined.jl | 2 +- src/combinators/product_transport.jl | 12 ++++++------ src/combinators/transformedmeasure.jl | 2 +- src/combinators/weighted.jl | 2 +- src/density-batched.jl | 6 +++--- src/primitive.jl | 2 +- src/primitives/counting.jl | 2 +- src/primitives/dirac.jl | 4 ++-- src/primitives/lebesgue.jl | 4 ++-- src/utils.jl | 8 ++++---- 11 files changed, 26 insertions(+), 26 deletions(-) diff --git a/src/collection_utils.jl b/src/collection_utils.jl index a4c5e979..5b4d915a 100644 --- a/src/collection_utils.jl +++ b/src/collection_utils.jl @@ -84,12 +84,12 @@ _fill_value(x::FillArrays.Fill) = x.value _fill_axes(x::FillArrays.Fill) = x.axes -_flatten_to_rv(VV::AbstractVector{<:AbstractVector{<:Real}}) = flatview(VectorOfArrays(VV)) -_flatten_to_rv(VV::AbstractVector{<:StaticVector{N,<:Real}}) where {N} = +_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{<:Real}) = flatview(VV) -_flatten_to_rv(VV::VectorOfVectors{<:Real}) = flatview(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...) diff --git a/src/combinators/combined.jl b/src/combinators/combined.jl index bf4e4291..d425ae0a 100644 --- a/src/combinators/combined.jl +++ b/src/combinators/combined.jl @@ -33,7 +33,7 @@ end _split_variate_byvalue(::typeof(vcat), test_a::AbstractVector, ab::AbstractVector) = _split_after(ab, length(test_a)) -_split_variate_byvalue(::typeof(vcat), ::Real, ab::AbstractVector) = +_split_variate_byvalue(::typeof(vcat), ::Number, ab::AbstractVector) = _consume_from_stream(ab, ()) _split_variate_byvalue(::typeof(vcat), ::NTuple{N,Any}, ab::Tuple) where {N} = diff --git a/src/combinators/product_transport.jl b/src/combinators/product_transport.jl index da2d3461..a534300b 100644 --- a/src/combinators/product_transport.jl +++ b/src/combinators/product_transport.jl @@ -372,7 +372,7 @@ end function _split_x_by_marginals_with_rest( dofs::Union{Tuple,AbstractVector}, - x::AbstractVector{<:Real}, + x::AbstractVector{<:Number}, ) x_idxs = maybestatic_eachindex(x) first_idxs = _dof_access_firstidxs(dofs, maybestatic_first(x_idxs)) @@ -385,7 +385,7 @@ function _marginals_from_mvstd_with_rest( νs, dofs::_KnownDOFs, μ_inner::StdMeasure, - x::AbstractVector{<:Real}, + x::AbstractVector{<:Number}, ) xs, x_rest = _split_x_by_marginals_with_rest(dofs, x) μs = map(n -> μ_inner^n, dofs) @@ -397,7 +397,7 @@ function _marginals_from_mvstd_with_rest( νs, dofs, μ_inner::StdMeasure, - x::AbstractVector{<:Real}, + x::AbstractVector{<:Number}, ) _marginals_from_mvstd_with_rest_nodof(νs, μ_inner, x) end @@ -405,7 +405,7 @@ end function _marginals_from_mvstd_with_rest_nodof( νs::Tuple{Vararg{AbstractMeasure}}, μ_inner::StdMeasure, - x::AbstractVector{<:Real}, + x::AbstractVector{<:Number}, ) # ToDo: Check for type stability, may need a generated function: y1, x_rest = transport_from_mvstd_with_rest(νs[1], μ_inner, x) @@ -413,13 +413,13 @@ function _marginals_from_mvstd_with_rest_nodof( return (y1, y2_end...), x_final_rest end -_marginals_from_mvstd_with_rest_nodof(::Tuple{}, ::StdMeasure, x::AbstractVector{<:Real}) = +_marginals_from_mvstd_with_rest_nodof(::Tuple{}, ::StdMeasure, x::AbstractVector{<:Number}) = (), x function _marginals_from_mvstd_with_rest_nodof( νs::AbstractVector{M}, μ_inner::StdMeasure, - x::AbstractVector{<:Real}, + x::AbstractVector{<:Number}, ) where {M<:AbstractMeasure} if isconcretetype(M) # Marginals of concrete type produce variates of uniform type, so diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl index ee960f9a..8c8ff9e9 100644 --- a/src/combinators/transformedmeasure.jl +++ b/src/combinators/transformedmeasure.jl @@ -118,7 +118,7 @@ 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) diff --git a/src/combinators/weighted.jl b/src/combinators/weighted.jl index 124662b6..e70b8e86 100644 --- a/src/combinators/weighted.jl +++ b/src/combinators/weighted.jl @@ -44,7 +44,7 @@ function Base.:*(k::T, m::AbstractMeasure) where {T<:Number} return weightedmeasure(logk, m) end -Base.:*(m::AbstractMeasure, k::Real) = k * m +Base.:*(m::AbstractMeasure, k::Number) = k * m gentype(μ::WeightedMeasure) = gentype(μ.base) diff --git a/src/density-batched.jl b/src/density-batched.jl index 6a0aac44..a29effd4 100644 --- a/src/density-batched.jl +++ b/src/density-batched.jl @@ -86,7 +86,7 @@ end function _logdensities_stripped( f::F, μ, - X::ArrayOfSimilarArrays{<:Real}, + X::ArrayOfSimilarArrays{<:Number}, p1, powers::Vararg{Any,N}, ) where {F,N} @@ -110,7 +110,7 @@ end @inline function _logdensities_byelsize( f::F, μ, - X::AbstractArray{<:Real}, + X::AbstractArray{<:Number}, ::Tuple{}, ) where {F} broadcast(Base.Fix1(f, μ), X) @@ -126,7 +126,7 @@ end function _logdensities_fused( f::F, μ, - X::ArrayOfSimilarArrays{<:Real,M}, + X::ArrayOfSimilarArrays{<:Number,M}, ::Tuple{}, powers::Vararg{Any,N}, ) where {F,M,N} diff --git a/src/primitive.jl b/src/primitive.jl index 80f24847..0a2aba37 100644 --- a/src/primitive.jl +++ b/src/primitive.jl @@ -19,7 +19,7 @@ basemeasure(μ::PrimitiveMeasure) = μ @inline basemeasure_depth(::PrimitiveMeasure) = static(0) -@inline logdensityof_impl(::PrimitiveMeasure, x::Real) = zero(float(typeof(x))) +@inline logdensityof_impl(::PrimitiveMeasure, x::Number) = zero(float(typeof(x))) @inline logdensityof_impl(::PrimitiveMeasure, x) = static(0.0) logdensity_def(::PrimitiveMeasure, x) = static(0.0) diff --git a/src/primitives/counting.jl b/src/primitives/counting.jl index 5a7a998e..776554c5 100644 --- a/src/primitives/counting.jl +++ b/src/primitives/counting.jl @@ -12,7 +12,7 @@ struct Counting{T} <: AbstractMeasure Counting(supp) = new{Core.Typeof(supp)}(supp) end -@inline function logdensityof_impl(μ::Counting, x::Real) +@inline function logdensityof_impl(μ::Counting, x::Number) R = float(typeof(x)) insupport(μ, x) ? zero(R) : R(-Inf) end diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl index e4e64a50..14036dff 100644 --- a/src/primitives/dirac.jl +++ b/src/primitives/dirac.jl @@ -23,14 +23,14 @@ basemeasure(d::Dirac) = CountingBase() massof(::Dirac) = static(1.0) -function logdensityof_impl(μ::Dirac, x::Real) +function logdensityof_impl(μ::Dirac, x::Number) R = float(typeof(x)) insupport(μ, x) ? zero(R) : R(-Inf) end logdensityof_impl(μ::Dirac, x) = insupport(μ, x) ? 0.0 : -Inf -logdensity_def(::Dirac, x::Real) = zero(float(typeof(x))) +logdensity_def(::Dirac, x::Number) = zero(float(typeof(x))) logdensity_def(::Dirac, x) = 0.0 Base.rand(::Random.AbstractRNG, T::Type, μ::Dirac) = μ.x diff --git a/src/primitives/lebesgue.jl b/src/primitives/lebesgue.jl index 7bfdf715..c66536bd 100644 --- a/src/primitives/lebesgue.jl +++ b/src/primitives/lebesgue.jl @@ -18,7 +18,7 @@ logdensity_rel_def(::CountingBase, ::LebesgueBase, x) = Inf @inline getdof(::LebesgueBase) = static(1) -@inline checked_arg(::LebesgueBase, x::Real) = x +@inline checked_arg(::LebesgueBase, x::Number) = x @propagate_inbounds function checked_arg(::LebesgueBase, x::Any) @boundscheck throw(ArgumentError("Invalid variate type for measure")) @@ -63,7 +63,7 @@ insupport(μ::Lebesgue, x) = x ∈ μ.support insupport(::Lebesgue{RealValues}, ::Real) = true -@inline function logdensityof_impl(μ::Lebesgue, x::Real) +@inline function logdensityof_impl(μ::Lebesgue, x::Number) R = float(typeof(x)) insupport(μ, x) ? zero(R) : R(-Inf) end diff --git a/src/utils.jl b/src/utils.jl index 451d5045..b6035e96 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -138,7 +138,7 @@ fcomp(::typeof(identity), g) = g fcomp(f, ::typeof(identity)) = f fcomp(::typeof(identity), ::typeof(identity)) = identity -near_neg_inf(::Type{T}) where {T<:Real} = T(-1E38) # Still fits into Float32 +near_neg_inf(::Type{T}) where {T<:Number} = T(-1E38) # Still fits into Float32 isneginf(x) = isinf(x) && x < zero(x) isposinf(x) = isinf(x) && x > zero(x) @@ -149,7 +149,7 @@ isapproxzero(A::AbstractArray) = all(isapproxzero, A) isapproxone(x::T) where {T<:Real} = x ≈ one(T) isapproxone(A::AbstractArray) = all(isapproxone, A) -containsnan(x::Real) = isnan(x) +containsnan(x::Number) = isnan(x) containsnan(x) = any(containsnan, x) @@ -176,8 +176,8 @@ function convert_realtype end @inline convert_realtype(::Type{T}, x::T) where {T<:Real} = x @inline convert_realtype(::Type{T}, x::AbstractArray{T}) where {T<:Real} = x -@inline convert_realtype(::Type{T}, x::U) where {T<:Real,U<:Real} = T(x) -convert_realtype(::Type{T}, x::AbstractArray{U}) where {T<:Real,U<:Real} = T.(x) +@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} = From 36293fb7d1af19c17f8280e6c6a90d6c8ce569fe Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Fri, 10 Jul 2026 12:40:21 +0200 Subject: [PATCH 72/75] Make density kernels and support checks branch-free Replaces insupport ternaries in the primitive and standard density kernels with _checksupport (ifelse-based) and the value branches in _combine_logd_with_ladj with nested ifelse. StdUniform's insupport uses non-short-circuiting comparisons. Required for traced values (Reactant), where control flow must not depend on runtime values. Created by generative AI. --- src/combinators/transformedmeasure.jl | 19 ++++++++----------- src/primitives/counting.jl | 4 ++-- src/primitives/dirac.jl | 4 ++-- src/primitives/lebesgue.jl | 4 ++-- src/standard/stdexponential.jl | 4 ++-- src/standard/stduniform.jl | 6 +++--- 6 files changed, 19 insertions(+), 22 deletions(-) diff --git a/src/combinators/transformedmeasure.jl b/src/combinators/transformedmeasure.jl index 8c8ff9e9..706784e6 100644 --- a/src/combinators/transformedmeasure.jl +++ b/src/combinators/transformedmeasure.jl @@ -122,17 +122,14 @@ 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_impl( diff --git a/src/primitives/counting.jl b/src/primitives/counting.jl index 776554c5..74e7f8c2 100644 --- a/src/primitives/counting.jl +++ b/src/primitives/counting.jl @@ -14,10 +14,10 @@ end @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_impl(μ::Counting, x) = insupport(μ, x) ? 0.0 : -Inf +@inline logdensityof_impl(μ::Counting, x) = _checksupport(insupport(μ, x), 0.0) @inline logdensity_def(μ::Counting, x) = logdensityof(μ, x) diff --git a/src/primitives/dirac.jl b/src/primitives/dirac.jl index 14036dff..96077696 100644 --- a/src/primitives/dirac.jl +++ b/src/primitives/dirac.jl @@ -25,10 +25,10 @@ massof(::Dirac) = static(1.0) function logdensityof_impl(μ::Dirac, x::Number) R = float(typeof(x)) - insupport(μ, x) ? zero(R) : R(-Inf) + _checksupport(insupport(μ, x), zero(R)) end -logdensityof_impl(μ::Dirac, x) = insupport(μ, x) ? 0.0 : -Inf +logdensityof_impl(μ::Dirac, x) = _checksupport(insupport(μ, x), 0.0) logdensity_def(::Dirac, x::Number) = zero(float(typeof(x))) logdensity_def(::Dirac, x) = 0.0 diff --git a/src/primitives/lebesgue.jl b/src/primitives/lebesgue.jl index c66536bd..770ed03e 100644 --- a/src/primitives/lebesgue.jl +++ b/src/primitives/lebesgue.jl @@ -65,10 +65,10 @@ insupport(::Lebesgue{RealValues}, ::Real) = true @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_impl(μ::Lebesgue, x) = insupport(μ, x) ? 0.0 : -Inf +@inline logdensityof_impl(μ::Lebesgue, x) = _checksupport(insupport(μ, x), 0.0) massof(::Lebesgue{RealValues}, s::Interval) = width(s) diff --git a/src/standard/stdexponential.jl b/src/standard/stdexponential.jl index fda5b53a..1e10bb88 100644 --- a/src/standard/stdexponential.jl +++ b/src/standard/stdexponential.jl @@ -4,9 +4,9 @@ export StdExponential insupport(::StdExponential, x) = x ≥ zero(x) -@inline function logdensityof_impl(::StdExponential, x) +@inline function logdensityof_impl(d::StdExponential, x) R = float(typeof(x)) - x ≥ zero(R) ? convert(R, -x) : R(-Inf) + _checksupport(insupport(d, x), convert(R, -x)) end @inline logdensity_def(::StdExponential, x) = -x diff --git a/src/standard/stduniform.jl b/src/standard/stduniform.jl index d0fc236d..f9b07bf1 100644 --- a/src/standard/stduniform.jl +++ b/src/standard/stduniform.jl @@ -2,11 +2,11 @@ struct StdUniform <: StdMeasure end export StdUniform -insupport(::StdUniform, x) = zero(x) ≤ x ≤ one(x) +insupport(::StdUniform, x) = (zero(x) ≤ x) & (x ≤ one(x)) -@inline function logdensityof_impl(::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) From 71236f81cb21cc2c4d3ac57d3d0a16b7abe55a37 Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Fri, 10 Jul 2026 12:40:21 +0200 Subject: [PATCH 73/75] Add direct logdensityof_impl for weighted measures The weight-shifted density of a support-safe base density is support-safe, so weighted measures need no explicit support check. Removes a redundant per-point insupport sweep over the base measure. Created by generative AI. --- src/combinators/weighted.jl | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/combinators/weighted.jl b/src/combinators/weighted.jl index e70b8e86..033e334f 100644 --- a/src/combinators/weighted.jl +++ b/src/combinators/weighted.jl @@ -16,6 +16,12 @@ _logweight(::AbstractMeasure) = 0 d.logweight end +# The weight-shifted density of a support-safe base density is support-safe, +# no explicit support check required: +@inline function logdensityof_impl(d::AbstractWeightedMeasure, x) + d.logweight + logdensityof_impl(basemeasure(d), x) +end + function Base.rand(rng::AbstractRNG, ::Type{T}, μ::AbstractWeightedMeasure) where {T} rand(rng, T, basemeasure(μ)) end From bf8f7b29d9270f37b35d2cf0ef05899a44446b7a Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Fri, 10 Jul 2026 12:40:21 +0200 Subject: [PATCH 74/75] Make Distributions univariate transport traced-value compatible Widens Real argument types to Number in the cdf/quantile/affine transport machinery and makes the StdUniform transport gateways branch-free (out-of-support results are masked via ifelse, the quantile argument is clamped to keep eager evaluation valid). Distributions with standard-measure or affine transport origins now work with traced values; families that require inverse incomplete beta/gamma functions remain host-only. Created by generative AI. --- ext/MeasureBaseDistributionsExt/univariate.jl | 39 ++++++++----------- 1 file changed, 17 insertions(+), 22 deletions(-) diff --git a/ext/MeasureBaseDistributionsExt/univariate.jl b/ext/MeasureBaseDistributionsExt/univariate.jl index 126fe36d..f607eac3 100644 --- a/ext/MeasureBaseDistributionsExt/univariate.jl +++ b/ext/MeasureBaseDistributionsExt/univariate.jl @@ -12,21 +12,21 @@ _dist_params_numtype(d::Distribution) = real_numtype(typeof(Distributions.params(d))) -@inline _trafo_cdf(d::Distribution{Univariate,Continuous}, x::Real) = +@inline _trafo_cdf(d::Distribution{Univariate,Continuous}, x::Number) = _trafo_cdf_impl(_dist_params_numtype(d), d, x) -@inline _trafo_cdf_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, x::Real) = +@inline _trafo_cdf_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, x::Number) = Distributions.cdf(d, x) -@inline _trafo_quantile(d::Distribution{Univariate,Continuous}, u::Real) = +@inline _trafo_quantile(d::Distribution{Univariate,Continuous}, u::Number) = _trafo_quantile_impl(_dist_params_numtype(d), d, u) -@inline _trafo_quantile_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, u::Real) = +@inline _trafo_quantile_impl(::Type{<:Real}, d::Distribution{Univariate,Continuous}, u::Number) = _trafo_quantile_impl_generic(d, u) -@inline _trafo_quantile_impl_generic(d::Distribution{Univariate,Continuous}, u::Real) = +@inline _trafo_quantile_impl_generic(d::Distribution{Univariate,Continuous}, u::Number) = Distributions.quantile(d, u) # Workaround for Beta dist, current quantile implementation only supports Float64: @@ -50,45 +50,40 @@ end end -@inline function _result_numtype(d::Distribution{Univariate}, x::T) where {T<:Real} +@inline function _result_numtype(d::Distribution{Univariate}, x::T) where {T<:Number} float(promote_type(T, _dist_params_numtype(d))) end @inline function MeasureBase.transport_def(::StdUniform, μ::Distribution{Univariate,Continuous}, x) R = _result_numtype(μ, x) - if Distributions.insupport(μ, x) - y = _trafo_cdf(μ, x) - convert(R, y) - else - convert(R, NaN) - end + y = _trafo_cdf(μ, x) + ifelse(Distributions.insupport(μ, x), convert(R, y), convert(R, NaN)) end @inline function MeasureBase.transport_def(ν::Distribution{Univariate,Continuous}, ::StdUniform, x::T) where {T} R = _result_numtype(ν, x) TF = float(T) - if 0 <= x <= 1 - # Avoid x ≈ 0 and x ≈ 1 to avoid infinite variate values for target distributions with infinite support: - mod_x = ifelse(x ≈ 0, zero(TF) + eps(TF), ifelse(x ≈ 1, one(TF) - eps(TF), convert(TF, x))) - y = _trafo_quantile(ν, mod_x) - convert(R, y) - else - convert(R, NaN) - end + # Avoid x ≈ 0 and x ≈ 1 to avoid infinite variate values for target + # distributions with infinite support, keep the quantile argument valid + # for out-of-range x (the result is masked to NaN then): + clamped_x = clamp(convert(TF, x), zero(TF), one(TF)) + mod_x = ifelse(x ≈ 0, zero(TF) + eps(TF), ifelse(x ≈ 1, one(TF) - eps(TF), clamped_x)) + y = _trafo_quantile(ν, mod_x) + ifelse((zero(x) <= x) & (x <= one(x)), convert(R, y), convert(R, NaN)) end # Use standard measures as transformation origin for scaled/translated equivalents: -function _origin_to_affine(ν::Distribution{Univariate}, y::T) where {T<:Real} +function _origin_to_affine(ν::Distribution{Univariate}, y::T) where {T<:Number} trg_offs, trg_scale = Distributions.location(ν), Distributions.scale(ν) x = muladd(y, trg_scale, trg_offs) convert(_result_numtype(ν, y), x) end -function _affine_to_origin(μ::Distribution{Univariate}, x::T) where {T<:Real} +function _affine_to_origin(μ::Distribution{Univariate}, x::T) where {T<:Number} src_offs, src_scale = Distributions.location(μ), Distributions.scale(μ) y = (x - src_offs) / src_scale convert(_result_numtype(μ, x), y) From 6be425ea0d36c6bb06560782c737f76464fc7deb Mon Sep 17 00:00:00 2001 From: Oliver Schulz Date: Fri, 10 Jul 2026 12:40:21 +0200 Subject: [PATCH 75/75] Add MeasureBaseReactantExt Domain membership methods for traced numbers: Reactant's traced scalars subtype Number, not Real or Integer, so membership in RealValues and IntegerValues is decided by their value type parameter. Created by generative AI. --- Project.toml | 3 +++ ext/MeasureBaseReactantExt.jl | 12 ++++++++++++ 2 files changed, 15 insertions(+) create mode 100644 ext/MeasureBaseReactantExt.jl diff --git a/Project.toml b/Project.toml index cdd9199c..8700a16a 100644 --- a/Project.toml +++ b/Project.toml @@ -43,6 +43,7 @@ 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" @@ -56,6 +57,7 @@ MeasureBaseDistributionsMooncakeExt = ["Distributions", "Mooncake"] MeasureBaseForwardDiffExt = "ForwardDiff" MeasureBaseForwardDiffPullbacksExt = "ForwardDiffPullbacks" MeasureBaseMooncakeExt = "Mooncake" +MeasureBaseReactantExt = "Reactant" [compat] ArgCheck = "1, 2" @@ -87,6 +89,7 @@ PDMats = "0.11" PrettyPrinting = "0.3, 0.4" PropertyFunctions = "0.2.2" Random = "1" +Reactant = "0.2" Reexport = "1" SpecialFunctions = "2.1.4" Static = "0.8, 1" diff --git a/ext/MeasureBaseReactantExt.jl b/ext/MeasureBaseReactantExt.jl new file mode 100644 index 00000000..53bc6b4b --- /dev/null +++ b/ext/MeasureBaseReactantExt.jl @@ -0,0 +1,12 @@ +# This file is a part of MeasureBase.jl, licensed under the MIT License (MIT). + +module MeasureBaseReactantExt + +using Reactant: TracedRNumber +import MeasureBase +using MeasureBase: RealValues, IntegerValues + +Base.in(::TracedRNumber{<:Real}, ::RealValues) = true +Base.in(::TracedRNumber{<:Integer}, ::IntegerValues) = true + +end # module MeasureBaseReactantExt