From b1b2601f61369b5bbaa242e38a1007385d6db01d Mon Sep 17 00:00:00 2001 From: Eliassj Date: Tue, 31 Mar 2026 23:42:05 +0200 Subject: [PATCH 1/8] Add tests --- Project.toml | 12 +- src/AdaptiveRejectionSampling.jl | 730 +++++++++++++++++++++---------- test/runtests.jl | 265 ++++++++--- 3 files changed, 708 insertions(+), 299 deletions(-) diff --git a/Project.toml b/Project.toml index 3e30ade..e33cf41 100644 --- a/Project.toml +++ b/Project.toml @@ -3,14 +3,20 @@ uuid = "c75e803d-635f-53bd-ab7d-544e482d8c75" version = "0.2.1" [deps] +Compat = "34da2185-b29b-5c13-b0c7-acf172513d20" +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" [compat] -ForwardDiff = "0.10.1" -StatsBase = "0.24, 0.34" -julia = "0.7.0, 1" +Compat = "4.18.1" +DifferentiationInterface = "0.7" +ForwardDiff = "1" +SpecialFunctions = "2" +StatsBase = "0.34.10" +julia = "1" [extras] Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/src/AdaptiveRejectionSampling.jl b/src/AdaptiveRejectionSampling.jl index 386f0f8..c295272 100644 --- a/src/AdaptiveRejectionSampling.jl +++ b/src/AdaptiveRejectionSampling.jl @@ -1,284 +1,556 @@ - """ A log conconcave function is majorized with a piecewise envelop, which on the original scale is piecewise exponential. As the resulting extremely precise envelop adapts, the rejection rate dramatically decreases. """ module AdaptiveRejectionSampling -# ------------------------------ -using Random # Random stdlib -# ------------------------------ -import ForwardDiff: derivative -import StatsBase: sample, weights -# ------------------------------ -export Line, Objective, Envelop, RejectionSampler # Structures/classes -export run_sampler!, sample_envelop, eval_envelop, add_segment! # Methods -# ------------------------------ - -""" - Line(slope::Float64, intercept::Float64) -Basic ensamble-unit for an envelop. -""" -mutable struct Line - slope::Float64 - intercept::Float64 + +import SpecialFunctions: loggamma +using StatsBase +import Random: default_rng, AbstractRNG +using DifferentiationInterface +import ForwardDiff +using Compat + +const DEFAULT_MAX_SEGMENTS::Int = 25 + + +""" + AllocFreeWeights{S <: Real, T <: Number, V <: AbstractVector{T}} <: AbstractWeights{S, T, V} + +Non-mutable weights used in the `sample` function in this package in order to avoid allocations. +""" +struct AllocFreeWeights{S <: Real, T <: Number, V <: AbstractVector{T}} <: AbstractWeights{S, T, V} + values::V + sum::S + + function AllocFreeWeights(ws::V, sum::S) where {S <: Real, T <: Number, V <: AbstractVector{T}} + return new{S, T, V}(ws, sum) + end +end + +function AllocFreeWeights(ws::V) where {T <: Number, V <: AbstractVector{T}} + return AllocFreeWeights(ws, sum(ws)) +end + + +struct Objective{T <: Number, F <: Function, G <: Function} + "The (log-concave) function defining the density `f(x) -> y`" + f::F + "Its gradient `grad(x) -> y'`" + grad::G + + @doc """ + Objective(f::Function, grad::Function, ::Type{T} = Float64) where {T <: Number} + + Create an `Objective` directly defined by its function `f` and custom gradient `grad`. + The parameter `T` represents the type of the expected input and is utilized when preparing gradients using autodiff. + By default `T = Float64`. + + + !!! warning + Observe that `f` should be in its + log-concave form and that no checks are performed in order to verify this. + """ + function Objective(f::Function, grad::Function, ::Type{T} = Float64) where {T <: Number} + return new{T, typeof(f), typeof(grad)}(f, grad) + end end """ - intersection(l1::Line, l2::Line) -Finds the horizontal coordinate of the intersection between lines + Objective(f::Function, init = one(Float64); adbackend = AutoForwardDiff()) + +Create an `Objective` for a function automatically generating its gradient. +Gradients are calculated using `DifferentiationInterface.jl` using the backend of choice +with `ForwardDiff.jl` being the default. In order to prepare the gradient an initial value is +required. By default this is `one(Float64)`. If a gradient for a different type is +desired, it should be specified through this initial value. + +If a custom/manual gradient is available it may instead be provided. + +!!! warning + Observe that `f` should be in its + log-concave form and that no checks are performed in order to verify this. + +# Example + +```julia +# Create an objective using `Mooncake.jl` autodiff and Float32 as its type. +ARS.Objective(somefun, one(Float32); adbackend = AutoMooncake(; config=nothing)) +# Create an objective using the defaults `AutoDiff.jl` autodiff and Float64 as its type. +ARS.Objective(somefun) +``` """ -function intersection(l1::Line, l2::Line) - @assert l1.slope != l2.slope "slopes should be different" - -(l2.intercept - l1.intercept) / (l2.slope - l1.slope) +function Objective(f::Function, init = one(Float64); adbackend = AutoForwardDiff()) + gradprep = prepare_gradient(f, adbackend, init) + obj = Objective( + f, + x -> gradient(f, gradprep, adbackend, x), + typeof(init) + ) + return obj end +abstract type AbstractHull end + """ - exp_integral(l::Line, x1::Float64, x2::Float64) -Computes the integral - ``LaTeX \\int_{x_1} ^ {x_2} \\exp\\{ax + b\\} dx. `` -The resulting value is the weight assigned to the segment [x1, x2] in the envelop -""" -@inline function exp_integral(l::Line, x1::Float64, x2::Float64) - a, b = l.slope, l.intercept - v1, v2 = a*x1, a*x2 - if v1 > 25.0 || v2 > 25.0 || a == 0.0 || b > 25.0 - @warn "exp_integral: numerical instability, truncating, check for under/overflow, consider truncating logf" - v1 = min(v1, 25.0) - v2 = min(v2, 25.0) - b = min(b, 25.0) - if a == 0 - a = (v2 - v1) * 1e-6 - end + struct UpperHull{T} <: AbstractHull + intercepts::Vector{T} + slopes::Vector{T} + intersections::Vector{T} + abscissae::Vector{T} + segment_weights::Vector{T} + domain::Tuple{T, T} end - exp(b) * (exp(v2) - exp(v1)) / a + +Upper hull of the enveloping function. +""" +struct UpperHull{T} <: AbstractHull + intercepts::Vector{T} + slopes::Vector{T} + intersections::Vector{T} + abscissae::Vector{T} + segment_weights::Vector{T} + domain::Tuple{T, T} end +Base.broadcastable(h::UpperHull) = Ref(h) + +slopes(h::UpperHull) = h.slopes +intercepts(h::UpperHull) = h.intercepts +intersections(h::UpperHull) = h.intersections +abscissae(h::UpperHull) = h.abscissae +line(h::UpperHull, i::Integer) = (h.slopes[i], h.intercepts[i]) +lineinds(h::UpperHull) = 1:(length(slopes(h))) +n_lines(h::UpperHull) = length(slopes(h)) +segment_weights(h::UpperHull) = h.segment_weights + """ - Envelop(lines::Vector{Line}, support::Tuple{Float64, Float64}) -A piecewise linear function with k segments defined by the lines `L_1, ..., L_k` and cutpoints -`c_1, ..., c_k+1` with `c1 = support[1]` and `c2 = support[2]`. A line L_k is active in the segment -[c_k, c_k+1], and it's assigned a weight w_k based on [exp_integral](@exp_integral). The weighted integral -over c_1 to c_k+1 is one, so that the envelop is interpreted as a density. -""" -mutable struct Envelop - lines::Vector{Line} - cutpoints::AbstractVector{Float64} - weights::AbstractVector{Float64} - size::Int - - Envelop(lines::Vector{Line}, support::Tuple{Float64,Float64}) = begin - @assert issorted([l.slope for l in lines], rev=true) "line slopes must be decreasing" - intersections = [intersection(lines[i], lines[i+1]) for i in 1:(length(lines)-1)] - cutpoints = [support[1]; intersections; support[2]] - @assert issorted(cutpoints) "cutpoints must be ordered" - @assert length(unique(cutpoints)) == length(cutpoints) "cutpoints can't have duplicates" - weights = [exp_integral(l, cutpoints[i], cutpoints[i+1]) for (i, l) in enumerate(lines)] - @assert Inf ∉ weights "Overflow in assigning weights" - new(lines, cutpoints, weights, length(lines)) - end + eval_hull(h::UpperHull, x) + +Eval hull `h` at `x`. +""" +function eval_hull(h::UpperHull, x) + i = searchsortedfirst(intersections(h), x) - 1 + sl, int = line(h, i) + return sl * x + int end """ - add_segment!(e::Envelop, l::Line) -Adds a new line segment to an envelop based on the value of its slope (slopes must be decreasing -always in the envelop). The cutpoints are automatically determined by intersecting the line with -the adjacent lines. + intersection(slope1::T, intercept1::T, slope2::T, intercept2::T) where {T} + +Returns the intersection abscissa between 2 lines as defined by their slopes and intercepts. Returns NaN if the lines are paralell. """ -function add_segment!(e::Envelop, l::Line) - # Find the position in sorted array with binary search - pos = searchsortedfirst([-line.slope for line in e.lines], -l.slope) - # Find the new cutpoints - if pos == 1 - new_cut = intersection(l, e.lines[pos]) - # Insert in second position, first one is the support bound - insert!(e.cutpoints, pos + 1, new_cut) - elseif pos == e.size + 1 - new_cut = intersection(l, e.lines[pos-1]) - insert!(e.cutpoints, pos, new_cut) - else - new_cut1 = intersection(l, e.lines[pos-1]) - new_cut2 = intersection(l, e.lines[pos]) - splice!(e.cutpoints, pos, [new_cut1, new_cut2]) - @assert issorted(e.cutpoints) "incompatible line: resulting intersection points aren't sorted" +function intersection(slope1::T, intercept1::T, slope2::T, intercept2::T) where {T} + return (intercept2 - intercept1) / (slope1 - slope2) +end + +""" + calc_intersects!(out::AbstractVector{T}, slopes::AbstractVector{T}, intercepts::AbstractVector{T}) where {T} + +Calculate intersections of a series of lines defined by their `slopes` and `intercepts`, storing the result in `out`. +""" +function calc_intersects!(out::AbstractVector{T}, slopes::AbstractVector{T}, intercepts::AbstractVector{T}) where {T} + for i in eachindex(out) + out[i] = intersection(slopes[i], intercepts[i], slopes[i + 1], intercepts[i + 1]) end - # Insert the new line - insert!(e.lines, pos, l) - e.size += 1 - # Recompute weights (this could be done more efficiently in the future by updating the neccesary ones only) - e.weights = [exp_integral(line, e.cutpoints[i], e.cutpoints[i+1]) for (i, line) in enumerate(e.lines)] + return nothing end """ - sample_envelop(rng::AbstractRNG, e::Envelop) - sample_envelop(e::Envelop) -Samples an element from the density defined by the envelop `e` with it's exponential weights. -See [`Envelop`](@Envelop) for details. + calc_intersects(slopes::AbstractVector{T}, intercepts::AbstractVector{T}) where {T} + +Calculate intersections of a series of lines defined by their `slopes` and `intercepts`. """ -function sample_envelop(rng::AbstractRNG, e::Envelop) - # Randomly select lines based on envelop weights - i = sample(rng, 1:e.size, weights(e.weights)) - a, b = e.lines[i].slope, e.lines[i].intercept - # Use the inverse CDF method for sampling - log(exp(-b) * rand(rng) * e.weights[i] * a + exp(a * e.cutpoints[i])) / a +function calc_intersects(slopes::AbstractVector{T}, intercepts::AbstractVector{T}) where {T} + # return [ + # intersection(slopes[i], intercepts[i], slopes[i + 1], intercepts[i + 1]) + # for i in 1:(length(slopes) - 1) + # ] + + out = Vector{T}(undef, length(slopes) - 1) + calc_intersects!(out, slopes, intercepts) + return out end -function sample_envelop(e::Envelop) - sample_envelop(Random.GLOBAL_RNG, e) + +""" + calc_slopes_and_intercepts(obj::Objective, abscissae::AbstractVector{T}) where {T} + +Calculates slopes and intercepts of the lines tangential to the objective function at `abscissae`. +""" +function calc_slopes_and_intercepts(obj::Objective, abscissae::AbstractVector{T}) where {T} + issorted(abscissae) || + throw(ArgumentError("`abscissae` should be sorted in ascending order.")) + slopes = [obj.grad(abscissae[i]) for i in eachindex(abscissae)] + intercepts = [ + slopes[i] * (-abscissae[i]) + obj.f(abscissae[i]) + for i in eachindex(abscissae) + ] + return slopes, intercepts end """ - eval_envelop(e::Envelop, x::Float64) -Eval point a point `x` in the piecewise linear function defined by `e`. Necessary for evaluating -the density assigned to the point `x`. + exp_integral_line(slope, intercept, x1, x2) + +Calculate the integral of `exp(slope * x + intercept)` between `x1` and `x2`. """ -function eval_envelop(e::Envelop, x::Float64) - # searchsortedfirst is the proper method for and ordered list - pos = searchsortedfirst(e.cutpoints, x) - if pos == 1 || pos == length(e.cutpoints) + 1 - return 0.0 +function exp_integral_line(slope, intercept, x1, x2) + return if !iszero(slope) + (exp(slope * x2 + intercept) - exp(slope * x1 + intercept)) / slope else - a, b = e.lines[pos-1].slope, e.lines[pos-1].intercept - return exp(a * x + b) + (x2 - x1) * exp(intercept) end end -# -------------------------------- - """ - Objective(logf::Function, support:) - Objective(logf::Function, grad::Function) -Convenient structure to store the objective function to be sampled. It must receive the logarithm of -f and not f directly. It uses automatic differentiation by default, but the user can provide the derivative optionally. + UpperHull(obj::Objective, abscissae::Vector{T}, domain::Tuple{T, T}) where {T} + +Create an `UpperHull` based on the objective function, initial abscissae and its domain. """ -struct Objective - logf::Function - grad::Function - Objective(logf::Function) = begin - # Automatic differentiation - grad(x) = derivative(logf, x) - new(logf, grad) +function UpperHull(obj::Objective{T}, abscissae::Vector{T}, domain::Tuple{T, T}) where {T} + slopes, intercepts = calc_slopes_and_intercepts(obj, abscissae) + + intersects = [domain[1]; calc_intersects(slopes, intercepts); domain[2]] + + wgts = Vector{T}(undef, length(slopes)) + for i in 1:length(wgts) + wgts[i] = exp_integral_line(slopes[i], intercepts[i], intersects[i], intersects[i + 1]) end - Objective(logf::Function, grad::Function) = new(logf, grad) + + return UpperHull(intercepts, slopes, intersects, abscissae, wgts, domain) end +# Calculate the inverse CDF of a segment, handling a slope of zero. """ -RejectionSampler(f::Function, support::Tuple{Float64, Float64}, init::Tuple{Float64, Float64}) - RejectionSampler(f::Function, support::Tuple{Float64, Float64}[ ,δ::Float64]) -An adaptive rejection sampler to obtain iid samples from a logconcave function supported in -`support = (support[1], support[2])`. f can either be the either probability density of the -function to be sampled, or its logarithm. For the latter, use the keyword argument `log=true`. -The functions can be unnormalized in the sense that the probability density can be specified up to a constant. -The adaptive rejection samplings algorithm requires two initial points `init = init[1], init[2]` -such that (d/dx)logp(init[1]) > 0 and (d/dx)logp(init[2]) < 0. These points can be provided directly -(typically, any point left and right of the mode will do). It is also possibe to specify and search -range and delta for a greedy search of the initial points. -The `support` must be of the form `(-Inf, Inf), (-Inf, a), (b, Inf), (a,b)`, and it represent the -interval in which f has positive value, and zero elsewhere. - -The alternative constructor uses a search_range, a value δ for the distance between points in the search, -and min/max slope values in absolute terms. - -## Keyword arguments -- `max_segments::Int = 10` : max size of envelop, the rejection-rate is usually slow with a small number of segments -- `max_failed_factor::Float64 = 0.001`: level at which throw an error if one single sample has a rejection rate - exceeding this value -- `logdensity::Bool = false`: indicator fo whether `f` is the log of the probability density, up to a normalization constant. -- `search_range::Tuple{Float64,Float64} = (-10.0, 10.0)`: range in which to search for initial points -- `min_slope::Float64 = 1e-6`: minimum slope in absolute value of logf at the initial/found points -- `max_slope::Float64 = Inf: maximum slope in absolute value of logf at the initial/found points -""" -struct RejectionSampler - objective::Objective - envelop::Envelop - max_segments::Int - max_failed_rate::Float64 - # Constructor when initial points are provided - RejectionSampler( - f::Function, - support::Tuple{Float64,Float64}, - init::Tuple{Float64,Float64}; - max_segments::Int=25, - logdensity::Bool=false, - max_failed_rate::Float64=0.001, - ) = begin - @assert support[1] < support[2] "invalid support, not an interval" - if logdensity - objective = Objective(f) - else - objective = Objective(x -> log(f(x))) - end - x1, x2 = init - @assert x1 < x2 "cutpoints must be ordered" - a1, a2 = objective.grad(x1), objective.grad(x2) - @assert 0.0 < a1 "logf must have positive slope at initial cutpoint 1" - @assert a2 < 0.0 "logf must have negative slope at initial cutpoint 2" - b1, b2 = objective.logf(x1) - a1 * x1, objective.logf(x2) - a2 * x2 - line1, line2 = Line(a1, b1), Line(a2, b2) - envelop = Envelop([line1, line2], support) - new(objective, envelop, max_segments, max_failed_rate) + inv_cdf_seg(slope, intercept, r, w, intersect, intersect2) + +Calculate the inverse cumulative density of a hull segment at `r` based on its... + +- `slope` +- `intercept` +- `w` (weight) +- `intersect`, `intersect2`, its intersections with previous/next segments +""" +@inline function inv_cdf_seg(slope, intercept, r, w, intersect, intersect2) + return if !iszero(slope) + log(exp(-intercept) * r * w * slope + exp(slope * intersect)) / slope + else + r * (intersect2 - intersect) + intersect end -"" - RejectionSampler( - f::Function, - support::Tuple{Float64,Float64}, - δ::Float64=0.5; - search_range::Tuple{Float64,Float64}=(-10.0, 10.0), - min_slope::Float64=1e-6, - max_slope::Float64=10.0, - logdensity::Bool=false, - kwargs... - ) = begin - if logdensity - logf = f - else - logf = (x -> log(f(x))) - end - grad(x) = derivative(logf, x) - grid_lims = max(search_range[1], support[1]), min(search_range[2], support[2]) - grid = grid_lims[1]:δ:grid_lims[2] - grads = grad.(grid) - i1 = findfirst(min_slope .< grads .< max_slope) - i2 = findlast(min_slope .< - grads .< max_slope) - @assert (i1 !== nothing) && (i2 !== nothing) "couldn't find initial points, please provide them or change `search_range`" - @assert i1 < i2 "function is not logconcave, first index with grad>0 must be smaller than first index with grad<0" - x1, x2 = grid[i1], grid[i2] - @info "initial points found at $(x1), $(x2) with grads $(grads[i1]), $(grads[i2])" - RejectionSampler(f, support, (x1, x2); logdensity=logdensity, kwargs...) +end + +# Draw a single sample from `h`. +function sample_hull(rng::AbstractRNG, h::UpperHull) + ws = segment_weights(h) + ind = sample(rng, AllocFreeWeights(ws, sum(ws))) + sl, int = line(h, ind) + return inv_cdf_seg(sl, int, rand(rng), segment_weights(h)[ind], intersections(h)[ind], intersections(h)[ind + 1]) +end +sample_hull(h::UpperHull) = sample_hull(default_rng(), h) + + +# Draw `n` samples from `h`, storing the result in `out`. +function sample_hull!(rng::AbstractRNG, out::AbstractVector{T}, h::UpperHull{T}, n::Integer) where {T} + seg_wgts = segment_weights(h) + inds = sample(rng, 1:n_lines(h), AllocFreeWeights(seg_wgts, sum(seg_wgts)), n) + rands = rand(rng, n) + for i in eachindex(inds, out, rands) + ind = inds[i] + sl, int = line(h, ind) + out[i] = inv_cdf_seg(sl, int, rands[i], segment_weights(h)[ind], intersections(h)[ind], intersections(h)[ind + 1]) end + return end +sample_hull!(out::AbstractVector{T}, h::UpperHull{T}, n::Integer) where {T} = sample_hull!(default_rng(), out, h, n) + +function sample_hull!(rng, out::AbstractVector{T}, h::UpperHull) where {T} + return sample_hull!(rng, out, h, length(out)) +end +sample_hull!(out::AbstractVector{T}, h::UpperHull) where {T} = sample_hull!(default_rng(), out, h) + +# Draw `n` samples from `h`. +function sample_hull(rng::AbstractRNG, h::UpperHull{T}, n::Integer) where {T} + out = Vector{T}(undef, n) + sample_hull!(rng, out, h, n) + return out +end +sample_hull(h::UpperHull, n::Integer) = sample_hull(default_rng(), h, n) + +struct LowerHull{T} + intercepts::Vector{T} + slopes::Vector{T} + intersections::Vector{T} +end + + +Base.broadcastable(h::LowerHull) = Ref(h) + +intercepts(h::LowerHull) = h.intercepts +slopes(h::LowerHull) = h.slopes +intersections(h::LowerHull) = h.intersections +line(h::LowerHull, i) = (h.slopes[i], h.intercepts[i]) + +function calc_lower_slopes_and_intercepts!( + slout::AbstractVector{T}, intout::AbstractVector{T}, + intersections::AbstractVector{T}, f::Function + ) where {T} + + for i in eachindex(slout, intout) + slout[i] = (f(intersections[i + 1]) - f(intersections[i])) / + (intersections[i + 1] - intersections[i]) + intout[i] = slout[i] * (-intersections[i]) + f(intersections[i]) + end + return nothing +end + + +function LowerHull(upper::UpperHull{T}, obj::Objective{T}) where {T} + intersections = abscissae(upper) # NOTE: This makes them alias each other! + n_segs = length(intersections) - 1 + sl = Vector{T}(undef, n_segs) + int = Vector{T}(undef, n_segs) + + calc_lower_slopes_and_intercepts!(sl, int, intersections, obj.f) + + return LowerHull(int, sl, intersections) +end + +function eval_hull(h::LowerHull, x) + i = searchsortedfirst(intersections(h), x) + if isone(i) || i == lastindex(intersections(h)) + 1 + return -Inf + end + sl, int = line(h, i - 1) + return sl * x + int +end + + +""" + ARSampler{T, F, G} + +Adaptive rejection sampler containing the objective function, its gradient and the +upper/lower hull of the piecewise linear envelope. +""" +struct ARSampler{T, F, G} + objective::Objective{T, F, G} + upper_hull::UpperHull{T} + lower_hull::LowerHull{T} +end + +""" + ARSampler( + obj::Objective{T, F, G}, + domain::Tuple{T, T}, + initial_points::Vector{T} + ) where {T <: AbstractFloat, F <: Function, G <: Function} + + ARSampler( + obj::Objective{T, F, G}, + domain::Tuple{T, T}, + search_range::Tuple{T, T} = (-10.0, 10.0), + search_step = 0.1 + ) where {T <: AbstractFloat, F <: Function, G <: Function} + +Initialize an adaptive rejection sampler over an objective function from `obj` ([`ARSampling.Objective`](@ref)). +`initial_points` should be a vector defining the abscissae of the initial segments of the sampler. + At least 2 of the points should be on opposite sides of the objective function's maximum. + +If no `initial_points` an attempt to find suitable ones will be made by searching for +the first negative/positive slope of `obj`. The default search range is `-10:10` with a +step of `0.1`. + +!!! warning + As it currently stands, finding inital points will fail for distributions bounded + to the left/right of their maximum. In these cases the initial point(s) need to + be provided manually (however, only a single initial point has to be provided). +""" +function ARSampler end + +function ARSampler( + obj::Objective{T, F, G}, + domain::Tuple{T, T}, + initial_points::Vector{T} + ) where {T <: AbstractFloat, F <: Function, G <: Function} + + u = UpperHull(obj, initial_points, domain) + l = LowerHull(u, obj) + + return ARSampler{T, F, G}(obj, u, l) +end + + +function ARSampler(obj::Objective{T, F, G}, domain::Tuple{<:Number, <:Number}, search_range::Tuple{<:Number, <:Number} = (-10.0, 10.0); search_step = 0.5) where {T <: AbstractFloat, F <: Function, G <: Function} + domain_T = T.(domain) + search_range_T = T.(search_range) + lbs = T(max(domain_T[1], search_range_T[1])) + ubs = T(min(domain_T[2], search_range_T[2])) + search_step_T = T(search_step) + initial_points = find_initial_points(obj, lbs, ubs, search_step_T) + + return ARSampler(obj, domain_T, initial_points) +end + +""" + find_initial_points(obj, lbs, ubs, search_step) + +Attempts to find suitable initial points for abscissae based on the objective functions by looking for +the first points with positive and negative slopes respectively. Simply searches between the bounds specified +by `lbs` and `ubs` with step size `search_step`. """ - run_sampler!(rng::AbstractRNG, sampler::RejectionSampler, n::Int) - run_sampler!(sampler::RejectionSampler, n::Int) -It draws `n` iid samples of the objective function of `sampler`, and at each iteration it adapts the envelop -of `sampler` by adding new segments to its envelop. -""" -function run_sampler!(rng::AbstractRNG, s::RejectionSampler, n::Int) - i = 0 - failed, max_failed = 0, trunc(Int, n / s.max_failed_rate) - out = zeros(n) - while i < n - candidate = sample_envelop(rng, s.envelop) - acceptance_ratio = exp(s.objective.logf(candidate)) / eval_envelop(s.envelop, candidate) - if rand(rng) < acceptance_ratio - i += 1 - out[i] = candidate - else - if length(s.envelop.lines) <= s.max_segments - a = s.objective.grad(candidate) - b = s.objective.logf(candidate) - a * candidate - add_segment!(s.envelop, Line(a, b)) +function find_initial_points(obj, lbs, ubs, search_step) + grad = obj.grad + r = lbs:search_step:ubs + g = (grad(x) for x in r) + l = findfirst(>(0), g) + u = findfirst(<(0), g) + isnothing(l) && throw( + ErrorException( + lazy""" + Could not find the left initial segment. + Search range: $(lbs) -- $(ubs) + Search step: $(search_step) + """ + ) + ) + isnothing(u) && throw( + ErrorException( + lazy""" + Could not find the right initial segment. + Search range: $(lbs) -- $(ubs) + Search step: $(search_step) + """ + ) + ) + return [r[l], r[u]] +end + + +function n_segments(s::ARSampler) + return n_lines(s.upper_hull) +end + +# Adds a segment with abscissa at `x` to `s` +""" + add_segment!(s::ARSampler{T}, x::T) where {T <: AbstractFloat} + +Modifies the hulls of `s`, adding a segment with abscissa at `x`. +""" +function add_segment!(s::ARSampler{T}, x::T) where {T <: AbstractFloat} + + # Calculate slope, intercept and index of new segment + new_slope = s.objective.grad(x) + new_intercept = (new_slope * -x) + s.objective.f(x) + new_ind = searchsortedfirst(abscissae(s.upper_hull), x) + + # Insert new slope and intercept + all_inters = intersections(s.upper_hull) + all_slopes = slopes(s.upper_hull) + all_intercepts = intercepts(s.upper_hull) + insert!(all_slopes, new_ind, new_slope) + insert!(all_intercepts, new_ind, new_intercept) + + # Insert new abscissa + insert!(abscissae(s.upper_hull), new_ind, x) + + # TODO: Only calculate the intersections and weights that actually change. + # Recalculate intersection points for segments, given the new segment + push!(all_inters, zero(T)) # Extend intercepts by one + calc_intersects!(@view(all_inters[(begin + 1):(end - 1)]), all_slopes, all_intercepts) + + # Recalculate segment weights + all_weights = segment_weights(s.upper_hull) + push!(all_weights, zero(T)) + + for i in eachindex(all_weights) + all_weights[i] = exp_integral_line(all_slopes[i], all_intercepts[i], all_inters[i], all_inters[i + 1]) + end + + #= Lower hull time =# + # We don't need to add anything to the lower intersections as it is the same vector used for abscissae + # in the upper hull + + all_inters_lower = intersections(s.lower_hull) + lowerslopes = slopes(s.lower_hull) + lowerintercepts = intercepts(s.lower_hull) + push!(lowerslopes, zero(T)) + push!(lowerintercepts, zero(T)) + calc_lower_slopes_and_intercepts!( + lowerslopes, + lowerintercepts, + all_inters_lower, + s.objective.f + ) + + return nothing +end + +# TODO: Make this prepopulate `out` using `sample_hull!`? +# NOTE: The above might be inefficient initially if most of the samples get rejected anyway? +""" + __sample!(rng::AbstractRNG, out::AbstractVector{T}, s::ARSampler{T}, add_segments::Bool, max_segments) where {T<:AbstractFloat} + + +All other `__sample!` methods call this one. +""" +function __sample!(rng::AbstractRNG, out::AbstractVector{T}, s::ARSampler{T}, add_segments::Bool, max_segments) where {T <: AbstractFloat} + n = length(out) + n_accepted = 0 + while n_accepted < n + x = sample_hull(rng, s.upper_hull) + up = eval_hull(s.upper_hull, x) + lo = eval_hull(s.lower_hull, x) + w = rand(rng) + # Squeeze test + if w <= exp(lo - up) + n_accepted += 1 + out[n_accepted] = x + elseif w <= exp(s.objective.f(x) - up) + # Accept sample i + n_accepted += 1 + out[n_accepted] = x + if add_segments && n_segments(s) < max_segments + add_segment!(s, x) end - failed += 1 - @assert failed < max_failed "max_failed_factor reached" + elseif add_segments && n_segments(s) < max_segments + add_segment!(s, x) end end - out + return nothing +end +__sample!(out::Vector{T}, s::ARSampler{T}, add_segments::Bool, max_segments) where {T} = __sample!(default_rng(), out, s, add_segments, max_segments) + +# Draw `n` samples, returning a newly allocated vector. +function __sample!(rng::AbstractRNG, s::ARSampler{T}, n::Integer, add_segments::Bool, max_segments) where {T <: AbstractFloat} + out = Vector{T}(undef, n) + __sample!(rng, out, s, add_segments, max_segments) + return out +end +__sample!(s::ARSampler{T}, n::Integer, add_segments::Bool, max_segments) where {T <: AbstractFloat} = __sample!(default_rng(), s, n, add_segments, max_segments) + + +""" + sample!([rng=default_rng()], s::ARSampler, n::Integer, add_segments::Bool=true) + sample!([rng=default_rng()], v::AbstractVector, s::ARSampler, add_segments::Bool=true) + +Draw samples from `s`. If supplied, a vector `v` will be filled with samples. +Otherwise the number of samples is specified with `n`. + +""" +function sample! end + +function sample!(rng::AbstractRNG, s::ARSampler{T}, n::Integer, add_segments::Bool = true, max_segments = DEFAULT_MAX_SEGMENTS) where {T <: AbstractFloat} + return __sample!(rng, s, n, add_segments, max_segments) end +sample!(s::ARSampler{T}, n::Integer, add_segments::Bool = true, max_segments = DEFAULT_MAX_SEGMENTS) where {T <: AbstractFloat} = sample!(default_rng(), s, n, add_segments, max_segments) -function run_sampler!(s::RejectionSampler, n::Int) - run_sampler!(Random.GLOBAL_RNG, s, n) +function sample!(rng::AbstractRNG, v::AbstractVector{T}, s::ARSampler{T}, add_segments::Bool = true, max_segments = DEFAULT_MAX_SEGMENTS) where {T} + __sample!(rng, v, s, add_segments, max_segments) + return nothing end +sample!(v::AbstractVector{T}, s::ARSampler{T}, add_segments::Bool = true, max_segments = DEFAULT_MAX_SEGMENTS) where {T} = sample!(default_rng(), v, s, add_segments, max_segments) + + end # module AdaptiveRejectionSampling diff --git a/test/runtests.jl b/test/runtests.jl index 75e247f..166d2c0 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,81 +1,212 @@ -import AdaptiveRejectionSampling -ars = AdaptiveRejectionSampling - +using AdaptiveRejectionSampling: ARSampler, Objective, sample!, AllocFreeWeights +using DifferentiationInterface +using ForwardDiff +using Mooncake +using Enzyme +using SpecialFunctions: loggamma +using Random +using Distributions using Test +using JET + +include("wrappedallocs.jl") + +normal(x, mu, sigma) = pdf(Normal(mu, sigma), x) +normal_log(x, mu, sigma) = logpdf(Normal(mu, sigma), x) +score_normal(x, mu, sigma) = -((x - mu) / sigma^2) -@testset "Line" begin - @test ars.Line(2.0, 3) isa ars.Line - @test ars.intersection(ars.Line(-1.0, 1.0), ars.Line(1.0, -1.0)) == 1.0 - @test_throws AssertionError ars.intersection(ars.Line(-1.0, 1.0), ars.Line(-1.0, -1.0)) +function test_density(x, k = 3, n = 20) + alpha = exp(x) + return x + + x * (k - 3 / 2) + + (-1 / (2 * alpha)) + + loggamma(alpha) - + loggamma(n + alpha) end -@testset "Envelop" begin - @test begin - l1 = ars.Line(1.0, 1.0) - l2 = ars.Line(-3.0, 2.0) - support = (-Inf, Inf) - ars.Envelop([l1, l2], support) - end isa ars.Envelop - @test_throws AssertionError begin - l1 = ars.Line(1.0, 1.0) - l2 = ars.Line(-3.0, 2.0) - ars.Envelop([l2, l1], (-Inf, Inf)) - end - @test_throws AssertionError begin - l1 = ars.Line(1.0, 1.0) - l2 = ars.Line(-3.0, 2.0) - ars.Envelop([l2, l1], (1.0, -1.0)) - end - @test begin - l1 = ars.Line(1.0, 1.0) - l2 = ars.Line(-3.0, 2.0) - l3 = ars.Line(-5.0, 5.0) - e = ars.Envelop([l1, l2, l3], (-Inf, Inf)) - e.cutpoints == [-Inf, 0.25, 1.5, Inf] - end - @test begin - l1 = ars.Line(1.0, 1.0) - newline = ars.Line(0.0, 0.0) - l3 = ars.Line(-1.0, 1.0) - e = ars.Envelop([l1, l3], (-Inf, Inf)) - ars.add_segment!(e, newline) - e.lines[2] == newline - end +@testset "AllocFreeWeights" begin + v = rand(1000) + + @test @wrappedallocs(AllocFreeWeights(v)) == 0 + + w = AllocFreeWeights(v) + + @test w.sum == sum(v) + @test w.values == v end +fn(x) = normal_log(x, 3.14, 3.0) +obj = Objective(fn) +x = 2.4 + @testset "Objective" begin - @test ars.Objective(x -> exp(-0.5 * x^2) / sqrt(2pi)) isa ars.Objective - @test begin - objective = ars.Objective(x -> exp(-0.5 * x^2) / sqrt(2pi)) - objective.logf(0.0) ≈ 0.3989422804014327 - end - @test begin - objective = ars.Objective(x -> exp(-0.5 * x^2) / sqrt(2pi)) - objective.grad(0.0) ≈ 0.0 - end - @test begin - objective = ars.Objective(x -> exp(-0.5 * x^2) / sqrt(2pi)) - objective.logf(Inf) == 0.0 && objective.logf(Inf) == 0.0 - end + @test obj isa Objective{Float64, <:Function, <:Function} + @test obj.grad(x) ≈ score_normal(x, 3.14, 3.0) + @test obj.f(x) ≈ normal_log(x, 3.14, 3.0) + @test obj.grad(3.14) == 0.0 end -@testset "RejectionSampler" begin - @test ars.RejectionSampler(x -> exp(-0.5 * x^2) / sqrt(2pi), (-Inf, Inf), (-1.0, 1.0)) isa ars.RejectionSampler - @test begin - sampler = ars.RejectionSampler(x -> exp(-0.5 * x^2) / sqrt(2pi), (-Inf, Inf), (-1.0, 1.0)) - ars.run_sampler!(sampler, 5) isa Vector{T} where T <: AbstractFloat - end - @test begin - sampler = ars.RejectionSampler(x -> exp(-0.5 * x^2) / sqrt(2pi), (-Inf, Inf)) - ars.run_sampler!(sampler, 5) isa Vector{T} where T <: AbstractFloat +const n_samples = 100_000 + +@testset "Autodiff backends" begin + @testset "ForwardDiff" begin + let + sam = ARSampler(Objective(test_density; adbackend = AutoForwardDiff()), (-Inf, Inf)) + obj_big = Objective(test_density, one(BigFloat); adbackend = AutoForwardDiff()) + sam_big = ARSampler(obj_big, (-Inf, Inf)) + # Setup samples + sam_2 = deepcopy(sam) + Random.seed!(1) + samples = sample!(sam, n_samples) + Random.seed!(1) + samples_2 = sample!(sam_2, n_samples) + + # Tests + @test samples isa Vector{Float64} + @test length(samples) == n_samples + @test samples == samples_2 # Try to catch uninitialized memory + + # JET + @test_opt Objective(test_density; adbackend = AutoForwardDiff()) + obj_jet = Objective(test_density; adbackend = AutoForwardDiff()) + @test_opt ARSampler(obj_jet, (-Inf, Inf)) + @test_opt sample!(sam, n_samples) + + + # Test sampling BigFloat + @test sample!(sam_big, 10_000) isa Vector{BigFloat} + end end - @test_throws AssertionError begin - sampler = ars.RejectionSampler(x -> exp(-0.5 * x) / sqrt(2pi), (-Inf, Inf)) - ars.run_sampler!(sampler, 5) isa Vector{T} where T <: AbstractFloat + + + @testset "Mooncake" begin + let + sam = ARSampler(Objective(test_density; adbackend = AutoMooncake()), (-Inf, Inf)) + # Setup samples + sam_2 = deepcopy(sam) + Random.seed!(1) + samples = sample!(sam, n_samples) + Random.seed!(1) + samples_2 = sample!(sam_2, n_samples) + + # Tests + @test samples isa Vector{Float64} + @test length(samples) == n_samples + @test samples == samples_2 # Try to catch uninitialized memory + + # JET, errors from within Mooncake + # @test_opt Objective(test_density; adbackend = AutoMooncake()) + # obj_jet = Objective(test_density; adbackend = AutoMooncake()) + # @test_opt ARSampler(obj_jet, (-Inf, Inf)) + # @test_opt sample!(sam, n_samples) + end end - @test begin - sampler = ars.RejectionSampler(x -> -abs(x), (-Inf, Inf), logdensity=true) - ars.run_sampler!(sampler, 5) isa Vector{T} where {T<:AbstractFloat} + + + @testset "Enzyme" begin + let + sam = ARSampler(Objective(test_density; adbackend = AutoEnzyme()), (-Inf, Inf)) + # Setup samples + sam_2 = deepcopy(sam) + Random.seed!(1) + samples = sample!(sam, n_samples) + Random.seed!(1) + samples_2 = sample!(sam_2, n_samples) + + # Tests + @test samples isa Vector{Float64} + @test length(samples) == n_samples + @test samples == samples_2 # Try to catch uninitialized memory + + # JET + @test_opt Objective(test_density; adbackend = AutoEnzyme()) + obj_jet = Objective(test_density; adbackend = AutoEnzyme()) + @test_opt ARSampler(obj_jet, (-Inf, Inf)) + @test_opt sample!(sam, n_samples) + end end end + +# import AdaptiveRejectionSampling +# ars = AdaptiveRejectionSampling + +# using Test + + + +# @testset "Line" begin +# @test ars.Line(2.0, 3) isa ars.Line +# @test ars.intersection(ars.Line(-1.0, 1.0), ars.Line(1.0, -1.0)) == 1.0 +# @test_throws AssertionError ars.intersection(ars.Line(-1.0, 1.0), ars.Line(-1.0, -1.0)) +# end + +# @testset "Envelop" begin +# @test begin +# l1 = ars.Line(1.0, 1.0) +# l2 = ars.Line(-3.0, 2.0) +# support = (-Inf, Inf) +# ars.Envelop([l1, l2], support) +# end isa ars.Envelop +# @test_throws AssertionError begin +# l1 = ars.Line(1.0, 1.0) +# l2 = ars.Line(-3.0, 2.0) +# ars.Envelop([l2, l1], (-Inf, Inf)) +# end +# @test_throws AssertionError begin +# l1 = ars.Line(1.0, 1.0) +# l2 = ars.Line(-3.0, 2.0) +# ars.Envelop([l2, l1], (1.0, -1.0)) +# end +# @test begin +# l1 = ars.Line(1.0, 1.0) +# l2 = ars.Line(-3.0, 2.0) +# l3 = ars.Line(-5.0, 5.0) +# e = ars.Envelop([l1, l2, l3], (-Inf, Inf)) +# e.cutpoints == [-Inf, 0.25, 1.5, Inf] +# end +# @test begin +# l1 = ars.Line(1.0, 1.0) +# newline = ars.Line(0.0, 0.0) +# l3 = ars.Line(-1.0, 1.0) +# e = ars.Envelop([l1, l3], (-Inf, Inf)) +# ars.add_segment!(e, newline) +# e.lines[2] == newline +# end +# end + +# @testset "Objective" begin +# @test ars.Objective(x -> exp(-0.5 * x^2) / sqrt(2pi)) isa ars.Objective +# @test begin +# objective = ars.Objective(x -> exp(-0.5 * x^2) / sqrt(2pi)) +# objective.logf(0.0) ≈ 0.3989422804014327 +# end +# @test begin +# objective = ars.Objective(x -> exp(-0.5 * x^2) / sqrt(2pi)) +# objective.grad(0.0) ≈ 0.0 +# end +# @test begin +# objective = ars.Objective(x -> exp(-0.5 * x^2) / sqrt(2pi)) +# objective.logf(Inf) == 0.0 && objective.logf(Inf) == 0.0 +# end +# end + +# @testset "RejectionSampler" begin +# @test ars.RejectionSampler(x -> exp(-0.5 * x^2) / sqrt(2pi), (-Inf, Inf), (-1.0, 1.0)) isa ars.RejectionSampler +# @test begin +# sampler = ars.RejectionSampler(x -> exp(-0.5 * x^2) / sqrt(2pi), (-Inf, Inf), (-1.0, 1.0)) +# ars.run_sampler!(sampler, 5) isa Vector{T} where T <: AbstractFloat +# end +# @test begin +# sampler = ars.RejectionSampler(x -> exp(-0.5 * x^2) / sqrt(2pi), (-Inf, Inf)) +# ars.run_sampler!(sampler, 5) isa Vector{T} where T <: AbstractFloat +# end +# @test_throws AssertionError begin +# sampler = ars.RejectionSampler(x -> exp(-0.5 * x) / sqrt(2pi), (-Inf, Inf)) +# ars.run_sampler!(sampler, 5) isa Vector{T} where T <: AbstractFloat +# end +# @test begin +# sampler = ars.RejectionSampler(x -> -abs(x), (-Inf, Inf), logdensity=true) +# ars.run_sampler!(sampler, 5) isa Vector{T} where {T<:AbstractFloat} +# end +# end From 8d906f89aebf6dfba51bc33acbaa38658ba3240d Mon Sep 17 00:00:00 2001 From: Eliassj Date: Fri, 3 Apr 2026 10:59:19 +0200 Subject: [PATCH 2/8] Add some Supposition.jl to tests --- test/runtests.jl | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/test/runtests.jl b/test/runtests.jl index 166d2c0..83ee1a9 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -8,6 +8,8 @@ using Random using Distributions using Test using JET +using Supposition +using Supposition: @check include("wrappedallocs.jl") @@ -39,13 +41,25 @@ fn(x) = normal_log(x, 3.14, 3.0) obj = Objective(fn) x = 2.4 +const f64gen = Data.Floats{Float64}(nans = false) +const f64genmu = Data.Floats{Float64}(nans = false) +const f64gensigma = Data.Floats{Float64}(nans = false, minimum = 0.0001) + @testset "Objective" begin @test obj isa Objective{Float64, <:Function, <:Function} @test obj.grad(x) ≈ score_normal(x, 3.14, 3.0) - @test obj.f(x) ≈ normal_log(x, 3.14, 3.0) + @check function objective_f64_grad(v = f64gen, mu = f64genmu, sigma = f64gensigma) + o = Objective(x -> normal_log(x, mu, sigma)) + ≈(score_normal(v, mu, sigma), o.grad(v), rtol = 0.01, atol = 0.0001) + end + @check function objective_f64_f(v = f64gen, mu = f64genmu, sigma = f64gensigma) + o = Objective(x -> normal_log(x, mu, sigma)) + o.f(v) ≈ normal_log(v, mu, sigma) + end @test obj.grad(3.14) == 0.0 end + const n_samples = 100_000 @testset "Autodiff backends" begin @@ -134,7 +148,6 @@ end # using Test - # @testset "Line" begin # @test ars.Line(2.0, 3) isa ars.Line # @test ars.intersection(ars.Line(-1.0, 1.0), ars.Line(1.0, -1.0)) == 1.0 From 412c0ccf5fceecfd726e4f475ac5473af427eb14 Mon Sep 17 00:00:00 2001 From: Eliassj Date: Wed, 8 Apr 2026 17:16:32 +0200 Subject: [PATCH 3/8] Add tests and update README --- .github/workflows/Docs.yml | 62 ++++ .github/workflows/ReusableTest.yml | 52 +++ .github/workflows/Test.yml | 35 ++ img/example1.svg | 315 ++++++++++++++++ img/example2.svg | 306 ++++++++++++++++ src/sampling.jl | 553 +++++++++++++++++++++++++++++ test/Project.toml | 12 + test/wrappedallocs.jl | 35 ++ 8 files changed, 1370 insertions(+) create mode 100644 .github/workflows/Docs.yml create mode 100644 .github/workflows/ReusableTest.yml create mode 100644 .github/workflows/Test.yml create mode 100644 img/example1.svg create mode 100644 img/example2.svg create mode 100644 src/sampling.jl create mode 100644 test/Project.toml create mode 100644 test/wrappedallocs.jl diff --git a/.github/workflows/Docs.yml b/.github/workflows/Docs.yml new file mode 100644 index 0000000..dc2659a --- /dev/null +++ b/.github/workflows/Docs.yml @@ -0,0 +1,62 @@ +name: Docs + +on: + push: + branches: + - master + paths: + - "docs/**" + - "src/**" + - "*.toml" + tags: ["*"] + pull_request: + branches: + - master + paths: + - "docs/**" + - "src/**" + - "*.toml" + types: [opened, synchronize, reopened] + +concurrency: + # Skip intermediate builds: always. + # Cancel intermediate builds: only if it is a pull request build. + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }} + +jobs: + docs: + permissions: + actions: write + contents: write + pull-requests: read + statuses: write + name: Documentation + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + - uses: julia-actions/setup-julia@v2 + with: + version: "1" + - name: Use Julia cache + uses: julia-actions/cache@v2 + - name: Instantiate environment with development version of the package + run: | + julia --project=docs -e ' + using Pkg + Pkg.develop(PackageSpec(path=pwd())) + Pkg.instantiate()' + - name: Run doctest + run: | + julia --project=docs -e ' + using Documenter: DocMeta, doctest + using ARSampling + DocMeta.setdocmeta!(ARSampling, :DocTestSetup, :(using ARSampling); recursive=true) + doctest(ARSampling)' + - name: Generate and deploy documentation + run: julia --project=docs docs/make.jl + env: + JULIA_PKG_SERVER: "" + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }} + GKSwstype: "100" # https://discourse.julialang.org/t/generation-of-documentation-fails-qt-qpa-xcb-could-not-connect-to-display/60988 diff --git a/.github/workflows/ReusableTest.yml b/.github/workflows/ReusableTest.yml new file mode 100644 index 0000000..ec97196 --- /dev/null +++ b/.github/workflows/ReusableTest.yml @@ -0,0 +1,52 @@ +name: Reusable test + +on: + workflow_call: + inputs: + version: + required: false + type: string + default: "1" + os: + required: false + type: string + default: ubuntu-latest + arch: + required: false + type: string + default: x64 + allow_failure: + required: false + type: boolean + default: false + run_codecov: + required: false + type: boolean + default: false + secrets: + codecov_token: + required: true + +jobs: + test: + name: Julia ${{ inputs.version }} - ${{ inputs.os }} - ${{ inputs.arch }} - ${{ github.event_name }} + runs-on: ${{ inputs.os }} + continue-on-error: ${{ inputs.allow_failure }} + + steps: + - uses: actions/checkout@v6 + - uses: julia-actions/setup-julia@v2 + with: + version: ${{ inputs.version }} + arch: ${{ inputs.arch }} + - name: Use Julia cache + uses: julia-actions/cache@v2 + - uses: julia-actions/julia-buildpkg@v1 + - uses: julia-actions/julia-runtest@v1 + - uses: julia-actions/julia-processcoverage@v1 + if: ${{ inputs.run_codecov }} + - uses: codecov/codecov-action@v4 + if: ${{ inputs.run_codecov }} + with: + file: lcov.info + token: ${{ secrets.codecov_token }} diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml new file mode 100644 index 0000000..4d5bcc3 --- /dev/null +++ b/.github/workflows/Test.yml @@ -0,0 +1,35 @@ +name: Test + +on: + push: + branches: + - master + tags: ["*"] + workflow_dispatch: + +jobs: + test: + uses: ./.github/workflows/ReusableTest.yml + with: + os: ${{ matrix.os }} + version: ${{ matrix.version }} + arch: ${{ matrix.arch }} + allow_failure: ${{ matrix.allow_failure }} + run_codecov: ${{ matrix.version == '1' && matrix.os == 'ubuntu-latest' }} + secrets: + codecov_token: ${{ secrets.CODECOV_TOKEN }} + strategy: + fail-fast: false + matrix: + version: + - "lts" + - "1" + os: + - ubuntu-latest + - macOS-latest + + - windows-latest + + arch: + - x64 + allow_failure: [false] diff --git a/img/example1.svg b/img/example1.svg new file mode 100644 index 0000000..d7ee899 --- /dev/null +++ b/img/example1.svg @@ -0,0 +1,315 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/img/example2.svg b/img/example2.svg new file mode 100644 index 0000000..acf0282 --- /dev/null +++ b/img/example2.svg @@ -0,0 +1,306 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/sampling.jl b/src/sampling.jl new file mode 100644 index 0000000..784e24a --- /dev/null +++ b/src/sampling.jl @@ -0,0 +1,553 @@ +const DEFAULT_MAX_SEGMENTS::Int = 25 + + +""" + AllocFreeWeights{S <: Real, T <: Number, V <: AbstractVector{T}} <: AbstractWeights{S, T, V} + +Non-mutable weights used in the `sample` function in this package in order to avoid allocations. +""" +struct AllocFreeWeights{S <: Real, T <: Number, V <: AbstractVector{T}} <: AbstractWeights{S, T, V} + values::V + sum::S + + function AllocFreeWeights(ws::V, sum::S) where {S <: Real, T <: Number, V <: AbstractVector{T}} + return new{S, T, V}(ws, sum) + end +end + +function AllocFreeWeights(ws::V) where {T <: Number, V <: AbstractVector{T}} + return AllocFreeWeights(ws, sum(ws)) +end + + +struct Objective{T <: Number, F <: Function, G <: Function} + "The (log-concave) function defining the density `f(x) -> y`" + f::F + "Its gradient `grad(x) -> y'`" + grad::G + + @doc """ + Objective(f::Function, grad::Function, ::Type{T} = Float64) where {T <: Number} + + Create an `Objective` directly defined by its function `f` and custom gradient `grad`. + The parameter `T` represents the type of the expected input and is utilized when preparing gradients using autodiff. + By default `T = Float64`. + + + !!! warning + Observe that `f` should be in its + log-concave form and that no checks are performed in order to verify this. + """ + function Objective(f::Function, grad::Function, ::Type{T} = Float64) where {T <: Number} + return new{T, typeof(f), typeof(grad)}(f, grad) + end +end + +""" + Objective(f::Function, init = one(Float64); adbackend = AutoForwardDiff()) + +Create an `Objective` for a function automatically generating its gradient. +Gradients are calculated using `DifferentiationInterface.jl` using the backend of choice +with `ForwardDiff.jl` being the default. In order to prepare the gradient an initial value is +required. By default this is `one(Float64)`. If a gradient for a different type is +desired, it should be specified through this initial value. + +If a custom/manual gradient is available it may instead be provided. + +!!! warning + Observe that `f` should be in its + log-concave form and that no checks are performed in order to verify this. + +# Example + +```julia +# Create an objective using `Mooncake.jl` autodiff and Float32 as its type. +ARS.Objective(somefun, one(Float32); adbackend = AutoMooncake(; config=nothing)) +# Create an objective using the defaults `AutoDiff.jl` autodiff and Float64 as its type. +ARS.Objective(somefun) +``` +""" +function Objective(f::Function, init = one(Float64); adbackend = AutoForwardDiff()) + gradprep = prepare_gradient(f, adbackend, init) + obj = Objective( + f, + x -> gradient(f, gradprep, adbackend, x), + typeof(init) + ) + return obj +end + +abstract type AbstractHull end + +""" + struct UpperHull{T} <: AbstractHull + intercepts::Vector{T} + slopes::Vector{T} + intersections::Vector{T} + abscissae::Vector{T} + segment_weights::Vector{T} + domain::Tuple{T, T} + end + +Upper hull of the enveloping function. +""" +struct UpperHull{T} <: AbstractHull + intercepts::Vector{T} + slopes::Vector{T} + intersections::Vector{T} + abscissae::Vector{T} + segment_weights::Vector{T} + domain::Tuple{T, T} +end + +Base.broadcastable(h::UpperHull) = Ref(h) + +slopes(h::UpperHull) = h.slopes +intercepts(h::UpperHull) = h.intercepts +intersections(h::UpperHull) = h.intersections +abscissae(h::UpperHull) = h.abscissae +line(h::UpperHull, i::Integer) = (h.slopes[i], h.intercepts[i]) +lineinds(h::UpperHull) = 1:(length(slopes(h))) +n_lines(h::UpperHull) = length(slopes(h)) +segment_weights(h::UpperHull) = h.segment_weights + +""" + eval_hull(h::UpperHull, x) + +Eval hull `h` at `x`. +""" +function eval_hull(h::UpperHull{T}, x::T) where {T} + ints = intersections(h) + i = searchsortedfirst(ints, x) - 1 + i = i == length(ints) + 1 ? i - 1 : i + i = iszero(i) ? i + 1 : i + sl, int = line(h, i) + return sl * x + int +end + + +""" + intersection(slope1::T, intercept1::T, slope2::T, intercept2::T) where {T} + +Returns the intersection abscissa between 2 lines as defined by their slopes and intercepts. Returns NaN if the lines are paralell. +""" +function intersection(slope1::T, intercept1::T, slope2::T, intercept2::T) where {T} + return (intercept2 - intercept1) / (slope1 - slope2) +end + +""" + calc_intersects!(out::AbstractVector{T}, slopes::AbstractVector{T}, intercepts::AbstractVector{T}) where {T} + +Calculate intersections of a series of lines defined by their `slopes` and `intercepts`, storing the result in `out`. +""" +function calc_intersects!(out::AbstractVector{T}, slopes::AbstractVector{T}, intercepts::AbstractVector{T}) where {T} + for i in eachindex(out) + out[i] = intersection(slopes[i], intercepts[i], slopes[i + 1], intercepts[i + 1]) + end + return nothing +end + +""" + calc_intersects(slopes::AbstractVector{T}, intercepts::AbstractVector{T}) where {T} + +Calculate intersections of a series of lines defined by their `slopes` and `intercepts`. +""" +function calc_intersects(slopes::AbstractVector{T}, intercepts::AbstractVector{T}) where {T} + # return [ + # intersection(slopes[i], intercepts[i], slopes[i + 1], intercepts[i + 1]) + # for i in 1:(length(slopes) - 1) + # ] + + out = Vector{T}(undef, length(slopes) - 1) + calc_intersects!(out, slopes, intercepts) + return out +end + + +""" + calc_slopes_and_intercepts(obj::Objective, abscissae::AbstractVector{T}) where {T} + +Calculates slopes and intercepts of the lines tangential to the objective function at `abscissae`. +""" +function calc_slopes_and_intercepts(obj::Objective, abscissae::AbstractVector{T}) where {T} + issorted(abscissae) || + throw(ArgumentError("`abscissae` should be sorted in ascending order.")) + slopes = [obj.grad(abscissae[i]) for i in eachindex(abscissae)] + intercepts = [ + slopes[i] * (-abscissae[i]) + obj.f(abscissae[i]) + for i in eachindex(abscissae) + ] + return slopes, intercepts +end + + +""" + exp_integral_line(slope, intercept, x1, x2) + +Calculate the integral of `exp(slope * x + intercept)` between `x1` and `x2`. +""" +function exp_integral_line(slope, intercept, x1, x2) + return if !iszero(slope) + (exp(slope * x2 + intercept) - exp(slope * x1 + intercept)) / slope + else + (x2 - x1) * exp(intercept) + end +end + +""" + UpperHull(obj::Objective, abscissae::Vector{T}, domain::Tuple{T, T}) where {T} + +Create an `UpperHull` based on the objective function, initial abscissae and its domain. +""" +function UpperHull(obj::Objective{T}, abscissae::Vector{T}, domain::Tuple{T, T}) where {T} + slopes, intercepts = calc_slopes_and_intercepts(obj, abscissae) + + intersects = [domain[1]; calc_intersects(slopes, intercepts); domain[2]] + + wgts = Vector{T}(undef, length(slopes)) + for i in 1:length(wgts) + wgts[i] = exp_integral_line(slopes[i], intercepts[i], intersects[i], intersects[i + 1]) + end + + return UpperHull(intercepts, slopes, intersects, abscissae, wgts, domain) +end + +# Calculate the inverse CDF of a segment, handling a slope of zero. +""" + inv_cdf_seg(slope, intercept, r, w, intersect, intersect2) + +Calculate the inverse cumulative density of a hull segment at `r` based on its... + +- `slope` +- `intercept` +- `w` (weight) +- `intersect`, `intersect2`, its intersections with previous/next segments +""" +@inline function inv_cdf_seg(slope, intercept, r, w, intersect, intersect2) + return if !iszero(slope) + log(exp(-intercept) * r * w * slope + exp(slope * intersect)) / slope + else + r * (intersect2 - intersect) + intersect + end +end + +# Draw a single sample from `h`. +function sample_hull(rng::AbstractRNG, h::UpperHull) + ws = segment_weights(h) + ind = sample(rng, AllocFreeWeights(ws, sum(ws))) + sl, int = line(h, ind) + return inv_cdf_seg(sl, int, rand(rng), segment_weights(h)[ind], intersections(h)[ind], intersections(h)[ind + 1]) +end +sample_hull(h::UpperHull) = sample_hull(default_rng(), h) + + +# Draw `n` samples from `h`, storing the result in `out`. +function sample_hull!(rng::AbstractRNG, out::AbstractVector{T}, h::UpperHull{T}, n::Integer) where {T} + seg_wgts = segment_weights(h) + inds = sample(rng, 1:n_lines(h), AllocFreeWeights(seg_wgts, sum(seg_wgts)), n) + rands = rand(rng, n) + for i in eachindex(inds, out, rands) + ind = inds[i] + sl, int = line(h, ind) + out[i] = inv_cdf_seg(sl, int, rands[i], segment_weights(h)[ind], intersections(h)[ind], intersections(h)[ind + 1]) + end + return +end +sample_hull!(out::AbstractVector{T}, h::UpperHull{T}, n::Integer) where {T} = sample_hull!(default_rng(), out, h, n) + +function sample_hull!(rng, out::AbstractVector{T}, h::UpperHull) where {T} + return sample_hull!(rng, out, h, length(out)) +end +sample_hull!(out::AbstractVector{T}, h::UpperHull) where {T} = sample_hull!(default_rng(), out, h) + +# Draw `n` samples from `h`. +function sample_hull(rng::AbstractRNG, h::UpperHull{T}, n::Integer) where {T} + out = Vector{T}(undef, n) + sample_hull!(rng, out, h, n) + return out +end +sample_hull(h::UpperHull, n::Integer) = sample_hull(default_rng(), h, n) + +struct LowerHull{T} + intercepts::Vector{T} + slopes::Vector{T} + intersections::Vector{T} +end + + +Base.broadcastable(h::LowerHull) = Ref(h) + + +intercepts(h::LowerHull) = h.intercepts +slopes(h::LowerHull) = h.slopes +intersections(h::LowerHull) = h.intersections +line(h::LowerHull, i) = (h.slopes[i], h.intercepts[i]) + +function calc_lower_slopes_and_intercepts!( + slout::AbstractVector{T}, intout::AbstractVector{T}, + intersections::AbstractVector{T}, f::Function + ) where {T} + + for i in eachindex(slout, intout) + slout[i] = (f(intersections[i + 1]) - f(intersections[i])) / + (intersections[i + 1] - intersections[i]) + intout[i] = slout[i] * (-intersections[i]) + f(intersections[i]) + end + return nothing +end + + +function LowerHull(upper::UpperHull{T}, obj::Objective{T}) where {T} + intersections = abscissae(upper) # NOTE: This makes them alias each other! + n_segs = length(intersections) - 1 + sl = Vector{T}(undef, n_segs) + int = Vector{T}(undef, n_segs) + + calc_lower_slopes_and_intercepts!(sl, int, intersections, obj.f) + + return LowerHull(int, sl, intersections) +end + +function eval_hull(h::LowerHull, x) + i = searchsortedfirst(intersections(h), x) + if isone(i) || i == lastindex(intersections(h)) + 1 + return -Inf + end + sl, int = line(h, i - 1) + return sl * x + int +end + + +""" + ARSampler{T, F, G} + +Adaptive rejection sampler containing the objective function, its gradient and the +upper/lower hull of the piecewise linear envelope. +""" +struct ARSampler{T, F, G} + objective::Objective{T, F, G} + upper_hull::UpperHull{T} + lower_hull::LowerHull{T} +end + +""" + ARSampler( + obj::Objective{T, F, G}, + domain::Tuple{T, T}, + initial_points::Vector{T} + ) where {T <: AbstractFloat, F <: Function, G <: Function} + + ARSampler( + obj::Objective{T, F, G}, + domain::Tuple{T, T}, + search_range::Tuple{T, T} = (-10.0, 10.0), + search_step = 0.1 + ) where {T <: AbstractFloat, F <: Function, G <: Function} + +Initialize an adaptive rejection sampler over an objective function from `obj` ([`AdaptiveRejectionSampling.Objective`](@ref)). +`initial_points` should be a vector defining the abscissae of the initial segments of the sampler. + At least 2 of the points should be on opposite sides of the objective function's maximum. + +If no `initial_points` an attempt to find suitable ones will be made by searching for +the first negative/positive slope of `obj`. The default search range is `-10:10` with a +step of `0.1`. + +!!! warning + As it currently stands, finding inital points will fail for distributions bounded + to the left/right of their maximum. In these cases the initial point(s) need to + be provided manually (however, only a single initial point has to be provided). +""" +function ARSampler end + +function ARSampler( + obj::Objective{T, F, G}, + domain::Tuple{T, T}, + initial_points::Vector{T} + ) where {T <: AbstractFloat, F <: Function, G <: Function} + + u = UpperHull(obj, initial_points, domain) + l = LowerHull(u, obj) + + return ARSampler{T, F, G}(obj, u, l) +end + + +function ARSampler(obj::Objective{T, F, G}, domain::Tuple{<:Number, <:Number}, search_range::Tuple{<:Number, <:Number} = (-10.0, 10.0); search_step = 0.5) where {T <: AbstractFloat, F <: Function, G <: Function} + domain_T = T.(domain) + search_range_T = T.(search_range) + lbs = T(max(domain_T[1], search_range_T[1])) + ubs = T(min(domain_T[2], search_range_T[2])) + search_step_T = T(search_step) + initial_points = find_initial_points(obj, lbs, ubs, search_step_T) + + return ARSampler(obj, domain_T, initial_points) +end + +support(s::ARSampler) = s.upper_hull.domain +abscissae(s::ARSampler) = abscissae(s.upper_hull) +slopes(s::ARSampler) = slopes(s.upper_hull) +intercepts(s::ARSampler) = intercepts(s.upper_hull) + +""" + find_initial_points(obj, lbs, ubs, search_step) + +Attempts to find suitable initial points for abscissae based on the objective functions by looking for +the first points with positive and negative slopes respectively. Simply searches between the bounds specified +by `lbs` and `ubs` with step size `search_step`. +""" +function find_initial_points(obj, lbs, ubs, search_step) + grad = obj.grad + r = lbs:search_step:ubs + g = (grad(x) for x in r) + l = findfirst(>(0), g) + u = findfirst(<(0), g) + isnothing(l) && throw( + ErrorException( + lazy""" + Could not find the left initial segment. + Search range: $(lbs) -- $(ubs) + Search step: $(search_step) + """ + ) + ) + isnothing(u) && throw( + ErrorException( + lazy""" + Could not find the right initial segment. + Search range: $(lbs) -- $(ubs) + Search step: $(search_step) + """ + ) + ) + return [r[l], r[u]] +end + + +function n_segments(s::ARSampler) + return n_lines(s.upper_hull) +end + +# Adds a segment with abscissa at `x` to `s` +""" + add_segment!(s::ARSampler{T}, x::T) where {T <: AbstractFloat} + +Modifies the hulls of `s`, adding a segment with abscissa at `x`. +""" +function add_segment!(s::ARSampler{T}, x::T) where {T <: AbstractFloat} + + # Calculate slope, intercept and index of new segment + new_slope = s.objective.grad(x) + new_intercept = (new_slope * -x) + s.objective.f(x) + new_ind = searchsortedfirst(abscissae(s.upper_hull), x) + + # Insert new slope and intercept + all_intersections = intersections(s.upper_hull) + all_slopes = slopes(s.upper_hull) + all_intercepts = intercepts(s.upper_hull) + insert!(all_slopes, new_ind, new_slope) + insert!(all_intercepts, new_ind, new_intercept) + + # Insert new abscissa + insert!(abscissae(s.upper_hull), new_ind, x) + + # TODO: Only calculate the intersections and weights that actually change. + # Recalculate intersection points for segments, given the new segment + push!(all_intersections, all_intersections[end]) # Extend intercepts by one + calc_intersects!(@view(all_intersections[(begin + 1):(end - 1)]), all_slopes, all_intercepts) + + # Recalculate segment weights + all_weights = segment_weights(s.upper_hull) + push!(all_weights, zero(T)) + + for i in eachindex(all_weights) + all_weights[i] = exp_integral_line(all_slopes[i], all_intercepts[i], all_intersections[i], all_intersections[i + 1]) + end + + #= Lower hull time =# + # We don't need to add anything to the lower intersections as it is the same vector used for abscissae + # in the upper hull + + all_inters_lower = intersections(s.lower_hull) + lowerslopes = slopes(s.lower_hull) + lowerintercepts = intercepts(s.lower_hull) + push!(lowerslopes, zero(T)) + push!(lowerintercepts, zero(T)) + calc_lower_slopes_and_intercepts!( + lowerslopes, + lowerintercepts, + all_inters_lower, + s.objective.f + ) + + return nothing +end + +# TODO: Make this prepopulate `out` using `sample_hull!`? +# NOTE: The above might be inefficient initially if most of the samples get rejected anyway? +""" + __sample!(rng::AbstractRNG, out::AbstractVector{T}, s::ARSampler{T}, add_segments::Bool, max_segments) where {T<:AbstractFloat} + + +All other `__sample!` methods call this one. +""" +function __sample!(rng::AbstractRNG, out::AbstractVector{T}, s::ARSampler{T}, add_segments::Bool, max_segments) where {T <: AbstractFloat} + n = length(out) + n_accepted = 0 + while n_accepted < n + x = sample_hull(rng, s.upper_hull) + up = eval_hull(s.upper_hull, x) + lo = eval_hull(s.lower_hull, x) + w = rand(rng) + # Squeeze test + if w <= exp(lo - up) + n_accepted += 1 + out[n_accepted] = x + elseif w <= exp(s.objective.f(x) - up) + # Accept sample i + n_accepted += 1 + out[n_accepted] = x + if add_segments && n_segments(s) < max_segments + add_segment!(s, x) + end + elseif add_segments && n_segments(s) < max_segments + add_segment!(s, x) + end + end + return nothing +end +__sample!(out::Vector{T}, s::ARSampler{T}, add_segments::Bool, max_segments) where {T} = __sample!(default_rng(), out, s, add_segments, max_segments) + +# Draw `n` samples, returning a newly allocated vector. +function __sample!(rng::AbstractRNG, s::ARSampler{T}, n::Integer, add_segments::Bool, max_segments) where {T <: AbstractFloat} + out = Vector{T}(undef, n) + __sample!(rng, out, s, add_segments, max_segments) + return out +end +__sample!(s::ARSampler{T}, n::Integer, add_segments::Bool, max_segments) where {T <: AbstractFloat} = __sample!(default_rng(), s, n, add_segments, max_segments) + + +""" + sample!([rng=default_rng()], s::ARSampler, n::Integer; add_segments::Bool=true, max_segments::Integer = DEFAULT_MAX_SEGMENTS) + sample!([rng=default_rng()], v::AbstractVector, s::ARSampler; add_segments::Bool=true, max_segments::Integer = DEFAULT_MAX_SEGMENTS) + +Draw samples from `s`. If supplied, a vector `v` will be filled with samples. +Otherwise the number of samples is specified with `n`. +""" +function sample! end + +function sample!(rng::AbstractRNG, s::ARSampler{T}, n::Integer; add_segments::Bool = true, max_segments = DEFAULT_MAX_SEGMENTS) where {T <: AbstractFloat} + return __sample!(rng, s, n, add_segments, max_segments) +end + +function sample!(s::ARSampler{T}, n::Integer; add_segments::Bool = true, max_segments = DEFAULT_MAX_SEGMENTS) where {T <: AbstractFloat} + sample!(default_rng(), s, n, add_segments = add_segments, max_segments = max_segments) +end + +function sample!(rng::AbstractRNG, v::AbstractVector{T}, s::ARSampler{T}; add_segments::Bool = true, max_segments = DEFAULT_MAX_SEGMENTS) where {T} + __sample!(rng, v, s, add_segments, max_segments) + return nothing +end + +function sample!(v::AbstractVector{T}, s::ARSampler{T}; add_segments::Bool = true, max_segments = DEFAULT_MAX_SEGMENTS) where {T} + sample!(default_rng(), v, s, add_segments = add_segments, max_segments = max_segments) +end diff --git a/test/Project.toml b/test/Project.toml new file mode 100644 index 0000000..db8678f --- /dev/null +++ b/test/Project.toml @@ -0,0 +1,12 @@ +[deps] +AdaptiveRejectionSampling = "c75e803d-635f-53bd-ab7d-544e482d8c75" +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" +Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" +Supposition = "5a0628fe-1738-4658-9b6d-0b7605a9755b" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/test/wrappedallocs.jl b/test/wrappedallocs.jl new file mode 100644 index 0000000..0c24a7d --- /dev/null +++ b/test/wrappedallocs.jl @@ -0,0 +1,35 @@ +# https://github.com/JuliaAlgebra/TypedPolynomials.jl/blob/master/test/runtests.jl +""" + @wrappedallocs(expr) + +Given an expression, this macro wraps that expression inside a new function +which will evaluate that expression and measure the amount of memory allocated +by the expression. Wrapping the expression in a new function allows for more +accurate memory allocation detection when using global variables (e.g. when +at the REPL). + +For example, `@wrappedallocs(x + y)` produces: + +```julia +function g(x1, x2) + @allocated x1 + x2 +end +g(x, y) +``` + +You can use this macro in a unit test to verify that a function does not +allocate: + +``` +@test @wrappedallocs(x + y) == 0 +``` +""" +macro wrappedallocs(expr) + argnames = [gensym() for a in expr.args] + quote + function g($(argnames...)) + @allocated $(Expr(expr.head, argnames...)) + end + $(Expr(:call, :g, [esc(a) for a in expr.args]...)) + end +end From eab2d2b0b9709e3918b226b0de1df1c2cf3e85f6 Mon Sep 17 00:00:00 2001 From: Eliassj Date: Wed, 8 Apr 2026 17:17:30 +0200 Subject: [PATCH 4/8] Actually stage readme --- README.md | 87 +++-- src/AdaptiveRejectionSampling.jl | 542 +------------------------------ 2 files changed, 56 insertions(+), 573 deletions(-) diff --git a/README.md b/README.md index 5461761..82b4a8c 100644 --- a/README.md +++ b/README.md @@ -12,43 +12,62 @@ This package is useful for efficient sampling from log-concave univariate densit ## Examples ```julia -using AdaptiveRejectionSampling -using Plots +using AdaptiveRejectionSampling: Objective, ARSampler, sample!, eval_hull, abscissae +using SpecialFunctions: gamma +using CairoMakie ``` - ### Sampling from a shifted normal distribution ```julia -# Define function to be sampled -μ, σ = 1.0, 2.0 -f(x) = exp(-0.5(x - μ)^2 / σ^2) / sqrt(2pi * σ^2) -support = (-Inf, Inf) +const μ, σ = 1.0, 2.0 -# Build the sampler and simulate 10,000 samples -sampler = RejectionSampler(f, support, max_segments = 5) -@time sim = run_sampler!(sampler, 10000); +function bench() + # Define function to be sampled + f(x) = log(exp(-0.5(x - μ)^2 / σ^2) / sqrt(2pi * σ^2)) + support = (-10., 10.) + + # Build the sampler and simulate 10,000 samples + obj = Objective(f) + sampler = ARSampler(obj, support) + b = @be deepcopy(sampler) sample!(_, 10000, true, 10); +return sampler, b +end +s, b = bench(); +b ``` - 0.010434 seconds (192.15 k allocations: 3.173 MiB) - + Benchmark: 126 samples with 1 evaluation + min 616.951 μs (17 allocs: 80.523 KiB) + median 641.532 μs (17 allocs: 80.523 KiB) + mean 749.812 μs (17.10 allocs: 80.625 KiB, 0.75% gc time) + max 13.744 ms (20 allocs: 83.711 KiB, 94.06% gc time) Let's verify the result ```julia # Plot the results and compare to target distribution +f(x) = log(exp(-0.5(x - μ)^2 / σ^2) / sqrt(2pi * σ^2)) x = range(-10.0, 10.0, length=100) -envelop = [eval_envelop(sampler.envelop, xi) for xi in x] -target = [f(xi) for xi in x] - -histogram(sim, normalize = true, label = "Histogram") -plot!(x, [target envelop], width = 2, label = ["Normal(μ, σ)" "Envelop"]) +envelop = eval_hull.(s.upper_hull, x) +target = f.(x) +samples = sample!(s, 10000, max_segments = 5) + +fig, ax, p = hist(samples, bins=25, normalization = :pdf, label = "Samples"); +lines!(ax, -10..10, x -> exp(f(x)), label = "Target", color = :orange) +lines!(ax, x, exp.(envelop), label = "Envelop", color = :black, alpha = 0.8) +absc = ARS.abscissae(s.upper_hull) +scatter!(ax, absc, exp.(eval_hull.(s.upper_hull, absc)), label = "Abscissae", color = :red) +axislegend(ax) +fig +# histogram(sim, normalize = true, label = "Histogram") +# plot!(x, [target envelop], width = 2, label = ["Normal(μ, σ)" "Envelop"]) ``` -![](img/example1.png) +![](img/example1.svg) ### Let's try a Gamma @@ -59,24 +78,28 @@ plot!(x, [target envelop], width = 2, label = ["Normal(μ, σ)" "Envelop"]) f(x) = β^α * x^(α-1) * exp(-β*x) / gamma(α) support = (0.0, Inf) -# Build the sampler and simulate 10,000 samples -sampler = RejectionSampler(f, support) -@time sim = run_sampler!(sampler, 10000) - -# Plot the results and compare to target distribution -x = range(0.0, 5.0, length=100) -envelop = [eval_envelop(sampler.envelop, xi) for xi in x] -target = [f(xi) for xi in x] - -histogram(sim, normalize = true, label = "Histogram") -plot!(x, [target envelop], width = 2, label = ["Gamma(α, β)" "Envelop"]) +obj = Objective(x -> log(f(x))) +sam = ARSampler(obj, support) +@time sim = sample!(sam, 10000, max_segments = 5) + +# Verify result +x = range(0.0, 8.0, length=100) +envelop = eval_hull.(sam.upper_hull, x) +target = f.(x) +mx = maximum(sim) +fig, ax, p = hist(sim, bins=50, normalization = :pdf, label = "Samples"); +lines!(ax, 0..mx, y -> f(y), label = "Target", color = :orange) +lines!(ax, 0..mx, y -> exp(eval_hull(sam.upper_hull, y)), label = "Envelop", color = :black, alpha = 0.8) +absc = abscissae(sam.upper_hull) +scatter!(ax, absc, exp.(eval_hull.(sam.upper_hull, absc)), label = "Abscissae", color = :red) +axislegend(ax) +fig ``` - 0.007299 seconds (182.00 k allocations: 3.027 MiB) - + 0.000772 seconds (100 allocations: 85.211 KiB) -![](img/example2.png) +![](img/example2.svg) ### Truncated distributions and unknown normalization constant diff --git a/src/AdaptiveRejectionSampling.jl b/src/AdaptiveRejectionSampling.jl index c295272..43cad67 100644 --- a/src/AdaptiveRejectionSampling.jl +++ b/src/AdaptiveRejectionSampling.jl @@ -11,546 +11,6 @@ using DifferentiationInterface import ForwardDiff using Compat -const DEFAULT_MAX_SEGMENTS::Int = 25 - - -""" - AllocFreeWeights{S <: Real, T <: Number, V <: AbstractVector{T}} <: AbstractWeights{S, T, V} - -Non-mutable weights used in the `sample` function in this package in order to avoid allocations. -""" -struct AllocFreeWeights{S <: Real, T <: Number, V <: AbstractVector{T}} <: AbstractWeights{S, T, V} - values::V - sum::S - - function AllocFreeWeights(ws::V, sum::S) where {S <: Real, T <: Number, V <: AbstractVector{T}} - return new{S, T, V}(ws, sum) - end -end - -function AllocFreeWeights(ws::V) where {T <: Number, V <: AbstractVector{T}} - return AllocFreeWeights(ws, sum(ws)) -end - - -struct Objective{T <: Number, F <: Function, G <: Function} - "The (log-concave) function defining the density `f(x) -> y`" - f::F - "Its gradient `grad(x) -> y'`" - grad::G - - @doc """ - Objective(f::Function, grad::Function, ::Type{T} = Float64) where {T <: Number} - - Create an `Objective` directly defined by its function `f` and custom gradient `grad`. - The parameter `T` represents the type of the expected input and is utilized when preparing gradients using autodiff. - By default `T = Float64`. - - - !!! warning - Observe that `f` should be in its - log-concave form and that no checks are performed in order to verify this. - """ - function Objective(f::Function, grad::Function, ::Type{T} = Float64) where {T <: Number} - return new{T, typeof(f), typeof(grad)}(f, grad) - end -end - -""" - Objective(f::Function, init = one(Float64); adbackend = AutoForwardDiff()) - -Create an `Objective` for a function automatically generating its gradient. -Gradients are calculated using `DifferentiationInterface.jl` using the backend of choice -with `ForwardDiff.jl` being the default. In order to prepare the gradient an initial value is -required. By default this is `one(Float64)`. If a gradient for a different type is -desired, it should be specified through this initial value. - -If a custom/manual gradient is available it may instead be provided. - -!!! warning - Observe that `f` should be in its - log-concave form and that no checks are performed in order to verify this. - -# Example - -```julia -# Create an objective using `Mooncake.jl` autodiff and Float32 as its type. -ARS.Objective(somefun, one(Float32); adbackend = AutoMooncake(; config=nothing)) -# Create an objective using the defaults `AutoDiff.jl` autodiff and Float64 as its type. -ARS.Objective(somefun) -``` -""" -function Objective(f::Function, init = one(Float64); adbackend = AutoForwardDiff()) - gradprep = prepare_gradient(f, adbackend, init) - obj = Objective( - f, - x -> gradient(f, gradprep, adbackend, x), - typeof(init) - ) - return obj -end - -abstract type AbstractHull end - -""" - struct UpperHull{T} <: AbstractHull - intercepts::Vector{T} - slopes::Vector{T} - intersections::Vector{T} - abscissae::Vector{T} - segment_weights::Vector{T} - domain::Tuple{T, T} - end - -Upper hull of the enveloping function. -""" -struct UpperHull{T} <: AbstractHull - intercepts::Vector{T} - slopes::Vector{T} - intersections::Vector{T} - abscissae::Vector{T} - segment_weights::Vector{T} - domain::Tuple{T, T} -end - -Base.broadcastable(h::UpperHull) = Ref(h) - -slopes(h::UpperHull) = h.slopes -intercepts(h::UpperHull) = h.intercepts -intersections(h::UpperHull) = h.intersections -abscissae(h::UpperHull) = h.abscissae -line(h::UpperHull, i::Integer) = (h.slopes[i], h.intercepts[i]) -lineinds(h::UpperHull) = 1:(length(slopes(h))) -n_lines(h::UpperHull) = length(slopes(h)) -segment_weights(h::UpperHull) = h.segment_weights - -""" - eval_hull(h::UpperHull, x) - -Eval hull `h` at `x`. -""" -function eval_hull(h::UpperHull, x) - i = searchsortedfirst(intersections(h), x) - 1 - sl, int = line(h, i) - return sl * x + int -end - - -""" - intersection(slope1::T, intercept1::T, slope2::T, intercept2::T) where {T} - -Returns the intersection abscissa between 2 lines as defined by their slopes and intercepts. Returns NaN if the lines are paralell. -""" -function intersection(slope1::T, intercept1::T, slope2::T, intercept2::T) where {T} - return (intercept2 - intercept1) / (slope1 - slope2) -end - -""" - calc_intersects!(out::AbstractVector{T}, slopes::AbstractVector{T}, intercepts::AbstractVector{T}) where {T} - -Calculate intersections of a series of lines defined by their `slopes` and `intercepts`, storing the result in `out`. -""" -function calc_intersects!(out::AbstractVector{T}, slopes::AbstractVector{T}, intercepts::AbstractVector{T}) where {T} - for i in eachindex(out) - out[i] = intersection(slopes[i], intercepts[i], slopes[i + 1], intercepts[i + 1]) - end - return nothing -end - -""" - calc_intersects(slopes::AbstractVector{T}, intercepts::AbstractVector{T}) where {T} - -Calculate intersections of a series of lines defined by their `slopes` and `intercepts`. -""" -function calc_intersects(slopes::AbstractVector{T}, intercepts::AbstractVector{T}) where {T} - # return [ - # intersection(slopes[i], intercepts[i], slopes[i + 1], intercepts[i + 1]) - # for i in 1:(length(slopes) - 1) - # ] - - out = Vector{T}(undef, length(slopes) - 1) - calc_intersects!(out, slopes, intercepts) - return out -end - - -""" - calc_slopes_and_intercepts(obj::Objective, abscissae::AbstractVector{T}) where {T} - -Calculates slopes and intercepts of the lines tangential to the objective function at `abscissae`. -""" -function calc_slopes_and_intercepts(obj::Objective, abscissae::AbstractVector{T}) where {T} - issorted(abscissae) || - throw(ArgumentError("`abscissae` should be sorted in ascending order.")) - slopes = [obj.grad(abscissae[i]) for i in eachindex(abscissae)] - intercepts = [ - slopes[i] * (-abscissae[i]) + obj.f(abscissae[i]) - for i in eachindex(abscissae) - ] - return slopes, intercepts -end - - -""" - exp_integral_line(slope, intercept, x1, x2) - -Calculate the integral of `exp(slope * x + intercept)` between `x1` and `x2`. -""" -function exp_integral_line(slope, intercept, x1, x2) - return if !iszero(slope) - (exp(slope * x2 + intercept) - exp(slope * x1 + intercept)) / slope - else - (x2 - x1) * exp(intercept) - end -end - -""" - UpperHull(obj::Objective, abscissae::Vector{T}, domain::Tuple{T, T}) where {T} - -Create an `UpperHull` based on the objective function, initial abscissae and its domain. -""" -function UpperHull(obj::Objective{T}, abscissae::Vector{T}, domain::Tuple{T, T}) where {T} - slopes, intercepts = calc_slopes_and_intercepts(obj, abscissae) - - intersects = [domain[1]; calc_intersects(slopes, intercepts); domain[2]] - - wgts = Vector{T}(undef, length(slopes)) - for i in 1:length(wgts) - wgts[i] = exp_integral_line(slopes[i], intercepts[i], intersects[i], intersects[i + 1]) - end - - return UpperHull(intercepts, slopes, intersects, abscissae, wgts, domain) -end - -# Calculate the inverse CDF of a segment, handling a slope of zero. -""" - inv_cdf_seg(slope, intercept, r, w, intersect, intersect2) - -Calculate the inverse cumulative density of a hull segment at `r` based on its... - -- `slope` -- `intercept` -- `w` (weight) -- `intersect`, `intersect2`, its intersections with previous/next segments -""" -@inline function inv_cdf_seg(slope, intercept, r, w, intersect, intersect2) - return if !iszero(slope) - log(exp(-intercept) * r * w * slope + exp(slope * intersect)) / slope - else - r * (intersect2 - intersect) + intersect - end -end - -# Draw a single sample from `h`. -function sample_hull(rng::AbstractRNG, h::UpperHull) - ws = segment_weights(h) - ind = sample(rng, AllocFreeWeights(ws, sum(ws))) - sl, int = line(h, ind) - return inv_cdf_seg(sl, int, rand(rng), segment_weights(h)[ind], intersections(h)[ind], intersections(h)[ind + 1]) -end -sample_hull(h::UpperHull) = sample_hull(default_rng(), h) - - -# Draw `n` samples from `h`, storing the result in `out`. -function sample_hull!(rng::AbstractRNG, out::AbstractVector{T}, h::UpperHull{T}, n::Integer) where {T} - seg_wgts = segment_weights(h) - inds = sample(rng, 1:n_lines(h), AllocFreeWeights(seg_wgts, sum(seg_wgts)), n) - rands = rand(rng, n) - for i in eachindex(inds, out, rands) - ind = inds[i] - sl, int = line(h, ind) - out[i] = inv_cdf_seg(sl, int, rands[i], segment_weights(h)[ind], intersections(h)[ind], intersections(h)[ind + 1]) - end - return -end -sample_hull!(out::AbstractVector{T}, h::UpperHull{T}, n::Integer) where {T} = sample_hull!(default_rng(), out, h, n) - -function sample_hull!(rng, out::AbstractVector{T}, h::UpperHull) where {T} - return sample_hull!(rng, out, h, length(out)) -end -sample_hull!(out::AbstractVector{T}, h::UpperHull) where {T} = sample_hull!(default_rng(), out, h) - -# Draw `n` samples from `h`. -function sample_hull(rng::AbstractRNG, h::UpperHull{T}, n::Integer) where {T} - out = Vector{T}(undef, n) - sample_hull!(rng, out, h, n) - return out -end -sample_hull(h::UpperHull, n::Integer) = sample_hull(default_rng(), h, n) - -struct LowerHull{T} - intercepts::Vector{T} - slopes::Vector{T} - intersections::Vector{T} -end - - -Base.broadcastable(h::LowerHull) = Ref(h) - - -intercepts(h::LowerHull) = h.intercepts -slopes(h::LowerHull) = h.slopes -intersections(h::LowerHull) = h.intersections -line(h::LowerHull, i) = (h.slopes[i], h.intercepts[i]) - -function calc_lower_slopes_and_intercepts!( - slout::AbstractVector{T}, intout::AbstractVector{T}, - intersections::AbstractVector{T}, f::Function - ) where {T} - - for i in eachindex(slout, intout) - slout[i] = (f(intersections[i + 1]) - f(intersections[i])) / - (intersections[i + 1] - intersections[i]) - intout[i] = slout[i] * (-intersections[i]) + f(intersections[i]) - end - return nothing -end - - -function LowerHull(upper::UpperHull{T}, obj::Objective{T}) where {T} - intersections = abscissae(upper) # NOTE: This makes them alias each other! - n_segs = length(intersections) - 1 - sl = Vector{T}(undef, n_segs) - int = Vector{T}(undef, n_segs) - - calc_lower_slopes_and_intercepts!(sl, int, intersections, obj.f) - - return LowerHull(int, sl, intersections) -end - -function eval_hull(h::LowerHull, x) - i = searchsortedfirst(intersections(h), x) - if isone(i) || i == lastindex(intersections(h)) + 1 - return -Inf - end - sl, int = line(h, i - 1) - return sl * x + int -end - - -""" - ARSampler{T, F, G} - -Adaptive rejection sampler containing the objective function, its gradient and the -upper/lower hull of the piecewise linear envelope. -""" -struct ARSampler{T, F, G} - objective::Objective{T, F, G} - upper_hull::UpperHull{T} - lower_hull::LowerHull{T} -end - -""" - ARSampler( - obj::Objective{T, F, G}, - domain::Tuple{T, T}, - initial_points::Vector{T} - ) where {T <: AbstractFloat, F <: Function, G <: Function} - - ARSampler( - obj::Objective{T, F, G}, - domain::Tuple{T, T}, - search_range::Tuple{T, T} = (-10.0, 10.0), - search_step = 0.1 - ) where {T <: AbstractFloat, F <: Function, G <: Function} - -Initialize an adaptive rejection sampler over an objective function from `obj` ([`ARSampling.Objective`](@ref)). -`initial_points` should be a vector defining the abscissae of the initial segments of the sampler. - At least 2 of the points should be on opposite sides of the objective function's maximum. - -If no `initial_points` an attempt to find suitable ones will be made by searching for -the first negative/positive slope of `obj`. The default search range is `-10:10` with a -step of `0.1`. - -!!! warning - As it currently stands, finding inital points will fail for distributions bounded - to the left/right of their maximum. In these cases the initial point(s) need to - be provided manually (however, only a single initial point has to be provided). -""" -function ARSampler end - -function ARSampler( - obj::Objective{T, F, G}, - domain::Tuple{T, T}, - initial_points::Vector{T} - ) where {T <: AbstractFloat, F <: Function, G <: Function} - - u = UpperHull(obj, initial_points, domain) - l = LowerHull(u, obj) - - return ARSampler{T, F, G}(obj, u, l) -end - - -function ARSampler(obj::Objective{T, F, G}, domain::Tuple{<:Number, <:Number}, search_range::Tuple{<:Number, <:Number} = (-10.0, 10.0); search_step = 0.5) where {T <: AbstractFloat, F <: Function, G <: Function} - domain_T = T.(domain) - search_range_T = T.(search_range) - lbs = T(max(domain_T[1], search_range_T[1])) - ubs = T(min(domain_T[2], search_range_T[2])) - search_step_T = T(search_step) - initial_points = find_initial_points(obj, lbs, ubs, search_step_T) - - return ARSampler(obj, domain_T, initial_points) -end - -""" - find_initial_points(obj, lbs, ubs, search_step) - -Attempts to find suitable initial points for abscissae based on the objective functions by looking for -the first points with positive and negative slopes respectively. Simply searches between the bounds specified -by `lbs` and `ubs` with step size `search_step`. -""" -function find_initial_points(obj, lbs, ubs, search_step) - grad = obj.grad - r = lbs:search_step:ubs - g = (grad(x) for x in r) - l = findfirst(>(0), g) - u = findfirst(<(0), g) - isnothing(l) && throw( - ErrorException( - lazy""" - Could not find the left initial segment. - Search range: $(lbs) -- $(ubs) - Search step: $(search_step) - """ - ) - ) - isnothing(u) && throw( - ErrorException( - lazy""" - Could not find the right initial segment. - Search range: $(lbs) -- $(ubs) - Search step: $(search_step) - """ - ) - ) - return [r[l], r[u]] -end - - -function n_segments(s::ARSampler) - return n_lines(s.upper_hull) -end - -# Adds a segment with abscissa at `x` to `s` -""" - add_segment!(s::ARSampler{T}, x::T) where {T <: AbstractFloat} - -Modifies the hulls of `s`, adding a segment with abscissa at `x`. -""" -function add_segment!(s::ARSampler{T}, x::T) where {T <: AbstractFloat} - - # Calculate slope, intercept and index of new segment - new_slope = s.objective.grad(x) - new_intercept = (new_slope * -x) + s.objective.f(x) - new_ind = searchsortedfirst(abscissae(s.upper_hull), x) - - # Insert new slope and intercept - all_inters = intersections(s.upper_hull) - all_slopes = slopes(s.upper_hull) - all_intercepts = intercepts(s.upper_hull) - insert!(all_slopes, new_ind, new_slope) - insert!(all_intercepts, new_ind, new_intercept) - - # Insert new abscissa - insert!(abscissae(s.upper_hull), new_ind, x) - - # TODO: Only calculate the intersections and weights that actually change. - # Recalculate intersection points for segments, given the new segment - push!(all_inters, zero(T)) # Extend intercepts by one - calc_intersects!(@view(all_inters[(begin + 1):(end - 1)]), all_slopes, all_intercepts) - - # Recalculate segment weights - all_weights = segment_weights(s.upper_hull) - push!(all_weights, zero(T)) - - for i in eachindex(all_weights) - all_weights[i] = exp_integral_line(all_slopes[i], all_intercepts[i], all_inters[i], all_inters[i + 1]) - end - - #= Lower hull time =# - # We don't need to add anything to the lower intersections as it is the same vector used for abscissae - # in the upper hull - - all_inters_lower = intersections(s.lower_hull) - lowerslopes = slopes(s.lower_hull) - lowerintercepts = intercepts(s.lower_hull) - push!(lowerslopes, zero(T)) - push!(lowerintercepts, zero(T)) - calc_lower_slopes_and_intercepts!( - lowerslopes, - lowerintercepts, - all_inters_lower, - s.objective.f - ) - - return nothing -end - -# TODO: Make this prepopulate `out` using `sample_hull!`? -# NOTE: The above might be inefficient initially if most of the samples get rejected anyway? -""" - __sample!(rng::AbstractRNG, out::AbstractVector{T}, s::ARSampler{T}, add_segments::Bool, max_segments) where {T<:AbstractFloat} - - -All other `__sample!` methods call this one. -""" -function __sample!(rng::AbstractRNG, out::AbstractVector{T}, s::ARSampler{T}, add_segments::Bool, max_segments) where {T <: AbstractFloat} - n = length(out) - n_accepted = 0 - while n_accepted < n - x = sample_hull(rng, s.upper_hull) - up = eval_hull(s.upper_hull, x) - lo = eval_hull(s.lower_hull, x) - w = rand(rng) - # Squeeze test - if w <= exp(lo - up) - n_accepted += 1 - out[n_accepted] = x - elseif w <= exp(s.objective.f(x) - up) - # Accept sample i - n_accepted += 1 - out[n_accepted] = x - if add_segments && n_segments(s) < max_segments - add_segment!(s, x) - end - elseif add_segments && n_segments(s) < max_segments - add_segment!(s, x) - end - end - return nothing -end -__sample!(out::Vector{T}, s::ARSampler{T}, add_segments::Bool, max_segments) where {T} = __sample!(default_rng(), out, s, add_segments, max_segments) - -# Draw `n` samples, returning a newly allocated vector. -function __sample!(rng::AbstractRNG, s::ARSampler{T}, n::Integer, add_segments::Bool, max_segments) where {T <: AbstractFloat} - out = Vector{T}(undef, n) - __sample!(rng, out, s, add_segments, max_segments) - return out -end -__sample!(s::ARSampler{T}, n::Integer, add_segments::Bool, max_segments) where {T <: AbstractFloat} = __sample!(default_rng(), s, n, add_segments, max_segments) - - -""" - sample!([rng=default_rng()], s::ARSampler, n::Integer, add_segments::Bool=true) - sample!([rng=default_rng()], v::AbstractVector, s::ARSampler, add_segments::Bool=true) - -Draw samples from `s`. If supplied, a vector `v` will be filled with samples. -Otherwise the number of samples is specified with `n`. - -""" -function sample! end - -function sample!(rng::AbstractRNG, s::ARSampler{T}, n::Integer, add_segments::Bool = true, max_segments = DEFAULT_MAX_SEGMENTS) where {T <: AbstractFloat} - return __sample!(rng, s, n, add_segments, max_segments) -end -sample!(s::ARSampler{T}, n::Integer, add_segments::Bool = true, max_segments = DEFAULT_MAX_SEGMENTS) where {T <: AbstractFloat} = sample!(default_rng(), s, n, add_segments, max_segments) - -function sample!(rng::AbstractRNG, v::AbstractVector{T}, s::ARSampler{T}, add_segments::Bool = true, max_segments = DEFAULT_MAX_SEGMENTS) where {T} - __sample!(rng, v, s, add_segments, max_segments) - return nothing -end -sample!(v::AbstractVector{T}, s::ARSampler{T}, add_segments::Bool = true, max_segments = DEFAULT_MAX_SEGMENTS) where {T} = sample!(default_rng(), v, s, add_segments, max_segments) - +include("sampling.jl") end # module AdaptiveRejectionSampling From 75fc4d5203b0a2ed23e7bff07ff2688e085fb26e Mon Sep 17 00:00:00 2001 From: Eliassj Date: Fri, 10 Apr 2026 16:14:58 +0200 Subject: [PATCH 5/8] Add recipe for plotting hull --- src/AdaptiveRejectionSampling.jl | 2 +- src/plot.jl | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) create mode 100644 src/plot.jl diff --git a/src/AdaptiveRejectionSampling.jl b/src/AdaptiveRejectionSampling.jl index 43cad67..344632e 100644 --- a/src/AdaptiveRejectionSampling.jl +++ b/src/AdaptiveRejectionSampling.jl @@ -12,5 +12,5 @@ import ForwardDiff using Compat include("sampling.jl") - +include("plot.jl") end # module AdaptiveRejectionSampling diff --git a/src/plot.jl b/src/plot.jl new file mode 100644 index 0000000..eac0238 --- /dev/null +++ b/src/plot.jl @@ -0,0 +1,4 @@ +# Function definitions for use in Makie extension + +function hullplot end +function hullplot! end From 3bd69464d5752f746677f79393b46bbc11315fbb Mon Sep 17 00:00:00 2001 From: Eliassj Date: Fri, 10 Apr 2026 16:15:10 +0200 Subject: [PATCH 6/8] Improve sampling from hull --- Project.toml | 7 +++++++ src/sampling.jl | 43 ++++++++++++++++++++++++++++++++++++++----- 2 files changed, 45 insertions(+), 5 deletions(-) diff --git a/Project.toml b/Project.toml index e33cf41..dc67627 100644 --- a/Project.toml +++ b/Project.toml @@ -10,14 +10,21 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" +[weakdeps] +Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a" + [compat] Compat = "4.18.1" DifferentiationInterface = "0.7" ForwardDiff = "1" +Makie = "0.24.9" SpecialFunctions = "2" StatsBase = "0.34.10" julia = "1" +[extensions] +MakieExtension = "Makie" + [extras] Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/src/sampling.jl b/src/sampling.jl index 784e24a..357d57b 100644 --- a/src/sampling.jl +++ b/src/sampling.jl @@ -97,6 +97,7 @@ struct UpperHull{T} <: AbstractHull intersections::Vector{T} abscissae::Vector{T} segment_weights::Vector{T} + weights_cumsum::Vector{T} domain::Tuple{T, T} end @@ -110,6 +111,7 @@ line(h::UpperHull, i::Integer) = (h.slopes[i], h.intercepts[i]) lineinds(h::UpperHull) = 1:(length(slopes(h))) n_lines(h::UpperHull) = length(slopes(h)) segment_weights(h::UpperHull) = h.segment_weights +weights_cumsum(h::UpperHull) = h.weights_cumsum """ eval_hull(h::UpperHull, x) @@ -202,14 +204,22 @@ Create an `UpperHull` based on the objective function, initial abscissae and its function UpperHull(obj::Objective{T}, abscissae::Vector{T}, domain::Tuple{T, T}) where {T} slopes, intercepts = calc_slopes_and_intercepts(obj, abscissae) + intersects = [domain[1]; calc_intersects(slopes, intercepts); domain[2]] wgts = Vector{T}(undef, length(slopes)) for i in 1:length(wgts) wgts[i] = exp_integral_line(slopes[i], intercepts[i], intersects[i], intersects[i + 1]) end - - return UpperHull(intercepts, slopes, intersects, abscissae, wgts, domain) + wgts_cumsum = cumsum(wgts) + # Sizehinting + sizehint!(slopes, 50) + sizehint!(intercepts, 50) + sizehint!(intersects, 50) + sizehint!(wgts, 50) + sizehint!(wgts_cumsum, 50) + sizehint!(abscissae, 50) + return UpperHull(intercepts, slopes, intersects, abscissae, wgts, wgts_cumsum, domain) end # Calculate the inverse CDF of a segment, handling a slope of zero. @@ -232,9 +242,16 @@ Calculate the inverse cumulative density of a hull segment at `r` based on its.. end # Draw a single sample from `h`. -function sample_hull(rng::AbstractRNG, h::UpperHull) - ws = segment_weights(h) - ind = sample(rng, AllocFreeWeights(ws, sum(ws))) +function sample_hull(rng::AbstractRNG, h::UpperHull{T}) where {T} + s = weights_cumsum(h) + u = rand(rng, T) * s[end] + k = length(s) + + ind = 1 + while u > s[ind] && ind < k + ind += 1 + end + sl, int = line(h, ind) return inv_cdf_seg(sl, int, rand(rng), segment_weights(h)[ind], intersections(h)[ind], intersections(h)[ind + 1]) end @@ -305,6 +322,10 @@ function LowerHull(upper::UpperHull{T}, obj::Objective{T}) where {T} calc_lower_slopes_and_intercepts!(sl, int, intersections, obj.f) + # Sizehinting + sizehint!(sl, 50) + sizehint!(int, 50) + return LowerHull(int, sl, intersections) end @@ -445,10 +466,13 @@ function add_segment!(s::ARSampler{T}, x::T) where {T <: AbstractFloat} all_slopes = slopes(s.upper_hull) all_intercepts = intercepts(s.upper_hull) insert!(all_slopes, new_ind, new_slope) + # simple_insert!(all_slopes, new_ind, new_slope) insert!(all_intercepts, new_ind, new_intercept) + # simple_insert!(all_intercepts, new_ind, new_intercept) # Insert new abscissa insert!(abscissae(s.upper_hull), new_ind, x) + # simple_insert!(abscissae(s.upper_hull), new_ind, x) # TODO: Only calculate the intersections and weights that actually change. # Recalculate intersection points for segments, given the new segment @@ -463,6 +487,12 @@ function add_segment!(s::ARSampler{T}, x::T) where {T <: AbstractFloat} all_weights[i] = exp_integral_line(all_slopes[i], all_intercepts[i], all_intersections[i], all_intersections[i + 1]) end + # Recalculate cumsum of weights + w_cumsum = weights_cumsum(s.upper_hull) + push!(w_cumsum, zero(T)) + cumsum!(w_cumsum, all_weights) + + #= Lower hull time =# # We don't need to add anything to the lower intersections as it is the same vector used for abscissae # in the upper hull @@ -482,6 +512,9 @@ function add_segment!(s::ARSampler{T}, x::T) where {T <: AbstractFloat} return nothing end +function __sample_single!(rng::AbstractRNG, s::ARSampler{T}, add_segments::Bool, max_segments::Integer) where {T <: AbstractFloat} +end + # TODO: Make this prepopulate `out` using `sample_hull!`? # NOTE: The above might be inefficient initially if most of the samples get rejected anyway? """ From 19ab62859702558967ab0ea7f8eceebe5c6670d0 Mon Sep 17 00:00:00 2001 From: Eliassj Date: Fri, 10 Apr 2026 16:50:34 +0200 Subject: [PATCH 7/8] Commit actual plotting code --- ext/MakieExtension.jl | 81 +++++++++++++++++++++++++++++++++++++++++++ src/sampling.jl | 19 +++++++++- 2 files changed, 99 insertions(+), 1 deletion(-) create mode 100644 ext/MakieExtension.jl diff --git a/ext/MakieExtension.jl b/ext/MakieExtension.jl new file mode 100644 index 0000000..88639ee --- /dev/null +++ b/ext/MakieExtension.jl @@ -0,0 +1,81 @@ +module MakieExtension +using Makie +import Makie: mixin_colormap_attributes +using AdaptiveRejectionSampling +import AdaptiveRejectionSampling: hullplot, hullplot!, eval_hull, abscissae, ARSampler + +@recipe Hullplot begin + abscissae = true + target = true + upper_hull = true + lower_hull = true + target_linewidth = @inherit linewidth + upper_hull_linewidth = @inherit linewidth + lower_hull_linewidth = @inherit linewidth + target_label = "Target" + upper_hull_label = "Upper hull" + lower_hull_label = "Lower hull" + abscissae_label = "Abscissae" +end + + +function Makie.plot!(p::Hullplot{<:Tuple{<:Any, <:ARSampler}}) + + sam = p[2][] + upper_hull = sam.upper_hull + lower_hull = sam.lower_hull + obj = sam.objective + xs = p[1] + + i = 1 + if p.lower_hull[] + lines!( + p, + xs, + x -> exp(eval_hull(lower_hull, x)), + color=Cycled(i), + linewidth = p.lower_hull_linewidth, + label = p.lower_hull_label + ) + i += 1 + end + if p.upper_hull[] + lines!( + p, + xs, + x -> exp(eval_hull(upper_hull, x)), + linewidth = p.upper_hull_linewidth, + color=Cycled(i), + label = p.upper_hull_label + ) + i += 1 + end + if p.target[] + lines!( + p, + xs, + x -> exp(obj.f(x)), + color=Cycled(i), + label = p.target_label, + linewidth = p.target_linewidth + ) + i += 1 + end + if p.abscissae[] + absc = abscissae(upper_hull) + scatter!( + p, + absc, + exp.(eval_hull.(upper_hull, absc)), + color = Cycled(i), + label = p.abscissae_label + ) + end + +end + +function Makie.get_plots(plot::Hullplot) + return plot.plots +end + +end # End of module diff --git a/src/sampling.jl b/src/sampling.jl index 357d57b..8011ec6 100644 --- a/src/sampling.jl +++ b/src/sampling.jl @@ -512,7 +512,24 @@ function add_segment!(s::ARSampler{T}, x::T) where {T <: AbstractFloat} return nothing end -function __sample_single!(rng::AbstractRNG, s::ARSampler{T}, add_segments::Bool, max_segments::Integer) where {T <: AbstractFloat} +function __sample_single!(rng::AbstractRNG, s::ARSampler{T}, n_accepted::Integer, add_segments::Bool, max_segments::Integer) where {T <: AbstractFloat} + while true + x = sample_hull(rng, s.upper_hull) + up = eval_hull(s.upper_hull, x) + lo = eval_hull(s.lower_hull, x) + w = rand(rng) + + if w <= exp(lo - up) + return x + elseif w <= exp(s.objective.f(x) - up) + if add_segments && n_segments(s) < max_segments + add_segment!(s, x) + end + return x + elseif add_segments && n_segments(s) < max_segments + add_segment!(s, x) + end + end end # TODO: Make this prepopulate `out` using `sample_hull!`? From 382c4a52040baffce6d9fe874c609a5fb654b40b Mon Sep 17 00:00:00 2001 From: Eliassj Date: Fri, 10 Apr 2026 22:24:00 +0200 Subject: [PATCH 8/8] Add tests for lb/ub --- src/sampling.jl | 1 - test/runtests.jl | 21 ++++++++++++++++++++- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/src/sampling.jl b/src/sampling.jl index 8011ec6..22c21df 100644 --- a/src/sampling.jl +++ b/src/sampling.jl @@ -455,7 +455,6 @@ end Modifies the hulls of `s`, adding a segment with abscissa at `x`. """ function add_segment!(s::ARSampler{T}, x::T) where {T <: AbstractFloat} - # Calculate slope, intercept and index of new segment new_slope = s.objective.grad(x) new_intercept = (new_slope * -x) + s.objective.f(x) diff --git a/test/runtests.jl b/test/runtests.jl index 83ee1a9..7cf91c4 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -42,7 +42,7 @@ obj = Objective(fn) x = 2.4 const f64gen = Data.Floats{Float64}(nans = false) -const f64genmu = Data.Floats{Float64}(nans = false) +const f64genmu = Data.Floats{Float64}(nans = false, minimum = -1.0e8, maximum = 1.0e8) const f64gensigma = Data.Floats{Float64}(nans = false, minimum = 0.0001) @testset "Objective" begin @@ -60,6 +60,25 @@ const f64gensigma = Data.Floats{Float64}(nans = false, minimum = 0.0001) end +@testset "Truncated" begin + f_gamma(x, α, β) = β^α * x^(α - 1) * exp(-β * x) / gamma(α) + f_log_gamma(x, α, β) = α * log(β) + (α - 1) * log(x) - β * x - loggamma(α) + obj = Objective(x -> f_log_gamma(x, 4, 2)) + lb = 4.0 + ub = 8.0 + sam = ARSampler(obj, (lb, ub), [lb, ub]) + v = sample!(sam, 10_000) + @test all(lb .< v .< ub) + + obj = Objective(x -> f_log_gamma(x, 4, 2)) + lb = 9.5 + ub = 10.0 + sam = ARSampler(obj, (lb, ub), [(lb + ub) / 2]) + v = sample!(sam, 10_000) + @test all(lb .< v .< ub) +end + + const n_samples = 100_000 @testset "Autodiff backends" begin