Improve pushforward implementation

Rename TransformVolCorr and subtypes (with backward compatibility).

Author: oschulz

src/MeasureBase.jl

@@ -21,6 +21,7 @@ using DensityInterface: FuncDensity, LogFuncDensity
 using DensityInterface
 
 using InverseFunctions
+using InverseFunctions: FunctionWithInverse
 using ChangesOfVariables
 
 import Base.iterate

src/combinators/transformedmeasure.jl

@@ -1,6 +1,46 @@
-# TODO: Compare with ChangesOfVariables.jl
+"""
+    abstract type PushFwdStyle
+
+Provides the behavior of a measure's [`rootmeasure`](@ref) under a
+pushforward. Either [`AdaptRootMeasure()`](@ref) or
+[`PushfwdRootMeasure()`](@ref)
+"""
+abstract type PushFwdStyle end
+export PushFwdStyle
+
+const TransformVolCorr = PushFwdStyle
+
+"""
+    AdaptRootMeasure()
 
-using InverseFunctions: FunctionWithInverse
+Indicates that when applying a pushforward to a measure, it's
+[`rootmeasure`](@ref) not not be pushed forward. Instead, the root measure
+should be kept just "reshaped" to the new measurable space if necessary.
+
+Density calculations for pushforward measures constructed with
+`AdaptRootMeasure()` will take take the volume element of variate
+transform (typically via the log-abs-det-Jacobian of the transform) into
+account.
+"""
+struct AdaptRootMeasure <: TransformVolCorr end
+export AdaptRootMeasure
+
+const WithVolCorr = AdaptRootMeasure
+
+"""
+    PushfwdRootMeasure()
+
+Indicates than when applying a pushforward to a measure, it's
+[`rootmeasure`](@ref) should be pushed forward with the same function.
+
+Density calculations for pushforward measures constructed with
+`PushfwdRootMeasure()` will ignore the volume element of the variate
+transform.
+"""
+struct PushfwdRootMeasure <: TransformVolCorr end
+export PushfwdRootMeasure
+
+const NoVolCorr = PushfwdRootMeasure
 
 abstract type AbstractTransformedMeasure <: AbstractMeasure end
 
@@ -19,23 +59,37 @@ function parent(::AbstractTransformedMeasure) end
 export PushforwardMeasure
 
 """
-    struct PushforwardMeasure{F,I,M,VC<:TransformVolCorr} <: AbstractPushforward
+    struct PushforwardMeasure{F,I,M,S<:PushFwdStyle} <: AbstractPushforward
         f :: F
         finv :: I
         origin :: M
-        volcorr :: VC
+        style :: S
     end
 
     Users should not call `PushforwardMeasure` directly. Instead call or add
     methods to `pushfwd`.
 """
-struct PushforwardMeasure{F,I,M,VC<:TransformVolCorr} <: AbstractPushforward
+struct PushforwardMeasure{F,I,M,S<:PushFwdStyle} <: AbstractPushforward
     f::F
     finv::I
     origin::M
-    volcorr::VC
+    style::S
+
+    function PushforwardMeasure{F,I,M,S}(f::F, finv::I, origin::M, style::S) where {F,I,M,S<:PushFwdStyle}
+        new{F,I,M,S}(f, finv, origin, style)
+    end
+
+    function PushforwardMeasure(f, finv, origin::M, style::S) where {M,S<:PushFwdStyle}
+        new{Core.Typeof(f),Core.Typeof(finv),M,S}(f, finv, origin, style)
+    end
 end
 
+const _NonBijectivePusfwdMeasure{M<:PushforwardMeasure,S<:PushFwdStyle} = Union{
+    PushforwardMeasure{<:Any,<:NoInverse,M,S},
+    PushforwardMeasure{<:NoInverse,<:Any,M,S},
+    PushforwardMeasure{<:NoInverse,<:NoInverse,M,S},
+}
+
 gettransform(ν::PushforwardMeasure) = ν.f
 parent(ν::PushforwardMeasure) = ν.origin
 
@@ -45,55 +99,94 @@ end
 
 # TODO: THIS IS ALMOST CERTAINLY WRONG 
 # @inline function logdensity_rel(
-#     ν::PushforwardMeasure{FF1,IF1,M1,<:WithVolCorr},
-#     β::PushforwardMeasure{FF2,IF2,M2,<:WithVolCorr},
+#     ν::PushforwardMeasure{FF1,IF1,M1,<:AdaptRootMeasure},
+#     β::PushforwardMeasure{FF2,IF2,M2,<:AdaptRootMeasure},
 #     y,
 # ) where {FF1,IF1,M1,FF2,IF2,M2}
 #     x = β.inv_f(y)
 #     f = ν.inv_f ∘ β.f
 #     inv_f = β.inv_f ∘ ν.f
-#     logdensity_rel(pushfwd(f, inv_f, ν.origin, WithVolCorr()), β.origin, x)
+#     logdensity_rel(pushfwd(f, inv_f, ν.origin, AdaptRootMeasure()), β.origin, x)
 # end
 
+# TODO: Would profit from custom pullback:
+function _combine_logd_with_ladj(logd_orig::Real, ladj::Real)
+    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
+end
+
+function logdensityof(
+    @nospecialize(μ::_NonBijectivePusfwdMeasure{M,<:PushfwdRootMeasure}),
+    @nospecialize(v::Any)
+) where {M}
+    throw(
+        ArgumentError(
+            "Can't calculate densities for non-bijective pushforward measure $(nameof(M))",
+        ),
+    )
+end
+
+function logdensityof(
+    @nospecialize(μ::_NonBijectivePusfwdMeasure{M,<:AdaptRootMeasure}),
+    @nospecialize(v::Any)
+) where {M}
+    throw(
+        ArgumentError(
+            "Can't calculate densities for non-bijective pushforward measure $(nameof(M))",
+        ),
+    )
+end
+
 for func in [:logdensityof, :logdensity_def]
-    @eval @inline function $func(ν::PushforwardMeasure{F,I,M,<:WithVolCorr}, y) where {F,I,M}
-        f = ν.f
-        finv = ν.finv
-        x_orig, inv_ladj = with_logabsdet_jacobian(unwrap(finv), y)
-        logd_orig = $func(ν.origin, x_orig)
-        logd = float(logd_orig + inv_ladj)
-        neginf = oftype(logd, -Inf)
-        return ifelse(
-            # Zero density wins against infinite volume:
-            (isnan(logd) && logd_orig == -Inf && inv_ladj == +Inf) ||
-            # Maybe  also for (logd_orig == -Inf) && isfinite(inv_ladj) ?
-            # Return constant -Inf to prevent problems with ForwardDiff:
-            (isfinite(logd_orig) && (inv_ladj == -Inf)),
-            neginf,
-            logd,
-        )
+    @eval function $func(ν::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 @inline function $func(ν::PushforwardMeasure{F,I,M,<:NoVolCorr}, y) where {F,I,M}
-        x = ν.finv(y)
-        return $func(ν.origin, x)
+    @eval function $func(ν::PushforwardMeasure{F,I,M,<:PushfwdRootMeasure}, y) where {F,I,M}
+        f_inv = unwrap(ν.finv)
+        x = f_inv(y)
+        logd_orig = $func(ν.origin, x)
+        return logd_orig
     end
 end
 
-insupport(ν::PushforwardMeasure, y) = insupport(ν.origin, ν.finv(y))
+insupport(m::PushforwardMeasure, x) = insupport(transport_origin(m), to_origin(m, x))
 
 function testvalue(::Type{T}, ν::PushforwardMeasure) where {T}
     ν.f(testvalue(T, parent(ν)))
 end
 
 @inline function basemeasure(ν::PushforwardMeasure)
-    pushfwd(ν.f, basemeasure(parent(ν)), NoVolCorr())
+    pushfwd(ν.f, basemeasure(parent(ν)), PushfwdRootMeasure())
+end
+
+function rootmeasure(m::PushforwardMeasure{F,I,M,PushfwdRootMeasure}) where {F,I,M}
+    pushfwd(m.f, rootmeasure(m.origin))
+end
+function rootmeasure(m::PushforwardMeasure{F,I,M,AdaptRootMeasure}) where {F,I,M}
+    rootmeasure(m.origin)
 end
 
 _pushfwd_dof(::Type{MU}, ::Type, dof) where {MU} = NoDOF{MU}()
 _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)}()
 
 # Bypass `checked_arg`, would require potentially costly transformation:
 @inline checked_arg(::PushforwardMeasure, x) = x
@@ -102,47 +195,53 @@ _pushfwd_dof(::Type{MU}, ::Type{<:Tuple{Any,Real}}, dof) where {MU} = dof
 @inline from_origin(ν::PushforwardMeasure, x) = ν.f(x)
 @inline to_origin(ν::PushforwardMeasure, y) = ν.finv(y)
 
+massof(m::PushforwardMeasure) = massof(transport_origin(m))
+
 function Base.rand(rng::AbstractRNG, ::Type{T}, ν::PushforwardMeasure) where {T}
-    return ν.f(rand(rng, T, parent(ν)))
+    return ν.f(rand(rng, T, ν.origin))
 end
 
 ###############################################################################
 # pushfwd
 
-export pushfwd
-
 """
-    pushfwd(f, μ, volcorr = WithVolCorr())
+    pushfwd(f, μ, style = AdaptRootMeasure())
 
 Return the [pushforward
 measure](https://en.wikipedia.org/wiki/Pushforward_measure) from `μ` the
 [measurable function](https://en.wikipedia.org/wiki/Measurable_function) `f`.
 
 To manually specify an inverse, call 
-`pushfwd(InverseFunctions.setinverse(f, finv), μ, volcorr)`.
+`pushfwd(InverseFunctions.setinverse(f, finv), μ, style)`.
 """
-function pushfwd(f, μ, volcorr::TransformVolCorr = WithVolCorr())
-    PushforwardMeasure(f, inverse(f), μ, volcorr)
-end
-
-function pushfwd(f, μ::PushforwardMeasure, volcorr::TransformVolCorr = WithVolCorr())
-    _pushfwd_of_pushfwd(f, μ, μ.volcorr, volcorr)
-end
+function pushfwd end
+export pushfwd
 
-# Either both WithVolCorr or both NoVolCorr, so we can merge them
-function _pushfwd_of_pushfwd(f, μ::PushforwardMeasure, ::V, v::V) where {V}
-    pushfwd(fchain((μ.f, f)), μ.origin, v)
+@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)
+
+_pushfwd_impl(f, μ, style) = PushforwardMeasure(f, inverse(f), μ, style)
+
+function _pushfwd_impl(
+    f,
+    μ::PushforwardMeasure{F,I,M,S},
+    style::S,
+) where {F,I,M,S<:PushFwdStyle}
+    orig_μ = μ.origin
+    new_f = fcomp(f, μ.f)
+    new_f_inv = fcomp(μ.finv, inverse(f))
+    PushforwardMeasure(new_f, new_f_inv, orig_μ, style)
 end
 
-function _pushfwd_of_pushfwd(f, μ::PushforwardMeasure, _, v)
-    PushforwardMeasure(f, inverse(f), μ, v)
-end
+_pushfwd_impl(::typeof(identity), μ, ::AdaptRootMeasure) = μ
+_pushfwd_impl(::typeof(identity), μ, ::PushfwdRootMeasure) = μ
 
 ###############################################################################
 # pullback
 
 """
-    pullback(f, μ, volcorr = WithVolCorr())
+    pullbck(f, μ, style = AdaptRootMeasure())
 
 A _pullback_ is a dual concept to a _pushforward_. While a pushforward needs a
 map _from_ the support of a measure, a pullback requires a map _into_ the
@@ -154,8 +253,17 @@ in terms of the inverse function; the "forward" function is not used at all. In
 some cases, we may be focusing on log-density (and not, for example, sampling).
 
 To manually specify an inverse, call 
-`pullback(InverseFunctions.setinverse(f, finv), μ, volcorr)`.
+`pullbck(InverseFunctions.setinverse(f, finv), μ, style)`.
 """
-function pullback(f, μ, volcorr::TransformVolCorr = WithVolCorr())
-    pushfwd(setinverse(inverse(f), f), μ, volcorr)
+function pullback end
+export pullback
+
+@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)
+
+function _pullback_impl(f, μ, style = AdaptRootMeasure())
+    pushfwd(setinverse(inverse(f), f), μ, style)
 end
+
+@deprecate pullback(f, μ, style::PushFwdStyle = AdaptRootMeasure()) pullbck(f, μ, style)

src/transport.jl

@@ -274,30 +274,3 @@ function Base.show(io::IO, f::TransportFunction)
 end
 
 Base.show(io::IO, M::MIME"text/plain", f::TransportFunction) = show(io, f)
-
-"""
-    abstract type TransformVolCorr
-
-Provides control over density correction by transform volume element.
-Either [`NoVolCorr()`](@ref) or [`WithVolCorr()`](@ref)
-"""
-abstract type TransformVolCorr end
-
-"""
-    NoVolCorr()
-
-Indicate that density calculations should ignore the volume element of
-variate transformations. Should only be used in special cases in which
-the volume element has already been taken into account in a different
-way.
-"""
-struct NoVolCorr <: TransformVolCorr end
-
-"""
-    WithVolCorr()
-
-Indicate that density calculations should take the volume element of
-variate transformations into account (typically via the
-log-abs-det-Jacobian of the transform).
-"""
-struct WithVolCorr <: TransformVolCorr end

src/utils.jl

@@ -164,3 +164,21 @@ using InverseFunctions: FunctionWithInverse
 
 unwrap(f) = f
 unwrap(f::FunctionWithInverse) = f.f
+
+
+fcomp(f, g) = fchain(g, f)
+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
+
+isneginf(x) = isinf(x) && x < 0 
+isposinf(x) = isinf(x) && x > 0 
+
+isapproxzero(x::T) where T<:Real = x ≈ zero(T)
+isapproxzero(A::AbstractArray) = all(isapproxzero, A)
+
+isapproxone(x::T) where T<:Real = x ≈ one(T)
+isapproxone(A::AbstractArray) = all(isapproxone, A)