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/Project.toml b/Project.toml index 3e30ade..dc67627 100644 --- a/Project.toml +++ b/Project.toml @@ -3,14 +3,27 @@ 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" +[weakdeps] +Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a" + [compat] -ForwardDiff = "0.10.1" -StatsBase = "0.24, 0.34" -julia = "0.7.0, 1" +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/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/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/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/AdaptiveRejectionSampling.jl b/src/AdaptiveRejectionSampling.jl index 386f0f8..344632e 100644 --- a/src/AdaptiveRejectionSampling.jl +++ b/src/AdaptiveRejectionSampling.jl @@ -1,284 +1,16 @@ - """ 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 -end - -""" - intersection(l1::Line, l2::Line) -Finds the horizontal coordinate of the intersection between lines -""" -function intersection(l1::Line, l2::Line) - @assert l1.slope != l2.slope "slopes should be different" - -(l2.intercept - l1.intercept) / (l2.slope - l1.slope) -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 - end - exp(b) * (exp(v2) - exp(v1)) / a -end -""" - 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 -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. -""" -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" - 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)] -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. -""" -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 -end - -function sample_envelop(e::Envelop) - sample_envelop(Random.GLOBAL_RNG, e) -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`. -""" -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 - else - a, b = e.lines[pos-1].slope, e.lines[pos-1].intercept - return exp(a * x + b) - 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. -""" -struct Objective - logf::Function - grad::Function - Objective(logf::Function) = begin - # Automatic differentiation - grad(x) = derivative(logf, x) - new(logf, grad) - end - Objective(logf::Function, grad::Function) = new(logf, grad) -end - -""" -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) - 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 -end - -""" - 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)) - end - failed += 1 - @assert failed < max_failed "max_failed_factor reached" - end - end - out -end +import SpecialFunctions: loggamma +using StatsBase +import Random: default_rng, AbstractRNG +using DifferentiationInterface +import ForwardDiff +using Compat -function run_sampler!(s::RejectionSampler, n::Int) - run_sampler!(Random.GLOBAL_RNG, s, n) -end +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 diff --git a/src/sampling.jl b/src/sampling.jl new file mode 100644 index 0000000..22c21df --- /dev/null +++ b/src/sampling.jl @@ -0,0 +1,602 @@ +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} + weights_cumsum::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 +weights_cumsum(h::UpperHull) = h.weights_cumsum + +""" + 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 + 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. +""" + 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{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 +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) + + # Sizehinting + sizehint!(sl, 50) + sizehint!(int, 50) + + 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) + # 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 + 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 + + # 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 + + 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 + +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!`? +# 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/runtests.jl b/test/runtests.jl index 75e247f..7cf91c4 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,81 +1,244 @@ -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 +using Supposition +using Supposition: @check + +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 + +const f64gen = 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 - @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 + @test obj isa Objective{Float64, <:Function, <:Function} + @test obj.grad(x) ≈ score_normal(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 - @test begin - objective = ars.Objective(x -> exp(-0.5 * x^2) / sqrt(2pi)) - objective.logf(Inf) == 0.0 && objective.logf(Inf) == 0.0 + @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 -@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 + +@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 + @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 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