Special case Geometric(OneHalf())

[LilithHafner]

Before

julia> @b Geometric() mean
2.097 ns

julia> @b Geometric() rand
10.762 ns

julia> @b Geometric() var
2.102 ns

julia> @b Geometric() entropy
8.543 ns

julia> @b Geometric() rand(_, 1000)
9.264 μs (3 allocs: 7.875 KiB)

1331649

julia> @b Geometric() mean
1.199 ns

julia> @b Geometric() rand
2.409 ns

julia> @b Geometric() var
1.195 ns

julia> @b Geometric() entropy
1.197 ns

julia> @b Geometric() rand(_, 1000)
1.118 μs (3 allocs: 7.875 KiB)

After

julia> @b Geometric() mean
1.990 ns

julia> @b Geometric() rand
3.516 ns

julia> @b Geometric() var
2.047 ns

julia> @b Geometric() entropy
8.339 ns

julia> @b Geometric() rand(_, 1000)
1.193 μs (3 allocs: 7.875 KiB)

In the PR, rand uses a substantially different (and simpler) algorithm proposed in #1933 for the p=0.5 case. mean, var, and entropy constant propagate. This latter impact is more of a convenient result than an intended design as those methods should never be bottlenecks for Geometric()

[devmotion]

The common approach in Distributions is to branch depending on the values of the parameters if a faster or more accurate sampler exists. That is, one would just check if p == 1//2 or something similar inside of rand and sampler (if it is specialized).

[LilithHafner]

There's some runtime cost there but if you'd prefer runtime branching over compile time branching I can switch the implementation. Here's an idea of the runtime cost:

julia> using Random, Chairmarks

julia> current(p) = floor(Int,-randexp() / log1p(-p))
current (generic function with 1 method)

julia> pr() = leading_zeros(rand(UInt))
pr (generic function with 1 method)

julia> proposed(p) = p == 1//2 ? pr() : current(p)
proposed (generic function with 1 method)

julia> @b .2 current,proposed
(8.764 ns, 9.627 ns)

julia> @b .5 current,proposed
(8.582 ns, 2.584 ns)

julia> @b pr
2.280 ns

[devmotion]

I think the runtime cost is fine (overhead of ~0.3 ns it seems?). It's still a significant improvement and it would be consistent with existing samplers and rand implementations in Distributions which would also be a bit simpler for users I assume.

[LilithHafner]

Done. The runtime comparison was pretty substantial in bulk generation but inlining allows the compiler to hoist it, giving another 3x speedup:

julia> using Distributions

julia> @b Geometric() rand(_, 100)
891.906 ns (2 allocs: 928 bytes)

julia> @eval Distributions function rand(rng::AbstractRNG, d::Geometric)
           if d.p == 0.5
               leading_zeros(rand(rng, UInt)) # This branch is a performance optimization
           else
               floor(Int,-randexp(rng) / log1p(-d.p))
           end
       end
rand (generic function with 181 methods)

julia> @b Geometric() rand(_, 100)
342.896 ns (2 allocs: 928 bytes)

julia> @eval Distributions @inline function rand(rng::AbstractRNG, d::Geometric)
           if d.p == 0.5
               leading_zeros(rand(rng, UInt)) # This branch is a performance optimization
           else
               floor(Int,-randexp(rng) / log1p(-d.p))
           end
       end
rand (generic function with 181 methods)

julia> @b Geometric() rand(_, 100)
110.390 ns (2 allocs: 928 bytes)

(@v1.11) pkg> st Distributions
Status `~/.julia/environments/v1.11/Project.toml`
  [31c24e10] Distributions v0.25.115

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 86.02%. Comparing base (ceb6343) to head (55be7f6).

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #1934   +/-   ##
=======================================
  Coverage   86.01%   86.02%           
=======================================
  Files         144      144           
  Lines        8696     8699    +3     
=======================================
+ Hits         7480     7483    +3     
  Misses       1216     1216           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.