Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Special case Geometric(OneHalf()) #1934

Merged
merged 3 commits into from
Jan 8, 2025
Merged

Special case Geometric(OneHalf()) #1934

Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
1331649
Special case Geometric(OneHalf())
LilithHafner Dec 19, 2024
44bbde7
Remove OneHalf and branch at sample time
LilithHafner Dec 20, 2024
55be7f6
Use 1//2 instead of 0.5
LilithHafner Dec 20, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion src/univariate/discrete/geometric.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,9 @@ function Geometric(p::Real; check_args::Bool=true)
return Geometric{typeof(p)}(p)
end

Geometric() = Geometric{Float64}(0.5)
struct OneHalf <: Real end
Geometric() = Geometric{OneHalf}(OneHalf())
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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).

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Base.getproperty(d::Geometric{OneHalf}, s::Symbol) = s == :p ? 0.5 : getfield(d, s)

@distr_support Geometric 0 Inf

Expand Down Expand Up @@ -137,6 +139,7 @@ cf(d::Geometric, t::Real) = laplace_transform(d, -t*im)
### Sampling

rand(rng::AbstractRNG, d::Geometric) = floor(Int,-randexp(rng) / log1p(-d.p))
rand(rng::AbstractRNG, ::Geometric{OneHalf}) = leading_zeros(rand(rng, UInt))

### Model Fitting

Expand Down
Loading