Add InplaceLogpdf wrapper type

docs/src/index.md

@@ -113,4 +113,5 @@ BayesBase.distribution_typewrapper
 BayesBase.CountingReal
 BayesBase.Infinity
 BayesBase.MinusInfinity
+BayesBase.InplaceLogpdf
 ```

src/statsfuns.jl

@@ -326,3 +326,72 @@
 function Base.:(==)(left::CountingReal{T}, right::CountingReal{T}) where {T}
  return (value(left) == value(right)) && (infinities(left) == infinities(right))
 end
+
+"""
+ InplaceLogpdf(logpdf!)
+
+Wraps a `logpdf!` function in a type that can later on be used for dispatch. 
+The sole purpose of this wrapper type is to allow for in-place logpdf operation on a batch of samples.
+Accepts a function `logpdf!` that takes two arguments: `out` and `sample` and writes the logpdf of the sample to the `out` array.
+A regular `logpdf` function can be converted to `logpdf!` by using `convert(InplaceLogpdf, logpdf)`.
+
+```jldoctest
+julia> using Distributions, BayesBase
+
+julia> d = Beta(2, 3);
+
+julia> inplace = convert(BayesBase.InplaceLogpdf, (sample) -> logpdf(d, sample));
+
+julia> out = zeros(9);
+
+julia> inplace(out, 0.1:0.1:0.9)
+9-element Vector{Float64}:
+ -0.028399474521697776
+ 0.42918163472548043
+ 0.5675839575845996
+ 0.5469646703818638
+ 0.4054651081081646
+ 0.14149956227369964
+ -0.2797139028026039
+ -0.9571127263944104
+ -2.2256240518579173
+```
+
+```jldoctest
+julia> using Distributions, BayesBase
+
+julia> d = Beta(2, 3);
+
+julia> inplace = BayesBase.InplaceLogpdf((out, sample) -> logpdf!(out, d, sample));
+
+julia> out = zeros(9);
+
+julia> inplace(out, 0.1:0.1:0.9)
+9-element Vector{Float64}:
+ -0.028399474521697776
+ 0.42918163472548043
+ 0.5675839575845996
+ 0.5469646703818638
+ 0.4054651081081646
+ 0.14149956227369964
+ -0.2797139028026039
+ -0.9571127263944104
+ -2.2256240518579173
+```
+"""
+struct InplaceLogpdf{F}
+ logpdf!::F
+end
+
+function (inplace::InplaceLogpdf)(out, x)
+ inplace.logpdf!(out, x)
+ return out
+end
+
+function Base.convert(::Type{InplaceLogpdf}, something)
+ return InplaceLogpdf((out, x) -> map!(something, out, x))
+end
+
+function Base.convert(::Type{InplaceLogpdf}, inplace::InplaceLogpdf)
+ return inplace
+end

test/statsfuns_tests.jl

@@ -26,10 +26,10 @@ end
  end
 end
 
-@testitem "dtanh" begin 
+@testitem "dtanh" begin
  for T in (Float32, Float64, BigFloat)
  foreach(rand(T, 10)) do number
- @test dtanh(number) ≈ 1 - tanh(number) ^ 2
+ @test dtanh(number) ≈ 1 - tanh(number)^2
  end
  end
 end
@@ -87,6 +87,62 @@ end
 
  @test float(convert(CountingReal, r)) ≈ zero(T)
  @test float(convert(CountingReal{Float64}, r)) ≈ zero(Float64)
+ end
+end
+
+@testitem "InplaceLogpdf" begin
+ import BayesBase: InplaceLogpdf
+ using Distributions, LinearAlgebra, StableRNGs
+
+ @testset "Vector based samples" begin
+ distribution = Beta(10, 10)
+ fn = (x) -> logpdf(distribution, x)
+ inplacefn = convert(InplaceLogpdf, fn)
+
+ @test fn !== inplacefn
+
+ rng = StableRNG(42)
+ samples = rand(rng, distribution, 100)
+ evaluated = map(fn, samples)
+
+ container = similar(evaluated)
+ inplacefn(container, samples)
+
+ @test evaluated == container
+ end
+
+ @testset "Matrix based samples" begin
+ distribution = MvNormal(ones(2), ones(2))
+ fn = (x) -> logpdf(distribution, x)
+ inplacefn = convert(InplaceLogpdf, fn)
+
+ @test inplacefn !== fn
+
+ rng = StableRNG(42)
+ samples = rand(rng, distribution, 100)
+ evaluated = map(fn, eachcol(samples))
+
+ container = similar(evaluated)
+ inplacefn(container, eachcol(samples))
+
+ @test evaluated == container
+ end
+
+ @testset "Do not convert already inplace version" begin
+ distribution = MvNormal(ones(2), ones(2))
+ fn = InplaceLogpdf((out, x) -> logpdf!(out, distribution, x))
+ inplacefn = convert(InplaceLogpdf, fn)
+
+ @test inplacefn === fn
+
+ rng = StableRNG(42)
+ samples = rand(rng, distribution, 100)
+ evaluated = zeros(100)
+ fn(evaluated, eachcol(samples))
+
+ container = similar(evaluated)
+ inplacefn(container, eachcol(samples))
 
+ @test evaluated == container
  end
 end