Unify sample transformation and improve rand for measures

Author: oschulz

src/BAT.jl

@@ -67,7 +67,7 @@ using Accessors: @set
 
 import HeterogeneousComputing
 using HeterogeneousComputing: AbstractComputeUnit, CPUnit
-using HeterogeneousComputing: GenContext, get_rng, get_precision, get_compute_unit, get_gencontext
+using HeterogeneousComputing: GenContext, get_rng, get_precision, get_compute_unit, get_gencontext, allocate_array
 
 import MeasureBase
 using MeasureBase: AbstractMeasure, DensityMeasure, Likelihood

src/algotypes/transform_algorithm.jl

@@ -277,7 +277,7 @@ end
 struct SampleTransformation <: TransformAlgorithm end
 
 function bat_transform_impl(f::Function, smpls::DensitySampleVector, ::SampleTransformation, context::BATContext)
-    (result = broadcast_arbitrary_trafo(f, smpls), f_transform = f)
+    (result = transform_samples(f, smpls), f_transform = f)
 end
 
 function bat_transform_impl(shp::AbstractValueShape, smpls::DensitySampleVector, ::SampleTransformation, context::BATContext)

src/initvals/initvals.jl

@@ -73,7 +73,7 @@ end
 
 function bat_initval_impl(target::BATPushFwdMeasure, n::Integer, algorithm::InitFromTarget, context::BATContext)
     vs_orig = bat_initval_impl(target.orig, n, algorithm, context).result
-    vs = BAT.broadcast_trafo(gettransform(target), vs_orig)
+    vs = BAT.transform_samples(gettransform(target), vs_orig)
     (result = vs,)
 end
 
@@ -164,7 +164,7 @@ end
 
 
 function apply_trafo_to_init(f_transform::Function, initalg::ExplicitInit)
-    xs_tr = broadcast_trafo(f_transform, initalg.xs)
+    xs_tr = transform_samples(f_transform, initalg.xs)
     ExplicitInit(xs_tr)
 end
 

src/measures/bat_pushfwd_measure.jl

@@ -135,9 +135,15 @@ function checked_logdensityof(m::BATPushFwdMeasure{F,I,M,KeepRootMeasure}, v::An
 end
 
 
-Random.rand(rng::AbstractRNG, ::Type{T}, m::BATPushFwdMeasure) where {T<:Real} = m.f(rand(rng, T, m.origin))
-
-Random.rand(rng::AbstractRNG, m::BATPushFwdMeasure) = m.f(rand(rng, m.origin))
+Random.rand(gen::GenContext, m::BATPushFwdMeasure) = m.f(rand(gen, m.origin))
+
+function Base.rand(gen::GenContext, m::BATPwrMeasure{<:BATPushFwdMeasure})
+    m_nonpwr, sz = m.parent, m.sz
+    f = m_nonpwr.f
+    m_origin = m_nonpwr.origin ^ sz
+    X_origin = rand(gen, m_origin)
+    return transform_samples(f, X_origin)
+end
 
 supports_rand(m::BATPushFwdMeasure) = supports_rand(m.origin)
 

src/measures/bat_pwr_measure.jl

@@ -34,6 +34,7 @@ function _cartidxs(axs::Tuple{Vararg{AbstractUnitRange,N}}) where {N}
     CartesianIndices(map(_dynamic, axs))
 end
 
+
 function Base.rand(gen::GenContext, m::BATPwrMeasure)
     cunit = get_compute_unit(gen)
     rng = get_rng(gen)
@@ -47,12 +48,6 @@ function Base.rand(gen::GenContext, m::BATPwrMeasure{<:BATDistMeasure})
     gen_adapt(gen, reshaped_X)
 end
 
-function Base.rand(gen::GenContext, m::BATPwrMeasure{<:BATDistMeasure})
-    X = rand(get_rng(gen), m.parent.dist, size(marginals(m))...)
-    reshaped_X = _reshape_rand_n_output(X)
-    gen_adapt(gen, reshaped_X)
-end
-
 function Base.rand(gen::GenContext, m::BATPwrMeasure{<:DensitySampleMeasure})
     # Always generate R on CPU for now:
     R = rand(get_rng(gen), size(marginals(m))...)
@@ -75,7 +70,7 @@ MeasureBase.marginals(m::BATPwrMeasure) = Fill(_pwr_base(m), _pwr_size(m))
     @assert size(x) == _pwr_size(m)
     m_base = _pwr_base(m)
     sum(x) do x_i
-        logdensity_def(m_base, x_i)
+        logdensityof(m_base, x_i)
     end
 end
 

src/measures/bat_weighted_measure.jl

@@ -50,6 +50,13 @@ end
 
 Base.rand(gen::GenContext, m::BATWeightedMeasure) = rand(gen, m.base)
 
+function Base.rand(gen::GenContext, m::BATPwrMeasure{<:BATWeightedMeasure})
+    m_nonpwr, sz = m.parent, m.sz
+    m_origin = m_nonpwr.base ^ sz
+    X_origin = rand(gen, m_origin)
+    return X_origin
+end
+
 supports_rand(m::BATWeightedMeasure) = supports_rand(m.origin)
 
 

src/transforms/distribution_transform.jl

@@ -112,12 +112,6 @@ struct DistributionTransform{
 end
 
 
-# ToDo: Unify with broadcast_trafo
-function broadcast_arbitrary_trafo(f_transform::DistributionTransform, smpls::DensitySampleVector)
-    broadcast_trafo(f_transform, smpls)
-end
-
-
 function _distrafo_ctor_impl(target_dist::DT, source_dist::DF) where {DT<:ContinuousDistribution,DF<:ContinuousDistribution}
     @argcheck eff_totalndof(target_dist) == eff_totalndof(source_dist)
     DistributionTransform{DT,DF}(target_dist, source_dist)
@@ -177,7 +171,7 @@ function Base.Broadcast.broadcasted(
     f_transform::DistributionTransform,
     v_src::Union{ArrayOfSimilarVectors{<:Real},ShapedAsNTArray}
 )
-    broadcast_trafo(f_transform, v_src)
+    transform_samples(f_transform, v_src)
 end
 
 
@@ -192,7 +186,7 @@ function Base.Broadcast.broadcasted(
     f_transform::DistributionTransform,
     s_src::DensitySampleVector
 )
-    broadcast_trafo(f_transform, s_src)
+    transform_samples(f_transform, s_src)
 end